Citus Colocation Violation Guard

id: citus-colocation-violation-guard
namespace: company.team

description: |
  Two Citus distributed tables join locally on each worker only while they sit in the same
  colocation group. Create one with colocate_with none or load two tables through different
  paths and they land in different groups. Nothing errors at create time. Every join between
  them then repartitions across workers and a foreign key between them is rejected outright.
  This flow compares a declared set of tables that should colocate against colocationid in
  pg_dist_partition, reports the pairs that drifted apart, then calls
  update_distributed_table_colocation behind a dry run and an approval, re-reads colocationid
  and fails the run if the pair is still split.

inputs:
  - id: jdbc_url
    type: STRING
    displayName: Coordinator JDBC URL
    description: The Citus coordinator. All colocation metadata lives there.
    defaults: "jdbc:postgresql://localhost:5432/citus"

  - id: expected_groups
    type: JSON
    displayName: Tables that should colocate
    description: >-
      A list of groups. Each group is a list of schema qualified table names
      whose first entry is the anchor. Every later entry in a group should share
      the anchor's colocation group. Tables you never join together belong in
      different groups, so leave them out.
    defaults:
      [
        [ "public.orders", "public.order_items" ],
        [ "public.customers", "public.customer_events" ]
      ]

  - id: auto_remediate
    type: BOOL
    displayName: Auto remediate
    description: Allow the flow to change a colocation group. Off means report only.
    defaults: false

  - id: dry_run
    type: BOOL
    displayName: Dry run
    description: Must be false, together with auto_remediate, before any colocation
      group is changed.
    defaults: true

variables:
  audit_cte: |
    WITH decl AS (
      SELECT g.ord::int AS group_no,
             (g.grp -> 0) #>> '{}' AS anchor_table,
             m.tbl AS member_table
        FROM json_array_elements('{{ inputs.expected_groups }}'::json) WITH ORDINALITY AS g(grp, ord)
        CROSS JOIN LATERAL (
          SELECT e.value #>> '{}' AS tbl, (e.ordinality - 1)::int AS pos
            FROM json_array_elements(g.grp) WITH ORDINALITY AS e(value, ordinality)
        ) m
       WHERE m.pos > 0
    ), part AS (
      SELECT p.logicalrelid AS reloid,
             p.colocationid::int AS colocationid,
             column_to_column_name(p.logicalrelid, p.partkey::text) AS dist_column,
             (SELECT count(*) FROM pg_dist_shard s WHERE s.logicalrelid = p.logicalrelid)::int AS shard_count
        FROM pg_dist_partition p
       WHERE p.partmethod = 'h'
    ), typed AS (
      SELECT t.reloid, t.colocationid, t.dist_column, t.shard_count,
             format_type(a.atttypid, a.atttypmod) AS dist_type
        FROM part t
        JOIN pg_attribute a ON a.attrelid = t.reloid AND a.attname = t.dist_column
                           AND a.attnum > 0 AND NOT a.attisdropped
    ), pair AS (
      SELECT d.group_no, d.anchor_table, d.member_table,
             ta.colocationid AS anchor_colocationid, tm.colocationid AS member_colocationid,
             ta.dist_type AS anchor_dist_type, tm.dist_type AS member_dist_type,
             ta.dist_column AS anchor_dist_column, tm.dist_column AS member_dist_column,
             ta.shard_count AS anchor_shard_count, tm.shard_count AS member_shard_count
        FROM decl d
        LEFT JOIN typed ta ON ta.reloid = to_regclass(d.anchor_table)
        LEFT JOIN typed tm ON tm.reloid = to_regclass(d.member_table)
    ), audit AS (
      SELECT group_no, anchor_table, member_table,
             coalesce(anchor_colocationid, 0) AS anchor_colocationid,
             coalesce(member_colocationid, 0) AS member_colocationid,
             coalesce(anchor_dist_type, 'n/a') AS anchor_dist_type,
             coalesce(member_dist_type, 'n/a') AS member_dist_type,
             coalesce(anchor_dist_column, 'n/a') AS anchor_dist_column,
             coalesce(member_dist_column, 'n/a') AS member_dist_column,
             coalesce(anchor_shard_count, 0) AS anchor_shard_count,
             coalesce(member_shard_count, 0) AS member_shard_count,
             CASE
               WHEN anchor_colocationid IS NULL OR member_colocationid IS NULL THEN 'UNKNOWN_TABLE'
               WHEN anchor_colocationid = member_colocationid THEN 'COLOCATED'
               WHEN anchor_dist_type IS DISTINCT FROM member_dist_type THEN 'TYPE_MISMATCH'
               WHEN anchor_shard_count IS DISTINCT FROM member_shard_count THEN 'SHARD_COUNT_MISMATCH'
               ELSE 'VIOLATION'
             END AS verdict,
             CASE
               WHEN anchor_colocationid IS NULL AND member_colocationid IS NULL
                 THEN 'neither table is a hash distributed Citus table on this coordinator'
               WHEN anchor_colocationid IS NULL
                 THEN 'the anchor is not a hash distributed Citus table on this coordinator'
               WHEN member_colocationid IS NULL
                 THEN 'this table is not a hash distributed Citus table on this coordinator'
               WHEN anchor_colocationid = member_colocationid
                 THEN 'already sharing colocation group ' || member_colocationid
               WHEN anchor_dist_type IS DISTINCT FROM member_dist_type
                 THEN 'distribution column types differ (' || member_dist_type || ' against '
                      || anchor_dist_type || '), so Citus refuses to colocate them'
               WHEN anchor_shard_count IS DISTINCT FROM member_shard_count
                 THEN 'shard counts differ (' || member_shard_count || ' against '
                      || anchor_shard_count || '), so Citus refuses to colocate them'
               ELSE 'same ' || member_dist_type || ' distribution column over '
                    || member_shard_count || ' shards, but colocation group '
                    || member_colocationid || ' against ' || anchor_colocationid
             END AS reason
        FROM pair
    )

tasks:
  - id: audit
    type: io.kestra.plugin.jdbc.postgresql.Query
    description: >-
      Read-only. One row per declared pair, with each side's colocationid from
      pg_dist_partition, its distribution column and type plus its shard count.
      A pair is only a violation when the two sides agree on type and shard
      count, because those are what Citus itself requires before it will put two
      tables in one colocation group.
    url: "{{ inputs.jdbc_url }}"
    username: "{{ secret('POSTGRES_USERNAME') }}"
    password: "{{ secret('POSTGRES_PASSWORD') }}"
    fetchType: FETCH
    sql: |
      {{ render(vars.audit_cte) }}
      SELECT * FROM audit ORDER BY group_no, member_table

  - id: assess
    type: io.kestra.plugin.jdbc.postgresql.Query
    description: >-
      Read-only. Classify the whole declaration as VIOLATION, NOT_COLOCATABLE or
      OK, build the target list for the remediation loop and read
      citus.enable_repartition_joins, which decides whether a split pair makes
      joins slow or makes them fail outright.
    url: "{{ inputs.jdbc_url }}"
    username: "{{ secret('POSTGRES_USERNAME') }}"
    password: "{{ secret('POSTGRES_PASSWORD') }}"
    fetchType: FETCH_ONE
    sql: |
      {{ render(vars.audit_cte) }}
      SELECT
        CASE
          WHEN count(*) FILTER (WHERE verdict = 'VIOLATION') > 0 THEN 'VIOLATION'
          WHEN count(*) FILTER (WHERE verdict <> 'COLOCATED') > 0 THEN 'NOT_COLOCATABLE'
          ELSE 'OK'
        END AS state,
        count(*)::int AS pair_count,
        (count(*) FILTER (WHERE verdict = 'VIOLATION'))::int AS violation_count,
        (count(*) FILTER (WHERE verdict = 'COLOCATED'))::int AS colocated_count,
        (count(*) FILTER (WHERE verdict IN ('TYPE_MISMATCH', 'SHARD_COUNT_MISMATCH')))::int AS refused_count,
        (count(*) FILTER (WHERE verdict = 'UNKNOWN_TABLE'))::int AS unknown_count,
        count(DISTINCT colocationid_pair)::int AS group_spread,
        current_setting('citus.enable_repartition_joins') AS repartition_joins,
        coalesce((SELECT json_agg(json_build_object(
                    'member', member_table, 'anchor', anchor_table,
                    'before_colocationid', member_colocationid,
                    'anchor_colocationid', anchor_colocationid) ORDER BY member_table)::text
                    FROM audit WHERE verdict = 'VIOLATION'), '[]') AS targets_json,
        coalesce((SELECT string_agg(member_table || ' should join ' || anchor_table || ' ('
                    || reason || ')', '; ' ORDER BY member_table)
                    FROM audit WHERE verdict = 'VIOLATION'), '') AS violation_list,
        coalesce((SELECT string_agg('SELECT update_distributed_table_colocation('
                    || quote_literal(member_table) || ', colocate_with => '
                    || quote_literal(anchor_table) || ');', ' ' ORDER BY member_table)
                    FROM audit WHERE verdict = 'VIOLATION'), '') AS plan_sql,
        coalesce((SELECT string_agg(member_table || ' with ' || anchor_table || ' (' || reason || ')',
                    '; ' ORDER BY member_table)
                    FROM audit WHERE verdict IN ('TYPE_MISMATCH', 'SHARD_COUNT_MISMATCH', 'UNKNOWN_TABLE')), '') AS refused_list,
        coalesce((SELECT string_agg(member_table || ' with ' || anchor_table || ' in group '
                    || member_colocationid, '; ' ORDER BY member_table)
                    FROM audit WHERE verdict = 'COLOCATED'), '') AS colocated_list
        FROM (SELECT a.*, a.anchor_colocationid || ':' || a.member_colocationid AS colocationid_pair
                FROM audit a) a

  - id: report
    type: io.kestra.plugin.core.log.Log
    message: >-
      Citus colocation on {{ inputs.jdbc_url }} is {{ outputs.assess.row.state
      }}. {{ outputs.assess.row.pair_count }} declared pairs: {{
      outputs.assess.row.colocated_count }} colocated, {{
      outputs.assess.row.violation_count }} split while colocatable, {{
      outputs.assess.row.refused_count }} that Citus would refuse, {{
      outputs.assess.row.unknown_count }} naming a table that is not hash
      distributed here. citus.enable_repartition_joins is {{
      outputs.assess.row.repartition_joins }}, so a split pair {{
      outputs.assess.row.repartition_joins == 'on' ? 'joins through a
      repartition that ships data between workers' : 'cannot be joined at all
      (the planner raises "the query contains a join that requires
      repartitioning")' }}. {% for r in outputs.audit.rows %}[{{ r.member_table
      }} group={{ r.member_colocationid }} vs {{ r.anchor_table }} group={{
      r.anchor_colocationid }} on {{ r.member_dist_column }} {{
      r.member_dist_type }} x{{ r.member_shard_count }} shards: {{ r.verdict }}]
      {% endfor %}

  - id: route
    type: io.kestra.plugin.core.flow.If
    condition: "{{ outputs.assess.row.state != 'OK' }}"
    then:
      - id: alert
        type: io.kestra.plugin.slack.notifications.SlackIncomingWebhook
        url: "{{ secret('SLACK_WEBHOOK_URL') }}"
        payload: |
          {{ {'text': 'Citus colocation on ' ~ inputs.jdbc_url ~ ' is ' ~ outputs.assess.row.state ~ '. ' ~ outputs.assess.row.violation_count ~ ' of ' ~ outputs.assess.row.pair_count ~ ' declared pairs are in different colocation groups while being colocatable, so their joins repartition and a foreign key between them is rejected. ' ~ (outputs.assess.row.violation_list == '' ? '' : 'Split: ' ~ outputs.assess.row.violation_list ~ '. ') ~ (outputs.assess.row.refused_list == '' ? '' : 'Citus would refuse: ' ~ outputs.assess.row.refused_list ~ '. ') ~ (inputs.auto_remediate and not inputs.dry_run ? 'Re-colocation pending approval.' : 'No change applied (auto_remediate=' ~ inputs.auto_remediate ~ ', dry_run=' ~ inputs.dry_run ~ ').')} | toJson }}

      - id: fixable
        type: io.kestra.plugin.core.flow.If
        condition: "{{ outputs.assess.row.violation_count > 0 }}"
        then:
          - id: log_plan
            type: io.kestra.plugin.core.log.Log
            message: >-
              The dry run. These are the exact statements this flow would run,
              one per split table: {{ outputs.assess.row.plan_sql }}
              update_distributed_table_colocation only rewrites
              pg_dist_partition.colocationid and the pg_dist_colocation row
              count. The shards already line up, so it moves no data and copies
              nothing between workers. It still changes how every future join
              between these tables is planned, which is why it waits for an
              approval. {{ outputs.assess.row.refused_list == '' ? '' : 'Left
              alone because Citus would refuse them: ' ~
              outputs.assess.row.refused_list ~ '.' }}

          - id: remediate_gate
            type: io.kestra.plugin.core.flow.If
            condition: "{{ inputs.auto_remediate and not inputs.dry_run }}"
            then:
              - id: approve
                type: io.kestra.plugin.core.flow.Pause
                description: >-
                  Resume the execution to run the statements listed by log_plan.
                  Kill it to cancel. The change is metadata only, so it is quick
                  and reversible, but it re-plans every join and every foreign
                  key between these tables, so it is never applied without this.
                onResume:
                  - id: approved
                    type: BOOL
                    description: Set true to apply the re-colocation. Anything else leaves the
                      groups as they are.
                    required: false

              - id: approval_result
                type: io.kestra.plugin.core.flow.If
                condition: "{{ outputs.approve.onResume.approved ?? false }}"
                then:
                  - id: recolocate
                    type: io.kestra.plugin.core.flow.Loop
                    description: >-
                      Remediation loop. One update_distributed_table_colocation
                      per split table, concurrencyLimit 1 so two statements
                      never rewrite the same pg_dist_colocation row at once.
                      allowFailure on the statement, because a refusal from
                      Citus must not stop the pass: the re-read below is the
                      only thing that decides the outcome.
                    concurrencyLimit: 1
                    values: "{{ outputs.assess.row.targets_json }}"
                    tasks:
                      - id: set_colocation
                        type: io.kestra.plugin.jdbc.postgresql.Query
                        description: >-
                          Move one table into the anchor's colocation group.
                          Citus validates the distribution column type and the
                          shard count itself and raises if they do not match,
                          which is why the audit filters those pairs out first.
                        allowFailure: true
                        url: "{{ inputs.jdbc_url }}"
                        username: "{{ secret('POSTGRES_USERNAME') }}"
                        password: "{{ secret('POSTGRES_PASSWORD') }}"
                        fetchType: NONE
                        sql: >-
                          SELECT update_distributed_table_colocation('{{
                          fromJson(item.value).member }}', colocate_with => '{{
                          fromJson(item.value).anchor }}')

                  - id: recheck
                    type: io.kestra.plugin.jdbc.postgresql.Query
                    description: >-
                      The gate. Re-read colocationid from pg_dist_partition for
                      exactly the tables that were targeted and compare it with
                      the anchor again. The function returns void and raises
                      nothing when it silently achieves nothing, so its return
                      is worth less than this read.
                    url: "{{ inputs.jdbc_url }}"
                    username: "{{ secret('POSTGRES_USERNAME') }}"
                    password: "{{ secret('POSTGRES_PASSWORD') }}"
                    fetchType: FETCH_ONE
                    sql: |
                      {{ render(vars.audit_cte) }}, target AS (
                        SELECT (e.value -> 'member') #>> '{}' AS member_table,
                               ((e.value -> 'before_colocationid') #>> '{}')::int AS before_colocationid
                          FROM json_array_elements('{{ outputs.assess.row.targets_json }}'::json) AS e(value)
                      ), outcome AS (
                        SELECT t.member_table, t.before_colocationid,
                               a.member_colocationid AS after_colocationid,
                               a.anchor_table, a.anchor_colocationid, a.verdict
                          FROM target t LEFT JOIN audit a ON a.member_table = t.member_table
                      )
                      SELECT count(*)::int AS targeted,
                             (count(*) FILTER (WHERE verdict = 'COLOCATED'))::int AS now_colocated,
                             (count(*) FILTER (WHERE verdict IS DISTINCT FROM 'COLOCATED'))::int AS still_split,
                             (count(*) FILTER (WHERE verdict IS DISTINCT FROM 'COLOCATED')) = 0 AS fixed,
                             coalesce(string_agg(member_table || ' group ' || before_colocationid
                               || ' to ' || coalesce(after_colocationid::text, 'gone') || ' (anchor '
                               || coalesce(anchor_table, 'unknown') || ' group '
                               || coalesce(anchor_colocationid::text, 'none') || ')', '; '
                               ORDER BY member_table), '') AS movement,
                             coalesce((SELECT string_agg(member_table || ' still in group '
                               || coalesce(after_colocationid::text, 'none') || ' against anchor group '
                               || coalesce(anchor_colocationid::text, 'none'), '; ' ORDER BY member_table)
                               FROM outcome WHERE verdict IS DISTINCT FROM 'COLOCATED'), '') AS survivors
                        FROM outcome

                  - id: verify
                    type: io.kestra.plugin.core.flow.If
                    condition: "{{ outputs.recheck.row.fixed }}"
                    then:
                      - id: recolocated
                        type: io.kestra.plugin.core.log.Log
                        message: >-
                          {{ outputs.recheck.row.now_colocated }} of {{
                          outputs.recheck.row.targeted }} tables now share their
                          anchor's colocation group, read back from
                          pg_dist_partition after the change: {{
                          outputs.recheck.row.movement }}. Joins between them
                          plan locally per shard again, a foreign key on the
                          distribution column is accepted again and no data
                          moved to get there. {{
                          outputs.assess.row.colocated_list == '' ? '' : 'Pairs
                          that were already correct and were not touched: ' ~
                          outputs.assess.row.colocated_list ~ '.' }}
                    else:
                      - id: remediation_failed
                        type: io.kestra.plugin.core.execution.Fail
                        errorMessage: >-
                          update_distributed_table_colocation ran without
                          raising, but {{ outputs.recheck.row.still_split }} of
                          {{ outputs.recheck.row.targeted }} targeted tables are
                          still outside their anchor's colocation group: {{
                          outputs.recheck.row.survivors }}. Movement seen: {{
                          outputs.recheck.row.movement }}. The call needs the
                          table owner or a superuser, it raises on a shard count
                          or distribution type it cannot match, and a concurrent
                          create_distributed_table can move the anchor out from
                          under it. Read pg_dist_partition and
                          pg_dist_colocation on the coordinator.
                else:
                  - id: declined
                    type: io.kestra.plugin.core.log.Log
                    message: >-
                      The approval was resumed without approved set to true, so
                      nothing was changed. {{ outputs.assess.row.violation_count
                      }} pairs are still split: {{
                      outputs.assess.row.violation_list }}.
            else:
              - id: reported_only
                type: io.kestra.plugin.core.log.Log
                message: >-
                  {{ outputs.assess.row.violation_count }} colocatable pairs
                  reported, none changed (auto_remediate={{
                  inputs.auto_remediate }}, dry_run={{ inputs.dry_run }}). {{
                  outputs.assess.row.violation_list }}.
        else:
          - id: cannot_colocate
            type: io.kestra.plugin.core.log.Log
            message: >-
              Nothing here can be fixed by changing a colocation group. {{
              outputs.assess.row.refused_count }} declared pairs disagree on
              distribution column type or shard count, which Citus checks itself
              and refuses. {{ outputs.assess.row.unknown_count }} declared pairs
              name a table that is not hash distributed on this coordinator: {{
              outputs.assess.row.refused_list }}. Forcing these is not possible
              through update_distributed_table_colocation, so this flow does not
              try. Fixing them means an alter_distributed_table to the same
              shard count or distribution column type, which does rewrite the
              data. The other option is correcting the declaration because the
              two tables were never meant to join.
    else:
      - id: colocated
        type: io.kestra.plugin.core.log.Log
        message: >-
          All {{ outputs.assess.row.pair_count }} declared pairs share their
          anchor's colocation group, so their joins plan locally per shard: {{
          outputs.assess.row.colocated_list }}.

errors:
  - id: alert_flow_failure
    type: io.kestra.plugin.slack.notifications.SlackIncomingWebhook
    url: "{{ secret('SLACK_WEBHOOK_URL') }}"
    payload: |
      {{ {'text': 'Citus colocation violation guard failed on ' ~ inputs.jdbc_url ~ ' (execution ' ~ execution.id ~ '). Check the logs.'} | toJson }}

triggers:
  - id: daily
    type: io.kestra.plugin.core.trigger.Schedule
    description: Shipped disabled and always in safe mode.
    cron: "0 3 * * *"
    disabled: true
    inputs:
      auto_remediate: false
      dry_run: true