PostgreSQL Slow Query AI Index Healer with DuckDB and HypoPG

id: postgres-slow-query-ai-index-healer
namespace: company.data

description: |
  Continuously profiles PostgreSQL for slow SELECT queries via pg_stat_statements, scores and ranks
  bottlenecks in-memory using DuckDB, synthesizes non-locking B-tree indexes with an LLM, verifies
  cost reductions via EXPLAIN and HypoPG, and dispatches GitHub issues and Slack notifications.

variables:
  llm_endpoint: https://api.groq.com/openai/v1/chat/completions

inputs:
  - id: pg_url
    type: STRING
    defaults: "jdbc:postgresql://postgres-test:5432/postgres"
    description: "JDBC connection string for the target PostgreSQL instance."

  - id: pg_user
    type: STRING
    defaults: "postgres"
    description: "PostgreSQL role with pg_read_all_stats privileges."

  - id: min_mean_exec_time_ms
    type: FLOAT
    defaults: 1.0
    description: "Minimum mean execution latency (in ms) required to flag a slow query."

  - id: min_calls
    type: INT
    defaults: 25
    description: "Minimum call count required before evaluating a query."

  - id: max_queries_to_analyze
    type: INT
    defaults: 3
    description: "Maximum number of top-ranked slow queries to remediate per run."

  - id: github_repository
    type: STRING
    defaults: "your-org/your-repo"
    description: "Target GitHub repository (owner/repo) for remediation issues."

  - id: llm_model
    type: STRING
    defaults: "openai/gpt-oss-20b"
    description: "OpenAI-compatible model identifier served by the LLM endpoint."

tasks:
  - id: extract_slow_queries
    type: io.kestra.plugin.jdbc.postgresql.Query
    description: "Extract top slow SELECT statements from pg_stat_statements."
    url: "{{ inputs.pg_url }}"
    username: "{{ inputs.pg_user }}"
    password: "{{ secret('POSTGRES_PASSWORD') }}"
    fetchType: FETCH
    sql: |
      SELECT
        CAST(queryid AS text) AS query_id,
        LEFT(regexp_replace(query, '\s+', ' ', 'g'), 8000) AS query,
        calls,
        CAST(ROUND(CAST(mean_exec_time AS numeric), 2) AS double precision) AS mean_exec_time_ms,
        CAST(ROUND(CAST(stddev_exec_time AS numeric), 2) AS double precision) AS stddev_exec_time_ms,
        CAST(ROUND(CAST(total_exec_time AS numeric), 2) AS double precision) AS total_exec_time_ms,
        rows AS rows_returned,
        shared_blks_hit,
        shared_blks_read
      FROM pg_stat_statements
      WHERE dbid = (SELECT oid FROM pg_database WHERE datname = current_database())
        AND query ILIKE 'select%'
        AND query NOT ILIKE '%pg_stat_statements%'
        AND query NOT ILIKE '%pg_catalog.%'
        AND mean_exec_time >= {{ inputs.min_mean_exec_time_ms }}
        AND calls >= {{ inputs.min_calls }}
      ORDER BY total_exec_time DESC
      LIMIT 50

  - id: check_slow_queries_exist
    type: io.kestra.plugin.core.flow.If
    description: "Proceed with catalog extraction and ranking only if slow queries exist."
    condition: "{{ outputs.extract_slow_queries.size > 0 }}"
    then:
      - id: fetch_index_catalog
        type: io.kestra.plugin.jdbc.postgresql.Query
        description: "Fetch existing user-defined indexes to prevent duplicate index
          proposals."
        url: "{{ inputs.pg_url }}"
        username: "{{ inputs.pg_user }}"
        password: "{{ secret('POSTGRES_PASSWORD') }}"
        fetchType: FETCH_ONE
        sql: |
          SELECT LEFT(COALESCE(string_agg(indexdef, chr(10) ORDER BY tablename, indexname), ''), 6000) AS index_catalog
          FROM pg_indexes
          WHERE schemaname NOT IN ('pg_catalog', 'information_schema', 'pg_toast')

      - id: rank_bottlenecks
        type: io.kestra.plugin.jdbc.duckdb.Query
        description: "Compute composite severity scores in DuckDB using latency, buffer
          misses, and jitter."
        fetchType: FETCH
        inputFiles:
          slow_queries.json: "{{ outputs.extract_slow_queries.rows | toJson }}"
        sql: |
          WITH metrics AS (
            SELECT
              CAST(query_id AS VARCHAR) AS query_id,
              query,
              CAST(calls AS BIGINT) AS calls,
              mean_exec_time_ms,
              stddev_exec_time_ms,
              total_exec_time_ms,
              rows_returned,
              CASE WHEN shared_blks_hit + shared_blks_read = 0 THEN 1.0
                   ELSE shared_blks_hit * 1.0 / (shared_blks_hit + shared_blks_read) END AS cache_hit_ratio,
              CASE WHEN mean_exec_time_ms = 0 THEN 0.0
                   ELSE stddev_exec_time_ms / mean_exec_time_ms END AS latency_variability
            FROM read_json_auto('{{ workingDir }}/slow_queries.json')
          ),
          scored AS (
            SELECT
              *,
              ROUND(
                (mean_exec_time_ms * calls / 1000.0)
                * (1 + (1 - cache_hit_ratio))
                * (1 + LEAST(latency_variability, 1.0)),
                2
              ) AS severity_score
            FROM metrics
          )
          SELECT
            query_id,
            query,
            calls,
            mean_exec_time_ms,
            CAST(ROUND(mean_exec_time_ms) AS BIGINT) AS mean_ms_rounded,
            stddev_exec_time_ms,
            total_exec_time_ms,
            rows_returned,
            ROUND(cache_hit_ratio, 3) AS cache_hit_ratio,
            severity_score,
            ROW_NUMBER() OVER (ORDER BY severity_score DESC) AS severity_rank
          FROM scored
          ORDER BY severity_score DESC
          LIMIT {{ inputs.max_queries_to_analyze }}

      - id: heal_each_bottleneck
        type: io.kestra.plugin.core.flow.Loop
        description: "Iterate over top-ranked bottlenecks with bounded concurrency."
        values: "{{ outputs.rank_bottlenecks.rows }}"
        concurrencyLimit: 2
        tasks:
          - id: check_already_tracked
            type: io.kestra.plugin.core.kv.Get
            description: "Check Kestra KV Store to avoid duplicate alerts for the same query
              ID."
            key: "pg_healer_{{ fromJson(item.value).query_id }}"
            errorOnMissing: false

          - id: gate_new_bottleneck
            type: io.kestra.plugin.core.flow.If
            condition: "{{ (outputs.check_already_tracked.value ?? '') == '' }}"
            then:
              - id: ai_synthesis
                type: io.kestra.plugin.core.http.Request
                description: "Request a non-locking B-tree index DDL and rollback plan from the
                  LLM."
                uri: "{{ vars.llm_endpoint }}"
                method: POST
                contentType: application/json
                timeout: PT2M
                headers:
                  Authorization: "Bearer {{ secret('LLM_API_KEY') }}"
                body: |
                  {%- set q = fromJson(item.value) -%}
                  {{ {
                    "model": inputs.llm_model,
                    "temperature": 0.1,
                    "max_tokens": 900,
                    "response_format": {"type": "json_object"},
                    "messages": [
                      {
                        "role": "system",
                        "content": "You are a Principal PostgreSQL Database Engineer. Design the single best B-tree index to fix the slow query. Respond with ONE JSON object only, using keys: bottleneck_reason (string), index_ddl (string, non-locking CREATE INDEX CONCURRENTLY IF NOT EXISTS statement), rollback_ddl (string, DROP INDEX CONCURRENTLY IF EXISTS statement), expected_impact (string)."
                      },
                      {
                        "role": "user",
                        "content": "Query ID: " ~ q.query_id ~ ". Mean latency: " ~ q.mean_exec_time_ms ~ " ms over " ~ q.calls ~ " calls. SQL: " ~ q.query ~ " EXISTING INDEXES: " ~ (outputs.fetch_index_catalog.row.index_catalog | default('none'))
                      }
                    ]
                  } | toJson }}

              - id: parse_ai_response
                type: io.kestra.plugin.core.output.OutputValues
                description: "Validate that both creation and rollback DDL statements use
                  CONCURRENTLY."
                values:
                  fix: |-
                    {%- set raw = outputs.ai_synthesis.body -%}
                    {%- set envelope = raw is map ? raw : fromJson(raw) -%}
                    {%- set content = envelope.choices[0].message.content -%}
                    {%- set answer = content is map ? content : fromJson(content) -%}
                    {%- set ddl = (answer.index_ddl ?? '') | trim -%}
                    {%- set rollback = (answer.rollback_ddl ?? '') | trim -%}
                    {%- set safe = (ddl | upper | startsWith('CREATE INDEX CONCURRENTLY ')) and (rollback | upper | startsWith('DROP INDEX CONCURRENTLY ')) -%}
                    {{ {
                      "bottleneck_reason": answer.bottleneck_reason ?? 'n/a',
                      "index_ddl": ddl,
                      "rollback_ddl": rollback,
                      "expected_impact": answer.expected_impact ?? 'n/a',
                      "guardrails_passed": safe
                    } | toJson }}

              - id: verify_with_explain
                type: io.kestra.plugin.jdbc.postgresql.Queries
                description: "Compare generic query plans before and after virtual index
                  creation via HypoPG."
                runIf: "{{ fromJson(outputs.parse_ai_response.values.fix).guardrails_passed }}"
                allowFailure: true
                url: "{{ inputs.pg_url }}"
                username: "{{ inputs.pg_user }}"
                password: "{{ secret('POSTGRES_PASSWORD') }}"
                fetchType: FETCH
                sql: |
                  {%- set q = fromJson(item.value) -%}
                  {%- set fix = fromJson(outputs.parse_ai_response.values.fix) -%}
                  EXPLAIN (GENERIC_PLAN) {{ q.query | trim | replace({';': ''}) }};
                  SELECT * FROM hypopg_create_index('{{ fix.index_ddl | replace({'CONCURRENTLY ': '', 'IF NOT EXISTS ': '', ';': '', "'": "''"}) }}');
                  EXPLAIN (GENERIC_PLAN) {{ q.query | trim | replace({';': ''}) }};

              - id: score_verification
                type: io.kestra.plugin.jdbc.duckdb.Query
                description: "Parse EXPLAIN cost estimates in DuckDB to confirm >= 10% planner
                  cost reduction."
                runIf: "{{ (outputs.verify_with_explain.outputs ?? []) | length == 3 }}"
                allowFailure: true
                fetchType: FETCH
                inputFiles:
                  plans.json: |
                    {{ {
                      "before_rows": outputs.verify_with_explain.outputs[0].rows,
                      "after_rows": outputs.verify_with_explain.outputs[2].rows
                    } | toJson }}
                sql: |
                  WITH p AS (
                    SELECT
                      CAST(before_rows[1]."QUERY PLAN" AS VARCHAR) AS plan_before,
                      CAST(after_rows[1]."QUERY PLAN" AS VARCHAR) AS plan_after
                    FROM read_json_auto('{{ workingDir }}/plans.json')
                  ),
                  c AS (
                    SELECT
                      plan_before,
                      plan_after,
                      TRY_CAST(regexp_extract(plan_before, 'cost=[0-9.]+\.\.([0-9.]+)', 1) AS DOUBLE) AS cost_before,
                      TRY_CAST(regexp_extract(plan_after, 'cost=[0-9.]+\.\.([0-9.]+)', 1) AS DOUBLE) AS cost_after
                    FROM p
                  )
                  SELECT
                    plan_before,
                    plan_after,
                    cost_before,
                    cost_after,
                    ROUND(100.0 * (cost_before - cost_after) / NULLIF(cost_before, 0), 1) AS improvement_pct,
                    CASE
                      WHEN cost_before IS NULL OR cost_after IS NULL THEN 'UNVERIFIED'
                      WHEN cost_after < cost_before * 0.9 THEN 'CONFIRMED'
                      ELSE 'NOT_CONFIRMED'
                    END AS verdict
                  FROM c

              - id: route_remediation
                type: io.kestra.plugin.core.flow.If
                condition: "{{ fromJson(outputs.parse_ai_response.values.fix).guardrails_passed
                  }}"
                then:
                  - id: fan_out_remediation
                    type: io.kestra.plugin.core.flow.Parallel
                    tasks:
                      - id: open_github_issue
                        type: io.kestra.plugin.core.http.Request
                        description: "Open a GitHub issue with telemetry, AI diagnosis, and verified
                          DDL."
                        uri: "https://api.github.com/repos/{{ inputs.github_repository }}/issues"
                        method: POST
                        contentType: application/json
                        headers:
                          Accept: application/vnd.github+json
                          Authorization: "Bearer {{ secret('GITHUB_TOKEN') }}"
                          X-GitHub-Api-Version: "2022-11-28"
                          User-Agent: kestra-postgres-index-healer
                        body: |
                          {%- set q = fromJson(item.value) -%}
                          {%- set fix = fromJson(outputs.parse_ai_response.values.fix) -%}
                          {%- set vr = outputs.score_verification.rows[0] ?? null -%}
                          {%- set verdict = vr is null ? 'UNVERIFIED' : vr.verdict -%}
                          {
                            "title": {{ ('[DB Healer] Slow Query Remediation (Query ' ~ q.query_id ~ ' - ' ~ q.mean_ms_rounded ~ 'ms)') | toJson }},
                            "labels": ["db-healer", "performance", "automated"],
                            "body": {% filter toJson %}## Slow query remediation proposal

                          - Query ID: `{{ q.query_id }}`
                          - Mean latency: **{{ q.mean_exec_time_ms }} ms**
                          - Calls: {{ q.calls }}

                          ### Slow query
                          ```sql
                          {{ q.query }}
                          ```

                          ### AI Diagnosis
                          {{ fix.bottleneck_reason }}

                          ### Migration DDL
                          ```sql
                          {{ fix.index_ddl }}
                          ```

                          ### Rollback DDL
                          ```sql
                          {{ fix.rollback_ddl }}
                          ```

                          **Verdict: {{ verdict }}**
                          {% endfilter %}
                          }

                      - id: notify_slack
                        type: io.kestra.plugin.slack.notifications.SlackIncomingWebhook
                        description: "Send real-time remediation alert to Slack."
                        allowFailure: true
                        url: "{{ secret('SLACK_WEBHOOK_URL') }}"
                        payload: |
                          {%- set q = fromJson(item.value) -%}
                          {%- set fix = fromJson(outputs.parse_ai_response.values.fix) -%}
                          {{ {
                            "text": "⚡ PostgreSQL Slow Query Remediated - Query " ~ q.query_id ~ " (AI DDL: `" ~ fix.index_ddl ~ "`)"
                          } | toJson }}

                  - id: remember_query
                    type: io.kestra.plugin.core.kv.Set
                    description: "Record query_id in KV Store with a 7-day TTL to prevent duplicate
                      issues."
                    key: "pg_healer_{{ fromJson(item.value).query_id }}"
                    value: "{{ execution.id }}"
                    kvType: STRING
                    ttl: P7D

    else:
      - id: log_database_healthy
        type: io.kestra.plugin.core.log.Log
        message: "All PostgreSQL queries are executing within SLA threshold."

triggers:
  - id: hourly_schedule
    type: io.kestra.plugin.core.trigger.Schedule
    cron: "0 * * * *"