diff --git a/Cargo.toml b/Cargo.toml index b5a0240b7..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 = { version = "0.27.0", 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 } 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/Cargo.toml b/etl-destinations/Cargo.toml index 910f238d0..52294becf 100644 --- a/etl-destinations/Cargo.toml +++ b/etl-destinations/Cargo.toml @@ -58,6 +58,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 } chrono = { workspace = true } diff --git a/etl-destinations/src/bigquery/client.rs b/etl-destinations/src/bigquery/client.rs index 73c90fb7b..cce746aef 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. @@ -133,7 +149,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 @@ -267,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. @@ -286,6 +323,149 @@ 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?; + + info!( + "added column {} to table {full_table_name} in BigQuery", + column_schema.name + ); + + 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 = Self::quote_identifier(column_name); + let query = + format!("alter table {full_table_name} drop column if exists {column_identifier}"); + + let _ = self.query(QueryRequest::new(query)).await?; + + info!( + "dropped column {} from table {full_table_name} in BigQuery", + column_name + ); + + Ok(()) + } + + /// Alters the data type of the 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 = 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}" + ); + + let _ = self.query(QueryRequest::new(query)).await?; + + info!( + "altered column type {} in table {full_table_name} in BigQuery", + column_schema.name + ); + + Ok(()) + } + + /// 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, + 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 has_primary_key { + self.drop_primary_key(dataset_id, table_id).await?; + } + + if primary_columns.is_empty() { + return Ok(()); + } + + let columns = primary_columns + .iter() + .map(|column| Self::quote_identifier(&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?; + + info!("synced primary key for table {full_table_name} in BigQuery"); + + Ok(()) + } + + async fn has_primary_key( + &self, + dataset_id: &BigQueryDatasetId, + table_id: &BigQueryTableId, + ) -> EtlResult { + 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'", + ); + + 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?; + + info!("dropped primary key for table {full_table_name} in BigQuery"); + + Ok(()) + } + /// Checks whether a table exists in the BigQuery dataset. /// /// Returns `true` if the table exists, `false` otherwise. @@ -318,16 +498,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 ); @@ -400,16 +575,16 @@ 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, table_descriptor, - validated_rows, + validated_rows.into(), )) } @@ -428,8 +603,8 @@ 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::quote_identifier(&column_schema.name), Self::postgres_to_bigquery_type(&column_schema.typ) ); @@ -447,7 +622,7 @@ impl BigQueryClient { let identity_columns: Vec = column_schemas .iter() .filter(|s| s.primary) - .map(|c| format!("`{}`", c.name)) + .map(|c| Self::quote_identifier(&c.name)) .collect(); if identity_columns.is_empty() { @@ -525,7 +700,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, @@ -631,6 +806,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(_) => ( @@ -726,16 +911,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)] @@ -800,16 +999,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, 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, true); + 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, 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"); } @@ -817,16 +1015,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, 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, 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, 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); @@ -836,8 +1034,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, 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, ""); @@ -846,14 +1044,14 @@ 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, 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); 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)" ); } @@ -866,10 +1064,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, 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), ]; let descriptor = BigQueryClient::column_schemas_to_table_descriptor(&columns, true); @@ -884,7 +1082,7 @@ mod tests { )); assert!(matches!( descriptor.field_descriptors[0].mode, - ColumnMode::Required + ColumnMode::Nullable )); assert_eq!(descriptor.field_descriptors[1].name, "name"); @@ -904,7 +1102,7 @@ mod tests { )); assert!(matches!( descriptor.field_descriptors[2].mode, - ColumnMode::Required + ColumnMode::Nullable )); // Check array column @@ -949,12 +1147,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, 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); @@ -1006,8 +1204,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, true), + ColumnSchema::new("name".to_string(), Type::TEXT, -1, false), ]; // Simulate the query generation logic @@ -1015,7 +1213,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); } @@ -1026,13 +1224,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, true)]; // Simulate the query generation logic with staleness let full_table_name = format!("`{project_id}.{dataset_id}.{table_id}`"); @@ -1042,7 +1234,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-destinations/src/bigquery/core.rs b/etl-destinations/src/bigquery/core.rs index 0b3231fc0..f53f2ec96 100644 --- a/etl-destinations/src/bigquery/core.rs +++ b/etl-destinations/src/bigquery/core.rs @@ -2,24 +2,35 @@ 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, TableId, TableName, TableRow, generate_sequence_number}; +use etl::types::{ + Cell, Event, PgLsn, RelationChange, RelationEvent, SchemaVersion, TableId, TableName, TableRow, + VersionedTableSchema, generate_sequence_number, +}; 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 = 10; +/// Delay in milliseconds between retry attempts triggered by BigQuery schema mismatches. +const SCHEMA_MISMATCH_RETRY_DELAY_MS: u64 = 1000; /// Returns the [`BigQueryTableId`] for a supplied [`TableName`]. /// @@ -52,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. @@ -64,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 { @@ -304,35 +310,33 @@ 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)> { + 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. let mut inner = self.inner.lock().await; // 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) - .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 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); @@ -365,7 +369,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, @@ -373,6 +377,19 @@ where ) .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, + schema_version: Option, + use_cdc_sequence_column: bool, + ) -> EtlResult<(SequencedBigQueryTableId, Arc)> { + 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, use_cdc_sequence_column, @@ -462,6 +479,47 @@ where Ok(true) } + /// 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, + 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 @@ -471,9 +529,10 @@ where table_id: TableId, 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?; + // 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, None, false) + .await?; // Add CDC operation type to all rows (no lock needed). for table_row in table_rows.iter_mut() { @@ -501,10 +560,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,132 +578,337 @@ where Ok(()) } - /// 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 event_iter = events.into_iter().peekable(); + /// Persists the accumulated CDC batches for each table to BigQuery. + async fn process_table_events( + &self, + table_batches_by_id: HashMap<(TableId, SchemaVersion), Vec>, + ) -> EtlResult<()> { + if table_batches_by_id.is_empty() { + return Ok(()); + } - while event_iter.peek().is_some() { - let mut table_id_to_table_rows = HashMap::new(); + 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; + } - // 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(_)) { - break; - } + let (sequenced_bigquery_table_id, table_descriptor) = self + .prepare_table_for_streaming(&table_id, Some(schema_version), true) + .await?; - 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).or_default(); - table_rows.push(insert.table_row); - } - 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).or_default(); - table_rows.push(update.table_row); - } - Event::Delete(delete) => { - let Some((_, mut 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).or_default(); - table_rows.push(old_table_row); + 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.stream_with_schema_retry(table_batches).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(()) + } + + /// 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, 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]>>, + { + 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); } - _ => { - // Every other event type is currently not supported. - debug!("skipping unsupported event in BigQuery"); + + 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; } } + } + } - // Process accumulated events for each table. - if !table_id_to_table_rows.is_empty() { - let mut table_batches = Vec::with_capacity(table_id_to_table_rows.len()); + /// 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() + .contains(&ErrorKind::DestinationSchemaMismatch) + } - 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?; + 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" + ); - 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); + 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::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(); - if !table_batches.is_empty() { - let (bytes_sent, bytes_received) = self - .client - .stream_table_batches_concurrent(table_batches, self.max_concurrent_streams) + self.client + .add_column(&self.dataset_id, &sequenced_table_name, &column_schema) .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" + debug!( + table = %sequenced_table_name, + column = %column_name, + "added column in BigQuery" ); } - } + RelationChange::DropColumn(column_schema) => { + let column_name = column_schema.name.clone(); - // 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_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 { - truncate_table_ids.insert(TableId::new(table_id)); + 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" + ); } } } + } + + // 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, + table = %sequenced_table_name, + "applied relation changes in BigQuery" + ); + + Ok(()) + } + + #[allow(clippy::too_many_arguments)] + #[inline] + fn push_dml_statement( + table_batches: &mut HashMap<(TableId, SchemaVersion), Vec>, + start_lsn: PgLsn, + commit_lsn: PgLsn, + table_id: TableId, + mut table_row: TableRow, + schema_version: SchemaVersion, + operation_type: BigQueryOperationType, + ) -> EtlResult<()> { + // 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)); + + let key = (table_id, schema_version); + if let Some(rows) = table_batches.get_mut(&key) { + rows.push(table_row); + return Ok(()); + } + + // 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}" + ) + ); + } - if !truncate_table_ids.is_empty() { - self.process_truncate_for_table_ids(truncate_table_ids.into_iter(), true) - .await?; + table_batches.insert(key, vec![table_row]); + + Ok(()) + } + + /// Flushes the batch of events. + async fn flush_batch( + &self, + table_batches_by_id: &mut HashMap<(TableId, SchemaVersion), Vec>, + ) -> EtlResult<()> { + if table_batches_by_id.is_empty() { + return Ok(()); + } + + 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. + /// + /// Groups streaming operations (insert/update/delete) by table and processes them together, + /// 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_batches_by_id: HashMap<(TableId, SchemaVersion), Vec> = + HashMap::new(); + + // Process stream. + for event in events { + match event { + // DML events. + Event::Insert(insert) => { + Self::push_dml_statement( + &mut table_batches_by_id, + insert.start_lsn, + insert.commit_lsn, + insert.table_id, + insert.table_row, + insert.schema_version, + BigQueryOperationType::Upsert, + )?; + } + Event::Update(update) => { + Self::push_dml_statement( + &mut table_batches_by_id, + 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_batches_by_id, + delete.start_lsn, + delete.commit_lsn, + delete.table_id, + old_row, + delete.schema_version, + BigQueryOperationType::Delete, + )?; + } else { + warn!("the `DELETE` event has no row, so it was skipped"); + } + } + + // Batch breaker events. + Event::Relation(relation) => { + // Finish the current batch before applying schema change. + self.flush_batch(&mut table_batches_by_id).await?; + + // 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 table_batches_by_id).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. + other => { + debug!("skipping unsupported event {other:?} in BigQuery"); + } } } + // Flush any trailing DML. + self.flush_batch(&mut table_batches_by_id).await?; + Ok(()) } @@ -658,35 +919,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 in table_ids { - let table_schema = self.store.get_table_schema(&table_id).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 { - 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; - } + for (table_id, schema_version) in table_ids { + // 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 = @@ -723,11 +999,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?; @@ -793,7 +1068,7 @@ where } async fn truncate_table(&self, table_id: TableId) -> EtlResult<()> { - self.process_truncate_for_table_ids(iter::once(table_id), false) + self.process_truncate_for_table_ids(iter::once((table_id, None)), false) .await } @@ -906,7 +1181,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); } @@ -914,7 +1189,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); } @@ -922,7 +1197,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); } @@ -930,7 +1205,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); } @@ -961,14 +1236,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); } @@ -979,7 +1254,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] @@ -988,7 +1263,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] @@ -997,25 +1272,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] @@ -1142,7 +1417,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); } 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 f24ebd948..022432c05 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::{ + BigQueryOrder, BigQueryUser, 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; @@ -8,20 +12,18 @@ 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::types::{EventType, PgNumeric, PipelineId}; +use etl::test_utils::test_schema::{ + TableSelection, insert_mock_data, insert_users_data, setup_test_database_schema, +}; +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; use rand::random; use std::str::FromStr; use std::time::Duration; use tokio::time::sleep; -use crate::support::bigquery::{ - BigQueryOrder, BigQueryUser, NonNullableColsScalar, NullableColsArray, NullableColsScalar, - parse_bigquery_table_rows, setup_bigquery_connection, -}; - mod support; #[tokio::test(flavor = "multi_thread")] @@ -494,6 +496,285 @@ async fn table_truncate_with_batching() { ); } +#[tokio::test(flavor = "multi_thread")] +async fn table_schema_changes_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; + + 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), + ] + ); + + // 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(); + assert_eq!(users_rows.len(), 3); + + // We perform schema changes. + database + .alter_table( + test_table_name("users"), + &[ + TableModification::DropColumn { name: "year" }, + // TODO: data type change doesn't work rn. + // 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", &(2i32)], + ) + .await + .expect("Failed to insert users"); + + insert_event_notify.notified().await; + + let users_rows = bigquery_database + .query_table(database_schema.users_schema().name) + .await + .unwrap(); + assert_eq!(users_rows.len(), 4); + + 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, 2)]) + .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(); @@ -1580,6 +1861,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_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 diff --git a/etl-destinations/tests/iceberg_destination.rs b/etl-destinations/tests/iceberg_destination.rs index d5246ece8..58101fe49 100644 --- a/etl-destinations/tests/iceberg_destination.rs +++ b/etl-destinations/tests/iceberg_destination.rs @@ -344,8 +344,8 @@ async fn cdc_streaming() { TableRow { values: vec![ Cell::I64(1), - Cell::String("".to_string()), - Cell::I32(0), + Cell::Null, + Cell::Null, 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(), ], }, diff --git a/etl-destinations/tests/support/bigquery.rs b/etl-destinations/tests/support/bigquery.rs index e46ce409c..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. @@ -204,17 +203,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; + // }); + // }); } } @@ -256,6 +255,43 @@ 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_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 + .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(|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 +357,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 cd540d653..ec514bcc5 100644 --- a/etl-postgres/src/replication/schema.rs +++ b/etl-postgres/src/replication/schema.rs @@ -4,7 +4,12 @@ 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, VersionedTableSchema, +}; + +/// The initial schema version number. +const STARTING_SCHEMA_VERSION: u64 = 0; macro_rules! define_type_mappings { ( @@ -141,20 +146,35 @@ define_type_mappings! { pub async fn store_table_schema( pool: &PgPool, pipeline_id: i64, - table_schema: &TableSchema, -) -> Result<(), sqlx::Error> { + table_schema: TableSchema, +) -> Result { let mut tx = pool.begin().await?; - // Insert or update table schema record + // 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) + 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?; + + // 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 = 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#" - 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 +182,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_schema_version as i64) .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,23 +210,24 @@ pub async fn store_table_schema( tx.commit().await?; - Ok(()) + 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 ts.table_id, ts.schema_name, ts.table_name, + ts.schema_version, tc.column_name, tc.column_type, tc.type_modifier, @@ -222,24 +237,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), VersionedTableSchema> = 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(|| { + VersionedTableSchema::new( + table_id, + TableName::new(schema_name.clone(), table_name.clone()), + schema_version, + vec![], + ) + }); entry.add_column_schema(parse_column_schema(&row)); } @@ -303,14 +331,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/tokio/test_utils.rs b/etl-postgres/src/tokio/test_utils.rs index 6692e486b..7db5db3cf 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(); @@ -474,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-postgres/src/types/schema.rs b/etl-postgres/src/types/schema.rs index c66ab41f1..f357d9778 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 @@ -54,48 +57,31 @@ 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, } } - - /// 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. @@ -185,15 +171,60 @@ 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 pub name: TableName, + /// Monotonically increasing schema version. + pub version: SchemaVersion, /// The schemas of all columns in the table pub column_schemas: Vec, } +impl VersionedTableSchema { + pub fn new( + id: TableId, + name: TableName, + version: SchemaVersion, + column_schemas: Vec, + ) -> Self { + Self { + id, + name, + version, + column_schemas, + } + } + + /// 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 VersionedTableSchema { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +impl Ord for VersionedTableSchema { + fn cmp(&self, other: &Self) -> Ordering { + self.name + .cmp(&other.name) + .then(self.version.cmp(&other.version)) + } +} + +/// Draft version of [`VersionedTableSchema`] used before a schema version is assigned. +#[derive(Debug, Clone, Eq, PartialEq)] +pub struct TableSchema { + pub id: TableId, + pub name: TableName, + pub column_schemas: Vec, +} + impl TableSchema { pub fn new(id: TableId, name: TableName, column_schemas: Vec) -> Self { Self { @@ -203,7 +234,14 @@ impl TableSchema { } } - /// Adds a new column schema to this [`TableSchema`]. + 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); } @@ -215,29 +253,7 @@ impl TableSchema { 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 { - fn partial_cmp(&self, other: &Self) -> Option { - Some(self.cmp(other)) - } -} - -impl Ord for TableSchema { - fn cmp(&self, other: &Self) -> Ordering { - self.name.cmp(&other.name) + pub fn into_versioned(self, version: SchemaVersion) -> VersionedTableSchema { + VersionedTableSchema::new(self.id, self.name, version, self.column_schemas) } } 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-replicator/migrations/20251001000100_add_schema_versions.sql b/etl-replicator/migrations/20251001000100_add_schema_versions.sql new file mode 100644 index 000000000..e17f45ceb --- /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 table_schemas_pipeline_id_table_id_key; + +alter table etl.table_schemas + add constraint table_schemas_pipeline_id_table_id_schema_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-replicator/src/main.rs b/etl-replicator/src/main.rs index 48eb323ec..4448b9323 100644 --- a/etl-replicator/src/main.rs +++ b/etl-replicator/src/main.rs @@ -4,8 +4,6 @@ //! 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; @@ -13,6 +11,9 @@ use etl_telemetry::tracing::init_tracing_with_top_level_fields; use std::sync::Arc; use tracing::{error, info}; +use crate::config::load_replicator_config; +use crate::core::start_replicator_with_config; + mod config; mod core; mod migrations; diff --git a/etl/src/conversions/event.rs b/etl/src/conversions/event.rs index 81c9867df..d8630d160 100644 --- a/etl/src/conversions/event.rs +++ b/etl/src/conversions/event.rs @@ -1,17 +1,16 @@ use core::str; use etl_postgres::types::{ - ColumnSchema, TableId, TableName, TableSchema, 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::conversions::text::{default_value_for_type, parse_cell_from_postgres_text}; -use crate::error::{ErrorKind, EtlResult}; -use crate::store::schema::SchemaStore; +use crate::error::{ErrorKind, EtlError, EtlResult}; use crate::types::{ - BeginEvent, Cell, CommitEvent, DeleteEvent, InsertEvent, RelationEvent, TableRow, - TruncateEvent, UpdateEvent, + BeginEvent, Cell, CommitEvent, DeleteEvent, InsertEvent, TableRow, TruncateEvent, UpdateEvent, }; use crate::{bail, etl_error}; @@ -50,54 +49,58 @@ pub fn parse_event_from_commit_message( } } -/// Creates a [`RelationEvent`] from Postgres protocol data. +/// Creates a [`TableSchema`] 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 fn parse_event_from_relation_message( - start_lsn: PgLsn, - commit_lsn: PgLsn, +/// 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 { - let table_name = TableName::new( - relation_body.namespace()?.to_string(), - relation_body.name()?.to_string(), - ); - let column_schemas = relation_body +) -> 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 .columns() .iter() .map(build_column_schema) - .collect::, _>>()?; - let table_schema = TableSchema::new( - TableId::new(relation_body.rel_id()), - table_name, - column_schemas, - ); + .collect::, EtlError>>()?; - Ok(RelationEvent { - start_lsn, - commit_lsn, - table_schema, - }) + Ok(TableSchema::new(table_id, table_name, new_column_schemas)) } /// 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(), @@ -108,7 +111,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, }) } @@ -119,18 +123,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(); @@ -158,7 +156,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, }) @@ -170,18 +169,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(); @@ -200,7 +193,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, }) } @@ -213,54 +207,31 @@ pub fn parse_event_from_truncate_message( start_lsn: PgLsn, commit_lsn: PgLsn, truncate_body: &protocol::TruncateBody, - overridden_rel_ids: Vec, + table_ids: Vec<(TableId, SchemaVersion)>, ) -> TruncateEvent { TruncateEvent { start_lsn, commit_lsn, options: truncate_body.options(), - rel_ids: overridden_rel_ids, + table_ids, } } -/// 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 [`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( + let column_schema = ColumnSchema::new( column.name()?.to_string(), convert_type_oid_to_type(column.type_id() as u32), column.type_modifier(), - // We do not have access to this information, so we default it to `false`. - // TODO: figure out how to fill this value correctly or how to handle the missing value - // better. - false, // Currently 1 means that the column is part of the primary key. column.flags() == 1, - )) + ); + + Ok(column_schema) } /// Converts Postgres tuple data into a [`TableRow`] using column schemas. @@ -300,7 +271,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", @@ -311,7 +282,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/conversions/table_row.rs b/etl/src/conversions/table_row.rs index 3fd95c6df..8bd8985c7 100644 --- a/etl/src/conversions/table_row.rs +++ b/etl/src/conversions/table_row.rs @@ -164,14 +164,14 @@ mod tests { fn create_test_schema() -> 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, true), + ColumnSchema::new("name".to_string(), Type::TEXT, -1, false), + 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/src/destination/memory.rs b/etl/src/destination/memory.rs index 06976fd74..fb918a00e 100644 --- a/etl/src/destination/memory.rs +++ b/etl/src/destination/memory.rs @@ -93,12 +93,16 @@ 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; } @@ -123,7 +127,11 @@ 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/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/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/src/replication/apply.rs b/etl/src/replication/apply.rs index 0cfe52371..37eeaf362 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; +use etl_postgres::types::{TableId, VersionedTableSchema}; use futures::StreamExt; use metrics::histogram; use postgres_replication::protocol; @@ -19,8 +19,8 @@ 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_event_from_update_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}; @@ -31,9 +31,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::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 @@ -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. @@ -643,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, @@ -662,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?); @@ -704,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 } } @@ -761,6 +747,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(()) } @@ -926,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, hook).await + handle_truncate_message(state, start_lsn, truncate_body, hook, schema_store).await } LogicalReplicationMessage::Origin(_) => { debug!("received unsupported ORIGIN message"); @@ -1115,37 +1104,33 @@ 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") - ) - })?; + // 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?; - // 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, - ); + // 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?; - return Ok(HandleMessageResult::finish_batch_and_exclude_event(error)); + // 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_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(new_table_schema).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))) } @@ -1169,16 +1154,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_latest_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))) } @@ -1203,16 +1191,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_latest_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))) } @@ -1237,16 +1228,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_latest_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))) } @@ -1257,13 +1251,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, hook: &T, + schema_store: &S, ) -> EtlResult where + S: SchemaStore + Clone + Send + 'static, T: ApplyLoopHook, { let Some(remote_final_lsn) = state.remote_final_lsn else { @@ -1274,24 +1270,56 @@ 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()); - for &table_id in message.rel_ids().iter() { + 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(TableId::new(table_id), remote_final_lsn) + .should_apply_changes(table_id, remote_final_lsn) .await? { - rel_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_latest_table_schema(schema_store, table_id).await?; + table_ids.push((table_id, table_schema.version)); } } + // If nothing to apply, skip conversion entirely - if rel_ids.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, rel_ids); + let event = parse_event_from_truncate_message(start_lsn, remote_final_lsn, message, table_ids); Ok(HandleMessageResult::return_event(Event::Truncate(event))) } + +/// 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, +{ + 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 722437505..202e2fecd 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::{ColumnSchema, TableId, TableName}; +use etl_postgres::types::{TableSchema, convert_type_oid_to_type}; use pg_escape::{quote_identifier, quote_literal}; use postgres_replication::LogicalReplicationStream; use rustls::ClientConfig; @@ -712,25 +712,18 @@ 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)) } } 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. diff --git a/etl/src/store/both/memory.rs b/etl/src/store/both/memory.rs index 88a2af733..b434b876b 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, VersionedTableSchema}; +use std::collections::{BTreeMap, HashMap}; use std::sync::Arc; use tokio::sync::Mutex; @@ -10,22 +10,20 @@ 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 { - /// 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_schemas: HashMap>, - /// Mapping from table IDs to human-readable table names for easier debugging - /// and logging. These mappings are established during schema discovery. + /// Table schema definitions for each table. Each schema has multiple versions which are created + /// when new schema changes are detected. + table_schemas: HashMap>>, + /// Mappings between source and destination tables. table_mappings: HashMap, } @@ -172,31 +170,63 @@ 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).cloned()) + Ok(inner + .table_schemas + .get(table_id) + .and_then(|table_schemas| table_schemas.get(&version).cloned())) } - async fn get_table_schemas(&self) -> EtlResult>> { + async fn get_latest_table_schema( + &self, + table_id: &TableId, + ) -> EtlResult>> { let inner = self.inner.lock().await; - Ok(inner.table_schemas.values().cloned().collect()) + 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 { let inner = self.inner.lock().await; - Ok(inner.table_schemas.len()) + Ok(inner + .table_schemas + .values() + .map(|table_schemas| table_schemas.len()) + .sum()) } - async fn store_table_schema(&self, table_schema: TableSchema) -> EtlResult<()> { + async fn store_table_schema( + &self, + table_schema: TableSchema, + ) -> EtlResult> { let mut inner = self.inner.lock().await; - inner + let table_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 = table_schemas + .keys() + .next_back() + .map(|version| version + 1) + .unwrap_or(STARTING_SCHEMA_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 427e25cd3..030c6d2e3 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, VersionedTableSchema}; use metrics::gauge; use sqlx::PgPool; +use std::{ + collections::{BTreeMap, HashMap}, + sync::Arc, +}; use tokio::sync::Mutex; use tracing::{debug, info}; @@ -17,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 @@ -136,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, } @@ -197,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. @@ -210,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() } } @@ -328,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) @@ -371,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) @@ -506,21 +518,30 @@ 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).cloned()) + Ok(inner + .table_schemas + .get(table_id) + .and_then(|schemas| schemas.get(&version).cloned())) } - /// 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>> { + /// Retrieves the latest table schema for a table. + async fn get_latest_table_schema( + &self, + table_id: &TableId, + ) -> EtlResult>> { let inner = self.inner.lock().await; - Ok(inner.table_schemas.values().cloned().collect()) + Ok(inner + .table_schemas + .get(table_id) + .and_then(|schemas| schemas.iter().next_back().map(|(_, schema)| schema.clone()))) } /// Loads table schemas from Postgres into memory cache. @@ -552,7 +573,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!( @@ -568,27 +591,36 @@ 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: TableSchema, + ) -> 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 stored_schema = Arc::new(stored_schema); inner .table_schemas - .insert(table_schema.id, Arc::new(table_schema)); + .entry(stored_schema.id) + .or_insert_with(BTreeMap::new) + .insert(stored_schema.version, stored_schema.clone()); - Ok(()) + Ok(stored_schema) } } diff --git a/etl/src/store/schema/base.rs b/etl/src/store/schema/base.rs index 2a52126b9..c14032fef 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, VersionedTableSchema}; use std::sync::Arc; use crate::error::EtlResult; @@ -10,18 +10,20 @@ 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( &self, table_id: &TableId, - ) -> impl Future>>> + Send; + version: SchemaVersion, + ) -> impl Future>>> + Send; - /// Returns all table schemas from the cache. - /// - /// Does not read from the persistent store. - fn get_table_schemas(&self) -> 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; /// Loads table schemas from the persistent state into the cache. /// @@ -32,5 +34,5 @@ pub trait SchemaStore { fn store_table_schema( &self, table_schema: TableSchema, - ) -> impl Future> + Send; + ) -> impl Future>> + Send; } diff --git a/etl/src/test_utils/notify.rs b/etl/src/test_utils/notify.rs index 6793066ea..1efdac168 100644 --- a/etl/src/test_utils/notify.rs +++ b/etl/src/test_utils/notify.rs @@ -1,6 +1,9 @@ -use std::{collections::HashMap, fmt, sync::Arc}; - -use etl_postgres::types::{TableId, TableSchema}; +use etl_postgres::types::{SchemaVersion, TableId, TableSchema, VersionedTableSchema}; +use std::{ + collections::{BTreeMap, HashMap}, + fmt, + sync::Arc, +}; use tokio::sync::{Notify, RwLock}; use crate::error::{ErrorKind, EtlResult}; @@ -29,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, @@ -103,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() } @@ -290,30 +298,60 @@ 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).cloned()) + Ok(inner + .table_schemas + .get(table_id) + .and_then(|schemas| schemas.get(&version).cloned())) } - async fn get_table_schemas(&self) -> EtlResult>> { + async fn get_latest_table_schema( + &self, + table_id: &TableId, + ) -> EtlResult>> { let inner = self.inner.read().await; - Ok(inner.table_schemas.values().cloned().collect()) + Ok(inner + .table_schemas + .get(table_id) + .and_then(|schemas| schemas.iter().next_back().map(|(_, schema)| schema.clone()))) } 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: TableSchema, + ) -> 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_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 a2c6d46ad..d3a7fae54 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 etl_postgres::types::{ColumnSchema, TableId, TableName, TableSchema, VersionedTableSchema}; use std::collections::HashMap; +/// Return the names of the column schema. +pub fn column_schema_names(table_schema: &VersionedTableSchema) -> 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/src/test_utils/test_destination_wrapper.rs b/etl/src/test_utils/test_destination_wrapper.rs index 58e107d07..ded01a4bf 100644 --- a/etl/src/test_utils/test_destination_wrapper.rs +++ b/etl/src/test_utils/test_destination_wrapper.rs @@ -168,12 +168,16 @@ 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, _)| 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/test_utils/test_schema.rs b/etl/src/test_utils/test_schema.rs index 8faa6449b..ea53dacc6 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), ], )); } @@ -236,9 +218,9 @@ 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 + left.options == right.options && left.table_ids == right.table_ids } (Event::Unsupported, Event::Unsupported) => true, _ => false, // Different event types @@ -264,6 +246,7 @@ pub fn build_expected_users_inserts( Cell::I32(age), ], }, + schema_version: 0, })); starting_id += 1; @@ -287,6 +270,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 144823ac8..848ad9dbc 100644 --- a/etl/src/types/event.rs +++ b/etl/src/types/event.rs @@ -1,5 +1,8 @@ -use etl_postgres::types::{TableId, TableSchema}; +use etl_postgres::types::{ColumnSchema, SchemaVersion, TableId, TableName, VersionedTableSchema}; +use std::collections::HashSet; use std::fmt; +use std::hash::Hash; +use std::sync::Arc; use tokio_postgres::types::PgLsn; use crate::types::TableRow; @@ -40,6 +43,25 @@ pub struct CommitEvent { pub timestamp: i64, } +/// 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. + 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), +} + /// Table schema definition event from Postgres logical replication. /// /// [`RelationEvent`] provides schema information for tables involved in replication. @@ -51,8 +73,70 @@ 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, + /// 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. 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 + .column_schemas + .iter() + .cloned() + .map(IndexedColumnSchema) + .collect::>(); + + // We process all the changes that we want to dispatch to the destination. + 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 != 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( + 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. @@ -67,6 +151,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, } @@ -84,6 +170,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. @@ -106,6 +194,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`) @@ -128,7 +218,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 table_ids: Vec<(TableId, SchemaVersion)>, } /// Represents a single replication event from Postgres logical replication. @@ -175,14 +265,8 @@ 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::Truncate(event) => { - let Some(_) = event.rel_ids.iter().find(|&&id| table_id.0 == id) else { - return false; - }; - - true - } + Event::Relation(relation_event) => relation_event.table_id == *table_id, + Event::Truncate(event) => event.table_ids.iter().any(|(id, _)| id == table_id), _ => false, } } @@ -248,3 +332,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/failpoints_pipeline.rs b/etl/tests/failpoints_pipeline.rs index 0f597b64d..5361b2596 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,13 +205,14 @@ 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 .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) ); } @@ -268,12 +269,13 @@ 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 .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/pipeline.rs b/etl/tests/pipeline.rs index ee6230460..a80c0f66d 100644 --- a/etl/tests/pipeline.rs +++ b/etl/tests/pipeline.rs @@ -7,16 +7,18 @@ 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, 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_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; 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; @@ -81,19 +83,19 @@ 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 .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. @@ -127,6 +129,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; @@ -370,7 +389,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)); @@ -396,7 +415,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)); @@ -489,6 +508,385 @@ 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::Both).await; + + // 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()); + + // 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; + 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(), 2); + 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"] + ); + 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 + .alter_table( + test_table_name("users"), + &[ + TableModification::AddColumn { + name: "year", + params: "integer", + }, + TableModification::RenameColumn { + name: "age", + new_name: "new_age", + }, + ], + ) + .await + .unwrap(); + database + .alter_table( + test_table_name("orders"), + &[TableModification::AddColumn { + name: "summary", + params: "text", + }], + ) + .await + .unwrap(); + + // Register notifications for the inserts. + 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_2", &(2i32), &(2025i32)], + ) + .await + .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(), 2); + 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"] + ); + 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; + 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 orders_inserts = grouped_events + .get(&(EventType::Insert, database_schema.orders_schema().id)) + .unwrap(); + assert_eq!(orders_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(); + 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 inserts. + let insert_event_notify = destination + .wait_for_events_count(vec![(EventType::Insert, 4)]) + .await; + + // We insert data. + database + .insert_values( + database_schema.users_schema().name.clone(), + &["name", "new_age"], + &[&"user_3", &(2f64)], + ) + .await + .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(), 2); + 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"] + ); + 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; + 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); + 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(); +} + +#[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", &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 + .get(&(EventType::Insert, database_schema.users_schema().id)) + .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) + .unwrap(); + assert_eq!(users_table_schema.version, 1); + assert_eq!(users_table_schema.name, test_table_name("customers")); + + // 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(); +} + #[tokio::test(flavor = "multi_thread")] async fn table_copy_and_sync_streams_new_data() { init_test_tracing(); @@ -794,146 +1192,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", - data_type: "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(); diff --git a/etl/tests/postgres_store.rs b/etl/tests/postgres_store.rs index 1038c591a..101ca0a6f 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, true), + ColumnSchema::new("name".to_string(), PgType::TEXT, -1, false), + 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, true), + ColumnSchema::new("description".to_string(), PgType::VARCHAR, 255, false), ]; TableSchema::new(table_id, table_name, columns) @@ -183,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()); @@ -193,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(); @@ -215,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")] @@ -231,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); @@ -253,17 +250,17 @@ 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 - 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(); - let all_schemas = store.get_table_schemas().await.unwrap(); + let all_schemas = store.get_all_table_schemas().await; assert_eq!(all_schemas.len(), 2); } @@ -275,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 - store - .store_table_schema(table_schema1.clone()) + let stored_schema_1 = store + .store_table_schema(table_schema_1.clone()) .await .unwrap(); - store - .store_table_schema(table_schema2.clone()) + let stored_schema_2 = store + .store_table_schema(table_schema_2.clone()) .await .unwrap(); @@ -292,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 @@ -300,20 +297,26 @@ 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_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).await.unwrap(); + let schema2 = new_store + .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")] @@ -327,28 +330,39 @@ 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(); + // 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(), PgType::TIMESTAMPTZ, -1, - true, false, )); // 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).await.unwrap(); + let schema = store + .get_table_schema(&table_schema.id, stored_schema_updated.version) + .await + .unwrap(); assert!(schema.is_some()); let schema = schema.unwrap(); assert_eq!(schema.column_schemas.len(), 4); // Original 3 + 1 new column @@ -366,55 +380,61 @@ 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 store1 = PostgresStore::new(pipeline_id1, database.config.clone()); - let store2 = 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; - 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(); + let table_schema_1 = create_sample_table_schema(); + let table_schema_2 = create_another_table_schema(); - store1 - .store_table_schema(table_schema1.clone()) + let stored_schema_1 = store_1 + .store_table_schema(table_schema_1.clone()) .await .unwrap(); - store2 - .store_table_schema(table_schema2.clone()) + let stored_schema_2 = store_2 + .store_table_schema(table_schema_2.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_schema_1.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_schema_2.id][&stored_schema_2.version], + stored_schema_2 + ); } #[tokio::test(flavor = "multi_thread")] @@ -542,11 +562,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(); @@ -554,33 +574,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())); } @@ -632,37 +652,37 @@ 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 store1 = PostgresStore::new(pipeline_id1, database.config.clone()); - let store2 = 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); // 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_id_1, 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())); } @@ -691,11 +711,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(); @@ -717,7 +737,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) @@ -737,7 +763,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) @@ -754,7 +786,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) @@ -779,7 +817,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() @@ -802,7 +840,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() @@ -835,7 +873,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 @@ -846,7 +890,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 @@ -862,6 +906,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()); } 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..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")]) @@ -261,7 +243,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 @@ -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();