Fix partitioned tables

Author: abhi-airspace-intelligence

etl/src/replication/client.rs

@@ -659,7 +659,8 @@ impl PgReplicationClient {
                         when 0 then true
                         else (a.attnum in (select * from pub_attrs))
                         end
-                    )",
+                    )"
+                    .to_string(),
                 )
             } else {
                 // Postgres 14 or earlier or unknown, fallback to no column-level filtering
@@ -674,20 +675,90 @@ impl PgReplicationClient {
                         )",
                         publication = quote_literal(publication),
                     ),
-                    "and (select count(*) from pub_table) > 0",
+                    format!(
+                        "and ((select count(*) from pub_table) > 0 or exists(
+                            -- Also allow if parent table is in publication (for partitioned tables)
+                            select 1 from pg_inherits i
+                            join pg_publication_rel r on r.prrelid = i.inhparent
+                            join pg_publication p on p.oid = r.prpubid
+                            where i.inhrelid = {table_id} and p.pubname = {publication}
+                        ))",
+                        publication = quote_literal(publication),
+                    ),
                 )
             }
         } else {
-            ("".into(), "")
+            ("".to_string(), "".to_string())
+        };
+
+        let has_pub_cte = !pub_cte.is_empty();
+
+        let cte_prefix = if has_pub_cte {
+            // If there's already a pub_cte WITH clause, add our CTEs to it with a comma
+            format!("{pub_cte},")
+        } else {
+            // If no pub_cte, start our own WITH clause (no need for RECURSIVE)
+            "with ".to_string()
         };
 
         let column_info_query = format!(
-            "{pub_cte}
-            select a.attname,
+            "{cte_prefix}
+            -- Find direct parent of current table (if it's a partition)
+            direct_parent as (
+                select i.inhparent as parent_oid
+                from pg_inherits i 
+                where i.inhrelid = {table_id}::oid
+                limit 1
+            ),
+            -- Get parent table's primary key columns
+            parent_pk_cols as (
+                select array_agg(a.attname order by x.n) as pk_column_names
+                from pg_constraint con
+                join unnest(con.conkey) with ordinality as x(attnum, n) on true
+                join pg_attribute a on a.attrelid = con.conrelid and a.attnum = x.attnum
+                join direct_parent dp on con.conrelid = dp.parent_oid
+                where con.contype = 'p'
+                group by con.conname
+            ),
+            -- Check if current table has a unique index on the parent PK columns
+            partition_has_pk_index as (
+                select case 
+                    when exists (select 1 from direct_parent) 
+                         and exists (select 1 from parent_pk_cols)
+                         and exists (
+                            -- Check if there's a unique, valid index on the parent PK columns
+                            select 1
+                            from pg_index ix
+                            cross join parent_pk_cols pk
+                            where ix.indrelid = {table_id}::oid
+                            and ix.indisunique = true
+                            and ix.indisvalid = true
+                            and array(
+                                select a.attname
+                                from unnest(ix.indkey) with ordinality k(attnum, ord)
+                                join pg_attribute a on a.attrelid = ix.indrelid and a.attnum = k.attnum
+                                where ord <= ix.indnkeyatts  -- exclude INCLUDE columns
+                                order by ord
+                            ) = pk.pk_column_names
+                        ) then true
+                    else false
+                end as has_inherited_pk
+            )
+            SELECT a.attname,
                 a.atttypid,
                 a.atttypmod,
                 a.attnotnull,
-                coalesce(i.indisprimary, false) as primary
+                case 
+                    -- First check for direct primary key
+                    when coalesce(i.indisprimary, false) = true then true
+                    -- Then check for inherited primary key from partitioned table parent
+                    when (select has_inherited_pk from partition_has_pk_index) = true 
+                         and exists (
+                            select 1 from parent_pk_cols pk 
+                            where a.attname = any(pk.pk_column_names)
+                         ) then true
+                    else false
+                end as primary
             from pg_attribute a
             left join pg_index i
                 on a.attrelid = i.indrelid

etl/src/test_utils/test_schema.rs

@@ -127,6 +127,77 @@ pub async fn setup_test_database_schema<G: GenericClient>(
     }
 }
 
+/// Creates a partitioned table with the given name and partitions.
+///
+/// This function creates:
+/// 1. A parent partitioned table with a primary key
+/// 2. Several child partitions based on the provided partition specifications
+///
+/// Returns the table ID of the parent table and a list of partition table IDs.
+pub async fn create_partitioned_table<G: GenericClient>(
+    database: &PgDatabase<G>,
+    table_name: TableName,
+    partition_specs: &[(&str, &str)], // (partition_name, partition_constraint)
+) -> Result<(TableId, Vec<TableId>), tokio_postgres::Error> {
+    let create_parent_query = format!(
+        "CREATE TABLE {} (
+            id bigserial,
+            data text NOT NULL,
+            partition_key integer NOT NULL,
+            PRIMARY KEY (id, partition_key)
+        ) PARTITION BY RANGE (partition_key)",
+        table_name.as_quoted_identifier()
+    );
+
+    database.run_sql(&create_parent_query).await?;
+
+    let parent_row = database
+        .client
+        .as_ref()
+        .unwrap()
+        .query_one(
+            "SELECT c.oid FROM pg_class c JOIN pg_namespace n ON n.oid = c.relnamespace 
+             WHERE n.nspname = $1 AND c.relname = $2",
+            &[&table_name.schema, &table_name.name],
+        )
+        .await?;
+
+    let parent_table_id: TableId = parent_row.get(0);
+    let mut partition_table_ids = Vec::new();
+
+    for (partition_name, partition_constraint) in partition_specs {
+        let partition_table_name = TableName::new(
+            table_name.schema.clone(),
+            format!("{}_{}", table_name.name, partition_name),
+        );
+
+        let create_partition_query = format!(
+            "CREATE TABLE {} PARTITION OF {} FOR VALUES {}",
+            partition_table_name.as_quoted_identifier(),
+            table_name.as_quoted_identifier(),
+            partition_constraint
+        );
+
+        database.run_sql(&create_partition_query).await?;
+
+        let partition_row = database
+            .client
+            .as_ref()
+            .unwrap()
+            .query_one(
+                "SELECT c.oid FROM pg_class c JOIN pg_namespace n ON n.oid = c.relnamespace 
+                 WHERE n.nspname = $1 AND c.relname = $2",
+                &[&partition_table_name.schema, &partition_table_name.name],
+            )
+            .await?;
+
+        let partition_table_id: TableId = partition_row.get(0);
+        partition_table_ids.push(partition_table_id);
+    }
+
+    Ok((parent_table_id, partition_table_ids))
+}
+
 /// Inserts users data into the database for testing purposes.
 pub async fn insert_users_data<G: GenericClient>(
     client: &mut PgDatabase<G>,