From 9b4feb6c4373778f37c3f6a28267ab9bd1603f47 Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Mon, 29 Sep 2025 08:52:19 +0200 Subject: [PATCH 01/45] feat(replication): Add support for relation changes --- etl-replicator/configuration/base.yaml | 4 ++-- etl/src/conversions/event.rs | 28 +++++++++-------------- etl/src/replication/apply.rs | 31 ++------------------------ etl/src/types/event.rs | 15 ++++++++++--- 4 files changed, 27 insertions(+), 51 deletions(-) diff --git a/etl-replicator/configuration/base.yaml b/etl-replicator/configuration/base.yaml index 7dd8facff..05dc474ca 100644 --- a/etl-replicator/configuration/base.yaml +++ b/etl-replicator/configuration/base.yaml @@ -1,6 +1,6 @@ pipeline: - id: 3 - publication_name: "replicator_publication" + id: 4 + publication_name: "test_pub" batch: max_size: 10000 max_fill_ms: 1000 diff --git a/etl/src/conversions/event.rs b/etl/src/conversions/event.rs index 81c9867df..16920e072 100644 --- a/etl/src/conversions/event.rs +++ b/etl/src/conversions/event.rs @@ -54,30 +54,24 @@ pub fn parse_event_from_commit_message( /// /// This method parses the replication protocol relation message and builds /// a complete table schema for use in interpreting subsequent data events. -pub fn parse_event_from_relation_message( +pub fn parse_event_from_relation_message( + schema_store: &S, start_lsn: PgLsn, commit_lsn: PgLsn, relation_body: &protocol::RelationBody, -) -> EtlResult { - let table_name = TableName::new( - relation_body.namespace()?.to_string(), - relation_body.name()?.to_string(), - ); - let column_schemas = relation_body - .columns() - .iter() - .map(build_column_schema) - .collect::, _>>()?; - let table_schema = TableSchema::new( - TableId::new(relation_body.rel_id()), - table_name, - column_schemas, - ); +) -> EtlResult +where + S: SchemaStore, +{ + let table_id = relation_body.rel_id(); + + let column_schemas = relation_body.columns().iter().map(build_column_schema); Ok(RelationEvent { start_lsn, commit_lsn, - table_schema, + changes: vec![], + table_id: TableId::new(table_id), }) } diff --git a/etl/src/replication/apply.rs b/etl/src/replication/apply.rs index 0cfe52371..301393ce4 100644 --- a/etl/src/replication/apply.rs +++ b/etl/src/replication/apply.rs @@ -1115,36 +1115,9 @@ where return Ok(HandleMessageResult::no_event()); } - // If no table schema is found, it means that something went wrong since we should have schemas - // ready before starting the apply loop. - let existing_table_schema = - schema_store - .get_table_schema(&table_id) - .await? - .ok_or_else(|| { - etl_error!( - ErrorKind::MissingTableSchema, - "Table not found in the schema cache", - format!("The table schema for table {table_id} was not found in the cache") - ) - })?; - // Convert event from the protocol message. - let event = parse_event_from_relation_message(start_lsn, remote_final_lsn, message)?; - - // We compare the table schema from the relation message with the existing schema (if any). - // The purpose of this comparison is that we want to throw an error and stop the processing - // of any table that incurs in a schema change after the initial table sync is performed. - if !existing_table_schema.partial_eq(&event.table_schema) { - let error = TableReplicationError::with_solution( - table_id, - format!("The schema for table {table_id} has changed during streaming"), - "ETL doesn't support schema changes at this point in time, rollback the schema", - RetryPolicy::ManualRetry, - ); - - return Ok(HandleMessageResult::finish_batch_and_exclude_event(error)); - } + let event = + parse_event_from_relation_message(schema_store, start_lsn, remote_final_lsn, message)?; Ok(HandleMessageResult::return_event(Event::Relation(event))) } diff --git a/etl/src/types/event.rs b/etl/src/types/event.rs index 144823ac8..aeac8c3df 100644 --- a/etl/src/types/event.rs +++ b/etl/src/types/event.rs @@ -1,4 +1,4 @@ -use etl_postgres::types::{TableId, TableSchema}; +use etl_postgres::types::{ColumnSchema, TableId, TableSchema}; use std::fmt; use tokio_postgres::types::PgLsn; @@ -40,6 +40,13 @@ pub struct CommitEvent { pub timestamp: i64, } +#[derive(Debug, Clone, PartialEq)] +pub enum RelationChange { + AddColumn(ColumnSchema), + DropColumn(ColumnSchema), + AlterColumn(ColumnSchema), +} + /// Table schema definition event from Postgres logical replication. /// /// [`RelationEvent`] provides schema information for tables involved in replication. @@ -51,8 +58,10 @@ pub struct RelationEvent { pub start_lsn: PgLsn, /// LSN position where the transaction of this event will commit. pub commit_lsn: PgLsn, - /// Complete table schema including columns and types. - pub table_schema: TableSchema, + /// Set of changes for this table (compared to the most recent version of the schema). + pub changes: Vec, + /// ID of the table of which this is a schema change. + pub table_id: TableId, } /// Row insertion event from Postgres logical replication. From 077d081f934cb059a9a58c45c19050a9adcc5b7e Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Tue, 30 Sep 2025 09:30:22 -0700 Subject: [PATCH 02/45] Improve --- etl-postgres/src/types/schema.rs | 28 ----------- etl/src/conversions/event.rs | 85 +++++++++++++++++++++++++++----- 2 files changed, 72 insertions(+), 41 deletions(-) diff --git a/etl-postgres/src/types/schema.rs b/etl-postgres/src/types/schema.rs index c66ab41f1..d2f6eeb27 100644 --- a/etl-postgres/src/types/schema.rs +++ b/etl-postgres/src/types/schema.rs @@ -82,20 +82,6 @@ impl ColumnSchema { primary, } } - - /// Compares two [`ColumnSchema`] instances, excluding the `nullable` field. - /// - /// Return `true` if all fields except `nullable` are equal, `false` otherwise. - /// - /// This method is used for comparing table schemas loaded via the initial table sync and the - /// relation messages received via CDC. The reason for skipping the `nullable` field is that - /// unfortunately Postgres doesn't seem to propagate nullable information of a column via - /// relation messages. The reason for skipping the `primary` field is that if the replica - /// identity of a table is set to full, the relation message sets all columns as primary - /// key, irrespective of what the actual primary key in the table is. - fn partial_eq(&self, other: &ColumnSchema) -> bool { - self.name == other.name && self.typ == other.typ && self.modifier == other.modifier - } } /// A type-safe wrapper for Postgres table OIDs. @@ -214,20 +200,6 @@ impl TableSchema { pub fn has_primary_keys(&self) -> bool { self.column_schemas.iter().any(|cs| cs.primary) } - - /// Compares two [`TableSchema`] instances, excluding the [`ColumnSchema`]'s `nullable` field. - /// - /// Return `true` if all fields except `nullable` are equal, `false` otherwise. - pub fn partial_eq(&self, other: &TableSchema) -> bool { - self.id == other.id - && self.name == other.name - && self.column_schemas.len() == other.column_schemas.len() - && self - .column_schemas - .iter() - .zip(other.column_schemas.iter()) - .all(|(c1, c2)| c1.partial_eq(c2)) - } } impl PartialOrd for TableSchema { diff --git a/etl/src/conversions/event.rs b/etl/src/conversions/event.rs index 16920e072..11fe07232 100644 --- a/etl/src/conversions/event.rs +++ b/etl/src/conversions/event.rs @@ -3,18 +3,43 @@ use etl_postgres::types::{ ColumnSchema, TableId, TableName, TableSchema, convert_type_oid_to_type, }; use postgres_replication::protocol; +use std::collections::HashSet; +use std::hash::Hash; use std::sync::Arc; use tokio_postgres::types::PgLsn; use crate::conversions::text::{default_value_for_type, parse_cell_from_postgres_text}; -use crate::error::{ErrorKind, EtlResult}; +use crate::error::{ErrorKind, EtlError, EtlResult}; use crate::store::schema::SchemaStore; -use crate::types::{ - BeginEvent, Cell, CommitEvent, DeleteEvent, InsertEvent, RelationEvent, TableRow, - TruncateEvent, UpdateEvent, -}; +use crate::types::{BeginEvent, Cell, CommitEvent, DeleteEvent, InsertEvent, RelationChange, RelationEvent, TableRow, TruncateEvent, UpdateEvent}; use crate::{bail, etl_error}; +#[derive(Debug, Clone)] +struct IndexedColumnSchema(ColumnSchema); + +impl IndexedColumnSchema { + + fn into_inner(self) -> ColumnSchema { + self.0 + } +} + +impl Eq for IndexedColumnSchema {} + +impl PartialEq for IndexedColumnSchema { + + fn eq(&self, other: &Self) -> bool { + self.0.name == other.0.name + } +} + +impl Hash for IndexedColumnSchema { + + fn hash(&self, state: &mut H) { + self.0.name.hash(state); + } +} + /// Creates a [`BeginEvent`] from Postgres protocol data. /// /// This method parses the replication protocol begin message and extracts @@ -54,7 +79,7 @@ pub fn parse_event_from_commit_message( /// /// This method parses the replication protocol relation message and builds /// a complete table schema for use in interpreting subsequent data events. -pub fn parse_event_from_relation_message( +pub async fn parse_event_from_relation_message( schema_store: &S, start_lsn: PgLsn, commit_lsn: PgLsn, @@ -63,15 +88,47 @@ pub fn parse_event_from_relation_message( where S: SchemaStore, { - let table_id = relation_body.rel_id(); + let table_id = TableId::new(relation_body.rel_id()); + + let Some(existing_column_schemas) = schema_store.get_table_schema(&table_id).await? else { + bail!( + ErrorKind::MissingTableSchema, + "Table not found in the schema cache", + format!("The table schema for table {table_id} was not found in the cache") + ) + }; + + let mut latest_column_schemas = relation_body + .columns() + .iter() + .map(build_indexed_column_schema) + .collect::, EtlError>>()?; + + let mut changes = vec![]; + for column_schema in existing_column_schemas { + let latest_column_schema = latest_column_schemas.take(column_schema); + match latest_column_schema { + Some(latest_column_schema) => { + changes.push(RelationChange::AlterColumn(latest_column_schema.into_inner())); + } + None => { + // If we don't find the column in the latest schema, we assume it was dropped even + // though it could be renamed. + changes.push(RelationChange::DropColumn(column_schema)); + } + } + } - let column_schemas = relation_body.columns().iter().map(build_column_schema); + // For the remaining columns that didn't match, we assume they were added. + for column_schema in latest_column_schemas { + changes.push(RelationChange::AddColumn(column_schema.into_inner())); + } Ok(RelationEvent { start_lsn, commit_lsn, changes: vec![], - table_id: TableId::new(table_id), + table_id, }) } @@ -238,13 +295,13 @@ where }) } -/// Constructs a [`ColumnSchema`] from Postgres protocol column data. +/// Constructs a [`IndexedColumnSchema`] from Postgres protocol column data. /// /// This helper method extracts column metadata from the replication protocol /// and converts it into the internal column schema representation. Some fields /// like nullable status have default values due to protocol limitations. -fn build_column_schema(column: &protocol::Column) -> EtlResult { - Ok(ColumnSchema::new( +fn build_indexed_column_schema(column: &protocol::Column) -> EtlResult { + let column_schema = ColumnSchema::new( column.name()?.to_string(), convert_type_oid_to_type(column.type_id() as u32), column.type_modifier(), @@ -254,7 +311,9 @@ fn build_column_schema(column: &protocol::Column) -> EtlResult { false, // Currently 1 means that the column is part of the primary key. column.flags() == 1, - )) + ); + + Ok(IndexedColumnSchema(column_schema)) } /// Converts Postgres tuple data into a [`TableRow`] using column schemas. From c42451a13149abbde8ccc9228be8bbfa4b1cd9e7 Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Tue, 30 Sep 2025 10:34:55 -0700 Subject: [PATCH 03/45] Improve --- etl-postgres/src/tokio/test_utils.rs | 63 ++++++------- etl/src/conversions/event.rs | 26 ++--- etl/src/replication/apply.rs | 21 +---- etl/src/test_utils/test_schema.rs | 2 +- etl/src/types/event.rs | 6 +- etl/tests/pipeline.rs | 136 ++++++++++++++++++++++++++- etl/tests/replica_identity.rs | 10 +- etl/tests/replication.rs | 2 +- 8 files changed, 194 insertions(+), 72 deletions(-) diff --git a/etl-postgres/src/tokio/test_utils.rs b/etl-postgres/src/tokio/test_utils.rs index 6692e486b..9100ccc7a 100644 --- a/etl-postgres/src/tokio/test_utils.rs +++ b/etl-postgres/src/tokio/test_utils.rs @@ -12,22 +12,15 @@ use crate::types::{ColumnSchema, TableId, TableName}; /// Table modification operations for ALTER TABLE statements. pub enum TableModification<'a> { /// Add a new column with specified name and data type. - AddColumn { - name: &'a str, - data_type: &'a str, - }, + AddColumn { name: &'a str, params: &'a str }, /// Drop an existing column by name. - DropColumn { - name: &'a str, - }, + DropColumn { name: &'a str }, /// Alter an existing column with the specified alteration. - AlterColumn { - name: &'a str, - alteration: &'a str, - }, - ReplicaIdentity { - value: &'a str, - }, + AlterColumn { name: &'a str, params: &'a str }, + /// Rename an existing column. + RenameColumn { name: &'a str, new_name: &'a str }, + /// Change the replica identity setting for the table. + ReplicaIdentity { value: &'a str }, } /// Postgres database wrapper for testing operations. @@ -183,35 +176,39 @@ impl PgDatabase { table_name: TableName, modifications: &[TableModification<'_>], ) -> Result<(), tokio_postgres::Error> { - let modifications_str = modifications - .iter() - .map(|modification| match modification { - TableModification::AddColumn { name, data_type } => { + for modification in modifications { + let modification_str = match modification { + TableModification::AddColumn { + name, + params: data_type, + } => { format!("add column {name} {data_type}") } TableModification::DropColumn { name } => { format!("drop column {name}") } - TableModification::AlterColumn { name, alteration } => { + TableModification::AlterColumn { + name, + params: alteration, + } => { format!("alter column {name} {alteration}") } + TableModification::RenameColumn { name, new_name } => { + format!("rename column {name} to {new_name}") + } TableModification::ReplicaIdentity { value } => { format!("replica identity {value}") } - }) - .collect::>() - .join(", "); + }; - let alter_table_query = format!( - "alter table {} {}", - table_name.as_quoted_identifier(), - modifications_str - ); - self.client - .as_ref() - .unwrap() - .execute(&alter_table_query, &[]) - .await?; + let query = format!( + "alter table {} {}", + table_name.as_quoted_identifier(), + modification_str + ); + + self.client.as_ref().unwrap().execute(&query, &[]).await?; + } Ok(()) } @@ -224,7 +221,7 @@ impl PgDatabase { &self, table_name: TableName, columns: &[&str], - values: &[&(dyn tokio_postgres::types::ToSql + Sync)], + values: &[&(dyn ToSql + Sync)], ) -> Result { let columns_str = columns.join(", "); let placeholders: Vec = (1..=values.len()).map(|i| format!("${i}")).collect(); diff --git a/etl/src/conversions/event.rs b/etl/src/conversions/event.rs index 11fe07232..5d7711021 100644 --- a/etl/src/conversions/event.rs +++ b/etl/src/conversions/event.rs @@ -1,7 +1,5 @@ use core::str; -use etl_postgres::types::{ - ColumnSchema, TableId, TableName, TableSchema, convert_type_oid_to_type, -}; +use etl_postgres::types::{ColumnSchema, TableId, TableSchema, convert_type_oid_to_type}; use postgres_replication::protocol; use std::collections::HashSet; use std::hash::Hash; @@ -11,14 +9,16 @@ use tokio_postgres::types::PgLsn; use crate::conversions::text::{default_value_for_type, parse_cell_from_postgres_text}; use crate::error::{ErrorKind, EtlError, EtlResult}; use crate::store::schema::SchemaStore; -use crate::types::{BeginEvent, Cell, CommitEvent, DeleteEvent, InsertEvent, RelationChange, RelationEvent, TableRow, TruncateEvent, UpdateEvent}; +use crate::types::{ + BeginEvent, Cell, CommitEvent, DeleteEvent, InsertEvent, RelationChange, RelationEvent, + TableRow, TruncateEvent, UpdateEvent, +}; use crate::{bail, etl_error}; #[derive(Debug, Clone)] struct IndexedColumnSchema(ColumnSchema); impl IndexedColumnSchema { - fn into_inner(self) -> ColumnSchema { self.0 } @@ -27,14 +27,12 @@ impl IndexedColumnSchema { impl Eq for IndexedColumnSchema {} impl PartialEq for IndexedColumnSchema { - fn eq(&self, other: &Self) -> bool { self.0.name == other.0.name } } impl Hash for IndexedColumnSchema { - fn hash(&self, state: &mut H) { self.0.name.hash(state); } @@ -105,16 +103,22 @@ where .collect::, EtlError>>()?; let mut changes = vec![]; - for column_schema in existing_column_schemas { - let latest_column_schema = latest_column_schemas.take(column_schema); + for column_schema in existing_column_schemas.column_schemas.iter() { + let column_schema = IndexedColumnSchema(column_schema.clone()); + let latest_column_schema = latest_column_schemas.take(&column_schema); match latest_column_schema { Some(latest_column_schema) => { - changes.push(RelationChange::AlterColumn(latest_column_schema.into_inner())); + // If we find a column with the same name, we assume it was changed. The only changes + // that we detect are changes to the column but with preserved name. + changes.push(RelationChange::AlterColumn( + column_schema.into_inner(), + latest_column_schema.into_inner(), + )); } None => { // If we don't find the column in the latest schema, we assume it was dropped even // though it could be renamed. - changes.push(RelationChange::DropColumn(column_schema)); + changes.push(RelationChange::DropColumn(column_schema.into_inner())); } } } diff --git a/etl/src/replication/apply.rs b/etl/src/replication/apply.rs index 301393ce4..1146aba61 100644 --- a/etl/src/replication/apply.rs +++ b/etl/src/replication/apply.rs @@ -13,6 +13,7 @@ use tokio::pin; use tokio_postgres::types::PgLsn; use tracing::{debug, info}; +use crate::bail; use crate::concurrency::shutdown::ShutdownRx; use crate::concurrency::signal::SignalRx; use crate::concurrency::stream::{TimeoutStream, TimeoutStreamResult}; @@ -31,10 +32,9 @@ use crate::metrics::{ }; use crate::replication::client::PgReplicationClient; use crate::replication::stream::EventsStream; -use crate::state::table::{RetryPolicy, TableReplicationError}; +use crate::state::table::TableReplicationError; use crate::store::schema::SchemaStore; use crate::types::{Event, PipelineId}; -use crate::{bail, etl_error}; /// The amount of milliseconds that pass between one refresh and the other of the system, in case no /// events or shutdown signal are received. @@ -210,8 +210,6 @@ impl StatusUpdate { enum EndBatch { /// The batch should include the last processed event and end. Inclusive, - /// The batch should exclude the last processed event and end. - Exclusive, } /// Result returned from `handle_replication_message` and related functions @@ -286,18 +284,6 @@ impl HandleMessageResult { ..Default::default() } } - - /// Creates a result that excludes the current event and requests batch termination. - /// - /// Used when the current message triggers a recoverable table-level error. - /// The error is propagated to be handled by the apply loop hook. - fn finish_batch_and_exclude_event(error: TableReplicationError) -> Self { - Self { - end_batch: Some(EndBatch::Exclusive), - table_replication_error: Some(error), - ..Default::default() - } - } } /// A shared state that is used throughout the apply loop to track progress. @@ -1117,7 +1103,8 @@ where // Convert event from the protocol message. let event = - parse_event_from_relation_message(schema_store, start_lsn, remote_final_lsn, message)?; + parse_event_from_relation_message(schema_store, start_lsn, remote_final_lsn, message) + .await?; Ok(HandleMessageResult::return_event(Event::Relation(event))) } diff --git a/etl/src/test_utils/test_schema.rs b/etl/src/test_utils/test_schema.rs index 8faa6449b..caf2d64bc 100644 --- a/etl/src/test_utils/test_schema.rs +++ b/etl/src/test_utils/test_schema.rs @@ -236,7 +236,7 @@ pub fn events_equal_excluding_fields(left: &Event, right: &Event) -> bool { (Event::Delete(left), Event::Delete(right)) => { left.table_id == right.table_id && left.old_table_row == right.old_table_row } - (Event::Relation(left), Event::Relation(right)) => left.table_schema == right.table_schema, + (Event::Relation(left), Event::Relation(right)) => left.table_id == right.table_id, (Event::Truncate(left), Event::Truncate(right)) => { left.options == right.options && left.rel_ids == right.rel_ids } diff --git a/etl/src/types/event.rs b/etl/src/types/event.rs index aeac8c3df..4c469dc6f 100644 --- a/etl/src/types/event.rs +++ b/etl/src/types/event.rs @@ -1,4 +1,4 @@ -use etl_postgres::types::{ColumnSchema, TableId, TableSchema}; +use etl_postgres::types::{ColumnSchema, TableId}; use std::fmt; use tokio_postgres::types::PgLsn; @@ -44,7 +44,7 @@ pub struct CommitEvent { pub enum RelationChange { AddColumn(ColumnSchema), DropColumn(ColumnSchema), - AlterColumn(ColumnSchema), + AlterColumn(ColumnSchema, ColumnSchema), } /// Table schema definition event from Postgres logical replication. @@ -184,7 +184,7 @@ impl Event { Event::Insert(insert_event) => insert_event.table_id == *table_id, Event::Update(update_event) => update_event.table_id == *table_id, Event::Delete(delete_event) => delete_event.table_id == *table_id, - Event::Relation(relation_event) => relation_event.table_schema.id == *table_id, + Event::Relation(relation_event) => relation_event.table_id == *table_id, Event::Truncate(event) => { let Some(_) = event.rel_ids.iter().find(|&&id| table_id.0 == id) else { return false; diff --git a/etl/tests/pipeline.rs b/etl/tests/pipeline.rs index ee6230460..ec5f38370 100644 --- a/etl/tests/pipeline.rs +++ b/etl/tests/pipeline.rs @@ -489,6 +489,140 @@ async fn table_copy_replicates_existing_data() { ); } +#[tokio::test(flavor = "multi_thread")] +async fn table_schema_changes_are_handled_correctly() { + init_test_tracing(); + let mut database = spawn_source_database().await; + let database_schema = setup_test_database_schema(&database, TableSelection::UsersOnly).await; + + // Insert initial users data. + insert_users_data(&mut database, &database_schema.users_schema().name, 1..=1).await; + + let store = NotifyingStore::new(); + let destination = TestDestinationWrapper::wrap(MemoryDestination::new()); + + // Start pipeline from scratch. + let pipeline_id: PipelineId = random(); + let mut pipeline = create_pipeline( + &database.config, + pipeline_id, + database_schema.publication_name(), + store.clone(), + destination.clone(), + ); + + // Register notifications for table copy completion. + let users_state_notify = store + .notify_on_table_state_type( + database_schema.users_schema().id, + TableReplicationPhaseType::SyncDone, + ) + .await; + + pipeline.start().await.unwrap(); + + users_state_notify.notified().await; + + // Check the initial schema. + let table_schemas = store.get_table_schemas().await; + assert_eq!(table_schemas.len(), 1); + assert!(table_schemas.contains_key(&database_schema.users_schema().id)); + + // Check the initial data. + let table_rows = destination.get_table_rows().await; + let users_table_rows = table_rows.get(&database_schema.users_schema().id).unwrap(); + assert_eq!(users_table_rows.len(), 1); + + // We perform schema changes. + database + .alter_table( + test_table_name("users"), + &[ + TableModification::AddColumn { + name: "year", + params: "integer", + }, + TableModification::RenameColumn { + name: "age", + new_name: "new_age", + }, + ], + ) + .await + .unwrap(); + + // Register notifications for the insert. + let insert_event_notify = destination + .wait_for_events_count(vec![(EventType::Insert, 1)]) + .await; + + // We insert data. + database + .insert_values( + database_schema.users_schema().name.clone(), + &["name", "new_age", "year"], + &[&"user_2", &(2i32), &(2025i32)], + ) + .await + .expect("Failed to insert users"); + + insert_event_notify.notified().await; + + // Check the updated schema. + let table_schemas = store.get_table_schemas().await; + assert_eq!(table_schemas.len(), 1); + assert!(table_schemas.contains_key(&database_schema.users_schema().id)); + + // Check the updated data. + let table_rows = destination.get_table_rows().await; + let users_table_rows = table_rows.get(&database_schema.users_schema().id).unwrap(); + assert_eq!(users_table_rows.len(), 1); + // + // // We perform schema changes. + // database + // .alter_table( + // test_table_name("users"), + // &[ + // TableModification::DropColumn { + // name: "year", + // }, + // TableModification::AlterColumn { + // name: "new_age", + // params: "type double precision using new_age::double precision", + // }, + // ], + // ) + // .await + // .unwrap(); + // + // // Register notifications for the insert. + // let insert_event_notify = destination.wait_for_events_count(vec![(EventType::Insert, 2)]).await; + // + // // We insert data. + // database + // .insert_values( + // database_schema.users_schema().name.clone(), + // &["name", "new_age", "year"], + // &[&"user_3", &(2i32), &(2025i32)], + // ) + // .await + // .expect("Failed to insert users"); + // + // insert_event_notify.notified().await; + // + // // Check the updated schema. + // let table_schemas = store.get_table_schemas().await; + // assert_eq!(table_schemas.len(), 1); + // assert!(table_schemas.contains_key(&database_schema.users_schema().id)); + // + // // Check the updated data. + // let table_rows = destination.get_table_rows().await; + // let users_table_rows = table_rows.get(&database_schema.users_schema().id).unwrap(); + // assert_eq!(users_table_rows.len(), 1); + // + // pipeline.shutdown_and_wait().await.unwrap(); +} + #[tokio::test(flavor = "multi_thread")] async fn table_copy_and_sync_streams_new_data() { init_test_tracing(); @@ -889,7 +1023,7 @@ async fn table_processing_with_schema_change_errors_table() { database_schema.orders_schema().name.clone(), &[TableModification::AddColumn { name: "date", - data_type: "integer", + params: "integer", }], ) .await diff --git a/etl/tests/replica_identity.rs b/etl/tests/replica_identity.rs index 29b8f8774..383eb9f21 100644 --- a/etl/tests/replica_identity.rs +++ b/etl/tests/replica_identity.rs @@ -115,7 +115,7 @@ async fn update_non_toast_values_with_default_replica_identity() { table_name.clone(), &[TableModification::AlterColumn { name: "large_text", - alteration: "set storage external", + params: "set storage external", }], ) .await @@ -237,7 +237,7 @@ async fn update_non_toast_values_with_full_replica_identity() { table_name.clone(), &[TableModification::AlterColumn { name: "large_text", - alteration: "set storage external", + params: "set storage external", }], ) .await @@ -370,7 +370,7 @@ async fn update_toast_values_with_default_replica_identity() { table_name.clone(), &[TableModification::AlterColumn { name: "large_text", - alteration: "set storage external", + params: "set storage external", }], ) .await @@ -493,7 +493,7 @@ async fn update_non_toast_values_with_none_replica_identity() { table_name.clone(), &[TableModification::AlterColumn { name: "large_text", - alteration: "set storage external", + params: "set storage external", }], ) .await @@ -630,7 +630,7 @@ async fn update_non_toast_values_with_unique_index_replica_identity() { table_name.clone(), &[TableModification::AlterColumn { name: "large_text", - alteration: "set storage external", + params: "set storage external", }], ) .await diff --git a/etl/tests/replication.rs b/etl/tests/replication.rs index 67092271c..9602a6cc5 100644 --- a/etl/tests/replication.rs +++ b/etl/tests/replication.rs @@ -261,7 +261,7 @@ async fn test_table_schema_copy_across_multiple_connections() { test_table_name("table_1"), &[TableModification::AddColumn { name: "year", - data_type: "integer", + params: "integer", }], ) .await From ea6087ffc3f1dcf23cf22b8f6b88ec9b55716450 Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Tue, 30 Sep 2025 11:05:38 -0700 Subject: [PATCH 04/45] Improve --- etl-destinations/src/bigquery/client.rs | 59 +++--- etl-destinations/tests/iceberg_pipeline.rs | 201 +++++++-------------- etl-postgres/src/replication/schema.rs | 2 - etl-postgres/src/types/schema.rs | 23 +-- etl/src/conversions/event.rs | 23 +-- etl/src/conversions/table_row.rs | 34 ++-- etl/tests/postgres_store.rs | 17 +- 7 files changed, 137 insertions(+), 222 deletions(-) diff --git a/etl-destinations/src/bigquery/client.rs b/etl-destinations/src/bigquery/client.rs index b88b286a3..975acc744 100644 --- a/etl-destinations/src/bigquery/client.rs +++ b/etl-destinations/src/bigquery/client.rs @@ -773,16 +773,15 @@ mod tests { #[test] fn test_column_spec() { - let column_schema = ColumnSchema::new("test_col".to_string(), Type::TEXT, -1, true, false); + let column_schema = ColumnSchema::new("test_col".to_string(), Type::TEXT, -1, true); let spec = BigQueryClient::column_spec(&column_schema); assert_eq!(spec, "`test_col` string"); - let not_null_column = ColumnSchema::new("id".to_string(), Type::INT4, -1, false, true); + let not_null_column = ColumnSchema::new("id".to_string(), Type::INT4, -1, false); let not_null_spec = BigQueryClient::column_spec(¬_null_column); assert_eq!(not_null_spec, "`id` int64 not null"); - let array_column = - ColumnSchema::new("tags".to_string(), Type::TEXT_ARRAY, -1, false, false); + let array_column = ColumnSchema::new("tags".to_string(), Type::TEXT_ARRAY, -1, false); let array_spec = BigQueryClient::column_spec(&array_column); assert_eq!(array_spec, "`tags` array"); } @@ -790,16 +789,16 @@ mod tests { #[test] fn test_add_primary_key_clause() { let columns_with_pk = vec![ - ColumnSchema::new("id".to_string(), Type::INT4, -1, false, true), - ColumnSchema::new("name".to_string(), Type::TEXT, -1, true, false), + ColumnSchema::new("id".to_string(), Type::INT4, -1, false), + ColumnSchema::new("name".to_string(), Type::TEXT, -1, true), ]; let pk_clause = BigQueryClient::add_primary_key_clause(&columns_with_pk); assert_eq!(pk_clause, ", primary key (`id`) not enforced"); let columns_with_composite_pk = vec![ - ColumnSchema::new("tenant_id".to_string(), Type::INT4, -1, false, true), - ColumnSchema::new("id".to_string(), Type::INT4, -1, false, true), - ColumnSchema::new("name".to_string(), Type::TEXT, -1, true, false), + ColumnSchema::new("tenant_id".to_string(), Type::INT4, -1, false), + ColumnSchema::new("id".to_string(), Type::INT4, -1, false), + ColumnSchema::new("name".to_string(), Type::TEXT, -1, true), ]; let composite_pk_clause = BigQueryClient::add_primary_key_clause(&columns_with_composite_pk); @@ -809,8 +808,8 @@ mod tests { ); let columns_no_pk = vec![ - ColumnSchema::new("name".to_string(), Type::TEXT, -1, true, false), - ColumnSchema::new("age".to_string(), Type::INT4, -1, true, false), + ColumnSchema::new("name".to_string(), Type::TEXT, -1, true), + ColumnSchema::new("age".to_string(), Type::INT4, -1, true), ]; let no_pk_clause = BigQueryClient::add_primary_key_clause(&columns_no_pk); assert_eq!(no_pk_clause, ""); @@ -819,9 +818,9 @@ mod tests { #[test] fn test_create_columns_spec() { let columns = vec![ - ColumnSchema::new("id".to_string(), Type::INT4, -1, false, true), - ColumnSchema::new("name".to_string(), Type::TEXT, -1, true, false), - ColumnSchema::new("active".to_string(), Type::BOOL, -1, false, false), + ColumnSchema::new("id".to_string(), Type::INT4, -1, false), + ColumnSchema::new("name".to_string(), Type::TEXT, -1, true), + ColumnSchema::new("active".to_string(), Type::BOOL, -1, false), ]; let spec = BigQueryClient::create_columns_spec(&columns); assert_eq!( @@ -839,10 +838,10 @@ mod tests { #[test] fn test_column_schemas_to_table_descriptor() { let columns = vec![ - ColumnSchema::new("id".to_string(), Type::INT4, -1, false, true), - ColumnSchema::new("name".to_string(), Type::TEXT, -1, true, false), - ColumnSchema::new("active".to_string(), Type::BOOL, -1, false, false), - ColumnSchema::new("tags".to_string(), Type::TEXT_ARRAY, -1, false, false), + ColumnSchema::new("id".to_string(), Type::INT4, -1, false), + ColumnSchema::new("name".to_string(), Type::TEXT, -1, true), + ColumnSchema::new("active".to_string(), Type::BOOL, -1, false), + ColumnSchema::new("tags".to_string(), Type::TEXT_ARRAY, -1, false), ]; let descriptor = BigQueryClient::column_schemas_to_table_descriptor(&columns, true); @@ -922,12 +921,12 @@ mod tests { #[test] fn test_column_schemas_to_table_descriptor_complex_types() { let columns = vec![ - ColumnSchema::new("uuid_col".to_string(), Type::UUID, -1, true, false), - ColumnSchema::new("json_col".to_string(), Type::JSON, -1, true, false), - ColumnSchema::new("bytea_col".to_string(), Type::BYTEA, -1, true, false), - ColumnSchema::new("numeric_col".to_string(), Type::NUMERIC, -1, true, false), - ColumnSchema::new("date_col".to_string(), Type::DATE, -1, true, false), - ColumnSchema::new("time_col".to_string(), Type::TIME, -1, true, false), + ColumnSchema::new("uuid_col".to_string(), Type::UUID, -1, true), + ColumnSchema::new("json_col".to_string(), Type::JSON, -1, true), + ColumnSchema::new("bytea_col".to_string(), Type::BYTEA, -1, true), + ColumnSchema::new("numeric_col".to_string(), Type::NUMERIC, -1, true), + ColumnSchema::new("date_col".to_string(), Type::DATE, -1, true), + ColumnSchema::new("time_col".to_string(), Type::TIME, -1, true), ]; let descriptor = BigQueryClient::column_schemas_to_table_descriptor(&columns, true); @@ -979,8 +978,8 @@ mod tests { let table_id = "test_table"; let columns = vec![ - ColumnSchema::new("id".to_string(), Type::INT4, -1, false, true), - ColumnSchema::new("name".to_string(), Type::TEXT, -1, true, false), + ColumnSchema::new("id".to_string(), Type::INT4, -1, false), + ColumnSchema::new("name".to_string(), Type::TEXT, -1, true), ]; // Simulate the query generation logic @@ -999,13 +998,7 @@ mod tests { let table_id = "test_table"; let max_staleness_mins = 15; - let columns = vec![ColumnSchema::new( - "id".to_string(), - Type::INT4, - -1, - false, - true, - )]; + let columns = vec![ColumnSchema::new("id".to_string(), Type::INT4, -1, false)]; // Simulate the query generation logic with staleness let full_table_name = format!("`{project_id}.{dataset_id}.{table_id}`"); diff --git a/etl-destinations/tests/iceberg_pipeline.rs b/etl-destinations/tests/iceberg_pipeline.rs index f552677ca..8e574227b 100644 --- a/etl-destinations/tests/iceberg_pipeline.rs +++ b/etl-destinations/tests/iceberg_pipeline.rs @@ -126,197 +126,180 @@ async fn create_table_if_missing() { let table_name_struct = TableName::new("test_schema".to_string(), table_name.clone()); let columns = vec![ // Primary key - ColumnSchema::new("id".to_string(), Type::INT4, -1, false, true), + ColumnSchema::new("id".to_string(), Type::INT4, -1, false), // Boolean types - ColumnSchema::new("bool_col".to_string(), Type::BOOL, -1, true, false), + ColumnSchema::new("bool_col".to_string(), Type::BOOL, -1, true), // String types - ColumnSchema::new("char_col".to_string(), Type::CHAR, -1, true, false), - ColumnSchema::new("bpchar_col".to_string(), Type::BPCHAR, -1, true, false), - ColumnSchema::new("varchar_col".to_string(), Type::VARCHAR, -1, true, false), - ColumnSchema::new("name_col".to_string(), Type::NAME, -1, true, false), - ColumnSchema::new("text_col".to_string(), Type::TEXT, -1, true, false), + ColumnSchema::new("char_col".to_string(), Type::CHAR, -1, true), + ColumnSchema::new("bpchar_col".to_string(), Type::BPCHAR, -1, true), + ColumnSchema::new("varchar_col".to_string(), Type::VARCHAR, -1, true), + ColumnSchema::new("name_col".to_string(), Type::NAME, -1, true), + ColumnSchema::new("text_col".to_string(), Type::TEXT, -1, true), // Integer types - ColumnSchema::new("int2_col".to_string(), Type::INT2, -1, true, false), - ColumnSchema::new("int4_col".to_string(), Type::INT4, -1, true, false), - ColumnSchema::new("int8_col".to_string(), Type::INT8, -1, true, false), + ColumnSchema::new("int2_col".to_string(), Type::INT2, -1, true), + ColumnSchema::new("int4_col".to_string(), Type::INT4, -1, true), + ColumnSchema::new("int8_col".to_string(), Type::INT8, -1, true), // Float types - ColumnSchema::new("float4_col".to_string(), Type::FLOAT4, -1, true, false), - ColumnSchema::new("float8_col".to_string(), Type::FLOAT8, -1, true, false), + ColumnSchema::new("float4_col".to_string(), Type::FLOAT4, -1, true), + ColumnSchema::new("float8_col".to_string(), Type::FLOAT8, -1, true), // Numeric type - ColumnSchema::new("numeric_col".to_string(), Type::NUMERIC, -1, true, false), + ColumnSchema::new("numeric_col".to_string(), Type::NUMERIC, -1, true), // Date/Time types - ColumnSchema::new("date_col".to_string(), Type::DATE, -1, true, false), - ColumnSchema::new("time_col".to_string(), Type::TIME, -1, true, false), + ColumnSchema::new("date_col".to_string(), Type::DATE, -1, true), + ColumnSchema::new("time_col".to_string(), Type::TIME, -1, true), ColumnSchema::new( "timestamp_col".to_string(), Type::TIMESTAMP, -1, true, - false, + ), ColumnSchema::new( "timestamptz_col".to_string(), Type::TIMESTAMPTZ, -1, true, - false, + ), // UUID type - ColumnSchema::new("uuid_col".to_string(), Type::UUID, -1, true, false), + ColumnSchema::new("uuid_col".to_string(), Type::UUID, -1, true), // JSON types - ColumnSchema::new("json_col".to_string(), Type::JSON, -1, true, false), - ColumnSchema::new("jsonb_col".to_string(), Type::JSONB, -1, true, false), + ColumnSchema::new("json_col".to_string(), Type::JSON, -1, true), + ColumnSchema::new("jsonb_col".to_string(), Type::JSONB, -1, true), // OID type - ColumnSchema::new("oid_col".to_string(), Type::OID, -1, true, false), + ColumnSchema::new("oid_col".to_string(), Type::OID, -1, true), // Binary type - ColumnSchema::new("bytea_col".to_string(), Type::BYTEA, -1, true, false), + ColumnSchema::new("bytea_col".to_string(), Type::BYTEA, -1, true), // Array types ColumnSchema::new( "bool_array_col".to_string(), Type::BOOL_ARRAY, -1, true, - false, + ), ColumnSchema::new( "char_array_col".to_string(), Type::CHAR_ARRAY, -1, true, - false, + ), ColumnSchema::new( "bpchar_array_col".to_string(), Type::BPCHAR_ARRAY, -1, true, - false, + ), ColumnSchema::new( "varchar_array_col".to_string(), Type::VARCHAR_ARRAY, -1, true, - false, + ), ColumnSchema::new( "name_array_col".to_string(), Type::NAME_ARRAY, -1, true, - false, ), ColumnSchema::new( "text_array_col".to_string(), Type::TEXT_ARRAY, -1, true, - false, ), ColumnSchema::new( "int2_array_col".to_string(), Type::INT2_ARRAY, -1, true, - false, ), ColumnSchema::new( "int4_array_col".to_string(), Type::INT4_ARRAY, -1, true, - false, ), ColumnSchema::new( "int8_array_col".to_string(), Type::INT8_ARRAY, -1, true, - false, ), ColumnSchema::new( "float4_array_col".to_string(), Type::FLOAT4_ARRAY, -1, true, - false, ), ColumnSchema::new( "float8_array_col".to_string(), Type::FLOAT8_ARRAY, -1, true, - false, ), ColumnSchema::new( "numeric_array_col".to_string(), Type::NUMERIC_ARRAY, -1, true, - false, ), ColumnSchema::new( "date_array_col".to_string(), Type::DATE_ARRAY, -1, true, - false, ), ColumnSchema::new( "time_array_col".to_string(), Type::TIME_ARRAY, -1, true, - false, ), ColumnSchema::new( "timestamp_array_col".to_string(), Type::TIMESTAMP_ARRAY, -1, true, - false, ), ColumnSchema::new( "timestamptz_array_col".to_string(), Type::TIMESTAMPTZ_ARRAY, -1, true, - false, ), ColumnSchema::new( "uuid_array_col".to_string(), Type::UUID_ARRAY, -1, true, - false, ), ColumnSchema::new( "json_array_col".to_string(), Type::JSON_ARRAY, -1, true, - false, ), ColumnSchema::new( "jsonb_array_col".to_string(), Type::JSONB_ARRAY, -1, true, - false, ), ColumnSchema::new( "oid_array_col".to_string(), Type::OID_ARRAY, -1, true, - false, ), ColumnSchema::new( "bytea_array_col".to_string(), Type::BYTEA_ARRAY, -1, true, - false, ), ]; let table_schema = TableSchema::new(table_id, table_name_struct, columns); @@ -388,50 +371,50 @@ async fn insert_nullable_scalars() { let table_name_struct = TableName::new("test_schema".to_string(), table_name.clone()); let columns = vec![ // Primary key - ColumnSchema::new("id".to_string(), Type::INT4, -1, false, true), + ColumnSchema::new("id".to_string(), Type::INT4, -1, false), // Boolean types - ColumnSchema::new("bool_col".to_string(), Type::BOOL, -1, true, false), + ColumnSchema::new("bool_col".to_string(), Type::BOOL, -1, true), // String types - ColumnSchema::new("char_col".to_string(), Type::CHAR, -1, true, false), - ColumnSchema::new("bpchar_col".to_string(), Type::BPCHAR, -1, true, false), - ColumnSchema::new("varchar_col".to_string(), Type::VARCHAR, -1, true, false), - ColumnSchema::new("name_col".to_string(), Type::NAME, -1, true, false), - ColumnSchema::new("text_col".to_string(), Type::TEXT, -1, true, false), + ColumnSchema::new("char_col".to_string(), Type::CHAR, -1, true), + ColumnSchema::new("bpchar_col".to_string(), Type::BPCHAR, -1, true), + ColumnSchema::new("varchar_col".to_string(), Type::VARCHAR, -1, true), + ColumnSchema::new("name_col".to_string(), Type::NAME, -1, true), + ColumnSchema::new("text_col".to_string(), Type::TEXT, -1, true), // Integer types - ColumnSchema::new("int2_col".to_string(), Type::INT2, -1, true, false), - ColumnSchema::new("int4_col".to_string(), Type::INT4, -1, true, false), - ColumnSchema::new("int8_col".to_string(), Type::INT8, -1, true, false), + ColumnSchema::new("int2_col".to_string(), Type::INT2, -1, true), + ColumnSchema::new("int4_col".to_string(), Type::INT4, -1, true), + ColumnSchema::new("int8_col".to_string(), Type::INT8, -1, true), // Float types - ColumnSchema::new("float4_col".to_string(), Type::FLOAT4, -1, true, false), - ColumnSchema::new("float8_col".to_string(), Type::FLOAT8, -1, true, false), + ColumnSchema::new("float4_col".to_string(), Type::FLOAT4, -1, true), + ColumnSchema::new("float8_col".to_string(), Type::FLOAT8, -1, true), // Numeric type - ColumnSchema::new("numeric_col".to_string(), Type::NUMERIC, -1, true, false), + ColumnSchema::new("numeric_col".to_string(), Type::NUMERIC, -1, true), // Date/Time types - ColumnSchema::new("date_col".to_string(), Type::DATE, -1, true, false), - ColumnSchema::new("time_col".to_string(), Type::TIME, -1, true, false), + ColumnSchema::new("date_col".to_string(), Type::DATE, -1, true), + ColumnSchema::new("time_col".to_string(), Type::TIME, -1, true), ColumnSchema::new( "timestamp_col".to_string(), Type::TIMESTAMP, -1, true, - false, + ), ColumnSchema::new( "timestamptz_col".to_string(), Type::TIMESTAMPTZ, -1, true, - false, + ), // UUID type - ColumnSchema::new("uuid_col".to_string(), Type::UUID, -1, true, false), + ColumnSchema::new("uuid_col".to_string(), Type::UUID, -1, true), // JSON types - ColumnSchema::new("json_col".to_string(), Type::JSON, -1, true, false), - ColumnSchema::new("jsonb_col".to_string(), Type::JSONB, -1, true, false), + ColumnSchema::new("json_col".to_string(), Type::JSON, -1, true), + ColumnSchema::new("jsonb_col".to_string(), Type::JSONB, -1, true), // OID type - ColumnSchema::new("oid_col".to_string(), Type::OID, -1, true, false), + ColumnSchema::new("oid_col".to_string(), Type::OID, -1, true), // Binary type - ColumnSchema::new("bytea_col".to_string(), Type::BYTEA, -1, true, false), + ColumnSchema::new("bytea_col".to_string(), Type::BYTEA, -1, true), ]; let table_schema = TableSchema::new(table_id, table_name_struct, columns); @@ -551,50 +534,48 @@ async fn insert_non_nullable_scalars() { let table_name_struct = TableName::new("test_schema".to_string(), table_name.clone()); let columns = vec![ // Primary key - ColumnSchema::new("id".to_string(), Type::INT4, -1, false, true), + ColumnSchema::new("id".to_string(), Type::INT4, -1, false), // Boolean types - ColumnSchema::new("bool_col".to_string(), Type::BOOL, -1, false, false), + ColumnSchema::new("bool_col".to_string(), Type::BOOL, -1, false), // String types - ColumnSchema::new("char_col".to_string(), Type::CHAR, -1, false, false), - ColumnSchema::new("bpchar_col".to_string(), Type::BPCHAR, -1, false, false), - ColumnSchema::new("varchar_col".to_string(), Type::VARCHAR, -1, false, false), - ColumnSchema::new("name_col".to_string(), Type::NAME, -1, false, false), - ColumnSchema::new("text_col".to_string(), Type::TEXT, -1, false, false), + ColumnSchema::new("char_col".to_string(), Type::CHAR, -1, false), + ColumnSchema::new("bpchar_col".to_string(), Type::BPCHAR, -1, false), + ColumnSchema::new("varchar_col".to_string(), Type::VARCHAR, -1, false), + ColumnSchema::new("name_col".to_string(), Type::NAME, -1, false), + ColumnSchema::new("text_col".to_string(), Type::TEXT, -1, false), // Integer types - ColumnSchema::new("int2_col".to_string(), Type::INT2, -1, false, false), - ColumnSchema::new("int4_col".to_string(), Type::INT4, -1, false, false), - ColumnSchema::new("int8_col".to_string(), Type::INT8, -1, false, false), + ColumnSchema::new("int2_col".to_string(), Type::INT2, -1, false), + ColumnSchema::new("int4_col".to_string(), Type::INT4, -1, false), + ColumnSchema::new("int8_col".to_string(), Type::INT8, -1, false), // Float types - ColumnSchema::new("float4_col".to_string(), Type::FLOAT4, -1, false, false), - ColumnSchema::new("float8_col".to_string(), Type::FLOAT8, -1, false, false), + ColumnSchema::new("float4_col".to_string(), Type::FLOAT4, -1, false,), + ColumnSchema::new("float8_col".to_string(), Type::FLOAT8, -1, false), // Numeric type - ColumnSchema::new("numeric_col".to_string(), Type::NUMERIC, -1, false, false), + ColumnSchema::new("numeric_col".to_string(), Type::NUMERIC, -1, false), // Date/Time types - ColumnSchema::new("date_col".to_string(), Type::DATE, -1, false, false), - ColumnSchema::new("time_col".to_string(), Type::TIME, -1, false, false), + ColumnSchema::new("date_col".to_string(), Type::DATE, -1, false), + ColumnSchema::new("time_col".to_string(), Type::TIME, -1, false), ColumnSchema::new( "timestamp_col".to_string(), Type::TIMESTAMP, -1, false, - false, ), ColumnSchema::new( "timestamptz_col".to_string(), Type::TIMESTAMPTZ, -1, false, - false, ), // UUID type - ColumnSchema::new("uuid_col".to_string(), Type::UUID, -1, false, false), + ColumnSchema::new("uuid_col".to_string(), Type::UUID, -1, false), // JSON types - ColumnSchema::new("json_col".to_string(), Type::JSON, -1, false, false), - ColumnSchema::new("jsonb_col".to_string(), Type::JSONB, -1, false, false), + ColumnSchema::new("json_col".to_string(), Type::JSON, -1, false), + ColumnSchema::new("jsonb_col".to_string(), Type::JSONB, -1, false), // OID type - ColumnSchema::new("oid_col".to_string(), Type::OID, -1, false, false), + ColumnSchema::new("oid_col".to_string(), Type::OID, -1, false), // Binary type - ColumnSchema::new("bytea_col".to_string(), Type::BYTEA, -1, false, false), + ColumnSchema::new("bytea_col".to_string(), Type::BYTEA, -1, false), ]; let table_schema = TableSchema::new(table_id, table_name_struct, columns); @@ -688,14 +669,13 @@ async fn insert_nullable_array() { let table_name_struct = TableName::new("test_schema".to_string(), table_name.clone()); let columns = vec![ // Primary key - ColumnSchema::new("id".to_string(), Type::INT4, -1, false, true), + ColumnSchema::new("id".to_string(), Type::INT4, -1, false), // Boolean array type ColumnSchema::new( "bool_array_col".to_string(), Type::BOOL_ARRAY, -1, true, - false, ), // String array types ColumnSchema::new( @@ -703,35 +683,30 @@ async fn insert_nullable_array() { Type::CHAR_ARRAY, -1, true, - false, ), ColumnSchema::new( "bpchar_array_col".to_string(), Type::BPCHAR_ARRAY, -1, true, - false, ), ColumnSchema::new( "varchar_array_col".to_string(), Type::VARCHAR_ARRAY, -1, true, - false, ), ColumnSchema::new( "name_array_col".to_string(), Type::NAME_ARRAY, -1, true, - false, ), ColumnSchema::new( "text_array_col".to_string(), Type::TEXT_ARRAY, -1, true, - false, ), // Integer array types ColumnSchema::new( @@ -739,21 +714,18 @@ async fn insert_nullable_array() { Type::INT2_ARRAY, -1, true, - false, ), ColumnSchema::new( "int4_array_col".to_string(), Type::INT4_ARRAY, -1, true, - false, ), ColumnSchema::new( "int8_array_col".to_string(), Type::INT8_ARRAY, -1, true, - false, ), // Float array types ColumnSchema::new( @@ -761,14 +733,12 @@ async fn insert_nullable_array() { Type::FLOAT4_ARRAY, -1, true, - false, ), ColumnSchema::new( "float8_array_col".to_string(), Type::FLOAT8_ARRAY, -1, true, - false, ), // Numeric array type ColumnSchema::new( @@ -776,7 +746,6 @@ async fn insert_nullable_array() { Type::NUMERIC_ARRAY, -1, true, - false, ), // Date/Time array types ColumnSchema::new( @@ -784,28 +753,24 @@ async fn insert_nullable_array() { Type::DATE_ARRAY, -1, true, - false, ), ColumnSchema::new( "time_array_col".to_string(), Type::TIME_ARRAY, -1, true, - false, ), ColumnSchema::new( "timestamp_array_col".to_string(), Type::TIMESTAMP_ARRAY, -1, true, - false, ), ColumnSchema::new( "timestamptz_array_col".to_string(), Type::TIMESTAMPTZ_ARRAY, -1, true, - false, ), // UUID array type ColumnSchema::new( @@ -813,7 +778,6 @@ async fn insert_nullable_array() { Type::UUID_ARRAY, -1, true, - false, ), // JSON array types ColumnSchema::new( @@ -821,14 +785,12 @@ async fn insert_nullable_array() { Type::JSON_ARRAY, -1, true, - false, ), ColumnSchema::new( "jsonb_array_col".to_string(), Type::JSONB_ARRAY, -1, true, - false, ), // OID array type ColumnSchema::new( @@ -836,7 +798,6 @@ async fn insert_nullable_array() { Type::OID_ARRAY, -1, true, - false, ), // Binary array type ColumnSchema::new( @@ -844,7 +805,6 @@ async fn insert_nullable_array() { Type::BYTEA_ARRAY, -1, true, - false, ), ]; let table_schema = TableSchema::new(table_id, table_name_struct, columns); @@ -1053,14 +1013,13 @@ async fn insert_non_nullable_array() { let table_name_struct = TableName::new("test_schema".to_string(), table_name.clone()); let columns = vec![ // Primary key - ColumnSchema::new("id".to_string(), Type::INT4, -1, false, true), + ColumnSchema::new("id".to_string(), Type::INT4, -1, false), // Boolean array type ColumnSchema::new( "bool_array_col".to_string(), Type::BOOL_ARRAY, -1, false, - false, ), // String array types ColumnSchema::new( @@ -1068,35 +1027,30 @@ async fn insert_non_nullable_array() { Type::CHAR_ARRAY, -1, false, - false, ), ColumnSchema::new( "bpchar_array_col".to_string(), Type::BPCHAR_ARRAY, -1, false, - false, ), ColumnSchema::new( "varchar_array_col".to_string(), Type::VARCHAR_ARRAY, -1, false, - false, ), ColumnSchema::new( "name_array_col".to_string(), Type::NAME_ARRAY, -1, false, - false, ), ColumnSchema::new( "text_array_col".to_string(), Type::TEXT_ARRAY, -1, false, - false, ), // Integer array types ColumnSchema::new( @@ -1104,21 +1058,18 @@ async fn insert_non_nullable_array() { Type::INT2_ARRAY, -1, false, - false, ), ColumnSchema::new( "int4_array_col".to_string(), Type::INT4_ARRAY, -1, false, - false, ), ColumnSchema::new( "int8_array_col".to_string(), Type::INT8_ARRAY, -1, false, - false, ), // Float array types ColumnSchema::new( @@ -1126,14 +1077,12 @@ async fn insert_non_nullable_array() { Type::FLOAT4_ARRAY, -1, false, - false, ), ColumnSchema::new( "float8_array_col".to_string(), Type::FLOAT8_ARRAY, -1, false, - false, ), // Numeric array type ColumnSchema::new( @@ -1141,7 +1090,6 @@ async fn insert_non_nullable_array() { Type::NUMERIC_ARRAY, -1, false, - false, ), // Date/Time array types ColumnSchema::new( @@ -1149,28 +1097,24 @@ async fn insert_non_nullable_array() { Type::DATE_ARRAY, -1, false, - false, ), ColumnSchema::new( "time_array_col".to_string(), Type::TIME_ARRAY, -1, false, - false, ), ColumnSchema::new( "timestamp_array_col".to_string(), Type::TIMESTAMP_ARRAY, -1, false, - false, ), ColumnSchema::new( "timestamptz_array_col".to_string(), Type::TIMESTAMPTZ_ARRAY, -1, false, - false, ), // UUID array type ColumnSchema::new( @@ -1178,7 +1122,6 @@ async fn insert_non_nullable_array() { Type::UUID_ARRAY, -1, false, - false, ), // JSON array types ColumnSchema::new( @@ -1186,14 +1129,12 @@ async fn insert_non_nullable_array() { Type::JSON_ARRAY, -1, false, - false, ), ColumnSchema::new( "jsonb_array_col".to_string(), Type::JSONB_ARRAY, -1, false, - false, ), // OID array type ColumnSchema::new( @@ -1201,7 +1142,6 @@ async fn insert_non_nullable_array() { Type::OID_ARRAY, -1, false, - false, ), // Binary array type ColumnSchema::new( @@ -1209,7 +1149,6 @@ async fn insert_non_nullable_array() { Type::BYTEA_ARRAY, -1, false, - false, ), ]; let table_schema = TableSchema::new(table_id, table_name_struct, columns); diff --git a/etl-postgres/src/replication/schema.rs b/etl-postgres/src/replication/schema.rs index cd540d653..336ff166c 100644 --- a/etl-postgres/src/replication/schema.rs +++ b/etl-postgres/src/replication/schema.rs @@ -303,14 +303,12 @@ fn parse_column_schema(row: &PgRow) -> ColumnSchema { let column_name: String = row.get("column_name"); let column_type: String = row.get("column_type"); let type_modifier: i32 = row.get("type_modifier"); - let nullable: bool = row.get("nullable"); let primary_key: bool = row.get("primary_key"); ColumnSchema::new( column_name, string_to_postgres_type(&column_type), type_modifier, - nullable, primary_key, ) } diff --git a/etl-postgres/src/types/schema.rs b/etl-postgres/src/types/schema.rs index d2f6eeb27..2b64baa65 100644 --- a/etl-postgres/src/types/schema.rs +++ b/etl-postgres/src/types/schema.rs @@ -54,31 +54,28 @@ type TypeModifier = i32; /// type modifier, nullability, and whether it's part of the primary key. #[derive(Debug, Clone, Eq, PartialEq)] pub struct ColumnSchema { - /// The name of the column + /// The name of the column. pub name: String, - /// The Postgres data type of the column + /// The Postgres data type of the column. pub typ: Type, - /// Type-specific modifier value (e.g., length for varchar) + /// Type-specific modifier value (e.g., length for varchar). pub modifier: TypeModifier, - /// Whether the column can contain NULL values + /// Whether the column can contain NULL values. pub nullable: bool, - /// Whether the column is part of the table's primary key + /// Whether the column is part of the table's primary key. pub primary: bool, } impl ColumnSchema { - pub fn new( - name: String, - typ: Type, - modifier: TypeModifier, - nullable: bool, - primary: bool, - ) -> ColumnSchema { + pub fn new(name: String, typ: Type, modifier: TypeModifier, primary: bool) -> ColumnSchema { Self { name, typ, modifier, - nullable, + // For now, we assume all columns are nullable. The rationale behind this is that we + // do not have access to nullability information on relation messages, so to support schema + // evolution, we assume all columns are nullable. + nullable: true, primary, } } diff --git a/etl/src/conversions/event.rs b/etl/src/conversions/event.rs index 5d7711021..1861a88bf 100644 --- a/etl/src/conversions/event.rs +++ b/etl/src/conversions/event.rs @@ -108,12 +108,17 @@ where let latest_column_schema = latest_column_schemas.take(&column_schema); match latest_column_schema { Some(latest_column_schema) => { - // If we find a column with the same name, we assume it was changed. The only changes - // that we detect are changes to the column but with preserved name. - changes.push(RelationChange::AlterColumn( - column_schema.into_inner(), - latest_column_schema.into_inner(), - )); + let column_schema = column_schema.into_inner(); + let latest_column_schema = latest_column_schema.into_inner(); + + if column_schema.name != latest_column_schema.name { + // If we find a column with the same name but different fields, we assume it was changed. The only changes + // that we detect are changes to the column but with preserved name. + changes.push(RelationChange::AlterColumn( + column_schema, + latest_column_schema, + )); + } } None => { // If we don't find the column in the latest schema, we assume it was dropped even @@ -131,7 +136,7 @@ where Ok(RelationEvent { start_lsn, commit_lsn, - changes: vec![], + changes, table_id, }) } @@ -309,10 +314,6 @@ fn build_indexed_column_schema(column: &protocol::Column) -> EtlResult Vec { vec![ - ColumnSchema::new("id".to_string(), Type::INT4, -1, false, true), - ColumnSchema::new("name".to_string(), Type::TEXT, -1, true, false), - ColumnSchema::new("active".to_string(), Type::BOOL, -1, false, false), + ColumnSchema::new("id".to_string(), Type::INT4, -1, false), + ColumnSchema::new("name".to_string(), Type::TEXT, -1, true), + ColumnSchema::new("active".to_string(), Type::BOOL, -1, false), ] } fn create_single_column_schema(name: &str, typ: Type) -> Vec { - vec![ColumnSchema::new(name.to_string(), typ, -1, false, false)] + vec![ColumnSchema::new(name.to_string(), typ, -1, false)] } #[test] @@ -227,10 +227,10 @@ mod tests { #[test] fn try_from_multiple_columns_different_types() { let schema = vec![ - ColumnSchema::new("int_col".to_string(), Type::INT4, -1, false, false), - ColumnSchema::new("float_col".to_string(), Type::FLOAT8, -1, false, false), - ColumnSchema::new("text_col".to_string(), Type::TEXT, -1, false, false), - ColumnSchema::new("bool_col".to_string(), Type::BOOL, -1, false, false), + ColumnSchema::new("int_col".to_string(), Type::INT4, -1, false), + ColumnSchema::new("float_col".to_string(), Type::FLOAT8, -1, false), + ColumnSchema::new("text_col".to_string(), Type::TEXT, -1, false), + ColumnSchema::new("bool_col".to_string(), Type::BOOL, -1, false), ]; let row_data = b"123\t3.15\tHello World\tt\n"; @@ -333,13 +333,7 @@ mod tests { let mut expected_row = String::new(); for i in 0..50 { - schema.push(ColumnSchema::new( - format!("col{i}"), - Type::INT4, - -1, - false, - false, - )); + schema.push(ColumnSchema::new(format!("col{i}"), Type::INT4, -1, false)); if i > 0 { expected_row.push('\t'); } @@ -369,8 +363,8 @@ mod tests { #[test] fn try_from_postgres_delimiter_escaping() { let schema = vec![ - ColumnSchema::new("col1".to_string(), Type::TEXT, -1, false, false), - ColumnSchema::new("col2".to_string(), Type::TEXT, -1, false, false), + ColumnSchema::new("col1".to_string(), Type::TEXT, -1, false), + ColumnSchema::new("col2".to_string(), Type::TEXT, -1, false), ]; // Postgres escapes tab characters in data with \\t @@ -387,9 +381,9 @@ mod tests { #[test] fn try_from_postgres_escape_at_field_boundaries() { let schema = vec![ - ColumnSchema::new("col1".to_string(), Type::TEXT, -1, false, false), - ColumnSchema::new("col2".to_string(), Type::TEXT, -1, false, false), - ColumnSchema::new("col3".to_string(), Type::TEXT, -1, false, false), + ColumnSchema::new("col1".to_string(), Type::TEXT, -1, false), + ColumnSchema::new("col2".to_string(), Type::TEXT, -1, false), + ColumnSchema::new("col3".to_string(), Type::TEXT, -1, false), ]; // Escapes at the beginning, middle, and end of fields diff --git a/etl/tests/postgres_store.rs b/etl/tests/postgres_store.rs index 1038c591a..afdcf4960 100644 --- a/etl/tests/postgres_store.rs +++ b/etl/tests/postgres_store.rs @@ -16,15 +16,9 @@ fn create_sample_table_schema() -> TableSchema { let table_id = TableId::new(12345); let table_name = TableName::new("public".to_string(), "test_table".to_string()); let columns = vec![ - ColumnSchema::new("id".to_string(), PgType::INT4, -1, false, true), - ColumnSchema::new("name".to_string(), PgType::TEXT, -1, true, false), - ColumnSchema::new( - "created_at".to_string(), - PgType::TIMESTAMPTZ, - -1, - false, - false, - ), + ColumnSchema::new("id".to_string(), PgType::INT4, -1, false), + ColumnSchema::new("name".to_string(), PgType::TEXT, -1, true), + ColumnSchema::new("created_at".to_string(), PgType::TIMESTAMPTZ, -1, false), ]; TableSchema::new(table_id, table_name, columns) @@ -34,8 +28,8 @@ fn create_another_table_schema() -> TableSchema { let table_id = TableId::new(67890); let table_name = TableName::new("public".to_string(), "another_table".to_string()); let columns = vec![ - ColumnSchema::new("id".to_string(), PgType::INT8, -1, false, true), - ColumnSchema::new("description".to_string(), PgType::VARCHAR, 255, true, false), + ColumnSchema::new("id".to_string(), PgType::INT8, -1, false), + ColumnSchema::new("description".to_string(), PgType::VARCHAR, 255, true), ]; TableSchema::new(table_id, table_name, columns) @@ -338,7 +332,6 @@ async fn test_schema_store_update_existing() { PgType::TIMESTAMPTZ, -1, true, - false, )); // Store updated schema From 0e2bd9d7208e0f69572303e9d1f325d0b7e93662 Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Tue, 30 Sep 2025 11:34:35 -0700 Subject: [PATCH 05/45] Improve --- etl-replicator/src/main.rs | 8 ++++---- etl/src/conversions/event.rs | 38 ++++++++++++++++++++++++++++++++---- 2 files changed, 38 insertions(+), 8 deletions(-) diff --git a/etl-replicator/src/main.rs b/etl-replicator/src/main.rs index fd111c8c8..4448b9323 100644 --- a/etl-replicator/src/main.rs +++ b/etl-replicator/src/main.rs @@ -4,16 +4,16 @@ //! and routes data to configured destinations. Includes telemetry, error handling, and //! graceful shutdown capabilities. -use crate::config::load_replicator_config; -use crate::core::start_replicator_with_config; use etl_config::Environment; use etl_config::shared::ReplicatorConfig; use etl_telemetry::metrics::init_metrics; use etl_telemetry::tracing::init_tracing_with_top_level_fields; use std::sync::Arc; -use thiserror::__private::AsDynError; use tracing::{error, info}; +use crate::config::load_replicator_config; +use crate::core::start_replicator_with_config; + mod config; mod core; mod migrations; @@ -65,7 +65,7 @@ fn main() -> anyhow::Result<()> { async fn async_main(replicator_config: ReplicatorConfig) -> anyhow::Result<()> { // We start the replicator and catch any errors. if let Err(err) = start_replicator_with_config(replicator_config).await { - sentry::capture_error(err.as_dyn_error()); + sentry::capture_error(&*err); error!("an error occurred in the replicator: {err}"); return Err(err); diff --git a/etl/src/conversions/event.rs b/etl/src/conversions/event.rs index 1861a88bf..1f0498cdb 100644 --- a/etl/src/conversions/event.rs +++ b/etl/src/conversions/event.rs @@ -1,7 +1,8 @@ use core::str; +use std::cmp::Ordering; use etl_postgres::types::{ColumnSchema, TableId, TableSchema, convert_type_oid_to_type}; use postgres_replication::protocol; -use std::collections::HashSet; +use std::collections::{BTreeSet, HashSet}; use std::hash::Hash; use std::sync::Arc; use tokio_postgres::types::PgLsn; @@ -32,6 +33,19 @@ impl PartialEq for IndexedColumnSchema { } } +impl Ord for IndexedColumnSchema { + fn cmp(&self, other: &Self) -> Ordering { + self.0.name.cmp(&other.0.name) + } +} + +impl PartialOrd for IndexedColumnSchema { + + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.0.name.cmp(&other.0.name)) + } +} + impl Hash for IndexedColumnSchema { fn hash(&self, state: &mut H) { self.0.name.hash(state); @@ -88,7 +102,7 @@ where { let table_id = TableId::new(relation_body.rel_id()); - let Some(existing_column_schemas) = schema_store.get_table_schema(&table_id).await? else { + let Some(table_schema) = schema_store.get_table_schema(&table_id).await? else { bail!( ErrorKind::MissingTableSchema, "Table not found in the schema cache", @@ -96,14 +110,30 @@ where ) }; + // We build a set of the new column chemas from the relation message. The rationale for using a + // BTreeSet is that we want to preserve the order of columns in the schema, which is important, + // and also we can reuse the same set to build the vector of column schemas needed for the table + // schema. let mut latest_column_schemas = relation_body .columns() .iter() .map(build_indexed_column_schema) - .collect::, EtlError>>()?; + .collect::, EtlError>>()?; + + // We build the updated table schema to store in the schema store. + let mut latest_table_schema = TableSchema::new( + table_schema.id, + table_schema.name.clone(), + Vec::with_capacity(latest_column_schemas.len()), + ); + for column_schema in latest_column_schemas.iter() { + latest_table_schema.add_column_schema(column_schema.clone().into_inner()); + } + schema_store.store_table_schema(latest_table_schema).await?; + // We process all the changes that we want to dispatch to the destination. let mut changes = vec![]; - for column_schema in existing_column_schemas.column_schemas.iter() { + for column_schema in table_schema.column_schemas.iter() { let column_schema = IndexedColumnSchema(column_schema.clone()); let latest_column_schema = latest_column_schemas.take(&column_schema); match latest_column_schema { From d4701be0c5a782d3cb107028a7906d9092882205 Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Tue, 30 Sep 2025 11:49:58 -0700 Subject: [PATCH 06/45] Improve --- etl-destinations/src/bigquery/client.rs | 42 +-- etl-destinations/tests/iceberg_pipeline.rs | 388 +++------------------ etl-postgres/src/tokio/test_utils.rs | 8 +- etl/src/conversions/event.rs | 3 +- etl/src/conversions/table_row.rs | 4 +- etl/src/replication/client.rs | 10 +- etl/src/test_utils/test_schema.rs | 24 +- etl/tests/pipeline.rs | 10 +- etl/tests/replication.rs | 32 +- 9 files changed, 96 insertions(+), 425 deletions(-) diff --git a/etl-destinations/src/bigquery/client.rs b/etl-destinations/src/bigquery/client.rs index 975acc744..96d8b6cf6 100644 --- a/etl-destinations/src/bigquery/client.rs +++ b/etl-destinations/src/bigquery/client.rs @@ -773,11 +773,11 @@ mod tests { #[test] fn test_column_spec() { - let column_schema = ColumnSchema::new("test_col".to_string(), Type::TEXT, -1, true); + let column_schema = ColumnSchema::new("test_col".to_string(), Type::TEXT, -1, false); let spec = BigQueryClient::column_spec(&column_schema); assert_eq!(spec, "`test_col` string"); - let not_null_column = ColumnSchema::new("id".to_string(), Type::INT4, -1, false); + let not_null_column = ColumnSchema::new("id".to_string(), Type::INT4, -1, true); let not_null_spec = BigQueryClient::column_spec(¬_null_column); assert_eq!(not_null_spec, "`id` int64 not null"); @@ -789,16 +789,16 @@ mod tests { #[test] fn test_add_primary_key_clause() { let columns_with_pk = vec![ - ColumnSchema::new("id".to_string(), Type::INT4, -1, false), - ColumnSchema::new("name".to_string(), Type::TEXT, -1, true), + ColumnSchema::new("id".to_string(), Type::INT4, -1, true), + ColumnSchema::new("name".to_string(), Type::TEXT, -1, false), ]; let pk_clause = BigQueryClient::add_primary_key_clause(&columns_with_pk); assert_eq!(pk_clause, ", primary key (`id`) not enforced"); let columns_with_composite_pk = vec![ - ColumnSchema::new("tenant_id".to_string(), Type::INT4, -1, false), - ColumnSchema::new("id".to_string(), Type::INT4, -1, false), - ColumnSchema::new("name".to_string(), Type::TEXT, -1, true), + ColumnSchema::new("tenant_id".to_string(), Type::INT4, -1, true), + ColumnSchema::new("id".to_string(), Type::INT4, -1, true), + ColumnSchema::new("name".to_string(), Type::TEXT, -1, false), ]; let composite_pk_clause = BigQueryClient::add_primary_key_clause(&columns_with_composite_pk); @@ -808,8 +808,8 @@ mod tests { ); let columns_no_pk = vec![ - ColumnSchema::new("name".to_string(), Type::TEXT, -1, true), - ColumnSchema::new("age".to_string(), Type::INT4, -1, true), + ColumnSchema::new("name".to_string(), Type::TEXT, -1, false), + ColumnSchema::new("age".to_string(), Type::INT4, -1, false), ]; let no_pk_clause = BigQueryClient::add_primary_key_clause(&columns_no_pk); assert_eq!(no_pk_clause, ""); @@ -818,8 +818,8 @@ mod tests { #[test] fn test_create_columns_spec() { let columns = vec![ - ColumnSchema::new("id".to_string(), Type::INT4, -1, false), - ColumnSchema::new("name".to_string(), Type::TEXT, -1, true), + ColumnSchema::new("id".to_string(), Type::INT4, -1, true), + ColumnSchema::new("name".to_string(), Type::TEXT, -1, false), ColumnSchema::new("active".to_string(), Type::BOOL, -1, false), ]; let spec = BigQueryClient::create_columns_spec(&columns); @@ -838,8 +838,8 @@ mod tests { #[test] fn test_column_schemas_to_table_descriptor() { let columns = vec![ - ColumnSchema::new("id".to_string(), Type::INT4, -1, false), - ColumnSchema::new("name".to_string(), Type::TEXT, -1, true), + ColumnSchema::new("id".to_string(), Type::INT4, -1, true), + ColumnSchema::new("name".to_string(), Type::TEXT, -1, false), ColumnSchema::new("active".to_string(), Type::BOOL, -1, false), ColumnSchema::new("tags".to_string(), Type::TEXT_ARRAY, -1, false), ]; @@ -921,12 +921,12 @@ mod tests { #[test] fn test_column_schemas_to_table_descriptor_complex_types() { let columns = vec![ - ColumnSchema::new("uuid_col".to_string(), Type::UUID, -1, true), - ColumnSchema::new("json_col".to_string(), Type::JSON, -1, true), - ColumnSchema::new("bytea_col".to_string(), Type::BYTEA, -1, true), - ColumnSchema::new("numeric_col".to_string(), Type::NUMERIC, -1, true), - ColumnSchema::new("date_col".to_string(), Type::DATE, -1, true), - ColumnSchema::new("time_col".to_string(), Type::TIME, -1, true), + ColumnSchema::new("uuid_col".to_string(), Type::UUID, -1, false), + ColumnSchema::new("json_col".to_string(), Type::JSON, -1, false), + ColumnSchema::new("bytea_col".to_string(), Type::BYTEA, -1, false), + ColumnSchema::new("numeric_col".to_string(), Type::NUMERIC, -1, false), + ColumnSchema::new("date_col".to_string(), Type::DATE, -1, false), + ColumnSchema::new("time_col".to_string(), Type::TIME, -1, false), ]; let descriptor = BigQueryClient::column_schemas_to_table_descriptor(&columns, true); @@ -978,8 +978,8 @@ mod tests { let table_id = "test_table"; let columns = vec![ - ColumnSchema::new("id".to_string(), Type::INT4, -1, false), - ColumnSchema::new("name".to_string(), Type::TEXT, -1, true), + ColumnSchema::new("id".to_string(), Type::INT4, -1, true), + ColumnSchema::new("name".to_string(), Type::TEXT, -1, false), ]; // Simulate the query generation logic diff --git a/etl-destinations/tests/iceberg_pipeline.rs b/etl-destinations/tests/iceberg_pipeline.rs index 8e574227b..653189b19 100644 --- a/etl-destinations/tests/iceberg_pipeline.rs +++ b/etl-destinations/tests/iceberg_pipeline.rs @@ -147,20 +147,8 @@ async fn create_table_if_missing() { // Date/Time types ColumnSchema::new("date_col".to_string(), Type::DATE, -1, true), ColumnSchema::new("time_col".to_string(), Type::TIME, -1, true), - ColumnSchema::new( - "timestamp_col".to_string(), - Type::TIMESTAMP, - -1, - true, - - ), - ColumnSchema::new( - "timestamptz_col".to_string(), - Type::TIMESTAMPTZ, - -1, - true, - - ), + ColumnSchema::new("timestamp_col".to_string(), Type::TIMESTAMP, -1, true), + ColumnSchema::new("timestamptz_col".to_string(), Type::TIMESTAMPTZ, -1, true), // UUID type ColumnSchema::new("uuid_col".to_string(), Type::UUID, -1, true), // JSON types @@ -171,94 +159,30 @@ async fn create_table_if_missing() { // Binary type ColumnSchema::new("bytea_col".to_string(), Type::BYTEA, -1, true), // Array types - ColumnSchema::new( - "bool_array_col".to_string(), - Type::BOOL_ARRAY, - -1, - true, - - ), - ColumnSchema::new( - "char_array_col".to_string(), - Type::CHAR_ARRAY, - -1, - true, - - ), - ColumnSchema::new( - "bpchar_array_col".to_string(), - Type::BPCHAR_ARRAY, - -1, - true, - - ), + ColumnSchema::new("bool_array_col".to_string(), Type::BOOL_ARRAY, -1, true), + ColumnSchema::new("char_array_col".to_string(), Type::CHAR_ARRAY, -1, true), + ColumnSchema::new("bpchar_array_col".to_string(), Type::BPCHAR_ARRAY, -1, true), ColumnSchema::new( "varchar_array_col".to_string(), Type::VARCHAR_ARRAY, -1, true, - - ), - ColumnSchema::new( - "name_array_col".to_string(), - Type::NAME_ARRAY, - -1, - true, - ), - ColumnSchema::new( - "text_array_col".to_string(), - Type::TEXT_ARRAY, - -1, - true, - ), - ColumnSchema::new( - "int2_array_col".to_string(), - Type::INT2_ARRAY, - -1, - true, - ), - ColumnSchema::new( - "int4_array_col".to_string(), - Type::INT4_ARRAY, - -1, - true, - ), - ColumnSchema::new( - "int8_array_col".to_string(), - Type::INT8_ARRAY, - -1, - true, - ), - ColumnSchema::new( - "float4_array_col".to_string(), - Type::FLOAT4_ARRAY, - -1, - true, - ), - ColumnSchema::new( - "float8_array_col".to_string(), - Type::FLOAT8_ARRAY, - -1, - true, ), + ColumnSchema::new("name_array_col".to_string(), Type::NAME_ARRAY, -1, true), + ColumnSchema::new("text_array_col".to_string(), Type::TEXT_ARRAY, -1, true), + ColumnSchema::new("int2_array_col".to_string(), Type::INT2_ARRAY, -1, true), + ColumnSchema::new("int4_array_col".to_string(), Type::INT4_ARRAY, -1, true), + ColumnSchema::new("int8_array_col".to_string(), Type::INT8_ARRAY, -1, true), + ColumnSchema::new("float4_array_col".to_string(), Type::FLOAT4_ARRAY, -1, true), + ColumnSchema::new("float8_array_col".to_string(), Type::FLOAT8_ARRAY, -1, true), ColumnSchema::new( "numeric_array_col".to_string(), Type::NUMERIC_ARRAY, -1, true, ), - ColumnSchema::new( - "date_array_col".to_string(), - Type::DATE_ARRAY, - -1, - true, - ), - ColumnSchema::new( - "time_array_col".to_string(), - Type::TIME_ARRAY, - -1, - true, - ), + ColumnSchema::new("date_array_col".to_string(), Type::DATE_ARRAY, -1, true), + ColumnSchema::new("time_array_col".to_string(), Type::TIME_ARRAY, -1, true), ColumnSchema::new( "timestamp_array_col".to_string(), Type::TIMESTAMP_ARRAY, @@ -271,36 +195,11 @@ async fn create_table_if_missing() { -1, true, ), - ColumnSchema::new( - "uuid_array_col".to_string(), - Type::UUID_ARRAY, - -1, - true, - ), - ColumnSchema::new( - "json_array_col".to_string(), - Type::JSON_ARRAY, - -1, - true, - ), - ColumnSchema::new( - "jsonb_array_col".to_string(), - Type::JSONB_ARRAY, - -1, - true, - ), - ColumnSchema::new( - "oid_array_col".to_string(), - Type::OID_ARRAY, - -1, - true, - ), - ColumnSchema::new( - "bytea_array_col".to_string(), - Type::BYTEA_ARRAY, - -1, - true, - ), + ColumnSchema::new("uuid_array_col".to_string(), Type::UUID_ARRAY, -1, true), + ColumnSchema::new("json_array_col".to_string(), Type::JSON_ARRAY, -1, true), + ColumnSchema::new("jsonb_array_col".to_string(), Type::JSONB_ARRAY, -1, true), + ColumnSchema::new("oid_array_col".to_string(), Type::OID_ARRAY, -1, true), + ColumnSchema::new("bytea_array_col".to_string(), Type::BYTEA_ARRAY, -1, true), ]; let table_schema = TableSchema::new(table_id, table_name_struct, columns); @@ -392,20 +291,8 @@ async fn insert_nullable_scalars() { // Date/Time types ColumnSchema::new("date_col".to_string(), Type::DATE, -1, true), ColumnSchema::new("time_col".to_string(), Type::TIME, -1, true), - ColumnSchema::new( - "timestamp_col".to_string(), - Type::TIMESTAMP, - -1, - true, - - ), - ColumnSchema::new( - "timestamptz_col".to_string(), - Type::TIMESTAMPTZ, - -1, - true, - - ), + ColumnSchema::new("timestamp_col".to_string(), Type::TIMESTAMP, -1, true), + ColumnSchema::new("timestamptz_col".to_string(), Type::TIMESTAMPTZ, -1, true), // UUID type ColumnSchema::new("uuid_col".to_string(), Type::UUID, -1, true), // JSON types @@ -548,25 +435,15 @@ async fn insert_non_nullable_scalars() { ColumnSchema::new("int4_col".to_string(), Type::INT4, -1, false), ColumnSchema::new("int8_col".to_string(), Type::INT8, -1, false), // Float types - ColumnSchema::new("float4_col".to_string(), Type::FLOAT4, -1, false,), + ColumnSchema::new("float4_col".to_string(), Type::FLOAT4, -1, false), ColumnSchema::new("float8_col".to_string(), Type::FLOAT8, -1, false), // Numeric type ColumnSchema::new("numeric_col".to_string(), Type::NUMERIC, -1, false), // Date/Time types ColumnSchema::new("date_col".to_string(), Type::DATE, -1, false), ColumnSchema::new("time_col".to_string(), Type::TIME, -1, false), - ColumnSchema::new( - "timestamp_col".to_string(), - Type::TIMESTAMP, - -1, - false, - ), - ColumnSchema::new( - "timestamptz_col".to_string(), - Type::TIMESTAMPTZ, - -1, - false, - ), + ColumnSchema::new("timestamp_col".to_string(), Type::TIMESTAMP, -1, false), + ColumnSchema::new("timestamptz_col".to_string(), Type::TIMESTAMPTZ, -1, false), // UUID type ColumnSchema::new("uuid_col".to_string(), Type::UUID, -1, false), // JSON types @@ -671,75 +548,25 @@ async fn insert_nullable_array() { // Primary key ColumnSchema::new("id".to_string(), Type::INT4, -1, false), // Boolean array type - ColumnSchema::new( - "bool_array_col".to_string(), - Type::BOOL_ARRAY, - -1, - true, - ), + ColumnSchema::new("bool_array_col".to_string(), Type::BOOL_ARRAY, -1, true), // String array types - ColumnSchema::new( - "char_array_col".to_string(), - Type::CHAR_ARRAY, - -1, - true, - ), - ColumnSchema::new( - "bpchar_array_col".to_string(), - Type::BPCHAR_ARRAY, - -1, - true, - ), + ColumnSchema::new("char_array_col".to_string(), Type::CHAR_ARRAY, -1, true), + ColumnSchema::new("bpchar_array_col".to_string(), Type::BPCHAR_ARRAY, -1, true), ColumnSchema::new( "varchar_array_col".to_string(), Type::VARCHAR_ARRAY, -1, true, ), - ColumnSchema::new( - "name_array_col".to_string(), - Type::NAME_ARRAY, - -1, - true, - ), - ColumnSchema::new( - "text_array_col".to_string(), - Type::TEXT_ARRAY, - -1, - true, - ), + ColumnSchema::new("name_array_col".to_string(), Type::NAME_ARRAY, -1, true), + ColumnSchema::new("text_array_col".to_string(), Type::TEXT_ARRAY, -1, true), // Integer array types - ColumnSchema::new( - "int2_array_col".to_string(), - Type::INT2_ARRAY, - -1, - true, - ), - ColumnSchema::new( - "int4_array_col".to_string(), - Type::INT4_ARRAY, - -1, - true, - ), - ColumnSchema::new( - "int8_array_col".to_string(), - Type::INT8_ARRAY, - -1, - true, - ), + ColumnSchema::new("int2_array_col".to_string(), Type::INT2_ARRAY, -1, true), + ColumnSchema::new("int4_array_col".to_string(), Type::INT4_ARRAY, -1, true), + ColumnSchema::new("int8_array_col".to_string(), Type::INT8_ARRAY, -1, true), // Float array types - ColumnSchema::new( - "float4_array_col".to_string(), - Type::FLOAT4_ARRAY, - -1, - true, - ), - ColumnSchema::new( - "float8_array_col".to_string(), - Type::FLOAT8_ARRAY, - -1, - true, - ), + ColumnSchema::new("float4_array_col".to_string(), Type::FLOAT4_ARRAY, -1, true), + ColumnSchema::new("float8_array_col".to_string(), Type::FLOAT8_ARRAY, -1, true), // Numeric array type ColumnSchema::new( "numeric_array_col".to_string(), @@ -748,18 +575,8 @@ async fn insert_nullable_array() { true, ), // Date/Time array types - ColumnSchema::new( - "date_array_col".to_string(), - Type::DATE_ARRAY, - -1, - true, - ), - ColumnSchema::new( - "time_array_col".to_string(), - Type::TIME_ARRAY, - -1, - true, - ), + ColumnSchema::new("date_array_col".to_string(), Type::DATE_ARRAY, -1, true), + ColumnSchema::new("time_array_col".to_string(), Type::TIME_ARRAY, -1, true), ColumnSchema::new( "timestamp_array_col".to_string(), Type::TIMESTAMP_ARRAY, @@ -773,39 +590,14 @@ async fn insert_nullable_array() { true, ), // UUID array type - ColumnSchema::new( - "uuid_array_col".to_string(), - Type::UUID_ARRAY, - -1, - true, - ), + ColumnSchema::new("uuid_array_col".to_string(), Type::UUID_ARRAY, -1, true), // JSON array types - ColumnSchema::new( - "json_array_col".to_string(), - Type::JSON_ARRAY, - -1, - true, - ), - ColumnSchema::new( - "jsonb_array_col".to_string(), - Type::JSONB_ARRAY, - -1, - true, - ), + ColumnSchema::new("json_array_col".to_string(), Type::JSON_ARRAY, -1, true), + ColumnSchema::new("jsonb_array_col".to_string(), Type::JSONB_ARRAY, -1, true), // OID array type - ColumnSchema::new( - "oid_array_col".to_string(), - Type::OID_ARRAY, - -1, - true, - ), + ColumnSchema::new("oid_array_col".to_string(), Type::OID_ARRAY, -1, true), // Binary array type - ColumnSchema::new( - "bytea_array_col".to_string(), - Type::BYTEA_ARRAY, - -1, - true, - ), + ColumnSchema::new("bytea_array_col".to_string(), Type::BYTEA_ARRAY, -1, true), ]; let table_schema = TableSchema::new(table_id, table_name_struct, columns); @@ -1015,19 +807,9 @@ async fn insert_non_nullable_array() { // Primary key ColumnSchema::new("id".to_string(), Type::INT4, -1, false), // Boolean array type - ColumnSchema::new( - "bool_array_col".to_string(), - Type::BOOL_ARRAY, - -1, - false, - ), + ColumnSchema::new("bool_array_col".to_string(), Type::BOOL_ARRAY, -1, false), // String array types - ColumnSchema::new( - "char_array_col".to_string(), - Type::CHAR_ARRAY, - -1, - false, - ), + ColumnSchema::new("char_array_col".to_string(), Type::CHAR_ARRAY, -1, false), ColumnSchema::new( "bpchar_array_col".to_string(), Type::BPCHAR_ARRAY, @@ -1040,37 +822,12 @@ async fn insert_non_nullable_array() { -1, false, ), - ColumnSchema::new( - "name_array_col".to_string(), - Type::NAME_ARRAY, - -1, - false, - ), - ColumnSchema::new( - "text_array_col".to_string(), - Type::TEXT_ARRAY, - -1, - false, - ), + ColumnSchema::new("name_array_col".to_string(), Type::NAME_ARRAY, -1, false), + ColumnSchema::new("text_array_col".to_string(), Type::TEXT_ARRAY, -1, false), // Integer array types - ColumnSchema::new( - "int2_array_col".to_string(), - Type::INT2_ARRAY, - -1, - false, - ), - ColumnSchema::new( - "int4_array_col".to_string(), - Type::INT4_ARRAY, - -1, - false, - ), - ColumnSchema::new( - "int8_array_col".to_string(), - Type::INT8_ARRAY, - -1, - false, - ), + ColumnSchema::new("int2_array_col".to_string(), Type::INT2_ARRAY, -1, false), + ColumnSchema::new("int4_array_col".to_string(), Type::INT4_ARRAY, -1, false), + ColumnSchema::new("int8_array_col".to_string(), Type::INT8_ARRAY, -1, false), // Float array types ColumnSchema::new( "float4_array_col".to_string(), @@ -1092,18 +849,8 @@ async fn insert_non_nullable_array() { false, ), // Date/Time array types - ColumnSchema::new( - "date_array_col".to_string(), - Type::DATE_ARRAY, - -1, - false, - ), - ColumnSchema::new( - "time_array_col".to_string(), - Type::TIME_ARRAY, - -1, - false, - ), + ColumnSchema::new("date_array_col".to_string(), Type::DATE_ARRAY, -1, false), + ColumnSchema::new("time_array_col".to_string(), Type::TIME_ARRAY, -1, false), ColumnSchema::new( "timestamp_array_col".to_string(), Type::TIMESTAMP_ARRAY, @@ -1117,39 +864,14 @@ async fn insert_non_nullable_array() { false, ), // UUID array type - ColumnSchema::new( - "uuid_array_col".to_string(), - Type::UUID_ARRAY, - -1, - false, - ), + ColumnSchema::new("uuid_array_col".to_string(), Type::UUID_ARRAY, -1, false), // JSON array types - ColumnSchema::new( - "json_array_col".to_string(), - Type::JSON_ARRAY, - -1, - false, - ), - ColumnSchema::new( - "jsonb_array_col".to_string(), - Type::JSONB_ARRAY, - -1, - false, - ), + ColumnSchema::new("json_array_col".to_string(), Type::JSON_ARRAY, -1, false), + ColumnSchema::new("jsonb_array_col".to_string(), Type::JSONB_ARRAY, -1, false), // OID array type - ColumnSchema::new( - "oid_array_col".to_string(), - Type::OID_ARRAY, - -1, - false, - ), + ColumnSchema::new("oid_array_col".to_string(), Type::OID_ARRAY, -1, false), // Binary array type - ColumnSchema::new( - "bytea_array_col".to_string(), - Type::BYTEA_ARRAY, - -1, - false, - ), + ColumnSchema::new("bytea_array_col".to_string(), Type::BYTEA_ARRAY, -1, false), ]; let table_schema = TableSchema::new(table_id, table_name_struct, columns); diff --git a/etl-postgres/src/tokio/test_utils.rs b/etl-postgres/src/tokio/test_utils.rs index 9100ccc7a..7db5db3cf 100644 --- a/etl-postgres/src/tokio/test_utils.rs +++ b/etl-postgres/src/tokio/test_utils.rs @@ -471,13 +471,7 @@ impl Drop for PgDatabase { /// Creates a [`ColumnSchema`] for a non-nullable, primary key column named "id" /// of type `INT8` that is added by default to tables created by [`PgDatabase`]. pub fn id_column_schema() -> ColumnSchema { - ColumnSchema { - name: "id".to_string(), - typ: Type::INT8, - modifier: -1, - nullable: false, - primary: true, - } + ColumnSchema::new("id".to_string(), Type::INT8, -1, true) } /// Creates a new Postgres database and returns a connected client. diff --git a/etl/src/conversions/event.rs b/etl/src/conversions/event.rs index 1f0498cdb..4b28705f1 100644 --- a/etl/src/conversions/event.rs +++ b/etl/src/conversions/event.rs @@ -1,7 +1,7 @@ use core::str; -use std::cmp::Ordering; use etl_postgres::types::{ColumnSchema, TableId, TableSchema, convert_type_oid_to_type}; use postgres_replication::protocol; +use std::cmp::Ordering; use std::collections::{BTreeSet, HashSet}; use std::hash::Hash; use std::sync::Arc; @@ -40,7 +40,6 @@ impl Ord for IndexedColumnSchema { } impl PartialOrd for IndexedColumnSchema { - fn partial_cmp(&self, other: &Self) -> Option { Some(self.0.name.cmp(&other.0.name)) } diff --git a/etl/src/conversions/table_row.rs b/etl/src/conversions/table_row.rs index 50e84c1da..8bd8985c7 100644 --- a/etl/src/conversions/table_row.rs +++ b/etl/src/conversions/table_row.rs @@ -164,8 +164,8 @@ mod tests { fn create_test_schema() -> Vec { vec![ - ColumnSchema::new("id".to_string(), Type::INT4, -1, false), - ColumnSchema::new("name".to_string(), Type::TEXT, -1, true), + ColumnSchema::new("id".to_string(), Type::INT4, -1, true), + ColumnSchema::new("name".to_string(), Type::TEXT, -1, false), ColumnSchema::new("active".to_string(), Type::BOOL, -1, false), ] } diff --git a/etl/src/replication/client.rs b/etl/src/replication/client.rs index 722437505..98cb0a878 100644 --- a/etl/src/replication/client.rs +++ b/etl/src/replication/client.rs @@ -712,20 +712,12 @@ impl PgReplicationClient { let type_oid = Self::get_row_value::(&row, "atttypid", "pg_attribute").await?; let modifier = Self::get_row_value::(&row, "atttypmod", "pg_attribute").await?; - let nullable = - Self::get_row_value::(&row, "attnotnull", "pg_attribute").await? == "f"; let primary = Self::get_row_value::(&row, "primary", "pg_index").await? == "t"; let typ = convert_type_oid_to_type(type_oid); - column_schemas.push(ColumnSchema { - name, - typ, - modifier, - nullable, - primary, - }) + column_schemas.push(ColumnSchema::new(name, typ, modifier, primary)) } } diff --git a/etl/src/test_utils/test_schema.rs b/etl/src/test_utils/test_schema.rs index caf2d64bc..392077fd0 100644 --- a/etl/src/test_utils/test_schema.rs +++ b/etl/src/test_utils/test_schema.rs @@ -66,20 +66,8 @@ pub async fn setup_test_database_schema( users_table_name, vec![ id_column_schema(), - ColumnSchema { - name: "name".to_string(), - typ: Type::TEXT, - modifier: -1, - nullable: false, - primary: false, - }, - ColumnSchema { - name: "age".to_string(), - typ: Type::INT4, - modifier: -1, - nullable: false, - primary: false, - }, + ColumnSchema::new("name".to_string(), Type::TEXT, -1, false), + ColumnSchema::new("age".to_string(), Type::INT4, -1, false), ], )); } @@ -102,13 +90,7 @@ pub async fn setup_test_database_schema( orders_table_name, vec![ id_column_schema(), - ColumnSchema { - name: "description".to_string(), - typ: Type::TEXT, - modifier: -1, - nullable: false, - primary: false, - }, + ColumnSchema::new("description".to_string(), Type::TEXT, -1, false), ], )); } diff --git a/etl/tests/pipeline.rs b/etl/tests/pipeline.rs index ec5f38370..9c611f07e 100644 --- a/etl/tests/pipeline.rs +++ b/etl/tests/pipeline.rs @@ -526,7 +526,10 @@ async fn table_schema_changes_are_handled_correctly() { // Check the initial schema. let table_schemas = store.get_table_schemas().await; assert_eq!(table_schemas.len(), 1); - assert!(table_schemas.contains_key(&database_schema.users_schema().id)); + let users_table_schema = table_schemas + .get(&database_schema.users_schema().id) + .unwrap(); + println!("{:?}", users_table_schema); // Check the initial data. let table_rows = destination.get_table_rows().await; @@ -571,7 +574,10 @@ async fn table_schema_changes_are_handled_correctly() { // Check the updated schema. let table_schemas = store.get_table_schemas().await; assert_eq!(table_schemas.len(), 1); - assert!(table_schemas.contains_key(&database_schema.users_schema().id)); + let users_table_schema = table_schemas + .get(&database_schema.users_schema().id) + .unwrap(); + println!("{:?}", users_table_schema); // Check the updated data. let table_rows = destination.get_table_rows().await; diff --git a/etl/tests/replication.rs b/etl/tests/replication.rs index 9602a6cc5..fec50b3d5 100644 --- a/etl/tests/replication.rs +++ b/etl/tests/replication.rs @@ -167,13 +167,7 @@ async fn test_table_schema_copy_is_consistent() { .await .unwrap(); - let age_schema = ColumnSchema { - name: "age".to_string(), - typ: Type::INT4, - modifier: -1, - nullable: true, - primary: false, - }; + let age_schema = ColumnSchema::new("age".to_string(), Type::INT4, -1, false); let table_1_id = database .create_table(test_table_name("table_1"), true, &[("age", "integer")]) @@ -212,20 +206,8 @@ async fn test_table_schema_copy_across_multiple_connections() { .await .unwrap(); - let age_schema = ColumnSchema { - name: "age".to_string(), - typ: Type::INT4, - modifier: -1, - nullable: true, - primary: false, - }; - let year_schema = ColumnSchema { - name: "year".to_string(), - typ: Type::INT4, - modifier: -1, - nullable: true, - primary: false, - }; + let age_schema = ColumnSchema::new("age".to_string(), Type::INT4, -1, false); + let year_schema = ColumnSchema::new("year".to_string(), Type::INT4, -1, false); let table_1_id = database .create_table(test_table_name("table_1"), true, &[("age", "integer")]) @@ -341,13 +323,7 @@ async fn test_table_copy_stream_is_consistent() { let stream = transaction .get_table_copy_stream( table_1_id, - &[ColumnSchema { - name: "age".to_string(), - typ: Type::INT4, - modifier: -1, - nullable: true, - primary: false, - }], + &[ColumnSchema::new("age".to_string(), Type::INT4, -1, false)], ) .await .unwrap(); From 14c364f94f6459016b363be3c276da7715561cfc Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Tue, 30 Sep 2025 11:52:54 -0700 Subject: [PATCH 07/45] Improve --- etl/tests/postgres_store.rs | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/etl/tests/postgres_store.rs b/etl/tests/postgres_store.rs index afdcf4960..f333e9c18 100644 --- a/etl/tests/postgres_store.rs +++ b/etl/tests/postgres_store.rs @@ -16,8 +16,8 @@ fn create_sample_table_schema() -> TableSchema { let table_id = TableId::new(12345); let table_name = TableName::new("public".to_string(), "test_table".to_string()); let columns = vec![ - ColumnSchema::new("id".to_string(), PgType::INT4, -1, false), - ColumnSchema::new("name".to_string(), PgType::TEXT, -1, true), + ColumnSchema::new("id".to_string(), PgType::INT4, -1, true), + ColumnSchema::new("name".to_string(), PgType::TEXT, -1, false), ColumnSchema::new("created_at".to_string(), PgType::TIMESTAMPTZ, -1, false), ]; @@ -28,8 +28,8 @@ fn create_another_table_schema() -> TableSchema { let table_id = TableId::new(67890); let table_name = TableName::new("public".to_string(), "another_table".to_string()); let columns = vec![ - ColumnSchema::new("id".to_string(), PgType::INT8, -1, false), - ColumnSchema::new("description".to_string(), PgType::VARCHAR, 255, true), + ColumnSchema::new("id".to_string(), PgType::INT8, -1, true), + ColumnSchema::new("description".to_string(), PgType::VARCHAR, 255, false), ]; TableSchema::new(table_id, table_name, columns) @@ -331,7 +331,7 @@ async fn test_schema_store_update_existing() { "updated_at".to_string(), PgType::TIMESTAMPTZ, -1, - true, + false, )); // Store updated schema From ecfbb09690a6fe2b146df27a60135d114c264845 Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Tue, 30 Sep 2025 11:57:43 -0700 Subject: [PATCH 08/45] Improve --- etl-destinations/tests/iceberg_pipeline.rs | 208 ++++++++++++--------- etl/src/conversions/event.rs | 9 +- 2 files changed, 120 insertions(+), 97 deletions(-) diff --git a/etl-destinations/tests/iceberg_pipeline.rs b/etl-destinations/tests/iceberg_pipeline.rs index 653189b19..84aa0e407 100644 --- a/etl-destinations/tests/iceberg_pipeline.rs +++ b/etl-destinations/tests/iceberg_pipeline.rs @@ -126,80 +126,95 @@ async fn create_table_if_missing() { let table_name_struct = TableName::new("test_schema".to_string(), table_name.clone()); let columns = vec![ // Primary key - ColumnSchema::new("id".to_string(), Type::INT4, -1, false), + ColumnSchema::new("id".to_string(), Type::INT4, -1, true), // Boolean types - ColumnSchema::new("bool_col".to_string(), Type::BOOL, -1, true), + ColumnSchema::new("bool_col".to_string(), Type::BOOL, -1, false), // String types - ColumnSchema::new("char_col".to_string(), Type::CHAR, -1, true), - ColumnSchema::new("bpchar_col".to_string(), Type::BPCHAR, -1, true), - ColumnSchema::new("varchar_col".to_string(), Type::VARCHAR, -1, true), - ColumnSchema::new("name_col".to_string(), Type::NAME, -1, true), - ColumnSchema::new("text_col".to_string(), Type::TEXT, -1, true), + ColumnSchema::new("char_col".to_string(), Type::CHAR, -1, false), + ColumnSchema::new("bpchar_col".to_string(), Type::BPCHAR, -1, false), + ColumnSchema::new("varchar_col".to_string(), Type::VARCHAR, -1, false), + ColumnSchema::new("name_col".to_string(), Type::NAME, -1, false), + ColumnSchema::new("text_col".to_string(), Type::TEXT, -1, false), // Integer types - ColumnSchema::new("int2_col".to_string(), Type::INT2, -1, true), - ColumnSchema::new("int4_col".to_string(), Type::INT4, -1, true), - ColumnSchema::new("int8_col".to_string(), Type::INT8, -1, true), + ColumnSchema::new("int2_col".to_string(), Type::INT2, -1, false), + ColumnSchema::new("int4_col".to_string(), Type::INT4, -1, false), + ColumnSchema::new("int8_col".to_string(), Type::INT8, -1, false), // Float types - ColumnSchema::new("float4_col".to_string(), Type::FLOAT4, -1, true), - ColumnSchema::new("float8_col".to_string(), Type::FLOAT8, -1, true), + ColumnSchema::new("float4_col".to_string(), Type::FLOAT4, -1, false), + ColumnSchema::new("float8_col".to_string(), Type::FLOAT8, -1, false), // Numeric type - ColumnSchema::new("numeric_col".to_string(), Type::NUMERIC, -1, true), + ColumnSchema::new("numeric_col".to_string(), Type::NUMERIC, -1, false), // Date/Time types - ColumnSchema::new("date_col".to_string(), Type::DATE, -1, true), - ColumnSchema::new("time_col".to_string(), Type::TIME, -1, true), - ColumnSchema::new("timestamp_col".to_string(), Type::TIMESTAMP, -1, true), - ColumnSchema::new("timestamptz_col".to_string(), Type::TIMESTAMPTZ, -1, true), + ColumnSchema::new("date_col".to_string(), Type::DATE, -1, false), + ColumnSchema::new("time_col".to_string(), Type::TIME, -1, false), + ColumnSchema::new("timestamp_col".to_string(), Type::TIMESTAMP, -1, false), + ColumnSchema::new("timestamptz_col".to_string(), Type::TIMESTAMPTZ, -1, false), // UUID type - ColumnSchema::new("uuid_col".to_string(), Type::UUID, -1, true), + ColumnSchema::new("uuid_col".to_string(), Type::UUID, -1, false), // JSON types - ColumnSchema::new("json_col".to_string(), Type::JSON, -1, true), - ColumnSchema::new("jsonb_col".to_string(), Type::JSONB, -1, true), + ColumnSchema::new("json_col".to_string(), Type::JSON, -1, false), + ColumnSchema::new("jsonb_col".to_string(), Type::JSONB, -1, false), // OID type - ColumnSchema::new("oid_col".to_string(), Type::OID, -1, true), + ColumnSchema::new("oid_col".to_string(), Type::OID, -1, false), // Binary type - ColumnSchema::new("bytea_col".to_string(), Type::BYTEA, -1, true), + ColumnSchema::new("bytea_col".to_string(), Type::BYTEA, -1, false), // Array types - ColumnSchema::new("bool_array_col".to_string(), Type::BOOL_ARRAY, -1, true), - ColumnSchema::new("char_array_col".to_string(), Type::CHAR_ARRAY, -1, true), - ColumnSchema::new("bpchar_array_col".to_string(), Type::BPCHAR_ARRAY, -1, true), + ColumnSchema::new("bool_array_col".to_string(), Type::BOOL_ARRAY, -1, false), + ColumnSchema::new("char_array_col".to_string(), Type::CHAR_ARRAY, -1, false), + ColumnSchema::new( + "bpchar_array_col".to_string(), + Type::BPCHAR_ARRAY, + -1, + false, + ), ColumnSchema::new( "varchar_array_col".to_string(), Type::VARCHAR_ARRAY, -1, - true, + false, + ), + ColumnSchema::new("name_array_col".to_string(), Type::NAME_ARRAY, -1, false), + ColumnSchema::new("text_array_col".to_string(), Type::TEXT_ARRAY, -1, false), + ColumnSchema::new("int2_array_col".to_string(), Type::INT2_ARRAY, -1, false), + ColumnSchema::new("int4_array_col".to_string(), Type::INT4_ARRAY, -1, false), + ColumnSchema::new("int8_array_col".to_string(), Type::INT8_ARRAY, -1, false), + ColumnSchema::new( + "float4_array_col".to_string(), + Type::FLOAT4_ARRAY, + -1, + false, + ), + ColumnSchema::new( + "float8_array_col".to_string(), + Type::FLOAT8_ARRAY, + -1, + false, ), - ColumnSchema::new("name_array_col".to_string(), Type::NAME_ARRAY, -1, true), - ColumnSchema::new("text_array_col".to_string(), Type::TEXT_ARRAY, -1, true), - ColumnSchema::new("int2_array_col".to_string(), Type::INT2_ARRAY, -1, true), - ColumnSchema::new("int4_array_col".to_string(), Type::INT4_ARRAY, -1, true), - ColumnSchema::new("int8_array_col".to_string(), Type::INT8_ARRAY, -1, true), - ColumnSchema::new("float4_array_col".to_string(), Type::FLOAT4_ARRAY, -1, true), - ColumnSchema::new("float8_array_col".to_string(), Type::FLOAT8_ARRAY, -1, true), ColumnSchema::new( "numeric_array_col".to_string(), Type::NUMERIC_ARRAY, -1, - true, + false, ), - ColumnSchema::new("date_array_col".to_string(), Type::DATE_ARRAY, -1, true), - ColumnSchema::new("time_array_col".to_string(), Type::TIME_ARRAY, -1, true), + ColumnSchema::new("date_array_col".to_string(), Type::DATE_ARRAY, -1, false), + ColumnSchema::new("time_array_col".to_string(), Type::TIME_ARRAY, -1, false), ColumnSchema::new( "timestamp_array_col".to_string(), Type::TIMESTAMP_ARRAY, -1, - true, + false, ), ColumnSchema::new( "timestamptz_array_col".to_string(), Type::TIMESTAMPTZ_ARRAY, -1, - true, + false, ), - ColumnSchema::new("uuid_array_col".to_string(), Type::UUID_ARRAY, -1, true), - ColumnSchema::new("json_array_col".to_string(), Type::JSON_ARRAY, -1, true), - ColumnSchema::new("jsonb_array_col".to_string(), Type::JSONB_ARRAY, -1, true), - ColumnSchema::new("oid_array_col".to_string(), Type::OID_ARRAY, -1, true), - ColumnSchema::new("bytea_array_col".to_string(), Type::BYTEA_ARRAY, -1, true), + ColumnSchema::new("uuid_array_col".to_string(), Type::UUID_ARRAY, -1, false), + ColumnSchema::new("json_array_col".to_string(), Type::JSON_ARRAY, -1, false), + ColumnSchema::new("jsonb_array_col".to_string(), Type::JSONB_ARRAY, -1, false), + ColumnSchema::new("oid_array_col".to_string(), Type::OID_ARRAY, -1, false), + ColumnSchema::new("bytea_array_col".to_string(), Type::BYTEA_ARRAY, -1, false), ]; let table_schema = TableSchema::new(table_id, table_name_struct, columns); @@ -270,38 +285,38 @@ async fn insert_nullable_scalars() { let table_name_struct = TableName::new("test_schema".to_string(), table_name.clone()); let columns = vec![ // Primary key - ColumnSchema::new("id".to_string(), Type::INT4, -1, false), + ColumnSchema::new("id".to_string(), Type::INT4, -1, true), // Boolean types - ColumnSchema::new("bool_col".to_string(), Type::BOOL, -1, true), + ColumnSchema::new("bool_col".to_string(), Type::BOOL, -1, false), // String types - ColumnSchema::new("char_col".to_string(), Type::CHAR, -1, true), - ColumnSchema::new("bpchar_col".to_string(), Type::BPCHAR, -1, true), - ColumnSchema::new("varchar_col".to_string(), Type::VARCHAR, -1, true), - ColumnSchema::new("name_col".to_string(), Type::NAME, -1, true), - ColumnSchema::new("text_col".to_string(), Type::TEXT, -1, true), + ColumnSchema::new("char_col".to_string(), Type::CHAR, -1, false), + ColumnSchema::new("bpchar_col".to_string(), Type::BPCHAR, -1, false), + ColumnSchema::new("varchar_col".to_string(), Type::VARCHAR, -1, false), + ColumnSchema::new("name_col".to_string(), Type::NAME, -1, false), + ColumnSchema::new("text_col".to_string(), Type::TEXT, -1, false), // Integer types - ColumnSchema::new("int2_col".to_string(), Type::INT2, -1, true), - ColumnSchema::new("int4_col".to_string(), Type::INT4, -1, true), - ColumnSchema::new("int8_col".to_string(), Type::INT8, -1, true), + ColumnSchema::new("int2_col".to_string(), Type::INT2, -1, false), + ColumnSchema::new("int4_col".to_string(), Type::INT4, -1, false), + ColumnSchema::new("int8_col".to_string(), Type::INT8, -1, false), // Float types - ColumnSchema::new("float4_col".to_string(), Type::FLOAT4, -1, true), - ColumnSchema::new("float8_col".to_string(), Type::FLOAT8, -1, true), + ColumnSchema::new("float4_col".to_string(), Type::FLOAT4, -1, false), + ColumnSchema::new("float8_col".to_string(), Type::FLOAT8, -1, false), // Numeric type - ColumnSchema::new("numeric_col".to_string(), Type::NUMERIC, -1, true), + ColumnSchema::new("numeric_col".to_string(), Type::NUMERIC, -1, false), // Date/Time types - ColumnSchema::new("date_col".to_string(), Type::DATE, -1, true), - ColumnSchema::new("time_col".to_string(), Type::TIME, -1, true), - ColumnSchema::new("timestamp_col".to_string(), Type::TIMESTAMP, -1, true), - ColumnSchema::new("timestamptz_col".to_string(), Type::TIMESTAMPTZ, -1, true), + ColumnSchema::new("date_col".to_string(), Type::DATE, -1, false), + ColumnSchema::new("time_col".to_string(), Type::TIME, -1, false), + ColumnSchema::new("timestamp_col".to_string(), Type::TIMESTAMP, -1, false), + ColumnSchema::new("timestamptz_col".to_string(), Type::TIMESTAMPTZ, -1, false), // UUID type - ColumnSchema::new("uuid_col".to_string(), Type::UUID, -1, true), + ColumnSchema::new("uuid_col".to_string(), Type::UUID, -1, false), // JSON types - ColumnSchema::new("json_col".to_string(), Type::JSON, -1, true), - ColumnSchema::new("jsonb_col".to_string(), Type::JSONB, -1, true), + ColumnSchema::new("json_col".to_string(), Type::JSON, -1, false), + ColumnSchema::new("jsonb_col".to_string(), Type::JSONB, -1, false), // OID type - ColumnSchema::new("oid_col".to_string(), Type::OID, -1, true), + ColumnSchema::new("oid_col".to_string(), Type::OID, -1, false), // Binary type - ColumnSchema::new("bytea_col".to_string(), Type::BYTEA, -1, true), + ColumnSchema::new("bytea_col".to_string(), Type::BYTEA, -1, false), ]; let table_schema = TableSchema::new(table_id, table_name_struct, columns); @@ -421,7 +436,7 @@ async fn insert_non_nullable_scalars() { let table_name_struct = TableName::new("test_schema".to_string(), table_name.clone()); let columns = vec![ // Primary key - ColumnSchema::new("id".to_string(), Type::INT4, -1, false), + ColumnSchema::new("id".to_string(), Type::INT4, -1, true), // Boolean types ColumnSchema::new("bool_col".to_string(), Type::BOOL, -1, false), // String types @@ -546,58 +561,73 @@ async fn insert_nullable_array() { let table_name_struct = TableName::new("test_schema".to_string(), table_name.clone()); let columns = vec![ // Primary key - ColumnSchema::new("id".to_string(), Type::INT4, -1, false), + ColumnSchema::new("id".to_string(), Type::INT4, -1, true), // Boolean array type - ColumnSchema::new("bool_array_col".to_string(), Type::BOOL_ARRAY, -1, true), + ColumnSchema::new("bool_array_col".to_string(), Type::BOOL_ARRAY, -1, false), // String array types - ColumnSchema::new("char_array_col".to_string(), Type::CHAR_ARRAY, -1, true), - ColumnSchema::new("bpchar_array_col".to_string(), Type::BPCHAR_ARRAY, -1, true), + ColumnSchema::new("char_array_col".to_string(), Type::CHAR_ARRAY, -1, false), + ColumnSchema::new( + "bpchar_array_col".to_string(), + Type::BPCHAR_ARRAY, + -1, + false, + ), ColumnSchema::new( "varchar_array_col".to_string(), Type::VARCHAR_ARRAY, -1, - true, + false, ), - ColumnSchema::new("name_array_col".to_string(), Type::NAME_ARRAY, -1, true), - ColumnSchema::new("text_array_col".to_string(), Type::TEXT_ARRAY, -1, true), + ColumnSchema::new("name_array_col".to_string(), Type::NAME_ARRAY, -1, false), + ColumnSchema::new("text_array_col".to_string(), Type::TEXT_ARRAY, -1, false), // Integer array types - ColumnSchema::new("int2_array_col".to_string(), Type::INT2_ARRAY, -1, true), - ColumnSchema::new("int4_array_col".to_string(), Type::INT4_ARRAY, -1, true), - ColumnSchema::new("int8_array_col".to_string(), Type::INT8_ARRAY, -1, true), + ColumnSchema::new("int2_array_col".to_string(), Type::INT2_ARRAY, -1, false), + ColumnSchema::new("int4_array_col".to_string(), Type::INT4_ARRAY, -1, false), + ColumnSchema::new("int8_array_col".to_string(), Type::INT8_ARRAY, -1, false), // Float array types - ColumnSchema::new("float4_array_col".to_string(), Type::FLOAT4_ARRAY, -1, true), - ColumnSchema::new("float8_array_col".to_string(), Type::FLOAT8_ARRAY, -1, true), + ColumnSchema::new( + "float4_array_col".to_string(), + Type::FLOAT4_ARRAY, + -1, + false, + ), + ColumnSchema::new( + "float8_array_col".to_string(), + Type::FLOAT8_ARRAY, + -1, + false, + ), // Numeric array type ColumnSchema::new( "numeric_array_col".to_string(), Type::NUMERIC_ARRAY, -1, - true, + false, ), // Date/Time array types - ColumnSchema::new("date_array_col".to_string(), Type::DATE_ARRAY, -1, true), - ColumnSchema::new("time_array_col".to_string(), Type::TIME_ARRAY, -1, true), + ColumnSchema::new("date_array_col".to_string(), Type::DATE_ARRAY, -1, false), + ColumnSchema::new("time_array_col".to_string(), Type::TIME_ARRAY, -1, false), ColumnSchema::new( "timestamp_array_col".to_string(), Type::TIMESTAMP_ARRAY, -1, - true, + false, ), ColumnSchema::new( "timestamptz_array_col".to_string(), Type::TIMESTAMPTZ_ARRAY, -1, - true, + false, ), // UUID array type - ColumnSchema::new("uuid_array_col".to_string(), Type::UUID_ARRAY, -1, true), + ColumnSchema::new("uuid_array_col".to_string(), Type::UUID_ARRAY, -1, false), // JSON array types - ColumnSchema::new("json_array_col".to_string(), Type::JSON_ARRAY, -1, true), - ColumnSchema::new("jsonb_array_col".to_string(), Type::JSONB_ARRAY, -1, true), + ColumnSchema::new("json_array_col".to_string(), Type::JSON_ARRAY, -1, false), + ColumnSchema::new("jsonb_array_col".to_string(), Type::JSONB_ARRAY, -1, false), // OID array type - ColumnSchema::new("oid_array_col".to_string(), Type::OID_ARRAY, -1, true), + ColumnSchema::new("oid_array_col".to_string(), Type::OID_ARRAY, -1, false), // Binary array type - ColumnSchema::new("bytea_array_col".to_string(), Type::BYTEA_ARRAY, -1, true), + ColumnSchema::new("bytea_array_col".to_string(), Type::BYTEA_ARRAY, -1, false), ]; let table_schema = TableSchema::new(table_id, table_name_struct, columns); @@ -805,7 +835,7 @@ async fn insert_non_nullable_array() { let table_name_struct = TableName::new("test_schema".to_string(), table_name.clone()); let columns = vec![ // Primary key - ColumnSchema::new("id".to_string(), Type::INT4, -1, false), + ColumnSchema::new("id".to_string(), Type::INT4, -1, true), // Boolean array type ColumnSchema::new("bool_array_col".to_string(), Type::BOOL_ARRAY, -1, false), // String array types diff --git a/etl/src/conversions/event.rs b/etl/src/conversions/event.rs index 4b28705f1..592a70e8b 100644 --- a/etl/src/conversions/event.rs +++ b/etl/src/conversions/event.rs @@ -2,8 +2,7 @@ use core::str; use etl_postgres::types::{ColumnSchema, TableId, TableSchema, convert_type_oid_to_type}; use postgres_replication::protocol; use std::cmp::Ordering; -use std::collections::{BTreeSet, HashSet}; -use std::hash::Hash; +use std::collections::BTreeSet; use std::sync::Arc; use tokio_postgres::types::PgLsn; @@ -45,12 +44,6 @@ impl PartialOrd for IndexedColumnSchema { } } -impl Hash for IndexedColumnSchema { - fn hash(&self, state: &mut H) { - self.0.name.hash(state); - } -} - /// Creates a [`BeginEvent`] from Postgres protocol data. /// /// This method parses the replication protocol begin message and extracts From 15ffe5ec71f65c0d58b0de0c2a48e8836d4cba92 Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Tue, 30 Sep 2025 12:04:28 -0700 Subject: [PATCH 09/45] Improve --- etl/tests/pipeline.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/etl/tests/pipeline.rs b/etl/tests/pipeline.rs index 9c611f07e..855981b3a 100644 --- a/etl/tests/pipeline.rs +++ b/etl/tests/pipeline.rs @@ -529,7 +529,7 @@ async fn table_schema_changes_are_handled_correctly() { let users_table_schema = table_schemas .get(&database_schema.users_schema().id) .unwrap(); - println!("{:?}", users_table_schema); + println!("{users_table_schema:?}"); // Check the initial data. let table_rows = destination.get_table_rows().await; @@ -577,7 +577,7 @@ async fn table_schema_changes_are_handled_correctly() { let users_table_schema = table_schemas .get(&database_schema.users_schema().id) .unwrap(); - println!("{:?}", users_table_schema); + println!("{users_table_schema:?}"); // Check the updated data. let table_rows = destination.get_table_rows().await; From e2ef40cc44aa9d87b3e26aecf206e21af024ec83 Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Tue, 30 Sep 2025 12:58:00 -0700 Subject: [PATCH 10/45] Improve --- etl-destinations/src/bigquery/client.rs | 169 ++++++++++++++++++++++ etl-destinations/src/bigquery/core.rs | 185 +++++++++++++++++++++++- etl/Cargo.toml | 1 + etl/src/conversions/event.rs | 5 +- etl/src/pipeline.rs | 6 +- etl/tests/pipeline.rs | 154 ++++++++++---------- 6 files changed, 437 insertions(+), 83 deletions(-) diff --git a/etl-destinations/src/bigquery/client.rs b/etl-destinations/src/bigquery/client.rs index 96d8b6cf6..9d0168f22 100644 --- a/etl-destinations/src/bigquery/client.rs +++ b/etl-destinations/src/bigquery/client.rs @@ -259,6 +259,175 @@ impl BigQueryClient { Ok(()) } + /// Adds a column to an existing BigQuery table if it does not already exist. + pub async fn add_column( + &self, + dataset_id: &BigQueryDatasetId, + table_id: &BigQueryTableId, + column_schema: &ColumnSchema, + ) -> EtlResult<()> { + let full_table_name = self.full_table_name(dataset_id, table_id); + let column_definition = Self::column_spec(column_schema); + let query = + format!("alter table {full_table_name} add column if not exists {column_definition}"); + + let _ = self.query(QueryRequest::new(query)).await?; + + Ok(()) + } + + /// Drops a column from an existing BigQuery table if it exists. + pub async fn drop_column( + &self, + dataset_id: &BigQueryDatasetId, + table_id: &BigQueryTableId, + column_name: &str, + ) -> EtlResult<()> { + let full_table_name = self.full_table_name(dataset_id, table_id); + let column_identifier = column_name; + let query = + format!("alter table {full_table_name} drop column if exists {column_identifier}"); + + let _ = self.query(QueryRequest::new(query)).await?; + + Ok(()) + } + + /// Renames a column in an existing BigQuery table. + pub async fn rename_column( + &self, + dataset_id: &BigQueryDatasetId, + table_id: &BigQueryTableId, + old_name: &str, + new_name: &str, + ) -> EtlResult<()> { + let full_table_name = self.full_table_name(dataset_id, table_id); + let old_identifier = old_name; + let new_identifier = new_name; + let query = format!( + "alter table {full_table_name} rename column {old_identifier} to {new_identifier}" + ); + + let _ = self.query(QueryRequest::new(query)).await?; + + Ok(()) + } + + /// Alters the data type of an existing column in a BigQuery table. + pub async fn alter_column_type( + &self, + dataset_id: &BigQueryDatasetId, + table_id: &BigQueryTableId, + column_schema: &ColumnSchema, + ) -> EtlResult<()> { + let full_table_name = self.full_table_name(dataset_id, table_id); + let column_identifier = &column_schema.name; + let column_type = Self::postgres_to_bigquery_type(&column_schema.typ); + let query = format!( + "alter table {full_table_name} alter column {column_identifier} set data type {column_type}" + ); + + let _ = self.query(QueryRequest::new(query)).await?; + + Ok(()) + } + + /// Updates the nullability of an existing column in a BigQuery table. + pub async fn alter_column_nullability( + &self, + dataset_id: &BigQueryDatasetId, + table_id: &BigQueryTableId, + column_name: &str, + nullable: bool, + ) -> EtlResult<()> { + let full_table_name = self.full_table_name(dataset_id, table_id); + let column_identifier = column_name; + let clause = if nullable { + "drop not null" + } else { + "set not null" + }; + let query = + format!("alter table {full_table_name} alter column {column_identifier} {clause}"); + + let _ = self.query(QueryRequest::new(query)).await?; + + Ok(()) + } + + /// Synchronizes the primary key definition for a BigQuery table with the provided schema. + pub async fn sync_primary_key( + &self, + dataset_id: &BigQueryDatasetId, + table_id: &BigQueryTableId, + column_schemas: &[ColumnSchema], + ) -> EtlResult<()> { + let primary_columns: Vec<&ColumnSchema> = column_schemas + .iter() + .filter(|column| column.primary) + .collect(); + + let has_primary_key = self.has_primary_key(dataset_id, table_id).await?; + + if primary_columns.is_empty() { + if has_primary_key { + self.drop_primary_key(dataset_id, table_id).await?; + } + return Ok(()); + } + + if has_primary_key { + self.drop_primary_key(dataset_id, table_id).await?; + } + + let columns = primary_columns + .iter() + .map(|column| column.name) + .collect::>() + .join(","); + + let full_table_name = self.full_table_name(dataset_id, table_id); + let query = + format!("alter table {full_table_name} add primary key ({columns}) not enforced"); + + let _ = self.query(QueryRequest::new(query)).await?; + + Ok(()) + } + + async fn has_primary_key( + &self, + dataset_id: &BigQueryDatasetId, + table_id: &BigQueryTableId, + ) -> EtlResult { + let info_schema_table = format!( + "`{}.{}`.INFORMATION_SCHEMA.TABLE_CONSTRAINTS`", + &self.project_id, + dataset_id + ); + let table_literal = table_id; + let query = format!( + "select constraint_name from {info_schema_table} where table_name = '{table_literal}' and constraint_type = 'PRIMARY KEY'", + ); + + let result_set = self.query(QueryRequest::new(query)).await?; + + Ok(result_set.row_count() > 0) + } + + async fn drop_primary_key( + &self, + dataset_id: &BigQueryDatasetId, + table_id: &BigQueryTableId, + ) -> EtlResult<()> { + let full_table_name = self.full_table_name(dataset_id, table_id); + let query = format!("alter table {full_table_name} drop primary key"); + + let _ = self.query(QueryRequest::new(query)).await?; + + Ok(()) + } + /// Checks whether a table exists in the BigQuery dataset. /// /// Returns `true` if the table exists, `false` otherwise. diff --git a/etl-destinations/src/bigquery/core.rs b/etl-destinations/src/bigquery/core.rs index 27b6fc14b..923193a4a 100644 --- a/etl-destinations/src/bigquery/core.rs +++ b/etl-destinations/src/bigquery/core.rs @@ -2,7 +2,7 @@ use etl::destination::Destination; use etl::error::{ErrorKind, EtlError, EtlResult}; use etl::store::schema::SchemaStore; use etl::store::state::StateStore; -use etl::types::{Cell, Event, PgLsn, TableId, TableName, TableRow}; +use etl::types::{Cell, Event, PgLsn, RelationChange, RelationEvent, TableId, TableName, TableRow}; use etl::{bail, etl_error}; use gcp_bigquery_client::storage::TableDescriptor; use std::collections::{HashMap, HashSet}; @@ -477,6 +477,186 @@ where Ok(()) } + async fn apply_relation_event(&self, relation_event: RelationEvent) -> EtlResult<()> { + if relation_event.changes.is_empty() { + debug!( + table_id = %relation_event.table_id, + "relation event contained no schema changes; skipping" + ); + + return Ok(()); + } + + let Some(table_schema) = self + .store + .get_table_schema(&relation_event.table_id) + .await? + else { + bail!( + ErrorKind::MissingTableSchema, + "Table not found in the schema store", + format!( + "The table schema for table {} was not found while applying relation changes to BigQuery", + relation_event.table_id + ) + ); + }; + + let bigquery_table_id = table_name_to_bigquery_table_id(&table_schema.name); + let sequenced_bigquery_table_id = self + .get_or_create_sequenced_bigquery_table_id(&relation_event.table_id, &bigquery_table_id) + .await?; + + { + let mut inner = self.inner.lock().await; + if !inner.created_tables.contains(&sequenced_bigquery_table_id) { + self.client + .create_table_if_missing( + &self.dataset_id, + &sequenced_bigquery_table_id.to_string(), + &table_schema.column_schemas, + self.max_staleness_mins, + ) + .await?; + + Self::add_to_created_tables_cache(&mut inner, &sequenced_bigquery_table_id); + + self.ensure_view_points_to_table( + &mut inner, + &bigquery_table_id, + &sequenced_bigquery_table_id, + ) + .await?; + } + } + + let sequenced_table_name = sequenced_bigquery_table_id.to_string(); + let mut primary_key_dirty = false; + + for change in relation_event.changes { + match change { + RelationChange::AddColumn(column_schema) => { + let column_name = column_schema.name.clone(); + let is_primary = column_schema.primary; + + self.client + .add_column(&self.dataset_id, &sequenced_table_name, &column_schema) + .await?; + + debug!( + table = %sequenced_table_name, + column = %column_name, + "added column in BigQuery" + ); + + if is_primary { + primary_key_dirty = true; + } + } + RelationChange::DropColumn(column_schema) => { + let column_name = column_schema.name.clone(); + let was_primary = column_schema.primary; + + self.client + .drop_column(&self.dataset_id, &sequenced_table_name, &column_schema.name) + .await?; + + debug!( + table = %sequenced_table_name, + column = %column_name, + "dropped column in BigQuery" + ); + + if was_primary { + primary_key_dirty = true; + } + } + RelationChange::AlterColumn(previous, latest) => { + let old_name = previous.name.clone(); + let new_name = latest.name.clone(); + let renamed = old_name != new_name; + + if renamed { + self.client + .rename_column( + &self.dataset_id, + &sequenced_table_name, + &previous.name, + &latest.name, + ) + .await?; + + debug!( + table = %sequenced_table_name, + old_column = %old_name, + new_column = %new_name, + "renamed column in BigQuery" + ); + } + + if previous.typ != latest.typ { + self.client + .alter_column_type(&self.dataset_id, &sequenced_table_name, &latest) + .await?; + + debug!( + table = %sequenced_table_name, + column = %new_name, + "updated column type in BigQuery" + ); + } + + if previous.nullable != latest.nullable { + self.client + .alter_column_nullability( + &self.dataset_id, + &sequenced_table_name, + &latest.name, + latest.nullable, + ) + .await?; + + debug!( + table = %sequenced_table_name, + column = %new_name, + nullable = latest.nullable, + "updated column nullability in BigQuery" + ); + } + + if previous.primary != latest.primary + || (renamed && (previous.primary || latest.primary)) + { + primary_key_dirty = true; + } + } + } + } + + if primary_key_dirty { + self.client + .sync_primary_key( + &self.dataset_id, + &sequenced_table_name, + &table_schema.column_schemas, + ) + .await?; + + debug!( + table = %sequenced_table_name, + "synchronized primary key definition in BigQuery" + ); + } + + info!( + table_id = %relation_event.table_id, + table = %sequenced_table_name, + "applied relation changes in BigQuery" + ); + + Ok(()) + } + /// Processes CDC events in batches with proper ordering and truncate handling. /// /// Groups streaming operations (insert/update/delete) by table and processes them together, @@ -538,6 +718,9 @@ where table_id_to_table_rows.entry(delete.table_id).or_default(); table_rows.push(old_table_row); } + Event::Relation(relation_event) => { + self.apply_relation_event(relation_event).await?; + } _ => { // Every other event type is currently not supported. debug!("skipping unsupported event in BigQuery"); diff --git a/etl/Cargo.toml b/etl/Cargo.toml index b6b94eff2..8f758e07a 100644 --- a/etl/Cargo.toml +++ b/etl/Cargo.toml @@ -43,6 +43,7 @@ tokio-rustls = { workspace = true, default-features = false } tracing = { workspace = true, default-features = true } uuid = { workspace = true, features = ["v4"] } x509-cert = { workspace = true, default-features = false } +log = "0.4.28" [dev-dependencies] etl-postgres = { workspace = true, features = [ diff --git a/etl/src/conversions/event.rs b/etl/src/conversions/event.rs index 592a70e8b..93b6ee29a 100644 --- a/etl/src/conversions/event.rs +++ b/etl/src/conversions/event.rs @@ -5,6 +5,7 @@ use std::cmp::Ordering; use std::collections::BTreeSet; use std::sync::Arc; use tokio_postgres::types::PgLsn; +use tracing::info; use crate::conversions::text::{default_value_for_type, parse_cell_from_postgres_text}; use crate::error::{ErrorKind, EtlError, EtlResult}; @@ -380,7 +381,7 @@ pub fn convert_tuple_to_row( } else if use_default_for_missing_cols { default_value_for_type(&column_schema.typ)? } else { - // This is protocol level error, so we panic instead of carrying on + // This is a protocol level error, so we panic instead of carrying on // with incorrect data to avoid corruption downstream. panic!( "A required column {} was missing from the tuple", @@ -391,7 +392,7 @@ pub fn convert_tuple_to_row( protocol::TupleData::UnchangedToast => { // For unchanged toast values we try to use the value from the old row if it is present // but only if it is not null. In all other cases we send the default value for - // consistency. As a bit of a practical hack we take the value out of the old row and + // consistency. As a bit of a practical hack, we take the value out of the old row and // move a null value in its place to avoid a clone because toast values tend to be large. if let Some(row) = old_table_row { let old_row_value = std::mem::replace(&mut row.values[i], Cell::Null); diff --git a/etl/src/pipeline.rs b/etl/src/pipeline.rs index e021886d4..d84a582ea 100644 --- a/etl/src/pipeline.rs +++ b/etl/src/pipeline.rs @@ -22,7 +22,7 @@ use etl_postgres::types::TableId; use std::collections::HashSet; use std::sync::Arc; use tokio::sync::Semaphore; -use tracing::{error, info}; +use tracing::{error, info, warn}; /// Internal state tracking for pipeline lifecycle. /// @@ -204,7 +204,7 @@ where // it means that no table sync workers are running, which is fine. let _ = self.shutdown_tx.shutdown(); - info!("apply worker completed with an error, shutting down table sync workers"); + warn!("apply worker completed with an error, shutting down table sync workers"); } info!("waiting for table sync workers to complete"); @@ -217,7 +217,7 @@ where errors.push(err); - info!("{} table sync workers failed with an error", errors_number); + warn!("{} table sync workers failed with an error", errors_number); } if !errors.is_empty() { diff --git a/etl/tests/pipeline.rs b/etl/tests/pipeline.rs index 855981b3a..311bdca5b 100644 --- a/etl/tests/pipeline.rs +++ b/etl/tests/pipeline.rs @@ -20,6 +20,7 @@ use etl_postgres::tokio::test_utils::TableModification; use etl_telemetry::tracing::init_test_tracing; use rand::random; use std::time::Duration; +use sqlx::__rt::timeout; use tokio::time::sleep; #[tokio::test(flavor = "multi_thread")] @@ -717,84 +718,83 @@ async fn table_copy_and_sync_streams_new_data() { ) .await; - users_state_notify.notified().await; - orders_state_notify.notified().await; - events_notify.notified().await; - + // users_state_notify.notified().await; + // orders_state_notify.notified().await; + timeout(Duration::from_secs(2), events_notify.notified()).await; pipeline.shutdown_and_wait().await.unwrap(); - - // Verify initial table copy data. - let table_rows = destination.get_table_rows().await; - let users_table_rows = table_rows.get(&database_schema.users_schema().id).unwrap(); - let orders_table_rows = table_rows.get(&database_schema.orders_schema().id).unwrap(); - assert_eq!(users_table_rows.len(), rows_inserted); - assert_eq!(orders_table_rows.len(), rows_inserted); - - // Verify age sum calculation. - let expected_age_sum = get_n_integers_sum(rows_inserted); - let age_sum = - get_users_age_sum_from_rows(&destination, database_schema.users_schema().id).await; - assert_eq!(age_sum, expected_age_sum); - - // Get all the events that were produced to the destination and assert them individually by table - // since the only thing we are guaranteed is that the order of operations is preserved within the - // same table but not across tables given the asynchronous nature of the pipeline (e.g., we could - // start streaming earlier on a table for data which was inserted after another table which was - // modified before this one) - let events = destination.get_events().await; - let grouped_events = group_events_by_type_and_table_id(&events); - let users_inserts = grouped_events - .get(&(EventType::Insert, database_schema.users_schema().id)) - .unwrap(); - let orders_inserts = grouped_events - .get(&(EventType::Insert, database_schema.orders_schema().id)) - .unwrap(); - - // Build expected events for verification - let expected_users_inserts = build_expected_users_inserts( - 11, - database_schema.users_schema().id, - vec![ - ("user_11", 11), - ("user_12", 12), - ("user_13", 13), - ("user_14", 14), - ], - ); - let expected_orders_inserts = build_expected_orders_inserts( - 11, - database_schema.orders_schema().id, - vec![ - "description_11", - "description_12", - "description_13", - "description_14", - ], - ); - assert_events_equal(users_inserts, &expected_users_inserts); - assert_events_equal(orders_inserts, &expected_orders_inserts); - - // Check that the replication slots for the two tables have been removed. - let users_replication_slot: String = - EtlReplicationSlot::for_table_sync_worker(pipeline_id, database_schema.users_schema().id) - .try_into() - .unwrap(); - let orders_replication_slot: String = - EtlReplicationSlot::for_table_sync_worker(pipeline_id, database_schema.orders_schema().id) - .try_into() - .unwrap(); - assert!( - !database - .replication_slot_exists(&users_replication_slot) - .await - .unwrap() - ); - assert!( - !database - .replication_slot_exists(&orders_replication_slot) - .await - .unwrap() - ); + // + // // Verify initial table copy data. + // let table_rows = destination.get_table_rows().await; + // let users_table_rows = table_rows.get(&database_schema.users_schema().id).unwrap(); + // let orders_table_rows = table_rows.get(&database_schema.orders_schema().id).unwrap(); + // assert_eq!(users_table_rows.len(), rows_inserted); + // assert_eq!(orders_table_rows.len(), rows_inserted); + // + // // Verify age sum calculation. + // let expected_age_sum = get_n_integers_sum(rows_inserted); + // let age_sum = + // get_users_age_sum_from_rows(&destination, database_schema.users_schema().id).await; + // assert_eq!(age_sum, expected_age_sum); + // + // // Get all the events that were produced to the destination and assert them individually by table + // // since the only thing we are guaranteed is that the order of operations is preserved within the + // // same table but not across tables given the asynchronous nature of the pipeline (e.g., we could + // // start streaming earlier on a table for data which was inserted after another table which was + // // modified before this one) + // let events = destination.get_events().await; + // let grouped_events = group_events_by_type_and_table_id(&events); + // let users_inserts = grouped_events + // .get(&(EventType::Insert, database_schema.users_schema().id)) + // .unwrap(); + // let orders_inserts = grouped_events + // .get(&(EventType::Insert, database_schema.orders_schema().id)) + // .unwrap(); + // + // // Build expected events for verification + // let expected_users_inserts = build_expected_users_inserts( + // 11, + // database_schema.users_schema().id, + // vec![ + // ("user_11", 11), + // ("user_12", 12), + // ("user_13", 13), + // ("user_14", 14), + // ], + // ); + // let expected_orders_inserts = build_expected_orders_inserts( + // 11, + // database_schema.orders_schema().id, + // vec![ + // "description_11", + // "description_12", + // "description_13", + // "description_14", + // ], + // ); + // assert_events_equal(users_inserts, &expected_users_inserts); + // assert_events_equal(orders_inserts, &expected_orders_inserts); + // + // // Check that the replication slots for the two tables have been removed. + // let users_replication_slot: String = + // EtlReplicationSlot::for_table_sync_worker(pipeline_id, database_schema.users_schema().id) + // .try_into() + // .unwrap(); + // let orders_replication_slot: String = + // EtlReplicationSlot::for_table_sync_worker(pipeline_id, database_schema.orders_schema().id) + // .try_into() + // .unwrap(); + // assert!( + // !database + // .replication_slot_exists(&users_replication_slot) + // .await + // .unwrap() + // ); + // assert!( + // !database + // .replication_slot_exists(&orders_replication_slot) + // .await + // .unwrap() + // ); } #[tokio::test(flavor = "multi_thread")] From 4d4177f36b907485fe1feeb3a201d65ca938cb49 Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Tue, 30 Sep 2025 13:17:39 -0700 Subject: [PATCH 11/45] Improve --- etl-destinations/src/bigquery/client.rs | 5 +- etl/src/conversions/event.rs | 47 ++++---- etl/src/error.rs | 1 + etl/src/test_utils/notify.rs | 4 +- etl/tests/pipeline.rs | 154 ++++++++++++------------ 5 files changed, 106 insertions(+), 105 deletions(-) diff --git a/etl-destinations/src/bigquery/client.rs b/etl-destinations/src/bigquery/client.rs index 9d0168f22..7af9baddf 100644 --- a/etl-destinations/src/bigquery/client.rs +++ b/etl-destinations/src/bigquery/client.rs @@ -382,7 +382,7 @@ impl BigQueryClient { let columns = primary_columns .iter() - .map(|column| column.name) + .map(|column| format!("`{}`", column.name)) .collect::>() .join(","); @@ -402,8 +402,7 @@ impl BigQueryClient { ) -> EtlResult { let info_schema_table = format!( "`{}.{}`.INFORMATION_SCHEMA.TABLE_CONSTRAINTS`", - &self.project_id, - dataset_id + &self.project_id, dataset_id ); let table_literal = table_id; let query = format!( diff --git a/etl/src/conversions/event.rs b/etl/src/conversions/event.rs index 93b6ee29a..9223f1f06 100644 --- a/etl/src/conversions/event.rs +++ b/etl/src/conversions/event.rs @@ -5,7 +5,6 @@ use std::cmp::Ordering; use std::collections::BTreeSet; use std::sync::Arc; use tokio_postgres::types::PgLsn; -use tracing::info; use crate::conversions::text::{default_value_for_type, parse_cell_from_postgres_text}; use crate::error::{ErrorKind, EtlError, EtlResult}; @@ -103,32 +102,26 @@ where ) }; - // We build a set of the new column chemas from the relation message. The rationale for using a - // BTreeSet is that we want to preserve the order of columns in the schema, which is important, - // and also we can reuse the same set to build the vector of column schemas needed for the table - // schema. - let mut latest_column_schemas = relation_body + // We construct the latest column schemas in order. The order is important since the table schema + // relies on the right ordering to interpret the Postgres correctly. + let latest_column_schemas = relation_body .columns() .iter() - .map(build_indexed_column_schema) - .collect::, EtlError>>()?; + .map(build_column_schema) + .collect::, EtlError>>()?; - // We build the updated table schema to store in the schema store. - let mut latest_table_schema = TableSchema::new( - table_schema.id, - table_schema.name.clone(), - Vec::with_capacity(latest_column_schemas.len()), - ); - for column_schema in latest_column_schemas.iter() { - latest_table_schema.add_column_schema(column_schema.clone().into_inner()); - } - schema_store.store_table_schema(latest_table_schema).await?; + // We build a lookup set for the latest column schemas for quick change detection. + let mut latest_indexed_column_schemas = latest_column_schemas + .iter() + .cloned() + .map(IndexedColumnSchema) + .collect::>(); // We process all the changes that we want to dispatch to the destination. let mut changes = vec![]; for column_schema in table_schema.column_schemas.iter() { let column_schema = IndexedColumnSchema(column_schema.clone()); - let latest_column_schema = latest_column_schemas.take(&column_schema); + let latest_column_schema = latest_indexed_column_schemas.take(&column_schema); match latest_column_schema { Some(latest_column_schema) => { let column_schema = column_schema.into_inner(); @@ -152,10 +145,19 @@ where } // For the remaining columns that didn't match, we assume they were added. - for column_schema in latest_column_schemas { + for column_schema in latest_indexed_column_schemas { changes.push(RelationChange::AddColumn(column_schema.into_inner())); } + // We build the updated table schema to store in the schema store. + schema_store + .store_table_schema(TableSchema::new( + table_schema.id, + table_schema.name.clone(), + latest_column_schemas, + )) + .await?; + Ok(RelationEvent { start_lsn, commit_lsn, @@ -181,6 +183,7 @@ where let table_id = insert_body.rel_id(); let table_schema = get_table_schema(schema_store, TableId::new(table_id)).await?; + println!("INSERT DATA {:?}", insert_body.tuple().tuple_data()); let table_row = convert_tuple_to_row( &table_schema.column_schemas, insert_body.tuple().tuple_data(), @@ -332,7 +335,7 @@ where /// This helper method extracts column metadata from the replication protocol /// and converts it into the internal column schema representation. Some fields /// like nullable status have default values due to protocol limitations. -fn build_indexed_column_schema(column: &protocol::Column) -> EtlResult { +fn build_column_schema(column: &protocol::Column) -> EtlResult { let column_schema = ColumnSchema::new( column.name()?.to_string(), convert_type_oid_to_type(column.type_id() as u32), @@ -341,7 +344,7 @@ fn build_indexed_column_schema(column: &protocol::Column) -> EtlResult Date: Tue, 30 Sep 2025 13:19:55 -0700 Subject: [PATCH 12/45] Improve --- etl/src/conversions/event.rs | 1 - etl/src/error.rs | 1 - etl/src/test_utils/notify.rs | 1 - etl/tests/pipeline.rs | 4 +++- 4 files changed, 3 insertions(+), 4 deletions(-) diff --git a/etl/src/conversions/event.rs b/etl/src/conversions/event.rs index 9223f1f06..1b046b750 100644 --- a/etl/src/conversions/event.rs +++ b/etl/src/conversions/event.rs @@ -183,7 +183,6 @@ where let table_id = insert_body.rel_id(); let table_schema = get_table_schema(schema_store, TableId::new(table_id)).await?; - println!("INSERT DATA {:?}", insert_body.tuple().tuple_data()); let table_row = convert_tuple_to_row( &table_schema.column_schemas, insert_body.tuple().tuple_data(), diff --git a/etl/src/error.rs b/etl/src/error.rs index 434ca9575..02db592ae 100644 --- a/etl/src/error.rs +++ b/etl/src/error.rs @@ -4,7 +4,6 @@ //! for ETL pipeline operations. The [`EtlError`] type supports single errors, errors with additional detail, //! and multiple aggregated errors for complex failure scenarios. -use std::backtrace::Backtrace; use std::error; use std::fmt; diff --git a/etl/src/test_utils/notify.rs b/etl/src/test_utils/notify.rs index 7a720db50..09fff96b3 100644 --- a/etl/src/test_utils/notify.rs +++ b/etl/src/test_utils/notify.rs @@ -1,5 +1,4 @@ use etl_postgres::types::{TableId, TableSchema}; -use log::info; use std::{collections::HashMap, fmt, sync::Arc}; use tokio::sync::{Notify, RwLock}; diff --git a/etl/tests/pipeline.rs b/etl/tests/pipeline.rs index 2a2b4294d..855981b3a 100644 --- a/etl/tests/pipeline.rs +++ b/etl/tests/pipeline.rs @@ -19,7 +19,6 @@ use etl_postgres::replication::slots::EtlReplicationSlot; use etl_postgres::tokio::test_utils::TableModification; use etl_telemetry::tracing::init_test_tracing; use rand::random; -use sqlx::__rt::timeout; use std::time::Duration; use tokio::time::sleep; @@ -720,6 +719,9 @@ async fn table_copy_and_sync_streams_new_data() { users_state_notify.notified().await; orders_state_notify.notified().await; + events_notify.notified().await; + + pipeline.shutdown_and_wait().await.unwrap(); // Verify initial table copy data. let table_rows = destination.get_table_rows().await; From 7d6fd102de67a4bdf0b6d6c3c29ee7796f1c79df Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Tue, 30 Sep 2025 15:24:48 -0700 Subject: [PATCH 13/45] Improve --- etl/src/test_utils/table.rs | 9 ++ etl/tests/pipeline.rs | 264 ++++++++++-------------------------- 2 files changed, 79 insertions(+), 194 deletions(-) diff --git a/etl/src/test_utils/table.rs b/etl/src/test_utils/table.rs index a2c6d46ad..50f175fed 100644 --- a/etl/src/test_utils/table.rs +++ b/etl/src/test_utils/table.rs @@ -1,6 +1,15 @@ use etl_postgres::types::{ColumnSchema, TableId, TableName, TableSchema}; use std::collections::HashMap; +/// Return the names of the column schema. +pub fn column_schema_names(table_schema: &TableSchema) -> Vec { + table_schema + .column_schemas + .iter() + .map(|c| c.name.clone()) + .collect() +} + /// Asserts that a table schema matches the expected schema. /// /// Compares all aspects of the table schema including table ID, name, and column diff --git a/etl/tests/pipeline.rs b/etl/tests/pipeline.rs index 855981b3a..dd683925b 100644 --- a/etl/tests/pipeline.rs +++ b/etl/tests/pipeline.rs @@ -7,6 +7,7 @@ use etl::test_utils::database::{spawn_source_database, test_table_name}; use etl::test_utils::event::group_events_by_type_and_table_id; use etl::test_utils::notify::NotifyingStore; use etl::test_utils::pipeline::{create_pipeline, create_pipeline_with}; +use etl::test_utils::table::column_schema_names; use etl::test_utils::test_destination_wrapper::TestDestinationWrapper; use etl::test_utils::test_schema::{ TableSelection, assert_events_equal, build_expected_orders_inserts, @@ -526,10 +527,12 @@ async fn table_schema_changes_are_handled_correctly() { // Check the initial schema. let table_schemas = store.get_table_schemas().await; assert_eq!(table_schemas.len(), 1); - let users_table_schema = table_schemas - .get(&database_schema.users_schema().id) - .unwrap(); - println!("{users_table_schema:?}"); + let users_table_schema = column_schema_names( + table_schemas + .get(&database_schema.users_schema().id) + .unwrap(), + ); + assert_eq!(users_table_schema, vec!["id", "name", "age"]); // Check the initial data. let table_rows = destination.get_table_rows().await; @@ -574,59 +577,72 @@ async fn table_schema_changes_are_handled_correctly() { // Check the updated schema. let table_schemas = store.get_table_schemas().await; assert_eq!(table_schemas.len(), 1); - let users_table_schema = table_schemas - .get(&database_schema.users_schema().id) + let users_table_schema = column_schema_names( + table_schemas + .get(&database_schema.users_schema().id) + .unwrap(), + ); + assert_eq!(users_table_schema, vec!["id", "name", "new_age", "year"]); + + // Check the updated data. + let events = destination.get_events().await; + let grouped_events = group_events_by_type_and_table_id(&events); + let users_inserts = grouped_events + .get(&(EventType::Insert, database_schema.users_schema().id)) .unwrap(); - println!("{users_table_schema:?}"); + assert_eq!(users_inserts.len(), 1); + + // We perform schema changes. + database + .alter_table( + test_table_name("users"), + &[ + TableModification::DropColumn { name: "year" }, + TableModification::AlterColumn { + name: "new_age", + params: "type double precision using new_age::double precision", + }, + ], + ) + .await + .unwrap(); + + // Register notifications for the insert. + let insert_event_notify = destination + .wait_for_events_count(vec![(EventType::Insert, 2)]) + .await; + + // We insert data. + database + .insert_values( + database_schema.users_schema().name.clone(), + &["name", "new_age"], + &[&"user_3", &(2f64)], + ) + .await + .expect("Failed to insert users"); + + insert_event_notify.notified().await; + + // Check the updated schema. + let table_schemas = store.get_table_schemas().await; + assert_eq!(table_schemas.len(), 1); + let users_table_schema = column_schema_names( + table_schemas + .get(&database_schema.users_schema().id) + .unwrap(), + ); + assert_eq!(users_table_schema, vec!["id", "name", "new_age"]); // Check the updated data. - let table_rows = destination.get_table_rows().await; - let users_table_rows = table_rows.get(&database_schema.users_schema().id).unwrap(); - assert_eq!(users_table_rows.len(), 1); - // - // // We perform schema changes. - // database - // .alter_table( - // test_table_name("users"), - // &[ - // TableModification::DropColumn { - // name: "year", - // }, - // TableModification::AlterColumn { - // name: "new_age", - // params: "type double precision using new_age::double precision", - // }, - // ], - // ) - // .await - // .unwrap(); - // - // // Register notifications for the insert. - // let insert_event_notify = destination.wait_for_events_count(vec![(EventType::Insert, 2)]).await; - // - // // We insert data. - // database - // .insert_values( - // database_schema.users_schema().name.clone(), - // &["name", "new_age", "year"], - // &[&"user_3", &(2i32), &(2025i32)], - // ) - // .await - // .expect("Failed to insert users"); - // - // insert_event_notify.notified().await; - // - // // Check the updated schema. - // let table_schemas = store.get_table_schemas().await; - // assert_eq!(table_schemas.len(), 1); - // assert!(table_schemas.contains_key(&database_schema.users_schema().id)); - // - // // Check the updated data. - // let table_rows = destination.get_table_rows().await; - // let users_table_rows = table_rows.get(&database_schema.users_schema().id).unwrap(); - // assert_eq!(users_table_rows.len(), 1); - // - // pipeline.shutdown_and_wait().await.unwrap(); + let events = destination.get_events().await; + let grouped_events = group_events_by_type_and_table_id(&events); + let users_inserts = grouped_events + .get(&(EventType::Insert, database_schema.users_schema().id)) + .unwrap(); + assert_eq!(users_inserts.len(), 2); + + pipeline.shutdown_and_wait().await.unwrap(); } #[tokio::test(flavor = "multi_thread")] @@ -934,146 +950,6 @@ async fn table_processing_converges_to_apply_loop_with_no_events_coming() { assert_eq!(age_sum, expected_age_sum); } -#[tokio::test(flavor = "multi_thread")] -async fn table_processing_with_schema_change_errors_table() { - init_test_tracing(); - let database = spawn_source_database().await; - let database_schema = setup_test_database_schema(&database, TableSelection::OrdersOnly).await; - - // Insert data in the table. - database - .insert_values( - database_schema.orders_schema().name.clone(), - &["description"], - &[&"description_1"], - ) - .await - .unwrap(); - - let store = NotifyingStore::new(); - let destination = TestDestinationWrapper::wrap(MemoryDestination::new()); - - // Start pipeline from scratch. - let pipeline_id: PipelineId = random(); - let mut pipeline = create_pipeline( - &database.config, - pipeline_id, - database_schema.publication_name(), - store.clone(), - destination.clone(), - ); - - // Register notifications for initial table copy completion. - let orders_state_notify = store - .notify_on_table_state_type( - database_schema.orders_schema().id, - TableReplicationPhaseType::FinishedCopy, - ) - .await; - - pipeline.start().await.unwrap(); - - orders_state_notify.notified().await; - - // Register notification for the sync done state. - let orders_state_notify = store - .notify_on_table_state_type( - database_schema.orders_schema().id, - TableReplicationPhaseType::SyncDone, - ) - .await; - - // Insert new data in the table. - database - .insert_values( - database_schema.orders_schema().name.clone(), - &["description"], - &[&"description_2"], - ) - .await - .unwrap(); - - orders_state_notify.notified().await; - - // Register notification for the ready state. - let orders_state_notify = store - .notify_on_table_state_type( - database_schema.orders_schema().id, - TableReplicationPhaseType::Ready, - ) - .await; - - // Insert new data in the table. - database - .insert_values( - database_schema.orders_schema().name.clone(), - &["description"], - &[&"description_3"], - ) - .await - .unwrap(); - - orders_state_notify.notified().await; - - // Register notification for the errored state. - let orders_state_notify = store - .notify_on_table_state_type( - database_schema.orders_schema().id, - TableReplicationPhaseType::Errored, - ) - .await; - - // Change the schema of orders by adding a new column. - database - .alter_table( - database_schema.orders_schema().name.clone(), - &[TableModification::AddColumn { - name: "date", - params: "integer", - }], - ) - .await - .unwrap(); - - // Insert new data in the table. - database - .insert_values( - database_schema.orders_schema().name.clone(), - &["description", "date"], - &[&"description_with_date", &10], - ) - .await - .unwrap(); - - orders_state_notify.notified().await; - - pipeline.shutdown_and_wait().await.unwrap(); - - // We assert that the schema is the initial one. - let table_schemas = store.get_table_schemas().await; - assert_eq!(table_schemas.len(), 1); - assert_eq!( - *table_schemas - .get(&database_schema.orders_schema().id) - .unwrap(), - database_schema.orders_schema() - ); - - // We check that we got the insert events after the first data of the table has been copied. - let events = destination.get_events().await; - let grouped_events = group_events_by_type_and_table_id(&events); - let orders_inserts = grouped_events - .get(&(EventType::Insert, database_schema.orders_schema().id)) - .unwrap(); - - let expected_orders_inserts = build_expected_orders_inserts( - 2, - database_schema.orders_schema().id, - vec!["description_2", "description_3"], - ); - assert_events_equal(orders_inserts, &expected_orders_inserts); -} - #[tokio::test(flavor = "multi_thread")] async fn table_without_primary_key_is_errored() { init_test_tracing(); From 8505998e98c343cd3d6059e9037435a5f10480f9 Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Tue, 30 Sep 2025 17:03:34 -0700 Subject: [PATCH 14/45] Improve --- etl-destinations/src/bigquery/client.rs | 4 +- etl-destinations/src/bigquery/core.rs | 177 +++++++++++------------- etl/src/types/event.rs | 9 +- 3 files changed, 93 insertions(+), 97 deletions(-) diff --git a/etl-destinations/src/bigquery/client.rs b/etl-destinations/src/bigquery/client.rs index 7af9baddf..f9a9bfcf7 100644 --- a/etl-destinations/src/bigquery/client.rs +++ b/etl-destinations/src/bigquery/client.rs @@ -569,7 +569,7 @@ impl BigQueryClient { /// Generates SQL column specification for CREATE TABLE statements. fn column_spec(column_schema: &ColumnSchema) -> String { let mut column_spec = format!( - "`{}` {}", + "{} {}", column_schema.name, Self::postgres_to_bigquery_type(&column_schema.typ) ); @@ -588,7 +588,7 @@ impl BigQueryClient { let identity_columns: Vec = column_schemas .iter() .filter(|s| s.primary) - .map(|c| format!("`{}`", c.name)) + .map(|c| c.name.clone()) .collect(); if identity_columns.is_empty() { diff --git a/etl-destinations/src/bigquery/core.rs b/etl-destinations/src/bigquery/core.rs index 923193a4a..1072efac2 100644 --- a/etl-destinations/src/bigquery/core.rs +++ b/etl-destinations/src/bigquery/core.rs @@ -2,12 +2,13 @@ use etl::destination::Destination; use etl::error::{ErrorKind, EtlError, EtlResult}; use etl::store::schema::SchemaStore; use etl::store::state::StateStore; -use etl::types::{Cell, Event, PgLsn, RelationChange, RelationEvent, TableId, TableName, TableRow}; +use etl::types::{Cell, Event, PgLsn, RelationChange, RelationEvent, TableId, TableName, TableRow, TableSchema}; use etl::{bail, etl_error}; use gcp_bigquery_client::storage::TableDescriptor; use std::collections::{HashMap, HashSet}; use std::fmt::Display; use std::iter; +use std::mem; use std::str::FromStr; use std::sync::Arc; use tokio::sync::Mutex; @@ -259,16 +260,11 @@ where }) } - /// Prepares a table for CDC streaming operations with schema-aware table creation. - /// - /// Retrieves the table schema from the store, creates or verifies the BigQuery table exists, - /// and ensures the view points to the current versioned table. Uses caching to avoid - /// redundant table creation checks. - async fn prepare_table_for_streaming( + /// Prepares a table for any operations. + async fn prepare_table( &self, table_id: &TableId, - use_cdc_sequence_column: bool, - ) -> EtlResult<(SequencedBigQueryTableId, Arc)> { + ) -> EtlResult<(SequencedBigQueryTableId, Arc)> { // We hold the lock for the entire preparation to avoid race conditions since the consistency // of this code path is critical. let mut inner = self.inner.lock().await; @@ -326,7 +322,18 @@ where &bigquery_table_id, &sequenced_bigquery_table_id, ) - .await?; + .await?; + + Ok((sequenced_bigquery_table_id, table_schema)) + } + + /// Prepares a table for CDC streaming operations with the table descriptor. + async fn prepare_table_for_streaming( + &self, + table_id: &TableId, + use_cdc_sequence_column: bool, + ) -> EtlResult<(SequencedBigQueryTableId, Arc)> { + let (sequenced_bigquery_table_id, table_schema) = self.prepare_table(table_id).await?; let table_descriptor = BigQueryClient::column_schemas_to_table_descriptor( &table_schema.column_schemas, @@ -477,6 +484,50 @@ where Ok(()) } + /// Persists the accumulated CDC batches for each table to BigQuery. + async fn process_table_events( + &self, + table_id_to_table_rows: HashMap>, + ) -> EtlResult<()> { + if table_id_to_table_rows.is_empty() { + return Ok(()); + } + + let mut table_batches = Vec::with_capacity(table_id_to_table_rows.len()); + for (table_id, table_rows) in table_id_to_table_rows { + let (sequenced_bigquery_table_id, table_descriptor) = + self.prepare_table_for_streaming(&table_id, true).await?; + + let table_batch = self.client.create_table_batch( + &self.dataset_id, + &sequenced_bigquery_table_id.to_string(), + table_descriptor.clone(), + table_rows, + )?; + table_batches.push(table_batch); + } + + if table_batches.is_empty() { + return Ok(()); + } + + let (bytes_sent, bytes_received) = self + .client + .stream_table_batches_concurrent(table_batches, self.max_concurrent_streams) + .await?; + + // Logs with egress_metric = true can be used to identify egress logs. + info!( + bytes_sent, + bytes_received, + phase = "apply", + egress_metric = true, + "wrote cdc events to bigquery" + ); + + Ok(()) + } + async fn apply_relation_event(&self, relation_event: RelationEvent) -> EtlResult<()> { if relation_event.changes.is_empty() { debug!( @@ -487,52 +538,10 @@ where return Ok(()); } - let Some(table_schema) = self - .store - .get_table_schema(&relation_event.table_id) - .await? - else { - bail!( - ErrorKind::MissingTableSchema, - "Table not found in the schema store", - format!( - "The table schema for table {} was not found while applying relation changes to BigQuery", - relation_event.table_id - ) - ); - }; - - let bigquery_table_id = table_name_to_bigquery_table_id(&table_schema.name); - let sequenced_bigquery_table_id = self - .get_or_create_sequenced_bigquery_table_id(&relation_event.table_id, &bigquery_table_id) - .await?; - - { - let mut inner = self.inner.lock().await; - if !inner.created_tables.contains(&sequenced_bigquery_table_id) { - self.client - .create_table_if_missing( - &self.dataset_id, - &sequenced_bigquery_table_id.to_string(), - &table_schema.column_schemas, - self.max_staleness_mins, - ) - .await?; - - Self::add_to_created_tables_cache(&mut inner, &sequenced_bigquery_table_id); - - self.ensure_view_points_to_table( - &mut inner, - &bigquery_table_id, - &sequenced_bigquery_table_id, - ) - .await?; - } - } + let (sequenced_bigquery_table_id, table_schema) = self.prepare_table(&relation_event.table_id).await?; let sequenced_table_name = sequenced_bigquery_table_id.to_string(); let mut primary_key_dirty = false; - for change in relation_event.changes { match change { RelationChange::AddColumn(column_schema) => { @@ -657,6 +666,16 @@ where Ok(()) } + /// Returns `true` whether the event should break the batch that is being built up when + /// writing events. `false` otherwise. + fn is_batch_breaker(event: &Event) -> bool { + match event { + Event::Truncate(_) => true, + Event::Relation(_) => true, + _ => false, + } + } + /// Processes CDC events in batches with proper ordering and truncate handling. /// /// Groups streaming operations (insert/update/delete) by table and processes them together, @@ -669,7 +688,7 @@ where // Process events until we hit a truncate event or run out of events while let Some(event) = event_iter.peek() { - if matches!(event, Event::Truncate(_)) { + if Self::is_batch_breaker(event) { break; } @@ -718,51 +737,16 @@ where table_id_to_table_rows.entry(delete.table_id).or_default(); table_rows.push(old_table_row); } - Event::Relation(relation_event) => { - self.apply_relation_event(relation_event).await?; - } _ => { - // Every other event type is currently not supported. debug!("skipping unsupported event in BigQuery"); } } } - // Process accumulated events for each table. + // Process accumulated events for each table before a batch breaker is encountered. if !table_id_to_table_rows.is_empty() { - let mut table_batches = Vec::with_capacity(table_id_to_table_rows.len()); - - for (table_id, table_rows) in table_id_to_table_rows { - let (sequenced_bigquery_table_id, table_descriptor) = - self.prepare_table_for_streaming(&table_id, true).await?; - - let table_batch = self.client.create_table_batch( - &self.dataset_id, - &sequenced_bigquery_table_id.to_string(), - table_descriptor.clone(), - table_rows, - )?; - table_batches.push(table_batch); - } - - if !table_batches.is_empty() { - let (bytes_sent, bytes_received) = self - .client - .stream_table_batches_concurrent(table_batches, self.max_concurrent_streams) - .await?; - - // Logs with egress_metric = true can be used to identify egress logs. - // This can e.g. be used to send egress logs to a location different - // than the other logs. These logs should also have bytes_sent set to - // the number of bytes sent to the destination. - info!( - bytes_sent, - bytes_received, - phase = "apply", - egress_metric = true, - "wrote cdc events to bigquery" - ); - } + let pending_rows = mem::take(&mut table_id_to_table_rows); + self.process_table_events(pending_rows).await?; } // Collect and deduplicate all table IDs from all truncate events. @@ -771,7 +755,6 @@ where // row without applying other events in the meanwhile, it doesn't make any sense to create // new empty tables for each of them. let mut truncate_table_ids = HashSet::new(); - while let Some(Event::Truncate(_)) = event_iter.peek() { if let Some(Event::Truncate(truncate_event)) = event_iter.next() { for table_id in truncate_event.rel_ids { @@ -779,11 +762,17 @@ where } } } - if !truncate_table_ids.is_empty() { self.process_truncate_for_table_ids(truncate_table_ids.into_iter(), true) .await?; } + + // Collect all the relation events in a sequence and apply them. + while let Some(Event::Relation(_)) = event_iter.peek() { + if let Some(Event::Relation(relation_event)) = event_iter.next() { + self.apply_relation_event(relation_event).await?; + } + } } Ok(()) diff --git a/etl/src/types/event.rs b/etl/src/types/event.rs index 4c469dc6f..44010552a 100644 --- a/etl/src/types/event.rs +++ b/etl/src/types/event.rs @@ -1,4 +1,4 @@ -use etl_postgres::types::{ColumnSchema, TableId}; +use etl_postgres::types::{ColumnSchema, TableId, TableName}; use std::fmt; use tokio_postgres::types::PgLsn; @@ -40,10 +40,17 @@ pub struct CommitEvent { pub timestamp: i64, } +/// A change in a relation. #[derive(Debug, Clone, PartialEq)] pub enum RelationChange { + /// A change that describes adding a new column. AddColumn(ColumnSchema), + /// A change that describes dropping an existing column. DropColumn(ColumnSchema), + /// A change that describes altering an existing column. + /// + /// The alteration of a column is defined as any modifications to a column + /// while keeping the same name. AlterColumn(ColumnSchema, ColumnSchema), } From 791257b90f06d235a323e254e95edd1e842afe51 Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Tue, 30 Sep 2025 21:19:20 -0700 Subject: [PATCH 15/45] Improve --- etl-destinations/src/bigquery/core.rs | 101 ++++++--- etl-destinations/tests/iceberg_pipeline.rs | 10 +- etl-postgres/src/replication/schema.rs | 72 +++--- etl-postgres/src/types/schema.rs | 65 +++++- .../20251001000100_add_schema_versions.sql | 21 ++ etl/src/conversions/event.rs | 209 ++++-------------- etl/src/replication/apply.rs | 90 ++++++-- etl/src/replication/client.rs | 12 +- etl/src/store/both/memory.rs | 62 +++++- etl/src/store/both/postgres.rs | 80 +++++-- etl/src/store/schema/base.rs | 17 +- etl/src/test_utils/notify.rs | 66 +++++- etl/src/test_utils/test_schema.rs | 2 + etl/src/types/event.rs | 99 ++++++++- etl/tests/postgres_store.rs | 4 +- 15 files changed, 589 insertions(+), 321 deletions(-) create mode 100644 etl-replicator/migrations/20251001000100_add_schema_versions.sql diff --git a/etl-destinations/src/bigquery/core.rs b/etl-destinations/src/bigquery/core.rs index 1072efac2..2ce00df47 100644 --- a/etl-destinations/src/bigquery/core.rs +++ b/etl-destinations/src/bigquery/core.rs @@ -2,7 +2,10 @@ use etl::destination::Destination; use etl::error::{ErrorKind, EtlError, EtlResult}; use etl::store::schema::SchemaStore; use etl::store::state::StateStore; -use etl::types::{Cell, Event, PgLsn, RelationChange, RelationEvent, TableId, TableName, TableRow, TableSchema}; +use etl::types::{ + Cell, Event, PgLsn, RelationChange, RelationEvent, SchemaVersion, TableId, TableName, TableRow, + TableSchema, +}; use etl::{bail, etl_error}; use gcp_bigquery_client::storage::TableDescriptor; use std::collections::{HashMap, HashSet}; @@ -264,6 +267,7 @@ where async fn prepare_table( &self, table_id: &TableId, + schema_version: SchemaVersion, ) -> EtlResult<(SequencedBigQueryTableId, Arc)> { // We hold the lock for the entire preparation to avoid race conditions since the consistency // of this code path is critical. @@ -273,14 +277,14 @@ where // and also prepare the table descriptor for CDC streaming. let table_schema = self .store - .get_table_schema(table_id) + .get_table_schema(table_id, schema_version) .await? .ok_or_else(|| { etl_error!( ErrorKind::MissingTableSchema, "Table not found in the schema store", format!( - "The table schema for table {table_id} was not found in the schema store" + "The table schema for table {table_id} version {schema_version} was not found in the schema store" ) ) })?; @@ -322,7 +326,7 @@ where &bigquery_table_id, &sequenced_bigquery_table_id, ) - .await?; + .await?; Ok((sequenced_bigquery_table_id, table_schema)) } @@ -331,9 +335,11 @@ where async fn prepare_table_for_streaming( &self, table_id: &TableId, + schema_version: SchemaVersion, use_cdc_sequence_column: bool, ) -> EtlResult<(SequencedBigQueryTableId, Arc)> { - let (sequenced_bigquery_table_id, table_schema) = self.prepare_table(table_id).await?; + let (sequenced_bigquery_table_id, table_schema) = + self.prepare_table(table_id, schema_version).await?; let table_descriptor = BigQueryClient::column_schemas_to_table_descriptor( &table_schema.column_schemas, @@ -434,8 +440,23 @@ where mut table_rows: Vec, ) -> EtlResult<()> { // Prepare table for streaming. - let (sequenced_bigquery_table_id, table_descriptor) = - self.prepare_table_for_streaming(&table_id, false).await?; + let latest_schema = self + .store + .get_latest_table_schema(&table_id) + .await? + .ok_or_else(|| { + etl_error!( + ErrorKind::MissingTableSchema, + "Table not found in the schema store", + format!( + "The table schema for table {table_id} was not found in the schema store" + ) + ) + })?; + + let (sequenced_bigquery_table_id, table_descriptor) = self + .prepare_table_for_streaming(&table_id, latest_schema.version, false) + .await?; // Add CDC operation type to all rows (no lock needed). for table_row in table_rows.iter_mut() { @@ -487,16 +508,17 @@ where /// Persists the accumulated CDC batches for each table to BigQuery. async fn process_table_events( &self, - table_id_to_table_rows: HashMap>, + table_id_to_table_rows: HashMap<(TableId, SchemaVersion), Vec>, ) -> EtlResult<()> { if table_id_to_table_rows.is_empty() { return Ok(()); } let mut table_batches = Vec::with_capacity(table_id_to_table_rows.len()); - for (table_id, table_rows) in table_id_to_table_rows { - let (sequenced_bigquery_table_id, table_descriptor) = - self.prepare_table_for_streaming(&table_id, true).await?; + for ((table_id, schema_version), table_rows) in table_id_to_table_rows { + let (sequenced_bigquery_table_id, table_descriptor) = self + .prepare_table_for_streaming(&table_id, schema_version, true) + .await?; let table_batch = self.client.create_table_batch( &self.dataset_id, @@ -538,7 +560,9 @@ where return Ok(()); } - let (sequenced_bigquery_table_id, table_schema) = self.prepare_table(&relation_event.table_id).await?; + let (sequenced_bigquery_table_id, table_schema) = self + .prepare_table(&relation_event.table_id, relation_event.new_schema_version) + .await?; let sequenced_table_name = sequenced_bigquery_table_id.to_string(); let mut primary_key_dirty = false; @@ -684,7 +708,8 @@ where let mut event_iter = events.into_iter().peekable(); while event_iter.peek().is_some() { - let mut table_id_to_table_rows = HashMap::new(); + let mut table_id_to_table_rows: HashMap<(TableId, SchemaVersion), Vec> = + HashMap::new(); // Process events until we hit a truncate event or run out of events while let Some(event) = event_iter.peek() { @@ -703,8 +728,9 @@ where .push(BigQueryOperationType::Upsert.into_cell()); insert.table_row.values.push(Cell::String(sequence_number)); - let table_rows: &mut Vec = - table_id_to_table_rows.entry(insert.table_id).or_default(); + let table_rows: &mut Vec = table_id_to_table_rows + .entry((insert.table_id, insert.schema_version)) + .or_default(); table_rows.push(insert.table_row); } Event::Update(mut update) => { @@ -716,8 +742,9 @@ where .push(BigQueryOperationType::Upsert.into_cell()); update.table_row.values.push(Cell::String(sequence_number)); - let table_rows: &mut Vec = - table_id_to_table_rows.entry(update.table_id).or_default(); + let table_rows: &mut Vec = table_id_to_table_rows + .entry((update.table_id, update.schema_version)) + .or_default(); table_rows.push(update.table_row); } Event::Delete(delete) => { @@ -733,8 +760,9 @@ where .push(BigQueryOperationType::Delete.into_cell()); old_table_row.values.push(Cell::String(sequence_number)); - let table_rows: &mut Vec = - table_id_to_table_rows.entry(delete.table_id).or_default(); + let table_rows: &mut Vec = table_id_to_table_rows + .entry((delete.table_id, delete.schema_version)) + .or_default(); table_rows.push(old_table_row); } _ => { @@ -754,16 +782,16 @@ where // This is done as an optimization since if we have multiple table ids being truncated in a // row without applying other events in the meanwhile, it doesn't make any sense to create // new empty tables for each of them. - let mut truncate_table_ids = HashSet::new(); + let mut truncate_tables = Vec::new(); while let Some(Event::Truncate(_)) = event_iter.peek() { if let Some(Event::Truncate(truncate_event)) = event_iter.next() { - for table_id in truncate_event.rel_ids { - truncate_table_ids.insert(TableId::new(table_id)); + for (table_id, schema_version) in truncate_event.relations { + truncate_tables.push((table_id, schema_version)); } } } - if !truncate_table_ids.is_empty() { - self.process_truncate_for_table_ids(truncate_table_ids.into_iter(), true) + if !truncate_tables.is_empty() { + self.process_truncate_for_table_ids(truncate_tables.into_iter(), true) .await?; } @@ -785,15 +813,18 @@ where /// to optimize multiple truncates of the same table. async fn process_truncate_for_table_ids( &self, - table_ids: impl IntoIterator, + table_ids: impl IntoIterator, is_cdc_truncate: bool, ) -> EtlResult<()> { // We want to lock for the entire processing to ensure that we don't have any race conditions // and possible errors are easier to reason about. let mut inner = self.inner.lock().await; - for table_id in table_ids { - let table_schema = self.store.get_table_schema(&table_id).await?; + for (table_id, schema_version) in table_ids { + let table_schema = self + .store + .get_table_schema(&table_id, schema_version) + .await?; // If we are not doing CDC, it means that this truncation has been issued while recovering // from a failed data sync operation. In that case, we could have failed before table schemas // were stored in the schema store, so we just continue and emit a warning. If we are doing @@ -920,7 +951,21 @@ where } async fn truncate_table(&self, table_id: TableId) -> EtlResult<()> { - self.process_truncate_for_table_ids(iter::once(table_id), false) + let latest_schema = self + .store + .get_latest_table_schema(&table_id) + .await? + .ok_or_else(|| { + etl_error!( + ErrorKind::MissingTableSchema, + "Table not found in the schema store", + format!( + "The table schema for table {table_id} was not found in the schema store" + ) + ) + })?; + + self.process_truncate_for_table_ids(iter::once((table_id, latest_schema.version)), false) .await } diff --git a/etl-destinations/tests/iceberg_pipeline.rs b/etl-destinations/tests/iceberg_pipeline.rs index 84aa0e407..c4c3bc213 100644 --- a/etl-destinations/tests/iceberg_pipeline.rs +++ b/etl-destinations/tests/iceberg_pipeline.rs @@ -216,7 +216,7 @@ async fn create_table_if_missing() { ColumnSchema::new("oid_array_col".to_string(), Type::OID_ARRAY, -1, false), ColumnSchema::new("bytea_array_col".to_string(), Type::BYTEA_ARRAY, -1, false), ]; - let table_schema = TableSchema::new(table_id, table_name_struct, columns); + let table_schema = TableSchema::new(table_id, table_name_struct, 0, columns); // table doesn't exist yet assert!( @@ -318,7 +318,7 @@ async fn insert_nullable_scalars() { // Binary type ColumnSchema::new("bytea_col".to_string(), Type::BYTEA, -1, false), ]; - let table_schema = TableSchema::new(table_id, table_name_struct, columns); + let table_schema = TableSchema::new(table_id, table_name_struct, 0, columns); client .create_table_if_missing(namespace, table_name.clone(), &table_schema) @@ -469,7 +469,7 @@ async fn insert_non_nullable_scalars() { // Binary type ColumnSchema::new("bytea_col".to_string(), Type::BYTEA, -1, false), ]; - let table_schema = TableSchema::new(table_id, table_name_struct, columns); + let table_schema = TableSchema::new(table_id, table_name_struct, 0, columns); client .create_table_if_missing(namespace, table_name.clone(), &table_schema) @@ -629,7 +629,7 @@ async fn insert_nullable_array() { // Binary array type ColumnSchema::new("bytea_array_col".to_string(), Type::BYTEA_ARRAY, -1, false), ]; - let table_schema = TableSchema::new(table_id, table_name_struct, columns); + let table_schema = TableSchema::new(table_id, table_name_struct, 0, columns); client .create_table_if_missing(namespace, table_name.clone(), &table_schema) @@ -903,7 +903,7 @@ async fn insert_non_nullable_array() { // Binary array type ColumnSchema::new("bytea_array_col".to_string(), Type::BYTEA_ARRAY, -1, false), ]; - let table_schema = TableSchema::new(table_id, table_name_struct, columns); + let table_schema = TableSchema::new(table_id, table_name_struct, 0, columns); client .create_table_if_missing(namespace, table_name.clone(), &table_schema) diff --git a/etl-postgres/src/replication/schema.rs b/etl-postgres/src/replication/schema.rs index 336ff166c..8502872e5 100644 --- a/etl-postgres/src/replication/schema.rs +++ b/etl-postgres/src/replication/schema.rs @@ -4,7 +4,9 @@ use sqlx::{PgExecutor, PgPool, Row}; use std::collections::HashMap; use tokio_postgres::types::Type as PgType; -use crate::types::{ColumnSchema, TableId, TableName, TableSchema}; +use crate::types::{ + ColumnSchema, SchemaVersion, TableId, TableName, TableSchema, TableSchemaDraft, +}; macro_rules! define_type_mappings { ( @@ -141,20 +143,32 @@ define_type_mappings! { pub async fn store_table_schema( pool: &PgPool, pipeline_id: i64, - table_schema: &TableSchema, -) -> Result<(), sqlx::Error> { + table_schema: TableSchemaDraft, +) -> Result { let mut tx = pool.begin().await?; - // Insert or update table schema record + let current_version: Option = sqlx::query_scalar( + r#" + select max(schema_version) + from etl.table_schemas + where pipeline_id = $1 + and table_id = $2 + "#, + ) + .bind(pipeline_id) + .bind(table_schema.id.into_inner() as i64) + .fetch_one(&mut *tx) + .await?; + + let next_version = current_version.unwrap_or(-1) + 1; + let schema_version: SchemaVersion = next_version + .try_into() + .expect("schema version should not overflow u64"); + let table_schema_id: i64 = sqlx::query( r#" - insert into etl.table_schemas (pipeline_id, table_id, schema_name, table_name) - values ($1, $2, $3, $4) - on conflict (pipeline_id, table_id) - do update set - schema_name = excluded.schema_name, - table_name = excluded.table_name, - updated_at = now() + insert into etl.table_schemas (pipeline_id, table_id, schema_name, table_name, schema_version) + values ($1, $2, $3, $4, $5) returning id "#, ) @@ -162,17 +176,11 @@ pub async fn store_table_schema( .bind(table_schema.id.into_inner() as i64) .bind(&table_schema.name.schema) .bind(&table_schema.name.name) + .bind(next_version) .fetch_one(&mut *tx) .await? .get(0); - // Delete existing columns for this table schema to handle schema changes - sqlx::query("delete from etl.table_columns where table_schema_id = $1") - .bind(table_schema_id) - .execute(&mut *tx) - .await?; - - // Insert all columns for (column_order, column_schema) in table_schema.column_schemas.iter().enumerate() { let column_type_str = postgres_type_to_string(&column_schema.typ); @@ -196,7 +204,7 @@ pub async fn store_table_schema( tx.commit().await?; - Ok(()) + Ok(table_schema.into_table_schema(schema_version)) } /// Loads all table schemas for a pipeline from the database. @@ -213,6 +221,7 @@ pub async fn load_table_schemas( ts.table_id, ts.schema_name, ts.table_name, + ts.schema_version, tc.column_name, tc.column_type, tc.type_modifier, @@ -222,24 +231,37 @@ pub async fn load_table_schemas( from etl.table_schemas ts inner join etl.table_columns tc on ts.id = tc.table_schema_id where ts.pipeline_id = $1 - order by ts.table_id, tc.column_order + order by ts.table_id, ts.schema_version, tc.column_order "#, ) .bind(pipeline_id) .fetch_all(pool) .await?; - let mut table_schemas = HashMap::new(); + let mut table_schemas: HashMap<(TableId, SchemaVersion), TableSchema> = HashMap::new(); for row in rows { 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"); - - let entry = table_schemas.entry(table_id).or_insert_with(|| { - TableSchema::new(table_id, TableName::new(schema_name, table_name), vec![]) - }); + let schema_version: SchemaVersion = { + let value: i64 = row.get("schema_version"); + value + .try_into() + .expect("schema_version should fit into SchemaVersion") + }; + + let entry = table_schemas + .entry((table_id, schema_version)) + .or_insert_with(|| { + TableSchema::new( + table_id, + TableName::new(schema_name.clone(), table_name.clone()), + schema_version, + vec![], + ) + }); entry.add_column_schema(parse_column_schema(&row)); } diff --git a/etl-postgres/src/types/schema.rs b/etl-postgres/src/types/schema.rs index 2b64baa65..ea01bf275 100644 --- a/etl-postgres/src/types/schema.rs +++ b/etl-postgres/src/types/schema.rs @@ -7,6 +7,9 @@ use tokio_postgres::types::{FromSql, ToSql, Type}; /// An object identifier in Postgres. type Oid = u32; +/// Version number for a table schema. +pub type SchemaVersion = u64; + /// A fully qualified Postgres table name consisting of a schema and table name. /// /// This type represents a table identifier in Postgres, which requires both a schema name @@ -173,15 +176,23 @@ pub struct TableSchema { pub id: TableId, /// The fully qualified name of the table pub name: TableName, + /// Monotonically increasing schema version. + pub version: SchemaVersion, /// The schemas of all columns in the table pub column_schemas: Vec, } impl TableSchema { - pub fn new(id: TableId, name: TableName, column_schemas: Vec) -> Self { + pub fn new( + id: TableId, + name: TableName, + version: SchemaVersion, + column_schemas: Vec, + ) -> Self { Self { id, name, + version, column_schemas, } } @@ -190,13 +201,6 @@ impl TableSchema { pub fn add_column_schema(&mut self, column_schema: ColumnSchema) { self.column_schemas.push(column_schema); } - - /// Returns whether the table has any primary key columns. - /// - /// This method checks if any column in the table is marked as part of the primary key. - pub fn has_primary_keys(&self) -> bool { - self.column_schemas.iter().any(|cs| cs.primary) - } } impl PartialOrd for TableSchema { @@ -207,6 +211,49 @@ impl PartialOrd for TableSchema { impl Ord for TableSchema { fn cmp(&self, other: &Self) -> Ordering { - self.name.cmp(&other.name) + self.name + .cmp(&other.name) + .then(self.version.cmp(&other.version)) + } +} + +/// Draft version of [`TableSchema`] used before a schema version is assigned. +#[derive(Debug, Clone, Eq, PartialEq)] +pub struct TableSchemaDraft { + pub id: TableId, + pub name: TableName, + pub column_schemas: Vec, +} + +impl TableSchemaDraft { + pub fn new(id: TableId, name: TableName, column_schemas: Vec) -> Self { + Self { + id, + name, + column_schemas, + } + } + + pub fn with_capacity(id: TableId, name: TableName, capacity: usize) -> Self { + Self { + id, + name, + column_schemas: Vec::with_capacity(capacity), + } + } + + pub fn add_column_schema(&mut self, column_schema: ColumnSchema) { + self.column_schemas.push(column_schema); + } + + /// Returns whether the table has any primary key columns. + /// + /// This method checks if any column in the table is marked as part of the primary key. + pub fn has_primary_keys(&self) -> bool { + self.column_schemas.iter().any(|cs| cs.primary) + } + + pub fn into_table_schema(self, version: SchemaVersion) -> TableSchema { + TableSchema::new(self.id, self.name, version, self.column_schemas) } } diff --git a/etl-replicator/migrations/20251001000100_add_schema_versions.sql b/etl-replicator/migrations/20251001000100_add_schema_versions.sql new file mode 100644 index 000000000..09ded33e1 --- /dev/null +++ b/etl-replicator/migrations/20251001000100_add_schema_versions.sql @@ -0,0 +1,21 @@ +-- Add schema_version to table schemas and support versioned definitions +alter table etl.table_schemas + add column if not exists schema_version bigint not null default 0; + +-- Adjust unique constraint to account for schema versions +alter table etl.table_schemas + drop constraint if exists etl_table_schemas_pipeline_id_table_id_key; + +alter table etl.table_schemas + add constraint etl_table_schemas_pipeline_id_table_id_version_key + unique (pipeline_id, table_id, schema_version); + +-- Refresh supporting indexes +drop index if exists idx_table_schemas_pipeline_table; + +create index if not exists idx_table_schemas_pipeline_table_version + on etl.table_schemas (pipeline_id, table_id, schema_version); + +-- Remove default now that legacy rows have been backfilled +alter table etl.table_schemas + alter column schema_version drop default; diff --git a/etl/src/conversions/event.rs b/etl/src/conversions/event.rs index 1b046b750..14f55d4cd 100644 --- a/etl/src/conversions/event.rs +++ b/etl/src/conversions/event.rs @@ -1,48 +1,18 @@ use core::str; -use etl_postgres::types::{ColumnSchema, TableId, TableSchema, convert_type_oid_to_type}; +use etl_postgres::types::{ + ColumnSchema, TableId, TableSchema, TableSchemaDraft, convert_type_oid_to_type, +}; use postgres_replication::protocol; -use std::cmp::Ordering; -use std::collections::BTreeSet; use std::sync::Arc; use tokio_postgres::types::PgLsn; +use crate::bail; use crate::conversions::text::{default_value_for_type, parse_cell_from_postgres_text}; use crate::error::{ErrorKind, EtlError, EtlResult}; -use crate::store::schema::SchemaStore; use crate::types::{ - BeginEvent, Cell, CommitEvent, DeleteEvent, InsertEvent, RelationChange, RelationEvent, - TableRow, TruncateEvent, UpdateEvent, + BeginEvent, Cell, CommitEvent, DeleteEvent, InsertEvent, RelationEventDraft, TableRow, + TruncateEvent, UpdateEvent, }; -use crate::{bail, etl_error}; - -#[derive(Debug, Clone)] -struct IndexedColumnSchema(ColumnSchema); - -impl IndexedColumnSchema { - fn into_inner(self) -> ColumnSchema { - self.0 - } -} - -impl Eq for IndexedColumnSchema {} - -impl PartialEq for IndexedColumnSchema { - fn eq(&self, other: &Self) -> bool { - self.0.name == other.0.name - } -} - -impl Ord for IndexedColumnSchema { - fn cmp(&self, other: &Self) -> Ordering { - self.0.name.cmp(&other.0.name) - } -} - -impl PartialOrd for IndexedColumnSchema { - fn partial_cmp(&self, other: &Self) -> Option { - Some(self.0.name.cmp(&other.0.name)) - } -} /// Creates a [`BeginEvent`] from Postgres protocol data. /// @@ -79,110 +49,45 @@ pub fn parse_event_from_commit_message( } } -/// Creates a [`RelationEvent`] from Postgres protocol data. +/// Creates a [`RelationEventDraft`] from Postgres protocol data. /// /// This method parses the replication protocol relation message and builds -/// a complete table schema for use in interpreting subsequent data events. -pub async fn parse_event_from_relation_message( - schema_store: &S, - start_lsn: PgLsn, - commit_lsn: PgLsn, +/// a complete table schema for use in interpreting later data events. +/// +/// The method returns a draft since the actual full event needs the table schema version +/// from the schema store, and we don't want to pollute the logic of this module with schema +/// store logic. +pub async fn parse_event_from_relation_message( + old_table_schema: Arc, relation_body: &protocol::RelationBody, -) -> EtlResult -where - S: SchemaStore, -{ - let table_id = TableId::new(relation_body.rel_id()); - - let Some(table_schema) = schema_store.get_table_schema(&table_id).await? else { - bail!( - ErrorKind::MissingTableSchema, - "Table not found in the schema cache", - format!("The table schema for table {table_id} was not found in the cache") - ) - }; - - // We construct the latest column schemas in order. The order is important since the table schema +) -> EtlResult { + // We construct the new column schemas in order. The order is important since the table schema // relies on the right ordering to interpret the Postgres correctly. - let latest_column_schemas = relation_body + let new_column_schemas = relation_body .columns() .iter() .map(build_column_schema) .collect::, EtlError>>()?; - // We build a lookup set for the latest column schemas for quick change detection. - let mut latest_indexed_column_schemas = latest_column_schemas - .iter() - .cloned() - .map(IndexedColumnSchema) - .collect::>(); - - // We process all the changes that we want to dispatch to the destination. - let mut changes = vec![]; - for column_schema in table_schema.column_schemas.iter() { - let column_schema = IndexedColumnSchema(column_schema.clone()); - let latest_column_schema = latest_indexed_column_schemas.take(&column_schema); - match latest_column_schema { - Some(latest_column_schema) => { - let column_schema = column_schema.into_inner(); - let latest_column_schema = latest_column_schema.into_inner(); - - if column_schema.name != latest_column_schema.name { - // If we find a column with the same name but different fields, we assume it was changed. The only changes - // that we detect are changes to the column but with preserved name. - changes.push(RelationChange::AlterColumn( - column_schema, - latest_column_schema, - )); - } - } - None => { - // If we don't find the column in the latest schema, we assume it was dropped even - // though it could be renamed. - changes.push(RelationChange::DropColumn(column_schema.into_inner())); - } - } - } - - // For the remaining columns that didn't match, we assume they were added. - for column_schema in latest_indexed_column_schemas { - changes.push(RelationChange::AddColumn(column_schema.into_inner())); - } - - // We build the updated table schema to store in the schema store. - schema_store - .store_table_schema(TableSchema::new( - table_schema.id, - table_schema.name.clone(), - latest_column_schemas, - )) - .await?; + let new_table_schema = TableSchemaDraft::new( + old_table_schema.id, + old_table_schema.name.clone(), + new_column_schemas, + ); - Ok(RelationEvent { - start_lsn, - commit_lsn, - changes, - table_id, - }) + Ok(new_table_schema) } /// Converts a Postgres insert message into an [`InsertEvent`]. /// -/// This function processes an insert operation from the replication stream, -/// retrieves the table schema from the store, and constructs a complete -/// insert event with the new row data ready for ETL processing. -pub async fn parse_event_from_insert_message( - schema_store: &S, +/// This function processes an insert operation from the replication stream +/// using the supplied table schema version to build the resulting event. +pub fn parse_event_from_insert_message( + table_schema: Arc, start_lsn: PgLsn, commit_lsn: PgLsn, insert_body: &protocol::InsertBody, -) -> EtlResult -where - S: SchemaStore, -{ - let table_id = insert_body.rel_id(); - let table_schema = get_table_schema(schema_store, TableId::new(table_id)).await?; - +) -> EtlResult { let table_row = convert_tuple_to_row( &table_schema.column_schemas, insert_body.tuple().tuple_data(), @@ -193,7 +98,8 @@ where Ok(InsertEvent { start_lsn, commit_lsn, - table_id: TableId::new(table_id), + schema_version: table_schema.version, + table_id: table_schema.id, table_row, }) } @@ -204,18 +110,12 @@ where /// handling both the old and new row data. The old row data may be either /// the complete row or just the key columns, depending on the table's /// `REPLICA IDENTITY` setting in Postgres. -pub async fn parse_event_from_update_message( - schema_store: &S, +pub fn parse_event_from_update_message( + table_schema: Arc, start_lsn: PgLsn, commit_lsn: PgLsn, update_body: &protocol::UpdateBody, -) -> EtlResult -where - S: SchemaStore, -{ - let table_id = update_body.rel_id(); - let table_schema = get_table_schema(schema_store, TableId::new(table_id)).await?; - +) -> EtlResult { // We try to extract the old tuple by either taking the entire old tuple or the key of the old // tuple. let is_key = update_body.old_tuple().is_none(); @@ -243,7 +143,8 @@ where Ok(UpdateEvent { start_lsn, commit_lsn, - table_id: TableId::new(table_id), + schema_version: table_schema.version, + table_id: table_schema.id, table_row, old_table_row, }) @@ -255,18 +156,12 @@ where /// extracting the old row data that was deleted. The old row data may be /// either the complete row or just the key columns, depending on the table's /// `REPLICA IDENTITY` setting in Postgres. -pub async fn parse_event_from_delete_message( - schema_store: &S, +pub fn parse_event_from_delete_message( + table_schema: Arc, start_lsn: PgLsn, commit_lsn: PgLsn, delete_body: &protocol::DeleteBody, -) -> EtlResult -where - S: SchemaStore, -{ - let table_id = delete_body.rel_id(); - let table_schema = get_table_schema(schema_store, TableId::new(table_id)).await?; - +) -> EtlResult { // We try to extract the old tuple by either taking the entire old tuple or the key of the old // tuple. let is_key = delete_body.old_tuple().is_none(); @@ -285,7 +180,8 @@ where Ok(DeleteEvent { start_lsn, commit_lsn, - table_id: TableId::new(table_id), + schema_version: table_schema.version, + table_id: table_schema.id, old_table_row, }) } @@ -298,37 +194,16 @@ pub fn parse_event_from_truncate_message( start_lsn: PgLsn, commit_lsn: PgLsn, truncate_body: &protocol::TruncateBody, - overridden_rel_ids: Vec, + relations: Vec, ) -> TruncateEvent { TruncateEvent { start_lsn, commit_lsn, options: truncate_body.options(), - rel_ids: overridden_rel_ids, + relations, } } -/// Retrieves a table schema from the schema store by table ID. -/// -/// This function looks up the table schema for the specified table ID in the -/// schema store. If the schema is not found, it returns an error indicating -/// that the table is missing from the cache. -async fn get_table_schema(schema_store: &S, table_id: TableId) -> EtlResult> -where - S: SchemaStore, -{ - schema_store - .get_table_schema(&table_id) - .await? - .ok_or_else(|| { - etl_error!( - ErrorKind::MissingTableSchema, - "Table not found in the schema cache", - format!("The table schema for table {table_id} was not found in the cache") - ) - }) -} - /// Constructs a [`IndexedColumnSchema`] from Postgres protocol column data. /// /// This helper method extracts column metadata from the replication protocol diff --git a/etl/src/replication/apply.rs b/etl/src/replication/apply.rs index 1146aba61..f65cf5b1f 100644 --- a/etl/src/replication/apply.rs +++ b/etl/src/replication/apply.rs @@ -1,10 +1,11 @@ use etl_config::shared::PipelineConfig; use etl_postgres::replication::worker::WorkerType; -use etl_postgres::types::TableId; +use etl_postgres::types::{SchemaVersion, TableId, TableSchema}; use futures::StreamExt; use metrics::histogram; use postgres_replication::protocol; use postgres_replication::protocol::{LogicalReplicationMessage, ReplicationMessage}; +use std::collections::HashMap; use std::future::Future; use std::pin::Pin; use std::sync::Arc; @@ -13,7 +14,6 @@ use tokio::pin; use tokio_postgres::types::PgLsn; use tracing::{debug, info}; -use crate::bail; use crate::concurrency::shutdown::ShutdownRx; use crate::concurrency::signal::SignalRx; use crate::concurrency::stream::{TimeoutStream, TimeoutStreamResult}; @@ -34,7 +34,8 @@ use crate::replication::client::PgReplicationClient; use crate::replication::stream::EventsStream; use crate::state::table::TableReplicationError; use crate::store::schema::SchemaStore; -use crate::types::{Event, PipelineId}; +use crate::types::{Event, PipelineId, RelationEvent}; +use crate::{bail, etl_error}; /// The amount of milliseconds that pass between one refresh and the other of the system, in case no /// events or shutdown signal are received. @@ -747,6 +748,9 @@ where // is restarted. events_stream.mark_reset_timer(); + // TODO: once the batch is sent, we should also remove previous schema versions that are not needed + // anymore. + Ok(()) } @@ -912,7 +916,7 @@ where handle_delete_message(state, start_lsn, delete_body, hook, schema_store).await } LogicalReplicationMessage::Truncate(truncate_body) => { - handle_truncate_message(state, start_lsn, truncate_body, hook).await + handle_truncate_message(state, start_lsn, truncate_body, schema_store, hook).await } LogicalReplicationMessage::Origin(_) => { debug!("received unsupported ORIGIN message"); @@ -1102,9 +1106,22 @@ where } // Convert event from the protocol message. - let event = - parse_event_from_relation_message(schema_store, start_lsn, remote_final_lsn, message) - .await?; + let old_table_schema = load_table_schema(schema_store, table_id).await?; + let new_table_schema_draft = + parse_event_from_relation_message(old_table_schema.clone(), message).await?; + + // We store the new schema in the store and build the final relation event. + let new_table_schema = schema_store + .store_table_schema(new_table_schema_draft) + .await?; + + let event = RelationEvent { + start_lsn, + commit_lsn: remote_final_lsn, + table_id, + old_table_schema, + new_table_schema, + }; Ok(HandleMessageResult::return_event(Event::Relation(event))) } @@ -1129,16 +1146,19 @@ where ); }; + let table_id = TableId::new(message.rel_id()); + if !hook - .should_apply_changes(TableId::new(message.rel_id()), remote_final_lsn) + .should_apply_changes(table_id, remote_final_lsn) .await? { return Ok(HandleMessageResult::no_event()); } // Convert event from the protocol message. + let table_schema = load_table_schema(schema_store, table_id).await?; let event = - parse_event_from_insert_message(schema_store, start_lsn, remote_final_lsn, message).await?; + parse_event_from_insert_message(table_schema, start_lsn, remote_final_lsn, message)?; Ok(HandleMessageResult::return_event(Event::Insert(event))) } @@ -1163,16 +1183,19 @@ where ); }; + let table_id = TableId::new(message.rel_id()); + if !hook - .should_apply_changes(TableId::new(message.rel_id()), remote_final_lsn) + .should_apply_changes(table_id, remote_final_lsn) .await? { return Ok(HandleMessageResult::no_event()); } // Convert event from the protocol message. + let table_schema = load_table_schema(schema_store, table_id).await?; let event = - parse_event_from_update_message(schema_store, start_lsn, remote_final_lsn, message).await?; + parse_event_from_update_message(table_schema, start_lsn, remote_final_lsn, message)?; Ok(HandleMessageResult::return_event(Event::Update(event))) } @@ -1197,16 +1220,19 @@ where ); }; + let table_id = TableId::new(message.rel_id()); + if !hook - .should_apply_changes(TableId::new(message.rel_id()), remote_final_lsn) + .should_apply_changes(table_id, remote_final_lsn) .await? { return Ok(HandleMessageResult::no_event()); } // Convert event from the protocol message. + let table_schema = load_table_schema(schema_store, table_id).await?; let event = - parse_event_from_delete_message(schema_store, start_lsn, remote_final_lsn, message).await?; + parse_event_from_delete_message(table_schema, start_lsn, remote_final_lsn, message)?; Ok(HandleMessageResult::return_event(Event::Delete(event))) } @@ -1217,13 +1243,15 @@ where /// ensuring transaction context, and filtering the affected table list based on /// hook decisions. Since TRUNCATE can affect multiple tables simultaneously, /// it evaluates each table individually. -async fn handle_truncate_message( +async fn handle_truncate_message( state: &mut ApplyLoopState, start_lsn: PgLsn, message: &protocol::TruncateBody, + schema_store: &S, hook: &T, ) -> EtlResult where + S: SchemaStore + Clone + Send + 'static, T: ApplyLoopHook, { let Some(remote_final_lsn) = state.remote_final_lsn else { @@ -1234,24 +1262,46 @@ where ); }; - // We collect only the relation ids for which we are allow to apply changes, thus in this case + // We collect only the relation ids for which we are allowed to apply changes, thus in this case // the truncation. - let mut rel_ids = Vec::with_capacity(message.rel_ids().len()); + let mut relations = Vec::with_capacity(message.rel_ids().len()); for &table_id in message.rel_ids().iter() { + let table_id = TableId::new(table_id); if hook - .should_apply_changes(TableId::new(table_id), remote_final_lsn) + .should_apply_changes(table_id, remote_final_lsn) .await? { - rel_ids.push(table_id) + relations.push(table_id); } } // If nothing to apply, skip conversion entirely - if rel_ids.is_empty() { + if relations.is_empty() { return Ok(HandleMessageResult::no_event()); } // Convert event from the protocol message. - let event = parse_event_from_truncate_message(start_lsn, remote_final_lsn, message, rel_ids); + let event = parse_event_from_truncate_message(start_lsn, remote_final_lsn, message, relations); Ok(HandleMessageResult::return_event(Event::Truncate(event))) } + +/// Loads the table schema for processing a specific event. +async fn load_table_schema(schema_store: &S, table_id: TableId) -> EtlResult> +where + S: SchemaStore + Clone + Send + 'static, +{ + let table_schema = schema_store + .get_latest_table_schema(&table_id) + .await? + .ok_or_else(|| { + etl_error!( + ErrorKind::MissingTableSchema, + "Table not found in the schema store", + format!( + "The latest table schema for table {table_id} was not found in the schema store" + ) + ) + })?; + + Ok(table_schema) +} diff --git a/etl/src/replication/client.rs b/etl/src/replication/client.rs index 98cb0a878..fc41f3168 100644 --- a/etl/src/replication/client.rs +++ b/etl/src/replication/client.rs @@ -3,8 +3,8 @@ use crate::utils::tokio::MakeRustlsConnect; use crate::{bail, etl_error}; use etl_config::shared::{IntoConnectOptions, PgConnectionConfig}; use etl_postgres::replication::extract_server_version; -use etl_postgres::types::convert_type_oid_to_type; use etl_postgres::types::{ColumnSchema, TableId, TableName, TableSchema}; +use etl_postgres::types::{TableSchemaDraft, convert_type_oid_to_type}; use pg_escape::{quote_identifier, quote_literal}; use postgres_replication::LogicalReplicationStream; use rustls::ClientConfig; @@ -117,7 +117,7 @@ impl PgReplicationSlotTransaction { &self, table_ids: &[TableId], publication_name: Option<&str>, - ) -> EtlResult> { + ) -> EtlResult> { self.client .get_table_schemas(table_ids, publication_name) .await @@ -131,7 +131,7 @@ impl PgReplicationSlotTransaction { &self, table_id: TableId, publication: Option<&str>, - ) -> EtlResult { + ) -> EtlResult { self.client.get_table_schema(table_id, publication).await } @@ -558,7 +558,7 @@ impl PgReplicationClient { &self, table_ids: &[TableId], publication_name: Option<&str>, - ) -> EtlResult> { + ) -> EtlResult> { let mut table_schemas = HashMap::new(); // TODO: consider if we want to fail when at least one table was missing or not. @@ -589,11 +589,11 @@ impl PgReplicationClient { &self, table_id: TableId, publication: Option<&str>, - ) -> EtlResult { + ) -> EtlResult { let table_name = self.get_table_name(table_id).await?; let column_schemas = self.get_column_schemas(table_id, publication).await?; - Ok(TableSchema { + Ok(TableSchemaDraft { name: table_name, id: table_id, column_schemas, diff --git a/etl/src/store/both/memory.rs b/etl/src/store/both/memory.rs index 88a2af733..3fccef886 100644 --- a/etl/src/store/both/memory.rs +++ b/etl/src/store/both/memory.rs @@ -1,5 +1,5 @@ -use etl_postgres::types::{TableId, TableSchema}; -use std::collections::HashMap; +use etl_postgres::types::{SchemaVersion, TableId, TableSchema, TableSchemaDraft}; +use std::collections::{BTreeMap, HashMap}; use std::sync::Arc; use tokio::sync::Mutex; @@ -23,7 +23,7 @@ struct Inner { /// Cached table schema definitions, reference-counted for efficient sharing. /// Schemas are expensive to fetch from Postgres, so they're cached here /// once retrieved and shared via Arc across the application. - table_schemas: HashMap>, + table_schemas: HashMap>>, /// Mapping from table IDs to human-readable table names for easier debugging /// and logging. These mappings are established during schema discovery. table_mappings: HashMap, @@ -172,31 +172,71 @@ impl StateStore for MemoryStore { } impl SchemaStore for MemoryStore { - async fn get_table_schema(&self, table_id: &TableId) -> EtlResult>> { + async fn get_table_schema( + &self, + table_id: &TableId, + version: SchemaVersion, + ) -> EtlResult>> { + let inner = self.inner.lock().await; + + Ok(inner + .table_schemas + .get(table_id) + .and_then(|schemas| schemas.get(&version).cloned())) + } + + async fn get_latest_table_schema( + &self, + table_id: &TableId, + ) -> EtlResult>> { let inner = self.inner.lock().await; - Ok(inner.table_schemas.get(table_id).cloned()) + Ok(inner + .table_schemas + .get(table_id) + .and_then(|schemas| schemas.iter().next_back().map(|(_, schema)| schema.clone()))) } async fn get_table_schemas(&self) -> EtlResult>> { let inner = self.inner.lock().await; - Ok(inner.table_schemas.values().cloned().collect()) + Ok(inner + .table_schemas + .values() + .filter_map(|schemas| schemas.iter().next_back().map(|(_, schema)| schema.clone())) + .collect()) } async fn load_table_schemas(&self) -> EtlResult { let inner = self.inner.lock().await; - Ok(inner.table_schemas.len()) + Ok(inner + .table_schemas + .values() + .map(|schemas| schemas.len()) + .sum()) } - async fn store_table_schema(&self, table_schema: TableSchema) -> EtlResult<()> { + async fn store_table_schema( + &self, + table_schema: TableSchemaDraft, + ) -> EtlResult> { let mut inner = self.inner.lock().await; - inner + let schemas = inner .table_schemas - .insert(table_schema.id, Arc::new(table_schema)); + .entry(table_schema.id) + .or_insert_with(BTreeMap::new); - Ok(()) + let next_version = schemas + .keys() + .next_back() + .map(|version| version + 1) + .unwrap_or(0); + + let schema = Arc::new(table_schema.into_table_schema(next_version)); + schemas.insert(next_version, Arc::clone(&schema)); + + Ok(schema) } } diff --git a/etl/src/store/both/postgres.rs b/etl/src/store/both/postgres.rs index 427e25cd3..5db97f4a6 100644 --- a/etl/src/store/both/postgres.rs +++ b/etl/src/store/both/postgres.rs @@ -1,10 +1,12 @@ -use std::{collections::HashMap, sync::Arc}; - use etl_config::shared::PgConnectionConfig; use etl_postgres::replication::{connect_to_source_database, schema, state, table_mappings}; -use etl_postgres::types::{TableId, TableSchema}; +use etl_postgres::types::{SchemaVersion, TableId, TableSchema, TableSchemaDraft}; use metrics::gauge; use sqlx::PgPool; +use std::{ + collections::{BTreeMap, HashMap}, + sync::Arc, +}; use tokio::sync::Mutex; use tracing::{debug, info}; @@ -136,7 +138,7 @@ struct Inner { /// Cached table replication states indexed by table ID. table_states: HashMap, /// Cached table schemas indexed by table ID. - table_schemas: HashMap>, + table_schemas: HashMap>>, /// Cached table mappings from source table ID to destination table name. table_mappings: HashMap, } @@ -506,10 +508,29 @@ impl SchemaStore for PostgresStore { /// This method provides fast access to cached table schemas, which are /// essential for processing replication events. Schemas are loaded during /// startup and cached for the lifetime of the pipeline. - async fn get_table_schema(&self, table_id: &TableId) -> EtlResult>> { + async fn get_table_schema( + &self, + table_id: &TableId, + version: SchemaVersion, + ) -> EtlResult>> { + let inner = self.inner.lock().await; + + Ok(inner + .table_schemas + .get(table_id) + .and_then(|schemas| schemas.get(&version).cloned())) + } + + async fn get_latest_table_schema( + &self, + table_id: &TableId, + ) -> EtlResult>> { let inner = self.inner.lock().await; - Ok(inner.table_schemas.get(table_id).cloned()) + Ok(inner + .table_schemas + .get(table_id) + .and_then(|schemas| schemas.iter().next_back().map(|(_, schema)| schema.clone()))) } /// Retrieves all cached table schemas as a vector. @@ -520,7 +541,11 @@ impl SchemaStore for PostgresStore { async fn get_table_schemas(&self) -> EtlResult>> { let inner = self.inner.lock().await; - Ok(inner.table_schemas.values().cloned().collect()) + Ok(inner + .table_schemas + .values() + .filter_map(|schemas| schemas.iter().next_back().map(|(_, schema)| schema.clone())) + .collect()) } /// Loads table schemas from Postgres into memory cache. @@ -552,7 +577,9 @@ impl SchemaStore for PostgresStore { for table_schema in table_schemas { inner .table_schemas - .insert(table_schema.id, Arc::new(table_schema)); + .entry(table_schema.id) + .or_insert_with(BTreeMap::new) + .insert(table_schema.version, Arc::new(table_schema)); } info!( @@ -560,6 +587,12 @@ impl SchemaStore for PostgresStore { table_schemas_len ); + emit_table_metrics( + self.pipeline_id, + inner.table_states.keys().len(), + inner.table_schemas.keys().len(), + ); + Ok(table_schemas_len) } @@ -568,27 +601,34 @@ impl SchemaStore for PostgresStore { /// This method persists a table schema to the database and updates the /// in-memory cache atomically. Used when new tables are discovered during /// replication or when schema definitions need to be updated. - async fn store_table_schema(&self, table_schema: TableSchema) -> EtlResult<()> { + async fn store_table_schema( + &self, + table_schema: TableSchemaDraft, + ) -> EtlResult> { debug!("storing table schema for table '{}'", table_schema.name); let pool = self.connect_to_source().await?; // We also lock the entire section to be consistent. let mut inner = self.inner.lock().await; - schema::store_table_schema(&pool, self.pipeline_id as i64, &table_schema) - .await - .map_err(|err| { - etl_error!( - ErrorKind::SourceQueryFailed, - "Failed to store table schema", - format!("Failed to store table schema in postgres: {err}") - ) - })?; + let stored_schema = + schema::store_table_schema(&pool, self.pipeline_id as i64, table_schema) + .await + .map_err(|err| { + etl_error!( + ErrorKind::SourceQueryFailed, + "Failed to store table schema", + format!("Failed to store table schema in postgres: {err}") + ) + })?; + let schema_arc = Arc::new(stored_schema); inner .table_schemas - .insert(table_schema.id, Arc::new(table_schema)); + .entry(schema_arc.id) + .or_insert_with(BTreeMap::new) + .insert(schema_arc.version, Arc::clone(&schema_arc)); - Ok(()) + Ok(schema_arc) } } diff --git a/etl/src/store/schema/base.rs b/etl/src/store/schema/base.rs index 2a52126b9..227202bda 100644 --- a/etl/src/store/schema/base.rs +++ b/etl/src/store/schema/base.rs @@ -1,4 +1,4 @@ -use etl_postgres::types::{TableId, TableSchema}; +use etl_postgres::types::{SchemaVersion, TableId, TableSchema, TableSchemaDraft}; use std::sync::Arc; use crate::error::EtlResult; @@ -16,11 +16,16 @@ pub trait SchemaStore { fn get_table_schema( &self, table_id: &TableId, + version: SchemaVersion, ) -> impl Future>>> + Send; - /// Returns all table schemas from the cache. - /// - /// Does not read from the persistent store. + /// Returns the latest table schema version for table with id `table_id` from the cache. + fn get_latest_table_schema( + &self, + table_id: &TableId, + ) -> impl Future>>> + Send; + + /// Returns the latest table schema for all tables from the cache. fn get_table_schemas(&self) -> impl Future>>> + Send; /// Loads table schemas from the persistent state into the cache. @@ -31,6 +36,6 @@ pub trait SchemaStore { /// Stores a table schema in both the cache and the persistent store. fn store_table_schema( &self, - table_schema: TableSchema, - ) -> impl Future> + Send; + table_schema: TableSchemaDraft, + ) -> impl Future>> + Send; } diff --git a/etl/src/test_utils/notify.rs b/etl/src/test_utils/notify.rs index 09fff96b3..44849455f 100644 --- a/etl/src/test_utils/notify.rs +++ b/etl/src/test_utils/notify.rs @@ -1,5 +1,9 @@ -use etl_postgres::types::{TableId, TableSchema}; -use std::{collections::HashMap, fmt, sync::Arc}; +use etl_postgres::types::{SchemaVersion, TableId, TableSchema, TableSchemaDraft}; +use std::{ + collections::{BTreeMap, HashMap}, + fmt, + sync::Arc, +}; use tokio::sync::{Notify, RwLock}; use crate::error::{ErrorKind, EtlResult}; @@ -28,7 +32,7 @@ type TableStateCondition = ( struct Inner { table_replication_states: HashMap, table_state_history: HashMap>, - table_schemas: HashMap>, + table_schemas: HashMap>>, table_mappings: HashMap, table_state_type_conditions: Vec, table_state_conditions: Vec, @@ -289,30 +293,70 @@ impl StateStore for NotifyingStore { } impl SchemaStore for NotifyingStore { - async fn get_table_schema(&self, table_id: &TableId) -> EtlResult>> { + async fn get_table_schema( + &self, + table_id: &TableId, + version: SchemaVersion, + ) -> EtlResult>> { + let inner = self.inner.read().await; + + Ok(inner + .table_schemas + .get(table_id) + .and_then(|schemas| schemas.get(&version).cloned())) + } + + async fn get_latest_table_schema( + &self, + table_id: &TableId, + ) -> EtlResult>> { let inner = self.inner.read().await; - Ok(inner.table_schemas.get(table_id).cloned()) + Ok(inner + .table_schemas + .get(table_id) + .and_then(|schemas| schemas.iter().next_back().map(|(_, schema)| schema.clone()))) } async fn get_table_schemas(&self) -> EtlResult>> { let inner = self.inner.read().await; - Ok(inner.table_schemas.values().cloned().collect()) + Ok(inner + .table_schemas + .values() + .filter_map(|schemas| schemas.iter().next_back().map(|(_, schema)| schema.clone())) + .collect()) } async fn load_table_schemas(&self) -> EtlResult { let inner = self.inner.read().await; - Ok(inner.table_schemas.len()) + Ok(inner + .table_schemas + .values() + .map(|schemas| schemas.len()) + .sum()) } - async fn store_table_schema(&self, table_schema: TableSchema) -> EtlResult<()> { + async fn store_table_schema( + &self, + table_schema: TableSchemaDraft, + ) -> EtlResult> { let mut inner = self.inner.write().await; - inner + let schemas = inner .table_schemas - .insert(table_schema.id, Arc::new(table_schema)); + .entry(table_schema.id) + .or_insert_with(BTreeMap::new); - Ok(()) + let next_version = schemas + .keys() + .next_back() + .map(|version| version + 1) + .unwrap_or(0); + + let schema = Arc::new(table_schema.into_table_schema(next_version)); + schemas.insert(next_version, Arc::clone(&schema)); + + Ok(schema) } } diff --git a/etl/src/test_utils/test_schema.rs b/etl/src/test_utils/test_schema.rs index 392077fd0..4d266ba0d 100644 --- a/etl/src/test_utils/test_schema.rs +++ b/etl/src/test_utils/test_schema.rs @@ -64,6 +64,7 @@ pub async fn setup_test_database_schema( users_table_schema = Some(TableSchema::new( users_table_id, users_table_name, + 0, vec![ id_column_schema(), ColumnSchema::new("name".to_string(), Type::TEXT, -1, false), @@ -88,6 +89,7 @@ pub async fn setup_test_database_schema( orders_table_schema = Some(TableSchema::new( orders_table_id, orders_table_name, + 0, vec![ id_column_schema(), ColumnSchema::new("description".to_string(), Type::TEXT, -1, false), diff --git a/etl/src/types/event.rs b/etl/src/types/event.rs index 44010552a..383803a61 100644 --- a/etl/src/types/event.rs +++ b/etl/src/types/event.rs @@ -1,5 +1,8 @@ -use etl_postgres::types::{ColumnSchema, TableId, TableName}; +use etl_postgres::types::{ColumnSchema, SchemaVersion, TableId, TableSchema, TableSchemaDraft}; +use std::collections::HashSet; use std::fmt; +use std::hash::Hash; +use std::sync::Arc; use tokio_postgres::types::PgLsn; use crate::types::TableRow; @@ -65,10 +68,61 @@ pub struct RelationEvent { pub start_lsn: PgLsn, /// LSN position where the transaction of this event will commit. pub commit_lsn: PgLsn, - /// Set of changes for this table (compared to the most recent version of the schema). - pub changes: Vec, /// ID of the table of which this is a schema change. pub table_id: TableId, + /// The old table schema. + pub old_table_schema: Arc, + /// The new table schema. + pub new_table_schema: Arc, +} + +impl RelationEvent { + /// Builds a list of [`RelationChange`]s that describe the changes between the old and new table + /// schemas. + pub fn changes(&self) -> Vec { + // We build a lookup set for the new column schemas for quick change detection. + let mut new_indexed_column_schemas = self + .new_table_schema + .column_schemas + .iter() + .cloned() + .map(IndexedColumnSchema) + .collect::>(); + + // We process all the changes that we want to dispatch to the destination. + let mut changes = vec![]; + for column_schema in self.old_table_schema.column_schemas.iter() { + let column_schema = IndexedColumnSchema(column_schema.clone()); + let latest_column_schema = new_indexed_column_schemas.take(&column_schema); + match latest_column_schema { + Some(latest_column_schema) => { + let column_schema = column_schema.into_inner(); + let latest_column_schema = latest_column_schema.into_inner(); + + if column_schema.name != latest_column_schema.name { + // If we find a column with the same name but different fields, we assume it was changed. The only changes + // that we detect are changes to the column but with preserved name. + changes.push(RelationChange::AlterColumn( + column_schema, + latest_column_schema, + )); + } + } + None => { + // If we don't find the column in the latest schema, we assume it was dropped even + // though it could be renamed. + changes.push(RelationChange::DropColumn(column_schema.into_inner())); + } + } + } + + // For the remaining columns that didn't match, we assume they were added. + for column_schema in new_indexed_column_schemas { + changes.push(RelationChange::AddColumn(column_schema.into_inner())); + } + + changes + } } /// Row insertion event from Postgres logical replication. @@ -83,6 +137,8 @@ pub struct InsertEvent { pub commit_lsn: PgLsn, /// ID of the table where the row was inserted. pub table_id: TableId, + /// Schema version that should be used to interpret this row. + pub schema_version: SchemaVersion, /// Complete row data for the inserted row. pub table_row: TableRow, } @@ -100,6 +156,8 @@ pub struct UpdateEvent { pub commit_lsn: PgLsn, /// ID of the table where the row was updated. pub table_id: TableId, + /// Schema version that should be used to interpret this row. + pub schema_version: SchemaVersion, /// New row data after the update. pub table_row: TableRow, /// Previous row data before the update. @@ -122,6 +180,8 @@ pub struct DeleteEvent { pub commit_lsn: PgLsn, /// ID of the table where the row was deleted. pub table_id: TableId, + /// Schema version that should be used to interpret this row. + pub schema_version: SchemaVersion, /// Data from the deleted row. /// /// The boolean indicates whether the row contains only key columns (`true`) @@ -144,7 +204,7 @@ pub struct TruncateEvent { /// Truncate operation options from Postgres. pub options: i8, /// List of table IDs that were truncated in this operation. - pub rel_ids: Vec, + pub relations: Vec, } /// Represents a single replication event from Postgres logical replication. @@ -192,13 +252,7 @@ impl Event { Event::Update(update_event) => update_event.table_id == *table_id, Event::Delete(delete_event) => delete_event.table_id == *table_id, Event::Relation(relation_event) => relation_event.table_id == *table_id, - Event::Truncate(event) => { - let Some(_) = event.rel_ids.iter().find(|&&id| table_id.0 == id) else { - return false; - }; - - true - } + Event::Truncate(event) => event.relations.iter().any(|rel_id| rel_id == table_id), _ => false, } } @@ -264,3 +318,26 @@ impl From for EventType { (&event).into() } } + +#[derive(Debug, Clone)] +struct IndexedColumnSchema(ColumnSchema); + +impl IndexedColumnSchema { + fn into_inner(self) -> ColumnSchema { + self.0 + } +} + +impl Eq for IndexedColumnSchema {} + +impl PartialEq for IndexedColumnSchema { + fn eq(&self, other: &Self) -> bool { + self.0.name == other.0.name + } +} + +impl Hash for IndexedColumnSchema { + fn hash(&self, state: &mut H) { + self.0.name.hash(state); + } +} diff --git a/etl/tests/postgres_store.rs b/etl/tests/postgres_store.rs index f333e9c18..519438693 100644 --- a/etl/tests/postgres_store.rs +++ b/etl/tests/postgres_store.rs @@ -21,7 +21,7 @@ fn create_sample_table_schema() -> TableSchema { ColumnSchema::new("created_at".to_string(), PgType::TIMESTAMPTZ, -1, false), ]; - TableSchema::new(table_id, table_name, columns) + TableSchema::new(table_id, table_name, 0, columns) } fn create_another_table_schema() -> TableSchema { @@ -32,7 +32,7 @@ fn create_another_table_schema() -> TableSchema { ColumnSchema::new("description".to_string(), PgType::VARCHAR, 255, false), ]; - TableSchema::new(table_id, table_name, columns) + TableSchema::new(table_id, table_name, 0, columns) } #[tokio::test(flavor = "multi_thread")] From 0d079e07f7a3b553374ee77a6e8aa1478a0b37ce Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Tue, 30 Sep 2025 21:31:15 -0700 Subject: [PATCH 16/45] Improve --- etl-destinations/src/bigquery/core.rs | 4 ++-- etl/src/conversions/event.rs | 6 +++--- etl/src/destination/memory.rs | 6 +++--- etl/src/replication/apply.rs | 17 ++++++++--------- etl/src/replication/client.rs | 2 +- etl/src/test_utils/test_destination_wrapper.rs | 6 +++--- etl/src/test_utils/test_schema.rs | 2 +- etl/src/types/event.rs | 6 +++--- 8 files changed, 24 insertions(+), 25 deletions(-) diff --git a/etl-destinations/src/bigquery/core.rs b/etl-destinations/src/bigquery/core.rs index 2ce00df47..67ebf598a 100644 --- a/etl-destinations/src/bigquery/core.rs +++ b/etl-destinations/src/bigquery/core.rs @@ -785,8 +785,8 @@ where let mut truncate_tables = Vec::new(); while let Some(Event::Truncate(_)) = event_iter.peek() { if let Some(Event::Truncate(truncate_event)) = event_iter.next() { - for (table_id, schema_version) in truncate_event.relations { - truncate_tables.push((table_id, schema_version)); + for table_id in truncate_event.table_ids { + truncate_tables.push(table_id); } } } diff --git a/etl/src/conversions/event.rs b/etl/src/conversions/event.rs index 14f55d4cd..271db8026 100644 --- a/etl/src/conversions/event.rs +++ b/etl/src/conversions/event.rs @@ -10,7 +10,7 @@ use crate::bail; use crate::conversions::text::{default_value_for_type, parse_cell_from_postgres_text}; use crate::error::{ErrorKind, EtlError, EtlResult}; use crate::types::{ - BeginEvent, Cell, CommitEvent, DeleteEvent, InsertEvent, RelationEventDraft, TableRow, + BeginEvent, Cell, CommitEvent, DeleteEvent, InsertEvent, TableRow, TruncateEvent, UpdateEvent, }; @@ -194,13 +194,13 @@ pub fn parse_event_from_truncate_message( start_lsn: PgLsn, commit_lsn: PgLsn, truncate_body: &protocol::TruncateBody, - relations: Vec, + table_ids: Vec, ) -> TruncateEvent { TruncateEvent { start_lsn, commit_lsn, options: truncate_body.options(), - relations, + table_ids, } } diff --git a/etl/src/destination/memory.rs b/etl/src/destination/memory.rs index 06976fd74..3bc4565fa 100644 --- a/etl/src/destination/memory.rs +++ b/etl/src/destination/memory.rs @@ -93,12 +93,12 @@ impl Destination for MemoryDestination { if let Event::Truncate(event) = event && has_table_id { - let Some(index) = event.rel_ids.iter().position(|&id| table_id.0 == id) else { + let Some(index) = event.table_ids.iter().position(|&id| id.0 == table_id.0) else { return true; }; - event.rel_ids.remove(index); - if event.rel_ids.is_empty() { + event.table_ids.remove(index); + if event.table_ids.is_empty() { return false; } diff --git a/etl/src/replication/apply.rs b/etl/src/replication/apply.rs index f65cf5b1f..7ccdcfe97 100644 --- a/etl/src/replication/apply.rs +++ b/etl/src/replication/apply.rs @@ -1,11 +1,10 @@ use etl_config::shared::PipelineConfig; use etl_postgres::replication::worker::WorkerType; -use etl_postgres::types::{SchemaVersion, TableId, TableSchema}; +use etl_postgres::types::{TableId, TableSchema}; use futures::StreamExt; use metrics::histogram; use postgres_replication::protocol; use postgres_replication::protocol::{LogicalReplicationMessage, ReplicationMessage}; -use std::collections::HashMap; use std::future::Future; use std::pin::Pin; use std::sync::Arc; @@ -1247,7 +1246,6 @@ async fn handle_truncate_message( state: &mut ApplyLoopState, start_lsn: PgLsn, message: &protocol::TruncateBody, - schema_store: &S, hook: &T, ) -> EtlResult where @@ -1264,23 +1262,24 @@ where // We collect only the relation ids for which we are allowed to apply changes, thus in this case // the truncation. - let mut relations = Vec::with_capacity(message.rel_ids().len()); - for &table_id in message.rel_ids().iter() { - let table_id = TableId::new(table_id); + let mut table_ids = Vec::with_capacity(message.rel_ids().len()); + for table_id in message.rel_ids().iter() { + let table_id = TableId::new(*table_id); if hook .should_apply_changes(table_id, remote_final_lsn) .await? { - relations.push(table_id); + table_ids.push(table_id); } } + // If nothing to apply, skip conversion entirely - if relations.is_empty() { + if table_ids.is_empty() { return Ok(HandleMessageResult::no_event()); } // Convert event from the protocol message. - let event = parse_event_from_truncate_message(start_lsn, remote_final_lsn, message, relations); + let event = parse_event_from_truncate_message(start_lsn, remote_final_lsn, message, table_ids); Ok(HandleMessageResult::return_event(Event::Truncate(event))) } diff --git a/etl/src/replication/client.rs b/etl/src/replication/client.rs index fc41f3168..8d42736ea 100644 --- a/etl/src/replication/client.rs +++ b/etl/src/replication/client.rs @@ -3,7 +3,7 @@ use crate::utils::tokio::MakeRustlsConnect; use crate::{bail, etl_error}; use etl_config::shared::{IntoConnectOptions, PgConnectionConfig}; use etl_postgres::replication::extract_server_version; -use etl_postgres::types::{ColumnSchema, TableId, TableName, TableSchema}; +use etl_postgres::types::{ColumnSchema, TableId, TableName}; use etl_postgres::types::{TableSchemaDraft, convert_type_oid_to_type}; use pg_escape::{quote_identifier, quote_literal}; use postgres_replication::LogicalReplicationStream; diff --git a/etl/src/test_utils/test_destination_wrapper.rs b/etl/src/test_utils/test_destination_wrapper.rs index 3ec76c147..f4e17c60e 100644 --- a/etl/src/test_utils/test_destination_wrapper.rs +++ b/etl/src/test_utils/test_destination_wrapper.rs @@ -161,12 +161,12 @@ where if let Event::Truncate(event) = event && has_table_id { - let Some(index) = event.rel_ids.iter().position(|&id| table_id.0 == id) else { + let Some(index) = event.table_ids.iter().position(|&id| table_id.0 == id) else { return true; }; - event.rel_ids.remove(index); - if event.rel_ids.is_empty() { + event.table_ids.remove(index); + if event.table_ids.is_empty() { return false; } diff --git a/etl/src/test_utils/test_schema.rs b/etl/src/test_utils/test_schema.rs index 4d266ba0d..28059760f 100644 --- a/etl/src/test_utils/test_schema.rs +++ b/etl/src/test_utils/test_schema.rs @@ -222,7 +222,7 @@ pub fn events_equal_excluding_fields(left: &Event, right: &Event) -> bool { } (Event::Relation(left), Event::Relation(right)) => left.table_id == right.table_id, (Event::Truncate(left), Event::Truncate(right)) => { - left.options == right.options && left.rel_ids == right.rel_ids + left.options == right.options && left.table_ids == right.table_ids } (Event::Unsupported, Event::Unsupported) => true, _ => false, // Different event types diff --git a/etl/src/types/event.rs b/etl/src/types/event.rs index 383803a61..01e2b4659 100644 --- a/etl/src/types/event.rs +++ b/etl/src/types/event.rs @@ -1,4 +1,4 @@ -use etl_postgres::types::{ColumnSchema, SchemaVersion, TableId, TableSchema, TableSchemaDraft}; +use etl_postgres::types::{ColumnSchema, SchemaVersion, TableId, TableSchema}; use std::collections::HashSet; use std::fmt; use std::hash::Hash; @@ -204,7 +204,7 @@ pub struct TruncateEvent { /// Truncate operation options from Postgres. pub options: i8, /// List of table IDs that were truncated in this operation. - pub relations: Vec, + pub table_ids: Vec, } /// Represents a single replication event from Postgres logical replication. @@ -252,7 +252,7 @@ impl Event { Event::Update(update_event) => update_event.table_id == *table_id, Event::Delete(delete_event) => delete_event.table_id == *table_id, Event::Relation(relation_event) => relation_event.table_id == *table_id, - Event::Truncate(event) => event.relations.iter().any(|rel_id| rel_id == table_id), + Event::Truncate(event) => event.table_ids.iter().any(|id| id == table_id), _ => false, } } From 8cec27443270df4d7c4d45d3ebffa4d05d00df5c Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Tue, 30 Sep 2025 21:48:16 -0700 Subject: [PATCH 17/45] Improve --- etl/src/replication/apply.rs | 10 +++++++--- etl/src/types/event.rs | 4 ++-- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/etl/src/replication/apply.rs b/etl/src/replication/apply.rs index 7ccdcfe97..df77c1f43 100644 --- a/etl/src/replication/apply.rs +++ b/etl/src/replication/apply.rs @@ -915,7 +915,7 @@ where handle_delete_message(state, start_lsn, delete_body, hook, schema_store).await } LogicalReplicationMessage::Truncate(truncate_body) => { - handle_truncate_message(state, start_lsn, truncate_body, schema_store, hook).await + handle_truncate_message(state, start_lsn, truncate_body, hook, schema_store).await } LogicalReplicationMessage::Origin(_) => { debug!("received unsupported ORIGIN message"); @@ -1247,9 +1247,10 @@ async fn handle_truncate_message( start_lsn: PgLsn, message: &protocol::TruncateBody, hook: &T, + schema_store: &S, ) -> EtlResult where - S: SchemaStore + Clone + Send + 'static, +S: SchemaStore + Clone + Send + 'static, T: ApplyLoopHook, { let Some(remote_final_lsn) = state.remote_final_lsn else { @@ -1269,7 +1270,10 @@ where .should_apply_changes(table_id, remote_final_lsn) .await? { - table_ids.push(table_id); + // We load the last table schema available at the time of the truncation, this way we + // can bind the schema version to use when processing the truncation in the destination. + let table_schema = load_table_schema(schema_store, table_id).await?; + table_ids.push((table_id, table_schema.version)); } } diff --git a/etl/src/types/event.rs b/etl/src/types/event.rs index 01e2b4659..041f12b40 100644 --- a/etl/src/types/event.rs +++ b/etl/src/types/event.rs @@ -204,7 +204,7 @@ pub struct TruncateEvent { /// Truncate operation options from Postgres. pub options: i8, /// List of table IDs that were truncated in this operation. - pub table_ids: Vec, + pub table_ids: Vec<(TableId, SchemaVersion)>, } /// Represents a single replication event from Postgres logical replication. @@ -252,7 +252,7 @@ impl Event { Event::Update(update_event) => update_event.table_id == *table_id, Event::Delete(delete_event) => delete_event.table_id == *table_id, Event::Relation(relation_event) => relation_event.table_id == *table_id, - Event::Truncate(event) => event.table_ids.iter().any(|id| id == table_id), + Event::Truncate(event) => event.table_ids.iter().any(|(id, _)| id == table_id), _ => false, } } From 3871267fa3de6f36c3715759112bf795011f3550 Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Wed, 1 Oct 2025 10:57:30 -0700 Subject: [PATCH 18/45] Improve --- etl-destinations/src/bigquery/core.rs | 473 +++++++++--------- etl/src/conversions/event.rs | 7 +- etl/src/destination/memory.rs | 6 +- etl/src/replication/apply.rs | 22 +- etl/src/store/both/memory.rs | 25 +- etl/src/store/both/postgres.rs | 74 ++- etl/src/store/schema/base.rs | 5 +- etl/src/test_utils/notify.rs | 19 +- etl/src/test_utils/table.rs | 4 +- .../test_utils/test_destination_wrapper.rs | 6 +- etl/src/test_utils/test_schema.rs | 2 + etl/src/types/event.rs | 2 +- etl/tests/failpoints_pipeline.rs | 8 +- etl/tests/pipeline.rs | 12 +- etl/tests/postgres_store.rs | 182 ++++--- 15 files changed, 437 insertions(+), 410 deletions(-) diff --git a/etl-destinations/src/bigquery/core.rs b/etl-destinations/src/bigquery/core.rs index 67ebf598a..f2b8b0571 100644 --- a/etl-destinations/src/bigquery/core.rs +++ b/etl-destinations/src/bigquery/core.rs @@ -3,8 +3,7 @@ use etl::error::{ErrorKind, EtlError, EtlResult}; use etl::store::schema::SchemaStore; use etl::store::state::StateStore; use etl::types::{ - Cell, Event, PgLsn, RelationChange, RelationEvent, SchemaVersion, TableId, TableName, TableRow, - TableSchema, + Cell, Event, PgLsn, RelationEvent, SchemaVersion, TableId, TableName, TableRow, TableSchema, }; use etl::{bail, etl_error}; use gcp_bigquery_client::storage::TableDescriptor; @@ -267,7 +266,7 @@ where async fn prepare_table( &self, table_id: &TableId, - schema_version: SchemaVersion, + schema_version: Option, ) -> EtlResult<(SequencedBigQueryTableId, Arc)> { // We hold the lock for the entire preparation to avoid race conditions since the consistency // of this code path is critical. @@ -275,19 +274,21 @@ where // We load the schema of the table, if present. This is needed to create the table in BigQuery // and also prepare the table descriptor for CDC streaming. - let table_schema = self - .store - .get_table_schema(table_id, schema_version) - .await? - .ok_or_else(|| { - etl_error!( - ErrorKind::MissingTableSchema, - "Table not found in the schema store", - format!( - "The table schema for table {table_id} version {schema_version} was not found in the schema store" - ) - ) - })?; + let table_schema = match schema_version { + Some(schema_version) => { + self.store + .get_table_schema(table_id, schema_version) + .await? + } + None => self.store.get_latest_table_schema(table_id).await?, + } + .ok_or_else(|| { + etl_error!( + ErrorKind::MissingTableSchema, + "Table not found in the schema store", + format!("The table schema for table {table_id} was not found in the schema store") + ) + })?; // We determine the BigQuery table ID for the table together with the current sequence number. let bigquery_table_id = table_name_to_bigquery_table_id(&table_schema.name); @@ -320,7 +321,7 @@ where ); } - // Ensure view points to this sequenced table (uses cache to avoid redundant operations) + // Ensure view points to this sequenced table. self.ensure_view_points_to_table( &mut inner, &bigquery_table_id, @@ -335,7 +336,7 @@ where async fn prepare_table_for_streaming( &self, table_id: &TableId, - schema_version: SchemaVersion, + schema_version: Option, use_cdc_sequence_column: bool, ) -> EtlResult<(SequencedBigQueryTableId, Arc)> { let (sequenced_bigquery_table_id, table_schema) = @@ -439,23 +440,9 @@ where table_id: TableId, mut table_rows: Vec, ) -> EtlResult<()> { - // Prepare table for streaming. - let latest_schema = self - .store - .get_latest_table_schema(&table_id) - .await? - .ok_or_else(|| { - etl_error!( - ErrorKind::MissingTableSchema, - "Table not found in the schema store", - format!( - "The table schema for table {table_id} was not found in the schema store" - ) - ) - })?; - + // For table rows copy, we load the last table schema that we have available. let (sequenced_bigquery_table_id, table_descriptor) = self - .prepare_table_for_streaming(&table_id, latest_schema.version, false) + .prepare_table_for_streaming(&table_id, None, false) .await?; // Add CDC operation type to all rows (no lock needed). @@ -508,16 +495,17 @@ where /// Persists the accumulated CDC batches for each table to BigQuery. async fn process_table_events( &self, - table_id_to_table_rows: HashMap<(TableId, SchemaVersion), Vec>, + schema_version: SchemaVersion, + table_id_to_table_rows: HashMap>, ) -> EtlResult<()> { if table_id_to_table_rows.is_empty() { return Ok(()); } let mut table_batches = Vec::with_capacity(table_id_to_table_rows.len()); - for ((table_id, schema_version), table_rows) in table_id_to_table_rows { + for (table_id, table_rows) in table_id_to_table_rows { let (sequenced_bigquery_table_id, table_descriptor) = self - .prepare_table_for_streaming(&table_id, schema_version, true) + .prepare_table_for_streaming(&table_id, Some(schema_version), true) .await?; let table_batch = self.client.create_table_batch( @@ -551,141 +539,142 @@ where } async fn apply_relation_event(&self, relation_event: RelationEvent) -> EtlResult<()> { - if relation_event.changes.is_empty() { - debug!( - table_id = %relation_event.table_id, - "relation event contained no schema changes; skipping" - ); - - return Ok(()); - } - - let (sequenced_bigquery_table_id, table_schema) = self - .prepare_table(&relation_event.table_id, relation_event.new_schema_version) - .await?; - - let sequenced_table_name = sequenced_bigquery_table_id.to_string(); - let mut primary_key_dirty = false; - for change in relation_event.changes { - match change { - RelationChange::AddColumn(column_schema) => { - let column_name = column_schema.name.clone(); - let is_primary = column_schema.primary; - - self.client - .add_column(&self.dataset_id, &sequenced_table_name, &column_schema) - .await?; - - debug!( - table = %sequenced_table_name, - column = %column_name, - "added column in BigQuery" - ); - - if is_primary { - primary_key_dirty = true; - } - } - RelationChange::DropColumn(column_schema) => { - let column_name = column_schema.name.clone(); - let was_primary = column_schema.primary; - - self.client - .drop_column(&self.dataset_id, &sequenced_table_name, &column_schema.name) - .await?; - - debug!( - table = %sequenced_table_name, - column = %column_name, - "dropped column in BigQuery" - ); - - if was_primary { - primary_key_dirty = true; - } - } - RelationChange::AlterColumn(previous, latest) => { - let old_name = previous.name.clone(); - let new_name = latest.name.clone(); - let renamed = old_name != new_name; - - if renamed { - self.client - .rename_column( - &self.dataset_id, - &sequenced_table_name, - &previous.name, - &latest.name, - ) - .await?; - - debug!( - table = %sequenced_table_name, - old_column = %old_name, - new_column = %new_name, - "renamed column in BigQuery" - ); - } - - if previous.typ != latest.typ { - self.client - .alter_column_type(&self.dataset_id, &sequenced_table_name, &latest) - .await?; - - debug!( - table = %sequenced_table_name, - column = %new_name, - "updated column type in BigQuery" - ); - } - - if previous.nullable != latest.nullable { - self.client - .alter_column_nullability( - &self.dataset_id, - &sequenced_table_name, - &latest.name, - latest.nullable, - ) - .await?; - - debug!( - table = %sequenced_table_name, - column = %new_name, - nullable = latest.nullable, - "updated column nullability in BigQuery" - ); - } - - if previous.primary != latest.primary - || (renamed && (previous.primary || latest.primary)) - { - primary_key_dirty = true; - } - } - } - } - - if primary_key_dirty { - self.client - .sync_primary_key( - &self.dataset_id, - &sequenced_table_name, - &table_schema.column_schemas, - ) - .await?; - - debug!( - table = %sequenced_table_name, - "synchronized primary key definition in BigQuery" - ); - } - - info!( - table_id = %relation_event.table_id, - table = %sequenced_table_name, - "applied relation changes in BigQuery" - ); + // TODO: implement the relation event. + // if relation_event.changes.is_empty() { + // debug!( + // table_id = %relation_event.table_id, + // "relation event contained no schema changes; skipping" + // ); + // + // return Ok(()); + // } + // + // let (sequenced_bigquery_table_id, table_schema) = self + // .prepare_table(&relation_event.table_id, relation_event.new_schema_version) + // .await?; + // + // let sequenced_table_name = sequenced_bigquery_table_id.to_string(); + // let mut primary_key_dirty = false; + // for change in relation_event.changes { + // match change { + // RelationChange::AddColumn(column_schema) => { + // let column_name = column_schema.name.clone(); + // let is_primary = column_schema.primary; + // + // self.client + // .add_column(&self.dataset_id, &sequenced_table_name, &column_schema) + // .await?; + // + // debug!( + // table = %sequenced_table_name, + // column = %column_name, + // "added column in BigQuery" + // ); + // + // if is_primary { + // primary_key_dirty = true; + // } + // } + // RelationChange::DropColumn(column_schema) => { + // let column_name = column_schema.name.clone(); + // let was_primary = column_schema.primary; + // + // self.client + // .drop_column(&self.dataset_id, &sequenced_table_name, &column_schema.name) + // .await?; + // + // debug!( + // table = %sequenced_table_name, + // column = %column_name, + // "dropped column in BigQuery" + // ); + // + // if was_primary { + // primary_key_dirty = true; + // } + // } + // RelationChange::AlterColumn(previous, latest) => { + // let old_name = previous.name.clone(); + // let new_name = latest.name.clone(); + // let renamed = old_name != new_name; + // + // if renamed { + // self.client + // .rename_column( + // &self.dataset_id, + // &sequenced_table_name, + // &previous.name, + // &latest.name, + // ) + // .await?; + // + // debug!( + // table = %sequenced_table_name, + // old_column = %old_name, + // new_column = %new_name, + // "renamed column in BigQuery" + // ); + // } + // + // if previous.typ != latest.typ { + // self.client + // .alter_column_type(&self.dataset_id, &sequenced_table_name, &latest) + // .await?; + // + // debug!( + // table = %sequenced_table_name, + // column = %new_name, + // "updated column type in BigQuery" + // ); + // } + // + // if previous.nullable != latest.nullable { + // self.client + // .alter_column_nullability( + // &self.dataset_id, + // &sequenced_table_name, + // &latest.name, + // latest.nullable, + // ) + // .await?; + // + // debug!( + // table = %sequenced_table_name, + // column = %new_name, + // nullable = latest.nullable, + // "updated column nullability in BigQuery" + // ); + // } + // + // if previous.primary != latest.primary + // || (renamed && (previous.primary || latest.primary)) + // { + // primary_key_dirty = true; + // } + // } + // } + // } + // + // if primary_key_dirty { + // self.client + // .sync_primary_key( + // &self.dataset_id, + // &sequenced_table_name, + // &table_schema.column_schemas, + // ) + // .await?; + // + // debug!( + // table = %sequenced_table_name, + // "synchronized primary key definition in BigQuery" + // ); + // } + // + // info!( + // table_id = %relation_event.table_id, + // table = %sequenced_table_name, + // "applied relation changes in BigQuery" + // ); Ok(()) } @@ -693,11 +682,7 @@ where /// Returns `true` whether the event should break the batch that is being built up when /// writing events. `false` otherwise. fn is_batch_breaker(event: &Event) -> bool { - match event { - Event::Truncate(_) => true, - Event::Relation(_) => true, - _ => false, - } + matches!(event, Event::Truncate(_) | Event::Relation(_)) } /// Processes CDC events in batches with proper ordering and truncate handling. @@ -705,76 +690,86 @@ where /// Groups streaming operations (insert/update/delete) by table and processes them together, /// then handles truncate events separately by creating new versioned tables. async fn write_events(&self, events: Vec) -> EtlResult<()> { - let mut event_iter = events.into_iter().peekable(); - - while event_iter.peek().is_some() { - let mut table_id_to_table_rows: HashMap<(TableId, SchemaVersion), Vec> = - HashMap::new(); - - // Process events until we hit a truncate event or run out of events - while let Some(event) = event_iter.peek() { + let mut events = events.into_iter().peekable(); + + while events.peek().is_some() { + let mut batch_schema_version = None; + let mut table_id_to_table_rows = HashMap::new(); + + let mut handle_dml_statement = + |start_lsn: PgLsn, + commit_lsn: PgLsn, + table_id: TableId, + mut table_row: TableRow, + schema_version: SchemaVersion| { + let sequence_number = generate_sequence_number(start_lsn, commit_lsn); + table_row + .values + .push(BigQueryOperationType::Upsert.into_cell()); + table_row.values.push(Cell::String(sequence_number)); + + let table_rows: &mut Vec = + table_id_to_table_rows.entry(table_id).or_default(); + table_rows.push(table_row); + + batch_schema_version = Some(schema_version); + }; + + // Process events until we hit a batch breaker, or we run out of events. + while let Some(event) = events.peek() { if Self::is_batch_breaker(event) { break; } - let event = event_iter.next().unwrap(); - match event { - Event::Insert(mut insert) => { - let sequence_number = - generate_sequence_number(insert.start_lsn, insert.commit_lsn); - insert - .table_row - .values - .push(BigQueryOperationType::Upsert.into_cell()); - insert.table_row.values.push(Cell::String(sequence_number)); - - let table_rows: &mut Vec = table_id_to_table_rows - .entry((insert.table_id, insert.schema_version)) - .or_default(); - table_rows.push(insert.table_row); + match events.next().unwrap() { + Event::Insert(insert) => { + handle_dml_statement( + insert.start_lsn, + insert.commit_lsn, + insert.table_id, + insert.table_row, + insert.schema_version, + ); } - Event::Update(mut update) => { - let sequence_number = - generate_sequence_number(update.start_lsn, update.commit_lsn); - update - .table_row - .values - .push(BigQueryOperationType::Upsert.into_cell()); - update.table_row.values.push(Cell::String(sequence_number)); - - let table_rows: &mut Vec = table_id_to_table_rows - .entry((update.table_id, update.schema_version)) - .or_default(); - table_rows.push(update.table_row); + Event::Update(update) => { + handle_dml_statement( + update.start_lsn, + update.commit_lsn, + update.table_id, + update.table_row, + update.schema_version, + ); } Event::Delete(delete) => { - let Some((_, mut old_table_row)) = delete.old_table_row else { + let Some((_, old_table_row)) = delete.old_table_row else { info!("the `DELETE` event has no row, so it was skipped"); continue; }; - let sequence_number = - generate_sequence_number(delete.start_lsn, delete.commit_lsn); - old_table_row - .values - .push(BigQueryOperationType::Delete.into_cell()); - old_table_row.values.push(Cell::String(sequence_number)); - - let table_rows: &mut Vec = table_id_to_table_rows - .entry((delete.table_id, delete.schema_version)) - .or_default(); - table_rows.push(old_table_row); + handle_dml_statement( + delete.start_lsn, + delete.commit_lsn, + delete.table_id, + old_table_row, + delete.schema_version, + ); } - _ => { - debug!("skipping unsupported event in BigQuery"); + event => { + debug!("skipping unsupported event {event:?} in BigQuery"); } } } - - // Process accumulated events for each table before a batch breaker is encountered. if !table_id_to_table_rows.is_empty() { - let pending_rows = mem::take(&mut table_id_to_table_rows); - self.process_table_events(pending_rows).await?; + let Some(schema_version) = batch_schema_version.take() else { + bail!( + ErrorKind::InvalidState, + "Missing schema version", + "Missing schema version before writing events in BigQuery" + ) + }; + let table_id_to_table_rows = mem::take(&mut table_id_to_table_rows); + self.process_table_events(schema_version, table_id_to_table_rows) + .await?; } // Collect and deduplicate all table IDs from all truncate events. @@ -783,8 +778,8 @@ where // row without applying other events in the meanwhile, it doesn't make any sense to create // new empty tables for each of them. let mut truncate_tables = Vec::new(); - while let Some(Event::Truncate(_)) = event_iter.peek() { - if let Some(Event::Truncate(truncate_event)) = event_iter.next() { + while let Some(Event::Truncate(_)) = events.peek() { + if let Some(Event::Truncate(truncate_event)) = events.next() { for table_id in truncate_event.table_ids { truncate_tables.push(table_id); } @@ -796,8 +791,8 @@ where } // Collect all the relation events in a sequence and apply them. - while let Some(Event::Relation(_)) = event_iter.peek() { - if let Some(Event::Relation(relation_event)) = event_iter.next() { + while let Some(Event::Relation(_)) = events.peek() { + if let Some(Event::Relation(relation_event)) = events.next() { self.apply_relation_event(relation_event).await?; } } @@ -825,11 +820,13 @@ where .store .get_table_schema(&table_id, schema_version) .await?; + // If we are not doing CDC, it means that this truncation has been issued while recovering // from a failed data sync operation. In that case, we could have failed before table schemas - // were stored in the schema store, so we just continue and emit a warning. If we are doing - // CDC, it's a problem if the schema disappears while streaming, so we error out. - if !is_cdc_truncate { + // were stored in the schema store, so if we don't find a table schema we just continue + // and emit a warning. If we are doing CDC, it's a problem if the schema disappears while + // streaming, so we error out. + if table_schema.is_none() && !is_cdc_truncate { warn!( "the table schema for table {table_id} was not found in the schema store while processing truncate events for BigQuery", table_id = table_id.to_string() diff --git a/etl/src/conversions/event.rs b/etl/src/conversions/event.rs index 271db8026..9a7acc2cc 100644 --- a/etl/src/conversions/event.rs +++ b/etl/src/conversions/event.rs @@ -1,6 +1,6 @@ use core::str; use etl_postgres::types::{ - ColumnSchema, TableId, TableSchema, TableSchemaDraft, convert_type_oid_to_type, + ColumnSchema, SchemaVersion, TableId, TableSchema, TableSchemaDraft, convert_type_oid_to_type, }; use postgres_replication::protocol; use std::sync::Arc; @@ -10,8 +10,7 @@ use crate::bail; use crate::conversions::text::{default_value_for_type, parse_cell_from_postgres_text}; use crate::error::{ErrorKind, EtlError, EtlResult}; use crate::types::{ - BeginEvent, Cell, CommitEvent, DeleteEvent, InsertEvent, TableRow, - TruncateEvent, UpdateEvent, + BeginEvent, Cell, CommitEvent, DeleteEvent, InsertEvent, TableRow, TruncateEvent, UpdateEvent, }; /// Creates a [`BeginEvent`] from Postgres protocol data. @@ -194,7 +193,7 @@ pub fn parse_event_from_truncate_message( start_lsn: PgLsn, commit_lsn: PgLsn, truncate_body: &protocol::TruncateBody, - table_ids: Vec, + table_ids: Vec<(TableId, SchemaVersion)>, ) -> TruncateEvent { TruncateEvent { start_lsn, diff --git a/etl/src/destination/memory.rs b/etl/src/destination/memory.rs index 3bc4565fa..c7cd4529f 100644 --- a/etl/src/destination/memory.rs +++ b/etl/src/destination/memory.rs @@ -93,7 +93,11 @@ impl Destination for MemoryDestination { if let Event::Truncate(event) = event && has_table_id { - let Some(index) = event.table_ids.iter().position(|&id| id.0 == table_id.0) else { + let Some(index) = event + .table_ids + .iter() + .position(|(id, _)| id.0 == table_id.0) + else { return true; }; diff --git a/etl/src/replication/apply.rs b/etl/src/replication/apply.rs index df77c1f43..756e26406 100644 --- a/etl/src/replication/apply.rs +++ b/etl/src/replication/apply.rs @@ -1105,7 +1105,7 @@ where } // Convert event from the protocol message. - let old_table_schema = load_table_schema(schema_store, table_id).await?; + let old_table_schema = load_latest_table_schema(schema_store, table_id).await?; let new_table_schema_draft = parse_event_from_relation_message(old_table_schema.clone(), message).await?; @@ -1155,7 +1155,7 @@ where } // Convert event from the protocol message. - let table_schema = load_table_schema(schema_store, table_id).await?; + let table_schema = load_latest_table_schema(schema_store, table_id).await?; let event = parse_event_from_insert_message(table_schema, start_lsn, remote_final_lsn, message)?; @@ -1192,7 +1192,7 @@ where } // Convert event from the protocol message. - let table_schema = load_table_schema(schema_store, table_id).await?; + let table_schema = load_latest_table_schema(schema_store, table_id).await?; let event = parse_event_from_update_message(table_schema, start_lsn, remote_final_lsn, message)?; @@ -1229,7 +1229,7 @@ where } // Convert event from the protocol message. - let table_schema = load_table_schema(schema_store, table_id).await?; + let table_schema = load_latest_table_schema(schema_store, table_id).await?; let event = parse_event_from_delete_message(table_schema, start_lsn, remote_final_lsn, message)?; @@ -1250,7 +1250,7 @@ async fn handle_truncate_message( schema_store: &S, ) -> EtlResult where -S: SchemaStore + Clone + Send + 'static, + S: SchemaStore + Clone + Send + 'static, T: ApplyLoopHook, { let Some(remote_final_lsn) = state.remote_final_lsn else { @@ -1272,7 +1272,7 @@ S: SchemaStore + Clone + Send + 'static, { // We load the last table schema available at the time of the truncation, this way we // can bind the schema version to use when processing the truncation in the destination. - let table_schema = load_table_schema(schema_store, table_id).await?; + let table_schema = load_latest_table_schema(schema_store, table_id).await?; table_ids.push((table_id, table_schema.version)); } } @@ -1288,8 +1288,14 @@ S: SchemaStore + Clone + Send + 'static, Ok(HandleMessageResult::return_event(Event::Truncate(event))) } -/// Loads the table schema for processing a specific event. -async fn load_table_schema(schema_store: &S, table_id: TableId) -> EtlResult> +/// Loads the latest table schema for processing a specific event. +/// +/// The latest table schema is defined as the schema with the highest version at a specific point +/// of processing the WAL. +async fn load_latest_table_schema( + schema_store: &S, + table_id: TableId, +) -> EtlResult> where S: SchemaStore + Clone + Send + 'static, { diff --git a/etl/src/store/both/memory.rs b/etl/src/store/both/memory.rs index 3fccef886..16f83970c 100644 --- a/etl/src/store/both/memory.rs +++ b/etl/src/store/both/memory.rs @@ -13,19 +13,14 @@ use crate::store::state::StateStore; /// Inner state of [`MemoryStore`] #[derive(Debug)] struct Inner { - /// Current replication state for each table - this is the authoritative source of truth - /// for table states. Every table being replicated must have an entry here. + /// Current replication state for each table. table_replication_states: HashMap, - /// Complete history of state transitions for each table, used for debugging and auditing. - /// This is an append-only log that grows over time and provides visibility into - /// table state evolution. Entries are chronologically ordered. + /// Complete history of state transitions for each table which is used for state rollbacks. table_state_history: HashMap>, - /// Cached table schema definitions, reference-counted for efficient sharing. - /// Schemas are expensive to fetch from Postgres, so they're cached here - /// once retrieved and shared via Arc across the application. + /// Table schema definitions for each table. Each schema has multiple versions which are created + /// when new schema changes are detected. table_schemas: HashMap>>, - /// Mapping from table IDs to human-readable table names for easier debugging - /// and logging. These mappings are established during schema discovery. + /// Mappings between source and destination tables. table_mappings: HashMap, } @@ -197,16 +192,6 @@ impl SchemaStore for MemoryStore { .and_then(|schemas| schemas.iter().next_back().map(|(_, schema)| schema.clone()))) } - async fn get_table_schemas(&self) -> EtlResult>> { - let inner = self.inner.lock().await; - - Ok(inner - .table_schemas - .values() - .filter_map(|schemas| schemas.iter().next_back().map(|(_, schema)| schema.clone())) - .collect()) - } - async fn load_table_schemas(&self) -> EtlResult { let inner = self.inner.lock().await; diff --git a/etl/src/store/both/postgres.rs b/etl/src/store/both/postgres.rs index 5db97f4a6..d88a2d387 100644 --- a/etl/src/store/both/postgres.rs +++ b/etl/src/store/both/postgres.rs @@ -19,8 +19,28 @@ use crate::store::state::StateStore; use crate::types::PipelineId; use crate::{bail, etl_error}; +/// Default number of connections in the pool used by the [`PostgresStore`] to connect to the +/// source database. const NUM_POOL_CONNECTIONS: u32 = 1; +/// Emits table-related metrics. +fn emit_table_metrics( + pipeline_id: PipelineId, + total_tables: usize, + phase_counts: &HashMap<&'static str, u64>, +) { + gauge!(ETL_TABLES_TOTAL, PIPELINE_ID_LABEL => pipeline_id.to_string()).set(total_tables as f64); + + for (phase, count) in phase_counts { + gauge!( + ETL_TABLES_TOTAL, + PIPELINE_ID_LABEL => pipeline_id.to_string(), + PHASE_LABEL => *phase + ) + .set(*count as f64); + } +} + /// Converts ETL table replication phases to Postgres database state format. /// /// This conversion transforms internal ETL replication states into the format @@ -199,7 +219,7 @@ impl PostgresStore { /// than maintaining persistent connections. This approach trades connection /// setup overhead for reduced resource usage during periods of low activity. async fn connect_to_source(&self) -> Result { - // We connect to source database each time we update because we assume that + // We connect to the source database each time we update because we assume that // these updates will be infrequent. It has some overhead to establish a // connection, but it's better than holding a connection open for long periods // when there's little activity on it. @@ -212,23 +232,14 @@ impl PostgresStore { Ok(pool) } -} - -/// Emits table related metrics. -fn emit_table_metrics( - pipeline_id: PipelineId, - total_tables: usize, - phase_counts: &HashMap<&'static str, u64>, -) { - gauge!(ETL_TABLES_TOTAL, PIPELINE_ID_LABEL => pipeline_id.to_string()).set(total_tables as f64); - for (phase, count) in phase_counts { - gauge!( - ETL_TABLES_TOTAL, - PIPELINE_ID_LABEL => pipeline_id.to_string(), - PHASE_LABEL => *phase - ) - .set(*count as f64); + /// Returns all the table schemas stored in this store. + #[cfg(feature = "test-utils")] + pub async fn get_all_table_schemas( + &self, + ) -> HashMap>> { + let inner = self.inner.lock().await; + inner.table_schemas.clone() } } @@ -330,8 +341,7 @@ impl StateStore for PostgresStore { let mut inner = self.inner.lock().await; state::update_replication_state(&pool, self.pipeline_id as i64, table_id, db_state).await?; - // Compute which phases need to be increment and decremented to - // keep table metrics updated + // Compute which phases need to be increment and decremented to keep table metrics updated. let phase_to_decrement = inner .table_states .get(&table_id) @@ -373,8 +383,8 @@ impl StateStore for PostgresStore { let mut inner = self.inner.lock().await; match state::rollback_replication_state(&pool, self.pipeline_id as i64, table_id).await? { Some(restored_row) => { - // Compute which phases need to be increment and decremented to - // keep table metrics updated + // Compute which phases need to be increment and decremented to keep table metrics + // updated. let phase_to_decrement = inner .table_states .get(&table_id) @@ -521,6 +531,7 @@ impl SchemaStore for PostgresStore { .and_then(|schemas| schemas.get(&version).cloned())) } + /// Retrieves the latest table schema for a table. async fn get_latest_table_schema( &self, table_id: &TableId, @@ -533,21 +544,6 @@ impl SchemaStore for PostgresStore { .and_then(|schemas| schemas.iter().next_back().map(|(_, schema)| schema.clone()))) } - /// Retrieves all cached table schemas as a vector. - /// - /// This method returns all currently cached table schemas, providing a - /// complete view of the schema information available to the pipeline. - /// Useful for operations that need to process or analyze all table schemas. - async fn get_table_schemas(&self) -> EtlResult>> { - let inner = self.inner.lock().await; - - Ok(inner - .table_schemas - .values() - .filter_map(|schemas| schemas.iter().next_back().map(|(_, schema)| schema.clone())) - .collect()) - } - /// Loads table schemas from Postgres into memory cache. /// /// This method connects to the source database, retrieves schema information @@ -587,12 +583,6 @@ impl SchemaStore for PostgresStore { table_schemas_len ); - emit_table_metrics( - self.pipeline_id, - inner.table_states.keys().len(), - inner.table_schemas.keys().len(), - ); - Ok(table_schemas_len) } diff --git a/etl/src/store/schema/base.rs b/etl/src/store/schema/base.rs index 227202bda..e437150a4 100644 --- a/etl/src/store/schema/base.rs +++ b/etl/src/store/schema/base.rs @@ -10,7 +10,7 @@ use crate::error::EtlResult; /// /// Implementations should ensure thread-safety and handle concurrent access to the data. pub trait SchemaStore { - /// Returns table schema for table with id `table_id` from the cache. + /// Returns table schema for table with a specific schema version. /// /// Does not load any new data into the cache. fn get_table_schema( @@ -25,9 +25,6 @@ pub trait SchemaStore { table_id: &TableId, ) -> impl Future>>> + Send; - /// Returns the latest table schema for all tables from the cache. - fn get_table_schemas(&self) -> impl Future>>> + Send; - /// Loads table schemas from the persistent state into the cache. /// /// This should be called once the program starts to load the schemas into the cache. diff --git a/etl/src/test_utils/notify.rs b/etl/src/test_utils/notify.rs index 44849455f..0cad5bbfc 100644 --- a/etl/src/test_utils/notify.rs +++ b/etl/src/test_utils/notify.rs @@ -106,12 +106,17 @@ impl NotifyingStore { inner.table_replication_states.clone() } - pub async fn get_table_schemas(&self) -> HashMap { + pub async fn get_latest_table_schemas(&self) -> HashMap { let inner = self.inner.read().await; inner .table_schemas .iter() - .map(|(id, schema)| (*id, Arc::as_ref(schema).clone())) + .filter_map(|(table_id, table_schemas)| { + table_schemas + .iter() + .next_back() + .map(|(_, schema)| (*table_id, schema.as_ref().clone())) + }) .collect() } @@ -318,16 +323,6 @@ impl SchemaStore for NotifyingStore { .and_then(|schemas| schemas.iter().next_back().map(|(_, schema)| schema.clone()))) } - async fn get_table_schemas(&self) -> EtlResult>> { - let inner = self.inner.read().await; - - Ok(inner - .table_schemas - .values() - .filter_map(|schemas| schemas.iter().next_back().map(|(_, schema)| schema.clone())) - .collect()) - } - async fn load_table_schemas(&self) -> EtlResult { let inner = self.inner.read().await; Ok(inner diff --git a/etl/src/test_utils/table.rs b/etl/src/test_utils/table.rs index 50f175fed..af53a086f 100644 --- a/etl/src/test_utils/table.rs +++ b/etl/src/test_utils/table.rs @@ -1,4 +1,4 @@ -use etl_postgres::types::{ColumnSchema, TableId, TableName, TableSchema}; +use etl_postgres::types::{ColumnSchema, TableId, TableName, TableSchema, TableSchemaDraft}; use std::collections::HashMap; /// Return the names of the column schema. @@ -21,7 +21,7 @@ pub fn column_schema_names(table_schema: &TableSchema) -> Vec { /// Panics if the table ID doesn't exist in the provided schemas, or if any aspect /// of the schema doesn't match the expected values. pub fn assert_table_schema( - table_schemas: &HashMap, + table_schemas: &HashMap, table_id: TableId, expected_table_name: TableName, expected_columns: &[ColumnSchema], diff --git a/etl/src/test_utils/test_destination_wrapper.rs b/etl/src/test_utils/test_destination_wrapper.rs index f4e17c60e..4d1a6b17e 100644 --- a/etl/src/test_utils/test_destination_wrapper.rs +++ b/etl/src/test_utils/test_destination_wrapper.rs @@ -161,7 +161,11 @@ where if let Event::Truncate(event) = event && has_table_id { - let Some(index) = event.table_ids.iter().position(|&id| table_id.0 == id) else { + let Some(index) = event + .table_ids + .iter() + .position(|(id, _)| id.0 == table_id.0) + else { return true; }; diff --git a/etl/src/test_utils/test_schema.rs b/etl/src/test_utils/test_schema.rs index 28059760f..b6b681de5 100644 --- a/etl/src/test_utils/test_schema.rs +++ b/etl/src/test_utils/test_schema.rs @@ -248,6 +248,7 @@ pub fn build_expected_users_inserts( Cell::I32(age), ], }, + schema_version: 0, })); starting_id += 1; @@ -271,6 +272,7 @@ pub fn build_expected_orders_inserts( table_row: TableRow { values: vec![Cell::I64(starting_id), Cell::String(name.to_owned())], }, + schema_version: 0, })); starting_id += 1; diff --git a/etl/src/types/event.rs b/etl/src/types/event.rs index 041f12b40..682cf1a7b 100644 --- a/etl/src/types/event.rs +++ b/etl/src/types/event.rs @@ -79,7 +79,7 @@ pub struct RelationEvent { impl RelationEvent { /// Builds a list of [`RelationChange`]s that describe the changes between the old and new table /// schemas. - pub fn changes(&self) -> Vec { + pub fn build_changes(&self) -> Vec { // We build a lookup set for the new column schemas for quick change detection. let mut new_indexed_column_schemas = self .new_table_schema diff --git a/etl/tests/failpoints_pipeline.rs b/etl/tests/failpoints_pipeline.rs index 0f597b64d..8a0efd20f 100644 --- a/etl/tests/failpoints_pipeline.rs +++ b/etl/tests/failpoints_pipeline.rs @@ -74,7 +74,7 @@ async fn table_copy_fails_after_data_sync_threw_an_error_with_no_retry() { assert!(table_rows.is_empty()); // Verify table schemas were correctly stored. - let table_schemas = store.get_table_schemas().await; + let table_schemas = store.get_latest_table_schemas().await; assert!(table_schemas.is_empty()); } @@ -144,7 +144,7 @@ async fn table_copy_fails_after_timed_retry_exceeded_max_attempts() { assert!(table_rows.is_empty()); // Verify table schemas were correctly stored. - let table_schemas = store.get_table_schemas().await; + let table_schemas = store.get_latest_table_schemas().await; assert!(table_schemas.is_empty()); } @@ -205,7 +205,7 @@ async fn table_copy_is_consistent_after_data_sync_threw_an_error_with_timed_retr assert_eq!(users_table_rows.len(), rows_inserted); // Verify table schemas were correctly stored. - let table_schemas = store.get_table_schemas().await; + let table_schemas = store.get_latest_table_schemas().await; assert_eq!(table_schemas.len(), 1); assert_eq!( *table_schemas @@ -268,7 +268,7 @@ async fn table_copy_is_consistent_during_data_sync_threw_an_error_with_timed_ret assert_eq!(users_table_rows.len(), rows_inserted); // Verify table schemas were correctly stored. - let table_schemas = store.get_table_schemas().await; + let table_schemas = store.get_latest_table_schemas().await; assert_eq!(table_schemas.len(), 1); assert_eq!( *table_schemas diff --git a/etl/tests/pipeline.rs b/etl/tests/pipeline.rs index dd683925b..c21242f06 100644 --- a/etl/tests/pipeline.rs +++ b/etl/tests/pipeline.rs @@ -82,7 +82,7 @@ async fn table_schema_copy_survives_pipeline_restarts() { ); // We check that the table schemas have been stored. - let table_schemas = store.get_table_schemas().await; + let table_schemas = store.get_latest_table_schemas().await; assert_eq!(table_schemas.len(), 2); assert_eq!( *table_schemas @@ -371,7 +371,7 @@ async fn publication_for_all_tables_in_schema_ignores_new_tables_until_restart() pipeline.shutdown_and_wait().await.unwrap(); // Check that only the schemas of the first table were stored. - let table_schemas = store.get_table_schemas().await; + let table_schemas = store.get_latest_table_schemas().await; assert_eq!(table_schemas.len(), 1); assert!(table_schemas.contains_key(&table_1_id)); assert!(!table_schemas.contains_key(&table_2_id)); @@ -397,7 +397,7 @@ async fn publication_for_all_tables_in_schema_ignores_new_tables_until_restart() pipeline.shutdown_and_wait().await.unwrap(); // Check that both schemas exist. - let table_schemas = store.get_table_schemas().await; + let table_schemas = store.get_latest_table_schemas().await; assert_eq!(table_schemas.len(), 2); assert!(table_schemas.contains_key(&table_1_id)); assert!(table_schemas.contains_key(&table_2_id)); @@ -525,7 +525,7 @@ async fn table_schema_changes_are_handled_correctly() { users_state_notify.notified().await; // Check the initial schema. - let table_schemas = store.get_table_schemas().await; + let table_schemas = store.get_latest_table_schemas().await; assert_eq!(table_schemas.len(), 1); let users_table_schema = column_schema_names( table_schemas @@ -575,7 +575,7 @@ async fn table_schema_changes_are_handled_correctly() { insert_event_notify.notified().await; // Check the updated schema. - let table_schemas = store.get_table_schemas().await; + let table_schemas = store.get_latest_table_schemas().await; assert_eq!(table_schemas.len(), 1); let users_table_schema = column_schema_names( table_schemas @@ -625,7 +625,7 @@ async fn table_schema_changes_are_handled_correctly() { insert_event_notify.notified().await; // Check the updated schema. - let table_schemas = store.get_table_schemas().await; + let table_schemas = store.get_latest_table_schemas().await; assert_eq!(table_schemas.len(), 1); let users_table_schema = column_schema_names( table_schemas diff --git a/etl/tests/postgres_store.rs b/etl/tests/postgres_store.rs index 519438693..ddc81e913 100644 --- a/etl/tests/postgres_store.rs +++ b/etl/tests/postgres_store.rs @@ -7,12 +7,12 @@ use etl::store::schema::SchemaStore; use etl::store::state::StateStore; use etl::test_utils::database::spawn_source_database_for_store; use etl_postgres::replication::connect_to_source_database; -use etl_postgres::types::{ColumnSchema, TableId, TableName, TableSchema}; +use etl_postgres::types::{ColumnSchema, TableId, TableName, TableSchemaDraft}; use etl_telemetry::tracing::init_test_tracing; use sqlx::postgres::types::Oid as SqlxTableId; use tokio_postgres::types::{PgLsn, Type as PgType}; -fn create_sample_table_schema() -> TableSchema { +fn create_sample_table_schema() -> TableSchemaDraft { let table_id = TableId::new(12345); let table_name = TableName::new("public".to_string(), "test_table".to_string()); let columns = vec![ @@ -21,10 +21,10 @@ fn create_sample_table_schema() -> TableSchema { ColumnSchema::new("created_at".to_string(), PgType::TIMESTAMPTZ, -1, false), ]; - TableSchema::new(table_id, table_name, 0, columns) + TableSchemaDraft::new(table_id, table_name, columns) } -fn create_another_table_schema() -> TableSchema { +fn create_another_table_schema() -> TableSchemaDraft { let table_id = TableId::new(67890); let table_name = TableName::new("public".to_string(), "another_table".to_string()); let columns = vec![ @@ -32,7 +32,7 @@ fn create_another_table_schema() -> TableSchema { ColumnSchema::new("description".to_string(), PgType::VARCHAR, 255, false), ]; - TableSchema::new(table_id, table_name, 0, columns) + TableSchemaDraft::new(table_id, table_name, columns) } #[tokio::test(flavor = "multi_thread")] @@ -177,8 +177,8 @@ async fn test_state_store_load_states() { let database = spawn_source_database_for_store().await; let pipeline_id = 1; - let table_id1 = TableId::new(12345); - let table_id2 = TableId::new(67890); + let table_id_1 = TableId::new(12345); + let table_id_2 = TableId::new(67890); let store = PostgresStore::new(pipeline_id, database.config.clone()); @@ -187,11 +187,11 @@ async fn test_state_store_load_states() { let data_sync_phase = TableReplicationPhase::DataSync; store - .update_table_replication_state(table_id1, init_phase.clone()) + .update_table_replication_state(table_id_1, init_phase.clone()) .await .unwrap(); store - .update_table_replication_state(table_id2, data_sync_phase.clone()) + .update_table_replication_state(table_id_2, data_sync_phase.clone()) .await .unwrap(); @@ -209,8 +209,8 @@ async fn test_state_store_load_states() { // Verify loaded states let states = new_store.get_table_replication_states().await.unwrap(); assert_eq!(states.len(), 2); - assert_eq!(states.get(&table_id1), Some(&init_phase)); - assert_eq!(states.get(&table_id2), Some(&data_sync_phase)); + assert_eq!(states.get(&table_id_1), Some(&init_phase)); + assert_eq!(states.get(&table_id_2), Some(&data_sync_phase)); } #[tokio::test(flavor = "multi_thread")] @@ -225,19 +225,22 @@ async fn test_schema_store_operations() { let table_id = table_schema.id; // Test initial state - should be empty - let schema = store.get_table_schema(&table_id).await.unwrap(); + let schema = store.get_table_schema(&table_id, 0).await.unwrap(); assert!(schema.is_none()); - let all_schemas = store.get_table_schemas().await.unwrap(); + let all_schemas = store.get_all_table_schemas().await; assert!(all_schemas.is_empty()); // Test storing schema - store + let stored_schema = store .store_table_schema(table_schema.clone()) .await .unwrap(); - let schema = store.get_table_schema(&table_id).await.unwrap(); + let schema = store + .get_table_schema(&table_id, stored_schema.version) + .await + .unwrap(); assert!(schema.is_some()); let schema = schema.unwrap(); assert_eq!(schema.id, table_schema.id); @@ -247,7 +250,7 @@ async fn test_schema_store_operations() { table_schema.column_schemas.len() ); - let all_schemas = store.get_table_schemas().await.unwrap(); + let all_schemas = store.get_all_table_schemas().await; assert_eq!(all_schemas.len(), 1); // Test storing another schema @@ -257,7 +260,7 @@ async fn test_schema_store_operations() { .await .unwrap(); - let all_schemas = store.get_table_schemas().await.unwrap(); + let all_schemas = store.get_all_table_schemas().await; assert_eq!(all_schemas.len(), 2); } @@ -273,11 +276,11 @@ async fn test_schema_store_load_schemas() { let table_schema2 = create_another_table_schema(); // Store schemas - store + let stored_schema_1 = store .store_table_schema(table_schema1.clone()) .await .unwrap(); - store + let stored_schema_2 = store .store_table_schema(table_schema2.clone()) .await .unwrap(); @@ -286,7 +289,7 @@ async fn test_schema_store_load_schemas() { let new_store = PostgresStore::new(pipeline_id, database.config.clone()); // Initially empty (not loaded yet) - let schemas = new_store.get_table_schemas().await.unwrap(); + let schemas = new_store.get_all_table_schemas().await; assert!(schemas.is_empty()); // Load schemas from database @@ -294,16 +297,22 @@ async fn test_schema_store_load_schemas() { assert_eq!(loaded_count, 2); // Verify loaded schemas - let schemas = new_store.get_table_schemas().await.unwrap(); + let schemas = new_store.get_all_table_schemas().await; assert_eq!(schemas.len(), 2); - let schema1 = new_store.get_table_schema(&table_schema1.id).await.unwrap(); + let schema1 = new_store + .get_table_schema(&table_schema1.id, stored_schema_1.version) + .await + .unwrap(); assert!(schema1.is_some()); let schema1 = schema1.unwrap(); assert_eq!(schema1.id, table_schema1.id); assert_eq!(schema1.name, table_schema1.name); - let schema2 = new_store.get_table_schema(&table_schema2.id).await.unwrap(); + let schema2 = new_store + .get_table_schema(&table_schema2.id, stored_schema_2.version) + .await + .unwrap(); assert!(schema2.is_some()); let schema2 = schema2.unwrap(); assert_eq!(schema2.id, table_schema2.id); @@ -321,7 +330,7 @@ async fn test_schema_store_update_existing() { let mut table_schema = create_sample_table_schema(); // Store initial schema - store + let stored_schema = store .store_table_schema(table_schema.clone()) .await .unwrap(); @@ -341,7 +350,10 @@ async fn test_schema_store_update_existing() { .unwrap(); // Verify updated schema - let schema = store.get_table_schema(&table_schema.id).await.unwrap(); + let schema = store + .get_table_schema(&table_schema.id, stored_schema.version) + .await + .unwrap(); assert!(schema.is_some()); let schema = schema.unwrap(); assert_eq!(schema.column_schemas.len(), 4); // Original 3 + 1 new column @@ -363,51 +375,57 @@ async fn test_multiple_pipelines_isolation() { let pipeline_id2 = 2; let table_id = TableId::new(12345); - let store1 = PostgresStore::new(pipeline_id1, database.config.clone()); - let store2 = PostgresStore::new(pipeline_id2, database.config.clone()); + let store_1 = PostgresStore::new(pipeline_id1, database.config.clone()); + let store_2 = PostgresStore::new(pipeline_id2, database.config.clone()); // Add state to pipeline 1 let init_phase = TableReplicationPhase::Init; - store1 + store_1 .update_table_replication_state(table_id, init_phase.clone()) .await .unwrap(); // Add different state to pipeline 2 for the same table let data_sync_phase = TableReplicationPhase::DataSync; - store2 + store_2 .update_table_replication_state(table_id, data_sync_phase.clone()) .await .unwrap(); // Verify isolation - each pipeline sees only its own state - let state1 = store1.get_table_replication_state(table_id).await.unwrap(); + let state1 = store_1.get_table_replication_state(table_id).await.unwrap(); assert_eq!(state1, Some(init_phase)); - let state2 = store2.get_table_replication_state(table_id).await.unwrap(); + let state2 = store_2.get_table_replication_state(table_id).await.unwrap(); assert_eq!(state2, Some(data_sync_phase)); // Test schema isolation let table_schema1 = create_sample_table_schema(); let table_schema2 = create_another_table_schema(); - store1 + let stored_schema_1 = store_1 .store_table_schema(table_schema1.clone()) .await .unwrap(); - store2 + let stored_schema_2 = store_2 .store_table_schema(table_schema2.clone()) .await .unwrap(); // Each pipeline sees only its own schemas - let schemas1 = store1.get_table_schemas().await.unwrap(); - assert_eq!(schemas1.len(), 1); - assert_eq!(schemas1[0].id, table_schema1.id); + let schemas_1 = store_1.get_all_table_schemas().await; + assert_eq!(schemas_1.len(), 1); + assert_eq!( + schemas_1[&table_schema1.id][&stored_schema_1.version], + stored_schema_1 + ); - let schemas2 = store2.get_table_schemas().await.unwrap(); - assert_eq!(schemas2.len(), 1); - assert_eq!(schemas2[0].id, table_schema2.id); + let schemas_2 = store_2.get_all_table_schemas().await; + assert_eq!(schemas_2.len(), 1); + assert_eq!( + schemas_2[&table_schema1.id][&stored_schema_2.version], + stored_schema_2 + ); } #[tokio::test(flavor = "multi_thread")] @@ -535,11 +553,11 @@ async fn test_table_mappings_basic_operations() { let store = PostgresStore::new(pipeline_id, database.config.clone()); - let table_id1 = TableId::new(12345); - let table_id2 = TableId::new(67890); + let table_id_1 = TableId::new(12345); + let table_id_2 = TableId::new(67890); // Test initial state - should be empty - let mapping = store.get_table_mapping(&table_id1).await.unwrap(); + let mapping = store.get_table_mapping(&table_id_1).await.unwrap(); assert!(mapping.is_none()); let all_mappings = store.get_table_mappings().await.unwrap(); @@ -547,33 +565,33 @@ async fn test_table_mappings_basic_operations() { // Test storing and retrieving mappings store - .store_table_mapping(table_id1, "public_users_1".to_string()) + .store_table_mapping(table_id_1, "public_users_1".to_string()) .await .unwrap(); store - .store_table_mapping(table_id2, "public_orders_2".to_string()) + .store_table_mapping(table_id_2, "public_orders_2".to_string()) .await .unwrap(); let all_mappings = store.get_table_mappings().await.unwrap(); assert_eq!(all_mappings.len(), 2); assert_eq!( - all_mappings.get(&table_id1), + all_mappings.get(&table_id_1), Some(&"public_users_1".to_string()) ); assert_eq!( - all_mappings.get(&table_id2), + all_mappings.get(&table_id_2), Some(&"public_orders_2".to_string()) ); // Test updating an existing mapping (upsert) store - .store_table_mapping(table_id1, "public_users_1_updated".to_string()) + .store_table_mapping(table_id_1, "public_users_1_updated".to_string()) .await .unwrap(); - let mapping = store.get_table_mapping(&table_id1).await.unwrap(); + let mapping = store.get_table_mapping(&table_id_1).await.unwrap(); assert_eq!(mapping, Some("public_users_1_updated".to_string())); } @@ -628,34 +646,34 @@ async fn test_table_mappings_pipeline_isolation() { let pipeline_id1 = 1; let pipeline_id2 = 2; - let store1 = PostgresStore::new(pipeline_id1, database.config.clone()); - let store2 = PostgresStore::new(pipeline_id2, database.config.clone()); + let store_1 = PostgresStore::new(pipeline_id1, database.config.clone()); + let store_2 = PostgresStore::new(pipeline_id2, database.config.clone()); let table_id = TableId::new(12345); // Store different mappings for the same table ID in different pipelines - store1 + store_1 .store_table_mapping(table_id, "pipeline1_table".to_string()) .await .unwrap(); - store2 + store_2 .store_table_mapping(table_id, "pipeline2_table".to_string()) .await .unwrap(); // Verify isolation - each pipeline sees only its own mapping - let mapping1 = store1.get_table_mapping(&table_id).await.unwrap(); + let mapping1 = store_1.get_table_mapping(&table_id).await.unwrap(); assert_eq!(mapping1, Some("pipeline1_table".to_string())); - let mapping2 = store2.get_table_mapping(&table_id).await.unwrap(); + let mapping2 = store_2.get_table_mapping(&table_id).await.unwrap(); assert_eq!(mapping2, Some("pipeline2_table".to_string())); // Verify isolation persists after loading from database - let new_store1 = PostgresStore::new(pipeline_id1, database.config.clone()); - new_store1.load_table_mappings().await.unwrap(); + let new_store_1 = PostgresStore::new(pipeline_id1, database.config.clone()); + new_store_1.load_table_mappings().await.unwrap(); - let loaded_mapping1 = new_store1.get_table_mapping(&table_id).await.unwrap(); + let loaded_mapping1 = new_store_1.get_table_mapping(&table_id).await.unwrap(); assert_eq!(loaded_mapping1, Some("pipeline1_table".to_string())); } @@ -684,11 +702,11 @@ async fn test_cleanup_deletes_state_schema_and_mapping_for_table() { .await .unwrap(); - store + let stored_schema_1 = store .store_table_schema(table_1_schema.clone()) .await .unwrap(); - store + let stored_schema_2 = store .store_table_schema(table_2_schema.clone()) .await .unwrap(); @@ -710,7 +728,13 @@ async fn test_cleanup_deletes_state_schema_and_mapping_for_table() { .unwrap() .is_some() ); - assert!(store.get_table_schema(&table_1_id).await.unwrap().is_some()); + assert!( + store + .get_table_schema(&table_1_id, stored_schema_1.version) + .await + .unwrap() + .is_some() + ); assert!( store .get_table_mapping(&table_1_id) @@ -730,7 +754,13 @@ async fn test_cleanup_deletes_state_schema_and_mapping_for_table() { .unwrap() .is_none() ); - assert!(store.get_table_schema(&table_1_id).await.unwrap().is_none()); + assert!( + store + .get_table_schema(&table_1_id, stored_schema_1.version) + .await + .unwrap() + .is_none() + ); assert!( store .get_table_mapping(&table_1_id) @@ -747,7 +777,13 @@ async fn test_cleanup_deletes_state_schema_and_mapping_for_table() { .unwrap() .is_some() ); - assert!(store.get_table_schema(&table_2_id).await.unwrap().is_some()); + assert!( + store + .get_table_schema(&table_2_id, stored_schema_2.version) + .await + .unwrap() + .is_some() + ); assert!( store .get_table_mapping(&table_2_id) @@ -772,7 +808,7 @@ async fn test_cleanup_deletes_state_schema_and_mapping_for_table() { ); assert!( new_store - .get_table_schema(&table_1_id) + .get_table_schema(&table_1_id, stored_schema_1.version) .await .unwrap() .is_none() @@ -795,7 +831,7 @@ async fn test_cleanup_deletes_state_schema_and_mapping_for_table() { ); assert!( new_store - .get_table_schema(&table_2_id) + .get_table_schema(&table_2_id, stored_schema_2.version) .await .unwrap() .is_some() @@ -828,7 +864,13 @@ async fn test_cleanup_idempotent_when_no_state_present() { .unwrap() .is_none() ); - assert!(store.get_table_schema(&table_id).await.unwrap().is_none()); + assert!( + store + .get_table_schema(&table_id, 0) + .await + .unwrap() + .is_none() + ); assert!(store.get_table_mapping(&table_id).await.unwrap().is_none()); // Calling cleanup should succeed even if nothing exists @@ -839,7 +881,7 @@ async fn test_cleanup_idempotent_when_no_state_present() { .update_table_replication_state(table_id, TableReplicationPhase::Init) .await .unwrap(); - store.store_table_schema(table_schema).await.unwrap(); + let stored_schema = store.store_table_schema(table_schema).await.unwrap(); store .store_table_mapping(table_id, "dest_table".to_string()) .await @@ -855,6 +897,12 @@ async fn test_cleanup_idempotent_when_no_state_present() { .unwrap() .is_none() ); - assert!(store.get_table_schema(&table_id).await.unwrap().is_none()); + assert!( + store + .get_table_schema(&table_id, stored_schema.version) + .await + .unwrap() + .is_none() + ); assert!(store.get_table_mapping(&table_id).await.unwrap().is_none()); } From 3aeab781ec98190635cf0fd571159231e7785f53 Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Wed, 1 Oct 2025 12:31:11 -0700 Subject: [PATCH 19/45] Improve --- etl-destinations/src/bigquery/core.rs | 238 ++++++++++-------- etl-postgres/src/replication/schema.rs | 20 +- .../20251001000100_add_schema_versions.sql | 4 +- etl/src/store/both/postgres.rs | 10 +- etl/tests/postgres_store.rs | 67 ++--- 5 files changed, 195 insertions(+), 144 deletions(-) diff --git a/etl-destinations/src/bigquery/core.rs b/etl-destinations/src/bigquery/core.rs index f2b8b0571..80a33cfa7 100644 --- a/etl-destinations/src/bigquery/core.rs +++ b/etl-destinations/src/bigquery/core.rs @@ -538,7 +538,7 @@ where Ok(()) } - async fn apply_relation_event(&self, relation_event: RelationEvent) -> EtlResult<()> { + async fn apply_relation_event(&self, _relation_event: RelationEvent) -> EtlResult<()> { // TODO: implement the relation event. // if relation_event.changes.is_empty() { // debug!( @@ -679,125 +679,163 @@ where Ok(()) } - /// Returns `true` whether the event should break the batch that is being built up when - /// writing events. `false` otherwise. - fn is_batch_breaker(event: &Event) -> bool { - matches!(event, Event::Truncate(_) | Event::Relation(_)) + #[inline] + fn push_dml_statement( + table_id_to_table_rows: &mut HashMap>, + batch_schema_version: &mut Option, + start_lsn: PgLsn, + commit_lsn: PgLsn, + table_id: TableId, + mut table_row: TableRow, + schema_version: SchemaVersion, + operation_type: BigQueryOperationType, + ) -> EtlResult<()> { + // BigQuery CDC extras. + let sequence_number = generate_sequence_number(start_lsn, commit_lsn); + table_row.values.push(operation_type.into_cell()); + table_row.values.push(Cell::String(sequence_number)); + + // Preserve per-table ordering. + table_id_to_table_rows + .entry(table_id) + .or_default() + .push(table_row); + + // Ensure a single schema version per batch. + // + // We need to do this since we don't make any assumptions on relation events being there + // so we use the schema version of the first element that we find. + // + // The invariant that must be upheld is that for all events in a batch, they must all have + // the same schema version. + match batch_schema_version { + Some(batch_schema_version) => { + if schema_version != *batch_schema_version { + bail!( + ErrorKind::InvalidState, + "Multiple schema versions in the same batch", + "Multiple schema versions in the same batch were found while processing events for BigQuery" + ) + } + } + None => { + *batch_schema_version = Some(schema_version); + } + } + + Ok(()) + } + + async fn flush_batch( + &self, + batch_schema_version: &mut Option, + table_id_to_table_rows: &mut HashMap>, + ) -> EtlResult<()> { + if table_id_to_table_rows.is_empty() { + return Ok(()); + } + + let Some(schema_version) = batch_schema_version.take() else { + bail!( + ErrorKind::InvalidState, + "Missing schema version", + "Missing schema version before writing events in BigQuery" + ); + }; + + let rows = mem::take(table_id_to_table_rows); + self.process_table_events(schema_version, rows).await } /// Processes CDC events in batches with proper ordering and truncate handling. /// /// Groups streaming operations (insert/update/delete) by table and processes them together, /// then handles truncate events separately by creating new versioned tables. - async fn write_events(&self, events: Vec) -> EtlResult<()> { - let mut events = events.into_iter().peekable(); - - while events.peek().is_some() { - let mut batch_schema_version = None; - let mut table_id_to_table_rows = HashMap::new(); - - let mut handle_dml_statement = - |start_lsn: PgLsn, - commit_lsn: PgLsn, - table_id: TableId, - mut table_row: TableRow, - schema_version: SchemaVersion| { - let sequence_number = generate_sequence_number(start_lsn, commit_lsn); - table_row - .values - .push(BigQueryOperationType::Upsert.into_cell()); - table_row.values.push(Cell::String(sequence_number)); - - let table_rows: &mut Vec = - table_id_to_table_rows.entry(table_id).or_default(); - table_rows.push(table_row); - - batch_schema_version = Some(schema_version); - }; - - // Process events until we hit a batch breaker, or we run out of events. - while let Some(event) = events.peek() { - if Self::is_batch_breaker(event) { - break; + pub async fn write_events(&self, events: Vec) -> EtlResult<()> { + // Accumulates rows for the current batch, grouped by table. + let mut table_id_to_table_rows: HashMap> = HashMap::new(); + + // The schema version used for the *current* batch (must be consistent within a batch). + let mut batch_schema_version: Option = None; + + // Process stream. + for event in events { + match event { + // DML events. + Event::Insert(insert) => { + Self::push_dml_statement( + &mut table_id_to_table_rows, + &mut batch_schema_version, + insert.start_lsn, + insert.commit_lsn, + insert.table_id, + insert.table_row, + insert.schema_version, + BigQueryOperationType::Upsert, + )?; } - - match events.next().unwrap() { - Event::Insert(insert) => { - handle_dml_statement( - insert.start_lsn, - insert.commit_lsn, - insert.table_id, - insert.table_row, - insert.schema_version, - ); - } - Event::Update(update) => { - handle_dml_statement( - update.start_lsn, - update.commit_lsn, - update.table_id, - update.table_row, - update.schema_version, - ); - } - Event::Delete(delete) => { - let Some((_, old_table_row)) = delete.old_table_row else { - info!("the `DELETE` event has no row, so it was skipped"); - continue; - }; - - handle_dml_statement( + Event::Update(update) => { + Self::push_dml_statement( + &mut table_id_to_table_rows, + &mut batch_schema_version, + update.start_lsn, + update.commit_lsn, + update.table_id, + update.table_row, + update.schema_version, + BigQueryOperationType::Upsert, + )?; + } + Event::Delete(delete) => { + if let Some((_, old_row)) = delete.old_table_row { + Self::push_dml_statement( + &mut table_id_to_table_rows, + &mut batch_schema_version, delete.start_lsn, delete.commit_lsn, delete.table_id, - old_table_row, + old_row, delete.schema_version, - ); - } - event => { - debug!("skipping unsupported event {event:?} in BigQuery"); + BigQueryOperationType::Delete, + )?; + } else { + warn!("the `DELETE` event has no row, so it was skipped"); } } - } - if !table_id_to_table_rows.is_empty() { - let Some(schema_version) = batch_schema_version.take() else { - bail!( - ErrorKind::InvalidState, - "Missing schema version", - "Missing schema version before writing events in BigQuery" - ) - }; - let table_id_to_table_rows = mem::take(&mut table_id_to_table_rows); - self.process_table_events(schema_version, table_id_to_table_rows) - .await?; - } - // Collect and deduplicate all table IDs from all truncate events. - // - // This is done as an optimization since if we have multiple table ids being truncated in a - // row without applying other events in the meanwhile, it doesn't make any sense to create - // new empty tables for each of them. - let mut truncate_tables = Vec::new(); - while let Some(Event::Truncate(_)) = events.peek() { - if let Some(Event::Truncate(truncate_event)) = events.next() { - for table_id in truncate_event.table_ids { - truncate_tables.push(table_id); - } + // Batch breaker events. + Event::Relation(relation) => { + // Finish current batch before applying schema change. + self.flush_batch(&mut batch_schema_version, &mut table_id_to_table_rows) + .await?; + + // We mark the new batch schema version with the relation schema version, since + // after a relation message a new schema is meant to be stored in the schema store. + batch_schema_version = Some(relation.new_table_schema.version); + + // Apply relation change, then prime the next batch with the new schema version. + self.apply_relation_event(relation).await?; + } + Event::Truncate(truncate) => { + // Finish current batch before a TRUNCATE (it affects table state). + self.flush_batch(&mut batch_schema_version, &mut table_id_to_table_rows) + .await?; + + self.process_truncate_for_table_ids(truncate.table_ids.into_iter(), true) + .await?; } - } - if !truncate_tables.is_empty() { - self.process_truncate_for_table_ids(truncate_tables.into_iter(), true) - .await?; - } - // Collect all the relation events in a sequence and apply them. - while let Some(Event::Relation(_)) = events.peek() { - if let Some(Event::Relation(relation_event)) = events.next() { - self.apply_relation_event(relation_event).await?; + // Unsupported events. + other => { + debug!("skipping unsupported event {other:?} in BigQuery"); } } } + // Flush any trailing DML. + self.flush_batch(&mut batch_schema_version, &mut table_id_to_table_rows) + .await?; + Ok(()) } diff --git a/etl-postgres/src/replication/schema.rs b/etl-postgres/src/replication/schema.rs index 8502872e5..02556eb54 100644 --- a/etl-postgres/src/replication/schema.rs +++ b/etl-postgres/src/replication/schema.rs @@ -8,6 +8,9 @@ use crate::types::{ ColumnSchema, SchemaVersion, TableId, TableName, TableSchema, TableSchemaDraft, }; +/// The initial schema version number. +const STARTING_SCHEMA_VERSION: u64 = 0; + macro_rules! define_type_mappings { ( $( @@ -147,12 +150,11 @@ pub async fn store_table_schema( ) -> Result { let mut tx = pool.begin().await?; - let current_version: Option = sqlx::query_scalar( + let current_schema_version: Option = sqlx::query_scalar( r#" select max(schema_version) from etl.table_schemas - where pipeline_id = $1 - and table_id = $2 + where pipeline_id = $1 and table_id = $2 "#, ) .bind(pipeline_id) @@ -160,10 +162,10 @@ pub async fn store_table_schema( .fetch_one(&mut *tx) .await?; - let next_version = current_version.unwrap_or(-1) + 1; - let schema_version: SchemaVersion = next_version - .try_into() - .expect("schema version should not overflow u64"); + // We case to `u64` without checks. This is fine since we control the database, but we might want + // to be more defensive in the future. + let next_schema_version: u64 = + current_schema_version.map_or(STARTING_SCHEMA_VERSION, |v| v as u64) + 1; let table_schema_id: i64 = sqlx::query( r#" @@ -176,7 +178,7 @@ pub async fn store_table_schema( .bind(table_schema.id.into_inner() as i64) .bind(&table_schema.name.schema) .bind(&table_schema.name.name) - .bind(next_version) + .bind(next_schema_version as i64) .fetch_one(&mut *tx) .await? .get(0); @@ -204,7 +206,7 @@ pub async fn store_table_schema( tx.commit().await?; - Ok(table_schema.into_table_schema(schema_version)) + Ok(table_schema.into_table_schema(next_schema_version)) } /// Loads all table schemas for a pipeline from the database. diff --git a/etl-replicator/migrations/20251001000100_add_schema_versions.sql b/etl-replicator/migrations/20251001000100_add_schema_versions.sql index 09ded33e1..e17f45ceb 100644 --- a/etl-replicator/migrations/20251001000100_add_schema_versions.sql +++ b/etl-replicator/migrations/20251001000100_add_schema_versions.sql @@ -4,10 +4,10 @@ alter table etl.table_schemas -- Adjust unique constraint to account for schema versions alter table etl.table_schemas - drop constraint if exists etl_table_schemas_pipeline_id_table_id_key; + drop constraint if exists table_schemas_pipeline_id_table_id_key; alter table etl.table_schemas - add constraint etl_table_schemas_pipeline_id_table_id_version_key + add constraint table_schemas_pipeline_id_table_id_schema_version_key unique (pipeline_id, table_id, schema_version); -- Refresh supporting indexes diff --git a/etl/src/store/both/postgres.rs b/etl/src/store/both/postgres.rs index d88a2d387..b23d410ae 100644 --- a/etl/src/store/both/postgres.rs +++ b/etl/src/store/both/postgres.rs @@ -601,6 +601,7 @@ impl SchemaStore for PostgresStore { // We also lock the entire section to be consistent. let mut inner = self.inner.lock().await; + let stored_schema = schema::store_table_schema(&pool, self.pipeline_id as i64, table_schema) .await @@ -611,14 +612,15 @@ impl SchemaStore for PostgresStore { format!("Failed to store table schema in postgres: {err}") ) })?; - let schema_arc = Arc::new(stored_schema); + + let stored_schema = Arc::new(stored_schema); inner .table_schemas - .entry(schema_arc.id) + .entry(stored_schema.id) .or_insert_with(BTreeMap::new) - .insert(schema_arc.version, Arc::clone(&schema_arc)); + .insert(stored_schema.version, stored_schema.clone()); - Ok(schema_arc) + Ok(stored_schema) } } diff --git a/etl/tests/postgres_store.rs b/etl/tests/postgres_store.rs index ddc81e913..2d8755a65 100644 --- a/etl/tests/postgres_store.rs +++ b/etl/tests/postgres_store.rs @@ -254,9 +254,9 @@ async fn test_schema_store_operations() { assert_eq!(all_schemas.len(), 1); // Test storing another schema - let table_schema2 = create_another_table_schema(); + let table_schema_2 = create_another_table_schema(); store - .store_table_schema(table_schema2.clone()) + .store_table_schema(table_schema_2.clone()) .await .unwrap(); @@ -272,16 +272,16 @@ async fn test_schema_store_load_schemas() { let pipeline_id = 1; let store = PostgresStore::new(pipeline_id, database.config.clone()); - let table_schema1 = create_sample_table_schema(); - let table_schema2 = create_another_table_schema(); + let table_schema_1 = create_sample_table_schema(); + let table_schema_2 = create_another_table_schema(); // Store schemas let stored_schema_1 = store - .store_table_schema(table_schema1.clone()) + .store_table_schema(table_schema_1.clone()) .await .unwrap(); let stored_schema_2 = store - .store_table_schema(table_schema2.clone()) + .store_table_schema(table_schema_2.clone()) .await .unwrap(); @@ -301,22 +301,22 @@ async fn test_schema_store_load_schemas() { assert_eq!(schemas.len(), 2); let schema1 = new_store - .get_table_schema(&table_schema1.id, stored_schema_1.version) + .get_table_schema(&table_schema_1.id, stored_schema_1.version) .await .unwrap(); assert!(schema1.is_some()); let schema1 = schema1.unwrap(); - assert_eq!(schema1.id, table_schema1.id); - assert_eq!(schema1.name, table_schema1.name); + assert_eq!(schema1.id, table_schema_1.id); + assert_eq!(schema1.name, table_schema_1.name); let schema2 = new_store - .get_table_schema(&table_schema2.id, stored_schema_2.version) + .get_table_schema(&table_schema_2.id, stored_schema_2.version) .await .unwrap(); assert!(schema2.is_some()); let schema2 = schema2.unwrap(); - assert_eq!(schema2.id, table_schema2.id); - assert_eq!(schema2.name, table_schema2.name); + assert_eq!(schema2.id, table_schema_2.id); + assert_eq!(schema2.name, table_schema_2.name); } #[tokio::test(flavor = "multi_thread")] @@ -335,6 +335,15 @@ async fn test_schema_store_update_existing() { .await .unwrap(); + // Verify schema + let schema = store + .get_table_schema(&table_schema.id, stored_schema.version) + .await + .unwrap(); + assert!(schema.is_some()); + let schema = schema.unwrap(); + assert_eq!(schema.column_schemas.len(), 3); + // Update schema by adding a column table_schema.add_column_schema(ColumnSchema::new( "updated_at".to_string(), @@ -344,14 +353,14 @@ async fn test_schema_store_update_existing() { )); // Store updated schema - store + let stored_schema_updated = store .store_table_schema(table_schema.clone()) .await .unwrap(); // Verify updated schema let schema = store - .get_table_schema(&table_schema.id, stored_schema.version) + .get_table_schema(&table_schema.id, stored_schema_updated.version) .await .unwrap(); assert!(schema.is_some()); @@ -371,12 +380,12 @@ async fn test_multiple_pipelines_isolation() { init_test_tracing(); let database = spawn_source_database_for_store().await; - let pipeline_id1 = 1; - let pipeline_id2 = 2; + let pipeline_id_1 = 1; + let pipeline_id_2 = 2; let table_id = TableId::new(12345); - let store_1 = PostgresStore::new(pipeline_id1, database.config.clone()); - let store_2 = PostgresStore::new(pipeline_id2, database.config.clone()); + let store_1 = PostgresStore::new(pipeline_id_1, database.config.clone()); + let store_2 = PostgresStore::new(pipeline_id_2, database.config.clone()); // Add state to pipeline 1 let init_phase = TableReplicationPhase::Init; @@ -400,15 +409,15 @@ async fn test_multiple_pipelines_isolation() { assert_eq!(state2, Some(data_sync_phase)); // Test schema isolation - let table_schema1 = create_sample_table_schema(); - let table_schema2 = create_another_table_schema(); + let table_schema_1 = create_sample_table_schema(); + let table_schema_2 = create_another_table_schema(); let stored_schema_1 = store_1 - .store_table_schema(table_schema1.clone()) + .store_table_schema(table_schema_1.clone()) .await .unwrap(); let stored_schema_2 = store_2 - .store_table_schema(table_schema2.clone()) + .store_table_schema(table_schema_2.clone()) .await .unwrap(); @@ -416,14 +425,14 @@ async fn test_multiple_pipelines_isolation() { let schemas_1 = store_1.get_all_table_schemas().await; assert_eq!(schemas_1.len(), 1); assert_eq!( - schemas_1[&table_schema1.id][&stored_schema_1.version], + schemas_1[&table_schema_1.id][&stored_schema_1.version], stored_schema_1 ); let schemas_2 = store_2.get_all_table_schemas().await; assert_eq!(schemas_2.len(), 1); assert_eq!( - schemas_2[&table_schema1.id][&stored_schema_2.version], + schemas_2[&table_schema_2.id][&stored_schema_2.version], stored_schema_2 ); } @@ -643,11 +652,11 @@ async fn test_table_mappings_pipeline_isolation() { init_test_tracing(); let database = spawn_source_database_for_store().await; - let pipeline_id1 = 1; - let pipeline_id2 = 2; + let pipeline_id_1 = 1; + let pipeline_id_2 = 2; - let store_1 = PostgresStore::new(pipeline_id1, database.config.clone()); - let store_2 = PostgresStore::new(pipeline_id2, database.config.clone()); + let store_1 = PostgresStore::new(pipeline_id_1, database.config.clone()); + let store_2 = PostgresStore::new(pipeline_id_2, database.config.clone()); let table_id = TableId::new(12345); @@ -670,7 +679,7 @@ async fn test_table_mappings_pipeline_isolation() { assert_eq!(mapping2, Some("pipeline2_table".to_string())); // Verify isolation persists after loading from database - let new_store_1 = PostgresStore::new(pipeline_id1, database.config.clone()); + let new_store_1 = PostgresStore::new(pipeline_id_1, database.config.clone()); new_store_1.load_table_mappings().await.unwrap(); let loaded_mapping1 = new_store_1.get_table_mapping(&table_id).await.unwrap(); From 6912b1c356c838aca7fd072571dd10e092527866 Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Wed, 1 Oct 2025 13:49:03 -0700 Subject: [PATCH 20/45] Improve --- etl/src/store/both/memory.rs | 16 ++++++++-------- etl/tests/pipeline.rs | 33 +++++++++++++++------------------ 2 files changed, 23 insertions(+), 26 deletions(-) diff --git a/etl/src/store/both/memory.rs b/etl/src/store/both/memory.rs index 16f83970c..92eda8d64 100644 --- a/etl/src/store/both/memory.rs +++ b/etl/src/store/both/memory.rs @@ -177,7 +177,7 @@ impl SchemaStore for MemoryStore { Ok(inner .table_schemas .get(table_id) - .and_then(|schemas| schemas.get(&version).cloned())) + .and_then(|table_schemas| table_schemas.get(&version).cloned())) } async fn get_latest_table_schema( @@ -189,7 +189,7 @@ impl SchemaStore for MemoryStore { Ok(inner .table_schemas .get(table_id) - .and_then(|schemas| schemas.iter().next_back().map(|(_, schema)| schema.clone()))) + .and_then(|table_schemas| table_schemas.iter().next_back().map(|(_, schema)| schema.clone()))) } async fn load_table_schemas(&self) -> EtlResult { @@ -198,7 +198,7 @@ impl SchemaStore for MemoryStore { Ok(inner .table_schemas .values() - .map(|schemas| schemas.len()) + .map(|table_schemas| table_schemas.len()) .sum()) } @@ -207,21 +207,21 @@ impl SchemaStore for MemoryStore { table_schema: TableSchemaDraft, ) -> EtlResult> { let mut inner = self.inner.lock().await; - let schemas = inner + let table_schemas = inner .table_schemas .entry(table_schema.id) .or_insert_with(BTreeMap::new); - let next_version = schemas + let next_version = table_schemas .keys() .next_back() .map(|version| version + 1) .unwrap_or(0); - let schema = Arc::new(table_schema.into_table_schema(next_version)); - schemas.insert(next_version, Arc::clone(&schema)); + let table_schema = Arc::new(table_schema.into_table_schema(next_version)); + table_schemas.insert(next_version, table_schema.clone()); - Ok(schema) + Ok(table_schema) } } diff --git a/etl/tests/pipeline.rs b/etl/tests/pipeline.rs index c21242f06..791263122 100644 --- a/etl/tests/pipeline.rs +++ b/etl/tests/pipeline.rs @@ -527,12 +527,11 @@ async fn table_schema_changes_are_handled_correctly() { // Check the initial schema. let table_schemas = store.get_latest_table_schemas().await; assert_eq!(table_schemas.len(), 1); - let users_table_schema = column_schema_names( - table_schemas - .get(&database_schema.users_schema().id) - .unwrap(), - ); - assert_eq!(users_table_schema, vec!["id", "name", "age"]); + let users_table_schema = table_schemas + .get(&database_schema.users_schema().id) + .unwrap(); + assert_eq!(users_table_schema.version, 0); + assert_eq!(column_schema_names(users_table_schema), vec!["id", "name", "age"]); // Check the initial data. let table_rows = destination.get_table_rows().await; @@ -577,12 +576,11 @@ async fn table_schema_changes_are_handled_correctly() { // Check the updated schema. let table_schemas = store.get_latest_table_schemas().await; assert_eq!(table_schemas.len(), 1); - let users_table_schema = column_schema_names( - table_schemas - .get(&database_schema.users_schema().id) - .unwrap(), - ); - assert_eq!(users_table_schema, vec!["id", "name", "new_age", "year"]); + let users_table_schema = table_schemas + .get(&database_schema.users_schema().id) + .unwrap(); + assert_eq!(users_table_schema.version, 1); + assert_eq!(column_schema_names(users_table_schema), vec!["id", "name", "new_age", "year"]); // Check the updated data. let events = destination.get_events().await; @@ -627,12 +625,11 @@ async fn table_schema_changes_are_handled_correctly() { // Check the updated schema. let table_schemas = store.get_latest_table_schemas().await; assert_eq!(table_schemas.len(), 1); - let users_table_schema = column_schema_names( - table_schemas - .get(&database_schema.users_schema().id) - .unwrap(), - ); - assert_eq!(users_table_schema, vec!["id", "name", "new_age"]); + let users_table_schema = table_schemas + .get(&database_schema.users_schema().id) + .unwrap(); + assert_eq!(users_table_schema.version, 2); + assert_eq!(column_schema_names(users_table_schema), vec!["id", "name", "new_age"]); // Check the updated data. let events = destination.get_events().await; From 76e4a5c4db9a9fd21c930f8549c18f8f694f46ec Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Wed, 1 Oct 2025 14:31:09 -0700 Subject: [PATCH 21/45] Improve --- etl-destinations/src/bigquery/client.rs | 62 +++++++++++++++++-------- etl-postgres/src/replication/schema.rs | 8 +++- etl/Cargo.toml | 1 - etl/src/store/both/memory.rs | 15 ++++-- etl/tests/pipeline.rs | 15 ++++-- 5 files changed, 70 insertions(+), 31 deletions(-) diff --git a/etl-destinations/src/bigquery/client.rs b/etl-destinations/src/bigquery/client.rs index f9a9bfcf7..edc4ce281 100644 --- a/etl-destinations/src/bigquery/client.rs +++ b/etl-destinations/src/bigquery/client.rs @@ -67,6 +67,22 @@ pub struct BigQueryClient { } impl BigQueryClient { + /// Quotes a single BigQuery identifier, escaping embedded backticks. + fn quote_identifier(identifier: &str) -> String { + format!("`{}`", identifier.replace('`', "``")) + } + + /// Quotes a dotted BigQuery path (e.g. project.dataset.table), escaping each segment. + fn quote_qualified_identifiers(parts: &[&str]) -> String { + let escaped = parts + .iter() + .map(|part| part.replace('`', "``")) + .collect::>() + .join("."); + + format!("`{escaped}`") + } + /// Creates a new [`BigQueryClient`] from a service account key file. /// /// Authenticates with BigQuery using the service account key at the specified file path. @@ -106,7 +122,11 @@ impl BigQueryClient { dataset_id: &BigQueryDatasetId, table_id: &BigQueryTableId, ) -> String { - format!("`{}.{}.{}`", self.project_id, dataset_id, table_id) + Self::quote_qualified_identifiers(&[ + &self.project_id, + dataset_id.as_str(), + table_id.as_str(), + ]) } /// Creates a table in BigQuery if it doesn't already exist, otherwise efficiently truncates @@ -284,7 +304,7 @@ impl BigQueryClient { column_name: &str, ) -> EtlResult<()> { let full_table_name = self.full_table_name(dataset_id, table_id); - let column_identifier = column_name; + let column_identifier = Self::quote_identifier(column_name); let query = format!("alter table {full_table_name} drop column if exists {column_identifier}"); @@ -302,8 +322,8 @@ impl BigQueryClient { new_name: &str, ) -> EtlResult<()> { let full_table_name = self.full_table_name(dataset_id, table_id); - let old_identifier = old_name; - let new_identifier = new_name; + let old_identifier = Self::quote_identifier(old_name); + let new_identifier = Self::quote_identifier(new_name); let query = format!( "alter table {full_table_name} rename column {old_identifier} to {new_identifier}" ); @@ -321,7 +341,7 @@ impl BigQueryClient { column_schema: &ColumnSchema, ) -> EtlResult<()> { let full_table_name = self.full_table_name(dataset_id, table_id); - let column_identifier = &column_schema.name; + let column_identifier = Self::quote_identifier(&column_schema.name); let column_type = Self::postgres_to_bigquery_type(&column_schema.typ); let query = format!( "alter table {full_table_name} alter column {column_identifier} set data type {column_type}" @@ -341,7 +361,7 @@ impl BigQueryClient { nullable: bool, ) -> EtlResult<()> { let full_table_name = self.full_table_name(dataset_id, table_id); - let column_identifier = column_name; + let column_identifier = Self::quote_identifier(column_name); let clause = if nullable { "drop not null" } else { @@ -382,7 +402,7 @@ impl BigQueryClient { let columns = primary_columns .iter() - .map(|column| format!("`{}`", column.name)) + .map(|column| Self::quote_identifier(&column.name)) .collect::>() .join(","); @@ -400,10 +420,12 @@ impl BigQueryClient { dataset_id: &BigQueryDatasetId, table_id: &BigQueryTableId, ) -> EtlResult { - let info_schema_table = format!( - "`{}.{}`.INFORMATION_SCHEMA.TABLE_CONSTRAINTS`", - &self.project_id, dataset_id - ); + let info_schema_table = Self::quote_qualified_identifiers(&[ + &self.project_id, + dataset_id.as_str(), + "INFORMATION_SCHEMA", + "TABLE_CONSTRAINTS", + ]); let table_literal = table_id; let query = format!( "select constraint_name from {info_schema_table} where table_name = '{table_literal}' and constraint_type = 'PRIMARY KEY'", @@ -570,7 +592,7 @@ impl BigQueryClient { fn column_spec(column_schema: &ColumnSchema) -> String { let mut column_spec = format!( "{} {}", - column_schema.name, + Self::quote_identifier(&column_schema.name), Self::postgres_to_bigquery_type(&column_schema.typ) ); @@ -588,7 +610,7 @@ impl BigQueryClient { let identity_columns: Vec = column_schemas .iter() .filter(|s| s.primary) - .map(|c| c.name.clone()) + .map(|c| Self::quote_identifier(&c.name)) .collect(); if identity_columns.is_empty() { @@ -947,7 +969,7 @@ mod tests { let not_null_column = ColumnSchema::new("id".to_string(), Type::INT4, -1, true); let not_null_spec = BigQueryClient::column_spec(¬_null_column); - assert_eq!(not_null_spec, "`id` int64 not null"); + assert_eq!(not_null_spec, "`id` int64"); let array_column = ColumnSchema::new("tags".to_string(), Type::TEXT_ARRAY, -1, false); let array_spec = BigQueryClient::column_spec(&array_column); @@ -993,7 +1015,7 @@ mod tests { let spec = BigQueryClient::create_columns_spec(&columns); assert_eq!( spec, - "(`id` int64 not null,`name` string,`active` bool not null, primary key (`id`) not enforced)" + "(`id` int64,`name` string,`active` bool, primary key (`id`) not enforced)" ); } @@ -1024,7 +1046,7 @@ mod tests { )); assert!(matches!( descriptor.field_descriptors[0].mode, - ColumnMode::Required + ColumnMode::Nullable )); assert_eq!(descriptor.field_descriptors[1].name, "name"); @@ -1044,7 +1066,7 @@ mod tests { )); assert!(matches!( descriptor.field_descriptors[2].mode, - ColumnMode::Required + ColumnMode::Nullable )); // Check array column @@ -1155,7 +1177,7 @@ mod tests { let columns_spec = BigQueryClient::create_columns_spec(&columns); let query = format!("create or replace table {full_table_name} {columns_spec}"); - let expected_query = "create or replace table `test-project.test_dataset.test_table` (`id` int64 not null,`name` string, primary key (`id`) not enforced)"; + let expected_query = "create or replace table `test-project.test_dataset.test_table` (`id` int64,`name` string, primary key (`id`) not enforced)"; assert_eq!(query, expected_query); } @@ -1166,7 +1188,7 @@ mod tests { let table_id = "test_table"; let max_staleness_mins = 15; - let columns = vec![ColumnSchema::new("id".to_string(), Type::INT4, -1, false)]; + let columns = vec![ColumnSchema::new("id".to_string(), Type::INT4, -1, true)]; // Simulate the query generation logic with staleness let full_table_name = format!("`{project_id}.{dataset_id}.{table_id}`"); @@ -1176,7 +1198,7 @@ mod tests { "create or replace table {full_table_name} {columns_spec} {max_staleness_option}" ); - let expected_query = "create or replace table `test-project.test_dataset.test_table` (`id` int64 not null, primary key (`id`) not enforced) options (max_staleness = interval 15 minute)"; + let expected_query = "create or replace table `test-project.test_dataset.test_table` (`id` int64, primary key (`id`) not enforced) options (max_staleness = interval 15 minute)"; assert_eq!(query, expected_query); } } diff --git a/etl-postgres/src/replication/schema.rs b/etl-postgres/src/replication/schema.rs index 02556eb54..cb56dc748 100644 --- a/etl-postgres/src/replication/schema.rs +++ b/etl-postgres/src/replication/schema.rs @@ -150,6 +150,8 @@ pub async fn store_table_schema( ) -> Result { let mut tx = pool.begin().await?; + // We are fine with running this query for every table schema store since we assume that + // it's not a frequent operation. let current_schema_version: Option = sqlx::query_scalar( r#" select max(schema_version) @@ -164,8 +166,10 @@ pub async fn store_table_schema( // We case to `u64` without checks. This is fine since we control the database, but we might want // to be more defensive in the future. - let next_schema_version: u64 = - current_schema_version.map_or(STARTING_SCHEMA_VERSION, |v| v as u64) + 1; + let next_schema_version: u64 = match current_schema_version { + Some(current_schema_version) => current_schema_version as u64 + 1, + None => STARTING_SCHEMA_VERSION, + }; let table_schema_id: i64 = sqlx::query( r#" diff --git a/etl/Cargo.toml b/etl/Cargo.toml index 8f758e07a..b6b94eff2 100644 --- a/etl/Cargo.toml +++ b/etl/Cargo.toml @@ -43,7 +43,6 @@ tokio-rustls = { workspace = true, default-features = false } tracing = { workspace = true, default-features = true } uuid = { workspace = true, features = ["v4"] } x509-cert = { workspace = true, default-features = false } -log = "0.4.28" [dev-dependencies] etl-postgres = { workspace = true, features = [ diff --git a/etl/src/store/both/memory.rs b/etl/src/store/both/memory.rs index 92eda8d64..cb8890537 100644 --- a/etl/src/store/both/memory.rs +++ b/etl/src/store/both/memory.rs @@ -10,6 +10,9 @@ use crate::store::cleanup::CleanupStore; use crate::store::schema::SchemaStore; use crate::store::state::StateStore; +/// The initial schema version number. +const STARTING_SCHEMA_VERSION: u64 = 0; + /// Inner state of [`MemoryStore`] #[derive(Debug)] struct Inner { @@ -186,10 +189,12 @@ impl SchemaStore for MemoryStore { ) -> EtlResult>> { let inner = self.inner.lock().await; - Ok(inner - .table_schemas - .get(table_id) - .and_then(|table_schemas| table_schemas.iter().next_back().map(|(_, schema)| schema.clone()))) + Ok(inner.table_schemas.get(table_id).and_then(|table_schemas| { + table_schemas + .iter() + .next_back() + .map(|(_, schema)| schema.clone()) + })) } async fn load_table_schemas(&self) -> EtlResult { @@ -216,7 +221,7 @@ impl SchemaStore for MemoryStore { .keys() .next_back() .map(|version| version + 1) - .unwrap_or(0); + .unwrap_or(STARTING_SCHEMA_VERSION); let table_schema = Arc::new(table_schema.into_table_schema(next_version)); table_schemas.insert(next_version, table_schema.clone()); diff --git a/etl/tests/pipeline.rs b/etl/tests/pipeline.rs index 791263122..db61edc50 100644 --- a/etl/tests/pipeline.rs +++ b/etl/tests/pipeline.rs @@ -531,7 +531,10 @@ async fn table_schema_changes_are_handled_correctly() { .get(&database_schema.users_schema().id) .unwrap(); assert_eq!(users_table_schema.version, 0); - assert_eq!(column_schema_names(users_table_schema), vec!["id", "name", "age"]); + assert_eq!( + column_schema_names(users_table_schema), + vec!["id", "name", "age"] + ); // Check the initial data. let table_rows = destination.get_table_rows().await; @@ -580,7 +583,10 @@ async fn table_schema_changes_are_handled_correctly() { .get(&database_schema.users_schema().id) .unwrap(); assert_eq!(users_table_schema.version, 1); - assert_eq!(column_schema_names(users_table_schema), vec!["id", "name", "new_age", "year"]); + assert_eq!( + column_schema_names(users_table_schema), + vec!["id", "name", "new_age", "year"] + ); // Check the updated data. let events = destination.get_events().await; @@ -629,7 +635,10 @@ async fn table_schema_changes_are_handled_correctly() { .get(&database_schema.users_schema().id) .unwrap(); assert_eq!(users_table_schema.version, 2); - assert_eq!(column_schema_names(users_table_schema), vec!["id", "name", "new_age"]); + assert_eq!( + column_schema_names(users_table_schema), + vec!["id", "name", "new_age"] + ); // Check the updated data. let events = destination.get_events().await; From 431c10f324dee1caa53f39ee947801a711dc8b35 Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Wed, 1 Oct 2025 15:01:42 -0700 Subject: [PATCH 22/45] Improve --- etl-destinations/src/bigquery/client.rs | 57 +----- etl-destinations/src/bigquery/core.rs | 231 ++++++++++-------------- etl/src/replication/client.rs | 1 + 3 files changed, 98 insertions(+), 191 deletions(-) diff --git a/etl-destinations/src/bigquery/client.rs b/etl-destinations/src/bigquery/client.rs index edc4ce281..9bc96e30d 100644 --- a/etl-destinations/src/bigquery/client.rs +++ b/etl-destinations/src/bigquery/client.rs @@ -313,27 +313,7 @@ impl BigQueryClient { Ok(()) } - /// Renames a column in an existing BigQuery table. - pub async fn rename_column( - &self, - dataset_id: &BigQueryDatasetId, - table_id: &BigQueryTableId, - old_name: &str, - new_name: &str, - ) -> EtlResult<()> { - let full_table_name = self.full_table_name(dataset_id, table_id); - let old_identifier = Self::quote_identifier(old_name); - let new_identifier = Self::quote_identifier(new_name); - let query = format!( - "alter table {full_table_name} rename column {old_identifier} to {new_identifier}" - ); - - let _ = self.query(QueryRequest::new(query)).await?; - - Ok(()) - } - - /// Alters the data type of an existing column in a BigQuery table. + /// Alters the data type of the existing column in a BigQuery table. pub async fn alter_column_type( &self, dataset_id: &BigQueryDatasetId, @@ -352,29 +332,6 @@ impl BigQueryClient { Ok(()) } - /// Updates the nullability of an existing column in a BigQuery table. - pub async fn alter_column_nullability( - &self, - dataset_id: &BigQueryDatasetId, - table_id: &BigQueryTableId, - column_name: &str, - nullable: bool, - ) -> EtlResult<()> { - let full_table_name = self.full_table_name(dataset_id, table_id); - let column_identifier = Self::quote_identifier(column_name); - let clause = if nullable { - "drop not null" - } else { - "set not null" - }; - let query = - format!("alter table {full_table_name} alter column {column_identifier} {clause}"); - - let _ = self.query(QueryRequest::new(query)).await?; - - Ok(()) - } - /// Synchronizes the primary key definition for a BigQuery table with the provided schema. pub async fn sync_primary_key( &self, @@ -388,18 +345,14 @@ impl BigQueryClient { .collect(); let has_primary_key = self.has_primary_key(dataset_id, table_id).await?; + if has_primary_key { + self.drop_primary_key(dataset_id, table_id).await?; + } if primary_columns.is_empty() { - if has_primary_key { - self.drop_primary_key(dataset_id, table_id).await?; - } return Ok(()); } - if has_primary_key { - self.drop_primary_key(dataset_id, table_id).await?; - } - let columns = primary_columns .iter() .map(|column| Self::quote_identifier(&column.name)) @@ -688,7 +641,7 @@ impl BigQueryClient { /// Converts Postgres column schemas to a BigQuery [`TableDescriptor`]. /// /// Maps data types and nullability to BigQuery column specifications, setting - /// appropriate column modes and automatically adding CDC special columns. + /// appropriate column modes, and automatically adding CDC special columns. pub fn column_schemas_to_table_descriptor( column_schemas: &[ColumnSchema], use_cdc_sequence_column: bool, diff --git a/etl-destinations/src/bigquery/core.rs b/etl-destinations/src/bigquery/core.rs index 80a33cfa7..78acdf0a9 100644 --- a/etl-destinations/src/bigquery/core.rs +++ b/etl-destinations/src/bigquery/core.rs @@ -3,7 +3,8 @@ use etl::error::{ErrorKind, EtlError, EtlResult}; use etl::store::schema::SchemaStore; use etl::store::state::StateStore; use etl::types::{ - Cell, Event, PgLsn, RelationEvent, SchemaVersion, TableId, TableName, TableRow, TableSchema, + Cell, Event, PgLsn, RelationChange, RelationEvent, SchemaVersion, TableId, TableName, TableRow, + TableSchema, }; use etl::{bail, etl_error}; use gcp_bigquery_client::storage::TableDescriptor; @@ -538,143 +539,95 @@ where Ok(()) } - async fn apply_relation_event(&self, _relation_event: RelationEvent) -> EtlResult<()> { - // TODO: implement the relation event. - // if relation_event.changes.is_empty() { - // debug!( - // table_id = %relation_event.table_id, - // "relation event contained no schema changes; skipping" - // ); - // - // return Ok(()); - // } - // - // let (sequenced_bigquery_table_id, table_schema) = self - // .prepare_table(&relation_event.table_id, relation_event.new_schema_version) - // .await?; - // - // let sequenced_table_name = sequenced_bigquery_table_id.to_string(); - // let mut primary_key_dirty = false; - // for change in relation_event.changes { - // match change { - // RelationChange::AddColumn(column_schema) => { - // let column_name = column_schema.name.clone(); - // let is_primary = column_schema.primary; - // - // self.client - // .add_column(&self.dataset_id, &sequenced_table_name, &column_schema) - // .await?; - // - // debug!( - // table = %sequenced_table_name, - // column = %column_name, - // "added column in BigQuery" - // ); - // - // if is_primary { - // primary_key_dirty = true; - // } - // } - // RelationChange::DropColumn(column_schema) => { - // let column_name = column_schema.name.clone(); - // let was_primary = column_schema.primary; - // - // self.client - // .drop_column(&self.dataset_id, &sequenced_table_name, &column_schema.name) - // .await?; - // - // debug!( - // table = %sequenced_table_name, - // column = %column_name, - // "dropped column in BigQuery" - // ); - // - // if was_primary { - // primary_key_dirty = true; - // } - // } - // RelationChange::AlterColumn(previous, latest) => { - // let old_name = previous.name.clone(); - // let new_name = latest.name.clone(); - // let renamed = old_name != new_name; - // - // if renamed { - // self.client - // .rename_column( - // &self.dataset_id, - // &sequenced_table_name, - // &previous.name, - // &latest.name, - // ) - // .await?; - // - // debug!( - // table = %sequenced_table_name, - // old_column = %old_name, - // new_column = %new_name, - // "renamed column in BigQuery" - // ); - // } - // - // if previous.typ != latest.typ { - // self.client - // .alter_column_type(&self.dataset_id, &sequenced_table_name, &latest) - // .await?; - // - // debug!( - // table = %sequenced_table_name, - // column = %new_name, - // "updated column type in BigQuery" - // ); - // } - // - // if previous.nullable != latest.nullable { - // self.client - // .alter_column_nullability( - // &self.dataset_id, - // &sequenced_table_name, - // &latest.name, - // latest.nullable, - // ) - // .await?; - // - // debug!( - // table = %sequenced_table_name, - // column = %new_name, - // nullable = latest.nullable, - // "updated column nullability in BigQuery" - // ); - // } - // - // if previous.primary != latest.primary - // || (renamed && (previous.primary || latest.primary)) - // { - // primary_key_dirty = true; - // } - // } - // } - // } - // - // if primary_key_dirty { - // self.client - // .sync_primary_key( - // &self.dataset_id, - // &sequenced_table_name, - // &table_schema.column_schemas, - // ) - // .await?; - // - // debug!( - // table = %sequenced_table_name, - // "synchronized primary key definition in BigQuery" - // ); - // } - // - // info!( - // table_id = %relation_event.table_id, - // table = %sequenced_table_name, - // "applied relation changes in BigQuery" - // ); + async fn apply_relation_event_changes(&self, relation_event: RelationEvent) -> EtlResult<()> { + // We build the list of changes for this relation event, this way, we can express the event + // in terms of a minimal set of operations to apply on the destination table schema. + let changes = relation_event.build_changes(); + if changes.is_empty() { + debug!( + table_id = %relation_event.table_id, + "relation event contained no schema changes; skipping" + ); + + return Ok(()); + } + + // We prepare the table for the changes. We don't make any assumptions on the table state, we + // just want to work on an existing table and apply changes to it. + let (sequenced_bigquery_table_id, _table_schema) = self + .prepare_table( + &relation_event.table_id, + Some(relation_event.new_table_schema.version), + ) + .await?; + + let sequenced_table_name = sequenced_bigquery_table_id.to_string(); + for change in changes { + match change { + RelationChange::AddColumn(column_schema) => { + let column_name = column_schema.name.clone(); + + self.client + .add_column(&self.dataset_id, &sequenced_table_name, &column_schema) + .await?; + + debug!( + table = %sequenced_table_name, + column = %column_name, + "added column in BigQuery" + ); + } + RelationChange::DropColumn(column_schema) => { + let column_name = column_schema.name.clone(); + + self.client + .drop_column(&self.dataset_id, &sequenced_table_name, &column_schema.name) + .await?; + + debug!( + table = %sequenced_table_name, + column = %column_name, + "dropped column in BigQuery" + ); + } + RelationChange::AlterColumn(previous_column_schema, latest_column_schema) => { + if previous_column_schema.typ != latest_column_schema.typ { + self.client + .alter_column_type( + &self.dataset_id, + &sequenced_table_name, + &latest_column_schema, + ) + .await?; + + debug!( + table = %sequenced_table_name, + column = %latest_column_schema.name, + "updated column type in BigQuery" + ); + } + } + } + } + + self.client + .sync_primary_key( + &self.dataset_id, + &sequenced_table_name, + &relation_event.new_table_schema.column_schemas, + ) + .await?; + + debug!( + table = %sequenced_table_name, + "synchronized primary key definition in BigQuery" + ); + + info!( + table_id = %relation_event.table_id, + table = %sequenced_table_name, + "applied relation changes in BigQuery" + ); Ok(()) } @@ -814,7 +767,7 @@ where batch_schema_version = Some(relation.new_table_schema.version); // Apply relation change, then prime the next batch with the new schema version. - self.apply_relation_event(relation).await?; + self.apply_relation_event_changes(relation).await?; } Event::Truncate(truncate) => { // Finish current batch before a TRUNCATE (it affects table state). diff --git a/etl/src/replication/client.rs b/etl/src/replication/client.rs index 8d42736ea..b339b72e8 100644 --- a/etl/src/replication/client.rs +++ b/etl/src/replication/client.rs @@ -723,6 +723,7 @@ impl PgReplicationClient { Ok(column_schemas) } + /// Creates a COPY stream for reading data from a table using its OID. /// /// The stream will include only the specified columns and use text format. From f21a74ffd0f851419c7d649b25809a40f1808fe3 Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Wed, 1 Oct 2025 15:08:34 -0700 Subject: [PATCH 23/45] Improve --- etl-api/tests/pipelines.rs | 4 ++-- etl-destinations/src/bigquery/core.rs | 1 + 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/etl-api/tests/pipelines.rs b/etl-api/tests/pipelines.rs index e12712dff..34b4d3f6c 100644 --- a/etl-api/tests/pipelines.rs +++ b/etl-api/tests/pipelines.rs @@ -1209,11 +1209,11 @@ async fn deleting_pipeline_removes_table_schemas_from_source_database() { // Insert table schemas using production schema let table_schema_id_1 = sqlx::query_scalar::<_, i64>( - "INSERT INTO etl.table_schemas (pipeline_id, table_id, schema_name, table_name) VALUES ($1, $2, 'public', 'test_users') RETURNING id" + "INSERT INTO etl.table_schemas (pipeline_id, table_id, schema_name, table_name, schema_version) VALUES ($1, $2, 'public', 'test_users', 0) RETURNING id" ).bind(pipeline_id).bind(table1_oid).fetch_one(&source_db_pool).await.unwrap(); let table_schema_id_2 = sqlx::query_scalar::<_, i64>( - "INSERT INTO etl.table_schemas (pipeline_id, table_id, schema_name, table_name) VALUES ($1, $2, 'public', 'test_orders') RETURNING id" + "INSERT INTO etl.table_schemas (pipeline_id, table_id, schema_name, table_name, schema_version) VALUES ($1, $2, 'public', 'test_orders', 0) RETURNING id" ).bind(pipeline_id).bind(table2_oid).fetch_one(&source_db_pool).await.unwrap(); // Insert multiple columns for each table to test CASCADE behavior diff --git a/etl-destinations/src/bigquery/core.rs b/etl-destinations/src/bigquery/core.rs index 78acdf0a9..4ec559322 100644 --- a/etl-destinations/src/bigquery/core.rs +++ b/etl-destinations/src/bigquery/core.rs @@ -632,6 +632,7 @@ where Ok(()) } + #[allow(clippy::too_many_arguments)] #[inline] fn push_dml_statement( table_id_to_table_rows: &mut HashMap>, From 9bf2a347fdae8576f615192a251f8e6a227576b8 Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Wed, 1 Oct 2025 15:20:07 -0700 Subject: [PATCH 24/45] Improve --- etl/src/conversions/event.rs | 25 +++++++------------------ etl/src/replication/apply.rs | 29 +++++++++++++++++++++-------- 2 files changed, 28 insertions(+), 26 deletions(-) diff --git a/etl/src/conversions/event.rs b/etl/src/conversions/event.rs index 9a7acc2cc..00c8c45e8 100644 --- a/etl/src/conversions/event.rs +++ b/etl/src/conversions/event.rs @@ -1,6 +1,6 @@ use core::str; use etl_postgres::types::{ - ColumnSchema, SchemaVersion, TableId, TableSchema, TableSchemaDraft, convert_type_oid_to_type, + ColumnSchema, SchemaVersion, TableId, TableSchema, convert_type_oid_to_type, }; use postgres_replication::protocol; use std::sync::Arc; @@ -48,18 +48,13 @@ pub fn parse_event_from_commit_message( } } -/// Creates a [`RelationEventDraft`] from Postgres protocol data. +/// Creates a [`Vec`] from Postgres protocol data. /// -/// This method parses the replication protocol relation message and builds -/// a complete table schema for use in interpreting later data events. -/// -/// The method returns a draft since the actual full event needs the table schema version -/// from the schema store, and we don't want to pollute the logic of this module with schema -/// store logic. -pub async fn parse_event_from_relation_message( - old_table_schema: Arc, +/// This method parses the replication protocol relation message and builds a vector of all the +/// columns that were received in the relation message. +pub async fn parse_column_schemas_from_relation_message( relation_body: &protocol::RelationBody, -) -> EtlResult { +) -> EtlResult> { // We construct the new column schemas in order. The order is important since the table schema // relies on the right ordering to interpret the Postgres correctly. let new_column_schemas = relation_body @@ -68,13 +63,7 @@ pub async fn parse_event_from_relation_message( .map(build_column_schema) .collect::, EtlError>>()?; - let new_table_schema = TableSchemaDraft::new( - old_table_schema.id, - old_table_schema.name.clone(), - new_column_schemas, - ); - - Ok(new_table_schema) + Ok(new_column_schemas) } /// Converts a Postgres insert message into an [`InsertEvent`]. diff --git a/etl/src/replication/apply.rs b/etl/src/replication/apply.rs index 756e26406..ddeb8ff0e 100644 --- a/etl/src/replication/apply.rs +++ b/etl/src/replication/apply.rs @@ -1,6 +1,6 @@ use etl_config::shared::PipelineConfig; use etl_postgres::replication::worker::WorkerType; -use etl_postgres::types::{TableId, TableSchema}; +use etl_postgres::types::{TableId, TableSchema, TableSchemaDraft}; use futures::StreamExt; use metrics::histogram; use postgres_replication::protocol; @@ -17,9 +17,9 @@ use crate::concurrency::shutdown::ShutdownRx; use crate::concurrency::signal::SignalRx; use crate::concurrency::stream::{TimeoutStream, TimeoutStreamResult}; use crate::conversions::event::{ - parse_event_from_begin_message, parse_event_from_commit_message, - parse_event_from_delete_message, parse_event_from_insert_message, - parse_event_from_relation_message, parse_event_from_truncate_message, + parse_column_schemas_from_relation_message, parse_event_from_begin_message, + parse_event_from_commit_message, parse_event_from_delete_message, + parse_event_from_insert_message, parse_event_from_truncate_message, parse_event_from_update_message, }; use crate::destination::Destination; @@ -1104,14 +1104,27 @@ where return Ok(HandleMessageResult::no_event()); } - // Convert event from the protocol message. + // Parse the relation message columns into column schemas. + let new_column_schemas = parse_column_schemas_from_relation_message(message).await?; + + // We load the latest table schema before this relation message, which contains the last known + // schema. let old_table_schema = load_latest_table_schema(schema_store, table_id).await?; - let new_table_schema_draft = - parse_event_from_relation_message(old_table_schema.clone(), message).await?; + + // If the column schemas are the same, we treat the relation message as a no-op. This is pretty + // common since Postgres will send a `Relation` message as the first message for every new + // connection even if the table schema hasn't changed. + if new_column_schemas == old_table_schema.column_schemas { + return Ok(HandleMessageResult::no_event()); + } // We store the new schema in the store and build the final relation event. let new_table_schema = schema_store - .store_table_schema(new_table_schema_draft) + .store_table_schema(TableSchemaDraft::new( + table_id, + old_table_schema.name.clone(), + new_column_schemas, + )) .await?; let event = RelationEvent { From f642630159a9aa754780ec3fdc3da4967e94c432 Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Wed, 1 Oct 2025 16:04:15 -0700 Subject: [PATCH 25/45] Improve --- etl-destinations/Cargo.toml | 1 + etl-destinations/src/bigquery/core.rs | 4 +- etl-destinations/src/iceberg/client.rs | 4 +- etl-destinations/src/iceberg/schema.rs | 4 +- etl-destinations/tests/bigquery_pipeline.rs | 153 +++++++++++++++++++- etl-destinations/tests/iceberg_pipeline.rs | 12 +- etl-destinations/tests/support/bigquery.rs | 72 +++++++++ etl-postgres/src/replication/schema.rs | 16 +- etl-postgres/src/types/schema.rs | 20 +-- etl/src/conversions/event.rs | 8 +- etl/src/replication/apply.rs | 6 +- etl/src/replication/client.rs | 12 +- etl/src/store/both/memory.rs | 14 +- etl/src/store/both/postgres.rs | 14 +- etl/src/store/schema/base.rs | 10 +- etl/src/test_utils/notify.rs | 16 +- etl/src/test_utils/table.rs | 6 +- etl/src/test_utils/test_schema.rs | 4 +- etl/src/types/event.rs | 6 +- etl/tests/failpoints_pipeline.rs | 6 +- etl/tests/postgres_store.rs | 10 +- 21 files changed, 307 insertions(+), 91 deletions(-) diff --git a/etl-destinations/Cargo.toml b/etl-destinations/Cargo.toml index 062ac3926..926b8dece 100644 --- a/etl-destinations/Cargo.toml +++ b/etl-destinations/Cargo.toml @@ -46,6 +46,7 @@ uuid = { workspace = true, optional = true, features = ["v4"] } [dev-dependencies] etl = { workspace = true, features = ["test-utils"] } +etl-postgres = { workspace = true, features = ["test-utils", "sqlx"] } etl-telemetry = { workspace = true } base64 = { workspace = true } diff --git a/etl-destinations/src/bigquery/core.rs b/etl-destinations/src/bigquery/core.rs index 4ec559322..01ee76975 100644 --- a/etl-destinations/src/bigquery/core.rs +++ b/etl-destinations/src/bigquery/core.rs @@ -4,7 +4,7 @@ use etl::store::schema::SchemaStore; use etl::store::state::StateStore; use etl::types::{ Cell, Event, PgLsn, RelationChange, RelationEvent, SchemaVersion, TableId, TableName, TableRow, - TableSchema, + VersionedTableSchema, }; use etl::{bail, etl_error}; use gcp_bigquery_client::storage::TableDescriptor; @@ -268,7 +268,7 @@ where &self, table_id: &TableId, schema_version: Option, - ) -> EtlResult<(SequencedBigQueryTableId, Arc)> { + ) -> EtlResult<(SequencedBigQueryTableId, Arc)> { // We hold the lock for the entire preparation to avoid race conditions since the consistency // of this code path is critical. let mut inner = self.inner.lock().await; diff --git a/etl-destinations/src/iceberg/client.rs b/etl-destinations/src/iceberg/client.rs index 07e1b474a..185a78957 100644 --- a/etl-destinations/src/iceberg/client.rs +++ b/etl-destinations/src/iceberg/client.rs @@ -3,7 +3,7 @@ use std::{collections::HashMap, sync::Arc}; use arrow::array::RecordBatch; use etl::{ error::EtlResult, - types::{TableRow, TableSchema}, + types::{TableRow, VersionedTableSchema}, }; use iceberg::{ Catalog, NamespaceIdent, TableCreation, TableIdent, @@ -74,7 +74,7 @@ impl IcebergClient { &self, namespace: &str, table_name: String, - table_schema: &TableSchema, + table_schema: &VersionedTableSchema, ) -> Result<(), iceberg::Error> { let namespace_ident = NamespaceIdent::from_strs(namespace.split('.'))?; let table_ident = TableIdent::new(namespace_ident.clone(), table_name.clone()); diff --git a/etl-destinations/src/iceberg/schema.rs b/etl-destinations/src/iceberg/schema.rs index 54d4eb4e4..1fe0efa19 100644 --- a/etl-destinations/src/iceberg/schema.rs +++ b/etl-destinations/src/iceberg/schema.rs @@ -1,6 +1,6 @@ use std::sync::Arc; -use etl::types::{TableSchema, Type, is_array_type}; +use etl::types::{VersionedTableSchema, Type, is_array_type}; use iceberg::spec::{ ListType, NestedField, PrimitiveType, Schema as IcebergSchema, Type as IcebergType, }; @@ -79,7 +79,7 @@ fn create_iceberg_list_type(element_type: PrimitiveType, field_id: i32) -> Icebe } /// Converts a Postgres table schema to an Iceberg schema. -pub fn postgres_to_iceberg_schema(schema: &TableSchema) -> Result { +pub fn postgres_to_iceberg_schema(schema: &VersionedTableSchema) -> Result { let mut fields = Vec::new(); let mut field_id = 1; diff --git a/etl-destinations/tests/bigquery_pipeline.rs b/etl-destinations/tests/bigquery_pipeline.rs index f24ebd948..bc0799c62 100644 --- a/etl-destinations/tests/bigquery_pipeline.rs +++ b/etl-destinations/tests/bigquery_pipeline.rs @@ -8,18 +8,23 @@ 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}; use etl::test_utils::test_destination_wrapper::TestDestinationWrapper; -use etl::test_utils::test_schema::{TableSelection, insert_mock_data, setup_test_database_schema}; +use etl::test_utils::test_schema::{ + TableSelection, insert_mock_data, insert_users_data, setup_test_database_schema, +}; use etl::types::{EventType, PgNumeric, PipelineId}; -use etl_destinations::bigquery::install_crypto_provider_for_bigquery; +use etl_destinations::bigquery::{ + install_crypto_provider_for_bigquery, table_name_to_bigquery_table_id, +}; use etl_telemetry::tracing::init_test_tracing; use rand::random; use std::str::FromStr; use std::time::Duration; use tokio::time::sleep; - +use tracing::debug; +use etl_postgres::tokio::test_utils::TableModification; use crate::support::bigquery::{ - BigQueryOrder, BigQueryUser, NonNullableColsScalar, NullableColsArray, NullableColsScalar, - parse_bigquery_table_rows, setup_bigquery_connection, + BigQueryDatabase, BigQueryOrder, BigQueryUser, BigQueryUserWithTenant, NonNullableColsScalar, + NullableColsArray, NullableColsScalar, parse_bigquery_table_rows, setup_bigquery_connection, }; mod support; @@ -494,6 +499,143 @@ async fn table_truncate_with_batching() { ); } +#[tokio::test(flavor = "multi_thread")] +async fn relation_event_primary_key_syncs_in_bigquery() { + init_test_tracing(); + install_crypto_provider_for_bigquery(); + + let mut database = spawn_source_database().await; + let database_schema = setup_test_database_schema(&database, TableSelection::UsersOnly).await; + + let users_schema = database_schema.users_schema(); + insert_users_data(&mut database, &users_schema.name, 1..=2).await; + + let bigquery_database = setup_bigquery_connection().await; + + let store = NotifyingStore::new(); + let pipeline_id: PipelineId = random(); + let raw_destination = bigquery_database.build_destination(store.clone()).await; + let destination = TestDestinationWrapper::wrap(raw_destination); + + let mut pipeline = create_pipeline( + &database.config, + pipeline_id, + database_schema.publication_name(), + store.clone(), + destination.clone(), + ); + + let users_state_notify = store + .notify_on_table_state_type(users_schema.id, TableReplicationPhaseType::SyncDone) + .await; + + pipeline.start().await.unwrap(); + + users_state_notify.notified().await; + + // We check the initial BigQuery data. + let users_rows = bigquery_database + .query_table(database_schema.users_schema().name) + .await + .unwrap(); + let parsed_users_rows = parse_bigquery_table_rows::(users_rows); + assert_eq!( + parsed_users_rows, + vec![ + BigQueryUser::new(1, "user_1", 1), + BigQueryUser::new(2, "user_2", 2), + ] + ); + + // let base_table_id = table_name_to_bigquery_table_id(&users_schema.name); + // let sequenced_table_id = format!("{base_table_id}_0"); + // + // async fn wait_for_primary_key_columns( + // bigquery_database: &BigQueryDatabase, + // table_id: &str, + // expected_columns: &[&str], + // ) { + // let expected = expected_columns + // .iter() + // .map(|column| column.to_string()) + // .collect::>(); + // + // for attempt in 0..30 { + // let actual = bigquery_database.primary_key_columns(table_id).await; + // if actual == expected { + // return; + // } + // + // debug!( + // table = table_id, + // ?actual, + // ?expected, + // attempt, + // "waiting for primary key columns to sync" + // ); + // sleep(Duration::from_millis(500)).await; + // } + // + // panic!( + // "primary key columns for {table_id} did not match expected {:?}", + // expected_columns + // ); + // } + // wait_for_primary_key_columns(&bigquery_database, &sequenced_table_id, &["id"]).await; + + // We perform schema changes. + database + .alter_table( + test_table_name("users"), + &[ + TableModification::AddColumn { + name: "year", + params: "integer", + }, + TableModification::RenameColumn { + name: "age", + new_name: "new_age", + }, + ], + ) + .await + .unwrap(); + + // Register notifications for the insert. + let insert_event_notify = destination + .wait_for_events_count(vec![(EventType::Insert, 1)]) + .await; + + // We insert data. + database + .insert_values( + database_schema.users_schema().name.clone(), + &["name", "new_age", "year"], + &[&"user_3", &(3i32), &(2025i32)], + ) + .await + .expect("Failed to insert users"); + + insert_event_notify.notified().await; + + // We check the BigQuery data after the first schema change. + let users_rows = bigquery_database + .query_table(database_schema.users_schema().name) + .await + .unwrap(); + println!("USERS ROWS {:?}", users_rows); + // let parsed_users_rows = parse_bigquery_table_rows::(users_rows); + // assert_eq!( + // parsed_users_rows, + // vec![ + // BigQueryUser::new(1, "user_2", 1), + // BigQueryUser::new(2, "user_2", 2), + // ] + // ); + + pipeline.shutdown_and_wait().await.unwrap(); +} + #[tokio::test(flavor = "multi_thread")] async fn table_nullable_scalar_columns() { init_test_tracing(); @@ -1580,6 +1722,7 @@ async fn table_array_with_null_values() { // We have to reset the state of the table and copy it from scratch, otherwise the CDC will contain // the inserts and deletes, failing again. store.reset_table_state(table_id).await.unwrap(); + // We also clear the events so that it's more idiomatic to wait for them, since we don't have // the insert of before. destination.clear_events().await; diff --git a/etl-destinations/tests/iceberg_pipeline.rs b/etl-destinations/tests/iceberg_pipeline.rs index c4c3bc213..e3f20b441 100644 --- a/etl-destinations/tests/iceberg_pipeline.rs +++ b/etl-destinations/tests/iceberg_pipeline.rs @@ -3,7 +3,7 @@ use std::collections::HashMap; use chrono::{DateTime, NaiveDate, NaiveDateTime, NaiveTime, Utc}; -use etl::types::{ArrayCell, Cell, ColumnSchema, TableId, TableName, TableRow, TableSchema, Type}; +use etl::types::{ArrayCell, Cell, ColumnSchema, TableId, TableName, TableRow, VersionedTableSchema, Type}; use etl_destinations::iceberg::IcebergClient; use etl_telemetry::tracing::init_test_tracing; use iceberg::io::{S3_ACCESS_KEY_ID, S3_ENDPOINT, S3_SECRET_ACCESS_KEY}; @@ -216,7 +216,7 @@ async fn create_table_if_missing() { ColumnSchema::new("oid_array_col".to_string(), Type::OID_ARRAY, -1, false), ColumnSchema::new("bytea_array_col".to_string(), Type::BYTEA_ARRAY, -1, false), ]; - let table_schema = TableSchema::new(table_id, table_name_struct, 0, columns); + let table_schema = VersionedTableSchema::new(table_id, table_name_struct, 0, columns); // table doesn't exist yet assert!( @@ -318,7 +318,7 @@ async fn insert_nullable_scalars() { // Binary type ColumnSchema::new("bytea_col".to_string(), Type::BYTEA, -1, false), ]; - let table_schema = TableSchema::new(table_id, table_name_struct, 0, columns); + let table_schema = VersionedTableSchema::new(table_id, table_name_struct, 0, columns); client .create_table_if_missing(namespace, table_name.clone(), &table_schema) @@ -469,7 +469,7 @@ async fn insert_non_nullable_scalars() { // Binary type ColumnSchema::new("bytea_col".to_string(), Type::BYTEA, -1, false), ]; - let table_schema = TableSchema::new(table_id, table_name_struct, 0, columns); + let table_schema = VersionedTableSchema::new(table_id, table_name_struct, 0, columns); client .create_table_if_missing(namespace, table_name.clone(), &table_schema) @@ -629,7 +629,7 @@ async fn insert_nullable_array() { // Binary array type ColumnSchema::new("bytea_array_col".to_string(), Type::BYTEA_ARRAY, -1, false), ]; - let table_schema = TableSchema::new(table_id, table_name_struct, 0, columns); + let table_schema = VersionedTableSchema::new(table_id, table_name_struct, 0, columns); client .create_table_if_missing(namespace, table_name.clone(), &table_schema) @@ -903,7 +903,7 @@ async fn insert_non_nullable_array() { // Binary array type ColumnSchema::new("bytea_array_col".to_string(), Type::BYTEA_ARRAY, -1, false), ]; - let table_schema = TableSchema::new(table_id, table_name_struct, 0, columns); + let table_schema = VersionedTableSchema::new(table_id, table_name_struct, 0, columns); client .create_table_if_missing(namespace, table_name.clone(), &table_schema) diff --git a/etl-destinations/tests/support/bigquery.rs b/etl-destinations/tests/support/bigquery.rs index e46ce409c..e586a716e 100644 --- a/etl-destinations/tests/support/bigquery.rs +++ b/etl-destinations/tests/support/bigquery.rs @@ -256,6 +256,46 @@ pub async fn setup_bigquery_connection() -> BigQueryDatabase { BigQueryDatabase::new_real(sa_key_path).await } +impl BigQueryDatabase { + /// Fetches primary key column names for a BigQuery table ordered by ordinal position. + pub async fn primary_key_columns(&self, table_id: &str) -> Vec { + let client = self.client().expect("BigQuery client not available"); + let project_id = self.project_id(); + let dataset_id = self.dataset_id(); + + let query = format!( + "SELECT column_name \ + FROM `{project}.{dataset}.INFORMATION_SCHEMA.KEY_COLUMN_USAGE` \ + WHERE table_name = '{table}' AND constraint_name = 'PRIMARY KEY' \ + ORDER BY ordinal_position", + project = project_id, + dataset = dataset_id, + table = table_id + ); + + let response = client + .job() + .query(project_id, QueryRequest::new(query)) + .await + .expect("failed to query BigQuery primary key metadata"); + + response + .rows + .unwrap_or_default() + .into_iter() + .filter_map(|row| { + row.columns.and_then(|mut columns| { + columns + .into_iter() + .next() + .and_then(|cell| cell.value) + .and_then(|value| value.as_str().map(|s| s.to_owned())) + }) + }) + .collect() + } +} + pub fn parse_table_cell(table_cell: TableCell) -> Option where O: FromStr, @@ -321,6 +361,38 @@ impl From for BigQueryOrder { } } +#[derive(Debug, Eq, PartialEq, Ord, PartialOrd)] +pub struct BigQueryUserWithTenant { + id: i32, + name: String, + age: i32, + tenant_id: i32, +} + +impl BigQueryUserWithTenant { + pub fn new(id: i32, name: &str, age: i32, tenant_id: i32) -> Self { + Self { + id, + name: name.to_owned(), + age, + tenant_id, + } + } +} + +impl From for BigQueryUserWithTenant { + fn from(value: TableRow) -> Self { + let columns = value.columns.unwrap(); + + BigQueryUserWithTenant { + id: parse_table_cell(columns[0].clone()).unwrap(), + name: parse_table_cell(columns[1].clone()).unwrap(), + age: parse_table_cell(columns[2].clone()).unwrap(), + tenant_id: parse_table_cell(columns[3].clone()).unwrap(), + } + } +} + #[derive(Debug)] pub struct NullableColsScalar { id: i32, diff --git a/etl-postgres/src/replication/schema.rs b/etl-postgres/src/replication/schema.rs index cb56dc748..023634950 100644 --- a/etl-postgres/src/replication/schema.rs +++ b/etl-postgres/src/replication/schema.rs @@ -5,7 +5,7 @@ use std::collections::HashMap; use tokio_postgres::types::Type as PgType; use crate::types::{ - ColumnSchema, SchemaVersion, TableId, TableName, TableSchema, TableSchemaDraft, + ColumnSchema, SchemaVersion, TableId, TableName, VersionedTableSchema, TableSchema, }; /// The initial schema version number. @@ -146,8 +146,8 @@ define_type_mappings! { pub async fn store_table_schema( pool: &PgPool, pipeline_id: i64, - table_schema: TableSchemaDraft, -) -> Result { + table_schema: TableSchema, +) -> Result { let mut tx = pool.begin().await?; // We are fine with running this query for every table schema store since we assume that @@ -210,17 +210,17 @@ pub async fn store_table_schema( tx.commit().await?; - Ok(table_schema.into_table_schema(next_schema_version)) + Ok(table_schema.into_versioned(next_schema_version)) } /// Loads all table schemas for a pipeline from the database. /// /// Retrieves table schemas and columns from schema storage tables, -/// reconstructing complete [`TableSchema`] objects. +/// reconstructing complete [`VersionedTableSchema`] objects. pub async fn load_table_schemas( pool: &PgPool, pipeline_id: i64, -) -> Result, sqlx::Error> { +) -> Result, sqlx::Error> { let rows = sqlx::query( r#" select @@ -244,7 +244,7 @@ pub async fn load_table_schemas( .fetch_all(pool) .await?; - let mut table_schemas: HashMap<(TableId, SchemaVersion), TableSchema> = HashMap::new(); + let mut table_schemas: HashMap<(TableId, SchemaVersion), VersionedTableSchema> = HashMap::new(); for row in rows { let table_oid: SqlxTableId = row.get("table_id"); @@ -261,7 +261,7 @@ pub async fn load_table_schemas( let entry = table_schemas .entry((table_id, schema_version)) .or_insert_with(|| { - TableSchema::new( + VersionedTableSchema::new( table_id, TableName::new(schema_name.clone(), table_name.clone()), schema_version, diff --git a/etl-postgres/src/types/schema.rs b/etl-postgres/src/types/schema.rs index ea01bf275..f357d9778 100644 --- a/etl-postgres/src/types/schema.rs +++ b/etl-postgres/src/types/schema.rs @@ -171,7 +171,7 @@ impl ToSql for TableId { /// This type contains all metadata about a table including its name, OID, /// and the schemas of all its columns. #[derive(Debug, Clone, Eq, PartialEq)] -pub struct TableSchema { +pub struct VersionedTableSchema { /// The Postgres OID of the table pub id: TableId, /// The fully qualified name of the table @@ -182,7 +182,7 @@ pub struct TableSchema { pub column_schemas: Vec, } -impl TableSchema { +impl VersionedTableSchema { pub fn new( id: TableId, name: TableName, @@ -197,19 +197,19 @@ impl TableSchema { } } - /// Adds a new column schema to this [`TableSchema`]. + /// Adds a new column schema to this [`VersionedTableSchema`]. pub fn add_column_schema(&mut self, column_schema: ColumnSchema) { self.column_schemas.push(column_schema); } } -impl PartialOrd for TableSchema { +impl PartialOrd for VersionedTableSchema { fn partial_cmp(&self, other: &Self) -> Option { Some(self.cmp(other)) } } -impl Ord for TableSchema { +impl Ord for VersionedTableSchema { fn cmp(&self, other: &Self) -> Ordering { self.name .cmp(&other.name) @@ -217,15 +217,15 @@ impl Ord for TableSchema { } } -/// Draft version of [`TableSchema`] used before a schema version is assigned. +/// Draft version of [`VersionedTableSchema`] used before a schema version is assigned. #[derive(Debug, Clone, Eq, PartialEq)] -pub struct TableSchemaDraft { +pub struct TableSchema { pub id: TableId, pub name: TableName, pub column_schemas: Vec, } -impl TableSchemaDraft { +impl TableSchema { pub fn new(id: TableId, name: TableName, column_schemas: Vec) -> Self { Self { id, @@ -253,7 +253,7 @@ impl TableSchemaDraft { self.column_schemas.iter().any(|cs| cs.primary) } - pub fn into_table_schema(self, version: SchemaVersion) -> TableSchema { - TableSchema::new(self.id, self.name, version, self.column_schemas) + pub fn into_versioned(self, version: SchemaVersion) -> VersionedTableSchema { + VersionedTableSchema::new(self.id, self.name, version, self.column_schemas) } } diff --git a/etl/src/conversions/event.rs b/etl/src/conversions/event.rs index 00c8c45e8..97cb57a45 100644 --- a/etl/src/conversions/event.rs +++ b/etl/src/conversions/event.rs @@ -1,6 +1,6 @@ use core::str; use etl_postgres::types::{ - ColumnSchema, SchemaVersion, TableId, TableSchema, convert_type_oid_to_type, + ColumnSchema, SchemaVersion, TableId, VersionedTableSchema, convert_type_oid_to_type, }; use postgres_replication::protocol; use std::sync::Arc; @@ -71,7 +71,7 @@ pub async fn parse_column_schemas_from_relation_message( /// This function processes an insert operation from the replication stream /// using the supplied table schema version to build the resulting event. pub fn parse_event_from_insert_message( - table_schema: Arc, + table_schema: Arc, start_lsn: PgLsn, commit_lsn: PgLsn, insert_body: &protocol::InsertBody, @@ -99,7 +99,7 @@ pub fn parse_event_from_insert_message( /// the complete row or just the key columns, depending on the table's /// `REPLICA IDENTITY` setting in Postgres. pub fn parse_event_from_update_message( - table_schema: Arc, + table_schema: Arc, start_lsn: PgLsn, commit_lsn: PgLsn, update_body: &protocol::UpdateBody, @@ -145,7 +145,7 @@ pub fn parse_event_from_update_message( /// either the complete row or just the key columns, depending on the table's /// `REPLICA IDENTITY` setting in Postgres. pub fn parse_event_from_delete_message( - table_schema: Arc, + table_schema: Arc, start_lsn: PgLsn, commit_lsn: PgLsn, delete_body: &protocol::DeleteBody, diff --git a/etl/src/replication/apply.rs b/etl/src/replication/apply.rs index ddeb8ff0e..831a5a7a1 100644 --- a/etl/src/replication/apply.rs +++ b/etl/src/replication/apply.rs @@ -1,6 +1,6 @@ use etl_config::shared::PipelineConfig; use etl_postgres::replication::worker::WorkerType; -use etl_postgres::types::{TableId, TableSchema, TableSchemaDraft}; +use etl_postgres::types::{TableId, VersionedTableSchema, TableSchema}; use futures::StreamExt; use metrics::histogram; use postgres_replication::protocol; @@ -1120,7 +1120,7 @@ where // We store the new schema in the store and build the final relation event. let new_table_schema = schema_store - .store_table_schema(TableSchemaDraft::new( + .store_table_schema(TableSchema::new( table_id, old_table_schema.name.clone(), new_column_schemas, @@ -1308,7 +1308,7 @@ where async fn load_latest_table_schema( schema_store: &S, table_id: TableId, -) -> EtlResult> +) -> EtlResult> where S: SchemaStore + Clone + Send + 'static, { diff --git a/etl/src/replication/client.rs b/etl/src/replication/client.rs index b339b72e8..202e2fecd 100644 --- a/etl/src/replication/client.rs +++ b/etl/src/replication/client.rs @@ -4,7 +4,7 @@ use crate::{bail, etl_error}; use etl_config::shared::{IntoConnectOptions, PgConnectionConfig}; use etl_postgres::replication::extract_server_version; use etl_postgres::types::{ColumnSchema, TableId, TableName}; -use etl_postgres::types::{TableSchemaDraft, convert_type_oid_to_type}; +use etl_postgres::types::{TableSchema, convert_type_oid_to_type}; use pg_escape::{quote_identifier, quote_literal}; use postgres_replication::LogicalReplicationStream; use rustls::ClientConfig; @@ -117,7 +117,7 @@ impl PgReplicationSlotTransaction { &self, table_ids: &[TableId], publication_name: Option<&str>, - ) -> EtlResult> { + ) -> EtlResult> { self.client .get_table_schemas(table_ids, publication_name) .await @@ -131,7 +131,7 @@ impl PgReplicationSlotTransaction { &self, table_id: TableId, publication: Option<&str>, - ) -> EtlResult { + ) -> EtlResult { self.client.get_table_schema(table_id, publication).await } @@ -558,7 +558,7 @@ impl PgReplicationClient { &self, table_ids: &[TableId], publication_name: Option<&str>, - ) -> EtlResult> { + ) -> EtlResult> { let mut table_schemas = HashMap::new(); // TODO: consider if we want to fail when at least one table was missing or not. @@ -589,11 +589,11 @@ impl PgReplicationClient { &self, table_id: TableId, publication: Option<&str>, - ) -> EtlResult { + ) -> EtlResult { let table_name = self.get_table_name(table_id).await?; let column_schemas = self.get_column_schemas(table_id, publication).await?; - Ok(TableSchemaDraft { + Ok(TableSchema { name: table_name, id: table_id, column_schemas, diff --git a/etl/src/store/both/memory.rs b/etl/src/store/both/memory.rs index cb8890537..82cc95b35 100644 --- a/etl/src/store/both/memory.rs +++ b/etl/src/store/both/memory.rs @@ -1,4 +1,4 @@ -use etl_postgres::types::{SchemaVersion, TableId, TableSchema, TableSchemaDraft}; +use etl_postgres::types::{SchemaVersion, TableId, VersionedTableSchema, TableSchema}; use std::collections::{BTreeMap, HashMap}; use std::sync::Arc; use tokio::sync::Mutex; @@ -22,7 +22,7 @@ struct Inner { table_state_history: HashMap>, /// Table schema definitions for each table. Each schema has multiple versions which are created /// when new schema changes are detected. - table_schemas: HashMap>>, + table_schemas: HashMap>>, /// Mappings between source and destination tables. table_mappings: HashMap, } @@ -174,7 +174,7 @@ impl SchemaStore for MemoryStore { &self, table_id: &TableId, version: SchemaVersion, - ) -> EtlResult>> { + ) -> EtlResult>> { let inner = self.inner.lock().await; Ok(inner @@ -186,7 +186,7 @@ impl SchemaStore for MemoryStore { async fn get_latest_table_schema( &self, table_id: &TableId, - ) -> EtlResult>> { + ) -> EtlResult>> { let inner = self.inner.lock().await; Ok(inner.table_schemas.get(table_id).and_then(|table_schemas| { @@ -209,8 +209,8 @@ impl SchemaStore for MemoryStore { async fn store_table_schema( &self, - table_schema: TableSchemaDraft, - ) -> EtlResult> { + table_schema: TableSchema, + ) -> EtlResult> { let mut inner = self.inner.lock().await; let table_schemas = inner .table_schemas @@ -223,7 +223,7 @@ impl SchemaStore for MemoryStore { .map(|version| version + 1) .unwrap_or(STARTING_SCHEMA_VERSION); - let table_schema = Arc::new(table_schema.into_table_schema(next_version)); + let table_schema = Arc::new(table_schema.into_versioned(next_version)); table_schemas.insert(next_version, table_schema.clone()); Ok(table_schema) diff --git a/etl/src/store/both/postgres.rs b/etl/src/store/both/postgres.rs index b23d410ae..ec7d43e93 100644 --- a/etl/src/store/both/postgres.rs +++ b/etl/src/store/both/postgres.rs @@ -1,6 +1,6 @@ use etl_config::shared::PgConnectionConfig; use etl_postgres::replication::{connect_to_source_database, schema, state, table_mappings}; -use etl_postgres::types::{SchemaVersion, TableId, TableSchema, TableSchemaDraft}; +use etl_postgres::types::{SchemaVersion, TableId, VersionedTableSchema, TableSchema}; use metrics::gauge; use sqlx::PgPool; use std::{ @@ -158,7 +158,7 @@ struct Inner { /// Cached table replication states indexed by table ID. table_states: HashMap, /// Cached table schemas indexed by table ID. - table_schemas: HashMap>>, + table_schemas: HashMap>>, /// Cached table mappings from source table ID to destination table name. table_mappings: HashMap, } @@ -237,7 +237,7 @@ impl PostgresStore { #[cfg(feature = "test-utils")] pub async fn get_all_table_schemas( &self, - ) -> HashMap>> { + ) -> HashMap>> { let inner = self.inner.lock().await; inner.table_schemas.clone() } @@ -522,7 +522,7 @@ impl SchemaStore for PostgresStore { &self, table_id: &TableId, version: SchemaVersion, - ) -> EtlResult>> { + ) -> EtlResult>> { let inner = self.inner.lock().await; Ok(inner @@ -535,7 +535,7 @@ impl SchemaStore for PostgresStore { async fn get_latest_table_schema( &self, table_id: &TableId, - ) -> EtlResult>> { + ) -> EtlResult>> { let inner = self.inner.lock().await; Ok(inner @@ -593,8 +593,8 @@ impl SchemaStore for PostgresStore { /// replication or when schema definitions need to be updated. async fn store_table_schema( &self, - table_schema: TableSchemaDraft, - ) -> EtlResult> { + table_schema: TableSchema, + ) -> EtlResult> { debug!("storing table schema for table '{}'", table_schema.name); let pool = self.connect_to_source().await?; diff --git a/etl/src/store/schema/base.rs b/etl/src/store/schema/base.rs index e437150a4..537292518 100644 --- a/etl/src/store/schema/base.rs +++ b/etl/src/store/schema/base.rs @@ -1,4 +1,4 @@ -use etl_postgres::types::{SchemaVersion, TableId, TableSchema, TableSchemaDraft}; +use etl_postgres::types::{SchemaVersion, TableId, VersionedTableSchema, TableSchema}; use std::sync::Arc; use crate::error::EtlResult; @@ -17,13 +17,13 @@ pub trait SchemaStore { &self, table_id: &TableId, version: SchemaVersion, - ) -> impl Future>>> + Send; + ) -> impl Future>>> + Send; /// Returns the latest table schema version for table with id `table_id` from the cache. fn get_latest_table_schema( &self, table_id: &TableId, - ) -> impl Future>>> + Send; + ) -> impl Future>>> + Send; /// Loads table schemas from the persistent state into the cache. /// @@ -33,6 +33,6 @@ pub trait SchemaStore { /// Stores a table schema in both the cache and the persistent store. fn store_table_schema( &self, - table_schema: TableSchemaDraft, - ) -> impl Future>> + Send; + table_schema: TableSchema, + ) -> impl Future>> + Send; } diff --git a/etl/src/test_utils/notify.rs b/etl/src/test_utils/notify.rs index 0cad5bbfc..547f904a2 100644 --- a/etl/src/test_utils/notify.rs +++ b/etl/src/test_utils/notify.rs @@ -1,4 +1,4 @@ -use etl_postgres::types::{SchemaVersion, TableId, TableSchema, TableSchemaDraft}; +use etl_postgres::types::{SchemaVersion, TableId, VersionedTableSchema, TableSchema}; use std::{ collections::{BTreeMap, HashMap}, fmt, @@ -32,7 +32,7 @@ type TableStateCondition = ( struct Inner { table_replication_states: HashMap, table_state_history: HashMap>, - table_schemas: HashMap>>, + table_schemas: HashMap>>, table_mappings: HashMap, table_state_type_conditions: Vec, table_state_conditions: Vec, @@ -106,7 +106,7 @@ impl NotifyingStore { inner.table_replication_states.clone() } - pub async fn get_latest_table_schemas(&self) -> HashMap { + pub async fn get_latest_table_schemas(&self) -> HashMap { let inner = self.inner.read().await; inner .table_schemas @@ -302,7 +302,7 @@ impl SchemaStore for NotifyingStore { &self, table_id: &TableId, version: SchemaVersion, - ) -> EtlResult>> { + ) -> EtlResult>> { let inner = self.inner.read().await; Ok(inner @@ -314,7 +314,7 @@ impl SchemaStore for NotifyingStore { async fn get_latest_table_schema( &self, table_id: &TableId, - ) -> EtlResult>> { + ) -> EtlResult>> { let inner = self.inner.read().await; Ok(inner @@ -334,8 +334,8 @@ impl SchemaStore for NotifyingStore { async fn store_table_schema( &self, - table_schema: TableSchemaDraft, - ) -> EtlResult> { + table_schema: TableSchema, + ) -> EtlResult> { let mut inner = self.inner.write().await; let schemas = inner .table_schemas @@ -348,7 +348,7 @@ impl SchemaStore for NotifyingStore { .map(|version| version + 1) .unwrap_or(0); - let schema = Arc::new(table_schema.into_table_schema(next_version)); + let schema = Arc::new(table_schema.into_versioned(next_version)); schemas.insert(next_version, Arc::clone(&schema)); Ok(schema) diff --git a/etl/src/test_utils/table.rs b/etl/src/test_utils/table.rs index af53a086f..2a6781d59 100644 --- a/etl/src/test_utils/table.rs +++ b/etl/src/test_utils/table.rs @@ -1,8 +1,8 @@ -use etl_postgres::types::{ColumnSchema, TableId, TableName, TableSchema, TableSchemaDraft}; +use etl_postgres::types::{ColumnSchema, TableId, TableName, VersionedTableSchema, TableSchema}; use std::collections::HashMap; /// Return the names of the column schema. -pub fn column_schema_names(table_schema: &TableSchema) -> Vec { +pub fn column_schema_names(table_schema: &VersionedTableSchema) -> Vec { table_schema .column_schemas .iter() @@ -21,7 +21,7 @@ pub fn column_schema_names(table_schema: &TableSchema) -> Vec { /// Panics if the table ID doesn't exist in the provided schemas, or if any aspect /// of the schema doesn't match the expected values. pub fn assert_table_schema( - table_schemas: &HashMap, + table_schemas: &HashMap, table_id: TableId, expected_table_name: TableName, expected_columns: &[ColumnSchema], diff --git a/etl/src/test_utils/test_schema.rs b/etl/src/test_utils/test_schema.rs index b6b681de5..77b375d9a 100644 --- a/etl/src/test_utils/test_schema.rs +++ b/etl/src/test_utils/test_schema.rs @@ -1,5 +1,5 @@ use etl_postgres::tokio::test_utils::{PgDatabase, id_column_schema}; -use etl_postgres::types::{ColumnSchema, TableId, TableName, TableSchema}; +use etl_postgres::types::{ColumnSchema, TableId, TableName, VersionedTableSchema, TableSchema}; use std::ops::RangeInclusive; use tokio_postgres::types::{PgLsn, Type}; use tokio_postgres::{Client, GenericClient}; @@ -64,7 +64,6 @@ pub async fn setup_test_database_schema( users_table_schema = Some(TableSchema::new( users_table_id, users_table_name, - 0, vec![ id_column_schema(), ColumnSchema::new("name".to_string(), Type::TEXT, -1, false), @@ -89,7 +88,6 @@ pub async fn setup_test_database_schema( orders_table_schema = Some(TableSchema::new( orders_table_id, orders_table_name, - 0, vec![ id_column_schema(), ColumnSchema::new("description".to_string(), Type::TEXT, -1, false), diff --git a/etl/src/types/event.rs b/etl/src/types/event.rs index 682cf1a7b..3e26d3908 100644 --- a/etl/src/types/event.rs +++ b/etl/src/types/event.rs @@ -1,4 +1,4 @@ -use etl_postgres::types::{ColumnSchema, SchemaVersion, TableId, TableSchema}; +use etl_postgres::types::{ColumnSchema, SchemaVersion, TableId, VersionedTableSchema}; use std::collections::HashSet; use std::fmt; use std::hash::Hash; @@ -71,9 +71,9 @@ pub struct RelationEvent { /// ID of the table of which this is a schema change. pub table_id: TableId, /// The old table schema. - pub old_table_schema: Arc, + pub old_table_schema: Arc, /// The new table schema. - pub new_table_schema: Arc, + pub new_table_schema: Arc, } impl RelationEvent { diff --git a/etl/tests/failpoints_pipeline.rs b/etl/tests/failpoints_pipeline.rs index 8a0efd20f..5361b2596 100644 --- a/etl/tests/failpoints_pipeline.rs +++ b/etl/tests/failpoints_pipeline.rs @@ -211,7 +211,8 @@ async fn table_copy_is_consistent_after_data_sync_threw_an_error_with_timed_retr *table_schemas .get(&database_schema.users_schema().id) .unwrap(), - database_schema.users_schema() + // We expect version 0 since the schema is copied only once. + database_schema.users_schema().into_versioned(0) ); } @@ -274,6 +275,7 @@ async fn table_copy_is_consistent_during_data_sync_threw_an_error_with_timed_ret *table_schemas .get(&database_schema.users_schema().id) .unwrap(), - database_schema.users_schema() + // We expect version 1 since the schema is copied twice. + database_schema.users_schema().into_versioned(1) ); } diff --git a/etl/tests/postgres_store.rs b/etl/tests/postgres_store.rs index 2d8755a65..101ca0a6f 100644 --- a/etl/tests/postgres_store.rs +++ b/etl/tests/postgres_store.rs @@ -7,12 +7,12 @@ use etl::store::schema::SchemaStore; use etl::store::state::StateStore; use etl::test_utils::database::spawn_source_database_for_store; use etl_postgres::replication::connect_to_source_database; -use etl_postgres::types::{ColumnSchema, TableId, TableName, TableSchemaDraft}; +use etl_postgres::types::{ColumnSchema, TableId, TableName, TableSchema}; use etl_telemetry::tracing::init_test_tracing; use sqlx::postgres::types::Oid as SqlxTableId; use tokio_postgres::types::{PgLsn, Type as PgType}; -fn create_sample_table_schema() -> TableSchemaDraft { +fn create_sample_table_schema() -> TableSchema { let table_id = TableId::new(12345); let table_name = TableName::new("public".to_string(), "test_table".to_string()); let columns = vec![ @@ -21,10 +21,10 @@ fn create_sample_table_schema() -> TableSchemaDraft { ColumnSchema::new("created_at".to_string(), PgType::TIMESTAMPTZ, -1, false), ]; - TableSchemaDraft::new(table_id, table_name, columns) + TableSchema::new(table_id, table_name, columns) } -fn create_another_table_schema() -> TableSchemaDraft { +fn create_another_table_schema() -> TableSchema { let table_id = TableId::new(67890); let table_name = TableName::new("public".to_string(), "another_table".to_string()); let columns = vec![ @@ -32,7 +32,7 @@ fn create_another_table_schema() -> TableSchemaDraft { ColumnSchema::new("description".to_string(), PgType::VARCHAR, 255, false), ]; - TableSchemaDraft::new(table_id, table_name, columns) + TableSchema::new(table_id, table_name, columns) } #[tokio::test(flavor = "multi_thread")] From 88c8d006fed9e1ab6eec4a1de8437770038ecc16 Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Thu, 2 Oct 2025 08:46:37 -0700 Subject: [PATCH 26/45] Improve --- etl-destinations/src/bigquery/core.rs | 35 ++++----- etl-destinations/src/iceberg/schema.rs | 6 +- etl-destinations/tests/bigquery_pipeline.rs | 84 ++++++--------------- etl-destinations/tests/iceberg_pipeline.rs | 4 +- etl-postgres/src/replication/schema.rs | 2 +- etl/src/replication/apply.rs | 2 +- etl/src/store/both/memory.rs | 2 +- etl/src/store/both/postgres.rs | 2 +- etl/src/store/schema/base.rs | 2 +- etl/src/test_utils/notify.rs | 2 +- etl/src/test_utils/table.rs | 2 +- etl/src/test_utils/test_schema.rs | 2 +- 12 files changed, 58 insertions(+), 87 deletions(-) diff --git a/etl-destinations/src/bigquery/core.rs b/etl-destinations/src/bigquery/core.rs index 01ee76975..30b062301 100644 --- a/etl-destinations/src/bigquery/core.rs +++ b/etl-destinations/src/bigquery/core.rs @@ -610,18 +610,19 @@ where } } - self.client - .sync_primary_key( - &self.dataset_id, - &sequenced_table_name, - &relation_event.new_table_schema.column_schemas, - ) - .await?; - - debug!( - table = %sequenced_table_name, - "synchronized primary key definition in BigQuery" - ); + // TODO: implement primary key synchronization. + // self.client + // .sync_primary_key( + // &self.dataset_id, + // &sequenced_table_name, + // &relation_event.new_table_schema.column_schemas, + // ) + // .await?; + // + // debug!( + // table = %sequenced_table_name, + // "synchronized primary key definition in BigQuery" + // ); info!( table_id = %relation_event.table_id, @@ -689,7 +690,7 @@ where return Ok(()); } - let Some(schema_version) = batch_schema_version.take() else { + let Some(batch_schema_version) = batch_schema_version.take() else { bail!( ErrorKind::InvalidState, "Missing schema version", @@ -698,7 +699,7 @@ where }; let rows = mem::take(table_id_to_table_rows); - self.process_table_events(schema_version, rows).await + self.process_table_events(batch_schema_version, rows).await } /// Processes CDC events in batches with proper ordering and truncate handling. @@ -759,19 +760,19 @@ where // Batch breaker events. Event::Relation(relation) => { - // Finish current batch before applying schema change. + // Finish the current batch before applying schema change. self.flush_batch(&mut batch_schema_version, &mut table_id_to_table_rows) .await?; // We mark the new batch schema version with the relation schema version, since - // after a relation message a new schema is meant to be stored in the schema store. + // after a relation message, a new schema is meant to be stored in the schema store. batch_schema_version = Some(relation.new_table_schema.version); // Apply relation change, then prime the next batch with the new schema version. self.apply_relation_event_changes(relation).await?; } Event::Truncate(truncate) => { - // Finish current batch before a TRUNCATE (it affects table state). + // Finish the current batch before a TRUNCATE (it affects the table state). self.flush_batch(&mut batch_schema_version, &mut table_id_to_table_rows) .await?; diff --git a/etl-destinations/src/iceberg/schema.rs b/etl-destinations/src/iceberg/schema.rs index 1fe0efa19..d56fc4a3e 100644 --- a/etl-destinations/src/iceberg/schema.rs +++ b/etl-destinations/src/iceberg/schema.rs @@ -1,6 +1,6 @@ use std::sync::Arc; -use etl::types::{VersionedTableSchema, Type, is_array_type}; +use etl::types::{Type, VersionedTableSchema, is_array_type}; use iceberg::spec::{ ListType, NestedField, PrimitiveType, Schema as IcebergSchema, Type as IcebergType, }; @@ -79,7 +79,9 @@ fn create_iceberg_list_type(element_type: PrimitiveType, field_id: i32) -> Icebe } /// Converts a Postgres table schema to an Iceberg schema. -pub fn postgres_to_iceberg_schema(schema: &VersionedTableSchema) -> Result { +pub fn postgres_to_iceberg_schema( + schema: &VersionedTableSchema, +) -> Result { let mut fields = Vec::new(); let mut field_id = 1; diff --git a/etl-destinations/tests/bigquery_pipeline.rs b/etl-destinations/tests/bigquery_pipeline.rs index bc0799c62..689c1e860 100644 --- a/etl-destinations/tests/bigquery_pipeline.rs +++ b/etl-destinations/tests/bigquery_pipeline.rs @@ -1,5 +1,9 @@ #![cfg(feature = "bigquery")] +use crate::support::bigquery::{ + BigQueryDatabase, BigQueryOrder, BigQueryUser, BigQueryUserWithTenant, NonNullableColsScalar, + NullableColsArray, NullableColsScalar, parse_bigquery_table_rows, setup_bigquery_connection, +}; use chrono::{DateTime, NaiveDate, NaiveDateTime, NaiveTime, Utc}; use etl::config::BatchConfig; use etl::error::ErrorKind; @@ -15,17 +19,13 @@ use etl::types::{EventType, PgNumeric, PipelineId}; use etl_destinations::bigquery::{ install_crypto_provider_for_bigquery, table_name_to_bigquery_table_id, }; +use etl_postgres::tokio::test_utils::TableModification; use etl_telemetry::tracing::init_test_tracing; use rand::random; use std::str::FromStr; use std::time::Duration; -use tokio::time::sleep; +use tokio::time::{sleep, timeout}; use tracing::debug; -use etl_postgres::tokio::test_utils::TableModification; -use crate::support::bigquery::{ - BigQueryDatabase, BigQueryOrder, BigQueryUser, BigQueryUserWithTenant, NonNullableColsScalar, - NullableColsArray, NullableColsScalar, parse_bigquery_table_rows, setup_bigquery_connection, -}; mod support; @@ -547,42 +547,6 @@ async fn relation_event_primary_key_syncs_in_bigquery() { ] ); - // let base_table_id = table_name_to_bigquery_table_id(&users_schema.name); - // let sequenced_table_id = format!("{base_table_id}_0"); - // - // async fn wait_for_primary_key_columns( - // bigquery_database: &BigQueryDatabase, - // table_id: &str, - // expected_columns: &[&str], - // ) { - // let expected = expected_columns - // .iter() - // .map(|column| column.to_string()) - // .collect::>(); - // - // for attempt in 0..30 { - // let actual = bigquery_database.primary_key_columns(table_id).await; - // if actual == expected { - // return; - // } - // - // debug!( - // table = table_id, - // ?actual, - // ?expected, - // attempt, - // "waiting for primary key columns to sync" - // ); - // sleep(Duration::from_millis(500)).await; - // } - // - // panic!( - // "primary key columns for {table_id} did not match expected {:?}", - // expected_columns - // ); - // } - // wait_for_primary_key_columns(&bigquery_database, &sequenced_table_id, &["id"]).await; - // We perform schema changes. database .alter_table( @@ -601,34 +565,36 @@ async fn relation_event_primary_key_syncs_in_bigquery() { .await .unwrap(); - // Register notifications for the insert. - let insert_event_notify = destination - .wait_for_events_count(vec![(EventType::Insert, 1)]) - .await; - - // We insert data. - database - .insert_values( - database_schema.users_schema().name.clone(), - &["name", "new_age", "year"], - &[&"user_3", &(3i32), &(2025i32)], - ) - .await - .expect("Failed to insert users"); + // // Register notifications for the insert. + // let insert_event_notify = destination + // .wait_for_events_count(vec![(EventType::Insert, 1)]) + // .await; + // + // // We insert data. + // database + // .insert_values( + // database_schema.users_schema().name.clone(), + // &["name", "new_age", "year"], + // &[&"user_3", &(3i32), &(2025i32)], + // ) + // .await + // .expect("Failed to insert users"); + // + // timeout(Duration::from_secs(2), insert_event_notify.notified()).await; - insert_event_notify.notified().await; + sleep(Duration::from_secs(2)).await; // We check the BigQuery data after the first schema change. let users_rows = bigquery_database .query_table(database_schema.users_schema().name) .await .unwrap(); - println!("USERS ROWS {:?}", users_rows); + println!("USERS {:?}", users_rows); // let parsed_users_rows = parse_bigquery_table_rows::(users_rows); // assert_eq!( // parsed_users_rows, // vec![ - // BigQueryUser::new(1, "user_2", 1), + // BigQueryUser::new(1, "user_1", 1), // BigQueryUser::new(2, "user_2", 2), // ] // ); diff --git a/etl-destinations/tests/iceberg_pipeline.rs b/etl-destinations/tests/iceberg_pipeline.rs index e3f20b441..c26266e91 100644 --- a/etl-destinations/tests/iceberg_pipeline.rs +++ b/etl-destinations/tests/iceberg_pipeline.rs @@ -3,7 +3,9 @@ use std::collections::HashMap; use chrono::{DateTime, NaiveDate, NaiveDateTime, NaiveTime, Utc}; -use etl::types::{ArrayCell, Cell, ColumnSchema, TableId, TableName, TableRow, VersionedTableSchema, Type}; +use etl::types::{ + ArrayCell, Cell, ColumnSchema, TableId, TableName, TableRow, Type, VersionedTableSchema, +}; use etl_destinations::iceberg::IcebergClient; use etl_telemetry::tracing::init_test_tracing; use iceberg::io::{S3_ACCESS_KEY_ID, S3_ENDPOINT, S3_SECRET_ACCESS_KEY}; diff --git a/etl-postgres/src/replication/schema.rs b/etl-postgres/src/replication/schema.rs index 023634950..ec514bcc5 100644 --- a/etl-postgres/src/replication/schema.rs +++ b/etl-postgres/src/replication/schema.rs @@ -5,7 +5,7 @@ use std::collections::HashMap; use tokio_postgres::types::Type as PgType; use crate::types::{ - ColumnSchema, SchemaVersion, TableId, TableName, VersionedTableSchema, TableSchema, + ColumnSchema, SchemaVersion, TableId, TableName, TableSchema, VersionedTableSchema, }; /// The initial schema version number. diff --git a/etl/src/replication/apply.rs b/etl/src/replication/apply.rs index 831a5a7a1..cf1f09a88 100644 --- a/etl/src/replication/apply.rs +++ b/etl/src/replication/apply.rs @@ -1,6 +1,6 @@ use etl_config::shared::PipelineConfig; use etl_postgres::replication::worker::WorkerType; -use etl_postgres::types::{TableId, VersionedTableSchema, TableSchema}; +use etl_postgres::types::{TableId, TableSchema, VersionedTableSchema}; use futures::StreamExt; use metrics::histogram; use postgres_replication::protocol; diff --git a/etl/src/store/both/memory.rs b/etl/src/store/both/memory.rs index 82cc95b35..b434b876b 100644 --- a/etl/src/store/both/memory.rs +++ b/etl/src/store/both/memory.rs @@ -1,4 +1,4 @@ -use etl_postgres::types::{SchemaVersion, TableId, VersionedTableSchema, TableSchema}; +use etl_postgres::types::{SchemaVersion, TableId, TableSchema, VersionedTableSchema}; use std::collections::{BTreeMap, HashMap}; use std::sync::Arc; use tokio::sync::Mutex; diff --git a/etl/src/store/both/postgres.rs b/etl/src/store/both/postgres.rs index ec7d43e93..030c6d2e3 100644 --- a/etl/src/store/both/postgres.rs +++ b/etl/src/store/both/postgres.rs @@ -1,6 +1,6 @@ use etl_config::shared::PgConnectionConfig; use etl_postgres::replication::{connect_to_source_database, schema, state, table_mappings}; -use etl_postgres::types::{SchemaVersion, TableId, VersionedTableSchema, TableSchema}; +use etl_postgres::types::{SchemaVersion, TableId, TableSchema, VersionedTableSchema}; use metrics::gauge; use sqlx::PgPool; use std::{ diff --git a/etl/src/store/schema/base.rs b/etl/src/store/schema/base.rs index 537292518..c14032fef 100644 --- a/etl/src/store/schema/base.rs +++ b/etl/src/store/schema/base.rs @@ -1,4 +1,4 @@ -use etl_postgres::types::{SchemaVersion, TableId, VersionedTableSchema, TableSchema}; +use etl_postgres::types::{SchemaVersion, TableId, TableSchema, VersionedTableSchema}; use std::sync::Arc; use crate::error::EtlResult; diff --git a/etl/src/test_utils/notify.rs b/etl/src/test_utils/notify.rs index 547f904a2..1efdac168 100644 --- a/etl/src/test_utils/notify.rs +++ b/etl/src/test_utils/notify.rs @@ -1,4 +1,4 @@ -use etl_postgres::types::{SchemaVersion, TableId, VersionedTableSchema, TableSchema}; +use etl_postgres::types::{SchemaVersion, TableId, TableSchema, VersionedTableSchema}; use std::{ collections::{BTreeMap, HashMap}, fmt, diff --git a/etl/src/test_utils/table.rs b/etl/src/test_utils/table.rs index 2a6781d59..d3a7fae54 100644 --- a/etl/src/test_utils/table.rs +++ b/etl/src/test_utils/table.rs @@ -1,4 +1,4 @@ -use etl_postgres::types::{ColumnSchema, TableId, TableName, VersionedTableSchema, TableSchema}; +use etl_postgres::types::{ColumnSchema, TableId, TableName, TableSchema, VersionedTableSchema}; use std::collections::HashMap; /// Return the names of the column schema. diff --git a/etl/src/test_utils/test_schema.rs b/etl/src/test_utils/test_schema.rs index 77b375d9a..fcce4bbe7 100644 --- a/etl/src/test_utils/test_schema.rs +++ b/etl/src/test_utils/test_schema.rs @@ -1,5 +1,5 @@ use etl_postgres::tokio::test_utils::{PgDatabase, id_column_schema}; -use etl_postgres::types::{ColumnSchema, TableId, TableName, VersionedTableSchema, TableSchema}; +use etl_postgres::types::{ColumnSchema, TableId, TableName, TableSchema, VersionedTableSchema}; use std::ops::RangeInclusive; use tokio_postgres::types::{PgLsn, Type}; use tokio_postgres::{Client, GenericClient}; From 6bee40988a93236874c9755609df1c951ebd231c Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Thu, 2 Oct 2025 09:14:52 -0700 Subject: [PATCH 27/45] Improve --- etl-destinations/src/bigquery/client.rs | 10 ++++++ etl-destinations/tests/bigquery_pipeline.rs | 34 ++++++++++----------- 2 files changed, 26 insertions(+), 18 deletions(-) diff --git a/etl-destinations/src/bigquery/client.rs b/etl-destinations/src/bigquery/client.rs index 9bc96e30d..77a9594bb 100644 --- a/etl-destinations/src/bigquery/client.rs +++ b/etl-destinations/src/bigquery/client.rs @@ -293,6 +293,8 @@ impl BigQueryClient { let _ = self.query(QueryRequest::new(query)).await?; + info!("added column {} to table {full_table_name} in BigQuery", column_schema.name); + Ok(()) } @@ -310,6 +312,8 @@ impl BigQueryClient { let _ = self.query(QueryRequest::new(query)).await?; + info!("dropped column {} from table {full_table_name} in BigQuery", column_name); + Ok(()) } @@ -329,6 +333,8 @@ impl BigQueryClient { let _ = self.query(QueryRequest::new(query)).await?; + info!("altered column {} in table {full_table_name} in BigQuery", column_schema.name); + Ok(()) } @@ -365,6 +371,8 @@ impl BigQueryClient { let _ = self.query(QueryRequest::new(query)).await?; + info!("synced primary key for table {full_table_name} in BigQuery"); + Ok(()) } @@ -399,6 +407,8 @@ impl BigQueryClient { let _ = self.query(QueryRequest::new(query)).await?; + info!("dropped primary key for table {full_table_name} in BigQuery"); + Ok(()) } diff --git a/etl-destinations/tests/bigquery_pipeline.rs b/etl-destinations/tests/bigquery_pipeline.rs index 689c1e860..2343c0891 100644 --- a/etl-destinations/tests/bigquery_pipeline.rs +++ b/etl-destinations/tests/bigquery_pipeline.rs @@ -565,24 +565,22 @@ async fn relation_event_primary_key_syncs_in_bigquery() { .await .unwrap(); - // // Register notifications for the insert. - // let insert_event_notify = destination - // .wait_for_events_count(vec![(EventType::Insert, 1)]) - // .await; - // - // // We insert data. - // database - // .insert_values( - // database_schema.users_schema().name.clone(), - // &["name", "new_age", "year"], - // &[&"user_3", &(3i32), &(2025i32)], - // ) - // .await - // .expect("Failed to insert users"); - // - // timeout(Duration::from_secs(2), insert_event_notify.notified()).await; - - sleep(Duration::from_secs(2)).await; + // Register notifications for the insert. + let insert_event_notify = destination + .wait_for_events_count(vec![(EventType::Insert, 1)]) + .await; + + // We insert data. + database + .insert_values( + database_schema.users_schema().name.clone(), + &["name", "new_age", "year"], + &[&"user_3", &(3i32), &(2025i32)], + ) + .await + .expect("Failed to insert users"); + + timeout(Duration::from_secs(2), insert_event_notify.notified()).await; // We check the BigQuery data after the first schema change. let users_rows = bigquery_database From 97d5c9bc822de81f54ec50736a750a49bf81eb43 Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Thu, 2 Oct 2025 09:20:57 -0700 Subject: [PATCH 28/45] Improve --- etl-destinations/src/bigquery/client.rs | 2 +- etl-destinations/src/bigquery/core.rs | 4 ++-- etl/tests/pipeline.rs | 4 ++-- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/etl-destinations/src/bigquery/client.rs b/etl-destinations/src/bigquery/client.rs index 77a9594bb..33d1359cd 100644 --- a/etl-destinations/src/bigquery/client.rs +++ b/etl-destinations/src/bigquery/client.rs @@ -291,7 +291,7 @@ impl BigQueryClient { let query = format!("alter table {full_table_name} add column if not exists {column_definition}"); - let _ = self.query(QueryRequest::new(query)).await?; + let x = self.query(QueryRequest::new(query)).await?; info!("added column {} to table {full_table_name} in BigQuery", column_schema.name); diff --git a/etl-destinations/src/bigquery/core.rs b/etl-destinations/src/bigquery/core.rs index 30b062301..7af5bca44 100644 --- a/etl-destinations/src/bigquery/core.rs +++ b/etl-destinations/src/bigquery/core.rs @@ -698,8 +698,8 @@ where ); }; - let rows = mem::take(table_id_to_table_rows); - self.process_table_events(batch_schema_version, rows).await + let table_id_to_table_rows = mem::take(table_id_to_table_rows); + self.process_table_events(batch_schema_version, table_id_to_table_rows).await } /// Processes CDC events in batches with proper ordering and truncate handling. diff --git a/etl/tests/pipeline.rs b/etl/tests/pipeline.rs index db61edc50..25474acf6 100644 --- a/etl/tests/pipeline.rs +++ b/etl/tests/pipeline.rs @@ -88,13 +88,13 @@ async fn table_schema_copy_survives_pipeline_restarts() { *table_schemas .get(&database_schema.users_schema().id) .unwrap(), - database_schema.users_schema() + database_schema.users_schema().into_versioned(0) ); assert_eq!( *table_schemas .get(&database_schema.orders_schema().id) .unwrap(), - database_schema.orders_schema() + database_schema.orders_schema().into_versioned(0) ); // We recreate a pipeline, assuming the other one was stopped, using the same state and destination. From dd471609f279e114f3f59a4a16a94ea98d84810f Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Thu, 2 Oct 2025 09:54:16 -0700 Subject: [PATCH 29/45] Improve --- etl-destinations/src/bigquery/client.rs | 2 +- etl-destinations/tests/support/bigquery.rs | 22 +++++++++++----------- 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/etl-destinations/src/bigquery/client.rs b/etl-destinations/src/bigquery/client.rs index 33d1359cd..c93de5648 100644 --- a/etl-destinations/src/bigquery/client.rs +++ b/etl-destinations/src/bigquery/client.rs @@ -333,7 +333,7 @@ impl BigQueryClient { let _ = self.query(QueryRequest::new(query)).await?; - info!("altered column {} in table {full_table_name} in BigQuery", column_schema.name); + info!("altered column type {} in table {full_table_name} in BigQuery", column_schema.name); Ok(()) } diff --git a/etl-destinations/tests/support/bigquery.rs b/etl-destinations/tests/support/bigquery.rs index e586a716e..d3c858549 100644 --- a/etl-destinations/tests/support/bigquery.rs +++ b/etl-destinations/tests/support/bigquery.rs @@ -204,17 +204,17 @@ impl Drop for BigQueryDatabase { fn drop(&mut self) { // We take out the client since during destruction we know that the struct won't be used // anymore. - let Some(client) = self.take_client() else { - return; - }; - - // To use `block_in_place,` we need a multithreaded runtime since when a blocking - // task is issued, the runtime will offload existing tasks to another worker. - tokio::task::block_in_place(move || { - Handle::current().block_on(async move { - destroy_bigquery(&client, self.project_id(), self.dataset_id()).await; - }); - }); + // let Some(client) = self.take_client() else { + // return; + // }; + // + // // To use `block_in_place,` we need a multithreaded runtime since when a blocking + // // task is issued, the runtime will offload existing tasks to another worker. + // tokio::task::block_in_place(move || { + // Handle::current().block_on(async move { + // destroy_bigquery(&client, self.project_id(), self.dataset_id()).await; + // }); + // }); } } From bfd08b3e79ef5b9793f75a337ad5085b06e73b0e Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Thu, 2 Oct 2025 16:59:29 -0700 Subject: [PATCH 30/45] Improve --- Cargo.toml | 2 +- etl-destinations/src/bigquery/client.rs | 66 +++++++++++++------ etl-destinations/src/bigquery/core.rs | 73 ++++++++++++++++++--- etl-destinations/tests/bigquery_pipeline.rs | 42 +++++++++++- etl/src/error.rs | 1 + etl/src/test_utils/test_schema.rs | 2 +- etl/tests/pipeline.rs | 17 +++++ 7 files changed, 170 insertions(+), 33 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index b5a0240b7..d7242f94b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -45,7 +45,7 @@ const-oid = { version = "0.9.6", default-features = false } constant_time_eq = { version = "0.4.2" } fail = { version = "0.5.1", default-features = false } futures = { version = "0.3.31", default-features = false } -gcp-bigquery-client = { version = "0.27.0", default-features = false } +gcp-bigquery-client = { path = "../gcp-bigquery-client", default-features = false } iceberg = { version = "0.6.0", default-features = false } iceberg-catalog-rest = { version = "0.6.0", default-features = false } insta = { version = "1.43.1", default-features = false } diff --git a/etl-destinations/src/bigquery/client.rs b/etl-destinations/src/bigquery/client.rs index c93de5648..e38106f35 100644 --- a/etl-destinations/src/bigquery/client.rs +++ b/etl-destinations/src/bigquery/client.rs @@ -291,9 +291,12 @@ impl BigQueryClient { let query = format!("alter table {full_table_name} add column if not exists {column_definition}"); - let x = self.query(QueryRequest::new(query)).await?; + let _ = self.query(QueryRequest::new(query)).await?; - info!("added column {} to table {full_table_name} in BigQuery", column_schema.name); + info!( + "added column {} to table {full_table_name} in BigQuery", + column_schema.name + ); Ok(()) } @@ -312,7 +315,10 @@ impl BigQueryClient { let _ = self.query(QueryRequest::new(query)).await?; - info!("dropped column {} from table {full_table_name} in BigQuery", column_name); + info!( + "dropped column {} from table {full_table_name} in BigQuery", + column_name + ); Ok(()) } @@ -333,7 +339,10 @@ impl BigQueryClient { let _ = self.query(QueryRequest::new(query)).await?; - info!("altered column type {} in table {full_table_name} in BigQuery", column_schema.name); + info!( + "altered column type {} in table {full_table_name} in BigQuery", + column_schema.name + ); Ok(()) } @@ -444,16 +453,11 @@ impl BigQueryClient { /// in a single batch. pub async fn stream_table_batches_concurrent( &self, - table_batches: Vec>, + table_batches: Arc<[TableBatch]>, max_concurrent_streams: usize, ) -> EtlResult<(usize, usize)> { - if table_batches.is_empty() { - return Ok((0, 0)); - } - debug!( - "streaming {:?} table batches concurrently with maximum {:?} concurrent streams", - table_batches.len(), + "streaming table batches concurrently with maximum {:?} concurrent streams", max_concurrent_streams ); @@ -526,11 +530,11 @@ impl BigQueryClient { // We want to use the default stream from BigQuery since it allows multiple connections to // send data to it. In addition, it's available by default for every table, so it also reduces // complexity. - let stream_name = StreamName::new_default( + let stream_name = Arc::new(StreamName::new_default( self.project_id.clone(), dataset_id.to_string(), table_id.to_string(), - ); + )); Ok(TableBatch::new( stream_name, @@ -757,6 +761,16 @@ impl fmt::Debug for BigQueryClient { fn bq_error_to_etl_error(err: BQError) -> EtlError { use BQError; + let error_message = err.to_string(); + + if is_schema_mismatch_message(&error_message) { + return etl_error!( + ErrorKind::DestinationSchemaMismatch, + "BigQuery schema mismatch error", + error_message + ); + } + let (kind, description) = match &err { // Authentication related errors BQError::InvalidServiceAccountKey(_) => ( @@ -852,16 +866,30 @@ fn bq_error_to_etl_error(err: BQError) -> EtlError { ), }; - etl_error!(kind, description, err.to_string()) + etl_error!(kind, description, error_message) } /// Converts BigQuery row errors to ETL destination errors. fn row_error_to_etl_error(err: RowError) -> EtlError { - etl_error!( - ErrorKind::DestinationError, - "BigQuery row error", - format!("{err:?}") - ) + let detail = format!("{err:?}"); + + if is_schema_mismatch_message(&detail) { + return etl_error!( + ErrorKind::DestinationSchemaMismatch, + "BigQuery schema mismatch error", + detail + ); + } + + etl_error!(ErrorKind::DestinationError, "BigQuery row error", detail) +} + +/// Returns `true` when the provided message indicates a BigQuery schema mismatch. +fn is_schema_mismatch_message(message: &str) -> bool { + let lowercase = message.to_ascii_lowercase(); + + lowercase.contains("input schema has more fields than bigquery schema") + || lowercase.contains("schema_mismatch_extra_field") } #[cfg(test)] diff --git a/etl-destinations/src/bigquery/core.rs b/etl-destinations/src/bigquery/core.rs index 7af5bca44..815d9cc64 100644 --- a/etl-destinations/src/bigquery/core.rs +++ b/etl-destinations/src/bigquery/core.rs @@ -7,23 +7,30 @@ use etl::types::{ VersionedTableSchema, }; use etl::{bail, etl_error}; -use gcp_bigquery_client::storage::TableDescriptor; +use gcp_bigquery_client::storage::{TableBatch, TableDescriptor}; use std::collections::{HashMap, HashSet}; use std::fmt::Display; use std::iter; use std::mem; use std::str::FromStr; use std::sync::Arc; +use std::time::Duration; use tokio::sync::Mutex; +use tokio::time::sleep; use tracing::{debug, info, warn}; use crate::bigquery::client::{BigQueryClient, BigQueryOperationType}; +use crate::bigquery::encoding::BigQueryTableRow; use crate::bigquery::{BigQueryDatasetId, BigQueryTableId}; /// Delimiter separating schema from table name in BigQuery table identifiers. const BIGQUERY_TABLE_ID_DELIMITER: &str = "_"; /// Replacement string for escaping underscores in Postgres names. const BIGQUERY_TABLE_ID_DELIMITER_ESCAPE_REPLACEMENT: &str = "__"; +/// Maximum number of BigQuery streaming attempts when schema propagation lags behind. +const MAX_SCHEMA_MISMATCH_ATTEMPTS: usize = 5; +/// Delay in milliseconds between retry attempts triggered by BigQuery schema mismatches. +const SCHEMA_MISMATCH_RETRY_DELAY_MS: u64 = 500; /// Creates a hex-encoded sequence number from Postgres LSNs to ensure correct event ordering. /// @@ -472,10 +479,7 @@ where // Stream all the batches concurrently. if !table_batches.is_empty() { - let (bytes_sent, bytes_received) = self - .client - .stream_table_batches_concurrent(table_batches, self.max_concurrent_streams) - .await?; + let (bytes_sent, bytes_received) = self.stream_with_schema_retry(table_batches).await?; // Logs with egress_metric = true can be used to identify egress logs. // This can e.g. be used to send egress logs to a location different @@ -522,10 +526,7 @@ where return Ok(()); } - let (bytes_sent, bytes_received) = self - .client - .stream_table_batches_concurrent(table_batches, self.max_concurrent_streams) - .await?; + let (bytes_sent, bytes_received) = self.stream_with_schema_retry(table_batches).await?; // Logs with egress_metric = true can be used to identify egress logs. info!( @@ -539,6 +540,57 @@ where Ok(()) } + /// Streams table batches to BigQuery, retrying when schema mismatch errors occur. + /// + /// The rationale is that per BigQuery docs, the Storage Write API detects schema changes after + /// a short time, on the order of minutes. + async fn stream_with_schema_retry(&self, table_batches: T) -> EtlResult<(usize, usize)> + where + T: Into]>>, + { + let retry_delay = Duration::from_millis(SCHEMA_MISMATCH_RETRY_DELAY_MS); + let mut attempts = 0; + + let table_batches = table_batches.into(); + loop { + match self + .client + .stream_table_batches_concurrent(table_batches.clone(), self.max_concurrent_streams) + .await + { + Ok(result) => return Ok(result), + Err(error) => { + if !Self::is_schema_mismatch_error(&error) { + return Err(error); + } + + attempts += 1; + + if attempts >= MAX_SCHEMA_MISMATCH_ATTEMPTS { + return Err(error); + } + + warn!( + attempt = attempts, + max_attempts = MAX_SCHEMA_MISMATCH_ATTEMPTS, + error = %error, + "schema mismatch detected while streaming to BigQuery; retrying" + ); + + sleep(retry_delay).await; + } + } + } + } + + /// Returns `true` when the error or one of its aggregated errors indicates a schema mismatch. + fn is_schema_mismatch_error(error: &EtlError) -> bool { + error + .kinds() + .iter() + .any(|kind| *kind == ErrorKind::DestinationSchemaMismatch) + } + async fn apply_relation_event_changes(&self, relation_event: RelationEvent) -> EtlResult<()> { // We build the list of changes for this relation event, this way, we can express the event // in terms of a minimal set of operations to apply on the destination table schema. @@ -699,7 +751,8 @@ where }; let table_id_to_table_rows = mem::take(table_id_to_table_rows); - self.process_table_events(batch_schema_version, table_id_to_table_rows).await + self.process_table_events(batch_schema_version, table_id_to_table_rows) + .await } /// Processes CDC events in batches with proper ordering and truncate handling. diff --git a/etl-destinations/tests/bigquery_pipeline.rs b/etl-destinations/tests/bigquery_pipeline.rs index 2343c0891..e752b59b4 100644 --- a/etl-destinations/tests/bigquery_pipeline.rs +++ b/etl-destinations/tests/bigquery_pipeline.rs @@ -580,15 +580,15 @@ async fn relation_event_primary_key_syncs_in_bigquery() { .await .expect("Failed to insert users"); - timeout(Duration::from_secs(2), insert_event_notify.notified()).await; + insert_event_notify.notified().await; // We check the BigQuery data after the first schema change. let users_rows = bigquery_database .query_table(database_schema.users_schema().name) .await .unwrap(); - println!("USERS {:?}", users_rows); // let parsed_users_rows = parse_bigquery_table_rows::(users_rows); + // println!("{:#?}", parsed_users_rows); // assert_eq!( // parsed_users_rows, // vec![ @@ -597,6 +597,44 @@ async fn relation_event_primary_key_syncs_in_bigquery() { // ] // ); + // We perform schema changes. + database + .alter_table( + test_table_name("users"), + &[ + TableModification::DropColumn { name: "year" }, + TableModification::AlterColumn { + name: "new_age", + params: "type double precision using new_age::double precision", + }, + ], + ) + .await + .unwrap(); + + // Register notifications for the insert. + let insert_event_notify = destination + .wait_for_events_count(vec![(EventType::Insert, 2)]) + .await; + + // We insert data. + database + .insert_values( + database_schema.users_schema().name.clone(), + &["name", "new_age"], + &[&"user_3", &(2f64)], + ) + .await + .expect("Failed to insert users"); + + timeout(Duration::from_secs(2), insert_event_notify.notified()).await; + + let users_rows = bigquery_database + .query_table(database_schema.users_schema().name) + .await + .unwrap(); + println!("users_rows: {:#?}", users_rows); + pipeline.shutdown_and_wait().await.unwrap(); } diff --git a/etl/src/error.rs b/etl/src/error.rs index 02db592ae..46c8bff7c 100644 --- a/etl/src/error.rs +++ b/etl/src/error.rs @@ -108,6 +108,7 @@ pub enum ErrorKind { // General Errors SourceError, DestinationError, + DestinationSchemaMismatch, // Unknown / Uncategorized Unknown, diff --git a/etl/src/test_utils/test_schema.rs b/etl/src/test_utils/test_schema.rs index fcce4bbe7..ea53dacc6 100644 --- a/etl/src/test_utils/test_schema.rs +++ b/etl/src/test_utils/test_schema.rs @@ -1,5 +1,5 @@ use etl_postgres::tokio::test_utils::{PgDatabase, id_column_schema}; -use etl_postgres::types::{ColumnSchema, TableId, TableName, TableSchema, VersionedTableSchema}; +use etl_postgres::types::{ColumnSchema, TableId, TableName, TableSchema}; use std::ops::RangeInclusive; use tokio_postgres::types::{PgLsn, Type}; use tokio_postgres::{Client, GenericClient}; diff --git a/etl/tests/pipeline.rs b/etl/tests/pipeline.rs index 25474acf6..ef259e4d7 100644 --- a/etl/tests/pipeline.rs +++ b/etl/tests/pipeline.rs @@ -128,6 +128,23 @@ async fn table_schema_copy_survives_pipeline_restarts() { pipeline.shutdown_and_wait().await.unwrap(); + // We want to make sure that during a restart, a relation message with the same schema doesn't + // cause a new schema version to be created. + let table_schemas = store.get_latest_table_schemas().await; + assert_eq!(table_schemas.len(), 2); + assert_eq!( + *table_schemas + .get(&database_schema.users_schema().id) + .unwrap(), + database_schema.users_schema().into_versioned(0) + ); + assert_eq!( + *table_schemas + .get(&database_schema.orders_schema().id) + .unwrap(), + database_schema.orders_schema().into_versioned(0) + ); + // We check that both inserts were received, and we know that we can receive them only when the table // schemas are available. let events = destination.get_events().await; From fff89cd4037721b13a8abbf931a13d051d79f309 Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Mon, 6 Oct 2025 09:26:32 -0700 Subject: [PATCH 31/45] Improve --- etl-destinations/src/bigquery/core.rs | 114 +++++++++++--------------- etl/src/destination/memory.rs | 2 +- etl/src/types/event.rs | 2 +- 3 files changed, 48 insertions(+), 70 deletions(-) diff --git a/etl-destinations/src/bigquery/core.rs b/etl-destinations/src/bigquery/core.rs index 815d9cc64..d2c163755 100644 --- a/etl-destinations/src/bigquery/core.rs +++ b/etl-destinations/src/bigquery/core.rs @@ -500,15 +500,18 @@ where /// Persists the accumulated CDC batches for each table to BigQuery. async fn process_table_events( &self, - schema_version: SchemaVersion, - table_id_to_table_rows: HashMap>, + table_batches_by_id: HashMap<(TableId, SchemaVersion), Vec>, ) -> EtlResult<()> { - if table_id_to_table_rows.is_empty() { + if table_batches_by_id.is_empty() { return Ok(()); } - let mut table_batches = Vec::with_capacity(table_id_to_table_rows.len()); - for (table_id, table_rows) in table_id_to_table_rows { + let mut table_batches = Vec::with_capacity(table_batches_by_id.len()); + for ((table_id, schema_version), table_rows) in table_batches_by_id { + if table_rows.is_empty() { + continue; + } + let (sequenced_bigquery_table_id, table_descriptor) = self .prepare_table_for_streaming(&table_id, Some(schema_version), true) .await?; @@ -688,8 +691,7 @@ where #[allow(clippy::too_many_arguments)] #[inline] fn push_dml_statement( - table_id_to_table_rows: &mut HashMap>, - batch_schema_version: &mut Option, + table_batches: &mut HashMap<(TableId, SchemaVersion), Vec>, start_lsn: PgLsn, commit_lsn: PgLsn, table_id: TableId, @@ -697,62 +699,50 @@ where schema_version: SchemaVersion, operation_type: BigQueryOperationType, ) -> EtlResult<()> { - // BigQuery CDC extras. + // BigQuery CDC extra fields. let sequence_number = generate_sequence_number(start_lsn, commit_lsn); table_row.values.push(operation_type.into_cell()); table_row.values.push(Cell::String(sequence_number)); - // Preserve per-table ordering. - table_id_to_table_rows - .entry(table_id) - .or_default() - .push(table_row); + let key = (table_id, schema_version); + if let Some(rows) = table_batches.get_mut(&key) { + rows.push(table_row); + return Ok(()); + } - // Ensure a single schema version per batch. - // - // We need to do this since we don't make any assumptions on relation events being there - // so we use the schema version of the first element that we find. - // - // The invariant that must be upheld is that for all events in a batch, they must all have - // the same schema version. - match batch_schema_version { - Some(batch_schema_version) => { - if schema_version != *batch_schema_version { - bail!( - ErrorKind::InvalidState, - "Multiple schema versions in the same batch", - "Multiple schema versions in the same batch were found while processing events for BigQuery" - ) - } - } - None => { - *batch_schema_version = Some(schema_version); - } + // TODO: maybe we want to remove this check and always postively assume that it works. + // Preserve per-table ordering and enforce a consistent schema version per table. + if table_batches + .keys() + .any(|(existing_table_id, existing_schema_version)| { + *existing_table_id == table_id && *existing_schema_version != schema_version + }) + { + bail!( + ErrorKind::InvalidState, + "Multiple schema versions for table in batch", + format!( + "Encountered schema version {schema_version} after a different version for table {table_id}" + ) + ); } + table_batches.insert(key, vec![table_row]); + Ok(()) } + /// Flushes the batch of events. async fn flush_batch( &self, - batch_schema_version: &mut Option, - table_id_to_table_rows: &mut HashMap>, + table_batches_by_id: &mut HashMap<(TableId, SchemaVersion), Vec>, ) -> EtlResult<()> { - if table_id_to_table_rows.is_empty() { + if table_batches_by_id.is_empty() { return Ok(()); } - let Some(batch_schema_version) = batch_schema_version.take() else { - bail!( - ErrorKind::InvalidState, - "Missing schema version", - "Missing schema version before writing events in BigQuery" - ); - }; - - let table_id_to_table_rows = mem::take(table_id_to_table_rows); - self.process_table_events(batch_schema_version, table_id_to_table_rows) - .await + let table_batches_by_id = mem::take(table_batches_by_id); + self.process_table_events(table_batches_by_id).await } /// Processes CDC events in batches with proper ordering and truncate handling. @@ -761,10 +751,8 @@ where /// then handles truncate events separately by creating new versioned tables. pub async fn write_events(&self, events: Vec) -> EtlResult<()> { // Accumulates rows for the current batch, grouped by table. - let mut table_id_to_table_rows: HashMap> = HashMap::new(); - - // The schema version used for the *current* batch (must be consistent within a batch). - let mut batch_schema_version: Option = None; + let mut table_batches_by_id: HashMap<(TableId, SchemaVersion), Vec> = + HashMap::new(); // Process stream. for event in events { @@ -772,8 +760,7 @@ where // DML events. Event::Insert(insert) => { Self::push_dml_statement( - &mut table_id_to_table_rows, - &mut batch_schema_version, + &mut table_batches_by_id, insert.start_lsn, insert.commit_lsn, insert.table_id, @@ -784,8 +771,7 @@ where } Event::Update(update) => { Self::push_dml_statement( - &mut table_id_to_table_rows, - &mut batch_schema_version, + &mut table_batches_by_id, update.start_lsn, update.commit_lsn, update.table_id, @@ -797,8 +783,7 @@ where Event::Delete(delete) => { if let Some((_, old_row)) = delete.old_table_row { Self::push_dml_statement( - &mut table_id_to_table_rows, - &mut batch_schema_version, + &mut table_batches_by_id, delete.start_lsn, delete.commit_lsn, delete.table_id, @@ -814,20 +799,14 @@ where // Batch breaker events. Event::Relation(relation) => { // Finish the current batch before applying schema change. - self.flush_batch(&mut batch_schema_version, &mut table_id_to_table_rows) - .await?; + self.flush_batch(&mut table_batches_by_id).await?; - // We mark the new batch schema version with the relation schema version, since - // after a relation message, a new schema is meant to be stored in the schema store. - batch_schema_version = Some(relation.new_table_schema.version); - - // Apply relation change, then prime the next batch with the new schema version. + // Apply relation change before processing subsequent DML. self.apply_relation_event_changes(relation).await?; } Event::Truncate(truncate) => { // Finish the current batch before a TRUNCATE (it affects the table state). - self.flush_batch(&mut batch_schema_version, &mut table_id_to_table_rows) - .await?; + self.flush_batch(&mut table_batches_by_id).await?; self.process_truncate_for_table_ids(truncate.table_ids.into_iter(), true) .await?; @@ -841,8 +820,7 @@ where } // Flush any trailing DML. - self.flush_batch(&mut batch_schema_version, &mut table_id_to_table_rows) - .await?; + self.flush_batch(&mut table_batches_by_id).await?; Ok(()) } diff --git a/etl/src/destination/memory.rs b/etl/src/destination/memory.rs index c7cd4529f..4c3ca9f48 100644 --- a/etl/src/destination/memory.rs +++ b/etl/src/destination/memory.rs @@ -127,7 +127,7 @@ impl Destination for MemoryDestination { for table_row in &table_rows { info!(" {:?}", table_row); } - inner.table_rows.insert(table_id, table_rows); + inner.table_rows.entry(table_id).or_default().extend(table_rows); Ok(()) } diff --git a/etl/src/types/event.rs b/etl/src/types/event.rs index 3e26d3908..519f91121 100644 --- a/etl/src/types/event.rs +++ b/etl/src/types/event.rs @@ -99,7 +99,7 @@ impl RelationEvent { let column_schema = column_schema.into_inner(); let latest_column_schema = latest_column_schema.into_inner(); - if column_schema.name != latest_column_schema.name { + if column_schema != latest_column_schema { // If we find a column with the same name but different fields, we assume it was changed. The only changes // that we detect are changes to the column but with preserved name. changes.push(RelationChange::AlterColumn( From 172ddc7fe7f7494030a25730c5907a90f0cd757a Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Mon, 6 Oct 2025 09:47:06 -0700 Subject: [PATCH 32/45] Improve --- etl/src/replication/apply.rs | 10 ++-- etl/tests/pipeline.rs | 111 ++++++++++++++++++++++++++++++----- 2 files changed, 100 insertions(+), 21 deletions(-) diff --git a/etl/src/replication/apply.rs b/etl/src/replication/apply.rs index cf1f09a88..3d2f7d29f 100644 --- a/etl/src/replication/apply.rs +++ b/etl/src/replication/apply.rs @@ -629,8 +629,8 @@ where if state.events_batch.len() >= max_batch_size || result.end_batch.is_some() { // We check if the batch has elements. It can be that a batch has no elements when // the batch is ended prematurely, and it contains only skipped events. In this case, - // we don't produce any events to the destination but downstream code treats it as if - // those evets are "persisted". + // we don't produce any events to the destination, but downstream code treats it as if + // those events are "persisted". if !state.events_batch.is_empty() { send_batch( state, @@ -648,13 +648,13 @@ where // be reprocessed, even the events before the failure will be skipped. // // Usually in the apply loop, errors are propagated upstream and handled based on if - // we are in a table sync worker or apply worker, however we have an edge case (for + // we are in a table sync worker or apply worker, however, we have an edge case (for // relation messages that change the schema) where we want to mark a table as errored // manually, not propagating the error outside the loop, which is going to be handled // differently based on the worker: // - Apply worker -> will continue the loop skipping the table. // - Table sync worker -> will stop the work (as if it had a normal uncaught error). - // Ideally we would get rid of this since it's an anomalous case which adds unnecessary + // Ideally, we would get rid of this since it's an anomalous case that adds unnecessary // complexity. if let Some(error) = result.table_replication_error { action = action.merge(hook.mark_table_errored(error).await?); @@ -690,7 +690,7 @@ where .await?; } - // We perform synchronization, to make sure that tables are synced. + // We perform synchronization to make sure that tables are synced. synchronize(state, hook).await } } diff --git a/etl/tests/pipeline.rs b/etl/tests/pipeline.rs index ef259e4d7..0250a8634 100644 --- a/etl/tests/pipeline.rs +++ b/etl/tests/pipeline.rs @@ -9,11 +9,7 @@ use etl::test_utils::notify::NotifyingStore; use etl::test_utils::pipeline::{create_pipeline, create_pipeline_with}; use etl::test_utils::table::column_schema_names; use etl::test_utils::test_destination_wrapper::TestDestinationWrapper; -use etl::test_utils::test_schema::{ - TableSelection, assert_events_equal, build_expected_orders_inserts, - build_expected_users_inserts, get_n_integers_sum, get_users_age_sum_from_rows, - insert_mock_data, insert_users_data, setup_test_database_schema, -}; +use etl::test_utils::test_schema::{TableSelection, assert_events_equal, build_expected_orders_inserts, build_expected_users_inserts, get_n_integers_sum, get_users_age_sum_from_rows, insert_mock_data, insert_users_data, setup_test_database_schema, insert_orders_data}; use etl::types::{EventType, PipelineId}; use etl_config::shared::BatchConfig; use etl_postgres::replication::slots::EtlReplicationSlot; @@ -511,10 +507,11 @@ async fn table_copy_replicates_existing_data() { async fn table_schema_changes_are_handled_correctly() { init_test_tracing(); let mut database = spawn_source_database().await; - let database_schema = setup_test_database_schema(&database, TableSelection::UsersOnly).await; + let database_schema = setup_test_database_schema(&database, TableSelection::Both).await; - // Insert initial users data. + // Insert initial users/orders data. insert_users_data(&mut database, &database_schema.users_schema().name, 1..=1).await; + insert_orders_data(&mut database, &database_schema.orders_schema().name, 1..=1).await; let store = NotifyingStore::new(); let destination = TestDestinationWrapper::wrap(MemoryDestination::new()); @@ -536,14 +533,21 @@ async fn table_schema_changes_are_handled_correctly() { TableReplicationPhaseType::SyncDone, ) .await; + let orders_state_notify = store + .notify_on_table_state_type( + database_schema.orders_schema().id, + TableReplicationPhaseType::SyncDone, + ) + .await; pipeline.start().await.unwrap(); users_state_notify.notified().await; + orders_state_notify.notified().await; // Check the initial schema. let table_schemas = store.get_latest_table_schemas().await; - assert_eq!(table_schemas.len(), 1); + assert_eq!(table_schemas.len(), 2); let users_table_schema = table_schemas .get(&database_schema.users_schema().id) .unwrap(); @@ -552,11 +556,21 @@ async fn table_schema_changes_are_handled_correctly() { column_schema_names(users_table_schema), vec!["id", "name", "age"] ); + let orders_table_schema = table_schemas + .get(&database_schema.orders_schema().id) + .unwrap(); + assert_eq!(orders_table_schema.version, 0); + assert_eq!( + column_schema_names(orders_table_schema), + vec!["id", "description"] + ); // Check the initial data. let table_rows = destination.get_table_rows().await; let users_table_rows = table_rows.get(&database_schema.users_schema().id).unwrap(); assert_eq!(users_table_rows.len(), 1); + let orders_table_rows = table_rows.get(&database_schema.orders_schema().id).unwrap(); + assert_eq!(orders_table_rows.len(), 1); // We perform schema changes. database @@ -575,10 +589,22 @@ async fn table_schema_changes_are_handled_correctly() { ) .await .unwrap(); + database + .alter_table( + test_table_name("orders"), + &[ + TableModification::AddColumn { + name: "summary", + params: "text", + }, + ], + ) + .await + .unwrap(); - // Register notifications for the insert. + // Register notifications for the inserts. let insert_event_notify = destination - .wait_for_events_count(vec![(EventType::Insert, 1)]) + .wait_for_events_count(vec![(EventType::Insert, 2)]) .await; // We insert data. @@ -589,13 +615,21 @@ async fn table_schema_changes_are_handled_correctly() { &[&"user_2", &(2i32), &(2025i32)], ) .await - .expect("Failed to insert users"); + .expect("Failed to insert user"); + database + .insert_values( + database_schema.orders_schema().name.clone(), + &["description", "summary"], + &[&"order_2", &"order_2_summary"], + ) + .await + .expect("Failed to insert order"); insert_event_notify.notified().await; // Check the updated schema. let table_schemas = store.get_latest_table_schemas().await; - assert_eq!(table_schemas.len(), 1); + assert_eq!(table_schemas.len(), 2); let users_table_schema = table_schemas .get(&database_schema.users_schema().id) .unwrap(); @@ -604,6 +638,14 @@ async fn table_schema_changes_are_handled_correctly() { column_schema_names(users_table_schema), vec!["id", "name", "new_age", "year"] ); + let orders_table_schema = table_schemas + .get(&database_schema.orders_schema().id) + .unwrap(); + assert_eq!(orders_table_schema.version, 1); + assert_eq!( + column_schema_names(orders_table_schema), + vec!["id", "description", "summary"] + ); // Check the updated data. let events = destination.get_events().await; @@ -612,6 +654,10 @@ async fn table_schema_changes_are_handled_correctly() { .get(&(EventType::Insert, database_schema.users_schema().id)) .unwrap(); assert_eq!(users_inserts.len(), 1); + let orders_inserts = grouped_events + .get(&(EventType::Insert, database_schema.orders_schema().id)) + .unwrap(); + assert_eq!(orders_inserts.len(), 1); // We perform schema changes. database @@ -627,10 +673,23 @@ async fn table_schema_changes_are_handled_correctly() { ) .await .unwrap(); + database + .alter_table( + test_table_name("orders"), + &[ + TableModification::DropColumn { name: "summary" }, + TableModification::RenameColumn { + name: "description", + new_name: "new_description", + }, + ], + ) + .await + .unwrap(); - // Register notifications for the insert. + // Register notifications for the inserts. let insert_event_notify = destination - .wait_for_events_count(vec![(EventType::Insert, 2)]) + .wait_for_events_count(vec![(EventType::Insert, 4)]) .await; // We insert data. @@ -641,13 +700,21 @@ async fn table_schema_changes_are_handled_correctly() { &[&"user_3", &(2f64)], ) .await - .expect("Failed to insert users"); + .expect("Failed to insert user"); + database + .insert_values( + database_schema.orders_schema().name.clone(), + &["new_description"], + &[&"order_3"], + ) + .await + .expect("Failed to insert order"); insert_event_notify.notified().await; // Check the updated schema. let table_schemas = store.get_latest_table_schemas().await; - assert_eq!(table_schemas.len(), 1); + assert_eq!(table_schemas.len(), 2); let users_table_schema = table_schemas .get(&database_schema.users_schema().id) .unwrap(); @@ -656,6 +723,14 @@ async fn table_schema_changes_are_handled_correctly() { column_schema_names(users_table_schema), vec!["id", "name", "new_age"] ); + let orders_table_schema = table_schemas + .get(&database_schema.orders_schema().id) + .unwrap(); + assert_eq!(orders_table_schema.version, 2); + assert_eq!( + column_schema_names(orders_table_schema), + vec!["id", "new_description"] + ); // Check the updated data. let events = destination.get_events().await; @@ -664,6 +739,10 @@ async fn table_schema_changes_are_handled_correctly() { .get(&(EventType::Insert, database_schema.users_schema().id)) .unwrap(); assert_eq!(users_inserts.len(), 2); + let orders_inserts = grouped_events + .get(&(EventType::Insert, database_schema.orders_schema().id)) + .unwrap(); + assert_eq!(orders_inserts.len(), 2); pipeline.shutdown_and_wait().await.unwrap(); } From 794f1488afe77b528a5769b701360bdf2f337b65 Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Mon, 6 Oct 2025 10:33:36 -0700 Subject: [PATCH 33/45] Improve --- etl-destinations/src/bigquery/client.rs | 3 +- etl-destinations/src/bigquery/core.rs | 3 +- etl-destinations/tests/bigquery_pipeline.rs | 32 ++++++++------------- etl-destinations/tests/support/bigquery.rs | 12 +++----- etl/src/destination/memory.rs | 6 +++- etl/tests/pipeline.rs | 16 ++++++----- 6 files changed, 33 insertions(+), 39 deletions(-) diff --git a/etl-destinations/src/bigquery/client.rs b/etl-destinations/src/bigquery/client.rs index e38106f35..8724c563e 100644 --- a/etl-destinations/src/bigquery/client.rs +++ b/etl-destinations/src/bigquery/client.rs @@ -348,6 +348,7 @@ impl BigQueryClient { } /// Synchronizes the primary key definition for a BigQuery table with the provided schema. + #[allow(dead_code)] pub async fn sync_primary_key( &self, dataset_id: &BigQueryDatasetId, @@ -539,7 +540,7 @@ impl BigQueryClient { Ok(TableBatch::new( stream_name, table_descriptor, - validated_rows, + validated_rows.into(), )) } diff --git a/etl-destinations/src/bigquery/core.rs b/etl-destinations/src/bigquery/core.rs index d2c163755..b63466cc6 100644 --- a/etl-destinations/src/bigquery/core.rs +++ b/etl-destinations/src/bigquery/core.rs @@ -590,8 +590,7 @@ where fn is_schema_mismatch_error(error: &EtlError) -> bool { error .kinds() - .iter() - .any(|kind| *kind == ErrorKind::DestinationSchemaMismatch) + .contains(&ErrorKind::DestinationSchemaMismatch) } async fn apply_relation_event_changes(&self, relation_event: RelationEvent) -> EtlResult<()> { diff --git a/etl-destinations/tests/bigquery_pipeline.rs b/etl-destinations/tests/bigquery_pipeline.rs index e752b59b4..3a83a4f2d 100644 --- a/etl-destinations/tests/bigquery_pipeline.rs +++ b/etl-destinations/tests/bigquery_pipeline.rs @@ -1,7 +1,7 @@ #![cfg(feature = "bigquery")] use crate::support::bigquery::{ - BigQueryDatabase, BigQueryOrder, BigQueryUser, BigQueryUserWithTenant, NonNullableColsScalar, + BigQueryOrder, BigQueryUser, NonNullableColsScalar, NullableColsArray, NullableColsScalar, parse_bigquery_table_rows, setup_bigquery_connection, }; use chrono::{DateTime, NaiveDate, NaiveDateTime, NaiveTime, Utc}; @@ -17,15 +17,14 @@ use etl::test_utils::test_schema::{ }; use etl::types::{EventType, PgNumeric, PipelineId}; use etl_destinations::bigquery::{ - install_crypto_provider_for_bigquery, table_name_to_bigquery_table_id, + install_crypto_provider_for_bigquery, }; use etl_postgres::tokio::test_utils::TableModification; use etl_telemetry::tracing::init_test_tracing; use rand::random; use std::str::FromStr; use std::time::Duration; -use tokio::time::{sleep, timeout}; -use tracing::debug; +use tokio::time::sleep; mod support; @@ -587,15 +586,7 @@ async fn relation_event_primary_key_syncs_in_bigquery() { .query_table(database_schema.users_schema().name) .await .unwrap(); - // let parsed_users_rows = parse_bigquery_table_rows::(users_rows); - // println!("{:#?}", parsed_users_rows); - // assert_eq!( - // parsed_users_rows, - // vec![ - // BigQueryUser::new(1, "user_1", 1), - // BigQueryUser::new(2, "user_2", 2), - // ] - // ); + assert_eq!(users_rows.len(), 3); // We perform schema changes. database @@ -603,10 +594,11 @@ async fn relation_event_primary_key_syncs_in_bigquery() { test_table_name("users"), &[ TableModification::DropColumn { name: "year" }, - TableModification::AlterColumn { - name: "new_age", - params: "type double precision using new_age::double precision", - }, + // TODO: data type change doesn't work rn. + // TableModification::AlterColumn { + // name: "new_age", + // params: "type double precision using new_age::double precision", + // }, ], ) .await @@ -622,18 +614,18 @@ async fn relation_event_primary_key_syncs_in_bigquery() { .insert_values( database_schema.users_schema().name.clone(), &["name", "new_age"], - &[&"user_3", &(2f64)], + &[&"user_3", &(2i32)], ) .await .expect("Failed to insert users"); - timeout(Duration::from_secs(2), insert_event_notify.notified()).await; + insert_event_notify.notified().await; let users_rows = bigquery_database .query_table(database_schema.users_schema().name) .await .unwrap(); - println!("users_rows: {:#?}", users_rows); + assert_eq!(users_rows.len(), 4); pipeline.shutdown_and_wait().await.unwrap(); } diff --git a/etl-destinations/tests/support/bigquery.rs b/etl-destinations/tests/support/bigquery.rs index d3c858549..c8855666f 100644 --- a/etl-destinations/tests/support/bigquery.rs +++ b/etl-destinations/tests/support/bigquery.rs @@ -16,7 +16,6 @@ use gcp_bigquery_client::model::table_cell::TableCell; use gcp_bigquery_client::model::table_row::TableRow; use std::fmt; use std::str::FromStr; -use tokio::runtime::Handle; use uuid::Uuid; /// Environment variable name for the BigQuery project id. @@ -265,12 +264,9 @@ impl BigQueryDatabase { let query = format!( "SELECT column_name \ - FROM `{project}.{dataset}.INFORMATION_SCHEMA.KEY_COLUMN_USAGE` \ - WHERE table_name = '{table}' AND constraint_name = 'PRIMARY KEY' \ - ORDER BY ordinal_position", - project = project_id, - dataset = dataset_id, - table = table_id + FROM `{project_id}.{dataset_id}.INFORMATION_SCHEMA.KEY_COLUMN_USAGE` \ + WHERE table_name = '{table_id}' AND constraint_name = 'PRIMARY KEY' \ + ORDER BY ordinal_position" ); let response = client @@ -284,7 +280,7 @@ impl BigQueryDatabase { .unwrap_or_default() .into_iter() .filter_map(|row| { - row.columns.and_then(|mut columns| { + row.columns.and_then(|columns| { columns .into_iter() .next() diff --git a/etl/src/destination/memory.rs b/etl/src/destination/memory.rs index 4c3ca9f48..fb918a00e 100644 --- a/etl/src/destination/memory.rs +++ b/etl/src/destination/memory.rs @@ -127,7 +127,11 @@ impl Destination for MemoryDestination { for table_row in &table_rows { info!(" {:?}", table_row); } - inner.table_rows.entry(table_id).or_default().extend(table_rows); + inner + .table_rows + .entry(table_id) + .or_default() + .extend(table_rows); Ok(()) } diff --git a/etl/tests/pipeline.rs b/etl/tests/pipeline.rs index 0250a8634..bcbfd2aa9 100644 --- a/etl/tests/pipeline.rs +++ b/etl/tests/pipeline.rs @@ -9,7 +9,11 @@ use etl::test_utils::notify::NotifyingStore; use etl::test_utils::pipeline::{create_pipeline, create_pipeline_with}; use etl::test_utils::table::column_schema_names; use etl::test_utils::test_destination_wrapper::TestDestinationWrapper; -use etl::test_utils::test_schema::{TableSelection, assert_events_equal, build_expected_orders_inserts, build_expected_users_inserts, get_n_integers_sum, get_users_age_sum_from_rows, insert_mock_data, insert_users_data, setup_test_database_schema, insert_orders_data}; +use etl::test_utils::test_schema::{ + TableSelection, assert_events_equal, build_expected_orders_inserts, + build_expected_users_inserts, get_n_integers_sum, get_users_age_sum_from_rows, + insert_mock_data, insert_orders_data, insert_users_data, setup_test_database_schema, +}; use etl::types::{EventType, PipelineId}; use etl_config::shared::BatchConfig; use etl_postgres::replication::slots::EtlReplicationSlot; @@ -592,12 +596,10 @@ async fn table_schema_changes_are_handled_correctly() { database .alter_table( test_table_name("orders"), - &[ - TableModification::AddColumn { - name: "summary", - params: "text", - }, - ], + &[TableModification::AddColumn { + name: "summary", + params: "text", + }], ) .await .unwrap(); From 2873396ce15ada96a7b76d43c52597982bc2efdc Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Mon, 6 Oct 2025 10:53:14 -0700 Subject: [PATCH 34/45] Improve --- etl-destinations/src/bigquery/core.rs | 3 +- etl-destinations/src/iceberg/destination.rs | 17 +- etl-destinations/tests/bigquery_pipeline.rs | 8 +- etl-destinations/tests/iceberg_client.rs | 529 ++++---------------- 4 files changed, 123 insertions(+), 434 deletions(-) diff --git a/etl-destinations/src/bigquery/core.rs b/etl-destinations/src/bigquery/core.rs index f3c6720ff..11208f638 100644 --- a/etl-destinations/src/bigquery/core.rs +++ b/etl-destinations/src/bigquery/core.rs @@ -4,9 +4,8 @@ use etl::store::schema::SchemaStore; use etl::store::state::StateStore; use etl::types::{ Cell, Event, PgLsn, RelationChange, RelationEvent, SchemaVersion, TableId, TableName, TableRow, - VersionedTableSchema, + VersionedTableSchema, generate_sequence_number, }; -use etl::types::{Cell, Event, TableId, TableName, TableRow, generate_sequence_number}; use etl::{bail, etl_error}; use gcp_bigquery_client::storage::{TableBatch, TableDescriptor}; use std::collections::{HashMap, HashSet}; diff --git a/etl-destinations/src/iceberg/destination.rs b/etl-destinations/src/iceberg/destination.rs index 808c99393..fee3877e0 100644 --- a/etl-destinations/src/iceberg/destination.rs +++ b/etl-destinations/src/iceberg/destination.rs @@ -4,21 +4,22 @@ use std::{ sync::Arc, }; -use crate::iceberg::IcebergClient; -use crate::iceberg::error::iceberg_error_to_etl_error; use etl::destination::Destination; use etl::error::{ErrorKind, EtlResult}; use etl::etl_error; use etl::store::schema::SchemaStore; use etl::store::state::StateStore; use etl::types::{ - Cell, ColumnSchema, Event, TableId, TableName, TableRow, TableSchema, Type, + Cell, ColumnSchema, Event, TableId, TableName, TableRow, Type, VersionedTableSchema, generate_sequence_number, }; use tokio::sync::Mutex; use tokio::task::JoinSet; use tracing::{debug, info}; +use crate::iceberg::IcebergClient; +use crate::iceberg::error::iceberg_error_to_etl_error; + /// CDC operation types for Iceberg changelog tables. /// /// Represents the type of change operation that occurred in the source database. @@ -307,8 +308,8 @@ where while let Some(Event::Truncate(_)) = event_iter.peek() { if let Some(Event::Truncate(truncate_event)) = event_iter.next() { - for table_id in truncate_event.rel_ids { - truncate_table_ids.insert(TableId::new(table_id)); + for (table_id, _schema_version) in truncate_event.table_ids { + truncate_table_ids.insert(table_id); } } } @@ -334,7 +335,7 @@ where let table_schema = self .store - .get_table_schema(&table_id) + .get_latest_table_schema(&table_id) .await? .ok_or_else(|| { etl_error!( @@ -344,7 +345,7 @@ where ) })?; - let table_schema = Self::add_cdc_columns(&table_schema); + let table_schema = Self::add_cdc_columns(table_schema.as_ref()); let iceberg_table_name = table_name_to_iceberg_table_name(&table_schema.name); let iceberg_table_name = self @@ -378,7 +379,7 @@ where /// /// These columns enable CDC consumers to understand the chronological order /// of changes and distinguish between different types of operations. - fn add_cdc_columns(table_schema: &TableSchema) -> TableSchema { + fn add_cdc_columns(table_schema: &VersionedTableSchema) -> VersionedTableSchema { let mut final_schema = table_schema.clone(); // Add cdc specific columns diff --git a/etl-destinations/tests/bigquery_pipeline.rs b/etl-destinations/tests/bigquery_pipeline.rs index 3a83a4f2d..bc5eb037a 100644 --- a/etl-destinations/tests/bigquery_pipeline.rs +++ b/etl-destinations/tests/bigquery_pipeline.rs @@ -1,8 +1,8 @@ #![cfg(feature = "bigquery")] use crate::support::bigquery::{ - BigQueryOrder, BigQueryUser, NonNullableColsScalar, - NullableColsArray, NullableColsScalar, parse_bigquery_table_rows, setup_bigquery_connection, + BigQueryOrder, BigQueryUser, NonNullableColsScalar, NullableColsArray, NullableColsScalar, + parse_bigquery_table_rows, setup_bigquery_connection, }; use chrono::{DateTime, NaiveDate, NaiveDateTime, NaiveTime, Utc}; use etl::config::BatchConfig; @@ -16,9 +16,7 @@ use etl::test_utils::test_schema::{ TableSelection, insert_mock_data, insert_users_data, setup_test_database_schema, }; use etl::types::{EventType, PgNumeric, PipelineId}; -use etl_destinations::bigquery::{ - install_crypto_provider_for_bigquery, -}; +use etl_destinations::bigquery::install_crypto_provider_for_bigquery; use etl_postgres::tokio::test_utils::TableModification; use etl_telemetry::tracing::init_test_tracing; use rand::random; diff --git a/etl-destinations/tests/iceberg_client.rs b/etl-destinations/tests/iceberg_client.rs index e716b3aa5..7de9a0974 100644 --- a/etl-destinations/tests/iceberg_client.rs +++ b/etl-destinations/tests/iceberg_client.rs @@ -105,198 +105,95 @@ async fn create_table_if_missing() { let table_name = "test_table".to_string(); let column_schemas = vec![ // Primary key - ColumnSchema::new("id".to_string(), Type::INT4, -1, false, true), + ColumnSchema::new("id".to_string(), Type::INT4, -1, true), // Boolean types - ColumnSchema::new("bool_col".to_string(), Type::BOOL, -1, true, false), + ColumnSchema::new("bool_col".to_string(), Type::BOOL, -1, false), // String types - ColumnSchema::new("char_col".to_string(), Type::CHAR, -1, true, false), - ColumnSchema::new("bpchar_col".to_string(), Type::BPCHAR, -1, true, false), - ColumnSchema::new("varchar_col".to_string(), Type::VARCHAR, -1, true, false), - ColumnSchema::new("name_col".to_string(), Type::NAME, -1, true, false), - ColumnSchema::new("text_col".to_string(), Type::TEXT, -1, true, false), + ColumnSchema::new("char_col".to_string(), Type::CHAR, -1, false), + ColumnSchema::new("bpchar_col".to_string(), Type::BPCHAR, -1, false), + ColumnSchema::new("varchar_col".to_string(), Type::VARCHAR, -1, false), + ColumnSchema::new("name_col".to_string(), Type::NAME, -1, false), + ColumnSchema::new("text_col".to_string(), Type::TEXT, -1, false), // Integer types - ColumnSchema::new("int2_col".to_string(), Type::INT2, -1, true, false), - ColumnSchema::new("int4_col".to_string(), Type::INT4, -1, true, false), - ColumnSchema::new("int8_col".to_string(), Type::INT8, -1, true, false), + ColumnSchema::new("int2_col".to_string(), Type::INT2, -1, false), + ColumnSchema::new("int4_col".to_string(), Type::INT4, -1, false), + ColumnSchema::new("int8_col".to_string(), Type::INT8, -1, false), // Float types - ColumnSchema::new("float4_col".to_string(), Type::FLOAT4, -1, true, false), - ColumnSchema::new("float8_col".to_string(), Type::FLOAT8, -1, true, false), + ColumnSchema::new("float4_col".to_string(), Type::FLOAT4, -1, false), + ColumnSchema::new("float8_col".to_string(), Type::FLOAT8, -1, false), // Numeric type - ColumnSchema::new("numeric_col".to_string(), Type::NUMERIC, -1, true, false), + ColumnSchema::new("numeric_col".to_string(), Type::NUMERIC, -1, false), // Date/Time types - ColumnSchema::new("date_col".to_string(), Type::DATE, -1, true, false), - ColumnSchema::new("time_col".to_string(), Type::TIME, -1, true, false), - ColumnSchema::new( - "timestamp_col".to_string(), - Type::TIMESTAMP, - -1, - true, - false, - ), - ColumnSchema::new( - "timestamptz_col".to_string(), - Type::TIMESTAMPTZ, - -1, - true, - false, - ), + ColumnSchema::new("date_col".to_string(), Type::DATE, -1, false), + ColumnSchema::new("time_col".to_string(), Type::TIME, -1, false), + ColumnSchema::new("timestamp_col".to_string(), Type::TIMESTAMP, -1, false), + ColumnSchema::new("timestamptz_col".to_string(), Type::TIMESTAMPTZ, -1, false), // UUID type - ColumnSchema::new("uuid_col".to_string(), Type::UUID, -1, true, false), + ColumnSchema::new("uuid_col".to_string(), Type::UUID, -1, false), // JSON types - ColumnSchema::new("json_col".to_string(), Type::JSON, -1, true, false), - ColumnSchema::new("jsonb_col".to_string(), Type::JSONB, -1, true, false), + ColumnSchema::new("json_col".to_string(), Type::JSON, -1, false), + ColumnSchema::new("jsonb_col".to_string(), Type::JSONB, -1, false), // OID type - ColumnSchema::new("oid_col".to_string(), Type::OID, -1, true, false), + ColumnSchema::new("oid_col".to_string(), Type::OID, -1, false), // Binary type - ColumnSchema::new("bytea_col".to_string(), Type::BYTEA, -1, true, false), + ColumnSchema::new("bytea_col".to_string(), Type::BYTEA, -1, false), // Array types - ColumnSchema::new( - "bool_array_col".to_string(), - Type::BOOL_ARRAY, - -1, - true, - false, - ), - ColumnSchema::new( - "char_array_col".to_string(), - Type::CHAR_ARRAY, - -1, - true, - false, - ), + ColumnSchema::new("bool_array_col".to_string(), Type::BOOL_ARRAY, -1, false), + ColumnSchema::new("char_array_col".to_string(), Type::CHAR_ARRAY, -1, false), ColumnSchema::new( "bpchar_array_col".to_string(), Type::BPCHAR_ARRAY, -1, - true, false, ), ColumnSchema::new( "varchar_array_col".to_string(), Type::VARCHAR_ARRAY, -1, - true, - false, - ), - ColumnSchema::new( - "name_array_col".to_string(), - Type::NAME_ARRAY, - -1, - true, - false, - ), - ColumnSchema::new( - "text_array_col".to_string(), - Type::TEXT_ARRAY, - -1, - true, - false, - ), - ColumnSchema::new( - "int2_array_col".to_string(), - Type::INT2_ARRAY, - -1, - true, - false, - ), - ColumnSchema::new( - "int4_array_col".to_string(), - Type::INT4_ARRAY, - -1, - true, - false, - ), - ColumnSchema::new( - "int8_array_col".to_string(), - Type::INT8_ARRAY, - -1, - true, false, ), + ColumnSchema::new("name_array_col".to_string(), Type::NAME_ARRAY, -1, false), + ColumnSchema::new("text_array_col".to_string(), Type::TEXT_ARRAY, -1, false), + ColumnSchema::new("int2_array_col".to_string(), Type::INT2_ARRAY, -1, false), + ColumnSchema::new("int4_array_col".to_string(), Type::INT4_ARRAY, -1, false), + ColumnSchema::new("int8_array_col".to_string(), Type::INT8_ARRAY, -1, false), ColumnSchema::new( "float4_array_col".to_string(), Type::FLOAT4_ARRAY, -1, - true, false, ), ColumnSchema::new( "float8_array_col".to_string(), Type::FLOAT8_ARRAY, -1, - true, false, ), ColumnSchema::new( "numeric_array_col".to_string(), Type::NUMERIC_ARRAY, -1, - true, - false, - ), - ColumnSchema::new( - "date_array_col".to_string(), - Type::DATE_ARRAY, - -1, - true, - false, - ), - ColumnSchema::new( - "time_array_col".to_string(), - Type::TIME_ARRAY, - -1, - true, false, ), + ColumnSchema::new("date_array_col".to_string(), Type::DATE_ARRAY, -1, false), + ColumnSchema::new("time_array_col".to_string(), Type::TIME_ARRAY, -1, false), ColumnSchema::new( "timestamp_array_col".to_string(), Type::TIMESTAMP_ARRAY, -1, - true, false, ), ColumnSchema::new( "timestamptz_array_col".to_string(), Type::TIMESTAMPTZ_ARRAY, -1, - true, - false, - ), - ColumnSchema::new( - "uuid_array_col".to_string(), - Type::UUID_ARRAY, - -1, - true, - false, - ), - ColumnSchema::new( - "json_array_col".to_string(), - Type::JSON_ARRAY, - -1, - true, - false, - ), - ColumnSchema::new( - "jsonb_array_col".to_string(), - Type::JSONB_ARRAY, - -1, - true, - false, - ), - ColumnSchema::new( - "oid_array_col".to_string(), - Type::OID_ARRAY, - -1, - true, - false, - ), - ColumnSchema::new( - "bytea_array_col".to_string(), - Type::BYTEA_ARRAY, - -1, - true, false, ), + ColumnSchema::new("uuid_array_col".to_string(), Type::UUID_ARRAY, -1, false), + ColumnSchema::new("json_array_col".to_string(), Type::JSON_ARRAY, -1, false), + ColumnSchema::new("jsonb_array_col".to_string(), Type::JSONB_ARRAY, -1, false), + ColumnSchema::new("oid_array_col".to_string(), Type::OID_ARRAY, -1, false), + ColumnSchema::new("bytea_array_col".to_string(), Type::BYTEA_ARRAY, -1, false), ]; // table doesn't exist yet @@ -364,50 +261,38 @@ async fn insert_nullable_scalars() { let table_name = "test_table".to_string(); let column_schemas = vec![ // Primary key - ColumnSchema::new("id".to_string(), Type::INT4, -1, false, true), + ColumnSchema::new("id".to_string(), Type::INT4, -1, true), // Boolean types - ColumnSchema::new("bool_col".to_string(), Type::BOOL, -1, true, false), + ColumnSchema::new("bool_col".to_string(), Type::BOOL, -1, false), // String types - ColumnSchema::new("char_col".to_string(), Type::CHAR, -1, true, false), - ColumnSchema::new("bpchar_col".to_string(), Type::BPCHAR, -1, true, false), - ColumnSchema::new("varchar_col".to_string(), Type::VARCHAR, -1, true, false), - ColumnSchema::new("name_col".to_string(), Type::NAME, -1, true, false), - ColumnSchema::new("text_col".to_string(), Type::TEXT, -1, true, false), + ColumnSchema::new("char_col".to_string(), Type::CHAR, -1, false), + ColumnSchema::new("bpchar_col".to_string(), Type::BPCHAR, -1, false), + ColumnSchema::new("varchar_col".to_string(), Type::VARCHAR, -1, false), + ColumnSchema::new("name_col".to_string(), Type::NAME, -1, false), + ColumnSchema::new("text_col".to_string(), Type::TEXT, -1, false), // Integer types - ColumnSchema::new("int2_col".to_string(), Type::INT2, -1, true, false), - ColumnSchema::new("int4_col".to_string(), Type::INT4, -1, true, false), - ColumnSchema::new("int8_col".to_string(), Type::INT8, -1, true, false), + ColumnSchema::new("int2_col".to_string(), Type::INT2, -1, false), + ColumnSchema::new("int4_col".to_string(), Type::INT4, -1, false), + ColumnSchema::new("int8_col".to_string(), Type::INT8, -1, false), // Float types - ColumnSchema::new("float4_col".to_string(), Type::FLOAT4, -1, true, false), - ColumnSchema::new("float8_col".to_string(), Type::FLOAT8, -1, true, false), + ColumnSchema::new("float4_col".to_string(), Type::FLOAT4, -1, false), + ColumnSchema::new("float8_col".to_string(), Type::FLOAT8, -1, false), // Numeric type - ColumnSchema::new("numeric_col".to_string(), Type::NUMERIC, -1, true, false), + ColumnSchema::new("numeric_col".to_string(), Type::NUMERIC, -1, false), // Date/Time types - ColumnSchema::new("date_col".to_string(), Type::DATE, -1, true, false), - ColumnSchema::new("time_col".to_string(), Type::TIME, -1, true, false), - ColumnSchema::new( - "timestamp_col".to_string(), - Type::TIMESTAMP, - -1, - true, - false, - ), - ColumnSchema::new( - "timestamptz_col".to_string(), - Type::TIMESTAMPTZ, - -1, - true, - false, - ), + ColumnSchema::new("date_col".to_string(), Type::DATE, -1, false), + ColumnSchema::new("time_col".to_string(), Type::TIME, -1, false), + ColumnSchema::new("timestamp_col".to_string(), Type::TIMESTAMP, -1, false), + ColumnSchema::new("timestamptz_col".to_string(), Type::TIMESTAMPTZ, -1, false), // UUID type - ColumnSchema::new("uuid_col".to_string(), Type::UUID, -1, true, false), + ColumnSchema::new("uuid_col".to_string(), Type::UUID, -1, false), // JSON types - ColumnSchema::new("json_col".to_string(), Type::JSON, -1, true, false), - ColumnSchema::new("jsonb_col".to_string(), Type::JSONB, -1, true, false), + ColumnSchema::new("json_col".to_string(), Type::JSON, -1, false), + ColumnSchema::new("jsonb_col".to_string(), Type::JSONB, -1, false), // OID type - ColumnSchema::new("oid_col".to_string(), Type::OID, -1, true, false), + ColumnSchema::new("oid_col".to_string(), Type::OID, -1, false), // Binary type - ColumnSchema::new("bytea_col".to_string(), Type::BYTEA, -1, true, false), + ColumnSchema::new("bytea_col".to_string(), Type::BYTEA, -1, false), ]; client @@ -528,50 +413,38 @@ async fn insert_non_nullable_scalars() { let table_name = "test_table".to_string(); let column_schemas = vec![ // Primary key - ColumnSchema::new("id".to_string(), Type::INT4, -1, false, true), + ColumnSchema::new("id".to_string(), Type::INT4, -1, true), // Boolean types - ColumnSchema::new("bool_col".to_string(), Type::BOOL, -1, false, false), + ColumnSchema::new("bool_col".to_string(), Type::BOOL, -1, false), // String types - ColumnSchema::new("char_col".to_string(), Type::CHAR, -1, false, false), - ColumnSchema::new("bpchar_col".to_string(), Type::BPCHAR, -1, false, false), - ColumnSchema::new("varchar_col".to_string(), Type::VARCHAR, -1, false, false), - ColumnSchema::new("name_col".to_string(), Type::NAME, -1, false, false), - ColumnSchema::new("text_col".to_string(), Type::TEXT, -1, false, false), + ColumnSchema::new("char_col".to_string(), Type::CHAR, -1, false), + ColumnSchema::new("bpchar_col".to_string(), Type::BPCHAR, -1, false), + ColumnSchema::new("varchar_col".to_string(), Type::VARCHAR, -1, false), + ColumnSchema::new("name_col".to_string(), Type::NAME, -1, false), + ColumnSchema::new("text_col".to_string(), Type::TEXT, -1, false), // Integer types - ColumnSchema::new("int2_col".to_string(), Type::INT2, -1, false, false), - ColumnSchema::new("int4_col".to_string(), Type::INT4, -1, false, false), - ColumnSchema::new("int8_col".to_string(), Type::INT8, -1, false, false), + ColumnSchema::new("int2_col".to_string(), Type::INT2, -1, false), + ColumnSchema::new("int4_col".to_string(), Type::INT4, -1, false), + ColumnSchema::new("int8_col".to_string(), Type::INT8, -1, false), // Float types - ColumnSchema::new("float4_col".to_string(), Type::FLOAT4, -1, false, false), - ColumnSchema::new("float8_col".to_string(), Type::FLOAT8, -1, false, false), + ColumnSchema::new("float4_col".to_string(), Type::FLOAT4, -1, false), + ColumnSchema::new("float8_col".to_string(), Type::FLOAT8, -1, false), // Numeric type - ColumnSchema::new("numeric_col".to_string(), Type::NUMERIC, -1, false, false), + ColumnSchema::new("numeric_col".to_string(), Type::NUMERIC, -1, false), // Date/Time types - ColumnSchema::new("date_col".to_string(), Type::DATE, -1, false, false), - ColumnSchema::new("time_col".to_string(), Type::TIME, -1, false, false), - ColumnSchema::new( - "timestamp_col".to_string(), - Type::TIMESTAMP, - -1, - false, - false, - ), - ColumnSchema::new( - "timestamptz_col".to_string(), - Type::TIMESTAMPTZ, - -1, - false, - false, - ), + ColumnSchema::new("date_col".to_string(), Type::DATE, -1, false), + ColumnSchema::new("time_col".to_string(), Type::TIME, -1, false), + ColumnSchema::new("timestamp_col".to_string(), Type::TIMESTAMP, -1, false), + ColumnSchema::new("timestamptz_col".to_string(), Type::TIMESTAMPTZ, -1, false), // UUID type - ColumnSchema::new("uuid_col".to_string(), Type::UUID, -1, false, false), + ColumnSchema::new("uuid_col".to_string(), Type::UUID, -1, false), // JSON types - ColumnSchema::new("json_col".to_string(), Type::JSON, -1, false, false), - ColumnSchema::new("jsonb_col".to_string(), Type::JSONB, -1, false, false), + ColumnSchema::new("json_col".to_string(), Type::JSON, -1, false), + ColumnSchema::new("jsonb_col".to_string(), Type::JSONB, -1, false), // OID type - ColumnSchema::new("oid_col".to_string(), Type::OID, -1, false, false), + ColumnSchema::new("oid_col".to_string(), Type::OID, -1, false), // Binary type - ColumnSchema::new("bytea_col".to_string(), Type::BYTEA, -1, false, false), + ColumnSchema::new("bytea_col".to_string(), Type::BYTEA, -1, false), ]; client @@ -666,86 +539,40 @@ async fn insert_nullable_array() { let table_name = "test_array_table".to_string(); let column_schemas = vec![ // Primary key - ColumnSchema::new("id".to_string(), Type::INT4, -1, false, true), + ColumnSchema::new("id".to_string(), Type::INT4, -1, true), // Boolean array type - ColumnSchema::new( - "bool_array_col".to_string(), - Type::BOOL_ARRAY, - -1, - true, - false, - ), + ColumnSchema::new("bool_array_col".to_string(), Type::BOOL_ARRAY, -1, false), // String array types - ColumnSchema::new( - "char_array_col".to_string(), - Type::CHAR_ARRAY, - -1, - true, - false, - ), + ColumnSchema::new("char_array_col".to_string(), Type::CHAR_ARRAY, -1, false), ColumnSchema::new( "bpchar_array_col".to_string(), Type::BPCHAR_ARRAY, -1, - true, false, ), ColumnSchema::new( "varchar_array_col".to_string(), Type::VARCHAR_ARRAY, -1, - true, - false, - ), - ColumnSchema::new( - "name_array_col".to_string(), - Type::NAME_ARRAY, - -1, - true, - false, - ), - ColumnSchema::new( - "text_array_col".to_string(), - Type::TEXT_ARRAY, - -1, - true, false, ), + ColumnSchema::new("name_array_col".to_string(), Type::NAME_ARRAY, -1, false), + ColumnSchema::new("text_array_col".to_string(), Type::TEXT_ARRAY, -1, false), // Integer array types - ColumnSchema::new( - "int2_array_col".to_string(), - Type::INT2_ARRAY, - -1, - true, - false, - ), - ColumnSchema::new( - "int4_array_col".to_string(), - Type::INT4_ARRAY, - -1, - true, - false, - ), - ColumnSchema::new( - "int8_array_col".to_string(), - Type::INT8_ARRAY, - -1, - true, - false, - ), + ColumnSchema::new("int2_array_col".to_string(), Type::INT2_ARRAY, -1, false), + ColumnSchema::new("int4_array_col".to_string(), Type::INT4_ARRAY, -1, false), + ColumnSchema::new("int8_array_col".to_string(), Type::INT8_ARRAY, -1, false), // Float array types ColumnSchema::new( "float4_array_col".to_string(), Type::FLOAT4_ARRAY, -1, - true, false, ), ColumnSchema::new( "float8_array_col".to_string(), Type::FLOAT8_ARRAY, -1, - true, false, ), // Numeric array type @@ -753,77 +580,32 @@ async fn insert_nullable_array() { "numeric_array_col".to_string(), Type::NUMERIC_ARRAY, -1, - true, false, ), // Date/Time array types - ColumnSchema::new( - "date_array_col".to_string(), - Type::DATE_ARRAY, - -1, - true, - false, - ), - ColumnSchema::new( - "time_array_col".to_string(), - Type::TIME_ARRAY, - -1, - true, - false, - ), + ColumnSchema::new("date_array_col".to_string(), Type::DATE_ARRAY, -1, false), + ColumnSchema::new("time_array_col".to_string(), Type::TIME_ARRAY, -1, false), ColumnSchema::new( "timestamp_array_col".to_string(), Type::TIMESTAMP_ARRAY, -1, - true, false, ), ColumnSchema::new( "timestamptz_array_col".to_string(), Type::TIMESTAMPTZ_ARRAY, -1, - true, false, ), // UUID array type - ColumnSchema::new( - "uuid_array_col".to_string(), - Type::UUID_ARRAY, - -1, - true, - false, - ), + ColumnSchema::new("uuid_array_col".to_string(), Type::UUID_ARRAY, -1, false), // JSON array types - ColumnSchema::new( - "json_array_col".to_string(), - Type::JSON_ARRAY, - -1, - true, - false, - ), - ColumnSchema::new( - "jsonb_array_col".to_string(), - Type::JSONB_ARRAY, - -1, - true, - false, - ), + ColumnSchema::new("json_array_col".to_string(), Type::JSON_ARRAY, -1, false), + ColumnSchema::new("jsonb_array_col".to_string(), Type::JSONB_ARRAY, -1, false), // OID array type - ColumnSchema::new( - "oid_array_col".to_string(), - Type::OID_ARRAY, - -1, - true, - false, - ), + ColumnSchema::new("oid_array_col".to_string(), Type::OID_ARRAY, -1, false), // Binary array type - ColumnSchema::new( - "bytea_array_col".to_string(), - Type::BYTEA_ARRAY, - -1, - true, - false, - ), + ColumnSchema::new("bytea_array_col".to_string(), Type::BYTEA_ARRAY, -1, false), ]; client @@ -1032,87 +814,41 @@ async fn insert_non_nullable_array() { let table_name = "test_non_nullable_array_table".to_string(); let column_schemas = vec![ // Primary key - ColumnSchema::new("id".to_string(), Type::INT4, -1, false, true), + ColumnSchema::new("id".to_string(), Type::INT4, -1, true), // Boolean array type - ColumnSchema::new( - "bool_array_col".to_string(), - Type::BOOL_ARRAY, - -1, - false, - false, - ), + ColumnSchema::new("bool_array_col".to_string(), Type::BOOL_ARRAY, -1, false), // String array types - ColumnSchema::new( - "char_array_col".to_string(), - Type::CHAR_ARRAY, - -1, - false, - false, - ), + ColumnSchema::new("char_array_col".to_string(), Type::CHAR_ARRAY, -1, false), ColumnSchema::new( "bpchar_array_col".to_string(), Type::BPCHAR_ARRAY, -1, false, - false, ), ColumnSchema::new( "varchar_array_col".to_string(), Type::VARCHAR_ARRAY, -1, false, - false, - ), - ColumnSchema::new( - "name_array_col".to_string(), - Type::NAME_ARRAY, - -1, - false, - false, - ), - ColumnSchema::new( - "text_array_col".to_string(), - Type::TEXT_ARRAY, - -1, - false, - false, ), + ColumnSchema::new("name_array_col".to_string(), Type::NAME_ARRAY, -1, false), + ColumnSchema::new("text_array_col".to_string(), Type::TEXT_ARRAY, -1, false), // Integer array types - ColumnSchema::new( - "int2_array_col".to_string(), - Type::INT2_ARRAY, - -1, - false, - false, - ), - ColumnSchema::new( - "int4_array_col".to_string(), - Type::INT4_ARRAY, - -1, - false, - false, - ), - ColumnSchema::new( - "int8_array_col".to_string(), - Type::INT8_ARRAY, - -1, - false, - false, - ), + ColumnSchema::new("int2_array_col".to_string(), Type::INT2_ARRAY, -1, false), + ColumnSchema::new("int4_array_col".to_string(), Type::INT4_ARRAY, -1, false), + ColumnSchema::new("int8_array_col".to_string(), Type::INT8_ARRAY, -1, false), // Float array types ColumnSchema::new( "float4_array_col".to_string(), Type::FLOAT4_ARRAY, -1, false, - false, ), ColumnSchema::new( "float8_array_col".to_string(), Type::FLOAT8_ARRAY, -1, false, - false, ), // Numeric array type ColumnSchema::new( @@ -1120,76 +856,31 @@ async fn insert_non_nullable_array() { Type::NUMERIC_ARRAY, -1, false, - false, ), // Date/Time array types - ColumnSchema::new( - "date_array_col".to_string(), - Type::DATE_ARRAY, - -1, - false, - false, - ), - ColumnSchema::new( - "time_array_col".to_string(), - Type::TIME_ARRAY, - -1, - false, - false, - ), + ColumnSchema::new("date_array_col".to_string(), Type::DATE_ARRAY, -1, false), + ColumnSchema::new("time_array_col".to_string(), Type::TIME_ARRAY, -1, false), ColumnSchema::new( "timestamp_array_col".to_string(), Type::TIMESTAMP_ARRAY, -1, false, - false, ), ColumnSchema::new( "timestamptz_array_col".to_string(), Type::TIMESTAMPTZ_ARRAY, -1, false, - false, ), // UUID array type - ColumnSchema::new( - "uuid_array_col".to_string(), - Type::UUID_ARRAY, - -1, - false, - false, - ), + ColumnSchema::new("uuid_array_col".to_string(), Type::UUID_ARRAY, -1, false), // JSON array types - ColumnSchema::new( - "json_array_col".to_string(), - Type::JSON_ARRAY, - -1, - false, - false, - ), - ColumnSchema::new( - "jsonb_array_col".to_string(), - Type::JSONB_ARRAY, - -1, - false, - false, - ), + ColumnSchema::new("json_array_col".to_string(), Type::JSON_ARRAY, -1, false), + ColumnSchema::new("jsonb_array_col".to_string(), Type::JSONB_ARRAY, -1, false), // OID array type - ColumnSchema::new( - "oid_array_col".to_string(), - Type::OID_ARRAY, - -1, - false, - false, - ), + ColumnSchema::new("oid_array_col".to_string(), Type::OID_ARRAY, -1, false), // Binary array type - ColumnSchema::new( - "bytea_array_col".to_string(), - Type::BYTEA_ARRAY, - -1, - false, - false, - ), + ColumnSchema::new("bytea_array_col".to_string(), Type::BYTEA_ARRAY, -1, false), ]; client From a400b7342a674b2572770ef72c24523ae0de7393 Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Mon, 6 Oct 2025 11:22:17 -0700 Subject: [PATCH 35/45] Improve --- Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Cargo.toml b/Cargo.toml index d7242f94b..728c0127c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -45,7 +45,7 @@ const-oid = { version = "0.9.6", default-features = false } constant_time_eq = { version = "0.4.2" } fail = { version = "0.5.1", default-features = false } futures = { version = "0.3.31", default-features = false } -gcp-bigquery-client = { path = "../gcp-bigquery-client", default-features = false } +gcp-bigquery-client = { git = "https://github.com/iambriccardo/gcp-bigquery-client", default-features = false, rev = "4759f728b9083f2288d44bec9338207d8d54e5ec" } iceberg = { version = "0.6.0", default-features = false } iceberg-catalog-rest = { version = "0.6.0", default-features = false } insta = { version = "1.43.1", default-features = false } From 84264023d45949629bd1993fbb5f35f261672e7c Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Mon, 6 Oct 2025 13:30:33 -0700 Subject: [PATCH 36/45] Improve --- etl-destinations/src/bigquery/client.rs | 17 ++ etl-destinations/src/bigquery/core.rs | 48 +++++- etl/src/conversions/event.rs | 41 ++++- etl/src/replication/apply.rs | 28 ++-- etl/src/types/event.rs | 20 ++- etl/tests/pipeline.rs | 210 +++++++++++++++++++++++- 6 files changed, 334 insertions(+), 30 deletions(-) diff --git a/etl-destinations/src/bigquery/client.rs b/etl-destinations/src/bigquery/client.rs index a2d546588..cce746aef 100644 --- a/etl-destinations/src/bigquery/client.rs +++ b/etl-destinations/src/bigquery/client.rs @@ -287,6 +287,23 @@ impl BigQueryClient { Ok(()) } + /// Drops a view from BigQuery if it exists. + pub async fn drop_view( + &self, + dataset_id: &BigQueryDatasetId, + view_name: &BigQueryTableId, + ) -> EtlResult<()> { + let full_view_name = self.full_table_name(dataset_id, view_name); + + info!("dropping view {full_view_name} from bigquery"); + + let query = format!("drop view if exists {full_view_name}"); + + let _ = self.query(QueryRequest::new(query)).await?; + + Ok(()) + } + /// Drops a table from BigQuery. /// /// Executes a DROP TABLE statement to remove the table and all its data. diff --git a/etl-destinations/src/bigquery/core.rs b/etl-destinations/src/bigquery/core.rs index 11208f638..989ff160c 100644 --- a/etl-destinations/src/bigquery/core.rs +++ b/etl-destinations/src/bigquery/core.rs @@ -484,6 +484,45 @@ where Ok(true) } + /// Handles a table rename by dropping the old view and updating view caches. + async fn handle_table_rename( + &self, + sequenced_table_id: &SequencedBigQueryTableId, + old_name: &TableName, + new_name: &TableName, + ) -> EtlResult<()> { + let old_view_name = table_name_to_bigquery_table_id(old_name); + let new_view_name = table_name_to_bigquery_table_id(new_name); + + if old_view_name == new_view_name { + return Ok(()); + } + + debug!( + table = %sequenced_table_id, + old_view = %old_view_name, + new_view = %new_view_name, + "renaming BigQuery table view" + ); + + self.client + .drop_view(&self.dataset_id, &old_view_name) + .await?; + + let target_table_id = sequenced_table_id.to_string(); + self.client + .create_or_replace_view(&self.dataset_id, &new_view_name, &target_table_id) + .await?; + + let mut inner = self.inner.lock().await; + inner.created_views.remove(&old_view_name); + inner + .created_views + .insert(new_view_name, sequenced_table_id.clone()); + + Ok(()) + } + /// Writes table rows with CDC metadata for non-event streaming operations. /// /// Adds an `Upsert` operation type to each row, splits them into optimal batches based on @@ -663,6 +702,10 @@ where let sequenced_table_name = sequenced_bigquery_table_id.to_string(); for change in changes { match change { + RelationChange::RenameTable { old_name, new_name } => { + self.handle_table_rename(&sequenced_bigquery_table_id, &old_name, &new_name) + .await?; + } RelationChange::AddColumn(column_schema) => { let column_name = column_schema.name.clone(); @@ -946,11 +989,10 @@ where Self::add_to_created_tables_cache(&mut inner, &next_sequenced_bigquery_table_id); // Update the view to point to the new table. + let view_name = table_name_to_bigquery_table_id(&table_schema.name); self.ensure_view_points_to_table( &mut inner, - // We convert the sequenced table ID to a BigQuery table ID since the view will have - // the name of the BigQuery table id (without the sequence number). - &sequenced_bigquery_table_id.to_bigquery_table_id(), + &view_name, &next_sequenced_bigquery_table_id, ) .await?; diff --git a/etl/src/conversions/event.rs b/etl/src/conversions/event.rs index 97cb57a45..d8630d160 100644 --- a/etl/src/conversions/event.rs +++ b/etl/src/conversions/event.rs @@ -1,17 +1,18 @@ use core::str; use etl_postgres::types::{ - ColumnSchema, SchemaVersion, TableId, VersionedTableSchema, convert_type_oid_to_type, + ColumnSchema, SchemaVersion, TableId, TableName, TableSchema, VersionedTableSchema, + convert_type_oid_to_type, }; use postgres_replication::protocol; use std::sync::Arc; use tokio_postgres::types::PgLsn; -use crate::bail; use crate::conversions::text::{default_value_for_type, parse_cell_from_postgres_text}; use crate::error::{ErrorKind, EtlError, EtlResult}; use crate::types::{ BeginEvent, Cell, CommitEvent, DeleteEvent, InsertEvent, TableRow, TruncateEvent, UpdateEvent, }; +use crate::{bail, etl_error}; /// Creates a [`BeginEvent`] from Postgres protocol data. /// @@ -48,13 +49,37 @@ pub fn parse_event_from_commit_message( } } -/// Creates a [`Vec`] from Postgres protocol data. +/// Creates a [`TableSchema`] from Postgres protocol data. /// -/// This method parses the replication protocol relation message and builds a vector of all the -/// columns that were received in the relation message. -pub async fn parse_column_schemas_from_relation_message( +/// This method parses the replication protocol relation message and builds the table schema that +/// it represents. +pub async fn parse_table_schema_from_relation_message( relation_body: &protocol::RelationBody, -) -> EtlResult> { +) -> EtlResult { + let table_id = TableId::new(relation_body.rel_id()); + let schema = relation_body.namespace().map_err(|error| { + etl_error!( + ErrorKind::InvalidData, + "Invalid namespace in relation message", + format!( + "Failed to decode namespace for relation {}: {error}", + relation_body.rel_id() + ) + ) + })?; + let table = relation_body.name().map_err(|error| { + etl_error!( + ErrorKind::InvalidData, + "Invalid table name in relation message", + format!( + "Failed to decode table name for relation {}: {error}", + relation_body.rel_id() + ) + ) + })?; + + let table_name = TableName::new(schema.to_string(), table.to_string()); + // We construct the new column schemas in order. The order is important since the table schema // relies on the right ordering to interpret the Postgres correctly. let new_column_schemas = relation_body @@ -63,7 +88,7 @@ pub async fn parse_column_schemas_from_relation_message( .map(build_column_schema) .collect::, EtlError>>()?; - Ok(new_column_schemas) + Ok(TableSchema::new(table_id, table_name, new_column_schemas)) } /// Converts a Postgres insert message into an [`InsertEvent`]. diff --git a/etl/src/replication/apply.rs b/etl/src/replication/apply.rs index 3d2f7d29f..bbc71384a 100644 --- a/etl/src/replication/apply.rs +++ b/etl/src/replication/apply.rs @@ -1,6 +1,6 @@ use etl_config::shared::PipelineConfig; use etl_postgres::replication::worker::WorkerType; -use etl_postgres::types::{TableId, TableSchema, VersionedTableSchema}; +use etl_postgres::types::{TableId, VersionedTableSchema}; use futures::StreamExt; use metrics::histogram; use postgres_replication::protocol; @@ -17,10 +17,10 @@ use crate::concurrency::shutdown::ShutdownRx; use crate::concurrency::signal::SignalRx; use crate::concurrency::stream::{TimeoutStream, TimeoutStreamResult}; use crate::conversions::event::{ - parse_column_schemas_from_relation_message, parse_event_from_begin_message, - parse_event_from_commit_message, parse_event_from_delete_message, - parse_event_from_insert_message, parse_event_from_truncate_message, - parse_event_from_update_message, + parse_event_from_begin_message, parse_event_from_commit_message, + parse_event_from_delete_message, parse_event_from_insert_message, + parse_event_from_truncate_message, parse_event_from_update_message, + parse_table_schema_from_relation_message, }; use crate::destination::Destination; use crate::error::{ErrorKind, EtlResult}; @@ -1104,28 +1104,26 @@ where return Ok(HandleMessageResult::no_event()); } - // Parse the relation message columns into column schemas. - let new_column_schemas = parse_column_schemas_from_relation_message(message).await?; + // Parse the relation message columns into column schemas and resolve the table name. + let new_table_schema = parse_table_schema_from_relation_message(message).await?; // We load the latest table schema before this relation message, which contains the last known // schema. let old_table_schema = load_latest_table_schema(schema_store, table_id).await?; + info!("SCHEMAS {:?} {:?}", new_table_schema, old_table_schema); + // If the column schemas are the same, we treat the relation message as a no-op. This is pretty // common since Postgres will send a `Relation` message as the first message for every new // connection even if the table schema hasn't changed. - if new_column_schemas == old_table_schema.column_schemas { + if new_table_schema.name == old_table_schema.name + && new_table_schema.column_schemas == old_table_schema.column_schemas + { return Ok(HandleMessageResult::no_event()); } // We store the new schema in the store and build the final relation event. - let new_table_schema = schema_store - .store_table_schema(TableSchema::new( - table_id, - old_table_schema.name.clone(), - new_column_schemas, - )) - .await?; + let new_table_schema = schema_store.store_table_schema(new_table_schema).await?; let event = RelationEvent { start_lsn, diff --git a/etl/src/types/event.rs b/etl/src/types/event.rs index 519f91121..848ad9dbc 100644 --- a/etl/src/types/event.rs +++ b/etl/src/types/event.rs @@ -1,4 +1,4 @@ -use etl_postgres::types::{ColumnSchema, SchemaVersion, TableId, VersionedTableSchema}; +use etl_postgres::types::{ColumnSchema, SchemaVersion, TableId, TableName, VersionedTableSchema}; use std::collections::HashSet; use std::fmt; use std::hash::Hash; @@ -46,6 +46,11 @@ pub struct CommitEvent { /// A change in a relation. #[derive(Debug, Clone, PartialEq)] pub enum RelationChange { + /// A change that describes renaming the table or moving it across schemas. + RenameTable { + old_name: TableName, + new_name: TableName, + }, /// A change that describes adding a new column. AddColumn(ColumnSchema), /// A change that describes dropping an existing column. @@ -78,8 +83,18 @@ pub struct RelationEvent { impl RelationEvent { /// Builds a list of [`RelationChange`]s that describe the changes between the old and new table - /// schemas. + /// schemas. Table-level operations (such as renames) are emitted before column-level operations + /// so destinations can adjust object names before applying column mutations. pub fn build_changes(&self) -> Vec { + let mut changes = Vec::new(); + + if self.old_table_schema.name != self.new_table_schema.name { + changes.push(RelationChange::RenameTable { + old_name: self.old_table_schema.name.clone(), + new_name: self.new_table_schema.name.clone(), + }); + } + // We build a lookup set for the new column schemas for quick change detection. let mut new_indexed_column_schemas = self .new_table_schema @@ -90,7 +105,6 @@ impl RelationEvent { .collect::>(); // We process all the changes that we want to dispatch to the destination. - let mut changes = vec![]; for column_schema in self.old_table_schema.column_schemas.iter() { let column_schema = IndexedColumnSchema(column_schema.clone()); let latest_column_schema = new_indexed_column_schemas.take(&column_schema); diff --git a/etl/tests/pipeline.rs b/etl/tests/pipeline.rs index bcbfd2aa9..678220d8d 100644 --- a/etl/tests/pipeline.rs +++ b/etl/tests/pipeline.rs @@ -14,10 +14,11 @@ use etl::test_utils::test_schema::{ build_expected_users_inserts, get_n_integers_sum, get_users_age_sum_from_rows, insert_mock_data, insert_orders_data, insert_users_data, setup_test_database_schema, }; -use etl::types::{EventType, PipelineId}; +use etl::types::{Event, EventType, PipelineId, RelationChange}; use etl_config::shared::BatchConfig; use etl_postgres::replication::slots::EtlReplicationSlot; use etl_postgres::tokio::test_utils::TableModification; +use etl_postgres::types::TableName; use etl_telemetry::tracing::init_test_tracing; use rand::random; use std::time::Duration; @@ -749,6 +750,213 @@ async fn table_schema_changes_are_handled_correctly() { pipeline.shutdown_and_wait().await.unwrap(); } +#[tokio::test(flavor = "multi_thread")] +async fn table_renames_are_handled_correctly() { + init_test_tracing(); + let mut database = spawn_source_database().await; + let database_schema = setup_test_database_schema(&database, TableSelection::UsersOnly).await; + + // Seed a single row so the initial copy has data. + insert_users_data(&mut database, &database_schema.users_schema().name, 1..=1).await; + + let store = NotifyingStore::new(); + let destination = TestDestinationWrapper::wrap(MemoryDestination::new()); + + let pipeline_id: PipelineId = random(); + let mut pipeline = create_pipeline( + &database.config, + pipeline_id, + database_schema.publication_name(), + store.clone(), + destination.clone(), + ); + + let users_state_notify = store + .notify_on_table_state_type( + database_schema.users_schema().id, + TableReplicationPhaseType::SyncDone, + ) + .await; + + pipeline.start().await.unwrap(); + + users_state_notify.notified().await; + + // Validate the initial schema snapshot. + let initial_table_schemas = store.get_latest_table_schemas().await; + let initial_users_schema = initial_table_schemas + .get(&database_schema.users_schema().id) + .expect("users table schema not stored"); + assert_eq!(initial_users_schema.version, 0); + assert_eq!(initial_users_schema.name, test_table_name("users")); + + // Create a new schema for the future table to be moved here. + database + .client + .as_ref() + .unwrap() + .execute("create schema renamed", &[]) + .await + .unwrap(); + + // Rename table. + database + .client + .as_ref() + .unwrap() + .execute("alter table test.users rename to customers", &[]) + .await + .unwrap(); + + // Register notifications for the inserts. + let insert_event_notify = destination + .wait_for_events_count(vec![(EventType::Insert, 1)]) + .await; + + database + .insert_values( + TableName::new("test".to_owned(), "customers".to_owned()), + &["name", "age"], + &[&"customer_2", &25i32], + ) + .await + .expect("Failed to insert user"); + + insert_event_notify.notified().await; + + let events = destination.get_events().await; + let grouped_events = group_events_by_type_and_table_id(&events); + let users_inserts = grouped_events + .get(&(EventType::Insert, database_schema.users_schema().id)) + .unwrap(); + assert_eq!(users_inserts.len(), 1); + + let table_schemas = store.get_latest_table_schemas().await; + let users_table_schema = table_schemas + .get(&database_schema.users_schema().id) + .expect("users schema missing after rename"); + assert_eq!(users_table_schema.version, 1); + assert_eq!(users_table_schema.name, test_table_name("customers")); + + // assert_eq!( + // users_table_schema.version, + // rename_event.new_table_schema.version + // ); + // assert_eq!(users_table_schema.name, renamed_users_name); + + // // Insert data into the renamed table and ensure it streams correctly. + // let insert_notify = destination + // .wait_for_events_count(vec![(EventType::Insert, 1)]) + // .await; + // database + // .insert_values( + // renamed_users_name.clone(), + // &["name", "age"], + // &[&"user_after_name_rename", &(42i32)], + // ) + // .await + // .expect("failed to insert after table rename"); + // insert_notify.notified().await; + // + // let post_rename_events = destination.get_events().await; + // let post_rename_insert = post_rename_events + // .iter() + // .find_map(|event| { + // if let Event::Insert(insert) = event { + // Some(insert.clone()) + // } else { + // None + // } + // }) + // .expect("expected insert event after table rename"); + // assert_eq!( + // post_rename_insert.schema_version, + // rename_event.new_table_schema.version + // ); + // + // destination.clear_events().await; + // + // // --- Move the table to a different schema --- + // let relation_notify = destination + // .wait_for_events_count(vec![(EventType::Relation, 1)]) + // .await; + // + // database + // .client + // .as_ref() + // .unwrap() + // .execute("alter table test.customers set schema renamed", &[]) + // .await + // .unwrap(); + // + // relation_notify.notified().await; + // + // let relation_events = destination.get_events().await; + // let move_event = relation_events + // .iter() + // .find_map(|event| { + // if let Event::Relation(relation) = event { + // Some(relation.clone()) + // } else { + // None + // } + // }) + // .expect("expected relation event for schema move"); + // + // let moved_users_name = TableName::new("renamed".to_owned(), "customers".to_owned()); + // assert_eq!(move_event.new_table_schema.name, moved_users_name); + // assert_eq!(move_event.old_table_schema.name, renamed_users_name); + // let changes = move_event.build_changes(); + // assert_eq!(changes.len(), 1); + // assert!(matches!( + // &changes[0], + // RelationChange::RenameTable { new_name, .. } if new_name == &moved_users_name + // )); + // + // let table_schemas = store.get_latest_table_schemas().await; + // let latest_users_schema = table_schemas + // .get(&database_schema.users_schema().id) + // .expect("users schema missing after schema move"); + // assert_eq!( + // latest_users_schema.version, + // move_event.new_table_schema.version + // ); + // assert_eq!(latest_users_schema.name, moved_users_name); + // + // destination.clear_events().await; + // + // let insert_notify = destination + // .wait_for_events_count(vec![(EventType::Insert, 1)]) + // .await; + // database + // .insert_values( + // moved_users_name.clone(), + // &["name", "age"], + // &[&"user_after_schema_move", &(43i32)], + // ) + // .await + // .expect("failed to insert after schema move"); + // insert_notify.notified().await; + // + // let post_move_events = destination.get_events().await; + // let post_move_insert = post_move_events + // .iter() + // .find_map(|event| { + // if let Event::Insert(insert) = event { + // Some(insert.clone()) + // } else { + // None + // } + // }) + // .expect("expected insert event after schema move"); + // assert_eq!( + // post_move_insert.schema_version, + // move_event.new_table_schema.version + // ); + + pipeline.shutdown_and_wait().await.unwrap(); +} + #[tokio::test(flavor = "multi_thread")] async fn table_copy_and_sync_streams_new_data() { init_test_tracing(); From 5c8f4db8757909f851a6570c1b054142135a6a33 Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Mon, 6 Oct 2025 13:34:47 -0700 Subject: [PATCH 37/45] Improve --- etl/tests/pipeline.rs | 168 ++++++++++++------------------------------ 1 file changed, 49 insertions(+), 119 deletions(-) diff --git a/etl/tests/pipeline.rs b/etl/tests/pipeline.rs index 678220d8d..a80c0f66d 100644 --- a/etl/tests/pipeline.rs +++ b/etl/tests/pipeline.rs @@ -14,7 +14,7 @@ use etl::test_utils::test_schema::{ build_expected_users_inserts, get_n_integers_sum, get_users_age_sum_from_rows, insert_mock_data, insert_orders_data, insert_users_data, setup_test_database_schema, }; -use etl::types::{Event, EventType, PipelineId, RelationChange}; +use etl::types::{EventType, PipelineId}; use etl_config::shared::BatchConfig; use etl_postgres::replication::slots::EtlReplicationSlot; use etl_postgres::tokio::test_utils::TableModification; @@ -817,13 +817,14 @@ async fn table_renames_are_handled_correctly() { .insert_values( TableName::new("test".to_owned(), "customers".to_owned()), &["name", "age"], - &[&"customer_2", &25i32], + &[&"customer_2", &2i32], ) .await .expect("Failed to insert user"); insert_event_notify.notified().await; + // Check the new event. let events = destination.get_events().await; let grouped_events = group_events_by_type_and_table_id(&events); let users_inserts = grouped_events @@ -831,128 +832,57 @@ async fn table_renames_are_handled_correctly() { .unwrap(); assert_eq!(users_inserts.len(), 1); + // Check updated schema. let table_schemas = store.get_latest_table_schemas().await; let users_table_schema = table_schemas .get(&database_schema.users_schema().id) - .expect("users schema missing after rename"); + .unwrap(); assert_eq!(users_table_schema.version, 1); assert_eq!(users_table_schema.name, test_table_name("customers")); - - // assert_eq!( - // users_table_schema.version, - // rename_event.new_table_schema.version - // ); - // assert_eq!(users_table_schema.name, renamed_users_name); - - // // Insert data into the renamed table and ensure it streams correctly. - // let insert_notify = destination - // .wait_for_events_count(vec![(EventType::Insert, 1)]) - // .await; - // database - // .insert_values( - // renamed_users_name.clone(), - // &["name", "age"], - // &[&"user_after_name_rename", &(42i32)], - // ) - // .await - // .expect("failed to insert after table rename"); - // insert_notify.notified().await; - // - // let post_rename_events = destination.get_events().await; - // let post_rename_insert = post_rename_events - // .iter() - // .find_map(|event| { - // if let Event::Insert(insert) = event { - // Some(insert.clone()) - // } else { - // None - // } - // }) - // .expect("expected insert event after table rename"); - // assert_eq!( - // post_rename_insert.schema_version, - // rename_event.new_table_schema.version - // ); - // - // destination.clear_events().await; - // - // // --- Move the table to a different schema --- - // let relation_notify = destination - // .wait_for_events_count(vec![(EventType::Relation, 1)]) - // .await; - // - // database - // .client - // .as_ref() - // .unwrap() - // .execute("alter table test.customers set schema renamed", &[]) - // .await - // .unwrap(); - // - // relation_notify.notified().await; - // - // let relation_events = destination.get_events().await; - // let move_event = relation_events - // .iter() - // .find_map(|event| { - // if let Event::Relation(relation) = event { - // Some(relation.clone()) - // } else { - // None - // } - // }) - // .expect("expected relation event for schema move"); - // - // let moved_users_name = TableName::new("renamed".to_owned(), "customers".to_owned()); - // assert_eq!(move_event.new_table_schema.name, moved_users_name); - // assert_eq!(move_event.old_table_schema.name, renamed_users_name); - // let changes = move_event.build_changes(); - // assert_eq!(changes.len(), 1); - // assert!(matches!( - // &changes[0], - // RelationChange::RenameTable { new_name, .. } if new_name == &moved_users_name - // )); - // - // let table_schemas = store.get_latest_table_schemas().await; - // let latest_users_schema = table_schemas - // .get(&database_schema.users_schema().id) - // .expect("users schema missing after schema move"); - // assert_eq!( - // latest_users_schema.version, - // move_event.new_table_schema.version - // ); - // assert_eq!(latest_users_schema.name, moved_users_name); - // - // destination.clear_events().await; - // - // let insert_notify = destination - // .wait_for_events_count(vec![(EventType::Insert, 1)]) - // .await; - // database - // .insert_values( - // moved_users_name.clone(), - // &["name", "age"], - // &[&"user_after_schema_move", &(43i32)], - // ) - // .await - // .expect("failed to insert after schema move"); - // insert_notify.notified().await; - // - // let post_move_events = destination.get_events().await; - // let post_move_insert = post_move_events - // .iter() - // .find_map(|event| { - // if let Event::Insert(insert) = event { - // Some(insert.clone()) - // } else { - // None - // } - // }) - // .expect("expected insert event after schema move"); - // assert_eq!( - // post_move_insert.schema_version, - // move_event.new_table_schema.version - // ); + + // Move the table to a different schema. + database + .client + .as_ref() + .unwrap() + .execute("alter table test.customers set schema renamed", &[]) + .await + .unwrap(); + + // Register notifications for the inserts. + let insert_event_notify = destination + .wait_for_events_count(vec![(EventType::Insert, 2)]) + .await; + + database + .insert_values( + TableName::new("renamed".to_owned(), "customers".to_owned()), + &["name", "age"], + &[&"customer_3", &3i32], + ) + .await + .expect("Failed to insert user"); + + insert_event_notify.notified().await; + + // Check the new event. + let events = destination.get_events().await; + let grouped_events = group_events_by_type_and_table_id(&events); + let users_inserts = grouped_events + .get(&(EventType::Insert, database_schema.users_schema().id)) + .unwrap(); + assert_eq!(users_inserts.len(), 2); + + // Check updated schema. + let table_schemas = store.get_latest_table_schemas().await; + let users_table_schema = table_schemas + .get(&database_schema.users_schema().id) + .unwrap(); + assert_eq!(users_table_schema.version, 2); + assert_eq!( + users_table_schema.name, + TableName::new("renamed".to_owned(), "customers".to_owned()) + ); pipeline.shutdown_and_wait().await.unwrap(); } From a837b89cdb2890d59f4df59a3f9a7c8b30541fd9 Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Mon, 6 Oct 2025 13:47:31 -0700 Subject: [PATCH 38/45] Improve --- etl-destinations/src/bigquery/core.rs | 111 ++++++++++++-------------- 1 file changed, 51 insertions(+), 60 deletions(-) diff --git a/etl-destinations/src/bigquery/core.rs b/etl-destinations/src/bigquery/core.rs index 989ff160c..c5ad00ce8 100644 --- a/etl-destinations/src/bigquery/core.rs +++ b/etl-destinations/src/bigquery/core.rs @@ -63,7 +63,7 @@ pub fn table_name_to_bigquery_table_id(table_name: &TableName) -> BigQueryTableI /// Combines a base table name with a sequence number to enable versioned tables. /// Used for truncate handling where each truncate creates a new table version. #[derive(Debug, Clone, Eq, PartialEq, Hash)] -struct SequencedBigQueryTableId(BigQueryTableId, u64); +struct SequencedBigQueryTableId(pub BigQueryTableId, pub u64); impl SequencedBigQueryTableId { /// Creates a new sequenced table ID starting at version 0. @@ -75,11 +75,6 @@ impl SequencedBigQueryTableId { pub fn next(&self) -> Self { Self(self.0.clone(), self.1 + 1) } - - /// Extracts the base BigQuery table ID without the sequence number. - pub fn to_bigquery_table_id(&self) -> BigQueryTableId { - self.0.clone() - } } impl FromStr for SequencedBigQueryTableId { @@ -919,40 +914,50 @@ where /// to optimize multiple truncates of the same table. async fn process_truncate_for_table_ids( &self, - table_ids: impl IntoIterator, - is_cdc_truncate: bool, + table_ids: impl IntoIterator)>, + is_truncate_event: bool, ) -> EtlResult<()> { // We want to lock for the entire processing to ensure that we don't have any race conditions // and possible errors are easier to reason about. let mut inner = self.inner.lock().await; for (table_id, schema_version) in table_ids { - let table_schema = self - .store - .get_table_schema(&table_id, schema_version) - .await?; - - // If we are not doing CDC, it means that this truncation has been issued while recovering - // from a failed data sync operation. In that case, we could have failed before table schemas - // were stored in the schema store, so if we don't find a table schema we just continue - // and emit a warning. If we are doing CDC, it's a problem if the schema disappears while - // streaming, so we error out. - if table_schema.is_none() && !is_cdc_truncate { - warn!( - "the table schema for table {table_id} was not found in the schema store while processing truncate events for BigQuery", - table_id = table_id.to_string() - ); - - continue; - } + // If a schema version is supplied, we use it to create a new table. Otherwise, we use + // the latest schema version since that is what we have available. + let table_schema = match schema_version { + Some(schema_version) => { + self.store + .get_table_schema(&table_id, schema_version) + .await? + } + None => self.store.get_latest_table_schema(&table_id).await?, + }; + + let table_schema = match table_schema { + Some(table_schema) => table_schema, + // If we are not doing CDC, it means that this truncation has been issued while recovering + // from a failed data sync operation. In that case, we could have failed before table schemas + // were stored in the schema store, so if we don't find a table schema we just continue + // and emit a warning. If we are doing CDC, it's a problem if the schema disappears while + // streaming, so we error out. + None if !is_truncate_event => { + warn!( + "the table schema for table {table_id} was not found in the schema store while processing truncate events for BigQuery", + table_id = table_id.to_string() + ); - let table_schema = table_schema.ok_or_else(|| etl_error!( - ErrorKind::MissingTableSchema, - "Table not found in the schema store", - format!( - "The table schema for table {table_id} was not found in the schema store while processing truncate events for BigQuery" + continue; + } + None => { + bail!( + ErrorKind::MissingTableSchema, + "Table not found in the schema store", + format!( + "The table schema for table {table_id} was not found in the schema store while processing truncate events for BigQuery" + ) ) - ))?; + } + }; // We need to determine the current sequenced table ID for this table. let sequenced_bigquery_table_id = @@ -1058,21 +1063,7 @@ where } async fn truncate_table(&self, table_id: TableId) -> EtlResult<()> { - let latest_schema = self - .store - .get_latest_table_schema(&table_id) - .await? - .ok_or_else(|| { - etl_error!( - ErrorKind::MissingTableSchema, - "Table not found in the schema store", - format!( - "The table schema for table {table_id} was not found in the schema store" - ) - ) - })?; - - self.process_truncate_for_table_ids(iter::once((table_id, latest_schema.version)), false) + self.process_truncate_for_table_ids(iter::once((table_id, None)), false) .await } @@ -1185,7 +1176,7 @@ mod tests { fn test_sequenced_bigquery_table_id_from_str_valid() { let table_id = "users_table_123"; let parsed = table_id.parse::().unwrap(); - assert_eq!(parsed.to_bigquery_table_id(), "users_table"); + assert_eq!(parsed.0, "users_table"); assert_eq!(parsed.1, 123); } @@ -1193,7 +1184,7 @@ mod tests { fn test_sequenced_bigquery_table_id_from_str_zero_sequence() { let table_id = "simple_table_0"; let parsed = table_id.parse::().unwrap(); - assert_eq!(parsed.to_bigquery_table_id(), "simple_table"); + assert_eq!(parsed.0, "simple_table"); assert_eq!(parsed.1, 0); } @@ -1201,7 +1192,7 @@ mod tests { fn test_sequenced_bigquery_table_id_from_str_large_sequence() { let table_id = "test_table_18446744073709551615"; // u64::MAX let parsed = table_id.parse::().unwrap(); - assert_eq!(parsed.to_bigquery_table_id(), "test_table"); + assert_eq!(parsed.0, "test_table"); assert_eq!(parsed.1, u64::MAX); } @@ -1209,7 +1200,7 @@ mod tests { fn test_sequenced_bigquery_table_id_from_str_escaped_underscores() { let table_id = "a__b_c__d_42"; let parsed = table_id.parse::().unwrap(); - assert_eq!(parsed.to_bigquery_table_id(), "a__b_c__d"); + assert_eq!(parsed.0, "a__b_c__d"); assert_eq!(parsed.1, 42); } @@ -1240,14 +1231,14 @@ mod tests { #[test] fn test_sequenced_bigquery_table_id_new() { let table_id = SequencedBigQueryTableId::new("users_table".to_string()); - assert_eq!(table_id.to_bigquery_table_id(), "users_table"); + assert_eq!(table_id.0, "users_table"); assert_eq!(table_id.1, 0); } #[test] fn test_sequenced_bigquery_table_id_new_with_underscores() { let table_id = SequencedBigQueryTableId::new("a__b_c__d".to_string()); - assert_eq!(table_id.to_bigquery_table_id(), "a__b_c__d"); + assert_eq!(table_id.0, "a__b_c__d"); assert_eq!(table_id.1, 0); } @@ -1258,7 +1249,7 @@ mod tests { assert_eq!(table_id.1, 0); assert_eq!(next_table_id.1, 1); - assert_eq!(next_table_id.to_bigquery_table_id(), "users_table"); + assert_eq!(next_table_id.0, "users_table"); } #[test] @@ -1267,7 +1258,7 @@ mod tests { let next_table_id = table_id.next(); assert_eq!(next_table_id.1, 43); - assert_eq!(next_table_id.to_bigquery_table_id(), "test_table"); + assert_eq!(next_table_id.0, "test_table"); } #[test] @@ -1276,25 +1267,25 @@ mod tests { let next_table_id = table_id.next(); assert_eq!(next_table_id.1, u64::MAX); - assert_eq!(next_table_id.to_bigquery_table_id(), "test_table"); + assert_eq!(next_table_id.0, "test_table"); } #[test] fn test_sequenced_bigquery_table_id_to_bigquery_table_id() { let table_id = SequencedBigQueryTableId("users_table".to_string(), 123); - assert_eq!(table_id.to_bigquery_table_id(), "users_table"); + assert_eq!(table_id.0, "users_table"); } #[test] fn test_sequenced_bigquery_table_id_to_bigquery_table_id_with_underscores() { let table_id = SequencedBigQueryTableId("a__b_c__d".to_string(), 42); - assert_eq!(table_id.to_bigquery_table_id(), "a__b_c__d"); + assert_eq!(table_id.0, "a__b_c__d"); } #[test] fn test_sequenced_bigquery_table_id_to_bigquery_table_id_zero_sequence() { let table_id = SequencedBigQueryTableId("simple_table".to_string(), 0); - assert_eq!(table_id.to_bigquery_table_id(), "simple_table"); + assert_eq!(table_id.0, "simple_table"); } #[test] @@ -1421,7 +1412,7 @@ mod tests { let parsed = original.parse::().unwrap(); let formatted = parsed.to_string(); assert_eq!(original, formatted); - assert_eq!(parsed.to_bigquery_table_id(), "a__b_c__d"); + assert_eq!(parsed.0, "a__b_c__d"); assert_eq!(parsed.1, 999); } From b7418b8ca6c6a89485757b6b271f307641df3dd0 Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Mon, 6 Oct 2025 13:50:03 -0700 Subject: [PATCH 39/45] Improve --- etl-destinations/src/bigquery/core.rs | 7 +++++-- etl/src/replication/apply.rs | 2 -- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/etl-destinations/src/bigquery/core.rs b/etl-destinations/src/bigquery/core.rs index c5ad00ce8..b2d04693b 100644 --- a/etl-destinations/src/bigquery/core.rs +++ b/etl-destinations/src/bigquery/core.rs @@ -890,8 +890,11 @@ where // Finish the current batch before a TRUNCATE (it affects the table state). self.flush_batch(&mut table_batches_by_id).await?; - self.process_truncate_for_table_ids(truncate.table_ids.into_iter(), true) - .await?; + let table_ids = truncate + .table_ids + .into_iter() + .map(|(table_id, schema_version)| (table_id, Some(schema_version))); + self.process_truncate_for_table_ids(table_ids, true).await?; } // Unsupported events. diff --git a/etl/src/replication/apply.rs b/etl/src/replication/apply.rs index bbc71384a..37eeaf362 100644 --- a/etl/src/replication/apply.rs +++ b/etl/src/replication/apply.rs @@ -1111,8 +1111,6 @@ where // schema. let old_table_schema = load_latest_table_schema(schema_store, table_id).await?; - info!("SCHEMAS {:?} {:?}", new_table_schema, old_table_schema); - // If the column schemas are the same, we treat the relation message as a no-op. This is pretty // common since Postgres will send a `Relation` message as the first message for every new // connection even if the table schema hasn't changed. From 8eea075e0dc6d9420f780e7921b4676ce608b3cb Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Mon, 6 Oct 2025 14:20:16 -0700 Subject: [PATCH 40/45] Improve --- etl-destinations/tests/iceberg_destination.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/etl-destinations/tests/iceberg_destination.rs b/etl-destinations/tests/iceberg_destination.rs index d5246ece8..43e9755b5 100644 --- a/etl-destinations/tests/iceberg_destination.rs +++ b/etl-destinations/tests/iceberg_destination.rs @@ -344,7 +344,7 @@ async fn cdc_streaming() { TableRow { values: vec![ Cell::I64(1), - Cell::String("".to_string()), + Cell::Null, Cell::I32(0), IcebergOperationType::Delete.into(), ], @@ -405,7 +405,7 @@ async fn cdc_streaming() { TableRow { values: vec![ Cell::I64(2), - Cell::String("".to_string()), + Cell::Null, IcebergOperationType::Delete.into(), ], }, From c63e2f878f7e16f98a2f020bd38abeb59b32dd5c Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Mon, 6 Oct 2025 14:27:49 -0700 Subject: [PATCH 41/45] Improve --- etl-destinations/tests/iceberg_destination.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/etl-destinations/tests/iceberg_destination.rs b/etl-destinations/tests/iceberg_destination.rs index 43e9755b5..58101fe49 100644 --- a/etl-destinations/tests/iceberg_destination.rs +++ b/etl-destinations/tests/iceberg_destination.rs @@ -345,7 +345,7 @@ async fn cdc_streaming() { values: vec![ Cell::I64(1), Cell::Null, - Cell::I32(0), + Cell::Null, IcebergOperationType::Delete.into(), ], }, From 92cf8e875cf3b2c43bda806139eb9cbbc25b59e2 Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Mon, 6 Oct 2025 14:44:51 -0700 Subject: [PATCH 42/45] Improve --- etl-destinations/src/bigquery/core.rs | 6 +- etl-destinations/tests/bigquery_pipeline.rs | 156 +++++++++++++++++++- 2 files changed, 157 insertions(+), 5 deletions(-) diff --git a/etl-destinations/src/bigquery/core.rs b/etl-destinations/src/bigquery/core.rs index b2d04693b..592271f9c 100644 --- a/etl-destinations/src/bigquery/core.rs +++ b/etl-destinations/src/bigquery/core.rs @@ -28,9 +28,9 @@ const BIGQUERY_TABLE_ID_DELIMITER: &str = "_"; /// Replacement string for escaping underscores in Postgres names. const BIGQUERY_TABLE_ID_DELIMITER_ESCAPE_REPLACEMENT: &str = "__"; /// Maximum number of BigQuery streaming attempts when schema propagation lags behind. -const MAX_SCHEMA_MISMATCH_ATTEMPTS: usize = 5; +const MAX_SCHEMA_MISMATCH_ATTEMPTS: usize = 10; /// Delay in milliseconds between retry attempts triggered by BigQuery schema mismatches. -const SCHEMA_MISMATCH_RETRY_DELAY_MS: u64 = 500; +const SCHEMA_MISMATCH_RETRY_DELAY_MS: u64 = 1000; /// Returns the [`BigQueryTableId`] for a supplied [`TableName`]. /// @@ -625,7 +625,7 @@ where /// Streams table batches to BigQuery, retrying when schema mismatch errors occur. /// /// The rationale is that per BigQuery docs, the Storage Write API detects schema changes after - /// a short time, on the order of minutes. + /// a short time, on the order of minutes, thus we want to retry for a bit until we succeed. async fn stream_with_schema_retry(&self, table_batches: T) -> EtlResult<(usize, usize)> where T: Into]>>, diff --git a/etl-destinations/tests/bigquery_pipeline.rs b/etl-destinations/tests/bigquery_pipeline.rs index bc5eb037a..acec8a9c0 100644 --- a/etl-destinations/tests/bigquery_pipeline.rs +++ b/etl-destinations/tests/bigquery_pipeline.rs @@ -15,7 +15,7 @@ use etl::test_utils::test_destination_wrapper::TestDestinationWrapper; use etl::test_utils::test_schema::{ TableSelection, insert_mock_data, insert_users_data, setup_test_database_schema, }; -use etl::types::{EventType, PgNumeric, PipelineId}; +use etl::types::{EventType, PgNumeric, PipelineId, TableName}; use etl_destinations::bigquery::install_crypto_provider_for_bigquery; use etl_postgres::tokio::test_utils::TableModification; use etl_telemetry::tracing::init_test_tracing; @@ -497,7 +497,7 @@ async fn table_truncate_with_batching() { } #[tokio::test(flavor = "multi_thread")] -async fn relation_event_primary_key_syncs_in_bigquery() { +async fn table_schema_changes_are_handled_correctly() { init_test_tracing(); install_crypto_provider_for_bigquery(); @@ -628,6 +628,158 @@ async fn relation_event_primary_key_syncs_in_bigquery() { pipeline.shutdown_and_wait().await.unwrap(); } +#[tokio::test(flavor = "multi_thread")] +async fn table_renames_are_handled_correctly() { + init_test_tracing(); + install_crypto_provider_for_bigquery(); + + let mut database = spawn_source_database().await; + let database_schema = setup_test_database_schema(&database, TableSelection::UsersOnly).await; + + insert_users_data( + &mut database, + &database_schema.users_schema().name, + 1..=1, + ) + .await; + + let bigquery_database = setup_bigquery_connection().await; + + let store = NotifyingStore::new(); + let pipeline_id: PipelineId = random(); + let raw_destination = bigquery_database.build_destination(store.clone()).await; + let destination = TestDestinationWrapper::wrap(raw_destination); + + let mut pipeline = create_pipeline( + &database.config, + pipeline_id, + database_schema.publication_name(), + store.clone(), + destination.clone(), + ); + + let users_state_notify = store + .notify_on_table_state_type( + database_schema.users_schema().id, + TableReplicationPhaseType::SyncDone, + ) + .await; + + pipeline.start().await.unwrap(); + + users_state_notify.notified().await; + + let initial_rows = bigquery_database + .query_table(database_schema.users_schema().name.clone()) + .await + .unwrap(); + let parsed_initial_rows = parse_bigquery_table_rows::(initial_rows); + assert_eq!(parsed_initial_rows, vec![BigQueryUser::new(1, "user_1", 1)]); + + database + .client + .as_ref() + .unwrap() + .execute("create schema renamed", &[]) + .await + .unwrap(); + + database + .client + .as_ref() + .unwrap() + .execute("alter table test.users rename to customers", &[]) + .await + .unwrap(); + + let insert_event_notify = destination + .wait_for_events_count(vec![(EventType::Insert, 1)]) + .await; + + database + .insert_values( + TableName::new("test".to_owned(), "customers".to_owned()), + &["name", "age"], + &[&"customer_2", &2i32], + ) + .await + .unwrap(); + + insert_event_notify.notified().await; + + let renamed_rows = bigquery_database + .query_table(TableName::new("test".to_owned(), "customers".to_owned())) + .await + .unwrap(); + let parsed_renamed_rows = parse_bigquery_table_rows::(renamed_rows); + assert_eq!( + parsed_renamed_rows, + vec![ + BigQueryUser::new(1, "user_1", 1), + BigQueryUser::new(2, "customer_2", 2), + ], + ); + + let table_schemas = store.get_latest_table_schemas().await; + let users_schema = table_schemas + .get(&database_schema.users_schema().id) + .unwrap(); + assert_eq!(users_schema.version, 1); + assert_eq!( + users_schema.name, + TableName::new("test".to_owned(), "customers".to_owned()) + ); + + database + .client + .as_ref() + .unwrap() + .execute("alter table test.customers set schema renamed", &[]) + .await + .unwrap(); + + let insert_event_notify = destination + .wait_for_events_count(vec![(EventType::Insert, 1)]) + .await; + + database + .insert_values( + TableName::new("renamed".to_owned(), "customers".to_owned()), + &["name", "age"], + &[&"customer_3", &3i32], + ) + .await + .unwrap(); + + insert_event_notify.notified().await; + + let moved_rows = bigquery_database + .query_table(TableName::new("renamed".to_owned(), "customers".to_owned())) + .await + .unwrap(); + let parsed_moved_rows = parse_bigquery_table_rows::(moved_rows); + assert_eq!( + parsed_moved_rows, + vec![ + BigQueryUser::new(1, "user_1", 1), + BigQueryUser::new(2, "customer_2", 2), + BigQueryUser::new(3, "customer_3", 3), + ], + ); + + let table_schemas = store.get_latest_table_schemas().await; + let users_schema = table_schemas + .get(&database_schema.users_schema().id) + .unwrap(); + assert_eq!(users_schema.version, 2); + assert_eq!( + users_schema.name, + TableName::new("renamed".to_owned(), "customers".to_owned()) + ); + + pipeline.shutdown_and_wait().await.unwrap(); +} + #[tokio::test(flavor = "multi_thread")] async fn table_nullable_scalar_columns() { init_test_tracing(); From 3f0a917b0de7c0cf993a01e4cc9475493d578c14 Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Mon, 6 Oct 2025 14:49:15 -0700 Subject: [PATCH 43/45] Improve --- etl-destinations/tests/bigquery_pipeline.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/etl-destinations/tests/bigquery_pipeline.rs b/etl-destinations/tests/bigquery_pipeline.rs index acec8a9c0..977b30b8c 100644 --- a/etl-destinations/tests/bigquery_pipeline.rs +++ b/etl-destinations/tests/bigquery_pipeline.rs @@ -739,7 +739,7 @@ async fn table_renames_are_handled_correctly() { .unwrap(); let insert_event_notify = destination - .wait_for_events_count(vec![(EventType::Insert, 1)]) + .wait_for_events_count(vec![(EventType::Insert, 2)]) .await; database From dae764dd7b67569f41d38b0b8a58afba5bc7a44c Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Mon, 6 Oct 2025 14:52:48 -0700 Subject: [PATCH 44/45] Improve --- etl-destinations/src/bigquery/core.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/etl-destinations/src/bigquery/core.rs b/etl-destinations/src/bigquery/core.rs index 592271f9c..f53f2ec96 100644 --- a/etl-destinations/src/bigquery/core.rs +++ b/etl-destinations/src/bigquery/core.rs @@ -479,7 +479,9 @@ where Ok(true) } - /// Handles a table rename by dropping the old view and updating view caches. + /// Handles a table rename by dropping the old view, creating a new one and updating view caches. + /// + /// For performance reasons, we only rename the view and not the underlying table. async fn handle_table_rename( &self, sequenced_table_id: &SequencedBigQueryTableId, From 5ecac8444bb1c94c070fb12a48860c53e626896a Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Mon, 6 Oct 2025 14:54:45 -0700 Subject: [PATCH 45/45] Improve --- etl-destinations/tests/bigquery_pipeline.rs | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/etl-destinations/tests/bigquery_pipeline.rs b/etl-destinations/tests/bigquery_pipeline.rs index 977b30b8c..022432c05 100644 --- a/etl-destinations/tests/bigquery_pipeline.rs +++ b/etl-destinations/tests/bigquery_pipeline.rs @@ -636,12 +636,7 @@ async fn table_renames_are_handled_correctly() { let mut database = spawn_source_database().await; let database_schema = setup_test_database_schema(&database, TableSelection::UsersOnly).await; - insert_users_data( - &mut database, - &database_schema.users_schema().name, - 1..=1, - ) - .await; + insert_users_data(&mut database, &database_schema.users_schema().name, 1..=1).await; let bigquery_database = setup_bigquery_connection().await;