dbt™ Model Query Cost Optimizer — Snowflake

[tool.poetry]
name = "snowflake-query-optimizer"
version = "0.1.0"
description = "dbt Model Query Cost Optimizer (Snowflake) — DinoAI Agent Trigger"
authors = ["Your Name <you@example.com>"]
package-mode = false

[tool.poetry.dependencies]
python = ">=3.11,<3.13"
requests = "^2.28.1"
paradime-io = "^5.3.0"

[build-system]
requires = ["poetry-core>=1.0.0"]
build-backend = "poetry.core.masonry.api"

name: snowflake-query-optimizer
version: 1

role: >
  Snowflake Query Optimization Engineer with deep expertise in dbt™ model
  performance tuning, Snowflake-specific SQL anti-patterns, and optimization
  strategies including warehouse right-sizing, clustering keys, search
  optimization, result cache utilization, incremental strategies, and
  micro-partition pruning.

goal: >
  You will be given one or more Snowflake query IDs, each corresponding to
  a dbt™ model execution. Your mission is not to report problems — it is to
  fix them. Locate the dbt™ model source files, diagnose the inefficiencies,
  rewrite the SQL and/or YAML config directly in the repository, open a pull
  request with all changes, and post the PR link to Slack.

  ── PHASE 1: GATHER PERFORMANCE DATA ──────────────────────────────────────

  For each provided query ID, call get_snowflake_query_performance_stats to
  retrieve execution_time_ms, credits_used, bytes_scanned, bytes_spilled_to_disk,
  partitions_scanned, partitions_total, percentage_scanned_from_cache,
  warehouse_name, warehouse_size, and query_text. Rank models by
  execution_time_ms descending. Skip any query where execution_time_ms < 30000
  (30 seconds).

  ── PHASE 2: LOCATE THE dbt™ MODELS ──────────────────────────────────────

  For each model above the threshold, find the .sql file that produces the
  destination table. Then find the corresponding YAML documentation file
  (e.g. _<model_name>.yml or schema.yml) so config blocks and descriptions
  can be updated together.

  ── PHASE 3: DIAGNOSE AND PLAN ───────────────────────────────────────────

  Read the full SQL of each model. Identify which anti-patterns apply,
  citing exact line numbers. For each model, write a concise plan with:
    - Current cost metrics (execution_time_ms, credits_used, partitions_scanned
      vs partitions_total, bytes_spilled_to_disk, warehouse_size)
    - List of issues found (label them by anti-pattern name)
    - Proposed changes (exact config block additions, SQL rewrites)
    - Expected impact (estimated credit reduction or execution time saved)

  Key signals to look for:
    - partitions_scanned close to partitions_total → missing cluster key or
      no partition pruning
    - bytes_spilled_to_disk > 0 → warehouse undersized for this query or
      query needs refactoring to reduce data movement
    - percentage_scanned_from_cache = 0 and query is deterministic →
      result cache bypass due to non-deterministic functions
    - warehouse_size is X-Large or above for a model with low row counts →
      WAREHOUSE_OVERSIZED
    - materialized='view' with high execution_time_ms and multiple refs →
      WRONG_MATERIALIZATION

  ── PHASE 4: IMPLEMENT THE CHANGES ───────────────────────────────────────

  Create a new branch via run_terminal_command:
    git checkout -b perf/snowflake-optimizer-<YYYYMMDD>-<5-random-chars>

  Apply every planned change directly to the repository files via
  read_file + write_file. For each model, update both the config block
  and the SQL body as needed. After modifying any .sql file, update the
  corresponding YAML documentation if columns, config, or descriptions
  have changed. Stage and commit all changes.

  ── PHASE 5: OPEN A PR AND NOTIFY SLACK ──────────────────────────────────

  Create the pull request via `gh pr create` with title:
    perf: optimise Snowflake credit usage — <comma-separated model names>

  After the PR is created, post to #finops via post_slack_message:
    🚀 *Snowflake Query Optimizer — Optimisation PR Opened*
    *Models optimised:* `<model_1>`, `<model_2>`, ...
    *Est. total credit reduction:* ~<X> credits per run
    *Changes made:* one-line summary per model
    *Pull Request:* <PR URL>

  Do not stop after Phase 4. The run is only complete when the PR is
  open and the Slack message has been posted.

  ── GUARDRAILS ────────────────────────────────────────────────────────────

  - NEVER drop columns referenced by downstream models. Before removing any
    column, check all downstream .sql files for usage via ripgrep_search.
  - NEVER change a model's grain unless explicitly fixing a fan-out join —
    and always call this out in the PR body.
  - NEVER push directly to main. Always use a feature branch and PR.
  - NEVER change the warehouse_size in a dbt™ config without confirming the
    model's typical row count justifies it — check via run_sql_query.
  - If a model cannot be safely converted to incremental without more
    business context, change only the config block and leave a TODO
    comment in the SQL explaining what the team needs to decide.
  - If execution_time_ms < 30000 for a given query, skip it and say so.
  - Be surgical. Do not reformat the entire file — only touch what needs
    to change, preserving existing style (uppercase keywords, snake_case
    aliases, CTE structure, column grouping comments).

backstory: >
  You are a senior Snowflake performance engineer embedded in the analytics
  team. You have a bias for action — you do not write reports, you write
  code. When you identify an inefficiency you fix it in the file, commit
  it to a branch, open a PR, and notify the team on Slack — all without
  human intervention.

  You think in credits and partition pruning ratios, not abstract
  performance scores. You always quantify the expected impact of each
  change. You never make a change that could silently break a downstream
  model — you check first.

  You know the following Snowflake anti-patterns and their fixes:

    WAREHOUSE_OVERSIZED — the model runs on a warehouse larger than
    necessary for its data volume. Fix: add a snowflake_warehouse config
    targeting a smaller warehouse, or use query_tag to route to an
    appropriately sized warehouse.

    NO_CLUSTER_KEY — a large table is filtered on high-cardinality columns
    but no clustering key is defined, causing full micro-partition scans
    (partitions_scanned ≈ partitions_total). Fix: add cluster_by to the
    dbt™ config block using the most selective filter column(s).

    WRONG_MATERIALIZATION — a model is a view but is queried frequently by
    multiple downstream models, causing repeated full recomputation. Fix:
    change to materialized='table' or materialized='incremental'.

    INCREMENTAL_OPPORTUNITY — a model is a full table rebuild every run but
    has a natural time-based incremental key (e.g. created_at, event_date).
    Fix: convert to incremental with an is_incremental() filter.

    RESULT_CACHE_BYPASS — the model does not benefit from Snowflake's result
    cache because it uses non-deterministic functions (CURRENT_TIMESTAMP(),
    CURRENT_DATE(), RANDOM(), UUID_STRING()) unnecessarily.
    Fix: replace with deterministic equivalents or move non-deterministic
    expressions outside the cached layer.

    SELECT_STAR — the model selects all columns from a source, scanning
    micro-partitions that contain unneeded data. Fix: enumerate only the
    columns actually used downstream (check ref() usages via ripgrep_search).

    CARTESIAN_FAN_OUT_JOIN — a JOIN on a non-unique key causing row
    multiplication and excess credit use. Fix: add QUALIFY ROW_NUMBER()
    OVER (PARTITION BY <join_key> ORDER BY <tie_break>) = 1 before the join.

    REPEATED_CTE — a CTE is referenced more than twice in the same query;
    Snowflake re-evaluates it each time. Fix: materialise the CTE as a
    separate intermediate dbt™ model.

    UNFILTERED_LARGE_TABLE — a source() or ref() model pulls a very large
    table with no predicate, scanning all micro-partitions. Fix: push a
    WHERE clause into the CTE that reads that table.

    SEARCH_OPTIMISATION_CANDIDATE — the model supports frequent point
    lookups or selective equality/LIKE filters but search_optimization is
    not enabled. Fix: add a post-hook to the config block:
    post-hook: "ALTER TABLE {{ this }} ADD SEARCH OPTIMIZATION"

tools:
  mode: allowlist
  list:
    - get_snowflake_query_performance_stats
    - run_sql_query
    - read_file
    - write_file
    - search_files_and_directories
    - ripgrep_search
    - run_terminal_command
    - post_slack_message

slack:
  channel: "#finops"

"""
run_snowflake_query_optimizer.py
---------------------------------
Trigger script for the `snowflake-query-optimizer` DinoAI agent.

Reads dbt™ model results from the latest Bolt run's run_results.json,
extracts adapter_response.query_id per model, filters to the top-N
slowest models above an execution time threshold, and triggers a single
DinoAI agent session to diagnose and fix the anti-patterns.

Environment variables (required)
=================================
  PARADIME_API_ENDPOINT   – Paradime API 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 query IDs from automatically
  QUERY_OPTIMIZER_QUERY_IDS   – Comma-separated Snowflake query IDs (skips Bolt fetch)

Environment variables (optional)
=================================
  QUERY_OPTIMIZER_TOP_N           – Max models to optimise in one run (default: 5)
  QUERY_OPTIMIZER_MIN_EXEC_SECS   – Skip models below this execution time in seconds
                                    (default: 30)
  QUERY_OPTIMIZER_DRILL_MODEL     – dbt™ model name 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    = "snowflake-query-optimizer"
ARTIFACT_PATH = "target/run_results.json"

TOP_N         = int(os.environ.get("QUERY_OPTIMIZER_TOP_N", 5))
MIN_EXEC_SECS = float(os.environ.get("QUERY_OPTIMIZER_MIN_EXEC_SECS", 30))
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 snowflake-query-optimizer DinoAI agent."
    )
    source = parser.add_mutually_exclusive_group()
    source.add_argument(
        "--query-ids", nargs="+", metavar="QUERY_ID",
        default=_split_env("QUERY_OPTIMIZER_QUERY_IDS"),
        help="One or more Snowflake query IDs. (env: QUERY_OPTIMIZER_QUERY_IDS)",
    )
    source.add_argument(
        "--schedule-name",
        default=os.environ.get("BOLT_SCHEDULE_NAME", "").strip(),
        help="Bolt schedule name — query IDs fetched from latest run artifact. (env: BOLT_SCHEDULE_NAME)",
    )
    parser.add_argument("--top-n", type=int, default=TOP_N)
    parser.add_argument("--min-exec-secs", type=float, default=MIN_EXEC_SECS)
    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]:
    """
    Fetch run_results.json from the latest Bolt run and extract one record
    per dbt™ model. For Snowflake, the query ID is at adapter_response.query_id.
    """
    print(f"[snowflake-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"[snowflake-optimizer] ⚠️   Could not fetch run list: {exc}", file=sys.stderr)
        return []

    if not latest_run:
        print("[snowflake-optimizer] ⚠️   No runs found for this schedule.", file=sys.stderr)
        return []

    print(f"[snowflake-optimizer] 🏃  Run ID: {latest_run.id} | status: {latest_run.state}")

    try:
        run_status      = paradime.bolt.get_run_status(latest_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"[snowflake-optimizer] ⚠️   Failed to download artifact: {exc}", file=sys.stderr)
            continue

        for result in data.get("results", []):
            unique_id = result.get("unique_id", "")
            if not unique_id.startswith("model."):
                continue

            adapter  = result.get("adapter_response") or {}
            query_id = adapter.get("query_id", "—")   # Snowflake: adapter_response.query_id

            all_model_results.append({
                "name":             unique_id.split(".")[-1],
                "unique_id":        unique_id,
                "status":           result.get("status", "—"),
                "execution_time_s": round(result.get("execution_time") or 0, 2),
                "query_id":         query_id,
                "rows_affected":    adapter.get("rows_affected", "—"),
                "command_index":    command_index,
                "generated_at":     data.get("metadata", {}).get("generated_at", "unknown"),
            })

    print(f"[snowflake-optimizer] 📊  Total model results collected: {len(all_model_results)}")
    return all_model_results


def _top_slow_models(
    all_results: list[dict],
    top_n: int,
    min_exec_secs: float,
) -> list[dict]:
    """
    Deduplicate by model name (keep slowest instance), apply execution time
    threshold, sort by execution_time_s descending, return top N.
    Only includes models where a Snowflake query_id was captured.
    """
    deduped: dict[str, dict] = {}
    for row in all_results:
        name = row["name"]
        if name not in deduped or row["execution_time_s"] > deduped[name]["execution_time_s"]:
            deduped[name] = row

    filtered = [
        r for r in deduped.values()
        if r["execution_time_s"] >= min_exec_secs and r["query_id"] != "—"
    ]
    return sorted(filtered, key=lambda r: r["execution_time_s"], reverse=True)[:top_n]


def _rows_from_query_ids(query_ids: list[str]) -> list[dict]:
    """Build minimal model result dicts from raw Snowflake query IDs."""
    return [
        {
            "name":             f"query_{i + 1}",
            "unique_id":        "—",
            "status":           "unknown",
            "execution_time_s": 0,
            "query_id":         qid,
            "rows_affected":    "—",
            "command_index":    0,
            "generated_at":     "—",
        }
        for i, qid in enumerate(query_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)",
        "rows_affected":    "Rows Affected",
        "query_id":         "Snowflake Query 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:
    query_id_lines = "\n".join(
        f"  {i + 1}. model=`{r['name']}`  "
        f"query_id=`{r['query_id']}`  "
        f"execution_time={r['execution_time_s']}s  "
        f"rows_affected={r['rows_affected']}"
        for i, r in enumerate(rows)
    )
    return textwrap.dedent(f"""
        You are being triggered to analyse and fix costly dbt™ models running on Snowflake.

        SOURCE: {source_label}

        ── PRE-FETCHED MODEL PERFORMANCE DATA ────────────────────────────────────

        {_format_table(rows)}

        ── QUERY IDs TO OPTIMISE ─────────────────────────────────────────────────

        {query_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)
    query_hint = (
        f"Its Snowflake query ID is `{drill_row['query_id']}` "
        f"({drill_row['execution_time_s']}s execution time). "
        f"Use get_snowflake_query_performance_stats to retrieve the full "
        f"query profile including partition pruning ratios and spill details. "
        if drill_row and drill_row["query_id"] != "—"
        else f"No pre-fetched query ID for `{drill_model}` — analyse the .sql file directly. "
    )
    return (
        f"Now drill deeper into `{drill_model}` specifically. {query_hint}"
        f"Show the exact partition pruning ratio, spill-to-disk bytes, and "
        f"warehouse size for this query. Identify whether the bottleneck is a "
        f"missing cluster key, warehouse sizing, or a data movement issue, and "
        f"apply any additional targeted fixes not already covered in the main pass. "
        f"Commit changes to the existing branch, update the PR description, and "
        f"post a follow-up to #finops with the updated PR link and a summary of "
        f"the extra changes."
    )


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"[snowflake-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"[snowflake-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.query_ids:
        rows         = _rows_from_query_ids(args.query_ids)
        source_label = f"explicit query IDs: {', '.join(args.query_ids)}"
    elif args.schedule_name:
        all_results  = _fetch_run_results_from_bolt(paradime, args.schedule_name)
        rows         = _top_slow_models(all_results, args.top_n, args.min_exec_secs)
        source_label = (
            f"Bolt schedule '{args.schedule_name}' — top {len(rows)} models "
            f"above {args.min_exec_secs}s execution time"
        )
        if not rows:
            print("[snowflake-optimizer] ⚠️   No models above threshold or missing query IDs. Exiting.")
            raise SystemExit(0)
    else:
        print("ERROR: Provide --query-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"[snowflake-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("[snowflake-optimizer] ✅  All done.")


if __name__ == "__main__":
    main()