feat(publications): Handle tables removed from the publication (#300)

Author: iambriccardo

AGENTS.md

@@ -19,6 +19,7 @@
 - Link types and methods as [`Type`], [`Type::method`].
 - Keep wording concise, correct, and punctuated; reword for clarity while preserving intent.
 - No code examples; include private helpers for maintainers; apply to modules, types, traits, impls, and functions.
+- Normal comments, should always finish with `.`.
 
 ## Rust Tests Execution
 - If output shows "0 passed; 0 failed; 0 ignored; n filtered out", tests did not run; treat as failure.

etl-api/src/db/pipelines.rs

@@ -220,11 +220,15 @@ pub async fn delete_pipeline_cascading(
 
     // Delete state, schema, and table mappings from the source database, only if ETL tables exist.
     if etl_present {
-        let _ =
-            state::delete_pipeline_replication_state(source_txn.deref_mut(), pipeline.id).await?;
-        let _ = schema::delete_pipeline_table_schemas(source_txn.deref_mut(), pipeline.id).await?;
-        let _ = table_mappings::delete_pipeline_table_mappings(source_txn.deref_mut(), pipeline.id)
+        let _ = state::delete_replication_state_for_all_tables(source_txn.deref_mut(), pipeline.id)
             .await?;
+        let _ = schema::delete_table_schemas_for_all_tables(source_txn.deref_mut(), pipeline.id)
+            .await?;
+        let _ = table_mappings::delete_table_mappings_for_all_tables(
+            source_txn.deref_mut(),
+            pipeline.id,
+        )
+        .await?;
     }
 
     // Here we finish `txn` before `source_txn` since we want the guarantee that the pipeline has

etl-postgres/src/replication/schema.rs

@@ -1,5 +1,5 @@
 use sqlx::postgres::PgRow;
-use sqlx::postgres::types::Oid;
+use sqlx::postgres::types::Oid as SqlxTableId;
 use sqlx::{PgExecutor, PgPool, Row};
 use std::collections::HashMap;
 use tokio_postgres::types::Type as PgType;
@@ -232,7 +232,7 @@ pub async fn load_table_schemas(
     let mut table_schemas = HashMap::new();
 
     for row in rows {
-        let table_oid: Oid = row.get("table_id");
+        let table_oid: SqlxTableId = row.get("table_id");
         let table_id = TableId::new(table_oid.0);
         let schema_name: String = row.get("schema_name");
         let table_name: String = row.get("table_name");
@@ -251,7 +251,7 @@ pub async fn load_table_schemas(
 ///
 /// Removes all table schema records and associated columns for the specified
 /// pipeline, using CASCADE delete for automatic cleanup of related column records.
-pub async fn delete_pipeline_table_schemas<'c, E>(
+pub async fn delete_table_schemas_for_all_tables<'c, E>(
     executor: E,
     pipeline_id: i64,
 ) -> Result<u64, sqlx::Error>
@@ -271,6 +271,31 @@ where
     Ok(result.rows_affected())
 }
 
+/// Deletes the table schema for a specific table in a pipeline.
+///
+/// Also cascades to `etl.table_columns` via FK.
+pub async fn delete_table_schema_for_table<'c, E>(
+    executor: E,
+    pipeline_id: i64,
+    table_id: TableId,
+) -> Result<u64, sqlx::Error>
+where
+    E: PgExecutor<'c>,
+{
+    let result = sqlx::query(
+        r#"
+        delete from etl.table_schemas
+        where pipeline_id = $1 and table_id = $2
+        "#,
+    )
+    .bind(pipeline_id)
+    .bind(SqlxTableId(table_id.into_inner()))
+    .execute(executor)
+    .await?;
+
+    Ok(result.rows_affected())
+}
+
 /// Builds a [`ColumnSchema`] from a database row.
 ///
 /// Assumes all required fields are present in the row after appropriate joins.

etl-postgres/src/replication/state.rs

@@ -328,7 +328,7 @@ pub async fn reset_replication_state(
 ///
 /// Removes all replication state records including historical entries
 /// for the specified pipeline. Used during pipeline cleanup.
-pub async fn delete_pipeline_replication_state<'c, E>(
+pub async fn delete_replication_state_for_all_tables<'c, E>(
     executor: E,
     pipeline_id: i64,
 ) -> sqlx::Result<u64>
@@ -348,6 +348,29 @@ where
     Ok(result.rows_affected())
 }
 
+/// Deletes all replication state entries for a specific table in a pipeline.
+pub async fn delete_replication_state_for_table<'c, E>(
+    executor: E,
+    pipeline_id: i64,
+    table_id: TableId,
+) -> sqlx::Result<u64>
+where
+    E: PgExecutor<'c>,
+{
+    let result = sqlx::query(
+        r#"
+        delete from etl.replication_state
+        where pipeline_id = $1 and table_id = $2
+        "#,
+    )
+    .bind(pipeline_id)
+    .bind(SqlxTableId(table_id.into_inner()))
+    .execute(executor)
+    .await?;
+
+    Ok(result.rows_affected())
+}
+
 /// Gets all table IDs that have replication state for a given pipeline.
 ///
 /// Returns a vector of table IDs that are currently being replicated for the specified pipeline.

etl-postgres/src/replication/table_mappings.rs

@@ -65,7 +65,7 @@ pub async fn load_table_mappings(
 ///
 /// Removes all table mapping records for the specified pipeline.
 /// Used during pipeline cleanup.
-pub async fn delete_pipeline_table_mappings<'c, E>(
+pub async fn delete_table_mappings_for_all_tables<'c, E>(
     executor: E,
     pipeline_id: i64,
 ) -> Result<u64, sqlx::Error>
@@ -84,3 +84,26 @@ where
 
     Ok(result.rows_affected())
 }
+
+/// Deletes a single table mapping for a given pipeline and source table id.
+pub async fn delete_table_mappings_for_table<'c, E>(
+    executor: E,
+    pipeline_id: i64,
+    source_table_id: &TableId,
+) -> Result<u64, sqlx::Error>
+where
+    E: PgExecutor<'c>,
+{
+    let result = sqlx::query(
+        r#"
+        delete from etl.table_mappings
+        where pipeline_id = $1 and source_table_id = $2
+        "#,
+    )
+    .bind(pipeline_id)
+    .bind(SqlxTableId(source_table_id.into_inner()))
+    .execute(executor)
+    .await?;
+
+    Ok(result.rows_affected())
+}

etl-replicator/src/core.rs

@@ -1,7 +1,9 @@
+use crate::migrations::migrate_state_store;
 use etl::destination::Destination;
 use etl::destination::memory::MemoryDestination;
 use etl::pipeline::Pipeline;
 use etl::store::both::postgres::PostgresStore;
+use etl::store::cleanup::CleanupStore;
 use etl::store::schema::SchemaStore;
 use etl::store::state::StateStore;
 use etl::types::PipelineId;
@@ -13,8 +15,6 @@ use secrecy::ExposeSecret;
 use tokio::signal::unix::{SignalKind, signal};
 use tracing::{debug, info, warn};
 
-use crate::migrations::migrate_state_store;
-
 /// Starts the replicator service with the provided configuration.
 ///
 /// Initializes the state store, creates the appropriate destination based on
@@ -138,7 +138,7 @@ fn log_batch_config(config: &BatchConfig) {
 async fn init_store(
     pipeline_id: PipelineId,
     pg_connection_config: PgConnectionConfig,
-) -> anyhow::Result<impl StateStore + SchemaStore + Clone> {
+) -> anyhow::Result<impl StateStore + SchemaStore + CleanupStore + Clone> {
     migrate_state_store(&pg_connection_config).await?;
 
     Ok(PostgresStore::new(pipeline_id, pg_connection_config))
@@ -152,7 +152,7 @@ async fn init_store(
 #[tracing::instrument(skip(pipeline), fields(pipeline_id = pipeline.id()))]
 async fn start_pipeline<S, D>(mut pipeline: Pipeline<S, D>) -> anyhow::Result<()>
 where
-    S: StateStore + SchemaStore + Clone + Send + Sync + 'static,
+    S: StateStore + SchemaStore + CleanupStore + Clone + Send + Sync + 'static,
     D: Destination + Clone + Send + Sync + 'static,
 {
     // Start the pipeline.

etl/src/failpoints.rs

@@ -9,7 +9,7 @@ use crate::bail;
 use crate::error::{ErrorKind, EtlError, EtlResult};
 
 pub const START_TABLE_SYNC__BEFORE_DATA_SYNC_SLOT_CREATION: &str =
-    "start_table_sync.befor_data_sync_slot_creation";
+    "start_table_sync.before_data_sync_slot_creation";
 pub const START_TABLE_SYNC__DURING_DATA_SYNC: &str = "start_table_sync.during_data_sync";
 
 /// Executes a configurable failpoint for testing error scenarios.

etl/src/pipeline.rs

@@ -3,24 +3,26 @@
 //! Contains the main [`Pipeline`] struct that coordinates Postgres logical replication
 //! with destination systems. Manages worker lifecycles, shutdown coordination, and error handling.
 
-use etl_config::shared::PipelineConfig;
-use std::sync::Arc;
-use tokio::sync::Semaphore;
-use tracing::{error, info};
-
 use crate::bail;
 use crate::concurrency::shutdown::{ShutdownTx, create_shutdown_channel};
 use crate::destination::Destination;
 use crate::error::{ErrorKind, EtlError, EtlResult};
 use crate::metrics::register_metrics;
 use crate::replication::client::PgReplicationClient;
 use crate::state::table::TableReplicationPhase;
+use crate::store::cleanup::CleanupStore;
 use crate::store::schema::SchemaStore;
 use crate::store::state::StateStore;
 use crate::types::PipelineId;
 use crate::workers::apply::{ApplyWorker, ApplyWorkerHandle};
 use crate::workers::base::{Worker, WorkerHandle};
 use crate::workers::pool::TableSyncWorkerPool;
+use etl_config::shared::PipelineConfig;
+use etl_postgres::types::TableId;
+use std::collections::HashSet;
+use std::sync::Arc;
+use tokio::sync::Semaphore;
+use tracing::{error, info};
 
 /// Internal state tracking for pipeline lifecycle.
 ///
@@ -62,7 +64,7 @@ pub struct Pipeline<S, D> {
 
 impl<S, D> Pipeline<S, D>
 where
-    S: StateStore + SchemaStore + Clone + Send + Sync + 'static,
+    S: StateStore + SchemaStore + CleanupStore + Clone + Send + Sync + 'static,
     D: Destination + Clone + Send + Sync + 'static,
 {
     /// Creates a new pipeline with the given configuration.
@@ -252,11 +254,17 @@ where
         self.wait().await
     }
 
-    /// Initializes table replication states for all tables in the publication.
+    /// Initializes table replication states for tables in the publication and
+    /// purges state for tables removed from it.
+    ///
+    /// Ensures each table currently in the Postgres publication has a
+    /// corresponding replication state; tables without existing states are
+    /// initialized to [`TableReplicationPhase::Init`].
     ///
-    /// 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.
+    /// Also detects tables for which we have stored state but are no longer
+    /// part of the publication, and deletes their stored state (replication
+    /// state, table mappings, and table schemas) without touching the actual
+    /// destination tables.
     async fn initialize_table_states(
         &self,
         replication_client: &PgReplicationClient,
@@ -281,26 +289,43 @@ where
             );
         }
 
-        let table_ids = replication_client
+        let publication_table_ids = replication_client
             .get_publication_table_ids(&self.config.publication_name)
             .await?;
 
         info!(
             "the publication '{}' contains {} tables",
             self.config.publication_name,
-            table_ids.len()
+            publication_table_ids.len()
         );
 
         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) {
+        let table_replication_states = self.store.get_table_replication_states().await?;
+
+        // Initialize states for newly added tables in the publication
+        for table_id in &publication_table_ids {
+            if !table_replication_states.contains_key(table_id) {
                 self.store
-                    .update_table_replication_state(table_id, TableReplicationPhase::Init)
+                    .update_table_replication_state(*table_id, TableReplicationPhase::Init)
                     .await?;
             }
         }
 
+        // Detect and purge tables that have been removed from the publication.
+        //
+        // We must not delete the destination table, only the internal state.
+        let publication_set: HashSet<TableId> = publication_table_ids.iter().copied().collect();
+        for (table_id, _) in table_replication_states {
+            if !publication_set.contains(&table_id) {
+                info!(
+                    "table {} removed from publication, purging stored state",
+                    table_id
+                );
+
+                self.store.cleanup_table_state(table_id).await?;
+            }
+        }
+
         Ok(())
     }
 }

etl/src/store/both/memory.rs

@@ -6,6 +6,7 @@ use tokio::sync::Mutex;
 use crate::error::{ErrorKind, EtlError, EtlResult};
 use crate::etl_error;
 use crate::state::table::TableReplicationPhase;
+use crate::store::cleanup::CleanupStore;
 use crate::store::schema::SchemaStore;
 use crate::store::state::StateStore;
 
@@ -198,3 +199,16 @@ impl SchemaStore for MemoryStore {
         Ok(())
     }
 }
+
+impl CleanupStore for MemoryStore {
+    async fn cleanup_table_state(&self, table_id: TableId) -> EtlResult<()> {
+        let mut inner = self.inner.lock().await;
+
+        inner.table_replication_states.remove(&table_id);
+        inner.table_state_history.remove(&table_id);
+        inner.table_schemas.remove(&table_id);
+        inner.table_mappings.remove(&table_id);
+
+        Ok(())
+    }
+}