dbt Test Failure AI Root Cause and Remediation Advisor

id: dbt-test-failure-ai-root-cause-advisor
namespace: company.analytics

description: |
  Automated analytics engineering triage pipeline that ingests dbt test failure artifacts,
  diagnoses root causes such as unhandled nulls, duplicate keys, and join fan-outs using Google Gemini,
  generates SQL remediation diffs, and alerts data teams in Slack.

triggers:
  - id: scheduled_sweep
    type: io.kestra.plugin.core.trigger.Schedule
    description: Morning test failure sweep following nightly dbt transformations.
      Shipped disabled by default.
    cron: "0 7 * * 1-5"
    disabled: true

  - id: webhook_dbt_test_failed
    type: io.kestra.plugin.core.trigger.Webhook
    description: Authenticated webhook for CI/CD or dbt Cloud failure notifications.
    key: "{{ secret('WEBHOOK_KEY') }}"

inputs:
  - id: model_name
    type: STRING
    defaults: fct_customer_orders
    description: Target dbt model that failed validation.

  - id: test_name
    type: STRING
    defaults: unique_fct_customer_orders_order_id
    description: Name of the failed dbt generic or singular test.

  - id: failure_count
    type: INT
    defaults: 14
    description: Number of failing records returned by the dbt test query.

  - id: compiled_test_sql
    type: STRING
    defaults: |
      select order_id, count(*) as n_records
      from analytics.fct_customer_orders
      group by order_id
      having count(*) > 1
    description: Compiled SQL query executed by dbt that yielded failing rows.

  - id: upstream_model_sql
    type: STRING
    defaults: |
      select
          o.order_id,
          o.customer_id,
          o.order_date,
          p.payment_method,
          o.total_amount
      from raw.orders o
      left join raw.payments p on o.order_id = p.order_id
    description: Upstream model SQL query containing the source transformation logic.

tasks:
  - id: diagnose_dbt_failure
    type: io.kestra.plugin.ai.agent.AIAgent
    description: Uses Gemini 2.5 Flash to analyze test compilation and upstream SQL,
      determining the exact cause of data duplication or null anomalies.
    provider:
      type: io.kestra.plugin.ai.provider.GoogleGemini
      apiKey: "{{ secret('GEMINI_API_KEY') }}"
      modelName: gemini-2.5-flash
    configuration:
      temperature: 0.1
      maxToken: 4096
      responseFormat:
        type: JSON
        jsonSchema:
          type: object
          required:
            - model_name
            - test_name
            - severity
            - failure_category
            - root_cause_explanation
            - recommended_sql_diff
            - prevention_guidance
          properties:
            model_name:
              type: string
            test_name:
              type: string
            severity:
              type: string
              enum:
                - CRITICAL
                - HIGH
                - MEDIUM
            failure_category:
              type: string
              enum:
                - JOIN_FANOUT
                - NULL_SPIKE
                - RELATIONSHIP_MISMATCH
                - SYNTAX_OR_SCHEMA_CHANGE
            root_cause_explanation:
              type: string
            recommended_sql_diff:
              type: string
            prevention_guidance:
              type: string
    prompt: |
      You are a Staff Analytics Engineer diagnosing a failing dbt test in a production data warehouse.
      Analyze the failed test execution and upstream model code below:

      Model: {{ inputs.model_name }}
      Failed Test: {{ inputs.test_name }}
      Failing Record Count: {{ inputs.failure_count }}

      Compiled Test SQL:
      \"\"\"
      {{ inputs.compiled_test_sql }}
      \"\"\"

      Upstream Model SQL:
      \"\"\"
      {{ inputs.upstream_model_sql }}
      \"\"\"

      Tasks:
      1. Pinpoint why the test failed (e.g. one-to-many join fan-out between orders and payments causing order_id duplication).
      2. Categorize the failure and assign an objective severity rating based on downstream BI impact.
      3. Provide the corrected SQL transformation code that resolves the cardinality or null issue.
      4. Suggest a prevention strategy (such as dbt unit tests or pre-join CTE aggregations).

  - id: render_remediation_brief
    type: io.kestra.plugin.scripts.python.Script
    description: Formats structured AI diagnosis and recommended SQL diff into a
      markdown engineering brief stored in execution artifacts.
    taskRunner:
      type: io.kestra.plugin.core.runner.Process
    script: |
      import json

      raw_response = """{{ outputs.diagnose_dbt_failure.text }}"""
      data = json.loads(raw_response, strict=False)

      model = data.get("model_name", "{{ inputs.model_name }}")
      test = data.get("test_name", "{{ inputs.test_name }}")
      sev = data.get("severity", "HIGH")
      cat = data.get("failure_category", "UNKNOWN")
      explanation = data.get("root_cause_explanation", "")
      sql_diff = data.get("recommended_sql_diff", "")
      prevention = data.get("prevention_guidance", "")

      markdown = f"""# DBT TEST FAILURE ROOT CAUSE & REMEDIATION BRIEF
      ## Incident Metadata
      - **Target Model**: `{model}`
      - **Failed Test**: `{test}`
      - **Severity**: **{sev}**
      - **Category**: `{cat}`

      ## Root Cause Diagnosis
      {explanation}

      ## Recommended SQL Remediation
      ```sql
      {sql_diff}
      ```

      ## Prevention & Best Practices
      {prevention}
      """

      with open("dbt-remediation-brief.md", "w") as f:
          f.write(markdown)
      print("Generated dbt-remediation-brief.md")
    outputFiles:
      - dbt-remediation-brief.md

  - id: triage_routing_gate
    type: io.kestra.plugin.core.flow.If
    description: Escalates high and critical severity model test failures to urgent
      Slack triage channels.
    condition: "{{ json(outputs.diagnose_dbt_failure.text).severity == 'CRITICAL' ||
      json(outputs.diagnose_dbt_failure.text).severity == 'HIGH' }}"
    then:
      - id: alert_urgent_slack
        type: io.kestra.plugin.slack.notifications.SlackIncomingWebhook
        description: Alerts urgent analytics engineering channel with root cause and
          remediation diff.
        url: "{{ secret('SLACK_WEBHOOK_URL') }}"
        payload: |
          {
            "channel": "#analytics-alerts-urgent",
            "text": "🚨 *dbt Test Failure*: Model `{{ inputs.model_name }}` failed test `{{ inputs.test_name }}` ({{ inputs.failure_count }} failing rows).\n*Category*: `{{ json(outputs.diagnose_dbt_failure.text).failure_category }}` (Severity: {{ json(outputs.diagnose_dbt_failure.text).severity }})\n*Diagnosis*: {{ json(outputs.diagnose_dbt_failure.text).root_cause_explanation }}\nRemediation brief generated in Kestra storage."
          }
    else:
      - id: notify_standard_feed
        type: io.kestra.plugin.slack.notifications.SlackIncomingWebhook
        description: Dispatches standard priority notice to data team log channel.
        url: "{{ secret('SLACK_WEBHOOK_URL') }}"
        payload: |
          {
            "channel": "#data-team-feed",
            "text": "⚠️ *dbt Test Warning*: Model `{{ inputs.model_name }}` test failed. Diagnosis completed with severity `{{ json(outputs.diagnose_dbt_failure.text).severity }}`."
          }

errors:
  - id: alert_on_failure
    type: io.kestra.plugin.slack.notifications.SlackIncomingWebhook
    description: Alerts on-call engineering channel if diagnosis or report rendering
      encounters an unhandled error.
    url: "{{ secret('SLACK_WEBHOOK_URL') }}"
    payload: |
      {
        "channel": "#analytics-ops",
        "text": "⚠️ *Pipeline Error*: dbt root cause advisor `{{ flow.id }}` failed on execution `{{ execution.id }}`."
      }