diff --git a/pipelineprobe/cli.py b/pipelineprobe/cli.py index 64de17e..ab3e729 100644 --- a/pipelineprobe/cli.py +++ b/pipelineprobe/cli.py @@ -38,7 +38,11 @@ def version_callback(value: bool): @app.callback() def main( version: bool = typer.Option( - None, "--version", callback=version_callback, is_eager=True, help="Show version and exit" + None, + "--version", + callback=version_callback, + is_eager=True, + help="Show version and exit", ), ): pass @@ -146,12 +150,18 @@ def audit( if cfg.report.include_cost_section: typer.echo("Fetching cost insights...") - if cfg.warehouse.type == "bigquery" and isinstance(warehouse_conn, BigQueryConnector): + if cfg.warehouse.type == "bigquery" and isinstance( + warehouse_conn, BigQueryConnector + ): bq_cost_insights = warehouse_conn.get_cost_insights_sync() typer.echo(f" BigQuery: {len(bq_cost_insights)} tables with cost data.") - elif cfg.warehouse.type == "snowflake" and isinstance(warehouse_conn, SnowflakeConnector): + elif cfg.warehouse.type == "snowflake" and isinstance( + warehouse_conn, SnowflakeConnector + ): sf_cost_insights = warehouse_conn.get_cost_insights_sync() - typer.echo(f" Snowflake: {len(sf_cost_insights)} warehouses with credit data.") + typer.echo( + f" Snowflake: {len(sf_cost_insights)} warehouses with credit data." + ) context = { # The existing Airflow rules operate on "airflow_dags" / "airflow_tasks". @@ -437,10 +447,12 @@ def doctor( 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 + 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 @@ -474,9 +486,7 @@ def doctor( ) except Exception as exc: - typer.secho( - f" {cfg.warehouse.type}: FAIL ({exc})", fg=typer.colors.RED - ) + typer.secho(f" {cfg.warehouse.type}: FAIL ({exc})", fg=typer.colors.RED) all_ok = False # ------------------------------------------------------------------ @@ -625,7 +635,8 @@ def _fp(issue: dict) -> str: ) for issue in improvements.values(): typer.secho( - f" [{issue['severity'].upper()}] {issue['summary']}", fg=typer.colors.GREEN + f" [{issue['severity'].upper()}] {issue['summary']}", + fg=typer.colors.GREEN, ) if not regressions and not improvements: diff --git a/pipelineprobe/connectors/airflow.py b/pipelineprobe/connectors/airflow.py index aba16cf..c368fe4 100644 --- a/pipelineprobe/connectors/airflow.py +++ b/pipelineprobe/connectors/airflow.py @@ -136,9 +136,7 @@ def _parse_dag_runs(runs_data: list) -> List[DagRun]: 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"]) + 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)) diff --git a/pipelineprobe/connectors/bigquery.py b/pipelineprobe/connectors/bigquery.py index c8485b0..cc60ead 100644 --- a/pipelineprobe/connectors/bigquery.py +++ b/pipelineprobe/connectors/bigquery.py @@ -108,7 +108,7 @@ def get_cost_insights_sync(self) -> List[Dict[str, Any]]: { "table_id": r.table_id, "total_bytes_billed": billed, - "total_gb_billed": round(billed / (1024 ** 3), 2), + "total_gb_billed": round(billed / (1024**3), 2), "query_count": int(r.query_count or 0), } ) diff --git a/pipelineprobe/connectors/dagster.py b/pipelineprobe/connectors/dagster.py index b36ce98..8eed00c 100644 --- a/pipelineprobe/connectors/dagster.py +++ b/pipelineprobe/connectors/dagster.py @@ -139,14 +139,14 @@ def _fetch_runs(self) -> List[Dict[str, Any]]: resp.raise_for_status() body = resp.json() - runs_or_error = ( - body.get("data", {}).get("runsOrError", {}) - ) + runs_or_error = body.get("data", {}).get("runsOrError", {}) typename = runs_or_error.get("__typename", "") if typename != "Runs": msg = runs_or_error.get("message", "unknown error") - logger.error("DagsterConnector: runsOrError returned %s: %s", typename, msg) + logger.error( + "DagsterConnector: runsOrError returned %s: %s", typename, msg + ) break page: List[Dict[str, Any]] = runs_or_error.get("results", []) diff --git a/pipelineprobe/connectors/prefect.py b/pipelineprobe/connectors/prefect.py index 0a497af..bfb48b5 100644 --- a/pipelineprobe/connectors/prefect.py +++ b/pipelineprobe/connectors/prefect.py @@ -126,7 +126,9 @@ def _parse_run(raw: Dict[str, Any]) -> DagRun: state_name: str = (raw.get("state") or {}).get("type", "UNKNOWN") state = _STATE_MAP.get(state_name.upper(), "unknown") - start_time = _parse_prefect_dt(raw.get("start_time") or raw.get("expected_start_time")) + start_time = _parse_prefect_dt( + raw.get("start_time") or raw.get("expected_start_time") + ) end_time = _parse_prefect_dt(raw.get("end_time")) return DagRun(state=state, start_time=start_time, end_time=end_time) diff --git a/pipelineprobe/rules/cost_rules.py b/pipelineprobe/rules/cost_rules.py index db733a4..cc3cba6 100644 --- a/pipelineprobe/rules/cost_rules.py +++ b/pipelineprobe/rules/cost_rules.py @@ -15,8 +15,8 @@ # Thresholds — intentionally conservative defaults so that the rules are # actionable for small and large organisations alike. -_BQ_WARN_GB = 500 # warn when a table accounts for >= 500 GiB billed / 30 days -_BQ_CRIT_GB = 5_000 # critical when >= 5 TiB billed +_BQ_WARN_GB = 500 # warn when a table accounts for >= 500 GiB billed / 30 days +_BQ_CRIT_GB = 5_000 # critical when >= 5 TiB billed _SF_WARN_CREDITS = 500 # warn when a warehouse consumes >= 500 credits / 30 days _SF_CRIT_CREDITS = 5_000 # critical when >= 5 000 credits diff --git a/tests/test_cost_rules.py b/tests/test_cost_rules.py index 62f68b7..51b592c 100644 --- a/tests/test_cost_rules.py +++ b/tests/test_cost_rules.py @@ -16,7 +16,12 @@ def test_skips_non_bigquery(self): context = { "warehouse_type": "postgres", "bq_cost_insights": [ - {"table_id": "p.d.t", "total_bytes_billed": 999_999_999_999, "total_gb_billed": 930.0, "query_count": 5} + { + "table_id": "p.d.t", + "total_bytes_billed": 999_999_999_999, + "total_gb_billed": 930.0, + "query_count": 5, + } ], } assert check_expensive_bq_tables(context) == [] @@ -29,7 +34,12 @@ def test_no_issues_below_threshold(self): context = { "warehouse_type": "bigquery", "bq_cost_insights": [ - {"table_id": "p.d.small", "total_bytes_billed": 1_073_741_824, "total_gb_billed": 1.0, "query_count": 10} + { + "table_id": "p.d.small", + "total_bytes_billed": 1_073_741_824, + "total_gb_billed": 1.0, + "query_count": 10, + } ], } issues = check_expensive_bq_tables(context) @@ -40,7 +50,12 @@ def test_warning_at_warn_threshold(self): "warehouse_type": "bigquery", "bq_cost_insights": [ # 600 GiB → warn (>= 500 GiB, < 5000 GiB) - {"table_id": "p.d.medium_table", "total_bytes_billed": 644_245_094_400, "total_gb_billed": 600.0, "query_count": 50} + { + "table_id": "p.d.medium_table", + "total_bytes_billed": 644_245_094_400, + "total_gb_billed": 600.0, + "query_count": 50, + } ], } issues = check_expensive_bq_tables(context) @@ -53,7 +68,12 @@ def test_critical_at_crit_threshold(self): "warehouse_type": "bigquery", "bq_cost_insights": [ # 6000 GiB → critical (>= 5000 GiB) - {"table_id": "p.d.huge_table", "total_bytes_billed": 6_442_450_944_000, "total_gb_billed": 6000.0, "query_count": 500} + { + "table_id": "p.d.huge_table", + "total_bytes_billed": 6_442_450_944_000, + "total_gb_billed": 6000.0, + "query_count": 500, + } ], } issues = check_expensive_bq_tables(context) @@ -65,9 +85,24 @@ def test_mixed_thresholds(self): context = { "warehouse_type": "bigquery", "bq_cost_insights": [ - {"table_id": "p.d.a", "total_bytes_billed": 5_000 * 1024**3, "total_gb_billed": 5000.0, "query_count": 100}, - {"table_id": "p.d.b", "total_bytes_billed": 600 * 1024**3, "total_gb_billed": 600.0, "query_count": 20}, - {"table_id": "p.d.c", "total_bytes_billed": 10 * 1024**3, "total_gb_billed": 10.0, "query_count": 5}, + { + "table_id": "p.d.a", + "total_bytes_billed": 5_000 * 1024**3, + "total_gb_billed": 5000.0, + "query_count": 100, + }, + { + "table_id": "p.d.b", + "total_bytes_billed": 600 * 1024**3, + "total_gb_billed": 600.0, + "query_count": 20, + }, + { + "table_id": "p.d.c", + "total_bytes_billed": 10 * 1024**3, + "total_gb_billed": 10.0, + "query_count": 5, + }, ], } issues = check_expensive_bq_tables(context) @@ -87,7 +122,11 @@ def test_skips_non_snowflake(self): context = { "warehouse_type": "bigquery", "sf_cost_insights": [ - {"warehouse_name": "WH", "total_credits": 9999.0, "cloud_services": 100.0} + { + "warehouse_name": "WH", + "total_credits": 9999.0, + "cloud_services": 100.0, + } ], } assert check_snowflake_credit_spenders(context) == [] @@ -100,7 +139,11 @@ def test_no_issues_below_threshold(self): context = { "warehouse_type": "snowflake", "sf_cost_insights": [ - {"warehouse_name": "SMALL_WH", "total_credits": 100.0, "cloud_services": 5.0} + { + "warehouse_name": "SMALL_WH", + "total_credits": 100.0, + "cloud_services": 5.0, + } ], } assert check_snowflake_credit_spenders(context) == [] @@ -110,7 +153,11 @@ def test_warning_at_warn_threshold(self): "warehouse_type": "snowflake", "sf_cost_insights": [ # 750 credits → warn (>= 500, < 5000) - {"warehouse_name": "MEDIUM_WH", "total_credits": 750.0, "cloud_services": 30.0} + { + "warehouse_name": "MEDIUM_WH", + "total_credits": 750.0, + "cloud_services": 30.0, + } ], } issues = check_snowflake_credit_spenders(context) @@ -123,7 +170,11 @@ def test_critical_at_crit_threshold(self): "warehouse_type": "snowflake", "sf_cost_insights": [ # 6000 credits → critical (>= 5000) - {"warehouse_name": "BIG_WH", "total_credits": 6000.0, "cloud_services": 200.0} + { + "warehouse_name": "BIG_WH", + "total_credits": 6000.0, + "cloud_services": 200.0, + } ], } issues = check_snowflake_credit_spenders(context) @@ -135,9 +186,21 @@ def test_multiple_warehouses(self): context = { "warehouse_type": "snowflake", "sf_cost_insights": [ - {"warehouse_name": "GIANT_WH", "total_credits": 8000.0, "cloud_services": 300.0}, - {"warehouse_name": "MED_WH", "total_credits": 600.0, "cloud_services": 20.0}, - {"warehouse_name": "TINY_WH", "total_credits": 50.0, "cloud_services": 2.0}, + { + "warehouse_name": "GIANT_WH", + "total_credits": 8000.0, + "cloud_services": 300.0, + }, + { + "warehouse_name": "MED_WH", + "total_credits": 600.0, + "cloud_services": 20.0, + }, + { + "warehouse_name": "TINY_WH", + "total_credits": 50.0, + "cloud_services": 2.0, + }, ], } issues = check_snowflake_credit_spenders(context) diff --git a/tests/test_new_connectors.py b/tests/test_new_connectors.py index e014d87..552f21e 100644 --- a/tests/test_new_connectors.py +++ b/tests/test_new_connectors.py @@ -58,9 +58,9 @@ def test_get_dags_success(self, mock_client_cls): ] mock_client_instance.post.side_effect = [ - flows_resp, # /api/flows/filter - runs_resp, # /api/flow_runs/filter for flow-uuid-1 - runs_resp, # /api/flow_runs/filter for flow-uuid-2 + flows_resp, # /api/flows/filter + runs_resp, # /api/flow_runs/filter for flow-uuid-1 + runs_resp, # /api/flow_runs/filter for flow-uuid-2 ] dags = connector.get_dags()