Detect distribution column joins when RLS policies split equivalence …

[codeforall]

Overview

This PR fixes a bug where Citus incorrectly rejects valid colocated joins when Row Level Security (RLS) policies use volatile functions like current_setting(). The issue occurs because PostgreSQL's planner splits equivalence classes in the presence of volatile RLS expressions, preventing Citus from recognizing that tables are joined on their distribution columns.

Problem

When an RLS policy uses a volatile function, PostgreSQL creates separate equivalence classes for each table instead of a unified equivalence class:

-- Query
SELECT * FROM table_a 
JOIN table_b ON table_a.tenant_id = table_b.tenant_id;

-- With RLS policy: USING (tenant_id = current_setting('app.tenant_id')::uuid)
-- PostgreSQL creates:
EC1: [table_a.tenant_id, current_setting('app.tenant_id')::uuid]
EC2: [table_b.tenant_id, current_setting('app.tenant_id')::uuid]

-- Instead of:
EC: [table_a.tenant_id, table_b.tenant_id, current_setting('app.tenant_id')::uuid]

This split prevents Citus's GenerateAttributeEquivalencesForRelationRestrictions() from detecting that the tables are joined on tenant_id.

Solution

Core Algorithm: Implements a two-pass approach in GenerateAttributeEquivalencesForRelationRestrictions():

First Pass: Standard EC processing + RLS pattern detection (EC with both Var and non-Var members) Second Pass (conditional):

• Groups ECs by their shared volatile expressions using hash-based comparison
• Merges Var members from ECs sharing identical expressions
• Creates new AttributeEquivalenceClass objects representing the implicit joins

PR includes also includes a test case to verify the fix

[codecov[bot]]

Codecov Report

:x: Patch coverage is 0% with 93 lines in your changes missing coverage. Please review. :white_check_mark: Project coverage is 29.42%. Comparing base (662b724) to head (467e74a). :warning: Report is 2 commits behind head on main.

:x: Your patch check has failed because the patch coverage (0.00%) is below the target coverage (75.00%). You can increase the patch coverage or adjust the target coverage. :x: Your project check has failed because the head coverage (29.42%) is below the target coverage (87.50%). You can increase the head coverage or adjust the target coverage.

:exclamation: There is a different number of reports uploaded between BASE (662b724) and HEAD (467e74a). Click for more details.

HEAD has 80 uploads less than BASE

Flag	BASE (662b724)	HEAD (467e74a)
15_citus_upgrade	1	0
15_regress_check-columnar-isolation	1	0
16_regress_check-follower-cluster	1	0
17_regress_check-columnar-isolation	1	0
15_16_upgrade	1	0
16_regress_check-add-backup-node	1	0
15_17_upgrade	1	0
16_regress_check-enterprise-isolation-logicalrep-2	1	0
16_regress_check-query-generator	1	0
16_regress_check-enterprise-isolation-logicalrep-3	1	0
16_regress_check-enterprise-failure	1	0
15_regress_check-add-backup-node	1	0
15_regress_check-split	1	0
16_regress_check-split	1	0
16_regress_check-columnar	1	0
15_regress_check-enterprise-isolation-logicalrep-2	1	0
17_regress_check-enterprise-isolation	1	0
15_regress_check-enterprise-isolation	1	0
16_regress_check-multi-mx	1	0
17_regress_check-multi-mx	1	0
16_regress_check-enterprise-isolation	1	0
16_regress_check-enterprise-isolation-logicalrep-1	1	0
17_regress_check-failure	1	0
16_regress_check-failure	1	0
17_cdc_installcheck	1	0
16_cdc_installcheck	1	0
16_regress_check-operations	1	0
17_regress_check-operations	1	0
17_regress_check-isolation	1	0
15_regress_check-isolation	1	0
16_regress_check-isolation	1	0
15_regress_check-enterprise-isolation-logicalrep-1	1	0
15_regress_check-query-generator	1	0
15_regress_check-vanilla	1	0
15_regress_check-enterprise	1	0
16_regress_check-vanilla	1	0
16_regress_check-enterprise	1	0
17_regress_check-vanilla	1	0
16_arbitrary_configs_3	1	0
15_arbitrary_configs_3	1	0
17_arbitrary_configs_3	1	0
16_arbitrary_configs_5	1	0
15_regress_check-operations	1	0
16_regress_check-multi-1	1	0
17_regress_check-multi-1	1	0
15_arbitrary_configs_5	1	0
15_regress_check-multi-1	1	0
15_arbitrary_configs_2	1	0
16_arbitrary_configs_2	1	0
17_arbitrary_configs_2	1	0
17_arbitrary_configs_4	1	0
15_regress_check-enterprise-failure	1	0
15_regress_check-columnar	1	0
17_regress_check-enterprise-failure	1	0
17_regress_check-split	1	0
17_regress_check-columnar	1	0
16_arbitrary_configs_4	1	0
15_arbitrary_configs_0	1	0
17_arbitrary_configs_5	1	0
17_regress_check-enterprise-isolation-logicalrep-3	1	0
17_regress_check-enterprise-isolation-logicalrep-2	1	0
15_regress_check-enterprise-isolation-logicalrep-3	1	0
15_arbitrary_configs_1	1	0
16_arbitrary_configs_1	1	0
17_arbitrary_configs_1	1	0
15_cdc_installcheck	1	0
17_regress_check-query-generator	1	0
17_regress_check-multi	1	0
17_regress_check-enterprise	1	0
15_arbitrary_configs_4	1	0
15_regress_check-failure	1	0
17_regress_check-enterprise-isolation-logicalrep-1	1	0
15_regress_check-multi-mx	1	0
15_regress_check-follower-cluster	1	0
16_regress_check-columnar-isolation	1	0
15_regress_check-multi	1	0
16_regress_check-multi	1	0
17_regress_check-add-backup-node	1	0
16_17_upgrade	1	0
17_regress_check-follower-cluster	1	0

Additional details and impacted files

@@             Coverage Diff             @@
##             main    #8357       +/-   ##
===========================================
- Coverage   88.95%   29.42%   -59.53%     
===========================================
  Files         287      287               
  Lines       63151    62906      -245     
  Branches     7942     7868       -74     
===========================================
- Hits        56176    18512    -37664     
- Misses       4667    42233    +37566     
+ Partials     2308     2161      -147     

:rocket: New features to boost your workflow:

• :snowflake: Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[colm-mchugh]

+1 nice fix, had some comments/questions on some of the details. And more tests please.