Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 22 additions & 11 deletions pipelineprobe/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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".
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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

# ------------------------------------------------------------------
Expand Down Expand Up @@ -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:
Expand Down
4 changes: 1 addition & 3 deletions pipelineprobe/connectors/airflow.py
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down
2 changes: 1 addition & 1 deletion pipelineprobe/connectors/bigquery.py
Original file line number Diff line number Diff line change
Expand Up @@ -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),
}
)
Expand Down
8 changes: 4 additions & 4 deletions pipelineprobe/connectors/dagster.py
Original file line number Diff line number Diff line change
Expand Up @@ -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", [])
Expand Down
4 changes: 3 additions & 1 deletion pipelineprobe/connectors/prefect.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
4 changes: 2 additions & 2 deletions pipelineprobe/rules/cost_rules.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
91 changes: 77 additions & 14 deletions tests/test_cost_rules.py
Original file line number Diff line number Diff line change
Expand Up @@ -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) == []
Expand All @@ -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)
Expand All @@ -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)
Expand All @@ -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)
Expand All @@ -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)
Expand All @@ -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) == []
Expand All @@ -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) == []
Expand All @@ -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)
Expand All @@ -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)
Expand All @@ -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)
Expand Down
6 changes: 3 additions & 3 deletions tests/test_new_connectors.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down
Loading