fix(api): Make sure ETL tables are there before performing any action on the source etl schema (supabase#286)

Author: iambriccardo

etl-api/src/configs/pipeline.rs

@@ -6,6 +6,7 @@ use utoipa::ToSchema;
 const DEFAULT_BATCH_MAX_SIZE: usize = 1000000;
 const DEFAULT_BATCH_MAX_FILL_MS: u64 = 10000;
 const DEFAULT_TABLE_ERROR_RETRY_DELAY_MS: u64 = 10000;
+const DEFAULT_MAX_TABLE_SYNC_WORKERS: u16 = 4;
 
 #[derive(Debug, Clone, Serialize, Deserialize, ToSchema)]
 pub struct FullApiPipelineConfig {
@@ -104,7 +105,9 @@ impl From<FullApiPipelineConfig> for StoredPipelineConfig {
             table_error_retry_delay_ms: value
                 .table_error_retry_delay_ms
                 .unwrap_or(DEFAULT_TABLE_ERROR_RETRY_DELAY_MS),
-            max_table_sync_workers: value.max_table_sync_workers.unwrap_or_default(),
+            max_table_sync_workers: value
+                .max_table_sync_workers
+                .unwrap_or(DEFAULT_MAX_TABLE_SYNC_WORKERS),
         }
     }
 }
@@ -173,7 +176,10 @@ mod tests {
             stored.table_error_retry_delay_ms,
             DEFAULT_TABLE_ERROR_RETRY_DELAY_MS
         );
-        assert_eq!(stored.max_table_sync_workers, 0);
+        assert_eq!(
+            stored.max_table_sync_workers,
+            DEFAULT_MAX_TABLE_SYNC_WORKERS
+        );
     }
 
     #[test]

etl-api/src/db/pipelines.rs

@@ -1,4 +1,4 @@
-use etl_postgres::replication::{schema, slots, state, table_mappings};
+use etl_postgres::replication::{health, schema, slots, state, table_mappings};
 use sqlx::{PgExecutor, PgTransaction};
 use std::ops::DerefMut;
 use thiserror::Error;
@@ -210,22 +210,34 @@ pub async fn delete_pipeline_cascading(
         db::destinations::delete_destination(txn.deref_mut(), tenant_id, destination.id).await?;
     }
 
-    // Get all table IDs for this pipeline before deleting state.
-    let table_ids = state::get_pipeline_table_ids(source_txn.deref_mut(), pipeline.id).await?;
+    // Get all table IDs for this pipeline before deleting state (only if all ETL tables exist).
+    let etl_present = health::etl_tables_present(source_txn.deref_mut()).await?;
+    let table_ids = if etl_present {
+        Some(state::get_pipeline_table_ids(source_txn.deref_mut(), pipeline.id).await?)
+    } else {
+        None
+    };
 
-    // Delete state, schema, and table mappings from the source database
-    state::delete_pipeline_replication_state(source_txn.deref_mut(), pipeline.id).await?;
-    schema::delete_pipeline_table_schemas(source_txn.deref_mut(), pipeline.id).await?;
-    table_mappings::delete_pipeline_table_mappings(source_txn.deref_mut(), pipeline.id).await?;
+    // 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)
+            .await?;
+    }
 
     // Here we finish `txn` before `source_txn` since we want the guarantee that the pipeline has
     // been deleted before committing the state and slots deletions.
     txn.commit().await?;
     source_txn.commit().await?;
 
-    // If we succeeded to commit both transactions, we are safe to delete the slots. The reason for
-    // not deleting slots in the transaction is that `pg_drop_replication_slot(...)` is not transactional.
-    slots::delete_pipeline_replication_slots(&source_pool, pipeline.id as u64, &table_ids).await?;
+    if let Some(table_ids) = table_ids {
+        // If we succeeded to commit both transactions, we are safe to delete the slots. The reason for
+        // not deleting slots in the transaction is that `pg_drop_replication_slot(...)` is not transactional.
+        slots::delete_pipeline_replication_slots(&source_pool, pipeline.id as u64, &table_ids)
+            .await?;
+    }
 
     Ok(())
 }

etl-api/src/routes/pipelines.rs

@@ -5,7 +5,7 @@ use actix_web::{
     web::{Data, Json, Path},
 };
 use etl_config::shared::{ReplicatorConfig, SupabaseConfig, TlsConfig};
-use etl_postgres::replication::{TableLookupError, get_table_name_from_oid, state};
+use etl_postgres::replication::{TableLookupError, get_table_name_from_oid, health, state};
 use etl_postgres::schema::TableId;
 use secrecy::ExposeSecret;
 use serde::{Deserialize, Serialize};
@@ -62,6 +62,9 @@ pub enum PipelineError {
     #[error("The table state is not rollbackable: {0}")]
     NotRollbackable(String),
 
+    #[error("The ETL table state has not been initialized first")]
+    EtlStateNotInitialized,
+
     #[error("invalid destination config")]
     InvalidConfig(#[from] serde_json::Error),
 
@@ -145,9 +148,9 @@ impl ResponseError for PipelineError {
             | PipelineError::InvalidTableReplicationState(_)
             | PipelineError::MissingTableReplicationState => StatusCode::INTERNAL_SERVER_ERROR,
             PipelineError::NotRollbackable(_) => StatusCode::BAD_REQUEST,
-            PipelineError::PipelineNotFound(_) | PipelineError::ImageNotFoundById(_) => {
-                StatusCode::NOT_FOUND
-            }
+            PipelineError::PipelineNotFound(_)
+            | PipelineError::ImageNotFoundById(_)
+            | PipelineError::EtlStateNotInitialized => StatusCode::NOT_FOUND,
             PipelineError::TenantId(_)
             | PipelineError::SourceNotFound(_)
             | PipelineError::DestinationNotFound(_) => StatusCode::BAD_REQUEST,
@@ -811,6 +814,11 @@ pub async fn get_pipeline_replication_status(
     let source_pool =
         connect_to_source_database_with_defaults(&source.config.into_connection_config()).await?;
 
+    // Ensure ETL tables exist in the source DB
+    if !health::etl_tables_present(&source_pool).await? {
+        return Err(PipelineError::EtlStateNotInitialized);
+    }
+
     // Fetch replication state for all tables in this pipeline
     let state_rows = state::get_table_replication_state_rows(&source_pool, pipeline_id).await?;
 
@@ -894,6 +902,11 @@ pub async fn rollback_table_state(
     let source_pool =
         connect_to_source_database_with_defaults(&source.config.into_connection_config()).await?;
 
+    // Ensure ETL tables exist in the source DB
+    if !health::etl_tables_present(&source_pool).await? {
+        return Err(PipelineError::EtlStateNotInitialized);
+    }
+
     // First, check current state to ensure it's rollbackable (manual retry policy)
     let state_rows = state::get_table_replication_state_rows(&source_pool, pipeline_id).await?;
     let current_row = state_rows

etl-api/tests/common/test_app.rs

@@ -505,7 +505,7 @@ pub async fn spawn_test_app() -> TestApp {
     // We use a random database name.
     config.database.name = Uuid::new_v4().to_string();
 
-    let connection_pool = create_etl_api_database(&config.database).await;
+    let api_db_pool = create_etl_api_database(&config.database).await;
 
     let key = generate_random_key::<32>().expect("failed to generate random key");
     let encryption_key = encryption::EncryptionKey { id: 0, key };
@@ -515,7 +515,7 @@ pub async fn spawn_test_app() -> TestApp {
     let server = run(
         config.clone(),
         listener,
-        connection_pool,
+        api_db_pool,
         encryption_key,
         k8s_client,
     )

etl-api/tests/integration/etl_tables_missing_test.rs

@@ -0,0 +1,108 @@
+use etl_api::routes::pipelines::{
+    CreatePipelineRequest, CreatePipelineResponse, RollbackTableStateRequest, RollbackType,
+};
+use etl_config::shared::PgConnectionConfig;
+use etl_postgres::sqlx::test_utils::drop_pg_database;
+use etl_telemetry::tracing::init_test_tracing;
+use reqwest::StatusCode;
+use sqlx::PgPool;
+
+use crate::common::database::create_test_source_database;
+use crate::common::test_app::{TestApp, spawn_test_app};
+use crate::integration::destination_test::create_destination;
+use crate::integration::images_test::create_default_image;
+use crate::integration::tenants_test::create_tenant;
+
+async fn create_pipeline_with_unmigrated_source_db(
+    app: &TestApp,
+    tenant_id: &str,
+) -> (i64, PgPool, PgConnectionConfig) {
+    let (source_pool, source_id, source_db_config) =
+        create_test_source_database(app, tenant_id).await;
+    let destination_id = create_destination(app, tenant_id).await;
+    create_default_image(app).await;
+
+    let req = CreatePipelineRequest {
+        source_id,
+        destination_id,
+        config: crate::integration::pipelines_test::new_pipeline_config(),
+    };
+
+    let response = app.create_pipeline(tenant_id, &req).await;
+    let response: CreatePipelineResponse = response
+        .json()
+        .await
+        .expect("failed to deserialize response");
+
+    (response.id, source_pool, source_db_config)
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn replication_status_fails_when_etl_tables_missing() {
+    init_test_tracing();
+    let app = spawn_test_app().await;
+    let tenant_id = create_tenant(&app).await;
+
+    let (pipeline_id, _source_pool, source_db_config) =
+        create_pipeline_with_unmigrated_source_db(&app, &tenant_id).await;
+
+    let response = app
+        .get_pipeline_replication_status(&tenant_id, pipeline_id)
+        .await;
+    assert_eq!(response.status(), StatusCode::NOT_FOUND);
+    assert!(
+        response
+            .text()
+            .await
+            .unwrap()
+            .contains("The ETL table state has not been initialized first")
+    );
+
+    drop_pg_database(&source_db_config).await;
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn rollback_fails_when_etl_tables_missing() {
+    init_test_tracing();
+    let app = spawn_test_app().await;
+    let tenant_id = create_tenant(&app).await;
+
+    let (pipeline_id, _source_pool, source_db_config) =
+        create_pipeline_with_unmigrated_source_db(&app, &tenant_id).await;
+
+    let req = RollbackTableStateRequest {
+        table_id: 1,
+        rollback_type: RollbackType::Individual,
+    };
+    let response = app
+        .rollback_table_state(&tenant_id, pipeline_id, &req)
+        .await;
+    assert_eq!(response.status(), StatusCode::NOT_FOUND);
+    assert!(
+        response
+            .text()
+            .await
+            .unwrap()
+            .contains("The ETL table state has not been initialized first")
+    );
+
+    drop_pg_database(&source_db_config).await;
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn deleting_pipeline_succeeds_when_etl_tables_missing() {
+    init_test_tracing();
+    let app = spawn_test_app().await;
+    let tenant_id = create_tenant(&app).await;
+
+    let (pipeline_id, _source_pool, source_db_config) =
+        create_pipeline_with_unmigrated_source_db(&app, &tenant_id).await;
+
+    let response = app.delete_pipeline(&tenant_id, pipeline_id).await;
+    assert_eq!(response.status(), StatusCode::OK);
+
+    let response = app.read_pipeline(&tenant_id, pipeline_id).await;
+    assert_eq!(response.status(), StatusCode::NOT_FOUND);
+
+    drop_pg_database(&source_db_config).await;
+}

etl-api/tests/integration/mod.rs

@@ -1,5 +1,6 @@
 mod destination_test;
 mod destinations_pipelines_test;
+mod etl_tables_missing_test;
 mod health_check_test;
 mod images_test;
 mod pipelines_test;