Automated Data Contract Diff and Schema Migration Generator

id: data-contract-schema-migration-generator
namespace: company.data

description: |
  Ingest data contract specifications, compute schema drift against target PostgreSQL tables,
  synthesize backward-compatible SQL migrations and dbt models with AI, and validate syntax
  in transactional dry-runs before publishing deployment artifacts.

concurrency:
  limit: 1
  behavior: QUEUE

inputs:
  - id: schema_name
    type: STRING
    defaults: "public"
    description: Target PostgreSQL database schema holding the table.

  - id: table_name
    type: STRING
    defaults: "orders"
    description: Target table name governed by the data contract.

  - id: data_contract
    type: JSON
    defaults:
      table: "orders"
      version: "2.1.0"
      description: "Customer transaction records with fulfillment attribution"
      columns:
        - name: "id"
          type: "bigint"
          description: "Primary unique order identifier"
          nullable: false
        - name: "customer_id"
          type: "integer"
          description: "Foreign key reference to registered customer"
          nullable: false
        - name: "total_amount"
          type: "numeric(10,2)"
          description: "Total order financial amount in transaction currency"
          nullable: false
        - name: "status"
          type: "varchar(32)"
          description: "Order fulfillment lifecycle state"
          nullable: false
          defaults: "'pending'"
        - name: "currency"
          type: "varchar(3)"
          description: "ISO 4217 three-letter currency code"
          nullable: false
          defaults: "'USD'"
        - name: "metadata"
          type: "jsonb"
          description: "Flexible unstructured payload for tags and channel attribution"
          nullable: true
          defaults: "'{}'::jsonb"
        - name: "created_at"
          type: "timestamptz"
          description: "Order creation timestamp in UTC"
          nullable: false
          defaults: "now()"
    description: Target data contract schema definition specifying required fields
      and constraints.

  - id: postgres_url
    type: STRING
    defaults: "jdbc:postgresql://127.0.0.1:5432/postgres"
    description: JDBC connection URL to the target PostgreSQL database.

  - id: postgres_user
    type: STRING
    defaults: "postgres"
    description: PostgreSQL user with schema inspection and DDL privileges.

  - id: ai_model
    type: STRING
    defaults: "gemini-2.5-flash"
    description: Google Gemini model used for backward-compatible migration and dbt
      model synthesis.

  - id: dry_run
    type: BOOL
    defaults: true
    description: When true, audits drift and validates migration syntax in
      rolled-back transactions without applying persistent DDL changes.

tasks:
  - id: inspect_target_table_schema
    type: io.kestra.plugin.jdbc.postgresql.Query
    description: Queries PostgreSQL information_schema to discover current table
      columns, types, and defaults.
    url: "{{ inputs.postgres_url }}"
    username: "{{ inputs.postgres_user }}"
    password: "{{ secret('POSTGRES_PASSWORD') }}"
    sql: |
      SELECT
        column_name,
        data_type,
        udt_name,
        is_nullable,
        column_default
      FROM information_schema.columns
      WHERE table_schema = '{{ inputs.schema_name }}'
        AND table_name = '{{ inputs.table_name }}'
      ORDER BY ordinal_position;
    fetchType: FETCH

  - id: compute_contract_diff
    type: io.kestra.plugin.scripts.python.Script
    description: Diffs target data contract against live PostgreSQL catalog and
      categorizes schema evolution.
    containerImage: python:3.11-slim
    taskRunner:
      type: io.kestra.plugin.scripts.runner.docker.Docker
    beforeCommands:
      - pip install -q kestra
    outputFiles:
      - contract_diff.json
    env:
      CONTRACT_JSON: "{{ inputs.data_contract | json }}"
      LIVE_COLUMNS_JSON: "{{ outputs.inspect_target_table_schema.rows | json }}"
      TABLE_NAME: "{{ inputs.table_name }}"
    script: |
      import os
      import json
      from kestra import Kestra

      contract_str = os.environ.get("CONTRACT_JSON", "{}")
      live_columns_str = os.environ.get("LIVE_COLUMNS_JSON", "[]")
      table_name = os.environ.get("TABLE_NAME", "orders")

      contract = json.loads(contract_str)
      live_columns = json.loads(live_columns_str) if live_columns_str else []

      contract_columns = {col["name"].lower(): col for col in contract.get("columns", [])}
      live_col_map = {row["column_name"].lower(): row for row in live_columns}

      added_columns = []
      modified_columns = []
      deprecated_columns = []

      for col_name, col_def in contract_columns.items():
          if col_name not in live_col_map:
              added_columns.append(col_def)
          else:
              live_col = live_col_map[col_name]
              contract_type = col_def.get("type", "").lower()
              live_type = live_col.get("data_type", "").lower()
              udt_name = live_col.get("udt_name", "").lower()

              type_mismatch = (contract_type not in live_type) and (contract_type not in udt_name)
              if type_mismatch:
                  modified_columns.append({
                      "name": col_name,
                      "contract_type": col_def.get("type"),
                      "database_type": live_type,
                      "description": col_def.get("description", "")
                  })

      for col_name, live_col in live_col_map.items():
          if col_name not in contract_columns:
              deprecated_columns.append({
                  "name": col_name,
                  "database_type": live_col.get("data_type")
              })

      drift_detected = len(added_columns) > 0 or len(modified_columns) > 0

      diff_report = {
          "table": table_name,
          "drift_detected": drift_detected,
          "added_columns": added_columns,
          "modified_columns": modified_columns,
          "deprecated_columns": deprecated_columns,
          "summary": {
              "added_count": len(added_columns),
              "modified_count": len(modified_columns),
              "deprecated_count": len(deprecated_columns)
          }
      }

      with open("contract_diff.json", "w") as f:
          json.dump(diff_report, f, indent=2)

      Kestra.outputs({
          "drift_detected": drift_detected,
          "added_count": len(added_columns),
          "modified_count": len(modified_columns),
          "deprecated_count": len(deprecated_columns),
          "added_columns": added_columns,
          "modified_columns": modified_columns
      })

  - id: evaluate_contract_evolution
    type: io.kestra.plugin.core.flow.If
    description: Triggers AI migration synthesis and dry-run compiler verification
      if schema drift is detected.
    condition: "{{ outputs.compute_contract_diff.vars.drift_detected }}"
    then:
      - id: generate_migration_and_dbt_models
        type: io.kestra.plugin.ai.completion.ChatCompletion
        description: Synthesizes backward-compatible PostgreSQL migration SQL and
          updated dbt schema YAML.
        provider:
          type: io.kestra.plugin.ai.provider.GoogleGemini
          apiKey: "{{ secret('GEMINI_API_KEY') }}"
          modelName: "{{ inputs.ai_model }}"
        messages:
          - type: SYSTEM
            content: |
              You are a Principal Database Reliability Engineer and dbt Analytics Architect.
              Your task is to resolve schema drift between an authoritative Data Contract and a live PostgreSQL table.

              Guidelines for migration_sql:
              1. Must be strictly backward-compatible. Never drop existing columns or add NOT NULL constraints without safe DEFAULT values.
              2. Use idempotent syntax: ALTER TABLE <schema>.<table> ADD COLUMN IF NOT EXISTS <name> <type> DEFAULT <default>;
              3. Do NOT include BEGIN or COMMIT in migration_sql (these are managed by the execution runner).
              4. Ensure all statements are safe for PostgreSQL 14+.

              Guidelines for dbt_schema_yaml:
              1. Generate valid dbt schema YAML (version: 2, models block).
              2. Include all contract columns, their descriptions, and data tests (e.g. not_null, unique).

              Return valid JSON strictly conforming to the requested schema.
          - type: USER
            content: |
              Target Table: {{ inputs.schema_name }}.{{ inputs.table_name }}
              Added Columns:
              {{ outputs.compute_contract_diff.vars.added_columns | json }}
              Modified Columns:
              {{ outputs.compute_contract_diff.vars.modified_columns | json }}
              Full Contract Specification:
              {{ inputs.data_contract | json }}
        configuration:
          responseFormat:
            type: JSON
            jsonSchema:
              type: object
              properties:
                migration_sql:
                  type: string
                dbt_schema_yaml:
                  type: string
                backward_compatibility_notes:
                  type: string
                risk_level:
                  type: string
              required:
                - migration_sql
                - dbt_schema_yaml
                - backward_compatibility_notes
                - risk_level

      - id: validate_migration_syntax
        type: io.kestra.plugin.jdbc.postgresql.Queries
        description: Executes transactional dry-run against PostgreSQL compiler to
          verify syntax without committing table mutations.
        url: "{{ inputs.postgres_url }}"
        username: "{{ inputs.postgres_user }}"
        password: "{{ secret('POSTGRES_PASSWORD') }}"
        sql: |
          BEGIN;
          {{ outputs.generate_migration_and_dbt_models.jsonOutput.migration_sql }}
          ROLLBACK;

      - id: persist_migration_sql
        type: io.kestra.plugin.core.storage.Write
        description: Writes backward-compatible SQL migration script to Kestra internal
          storage.
        extension: ".sql"
        content: "{{ outputs.generate_migration_and_dbt_models.jsonOutput.migration_sql
          }}"

      - id: persist_dbt_schema
        type: io.kestra.plugin.core.storage.Write
        description: Writes updated dbt model schema YAML definition to Kestra internal
          storage.
        extension: ".yml"
        content: "{{
          outputs.generate_migration_and_dbt_models.jsonOutput.dbt_schema_yaml
          }}"

      - id: persist_migration_summary
        type: io.kestra.plugin.core.storage.Write
        description: Writes executive schema evolution summary markdown to Kestra
          internal storage.
        extension: ".md"
        content: |
          # Data Contract Schema Migration Summary: {{ inputs.table_name }}

          - **Target Table:** `{{ inputs.schema_name }}.{{ inputs.table_name }}`
          - **Drift Detected:** {{ outputs.compute_contract_diff.vars.drift_detected }}
          - **Added Columns Count:** {{ outputs.compute_contract_diff.vars.added_count }}
          - **Modified Columns Count:** {{ outputs.compute_contract_diff.vars.modified_count }}
          - **Risk Classification:** `{{ outputs.generate_migration_and_dbt_models.jsonOutput.risk_level }}`
          - **Compiler Dry-Run:** Passed (`BEGIN ... ROLLBACK` verified against PostgreSQL)

          ## Backward Compatibility Assessment
          {{ outputs.generate_migration_and_dbt_models.jsonOutput.backward_compatibility_notes }}

          ## Generated PostgreSQL Migration SQL
          ```sql
          {{ outputs.generate_migration_and_dbt_models.jsonOutput.migration_sql }}
          ```

          ## Generated dbt Model Specification
          ```yaml
          {{ outputs.generate_migration_and_dbt_models.jsonOutput.dbt_schema_yaml }}
          ```

      - id: notify_slack_summary
        type: io.kestra.plugin.slack.notifications.SlackIncomingWebhook
        description: Dispatches schema drift notification card with SQL preview and
          artifact download links to Slack.
        url: "{{ secret('SLACK_WEBHOOK_URL') }}"
        messageText: |
          :gear: *Data Contract Drift Resolved for {{ inputs.schema_name }}.{{ inputs.table_name }}*
          *Added Columns:* `{{ outputs.compute_contract_diff.vars.added_count }}` | *Risk Level:* `{{ outputs.generate_migration_and_dbt_models.jsonOutput.risk_level }}`
          *Syntax Validation:* Verified in transactional dry-run on PostgreSQL.
          *Generated SQL Migration:* `{{ outputs.persist_migration_sql.uri }}`
          *Updated dbt Model:* `{{ outputs.persist_dbt_schema.uri }}`
          ```sql
          {{ outputs.generate_migration_and_dbt_models.jsonOutput.migration_sql }}
          ```
    else:
      - id: log_contract_in_sync
        type: io.kestra.plugin.core.log.Log
        description: Records confirmation when live database schema matches data contract.
        message: "Data contract for {{ inputs.schema_name }}.{{ inputs.table_name }} is
          in full alignment with PostgreSQL catalog. Zero schema drift
          detected."

triggers:
  - id: contract_webhook_trigger
    type: io.kestra.plugin.core.trigger.Webhook
    description: Webhook triggered by Git repository push on contract schema file changes.
    key: "{{ secret('CONTRACT_WEBHOOK_KEY') }}"

  - id: daily_contract_reconciliation
    type: io.kestra.plugin.core.trigger.Schedule
    description: Scheduled daily contract reconciliation check. Shipped disabled by default.
    cron: "0 4 * * 1-5"
    disabled: true

errors:
  - id: alert_on_failure
    type: io.kestra.plugin.slack.notifications.SlackIncomingWebhook
    description: Alerts data platform team if contract inspection or syntax
      validation fails.
    url: "{{ secret('SLACK_WEBHOOK_URL') }}"
    messageText: |
      :warning: *Data Contract Migration Generator FAILED for {{ inputs.table_name }}*
      Execution: `{{ execution.id }}`
      Check target PostgreSQL connection and AI provider credentials.