Fix race condition related to schema initialization

Author: abhi-airspace-intelligence

etl/src/pipeline.rs

@@ -257,6 +257,9 @@ where
     /// This private method ensures that each table in the Postgres publication has
     /// a corresponding replication state record. Tables without existing states are
     /// initialized to the [`TableReplicationPhase::Init`] phase.
+    ///
+    /// It also pre-loads schemas for all discovered tables to ensure they are
+    /// available in the cache before the apply worker starts processing.
     async fn initialize_table_states(
         &self,
         replication_client: &PgReplicationClient,
@@ -285,12 +288,25 @@ where
             .get_publication_table_ids(&self.config.publication_name)
             .await?;
 
+        // Load and store schemas for all discovered tables before starting the apply worker.
+        // This ensures that all table schemas are available in the cache when the apply worker
+        // starts processing the replication stream. Otherwise, you can get a pretty nasty race
+        // condition. However, this is probably too aggressive of a fix since it blocks _all_ apply workers
+        // from kicking off by design. It should perhaps be done on timeout or something similar.
+        let table_schemas = replication_client
+            .get_table_schemas(&table_ids, Some(&self.config.publication_name))
+            .await?;
+
+        for (_, table_schema) in table_schemas {
+            self.store.store_table_schema(table_schema).await?;
+        }
+
         self.store.load_table_replication_states().await?;
         let states = self.store.get_table_replication_states().await?;
-        for table_id in table_ids {
-            if !states.contains_key(&table_id) {
+        for table_id in &table_ids {
+            if !states.contains_key(table_id) {
                 self.store
-                    .update_table_replication_state(table_id, TableReplicationPhase::Init)
+                    .update_table_replication_state(*table_id, TableReplicationPhase::Init)
                     .await?;
             }
         }

etl/src/replication/client.rs

@@ -372,12 +372,19 @@ impl PgReplicationClient {
     }
 
     /// Retrieves the names of all tables included in a publication.
+    ///
+    /// Note: This excludes parent partitioned tables since they contain no data
+    /// and should not be processed. Only child partitions and regular tables are returned.
     pub async fn get_publication_table_names(
         &self,
         publication_name: &str,
     ) -> EtlResult<Vec<TableName>> {
         let publication_query = format!(
-            "select schemaname, tablename from pg_publication_tables where pubname = {};",
+            "select pt.schemaname, pt.tablename 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 = {}
+         and c.relkind != 'p'", /* Exclude parent partitioned tables */
             quote_literal(publication_name)
         );
 
@@ -399,6 +406,10 @@ impl PgReplicationClient {
     }
 
     /// Retrieves the OIDs of all tables included in a publication.
+    ///
+    /// Note: This excludes parent partitioned tables (relkind = 'p') since they
+    /// contain no data and should not be processed. Only child partitions and
+    /// regular tables are returned.
     pub async fn get_publication_table_ids(
         &self,
         publication_name: &str,
@@ -407,7 +418,8 @@ impl PgReplicationClient {
             "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 = {};",
+         where pt.pubname = {}
+         and c.relkind != 'p'", /* Exclude parent partitioned tables */
             quote_literal(publication_name)
         );
 
@@ -550,7 +562,7 @@ impl PgReplicationClient {
     /// Retrieves schema information for multiple tables.
     ///
     /// Tables without primary keys will be skipped and logged with a warning.
-    async fn get_table_schemas(
+    pub async fn get_table_schemas(
         &self,
         table_ids: &[TableId],
         publication_name: Option<&str>,

etl/tests/integration/partitioned_table_test.rs

@@ -22,7 +22,7 @@ async fn partitioned_table_sync_succeeds_with_inherited_primary_keys() {
         ("p3", "FROM (200) TO (300)"),
     ];
 
-    let (_parent_table_id, partition_table_ids) =
+    let (parent_table_id, partition_table_ids) =
         create_partitioned_table(&database, table_name.clone(), &partition_specs)
             .await
             .expect("Failed to create partitioned table");
@@ -80,17 +80,44 @@ async fn partitioned_table_sync_succeeds_with_inherited_primary_keys() {
     );
 
     let table_states = state_store.get_table_replication_states().await;
-    for (table_id, state) in &table_states {
-        if partition_table_ids.contains(table_id) {
-            assert!(
-                matches!(
-                    state.as_type(),
-                    TableReplicationPhaseType::SyncDone | TableReplicationPhaseType::Ready
-                ),
-                "Partition {} should be in SyncDone or Ready state, but was in {:?}",
-                table_id,
-                state.as_type()
-            );
-        }
+
+    assert_eq!(
+        table_states.len(),
+        partition_table_ids.len(),
+        "Expected {} partition states, but found {}",
+        partition_table_ids.len(),
+        table_states.len()
+    );
+
+    for &partition_id in &partition_table_ids {
+        let state = table_states
+            .get(&partition_id)
+            .expect(&format!("Partition {} should have a state", partition_id));
+        assert!(
+            matches!(
+                state.as_type(),
+                TableReplicationPhaseType::SyncDone | TableReplicationPhaseType::Ready
+            ),
+            "Partition {} should be in SyncDone or Ready state, but was in {:?}",
+            partition_id,
+            state.as_type()
+        );
     }
+
+    assert!(
+        !table_states.contains_key(&parent_table_id),
+        "Parent table {} should not be tracked since parent partitioned tables are excluded from processing",
+        parent_table_id
+    );
+
+    let parent_table_rows = table_rows
+        .iter()
+        .filter(|(table_id, _)| **table_id == parent_table_id)
+        .map(|(_, rows)| rows.len())
+        .sum::<usize>();
+    assert_eq!(
+        parent_table_rows, 0,
+        "Parent table {} should have no data since it's excluded from processing and all data goes to partitions, but found {} rows",
+        parent_table_id, parent_table_rows
+    );
 }