dbt™ Test Maintainer

name: test-maintainer
version: 1

role: >
  dbt™ Test Specialist responsible for writing and updating tests across
  analytics models.

goal: >
  Ensure the target model has appropriate, passing tests committed to the
  repo. Always write meaningful tests — not just for coverage.

  ── STEP 1: READ THE MODEL ────────────────────────────────────────────────

  Read the target model's .sql file in full using read_file. Then find its
  schema YAML (e.g. _<model_name>.yml or schema.yml in the same directory)
  and read the existing column and test definitions.

  ── STEP 2: INFER TEST CANDIDATES ────────────────────────────────────────

  Inspect the model SQL and upstream sources to identify what tests are
  appropriate. For each column, consider:
    - Does the name suggest a primary key? (_id, _key, _sk, _pk) → unique + not_null
    - Is the column used in a JOIN or GROUP BY in a downstream model? → not_null
    - Does a CASE WHEN or WHERE clause imply a finite value set? → accepted_values
    - Does the column reference a known ref() or source()? → relationships

  Use run_sql_query to inspect the actual data distribution for ambiguous
  columns before deciding. Never write a test you cannot justify from the
  SQL or the data.

  ── STEP 3: WRITE THE TESTS ──────────────────────────────────────────────

  Write the inferred tests into the model's schema YAML using write_file.
  Preserve all existing formatting and ordering. Only add — never remove
  or reorder existing test entries. If a test candidate is ambiguous,
  add a comment `# TODO: confirm test logic with owner` on the line above
  the test entry rather than omitting it entirely.

  ── STEP 4: RUN AND ITERATE ──────────────────────────────────────────────

  Run the tests via run_terminal_command:
    dbt test --select <model_name>

  If any tests fail:
    1. Read the failure output carefully.
    2. Diagnose whether the failure is a bad test definition or a genuine
       data quality issue.
    3. If a bad definition: fix the test in the YAML and re-run.
    4. If a genuine data quality issue: add a severity: warn to the test
       config so it does not block the pipeline, and add a comment
       `# DATA ISSUE: <description> — investigate before promoting to error`
       above the test entry.
    5. Iterate until all tests are either passing or set to severity: warn
       with an explanation comment.

  NEVER commit tests that are failing at severity: error.

  ── STEP 5: COMMIT AND NOTIFY ────────────────────────────────────────────

  Stage and commit all changed YAML files to the PR branch:
    git add <changed_yaml_files>
    git commit -m "test: add missing tests for <model_name> [DinoAI]"
    git push origin HEAD:<branch_name>

  Always use `git push origin HEAD:<branch_name>` — never plain `git push`.

  Then post a summary to #data-quality via post_slack_message:

    ✅ *Test Maintainer — <model_name>*
    *Tests added:* <N> (<list of test names>)
    *Tests passing:* <N> / <N>
    *Data quality warnings:* <N> (severity: warn — investigate separately)
    *Committed to branch:* `<branch_name>`

  If invoked as a sub-agent, call notify_parent_session with the same
  summary so the parent agent (e.g. pr-reviewer) can include the result
  in its own report.

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

  - NEVER commit tests that are failing at severity: error.
  - NEVER remove or modify existing passing tests.
  - NEVER guess at accepted_values without checking the actual data via
    run_sql_query first.
  - NEVER reformat or reorder unchanged YAML sections.
  - If dbt test output is ambiguous, re-run with --debug before concluding.
  - If the schema YAML does not exist yet, create it following the dbt™
    schema YAML convention for the project (check an existing file first).

backstory: >
  You are precise and methodical. You read the model and its upstream
  sources before writing tests. You always run dbt test before reporting
  completion. You write tests that mean something — not just not_null on
  every column for the sake of coverage. A test that fires on real bad
  data is worth ten cosmetic tests.

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

slack:
  channel: "#data-quality"

"""
test_maintainer.py
------------------
GitHub Actions driver for the `test-maintainer` DinoAI agent.

Fires one agent session per changed dbt™ model in parallel, polls until
all sessions complete, and posts a summary back to the PR.

Environment variables (required):
    PARADIME_API_ENDPOINT  – Paradime API endpoint
    PARADIME_API_KEY       – Paradime API key
    PARADIME_API_SECRET    – Paradime API secret
    GITHUB_EVENT_PATH      – Set automatically by GitHub Actions
    GITHUB_TOKEN           – Set automatically by GitHub Actions
    GITHUB_REPOSITORY      – Set automatically by GitHub Actions
"""

import json
import os
import sys
import subprocess
import time
import urllib.request
from concurrent.futures import ThreadPoolExecutor, as_completed

from paradime import Paradime
from paradime.apis.dinoai_agents.types import DinoaiAgentRunStatus

AGENT_NAME    = "test-maintainer"
TIMEOUT       = 1800   # 30 minutes per session
POLL_INTERVAL = 20


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 _post_pr_comment(repo: str, pr_number: int, body: str, gh_token: str) -> None:
    req = urllib.request.Request(
        f"https://api.github.com/repos/{repo}/issues/{pr_number}/comments",
        data=json.dumps({"body": body}).encode(),
        headers={
            "Authorization": f"Bearer {gh_token}",
            "Accept":        "application/vnd.github+json",
        },
        method="POST",
    )
    urllib.request.urlopen(req).read()


def _model_name_from_path(path: str) -> str:
    """Extract the dbt™ model name from a file path, e.g. models/marts/fct_orders.sql → fct_orders."""
    return os.path.splitext(os.path.basename(path))[0]


def _run_single_model(
    paradime: Paradime,
    model_name: str,
    sql_path: str,
    branch_name: str,
) -> dict:
    """
    Trigger a test-maintainer session for one model, poll until completion,
    and return a result dict with model name, status, and final output.
    """
    message = (
        f"Add missing dbt™ tests for model `{model_name}` (file: `{sql_path}`).\n\n"
        f"Branch: {branch_name}\n\n"
        f"Read the model SQL and its schema YAML, infer appropriate tests "
        f"(unique, not_null, accepted_values, relationships), write them, "
        f"run `dbt test --select {model_name}` and iterate until all tests "
        f"pass or are set to severity: warn with an explanation. "
        f"Commit the changes to branch `{branch_name}` using "
        f"`git push origin HEAD:{branch_name}` and post a summary to #data-quality."
    )

    trigger    = paradime.dinoai_agents.trigger_run(agent=AGENT_NAME, message=message)
    session_id = trigger.agent_session_id
    print(f"[test-maintainer] 🚀  {model_name} → session {session_id}")

    start = time.time()
    while True:
        run     = paradime.dinoai_agents.get_run(agent_session_id=session_id)
        elapsed = int(time.time() - start)

        if run.status == DinoaiAgentRunStatus.COMPLETED:
            print(f"[test-maintainer] ✅  {model_name} completed ({elapsed}s)")
            return {
                "model":   model_name,
                "status":  "completed",
                "output":  run.messages[-1].content if run.messages else "(no output)",
                "session": session_id,
            }

        if run.status == DinoaiAgentRunStatus.FAILED:
            last = run.messages[-1].content if run.messages else "no output"
            print(f"[test-maintainer] ❌  {model_name} failed ({elapsed}s)")
            return {
                "model":   model_name,
                "status":  "failed",
                "output":  last,
                "session": session_id,
            }

        if elapsed > TIMEOUT:
            print(f"[test-maintainer] ⏱  {model_name} timed out")
            return {
                "model":   model_name,
                "status":  "timeout",
                "output":  f"Timed out after {TIMEOUT}s",
                "session": session_id,
            }

        print(f"[test-maintainer]    {model_name} — {run.status.value} ({elapsed}s) … {POLL_INTERVAL}s")
        time.sleep(POLL_INTERVAL)


def main() -> None:
    paradime  = _build_client()
    gh_token  = os.environ.get("GITHUB_TOKEN", "")
    repo      = os.environ.get("GITHUB_REPOSITORY", "")

    # 1. Read PR context
    with open(os.environ["GITHUB_EVENT_PATH"]) as f:
        event = json.load(f)

    pr          = event["pull_request"]
    pr_number   = pr["number"]
    pr_title    = pr["title"]
    branch_name = pr["head"]["ref"]
    base_sha    = pr["base"]["sha"]
    head_sha    = pr["head"]["sha"]

    # 2. Find changed .sql files
    changed = subprocess.check_output(
        ["git", "diff", "--name-only", f"{base_sha}...{head_sha}"]
    ).decode().splitlines()

    sql_files = [f for f in changed if f.endswith(".sql") and "models/" in f]

    if not sql_files:
        print("[test-maintainer] No dbt™ model SQL files changed. Skipping.")
        raise SystemExit(0)

    models = [(_model_name_from_path(p), p) for p in sql_files]

    print(f"[test-maintainer] Found {len(models)} changed model(s):")
    for name, path in models:
        print(f"  - {name}  ({path})")

    # 3. Post "started" comment to PR
    if gh_token:
        model_list = "\n".join(f"- `{name}` (`{path}`)" for name, path in models)
        _post_pr_comment(
            repo, pr_number,
            f"🧪 **DinoAI test maintainer started** — PR #{pr_number}: {pr_title}\n\n"
            f"Writing and running missing tests for {len(models)} model(s):\n{model_list}\n\n"
            f"Changes will be committed to `{branch_name}`. Check `#data-quality` for live updates.",
            gh_token,
        )

    # 4. Trigger all sessions in parallel
    results = []
    with ThreadPoolExecutor(max_workers=len(models)) as executor:
        futures = {
            executor.submit(_run_single_model, paradime, name, path, branch_name): name
            for name, path in models
        }
        for future in as_completed(futures):
            results.append(future.result())

    # 5. Post completion comment to PR
    if gh_token:
        completed = [r for r in results if r["status"] == "completed"]
        failed    = [r for r in results if r["status"] != "completed"]

        sections = []
        for r in results:
            icon = "✅" if r["status"] == "completed" else "❌"
            sections.append(
                f"### {icon} `{r['model']}`\n\n"
                f"{r['output']}\n\n"
                f"*Session: `{r['session']}`*"
            )

        footer = (
            f"**{len(completed)}/{len(results)} models** had tests written and committed to `{branch_name}`."
        )
        if failed:
            footer += f"\n\n⚠️ {len(failed)} session(s) did not complete — check `#data-quality` for details."

        _post_pr_comment(
            repo, pr_number,
            f"🧪 **DinoAI test maintainer — complete**\n\n"
            + "\n\n---\n\n".join(sections)
            + f"\n\n---\n\n{footer}",
            gh_token,
        )

    # Exit with error if any session failed so the workflow is marked red
    if any(r["status"] != "completed" for r in results):
        raise SystemExit(1)


if __name__ == "__main__":
    main()

name: DinoAI dbt™ test maintainer

on:
  pull_request:
    types: [opened, synchronize, ready_for_review]
    paths:
      - "models/**/*.sql"

permissions:
  pull-requests: write
  contents: write       # required so the agent can push test commits to the PR branch

jobs:
  write-tests:
    runs-on: ubuntu-latest
    timeout-minutes: 35
    steps:
      - uses: actions/checkout@v4
        with:
          fetch-depth: 0
          ref: ${{ github.head_ref }}   # check out the PR branch so the agent can push to it

      - uses: actions/setup-python@v5
        with:
          python-version: "3.11"

      - name: Install Paradime SDK
        run: pip install paradime-io

      - name: Configure git for agent commits
        run: |
          git config user.name  "DinoAI Test Maintainer"
          git config user.email "dinoai-tests@noreply.github.com"

      - name: Run test maintainer
        env:
          PARADIME_API_ENDPOINT: ${{ secrets.PARADIME_API_ENDPOINT }}
          PARADIME_API_KEY:      ${{ secrets.PARADIME_API_KEY }}
          PARADIME_API_SECRET:   ${{ secrets.PARADIME_API_SECRET }}
          GITHUB_TOKEN:          ${{ secrets.GITHUB_TOKEN }}
        run: python scripts/test_maintainer.py