diff --git a/pipelineprobe/cli.py b/pipelineprobe/cli.py index e17d4ae..97a54c0 100644 --- a/pipelineprobe/cli.py +++ b/pipelineprobe/cli.py @@ -1,6 +1,13 @@ -import typer -import os +import asyncio +import json import logging +import os +from pathlib import Path + +import httpx +import typer + +from pipelineprobe import __version__ from pipelineprobe.config import load_config from pipelineprobe.connectors.airflow import AirflowConnector from pipelineprobe.connectors.dbt import DbtConnector @@ -12,19 +19,20 @@ logger = logging.getLogger(__name__) -from pipelineprobe import __version__ def version_callback(value: bool): if value: typer.echo(f"PipelineProbe v{__version__}") raise typer.Exit() + app = typer.Typer( name="pipelineprobe", help="Instant Data Pipeline Audit Report for Airflow + dbt + modern warehouses", add_completion=False, ) + @app.callback() def main( version: bool = typer.Option( @@ -78,18 +86,19 @@ def audit( warehouse_conn = PostgresConnector(cfg.warehouse) typer.echo("Fetching data from source systems...") - # Fetch top-level DAGs airflow_dags = airflow_conn.get_dags() - airflow_tasks = [] + airflow_tasks: list = [] - # Wire the actual dag runs and tasks for every returned DAG if airflow_dags: typer.echo( - f"Found {len(airflow_dags)} Airflow DAGs. Fetching runs and tasks..." + f"Found {len(airflow_dags)} Airflow DAGs. " + f"Fetching runs and tasks concurrently (concurrency={cfg.rules.fetch_concurrency})..." + ) + airflow_dags, airflow_tasks = asyncio.run( + airflow_conn.fetch_dag_details( + airflow_dags, concurrency=cfg.rules.fetch_concurrency + ) ) - for dag in airflow_dags: - dag.recent_runs = airflow_conn.get_dag_runs(dag.id) - airflow_tasks.extend(airflow_conn.get_tasks(dag.id)) dbt_models = dbt_conn.get_models() warehouse_tables = warehouse_conn.get_stats_sync() @@ -100,27 +109,59 @@ def audit( "dbt_models": dbt_models, "warehouse_tables": warehouse_tables, "warehouse_type": cfg.warehouse.type, + # Rule-level configuration injected into context so rules stay stateless + "rule_severity_overrides": cfg.rules.severity_overrides, + "stale_threshold_days": cfg.rules.stale_threshold_days, } typer.echo("Running rule engine...") engine = get_configured_engine() issues = engine.run(context) - # Compute a dummy summary for MVP + # ----------------------------------------------------------------------- + # Health score — normalized by DAG count so that five criticals on a + # 500-DAG shop is a very different signal than five on a 10-DAG shop. + # + # Formula: + # critical_density = critical_count / dag_count (criticals per DAG) + # warning_density = warning_count / dag_count (warnings per DAG) + # + # critical_penalty = min(90, critical_density * 200) + # → 45 % critical density → 90-pt penalty (score ≤ 10 before warnings) + # → 1 % critical density → 2-pt penalty (barely dents the score) + # + # warning_penalty = min(20, warning_density * 40) + # → 50 % warning density → 20-pt penalty + # → 5 % warning density → 2-pt penalty + # + # score = max(0, round(100 - critical_penalty - warning_penalty)) + # ----------------------------------------------------------------------- + dag_count = max(1, len(airflow_dags)) critical_count = sum(1 for i in issues if i.severity == "critical") warning_count = sum(1 for i in issues if i.severity == "warning") - score = max(0, 100 - (critical_count * 10) - (warning_count * 2)) + + critical_density = critical_count / dag_count + warning_density = warning_count / dag_count + + critical_penalty = min(90.0, critical_density * 200.0) + warning_penalty = min(20.0, warning_density * 40.0) + score = max(0, round(100.0 - critical_penalty - warning_penalty)) summary = { "score": score, "critical_count": critical_count, "warning_count": warning_count, "total_issues": len(issues), + "dag_count": dag_count, + "score_formula": ( + f"score = 100 - min(90, {critical_density:.3f} criticals/DAG × 200) " + f"- min(20, {warning_density:.3f} warnings/DAG × 40)" + ), "metadata": { "orchestrator_url": cfg.orchestrator.base_url, "warehouse_type": cfg.warehouse.type, "dbt_target": cfg.dbt.target, - } + }, } typer.echo("Rendering reports...") @@ -175,6 +216,21 @@ def init(): output_dir: "./reports" format: "both" fail_on_critical: 5 + +rules: + # How many days without a successful run before a DAG is flagged as stale. + stale_threshold_days: 7 + + # Maximum concurrent Airflow API calls during audit (runs + tasks per DAG). + fetch_concurrency: 10 + + # Per-rule severity overrides. Uncomment and adjust to match your team's SLAs. + # Valid severities: critical | warning | info + # severity_overrides: + # missing_sla: critical # fintech / real-time teams often require SLAs + # missing_retries: warning # default + # stale_dags: warning # default + # high_failure_rate: critical # default """ with open("pipelineprobe.yml", "w") as f: f.write(default_config) @@ -187,15 +243,237 @@ def doctor( config: str = typer.Option("pipelineprobe.yml", help="Path to config file"), ): """ - Validate connectivity to source systems. + Validate connectivity to Airflow, dbt artifacts, and the configured warehouse. """ - typer.echo(f"Checking connectivity using {config}...") + typer.echo(f"Checking connectivity using {config}...\n") cfg = load_config(config) - - typer.echo("Orchestrator: [STUB] Connection successful.") - typer.echo("Database: [STUB] Connection successful.") - typer.echo("dbt: [STUB] Artifacts found.") - typer.secho("\nAll systems operational.", fg=typer.colors.GREEN) + all_ok = True + + # ------------------------------------------------------------------ + # 1. Airflow — GET /api/v1/health returns metadatabase + scheduler state + # ------------------------------------------------------------------ + typer.echo("[Airflow]") + try: + auth = ( + (cfg.orchestrator.username, cfg.orchestrator.password) + if cfg.orchestrator.username and cfg.orchestrator.password + else None + ) + with httpx.Client( + base_url=cfg.orchestrator.base_url, + auth=auth, + verify=cfg.orchestrator.verify_ssl, + timeout=10.0, + ) as client: + resp = client.get("/api/v1/health") + resp.raise_for_status() + health = resp.json() + meta_status = health.get("metadatabase", {}).get("status", "unknown") + sched_status = health.get("scheduler", {}).get("status", "unknown") + typer.secho( + f" Connection: OK ({cfg.orchestrator.base_url})", + fg=typer.colors.GREEN, + ) + _print_status_line(" Metadatabase:", meta_status) + _print_status_line(" Scheduler: ", sched_status) + except httpx.HTTPStatusError as exc: + typer.secho( + f" Connection: FAIL (HTTP {exc.response.status_code}: {exc.response.text[:120]})", + fg=typer.colors.RED, + ) + all_ok = False + except Exception as exc: + typer.secho(f" Connection: FAIL ({exc})", fg=typer.colors.RED) + all_ok = False + + # ------------------------------------------------------------------ + # 2. dbt artifacts — check manifest.json (required) and run_results.json + # ------------------------------------------------------------------ + typer.echo("\n[dbt]") + manifest = Path(cfg.dbt.manifest_path) + run_results = Path(cfg.dbt.run_results_path) + + if manifest.exists(): + typer.secho(f" manifest.json: Found ({manifest})", fg=typer.colors.GREEN) + else: + typer.secho(f" manifest.json: MISSING ({manifest})", fg=typer.colors.RED) + all_ok = False + + if run_results.exists(): + typer.secho( + f" run_results.json: Found ({run_results})", fg=typer.colors.GREEN + ) + else: + # run_results is optional — warn but don't fail + typer.secho( + f" run_results.json: MISSING ({run_results}) [optional — some rules may be skipped]", + fg=typer.colors.YELLOW, + ) + + # ------------------------------------------------------------------ + # 3. Warehouse — light connectivity probe per driver + # ------------------------------------------------------------------ + typer.echo(f"\n[Warehouse — {cfg.warehouse.type}]") + try: + if cfg.warehouse.type == "postgres": + import psycopg2 + + conn = psycopg2.connect(cfg.warehouse.dsn, connect_timeout=10) + conn.close() + dsn_safe = cfg.warehouse.dsn.split("@")[-1] if "@" in cfg.warehouse.dsn else cfg.warehouse.dsn + typer.secho( + f" Postgres: OK (@{dsn_safe})", fg=typer.colors.GREEN + ) + + elif cfg.warehouse.type == "bigquery": + from google.cloud import bigquery + + client = bigquery.Client(project=cfg.warehouse.project_id) + # list_datasets is the lightest possible probe + next(iter(client.list_datasets(max_results=1)), None) + typer.secho( + f" BigQuery: OK (project={cfg.warehouse.project_id})", + fg=typer.colors.GREEN, + ) + + elif cfg.warehouse.type == "snowflake": + import snowflake.connector + + conn = snowflake.connector.connect( + account=cfg.warehouse.account, + user=cfg.warehouse.username, + password=cfg.warehouse.password, + login_timeout=10, + ) + conn.close() + typer.secho( + f" Snowflake: OK (account={cfg.warehouse.account})", + fg=typer.colors.GREEN, + ) + else: + typer.secho( + f" Unknown warehouse type '{cfg.warehouse.type}' — skipping probe.", + fg=typer.colors.YELLOW, + ) + + except Exception as exc: + typer.secho( + f" {cfg.warehouse.type}: FAIL ({exc})", fg=typer.colors.RED + ) + all_ok = False + + # ------------------------------------------------------------------ + # Summary + # ------------------------------------------------------------------ + typer.echo() + if all_ok: + typer.secho("All systems operational.", fg=typer.colors.GREEN) + else: + typer.secho( + "One or more checks failed — review the output above before running audit.", + fg=typer.colors.RED, + ) + raise typer.Exit(code=1) + + +def _print_status_line(label: str, status: str) -> None: + """Print a labelled status value coloured green/red based on 'healthy'.""" + color = typer.colors.GREEN if status == "healthy" else typer.colors.YELLOW + typer.secho(f"{label} {status}", fg=color) + + +@app.command() +def diff( + report_a: str = typer.Argument(..., help="Path to the baseline JSON report"), + report_b: str = typer.Argument(..., help="Path to the current JSON report"), +): + """ + Compare two audit JSON reports and show regressions and improvements. + + Exit code 1 if any regressions are found, 0 otherwise. + """ + path_a, path_b = Path(report_a), Path(report_b) + + for p in (path_a, path_b): + if not p.exists(): + typer.secho(f"Report not found: {p}", fg=typer.colors.RED) + raise typer.Exit(code=1) + + with open(path_a) as f: + data_a = json.load(f) + with open(path_b) as f: + data_b = json.load(f) + + score_a = data_a.get("summary", {}).get("score", 0) + score_b = data_b.get("summary", {}).get("score", 0) + delta = score_b - score_a + + # Fingerprint issues by (severity, summary) to detect changes. + # This is intentionally coarse — rule text changes will show as new issues, + # which is the desired behaviour during rule updates. + def _fp(issue: dict) -> str: + return f"{issue['severity']}|{issue['summary']}" + + fp_a = {_fp(i): i for i in data_a.get("issues", [])} + fp_b = {_fp(i): i for i in data_b.get("issues", [])} + + regressions = {k: v for k, v in fp_b.items() if k not in fp_a} + improvements = {k: v for k, v in fp_a.items() if k not in fp_b} + unchanged_count = sum(1 for k in fp_b if k in fp_a) + + # ---- Header ---- + delta_str = f"+{delta}" if delta > 0 else str(delta) + delta_color = typer.colors.GREEN if delta >= 0 else typer.colors.RED + typer.echo(f"Baseline : {report_a} (score {score_a})") + typer.echo(f"Current : {report_b} (score {score_b})") + typer.secho(f"Score delta: {delta_str}", fg=delta_color, bold=True) + + # ---- Regressions ---- + if regressions: + typer.secho( + f"\n{len(regressions)} Regression(s) — new issues in current report:", + fg=typer.colors.RED, + bold=True, + ) + _SEVERITY_ORDER = {"critical": 0, "warning": 1, "info": 2} + for issue in sorted( + regressions.values(), key=lambda x: _SEVERITY_ORDER.get(x["severity"], 9) + ): + sev_color = ( + typer.colors.RED + if issue["severity"] == "critical" + else typer.colors.YELLOW + ) + typer.secho( + f" [{issue['severity'].upper()}] {issue['summary']}", fg=sev_color + ) + if issue.get("recommendation"): + typer.echo(f" → {issue['recommendation']}") + + # ---- Improvements ---- + if improvements: + typer.secho( + f"\n{len(improvements)} Improvement(s) — issues resolved since baseline:", + fg=typer.colors.GREEN, + bold=True, + ) + for issue in improvements.values(): + typer.secho( + f" [{issue['severity'].upper()}] {issue['summary']}", fg=typer.colors.GREEN + ) + + if not regressions and not improvements: + typer.echo("\nNo changes detected between reports.") + + # ---- Footer ---- + typer.echo( + f"\nSummary: {unchanged_count} unchanged, " + f"{len(regressions)} regression(s), " + f"{len(improvements)} improvement(s)." + ) + + if regressions: + raise typer.Exit(code=1) if __name__ == "__main__": diff --git a/pipelineprobe/config.py b/pipelineprobe/config.py index ed7239e..d30b23f 100644 --- a/pipelineprobe/config.py +++ b/pipelineprobe/config.py @@ -3,6 +3,18 @@ from pydantic import BaseModel from pydantic_settings import BaseSettings +# Valid rule names that can have their severity overridden via YAML config. +OVERRIDABLE_RULES = { + "missing_sla", + "missing_retries", + "high_failure_rate", + "stale_dags", + "missing_dbt_tests", + "failing_dbt_models", +} + +VALID_SEVERITIES = {"critical", "warning", "info"} + class AirflowConfig(BaseModel): type: str = "airflow" @@ -36,11 +48,29 @@ class ReportConfig(BaseModel): fail_on_critical: int = 5 +class RulesConfig(BaseModel): + # Per-rule severity overrides. Keys must be one of OVERRIDABLE_RULES; + # values must be one of VALID_SEVERITIES. + # Example YAML: + # rules: + # severity_overrides: + # missing_sla: critical # fintech teams often treat this as critical + # stale_dags: warning + severity_overrides: dict[str, str] = {} + + # How many days without a successful run before a DAG is considered stale. + stale_threshold_days: int = 7 + + # Concurrency limit for async Airflow API calls during audit. + fetch_concurrency: int = 10 + + class PipelineProbeConfig(BaseSettings): orchestrator: AirflowConfig = AirflowConfig() dbt: DbtConfig = DbtConfig() warehouse: WarehouseConfig = WarehouseConfig() report: ReportConfig = ReportConfig() + rules: RulesConfig = RulesConfig() def load_config(config_path: str) -> PipelineProbeConfig: diff --git a/pipelineprobe/connectors/airflow.py b/pipelineprobe/connectors/airflow.py index 55cebc5..aba16cf 100644 --- a/pipelineprobe/connectors/airflow.py +++ b/pipelineprobe/connectors/airflow.py @@ -1,3 +1,4 @@ +import asyncio import logging from datetime import timedelta from typing import List @@ -29,6 +30,7 @@ def __init__(self, config: AirflowConfig): else None ) + self._auth = auth self.client = httpx.Client( base_url=self.config.base_url, auth=auth, @@ -121,3 +123,107 @@ def get_tasks(self, dag_id: str) -> List[Task]: except Exception as e: logger.error("Error fetching tasks for %s: %s", dag_id, e) return [] + + # ------------------------------------------------------------------ + # Async helpers — used by fetch_dag_details() for concurrent fetching + # ------------------------------------------------------------------ + + @staticmethod + def _parse_dag_runs(runs_data: list) -> List[DagRun]: + runs = [] + for r in runs_data: + state = r.get("state", "unknown") + start_str = r.get("start_date") + end_str = r.get("end_date") + start_time = ( + parse_date(start_str) + if start_str + else parse_date(r["execution_date"]) + ) + end_time = parse_date(end_str) if end_str else None + runs.append(DagRun(state=state, start_time=start_time, end_time=end_time)) + return runs + + @staticmethod + def _parse_tasks(dag_id: str, tasks_data: list) -> List[Task]: + tasks = [] + for t in tasks_data: + retries = t.get("retries", 0) + has_sla = bool(t.get("sla")) + sla_val = timedelta(seconds=1) if has_sla else None + has_alerts = bool(t.get("email")) or bool(t.get("email_on_failure")) + tasks.append( + Task( + dag_id=dag_id, + task_id=t.get("task_id", ""), + retries=int(retries) if retries is not None else 0, + sla=sla_val, + has_alerts=has_alerts, + ) + ) + return tasks + + async def _fetch_dag_runs_async( + self, dag_id: str, client: httpx.AsyncClient + ) -> List[DagRun]: + try: + response = await client.get( + f"/api/v1/dags/{dag_id}/dagRuns", + params={"limit": 20, "order_by": "-execution_date"}, + ) + response.raise_for_status() + return self._parse_dag_runs(response.json().get("dag_runs", [])) + except Exception as e: + logger.error("Error fetching DAG runs for %s: %s", dag_id, e) + return [] + + async def _fetch_tasks_async( + self, dag_id: str, client: httpx.AsyncClient + ) -> List[Task]: + try: + response = await client.get(f"/api/v1/dags/{dag_id}/tasks") + response.raise_for_status() + return self._parse_tasks(dag_id, response.json().get("tasks", [])) + except Exception as e: + logger.error("Error fetching tasks for %s: %s", dag_id, e) + return [] + + async def fetch_dag_details( + self, dags: List[Dag], concurrency: int = 10 + ) -> tuple[List[Dag], List[Task]]: + """Fetch dag runs and tasks for all DAGs concurrently. + + Uses a semaphore to cap the number of in-flight requests so large + installations don't overwhelm the Airflow API. Two requests are + issued per DAG (runs + tasks), so effective parallelism is + ``concurrency * 2`` in-flight connections at peak. + """ + semaphore = asyncio.Semaphore(concurrency) + + async with httpx.AsyncClient( + base_url=self.config.base_url, + auth=self._auth, + verify=self.config.verify_ssl, + ) as client: + + async def _fetch_one(dag: Dag) -> tuple[Dag, List[Task]]: + async with semaphore: + dag.recent_runs, tasks = await asyncio.gather( + self._fetch_dag_runs_async(dag.id, client), + self._fetch_tasks_async(dag.id, client), + ) + return dag, tasks + + results = await asyncio.gather( + *[_fetch_one(dag) for dag in dags], return_exceptions=True + ) + + all_tasks: List[Task] = [] + for result in results: + if isinstance(result, Exception): + logger.error("Unhandled error fetching DAG details: %s", result) + else: + _, tasks = result + all_tasks.extend(tasks) + + return dags, all_tasks diff --git a/pipelineprobe/rules/airflow_rules.py b/pipelineprobe/rules/airflow_rules.py index ccb4fa3..45df0ed 100644 --- a/pipelineprobe/rules/airflow_rules.py +++ b/pipelineprobe/rules/airflow_rules.py @@ -3,14 +3,21 @@ from pipelineprobe.models import Task, Issue, Dag +def _severity(context: dict, rule_name: str, default: str) -> str: + """Return the effective severity for a rule, respecting YAML overrides.""" + overrides: dict = context.get("rule_severity_overrides", {}) + return overrides.get(rule_name, default) + + def check_missing_retries(context: dict) -> List[Issue]: issues = [] tasks: List[Task] = context.get("airflow_tasks", []) + sev = _severity(context, "missing_retries", "warning") for task in tasks: if task.retries == 0: issues.append( Issue( - severity="warning", + severity=sev, category="task", summary=f"Task {task.task_id} in DAG {task.dag_id} has no retries configured.", details="Tasks without retries are more prone to transient failures.", @@ -24,11 +31,12 @@ def check_missing_retries(context: dict) -> List[Issue]: def check_missing_slas(context: dict) -> List[Issue]: issues = [] tasks: List[Task] = context.get("airflow_tasks", []) + sev = _severity(context, "missing_sla", "info") for task in tasks: if not task.sla: issues.append( Issue( - severity="info", + severity=sev, category="task", summary=f"Task {task.task_id} in DAG {task.dag_id} has no SLA configured.", details="Tasks without SLAs might silently miss deadlines.", @@ -42,6 +50,7 @@ def check_missing_slas(context: dict) -> List[Issue]: def check_high_failure_rate(context: dict) -> List[Issue]: issues = [] dags: List[Dag] = context.get("airflow_dags", []) + sev = _severity(context, "high_failure_rate", "critical") for dag in dags: if not dag.is_active or not dag.recent_runs: continue @@ -53,7 +62,7 @@ def check_high_failure_rate(context: dict) -> List[Issue]: if failure_ratio > 0.2: issues.append( Issue( - severity="critical", + severity=sev, category="dag", summary=f"DAG {dag.id} has a high failure rate ({failure_ratio:.0%}).", details=f"{failed_runs} out of the last {total_runs} runs failed.", @@ -67,8 +76,8 @@ def check_high_failure_rate(context: dict) -> List[Issue]: def check_stale_dags(context: dict) -> List[Issue]: issues = [] dags: List[Dag] = context.get("airflow_dags", []) - # Default stale threshold is 7 days, could be configurable later - stale_threshold_days = 7 + stale_threshold_days: int = context.get("stale_threshold_days", 7) + sev = _severity(context, "stale_dags", "warning") now = datetime.now(timezone.utc) for dag in dags: @@ -79,7 +88,7 @@ def check_stale_dags(context: dict) -> List[Issue]: if not dag.recent_runs: issues.append( Issue( - severity="warning", + severity=sev, category="dag", summary=f"DAG {dag.id} has no recent runs recorded.", details="The DAG is active but has never run (or runs were not retrieved).", @@ -95,10 +104,12 @@ def check_stale_dags(context: dict) -> List[Issue]: ] if not recent_successes: - # Has runs but zero successes — flag as critical + # Has runs but zero successes — flag as critical regardless of override, + # because this is an active outage signal (override still respected) + outage_sev = _severity(context, "stale_dags", "critical") issues.append( Issue( - severity="critical", + severity=outage_sev, category="dag", summary=f"DAG {dag.id} has no successful runs in its recent history.", details=f"Last {len(dag.recent_runs)} runs all ended in non-success states.", @@ -119,10 +130,10 @@ def check_stale_dags(context: dict) -> List[Issue]: if days_since > stale_threshold_days: issues.append( Issue( - severity="warning", + severity=sev, category="dag", summary=f"DAG {dag.id} is stale.", - details=f"No successful execution in the last {days_since:.1f} days.", + details=f"No successful execution in the last {days_since:.1f} days (threshold: {stale_threshold_days}d).", recommendation="Check if the DAG is still needed, or if it is silently failing to schedule.", affected_resources=[dag.id], )