Improve

iambriccardo · iambriccardo · commit 9aa8a8e28c72 · 2025-12-15T12:08:06.000+01:00
diff --git a/etl-destinations/src/bigquery/core.rs b/etl-destinations/src/bigquery/core.rs
@@ -1,7 +1,7 @@
 use etl::destination::Destination;
 use etl::error::{ErrorKind, EtlError, EtlResult};
 use etl::store::schema::SchemaStore;
-use etl::store::state::{DestinationSchemaState, DestinationSchemaStateType, StateStore};
+use etl::store::state::{DestinationSchemaState, StateStore};
 use etl::types::{
     Cell, Event, ReplicatedTableSchema, SchemaDiff, TableId, TableName, TableRow,
     generate_sequence_number,
@@ -346,6 +346,18 @@ where
         // Note that if the table is deleted outside ETL and the cache marks it as created, the
         // inserts will fail because the table will be missing and won't be created.
         if !inner.created_tables.contains(&sequenced_bigquery_table_id) {
+            let snapshot_id = replicated_table_schema.get_inner().snapshot_id;
+            let replication_mask = replicated_table_schema.replication_mask().clone();
+
+            // Mark the state as applying BEFORE creating the table. This ensures that if the
+            // process crashes during table creation, we'll detect the incomplete state on restart.
+            self.state_store
+                .store_destination_schema_state(
+                    table_id,
+                    DestinationSchemaState::applying(snapshot_id, replication_mask.clone()),
+                )
+                .await?;
+
             self.client
                 .create_table_if_missing(
                     &self.dataset_id,
@@ -356,17 +368,11 @@ where
                 )
                 .await?;
 
-            // Store the initial destination schema state with the current snapshot and mask.
-            // This is needed so that when schema changes occur, we can reconstruct the old
-            // ReplicatedTableSchema for accurate diffing.
+            // Mark as applied after successful table creation.
             self.state_store
                 .store_destination_schema_state(
                     table_id,
-                    DestinationSchemaState {
-                        state: DestinationSchemaStateType::Applied,
-                        snapshot_id: replicated_table_schema.get_inner().snapshot_id,
-                        replication_mask: replicated_table_schema.replication_mask().clone(),
-                    },
+                    DestinationSchemaState::applied(snapshot_id, replication_mask),
                 )
                 .await?;
 
@@ -572,7 +578,7 @@ where
                     )
                 );
             }
-            Some(state) if state.state == DestinationSchemaStateType::Applying => {
+            Some(state) if state.is_applying() => {
                 // A previous schema change was interrupted, require manual intervention.
                 // The previous valid snapshot_id can be derived from table_schemas table
                 // by finding the second-highest snapshot_id for this table.
@@ -587,7 +593,7 @@ where
                     )
                 );
             }
-            Some(state) if state.state == DestinationSchemaStateType::Applied => {
+            Some(state) if state.is_applied() => {
                 let current_snapshot_id = state.snapshot_id;
                 let current_replication_mask = state.replication_mask.clone();
 
@@ -626,15 +632,16 @@ where
                 let old_schema =
                     ReplicatedTableSchema::from_mask(old_table_schema, current_replication_mask);
 
+                let new_replication_mask = new_schema.replication_mask().clone();
+
                 // Mark as applying before making changes (with the NEW snapshot_id and mask).
                 self.state_store
                     .store_destination_schema_state(
                         table_id,
-                        DestinationSchemaState {
-                            state: DestinationSchemaStateType::Applying,
-                            snapshot_id: new_snapshot_id,
-                            replication_mask: new_schema.replication_mask().clone(),
-                        },
+                        DestinationSchemaState::applying(
+                            new_snapshot_id,
+                            new_replication_mask.clone(),
+                        ),
                     )
                     .await?;
 
@@ -652,11 +659,7 @@ where
                 self.state_store
                     .store_destination_schema_state(
                         table_id,
-                        DestinationSchemaState {
-                            state: DestinationSchemaStateType::Applied,
-                            snapshot_id: new_snapshot_id,
-                            replication_mask: new_schema.replication_mask().clone(),
-                        },
+                        DestinationSchemaState::applied(new_snapshot_id, new_replication_mask),
                     )
                     .await?;
 
diff --git a/etl-destinations/tests/bigquery_pipeline.rs b/etl-destinations/tests/bigquery_pipeline.rs
@@ -4,6 +4,7 @@ use chrono::{DateTime, NaiveDate, NaiveDateTime, NaiveTime, Utc};
 use etl::config::BatchConfig;
 use etl::error::ErrorKind;
 use etl::state::table::TableReplicationPhaseType;
+use etl::store::state::StateStore;
 use etl::test_utils::database::{spawn_source_database, test_table_name};
 use etl::test_utils::notify::NotifyingStore;
 use etl::test_utils::pipeline::{create_pipeline, create_pipeline_with};
@@ -1776,13 +1777,6 @@ async fn table_validation_out_of_bounds_values() {
     );
 }
 
-/// Tests schema change support (ADD COLUMN, DROP COLUMN, RENAME COLUMN) in BigQuery.
-///
-/// This test verifies that DDL changes are correctly replicated to BigQuery by applying
-/// multiple schema changes (rename, drop, add) and verifying the final schema matches.
-/// The Storage Write API caches schema information, so after DDL changes, initial inserts
-/// may fail with "extra fields" errors until the cache refreshes. The retry logic in
-/// `stream_with_schema_retry` handles this by retrying on schema mismatch.
 #[tokio::test(flavor = "multi_thread")]
 async fn table_schema_change() {
     init_test_tracing();
@@ -1830,7 +1824,7 @@ async fn table_schema_change() {
     pipeline.start().await.unwrap();
     table_sync_done.notified().await;
 
-    // Insert initial row.
+    // Insert the initial row.
     let event_notify = destination
         .wait_for_events_count(vec![(EventType::Insert, 1)])
         .await;
@@ -1853,9 +1847,21 @@ async fn table_schema_change() {
         .unwrap();
     initial_schema.assert_columns(&["id", "name", "age", "status"]);
 
+    // Verify destination schema state is applied after initial table creation.
+    let initial_state = store
+        .get_destination_schema_state(&table_id)
+        .await
+        .unwrap()
+        .expect("destination schema state should exist after table creation");
+    assert!(
+        initial_state.is_applied(),
+        "initial destination schema state should be applied"
+    );
+    let initial_snapshot_id = initial_state.snapshot_id;
+
     // Apply multiple schema changes:
     // 1. Rename name -> full_name
-    // 2. Drop status column
+    // 2. Drop the status column
     // 3. Add email column
     //
     // Note: Each DDL change is captured via the DDL event trigger and stored in the schema
@@ -1913,7 +1919,7 @@ async fn table_schema_change() {
 
     pipeline.shutdown_and_wait().await.unwrap();
 
-    // Verify final schema:
+    // Verify the final schema:
     // - name should be renamed to full_name
     // - status should be dropped
     // - email should be added
@@ -1925,6 +1931,21 @@ async fn table_schema_change() {
     final_schema.assert_no_column("name");
     final_schema.assert_no_column("status");
 
+    // Verify destination schema state is applied after schema changes.
+    let final_state = store
+        .get_destination_schema_state(&table_id)
+        .await
+        .unwrap()
+        .expect("destination schema state should exist after schema change");
+    assert!(
+        final_state.is_applied(),
+        "final destination schema state should be applied"
+    );
+    assert!(
+        final_state.snapshot_id > initial_snapshot_id,
+        "snapshot_id should have increased after schema change"
+    );
+
     // Verify data was inserted correctly.
     let rows = bigquery_database.query_table(table_name).await.unwrap();
     assert_eq!(rows.len(), 2);
diff --git a/etl-destinations/tests/support/bigquery.rs b/etl-destinations/tests/support/bigquery.rs
@@ -166,10 +166,7 @@ impl BigQueryDatabase {
     /// Returns the column names and data types from INFORMATION_SCHEMA.COLUMNS.
     /// The table name pattern matches using REGEXP_CONTAINS to match the sequenced
     /// table name format: `{table_id}_{sequence_number}`.
-    pub async fn query_table_schema(
-        &self,
-        table_name: TableName,
-    ) -> Option<BigQueryTableSchema> {
+    pub async fn query_table_schema(&self, table_name: TableName) -> Option<BigQueryTableSchema> {
         let client = self.client().unwrap();
 
         let project_id = self.project_id();
@@ -1129,8 +1126,7 @@ impl BigQueryTableSchema {
 
         assert_eq!(
             actual_sorted, expected_sorted,
-            "schema columns mismatch. Expected: {:?}, Actual: {:?}",
-            expected_sorted, actual_sorted
+            "schema columns mismatch. Expected: {expected_sorted:?}, Actual: {actual_sorted:?}"
         );
     }
 
diff --git a/etl-postgres/src/types/schema.rs b/etl-postgres/src/types/schema.rs
@@ -252,7 +252,7 @@ impl TableSchema {
 /// Each element is either 0 (not replicated) or 1 (replicated), with indices
 /// corresponding to the columns in the table schema. Wrapped in [`Arc`] for
 /// efficient sharing across multiple events.
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq, Eq)]
 pub struct ReplicationMask(Arc<Vec<u8>>);
 
 impl ReplicationMask {
diff --git a/etl/src/store/state/base.rs b/etl/src/store/state/base.rs
@@ -1,43 +1,10 @@
-use etl_postgres::types::{ReplicationMask, SnapshotId, TableId};
+use etl_postgres::types::TableId;
 use std::{collections::HashMap, future::Future};
 
 use crate::error::EtlResult;
 use crate::state::table::TableReplicationPhase;
 
-/// The state of a schema change operation.
-#[derive(Debug, Clone, Copy, PartialEq, Eq)]
-pub enum DestinationSchemaStateType {
-    /// A schema change is currently being applied.
-    ///
-    /// If the system restarts and finds this state, it indicates that a previous
-    /// schema change was interrupted and manual intervention may be required.
-    /// The previous valid snapshot_id can be derived from table_schemas table.
-    Applying,
-    /// The schema has been successfully applied.
-    Applied,
-}
-
-/// Represents the state of the schema at a destination.
-///
-/// Used to track which schema version is currently applied at a destination
-/// and to detect interrupted schema changes that require recovery.
-///
-/// This structure tracks both the snapshot_id and the replication mask, which
-/// is needed to correctly reconstruct the [`ReplicatedTableSchema`] for diffing
-/// when schema changes occur.
-#[derive(Debug, Clone)]
-pub struct DestinationSchemaState {
-    /// The current state of the schema change operation.
-    pub state: DestinationSchemaStateType,
-    /// The current snapshot_id at the destination.
-    pub snapshot_id: SnapshotId,
-    /// The replication mask indicating which columns are replicated.
-    ///
-    /// This is stored alongside the snapshot_id so that when a schema change
-    /// occurs, we can reconstruct the old [`ReplicatedTableSchema`] with the
-    /// correct mask for accurate diffing.
-    pub replication_mask: ReplicationMask,
-}
+use super::DestinationSchemaState;
 
 /// Trait for storing and retrieving table replication state and mapping information.
 ///
diff --git a/etl/src/store/state/destination_state.rs b/etl/src/store/state/destination_state.rs
@@ -0,0 +1,74 @@
+use etl_postgres::types::{ReplicationMask, SnapshotId};
+
+/// The state of a schema change operation.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum DestinationSchemaStateType {
+    /// A schema change is currently being applied.
+    ///
+    /// If the system restarts and finds this state, it indicates that a previous
+    /// schema change was interrupted and manual intervention may be required.
+    /// The previous valid snapshot_id can be derived from table_schemas table.
+    Applying,
+    /// The schema has been successfully applied.
+    Applied,
+}
+
+/// Represents the state of the schema at a destination.
+///
+/// Used to track which schema version is currently applied at a destination
+/// and to detect interrupted schema changes that require recovery.
+///
+/// This structure tracks both the snapshot_id and the replication mask, which
+/// is needed to correctly reconstruct the [`ReplicatedTableSchema`] for diffing
+/// when schema changes occur.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct DestinationSchemaState {
+    /// The current state of the schema change operation.
+    pub state: DestinationSchemaStateType,
+    /// The current snapshot_id at the destination.
+    pub snapshot_id: SnapshotId,
+    /// The replication mask indicating which columns are replicated.
+    ///
+    /// This is stored alongside the snapshot_id so that when a schema change
+    /// occurs, we can reconstruct the old [`ReplicatedTableSchema`] with the
+    /// correct mask for accurate diffing.
+    pub replication_mask: ReplicationMask,
+}
+
+impl DestinationSchemaState {
+    /// Creates a new state indicating a schema change is being applied.
+    pub fn applying(snapshot_id: SnapshotId, replication_mask: ReplicationMask) -> Self {
+        Self {
+            state: DestinationSchemaStateType::Applying,
+            snapshot_id,
+            replication_mask,
+        }
+    }
+
+    /// Creates a new state indicating a schema has been successfully applied.
+    pub fn applied(snapshot_id: SnapshotId, replication_mask: ReplicationMask) -> Self {
+        Self {
+            state: DestinationSchemaStateType::Applied,
+            snapshot_id,
+            replication_mask,
+        }
+    }
+
+    /// Returns true if the state indicates a schema change is in progress.
+    pub fn is_applying(&self) -> bool {
+        self.state == DestinationSchemaStateType::Applying
+    }
+
+    /// Returns true if the state indicates the schema has been applied.
+    pub fn is_applied(&self) -> bool {
+        self.state == DestinationSchemaStateType::Applied
+    }
+
+    /// Transitions this state to applied, keeping the same snapshot_id and mask.
+    pub fn to_applied(self) -> Self {
+        Self {
+            state: DestinationSchemaStateType::Applied,
+            ..self
+        }
+    }
+}
diff --git a/etl/src/store/state/mod.rs b/etl/src/store/state/mod.rs
@@ -1,3 +1,5 @@
 mod base;
+mod destination_state;
 
 pub use base::*;
+pub use destination_state::*;
