"""
run_query_optimizer.py
----------------------
Trigger script for the `query-optimizer` DinoAI agent.
Environment variables (required)
=================================
PARADIME_API_ENDPOINT – Paradime GraphQL endpoint
PARADIME_API_KEY – Paradime API key
PARADIME_API_SECRET – Paradime API secret
At least ONE of the following must also be set (or passed via CLI):
BOLT_SCHEDULE_NAME – Bolt schedule to pull job IDs from automatically
QUERY_OPTIMIZER_JOB_IDS – Comma-separated BigQuery job IDs (skips Bolt fetch)
Environment variables (optional)
=================================
QUERY_OPTIMIZER_TOP_N – Max models to optimise (default: 5)
QUERY_OPTIMIZER_MIN_BYTES_GB – Skip jobs below this GB threshold (default: 10)
QUERY_OPTIMIZER_DRILL_MODEL – dbt™ model to drill into after the main pass
QUERY_OPTIMIZER_TIMEOUT – Max seconds to wait for the agent (default: 3600)
QUERY_OPTIMIZER_POLL_INTERVAL – Polling cadence in seconds (default: 20)
"""
import argparse
import os
import sys
import textwrap
import time
import requests
from paradime import Paradime
from paradime.apis.dinoai_agents.exception import DinoaiAgentRunFailedException
from paradime.apis.dinoai_agents.types import DinoaiAgentRunStatus
AGENT_NAME = "query-optimizer"
ARTIFACT_PATH = "target/run_results.json"
_GB = 1024 ** 3
TOP_N = int(os.environ.get("QUERY_OPTIMIZER_TOP_N", 5))
MIN_BYTES_GB = float(os.environ.get("QUERY_OPTIMIZER_MIN_BYTES_GB", 10))
DRILL_MODEL = os.environ.get("QUERY_OPTIMIZER_DRILL_MODEL", "").strip()
TIMEOUT = int(os.environ.get("QUERY_OPTIMIZER_TIMEOUT", 3600))
POLL_INTERVAL = int(os.environ.get("QUERY_OPTIMIZER_POLL_INTERVAL", 20))
def _parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(
description="Trigger the query-optimizer DinoAI agent."
)
source = parser.add_mutually_exclusive_group()
source.add_argument(
"--job-ids", nargs="+", metavar="JOB_ID",
default=_split_env("QUERY_OPTIMIZER_JOB_IDS"),
help="One or more BigQuery job IDs. (env: QUERY_OPTIMIZER_JOB_IDS)",
)
source.add_argument(
"--schedule-name",
default=os.environ.get("BOLT_SCHEDULE_NAME", "").strip(),
help="Bolt schedule name — job IDs fetched from latest run artifact. (env: BOLT_SCHEDULE_NAME)",
)
parser.add_argument("--top-n", type=int, default=TOP_N)
parser.add_argument("--min-bytes-gb", type=float, default=MIN_BYTES_GB)
parser.add_argument("--drill-model", default=DRILL_MODEL)
return parser.parse_args()
def _split_env(var: str) -> list[str]:
raw = os.environ.get(var, "").strip()
return [v.strip() for v in raw.split(",") if v.strip()] if raw else []
def _build_client() -> Paradime:
required = {
"PARADIME_API_ENDPOINT": os.environ.get("PARADIME_API_ENDPOINT"),
"PARADIME_API_KEY": os.environ.get("PARADIME_API_KEY"),
"PARADIME_API_SECRET": os.environ.get("PARADIME_API_SECRET"),
}
missing = [k for k, v in required.items() if not v]
if missing:
print(f"ERROR: Missing env vars: {', '.join(missing)}", file=sys.stderr)
raise SystemExit(1)
return Paradime(
api_endpoint=required["PARADIME_API_ENDPOINT"],
api_key=required["PARADIME_API_KEY"],
api_secret=required["PARADIME_API_SECRET"],
)
def _fetch_run_results_from_bolt(paradime: Paradime, schedule_name: str) -> list[dict]:
print(f"[query-optimizer] 🔍 Fetching latest Bolt run for '{schedule_name}' …")
try:
latest_runs = paradime.bolt.list_runs(schedule_name=schedule_name, limit=1)
latest_run = latest_runs.runs[0] if latest_runs.runs else None
except Exception as exc:
print(f"[query-optimizer] ⚠️ Could not fetch run list: {exc}", file=sys.stderr)
return []
if not latest_run:
print("[query-optimizer] ⚠️ No runs found for this schedule.", file=sys.stderr)
return []
run_id = latest_run.id
print(f"[query-optimizer] 🏃 Run ID: {run_id} | status: {latest_run.state}")
try:
run_status = paradime.bolt.get_run_status(run_id)
commands = getattr(run_status, "commands", None) or []
command_indices = list(range(len(commands))) if commands else list(range(5))
except Exception:
command_indices = list(range(5))
all_model_results: list[dict] = []
seen_urls: set[str] = set()
for command_index in command_indices:
try:
url = paradime.bolt.get_latest_artifact_url(
schedule_name=schedule_name,
artifact_path=ARTIFACT_PATH,
command_index=command_index,
)
except Exception:
continue
if not url or url in seen_urls:
continue
seen_urls.add(url)
try:
resp = requests.get(url, timeout=30)
resp.raise_for_status()
data = resp.json()
except Exception as exc:
print(f"[query-optimizer] ⚠️ Failed to download artifact: {exc}", file=sys.stderr)
continue
for result in data.get("results", []):
if not result.get("unique_id", "").startswith("model."):
continue
adapter = result.get("adapter_response") or {}
all_model_results.append({
"name": result["unique_id"].split(".")[-1],
"unique_id": result["unique_id"],
"status": result.get("status", "—"),
"execution_time_s": round(result.get("execution_time") or 0, 2),
"job_id": adapter.get("job_id", "—"),
"bytes_billed": adapter.get("bytes_billed", 0),
"bytes_processed": adapter.get("bytes_processed", 0),
"rows_affected": adapter.get("rows_affected", "—"),
"command_index": command_index,
"generated_at": data.get("metadata", {}).get("generated_at", "unknown"),
})
return all_model_results
def _top_costly_models(all_results: list[dict], top_n: int, min_bytes_gb: float) -> list[dict]:
deduped: dict[str, dict] = {}
for row in all_results:
name = row["name"]
if name not in deduped or row["bytes_processed"] > deduped[name]["bytes_processed"]:
deduped[name] = row
min_bytes = int(min_bytes_gb * _GB)
filtered = [r for r in deduped.values() if r["bytes_processed"] >= min_bytes]
return sorted(filtered, key=lambda r: r["bytes_processed"], reverse=True)[:top_n]
def _rows_from_job_ids(job_ids: list[str]) -> list[dict]:
return [
{
"name": f"job_{i + 1}", "unique_id": "—", "status": "unknown",
"execution_time_s": 0, "job_id": jid, "bytes_billed": 0,
"bytes_processed": 2 ** 63, "rows_affected": "—",
"command_index": 0, "generated_at": "—",
}
for i, jid in enumerate(job_ids)
]
def _format_table(rows: list[dict]) -> str:
if not rows:
return "(no model results to display)"
headers = {
"name": "Model", "execution_time_s": "Time (s)",
"bytes_processed": "Bytes Processed", "bytes_billed": "Bytes Billed",
"job_id": "BigQuery Job ID", "command_index": "Cmd",
}
widths = {col: len(label) for col, label in headers.items()}
for row in rows:
for col in headers:
widths[col] = max(widths[col], len(str(row.get(col, ""))))
def _row_str(values):
return " ".join(str(values.get(col, "")).ljust(widths[col]) for col in headers)
header_line = _row_str(headers)
return "\n".join([header_line, "-" * len(header_line)] + [_row_str(r) for r in rows])
def _build_opening_message(rows: list[dict], source_label: str) -> str:
job_id_lines = "\n".join(
f" {i + 1}. model=`{r['name']}` job_id=`{r['job_id']}` "
f"bytes_processed={r['bytes_processed']:,} bytes_billed={r['bytes_billed']:,} "
f"execution_time={r['execution_time_s']}s"
for i, r in enumerate(rows)
)
return textwrap.dedent(f"""
You are being triggered to analyse and fix costly dbt™ models running on BigQuery.
SOURCE: {source_label}
── PRE-FETCHED MODEL PERFORMANCE DATA ────────────────────────────────────
{_format_table(rows)}
── JOB IDs TO OPTIMISE ───────────────────────────────────────────────────
{job_id_lines}
Work through all five phases (GATHER → LOCATE → DIAGNOSE → IMPLEMENT →
PR + SLACK) in order without stopping. The run is only complete when the
PR is open and the Slack message has been posted to #finops.
""").strip()
def _build_drill_message(drill_model: str, rows: list[dict]) -> str:
drill_row = next((r for r in rows if r["name"] == drill_model), None)
job_hint = (
f"Its BigQuery job ID is `{drill_row['job_id']}` ({drill_row['execution_time_s']}s). "
f"Use get_bigquery_query_performance_stats to retrieve the full query plan. "
if drill_row and drill_row["job_id"] != "—"
else f"No pre-fetched job ID for `{drill_model}` — analyse the .sql file directly. "
)
return (
f"Now drill deeper into `{drill_model}` specifically. {job_hint}"
f"Show the exact query plan stage(s) consuming the most slot ms, identify "
f"the bottleneck type, and apply any additional targeted fixes not already "
f"covered in the main pass. Commit changes to the existing branch, update "
f"the PR description, and post a follow-up to #finops with the updated PR link."
)
def _poll_until_done(paradime: Paradime, session_id: str, label: str) -> object:
start = time.time()
while True:
run = paradime.dinoai_agents.get_run(agent_session_id=session_id)
if run.status == DinoaiAgentRunStatus.COMPLETED:
print(f"[query-optimizer] ✅ {label} completed")
return run
if run.status == DinoaiAgentRunStatus.FAILED:
last_msg = run.messages[-1].content if run.messages else "no messages"
raise DinoaiAgentRunFailedException(f"{label} FAILED. Last: {last_msg}")
elapsed = time.time() - start
if elapsed > TIMEOUT:
raise TimeoutError(f"Timed out after {TIMEOUT}s for session {session_id}")
print(f"[query-optimizer] {label} — {run.status.value} ({int(elapsed)}s) … {POLL_INTERVAL}s")
time.sleep(POLL_INTERVAL)
def main() -> None:
args = _parse_args()
paradime = _build_client()
if args.job_ids:
rows = _rows_from_job_ids(args.job_ids)
source_label = f"explicit job IDs: {', '.join(args.job_ids)}"
elif args.schedule_name:
all_results = _fetch_run_results_from_bolt(paradime, args.schedule_name)
rows = _top_costly_models(all_results, args.top_n, args.min_bytes_gb)
source_label = f"Bolt schedule '{args.schedule_name}' — top {len(rows)} models above {args.min_bytes_gb} GB"
if not rows:
print("[query-optimizer] ⚠️ No models above threshold. Exiting.")
raise SystemExit(0)
else:
print("ERROR: Provide --job-ids or --schedule-name.", file=sys.stderr)
raise SystemExit(1)
print("\n" + "=" * 72)
print(_format_table(rows))
print("=" * 72 + "\n")
trigger = paradime.dinoai_agents.trigger_run(
agent=AGENT_NAME,
message=_build_opening_message(rows, source_label),
)
session_id = trigger.agent_session_id
print(f"[query-optimizer] 📋 Session ID: {session_id}")
main_run = _poll_until_done(paradime, session_id, "Main optimisation pass")
print(main_run.messages[-1].content if main_run.messages else "(no report)")
if args.drill_model:
paradime.dinoai_agents.send_message(
agent_session_id=session_id,
message=_build_drill_message(args.drill_model, rows),
)
drill_run = _poll_until_done(paradime, session_id, f"Drill-down on '{args.drill_model}'")
print(drill_run.messages[-1].content if drill_run.messages else "(no drill report)")
print("[query-optimizer] ✅ All done.")
if __name__ == "__main__":
main()
