Use recursive query to get oids and write test

Signed-off-by: Abhi Agarwal <[email protected]>

Author: abhi-airspace-intelligence

etl/src/replication/client.rs

@@ -410,47 +410,53 @@ impl PgReplicationClient {
         &self,
         publication_name: &str,
     ) -> EtlResult<Vec<TableId>> {
-        // Prefer pg_publication_rel (explicit tables in the publication, including partition roots)
-        let rel_query = format!(
-            r#"select r.prrelid as oid
-               from pg_publication_rel r
-               join pg_publication p on p.oid = r.prpubid
-              where p.pubname = {}"#,
-            quote_literal(publication_name)
-        );
-
-        let mut table_ids = vec![];
-        let mut has_rows = false;
-        for msg in self.client.simple_query(&rel_query).await? {
-            if let SimpleQueryMessage::Row(row) = msg {
-                has_rows = true;
-                let table_id =
-                    Self::get_row_value::<TableId>(&row, "oid", "pg_publication_rel").await?;
-                table_ids.push(table_id);
-            }
-        }
-
-        if has_rows {
-            return Ok(table_ids);
-        }
-
-        // Fallback to pg_publication_tables (expanded view), used for publications like FOR ALL TABLES
-        let publication_query = format!(
-            "select c.oid from pg_publication_tables pt 
-         join pg_class c on c.relname = pt.tablename 
-         join pg_namespace n on n.oid = c.relnamespace AND n.nspname = pt.schemaname 
-         where pt.pubname = {};",
-            quote_literal(publication_name)
+        let query = format!(
+            r#"
+            with recursive has_rel as (
+                select exists(
+                    select 1
+                    from pg_publication_rel r
+                    join pg_publication p on p.oid = r.prpubid
+                    where p.pubname = {pub}
+                ) as has
+            ),
+            pub_tables as (
+                select r.prrelid as oid
+                from pg_publication_rel r
+                join pg_publication p on p.oid = r.prpubid
+                where p.pubname = {pub} and (select has from has_rel)
+                union all
+                select c.oid
+                from pg_publication_tables pt
+                join pg_class c on c.relname = pt.tablename
+                join pg_namespace n on n.oid = c.relnamespace and n.nspname = pt.schemaname
+                where pt.pubname = {pub} and not (select has from has_rel)
+            ),
+            recurse(relid) as (
+                select oid from pub_tables
+                union all
+                select i.inhparent
+                from pg_inherits i
+                join recurse r on r.relid = i.inhrelid
+            )
+            select distinct relid as oid
+            from recurse r
+            where not exists (
+                select 1 from pg_inherits i where i.inhrelid = r.relid
+            );
+            "#,
+            pub = quote_literal(publication_name)
         );
 
-        for msg in self.client.simple_query(&publication_query).await? {
+        let mut roots = vec![];
+        for msg in self.client.simple_query(&query).await? {
             if let SimpleQueryMessage::Row(row) = msg {
                 let table_id = Self::get_row_value::<TableId>(&row, "oid", "pg_class").await?;
-                table_ids.push(table_id);
+                roots.push(table_id);
             }
         }
 
-        Ok(table_ids)
+        Ok(roots)
     }
 
     /// Starts a logical replication stream from the specified publication and slot.

etl/tests/replication.rs

@@ -1,10 +1,13 @@
 #![cfg(feature = "test-utils")]
 
+use std::collections::HashSet;
+
 use etl::error::ErrorKind;
 use etl::replication::client::PgReplicationClient;
 use etl::test_utils::database::{spawn_source_database, test_table_name};
 use etl::test_utils::pipeline::test_slot_name;
 use etl::test_utils::table::assert_table_schema;
+use etl::test_utils::test_schema::create_partitioned_table;
 use etl_postgres::tokio::test_utils::{TableModification, id_column_schema};
 use etl_postgres::types::ColumnSchema;
 use etl_telemetry::tracing::init_test_tracing;
@@ -405,11 +408,47 @@ async fn test_publication_creation_and_check() {
     );
 
     // We check the table ids of the tables in the publication.
-    let table_ids = parent_client
+    let table_ids: HashSet<_> = parent_client
         .get_publication_table_ids("my_publication")
+        .await
+        .unwrap()
+        .into_iter()
+        .collect();
+    assert_eq!(table_ids, HashSet::from([table_1_id, table_2_id]));
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_publication_table_ids_collapse_partitioned_root() {
+    init_test_tracing();
+    let database = spawn_source_database().await;
+
+    let client = PgReplicationClient::connect(database.config.clone())
         .await
         .unwrap();
-    assert_eq!(table_ids, vec![table_1_id, table_2_id]);
+
+    // We create a partitioned parent with two child partitions.
+    let table_name = test_table_name("part_parent");
+    let (parent_table_id, _children) = create_partitioned_table(
+        &database,
+        table_name.clone(),
+        &[("p1", "from (1) to (100)"), ("p2", "from (100) to (200)")],
+    )
+    .await
+    .unwrap();
+
+    let publication_name = "pub_part_root";
+    database
+        .create_publication(publication_name, std::slice::from_ref(&table_name))
+        .await
+        .unwrap();
+
+    let id = client
+        .get_publication_table_ids(publication_name)
+        .await
+        .unwrap();
+
+    // We expect to get only the parent table id.
+    assert_eq!(id, vec![parent_table_id]);
 }
 
 #[tokio::test(flavor = "multi_thread")]