Restore PostgreSQL foreign-key triggers after interrupted loads

# SPDX-License-Identifier: Apache-2.0
id: postgres-bulk-load-disabled-constraint-trigger-restore-guard
namespace: company.team
description: Restore disabled integer foreign-key triggers without changing rows
  loaded while enforcement was off.

concurrency:
  limit: 1
  behavior: QUEUE

inputs:
  - id: postgres_url
    type: STRING
    defaults: jdbc:postgresql://localhost:5432/app
  - id: postgres_username
    type: STRING
    defaults: "{{ secret('POSTGRES_USERNAME') }}"
  - id: postgres_password
    type: SECRET
    defaults: "{{ secret('POSTGRES_PASSWORD') }}"
  - id: schema_name
    type: STRING
    defaults: public
    description: One exact user schema. Both tables must be in this schema.
  - id: protected_constraints
    type: ARRAY
    itemType: STRING
    defaults: []
    description: Constraint names whose disabled triggers are intentional.
  - id: max_pair_mb
    type: INT
    defaults: 16
    min: 1
    max: 128
    description: Combined table and index size ceiling before any row-count audit.
  - id: max_repairs
    type: INT
    defaults: 2
    min: 1
    max: 5
  - id: dry_run
    type: BOOL
    defaults: true
  - id: auto_remediate
    type: BOOL
    defaults: false
  - id: slack_webhook_url
    type: SECRET
    defaults: "{{ secret('SLACK_WEBHOOK_URL') }}"

variables:
  jdbc_url: "{{ inputs.postgres_url }}{{ inputs.postgres_url contains '?' ? '&' :
    '?'
    }}connectTimeout=5&socketTimeout=30&options=-c%20statement_timeout%3D20000%\
    20-c%20lock_timeout%3D2000"
  audit_sql: |
    WITH candidates AS MATERIALIZED (
      SELECT k.oid AS constraint_oid, k.conrelid, k.confrelid, k.conname,
        n.nspname, c.relname, p.relname AS parent_name,
        pg_get_constraintdef(k.oid) AS definition,
        jsonb_agg(jsonb_build_object('oid', CAST(t.oid AS bigint), 'name', t.tgname, 'state', CAST(t.tgenabled AS text)) ORDER BY t.oid) AS triggers,
        bool_or(t.tgenabled = 'D') AS disabled,
        c.relkind = 'r' AND p.relkind = 'r' AND NOT c.relrowsecurity AND NOT p.relrowsecurity
          AND NOT EXISTS (SELECT 1 FROM pg_inherits i WHERE i.inhrelid IN (c.oid,p.oid) OR i.inhparent IN (c.oid,p.oid))
          AND cardinality(k.conkey) = 1 AND cardinality(k.confkey) = 1
          AND NOT k.condeferrable AND k.confmatchtype = 's'
          AND a.atttypid IN (20,21,23) AND b.atttypid IN (20,21,23)
          AND op.oprnamespace = CAST('pg_catalog' AS regnamespace)
          AND bool_and(t.tgenabled IN ('O','D')) AS supported,
        a.attname AS child_column, b.attname AS parent_column,
        pg_total_relation_size(c.oid) + pg_total_relation_size(p.oid) AS pair_bytes
      FROM pg_constraint k
      JOIN pg_class c ON c.oid = k.conrelid JOIN pg_namespace n ON n.oid = c.relnamespace
      JOIN pg_class p ON p.oid = k.confrelid
      JOIN pg_trigger t ON t.tgconstraint = k.oid AND t.tgisinternal
      JOIN pg_attribute a ON a.attrelid = c.oid AND a.attnum = k.conkey[1]
      JOIN pg_attribute b ON b.attrelid = p.oid AND b.attnum = k.confkey[1]
      JOIN pg_operator op ON op.oid = k.conpfeqop[1]
      WHERE k.contype = 'f' AND n.nspname = CAST(:schema AS text) AND p.relnamespace = c.relnamespace
        AND NOT EXISTS (SELECT 1 FROM jsonb_array_elements_text(CAST(:protected AS jsonb)) x WHERE x = k.conname)
      GROUP BY k.oid,n.nspname,c.oid,p.oid,a.attname,b.attname,a.atttypid,b.atttypid,op.oprnamespace
      HAVING bool_or(t.tgenabled = 'D')
      ORDER BY k.oid LIMIT 101
    ), counted AS MATERIALIZED (
      SELECT *, CASE WHEN supported AND pair_bytes <= CAST(:cap AS bigint)*1048576 THEN
        format('SELECT count(*) AS orphans FROM ONLY %I.%I c WHERE c.%I IS NOT NULL AND NOT EXISTS (SELECT 1 FROM ONLY %I.%I p WHERE p.%I = c.%I)',
          nspname,relname,child_column,nspname,parent_name,parent_column,child_column)
      END AS orphan_sql FROM candidates
    )
    SELECT CAST(constraint_oid AS bigint), CAST(conrelid AS bigint), CAST(confrelid AS bigint), conname, nspname, relname, parent_name,
      definition, CAST(triggers AS text) AS triggers_json, supported, pair_bytes, orphan_sql,
      CASE WHEN orphan_sql IS NOT NULL THEN
        CAST(CAST((xpath('/row/orphans/text()', query_to_xml(orphan_sql,false,true,'')))[1] AS text) AS bigint)
      END AS orphans,
      orphan_sql IS NOT NULL AS eligible
    FROM counted ORDER BY constraint_oid

tasks:
  - id: validate_scope
    type: io.kestra.plugin.core.execution.Assert
    conditions:
      - "{{ inputs.schema_name != '' and not (inputs.schema_name contains 'pg_')
        and inputs.schema_name != 'information_schema' }}"
    errorMessage: Choose an exact user schema.
  - id: audit
    type: io.kestra.plugin.jdbc.postgresql.Query
    connectionPooling: false
    url: "{{ render(vars.jdbc_url) }}"
    username: "{{ inputs.postgres_username }}"
    password: "{{ inputs.postgres_password }}"
    fetchType: FETCH
    parameters:
      schema: "{{ inputs.schema_name }}"
      protected: "{{ inputs.protected_constraints | toJson }}"
      cap: "{{ inputs.max_pair_mb }}"
    sql: "{{ render(vars.audit_sql) }}"
  - id: bound_scan
    type: io.kestra.plugin.core.execution.Assert
    conditions:
      - "{{ outputs.audit.size <= 100 }}"
    errorMessage: More than 100 disabled constraints. Narrow the schema scope.
  - id: plan
    type: io.kestra.plugin.core.output.OutputValues
    values:
      targets: >-
        {{ {'rows': outputs.audit.rows, 'limit': inputs.max_repairs} | toJson |
        jq('.limit as $n | [.rows[] | select(.eligible)][:($n|tonumber)]') |
        first | toJson }}
  - id: report
    type: io.kestra.plugin.core.log.Log
    message: "Disabled foreign keys: {{ outputs.audit.rows | toJson }}. Approved
      batch candidate: {{ outputs.plan.values.targets }}. Existing orphan rows
      are never repaired here."
  - id: require_counted_plan
    type: io.kestra.plugin.core.execution.Assert
    conditions:
      - "{{ (outputs.plan.values.targets | jq('[.[] | select(.orphans == null or
        .orphans < 0)] | length') | first) == 0 }}"
    errorMessage: An eligible constraint has no reliable orphan count. Refuse the plan.
  - id: findings
    type: io.kestra.plugin.core.flow.If
    condition: "{{ outputs.audit.size > 0 }}"
    then:
      - id: notify_findings
        type: io.kestra.plugin.slack.notifications.SlackIncomingWebhook
        url: "{{ inputs.slack_webhook_url }}"
        payload: "{{ {'text': 'Disabled PostgreSQL foreign-key triggers detected. Review
          execution ' ~ execution.id ~ '. Existing row violations are counted
          but never changed.'} | toJson }}"
      - id: live_mode
        type: io.kestra.plugin.core.flow.If
        condition: "{{ inputs.auto_remediate and not inputs.dry_run and
          (outputs.plan.values.targets | jq('length') | first) > 0 }}"
        then:
          - id: approval
            type: io.kestra.plugin.core.flow.Pause
            pauseDuration: PT1H
            behavior: FAIL
            description: Confirm these disables are accidental. Restoring internal triggers
              requires superuser and briefly locks both tables. Existing invalid
              rows remain.
            onResume:
              - id: approved
                type: BOOL
                required: false
          - id: approved_only
            type: io.kestra.plugin.core.flow.If
            condition: "{{ outputs.approval.onResume.approved ?? false }}"
            then:
              - id: remediate_constraints
                type: io.kestra.plugin.core.flow.Loop
                concurrencyLimit: 1
                values: "{{ outputs.plan.values.targets }}"
                tasks:
                  - id: restore_one
                    type: io.kestra.plugin.jdbc.postgresql.Query
                    connectionPooling: false
                    url: "{{ render(vars.jdbc_url) }}"
                    username: "{{ inputs.postgres_username }}"
                    password: "{{ inputs.postgres_password }}"
                    fetchType: NONE
                    sql: |
                      DO $restore$
                      DECLARE plan jsonb := convert_from(decode('{{ item.value | base64encode }}','base64'),'UTF8')::jsonb;
                        k record; r record; current_triggers jsonb; before_count bigint; after_count bigint;
                      BEGIN
                        FOR r IN SELECT c.oid,n.nspname,c.relname FROM pg_class c JOIN pg_namespace n ON n.oid=c.relnamespace
                          WHERE c.oid IN ((plan->>'conrelid')::oid,(plan->>'confrelid')::oid) ORDER BY c.oid
                        LOOP EXECUTE format('LOCK TABLE ONLY %I.%I IN SHARE ROW EXCLUSIVE MODE',r.nspname,r.relname); END LOOP;
                        SELECT x.*,c.relname,n.nspname,p.relname AS parent_name INTO k FROM pg_constraint x
                          JOIN pg_class c ON c.oid=x.conrelid JOIN pg_namespace n ON n.oid=c.relnamespace
                          JOIN pg_class p ON p.oid=x.confrelid
                          WHERE x.oid=(plan->>'constraint_oid')::oid AND x.contype='f'
                            AND x.conrelid=(plan->>'conrelid')::oid AND x.confrelid=(plan->>'confrelid')::oid
                            AND n.nspname=plan->>'nspname' AND c.relname=plan->>'relname' AND p.relname=plan->>'parent_name'
                            AND c.relkind='r' AND p.relkind='r' AND NOT c.relrowsecurity AND NOT p.relrowsecurity
                            AND NOT x.condeferrable AND x.confmatchtype='s' AND cardinality(x.conkey)=1 AND cardinality(x.confkey)=1
                            AND EXISTS (SELECT 1 FROM pg_attribute a WHERE a.attrelid=c.oid AND a.attnum=x.conkey[1] AND a.atttypid IN (20,21,23))
                            AND EXISTS (SELECT 1 FROM pg_attribute a WHERE a.attrelid=p.oid AND a.attnum=x.confkey[1] AND a.atttypid IN (20,21,23))
                            AND EXISTS (SELECT 1 FROM pg_operator op WHERE op.oid=x.conpfeqop[1] AND op.oprnamespace='pg_catalog'::regnamespace)
                            AND NOT EXISTS (SELECT 1 FROM pg_inherits i WHERE i.inhrelid IN (c.oid,p.oid) OR i.inhparent IN (c.oid,p.oid))
                            AND pg_total_relation_size(c.oid)+pg_total_relation_size(p.oid) <= {{ inputs.max_pair_mb }}*1048576;
                        IF NOT FOUND OR pg_get_constraintdef(k.oid) IS DISTINCT FROM plan->>'definition'
                          OR k.conname IS DISTINCT FROM plan->>'conname' THEN RAISE EXCEPTION 'Approved constraint identity or shape changed'; END IF;
                        SELECT jsonb_agg(jsonb_build_object('oid',t.oid::bigint,'name',t.tgname,'state',t.tgenabled::text) ORDER BY t.oid)
                          INTO current_triggers FROM pg_trigger t WHERE t.tgconstraint=k.oid AND t.tgisinternal;
                        IF current_triggers IS DISTINCT FROM (plan->>'triggers_json')::jsonb THEN RAISE EXCEPTION 'Approved trigger snapshot changed'; END IF;
                        EXECUTE plan->>'orphan_sql' INTO before_count;
                        IF before_count IS DISTINCT FROM (plan->>'orphans')::bigint THEN RAISE EXCEPTION 'Approved orphan count changed'; END IF;
                        FOR r IN SELECT t.tgname,c.relname,n.nspname FROM pg_trigger t JOIN pg_class c ON c.oid=t.tgrelid
                          JOIN pg_namespace n ON n.oid=c.relnamespace WHERE t.tgconstraint=k.oid AND t.tgisinternal AND t.tgenabled='D' ORDER BY t.oid
                        LOOP EXECUTE format('ALTER TABLE ONLY %I.%I ENABLE TRIGGER %I',r.nspname,r.relname,r.tgname); END LOOP;
                        IF EXISTS (SELECT 1 FROM pg_trigger WHERE tgconstraint=k.oid AND tgisinternal AND tgenabled='D') THEN
                          RAISE EXCEPTION 'A disabled trigger survived restoration'; END IF;
                        EXECUTE plan->>'orphan_sql' INTO after_count;
                        IF after_count IS DISTINCT FROM before_count THEN RAISE EXCEPTION 'Restoration changed the orphan count'; END IF;
                      END $restore$;
              - id: verify
                type: io.kestra.plugin.jdbc.postgresql.Query
                connectionPooling: false
                url: "{{ render(vars.jdbc_url) }}"
                username: "{{ inputs.postgres_username }}"
                password: "{{ inputs.postgres_password }}"
                fetchType: FETCH
                parameters:
                  targets: "{{ outputs.plan.values.targets }}"
                sql: |
                  SELECT p->>'conname' AS constraint_name,
                    EXISTS (SELECT 1 FROM pg_constraint k WHERE k.oid=CAST(p->>'constraint_oid' AS oid)
                      AND k.conrelid=CAST(p->>'conrelid' AS oid) AND k.confrelid=CAST(p->>'confrelid' AS oid)
                      AND pg_get_constraintdef(k.oid)=p->>'definition')
                    AND NOT EXISTS (SELECT 1 FROM pg_trigger t WHERE t.tgconstraint=CAST(p->>'constraint_oid' AS oid) AND t.tgisinternal AND t.tgenabled='D') AS restored
                  FROM jsonb_array_elements(CAST(:targets AS jsonb)) p
              - id: require_restored
                type: io.kestra.plugin.core.execution.Assert
                conditions:
                  - "{{ (outputs.verify.rows | toJson | jq('[.[] |
                    select(.restored != true)] | length') | first) == 0 }}"
                errorMessage: An approved foreign key is missing or still has disabled internal
                  triggers.
            else:
              - id: declined
                type: io.kestra.plugin.core.log.Log
                message: Approval declined. PostgreSQL was not changed.

errors:
  - id: notify_failure
    type: io.kestra.plugin.slack.notifications.SlackIncomingWebhook
    url: "{{ inputs.slack_webhook_url }}"
    payload: "{{ {'text': 'PostgreSQL trigger restoration failed in execution ' ~
      execution.id ~ '. Re-audit before retrying. Completed constraints stay
      enabled.'} | toJson }}"

triggers:
  - id: daily_check
    type: io.kestra.plugin.core.trigger.Schedule
    cron: "0 4 * * *"
    disabled: true
    inputs:
      dry_run: true
      auto_remediate: false