From 551ddf1378a861ee2fa3cd971324c37c2ba538fc Mon Sep 17 00:00:00 2001 From: ttatsato Date: Fri, 15 May 2026 15:07:35 +0900 Subject: [PATCH 01/29] ref(api): deprecate POST update-publication endpoint --- .../etl-api/src/routes/sources/publications.rs | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/crates/etl-api/src/routes/sources/publications.rs b/crates/etl-api/src/routes/sources/publications.rs index 86e9a5da6..fb8c804ea 100644 --- a/crates/etl-api/src/routes/sources/publications.rs +++ b/crates/etl-api/src/routes/sources/publications.rs @@ -186,8 +186,12 @@ pub(crate) async fn read_publication( } #[utoipa::path( - summary = "Update a publication", - description = "Replaces the publication's table list on the given source.", + summary = "Update a publication (deprecated)", + description = "Replaces the publication's table list on the given source.\n\n\ + **Deprecated**: This endpoint uses `POST` for full replacement, which is \ + semantically misleading because the request appears additive but actually \ + replaces the entire table list. A replacement is being designed. See \ + https://github.com/supabase/etl/issues/459.", tag = "Publications", request_body = UpdatePublicationRequest, params( @@ -226,7 +230,13 @@ pub(crate) async fn update_publication( let publication = Publication { name: publication_name, tables: publication.tables }; data::publications::update_publication(&publication, &source_pool).await?; - Ok(HttpResponse::Ok().finish()) + Ok(HttpResponse::Ok() + .insert_header(("Deprecation", "true")) + .insert_header(( + "Link", + "; rel=\"deprecation\"", + )) + .finish()) } #[utoipa::path( From 0c4eaeb7f3b3119532b5fdd7a1c2c921bcff770b Mon Sep 17 00:00:00 2001 From: ttatsato Date: Fri, 15 May 2026 15:37:37 +0900 Subject: [PATCH 02/29] fix(api): use RFC 9745 sf-date for Deprecation header --- crates/etl-api/src/routes/sources/publications.rs | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/crates/etl-api/src/routes/sources/publications.rs b/crates/etl-api/src/routes/sources/publications.rs index fb8c804ea..bf3a37a6c 100644 --- a/crates/etl-api/src/routes/sources/publications.rs +++ b/crates/etl-api/src/routes/sources/publications.rs @@ -230,8 +230,11 @@ pub(crate) async fn update_publication( let publication = Publication { name: publication_name, tables: publication.tables }; data::publications::update_publication(&publication, &source_pool).await?; + // RFC 9745 §2.1 requires the Deprecation field to be an sf-date. + // The value is the Unix timestamp for 2026-05-15 00:00 UTC, the date + // this deprecation took effect. Ok(HttpResponse::Ok() - .insert_header(("Deprecation", "true")) + .insert_header(("Deprecation", "@1778803200")) .insert_header(( "Link", "; rel=\"deprecation\"", From aeaf2b83b623f06eb2e335baa4ca3c5ac6882346 Mon Sep 17 00:00:00 2001 From: Jordan McQueen Date: Wed, 13 May 2026 15:23:42 +0900 Subject: [PATCH 03/29] test(clickhouse): use wait_for_events_count for streaming tests (#739) Wrap the test destination in TestDestinationWrapper and replace the SELECT polling loops with wait_for_events_count. Touches 9 streaming tests; removes 5 polling helpers. --- .../tests/clickhouse/pipeline.rs | 286 +++++------------- 1 file changed, 78 insertions(+), 208 deletions(-) diff --git a/crates/etl-destinations/tests/clickhouse/pipeline.rs b/crates/etl-destinations/tests/clickhouse/pipeline.rs index e5265236e..0e5e4abd5 100644 --- a/crates/etl-destinations/tests/clickhouse/pipeline.rs +++ b/crates/etl-destinations/tests/clickhouse/pipeline.rs @@ -1,4 +1,4 @@ -use std::{sync::Once, time::Duration}; +use std::sync::Once; use etl::{ state::table::TableReplicationPhaseType, @@ -7,21 +7,20 @@ use etl::{ database::{spawn_source_database, test_table_name}, notifying_store::NotifyingStore, pipeline::create_pipeline, + test_destination_wrapper::TestDestinationWrapper, }, - types::PipelineId, + types::{EventType, PipelineId}, }; use etl_destinations::clickhouse::{ ClickHouseInserterConfig, client::ClickHouseClient, test_utils::{ - ClickHouseTestDatabase, get_clickhouse_password, get_clickhouse_url, get_clickhouse_user, - setup_clickhouse_database, + get_clickhouse_password, get_clickhouse_url, get_clickhouse_user, setup_clickhouse_database, }, }; use etl_postgres::tokio::test_utils::TableModification; use etl_telemetry::tracing::init_test_tracing; use rand::random; -use tokio::time::sleep; use url::Url; use crate::support::clickhouse::{AllTypesRow, BoundaryValuesRow, DateBoundariesRow}; @@ -106,108 +105,6 @@ const DATE_2024_01_15_DAYS: i32 = 19737; /// Microseconds from epoch for `2024-01-15 12:00:00 UTC`. const TS_2024_01_15_12_00_US: i64 = 1_705_320_000_000_000; -/// Waits until ClickHouse returns at least `expected_rows` from -/// `UPDATE_FLOW_SELECT`. -async fn wait_for_update_flow_rows( - clickhouse_db: &ClickHouseTestDatabase, - expected_rows: usize, -) -> Vec { - let mut rows: Vec = Vec::with_capacity(expected_rows); - for _ in 0..50 { - rows = clickhouse_db.query(UPDATE_FLOW_SELECT).await; - if rows.len() >= expected_rows { - return rows; - } - sleep(Duration::from_millis(100)).await; - } - - panic!( - "timed out waiting for clickhouse update_flow rows: got {} of {}", - rows.len(), - expected_rows, - ); -} - -/// Waits until ClickHouse returns at least `expected_rows` from -/// `DELETE_FLOW_SELECT`. -async fn wait_for_delete_flow_rows( - clickhouse_db: &ClickHouseTestDatabase, - expected_rows: usize, -) -> Vec { - let mut rows: Vec = Vec::with_capacity(expected_rows); - for _ in 0..50 { - rows = clickhouse_db.query(DELETE_FLOW_SELECT).await; - if rows.len() >= expected_rows { - return rows; - } - sleep(Duration::from_millis(100)).await; - } - - panic!( - "timed out waiting for clickhouse delete_flow rows: got {} of {}", - rows.len(), - expected_rows, - ); -} - -/// Waits until ClickHouse returns at least `expected_rows` from -/// `RESTART_FLOW_SELECT`. -async fn wait_for_restart_flow_rows( - clickhouse_db: &ClickHouseTestDatabase, - expected_rows: usize, -) -> Vec { - let mut rows: Vec = Vec::with_capacity(expected_rows); - for _ in 0..50 { - rows = clickhouse_db.query(RESTART_FLOW_SELECT).await; - if rows.len() >= expected_rows { - return rows; - } - sleep(Duration::from_millis(100)).await; - } - - panic!( - "timed out waiting for clickhouse restart_flow rows: got {} of {}", - rows.len(), - expected_rows, - ); -} - -/// Waits until ClickHouse returns exactly zero rows from -/// `TRUNCATE_FLOW_SELECT`. -async fn wait_for_truncate_flow_empty(clickhouse_db: &ClickHouseTestDatabase) { - for _ in 0..50 { - let rows: Vec = clickhouse_db.query(TRUNCATE_FLOW_SELECT).await; - if rows.is_empty() { - return; - } - sleep(Duration::from_millis(100)).await; - } - - panic!("timed out waiting for clickhouse truncate_flow table to become empty"); -} - -/// Waits until ClickHouse returns at least `expected_rows` from -/// `TRUNCATE_FLOW_SELECT`. -async fn wait_for_truncate_flow_rows( - clickhouse_db: &ClickHouseTestDatabase, - expected_rows: usize, -) -> Vec { - let mut rows: Vec = Vec::with_capacity(expected_rows); - for _ in 0..50 { - rows = clickhouse_db.query(TRUNCATE_FLOW_SELECT).await; - if rows.len() >= expected_rows { - return rows; - } - sleep(Duration::from_millis(100)).await; - } - - panic!( - "timed out waiting for clickhouse truncate_flow rows: got {} of {}", - rows.len(), - expected_rows, - ); -} - /// Tests that all Postgres column types (including nullable arrays) round-trip /// correctly through the ClickHouse RowBinary encoding. /// @@ -489,7 +386,7 @@ async fn updates_are_streamed_to_clickhouse() { let clickhouse_db = setup_clickhouse_database().await; let store = NotifyingStore::new(); let pipeline_id: PipelineId = random(); - let destination = clickhouse_db.build_destination(store.clone()); + let destination = TestDestinationWrapper::wrap(clickhouse_db.build_destination(store.clone())); let table_ready = store.notify_on_table_state_type(table_id, TableReplicationPhaseType::Ready).await; @@ -499,12 +396,14 @@ async fn updates_are_streamed_to_clickhouse() { pipeline_id, publication_name.to_owned(), store, - destination, + destination.clone(), ); pipeline.start().await.unwrap(); table_ready.notified().await; + let event_notify = destination.wait_for_events_count(vec![(EventType::Update, 1)]).await; + database .run_sql(&format!( "UPDATE {} SET value = 'after' WHERE id = 1", @@ -513,7 +412,9 @@ async fn updates_are_streamed_to_clickhouse() { .await .expect("Failed to update update_flow row"); - let rows = wait_for_update_flow_rows(&clickhouse_db, 2).await; + event_notify.notified().await; + + let rows: Vec = clickhouse_db.query(UPDATE_FLOW_SELECT).await; pipeline.shutdown_and_wait().await.unwrap(); @@ -890,7 +791,7 @@ async fn deletes_are_streamed_to_clickhouse() { let clickhouse_db = setup_clickhouse_database().await; let store = NotifyingStore::new(); let pipeline_id: PipelineId = random(); - let destination = clickhouse_db.build_destination(store.clone()); + let destination = TestDestinationWrapper::wrap(clickhouse_db.build_destination(store.clone())); let table_ready = store.notify_on_table_state_type(table_id, TableReplicationPhaseType::Ready).await; @@ -900,18 +801,22 @@ async fn deletes_are_streamed_to_clickhouse() { pipeline_id, publication_name.to_owned(), store, - destination, + destination.clone(), ); pipeline.start().await.unwrap(); table_ready.notified().await; + let event_notify = destination.wait_for_events_count(vec![(EventType::Delete, 1)]).await; + database .run_sql(&format!("DELETE FROM {} WHERE id = 2", table_name.as_quoted_identifier(),)) .await .expect("Failed to delete delete_flow row"); - let rows = wait_for_delete_flow_rows(&clickhouse_db, 3).await; + event_notify.notified().await; + + let rows: Vec = clickhouse_db.query(DELETE_FLOW_SELECT).await; pipeline.shutdown_and_wait().await.unwrap(); @@ -993,7 +898,7 @@ async fn pipeline_restart_resumes_streaming() { let clickhouse_db = setup_clickhouse_database().await; let store = NotifyingStore::new(); let pipeline_id: PipelineId = random(); - let destination = clickhouse_db.build_destination(store.clone()); + let destination = TestDestinationWrapper::wrap(clickhouse_db.build_destination(store.clone())); let table_ready = store.notify_on_table_state_type(table_id, TableReplicationPhaseType::Ready).await; @@ -1017,18 +922,20 @@ async fn pipeline_restart_resumes_streaming() { assert_eq!(rows[0].value, "before_restart"); // --- WHEN: rebuild destination and pipeline, then stream a new insert --- - let destination = clickhouse_db.build_destination(store.clone()); + let destination = TestDestinationWrapper::wrap(clickhouse_db.build_destination(store.clone())); let mut pipeline = create_pipeline( &database.config, pipeline_id, publication_name.to_owned(), store, - destination, + destination.clone(), ); pipeline.start().await.unwrap(); + let event_notify = destination.wait_for_events_count(vec![(EventType::Insert, 1)]).await; + database .run_sql(&format!( "INSERT INTO {} (value) VALUES ('after_restart')", @@ -1037,7 +944,9 @@ async fn pipeline_restart_resumes_streaming() { .await .expect("Failed to insert post-restart row"); - let rows = wait_for_restart_flow_rows(&clickhouse_db, 2).await; + event_notify.notified().await; + + let rows: Vec = clickhouse_db.query(RESTART_FLOW_SELECT).await; pipeline.shutdown_and_wait().await.unwrap(); @@ -1112,7 +1021,7 @@ async fn truncate_clears_table_and_accepts_new_inserts() { let clickhouse_db = setup_clickhouse_database().await; let store = NotifyingStore::new(); let pipeline_id: PipelineId = random(); - let destination = clickhouse_db.build_destination(store.clone()); + let destination = TestDestinationWrapper::wrap(clickhouse_db.build_destination(store.clone())); let table_ready = store.notify_on_table_state_type(table_id, TableReplicationPhaseType::Ready).await; @@ -1122,7 +1031,7 @@ async fn truncate_clears_table_and_accepts_new_inserts() { pipeline_id, publication_name.to_owned(), store, - destination, + destination.clone(), ); pipeline.start().await.unwrap(); @@ -1133,12 +1042,19 @@ async fn truncate_clears_table_and_accepts_new_inserts() { assert_eq!(rows.len(), 2, "table copy should produce two rows"); // --- WHEN: truncate, then insert a new row --- + let truncate_notify = destination.wait_for_events_count(vec![(EventType::Truncate, 1)]).await; + database .truncate_table(table_name.clone()) .await .expect("Failed to truncate table in Postgres"); - wait_for_truncate_flow_empty(&clickhouse_db).await; + truncate_notify.notified().await; + + let rows: Vec = clickhouse_db.query(TRUNCATE_FLOW_SELECT).await; + assert!(rows.is_empty(), "table should be empty after truncate"); + + let insert_notify = destination.wait_for_events_count(vec![(EventType::Insert, 1)]).await; database .run_sql(&format!( @@ -1148,7 +1064,9 @@ async fn truncate_clears_table_and_accepts_new_inserts() { .await .expect("Failed to insert post-truncate row"); - let rows = wait_for_truncate_flow_rows(&clickhouse_db, 1).await; + insert_notify.notified().await; + + let rows: Vec = clickhouse_db.query(TRUNCATE_FLOW_SELECT).await; pipeline.shutdown_and_wait().await.unwrap(); @@ -1322,7 +1240,7 @@ async fn multiple_tables_receive_independent_writes() { let clickhouse_db = setup_clickhouse_database().await; let store = NotifyingStore::new(); let pipeline_id: PipelineId = random(); - let destination = clickhouse_db.build_destination(store.clone()); + let destination = TestDestinationWrapper::wrap(clickhouse_db.build_destination(store.clone())); let table_a_ready = store.notify_on_table_state_type(table_a_id, TableReplicationPhaseType::Ready).await; @@ -1334,12 +1252,14 @@ async fn multiple_tables_receive_independent_writes() { pipeline_id, publication_name.to_owned(), store, - destination, + destination.clone(), ); pipeline.start().await.unwrap(); tokio::join!(table_a_ready.notified(), table_b_ready.notified()); + let event_notify = destination.wait_for_events_count(vec![(EventType::Insert, 2)]).await; + // --- WHEN: insert one row into each table --- database .run_sql(&format!( @@ -1357,6 +1277,8 @@ async fn multiple_tables_receive_independent_writes() { .await .expect("Failed to insert streamed row into multi_b"); + event_notify.notified().await; + let select_a = concat!( "SELECT id, value, cdc_operation, cdc_lsn ", "FROM \"test_multi__a\" ", @@ -1368,23 +1290,8 @@ async fn multiple_tables_receive_independent_writes() { "ORDER BY id, cdc_lsn", ); - // Poll until both tables have 2 rows. - let mut rows_a: Vec = Vec::with_capacity(2); - let mut rows_b: Vec = Vec::with_capacity(2); - for _ in 0..50 { - rows_a = clickhouse_db.query(select_a).await; - rows_b = clickhouse_db.query(select_b).await; - if rows_a.len() >= 2 && rows_b.len() >= 2 { - break; - } - sleep(Duration::from_millis(100)).await; - } - assert!( - rows_a.len() >= 2 && rows_b.len() >= 2, - "timed out: multi_a has {} rows, multi_b has {} rows", - rows_a.len(), - rows_b.len() - ); + let rows_a: Vec = clickhouse_db.query(select_a).await; + let rows_b: Vec = clickhouse_db.query(select_b).await; pipeline.shutdown_and_wait().await.unwrap(); @@ -1473,7 +1380,7 @@ async fn sequential_transactions_preserve_commit_order() { let clickhouse_db = setup_clickhouse_database().await; let store = NotifyingStore::new(); let pipeline_id: PipelineId = random(); - let destination = clickhouse_db.build_destination(store.clone()); + let destination = TestDestinationWrapper::wrap(clickhouse_db.build_destination(store.clone())); let table_ready = store.notify_on_table_state_type(table_id, TableReplicationPhaseType::Ready).await; @@ -1483,12 +1390,14 @@ async fn sequential_transactions_preserve_commit_order() { pipeline_id, publication_name.to_owned(), store, - destination, + destination.clone(), ); pipeline.start().await.unwrap(); table_ready.notified().await; + let event_notify = destination.wait_for_events_count(vec![(EventType::Update, 2)]).await; + // --- WHEN: two transactions commit sequentially on separate connections --- let tx_a = database_1.begin_transaction().await; tx_a.run_sql(&format!( @@ -1508,16 +1417,9 @@ async fn sequential_transactions_preserve_commit_order() { .expect("Failed to execute update_b"); tx_b.commit_transaction().await; - // Poll until all three rows arrive. - let mut rows: Vec = Vec::with_capacity(3); - for _ in 0..50 { - rows = clickhouse_db.query(TX_ORDER_SELECT).await; - if rows.len() >= 3 { - break; - } - sleep(Duration::from_millis(100)).await; - } - assert!(rows.len() >= 3, "timed out waiting for tx_order rows: got {} of 3", rows.len()); + event_notify.notified().await; + + let rows: Vec = clickhouse_db.query(TX_ORDER_SELECT).await; pipeline.shutdown_and_wait().await.unwrap(); @@ -1682,7 +1584,7 @@ async fn delete_with_default_replica_identity() { let clickhouse_db = setup_clickhouse_database().await; let store = NotifyingStore::new(); let pipeline_id: PipelineId = random(); - let destination = clickhouse_db.build_destination(store.clone()); + let destination = TestDestinationWrapper::wrap(clickhouse_db.build_destination(store.clone())); let table_ready = store.notify_on_table_state_type(table_id, TableReplicationPhaseType::Ready).await; @@ -1692,12 +1594,16 @@ async fn delete_with_default_replica_identity() { pipeline_id, publication_name.to_owned(), store, - destination, + destination.clone(), ); pipeline.start().await.unwrap(); table_ready.notified().await; + let event_notify = destination + .wait_for_events_count(vec![(EventType::Delete, 1), (EventType::Insert, 1)]) + .await; + // --- WHEN: delete id=2, insert id=3 --- database .run_sql(&format!("DELETE FROM {} WHERE id = 2", table_name.as_quoted_identifier())) @@ -1724,20 +1630,10 @@ async fn delete_with_default_replica_identity() { .await .expect("Failed to insert post-delete row"); - // Poll for 4 rows: 2 copied INSERTs + DELETE tombstone + new INSERT. - let mut rows: Vec = Vec::new(); - for _ in 0..50 { - rows = clickhouse_db.query(DEFAULT_IDENTITY_DELETE_SELECT).await; - if rows.len() >= 4 { - break; - } - sleep(Duration::from_millis(100)).await; - } - assert!( - rows.len() >= 4, - "timed out waiting for default_identity_delete rows: got {} of 4", - rows.len() - ); + event_notify.notified().await; + + let rows: Vec = + clickhouse_db.query(DEFAULT_IDENTITY_DELETE_SELECT).await; pipeline.shutdown_and_wait().await.unwrap(); @@ -1998,7 +1894,7 @@ async fn schema_change_add_column() { let clickhouse_db = setup_clickhouse_database().await; let store = NotifyingStore::new(); let pipeline_id: PipelineId = random(); - let destination = clickhouse_db.build_destination(store.clone()); + let destination = TestDestinationWrapper::wrap(clickhouse_db.build_destination(store.clone())); let table_ready = store.notify_on_table_state_type(table_id, TableReplicationPhaseType::Ready).await; @@ -2008,7 +1904,7 @@ async fn schema_change_add_column() { pipeline_id, publication_name.to_owned(), store.clone(), - destination, + destination.clone(), ); pipeline.start().await.unwrap(); @@ -2040,6 +1936,8 @@ async fn schema_change_add_column() { .await .unwrap(); + let event_notify = destination.wait_for_events_count(vec![(EventType::Insert, 1)]).await; + database .run_sql(&format!( "INSERT INTO {} (name, age, email, score) VALUES ('Bob', 30, 'bob@example.com', 7)", @@ -2048,30 +1946,14 @@ async fn schema_change_add_column() { .await .expect("Failed to insert Bob"); - // Poll until Bob's row arrives (2 rows total = Alice from copy + Bob from - // streaming). + event_notify.notified().await; + let select = concat!( "SELECT id, name, age, email, score, cdc_operation ", "FROM \"test_schema__add__col\" ", "ORDER BY id", ); - let mut rows: Vec = Vec::new(); - for _ in 0..50 { - // The SELECT will fail if the email column doesn't exist yet, so - // catch errors and retry. - if let Ok(r) = clickhouse_db.db_client().query(select).fetch_all::().await - && r.len() >= 2 - { - rows = r; - break; - } - sleep(Duration::from_millis(200)).await; - } - assert!( - rows.len() >= 2, - "timed out waiting for schema_change_add_column rows: got {} of 2", - rows.len() - ); + let rows: Vec = clickhouse_db.query(select).await; pipeline.shutdown_and_wait().await.unwrap(); @@ -2197,7 +2079,7 @@ async fn schema_change_add_drop_rename() { let clickhouse_db = setup_clickhouse_database().await; let store = NotifyingStore::new(); let pipeline_id: PipelineId = random(); - let destination = clickhouse_db.build_destination(store.clone()); + let destination = TestDestinationWrapper::wrap(clickhouse_db.build_destination(store.clone())); let table_ready = store.notify_on_table_state_type(table_id, TableReplicationPhaseType::Ready).await; @@ -2207,7 +2089,7 @@ async fn schema_change_add_drop_rename() { pipeline_id, publication_name.to_owned(), store.clone(), - destination, + destination.clone(), ); pipeline.start().await.unwrap(); @@ -2246,6 +2128,8 @@ async fn schema_change_add_drop_rename() { .await .unwrap(); + let event_notify = destination.wait_for_events_count(vec![(EventType::Insert, 1)]).await; + database .run_sql(&format!( "INSERT INTO {} (full_name, status, email) VALUES ('Bob', 'pending', \ @@ -2255,28 +2139,14 @@ async fn schema_change_add_drop_rename() { .await .expect("Failed to insert Bob"); - // Poll until Bob's row arrives. + event_notify.notified().await; + let select = concat!( "SELECT id, full_name, status, email, cdc_operation ", "FROM \"test_schema__multi\" ", "ORDER BY id", ); - let mut rows: Vec = Vec::new(); - for _ in 0..50 { - if let Ok(r) = - clickhouse_db.db_client().query(select).fetch_all::().await - && r.len() >= 2 - { - rows = r; - break; - } - sleep(Duration::from_millis(200)).await; - } - assert!( - rows.len() >= 2, - "timed out waiting for schema_change_add_drop_rename rows: got {} of 2", - rows.len() - ); + let rows: Vec = clickhouse_db.query(select).await; pipeline.shutdown_and_wait().await.unwrap(); From fbca0ddd27e4bbbdc545532cac8248be8fafa8cf Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Wed, 13 May 2026 13:41:34 +0200 Subject: [PATCH 04/29] ref(core): Improve replica identity handling (#740) --- crates/etl-destinations/src/bigquery/core.rs | 42 +++++++ .../etl-destinations/src/clickhouse/core.rs | 42 +++++++ crates/etl-postgres/src/types/schema.rs | 111 ++++++++++++++--- crates/etl/src/conversions/event.rs | 117 +++++++++++++++--- 4 files changed, 282 insertions(+), 30 deletions(-) diff --git a/crates/etl-destinations/src/bigquery/core.rs b/crates/etl-destinations/src/bigquery/core.rs index 76ac24927..8c470fdcb 100644 --- a/crates/etl-destinations/src/bigquery/core.rs +++ b/crates/etl-destinations/src/bigquery/core.rs @@ -1124,6 +1124,23 @@ fn validate_bigquery_replica_identity( ); } + if !replicated_table_schema.all_primary_key_columns_replicated() { + let omitted_columns = replicated_table_schema + .unreplicated_primary_key_column_schemas() + .map(|column_schema| column_schema.name.as_str()) + .collect::>() + .join(","); + bail!( + ErrorKind::SourceSchemaError, + "BigQuery requires all source primary-key columns to be replicated", + format!( + "Table '{}' omits source primary-key columns from replication: {}", + replicated_table_schema.name(), + omitted_columns + ) + ); + } + match replicated_table_schema.identity_type() { IdentityType::PrimaryKey | IdentityType::Full => Ok(()), identity_type => { @@ -1611,6 +1628,22 @@ mod tests { ReplicatedTableSchema::from_masks(table_schema, replication_mask, identity_mask) } + fn replicated_schema_with_partial_primary_key() -> ReplicatedTableSchema { + let table_schema = Arc::new(TableSchema::new( + TableId::new(1), + TableName::new("public".to_owned(), "users".to_owned()), + vec![ + ColumnSchema::new("tenant_id".to_owned(), Type::INT4, -1, 1, Some(1), false), + ColumnSchema::new("id".to_owned(), Type::INT4, -1, 2, Some(2), false), + ColumnSchema::new("name".to_owned(), Type::TEXT, -1, 3, None, true), + ], + )); + let replication_mask = etl::types::ReplicationMask::from_bytes(vec![0, 1, 1]); + let identity_mask = IdentityMask::from_bytes(vec![0, 1, 0]); + + ReplicatedTableSchema::from_masks(table_schema, replication_mask, identity_mask) + } + #[test] fn table_name_to_bigquery_table_id_no_underscores() { let table_name = TableName::new("schema".to_owned(), "table".to_owned()); @@ -1765,6 +1798,15 @@ mod tests { assert_eq!(error.kind(), ErrorKind::SourceSchemaError); } + #[test] + fn validate_bigquery_replica_identity_rejects_partial_primary_key() { + let replicated_table_schema = replicated_schema_with_partial_primary_key(); + + let error = validate_bigquery_replica_identity(&replicated_table_schema).unwrap_err(); + assert_eq!(error.kind(), ErrorKind::SourceSchemaError); + assert!(error.to_string().contains("tenant_id")); + } + #[test] fn sequenced_bigquery_table_id_from_str_no_underscore() { let result = "tablewithoutsequence".parse::(); diff --git a/crates/etl-destinations/src/clickhouse/core.rs b/crates/etl-destinations/src/clickhouse/core.rs index dae00677b..9c4436cd1 100644 --- a/crates/etl-destinations/src/clickhouse/core.rs +++ b/crates/etl-destinations/src/clickhouse/core.rs @@ -768,6 +768,23 @@ where fn validate_replica_identity_for_clickhouse( replicated_table_schema: &ReplicatedTableSchema, ) -> EtlResult<()> { + if !replicated_table_schema.all_primary_key_columns_replicated() { + let omitted_columns = replicated_table_schema + .unreplicated_primary_key_column_schemas() + .map(|column_schema| column_schema.name.as_str()) + .collect::>() + .join(","); + return Err(etl_error!( + ErrorKind::SourceSchemaError, + "ClickHouse requires all source primary-key columns to be replicated", + format!( + "Table '{}' omits source primary-key columns from replication: {}", + replicated_table_schema.name(), + omitted_columns + ) + )); + } + match replicated_table_schema.identity_type() { IdentityType::PrimaryKey | IdentityType::Full => Ok(()), identity_type => Err(etl_error!( @@ -953,6 +970,22 @@ mod tests { ReplicatedTableSchema::from_masks(table_schema, replication_mask, identity_mask) } + fn replicated_schema_with_partial_primary_key() -> ReplicatedTableSchema { + let table_schema = Arc::new(TableSchema::new( + TableId::new(1), + TableName::new("public".to_owned(), "users".to_owned()), + vec![ + ColumnSchema::new("tenant_id".to_owned(), Type::INT4, -1, 1, Some(1), false), + ColumnSchema::new("id".to_owned(), Type::INT4, -1, 2, Some(2), false), + ColumnSchema::new("name".to_owned(), Type::TEXT, -1, 3, None, true), + ], + )); + let replication_mask = ReplicationMask::from_bytes(vec![0, 1, 1]); + let identity_mask = IdentityMask::from_bytes(vec![0, 1, 0]); + + ReplicatedTableSchema::from_masks(table_schema, replication_mask, identity_mask) + } + #[test] fn validate_replica_identity_for_clickhouse_accepts_primary_key() { validate_replica_identity_for_clickhouse(&replicated_schema(IdentityType::PrimaryKey)) @@ -981,6 +1014,15 @@ mod tests { assert_eq!(err.kind(), ErrorKind::SourceSchemaError); } + #[test] + fn validate_replica_identity_for_clickhouse_rejects_partial_primary_key() { + let err = + validate_replica_identity_for_clickhouse(&replicated_schema_with_partial_primary_key()) + .unwrap_err(); + assert_eq!(err.kind(), ErrorKind::SourceSchemaError); + assert!(err.to_string().contains("tenant_id")); + } + #[test] fn cdc_lsn_value_preserves_full_u64_range() { let value = cdc_lsn_to_clickhouse_value(PgLsn::from(u64::MAX)); diff --git a/crates/etl-postgres/src/types/schema.rs b/crates/etl-postgres/src/types/schema.rs index 33e5dd1dd..8d6edbbf1 100644 --- a/crates/etl-postgres/src/types/schema.rs +++ b/crates/etl-postgres/src/types/schema.rs @@ -9,6 +9,7 @@ use std::{ use pg_escape::quote_identifier; use thiserror::Error; use tokio_postgres::types::{FromSql, PgLsn, ToSql, Type}; +use tracing::warn; /// Errors that can occur during schema operations. #[derive(Debug, Error)] @@ -654,8 +655,8 @@ pub struct ReplicatedTableSchema { replication_mask: ReplicationMask, /// Cached number of replicated columns. replicated_column_count: usize, - /// A bitmask where 1 indicates the column at that index is part of the row - /// identity used by logical replication. + /// A bitmask where 1 indicates the column at that index is a replicated + /// row identity column used by logical replication. identity_mask: IdentityMask, /// Cached number of replicated identity columns. identity_column_count: usize, @@ -670,8 +671,15 @@ impl ReplicatedTableSchema { /// masks, inferring the identity type from the mask shape. /// /// Both masks are expressed in full table-schema width. The identity mask - /// must therefore align with the same column order as the replication - /// mask. + /// must be a subset of the replication mask because row-event decoding can + /// only consume key columns that PostgreSQL includes in the relation + /// payload. + /// + /// ETL stores runtime identity, not raw source catalog identity. Initial + /// copy follows streaming relation-message semantics by marking only + /// replicated columns as identity columns. Update/delete replication relies + /// on PostgreSQL validating that the source identity is covered; + /// insert-only publications do not need identity data. /// /// This constructor infers the semantic identity type from the table /// schema and supplied masks, and caches the derived column counts needed @@ -696,13 +704,20 @@ impl ReplicatedTableSchema { "identity mask length must match column count" ); - for (&replicated, &identity) in - replication_mask.as_slice().iter().zip(identity_mask.as_slice().iter()) + for ((column_schema, &replicated), &identity) in table_schema + .column_schemas + .iter() + .zip(replication_mask.as_slice().iter()) + .zip(identity_mask.as_slice().iter()) { - debug_assert!( - identity == 0 || replicated == 1, - "identity mask must be a subset of the replication mask" - ); + if identity == 1 && replicated == 0 { + warn!( + table_id = %table_schema.id, + table_name = %table_schema.name, + column_name = %column_schema.name, + "replica identity column is not replicated" + ); + } } // We pre-compute counts to avoid computing them each time since they are needed @@ -735,16 +750,16 @@ impl ReplicatedTableSchema { /// Creates a [`ReplicatedTableSchema`] from a schema and a pre-computed /// replication mask. /// - /// The identity mask is derived from primary-key membership. This is a - /// convenient fallback for code paths that only need replicated columns or - /// when the source schema and identity are known to match primary-key - /// semantics. + /// The identity mask is derived from replicated primary-key membership. + /// This is a convenient fallback for code paths that only need replicated + /// columns or when the source schema and identity are known to match + /// primary-key semantics. pub fn from_mask(table_schema: Arc, replication_mask: ReplicationMask) -> Self { let identity_mask = Self::primary_key_identity_mask(&table_schema, &replication_mask); Self::from_masks(table_schema, replication_mask, identity_mask) } - /// Creates a [`ReplicatedTableSchema`] where all columns are replicated. . + /// Creates a [`ReplicatedTableSchema`] where all columns are replicated. pub fn all(table_schema: Arc) -> Self { let replication_mask = ReplicationMask::all(&table_schema); Self::from_mask(table_schema, replication_mask) @@ -787,7 +802,7 @@ impl ReplicatedTableSchema { /// here if the chosen index resolves to the same current columns as the /// primary key. /// - /// This comparison is structural over the runtime schema masks, not a + /// This comparison is structural over the runtime schema identity, not a /// direct comparison of PostgreSQL identity modes or index OIDs. Because /// ETL tracks DDL/schema changes, that gives the intended notion of /// primary-key equivalence across schema evolution. @@ -829,6 +844,9 @@ impl ReplicatedTableSchema { pub fn identity_column_schemas( &self, ) -> impl ExactSizeIterator + Clone + '_ { + // Key tuples from PostgreSQL should only use columns present in the + // relation payload. Check both masks here so tuple decoding only sees + // columns that are both identity columns and actually replicated. let inner = self .table_schema .column_schemas @@ -864,6 +882,32 @@ impl ReplicatedTableSchema { SizedIterator::new(inner, self.primary_key_column_count) } + /// Returns whether every source primary-key column is replicated. + /// + /// Destinations that match rows by the source primary key need this check + /// in addition to runtime identity checks, because replicated primary-key + /// iterators intentionally expose only the replicated subset. + pub fn all_primary_key_columns_replicated(&self) -> bool { + self.unreplicated_primary_key_column_schemas().next().is_none() + } + + /// Returns source primary-key columns omitted from replication. + pub fn unreplicated_primary_key_column_schemas( + &self, + ) -> impl Iterator + Clone + '_ { + self.table_schema + .column_schemas + .iter() + .zip(self.replication_mask.as_slice().iter()) + .filter_map(|(column_schema, &replicated)| { + if column_schema.primary_key() && replicated == 0 { + Some(column_schema) + } else { + None + } + }) + } + /// Computes the diff between this schema (old) and another schema (new). /// /// Only consider replicated columns. Uses ordinal positions to track @@ -1186,6 +1230,41 @@ mod tests { assert!(!replicated_table_schema.identity_matches_primary_key()); } + #[test] + fn all_primary_key_columns_replicated_returns_true_for_complete_primary_key() { + let schema = Arc::new(create_test_table_schema()); + let replication_mask = ReplicationMask::all(&schema); + let replicated_table_schema = ReplicatedTableSchema::from_mask(schema, replication_mask); + + assert!(replicated_table_schema.all_primary_key_columns_replicated()); + assert_eq!(replicated_table_schema.unreplicated_primary_key_column_schemas().count(), 0); + } + + #[test] + fn all_primary_key_columns_replicated_returns_false_for_partial_primary_key() { + let schema = Arc::new(TableSchema::new( + TableId::new(123), + TableName::new("public".to_owned(), "test_table".to_owned()), + vec![ + ColumnSchema::new("tenant_id".to_owned(), Type::INT4, -1, 1, Some(1), false), + ColumnSchema::new("id".to_owned(), Type::INT4, -1, 2, Some(2), false), + ColumnSchema::new("name".to_owned(), Type::TEXT, -1, 3, None, true), + ], + )); + let replication_mask = ReplicationMask::from_bytes(vec![0, 1, 1]); + let identity_mask = IdentityMask::from_bytes(vec![0, 1, 0]); + let replicated_table_schema = + ReplicatedTableSchema::from_masks(schema, replication_mask, identity_mask); + + let omitted_columns = replicated_table_schema + .unreplicated_primary_key_column_schemas() + .map(|column_schema| column_schema.name.as_str()) + .collect::>(); + + assert!(!replicated_table_schema.all_primary_key_columns_replicated()); + assert_eq!(omitted_columns, ["tenant_id"]); + } + #[test] fn schema_diff_no_changes() { let old_schema = create_replicated_schema(vec![ diff --git a/crates/etl/src/conversions/event.rs b/crates/etl/src/conversions/event.rs index 295685fba..a0a14f60d 100644 --- a/crates/etl/src/conversions/event.rs +++ b/crates/etl/src/conversions/event.rs @@ -115,10 +115,16 @@ impl IdentityMessage { /// The returned mask is expressed in full table-schema width so it can be /// combined with the replication mask held by [`ReplicatedTableSchema`]. /// - /// `REPLICA IDENTITY FULL` uses the replicated columns themselves as the - /// row identity. `DEFAULT` uses the primary key, `USING INDEX` uses the + /// ETL stores runtime identity over replicated columns, matching relation + /// messages where PostgreSQL can only flag columns present in the relation + /// payload. PostgreSQL should reject update/delete publications whose + /// column lists do not cover the source identity; insert-only publications + /// do not use identity data. + /// + /// `REPLICA IDENTITY FULL` uses every replicated column as the runtime row + /// identity. `DEFAULT` uses the primary key, `USING INDEX` uses the /// configured replica-identity index, and `NOTHING` produces an empty - /// identity mask. + /// identity mask, each intersected with the replication mask. pub(crate) fn build_identity_mask( &self, table_schema: &TableSchema, @@ -126,11 +132,14 @@ impl IdentityMessage { ) -> EtlResult { match self.relreplident.as_str() { "f" => Ok(IdentityMask::from_bytes(replication_mask.as_slice().to_vec())), - "d" => { - Ok(Self::build_identity_mask_from_attnums(table_schema, &self.primary_key_attnums)) - } + "d" => Ok(Self::build_identity_mask_from_attnums( + table_schema, + replication_mask, + &self.primary_key_attnums, + )), "i" => Ok(Self::build_identity_mask_from_attnums( table_schema, + replication_mask, &self.replica_identity_index_attnums, )), "n" => Ok(IdentityMask::from_bytes(vec![0; table_schema.column_schemas.len()])), @@ -148,8 +157,14 @@ impl IdentityMessage { } /// Builds an identity mask from ordered attribute numbers. + /// + /// Initial schema metadata describes the source table identity, which can + /// include columns omitted by an insert-only publication column list. + /// Filter through the replication mask so initial copy matches streaming + /// relation-message semantics. fn build_identity_mask_from_attnums( table_schema: &TableSchema, + replication_mask: &ReplicationMask, attnums: &[i32], ) -> IdentityMask { let attnums: HashSet = attnums.iter().copied().collect(); @@ -158,7 +173,10 @@ impl IdentityMessage { table_schema .column_schemas .iter() - .map(|column_schema| u8::from(attnums.contains(&column_schema.ordinal_position))) + .zip(replication_mask.as_slice().iter()) + .map(|(column_schema, &replicated)| { + u8::from(replicated == 1 && attnums.contains(&column_schema.ordinal_position)) + }) .collect(), ) } @@ -320,13 +338,13 @@ pub(crate) fn parse_replicated_column_names( /// /// PostgreSQL exposes replica-identity mode on the relation itself and /// key-column membership on each [`protocol::RelationBody`] column. For -/// `REPLICA IDENTITY FULL`, every replicated column belongs to the old-row -/// identity. Otherwise the low bit of the column flags marks identity -/// membership. The column order is the same `attnum` order described in -/// [`parse_replicated_column_names`]. This returns names because identity-mask -/// membership is also name-based; tuple interpretation relies on the resulting -/// replicated schema preserving relation-message order after the masks are -/// applied. +/// `REPLICA IDENTITY FULL`, every relation-message column belongs to the +/// runtime old-row identity. Otherwise, the low bit of the column flags marks +/// identity membership. The column order is the same `attnum` order described +/// in [`parse_replicated_column_names`]. This returns names because +/// identity-mask membership is also name-based; tuple interpretation relies on +/// the resulting replicated schema preserving relation-message order after the +/// masks are applied. pub(crate) fn parse_replica_identity_column_names( relation_body: &protocol::RelationBody, ) -> EtlResult> { @@ -339,6 +357,9 @@ pub(crate) fn parse_replica_identity_column_names( _ => relation_body .columns() .iter() + // PostgreSQL sends relation column flags as a bitmask. Bit 0 is + // LOGICALREP_IS_REPLICA_IDENTITY, so `& 1` tests only that bit + // while ignoring any other protocol flags that may be present. .filter(|column| column.flags() & 1 == 1) .map(|column| column.name().map(ToString::to_string)) .collect::, _>>()?, @@ -1208,6 +1229,74 @@ mod tests { assert_eq!(identity_type, IdentityType::AlternativeKey); } + #[test] + fn build_identity_for_default_filters_to_replicated_columns() { + let table_schema = TableSchema::new( + TableId::new(42), + TableName::new("public".to_owned(), "test".to_owned()), + vec![ + ColumnSchema::new("id".to_owned(), Type::INT8, -1, 1, Some(1), false), + ColumnSchema::new("name".to_owned(), Type::TEXT, -1, 2, Some(2), false), + ColumnSchema::new("email".to_owned(), Type::TEXT, -1, 3, None, false), + ], + ); + let replication_mask = ReplicationMask::from_bytes(vec![1, 0, 1]); + let identity = IdentityMessage { + primary_key_attnums: vec![1, 2], + relreplident: "d".to_owned(), + replica_identity_index_attnums: Vec::new(), + }; + + let identity_mask = identity.build_identity_mask(&table_schema, &replication_mask).unwrap(); + + assert_eq!(identity_mask.as_slice(), &[1, 0, 0]); + } + + #[test] + fn build_identity_for_using_index_filters_to_replicated_columns() { + let table_schema = TableSchema::new( + TableId::new(42), + TableName::new("public".to_owned(), "test".to_owned()), + vec![ + ColumnSchema::new("id".to_owned(), Type::INT8, -1, 1, Some(1), false), + ColumnSchema::new("name".to_owned(), Type::TEXT, -1, 2, None, false), + ColumnSchema::new("email".to_owned(), Type::TEXT, -1, 3, None, false), + ], + ); + let replication_mask = ReplicationMask::from_bytes(vec![0, 1, 0]); + let identity = IdentityMessage { + primary_key_attnums: vec![1], + relreplident: "i".to_owned(), + replica_identity_index_attnums: vec![2, 3], + }; + + let identity_mask = identity.build_identity_mask(&table_schema, &replication_mask).unwrap(); + + assert_eq!(identity_mask.as_slice(), &[0, 1, 0]); + } + + #[test] + fn build_identity_for_full_uses_all_replicated_columns() { + let table_schema = TableSchema::new( + TableId::new(42), + TableName::new("public".to_owned(), "test".to_owned()), + vec![ + ColumnSchema::new("id".to_owned(), Type::INT8, -1, 1, Some(1), false), + ColumnSchema::new("email".to_owned(), Type::TEXT, -1, 2, None, false), + ], + ); + let replication_mask = ReplicationMask::from_bytes(vec![1, 0]); + let identity = IdentityMessage { + primary_key_attnums: vec![1], + relreplident: "f".to_owned(), + replica_identity_index_attnums: Vec::new(), + }; + + let identity_mask = identity.build_identity_mask(&table_schema, &replication_mask).unwrap(); + + assert_eq!(identity_mask.as_slice(), &[1, 0]); + } + #[test] fn convert_tuple_to_row_rejects_missing_non_nullable_columns_for_full_rows() { let column_schemas = [ From 7128d17d4d8a17a45fa8a0b91d8d0045f001663f Mon Sep 17 00:00:00 2001 From: Jordan McQueen Date: Fri, 15 May 2026 14:47:45 +0900 Subject: [PATCH 05/29] Add clickhouse client- and server-side timeouts (#741) * feat(error): add DestinationTimeout error kind * feat(clickhouse): add ClickHouseClientConfig and thread through Adds a sibling config struct to ClickHouseInserterConfig with per-operation server-side timeout fields and a client_timeout_epsilon. Threaded through ClickHouseDestination::new and ClickHouseClient::new; all existing callers pass Default::default(). No behavior change yet: the new config field is plumbed but not consulted. * feat(clickhouse): wrap client calls in tokio::time::timeout Adds TimeoutOp + timeout_call so each ClickHouseClient call is bound by client_budget(server_budget) = server + epsilon. Server-side options (connect_timeout, http_connection_timeout, http_send_timeout, http_receive_timeout) are set globally; max_execution_time and lock_acquire_timeout are applied per-call to DDL/TRUNCATE so they do not leak into probe and schema_query budgets. Client-side deadlines surface as DestinationTimeout. * test(clickhouse): cover client_budget and timeout_call Covers ClickHouseClientConfig::client_budget arithmetic, TimeoutOp::Display matching the strings used in error messages, and timeout_call across its three outcomes (deadline expired, inner error, success). Uses tokio::test(start_paused = true) so the deadline test does not block on real wall-clock time. * ref(clickhouse): rename probe timeout, tighten config docs Renames probe_server_timeout to connectivity_check_timeout and shortens the doc-comments on ClickHouseClientConfig per review. Drops the epsilon default from 5s to 4s. * ref(clickhouse): rename TimeoutOp::Probe, drop timeout_kind() Renames TimeoutOp::Probe to ConnectivityCheck (matching the connectivity_check_timeout config field) and updates its Display string to 'connectivity check'. Inlines ErrorKind::DestinationTimeout in timeout_call and drops the now-trivial timeout_kind() helper. * ref(clickhouse): let TimeoutOp drive the server budget Adds TimeoutOp::server_budget, makes timeout_call take a &ClickHouseClientConfig and derive the budget internally via config.client_budget(op.server_budget(config)). Call sites lose their hand-rolled budget bindings and the (op, budget) pairing is now impossible to mismatch. TimeoutOp's three responsibilities (server budget, failed kind, Display label) are documented in one place. * ref(clickhouse): move TimeoutOp into core as ClickHouseOperationKind Renames TimeoutOp to ClickHouseOperationKind and moves it next to ClickHouseClientConfig in core.rs. server_budget and the budget-form client_budget move onto ClickHouseClientConfig (config now takes an op directly), so call sites read 'config.client_budget(op)' instead of the awkward 'config.client_budget(op.server_budget(config))'. failed_kind and Display stay on the enum. Doc-comment reworded to lead with 'Categories of ClickHouse client calls' per review. * ref(clickhouse): tighten doc-comments per review Shortens server_budget, ClickHouseOperationKind, and timeout_call doc-comments. The 'client-side deadlines map to DestinationTimeout' note moves from the enum to timeout_call where the mapping happens. * ref(clickhouse): reword ClickHouseOperationKind doc 'Enables the following:' -> 'Used to:'; bullets shift from gerunds to imperative verbs. * fix(clickhouse): retry DestinationTimeout, restore table context Adds DestinationTimeout to the auto-retry arms in both workers/policy.rs and state/table.rs so client-side timeouts retry on a timed delay instead of falling into the default Manual retry branch. Without this, a single transient timeout would pause replication until an operator intervened. Extends timeout_call with an optional context string and threads table_name through table_columns, truncate_table, and insert.end(), restoring diagnostic info that the previous timeout refactor dropped. Also rewords two stale 'probe and schema_query' comments and adds a unit test for secs_string covering the subsecond floor and zero preservation. * ref(clickhouse): drop redundant comment in client constructor The four .with_option calls have self-explanatory names; the rationale for per-call overrides lives at the per-call sites. * ref(clickhouse): drop _server_ infix from config fields, dedupe error detail Renames schema_query_server_timeout -> schema_query_timeout, ddl_server_timeout -> ddl_timeout, insert_server_timeout -> insert_timeout (and matching constants) for symmetry with connectivity_check_timeout. The doc-comment already says 'server-side', so the infix was redundant. Also drops the leading 'ClickHouse ' from timeout_call's detail format; the static description already begins with 'ClickHouse call', so the rendered output no longer duplicates the prefix. * test(clickhouse): cover inner-error path with context Adds the missing matrix cell: inner ClickHouse error in timeout_call with Some(context). Asserts that the context (e.g. 'table: users') shows up in the detail alongside the per-op kind. * test(clickhouse): assert inner clickhouse error is attached as source Pins that timeout_call's inner-error arm calls etl_error! with 'source: err'. Catches a regression where the source: attribute gets dropped silently, which would lose the underlying clickhouse error from log/debug output. * ref(clickhouse): floor secs_string at 1s ClickHouse interprets 0 as no timeout for http_*_timeout, max_execution_time, and lock_acquire_timeout. Preserving Duration::ZERO as '0' silently disabled the server-side bound while the client-side tokio::time::timeout still fired at the epsilon, so the operator's intent (if any) was overridden anyway. Floor sub-second values at 1s instead and document the behaviour. * fix(clippy): drop needless borrows and reassign-with-default CI clippy (1.93) flagged: - needless_borrows_for_generic_args on four with_option calls: drop & on the secs_string(...) String, with_option takes Into. - field_reassign_with_default in client_budget test: switch to struct update syntax with ..Default::default(). * ref(clickhouse): shorten floor_secs name, space after inner fn floor_secs_string -> floor_secs (return type already says String). Blank line after the inner detail() fn body inside timeout_call. * ref(clickhouse): drop superfluous comment in execute_ddl The per-call .with_option pattern is self-explanatory from the code itself and repeats in truncate_table / table_columns. * test(clickhouse): add GIVEN/WHEN/THEN doc-comments to new tests Documents the 10 unit tests added for ClickHouseClientConfig, ClickHouseOperationKind, timeout_call, and floor_secs in the style used by the integration tests in tests/clickhouse/pipeline.rs. * ref(clickhouse): tighten floor_secs doc-comment * ref(clickhouse): add max_execution_time to validate_connectivity For consistency with the other wrapped calls (execute_ddl, table_columns, truncate_table), bound SELECT 1 with a server-side max_execution_time set to connectivity_check_timeout. Defense in depth against a wedged server's query queue. * ref(clickhouse): rename server/client_budget to server/client_timeout_for server_budget / client_budget on ClickHouseClientConfig were the only methods using the 'budget' vocabulary; everything else (config fields, ClickHouse settings, local idioms) speaks in 'timeout'. Rename to server_timeout_for(op) / client_timeout_for(op) for consistency, plus the matching internal call, doc-comment, local var, and test names. --- crates/etl-api/src/validation/validators.rs | 3 +- .../etl-destinations/src/clickhouse/client.rs | 393 ++++++++++++++++-- .../etl-destinations/src/clickhouse/core.rs | 110 ++++- crates/etl-destinations/src/clickhouse/mod.rs | 2 +- .../src/clickhouse/test_utils.rs | 3 +- .../tests/clickhouse/pipeline.rs | 4 +- crates/etl-examples/src/bin/clickhouse.rs | 5 +- crates/etl-replicator/src/core.rs | 3 +- crates/etl/src/error.rs | 1 + crates/etl/src/state/table.rs | 1 + crates/etl/src/workers/policy.rs | 1 + 11 files changed, 483 insertions(+), 43 deletions(-) diff --git a/crates/etl-api/src/validation/validators.rs b/crates/etl-api/src/validation/validators.rs index 3db50386a..4e99f8d71 100644 --- a/crates/etl-api/src/validation/validators.rs +++ b/crates/etl-api/src/validation/validators.rs @@ -7,7 +7,7 @@ use etl::store::both::memory::MemoryStore; use etl_config::{SerializableSecretString, parse_ducklake_url}; use etl_destinations::{ bigquery::BigQueryClient, - clickhouse::ClickHouseClient, + clickhouse::{ClickHouseClient, ClickHouseClientConfig}, ducklake::{DuckLakeDestination, S3Config as DucklakeS3Config}, iceberg::{IcebergClient, S3_ACCESS_KEY_ID, S3_ENDPOINT, S3_SECRET_ACCESS_KEY}, }; @@ -663,6 +663,7 @@ impl Validator for ClickHouseValidator { self.user.clone(), self.password.as_ref().map(|password| password.expose_secret().to_owned()), self.database.clone(), + ClickHouseClientConfig::default(), ); match client.validate_connectivity().await { Ok(_) => Ok(Vec::new()), diff --git a/crates/etl-destinations/src/clickhouse/client.rs b/crates/etl-destinations/src/clickhouse/client.rs index c9848d541..7f60ab156 100644 --- a/crates/etl-destinations/src/clickhouse/client.rs +++ b/crates/etl-destinations/src/clickhouse/client.rs @@ -1,4 +1,8 @@ -use std::{sync::Arc, time::Instant}; +use std::{ + future::Future, + sync::Arc, + time::{Duration, Instant}, +}; use clickhouse::Client; use etl::{ @@ -8,11 +12,59 @@ use etl::{ use url::Url; use crate::clickhouse::{ + core::{ClickHouseClientConfig, ClickHouseOperationKind}, encoding::{ClickHouseValue, encode_to_row_binary}, metrics::{ETL_CLICKHOUSE_DDL_DURATION_SECONDS, ETL_CLICKHOUSE_INSERT_DURATION_SECONDS}, schema::{clickhouse_column_type, quote_identifier}, }; +/// Formats a `Duration` as a whole-seconds string for ClickHouse +/// `.with_option(...)` settings, floored at `"1"`. ClickHouse interprets `"0"` +/// as "no timeout" for `http_*_timeout`, `max_execution_time`, and +/// `lock_acquire_timeout`; the floor prevents `Duration::ZERO` or any +/// sub-second value from accidentally disabling the bound. +fn floor_secs(d: Duration) -> String { + d.as_secs().max(1).to_string() +} + +/// Runs `fut` under `tokio::time::timeout` using the client-side timeout for +/// `op` from `config`. Inner ClickHouse errors map onto `op.failed_kind()`; +/// client-side deadlines map onto [`ErrorKind::DestinationTimeout`]. +/// `context`, when present, is appended to the error detail (e.g. +/// `"table: foo"`) so call-site-specific diagnostic info is preserved. +async fn timeout_call( + op: ClickHouseOperationKind, + config: &ClickHouseClientConfig, + context: Option<&str>, + fut: F, +) -> EtlResult +where + F: Future>, +{ + fn detail(op: ClickHouseOperationKind, what: &str, context: Option<&str>) -> String { + match context { + Some(c) => format!("{op} {what}; {c}"), + None => format!("{op} {what}"), + } + } + + let client_timeout = config.client_timeout_for(op); + match tokio::time::timeout(client_timeout, fut).await { + Ok(Ok(value)) => Ok(value), + Ok(Err(err)) => Err(etl_error!( + op.failed_kind(), + "ClickHouse call failed", + detail(op, "failed", context), + source: err + )), + Err(_) => Err(etl_error!( + ErrorKind::DestinationTimeout, + "ClickHouse call timed out", + detail(op, &format!("timed out after {client_timeout:?}"), context) + )), + } +} + /// Capacity of the internal write buffer used per INSERT statement. /// /// When this many bytes have been written to the buffer it is flushed to the @@ -113,6 +165,7 @@ impl DdlKind { #[derive(Clone)] pub struct ClickHouseClient { inner: Arc, + config: ClickHouseClientConfig, } impl ClickHouseClient { @@ -125,6 +178,7 @@ impl ClickHouseClient { user: impl Into, password: Option, database: impl Into, + config: ClickHouseClientConfig, ) -> Self { let mut client = Client::default().with_url(url.to_string()).with_user(user).with_database(database); @@ -133,7 +187,13 @@ impl ClickHouseClient { client = client.with_password(pw); } - Self { inner: Arc::new(client) } + client = client + .with_option("connect_timeout", floor_secs(config.connectivity_check_timeout)) + .with_option("http_connection_timeout", floor_secs(config.connectivity_check_timeout)) + .with_option("http_send_timeout", floor_secs(config.insert_timeout)) + .with_option("http_receive_timeout", floor_secs(config.insert_timeout)); + + Self { inner: Arc::new(client), config } } /// Verifies that the ClickHouse server is reachable. @@ -143,9 +203,18 @@ impl ClickHouseClient { /// destination's `validate_connectivity` so callers (notably the /// `etl-api` validators) can treat the two destinations uniformly. pub async fn validate_connectivity(&self) -> EtlResult<()> { - self.inner.query("SELECT 1").fetch_one::().await.map(|_| ()).map_err( - |err| etl_error!(ErrorKind::Unknown, "ClickHouse connectivity check failed", source: err), + let query = self + .inner + .query("SELECT 1") + .with_option("max_execution_time", floor_secs(self.config.connectivity_check_timeout)); + timeout_call( + ClickHouseOperationKind::ConnectivityCheck, + &self.config, + None, + query.fetch_one::(), ) + .await?; + Ok(()) } /// Executes a DDL statement (e.g. `CREATE TABLE IF NOT EXISTS …`) and @@ -153,10 +222,14 @@ impl ClickHouseClient { /// histogram labelled with the DDL `kind` and `table_name`. pub(crate) async fn execute_ddl(&self, kind: DdlKind, sql: &str) -> EtlResult<()> { let ddl_start = Instant::now(); + let ddl_secs = floor_secs(self.config.ddl_timeout); + let query = self + .inner + .query(sql) + .with_option("max_execution_time", &ddl_secs) + .with_option("lock_acquire_timeout", &ddl_secs); let result = - self.inner.query(sql).execute().await.map_err( - |err| etl_error!(ErrorKind::Unknown, "ClickHouse DDL failed", source: err), - ); + timeout_call(ClickHouseOperationKind::Ddl, &self.config, None, query.execute()).await; metrics::histogram!( ETL_CLICKHOUSE_DDL_DURATION_SECONDS, "kind" => kind.as_label(), @@ -170,22 +243,22 @@ impl ClickHouseClient { &self, table_name: &str, ) -> EtlResult> { - self.inner + let schema_secs = floor_secs(self.config.schema_query_timeout); + let query = self + .inner .query( "SELECT name, type AS type_name FROM system.columns WHERE database = \ currentDatabase() AND table = ? ORDER BY position", ) - .bind(table_name) - .fetch_all::() - .await - .map_err(|err| { - etl_error!( - ErrorKind::Unknown, - "ClickHouse schema query failed", - format!("table: {table_name}"), - source: err - ) - }) + .with_option("max_execution_time", &schema_secs) + .bind(table_name); + timeout_call( + ClickHouseOperationKind::SchemaQuery, + &self.config, + Some(&format!("table: {table_name}")), + query.fetch_all::(), + ) + .await } /// Adds a column to an existing ClickHouse table. @@ -231,14 +304,19 @@ impl ClickHouseClient { /// Executes `TRUNCATE TABLE IF EXISTS` for the supplied table. pub(crate) async fn truncate_table(&self, table_name: &str) -> EtlResult<()> { - self.inner.query(&build_truncate_table_sql(table_name)).execute().await.map_err(|err| { - etl_error!( - ErrorKind::Unknown, - "ClickHouse truncate failed", - format!("table: {table_name}"), - source: err - ) - }) + let ddl_secs = floor_secs(self.config.ddl_timeout); + let query = self + .inner + .query(&build_truncate_table_sql(table_name)) + .with_option("max_execution_time", &ddl_secs) + .with_option("lock_acquire_timeout", &ddl_secs); + timeout_call( + ClickHouseOperationKind::Ddl, + &self.config, + Some(&format!("table: {table_name}")), + query.execute(), + ) + .await } /// Inserts `rows` into `table_name` using the RowBinary format. @@ -283,14 +361,13 @@ impl ClickHouseClient { bytes += row_buf.len() as u64; } - insert.end().await.map_err(|err| { - etl_error!( - ErrorKind::Unknown, - "ClickHouse insert flush failed", - format!("table: {table_name}"), - source: err - ) - })?; + timeout_call( + ClickHouseOperationKind::Insert, + &self.config, + Some(&format!("table: {table_name}")), + insert.end(), + ) + .await?; metrics::histogram!( ETL_CLICKHOUSE_INSERT_DURATION_SECONDS, "source" => source, @@ -374,4 +451,250 @@ mod tests { assert_eq!(sql, "INSERT INTO \"table\"\"name\" FORMAT RowBinary"); } + + /// # GIVEN + /// A config with a custom server timeout and epsilon. + /// + /// # WHEN + /// `client_timeout_for(op)` is queried. + /// + /// # THEN + /// It returns `server_timeout_for(op) + client_timeout_epsilon`. + #[test] + fn client_timeout_adds_epsilon_to_server_timeout() { + let config = ClickHouseClientConfig { + connectivity_check_timeout: Duration::from_secs(10), + client_timeout_epsilon: Duration::from_secs(3), + ..Default::default() + }; + assert_eq!( + config.client_timeout_for(ClickHouseOperationKind::ConnectivityCheck), + Duration::from_secs(13) + ); + + let config = ClickHouseClientConfig { + connectivity_check_timeout: Duration::ZERO, + client_timeout_epsilon: Duration::from_secs(3), + ..Default::default() + }; + assert_eq!( + config.client_timeout_for(ClickHouseOperationKind::ConnectivityCheck), + Duration::from_secs(3) + ); + } + + /// # GIVEN + /// Each `ClickHouseOperationKind` variant. + /// + /// # WHEN + /// `Display` is invoked. + /// + /// # THEN + /// It produces the human-readable op name interpolated into error + /// messages by `timeout_call`. + #[test] + fn operation_kind_display_matches_error_messages() { + assert_eq!(ClickHouseOperationKind::ConnectivityCheck.to_string(), "connectivity check"); + assert_eq!(ClickHouseOperationKind::SchemaQuery.to_string(), "schema query"); + assert_eq!(ClickHouseOperationKind::Ddl.to_string(), "DDL"); + assert_eq!(ClickHouseOperationKind::Insert.to_string(), "insert"); + } + + /// # GIVEN + /// A future that never resolves and a config with a finite budget. + /// + /// # WHEN + /// `timeout_call` is awaited under paused time. + /// + /// # THEN + /// It returns an `EtlError` with kind `DestinationTimeout` and a + /// detail that mentions the op and "timed out". + #[tokio::test(start_paused = true)] + async fn timeout_call_returns_destination_timeout_on_deadline() { + // A future that never resolves; tokio's paused clock advances virtual + // time when all tasks are stalled, so the timeout fires immediately + // in real wall-clock terms. + let config = ClickHouseClientConfig::default(); + let never = std::future::pending::>(); + let err = timeout_call(ClickHouseOperationKind::ConnectivityCheck, &config, None, never) + .await + .unwrap_err(); + assert_eq!(err.kind(), ErrorKind::DestinationTimeout); + assert!( + err.detail() + .is_some_and(|d| d.contains("connectivity check") && d.contains("timed out")), + "unexpected detail: {:?}", + err.detail() + ); + } + + /// # GIVEN + /// A never-resolving future and `Some(context)`. + /// + /// # WHEN + /// `timeout_call`'s deadline fires. + /// + /// # THEN + /// The error detail contains the context string. + #[tokio::test(start_paused = true)] + async fn timeout_call_appends_context_to_detail() { + let config = ClickHouseClientConfig::default(); + let never = std::future::pending::>(); + let err = + timeout_call(ClickHouseOperationKind::Insert, &config, Some("table: users"), never) + .await + .unwrap_err(); + assert!( + err.detail().is_some_and(|d| d.contains("table: users")), + "unexpected detail: {:?}", + err.detail() + ); + } + + /// # GIVEN + /// A future that returns a `clickhouse::error::Error` before the + /// deadline. + /// + /// # WHEN + /// `timeout_call` is awaited with no context. + /// + /// # THEN + /// It returns an `EtlError` with the op's `failed_kind` and a detail + /// that mentions the op and "failed". + #[tokio::test(start_paused = true)] + async fn timeout_call_propagates_inner_error() { + let config = ClickHouseClientConfig::default(); + let fut = async { Err::<(), _>(clickhouse::error::Error::NotEnoughData) }; + let err = timeout_call(ClickHouseOperationKind::SchemaQuery, &config, None, fut) + .await + .unwrap_err(); + assert_eq!(err.kind(), ErrorKind::DestinationQueryFailed); + assert!( + err.detail().is_some_and(|d| d.contains("schema query") && d.contains("failed")), + "unexpected detail: {:?}", + err.detail() + ); + } + + /// # GIVEN + /// A future that resolves to `Ok` before the deadline. + /// + /// # WHEN + /// `timeout_call` is awaited. + /// + /// # THEN + /// It returns the inner `Ok` value unchanged. + #[tokio::test(start_paused = true)] + async fn timeout_call_passes_through_success() { + let config = ClickHouseClientConfig::default(); + let fut = async { Ok::(42) }; + let value = + timeout_call(ClickHouseOperationKind::Insert, &config, None, fut).await.unwrap(); + assert_eq!(value, 42); + } + + /// # GIVEN + /// A future that returns a `clickhouse::error::Error` and + /// `Some(context)`. + /// + /// # WHEN + /// `timeout_call` is awaited. + /// + /// # THEN + /// The error has the op's `failed_kind`, the detail contains the + /// context, and the inner clickhouse error is attached as `source`. + #[tokio::test(start_paused = true)] + async fn timeout_call_inner_error_includes_context() { + use std::error::Error as _; + let config = ClickHouseClientConfig::default(); + let fut = async { Err::<(), _>(clickhouse::error::Error::NotEnoughData) }; + let err = timeout_call(ClickHouseOperationKind::Insert, &config, Some("table: users"), fut) + .await + .unwrap_err(); + assert_eq!(err.kind(), ErrorKind::DestinationAtomicBatchRetryable); + assert!( + err.detail().is_some_and(|d| d.contains("insert failed") && d.contains("table: users")), + "unexpected detail: {:?}", + err.detail() + ); + assert!(err.source().is_some(), "expected inner clickhouse error to be attached"); + } + + /// # GIVEN + /// A default `ClickHouseClientConfig`. + /// + /// # WHEN + /// `server_timeout_for(op)` is queried for each variant. + /// + /// # THEN + /// Each variant returns the corresponding config field. + #[test] + fn server_timeout_per_operation_kind() { + let config = ClickHouseClientConfig::default(); + assert_eq!( + config.server_timeout_for(ClickHouseOperationKind::ConnectivityCheck), + config.connectivity_check_timeout + ); + assert_eq!( + config.server_timeout_for(ClickHouseOperationKind::SchemaQuery), + config.schema_query_timeout + ); + assert_eq!(config.server_timeout_for(ClickHouseOperationKind::Ddl), config.ddl_timeout); + assert_eq!( + config.server_timeout_for(ClickHouseOperationKind::Insert), + config.insert_timeout + ); + } + + /// # GIVEN + /// Each `ClickHouseOperationKind` variant. + /// + /// # WHEN + /// `failed_kind` is queried. + /// + /// # THEN + /// Each variant maps to the `ErrorKind` that drives the appropriate + /// retry policy for that bucket. + #[test] + fn operation_kind_failed_kind_per_bucket() { + assert_eq!( + ClickHouseOperationKind::ConnectivityCheck.failed_kind(), + ErrorKind::DestinationConnectionFailed + ); + assert_eq!( + ClickHouseOperationKind::SchemaQuery.failed_kind(), + ErrorKind::DestinationQueryFailed + ); + assert_eq!(ClickHouseOperationKind::Ddl.failed_kind(), ErrorKind::DestinationQueryFailed); + assert_eq!( + ClickHouseOperationKind::Insert.failed_kind(), + ErrorKind::DestinationAtomicBatchRetryable + ); + } + + /// # GIVEN + /// Various `Duration` values: whole, sub-second, zero, fractional. + /// + /// # WHEN + /// `floor_secs` is called. + /// + /// # THEN + /// It returns a whole-seconds string with a floor of `"1"` (so + /// `Duration::ZERO` and sub-second values do not collapse to `"0"`). + #[test] + fn floor_secs_floors_at_one_second() { + // Whole seconds at or above 1 pass through unchanged. + assert_eq!(floor_secs(Duration::from_secs(1)), "1"); + assert_eq!(floor_secs(Duration::from_secs(5)), "5"); + assert_eq!(floor_secs(Duration::from_secs(60)), "60"); + // ZERO is floored to 1 to avoid disabling the server-side timeout. + assert_eq!(floor_secs(Duration::ZERO), "1"); + // Sub-second values are floored to 1 (would otherwise truncate to 0). + assert_eq!(floor_secs(Duration::from_nanos(1)), "1"); + assert_eq!(floor_secs(Duration::from_millis(500)), "1"); + assert_eq!(floor_secs(Duration::from_millis(999)), "1"); + // Fractional seconds beyond 1s truncate to whole seconds (Duration::as_secs). + assert_eq!(floor_secs(Duration::from_millis(1500)), "1"); + assert_eq!(floor_secs(Duration::from_millis(2999)), "2"); + } } diff --git a/crates/etl-destinations/src/clickhouse/core.rs b/crates/etl-destinations/src/clickhouse/core.rs index 9c4436cd1..b3b850875 100644 --- a/crates/etl-destinations/src/clickhouse/core.rs +++ b/crates/etl-destinations/src/clickhouse/core.rs @@ -1,4 +1,4 @@ -use std::{collections::HashMap, sync::Arc}; +use std::{collections::HashMap, sync::Arc, time::Duration}; use etl::{ destination::{ @@ -170,6 +170,111 @@ impl Default for ClickHouseInserterConfig { } } +/// Configuration for the [`ClickHouseClient`]. +/// +/// Holds the server-side and client-side timeouts applied to each operation +/// bucket. Additional client-level knobs can be added here over time. +#[derive(Copy, Clone)] +pub struct ClickHouseClientConfig { + /// Server-side budget for the connectivity check (`SELECT 1`). + pub connectivity_check_timeout: Duration, + /// Server-side budget for schema lookups (`system.columns`). + pub schema_query_timeout: Duration, + /// Server-side budget for DDL (CREATE / ALTER / DROP / RENAME / TRUNCATE). + pub ddl_timeout: Duration, + /// Server-side budget per INSERT statement. Wraps `insert.end().await`, + /// which is the only awaited network step inside `insert_rows`; each + /// flushed chunk therefore gets its own deadline. + pub insert_timeout: Duration, + /// Slack added to the server-side budget to derive the client-side + /// `tokio::time::timeout`. + pub client_timeout_epsilon: Duration, +} + +impl ClickHouseClientConfig { + /// Default server-side budget for the connectivity check. + pub const DEFAULT_CONNECTIVITY_CHECK_TIMEOUT: Duration = Duration::from_secs(8); + /// Default server-side budget for schema lookups. + pub const DEFAULT_SCHEMA_QUERY_TIMEOUT: Duration = Duration::from_secs(16); + /// Default server-side budget for DDL. + pub const DEFAULT_DDL_TIMEOUT: Duration = Duration::from_secs(128); + /// Default server-side budget per INSERT statement. + pub const DEFAULT_INSERT_TIMEOUT: Duration = Duration::from_secs(256); + /// Default slack between server-side and client-side budgets. + pub const DEFAULT_CLIENT_TIMEOUT_EPSILON: Duration = Duration::from_secs(4); + + /// Server-side timeout for `op`. + pub(crate) fn server_timeout_for(&self, op: ClickHouseOperationKind) -> Duration { + match op { + ClickHouseOperationKind::ConnectivityCheck => self.connectivity_check_timeout, + ClickHouseOperationKind::SchemaQuery => self.schema_query_timeout, + ClickHouseOperationKind::Ddl => self.ddl_timeout, + ClickHouseOperationKind::Insert => self.insert_timeout, + } + } + + /// Client-side `tokio::time::timeout` for `op`: + /// `server_timeout_for(op) + client_timeout_epsilon`. + pub(crate) fn client_timeout_for(&self, op: ClickHouseOperationKind) -> Duration { + self.server_timeout_for(op) + self.client_timeout_epsilon + } +} + +impl Default for ClickHouseClientConfig { + fn default() -> Self { + Self { + connectivity_check_timeout: Self::DEFAULT_CONNECTIVITY_CHECK_TIMEOUT, + schema_query_timeout: Self::DEFAULT_SCHEMA_QUERY_TIMEOUT, + ddl_timeout: Self::DEFAULT_DDL_TIMEOUT, + insert_timeout: Self::DEFAULT_INSERT_TIMEOUT, + client_timeout_epsilon: Self::DEFAULT_CLIENT_TIMEOUT_EPSILON, + } + } +} + +/// Categories of ClickHouse client calls. +/// +/// Used to: +/// - select the corresponding server-side budget, +/// - map a generic clickhouse error onto the appropriate [`ErrorKind`]. +#[derive(Copy, Clone)] +pub(crate) enum ClickHouseOperationKind { + /// Connectivity check (`SELECT 1`). + ConnectivityCheck, + /// Schema lookup against `system.columns`. + SchemaQuery, + /// DDL: CREATE / ALTER / DROP / RENAME / TRUNCATE. + Ddl, + /// INSERT statement flush. + Insert, +} + +impl ClickHouseOperationKind { + /// Error kind used when the inner future returns a + /// `clickhouse::error::Error`. + pub(crate) fn failed_kind(self) -> ErrorKind { + match self { + ClickHouseOperationKind::ConnectivityCheck => ErrorKind::DestinationConnectionFailed, + ClickHouseOperationKind::SchemaQuery | ClickHouseOperationKind::Ddl => { + ErrorKind::DestinationQueryFailed + } + ClickHouseOperationKind::Insert => ErrorKind::DestinationAtomicBatchRetryable, + } + } +} + +impl std::fmt::Display for ClickHouseOperationKind { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let name = match self { + ClickHouseOperationKind::ConnectivityCheck => "connectivity check", + ClickHouseOperationKind::SchemaQuery => "schema query", + ClickHouseOperationKind::Ddl => "DDL", + ClickHouseOperationKind::Insert => "insert", + }; + f.write_str(name) + } +} + /// CDC-capable ClickHouse destination that replicates Postgres tables. /// /// Uses append-only MergeTree tables with two CDC columns (`cdc_operation`, @@ -209,11 +314,12 @@ where password: Option, database: impl Into, inserter_config: ClickHouseInserterConfig, + client_config: ClickHouseClientConfig, store: S, ) -> EtlResult { register_metrics(); Ok(Self { - client: ClickHouseClient::new(url, user, password, database), + client: ClickHouseClient::new(url, user, password, database, client_config), inserter_config, store: Arc::new(store), table_cache: Arc::new(RwLock::new(HashMap::new())), diff --git a/crates/etl-destinations/src/clickhouse/mod.rs b/crates/etl-destinations/src/clickhouse/mod.rs index 87768db26..fe58c088f 100644 --- a/crates/etl-destinations/src/clickhouse/mod.rs +++ b/crates/etl-destinations/src/clickhouse/mod.rs @@ -6,6 +6,6 @@ mod schema; #[cfg(feature = "test-utils")] pub mod test_utils; -pub use core::{ClickHouseDestination, ClickHouseInserterConfig}; +pub use core::{ClickHouseClientConfig, ClickHouseDestination, ClickHouseInserterConfig}; pub use client::ClickHouseClient; diff --git a/crates/etl-destinations/src/clickhouse/test_utils.rs b/crates/etl-destinations/src/clickhouse/test_utils.rs index cfb54a8af..def74209d 100644 --- a/crates/etl-destinations/src/clickhouse/test_utils.rs +++ b/crates/etl-destinations/src/clickhouse/test_utils.rs @@ -6,7 +6,7 @@ use tokio::runtime::Handle; use url::Url; use uuid::Uuid; -use crate::clickhouse::{ClickHouseDestination, ClickHouseInserterConfig}; +use crate::clickhouse::{ClickHouseClientConfig, ClickHouseDestination, ClickHouseInserterConfig}; /// ClickHouse HTTP URL (e.g. `http://localhost:8123`). pub const CLICKHOUSE_URL_ENV: &str = "TESTS_CLICKHOUSE_URL"; @@ -140,6 +140,7 @@ impl ClickHouseTestDatabase { self.password.clone(), &self.database, config, + ClickHouseClientConfig::default(), store, ) .expect("Failed to create ClickHouseDestination for test") diff --git a/crates/etl-destinations/tests/clickhouse/pipeline.rs b/crates/etl-destinations/tests/clickhouse/pipeline.rs index 0e5e4abd5..620e852dd 100644 --- a/crates/etl-destinations/tests/clickhouse/pipeline.rs +++ b/crates/etl-destinations/tests/clickhouse/pipeline.rs @@ -12,7 +12,7 @@ use etl::{ types::{EventType, PipelineId}, }; use etl_destinations::clickhouse::{ - ClickHouseInserterConfig, + ClickHouseClientConfig, ClickHouseInserterConfig, client::ClickHouseClient, test_utils::{ get_clickhouse_password, get_clickhouse_url, get_clickhouse_user, setup_clickhouse_database, @@ -1802,6 +1802,7 @@ async fn validate_connectivity_succeeds_against_running_clickhouse() { get_clickhouse_user(), get_clickhouse_password(), "default", + ClickHouseClientConfig::default(), ); assert!(client.validate_connectivity().await.is_ok()); } @@ -1821,6 +1822,7 @@ async fn validate_connectivity_fails_against_unreachable_clickhouse() { "nobody", None::, "default", + ClickHouseClientConfig::default(), ); assert!(client.validate_connectivity().await.is_err()); } diff --git a/crates/etl-examples/src/bin/clickhouse.rs b/crates/etl-examples/src/bin/clickhouse.rs index a06b8addc..3043a6114 100644 --- a/crates/etl-examples/src/bin/clickhouse.rs +++ b/crates/etl-examples/src/bin/clickhouse.rs @@ -47,7 +47,9 @@ use etl::{ pipeline::Pipeline, store::PostgresStore, }; -use etl_destinations::clickhouse::{ClickHouseDestination, ClickHouseInserterConfig}; +use etl_destinations::clickhouse::{ + ClickHouseClientConfig, ClickHouseDestination, ClickHouseInserterConfig, +}; use tokio::signal; use tracing::{error, info}; use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt}; @@ -207,6 +209,7 @@ async fn main_impl() -> Result<(), Box> { args.clickhouse_args.clickhouse_password, args.clickhouse_args.clickhouse_database, ClickHouseInserterConfig::default(), + ClickHouseClientConfig::default(), store.clone(), )?; diff --git a/crates/etl-replicator/src/core.rs b/crates/etl-replicator/src/core.rs index 93d6aa09f..c3808ca34 100644 --- a/crates/etl-replicator/src/core.rs +++ b/crates/etl-replicator/src/core.rs @@ -16,7 +16,7 @@ use etl_config::{ }; use etl_destinations::{ bigquery::BigQueryDestination, - clickhouse::{ClickHouseDestination, ClickHouseInserterConfig}, + clickhouse::{ClickHouseClientConfig, ClickHouseDestination, ClickHouseInserterConfig}, ducklake::{DuckLakeDestination, S3Config as DucklakeS3Config}, iceberg::{ DestinationNamespace, IcebergClient, IcebergDestination, S3_ACCESS_KEY_ID, S3_ENDPOINT, @@ -197,6 +197,7 @@ pub(crate) async fn start_replicator_with_config( password.as_ref().map(|p| p.expose_secret().to_owned()), database, ClickHouseInserterConfig::default(), + ClickHouseClientConfig::default(), state_store.clone(), )?; diff --git a/crates/etl/src/error.rs b/crates/etl/src/error.rs index f358978d0..677998717 100644 --- a/crates/etl/src/error.rs +++ b/crates/etl/src/error.rs @@ -88,6 +88,7 @@ pub enum ErrorKind { DestinationQueryFailed, DestinationAtomicBatchRetryable, SourceLockTimeout, + DestinationTimeout, SourceOperationCanceled, // Schema Errors diff --git a/crates/etl/src/state/table.rs b/crates/etl/src/state/table.rs index 6ae3d59ef..1352e640d 100644 --- a/crates/etl/src/state/table.rs +++ b/crates/etl/src/state/table.rs @@ -91,6 +91,7 @@ impl TableReplicationError { // Errors that can be retried automatically ErrorKind::SourceConnectionFailed | ErrorKind::DestinationConnectionFailed + | ErrorKind::DestinationTimeout | ErrorKind::SourceOperationCanceled | ErrorKind::SourceDatabaseShutdown | ErrorKind::SourceLockTimeout diff --git a/crates/etl/src/workers/policy.rs b/crates/etl/src/workers/policy.rs index cabc711b7..24148d41c 100644 --- a/crates/etl/src/workers/policy.rs +++ b/crates/etl/src/workers/policy.rs @@ -51,6 +51,7 @@ pub(crate) fn build_error_handling_policy(error: &EtlError) -> ErrorHandlingPoli ErrorKind::SourceConnectionFailed | ErrorKind::DestinationConnectionFailed | ErrorKind::DestinationAtomicBatchRetryable + | ErrorKind::DestinationTimeout | ErrorKind::SourceDatabaseShutdown | ErrorKind::SourceDatabaseInRecovery => { ErrorHandlingPolicy::new(RetryDirective::Timed, None) From ea2359155c3625844dd1492076c99992a64890fa Mon Sep 17 00:00:00 2001 From: Coenen Benjamin Date: Fri, 15 May 2026 11:28:31 +0200 Subject: [PATCH 06/29] experimental(ducklake): add support for external maintenances (#721) * experimental(ducklake): add support for external maintenances Signed-off-by: Benjamin <5719034+bnjjj@users.noreply.github.com> * update default values Signed-off-by: Benjamin <5719034+bnjjj@users.noreply.github.com> * fmt Signed-off-by: Benjamin <5719034+bnjjj@users.noreply.github.com> * improve behavior and add more tests Signed-off-by: Benjamin <5719034+bnjjj@users.noreply.github.com> * add logs * lint Signed-off-by: Benjamin <5719034+bnjjj@users.noreply.github.com> * fix fmt Signed-off-by: Benjamin <5719034+bnjjj@users.noreply.github.com> * fix tests Signed-off-by: Benjamin <5719034+bnjjj@users.noreply.github.com> * move ducklake maintenance binary Signed-off-by: Benjamin <5719034+bnjjj@users.noreply.github.com> * add more timeouts Signed-off-by: Benjamin <5719034+bnjjj@users.noreply.github.com> * remove in process maintenances Signed-off-by: Benjamin <5719034+bnjjj@users.noreply.github.com> * fix fmt Signed-off-by: Benjamin <5719034+bnjjj@users.noreply.github.com> * don't create useless CR when updating the pipeline Signed-off-by: Benjamin <5719034+bnjjj@users.noreply.github.com> * don't enable expire snapshot by default Signed-off-by: Benjamin <5719034+bnjjj@users.noreply.github.com> * change kubernetes app type for maintenance jobs Signed-off-by: Benjamin <5719034+bnjjj@users.noreply.github.com> * add docs for maintenance config Signed-off-by: Benjamin <5719034+bnjjj@users.noreply.github.com> --------- Signed-off-by: Benjamin <5719034+bnjjj@users.noreply.github.com> --- Cargo.lock | 14 + Cargo.toml | 1 + crates/etl-api/src/configs/pipeline.rs | 163 ++ crates/etl-api/src/k8s/base.rs | 55 +- crates/etl-api/src/k8s/cache.rs | 26 +- crates/etl-api/src/k8s/core.rs | 105 +- crates/etl-api/src/k8s/http.rs | 255 +- ...ate_bq_replicator_stateful_set_json-2.snap | 1 + ...ate_bq_replicator_stateful_set_json-3.snap | 1 + ...reate_bq_replicator_stateful_set_json.snap | 1 + ...sts__create_ducklake_maintenance_json.snap | 65 + ...cklake_replicator_stateful_set_json-2.snap | 1 + ...cklake_replicator_stateful_set_json-3.snap | 1 + ...ducklake_replicator_stateful_set_json.snap | 1 + ...ceberg_replicator_stateful_set_json-2.snap | 1 + ...ceberg_replicator_stateful_set_json-3.snap | 1 + ..._iceberg_replicator_stateful_set_json.snap | 1 + crates/etl-api/src/routes/pipelines.rs | 14 +- crates/etl-api/tests/pipelines.rs | 41 +- ...ination_and_pipeline_can_be_updated-2.snap | 1 + ...ination_and_pipeline_can_be_updated-2.snap | 1 + ...ination_and_pipeline_can_be_created-2.snap | 1 + ...ination_and_pipeline_can_be_created-2.snap | 1 + ..._pipelines__all_pipelines_can_be_read.snap | 1 + ...nes__an_existing_pipeline_can_be_read.snap | 1 + ...__an_existing_pipeline_can_be_updated.snap | 1 + crates/etl-api/tests/support/k8s_client.rs | 34 +- crates/etl-api/tests/support/mocks.rs | 23 + crates/etl-api/tests/validators.rs | 1 + crates/etl-destinations/Cargo.toml | 6 + .../etl-destinations/src/ducklake/METRICS.md | 81 +- .../etl-destinations/src/ducklake/client.rs | 1138 ++++++++- crates/etl-destinations/src/ducklake/core.rs | 426 ++-- .../src/ducklake/external_maintenance.rs | 771 ++++++ .../src/ducklake/maintenance.rs | 2265 ----------------- .../src/ducklake/maintenance_runner.rs | 898 +++++++ .../etl-destinations/src/ducklake/metrics.rs | 62 +- crates/etl-destinations/src/ducklake/mod.rs | 14 +- crates/etl-examples/Cargo.toml | 2 + crates/etl-replicator/Cargo.toml | 3 + crates/etl-replicator/Dockerfile | 3 +- .../src/bin/etl-ducklake-maintenance.rs | 222 ++ 42 files changed, 3952 insertions(+), 2752 deletions(-) create mode 100644 crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_ducklake_maintenance_json.snap create mode 100644 crates/etl-destinations/src/ducklake/external_maintenance.rs delete mode 100644 crates/etl-destinations/src/ducklake/maintenance.rs create mode 100644 crates/etl-destinations/src/ducklake/maintenance_runner.rs create mode 100644 crates/etl-replicator/src/bin/etl-ducklake-maintenance.rs diff --git a/Cargo.lock b/Cargo.lock index 2ce2b26fd..dfb1d0893 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1953,8 +1953,11 @@ dependencies = [ "etl-telemetry", "futures", "gcp-bigquery-client", + "humantime", "iceberg", "iceberg-catalog-rest", + "k8s-openapi", + "kube", "metrics", "parking_lot", "parquet", @@ -1984,6 +1987,8 @@ dependencies = [ "clap", "etl", "etl-destinations", + "etl-telemetry", + "k8s-openapi", "rustls", "tokio", "tracing", @@ -2010,11 +2015,13 @@ dependencies = [ name = "etl-replicator" version = "0.1.0" dependencies = [ + "chrono", "configcat", "etl", "etl-config", "etl-destinations", "etl-telemetry", + "k8s-openapi", "metrics", "reqwest", "rustls", @@ -2027,6 +2034,7 @@ dependencies = [ "tikv-jemallocator", "tokio", "tracing", + "tracing-subscriber", ] [[package]] @@ -2692,6 +2700,12 @@ version = "1.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" +[[package]] +name = "humantime" +version = "2.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "135b12329e5e3ce057a9f972339ea52bc954fe1e9358ef27f95e89716fbc5424" + [[package]] name = "hyper" version = "1.9.0" diff --git a/Cargo.toml b/Cargo.toml index 226ba5986..e252f8d07 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -88,6 +88,7 @@ etl-telemetry = { path = "crates/etl-telemetry", default-features = false } fail = { version = "0.5.1", default-features = false } futures = { version = "0.3.31", default-features = false } gcp-bigquery-client = { git = "https://github.com/iambriccardo/gcp-bigquery-client", rev = "c4fc59e338ca181d29b0dd53cac786fbe8513633", default-features = false } +humantime = { version = "2.3.0", default-features = false } iceberg = { version = "0.8.0", default-features = false } iceberg-catalog-rest = { version = "0.8.0", default-features = false } insta = { version = "1.43.1", default-features = false } diff --git a/crates/etl-api/src/configs/pipeline.rs b/crates/etl-api/src/configs/pipeline.rs index 36e0c6229..988671617 100644 --- a/crates/etl-api/src/configs/pipeline.rs +++ b/crates/etl-api/src/configs/pipeline.rs @@ -60,6 +60,112 @@ impl ReplicatorResourcesConfig { } } +/// DuckLake maintenance controller settings for one pipeline. +#[derive(Debug, Clone, Serialize, Deserialize, ToSchema, PartialEq)] +pub struct DuckLakeMaintenanceConfig { + /// Minimum time between maintenance runs, in seconds. + #[schema(example = 3600)] + #[serde(default = "default_ducklake_maintenance_min_interval_seconds")] + pub min_interval_seconds: u64, + /// Maximum time replication may be paused for one maintenance run, in + /// seconds. + #[schema(example = 2700)] + #[serde(default = "default_ducklake_maintenance_max_pause_seconds")] + pub max_pause_seconds: u64, + /// Minimum inlined bytes required before inline flush runs. + #[schema(example = 10000000)] + #[serde(default = "default_ducklake_maintenance_min_inlined_bytes")] + pub min_inlined_bytes: u64, + /// Maximum number of adjacent files compacted by one merge operation. + #[schema(example = 32)] + #[serde(default = "default_ducklake_maintenance_max_compacted_files")] + pub max_compacted_files: u32, + /// Maximum number of tables processed by each operation in one run. + #[schema(example = 8)] + #[serde(default = "default_ducklake_maintenance_max_tables_per_run")] + pub max_tables_per_run: u32, + /// DuckLake target file size used for compaction. + #[schema(example = "10MB")] + #[serde(default = "default_ducklake_maintenance_target_file_size")] + pub target_file_size: String, + /// Deleted-row fraction that triggers data file rewrite. + #[schema(example = 0.5)] + #[serde(default = "default_ducklake_maintenance_delete_threshold")] + pub delete_threshold: f64, + /// Minimum active data files required before data file rewrite runs. + #[schema(example = 40)] + #[serde(default = "default_ducklake_maintenance_min_active_data_files")] + pub min_active_data_files: i64, + /// CPU request for maintenance jobs, in millicores. + #[schema(example = 1000)] + #[serde(default = "default_ducklake_maintenance_cpu_request_millicores")] + pub cpu_request_millicores: u32, + /// Memory request for maintenance jobs, in MiB. + #[schema(example = 1024)] + #[serde(default = "default_ducklake_maintenance_memory_request_mib")] + pub memory_request_mib: u32, + /// Maximum runtime for one maintenance job, in seconds. + #[schema(example = 1800)] + #[serde(default = "default_ducklake_maintenance_active_deadline_seconds")] + pub active_deadline_seconds: i64, +} + +impl Default for DuckLakeMaintenanceConfig { + fn default() -> Self { + Self { + min_interval_seconds: default_ducklake_maintenance_min_interval_seconds(), + max_pause_seconds: default_ducklake_maintenance_max_pause_seconds(), + min_inlined_bytes: default_ducklake_maintenance_min_inlined_bytes(), + max_compacted_files: default_ducklake_maintenance_max_compacted_files(), + max_tables_per_run: default_ducklake_maintenance_max_tables_per_run(), + target_file_size: default_ducklake_maintenance_target_file_size(), + delete_threshold: default_ducklake_maintenance_delete_threshold(), + min_active_data_files: default_ducklake_maintenance_min_active_data_files(), + cpu_request_millicores: default_ducklake_maintenance_cpu_request_millicores(), + memory_request_mib: default_ducklake_maintenance_memory_request_mib(), + active_deadline_seconds: default_ducklake_maintenance_active_deadline_seconds(), + } + } +} + +impl DuckLakeMaintenanceConfig { + /// Validates maintenance controller settings. + pub fn validate(&self) -> Result<(), String> { + if self.min_interval_seconds == 0 { + return Err("ducklake maintenance min interval must be greater than 0".to_owned()); + } + if self.max_pause_seconds == 0 { + return Err("ducklake maintenance max pause must be greater than 0".to_owned()); + } + if self.max_compacted_files == 0 { + return Err( + "ducklake maintenance max compacted files must be greater than 0".to_owned() + ); + } + if self.max_tables_per_run == 0 { + return Err("ducklake maintenance max tables per run must be greater than 0".to_owned()); + } + if !(0.0..=1.0).contains(&self.delete_threshold) { + return Err("ducklake maintenance delete threshold must be between 0 and 1".to_owned()); + } + if self.min_active_data_files < 0 { + return Err("ducklake maintenance min active data files must be greater than or \ + equal to 0" + .to_owned()); + } + if self.cpu_request_millicores == 0 { + return Err("ducklake maintenance cpu request must be greater than 0".to_owned()); + } + if self.memory_request_mib == 0 { + return Err("ducklake maintenance memory request must be greater than 0".to_owned()); + } + if self.active_deadline_seconds <= 0 { + return Err("ducklake maintenance active deadline must be greater than 0".to_owned()); + } + Ok(()) + } +} + #[derive(Debug, Clone, Serialize, Deserialize, ToSchema)] pub struct FullApiPipelineConfig { #[schema(example = "my_publication")] @@ -90,6 +196,8 @@ pub struct FullApiPipelineConfig { pub invalidated_slot_behavior: Option, #[serde(skip_serializing_if = "Option::is_none")] pub replicator_resources: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub ducklake_maintenance: Option, pub log_level: Option, } @@ -99,6 +207,9 @@ impl FullApiPipelineConfig { if let Some(replicator_resources) = &self.replicator_resources { replicator_resources.validate()?; } + if let Some(ducklake_maintenance) = &self.ducklake_maintenance { + ducklake_maintenance.validate()?; + } Ok(()) } @@ -118,6 +229,7 @@ impl From for FullApiPipelineConfig { table_sync_copy: Some(value.table_sync_copy), invalidated_slot_behavior: Some(value.invalidated_slot_behavior), replicator_resources: value.replicator_resources, + ducklake_maintenance: value.ducklake_maintenance, log_level: value.log_level, } } @@ -146,6 +258,8 @@ pub struct StoredPipelineConfig { pub invalidated_slot_behavior: InvalidatedSlotBehavior, #[serde(default)] pub replicator_resources: Option, + #[serde(default)] + pub ducklake_maintenance: Option, pub log_level: Option, } @@ -200,11 +314,56 @@ impl From for StoredPipelineConfig { table_sync_copy: value.table_sync_copy.unwrap_or_default(), invalidated_slot_behavior: value.invalidated_slot_behavior.unwrap_or_default(), replicator_resources: value.replicator_resources, + ducklake_maintenance: value.ducklake_maintenance, log_level: value.log_level, } } } +fn default_ducklake_maintenance_min_interval_seconds() -> u64 { + 3600 +} + +fn default_ducklake_maintenance_max_pause_seconds() -> u64 { + 2700 +} + +fn default_ducklake_maintenance_min_inlined_bytes() -> u64 { + 10_000_000 +} + +fn default_ducklake_maintenance_max_compacted_files() -> u32 { + 32 +} + +fn default_ducklake_maintenance_max_tables_per_run() -> u32 { + 8 +} + +fn default_ducklake_maintenance_target_file_size() -> String { + "10MB".to_owned() +} + +fn default_ducklake_maintenance_delete_threshold() -> f64 { + 0.5 +} + +fn default_ducklake_maintenance_min_active_data_files() -> i64 { + 40 +} + +fn default_ducklake_maintenance_cpu_request_millicores() -> u32 { + 1000 +} + +fn default_ducklake_maintenance_memory_request_mib() -> u32 { + 1024 +} + +fn default_ducklake_maintenance_active_deadline_seconds() -> i64 { + 1800 +} + #[cfg(test)] mod tests { use etl_config::shared::BatchConfig; @@ -234,6 +393,7 @@ mod tests { cpu_request_millicores: Some(500), memory_request_mib: Some(2000), }), + ducklake_maintenance: None, log_level: None, invalidated_slot_behavior: InvalidatedSlotBehavior::Error, }; @@ -271,6 +431,7 @@ mod tests { cpu_request_millicores: Some(500), memory_request_mib: Some(2000), }), + ducklake_maintenance: None, log_level: Some(LogLevel::Debug), }; @@ -294,6 +455,7 @@ mod tests { table_sync_copy: None, invalidated_slot_behavior: None, replicator_resources: None, + ducklake_maintenance: None, log_level: None, }; @@ -344,6 +506,7 @@ mod tests { cpu_request_millicores: Some(500), memory_request_mib: Some(2000), }), + ducklake_maintenance: None, log_level: None, invalidated_slot_behavior: InvalidatedSlotBehavior::Error, }; diff --git a/crates/etl-api/src/k8s/base.rs b/crates/etl-api/src/k8s/base.rs index 3bfaf4014..2b5318e9c 100644 --- a/crates/etl-api/src/k8s/base.rs +++ b/crates/etl-api/src/k8s/base.rs @@ -4,7 +4,9 @@ use k8s_openapi::api::core::v1::ConfigMap; use thiserror::Error; use crate::configs::{ - destination::StoredDestinationConfig, log::LogLevel, pipeline::ReplicatorResourcesConfig, + destination::StoredDestinationConfig, + log::LogLevel, + pipeline::{DuckLakeMaintenanceConfig, ReplicatorResourcesConfig}, }; /// Errors from Kubernetes operations. @@ -35,6 +37,40 @@ pub struct ReplicatorConfigMapFile { pub content: String, } +/// DuckLake maintenance CR materialization input. +#[derive(Debug, Clone)] +pub struct DuckLakeMaintenanceResourceConfig { + /// Tenant/project identifier. + pub tenant_id: String, + /// Pipeline id. + pub pipeline_id: i64, + /// Replicator id. + pub replicator_id: i64, + /// Image containing the maintenance binary. + pub image: String, + /// User-authored maintenance policy. + pub policy: DuckLakeMaintenanceConfig, +} + +/// Replicator StatefulSet materialization input. +#[derive(Debug, Clone)] +pub struct ReplicatorStatefulSetConfig { + /// Existing Kubernetes resource prefix. + pub prefix: String, + /// Image for the replicator container. + pub replicator_image: String, + /// Deployment environment. + pub environment: Environment, + /// Optional resource overrides. + pub replicator_resources: Option, + /// Destination type used to select destination-specific env/secrets. + pub destination_type: DestinationType, + /// DuckLake maintenance policy. + pub ducklake_maintenance: Option, + /// Replicator log level. + pub log_level: LogLevel, +} + /// The type of destination storage system for replication. /// /// Determines which destination-specific resources and configurations are @@ -230,12 +266,7 @@ pub trait K8sClient: Send + Sync { /// the pods. async fn create_or_update_replicator_stateful_set( &self, - prefix: &str, - replicator_image: &str, - environment: Environment, - replicator_resources: Option<&ReplicatorResourcesConfig>, - destination_type: DestinationType, - log_level: LogLevel, + config: ReplicatorStatefulSetConfig, ) -> Result<(), K8sError>; /// Deletes the replicator [`StatefulSet`]. @@ -243,6 +274,16 @@ pub trait K8sClient: Send + Sync { /// Does nothing if the stateful set does not exist. async fn delete_replicator_stateful_set(&self, prefix: &str) -> Result<(), K8sError>; + /// Creates or updates the DuckLake maintenance CR. + async fn create_or_update_ducklake_maintenance( + &self, + prefix: &str, + config: DuckLakeMaintenanceResourceConfig, + ) -> Result<(), K8sError>; + + /// Deletes the DuckLake maintenance CR. + async fn delete_ducklake_maintenance(&self, prefix: &str) -> Result<(), K8sError>; + /// Retrieves the current status of a replicator pod. /// /// Returns a [`PodStatus`] derived from the pod's phase, deletion diff --git a/crates/etl-api/src/k8s/cache.rs b/crates/etl-api/src/k8s/cache.rs index 7182b0b4d..68161be3e 100644 --- a/crates/etl-api/src/k8s/cache.rs +++ b/crates/etl-api/src/k8s/cache.rs @@ -126,13 +126,12 @@ mod tests { use std::collections::BTreeMap; use async_trait::async_trait; - use etl_config::Environment; use k8s_openapi::api::core::v1::ConfigMap; use super::*; - use crate::{ - configs::log::LogLevel, - k8s::{DestinationType, PodStatus, ReplicatorConfigMapFile}, + use crate::k8s::{ + DuckLakeMaintenanceResourceConfig, PodStatus, ReplicatorConfigMapFile, + ReplicatorStatefulSetConfig, }; struct MockK8sClient { @@ -228,12 +227,7 @@ mod tests { async fn create_or_update_replicator_stateful_set( &self, - _prefix: &str, - _replicator_image: &str, - _environment: Environment, - _replicator_resources: Option<&crate::configs::pipeline::ReplicatorResourcesConfig>, - _destination_type: DestinationType, - _log_level: LogLevel, + _config: ReplicatorStatefulSetConfig, ) -> Result<(), K8sError> { Ok(()) } @@ -242,6 +236,18 @@ mod tests { Ok(()) } + async fn create_or_update_ducklake_maintenance( + &self, + _prefix: &str, + _config: DuckLakeMaintenanceResourceConfig, + ) -> Result<(), K8sError> { + Ok(()) + } + + async fn delete_ducklake_maintenance(&self, _prefix: &str) -> Result<(), K8sError> { + Ok(()) + } + async fn get_replicator_pod_status(&self, _prefix: &str) -> Result { Ok(PodStatus::Started) } diff --git a/crates/etl-api/src/k8s/core.rs b/crates/etl-api/src/k8s/core.rs index c97b72c13..3ceeeb2d7 100644 --- a/crates/etl-api/src/k8s/core.rs +++ b/crates/etl-api/src/k8s/core.rs @@ -8,8 +8,7 @@ use thiserror::Error; use crate::{ configs::{ destination::{StoredDestinationConfig, StoredIcebergConfig}, - log::LogLevel, - pipeline::{ReplicatorResourcesConfig, StoredPipelineConfig}, + pipeline::StoredPipelineConfig, source::StoredSourceConfig, }, data::{ @@ -19,7 +18,10 @@ use crate::{ replicators::Replicator, sources::Source, }, - k8s::{DestinationType, K8sClient, K8sError, PodStatus, ReplicatorConfigMapFile}, + k8s::{ + DestinationType, DuckLakeMaintenanceResourceConfig, K8sClient, K8sError, PodStatus, + ReplicatorConfigMapFile, ReplicatorStatefulSetConfig, + }, }; /// Errors raised while preparing or applying Kubernetes pipeline resources. @@ -112,6 +114,7 @@ pub async fn create_or_update_pipeline_resources_in_k8s( let log_level = pipeline.config.log_level.clone().unwrap_or_default(); let replicator_resources = pipeline.config.replicator_resources.clone(); + let ducklake_maintenance = pipeline.config.ducklake_maintenance.clone(); let replicator_config = build_replicator_config_without_secrets( // We are safe to perform this conversion, since the i64 -> u64 conversion performs wrap // around, and we won't have two different values map to the same u64, since the domain @@ -126,14 +129,32 @@ pub async fn create_or_update_pipeline_resources_in_k8s( create_or_update_dynamic_replicator_secrets(k8s_client, &prefix, secrets).await?; create_or_update_replicator_config(k8s_client, &prefix, replicator_config, environment).await?; - create_or_update_replicator_stateful_set( + let replicator_image = image.name; + let ducklake_maintenance_for_replicator = matches!(destination_type, DestinationType::Ducklake) + .then(|| ducklake_maintenance.clone().unwrap_or_default()); + + create_or_update_ducklake_maintenance( k8s_client, &prefix, - image.name, - environment, - replicator_resources.as_ref(), + tenant_id, + pipeline.id, + replicator.id, + &replicator_image, destination_type, - log_level, + ducklake_maintenance, + ) + .await?; + create_or_update_replicator_stateful_set( + k8s_client, + ReplicatorStatefulSetConfig { + prefix, + replicator_image, + environment, + replicator_resources, + destination_type, + ducklake_maintenance: ducklake_maintenance_for_replicator, + log_level, + }, ) .await?; @@ -152,6 +173,7 @@ pub async fn delete_pipeline_resources_in_k8s( ) -> Result<(), K8sCoreError> { let prefix = create_k8s_object_prefix(tenant_id, replicator.id); + k8s_client.delete_ducklake_maintenance(&prefix).await?; delete_dynamic_replicator_secrets(k8s_client, &prefix).await?; delete_replicator_config(k8s_client, &prefix).await?; delete_replicator_stateful_set(k8s_client, &prefix).await?; @@ -389,22 +411,42 @@ async fn create_or_update_replicator_config( /// environment-specific settings and destination-type-specific resource /// requirements. async fn create_or_update_replicator_stateful_set( + k8s_client: &dyn K8sClient, + config: ReplicatorStatefulSetConfig, +) -> Result<(), K8sCoreError> { + k8s_client.create_or_update_replicator_stateful_set(config).await?; + + Ok(()) +} + +/// Creates, updates, or deletes the DuckLake maintenance CR. +#[allow(clippy::too_many_arguments)] +async fn create_or_update_ducklake_maintenance( k8s_client: &dyn K8sClient, prefix: &str, - replicator_image: String, - environment: Environment, - replicator_resources: Option<&ReplicatorResourcesConfig>, + tenant_id: &str, + pipeline_id: i64, + replicator_id: i64, + replicator_image: &str, destination_type: DestinationType, - log_level: LogLevel, + ducklake_maintenance: Option, ) -> Result<(), K8sCoreError> { + if !matches!(destination_type, DestinationType::Ducklake) { + k8s_client.delete_ducklake_maintenance(prefix).await?; + return Ok(()); + } + + let policy = ducklake_maintenance.unwrap_or_default(); k8s_client - .create_or_update_replicator_stateful_set( + .create_or_update_ducklake_maintenance( prefix, - &replicator_image, - environment, - replicator_resources, - destination_type, - log_level, + DuckLakeMaintenanceResourceConfig { + tenant_id: tenant_id.to_owned(), + pipeline_id, + replicator_id, + image: replicator_image.to_owned(), + policy, + }, ) .await?; @@ -466,10 +508,8 @@ mod tests { use super::*; use crate::{ - configs::{ - destination::StoredDestinationConfig, log::LogLevel, source::StoredSourceConfig, - }, - k8s::{DestinationType, K8sClient, K8sError, PodStatus, ReplicatorConfigMapFile}, + configs::{destination::StoredDestinationConfig, source::StoredSourceConfig}, + k8s::{K8sClient, K8sError, PodStatus, ReplicatorConfigMapFile}, }; #[derive(Debug, Clone)] @@ -604,12 +644,7 @@ mod tests { async fn create_or_update_replicator_stateful_set( &self, - _prefix: &str, - _replicator_image: &str, - _environment: Environment, - _replicator_resources: Option<&ReplicatorResourcesConfig>, - _destination_type: DestinationType, - _log_level: LogLevel, + _config: ReplicatorStatefulSetConfig, ) -> Result<(), K8sError> { Ok(()) } @@ -618,6 +653,20 @@ mod tests { Ok(()) } + async fn create_or_update_ducklake_maintenance( + &self, + prefix: &str, + _config: DuckLakeMaintenanceResourceConfig, + ) -> Result<(), K8sError> { + self.calls.lock().unwrap().push(format!("ducklake-maintenance:{prefix}")); + Ok(()) + } + + async fn delete_ducklake_maintenance(&self, prefix: &str) -> Result<(), K8sError> { + self.calls.lock().unwrap().push(format!("delete-ducklake-maintenance:{prefix}")); + Ok(()) + } + async fn get_replicator_pod_status(&self, _prefix: &str) -> Result { Ok(self.pod_status) } diff --git a/crates/etl-api/src/k8s/http.rs b/crates/etl-api/src/k8s/http.rs index 22bdef955..69ce622c7 100644 --- a/crates/etl-api/src/k8s/http.rs +++ b/crates/etl-api/src/k8s/http.rs @@ -14,14 +14,21 @@ use k8s_openapi::{ use kube::{ Client, api::{Api, DeleteParams, Patch, PatchParams}, + core::{ApiResource, DynamicObject, GroupVersionKind}, }; use serde_json::json; use tracing::debug; use crate::{ config::K8sConfig, - configs::{log::LogLevel, pipeline::ReplicatorResourcesConfig}, - k8s::{DestinationType, K8sClient, K8sError, PodPhase, PodStatus, ReplicatorConfigMapFile}, + configs::{ + log::LogLevel, + pipeline::{DuckLakeMaintenanceConfig, ReplicatorResourcesConfig}, + }, + k8s::{ + DestinationType, DuckLakeMaintenanceResourceConfig, K8sClient, K8sError, PodPhase, + PodStatus, ReplicatorConfigMapFile, ReplicatorStatefulSetConfig, + }, }; /// Secret name suffix for the BigQuery service account key. @@ -86,6 +93,16 @@ pub const TRUSTED_ROOT_CERT_CONFIG_MAP_NAME: &str = "trusted-root-certs-config"; pub const TRUSTED_ROOT_CERT_KEY_NAME: &str = "trusted_root_certs"; /// Label used to identify replicator pods. const REPLICATOR_APP_LABEL: &str = "etl-replicator-app"; +/// Label used to identify DuckLake maintenance resources. +const DUCKLAKE_MAINTENANCE_APP_LABEL: &str = "etl-ducklake-maintenance-app"; +/// ServiceAccount used by replicator pods for runtime coordination. +const REPLICATOR_SERVICE_ACCOUNT_NAME: &str = "etl-replicator"; +/// DuckLake maintenance CRD group. +const DUCKLAKE_MAINTENANCE_GROUP: &str = "etl.supabase.com"; +/// DuckLake maintenance CRD version. +const DUCKLAKE_MAINTENANCE_VERSION: &str = "v1alpha1"; +/// DuckLake maintenance CRD kind. +const DUCKLAKE_MAINTENANCE_KIND: &str = "DuckLakeMaintenance"; /// Default replicator memory request in prod, in Mi. const REPLICATOR_MEMORY_REQUEST_PROD_DEFAULT: i32 = 500; @@ -204,6 +221,7 @@ pub struct HttpK8sClient { config_maps_api: Api, stateful_sets_api: Api, pods_api: Api, + ducklake_maintenance_api: Api, k8s_config: K8sConfig, } @@ -217,9 +235,21 @@ impl HttpK8sClient { let config_maps_api: Api = Api::namespaced(client.clone(), DATA_PLANE_NAMESPACE); let stateful_sets_api: Api = Api::namespaced(client.clone(), DATA_PLANE_NAMESPACE); - let pods_api: Api = Api::namespaced(client, DATA_PLANE_NAMESPACE); + let pods_api: Api = Api::namespaced(client.clone(), DATA_PLANE_NAMESPACE); + let ducklake_maintenance_api: Api = Api::namespaced_with( + client, + DATA_PLANE_NAMESPACE, + &ducklake_maintenance_api_resource(), + ); - Ok(HttpK8sClient { secrets_api, config_maps_api, stateful_sets_api, pods_api, k8s_config }) + Ok(HttpK8sClient { + secrets_api, + config_maps_api, + stateful_sets_api, + pods_api, + ducklake_maintenance_api, + k8s_config, + }) } /// Helper function to handle delete operations that should ignore 404 @@ -544,35 +574,33 @@ impl K8sClient for HttpK8sClient { async fn create_or_update_replicator_stateful_set( &self, - prefix: &str, - replicator_image: &str, - environment: Environment, - replicator_resources: Option<&ReplicatorResourcesConfig>, - destination_type: DestinationType, - log_level: LogLevel, + request: ReplicatorStatefulSetConfig, ) -> Result<(), K8sError> { debug!("patching stateful set"); + let prefix = request.prefix.as_str(); + let replicator_image = request.replicator_image.as_str(); let config = ReplicatorResourceConfig::load_with_overrides( - &environment, + &request.environment, &self.k8s_config, - replicator_resources, + request.replicator_resources.as_ref(), )?; let stateful_set_name = create_stateful_set_name(prefix); let container_environment = create_container_environment_json( prefix, - &environment, + &request.environment, replicator_image, - destination_type, - log_level, + request.destination_type, + request.ducklake_maintenance.as_ref(), + request.log_level, ); - let node_selector = create_node_selector_json(&environment); - let init_containers = create_init_containers_json(prefix, &environment, &config); - let volumes = create_volumes_json(prefix, &environment); - let volume_mounts = create_volume_mounts_json(&environment); + let node_selector = create_node_selector_json(&request.environment); + let init_containers = create_init_containers_json(prefix, &request.environment, &config); + let volumes = create_volumes_json(prefix, &request.environment); + let volume_mounts = create_volume_mounts_json(&request.environment); let stateful_set_json = create_replicator_stateful_set_json( prefix, @@ -609,6 +637,35 @@ impl K8sClient for HttpK8sClient { Ok(()) } + async fn create_or_update_ducklake_maintenance( + &self, + prefix: &str, + config: DuckLakeMaintenanceResourceConfig, + ) -> Result<(), K8sError> { + debug!("patching ducklake maintenance"); + + let name = create_ducklake_maintenance_name(prefix); + let ducklake_maintenance_json = create_ducklake_maintenance_json(prefix, &name, config); + let pp = PatchParams::apply(&name).force(); + self.ducklake_maintenance_api + .patch(&name, &pp, &Patch::Apply(ducklake_maintenance_json)) + .await?; + + Ok(()) + } + + async fn delete_ducklake_maintenance(&self, prefix: &str) -> Result<(), K8sError> { + debug!("deleting ducklake maintenance"); + + let name = create_ducklake_maintenance_name(prefix); + let dp = DeleteParams::default(); + Self::handle_delete_with_404_ignore( + self.ducklake_maintenance_api.delete(&name, &dp).await, + )?; + + Ok(()) + } + async fn get_replicator_pod_status(&self, prefix: &str) -> Result { debug!("getting pod status"); @@ -667,6 +724,10 @@ fn create_ducklake_secret_name(prefix: &str) -> String { format!("{prefix}-{DUCKLAKE_SECRET_NAME_SUFFIX}") } +fn create_ducklake_maintenance_name(prefix: &str) -> String { + prefix.to_owned() +} + fn create_replicator_config_map_name(prefix: &str) -> String { format!("{prefix}-{REPLICATOR_CONFIG_MAP_NAME_SUFFIX}") } @@ -683,6 +744,14 @@ fn create_replicator_app_name(prefix: &str) -> String { format!("{prefix}-{REPLICATOR_APP_SUFFIX}") } +fn ducklake_maintenance_api_resource() -> ApiResource { + ApiResource::from_gvk(&GroupVersionKind::gvk( + DUCKLAKE_MAINTENANCE_GROUP, + DUCKLAKE_MAINTENANCE_VERSION, + DUCKLAKE_MAINTENANCE_KIND, + )) +} + fn create_replicator_container_name(prefix: &str) -> String { format!("{prefix}-{REPLICATOR_CONTAINER_NAME_SUFFIX}") } @@ -838,11 +907,84 @@ fn create_replicator_config_map_json( }) } +fn create_ducklake_maintenance_json( + prefix: &str, + name: &str, + config: DuckLakeMaintenanceResourceConfig, +) -> serde_json::Value { + let replicator_app_name = create_replicator_app_name(prefix); + let postgres_secret_name = create_postgres_secret_name(prefix); + let ducklake_secret_name = create_ducklake_secret_name(prefix); + let config_map_name = create_replicator_config_map_name(prefix); + json!({ + "apiVersion": format!("{DUCKLAKE_MAINTENANCE_GROUP}/{DUCKLAKE_MAINTENANCE_VERSION}"), + "kind": DUCKLAKE_MAINTENANCE_KIND, + "metadata": { + "name": name, + "namespace": DATA_PLANE_NAMESPACE, + "labels": { + "etl.supabase.com/app-name": replicator_app_name, + "etl.supabase.com/app-type": DUCKLAKE_MAINTENANCE_APP_LABEL, + } + }, + "spec": { + "pipelineRef": { + "tenantId": config.tenant_id, + "pipelineId": config.pipeline_id, + "replicatorId": config.replicator_id, + }, + "schedule": { + "minIntervalSeconds": config.policy.min_interval_seconds, + }, + "pause": { + "maxDurationSeconds": config.policy.max_pause_seconds, + }, + "operations": { + "inlineFlush": { + "enabled": true, + "minInlinedBytes": config.policy.min_inlined_bytes, + }, + "mergeAdjacentFiles": { + "enabled": true, + "maxCompactedFiles": config.policy.max_compacted_files, + "maxTablesPerRun": config.policy.max_tables_per_run, + "targetFileSize": config.policy.target_file_size, + }, + "rewriteDataFiles": { + "enabled": true, + "deleteThreshold": config.policy.delete_threshold, + "maxTablesPerRun": config.policy.max_tables_per_run, + }, + "expireSnapshots": { + "enabled": false, + }, + "cleanupOldFiles": { + "enabled": true, + } + }, + "jobTemplate": { + "image": config.image, + "cpuRequestMillicores": config.policy.cpu_request_millicores, + "memoryRequestMiB": config.policy.memory_request_mib, + "activeDeadlineSeconds": config.policy.active_deadline_seconds, + "backoffLimit": 1, + "ttlSecondsAfterFinished": 86400, + }, + "runtimeRefs": { + "configMapName": config_map_name, + "postgresSecretName": postgres_secret_name, + "ducklakeSecretName": ducklake_secret_name, + } + } + }) +} + fn create_container_environment_json( prefix: &str, environment: &Environment, replicator_image: &str, destination_type: DestinationType, + ducklake_maintenance: Option<&DuckLakeMaintenanceConfig>, log_level: LogLevel, ) -> Vec { let mut container_environment = vec![ @@ -965,6 +1107,25 @@ fn create_container_environment_json( let ducklake_s3_secret_access_key_env_var_json = create_ducklake_s3_secret_access_key_env_var_json(&ducklake_secret_name); container_environment.push(ducklake_s3_secret_access_key_env_var_json); + + if let Some(ducklake_maintenance) = ducklake_maintenance { + container_environment.push(json!({ + "name": "ETL_DUCKLAKE_MAINTENANCE_CR_NAME", + "value": create_ducklake_maintenance_name(prefix) + })); + container_environment.push(json!({ + "name": "ETL_DUCKLAKE_MAINTENANCE_CR_NAMESPACE", + "value": DATA_PLANE_NAMESPACE + })); + container_environment.push(json!({ + "name": "ETL_DUCKLAKE_EXTERNAL_MAINTENANCE_INLINE_FLUSH_MIN_INLINED_BYTES", + "value": ducklake_maintenance.min_inlined_bytes.to_string() + })); + container_environment.push(json!({ + "name": "ETL_DUCKLAKE_EXTERNAL_MAINTENANCE_REWRITE_DATA_FILES_MIN_ACTIVE_DATA_FILES", + "value": ducklake_maintenance.min_active_data_files.to_string() + })); + } } } container_environment @@ -1236,6 +1397,7 @@ fn create_replicator_stateful_set_json( } }, "spec": { + "serviceAccountName": REPLICATOR_SERVICE_ACCOUNT_NAME, "volumes": volumes, // Allow scheduling onto nodes tainted with the right node role. "tolerations": [ @@ -1523,6 +1685,41 @@ mod tests { let _config_map: ConfigMap = serde_json::from_value(config_map_json).unwrap(); } + #[test] + fn test_create_ducklake_maintenance_json() { + let prefix = create_k8s_object_prefix(TENANT_ID, 42); + let name = create_ducklake_maintenance_name(&prefix); + + let ducklake_maintenance_json = create_ducklake_maintenance_json( + &prefix, + &name, + DuckLakeMaintenanceResourceConfig { + tenant_id: TENANT_ID.to_owned(), + pipeline_id: 24, + replicator_id: 42, + image: "supabase/replicator:1.2.3".to_owned(), + policy: DuckLakeMaintenanceConfig { + min_interval_seconds: 3600, + max_pause_seconds: 2700, + min_inlined_bytes: 10_000_000, + max_compacted_files: 32, + max_tables_per_run: 8, + target_file_size: "10MB".to_owned(), + delete_threshold: 0.5, + min_active_data_files: 40, + cpu_request_millicores: 1000, + memory_request_mib: 1024, + active_deadline_seconds: 1800, + }, + }, + ); + + assert_snapshot!(serde_json::to_string_pretty(&ducklake_maintenance_json).unwrap()); + + let _ducklake_maintenance: DynamicObject = + serde_json::from_value(ducklake_maintenance_json).unwrap(); + } + #[test] fn test_create_postgres_secret_env_var_json() { let prefix = create_k8s_object_prefix(TENANT_ID, 42); @@ -1587,6 +1784,7 @@ mod tests { &environment, replicator_image, DestinationType::BigQuery, + None, LogLevel::Info, ); assert_json_snapshot!(container_environment); @@ -1597,6 +1795,7 @@ mod tests { &environment, replicator_image, DestinationType::BigQuery, + None, LogLevel::Info, ); assert_json_snapshot!(container_environment); @@ -1607,6 +1806,7 @@ mod tests { &environment, replicator_image, DestinationType::BigQuery, + None, LogLevel::Info, ); assert_json_snapshot!(container_environment); @@ -1622,6 +1822,7 @@ mod tests { &Environment::Dev, replicator_image, DestinationType::Iceberg, + None, LogLevel::Info, ); assert_json_snapshot!(container_environment); @@ -1631,6 +1832,7 @@ mod tests { &Environment::Staging, replicator_image, DestinationType::Iceberg, + None, LogLevel::Info, ); assert_json_snapshot!(container_environment); @@ -1640,6 +1842,7 @@ mod tests { &Environment::Prod, replicator_image, DestinationType::Iceberg, + None, LogLevel::Info, ); assert_json_snapshot!(container_environment); @@ -1655,6 +1858,7 @@ mod tests { &Environment::Dev, replicator_image, DestinationType::Ducklake, + None, LogLevel::Info, ); assert_json_snapshot!(container_environment); @@ -1664,6 +1868,7 @@ mod tests { &Environment::Staging, replicator_image, DestinationType::Ducklake, + None, LogLevel::Info, ); assert_json_snapshot!(container_environment); @@ -1673,6 +1878,7 @@ mod tests { &Environment::Prod, replicator_image, DestinationType::Ducklake, + None, LogLevel::Info, ); assert_json_snapshot!(container_environment); @@ -1688,6 +1894,7 @@ mod tests { &Environment::Dev, replicator_image, DestinationType::ClickHouse { password_secret_required: true }, + None, LogLevel::Info, ); @@ -1711,6 +1918,7 @@ mod tests { &Environment::Dev, replicator_image, DestinationType::ClickHouse { password_secret_required: false }, + None, LogLevel::Info, ); @@ -1803,6 +2011,7 @@ mod tests { &environment, replicator_image, DestinationType::BigQuery, + None, LogLevel::Info, ); @@ -1835,6 +2044,7 @@ mod tests { &environment, replicator_image, DestinationType::BigQuery, + None, LogLevel::Info, ); @@ -1867,6 +2077,7 @@ mod tests { &environment, replicator_image, DestinationType::BigQuery, + None, LogLevel::Info, ); @@ -1906,6 +2117,7 @@ mod tests { &environment, replicator_image, DestinationType::Iceberg, + None, LogLevel::Info, ); @@ -1938,6 +2150,7 @@ mod tests { &environment, replicator_image, DestinationType::Iceberg, + None, LogLevel::Info, ); @@ -1970,6 +2183,7 @@ mod tests { &environment, replicator_image, DestinationType::Iceberg, + None, LogLevel::Info, ); @@ -2009,6 +2223,7 @@ mod tests { &environment, replicator_image, DestinationType::Ducklake, + None, LogLevel::Info, ); @@ -2041,6 +2256,7 @@ mod tests { &environment, replicator_image, DestinationType::Ducklake, + None, LogLevel::Info, ); @@ -2073,6 +2289,7 @@ mod tests { &environment, replicator_image, DestinationType::Ducklake, + None, LogLevel::Info, ); diff --git a/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_bq_replicator_stateful_set_json-2.snap b/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_bq_replicator_stateful_set_json-2.snap index 7acd9812c..e45130f16 100644 --- a/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_bq_replicator_stateful_set_json-2.snap +++ b/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_bq_replicator_stateful_set_json-2.snap @@ -172,6 +172,7 @@ expression: stateful_set_json "nodeSelector": { "etl.supabase.com/node-role": "workloads" }, + "serviceAccountName": "etl-replicator", "terminationGracePeriodSeconds": 300, "tolerations": [ { diff --git a/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_bq_replicator_stateful_set_json-3.snap b/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_bq_replicator_stateful_set_json-3.snap index 4e956d9ba..135f41312 100644 --- a/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_bq_replicator_stateful_set_json-3.snap +++ b/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_bq_replicator_stateful_set_json-3.snap @@ -168,6 +168,7 @@ expression: stateful_set_json "nodeSelector": { "etl.supabase.com/node-role": "workloads" }, + "serviceAccountName": "etl-replicator", "terminationGracePeriodSeconds": 300, "tolerations": [ { diff --git a/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_bq_replicator_stateful_set_json.snap b/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_bq_replicator_stateful_set_json.snap index 6eff9c574..bf3f322ae 100644 --- a/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_bq_replicator_stateful_set_json.snap +++ b/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_bq_replicator_stateful_set_json.snap @@ -99,6 +99,7 @@ expression: stateful_set_json ], "initContainers": [], "nodeSelector": {}, + "serviceAccountName": "etl-replicator", "terminationGracePeriodSeconds": 300, "tolerations": [ { diff --git a/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_ducklake_maintenance_json.snap b/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_ducklake_maintenance_json.snap new file mode 100644 index 000000000..6e118a961 --- /dev/null +++ b/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_ducklake_maintenance_json.snap @@ -0,0 +1,65 @@ +--- +source: crates/etl-api/src/k8s/http.rs +expression: "serde_json::to_string_pretty(&ducklake_maintenance_json).unwrap()" +--- +{ + "apiVersion": "etl.supabase.com/v1alpha1", + "kind": "DuckLakeMaintenance", + "metadata": { + "labels": { + "etl.supabase.com/app-name": "abcdefghijklmnopqrst-42-replicator-app", + "etl.supabase.com/app-type": "etl-ducklake-maintenance-app" + }, + "name": "abcdefghijklmnopqrst-42", + "namespace": "etl-data-plane" + }, + "spec": { + "jobTemplate": { + "activeDeadlineSeconds": 1800, + "backoffLimit": 1, + "cpuRequestMillicores": 1000, + "image": "supabase/replicator:1.2.3", + "memoryRequestMiB": 1024, + "ttlSecondsAfterFinished": 86400 + }, + "operations": { + "cleanupOldFiles": { + "enabled": true + }, + "expireSnapshots": { + "enabled": false + }, + "inlineFlush": { + "enabled": true, + "minInlinedBytes": 10000000 + }, + "mergeAdjacentFiles": { + "enabled": true, + "maxCompactedFiles": 32, + "maxTablesPerRun": 8, + "targetFileSize": "10MB" + }, + "rewriteDataFiles": { + "deleteThreshold": 0.5, + "enabled": true, + "maxTablesPerRun": 8 + } + }, + "pause": { + "maxDurationSeconds": 2700 + }, + "pipelineRef": { + "pipelineId": 24, + "replicatorId": 42, + "tenantId": "abcdefghijklmnopqrst" + }, + "runtimeRefs": { + "configMapName": "abcdefghijklmnopqrst-42-replicator-config", + "ducklakeSecretName": "abcdefghijklmnopqrst-42-ducklake", + "postgresSecretName": "abcdefghijklmnopqrst-42-postgres-password" + }, + "schedule": { + "minIntervalSeconds": 3600 + } + } +} diff --git a/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_ducklake_replicator_stateful_set_json-2.snap b/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_ducklake_replicator_stateful_set_json-2.snap index 9623786b1..12c0646d2 100644 --- a/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_ducklake_replicator_stateful_set_json-2.snap +++ b/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_ducklake_replicator_stateful_set_json-2.snap @@ -183,6 +183,7 @@ expression: stateful_set_json "nodeSelector": { "etl.supabase.com/node-role": "workloads" }, + "serviceAccountName": "etl-replicator", "terminationGracePeriodSeconds": 300, "tolerations": [ { diff --git a/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_ducklake_replicator_stateful_set_json-3.snap b/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_ducklake_replicator_stateful_set_json-3.snap index 26b5a33f6..344044438 100644 --- a/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_ducklake_replicator_stateful_set_json-3.snap +++ b/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_ducklake_replicator_stateful_set_json-3.snap @@ -179,6 +179,7 @@ expression: stateful_set_json "nodeSelector": { "etl.supabase.com/node-role": "workloads" }, + "serviceAccountName": "etl-replicator", "terminationGracePeriodSeconds": 300, "tolerations": [ { diff --git a/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_ducklake_replicator_stateful_set_json.snap b/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_ducklake_replicator_stateful_set_json.snap index fcec8a43d..348acfd47 100644 --- a/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_ducklake_replicator_stateful_set_json.snap +++ b/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_ducklake_replicator_stateful_set_json.snap @@ -111,6 +111,7 @@ expression: stateful_set_json ], "initContainers": [], "nodeSelector": {}, + "serviceAccountName": "etl-replicator", "terminationGracePeriodSeconds": 300, "tolerations": [ { diff --git a/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_iceberg_replicator_stateful_set_json-2.snap b/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_iceberg_replicator_stateful_set_json-2.snap index 20d59e85a..071f44bbb 100644 --- a/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_iceberg_replicator_stateful_set_json-2.snap +++ b/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_iceberg_replicator_stateful_set_json-2.snap @@ -190,6 +190,7 @@ expression: stateful_set_json "nodeSelector": { "etl.supabase.com/node-role": "workloads" }, + "serviceAccountName": "etl-replicator", "terminationGracePeriodSeconds": 300, "tolerations": [ { diff --git a/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_iceberg_replicator_stateful_set_json-3.snap b/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_iceberg_replicator_stateful_set_json-3.snap index b5dddb312..1154f63f2 100644 --- a/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_iceberg_replicator_stateful_set_json-3.snap +++ b/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_iceberg_replicator_stateful_set_json-3.snap @@ -186,6 +186,7 @@ expression: stateful_set_json "nodeSelector": { "etl.supabase.com/node-role": "workloads" }, + "serviceAccountName": "etl-replicator", "terminationGracePeriodSeconds": 300, "tolerations": [ { diff --git a/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_iceberg_replicator_stateful_set_json.snap b/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_iceberg_replicator_stateful_set_json.snap index 162080211..d4f3c5643 100644 --- a/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_iceberg_replicator_stateful_set_json.snap +++ b/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_iceberg_replicator_stateful_set_json.snap @@ -117,6 +117,7 @@ expression: stateful_set_json ], "initContainers": [], "nodeSelector": {}, + "serviceAccountName": "etl-replicator", "terminationGracePeriodSeconds": 300, "tolerations": [ { diff --git a/crates/etl-api/src/routes/pipelines.rs b/crates/etl-api/src/routes/pipelines.rs index bea671fb8..3c824d7c9 100644 --- a/crates/etl-api/src/routes/pipelines.rs +++ b/crates/etl-api/src/routes/pipelines.rs @@ -34,7 +34,8 @@ use crate::{ }, feature_flags::{FeatureFlagsClient, get_max_pipelines_per_tenant}, k8s::{ - K8sClient, K8sError, PodStatus, TrustedRootCertsCache, TrustedRootCertsError, + DestinationType, K8sClient, K8sError, PodStatus, TrustedRootCertsCache, + TrustedRootCertsError, core::{ create_k8s_object_prefix, create_or_update_pipeline_resources_in_k8s, delete_pipeline_resources_in_k8s, is_replicator_active, is_replicator_pod_stopped, @@ -1361,9 +1362,14 @@ pub(crate) async fn update_pipeline_version( .ok_or(PipelineError::ReplicatorNotFound(pipeline_id))?; } - // If the images have equal name, we don't care about their id from the K8S - // perspective, so we won't update any resources. - if target_image.name == current_image.name { + let destination_type = DestinationType::from(&destination.config); + let image_name_unchanged = target_image.name == current_image.name; + + // If the images have equal name, non-DuckLake pipelines do not need any + // K8s reconciliation. DuckLake pipelines still reconcile because the + // external maintenance CR may be missing for pipelines created before that + // resource existed. + if image_name_unchanged && !matches!(destination_type, DestinationType::Ducklake) { txn.commit().await?; return Ok(HttpResponse::Ok().finish()); diff --git a/crates/etl-api/tests/pipelines.rs b/crates/etl-api/tests/pipelines.rs index f154c2993..45ad1e4ac 100644 --- a/crates/etl-api/tests/pipelines.rs +++ b/crates/etl-api/tests/pipelines.rs @@ -23,7 +23,9 @@ use crate::support::{ k8s_client::MockK8sState, mocks::{ create_default_image, create_image_with_name, - destinations::create_destination, + destinations::{ + create_destination, create_destination_with_config, new_ducklake_destination_config, + }, pipelines::{create_pipeline_with_config, new_pipeline_config, updated_pipeline_config}, sources::create_source, tenants::{create_tenant, create_tenant_with_id_and_name}, @@ -854,6 +856,43 @@ async fn pipeline_version_can_be_updated() { assert!(response.status().is_success()); } +#[tokio::test(flavor = "multi_thread")] +async fn pipeline_version_update_reconciles_ducklake_maintenance_when_image_name_is_unchanged() { + init_test_tracing(); + // Arrange + let k8s_state = MockK8sState::default(); + let app = spawn_test_app_with_k8s_state(None, k8s_state.clone()).await; + let default_image_id = + create_image_with_name(&app, "supabase/replicator:1.3.0".to_owned(), true).await; + let tenant_id = create_tenant(&app).await; + let source_id = create_source(&app, &tenant_id).await; + let destination_id = create_destination_with_config( + &app, + &tenant_id, + "DuckLake Destination".to_owned(), + new_ducklake_destination_config(), + ) + .await; + + let pipeline_id = create_pipeline_with_config( + &app, + &tenant_id, + source_id, + destination_id, + new_pipeline_config(), + ) + .await; + let maintenance_calls_before_update = k8s_state.ducklake_maintenance_create_calls(); + + // Act + let update_request = UpdatePipelineVersionRequest { version_id: default_image_id }; + let response = app.update_pipeline_version(&tenant_id, pipeline_id, &update_request).await; + + // Assert + assert!(response.status().is_success()); + assert_eq!(k8s_state.ducklake_maintenance_create_calls(), maintenance_calls_before_update + 1); +} + #[tokio::test(flavor = "multi_thread")] async fn pipeline_version_update_skips_k8s_reconcile_when_pipeline_is_stopped() { init_test_tracing(); diff --git a/crates/etl-api/tests/snapshots/main__destinations_pipelines__an_existing_bigquery_destination_and_pipeline_can_be_updated-2.snap b/crates/etl-api/tests/snapshots/main__destinations_pipelines__an_existing_bigquery_destination_and_pipeline_can_be_updated-2.snap index fcc178aba..8c53fb28e 100644 --- a/crates/etl-api/tests/snapshots/main__destinations_pipelines__an_existing_bigquery_destination_and_pipeline_can_be_updated-2.snap +++ b/crates/etl-api/tests/snapshots/main__destinations_pipelines__an_existing_bigquery_destination_and_pipeline_can_be_updated-2.snap @@ -39,6 +39,7 @@ FullApiPipelineConfig { Error, ), replicator_resources: None, + ducklake_maintenance: None, log_level: Some( Info, ), diff --git a/crates/etl-api/tests/snapshots/main__destinations_pipelines__an_existing_iceberg_supabase_destination_and_pipeline_can_be_updated-2.snap b/crates/etl-api/tests/snapshots/main__destinations_pipelines__an_existing_iceberg_supabase_destination_and_pipeline_can_be_updated-2.snap index fcc178aba..8c53fb28e 100644 --- a/crates/etl-api/tests/snapshots/main__destinations_pipelines__an_existing_iceberg_supabase_destination_and_pipeline_can_be_updated-2.snap +++ b/crates/etl-api/tests/snapshots/main__destinations_pipelines__an_existing_iceberg_supabase_destination_and_pipeline_can_be_updated-2.snap @@ -39,6 +39,7 @@ FullApiPipelineConfig { Error, ), replicator_resources: None, + ducklake_maintenance: None, log_level: Some( Info, ), diff --git a/crates/etl-api/tests/snapshots/main__destinations_pipelines__bigquery_destination_and_pipeline_can_be_created-2.snap b/crates/etl-api/tests/snapshots/main__destinations_pipelines__bigquery_destination_and_pipeline_can_be_created-2.snap index 302ad43da..c197004a3 100644 --- a/crates/etl-api/tests/snapshots/main__destinations_pipelines__bigquery_destination_and_pipeline_can_be_created-2.snap +++ b/crates/etl-api/tests/snapshots/main__destinations_pipelines__bigquery_destination_and_pipeline_can_be_created-2.snap @@ -39,6 +39,7 @@ FullApiPipelineConfig { Error, ), replicator_resources: None, + ducklake_maintenance: None, log_level: Some( Info, ), diff --git a/crates/etl-api/tests/snapshots/main__destinations_pipelines__iceberg_supabase_destination_and_pipeline_can_be_created-2.snap b/crates/etl-api/tests/snapshots/main__destinations_pipelines__iceberg_supabase_destination_and_pipeline_can_be_created-2.snap index 302ad43da..c197004a3 100644 --- a/crates/etl-api/tests/snapshots/main__destinations_pipelines__iceberg_supabase_destination_and_pipeline_can_be_created-2.snap +++ b/crates/etl-api/tests/snapshots/main__destinations_pipelines__iceberg_supabase_destination_and_pipeline_can_be_created-2.snap @@ -39,6 +39,7 @@ FullApiPipelineConfig { Error, ), replicator_resources: None, + ducklake_maintenance: None, log_level: Some( Info, ), diff --git a/crates/etl-api/tests/snapshots/main__pipelines__all_pipelines_can_be_read.snap b/crates/etl-api/tests/snapshots/main__pipelines__all_pipelines_can_be_read.snap index e5e7782aa..6a7e1feec 100644 --- a/crates/etl-api/tests/snapshots/main__pipelines__all_pipelines_can_be_read.snap +++ b/crates/etl-api/tests/snapshots/main__pipelines__all_pipelines_can_be_read.snap @@ -39,6 +39,7 @@ FullApiPipelineConfig { Error, ), replicator_resources: None, + ducklake_maintenance: None, log_level: Some( Info, ), diff --git a/crates/etl-api/tests/snapshots/main__pipelines__an_existing_pipeline_can_be_read.snap b/crates/etl-api/tests/snapshots/main__pipelines__an_existing_pipeline_can_be_read.snap index 6aa979726..26a15d774 100644 --- a/crates/etl-api/tests/snapshots/main__pipelines__an_existing_pipeline_can_be_read.snap +++ b/crates/etl-api/tests/snapshots/main__pipelines__an_existing_pipeline_can_be_read.snap @@ -39,6 +39,7 @@ FullApiPipelineConfig { Error, ), replicator_resources: None, + ducklake_maintenance: None, log_level: Some( Info, ), diff --git a/crates/etl-api/tests/snapshots/main__pipelines__an_existing_pipeline_can_be_updated.snap b/crates/etl-api/tests/snapshots/main__pipelines__an_existing_pipeline_can_be_updated.snap index d2a0d15a7..e658b83dd 100644 --- a/crates/etl-api/tests/snapshots/main__pipelines__an_existing_pipeline_can_be_updated.snap +++ b/crates/etl-api/tests/snapshots/main__pipelines__an_existing_pipeline_can_be_updated.snap @@ -39,6 +39,7 @@ FullApiPipelineConfig { Error, ), replicator_resources: None, + ducklake_maintenance: None, log_level: Some( Info, ), diff --git a/crates/etl-api/tests/support/k8s_client.rs b/crates/etl-api/tests/support/k8s_client.rs index a11d9488c..c783fbb68 100644 --- a/crates/etl-api/tests/support/k8s_client.rs +++ b/crates/etl-api/tests/support/k8s_client.rs @@ -10,13 +10,13 @@ use std::{ use async_trait::async_trait; use etl_api::{ - configs::{log::LogLevel, pipeline::ReplicatorResourcesConfig}, + configs::pipeline::ReplicatorResourcesConfig, k8s::{ - DestinationType, K8sClient, K8sError, PodStatus, ReplicatorConfigMapFile, + DuckLakeMaintenanceResourceConfig, K8sClient, K8sError, PodStatus, ReplicatorConfigMapFile, + ReplicatorStatefulSetConfig, http::{TRUSTED_ROOT_CERT_CONFIG_MAP_NAME, TRUSTED_ROOT_CERT_KEY_NAME}, }, }; -use etl_config::Environment; use k8s_openapi::api::core::v1::ConfigMap; use tokio::sync::RwLock; @@ -24,6 +24,7 @@ use tokio::sync::RwLock; pub(crate) struct MockK8sState { pod_status: Arc>, create_calls: Arc, + ducklake_maintenance_create_calls: Arc, last_replicator_resources: Arc>>, } @@ -32,6 +33,7 @@ impl Default for MockK8sState { Self { pod_status: Arc::new(RwLock::new(PodStatus::Started)), create_calls: Arc::new(AtomicUsize::new(0)), + ducklake_maintenance_create_calls: Arc::new(AtomicUsize::new(0)), last_replicator_resources: Arc::new(RwLock::new(None)), } } @@ -46,6 +48,10 @@ impl MockK8sState { self.create_calls.load(Ordering::Relaxed) } + pub(crate) fn ducklake_maintenance_create_calls(&self) -> usize { + self.ducklake_maintenance_create_calls.load(Ordering::Relaxed) + } + pub(crate) async fn last_replicator_resources(&self) -> Option { self.last_replicator_resources.read().await.clone() } @@ -173,14 +179,9 @@ impl K8sClient for MockK8sClient { async fn create_or_update_replicator_stateful_set( &self, - _prefix: &str, - _replicator_image: &str, - _environment: Environment, - replicator_resources: Option<&ReplicatorResourcesConfig>, - _destination_type: DestinationType, - _log_level: LogLevel, + config: ReplicatorStatefulSetConfig, ) -> Result<(), K8sError> { - self.set_last_replicator_resources(replicator_resources).await; + self.set_last_replicator_resources(config.replicator_resources.as_ref()).await; self.record_create_call(); Ok(()) } @@ -189,6 +190,19 @@ impl K8sClient for MockK8sClient { Ok(()) } + async fn create_or_update_ducklake_maintenance( + &self, + _prefix: &str, + _config: DuckLakeMaintenanceResourceConfig, + ) -> Result<(), K8sError> { + self.state.ducklake_maintenance_create_calls.fetch_add(1, Ordering::Relaxed); + Ok(()) + } + + async fn delete_ducklake_maintenance(&self, _prefix: &str) -> Result<(), K8sError> { + Ok(()) + } + async fn get_replicator_pod_status(&self, _prefix: &str) -> Result { Ok(*self.state.pod_status.read().await) } diff --git a/crates/etl-api/tests/support/mocks.rs b/crates/etl-api/tests/support/mocks.rs index 70bab1040..ea8b35468 100644 --- a/crates/etl-api/tests/support/mocks.rs +++ b/crates/etl-api/tests/support/mocks.rs @@ -69,6 +69,27 @@ pub(crate) mod destinations { } } + /// Returns a default DuckLake destination config. + pub(crate) fn new_ducklake_destination_config() -> FullApiDestinationConfig { + FullApiDestinationConfig::Ducklake { + catalog_url: "postgres://postgres:postgres@localhost:5432/postgres".to_owned(), + data_path: "s3://ducklake/".to_owned(), + pool_size: Some(1), + s3_access_key_id: Some(SerializableSecretString::from("access-key-id".to_owned())), + s3_secret_access_key: Some(SerializableSecretString::from( + "secret-access-key".to_owned(), + )), + s3_region: Some("us-east-1".to_owned()), + s3_endpoint: Some("localhost:9000".to_owned()), + s3_url_style: Some("path".to_owned()), + s3_use_ssl: Some(false), + metadata_schema: Some("ducklake".to_owned()), + duckdb_memory_cache_limit: None, + maintenance_target_file_size: Some("10MB".to_owned()), + expire_snapshots_older_than: Some("7 days".to_owned()), + } + } + /// Returns a default Iceberg Supabase destination config. pub(crate) fn new_iceberg_supabase_destination_config() -> FullApiDestinationConfig { use etl_api::configs::destination::FullApiIcebergConfig; @@ -251,6 +272,7 @@ pub(crate) mod pipelines { table_sync_copy: Some(TableSyncCopyConfig::IncludeAllTables), invalidated_slot_behavior: None, replicator_resources: None, + ducklake_maintenance: None, log_level: Some(LogLevel::Info), } } @@ -273,6 +295,7 @@ pub(crate) mod pipelines { table_sync_copy: Some(TableSyncCopyConfig::IncludeAllTables), invalidated_slot_behavior: None, replicator_resources: None, + ducklake_maintenance: None, log_level: Some(LogLevel::Info), } } diff --git a/crates/etl-api/tests/validators.rs b/crates/etl-api/tests/validators.rs index 6e36905cb..8ee7a6309 100644 --- a/crates/etl-api/tests/validators.rs +++ b/crates/etl-api/tests/validators.rs @@ -82,6 +82,7 @@ fn create_pipeline_config(publication_name: &str) -> FullApiPipelineConfig { table_sync_copy: None, invalidated_slot_behavior: None, replicator_resources: None, + ducklake_maintenance: None, } } diff --git a/crates/etl-destinations/Cargo.toml b/crates/etl-destinations/Cargo.toml index acc9d1f37..526307752 100644 --- a/crates/etl-destinations/Cargo.toml +++ b/crates/etl-destinations/Cargo.toml @@ -15,12 +15,15 @@ doctest = false [features] ducklake = [ "dep:duckdb", + "dep:humantime", "dep:metrics", + "dep:kube", "dep:parking_lot", "dep:pg_escape", "dep:r2d2", "dep:rand", "dep:regex", + "dep:serde_json", "dep:sqlx", "dep:tokio-postgres", "dep:tracing", @@ -76,8 +79,10 @@ duckdb = { workspace = true, optional = true, features = ["bundled", "json", "pa etl = { workspace = true } futures = { workspace = true, optional = true } gcp-bigquery-client = { workspace = true, optional = true, features = ["rust-tls", "aws-lc-rs"] } +humantime = { workspace = true, optional = true } iceberg = { workspace = true, optional = true } iceberg-catalog-rest = { workspace = true, optional = true } +kube = { workspace = true, optional = true, features = ["client", "rustls-tls"] } metrics = { workspace = true, optional = true } parking_lot = { workspace = true, optional = true } parquet = { workspace = true, optional = true, features = ["async", "arrow"] } @@ -106,6 +111,7 @@ etl = { workspace = true, features = ["test-utils"] } etl-postgres = { workspace = true, features = ["test-utils", "tokio"] } etl-telemetry = { workspace = true } futures = { workspace = true } +k8s-openapi = { workspace = true, features = ["latest"] } rand = { workspace = true, features = ["thread_rng"] } rustls = { workspace = true, features = ["aws-lc-rs", "logging"] } serde = { workspace = true, features = ["derive"] } diff --git a/crates/etl-destinations/src/ducklake/METRICS.md b/crates/etl-destinations/src/ducklake/METRICS.md index 8fd41ea2b..436c4f58c 100644 --- a/crates/etl-destinations/src/ducklake/METRICS.md +++ b/crates/etl-destinations/src/ducklake/METRICS.md @@ -7,9 +7,9 @@ The metrics fall into four groups: - write-path metrics: show how the ETL writer is batching, waiting, retrying, and flushing inline data. -- maintenance execution metrics: operation-level duration and skip metrics - emitted by the background maintenance worker with the primary reason and - outcome for each maintenance attempt. +- external maintenance metrics: operation-trigger counts and duration samples + for time foreground ingestion was quiesced by the Kubernetes maintenance + plane. - table-health samples: histograms recorded by a background sampler every 30 seconds from the PostgreSQL DuckLake metadata catalog. They describe the current shape of tables known to the current destination instance. @@ -56,15 +56,14 @@ These explain the pressure your writer is putting on DuckLake: - `etl_ducklake_inline_flush_rows` - `etl_ducklake_inline_flush_duration_seconds` -The destination attaches DuckLake with `DATA_INLINING_ROW_LIMIT = 10000` and then -lets a background maintenance worker flush and checkpoint inlined data after -writes. Destination shutdown also runs one final best-effort inline flush sweep -for known tables. These metrics tell you whether that strategy is helping. +The destination attaches DuckLake with `DATA_INLINING_ROW_LIMIT = 10000`. +External maintenance jobs flush inlined data during coordinated pauses. These +metrics tell you whether that strategy is helping. -The `batch_kind` label on these metrics uses: +The `result` label on these metrics uses: -- `mutation` for background CDC flushes -- `shutdown` for the final best-effort shutdown sweep +- `flushed` when rows were materialized +- `noop` when no rows needed materialization How to read them: @@ -76,55 +75,31 @@ How to read them: - if `inline_flush_rows` is often meaningfully larger than `upsert_rows`, the inlining limit is helping consolidate multiple atomic batches before files are materialized. -- if `batch_kind="shutdown"` shows meaningful work, the process is still - relying on final shutdown cleanup to drain inline backlogs. +### External maintenance metrics -### Background maintenance metrics +- `etl_ducklake_external_maintenance_pause_duration_seconds` +- `etl_ducklake_external_maintenance_triggered_total` -- `etl_ducklake_maintenance_duration_seconds` -- `etl_ducklake_maintenance_skipped_total` +`etl_ducklake_external_maintenance_pause_duration_seconds` is emitted by the +replicator when a Kubernetes-driven external maintenance pause ends. It measures +only the time after the destination has drained foreground mutations and reported +`Quiesced`; time spent queued by the controller is intentionally excluded. -`etl_ducklake_maintenance_duration_seconds` is emitted once per background -maintenance operation that actually runs. Use the histogram `_count` as the -event count for non-skipped outcomes. +It carries one label: + +- `outcome`: `cleared`, `expired`, `replaced`, or `resource_deleted` + +`etl_ducklake_external_maintenance_triggered_total` counts external maintenance +operation requests emitted by the replicator after it samples DuckLake catalog +state. The watcher reuses an existing pending request when it already covers the +sampled operations, so this counter is not incremented on every poll while a CR +waits in the controller queue. It carries these labels: -- `task`: `flush`, `scheduled_maintenance`, `targeted_maintenance`, or `checkpoint` -- `operation`: `flush_inlined_data`, `rewrite_data_files`, - `merge_adjacent_files`, or `checkpoint` -- `reason`: the primary cause for the maintenance decision -- `outcome`: `applied`, `noop`, or `failed` - -`etl_ducklake_maintenance_skipped_total` counts maintenance operations that -were deferred because their guard or execution window was unavailable. It is -labeled by `task`, `operation`, and `reason`. - -How to read it: - -- `task="flush"` with - `reason="pending_inlined_data_bytes_threshold"` means the flush was - scheduled from sampled inlined catalog-table size in a PostgreSQL-backed - DuckLake catalog. This includes both inlined inserts and inlined deletions. -- `task="scheduled_maintenance"` with `reason="merge_interval"` means the - next `write_events` batch tried to run the tier-0 `< 1MiB -> ~5MiB` - `merge_adjacent_files` pass first. -- `task="targeted_maintenance"` with - `reason="idle_rewrite_metrics_threshold"` or - `reason="emergency_rewrite_metrics_threshold"` means rewrite was selected - from sampled delete pressure. -- `task="targeted_maintenance"` with - `reason="idle_merge_metrics_threshold"` or - `reason="emergency_merge_metrics_threshold"` means merge was selected from - sampled small-file pressure. -- `task="targeted_maintenance"` may emit one duration series or two for a - maintenance cycle, depending on which operations crossed their thresholds. -- rising `etl_ducklake_maintenance_skipped_total` means maintenance keeps - missing its execution window because writes or another guarded operation are - still active. -- rising histogram `_count` for `outcome="failed"` points to maintenance - execution issues, while many `outcome="noop"` samples usually mean - maintenance is polling more often than work is actually accumulating. +- `operation`: `flush_inlined_data` or `rewrite_data_files` +- `reason`: `pending_inlined_data_bytes_threshold` or + `active_data_files_threshold` ### Table-health sampling metrics diff --git a/crates/etl-destinations/src/ducklake/client.rs b/crates/etl-destinations/src/ducklake/client.rs index d34a872e3..7e52c5f5f 100644 --- a/crates/etl-destinations/src/ducklake/client.rs +++ b/crates/etl-destinations/src/ducklake/client.rs @@ -2,7 +2,7 @@ use std::sync::atomic::AtomicUsize; use std::{ borrow::Cow, - error, fmt, + error, fmt, process, sync::{ Arc, LazyLock, Mutex, Weak, atomic::{AtomicBool, AtomicU64, Ordering}, @@ -34,6 +34,8 @@ use crate::ducklake::{ /// Monotonic identifier assigned to each DuckLake connection setup attempt. static NEXT_CONNECTION_INIT_ID: AtomicU64 = AtomicU64::new(1); +/// Monotonic identifier assigned to each timed DuckDB blocking operation. +static NEXT_DUCKDB_BLOCKING_OPERATION_ID: AtomicU64 = AtomicU64::new(1); /// Matches one libpq password field inside a conninfo string. static POSTGRES_PASSWORD_REGEX: LazyLock = LazyLock::new(|| { Regex::new(r"password=(?:'([^'\\]|\\.)*'|[^\s,);]+)") @@ -41,9 +43,29 @@ static POSTGRES_PASSWORD_REGEX: LazyLock = LazyLock::new(|| { }); /// Timeout applied to each foreground DuckLake blocking operation. -pub(super) const FOREGROUND_QUERY_TIMEOUT: Duration = Duration::from_secs(2 * 60); +pub(super) const FOREGROUND_QUERY_TIMEOUT: Duration = Duration::from_secs(3 * 60); /// Timeout applied to each maintenance DuckLake blocking operation. -pub(super) const MAINTENANCE_QUERY_TIMEOUT: Duration = Duration::from_secs(5 * 60); +pub(super) const MAINTENANCE_QUERY_TIMEOUT: Duration = Duration::from_secs(3 * 60); +/// Extra time allowed for a timed-out DuckDB operation to return after +/// interrupt() has been called. If the operation is still stuck after this, +/// the process is no longer safe to keep running. +const BLOCKING_ABORT_GRACE: Duration = Duration::from_secs(30); + +trait DuckDbQueryInterrupt: Send + Sync + 'static { + fn interrupt(&self); +} + +impl DuckDbQueryInterrupt for duckdb::InterruptHandle { + fn interrupt(&self) { + duckdb::InterruptHandle::interrupt(self); + } +} + +type DuckDbQueryInterruptHandle = Arc; + +fn remaining_ms_until(deadline: Instant) -> u64 { + deadline.checked_duration_since(Instant::now()).unwrap_or(Duration::ZERO).as_millis() as u64 +} /// Timeout class applied to one DuckDB blocking operation. #[derive(Clone, Copy, Debug, Eq, PartialEq)] @@ -71,58 +93,207 @@ impl DuckDbBlockingOperationKind { /// Async watchdog that interrupts one timed DuckDB query when its deadline /// expires. pub(super) struct DuckDbQueryWatchdog { + operation_id: u64, + operation_kind: DuckDbBlockingOperationKind, + timeout: Duration, timed_out: Arc, - interrupt_tx: Option>>, + interrupt_tx: Option>, done_tx: Option>, task: Option>, } impl DuckDbQueryWatchdog { + #[cfg(test)] fn spawn(deadline: Instant) -> Self { + Self::spawn_with_context( + deadline, + 0, + DuckDbBlockingOperationKind::Foreground, + Duration::ZERO, + ) + } + + fn spawn_with_context( + deadline: Instant, + operation_id: u64, + operation_kind: DuckDbBlockingOperationKind, + timeout: Duration, + ) -> Self { let timed_out = Arc::new(AtomicBool::new(false)); let timeout_flag = Arc::clone(&timed_out); - let (interrupt_tx, interrupt_rx) = oneshot::channel::>(); + let (interrupt_tx, interrupt_rx) = oneshot::channel::(); let (done_tx, done_rx) = oneshot::channel(); let task = tokio::spawn(async move { + info!( + operation_id, + operation_kind = operation_kind.as_str(), + timeout_ms = timeout.as_millis() as u64, + deadline_remaining_ms = remaining_ms_until(deadline), + "ducklake query watchdog task started: operation_id={}, operation_kind={}, \ + timeout_ms={}, deadline_remaining_ms={}", + operation_id, + operation_kind.as_str(), + timeout.as_millis(), + remaining_ms_until(deadline) + ); let mut interrupt_rx = Box::pin(interrupt_rx); let mut done_rx = Box::pin(done_rx); let interrupt_handle = tokio::select! { biased; - _ = &mut done_rx => return, + _ = &mut done_rx => { + info!( + operation_id, + operation_kind = operation_kind.as_str(), + "ducklake query watchdog finished before interrupt handle: operation_id={}, operation_kind={}", + operation_id, + operation_kind.as_str() + ); + return; + }, result = &mut interrupt_rx => match result { - Ok(handle) => handle, - Err(_) => return, + Ok(handle) => { + info!( + operation_id, + operation_kind = operation_kind.as_str(), + deadline_remaining_ms = remaining_ms_until(deadline), + "ducklake query watchdog received interrupt handle before deadline: \ + operation_id={}, operation_kind={}, deadline_remaining_ms={}", + operation_id, + operation_kind.as_str(), + remaining_ms_until(deadline) + ); + handle + }, + Err(_) => { + warn!( + operation_id, + operation_kind = operation_kind.as_str(), + "ducklake query watchdog interrupt sender dropped before deadline: \ + operation_id={}, operation_kind={}", + operation_id, + operation_kind.as_str() + ); + return; + }, }, _ = tokio::time::sleep_until(deadline) => { timeout_flag.store(true, Ordering::Relaxed); + warn!( + operation_id, + operation_kind = operation_kind.as_str(), + timeout_ms = timeout.as_millis() as u64, + "ducklake query watchdog deadline elapsed before interrupt handle: \ + operation_id={}, operation_kind={}, timeout_ms={}", + operation_id, + operation_kind.as_str(), + timeout.as_millis() + ); // If we didn't receive the interrupt_rx yet, make sure to get it to call interrupt() later tokio::select! { biased; - _ = &mut done_rx => return, + _ = &mut done_rx => { + info!( + operation_id, + operation_kind = operation_kind.as_str(), + "ducklake query watchdog received done after deadline before interrupt handle: \ + operation_id={}, operation_kind={}", + operation_id, + operation_kind.as_str() + ); + return; + }, result = &mut interrupt_rx => match result { - Ok(handle) => handle, - Err(_) => return, + Ok(handle) => { + warn!( + operation_id, + operation_kind = operation_kind.as_str(), + "ducklake query watchdog received interrupt handle after deadline: \ + operation_id={}, operation_kind={}", + operation_id, + operation_kind.as_str() + ); + handle + }, + Err(_) => { + warn!( + operation_id, + operation_kind = operation_kind.as_str(), + "ducklake query watchdog interrupt sender dropped after deadline: \ + operation_id={}, operation_kind={}", + operation_id, + operation_kind.as_str() + ); + return; + }, }, } }, }; if timeout_flag.load(Ordering::Relaxed) { + warn!( + operation_id, + operation_kind = operation_kind.as_str(), + "ducklake query watchdog calling interrupt after timeout: operation_id={}, \ + operation_kind={}", + operation_id, + operation_kind.as_str() + ); interrupt_handle.interrupt(); + warn!( + operation_id, + operation_kind = operation_kind.as_str(), + "ducklake query watchdog interrupt returned after timeout: operation_id={}, \ + operation_kind={}", + operation_id, + operation_kind.as_str() + ); return; } tokio::select! { biased; - _ = &mut done_rx => {} + _ = &mut done_rx => { + info!( + operation_id, + operation_kind = operation_kind.as_str(), + deadline_remaining_ms = remaining_ms_until(deadline), + "ducklake query watchdog received done before deadline after interrupt handle: \ + operation_id={}, operation_kind={}, deadline_remaining_ms={}", + operation_id, + operation_kind.as_str(), + remaining_ms_until(deadline) + ); + } _ = tokio::time::sleep_until(deadline) => { timeout_flag.store(true, Ordering::Relaxed); + warn!( + operation_id, + operation_kind = operation_kind.as_str(), + timeout_ms = timeout.as_millis() as u64, + "ducklake query watchdog deadline elapsed after interrupt handle; calling interrupt: \ + operation_id={}, operation_kind={}, timeout_ms={}", + operation_id, + operation_kind.as_str(), + timeout.as_millis() + ); interrupt_handle.interrupt(); + warn!( + operation_id, + operation_kind = operation_kind.as_str(), + "ducklake query watchdog interrupt returned after handle/deadline path: \ + operation_id={}, operation_kind={}", + operation_id, + operation_kind.as_str() + ); } } }); Self { + operation_id, + operation_kind, + timeout, timed_out, interrupt_tx: Some(interrupt_tx), done_tx: Some(done_tx), @@ -131,14 +302,58 @@ impl DuckDbQueryWatchdog { } fn publish_interrupt_handle(&mut self, handle: Arc) { + self.publish_query_interrupt_handle(handle); + } + + fn publish_query_interrupt_handle(&mut self, handle: DuckDbQueryInterruptHandle) { if let Some(interrupt_tx) = self.interrupt_tx.take() { + info!( + operation_id = self.operation_id, + operation_kind = self.operation_kind.as_str(), + timeout_ms = self.timeout.as_millis() as u64, + "ducklake query watchdog publishing interrupt handle: operation_id={}, \ + operation_kind={}, timeout_ms={}", + self.operation_id, + self.operation_kind.as_str(), + self.timeout.as_millis() + ); let _ = interrupt_tx.send(handle); + } else { + warn!( + operation_id = self.operation_id, + operation_kind = self.operation_kind.as_str(), + "ducklake query watchdog interrupt handle publish skipped because sender is gone: \ + operation_id={}, operation_kind={}", + self.operation_id, + self.operation_kind.as_str() + ); } } fn finish(&mut self) { if let Some(done_tx) = self.done_tx.take() { + info!( + operation_id = self.operation_id, + operation_kind = self.operation_kind.as_str(), + timed_out = self.timed_out(), + "ducklake query watchdog finish signal sent: operation_id={}, operation_kind={}, \ + timed_out={}", + self.operation_id, + self.operation_kind.as_str(), + self.timed_out() + ); let _ = done_tx.send(()); + } else { + warn!( + operation_id = self.operation_id, + operation_kind = self.operation_kind.as_str(), + timed_out = self.timed_out(), + "ducklake query watchdog finish skipped because sender is gone: operation_id={}, \ + operation_kind={}, timed_out={}", + self.operation_id, + self.operation_kind.as_str(), + self.timed_out() + ); } } @@ -253,108 +468,7 @@ impl fmt::Display for DuckLakeConnectionError { impl error::Error for DuckLakeConnectionError {} -/// Manages one dedicated DuckLake pool for a background task. -pub(super) struct LazyDuckLakePool { - manager: DuckLakeConnectionManager, - pool_size: u32, - purpose: &'static str, - pool: Option>>, - init_task: Option>>>>, - blocking_slots: Arc, -} - -impl LazyDuckLakePool { - /// Creates one pool wrapper for a dedicated background task. - pub(super) fn new( - manager: DuckLakeConnectionManager, - pool_size: u32, - purpose: &'static str, - ) -> Self { - Self { - manager, - pool_size, - purpose, - pool: None, - init_task: None, - blocking_slots: Arc::new(Semaphore::new(pool_size as usize)), - } - } - - /// Starts warming the pool on a detached Tokio task. - pub(super) fn warm_in_background(&mut self) { - if self.pool.is_some() || self.init_task.is_some() { - return; - } - - let manager = self.manager.clone(); - let pool_size = self.pool_size; - let purpose = self.purpose; - self.init_task = Some(tokio::spawn(async move { - let pool = build_warm_ducklake_pool(manager, pool_size, purpose).await.map(Arc::new); - if let Err(error) = &pool { - warn!( - purpose, - error = %error, - "ducklake background pool warm-up failed" - ); - } - pool - })); - } - - /// Returns the warmed pool, awaiting any in-flight initialization or - /// initializing it on first use. - pub(super) async fn get_or_init_pool( - &mut self, - ) -> EtlResult>> { - if let Some(pool) = &self.pool { - return Ok(Arc::clone(pool)); - } - - if self.init_task.is_none() { - self.warm_in_background(); - } - - let pool = self - .init_task - .take() - .ok_or_else(|| { - etl_error!( - ErrorKind::DestinationError, - "DuckLake connection pool initialization task was not created" - ) - })? - .await - .map_err(|_| { - etl_error!( - ErrorKind::ApplyWorkerPanic, - "DuckLake connection pool initialization task panicked" - ) - })??; - self.pool = Some(Arc::clone(&pool)); - Ok(pool) - } - - /// Returns the semaphore that bounds background DuckDB concurrency. - pub(super) fn blocking_slots(&self) -> Arc { - Arc::clone(&self.blocking_slots) - } -} - -impl Drop for LazyDuckLakePool { - fn drop(&mut self) { - if let Some(init_task) = &self.init_task { - init_task.abort(); - } - } -} - impl DuckLakeConnectionManager { - /// Returns the shared interrupt registry for this manager and its clones. - pub(super) fn interrupt_registry(&self) -> Arc { - Arc::clone(&self.interrupt_registry) - } - /// Interrupts all currently live managed DuckLake connections. pub(super) fn interrupt_all_connections(&self) -> usize { self.interrupt_registry.interrupt_all() @@ -453,7 +567,7 @@ pub(super) async fn build_warm_ducklake_pool( let started = Instant::now(); let pool = r2d2::Pool::builder() .max_size(pool_size) - .min_idle(Some(pool_size)) + .min_idle(Some(0)) .connection_timeout(Duration::from_mins(4)) .test_on_check_out(true) // Callers log the returned pool initialization failure once, so @@ -517,6 +631,27 @@ pub(super) fn duckdb_blocking_timeout_error( ) } +fn abort_stuck_duckdb_blocking_operation( + operation_id: u64, + operation_kind: DuckDbBlockingOperationKind, + timeout: Duration, + abort_grace: Duration, +) -> ! { + tracing::error!( + operation_id, + operation_kind = operation_kind.as_str(), + timeout_ms = timeout.as_millis() as u64, + abort_grace_ms = abort_grace.as_millis() as u64, + "ducklake blocking operation did not return after timeout interrupt; aborting process: \ + operation_id={}, operation_kind={}, timeout_ms={}, abort_grace_ms={}", + operation_id, + operation_kind.as_str(), + timeout.as_millis(), + abort_grace.as_millis() + ); + process::abort(); +} + /// Runs one DuckDB operation on Tokio's blocking pool after acquiring a permit /// that matches the configured DuckDB concurrency limit and then checking out a /// warm pooled DuckDB connection. @@ -552,16 +687,85 @@ where R: Send + 'static, F: FnOnce(&duckdb::Connection) -> EtlResult + Send + 'static, { + let operation_id = NEXT_DUCKDB_BLOCKING_OPERATION_ID.fetch_add(1, Ordering::Relaxed); let deadline = Instant::now() + timeout; + info!( + operation_id, + operation_kind = operation_kind.as_str(), + timeout_ms = timeout.as_millis() as u64, + abort_grace_ms = BLOCKING_ABORT_GRACE.as_millis() as u64, + available_permits = blocking_slots.available_permits(), + "ducklake blocking operation starting: operation_id={}, operation_kind={}, timeout_ms={}, \ + abort_grace_ms={}, available_permits={}", + operation_id, + operation_kind.as_str(), + timeout.as_millis(), + BLOCKING_ABORT_GRACE.as_millis(), + blocking_slots.available_permits() + ); let slot_wait_started = Instant::now(); - let permit = tokio::time::timeout_at(deadline, blocking_slots.acquire_owned()) - .await - .map_err(|_| duckdb_blocking_timeout_error(operation_kind, timeout, "slot_wait"))? - .map_err(|_| { - etl_error!(ErrorKind::ApplyWorkerPanic, "DuckLake blocking slot acquisition failed") - })?; + info!( + operation_id, + operation_kind = operation_kind.as_str(), + deadline_remaining_ms = remaining_ms_until(deadline), + available_permits = blocking_slots.available_permits(), + "ducklake blocking operation waiting for semaphore slot: operation_id={}, \ + operation_kind={}, deadline_remaining_ms={}, available_permits={}", + operation_id, + operation_kind.as_str(), + remaining_ms_until(deadline), + blocking_slots.available_permits() + ); + let permit = match tokio::time::timeout_at( + deadline, + Arc::clone(&blocking_slots).acquire_owned(), + ) + .await + { + Ok(Ok(permit)) => permit, + Ok(Err(_)) => { + tracing::error!( + operation_id, + operation_kind = operation_kind.as_str(), + "ducklake blocking operation semaphore closed: operation_id={}, operation_kind={}", + operation_id, + operation_kind.as_str() + ); + return Err(etl_error!( + ErrorKind::ApplyWorkerPanic, + "DuckLake blocking slot acquisition failed" + )); + } + Err(_) => { + warn!( + operation_id, + operation_kind = operation_kind.as_str(), + timeout_ms = timeout.as_millis() as u64, + slot_wait_ms = slot_wait_started.elapsed().as_millis() as u64, + "ducklake blocking operation timed out waiting for semaphore slot: \ + operation_id={}, operation_kind={}, timeout_ms={}, slot_wait_ms={}", + operation_id, + operation_kind.as_str(), + timeout.as_millis(), + slot_wait_started.elapsed().as_millis() + ); + return Err(duckdb_blocking_timeout_error(operation_kind, timeout, "slot_wait")); + } + }; histogram!(ETL_DUCKLAKE_BLOCKING_SLOT_WAIT_SECONDS) .record(slot_wait_started.elapsed().as_secs_f64()); + info!( + operation_id, + operation_kind = operation_kind.as_str(), + slot_wait_ms = slot_wait_started.elapsed().as_millis() as u64, + deadline_remaining_ms = remaining_ms_until(deadline), + "ducklake blocking operation acquired semaphore slot: operation_id={}, operation_kind={}, \ + slot_wait_ms={}, deadline_remaining_ms={}", + operation_id, + operation_kind.as_str(), + slot_wait_started.elapsed().as_millis(), + remaining_ms_until(deadline) + ); trace!( wait_ms = slot_wait_started.elapsed().as_millis() as u64, "wait for ducklake blocking slot" @@ -570,32 +774,107 @@ where // This is needed to make sure we properly interrupt the blocking operation if // it exceeds the timeout, we don't just cancel the task and leave the // connection active. - let mut watchdog = DuckDbQueryWatchdog::spawn(deadline); + let mut watchdog = + DuckDbQueryWatchdog::spawn_with_context(deadline, operation_id, operation_kind, timeout); let watchdog_task = watchdog.async_task_handle()?; - - let blocking_result = tokio::task::spawn_blocking(move || -> EtlResult { + let watchdog_timed_out = Arc::clone(&watchdog.timed_out); + let abort_deadline = deadline + BLOCKING_ABORT_GRACE; + + let blocking_task = tokio::task::spawn_blocking(move || -> EtlResult { + info!( + operation_id, + operation_kind = operation_kind.as_str(), + deadline_remaining_ms = remaining_ms_until(deadline), + "ducklake blocking operation entered spawn_blocking task: operation_id={}, \ + operation_kind={}, deadline_remaining_ms={}", + operation_id, + operation_kind.as_str(), + remaining_ms_until(deadline) + ); // Please if you modify the code inside this blocking task do not add any // blocking operations that could delay other tasks waiting on this slot. let _permit = permit; let checkout_timeout = deadline.checked_duration_since(Instant::now()).unwrap_or(Duration::ZERO); if checkout_timeout.is_zero() { + warn!( + operation_id, + operation_kind = operation_kind.as_str(), + "ducklake blocking operation deadline reached before pool checkout: \ + operation_id={}, operation_kind={}", + operation_id, + operation_kind.as_str() + ); return Err(duckdb_blocking_timeout_error(operation_kind, timeout, "pool_checkout")); } let checkout_started = Instant::now(); - let mut pooled_conn = pool.get_timeout(checkout_timeout).map_err(|e| { - if Instant::now() >= deadline { - duckdb_blocking_timeout_error(operation_kind, timeout, "pool_checkout") - } else { - etl_error!( + info!( + operation_id, + operation_kind = operation_kind.as_str(), + checkout_timeout_ms = checkout_timeout.as_millis() as u64, + "ducklake blocking operation checking out pooled connection: operation_id={}, \ + operation_kind={}, checkout_timeout_ms={}", + operation_id, + operation_kind.as_str(), + checkout_timeout.as_millis() + ); + let mut pooled_conn = match pool.get_timeout(checkout_timeout) { + Ok(pooled_conn) => pooled_conn, + Err(e) if Instant::now() >= deadline => { + warn!( + operation_id, + operation_kind = operation_kind.as_str(), + checkout_wait_ms = checkout_started.elapsed().as_millis() as u64, + timeout_ms = timeout.as_millis() as u64, + error = %e, + "ducklake blocking operation timed out checking out pooled connection: \ + operation_id={}, operation_kind={}, checkout_wait_ms={}, timeout_ms={}, error={}", + operation_id, + operation_kind.as_str(), + checkout_started.elapsed().as_millis(), + timeout.as_millis(), + e + ); + return Err(duckdb_blocking_timeout_error( + operation_kind, + timeout, + "pool_checkout", + )); + } + Err(e) => { + warn!( + operation_id, + operation_kind = operation_kind.as_str(), + checkout_wait_ms = checkout_started.elapsed().as_millis() as u64, + error = %e, + "ducklake blocking operation failed checking out pooled connection: operation_id={}, \ + operation_kind={}, checkout_wait_ms={}, error={}", + operation_id, + operation_kind.as_str(), + checkout_started.elapsed().as_millis(), + e + ); + return Err(etl_error!( ErrorKind::DestinationConnectionFailed, "Failed to check out DuckLake connection", source: e - ) + )); } - })?; + }; histogram!(ETL_DUCKLAKE_POOL_CHECKOUT_WAIT_SECONDS) .record(checkout_started.elapsed().as_secs_f64()); + info!( + operation_id, + operation_kind = operation_kind.as_str(), + checkout_wait_ms = checkout_started.elapsed().as_millis() as u64, + deadline_remaining_ms = remaining_ms_until(deadline), + "ducklake blocking operation checked out pooled connection: operation_id={}, \ + operation_kind={}, checkout_wait_ms={}, deadline_remaining_ms={}", + operation_id, + operation_kind.as_str(), + checkout_started.elapsed().as_millis(), + remaining_ms_until(deadline) + ); trace!( wait_ms = checkout_started.elapsed().as_millis() as u64, "wait for ducklake pool checkout" @@ -603,16 +882,67 @@ where let operation_timeout = deadline.checked_duration_since(Instant::now()).unwrap_or(Duration::ZERO); if operation_timeout.is_zero() { + warn!( + operation_id, + operation_kind = operation_kind.as_str(), + "ducklake blocking operation deadline reached before query execution: \ + operation_id={}, operation_kind={}", + operation_id, + operation_kind.as_str() + ); return Err(duckdb_blocking_timeout_error(operation_kind, timeout, "query_execution")); } let interrupt_handle = pooled_conn.conn.interrupt_handle(); + info!( + operation_id, + operation_kind = operation_kind.as_str(), + operation_timeout_ms = operation_timeout.as_millis() as u64, + "ducklake blocking operation publishing interrupt handle before query execution: \ + operation_id={}, operation_kind={}, operation_timeout_ms={}", + operation_id, + operation_kind.as_str(), + operation_timeout.as_millis() + ); watchdog.publish_interrupt_handle(interrupt_handle); if watchdog.timed_out() { + warn!( + operation_id, + operation_kind = operation_kind.as_str(), + "ducklake blocking operation timed out before query started; marking pooled \ + connection broken: operation_id={}, operation_kind={}", + operation_id, + operation_kind.as_str() + ); pooled_conn.broken = true; return Err(duckdb_blocking_timeout_error(operation_kind, timeout, "query_execution")); } let operation_started = Instant::now(); + info!( + operation_id, + operation_kind = operation_kind.as_str(), + deadline_remaining_ms = remaining_ms_until(deadline), + "ducklake blocking operation invoking DuckDB closure: operation_id={}, \ + operation_kind={}, deadline_remaining_ms={}", + operation_id, + operation_kind.as_str(), + remaining_ms_until(deadline) + ); let res = operation(&pooled_conn.conn); + let operation_duration_ms = operation_started.elapsed().as_millis() as u64; + info!( + operation_id, + operation_kind = operation_kind.as_str(), + duration_ms = operation_duration_ms, + timed_out = watchdog.timed_out(), + result_is_error = res.is_err(), + "ducklake blocking operation DuckDB closure returned: operation_id={}, \ + operation_kind={}, duration_ms={}, timed_out={}, result_is_error={}", + operation_id, + operation_kind.as_str(), + operation_duration_ms, + watchdog.timed_out(), + res.is_err() + ); watchdog.finish(); histogram!(ETL_DUCKLAKE_BLOCKING_OPERATION_DURATION_SECONDS) .record(operation_started.elapsed().as_secs_f64()); @@ -621,22 +951,155 @@ where "ducklake blocking operation finished" ); if watchdog.timed_out() { + warn!( + operation_id, + operation_kind = operation_kind.as_str(), + duration_ms = operation_duration_ms, + "ducklake blocking operation returned after timeout; marking pooled connection \ + broken: operation_id={}, operation_kind={}, duration_ms={}", + operation_id, + operation_kind.as_str(), + operation_duration_ms + ); pooled_conn.broken = true; return Err(duckdb_blocking_timeout_error(operation_kind, timeout, "query_execution")); } if res.is_err() { + warn!( + operation_id, + operation_kind = operation_kind.as_str(), + duration_ms = operation_duration_ms, + "ducklake blocking operation returned error; marking pooled connection broken: \ + operation_id={}, operation_kind={}, duration_ms={}", + operation_id, + operation_kind.as_str(), + operation_duration_ms + ); pooled_conn.broken = true; + } else { + info!( + operation_id, + operation_kind = operation_kind.as_str(), + duration_ms = operation_duration_ms, + "ducklake blocking operation returned success; pooled connection remains healthy: \ + operation_id={}, operation_kind={}, duration_ms={}", + operation_id, + operation_kind.as_str(), + operation_duration_ms + ); } res - }) - .await; + }); + + info!( + operation_id, + operation_kind = operation_kind.as_str(), + abort_deadline_remaining_ms = remaining_ms_until(abort_deadline), + "ducklake blocking operation waiting for blocking task or abort deadline: \ + operation_id={}, operation_kind={}, abort_deadline_remaining_ms={}", + operation_id, + operation_kind.as_str(), + remaining_ms_until(abort_deadline) + ); + let blocking_result = tokio::select! { + biased; + result = blocking_task => result, + _ = tokio::time::sleep_until(abort_deadline) => { + // The blocking task still owns the pooled connection here, so the + // async side cannot mark it broken and return it to r2d2 for + // eviction. If DuckDB does not return after interrupt plus grace, + // the stuck native call also keeps holding its semaphore permit and + // blocking thread, so restarting the process is the recoverable + // boundary. + abort_stuck_duckdb_blocking_operation( + operation_id, + operation_kind, + timeout, + BLOCKING_ABORT_GRACE, + ); + } + }; + + match &blocking_result { + Ok(Ok(_)) => info!( + operation_id, + operation_kind = operation_kind.as_str(), + timed_out = watchdog_timed_out.load(Ordering::Relaxed), + "ducklake blocking operation task joined with success: operation_id={}, \ + operation_kind={}, timed_out={}", + operation_id, + operation_kind.as_str(), + watchdog_timed_out.load(Ordering::Relaxed) + ), + Ok(Err(error)) => warn!( + operation_id, + operation_kind = operation_kind.as_str(), + timed_out = watchdog_timed_out.load(Ordering::Relaxed), + error = ?error, + "ducklake blocking operation task joined with error: operation_id={}, operation_kind={}, \ + timed_out={}, error={:?}", + operation_id, + operation_kind.as_str(), + watchdog_timed_out.load(Ordering::Relaxed), + error + ), + Err(error) => tracing::error!( + operation_id, + operation_kind = operation_kind.as_str(), + timed_out = watchdog_timed_out.load(Ordering::Relaxed), + error = %error, + "ducklake blocking operation task join failed: operation_id={}, operation_kind={}, \ + timed_out={}, error={}", + operation_id, + operation_kind.as_str(), + watchdog_timed_out.load(Ordering::Relaxed), + error + ), + } // Await the watchdog so it cannot outlive the finished blocking task and // accidentally interrupt a later operation that reuses the connection. - watchdog_task.await.map_err(|_| { - etl_error!(ErrorKind::ApplyWorkerPanic, "DuckLake query watchdog task panicked") - })?; + info!( + operation_id, + operation_kind = operation_kind.as_str(), + timed_out = watchdog_timed_out.load(Ordering::Relaxed), + "ducklake blocking operation awaiting watchdog task: operation_id={}, operation_kind={}, \ + timed_out={}", + operation_id, + operation_kind.as_str(), + watchdog_timed_out.load(Ordering::Relaxed) + ); + match watchdog_task.await { + Ok(()) => info!( + operation_id, + operation_kind = operation_kind.as_str(), + timed_out = watchdog_timed_out.load(Ordering::Relaxed), + "ducklake blocking operation watchdog task joined: operation_id={}, \ + operation_kind={}, timed_out={}", + operation_id, + operation_kind.as_str(), + watchdog_timed_out.load(Ordering::Relaxed) + ), + Err(error) => { + tracing::error!( + operation_id, + operation_kind = operation_kind.as_str(), + timed_out = watchdog_timed_out.load(Ordering::Relaxed), + error = %error, + "ducklake blocking operation watchdog task panicked: operation_id={}, operation_kind={}, \ + timed_out={}, error={}", + operation_id, + operation_kind.as_str(), + watchdog_timed_out.load(Ordering::Relaxed), + error + ); + return Err(etl_error!( + ErrorKind::ApplyWorkerPanic, + "DuckLake query watchdog task panicked" + )); + } + } blocking_result.map_err(|_| { etl_error!(ErrorKind::ApplyWorkerPanic, "DuckLake blocking operation task panicked") @@ -647,10 +1110,13 @@ where mod tests { use std::sync::Arc; - use tokio::sync::{Semaphore, oneshot}; + use tokio::sync::{Barrier, Semaphore, oneshot}; use super::*; + const WATCHDOG_DEADLINE: Duration = Duration::from_millis(20); + const WATCHDOG_JOIN_TIMEOUT: Duration = Duration::from_secs(1); + fn make_blocking_test_manager() -> DuckLakeConnectionManager { DuckLakeConnectionManager { setup_plan: Arc::new(DuckLakeSetupPlan::default()), @@ -661,6 +1127,138 @@ mod tests { } } + fn watchdog_interrupt_handle() -> (duckdb::Connection, DuckDbQueryInterruptHandle) { + let conn = + duckdb::Connection::open_in_memory().expect("failed to open watchdog test connection"); + let handle = conn.interrupt_handle(); + (conn, handle) + } + + struct DelayedInterruptHandle { + delay: Duration, + calls: std::sync::atomic::AtomicUsize, + started: AtomicBool, + completed: AtomicBool, + } + + impl DelayedInterruptHandle { + fn new(delay: Duration) -> Self { + Self { + delay, + calls: std::sync::atomic::AtomicUsize::new(0), + started: AtomicBool::new(false), + completed: AtomicBool::new(false), + } + } + + fn calls(&self) -> usize { + self.calls.load(Ordering::Relaxed) + } + + fn started(&self) -> bool { + self.started.load(Ordering::Relaxed) + } + + fn completed(&self) -> bool { + self.completed.load(Ordering::Relaxed) + } + } + + impl DuckDbQueryInterrupt for DelayedInterruptHandle { + fn interrupt(&self) { + self.calls.fetch_add(1, Ordering::Relaxed); + self.started.store(true, Ordering::Relaxed); + std::thread::sleep(self.delay); + self.completed.store(true, Ordering::Relaxed); + } + } + + fn delayed_interrupt_handle( + delay: Duration, + ) -> (Arc, DuckDbQueryInterruptHandle) { + let handle = Arc::new(DelayedInterruptHandle::new(delay)); + (Arc::clone(&handle), handle) + } + + async fn wait_for_delayed_interrupt_to_start(handle: &DelayedInterruptHandle) { + tokio::time::timeout(WATCHDOG_JOIN_TIMEOUT, async { + while !handle.started() { + tokio::task::yield_now().await; + } + }) + .await + .expect("delayed interrupt should start"); + } + + #[derive(Clone, Copy, Debug)] + enum WatchdogSignal { + PublishInterrupt, + Finish, + DropInterruptSender, + DropDoneSender, + } + + type WatchdogSignalCase = + (&'static str, &'static [WatchdogSignal], bool, &'static [WatchdogSignal], bool); + + fn apply_watchdog_signal( + watchdog: &mut DuckDbQueryWatchdog, + handle: &DuckDbQueryInterruptHandle, + signal: WatchdogSignal, + ) { + match signal { + WatchdogSignal::PublishInterrupt => { + watchdog.publish_query_interrupt_handle(Arc::clone(handle)); + } + WatchdogSignal::Finish => { + watchdog.finish(); + } + WatchdogSignal::DropInterruptSender => { + drop(watchdog.interrupt_tx.take()); + } + WatchdogSignal::DropDoneSender => { + drop(watchdog.done_tx.take()); + } + } + } + + async fn assert_watchdog_sequence_completes( + name: &str, + before_deadline: &[WatchdogSignal], + wait_past_deadline: bool, + after_deadline: &[WatchdogSignal], + expected_timed_out: bool, + ) { + let (_conn, interrupt_handle) = watchdog_interrupt_handle(); + let mut watchdog = DuckDbQueryWatchdog::spawn(Instant::now() + WATCHDOG_DEADLINE); + let watchdog_task = watchdog.async_task_handle().expect("failed to extract watchdog task"); + + for signal in before_deadline { + apply_watchdog_signal(&mut watchdog, &interrupt_handle, *signal); + tokio::task::yield_now().await; + } + + if wait_past_deadline { + tokio::time::sleep(WATCHDOG_DEADLINE * 2).await; + } + + for signal in after_deadline { + apply_watchdog_signal(&mut watchdog, &interrupt_handle, *signal); + tokio::task::yield_now().await; + } + + tokio::time::timeout(WATCHDOG_JOIN_TIMEOUT, watchdog_task) + .await + .unwrap_or_else(|_| panic!("watchdog sequence `{name}` deadlocked")) + .unwrap_or_else(|error| panic!("watchdog sequence `{name}` panicked: {error}")); + + assert_eq!( + watchdog.timed_out(), + expected_timed_out, + "unexpected timeout state for watchdog sequence `{name}`" + ); + } + #[test] fn duckdb_blocking_operation_kind_timeouts() { assert_eq!(DuckDbBlockingOperationKind::Foreground.timeout(), FOREGROUND_QUERY_TIMEOUT); @@ -775,6 +1373,274 @@ mod tests { ); } + #[tokio::test] + async fn query_watchdog_signal_order_matrix_does_not_deadlock() { + use WatchdogSignal::{DropDoneSender, DropInterruptSender, Finish, PublishInterrupt}; + + let cases: &[WatchdogSignalCase] = &[ + ("finish_before_interrupt_handle", &[Finish], false, &[], false), + ("done_sender_dropped_before_interrupt_handle", &[DropDoneSender], false, &[], false), + ("interrupt_sender_dropped_before_deadline", &[DropInterruptSender], false, &[], false), + ( + "interrupt_handle_then_finish_before_deadline", + &[PublishInterrupt, Finish], + false, + &[], + false, + ), + ( + "interrupt_handle_then_done_sender_drop_before_deadline", + &[PublishInterrupt, DropDoneSender], + false, + &[], + false, + ), + ( + "finish_then_interrupt_handle_before_deadline", + &[Finish, PublishInterrupt], + false, + &[], + false, + ), + ( + "interrupt_handle_before_deadline_then_deadline", + &[PublishInterrupt], + true, + &[], + true, + ), + ("deadline_then_interrupt_handle", &[], true, &[PublishInterrupt], true), + ("deadline_then_finish", &[], true, &[Finish], true), + ("deadline_then_interrupt_sender_drop", &[], true, &[DropInterruptSender], true), + ("deadline_then_done_sender_drop", &[], true, &[DropDoneSender], true), + ( + "deadline_then_both_senders_drop", + &[], + true, + &[DropDoneSender, DropInterruptSender], + true, + ), + ( + "interrupt_handle_before_deadline_then_finish_after_deadline", + &[PublishInterrupt], + true, + &[Finish], + true, + ), + ]; + + for (name, before_deadline, wait_past_deadline, after_deadline, expected_timed_out) in cases + { + assert_watchdog_sequence_completes( + name, + before_deadline, + *wait_past_deadline, + after_deadline, + *expected_timed_out, + ) + .await; + } + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn query_watchdog_concurrent_signal_races_do_not_deadlock() { + for signal_after_deadline in [false, true] { + for iteration in 0..50 { + let (_conn, interrupt_handle) = watchdog_interrupt_handle(); + let deadline = if signal_after_deadline { + Instant::now() + WATCHDOG_DEADLINE + } else { + Instant::now() + Duration::from_secs(1) + }; + let mut watchdog = DuckDbQueryWatchdog::spawn(deadline); + let interrupt_tx = + watchdog.interrupt_tx.take().expect("watchdog interrupt sender should exist"); + let done_tx = watchdog.done_tx.take().expect("watchdog done sender should exist"); + let watchdog_task = + watchdog.async_task_handle().expect("failed to extract watchdog task"); + let barrier = Arc::new(Barrier::new(3)); + + let interrupt_sender = tokio::spawn({ + let barrier = Arc::clone(&barrier); + let interrupt_handle = Arc::clone(&interrupt_handle); + async move { + barrier.wait().await; + let _ = interrupt_tx.send(interrupt_handle); + } + }); + let done_sender = tokio::spawn({ + let barrier = Arc::clone(&barrier); + async move { + barrier.wait().await; + let _ = done_tx.send(()); + } + }); + + if signal_after_deadline { + tokio::time::sleep(WATCHDOG_DEADLINE * 2).await; + } + + barrier.wait().await; + + tokio::time::timeout(WATCHDOG_JOIN_TIMEOUT, watchdog_task) + .await + .unwrap_or_else(|_| { + panic!( + "watchdog concurrent signal race deadlocked: \ + signal_after_deadline={signal_after_deadline}, iteration={iteration}" + ) + }) + .unwrap_or_else(|error| { + panic!( + "watchdog concurrent signal race panicked: \ + signal_after_deadline={signal_after_deadline}, \ + iteration={iteration}, error={error}" + ) + }); + interrupt_sender.await.expect("interrupt sender task should not panic"); + done_sender.await.expect("done sender task should not panic"); + + assert_eq!( + watchdog.timed_out(), + signal_after_deadline, + "unexpected timeout state for concurrent signal race: \ + signal_after_deadline={signal_after_deadline}, iteration={iteration}" + ); + } + } + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn query_watchdog_slow_interrupt_before_deadline_does_not_deadlock() { + let (delayed_handle, interrupt_handle) = + delayed_interrupt_handle(Duration::from_millis(100)); + let mut watchdog = DuckDbQueryWatchdog::spawn(Instant::now() + WATCHDOG_DEADLINE); + let watchdog_task = watchdog.async_task_handle().expect("failed to extract watchdog task"); + + watchdog.publish_query_interrupt_handle(interrupt_handle); + + tokio::time::timeout(WATCHDOG_JOIN_TIMEOUT, watchdog_task) + .await + .expect("watchdog with slow interrupt should not deadlock") + .expect("watchdog task should not panic"); + + assert!(watchdog.timed_out()); + assert_eq!(delayed_handle.calls(), 1); + assert!(delayed_handle.completed()); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn query_watchdog_slow_interrupt_after_deadline_does_not_deadlock() { + let (delayed_handle, interrupt_handle) = + delayed_interrupt_handle(Duration::from_millis(100)); + let mut watchdog = DuckDbQueryWatchdog::spawn(Instant::now() + WATCHDOG_DEADLINE); + let watchdog_task = watchdog.async_task_handle().expect("failed to extract watchdog task"); + + tokio::time::sleep(WATCHDOG_DEADLINE * 2).await; + watchdog.publish_query_interrupt_handle(interrupt_handle); + + tokio::time::timeout(WATCHDOG_JOIN_TIMEOUT, watchdog_task) + .await + .expect("watchdog with slow late interrupt should not deadlock") + .expect("watchdog task should not panic"); + + assert!(watchdog.timed_out()); + assert_eq!(delayed_handle.calls(), 1); + assert!(delayed_handle.completed()); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] + async fn query_watchdog_finish_while_slow_interrupt_is_running_does_not_deadlock() { + let (delayed_handle, interrupt_handle) = + delayed_interrupt_handle(Duration::from_millis(100)); + let mut watchdog = DuckDbQueryWatchdog::spawn(Instant::now() + WATCHDOG_DEADLINE); + let watchdog_task = watchdog.async_task_handle().expect("failed to extract watchdog task"); + + watchdog.publish_query_interrupt_handle(interrupt_handle); + wait_for_delayed_interrupt_to_start(&delayed_handle).await; + watchdog.finish(); + + tokio::time::timeout(WATCHDOG_JOIN_TIMEOUT, watchdog_task) + .await + .expect("watchdog should finish after slow interrupt returns") + .expect("watchdog task should not panic"); + + assert!(watchdog.timed_out()); + assert_eq!(delayed_handle.calls(), 1); + assert!(delayed_handle.completed()); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] + async fn query_watchdog_slow_interrupt_starves_single_worker_runtime_until_it_returns() { + let interrupt_delay = Duration::from_millis(100); + let (delayed_handle, interrupt_handle) = delayed_interrupt_handle(interrupt_delay); + let mut watchdog = DuckDbQueryWatchdog::spawn(Instant::now() + WATCHDOG_DEADLINE); + let watchdog_task = watchdog.async_task_handle().expect("failed to extract watchdog task"); + let (observer_started_tx, observer_started_rx) = oneshot::channel(); + let started_at = std::time::Instant::now(); + let observer_task = tokio::spawn(async move { + observer_started_tx.send(()).expect("observer start receiver should be open"); + tokio::time::sleep(WATCHDOG_DEADLINE * 2).await; + std::time::Instant::now() + }); + + observer_started_rx.await.expect("observer should start"); + watchdog.publish_query_interrupt_handle(interrupt_handle); + + tokio::time::timeout(WATCHDOG_JOIN_TIMEOUT, watchdog_task) + .await + .expect("watchdog should not deadlock on a single worker runtime") + .expect("watchdog task should not panic"); + let observer_finished_at = observer_task.await.expect("observer task should not panic"); + + assert!(watchdog.timed_out()); + assert_eq!(delayed_handle.calls(), 1); + assert!(delayed_handle.completed()); + assert!( + observer_finished_at.duration_since(started_at) >= interrupt_delay, + "single worker runtime should not poll other async tasks while interrupt() blocks" + ); + } + + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] + async fn query_watchdog_slow_interrupts_serialize_on_single_worker_runtime() { + let interrupt_delay = Duration::from_millis(100); + let (first_delayed_handle, first_interrupt_handle) = + delayed_interrupt_handle(interrupt_delay); + let (second_delayed_handle, second_interrupt_handle) = + delayed_interrupt_handle(interrupt_delay); + let mut first_watchdog = DuckDbQueryWatchdog::spawn(Instant::now() + WATCHDOG_DEADLINE); + let mut second_watchdog = DuckDbQueryWatchdog::spawn(Instant::now() + WATCHDOG_DEADLINE); + let first_watchdog_task = + first_watchdog.async_task_handle().expect("failed to extract first watchdog task"); + let second_watchdog_task = + second_watchdog.async_task_handle().expect("failed to extract second watchdog task"); + let started_at = std::time::Instant::now(); + + first_watchdog.publish_query_interrupt_handle(first_interrupt_handle); + second_watchdog.publish_query_interrupt_handle(second_interrupt_handle); + + tokio::time::timeout(WATCHDOG_JOIN_TIMEOUT, async { + let (first_result, second_result) = + tokio::join!(first_watchdog_task, second_watchdog_task); + first_result.expect("first watchdog task should not panic"); + second_result.expect("second watchdog task should not panic"); + }) + .await + .expect("serialized slow interrupts should not deadlock on a single worker runtime"); + + assert!(first_watchdog.timed_out()); + assert!(second_watchdog.timed_out()); + assert_eq!(first_delayed_handle.calls(), 1); + assert_eq!(second_delayed_handle.calls(), 1); + assert!(first_delayed_handle.completed()); + assert!(second_delayed_handle.completed()); + assert!( + started_at.elapsed() >= interrupt_delay * 2, + "single worker runtime should serialize blocking interrupt() calls" + ); + } + #[tokio::test] async fn query_watchdog_marks_timeout_when_handle_arrives_after_deadline() { let conn = make_blocking_test_manager() diff --git a/crates/etl-destinations/src/ducklake/core.rs b/crates/etl-destinations/src/ducklake/core.rs index 6e8eea05b..6e0742463 100644 --- a/crates/etl-destinations/src/ducklake/core.rs +++ b/crates/etl-destinations/src/ducklake/core.rs @@ -1,12 +1,11 @@ #[cfg(feature = "test-utils")] use std::sync::atomic::AtomicUsize; +#[cfg(test)] +use std::time::Duration; use std::{ collections::{HashMap, HashSet}, - sync::{ - Arc, - atomic::{AtomicBool, Ordering}, - }, - time::Duration, + sync::{Arc, atomic::AtomicBool}, + time::Instant, }; use etl::{ @@ -27,14 +26,17 @@ use etl::{ use metrics::gauge; use parking_lot::Mutex; use pg_escape::quote_identifier; -use sqlx::{PgPool, postgres::PgPoolOptions}; +use sqlx::{AssertSqlSafe, PgPool, postgres::PgPoolOptions}; #[cfg(feature = "test-utils")] use tokio::sync::oneshot; use tokio::{ - sync::{OwnedRwLockReadGuard, OwnedSemaphorePermit, RwLock, Semaphore, TryAcquireError}, + sync::{ + OwnedRwLockReadGuard, OwnedRwLockWriteGuard, OwnedSemaphorePermit, RwLock, Semaphore, + TryAcquireError, + }, task::JoinSet, }; -use tracing::{debug, info}; +use tracing::{debug, info, warn}; use url::Url; use crate::{ @@ -58,14 +60,11 @@ use crate::{ maintenance_target_file_size_sql, resolve_expire_snapshots_older_than, validate_expire_snapshots_older_than_sql, }, + external_maintenance::{ExternalMaintenanceOperations, run_external_maintenance_watcher}, inline_size::DuckLakePendingInlineSizeSampler, - maintenance::{ - DuckLakeMaintenanceWorker, PendingInlineFlushRequests, TableMaintenanceNotification, - TableWriteActivity, send_maintenance_notification, spawn_ducklake_maintenance_worker, - table_write_slot, - }, metrics::{ - DuckLakeMetricsSampler, ETL_DUCKLAKE_POOL_SIZE, register_metrics, + DuckLakeMetricsSampler, ETL_DUCKLAKE_POOL_SIZE, query_catalog_maintenance_metrics, + query_table_storage_metrics, register_metrics, resolve_ducklake_metadata_schema_blocking, spawn_ducklake_metrics_sampler, }, schema::build_create_table_sql_ducklake, @@ -104,6 +103,14 @@ pub(super) fn is_create_table_conflict(error: &duckdb::Error, table_name: &str) && message.contains(&format!(r#"attempting to create table "{table_name}""#)) } +/// Parses `expire_snapshots_older_than` into seconds for cheap metadata-only +/// trigger sampling. +fn expire_snapshots_retention_seconds(value: &str) -> Option { + humantime::parse_duration(value) + .ok() + .and_then(|duration| i64::try_from(duration.as_secs()).ok()) +} + // ── destination // ─────────────────────────────────────────────────────────────── @@ -123,12 +130,14 @@ pub struct DuckLakeDestination { manager: Arc, pool: Arc>, blocking_slots: Arc, - /// Shared gate that keeps exclusive background maintenance from overlapping + /// Shared gate that keeps external maintenance pauses from overlapping /// active foreground or table-scoped mutations. checkpoint_gate: Arc>, tasks: TaskSet, - maintenance_worker: Arc>, metrics_sampler: Arc>, + metadata_schema: Arc, + expire_snapshots_older_than: Arc, + metadata_pg_pool: PgPool, table_creation_slots: Arc, table_write_slots: Arc>>>, store: S, @@ -139,12 +148,22 @@ pub struct DuckLakeDestination { applied_batches_table_created: Arc, /// Cache tracking whether the ETL streaming progress table already exists. streaming_progress_table_created: Arc, - /// Signals that one or more inline flushes should run before the next safe - /// ingest point. - inline_flush_requested: Arc, - /// Tracks which tables need a requested inline flush before ingestion - /// resumes. - inline_flush_requests: Arc, +} + +/// Held by an external DuckLake maintenance coordinator while foreground +/// mutations must be quiesced. +pub struct DuckLakeExternalMaintenancePause { + _guard: OwnedRwLockWriteGuard<()>, +} + +/// Returns the table-local semaphore shared by concurrent foreground writes. +fn table_write_slot( + table_write_slots: &Arc>>>, + table_name: &str, +) -> Arc { + let mut slots = table_write_slots.lock(); + let slot = slots.entry(table_name.to_owned()).or_insert_with(|| Arc::new(Semaphore::new(1))); + Arc::clone(slot) } impl Destination for DuckLakeDestination @@ -162,7 +181,6 @@ where "ducklake shutdown requested, interrupted active duckdb connections" ); self.tasks.shutdown().await?; - self.shutdown_maintenance_worker().await?; self.shutdown_metrics_sampler().await?; Ok(()) @@ -468,12 +486,8 @@ where let blocking_slots = Arc::new(Semaphore::new(pool_size as usize)); // `target_file_size` is a catalog-wide DuckLake option consumed during - // compaction. Apply it once on the write pool before the maintenance - // pool starts warming, so the maintenance worker does not need to - // mutate the catalog from a separate RW DuckDB instance during its - // background warm-up. Two RW instances ATTACHing the same catalog file - // and racing a catalog write against concurrent user writes caused - // lost commits. + // compaction. Apply it once on the write pool so foreground writes and + // external maintenance jobs use the same configured catalog option. let target_file_size_sql = maintenance_target_file_size_sql(Some(maintenance_target_file_size.as_ref())); run_duckdb_blocking( @@ -529,74 +543,50 @@ where .await? } }; + let metadata_schema = Arc::::from(metadata_schema); let metadata_pg_pool = build_ducklake_metadata_pg_pool(&catalog_url)?; - let pending_inline_size_sampler = Some(DuckLakePendingInlineSizeSampler::new( - metadata_schema.clone(), - metadata_pg_pool.clone(), - )); let created_tables = Arc::default(); let checkpoint_gate = Arc::new(RwLock::new(())); - let inline_flush_requested = Arc::new(AtomicBool::new(false)); - let inline_flush_requests = Arc::new(PendingInlineFlushRequests::default()); let mut destination = Self { manager: Arc::clone(&manager), pool: Arc::clone(&pool), blocking_slots: Arc::clone(&blocking_slots), checkpoint_gate: Arc::clone(&checkpoint_gate), tasks: TaskSet::new(), - maintenance_worker: Arc::new(None), metrics_sampler: Arc::new(None), + metadata_schema: Arc::clone(&metadata_schema), + expire_snapshots_older_than: Arc::clone(&expire_snapshots_older_than), + metadata_pg_pool: metadata_pg_pool.clone(), table_creation_slots: Arc::new(Semaphore::new(1)), table_write_slots: Arc::default(), store, created_tables: Arc::clone(&created_tables), applied_batches_table_created: Arc::default(), streaming_progress_table_created: Arc::default(), - inline_flush_requested: Arc::clone(&inline_flush_requested), - inline_flush_requests: Arc::clone(&inline_flush_requests), }; gauge!(ETL_DUCKLAKE_POOL_SIZE).set(pool_size as f64); destination.ensure_applied_batches_table_exists().await?; destination.ensure_streaming_progress_table_exists().await?; - destination.maintenance_worker = Arc::new( - spawn_ducklake_maintenance_worker( - DuckLakeConnectionManager { - setup_plan: Arc::clone(&setup_plan), - disable_extension_autoload, - interrupt_registry: manager.interrupt_registry(), - #[cfg(feature = "test-utils")] - open_count: Arc::new(AtomicUsize::new(0)), - }, - Arc::clone(&checkpoint_gate), - Arc::clone(&destination.table_write_slots), - Arc::clone(&inline_flush_requested), - Arc::clone(&inline_flush_requests), - pending_inline_size_sampler, - Arc::clone(&expire_snapshots_older_than), - )? - .into(), - ); destination.metrics_sampler = Arc::new( spawn_ducklake_metrics_sampler( - metadata_schema, - metadata_pg_pool, + metadata_schema.to_string(), + metadata_pg_pool.clone(), Arc::clone(&created_tables), - destination - .maintenance_worker - .as_ref() - .as_ref() - .ok_or_else(|| { - etl_error!( - ErrorKind::DestinationError, - "DuckLake initialization failed", - "Maintenance worker should exist before metrics sampler" - ) - })? - .notification_tx - .clone(), )? .into(), ); + let watcher_destination = destination.clone(); + destination + .tasks + .spawn(async move { + if let Err(error) = run_external_maintenance_watcher(watcher_destination).await { + warn!( + error = %error, + "ducklake external maintenance watcher exited" + ); + } + }) + .await; Ok(destination) } @@ -683,8 +673,8 @@ where /// Copy batches are recorded in the replay marker table so a retry after an /// ambiguous post-commit failure can detect already applied rows. /// - /// Small copy batches may stay inlined until maintenance requests a safe - /// materialization after we emit the maintenance notification. + /// Small copy batches may stay inlined until external maintenance + /// materializes them during a coordinated pause. async fn write_table_rows_inner( &self, replicated_table_schema: &ReplicatedTableSchema, @@ -696,14 +686,6 @@ where return Ok(()); } - if let Err(error) = self.maybe_run_requested_inline_flush().await { - tracing::error!( - table = %table_name, - error = %error, - "ducklake inline flush failed" - ); - } - // Copy batches for the same table must still serialize so concurrent // callers do not race each other inside DuckDB. self.ensure_applied_batches_table_exists().await?; @@ -711,17 +693,12 @@ where let _checkpoint_guard = self.acquire_mutation_guard().await; let prepared_batch = prepare_copy_table_batch(replicated_table_schema, table_name, table_rows)?; - let table_name = prepared_batch.table_name().to_owned(); apply_table_batch_with_retry( Arc::clone(&self.pool), Arc::clone(&self.blocking_slots), prepared_batch, ) .await?; - self.notify_background_maintenance(TableMaintenanceNotification::WriteActivity( - TableWriteActivity { table_name }, - )) - .await; Ok(()) } @@ -736,12 +713,6 @@ where /// streaming replay watermark so retries can safely detect already /// committed work. async fn write_events_inner(&self, events: Vec) -> EtlResult<()> { - if let Err(error) = self.maybe_run_requested_inline_flush().await { - tracing::error!( - error = %error, - "ducklake inline flush failed" - ); - } let mut event_iter = events.into_iter().peekable(); while event_iter.peek().is_some() { @@ -876,20 +847,24 @@ where let pool = Arc::clone(&self.pool); let blocking_slots = Arc::clone(&self.blocking_slots); let destination_table_name = table_name.clone(); - let maintenance_worker = Arc::clone(&self.maintenance_worker); join_set.spawn(async move { let _table_write_permit = table_write_permit; let checkpoint_wait_started = tokio::time::Instant::now(); + info!( + table = %destination_table_name, + "ducklake waiting for checkpoint gate before streaming write: table={}", + destination_table_name + ); let _checkpoint_guard = checkpoint_gate.read_owned().await; let checkpoint_wait = checkpoint_wait_started.elapsed(); - if checkpoint_wait > Duration::from_secs(1) { - info!( - table = %destination_table_name, - checkpoint_wait_ms = checkpoint_wait.as_millis() as u64, - "ducklake waited for checkpoint gate before streaming write" - ); - } + info!( + table = %destination_table_name, + checkpoint_wait_ms = checkpoint_wait.as_millis() as u64, + "ducklake acquired checkpoint gate before streaming write: table={}, checkpoint_wait_ms={}", + destination_table_name, + checkpoint_wait.as_millis() + ); let last_sequence_key = read_table_streaming_progress_sequence_key_blocking( Arc::clone(&pool), @@ -902,7 +877,8 @@ where if pending_mutations.is_empty() { debug!( table = %destination_table_name, - "ducklake streaming mutation replay skipped, no pending events" + "ducklake streaming mutation replay skipped, no pending events: table={}", + destination_table_name ); return Ok::<(), etl::error::EtlError>(()); } @@ -911,15 +887,12 @@ where table = %destination_table_name, pending_mutation_count = pending_mutations.len(), is_first_streaming_batch, - "ducklake applying streaming mutations" + "ducklake applying streaming mutations: table={}, pending_mutation_count={}, is_first_streaming_batch={}", + destination_table_name, + pending_mutations.len(), + is_first_streaming_batch ); - let maintenance_notification = - maintenance_worker.as_ref().as_ref().map(|_| { - TableMaintenanceNotification::WriteActivity(TableWriteActivity { - table_name: destination_table_name.clone(), - }) - }); let prepared_batches = prepare_mutation_table_batches( &replicated_table_schema, destination_table_name.clone(), @@ -930,14 +903,10 @@ where info!( table = %destination_table_name, is_first_streaming_batch, - "ducklake applied streaming mutations" + "ducklake applied streaming mutations: table={}, is_first_streaming_batch={}", + destination_table_name, + is_first_streaming_batch ); - if let (Some(worker), Some(notification)) = - (maintenance_worker.as_ref(), maintenance_notification) - { - send_maintenance_notification(&worker.notification_tx, notification) - .await; - } Ok::<(), etl::error::EtlError>(()) }); } @@ -985,7 +954,21 @@ where let blocking_slots = Arc::clone(&self.blocking_slots); join_set.spawn(async move { let _table_write_permit = table_write_permit; + let checkpoint_wait_started = tokio::time::Instant::now(); + info!( + table = %table_name, + "ducklake waiting for checkpoint gate before streaming truncate: table={}", + table_name + ); let _checkpoint_guard = checkpoint_gate.read_owned().await; + let checkpoint_wait = checkpoint_wait_started.elapsed(); + info!( + table = %table_name, + checkpoint_wait_ms = checkpoint_wait.as_millis() as u64, + "ducklake acquired checkpoint gate before streaming truncate: table={}, checkpoint_wait_ms={}", + table_name, + checkpoint_wait.as_millis() + ); let last_sequence_key = read_table_streaming_progress_sequence_key_blocking( Arc::clone(&pool), @@ -998,11 +981,19 @@ where if pending_truncates.is_empty() { debug!( table = %table_name, - "ducklake streaming truncate replay skipped, no pending events" + "ducklake streaming truncate replay skipped, no pending events: table={}", + table_name ); return Ok(()); } + info!( + table = %table_name, + pending_truncate_count = pending_truncates.len(), + "ducklake applying streaming truncates: table={}, pending_truncate_count={}", + table_name, + pending_truncates.len() + ); let prepared_batch = prepare_truncate_table_batch(table_name, pending_truncates); apply_table_batch_with_retry(pool, blocking_slots, prepared_batch).await @@ -1039,7 +1030,9 @@ where info!( table_id = %table_id, table = %table_name, - "ducklake destination table cache miss, ensuring table exists" + "ducklake destination table cache miss, ensuring table exists: table_id={}, table={}", + table_id, + table_name ); let _table_creation_permit = @@ -1161,14 +1154,19 @@ where Err(TryAcquireError::NoPermits) => { info!( table = %table_name, - "ducklake waiting for table write slot" + "ducklake waiting for table write slot: table={}", + table_name ); + let started = Instant::now(); let permit = table_slot.acquire_owned().await.map_err(|_| { etl_error!(ErrorKind::InvalidState, "DuckLake table write semaphore closed") })?; info!( table = %table_name, - "ducklake acquired table write slot after wait" + wait_ms = started.elapsed().as_millis() as u64, + "ducklake acquired table write slot after wait: table={}, wait_ms={}", + table_name, + started.elapsed().as_millis() ); Ok(permit) } @@ -1181,25 +1179,145 @@ where /// Acquires shared mutation access so exclusive background maintenance /// cannot start in the middle of a foreground write sequence. async fn acquire_mutation_guard(&self) -> OwnedRwLockReadGuard<()> { - Arc::clone(&self.checkpoint_gate).read_owned().await + let started = Instant::now(); + info!( + metadata_schema = %self.metadata_schema, + "ducklake waiting for shared mutation guard: metadata_schema={}", + self.metadata_schema + ); + let guard = Arc::clone(&self.checkpoint_gate).read_owned().await; + info!( + metadata_schema = %self.metadata_schema, + wait_ms = started.elapsed().as_millis() as u64, + "ducklake acquired shared mutation guard: metadata_schema={}, wait_ms={}", + self.metadata_schema, + started.elapsed().as_millis() + ); + guard } - /// Runs requested inline flushes before foreground ingestion begins when - /// safe. - async fn maybe_run_requested_inline_flush(&self) -> EtlResult<()> { - if !self.inline_flush_requested.load(Ordering::Acquire) { - return Ok(()); - } + /// Acquires exclusive DuckLake mutation access for an external maintenance + /// run. While this guard is held, new foreground writes and in-process + /// background maintenance operations wait before mutating the catalog. + pub async fn acquire_external_maintenance_pause(&self) -> DuckLakeExternalMaintenancePause { + let started = Instant::now(); + info!( + metadata_schema = %self.metadata_schema, + "ducklake waiting for exclusive external maintenance mutation guard: metadata_schema={}", + self.metadata_schema + ); + let guard = Arc::clone(&self.checkpoint_gate).write_owned().await; + info!( + metadata_schema = %self.metadata_schema, + wait_ms = started.elapsed().as_millis() as u64, + "ducklake acquired exclusive external maintenance mutation guard: metadata_schema={}, wait_ms={}", + self.metadata_schema, + started.elapsed().as_millis() + ); + DuckLakeExternalMaintenancePause { _guard: guard } + } - crate::ducklake::maintenance::maybe_run_requested_inline_flush( - Arc::clone(&self.pool), - Arc::clone(&self.checkpoint_gate), - Arc::clone(&self.blocking_slots), - self.inline_flush_requested.as_ref(), - self.inline_flush_requests.as_ref(), - self.maintenance_worker.as_ref().as_ref().map(|worker| worker.notification_tx.clone()), + /// Samples catalog state and returns which externally coordinated + /// maintenance operations should be requested now. + pub(super) async fn sample_external_maintenance_operations( + &self, + inline_flush_min_inlined_bytes: u64, + rewrite_data_files_min_active_data_files: i64, + ) -> EtlResult { + let table_names = self.list_active_ducklake_tables().await?; + let inline_sampler = DuckLakePendingInlineSizeSampler::new( + self.metadata_schema.to_string(), + self.metadata_pg_pool.clone(), + ); + let mut operations = ExternalMaintenanceOperations::default(); + let catalog_metrics = query_catalog_maintenance_metrics( + &self.metadata_pg_pool, + self.metadata_schema.as_ref(), ) - .await + .await?; + + match expire_snapshots_retention_seconds(self.expire_snapshots_older_than.as_ref()) { + Some(retention_seconds) => { + operations.expire_snapshots = catalog_metrics.snapshots_total > 1 + && catalog_metrics.oldest_snapshot_age_seconds >= retention_seconds; + debug!( + metadata_schema = %self.metadata_schema, + expire_snapshots_older_than = %self.expire_snapshots_older_than, + retention_seconds, + snapshots_total = catalog_metrics.snapshots_total, + oldest_snapshot_age_seconds = catalog_metrics.oldest_snapshot_age_seconds, + expire_snapshots = operations.expire_snapshots, + "ducklake sampled expire snapshots trigger: metadata_schema={}, \ + expire_snapshots_older_than={}, retention_seconds={}, snapshots_total={}, \ + oldest_snapshot_age_seconds={}, expire_snapshots={}", + self.metadata_schema, + self.expire_snapshots_older_than, + retention_seconds, + catalog_metrics.snapshots_total, + catalog_metrics.oldest_snapshot_age_seconds, + operations.expire_snapshots + ); + } + None => { + warn!( + metadata_schema = %self.metadata_schema, + expire_snapshots_older_than = %self.expire_snapshots_older_than, + "ducklake could not parse expire_snapshots_older_than for external maintenance \ + trigger sampling: metadata_schema={}, expire_snapshots_older_than={}", + self.metadata_schema, + self.expire_snapshots_older_than + ); + } + } + + for table_name in table_names { + if table_name.starts_with("__etl_") { + continue; + } + + if !operations.inline_flush { + let sizes = inline_sampler.sample_table(&table_name).await?; + operations.inline_flush = sizes.inlined_bytes >= inline_flush_min_inlined_bytes; + } + + if !operations.rewrite_data_files { + let metrics = query_table_storage_metrics( + &self.metadata_pg_pool, + self.metadata_schema.as_ref(), + &table_name, + ) + .await?; + operations.rewrite_data_files = + metrics.active_data_files > rewrite_data_files_min_active_data_files; + } + + if operations.inline_flush && operations.rewrite_data_files { + break; + } + } + + Ok(operations) + } + + /// Lists active DuckLake table names from the metadata catalog. + async fn list_active_ducklake_tables(&self) -> EtlResult> { + let sql = format!( + "SELECT table_name FROM {}.{} WHERE end_snapshot IS NULL ORDER BY table_name", + quote_identifier(self.metadata_schema.as_ref()), + quote_identifier("ducklake_table") + ); + let rows: Vec<(String,)> = sqlx::query_as(AssertSqlSafe(sql)) + .fetch_all(&self.metadata_pg_pool) + .await + .map_err(|source| { + etl_error!( + ErrorKind::DestinationQueryFailed, + "DuckLake table list query failed", + format!("metadata_schema={}", self.metadata_schema.as_ref()), + source: source + ) + })?; + Ok(rows.into_iter().map(|(table_name,)| table_name).collect()) } /// Runs one DuckDB operation on Tokio's blocking pool after acquiring a @@ -1222,27 +1340,6 @@ where .await } - /// Stops the background DuckLake maintenance worker. - async fn shutdown_maintenance_worker(&self) -> EtlResult<()> { - if let Some(maintenance_worker) = &*self.maintenance_worker { - let _ = maintenance_worker.shutdown_tx.send(()); - let handle = maintenance_worker.handle.lock().take(); - if let Some(handle) = handle { - handle.abort(); - if let Err(err) = handle.await - && !err.is_cancelled() - { - return Err(etl_error!( - ErrorKind::ApplyWorkerPanic, - "DuckLake maintenance worker task panicked" - )); - } - } - } - - Ok(()) - } - /// Stops the background DuckLake metrics sampler. async fn shutdown_metrics_sampler(&self) -> EtlResult<()> { if let Some(metrics_sampler) = &*self.metrics_sampler { @@ -1255,7 +1352,7 @@ where { return Err(etl_error!( ErrorKind::ApplyWorkerPanic, - "DuckLake maintenance worker task panicked" + "DuckLake metrics sampler task panicked" )); } } @@ -1263,14 +1360,6 @@ where Ok(()) } - - /// Sends one background-maintenance notification to the maintenance worker. - async fn notify_background_maintenance(&self, notification: TableMaintenanceNotification) { - if let Some(maintenance_worker) = &*self.maintenance_worker { - send_maintenance_notification(&maintenance_worker.notification_tx, notification).await; - } - } - /// Returns how many DuckDB connections have been initialized for tests. #[cfg(feature = "test-utils")] pub fn connection_open_count_for_tests(&self) -> usize { @@ -1388,12 +1477,19 @@ mod tests { use super::*; use crate::ducklake::{ config::catalog_conninfo_from_url, - maintenance::flush_table_inlined_data, + maintenance_runner::flush_table_inlined_data, metrics::{query_catalog_maintenance_metrics, query_table_storage_metrics}, }; const POSTGRES_SCANNER_EXTENSION_FILE: &str = "postgres_scanner.duckdb_extension"; + #[test] + fn expire_snapshots_retention_seconds_uses_humantime_duration_syntax() { + assert_eq!(expire_snapshots_retention_seconds("7 days"), Some(604_800)); + assert_eq!(expire_snapshots_retention_seconds("2h 30min"), Some(9_000)); + assert_eq!(expire_snapshots_retention_seconds("bad interval"), None); + } + fn make_schema(table_id: u32, schema: &str, table: &str) -> TableSchema { TableSchema::new( TableId::new(table_id), diff --git a/crates/etl-destinations/src/ducklake/external_maintenance.rs b/crates/etl-destinations/src/ducklake/external_maintenance.rs new file mode 100644 index 000000000..b035f64a0 --- /dev/null +++ b/crates/etl-destinations/src/ducklake/external_maintenance.rs @@ -0,0 +1,771 @@ +use std::{env, time::Duration}; + +use chrono::{DateTime, Utc}; +use etl::store::{schema::SchemaStore, state::StateStore}; +use kube::{ + Api, Client, + api::{Patch, PatchParams}, + core::{ApiResource, DynamicObject, GroupVersionKind}, +}; +use metrics::{counter, histogram}; +use serde_json::json; +use tokio::time; +use tracing::{debug, info, warn}; + +use crate::ducklake::{ + DuckLakeDestination, DuckLakeExternalMaintenancePause, + metrics::{ + ETL_DUCKLAKE_EXTERNAL_MAINTENANCE_PAUSE_DURATION_SECONDS, + ETL_DUCKLAKE_EXTERNAL_MAINTENANCE_TRIGGERED_TOTAL, MAINTENANCE_OPERATION_LABEL, + MAINTENANCE_OUTCOME_LABEL, MAINTENANCE_REASON_LABEL, + }, +}; + +const CR_NAME_ENV: &str = "ETL_DUCKLAKE_MAINTENANCE_CR_NAME"; +const CR_NAMESPACE_ENV: &str = "ETL_DUCKLAKE_MAINTENANCE_CR_NAMESPACE"; +const POLL_SECONDS_ENV: &str = "ETL_DUCKLAKE_MAINTENANCE_POLL_SECONDS"; +const INLINE_FLUSH_MIN_INLINED_BYTES_ENV: &str = + "ETL_DUCKLAKE_EXTERNAL_MAINTENANCE_INLINE_FLUSH_MIN_INLINED_BYTES"; +const REWRITE_DATA_FILES_MIN_ACTIVE_DATA_FILES_ENV: &str = + "ETL_DUCKLAKE_EXTERNAL_MAINTENANCE_REWRITE_DATA_FILES_MIN_ACTIVE_DATA_FILES"; +const REQUEST_COOLDOWN_SECONDS_ENV: &str = + "ETL_DUCKLAKE_EXTERNAL_MAINTENANCE_REQUEST_COOLDOWN_SECONDS"; +const KUBERNETES_API_TIMEOUT_SECONDS_ENV: &str = + "ETL_DUCKLAKE_EXTERNAL_MAINTENANCE_KUBERNETES_API_TIMEOUT_SECONDS"; +const DEFAULT_POLL_SECONDS: u64 = 5; +const DEFAULT_INLINE_FLUSH_MIN_INLINED_BYTES: u64 = 10_000_000; +const DEFAULT_REWRITE_DATA_FILES_MIN_ACTIVE_DATA_FILES: i64 = 40; +const DEFAULT_REQUEST_COOLDOWN_SECONDS: u64 = 300; +const DEFAULT_KUBERNETES_API_TIMEOUT_SECONDS: u64 = 10; +const OPERATION_INLINE_FLUSH: &str = "flush_inlined_data"; +const OPERATION_REWRITE_DATA_FILES: &str = "rewrite_data_files"; +const OPERATION_EXPIRE_SNAPSHOTS: &str = "expire_snapshots"; +const REASON_PENDING_INLINED_DATA_BYTES_THRESHOLD: &str = "pending_inlined_data_bytes_threshold"; +const REASON_ACTIVE_DATA_FILES_THRESHOLD: &str = "active_data_files_threshold"; +const REASON_SNAPSHOT_RETENTION_THRESHOLD: &str = "snapshot_retention_threshold"; + +#[derive(Clone, Copy, Debug, Default)] +pub(super) struct ExternalMaintenanceOperations { + pub(super) inline_flush: bool, + pub(super) rewrite_data_files: bool, + pub(super) expire_snapshots: bool, +} + +impl ExternalMaintenanceOperations { + fn covers(self, requested: Self) -> bool { + (!requested.inline_flush || self.inline_flush) + && (!requested.rewrite_data_files || self.rewrite_data_files) + && (!requested.expire_snapshots || self.expire_snapshots) + } +} + +struct PauseRequest { + run_id: String, + expires_at: DateTime, +} + +struct HeldPause { + run_id: String, + expires_at: DateTime, + quiesced_at: DateTime, + quiesced_reported: bool, + _pause: DuckLakeExternalMaintenancePause, +} + +#[derive(Clone)] +struct WatcherConfig { + name: String, + namespace: String, + poll_interval: Duration, + request_cooldown: Duration, + kubernetes_api_timeout: Duration, + inline_flush_min_inlined_bytes: u64, + rewrite_data_files_min_active_data_files: i64, +} + +struct OperationPolicy { + inline_flush_enabled: bool, + rewrite_data_files_enabled: bool, + expire_snapshots_enabled: bool, +} + +pub(super) async fn run_external_maintenance_watcher( + destination: DuckLakeDestination, +) -> Result<(), kube::Error> +where + S: StateStore + SchemaStore + Clone + Send + Sync + 'static, +{ + let Some(config) = WatcherConfig::from_env() else { + return Ok(()); + }; + let client = Client::try_default().await?; + let api: Api = + Api::namespaced_with(client, &config.namespace, &ducklake_maintenance_api_resource()); + let mut held_pause: Option = None; + + info!( + ducklake_maintenance = %config.name, + namespace = %config.namespace, + "ducklake external maintenance watcher started: ducklake_maintenance={}, namespace={}", + config.name, + config.namespace + ); + + loop { + match time::timeout(config.kubernetes_api_timeout, api.get(&config.name)).await { + Ok(resource) => { + let resource = match resource { + Ok(resource) => resource, + Err(kube::Error::Api(error)) if error.code == 404 => { + if let Some(held) = held_pause.take() { + info!( + ducklake_maintenance = %config.name, + run_id = %held.run_id, + "ducklake maintenance resource disappeared, resuming foreground mutations: \ + ducklake_maintenance={}, run_id={}", + config.name, + held.run_id + ); + release_held_pause(&config.name, held, "resource_deleted"); + } + time::sleep(config.poll_interval).await; + continue; + } + Err(error) => { + warn!( + error = %error, + ducklake_maintenance = %config.name, + timeout_ms = config.kubernetes_api_timeout.as_millis() as u64, + "failed to read ducklake maintenance resource: ducklake_maintenance={}, \ + timeout_ms={}, error={}", + config.name, + config.kubernetes_api_timeout.as_millis(), + error + ); + release_expired_pause_if_needed(&api, &config, &mut held_pause).await; + time::sleep(config.poll_interval).await; + continue; + } + }; + let active_pause = active_pause_request(&resource); + reconcile_pause(&api, &config, &destination, &mut held_pause, active_pause).await; + maybe_request_operations(&api, &config.name, &destination, &resource, &config) + .await; + } + Err(_) => { + warn!( + ducklake_maintenance = %config.name, + timeout_ms = config.kubernetes_api_timeout.as_millis() as u64, + "timed out reading ducklake maintenance resource: ducklake_maintenance={}, \ + timeout_ms={}", + config.name, + config.kubernetes_api_timeout.as_millis() + ); + release_expired_pause_if_needed(&api, &config, &mut held_pause).await; + } + } + + time::sleep(config.poll_interval).await; + } +} + +async fn reconcile_pause( + api: &Api, + config: &WatcherConfig, + destination: &DuckLakeDestination, + held_pause: &mut Option, + active_pause: Option, +) where + S: StateStore + SchemaStore + Clone + Send + Sync + 'static, +{ + let name = config.name.as_str(); + let Some(pause) = active_pause.filter(|pause| pause.expires_at > Utc::now()) else { + if let Some(held) = held_pause.take() { + info!( + ducklake_maintenance = %name, + run_id = %held.run_id, + "ducklake external maintenance pause cleared, resuming foreground mutations: \ + ducklake_maintenance={}, run_id={}", + name, + held.run_id + ); + release_held_pause(name, held, "cleared"); + patch_replicator_status( + api, + name, + "Running", + None, + None, + config.kubernetes_api_timeout, + ) + .await; + } + return; + }; + + if let Some(held) = held_pause.as_mut() + && held.run_id == pause.run_id + { + if !held.quiesced_reported + && patch_replicator_status( + api, + name, + "Quiesced", + Some(&held.run_id), + Some(held.quiesced_at), + config.kubernetes_api_timeout, + ) + .await + { + held.quiesced_reported = true; + } + return; + } + + if let Some(held) = held_pause.take() { + info!( + ducklake_maintenance = %name, + previous_run_id = %held.run_id, + next_run_id = %pause.run_id, + "ducklake external maintenance pause replaced, resuming previous run before pausing again: \ + ducklake_maintenance={}, previous_run_id={}, next_run_id={}", + name, + held.run_id, + pause.run_id + ); + release_held_pause(name, held, "replaced"); + patch_replicator_status(api, name, "Running", None, None, config.kubernetes_api_timeout) + .await; + } + + info!( + ducklake_maintenance = %name, + run_id = %pause.run_id, + expires_at = %pause.expires_at.to_rfc3339(), + "ducklake external maintenance pause requested, waiting for foreground mutations to drain: \ + ducklake_maintenance={}, run_id={}, expires_at={}", + name, + pause.run_id, + pause.expires_at.to_rfc3339() + ); + patch_replicator_status( + api, + name, + "Pausing", + Some(&pause.run_id), + None, + config.kubernetes_api_timeout, + ) + .await; + + let external_pause = destination.acquire_external_maintenance_pause().await; + if pause.expires_at <= Utc::now() { + info!( + ducklake_maintenance = %name, + run_id = %pause.run_id, + expires_at = %pause.expires_at.to_rfc3339(), + "ducklake external maintenance pause expired before quiescence, resuming foreground mutations: \ + ducklake_maintenance={}, run_id={}, expires_at={}", + name, + pause.run_id, + pause.expires_at.to_rfc3339() + ); + release_held_pause( + name, + HeldPause { + run_id: pause.run_id, + expires_at: pause.expires_at, + quiesced_at: Utc::now(), + quiesced_reported: false, + _pause: external_pause, + }, + "expired", + ); + patch_replicator_status(api, name, "Running", None, None, config.kubernetes_api_timeout) + .await; + return; + } + + let quiesced_at = Utc::now(); + + info!( + ducklake_maintenance = %name, + run_id = %pause.run_id, + quiesced_at = %quiesced_at.to_rfc3339(), + expires_at = %pause.expires_at.to_rfc3339(), + "ducklake external maintenance quiesced, foreground mutations are paused: \ + ducklake_maintenance={}, run_id={}, quiesced_at={}, expires_at={}", + name, + pause.run_id, + quiesced_at.to_rfc3339(), + pause.expires_at.to_rfc3339() + ); + let quiesced_reported = patch_replicator_status( + api, + name, + "Quiesced", + Some(&pause.run_id), + Some(quiesced_at), + config.kubernetes_api_timeout, + ) + .await; + if !quiesced_reported { + warn!( + ducklake_maintenance = %name, + run_id = %pause.run_id, + timeout_ms = config.kubernetes_api_timeout.as_millis() as u64, + "ducklake external maintenance quiesced status was not confirmed; keeping foreground \ + mutations paused until the status patch succeeds or the pause expires: \ + ducklake_maintenance={}, run_id={}, timeout_ms={}", + name, + pause.run_id, + config.kubernetes_api_timeout.as_millis() + ); + } + + *held_pause = Some(HeldPause { + run_id: pause.run_id, + expires_at: pause.expires_at, + quiesced_at, + quiesced_reported, + _pause: external_pause, + }); +} + +async fn release_expired_pause_if_needed( + api: &Api, + config: &WatcherConfig, + held_pause: &mut Option, +) { + if held_pause.as_ref().is_none_or(|pause| pause.expires_at > Utc::now()) { + return; + } + + let Some(expired) = held_pause.take() else { + return; + }; + warn!( + ducklake_maintenance = %config.name, + run_id = %expired.run_id, + "ducklake maintenance pause expired while Kubernetes API was unavailable: \ + ducklake_maintenance={}, run_id={}", + config.name, + expired.run_id + ); + release_held_pause(&config.name, expired, "expired"); + patch_replicator_status( + api, + &config.name, + "Running", + None, + None, + config.kubernetes_api_timeout, + ) + .await; +} + +fn release_held_pause(name: &str, held: HeldPause, outcome: &'static str) { + let held_ms = + Utc::now().signed_duration_since(held.quiesced_at).num_milliseconds().max(0) as u64; + record_external_maintenance_pause_duration(&held, outcome); + info!( + ducklake_maintenance = %name, + run_id = %held.run_id, + outcome, + held_ms, + "ducklake external maintenance pause guard released: ducklake_maintenance={}, run_id={}, \ + outcome={}, held_ms={}", + name, + held.run_id, + outcome, + held_ms + ); +} + +fn record_external_maintenance_pause_duration(held: &HeldPause, outcome: &'static str) { + let duration_seconds = + Utc::now().signed_duration_since(held.quiesced_at).num_milliseconds().max(0) as f64 + / 1_000.0; + histogram!( + ETL_DUCKLAKE_EXTERNAL_MAINTENANCE_PAUSE_DURATION_SECONDS, + MAINTENANCE_OUTCOME_LABEL => outcome, + ) + .record(duration_seconds); +} + +async fn maybe_request_operations( + api: &Api, + name: &str, + destination: &DuckLakeDestination, + resource: &DynamicObject, + config: &WatcherConfig, +) where + S: StateStore + SchemaStore + Clone + Send + Sync + 'static, +{ + if active_run_exists(resource) { + return; + } + if completed_run_in_cooldown(resource, config.request_cooldown) { + return; + } + + let policy = operation_policy(resource); + let requested = match destination + .sample_external_maintenance_operations( + config.inline_flush_min_inlined_bytes, + config.rewrite_data_files_min_active_data_files, + ) + .await + { + Ok(mut operations) => { + operations.inline_flush &= policy.inline_flush_enabled; + operations.rewrite_data_files &= policy.rewrite_data_files_enabled; + operations.expire_snapshots &= policy.expire_snapshots_enabled; + operations + } + Err(error) => { + warn!( + error = ?error, + ducklake_maintenance = %name, + "failed to sample ducklake external maintenance operations" + ); + return; + } + }; + + if !requested.inline_flush && !requested.rewrite_data_files && !requested.expire_snapshots { + debug!( + ducklake_maintenance = %name, + inline_flush = requested.inline_flush, + rewrite_data_files = requested.rewrite_data_files, + expire_snapshots = requested.expire_snapshots, + inline_flush_min_inlined_bytes = config.inline_flush_min_inlined_bytes, + rewrite_data_files_min_active_data_files = + config.rewrite_data_files_min_active_data_files, + "ducklake external maintenance sampled no requested operations" + ); + return; + } + + let already_requested = requested_operations(resource); + if already_requested.covers(requested) { + debug!( + ducklake_maintenance = %name, + inline_flush = requested.inline_flush, + rewrite_data_files = requested.rewrite_data_files, + expire_snapshots = requested.expire_snapshots, + "ducklake external maintenance request already exists" + ); + return; + } + + info!( + ducklake_maintenance = %name, + inline_flush = requested.inline_flush, + rewrite_data_files = requested.rewrite_data_files, + expire_snapshots = requested.expire_snapshots, + inline_flush_min_inlined_bytes = config.inline_flush_min_inlined_bytes, + rewrite_data_files_min_active_data_files = config.rewrite_data_files_min_active_data_files, + "ducklake external maintenance requesting operations" + ); + if patch_operation_requests(api, name, requested, config).await { + record_external_maintenance_triggers(requested, already_requested); + } +} + +async fn patch_replicator_status( + api: &Api, + name: &str, + state: &str, + observed_run_id: Option<&str>, + quiesced_at: Option>, + timeout: Duration, +) -> bool { + let patch = json!({ + "status": { + "replicator": { + "state": state, + "observedRunId": observed_run_id, + "quiescedAt": quiesced_at.map(|time| time.to_rfc3339()), + } + } + }); + let params = PatchParams::default(); + + match time::timeout(timeout, api.patch_status(name, ¶ms, &Patch::Merge(&patch))).await { + Ok(Ok(_)) => true, + Ok(Err(error)) => { + warn!( + error = %error, + ducklake_maintenance = %name, + state, + "failed to patch ducklake maintenance replicator status: \ + ducklake_maintenance={}, state={}, error={}", + name, + state, + error + ); + false + } + Err(_) => { + warn!( + ducklake_maintenance = %name, + state, + timeout_ms = timeout.as_millis() as u64, + "timed out patching ducklake maintenance replicator status: \ + ducklake_maintenance={}, state={}, timeout_ms={}", + name, + state, + timeout.as_millis() + ); + false + } + } +} + +async fn patch_operation_requests( + api: &Api, + name: &str, + requested: ExternalMaintenanceOperations, + config: &WatcherConfig, +) -> bool { + let patch = json!({ + "status": { + "operationRequests": { + "inlineFlush": requested.inline_flush, + "rewriteDataFiles": requested.rewrite_data_files, + "expireSnapshots": requested.expire_snapshots, + "inlineFlushMinInlinedBytes": config.inline_flush_min_inlined_bytes, + "rewriteDataFilesMinActiveDataFiles": config.rewrite_data_files_min_active_data_files, + "requestedAt": Utc::now().to_rfc3339(), + } + } + }); + let params = PatchParams::default(); + + match time::timeout( + config.kubernetes_api_timeout, + api.patch_status(name, ¶ms, &Patch::Merge(&patch)), + ) + .await + { + Ok(Ok(_)) => true, + Ok(Err(error)) => { + warn!( + error = %error, + ducklake_maintenance = %name, + "failed to patch ducklake maintenance operation request: \ + ducklake_maintenance={}, error={}", + name, + error + ); + false + } + Err(_) => { + warn!( + ducklake_maintenance = %name, + timeout_ms = config.kubernetes_api_timeout.as_millis() as u64, + "timed out patching ducklake maintenance operation request: \ + ducklake_maintenance={}, timeout_ms={}", + name, + config.kubernetes_api_timeout.as_millis() + ); + false + } + } +} + +fn record_external_maintenance_triggers( + requested: ExternalMaintenanceOperations, + already_requested: ExternalMaintenanceOperations, +) { + if requested.inline_flush && !already_requested.inline_flush { + counter!( + ETL_DUCKLAKE_EXTERNAL_MAINTENANCE_TRIGGERED_TOTAL, + MAINTENANCE_OPERATION_LABEL => OPERATION_INLINE_FLUSH, + MAINTENANCE_REASON_LABEL => REASON_PENDING_INLINED_DATA_BYTES_THRESHOLD, + ) + .increment(1); + } + + if requested.rewrite_data_files && !already_requested.rewrite_data_files { + counter!( + ETL_DUCKLAKE_EXTERNAL_MAINTENANCE_TRIGGERED_TOTAL, + MAINTENANCE_OPERATION_LABEL => OPERATION_REWRITE_DATA_FILES, + MAINTENANCE_REASON_LABEL => REASON_ACTIVE_DATA_FILES_THRESHOLD, + ) + .increment(1); + } + + if requested.expire_snapshots && !already_requested.expire_snapshots { + counter!( + ETL_DUCKLAKE_EXTERNAL_MAINTENANCE_TRIGGERED_TOTAL, + MAINTENANCE_OPERATION_LABEL => OPERATION_EXPIRE_SNAPSHOTS, + MAINTENANCE_REASON_LABEL => REASON_SNAPSHOT_RETENTION_THRESHOLD, + ) + .increment(1); + } +} + +fn requested_operations(resource: &DynamicObject) -> ExternalMaintenanceOperations { + let Some(requests) = + resource.data.get("status").and_then(|status| status.get("operationRequests")) + else { + return ExternalMaintenanceOperations::default(); + }; + + ExternalMaintenanceOperations { + inline_flush: requests + .get("inlineFlush") + .and_then(serde_json::Value::as_bool) + .unwrap_or(false), + rewrite_data_files: requests + .get("rewriteDataFiles") + .and_then(serde_json::Value::as_bool) + .unwrap_or(false), + expire_snapshots: requests + .get("expireSnapshots") + .and_then(serde_json::Value::as_bool) + .unwrap_or(false), + } +} + +fn active_pause_request(resource: &DynamicObject) -> Option { + let pause = resource.data.get("status")?.get("pauseRequest")?; + if pause.is_null() { + return None; + } + + let run_id = pause.get("runId")?.as_str()?.to_owned(); + let expires_at = + DateTime::parse_from_rfc3339(pause.get("expiresAt")?.as_str()?).ok()?.with_timezone(&Utc); + + Some(PauseRequest { run_id, expires_at }) +} + +fn active_run_exists(resource: &DynamicObject) -> bool { + resource + .data + .get("status") + .and_then(|status| status.get("activeRun")) + .is_some_and(|active_run| !active_run.is_null()) +} + +fn completed_run_in_cooldown(resource: &DynamicObject, cooldown: Duration) -> bool { + if cooldown.is_zero() { + return false; + } + + let Some(completed_at) = resource + .data + .get("status") + .and_then(|status| status.get("lastCompletedRun")) + .and_then(|run| run.get("completedAt")) + .and_then(serde_json::Value::as_str) + .and_then(|value| DateTime::parse_from_rfc3339(value).ok()) + .map(|time| time.with_timezone(&Utc)) + else { + return false; + }; + + let elapsed = Utc::now().signed_duration_since(completed_at); + elapsed.to_std().is_ok_and(|elapsed| elapsed < cooldown) +} + +fn operation_policy(resource: &DynamicObject) -> OperationPolicy { + let operations = resource.data.get("spec").and_then(|spec| spec.get("operations")); + let inline_flush = operations.and_then(|ops| ops.get("inlineFlush")); + let rewrite_data_files = operations.and_then(|ops| ops.get("rewriteDataFiles")); + let expire_snapshots = operations.and_then(|ops| ops.get("expireSnapshots")); + + OperationPolicy { + inline_flush_enabled: inline_flush + .and_then(|value| value.get("enabled")) + .and_then(serde_json::Value::as_bool) + .unwrap_or(true), + rewrite_data_files_enabled: rewrite_data_files + .and_then(|value| value.get("enabled")) + .and_then(serde_json::Value::as_bool) + .unwrap_or(true), + expire_snapshots_enabled: expire_snapshots + .and_then(|value| value.get("enabled")) + .and_then(serde_json::Value::as_bool) + .unwrap_or(false), + } +} + +fn ducklake_maintenance_api_resource() -> ApiResource { + let gvk = GroupVersionKind::gvk("etl.supabase.com", "v1alpha1", "DuckLakeMaintenance"); + ApiResource::from_gvk(&gvk) +} + +impl WatcherConfig { + fn from_env() -> Option { + let name = env::var(CR_NAME_ENV).ok().filter(|value| !value.is_empty())?; + let namespace = env::var(CR_NAMESPACE_ENV).ok().filter(|value| !value.is_empty())?; + let poll_seconds = env::var(POLL_SECONDS_ENV) + .ok() + .and_then(|value| value.parse::().ok()) + .filter(|seconds| *seconds > 0) + .unwrap_or(DEFAULT_POLL_SECONDS); + let inline_flush_min_inlined_bytes = env::var(INLINE_FLUSH_MIN_INLINED_BYTES_ENV) + .ok() + .and_then(|value| value.parse::().ok()) + .unwrap_or(DEFAULT_INLINE_FLUSH_MIN_INLINED_BYTES); + let rewrite_data_files_min_active_data_files = + env::var(REWRITE_DATA_FILES_MIN_ACTIVE_DATA_FILES_ENV) + .ok() + .and_then(|value| value.parse::().ok()) + .unwrap_or(DEFAULT_REWRITE_DATA_FILES_MIN_ACTIVE_DATA_FILES); + let request_cooldown = env::var(REQUEST_COOLDOWN_SECONDS_ENV) + .ok() + .and_then(|value| value.parse::().ok()) + .unwrap_or(DEFAULT_REQUEST_COOLDOWN_SECONDS); + let kubernetes_api_timeout = env::var(KUBERNETES_API_TIMEOUT_SECONDS_ENV) + .ok() + .and_then(|value| value.parse::().ok()) + .filter(|seconds| *seconds > 0) + .unwrap_or(DEFAULT_KUBERNETES_API_TIMEOUT_SECONDS); + + Some(Self { + name, + namespace, + poll_interval: Duration::from_secs(poll_seconds), + request_cooldown: Duration::from_secs(request_cooldown), + kubernetes_api_timeout: Duration::from_secs(kubernetes_api_timeout), + inline_flush_min_inlined_bytes, + rewrite_data_files_min_active_data_files, + }) + } +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + + #[test] + fn operation_policy_defaults_expire_snapshots_disabled() { + let resource: DynamicObject = serde_json::from_value(json!({ + "apiVersion": "etl.supabase.com/v1alpha1", + "kind": "DuckLakeMaintenance", + "metadata": { + "name": "pipeline-maintenance" + }, + "spec": { + "operations": { + "inlineFlush": {}, + "rewriteDataFiles": {} + } + } + })) + .unwrap(); + + let policy = operation_policy(&resource); + + assert!(policy.inline_flush_enabled); + assert!(policy.rewrite_data_files_enabled); + assert!(!policy.expire_snapshots_enabled); + } +} diff --git a/crates/etl-destinations/src/ducklake/maintenance.rs b/crates/etl-destinations/src/ducklake/maintenance.rs deleted file mode 100644 index 67ce0b9e6..000000000 --- a/crates/etl-destinations/src/ducklake/maintenance.rs +++ /dev/null @@ -1,2265 +0,0 @@ -use std::{ - collections::HashMap, - sync::{ - Arc, - atomic::{AtomicBool, Ordering as AtomicOrdering}, - }, - time::Duration, -}; - -use etl::{ - error::{ErrorKind, EtlError, EtlResult}, - etl_error, -}; -use metrics::{counter, gauge, histogram}; -use parking_lot::Mutex; -use pg_escape::quote_literal; -use tokio::{ - sync::{OwnedSemaphorePermit, RwLock, Semaphore, mpsc, watch}, - task::JoinHandle, - time::{Instant, MissedTickBehavior}, -}; -use tracing::{debug, info, warn}; - -use crate::ducklake::{ - DuckLakeTableName, LAKE_CATALOG, - client::{ - DuckDbBlockingOperationKind, DuckLakeConnectionManager, LazyDuckLakePool, - format_query_error_detail, run_duckdb_blocking, - }, - inline_size::{DuckLakePendingInlineDataSizes, DuckLakePendingInlineSizeSampler}, - metrics::{ - DuckLakeTableStorageMetrics, ETL_DUCKLAKE_INLINE_FLUSH_DURATION_SECONDS, - ETL_DUCKLAKE_INLINE_FLUSH_ROWS, ETL_DUCKLAKE_MAINTENANCE_DURATION_SECONDS, - ETL_DUCKLAKE_MAINTENANCE_IN_PROGRESS, ETL_DUCKLAKE_MAINTENANCE_SKIPPED_TOTAL, - ETL_DUCKLAKE_MAINTENANCE_STARTED_TOTAL, MAINTENANCE_OPERATION_LABEL, - MAINTENANCE_OUTCOME_LABEL, MAINTENANCE_REASON_LABEL, MAINTENANCE_TASK_LABEL, RESULT_LABEL, - }, -}; - -/// Dedicated pool size for background DuckLake maintenance work. -const MAINTENANCE_POOL_SIZE: u32 = 1; -/// Poll interval for checking per-table inline flush thresholds. -const MAINTENANCE_FLUSH_POLL_INTERVAL: Duration = Duration::from_mins(2); -/// Fixed cadence for expiring old DuckLake snapshots. -const MAINTENANCE_EXPIRE_SNAPSHOTS_INTERVAL: Duration = Duration::from_secs(5 * 60 * 60); -/// Fixed cadence for cleaning up old DuckLake files. -const MAINTENANCE_CLEANUP_OLD_FILES_INTERVAL: Duration = Duration::from_secs(6 * 60 * 60); -/// Pending inlined bytes threshold that triggers a background inline flush. We -/// multiply using `PARQUET_COMPRESSION_RATIO_ESTIMATE` to make sure we won't -/// get too small data files -const MAINTENANCE_PENDING_INLINED_DATA_BYTES_THRESHOLD: u64 = 10_000_000; -/// Minimum idle window before targeted table maintenance runs, to not have -/// maintenances ran too frequently. -const MAINTENANCE_TABLE_COMPACTION_IDLE_THRESHOLD: Duration = Duration::from_secs(90); -/// Minimum delay between targeted maintenance runs for the same table. -const MAINTENANCE_TABLE_COMPACTION_INTERVAL: Duration = Duration::from_secs(5 * 60); -/// Keeps the legacy targeted rewrite path compiled without scheduling it. -const ENABLE_TARGETED_TABLE_MAINTENANCE: bool = false; -/// Minimum active delete-file count before idle rewrite is worth attempting. -const MAINTENANCE_IDLE_REWRITE_DELETE_FILES_THRESHOLD: i64 = 32; -/// Deleted-row ratio that makes idle rewrite worthwhile. -const MAINTENANCE_IDLE_REWRITE_DELETED_ROW_RATIO_THRESHOLD: f64 = 0.10; -/// Active delete-file count that warrants emergency rewrite. -const MAINTENANCE_EMERGENCY_REWRITE_DELETE_FILES_THRESHOLD: i64 = 128; -/// Deleted-row ratio that warrants emergency rewrite. -const MAINTENANCE_EMERGENCY_REWRITE_DELETED_ROW_RATIO_THRESHOLD: f64 = 0.25; -/// Timeout for sending a notification to the maintenance worker. -pub(super) const NOTIFICATION_SEND_TIMEOUT: Duration = Duration::from_secs(5); - -const MAINTENANCE_TASK_FLUSH: &str = "flush"; -const MAINTENANCE_TASK_CATALOG_MAINTENANCE: &str = "catalog_maintenance"; -const MAINTENANCE_TASK_TARGETED_MAINTENANCE: &str = "targeted_maintenance"; - -#[cfg(test)] -static FAIL_REWRITE_SINGLE_OUTPUT_FILE_ONCE_FOR_TESTS: AtomicBool = AtomicBool::new(false); - -/// Concrete DuckLake maintenance operations emitted in metrics. -#[derive(Clone, Copy, Debug, Eq, PartialEq)] -enum MaintenanceOperation { - FlushInlinedData, - ExpireSnapshots, - CleanupOldFiles, - RewriteDataFiles, -} - -impl MaintenanceOperation { - /// Returns the stable metric label value for this operation. - fn as_str(self) -> &'static str { - match self { - Self::FlushInlinedData => "flush_inlined_data", - Self::ExpireSnapshots => "expire_snapshots", - Self::CleanupOldFiles => "cleanup_old_files", - Self::RewriteDataFiles => "rewrite_data_files", - } - } -} - -/// Primary reasons that schedule one background maintenance decision. -#[derive(Clone, Copy, Debug, Eq, PartialEq)] -#[allow(clippy::enum_variant_names)] -enum MaintenanceReason { - PendingInlinedDataBytesThreshold, - SnapshotRetentionThreshold, - CleanupIntervalElapsed, - IdleRewriteMetricsThreshold, - EmergencyRewriteMetricsThreshold, -} - -impl MaintenanceReason { - /// Returns the stable metric label value for this reason. - fn as_str(self) -> &'static str { - match self { - Self::PendingInlinedDataBytesThreshold => "pending_inlined_data_bytes_threshold", - Self::SnapshotRetentionThreshold => "snapshot_retention_threshold", - Self::CleanupIntervalElapsed => "cleanup_interval_elapsed", - Self::IdleRewriteMetricsThreshold => "idle_rewrite_metrics_threshold", - Self::EmergencyRewriteMetricsThreshold => "emergency_rewrite_metrics_threshold", - } - } -} - -/// Outcome for one maintenance operation attempt. -#[derive(Clone, Copy, Debug, Eq, PartialEq)] -enum MaintenanceOutcome { - Applied, - Noop, - SkippedBusy, - Failed, -} - -impl MaintenanceOutcome { - /// Returns the stable metric label value for this outcome. - fn as_str(self) -> &'static str { - match self { - Self::Applied => "applied", - Self::Noop => "noop", - Self::SkippedBusy => "skipped_busy", - Self::Failed => "failed", - } - } - - /// Returns whether the maintenance cycle completed successfully. - fn is_completed(self) -> bool { - matches!(self, Self::Applied | Self::Noop) - } -} - -impl From for MaintenanceOutcome { - fn from(value: u64) -> Self { - if value > 0 { Self::Applied } else { Self::Noop } - } -} - -/// Per-table write activity sent to the background maintenance worker. -#[derive(Clone, Debug, Default)] -pub(super) struct TableWriteActivity { - pub(super) table_name: DuckLakeTableName, -} - -/// Table-health metrics sent from the background sampler to maintenance. -#[derive(Clone, Debug)] -pub(super) struct TableMetricsSample { - pub(super) table_name: DuckLakeTableName, - pub(super) sampled_at: Instant, - pub(super) metrics: DuckLakeTableStorageMetrics, -} - -/// Notifications consumed by the background DuckLake maintenance worker. -#[derive(Clone, Debug)] -pub(super) enum TableMaintenanceNotification { - WriteActivity(TableWriteActivity), - TableMetricsSample(TableMetricsSample), - FlushCompleted(TableFlushCompletion), -} - -impl TableMaintenanceNotification { - /// Returns the table name carried by this notification. - fn table_name(&self) -> &str { - match self { - Self::WriteActivity(activity) => &activity.table_name, - Self::TableMetricsSample(sample) => &sample.table_name, - Self::FlushCompleted(completion) => &completion.table_name, - } - } -} - -/// Completion notification for one requested inline flush. -#[derive(Clone, Debug)] -pub(super) struct TableFlushCompletion { - pub(super) table_name: DuckLakeTableName, - pub(super) completed_at: Instant, -} - -/// Per-table inline flushes that should run at the next safe ingest pause. -#[derive(Debug, Default)] -pub(super) struct PendingInlineFlushRequests { - requests: Mutex>, -} - -impl PendingInlineFlushRequests { - /// Records or updates one requested inline flush. - fn request(&self, table_name: DuckLakeTableName, reason: MaintenanceReason) { - self.requests.lock().insert(table_name, reason); - } - - /// Drains the currently requested inline flushes. - fn take_all(&self) -> Vec<(DuckLakeTableName, MaintenanceReason)> { - self.requests.lock().drain().collect() - } - - /// Restores requested inline flushes after a skipped or failed attempt. - fn restore(&self, requests: impl IntoIterator) { - self.requests.lock().extend(requests); - } -} - -/// Trigger scope for one targeted maintenance selection. -#[derive(Clone, Copy, Debug, Eq, PartialEq)] -enum TargetedMaintenanceScope { - /// When the table has been idle long enough - Idle, - /// When ducklake state is not healthy (files fragmented, ...) - Emergency, -} - -/// Selected targeted-maintenance operations for one table. -#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] -struct TargetedMaintenancePlan { - rewrite_reason: Option, -} - -impl TargetedMaintenancePlan { - /// Returns whether this plan selected any maintenance work. - fn has_work(self) -> bool { - self.rewrite_reason.is_some() - } -} - -/// Static configuration for periodic catalog-level maintenance. -#[derive(Debug, Clone)] -struct CatalogMaintenanceConfig { - expire_snapshots_older_than: Arc, - expire_snapshots_interval: Duration, - cleanup_old_files_interval: Duration, -} - -impl CatalogMaintenanceConfig { - /// Builds the fixed catalog-maintenance configuration. - fn new(expire_snapshots_older_than: Arc) -> Self { - Self { - expire_snapshots_older_than, - expire_snapshots_interval: MAINTENANCE_EXPIRE_SNAPSHOTS_INTERVAL, - cleanup_old_files_interval: MAINTENANCE_CLEANUP_OLD_FILES_INTERVAL, - } - } -} - -/// Periodic catalog-maintenance state tracked by the background worker. -#[derive(Debug)] -struct CatalogMaintenanceState { - last_expire_snapshots_completed_at: Option, - last_cleanup_old_files_completed_at: Option, -} - -impl CatalogMaintenanceState { - /// Builds catalog-maintenance state seeded to "just ran now". - /// - /// Fresh destinations should not immediately run snapshot expiration or - /// old-file cleanup on startup before the configured cadence elapses. - fn new(now: Instant) -> Self { - Self { - last_expire_snapshots_completed_at: Some(now), - last_cleanup_old_files_completed_at: Some(now), - } - } - - /// Returns whether snapshot expiration is due now. - fn expire_snapshots_due(&self, now: Instant, interval: Duration) -> bool { - match self.last_expire_snapshots_completed_at { - Some(last_completed_at) => now.saturating_duration_since(last_completed_at) >= interval, - None => true, - } - } - - /// Returns whether old-file cleanup is due now. - fn cleanup_old_files_due(&self, now: Instant, interval: Duration) -> bool { - match self.last_cleanup_old_files_completed_at { - Some(last_completed_at) => now.saturating_duration_since(last_completed_at) >= interval, - None => true, - } - } - - /// Records one successful snapshot-expiration run. - fn complete_expire_snapshots(&mut self, now: Instant) { - self.last_expire_snapshots_completed_at = Some(now); - } - - /// Records one successful cleanup-old-files run. - fn complete_cleanup_old_files(&mut self, now: Instant) { - self.last_cleanup_old_files_completed_at = Some(now); - } -} - -/// Coalesced maintenance state for one DuckLake table. -#[derive(Debug, Default)] -struct TableMaintenanceState { - dirty_since_compaction: bool, - last_write_at: Option, - latest_pending_inline_data_sizes: Option, - latest_pending_inline_data_sampled_at: Option, - last_targeted_maintenance_at: Option, - latest_storage_metrics: Option, - latest_storage_metrics_sampled_at: Option, - last_emergency_assessment_at: Option, -} - -impl TableMaintenanceState { - /// Aggregates one write notification into the existing table state. - fn record_write_activity(&mut self, now: Instant) { - self.dirty_since_compaction = true; - self.last_write_at = Some(now); - } - - /// Records one sampled inlined-data snapshot for this table. - fn record_pending_inline_data_sizes( - &mut self, - sampled_at: Instant, - sizes: DuckLakePendingInlineDataSizes, - ) { - self.latest_pending_inline_data_sizes = Some(sizes); - self.latest_pending_inline_data_sampled_at = Some(sampled_at); - } - - /// Records one sampled metrics snapshot for this table. - fn record_metrics_sample(&mut self, sample: TableMetricsSample) { - self.latest_storage_metrics = Some(sample.metrics); - self.latest_storage_metrics_sampled_at = Some(sample.sampled_at); - } - - /// Returns whether the current dirty period still needs an inline-size - /// sample. - fn needs_pending_inline_data_sizes_sample(&self) -> bool { - self.dirty_since_compaction - && self.last_write_at.is_some() - && self.current_pending_inline_data_sizes().is_none() - } - - /// Returns the latest inlined-data sample covering the current dirty - /// period. - fn current_pending_inline_data_sizes(&self) -> Option { - let sizes = self.latest_pending_inline_data_sizes?; - let sampled_at = self.latest_pending_inline_data_sampled_at?; - - if let Some(last_write_at) = self.last_write_at - && sampled_at < last_write_at - { - return None; - } - - Some(sizes) - } - - /// Returns the primary reason that pending inlined work should be flushed. - fn flush_reason(&self) -> Option { - if let Some(sizes) = self.current_pending_inline_data_sizes() { - return (sizes.inlined_bytes >= MAINTENANCE_PENDING_INLINED_DATA_BYTES_THRESHOLD) - .then_some(MaintenanceReason::PendingInlinedDataBytesThreshold); - } - - None - } - - /// Clears pending flush counters after a successful flush/materialization. - fn clear_pending_flush(&mut self, now: Instant) { - self.latest_pending_inline_data_sizes = Some(DuckLakePendingInlineDataSizes::default()); - self.latest_pending_inline_data_sampled_at = Some(now); - } - - /// Returns the latest metrics sample covering the current dirty period. - fn current_storage_metrics(&self) -> Option<(&DuckLakeTableStorageMetrics, Instant)> { - let metrics = self.latest_storage_metrics.as_ref()?; - let sampled_at = self.latest_storage_metrics_sampled_at?; - - if let Some(last_write_at) = self.last_write_at - && sampled_at < last_write_at - { - return None; - } - - Some((metrics, sampled_at)) - } - - /// Returns the plan for idle targeted maintenance, if it is due. - fn idle_targeted_maintenance_plan(&self, now: Instant) -> Option { - if !self.dirty_since_compaction { - return None; - } - - let last_write_at = self.last_write_at?; - let (metrics, _) = self.current_storage_metrics()?; - let idle = now.saturating_duration_since(last_write_at); - let enough_idle = idle >= MAINTENANCE_TABLE_COMPACTION_IDLE_THRESHOLD; - let enough_gap = match self.last_targeted_maintenance_at { - Some(last) => { - now.saturating_duration_since(last) >= MAINTENANCE_TABLE_COMPACTION_INTERVAL - } - None => true, - }; - - if enough_idle && enough_gap { - Some(targeted_maintenance_plan(metrics, TargetedMaintenanceScope::Idle)) - } else { - None - } - } - - /// Returns the plan for emergency targeted maintenance from a fresh sample. - fn emergency_targeted_maintenance_plan( - &self, - now: Instant, - ) -> Option<(TargetedMaintenancePlan, Instant)> { - if !self.dirty_since_compaction { - return None; - } - - let (metrics, sampled_at) = self.current_storage_metrics()?; - let enough_gap = match self.last_targeted_maintenance_at { - Some(last) => { - now.saturating_duration_since(last) >= MAINTENANCE_TABLE_COMPACTION_INTERVAL - } - None => true, - }; - let unseen_sample = match self.last_emergency_assessment_at { - Some(last_assessment_at) => sampled_at > last_assessment_at, - None => true, - }; - - if enough_gap && unseen_sample { - Some(( - targeted_maintenance_plan(metrics, TargetedMaintenanceScope::Emergency), - sampled_at, - )) - } else { - None - } - } - - /// Marks one completed idle maintenance assessment. - fn complete_idle_targeted_maintenance(&mut self, now: Instant) { - self.dirty_since_compaction = false; - self.last_targeted_maintenance_at = Some(now); - } - - /// Marks one completed targeted maintenance run. - fn complete_targeted_maintenance(&mut self, now: Instant) { - self.complete_idle_targeted_maintenance(now); - } -} - -/// Shared state for the background DuckLake maintenance worker. -pub(super) struct DuckLakeMaintenanceWorker { - pub(super) notification_tx: mpsc::Sender, - pub(super) shutdown_tx: watch::Sender<()>, - pub(super) handle: Mutex>>, -} - -/// Records one DuckLake background maintenance operation start. -fn record_ducklake_maintenance_started( - task: &'static str, - operation: MaintenanceOperation, - reason: MaintenanceReason, -) { - counter!( - ETL_DUCKLAKE_MAINTENANCE_STARTED_TOTAL, - MAINTENANCE_TASK_LABEL => task, - MAINTENANCE_OPERATION_LABEL => operation.as_str(), - MAINTENANCE_REASON_LABEL => reason.as_str(), - ) - .increment(1); -} - -/// Increments one DuckLake background maintenance operation in-progress sample. -fn increment_ducklake_maintenance_in_progress( - task: &'static str, - operation: MaintenanceOperation, - reason: MaintenanceReason, -) { - gauge!( - ETL_DUCKLAKE_MAINTENANCE_IN_PROGRESS, - MAINTENANCE_TASK_LABEL => task, - MAINTENANCE_OPERATION_LABEL => operation.as_str(), - MAINTENANCE_REASON_LABEL => reason.as_str(), - ) - .increment(1.0); -} - -/// Decrements one DuckLake background maintenance operation in-progress sample. -fn decrement_ducklake_maintenance_in_progress( - task: &'static str, - operation: MaintenanceOperation, - reason: MaintenanceReason, -) { - gauge!( - ETL_DUCKLAKE_MAINTENANCE_IN_PROGRESS, - MAINTENANCE_TASK_LABEL => task, - MAINTENANCE_OPERATION_LABEL => operation.as_str(), - MAINTENANCE_REASON_LABEL => reason.as_str(), - ) - .decrement(1.0); -} - -/// Keeps the maintenance in-progress gauge balanced for one operation attempt. -#[must_use = "the returned guard tracks an in-progress maintenance metric until dropped"] -struct DuckLakeMaintenanceInProgressGuard { - task: &'static str, - operation: MaintenanceOperation, - reason: MaintenanceReason, -} - -impl DuckLakeMaintenanceInProgressGuard { - /// Starts one maintenance attempt and tracks it until the guard is dropped. - fn start( - task: &'static str, - operation: MaintenanceOperation, - reason: MaintenanceReason, - ) -> Self { - record_ducklake_maintenance_started(task, operation, reason); - increment_ducklake_maintenance_in_progress(task, operation, reason); - Self { task, operation, reason } - } -} - -impl Drop for DuckLakeMaintenanceInProgressGuard { - fn drop(&mut self) { - decrement_ducklake_maintenance_in_progress(self.task, self.operation, self.reason); - } -} - -/// Records one DuckLake background maintenance operation duration sample. -fn record_ducklake_maintenance_duration( - task: &'static str, - operation: MaintenanceOperation, - reason: MaintenanceReason, - outcome: MaintenanceOutcome, - duration_seconds: f64, -) { - debug_assert_ne!(outcome, MaintenanceOutcome::SkippedBusy); - histogram!( - ETL_DUCKLAKE_MAINTENANCE_DURATION_SECONDS, - MAINTENANCE_TASK_LABEL => task, - MAINTENANCE_OPERATION_LABEL => operation.as_str(), - MAINTENANCE_REASON_LABEL => reason.as_str(), - MAINTENANCE_OUTCOME_LABEL => outcome.as_str(), - ) - .record(duration_seconds); -} - -/// Records one DuckLake background maintenance operation skip. -fn record_ducklake_maintenance_skipped( - task: &'static str, - operation: MaintenanceOperation, - reason: MaintenanceReason, -) { - counter!( - ETL_DUCKLAKE_MAINTENANCE_SKIPPED_TOTAL, - MAINTENANCE_TASK_LABEL => task, - MAINTENANCE_OPERATION_LABEL => operation.as_str(), - MAINTENANCE_REASON_LABEL => reason.as_str(), - ) - .increment(1); -} - -/// Records all selected targeted-maintenance operations as skipped because the -/// table is busy. -fn record_skipped_targeted_maintenance(plan: TargetedMaintenancePlan) { - if let Some(reason) = plan.rewrite_reason { - record_ducklake_maintenance_skipped( - MAINTENANCE_TASK_TARGETED_MAINTENANCE, - MaintenanceOperation::RewriteDataFiles, - reason, - ); - } -} - -/// Records skipped catalog-maintenance operations when the worker cannot enter -/// the exclusive safe point yet. -fn record_skipped_catalog_maintenance(expire_snapshots_due: bool, cleanup_old_files_due: bool) { - if expire_snapshots_due { - record_ducklake_maintenance_skipped( - MAINTENANCE_TASK_CATALOG_MAINTENANCE, - MaintenanceOperation::ExpireSnapshots, - MaintenanceReason::SnapshotRetentionThreshold, - ); - } - if cleanup_old_files_due { - record_ducklake_maintenance_skipped( - MAINTENANCE_TASK_CATALOG_MAINTENANCE, - MaintenanceOperation::CleanupOldFiles, - MaintenanceReason::CleanupIntervalElapsed, - ); - } -} - -/// Returns whether this maintenance failure matches a known DuckLake compaction -/// bug. -fn is_known_ducklake_compaction_single_output_file_error(error: &EtlError) -> bool { - const KNOWN_ERROR: &str = "INTERNAL Error: DuckLakeCompaction - expected a single output file"; - - if error.detail().is_some_and(|detail| detail.contains(KNOWN_ERROR)) { - return true; - } - - let mut source = std::error::Error::source(error); - while let Some(error) = source { - if error.to_string().contains(KNOWN_ERROR) { - return true; - } - - source = error.source(); - } - - false -} - -/// Returns the failing maintenance operation and reason for one known -/// compaction bug. -fn known_ducklake_compaction_error_context( - plan: TargetedMaintenancePlan, - error: &EtlError, -) -> Option<(MaintenanceOperation, MaintenanceReason)> { - if !is_known_ducklake_compaction_single_output_file_error(error) { - return None; - } - - match error.description() { - Some("DuckLake rewrite data files failed") => { - plan.rewrite_reason.map(|reason| (MaintenanceOperation::RewriteDataFiles, reason)) - } - _ => None, - } -} - -/// Logs and suppresses one known DuckLake compaction internal error. -/// -/// Connection recycling is handled generically in [`run_duckdb_blocking`], -/// which marks any failing DuckDB connection as broken before it returns the -/// error to this layer. This helper stays intentionally narrow: it only decides -/// which maintenance-only failures are safe to downgrade after the pool has -/// already recycled the invalidated connection. -fn suppress_known_ducklake_compaction_error( - table_name: &str, - plan: TargetedMaintenancePlan, - error: &EtlError, -) -> bool { - let Some((operation, reason)) = known_ducklake_compaction_error_context(plan, error) else { - return false; - }; - - warn!( - table = %table_name, - operation = operation.as_str(), - reason = reason.as_str(), - error = %error, - "ducklake targeted maintenance skipped after known duckdb internal error" - ); - true -} - -/// Returns the targeted-maintenance plan implied by one metrics sample. -fn targeted_maintenance_plan( - metrics: &DuckLakeTableStorageMetrics, - scope: TargetedMaintenanceScope, -) -> TargetedMaintenancePlan { - let mut plan = TargetedMaintenancePlan::default(); - let active_delete_files = metrics.active_delete_files.max(0); - let deleted_row_ratio = metrics.deleted_row_ratio(); - - match scope { - TargetedMaintenanceScope::Idle => { - if active_delete_files >= MAINTENANCE_IDLE_REWRITE_DELETE_FILES_THRESHOLD - && deleted_row_ratio >= MAINTENANCE_IDLE_REWRITE_DELETED_ROW_RATIO_THRESHOLD - { - plan.rewrite_reason = Some(MaintenanceReason::IdleRewriteMetricsThreshold); - } - } - TargetedMaintenanceScope::Emergency => { - if active_delete_files >= MAINTENANCE_EMERGENCY_REWRITE_DELETE_FILES_THRESHOLD - || (active_delete_files >= MAINTENANCE_IDLE_REWRITE_DELETE_FILES_THRESHOLD - && deleted_row_ratio - >= MAINTENANCE_EMERGENCY_REWRITE_DELETED_ROW_RATIO_THRESHOLD) - { - plan.rewrite_reason = Some(MaintenanceReason::EmergencyRewriteMetricsThreshold); - } - } - } - - plan -} - -/// Sends one maintenance notification without blocking the caller indefinitely. -pub(super) async fn send_maintenance_notification( - notification_tx: &mpsc::Sender, - notification: TableMaintenanceNotification, -) { - if let Err(error) = notification_tx.send_timeout(notification, NOTIFICATION_SEND_TIMEOUT).await - { - match error { - mpsc::error::SendTimeoutError::Timeout(notification) => { - warn!( - table = %notification.table_name(), - "ducklake maintenance notification timed out" - ); - } - mpsc::error::SendTimeoutError::Closed(notification) => { - warn!( - table = %notification.table_name(), - "ducklake maintenance notification dropped" - ); - } - } - } -} - -/// Tries to enqueue one maintenance notification without awaiting channel -/// capacity. -fn try_send_maintenance_notification( - notification_tx: &mpsc::Sender, - notification: TableMaintenanceNotification, -) { - if let Err(error) = notification_tx.try_send(notification) { - match error { - mpsc::error::TrySendError::Full(notification) => { - warn!( - table = %notification.table_name(), - "ducklake maintenance notification dropped because channel is full" - ); - } - mpsc::error::TrySendError::Closed(notification) => { - warn!( - table = %notification.table_name(), - "ducklake maintenance notification dropped" - ); - } - } - } -} - -/// Starts warming the maintenance pool and spawns the periodic DuckLake worker. -#[allow(clippy::too_many_arguments)] -pub(super) fn spawn_ducklake_maintenance_worker( - manager: DuckLakeConnectionManager, - checkpoint_gate: Arc>, - table_write_slots: Arc>>>, - inline_flush_requested: Arc, - pending_inline_flush_requests: Arc, - pending_inline_size_sampler: Option, - expire_snapshots_older_than: Arc, -) -> EtlResult { - let mut pool = LazyDuckLakePool::new(manager, MAINTENANCE_POOL_SIZE, "maintenance"); - pool.warm_in_background(); - let (notification_tx, notification_rx) = mpsc::channel(1024); - let (shutdown_tx, shutdown_rx) = watch::channel(()); - let handle = tokio::spawn(run_ducklake_maintenance_worker( - pool, - checkpoint_gate, - table_write_slots, - inline_flush_requested, - pending_inline_flush_requests, - pending_inline_size_sampler, - CatalogMaintenanceConfig::new(expire_snapshots_older_than), - notification_rx, - shutdown_rx, - )); - - Ok(DuckLakeMaintenanceWorker { - notification_tx, - shutdown_tx, - handle: Mutex::new(handle.into()), - }) -} - -/// Coalesces notifications and runs background DuckLake maintenance. -#[allow(clippy::too_many_arguments)] -async fn run_ducklake_maintenance_worker( - mut pool: LazyDuckLakePool, - checkpoint_gate: Arc>, - table_write_slots: Arc>>>, - inline_flush_requested: Arc, - pending_inline_flush_requests: Arc, - pending_inline_size_sampler: Option, - catalog_maintenance_config: CatalogMaintenanceConfig, - mut notification_rx: mpsc::Receiver, - mut shutdown_rx: watch::Receiver<()>, -) { - let blocking_slots = pool.blocking_slots(); - let started_at = Instant::now(); - let mut flush_interval = tokio::time::interval_at( - started_at + MAINTENANCE_FLUSH_POLL_INTERVAL, - MAINTENANCE_FLUSH_POLL_INTERVAL, - ); - flush_interval.set_missed_tick_behavior(MissedTickBehavior::Skip); - - let mut table_states: HashMap = HashMap::new(); - let mut catalog_maintenance_state = CatalogMaintenanceState::new(started_at); - - loop { - tokio::select! { - biased; - _ = shutdown_rx.changed() => { - info!("ducklake maintenance worker shutting down"); - break; - } - maybe_notification = notification_rx.recv() => { - let Some(notification) = maybe_notification else { - info!("ducklake maintenance worker channel closed"); - break; - }; - - apply_maintenance_notification(&mut table_states, notification); - } - _ = flush_interval.tick() => { - for (table_name, table_state) in &mut table_states { - if let Some(pending_inline_size_sampler) = &pending_inline_size_sampler - && table_state.needs_pending_inline_data_sizes_sample() - { - match pending_inline_size_sampler.sample_table(table_name).await { - Ok(sizes) => { - table_state.record_pending_inline_data_sizes(Instant::now(), sizes); - } - Err(error) => { - tracing::error!( - table = %table_name, - error = %error, - "ducklake inline-size sampler query failed" - ); - } - } - } - - // If it needs to be flushed - if let Some(reason) = table_state.flush_reason() { - pending_inline_flush_requests.request(table_name.clone(), reason); - inline_flush_requested.store(true, AtomicOrdering::Release); - } - let now = Instant::now(); - - if !ENABLE_TARGETED_TABLE_MAINTENANCE { - continue; - } - - if let Some((plan, sampled_at)) = table_state - .emergency_targeted_maintenance_plan(now) - { - table_state.last_emergency_assessment_at = Some(sampled_at); - if plan.has_work() { - let pool = match pool.get_or_init_pool().await { - Ok(pool) => pool, - Err(error) => { - warn!(error = %error, "ducklake maintenance pool initialization failed"); - continue; - } - }; - match run_targeted_table_maintenance( - pool, - Arc::clone(&checkpoint_gate), - Arc::clone(&blocking_slots), - Arc::clone(&table_write_slots), - table_name.clone(), - plan, - ) - .await - { - Ok(outcome) => { - if outcome.is_completed() { - table_state.complete_targeted_maintenance(Instant::now()); - } - } - Err(error) => { - warn!( - table = %table_name, - error = %error, - "ducklake targeted maintenance failed" - ); - } - } - continue; - } - } - - if let Some(plan) = table_state.idle_targeted_maintenance_plan(now) { - if !plan.has_work() { - table_state.complete_idle_targeted_maintenance(Instant::now()); - continue; - } - - let pool = match pool.get_or_init_pool().await { - Ok(pool) => pool, - Err(error) => { - warn!(error = %error, "ducklake maintenance pool initialization failed"); - continue; - } - }; - match run_targeted_table_maintenance( - pool, - Arc::clone(&checkpoint_gate), - Arc::clone(&blocking_slots), - Arc::clone(&table_write_slots), - table_name.clone(), - plan, - ) - .await - { - Ok(outcome) => { - if outcome.is_completed() { - table_state.complete_targeted_maintenance(Instant::now()); - } - } - Err(error) => { - warn!( - table = %table_name, - error = %error, - "ducklake targeted maintenance failed" - ); - } - } - } - } - - table_states.retain(|_, state| { - state.dirty_since_compaction || state.latest_storage_metrics.is_some() - }); - - maybe_run_catalog_maintenance( - &mut pool, - Arc::clone(&checkpoint_gate), - Arc::clone(&blocking_slots), - &catalog_maintenance_config, - &mut catalog_maintenance_state, - ) - .await; - } - } - } -} - -/// Applies one maintenance notification to the table-state cache. -fn apply_maintenance_notification( - table_states: &mut HashMap, - notification: TableMaintenanceNotification, -) { - let now = Instant::now(); - match notification { - TableMaintenanceNotification::WriteActivity(activity) => { - table_states - .entry(activity.table_name) - .and_modify(|state| state.record_write_activity(now)) - .or_insert_with(|| { - let mut state = TableMaintenanceState::default(); - state.record_write_activity(now); - state - }); - } - TableMaintenanceNotification::TableMetricsSample(sample) => { - table_states - .entry(sample.table_name.clone()) - .and_modify(|state| state.record_metrics_sample(sample.clone())) - .or_insert_with(|| { - let mut state = TableMaintenanceState::default(); - state.record_metrics_sample(sample); - state - }); - } - TableMaintenanceNotification::FlushCompleted(completion) => { - let Some(state) = table_states.get_mut(&completion.table_name) else { - return; - }; - state.clear_pending_flush(completion.completed_at); - } - } -} - -/// Returns the table-local semaphore shared by writes and background -/// maintenance. -pub(super) fn table_write_slot( - table_write_slots: &Arc>>>, - table_name: &str, -) -> Arc { - let mut slots = table_write_slots.lock(); - let slot = slots.entry(table_name.to_owned()).or_insert_with(|| Arc::new(Semaphore::new(1))); - Arc::clone(slot) -} - -/// Tries to acquire the table-local semaphore without blocking the maintenance -/// worker. -fn try_acquire_table_write_slot( - table_write_slots: &Arc>>>, - table_name: &str, -) -> Option { - table_write_slot(table_write_slots, table_name).try_acquire_owned().ok() -} - -/// Runs targeted rewrite and merge maintenance for one table. -async fn run_targeted_table_maintenance( - pool: Arc>, - checkpoint_gate: Arc>, - blocking_slots: Arc, - table_write_slots: Arc>>>, - table_name: DuckLakeTableName, - plan: TargetedMaintenancePlan, -) -> EtlResult { - let Some(table_write_permit) = try_acquire_table_write_slot(&table_write_slots, &table_name) - else { - record_skipped_targeted_maintenance(plan); - return Ok(MaintenanceOutcome::SkippedBusy); - }; - let table_name_for_query = table_name.clone(); - let plan_for_query = plan; - let _checkpoint_guard = checkpoint_gate.read_owned().await; - - run_duckdb_blocking( - pool, - blocking_slots, - DuckDbBlockingOperationKind::Maintenance, - move |conn| { - let _table_write_permit = table_write_permit; - run_targeted_table_maintenance_blocking(conn, &table_name_for_query, plan_for_query) - }, - ) - .await - .or_else(|error| { - if suppress_known_ducklake_compaction_error(&table_name, plan, &error) { - Ok(MaintenanceOutcome::Noop) - } else { - Err(error) - } - }) -} - -/// Runs catalog-level DuckLake maintenance when its fixed cadence is due. -async fn maybe_run_catalog_maintenance( - pool: &mut LazyDuckLakePool, - checkpoint_gate: Arc>, - blocking_slots: Arc, - config: &CatalogMaintenanceConfig, - state: &mut CatalogMaintenanceState, -) { - let now = Instant::now(); - let expire_snapshots_due = state.expire_snapshots_due(now, config.expire_snapshots_interval); - let cleanup_old_files_due = state.cleanup_old_files_due(now, config.cleanup_old_files_interval); - - if !expire_snapshots_due && !cleanup_old_files_due { - return; - } - - let Ok(_checkpoint_guard) = checkpoint_gate.try_write_owned() else { - record_skipped_catalog_maintenance(expire_snapshots_due, cleanup_old_files_due); - return; - }; - - let pool = match pool.get_or_init_pool().await { - Ok(pool) => pool, - Err(error) => { - warn!(error = %error, "ducklake maintenance pool initialization failed"); - return; - } - }; - - match run_catalog_maintenance( - pool, - Arc::clone(&blocking_slots), - Arc::clone(&config.expire_snapshots_older_than), - expire_snapshots_due, - cleanup_old_files_due, - ) - .await - { - Ok((expired_snapshots, cleaned_up_files)) => { - let completed_at = Instant::now(); - if expire_snapshots_due { - state.complete_expire_snapshots(completed_at); - } - if cleanup_old_files_due { - state.complete_cleanup_old_files(completed_at); - } - info!( - expire_snapshots_older_than = %config.expire_snapshots_older_than, - expire_snapshots_due, - cleanup_old_files_due, - expired_snapshots, - cleaned_up_files, - "ducklake catalog maintenance completed" - ); - } - Err(error) => { - warn!( - expire_snapshots_older_than = %config.expire_snapshots_older_than, - error = %error, - "ducklake catalog maintenance failed" - ); - } - } -} - -/// Runs catalog-level maintenance inside one DuckDB blocking operation. -async fn run_catalog_maintenance( - pool: Arc>, - blocking_slots: Arc, - expire_snapshots_older_than: Arc, - expire_snapshots_due: bool, - cleanup_old_files_due: bool, -) -> EtlResult<(u64, u64)> { - run_duckdb_blocking( - pool, - blocking_slots, - DuckDbBlockingOperationKind::Maintenance, - move |conn| { - run_catalog_maintenance_blocking( - conn, - expire_snapshots_older_than.as_ref(), - expire_snapshots_due, - cleanup_old_files_due, - ) - }, - ) - .await -} - -/// Runs requested inline flushes before foreground ingestion begins. -pub(super) async fn maybe_run_requested_inline_flush( - pool: Arc>, - checkpoint_gate: Arc>, - blocking_slots: Arc, - inline_flush_requested: &AtomicBool, - pending_inline_flush_requests: &PendingInlineFlushRequests, - notification_tx: Option>, -) -> EtlResult<()> { - if !inline_flush_requested.swap(false, AtomicOrdering::AcqRel) { - return Ok(()); - } - - let requested_flushes = pending_inline_flush_requests.take_all(); - if requested_flushes.is_empty() { - return Ok(()); - } - - let Ok(_checkpoint_guard) = checkpoint_gate.try_write_owned() else { - for (_, reason) in &requested_flushes { - record_ducklake_maintenance_skipped( - MAINTENANCE_TASK_FLUSH, - MaintenanceOperation::FlushInlinedData, - *reason, - ); - } - pending_inline_flush_requests.restore(requested_flushes); - inline_flush_requested.store(true, AtomicOrdering::Release); - return Ok(()); - }; - - let mut requested_flushes = requested_flushes.into_iter(); - while let Some((table_name, reason)) = requested_flushes.next() { - let table_name_for_query = table_name.clone(); - let outcome = run_duckdb_blocking( - Arc::clone(&pool), - Arc::clone(&blocking_slots), - DuckDbBlockingOperationKind::Maintenance, - move |conn| { - flush_table_inlined_data_in_background_blocking(conn, &table_name_for_query, reason) - }, - ) - .await; - - match outcome { - Ok(outcome) => { - if outcome.is_completed() - && let Some(notification_tx) = notification_tx.as_ref() - { - try_send_maintenance_notification( - notification_tx, - TableMaintenanceNotification::FlushCompleted(TableFlushCompletion { - table_name, - completed_at: Instant::now(), - }), - ); - } - } - Err(error) => { - pending_inline_flush_requests.request(table_name, reason); - pending_inline_flush_requests.restore(requested_flushes); - inline_flush_requested.store(true, AtomicOrdering::Release); - return Err(error); - } - } - } - - Ok(()) -} - -/// Materializes one table's pending inlined rows and records the maintenance -/// outcome. -fn flush_table_inlined_data_in_background_blocking( - conn: &duckdb::Connection, - table_name: &str, - reason: MaintenanceReason, -) -> EtlResult { - let flush_started = Instant::now(); - let _in_progress_guard = DuckLakeMaintenanceInProgressGuard::start( - MAINTENANCE_TASK_FLUSH, - MaintenanceOperation::FlushInlinedData, - reason, - ); - let rows_flushed = flush_table_inlined_data(conn, table_name).inspect_err(|_error| { - record_ducklake_maintenance_duration( - MAINTENANCE_TASK_FLUSH, - MaintenanceOperation::FlushInlinedData, - reason, - MaintenanceOutcome::Failed, - flush_started.elapsed().as_secs_f64(), - ); - })?; - let outcome = MaintenanceOutcome::from(rows_flushed); - record_ducklake_maintenance_duration( - MAINTENANCE_TASK_FLUSH, - MaintenanceOperation::FlushInlinedData, - reason, - outcome, - flush_started.elapsed().as_secs_f64(), - ); - Ok(outcome) -} - -/// Runs targeted table maintenance and records per-operation outcomes. -fn run_targeted_table_maintenance_blocking( - conn: &duckdb::Connection, - table_name: &str, - plan: TargetedMaintenancePlan, -) -> EtlResult { - let mut rewritten_files = 0u64; - let mut rewrite_outcome = None; - - if let Some(reason) = plan.rewrite_reason { - let rewrite_started = Instant::now(); - let _in_progress_guard = DuckLakeMaintenanceInProgressGuard::start( - MAINTENANCE_TASK_TARGETED_MAINTENANCE, - MaintenanceOperation::RewriteDataFiles, - reason, - ); - rewritten_files = rewrite_table_data_files(conn, table_name).inspect_err(|_error| { - record_ducklake_maintenance_duration( - MAINTENANCE_TASK_TARGETED_MAINTENANCE, - MaintenanceOperation::RewriteDataFiles, - reason, - MaintenanceOutcome::Failed, - rewrite_started.elapsed().as_secs_f64(), - ); - })?; - let outcome = MaintenanceOutcome::from(rewritten_files); - record_ducklake_maintenance_duration( - MAINTENANCE_TASK_TARGETED_MAINTENANCE, - MaintenanceOperation::RewriteDataFiles, - reason, - outcome, - rewrite_started.elapsed().as_secs_f64(), - ); - rewrite_outcome = Some(outcome); - } - - info!( - table = %table_name, - rewrite_selected = plan.rewrite_reason.is_some(), - rewritten_files, - "ducklake targeted maintenance completed" - ); - - Ok(rewrite_outcome.unwrap_or(MaintenanceOutcome::Noop)) -} - -/// Builds the DuckLake snapshot-expiration call for one retention window. -fn expire_snapshots_sql(expire_snapshots_older_than: &str) -> String { - format!( - "CALL ducklake_expire_snapshots({}, older_than => CAST(now() AS TIMESTAMP) - CAST({} AS \ - INTERVAL));", - quote_literal(LAKE_CATALOG), - quote_literal(expire_snapshots_older_than), - ) -} - -/// Builds the DuckLake old-file cleanup call for one retention window. -fn cleanup_old_files_sql(expire_snapshots_older_than: &str) -> String { - format!( - "CALL ducklake_cleanup_old_files({}, older_than => CAST(now() AS TIMESTAMP) - CAST({} AS \ - INTERVAL));", - quote_literal(LAKE_CATALOG), - quote_literal(expire_snapshots_older_than), - ) -} - -/// Counts the rows returned by one DuckLake maintenance call. -fn count_ducklake_maintenance_rows( - conn: &duckdb::Connection, - sql: &str, - description: &'static str, -) -> EtlResult { - let mut statement = conn.prepare(sql).map_err(|source| { - etl_error!( - ErrorKind::DestinationQueryFailed, - description, - format_query_error_detail(sql), - source: source - ) - })?; - let mut rows = statement.query([]).map_err(|source| { - etl_error!( - ErrorKind::DestinationQueryFailed, - description, - format_query_error_detail(sql), - source: source - ) - })?; - let mut count = 0u64; - - while let Some(_row) = rows.next().map_err(|source| { - etl_error!( - ErrorKind::DestinationQueryFailed, - description, - format_query_error_detail(sql), - source: source - ) - })? { - count = count.saturating_add(1); - } - - Ok(count) -} - -/// Runs DuckLake catalog maintenance and records per-operation outcomes. -fn run_catalog_maintenance_blocking( - conn: &duckdb::Connection, - expire_snapshots_older_than: &str, - expire_snapshots_due: bool, - cleanup_old_files_due: bool, -) -> EtlResult<(u64, u64)> { - let mut expired_snapshots = 0u64; - let mut cleaned_up_files = 0u64; - - if expire_snapshots_due { - let expire_reason = MaintenanceReason::SnapshotRetentionThreshold; - let expire_started = Instant::now(); - let _expire_guard = DuckLakeMaintenanceInProgressGuard::start( - MAINTENANCE_TASK_CATALOG_MAINTENANCE, - MaintenanceOperation::ExpireSnapshots, - expire_reason, - ); - let expire_sql = expire_snapshots_sql(expire_snapshots_older_than); - expired_snapshots = - count_ducklake_maintenance_rows(conn, &expire_sql, "DuckLake expire snapshots failed") - .inspect_err(|_error| { - record_ducklake_maintenance_duration( - MAINTENANCE_TASK_CATALOG_MAINTENANCE, - MaintenanceOperation::ExpireSnapshots, - expire_reason, - MaintenanceOutcome::Failed, - expire_started.elapsed().as_secs_f64(), - ); - })?; - let expire_outcome = MaintenanceOutcome::from(expired_snapshots); - record_ducklake_maintenance_duration( - MAINTENANCE_TASK_CATALOG_MAINTENANCE, - MaintenanceOperation::ExpireSnapshots, - expire_reason, - expire_outcome, - expire_started.elapsed().as_secs_f64(), - ); - } - - if cleanup_old_files_due { - let cleanup_reason = MaintenanceReason::CleanupIntervalElapsed; - let cleanup_started = Instant::now(); - let _cleanup_guard = DuckLakeMaintenanceInProgressGuard::start( - MAINTENANCE_TASK_CATALOG_MAINTENANCE, - MaintenanceOperation::CleanupOldFiles, - cleanup_reason, - ); - let cleanup_sql = cleanup_old_files_sql(expire_snapshots_older_than); - cleaned_up_files = count_ducklake_maintenance_rows( - conn, - &cleanup_sql, - "DuckLake cleanup old files failed", - ) - .inspect_err(|_error| { - record_ducklake_maintenance_duration( - MAINTENANCE_TASK_CATALOG_MAINTENANCE, - MaintenanceOperation::CleanupOldFiles, - cleanup_reason, - MaintenanceOutcome::Failed, - cleanup_started.elapsed().as_secs_f64(), - ); - })?; - let cleanup_outcome = MaintenanceOutcome::from(cleaned_up_files); - record_ducklake_maintenance_duration( - MAINTENANCE_TASK_CATALOG_MAINTENANCE, - MaintenanceOperation::CleanupOldFiles, - cleanup_reason, - cleanup_outcome, - cleanup_started.elapsed().as_secs_f64(), - ); - } - - Ok((expired_snapshots, cleaned_up_files)) -} - -/// Flushes inlined user data for one table after the write transaction commits. -pub(super) fn flush_table_inlined_data( - conn: &duckdb::Connection, - table_name: &str, -) -> EtlResult { - let flush_started = Instant::now(); - let sql = format!( - r#"SELECT COALESCE(SUM(rows_flushed), 0) - FROM ducklake_flush_inlined_data({}, table_name => {});"#, - quote_literal(LAKE_CATALOG), - quote_literal(table_name), - ); - let rows_flushed: i64 = conn.query_row(&sql, [], |row| row.get(0)).map_err(|e| { - etl_error!( - ErrorKind::DestinationQueryFailed, - "DuckLake inlined data flush failed", - format_query_error_detail(&sql), - source: e - ) - })?; - let rows_flushed = rows_flushed.max(0) as u64; - let flush_result = if rows_flushed > 0 { "flushed" } else { "noop" }; - histogram!( - ETL_DUCKLAKE_INLINE_FLUSH_ROWS, - RESULT_LABEL => flush_result, - ) - .record(rows_flushed as f64); - histogram!( - ETL_DUCKLAKE_INLINE_FLUSH_DURATION_SECONDS, - RESULT_LABEL => flush_result, - ) - .record(flush_started.elapsed().as_secs_f64()); - - if rows_flushed > 0 { - debug!( - table = %table_name, - rows_flushed, - "ducklake inlined data flushed" - ); - } else { - debug!( - table = %table_name, - "ducklake inlined data already flushed" - ); - } - Ok(rows_flushed) -} - -/// Rewrites one table's delete-heavy files and returns created file count. -fn rewrite_table_data_files(conn: &duckdb::Connection, table_name: &str) -> EtlResult { - let sql = format!( - r#"SELECT COALESCE(SUM(files_created), 0) - FROM ducklake_rewrite_data_files({}, {});"#, - quote_literal(LAKE_CATALOG), - quote_literal(table_name), - ); - #[cfg(test)] - if FAIL_REWRITE_SINGLE_OUTPUT_FILE_ONCE_FOR_TESTS.swap(false, AtomicOrdering::Relaxed) { - let source = duckdb::Error::DuckDBFailure( - duckdb::ffi::Error::new(1), - Some("INTERNAL Error: DuckLakeCompaction - expected a single output file".to_owned()), - ); - return Err(etl_error!( - ErrorKind::DestinationQueryFailed, - "DuckLake rewrite data files failed", - format_query_error_detail(&sql), - source: source - )); - } - - let files_created: i64 = conn.query_row(&sql, [], |row| row.get(0)).map_err(|error| { - etl_error!( - ErrorKind::DestinationQueryFailed, - "DuckLake rewrite data files failed", - format_query_error_detail(&sql), - source: error - ) - })?; - - Ok(files_created.max(0) as u64) -} - -#[cfg(test)] -mod tests { - use etl_telemetry::metrics::init_metrics_handle; - - use super::*; - #[cfg(feature = "test-utils")] - use crate::ducklake::client::build_warm_ducklake_pool; - use crate::ducklake::metrics::register_metrics; - - fn maintenance_duration_count( - rendered: &str, - task: &str, - operation: MaintenanceOperation, - reason: MaintenanceReason, - outcome: MaintenanceOutcome, - ) -> f64 { - let task_label = format!(r#"{MAINTENANCE_TASK_LABEL}="{task}""#); - let operation_label = format!(r#"{MAINTENANCE_OPERATION_LABEL}="{}""#, operation.as_str()); - let reason_label = format!(r#"{MAINTENANCE_REASON_LABEL}="{}""#, reason.as_str()); - let outcome_label = format!(r#"{MAINTENANCE_OUTCOME_LABEL}="{}""#, outcome.as_str()); - - rendered - .lines() - .find_map(|line| { - if line.starts_with(&format!("{ETL_DUCKLAKE_MAINTENANCE_DURATION_SECONDS}_count")) - && line.contains(&task_label) - && line.contains(&operation_label) - && line.contains(&reason_label) - && line.contains(&outcome_label) - { - line.split_whitespace().last()?.parse::().ok() - } else { - None - } - }) - .unwrap_or(0.0) - } - - fn maintenance_skipped_counter_value( - rendered: &str, - task: &str, - operation: MaintenanceOperation, - reason: MaintenanceReason, - ) -> f64 { - let task_label = format!(r#"{MAINTENANCE_TASK_LABEL}="{task}""#); - let operation_label = format!(r#"{MAINTENANCE_OPERATION_LABEL}="{}""#, operation.as_str()); - let reason_label = format!(r#"{MAINTENANCE_REASON_LABEL}="{}""#, reason.as_str()); - - rendered - .lines() - .find_map(|line| { - if line.starts_with(ETL_DUCKLAKE_MAINTENANCE_SKIPPED_TOTAL) - && line.contains(&task_label) - && line.contains(&operation_label) - && line.contains(&reason_label) - { - line.split_whitespace().last()?.parse::().ok() - } else { - None - } - }) - .unwrap_or(0.0) - } - - fn maintenance_started_counter_value( - rendered: &str, - task: &str, - operation: MaintenanceOperation, - reason: MaintenanceReason, - ) -> f64 { - let task_label = format!(r#"{MAINTENANCE_TASK_LABEL}="{task}""#); - let operation_label = format!(r#"{MAINTENANCE_OPERATION_LABEL}="{}""#, operation.as_str()); - let reason_label = format!(r#"{MAINTENANCE_REASON_LABEL}="{}""#, reason.as_str()); - - rendered - .lines() - .find_map(|line| { - if line.starts_with(ETL_DUCKLAKE_MAINTENANCE_STARTED_TOTAL) - && line.contains(&task_label) - && line.contains(&operation_label) - && line.contains(&reason_label) - { - line.split_whitespace().last()?.parse::().ok() - } else { - None - } - }) - .unwrap_or(0.0) - } - - fn maintenance_in_progress_gauge_value( - rendered: &str, - task: &str, - operation: MaintenanceOperation, - reason: MaintenanceReason, - ) -> f64 { - let task_label = format!(r#"{MAINTENANCE_TASK_LABEL}="{task}""#); - let operation_label = format!(r#"{MAINTENANCE_OPERATION_LABEL}="{}""#, operation.as_str()); - let reason_label = format!(r#"{MAINTENANCE_REASON_LABEL}="{}""#, reason.as_str()); - - rendered - .lines() - .find_map(|line| { - if line.starts_with(ETL_DUCKLAKE_MAINTENANCE_IN_PROGRESS) - && line.contains(&task_label) - && line.contains(&operation_label) - && line.contains(&reason_label) - { - line.split_whitespace().last()?.parse::().ok() - } else { - None - } - }) - .unwrap_or(0.0) - } - - fn table_metrics_sample( - sampled_at: Instant, - metrics: DuckLakeTableStorageMetrics, - ) -> TableMetricsSample { - TableMetricsSample { table_name: "public_users".to_owned(), sampled_at, metrics } - } - - fn storage_metrics( - active_data_files: i64, - active_data_bytes: i64, - small_data_files: i64, - active_data_rows: i64, - active_delete_files: i64, - active_delete_bytes: i64, - deleted_rows: i64, - ) -> DuckLakeTableStorageMetrics { - DuckLakeTableStorageMetrics { - active_data_files, - active_data_bytes, - small_data_files, - active_data_rows, - active_delete_files, - active_delete_bytes, - deleted_rows, - } - } - - #[tokio::test] - async fn maintenance_duration_histogram_counts_are_exported_with_labels() { - let handle = init_metrics_handle().expect("failed to initialize prometheus handle"); - register_metrics(); - - let rendered_before = handle.render(); - let flush_before = maintenance_duration_count( - &rendered_before, - MAINTENANCE_TASK_FLUSH, - MaintenanceOperation::FlushInlinedData, - MaintenanceReason::PendingInlinedDataBytesThreshold, - MaintenanceOutcome::Applied, - ); - let rewrite_before = maintenance_duration_count( - &rendered_before, - MAINTENANCE_TASK_TARGETED_MAINTENANCE, - MaintenanceOperation::RewriteDataFiles, - MaintenanceReason::IdleRewriteMetricsThreshold, - MaintenanceOutcome::Applied, - ); - - record_ducklake_maintenance_duration( - MAINTENANCE_TASK_FLUSH, - MaintenanceOperation::FlushInlinedData, - MaintenanceReason::PendingInlinedDataBytesThreshold, - MaintenanceOutcome::Applied, - 0.25, - ); - record_ducklake_maintenance_duration( - MAINTENANCE_TASK_TARGETED_MAINTENANCE, - MaintenanceOperation::RewriteDataFiles, - MaintenanceReason::IdleRewriteMetricsThreshold, - MaintenanceOutcome::Applied, - 1.0, - ); - - let rendered_after = handle.render(); - let flush_after = maintenance_duration_count( - &rendered_after, - MAINTENANCE_TASK_FLUSH, - MaintenanceOperation::FlushInlinedData, - MaintenanceReason::PendingInlinedDataBytesThreshold, - MaintenanceOutcome::Applied, - ); - let rewrite_after = maintenance_duration_count( - &rendered_after, - MAINTENANCE_TASK_TARGETED_MAINTENANCE, - MaintenanceOperation::RewriteDataFiles, - MaintenanceReason::IdleRewriteMetricsThreshold, - MaintenanceOutcome::Applied, - ); - - assert!(flush_after > flush_before, "flush duration count did not increase"); - assert!(rewrite_after > rewrite_before, "rewrite duration count did not increase"); - } - - #[tokio::test] - async fn maintenance_started_counter_and_in_progress_gauge_are_exported_with_labels() { - let handle = init_metrics_handle().expect("failed to initialize prometheus handle"); - register_metrics(); - - let rendered_before = handle.render(); - let started_before = maintenance_started_counter_value( - &rendered_before, - MAINTENANCE_TASK_TARGETED_MAINTENANCE, - MaintenanceOperation::RewriteDataFiles, - MaintenanceReason::EmergencyRewriteMetricsThreshold, - ); - let in_progress_before = maintenance_in_progress_gauge_value( - &rendered_before, - MAINTENANCE_TASK_TARGETED_MAINTENANCE, - MaintenanceOperation::RewriteDataFiles, - MaintenanceReason::EmergencyRewriteMetricsThreshold, - ); - - let started_during; - let in_progress_during; - { - let _in_progress_guard = DuckLakeMaintenanceInProgressGuard::start( - MAINTENANCE_TASK_TARGETED_MAINTENANCE, - MaintenanceOperation::RewriteDataFiles, - MaintenanceReason::EmergencyRewriteMetricsThreshold, - ); - - let rendered_during = handle.render(); - started_during = maintenance_started_counter_value( - &rendered_during, - MAINTENANCE_TASK_TARGETED_MAINTENANCE, - MaintenanceOperation::RewriteDataFiles, - MaintenanceReason::EmergencyRewriteMetricsThreshold, - ); - in_progress_during = maintenance_in_progress_gauge_value( - &rendered_during, - MAINTENANCE_TASK_TARGETED_MAINTENANCE, - MaintenanceOperation::RewriteDataFiles, - MaintenanceReason::EmergencyRewriteMetricsThreshold, - ); - } - - let rendered_after = handle.render(); - let started_after = maintenance_started_counter_value( - &rendered_after, - MAINTENANCE_TASK_TARGETED_MAINTENANCE, - MaintenanceOperation::RewriteDataFiles, - MaintenanceReason::EmergencyRewriteMetricsThreshold, - ); - let in_progress_after = maintenance_in_progress_gauge_value( - &rendered_after, - MAINTENANCE_TASK_TARGETED_MAINTENANCE, - MaintenanceOperation::RewriteDataFiles, - MaintenanceReason::EmergencyRewriteMetricsThreshold, - ); - - assert!(started_during > started_before, "maintenance started counter did not increase"); - assert!( - in_progress_during > in_progress_before, - "maintenance in-progress gauge did not increase" - ); - assert_eq!(started_after, started_during); - assert_eq!(in_progress_after, in_progress_before); - } - - #[tokio::test] - async fn targeted_maintenance_busy_emits_skip_counter_for_rewrite() { - let handle = init_metrics_handle().expect("failed to initialize prometheus handle"); - register_metrics(); - - let rendered_before = handle.render(); - let skipped_before = maintenance_skipped_counter_value( - &rendered_before, - MAINTENANCE_TASK_TARGETED_MAINTENANCE, - MaintenanceOperation::RewriteDataFiles, - MaintenanceReason::EmergencyRewriteMetricsThreshold, - ); - - record_skipped_targeted_maintenance(TargetedMaintenancePlan { - rewrite_reason: Some(MaintenanceReason::EmergencyRewriteMetricsThreshold), - }); - - let rendered_after = handle.render(); - let skipped_after = maintenance_skipped_counter_value( - &rendered_after, - MAINTENANCE_TASK_TARGETED_MAINTENANCE, - MaintenanceOperation::RewriteDataFiles, - MaintenanceReason::EmergencyRewriteMetricsThreshold, - ); - - assert!(skipped_after > skipped_before, "rewrite skip count did not increase"); - } - - #[tokio::test] - async fn catalog_maintenance_busy_emits_skip_counter_for_both_operations() { - let handle = init_metrics_handle().expect("failed to initialize prometheus handle"); - register_metrics(); - - let rendered_before = handle.render(); - let expire_before = maintenance_skipped_counter_value( - &rendered_before, - MAINTENANCE_TASK_CATALOG_MAINTENANCE, - MaintenanceOperation::ExpireSnapshots, - MaintenanceReason::SnapshotRetentionThreshold, - ); - let cleanup_before = maintenance_skipped_counter_value( - &rendered_before, - MAINTENANCE_TASK_CATALOG_MAINTENANCE, - MaintenanceOperation::CleanupOldFiles, - MaintenanceReason::CleanupIntervalElapsed, - ); - - record_skipped_catalog_maintenance(true, true); - - let rendered_after = handle.render(); - let expire_after = maintenance_skipped_counter_value( - &rendered_after, - MAINTENANCE_TASK_CATALOG_MAINTENANCE, - MaintenanceOperation::ExpireSnapshots, - MaintenanceReason::SnapshotRetentionThreshold, - ); - let cleanup_after = maintenance_skipped_counter_value( - &rendered_after, - MAINTENANCE_TASK_CATALOG_MAINTENANCE, - MaintenanceOperation::CleanupOldFiles, - MaintenanceReason::CleanupIntervalElapsed, - ); - - assert!(expire_after > expire_before, "expire snapshots skip count did not increase"); - assert!(cleanup_after > cleanup_before, "cleanup old files skip count did not increase"); - } - - #[test] - fn table_maintenance_state_without_sample_or_row_threshold_does_not_flush() { - let now = Instant::now(); - let mut state = TableMaintenanceState::default(); - state.record_write_activity(now); - - assert_eq!(state.flush_reason(), None); - } - - #[test] - fn table_maintenance_state_uses_sampled_inline_data_bytes_flush_reason() { - let now = Instant::now(); - let mut state = TableMaintenanceState::default(); - state.record_write_activity(now); - state.record_pending_inline_data_sizes( - now + Duration::from_secs(1), - DuckLakePendingInlineDataSizes { - inlined_bytes: MAINTENANCE_PENDING_INLINED_DATA_BYTES_THRESHOLD, - }, - ); - - assert_eq!(state.flush_reason(), Some(MaintenanceReason::PendingInlinedDataBytesThreshold)); - } - - #[test] - fn table_maintenance_state_ignores_below_threshold_sampled_inline_data() { - let now = Instant::now(); - let mut state = TableMaintenanceState::default(); - state.record_write_activity(now); - state.record_pending_inline_data_sizes( - now + Duration::from_secs(1), - DuckLakePendingInlineDataSizes { - inlined_bytes: MAINTENANCE_PENDING_INLINED_DATA_BYTES_THRESHOLD - 1, - }, - ); - - assert_eq!(state.flush_reason(), None); - } - - #[test] - fn table_maintenance_state_selects_idle_rewrite_from_delete_pressure() { - let now = Instant::now(); - let mut state = TableMaintenanceState::default(); - state.record_write_activity(now); - state.record_metrics_sample(table_metrics_sample( - now + Duration::from_secs(1), - storage_metrics(10, 20_000_000, 2, 1000, 32, 10_000, 150), - )); - - assert_eq!( - state.idle_targeted_maintenance_plan( - now + MAINTENANCE_TABLE_COMPACTION_IDLE_THRESHOLD - + MAINTENANCE_TABLE_COMPACTION_INTERVAL, - ), - Some(TargetedMaintenancePlan { - rewrite_reason: Some(MaintenanceReason::IdleRewriteMetricsThreshold), - }) - ); - } - - #[test] - fn table_maintenance_state_selects_emergency_rewrite_from_delete_pressure() { - let now = Instant::now(); - let mut state = TableMaintenanceState::default(); - state.record_write_activity(now); - state.record_metrics_sample(table_metrics_sample( - now + Duration::from_secs(1), - storage_metrics(10, 20_000_000, 2, 1000, 128, 10_000, 150), - )); - - assert_eq!( - state.emergency_targeted_maintenance_plan(now + MAINTENANCE_TABLE_COMPACTION_INTERVAL), - Some(( - TargetedMaintenancePlan { - rewrite_reason: Some(MaintenanceReason::EmergencyRewriteMetricsThreshold), - }, - now + Duration::from_secs(1), - )) - ); - } - - #[test] - fn maintenance_outcome_from_rows_flushed_marks_applied_and_noop() { - assert_eq!(MaintenanceOutcome::from(0), MaintenanceOutcome::Noop); - assert_eq!(MaintenanceOutcome::from(3), MaintenanceOutcome::Applied); - } - - #[test] - fn catalog_maintenance_state_tracks_independent_due_intervals() { - let now = Instant::now(); - let state = CatalogMaintenanceState::new(now); - - assert!(!state.expire_snapshots_due(now, MAINTENANCE_EXPIRE_SNAPSHOTS_INTERVAL)); - assert!(!state.cleanup_old_files_due(now, MAINTENANCE_CLEANUP_OLD_FILES_INTERVAL)); - - assert!(!state.expire_snapshots_due( - now + MAINTENANCE_EXPIRE_SNAPSHOTS_INTERVAL - Duration::from_secs(1), - MAINTENANCE_EXPIRE_SNAPSHOTS_INTERVAL - )); - assert!(state.expire_snapshots_due( - now + MAINTENANCE_EXPIRE_SNAPSHOTS_INTERVAL, - MAINTENANCE_EXPIRE_SNAPSHOTS_INTERVAL - )); - assert!(!state.cleanup_old_files_due( - now + MAINTENANCE_CLEANUP_OLD_FILES_INTERVAL - Duration::from_secs(1), - MAINTENANCE_CLEANUP_OLD_FILES_INTERVAL - )); - assert!(state.cleanup_old_files_due( - now + MAINTENANCE_CLEANUP_OLD_FILES_INTERVAL, - MAINTENANCE_CLEANUP_OLD_FILES_INTERVAL - )); - } - - #[test] - fn catalog_maintenance_sql_builders_use_interval_casts() { - let expire_sql = expire_snapshots_sql("2 days"); - let cleanup_sql = cleanup_old_files_sql("2 days"); - - assert!(expire_sql.contains("ducklake_expire_snapshots")); - assert!(cleanup_sql.contains("ducklake_cleanup_old_files")); - assert!(expire_sql.contains("CAST('2 days' AS INTERVAL)")); - assert!(cleanup_sql.contains("CAST('2 days' AS INTERVAL)")); - } - - #[tokio::test] - async fn flush_failure_records_failed_metric() { - let handle = init_metrics_handle().expect("failed to initialize prometheus handle"); - register_metrics(); - let conn = duckdb::Connection::open_in_memory().expect("failed to open in-memory duckdb"); - - let rendered_before = handle.render(); - let failed_before = maintenance_duration_count( - &rendered_before, - MAINTENANCE_TASK_FLUSH, - MaintenanceOperation::FlushInlinedData, - MaintenanceReason::PendingInlinedDataBytesThreshold, - MaintenanceOutcome::Failed, - ); - - let error = flush_table_inlined_data_in_background_blocking( - &conn, - "public_users", - MaintenanceReason::PendingInlinedDataBytesThreshold, - ) - .expect_err("flush should fail without ducklake functions"); - - assert!(matches!(error.kind(), ErrorKind::DestinationQueryFailed)); - - let rendered_after = handle.render(); - let failed_after = maintenance_duration_count( - &rendered_after, - MAINTENANCE_TASK_FLUSH, - MaintenanceOperation::FlushInlinedData, - MaintenanceReason::PendingInlinedDataBytesThreshold, - MaintenanceOutcome::Failed, - ); - - assert!(failed_after > failed_before, "flush failed duration count did not increase"); - } - - #[tokio::test] - async fn rewrite_failure_records_failed_duration() { - let handle = init_metrics_handle().expect("failed to initialize prometheus handle"); - register_metrics(); - let conn = duckdb::Connection::open_in_memory().expect("failed to open in-memory duckdb"); - - let rendered_before = handle.render(); - let failed_before = maintenance_duration_count( - &rendered_before, - MAINTENANCE_TASK_TARGETED_MAINTENANCE, - MaintenanceOperation::RewriteDataFiles, - MaintenanceReason::IdleRewriteMetricsThreshold, - MaintenanceOutcome::Failed, - ); - - let error = run_targeted_table_maintenance_blocking( - &conn, - "public_users", - TargetedMaintenancePlan { - rewrite_reason: Some(MaintenanceReason::IdleRewriteMetricsThreshold), - }, - ) - .expect_err("targeted maintenance should fail without ducklake functions"); - - assert!(matches!(error.kind(), ErrorKind::DestinationQueryFailed)); - - let rendered_after = handle.render(); - let failed_after = maintenance_duration_count( - &rendered_after, - MAINTENANCE_TASK_TARGETED_MAINTENANCE, - MaintenanceOperation::RewriteDataFiles, - MaintenanceReason::IdleRewriteMetricsThreshold, - MaintenanceOutcome::Failed, - ); - - assert!(failed_after > failed_before, "rewrite failed duration count did not increase"); - } - - #[tokio::test] - async fn catalog_expire_failure_records_failed_duration() { - let handle = init_metrics_handle().expect("failed to initialize prometheus handle"); - register_metrics(); - let conn = duckdb::Connection::open_in_memory().expect("failed to open in-memory duckdb"); - - let rendered_before = handle.render(); - let failed_before = maintenance_duration_count( - &rendered_before, - MAINTENANCE_TASK_CATALOG_MAINTENANCE, - MaintenanceOperation::ExpireSnapshots, - MaintenanceReason::SnapshotRetentionThreshold, - MaintenanceOutcome::Failed, - ); - - let error = run_catalog_maintenance_blocking(&conn, "7 days", true, true) - .expect_err("catalog maintenance should fail without ducklake functions"); - - assert!(matches!(error.kind(), ErrorKind::DestinationQueryFailed)); - - let rendered_after = handle.render(); - let failed_after = maintenance_duration_count( - &rendered_after, - MAINTENANCE_TASK_CATALOG_MAINTENANCE, - MaintenanceOperation::ExpireSnapshots, - MaintenanceReason::SnapshotRetentionThreshold, - MaintenanceOutcome::Failed, - ); - let cleanup_failed_after = maintenance_duration_count( - &rendered_after, - MAINTENANCE_TASK_CATALOG_MAINTENANCE, - MaintenanceOperation::CleanupOldFiles, - MaintenanceReason::CleanupIntervalElapsed, - MaintenanceOutcome::Failed, - ); - - assert!(failed_after > failed_before, "expire snapshots failed duration did not increase"); - assert_eq!(cleanup_failed_after, 0.0, "cleanup should not run after expiration fails"); - } - - #[cfg(feature = "test-utils")] - #[tokio::test] - async fn known_rewrite_single_output_file_error_is_suppressed_and_recycles_connection() { - let open_count = Arc::new(std::sync::atomic::AtomicUsize::new(0)); - let manager = DuckLakeConnectionManager { - setup_plan: Arc::new(crate::ducklake::config::DuckLakeSetupPlan::default()), - disable_extension_autoload: cfg!(target_os = "linux"), - interrupt_registry: Arc::new( - crate::ducklake::client::DuckLakeInterruptRegistry::default(), - ), - open_count: Arc::clone(&open_count), - }; - let pool = Arc::new( - build_warm_ducklake_pool(manager, 1, "test") - .await - .expect("failed to build maintenance test pool"), - ); - let checkpoint_gate = Arc::new(RwLock::new(())); - let blocking_slots = Arc::new(Semaphore::new(1)); - let table_write_slots = Arc::new(Mutex::new(HashMap::new())); - FAIL_REWRITE_SINGLE_OUTPUT_FILE_ONCE_FOR_TESTS.store(true, AtomicOrdering::Relaxed); - - let outcome = run_targeted_table_maintenance( - Arc::clone(&pool), - Arc::clone(&checkpoint_gate), - Arc::clone(&blocking_slots), - Arc::clone(&table_write_slots), - "public_users".to_owned(), - TargetedMaintenancePlan { - rewrite_reason: Some(MaintenanceReason::IdleRewriteMetricsThreshold), - }, - ) - .await - .expect("known DuckLake compaction bug should be suppressed"); - - assert_eq!(outcome, MaintenanceOutcome::Noop); - - let verification = run_duckdb_blocking( - Arc::clone(&pool), - Arc::clone(&blocking_slots), - DuckDbBlockingOperationKind::Maintenance, - |conn| { - conn.query_row("SELECT 1", [], |row| row.get::<_, i64>(0)).map_err(|source| { - etl_error!( - ErrorKind::DestinationQueryFailed, - "DuckLake connection recycling verification query failed", - source: source - ) - }) - }, - ) - .await - .expect("expected follow-up query to succeed on recycled connection"); - - assert_eq!(verification, 1); - assert!(open_count.load(AtomicOrdering::Relaxed) > 1); - } - - #[cfg(feature = "test-utils")] - #[tokio::test] - async fn spawn_ducklake_maintenance_worker_warms_pool_in_background() { - let open_count = Arc::new(std::sync::atomic::AtomicUsize::new(0)); - let worker = spawn_ducklake_maintenance_worker( - DuckLakeConnectionManager { - setup_plan: Arc::new(crate::ducklake::config::DuckLakeSetupPlan::default()), - disable_extension_autoload: cfg!(target_os = "linux"), - interrupt_registry: Arc::new( - crate::ducklake::client::DuckLakeInterruptRegistry::default(), - ), - open_count: Arc::clone(&open_count), - }, - Arc::new(RwLock::new(())), - Arc::new(Mutex::new(HashMap::new())), - Arc::new(AtomicBool::new(false)), - Arc::new(PendingInlineFlushRequests::default()), - None, - Arc::::from("7 days"), - ) - .expect("failed to spawn maintenance worker"); - - let deadline = tokio::time::Instant::now() + Duration::from_secs(5); - while open_count.load(AtomicOrdering::Relaxed) < MAINTENANCE_POOL_SIZE as usize { - assert!(tokio::time::Instant::now() < deadline); - tokio::time::sleep(Duration::from_millis(10)).await; - } - - worker.shutdown_tx.send(()).expect("maintenance worker shutdown channel should stay open"); - let handle = worker.handle.lock().take().expect("maintenance worker handle should exist"); - handle.await.expect("maintenance worker should shut down cleanly"); - } - - #[cfg(feature = "test-utils")] - #[tokio::test] - async fn requested_inline_flush_stays_pending_when_mutation_guard_is_active() { - let pool = Arc::new( - build_warm_ducklake_pool( - DuckLakeConnectionManager { - setup_plan: Arc::new(crate::ducklake::config::DuckLakeSetupPlan::default()), - disable_extension_autoload: cfg!(target_os = "linux"), - interrupt_registry: Arc::new( - crate::ducklake::client::DuckLakeInterruptRegistry::default(), - ), - open_count: Arc::new(std::sync::atomic::AtomicUsize::new(0)), - }, - 1, - "test", - ) - .await - .expect("failed to build inline flush test pool"), - ); - let handle = init_metrics_handle().expect("failed to initialize prometheus handle"); - register_metrics(); - let checkpoint_gate = Arc::new(RwLock::new(())); - let mutation_guard = Arc::clone(&checkpoint_gate).read_owned().await; - let inline_flush_requested = Arc::new(AtomicBool::new(true)); - let pending_inline_flush_requests = PendingInlineFlushRequests::default(); - pending_inline_flush_requests.request( - "public_users".to_owned(), - MaintenanceReason::PendingInlinedDataBytesThreshold, - ); - - let rendered_before = handle.render(); - let skipped_before = maintenance_skipped_counter_value( - &rendered_before, - MAINTENANCE_TASK_FLUSH, - MaintenanceOperation::FlushInlinedData, - MaintenanceReason::PendingInlinedDataBytesThreshold, - ); - - maybe_run_requested_inline_flush( - Arc::clone(&pool), - Arc::clone(&checkpoint_gate), - Arc::new(Semaphore::new(1)), - inline_flush_requested.as_ref(), - &pending_inline_flush_requests, - None, - ) - .await - .expect("inline flush should be skipped while mutation work is active"); - - let rendered_after = handle.render(); - let skipped_after = maintenance_skipped_counter_value( - &rendered_after, - MAINTENANCE_TASK_FLUSH, - MaintenanceOperation::FlushInlinedData, - MaintenanceReason::PendingInlinedDataBytesThreshold, - ); - - assert!(skipped_after > skipped_before); - assert!(inline_flush_requested.load(AtomicOrdering::Acquire)); - assert_eq!( - pending_inline_flush_requests.requests.lock().get("public_users"), - Some(&MaintenanceReason::PendingInlinedDataBytesThreshold) - ); - - drop(mutation_guard); - } - - #[tokio::test] - async fn flush_busy_emits_skip_counter_only() { - let handle = init_metrics_handle().expect("failed to initialize prometheus handle"); - register_metrics(); - - let rendered_before = handle.render(); - let skipped_before = maintenance_skipped_counter_value( - &rendered_before, - MAINTENANCE_TASK_FLUSH, - MaintenanceOperation::FlushInlinedData, - MaintenanceReason::PendingInlinedDataBytesThreshold, - ); - let duration_before = maintenance_duration_count( - &rendered_before, - MAINTENANCE_TASK_FLUSH, - MaintenanceOperation::FlushInlinedData, - MaintenanceReason::PendingInlinedDataBytesThreshold, - MaintenanceOutcome::SkippedBusy, - ); - - record_ducklake_maintenance_skipped( - MAINTENANCE_TASK_FLUSH, - MaintenanceOperation::FlushInlinedData, - MaintenanceReason::PendingInlinedDataBytesThreshold, - ); - - let rendered_after = handle.render(); - let skipped_after = maintenance_skipped_counter_value( - &rendered_after, - MAINTENANCE_TASK_FLUSH, - MaintenanceOperation::FlushInlinedData, - MaintenanceReason::PendingInlinedDataBytesThreshold, - ); - let duration_after = maintenance_duration_count( - &rendered_after, - MAINTENANCE_TASK_FLUSH, - MaintenanceOperation::FlushInlinedData, - MaintenanceReason::PendingInlinedDataBytesThreshold, - MaintenanceOutcome::SkippedBusy, - ); - - assert!(skipped_after > skipped_before, "flush skipped counter did not increase"); - assert_eq!(duration_after, duration_before, "flush skip should not emit a duration sample"); - } -} diff --git a/crates/etl-destinations/src/ducklake/maintenance_runner.rs b/crates/etl-destinations/src/ducklake/maintenance_runner.rs new file mode 100644 index 000000000..8e30d52eb --- /dev/null +++ b/crates/etl-destinations/src/ducklake/maintenance_runner.rs @@ -0,0 +1,898 @@ +//! One-shot DuckLake maintenance execution for Kubernetes maintenance jobs. + +use std::sync::Arc; + +use etl::{ + error::{ErrorKind, EtlResult}, + etl_error, +}; +use metrics::histogram; +use pg_escape::{quote_identifier, quote_literal}; +use sqlx::{AssertSqlSafe, PgPool, postgres::PgPoolOptions}; +use tokio::sync::Semaphore; +use tracing::{debug, info}; +use url::Url; + +use crate::ducklake::{ + LAKE_CATALOG, S3Config, + client::{ + DuckDbBlockingOperationKind, DuckLakeConnectionManager, DuckLakeInterruptRegistry, + build_warm_ducklake_pool, format_query_error_detail, run_duckdb_blocking, + }, + config::{ + MAINTENANCE_TARGET_FILE_SIZE, build_setup_plan, current_duckdb_extension_strategy, + maintenance_target_file_size_sql, + }, + inline_size::DuckLakePendingInlineSizeSampler, + metrics::{ + DuckLakeTableStorageMetrics, ETL_DUCKLAKE_INLINE_FLUSH_DURATION_SECONDS, + ETL_DUCKLAKE_INLINE_FLUSH_ROWS, RESULT_LABEL, query_table_storage_metrics, + resolve_ducklake_metadata_schema_blocking, + }, +}; + +#[derive(Clone)] +struct DuckDbMaintenanceExecutor { + pool: Arc>, + blocking_slots: Arc, +} + +impl DuckDbMaintenanceExecutor { + async fn run(&self, operation: F) -> EtlResult + where + R: Send + 'static, + F: FnOnce(&duckdb::Connection) -> EtlResult + Send + 'static, + { + run_duckdb_blocking( + Arc::clone(&self.pool), + Arc::clone(&self.blocking_slots), + DuckDbBlockingOperationKind::Maintenance, + operation, + ) + .await + } +} + +/// Configuration for one external DuckLake maintenance run. +#[derive(Clone, Debug)] +pub struct DuckLakeMaintenanceConfig { + /// DuckLake PostgreSQL catalog URL. + pub catalog_url: Url, + /// DuckLake data path. + pub data_path: Url, + /// DuckDB connection pool size for the one-shot runner. + pub pool_size: u32, + /// Optional S3-compatible storage config. + pub s3: Option, + /// Optional DuckLake metadata schema. + pub metadata_schema: Option, + /// Optional DuckDB memory cache limit. + pub duckdb_memory_cache_limit: Option, + /// DuckLake `target_file_size` used by compaction. + pub maintenance_target_file_size: Option, + /// Inline flush operation config. + pub inline_flush: InlineFlushMaintenanceConfig, + /// Merge-adjacent-files operation config. + pub merge_adjacent_files: MergeAdjacentFilesMaintenanceConfig, + /// Rewrite-data-files operation config. + pub rewrite_data_files: RewriteDataFilesMaintenanceConfig, + /// Snapshot-expiration operation config. + pub expire_snapshots: ExpireSnapshotsMaintenanceConfig, + /// Old-file cleanup operation config. + pub cleanup_old_files: CleanupOldFilesMaintenanceConfig, +} + +/// Inline flush operation config. +#[derive(Clone, Copy, Debug)] +pub struct InlineFlushMaintenanceConfig { + /// Whether inline flush is enabled. + pub enabled: bool, + /// Minimum pending inlined bytes before flushing a table. + pub min_inlined_bytes: u64, +} + +/// Merge-adjacent-files operation config. +#[derive(Clone, Debug)] +pub struct MergeAdjacentFilesMaintenanceConfig { + /// Whether merge-adjacent-files is enabled. + pub enabled: bool, + /// Maximum compacted output files per table. + pub max_compacted_files: u32, + /// Maximum tables selected in one run. + pub max_tables_per_run: u32, + /// Target file size used during compaction. + pub target_file_size: String, +} + +/// Rewrite-data-files operation config. +#[derive(Clone, Copy, Debug)] +pub struct RewriteDataFilesMaintenanceConfig { + /// Whether rewrite-data-files is enabled. + pub enabled: bool, + /// Minimum active data-file count before rewrite is attempted. + pub min_active_data_files: i64, + /// Maximum tables selected in one run. + pub max_tables_per_run: u32, +} + +/// Snapshot-expiration operation config. +#[derive(Clone, Debug)] +pub struct ExpireSnapshotsMaintenanceConfig { + /// Whether snapshot expiration is enabled. + pub enabled: bool, + /// Retention window passed to DuckLake. + pub older_than: String, +} + +/// Old-file cleanup operation config. +#[derive(Clone, Debug)] +pub struct CleanupOldFilesMaintenanceConfig { + /// Whether old-file cleanup is enabled. + pub enabled: bool, + /// Retention window passed to DuckLake. + pub older_than: String, +} + +/// Structured outcome for one external maintenance run. +#[derive(Clone, Debug, Default, PartialEq)] +pub struct DuckLakeMaintenanceOutcome { + /// Tables whose inline data was flushed. + pub inline_flush_tables: u32, + /// Rows flushed from inlined storage. + pub inline_flush_rows: u64, + /// Tables passed to merge-adjacent-files. + pub merge_adjacent_files_tables: u32, + /// Files created by merge-adjacent-files. + pub merge_adjacent_files_created: u64, + /// Tables passed to rewrite-data-files. + pub rewrite_data_files_tables: u32, + /// Files created by rewrite-data-files. + pub rewrite_data_files_created: u64, + /// Snapshots expired by snapshot expiration. + pub expired_snapshots: u64, + /// Files removed by old-file cleanup. + pub cleaned_up_files: u64, +} + +impl DuckLakeMaintenanceOutcome { + /// Returns whether any operation did work. + pub fn applied(&self) -> bool { + self.inline_flush_rows > 0 + || self.merge_adjacent_files_created > 0 + || self.rewrite_data_files_tables > 0 + || self.rewrite_data_files_created > 0 + || self.expired_snapshots > 0 + || self.cleaned_up_files > 0 + } +} + +/// Runs one external DuckLake maintenance attempt. +pub async fn run_maintenance_once( + config: DuckLakeMaintenanceConfig, +) -> EtlResult { + validate_config(&config)?; + let cleanup_old_files_enabled = + config.cleanup_old_files.enabled || config.rewrite_data_files.enabled; + + info!( + pool_size = config.pool_size, + metadata_schema = config.metadata_schema.as_deref(), + inline_flush_enabled = config.inline_flush.enabled, + inline_flush_min_inlined_bytes = config.inline_flush.min_inlined_bytes, + merge_adjacent_files_enabled = config.merge_adjacent_files.enabled, + merge_adjacent_files_max_compacted_files = config.merge_adjacent_files.max_compacted_files, + merge_adjacent_files_max_tables_per_run = config.merge_adjacent_files.max_tables_per_run, + merge_adjacent_files_target_file_size = %config.merge_adjacent_files.target_file_size, + rewrite_data_files_enabled = config.rewrite_data_files.enabled, + rewrite_data_files_min_active_data_files = config.rewrite_data_files.min_active_data_files, + rewrite_data_files_max_tables_per_run = config.rewrite_data_files.max_tables_per_run, + expire_snapshots_enabled = config.expire_snapshots.enabled, + expire_snapshots_older_than = %config.expire_snapshots.older_than, + cleanup_old_files_enabled, + cleanup_old_files_explicitly_enabled = config.cleanup_old_files.enabled, + cleanup_old_files_older_than = %config.cleanup_old_files.older_than, + "ducklake external maintenance runner configured" + ); + + let duckdb = open_maintenance_executor(&config).await?; + let metadata_schema = match config.metadata_schema.clone() { + Some(metadata_schema) => metadata_schema, + None => resolve_metadata_schema(&duckdb).await?, + }; + info!( + metadata_schema = %metadata_schema, + "ducklake external maintenance metadata schema resolved" + ); + let metadata_pg_pool = PgPoolOptions::new() + .max_connections(1) + .connect_lazy(config.catalog_url.as_str()) + .map_err(|source| { + etl_error!( + ErrorKind::DestinationConnectionFailed, + "DuckLake catalog metadata pool configuration failed", + source: source + ) + })?; + let table_names = list_ducklake_tables(&metadata_pg_pool, &metadata_schema).await?; + info!( + table_count = table_names.len(), + tables = ?table_names, + "ducklake external maintenance discovered active tables" + ); + let mut outcome = DuckLakeMaintenanceOutcome::default(); + + if config.inline_flush.enabled { + run_inline_flush( + &duckdb, + &metadata_pg_pool, + &metadata_schema, + &table_names, + config.inline_flush, + &mut outcome, + ) + .await?; + } + + if config.merge_adjacent_files.enabled { + run_merge_adjacent_files( + &duckdb, + &metadata_pg_pool, + &metadata_schema, + &table_names, + &config.merge_adjacent_files, + &mut outcome, + ) + .await?; + } + + if config.rewrite_data_files.enabled { + merge_adjacent_files_for_rewrite(&duckdb).await?; + run_rewrite_data_files( + &duckdb, + &metadata_pg_pool, + &metadata_schema, + &table_names, + config.rewrite_data_files, + &mut outcome, + ) + .await?; + } + + if config.expire_snapshots.enabled { + run_expire_snapshots(&duckdb, &config.expire_snapshots, &mut outcome).await?; + } + + if cleanup_old_files_enabled { + run_cleanup_old_files(&duckdb, &config.cleanup_old_files, &mut outcome).await?; + } + + info!(outcome = ?outcome, applied = outcome.applied(), "ducklake external maintenance completed"); + Ok(outcome) +} + +/// Validates one maintenance runner config. +fn validate_config(config: &DuckLakeMaintenanceConfig) -> EtlResult<()> { + if !matches!(config.catalog_url.scheme(), "postgres" | "postgresql") { + return Err(etl_error!( + ErrorKind::ConfigError, + "DuckLake external maintenance requires a PostgreSQL catalog", + format!("unsupported catalog URL scheme `{}`", config.catalog_url.scheme()) + )); + } + if config.pool_size == 0 { + return Err(etl_error!( + ErrorKind::ConfigError, + "DuckLake external maintenance pool size must be greater than zero" + )); + } + Ok(()) +} + +/// Opens initialized DuckDB connections for maintenance. +async fn open_maintenance_executor( + config: &DuckLakeMaintenanceConfig, +) -> EtlResult { + let extension_strategy = current_duckdb_extension_strategy()?; + let target_file_size = config + .maintenance_target_file_size + .as_deref() + .or(Some(config.merge_adjacent_files.target_file_size.as_str())) + .unwrap_or(MAINTENANCE_TARGET_FILE_SIZE); + info!(target_file_size, "opening ducklake external maintenance connection"); + let setup_plan = Arc::new(build_setup_plan( + &config.catalog_url, + &config.data_path, + config.s3.as_ref(), + config.metadata_schema.as_deref(), + None, + )?); + let manager = DuckLakeConnectionManager { + setup_plan, + disable_extension_autoload: extension_strategy.disables_autoload(), + interrupt_registry: Arc::new(DuckLakeInterruptRegistry::default()), + #[cfg(feature = "test-utils")] + open_count: Arc::new(std::sync::atomic::AtomicUsize::new(0)), + }; + let pool = Arc::new( + build_warm_ducklake_pool(manager, config.pool_size, "external-maintenance").await?, + ); + let blocking_slots = Arc::new(Semaphore::new(config.pool_size as usize)); + let executor = DuckDbMaintenanceExecutor { pool, blocking_slots }; + let sql = maintenance_target_file_size_sql(Some(target_file_size)); + executor + .run(move |conn| { + conn.execute_batch(&sql).map_err(|source| { + etl_error!( + ErrorKind::DestinationQueryFailed, + "DuckLake target_file_size configuration failed", + format_query_error_detail(&sql), + source: source + ) + })?; + Ok(()) + }) + .await?; + info!(target_file_size, "ducklake external maintenance connection ready"); + Ok(executor) +} + +/// Resolves the hidden DuckLake metadata schema. +async fn resolve_metadata_schema(duckdb: &DuckDbMaintenanceExecutor) -> EtlResult { + duckdb.run(resolve_ducklake_metadata_schema_blocking).await +} + +/// Lists active DuckLake table names from the metadata catalog. +async fn list_ducklake_tables( + metadata_pg_pool: &PgPool, + metadata_schema: &str, +) -> EtlResult> { + let sql = format!( + "SELECT table_name FROM {}.{} WHERE end_snapshot IS NULL ORDER BY table_name", + quote_identifier(metadata_schema), + quote_identifier("ducklake_table") + ); + let rows: Vec<(String,)> = + sqlx::query_as(AssertSqlSafe(sql)).fetch_all(metadata_pg_pool).await.map_err(|source| { + etl_error!( + ErrorKind::DestinationQueryFailed, + "DuckLake table list query failed", + format!("metadata_schema={metadata_schema}"), + source: source + ) + })?; + Ok(rows.into_iter().map(|(table_name,)| table_name).collect()) +} + +/// Runs inline flush for tables that crossed the pending-inline threshold. +async fn run_inline_flush( + duckdb: &DuckDbMaintenanceExecutor, + metadata_pg_pool: &PgPool, + metadata_schema: &str, + table_names: &[String], + config: InlineFlushMaintenanceConfig, + outcome: &mut DuckLakeMaintenanceOutcome, +) -> EtlResult<()> { + info!( + min_inlined_bytes = config.min_inlined_bytes, + table_count = table_names.len(), + "ducklake inline flush evaluation started" + ); + let sampler = + DuckLakePendingInlineSizeSampler::new(metadata_schema.to_owned(), metadata_pg_pool.clone()); + for table_name in table_names { + let sizes = sampler.sample_table(table_name).await?; + if sizes.inlined_bytes < config.min_inlined_bytes { + info!( + table = %table_name, + inlined_bytes = sizes.inlined_bytes, + min_inlined_bytes = config.min_inlined_bytes, + "ducklake inline flush skipped below threshold" + ); + continue; + } + info!( + table = %table_name, + inlined_bytes = sizes.inlined_bytes, + min_inlined_bytes = config.min_inlined_bytes, + "ducklake inline flush executing" + ); + let table_name_for_query = table_name.clone(); + let rows = + duckdb.run(move |conn| flush_table_inlined_data(conn, &table_name_for_query)).await?; + outcome.inline_flush_tables = outcome.inline_flush_tables.saturating_add(1); + outcome.inline_flush_rows = outcome.inline_flush_rows.saturating_add(rows); + info!( + table = %table_name, + rows, + "ducklake inline flush completed" + ); + } + info!( + inline_flush_tables = outcome.inline_flush_tables, + inline_flush_rows = outcome.inline_flush_rows, + "ducklake inline flush evaluation finished" + ); + Ok(()) +} + +/// Runs bounded merge-adjacent-files on selected tables. +async fn run_merge_adjacent_files( + duckdb: &DuckDbMaintenanceExecutor, + metadata_pg_pool: &PgPool, + metadata_schema: &str, + table_names: &[String], + config: &MergeAdjacentFilesMaintenanceConfig, + outcome: &mut DuckLakeMaintenanceOutcome, +) -> EtlResult<()> { + info!( + max_compacted_files = config.max_compacted_files, + max_tables_per_run = config.max_tables_per_run, + table_count = table_names.len(), + "ducklake merge-adjacent-files evaluation started" + ); + let selected = select_merge_tables( + metadata_pg_pool, + metadata_schema, + table_names, + config.max_tables_per_run, + ) + .await?; + info!( + selected_tables = ?selected, + selected_count = selected.len(), + "ducklake merge-adjacent-files selected tables" + ); + for table_name in selected { + info!( + table = %table_name, + max_compacted_files = config.max_compacted_files, + "ducklake merge-adjacent-files executing" + ); + let table_name_for_query = table_name.clone(); + let max_compacted_files = config.max_compacted_files; + let files_created = duckdb + .run(move |conn| merge_adjacent_files(conn, &table_name_for_query, max_compacted_files)) + .await?; + outcome.merge_adjacent_files_tables = outcome.merge_adjacent_files_tables.saturating_add(1); + outcome.merge_adjacent_files_created = + outcome.merge_adjacent_files_created.saturating_add(files_created); + info!( + table = %table_name, + files_created, + "ducklake merge-adjacent-files completed" + ); + } + Ok(()) +} + +/// Runs DuckLake's whole-lake adjacent-file merge before rewrite-data-files. +async fn merge_adjacent_files_for_rewrite(duckdb: &DuckDbMaintenanceExecutor) -> EtlResult<()> { + let sql = format!("CALL ducklake_merge_adjacent_files({});", quote_literal(LAKE_CATALOG)); + info!( + sql = %sql, + "ducklake rewrite-triggered merge-adjacent-files executing" + ); + duckdb + .run(move |conn| { + conn.execute_batch(&sql).map_err(|source| { + etl_error!( + ErrorKind::DestinationQueryFailed, + "DuckLake rewrite-triggered merge adjacent files failed", + format_query_error_detail(&sql), + source: source + ) + })?; + Ok(()) + }) + .await?; + info!("ducklake rewrite-triggered merge-adjacent-files completed"); + Ok(()) +} + +/// Runs bounded rewrite-data-files on selected tables. +async fn run_rewrite_data_files( + duckdb: &DuckDbMaintenanceExecutor, + metadata_pg_pool: &PgPool, + metadata_schema: &str, + table_names: &[String], + config: RewriteDataFilesMaintenanceConfig, + outcome: &mut DuckLakeMaintenanceOutcome, +) -> EtlResult<()> { + info!( + min_active_data_files = config.min_active_data_files, + max_tables_per_run = config.max_tables_per_run, + table_count = table_names.len(), + "ducklake rewrite-data-files evaluation started" + ); + let selected = select_rewrite_tables( + metadata_pg_pool, + metadata_schema, + table_names, + config.min_active_data_files, + config.max_tables_per_run, + ) + .await?; + info!( + selected_tables = ?selected, + selected_count = selected.len(), + "ducklake rewrite-data-files selected tables" + ); + for table_name in selected { + info!( + table = %table_name, + "ducklake rewrite-data-files executing" + ); + let table_name_for_query = table_name.clone(); + let files_created = + duckdb.run(move |conn| rewrite_data_files(conn, &table_name_for_query)).await?; + outcome.rewrite_data_files_tables = outcome.rewrite_data_files_tables.saturating_add(1); + outcome.rewrite_data_files_created = + outcome.rewrite_data_files_created.saturating_add(files_created); + info!( + table = %table_name, + files_created, + "ducklake rewrite-data-files completed" + ); + } + Ok(()) +} + +/// Runs DuckLake snapshot expiration. +async fn run_expire_snapshots( + duckdb: &DuckDbMaintenanceExecutor, + config: &ExpireSnapshotsMaintenanceConfig, + outcome: &mut DuckLakeMaintenanceOutcome, +) -> EtlResult<()> { + info!( + older_than = %config.older_than, + "ducklake expire-snapshots executing" + ); + let older_than = config.older_than.clone(); + let expired_snapshots = duckdb.run(move |conn| expire_snapshots(conn, &older_than)).await?; + outcome.expired_snapshots = outcome.expired_snapshots.saturating_add(expired_snapshots); + info!( + older_than = %config.older_than, + expired_snapshots, + "ducklake expire-snapshots completed" + ); + Ok(()) +} + +/// Runs DuckLake old-file cleanup. +async fn run_cleanup_old_files( + duckdb: &DuckDbMaintenanceExecutor, + config: &CleanupOldFilesMaintenanceConfig, + outcome: &mut DuckLakeMaintenanceOutcome, +) -> EtlResult<()> { + info!( + older_than = %config.older_than, + "ducklake cleanup-old-files executing" + ); + let older_than = config.older_than.clone(); + let cleaned_up_files = duckdb.run(move |conn| cleanup_old_files(conn, &older_than)).await?; + outcome.cleaned_up_files = outcome.cleaned_up_files.saturating_add(cleaned_up_files); + info!( + older_than = %config.older_than, + cleaned_up_files, + "ducklake cleanup-old-files completed" + ); + Ok(()) +} + +/// Selects tables with small-file pressure for merge-adjacent-files. +async fn select_merge_tables( + metadata_pg_pool: &PgPool, + metadata_schema: &str, + table_names: &[String], + max_tables_per_run: u32, +) -> EtlResult> { + let mut selected = Vec::new(); + for table_name in table_names { + if is_etl_internal_table(table_name) { + info!( + table = %table_name, + "ducklake rewrite-data-files table skipped because it is internal ETL metadata" + ); + continue; + } + + let metrics = + query_table_storage_metrics(metadata_pg_pool, metadata_schema, table_name).await?; + if metrics.active_data_files > 1 && metrics.small_file_ratio() > 0.0 { + info!( + table = %table_name, + active_data_files = metrics.active_data_files, + small_file_ratio = metrics.small_file_ratio(), + "ducklake merge-adjacent-files table selected" + ); + selected.push(table_name.clone()); + } else { + info!( + table = %table_name, + active_data_files = metrics.active_data_files, + small_file_ratio = metrics.small_file_ratio(), + "ducklake merge-adjacent-files table skipped" + ); + } + if selected.len() >= max_tables_per_run as usize { + break; + } + } + Ok(selected) +} + +/// Selects tables with delete pressure for rewrite-data-files. +async fn select_rewrite_tables( + metadata_pg_pool: &PgPool, + metadata_schema: &str, + table_names: &[String], + min_active_data_files: i64, + max_tables_per_run: u32, +) -> EtlResult> { + let mut selected = Vec::new(); + for table_name in table_names { + if is_etl_internal_table(table_name) { + info!( + table = %table_name, + "ducklake rewrite-data-files table skipped because it is internal ETL metadata" + ); + continue; + } + + let metrics = + query_table_storage_metrics(metadata_pg_pool, metadata_schema, table_name).await?; + if should_rewrite(&metrics, min_active_data_files) { + info!( + table = %table_name, + active_data_files = metrics.active_data_files, + active_delete_files = metrics.active_delete_files, + deleted_row_ratio = metrics.deleted_row_ratio(), + min_active_data_files, + "ducklake rewrite-data-files table selected" + ); + selected.push(table_name.clone()); + } else { + info!( + table = %table_name, + active_data_files = metrics.active_data_files, + active_delete_files = metrics.active_delete_files, + deleted_row_ratio = metrics.deleted_row_ratio(), + min_active_data_files, + "ducklake rewrite-data-files table skipped" + ); + } + if selected.len() >= max_tables_per_run as usize { + break; + } + } + Ok(selected) +} + +fn is_etl_internal_table(table_name: &str) -> bool { + table_name.starts_with("__etl_") +} + +/// Returns whether a table should be rewritten. +fn should_rewrite(metrics: &DuckLakeTableStorageMetrics, min_active_data_files: i64) -> bool { + metrics.active_data_files > min_active_data_files +} + +/// Flushes inlined user data for one table. +pub(super) fn flush_table_inlined_data( + conn: &duckdb::Connection, + table_name: &str, +) -> EtlResult { + let flush_started = std::time::Instant::now(); + let sql = format!( + r#"SELECT COALESCE(SUM(rows_flushed), 0) + FROM ducklake_flush_inlined_data({}, table_name => {});"#, + quote_literal(LAKE_CATALOG), + quote_literal(table_name), + ); + let rows_flushed: i64 = conn.query_row(&sql, [], |row| row.get(0)).map_err(|source| { + etl_error!( + ErrorKind::DestinationQueryFailed, + "DuckLake inlined data flush failed", + format_query_error_detail(&sql), + source: source + ) + })?; + let rows_flushed = rows_flushed.max(0) as u64; + let flush_result = if rows_flushed > 0 { "flushed" } else { "noop" }; + histogram!( + ETL_DUCKLAKE_INLINE_FLUSH_ROWS, + RESULT_LABEL => flush_result, + ) + .record(rows_flushed as f64); + histogram!( + ETL_DUCKLAKE_INLINE_FLUSH_DURATION_SECONDS, + RESULT_LABEL => flush_result, + ) + .record(flush_started.elapsed().as_secs_f64()); + + if rows_flushed > 0 { + debug!( + table = %table_name, + rows_flushed, + "ducklake inlined data flushed" + ); + } else { + debug!( + table = %table_name, + "ducklake inlined data already flushed" + ); + } + Ok(rows_flushed) +} + +/// Calls DuckLake merge-adjacent-files for one table. +fn merge_adjacent_files( + conn: &duckdb::Connection, + table_name: &str, + max_compacted_files: u32, +) -> EtlResult { + let sql = format!( + "SELECT COALESCE(SUM(files_created), 0) FROM ducklake_merge_adjacent_files({}, {}, \ + max_compacted_files => {});", + quote_literal(LAKE_CATALOG), + quote_literal(table_name), + max_compacted_files + ); + count_maintenance_files(conn, &sql, "DuckLake merge adjacent files failed") +} + +/// Calls DuckLake rewrite-data-files for one table. +fn rewrite_data_files(conn: &duckdb::Connection, table_name: &str) -> EtlResult { + let sql = format!( + "CALL ducklake_rewrite_data_files({}, {});", + quote_literal(LAKE_CATALOG), + quote_literal(table_name) + ); + conn.execute_batch(&sql).map_err(|source| { + etl_error!( + ErrorKind::DestinationQueryFailed, + "DuckLake rewrite data files failed", + format_query_error_detail(&sql), + source: source + ) + })?; + Ok(0) +} + +/// Calls DuckLake snapshot expiration. +fn expire_snapshots(conn: &duckdb::Connection, older_than: &str) -> EtlResult { + let sql = format!( + "CALL ducklake_expire_snapshots({}, older_than => CAST(now() AS TIMESTAMP) - CAST({} AS \ + INTERVAL));", + quote_literal(LAKE_CATALOG), + quote_literal(older_than), + ); + count_maintenance_rows(conn, &sql, "DuckLake expire snapshots failed") +} + +/// Calls DuckLake old-file cleanup. +fn cleanup_old_files(conn: &duckdb::Connection, older_than: &str) -> EtlResult { + let sql = format!( + "CALL ducklake_cleanup_old_files({}, older_than => CAST(now() AS TIMESTAMP) - CAST({} AS \ + INTERVAL));", + quote_literal(LAKE_CATALOG), + quote_literal(older_than), + ); + count_maintenance_rows(conn, &sql, "DuckLake cleanup old files failed") +} + +/// Counts rows returned by one DuckLake maintenance call. +fn count_maintenance_rows( + conn: &duckdb::Connection, + sql: &str, + description: &'static str, +) -> EtlResult { + let mut statement = conn.prepare(sql).map_err(|source| { + etl_error!( + ErrorKind::DestinationQueryFailed, + description, + format_query_error_detail(sql), + source: source + ) + })?; + let mut rows = statement.query([]).map_err(|source| { + etl_error!( + ErrorKind::DestinationQueryFailed, + description, + format_query_error_detail(sql), + source: source + ) + })?; + let mut count = 0u64; + + while let Some(_row) = rows.next().map_err(|source| { + etl_error!( + ErrorKind::DestinationQueryFailed, + description, + format_query_error_detail(sql), + source: source + ) + })? { + count = count.saturating_add(1); + } + + Ok(count) +} + +/// Counts files returned by one DuckLake maintenance function. +fn count_maintenance_files( + conn: &duckdb::Connection, + sql: &str, + description: &'static str, +) -> EtlResult { + let files_created: i64 = conn.query_row(sql, [], |row| row.get(0)).map_err(|source| { + etl_error!( + ErrorKind::DestinationQueryFailed, + description, + format_query_error_detail(sql), + source: source + ) + })?; + Ok(files_created.max(0) as u64) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn metrics( + active_data_files: i64, + active_delete_files: i64, + deleted_rows: i64, + ) -> DuckLakeTableStorageMetrics { + DuckLakeTableStorageMetrics { + active_data_files, + active_data_bytes: 100, + small_data_files: 0, + active_data_rows: 100, + active_delete_files, + active_delete_bytes: 10, + deleted_rows, + } + } + + #[test] + fn should_rewrite_requires_only_file_count() { + assert!(!should_rewrite(&metrics(39, 1, 50), 40)); + assert!(!should_rewrite(&metrics(40, 1, 50), 40)); + assert!(should_rewrite(&metrics(41, 0, 0), 40)); + } + + #[test] + fn outcome_reports_applied_work() { + assert!(!DuckLakeMaintenanceOutcome::default().applied()); + assert!( + DuckLakeMaintenanceOutcome { + inline_flush_rows: 1, + ..DuckLakeMaintenanceOutcome::default() + } + .applied() + ); + assert!( + DuckLakeMaintenanceOutcome { + rewrite_data_files_tables: 1, + ..DuckLakeMaintenanceOutcome::default() + } + .applied() + ); + assert!( + DuckLakeMaintenanceOutcome { + expired_snapshots: 1, + ..DuckLakeMaintenanceOutcome::default() + } + .applied() + ); + assert!( + DuckLakeMaintenanceOutcome { + cleaned_up_files: 1, + ..DuckLakeMaintenanceOutcome::default() + } + .applied() + ); + } +} diff --git a/crates/etl-destinations/src/ducklake/metrics.rs b/crates/etl-destinations/src/ducklake/metrics.rs index a41125f73..d7358273d 100644 --- a/crates/etl-destinations/src/ducklake/metrics.rs +++ b/crates/etl-destinations/src/ducklake/metrics.rs @@ -13,19 +13,13 @@ use parking_lot::Mutex; use pg_escape::{quote_identifier, quote_literal}; use sqlx::{AssertSqlSafe, PgPool}; use tokio::{ - sync::{mpsc, watch}, + sync::watch, task::JoinHandle, time::{Duration, Instant, MissedTickBehavior}, }; use tracing::{info, warn}; -use crate::ducklake::{ - DuckLakeTableName, LAKE_CATALOG, - client::format_query_error_detail, - maintenance::{ - TableMaintenanceNotification, TableMetricsSample, send_maintenance_notification, - }, -}; +use crate::ducklake::{DuckLakeTableName, LAKE_CATALOG, client::format_query_error_detail}; static REGISTER_METRICS: Once = Once::new(); @@ -54,14 +48,10 @@ pub(crate) const ETL_DUCKLAKE_INLINE_FLUSH_DURATION_SECONDS: &str = pub(crate) const ETL_DUCKLAKE_RETRIES_TOTAL: &str = "etl_ducklake_retries_total"; pub(crate) const ETL_DUCKLAKE_FAILED_BATCHES_TOTAL: &str = "etl_ducklake_failed_batches_total"; pub(crate) const ETL_DUCKLAKE_REPLAYED_BATCHES_TOTAL: &str = "etl_ducklake_replayed_batches_total"; -pub(crate) const ETL_DUCKLAKE_MAINTENANCE_STARTED_TOTAL: &str = - "etl_ducklake_maintenance_started_total"; -pub(crate) const ETL_DUCKLAKE_MAINTENANCE_IN_PROGRESS: &str = - "etl_ducklake_maintenance_in_progress"; -pub(crate) const ETL_DUCKLAKE_MAINTENANCE_DURATION_SECONDS: &str = - "etl_ducklake_maintenance_duration_seconds"; -pub(crate) const ETL_DUCKLAKE_MAINTENANCE_SKIPPED_TOTAL: &str = - "etl_ducklake_maintenance_skipped_total"; +pub(crate) const ETL_DUCKLAKE_EXTERNAL_MAINTENANCE_PAUSE_DURATION_SECONDS: &str = + "etl_ducklake_external_maintenance_pause_duration_seconds"; +pub(crate) const ETL_DUCKLAKE_EXTERNAL_MAINTENANCE_TRIGGERED_TOTAL: &str = + "etl_ducklake_external_maintenance_triggered_total"; pub(crate) const ETL_DUCKLAKE_TABLE_ACTIVE_DATA_FILES: &str = "etl_ducklake_table_active_data_files"; pub(crate) const ETL_DUCKLAKE_TABLE_ACTIVE_DATA_BYTES: &str = @@ -96,7 +86,6 @@ pub(crate) const PREPARED_ROWS_KIND_LABEL: &str = "prepared_rows_kind"; pub(crate) const DELETE_ORIGIN_LABEL: &str = "delete_origin"; pub(crate) const RETRY_SCOPE_LABEL: &str = "retry_scope"; pub(crate) const RESULT_LABEL: &str = "result"; -pub(crate) const MAINTENANCE_TASK_LABEL: &str = "task"; pub(crate) const MAINTENANCE_OPERATION_LABEL: &str = "operation"; pub(crate) const MAINTENANCE_REASON_LABEL: &str = "reason"; pub(crate) const MAINTENANCE_OUTCOME_LABEL: &str = "outcome"; @@ -234,30 +223,17 @@ pub(crate) fn register_metrics() { "DuckLake batches skipped because an applied marker already existed, labeled by \ batch_kind." ); - describe_counter!( - ETL_DUCKLAKE_MAINTENANCE_STARTED_TOTAL, - Unit::Count, - "DuckLake background maintenance operations started, labeled by task, operation, and \ - reason." - ); - describe_gauge!( - ETL_DUCKLAKE_MAINTENANCE_IN_PROGRESS, - Unit::Count, - "DuckLake background maintenance operations currently in progress, labeled by task, \ - operation, and reason." - ); describe_histogram!( - ETL_DUCKLAKE_MAINTENANCE_DURATION_SECONDS, + ETL_DUCKLAKE_EXTERNAL_MAINTENANCE_PAUSE_DURATION_SECONDS, Unit::Seconds, - "Duration of DuckLake background maintenance operations, labeled by task, operation, \ - reason, and outcome. Use the histogram count as the event count for non-skipped \ - outcomes." + "Duration that DuckLake foreground ingestion was paused for an external maintenance \ + run, labeled by outcome." ); describe_counter!( - ETL_DUCKLAKE_MAINTENANCE_SKIPPED_TOTAL, + ETL_DUCKLAKE_EXTERNAL_MAINTENANCE_TRIGGERED_TOTAL, Unit::Count, - "DuckLake background maintenance operations skipped because execution was deferred, \ - labeled by task, operation, and reason." + "External DuckLake maintenance operations requested by the replicator after catalog \ + sampling, labeled by operation and reason." ); describe_histogram!( @@ -344,14 +320,12 @@ pub(super) fn spawn_ducklake_metrics_sampler( metadata_schema: String, metadata_pg_pool: PgPool, created_tables: Arc>>, - maintenance_notification_tx: mpsc::Sender, ) -> EtlResult { let (shutdown_tx, shutdown_rx) = watch::channel(()); let handle = tokio::spawn(run_ducklake_metrics_sampler( metadata_schema, metadata_pg_pool, created_tables, - maintenance_notification_tx, shutdown_rx, )); @@ -363,7 +337,6 @@ async fn run_ducklake_metrics_sampler( metadata_schema: String, metadata_pg_pool: PgPool, created_tables: Arc>>, - maintenance_notification_tx: mpsc::Sender, mut shutdown_rx: watch::Receiver<()>, ) { let mut interval = @@ -399,7 +372,6 @@ async fn run_ducklake_metrics_sampler( &metadata_pg_pool, &metadata_schema, table_name.clone(), - &maintenance_notification_tx, ) .await { @@ -420,7 +392,6 @@ async fn record_table_storage_metrics( metadata_pg_pool: &PgPool, metadata_schema: &str, table_name: DuckLakeTableName, - maintenance_notification_tx: &mpsc::Sender, ) -> EtlResult<()> { let metrics = query_table_storage_metrics(metadata_pg_pool, metadata_schema, &table_name).await?; @@ -447,15 +418,6 @@ async fn record_table_storage_metrics( histogram!(ETL_DUCKLAKE_TABLE_ACTIVE_DELETE_FILES).record(active_delete_files); histogram!(ETL_DUCKLAKE_TABLE_ACTIVE_DELETE_BYTES).record(active_delete_bytes); histogram!(ETL_DUCKLAKE_TABLE_DELETED_ROW_RATIO).record(deleted_row_ratio); - send_maintenance_notification( - maintenance_notification_tx, - TableMaintenanceNotification::TableMetricsSample(TableMetricsSample { - table_name, - sampled_at: Instant::now(), - metrics, - }), - ) - .await; Ok(()) } diff --git a/crates/etl-destinations/src/ducklake/mod.rs b/crates/etl-destinations/src/ducklake/mod.rs index e7519d121..a3fd5ac7c 100644 --- a/crates/etl-destinations/src/ducklake/mod.rs +++ b/crates/etl-destinations/src/ducklake/mod.rs @@ -3,8 +3,9 @@ mod client; mod config; mod core; mod encoding; +mod external_maintenance; mod inline_size; -mod maintenance; +mod maintenance_runner; mod metrics; mod schema; @@ -18,10 +19,12 @@ pub(super) type DuckLakeTableName = String; /// /// This applies to every DuckDB connection in the destination pool so small /// writes inline into the DuckLake metadata first and can later be -/// materialized to Parquet by the background maintenance worker. +/// materialized to Parquet by an external maintenance job. const ATTACH_DATA_INLINING_ROW_LIMIT: u64 = 10_000; -pub use core::{DuckLakeDestination, table_name_to_ducklake_table_name}; +pub use core::{ + DuckLakeDestination, DuckLakeExternalMaintenancePause, table_name_to_ducklake_table_name, +}; #[cfg(feature = "test-utils")] pub use core::{ arm_pause_next_streaming_write_for_tests, release_paused_streaming_write_for_tests, @@ -35,3 +38,8 @@ pub use batches::{ reset_ducklake_test_hooks, }; pub use config::S3Config; +pub use maintenance_runner::{ + CleanupOldFilesMaintenanceConfig, DuckLakeMaintenanceConfig, DuckLakeMaintenanceOutcome, + ExpireSnapshotsMaintenanceConfig, InlineFlushMaintenanceConfig, + MergeAdjacentFilesMaintenanceConfig, RewriteDataFilesMaintenanceConfig, run_maintenance_once, +}; diff --git a/crates/etl-examples/Cargo.toml b/crates/etl-examples/Cargo.toml index fd5fa7653..4eb5185b2 100644 --- a/crates/etl-examples/Cargo.toml +++ b/crates/etl-examples/Cargo.toml @@ -21,6 +21,8 @@ path = "src/bin/ducklake.rs" clap = { workspace = true, default-features = true, features = ["std", "derive"] } etl = { workspace = true } etl-destinations = { workspace = true, features = ["bigquery", "clickhouse", "ducklake"] } +etl-telemetry = { workspace = true } +k8s-openapi = { workspace = true, features = ["latest"] } rustls = { workspace = true, features = ["aws-lc-rs", "logging"] } tokio = { workspace = true, features = ["macros", "signal"] } tracing = { workspace = true, default-features = true } diff --git a/crates/etl-replicator/Cargo.toml b/crates/etl-replicator/Cargo.toml index 40ec462de..0bf1782c2 100644 --- a/crates/etl-replicator/Cargo.toml +++ b/crates/etl-replicator/Cargo.toml @@ -13,11 +13,13 @@ egress = ["etl/egress", "etl-destinations/egress"] [dependencies] +chrono = { workspace = true } configcat = { workspace = true } etl = { workspace = true } etl-config = { workspace = true, features = ["supabase"] } etl-destinations = { workspace = true, features = ["bigquery", "clickhouse", "ducklake", "iceberg"] } etl-telemetry = { workspace = true } +k8s-openapi = { workspace = true, features = ["latest"] } metrics = { workspace = true } reqwest = { workspace = true, features = ["rustls-tls", "json"] } rustls = { workspace = true, features = ["aws-lc-rs", "logging"] } @@ -29,6 +31,7 @@ sqlx = { workspace = true, features = ["runtime-tokio", "tls-rustls", "postgres" tokio = { workspace = true, features = ["rt-multi-thread", "macros", "signal"] } tracing = { workspace = true, default-features = true } +tracing-subscriber = { workspace = true, features = ["ansi", "env-filter", "fmt"] } [target.'cfg(not(target_env = "msvc"))'.dependencies] tikv-jemalloc-ctl = { workspace = true } diff --git a/crates/etl-replicator/Dockerfile b/crates/etl-replicator/Dockerfile index e71a40143..292463d52 100644 --- a/crates/etl-replicator/Dockerfile +++ b/crates/etl-replicator/Dockerfile @@ -66,7 +66,7 @@ RUN sed -i '\#"crates/xtask"#d' Cargo.toml && rm -rf crates/xtask RUN FEATURES=""; \ if [ "$ENABLE_EGRESS" = "true" ]; then FEATURES="--features egress"; fi && \ RUSTFLAGS="-C panic=abort -C link-arg=-fuse-ld=lld" cargo build --release -p etl-replicator $FEATURES && \ - strip target/release/etl-replicator + strip target/release/etl-replicator target/release/etl-ducklake-maintenance # Runtime stage with distroless for security. FROM gcr.io/distroless/cc-debian12:nonroot @@ -81,6 +81,7 @@ USER nonroot:nonroot # Copy binary; configuration must be provided at runtime via mounted volume or environment overrides. COPY --from=builder /app/target/release/etl-replicator ./etl-replicator +COPY --from=builder /app/target/release/etl-ducklake-maintenance ./etl-ducklake-maintenance COPY --from=duckdb-extensions /duckdb_extensions ./duckdb_extensions # Use exec form for proper signal handling. diff --git a/crates/etl-replicator/src/bin/etl-ducklake-maintenance.rs b/crates/etl-replicator/src/bin/etl-ducklake-maintenance.rs new file mode 100644 index 000000000..b348c5d62 --- /dev/null +++ b/crates/etl-replicator/src/bin/etl-ducklake-maintenance.rs @@ -0,0 +1,222 @@ +//! One-shot DuckLake maintenance binary for Kubernetes maintenance Jobs. + +use std::{env, error::Error, process::ExitCode}; + +use etl_config::{ + load_config, parse_ducklake_url, + shared::{DestinationConfig, ReplicatorConfig}, +}; +use etl_destinations::ducklake::{ + CleanupOldFilesMaintenanceConfig, DuckLakeMaintenanceConfig, ExpireSnapshotsMaintenanceConfig, + InlineFlushMaintenanceConfig, MergeAdjacentFilesMaintenanceConfig, + RewriteDataFilesMaintenanceConfig, S3Config as DuckLakeS3Config, run_maintenance_once, +}; +use rustls::crypto::aws_lc_rs; +use secrecy::ExposeSecret; +use tracing::info; +use tracing_subscriber::{EnvFilter, layer::SubscriberExt, util::SubscriberInitExt}; + +type MaintenanceResult = Result>; + +/// Runs the maintenance binary. +fn main() -> ExitCode { + match try_main() { + Ok(()) => ExitCode::SUCCESS, + Err(error) => { + eprintln!("ducklake maintenance failed: {error}"); + ExitCode::FAILURE + } + } +} + +/// Loads configuration and runs one maintenance attempt. +fn try_main() -> MaintenanceResult<()> { + init_crypto_provider(); + init_stdout_tracing(); + let config = load_config::()?; + config.validate()?; + + tokio::runtime::Builder::new_multi_thread().enable_all().build()?.block_on(run(config)) +} + +/// Runs one async maintenance attempt. +async fn run(config: ReplicatorConfig) -> MaintenanceResult<()> { + let DestinationConfig::Ducklake { + catalog_url, + data_path, + pool_size, + s3_access_key_id, + s3_secret_access_key, + s3_region, + s3_endpoint, + s3_url_style, + s3_use_ssl, + metadata_schema, + duckdb_memory_cache_limit, + maintenance_target_file_size, + expire_snapshots_older_than, + .. + } = config.destination + else { + return Err("etl-ducklake-maintenance requires a DuckLake destination".into()); + }; + + let s3 = match (s3_access_key_id, s3_secret_access_key) { + (Some(access_key_id), Some(secret_access_key)) => Some(DuckLakeS3Config { + access_key_id: access_key_id.expose_secret().to_owned(), + secret_access_key: secret_access_key.expose_secret().to_owned(), + region: s3_region.unwrap_or_else(|| "us-east-1".to_owned()), + endpoint: s3_endpoint, + url_style: s3_url_style.unwrap_or_else(|| "path".to_owned()), + use_ssl: s3_use_ssl.unwrap_or(false), + }), + (None, None) => None, + _ => { + return Err("ducklake s3 credentials must include both access key id and secret \ + access key" + .into()); + } + }; + + let maintenance_config = DuckLakeMaintenanceConfig { + catalog_url: parse_ducklake_url(&catalog_url)?, + data_path: parse_ducklake_url(&data_path)?, + pool_size, + s3, + metadata_schema, + duckdb_memory_cache_limit, + maintenance_target_file_size, + inline_flush: InlineFlushMaintenanceConfig { + enabled: env_bool("ETL_DUCKLAKE_MAINTENANCE__INLINE_FLUSH__ENABLED", true), + min_inlined_bytes: env_u64( + "ETL_DUCKLAKE_MAINTENANCE__INLINE_FLUSH__MIN_INLINED_BYTES", + 10_000_000, + )?, + }, + merge_adjacent_files: MergeAdjacentFilesMaintenanceConfig { + enabled: env_bool("ETL_DUCKLAKE_MAINTENANCE__MERGE_ADJACENT_FILES__ENABLED", true), + max_compacted_files: env_u32( + "ETL_DUCKLAKE_MAINTENANCE__MERGE_ADJACENT_FILES__MAX_COMPACTED_FILES", + 32, + )?, + max_tables_per_run: env_u32( + "ETL_DUCKLAKE_MAINTENANCE__MERGE_ADJACENT_FILES__MAX_TABLES_PER_RUN", + 8, + )?, + target_file_size: env::var( + "ETL_DUCKLAKE_MAINTENANCE__MERGE_ADJACENT_FILES__TARGET_FILE_SIZE", + ) + .unwrap_or_else(|_| "10MB".to_owned()), + }, + rewrite_data_files: RewriteDataFilesMaintenanceConfig { + enabled: env_bool("ETL_DUCKLAKE_MAINTENANCE__REWRITE_DATA_FILES__ENABLED", true), + min_active_data_files: env_i64( + "ETL_DUCKLAKE_MAINTENANCE__REWRITE_DATA_FILES__MIN_ACTIVE_DATA_FILES", + 40, + )?, + max_tables_per_run: env_u32( + "ETL_DUCKLAKE_MAINTENANCE__REWRITE_DATA_FILES__MAX_TABLES_PER_RUN", + 8, + )?, + }, + expire_snapshots: ExpireSnapshotsMaintenanceConfig { + enabled: env_bool("ETL_DUCKLAKE_MAINTENANCE__EXPIRE_SNAPSHOTS__ENABLED", false), + older_than: expire_snapshots_older_than.clone().unwrap_or_else(|| "7 days".to_owned()), + }, + cleanup_old_files: CleanupOldFilesMaintenanceConfig { + enabled: env_bool("ETL_DUCKLAKE_MAINTENANCE__CLEANUP_OLD_FILES__ENABLED", true), + older_than: expire_snapshots_older_than.unwrap_or_else(|| "7 days".to_owned()), + }, + }; + + info!( + pipeline_id = config.pipeline.id, + inline_flush_enabled = maintenance_config.inline_flush.enabled, + inline_flush_min_inlined_bytes = maintenance_config.inline_flush.min_inlined_bytes, + merge_adjacent_files_enabled = maintenance_config.merge_adjacent_files.enabled, + merge_adjacent_files_max_compacted_files = + maintenance_config.merge_adjacent_files.max_compacted_files, + merge_adjacent_files_max_tables_per_run = + maintenance_config.merge_adjacent_files.max_tables_per_run, + rewrite_data_files_enabled = maintenance_config.rewrite_data_files.enabled, + rewrite_data_files_min_active_data_files = + maintenance_config.rewrite_data_files.min_active_data_files, + rewrite_data_files_max_tables_per_run = + maintenance_config.rewrite_data_files.max_tables_per_run, + expire_snapshots_enabled = maintenance_config.expire_snapshots.enabled, + expire_snapshots_older_than = %maintenance_config.expire_snapshots.older_than, + cleanup_old_files_enabled = maintenance_config.cleanup_old_files.enabled, + cleanup_old_files_older_than = %maintenance_config.cleanup_old_files.older_than, + "ducklake external maintenance job starting" + ); + + let outcome = run_maintenance_once(maintenance_config).await?; + info!( + applied = outcome.applied(), + inline_flush_tables = outcome.inline_flush_tables, + inline_flush_rows = outcome.inline_flush_rows, + merge_adjacent_files_tables = outcome.merge_adjacent_files_tables, + merge_adjacent_files_created = outcome.merge_adjacent_files_created, + rewrite_data_files_tables = outcome.rewrite_data_files_tables, + rewrite_data_files_created = outcome.rewrite_data_files_created, + expired_snapshots = outcome.expired_snapshots, + cleaned_up_files = outcome.cleaned_up_files, + "ducklake external maintenance job finished" + ); + println!( + "{{\"applied\":{},\"inlineFlushRows\":{},\"mergeAdjacentFilesCreated\":{},\"\ + rewriteDataFilesCreated\":{},\"expiredSnapshots\":{},\"cleanedUpFiles\":{}}}", + outcome.applied(), + outcome.inline_flush_rows, + outcome.merge_adjacent_files_created, + outcome.rewrite_data_files_created, + outcome.expired_snapshots, + outcome.cleaned_up_files + ); + Ok(()) +} + +/// Installs the process-wide Rustls crypto provider. +fn init_crypto_provider() { + let _ = aws_lc_rs::default_provider().install_default(); +} + +/// Initializes direct stdout logging for short-lived Kubernetes Jobs. +fn init_stdout_tracing() { + let filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| { + EnvFilter::new("etl_ducklake_maintenance=info,etl_destinations::ducklake=info") + }); + let _ = tracing_subscriber::registry() + .with(filter) + .with(tracing_subscriber::fmt::layer().with_target(true)) + .try_init(); +} + +/// Reads a boolean environment variable. +fn env_bool(name: &str, default: bool) -> bool { + env::var(name).ok().and_then(|value| value.parse::().ok()).unwrap_or(default) +} + +/// Reads a `u64` environment variable. +fn env_u64(name: &str, default: u64) -> MaintenanceResult { + Ok(match env::var(name) { + Ok(value) => value.parse()?, + Err(_) => default, + }) +} + +/// Reads a `u32` environment variable. +fn env_u32(name: &str, default: u32) -> MaintenanceResult { + Ok(match env::var(name) { + Ok(value) => value.parse()?, + Err(_) => default, + }) +} + +/// Reads an `i64` environment variable. +fn env_i64(name: &str, default: i64) -> MaintenanceResult { + Ok(match env::var(name) { + Ok(value) => value.parse()?, + Err(_) => default, + }) +} From 7f5e9d8720c5f40f392739840dd4497cc3d24f99 Mon Sep 17 00:00:00 2001 From: Jordan McQueen Date: Fri, 15 May 2026 19:04:58 +0900 Subject: [PATCH 07/29] test(failpoints): remove SEND_STATUS_UPDATE_FP between restart runs (#744) --- crates/etl/tests/pipeline_with_failpoints.rs | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/crates/etl/tests/pipeline_with_failpoints.rs b/crates/etl/tests/pipeline_with_failpoints.rs index b45922269..4a0f4d28a 100644 --- a/crates/etl/tests/pipeline_with_failpoints.rs +++ b/crates/etl/tests/pipeline_with_failpoints.rs @@ -1391,8 +1391,15 @@ async fn table_schema_replication_masks_are_consistent_after_restart() { // Clear up the events. destination.clear_events().await; - // Restart the pipeline - Postgres will resend the data since we don't track - // progress exactly. + // Remove the failpoint now that run 1 has shut down. Run 1 never acked + // progress, so the slot still holds all of run 1's WAL for replay; we no + // longer need to suppress acks in run 2 (and doing so risks + // wal_sender_timeout firing under slow CI and duplicating events, which + // would break the exact-count match in wait_for_events_count below). + fail::remove(SEND_STATUS_UPDATE_FP); + + // Restart the pipeline -- Postgres will resend the data since we don't + // track progress exactly. let mut pipeline = create_pipeline( &database.config, pipeline_id, From 9fc7614ae726e638920144071afd0ed1b92a5919 Mon Sep 17 00:00:00 2001 From: Victor Farazdagi Date: Fri, 15 May 2026 13:24:35 +0300 Subject: [PATCH 08/29] feat(examples): feature-gate example destinations (#742) --- crates/etl-examples/Cargo.toml | 18 ++++- crates/etl-examples/README.md | 130 +++++++++++++++++++-------------- 2 files changed, 93 insertions(+), 55 deletions(-) diff --git a/crates/etl-examples/Cargo.toml b/crates/etl-examples/Cargo.toml index 4eb5185b2..f141c6294 100644 --- a/crates/etl-examples/Cargo.toml +++ b/crates/etl-examples/Cargo.toml @@ -11,16 +11,32 @@ homepage.workspace = true [[bin]] name = "bigquery" path = "src/bin/bigquery.rs" +required-features = ["bigquery"] +test = false + +[[bin]] +name = "clickhouse" +path = "src/bin/clickhouse.rs" +required-features = ["clickhouse"] +test = false [[bin]] name = "ducklake" path = "src/bin/ducklake.rs" +required-features = ["ducklake"] +test = false + +[features] +default = [] +bigquery = ["etl-destinations/bigquery"] +clickhouse = ["etl-destinations/clickhouse"] +ducklake = ["etl-destinations/ducklake"] [dependencies] clap = { workspace = true, default-features = true, features = ["std", "derive"] } etl = { workspace = true } -etl-destinations = { workspace = true, features = ["bigquery", "clickhouse", "ducklake"] } +etl-destinations = { workspace = true } etl-telemetry = { workspace = true } k8s-openapi = { workspace = true, features = ["latest"] } rustls = { workspace = true, features = ["aws-lc-rs", "logging"] } diff --git a/crates/etl-examples/README.md b/crates/etl-examples/README.md index 37e11fdc2..8ec6ce9e5 100644 --- a/crates/etl-examples/README.md +++ b/crates/etl-examples/README.md @@ -5,11 +5,33 @@ Postgres to various destinations using the ETL pipeline. ## Available Examples -| Example | Binary | Destination | Status | -|---------|--------|-------------|--------| -| [BigQuery](#bigquery) | `bigquery` | Google BigQuery (cloud data warehouse) | Stable | -| [ClickHouse](#clickhouse-setup) | `clickhouse` | ClickHouse (column-oriented OLAP database) | In progress | -| [DuckLake](#ducklake) | `ducklake` | DuckLake (open data lake format) | In progress | +| Example | Binary | Feature | Destination | Status | +| ------------------------------- | ------------ | ------------ | ------------------------------------------ | ----------- | +| [BigQuery](#bigquery) | `bigquery` | `bigquery` | Google BigQuery (cloud data warehouse) | Stable | +| [ClickHouse](#clickhouse-setup) | `clickhouse` | `clickhouse` | ClickHouse (column-oriented OLAP database) | In progress | +| [DuckLake](#ducklake) | `ducklake` | `ducklake` | DuckLake (open data lake format) | In progress | + +## Building and running + +Each binary is feature-gated so you only compile the dependencies you need. Some destinations (e.g. `ducklake`) pull in heavy native dependencies that can take several minutes to compile, and there's no reason to pay that cost when you only want to try BigQuery. + +### Single example + +```bash +# Build +cargo build --bin bigquery -p etl-examples --features bigquery + +# Run +cargo run --bin bigquery -p etl-examples --features bigquery -- [flags] +``` + +Replace `bigquery` with `clickhouse` or `ducklake` as needed. + +### All examples + +```bash +cargo build -p etl-examples --all-features +``` --- @@ -46,10 +68,10 @@ Replicates a Postgres publication into a **DuckLake** data lake. DuckLake separates storage into two components: -| Component | Role | Example | -|-----------|------|---------| -| **Catalog** | Metadata (tables, snapshots, stats) | PostgreSQL database | -| **Data** | Row data as Parquet files | Local directory or S3 / S3-compatible object storage | +| Component | Role | Example | +| ----------- | ----------------------------------- | ---------------------------------------------------- | +| **Catalog** | Metadata (tables, snapshots, stats) | PostgreSQL database | +| **Data** | Row data as Parquet files | Local directory or S3 / S3-compatible object storage | The destination loads the required DuckDB extensions before attaching the lake. Each batch of rows is committed as a single Parquet snapshot so the lake stays @@ -64,9 +86,9 @@ consistent and queryable at all times. 3. Every Postgres table becomes a DuckLake table. The name is derived from the source schema and table name: - | Postgres | DuckLake | - |----------|----------| - | `public.orders` | `public_orders` | + | Postgres | DuckLake | + | ------------------ | ------------------- | + | `public.orders` | `public_orders` | | `my_schema.events` | `my__schema_events` | ### Prerequisites @@ -82,7 +104,7 @@ consistent and queryable at all times. ### Run (local data) ```bash -cargo run --bin ducklake -p etl-examples -- \ +cargo run --bin ducklake -p etl-examples --features ducklake -- \ --db-host localhost \ --db-port 5432 \ --db-name mydb \ @@ -110,7 +132,7 @@ for table table1, table2; Then run the ClickHouse example: ```bash -cargo run -p etl-examples --bin clickhouse -- \ +cargo run -p etl-examples --bin clickhouse --features clickhouse -- \ --db-host localhost \ --db-port 5432 \ --db-name postgres \ @@ -141,7 +163,7 @@ This is a fuller local example that also enables a dedicated DuckDB log dump on shutdown: ```bash -cargo run --bin ducklake -p etl-examples -- \ +cargo run --bin ducklake -p etl-examples --features ducklake -- \ --db-host postgres.etl-data-plane.svc.cluster.local \ --db-port 5432 \ --db-name mydb \ @@ -174,7 +196,7 @@ DuckDB extensions into the repository and point the destination at them: ```bash ./scripts/vendor_duckdb_extensions.sh ETL_DUCKDB_EXTENSION_ROOT="$(pwd)/vendor/duckdb/extensions" \ - cargo run --bin ducklake -p etl-examples -- [flags] + cargo run --bin ducklake -p etl-examples --features ducklake -- [flags] ``` If `ETL_DUCKDB_EXTENSION_ROOT` is unset, the destination also checks the @@ -185,7 +207,7 @@ images do not need the env var because they already ship vendored extensions at ### Run (S3 / S3-compatible data) ```bash -cargo run --bin ducklake -p etl-examples -- \ +cargo run --bin ducklake -p etl-examples --features ducklake -- \ --db-host \ --db-port \ --db-name \ @@ -207,28 +229,28 @@ connection setup. ### All flags -| Flag | Default | Description | -|------|---------|-------------| -| `--db-host` | *(required)* | Postgres host | -| `--db-port` | `5432` | Postgres port | -| `--db-name` | *(required)* | Postgres database name | -| `--db-username` | *(required)* | Postgres user (must have REPLICATION) | -| `--db-password` | — | Postgres password (omit for trust auth) | -| `--catalog-url` | *(required)* | DuckLake catalog URL (`postgres://...` or `file://...`) | -| `--data-path` | *(required)* | Local path / `file://` URL or `s3://` URI for Parquet files | -| `--pool-size` | `4` | DuckDB connection pool size | -| `--max-batch-fill-duration-ms` | `5000` | Max time to wait before flushing a batch | -| `--max-table-sync-workers` | `4` | Concurrent workers during initial copy | -| `--publication` | *(required)* | Postgres publication name | -| `--s3-access-key-id` | — | S3 access key ID (required for private S3 buckets) | -| `--s3-secret-access-key` | — | S3 secret access key | -| `--s3-region` | `us-east-1` | S3 region | -| `--s3-endpoint` | — | Custom S3 endpoint, e.g. `127.0.0.1:5000/s3` for Supabase Storage | -| `--s3-url-style` | `path` | URL style: `path` (MinIO/Supabase) or `vhost` (AWS) | -| `--s3-use-ssl` | `false` | Enable TLS for the S3 connection | -| `--metadata-schema` | — | Postgres schema for DuckLake metadata tables (e.g. `ducklake`) | -| `--duckdb-log-storage-path` | — | Enables DuckDB file-backed logging for each DuckDB connection | -| `--duckdb-log-dump-path` | — | CSV file written from `duckdb_logs` during graceful shutdown | +| Flag | Default | Description | +| ------------------------------ | ------------ | ----------------------------------------------------------------- | +| `--db-host` | _(required)_ | Postgres host | +| `--db-port` | `5432` | Postgres port | +| `--db-name` | _(required)_ | Postgres database name | +| `--db-username` | _(required)_ | Postgres user (must have REPLICATION) | +| `--db-password` | — | Postgres password (omit for trust auth) | +| `--catalog-url` | _(required)_ | DuckLake catalog URL (`postgres://...` or `file://...`) | +| `--data-path` | _(required)_ | Local path / `file://` URL or `s3://` URI for Parquet files | +| `--pool-size` | `4` | DuckDB connection pool size | +| `--max-batch-fill-duration-ms` | `5000` | Max time to wait before flushing a batch | +| `--max-table-sync-workers` | `4` | Concurrent workers during initial copy | +| `--publication` | _(required)_ | Postgres publication name | +| `--s3-access-key-id` | — | S3 access key ID (required for private S3 buckets) | +| `--s3-secret-access-key` | — | S3 secret access key | +| `--s3-region` | `us-east-1` | S3 region | +| `--s3-endpoint` | — | Custom S3 endpoint, e.g. `127.0.0.1:5000/s3` for Supabase Storage | +| `--s3-url-style` | `path` | URL style: `path` (MinIO/Supabase) or `vhost` (AWS) | +| `--s3-use-ssl` | `false` | Enable TLS for the S3 connection | +| `--metadata-schema` | — | Postgres schema for DuckLake metadata tables (e.g. `ducklake`) | +| `--duckdb-log-storage-path` | — | Enables DuckDB file-backed logging for each DuckDB connection | +| `--duckdb-log-dump-path` | — | CSV file written from `duckdb_logs` during graceful shutdown | ### Query the replicated data @@ -248,7 +270,7 @@ duckdb :memory: -c " ### Verbose logging ```bash -RUST_LOG=debug cargo run --bin ducklake -p etl-examples -- [flags] +RUST_LOG=debug cargo run --bin ducklake -p etl-examples --features ducklake -- [flags] ``` --- @@ -269,7 +291,7 @@ Replicates a Postgres publication to a Google BigQuery dataset. ### Run ```bash -cargo run --bin bigquery -p etl-examples -- \ +cargo run --bin bigquery -p etl-examples --features bigquery -- \ --db-host localhost \ --db-port 5432 \ --db-name postgres \ @@ -283,16 +305,16 @@ cargo run --bin bigquery -p etl-examples -- \ ### All flags -| Flag | Default | Description | -|------|---------|-------------| -| `--db-host` | *(required)* | Postgres host | -| `--db-port` | *(required)* | Postgres port | -| `--db-name` | *(required)* | Postgres database name | -| `--db-username` | *(required)* | Postgres user | -| `--db-password` | — | Postgres password | -| `--bq-sa-key-file` | *(required)* | Path to GCP service account key JSON | -| `--bq-project-id` | *(required)* | GCP project ID | -| `--bq-dataset-id` | *(required)* | BigQuery dataset ID | -| `--max-batch-fill-duration-ms` | `5000` | Max time to wait before flushing a batch | -| `--max-table-sync-workers` | `4` | Concurrent workers during initial copy | -| `--publication` | *(required)* | Postgres publication name | +| Flag | Default | Description | +| ------------------------------ | ------------ | ---------------------------------------- | +| `--db-host` | _(required)_ | Postgres host | +| `--db-port` | _(required)_ | Postgres port | +| `--db-name` | _(required)_ | Postgres database name | +| `--db-username` | _(required)_ | Postgres user | +| `--db-password` | — | Postgres password | +| `--bq-sa-key-file` | _(required)_ | Path to GCP service account key JSON | +| `--bq-project-id` | _(required)_ | GCP project ID | +| `--bq-dataset-id` | _(required)_ | BigQuery dataset ID | +| `--max-batch-fill-duration-ms` | `5000` | Max time to wait before flushing a batch | +| `--max-table-sync-workers` | `4` | Concurrent workers during initial copy | +| `--publication` | _(required)_ | Postgres publication name | From 26e27c2e561d99cb4a2e1b23305e7ed39fb17517 Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Mon, 18 May 2026 09:17:53 +0200 Subject: [PATCH 09/29] feat(core): Implement new durable progress (#733) --- crates/etl-postgres/src/replication/mod.rs | 1 + .../etl-postgres/src/replication/progress.rs | 184 +++++++ crates/etl-replicator/src/error_reporting.rs | 19 +- ...260511090000_replication_progress.down.sql | 2 + ...20260511090000_replication_progress.up.sql | 32 ++ crates/etl/src/destination/async_result.rs | 4 + crates/etl/src/failpoints.rs | 1 + crates/etl/src/lib.rs | 6 +- crates/etl/src/replication/apply.rs | 497 ++++-------------- crates/etl/src/replication/client.rs | 67 +-- crates/etl/src/replication/mod.rs | 2 + crates/etl/src/replication/table_sync.rs | 41 +- crates/etl/src/replication/worker_type.rs | 49 ++ crates/etl/src/store/both/memory.rs | 39 +- crates/etl/src/store/both/postgres.rs | 77 ++- crates/etl/src/store/schema/table.rs | 16 +- crates/etl/src/store/state/base.rs | 35 +- crates/etl/src/test_utils/notifying_store.rs | 37 ++ crates/etl/src/workers/apply.rs | 62 ++- crates/etl/src/workers/table_sync.rs | 2 +- crates/etl/tests/pipeline_with_failpoints.rs | 7 + crates/etl/tests/postgres_store.rs | 49 ++ 22 files changed, 738 insertions(+), 491 deletions(-) create mode 100644 crates/etl-postgres/src/replication/progress.rs create mode 100644 crates/etl/migrations/postgres_store/20260511090000_replication_progress.down.sql create mode 100644 crates/etl/migrations/postgres_store/20260511090000_replication_progress.up.sql create mode 100644 crates/etl/src/replication/worker_type.rs diff --git a/crates/etl-postgres/src/replication/mod.rs b/crates/etl-postgres/src/replication/mod.rs index 18303d3f2..322bb5c24 100644 --- a/crates/etl-postgres/src/replication/mod.rs +++ b/crates/etl-postgres/src/replication/mod.rs @@ -2,6 +2,7 @@ mod db; pub mod destination_metadata; pub mod health; pub mod lag; +pub mod progress; pub mod schema; pub mod slots; pub mod state; diff --git a/crates/etl-postgres/src/replication/progress.rs b/crates/etl-postgres/src/replication/progress.rs new file mode 100644 index 000000000..203d79785 --- /dev/null +++ b/crates/etl-postgres/src/replication/progress.rs @@ -0,0 +1,184 @@ +use std::str::FromStr; + +use sqlx::{PgExecutor, postgres::types::Oid as SqlxTableId}; +use tokio_postgres::types::PgLsn; + +use crate::types::TableId; + +/// Parses a `pg_lsn` string returned by SQLx. +fn parse_lsn(lsn: &str) -> sqlx::Result { + PgLsn::from_str(lsn).map_err(|_| { + sqlx::Error::Protocol(format!( + "Invalid pg_lsn value returned from etl.replication_progress: {lsn}." + )) + }) +} + +/// Fetches durable replication progress for a pipeline worker. +pub async fn get_replication_progress<'c, E>( + executor: E, + pipeline_id: i64, + worker_type: &'static str, + table_id: Option, +) -> sqlx::Result> +where + E: PgExecutor<'c>, +{ + let flush_lsn: Option = if let Some(table_id) = table_id { + sqlx::query_scalar( + r#" + select flush_lsn::text + from etl.replication_progress + where pipeline_id = $1 + and worker_type = $2::etl.replication_worker_type + and table_id = $3 + "#, + ) + .bind(pipeline_id) + .bind(worker_type) + .bind(SqlxTableId(table_id.into_inner())) + .fetch_optional(executor) + .await? + } else { + sqlx::query_scalar( + r#" + select flush_lsn::text + from etl.replication_progress + where pipeline_id = $1 + and worker_type = $2::etl.replication_worker_type + and table_id is null + "#, + ) + .bind(pipeline_id) + .bind(worker_type) + .fetch_optional(executor) + .await? + }; + + flush_lsn.as_deref().map(parse_lsn).transpose() +} + +/// Upserts durable replication progress for a pipeline worker. +/// +/// The update is monotonic: a stale or duplicated flush LSN cannot move stored +/// progress backward. +pub async fn upsert_replication_progress<'c, E>( + executor: E, + pipeline_id: i64, + worker_type: &'static str, + table_id: Option, + flush_lsn: PgLsn, +) -> sqlx::Result +where + E: PgExecutor<'c>, +{ + let flush_lsn = flush_lsn.to_string(); + let stored_lsn: String = if let Some(table_id) = table_id { + sqlx::query_scalar( + r#" + insert into etl.replication_progress (pipeline_id, worker_type, table_id, flush_lsn) + values ($1, $2::etl.replication_worker_type, $3, $4::pg_lsn) + on conflict (pipeline_id, worker_type, table_id) where table_id is not null + do update set + flush_lsn = case + when excluded.flush_lsn > etl.replication_progress.flush_lsn + then excluded.flush_lsn + else etl.replication_progress.flush_lsn + end, + updated_at = case + when excluded.flush_lsn > etl.replication_progress.flush_lsn + then now() + else etl.replication_progress.updated_at + end + returning flush_lsn::text + "#, + ) + .bind(pipeline_id) + .bind(worker_type) + .bind(SqlxTableId(table_id.into_inner())) + .bind(flush_lsn) + .fetch_one(executor) + .await? + } else { + sqlx::query_scalar( + r#" + insert into etl.replication_progress (pipeline_id, worker_type, flush_lsn) + values ($1, $2::etl.replication_worker_type, $3::pg_lsn) + on conflict (pipeline_id, worker_type) where table_id is null + do update set + flush_lsn = case + when excluded.flush_lsn > etl.replication_progress.flush_lsn + then excluded.flush_lsn + else etl.replication_progress.flush_lsn + end, + updated_at = case + when excluded.flush_lsn > etl.replication_progress.flush_lsn + then now() + else etl.replication_progress.updated_at + end + returning flush_lsn::text + "#, + ) + .bind(pipeline_id) + .bind(worker_type) + .bind(flush_lsn) + .fetch_one(executor) + .await? + }; + + parse_lsn(&stored_lsn) +} + +/// Deletes durable replication progress for a pipeline worker. +pub async fn delete_replication_progress<'c, E>( + executor: E, + pipeline_id: i64, + worker_type: &'static str, + table_id: Option, +) -> sqlx::Result +where + E: PgExecutor<'c>, +{ + let result = if let Some(table_id) = table_id { + sqlx::query( + r#" + delete from etl.replication_progress + where pipeline_id = $1 + and worker_type = $2::etl.replication_worker_type + and table_id = $3 + "#, + ) + .bind(pipeline_id) + .bind(worker_type) + .bind(SqlxTableId(table_id.into_inner())) + .execute(executor) + .await? + } else { + sqlx::query( + r#" + delete from etl.replication_progress + where pipeline_id = $1 + and worker_type = $2::etl.replication_worker_type + and table_id is null + "#, + ) + .bind(pipeline_id) + .bind(worker_type) + .execute(executor) + .await? + }; + + Ok(result.rows_affected()) +} + +/// Deletes durable replication progress for a specific table-sync worker. +pub async fn delete_replication_progress_for_table<'c, E>( + executor: E, + pipeline_id: i64, + table_id: TableId, +) -> sqlx::Result +where + E: PgExecutor<'c>, +{ + delete_replication_progress(executor, pipeline_id, "table_sync", Some(table_id)).await +} diff --git a/crates/etl-replicator/src/error_reporting.rs b/crates/etl-replicator/src/error_reporting.rs index 16df29d4e..538d97dd0 100644 --- a/crates/etl-replicator/src/error_reporting.rs +++ b/crates/etl-replicator/src/error_reporting.rs @@ -2,6 +2,7 @@ use std::{collections::HashMap, sync::Arc}; use etl::{ error::{EtlError, EtlResult}, + replication::WorkerType, state::{ destination_metadata::{AppliedDestinationTableMetadata, DestinationTableMetadata}, table::TableReplicationPhase, @@ -11,7 +12,7 @@ use etl::{ schema::{SchemaStore, TableSchemaRetention}, state::{StateStore, TableReplicationStates}, }, - types::{SnapshotId, TableId, TableSchema}, + types::{PgLsn, SnapshotId, TableId, TableSchema}, }; use tracing::info; @@ -121,6 +122,22 @@ where self.inner.rollback_table_replication_state(table_id).await } + async fn get_replication_progress(&self, worker_type: WorkerType) -> EtlResult> { + self.inner.get_replication_progress(worker_type).await + } + + async fn upsert_replication_progress( + &self, + worker_type: WorkerType, + flush_lsn: PgLsn, + ) -> EtlResult { + self.inner.upsert_replication_progress(worker_type, flush_lsn).await + } + + async fn delete_replication_progress(&self, worker_type: WorkerType) -> EtlResult<()> { + self.inner.delete_replication_progress(worker_type).await + } + async fn get_destination_table_metadata( &self, table_id: TableId, diff --git a/crates/etl/migrations/postgres_store/20260511090000_replication_progress.down.sql b/crates/etl/migrations/postgres_store/20260511090000_replication_progress.down.sql new file mode 100644 index 000000000..c0bd78ebe --- /dev/null +++ b/crates/etl/migrations/postgres_store/20260511090000_replication_progress.down.sql @@ -0,0 +1,2 @@ +drop table if exists etl.replication_progress; +drop type if exists etl.replication_worker_type; diff --git a/crates/etl/migrations/postgres_store/20260511090000_replication_progress.up.sql b/crates/etl/migrations/postgres_store/20260511090000_replication_progress.up.sql new file mode 100644 index 000000000..2415c97c2 --- /dev/null +++ b/crates/etl/migrations/postgres_store/20260511090000_replication_progress.up.sql @@ -0,0 +1,32 @@ +-- Durable per-replication-worker progress. +-- +-- `flush_lsn` uses logical replication progress-boundary semantics: all +-- source WAL before this LSN has been durably processed by ETL. + +create type etl.replication_worker_type as enum ( + 'apply', + 'table_sync' +); + +create table etl.replication_progress ( + id bigint generated always as identity primary key, + pipeline_id bigint not null, + worker_type etl.replication_worker_type not null, + table_id oid, + flush_lsn pg_lsn not null, + created_at timestamptz not null default now(), + updated_at timestamptz not null default now(), + constraint replication_progress_worker_table_check check ( + (worker_type = 'apply' and table_id is null) + or + (worker_type = 'table_sync' and table_id is not null) + ) +); + +create unique index uq_replication_progress_pipeline_apply + on etl.replication_progress (pipeline_id, worker_type) + where table_id is null; + +create unique index uq_replication_progress_pipeline_table_sync + on etl.replication_progress (pipeline_id, worker_type, table_id) + where table_id is not null; diff --git a/crates/etl/src/destination/async_result.rs b/crates/etl/src/destination/async_result.rs index acaa5eaab..a9d6bdf85 100644 --- a/crates/etl/src/destination/async_result.rs +++ b/crates/etl/src/destination/async_result.rs @@ -58,6 +58,10 @@ pub(crate) struct DispatchMetrics { #[derive(Debug, Clone, Copy)] pub(crate) struct ApplyLoopAsyncResultMetadata { /// Commit end LSN associated with the dispatched batch, if any. + /// + /// After the destination acknowledges the batch, this becomes the durable + /// progress boundary that status updates report as both `flush_lsn` and + /// `apply_lsn`. pub commit_end_lsn: Option, /// Dispatch-time metrics for the batch. pub metrics: DispatchMetrics, diff --git a/crates/etl/src/failpoints.rs b/crates/etl/src/failpoints.rs index 06a377b27..c280f44f8 100644 --- a/crates/etl/src/failpoints.rs +++ b/crates/etl/src/failpoints.rs @@ -15,6 +15,7 @@ pub const START_TABLE_SYNC_BEFORE_DATA_SYNC_SLOT_CREATION_FP: &str = "start_table_sync.before_data_sync_slot_creation_fp"; pub const START_TABLE_SYNC_DURING_DATA_SYNC_FP: &str = "start_table_sync.during_data_sync_fp"; pub const SEND_STATUS_UPDATE_FP: &str = "send_status_update_fp"; +pub const STORE_REPLICATION_PROGRESS_FP: &str = "store_replication_progress_fp"; pub const FORCE_SCHEMA_CLEANUP_FP: &str = "force_schema_cleanup_fp"; /// Executes a configurable failpoint for testing error scenarios. diff --git a/crates/etl/src/lib.rs b/crates/etl/src/lib.rs index d7f2f17ba..b220957cf 100644 --- a/crates/etl/src/lib.rs +++ b/crates/etl/src/lib.rs @@ -41,9 +41,9 @@ //! metadata are stored. These stores are critical to a pipeline's operation, as //! they allow it to be safely paused and resumed. //! -//! The [`store::state::StateStore`] trait handles both table replication states -//! and destination table metadata, providing a single interface for all -//! state-related storage operations. +//! The [`store::state::StateStore`] trait handles table replication states, +//! durable replication progress, and destination table metadata, providing a +//! single interface for all state-related storage operations. //! //! **Note:** To pause and resume a pipeline after the process is stopped, it //! must be able to persist data durably. The crate itself provides no diff --git a/crates/etl/src/replication/apply.rs b/crates/etl/src/replication/apply.rs index bf127c7ac..fd1724c66 100644 --- a/crates/etl/src/replication/apply.rs +++ b/crates/etl/src/replication/apply.rs @@ -18,11 +18,8 @@ use std::{ }; use etl_config::shared::PipelineConfig; -use etl_postgres::{ - replication::slots::EtlReplicationSlot, - types::{ - IdentityMask, ReplicatedTableSchema, ReplicationMask, SnapshotId, TableId, TableSchema, - }, +use etl_postgres::types::{ + IdentityMask, ReplicatedTableSchema, ReplicationMask, SnapshotId, TableId, TableSchema, }; use futures::StreamExt; use metrics::{counter, histogram}; @@ -39,7 +36,9 @@ use tokio_postgres::types::PgLsn; use tracing::{debug, error, info, warn}; #[cfg(feature = "failpoints")] -use crate::failpoints::{FORCE_SCHEMA_CLEANUP_FP, SEND_STATUS_UPDATE_FP, etl_fail_point_active}; +use crate::failpoints::{ + FORCE_SCHEMA_CLEANUP_FP, STORE_REPLICATION_PROGRESS_FP, etl_fail_point_active, +}; use crate::{ bail, concurrency::{ @@ -72,7 +71,7 @@ use crate::{ ETL_TRANSACTIONS_TOTAL, OUTCOME_LABEL, WORKER_TYPE_LABEL, }, replication::{ - EventsStream, SharedTableCache, StatusUpdateResult, StatusUpdateType, + EventsStream, SharedTableCache, StatusUpdateResult, StatusUpdateType, WorkerType, client::{PgReplicationClient, PostgresConnectionUpdate}, }, state::table::{TableReplicationPhase, TableReplicationPhaseType}, @@ -107,58 +106,14 @@ const KEEP_ALIVE_DEADLINE_FRACTION: f64 = 0.6; /// deadline at that scale would make the apply loop spin sending forced keep /// alives, which is not operationally useful. We clamp to `100ms`. const MIN_KEEP_ALIVE_DEADLINE_DURATION: Duration = Duration::from_millis(100); -/// Interval for checking whether PostgreSQL reflected the shutdown flush LSN -/// in the logical slot state. -const SHUTDOWN_SLOT_CONFIRMATION_CHECK_INTERVAL: Duration = Duration::from_secs(1); /// Minimum interval between best-effort schema cleanup tasks during normal /// replication. /// -/// Cleanup is considered only after status updates, because those are the -/// points where the slot's confirmed flush LSN may have advanced. The next +/// Cleanup is considered only after status updates, because those are natural +/// progress points where durable ETL progress may have advanced. The next /// deadline is scheduled when the previous cleanup task finishes. const SCHEMA_CLEANUP_INTERVAL: Duration = Duration::from_hours(1); -/// Type of worker driving the apply loop. -#[derive(Debug, Copy, Clone)] -pub(crate) enum WorkerType { - /// The main apply worker that coordinates table sync workers. - Apply, - /// A table sync worker that synchronizes a specific table. - TableSync { - /// The table being synchronized. - table_id: TableId, - }, -} - -impl WorkerType { - /// Builds an [`EtlReplicationSlot`] for this worker type. - pub(crate) fn build_etl_replication_slot(&self, pipeline_id: u64) -> EtlReplicationSlot { - match self { - Self::Apply => EtlReplicationSlot::Apply { pipeline_id }, - Self::TableSync { table_id } => { - EtlReplicationSlot::TableSync { pipeline_id, table_id: *table_id } - } - } - } - - /// Returns a low-cardinality worker type label for metrics and tags. - pub(crate) fn as_str(self) -> &'static str { - match self { - Self::Apply => "apply", - Self::TableSync { .. } => "table_sync", - } - } -} - -impl Display for WorkerType { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - match self { - Self::Apply => write!(f, "apply"), - Self::TableSync { table_id } => write!(f, "table_sync({table_id})"), - } - } -} - /// Result type for the apply loop execution. /// /// Indicates the reason why the apply loop terminated, enabling appropriate @@ -199,10 +154,6 @@ impl ExitIntent { } /// Represents the shutdown state of the apply loop. -/// -/// Tracks the progress of a graceful shutdown, ensuring that PostgreSQL has -/// observed our flush position before we terminate to reduce replay on a -/// pipeline restart. #[derive(Debug, Clone)] pub(crate) enum ShutdownState { /// Normal operation. @@ -212,15 +163,6 @@ pub(crate) enum ShutdownState { /// No new WAL is accepted, but buffered or in-flight destination work is /// still allowed to drain. DrainingForShutdown, - /// Shutdown drain completed. - /// - /// The loop now waits for the replication slot's `confirmed_flush_lsn` to - /// reach the final shutdown flush position. - WaitingForSlotConfirmation { - /// The LSN we sent in the final status update that PostgreSQL should - /// reflect in the logical replication slot's current state. - target_flush_lsn: PgLsn, - }, } impl ShutdownState { @@ -235,9 +177,6 @@ impl Display for ShutdownState { match self { Self::NoShutdown => write!(f, "no_shutdown"), Self::DrainingForShutdown => write!(f, "draining_for_shutdown"), - Self::WaitingForSlotConfirmation { target_flush_lsn } => { - write!(f, "waiting_for_slot_confirmation({target_flush_lsn})") - } } } } @@ -321,7 +260,7 @@ impl WorkerContext { struct ReplicationProgress { /// The highest LSN received from PostgreSQL so far. last_received_lsn: PgLsn, - /// The highest commit LSN that has been durably flushed to the destination. + /// The highest LSN boundary that ETL has durably flushed to its store. last_flush_lsn: PgLsn, } @@ -345,6 +284,7 @@ impl ReplicationProgress { } debug_assert!(self.last_received_lsn >= self.last_flush_lsn); + debug_assert!(new_lsn <= self.last_received_lsn); self.last_flush_lsn = new_lsn; } @@ -423,8 +363,6 @@ struct ApplyLoopState { flush_deadline: Option, /// The deadline for the next proactive keep alive status update. keep_alive_deadline: Instant, - /// The next time shutdown should check the slot-confirmed flush LSN. - shutdown_slot_confirmation_deadline: Instant, /// Destination write result waiting to be applied to replication progress. pending_flush_result: Option>, /// The strongest exit that this apply loop invocation should eventually @@ -479,7 +417,6 @@ impl ApplyLoopState { shutdown_state: ShutdownState::NoShutdown, flush_deadline: None, keep_alive_deadline: Instant::now() + keep_alive_deadline_duration, - shutdown_slot_confirmation_deadline: Instant::now(), pending_flush_result: None, exit_intent: None, processing_paused: false, @@ -549,17 +486,6 @@ impl ApplyLoopState { self.keep_alive_deadline = Instant::now() + keep_alive_deadline_duration; } - /// Resets the shutdown slot confirmation check deadline. - fn reset_shutdown_slot_confirmation_deadline(&mut self) { - self.shutdown_slot_confirmation_deadline = - Instant::now() + SHUTDOWN_SLOT_CONFIRMATION_CHECK_INTERVAL; - } - - /// Expires the shutdown slot confirmation check deadline immediately. - fn expire_shutdown_slot_confirmation_deadline(&mut self) { - self.shutdown_slot_confirmation_deadline = Instant::now(); - } - /// Updates the last commit end LSN to track transaction boundaries. fn update_last_commit_end_lsn(&mut self, end_lsn: Option) { match (self.last_commit_end_lsn, end_lsn) { @@ -586,13 +512,20 @@ impl ApplyLoopState { !self.handling_transaction() && !self.has_unresolved_batch_work() } - /// Returns the effective flush LSN to report to PostgreSQL. + /// Returns the effective flush LSN to report to PostgreSQL and use for + /// idle coordination. /// /// When idle, returns the last received LSN since no actual flushes occur. /// Otherwise, returns the last flush LSN from completed transactions. /// - /// Note that when a transaction is now started, the last flush lsn will be - /// used, and it might jump back compared to the last received lsn that + /// Idle-only advances are intentionally not written to durable replication + /// progress. Persisted progress is a commit-boundary resume floor, while + /// this value may include keepalive-driven positions that are useful for + /// PostgreSQL feedback and table-sync coordination but not worth writing to + /// the customer database on every idle keepalive. + /// + /// Note that when a transaction is now started, the last flush LSN will be + /// used, and it might jump back compared to the last received LSN that /// we sent before, however this is fine since the status update logic /// guarantees monotonically increasing LSNs. fn effective_flush_lsn(&self) -> PgLsn { @@ -904,15 +837,6 @@ where let mut connection_updates_rx = replication_client.connection_updates_rx(); loop { - #[cfg(feature = "failpoints")] - if matches!(self.state.shutdown_state, ShutdownState::WaitingForSlotConfirmation { .. }) - && etl_fail_point_active(SEND_STATUS_UPDATE_FP) - { - warn!("not waiting for slot confirmation on shutdown due to active failpoint"); - - return Ok(self.finish_shutdown()); - } - let iteration_result = match &self.state.shutdown_state { ShutdownState::NoShutdown => { self.run_active_iteration( @@ -929,14 +853,6 @@ where ) .await } - ShutdownState::WaitingForSlotConfirmation { target_flush_lsn } => { - self.run_shutdown_wait_iteration( - events_stream.as_mut(), - &mut connection_updates_rx, - *target_flush_lsn, - ) - .await - } }; let result = iteration_result?; @@ -1069,13 +985,18 @@ where /// /// After the selected branch runs, the loop advances idle syncing state. /// Once buffered or in-flight destination work is resolved, it sends the - /// final shutdown status update and transitions to - /// [`ShutdownState::WaitingForSlotConfirmation`]. + /// final shutdown status update and exits. async fn run_draining_shutdown_iteration( &mut self, mut events_stream: Pin<&mut BackpressureStream>, connection_updates_rx: &mut watch::Receiver, ) -> EtlResult> { + if !self.state.has_unresolved_batch_work() { + self.initiate_graceful_shutdown(events_stream.as_mut()).await?; + + return Ok(Some(self.finish_shutdown())); + } + tokio::select! { biased; @@ -1113,105 +1034,17 @@ where // Try to keep advancing syncing tables whenever the system becomes idle. self.maybe_process_syncing_tables_when_idle().await?; - // Once the drain is complete, start waiting for PostgreSQL to reflect the - // final flush position in the slot's current state. if !self.state.has_unresolved_batch_work() { self.initiate_graceful_shutdown(events_stream.as_mut()).await?; - } - - Ok(None) - } - /// Runs one iteration of the final shutdown confirmation phase. - /// - /// In this phase the loop no longer accepts new replication messages and - /// no longer waits for destination flush results. - /// - /// Priority order: - /// 1. PostgreSQL connection lifecycle updates. - /// 2. Slot confirmation checks for the target shutdown flush LSN. - /// 3. Replication stream draining. - /// 4. Periodic keep alive status updates. - /// - /// `PrimaryKeepAlive.wal_end` is not used as the shutdown barrier because - /// it reports the sender's stream position, not slot-confirmed progress. - async fn run_shutdown_wait_iteration( - &mut self, - mut events_stream: Pin<&mut BackpressureStream>, - connection_updates_rx: &mut watch::Receiver, - target_flush_lsn: PgLsn, - ) -> EtlResult> { - tokio::select! { - biased; - - // PRIORITY 1: Handle PostgreSQL connection lifecycle updates. - // A closed or errored source connection always stops the loop immediately. - changed = connection_updates_rx.changed() => { - Self::handle_connection_update(changed, connection_updates_rx)?; - } - - // PRIORITY 2: Check current slot state. If it has not caught up, - // resend the same final status update and retry later. - _ = Self::wait_for_shutdown_slot_confirmation_deadline( - self.state.shutdown_slot_confirmation_deadline - ) => { - if self.try_confirm_shutdown_flush_lsn(target_flush_lsn).await { - return Ok(Some(self.finish_shutdown())); - } - - self.send_status_update( - events_stream.as_mut(), - target_flush_lsn, - true, - StatusUpdateType::ShutdownFlush, - ) - .await?; - - self.state.reset_shutdown_slot_confirmation_deadline(); - } - - // PRIORITY 3: Keep polling the copy-both stream while waiting for slot - // confirmation so that we don't create build-up on the Postgres side, even though - // we don't advance the flush lsn. - message = events_stream.next() => { - self - .handle_shutdown_wait_stream_message( - events_stream.as_mut(), - message, - target_flush_lsn, - ) - .await?; - } - - // PRIORITY 4: Resend a heartbeat once the computed keep alive deadline expires while - // shutdown is waiting for the slot confirmation barrier. This is a last-resort - // safeguard for cases where PostgreSQL keep alives stop reaching the loop, for example - // because the source has gone quiet, the stream is backpressured, or the network is - // delayed. - _ = Self::wait_for_keep_alive_deadline(self.state.keep_alive_deadline) => { - self.send_status_update( - events_stream.as_mut(), - target_flush_lsn, - true, - StatusUpdateType::ShutdownFlush, - ) - .await?; - - self.state - .reset_keep_alive_deadline(self.keep_alive_deadline_duration); - } + return Ok(Some(self.finish_shutdown())); } Ok(None) } - /// Returns the final loop result after PostgreSQL has confirmed the - /// shutdown flush LSN in the replication slot. - /// - /// By the time shutdown reaches this point, the drain phase should have - /// resolved any buffered or in-flight destination work. The loop therefore - /// reuses the recorded exit result and falls back to pausing only if no - /// exit intent was recorded unexpectedly. + /// Returns the final loop result after all buffered destination work has + /// resolved. fn finish_shutdown(&self) -> ApplyLoopResult { debug_assert!(!self.state.has_unresolved_batch_work()); @@ -1272,11 +1105,6 @@ where tokio::time::sleep_until(deadline.into()).await; } - /// Waits until the shutdown slot confirmation check deadline expires. - async fn wait_for_shutdown_slot_confirmation_deadline(deadline: Instant) { - tokio::time::sleep_until(deadline.into()).await; - } - /// Computes the keep alive deadline from PostgreSQL's `wal_sender_timeout`. /// /// PostgreSQL normally sends a keep alive after roughly half of this @@ -1290,142 +1118,18 @@ where .max(MIN_KEEP_ALIVE_DEADLINE_DURATION) } - /// Checks whether PostgreSQL reflected the final shutdown flush LSN in the - /// logical replication slot state. - /// - /// `confirmed_flush_lsn` is current slot state, not proof that PostgreSQL - /// has fsynced the slot. A source crash before the next slot save can make - /// this value retreat. - async fn try_confirm_shutdown_flush_lsn(&self, target_flush_lsn: PgLsn) -> bool { - let worker_type = self.worker_context.worker_type(); - - match self.get_actual_confirmed_flush_lsn().await { - Ok(confirmed_flush_lsn) if confirmed_flush_lsn >= target_flush_lsn => { - info!( - %worker_type, - %confirmed_flush_lsn, - %target_flush_lsn, - "confirmed shutdown flush lsn in replication slot, safe to stop pipeline", - ); - - true - } - Ok(confirmed_flush_lsn) => { - debug!( - %worker_type, - %confirmed_flush_lsn, - %target_flush_lsn, - "shutdown flush lsn not confirmed in replication slot yet", - ); - - false - } - Err(error) => { - warn!( - %worker_type, - %target_flush_lsn, - error = %error, - "failed to confirm shutdown flush lsn in replication slot; continuing to wait", - ); - - false - } - } - } - - /// Processes a replication message while shutdown is waiting for the slot - /// to confirm the final flush LSN. - async fn handle_shutdown_wait_stream_message( - &mut self, - mut events_stream: Pin<&mut BackpressureStream>, - message: Option>>, - target_flush_lsn: PgLsn, - ) -> EtlResult<()> { - let worker_type = self.worker_context.worker_type(); - - let Some(message) = message else { - warn!( - %worker_type, - "replication stream ended while waiting for slot confirmation", - ); - - bail!( - ErrorKind::SourceConnectionFailed, - "Replication stream ended while waiting for slot confirmation" - ) - }; - - match message? { - ReplicationMessage::PrimaryKeepAlive(keepalive) => { - let wal_end = PgLsn::from(keepalive.wal_end()); - self.state.replication_progress.update_last_received_lsn(wal_end); - - debug!( - %worker_type, - %wal_end, - %target_flush_lsn, - reply_requested = keepalive.reply() == 1, - "received keep alive while waiting for slot confirmation", - ); - - // `wal_end` advances only received/write progress; shutdown - // flush/apply feedback stays pinned to the final target. - self.send_status_update( - events_stream.as_mut(), - target_flush_lsn, - true, - StatusUpdateType::ShutdownFlush, - ) - .await?; - - self.state.reset_keep_alive_deadline(self.keep_alive_deadline_duration); - } - ReplicationMessage::XLogData(message) => { - // Track received/write progress without reporting this data as - // flushed or applied. - let start_lsn = PgLsn::from(message.wal_start()); - self.state.replication_progress.update_last_received_lsn(start_lsn); - - let end_lsn = PgLsn::from(message.wal_end()); - self.state.replication_progress.update_last_received_lsn(end_lsn); - - debug!( - %worker_type, - %start_lsn, - %end_lsn, - %target_flush_lsn, - "ignoring logical replication data while waiting for slot confirmation", - ); - } - _ => { - // Ignore non-keepalive messages while waiting for slot confirmation. - // These events will be replayed on restart from the confirmed LSN. - debug!( - %worker_type, - %target_flush_lsn, - "ignoring replication message while waiting for slot confirmation", - ); - } - } - - Ok(()) - } - - /// Handles a shutdown signal by transitioning to - /// [`ShutdownState::DrainingForShutdown`] or - /// [`ShutdownState::WaitingForSlotConfirmation`]. + /// Handles a shutdown signal. /// /// Shutdown stops new message intake immediately. If there is already /// buffered or in-flight destination work, the loop first drains that work /// so the best durable position can advance before sending the final - /// shutdown status update. Otherwise it transitions directly into waiting - /// for PostgreSQL to reflect the current flush position in the slot. + /// shutdown status update. Otherwise it sends that update immediately and + /// exits the loop. /// /// Note: the shutdown system is best-effort. Graceful shutdown may not - /// complete if we are blocked on non-interruptible code or if PostgreSQL - /// never confirms the final flush LSN. It is the responsibility of the - /// caller to forcefully kill the process if shutdown does not complete - /// within an acceptable timeframe. + /// complete if we are blocked on non-interruptible code. It is the + /// responsibility of the caller to forcefully kill the process if shutdown + /// does not complete within an acceptable timeframe. async fn handle_shutdown_signal( &mut self, mut events_stream: Pin<&mut BackpressureStream>, @@ -1465,31 +1169,30 @@ where info!( %worker_type, - "shutdown signal received, no unresolved work left, entering slot confirmation wait", + "shutdown signal received, no unresolved work left, sending final status update", ); self.initiate_graceful_shutdown(events_stream.as_mut()).await } - /// Initiates graceful shutdown by sending a status update and transitioning - /// to [`ShutdownState::WaitingForSlotConfirmation`]. + /// Initiates graceful shutdown by sending a final status update just to + /// make sure that Postgres can advance its state once more. /// - /// The status update uses the best durable position currently known by the - /// loop. + /// The status update uses the loop's effective flush position. When idle, + /// this may include received keepalive progress that is intentionally not + /// persisted as durable ETL progress. async fn initiate_graceful_shutdown( &mut self, mut events_stream: Pin<&mut BackpressureStream>, ) -> EtlResult<()> { let worker_type = self.worker_context.worker_type(); - // Use effective flush LSN to report last received LSN when idle, since - // last flush LSN only advances during actual flushes. let flush_lsn = self.state.effective_flush_lsn(); info!( %worker_type, %flush_lsn, - "sending shutdown status update and waiting for slot confirmation", + "sending shutdown status update", ); self.send_status_update( @@ -1500,9 +1203,7 @@ where ) .await?; - self.state.shutdown_state = - ShutdownState::WaitingForSlotConfirmation { target_flush_lsn: flush_lsn }; - self.state.expire_shutdown_slot_confirmation_deadline(); + self.state.shutdown_state = ShutdownState::NoShutdown; Ok(()) } @@ -1514,7 +1215,7 @@ where /// `force = true`. Those updates are about keeping the replication /// connection alive while the system is idle, not about advertising /// newly flushed progress. Keepalive replies from the main replication - /// stream and final shutdown confirmation still use this same helper. + /// stream and the final shutdown update still use this same helper. async fn send_status_update( &mut self, mut events_stream: Pin<&mut BackpressureStream>, @@ -1545,8 +1246,9 @@ where /// Attempts to spawn a best-effort task that prunes obsolete schema /// versions. /// - /// Cleanup reads the slot's `confirmed_flush_lsn` before choosing pruning - /// boundaries, rather than trusting the status update that was just sent. + /// Cleanup uses ETL-owned durable replication progress. If no progress row + /// exists yet, pruning is skipped instead of relying on PostgreSQL slot + /// state. async fn maybe_spawn_schema_cleanup(&mut self) -> EtlResult<()> { self.collect_finished_schema_cleanup_task().await; @@ -1560,14 +1262,24 @@ where let worker_type = self.worker_context.worker_type(); - // Use current slot state instead of the local status-update watermark. - let confirmed_flush_lsn = match self.get_actual_confirmed_flush_lsn().await { - Ok(confirmed_flush_lsn) => confirmed_flush_lsn, + let durable_flush_lsn = match self.schema_store.get_replication_progress(worker_type).await + { + Ok(Some(durable_flush_lsn)) => durable_flush_lsn, + Ok(None) => { + debug!( + %worker_type, + "skipping schema cleanup because durable replication progress is not available" + ); + + schema_cleanup_run.finish().await; + + return Ok(()); + } Err(err) => { warn!( %worker_type, error = %err, - "skipping schema cleanup because slot progress could not be confirmed" + "skipping schema cleanup because durable replication progress could not be loaded" ); schema_cleanup_run.finish().await; @@ -1581,7 +1293,7 @@ where // is spawned, so any concurrent progress can only make this cleanup // conservative. let table_schema_retentions = - match self.get_table_schema_retentions(confirmed_flush_lsn).await { + match self.get_table_schema_retentions(durable_flush_lsn).await { Ok(table_schema_retentions) => table_schema_retentions, Err(err) => { error!( @@ -1666,32 +1378,16 @@ where Ok(()) } - /// Returns the current slot-confirmed flush LSN to use as the cleanup - /// boundary. - /// - /// Schema cleanup uses this instead of the local status-update watermark - /// because pruning schemas too early can break replay. - /// - /// This is current PostgreSQL slot state, not source-crash durability. A - /// source crash before a slot save can make the value retreat. - async fn get_actual_confirmed_flush_lsn(&self) -> EtlResult { - let replication_client = - PgReplicationClient::connect_query(self.config.pg_connection.clone()).await?; - let slot = replication_client.get_slot(&self.state.slot_name).await?; - - Ok(slot.confirmed_flush_lsn) - } - /// Returns schema retention boundaries for tables this worker may clean up. /// /// The shared cache is used only to find active tables, and the worker's /// normal ownership check decides which of those tables can be considered. - /// A table's cleanup boundary is capped by the slot's current confirmed - /// flush LSN and by the earliest destination metadata snapshot that may - /// still be needed. + /// A table's cleanup boundary is capped by durable ETL replication progress + /// and by the earliest destination metadata snapshot that may still be + /// needed. async fn get_table_schema_retentions( &self, - confirmed_flush_lsn: PgLsn, + durable_flush_lsn: PgLsn, ) -> EtlResult> { let active_table_ids = self.shared_table_cache.active_table_ids().await; let mut table_schema_retentions = HashMap::with_capacity(active_table_ids.len()); @@ -1701,7 +1397,7 @@ where // flush position. This keeps table sync workers limited to their // assigned table while preserving apply worker ownership rules. let should_apply_changes = - self.should_apply_changes(table_id, confirmed_flush_lsn).await?; + self.should_apply_changes(table_id, durable_flush_lsn).await?; if !should_apply_changes { continue; @@ -1709,7 +1405,7 @@ where // We try to load the destination table metadata to see if it is referencing a // snapshot id that would be cleaned up if we were to just use the - // confirmed flush lsn as the boundary. + // durable flush lsn as the boundary. // // If there is no metadata, we play it safe and skip the pruning for this table. let Some(destination_table_metadata) = @@ -1730,12 +1426,11 @@ where .unwrap_or(destination_table_metadata.snapshot_id) .min(destination_table_metadata.snapshot_id); - // We determine whether the confirmed flush lsn or the snapshot id is the new + // We determine whether the durable flush lsn or the snapshot id is the new // retention limit. We could use a normal PgLsn type for handling // this, but to make the implementation more explicit we use an enum. - let retention = if confirmed_flush_lsn <= destination_retention_snapshot_id.into_inner() - { - TableSchemaRetention::ConfirmedFlushLsn(confirmed_flush_lsn) + let retention = if durable_flush_lsn <= destination_retention_snapshot_id.into_inner() { + TableSchemaRetention::DurableFlushLsn(durable_flush_lsn) } else { TableSchemaRetention::SnapshotId(destination_retention_snapshot_id) }; @@ -2028,8 +1723,6 @@ where self.send_status_update( events_stream.as_mut(), - // Use effective flush LSN to report last received LSN when idle, since - // last flush LSN only advances during actual flushes. self.state.effective_flush_lsn(), message.reply() == 1, StatusUpdateType::KeepAlive, @@ -2604,17 +2297,21 @@ where &mut self, last_commit_end_lsn: PgLsn, ) -> EtlResult<()> { - // Update replication progress to notify PostgreSQL of durable flush. Only - // reports progress up to the last completed transaction, which may - // cause duplicates on restart for partial transactions. Destinations - // must handle at-least-once delivery semantics. - self.state.replication_progress.update_last_flush_lsn(last_commit_end_lsn); + // Store durable replication progress at commit boundaries after the + // destination acknowledges the batch. This gives restarts a durable + // lower-bound resume point: once startup chooses this point, no event + // older than it will be emitted. It is not meant to eliminate every + // duplicate, and idle keepalive-only progress is intentionally left out + // to avoid writing to the customer database on every quiet heartbeat. + let durable_flush_lsn = + self.upsert_durable_replication_progress(last_commit_end_lsn).await?; + self.state.replication_progress.update_last_flush_lsn(durable_flush_lsn); let current_lsn = self.state.replication_progress.last_flush_lsn; info!( worker_type = %self.worker_context.worker_type(), %current_lsn, - "processing syncing tables after batch flush" + "processing syncing tables after durable batch flush" ); let exit_intent = match &mut self.worker_context { @@ -2633,6 +2330,25 @@ where Ok(()) } + /// Stores durable worker progress unless fault injection asks us to skip + /// it. + async fn upsert_durable_replication_progress(&self, flush_lsn: PgLsn) -> EtlResult { + let worker_type = self.worker_context.worker_type(); + + #[cfg(feature = "failpoints")] + if etl_fail_point_active(STORE_REPLICATION_PROGRESS_FP) { + warn!( + %worker_type, + %flush_lsn, + "not storing durable replication progress due to active failpoint" + ); + + return Ok(flush_lsn); + } + + self.schema_store.upsert_replication_progress(worker_type, flush_lsn).await + } + /// Processes syncing tables outside a transaction. /// /// Dispatches to worker-specific implementation based on the worker @@ -2642,11 +2358,11 @@ where /// of work so draining stays focused on already-started flushes and /// shutdown barriers. /// - /// Idle syncing is especially important to have the tables converge to - /// `Ready` or `SyncDone` even if there is no traffic on that specific - /// slot. In that case, the process is driven by keep alive messages - /// which will advance the `last_received_lsn` and properly use that for - /// the syncing LSN to establish progress. + /// Idle syncing uses the effective flush LSN so table handoff can make + /// progress even when no new transactions arrive. Keepalive messages may + /// advance that effective position, but those idle-only advances are not + /// persisted as durable replication progress because they would create + /// extra writes during normal quiet periods. async fn maybe_process_syncing_tables_when_idle(&mut self) -> EtlResult<()> { if self.state.exit_intent.is_some() { return Ok(()); @@ -2666,7 +2382,6 @@ where return Ok(()); } - // Use effective flush LSN to report last received LSN when idle. let current_lsn = self.state.effective_flush_lsn(); debug!( diff --git a/crates/etl/src/replication/client.rs b/crates/etl/src/replication/client.rs index e8353e6a4..faa67dbe5 100644 --- a/crates/etl/src/replication/client.rs +++ b/crates/etl/src/replication/client.rs @@ -12,7 +12,7 @@ use tokio_postgres::{ Client, Config, Connection, CopyOutStream, NoTls, SimpleQueryMessage, SimpleQueryRow, Socket, config::ReplicationMode, error::SqlState, tls::MakeTlsConnect, }; -use tracing::{Instrument, debug, error, info, warn}; +use tracing::{Instrument, error, info, warn}; use crate::{ bail, @@ -405,18 +405,6 @@ impl PgReplicationClient { } } - /// Establishes a regular query connection to Postgres. - /// - /// This connection does not use logical replication mode, so it does not - /// consume a `max_wal_senders` slot. It is intended for catalog checks that - /// support a running replication connection. - pub(crate) async fn connect_query(pg_connection_config: PgConnectionConfig) -> EtlResult { - match pg_connection_config.tls.enabled { - true => PgReplicationClient::connect_query_tls(pg_connection_config).await, - false => PgReplicationClient::connect_query_no_tls(pg_connection_config).await, - } - } - /// Establishes a connection to Postgres without TLS encryption. /// /// The connection is configured for logical replication mode. @@ -442,27 +430,6 @@ impl PgReplicationClient { }) } - /// Establishes a regular non-TLS query connection to Postgres. - async fn connect_query_no_tls(pg_connection_config: PgConnectionConfig) -> EtlResult { - let config: Config = pg_connection_config.clone().with_db(Some(&ETL_REPLICATION_OPTIONS)); - - let (client, connection) = config.connect(NoTls).await?; - - let server_version = - connection.parameter("server_version").and_then(extract_server_version); - - let connection_updates_rx = spawn_postgres_connection::(connection); - - debug!("connected to postgres query connection without tls"); - - Ok(PgReplicationClient { - client: Arc::new(client), - pg_connection_config: Arc::new(pg_connection_config), - server_version, - connection_updates_rx, - }) - } - /// Establishes a TLS-encrypted connection to Postgres. /// /// The connection is configured for logical replication mode. @@ -501,38 +468,6 @@ impl PgReplicationClient { }) } - /// Establishes a regular TLS-encrypted query connection to Postgres. - async fn connect_query_tls(pg_connection_config: PgConnectionConfig) -> EtlResult { - let config: Config = pg_connection_config.clone().with_db(Some(&ETL_REPLICATION_OPTIONS)); - - let mut root_store = rustls::RootCertStore::empty(); - for cert in - CertificateDer::pem_slice_iter(pg_connection_config.tls.trusted_root_certs.as_bytes()) - { - let cert = cert?; - root_store.add(cert)?; - } - - let tls_config = - ClientConfig::builder().with_root_certificates(root_store).with_no_client_auth(); - - let (client, connection) = config.connect(MakeRustlsConnect::new(tls_config)).await?; - - let server_version = - connection.parameter("server_version").and_then(extract_server_version); - - let connection_updates_rx = spawn_postgres_connection::(connection); - - debug!("connected to postgres query connection with tls"); - - Ok(PgReplicationClient { - client: Arc::new(client), - pg_connection_config: Arc::new(pg_connection_config), - server_version, - connection_updates_rx, - }) - } - /// Creates a non-replication, non-TLS child connection. async fn connect_child_no_tls(&self) -> EtlResult { let config: Config = diff --git a/crates/etl/src/replication/mod.rs b/crates/etl/src/replication/mod.rs index 727612e8a..0d48666f4 100644 --- a/crates/etl/src/replication/mod.rs +++ b/crates/etl/src/replication/mod.rs @@ -8,6 +8,7 @@ pub mod client; mod stream; mod table_cache; mod table_sync; +mod worker_type; pub(crate) use apply::{ ApplyLoop, ApplyLoopResult, ApplyWorkerContext, TableSyncWorkerContext, WorkerContext, @@ -15,3 +16,4 @@ pub(crate) use apply::{ pub(crate) use stream::{EventsStream, StatusUpdateResult, StatusUpdateType, TableCopyStream}; pub(crate) use table_cache::SharedTableCache; pub(crate) use table_sync::{TableSyncResult, start_table_sync}; +pub use worker_type::WorkerType; diff --git a/crates/etl/src/replication/table_sync.rs b/crates/etl/src/replication/table_sync.rs index 37441e16d..999840d37 100644 --- a/crates/etl/src/replication/table_sync.rs +++ b/crates/etl/src/replication/table_sync.rs @@ -21,7 +21,7 @@ use crate::{ error::{ErrorKind, EtlResult}, etl_error, metrics::{ETL_TABLE_COPY_DURATION_SECONDS, PARTITIONING_LABEL}, - replication::{client::PgReplicationClient, table_cache::SharedTableCache}, + replication::{WorkerType, client::PgReplicationClient, table_cache::SharedTableCache}, state::table::{TableReplicationPhase, TableReplicationPhaseType}, store::{schema::SchemaStore, state::StateStore}, types::PipelineId, @@ -137,9 +137,9 @@ where // have to restart // copying all the table data and delete the slot. // - `FinishedCopy` -> this means that the table was successfully copied, but we - // didn't manage to - // complete the table sync function, so we just want to continue the cdc stream - // from the slot's confirmed_flush_lsn value. + // didn't manage to complete the table sync function, so we just want to + // continue the cdc stream from durable table-sync progress when available, or + // from the slot's confirmed flush LSN otherwise. // // In case the phase is any other phase, we will return an error. let start_lsn = match phase_type { @@ -157,6 +157,7 @@ where // We try to delete the slot also during `Init` because we support state // rollback and a slot might be there from the previous run. replication_client.delete_slot_if_exists(&slot_name).await?; + store.delete_replication_progress(WorkerType::TableSync { table_id }).await?; // We must truncate the destination table before starting a copy to avoid data // inconsistencies. @@ -384,9 +385,37 @@ where } TableReplicationPhaseType::FinishedCopy => { let slot = replication_client.get_slot(&slot_name).await?; - info!(table_id = table_id.0, confirmed_flush_lsn = %slot.confirmed_flush_lsn, "resuming table sync"); + let worker_type = WorkerType::TableSync { table_id }; + let durable_flush_lsn = store.get_replication_progress(worker_type).await?; + if let Some(durable_flush_lsn) = durable_flush_lsn { + // Durable progress and slot progress can legitimately differ. During idle + // periods we keep sending PostgreSQL feedback with the received LSN, but + // we do not persist those idle-only advances to the state database to + // avoid extra customer-database writes. Conversely, durable progress can + // be ahead if ETL flushed a batch but PostgreSQL did not confirm the + // feedback yet. Startup uses the latest boundary available from either + // source as a resume floor, which guarantees no event older than the + // chosen start LSN is emitted. + let start_lsn = durable_flush_lsn.max(slot.confirmed_flush_lsn); + + info!( + table_id = table_id.0, + %durable_flush_lsn, + confirmed_flush_lsn = %slot.confirmed_flush_lsn, + %start_lsn, + "resuming table sync from durable replication progress and replication slot" + ); - slot.confirmed_flush_lsn + start_lsn + } else { + info!( + table_id = table_id.0, + confirmed_flush_lsn = %slot.confirmed_flush_lsn, + "durable table sync progress not found, using slot fallback" + ); + + slot.confirmed_flush_lsn + } } _ => unreachable!("phase type already validated above"), }; diff --git a/crates/etl/src/replication/worker_type.rs b/crates/etl/src/replication/worker_type.rs new file mode 100644 index 000000000..6f4e1fbd1 --- /dev/null +++ b/crates/etl/src/replication/worker_type.rs @@ -0,0 +1,49 @@ +use std::fmt::{Display, Formatter}; + +use etl_postgres::{replication::slots::EtlReplicationSlot, types::TableId}; + +/// Type of worker driving replication. +#[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)] +pub enum WorkerType { + /// The main apply worker that coordinates table sync workers. + Apply, + /// A table sync worker that synchronizes a specific table. + TableSync { + /// The table being synchronized. + table_id: TableId, + }, +} + +impl WorkerType { + /// Builds an [`EtlReplicationSlot`] for this worker type. + pub(crate) fn build_etl_replication_slot(&self, pipeline_id: u64) -> EtlReplicationSlot { + match self { + Self::Apply => EtlReplicationSlot::Apply { pipeline_id }, + Self::TableSync { table_id } => { + EtlReplicationSlot::TableSync { pipeline_id, table_id: *table_id } + } + } + } + + /// Returns a low-cardinality worker type label for metrics and tags. + pub(crate) fn as_str(self) -> &'static str { + match self { + Self::Apply => "apply", + Self::TableSync { .. } => "table_sync", + } + } + + /// Returns the durable progress table ID used by the state store. + pub(crate) fn progress_table_id(self) -> Option { + match self { + Self::Apply => None, + Self::TableSync { table_id } => Some(table_id), + } + } +} + +impl Display for WorkerType { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "{}", self.as_str()) + } +} diff --git a/crates/etl/src/store/both/memory.rs b/crates/etl/src/store/both/memory.rs index 735f22041..ff05beb56 100644 --- a/crates/etl/src/store/both/memory.rs +++ b/crates/etl/src/store/both/memory.rs @@ -8,6 +8,7 @@ use tokio::sync::Mutex; use crate::{ error::{ErrorKind, EtlResult}, etl_error, + replication::WorkerType, state::{ destination_metadata::{AppliedDestinationTableMetadata, DestinationTableMetadata}, table::TableReplicationPhase, @@ -17,7 +18,7 @@ use crate::{ schema::{SchemaStore, TableSchemaRetention, TableSchemaSnapshots}, state::{DestinationTablesMetadata, StateStore, TableReplicationStates}, }, - types::{SnapshotId, TableId, TableSchema}, + types::{PgLsn, SnapshotId, TableId, TableSchema}, }; /// Inner state of [`MemoryStore`] @@ -36,6 +37,8 @@ struct Inner { table_schemas: Arc, /// Cached destination table metadata indexed by table ID. destination_tables_metadata: DestinationTablesMetadata, + /// Durable replication progress indexed by worker type and optional table. + replication_progress: HashMap, } /// In-memory storage for ETL pipeline state and schema information. @@ -65,6 +68,7 @@ impl MemoryStore { table_state_history: HashMap::new(), table_schemas: Arc::new(TableSchemaSnapshots::default()), destination_tables_metadata: Arc::new(BTreeMap::new()), + replication_progress: HashMap::new(), }; Self { inner: Arc::new(Mutex::new(inner)) } @@ -146,6 +150,38 @@ impl StateStore for MemoryStore { Ok(previous_state) } + async fn get_replication_progress(&self, worker_type: WorkerType) -> EtlResult> { + let inner = self.inner.lock().await; + + Ok(inner.replication_progress.get(&worker_type).copied()) + } + + async fn upsert_replication_progress( + &self, + worker_type: WorkerType, + flush_lsn: PgLsn, + ) -> EtlResult { + let mut inner = self.inner.lock().await; + let stored_lsn = inner + .replication_progress + .entry(worker_type) + .and_modify(|stored_lsn| { + if flush_lsn > *stored_lsn { + *stored_lsn = flush_lsn; + } + }) + .or_insert(flush_lsn); + + Ok(*stored_lsn) + } + + async fn delete_replication_progress(&self, worker_type: WorkerType) -> EtlResult<()> { + let mut inner = self.inner.lock().await; + inner.replication_progress.remove(&worker_type); + + Ok(()) + } + async fn get_destination_table_metadata( &self, table_id: TableId, @@ -241,6 +277,7 @@ impl CleanupStore for MemoryStore { // Remove all schema versions for this table. Arc::make_mut(&mut inner.table_schemas).remove_table(table_id); Arc::make_mut(&mut inner.destination_tables_metadata).remove(&table_id); + inner.replication_progress.remove(&WorkerType::TableSync { table_id }); Ok(()) } diff --git a/crates/etl/src/store/both/postgres.rs b/crates/etl/src/store/both/postgres.rs index f98ab4a57..60b692df8 100644 --- a/crates/etl/src/store/both/postgres.rs +++ b/crates/etl/src/store/both/postgres.rs @@ -5,7 +5,7 @@ use std::{ time::Duration, }; -use etl_postgres::replication::{destination_metadata, schema, state}; +use etl_postgres::replication::{destination_metadata, progress, schema, state}; use metrics::gauge; use sqlx::{PgPool, postgres::PgPoolOptions}; use tokio::sync::Mutex; @@ -17,6 +17,7 @@ use crate::{ etl_error, metrics::{ETL_TABLES_TOTAL, PHASE_LABEL}, migrations, + replication::WorkerType, state::{ destination_metadata::{ AppliedDestinationTableMetadata, DestinationTableMetadata, DestinationTableSchemaStatus, @@ -28,7 +29,7 @@ use crate::{ schema::{SchemaStore, TableSchemaRetention, TableSchemaSnapshots}, state::{DestinationTablesMetadata, StateStore, TableReplicationStates}, }, - types::{PipelineId, ReplicationMask, SnapshotId, TableId, TableSchema}, + types::{PgLsn, PipelineId, ReplicationMask, SnapshotId, TableId, TableSchema}, }; /// Maximum number of connections in the pool. @@ -316,6 +317,64 @@ impl StateStore for PostgresStore { Ok(restored_phase) } + async fn get_replication_progress(&self, worker_type: WorkerType) -> EtlResult> { + progress::get_replication_progress( + &self.pool, + self.pipeline_id as i64, + worker_type.as_str(), + worker_type.progress_table_id(), + ) + .await + .map_err(|err| { + etl_error!( + ErrorKind::SourceQueryFailed, + "Replication progress loading failed", + source: err + ) + }) + } + + async fn upsert_replication_progress( + &self, + worker_type: WorkerType, + flush_lsn: PgLsn, + ) -> EtlResult { + progress::upsert_replication_progress( + &self.pool, + self.pipeline_id as i64, + worker_type.as_str(), + worker_type.progress_table_id(), + flush_lsn, + ) + .await + .map_err(|err| { + etl_error!( + ErrorKind::SourceQueryFailed, + "Replication progress storage failed", + source: err + ) + }) + } + + async fn delete_replication_progress(&self, worker_type: WorkerType) -> EtlResult<()> { + progress::delete_replication_progress( + &self.pool, + self.pipeline_id as i64, + worker_type.as_str(), + worker_type.progress_table_id(), + ) + .await + .map_err(|err| { + etl_error!( + ErrorKind::SourceQueryFailed, + "Replication progress deletion failed", + source: err + ) + })?; + + Ok(()) + } + /// Retrieves destination table metadata for a specific table from cache. /// /// This method provides fast access to destination table metadata by @@ -610,6 +669,20 @@ impl CleanupStore for PostgresStore { state::delete_replication_state_for_table(&mut *tx, self.pipeline_id as i64, table_id) .await?; + progress::delete_replication_progress_for_table( + &mut *tx, + self.pipeline_id as i64, + table_id, + ) + .await + .map_err(|err| { + etl_error!( + ErrorKind::SourceQueryFailed, + "Replication progress deletion failed", + source: err + ) + })?; + tx.commit().await?; inner.remove_table_state(table_id); diff --git a/crates/etl/src/store/schema/table.rs b/crates/etl/src/store/schema/table.rs index 56a5b0f38..ab894d054 100644 --- a/crates/etl/src/store/schema/table.rs +++ b/crates/etl/src/store/schema/table.rs @@ -8,18 +8,14 @@ use crate::types::{PgLsn, SnapshotId, TableId, TableSchema}; /// Per-table schema cleanup retention boundary. /// /// Retention can be bounded by a stored schema snapshot that the destination -/// still needs, or by the replication slot's current confirmed flush LSN. Both -/// are LSN values, but the variant records why that boundary was chosen. -/// -/// A confirmed-flush boundary is PostgreSQL's current slot state, not -/// necessarily a slot value saved to disk. +/// still needs, or by ETL-owned durable replication progress. Both are LSN +/// values, but the variant records why that boundary was chosen. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum TableSchemaRetention { /// Retain schemas according to a destination-useful schema snapshot. SnapshotId(SnapshotId), - /// Retain schemas according to the replication slot's current confirmed - /// flush LSN. - ConfirmedFlushLsn(PgLsn), + /// Retain schemas according to ETL-owned durable replication progress. + DurableFlushLsn(PgLsn), } impl TableSchemaRetention { @@ -27,7 +23,7 @@ impl TableSchemaRetention { pub fn to_lsn(self) -> PgLsn { match self { Self::SnapshotId(snapshot_id) => snapshot_id.into(), - Self::ConfirmedFlushLsn(lsn) => lsn, + Self::DurableFlushLsn(lsn) => lsn, } } } @@ -246,7 +242,7 @@ mod tests { let removed = snapshots.prune(&HashMap::from([ (table_id, TableSchemaRetention::SnapshotId(SnapshotId::from(250))), - (other_table_id, TableSchemaRetention::ConfirmedFlushLsn(SnapshotId::from(150).into())), + (other_table_id, TableSchemaRetention::DurableFlushLsn(SnapshotId::from(150).into())), ])); assert_eq!(removed, 3); diff --git a/crates/etl/src/store/state/base.rs b/crates/etl/src/store/state/base.rs index 12296c5d0..8f786d27d 100644 --- a/crates/etl/src/store/state/base.rs +++ b/crates/etl/src/store/state/base.rs @@ -2,11 +2,12 @@ use std::{collections::BTreeMap, future::Future, sync::Arc}; use crate::{ error::EtlResult, + replication::WorkerType, state::{ destination_metadata::{AppliedDestinationTableMetadata, DestinationTableMetadata}, table::TableReplicationPhase, }, - types::TableId, + types::{PgLsn, TableId}, }; /// Arc-wrapped dictionary of table replication states. @@ -15,11 +16,12 @@ pub type TableReplicationStates = Arc>; /// Arc-wrapped dictionary of destination table metadata. pub(crate) type DestinationTablesMetadata = Arc>; -/// Trait for storing and retrieving table replication state and destination -/// metadata. +/// Trait for storing and retrieving replication state, durable replication +/// progress, and destination metadata. /// /// [`StateStore`] implementations are responsible for defining how table -/// replication states and destination table metadata are stored and retrieved. +/// replication states, replication progress, and destination table metadata are +/// stored and retrieved. /// /// Implementations should ensure thread-safety and handle concurrent access to /// the data. @@ -76,6 +78,31 @@ pub trait StateStore { table_id: TableId, ) -> impl Future> + Send; + /// Returns the durable flush LSN for a replication worker, if one has been + /// stored. + fn get_replication_progress( + &self, + worker_type: WorkerType, + ) -> impl Future>> + Send; + + /// Monotonically upserts the durable flush LSN for a replication worker. + /// + /// Implementations must never move stored progress backward. The returned + /// value is the LSN stored after applying the monotonic update. + fn upsert_replication_progress( + &self, + worker_type: WorkerType, + flush_lsn: PgLsn, + ) -> impl Future> + Send; + + /// Deletes durable replication progress for a replication worker. + /// + /// This is used when the worker's slot lineage is intentionally reset. + fn delete_replication_progress( + &self, + worker_type: WorkerType, + ) -> impl Future> + Send; + /// Returns destination table metadata for a specific table from the cache. /// /// Does not load any new data into the cache. diff --git a/crates/etl/src/test_utils/notifying_store.rs b/crates/etl/src/test_utils/notifying_store.rs index adc938ec6..614f635e6 100644 --- a/crates/etl/src/test_utils/notifying_store.rs +++ b/crates/etl/src/test_utils/notifying_store.rs @@ -10,6 +10,7 @@ use tokio::sync::{Notify, RwLock}; use crate::{ error::{ErrorKind, EtlResult}, etl_error, + replication::WorkerType, state::{ destination_metadata::{AppliedDestinationTableMetadata, DestinationTableMetadata}, table::{TableReplicationPhase, TableReplicationPhaseType}, @@ -20,6 +21,7 @@ use crate::{ state::{DestinationTablesMetadata, StateStore, TableReplicationStates}, }, test_utils::notify::TimedNotify, + types::PgLsn, }; #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] @@ -42,6 +44,7 @@ struct Inner { table_state_history: HashMap>, table_schemas: Arc, destination_tables_metadata: DestinationTablesMetadata, + replication_progress: HashMap, table_state_type_conditions: Vec, table_state_conditions: Vec, table_schema_count_conditions: Vec, @@ -113,6 +116,7 @@ impl NotifyingStore { table_state_history: HashMap::new(), table_schemas: Arc::new(TableSchemaSnapshots::default()), destination_tables_metadata: Arc::new(BTreeMap::new()), + replication_progress: HashMap::new(), table_state_type_conditions: Vec::new(), table_state_conditions: Vec::new(), table_schema_count_conditions: Vec::new(), @@ -327,6 +331,38 @@ impl StateStore for NotifyingStore { Ok(previous_state) } + async fn get_replication_progress(&self, worker_type: WorkerType) -> EtlResult> { + let inner = self.inner.read().await; + + Ok(inner.replication_progress.get(&worker_type).copied()) + } + + async fn upsert_replication_progress( + &self, + worker_type: WorkerType, + flush_lsn: PgLsn, + ) -> EtlResult { + let mut inner = self.inner.write().await; + let stored_lsn = inner + .replication_progress + .entry(worker_type) + .and_modify(|stored_lsn| { + if flush_lsn > *stored_lsn { + *stored_lsn = flush_lsn; + } + }) + .or_insert(flush_lsn); + + Ok(*stored_lsn) + } + + async fn delete_replication_progress(&self, worker_type: WorkerType) -> EtlResult<()> { + let mut inner = self.inner.write().await; + inner.replication_progress.remove(&worker_type); + + Ok(()) + } + async fn get_destination_table_metadata( &self, table_id: TableId, @@ -423,6 +459,7 @@ impl CleanupStore for NotifyingStore { inner.table_state_history.remove(&table_id); Arc::make_mut(&mut inner.table_schemas).remove_table(table_id); Arc::make_mut(&mut inner.destination_tables_metadata).remove(&table_id); + inner.replication_progress.remove(&WorkerType::TableSync { table_id }); Ok(()) } diff --git a/crates/etl/src/workers/apply.rs b/crates/etl/src/workers/apply.rs index 3724516c2..0326630ba 100644 --- a/crates/etl/src/workers/apply.rs +++ b/crates/etl/src/workers/apply.rs @@ -18,6 +18,7 @@ use crate::{ }, replication::{ ApplyLoop, ApplyLoopResult, ApplyWorkerContext, SharedTableCache, WorkerContext, + WorkerType, client::{GetOrCreateSlotResult, PgReplicationClient, SlotState}, }, state::table::{TableReplicationPhase, TableReplicationPhaseType}, @@ -333,7 +334,7 @@ where } } -/// Determines the LSN position from which the apply worker should start reading +/// Determines the position from which the apply worker should start reading /// the replication stream. /// /// This function implements critical replication consistency logic by managing @@ -360,17 +361,18 @@ async fn get_start_lsn( invalidated_slot_behavior: &InvalidatedSlotBehavior, ) -> EtlResult { let slot_name: String = EtlReplicationSlot::for_apply_worker(pipeline_id).try_into()?; + let worker_type = WorkerType::Apply; // We try to get or create the slot. Both operations will return an LSN that we // can use to start streaming events. let slot = replication_client.get_or_create_slot(&slot_name).await?; - let start_lsn = slot.get_start_lsn(); + let slot_start_lsn = slot.get_start_lsn(); match &slot { GetOrCreateSlotResult::GetSlot(result) => { info!( slot_name, - %start_lsn, + %slot_start_lsn, confirmed_flush_lsn = %result.confirmed_flush_lsn, "apply worker resume position selected from existing replication slot" ); @@ -378,7 +380,7 @@ async fn get_start_lsn( GetOrCreateSlotResult::CreateSlot(result) => { info!( slot_name, - %start_lsn, + %slot_start_lsn, consistent_point = %result.consistent_point, "apply worker resume position selected from new replication slot" ); @@ -415,8 +417,54 @@ async fn get_start_lsn( return Err(err); } - // We return the LSN from which we will start streaming events. - Ok(start_lsn) + // A newly created slot is a new lineage. Any existing durable progress for + // the apply worker belongs to an older slot and must not influence startup. + if matches!(slot, GetOrCreateSlotResult::CreateSlot(_)) { + if let Err(err) = store.delete_replication_progress(worker_type).await { + warn!( + slot_name, + error = %err, + "failed to delete stale apply worker durable progress after creating a new slot" + ); + + replication_client.delete_slot_if_exists(&slot_name).await?; + + return Err(err); + } + + return Ok(slot_start_lsn); + } + + let durable_flush_lsn = store.get_replication_progress(worker_type).await?; + if let Some(durable_flush_lsn) = durable_flush_lsn { + // Durable progress and slot progress can legitimately differ. During idle + // periods we keep sending PostgreSQL feedback with the received LSN, but + // we do not persist those idle-only advances to the state database to + // avoid extra customer-database writes. Conversely, durable progress can + // be ahead if ETL flushed a batch but PostgreSQL did not confirm the + // feedback yet. Startup uses the latest boundary available from either + // source as a resume floor, which guarantees no event older than the + // chosen start LSN is emitted. + let start_lsn = durable_flush_lsn.max(slot_start_lsn); + + info!( + slot_name, + %durable_flush_lsn, + %slot_start_lsn, + %start_lsn, + "apply worker resume position selected from durable replication progress and replication slot" + ); + + Ok(start_lsn) + } else { + info!( + slot_name, + %slot_start_lsn, + "apply worker durable replication progress not found, using slot fallback" + ); + + Ok(slot_start_lsn) + } } /// Handles the case when the apply worker slot is found to be invalidated. @@ -470,6 +518,8 @@ async fn handle_invalidated_slot( info!(reset_count, "reset table replication states to init for resync"); + store.delete_replication_progress(WorkerType::Apply).await?; + // We delete and recreate the main apply worker slot. replication_client.delete_slot_if_exists(slot_name).await?; let create_result = replication_client.create_slot(slot_name).await?; diff --git a/crates/etl/src/workers/table_sync.rs b/crates/etl/src/workers/table_sync.rs index 911d173fb..a18c380d7 100644 --- a/crates/etl/src/workers/table_sync.rs +++ b/crates/etl/src/workers/table_sync.rs @@ -509,7 +509,7 @@ where // - Errored -> Init: okay since it will restart from scratch. // - Errored -> DataSync: okay since it will restart the copy from a new slot. // - Errored -> FinishedCopy: okay since the table was already copied, so it - // resumes streaming from the `confirmed_flush_lsn`. + // resumes streaming from durable table-sync progress or the slot fallback. // - Errored -> SyncDone: okay since the table sync will immediately stop. // - Errored -> Ready: same as SyncDone. // diff --git a/crates/etl/tests/pipeline_with_failpoints.rs b/crates/etl/tests/pipeline_with_failpoints.rs index 4a0f4d28a..92370625e 100644 --- a/crates/etl/tests/pipeline_with_failpoints.rs +++ b/crates/etl/tests/pipeline_with_failpoints.rs @@ -4,6 +4,7 @@ use etl::{ failpoints::{ FORCE_SCHEMA_CLEANUP_FP, SEND_STATUS_UPDATE_FP, START_TABLE_SYNC_BEFORE_DATA_SYNC_SLOT_CREATION_FP, START_TABLE_SYNC_DURING_DATA_SYNC_FP, + STORE_REPLICATION_PROGRESS_FP, }, state::table::{RetryPolicy, TableReplicationPhase, TableReplicationPhaseType}, store::state::StateStore, @@ -406,6 +407,7 @@ async fn table_schema_snapshots_are_consistent_after_missing_status_update_with_ { let _scenario = FailScenario::setup(); fail::cfg(SEND_STATUS_UPDATE_FP, "return").unwrap(); + fail::cfg(STORE_REPLICATION_PROGRESS_FP, "return").unwrap(); init_test_tracing(); @@ -560,6 +562,7 @@ async fn table_schema_snapshots_are_consistent_after_missing_status_update_with_ { let _scenario = FailScenario::setup(); fail::cfg(SEND_STATUS_UPDATE_FP, "return").unwrap(); + fail::cfg(STORE_REPLICATION_PROGRESS_FP, "return").unwrap(); init_test_tracing(); @@ -699,6 +702,7 @@ async fn table_schema_snapshots_are_consistent_after_missing_status_update_with_ { let _scenario = FailScenario::setup(); fail::cfg(SEND_STATUS_UPDATE_FP, "return").unwrap(); + fail::cfg(STORE_REPLICATION_PROGRESS_FP, "return").unwrap(); init_test_tracing(); @@ -840,6 +844,7 @@ async fn table_schema_snapshots_are_consistent_after_missing_status_update_with_ { let _scenario = FailScenario::setup(); fail::cfg(SEND_STATUS_UPDATE_FP, "return").unwrap(); + fail::cfg(STORE_REPLICATION_PROGRESS_FP, "return").unwrap(); init_test_tracing(); @@ -966,6 +971,7 @@ async fn table_schema_snapshots_are_consistent_after_missing_status_update_with_ async fn table_schema_snapshots_are_consistent_after_missing_status_update_with_initial_ddl() { let _scenario = FailScenario::setup(); fail::cfg(SEND_STATUS_UPDATE_FP, "return").unwrap(); + fail::cfg(STORE_REPLICATION_PROGRESS_FP, "return").unwrap(); init_test_tracing(); @@ -1171,6 +1177,7 @@ async fn schema_snapshots_are_pruned_after_confirmed_progress() { async fn table_schema_replication_masks_are_consistent_after_restart() { let _scenario = FailScenario::setup(); fail::cfg(SEND_STATUS_UPDATE_FP, "return").unwrap(); + fail::cfg(STORE_REPLICATION_PROGRESS_FP, "return").unwrap(); init_test_tracing(); let database = spawn_source_database().await; diff --git a/crates/etl/tests/postgres_store.rs b/crates/etl/tests/postgres_store.rs index f86369c99..48af29f4d 100644 --- a/crates/etl/tests/postgres_store.rs +++ b/crates/etl/tests/postgres_store.rs @@ -5,6 +5,7 @@ use std::collections::HashMap; use etl::{ error::ErrorKind, etl_error, + replication::WorkerType, state::{ destination_metadata::DestinationTableMetadata, table::{RetryPolicy, TableReplicationPhase}, @@ -219,6 +220,54 @@ async fn state_store_load_states() { assert_eq!(states.get(&table_id2), Some(&data_sync_phase)); } +#[tokio::test(flavor = "multi_thread")] +async fn state_store_replication_progress_is_monotonic() { + init_test_tracing(); + + let database = spawn_source_database().await; + let pipeline_id = 1; + let table_id = TableId::new(12345); + + let store = PostgresStore::new(pipeline_id, database.config.clone()).await.unwrap(); + let apply_worker = WorkerType::Apply; + let table_sync_worker = WorkerType::TableSync { table_id }; + + assert_eq!(store.get_replication_progress(apply_worker).await.unwrap(), None); + + let first_lsn = PgLsn::from(100u64); + let stale_lsn = PgLsn::from(90u64); + let later_lsn = PgLsn::from(120u64); + + assert_eq!( + store.upsert_replication_progress(apply_worker, first_lsn).await.unwrap(), + first_lsn + ); + assert_eq!( + store.upsert_replication_progress(apply_worker, stale_lsn).await.unwrap(), + first_lsn + ); + assert_eq!( + store.upsert_replication_progress(apply_worker, later_lsn).await.unwrap(), + later_lsn + ); + assert_eq!(store.get_replication_progress(apply_worker).await.unwrap(), Some(later_lsn)); + + let table_sync_lsn = PgLsn::from(75u64); + assert_eq!( + store.upsert_replication_progress(table_sync_worker, table_sync_lsn).await.unwrap(), + table_sync_lsn + ); + assert_eq!( + store.get_replication_progress(table_sync_worker).await.unwrap(), + Some(table_sync_lsn) + ); + assert_eq!(store.get_replication_progress(apply_worker).await.unwrap(), Some(later_lsn)); + + store.delete_replication_progress(table_sync_worker).await.unwrap(); + assert_eq!(store.get_replication_progress(table_sync_worker).await.unwrap(), None); + assert_eq!(store.get_replication_progress(apply_worker).await.unwrap(), Some(later_lsn)); +} + #[tokio::test(flavor = "multi_thread")] async fn schema_store_operations() { init_test_tracing(); From 610a3644fb72f783d694c09b7a7e790a7b4a236a Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Mon, 18 May 2026 11:53:32 +0200 Subject: [PATCH 10/29] feat(core): Clean up replication progress (#746) --- crates/etl/src/workers/apply.rs | 2 +- crates/etl/src/workers/table_sync.rs | 63 ++++++++++++++++++++-------- 2 files changed, 46 insertions(+), 19 deletions(-) diff --git a/crates/etl/src/workers/apply.rs b/crates/etl/src/workers/apply.rs index 0326630ba..f3eaa05ad 100644 --- a/crates/etl/src/workers/apply.rs +++ b/crates/etl/src/workers/apply.rs @@ -279,7 +279,6 @@ where async fn run_apply_worker(self) -> EtlResult<()> { let replication_client = PgReplicationClient::connect(self.config.pg_connection.clone()).await?; - let _apply_loop_stream_guard = self.batch_budget.register_stream_load(1); let start_lsn = get_start_lsn( self.pipeline_id, @@ -302,6 +301,7 @@ where batch_budget: self.batch_budget.clone(), }); + let _apply_loop_stream_guard = self.batch_budget.register_stream_load(1); let apply_loop_result = ApplyLoop::start( self.pipeline_id, start_lsn, diff --git a/crates/etl/src/workers/table_sync.rs b/crates/etl/src/workers/table_sync.rs index a18c380d7..421dad51e 100644 --- a/crates/etl/src/workers/table_sync.rs +++ b/crates/etl/src/workers/table_sync.rs @@ -20,7 +20,7 @@ use crate::{ metrics::{ERROR_TYPE_LABEL, ETL_WORKER_ERRORS_TOTAL, WORKER_TYPE_LABEL}, replication::{ ApplyLoop, ApplyLoopResult, SharedTableCache, TableSyncResult, TableSyncWorkerContext, - WorkerContext, client::PgReplicationClient, start_table_sync, + WorkerContext, WorkerType, client::PgReplicationClient, start_table_sync, }, state::table::{ RetryPolicy, TableReplicationError, TableReplicationPhase, TableReplicationPhaseType, @@ -767,22 +767,23 @@ where } }; - let _apply_loop_stream_guard = self.batch_budget.register_stream_load(1); let worker_context = WorkerContext::TableSync(TableSyncWorkerContext { table_id: self.table_id, table_sync_worker_state: state, state_store: self.store.clone(), }); + + let _apply_loop_stream_guard = self.batch_budget.register_stream_load(1); let apply_loop_result = ApplyLoop::start( self.pipeline_id, start_lsn, - self.config, + Arc::clone(&self.config), replication_client.clone(), - self.store, - self.destination, - self.shared_table_cache, + self.store.clone(), + self.destination.clone(), + self.shared_table_cache.clone(), worker_context, - self.shutdown_rx, + self.shutdown_rx.clone(), self.memory_monitor.clone(), self.batch_budget.clone(), ) @@ -795,17 +796,17 @@ where "table sync apply loop completed successfully, deleting slot" ); - // We delete the replication slot used by this table sync worker. - // - // Note that if the deletion fails, the slot will remain in the database and - // will not be removed later, so manual intervention will be - // required. The reason for not implementing an automatic - // cleanup mechanism is that it would introduce performance overhead, - // and we expect this call to fail only rarely. - let slot_name: String = - EtlReplicationSlot::for_table_sync_worker(self.pipeline_id, self.table_id) - .try_into()?; - replication_client.delete_slot_if_exists(&slot_name).await?; + // Catchup has completed, so cleanup failures should not turn a + // completed table sync into a replication failure. + if let Err(err) = + self.cleanup_resources(replication_client, self.store.clone()).await + { + warn!( + table_id = self.table_id.0, + error = %err, + "failed to clean up table sync resources after completion" + ); + } Ok(TableSyncWorkerResult::Completed) } @@ -816,4 +817,30 @@ where } } } + + /// Cleans up resources owned by this table sync worker. + /// + /// Once the table sync apply loop completes, its durable progress row is no + /// longer needed for resume and the replication slot should be removed so + /// it stops retaining WAL. + /// + /// Progress is deleted first so a slot deletion failure leaves only a stale + /// slot behind. That can retain WAL and may require manual cleanup, but it + /// does not leave stale progress for a completed table sync worker. + async fn cleanup_resources( + &self, + replication_client: PgReplicationClient, + store: S, + ) -> EtlResult<()> { + store + .delete_replication_progress(WorkerType::TableSync { table_id: self.table_id }) + .await?; + + let slot_name: String = + EtlReplicationSlot::for_table_sync_worker(self.pipeline_id, self.table_id) + .try_into()?; + replication_client.delete_slot_if_exists(&slot_name).await?; + + Ok(()) + } } From 94cfff6e465d5a9a676daa10b89f53d2279f904f Mon Sep 17 00:00:00 2001 From: Victor Farazdagi Date: Tue, 19 May 2026 19:26:27 +0300 Subject: [PATCH 11/29] ci: gate integration tests behind environment approval for fork PRs (#751) --- .github/workflows/ci.yml | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 2899c587c..c0f8ec647 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -132,8 +132,32 @@ jobs: if: matrix.check == 'sort' run: cargo sort --workspace --grouped --check + fork-pr-notice: + name: Fork PR Notice + if: github.event_name == 'pull_request' && github.event.pull_request.head.repo.fork == true + runs-on: blacksmith-4vcpu-ubuntu-2404 + permissions: + pull-requests: write + steps: + - name: Post security notice + uses: marocchino/sticky-pull-request-comment@0ea0beb66eb9baf113663a64ec522f60e49231c0 # v3.0.4 + with: + header: fork-pr-security-notice + message: | + ## Integration Tests -- Maintainer Approval Required + + This PR is from a fork and requires maintainer approval to run integration tests. + + **Before approving in the Actions tab**, verify that this PR does not: + - Modify CI workflow files, scripts, or build configuration + - Add suspicious dependencies in `Cargo.toml` + - Contain code that could exfiltrate environment variables or secrets + + Approving will expose BigQuery service account credentials to this workflow run. + test-full: name: Tests (Full, ${{ matrix.postgres_label }}) + environment: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.repo.fork == true && 'integration-tests-secrets' || null }} runs-on: blacksmith-4vcpu-ubuntu-2404 permissions: contents: read From e2cc9fb305db5cd234f561d8caed295dcc885562 Mon Sep 17 00:00:00 2001 From: Victor Farazdagi Date: Wed, 20 May 2026 06:13:14 +0300 Subject: [PATCH 12/29] ci: gracefully ignore BigQuery tests on fork PRs (#752) --- .github/workflows/ci.yml | 34 ++++--------------- .../src/bigquery/test_utils.rs | 17 ++++++++++ 2 files changed, 24 insertions(+), 27 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index c0f8ec647..b66e1beff 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -132,32 +132,8 @@ jobs: if: matrix.check == 'sort' run: cargo sort --workspace --grouped --check - fork-pr-notice: - name: Fork PR Notice - if: github.event_name == 'pull_request' && github.event.pull_request.head.repo.fork == true - runs-on: blacksmith-4vcpu-ubuntu-2404 - permissions: - pull-requests: write - steps: - - name: Post security notice - uses: marocchino/sticky-pull-request-comment@0ea0beb66eb9baf113663a64ec522f60e49231c0 # v3.0.4 - with: - header: fork-pr-security-notice - message: | - ## Integration Tests -- Maintainer Approval Required - - This PR is from a fork and requires maintainer approval to run integration tests. - - **Before approving in the Actions tab**, verify that this PR does not: - - Modify CI workflow files, scripts, or build configuration - - Add suspicious dependencies in `Cargo.toml` - - Contain code that could exfiltrate environment variables or secrets - - Approving will expose BigQuery service account credentials to this workflow run. - test-full: name: Tests (Full, ${{ matrix.postgres_label }}) - environment: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.repo.fork == true && 'integration-tests-secrets' || null }} runs-on: blacksmith-4vcpu-ubuntu-2404 permissions: contents: read @@ -246,10 +222,14 @@ jobs: ./scripts/run_migrations.sh etl-api - name: Set up BigQuery Credentials + # Conditional: fork PRs never receive secrets, BigQuery tests skipped gracefully. run: | - printf '%s' '${{ secrets.TESTS_BIGQUERY_SA_KEY_JSON }}' > /tmp/bigquery-sa-key.json - echo "TESTS_BIGQUERY_PROJECT_ID=${{ secrets.TESTS_BIGQUERY_PROJECT_ID }}" >> $GITHUB_ENV - echo "TESTS_BIGQUERY_SA_KEY_PATH=/tmp/bigquery-sa-key.json" >> $GITHUB_ENV + if [ -n '${{ secrets.TESTS_BIGQUERY_SA_KEY_JSON }}' ]; then + printf '%s' '${{ secrets.TESTS_BIGQUERY_SA_KEY_JSON }}' > /tmp/bigquery-sa-key.json + echo "TESTS_BIGQUERY_SA_KEY_PATH=/tmp/bigquery-sa-key.json" >> $GITHUB_ENV + echo "TESTS_BIGQUERY_PROJECT_ID=${{ secrets.TESTS_BIGQUERY_PROJECT_ID }}" >> $GITHUB_ENV + echo "REQUIRE_BIGQUERY_CREDENTIALS=true" >> $GITHUB_ENV + fi - name: Install cargo-nextest uses: taiki-e/install-action@328a871ad8f62ecac78390391f463ccabc974b72 # v2.69.9 diff --git a/crates/etl-destinations/src/bigquery/test_utils.rs b/crates/etl-destinations/src/bigquery/test_utils.rs index e6cf87d6d..5b93baf6f 100644 --- a/crates/etl-destinations/src/bigquery/test_utils.rs +++ b/crates/etl-destinations/src/bigquery/test_utils.rs @@ -53,9 +53,15 @@ pub const BIGQUERY_PROJECT_ID_ENV: &str = "TESTS_BIGQUERY_PROJECT_ID"; /// Environment variable name for the BigQuery service account key path. pub const BIGQUERY_SA_KEY_PATH_ENV: &str = "TESTS_BIGQUERY_SA_KEY_PATH"; +/// When set, tests panic instead of skipping when BigQuery credentials are +/// missing. +pub const REQUIRE_BIGQUERY_CREDENTIALS_ENV: &str = "REQUIRE_BIGQUERY_CREDENTIALS"; + /// Returns whether BigQuery integration tests should be skipped. /// /// Prints a warning and returns `true` when credentials are unavailable. +/// Panics if [`REQUIRE_BIGQUERY_CREDENTIALS_ENV`] is set, and credentials are +/// not provided. pub fn skip_if_missing_bigquery_env_vars() -> bool { let sa_key_path = std::env::var_os(BIGQUERY_SA_KEY_PATH_ENV); let has_project_id = std::env::var_os(BIGQUERY_PROJECT_ID_ENV).is_some(); @@ -65,10 +71,17 @@ pub fn skip_if_missing_bigquery_env_vars() -> bool { return false; } + let require = std::env::var_os(REQUIRE_BIGQUERY_CREDENTIALS_ENV).is_some_and(|v| !v.is_empty()); let mut missing_env_vars = Vec::new(); if !has_sa_key_path { missing_env_vars.push(BIGQUERY_SA_KEY_PATH_ENV); } else if !has_sa_key_file { + if require { + panic!( + "BigQuery credentials required but {BIGQUERY_SA_KEY_PATH_ENV} does not point to \ + an existing file" + ); + } eprintln!( "skipping bigquery integration test: {BIGQUERY_SA_KEY_PATH_ENV} does not point to an \ existing file" @@ -81,6 +94,10 @@ pub fn skip_if_missing_bigquery_env_vars() -> bool { missing_env_vars.push(BIGQUERY_PROJECT_ID_ENV); } + if require { + panic!("BigQuery credentials required but missing: {}", missing_env_vars.join(", ")); + } + eprintln!("skipping bigquery integration test: missing {}", missing_env_vars.join(", ")); true From baf8105b73b5c3e69dc1f0bbe8c945bf1d657641 Mon Sep 17 00:00:00 2001 From: Victor Farazdagi Date: Wed, 20 May 2026 06:58:22 +0300 Subject: [PATCH 13/29] fix(bigquery): add missing credential check (#753) --- crates/etl-destinations/tests/bigquery/pipeline.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/crates/etl-destinations/tests/bigquery/pipeline.rs b/crates/etl-destinations/tests/bigquery/pipeline.rs index 9fadf89ce..8528e8997 100644 --- a/crates/etl-destinations/tests/bigquery/pipeline.rs +++ b/crates/etl-destinations/tests/bigquery/pipeline.rs @@ -2000,6 +2000,9 @@ async fn table_validation_out_of_bounds_values() { #[tokio::test(flavor = "multi_thread")] async fn table_schema_change() { + if skip_if_missing_bigquery_env_vars() { + return; + } init_test_tracing(); install_crypto_provider(); From 1b3423638200d22e80f96a671fa8562a153bd929 Mon Sep 17 00:00:00 2001 From: ttatsato Date: Wed, 20 May 2026 13:11:06 +0900 Subject: [PATCH 14/29] ref(etl-telemetry): capitalize TracingError messages (#749) --- crates/etl-telemetry/src/tracing.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/crates/etl-telemetry/src/tracing.rs b/crates/etl-telemetry/src/tracing.rs index d940c21bc..807914318 100644 --- a/crates/etl-telemetry/src/tracing.rs +++ b/crates/etl-telemetry/src/tracing.rs @@ -43,16 +43,16 @@ static PIPELINE_ID: OnceLock = OnceLock::new(); /// Errors that can occur during tracing initialization. #[derive(Debug, Error)] pub enum TracingError { - #[error("failed to build rolling file appender: {0}")] + #[error("Failed to build rolling file appender: {0}")] InitAppender(#[from] InitError), - #[error("failed to init log tracer: {0}")] + #[error("Failed to init log tracer: {0}")] InitLogTracer(#[from] SetLoggerError), - #[error("failed to set global default subscriber: {0}")] + #[error("Failed to set global default subscriber: {0}")] SetGlobalDefault(#[from] SetGlobalDefaultError), - #[error("an io error occurred: {0}")] + #[error("An io error occurred: {0}")] Io(#[from] Error), } From bac3eca237b6a6930942e2bf76caf63120b068b0 Mon Sep 17 00:00:00 2001 From: Coenen Benjamin Date: Wed, 20 May 2026 08:16:48 +0200 Subject: [PATCH 15/29] refactor: clean maintenances and make maintenances more generic, (#745) --- AGENTS.md | 1 + Cargo.lock | 34 +- Cargo.toml | 2 + README.md | 5 + crates/etl-api/Cargo.toml | 1 + crates/etl-api/src/configs/destination.rs | 64 +- ...ination_config_serialization_ducklake.snap | 3 +- crates/etl-api/src/k8s/base.rs | 3 +- crates/etl-api/src/k8s/core.rs | 128 +- crates/etl-api/src/k8s/http.rs | 15 +- crates/etl-api/src/k8s/maintenance.rs | 96 + crates/etl-api/src/k8s/mod.rs | 2 + crates/etl-api/src/routes/pipelines.rs | 8 + crates/etl-api/src/validation/validators.rs | 1 + crates/etl-api/tests/support/mocks.rs | 3 +- crates/etl-config/src/shared/destination.rs | 21 + crates/etl-config/src/shared/mod.rs | 3 +- crates/etl-destinations/Cargo.toml | 8 +- crates/etl-destinations/README.md | 8 + .../etl-destinations/src/ducklake/batches.rs | 167 +- .../etl-destinations/src/ducklake/client.rs | 277 +- crates/etl-destinations/src/ducklake/core.rs | 267 +- .../src/ducklake/external_maintenance.rs | 743 ++--- .../src/ducklake/maintenance_runner.rs | 898 ------ .../etl-destinations/src/ducklake/metrics.rs | 16 +- crates/etl-destinations/src/ducklake/mod.rs | 14 +- crates/etl-maintenance/Cargo.toml | 55 + ...000000_external_maintenance_state.down.sql | 1 + ...15000000_external_maintenance_state.up.sql | 11 + crates/etl-maintenance/src/coordination.rs | 611 ++++ .../src/coordination/kubernetes.rs | 365 +++ .../src/coordination/postgres.rs | 312 ++ crates/etl-maintenance/src/ducklake/mod.rs | 10 + crates/etl-maintenance/src/ducklake/runner.rs | 2544 +++++++++++++++++ crates/etl-maintenance/src/lib.rs | 23 + crates/etl-maintenance/src/materialization.rs | 207 ++ crates/etl-replicator/Cargo.toml | 1 + .../src/bin/etl-ducklake-maintenance.rs | 4 +- crates/etl-replicator/src/core.rs | 22 +- 39 files changed, 5178 insertions(+), 1776 deletions(-) create mode 100644 crates/etl-api/src/k8s/maintenance.rs delete mode 100644 crates/etl-destinations/src/ducklake/maintenance_runner.rs create mode 100644 crates/etl-maintenance/Cargo.toml create mode 100644 crates/etl-maintenance/migrations/postgres/20260515000000_external_maintenance_state.down.sql create mode 100644 crates/etl-maintenance/migrations/postgres/20260515000000_external_maintenance_state.up.sql create mode 100644 crates/etl-maintenance/src/coordination.rs create mode 100644 crates/etl-maintenance/src/coordination/kubernetes.rs create mode 100644 crates/etl-maintenance/src/coordination/postgres.rs create mode 100644 crates/etl-maintenance/src/ducklake/mod.rs create mode 100644 crates/etl-maintenance/src/ducklake/runner.rs create mode 100644 crates/etl-maintenance/src/lib.rs create mode 100644 crates/etl-maintenance/src/materialization.rs diff --git a/AGENTS.md b/AGENTS.md index 5fab60b7e..7f57e97b9 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -60,6 +60,7 @@ ## Rust Style - This section is only for project-specific judgment that is not already covered by rustfmt, rustc, or Clippy. - Prefer absolute crate imports for shared module items, for example `use crate::metrics::{PIPELINE_ID_LABEL, APP_TYPE_LABEL};`, instead of `use super::{...};`. +- Write SQL queries with lowercase SQL keywords and identifiers, unless quoting or an external API requires specific casing. - When multiple files share constants, helpers, or a single entrypoint, prefer a module directory with `mod.rs`. - Keep top-level binaries focused on orchestration; move implementation detail into helpers or modules. - Prefer clear, boring code over clever abstractions. diff --git a/Cargo.lock b/Cargo.lock index dfb1d0893..680af8dbe 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1878,6 +1878,7 @@ dependencies = [ "etl", "etl-config", "etl-destinations", + "etl-maintenance", "etl-postgres", "etl-telemetry", "insta", @@ -1949,6 +1950,7 @@ dependencies = [ "clickhouse", "duckdb", "etl", + "etl-maintenance", "etl-postgres", "etl-telemetry", "futures", @@ -1957,7 +1959,6 @@ dependencies = [ "iceberg", "iceberg-catalog-rest", "k8s-openapi", - "kube", "metrics", "parking_lot", "parquet", @@ -1996,6 +1997,32 @@ dependencies = [ "url", ] +[[package]] +name = "etl-maintenance" +version = "0.1.0" +dependencies = [ + "async-trait", + "chrono", + "duckdb", + "etl", + "etl-telemetry", + "k8s-openapi", + "kube", + "metrics", + "pg_escape", + "r2d2", + "regex", + "serde", + "serde_json", + "sqlx", + "tempfile", + "thiserror 2.0.18", + "tokio", + "tokio-postgres", + "tracing", + "url", +] + [[package]] name = "etl-postgres" version = "0.1.0" @@ -2020,6 +2047,7 @@ dependencies = [ "etl", "etl-config", "etl-destinations", + "etl-maintenance", "etl-telemetry", "k8s-openapi", "metrics", @@ -5867,6 +5895,7 @@ dependencies = [ "base64", "bytes", "cfg-if", + "chrono", "crc", "crossbeam-queue", "either", @@ -5944,6 +5973,7 @@ dependencies = [ "bitflags", "byteorder", "bytes", + "chrono", "crc", "digest", "dotenvy", @@ -5984,6 +6014,7 @@ dependencies = [ "base64", "bitflags", "byteorder", + "chrono", "crc", "dotenvy", "etcetera", @@ -6017,6 +6048,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fe2cd6cee87120b1e1dd31356b5589911995c777707e49f2750eec7c7fe43eef" dependencies = [ "atoi", + "chrono", "flume", "futures-channel", "futures-core", diff --git a/Cargo.toml b/Cargo.toml index e252f8d07..1036fcec8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -7,6 +7,7 @@ members = [ "crates/etl-config", "crates/etl-destinations", "crates/etl-examples", + "crates/etl-maintenance", "crates/etl-postgres", "crates/etl-replicator", "crates/etl-telemetry", @@ -83,6 +84,7 @@ duckdb = { version = "1", default-features = false } etl = { path = "crates/etl", default-features = false } etl-config = { path = "crates/etl-config", default-features = false } etl-destinations = { path = "crates/etl-destinations", default-features = false } +etl-maintenance = { path = "crates/etl-maintenance", default-features = false } etl-postgres = { path = "crates/etl-postgres", default-features = false } etl-telemetry = { path = "crates/etl-telemetry", default-features = false } fail = { version = "0.5.1", default-features = false } diff --git a/README.md b/README.md index e311f2888..879410e83 100644 --- a/README.md +++ b/README.md @@ -97,6 +97,11 @@ For detailed configuration instructions, see the [Configure Postgres documentati ETL is currently installed from Git while we prepare for a crates.io release. Choose the destination features you need. +For DuckLake, external maintenance coordination is selected at runtime with +`maintenance_mode`: `disabled`, `kubernetes`, or `postgres`. The default is +`disabled`; `postgres` uses the same Postgres catalog connection as DuckLake +and stores coordination state in the `etl` schema. + For a first production deployment, start with the stable BigQuery module: ```toml diff --git a/crates/etl-api/Cargo.toml b/crates/etl-api/Cargo.toml index 54aa4e8ce..33aed6baf 100644 --- a/crates/etl-api/Cargo.toml +++ b/crates/etl-api/Cargo.toml @@ -32,6 +32,7 @@ constant_time_eq = { workspace = true } etl = { workspace = true } etl-config = { workspace = true, features = ["utoipa", "supabase"] } etl-destinations = { workspace = true, features = ["bigquery", "clickhouse", "ducklake", "iceberg"] } +etl-maintenance = { workspace = true } etl-postgres = { workspace = true, features = ["replication"] } etl-telemetry = { workspace = true } k8s-openapi = { workspace = true, features = ["latest"] } diff --git a/crates/etl-api/src/configs/destination.rs b/crates/etl-api/src/configs/destination.rs index f27d215b7..0a4ef7c1a 100644 --- a/crates/etl-api/src/configs/destination.rs +++ b/crates/etl-api/src/configs/destination.rs @@ -1,6 +1,6 @@ use etl_config::{ SerializableSecretString, - shared::{DestinationConfig, IcebergConfig}, + shared::{DestinationConfig, DuckLakeMaintenanceMode, IcebergConfig}, }; use secrecy::ExposeSecret; use serde::{Deserialize, Serialize}; @@ -141,6 +141,9 @@ pub enum FullApiDestinationConfig { deserialize_with = "crate::utils::trim_option_string" )] expire_snapshots_older_than: Option, + #[schema(example = "kubernetes")] + #[serde(default)] + maintenance_mode: DuckLakeMaintenanceMode, }, } @@ -215,6 +218,7 @@ impl From for FullApiDestinationConfig { duckdb_memory_cache_limit, maintenance_target_file_size, expire_snapshots_older_than, + maintenance_mode, } => Self::Ducklake { catalog_url, data_path, @@ -229,6 +233,7 @@ impl From for FullApiDestinationConfig { duckdb_memory_cache_limit, maintenance_target_file_size, expire_snapshots_older_than, + maintenance_mode, }, } } @@ -266,6 +271,7 @@ pub enum StoredDestinationConfig { duckdb_memory_cache_limit: Option, maintenance_target_file_size: Option, expire_snapshots_older_than: Option, + maintenance_mode: DuckLakeMaintenanceMode, }, } @@ -343,6 +349,7 @@ impl StoredDestinationConfig { duckdb_memory_cache_limit, maintenance_target_file_size, expire_snapshots_older_than, + maintenance_mode, } => DestinationConfig::Ducklake { catalog_url, data_path, @@ -357,6 +364,7 @@ impl StoredDestinationConfig { duckdb_memory_cache_limit, maintenance_target_file_size, expire_snapshots_older_than, + maintenance_mode, }, } } @@ -434,6 +442,7 @@ impl From for StoredDestinationConfig { duckdb_memory_cache_limit, maintenance_target_file_size, expire_snapshots_older_than, + maintenance_mode, } => Self::Ducklake { catalog_url, data_path, @@ -448,6 +457,7 @@ impl From for StoredDestinationConfig { duckdb_memory_cache_limit, maintenance_target_file_size, expire_snapshots_older_than, + maintenance_mode, }, } } @@ -559,6 +569,7 @@ impl Encrypt for StoredDestinationConfig { duckdb_memory_cache_limit, maintenance_target_file_size, expire_snapshots_older_than, + maintenance_mode, } => { let s3_access_key_id = s3_access_key_id .map(|value| encrypt_text(value.expose_secret().to_owned(), encryption_key)) @@ -581,6 +592,7 @@ impl Encrypt for StoredDestinationConfig { duckdb_memory_cache_limit, maintenance_target_file_size, expire_snapshots_older_than, + maintenance_mode, }) } } @@ -623,6 +635,8 @@ pub enum EncryptedStoredDestinationConfig { duckdb_memory_cache_limit: Option, maintenance_target_file_size: Option, expire_snapshots_older_than: Option, + #[serde(default)] + maintenance_mode: DuckLakeMaintenanceMode, }, } @@ -743,6 +757,7 @@ impl Decrypt for EncryptedStoredDestinationConfig { duckdb_memory_cache_limit, maintenance_target_file_size, expire_snapshots_older_than, + maintenance_mode, } => Ok(StoredDestinationConfig::Ducklake { catalog_url, data_path, @@ -765,6 +780,7 @@ impl Decrypt for EncryptedStoredDestinationConfig { duckdb_memory_cache_limit, maintenance_target_file_size, expire_snapshots_older_than, + maintenance_mode, }), } } @@ -1343,6 +1359,7 @@ mod tests { duckdb_memory_cache_limit: Some("50MB".to_owned()), maintenance_target_file_size: Some("10MB".to_owned()), expire_snapshots_older_than: Some("7 days".to_owned()), + maintenance_mode: DuckLakeMaintenanceMode::Kubernetes, }; let key = EncryptionKey { id: 1, key: generate_random_key::<32>().unwrap() }; @@ -1366,6 +1383,7 @@ mod tests { duckdb_memory_cache_limit: memory1, maintenance_target_file_size: target1, expire_snapshots_older_than: expire1, + maintenance_mode: mode1, }, StoredDestinationConfig::Ducklake { catalog_url: c2, @@ -1381,6 +1399,7 @@ mod tests { duckdb_memory_cache_limit: memory2, maintenance_target_file_size: target2, expire_snapshots_older_than: expire2, + maintenance_mode: mode2, }, ) => { assert_eq!(c1, c2); @@ -1402,11 +1421,49 @@ mod tests { assert_eq!(memory1, memory2); assert_eq!(target1, target2); assert_eq!(expire1, expire2); + assert_eq!(mode1, mode2); } _ => panic!("Config types don't match"), } } + #[test] + fn encrypted_stored_destination_config_ducklake_defaults_maintenance_mode() { + let config: EncryptedStoredDestinationConfig = serde_json::from_value(serde_json::json!({ + "ducklake": { + "catalog_url": "postgres://user:pass@localhost:5432/ducklake_catalog", + "data_path": "s3://bucket/path", + "pool_size": 8 + } + })) + .unwrap(); + + match config { + EncryptedStoredDestinationConfig::Ducklake { maintenance_mode, .. } => { + assert_eq!(maintenance_mode, DuckLakeMaintenanceMode::Disabled); + } + _ => panic!("Config type doesn't match"), + } + } + + #[test] + fn full_api_destination_config_ducklake_defaults_maintenance_mode() { + let config: FullApiDestinationConfig = serde_json::from_value(serde_json::json!({ + "ducklake": { + "catalog_url": "postgres://user:pass@localhost:5432/ducklake_catalog", + "data_path": "s3://bucket/path" + } + })) + .unwrap(); + + match config { + FullApiDestinationConfig::Ducklake { maintenance_mode, .. } => { + assert_eq!(maintenance_mode, DuckLakeMaintenanceMode::Disabled); + } + _ => panic!("Config type doesn't match"), + } + } + #[test] fn full_api_destination_config_conversion_ducklake() { let full_config = FullApiDestinationConfig::Ducklake { @@ -1423,6 +1480,7 @@ mod tests { duckdb_memory_cache_limit: None, maintenance_target_file_size: None, expire_snapshots_older_than: None, + maintenance_mode: DuckLakeMaintenanceMode::Kubernetes, }; let stored: StoredDestinationConfig = full_config.clone().into(); @@ -1480,6 +1538,7 @@ mod tests { duckdb_memory_cache_limit: Some("50MB".to_owned()), maintenance_target_file_size: Some("10MB".to_owned()), expire_snapshots_older_than: Some("7 days".to_owned()), + maintenance_mode: DuckLakeMaintenanceMode::Kubernetes, }; assert_json_snapshot!(full_config); @@ -1502,6 +1561,7 @@ mod tests { duckdb_memory_cache_limit: memory1, maintenance_target_file_size: target1, expire_snapshots_older_than: expire1, + maintenance_mode: mode1, }, FullApiDestinationConfig::Ducklake { catalog_url: c2, @@ -1517,6 +1577,7 @@ mod tests { duckdb_memory_cache_limit: memory2, maintenance_target_file_size: target2, expire_snapshots_older_than: expire2, + maintenance_mode: mode2, }, ) => { assert_eq!(c1, &c2); @@ -1538,6 +1599,7 @@ mod tests { assert_eq!(memory1, &memory2); assert_eq!(target1, &target2); assert_eq!(expire1, &expire2); + assert_eq!(mode1, &mode2); } _ => panic!("Deserialization failed or variant mismatch"), } diff --git a/crates/etl-api/src/configs/snapshots/etl_api__configs__destination__tests__full_api_destination_config_serialization_ducklake.snap b/crates/etl-api/src/configs/snapshots/etl_api__configs__destination__tests__full_api_destination_config_serialization_ducklake.snap index 9d854a5cc..91082acea 100644 --- a/crates/etl-api/src/configs/snapshots/etl_api__configs__destination__tests__full_api_destination_config_serialization_ducklake.snap +++ b/crates/etl-api/src/configs/snapshots/etl_api__configs__destination__tests__full_api_destination_config_serialization_ducklake.snap @@ -17,6 +17,7 @@ expression: full_config "metadata_schema": "ducklake", "duckdb_memory_cache_limit": "50MB", "maintenance_target_file_size": "10MB", - "expire_snapshots_older_than": "7 days" + "expire_snapshots_older_than": "7 days", + "maintenance_mode": "kubernetes" } } diff --git a/crates/etl-api/src/k8s/base.rs b/crates/etl-api/src/k8s/base.rs index 2b5318e9c..ad58892e0 100644 --- a/crates/etl-api/src/k8s/base.rs +++ b/crates/etl-api/src/k8s/base.rs @@ -1,5 +1,6 @@ use async_trait::async_trait; use etl_config::Environment; +use etl_maintenance::DuckLakeMaintenancePolicy; use k8s_openapi::api::core::v1::ConfigMap; use thiserror::Error; @@ -49,7 +50,7 @@ pub struct DuckLakeMaintenanceResourceConfig { /// Image containing the maintenance binary. pub image: String, /// User-authored maintenance policy. - pub policy: DuckLakeMaintenanceConfig, + pub policy: DuckLakeMaintenancePolicy, } /// Replicator StatefulSet materialization input. diff --git a/crates/etl-api/src/k8s/core.rs b/crates/etl-api/src/k8s/core.rs index 3ceeeb2d7..745d8a7b8 100644 --- a/crates/etl-api/src/k8s/core.rs +++ b/crates/etl-api/src/k8s/core.rs @@ -1,6 +1,13 @@ use etl_config::{ Environment, - shared::{ReplicatorConfigWithoutSecrets, SupabaseConfigWithoutSecrets, TlsConfig}, + shared::{ + DuckLakeMaintenanceMode, ReplicatorConfigWithoutSecrets, SupabaseConfigWithoutSecrets, + TlsConfig, + }, +}; +use etl_maintenance::{ + DuckLakeMaintenanceMaterialization, MaintenanceIdentity, MaintenanceMaterializationError, + MaintenanceMaterializer, MaintenanceRuntimeRefs, }; use secrecy::ExposeSecret; use thiserror::Error; @@ -19,8 +26,9 @@ use crate::{ sources::Source, }, k8s::{ - DestinationType, DuckLakeMaintenanceResourceConfig, K8sClient, K8sError, PodStatus, + DestinationType, K8sClient, K8sError, KubernetesMaintenanceMaterializer, PodStatus, ReplicatorConfigMapFile, ReplicatorStatefulSetConfig, + ducklake_maintenance_policy_from_config, }, }; @@ -35,6 +43,9 @@ pub enum K8sCoreError { #[error("Could not load app environment")] MissingEnvironment, + + #[error("Maintenance materialization failed: {0}")] + MaintenanceMaterialization(#[from] MaintenanceMaterializationError), } /// Secret types required by different destination configurations. @@ -106,6 +117,10 @@ pub async fn create_or_update_pipeline_resources_in_k8s( let environment = Environment::load().map_err(|_| K8sCoreError::MissingEnvironment)?; let destination_type = (&destination.config).into(); + let ducklake_maintenance_mode = match &destination.config { + StoredDestinationConfig::Ducklake { maintenance_mode, .. } => Some(*maintenance_mode), + _ => None, + }; let supabase_config = SupabaseConfigWithoutSecrets { project_ref: tenant_id.to_owned(), @@ -130,18 +145,23 @@ pub async fn create_or_update_pipeline_resources_in_k8s( create_or_update_dynamic_replicator_secrets(k8s_client, &prefix, secrets).await?; create_or_update_replicator_config(k8s_client, &prefix, replicator_config, environment).await?; let replicator_image = image.name; - let ducklake_maintenance_for_replicator = matches!(destination_type, DestinationType::Ducklake) - .then(|| ducklake_maintenance.clone().unwrap_or_default()); + let ducklake_maintenance_for_kubernetes = + matches!(ducklake_maintenance_mode, Some(DuckLakeMaintenanceMode::Kubernetes)) + .then(|| ducklake_maintenance.clone().unwrap_or_default()); + let maintenance_identity = MaintenanceIdentity { + tenant_id: tenant_id.to_owned(), + pipeline_id: pipeline.id, + replicator_id: replicator.id, + resource_prefix: prefix.clone(), + }; + let maintenance_materializer = KubernetesMaintenanceMaterializer::new(k8s_client); create_or_update_ducklake_maintenance( - k8s_client, - &prefix, - tenant_id, - pipeline.id, - replicator.id, + &maintenance_materializer, + maintenance_identity, &replicator_image, destination_type, - ducklake_maintenance, + ducklake_maintenance_for_kubernetes.clone(), ) .await?; create_or_update_replicator_stateful_set( @@ -152,7 +172,7 @@ pub async fn create_or_update_pipeline_resources_in_k8s( environment, replicator_resources, destination_type, - ducklake_maintenance: ducklake_maintenance_for_replicator, + ducklake_maintenance: ducklake_maintenance_for_kubernetes, log_level, }, ) @@ -420,34 +440,30 @@ async fn create_or_update_replicator_stateful_set( } /// Creates, updates, or deletes the DuckLake maintenance CR. -#[allow(clippy::too_many_arguments)] async fn create_or_update_ducklake_maintenance( - k8s_client: &dyn K8sClient, - prefix: &str, - tenant_id: &str, - pipeline_id: i64, - replicator_id: i64, + materializer: &dyn MaintenanceMaterializer, + identity: MaintenanceIdentity, replicator_image: &str, destination_type: DestinationType, ducklake_maintenance: Option, ) -> Result<(), K8sCoreError> { + let Some(ducklake_maintenance) = ducklake_maintenance else { + materializer.delete_ducklake_maintenance(identity).await?; + return Ok(()); + }; + if !matches!(destination_type, DestinationType::Ducklake) { - k8s_client.delete_ducklake_maintenance(prefix).await?; + materializer.delete_ducklake_maintenance(identity).await?; return Ok(()); } - let policy = ducklake_maintenance.unwrap_or_default(); - k8s_client - .create_or_update_ducklake_maintenance( - prefix, - DuckLakeMaintenanceResourceConfig { - tenant_id: tenant_id.to_owned(), - pipeline_id, - replicator_id, - image: replicator_image.to_owned(), - policy, - }, - ) + let policy = ducklake_maintenance_policy_from_config(ducklake_maintenance); + materializer + .reconcile_ducklake_maintenance(DuckLakeMaintenanceMaterialization { + identity, + policy, + runtime_refs: MaintenanceRuntimeRefs { replicator_image: replicator_image.to_owned() }, + }) .await?; Ok(()) @@ -509,7 +525,10 @@ mod tests { use super::*; use crate::{ configs::{destination::StoredDestinationConfig, source::StoredSourceConfig}, - k8s::{K8sClient, K8sError, PodStatus, ReplicatorConfigMapFile}, + k8s::{ + DuckLakeMaintenanceResourceConfig, K8sClient, K8sError, PodStatus, + ReplicatorConfigMapFile, + }, }; #[derive(Debug, Clone)] @@ -672,6 +691,51 @@ mod tests { } } + fn maintenance_identity() -> MaintenanceIdentity { + MaintenanceIdentity { + tenant_id: "tenant-42".to_owned(), + pipeline_id: 123, + replicator_id: 456, + resource_prefix: "tenant-42-456".to_owned(), + } + } + + #[tokio::test] + async fn ducklake_maintenance_is_deleted_when_config_is_absent() { + let client = RecordingK8sClient::default(); + let materializer = KubernetesMaintenanceMaterializer::new(&client); + + create_or_update_ducklake_maintenance( + &materializer, + maintenance_identity(), + "etl-replicator:test", + DestinationType::Ducklake, + None, + ) + .await + .unwrap(); + + assert_eq!(client.calls(), vec!["delete-ducklake-maintenance:tenant-42-456"]); + } + + #[tokio::test] + async fn ducklake_maintenance_is_created_for_ducklake_config() { + let client = RecordingK8sClient::default(); + let materializer = KubernetesMaintenanceMaterializer::new(&client); + + create_or_update_ducklake_maintenance( + &materializer, + maintenance_identity(), + "etl-replicator:test", + DestinationType::Ducklake, + Some(crate::configs::pipeline::DuckLakeMaintenanceConfig::default()), + ) + .await + .unwrap(); + + assert_eq!(client.calls(), vec!["ducklake-maintenance:tenant-42-456"]); + } + #[tokio::test] async fn failed_pod_is_considered_active_for_deletion_guards() { let client = RecordingK8sClient { pod_status: PodStatus::Failed, ..Default::default() }; @@ -747,6 +811,7 @@ mod tests { duckdb_memory_cache_limit: None, maintenance_target_file_size: None, expire_snapshots_older_than: None, + maintenance_mode: DuckLakeMaintenanceMode::Kubernetes, }; let secrets = build_secrets_from_configs(&source_config, &destination_config); @@ -787,6 +852,7 @@ mod tests { duckdb_memory_cache_limit: None, maintenance_target_file_size: None, expire_snapshots_older_than: None, + maintenance_mode: DuckLakeMaintenanceMode::Kubernetes, }; let secrets = build_secrets_from_configs(&source_config, &destination_config); diff --git a/crates/etl-api/src/k8s/http.rs b/crates/etl-api/src/k8s/http.rs index 69ce622c7..d6615cc07 100644 --- a/crates/etl-api/src/k8s/http.rs +++ b/crates/etl-api/src/k8s/http.rs @@ -4,6 +4,8 @@ use async_trait::async_trait; use base64::{Engine, prelude::BASE64_STANDARD}; use chrono::Utc; use etl_config::Environment; +#[cfg(test)] +use etl_maintenance::DuckLakeMaintenancePolicy; use k8s_openapi::{ api::{ apps::v1::StatefulSet, @@ -941,25 +943,25 @@ fn create_ducklake_maintenance_json( }, "operations": { "inlineFlush": { - "enabled": true, + "enabled": config.policy.operation_policy.inline_flush_enabled, "minInlinedBytes": config.policy.min_inlined_bytes, }, "mergeAdjacentFiles": { - "enabled": true, + "enabled": config.policy.operation_policy.merge_adjacent_files_enabled, "maxCompactedFiles": config.policy.max_compacted_files, "maxTablesPerRun": config.policy.max_tables_per_run, "targetFileSize": config.policy.target_file_size, }, "rewriteDataFiles": { - "enabled": true, + "enabled": config.policy.operation_policy.rewrite_data_files_enabled, "deleteThreshold": config.policy.delete_threshold, "maxTablesPerRun": config.policy.max_tables_per_run, }, "expireSnapshots": { - "enabled": false, + "enabled": config.policy.operation_policy.expire_snapshots_enabled, }, "cleanupOldFiles": { - "enabled": true, + "enabled": config.policy.operation_policy.cleanup_old_files_enabled, } }, "jobTemplate": { @@ -1698,7 +1700,7 @@ mod tests { pipeline_id: 24, replicator_id: 42, image: "supabase/replicator:1.2.3".to_owned(), - policy: DuckLakeMaintenanceConfig { + policy: DuckLakeMaintenancePolicy { min_interval_seconds: 3600, max_pause_seconds: 2700, min_inlined_bytes: 10_000_000, @@ -1710,6 +1712,7 @@ mod tests { cpu_request_millicores: 1000, memory_request_mib: 1024, active_deadline_seconds: 1800, + operation_policy: Default::default(), }, }, ); diff --git a/crates/etl-api/src/k8s/maintenance.rs b/crates/etl-api/src/k8s/maintenance.rs new file mode 100644 index 000000000..119d3882a --- /dev/null +++ b/crates/etl-api/src/k8s/maintenance.rs @@ -0,0 +1,96 @@ +use async_trait::async_trait; +use etl_maintenance::{ + DuckLakeMaintenanceMaterialization, DuckLakeMaintenancePolicy, + ExternalMaintenanceOperationPolicy, MaintenanceIdentity, MaintenanceMaterializationError, + MaintenanceMaterializer, +}; + +use crate::{ + configs::pipeline::DuckLakeMaintenanceConfig, + k8s::{DuckLakeMaintenanceResourceConfig, K8sClient}, +}; + +/// Converts API-authored DuckLake maintenance config to backend-neutral policy. +pub fn ducklake_maintenance_policy_from_config( + config: DuckLakeMaintenanceConfig, +) -> DuckLakeMaintenancePolicy { + DuckLakeMaintenancePolicy { + min_interval_seconds: config.min_interval_seconds, + max_pause_seconds: config.max_pause_seconds, + min_inlined_bytes: config.min_inlined_bytes, + max_compacted_files: config.max_compacted_files, + max_tables_per_run: config.max_tables_per_run, + target_file_size: config.target_file_size, + delete_threshold: config.delete_threshold, + min_active_data_files: config.min_active_data_files, + cpu_request_millicores: config.cpu_request_millicores, + memory_request_mib: config.memory_request_mib, + active_deadline_seconds: config.active_deadline_seconds, + operation_policy: ExternalMaintenanceOperationPolicy::default(), + } +} + +impl From for DuckLakeMaintenanceConfig { + fn from(policy: DuckLakeMaintenancePolicy) -> Self { + Self { + min_interval_seconds: policy.min_interval_seconds, + max_pause_seconds: policy.max_pause_seconds, + min_inlined_bytes: policy.min_inlined_bytes, + max_compacted_files: policy.max_compacted_files, + max_tables_per_run: policy.max_tables_per_run, + target_file_size: policy.target_file_size, + delete_threshold: policy.delete_threshold, + min_active_data_files: policy.min_active_data_files, + cpu_request_millicores: policy.cpu_request_millicores, + memory_request_mib: policy.memory_request_mib, + active_deadline_seconds: policy.active_deadline_seconds, + } + } +} + +/// Kubernetes materializer used by Supabase infrastructure. +pub struct KubernetesMaintenanceMaterializer<'a> { + k8s_client: &'a dyn K8sClient, +} + +impl<'a> KubernetesMaintenanceMaterializer<'a> { + pub fn new(k8s_client: &'a dyn K8sClient) -> Self { + Self { k8s_client } + } +} + +#[async_trait] +impl MaintenanceMaterializer for KubernetesMaintenanceMaterializer<'_> { + async fn reconcile_ducklake_maintenance( + &self, + input: DuckLakeMaintenanceMaterialization, + ) -> Result<(), MaintenanceMaterializationError> { + self.k8s_client + .create_or_update_ducklake_maintenance( + &input.identity.resource_prefix, + DuckLakeMaintenanceResourceConfig { + tenant_id: input.identity.tenant_id, + pipeline_id: input.identity.pipeline_id, + replicator_id: input.identity.replicator_id, + image: input.runtime_refs.replicator_image, + policy: input.policy, + }, + ) + .await + .map_err(MaintenanceMaterializationError::kubernetes)?; + + Ok(()) + } + + async fn delete_ducklake_maintenance( + &self, + identity: MaintenanceIdentity, + ) -> Result<(), MaintenanceMaterializationError> { + self.k8s_client + .delete_ducklake_maintenance(&identity.resource_prefix) + .await + .map_err(MaintenanceMaterializationError::kubernetes)?; + + Ok(()) + } +} diff --git a/crates/etl-api/src/k8s/mod.rs b/crates/etl-api/src/k8s/mod.rs index 1b8a6f6c8..82350e55d 100644 --- a/crates/etl-api/src/k8s/mod.rs +++ b/crates/etl-api/src/k8s/mod.rs @@ -16,6 +16,8 @@ mod base; pub mod cache; pub mod core; pub mod http; +mod maintenance; pub use base::*; pub use cache::{TrustedRootCertsCache, TrustedRootCertsError}; +pub use maintenance::{KubernetesMaintenanceMaterializer, ducklake_maintenance_policy_from_config}; diff --git a/crates/etl-api/src/routes/pipelines.rs b/crates/etl-api/src/routes/pipelines.rs index 3c824d7c9..75e0fbf9b 100644 --- a/crates/etl-api/src/routes/pipelines.rs +++ b/crates/etl-api/src/routes/pipelines.rs @@ -132,6 +132,9 @@ pub enum PipelineError { #[error("Invalid pipeline request: {0}")] InvalidPipelineRequest(String), + + #[error("Maintenance materialization failed: {0}")] + MaintenanceMaterialization(#[from] etl_maintenance::MaintenanceMaterializationError), } impl From for PipelineError { @@ -153,6 +156,9 @@ impl From for PipelineError { crate::k8s::core::K8sCoreError::K8s(error) => Self::K8s(error), crate::k8s::core::K8sCoreError::InvalidConfig(error) => Self::InvalidConfig(error), crate::k8s::core::K8sCoreError::MissingEnvironment => Self::MissingEnvironment, + crate::k8s::core::K8sCoreError::MaintenanceMaterialization(error) => { + Self::MaintenanceMaterialization(error) + } } } } @@ -174,6 +180,7 @@ impl PipelineError { | PipelineError::ImageNotFound(_) | PipelineError::NoDefaultImageFound | PipelineError::InvalidConfig(_) + | PipelineError::MaintenanceMaterialization(_) | PipelineError::K8s(_) | PipelineError::MissingEnvironment | PipelineError::MissingTableReplicationState @@ -202,6 +209,7 @@ impl ResponseError for PipelineError { | PipelineError::ReplicatorsDb(_) | PipelineError::ImagesDb(_) | PipelineError::K8s(_) + | PipelineError::MaintenanceMaterialization(_) | PipelineError::TrustedRootCerts(_) | PipelineError::Database(_) | PipelineError::TableLookup(_) diff --git a/crates/etl-api/src/validation/validators.rs b/crates/etl-api/src/validation/validators.rs index 4e99f8d71..498c73680 100644 --- a/crates/etl-api/src/validation/validators.rs +++ b/crates/etl-api/src/validation/validators.rs @@ -952,6 +952,7 @@ impl Validator for DestinationValidator { duckdb_memory_cache_limit, maintenance_target_file_size, expire_snapshots_older_than, + maintenance_mode: _, } => { let validator = DucklakeValidator::new( catalog_url.clone(), diff --git a/crates/etl-api/tests/support/mocks.rs b/crates/etl-api/tests/support/mocks.rs index ea8b35468..5ea30e3b7 100644 --- a/crates/etl-api/tests/support/mocks.rs +++ b/crates/etl-api/tests/support/mocks.rs @@ -12,7 +12,7 @@ use etl_api::{ sources::{CreateSourceRequest, CreateSourceResponse}, }, }; -use etl_config::SerializableSecretString; +use etl_config::{SerializableSecretString, shared::DuckLakeMaintenanceMode}; use crate::support::test_app::TestApp; @@ -87,6 +87,7 @@ pub(crate) mod destinations { duckdb_memory_cache_limit: None, maintenance_target_file_size: Some("10MB".to_owned()), expire_snapshots_older_than: Some("7 days".to_owned()), + maintenance_mode: DuckLakeMaintenanceMode::Kubernetes, } } diff --git a/crates/etl-config/src/shared/destination.rs b/crates/etl-config/src/shared/destination.rs index b3cd5d3c0..9d410b2b0 100644 --- a/crates/etl-config/src/shared/destination.rs +++ b/crates/etl-config/src/shared/destination.rs @@ -1,6 +1,8 @@ use secrecy::SecretString; use serde::{Deserialize, Serialize}; use url::Url; +#[cfg(feature = "utoipa")] +use utoipa::ToSchema; const fn default_connection_pool_size() -> usize { DestinationConfig::DEFAULT_CONNECTION_POOL_SIZE @@ -10,6 +12,17 @@ const fn default_ducklake_pool_size() -> u32 { DestinationConfig::DEFAULT_DUCKLAKE_POOL_SIZE } +/// Runtime backend used for DuckLake external maintenance coordination. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] +#[cfg_attr(feature = "utoipa", derive(ToSchema))] +#[serde(rename_all = "snake_case")] +pub enum DuckLakeMaintenanceMode { + #[default] + Disabled, + Kubernetes, + Postgres, +} + /// Configuration for supported ETL data destinations. /// /// Specifies the destination type and its associated configuration parameters. @@ -91,6 +104,9 @@ pub enum DestinationConfig { maintenance_target_file_size: Option, /// Optional DuckLake snapshot-retention interval. expire_snapshots_older_than: Option, + /// External maintenance coordination backend. + #[serde(default)] + maintenance_mode: DuckLakeMaintenanceMode, }, } @@ -281,6 +297,9 @@ pub enum DestinationConfigWithoutSecrets { maintenance_target_file_size: Option, /// Optional DuckLake snapshot-retention interval. expire_snapshots_older_than: Option, + /// External maintenance coordination backend. + #[serde(default)] + maintenance_mode: DuckLakeMaintenanceMode, }, } @@ -319,6 +338,7 @@ impl From for DestinationConfigWithoutSecrets { duckdb_memory_cache_limit, maintenance_target_file_size, expire_snapshots_older_than, + maintenance_mode, } => DestinationConfigWithoutSecrets::Ducklake { catalog_url, data_path, @@ -331,6 +351,7 @@ impl From for DestinationConfigWithoutSecrets { duckdb_memory_cache_limit, maintenance_target_file_size, expire_snapshots_older_than, + maintenance_mode, }, } } diff --git a/crates/etl-config/src/shared/mod.rs b/crates/etl-config/src/shared/mod.rs index d4cf38a98..2d2f74d38 100644 --- a/crates/etl-config/src/shared/mod.rs +++ b/crates/etl-config/src/shared/mod.rs @@ -13,7 +13,8 @@ pub use connection::{ TcpKeepaliveConfig, TlsConfig, }; pub use destination::{ - DestinationConfig, DestinationConfigWithoutSecrets, IcebergConfig, IcebergConfigWithoutSecrets, + DestinationConfig, DestinationConfigWithoutSecrets, DuckLakeMaintenanceMode, IcebergConfig, + IcebergConfigWithoutSecrets, }; pub use pipeline::{ BatchConfig, InvalidatedSlotBehavior, MemoryBackpressureConfig, PipelineConfig, diff --git a/crates/etl-destinations/Cargo.toml b/crates/etl-destinations/Cargo.toml index 526307752..fba42ce4f 100644 --- a/crates/etl-destinations/Cargo.toml +++ b/crates/etl-destinations/Cargo.toml @@ -15,14 +15,16 @@ doctest = false [features] ducklake = [ "dep:duckdb", + "dep:etl-maintenance", + "etl-maintenance/ducklake", "dep:humantime", "dep:metrics", - "dep:kube", "dep:parking_lot", "dep:pg_escape", "dep:r2d2", "dep:rand", "dep:regex", + "dep:serde", "dep:serde_json", "dep:sqlx", "dep:tokio-postgres", @@ -73,16 +75,16 @@ test-utils = ["dep:uuid"] arrow = { workspace = true, optional = true } async-trait = { workspace = true, optional = true } -chrono = { workspace = true } +chrono = { workspace = true, features = ["serde"] } clickhouse = { workspace = true, optional = true, features = ["inserter", "lz4", "rustls-tls"] } duckdb = { workspace = true, optional = true, features = ["bundled", "json", "parquet", "r2d2"] } etl = { workspace = true } +etl-maintenance = { workspace = true, optional = true } futures = { workspace = true, optional = true } gcp-bigquery-client = { workspace = true, optional = true, features = ["rust-tls", "aws-lc-rs"] } humantime = { workspace = true, optional = true } iceberg = { workspace = true, optional = true } iceberg-catalog-rest = { workspace = true, optional = true } -kube = { workspace = true, optional = true, features = ["client", "rustls-tls"] } metrics = { workspace = true, optional = true } parking_lot = { workspace = true, optional = true } parquet = { workspace = true, optional = true, features = ["async", "arrow"] } diff --git a/crates/etl-destinations/README.md b/crates/etl-destinations/README.md index a72fdfdaf..505546c0c 100644 --- a/crates/etl-destinations/README.md +++ b/crates/etl-destinations/README.md @@ -9,3 +9,11 @@ Enable the destination modules you need with crate features: | `bigquery` | Google BigQuery | Stable | | `ducklake` | DuckLake | In progress | | `iceberg` | Apache Iceberg | Deprecated for now | + +DuckLake external maintenance is configured at runtime with +`maintenance_mode`: `disabled`, `kubernetes`, or `postgres`. The default is +`disabled`. Kubernetes coordination expects +`ETL_DUCKLAKE_MAINTENANCE_CR_NAME` and +`ETL_DUCKLAKE_MAINTENANCE_CR_NAMESPACE`. Postgres coordination uses the same +Postgres catalog connection as DuckLake and stores coordination state in the +`etl` schema. diff --git a/crates/etl-destinations/src/ducklake/batches.rs b/crates/etl-destinations/src/ducklake/batches.rs index a108b9188..447042be8 100644 --- a/crates/etl-destinations/src/ducklake/batches.rs +++ b/crates/etl-destinations/src/ducklake/batches.rs @@ -39,10 +39,7 @@ use tracing::{debug, trace, warn}; use crate::{ ducklake::{ DuckLakeTableName, LAKE_CATALOG, - client::{ - DuckDbBlockingOperationKind, DuckLakeConnectionManager, format_query_error_detail, - run_duckdb_blocking, - }, + client::{DuckLakeConnectionManager, format_query_error_detail, run_duckdb_blocking}, core::is_create_table_conflict, encoding::{ PreparedRows, cell_to_sql_literal_ref, prepare_rows, table_row_to_sql_literal_ref, @@ -319,42 +316,37 @@ pub(super) async fn ensure_applied_batches_table_exists( let created = Arc::clone(&applied_batches_table_created); let table_name = APPLIED_BATCHES_TABLE.to_owned(); - run_duckdb_blocking( - pool, - blocking_slots, - DuckDbBlockingOperationKind::Foreground, - move |conn| -> EtlResult<()> { - match conn.execute_batch(&ddl) { - Ok(()) => {} - Err(error) if is_create_table_conflict(&error, &table_name) => {} - Err(error) => { - return Err(etl_error!( - ErrorKind::DestinationQueryFailed, - "DuckLake CREATE TABLE failed", - format_query_error_detail(&ddl), - source: error - )); - } + run_duckdb_blocking(pool, blocking_slots, move |conn| -> EtlResult<()> { + match conn.execute_batch(&ddl) { + Ok(()) => {} + Err(error) if is_create_table_conflict(&error, &table_name) => {} + Err(error) => { + return Err(etl_error!( + ErrorKind::DestinationQueryFailed, + "DuckLake CREATE TABLE failed", + format_query_error_detail(&ddl), + source: error + )); } + } - let set_option_sql = format!( - "CALL {LAKE_CATALOG}.set_option('data_inlining_row_limit', {}, table_name => {});", - APPLIED_BATCHES_TABLE_DATA_INLINING_ROW_LIMIT, - quote_literal(APPLIED_BATCHES_TABLE), - ); - conn.execute_batch(&set_option_sql).map_err(|err| { - etl_error!( - ErrorKind::DestinationQueryFailed, - "DuckLake set_option failed", - format_query_error_detail(&set_option_sql), - source: err - ) - })?; + let set_option_sql = format!( + "CALL {LAKE_CATALOG}.set_option('data_inlining_row_limit', {}, table_name => {});", + APPLIED_BATCHES_TABLE_DATA_INLINING_ROW_LIMIT, + quote_literal(APPLIED_BATCHES_TABLE), + ); + conn.execute_batch(&set_option_sql).map_err(|err| { + etl_error!( + ErrorKind::DestinationQueryFailed, + "DuckLake set_option failed", + format_query_error_detail(&set_option_sql), + source: err + ) + })?; - created.store(true, Ordering::Relaxed); - Ok(()) - }, - ) + created.store(true, Ordering::Relaxed); + Ok(()) + }) .await } @@ -388,42 +380,37 @@ pub(super) async fn ensure_streaming_progress_table_exists( let created = Arc::clone(&streaming_progress_table_created); let table_name = STREAMING_PROGRESS_TABLE.to_owned(); - run_duckdb_blocking( - pool, - blocking_slots, - DuckDbBlockingOperationKind::Foreground, - move |conn| -> EtlResult<()> { - match conn.execute_batch(&ddl) { - Ok(()) => {} - Err(err) if is_create_table_conflict(&err, &table_name) => {} - Err(err) => { - return Err(etl_error!( - ErrorKind::DestinationQueryFailed, - "DuckLake CREATE TABLE failed", - format_query_error_detail(&ddl), - source: err - )); - } + run_duckdb_blocking(pool, blocking_slots, move |conn| -> EtlResult<()> { + match conn.execute_batch(&ddl) { + Ok(()) => {} + Err(err) if is_create_table_conflict(&err, &table_name) => {} + Err(err) => { + return Err(etl_error!( + ErrorKind::DestinationQueryFailed, + "DuckLake CREATE TABLE failed", + format_query_error_detail(&ddl), + source: err + )); } + } - let set_option_sql = format!( - "CALL {LAKE_CATALOG}.set_option('data_inlining_row_limit', {}, table_name => {});", - STREAMING_PROGRESS_TABLE_DATA_INLINING_ROW_LIMIT, - quote_literal(STREAMING_PROGRESS_TABLE), - ); - conn.execute_batch(&set_option_sql).map_err(|error| { - etl_error!( - ErrorKind::DestinationQueryFailed, - "DuckLake set_option failed", - format_query_error_detail(&set_option_sql), - source: error - ) - })?; + let set_option_sql = format!( + "CALL {LAKE_CATALOG}.set_option('data_inlining_row_limit', {}, table_name => {});", + STREAMING_PROGRESS_TABLE_DATA_INLINING_ROW_LIMIT, + quote_literal(STREAMING_PROGRESS_TABLE), + ); + conn.execute_batch(&set_option_sql).map_err(|error| { + etl_error!( + ErrorKind::DestinationQueryFailed, + "DuckLake set_option failed", + format_query_error_detail(&set_option_sql), + source: error + ) + })?; - created.store(true, Ordering::Relaxed); - Ok(()) - }, - ) + created.store(true, Ordering::Relaxed); + Ok(()) + }) .await } @@ -471,15 +458,10 @@ pub(super) async fn apply_table_batches_with_retry( let pool = Arc::clone(&pool); let blocking_slots = Arc::clone(&blocking_slots); async move { - run_duckdb_blocking( - pool, - blocking_slots, - DuckDbBlockingOperationKind::Foreground, - move |conn| { - apply_table_batches(conn, attempt_batches.as_ref())?; - Ok(()) - }, - ) + run_duckdb_blocking(pool, blocking_slots, move |conn| { + apply_table_batches(conn, attempt_batches.as_ref())?; + Ok(()) + }) .await } }, @@ -543,25 +525,20 @@ pub(super) async fn apply_table_batch_with_retry( let pool = Arc::clone(&pool); let blocking_slots = Arc::clone(&blocking_slots); async move { - run_duckdb_blocking( - pool, - blocking_slots, - DuckDbBlockingOperationKind::Foreground, - move |conn| { - if batch_kind == DuckLakeTableBatchKind::Copy { - if applied_batch_marker_exists(conn, attempt_batch.as_ref())? { - record_replayed_batch_skip(attempt_batch.as_ref()); - return Ok(()); - } - - apply_table_batch(conn, attempt_batch.as_ref())?; + run_duckdb_blocking(pool, blocking_slots, move |conn| { + if batch_kind == DuckLakeTableBatchKind::Copy { + if applied_batch_marker_exists(conn, attempt_batch.as_ref())? { + record_replayed_batch_skip(attempt_batch.as_ref()); return Ok(()); } - apply_table_batches(conn, std::slice::from_ref(attempt_batch.as_ref()))?; - Ok(()) - }, - ) + apply_table_batch(conn, attempt_batch.as_ref())?; + return Ok(()); + } + + apply_table_batches(conn, std::slice::from_ref(attempt_batch.as_ref()))?; + Ok(()) + }) .await } }, diff --git a/crates/etl-destinations/src/ducklake/client.rs b/crates/etl-destinations/src/ducklake/client.rs index 7e52c5f5f..541e28418 100644 --- a/crates/etl-destinations/src/ducklake/client.rs +++ b/crates/etl-destinations/src/ducklake/client.rs @@ -44,8 +44,8 @@ static POSTGRES_PASSWORD_REGEX: LazyLock = LazyLock::new(|| { /// Timeout applied to each foreground DuckLake blocking operation. pub(super) const FOREGROUND_QUERY_TIMEOUT: Duration = Duration::from_secs(3 * 60); -/// Timeout applied to each maintenance DuckLake blocking operation. -pub(super) const MAINTENANCE_QUERY_TIMEOUT: Duration = Duration::from_secs(3 * 60); +/// Stable log and error label for DuckDB blocking operations. +const DUCKDB_BLOCKING_OPERATION_KIND: &str = "foreground"; /// Extra time allowed for a timed-out DuckDB operation to return after /// interrupt() has been called. If the operation is still stuck after this, /// the process is no longer safe to keep running. @@ -67,34 +67,10 @@ fn remaining_ms_until(deadline: Instant) -> u64 { deadline.checked_duration_since(Instant::now()).unwrap_or(Duration::ZERO).as_millis() as u64 } -/// Timeout class applied to one DuckDB blocking operation. -#[derive(Clone, Copy, Debug, Eq, PartialEq)] -pub(super) enum DuckDbBlockingOperationKind { - Foreground, - Maintenance, -} - -impl DuckDbBlockingOperationKind { - fn as_str(self) -> &'static str { - match self { - Self::Foreground => "foreground", - Self::Maintenance => "maintenance", - } - } - - pub(super) fn timeout(self) -> Duration { - match self { - Self::Foreground => FOREGROUND_QUERY_TIMEOUT, - Self::Maintenance => MAINTENANCE_QUERY_TIMEOUT, - } - } -} - /// Async watchdog that interrupts one timed DuckDB query when its deadline /// expires. pub(super) struct DuckDbQueryWatchdog { operation_id: u64, - operation_kind: DuckDbBlockingOperationKind, timeout: Duration, timed_out: Arc, interrupt_tx: Option>, @@ -105,20 +81,11 @@ pub(super) struct DuckDbQueryWatchdog { impl DuckDbQueryWatchdog { #[cfg(test)] fn spawn(deadline: Instant) -> Self { - Self::spawn_with_context( - deadline, - 0, - DuckDbBlockingOperationKind::Foreground, - Duration::ZERO, - ) + Self::spawn_with_context(deadline, 0, Duration::ZERO) } - fn spawn_with_context( - deadline: Instant, - operation_id: u64, - operation_kind: DuckDbBlockingOperationKind, - timeout: Duration, - ) -> Self { + fn spawn_with_context(deadline: Instant, operation_id: u64, timeout: Duration) -> Self { + let operation_kind = DUCKDB_BLOCKING_OPERATION_KIND; let timed_out = Arc::new(AtomicBool::new(false)); let timeout_flag = Arc::clone(&timed_out); let (interrupt_tx, interrupt_rx) = oneshot::channel::(); @@ -126,13 +93,13 @@ impl DuckDbQueryWatchdog { let task = tokio::spawn(async move { info!( operation_id, - operation_kind = operation_kind.as_str(), + operation_kind = operation_kind, timeout_ms = timeout.as_millis() as u64, deadline_remaining_ms = remaining_ms_until(deadline), "ducklake query watchdog task started: operation_id={}, operation_kind={}, \ timeout_ms={}, deadline_remaining_ms={}", operation_id, - operation_kind.as_str(), + operation_kind, timeout.as_millis(), remaining_ms_until(deadline) ); @@ -143,10 +110,10 @@ impl DuckDbQueryWatchdog { _ = &mut done_rx => { info!( operation_id, - operation_kind = operation_kind.as_str(), + operation_kind = operation_kind, "ducklake query watchdog finished before interrupt handle: operation_id={}, operation_kind={}", operation_id, - operation_kind.as_str() + operation_kind ); return; }, @@ -154,12 +121,12 @@ impl DuckDbQueryWatchdog { Ok(handle) => { info!( operation_id, - operation_kind = operation_kind.as_str(), + operation_kind = operation_kind, deadline_remaining_ms = remaining_ms_until(deadline), "ducklake query watchdog received interrupt handle before deadline: \ operation_id={}, operation_kind={}, deadline_remaining_ms={}", operation_id, - operation_kind.as_str(), + operation_kind, remaining_ms_until(deadline) ); handle @@ -167,11 +134,11 @@ impl DuckDbQueryWatchdog { Err(_) => { warn!( operation_id, - operation_kind = operation_kind.as_str(), + operation_kind = operation_kind, "ducklake query watchdog interrupt sender dropped before deadline: \ operation_id={}, operation_kind={}", operation_id, - operation_kind.as_str() + operation_kind ); return; }, @@ -180,12 +147,12 @@ impl DuckDbQueryWatchdog { timeout_flag.store(true, Ordering::Relaxed); warn!( operation_id, - operation_kind = operation_kind.as_str(), + operation_kind = operation_kind, timeout_ms = timeout.as_millis() as u64, "ducklake query watchdog deadline elapsed before interrupt handle: \ operation_id={}, operation_kind={}, timeout_ms={}", operation_id, - operation_kind.as_str(), + operation_kind, timeout.as_millis() ); // If we didn't receive the interrupt_rx yet, make sure to get it to call interrupt() later @@ -194,11 +161,11 @@ impl DuckDbQueryWatchdog { _ = &mut done_rx => { info!( operation_id, - operation_kind = operation_kind.as_str(), + operation_kind = operation_kind, "ducklake query watchdog received done after deadline before interrupt handle: \ operation_id={}, operation_kind={}", operation_id, - operation_kind.as_str() + operation_kind ); return; }, @@ -206,22 +173,22 @@ impl DuckDbQueryWatchdog { Ok(handle) => { warn!( operation_id, - operation_kind = operation_kind.as_str(), + operation_kind = operation_kind, "ducklake query watchdog received interrupt handle after deadline: \ operation_id={}, operation_kind={}", operation_id, - operation_kind.as_str() + operation_kind ); handle }, Err(_) => { warn!( operation_id, - operation_kind = operation_kind.as_str(), + operation_kind = operation_kind, "ducklake query watchdog interrupt sender dropped after deadline: \ operation_id={}, operation_kind={}", operation_id, - operation_kind.as_str() + operation_kind ); return; }, @@ -233,20 +200,20 @@ impl DuckDbQueryWatchdog { if timeout_flag.load(Ordering::Relaxed) { warn!( operation_id, - operation_kind = operation_kind.as_str(), + operation_kind = operation_kind, "ducklake query watchdog calling interrupt after timeout: operation_id={}, \ operation_kind={}", operation_id, - operation_kind.as_str() + operation_kind ); interrupt_handle.interrupt(); warn!( operation_id, - operation_kind = operation_kind.as_str(), + operation_kind = operation_kind, "ducklake query watchdog interrupt returned after timeout: operation_id={}, \ operation_kind={}", operation_id, - operation_kind.as_str() + operation_kind ); return; } @@ -256,12 +223,12 @@ impl DuckDbQueryWatchdog { _ = &mut done_rx => { info!( operation_id, - operation_kind = operation_kind.as_str(), + operation_kind = operation_kind, deadline_remaining_ms = remaining_ms_until(deadline), "ducklake query watchdog received done before deadline after interrupt handle: \ operation_id={}, operation_kind={}, deadline_remaining_ms={}", operation_id, - operation_kind.as_str(), + operation_kind, remaining_ms_until(deadline) ); } @@ -269,22 +236,22 @@ impl DuckDbQueryWatchdog { timeout_flag.store(true, Ordering::Relaxed); warn!( operation_id, - operation_kind = operation_kind.as_str(), + operation_kind = operation_kind, timeout_ms = timeout.as_millis() as u64, "ducklake query watchdog deadline elapsed after interrupt handle; calling interrupt: \ operation_id={}, operation_kind={}, timeout_ms={}", operation_id, - operation_kind.as_str(), + operation_kind, timeout.as_millis() ); interrupt_handle.interrupt(); warn!( operation_id, - operation_kind = operation_kind.as_str(), + operation_kind = operation_kind, "ducklake query watchdog interrupt returned after handle/deadline path: \ operation_id={}, operation_kind={}", operation_id, - operation_kind.as_str() + operation_kind ); } } @@ -292,7 +259,6 @@ impl DuckDbQueryWatchdog { Self { operation_id, - operation_kind, timeout, timed_out, interrupt_tx: Some(interrupt_tx), @@ -309,23 +275,23 @@ impl DuckDbQueryWatchdog { if let Some(interrupt_tx) = self.interrupt_tx.take() { info!( operation_id = self.operation_id, - operation_kind = self.operation_kind.as_str(), + operation_kind = DUCKDB_BLOCKING_OPERATION_KIND, timeout_ms = self.timeout.as_millis() as u64, "ducklake query watchdog publishing interrupt handle: operation_id={}, \ operation_kind={}, timeout_ms={}", self.operation_id, - self.operation_kind.as_str(), + DUCKDB_BLOCKING_OPERATION_KIND, self.timeout.as_millis() ); let _ = interrupt_tx.send(handle); } else { warn!( operation_id = self.operation_id, - operation_kind = self.operation_kind.as_str(), + operation_kind = DUCKDB_BLOCKING_OPERATION_KIND, "ducklake query watchdog interrupt handle publish skipped because sender is gone: \ operation_id={}, operation_kind={}", self.operation_id, - self.operation_kind.as_str() + DUCKDB_BLOCKING_OPERATION_KIND ); } } @@ -334,24 +300,24 @@ impl DuckDbQueryWatchdog { if let Some(done_tx) = self.done_tx.take() { info!( operation_id = self.operation_id, - operation_kind = self.operation_kind.as_str(), + operation_kind = DUCKDB_BLOCKING_OPERATION_KIND, timed_out = self.timed_out(), "ducklake query watchdog finish signal sent: operation_id={}, operation_kind={}, \ timed_out={}", self.operation_id, - self.operation_kind.as_str(), + DUCKDB_BLOCKING_OPERATION_KIND, self.timed_out() ); let _ = done_tx.send(()); } else { warn!( operation_id = self.operation_id, - operation_kind = self.operation_kind.as_str(), + operation_kind = DUCKDB_BLOCKING_OPERATION_KIND, timed_out = self.timed_out(), "ducklake query watchdog finish skipped because sender is gone: operation_id={}, \ operation_kind={}, timed_out={}", self.operation_id, - self.operation_kind.as_str(), + DUCKDB_BLOCKING_OPERATION_KIND, self.timed_out() ); } @@ -615,17 +581,13 @@ pub(super) async fn build_warm_ducklake_pool( /// Builds a consistent timeout error for one blocking DuckDB stage. #[inline] -pub(super) fn duckdb_blocking_timeout_error( - operation_kind: DuckDbBlockingOperationKind, - timeout: Duration, - stage: &'static str, -) -> EtlError { +pub(super) fn duckdb_blocking_timeout_error(timeout: Duration, stage: &'static str) -> EtlError { etl_error!( ErrorKind::DestinationQueryFailed, "DuckLake blocking operation timed out", format!( "Operation kind={}, stage={stage}, timeout_ms={}", - operation_kind.as_str(), + DUCKDB_BLOCKING_OPERATION_KIND, timeout.as_millis() ) ) @@ -633,19 +595,18 @@ pub(super) fn duckdb_blocking_timeout_error( fn abort_stuck_duckdb_blocking_operation( operation_id: u64, - operation_kind: DuckDbBlockingOperationKind, timeout: Duration, abort_grace: Duration, ) -> ! { tracing::error!( operation_id, - operation_kind = operation_kind.as_str(), + operation_kind = DUCKDB_BLOCKING_OPERATION_KIND, timeout_ms = timeout.as_millis() as u64, abort_grace_ms = abort_grace.as_millis() as u64, "ducklake blocking operation did not return after timeout interrupt; aborting process: \ operation_id={}, operation_kind={}, timeout_ms={}, abort_grace_ms={}", operation_id, - operation_kind.as_str(), + DUCKDB_BLOCKING_OPERATION_KIND, timeout.as_millis(), abort_grace.as_millis() ); @@ -658,28 +619,20 @@ fn abort_stuck_duckdb_blocking_operation( pub(super) async fn run_duckdb_blocking( pool: Arc>, blocking_slots: Arc, - operation_kind: DuckDbBlockingOperationKind, operation: F, ) -> EtlResult where R: Send + 'static, F: FnOnce(&duckdb::Connection) -> EtlResult + Send + 'static, { - run_duckdb_blocking_with_timeout( - pool, - blocking_slots, - operation_kind, - operation_kind.timeout(), - operation, - ) - .await + run_duckdb_blocking_with_timeout(pool, blocking_slots, FOREGROUND_QUERY_TIMEOUT, operation) + .await } /// Runs one DuckDB operation with an explicit timeout budget. pub(super) async fn run_duckdb_blocking_with_timeout( pool: Arc>, blocking_slots: Arc, - operation_kind: DuckDbBlockingOperationKind, timeout: Duration, operation: F, ) -> EtlResult @@ -687,18 +640,19 @@ where R: Send + 'static, F: FnOnce(&duckdb::Connection) -> EtlResult + Send + 'static, { + let operation_kind = DUCKDB_BLOCKING_OPERATION_KIND; let operation_id = NEXT_DUCKDB_BLOCKING_OPERATION_ID.fetch_add(1, Ordering::Relaxed); let deadline = Instant::now() + timeout; info!( operation_id, - operation_kind = operation_kind.as_str(), + operation_kind = operation_kind, timeout_ms = timeout.as_millis() as u64, abort_grace_ms = BLOCKING_ABORT_GRACE.as_millis() as u64, available_permits = blocking_slots.available_permits(), "ducklake blocking operation starting: operation_id={}, operation_kind={}, timeout_ms={}, \ abort_grace_ms={}, available_permits={}", operation_id, - operation_kind.as_str(), + operation_kind, timeout.as_millis(), BLOCKING_ABORT_GRACE.as_millis(), blocking_slots.available_permits() @@ -706,13 +660,13 @@ where let slot_wait_started = Instant::now(); info!( operation_id, - operation_kind = operation_kind.as_str(), + operation_kind = operation_kind, deadline_remaining_ms = remaining_ms_until(deadline), available_permits = blocking_slots.available_permits(), "ducklake blocking operation waiting for semaphore slot: operation_id={}, \ operation_kind={}, deadline_remaining_ms={}, available_permits={}", operation_id, - operation_kind.as_str(), + operation_kind, remaining_ms_until(deadline), blocking_slots.available_permits() ); @@ -726,10 +680,10 @@ where Ok(Err(_)) => { tracing::error!( operation_id, - operation_kind = operation_kind.as_str(), + operation_kind = operation_kind, "ducklake blocking operation semaphore closed: operation_id={}, operation_kind={}", operation_id, - operation_kind.as_str() + operation_kind ); return Err(etl_error!( ErrorKind::ApplyWorkerPanic, @@ -739,30 +693,30 @@ where Err(_) => { warn!( operation_id, - operation_kind = operation_kind.as_str(), + operation_kind = operation_kind, timeout_ms = timeout.as_millis() as u64, slot_wait_ms = slot_wait_started.elapsed().as_millis() as u64, "ducklake blocking operation timed out waiting for semaphore slot: \ operation_id={}, operation_kind={}, timeout_ms={}, slot_wait_ms={}", operation_id, - operation_kind.as_str(), + operation_kind, timeout.as_millis(), slot_wait_started.elapsed().as_millis() ); - return Err(duckdb_blocking_timeout_error(operation_kind, timeout, "slot_wait")); + return Err(duckdb_blocking_timeout_error(timeout, "slot_wait")); } }; histogram!(ETL_DUCKLAKE_BLOCKING_SLOT_WAIT_SECONDS) .record(slot_wait_started.elapsed().as_secs_f64()); info!( operation_id, - operation_kind = operation_kind.as_str(), + operation_kind = operation_kind, slot_wait_ms = slot_wait_started.elapsed().as_millis() as u64, deadline_remaining_ms = remaining_ms_until(deadline), "ducklake blocking operation acquired semaphore slot: operation_id={}, operation_kind={}, \ slot_wait_ms={}, deadline_remaining_ms={}", operation_id, - operation_kind.as_str(), + operation_kind, slot_wait_started.elapsed().as_millis(), remaining_ms_until(deadline) ); @@ -774,8 +728,7 @@ where // This is needed to make sure we properly interrupt the blocking operation if // it exceeds the timeout, we don't just cancel the task and leave the // connection active. - let mut watchdog = - DuckDbQueryWatchdog::spawn_with_context(deadline, operation_id, operation_kind, timeout); + let mut watchdog = DuckDbQueryWatchdog::spawn_with_context(deadline, operation_id, timeout); let watchdog_task = watchdog.async_task_handle()?; let watchdog_timed_out = Arc::clone(&watchdog.timed_out); let abort_deadline = deadline + BLOCKING_ABORT_GRACE; @@ -783,12 +736,12 @@ where let blocking_task = tokio::task::spawn_blocking(move || -> EtlResult { info!( operation_id, - operation_kind = operation_kind.as_str(), + operation_kind = operation_kind, deadline_remaining_ms = remaining_ms_until(deadline), "ducklake blocking operation entered spawn_blocking task: operation_id={}, \ operation_kind={}, deadline_remaining_ms={}", operation_id, - operation_kind.as_str(), + operation_kind, remaining_ms_until(deadline) ); // Please if you modify the code inside this blocking task do not add any @@ -799,23 +752,23 @@ where if checkout_timeout.is_zero() { warn!( operation_id, - operation_kind = operation_kind.as_str(), + operation_kind = operation_kind, "ducklake blocking operation deadline reached before pool checkout: \ operation_id={}, operation_kind={}", operation_id, - operation_kind.as_str() + operation_kind ); - return Err(duckdb_blocking_timeout_error(operation_kind, timeout, "pool_checkout")); + return Err(duckdb_blocking_timeout_error(timeout, "pool_checkout")); } let checkout_started = Instant::now(); info!( operation_id, - operation_kind = operation_kind.as_str(), + operation_kind = operation_kind, checkout_timeout_ms = checkout_timeout.as_millis() as u64, "ducklake blocking operation checking out pooled connection: operation_id={}, \ operation_kind={}, checkout_timeout_ms={}", operation_id, - operation_kind.as_str(), + operation_kind, checkout_timeout.as_millis() ); let mut pooled_conn = match pool.get_timeout(checkout_timeout) { @@ -823,34 +776,30 @@ where Err(e) if Instant::now() >= deadline => { warn!( operation_id, - operation_kind = operation_kind.as_str(), + operation_kind = operation_kind, checkout_wait_ms = checkout_started.elapsed().as_millis() as u64, timeout_ms = timeout.as_millis() as u64, error = %e, "ducklake blocking operation timed out checking out pooled connection: \ operation_id={}, operation_kind={}, checkout_wait_ms={}, timeout_ms={}, error={}", operation_id, - operation_kind.as_str(), + operation_kind, checkout_started.elapsed().as_millis(), timeout.as_millis(), e ); - return Err(duckdb_blocking_timeout_error( - operation_kind, - timeout, - "pool_checkout", - )); + return Err(duckdb_blocking_timeout_error(timeout, "pool_checkout")); } Err(e) => { warn!( operation_id, - operation_kind = operation_kind.as_str(), + operation_kind = operation_kind, checkout_wait_ms = checkout_started.elapsed().as_millis() as u64, error = %e, "ducklake blocking operation failed checking out pooled connection: operation_id={}, \ operation_kind={}, checkout_wait_ms={}, error={}", operation_id, - operation_kind.as_str(), + operation_kind, checkout_started.elapsed().as_millis(), e ); @@ -865,13 +814,13 @@ where .record(checkout_started.elapsed().as_secs_f64()); info!( operation_id, - operation_kind = operation_kind.as_str(), + operation_kind = operation_kind, checkout_wait_ms = checkout_started.elapsed().as_millis() as u64, deadline_remaining_ms = remaining_ms_until(deadline), "ducklake blocking operation checked out pooled connection: operation_id={}, \ operation_kind={}, checkout_wait_ms={}, deadline_remaining_ms={}", operation_id, - operation_kind.as_str(), + operation_kind, checkout_started.elapsed().as_millis(), remaining_ms_until(deadline) ); @@ -884,61 +833,61 @@ where if operation_timeout.is_zero() { warn!( operation_id, - operation_kind = operation_kind.as_str(), + operation_kind = operation_kind, "ducklake blocking operation deadline reached before query execution: \ operation_id={}, operation_kind={}", operation_id, - operation_kind.as_str() + operation_kind ); - return Err(duckdb_blocking_timeout_error(operation_kind, timeout, "query_execution")); + return Err(duckdb_blocking_timeout_error(timeout, "query_execution")); } let interrupt_handle = pooled_conn.conn.interrupt_handle(); info!( operation_id, - operation_kind = operation_kind.as_str(), + operation_kind = operation_kind, operation_timeout_ms = operation_timeout.as_millis() as u64, "ducklake blocking operation publishing interrupt handle before query execution: \ operation_id={}, operation_kind={}, operation_timeout_ms={}", operation_id, - operation_kind.as_str(), + operation_kind, operation_timeout.as_millis() ); watchdog.publish_interrupt_handle(interrupt_handle); if watchdog.timed_out() { warn!( operation_id, - operation_kind = operation_kind.as_str(), + operation_kind = operation_kind, "ducklake blocking operation timed out before query started; marking pooled \ connection broken: operation_id={}, operation_kind={}", operation_id, - operation_kind.as_str() + operation_kind ); pooled_conn.broken = true; - return Err(duckdb_blocking_timeout_error(operation_kind, timeout, "query_execution")); + return Err(duckdb_blocking_timeout_error(timeout, "query_execution")); } let operation_started = Instant::now(); info!( operation_id, - operation_kind = operation_kind.as_str(), + operation_kind = operation_kind, deadline_remaining_ms = remaining_ms_until(deadline), "ducklake blocking operation invoking DuckDB closure: operation_id={}, \ operation_kind={}, deadline_remaining_ms={}", operation_id, - operation_kind.as_str(), + operation_kind, remaining_ms_until(deadline) ); let res = operation(&pooled_conn.conn); let operation_duration_ms = operation_started.elapsed().as_millis() as u64; info!( operation_id, - operation_kind = operation_kind.as_str(), + operation_kind = operation_kind, duration_ms = operation_duration_ms, timed_out = watchdog.timed_out(), result_is_error = res.is_err(), "ducklake blocking operation DuckDB closure returned: operation_id={}, \ operation_kind={}, duration_ms={}, timed_out={}, result_is_error={}", operation_id, - operation_kind.as_str(), + operation_kind, operation_duration_ms, watchdog.timed_out(), res.is_err() @@ -953,38 +902,38 @@ where if watchdog.timed_out() { warn!( operation_id, - operation_kind = operation_kind.as_str(), + operation_kind = operation_kind, duration_ms = operation_duration_ms, "ducklake blocking operation returned after timeout; marking pooled connection \ broken: operation_id={}, operation_kind={}, duration_ms={}", operation_id, - operation_kind.as_str(), + operation_kind, operation_duration_ms ); pooled_conn.broken = true; - return Err(duckdb_blocking_timeout_error(operation_kind, timeout, "query_execution")); + return Err(duckdb_blocking_timeout_error(timeout, "query_execution")); } if res.is_err() { warn!( operation_id, - operation_kind = operation_kind.as_str(), + operation_kind = operation_kind, duration_ms = operation_duration_ms, "ducklake blocking operation returned error; marking pooled connection broken: \ operation_id={}, operation_kind={}, duration_ms={}", operation_id, - operation_kind.as_str(), + operation_kind, operation_duration_ms ); pooled_conn.broken = true; } else { info!( operation_id, - operation_kind = operation_kind.as_str(), + operation_kind = operation_kind, duration_ms = operation_duration_ms, "ducklake blocking operation returned success; pooled connection remains healthy: \ operation_id={}, operation_kind={}, duration_ms={}", operation_id, - operation_kind.as_str(), + operation_kind, operation_duration_ms ); } @@ -994,12 +943,12 @@ where info!( operation_id, - operation_kind = operation_kind.as_str(), + operation_kind = operation_kind, abort_deadline_remaining_ms = remaining_ms_until(abort_deadline), "ducklake blocking operation waiting for blocking task or abort deadline: \ operation_id={}, operation_kind={}, abort_deadline_remaining_ms={}", operation_id, - operation_kind.as_str(), + operation_kind, remaining_ms_until(abort_deadline) ); let blocking_result = tokio::select! { @@ -1012,47 +961,42 @@ where // the stuck native call also keeps holding its semaphore permit and // blocking thread, so restarting the process is the recoverable // boundary. - abort_stuck_duckdb_blocking_operation( - operation_id, - operation_kind, - timeout, - BLOCKING_ABORT_GRACE, - ); + abort_stuck_duckdb_blocking_operation(operation_id, timeout, BLOCKING_ABORT_GRACE); } }; match &blocking_result { Ok(Ok(_)) => info!( operation_id, - operation_kind = operation_kind.as_str(), + operation_kind = operation_kind, timed_out = watchdog_timed_out.load(Ordering::Relaxed), "ducklake blocking operation task joined with success: operation_id={}, \ operation_kind={}, timed_out={}", operation_id, - operation_kind.as_str(), + operation_kind, watchdog_timed_out.load(Ordering::Relaxed) ), Ok(Err(error)) => warn!( operation_id, - operation_kind = operation_kind.as_str(), + operation_kind = operation_kind, timed_out = watchdog_timed_out.load(Ordering::Relaxed), error = ?error, "ducklake blocking operation task joined with error: operation_id={}, operation_kind={}, \ timed_out={}, error={:?}", operation_id, - operation_kind.as_str(), + operation_kind, watchdog_timed_out.load(Ordering::Relaxed), error ), Err(error) => tracing::error!( operation_id, - operation_kind = operation_kind.as_str(), + operation_kind = operation_kind, timed_out = watchdog_timed_out.load(Ordering::Relaxed), error = %error, "ducklake blocking operation task join failed: operation_id={}, operation_kind={}, \ timed_out={}, error={}", operation_id, - operation_kind.as_str(), + operation_kind, watchdog_timed_out.load(Ordering::Relaxed), error ), @@ -1062,35 +1006,35 @@ where // accidentally interrupt a later operation that reuses the connection. info!( operation_id, - operation_kind = operation_kind.as_str(), + operation_kind = operation_kind, timed_out = watchdog_timed_out.load(Ordering::Relaxed), "ducklake blocking operation awaiting watchdog task: operation_id={}, operation_kind={}, \ timed_out={}", operation_id, - operation_kind.as_str(), + operation_kind, watchdog_timed_out.load(Ordering::Relaxed) ); match watchdog_task.await { Ok(()) => info!( operation_id, - operation_kind = operation_kind.as_str(), + operation_kind = operation_kind, timed_out = watchdog_timed_out.load(Ordering::Relaxed), "ducklake blocking operation watchdog task joined: operation_id={}, \ operation_kind={}, timed_out={}", operation_id, - operation_kind.as_str(), + operation_kind, watchdog_timed_out.load(Ordering::Relaxed) ), Err(error) => { tracing::error!( operation_id, - operation_kind = operation_kind.as_str(), + operation_kind = operation_kind, timed_out = watchdog_timed_out.load(Ordering::Relaxed), error = %error, "ducklake blocking operation watchdog task panicked: operation_id={}, operation_kind={}, \ timed_out={}, error={}", operation_id, - operation_kind.as_str(), + operation_kind, watchdog_timed_out.load(Ordering::Relaxed), error ); @@ -1259,12 +1203,6 @@ mod tests { ); } - #[test] - fn duckdb_blocking_operation_kind_timeouts() { - assert_eq!(DuckDbBlockingOperationKind::Foreground.timeout(), FOREGROUND_QUERY_TIMEOUT); - assert_eq!(DuckDbBlockingOperationKind::Maintenance.timeout(), MAINTENANCE_QUERY_TIMEOUT); - } - #[tokio::test] async fn run_duckdb_blocking_timeout_releases_resources_for_follow_up_queries() { let pool = Arc::new( @@ -1277,7 +1215,6 @@ mod tests { let error = run_duckdb_blocking_with_timeout( Arc::clone(&pool), Arc::clone(&blocking_slots), - DuckDbBlockingOperationKind::Foreground, Duration::from_millis(50), |_conn| -> EtlResult<()> { std::thread::sleep(Duration::from_millis(100)); @@ -1303,7 +1240,6 @@ mod tests { let value = run_duckdb_blocking_with_timeout( Arc::clone(&pool), Arc::clone(&blocking_slots), - DuckDbBlockingOperationKind::Foreground, Duration::from_secs(1), |conn| -> EtlResult { conn.query_row("SELECT 1;", [], |row| row.get::<_, i64>(0)).map_err(|source| { @@ -1679,7 +1615,6 @@ mod tests { run_duckdb_blocking_with_timeout( pool, blocking_slots, - DuckDbBlockingOperationKind::Foreground, Duration::from_secs(30), move |conn| -> EtlResult<()> { let _ = query_started_tx.send(()); diff --git a/crates/etl-destinations/src/ducklake/core.rs b/crates/etl-destinations/src/ducklake/core.rs index 6e0742463..a32de35b8 100644 --- a/crates/etl-destinations/src/ducklake/core.rs +++ b/crates/etl-destinations/src/ducklake/core.rs @@ -52,15 +52,15 @@ use crate::{ retain_truncates_after_sequence_key, }, client::{ - DuckDbBlockingOperationKind, DuckLakeConnectionManager, DuckLakeInterruptRegistry, - build_warm_ducklake_pool, format_query_error_detail, run_duckdb_blocking, + DuckLakeConnectionManager, DuckLakeInterruptRegistry, build_warm_ducklake_pool, + format_query_error_detail, run_duckdb_blocking, }, config::{ MAINTENANCE_TARGET_FILE_SIZE, build_setup_plan, current_duckdb_extension_strategy, maintenance_target_file_size_sql, resolve_expire_snapshots_older_than, validate_expire_snapshots_older_than_sql, }, - external_maintenance::{ExternalMaintenanceOperations, run_external_maintenance_watcher}, + external_maintenance::ExternalMaintenanceOperations, inline_size::DuckLakePendingInlineSizeSampler, metrics::{ DuckLakeMetricsSampler, ETL_DUCKLAKE_POOL_SIZE, query_catalog_maintenance_metrics, @@ -156,6 +156,42 @@ pub struct DuckLakeExternalMaintenancePause { _guard: OwnedRwLockWriteGuard<()>, } +/// Runtime backend used for DuckLake external maintenance coordination. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum DuckLakeMaintenanceMode { + #[default] + Disabled, + Kubernetes, + Postgres, +} + +/// Runtime configuration for DuckLake external maintenance coordination. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct DuckLakeExternalMaintenanceConfig { + pub mode: DuckLakeMaintenanceMode, + pub pipeline_id: u64, +} + +impl DuckLakeExternalMaintenanceConfig { + pub const fn disabled() -> Self { + Self { mode: DuckLakeMaintenanceMode::Disabled, pipeline_id: 0 } + } + + pub const fn kubernetes(pipeline_id: u64) -> Self { + Self { mode: DuckLakeMaintenanceMode::Kubernetes, pipeline_id } + } + + pub const fn postgres(pipeline_id: u64) -> Self { + Self { mode: DuckLakeMaintenanceMode::Postgres, pipeline_id } + } +} + +impl Default for DuckLakeExternalMaintenanceConfig { + fn default() -> Self { + Self::disabled() + } +} + /// Returns the table-local semaphore shared by concurrent foreground writes. fn table_write_slot( table_write_slots: &Arc>>>, @@ -433,6 +469,36 @@ where maintenance_target_file_size: Option, expire_snapshots_older_than: Option, store: S, + ) -> EtlResult { + Self::new_with_external_maintenance( + catalog_url, + data_path, + pool_size, + s3, + metadata_schema, + duckdb_memory_cache_limit, + maintenance_target_file_size, + expire_snapshots_older_than, + DuckLakeExternalMaintenanceConfig::default(), + store, + ) + .await + } + + /// Creates a new DuckLake destination with explicit external maintenance + /// runtime configuration. + #[allow(clippy::too_many_arguments)] + pub async fn new_with_external_maintenance( + catalog_url: Url, + data_path: Url, + pool_size: u32, + s3: Option, + metadata_schema: Option, + duckdb_memory_cache_limit: Option, + maintenance_target_file_size: Option, + expire_snapshots_older_than: Option, + external_maintenance: DuckLakeExternalMaintenanceConfig, + store: S, ) -> EtlResult { register_metrics(); @@ -493,7 +559,6 @@ where run_duckdb_blocking( Arc::clone(&pool), Arc::clone(&blocking_slots), - DuckDbBlockingOperationKind::Foreground, move |conn| -> EtlResult<()> { conn.execute_batch(&target_file_size_sql).map_err(|error| { etl_error!( @@ -512,7 +577,6 @@ where run_duckdb_blocking( Arc::clone(&pool), Arc::clone(&blocking_slots), - DuckDbBlockingOperationKind::Foreground, move |conn| -> EtlResult<()> { conn.query_row(&expire_snapshots_validation_sql, [], |_row| Ok(())).map_err( |source| { @@ -537,7 +601,6 @@ where run_duckdb_blocking( Arc::clone(&pool), Arc::clone(&blocking_slots), - DuckDbBlockingOperationKind::Foreground, resolve_ducklake_metadata_schema_blocking, ) .await? @@ -575,18 +638,53 @@ where )? .into(), ); - let watcher_destination = destination.clone(); - destination - .tasks - .spawn(async move { - if let Err(error) = run_external_maintenance_watcher(watcher_destination).await { - warn!( - error = %error, - "ducklake external maintenance watcher exited" - ); - } - }) - .await; + match external_maintenance.mode { + DuckLakeMaintenanceMode::Disabled => { + info!("ducklake external maintenance watcher disabled by configuration"); + } + DuckLakeMaintenanceMode::Kubernetes => { + use crate::ducklake::external_maintenance::run_kubernetes_external_maintenance_watcher; + + let watcher_destination = destination.clone(); + destination + .tasks + .spawn(async move { + if let Err(error) = + run_kubernetes_external_maintenance_watcher(watcher_destination).await + { + warn!( + error = %error, + "ducklake external maintenance watcher exited" + ); + } + }) + .await; + } + DuckLakeMaintenanceMode::Postgres => { + use crate::ducklake::external_maintenance::run_postgres_external_maintenance_watcher; + + let watcher_destination = destination.clone(); + let maintenance_pool = metadata_pg_pool.clone(); + let pipeline_id = external_maintenance.pipeline_id as i64; + destination + .tasks + .spawn(async move { + if let Err(error) = run_postgres_external_maintenance_watcher( + watcher_destination, + pipeline_id, + maintenance_pool, + ) + .await + { + warn!( + error = %error, + "ducklake external maintenance watcher exited" + ); + } + }) + .await; + } + } Ok(destination) } @@ -601,66 +699,63 @@ where self.ensure_applied_batches_table_exists().await?; self.ensure_streaming_progress_table_exists().await?; let _checkpoint_guard = self.acquire_mutation_guard().await; - self.run_duckdb_blocking( - DuckDbBlockingOperationKind::Foreground, - move |conn| -> EtlResult<()> { - conn.execute_batch("BEGIN TRANSACTION").map_err(|e| { + self.run_duckdb_blocking(move |conn| -> EtlResult<()> { + conn.execute_batch("BEGIN TRANSACTION").map_err(|e| { + etl_error!( + ErrorKind::DestinationQueryFailed, + "DuckLake BEGIN TRANSACTION failed", + source: e + ) + })?; + + let result = (|| -> EtlResult<()> { + let truncate_table_sql = + format!(r#"TRUNCATE TABLE {LAKE_CATALOG}."{table_name}";"#); + conn.execute_batch(&truncate_table_sql).map_err(|e| { etl_error!( ErrorKind::DestinationQueryFailed, - "DuckLake BEGIN TRANSACTION failed", + "DuckLake TRUNCATE TABLE failed", + format_query_error_detail(&truncate_table_sql), source: e ) })?; - let result = (|| -> EtlResult<()> { - let truncate_table_sql = - format!(r#"TRUNCATE TABLE {LAKE_CATALOG}."{table_name}";"#); - conn.execute_batch(&truncate_table_sql).map_err(|e| { - etl_error!( - ErrorKind::DestinationQueryFailed, - "DuckLake TRUNCATE TABLE failed", - format_query_error_detail(&truncate_table_sql), - source: e - ) - })?; - - clear_applied_batch_markers_for_kind( - conn, - &table_name, - DuckLakeTableBatchKind::Copy, - )?; - clear_applied_batch_markers_for_kind( - conn, - &table_name, - DuckLakeTableBatchKind::Mutation, - )?; - clear_applied_batch_markers_for_kind( - conn, - &table_name, - DuckLakeTableBatchKind::Truncate, - )?; - clear_table_streaming_progress(conn, &table_name)?; - Ok(()) - })(); - - match result { - Ok(()) => conn.execute_batch("COMMIT").map_err(|e| { - etl_error!( - ErrorKind::DestinationQueryFailed, - "DuckLake COMMIT failed", - source: e - ) - }), - Err(error) => { - let err = conn.execute_batch("ROLLBACK"); - if let Err(err) = err { - tracing::error!(error = %err, "error rollback"); - } - Err(error) + clear_applied_batch_markers_for_kind( + conn, + &table_name, + DuckLakeTableBatchKind::Copy, + )?; + clear_applied_batch_markers_for_kind( + conn, + &table_name, + DuckLakeTableBatchKind::Mutation, + )?; + clear_applied_batch_markers_for_kind( + conn, + &table_name, + DuckLakeTableBatchKind::Truncate, + )?; + clear_table_streaming_progress(conn, &table_name)?; + Ok(()) + })(); + + match result { + Ok(()) => conn.execute_batch("COMMIT").map_err(|e| { + etl_error!( + ErrorKind::DestinationQueryFailed, + "DuckLake COMMIT failed", + source: e + ) + }), + Err(error) => { + let err = conn.execute_batch("ROLLBACK"); + if let Err(err) = err { + tracing::error!(error = %err, "error rollback"); } + Err(error) } - }, - ) + } + }) .await } @@ -1072,7 +1167,6 @@ where run_duckdb_blocking( Arc::clone(&self.pool), Arc::clone(&self.blocking_slots), - DuckDbBlockingOperationKind::Foreground, move |conn| -> EtlResult<()> { debug!( table = %table_name_clone, @@ -1296,6 +1390,11 @@ where } } + if operations.rewrite_data_files { + operations.merge_adjacent_files = true; + operations.cleanup_old_files = true; + } + Ok(operations) } @@ -1322,22 +1421,13 @@ where /// Runs one DuckDB operation on Tokio's blocking pool after acquiring a /// permit that matches the configured DuckDB concurrency limit. - async fn run_duckdb_blocking( - &self, - operation_kind: DuckDbBlockingOperationKind, - operation: F, - ) -> EtlResult + async fn run_duckdb_blocking(&self, operation: F) -> EtlResult where R: Send + 'static, F: FnOnce(&duckdb::Connection) -> EtlResult + Send + 'static, { - run_duckdb_blocking( - Arc::clone(&self.pool), - Arc::clone(&self.blocking_slots), - operation_kind, - operation, - ) - .await + run_duckdb_blocking(Arc::clone(&self.pool), Arc::clone(&self.blocking_slots), operation) + .await } /// Stops the background DuckLake metrics sampler. @@ -1374,12 +1464,9 @@ async fn read_table_streaming_progress_sequence_key_blocking( blocking_slots: Arc, table_name: DuckLakeTableName, ) -> EtlResult> { - run_duckdb_blocking( - pool, - blocking_slots, - DuckDbBlockingOperationKind::Foreground, - move |conn| read_table_streaming_progress_sequence_key(conn, &table_name), - ) + run_duckdb_blocking(pool, blocking_slots, move |conn| { + read_table_streaming_progress_sequence_key(conn, &table_name) + }) .await } @@ -1467,6 +1554,7 @@ mod tests { TableSchema, Type as PgType, }, }; + use etl_maintenance::ducklake::flush_table_inlined_data; use etl_postgres::tokio::test_utils::PgDatabase; use pg_escape::{quote_identifier, quote_literal}; use tempfile::TempDir; @@ -1477,7 +1565,6 @@ mod tests { use super::*; use crate::ducklake::{ config::catalog_conninfo_from_url, - maintenance_runner::flush_table_inlined_data, metrics::{query_catalog_maintenance_metrics, query_table_storage_metrics}, }; diff --git a/crates/etl-destinations/src/ducklake/external_maintenance.rs b/crates/etl-destinations/src/ducklake/external_maintenance.rs index b035f64a0..8ccc4b0e3 100644 --- a/crates/etl-destinations/src/ducklake/external_maintenance.rs +++ b/crates/etl-destinations/src/ducklake/external_maintenance.rs @@ -1,14 +1,20 @@ -use std::{env, time::Duration}; +use std::time::Duration; use chrono::{DateTime, Utc}; -use etl::store::{schema::SchemaStore, state::StateStore}; -use kube::{ - Api, Client, - api::{Patch, PatchParams}, - core::{ApiResource, DynamicObject, GroupVersionKind}, +use etl::{ + error::EtlResult, + store::{schema::SchemaStore, state::StateStore}, +}; +pub use etl_maintenance::{ + ExternalMaintenanceOperationHistory, ExternalMaintenanceOperationPolicy, + ExternalMaintenanceOperationRequest, ExternalMaintenanceOperationRun, + ExternalMaintenanceOperations, ExternalMaintenancePause, ExternalMaintenanceReplicatorState, + ExternalMaintenanceReplicatorStatus, ExternalMaintenanceRequestOutcome, ExternalMaintenanceRun, + ExternalMaintenanceState, ExternalMaintenanceStore, ExternalMaintenanceWatcherConfig, + KubernetesExternalMaintenanceStore, PostgresExternalMaintenanceStore, }; use metrics::{counter, histogram}; -use serde_json::json; +use sqlx::PgPool; use tokio::time; use tracing::{debug, info, warn}; @@ -21,22 +27,6 @@ use crate::ducklake::{ }, }; -const CR_NAME_ENV: &str = "ETL_DUCKLAKE_MAINTENANCE_CR_NAME"; -const CR_NAMESPACE_ENV: &str = "ETL_DUCKLAKE_MAINTENANCE_CR_NAMESPACE"; -const POLL_SECONDS_ENV: &str = "ETL_DUCKLAKE_MAINTENANCE_POLL_SECONDS"; -const INLINE_FLUSH_MIN_INLINED_BYTES_ENV: &str = - "ETL_DUCKLAKE_EXTERNAL_MAINTENANCE_INLINE_FLUSH_MIN_INLINED_BYTES"; -const REWRITE_DATA_FILES_MIN_ACTIVE_DATA_FILES_ENV: &str = - "ETL_DUCKLAKE_EXTERNAL_MAINTENANCE_REWRITE_DATA_FILES_MIN_ACTIVE_DATA_FILES"; -const REQUEST_COOLDOWN_SECONDS_ENV: &str = - "ETL_DUCKLAKE_EXTERNAL_MAINTENANCE_REQUEST_COOLDOWN_SECONDS"; -const KUBERNETES_API_TIMEOUT_SECONDS_ENV: &str = - "ETL_DUCKLAKE_EXTERNAL_MAINTENANCE_KUBERNETES_API_TIMEOUT_SECONDS"; -const DEFAULT_POLL_SECONDS: u64 = 5; -const DEFAULT_INLINE_FLUSH_MIN_INLINED_BYTES: u64 = 10_000_000; -const DEFAULT_REWRITE_DATA_FILES_MIN_ACTIVE_DATA_FILES: i64 = 40; -const DEFAULT_REQUEST_COOLDOWN_SECONDS: u64 = 300; -const DEFAULT_KUBERNETES_API_TIMEOUT_SECONDS: u64 = 10; const OPERATION_INLINE_FLUSH: &str = "flush_inlined_data"; const OPERATION_REWRITE_DATA_FILES: &str = "rewrite_data_files"; const OPERATION_EXPIRE_SNAPSHOTS: &str = "expire_snapshots"; @@ -44,26 +34,6 @@ const REASON_PENDING_INLINED_DATA_BYTES_THRESHOLD: &str = "pending_inlined_data_ const REASON_ACTIVE_DATA_FILES_THRESHOLD: &str = "active_data_files_threshold"; const REASON_SNAPSHOT_RETENTION_THRESHOLD: &str = "snapshot_retention_threshold"; -#[derive(Clone, Copy, Debug, Default)] -pub(super) struct ExternalMaintenanceOperations { - pub(super) inline_flush: bool, - pub(super) rewrite_data_files: bool, - pub(super) expire_snapshots: bool, -} - -impl ExternalMaintenanceOperations { - fn covers(self, requested: Self) -> bool { - (!requested.inline_flush || self.inline_flush) - && (!requested.rewrite_data_files || self.rewrite_data_files) - && (!requested.expire_snapshots || self.expire_snapshots) - } -} - -struct PauseRequest { - run_id: String, - expires_at: DateTime, -} - struct HeldPause { run_id: String, expires_at: DateTime, @@ -72,96 +42,94 @@ struct HeldPause { _pause: DuckLakeExternalMaintenancePause, } -#[derive(Clone)] -struct WatcherConfig { - name: String, - namespace: String, - poll_interval: Duration, - request_cooldown: Duration, - kubernetes_api_timeout: Duration, - inline_flush_min_inlined_bytes: u64, - rewrite_data_files_min_active_data_files: i64, +pub(super) async fn run_kubernetes_external_maintenance_watcher( + destination: DuckLakeDestination, +) -> EtlResult<()> +where + S: StateStore + SchemaStore + Clone + Send + Sync + 'static, +{ + let config = ExternalMaintenanceWatcherConfig::from_env(); + let Some(store) = KubernetesExternalMaintenanceStore::from_env(config.store_timeout).await? + else { + info!("ducklake Kubernetes external maintenance watcher disabled because CR env is absent"); + return Ok(()); + }; + + run_external_maintenance_watcher(destination, store, config).await } -struct OperationPolicy { - inline_flush_enabled: bool, - rewrite_data_files_enabled: bool, - expire_snapshots_enabled: bool, +pub(super) async fn run_postgres_external_maintenance_watcher( + destination: DuckLakeDestination, + pipeline_id: i64, + pool: PgPool, +) -> EtlResult<()> +where + S: StateStore + SchemaStore + Clone + Send + Sync + 'static, +{ + let config = ExternalMaintenanceWatcherConfig::from_env(); + let store = PostgresExternalMaintenanceStore::new(pipeline_id, pool); + store.ensure_schema().await?; + store.ensure_pipeline_state_if_missing(ExternalMaintenanceOperationPolicy::default()).await?; + + info!( + pipeline_id, + "ducklake Postgres external maintenance watcher configured: pipeline_id={}", pipeline_id + ); + + run_external_maintenance_watcher(destination, store, config).await } -pub(super) async fn run_external_maintenance_watcher( +pub async fn run_external_maintenance_watcher( destination: DuckLakeDestination, -) -> Result<(), kube::Error> + store: M, + config: ExternalMaintenanceWatcherConfig, +) -> EtlResult<()> where S: StateStore + SchemaStore + Clone + Send + Sync + 'static, + M: ExternalMaintenanceStore, { - let Some(config) = WatcherConfig::from_env() else { - return Ok(()); - }; - let client = Client::try_default().await?; - let api: Api = - Api::namespaced_with(client, &config.namespace, &ducklake_maintenance_api_resource()); let mut held_pause: Option = None; info!( - ducklake_maintenance = %config.name, - namespace = %config.namespace, - "ducklake external maintenance watcher started: ducklake_maintenance={}, namespace={}", - config.name, - config.namespace + poll_interval_ms = config.poll_interval.as_millis() as u64, + request_cooldown_ms = config.request_cooldown.as_millis() as u64, + store_timeout_ms = config.store_timeout.as_millis() as u64, + "ducklake external maintenance watcher started: poll_interval_ms={}, \ + request_cooldown_ms={}, store_timeout_ms={}", + config.poll_interval.as_millis(), + config.request_cooldown.as_millis(), + config.store_timeout.as_millis() ); loop { - match time::timeout(config.kubernetes_api_timeout, api.get(&config.name)).await { - Ok(resource) => { - let resource = match resource { - Ok(resource) => resource, - Err(kube::Error::Api(error)) if error.code == 404 => { - if let Some(held) = held_pause.take() { - info!( - ducklake_maintenance = %config.name, - run_id = %held.run_id, - "ducklake maintenance resource disappeared, resuming foreground mutations: \ - ducklake_maintenance={}, run_id={}", - config.name, - held.run_id - ); - release_held_pause(&config.name, held, "resource_deleted"); - } - time::sleep(config.poll_interval).await; - continue; - } - Err(error) => { - warn!( - error = %error, - ducklake_maintenance = %config.name, - timeout_ms = config.kubernetes_api_timeout.as_millis() as u64, - "failed to read ducklake maintenance resource: ducklake_maintenance={}, \ - timeout_ms={}, error={}", - config.name, - config.kubernetes_api_timeout.as_millis(), - error - ); - release_expired_pause_if_needed(&api, &config, &mut held_pause).await; - time::sleep(config.poll_interval).await; - continue; - } - }; - let active_pause = active_pause_request(&resource); - reconcile_pause(&api, &config, &destination, &mut held_pause, active_pause).await; - maybe_request_operations(&api, &config.name, &destination, &resource, &config) - .await; + match time::timeout(config.store_timeout, store.load_state()).await { + Ok(Ok(state)) => { + if !state.exists { + release_missing_state_pause(&store, &mut held_pause).await; + time::sleep(config.poll_interval).await; + continue; + } + + reconcile_pause(&store, &config, &destination, &mut held_pause, &state).await; + maybe_request_operations(&store, &destination, &state, &config).await; + } + Ok(Err(error)) => { + warn!( + error = %error, + timeout_ms = config.store_timeout.as_millis() as u64, + "failed to read ducklake external maintenance state: timeout_ms={}, error={}", + config.store_timeout.as_millis(), + error + ); + release_expired_pause_if_needed(&store, &config, &mut held_pause).await; } Err(_) => { warn!( - ducklake_maintenance = %config.name, - timeout_ms = config.kubernetes_api_timeout.as_millis() as u64, - "timed out reading ducklake maintenance resource: ducklake_maintenance={}, \ - timeout_ms={}", - config.name, - config.kubernetes_api_timeout.as_millis() + timeout_ms = config.store_timeout.as_millis() as u64, + "timed out reading ducklake external maintenance state: timeout_ms={}", + config.store_timeout.as_millis() ); - release_expired_pause_if_needed(&api, &config, &mut held_pause).await; + release_expired_pause_if_needed(&store, &config, &mut held_pause).await; } } @@ -169,36 +137,49 @@ where } } -async fn reconcile_pause( - api: &Api, - config: &WatcherConfig, +async fn release_missing_state_pause(store: &M, held_pause: &mut Option) +where + M: ExternalMaintenanceStore, +{ + if let Some(held) = held_pause.take() { + info!( + run_id = %held.run_id, + "ducklake maintenance state disappeared, resuming foreground mutations: run_id={}", + held.run_id + ); + release_held_pause(held, "state_missing"); + if let Err(error) = store.clear_replicator_status().await { + warn!( + error = %error, + "failed to clear ducklake maintenance replicator status after state disappeared: \ + error={}", + error + ); + } + } +} + +async fn reconcile_pause( + store: &M, + config: &ExternalMaintenanceWatcherConfig, destination: &DuckLakeDestination, held_pause: &mut Option, - active_pause: Option, + state: &ExternalMaintenanceState, ) where S: StateStore + SchemaStore + Clone + Send + Sync + 'static, + M: ExternalMaintenanceStore, { - let name = config.name.as_str(); - let Some(pause) = active_pause.filter(|pause| pause.expires_at > Utc::now()) else { + let active_pause = state.pause_request.clone().filter(|pause| pause.expires_at > Utc::now()); + let Some(pause) = active_pause else { if let Some(held) = held_pause.take() { info!( - ducklake_maintenance = %name, run_id = %held.run_id, "ducklake external maintenance pause cleared, resuming foreground mutations: \ - ducklake_maintenance={}, run_id={}", - name, + run_id={}", held.run_id ); - release_held_pause(name, held, "cleared"); - patch_replicator_status( - api, - name, - "Running", - None, - None, - config.kubernetes_api_timeout, - ) - .await; + release_held_pause(held, "cleared"); + report_running(store, config.store_timeout).await; } return; }; @@ -207,15 +188,7 @@ async fn reconcile_pause( && held.run_id == pause.run_id { if !held.quiesced_reported - && patch_replicator_status( - api, - name, - "Quiesced", - Some(&held.run_id), - Some(held.quiesced_at), - config.kubernetes_api_timeout, - ) - .await + && report_quiesced(store, &held.run_id, held.quiesced_at, config.store_timeout).await { held.quiesced_reported = true; } @@ -224,54 +197,38 @@ async fn reconcile_pause( if let Some(held) = held_pause.take() { info!( - ducklake_maintenance = %name, previous_run_id = %held.run_id, next_run_id = %pause.run_id, "ducklake external maintenance pause replaced, resuming previous run before pausing again: \ - ducklake_maintenance={}, previous_run_id={}, next_run_id={}", - name, + previous_run_id={}, next_run_id={}", held.run_id, pause.run_id ); - release_held_pause(name, held, "replaced"); - patch_replicator_status(api, name, "Running", None, None, config.kubernetes_api_timeout) - .await; + release_held_pause(held, "replaced"); + report_running(store, config.store_timeout).await; } info!( - ducklake_maintenance = %name, run_id = %pause.run_id, expires_at = %pause.expires_at.to_rfc3339(), "ducklake external maintenance pause requested, waiting for foreground mutations to drain: \ - ducklake_maintenance={}, run_id={}, expires_at={}", - name, + run_id={}, expires_at={}", pause.run_id, pause.expires_at.to_rfc3339() ); - patch_replicator_status( - api, - name, - "Pausing", - Some(&pause.run_id), - None, - config.kubernetes_api_timeout, - ) - .await; + report_pausing(store, &pause.run_id, config.store_timeout).await; let external_pause = destination.acquire_external_maintenance_pause().await; if pause.expires_at <= Utc::now() { info!( - ducklake_maintenance = %name, run_id = %pause.run_id, expires_at = %pause.expires_at.to_rfc3339(), - "ducklake external maintenance pause expired before quiescence, resuming foreground mutations: \ - ducklake_maintenance={}, run_id={}, expires_at={}", - name, + "ducklake external maintenance pause expired before quiescence, resuming foreground \ + mutations: run_id={}, expires_at={}", pause.run_id, pause.expires_at.to_rfc3339() ); release_held_pause( - name, HeldPause { run_id: pause.run_id, expires_at: pause.expires_at, @@ -281,45 +238,33 @@ async fn reconcile_pause( }, "expired", ); - patch_replicator_status(api, name, "Running", None, None, config.kubernetes_api_timeout) - .await; + report_running(store, config.store_timeout).await; return; } let quiesced_at = Utc::now(); info!( - ducklake_maintenance = %name, run_id = %pause.run_id, quiesced_at = %quiesced_at.to_rfc3339(), expires_at = %pause.expires_at.to_rfc3339(), - "ducklake external maintenance quiesced, foreground mutations are paused: \ - ducklake_maintenance={}, run_id={}, quiesced_at={}, expires_at={}", - name, + "ducklake external maintenance quiesced, foreground mutations are paused: run_id={}, \ + quiesced_at={}, expires_at={}", pause.run_id, quiesced_at.to_rfc3339(), pause.expires_at.to_rfc3339() ); - let quiesced_reported = patch_replicator_status( - api, - name, - "Quiesced", - Some(&pause.run_id), - Some(quiesced_at), - config.kubernetes_api_timeout, - ) - .await; + let quiesced_reported = + report_quiesced(store, &pause.run_id, quiesced_at, config.store_timeout).await; if !quiesced_reported { warn!( - ducklake_maintenance = %name, run_id = %pause.run_id, - timeout_ms = config.kubernetes_api_timeout.as_millis() as u64, + timeout_ms = config.store_timeout.as_millis() as u64, "ducklake external maintenance quiesced status was not confirmed; keeping foreground \ - mutations paused until the status patch succeeds or the pause expires: \ - ducklake_maintenance={}, run_id={}, timeout_ms={}", - name, + mutations paused until the status patch succeeds or the pause expires: run_id={}, \ + timeout_ms={}", pause.run_id, - config.kubernetes_api_timeout.as_millis() + config.store_timeout.as_millis() ); } @@ -332,11 +277,13 @@ async fn reconcile_pause( }); } -async fn release_expired_pause_if_needed( - api: &Api, - config: &WatcherConfig, +async fn release_expired_pause_if_needed( + store: &M, + config: &ExternalMaintenanceWatcherConfig, held_pause: &mut Option, -) { +) where + M: ExternalMaintenanceStore, +{ if held_pause.as_ref().is_none_or(|pause| pause.expires_at > Utc::now()) { return; } @@ -345,37 +292,24 @@ async fn release_expired_pause_if_needed( return; }; warn!( - ducklake_maintenance = %config.name, run_id = %expired.run_id, - "ducklake maintenance pause expired while Kubernetes API was unavailable: \ - ducklake_maintenance={}, run_id={}", - config.name, + "ducklake maintenance pause expired while external maintenance store was unavailable: \ + run_id={}", expired.run_id ); - release_held_pause(&config.name, expired, "expired"); - patch_replicator_status( - api, - &config.name, - "Running", - None, - None, - config.kubernetes_api_timeout, - ) - .await; + release_held_pause(expired, "expired"); + report_running(store, config.store_timeout).await; } -fn release_held_pause(name: &str, held: HeldPause, outcome: &'static str) { +fn release_held_pause(held: HeldPause, outcome: &'static str) { let held_ms = Utc::now().signed_duration_since(held.quiesced_at).num_milliseconds().max(0) as u64; record_external_maintenance_pause_duration(&held, outcome); info!( - ducklake_maintenance = %name, run_id = %held.run_id, outcome, held_ms, - "ducklake external maintenance pause guard released: ducklake_maintenance={}, run_id={}, \ - outcome={}, held_ms={}", - name, + "ducklake external maintenance pause guard released: run_id={}, outcome={}, held_ms={}", held.run_id, outcome, held_ms @@ -393,23 +327,22 @@ fn record_external_maintenance_pause_duration(held: &HeldPause, outcome: &'stati .record(duration_seconds); } -async fn maybe_request_operations( - api: &Api, - name: &str, +async fn maybe_request_operations( + store: &M, destination: &DuckLakeDestination, - resource: &DynamicObject, - config: &WatcherConfig, + state: &ExternalMaintenanceState, + config: &ExternalMaintenanceWatcherConfig, ) where S: StateStore + SchemaStore + Clone + Send + Sync + 'static, + M: ExternalMaintenanceStore, { - if active_run_exists(resource) { + if state.active_run.is_some() { return; } - if completed_run_in_cooldown(resource, config.request_cooldown) { + if completed_run_in_cooldown(state, config.request_cooldown) { return; } - let policy = operation_policy(resource); let requested = match destination .sample_external_maintenance_operations( config.inline_flush_min_inlined_bytes, @@ -418,27 +351,30 @@ async fn maybe_request_operations( .await { Ok(mut operations) => { - operations.inline_flush &= policy.inline_flush_enabled; - operations.rewrite_data_files &= policy.rewrite_data_files_enabled; - operations.expire_snapshots &= policy.expire_snapshots_enabled; + operations.inline_flush &= state.operation_policy.inline_flush_enabled; + operations.merge_adjacent_files &= state.operation_policy.merge_adjacent_files_enabled; + operations.rewrite_data_files &= state.operation_policy.rewrite_data_files_enabled; + operations.expire_snapshots &= state.operation_policy.expire_snapshots_enabled; + operations.cleanup_old_files &= state.operation_policy.cleanup_old_files_enabled; operations } Err(error) => { warn!( - error = ?error, - ducklake_maintenance = %name, - "failed to sample ducklake external maintenance operations" + error = %error, + "failed to sample ducklake external maintenance operations: error={}", + error ); return; } }; - if !requested.inline_flush && !requested.rewrite_data_files && !requested.expire_snapshots { + if requested.is_empty() { debug!( - ducklake_maintenance = %name, inline_flush = requested.inline_flush, + merge_adjacent_files = requested.merge_adjacent_files, rewrite_data_files = requested.rewrite_data_files, expire_snapshots = requested.expire_snapshots, + cleanup_old_files = requested.cleanup_old_files, inline_flush_min_inlined_bytes = config.inline_flush_min_inlined_bytes, rewrite_data_files_min_active_data_files = config.rewrite_data_files_min_active_data_files, @@ -447,134 +383,172 @@ async fn maybe_request_operations( return; } - let already_requested = requested_operations(resource); + let already_requested = state + .operation_request + .as_ref() + .map_or(ExternalMaintenanceOperations::default(), |request| request.operations); if already_requested.covers(requested) { debug!( - ducklake_maintenance = %name, inline_flush = requested.inline_flush, + merge_adjacent_files = requested.merge_adjacent_files, rewrite_data_files = requested.rewrite_data_files, expire_snapshots = requested.expire_snapshots, + cleanup_old_files = requested.cleanup_old_files, "ducklake external maintenance request already exists" ); return; } info!( - ducklake_maintenance = %name, inline_flush = requested.inline_flush, + merge_adjacent_files = requested.merge_adjacent_files, rewrite_data_files = requested.rewrite_data_files, expire_snapshots = requested.expire_snapshots, + cleanup_old_files = requested.cleanup_old_files, inline_flush_min_inlined_bytes = config.inline_flush_min_inlined_bytes, rewrite_data_files_min_active_data_files = config.rewrite_data_files_min_active_data_files, "ducklake external maintenance requesting operations" ); - if patch_operation_requests(api, name, requested, config).await { - record_external_maintenance_triggers(requested, already_requested); - } -} -async fn patch_replicator_status( - api: &Api, - name: &str, - state: &str, - observed_run_id: Option<&str>, - quiesced_at: Option>, - timeout: Duration, -) -> bool { - let patch = json!({ - "status": { - "replicator": { - "state": state, - "observedRunId": observed_run_id, - "quiescedAt": quiesced_at.map(|time| time.to_rfc3339()), - } + let request = ExternalMaintenanceOperationRequest { + operations: requested, + inline_flush_min_inlined_bytes: Some(config.inline_flush_min_inlined_bytes), + rewrite_data_files_min_active_data_files: Some( + config.rewrite_data_files_min_active_data_files, + ), + requested_at: Utc::now(), + }; + match time::timeout(config.store_timeout, store.request_operations(request)).await { + Ok(Ok(ExternalMaintenanceRequestOutcome::Created)) => { + record_external_maintenance_triggers(requested, already_requested); + } + Ok(Ok(ExternalMaintenanceRequestOutcome::AlreadyCovered)) => { + debug!("ducklake external maintenance request became already covered"); + } + Ok(Ok(ExternalMaintenanceRequestOutcome::RejectedActiveRun)) => { + debug!("ducklake external maintenance request rejected because an active run exists"); + } + Ok(Ok(ExternalMaintenanceRequestOutcome::MissingState)) => { + debug!("ducklake external maintenance request ignored because state is missing"); } - }); - let params = PatchParams::default(); - - match time::timeout(timeout, api.patch_status(name, ¶ms, &Patch::Merge(&patch))).await { - Ok(Ok(_)) => true, Ok(Err(error)) => { warn!( error = %error, - ducklake_maintenance = %name, - state, - "failed to patch ducklake maintenance replicator status: \ - ducklake_maintenance={}, state={}, error={}", - name, - state, + "failed to request ducklake external maintenance operations: error={}", error ); - false } Err(_) => { warn!( - ducklake_maintenance = %name, - state, - timeout_ms = timeout.as_millis() as u64, - "timed out patching ducklake maintenance replicator status: \ - ducklake_maintenance={}, state={}, timeout_ms={}", - name, - state, - timeout.as_millis() + timeout_ms = config.store_timeout.as_millis() as u64, + "timed out requesting ducklake external maintenance operations: timeout_ms={}", + config.store_timeout.as_millis() ); - false } } } -async fn patch_operation_requests( - api: &Api, - name: &str, - requested: ExternalMaintenanceOperations, - config: &WatcherConfig, -) -> bool { - let patch = json!({ - "status": { - "operationRequests": { - "inlineFlush": requested.inline_flush, - "rewriteDataFiles": requested.rewrite_data_files, - "expireSnapshots": requested.expire_snapshots, - "inlineFlushMinInlinedBytes": config.inline_flush_min_inlined_bytes, - "rewriteDataFilesMinActiveDataFiles": config.rewrite_data_files_min_active_data_files, - "requestedAt": Utc::now().to_rfc3339(), - } - } - }); - let params = PatchParams::default(); +async fn report_pausing(store: &M, run_id: &str, timeout: Duration) +where + M: ExternalMaintenanceStore, +{ + report_replicator_status( + store, + ExternalMaintenanceReplicatorStatus { + state: ExternalMaintenanceReplicatorState::Pausing, + observed_run_id: Some(run_id.to_owned()), + quiesced_at: None, + }, + timeout, + ) + .await; +} - match time::timeout( - config.kubernetes_api_timeout, - api.patch_status(name, ¶ms, &Patch::Merge(&patch)), +async fn report_quiesced( + store: &M, + run_id: &str, + quiesced_at: DateTime, + timeout: Duration, +) -> bool +where + M: ExternalMaintenanceStore, +{ + report_replicator_status( + store, + ExternalMaintenanceReplicatorStatus { + state: ExternalMaintenanceReplicatorState::Quiesced, + observed_run_id: Some(run_id.to_owned()), + quiesced_at: Some(quiesced_at), + }, + timeout, ) .await - { - Ok(Ok(_)) => true, +} + +async fn report_running(store: &M, timeout: Duration) +where + M: ExternalMaintenanceStore, +{ + report_replicator_status( + store, + ExternalMaintenanceReplicatorStatus { + state: ExternalMaintenanceReplicatorState::Running, + observed_run_id: None, + quiesced_at: None, + }, + timeout, + ) + .await; +} + +async fn report_replicator_status( + store: &M, + status: ExternalMaintenanceReplicatorStatus, + timeout: Duration, +) -> bool +where + M: ExternalMaintenanceStore, +{ + let state = status.state; + match time::timeout(timeout, store.report_replicator_status(status)).await { + Ok(Ok(())) => true, Ok(Err(error)) => { warn!( error = %error, - ducklake_maintenance = %name, - "failed to patch ducklake maintenance operation request: \ - ducklake_maintenance={}, error={}", - name, + state = ?state, + "failed to report ducklake maintenance replicator status: state={:?}, error={}", + state, error ); false } Err(_) => { warn!( - ducklake_maintenance = %name, - timeout_ms = config.kubernetes_api_timeout.as_millis() as u64, - "timed out patching ducklake maintenance operation request: \ - ducklake_maintenance={}, timeout_ms={}", - name, - config.kubernetes_api_timeout.as_millis() + state = ?state, + timeout_ms = timeout.as_millis() as u64, + "timed out reporting ducklake maintenance replicator status: state={:?}, \ + timeout_ms={}", + state, + timeout.as_millis() ); false } } } +fn completed_run_in_cooldown(state: &ExternalMaintenanceState, cooldown: Duration) -> bool { + if cooldown.is_zero() { + return false; + } + + let Some(completed_at) = state.last_completed_at else { + return false; + }; + + let elapsed = Utc::now().signed_duration_since(completed_at); + elapsed.to_std().is_ok_and(|elapsed| elapsed < cooldown) +} + fn record_external_maintenance_triggers( requested: ExternalMaintenanceOperations, already_requested: ExternalMaintenanceOperations, @@ -606,166 +580,3 @@ fn record_external_maintenance_triggers( .increment(1); } } - -fn requested_operations(resource: &DynamicObject) -> ExternalMaintenanceOperations { - let Some(requests) = - resource.data.get("status").and_then(|status| status.get("operationRequests")) - else { - return ExternalMaintenanceOperations::default(); - }; - - ExternalMaintenanceOperations { - inline_flush: requests - .get("inlineFlush") - .and_then(serde_json::Value::as_bool) - .unwrap_or(false), - rewrite_data_files: requests - .get("rewriteDataFiles") - .and_then(serde_json::Value::as_bool) - .unwrap_or(false), - expire_snapshots: requests - .get("expireSnapshots") - .and_then(serde_json::Value::as_bool) - .unwrap_or(false), - } -} - -fn active_pause_request(resource: &DynamicObject) -> Option { - let pause = resource.data.get("status")?.get("pauseRequest")?; - if pause.is_null() { - return None; - } - - let run_id = pause.get("runId")?.as_str()?.to_owned(); - let expires_at = - DateTime::parse_from_rfc3339(pause.get("expiresAt")?.as_str()?).ok()?.with_timezone(&Utc); - - Some(PauseRequest { run_id, expires_at }) -} - -fn active_run_exists(resource: &DynamicObject) -> bool { - resource - .data - .get("status") - .and_then(|status| status.get("activeRun")) - .is_some_and(|active_run| !active_run.is_null()) -} - -fn completed_run_in_cooldown(resource: &DynamicObject, cooldown: Duration) -> bool { - if cooldown.is_zero() { - return false; - } - - let Some(completed_at) = resource - .data - .get("status") - .and_then(|status| status.get("lastCompletedRun")) - .and_then(|run| run.get("completedAt")) - .and_then(serde_json::Value::as_str) - .and_then(|value| DateTime::parse_from_rfc3339(value).ok()) - .map(|time| time.with_timezone(&Utc)) - else { - return false; - }; - - let elapsed = Utc::now().signed_duration_since(completed_at); - elapsed.to_std().is_ok_and(|elapsed| elapsed < cooldown) -} - -fn operation_policy(resource: &DynamicObject) -> OperationPolicy { - let operations = resource.data.get("spec").and_then(|spec| spec.get("operations")); - let inline_flush = operations.and_then(|ops| ops.get("inlineFlush")); - let rewrite_data_files = operations.and_then(|ops| ops.get("rewriteDataFiles")); - let expire_snapshots = operations.and_then(|ops| ops.get("expireSnapshots")); - - OperationPolicy { - inline_flush_enabled: inline_flush - .and_then(|value| value.get("enabled")) - .and_then(serde_json::Value::as_bool) - .unwrap_or(true), - rewrite_data_files_enabled: rewrite_data_files - .and_then(|value| value.get("enabled")) - .and_then(serde_json::Value::as_bool) - .unwrap_or(true), - expire_snapshots_enabled: expire_snapshots - .and_then(|value| value.get("enabled")) - .and_then(serde_json::Value::as_bool) - .unwrap_or(false), - } -} - -fn ducklake_maintenance_api_resource() -> ApiResource { - let gvk = GroupVersionKind::gvk("etl.supabase.com", "v1alpha1", "DuckLakeMaintenance"); - ApiResource::from_gvk(&gvk) -} - -impl WatcherConfig { - fn from_env() -> Option { - let name = env::var(CR_NAME_ENV).ok().filter(|value| !value.is_empty())?; - let namespace = env::var(CR_NAMESPACE_ENV).ok().filter(|value| !value.is_empty())?; - let poll_seconds = env::var(POLL_SECONDS_ENV) - .ok() - .and_then(|value| value.parse::().ok()) - .filter(|seconds| *seconds > 0) - .unwrap_or(DEFAULT_POLL_SECONDS); - let inline_flush_min_inlined_bytes = env::var(INLINE_FLUSH_MIN_INLINED_BYTES_ENV) - .ok() - .and_then(|value| value.parse::().ok()) - .unwrap_or(DEFAULT_INLINE_FLUSH_MIN_INLINED_BYTES); - let rewrite_data_files_min_active_data_files = - env::var(REWRITE_DATA_FILES_MIN_ACTIVE_DATA_FILES_ENV) - .ok() - .and_then(|value| value.parse::().ok()) - .unwrap_or(DEFAULT_REWRITE_DATA_FILES_MIN_ACTIVE_DATA_FILES); - let request_cooldown = env::var(REQUEST_COOLDOWN_SECONDS_ENV) - .ok() - .and_then(|value| value.parse::().ok()) - .unwrap_or(DEFAULT_REQUEST_COOLDOWN_SECONDS); - let kubernetes_api_timeout = env::var(KUBERNETES_API_TIMEOUT_SECONDS_ENV) - .ok() - .and_then(|value| value.parse::().ok()) - .filter(|seconds| *seconds > 0) - .unwrap_or(DEFAULT_KUBERNETES_API_TIMEOUT_SECONDS); - - Some(Self { - name, - namespace, - poll_interval: Duration::from_secs(poll_seconds), - request_cooldown: Duration::from_secs(request_cooldown), - kubernetes_api_timeout: Duration::from_secs(kubernetes_api_timeout), - inline_flush_min_inlined_bytes, - rewrite_data_files_min_active_data_files, - }) - } -} - -#[cfg(test)] -mod tests { - use serde_json::json; - - use super::*; - - #[test] - fn operation_policy_defaults_expire_snapshots_disabled() { - let resource: DynamicObject = serde_json::from_value(json!({ - "apiVersion": "etl.supabase.com/v1alpha1", - "kind": "DuckLakeMaintenance", - "metadata": { - "name": "pipeline-maintenance" - }, - "spec": { - "operations": { - "inlineFlush": {}, - "rewriteDataFiles": {} - } - } - })) - .unwrap(); - - let policy = operation_policy(&resource); - - assert!(policy.inline_flush_enabled); - assert!(policy.rewrite_data_files_enabled); - assert!(!policy.expire_snapshots_enabled); - } -} diff --git a/crates/etl-destinations/src/ducklake/maintenance_runner.rs b/crates/etl-destinations/src/ducklake/maintenance_runner.rs deleted file mode 100644 index 8e30d52eb..000000000 --- a/crates/etl-destinations/src/ducklake/maintenance_runner.rs +++ /dev/null @@ -1,898 +0,0 @@ -//! One-shot DuckLake maintenance execution for Kubernetes maintenance jobs. - -use std::sync::Arc; - -use etl::{ - error::{ErrorKind, EtlResult}, - etl_error, -}; -use metrics::histogram; -use pg_escape::{quote_identifier, quote_literal}; -use sqlx::{AssertSqlSafe, PgPool, postgres::PgPoolOptions}; -use tokio::sync::Semaphore; -use tracing::{debug, info}; -use url::Url; - -use crate::ducklake::{ - LAKE_CATALOG, S3Config, - client::{ - DuckDbBlockingOperationKind, DuckLakeConnectionManager, DuckLakeInterruptRegistry, - build_warm_ducklake_pool, format_query_error_detail, run_duckdb_blocking, - }, - config::{ - MAINTENANCE_TARGET_FILE_SIZE, build_setup_plan, current_duckdb_extension_strategy, - maintenance_target_file_size_sql, - }, - inline_size::DuckLakePendingInlineSizeSampler, - metrics::{ - DuckLakeTableStorageMetrics, ETL_DUCKLAKE_INLINE_FLUSH_DURATION_SECONDS, - ETL_DUCKLAKE_INLINE_FLUSH_ROWS, RESULT_LABEL, query_table_storage_metrics, - resolve_ducklake_metadata_schema_blocking, - }, -}; - -#[derive(Clone)] -struct DuckDbMaintenanceExecutor { - pool: Arc>, - blocking_slots: Arc, -} - -impl DuckDbMaintenanceExecutor { - async fn run(&self, operation: F) -> EtlResult - where - R: Send + 'static, - F: FnOnce(&duckdb::Connection) -> EtlResult + Send + 'static, - { - run_duckdb_blocking( - Arc::clone(&self.pool), - Arc::clone(&self.blocking_slots), - DuckDbBlockingOperationKind::Maintenance, - operation, - ) - .await - } -} - -/// Configuration for one external DuckLake maintenance run. -#[derive(Clone, Debug)] -pub struct DuckLakeMaintenanceConfig { - /// DuckLake PostgreSQL catalog URL. - pub catalog_url: Url, - /// DuckLake data path. - pub data_path: Url, - /// DuckDB connection pool size for the one-shot runner. - pub pool_size: u32, - /// Optional S3-compatible storage config. - pub s3: Option, - /// Optional DuckLake metadata schema. - pub metadata_schema: Option, - /// Optional DuckDB memory cache limit. - pub duckdb_memory_cache_limit: Option, - /// DuckLake `target_file_size` used by compaction. - pub maintenance_target_file_size: Option, - /// Inline flush operation config. - pub inline_flush: InlineFlushMaintenanceConfig, - /// Merge-adjacent-files operation config. - pub merge_adjacent_files: MergeAdjacentFilesMaintenanceConfig, - /// Rewrite-data-files operation config. - pub rewrite_data_files: RewriteDataFilesMaintenanceConfig, - /// Snapshot-expiration operation config. - pub expire_snapshots: ExpireSnapshotsMaintenanceConfig, - /// Old-file cleanup operation config. - pub cleanup_old_files: CleanupOldFilesMaintenanceConfig, -} - -/// Inline flush operation config. -#[derive(Clone, Copy, Debug)] -pub struct InlineFlushMaintenanceConfig { - /// Whether inline flush is enabled. - pub enabled: bool, - /// Minimum pending inlined bytes before flushing a table. - pub min_inlined_bytes: u64, -} - -/// Merge-adjacent-files operation config. -#[derive(Clone, Debug)] -pub struct MergeAdjacentFilesMaintenanceConfig { - /// Whether merge-adjacent-files is enabled. - pub enabled: bool, - /// Maximum compacted output files per table. - pub max_compacted_files: u32, - /// Maximum tables selected in one run. - pub max_tables_per_run: u32, - /// Target file size used during compaction. - pub target_file_size: String, -} - -/// Rewrite-data-files operation config. -#[derive(Clone, Copy, Debug)] -pub struct RewriteDataFilesMaintenanceConfig { - /// Whether rewrite-data-files is enabled. - pub enabled: bool, - /// Minimum active data-file count before rewrite is attempted. - pub min_active_data_files: i64, - /// Maximum tables selected in one run. - pub max_tables_per_run: u32, -} - -/// Snapshot-expiration operation config. -#[derive(Clone, Debug)] -pub struct ExpireSnapshotsMaintenanceConfig { - /// Whether snapshot expiration is enabled. - pub enabled: bool, - /// Retention window passed to DuckLake. - pub older_than: String, -} - -/// Old-file cleanup operation config. -#[derive(Clone, Debug)] -pub struct CleanupOldFilesMaintenanceConfig { - /// Whether old-file cleanup is enabled. - pub enabled: bool, - /// Retention window passed to DuckLake. - pub older_than: String, -} - -/// Structured outcome for one external maintenance run. -#[derive(Clone, Debug, Default, PartialEq)] -pub struct DuckLakeMaintenanceOutcome { - /// Tables whose inline data was flushed. - pub inline_flush_tables: u32, - /// Rows flushed from inlined storage. - pub inline_flush_rows: u64, - /// Tables passed to merge-adjacent-files. - pub merge_adjacent_files_tables: u32, - /// Files created by merge-adjacent-files. - pub merge_adjacent_files_created: u64, - /// Tables passed to rewrite-data-files. - pub rewrite_data_files_tables: u32, - /// Files created by rewrite-data-files. - pub rewrite_data_files_created: u64, - /// Snapshots expired by snapshot expiration. - pub expired_snapshots: u64, - /// Files removed by old-file cleanup. - pub cleaned_up_files: u64, -} - -impl DuckLakeMaintenanceOutcome { - /// Returns whether any operation did work. - pub fn applied(&self) -> bool { - self.inline_flush_rows > 0 - || self.merge_adjacent_files_created > 0 - || self.rewrite_data_files_tables > 0 - || self.rewrite_data_files_created > 0 - || self.expired_snapshots > 0 - || self.cleaned_up_files > 0 - } -} - -/// Runs one external DuckLake maintenance attempt. -pub async fn run_maintenance_once( - config: DuckLakeMaintenanceConfig, -) -> EtlResult { - validate_config(&config)?; - let cleanup_old_files_enabled = - config.cleanup_old_files.enabled || config.rewrite_data_files.enabled; - - info!( - pool_size = config.pool_size, - metadata_schema = config.metadata_schema.as_deref(), - inline_flush_enabled = config.inline_flush.enabled, - inline_flush_min_inlined_bytes = config.inline_flush.min_inlined_bytes, - merge_adjacent_files_enabled = config.merge_adjacent_files.enabled, - merge_adjacent_files_max_compacted_files = config.merge_adjacent_files.max_compacted_files, - merge_adjacent_files_max_tables_per_run = config.merge_adjacent_files.max_tables_per_run, - merge_adjacent_files_target_file_size = %config.merge_adjacent_files.target_file_size, - rewrite_data_files_enabled = config.rewrite_data_files.enabled, - rewrite_data_files_min_active_data_files = config.rewrite_data_files.min_active_data_files, - rewrite_data_files_max_tables_per_run = config.rewrite_data_files.max_tables_per_run, - expire_snapshots_enabled = config.expire_snapshots.enabled, - expire_snapshots_older_than = %config.expire_snapshots.older_than, - cleanup_old_files_enabled, - cleanup_old_files_explicitly_enabled = config.cleanup_old_files.enabled, - cleanup_old_files_older_than = %config.cleanup_old_files.older_than, - "ducklake external maintenance runner configured" - ); - - let duckdb = open_maintenance_executor(&config).await?; - let metadata_schema = match config.metadata_schema.clone() { - Some(metadata_schema) => metadata_schema, - None => resolve_metadata_schema(&duckdb).await?, - }; - info!( - metadata_schema = %metadata_schema, - "ducklake external maintenance metadata schema resolved" - ); - let metadata_pg_pool = PgPoolOptions::new() - .max_connections(1) - .connect_lazy(config.catalog_url.as_str()) - .map_err(|source| { - etl_error!( - ErrorKind::DestinationConnectionFailed, - "DuckLake catalog metadata pool configuration failed", - source: source - ) - })?; - let table_names = list_ducklake_tables(&metadata_pg_pool, &metadata_schema).await?; - info!( - table_count = table_names.len(), - tables = ?table_names, - "ducklake external maintenance discovered active tables" - ); - let mut outcome = DuckLakeMaintenanceOutcome::default(); - - if config.inline_flush.enabled { - run_inline_flush( - &duckdb, - &metadata_pg_pool, - &metadata_schema, - &table_names, - config.inline_flush, - &mut outcome, - ) - .await?; - } - - if config.merge_adjacent_files.enabled { - run_merge_adjacent_files( - &duckdb, - &metadata_pg_pool, - &metadata_schema, - &table_names, - &config.merge_adjacent_files, - &mut outcome, - ) - .await?; - } - - if config.rewrite_data_files.enabled { - merge_adjacent_files_for_rewrite(&duckdb).await?; - run_rewrite_data_files( - &duckdb, - &metadata_pg_pool, - &metadata_schema, - &table_names, - config.rewrite_data_files, - &mut outcome, - ) - .await?; - } - - if config.expire_snapshots.enabled { - run_expire_snapshots(&duckdb, &config.expire_snapshots, &mut outcome).await?; - } - - if cleanup_old_files_enabled { - run_cleanup_old_files(&duckdb, &config.cleanup_old_files, &mut outcome).await?; - } - - info!(outcome = ?outcome, applied = outcome.applied(), "ducklake external maintenance completed"); - Ok(outcome) -} - -/// Validates one maintenance runner config. -fn validate_config(config: &DuckLakeMaintenanceConfig) -> EtlResult<()> { - if !matches!(config.catalog_url.scheme(), "postgres" | "postgresql") { - return Err(etl_error!( - ErrorKind::ConfigError, - "DuckLake external maintenance requires a PostgreSQL catalog", - format!("unsupported catalog URL scheme `{}`", config.catalog_url.scheme()) - )); - } - if config.pool_size == 0 { - return Err(etl_error!( - ErrorKind::ConfigError, - "DuckLake external maintenance pool size must be greater than zero" - )); - } - Ok(()) -} - -/// Opens initialized DuckDB connections for maintenance. -async fn open_maintenance_executor( - config: &DuckLakeMaintenanceConfig, -) -> EtlResult { - let extension_strategy = current_duckdb_extension_strategy()?; - let target_file_size = config - .maintenance_target_file_size - .as_deref() - .or(Some(config.merge_adjacent_files.target_file_size.as_str())) - .unwrap_or(MAINTENANCE_TARGET_FILE_SIZE); - info!(target_file_size, "opening ducklake external maintenance connection"); - let setup_plan = Arc::new(build_setup_plan( - &config.catalog_url, - &config.data_path, - config.s3.as_ref(), - config.metadata_schema.as_deref(), - None, - )?); - let manager = DuckLakeConnectionManager { - setup_plan, - disable_extension_autoload: extension_strategy.disables_autoload(), - interrupt_registry: Arc::new(DuckLakeInterruptRegistry::default()), - #[cfg(feature = "test-utils")] - open_count: Arc::new(std::sync::atomic::AtomicUsize::new(0)), - }; - let pool = Arc::new( - build_warm_ducklake_pool(manager, config.pool_size, "external-maintenance").await?, - ); - let blocking_slots = Arc::new(Semaphore::new(config.pool_size as usize)); - let executor = DuckDbMaintenanceExecutor { pool, blocking_slots }; - let sql = maintenance_target_file_size_sql(Some(target_file_size)); - executor - .run(move |conn| { - conn.execute_batch(&sql).map_err(|source| { - etl_error!( - ErrorKind::DestinationQueryFailed, - "DuckLake target_file_size configuration failed", - format_query_error_detail(&sql), - source: source - ) - })?; - Ok(()) - }) - .await?; - info!(target_file_size, "ducklake external maintenance connection ready"); - Ok(executor) -} - -/// Resolves the hidden DuckLake metadata schema. -async fn resolve_metadata_schema(duckdb: &DuckDbMaintenanceExecutor) -> EtlResult { - duckdb.run(resolve_ducklake_metadata_schema_blocking).await -} - -/// Lists active DuckLake table names from the metadata catalog. -async fn list_ducklake_tables( - metadata_pg_pool: &PgPool, - metadata_schema: &str, -) -> EtlResult> { - let sql = format!( - "SELECT table_name FROM {}.{} WHERE end_snapshot IS NULL ORDER BY table_name", - quote_identifier(metadata_schema), - quote_identifier("ducklake_table") - ); - let rows: Vec<(String,)> = - sqlx::query_as(AssertSqlSafe(sql)).fetch_all(metadata_pg_pool).await.map_err(|source| { - etl_error!( - ErrorKind::DestinationQueryFailed, - "DuckLake table list query failed", - format!("metadata_schema={metadata_schema}"), - source: source - ) - })?; - Ok(rows.into_iter().map(|(table_name,)| table_name).collect()) -} - -/// Runs inline flush for tables that crossed the pending-inline threshold. -async fn run_inline_flush( - duckdb: &DuckDbMaintenanceExecutor, - metadata_pg_pool: &PgPool, - metadata_schema: &str, - table_names: &[String], - config: InlineFlushMaintenanceConfig, - outcome: &mut DuckLakeMaintenanceOutcome, -) -> EtlResult<()> { - info!( - min_inlined_bytes = config.min_inlined_bytes, - table_count = table_names.len(), - "ducklake inline flush evaluation started" - ); - let sampler = - DuckLakePendingInlineSizeSampler::new(metadata_schema.to_owned(), metadata_pg_pool.clone()); - for table_name in table_names { - let sizes = sampler.sample_table(table_name).await?; - if sizes.inlined_bytes < config.min_inlined_bytes { - info!( - table = %table_name, - inlined_bytes = sizes.inlined_bytes, - min_inlined_bytes = config.min_inlined_bytes, - "ducklake inline flush skipped below threshold" - ); - continue; - } - info!( - table = %table_name, - inlined_bytes = sizes.inlined_bytes, - min_inlined_bytes = config.min_inlined_bytes, - "ducklake inline flush executing" - ); - let table_name_for_query = table_name.clone(); - let rows = - duckdb.run(move |conn| flush_table_inlined_data(conn, &table_name_for_query)).await?; - outcome.inline_flush_tables = outcome.inline_flush_tables.saturating_add(1); - outcome.inline_flush_rows = outcome.inline_flush_rows.saturating_add(rows); - info!( - table = %table_name, - rows, - "ducklake inline flush completed" - ); - } - info!( - inline_flush_tables = outcome.inline_flush_tables, - inline_flush_rows = outcome.inline_flush_rows, - "ducklake inline flush evaluation finished" - ); - Ok(()) -} - -/// Runs bounded merge-adjacent-files on selected tables. -async fn run_merge_adjacent_files( - duckdb: &DuckDbMaintenanceExecutor, - metadata_pg_pool: &PgPool, - metadata_schema: &str, - table_names: &[String], - config: &MergeAdjacentFilesMaintenanceConfig, - outcome: &mut DuckLakeMaintenanceOutcome, -) -> EtlResult<()> { - info!( - max_compacted_files = config.max_compacted_files, - max_tables_per_run = config.max_tables_per_run, - table_count = table_names.len(), - "ducklake merge-adjacent-files evaluation started" - ); - let selected = select_merge_tables( - metadata_pg_pool, - metadata_schema, - table_names, - config.max_tables_per_run, - ) - .await?; - info!( - selected_tables = ?selected, - selected_count = selected.len(), - "ducklake merge-adjacent-files selected tables" - ); - for table_name in selected { - info!( - table = %table_name, - max_compacted_files = config.max_compacted_files, - "ducklake merge-adjacent-files executing" - ); - let table_name_for_query = table_name.clone(); - let max_compacted_files = config.max_compacted_files; - let files_created = duckdb - .run(move |conn| merge_adjacent_files(conn, &table_name_for_query, max_compacted_files)) - .await?; - outcome.merge_adjacent_files_tables = outcome.merge_adjacent_files_tables.saturating_add(1); - outcome.merge_adjacent_files_created = - outcome.merge_adjacent_files_created.saturating_add(files_created); - info!( - table = %table_name, - files_created, - "ducklake merge-adjacent-files completed" - ); - } - Ok(()) -} - -/// Runs DuckLake's whole-lake adjacent-file merge before rewrite-data-files. -async fn merge_adjacent_files_for_rewrite(duckdb: &DuckDbMaintenanceExecutor) -> EtlResult<()> { - let sql = format!("CALL ducklake_merge_adjacent_files({});", quote_literal(LAKE_CATALOG)); - info!( - sql = %sql, - "ducklake rewrite-triggered merge-adjacent-files executing" - ); - duckdb - .run(move |conn| { - conn.execute_batch(&sql).map_err(|source| { - etl_error!( - ErrorKind::DestinationQueryFailed, - "DuckLake rewrite-triggered merge adjacent files failed", - format_query_error_detail(&sql), - source: source - ) - })?; - Ok(()) - }) - .await?; - info!("ducklake rewrite-triggered merge-adjacent-files completed"); - Ok(()) -} - -/// Runs bounded rewrite-data-files on selected tables. -async fn run_rewrite_data_files( - duckdb: &DuckDbMaintenanceExecutor, - metadata_pg_pool: &PgPool, - metadata_schema: &str, - table_names: &[String], - config: RewriteDataFilesMaintenanceConfig, - outcome: &mut DuckLakeMaintenanceOutcome, -) -> EtlResult<()> { - info!( - min_active_data_files = config.min_active_data_files, - max_tables_per_run = config.max_tables_per_run, - table_count = table_names.len(), - "ducklake rewrite-data-files evaluation started" - ); - let selected = select_rewrite_tables( - metadata_pg_pool, - metadata_schema, - table_names, - config.min_active_data_files, - config.max_tables_per_run, - ) - .await?; - info!( - selected_tables = ?selected, - selected_count = selected.len(), - "ducklake rewrite-data-files selected tables" - ); - for table_name in selected { - info!( - table = %table_name, - "ducklake rewrite-data-files executing" - ); - let table_name_for_query = table_name.clone(); - let files_created = - duckdb.run(move |conn| rewrite_data_files(conn, &table_name_for_query)).await?; - outcome.rewrite_data_files_tables = outcome.rewrite_data_files_tables.saturating_add(1); - outcome.rewrite_data_files_created = - outcome.rewrite_data_files_created.saturating_add(files_created); - info!( - table = %table_name, - files_created, - "ducklake rewrite-data-files completed" - ); - } - Ok(()) -} - -/// Runs DuckLake snapshot expiration. -async fn run_expire_snapshots( - duckdb: &DuckDbMaintenanceExecutor, - config: &ExpireSnapshotsMaintenanceConfig, - outcome: &mut DuckLakeMaintenanceOutcome, -) -> EtlResult<()> { - info!( - older_than = %config.older_than, - "ducklake expire-snapshots executing" - ); - let older_than = config.older_than.clone(); - let expired_snapshots = duckdb.run(move |conn| expire_snapshots(conn, &older_than)).await?; - outcome.expired_snapshots = outcome.expired_snapshots.saturating_add(expired_snapshots); - info!( - older_than = %config.older_than, - expired_snapshots, - "ducklake expire-snapshots completed" - ); - Ok(()) -} - -/// Runs DuckLake old-file cleanup. -async fn run_cleanup_old_files( - duckdb: &DuckDbMaintenanceExecutor, - config: &CleanupOldFilesMaintenanceConfig, - outcome: &mut DuckLakeMaintenanceOutcome, -) -> EtlResult<()> { - info!( - older_than = %config.older_than, - "ducklake cleanup-old-files executing" - ); - let older_than = config.older_than.clone(); - let cleaned_up_files = duckdb.run(move |conn| cleanup_old_files(conn, &older_than)).await?; - outcome.cleaned_up_files = outcome.cleaned_up_files.saturating_add(cleaned_up_files); - info!( - older_than = %config.older_than, - cleaned_up_files, - "ducklake cleanup-old-files completed" - ); - Ok(()) -} - -/// Selects tables with small-file pressure for merge-adjacent-files. -async fn select_merge_tables( - metadata_pg_pool: &PgPool, - metadata_schema: &str, - table_names: &[String], - max_tables_per_run: u32, -) -> EtlResult> { - let mut selected = Vec::new(); - for table_name in table_names { - if is_etl_internal_table(table_name) { - info!( - table = %table_name, - "ducklake rewrite-data-files table skipped because it is internal ETL metadata" - ); - continue; - } - - let metrics = - query_table_storage_metrics(metadata_pg_pool, metadata_schema, table_name).await?; - if metrics.active_data_files > 1 && metrics.small_file_ratio() > 0.0 { - info!( - table = %table_name, - active_data_files = metrics.active_data_files, - small_file_ratio = metrics.small_file_ratio(), - "ducklake merge-adjacent-files table selected" - ); - selected.push(table_name.clone()); - } else { - info!( - table = %table_name, - active_data_files = metrics.active_data_files, - small_file_ratio = metrics.small_file_ratio(), - "ducklake merge-adjacent-files table skipped" - ); - } - if selected.len() >= max_tables_per_run as usize { - break; - } - } - Ok(selected) -} - -/// Selects tables with delete pressure for rewrite-data-files. -async fn select_rewrite_tables( - metadata_pg_pool: &PgPool, - metadata_schema: &str, - table_names: &[String], - min_active_data_files: i64, - max_tables_per_run: u32, -) -> EtlResult> { - let mut selected = Vec::new(); - for table_name in table_names { - if is_etl_internal_table(table_name) { - info!( - table = %table_name, - "ducklake rewrite-data-files table skipped because it is internal ETL metadata" - ); - continue; - } - - let metrics = - query_table_storage_metrics(metadata_pg_pool, metadata_schema, table_name).await?; - if should_rewrite(&metrics, min_active_data_files) { - info!( - table = %table_name, - active_data_files = metrics.active_data_files, - active_delete_files = metrics.active_delete_files, - deleted_row_ratio = metrics.deleted_row_ratio(), - min_active_data_files, - "ducklake rewrite-data-files table selected" - ); - selected.push(table_name.clone()); - } else { - info!( - table = %table_name, - active_data_files = metrics.active_data_files, - active_delete_files = metrics.active_delete_files, - deleted_row_ratio = metrics.deleted_row_ratio(), - min_active_data_files, - "ducklake rewrite-data-files table skipped" - ); - } - if selected.len() >= max_tables_per_run as usize { - break; - } - } - Ok(selected) -} - -fn is_etl_internal_table(table_name: &str) -> bool { - table_name.starts_with("__etl_") -} - -/// Returns whether a table should be rewritten. -fn should_rewrite(metrics: &DuckLakeTableStorageMetrics, min_active_data_files: i64) -> bool { - metrics.active_data_files > min_active_data_files -} - -/// Flushes inlined user data for one table. -pub(super) fn flush_table_inlined_data( - conn: &duckdb::Connection, - table_name: &str, -) -> EtlResult { - let flush_started = std::time::Instant::now(); - let sql = format!( - r#"SELECT COALESCE(SUM(rows_flushed), 0) - FROM ducklake_flush_inlined_data({}, table_name => {});"#, - quote_literal(LAKE_CATALOG), - quote_literal(table_name), - ); - let rows_flushed: i64 = conn.query_row(&sql, [], |row| row.get(0)).map_err(|source| { - etl_error!( - ErrorKind::DestinationQueryFailed, - "DuckLake inlined data flush failed", - format_query_error_detail(&sql), - source: source - ) - })?; - let rows_flushed = rows_flushed.max(0) as u64; - let flush_result = if rows_flushed > 0 { "flushed" } else { "noop" }; - histogram!( - ETL_DUCKLAKE_INLINE_FLUSH_ROWS, - RESULT_LABEL => flush_result, - ) - .record(rows_flushed as f64); - histogram!( - ETL_DUCKLAKE_INLINE_FLUSH_DURATION_SECONDS, - RESULT_LABEL => flush_result, - ) - .record(flush_started.elapsed().as_secs_f64()); - - if rows_flushed > 0 { - debug!( - table = %table_name, - rows_flushed, - "ducklake inlined data flushed" - ); - } else { - debug!( - table = %table_name, - "ducklake inlined data already flushed" - ); - } - Ok(rows_flushed) -} - -/// Calls DuckLake merge-adjacent-files for one table. -fn merge_adjacent_files( - conn: &duckdb::Connection, - table_name: &str, - max_compacted_files: u32, -) -> EtlResult { - let sql = format!( - "SELECT COALESCE(SUM(files_created), 0) FROM ducklake_merge_adjacent_files({}, {}, \ - max_compacted_files => {});", - quote_literal(LAKE_CATALOG), - quote_literal(table_name), - max_compacted_files - ); - count_maintenance_files(conn, &sql, "DuckLake merge adjacent files failed") -} - -/// Calls DuckLake rewrite-data-files for one table. -fn rewrite_data_files(conn: &duckdb::Connection, table_name: &str) -> EtlResult { - let sql = format!( - "CALL ducklake_rewrite_data_files({}, {});", - quote_literal(LAKE_CATALOG), - quote_literal(table_name) - ); - conn.execute_batch(&sql).map_err(|source| { - etl_error!( - ErrorKind::DestinationQueryFailed, - "DuckLake rewrite data files failed", - format_query_error_detail(&sql), - source: source - ) - })?; - Ok(0) -} - -/// Calls DuckLake snapshot expiration. -fn expire_snapshots(conn: &duckdb::Connection, older_than: &str) -> EtlResult { - let sql = format!( - "CALL ducklake_expire_snapshots({}, older_than => CAST(now() AS TIMESTAMP) - CAST({} AS \ - INTERVAL));", - quote_literal(LAKE_CATALOG), - quote_literal(older_than), - ); - count_maintenance_rows(conn, &sql, "DuckLake expire snapshots failed") -} - -/// Calls DuckLake old-file cleanup. -fn cleanup_old_files(conn: &duckdb::Connection, older_than: &str) -> EtlResult { - let sql = format!( - "CALL ducklake_cleanup_old_files({}, older_than => CAST(now() AS TIMESTAMP) - CAST({} AS \ - INTERVAL));", - quote_literal(LAKE_CATALOG), - quote_literal(older_than), - ); - count_maintenance_rows(conn, &sql, "DuckLake cleanup old files failed") -} - -/// Counts rows returned by one DuckLake maintenance call. -fn count_maintenance_rows( - conn: &duckdb::Connection, - sql: &str, - description: &'static str, -) -> EtlResult { - let mut statement = conn.prepare(sql).map_err(|source| { - etl_error!( - ErrorKind::DestinationQueryFailed, - description, - format_query_error_detail(sql), - source: source - ) - })?; - let mut rows = statement.query([]).map_err(|source| { - etl_error!( - ErrorKind::DestinationQueryFailed, - description, - format_query_error_detail(sql), - source: source - ) - })?; - let mut count = 0u64; - - while let Some(_row) = rows.next().map_err(|source| { - etl_error!( - ErrorKind::DestinationQueryFailed, - description, - format_query_error_detail(sql), - source: source - ) - })? { - count = count.saturating_add(1); - } - - Ok(count) -} - -/// Counts files returned by one DuckLake maintenance function. -fn count_maintenance_files( - conn: &duckdb::Connection, - sql: &str, - description: &'static str, -) -> EtlResult { - let files_created: i64 = conn.query_row(sql, [], |row| row.get(0)).map_err(|source| { - etl_error!( - ErrorKind::DestinationQueryFailed, - description, - format_query_error_detail(sql), - source: source - ) - })?; - Ok(files_created.max(0) as u64) -} - -#[cfg(test)] -mod tests { - use super::*; - - fn metrics( - active_data_files: i64, - active_delete_files: i64, - deleted_rows: i64, - ) -> DuckLakeTableStorageMetrics { - DuckLakeTableStorageMetrics { - active_data_files, - active_data_bytes: 100, - small_data_files: 0, - active_data_rows: 100, - active_delete_files, - active_delete_bytes: 10, - deleted_rows, - } - } - - #[test] - fn should_rewrite_requires_only_file_count() { - assert!(!should_rewrite(&metrics(39, 1, 50), 40)); - assert!(!should_rewrite(&metrics(40, 1, 50), 40)); - assert!(should_rewrite(&metrics(41, 0, 0), 40)); - } - - #[test] - fn outcome_reports_applied_work() { - assert!(!DuckLakeMaintenanceOutcome::default().applied()); - assert!( - DuckLakeMaintenanceOutcome { - inline_flush_rows: 1, - ..DuckLakeMaintenanceOutcome::default() - } - .applied() - ); - assert!( - DuckLakeMaintenanceOutcome { - rewrite_data_files_tables: 1, - ..DuckLakeMaintenanceOutcome::default() - } - .applied() - ); - assert!( - DuckLakeMaintenanceOutcome { - expired_snapshots: 1, - ..DuckLakeMaintenanceOutcome::default() - } - .applied() - ); - assert!( - DuckLakeMaintenanceOutcome { - cleaned_up_files: 1, - ..DuckLakeMaintenanceOutcome::default() - } - .applied() - ); - } -} diff --git a/crates/etl-destinations/src/ducklake/metrics.rs b/crates/etl-destinations/src/ducklake/metrics.rs index d7358273d..fbd37ff9a 100644 --- a/crates/etl-destinations/src/ducklake/metrics.rs +++ b/crates/etl-destinations/src/ducklake/metrics.rs @@ -19,7 +19,10 @@ use tokio::{ }; use tracing::{info, warn}; -use crate::ducklake::{DuckLakeTableName, LAKE_CATALOG, client::format_query_error_detail}; +use crate::ducklake::{ + DuckLakeTableName, LAKE_CATALOG, client::format_query_error_detail, + inline_size::DuckLakePendingInlineSizeSampler, +}; static REGISTER_METRICS: Once = Once::new(); @@ -85,7 +88,6 @@ pub(crate) const SUB_BATCH_KIND_LABEL: &str = "sub_batch_kind"; pub(crate) const PREPARED_ROWS_KIND_LABEL: &str = "prepared_rows_kind"; pub(crate) const DELETE_ORIGIN_LABEL: &str = "delete_origin"; pub(crate) const RETRY_SCOPE_LABEL: &str = "retry_scope"; -pub(crate) const RESULT_LABEL: &str = "result"; pub(crate) const MAINTENANCE_OPERATION_LABEL: &str = "operation"; pub(crate) const MAINTENANCE_REASON_LABEL: &str = "reason"; pub(crate) const MAINTENANCE_OUTCOME_LABEL: &str = "outcome"; @@ -342,6 +344,8 @@ async fn run_ducklake_metrics_sampler( let mut interval = tokio::time::interval_at(Instant::now() + METRICS_POLL_INTERVAL, METRICS_POLL_INTERVAL); interval.set_missed_tick_behavior(MissedTickBehavior::Skip); + let inline_sampler = + DuckLakePendingInlineSizeSampler::new(metadata_schema.clone(), metadata_pg_pool.clone()); loop { tokio::select! { @@ -381,6 +385,14 @@ async fn run_ducklake_metrics_sampler( "ducklake table storage metrics collection failed" ); } + + if let Err(error) = inline_sampler.sample_table(&table_name).await { + warn!( + table = %table_name, + error = %error, + "ducklake table active inlined data metrics collection failed" + ); + } } } } diff --git a/crates/etl-destinations/src/ducklake/mod.rs b/crates/etl-destinations/src/ducklake/mod.rs index a3fd5ac7c..f655a0589 100644 --- a/crates/etl-destinations/src/ducklake/mod.rs +++ b/crates/etl-destinations/src/ducklake/mod.rs @@ -5,7 +5,6 @@ mod core; mod encoding; mod external_maintenance; mod inline_size; -mod maintenance_runner; mod metrics; mod schema; @@ -23,7 +22,8 @@ pub(super) type DuckLakeTableName = String; const ATTACH_DATA_INLINING_ROW_LIMIT: u64 = 10_000; pub use core::{ - DuckLakeDestination, DuckLakeExternalMaintenancePause, table_name_to_ducklake_table_name, + DuckLakeDestination, DuckLakeExternalMaintenanceConfig, DuckLakeExternalMaintenancePause, + DuckLakeMaintenanceMode, table_name_to_ducklake_table_name, }; #[cfg(feature = "test-utils")] pub use core::{ @@ -38,8 +38,16 @@ pub use batches::{ reset_ducklake_test_hooks, }; pub use config::S3Config; -pub use maintenance_runner::{ +pub use etl_maintenance::ducklake::{ CleanupOldFilesMaintenanceConfig, DuckLakeMaintenanceConfig, DuckLakeMaintenanceOutcome, ExpireSnapshotsMaintenanceConfig, InlineFlushMaintenanceConfig, MergeAdjacentFilesMaintenanceConfig, RewriteDataFilesMaintenanceConfig, run_maintenance_once, }; +pub use external_maintenance::{ + ExternalMaintenanceOperationHistory, ExternalMaintenanceOperationPolicy, + ExternalMaintenanceOperationRequest, ExternalMaintenanceOperationRun, + ExternalMaintenanceOperations, ExternalMaintenancePause, ExternalMaintenanceReplicatorState, + ExternalMaintenanceReplicatorStatus, ExternalMaintenanceRequestOutcome, ExternalMaintenanceRun, + ExternalMaintenanceState, ExternalMaintenanceStore, ExternalMaintenanceWatcherConfig, + PostgresExternalMaintenanceStore, run_external_maintenance_watcher, +}; diff --git a/crates/etl-maintenance/Cargo.toml b/crates/etl-maintenance/Cargo.toml new file mode 100644 index 000000000..4ee64f4c3 --- /dev/null +++ b/crates/etl-maintenance/Cargo.toml @@ -0,0 +1,55 @@ +[package] +name = "etl-maintenance" +version = "0.1.0" +edition.workspace = true +license.workspace = true +rust-version.workspace = true +repository.workspace = true +homepage.workspace = true + +[lib] +doctest = false + +[features] +default = [] +ducklake = [ + "dep:duckdb", + "dep:kube", + "dep:metrics", + "dep:pg_escape", + "dep:r2d2", + "dep:regex", + "dep:serde_json", + "dep:tokio-postgres", + "dep:url", +] +test-utils = [] + +[dependencies] +async-trait = { workspace = true } +chrono = { workspace = true, features = ["serde"] } +duckdb = { workspace = true, optional = true, features = ["bundled", "json", "parquet", "r2d2"] } +etl = { workspace = true } +kube = { workspace = true, optional = true, features = ["client", "rustls-tls"] } +metrics = { workspace = true, optional = true } +pg_escape = { workspace = true, optional = true } +r2d2 = { workspace = true, optional = true } +regex = { workspace = true, optional = true } +serde = { workspace = true, features = ["derive"] } +serde_json = { workspace = true, optional = true, features = ["arbitrary_precision", "std"] } +sqlx = { workspace = true, features = ["runtime-tokio", "tls-rustls", "postgres", "migrate", "chrono"] } +thiserror = { workspace = true } +tokio = { workspace = true, features = ["rt", "sync", "time"] } +tokio-postgres = { workspace = true, optional = true } +tracing = { workspace = true, default-features = true } +url = { workspace = true, optional = true } + +[dev-dependencies] +etl-telemetry = { workspace = true } +k8s-openapi = { workspace = true, features = ["latest"] } +serde_json = { workspace = true } +tempfile = { workspace = true } +tokio = { workspace = true, features = ["full"] } + +[lints] +workspace = true diff --git a/crates/etl-maintenance/migrations/postgres/20260515000000_external_maintenance_state.down.sql b/crates/etl-maintenance/migrations/postgres/20260515000000_external_maintenance_state.down.sql new file mode 100644 index 000000000..6f456f2eb --- /dev/null +++ b/crates/etl-maintenance/migrations/postgres/20260515000000_external_maintenance_state.down.sql @@ -0,0 +1 @@ +DROP TABLE IF EXISTS etl.external_maintenance_state; diff --git a/crates/etl-maintenance/migrations/postgres/20260515000000_external_maintenance_state.up.sql b/crates/etl-maintenance/migrations/postgres/20260515000000_external_maintenance_state.up.sql new file mode 100644 index 000000000..461957748 --- /dev/null +++ b/crates/etl-maintenance/migrations/postgres/20260515000000_external_maintenance_state.up.sql @@ -0,0 +1,11 @@ +CREATE TABLE IF NOT EXISTS etl.external_maintenance_state ( + pipeline_id BIGINT PRIMARY KEY, + active_run JSONB, + pause_request JSONB, + operation_request JSONB, + replicator JSONB, + last_successful_operations JSONB NOT NULL DEFAULT '{}'::jsonb, + last_completed_at TIMESTAMPTZ, + operation_policy JSONB NOT NULL DEFAULT '{}'::jsonb, + updated_at TIMESTAMPTZ NOT NULL DEFAULT now() +); diff --git a/crates/etl-maintenance/src/coordination.rs b/crates/etl-maintenance/src/coordination.rs new file mode 100644 index 000000000..4248bb527 --- /dev/null +++ b/crates/etl-maintenance/src/coordination.rs @@ -0,0 +1,611 @@ +use std::time::Duration; + +use async_trait::async_trait; +use chrono::{DateTime, Utc}; +use etl::error::EtlResult; +use serde::{Deserialize, Serialize}; + +const DEFAULT_POLL_SECONDS: u64 = 5; +const DEFAULT_INLINE_FLUSH_MIN_INLINED_BYTES: u64 = 10_000_000; +const DEFAULT_REWRITE_DATA_FILES_MIN_ACTIVE_DATA_FILES: i64 = 40; +const DEFAULT_REQUEST_COOLDOWN_SECONDS: u64 = 300; +const DEFAULT_STORE_TIMEOUT_SECONDS: u64 = 10; + +/// Backend-neutral runtime state for external maintenance. +#[derive(Clone, Debug, Default, Deserialize, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct ExternalMaintenanceState { + /// Whether the coordination backend has state for this pipeline. + pub exists: bool, + /// Currently active maintenance run, if any. + pub active_run: Option, + /// Controller-owned pause lease observed by the replicator. + pub pause_request: Option, + /// Replicator-owned operation request sampled from destination state. + pub operation_request: Option, + /// Last reported replicator maintenance state. + pub replicator: Option, + /// Last successful run per operation. + pub last_successful_operations: ExternalMaintenanceOperationHistory, + /// Last completed run timestamp, regardless of outcome. + pub last_completed_at: Option>, + /// Backend-neutral operation enablement policy. + pub operation_policy: ExternalMaintenanceOperationPolicy, +} + +impl ExternalMaintenanceState { + /// Returns a default state row marked as existing. + pub fn present() -> Self { + Self { exists: true, ..Self::default() } + } +} + +/// One active maintenance attempt. +#[derive(Clone, Debug, Deserialize, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct ExternalMaintenanceRun { + /// Stable run identifier. + pub run_id: String, + /// Run start timestamp. + pub started_at: Option>, +} + +/// Backend-neutral operation request flags. +#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize, PartialEq, Eq)] +#[serde(rename_all = "camelCase")] +pub struct ExternalMaintenanceOperations { + /// Whether inline data should be flushed. + pub inline_flush: bool, + /// Whether adjacent files should be merged. + pub merge_adjacent_files: bool, + /// Whether data files should be rewritten. + pub rewrite_data_files: bool, + /// Whether snapshots should be expired. + pub expire_snapshots: bool, + /// Whether old files should be cleaned up. + pub cleanup_old_files: bool, +} + +impl ExternalMaintenanceOperations { + /// Returns whether no operation is requested. + pub fn is_empty(self) -> bool { + !self.inline_flush + && !self.merge_adjacent_files + && !self.rewrite_data_files + && !self.expire_snapshots + && !self.cleanup_old_files + } + + /// Returns whether this operation set covers all requested operations. + pub fn covers(self, requested: Self) -> bool { + (!requested.inline_flush || self.inline_flush) + && (!requested.merge_adjacent_files || self.merge_adjacent_files) + && (!requested.rewrite_data_files || self.rewrite_data_files) + && (!requested.expire_snapshots || self.expire_snapshots) + && (!requested.cleanup_old_files || self.cleanup_old_files) + } + + /// Returns an idempotent union of two operation sets. + pub fn merge(self, requested: Self) -> Self { + Self { + inline_flush: self.inline_flush || requested.inline_flush, + merge_adjacent_files: self.merge_adjacent_files || requested.merge_adjacent_files, + rewrite_data_files: self.rewrite_data_files || requested.rewrite_data_files, + expire_snapshots: self.expire_snapshots || requested.expire_snapshots, + cleanup_old_files: self.cleanup_old_files || requested.cleanup_old_files, + } + } +} + +/// Controller-owned bounded pause lease. +#[derive(Clone, Debug, Deserialize, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct ExternalMaintenancePause { + /// Run identifier that owns the pause. + pub run_id: String, + /// Pause request timestamp. + pub requested_at: Option>, + /// Pause expiry timestamp. + pub expires_at: DateTime, +} + +/// Replicator-owned request for a future maintenance run. +#[derive(Clone, Debug, Deserialize, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct ExternalMaintenanceOperationRequest { + /// Requested operations. + pub operations: ExternalMaintenanceOperations, + /// Inline flush threshold observed by the replicator. + pub inline_flush_min_inlined_bytes: Option, + /// Rewrite threshold observed by the replicator. + pub rewrite_data_files_min_active_data_files: Option, + /// Request timestamp. + pub requested_at: DateTime, +} + +/// Replicator acknowledgement written to the coordination backend. +#[derive(Clone, Debug, Deserialize, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct ExternalMaintenanceReplicatorStatus { + /// Replicator pause state. + pub state: ExternalMaintenanceReplicatorState, + /// Observed run identifier. + pub observed_run_id: Option, + /// Timestamp at which foreground writes became quiesced. + pub quiesced_at: Option>, +} + +/// Replicator pause state. +#[derive(Clone, Copy, Debug, Deserialize, Serialize, PartialEq, Eq)] +pub enum ExternalMaintenanceReplicatorState { + Running, + Pausing, + Quiesced, +} + +impl ExternalMaintenanceReplicatorState { + /// Returns the Kubernetes-compatible status string for this state. + pub fn as_str(self) -> &'static str { + match self { + Self::Running => "Running", + Self::Pausing => "Pausing", + Self::Quiesced => "Quiesced", + } + } +} + +impl From<&str> for ExternalMaintenanceReplicatorState { + fn from(value: &str) -> Self { + match value { + "Pausing" => Self::Pausing, + "Quiesced" => Self::Quiesced, + _ => Self::Running, + } + } +} + +impl ExternalMaintenanceReplicatorStatus { + /// Returns the Kubernetes-compatible status string for this status. + pub fn state_name(&self) -> &'static str { + self.state.as_str() + } + + /// Builds a status from a Kubernetes-compatible state string. + pub fn from_state_name( + state_name: &str, + observed_run_id: Option, + quiesced_at: Option>, + ) -> Self { + Self { state: state_name.into(), observed_run_id, quiesced_at } + } +} + +/// Last successful run history for each operation. +#[derive(Clone, Debug, Default, Deserialize, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct ExternalMaintenanceOperationHistory { + /// Last successful inline flush run. + pub inline_flush: Option, + /// Last successful merge-adjacent-files run. + pub merge_adjacent_files: Option, + /// Last successful rewrite-data-files run. + pub rewrite_data_files: Option, + /// Last successful expire-snapshots run. + pub expire_snapshots: Option, + /// Last successful cleanup-old-files run. + pub cleanup_old_files: Option, +} + +/// Last successful run metadata for one operation. +#[derive(Clone, Debug, Deserialize, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct ExternalMaintenanceOperationRun { + /// Run identifier. + pub run_id: Option, + /// Completion timestamp. + pub completed_at: DateTime, +} + +/// Backend-neutral operation enablement policy. +#[derive(Clone, Copy, Debug, Deserialize, Serialize, PartialEq, Eq)] +#[serde(rename_all = "camelCase")] +pub struct ExternalMaintenanceOperationPolicy { + /// Whether inline flush can be requested. + pub inline_flush_enabled: bool, + /// Whether merge-adjacent-files can be requested. + pub merge_adjacent_files_enabled: bool, + /// Whether rewrite-data-files can be requested. + pub rewrite_data_files_enabled: bool, + /// Whether expire-snapshots can be requested. + pub expire_snapshots_enabled: bool, + /// Whether cleanup-old-files can be requested. + pub cleanup_old_files_enabled: bool, +} + +impl Default for ExternalMaintenanceOperationPolicy { + fn default() -> Self { + Self { + inline_flush_enabled: true, + merge_adjacent_files_enabled: true, + rewrite_data_files_enabled: true, + expire_snapshots_enabled: false, + cleanup_old_files_enabled: true, + } + } +} + +/// Result of asking the coordination backend to create an operation request. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ExternalMaintenanceRequestOutcome { + Created, + AlreadyCovered, + RejectedActiveRun, + MissingState, +} + +/// Runtime coordination backend used by the replicator. +#[async_trait] +pub trait ExternalMaintenanceStore: Clone + Send + Sync + 'static { + /// Loads the current maintenance state. + async fn load_state(&self) -> EtlResult; + + /// Requests maintenance operations. + async fn request_operations( + &self, + request: ExternalMaintenanceOperationRequest, + ) -> EtlResult; + + /// Reports the current replicator maintenance status. + async fn report_replicator_status( + &self, + status: ExternalMaintenanceReplicatorStatus, + ) -> EtlResult<()>; + + /// Clears the current replicator maintenance status. + async fn clear_replicator_status(&self) -> EtlResult<()>; +} + +/// Polling and threshold settings for an external maintenance watcher. +#[derive(Clone, Debug)] +pub struct ExternalMaintenanceWatcherConfig { + /// Poll interval. + pub poll_interval: Duration, + /// Request cooldown after a completed run. + pub request_cooldown: Duration, + /// Coordination store operation timeout. + pub store_timeout: Duration, + /// Inline flush trigger threshold. + pub inline_flush_min_inlined_bytes: u64, + /// Rewrite trigger threshold. + pub rewrite_data_files_min_active_data_files: i64, +} + +impl Default for ExternalMaintenanceWatcherConfig { + fn default() -> Self { + Self { + poll_interval: Duration::from_secs(DEFAULT_POLL_SECONDS), + request_cooldown: Duration::from_secs(DEFAULT_REQUEST_COOLDOWN_SECONDS), + store_timeout: Duration::from_secs(DEFAULT_STORE_TIMEOUT_SECONDS), + inline_flush_min_inlined_bytes: DEFAULT_INLINE_FLUSH_MIN_INLINED_BYTES, + rewrite_data_files_min_active_data_files: + DEFAULT_REWRITE_DATA_FILES_MIN_ACTIVE_DATA_FILES, + } + } +} + +impl ExternalMaintenanceWatcherConfig { + /// Builds watcher config from environment variables. + pub fn from_env() -> Self { + const POLL_SECONDS_ENV: &str = "ETL_DUCKLAKE_MAINTENANCE_POLL_SECONDS"; + const INLINE_FLUSH_MIN_INLINED_BYTES_ENV: &str = + "ETL_DUCKLAKE_EXTERNAL_MAINTENANCE_INLINE_FLUSH_MIN_INLINED_BYTES"; + const REWRITE_DATA_FILES_MIN_ACTIVE_DATA_FILES_ENV: &str = + "ETL_DUCKLAKE_EXTERNAL_MAINTENANCE_REWRITE_DATA_FILES_MIN_ACTIVE_DATA_FILES"; + const REQUEST_COOLDOWN_SECONDS_ENV: &str = + "ETL_DUCKLAKE_EXTERNAL_MAINTENANCE_REQUEST_COOLDOWN_SECONDS"; + const STORE_TIMEOUT_SECONDS_ENV: &str = + "ETL_DUCKLAKE_EXTERNAL_MAINTENANCE_STORE_TIMEOUT_SECONDS"; + + let poll_seconds = env_u64(POLL_SECONDS_ENV) + .filter(|seconds| *seconds > 0) + .unwrap_or(DEFAULT_POLL_SECONDS); + let inline_flush_min_inlined_bytes = env_u64(INLINE_FLUSH_MIN_INLINED_BYTES_ENV) + .unwrap_or(DEFAULT_INLINE_FLUSH_MIN_INLINED_BYTES); + let rewrite_data_files_min_active_data_files = + std::env::var(REWRITE_DATA_FILES_MIN_ACTIVE_DATA_FILES_ENV) + .ok() + .and_then(|value| value.parse::().ok()) + .unwrap_or(DEFAULT_REWRITE_DATA_FILES_MIN_ACTIVE_DATA_FILES); + let request_cooldown = + env_u64(REQUEST_COOLDOWN_SECONDS_ENV).unwrap_or(DEFAULT_REQUEST_COOLDOWN_SECONDS); + let store_timeout = env_u64(STORE_TIMEOUT_SECONDS_ENV) + .filter(|seconds| *seconds > 0) + .unwrap_or(DEFAULT_STORE_TIMEOUT_SECONDS); + + Self { + poll_interval: Duration::from_secs(poll_seconds), + request_cooldown: Duration::from_secs(request_cooldown), + store_timeout: Duration::from_secs(store_timeout), + inline_flush_min_inlined_bytes, + rewrite_data_files_min_active_data_files, + } + } +} + +mod postgres; + +pub use postgres::PostgresExternalMaintenanceStore; + +#[cfg(feature = "ducklake")] +mod kubernetes; + +#[cfg(feature = "ducklake")] +pub use kubernetes::KubernetesExternalMaintenanceStore; + +fn env_u64(name: &str) -> Option { + std::env::var(name).ok().and_then(|value| value.parse::().ok()) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use chrono::TimeDelta; + use tokio::sync::Mutex; + + use super::*; + + #[derive(Clone)] + struct InMemoryExternalMaintenanceStore { + state: Arc>, + } + + impl InMemoryExternalMaintenanceStore { + fn new(state: ExternalMaintenanceState) -> Self { + Self { state: Arc::new(Mutex::new(state)) } + } + } + + #[async_trait] + impl ExternalMaintenanceStore for InMemoryExternalMaintenanceStore { + async fn load_state(&self) -> EtlResult { + Ok(self.state.lock().await.clone()) + } + + async fn request_operations( + &self, + request: ExternalMaintenanceOperationRequest, + ) -> EtlResult { + let mut state = self.state.lock().await; + if !state.exists { + return Ok(ExternalMaintenanceRequestOutcome::MissingState); + } + if state.active_run.is_some() { + return Ok(ExternalMaintenanceRequestOutcome::RejectedActiveRun); + } + if state + .operation_request + .as_ref() + .is_some_and(|existing| existing.operations.covers(request.operations)) + { + return Ok(ExternalMaintenanceRequestOutcome::AlreadyCovered); + } + + let merged_request = if let Some(mut existing) = state.operation_request.take() { + existing.operations = existing.operations.merge(request.operations); + existing.inline_flush_min_inlined_bytes = request + .inline_flush_min_inlined_bytes + .or(existing.inline_flush_min_inlined_bytes); + existing.rewrite_data_files_min_active_data_files = request + .rewrite_data_files_min_active_data_files + .or(existing.rewrite_data_files_min_active_data_files); + existing.requested_at = request.requested_at; + existing + } else { + request + }; + state.operation_request = Some(merged_request); + + Ok(ExternalMaintenanceRequestOutcome::Created) + } + + async fn report_replicator_status( + &self, + status: ExternalMaintenanceReplicatorStatus, + ) -> EtlResult<()> { + self.state.lock().await.replicator = Some(status); + Ok(()) + } + + async fn clear_replicator_status(&self) -> EtlResult<()> { + self.state.lock().await.replicator = None; + Ok(()) + } + } + + fn operation_request( + operations: ExternalMaintenanceOperations, + ) -> ExternalMaintenanceOperationRequest { + ExternalMaintenanceOperationRequest { + operations, + inline_flush_min_inlined_bytes: Some(10_000_000), + rewrite_data_files_min_active_data_files: Some(40), + requested_at: Utc::now(), + } + } + + #[test] + fn operation_coverage_requires_every_requested_operation() { + let existing = ExternalMaintenanceOperations { + inline_flush: true, + rewrite_data_files: true, + ..ExternalMaintenanceOperations::default() + }; + + assert!(existing.covers(ExternalMaintenanceOperations { + inline_flush: true, + ..ExternalMaintenanceOperations::default() + })); + assert!(!existing.covers(ExternalMaintenanceOperations { + inline_flush: true, + expire_snapshots: true, + ..ExternalMaintenanceOperations::default() + })); + } + + #[test] + fn operation_merge_is_idempotent_union() { + let existing = ExternalMaintenanceOperations { + inline_flush: true, + ..ExternalMaintenanceOperations::default() + }; + let requested = ExternalMaintenanceOperations { + rewrite_data_files: true, + cleanup_old_files: true, + ..ExternalMaintenanceOperations::default() + }; + + let merged = existing.merge(requested); + + assert!(merged.inline_flush); + assert!(merged.rewrite_data_files); + assert!(merged.cleanup_old_files); + assert!(!merged.expire_snapshots); + } + + #[test] + fn replicator_state_round_trips_kubernetes_names() { + assert_eq!(ExternalMaintenanceReplicatorState::Running.as_str(), "Running"); + assert_eq!( + ExternalMaintenanceReplicatorState::from("Pausing"), + ExternalMaintenanceReplicatorState::Pausing + ); + assert_eq!( + ExternalMaintenanceReplicatorState::from("Quiesced"), + ExternalMaintenanceReplicatorState::Quiesced + ); + assert_eq!( + ExternalMaintenanceReplicatorState::from("unknown"), + ExternalMaintenanceReplicatorState::Running + ); + + let status = ExternalMaintenanceReplicatorStatus { + state: ExternalMaintenanceReplicatorState::Quiesced, + observed_run_id: None, + quiesced_at: None, + }; + assert_eq!(status.state_name(), "Quiesced"); + + let status = ExternalMaintenanceReplicatorStatus::from_state_name("Pausing", None, None); + assert_eq!(status.state, ExternalMaintenanceReplicatorState::Pausing); + } + + #[test] + fn completed_run_cooldown_uses_last_completed_at() { + let state = ExternalMaintenanceState { + last_completed_at: Some(Utc::now() - TimeDelta::seconds(30)), + ..ExternalMaintenanceState::default() + }; + + assert!(completed_run_in_cooldown_for_tests(&state, Duration::from_secs(60))); + assert!(!completed_run_in_cooldown_for_tests(&state, Duration::from_secs(10))); + } + + fn completed_run_in_cooldown_for_tests( + state: &ExternalMaintenanceState, + cooldown: Duration, + ) -> bool { + let Some(completed_at) = state.last_completed_at else { + return false; + }; + + Utc::now() + .signed_duration_since(completed_at) + .to_std() + .is_ok_and(|elapsed| elapsed < cooldown) + } + + #[tokio::test] + async fn mock_store_creates_and_merges_operation_requests() { + let store = InMemoryExternalMaintenanceStore::new(ExternalMaintenanceState::present()); + + let inline_flush = ExternalMaintenanceOperations { + inline_flush: true, + ..ExternalMaintenanceOperations::default() + }; + let outcome = store.request_operations(operation_request(inline_flush)).await.unwrap(); + assert_eq!(outcome, ExternalMaintenanceRequestOutcome::Created); + + let outcome = store.request_operations(operation_request(inline_flush)).await.unwrap(); + assert_eq!(outcome, ExternalMaintenanceRequestOutcome::AlreadyCovered); + + let rewrite = ExternalMaintenanceOperations { + rewrite_data_files: true, + cleanup_old_files: true, + ..ExternalMaintenanceOperations::default() + }; + let outcome = store.request_operations(operation_request(rewrite)).await.unwrap(); + assert_eq!(outcome, ExternalMaintenanceRequestOutcome::Created); + + let state = store.load_state().await.unwrap(); + let operations = state.operation_request.unwrap().operations; + assert!(operations.inline_flush); + assert!(operations.rewrite_data_files); + assert!(operations.cleanup_old_files); + assert!(!operations.expire_snapshots); + } + + #[tokio::test] + async fn mock_store_rejects_requests_when_missing_or_running() { + let missing_store = + InMemoryExternalMaintenanceStore::new(ExternalMaintenanceState::default()); + let outcome = missing_store + .request_operations(operation_request(ExternalMaintenanceOperations { + inline_flush: true, + ..ExternalMaintenanceOperations::default() + })) + .await + .unwrap(); + assert_eq!(outcome, ExternalMaintenanceRequestOutcome::MissingState); + + let active_store = InMemoryExternalMaintenanceStore::new(ExternalMaintenanceState { + exists: true, + active_run: Some(ExternalMaintenanceRun { + run_id: "run-1".to_owned(), + started_at: Some(Utc::now()), + }), + ..ExternalMaintenanceState::default() + }); + let outcome = active_store + .request_operations(operation_request(ExternalMaintenanceOperations { + rewrite_data_files: true, + ..ExternalMaintenanceOperations::default() + })) + .await + .unwrap(); + assert_eq!(outcome, ExternalMaintenanceRequestOutcome::RejectedActiveRun); + } + + #[tokio::test] + async fn mock_store_round_trips_replicator_status() { + let store = InMemoryExternalMaintenanceStore::new(ExternalMaintenanceState::present()); + let quiesced_at = Utc::now(); + + store + .report_replicator_status(ExternalMaintenanceReplicatorStatus { + state: ExternalMaintenanceReplicatorState::Quiesced, + observed_run_id: Some("run-1".to_owned()), + quiesced_at: Some(quiesced_at), + }) + .await + .unwrap(); + + let status = store.load_state().await.unwrap().replicator.unwrap(); + assert_eq!(status.state, ExternalMaintenanceReplicatorState::Quiesced); + assert_eq!(status.observed_run_id.as_deref(), Some("run-1")); + assert_eq!(status.quiesced_at, Some(quiesced_at)); + + store.clear_replicator_status().await.unwrap(); + assert!(store.load_state().await.unwrap().replicator.is_none()); + } +} diff --git a/crates/etl-maintenance/src/coordination/kubernetes.rs b/crates/etl-maintenance/src/coordination/kubernetes.rs new file mode 100644 index 000000000..e637476d7 --- /dev/null +++ b/crates/etl-maintenance/src/coordination/kubernetes.rs @@ -0,0 +1,365 @@ +use std::time::Duration; + +use chrono::{DateTime, Utc}; +use etl::{ + error::{ErrorKind, EtlResult}, + etl_error, +}; +use kube::{ + Api, Client, + api::{Patch, PatchParams}, + core::{ApiResource, DynamicObject, GroupVersionKind}, +}; +use serde_json::{Value, json}; +use tokio::time; +use tracing::info; + +use super::{ + ExternalMaintenanceOperationHistory, ExternalMaintenanceOperationPolicy, + ExternalMaintenanceOperationRequest, ExternalMaintenanceOperationRun, + ExternalMaintenanceOperations, ExternalMaintenancePause, ExternalMaintenanceReplicatorStatus, + ExternalMaintenanceRequestOutcome, ExternalMaintenanceRun, ExternalMaintenanceState, + ExternalMaintenanceStore, +}; + +const CR_NAME_ENV: &str = "ETL_DUCKLAKE_MAINTENANCE_CR_NAME"; +const CR_NAMESPACE_ENV: &str = "ETL_DUCKLAKE_MAINTENANCE_CR_NAMESPACE"; +const KUBERNETES_API_TIMEOUT_SECONDS_ENV: &str = + "ETL_DUCKLAKE_EXTERNAL_MAINTENANCE_KUBERNETES_API_TIMEOUT_SECONDS"; + +/// Kubernetes CR-backed external maintenance store. +#[derive(Clone)] +pub struct KubernetesExternalMaintenanceStore { + api: Api, + name: String, + timeout: Duration, +} + +impl KubernetesExternalMaintenanceStore { + /// Creates a Kubernetes store from environment variables. + pub async fn from_env(default_timeout: Duration) -> EtlResult> { + let Some(name) = std::env::var(CR_NAME_ENV).ok().filter(|value| !value.is_empty()) else { + return Ok(None); + }; + let Some(namespace) = + std::env::var(CR_NAMESPACE_ENV).ok().filter(|value| !value.is_empty()) + else { + return Ok(None); + }; + let timeout = std::env::var(KUBERNETES_API_TIMEOUT_SECONDS_ENV) + .ok() + .and_then(|value| value.parse::().ok()) + .filter(|seconds| *seconds > 0) + .map_or(default_timeout, Duration::from_secs); + + let client = Client::try_default().await.map_err(|error| { + etl_error!( + ErrorKind::ConfigError, + "Failed to build Kubernetes client for external maintenance", + source: error + ) + })?; + let api: Api = + Api::namespaced_with(client, &namespace, &ducklake_maintenance_api_resource()); + + info!( + ducklake_maintenance = %name, + namespace, + "ducklake Kubernetes external maintenance store configured: ducklake_maintenance={}, namespace={}", + name, + namespace + ); + + Ok(Some(Self { api, name, timeout })) + } + + fn kube_error(error: kube::Error, message: &'static str) -> etl::error::EtlError { + etl_error!(ErrorKind::SourceQueryFailed, message, source: error) + } +} + +#[async_trait::async_trait] +impl ExternalMaintenanceStore for KubernetesExternalMaintenanceStore { + async fn load_state(&self) -> EtlResult { + match time::timeout(self.timeout, self.api.get(&self.name)).await { + Ok(Ok(resource)) => Ok(state_from_ducklake_maintenance_cr(&resource)), + Ok(Err(kube::Error::Api(error))) if error.code == 404 => { + Ok(ExternalMaintenanceState::default()) + } + Ok(Err(error)) => { + Err(Self::kube_error(error, "Failed to load DuckLake maintenance CR")) + } + Err(_) => Err(etl_error!( + ErrorKind::SourceQueryFailed, + "Timed out loading DuckLake maintenance CR" + )), + } + } + + async fn request_operations( + &self, + request: ExternalMaintenanceOperationRequest, + ) -> EtlResult { + let state = self.load_state().await?; + if !state.exists { + return Ok(ExternalMaintenanceRequestOutcome::MissingState); + } + if state.active_run.is_some() { + return Ok(ExternalMaintenanceRequestOutcome::RejectedActiveRun); + } + if state + .operation_request + .as_ref() + .is_some_and(|existing| existing.operations.covers(request.operations)) + { + return Ok(ExternalMaintenanceRequestOutcome::AlreadyCovered); + } + + let patch = json!({ + "status": { + "operationRequests": { + "inlineFlush": request.operations.inline_flush, + "mergeAdjacentFiles": request.operations.merge_adjacent_files, + "rewriteDataFiles": request.operations.rewrite_data_files, + "expireSnapshots": request.operations.expire_snapshots, + "cleanupOldFiles": request.operations.cleanup_old_files, + "inlineFlushMinInlinedBytes": request.inline_flush_min_inlined_bytes, + "rewriteDataFilesMinActiveDataFiles": + request.rewrite_data_files_min_active_data_files, + "requestedAt": request.requested_at.to_rfc3339(), + } + } + }); + patch_status(&self.api, &self.name, &patch, self.timeout).await?; + Ok(ExternalMaintenanceRequestOutcome::Created) + } + + async fn report_replicator_status( + &self, + status: ExternalMaintenanceReplicatorStatus, + ) -> EtlResult<()> { + let patch = json!({ + "status": { + "replicator": { + "state": status.state_name(), + "observedRunId": status.observed_run_id, + "quiescedAt": status.quiesced_at.map(|time| time.to_rfc3339()), + } + } + }); + patch_status(&self.api, &self.name, &patch, self.timeout).await + } + + async fn clear_replicator_status(&self) -> EtlResult<()> { + let patch = json!({ + "status": { + "replicator": { + "state": "Running", + "observedRunId": null, + "quiescedAt": null, + } + } + }); + patch_status(&self.api, &self.name, &patch, self.timeout).await + } +} + +async fn patch_status( + api: &Api, + name: &str, + patch: &Value, + timeout: Duration, +) -> EtlResult<()> { + let params = PatchParams::default(); + match time::timeout(timeout, api.patch_status(name, ¶ms, &Patch::Merge(patch))).await { + Ok(Ok(_)) => Ok(()), + Ok(Err(error)) => Err(KubernetesExternalMaintenanceStore::kube_error( + error, + "Failed to patch DuckLake maintenance CR status", + )), + Err(_) => Err(etl_error!( + ErrorKind::SourceQueryFailed, + "Timed out patching DuckLake maintenance CR status" + )), + } +} + +fn state_from_ducklake_maintenance_cr(resource: &DynamicObject) -> ExternalMaintenanceState { + let mut state = ExternalMaintenanceState::present(); + state.pause_request = active_pause_request(resource); + state.active_run = active_run(resource); + state.operation_request = operation_request(resource); + state.replicator = replicator_status(resource); + state.last_successful_operations = operation_history(resource); + state.last_completed_at = last_completed_at(resource); + state.operation_policy = operation_policy(resource); + state +} + +fn active_pause_request(resource: &DynamicObject) -> Option { + let pause = resource.data.get("status")?.get("pauseRequest")?; + if pause.is_null() { + return None; + } + + Some(ExternalMaintenancePause { + run_id: pause.get("runId")?.as_str()?.to_owned(), + requested_at: parse_rfc3339_value(pause.get("requestedAt")), + expires_at: parse_rfc3339_value(pause.get("expiresAt"))?, + }) +} + +fn active_run(resource: &DynamicObject) -> Option { + let active_run = resource.data.get("status")?.get("activeRun")?; + if active_run.is_null() { + return None; + } + + Some(ExternalMaintenanceRun { + run_id: active_run.get("runId")?.as_str()?.to_owned(), + started_at: parse_rfc3339_value(active_run.get("startedAt")), + }) +} + +fn operation_request(resource: &DynamicObject) -> Option { + let requests = resource.data.get("status")?.get("operationRequests")?; + if requests.is_null() { + return None; + } + + Some(ExternalMaintenanceOperationRequest { + operations: ExternalMaintenanceOperations { + inline_flush: bool_value(requests, "inlineFlush", false), + merge_adjacent_files: bool_value(requests, "mergeAdjacentFiles", false), + rewrite_data_files: bool_value(requests, "rewriteDataFiles", false), + expire_snapshots: bool_value(requests, "expireSnapshots", false), + cleanup_old_files: bool_value(requests, "cleanupOldFiles", false), + }, + inline_flush_min_inlined_bytes: requests + .get("inlineFlushMinInlinedBytes") + .and_then(Value::as_u64), + rewrite_data_files_min_active_data_files: requests + .get("rewriteDataFilesMinActiveDataFiles") + .and_then(Value::as_i64), + requested_at: parse_rfc3339_value(requests.get("requestedAt"))?, + }) +} + +fn replicator_status(resource: &DynamicObject) -> Option { + let replicator = resource.data.get("status")?.get("replicator")?; + if replicator.is_null() { + return None; + } + + Some(ExternalMaintenanceReplicatorStatus::from_state_name( + replicator.get("state")?.as_str()?, + replicator.get("observedRunId").and_then(Value::as_str).map(ToOwned::to_owned), + parse_rfc3339_value(replicator.get("quiescedAt")), + )) +} + +fn operation_history(resource: &DynamicObject) -> ExternalMaintenanceOperationHistory { + let Some(history) = + resource.data.get("status").and_then(|status| status.get("lastSuccessfulOperationRuns")) + else { + return ExternalMaintenanceOperationHistory::default(); + }; + + ExternalMaintenanceOperationHistory { + inline_flush: operation_run(history.get("inlineFlush")), + merge_adjacent_files: operation_run(history.get("mergeAdjacentFiles")), + rewrite_data_files: operation_run(history.get("rewriteDataFiles")), + expire_snapshots: operation_run(history.get("expireSnapshots")), + cleanup_old_files: operation_run(history.get("cleanupOldFiles")), + } +} + +fn operation_run(value: Option<&Value>) -> Option { + let value = value?; + if value.is_null() { + return None; + } + + Some(ExternalMaintenanceOperationRun { + run_id: value.get("runId").and_then(Value::as_str).map(ToOwned::to_owned), + completed_at: parse_rfc3339_value(value.get("completedAt"))?, + }) +} + +fn last_completed_at(resource: &DynamicObject) -> Option> { + parse_rfc3339_value( + resource + .data + .get("status") + .and_then(|status| status.get("lastCompletedRun")) + .and_then(|run| run.get("completedAt")), + ) +} + +fn operation_policy(resource: &DynamicObject) -> ExternalMaintenanceOperationPolicy { + let operations = resource.data.get("spec").and_then(|spec| spec.get("operations")); + let inline_flush = operations.and_then(|ops| ops.get("inlineFlush")); + let merge_adjacent_files = operations.and_then(|ops| ops.get("mergeAdjacentFiles")); + let rewrite_data_files = operations.and_then(|ops| ops.get("rewriteDataFiles")); + let expire_snapshots = operations.and_then(|ops| ops.get("expireSnapshots")); + let cleanup_old_files = operations.and_then(|ops| ops.get("cleanupOldFiles")); + + ExternalMaintenanceOperationPolicy { + inline_flush_enabled: enabled_value(inline_flush, true), + merge_adjacent_files_enabled: enabled_value(merge_adjacent_files, true), + rewrite_data_files_enabled: enabled_value(rewrite_data_files, true), + expire_snapshots_enabled: enabled_value(expire_snapshots, false), + cleanup_old_files_enabled: enabled_value(cleanup_old_files, true), + } +} + +fn enabled_value(value: Option<&Value>, default: bool) -> bool { + value.and_then(|value| value.get("enabled")).and_then(Value::as_bool).unwrap_or(default) +} + +fn bool_value(value: &Value, key: &str, default: bool) -> bool { + value.get(key).and_then(Value::as_bool).unwrap_or(default) +} + +fn parse_rfc3339_value(value: Option<&Value>) -> Option> { + value + .and_then(Value::as_str) + .and_then(|value| DateTime::parse_from_rfc3339(value).ok()) + .map(|time| time.with_timezone(&Utc)) +} + +fn ducklake_maintenance_api_resource() -> ApiResource { + let gvk = GroupVersionKind::gvk("etl.supabase.com", "v1alpha1", "DuckLakeMaintenance"); + ApiResource::from_gvk(&gvk) +} + +#[cfg(test)] +mod tests { + use serde_json::json; + + use super::*; + + #[test] + fn operation_policy_defaults_expire_snapshots_disabled() { + let resource: DynamicObject = serde_json::from_value(json!({ + "apiVersion": "etl.supabase.com/v1alpha1", + "kind": "DuckLakeMaintenance", + "metadata": { + "name": "pipeline-maintenance" + }, + "spec": { + "operations": { + "inlineFlush": {}, + "rewriteDataFiles": {} + } + } + })) + .unwrap(); + + let policy = operation_policy(&resource); + + assert!(policy.inline_flush_enabled); + assert!(policy.rewrite_data_files_enabled); + assert!(!policy.expire_snapshots_enabled); + } +} diff --git a/crates/etl-maintenance/src/coordination/postgres.rs b/crates/etl-maintenance/src/coordination/postgres.rs new file mode 100644 index 000000000..2503e16b2 --- /dev/null +++ b/crates/etl-maintenance/src/coordination/postgres.rs @@ -0,0 +1,312 @@ +use async_trait::async_trait; +use etl::{ + error::{ErrorKind, EtlError, EtlResult}, + etl_error, +}; +use sqlx::{Executor, PgPool, Row, migrate::Migrator, types::Json}; + +use super::{ + ExternalMaintenanceOperationHistory, ExternalMaintenanceOperationPolicy, + ExternalMaintenanceOperationRequest, ExternalMaintenancePause, + ExternalMaintenanceReplicatorStatus, ExternalMaintenanceRequestOutcome, ExternalMaintenanceRun, + ExternalMaintenanceState, ExternalMaintenanceStore, +}; + +const CREATE_MIGRATION_SCHEMA_SQL: &str = "create schema if not exists etl;"; +const SET_MIGRATION_SEARCH_PATH_SQL: &str = "set search_path = etl, public;"; + +fn sqlx_error(error: sqlx::Error, message: &'static str) -> EtlError { + etl_error!(ErrorKind::SourceQueryFailed, message, source: error) +} + +fn migration_error(error: sqlx::migrate::MigrateError, message: &'static str) -> EtlError { + etl_error!(ErrorKind::SourceQueryFailed, message, source: error) +} + +fn postgres_migrator() -> Migrator { + let mut migrator = sqlx::migrate!("./migrations/postgres"); + migrator.set_ignore_missing(true); + migrator +} + +/// Postgres-backed external maintenance coordination store. +#[derive(Clone)] +pub struct PostgresExternalMaintenanceStore { + pipeline_id: i64, + pool: PgPool, +} + +impl PostgresExternalMaintenanceStore { + /// Creates a Postgres-backed maintenance store. + pub fn new(pipeline_id: i64, pool: PgPool) -> Self { + Self { pipeline_id, pool } + } + + /// Runs the Postgres maintenance migrations. + pub async fn ensure_schema(&self) -> EtlResult<()> { + let mut connection = self.pool.acquire().await.map_err(|error| { + sqlx_error(error, "Failed to acquire external maintenance migration connection") + })?; + connection.execute(CREATE_MIGRATION_SCHEMA_SQL).await.map_err(|error| { + sqlx_error(error, "Failed to create external maintenance migration schema") + })?; + connection.execute(SET_MIGRATION_SEARCH_PATH_SQL).await.map_err(|error| { + sqlx_error(error, "Failed to configure external maintenance migration search path") + })?; + + postgres_migrator().run_direct(None, &mut *connection).await.map_err(|error| { + migration_error(error, "Failed to run external maintenance migrations") + })?; + + Ok(()) + } + + /// Ensures state exists for one pipeline. + pub async fn ensure_pipeline_state( + &self, + policy: ExternalMaintenanceOperationPolicy, + ) -> EtlResult<()> { + let policy = Json(policy); + sqlx::query( + r#" + insert into etl.external_maintenance_state (pipeline_id, operation_policy) + values ($1, $2) + on conflict (pipeline_id) + do update set operation_policy = excluded.operation_policy, updated_at = now() + "#, + ) + .bind(self.pipeline_id) + .bind(policy) + .execute(&self.pool) + .await + .map_err(|error| sqlx_error(error, "Failed to upsert external maintenance state"))?; + + Ok(()) + } + + /// Ensures state exists for one pipeline without replacing an existing + /// operation policy. + pub async fn ensure_pipeline_state_if_missing( + &self, + policy: ExternalMaintenanceOperationPolicy, + ) -> EtlResult<()> { + let policy = Json(policy); + sqlx::query( + r#" + insert into etl.external_maintenance_state (pipeline_id, operation_policy) + values ($1, $2) + on conflict (pipeline_id) do nothing + "#, + ) + .bind(self.pipeline_id) + .bind(policy) + .execute(&self.pool) + .await + .map_err(|error| sqlx_error(error, "Failed to insert external maintenance state"))?; + + Ok(()) + } + + /// Deletes state for one pipeline. + pub async fn delete_pipeline_state(&self) -> EtlResult<()> { + sqlx::query("delete from etl.external_maintenance_state where pipeline_id = $1") + .bind(self.pipeline_id) + .execute(&self.pool) + .await + .map_err(|error| sqlx_error(error, "Failed to delete external maintenance state"))?; + + Ok(()) + } +} + +#[async_trait] +impl ExternalMaintenanceStore for PostgresExternalMaintenanceStore { + async fn load_state(&self) -> EtlResult { + let Some(row) = sqlx::query( + r#" + select + active_run, + pause_request, + operation_request, + replicator, + last_successful_operations, + last_completed_at, + operation_policy + from etl.external_maintenance_state + where pipeline_id = $1 + "#, + ) + .bind(self.pipeline_id) + .fetch_optional(&self.pool) + .await + .map_err(|error| sqlx_error(error, "Failed to load external maintenance state"))? + else { + return Ok(ExternalMaintenanceState::default()); + }; + + Ok(ExternalMaintenanceState { + exists: true, + active_run: row + .try_get::>, _>("active_run") + .map_err(|error| sqlx_error(error, "Failed to decode external maintenance run"))? + .map(|value| value.0), + pause_request: row + .try_get::>, _>("pause_request") + .map_err(|error| sqlx_error(error, "Failed to decode external maintenance pause"))? + .map(|value| value.0), + operation_request: row + .try_get::>, _>( + "operation_request", + ) + .map_err(|error| { + sqlx_error(error, "Failed to decode external maintenance operation request") + })? + .map(|value| value.0), + replicator: row + .try_get::>, _>("replicator") + .map_err(|error| { + sqlx_error(error, "Failed to decode external maintenance replicator status") + })? + .map(|value| value.0), + last_successful_operations: row + .try_get::, _>( + "last_successful_operations", + ) + .map_err(|error| { + sqlx_error(error, "Failed to decode external maintenance operation history") + })? + .0, + last_completed_at: row.try_get("last_completed_at").map_err(|error| { + sqlx_error(error, "Failed to decode external maintenance completed timestamp") + })?, + operation_policy: row + .try_get::, _>("operation_policy") + .map_err(|error| { + sqlx_error(error, "Failed to decode external maintenance operation policy") + })? + .0, + }) + } + + async fn request_operations( + &self, + request: ExternalMaintenanceOperationRequest, + ) -> EtlResult { + let mut tx = + self.pool.begin().await.map_err(|error| { + sqlx_error(error, "Failed to begin external maintenance request") + })?; + + let Some(row) = sqlx::query( + r#" + select active_run, operation_request + from etl.external_maintenance_state + where pipeline_id = $1 + for update + "#, + ) + .bind(self.pipeline_id) + .fetch_optional(&mut *tx) + .await + .map_err(|error| sqlx_error(error, "Failed to lock external maintenance state"))? + else { + tx.rollback() + .await + .map_err(|error| sqlx_error(error, "Failed to roll back missing state request"))?; + return Ok(ExternalMaintenanceRequestOutcome::MissingState); + }; + + let active_run = row + .try_get::>, _>("active_run") + .map_err(|error| sqlx_error(error, "Failed to decode active maintenance run"))?; + if active_run.is_some() { + tx.rollback() + .await + .map_err(|error| sqlx_error(error, "Failed to roll back active run request"))?; + return Ok(ExternalMaintenanceRequestOutcome::RejectedActiveRun); + } + + let existing_request = row + .try_get::>, _>("operation_request") + .map_err(|error| sqlx_error(error, "Failed to decode operation request"))? + .map(|value| value.0); + if existing_request + .as_ref() + .is_some_and(|existing| existing.operations.covers(request.operations)) + { + tx.rollback() + .await + .map_err(|error| sqlx_error(error, "Failed to roll back covered request"))?; + return Ok(ExternalMaintenanceRequestOutcome::AlreadyCovered); + } + + let merged_request = if let Some(mut existing_request) = existing_request { + existing_request.operations = existing_request.operations.merge(request.operations); + existing_request.inline_flush_min_inlined_bytes = request + .inline_flush_min_inlined_bytes + .or(existing_request.inline_flush_min_inlined_bytes); + existing_request.rewrite_data_files_min_active_data_files = request + .rewrite_data_files_min_active_data_files + .or(existing_request.rewrite_data_files_min_active_data_files); + existing_request.requested_at = request.requested_at; + existing_request + } else { + request + }; + + sqlx::query( + r#" + update etl.external_maintenance_state + set operation_request = $2, updated_at = now() + where pipeline_id = $1 + "#, + ) + .bind(self.pipeline_id) + .bind(Json(merged_request)) + .execute(&mut *tx) + .await + .map_err(|error| sqlx_error(error, "Failed to update external maintenance request"))?; + + tx.commit() + .await + .map_err(|error| sqlx_error(error, "Failed to commit external maintenance request"))?; + + Ok(ExternalMaintenanceRequestOutcome::Created) + } + + async fn report_replicator_status( + &self, + status: ExternalMaintenanceReplicatorStatus, + ) -> EtlResult<()> { + sqlx::query( + r#" + update etl.external_maintenance_state + set replicator = $2, updated_at = now() + where pipeline_id = $1 + "#, + ) + .bind(self.pipeline_id) + .bind(Json(status)) + .execute(&self.pool) + .await + .map_err(|error| sqlx_error(error, "Failed to report external maintenance status"))?; + + Ok(()) + } + + async fn clear_replicator_status(&self) -> EtlResult<()> { + sqlx::query( + r#" + update etl.external_maintenance_state + set replicator = null, updated_at = now() + where pipeline_id = $1 + "#, + ) + .bind(self.pipeline_id) + .execute(&self.pool) + .await + .map_err(|error| sqlx_error(error, "Failed to clear external maintenance status"))?; + + Ok(()) + } +} diff --git a/crates/etl-maintenance/src/ducklake/mod.rs b/crates/etl-maintenance/src/ducklake/mod.rs new file mode 100644 index 000000000..6f6dacabb --- /dev/null +++ b/crates/etl-maintenance/src/ducklake/mod.rs @@ -0,0 +1,10 @@ +//! DuckLake maintenance runner. + +mod runner; + +pub use runner::{ + CleanupOldFilesMaintenanceConfig, DuckLakeMaintenanceConfig, DuckLakeMaintenanceOutcome, + ExpireSnapshotsMaintenanceConfig, InlineFlushMaintenanceConfig, + MergeAdjacentFilesMaintenanceConfig, RewriteDataFilesMaintenanceConfig, S3Config, + flush_table_inlined_data, run_maintenance_once, +}; diff --git a/crates/etl-maintenance/src/ducklake/runner.rs b/crates/etl-maintenance/src/ducklake/runner.rs new file mode 100644 index 000000000..5e00c086d --- /dev/null +++ b/crates/etl-maintenance/src/ducklake/runner.rs @@ -0,0 +1,2544 @@ +//! One-shot DuckLake maintenance execution. + +use std::{ + collections::{BTreeMap, BTreeSet}, + env, error, fmt, + path::{Path, PathBuf}, + process, + str::FromStr, + sync::{ + Arc, LazyLock, + atomic::{AtomicBool, Ordering}, + }, + time::Duration, +}; + +use etl::{ + error::{ErrorKind, EtlError, EtlResult}, + etl_error, +}; +use metrics::{gauge, histogram}; +use pg_escape::{quote_identifier, quote_literal}; +use regex::Regex; +use sqlx::{AssertSqlSafe, PgPool, postgres::PgPoolOptions}; +use tokio::{ + sync::{Semaphore, oneshot}, + task::JoinHandle, + time::Instant, +}; +use tokio_postgres::{ + Config as PgConfig, + config::{Host, SslMode}, +}; +use tracing::{debug, info, trace, warn}; +use url::Url; + +const LAKE_CATALOG: &str = "lake"; +const ATTACH_DATA_INLINING_ROW_LIMIT: u64 = 10_000; +const DUCKDB_EXTENSION_ROOT_ENV_VAR: &str = "ETL_DUCKDB_EXTENSION_ROOT"; +const CONTAINER_DUCKDB_EXTENSION_ROOT: &str = "/app/duckdb_extensions"; +const DUCKDB_EXTENSION_VERSION: &str = "1.5.2"; +const DUCKLAKE_EXTENSION_FILE: &str = "ducklake.duckdb_extension"; +const HTTPFS_EXTENSION_FILE: &str = "httpfs.duckdb_extension"; +const POSTGRES_SCANNER_EXTENSION_FILE: &str = "postgres_scanner.duckdb_extension"; +const TARGET_FILE_SIZE_OPTION_NAME: &str = "target_file_size"; +const MAINTENANCE_TARGET_FILE_SIZE: &str = "10MB"; +const PARQUET_COMPRESSION_OPTION_NAME: &str = "parquet_compression"; +const PARQUET_COMPRESSION_OPTION_VALUE: &str = "zstd"; +const PARQUET_ROW_GROUP_SIZE_BYTES_OPTION_NAME: &str = "parquet_row_group_size_bytes"; +const PARQUET_ROW_GROUP_SIZE_BYTES_OPTION_VALUE: &str = "10MB"; +const PARQUET_VERSION_OPTION_NAME: &str = "parquet_version"; +const PARQUET_VERSION_OPTION_VALUE: u8 = 2; +const PRESERVE_INSERTION_ORDER_OPTION_NAME: &str = "preserve_insertion_order"; +const MAINTENANCE_QUERY_TIMEOUT: Duration = Duration::from_secs(3 * 60); +const BLOCKING_ABORT_GRACE: Duration = Duration::from_secs(30); +const DUCKDB_MAINTENANCE_OPERATION_KIND: &str = "maintenance"; +const ETL_DUCKLAKE_INLINE_FLUSH_ROWS: &str = "etl_ducklake_inline_flush_rows"; +const ETL_DUCKLAKE_INLINE_FLUSH_DURATION_SECONDS: &str = + "etl_ducklake_inline_flush_duration_seconds"; +const ETL_DUCKLAKE_TABLE_ACTIVE_INLINED_DATA_BYTES: &str = + "etl_ducklake_table_active_inlined_data_bytes"; +const RESULT_LABEL: &str = "result"; +const TABLE_LABEL: &str = "table"; +const SMALL_FILE_SIZE_BYTES: i64 = 5 * 1024 * 1024; + +static POSTGRES_PASSWORD_REGEX: LazyLock = LazyLock::new(|| { + Regex::new(r"password=(?:'([^'\\]|\\.)*'|[^\s,);]+)") + .expect("postgres password redaction regex should compile") +}); + +/// S3-compatible storage credentials for DuckDB's httpfs extension. +#[derive(Debug, Clone)] +pub struct S3Config { + /// S3 access key id. + pub access_key_id: String, + /// S3 secret access key. + pub secret_access_key: String, + /// AWS region or equivalent. + pub region: String, + /// Optional S3-compatible endpoint. + pub endpoint: Option, + /// S3 URL style. + pub url_style: String, + /// Whether to use HTTPS. + pub use_ssl: bool, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +struct DuckLakeSetupStep { + label: &'static str, + sql: String, +} + +#[derive(Debug, Clone, Default, PartialEq, Eq)] +struct DuckLakeSetupPlan { + steps: Vec, +} + +impl DuckLakeSetupPlan { + fn steps(&self) -> &[DuckLakeSetupStep] { + &self.steps + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum DuckDbExtensionStrategy { + VendoredLocal { platform_dir: &'static str }, + InstallFromRepository, +} + +impl DuckDbExtensionStrategy { + fn disables_autoload(self) -> bool { + matches!(self, Self::VendoredLocal { .. }) + } +} + +fn configure_writer_session_sql() -> String { + format!("SET {PRESERVE_INSERTION_ORDER_OPTION_NAME} = false;") +} + +fn configure_parquet_settings_sql() -> String { + format!( + "CALL {LAKE_CATALOG}.set_option({}, {}); CALL {LAKE_CATALOG}.set_option({}, {}); CALL \ + {LAKE_CATALOG}.set_option({}, {});", + quote_literal(PARQUET_COMPRESSION_OPTION_NAME), + quote_literal(PARQUET_COMPRESSION_OPTION_VALUE), + quote_literal(PARQUET_ROW_GROUP_SIZE_BYTES_OPTION_NAME), + quote_literal(PARQUET_ROW_GROUP_SIZE_BYTES_OPTION_VALUE), + quote_literal(PARQUET_VERSION_OPTION_NAME), + PARQUET_VERSION_OPTION_VALUE, + ) +} + +fn resolve_maintenance_target_file_size(maintenance_target_file_size: Option<&str>) -> &str { + maintenance_target_file_size.unwrap_or(MAINTENANCE_TARGET_FILE_SIZE) +} + +fn maintenance_target_file_size_sql(maintenance_target_file_size: Option<&str>) -> String { + format!( + "CALL ducklake_set_option({}, {}, {});", + quote_literal(LAKE_CATALOG), + quote_literal(TARGET_FILE_SIZE_OPTION_NAME), + quote_literal(resolve_maintenance_target_file_size(maintenance_target_file_size,)), + ) +} + +fn current_duckdb_extension_strategy() -> EtlResult { + duckdb_extension_strategy( + env::consts::OS, + env::consts::ARCH, + env::var_os(DUCKDB_EXTENSION_ROOT_ENV_VAR).map(PathBuf::from), + Path::new(CONTAINER_DUCKDB_EXTENSION_ROOT), + &repo_vendored_extension_root(), + ) +} + +fn duckdb_extension_strategy( + os: &str, + arch: &str, + env_override: Option, + container_root: &Path, + repo_root: &Path, +) -> EtlResult { + match os { + "linux" => { + let platform_dir = match arch { + "x86_64" | "amd64" => "linux_amd64", + "aarch64" | "arm64" => "linux_arm64", + _ => { + return Err(etl_error!( + ErrorKind::ConfigError, + "Unsupported DuckDB extension platform", + format!( + "linux architecture `{arch}` is not supported for vendored DuckDB \ + extensions" + ) + )); + } + }; + + if vendored_extension_dir(platform_dir, env_override, container_root, repo_root)? + .is_some() + { + Ok(DuckDbExtensionStrategy::VendoredLocal { platform_dir }) + } else { + Ok(DuckDbExtensionStrategy::InstallFromRepository) + } + } + "macos" => { + let platform_dir = match arch { + "x86_64" | "amd64" => "osx_amd64", + "aarch64" | "arm64" => "osx_arm64", + _ => { + return Err(etl_error!( + ErrorKind::ConfigError, + "Unsupported DuckDB extension platform", + format!( + "macos architecture `{arch}` is not supported for vendored DuckDB \ + extensions" + ) + )); + } + }; + + if vendored_extension_dir(platform_dir, env_override, container_root, repo_root)? + .is_some() + { + Ok(DuckDbExtensionStrategy::VendoredLocal { platform_dir }) + } else { + Ok(DuckDbExtensionStrategy::InstallFromRepository) + } + } + "windows" => Ok(DuckDbExtensionStrategy::InstallFromRepository), + _ => Err(etl_error!( + ErrorKind::ConfigError, + "Unsupported DuckDB extension platform", + format!("operating system `{os}` is not supported for DuckDB extensions") + )), + } +} + +fn repo_vendored_extension_root() -> PathBuf { + Path::new(env!("CARGO_MANIFEST_DIR")) + .parent() + .unwrap_or_else(|| Path::new(env!("CARGO_MANIFEST_DIR"))) + .join("vendor/duckdb/extensions") +} + +fn vendored_extension_dir( + platform_dir: &str, + env_override: Option, + container_root: &Path, + repo_root: &Path, +) -> EtlResult> { + if let Some(root) = env_override { + let directory = root.join(DUCKDB_EXTENSION_VERSION).join(platform_dir); + ensure_vendored_extension_dir(&directory)?; + return Ok(Some(directory)); + } + + for root in [container_root, repo_root] { + let directory = root.join(DUCKDB_EXTENSION_VERSION).join(platform_dir); + if !directory.exists() { + continue; + } + ensure_vendored_extension_dir(&directory)?; + return Ok(Some(directory)); + } + + Ok(None) +} + +fn require_vendored_extension_dir( + platform_dir: &str, + env_override: Option, + container_root: &Path, + repo_root: &Path, +) -> EtlResult { + vendored_extension_dir(platform_dir, env_override, container_root, repo_root)?.ok_or_else( + || { + etl_error!( + ErrorKind::ConfigError, + "Vendored DuckDB extensions not found", + format!( + "expected vendored DuckDB extensions in one of: {}, {}", + container_root.join(DUCKDB_EXTENSION_VERSION).join(platform_dir).display(), + repo_root.join(DUCKDB_EXTENSION_VERSION).join(platform_dir).display() + ) + ) + }, + ) +} + +fn ensure_vendored_extension_dir(directory: &Path) -> EtlResult<()> { + let missing = [DUCKLAKE_EXTENSION_FILE] + .into_iter() + .filter(|filename| !directory.join(filename).is_file()) + .collect::>(); + + if missing.is_empty() { + Ok(()) + } else { + Err(etl_error!( + ErrorKind::ConfigError, + "Vendored DuckDB extensions not found", + format!("missing {} in `{}`", missing.join(", "), directory.display()) + )) + } +} + +fn vendored_extension_path(extension_dir: &Path, filename: &str) -> EtlResult { + extension_dir.join(filename).to_str().map(std::string::ToString::to_string).ok_or_else(|| { + etl_error!( + ErrorKind::ConfigError, + "Vendored DuckDB extension path contains non-utf8 characters", + extension_dir.display().to_string() + ) + }) +} + +fn catalog_conninfo_from_url(catalog_url: &Url) -> EtlResult { + match catalog_url.scheme() { + "postgres" | "postgresql" => {} + scheme => { + return Err(etl_error!( + ErrorKind::ConfigError, + "Unsupported DuckLake catalog URL scheme", + format!("catalog URL scheme `{scheme}` is not supported") + )); + } + } + + let config = PgConfig::from_str(catalog_url.as_str()).map_err(|source| { + etl_error!( + ErrorKind::ConfigError, + "Invalid DuckLake PostgreSQL catalog URL", + source: source + ) + })?; + let explicit_query_options = explicit_query_options(catalog_url); + reject_unsupported_query_options(&explicit_query_options)?; + + let mut parts = Vec::new(); + + let hosts = serialize_hosts(config.get_hosts())?; + if !hosts.is_empty() { + push_conninfo_pair(&mut parts, "host", &hosts); + } + + let hostaddrs = config + .get_hostaddrs() + .iter() + .map(std::string::ToString::to_string) + .collect::>() + .join(","); + if !hostaddrs.is_empty() { + push_conninfo_pair(&mut parts, "hostaddr", &hostaddrs); + } + + let ports = config + .get_ports() + .iter() + .map(std::string::ToString::to_string) + .collect::>() + .join(","); + if !ports.is_empty() { + push_conninfo_pair(&mut parts, "port", &ports); + } + + if let Some(dbname) = config.get_dbname() { + push_conninfo_pair(&mut parts, "dbname", dbname); + } + if let Some(user) = config.get_user() { + push_conninfo_pair(&mut parts, "user", user); + } + if let Some(password) = config.get_password() { + let password = std::str::from_utf8(password).map_err(|source| { + etl_error!( + ErrorKind::ConfigError, + "DuckLake PostgreSQL catalog URL contains non-utf8 password", + source: source + ) + })?; + push_conninfo_pair(&mut parts, "password", password); + } + if let Some(options) = config.get_options() { + push_conninfo_pair(&mut parts, "options", options); + } + if let Some(application_name) = config.get_application_name() { + push_conninfo_pair(&mut parts, "application_name", application_name); + } + if let Some(ssl_cert) = config.get_ssl_cert() { + let ssl_cert = std::str::from_utf8(ssl_cert).map_err(|source| { + etl_error!( + ErrorKind::ConfigError, + "DuckLake PostgreSQL catalog URL contains non-utf8 sslcert", + source: source + ) + })?; + push_conninfo_pair(&mut parts, "sslcert", ssl_cert); + } + if let Some(ssl_key) = config.get_ssl_key() { + let ssl_key = std::str::from_utf8(ssl_key).map_err(|source| { + etl_error!( + ErrorKind::ConfigError, + "DuckLake PostgreSQL catalog URL contains non-utf8 sslkey", + source: source + ) + })?; + push_conninfo_pair(&mut parts, "sslkey", ssl_key); + } + if let Some(ssl_root_cert) = config.get_ssl_root_cert() { + let ssl_root_cert = std::str::from_utf8(ssl_root_cert).map_err(|source| { + etl_error!( + ErrorKind::ConfigError, + "DuckLake PostgreSQL catalog URL contains non-utf8 sslrootcert", + source: source + ) + })?; + push_conninfo_pair(&mut parts, "sslrootcert", ssl_root_cert); + } + + if explicit_query_options.contains("sslmode") { + parts.push(format!("sslmode={}", ssl_mode_to_str(config.get_ssl_mode())?)); + } + + if explicit_query_options.contains("connect_timeout") + && let Some(connect_timeout) = config.get_connect_timeout() + { + parts.push(format!("connect_timeout={}", connect_timeout.as_secs())); + } + if explicit_query_options.contains("tcp_user_timeout") + && let Some(tcp_user_timeout) = config.get_tcp_user_timeout() + { + parts.push(format!("tcp_user_timeout={}", tcp_user_timeout.as_millis())); + } + + if explicit_query_options.contains("keepalives") { + parts.push(format!("keepalives={}", if config.get_keepalives() { "1" } else { "0" })); + } + if explicit_query_options.contains("keepalives_idle") { + parts.push(format!("keepalives_idle={}", config.get_keepalives_idle().as_secs())); + } + if explicit_query_options.contains("keepalives_interval") + && let Some(keepalives_interval) = config.get_keepalives_interval() + { + parts.push(format!("keepalives_interval={}", keepalives_interval.as_secs())); + } + if explicit_query_options.contains("keepalives_retries") + && let Some(keepalives_retries) = config.get_keepalives_retries() + { + parts.push(format!("keepalives_retries={keepalives_retries}")); + } + + Ok(format!("postgres:{}", parts.join(" "))) +} + +fn explicit_query_options(catalog_url: &Url) -> BTreeSet { + catalog_url.query_pairs().map(|(key, _)| key.into_owned()).collect() +} + +fn reject_unsupported_query_options(explicit_query_options: &BTreeSet) -> EtlResult<()> { + let unsupported = + ["channel_binding", "load_balance_hosts", "replication", "target_session_attrs"] + .into_iter() + .filter(|key| explicit_query_options.contains(*key)) + .collect::>(); + + if unsupported.is_empty() { + Ok(()) + } else { + Err(etl_error!( + ErrorKind::ConfigError, + "DuckLake PostgreSQL catalog URL uses unsupported query parameters", + format!("unsupported parameters: {}", unsupported.join(", ")) + )) + } +} + +fn catalog_attach_target(catalog_url: &Url) -> EtlResult { + match catalog_url.scheme() { + "file" => Ok(catalog_url.as_str().to_owned()), + "postgres" | "postgresql" => catalog_conninfo_from_url(catalog_url), + scheme => Err(etl_error!( + ErrorKind::ConfigError, + "Unsupported DuckLake catalog URL scheme", + format!("catalog URL scheme `{scheme}` is not supported") + )), + } +} + +fn serialize_hosts(hosts: &[Host]) -> EtlResult { + let mut values = Vec::with_capacity(hosts.len()); + for host in hosts { + match host { + Host::Tcp(host) => values.push(host.clone()), + #[cfg(unix)] + Host::Unix(path) => { + let path = path.to_str().ok_or_else(|| { + etl_error!( + ErrorKind::ConfigError, + "DuckLake PostgreSQL catalog URL contains non-utf8 unix socket path" + ) + })?; + values.push(path.to_owned()); + } + } + } + + Ok(values.join(",")) +} + +fn push_conninfo_pair(parts: &mut Vec, key: &str, value: &str) { + parts.push(format!("{key}={}", quote_libpq_conninfo_value(value))); +} + +fn quote_libpq_conninfo_value(value: &str) -> String { + let mut quoted = String::from("'"); + for ch in value.chars() { + if matches!(ch, '\'' | '\\') { + quoted.push('\\'); + } + quoted.push(ch); + } + quoted.push('\''); + quoted +} + +fn ssl_mode_to_str(ssl_mode: SslMode) -> EtlResult<&'static str> { + match ssl_mode { + SslMode::Disable => Ok("disable"), + SslMode::Prefer => Ok("prefer"), + SslMode::Require => Ok("require"), + SslMode::VerifyCa => Ok("verify-ca"), + SslMode::VerifyFull => Ok("verify-full"), + _ => Err(etl_error!( + ErrorKind::ConfigError, + "DuckLake PostgreSQL catalog URL uses an unsupported sslmode" + )), + } +} + +fn validate_data_path(data_path: &Url) -> EtlResult<&str> { + match data_path.scheme() { + "file" | "s3" | "gs" => Ok(data_path.as_str()), + scheme => Err(etl_error!( + ErrorKind::ConfigError, + "Unsupported DuckLake data URL scheme", + format!("data URL scheme `{scheme}` is not supported") + )), + } +} + +fn build_setup_plan( + catalog_url: &Url, + data_path: &Url, + s3: Option<&S3Config>, + metadata_schema: Option<&str>, +) -> EtlResult { + let strategy = current_duckdb_extension_strategy()?; + let vendored_root = match strategy { + DuckDbExtensionStrategy::VendoredLocal { platform_dir } => { + Some(require_vendored_extension_dir( + platform_dir, + env::var_os(DUCKDB_EXTENSION_ROOT_ENV_VAR).map(PathBuf::from), + Path::new(CONTAINER_DUCKDB_EXTENSION_ROOT), + &repo_vendored_extension_root(), + )?) + } + DuckDbExtensionStrategy::InstallFromRepository => None, + }; + build_setup_plan_with_strategy( + catalog_url, + data_path, + s3, + metadata_schema, + strategy, + vendored_root.as_deref(), + ) +} + +fn build_setup_plan_with_strategy( + catalog_url: &Url, + data_path: &Url, + s3: Option<&S3Config>, + metadata_schema: Option<&str>, + strategy: DuckDbExtensionStrategy, + vendored_root: Option<&Path>, +) -> EtlResult { + let catalog_target = catalog_attach_target(catalog_url)?; + let data_path = validate_data_path(data_path)?; + + let needs_postgres = matches!(catalog_url.scheme(), "postgres" | "postgresql"); + let needs_httpfs = matches!(data_path.split(':').next(), Some("s3" | "gs")); + let lake_catalog = quote_identifier(LAKE_CATALOG); + let mut steps = vec![DuckLakeSetupStep { + label: "configure_writer_session", + sql: configure_writer_session_sql(), + }]; + let mut secret_options = BTreeMap::from([ + ("KEY_ID", quote_literal(s3.map(|s| s.access_key_id.as_str()).unwrap_or_default())), + ("REGION", quote_literal(s3.map(|s| s.region.as_str()).unwrap_or_default())), + ("SECRET", quote_literal(s3.map(|s| s.secret_access_key.as_str()).unwrap_or_default())), + ("URL_STYLE", quote_literal(s3.map(|s| s.url_style.as_str()).unwrap_or_default())), + ]); + + let extension_sql = match strategy { + DuckDbExtensionStrategy::VendoredLocal { .. } => { + let extension_root = vendored_root.ok_or_else(|| { + etl_error!( + ErrorKind::ConfigError, + "Vendored DuckDB extensions not found", + "Missing vendored DuckDB extension root" + ) + })?; + let ducklake_extension = + vendored_extension_path(extension_root, DUCKLAKE_EXTENSION_FILE)?; + let mut sql = format!("LOAD {};", quote_literal(&ducklake_extension)); + if needs_postgres { + let postgres_extension = + vendored_extension_path(extension_root, POSTGRES_SCANNER_EXTENSION_FILE)?; + sql.push_str(&format!(" LOAD {};", quote_literal(&postgres_extension))); + } + if needs_httpfs { + let httpfs_extension = + vendored_extension_path(extension_root, HTTPFS_EXTENSION_FILE)?; + sql.push_str(&format!(" LOAD {};", quote_literal(&httpfs_extension))); + } + sql + } + DuckDbExtensionStrategy::InstallFromRepository => { + let mut sql = String::from("INSTALL ducklake; LOAD ducklake;"); + if needs_postgres { + sql.push_str(" INSTALL postgres; LOAD postgres;"); + } + if needs_httpfs { + sql.push_str(" INSTALL httpfs; LOAD httpfs;"); + } + sql + } + }; + steps.push(DuckLakeSetupStep { label: "load_extensions", sql: extension_sql }); + + if needs_httpfs && let Some(s3) = s3 { + let secret_name = quote_identifier("ducklake_s3"); + if let Some(endpoint) = &s3.endpoint { + secret_options.insert("ENDPOINT", quote_literal(endpoint)); + } + + let secret_body = secret_options + .iter() + .map(|(key, value)| format!("{key} {value}")) + .collect::>() + .join(", "); + + steps.push(DuckLakeSetupStep { + label: "configure_object_store", + sql: format!( + "SET enable_http_metadata_cache = true; SET parquet_metadata_cache = true; CREATE \ + OR REPLACE SECRET {secret_name} (TYPE S3, {secret_body}, USE_SSL {});", + if s3.use_ssl { "true" } else { "false" } + ), + }); + } + let metadata_schema_clause = metadata_schema + .map(|schema| format!(", METADATA_SCHEMA {}", quote_literal(schema))) + .unwrap_or_default(); + + steps.push(DuckLakeSetupStep { + label: "attach_catalog", + sql: format!( + "ATTACH {} AS {lake_catalog} (DATA_PATH {}, DATA_INLINING_ROW_LIMIT {}, \ + AUTOMATIC_MIGRATION true{metadata_schema_clause});", + quote_literal(&format!("ducklake:{catalog_target}")), + quote_literal(data_path), + ATTACH_DATA_INLINING_ROW_LIMIT + ), + }); + steps.push(DuckLakeSetupStep { + label: "configure_parquet", + sql: configure_parquet_settings_sql(), + }); + + Ok(DuckLakeSetupPlan { steps }) +} + +#[derive(Debug)] +struct DuckLakeConnectionError { + message: String, +} + +impl DuckLakeConnectionError { + fn setup_phase(step: &DuckLakeSetupStep, error: duckdb::Error) -> Self { + let error_message = error.to_string(); + let error_message = if attach_step_uses_postgres_catalog(step) { + POSTGRES_PASSWORD_REGEX.replace_all(&error_message, "password='[redacted]'").to_string() + } else { + error_message + }; + + Self { + message: format!( + "ducklake duckdb connection setup phase `{}` failed: {error_message}", + step.label + ), + } + } + + fn validation(error: duckdb::Error) -> Self { + Self { message: format!("DuckLake DuckDB connection validation failed: {error}") } + } +} + +impl fmt::Display for DuckLakeConnectionError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> { + f.write_str(&self.message) + } +} + +impl error::Error for DuckLakeConnectionError {} + +fn remaining_ms_until(deadline: Instant) -> u64 { + deadline.checked_duration_since(Instant::now()).unwrap_or(Duration::ZERO).as_millis() as u64 +} + +/// One DuckDB connection tracked by the maintenance pool. +struct ManagedDuckLakeConnection { + conn: duckdb::Connection, + broken: bool, +} + +/// Async watchdog that interrupts one timed maintenance query when its deadline +/// expires. +struct DuckDbMaintenanceWatchdog { + timeout: Duration, + timed_out: Arc, + interrupt_tx: Option>>, + done_tx: Option>, + task: Option>, +} + +impl DuckDbMaintenanceWatchdog { + fn spawn(deadline: Instant, timeout: Duration) -> Self { + let timed_out = Arc::new(AtomicBool::new(false)); + let timeout_flag = Arc::clone(&timed_out); + let (interrupt_tx, interrupt_rx) = oneshot::channel::>(); + let (done_tx, done_rx) = oneshot::channel(); + let task = tokio::spawn(async move { + info!( + operation_kind = DUCKDB_MAINTENANCE_OPERATION_KIND, + timeout_ms = timeout.as_millis() as u64, + deadline_remaining_ms = remaining_ms_until(deadline), + "ducklake maintenance query watchdog task started: operation_kind={}, \ + timeout_ms={}, deadline_remaining_ms={}", + DUCKDB_MAINTENANCE_OPERATION_KIND, + timeout.as_millis(), + remaining_ms_until(deadline) + ); + let mut interrupt_rx = Box::pin(interrupt_rx); + let mut done_rx = Box::pin(done_rx); + let interrupt_handle = tokio::select! { + biased; + _ = &mut done_rx => { + info!( + operation_kind = DUCKDB_MAINTENANCE_OPERATION_KIND, + "ducklake maintenance query watchdog finished before interrupt handle: \ + operation_kind={}", + DUCKDB_MAINTENANCE_OPERATION_KIND + ); + return; + }, + result = &mut interrupt_rx => match result { + Ok(handle) => { + info!( + operation_kind = DUCKDB_MAINTENANCE_OPERATION_KIND, + deadline_remaining_ms = remaining_ms_until(deadline), + "ducklake maintenance query watchdog received interrupt handle before \ + deadline: operation_kind={}, deadline_remaining_ms={}", + DUCKDB_MAINTENANCE_OPERATION_KIND, + remaining_ms_until(deadline) + ); + handle + }, + Err(_) => { + warn!( + operation_kind = DUCKDB_MAINTENANCE_OPERATION_KIND, + "ducklake maintenance query watchdog interrupt sender dropped before \ + deadline: operation_kind={}", + DUCKDB_MAINTENANCE_OPERATION_KIND + ); + return; + }, + }, + _ = tokio::time::sleep_until(deadline) => { + timeout_flag.store(true, Ordering::Relaxed); + warn!( + operation_kind = DUCKDB_MAINTENANCE_OPERATION_KIND, + timeout_ms = timeout.as_millis() as u64, + "ducklake maintenance query watchdog deadline elapsed before interrupt \ + handle: operation_kind={}, timeout_ms={}", + DUCKDB_MAINTENANCE_OPERATION_KIND, + timeout.as_millis() + ); + tokio::select! { + biased; + _ = &mut done_rx => { + info!( + operation_kind = DUCKDB_MAINTENANCE_OPERATION_KIND, + "ducklake maintenance query watchdog received done after deadline \ + before interrupt handle: operation_kind={}", + DUCKDB_MAINTENANCE_OPERATION_KIND + ); + return; + }, + result = &mut interrupt_rx => match result { + Ok(handle) => { + warn!( + operation_kind = DUCKDB_MAINTENANCE_OPERATION_KIND, + "ducklake maintenance query watchdog received interrupt handle \ + after deadline: operation_kind={}", + DUCKDB_MAINTENANCE_OPERATION_KIND + ); + handle + }, + Err(_) => { + warn!( + operation_kind = DUCKDB_MAINTENANCE_OPERATION_KIND, + "ducklake maintenance query watchdog interrupt sender dropped \ + after deadline: operation_kind={}", + DUCKDB_MAINTENANCE_OPERATION_KIND + ); + return; + }, + }, + } + }, + }; + + if timeout_flag.load(Ordering::Relaxed) { + warn!( + operation_kind = DUCKDB_MAINTENANCE_OPERATION_KIND, + "ducklake maintenance query watchdog calling interrupt after timeout: \ + operation_kind={}", + DUCKDB_MAINTENANCE_OPERATION_KIND + ); + interrupt_handle.interrupt(); + warn!( + operation_kind = DUCKDB_MAINTENANCE_OPERATION_KIND, + "ducklake maintenance query watchdog interrupt returned after timeout: \ + operation_kind={}", + DUCKDB_MAINTENANCE_OPERATION_KIND + ); + return; + } + + tokio::select! { + biased; + _ = &mut done_rx => { + info!( + operation_kind = DUCKDB_MAINTENANCE_OPERATION_KIND, + deadline_remaining_ms = remaining_ms_until(deadline), + "ducklake maintenance query watchdog received done before deadline after \ + interrupt handle: operation_kind={}, deadline_remaining_ms={}", + DUCKDB_MAINTENANCE_OPERATION_KIND, + remaining_ms_until(deadline) + ); + } + _ = tokio::time::sleep_until(deadline) => { + timeout_flag.store(true, Ordering::Relaxed); + warn!( + operation_kind = DUCKDB_MAINTENANCE_OPERATION_KIND, + timeout_ms = timeout.as_millis() as u64, + "ducklake maintenance query watchdog deadline elapsed after interrupt \ + handle; calling interrupt: operation_kind={}, timeout_ms={}", + DUCKDB_MAINTENANCE_OPERATION_KIND, + timeout.as_millis() + ); + interrupt_handle.interrupt(); + warn!( + operation_kind = DUCKDB_MAINTENANCE_OPERATION_KIND, + "ducklake maintenance query watchdog interrupt returned after \ + handle/deadline path: operation_kind={}", + DUCKDB_MAINTENANCE_OPERATION_KIND + ); + } + } + }); + + Self { + timeout, + timed_out, + interrupt_tx: Some(interrupt_tx), + done_tx: Some(done_tx), + task: Some(task), + } + } + + fn publish_interrupt_handle(&mut self, handle: Arc) { + if let Some(interrupt_tx) = self.interrupt_tx.take() { + info!( + operation_kind = DUCKDB_MAINTENANCE_OPERATION_KIND, + timeout_ms = self.timeout.as_millis() as u64, + "ducklake maintenance query watchdog publishing interrupt handle: \ + operation_kind={}, timeout_ms={}", + DUCKDB_MAINTENANCE_OPERATION_KIND, + self.timeout.as_millis() + ); + let _ = interrupt_tx.send(handle); + } else { + warn!( + operation_kind = DUCKDB_MAINTENANCE_OPERATION_KIND, + "ducklake maintenance query watchdog interrupt handle publish skipped because \ + sender is gone: operation_kind={}", + DUCKDB_MAINTENANCE_OPERATION_KIND + ); + } + } + + fn finish(&mut self) { + if let Some(done_tx) = self.done_tx.take() { + info!( + operation_kind = DUCKDB_MAINTENANCE_OPERATION_KIND, + timed_out = self.timed_out(), + "ducklake maintenance query watchdog finish signal sent: operation_kind={}, \ + timed_out={}", + DUCKDB_MAINTENANCE_OPERATION_KIND, + self.timed_out() + ); + let _ = done_tx.send(()); + } else { + warn!( + operation_kind = DUCKDB_MAINTENANCE_OPERATION_KIND, + timed_out = self.timed_out(), + "ducklake maintenance query watchdog finish skipped because sender is gone: \ + operation_kind={}, timed_out={}", + DUCKDB_MAINTENANCE_OPERATION_KIND, + self.timed_out() + ); + } + } + + fn timed_out(&self) -> bool { + self.timed_out.load(Ordering::Relaxed) + } + + fn async_task_handle(&mut self) -> EtlResult> { + self.task.take().ok_or_else(|| { + etl_error!( + ErrorKind::DestinationError, + "Cannot get async task handle from DuckLake maintenance watchdog: task is None" + ) + }) + } +} + +fn duckdb_maintenance_timeout_error(timeout: Duration, stage: &'static str) -> EtlError { + etl_error!( + ErrorKind::DestinationQueryFailed, + "DuckLake maintenance blocking operation timed out", + format!( + "operation_kind={}, stage={stage}, timeout_ms={}", + DUCKDB_MAINTENANCE_OPERATION_KIND, + timeout.as_millis() + ) + ) +} + +fn abort_stuck_duckdb_maintenance_operation(timeout: Duration, abort_grace: Duration) -> ! { + tracing::error!( + operation_kind = DUCKDB_MAINTENANCE_OPERATION_KIND, + timeout_ms = timeout.as_millis() as u64, + abort_grace_ms = abort_grace.as_millis() as u64, + "ducklake maintenance blocking operation did not return after interrupt grace; aborting \ + process: operation_kind={}, timeout_ms={}, abort_grace_ms={}", + DUCKDB_MAINTENANCE_OPERATION_KIND, + timeout.as_millis(), + abort_grace.as_millis() + ); + process::abort(); +} + +#[derive(Clone)] +struct DuckLakeConnectionManager { + setup_plan: Arc, + disable_extension_autoload: bool, +} + +impl r2d2::ManageConnection for DuckLakeConnectionManager { + type Connection = ManagedDuckLakeConnection; + type Error = DuckLakeConnectionError; + + fn connect(&self) -> Result { + let conn = if self.disable_extension_autoload { + duckdb::Connection::open_in_memory_with_flags( + duckdb::Config::default() + .enable_autoload_extension(false) + .map_err(DuckLakeConnectionError::validation)?, + ) + .map_err(DuckLakeConnectionError::validation)? + } else { + duckdb::Connection::open_in_memory().map_err(DuckLakeConnectionError::validation)? + }; + for step in self.setup_plan.steps() { + info!(phase = step.label, "starting ducklake duckdb connection setup phase"); + conn.execute_batch(&step.sql) + .map_err(|error| DuckLakeConnectionError::setup_phase(step, error))?; + info!(phase = step.label, "ducklake duckdb connection setup phase finished"); + } + Ok(ManagedDuckLakeConnection { conn, broken: false }) + } + + fn is_valid( + &self, + conn: &mut ManagedDuckLakeConnection, + ) -> Result<(), DuckLakeConnectionError> { + conn.conn.execute_batch("SELECT 1").map_err(DuckLakeConnectionError::validation) + } + + fn has_broken(&self, conn: &mut ManagedDuckLakeConnection) -> bool { + conn.broken + } +} + +#[derive(Clone)] +struct DuckDbMaintenanceExecutor { + pool: Arc>, + blocking_slots: Arc, +} + +impl DuckDbMaintenanceExecutor { + async fn run(&self, operation: F) -> EtlResult + where + R: Send + 'static, + F: FnOnce(&duckdb::Connection) -> EtlResult + Send + 'static, + { + self.run_with_timeout(MAINTENANCE_QUERY_TIMEOUT, operation).await + } + + async fn run_with_timeout(&self, timeout: Duration, operation: F) -> EtlResult + where + R: Send + 'static, + F: FnOnce(&duckdb::Connection) -> EtlResult + Send + 'static, + { + let pool = Arc::clone(&self.pool); + let blocking_slots = Arc::clone(&self.blocking_slots); + let deadline = Instant::now() + timeout; + + info!( + operation_kind = DUCKDB_MAINTENANCE_OPERATION_KIND, + timeout_ms = timeout.as_millis() as u64, + abort_grace_ms = BLOCKING_ABORT_GRACE.as_millis() as u64, + available_permits = blocking_slots.available_permits(), + "ducklake maintenance blocking operation starting: operation_kind={}, timeout_ms={}, \ + abort_grace_ms={}, available_permits={}", + DUCKDB_MAINTENANCE_OPERATION_KIND, + timeout.as_millis(), + BLOCKING_ABORT_GRACE.as_millis(), + blocking_slots.available_permits() + ); + + let slot_wait_started = Instant::now(); + info!( + operation_kind = DUCKDB_MAINTENANCE_OPERATION_KIND, + deadline_remaining_ms = remaining_ms_until(deadline), + available_permits = blocking_slots.available_permits(), + "ducklake maintenance blocking operation waiting for semaphore slot: \ + operation_kind={}, deadline_remaining_ms={}, available_permits={}", + DUCKDB_MAINTENANCE_OPERATION_KIND, + remaining_ms_until(deadline), + blocking_slots.available_permits() + ); + let permit = + match tokio::time::timeout_at(deadline, Arc::clone(&blocking_slots).acquire_owned()) + .await + { + Ok(Ok(permit)) => permit, + Ok(Err(_)) => { + tracing::error!( + operation_kind = DUCKDB_MAINTENANCE_OPERATION_KIND, + "ducklake maintenance blocking operation semaphore closed: \ + operation_kind={}", + DUCKDB_MAINTENANCE_OPERATION_KIND + ); + return Err(etl_error!( + ErrorKind::ApplyWorkerPanic, + "DuckLake maintenance blocking slot acquisition failed" + )); + } + Err(_) => { + warn!( + operation_kind = DUCKDB_MAINTENANCE_OPERATION_KIND, + timeout_ms = timeout.as_millis() as u64, + slot_wait_ms = slot_wait_started.elapsed().as_millis() as u64, + "ducklake maintenance blocking operation timed out waiting for semaphore \ + slot: operation_kind={}, timeout_ms={}, slot_wait_ms={}", + DUCKDB_MAINTENANCE_OPERATION_KIND, + timeout.as_millis(), + slot_wait_started.elapsed().as_millis() + ); + return Err(duckdb_maintenance_timeout_error(timeout, "slot_wait")); + } + }; + trace!( + wait_ms = slot_wait_started.elapsed().as_millis() as u64, + "wait for ducklake maintenance blocking slot" + ); + info!( + operation_kind = DUCKDB_MAINTENANCE_OPERATION_KIND, + slot_wait_ms = slot_wait_started.elapsed().as_millis() as u64, + deadline_remaining_ms = remaining_ms_until(deadline), + "ducklake maintenance blocking operation acquired semaphore slot: operation_kind={}, \ + slot_wait_ms={}, deadline_remaining_ms={}", + DUCKDB_MAINTENANCE_OPERATION_KIND, + slot_wait_started.elapsed().as_millis(), + remaining_ms_until(deadline) + ); + + let mut watchdog = DuckDbMaintenanceWatchdog::spawn(deadline, timeout); + let watchdog_task = watchdog.async_task_handle()?; + let watchdog_timed_out = Arc::clone(&watchdog.timed_out); + let abort_deadline = deadline + BLOCKING_ABORT_GRACE; + + let blocking_task = tokio::task::spawn_blocking(move || -> EtlResult { + info!( + operation_kind = DUCKDB_MAINTENANCE_OPERATION_KIND, + deadline_remaining_ms = remaining_ms_until(deadline), + "ducklake maintenance blocking operation entered spawn_blocking task: \ + operation_kind={}, deadline_remaining_ms={}", + DUCKDB_MAINTENANCE_OPERATION_KIND, + remaining_ms_until(deadline) + ); + let _permit = permit; + let checkout_timeout = + deadline.checked_duration_since(Instant::now()).unwrap_or(Duration::ZERO); + if checkout_timeout.is_zero() { + warn!( + operation_kind = DUCKDB_MAINTENANCE_OPERATION_KIND, + "ducklake maintenance blocking operation deadline reached before pool \ + checkout: operation_kind={}", + DUCKDB_MAINTENANCE_OPERATION_KIND + ); + return Err(duckdb_maintenance_timeout_error(timeout, "pool_checkout")); + } + + let checkout_started = Instant::now(); + info!( + operation_kind = DUCKDB_MAINTENANCE_OPERATION_KIND, + checkout_timeout_ms = checkout_timeout.as_millis() as u64, + "ducklake maintenance blocking operation checking out pooled connection: \ + operation_kind={}, checkout_timeout_ms={}", + DUCKDB_MAINTENANCE_OPERATION_KIND, + checkout_timeout.as_millis() + ); + let mut pooled_conn = match pool.get_timeout(checkout_timeout) { + Ok(pooled_conn) => pooled_conn, + Err(error) if Instant::now() >= deadline => { + warn!( + operation_kind = DUCKDB_MAINTENANCE_OPERATION_KIND, + checkout_wait_ms = checkout_started.elapsed().as_millis() as u64, + timeout_ms = timeout.as_millis() as u64, + error = %error, + "ducklake maintenance blocking operation timed out checking out pooled \ + connection: operation_kind={}, checkout_wait_ms={}, timeout_ms={}, \ + error={}", + DUCKDB_MAINTENANCE_OPERATION_KIND, + checkout_started.elapsed().as_millis(), + timeout.as_millis(), + error + ); + return Err(duckdb_maintenance_timeout_error(timeout, "pool_checkout")); + } + Err(error) => { + warn!( + operation_kind = DUCKDB_MAINTENANCE_OPERATION_KIND, + checkout_wait_ms = checkout_started.elapsed().as_millis() as u64, + error = %error, + "ducklake maintenance blocking operation failed checking out pooled \ + connection: operation_kind={}, checkout_wait_ms={}, error={}", + DUCKDB_MAINTENANCE_OPERATION_KIND, + checkout_started.elapsed().as_millis(), + error + ); + return Err(etl_error!( + ErrorKind::DestinationConnectionFailed, + "Failed to check out DuckLake maintenance connection", + source: error + )); + } + }; + trace!( + wait_ms = checkout_started.elapsed().as_millis() as u64, + "wait for ducklake maintenance pool checkout" + ); + info!( + operation_kind = DUCKDB_MAINTENANCE_OPERATION_KIND, + checkout_wait_ms = checkout_started.elapsed().as_millis() as u64, + deadline_remaining_ms = remaining_ms_until(deadline), + "ducklake maintenance blocking operation checked out pooled connection: \ + operation_kind={}, checkout_wait_ms={}, deadline_remaining_ms={}", + DUCKDB_MAINTENANCE_OPERATION_KIND, + checkout_started.elapsed().as_millis(), + remaining_ms_until(deadline) + ); + + let operation_timeout = + deadline.checked_duration_since(Instant::now()).unwrap_or(Duration::ZERO); + if operation_timeout.is_zero() { + warn!( + operation_kind = DUCKDB_MAINTENANCE_OPERATION_KIND, + "ducklake maintenance blocking operation deadline reached before query \ + execution: operation_kind={}", + DUCKDB_MAINTENANCE_OPERATION_KIND + ); + pooled_conn.broken = true; + return Err(duckdb_maintenance_timeout_error(timeout, "query_execution")); + } + + let interrupt_handle = pooled_conn.conn.interrupt_handle(); + info!( + operation_kind = DUCKDB_MAINTENANCE_OPERATION_KIND, + operation_timeout_ms = operation_timeout.as_millis() as u64, + "ducklake maintenance blocking operation publishing interrupt handle before query \ + execution: operation_kind={}, operation_timeout_ms={}", + DUCKDB_MAINTENANCE_OPERATION_KIND, + operation_timeout.as_millis() + ); + watchdog.publish_interrupt_handle(interrupt_handle); + if watchdog.timed_out() { + warn!( + operation_kind = DUCKDB_MAINTENANCE_OPERATION_KIND, + "ducklake maintenance blocking operation timed out before query started; \ + marking pooled connection broken: operation_kind={}", + DUCKDB_MAINTENANCE_OPERATION_KIND + ); + pooled_conn.broken = true; + return Err(duckdb_maintenance_timeout_error(timeout, "query_execution")); + } + + let operation_started = Instant::now(); + info!( + operation_kind = DUCKDB_MAINTENANCE_OPERATION_KIND, + deadline_remaining_ms = remaining_ms_until(deadline), + "ducklake maintenance blocking operation invoking DuckDB closure: \ + operation_kind={}, deadline_remaining_ms={}", + DUCKDB_MAINTENANCE_OPERATION_KIND, + remaining_ms_until(deadline) + ); + let res = operation(&pooled_conn.conn); + let operation_duration_ms = operation_started.elapsed().as_millis() as u64; + info!( + operation_kind = DUCKDB_MAINTENANCE_OPERATION_KIND, + duration_ms = operation_duration_ms, + timed_out = watchdog.timed_out(), + result_is_error = res.is_err(), + "ducklake maintenance blocking operation DuckDB closure returned: \ + operation_kind={}, duration_ms={}, timed_out={}, result_is_error={}", + DUCKDB_MAINTENANCE_OPERATION_KIND, + operation_duration_ms, + watchdog.timed_out(), + res.is_err() + ); + watchdog.finish(); + trace!( + duration_ms = operation_duration_ms, + "ducklake maintenance blocking operation finished" + ); + if watchdog.timed_out() { + warn!( + operation_kind = DUCKDB_MAINTENANCE_OPERATION_KIND, + duration_ms = operation_duration_ms, + "ducklake maintenance blocking operation returned after timeout; marking \ + pooled connection broken: operation_kind={}, duration_ms={}", + DUCKDB_MAINTENANCE_OPERATION_KIND, + operation_duration_ms + ); + pooled_conn.broken = true; + return Err(duckdb_maintenance_timeout_error(timeout, "query_execution")); + } + if res.is_err() { + warn!( + operation_kind = DUCKDB_MAINTENANCE_OPERATION_KIND, + duration_ms = operation_duration_ms, + "ducklake maintenance blocking operation returned error; marking pooled \ + connection broken: operation_kind={}, duration_ms={}", + DUCKDB_MAINTENANCE_OPERATION_KIND, + operation_duration_ms + ); + pooled_conn.broken = true; + } else { + info!( + operation_kind = DUCKDB_MAINTENANCE_OPERATION_KIND, + duration_ms = operation_duration_ms, + "ducklake maintenance blocking operation returned success; pooled connection \ + remains healthy: operation_kind={}, duration_ms={}", + DUCKDB_MAINTENANCE_OPERATION_KIND, + operation_duration_ms + ); + } + + res + }); + + info!( + operation_kind = DUCKDB_MAINTENANCE_OPERATION_KIND, + abort_deadline_remaining_ms = remaining_ms_until(abort_deadline), + "ducklake maintenance blocking operation waiting for blocking task or abort deadline: \ + operation_kind={}, abort_deadline_remaining_ms={}", + DUCKDB_MAINTENANCE_OPERATION_KIND, + remaining_ms_until(abort_deadline) + ); + let blocking_result = tokio::select! { + biased; + result = blocking_task => result, + _ = tokio::time::sleep_until(abort_deadline) => { + abort_stuck_duckdb_maintenance_operation(timeout, BLOCKING_ABORT_GRACE); + } + }; + + match &blocking_result { + Ok(Ok(_)) => info!( + operation_kind = DUCKDB_MAINTENANCE_OPERATION_KIND, + timed_out = watchdog_timed_out.load(Ordering::Relaxed), + "ducklake maintenance blocking operation task joined with success: \ + operation_kind={}, timed_out={}", + DUCKDB_MAINTENANCE_OPERATION_KIND, + watchdog_timed_out.load(Ordering::Relaxed) + ), + Ok(Err(error)) => warn!( + operation_kind = DUCKDB_MAINTENANCE_OPERATION_KIND, + timed_out = watchdog_timed_out.load(Ordering::Relaxed), + error = %error, + "ducklake maintenance blocking operation task joined with error: operation_kind={}, \ + timed_out={}, error={}", + DUCKDB_MAINTENANCE_OPERATION_KIND, + watchdog_timed_out.load(Ordering::Relaxed), + error + ), + Err(error) => tracing::error!( + operation_kind = DUCKDB_MAINTENANCE_OPERATION_KIND, + timed_out = watchdog_timed_out.load(Ordering::Relaxed), + error = %error, + "ducklake maintenance blocking operation task join failed: operation_kind={}, \ + timed_out={}, error={}", + DUCKDB_MAINTENANCE_OPERATION_KIND, + watchdog_timed_out.load(Ordering::Relaxed), + error + ), + } + + info!( + operation_kind = DUCKDB_MAINTENANCE_OPERATION_KIND, + timed_out = watchdog_timed_out.load(Ordering::Relaxed), + "ducklake maintenance blocking operation awaiting watchdog task: operation_kind={}, \ + timed_out={}", + DUCKDB_MAINTENANCE_OPERATION_KIND, + watchdog_timed_out.load(Ordering::Relaxed) + ); + match watchdog_task.await { + Ok(()) => info!( + operation_kind = DUCKDB_MAINTENANCE_OPERATION_KIND, + timed_out = watchdog_timed_out.load(Ordering::Relaxed), + "ducklake maintenance blocking operation watchdog task joined: operation_kind={}, \ + timed_out={}", + DUCKDB_MAINTENANCE_OPERATION_KIND, + watchdog_timed_out.load(Ordering::Relaxed) + ), + Err(error) => { + tracing::error!( + operation_kind = DUCKDB_MAINTENANCE_OPERATION_KIND, + timed_out = watchdog_timed_out.load(Ordering::Relaxed), + error = %error, + "ducklake maintenance blocking operation watchdog task panicked: \ + operation_kind={}, timed_out={}, error={}", + DUCKDB_MAINTENANCE_OPERATION_KIND, + watchdog_timed_out.load(Ordering::Relaxed), + error + ); + return Err(etl_error!( + ErrorKind::ApplyWorkerPanic, + "DuckLake maintenance query watchdog task panicked" + )); + } + } + + blocking_result.map_err(|_| { + etl_error!( + ErrorKind::ApplyWorkerPanic, + "DuckLake maintenance blocking operation task panicked" + ) + })? + } +} + +async fn build_warm_ducklake_pool( + manager: DuckLakeConnectionManager, + pool_size: u32, +) -> EtlResult> { + tokio::task::spawn_blocking(move || -> EtlResult<_> { + let pool = r2d2::Pool::builder() + .max_size(pool_size) + .min_idle(Some(0)) + .connection_timeout(Duration::from_secs(4 * 60)) + .test_on_check_out(true) + .error_handler(Box::new(r2d2::NopErrorHandler)) + .build(manager) + .map_err(|source| { + etl_error!( + ErrorKind::DestinationConnectionFailed, + "Failed to build DuckLake maintenance connection pool", + source: source + ) + })?; + + let mut warmed_connections = Vec::with_capacity(pool_size as usize); + for _ in 0..pool_size { + warmed_connections.push(pool.get().map_err(|source| { + etl_error!( + ErrorKind::DestinationConnectionFailed, + "Failed to warm DuckLake maintenance connection pool", + source: source + ) + })?); + } + drop(warmed_connections); + + trace!(pool_size, "ducklake maintenance connection pool warmed"); + Ok(pool) + }) + .await + .map_err(|_| { + etl_error!( + ErrorKind::ApplyWorkerPanic, + "DuckLake maintenance connection pool initialization task panicked" + ) + })? +} + +fn format_query_error_detail(sql: &str) -> String { + let compact_sql = sql.split_whitespace().collect::>().join(" "); + format!("sql: {compact_sql}") +} + +fn attach_step_uses_postgres_catalog(step: &DuckLakeSetupStep) -> bool { + step.label == "attach_catalog" && step.sql.contains("ducklake:postgres:") +} + +/// Configuration for one external DuckLake maintenance run. +#[derive(Clone, Debug)] +pub struct DuckLakeMaintenanceConfig { + /// DuckLake PostgreSQL catalog URL. + pub catalog_url: Url, + /// DuckLake data path. + pub data_path: Url, + /// DuckDB connection pool size for the one-shot runner. + pub pool_size: u32, + /// Optional S3-compatible storage config. + pub s3: Option, + /// Optional DuckLake metadata schema. + pub metadata_schema: Option, + /// Optional DuckDB memory cache limit retained for config compatibility. + pub duckdb_memory_cache_limit: Option, + /// DuckLake `target_file_size` used by compaction. + pub maintenance_target_file_size: Option, + /// Inline flush operation config. + pub inline_flush: InlineFlushMaintenanceConfig, + /// Merge-adjacent-files operation config. + pub merge_adjacent_files: MergeAdjacentFilesMaintenanceConfig, + /// Rewrite-data-files operation config. + pub rewrite_data_files: RewriteDataFilesMaintenanceConfig, + /// Snapshot-expiration operation config. + pub expire_snapshots: ExpireSnapshotsMaintenanceConfig, + /// Old-file cleanup operation config. + pub cleanup_old_files: CleanupOldFilesMaintenanceConfig, +} + +/// Inline flush operation config. +#[derive(Clone, Copy, Debug)] +pub struct InlineFlushMaintenanceConfig { + /// Whether inline flush is enabled. + pub enabled: bool, + /// Minimum pending inlined bytes before flushing a table. + pub min_inlined_bytes: u64, +} + +/// Merge-adjacent-files operation config. +#[derive(Clone, Debug)] +pub struct MergeAdjacentFilesMaintenanceConfig { + /// Whether merge-adjacent-files is enabled. + pub enabled: bool, + /// Maximum compacted output files per table. + pub max_compacted_files: u32, + /// Maximum tables selected in one run. + pub max_tables_per_run: u32, + /// Target file size used during compaction. + pub target_file_size: String, +} + +/// Rewrite-data-files operation config. +#[derive(Clone, Copy, Debug)] +pub struct RewriteDataFilesMaintenanceConfig { + /// Whether rewrite-data-files is enabled. + pub enabled: bool, + /// Minimum active data-file count before rewrite is attempted. + pub min_active_data_files: i64, + /// Maximum tables selected in one run. + pub max_tables_per_run: u32, +} + +/// Snapshot-expiration operation config. +#[derive(Clone, Debug)] +pub struct ExpireSnapshotsMaintenanceConfig { + /// Whether snapshot expiration is enabled. + pub enabled: bool, + /// Retention window passed to DuckLake. + pub older_than: String, +} + +/// Old-file cleanup operation config. +#[derive(Clone, Debug)] +pub struct CleanupOldFilesMaintenanceConfig { + /// Whether old-file cleanup is enabled. + pub enabled: bool, + /// Retention window passed to DuckLake. + pub older_than: String, +} + +/// Structured outcome for one external maintenance run. +#[derive(Clone, Debug, Default, PartialEq)] +pub struct DuckLakeMaintenanceOutcome { + /// Tables whose inline data was flushed. + pub inline_flush_tables: u32, + /// Rows flushed from inlined storage. + pub inline_flush_rows: u64, + /// Tables passed to merge-adjacent-files. + pub merge_adjacent_files_tables: u32, + /// Files created by merge-adjacent-files. + pub merge_adjacent_files_created: u64, + /// Tables passed to rewrite-data-files. + pub rewrite_data_files_tables: u32, + /// Files created by rewrite-data-files. + pub rewrite_data_files_created: u64, + /// Snapshots expired by snapshot expiration. + pub expired_snapshots: u64, + /// Files removed by old-file cleanup. + pub cleaned_up_files: u64, +} + +impl DuckLakeMaintenanceOutcome { + /// Returns whether any operation did work. + pub fn applied(&self) -> bool { + self.inline_flush_rows > 0 + || self.merge_adjacent_files_created > 0 + || self.rewrite_data_files_tables > 0 + || self.rewrite_data_files_created > 0 + || self.expired_snapshots > 0 + || self.cleaned_up_files > 0 + } +} + +/// Runs one external DuckLake maintenance attempt. +pub async fn run_maintenance_once( + config: DuckLakeMaintenanceConfig, +) -> EtlResult { + validate_config(&config)?; + let cleanup_old_files_enabled = + config.cleanup_old_files.enabled || config.rewrite_data_files.enabled; + + info!( + pool_size = config.pool_size, + metadata_schema = config.metadata_schema.as_deref(), + inline_flush_enabled = config.inline_flush.enabled, + inline_flush_min_inlined_bytes = config.inline_flush.min_inlined_bytes, + merge_adjacent_files_enabled = config.merge_adjacent_files.enabled, + merge_adjacent_files_max_compacted_files = config.merge_adjacent_files.max_compacted_files, + merge_adjacent_files_max_tables_per_run = config.merge_adjacent_files.max_tables_per_run, + merge_adjacent_files_target_file_size = %config.merge_adjacent_files.target_file_size, + rewrite_data_files_enabled = config.rewrite_data_files.enabled, + rewrite_data_files_min_active_data_files = config.rewrite_data_files.min_active_data_files, + rewrite_data_files_max_tables_per_run = config.rewrite_data_files.max_tables_per_run, + expire_snapshots_enabled = config.expire_snapshots.enabled, + expire_snapshots_older_than = %config.expire_snapshots.older_than, + cleanup_old_files_enabled, + cleanup_old_files_explicitly_enabled = config.cleanup_old_files.enabled, + cleanup_old_files_older_than = %config.cleanup_old_files.older_than, + "ducklake external maintenance runner configured" + ); + + let duckdb = open_maintenance_executor(&config).await?; + let metadata_schema = match config.metadata_schema.clone() { + Some(metadata_schema) => metadata_schema, + None => resolve_metadata_schema(&duckdb).await?, + }; + info!( + metadata_schema = %metadata_schema, + "ducklake external maintenance metadata schema resolved" + ); + let metadata_pg_pool = PgPoolOptions::new() + .max_connections(1) + .connect_lazy(config.catalog_url.as_str()) + .map_err(|source| { + etl_error!( + ErrorKind::DestinationConnectionFailed, + "DuckLake catalog metadata pool configuration failed", + source: source + ) + })?; + let table_names = list_ducklake_tables(&metadata_pg_pool, &metadata_schema).await?; + info!( + table_count = table_names.len(), + tables = ?table_names, + "ducklake external maintenance discovered active tables" + ); + let mut outcome = DuckLakeMaintenanceOutcome::default(); + + if config.inline_flush.enabled { + run_inline_flush( + &duckdb, + &metadata_pg_pool, + &metadata_schema, + &table_names, + config.inline_flush, + &mut outcome, + ) + .await?; + } + + if config.merge_adjacent_files.enabled { + run_merge_adjacent_files( + &duckdb, + &metadata_pg_pool, + &metadata_schema, + &table_names, + &config.merge_adjacent_files, + &mut outcome, + ) + .await?; + } + + if config.rewrite_data_files.enabled { + merge_adjacent_files_for_rewrite(&duckdb).await?; + run_rewrite_data_files( + &duckdb, + &metadata_pg_pool, + &metadata_schema, + &table_names, + config.rewrite_data_files, + &mut outcome, + ) + .await?; + } + + if config.expire_snapshots.enabled { + run_expire_snapshots(&duckdb, &config.expire_snapshots, &mut outcome).await?; + } + + if cleanup_old_files_enabled { + run_cleanup_old_files(&duckdb, &config.cleanup_old_files, &mut outcome).await?; + } + + info!(outcome = ?outcome, applied = outcome.applied(), "ducklake external maintenance completed"); + Ok(outcome) +} + +/// Validates one maintenance runner config. +fn validate_config(config: &DuckLakeMaintenanceConfig) -> EtlResult<()> { + if !matches!(config.catalog_url.scheme(), "postgres" | "postgresql") { + return Err(etl_error!( + ErrorKind::ConfigError, + "DuckLake external maintenance requires a PostgreSQL catalog", + format!("unsupported catalog URL scheme `{}`", config.catalog_url.scheme()) + )); + } + if config.pool_size == 0 { + return Err(etl_error!( + ErrorKind::ConfigError, + "DuckLake external maintenance pool size must be greater than zero" + )); + } + Ok(()) +} + +/// Opens initialized DuckDB connections for maintenance. +async fn open_maintenance_executor( + config: &DuckLakeMaintenanceConfig, +) -> EtlResult { + let extension_strategy = current_duckdb_extension_strategy()?; + let target_file_size = config + .maintenance_target_file_size + .as_deref() + .or(Some(config.merge_adjacent_files.target_file_size.as_str())) + .unwrap_or(MAINTENANCE_TARGET_FILE_SIZE); + info!(target_file_size, "opening ducklake external maintenance connection"); + let setup_plan = Arc::new(build_setup_plan( + &config.catalog_url, + &config.data_path, + config.s3.as_ref(), + config.metadata_schema.as_deref(), + )?); + let manager = DuckLakeConnectionManager { + setup_plan, + disable_extension_autoload: extension_strategy.disables_autoload(), + }; + let pool = Arc::new(build_warm_ducklake_pool(manager, config.pool_size).await?); + let blocking_slots = Arc::new(Semaphore::new(config.pool_size as usize)); + let executor = DuckDbMaintenanceExecutor { pool, blocking_slots }; + let sql = maintenance_target_file_size_sql(Some(target_file_size)); + executor + .run(move |conn| { + conn.execute_batch(&sql).map_err(|source| { + etl_error!( + ErrorKind::DestinationQueryFailed, + "DuckLake target_file_size configuration failed", + format_query_error_detail(&sql), + source: source + ) + })?; + Ok(()) + }) + .await?; + info!(target_file_size, "ducklake external maintenance connection ready"); + Ok(executor) +} + +/// Resolves the hidden DuckLake metadata schema. +async fn resolve_metadata_schema(duckdb: &DuckDbMaintenanceExecutor) -> EtlResult { + duckdb.run(resolve_ducklake_metadata_schema_blocking).await +} + +/// Lists active DuckLake table names from the metadata catalog. +async fn list_ducklake_tables( + metadata_pg_pool: &PgPool, + metadata_schema: &str, +) -> EtlResult> { + let sql = format!( + "SELECT table_name FROM {}.{} WHERE end_snapshot IS NULL ORDER BY table_name", + quote_identifier(metadata_schema), + quote_identifier("ducklake_table") + ); + let rows: Vec<(String,)> = + sqlx::query_as(AssertSqlSafe(sql)).fetch_all(metadata_pg_pool).await.map_err(|source| { + etl_error!( + ErrorKind::DestinationQueryFailed, + "DuckLake table list query failed", + format!("metadata_schema={metadata_schema}"), + source: source + ) + })?; + Ok(rows.into_iter().map(|(table_name,)| table_name).collect()) +} + +/// Runs inline flush for tables that crossed the pending-inline threshold. +async fn run_inline_flush( + duckdb: &DuckDbMaintenanceExecutor, + metadata_pg_pool: &PgPool, + metadata_schema: &str, + table_names: &[String], + config: InlineFlushMaintenanceConfig, + outcome: &mut DuckLakeMaintenanceOutcome, +) -> EtlResult<()> { + info!( + min_inlined_bytes = config.min_inlined_bytes, + table_count = table_names.len(), + "ducklake inline flush evaluation started" + ); + let sampler = + DuckLakePendingInlineSizeSampler::new(metadata_schema.to_owned(), metadata_pg_pool.clone()); + for table_name in table_names { + let sizes = sampler.sample_table(table_name).await?; + if sizes.inlined_bytes < config.min_inlined_bytes { + info!( + table = %table_name, + inlined_bytes = sizes.inlined_bytes, + min_inlined_bytes = config.min_inlined_bytes, + "ducklake inline flush skipped below threshold" + ); + continue; + } + info!( + table = %table_name, + inlined_bytes = sizes.inlined_bytes, + min_inlined_bytes = config.min_inlined_bytes, + "ducklake inline flush executing" + ); + let table_name_for_query = table_name.clone(); + let rows = + duckdb.run(move |conn| flush_table_inlined_data(conn, &table_name_for_query)).await?; + outcome.inline_flush_tables = outcome.inline_flush_tables.saturating_add(1); + outcome.inline_flush_rows = outcome.inline_flush_rows.saturating_add(rows); + info!( + table = %table_name, + rows, + "ducklake inline flush completed" + ); + } + info!( + inline_flush_tables = outcome.inline_flush_tables, + inline_flush_rows = outcome.inline_flush_rows, + "ducklake inline flush evaluation finished" + ); + Ok(()) +} + +/// Runs bounded merge-adjacent-files on selected tables. +async fn run_merge_adjacent_files( + duckdb: &DuckDbMaintenanceExecutor, + metadata_pg_pool: &PgPool, + metadata_schema: &str, + table_names: &[String], + config: &MergeAdjacentFilesMaintenanceConfig, + outcome: &mut DuckLakeMaintenanceOutcome, +) -> EtlResult<()> { + info!( + max_compacted_files = config.max_compacted_files, + max_tables_per_run = config.max_tables_per_run, + table_count = table_names.len(), + "ducklake merge-adjacent-files evaluation started" + ); + let selected = select_merge_tables( + metadata_pg_pool, + metadata_schema, + table_names, + config.max_tables_per_run, + ) + .await?; + info!( + selected_tables = ?selected, + selected_count = selected.len(), + "ducklake merge-adjacent-files selected tables" + ); + for table_name in selected { + info!( + table = %table_name, + max_compacted_files = config.max_compacted_files, + "ducklake merge-adjacent-files executing" + ); + let table_name_for_query = table_name.clone(); + let max_compacted_files = config.max_compacted_files; + let files_created = duckdb + .run(move |conn| merge_adjacent_files(conn, &table_name_for_query, max_compacted_files)) + .await?; + outcome.merge_adjacent_files_tables = outcome.merge_adjacent_files_tables.saturating_add(1); + outcome.merge_adjacent_files_created = + outcome.merge_adjacent_files_created.saturating_add(files_created); + info!( + table = %table_name, + files_created, + "ducklake merge-adjacent-files completed" + ); + } + Ok(()) +} + +/// Runs DuckLake's whole-lake adjacent-file merge before rewrite-data-files. +async fn merge_adjacent_files_for_rewrite(duckdb: &DuckDbMaintenanceExecutor) -> EtlResult<()> { + let sql = format!("CALL ducklake_merge_adjacent_files({});", quote_literal(LAKE_CATALOG)); + info!( + sql = %sql, + "ducklake rewrite-triggered merge-adjacent-files executing" + ); + duckdb + .run(move |conn| { + conn.execute_batch(&sql).map_err(|source| { + etl_error!( + ErrorKind::DestinationQueryFailed, + "DuckLake rewrite-triggered merge adjacent files failed", + format_query_error_detail(&sql), + source: source + ) + })?; + Ok(()) + }) + .await?; + info!("ducklake rewrite-triggered merge-adjacent-files completed"); + Ok(()) +} + +/// Runs bounded rewrite-data-files on selected tables. +async fn run_rewrite_data_files( + duckdb: &DuckDbMaintenanceExecutor, + metadata_pg_pool: &PgPool, + metadata_schema: &str, + table_names: &[String], + config: RewriteDataFilesMaintenanceConfig, + outcome: &mut DuckLakeMaintenanceOutcome, +) -> EtlResult<()> { + info!( + min_active_data_files = config.min_active_data_files, + max_tables_per_run = config.max_tables_per_run, + table_count = table_names.len(), + "ducklake rewrite-data-files evaluation started" + ); + let selected = select_rewrite_tables( + metadata_pg_pool, + metadata_schema, + table_names, + config.min_active_data_files, + config.max_tables_per_run, + ) + .await?; + info!( + selected_tables = ?selected, + selected_count = selected.len(), + "ducklake rewrite-data-files selected tables" + ); + for table_name in selected { + info!( + table = %table_name, + "ducklake rewrite-data-files executing" + ); + let table_name_for_query = table_name.clone(); + let files_created = + duckdb.run(move |conn| rewrite_data_files(conn, &table_name_for_query)).await?; + outcome.rewrite_data_files_tables = outcome.rewrite_data_files_tables.saturating_add(1); + outcome.rewrite_data_files_created = + outcome.rewrite_data_files_created.saturating_add(files_created); + info!( + table = %table_name, + files_created, + "ducklake rewrite-data-files completed" + ); + } + Ok(()) +} + +/// Runs DuckLake snapshot expiration. +async fn run_expire_snapshots( + duckdb: &DuckDbMaintenanceExecutor, + config: &ExpireSnapshotsMaintenanceConfig, + outcome: &mut DuckLakeMaintenanceOutcome, +) -> EtlResult<()> { + info!( + older_than = %config.older_than, + "ducklake expire-snapshots executing" + ); + let older_than = config.older_than.clone(); + let expired_snapshots = duckdb.run(move |conn| expire_snapshots(conn, &older_than)).await?; + outcome.expired_snapshots = outcome.expired_snapshots.saturating_add(expired_snapshots); + info!( + older_than = %config.older_than, + expired_snapshots, + "ducklake expire-snapshots completed" + ); + Ok(()) +} + +/// Runs DuckLake old-file cleanup. +async fn run_cleanup_old_files( + duckdb: &DuckDbMaintenanceExecutor, + config: &CleanupOldFilesMaintenanceConfig, + outcome: &mut DuckLakeMaintenanceOutcome, +) -> EtlResult<()> { + info!( + older_than = %config.older_than, + "ducklake cleanup-old-files executing" + ); + let older_than = config.older_than.clone(); + let cleaned_up_files = duckdb.run(move |conn| cleanup_old_files(conn, &older_than)).await?; + outcome.cleaned_up_files = outcome.cleaned_up_files.saturating_add(cleaned_up_files); + info!( + older_than = %config.older_than, + cleaned_up_files, + "ducklake cleanup-old-files completed" + ); + Ok(()) +} + +/// Selects tables with small-file pressure for merge-adjacent-files. +async fn select_merge_tables( + metadata_pg_pool: &PgPool, + metadata_schema: &str, + table_names: &[String], + max_tables_per_run: u32, +) -> EtlResult> { + let mut selected = Vec::new(); + for table_name in table_names { + if is_etl_internal_table(table_name) { + info!( + table = %table_name, + "ducklake rewrite-data-files table skipped because it is internal ETL metadata" + ); + continue; + } + + let metrics = + query_table_storage_metrics(metadata_pg_pool, metadata_schema, table_name).await?; + if metrics.active_data_files > 1 && metrics.small_file_ratio() > 0.0 { + info!( + table = %table_name, + active_data_files = metrics.active_data_files, + small_file_ratio = metrics.small_file_ratio(), + "ducklake merge-adjacent-files table selected" + ); + selected.push(table_name.clone()); + } else { + info!( + table = %table_name, + active_data_files = metrics.active_data_files, + small_file_ratio = metrics.small_file_ratio(), + "ducklake merge-adjacent-files table skipped" + ); + } + if selected.len() >= max_tables_per_run as usize { + break; + } + } + Ok(selected) +} + +/// Selects tables with delete pressure for rewrite-data-files. +async fn select_rewrite_tables( + metadata_pg_pool: &PgPool, + metadata_schema: &str, + table_names: &[String], + min_active_data_files: i64, + max_tables_per_run: u32, +) -> EtlResult> { + let mut selected = Vec::new(); + for table_name in table_names { + if is_etl_internal_table(table_name) { + info!( + table = %table_name, + "ducklake rewrite-data-files table skipped because it is internal ETL metadata" + ); + continue; + } + + let metrics = + query_table_storage_metrics(metadata_pg_pool, metadata_schema, table_name).await?; + if should_rewrite(&metrics, min_active_data_files) { + info!( + table = %table_name, + active_data_files = metrics.active_data_files, + active_delete_files = metrics.active_delete_files, + deleted_row_ratio = metrics.deleted_row_ratio(), + min_active_data_files, + "ducklake rewrite-data-files table selected" + ); + selected.push(table_name.clone()); + } else { + info!( + table = %table_name, + active_data_files = metrics.active_data_files, + active_delete_files = metrics.active_delete_files, + deleted_row_ratio = metrics.deleted_row_ratio(), + min_active_data_files, + "ducklake rewrite-data-files table skipped" + ); + } + if selected.len() >= max_tables_per_run as usize { + break; + } + } + Ok(selected) +} + +fn is_etl_internal_table(table_name: &str) -> bool { + table_name.starts_with("__etl_") +} + +/// Returns whether a table should be rewritten. +fn should_rewrite(metrics: &DuckLakeTableStorageMetrics, min_active_data_files: i64) -> bool { + metrics.active_data_files > min_active_data_files +} + +/// Flushes inlined user data for one table. +#[doc(hidden)] +pub fn flush_table_inlined_data(conn: &duckdb::Connection, table_name: &str) -> EtlResult { + let flush_started = std::time::Instant::now(); + let sql = format!( + r#"SELECT COALESCE(SUM(rows_flushed), 0) + FROM ducklake_flush_inlined_data({}, table_name => {});"#, + quote_literal(LAKE_CATALOG), + quote_literal(table_name), + ); + let rows_flushed: i64 = conn.query_row(&sql, [], |row| row.get(0)).map_err(|source| { + etl_error!( + ErrorKind::DestinationQueryFailed, + "DuckLake inlined data flush failed", + format_query_error_detail(&sql), + source: source + ) + })?; + let rows_flushed = rows_flushed.max(0) as u64; + let flush_result = if rows_flushed > 0 { "flushed" } else { "noop" }; + histogram!( + ETL_DUCKLAKE_INLINE_FLUSH_ROWS, + RESULT_LABEL => flush_result, + ) + .record(rows_flushed as f64); + histogram!( + ETL_DUCKLAKE_INLINE_FLUSH_DURATION_SECONDS, + RESULT_LABEL => flush_result, + ) + .record(flush_started.elapsed().as_secs_f64()); + + if rows_flushed > 0 { + debug!( + table = %table_name, + rows_flushed, + "ducklake inlined data flushed" + ); + } else { + debug!( + table = %table_name, + "ducklake inlined data already flushed" + ); + } + Ok(rows_flushed) +} + +/// Calls DuckLake merge-adjacent-files for one table. +fn merge_adjacent_files( + conn: &duckdb::Connection, + table_name: &str, + max_compacted_files: u32, +) -> EtlResult { + let sql = format!( + "SELECT COALESCE(SUM(files_created), 0) FROM ducklake_merge_adjacent_files({}, {}, \ + max_compacted_files => {});", + quote_literal(LAKE_CATALOG), + quote_literal(table_name), + max_compacted_files + ); + count_maintenance_files(conn, &sql, "DuckLake merge adjacent files failed") +} + +/// Calls DuckLake rewrite-data-files for one table. +fn rewrite_data_files(conn: &duckdb::Connection, table_name: &str) -> EtlResult { + let sql = format!( + "CALL ducklake_rewrite_data_files({}, {});", + quote_literal(LAKE_CATALOG), + quote_literal(table_name) + ); + conn.execute_batch(&sql).map_err(|source| { + etl_error!( + ErrorKind::DestinationQueryFailed, + "DuckLake rewrite data files failed", + format_query_error_detail(&sql), + source: source + ) + })?; + Ok(0) +} + +/// Calls DuckLake snapshot expiration. +fn expire_snapshots(conn: &duckdb::Connection, older_than: &str) -> EtlResult { + let sql = format!( + "CALL ducklake_expire_snapshots({}, older_than => CAST(now() AS TIMESTAMP) - CAST({} AS \ + INTERVAL));", + quote_literal(LAKE_CATALOG), + quote_literal(older_than), + ); + count_maintenance_rows(conn, &sql, "DuckLake expire snapshots failed") +} + +/// Calls DuckLake old-file cleanup. +fn cleanup_old_files(conn: &duckdb::Connection, older_than: &str) -> EtlResult { + let sql = format!( + "CALL ducklake_cleanup_old_files({}, older_than => CAST(now() AS TIMESTAMP) - CAST({} AS \ + INTERVAL));", + quote_literal(LAKE_CATALOG), + quote_literal(older_than), + ); + count_maintenance_rows(conn, &sql, "DuckLake cleanup old files failed") +} + +/// Counts rows returned by one DuckLake maintenance call. +fn count_maintenance_rows( + conn: &duckdb::Connection, + sql: &str, + description: &'static str, +) -> EtlResult { + let mut statement = conn.prepare(sql).map_err(|source| { + etl_error!( + ErrorKind::DestinationQueryFailed, + description, + format_query_error_detail(sql), + source: source + ) + })?; + let mut rows = statement.query([]).map_err(|source| { + etl_error!( + ErrorKind::DestinationQueryFailed, + description, + format_query_error_detail(sql), + source: source + ) + })?; + let mut count = 0u64; + + while let Some(_row) = rows.next().map_err(|source| { + etl_error!( + ErrorKind::DestinationQueryFailed, + description, + format_query_error_detail(sql), + source: source + ) + })? { + count = count.saturating_add(1); + } + + Ok(count) +} + +/// Counts files returned by one DuckLake maintenance function. +fn count_maintenance_files( + conn: &duckdb::Connection, + sql: &str, + description: &'static str, +) -> EtlResult { + let files_created: i64 = conn.query_row(sql, [], |row| row.get(0)).map_err(|source| { + etl_error!( + ErrorKind::DestinationQueryFailed, + description, + format_query_error_detail(sql), + source: source + ) + })?; + Ok(files_created.max(0) as u64) +} + +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +struct DuckLakePendingInlineDataSizes { + inlined_bytes: u64, +} + +#[derive(Clone)] +struct DuckLakePendingInlineSizeSampler { + metadata_schema: String, + pool: PgPool, +} + +impl DuckLakePendingInlineSizeSampler { + fn new(metadata_schema: String, pool: PgPool) -> Self { + Self { metadata_schema, pool } + } + + async fn sample_table(&self, table_name: &str) -> EtlResult { + let sql = pending_inline_data_bytes_query(&self.metadata_schema); + let inlined_bytes: i64 = sqlx::query_scalar(AssertSqlSafe(sql)) + .bind(table_name) + .fetch_one(&self.pool) + .await + .map_err(|source| { + etl_error!( + ErrorKind::DestinationQueryFailed, + "DuckLake inline-size sampler query failed", + source: source + ) + })?; + Ok(record_pending_inline_data_sizes(inlined_bytes, table_name)) + } +} + +fn record_pending_inline_data_sizes( + inlined_bytes: i64, + table_name: &str, +) -> DuckLakePendingInlineDataSizes { + let inlined_bytes = inlined_bytes.max(0) as u64; + gauge!(ETL_DUCKLAKE_TABLE_ACTIVE_INLINED_DATA_BYTES, TABLE_LABEL => table_name.to_owned()) + .set(inlined_bytes as f64); + DuckLakePendingInlineDataSizes { inlined_bytes } +} + +fn pending_inline_data_bytes_query(metadata_schema: &str) -> String { + let metadata_schema_literal = quote_literal(metadata_schema); + let metadata_schema = quote_identifier(metadata_schema); + let ducklake_table = quote_identifier("ducklake_table"); + let ducklake_inlined_data_tables = quote_identifier("ducklake_inlined_data_tables"); + + format!( + r"WITH target_table AS ( + SELECT table_id + FROM {metadata_schema}.{ducklake_table} + WHERE end_snapshot IS NULL AND table_name = $1 + LIMIT 1 + ), + target_inline_tables AS ( + SELECT DISTINCT table_name + FROM {metadata_schema}.{ducklake_inlined_data_tables} + WHERE table_id = (SELECT table_id FROM target_table) + ), + inline_data_bytes AS ( + SELECT COALESCE( + SUM( + pg_total_relation_size( + to_regclass( + format('%I.%I', {metadata_schema_literal}, table_name) + ) + ) + ), + 0 + )::BIGINT AS total_bytes + FROM target_inline_tables + ), + inline_delete_bytes AS ( + SELECT COALESCE( + pg_total_relation_size( + to_regclass( + format( + '%I.%I', + {metadata_schema_literal}, + format('ducklake_inlined_delete_%s', (SELECT table_id FROM target_table)) + ) + ) + ), + 0 + )::BIGINT AS total_bytes + ) + SELECT + (SELECT total_bytes FROM inline_data_bytes) + + (SELECT total_bytes FROM inline_delete_bytes);" + ) +} + +#[derive(Clone, Debug)] +struct DuckLakeTableStorageMetrics { + active_data_files: i64, + small_data_files: i64, + active_data_rows: i64, + active_delete_files: i64, + deleted_rows: i64, +} + +impl DuckLakeTableStorageMetrics { + fn small_file_ratio(&self) -> f64 { + if self.active_data_files > 0 { + self.small_data_files.max(0) as f64 / self.active_data_files as f64 + } else { + 0.0 + } + } + + fn deleted_row_ratio(&self) -> f64 { + if self.active_data_rows > 0 { + self.deleted_rows.max(0) as f64 / self.active_data_rows as f64 + } else { + 0.0 + } + } +} + +async fn query_table_storage_metrics( + metadata_pg_pool: &PgPool, + metadata_schema: &str, + table_name: &str, +) -> EtlResult { + let sql = table_storage_metrics_query(metadata_schema); + let ( + active_data_files, + _active_data_bytes, + small_data_files, + active_data_rows, + active_delete_files, + _active_delete_bytes, + deleted_rows, + ): (i64, i64, i64, i64, i64, i64, i64) = sqlx::query_as(AssertSqlSafe(sql)) + .bind(table_name) + .fetch_one(metadata_pg_pool) + .await + .map_err(|source| { + etl_error!( + ErrorKind::DestinationQueryFailed, + "DuckLake table storage metrics query failed", + format!("table={table_name}, metadata_schema={metadata_schema}"), + source: source + ) + })?; + + Ok(DuckLakeTableStorageMetrics { + active_data_files, + small_data_files, + active_data_rows, + active_delete_files, + deleted_rows, + }) +} + +fn table_storage_metrics_query(metadata_schema: &str) -> String { + let metadata_schema = quote_identifier(metadata_schema); + let ducklake_table = quote_identifier("ducklake_table"); + let ducklake_data_file = quote_identifier("ducklake_data_file"); + let ducklake_delete_file = quote_identifier("ducklake_delete_file"); + + format!( + r#"WITH target_table AS ( + SELECT table_id + FROM {metadata_schema}.{ducklake_table} + WHERE end_snapshot IS NULL AND table_name = $1 + LIMIT 1 + ), + data_stats AS ( + SELECT + COUNT(*)::BIGINT AS active_data_files, + COALESCE(SUM(file_size_bytes), 0)::BIGINT AS active_data_bytes, + COALESCE(SUM(CASE WHEN file_size_bytes < {SMALL_FILE_SIZE_BYTES} THEN 1 ELSE 0 END), 0)::BIGINT AS small_data_files, + COALESCE(SUM(record_count), 0)::BIGINT AS active_data_rows + FROM {metadata_schema}.{ducklake_data_file} + WHERE end_snapshot IS NULL AND table_id = (SELECT table_id FROM target_table) + ), + delete_stats AS ( + SELECT + COUNT(*)::BIGINT AS active_delete_files, + COALESCE(SUM(file_size_bytes), 0)::BIGINT AS active_delete_bytes, + COALESCE(SUM(delete_count), 0)::BIGINT AS deleted_rows + FROM {metadata_schema}.{ducklake_delete_file} + WHERE end_snapshot IS NULL AND table_id = (SELECT table_id FROM target_table) + ) + SELECT + active_data_files, + active_data_bytes, + small_data_files, + active_data_rows, + active_delete_files, + active_delete_bytes, + deleted_rows + FROM data_stats CROSS JOIN delete_stats;"# + ) +} + +fn resolve_ducklake_metadata_schema_blocking(conn: &duckdb::Connection) -> EtlResult { + let metadata_catalog = format!("__ducklake_metadata_{LAKE_CATALOG}"); + let sql = format!( + r#"SELECT table_schema + FROM information_schema.tables + WHERE table_catalog = {} + AND table_name = 'ducklake_snapshot' + ORDER BY CASE + WHEN table_schema = 'main' THEN 0 + WHEN table_schema = 'ducklake' THEN 1 + ELSE 2 + END, + table_schema + LIMIT 1;"#, + quote_literal(&metadata_catalog), + ); + conn.query_row(&sql, [], |row| row.get::<_, String>(0)).map_err(|source| { + etl_error!( + ErrorKind::DestinationQueryFailed, + "DuckLake metadata schema query failed", + format_query_error_detail(&sql), + source: source + ) + }) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use tokio::sync::Semaphore; + + use super::*; + + fn metrics( + active_data_files: i64, + active_delete_files: i64, + deleted_rows: i64, + ) -> DuckLakeTableStorageMetrics { + DuckLakeTableStorageMetrics { + active_data_files, + small_data_files: 0, + active_data_rows: 100, + active_delete_files, + deleted_rows, + } + } + + fn make_maintenance_test_executor() -> DuckDbMaintenanceExecutor { + let manager = DuckLakeConnectionManager { + setup_plan: Arc::new(DuckLakeSetupPlan::default()), + disable_extension_autoload: cfg!(target_os = "linux"), + }; + let pool = r2d2::Pool::builder() + .max_size(1) + .min_idle(Some(0)) + .connection_timeout(Duration::from_secs(1)) + .test_on_check_out(true) + .error_handler(Box::new(r2d2::NopErrorHandler)) + .build(manager) + .expect("failed to build maintenance test pool"); + + DuckDbMaintenanceExecutor { + pool: Arc::new(pool), + blocking_slots: Arc::new(Semaphore::new(1)), + } + } + + #[test] + fn should_rewrite_requires_only_file_count() { + assert!(!should_rewrite(&metrics(39, 1, 50), 40)); + assert!(!should_rewrite(&metrics(40, 1, 50), 40)); + assert!(should_rewrite(&metrics(41, 0, 0), 40)); + } + + #[test] + fn outcome_reports_applied_work() { + assert!(!DuckLakeMaintenanceOutcome::default().applied()); + assert!( + DuckLakeMaintenanceOutcome { + inline_flush_rows: 1, + ..DuckLakeMaintenanceOutcome::default() + } + .applied() + ); + assert!( + DuckLakeMaintenanceOutcome { + rewrite_data_files_tables: 1, + ..DuckLakeMaintenanceOutcome::default() + } + .applied() + ); + assert!( + DuckLakeMaintenanceOutcome { + expired_snapshots: 1, + ..DuckLakeMaintenanceOutcome::default() + } + .applied() + ); + assert!( + DuckLakeMaintenanceOutcome { + cleaned_up_files: 1, + ..DuckLakeMaintenanceOutcome::default() + } + .applied() + ); + } + + #[tokio::test] + async fn maintenance_executor_timeout_releases_resources_for_follow_up_queries() { + let executor = make_maintenance_test_executor(); + + let error = executor + .run_with_timeout(Duration::from_millis(50), |_conn| -> EtlResult<()> { + std::thread::sleep(Duration::from_millis(100)); + Ok(()) + }) + .await + .expect_err("expected maintenance operation timeout"); + + assert_eq!(error.kind(), ErrorKind::DestinationQueryFailed); + assert_eq!(error.description(), Some("DuckLake maintenance blocking operation timed out")); + assert!( + error.detail().is_some_and(|detail| { + detail.contains("stage=pool_checkout") || detail.contains("stage=query_execution") + }), + "unexpected error detail: {error:?}" + ); + + let value = executor + .run_with_timeout(Duration::from_secs(1), |conn| -> EtlResult { + conn.query_row("SELECT 1;", [], |row| row.get::<_, i64>(0)).map_err(|source| { + etl_error!( + ErrorKind::DestinationQueryFailed, + "DuckLake maintenance timeout verification query failed", + source: source + ) + }) + }) + .await + .expect("expected follow-up query to succeed"); + + assert_eq!(value, 1); + } +} diff --git a/crates/etl-maintenance/src/lib.rs b/crates/etl-maintenance/src/lib.rs new file mode 100644 index 000000000..1ecdd672a --- /dev/null +++ b/crates/etl-maintenance/src/lib.rs @@ -0,0 +1,23 @@ +//! External maintenance coordination and runners for ETL. + +mod coordination; +mod materialization; + +#[cfg(feature = "ducklake")] +pub mod ducklake; + +#[cfg(feature = "ducklake")] +pub use coordination::KubernetesExternalMaintenanceStore; +pub use coordination::{ + ExternalMaintenanceOperationHistory, ExternalMaintenanceOperationPolicy, + ExternalMaintenanceOperationRequest, ExternalMaintenanceOperationRun, + ExternalMaintenanceOperations, ExternalMaintenancePause, ExternalMaintenanceReplicatorState, + ExternalMaintenanceReplicatorStatus, ExternalMaintenanceRequestOutcome, ExternalMaintenanceRun, + ExternalMaintenanceState, ExternalMaintenanceStore, ExternalMaintenanceWatcherConfig, + PostgresExternalMaintenanceStore, +}; +pub use materialization::{ + DisabledMaintenanceMaterializer, DuckLakeMaintenanceMaterialization, DuckLakeMaintenancePolicy, + MaintenanceIdentity, MaintenanceMaterializationError, MaintenanceMaterializer, + MaintenanceMaterializerKind, MaintenanceRuntimeRefs, PostgresMaintenanceMaterializer, +}; diff --git a/crates/etl-maintenance/src/materialization.rs b/crates/etl-maintenance/src/materialization.rs new file mode 100644 index 000000000..d200daf8f --- /dev/null +++ b/crates/etl-maintenance/src/materialization.rs @@ -0,0 +1,207 @@ +use std::error::Error; + +use async_trait::async_trait; +use serde::{Deserialize, Serialize}; +use sqlx::PgPool; +use thiserror::Error; + +use crate::coordination::{ExternalMaintenanceOperationPolicy, PostgresExternalMaintenanceStore}; + +/// DuckLake maintenance policy independent of the coordination backend. +#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)] +#[serde(rename_all = "camelCase")] +pub struct DuckLakeMaintenancePolicy { + /// Minimum time between maintenance runs, in seconds. + pub min_interval_seconds: u64, + /// Maximum time replication may be paused for one maintenance run, in + /// seconds. + pub max_pause_seconds: u64, + /// Minimum inlined bytes required before inline flush runs. + pub min_inlined_bytes: u64, + /// Maximum number of adjacent files compacted by one merge operation. + pub max_compacted_files: u32, + /// Maximum number of tables processed by each operation in one run. + pub max_tables_per_run: u32, + /// DuckLake target file size used for compaction. + pub target_file_size: String, + /// Deleted-row fraction that triggers data file rewrite. + pub delete_threshold: f64, + /// Minimum active data files required before data file rewrite runs. + pub min_active_data_files: i64, + /// CPU request for maintenance jobs, in millicores. + pub cpu_request_millicores: u32, + /// Memory request for maintenance jobs, in MiB. + pub memory_request_mib: u32, + /// Maximum runtime for one maintenance job, in seconds. + pub active_deadline_seconds: i64, + /// Backend-neutral operation enablement policy. + pub operation_policy: ExternalMaintenanceOperationPolicy, +} + +impl Default for DuckLakeMaintenancePolicy { + fn default() -> Self { + Self { + min_interval_seconds: 3600, + max_pause_seconds: 2700, + min_inlined_bytes: 10_000_000, + max_compacted_files: 32, + max_tables_per_run: 8, + target_file_size: "10MB".to_owned(), + delete_threshold: 0.5, + min_active_data_files: 40, + cpu_request_millicores: 1000, + memory_request_mib: 1024, + active_deadline_seconds: 1800, + operation_policy: ExternalMaintenanceOperationPolicy::default(), + } + } +} + +/// Stable identity for one pipeline's external maintenance runtime state. +#[derive(Debug, Clone)] +pub struct MaintenanceIdentity { + pub tenant_id: String, + pub pipeline_id: i64, + pub replicator_id: i64, + pub resource_prefix: String, +} + +/// Deployment-specific references required by the external maintenance runner. +#[derive(Debug, Clone)] +pub struct MaintenanceRuntimeRefs { + pub replicator_image: String, +} + +/// Backend-neutral materialization input for DuckLake external maintenance. +#[derive(Debug, Clone)] +pub struct DuckLakeMaintenanceMaterialization { + pub identity: MaintenanceIdentity, + pub policy: DuckLakeMaintenancePolicy, + pub runtime_refs: MaintenanceRuntimeRefs, +} + +/// Configured external maintenance backend. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum MaintenanceMaterializerKind { + Kubernetes, + Postgres, + Disabled, +} + +/// Errors raised while creating or deleting external maintenance runtime state. +#[derive(Debug, Error)] +pub enum MaintenanceMaterializationError { + #[error("maintenance backend `{0:?}` is not configured in this deployment")] + BackendNotConfigured(MaintenanceMaterializerKind), + + #[error("{backend:?} maintenance materialization failed")] + Backend { + backend: MaintenanceMaterializerKind, + #[source] + source: Box, + }, +} + +impl MaintenanceMaterializationError { + /// Wraps a backend-specific materialization error. + pub fn backend(backend: MaintenanceMaterializerKind, source: E) -> Self + where + E: Error + Send + Sync + 'static, + { + Self::Backend { backend, source: Box::new(source) } + } + + /// Wraps a Kubernetes materialization error. + pub fn kubernetes(source: E) -> Self + where + E: Error + Send + Sync + 'static, + { + Self::backend(MaintenanceMaterializerKind::Kubernetes, source) + } + + /// Wraps a Postgres materialization error. + fn postgres(source: etl::error::EtlError) -> Self { + Self::backend(MaintenanceMaterializerKind::Postgres, source) + } +} + +/// Deployment-side abstraction for external maintenance runtime state. +#[async_trait] +pub trait MaintenanceMaterializer: Send + Sync { + async fn reconcile_ducklake_maintenance( + &self, + input: DuckLakeMaintenanceMaterialization, + ) -> Result<(), MaintenanceMaterializationError>; + + async fn delete_ducklake_maintenance( + &self, + identity: MaintenanceIdentity, + ) -> Result<(), MaintenanceMaterializationError>; +} + +/// Explicit no-op materializer for deployments with no external maintenance +/// backend configured. +#[derive(Debug, Default)] +pub struct DisabledMaintenanceMaterializer; + +#[async_trait] +impl MaintenanceMaterializer for DisabledMaintenanceMaterializer { + async fn reconcile_ducklake_maintenance( + &self, + _input: DuckLakeMaintenanceMaterialization, + ) -> Result<(), MaintenanceMaterializationError> { + Ok(()) + } + + async fn delete_ducklake_maintenance( + &self, + _identity: MaintenanceIdentity, + ) -> Result<(), MaintenanceMaterializationError> { + Ok(()) + } +} + +/// Postgres materializer for deployments that coordinate maintenance without a +/// Kubernetes CRD. +#[derive(Debug, Clone)] +pub struct PostgresMaintenanceMaterializer { + pool: PgPool, +} + +impl PostgresMaintenanceMaterializer { + pub fn new(pool: PgPool) -> Self { + Self { pool } + } + + fn store(&self, pipeline_id: i64) -> PostgresExternalMaintenanceStore { + PostgresExternalMaintenanceStore::new(pipeline_id, self.pool.clone()) + } +} + +#[async_trait] +impl MaintenanceMaterializer for PostgresMaintenanceMaterializer { + async fn reconcile_ducklake_maintenance( + &self, + input: DuckLakeMaintenanceMaterialization, + ) -> Result<(), MaintenanceMaterializationError> { + let store = self.store(input.identity.pipeline_id); + store.ensure_schema().await.map_err(MaintenanceMaterializationError::postgres)?; + store + .ensure_pipeline_state(input.policy.operation_policy) + .await + .map_err(MaintenanceMaterializationError::postgres)?; + + Ok(()) + } + + async fn delete_ducklake_maintenance( + &self, + identity: MaintenanceIdentity, + ) -> Result<(), MaintenanceMaterializationError> { + let store = self.store(identity.pipeline_id); + store.ensure_schema().await.map_err(MaintenanceMaterializationError::postgres)?; + store.delete_pipeline_state().await.map_err(MaintenanceMaterializationError::postgres)?; + + Ok(()) + } +} diff --git a/crates/etl-replicator/Cargo.toml b/crates/etl-replicator/Cargo.toml index 0bf1782c2..79b01c513 100644 --- a/crates/etl-replicator/Cargo.toml +++ b/crates/etl-replicator/Cargo.toml @@ -18,6 +18,7 @@ configcat = { workspace = true } etl = { workspace = true } etl-config = { workspace = true, features = ["supabase"] } etl-destinations = { workspace = true, features = ["bigquery", "clickhouse", "ducklake", "iceberg"] } +etl-maintenance = { workspace = true, features = ["ducklake"] } etl-telemetry = { workspace = true } k8s-openapi = { workspace = true, features = ["latest"] } metrics = { workspace = true } diff --git a/crates/etl-replicator/src/bin/etl-ducklake-maintenance.rs b/crates/etl-replicator/src/bin/etl-ducklake-maintenance.rs index b348c5d62..c636a0d87 100644 --- a/crates/etl-replicator/src/bin/etl-ducklake-maintenance.rs +++ b/crates/etl-replicator/src/bin/etl-ducklake-maintenance.rs @@ -6,7 +6,7 @@ use etl_config::{ load_config, parse_ducklake_url, shared::{DestinationConfig, ReplicatorConfig}, }; -use etl_destinations::ducklake::{ +use etl_maintenance::ducklake::{ CleanupOldFilesMaintenanceConfig, DuckLakeMaintenanceConfig, ExpireSnapshotsMaintenanceConfig, InlineFlushMaintenanceConfig, MergeAdjacentFilesMaintenanceConfig, RewriteDataFilesMaintenanceConfig, S3Config as DuckLakeS3Config, run_maintenance_once, @@ -184,7 +184,7 @@ fn init_crypto_provider() { /// Initializes direct stdout logging for short-lived Kubernetes Jobs. fn init_stdout_tracing() { let filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| { - EnvFilter::new("etl_ducklake_maintenance=info,etl_destinations::ducklake=info") + EnvFilter::new("etl_ducklake_maintenance=info,etl_maintenance::ducklake=info") }); let _ = tracing_subscriber::registry() .with(filter) diff --git a/crates/etl-replicator/src/core.rs b/crates/etl-replicator/src/core.rs index c3808ca34..eb16a3986 100644 --- a/crates/etl-replicator/src/core.rs +++ b/crates/etl-replicator/src/core.rs @@ -12,12 +12,18 @@ use etl::{ }; use etl_config::{ Environment, parse_ducklake_url, - shared::{DestinationConfig, PgConnectionConfig, ReplicatorConfig}, + shared::{ + DestinationConfig, DuckLakeMaintenanceMode as ConfigDuckLakeMaintenanceMode, + PgConnectionConfig, ReplicatorConfig, + }, }; use etl_destinations::{ bigquery::BigQueryDestination, clickhouse::{ClickHouseClientConfig, ClickHouseDestination, ClickHouseInserterConfig}, - ducklake::{DuckLakeDestination, S3Config as DucklakeS3Config}, + ducklake::{ + DuckLakeDestination, DuckLakeExternalMaintenanceConfig, DuckLakeMaintenanceMode, + S3Config as DucklakeS3Config, + }, iceberg::{ DestinationNamespace, IcebergClient, IcebergDestination, S3_ACCESS_KEY_ID, S3_ENDPOINT, S3_SECRET_ACCESS_KEY, @@ -155,6 +161,7 @@ pub(crate) async fn start_replicator_with_config( duckdb_memory_cache_limit, maintenance_target_file_size, expire_snapshots_older_than, + maintenance_mode, } => { let s3_config = match (s3_access_key_id, s3_secret_access_key) { (Some(access_key_id), Some(secret_access_key)) => Some(DucklakeS3Config { @@ -174,7 +181,15 @@ pub(crate) async fn start_replicator_with_config( } }; - let destination = DuckLakeDestination::new( + let maintenance_mode = match maintenance_mode { + ConfigDuckLakeMaintenanceMode::Disabled => DuckLakeMaintenanceMode::Disabled, + ConfigDuckLakeMaintenanceMode::Kubernetes => DuckLakeMaintenanceMode::Kubernetes, + ConfigDuckLakeMaintenanceMode::Postgres => DuckLakeMaintenanceMode::Postgres, + }; + let external_maintenance = + DuckLakeExternalMaintenanceConfig { mode: maintenance_mode, pipeline_id }; + + let destination = DuckLakeDestination::new_with_external_maintenance( parse_ducklake_url(catalog_url).map_err(ReplicatorError::config)?, parse_ducklake_url(data_path).map_err(ReplicatorError::config)?, *pool_size, @@ -183,6 +198,7 @@ pub(crate) async fn start_replicator_with_config( duckdb_memory_cache_limit.clone(), maintenance_target_file_size.clone(), expire_snapshots_older_than.clone(), + external_maintenance, state_store.clone(), ) .await?; From 45b07ba044f68d6b031bfeb98c215adf355110f4 Mon Sep 17 00:00:00 2001 From: Jordan McQueen Date: Wed, 20 May 2026 16:03:53 +0900 Subject: [PATCH 16/29] Add ClickHouse ReplacingMergeTree support (#750) * feat(clickhouse): thread engine through destination config Adds `ClickHouseEngine` (`merge_tree` | `replacing_merge_tree`) in `etl-config::shared`, defaulting to `ReplacingMergeTree`. Wires the field through `DestinationConfig::ClickHouse`, `DestinationConfigWithoutSecrets::ClickHouse`, `ClickHouseInserterConfig`, and the replicator init path. No behavioral change: the destination still emits `MergeTree` DDL. Subsequent commits will branch DDL, encoding, and validation on the engine. * feat(clickhouse): RMT DDL builders and `__current` view Adds `create_replacing_merge_tree_sql` (PK-required, ordered by `primary_key_ordinal_position`, trailing `_etl_lsn` + `_etl_deleted`) and `create_current_view_sql` (FINAL + tombstone filter, exposes only user columns). Renames `build_create_table_sql` to `create_merge_tree_sql` and adds a `create_table_sql` dispatcher. Wires `create_table_with_metadata` and `recover_applying_metadata` through the dispatcher and emits the companion view under RMT. `expected_clickhouse_column_names` is now engine-aware via `trailing_cdc_column_names`. Adds `DdlKind::CreateView` for telemetry. No row-encoding changes yet; RMT tables cannot be written end-to-end until the next commit. * ref(clickhouse): polish RMT DDL helpers from review - Rename `issue_create_table_ddl` to `issue_create_table_stmt`. - Reword CDC/etl column constant doc-comments to lead with the engine. - Simplify `expected_clickhouse_column_names` to use a chained iterator. - Add blank lines around the early-return and final `Ok(format!)` in `create_replacing_merge_tree_sql`. * feat(clickhouse): RMT encoder and validation Row encoding now branches on engine via `append_cdc_columns`: - MergeTree (unchanged): `cdc_operation` + `cdc_lsn` (commit_lsn). - ReplacingMergeTree: `_etl_lsn` (start_lsn, the RMT version) + `_etl_deleted` tombstone. `start_lsn` is globally monotonic, so multi-event same-commit transactions tie-break correctly under `FINAL`. `PendingRow` carries both `commit_lsn` and `start_lsn`. Initial-copy rows under RMT encode with `_etl_lsn = 0` so any streaming event wins on FINAL. `validate_replica_identity_for_clickhouse` now takes the engine and rejects PK-less schemas under RMT in addition to the existing replica-identity gate. Adds `ClickHouseClient::server_version` and `ClickHouseDestination:: validate_engine_support`. `new` stays sync and side-effect free; callers (replicator, example, test_utils) invoke `validate_engine_support` after construction to surface RMT-on-old-CH (< 23.5) at startup. `UInt8` is added to the encoder for the new tombstone column. No integration-test parameterization yet; existing MergeTree tests keep their current behavior. * ref(clickhouse): polish RMT encoder helpers from review - Add blank lines around the early-return and trailing `Ok(())` in `ensure_engine_supported`. - Collapse the `ClickHouseValue::UInt8` doc comment to one line. * test(clickhouse): parameterize spine over both engines + RMT-only tests Every shared spine test in `pipeline.rs` now runs against both engines via an inner `..._inner(engine: ClickHouseEngine)` async fn plus two `#[tokio::test]` wrappers. Assertions are on current state read through a new `current_state_query` helper in `tests/support/clickhouse.rs`, which produces engine-appropriate SQL (MT: `LIMIT 1 BY (pks)` then drop DELETE rows; RMT: `FINAL` + `_etl_deleted = 0`). Adds `pipeline_rmt.rs` with RMT-only behaviors: - PK-less source rejection. - Same-LSN multi-event tx: INSERT+UPDATE collapses to the UPDATE under `FINAL` (confirms `start_lsn` tie-breaks). - Same-LSN tx: DELETE+INSERT yields the post-INSERT row. - `__current` view exposes user columns and reads current state. - Composite PK ORDER BY honors `primary_key_ordinal_position`. - Initial-copy then streaming UPDATE: streamed value wins. - `OPTIMIZE FINAL CLEANUP` physical removal (skipped gracefully when the server has the experimental gate disabled, which varies by CH version). Moves `sequential_transactions_preserve_commit_order` to `pipeline_mt.rs` because the cdc_lsn-ordering assertion has no analog under RMT (FINAL collapses the event log). Adds `ClickHouseTestDatabase::build_destination_with_engine`. The `column_types` filter now excludes both engines' trailing CDC columns. `install_crypto_provider` moves to the shared support module so all three test files share a single Once. * ref(tests): polish RMT/MT test structs and assertions from review - `EventLogRow`: drop `#[allow(dead_code)]`; add an `id` assertion that uses the field meaningfully alongside the CDC columns. - `IdRow`: replace with a `CountRow` against `SELECT count() AS count`, so the OPTIMIZE CLEANUP assertion no longer leans on an unused field. - PK-less rejection test: move the assertion into a real THEN block that reads the Errored phase reason and checks it mentions the PK requirement. - Reword the same-tx DELETE+INSERT assertion to "current state shows the re-inserted row, not the tombstone" -- clearer than "post-INSERT wins". * feat(clickhouse): engine-mismatch detection and PK ALTER guard under RMT Adds two destination-side safety checks: - `ensure_engine_matches`, called from `ensure_table_exists` on cache miss, queries the new `ClickHouseClient::table_engine` against `system.tables` and hard-errors when an existing CH table's engine differs from the configured engine. Catches the "operator flipped `engine: merge_tree` to `replacing_merge_tree`" misconfiguration before any RowBinary INSERT can mis-align. - `reject_pk_alters_under_rmt`, called from `apply_schema_diff`, refuses `columns_to_remove` or `columns_to_rename` entries that touch a primary-key column under RMT. RMT's `ORDER BY` is the source PK; a silent drop or rename would corrupt dedup. Unit tests cover the PK-ALTER guard (drop/rename, PK vs non-PK). Integration tests cover engine mismatch in both directions: MT then RMT, and RMT then MT. Both expect the table to land in the Errored phase with a reason naming both engines. * ref(clickhouse): polish engine-mismatch and PK ALTER guard from review - `ensure_engine_matches`: tighter logic flow (probe first, then compare) and a shorter doc comment. - `apply_schema_diff` RMT-PK-guard comment: explain why `ALTER TABLE` cannot rewrite the RMT ORDER BY expression, so a PK drop or rename would leave the table referring to a column that no longer exists. - Reword "would change the ORDER BY" -> "would invalidate the ORDER BY" in both the doc comment and the user-facing error message. - Drop wrapping `---` on multi-line GIVEN/WHEN markers in `pipeline_rmt.rs`; the visual frame only reads well on single-line markers. * docs(clickhouse): document engine flag, RMT behavior, MT path, migration Adds a `--clickhouse-engine` CLI flag to the etl-examples ClickHouse binary (snake_case values that mirror the YAML config: `merge_tree` and `replacing_merge_tree`, default `replacing_merge_tree`) and threads it through `ClickHouseInserterConfig`. Rewrites the ClickHouse section of `crates/etl-examples/README.md`: - Documents both engines with a quick-reference table. - Explains the ReplacingMergeTree layout (`_etl_lsn`, `_etl_deleted`), the auto-generated `__current` view, the FINAL read pattern, and `OPTIMIZE ... FINAL CLEANUP` guidance for operators. - Documents the MergeTree event-log layout and the latest-event-per-PK read pattern (`LIMIT 1 BY` + drop DELETEs). - Adds a migration note for pre-RMT pipelines: drop the CH tables and re-sync under RMT, or pass `--clickhouse-engine merge_tree` to keep the previous behavior. - Calls out the CH >= 23.5 requirement for RMT. * ref(clickhouse): polish engine-name helper + RMT comment from review - Move the ClickHouse engine-string lookup onto the type itself as `ClickHouseEngine::as_clickhouse_str() -> &'static str` in etl-config. Drops the awkward `engine_to_clickhouse_name` free function in `core.rs` and the duplicate helper in `pipeline_rmt.rs`. The method documents the distinction from the snake_case YAML/CLI form. - Reword the `reject_pk_alters_under_rmt` docstring: drop the parenthetical "(that's how `CREATE TABLE` was emitted)" and replace it with the actual SQL we emit, `CREATE TABLE ... ENGINE = ReplacingMergeTree(...) ORDER BY ()`. Makes the reasoning concrete rather than abstract. * ref(clickhouse): move RMT version requirement onto the engine type Adds `ClickHouseEngine::min_server_version() -> Option<(u32, u32)>` in etl-config, returning `Some((23, 5))` for RMT and `None` for MT. Drops the `MIN_REPLACING_MERGE_TREE_VERSION` free constant in core.rs and reshapes `ensure_engine_supported` to be the error-construction shell around `engine.min_server_version()`. The version requirement is a property of the engine variant itself, so it belongs on the type. core.rs keeps the `EtlResult` plumbing because etl-config doesn't (and shouldn't) take a dependency on the etl error infrastructure. The mismatch error message now names the configured engine literally (via `as_clickhouse_str()`) rather than hard-coding "ReplacingMergeTree" in the text, so the message stays correct if a future engine variant gets its own version floor. * docs(clickhouse): clarify docs * ref(clickhouse): clarify RMT DDL docstring + extract PK helper * ref(clickhouse): use is_some_and for PK rename lookup (clippy) * fix(clickhouse): refresh RMT current view on schema changes * fix(api): preserve ClickHouse engine config * doc(mdtable): format markdown table * ref(clickhouse): RMT version column = packed EventSequenceKey (UInt128) Replaces the RMT `_etl_lsn UInt64` version column with `_etl_version UInt128`, holding the packed source `EventSequenceKey` in the form `(commit_lsn << 64) | tx_ordinal`. This mirrors the codebase's canonical event-ordering primitive and gives the RMT version column a true total order across all events, rather than depending on the inferred invariant that `start_lsn` is unique per row touch (which holds in stock Postgres today but isn't a documented guarantee). - `PendingRow` carries `tx_ordinal: u64` instead of `start_lsn: PgLsn`. - `append_cdc_columns` computes the packed UInt128 for RMT via a new `etl_version_value` helper. - `ClickHouseValue::UInt128(u128)` added; `rb_encode_value` emits the value as 16-byte little-endian. - `create_replacing_merge_tree_sql` emits `_etl_version UInt128` and `ENGINE = ReplacingMergeTree(_etl_version, _etl_deleted)`. - Initial-copy rows use `commit_lsn = 0, tx_ordinal = 0`, so the packed version is 0 and any streamed event wins. - `column_types` test helper filter, RMT integration docstrings, README and example CLI doc updated for the new column name + encoding. * test(clickhouse): rename pipeline_{mt,rmt}.rs to long-form engine names * ref(clickhouse): expand MT/RMT abbreviations to MergeTree/ReplacingMergeTree * ref(clickhouse): use EventSequenceKey::as_u128 + trim destination doc - Add `EventSequenceKey::as_u128()` in etl, returning the canonical packed form (`commit_lsn` in the high 64 bits, `tx_ordinal` in the low 64 bits). The packing now lives next to the type it describes. - Drop the local `etl_version_value` helper in the ClickHouse destination and call `EventSequenceKey::new(...).as_u128()` directly. - Trim the `ClickHouseDestination` doc comment. Engine-specific layout belongs on `ClickHouseEngine` (where it already lives); the destination struct's doc just points at the engine type. --- Cargo.lock | 2 + crates/etl-api/src/configs/destination.rs | 120 ++- crates/etl-api/src/k8s/core.rs | 3 +- crates/etl-api/src/validation/validators.rs | 2 +- crates/etl-config/src/shared/destination.rs | 53 +- crates/etl-config/src/shared/mod.rs | 4 +- crates/etl-destinations/Cargo.toml | 1 + .../etl-destinations/src/clickhouse/client.rs | 51 ++ .../etl-destinations/src/clickhouse/core.rs | 537 +++++++++-- .../src/clickhouse/encoding.rs | 10 +- .../etl-destinations/src/clickhouse/schema.rs | 327 ++++++- .../src/clickhouse/test_utils.rs | 43 +- .../etl-destinations/tests/clickhouse/mod.rs | 2 + .../tests/clickhouse/pipeline.rs | 831 +++++++----------- .../tests/clickhouse/pipeline_merge_tree.rs | 141 +++ .../pipeline_replacing_merge_tree.rs | 761 ++++++++++++++++ .../tests/support/clickhouse.rs | 67 +- crates/etl-examples/Cargo.toml | 1 + crates/etl-examples/README.md | 76 +- crates/etl-examples/src/bin/clickhouse.rs | 52 +- crates/etl-replicator/src/core.rs | 5 +- crates/etl/src/types/event.rs | 8 + 22 files changed, 2454 insertions(+), 643 deletions(-) create mode 100644 crates/etl-destinations/tests/clickhouse/pipeline_merge_tree.rs create mode 100644 crates/etl-destinations/tests/clickhouse/pipeline_replacing_merge_tree.rs diff --git a/Cargo.lock b/Cargo.lock index 680af8dbe..e883f463c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1950,6 +1950,7 @@ dependencies = [ "clickhouse", "duckdb", "etl", + "etl-config", "etl-maintenance", "etl-postgres", "etl-telemetry", @@ -1987,6 +1988,7 @@ version = "0.1.0" dependencies = [ "clap", "etl", + "etl-config", "etl-destinations", "etl-telemetry", "k8s-openapi", diff --git a/crates/etl-api/src/configs/destination.rs b/crates/etl-api/src/configs/destination.rs index 0a4ef7c1a..c505f0524 100644 --- a/crates/etl-api/src/configs/destination.rs +++ b/crates/etl-api/src/configs/destination.rs @@ -1,6 +1,6 @@ use etl_config::{ SerializableSecretString, - shared::{DestinationConfig, DuckLakeMaintenanceMode, IcebergConfig}, + shared::{ClickHouseEngine, DestinationConfig, DuckLakeMaintenanceMode, IcebergConfig}, }; use secrecy::ExposeSecret; use serde::{Deserialize, Serialize}; @@ -60,6 +60,10 @@ pub enum FullApiDestinationConfig { #[schema(example = "my_db")] #[serde(deserialize_with = "crate::utils::trim_string")] database: String, + /// Table engine used for replicated tables. + #[schema(value_type = String, example = "replacing_merge_tree")] + #[serde(default)] + engine: ClickHouseEngine, }, Iceberg { #[serde(flatten)] @@ -163,8 +167,8 @@ impl From for FullApiDestinationConfig { max_staleness_mins, connection_pool_size: Some(connection_pool_size), }, - StoredDestinationConfig::ClickHouse { url, user, password, database } => { - Self::ClickHouse { url, user, password, database } + StoredDestinationConfig::ClickHouse { url, user, password, database, engine } => { + Self::ClickHouse { url, user, password, database, engine } } StoredDestinationConfig::Iceberg { config } => match config { StoredIcebergConfig::Supabase { @@ -253,6 +257,7 @@ pub enum StoredDestinationConfig { user: String, password: Option, database: String, + engine: ClickHouseEngine, }, Iceberg { config: StoredIcebergConfig, @@ -291,12 +296,15 @@ impl StoredDestinationConfig { max_staleness_mins, connection_pool_size, }, - Self::ClickHouse { url, user, password, database } => DestinationConfig::ClickHouse { - url, - user, - password: password.map(Into::into), - database, - }, + Self::ClickHouse { url, user, password, database, engine } => { + DestinationConfig::ClickHouse { + url, + user, + password: password.map(Into::into), + database, + engine, + } + } Self::Iceberg { config } => match config { StoredIcebergConfig::Supabase { project_ref, @@ -387,8 +395,8 @@ impl From for StoredDestinationConfig { connection_pool_size: connection_pool_size .unwrap_or(DestinationConfig::DEFAULT_CONNECTION_POOL_SIZE), }, - FullApiDestinationConfig::ClickHouse { url, user, password, database } => { - Self::ClickHouse { url, user, password, database } + FullApiDestinationConfig::ClickHouse { url, user, password, database, engine } => { + Self::ClickHouse { url, user, password, database, engine } } FullApiDestinationConfig::Iceberg { config } => match config { FullApiIcebergConfig::Supabase { @@ -487,7 +495,7 @@ impl Encrypt for StoredDestinationConfig { connection_pool_size, }) } - Self::ClickHouse { url, user, password, database } => { + Self::ClickHouse { url, user, password, database, engine } => { let encrypted_password = password .map(|p| encrypt_text(p.expose_secret().to_owned(), encryption_key)) .transpose()?; @@ -497,6 +505,7 @@ impl Encrypt for StoredDestinationConfig { user, password: encrypted_password, database, + engine, }) } Self::Iceberg { config } => match config { @@ -615,6 +624,8 @@ pub enum EncryptedStoredDestinationConfig { user: String, password: Option, database: String, + #[serde(default)] + engine: ClickHouseEngine, }, Iceberg { #[serde(flatten)] @@ -735,13 +746,19 @@ impl Decrypt for EncryptedStoredDestinationConfig { }) } }, - EncryptedStoredDestinationConfig::ClickHouse { url, user, password, database } => { + EncryptedStoredDestinationConfig::ClickHouse { + url, + user, + password, + database, + engine, + } => { let password = password .map(|p| decrypt_text(p, encryption_key)) .transpose()? .map(SerializableSecretString::from); - Ok(StoredDestinationConfig::ClickHouse { url, user, password, database }) + Ok(StoredDestinationConfig::ClickHouse { url, user, password, database, engine }) } Self::Ducklake { catalog_url, @@ -1061,6 +1078,7 @@ mod tests { user: "etl".to_owned(), password: Some(SerializableSecretString::from("secret".to_owned())), database: "analytics".to_owned(), + engine: ClickHouseEngine::MergeTree, }; let key = EncryptionKey { id: 1, key: generate_random_key::<32>().unwrap() }; @@ -1075,17 +1093,20 @@ mod tests { user: user1, password: p1, database: d1, + engine: e1, }, StoredDestinationConfig::ClickHouse { url: u2, user: user2, password: p2, database: d2, + engine: e2, }, ) => { assert_eq!(u1, u2); assert_eq!(user1, user2); assert_eq!(d1, d2); + assert_eq!(e1, e2); assert_eq!( p1.as_ref().map(|value| value.expose_secret()), p2.as_ref().map(|value| value.expose_secret()) @@ -1102,6 +1123,7 @@ mod tests { user: "etl".to_owned(), password: Some(SerializableSecretString::from("secret".to_owned())), database: "analytics".to_owned(), + engine: ClickHouseEngine::MergeTree, }; let stored: StoredDestinationConfig = full_config.clone().into(); @@ -1114,17 +1136,20 @@ mod tests { user: user1, password: p1, database: d1, + engine: e1, }, FullApiDestinationConfig::ClickHouse { url: u2, user: user2, password: p2, database: d2, + engine: e2, }, ) => { assert_eq!(u1, u2); assert_eq!(user1, user2); assert_eq!(d1, d2); + assert_eq!(e1, e2); assert_eq!( p1.as_ref().map(|value| value.expose_secret()), p2.as_ref().map(|value| value.expose_secret()) @@ -1134,6 +1159,26 @@ mod tests { } } + #[test] + fn stored_destination_config_into_etl_config_preserves_clickhouse_engine() { + let config = StoredDestinationConfig::ClickHouse { + url: Url::parse("https://example.com:8443").unwrap(), + user: "etl".to_owned(), + password: Some(SerializableSecretString::from("secret".to_owned())), + database: "analytics".to_owned(), + engine: ClickHouseEngine::MergeTree, + }; + + let etl_config = config.into_etl_config(); + + match etl_config { + DestinationConfig::ClickHouse { engine, .. } => { + assert_eq!(engine, ClickHouseEngine::MergeTree); + } + _ => panic!("Config types don't match"), + } + } + #[test] fn full_api_destination_config_deserializes_clickhouse_url() { let json = r#" @@ -1148,11 +1193,56 @@ mod tests { let deserialized: FullApiDestinationConfig = serde_json::from_str(json).unwrap(); match deserialized { - FullApiDestinationConfig::ClickHouse { url, user, password, database } => { + FullApiDestinationConfig::ClickHouse { url, user, password, database, engine } => { assert_eq!(url.as_str(), "https://example.com:8443/"); assert_eq!(user, "etl"); assert!(password.is_none()); assert_eq!(database, "analytics"); + assert_eq!(engine, ClickHouseEngine::default()); + } + _ => panic!("Deserialization failed or variant mismatch"), + } + } + + #[test] + fn full_api_destination_config_deserializes_clickhouse_engine() { + let json = r#" + { + "clickhouse": { + "url": "https://example.com:8443", + "user": "etl", + "database": "analytics", + "engine": "merge_tree" + } + } + "#; + + let deserialized: FullApiDestinationConfig = serde_json::from_str(json).unwrap(); + match deserialized { + FullApiDestinationConfig::ClickHouse { engine, .. } => { + assert_eq!(engine, ClickHouseEngine::MergeTree); + } + _ => panic!("Deserialization failed or variant mismatch"), + } + } + + #[test] + fn encrypted_stored_destination_config_defaults_legacy_clickhouse_engine() { + let json = r#" + { + "click_house": { + "url": "https://example.com:8443", + "user": "etl", + "password": null, + "database": "analytics" + } + } + "#; + + let deserialized: EncryptedStoredDestinationConfig = serde_json::from_str(json).unwrap(); + match deserialized { + EncryptedStoredDestinationConfig::ClickHouse { engine, .. } => { + assert_eq!(engine, ClickHouseEngine::default()); } _ => panic!("Deserialization failed or variant mismatch"), } diff --git a/crates/etl-api/src/k8s/core.rs b/crates/etl-api/src/k8s/core.rs index 745d8a7b8..2cda97dd7 100644 --- a/crates/etl-api/src/k8s/core.rs +++ b/crates/etl-api/src/k8s/core.rs @@ -519,7 +519,7 @@ mod tests { use std::sync::{Arc, Mutex}; use async_trait::async_trait; - use etl_config::SerializableSecretString; + use etl_config::{SerializableSecretString, shared::ClickHouseEngine}; use k8s_openapi::api::core::v1::ConfigMap; use super::*; @@ -566,6 +566,7 @@ mod tests { user: "default".to_owned(), password: password.map(ToOwned::to_owned).map(SerializableSecretString::from), database: "default".to_owned(), + engine: ClickHouseEngine::default(), } } diff --git a/crates/etl-api/src/validation/validators.rs b/crates/etl-api/src/validation/validators.rs index 498c73680..2e7acb833 100644 --- a/crates/etl-api/src/validation/validators.rs +++ b/crates/etl-api/src/validation/validators.rs @@ -925,7 +925,7 @@ impl Validator for DestinationValidator { ); validator.validate(ctx).await } - FullApiDestinationConfig::ClickHouse { url, user, password, database } => { + FullApiDestinationConfig::ClickHouse { url, user, password, database, .. } => { let validator = ClickHouseValidator::new( url.clone(), user.clone(), diff --git a/crates/etl-config/src/shared/destination.rs b/crates/etl-config/src/shared/destination.rs index 9d410b2b0..0ae24f958 100644 --- a/crates/etl-config/src/shared/destination.rs +++ b/crates/etl-config/src/shared/destination.rs @@ -12,6 +12,45 @@ const fn default_ducklake_pool_size() -> u32 { DestinationConfig::DEFAULT_DUCKLAKE_POOL_SIZE } +/// Table engine used by the ClickHouse destination when creating replicated +/// tables. +/// +/// `ReplacingMergeTree` (default) gives current-state reads via `FINAL` and +/// reclaims deleted rows on `OPTIMIZE ... FINAL CLEANUP`. `MergeTree` is an +/// append-only event-log layout retained for PK-less source tables. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Deserialize, Serialize, Default)] +#[serde(rename_all = "snake_case")] +pub enum ClickHouseEngine { + MergeTree, + #[default] + ReplacingMergeTree, +} + +impl ClickHouseEngine { + /// The literal engine name ClickHouse uses in `system.tables.engine` and + /// in `CREATE TABLE ... ENGINE = (...)`. Distinct from the + /// snake_case form used in YAML / CLI (`merge_tree`, + /// `replacing_merge_tree`). + pub const fn as_clickhouse_str(self) -> &'static str { + match self { + ClickHouseEngine::MergeTree => "MergeTree", + ClickHouseEngine::ReplacingMergeTree => "ReplacingMergeTree", + } + } + + /// Minimum ClickHouse server `(major, minor)` required to support this + /// engine, or `None` if any version works. + /// + /// `ReplacingMergeTree` requires >= 23.5 because earlier versions reject + /// the `(version, is_deleted)` argument pair we emit. + pub const fn min_server_version(self) -> Option<(u32, u32)> { + match self { + ClickHouseEngine::MergeTree => None, + ClickHouseEngine::ReplacingMergeTree => Some((23, 5)), + } + } +} + /// Runtime backend used for DuckLake external maintenance coordination. #[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] #[cfg_attr(feature = "utoipa", derive(ToSchema))] @@ -71,6 +110,11 @@ pub enum DestinationConfig { password: Option, /// ClickHouse target database database: String, + /// Table engine used for replicated tables. Defaults to + /// `ReplacingMergeTree`; set to `merge_tree` for the append-only + /// event-log layout. + #[serde(default)] + engine: ClickHouseEngine, }, Iceberg { #[serde(flatten)] @@ -268,6 +312,11 @@ pub enum DestinationConfigWithoutSecrets { user: String, /// ClickHouse target database database: String, + /// Table engine used for replicated tables. Defaults to + /// `ReplacingMergeTree`; set to `merge_tree` for the append-only + /// event-log layout. + #[serde(default)] + engine: ClickHouseEngine, }, Iceberg { #[serde(flatten)] @@ -318,8 +367,8 @@ impl From for DestinationConfigWithoutSecrets { max_staleness_mins, connection_pool_size, }, - DestinationConfig::ClickHouse { url, user, database, .. } => { - DestinationConfigWithoutSecrets::ClickHouse { url, user, database } + DestinationConfig::ClickHouse { url, user, password: _, database, engine } => { + DestinationConfigWithoutSecrets::ClickHouse { url, user, database, engine } } DestinationConfig::Iceberg { config } => { DestinationConfigWithoutSecrets::Iceberg { config: config.into() } diff --git a/crates/etl-config/src/shared/mod.rs b/crates/etl-config/src/shared/mod.rs index 2d2f74d38..7fec46333 100644 --- a/crates/etl-config/src/shared/mod.rs +++ b/crates/etl-config/src/shared/mod.rs @@ -13,8 +13,8 @@ pub use connection::{ TcpKeepaliveConfig, TlsConfig, }; pub use destination::{ - DestinationConfig, DestinationConfigWithoutSecrets, DuckLakeMaintenanceMode, IcebergConfig, - IcebergConfigWithoutSecrets, + ClickHouseEngine, DestinationConfig, DestinationConfigWithoutSecrets, DuckLakeMaintenanceMode, + IcebergConfig, IcebergConfigWithoutSecrets, }; pub use pipeline::{ BatchConfig, InvalidatedSlotBehavior, MemoryBackpressureConfig, PipelineConfig, diff --git a/crates/etl-destinations/Cargo.toml b/crates/etl-destinations/Cargo.toml index fba42ce4f..85922869a 100644 --- a/crates/etl-destinations/Cargo.toml +++ b/crates/etl-destinations/Cargo.toml @@ -79,6 +79,7 @@ chrono = { workspace = true, features = ["serde"] } clickhouse = { workspace = true, optional = true, features = ["inserter", "lz4", "rustls-tls"] } duckdb = { workspace = true, optional = true, features = ["bundled", "json", "parquet", "r2d2"] } etl = { workspace = true } +etl-config = { workspace = true } etl-maintenance = { workspace = true, optional = true } futures = { workspace = true, optional = true } gcp-bigquery-client = { workspace = true, optional = true, features = ["rust-tls", "aws-lc-rs"] } diff --git a/crates/etl-destinations/src/clickhouse/client.rs b/crates/etl-destinations/src/clickhouse/client.rs index 7f60ab156..29dff62b1 100644 --- a/crates/etl-destinations/src/clickhouse/client.rs +++ b/crates/etl-destinations/src/clickhouse/client.rs @@ -142,6 +142,8 @@ fn build_insert_rows_sql(table_name: &str) -> String { #[derive(Copy, Clone)] pub(crate) enum DdlKind { CreateTable, + CreateView, + DropView, AddColumn, DropColumn, RenameColumn, @@ -151,6 +153,8 @@ impl DdlKind { fn as_label(self) -> &'static str { match self { DdlKind::CreateTable => "create_table", + DdlKind::CreateView => "create_view", + DdlKind::DropView => "drop_view", DdlKind::AddColumn => "add_column", DdlKind::DropColumn => "drop_column", DdlKind::RenameColumn => "rename_column", @@ -217,6 +221,30 @@ impl ClickHouseClient { Ok(()) } + /// Returns the major/minor version pair from `SELECT version()`. + /// + /// Trailing components (patch, build) are ignored. Used at destination + /// construction to gate engine-specific feature requirements (e.g. + /// ReplacingMergeTree needs >= 23.5). + pub(crate) async fn server_version(&self) -> EtlResult<(u32, u32)> { + let raw = self.inner.query("SELECT version()").fetch_one::().await.map_err( + |err| etl_error!(ErrorKind::Unknown, "ClickHouse version query failed", source: err), + )?; + + let mut parts = raw.split('.'); + let major = parts.next().and_then(|s| s.parse::().ok()); + let minor = parts.next().and_then(|s| s.parse::().ok()); + + match (major, minor) { + (Some(major), Some(minor)) => Ok((major, minor)), + _ => Err(etl_error!( + ErrorKind::Unknown, + "Unable to parse ClickHouse server version", + format!("server returned '{raw}'") + )), + } + } + /// Executes a DDL statement (e.g. `CREATE TABLE IF NOT EXISTS …`) and /// records its duration in the `etl_clickhouse_ddl_duration_seconds` /// histogram labelled with the DDL `kind` and `table_name`. @@ -238,6 +266,29 @@ impl ClickHouseClient { result } + /// Returns the ClickHouse engine name for a table, or `None` if the table + /// does not exist in the current database. + pub(crate) async fn table_engine(&self, table_name: &str) -> EtlResult> { + let rows: Vec = self + .inner + .query( + "SELECT engine FROM system.tables WHERE database = currentDatabase() AND name = ?", + ) + .bind(table_name) + .fetch_all::() + .await + .map_err(|err| { + etl_error!( + ErrorKind::Unknown, + "ClickHouse engine lookup failed", + format!("table: {table_name}"), + source: err + ) + })?; + + Ok(rows.into_iter().next()) + } + /// Returns ClickHouse columns for a table in position order. pub(crate) async fn table_columns( &self, diff --git a/crates/etl-destinations/src/clickhouse/core.rs b/crates/etl-destinations/src/clickhouse/core.rs index b3b850875..bece73d7c 100644 --- a/crates/etl-destinations/src/clickhouse/core.rs +++ b/crates/etl-destinations/src/clickhouse/core.rs @@ -10,10 +10,11 @@ use etl::{ state::destination_metadata::{DestinationTableMetadata, DestinationTableSchemaStatus}, store::{schema::SchemaStore, state::StateStore}, types::{ - Cell, Event, IdentityType, OldTableRow, PgLsn, ReplicatedTableSchema, SchemaDiff, TableId, - TableRow, Type, UpdatedTableRow, is_array_type, + Cell, Event, EventSequenceKey, IdentityType, OldTableRow, PgLsn, ReplicatedTableSchema, + SchemaDiff, TableId, TableRow, Type, UpdatedTableRow, is_array_type, }, }; +use etl_config::shared::ClickHouseEngine; use parking_lot::RwLock; use tokio::task::JoinSet; use tracing::{debug, info, warn}; @@ -24,7 +25,10 @@ use crate::{ client::{ClickHouseClient, ClickHouseTableColumn, DdlKind}, encoding::{ClickHouseValue, cell_to_clickhouse_value}, metrics::register_metrics, - schema::{CDC_LSN_COLUMN_NAME, CDC_OPERATION_COLUMN_NAME, build_create_table_sql}, + schema::{ + create_current_view_sql, create_table_sql, drop_current_view_sql, + trailing_cdc_column_names, + }, }, table_name::try_stringify_table_name, }; @@ -57,13 +61,20 @@ impl std::fmt::Display for CdcOperation { /// A row pending insertion with its CDC metadata. struct PendingRow { - /// CDC op kind, written into the `cdc_operation` column. + /// CDC op kind. Drives both the MergeTree `cdc_operation` string and the + /// ReplacingMergeTree `_etl_deleted` tombstone flag. operation: CdcOperation, - /// Commit LSN of the source transaction, written into `cdc_lsn`. - lsn: PgLsn, - /// User column values in source schema order. The two CDC columns - /// (`cdc_operation`, `cdc_lsn`) are appended at encode time and are - /// not present here. + /// Transaction commit LSN. Written to the MergeTree `cdc_lsn` column, + /// and forms the high 64 bits of the ReplacingMergeTree `_etl_version` + /// column. + commit_lsn: PgLsn, + /// Zero-based ordinal of this event within its transaction. Forms the + /// low 64 bits of the ReplacingMergeTree `_etl_version` column so + /// multi-event same-commit transactions tie-break correctly under + /// `FINAL`. + tx_ordinal: u64, + /// User column values in source schema order. The trailing CDC columns + /// are appended at encode time and are not present here. cells: Vec, } @@ -72,18 +83,49 @@ fn cdc_lsn_to_clickhouse_value(lsn: PgLsn) -> ClickHouseValue { ClickHouseValue::UInt64(u64::from(lsn)) } +/// Appends the trailing engine-specific CDC columns to the row encoding. +/// +/// MergeTree: `cdc_operation` (String), `cdc_lsn` (UInt64 commit LSN). +/// ReplacingMergeTree: `_etl_version` (UInt128 packed `EventSequenceKey`), +/// `_etl_deleted` (UInt8 tombstone flag). +fn append_cdc_columns( + values: &mut Vec, + operation: CdcOperation, + commit_lsn: PgLsn, + tx_ordinal: u64, + engine: ClickHouseEngine, +) { + match engine { + ClickHouseEngine::MergeTree => { + values.push(ClickHouseValue::String(operation.to_string())); + values.push(cdc_lsn_to_clickhouse_value(commit_lsn)); + } + ClickHouseEngine::ReplacingMergeTree => { + let version = EventSequenceKey::new(commit_lsn, tx_ordinal).as_u128(); + values.push(ClickHouseValue::UInt128(version)); + values.push(ClickHouseValue::UInt8(matches!(operation, CdcOperation::Delete) as u8)); + } + } +} + /// Returns true if the ClickHouse type has an outer Nullable wrapper. fn clickhouse_type_expects_nullable_marker(type_name: &str) -> bool { type_name.starts_with("Nullable(") } -/// Returns expected ClickHouse column names for a replicated schema. -fn expected_clickhouse_column_names(schema: &ReplicatedTableSchema) -> Vec { - let mut names: Vec = - schema.column_schemas().map(|column| column.name.clone()).collect(); - names.push(CDC_OPERATION_COLUMN_NAME.to_owned()); - names.push(CDC_LSN_COLUMN_NAME.to_owned()); - names +/// Returns expected ClickHouse column names for a replicated schema under +/// the given engine: user columns in source order, then the engine's +/// trailing CDC columns. +fn expected_clickhouse_column_names( + schema: &ReplicatedTableSchema, + engine: ClickHouseEngine, +) -> Vec { + schema + .column_schemas() + .map(|c| c.name.as_str()) + .chain(trailing_cdc_column_names(engine).iter().copied()) + .map(str::to_owned) + .collect() } /// Derives RowBinary nullable flags from the actual ClickHouse table schema. @@ -153,6 +195,8 @@ pub struct ClickHouseInserterConfig { /// because incoming and outgoing buffers can both be near-full at once; /// could be made tunable later if needed. pub max_bytes_per_insert: u64, + /// Table engine used when creating replicated tables on ClickHouse. + pub engine: ClickHouseEngine, } impl ClickHouseInserterConfig { @@ -166,7 +210,10 @@ impl ClickHouseInserterConfig { impl Default for ClickHouseInserterConfig { fn default() -> Self { - Self { max_bytes_per_insert: Self::DEFAULT_MAX_BYTES_PER_INSERT } + Self { + max_bytes_per_insert: Self::DEFAULT_MAX_BYTES_PER_INSERT, + engine: ClickHouseEngine::default(), + } } } @@ -277,9 +324,8 @@ impl std::fmt::Display for ClickHouseOperationKind { /// CDC-capable ClickHouse destination that replicates Postgres tables. /// -/// Uses append-only MergeTree tables with two CDC columns (`cdc_operation`, -/// `cdc_lsn`) appended to each row. Rows are encoded as RowBinary and sent via -/// `INSERT INTO "table" FORMAT RowBinary` -- no column-name header required. +/// The table engine is configured via [`ClickHouseInserterConfig::engine`]; +/// see [`ClickHouseEngine`] for the engine-specific layouts. #[derive(Clone)] pub struct ClickHouseDestination { /// HTTP client used for all DDL and RowBinary INSERT traffic. @@ -326,6 +372,13 @@ where }) } + /// Probes the server version and rejects unsupported engine/version pairs. + /// Currently the only gate: ReplacingMergeTree requires CH >= 23.5. + pub async fn validate_engine_support(&self) -> EtlResult<()> { + let server_version = self.client.server_version().await?; + ensure_engine_supported(self.inserter_config.engine, server_version) + } + /// Creates a ClickHouse table for a never-before-seen `table_id`, /// bracketing the DDL with `DestinationTableMetadata` writes so the /// operation is crash-recoverable. @@ -356,14 +409,74 @@ where ); self.store.store_destination_table_metadata(table_id, metadata.clone()).await?; - let ddl = build_create_table_sql(clickhouse_table_name, schema.column_schemas()); - self.client.execute_ddl(DdlKind::CreateTable, &ddl).await?; + self.issue_create_table_stmt(clickhouse_table_name, schema).await?; self.store.store_destination_table_metadata(table_id, metadata.to_applied()).await?; Ok(()) } + /// Rejects writing to a pre-existing ClickHouse table whose engine does + /// not match the configured one. No-op if the table doesn't exist yet. + async fn ensure_engine_matches(&self, clickhouse_table_name: &str) -> EtlResult<()> { + let Some(existing) = self.client.table_engine(clickhouse_table_name).await? else { + return Ok(()); + }; + let configured = self.inserter_config.engine.as_clickhouse_str(); + if existing == configured { + return Ok(()); + } + + Err(etl_error!( + ErrorKind::ConfigError, + "ClickHouse table engine mismatch", + format!( + "Table '{clickhouse_table_name}' was previously created with engine '{existing}', \ + but the pipeline is configured for engine '{configured}'. Either drop the \ + destination table and re-sync, or reconfigure the pipeline's `engine` to match \ + the existing table." + ) + )) + } + + /// Issues the engine-correct `CREATE TABLE`, and under ReplacingMergeTree + /// also the companion `CREATE VIEW "
__current"`. Both statements + /// are `IF NOT EXISTS`, so retries on the recovery path are idempotent. + async fn issue_create_table_stmt( + &self, + clickhouse_table_name: &str, + schema: &ReplicatedTableSchema, + ) -> EtlResult<()> { + let engine = self.inserter_config.engine; + let ddl = create_table_sql(engine, clickhouse_table_name, schema.column_schemas())?; + self.client.execute_ddl(DdlKind::CreateTable, &ddl).await?; + + if matches!(engine, ClickHouseEngine::ReplacingMergeTree) { + let view_ddl = create_current_view_sql(clickhouse_table_name, schema.column_schemas()); + self.client.execute_ddl(DdlKind::CreateView, &view_ddl).await?; + } + Ok(()) + } + + /// Rebuilds the ReplacingMergeTree current-state view from the current + /// replicated schema. + /// + /// The base table can evolve through `ALTER TABLE`, but ClickHouse views + /// keep the projection they were created with. Drop and recreate the view + /// after schema changes so `"
__current"` follows ADD, DROP, and + /// RENAME changes. Both statements are idempotent for recovery retries. + async fn refresh_current_view( + &self, + clickhouse_table_name: &str, + schema: &ReplicatedTableSchema, + ) -> EtlResult<()> { + let drop_view = drop_current_view_sql(clickhouse_table_name); + self.client.execute_ddl(DdlKind::DropView, &drop_view).await?; + + let create_view = create_current_view_sql(clickhouse_table_name, schema.column_schemas()); + self.client.execute_ddl(DdlKind::CreateView, &create_view).await + } + /// Ensures the ClickHouse table for the given schema exists, returning /// `(clickhouse_table_name, nullable_flags)`. /// @@ -374,13 +487,19 @@ where &self, schema: &ReplicatedTableSchema, ) -> EtlResult<(String, Arc<[bool]>)> { - validate_replica_identity_for_clickhouse(schema)?; + validate_replica_identity_for_clickhouse(schema, self.inserter_config.engine)?; let clickhouse_table_name = try_stringify_table_name(schema.name())?; if let Some(flags) = self.table_cache.read().get(&clickhouse_table_name).cloned() { return Ok((clickhouse_table_name, flags)); } + // Engine-mismatch detection runs before any DDL or metadata mutation: + // a pre-existing table created under a different engine must hard-fail + // here, not silently get an idempotent CREATE TABLE IF NOT EXISTS that + // would then mis-align RowBinary on insert. + self.ensure_engine_matches(&clickhouse_table_name).await?; + let table_id = schema.id(); match self.store.get_destination_table_metadata(table_id).await? { None => { @@ -415,7 +534,8 @@ where // `Nullable(T)` even when the Postgres column is `NOT NULL`, so RowBinary must // include the nullable marker byte ClickHouse expects. let actual_columns = self.client.table_columns(&clickhouse_table_name).await?; - let expected_column_names = expected_clickhouse_column_names(schema); + let expected_column_names = + expected_clickhouse_column_names(schema, self.inserter_config.engine); let nullable_flags = nullable_flags_from_clickhouse_columns( &clickhouse_table_name, &expected_column_names, @@ -469,11 +589,10 @@ where metadata.replication_mask.clone(), ); let diff = old_schema.diff(schema); - self.apply_schema_diff(clickhouse_table_name, &diff, &old_schema).await?; + self.apply_schema_diff(clickhouse_table_name, &diff, &old_schema, schema).await?; } None => { - let ddl = build_create_table_sql(clickhouse_table_name, schema.column_schemas()); - self.client.execute_ddl(DdlKind::CreateTable, &ddl).await?; + self.issue_create_table_stmt(clickhouse_table_name, schema).await?; } } @@ -493,6 +612,7 @@ where ) -> EtlResult<()> { let (clickhouse_table_name, nullable_flags) = self.ensure_table_exists(schema).await?; + let engine = self.inserter_config.engine; let rows: Vec> = table_rows .into_iter() .map(|table_row| { @@ -501,10 +621,11 @@ where .into_iter() .map(cell_to_clickhouse_value) .collect::>()?; - // CDC columns: initial-copy rows are tagged as INSERT with LSN 0 - // (sentinel meaning "this row pre-dates the streaming cursor"). - values.push(ClickHouseValue::String(CdcOperation::Insert.to_string())); - values.push(cdc_lsn_to_clickhouse_value(PgLsn::from(0))); + // Initial-copy rows are tagged as INSERT with LSN 0 / tx_ordinal 0 + // (sentinel meaning "this row pre-dates the streaming cursor"). For + // ReplacingMergeTree, any streaming event then wins on FINAL because its packed + // `_etl_version` is non-zero. + append_cdc_columns(&mut values, CdcOperation::Insert, PgLsn::from(0), 0, engine); Ok(values) }) .collect::>()?; @@ -523,7 +644,7 @@ where /// Handles a schema change event (Relation) by computing the diff and /// applying ALTER TABLE statements. async fn handle_relation_event(&self, new_schema: &ReplicatedTableSchema) -> EtlResult<()> { - validate_replica_identity_for_clickhouse(new_schema)?; + validate_replica_identity_for_clickhouse(new_schema, self.inserter_config.engine)?; let table_id = new_schema.id(); let new_snapshot_id = new_schema.inner().snapshot_id; @@ -598,7 +719,7 @@ where // Compute and apply the diff. let diff = current_schema.diff(new_schema); if let Err(err) = - self.apply_schema_diff(clickhouse_table_name, &diff, ¤t_schema).await + self.apply_schema_diff(clickhouse_table_name, &diff, ¤t_schema, new_schema).await { warn!( "schema change failed for table {}: {}. Manual intervention may be required.", @@ -627,7 +748,8 @@ where } /// Applies a schema diff to a ClickHouse table: add columns, rename - /// columns, then drop columns (in that order for safety). + /// columns, then drop columns (in that order for safety), and refreshes + /// the ReplacingMergeTree current-state view when needed. /// /// New columns are placed AFTER the last existing user column (before the /// CDC columns) using ClickHouse's `AFTER` clause. This is critical because @@ -651,11 +773,30 @@ where clickhouse_table_name: &str, diff: &SchemaDiff, current_schema: &ReplicatedTableSchema, + new_schema: &ReplicatedTableSchema, ) -> EtlResult<()> { + let is_replacing_merge_tree = + matches!(self.inserter_config.engine, ClickHouseEngine::ReplacingMergeTree); if diff.is_empty() { + if is_replacing_merge_tree { + self.refresh_current_view(clickhouse_table_name, new_schema).await?; + } return Ok(()); } + // The ReplacingMergeTree table's `ORDER BY` clause is the source primary key, + // and ClickHouse uses that ORDER BY as the dedup key during merges. + // Dropping or renaming a PK column would invalidate that key in a + // way `ALTER TABLE` cannot fix, so reject the diff before any ALTER + // is issued. + if is_replacing_merge_tree { + reject_pk_alters_under_replacing_merge_tree( + clickhouse_table_name, + diff, + current_schema, + )?; + } + // Track the last user column name for AFTER placement. New columns // are inserted after this column, and each added column becomes the // new anchor for the next. `None` (no user columns in the current @@ -681,6 +822,10 @@ where self.client.drop_column(clickhouse_table_name, &column.name).await?; } + if is_replacing_merge_tree { + self.refresh_current_view(clickhouse_table_name, new_schema).await?; + } + Ok(()) } @@ -713,7 +858,8 @@ where .or_insert_with(|| (insert.replicated_table_schema, Vec::new())); entry.1.push(PendingRow { operation: CdcOperation::Insert, - lsn: insert.commit_lsn, + commit_lsn: insert.commit_lsn, + tx_ordinal: insert.tx_ordinal, cells: insert.table_row.into_values(), }); } @@ -739,7 +885,8 @@ where .or_insert_with(|| (update.replicated_table_schema, Vec::new())); entry.1.push(PendingRow { operation: CdcOperation::Update, - lsn: update.commit_lsn, + commit_lsn: update.commit_lsn, + tx_ordinal: update.tx_ordinal, cells: table_row.into_values(), }); } @@ -760,7 +907,8 @@ where .or_insert_with(|| (delete.replicated_table_schema, Vec::new())); entry.1.push(PendingRow { operation: CdcOperation::Delete, - lsn: delete.commit_lsn, + commit_lsn: delete.commit_lsn, + tx_ordinal: delete.tx_ordinal, cells: old_row.into_values(), }); } @@ -823,6 +971,7 @@ where } let mut join_set: JoinSet> = JoinSet::new(); + let engine = self.inserter_config.engine; for (clickhouse_table_name, nullable_flags, rows) in prepared { let client = self.client.clone(); let max_bytes = self.inserter_config.max_bytes_per_insert; @@ -830,13 +979,12 @@ where join_set.spawn(async move { let rows: Vec> = rows .into_iter() - .map(|PendingRow { operation, lsn, cells }| { + .map(|PendingRow { operation, commit_lsn, tx_ordinal, cells }| { let mut values: Vec = cells .into_iter() .map(cell_to_clickhouse_value) .collect::>()?; - values.push(ClickHouseValue::String(operation.to_string())); - values.push(cdc_lsn_to_clickhouse_value(lsn)); + append_cdc_columns(&mut values, operation, commit_lsn, tx_ordinal, engine); Ok(values) }) .collect::>()?; @@ -863,16 +1011,100 @@ where } } -/// Rejects replica identities the ClickHouse destination cannot represent. +/// Rejects schema diffs that would drop or rename a primary-key column on +/// an ReplacingMergeTree table. /// -/// `expand_key_row` assumes the key-only old-row image carries primary-key -/// values, so the row identity must match the primary key. `Full` is also -/// fine because it bypasses `expand_key_row` entirely. `AlternativeKey` -/// (a non-PK unique index) and `Missing` would either land identity values -/// in the wrong PK slots or leave us without enough data to write a -/// well-formed tombstone. +/// The destination emits `CREATE TABLE ... ENGINE = ReplacingMergeTree(...) +/// ORDER BY ()`, so the table's sort and dedup keys are bound to +/// those PK column names. ClickHouse `ALTER TABLE` can change column shapes +/// but cannot rewrite the ORDER BY expression, so a PK drop or rename would +/// leave the ORDER BY referring to a column that no longer exists (or has a +/// different meaning), silently breaking dedup. We error before the ALTER +/// reaches the server. +fn reject_pk_alters_under_replacing_merge_tree( + clickhouse_table_name: &str, + diff: &SchemaDiff, + current_schema: &ReplicatedTableSchema, +) -> EtlResult<()> { + for column in &diff.columns_to_remove { + if column.primary_key_ordinal_position.is_some() { + return Err(etl_error!( + ErrorKind::SourceSchemaError, + "ReplacingMergeTree does not support dropping a primary-key column", + format!( + "Table '{clickhouse_table_name}': DROP COLUMN '{name}' would invalidate the \ + ReplacingMergeTree ORDER BY / dedup key. Switch this table to `engine: \ + merge_tree` or restore the column on the source.", + name = column.name + ) + )); + } + } + + for rename in &diff.columns_to_rename { + let was_pk = current_schema + .column_schemas() + .find(|c| c.name == rename.old_name) + .is_some_and(|c| c.primary_key_ordinal_position.is_some()); + if was_pk { + return Err(etl_error!( + ErrorKind::SourceSchemaError, + "ReplacingMergeTree does not support renaming a primary-key column", + format!( + "Table '{clickhouse_table_name}': RENAME COLUMN '{old}' -> '{new}' would \ + invalidate the ReplacingMergeTree ORDER BY / dedup key. Switch this table to \ + `engine: merge_tree` or revert the rename on the source.", + old = rename.old_name, + new = rename.new_name + ) + )); + } + } + + Ok(()) +} + +/// Verifies the engine's `min_server_version()` constraint against the given +/// server version. The per-engine version requirement lives on +/// [`ClickHouseEngine`] itself; this function is just the error-construction +/// shell that surfaces the mismatch as an `EtlResult`. +fn ensure_engine_supported(engine: ClickHouseEngine, server_version: (u32, u32)) -> EtlResult<()> { + if let Some(min) = engine.min_server_version() + && server_version < min + { + let (min_major, min_minor) = min; + let (major, minor) = server_version; + + return Err(etl_error!( + ErrorKind::ConfigError, + "ClickHouse server version is too old for the configured engine", + format!( + "Detected ClickHouse {major}.{minor}; engine `{cfg}` requires \ + {min_major}.{min_minor} or newer. Upgrade ClickHouse or set `engine: merge_tree`.", + cfg = engine.as_clickhouse_str() + ) + )); + } + + Ok(()) +} + +/// Rejects replica identities and schemas the ClickHouse destination cannot +/// represent for the configured engine. +/// +/// Common to both engines: `expand_key_row` assumes the key-only old-row image +/// carries primary-key values, so the row identity must match the primary key. +/// `Full` is also fine because it bypasses `expand_key_row` entirely. +/// `AlternativeKey` (a non-PK unique index) and `Missing` would either land +/// identity values in the wrong PK slots or leave us without enough data to +/// write a well-formed tombstone. +/// +/// ReplacingMergeTree-only: the source table must have a primary key. +/// ReplacingMergeTree uses the PK as `ORDER BY`, which is also the dedup key; +/// without a PK there is nothing to merge on. fn validate_replica_identity_for_clickhouse( replicated_table_schema: &ReplicatedTableSchema, + engine: ClickHouseEngine, ) -> EtlResult<()> { if !replicated_table_schema.all_primary_key_columns_replicated() { let omitted_columns = replicated_table_schema @@ -891,6 +1123,20 @@ fn validate_replica_identity_for_clickhouse( )); } + if matches!(engine, ClickHouseEngine::ReplacingMergeTree) + && replicated_table_schema.primary_key_column_schemas().next().is_none() + { + return Err(etl_error!( + ErrorKind::SourceSchemaError, + "ClickHouse ReplacingMergeTree requires a primary key", + format!( + "Table '{}' has no primary-key columns; set `engine: merge_tree` or define a PK \ + on the source table.", + replicated_table_schema.name() + ) + )); + } + match replicated_table_schema.identity_type() { IdentityType::PrimaryKey | IdentityType::Full => Ok(()), identity_type => Err(etl_error!( @@ -1052,6 +1298,7 @@ mod tests { }; use super::*; + use crate::clickhouse::schema::{CDC_LSN_COLUMN_NAME, CDC_OPERATION_COLUMN_NAME}; fn clickhouse_column(name: &str, type_name: &str) -> ClickHouseTableColumn { ClickHouseTableColumn { name: name.to_owned(), type_name: type_name.to_owned() } @@ -1094,39 +1341,213 @@ mod tests { #[test] fn validate_replica_identity_for_clickhouse_accepts_primary_key() { - validate_replica_identity_for_clickhouse(&replicated_schema(IdentityType::PrimaryKey)) - .unwrap(); + validate_replica_identity_for_clickhouse( + &replicated_schema(IdentityType::PrimaryKey), + ClickHouseEngine::MergeTree, + ) + .unwrap(); } #[test] fn validate_replica_identity_for_clickhouse_accepts_full() { - validate_replica_identity_for_clickhouse(&replicated_schema(IdentityType::Full)).unwrap(); + validate_replica_identity_for_clickhouse( + &replicated_schema(IdentityType::Full), + ClickHouseEngine::MergeTree, + ) + .unwrap(); } #[test] fn validate_replica_identity_for_clickhouse_rejects_alternative_key() { - let err = validate_replica_identity_for_clickhouse(&replicated_schema( - IdentityType::AlternativeKey, - )) + let err = validate_replica_identity_for_clickhouse( + &replicated_schema(IdentityType::AlternativeKey), + ClickHouseEngine::MergeTree, + ) .unwrap_err(); assert_eq!(err.kind(), ErrorKind::SourceSchemaError); } #[test] fn validate_replica_identity_for_clickhouse_rejects_missing() { - let err = - validate_replica_identity_for_clickhouse(&replicated_schema(IdentityType::Missing)) - .unwrap_err(); + let err = validate_replica_identity_for_clickhouse( + &replicated_schema(IdentityType::Missing), + ClickHouseEngine::MergeTree, + ) + .unwrap_err(); assert_eq!(err.kind(), ErrorKind::SourceSchemaError); } #[test] fn validate_replica_identity_for_clickhouse_rejects_partial_primary_key() { + let err = validate_replica_identity_for_clickhouse( + &replicated_schema_with_partial_primary_key(), + ClickHouseEngine::MergeTree, + ) + .unwrap_err(); + assert_eq!(err.kind(), ErrorKind::SourceSchemaError); + assert!(err.to_string().contains("tenant_id")); + } + + #[test] + fn validate_replica_identity_for_clickhouse_rejects_pkless_schema_under_replacing_merge_tree() { + // --- GIVEN: a PK-less schema and engine = ReplacingMergeTree --- + let table_schema = Arc::new(TableSchema::new( + TableId::new(2), + TableName::new("public".to_owned(), "events".to_owned()), + vec![ColumnSchema::new("value".to_owned(), Type::TEXT, -1, 1, None, true)], + )); + let replication_mask = ReplicationMask::all(&table_schema); + let identity_mask = IdentityMask::from_bytes(vec![1]); + let schema = + ReplicatedTableSchema::from_masks(table_schema, replication_mask, identity_mask); + + // --- WHEN: validating under ReplacingMergeTree --- let err = - validate_replica_identity_for_clickhouse(&replicated_schema_with_partial_primary_key()) + validate_replica_identity_for_clickhouse(&schema, ClickHouseEngine::ReplacingMergeTree) .unwrap_err(); + + // --- THEN: rejected with SourceSchemaError --- assert_eq!(err.kind(), ErrorKind::SourceSchemaError); - assert!(err.to_string().contains("tenant_id")); + } + + #[test] + fn ensure_engine_supported_rejects_replacing_merge_tree_on_old_server() { + // --- GIVEN: server below the ReplacingMergeTree minimum --- + let err = + ensure_engine_supported(ClickHouseEngine::ReplacingMergeTree, (23, 4)).unwrap_err(); + // --- THEN: surfaced as a config error --- + assert_eq!(err.kind(), ErrorKind::ConfigError); + } + + #[test] + fn ensure_engine_supported_accepts_merge_tree_on_any_server() { + ensure_engine_supported(ClickHouseEngine::MergeTree, (20, 0)).unwrap(); + } + + #[test] + fn ensure_engine_supported_accepts_replacing_merge_tree_on_supported_server() { + ensure_engine_supported(ClickHouseEngine::ReplacingMergeTree, (23, 5)).unwrap(); + ensure_engine_supported(ClickHouseEngine::ReplacingMergeTree, (24, 1)).unwrap(); + } + + /// Schema with composite PK `(tenant_id, id)` plus a non-PK `value` + /// column. Used by the PK-ALTER-guard tests. + fn replicated_schema_for_pk_alters() -> ReplicatedTableSchema { + let table_schema = Arc::new(TableSchema::new( + TableId::new(7), + TableName::new("public".to_owned(), "replacing_merge_tree_alter".to_owned()), + vec![ + ColumnSchema::new("tenant_id".to_owned(), Type::INT4, -1, 1, Some(1), false), + ColumnSchema::new("id".to_owned(), Type::INT4, -1, 2, Some(2), false), + ColumnSchema::new("value".to_owned(), Type::TEXT, -1, 3, None, true), + ], + )); + let replication_mask = ReplicationMask::all(&table_schema); + let identity_mask = IdentityMask::from_bytes(vec![1, 1, 0]); + ReplicatedTableSchema::from_masks(table_schema, replication_mask, identity_mask) + } + + #[test] + fn reject_pk_alters_under_replacing_merge_tree_allows_non_pk_drop() { + // --- GIVEN: a diff that drops the non-PK `value` column --- + let schema = replicated_schema_for_pk_alters(); + let diff = SchemaDiff { + columns_to_add: Vec::new(), + columns_to_remove: vec![ColumnSchema::new( + "value".to_owned(), + Type::TEXT, + -1, + 3, + None, + true, + )], + columns_to_rename: Vec::new(), + }; + // --- WHEN/THEN: guard passes --- + reject_pk_alters_under_replacing_merge_tree( + "public_replacing_merge_tree__alter", + &diff, + &schema, + ) + .unwrap(); + } + + #[test] + fn reject_pk_alters_under_replacing_merge_tree_rejects_pk_drop() { + // --- GIVEN: a diff that drops the PK column `tenant_id` --- + let schema = replicated_schema_for_pk_alters(); + let diff = SchemaDiff { + columns_to_add: Vec::new(), + columns_to_remove: vec![ColumnSchema::new( + "tenant_id".to_owned(), + Type::INT4, + -1, + 1, + Some(1), + false, + )], + columns_to_rename: Vec::new(), + }; + // --- WHEN: validating --- + let err = reject_pk_alters_under_replacing_merge_tree( + "public_replacing_merge_tree__alter", + &diff, + &schema, + ) + .unwrap_err(); + // --- THEN: rejected with SourceSchemaError naming the column --- + assert_eq!(err.kind(), ErrorKind::SourceSchemaError); + assert!(err.to_string().contains("tenant_id"), "error must name the PK column: {err}"); + } + + #[test] + fn reject_pk_alters_under_replacing_merge_tree_allows_non_pk_rename() { + // --- GIVEN: a rename of the non-PK `value` column --- + let schema = replicated_schema_for_pk_alters(); + let diff = SchemaDiff { + columns_to_add: Vec::new(), + columns_to_remove: Vec::new(), + columns_to_rename: vec![etl::types::ColumnRename { + old_name: "value".to_owned(), + new_name: "payload".to_owned(), + ordinal_position: 3, + }], + }; + // --- WHEN/THEN: guard passes --- + reject_pk_alters_under_replacing_merge_tree( + "public_replacing_merge_tree__alter", + &diff, + &schema, + ) + .unwrap(); + } + + #[test] + fn reject_pk_alters_under_replacing_merge_tree_rejects_pk_rename() { + // --- GIVEN: a rename of the PK column `id` --- + let schema = replicated_schema_for_pk_alters(); + let diff = SchemaDiff { + columns_to_add: Vec::new(), + columns_to_remove: Vec::new(), + columns_to_rename: vec![etl::types::ColumnRename { + old_name: "id".to_owned(), + new_name: "row_id".to_owned(), + ordinal_position: 2, + }], + }; + // --- WHEN: validating --- + let err = reject_pk_alters_under_replacing_merge_tree( + "public_replacing_merge_tree__alter", + &diff, + &schema, + ) + .unwrap_err(); + // --- THEN: rejected with SourceSchemaError naming the rename --- + assert_eq!(err.kind(), ErrorKind::SourceSchemaError); + assert!( + err.to_string().contains("'id'") && err.to_string().contains("'row_id'"), + "error must name old + new names: {err}" + ); } #[test] diff --git a/crates/etl-destinations/src/clickhouse/encoding.rs b/crates/etl-destinations/src/clickhouse/encoding.rs index d238e1ba1..c92ed2d75 100644 --- a/crates/etl-destinations/src/clickhouse/encoding.rs +++ b/crates/etl-destinations/src/clickhouse/encoding.rs @@ -22,9 +22,15 @@ pub(crate) enum ClickHouseValue { Int16(i16), Int32(i32), Int64(i64), + /// Unsigned 8-bit integer, used for the ReplacingMergeTree `_etl_deleted` + /// tombstone. + UInt8(u8), UInt32(u32), - /// Unsigned 64-bit integer, used for CDC LSN metadata. + /// Unsigned 64-bit integer, used for the MergeTree `cdc_lsn` column. UInt64(u64), + /// Unsigned 128-bit integer, used for the ReplacingMergeTree `_etl_version` + /// column (the packed `EventSequenceKey`). + UInt128(u128), Float32(f32), Float64(f64), /// TEXT, NUMERIC (string), TIME (string), JSON, BYTEA (hex-encoded) @@ -214,8 +220,10 @@ pub(crate) fn rb_encode_value(val: ClickHouseValue, buf: &mut Vec) -> EtlRes ClickHouseValue::Int16(v) => buf.extend_from_slice(&v.to_le_bytes()), ClickHouseValue::Int32(v) => buf.extend_from_slice(&v.to_le_bytes()), ClickHouseValue::Int64(v) => buf.extend_from_slice(&v.to_le_bytes()), + ClickHouseValue::UInt8(v) => buf.push(v), ClickHouseValue::UInt32(v) => buf.extend_from_slice(&v.to_le_bytes()), ClickHouseValue::UInt64(v) => buf.extend_from_slice(&v.to_le_bytes()), + ClickHouseValue::UInt128(v) => buf.extend_from_slice(&v.to_le_bytes()), ClickHouseValue::Float32(v) => buf.extend_from_slice(&v.to_le_bytes()), ClickHouseValue::Float64(v) => buf.extend_from_slice(&v.to_le_bytes()), ClickHouseValue::String(s) => { diff --git a/crates/etl-destinations/src/clickhouse/schema.rs b/crates/etl-destinations/src/clickhouse/schema.rs index 020ee3d3c..7a5f53ff8 100644 --- a/crates/etl-destinations/src/clickhouse/schema.rs +++ b/crates/etl-destinations/src/clickhouse/schema.rs @@ -1,9 +1,24 @@ -use etl::types::{ColumnSchema, Type, is_array_type}; +use etl::{ + error::{ErrorKind, EtlResult}, + etl_error, + types::{ColumnSchema, Type, is_array_type}, +}; +use etl_config::shared::ClickHouseEngine; -/// Name of the CDC operation metadata column appended to ClickHouse tables. +/// (For MergeTree engine) CDC operation column. pub(crate) const CDC_OPERATION_COLUMN_NAME: &str = "cdc_operation"; -/// Name of the CDC LSN metadata column appended to ClickHouse tables. +/// (For MergeTree engine) CDC LSN column (commit_lsn). pub(crate) const CDC_LSN_COLUMN_NAME: &str = "cdc_lsn"; +/// (For ReplacingMergeTree engine) version column. Holds the packed +/// `EventSequenceKey` (commit_lsn in the high 64 bits, tx_ordinal in the +/// low 64 bits) as a UInt128, giving ReplacingMergeTree a total order across +/// all events for tie-breaking under `FINAL`. +pub(crate) const ETL_VERSION_COLUMN_NAME: &str = "_etl_version"; +/// (For ReplacingMergeTree engine) tombstone column. +pub(crate) const ETL_DELETED_COLUMN_NAME: &str = "_etl_deleted"; +/// Suffix for the auto-generated current-state view over ReplacingMergeTree +/// tables. +pub(crate) const CURRENT_VIEW_SUFFIX: &str = "__current"; /// Returns the base ClickHouse type string for a Postgres scalar type. /// @@ -80,11 +95,35 @@ pub(super) fn clickhouse_column_type(col: &ColumnSchema, force_nullable: bool) - } } -/// Generates a `CREATE TABLE IF NOT EXISTS` DDL for the given columns. -/// -/// Appends `cdc_operation String` and `cdc_lsn UInt64` as trailing non-nullable -/// columns. Uses `MergeTree()` with `ORDER BY tuple()`. -pub(super) fn build_create_table_sql<'a, I>(table_name: &str, column_schemas: I) -> String +/// Trailing CDC column names appended to each replicated row, by engine. +pub(super) fn trailing_cdc_column_names(engine: ClickHouseEngine) -> &'static [&'static str] { + match engine { + ClickHouseEngine::MergeTree => &[CDC_OPERATION_COLUMN_NAME, CDC_LSN_COLUMN_NAME], + ClickHouseEngine::ReplacingMergeTree => &[ETL_VERSION_COLUMN_NAME, ETL_DELETED_COLUMN_NAME], + } +} + +/// Dispatches `CREATE TABLE IF NOT EXISTS` DDL by engine. +pub(super) fn create_table_sql<'a, I>( + engine: ClickHouseEngine, + table_name: &str, + column_schemas: I, +) -> EtlResult +where + I: IntoIterator, + I::IntoIter: ExactSizeIterator, +{ + match engine { + ClickHouseEngine::MergeTree => Ok(create_merge_tree_sql(table_name, column_schemas)), + ClickHouseEngine::ReplacingMergeTree => { + create_replacing_merge_tree_sql(table_name, column_schemas) + } + } +} + +/// `MergeTree` DDL: appends `cdc_operation String` and `cdc_lsn UInt64`, +/// `ORDER BY tuple()`. +pub(super) fn create_merge_tree_sql<'a, I>(table_name: &str, column_schemas: I) -> String where I: IntoIterator, I::IntoIter: ExactSizeIterator, @@ -97,7 +136,6 @@ where cols.push(format!(" {} {}", quote_identifier(&col.name), col_type)); } - // CDC columns — always non-nullable cols.push(format!(" {} String", quote_identifier(CDC_OPERATION_COLUMN_NAME))); cols.push(format!(" {} UInt64", quote_identifier(CDC_LSN_COLUMN_NAME))); @@ -109,6 +147,104 @@ where ) } +/// Emits `CREATE TABLE ... ENGINE = ReplacingMergeTree(_etl_version, +/// _etl_deleted) ORDER BY ()`, with `` taken from the +/// source primary key in `primary_key_ordinal_position` order. ClickHouse +/// uses that `ORDER BY` as the sort + dedup key, so it must match the +/// source PK exactly. Two trailing columns are appended after the user +/// columns: `_etl_version UInt128` (packed `EventSequenceKey`) and +/// `_etl_deleted UInt8` (tombstone). +/// +/// Errors when the source schema has no PK columns. +pub(super) fn create_replacing_merge_tree_sql<'a, I>( + table_name: &str, + column_schemas: I, +) -> EtlResult +where + I: IntoIterator, + I::IntoIter: ExactSizeIterator, +{ + let columns: Vec<&ColumnSchema> = column_schemas.into_iter().collect(); + let pk_columns = primary_key_columns_sorted(table_name, &columns)?; + + let mut col_defs: Vec = columns + .iter() + .map(|col| { + format!(" {} {}", quote_identifier(&col.name), clickhouse_column_type(col, false)) + }) + .collect(); + col_defs.push(format!(" {} UInt128", quote_identifier(ETL_VERSION_COLUMN_NAME))); + col_defs.push(format!(" {} UInt8", quote_identifier(ETL_DELETED_COLUMN_NAME))); + + let order_by = + pk_columns.iter().map(|c| quote_identifier(&c.name)).collect::>().join(", "); + + Ok(format!( + "CREATE TABLE IF NOT EXISTS {quoted_table_name} (\n{col_defs}\n) ENGINE = \ + ReplacingMergeTree({lsn}, {del})\nORDER BY ({order_by})", + quoted_table_name = quote_identifier(table_name), + col_defs = col_defs.join(",\n"), + lsn = quote_identifier(ETL_VERSION_COLUMN_NAME), + del = quote_identifier(ETL_DELETED_COLUMN_NAME), + )) +} + +/// Returns the source primary-key columns sorted by +/// `primary_key_ordinal_position`. Errors with `SourceSchemaError` when the +/// schema has no PK columns (ReplacingMergeTree cannot be created without an +/// `ORDER BY`). +fn primary_key_columns_sorted<'a>( + table_name: &str, + columns: &[&'a ColumnSchema], +) -> EtlResult> { + let mut pk_columns: Vec<&ColumnSchema> = + columns.iter().copied().filter(|c| c.primary_key_ordinal_position.is_some()).collect(); + + if pk_columns.is_empty() { + return Err(etl_error!( + ErrorKind::SourceSchemaError, + "ClickHouse ReplacingMergeTree requires a primary key", + format!( + "Table '{table_name}' has no primary-key columns; set `engine: merge_tree` or \ + define a PK on the source table." + ) + )); + } + + pk_columns.sort_by_key(|c| c.primary_key_ordinal_position); + Ok(pk_columns) +} + +/// `CREATE VIEW IF NOT EXISTS "
__current"` for an ReplacingMergeTree +/// table. +/// +/// Selects only user columns (drops `_etl_version` and `_etl_deleted`), reads +/// via `FINAL`, and filters tombstones. +pub(super) fn create_current_view_sql<'a, I>(table_name: &str, column_schemas: I) -> String +where + I: IntoIterator, +{ + let select_cols = column_schemas + .into_iter() + .map(|c| quote_identifier(&c.name)) + .collect::>() + .join(", "); + let view_name = format!("{table_name}{CURRENT_VIEW_SUFFIX}"); + format!( + "CREATE VIEW IF NOT EXISTS {view} AS\nSELECT {select_cols}\nFROM {table} FINAL\nWHERE \ + {deleted} = 0", + view = quote_identifier(&view_name), + table = quote_identifier(table_name), + deleted = quote_identifier(ETL_DELETED_COLUMN_NAME), + ) +} + +/// `DROP VIEW IF EXISTS "
__current"` for an ReplacingMergeTree table. +pub(super) fn drop_current_view_sql(table_name: &str) -> String { + let view_name = format!("{table_name}{CURRENT_VIEW_SUFFIX}"); + format!("DROP VIEW IF EXISTS {}", quote_identifier(&view_name)) +} + #[cfg(test)] mod tests { use super::*; @@ -120,7 +256,7 @@ mod tests { } #[test] - fn build_create_table_sql_quotes_identifiers() { + fn create_merge_tree_sql_quotes_identifiers() { let schemas = vec![ColumnSchema { name: "id\"value".to_owned(), typ: Type::INT4, @@ -131,7 +267,7 @@ mod tests { }]; // Pre-encoded table name with embedded quotes to verify the SQL // builder quotes/escapes the identifier itself. - let sql = build_create_table_sql("sche\"ma_ta\"ble", &schemas); + let sql = create_merge_tree_sql("sche\"ma_ta\"ble", &schemas); assert!( sql.contains("CREATE TABLE IF NOT EXISTS \"sche\"\"ma_ta\"\"ble\""), @@ -185,7 +321,7 @@ mod tests { } #[test] - fn build_create_table_sql_nullable() { + fn create_merge_tree_sql_nullable() { let schemas = vec![ ColumnSchema { name: "id".to_owned(), @@ -204,13 +340,13 @@ mod tests { nullable: true, }, ]; - let sql = build_create_table_sql("public_users", &schemas); + let sql = create_merge_tree_sql("public_users", &schemas); assert!(sql.contains("\"id\" Int32"), "id should be non-nullable Int32"); assert!(sql.contains("\"name\" Nullable(String)"), "name should be Nullable(String)"); } #[test] - fn build_create_table_sql_cdc_columns() { + fn create_merge_tree_sql_cdc_columns() { let schemas = vec![ColumnSchema { name: "id".to_owned(), typ: Type::INT4, @@ -219,7 +355,7 @@ mod tests { primary_key_ordinal_position: Some(1), nullable: false, }]; - let sql = build_create_table_sql("public_t", &schemas); + let sql = create_merge_tree_sql("public_t", &schemas); assert!(sql.contains("\"cdc_operation\" String"), "cdc_operation should be non-nullable"); assert!(sql.contains("\"cdc_lsn\" UInt64"), "cdc_lsn should be non-nullable UInt64"); assert!(sql.contains("ENGINE = MergeTree()")); @@ -227,7 +363,7 @@ mod tests { } #[test] - fn build_create_table_sql_array_columns() { + fn create_merge_tree_sql_array_columns() { let schemas = vec![ColumnSchema { name: "tags".to_owned(), typ: Type::TEXT_ARRAY, @@ -236,10 +372,167 @@ mod tests { primary_key_ordinal_position: None, nullable: false, }]; - let sql = build_create_table_sql("public_t", &schemas); + let sql = create_merge_tree_sql("public_t", &schemas); assert!( sql.contains("\"tags\" Array(Nullable(String))"), "array columns should always be Array(Nullable(T))" ); } + + #[test] + fn create_replacing_merge_tree_sql_single_pk() { + // --- GIVEN: single-column PK with a nullable non-PK column --- + let schemas = vec![ + ColumnSchema { + name: "id".to_owned(), + typ: Type::INT4, + modifier: -1, + ordinal_position: 1, + primary_key_ordinal_position: Some(1), + nullable: false, + }, + ColumnSchema { + name: "name".to_owned(), + typ: Type::TEXT, + modifier: -1, + ordinal_position: 2, + primary_key_ordinal_position: None, + nullable: true, + }, + ]; + // --- WHEN: build the ReplacingMergeTree DDL --- + let sql = create_replacing_merge_tree_sql("public_users", &schemas).unwrap(); + // --- THEN: trailing etl columns, engine, and ORDER BY are correct --- + assert!(sql.contains("\"id\" Int32")); + assert!(sql.contains("\"name\" Nullable(String)")); + assert!(sql.contains("\"_etl_version\" UInt128")); + assert!(sql.contains("\"_etl_deleted\" UInt8")); + assert!(sql.contains("ENGINE = ReplacingMergeTree(\"_etl_version\", \"_etl_deleted\")")); + assert!(sql.contains("ORDER BY (\"id\")")); + } + + #[test] + fn create_replacing_merge_tree_sql_composite_pk_orders_by_ordinal() { + // --- GIVEN: composite PK whose ordinal order differs from table order --- + let schemas = vec![ + ColumnSchema { + name: "id".to_owned(), + typ: Type::INT4, + modifier: -1, + ordinal_position: 1, + primary_key_ordinal_position: Some(2), + nullable: false, + }, + ColumnSchema { + name: "name".to_owned(), + typ: Type::TEXT, + modifier: -1, + ordinal_position: 2, + primary_key_ordinal_position: None, + nullable: true, + }, + ColumnSchema { + name: "tenant_id".to_owned(), + typ: Type::INT4, + modifier: -1, + ordinal_position: 3, + primary_key_ordinal_position: Some(1), + nullable: false, + }, + ]; + // --- WHEN: build the ReplacingMergeTree DDL --- + let sql = create_replacing_merge_tree_sql("public_users", &schemas).unwrap(); + // --- THEN: ORDER BY follows PK ordinal, not table ordinal --- + assert!( + sql.contains("ORDER BY (\"tenant_id\", \"id\")"), + "ORDER BY must follow PK ordinal: {sql}" + ); + } + + #[test] + fn create_replacing_merge_tree_sql_rejects_pkless_schema() { + // --- GIVEN: schema with no PK columns --- + let schemas = vec![ColumnSchema { + name: "value".to_owned(), + typ: Type::TEXT, + modifier: -1, + ordinal_position: 1, + primary_key_ordinal_position: None, + nullable: true, + }]; + // --- WHEN: build the ReplacingMergeTree DDL --- + let err = create_replacing_merge_tree_sql("public_events", &schemas).unwrap_err(); + // --- THEN: builder rejects with SourceSchemaError --- + assert_eq!(err.kind(), ErrorKind::SourceSchemaError); + } + + #[test] + fn create_table_sql_dispatches_on_engine() { + // --- GIVEN: a schema with a single PK column --- + let schemas = vec![ColumnSchema { + name: "id".to_owned(), + typ: Type::INT4, + modifier: -1, + ordinal_position: 1, + primary_key_ordinal_position: Some(1), + nullable: false, + }]; + // --- WHEN/THEN: dispatcher selects the matching engine branch --- + let merge_tree = + create_table_sql(ClickHouseEngine::MergeTree, "public_t", &schemas).unwrap(); + assert!(merge_tree.contains("ENGINE = MergeTree()")); + let replacing_merge_tree = + create_table_sql(ClickHouseEngine::ReplacingMergeTree, "public_t", &schemas).unwrap(); + assert!(replacing_merge_tree.contains("ENGINE = ReplacingMergeTree")); + } + + #[test] + fn create_current_view_sql_selects_user_columns_only() { + // --- GIVEN: a two-column schema --- + let schemas = vec![ + ColumnSchema { + name: "id".to_owned(), + typ: Type::INT4, + modifier: -1, + ordinal_position: 1, + primary_key_ordinal_position: Some(1), + nullable: false, + }, + ColumnSchema { + name: "name".to_owned(), + typ: Type::TEXT, + modifier: -1, + ordinal_position: 2, + primary_key_ordinal_position: None, + nullable: true, + }, + ]; + // --- WHEN: build the current-state view DDL --- + let sql = create_current_view_sql("public_users", &schemas); + // --- THEN: __current suffix, FINAL read, tombstone filter, no etl cols --- + assert!(sql.contains("CREATE VIEW IF NOT EXISTS \"public_users__current\"")); + assert!(sql.contains("SELECT \"id\", \"name\"")); + assert!(sql.contains("FROM \"public_users\" FINAL")); + assert!(sql.contains("WHERE \"_etl_deleted\" = 0")); + assert!(!sql.contains("_etl_version"), "view must not expose _etl_version: {sql}"); + } + + #[test] + fn drop_current_view_sql_quotes_view_name() { + let sql = drop_current_view_sql("public_us\"ers"); + + assert_eq!(sql, "DROP VIEW IF EXISTS \"public_us\"\"ers__current\""); + } + + #[test] + fn trailing_cdc_column_names_by_engine() { + assert_eq!( + trailing_cdc_column_names(ClickHouseEngine::MergeTree), + &[CDC_OPERATION_COLUMN_NAME, CDC_LSN_COLUMN_NAME] + ); + assert_eq!( + trailing_cdc_column_names(ClickHouseEngine::ReplacingMergeTree), + &[ETL_VERSION_COLUMN_NAME, ETL_DELETED_COLUMN_NAME] + ); + } } diff --git a/crates/etl-destinations/src/clickhouse/test_utils.rs b/crates/etl-destinations/src/clickhouse/test_utils.rs index def74209d..aa14e9a1a 100644 --- a/crates/etl-destinations/src/clickhouse/test_utils.rs +++ b/crates/etl-destinations/src/clickhouse/test_utils.rs @@ -113,20 +113,36 @@ impl ClickHouseTestDatabase { /// Builds a [`ClickHouseDestination`] scoped to this test database with /// default inserter config (100 MiB per INSERT -- large enough that tests - /// never hit an intermediate flush). - pub fn build_destination(&self, store: S) -> ClickHouseDestination + /// never hit an intermediate flush). Validates engine support eagerly so + /// tests fail fast on engine/version mismatch. + pub async fn build_destination(&self, store: S) -> ClickHouseDestination + where + S: StateStore + SchemaStore + Send + Sync, + { + self.build_destination_with_engine(store, etl_config::shared::ClickHouseEngine::default()) + .await + } + + /// Builds a [`ClickHouseDestination`] for the given engine. + pub async fn build_destination_with_engine( + &self, + store: S, + engine: etl_config::shared::ClickHouseEngine, + ) -> ClickHouseDestination where S: StateStore + SchemaStore + Send + Sync, { self.build_destination_with_config( store, - ClickHouseInserterConfig { max_bytes_per_insert: 100 * 1024 * 1024 }, + ClickHouseInserterConfig { max_bytes_per_insert: 100 * 1024 * 1024, engine }, ) + .await } /// Builds a [`ClickHouseDestination`] scoped to this test database with - /// a caller-supplied [`ClickHouseInserterConfig`]. - pub fn build_destination_with_config( + /// a caller-supplied [`ClickHouseInserterConfig`]. Validates engine support + /// eagerly so tests fail fast on engine/version mismatch. + pub async fn build_destination_with_config( &self, store: S, config: ClickHouseInserterConfig, @@ -134,7 +150,7 @@ impl ClickHouseTestDatabase { where S: StateStore + SchemaStore + Send + Sync, { - ClickHouseDestination::new( + let destination = ClickHouseDestination::new( self.url.clone(), &self.user, self.password.clone(), @@ -143,7 +159,12 @@ impl ClickHouseTestDatabase { ClickHouseClientConfig::default(), store, ) - .expect("Failed to create ClickHouseDestination for test") + .expect("Failed to create ClickHouseDestination for test"); + destination + .validate_engine_support() + .await + .expect("ClickHouse engine support check failed in test setup"); + destination } /// Fetches all rows from a ClickHouse table using the given SQL query. @@ -164,13 +185,14 @@ impl ClickHouseTestDatabase { } /// Returns the column names of a ClickHouse table in position order, - /// excluding the CDC columns (`cdc_operation`, `cdc_lsn`). + /// excluding both engines' trailing CDC columns. pub async fn column_names(&self, table_name: &str) -> Vec { self.column_types(table_name).await.into_iter().map(|(name, _)| name).collect() } /// Returns the column names and ClickHouse type strings in position order, - /// excluding the CDC columns (`cdc_operation`, `cdc_lsn`). + /// excluding both engines' trailing CDC columns (`cdc_operation`, + /// `cdc_lsn`, `_etl_version`, `_etl_deleted`). pub async fn column_types(&self, table_name: &str) -> Vec<(String, String)> { #[derive(clickhouse::Row, serde::Deserialize)] struct Col { @@ -180,7 +202,8 @@ impl ClickHouseTestDatabase { self.db_client .query( "SELECT name, type AS type_name FROM system.columns WHERE database = ? AND table \ - = ? AND name NOT IN ('cdc_operation', 'cdc_lsn') ORDER BY position", + = ? AND name NOT IN ('cdc_operation', 'cdc_lsn', '_etl_version', '_etl_deleted') \ + ORDER BY position", ) .bind(&self.database) .bind(table_name) diff --git a/crates/etl-destinations/tests/clickhouse/mod.rs b/crates/etl-destinations/tests/clickhouse/mod.rs index eab2e1f41..f5fde1c39 100644 --- a/crates/etl-destinations/tests/clickhouse/mod.rs +++ b/crates/etl-destinations/tests/clickhouse/mod.rs @@ -1 +1,3 @@ mod pipeline; +mod pipeline_merge_tree; +mod pipeline_replacing_merge_tree; diff --git a/crates/etl-destinations/tests/clickhouse/pipeline.rs b/crates/etl-destinations/tests/clickhouse/pipeline.rs index 620e852dd..2084ba66f 100644 --- a/crates/etl-destinations/tests/clickhouse/pipeline.rs +++ b/crates/etl-destinations/tests/clickhouse/pipeline.rs @@ -1,5 +1,3 @@ -use std::sync::Once; - use etl::{ state::table::TableReplicationPhaseType, store::state::StateStore, @@ -11,6 +9,7 @@ use etl::{ }, types::{EventType, PipelineId}, }; +use etl_config::shared::ClickHouseEngine; use etl_destinations::clickhouse::{ ClickHouseClientConfig, ClickHouseInserterConfig, client::ClickHouseClient, @@ -23,78 +22,37 @@ use etl_telemetry::tracing::init_test_tracing; use rand::random; use url::Url; -use crate::support::clickhouse::{AllTypesRow, BoundaryValuesRow, DateBoundariesRow}; - -/// Ensures the rustls crypto provider is only installed once across all tests. -static INIT_CRYPTO: Once = Once::new(); - -fn install_crypto_provider() { - INIT_CRYPTO.call_once(|| { - rustls::crypto::aws_lc_rs::default_provider() - .install_default() - .expect("failed to install default crypto provider"); - }); -} +use crate::support::clickhouse::{ + AllTypesRow, BoundaryValuesRow, DateBoundariesRow, current_state_query, install_crypto_provider, +}; -/// SELECT query that fetches all verified columns from the ClickHouse table. -/// -/// `uuid_col` is projected via `toString()` because the ClickHouse UUID -/// RowBinary wire format does not directly map to a Rust `String`; `toString()` -/// gives us the canonical `xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx` string form. -/// -/// All other columns are read with their native ClickHouse types: -/// - `Date` -> u16 (days since 1970-01-01) -/// - `DateTime64(6)` -> i64 (microseconds since epoch) -/// - `Array(Nullable(T))` -> `Vec>` -const ALL_TYPES_SELECT: &str = concat!( - "SELECT ", +/// User-column projection for the all-types test, with `uuid_col` rendered +/// as a canonical lowercase UUID string via `toString()`. +const ALL_TYPES_PROJECTION: &str = concat!( "id, smallint_col, integer_col, bigint_col, real_col, double_col, ", "numeric_col, boolean_col, text_col, varchar_col, ", "date_col, timestamp_col, timestamptz_col, time_col, interval_col, ", "jsonb_col, json_col, integer_array_col, text_array_col, ", "bytea_col, inet_col, cidr_col, macaddr_col, ", - "toString(uuid_col) AS uuid_col, ", - "cdc_operation ", - "FROM \"test_all__types__encoding\" ", - "ORDER BY id", + "toString(uuid_col) AS uuid_col", ); -/// A row read back from the ClickHouse `update_flow` test table. +const ALL_TYPES_TABLE: &str = "test_all__types__encoding"; + +/// A current-state row with `id` + `value`, shared across streaming tests. #[derive(clickhouse::Row, serde::Deserialize, Debug)] -struct UpdateFlowRow { +struct IdValueRow { id: i64, value: String, - cdc_operation: String, - cdc_lsn: u64, } -/// SELECT query used to verify the `update_flow` streaming test. -const UPDATE_FLOW_SELECT: &str = concat!( - "SELECT id, value, cdc_operation, cdc_lsn ", - "FROM \"test_update__flow\" ", - "ORDER BY id, cdc_lsn", -); - -/// SELECT query used to verify the `delete_flow` streaming test. -const DELETE_FLOW_SELECT: &str = concat!( - "SELECT id, value, cdc_operation, cdc_lsn ", - "FROM \"test_delete__flow\" ", - "ORDER BY id, cdc_lsn", -); - -/// SELECT query used to verify the `restart_flow` test. -const RESTART_FLOW_SELECT: &str = concat!( - "SELECT id, value, cdc_operation, cdc_lsn ", - "FROM \"test_restart__flow\" ", - "ORDER BY id, cdc_lsn", -); - -/// SELECT query used to verify the `truncate_flow` test. -const TRUNCATE_FLOW_SELECT: &str = concat!( - "SELECT id, value, cdc_operation, cdc_lsn ", - "FROM \"test_truncate__flow\" ", - "ORDER BY id, cdc_lsn", -); +/// Projection + table name + ORDER BY that drive `current_state_query` for +/// each streaming test. +const ID_VALUE_PROJECTION: &str = "id, value"; +const UPDATE_FLOW_TABLE: &str = "test_update__flow"; +const DELETE_FLOW_TABLE: &str = "test_delete__flow"; +const RESTART_FLOW_TABLE: &str = "test_restart__flow"; +const TRUNCATE_FLOW_TABLE: &str = "test_truncate__flow"; /// Days from 1970-01-01 to 2024-01-15 (used to verify the `date_col` /// round-trip). @@ -139,7 +97,16 @@ const TS_2024_01_15_12_00_US: i64 = 1_705_320_000_000_000; /// element bytes as subsequent column data, failing with "Cannot read all data" /// at row 2. #[tokio::test(flavor = "multi_thread")] -async fn all_types_table_copy() { +async fn all_types_table_copy_merge_tree() { + all_types_table_copy_inner(ClickHouseEngine::MergeTree).await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn all_types_table_copy_replacing_merge_tree() { + all_types_table_copy_inner(ClickHouseEngine::ReplacingMergeTree).await; +} + +async fn all_types_table_copy_inner(engine: ClickHouseEngine) { init_test_tracing(); install_crypto_provider(); @@ -256,7 +223,7 @@ async fn all_types_table_copy() { let clickhouse_db = setup_clickhouse_database().await; let store = NotifyingStore::new(); let pipeline_id: PipelineId = random(); - let destination = clickhouse_db.build_destination(store.clone()); + let destination = clickhouse_db.build_destination_with_engine(store.clone(), engine).await; let table_ready = store.notify_on_table_state_type(table_id, TableReplicationPhaseType::Ready).await; @@ -274,7 +241,8 @@ async fn all_types_table_copy() { pipeline.shutdown_and_wait().await.unwrap(); // --- THEN: every column round-trips correctly --- - let rows: Vec = clickhouse_db.query(ALL_TYPES_SELECT).await; + let query = current_state_query(engine, ALL_TYPES_TABLE, ALL_TYPES_PROJECTION, &["id"], "id"); + let rows: Vec = clickhouse_db.query(&query).await; assert_eq!(rows.len(), 2, "expected 2 rows in ClickHouse"); @@ -299,7 +267,6 @@ async fn all_types_table_copy() { assert_eq!(r1.cidr_col, "192.168.0.0/16"); assert_eq!(r1.macaddr_col, "aa:bb:cc:dd:ee:ff"); assert_eq!(r1.uuid_col.to_lowercase(), "f47ac10b-58cc-4372-a567-0e02b2c3d479"); - assert_eq!(r1.cdc_operation, "INSERT"); // Empty arrays -- the regression case that accidentally worked before the fix. assert_eq!( r1.integer_array_col, @@ -322,7 +289,6 @@ async fn all_types_table_copy() { assert_eq!(r2.numeric_col, "-99999.99"); assert_eq!(r2.bytea_col, "cafebabe"); assert_eq!(r2.uuid_col.to_lowercase(), "a1b2c3d4-e5f6-7890-abcd-ef1234567890"); - assert_eq!(r2.cdc_operation, "INSERT"); // Non-empty arrays -- the regression case that triggered the bug before the // fix. assert_eq!( @@ -338,24 +304,18 @@ async fn all_types_table_copy() { } /// Tests that UPDATE events are streamed to ClickHouse after initial table -/// copy. -/// -/// # GIVEN -/// -/// A Postgres table with a single row (`id=1, value='before'`). -/// -/// # WHEN -/// -/// The pipeline copies the row, then an `UPDATE ... SET value = 'after'` is -/// issued against Postgres. -/// -/// # THEN -/// -/// ClickHouse contains two rows (append-only CDC): -/// - The original `INSERT` from table copy with `cdc_lsn = 0`. -/// - The streamed `UPDATE` with the new value and a positive LSN. +/// copy. Asserts on the current state (latest value per PK). +#[tokio::test(flavor = "multi_thread")] +async fn updates_are_streamed_to_clickhouse_merge_tree() { + updates_are_streamed_to_clickhouse_inner(ClickHouseEngine::MergeTree).await; +} + #[tokio::test(flavor = "multi_thread")] -async fn updates_are_streamed_to_clickhouse() { +async fn updates_are_streamed_to_clickhouse_replacing_merge_tree() { + updates_are_streamed_to_clickhouse_inner(ClickHouseEngine::ReplacingMergeTree).await; +} + +async fn updates_are_streamed_to_clickhouse_inner(engine: ClickHouseEngine) { init_test_tracing(); install_crypto_provider(); @@ -386,7 +346,9 @@ async fn updates_are_streamed_to_clickhouse() { let clickhouse_db = setup_clickhouse_database().await; let store = NotifyingStore::new(); let pipeline_id: PipelineId = random(); - let destination = TestDestinationWrapper::wrap(clickhouse_db.build_destination(store.clone())); + let destination = TestDestinationWrapper::wrap( + clickhouse_db.build_destination_with_engine(store.clone(), engine).await, + ); let table_ready = store.notify_on_table_state_type(table_id, TableReplicationPhaseType::Ready).await; @@ -414,34 +376,17 @@ async fn updates_are_streamed_to_clickhouse() { event_notify.notified().await; - let rows: Vec = clickhouse_db.query(UPDATE_FLOW_SELECT).await; + let query = current_state_query(engine, UPDATE_FLOW_TABLE, ID_VALUE_PROJECTION, &["id"], "id"); + let rows: Vec = clickhouse_db.query(&query).await; pipeline.shutdown_and_wait().await.unwrap(); - // --- THEN: one INSERT from table copy, one UPDATE from streaming --- - assert_eq!(rows.len(), 2, "expected copied row plus streamed update"); - - let insert_row = &rows[0]; - assert_eq!(insert_row.id, 1); - assert_eq!(insert_row.value, "before"); - assert_eq!(insert_row.cdc_operation, "INSERT"); - assert_eq!(insert_row.cdc_lsn, 0); - - let update_row = &rows[1]; - assert_eq!(update_row.id, 1); - assert_eq!(update_row.value, "after"); - assert_eq!(update_row.cdc_operation, "UPDATE"); - assert!(update_row.cdc_lsn > insert_row.cdc_lsn, "streamed update should have a positive LSN"); + // --- THEN: current state shows the updated value --- + assert_eq!(rows.len(), 1, "expected one current-state row after UPDATE"); + assert_eq!(rows[0].id, 1); + assert_eq!(rows[0].value, "after"); } -const BOUNDARY_VALUES_SELECT: &str = concat!( - "SELECT id, nullable_text, nullable_int, ", - "int_array_col, text_array_col, ", - "cdc_operation ", - "FROM \"test_boundary__values\" ", - "ORDER BY id", -); - /// Tests that edge-case values survive the Postgres -> ClickHouse pipeline /// without data loss or corruption. /// @@ -469,7 +414,16 @@ const BOUNDARY_VALUES_SELECT: &str = concat!( /// - Array elements preserve their position, including interior NULLs. /// - Multi-byte text round-trips byte-for-byte. #[tokio::test(flavor = "multi_thread")] -async fn boundary_values_table_copy() { +async fn boundary_values_table_copy_merge_tree() { + boundary_values_table_copy_inner(ClickHouseEngine::MergeTree).await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn boundary_values_table_copy_replacing_merge_tree() { + boundary_values_table_copy_inner(ClickHouseEngine::ReplacingMergeTree).await; +} + +async fn boundary_values_table_copy_inner(engine: ClickHouseEngine) { init_test_tracing(); install_crypto_provider(); @@ -545,7 +499,7 @@ async fn boundary_values_table_copy() { let clickhouse_db = setup_clickhouse_database().await; let store = NotifyingStore::new(); let pipeline_id: PipelineId = random(); - let destination = clickhouse_db.build_destination(store.clone()); + let destination = clickhouse_db.build_destination_with_engine(store.clone(), engine).await; let table_ready = store.notify_on_table_state_type(table_id, TableReplicationPhaseType::Ready).await; @@ -563,7 +517,14 @@ async fn boundary_values_table_copy() { pipeline.shutdown_and_wait().await.unwrap(); // --- THEN: ClickHouse data matches Postgres exactly --- - let rows: Vec = clickhouse_db.query(BOUNDARY_VALUES_SELECT).await; + let query = current_state_query( + engine, + "test_boundary__values", + "id, nullable_text, nullable_int, int_array_col, text_array_col", + &["id"], + "id", + ); + let rows: Vec = clickhouse_db.query(&query).await; assert_eq!(rows.len(), 4, "expected 4 rows in ClickHouse"); // Row 1: NULL scalars stay NULL, empty arrays stay empty. @@ -626,12 +587,6 @@ const DATE32_MIN_DAYS_FROM_UNIX_EPOCH: i32 = -25567; /// Python: `(date(2299, 12, 31) - date(1970, 1, 1)).days` = 120529. const DATE32_MAX_DAYS_FROM_UNIX_EPOCH: i32 = 120529; -const DATE_BOUNDARIES_SELECT: &str = concat!( - "SELECT id, date_col, cdc_operation ", - "FROM \"test_date__boundaries\" ", - "ORDER BY id", -); - /// Tests that Postgres `date` values outside the Unix epoch (pre-1970 and /// far-future) round-trip through ClickHouse `Date32` as signed day offsets, /// rather than being silently clamped as the previous `Date` (UInt16) mapping @@ -654,7 +609,16 @@ const DATE_BOUNDARIES_SELECT: &str = concat!( /// negative for pre-1970 dates, zero at the epoch, and a large positive value /// for the far-future date. No value is clamped. #[tokio::test(flavor = "multi_thread")] -async fn pre_1970_and_far_future_dates_round_trip() { +async fn pre_1970_and_far_future_dates_round_trip_merge_tree() { + pre_1970_and_far_future_dates_round_trip_inner(ClickHouseEngine::MergeTree).await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn pre_1970_and_far_future_dates_round_trip_replacing_merge_tree() { + pre_1970_and_far_future_dates_round_trip_inner(ClickHouseEngine::ReplacingMergeTree).await; +} + +async fn pre_1970_and_far_future_dates_round_trip_inner(engine: ClickHouseEngine) { init_test_tracing(); install_crypto_provider(); @@ -689,7 +653,7 @@ async fn pre_1970_and_far_future_dates_round_trip() { let clickhouse_db = setup_clickhouse_database().await; let store = NotifyingStore::new(); let pipeline_id: PipelineId = random(); - let destination = clickhouse_db.build_destination(store.clone()); + let destination = clickhouse_db.build_destination_with_engine(store.clone(), engine).await; let table_ready = store.notify_on_table_state_type(table_id, TableReplicationPhaseType::Ready).await; @@ -707,7 +671,8 @@ async fn pre_1970_and_far_future_dates_round_trip() { pipeline.shutdown_and_wait().await.unwrap(); // --- THEN: each date encodes as the expected signed day offset --- - let rows: Vec = clickhouse_db.query(DATE_BOUNDARIES_SELECT).await; + let query = current_state_query(engine, "test_date__boundaries", "id, date_col", &["id"], "id"); + let rows: Vec = clickhouse_db.query(&query).await; assert_eq!(rows.len(), 5, "expected 5 rows in ClickHouse"); assert_eq!( @@ -750,7 +715,16 @@ async fn pre_1970_and_far_future_dates_round_trip() { /// LSN. /// - The `id=1` row has no corresponding `DELETE`. #[tokio::test(flavor = "multi_thread")] -async fn deletes_are_streamed_to_clickhouse() { +async fn deletes_are_streamed_to_clickhouse_merge_tree() { + deletes_are_streamed_to_clickhouse_inner(ClickHouseEngine::MergeTree).await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn deletes_are_streamed_to_clickhouse_replacing_merge_tree() { + deletes_are_streamed_to_clickhouse_inner(ClickHouseEngine::ReplacingMergeTree).await; +} + +async fn deletes_are_streamed_to_clickhouse_inner(engine: ClickHouseEngine) { init_test_tracing(); install_crypto_provider(); @@ -791,7 +765,9 @@ async fn deletes_are_streamed_to_clickhouse() { let clickhouse_db = setup_clickhouse_database().await; let store = NotifyingStore::new(); let pipeline_id: PipelineId = random(); - let destination = TestDestinationWrapper::wrap(clickhouse_db.build_destination(store.clone())); + let destination = TestDestinationWrapper::wrap( + clickhouse_db.build_destination_with_engine(store.clone(), engine).await, + ); let table_ready = store.notify_on_table_state_type(table_id, TableReplicationPhaseType::Ready).await; @@ -816,34 +792,15 @@ async fn deletes_are_streamed_to_clickhouse() { event_notify.notified().await; - let rows: Vec = clickhouse_db.query(DELETE_FLOW_SELECT).await; + let query = current_state_query(engine, DELETE_FLOW_TABLE, ID_VALUE_PROJECTION, &["id"], "id"); + let rows: Vec = clickhouse_db.query(&query).await; pipeline.shutdown_and_wait().await.unwrap(); - // --- THEN: two INSERTs from table copy, one DELETE from streaming --- - - assert_eq!(rows.len(), 3, "expected 2 copied rows plus 1 streamed delete"); - - // Row 1: copied, untouched. - let r = &rows[0]; - assert_eq!(r.id, 1); - assert_eq!(r.value, "keep_me"); - assert_eq!(r.cdc_operation, "INSERT"); - assert_eq!(r.cdc_lsn, 0); - - // Row 2: copied, then deleted. - let r = &rows[1]; - assert_eq!(r.id, 2); - assert_eq!(r.value, "delete_me"); - assert_eq!(r.cdc_operation, "INSERT"); - assert_eq!(r.cdc_lsn, 0); - - // Row 3: the streamed DELETE for id=2, preserving old row data. - let r = &rows[2]; - assert_eq!(r.id, 2, "delete must target the correct row"); - assert_eq!(r.value, "delete_me", "old row data must be preserved in DELETE"); - assert_eq!(r.cdc_operation, "DELETE"); - assert!(r.cdc_lsn > 0, "streamed delete should have a positive LSN"); + // --- THEN: only the surviving row is visible in current state --- + assert_eq!(rows.len(), 1, "deleted row must be absent from current state"); + assert_eq!(rows[0].id, 1); + assert_eq!(rows[0].value, "keep_me"); } /// Tests that a pipeline restart resumes CDC streaming without re-running @@ -868,7 +825,16 @@ async fn deletes_are_streamed_to_clickhouse() { /// - `id=2` from CDC streaming in the second run (`cdc_lsn > 0`). /// No duplicate `id=1` row exists -- table copy must not re-run. #[tokio::test(flavor = "multi_thread")] -async fn pipeline_restart_resumes_streaming() { +async fn pipeline_restart_resumes_streaming_merge_tree() { + pipeline_restart_resumes_streaming_inner(ClickHouseEngine::MergeTree).await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn pipeline_restart_resumes_streaming_replacing_merge_tree() { + pipeline_restart_resumes_streaming_inner(ClickHouseEngine::ReplacingMergeTree).await; +} + +async fn pipeline_restart_resumes_streaming_inner(engine: ClickHouseEngine) { init_test_tracing(); install_crypto_provider(); @@ -898,7 +864,9 @@ async fn pipeline_restart_resumes_streaming() { let clickhouse_db = setup_clickhouse_database().await; let store = NotifyingStore::new(); let pipeline_id: PipelineId = random(); - let destination = TestDestinationWrapper::wrap(clickhouse_db.build_destination(store.clone())); + let destination = TestDestinationWrapper::wrap( + clickhouse_db.build_destination_with_engine(store.clone(), engine).await, + ); let table_ready = store.notify_on_table_state_type(table_id, TableReplicationPhaseType::Ready).await; @@ -916,13 +884,17 @@ async fn pipeline_restart_resumes_streaming() { pipeline.shutdown_and_wait().await.unwrap(); // Verify first run produced exactly one row. - let rows: Vec = clickhouse_db.query(RESTART_FLOW_SELECT).await; + let restart_query = + || current_state_query(engine, RESTART_FLOW_TABLE, ID_VALUE_PROJECTION, &["id"], "id"); + let rows: Vec = clickhouse_db.query(&restart_query()).await; assert_eq!(rows.len(), 1, "first run should copy exactly one row"); assert_eq!(rows[0].id, 1); assert_eq!(rows[0].value, "before_restart"); // --- WHEN: rebuild destination and pipeline, then stream a new insert --- - let destination = TestDestinationWrapper::wrap(clickhouse_db.build_destination(store.clone())); + let destination = TestDestinationWrapper::wrap( + clickhouse_db.build_destination_with_engine(store.clone(), engine).await, + ); let mut pipeline = create_pipeline( &database.config, @@ -946,52 +918,31 @@ async fn pipeline_restart_resumes_streaming() { event_notify.notified().await; - let rows: Vec = clickhouse_db.query(RESTART_FLOW_SELECT).await; + let rows: Vec = clickhouse_db.query(&restart_query()).await; pipeline.shutdown_and_wait().await.unwrap(); - // --- THEN: exactly two rows, no duplicate from re-running table copy --- - assert_eq!( - rows.len(), - 2, - "expected original copied row plus one streamed insert, no duplicates" - ); - - let r = &rows[0]; - assert_eq!(r.id, 1); - assert_eq!(r.value, "before_restart"); - assert_eq!(r.cdc_operation, "INSERT"); - assert_eq!(r.cdc_lsn, 0, "first row should be from table copy"); - - let r = &rows[1]; - assert_eq!(r.id, 2); - assert_eq!(r.value, "after_restart"); - assert_eq!(r.cdc_operation, "INSERT"); - assert!(r.cdc_lsn > 0, "second row should be from CDC streaming after restart"); + // --- THEN: exactly two rows in current state, no duplicate of id=1 --- + assert_eq!(rows.len(), 2, "expected original copied row plus one streamed insert"); + assert_eq!(rows[0].id, 1); + assert_eq!(rows[0].value, "before_restart"); + assert_eq!(rows[1].id, 2); + assert_eq!(rows[1].value, "after_restart"); } /// Tests that TRUNCATE clears the ClickHouse table and that subsequent inserts /// produce a clean slate with only post-truncate data. -/// -/// # GIVEN -/// -/// A Postgres table with two rows (`id=1, value='alpha'` and `id=2, -/// value='beta'`), copied to ClickHouse by the initial table copy. -/// -/// # WHEN -/// -/// 1. Postgres issues `TRUNCATE` on the table. -/// 2. After the table becomes empty in ClickHouse, a new row (`id=3, -/// value='gamma'`) is inserted into Postgres. -/// -/// # THEN -/// -/// After truncate, ClickHouse contains zero rows. -/// After the post-truncate insert, ClickHouse contains exactly one row: -/// - `id=3, value='gamma', cdc_operation='INSERT', cdc_lsn > 0`. -/// No pre-truncate rows survive. #[tokio::test(flavor = "multi_thread")] -async fn truncate_clears_table_and_accepts_new_inserts() { +async fn truncate_clears_table_and_accepts_new_inserts_merge_tree() { + truncate_clears_table_and_accepts_new_inserts_inner(ClickHouseEngine::MergeTree).await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn truncate_clears_table_and_accepts_new_inserts_replacing_merge_tree() { + truncate_clears_table_and_accepts_new_inserts_inner(ClickHouseEngine::ReplacingMergeTree).await; +} + +async fn truncate_clears_table_and_accepts_new_inserts_inner(engine: ClickHouseEngine) { init_test_tracing(); install_crypto_provider(); @@ -1021,7 +972,9 @@ async fn truncate_clears_table_and_accepts_new_inserts() { let clickhouse_db = setup_clickhouse_database().await; let store = NotifyingStore::new(); let pipeline_id: PipelineId = random(); - let destination = TestDestinationWrapper::wrap(clickhouse_db.build_destination(store.clone())); + let destination = TestDestinationWrapper::wrap( + clickhouse_db.build_destination_with_engine(store.clone(), engine).await, + ); let table_ready = store.notify_on_table_state_type(table_id, TableReplicationPhaseType::Ready).await; @@ -1038,7 +991,9 @@ async fn truncate_clears_table_and_accepts_new_inserts() { table_ready.notified().await; // Verify both rows arrived from table copy. - let rows: Vec = clickhouse_db.query(TRUNCATE_FLOW_SELECT).await; + let truncate_query = + || current_state_query(engine, TRUNCATE_FLOW_TABLE, ID_VALUE_PROJECTION, &["id"], "id"); + let rows: Vec = clickhouse_db.query(&truncate_query()).await; assert_eq!(rows.len(), 2, "table copy should produce two rows"); // --- WHEN: truncate, then insert a new row --- @@ -1051,7 +1006,7 @@ async fn truncate_clears_table_and_accepts_new_inserts() { truncate_notify.notified().await; - let rows: Vec = clickhouse_db.query(TRUNCATE_FLOW_SELECT).await; + let rows: Vec = clickhouse_db.query(&truncate_query()).await; assert!(rows.is_empty(), "table should be empty after truncate"); let insert_notify = destination.wait_for_events_count(vec![(EventType::Insert, 1)]).await; @@ -1066,45 +1021,29 @@ async fn truncate_clears_table_and_accepts_new_inserts() { insert_notify.notified().await; - let rows: Vec = clickhouse_db.query(TRUNCATE_FLOW_SELECT).await; + let rows: Vec = clickhouse_db.query(&truncate_query()).await; pipeline.shutdown_and_wait().await.unwrap(); // --- THEN: only the post-truncate row exists --- assert_eq!(rows.len(), 1, "only post-truncate row should exist"); - - let r = &rows[0]; - assert_eq!(r.id, 3, "post-truncate row should have id=3 (serial continues)"); - assert_eq!(r.value, "gamma"); - assert_eq!(r.cdc_operation, "INSERT"); - assert!(r.cdc_lsn > 0, "post-truncate insert should come from CDC streaming"); + assert_eq!(rows[0].id, 3, "post-truncate row should have id=3 (serial continues)"); + assert_eq!(rows[0].value, "gamma"); } -/// SELECT query used to verify the `flush_split` test. -const FLUSH_SPLIT_SELECT: &str = concat!( - "SELECT id, value, cdc_operation, cdc_lsn ", - "FROM \"test_flush__split\" ", - "ORDER BY id, cdc_lsn", -); - /// Tests that the intermediate INSERT flush (`max_bytes_per_insert`) does not /// lose rows when a batch is split across multiple INSERT statements. -/// -/// # GIVEN -/// -/// A Postgres table with 10 rows, and a ClickHouse destination configured with -/// `max_bytes_per_insert = 1` (forcing a new INSERT after every single row). -/// -/// # WHEN -/// -/// The pipeline runs initial table copy from Postgres to ClickHouse. -/// -/// # THEN -/// -/// All 10 rows arrive in ClickHouse despite being split across many INSERT -/// statements. No rows are lost at flush boundaries. #[tokio::test(flavor = "multi_thread")] -async fn intermediate_flush_preserves_all_rows() { +async fn intermediate_flush_preserves_all_rows_merge_tree() { + intermediate_flush_preserves_all_rows_inner(ClickHouseEngine::MergeTree).await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn intermediate_flush_preserves_all_rows_replacing_merge_tree() { + intermediate_flush_preserves_all_rows_inner(ClickHouseEngine::ReplacingMergeTree).await; +} + +async fn intermediate_flush_preserves_all_rows_inner(engine: ClickHouseEngine) { init_test_tracing(); install_crypto_provider(); @@ -1137,13 +1076,16 @@ async fn intermediate_flush_preserves_all_rows() { let clickhouse_db = setup_clickhouse_database().await; let store = NotifyingStore::new(); let pipeline_id: PipelineId = random(); - let destination = clickhouse_db.build_destination_with_config( - store.clone(), - ClickHouseInserterConfig { - // 1 byte -- forces a new INSERT after every row. - max_bytes_per_insert: 1, - }, - ); + let destination = clickhouse_db + .build_destination_with_config( + store.clone(), + ClickHouseInserterConfig { + // 1 byte -- forces a new INSERT after every row. + max_bytes_per_insert: 1, + engine, + }, + ) + .await; let table_ready = store.notify_on_table_state_type(table_id, TableReplicationPhaseType::Ready).await; @@ -1162,7 +1104,9 @@ async fn intermediate_flush_preserves_all_rows() { pipeline.shutdown_and_wait().await.unwrap(); // --- THEN: all rows arrive despite being split across many INSERTs --- - let rows: Vec = clickhouse_db.query(FLUSH_SPLIT_SELECT).await; + let query = + current_state_query(engine, "test_flush__split", ID_VALUE_PROJECTION, &["id"], "id"); + let rows: Vec = clickhouse_db.query(&query).await; assert_eq!(rows.len(), row_count, "all rows must survive intermediate flush splits"); for (i, r) in rows.iter().enumerate() { @@ -1170,8 +1114,6 @@ async fn intermediate_flush_preserves_all_rows() { let expected_value = format!("row_{}", i + 1); assert_eq!(r.id, expected_id, "row {} id mismatch", i + 1); assert_eq!(r.value, expected_value, "row {} value mismatch", i + 1); - assert_eq!(r.cdc_operation, "INSERT"); - assert_eq!(r.cdc_lsn, 0, "all rows should be from table copy"); } } @@ -1196,7 +1138,16 @@ async fn intermediate_flush_preserves_all_rows() { /// - One from table copy (`cdc_lsn = 0`) /// - One from CDC streaming (`cdc_lsn > 0`) #[tokio::test(flavor = "multi_thread")] -async fn multiple_tables_receive_independent_writes() { +async fn multiple_tables_receive_independent_writes_merge_tree() { + multiple_tables_receive_independent_writes_inner(ClickHouseEngine::MergeTree).await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn multiple_tables_receive_independent_writes_replacing_merge_tree() { + multiple_tables_receive_independent_writes_inner(ClickHouseEngine::ReplacingMergeTree).await; +} + +async fn multiple_tables_receive_independent_writes_inner(engine: ClickHouseEngine) { init_test_tracing(); install_crypto_provider(); @@ -1240,7 +1191,9 @@ async fn multiple_tables_receive_independent_writes() { let clickhouse_db = setup_clickhouse_database().await; let store = NotifyingStore::new(); let pipeline_id: PipelineId = random(); - let destination = TestDestinationWrapper::wrap(clickhouse_db.build_destination(store.clone())); + let destination = TestDestinationWrapper::wrap( + clickhouse_db.build_destination_with_engine(store.clone(), engine).await, + ); let table_a_ready = store.notify_on_table_state_type(table_a_id, TableReplicationPhaseType::Ready).await; @@ -1279,209 +1232,32 @@ async fn multiple_tables_receive_independent_writes() { event_notify.notified().await; - let select_a = concat!( - "SELECT id, value, cdc_operation, cdc_lsn ", - "FROM \"test_multi__a\" ", - "ORDER BY id, cdc_lsn", - ); - let select_b = concat!( - "SELECT id, value, cdc_operation, cdc_lsn ", - "FROM \"test_multi__b\" ", - "ORDER BY id, cdc_lsn", - ); - - let rows_a: Vec = clickhouse_db.query(select_a).await; - let rows_b: Vec = clickhouse_db.query(select_b).await; + let query_a = current_state_query(engine, "test_multi__a", ID_VALUE_PROJECTION, &["id"], "id"); + let query_b = current_state_query(engine, "test_multi__b", ID_VALUE_PROJECTION, &["id"], "id"); + let rows_a: Vec = clickhouse_db.query(&query_a).await; + let rows_b: Vec = clickhouse_db.query(&query_b).await; pipeline.shutdown_and_wait().await.unwrap(); - // --- THEN: each table has one copied row and one streamed row --- + // --- THEN: each table has both rows in current state --- assert_eq!(rows_a.len(), 2, "multi_a should have 2 rows"); assert_eq!(rows_b.len(), 2, "multi_b should have 2 rows"); - - assert_eq!(rows_a[0].id, 1); - assert_eq!(rows_a[0].value, "init_a"); - assert_eq!(rows_a[0].cdc_operation, "INSERT"); - assert_eq!(rows_a[0].cdc_lsn, 0); - - assert_eq!(rows_a[1].id, 2); - assert_eq!(rows_a[1].value, "streamed_a"); - assert_eq!(rows_a[1].cdc_operation, "INSERT"); - assert!(rows_a[1].cdc_lsn > 0); - - assert_eq!(rows_b[0].id, 1); - assert_eq!(rows_b[0].value, "init_b"); - assert_eq!(rows_b[0].cdc_operation, "INSERT"); - assert_eq!(rows_b[0].cdc_lsn, 0); - - assert_eq!(rows_b[1].id, 2); - assert_eq!(rows_b[1].value, "streamed_b"); - assert_eq!(rows_b[1].cdc_operation, "INSERT"); - assert!(rows_b[1].cdc_lsn > 0); + assert_eq!((rows_a[0].id, rows_a[0].value.as_str()), (1, "init_a")); + assert_eq!((rows_a[1].id, rows_a[1].value.as_str()), (2, "streamed_a")); + assert_eq!((rows_b[0].id, rows_b[0].value.as_str()), (1, "init_b")); + assert_eq!((rows_b[1].id, rows_b[1].value.as_str()), (2, "streamed_b")); } -/// SELECT query used to verify the `tx_order` test. -const TX_ORDER_SELECT: &str = concat!( - "SELECT id, value, cdc_operation, cdc_lsn ", - "FROM \"test_tx__order\" ", - "ORDER BY id, cdc_lsn", -); - -/// Tests that updates from separately committed transactions arrive in -/// ClickHouse with LSNs reflecting Postgres commit order. -/// -/// # GIVEN -/// -/// A Postgres table with one row (`id=1, value='original'`), copied to -/// ClickHouse. Two database connections to the same source. -/// -/// # WHEN -/// -/// Transaction A (on connection 1) updates the row to `'update_a'` and -/// commits. Then transaction B (on connection 2) updates the row to -/// `'update_b'` and commits. -/// -/// # THEN -/// -/// ClickHouse contains three rows (append-only CDC) with strictly -/// increasing `cdc_lsn`: -/// - `value='original'`, `cdc_operation='INSERT'`, `cdc_lsn=0` -/// - `value='update_a'`, `cdc_operation='UPDATE'`, `cdc_lsn > 0` -/// - `value='update_b'`, `cdc_operation='UPDATE'`, `cdc_lsn > update_a's lsn` -#[tokio::test(flavor = "multi_thread")] -async fn sequential_transactions_preserve_commit_order() { - init_test_tracing(); - install_crypto_provider(); - - // --- GIVEN: one row, two database connections --- - let mut database_1 = spawn_source_database().await; - let mut database_2 = database_1.duplicate().await; - let table_name = test_table_name("tx_order"); - - let table_id = database_1 - .create_table(table_name.clone(), true, &[("value", "text not null")]) - .await - .expect("Failed to create tx_order test table"); - - let publication_name = "test_pub_clickhouse_tx_order"; - database_1 - .create_publication(publication_name, std::slice::from_ref(&table_name)) - .await - .expect("Failed to create tx_order publication"); - - database_1 - .run_sql(&format!( - "INSERT INTO {} (value) VALUES ('original')", - table_name.as_quoted_identifier(), - )) - .await - .expect("Failed to insert initial tx_order row"); - - let clickhouse_db = setup_clickhouse_database().await; - let store = NotifyingStore::new(); - let pipeline_id: PipelineId = random(); - let destination = TestDestinationWrapper::wrap(clickhouse_db.build_destination(store.clone())); - - let table_ready = - store.notify_on_table_state_type(table_id, TableReplicationPhaseType::Ready).await; - - let mut pipeline = create_pipeline( - &database_1.config, - pipeline_id, - publication_name.to_owned(), - store, - destination.clone(), - ); - - pipeline.start().await.unwrap(); - table_ready.notified().await; - - let event_notify = destination.wait_for_events_count(vec![(EventType::Update, 2)]).await; - - // --- WHEN: two transactions commit sequentially on separate connections --- - let tx_a = database_1.begin_transaction().await; - tx_a.run_sql(&format!( - "UPDATE {} SET value = 'update_a' WHERE id = 1", - table_name.as_quoted_identifier(), - )) - .await - .expect("Failed to execute update_a"); - tx_a.commit_transaction().await; - - let tx_b = database_2.begin_transaction().await; - tx_b.run_sql(&format!( - "UPDATE {} SET value = 'update_b' WHERE id = 1", - table_name.as_quoted_identifier(), - )) - .await - .expect("Failed to execute update_b"); - tx_b.commit_transaction().await; - - event_notify.notified().await; - - let rows: Vec = clickhouse_db.query(TX_ORDER_SELECT).await; - - pipeline.shutdown_and_wait().await.unwrap(); - - // --- THEN: three rows with strictly increasing LSNs --- - assert_eq!(rows.len(), 3, "expected INSERT + two UPDATEs"); - - let r = &rows[0]; - assert_eq!(r.value, "original"); - assert_eq!(r.cdc_operation, "INSERT"); - assert_eq!(r.cdc_lsn, 0); - - let r = &rows[1]; - assert_eq!(r.value, "update_a"); - assert_eq!(r.cdc_operation, "UPDATE"); - assert!(r.cdc_lsn > 0); - - let r = &rows[2]; - assert_eq!(r.value, "update_b"); - assert_eq!(r.cdc_operation, "UPDATE"); - assert!(r.cdc_lsn > rows[1].cdc_lsn, "update_b must have a higher LSN than update_a"); -} - -/// Row struct for the wide default-identity delete test. +/// Current-state row for the wide default-identity delete test (user +/// columns only). #[derive(clickhouse::Row, serde::Deserialize, Debug)] -struct DefaultIdentityDeleteRow { +struct DefaultIdentityRow { id: i64, - smallint_col: i16, - integer_col: i32, - bigint_col: i64, - real_col: f32, - double_col: f64, - numeric_col: String, - boolean_col: bool, text_col: String, - varchar_col: String, - date_col: i32, - timestamp_col: i64, - timestamptz_col: i64, - time_col: String, - jsonb_col: String, - bytea_col: String, - uuid_col: String, - nullable_text: Option, - nullable_int: Option, - int_array_col: Vec>, - text_array_col: Vec>, - cdc_operation: String, - cdc_lsn: u64, } -const DEFAULT_IDENTITY_DELETE_SELECT: &str = concat!( - "SELECT id, ", - "smallint_col, integer_col, bigint_col, real_col, double_col, ", - "numeric_col, boolean_col, text_col, varchar_col, ", - "date_col, timestamp_col, timestamptz_col, time_col, ", - "jsonb_col, bytea_col, toString(uuid_col) AS uuid_col, ", - "nullable_text, nullable_int, ", - "int_array_col, text_array_col, ", - "cdc_operation, cdc_lsn ", - "FROM \"test_default__identity__delete\" ", - "ORDER BY id, cdc_lsn", -); +const DEFAULT_IDENTITY_PROJECTION: &str = "id, text_col"; +const DEFAULT_IDENTITY_TABLE: &str = "test_default__identity__delete"; /// Tests that a DELETE under default replica identity (PK only) produces a /// tombstone row with the correct PK and zero-value defaults for every @@ -1509,7 +1285,16 @@ const DEFAULT_IDENTITY_DELETE_SELECT: &str = concat!( /// - Nullable scalars filled with NULL /// - Arrays filled with empty arrays #[tokio::test(flavor = "multi_thread")] -async fn delete_with_default_replica_identity() { +async fn delete_with_default_replica_identity_merge_tree() { + delete_with_default_replica_identity_inner(ClickHouseEngine::MergeTree).await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn delete_with_default_replica_identity_replacing_merge_tree() { + delete_with_default_replica_identity_inner(ClickHouseEngine::ReplacingMergeTree).await; +} + +async fn delete_with_default_replica_identity_inner(engine: ClickHouseEngine) { init_test_tracing(); install_crypto_provider(); @@ -1584,7 +1369,9 @@ async fn delete_with_default_replica_identity() { let clickhouse_db = setup_clickhouse_database().await; let store = NotifyingStore::new(); let pipeline_id: PipelineId = random(); - let destination = TestDestinationWrapper::wrap(clickhouse_db.build_destination(store.clone())); + let destination = TestDestinationWrapper::wrap( + clickhouse_db.build_destination_with_engine(store.clone(), engine).await, + ); let table_ready = store.notify_on_table_state_type(table_id, TableReplicationPhaseType::Ready).await; @@ -1632,74 +1419,23 @@ async fn delete_with_default_replica_identity() { event_notify.notified().await; - let rows: Vec = - clickhouse_db.query(DEFAULT_IDENTITY_DELETE_SELECT).await; + let query = current_state_query( + engine, + DEFAULT_IDENTITY_TABLE, + DEFAULT_IDENTITY_PROJECTION, + &["id"], + "id", + ); + let rows: Vec = clickhouse_db.query(&query).await; pipeline.shutdown_and_wait().await.unwrap(); - // --- THEN: DELETE tombstone has zero-value defaults for all types --- - assert_eq!(rows.len(), 4, "expected 2 copied INSERTs + DELETE + new INSERT"); - - // Row 1: copied, untouched -- spot check. - let r = &rows[0]; - assert_eq!(r.id, 1); - assert_eq!(r.text_col, "keep"); - assert_eq!(r.integer_col, 10); - assert!(r.boolean_col); - assert_eq!(r.nullable_text, Some("present".to_owned())); - assert_eq!(r.int_array_col, vec![Some(1), Some(2), Some(3)]); - assert_eq!(r.cdc_operation, "INSERT"); - - // Row 2: copied, will be deleted. - let r = &rows[1]; - assert_eq!(r.id, 2); - assert_eq!(r.text_col, "delete_me"); - assert_eq!(r.cdc_operation, "INSERT"); - - // Row 3: DELETE tombstone -- every non-PK column type verified. - let r = &rows[2]; - assert_eq!(r.id, 2, "DELETE must target the correct row"); - assert_eq!(r.cdc_operation, "DELETE"); - assert!(r.cdc_lsn > 0); - // Non-nullable scalars -> zero values. - assert_eq!(r.smallint_col, 0, "smallint -> 0"); - assert_eq!(r.integer_col, 0, "integer -> 0"); - assert_eq!(r.bigint_col, 0, "bigint -> 0"); - assert!(r.real_col.abs() < 1e-6, "real -> 0.0"); - assert!(r.double_col.abs() < 1e-9, "double -> 0.0"); - assert_eq!(r.numeric_col, "", "numeric -> empty string"); - assert!(!r.boolean_col, "boolean -> false"); - assert_eq!(r.text_col, "", "text -> empty string"); - assert_eq!(r.varchar_col, "", "varchar -> empty string"); - assert_eq!(r.date_col, 0, "date -> 1970-01-01 (day 0)"); - assert_eq!(r.timestamp_col, 0, "timestamp -> unix epoch"); - assert_eq!(r.timestamptz_col, 0, "timestamptz -> unix epoch"); - assert_eq!(r.time_col, "", "time -> empty string (String-mapped)"); - assert_eq!(r.jsonb_col, "", "jsonb -> empty string (String-mapped)"); - assert_eq!(r.bytea_col, "", "bytea -> empty string"); - assert_eq!(r.uuid_col, "00000000-0000-0000-0000-000000000000", "uuid -> nil UUID"); - // Nullable scalars -> NULL. - assert_eq!(r.nullable_text, None, "nullable text -> NULL"); - assert_eq!(r.nullable_int, None, "nullable int -> NULL"); - // Arrays -> empty. - assert!(r.int_array_col.is_empty(), "int array -> empty"); - assert!(r.text_array_col.is_empty(), "text array -> empty"); - - // Row 4: post-delete INSERT proves pipeline continued. - let r = &rows[3]; - assert_eq!(r.id, 3); - assert_eq!(r.text_col, "after"); - assert_eq!(r.cdc_operation, "INSERT"); - assert!(r.cdc_lsn > 0); + // --- THEN: deleted row absent; surviving + post-insert rows present --- + assert_eq!(rows.len(), 2, "current state should show id=1 and id=3 only"); + assert_eq!((rows[0].id, rows[0].text_col.as_str()), (1, "keep")); + assert_eq!((rows[1].id, rows[1].text_col.as_str()), (3, "after")); } -/// SELECT query used to verify the `large_batch` test. -const LARGE_BATCH_SELECT: &str = concat!( - "SELECT id, value, cdc_operation, cdc_lsn ", - "FROM \"test_large__batch\" ", - "ORDER BY id, cdc_lsn", -); - /// Tests that a large table copy (1024 rows) completes without data loss or /// corruption. /// @@ -1718,7 +1454,16 @@ const LARGE_BATCH_SELECT: &str = concat!( /// (first, last, powers of two, and a few interior points) are spot-checked /// for correct id and value. #[tokio::test(flavor = "multi_thread")] -async fn exclusive_large_batch_table_copy() { +async fn exclusive_large_batch_table_copy_merge_tree() { + exclusive_large_batch_table_copy_inner(ClickHouseEngine::MergeTree).await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn exclusive_large_batch_table_copy_replacing_merge_tree() { + exclusive_large_batch_table_copy_inner(ClickHouseEngine::ReplacingMergeTree).await; +} + +async fn exclusive_large_batch_table_copy_inner(engine: ClickHouseEngine) { init_test_tracing(); install_crypto_provider(); @@ -1754,7 +1499,7 @@ async fn exclusive_large_batch_table_copy() { let clickhouse_db = setup_clickhouse_database().await; let store = NotifyingStore::new(); let pipeline_id: PipelineId = random(); - let destination = clickhouse_db.build_destination(store.clone()); + let destination = clickhouse_db.build_destination_with_engine(store.clone(), engine).await; let table_ready = store.notify_on_table_state_type(table_id, TableReplicationPhaseType::Ready).await; @@ -1773,7 +1518,9 @@ async fn exclusive_large_batch_table_copy() { pipeline.shutdown_and_wait().await.unwrap(); // --- THEN: all rows arrive, spot-check a sample --- - let rows: Vec = clickhouse_db.query(LARGE_BATCH_SELECT).await; + let query = + current_state_query(engine, "test_large__batch", ID_VALUE_PROJECTION, &["id"], "id"); + let rows: Vec = clickhouse_db.query(&query).await; assert_eq!(rows.len(), row_count, "all 1024 rows must arrive"); // Spot-check: first, last, powers of two, and a few interior points. @@ -1782,8 +1529,6 @@ async fn exclusive_large_batch_table_copy() { let r = &rows[id - 1]; assert_eq!(r.id, id as i64, "row {id} id mismatch"); assert_eq!(r.value, format!("val_{id:04}"), "row {id} value mismatch"); - assert_eq!(r.cdc_operation, "INSERT"); - assert_eq!(r.cdc_lsn, 0); } } @@ -1829,14 +1574,13 @@ async fn validate_connectivity_fails_against_unreachable_clickhouse() { /// Row struct for the ADD COLUMN test after schema change. /// Columns: id, name, age, email, score. -#[derive(clickhouse::Row, serde::Deserialize, Debug)] +#[derive(clickhouse::Row, serde::Deserialize, Debug, PartialEq, Eq)] struct AddColumnRow { id: i64, name: String, age: i32, email: Option, score: Option, - cdc_operation: String, } /// Tests that ALTER TABLE ADD COLUMN in Postgres propagates to ClickHouse @@ -1860,7 +1604,16 @@ struct AddColumnRow { /// rows. Bob's row has 'bob@example.com' and 7. The destination metadata /// snapshot_id has increased. #[tokio::test(flavor = "multi_thread")] -async fn schema_change_add_column() { +async fn schema_change_add_column_merge_tree() { + schema_change_add_column_inner(ClickHouseEngine::MergeTree).await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn schema_change_add_column_replacing_merge_tree() { + schema_change_add_column_inner(ClickHouseEngine::ReplacingMergeTree).await; +} + +async fn schema_change_add_column_inner(engine: ClickHouseEngine) { init_test_tracing(); install_crypto_provider(); @@ -1896,7 +1649,9 @@ async fn schema_change_add_column() { let clickhouse_db = setup_clickhouse_database().await; let store = NotifyingStore::new(); let pipeline_id: PipelineId = random(); - let destination = TestDestinationWrapper::wrap(clickhouse_db.build_destination(store.clone())); + let destination = TestDestinationWrapper::wrap( + clickhouse_db.build_destination_with_engine(store.clone(), engine).await, + ); let table_ready = store.notify_on_table_state_type(table_id, TableReplicationPhaseType::Ready).await; @@ -1950,12 +1705,25 @@ async fn schema_change_add_column() { event_notify.notified().await; - let select = concat!( - "SELECT id, name, age, email, score, cdc_operation ", - "FROM \"test_schema__add__col\" ", - "ORDER BY id", + let query = current_state_query( + engine, + clickhouse_table_name, + "id, name, age, email, score", + &["id"], + "id", ); - let rows: Vec = clickhouse_db.query(select).await; + let rows: Vec = clickhouse_db.query(&query).await; + let current_view_rows = if matches!(engine, ClickHouseEngine::ReplacingMergeTree) { + let rows: Vec = clickhouse_db + .query( + "SELECT id, name, age, email, score FROM \"test_schema__add__col__current\" ORDER \ + BY id", + ) + .await; + Some(rows) + } else { + None + }; pipeline.shutdown_and_wait().await.unwrap(); @@ -1983,7 +1751,6 @@ async fn schema_change_add_column() { assert_eq!(rows[0].age, 25); assert_eq!(rows[0].email, None, "Alice's email should be NULL (column added after her row)"); assert_eq!(rows[0].score, None, "Alice's score should be NULL (column added after her row)"); - assert_eq!(rows[0].cdc_operation, "INSERT"); // Bob: post-change row, added columns present. assert_eq!(rows[1].id, 2); @@ -1991,7 +1758,13 @@ async fn schema_change_add_column() { assert_eq!(rows[1].age, 30); assert_eq!(rows[1].email, Some("bob@example.com".to_owned())); assert_eq!(rows[1].score, Some(7)); - assert_eq!(rows[1].cdc_operation, "INSERT"); + + if let Some(view_rows) = current_view_rows { + assert_eq!( + view_rows, rows, + "__current view should match the evolved ReplacingMergeTree schema" + ); + } // Metadata snapshot_id should have advanced. let final_metadata = store @@ -2008,13 +1781,12 @@ async fn schema_change_add_column() { /// Row struct for the combined schema change test after all changes. /// Columns: id, full_name (renamed), status (kept), email (added). /// age is dropped. -#[derive(clickhouse::Row, serde::Deserialize, Debug)] +#[derive(clickhouse::Row, serde::Deserialize, Debug, PartialEq, Eq)] struct CombinedSchemaChangeRow { id: i64, full_name: String, status: Option, email: Option, - cdc_operation: String, } /// Tests that multiple schema changes (ADD, DROP, RENAME) in Postgres all @@ -2045,7 +1817,16 @@ struct CombinedSchemaChangeRow { /// email. Bob's row has the new values. /// The destination metadata snapshot_id has increased. #[tokio::test(flavor = "multi_thread")] -async fn schema_change_add_drop_rename() { +async fn schema_change_add_drop_rename_merge_tree() { + schema_change_add_drop_rename_inner(ClickHouseEngine::MergeTree).await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn schema_change_add_drop_rename_replacing_merge_tree() { + schema_change_add_drop_rename_inner(ClickHouseEngine::ReplacingMergeTree).await; +} + +async fn schema_change_add_drop_rename_inner(engine: ClickHouseEngine) { init_test_tracing(); install_crypto_provider(); @@ -2081,7 +1862,9 @@ async fn schema_change_add_drop_rename() { let clickhouse_db = setup_clickhouse_database().await; let store = NotifyingStore::new(); let pipeline_id: PipelineId = random(); - let destination = TestDestinationWrapper::wrap(clickhouse_db.build_destination(store.clone())); + let destination = TestDestinationWrapper::wrap( + clickhouse_db.build_destination_with_engine(store.clone(), engine).await, + ); let table_ready = store.notify_on_table_state_type(table_id, TableReplicationPhaseType::Ready).await; @@ -2143,12 +1926,25 @@ async fn schema_change_add_drop_rename() { event_notify.notified().await; - let select = concat!( - "SELECT id, full_name, status, email, cdc_operation ", - "FROM \"test_schema__multi\" ", - "ORDER BY id", + let query = current_state_query( + engine, + clickhouse_table_name, + "id, full_name, status, email", + &["id"], + "id", ); - let rows: Vec = clickhouse_db.query(select).await; + let rows: Vec = clickhouse_db.query(&query).await; + let current_view_rows = if matches!(engine, ClickHouseEngine::ReplacingMergeTree) { + let rows: Vec = clickhouse_db + .query( + "SELECT id, full_name, status, email FROM \"test_schema__multi__current\" ORDER \ + BY id", + ) + .await; + Some(rows) + } else { + None + }; pipeline.shutdown_and_wait().await.unwrap(); @@ -2163,14 +1959,19 @@ async fn schema_change_add_drop_rename() { assert_eq!(rows[0].full_name, "Alice", "renamed column should preserve data"); assert_eq!(rows[0].status, Some("active".to_owned())); assert_eq!(rows[0].email, None, "Alice's email should be NULL (added after her row)"); - assert_eq!(rows[0].cdc_operation, "INSERT"); // Bob: post-change row. assert_eq!(rows[1].id, 2); assert_eq!(rows[1].full_name, "Bob"); assert_eq!(rows[1].status, Some("pending".to_owned())); assert_eq!(rows[1].email, Some("bob@example.com".to_owned())); - assert_eq!(rows[1].cdc_operation, "INSERT"); + + if let Some(view_rows) = current_view_rows { + assert_eq!( + view_rows, rows, + "__current view should match the evolved ReplacingMergeTree schema" + ); + } // Metadata snapshot_id should have advanced. let final_metadata = store diff --git a/crates/etl-destinations/tests/clickhouse/pipeline_merge_tree.rs b/crates/etl-destinations/tests/clickhouse/pipeline_merge_tree.rs new file mode 100644 index 000000000..551c0e361 --- /dev/null +++ b/crates/etl-destinations/tests/clickhouse/pipeline_merge_tree.rs @@ -0,0 +1,141 @@ +//! MergeTree-only integration tests. These verify event-log semantics +//! (`cdc_operation` + `cdc_lsn`) that exist only on the MergeTree engine; the +//! parameterized spine in `pipeline.rs` covers current-state behavior on +//! both engines. + +use etl::{ + state::table::TableReplicationPhaseType, + test_utils::{ + database::{spawn_source_database, test_table_name}, + notifying_store::NotifyingStore, + pipeline::create_pipeline, + test_destination_wrapper::TestDestinationWrapper, + }, + types::{EventType, PipelineId}, +}; +use etl_config::shared::ClickHouseEngine; +use etl_destinations::clickhouse::test_utils::setup_clickhouse_database; +use etl_telemetry::tracing::init_test_tracing; +use rand::random; + +use crate::support::clickhouse::install_crypto_provider; + +/// MergeTree event-log row: includes CDC metadata. All three operations in this +/// test target the same source row, so `id` is asserted on alongside the +/// CDC columns. +#[derive(clickhouse::Row, serde::Deserialize, Debug)] +struct EventLogRow { + id: i64, + value: String, + cdc_operation: String, + cdc_lsn: u64, +} + +const TX_ORDER_SELECT: &str = concat!( + "SELECT id, value, cdc_operation, cdc_lsn ", + "FROM \"test_tx__order\" ", + "ORDER BY id, cdc_lsn", +); + +/// MergeTree-only: verifies that updates from separately committed transactions +/// arrive with strictly increasing `cdc_lsn` matching Postgres commit order. +/// +/// ReplacingMergeTree collapses the event log under `FINAL`, so this ordering +/// check has no analog on the ReplacingMergeTree side. +#[tokio::test(flavor = "multi_thread")] +async fn sequential_transactions_preserve_commit_order_merge_tree() { + init_test_tracing(); + install_crypto_provider(); + + // --- GIVEN: one row, two database connections --- + let mut database_1 = spawn_source_database().await; + let mut database_2 = database_1.duplicate().await; + let table_name = test_table_name("tx_order"); + + let table_id = database_1 + .create_table(table_name.clone(), true, &[("value", "text not null")]) + .await + .expect("Failed to create tx_order test table"); + + let publication_name = "test_pub_clickhouse_tx_order"; + database_1 + .create_publication(publication_name, std::slice::from_ref(&table_name)) + .await + .expect("Failed to create tx_order publication"); + + database_1 + .run_sql(&format!( + "INSERT INTO {} (value) VALUES ('original')", + table_name.as_quoted_identifier(), + )) + .await + .expect("Failed to insert initial tx_order row"); + + let clickhouse_db = setup_clickhouse_database().await; + let store = NotifyingStore::new(); + let pipeline_id: PipelineId = random(); + let destination = TestDestinationWrapper::wrap( + clickhouse_db + .build_destination_with_engine(store.clone(), ClickHouseEngine::MergeTree) + .await, + ); + + let table_ready = + store.notify_on_table_state_type(table_id, TableReplicationPhaseType::Ready).await; + + let mut pipeline = create_pipeline( + &database_1.config, + pipeline_id, + publication_name.to_owned(), + store, + destination.clone(), + ); + + pipeline.start().await.unwrap(); + table_ready.notified().await; + + let event_notify = destination.wait_for_events_count(vec![(EventType::Update, 2)]).await; + + // --- WHEN: two transactions commit sequentially on separate connections --- + let tx_a = database_1.begin_transaction().await; + tx_a.run_sql(&format!( + "UPDATE {} SET value = 'update_a' WHERE id = 1", + table_name.as_quoted_identifier(), + )) + .await + .expect("Failed to execute update_a"); + tx_a.commit_transaction().await; + + let tx_b = database_2.begin_transaction().await; + tx_b.run_sql(&format!( + "UPDATE {} SET value = 'update_b' WHERE id = 1", + table_name.as_quoted_identifier(), + )) + .await + .expect("Failed to execute update_b"); + tx_b.commit_transaction().await; + + event_notify.notified().await; + + let rows: Vec = clickhouse_db.query(TX_ORDER_SELECT).await; + + pipeline.shutdown_and_wait().await.unwrap(); + + // --- THEN: three rows on id=1 with strictly increasing LSNs --- + assert_eq!(rows.len(), 3, "expected INSERT + two UPDATEs"); + + assert_eq!(rows[0].id, 1); + assert_eq!(rows[0].value, "original"); + assert_eq!(rows[0].cdc_operation, "INSERT"); + assert_eq!(rows[0].cdc_lsn, 0); + + assert_eq!(rows[1].id, 1); + assert_eq!(rows[1].value, "update_a"); + assert_eq!(rows[1].cdc_operation, "UPDATE"); + assert!(rows[1].cdc_lsn > 0); + + assert_eq!(rows[2].id, 1); + assert_eq!(rows[2].value, "update_b"); + assert_eq!(rows[2].cdc_operation, "UPDATE"); + assert!(rows[2].cdc_lsn > rows[1].cdc_lsn, "update_b must have a higher LSN than update_a"); +} diff --git a/crates/etl-destinations/tests/clickhouse/pipeline_replacing_merge_tree.rs b/crates/etl-destinations/tests/clickhouse/pipeline_replacing_merge_tree.rs new file mode 100644 index 000000000..44f08e757 --- /dev/null +++ b/crates/etl-destinations/tests/clickhouse/pipeline_replacing_merge_tree.rs @@ -0,0 +1,761 @@ +//! ReplacingMergeTree-only integration tests. These verify +//! ReplacingMergeTree-specific semantics that have no analog under MergeTree: +//! same-LSN tie-break, FINAL reads, the `__current` view, `OPTIMIZE ... FINAL +//! CLEANUP` cleanup, PK-less source rejection, and composite PK ORDER BY. + +use etl::{ + state::table::{TableReplicationPhase, TableReplicationPhaseType}, + store::state::StateStore, + test_utils::{ + database::{spawn_source_database, test_table_name}, + notifying_store::NotifyingStore, + pipeline::create_pipeline, + test_destination_wrapper::TestDestinationWrapper, + }, + types::{EventType, PipelineId}, +}; +use etl_config::shared::ClickHouseEngine; +use etl_destinations::clickhouse::test_utils::setup_clickhouse_database; +use etl_telemetry::tracing::init_test_tracing; +use rand::random; + +use crate::support::clickhouse::{ + current_state_query, install_crypto_provider, optimize_final_cleanup_sql, table_engine_query, +}; + +#[derive(clickhouse::Row, serde::Deserialize, Debug)] +struct IdValueRow { + id: i64, + value: String, +} + +#[derive(clickhouse::Row, serde::Deserialize, Debug)] +struct CountRow { + count: u64, +} + +/// ReplacingMergeTree: source table must have a primary key. +/// `ensure_table_exists` rejects PK-less schemas under ReplacingMergeTree with +/// `SourceSchemaError`. +#[tokio::test(flavor = "multi_thread")] +async fn replacing_merge_tree_rejects_pkless_source_table() { + init_test_tracing(); + install_crypto_provider(); + + // --- GIVEN: a PK-less source table --- + let database = spawn_source_database().await; + let table_name = test_table_name("pkless_events"); + + let table_id = database + .create_table(table_name.clone(), false, &[("value", "text not null")]) + .await + .expect("Failed to create pkless table"); + + let publication_name = "test_pub_replacing_merge_tree_pkless"; + database + .create_publication(publication_name, std::slice::from_ref(&table_name)) + .await + .expect("Failed to create pkless publication"); + + database + .run_sql(&format!( + "INSERT INTO {} (value) VALUES ('seed')", + table_name.as_quoted_identifier(), + )) + .await + .expect("Failed to insert seed row"); + + // --- WHEN: pipeline runs under ReplacingMergeTree --- + let clickhouse_db = setup_clickhouse_database().await; + let store = NotifyingStore::new(); + let pipeline_id: PipelineId = random(); + let destination = clickhouse_db + .build_destination_with_engine(store.clone(), ClickHouseEngine::ReplacingMergeTree) + .await; + + let table_errored = + store.notify_on_table_state_type(table_id, TableReplicationPhaseType::Errored).await; + + let mut pipeline = create_pipeline( + &database.config, + pipeline_id, + publication_name.to_owned(), + store.clone(), + destination, + ); + + pipeline.start().await.unwrap(); + table_errored.notified().await; + pipeline.shutdown_and_wait().await.unwrap(); + + // --- THEN: the Errored phase reason names the PK-less rejection --- + let phase = store + .get_table_replication_state(table_id) + .await + .expect("state store should return the table's phase") + .expect("table should have a recorded replication phase"); + match phase { + TableReplicationPhase::Errored { reason, .. } => { + assert!( + reason.contains("primary key"), + "Errored reason should mention the PK requirement, got: {reason}" + ); + } + other => panic!("expected Errored phase, got {other:?}"), + } +} + +/// ReplacingMergeTree: a same-transaction INSERT followed by an UPDATE of the +/// same PK collapses under `FINAL` to the post-UPDATE value. Confirms that the +/// packed `_etl_version = (commit_lsn << 64) | tx_ordinal` tie-breaks +/// multi-event same-commit transactions correctly. +#[tokio::test(flavor = "multi_thread")] +async fn replacing_merge_tree_same_lsn_tx_insert_then_update_keeps_update() { + init_test_tracing(); + install_crypto_provider(); + + // --- GIVEN: an empty table copied to ClickHouse --- + let mut database = spawn_source_database().await; + let table_name = test_table_name("tie_break"); + + let table_id = database + .create_table(table_name.clone(), true, &[("value", "text not null")]) + .await + .expect("Failed to create tie_break table"); + + let publication_name = "test_pub_tie_break"; + database + .create_publication(publication_name, std::slice::from_ref(&table_name)) + .await + .expect("Failed to create publication"); + + let clickhouse_db = setup_clickhouse_database().await; + let store = NotifyingStore::new(); + let pipeline_id: PipelineId = random(); + let destination = TestDestinationWrapper::wrap( + clickhouse_db + .build_destination_with_engine(store.clone(), ClickHouseEngine::ReplacingMergeTree) + .await, + ); + + let table_ready = + store.notify_on_table_state_type(table_id, TableReplicationPhaseType::Ready).await; + + let mut pipeline = create_pipeline( + &database.config, + pipeline_id, + publication_name.to_owned(), + store, + destination.clone(), + ); + + pipeline.start().await.unwrap(); + table_ready.notified().await; + + // --- WHEN: INSERT + UPDATE in the same transaction --- + let event_notify = destination + .wait_for_events_count(vec![(EventType::Insert, 1), (EventType::Update, 1)]) + .await; + + let tx = database.begin_transaction().await; + tx.run_sql(&format!( + "INSERT INTO {} (value) VALUES ('initial')", + table_name.as_quoted_identifier(), + )) + .await + .expect("Failed to insert"); + tx.run_sql(&format!( + "UPDATE {} SET value = 'final' WHERE id = 1", + table_name.as_quoted_identifier(), + )) + .await + .expect("Failed to update"); + tx.commit_transaction().await; + + event_notify.notified().await; + + let query = current_state_query( + ClickHouseEngine::ReplacingMergeTree, + "test_tie__break", + "id, value", + &["id"], + "id", + ); + let rows: Vec = clickhouse_db.query(&query).await; + + pipeline.shutdown_and_wait().await.unwrap(); + + // --- THEN: the post-UPDATE value wins under FINAL --- + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].id, 1); + assert_eq!(rows[0].value, "final"); +} + +/// ReplacingMergeTree: a DELETE followed by an INSERT of the same PK in the +/// same transaction shows the post-INSERT row under FINAL (no lingering +/// tombstone). +#[tokio::test(flavor = "multi_thread")] +async fn replacing_merge_tree_same_lsn_tx_delete_then_insert_keeps_insert() { + init_test_tracing(); + install_crypto_provider(); + + // --- GIVEN: a row already replicated under REPLICA IDENTITY FULL --- + let mut database = spawn_source_database().await; + let table_name = test_table_name("del_ins"); + + let table_id = database + .create_table(table_name.clone(), true, &[("value", "text not null")]) + .await + .expect("Failed to create del_ins table"); + + database + .run_sql(&format!( + "ALTER TABLE {} REPLICA IDENTITY FULL", + table_name.as_quoted_identifier(), + )) + .await + .expect("Failed to set replica identity full"); + + let publication_name = "test_pub_del_ins"; + database + .create_publication(publication_name, std::slice::from_ref(&table_name)) + .await + .expect("Failed to create publication"); + + database + .run_sql(&format!( + "INSERT INTO {} (value) VALUES ('original')", + table_name.as_quoted_identifier(), + )) + .await + .expect("Failed to insert original row"); + + let clickhouse_db = setup_clickhouse_database().await; + let store = NotifyingStore::new(); + let pipeline_id: PipelineId = random(); + let destination = TestDestinationWrapper::wrap( + clickhouse_db + .build_destination_with_engine(store.clone(), ClickHouseEngine::ReplacingMergeTree) + .await, + ); + + let table_ready = + store.notify_on_table_state_type(table_id, TableReplicationPhaseType::Ready).await; + + let mut pipeline = create_pipeline( + &database.config, + pipeline_id, + publication_name.to_owned(), + store, + destination.clone(), + ); + + pipeline.start().await.unwrap(); + table_ready.notified().await; + + // --- WHEN: DELETE + INSERT of the same id in one transaction --- + let event_notify = destination + .wait_for_events_count(vec![(EventType::Delete, 1), (EventType::Insert, 1)]) + .await; + + let tx = database.begin_transaction().await; + tx.run_sql(&format!("DELETE FROM {} WHERE id = 1", table_name.as_quoted_identifier())) + .await + .expect("Failed to delete"); + tx.run_sql(&format!( + "INSERT INTO {} (id, value) VALUES (1, 'recreated')", + table_name.as_quoted_identifier(), + )) + .await + .expect("Failed to re-insert"); + tx.commit_transaction().await; + + event_notify.notified().await; + + let query = current_state_query( + ClickHouseEngine::ReplacingMergeTree, + "test_del__ins", + "id, value", + &["id"], + "id", + ); + let rows: Vec = clickhouse_db.query(&query).await; + + pipeline.shutdown_and_wait().await.unwrap(); + + // --- THEN: current state shows the re-inserted row, not the tombstone --- + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].id, 1); + assert_eq!(rows[0].value, "recreated"); +} + +/// ReplacingMergeTree: the auto-generated `__current` view exposes only user +/// columns and returns the current state of the table. +#[tokio::test(flavor = "multi_thread")] +async fn replacing_merge_tree_current_view_exposes_user_columns_and_current_state() { + init_test_tracing(); + install_crypto_provider(); + + // --- GIVEN: a row copied + then updated --- + let database = spawn_source_database().await; + let table_name = test_table_name("current_view"); + + let table_id = database + .create_table(table_name.clone(), true, &[("value", "text not null")]) + .await + .expect("Failed to create current_view table"); + + let publication_name = "test_pub_current_view"; + database + .create_publication(publication_name, std::slice::from_ref(&table_name)) + .await + .expect("Failed to create publication"); + + database + .run_sql(&format!( + "INSERT INTO {} (value) VALUES ('before')", + table_name.as_quoted_identifier(), + )) + .await + .expect("Failed to insert before-row"); + + let clickhouse_db = setup_clickhouse_database().await; + let store = NotifyingStore::new(); + let pipeline_id: PipelineId = random(); + let destination = TestDestinationWrapper::wrap( + clickhouse_db + .build_destination_with_engine(store.clone(), ClickHouseEngine::ReplacingMergeTree) + .await, + ); + + let table_ready = + store.notify_on_table_state_type(table_id, TableReplicationPhaseType::Ready).await; + + let mut pipeline = create_pipeline( + &database.config, + pipeline_id, + publication_name.to_owned(), + store, + destination.clone(), + ); + + pipeline.start().await.unwrap(); + table_ready.notified().await; + + let event_notify = destination.wait_for_events_count(vec![(EventType::Update, 1)]).await; + database + .run_sql(&format!( + "UPDATE {} SET value = 'after' WHERE id = 1", + table_name.as_quoted_identifier(), + )) + .await + .expect("Failed to update row"); + event_notify.notified().await; + + // --- WHEN: read via the __current view --- + let rows: Vec = clickhouse_db + .query("SELECT id, value FROM \"test_current__view__current\" ORDER BY id") + .await; + + pipeline.shutdown_and_wait().await.unwrap(); + + // --- THEN: view returns the post-UPDATE row --- + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].id, 1); + assert_eq!(rows[0].value, "after"); +} + +/// ReplacingMergeTree: composite PK source table is created with `ORDER BY` +/// matching `primary_key_ordinal_position`. +#[tokio::test(flavor = "multi_thread")] +async fn composite_pk_pk_order_by_matches_pk_ordinal() { + init_test_tracing(); + install_crypto_provider(); + + // GIVEN: a Postgres table with a composite PK whose ordinal order differs + // from table column order + let database = spawn_source_database().await; + let table_name = test_table_name("composite_pk"); + + // Source schema: (id serial PK is added by create_table). To get a true + // composite PK whose ordinal differs from table order, we drop the + // single-column PK and add a composite one. + let table_id = database + .create_table( + table_name.clone(), + true, + &[("tenant_id", "integer not null"), ("value", "text not null")], + ) + .await + .expect("Failed to create composite_pk table"); + + database + .run_sql(&format!( + "ALTER TABLE {} DROP CONSTRAINT {}_pkey", + table_name.as_quoted_identifier(), + table_name.name, + )) + .await + .expect("Failed to drop default pkey"); + + database + .run_sql(&format!( + "ALTER TABLE {} ADD PRIMARY KEY (tenant_id, id)", + table_name.as_quoted_identifier(), + )) + .await + .expect("Failed to add composite primary key"); + + let publication_name = "test_pub_composite_pk"; + database + .create_publication(publication_name, std::slice::from_ref(&table_name)) + .await + .expect("Failed to create publication"); + + database + .run_sql(&format!( + "INSERT INTO {} (tenant_id, value) VALUES (1, 'alpha')", + table_name.as_quoted_identifier(), + )) + .await + .expect("Failed to insert seed"); + + let clickhouse_db = setup_clickhouse_database().await; + let store = NotifyingStore::new(); + let pipeline_id: PipelineId = random(); + let destination = clickhouse_db + .build_destination_with_engine(store.clone(), ClickHouseEngine::ReplacingMergeTree) + .await; + + let table_ready = + store.notify_on_table_state_type(table_id, TableReplicationPhaseType::Ready).await; + + let mut pipeline = create_pipeline( + &database.config, + pipeline_id, + publication_name.to_owned(), + store, + destination, + ); + + pipeline.start().await.unwrap(); + table_ready.notified().await; + pipeline.shutdown_and_wait().await.unwrap(); + + // --- WHEN: query system.tables for the created ReplacingMergeTree ORDER BY + // clause --- + #[derive(clickhouse::Row, serde::Deserialize)] + struct EngineRow { + engine: String, + } + let engine_rows: Vec = + clickhouse_db.query(&table_engine_query("test_composite__pk")).await; + assert_eq!(engine_rows.len(), 1); + assert_eq!(engine_rows[0].engine, "ReplacingMergeTree"); + + // Read the sorting key from system.tables. + #[derive(clickhouse::Row, serde::Deserialize)] + struct SortingKeyRow { + sorting_key: String, + } + let sorting: Vec = clickhouse_db + .query( + "SELECT sorting_key FROM system.tables WHERE database = currentDatabase() AND name = \ + 'test_composite__pk'", + ) + .await; + + // --- THEN: sorting_key matches PK ordinal order (tenant_id, id) --- + assert_eq!(sorting.len(), 1); + assert_eq!(sorting[0].sorting_key, "tenant_id, id"); +} + +/// ReplacingMergeTree: an initial-copy row followed by a streaming UPDATE +/// collapses under `FINAL` to the streamed value. The initial-copy row has +/// `_etl_version = 0`; the streamed UPDATE has a non-zero packed +/// `_etl_version`, so the streamed row wins. +#[tokio::test(flavor = "multi_thread")] +async fn replacing_merge_tree_streamed_update_wins_over_initial_copy_row() { + init_test_tracing(); + install_crypto_provider(); + + // --- GIVEN: a row copied via initial copy --- + let database = spawn_source_database().await; + let table_name = test_table_name("copy_then_update"); + + let table_id = database + .create_table(table_name.clone(), true, &[("value", "text not null")]) + .await + .expect("Failed to create copy_then_update table"); + + let publication_name = "test_pub_copy_then_update"; + database + .create_publication(publication_name, std::slice::from_ref(&table_name)) + .await + .expect("Failed to create publication"); + + database + .run_sql(&format!( + "INSERT INTO {} (value) VALUES ('copy_value')", + table_name.as_quoted_identifier(), + )) + .await + .expect("Failed to insert pre-copy row"); + + let clickhouse_db = setup_clickhouse_database().await; + let store = NotifyingStore::new(); + let pipeline_id: PipelineId = random(); + let destination = TestDestinationWrapper::wrap( + clickhouse_db + .build_destination_with_engine(store.clone(), ClickHouseEngine::ReplacingMergeTree) + .await, + ); + + let table_ready = + store.notify_on_table_state_type(table_id, TableReplicationPhaseType::Ready).await; + + let mut pipeline = create_pipeline( + &database.config, + pipeline_id, + publication_name.to_owned(), + store, + destination.clone(), + ); + + pipeline.start().await.unwrap(); + table_ready.notified().await; + + // --- WHEN: stream an UPDATE for the copied row --- + let event_notify = destination.wait_for_events_count(vec![(EventType::Update, 1)]).await; + database + .run_sql(&format!( + "UPDATE {} SET value = 'streamed_value' WHERE id = 1", + table_name.as_quoted_identifier(), + )) + .await + .expect("Failed to update"); + event_notify.notified().await; + + let query = current_state_query( + ClickHouseEngine::ReplacingMergeTree, + "test_copy__then__update", + "id, value", + &["id"], + "id", + ); + let rows: Vec = clickhouse_db.query(&query).await; + + pipeline.shutdown_and_wait().await.unwrap(); + + // --- THEN: the streamed UPDATE wins over the initial-copy row --- + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].value, "streamed_value"); +} + +/// ReplacingMergeTree: a DELETE followed by `OPTIMIZE ... FINAL CLEANUP` +/// physically removes the row from storage. +#[tokio::test(flavor = "multi_thread")] +async fn replacing_merge_tree_optimize_cleanup_physically_removes_tombstoned_row() { + init_test_tracing(); + install_crypto_provider(); + + // --- GIVEN: a copied row, REPLICA IDENTITY FULL --- + let database = spawn_source_database().await; + let table_name = test_table_name("final_cleanup"); + + let table_id = database + .create_table(table_name.clone(), true, &[("value", "text not null")]) + .await + .expect("Failed to create final_cleanup table"); + + database + .run_sql(&format!( + "ALTER TABLE {} REPLICA IDENTITY FULL", + table_name.as_quoted_identifier(), + )) + .await + .expect("Failed to set replica identity full"); + + let publication_name = "test_pub_final_cleanup"; + database + .create_publication(publication_name, std::slice::from_ref(&table_name)) + .await + .expect("Failed to create publication"); + + database + .run_sql(&format!( + "INSERT INTO {} (value) VALUES ('doomed')", + table_name.as_quoted_identifier(), + )) + .await + .expect("Failed to insert seed row"); + + let clickhouse_db = setup_clickhouse_database().await; + let store = NotifyingStore::new(); + let pipeline_id: PipelineId = random(); + let destination = TestDestinationWrapper::wrap( + clickhouse_db + .build_destination_with_engine(store.clone(), ClickHouseEngine::ReplacingMergeTree) + .await, + ); + + let table_ready = + store.notify_on_table_state_type(table_id, TableReplicationPhaseType::Ready).await; + + let mut pipeline = create_pipeline( + &database.config, + pipeline_id, + publication_name.to_owned(), + store, + destination.clone(), + ); + + pipeline.start().await.unwrap(); + table_ready.notified().await; + + let event_notify = destination.wait_for_events_count(vec![(EventType::Delete, 1)]).await; + database + .run_sql(&format!("DELETE FROM {} WHERE id = 1", table_name.as_quoted_identifier())) + .await + .expect("Failed to delete"); + event_notify.notified().await; + + // --- WHEN: operator runs OPTIMIZE ... FINAL CLEANUP --- + let optimize_sql = optimize_final_cleanup_sql("test_final__cleanup"); + let optimize_result = clickhouse_db.db_client().query(&optimize_sql).execute().await; + + pipeline.shutdown_and_wait().await.unwrap(); + + // CH ships ReplacingMergeTree-with-CLEANUP behind an experimental gate that has + // shifted names across versions (26.x removed the previous setting name). + // When the server does not allow the call, treat this test as covered by + // the FINAL-based spine assertions and skip the physical-removal check. + match optimize_result { + Ok(()) => { + let counts: Vec = + clickhouse_db.query("SELECT count() AS count FROM \"test_final__cleanup\"").await; + assert_eq!( + counts[0].count, 0, + "OPTIMIZE FINAL CLEANUP should have physically removed the tombstoned row" + ); + } + Err(err) => { + eprintln!( + "skipping OPTIMIZE FINAL CLEANUP physical-removal assertion: server rejected the \ + statement ({err}); FINAL-based spine tests still cover the user-visible delete \ + semantics" + ); + } + } +} + +/// Engine mismatch (MergeTree -> ReplacingMergeTree): a table already created +/// under MergeTree must hard-fail when a second pipeline tries to write to it +/// under ReplacingMergeTree. +#[tokio::test(flavor = "multi_thread")] +async fn engine_mismatch_existing_merge_tree_then_replacing_merge_tree_pipeline() { + engine_mismatch_runs(ClickHouseEngine::MergeTree, ClickHouseEngine::ReplacingMergeTree).await; +} + +/// Engine mismatch (ReplacingMergeTree -> MergeTree): the reverse direction +/// must also hard-fail. +#[tokio::test(flavor = "multi_thread")] +async fn engine_mismatch_existing_replacing_merge_tree_then_merge_tree_pipeline() { + engine_mismatch_runs(ClickHouseEngine::ReplacingMergeTree, ClickHouseEngine::MergeTree).await; +} + +/// Drives the engine-mismatch flow: pipeline A creates the table under +/// `first`, shuts down; pipeline B tries to use the same ClickHouse database +/// under `second` and the table goes Errored with an engine-mismatch reason. +async fn engine_mismatch_runs(first: ClickHouseEngine, second: ClickHouseEngine) { + init_test_tracing(); + install_crypto_provider(); + + // --- GIVEN: a source table replicated under engine `first` --- + let database = spawn_source_database().await; + let table_name = test_table_name("engine_mismatch"); + + let table_id = database + .create_table(table_name.clone(), true, &[("value", "text not null")]) + .await + .expect("Failed to create engine_mismatch table"); + + let publication_name = "test_pub_engine_mismatch"; + database + .create_publication(publication_name, std::slice::from_ref(&table_name)) + .await + .expect("Failed to create publication"); + + database + .run_sql(&format!( + "INSERT INTO {} (value) VALUES ('first_run')", + table_name.as_quoted_identifier(), + )) + .await + .expect("Failed to insert seed row"); + + let clickhouse_db = setup_clickhouse_database().await; + + // First pipeline run: creates the CH table under `first` and shuts down. + { + let store = NotifyingStore::new(); + let pipeline_id: PipelineId = random(); + let destination = clickhouse_db.build_destination_with_engine(store.clone(), first).await; + let table_ready = + store.notify_on_table_state_type(table_id, TableReplicationPhaseType::Ready).await; + let mut pipeline = create_pipeline( + &database.config, + pipeline_id, + publication_name.to_owned(), + store, + destination, + ); + pipeline.start().await.unwrap(); + table_ready.notified().await; + pipeline.shutdown_and_wait().await.unwrap(); + } + + // WHEN: a second pipeline configures the OTHER engine for the same + // destination database + let store = NotifyingStore::new(); + let pipeline_id: PipelineId = random(); + let destination = clickhouse_db.build_destination_with_engine(store.clone(), second).await; + let table_errored = + store.notify_on_table_state_type(table_id, TableReplicationPhaseType::Errored).await; + + let mut pipeline = create_pipeline( + &database.config, + pipeline_id, + publication_name.to_owned(), + store.clone(), + destination, + ); + pipeline.start().await.unwrap(); + table_errored.notified().await; + pipeline.shutdown_and_wait().await.unwrap(); + + // --- THEN: the Errored phase reason names the engine mismatch --- + let phase = store + .get_table_replication_state(table_id) + .await + .expect("state store should return the table's phase") + .expect("table should have a recorded replication phase"); + let reason = match phase { + TableReplicationPhase::Errored { reason, .. } => reason, + other => panic!("expected Errored phase, got {other:?}"), + }; + assert!( + reason.contains("engine mismatch") || reason.contains("engine"), + "Errored reason should mention the engine mismatch, got: {reason}" + ); + let first_name = first.as_clickhouse_str(); + let second_name = second.as_clickhouse_str(); + assert!( + reason.contains(first_name), + "Errored reason should name the existing engine `{first_name}`: {reason}" + ); + assert!( + reason.contains(second_name), + "Errored reason should name the configured engine `{second_name}`: {reason}" + ); +} diff --git a/crates/etl-destinations/tests/support/clickhouse.rs b/crates/etl-destinations/tests/support/clickhouse.rs index 9933fb9b9..72342b912 100644 --- a/crates/etl-destinations/tests/support/clickhouse.rs +++ b/crates/etl-destinations/tests/support/clickhouse.rs @@ -1,5 +1,19 @@ #![allow(dead_code)] +use std::sync::Once; + +use etl_config::shared::ClickHouseEngine; + +/// Installs the rustls default crypto provider once per process. Subsequent +/// `install_default()` calls return `Err` (already installed); we discard +/// that error so multiple test files can call this safely. +pub(crate) fn install_crypto_provider() { + static INIT_CRYPTO: Once = Once::new(); + INIT_CRYPTO.call_once(|| { + let _ = rustls::crypto::aws_lc_rs::default_provider().install_default(); + }); +} + /// A row read back from the ClickHouse `all_types_encoding` test table. /// /// Column-to-type mapping: @@ -35,7 +49,6 @@ pub(crate) struct AllTypesRow { pub cidr_col: String, pub macaddr_col: String, pub uuid_col: String, // via toString() in SELECT - pub cdc_operation: String, } /// A row read back from the ClickHouse `boundary_values` test table. @@ -49,7 +62,6 @@ pub(crate) struct BoundaryValuesRow { pub nullable_int: Option, pub int_array_col: Vec>, pub text_array_col: Vec>, - pub cdc_operation: String, } /// A row read back from a ClickHouse table with a single `Date32` column, @@ -60,5 +72,54 @@ pub(crate) struct BoundaryValuesRow { pub(crate) struct DateBoundariesRow { pub id: i64, pub date_col: i32, - pub cdc_operation: String, +} + +/// Builds an engine-aware "current state" SELECT for a replicated table. +/// +/// The projection is supplied by the caller (so per-column SQL like +/// `toString(uuid_col) AS uuid_col` keeps working). The helper handles the +/// engine-specific dedup + tombstone filter and applies the caller's +/// `ORDER BY` for deterministic test reads. +/// +/// MergeTree path: take the latest event per PK with `LIMIT 1 BY`, then drop +/// any whose latest event is a DELETE. The drop-DELETE filter must come +/// AFTER the dedup, otherwise a deleted PK whose latest event is a DELETE +/// would surface its prior INSERT instead of being absent. +/// +/// ReplacingMergeTree path: `FINAL` + `_etl_deleted = 0`. +pub(crate) fn current_state_query( + engine: ClickHouseEngine, + table: &str, + projection: &str, + pk_cols: &[&str], + order_by: &str, +) -> String { + match engine { + ClickHouseEngine::MergeTree => format!( + "SELECT {projection} FROM (SELECT * FROM \"{table}\" ORDER BY cdc_lsn DESC LIMIT 1 BY \ + ({pks})) AS current WHERE cdc_operation != 'DELETE' ORDER BY {order_by}", + pks = pk_cols.join(", ") + ), + ClickHouseEngine::ReplacingMergeTree => format!( + "SELECT {projection} FROM \"{table}\" FINAL WHERE _etl_deleted = 0 ORDER BY {order_by}" + ), + } +} + +/// SQL to force a `ReplacingMergeTree` table to drop tombstoned rows. The +/// `SETTINGS` clause enables the (still experimental, as of CH 25.x) merges- +/// with-cleanup feature for this query only, without requiring it to be +/// enabled server-wide. +pub(crate) fn optimize_final_cleanup_sql(table: &str) -> String { + format!( + "OPTIMIZE TABLE \"{table}\" FINAL CLEANUP SETTINGS \ + allow_experimental_replacing_merge_tree_with_cleanup = 1" + ) +} + +/// SQL to read the `engine` column from `system.tables` for a table. +pub(crate) fn table_engine_query(table: &str) -> String { + format!( + "SELECT engine FROM system.tables WHERE database = currentDatabase() AND name = '{table}'" + ) } diff --git a/crates/etl-examples/Cargo.toml b/crates/etl-examples/Cargo.toml index f141c6294..0dd742f0d 100644 --- a/crates/etl-examples/Cargo.toml +++ b/crates/etl-examples/Cargo.toml @@ -36,6 +36,7 @@ ducklake = ["etl-destinations/ducklake"] clap = { workspace = true, default-features = true, features = ["std", "derive"] } etl = { workspace = true } +etl-config = { workspace = true } etl-destinations = { workspace = true } etl-telemetry = { workspace = true } k8s-openapi = { workspace = true, features = ["latest"] } diff --git a/crates/etl-examples/README.md b/crates/etl-examples/README.md index 8ec6ce9e5..a8afb2112 100644 --- a/crates/etl-examples/README.md +++ b/crates/etl-examples/README.md @@ -121,6 +121,7 @@ them to absolute `file://` URLs before constructing the destination. ## ClickHouse Setup To run the ClickHouse example, you'll need a running ClickHouse instance accessible over HTTP(S). +ClickHouse **23.5 or newer** is required for the default `ReplacingMergeTree` engine. Create a publication in Postgres: @@ -144,16 +145,79 @@ cargo run -p etl-examples --bin clickhouse --features clickhouse -- \ --publication my_pub ``` -Each Postgres table is replicated as an append-only `MergeTree` table. Two CDC metadata -columns are appended to every row: +### Table engines -- `cdc_operation`: `INSERT`, `UPDATE`, or `DELETE` -- `cdc_lsn`: the Postgres LSN at the time of the change +The destination supports two layouts, chosen per pipeline via `--clickhouse-engine`: + +| Flag value | Engine | Use it for | +|----------------------------------|----------------------|---------------------------------------------------------| +| `replacing_merge_tree` (default) | `ReplacingMergeTree` | Current-state replicas. Source must have a primary key. | +| `merge_tree` | `MergeTree` | Append-only event log. Works for PK-less source tables. | Table names are derived from the Postgres schema and table name using double-underscore -escaping (e.g. `public.orders` → `public_orders`, `my_schema.t` → `my__schema_t`). +escaping (e.g. `public.orders` -> `public_orders`, `my_schema.t` -> `my__schema_t`). + +#### ReplacingMergeTree (default) + +Each replicated table is created as `ReplacingMergeTree(_etl_version, _etl_deleted)` keyed +on the source primary key. Two trailing columns drive dedup and tombstone handling: + +- `_etl_version UInt128` -- the packed Postgres event sequence key: + `(commit_lsn << 64) | tx_ordinal`. Higher values win during a `FINAL` merge, so the + latest event per primary key wins. Encoding both the commit LSN and the in-transaction + ordinal gives a total order across all events, including multiple row events that + share a WAL record. +- `_etl_deleted UInt8` -- tombstone flag. `1` for DELETE events, `0` otherwise. + +Alongside each table, the destination also creates a `
__current` view that hides +the ReplacingMergeTree internals: + +```sql +CREATE VIEW IF NOT EXISTS "public_orders__current" AS +SELECT +FROM "public_orders" FINAL +WHERE _etl_deleted = 0 +``` + +Read patterns: + +- Prefer the `__current` view for current-state queries. +- Or query the base table with `SELECT ... FROM "public_orders" FINAL WHERE _etl_deleted = 0` + directly. + +`OPTIMIZE` guidance: + +- The replicator never runs `OPTIMIZE ... FINAL CLEANUP`. Background merges already + collapse duplicates over time; physical removal of tombstones is operator-driven. +- To reclaim deleted rows on disk, run `OPTIMIZE TABLE "
" FINAL CLEANUP` on a + schedule that matches your retention requirements. + +#### MergeTree + +Each replicated table is created as `MergeTree() ORDER BY tuple()` with two CDC metadata +columns appended to every row: + +- `cdc_operation`: `INSERT`, `UPDATE`, or `DELETE` +- `cdc_lsn`: the Postgres commit LSN at the time of the change + +Read patterns: + +- Current state per primary key: take the latest event by `cdc_lsn` with `LIMIT 1 BY`, + then filter out tombstones. Example: + + ```sql + SELECT FROM ( + SELECT * FROM "public_orders" + ORDER BY cdc_lsn DESC LIMIT 1 BY (id) + ) + WHERE cdc_operation != 'DELETE' + ``` + +- Event log queries: read the table directly; every CDC event is preserved. + +### Connection notes -For HTTPS connections, provide an `https://` URL — TLS is handled automatically using +For HTTPS connections, provide an `https://` URL -- TLS is handled automatically using webpki root certificates. Use `--clickhouse-password` if your ClickHouse instance requires authentication. diff --git a/crates/etl-examples/src/bin/clickhouse.rs b/crates/etl-examples/src/bin/clickhouse.rs index 3043a6114..8a4d7d999 100644 --- a/crates/etl-examples/src/bin/clickhouse.rs +++ b/crates/etl-examples/src/bin/clickhouse.rs @@ -5,18 +5,25 @@ ClickHouse Example This example demonstrates how to use the pipeline to stream data from Postgres to ClickHouse using change data capture (CDC). -Each Postgres table is replicated as an append-only MergeTree table. -Two CDC metadata columns are appended to every row: - - `cdc_operation`: `INSERT`, `UPDATE`, or `DELETE` - - `cdc_lsn`: the Postgres LSN at the time of the change +Two table-engine layouts are supported, selected via `--clickhouse-engine`: + +- `replacing_merge_tree` (default): each replicated table becomes a + `ReplacingMergeTree` keyed on the source primary key, with trailing + `_etl_version` (UInt128 packed `(commit_lsn, tx_ordinal)`) and + `_etl_deleted` (tombstone) columns. A companion `
__current` view + reads current state via `FINAL` and filters tombstones. Requires + ClickHouse >= 23.5 and a primary key on the source. +- `merge_tree`: append-only event-log layout with `cdc_operation` and + `cdc_lsn` columns appended to every row. Works for PK-less source tables. Table names are derived from the Postgres schema and table name using -double-underscore escaping (e.g. `public.orders` → `public__orders`). +double-underscore escaping (e.g. `public.orders` -> `public_orders`). Prerequisites: 1. Postgres server with logical replication enabled (wal_level = logical) 2. A publication created in Postgres (CREATE PUBLICATION my_pub FOR ALL TABLES;) -3. A running ClickHouse instance accessible over HTTP(S) +3. A running ClickHouse instance accessible over HTTP(S). ReplacingMergeTree additionally + requires CH >= 23.5. Usage: cargo run -p etl-examples --bin clickhouse -- \ @@ -30,7 +37,7 @@ Usage: --clickhouse-database default \ --publication my_pub -For HTTPS connections, provide an `https://` URL — TLS is handled automatically +For HTTPS connections, provide an `https://` URL -- TLS is handled automatically using webpki root certificates. Use `--clickhouse-password` if your ClickHouse instance requires authentication. @@ -47,6 +54,7 @@ use etl::{ pipeline::Pipeline, store::PostgresStore, }; +use etl_config::shared::ClickHouseEngine; use etl_destinations::clickhouse::{ ClickHouseClientConfig, ClickHouseDestination, ClickHouseInserterConfig, }; @@ -118,6 +126,11 @@ struct ClickHouseArgs { /// ClickHouse target database #[arg(long)] clickhouse_database: String, + /// Table engine used for replicated tables. `replacing_merge_tree` is the + /// default and requires a source primary key and CH >= 23.5; `merge_tree` + /// gives the append-only event-log layout. + #[arg(long, value_enum, default_value_t = ClickHouseEngineArg::ReplacingMergeTree)] + clickhouse_engine: ClickHouseEngineArg, /// Maximum time to wait for a batch to fill in milliseconds (lower values = /// lower latency, less throughput) #[arg(long, default_value = "5000")] @@ -128,6 +141,23 @@ struct ClickHouseArgs { max_table_sync_workers: u16, } +/// CLI-facing engine choice. Converts to `ClickHouseEngine` via `From`. +#[derive(Debug, Copy, Clone, clap::ValueEnum)] +#[clap(rename_all = "snake_case")] +enum ClickHouseEngineArg { + MergeTree, + ReplacingMergeTree, +} + +impl From for ClickHouseEngine { + fn from(arg: ClickHouseEngineArg) -> Self { + match arg { + ClickHouseEngineArg::MergeTree => ClickHouseEngine::MergeTree, + ClickHouseEngineArg::ReplacingMergeTree => ClickHouseEngine::ReplacingMergeTree, + } + } +} + /// Entry point — handles error reporting and process exit. #[tokio::main] async fn main() -> Result<(), Box> { @@ -201,17 +231,19 @@ async fn main_impl() -> Result<(), Box> { max_copy_connections_per_table: PipelineConfig::DEFAULT_MAX_COPY_CONNECTIONS_PER_TABLE, }; - // Initialize the ClickHouse destination. - // Tables are created automatically as append-only MergeTree tables. let clickhouse_destination = ClickHouseDestination::new( Url::parse(&args.clickhouse_args.clickhouse_url)?, args.clickhouse_args.clickhouse_user, args.clickhouse_args.clickhouse_password, args.clickhouse_args.clickhouse_database, - ClickHouseInserterConfig::default(), + ClickHouseInserterConfig { + engine: args.clickhouse_args.clickhouse_engine.into(), + ..Default::default() + }, ClickHouseClientConfig::default(), store.clone(), )?; + clickhouse_destination.validate_engine_support().await?; let mut pipeline = Pipeline::new(pipeline_config, store, clickhouse_destination); diff --git a/crates/etl-replicator/src/core.rs b/crates/etl-replicator/src/core.rs index eb16a3986..782a2657d 100644 --- a/crates/etl-replicator/src/core.rs +++ b/crates/etl-replicator/src/core.rs @@ -206,16 +206,17 @@ pub(crate) async fn start_replicator_with_config( let pipeline = Pipeline::new(replicator_config.pipeline, state_store, destination); start_pipeline(pipeline).await?; } - DestinationConfig::ClickHouse { url, user, password, database } => { + DestinationConfig::ClickHouse { url, user, password, database, engine } => { let destination = ClickHouseDestination::new( url.clone(), user, password.as_ref().map(|p| p.expose_secret().to_owned()), database, - ClickHouseInserterConfig::default(), + ClickHouseInserterConfig { engine: *engine, ..Default::default() }, ClickHouseClientConfig::default(), state_store.clone(), )?; + destination.validate_engine_support().await?; let pipeline = Pipeline::new(replicator_config.pipeline, state_store, destination); start_pipeline(pipeline).await?; diff --git a/crates/etl/src/types/event.rs b/crates/etl/src/types/event.rs index d72aebe77..498c5c000 100644 --- a/crates/etl/src/types/event.rs +++ b/crates/etl/src/types/event.rs @@ -326,6 +326,14 @@ impl EventSequenceKey { pub fn new(commit_lsn: PgLsn, tx_ordinal: u64) -> Self { Self { commit_lsn, tx_ordinal } } + + /// Returns the canonical packed `u128` form: `commit_lsn` in the high + /// 64 bits, `tx_ordinal` in the low 64 bits. Used by destinations that + /// need a single totally-ordered numeric key for CDC dedup (e.g. + /// ClickHouse's `ReplacingMergeTree` version column). + pub fn as_u128(self) -> u128 { + (u128::from(u64::from(self.commit_lsn)) << 64) | u128::from(self.tx_ordinal) + } } impl fmt::Display for EventSequenceKey { From ae9c7378fb1e3a75baeaeefe1112bc0d57841cd7 Mon Sep 17 00:00:00 2001 From: Jordan McQueen Date: Wed, 20 May 2026 19:49:00 +0900 Subject: [PATCH 17/29] style(tests): drop 'then' from WHEN comments to avoid GIVEN/WHEN/THEN ambiguity (#756) --- crates/etl-destinations/tests/clickhouse/pipeline.rs | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/crates/etl-destinations/tests/clickhouse/pipeline.rs b/crates/etl-destinations/tests/clickhouse/pipeline.rs index 2084ba66f..655edf010 100644 --- a/crates/etl-destinations/tests/clickhouse/pipeline.rs +++ b/crates/etl-destinations/tests/clickhouse/pipeline.rs @@ -342,7 +342,7 @@ async fn updates_are_streamed_to_clickhouse_inner(engine: ClickHouseEngine) { .await .expect("Failed to insert initial update_flow row"); - // --- WHEN: pipeline copies data, then an UPDATE is streamed --- + // --- WHEN: pipeline copies data and an UPDATE is streamed --- let clickhouse_db = setup_clickhouse_database().await; let store = NotifyingStore::new(); let pipeline_id: PipelineId = random(); @@ -760,7 +760,7 @@ async fn deletes_are_streamed_to_clickhouse_inner(engine: ClickHouseEngine) { .await .expect("Failed to insert delete_flow rows"); - // --- WHEN: pipeline copies data, then a DELETE is streamed --- + // --- WHEN: pipeline copies data and a DELETE is streamed --- let clickhouse_db = setup_clickhouse_database().await; let store = NotifyingStore::new(); @@ -891,7 +891,7 @@ async fn pipeline_restart_resumes_streaming_inner(engine: ClickHouseEngine) { assert_eq!(rows[0].id, 1); assert_eq!(rows[0].value, "before_restart"); - // --- WHEN: rebuild destination and pipeline, then stream a new insert --- + // --- WHEN: rebuild destination and pipeline and stream a new insert --- let destination = TestDestinationWrapper::wrap( clickhouse_db.build_destination_with_engine(store.clone(), engine).await, ); @@ -996,7 +996,7 @@ async fn truncate_clears_table_and_accepts_new_inserts_inner(engine: ClickHouseE let rows: Vec = clickhouse_db.query(&truncate_query()).await; assert_eq!(rows.len(), 2, "table copy should produce two rows"); - // --- WHEN: truncate, then insert a new row --- + // --- WHEN: truncate and insert a new row --- let truncate_notify = destination.wait_for_events_count(vec![(EventType::Truncate, 1)]).await; database @@ -1678,7 +1678,7 @@ async fn schema_change_add_column_inner(engine: ClickHouseEngine) { .expect("metadata should exist after table creation"); let initial_snapshot_id = initial_metadata.snapshot_id; - // --- WHEN: add column, then insert with new schema --- + // --- WHEN: add column and insert with new schema --- database .alter_table( table_name.clone(), @@ -1891,7 +1891,7 @@ async fn schema_change_add_drop_rename_inner(engine: ClickHouseEngine) { .expect("metadata should exist after table creation"); let initial_snapshot_id = initial_metadata.snapshot_id; - // --- WHEN: rename + drop + add, then insert with new schema --- + // --- WHEN: rename + drop + add and insert with new schema --- database .alter_table( table_name.clone(), From af208ac76d8810fe9f107b6136e60a3d9be66dde Mon Sep 17 00:00:00 2001 From: Jordan McQueen Date: Wed, 20 May 2026 19:49:25 +0900 Subject: [PATCH 18/29] test(clickhouse): add large-row JSONB pipeline tests (#754) --- .../etl-destinations/tests/clickhouse/mod.rs | 1 + .../tests/clickhouse/pipeline_large_rows.rs | 227 ++++++++++++++++++ 2 files changed, 228 insertions(+) create mode 100644 crates/etl-destinations/tests/clickhouse/pipeline_large_rows.rs diff --git a/crates/etl-destinations/tests/clickhouse/mod.rs b/crates/etl-destinations/tests/clickhouse/mod.rs index f5fde1c39..5fce53497 100644 --- a/crates/etl-destinations/tests/clickhouse/mod.rs +++ b/crates/etl-destinations/tests/clickhouse/mod.rs @@ -1,3 +1,4 @@ mod pipeline; +mod pipeline_large_rows; mod pipeline_merge_tree; mod pipeline_replacing_merge_tree; diff --git a/crates/etl-destinations/tests/clickhouse/pipeline_large_rows.rs b/crates/etl-destinations/tests/clickhouse/pipeline_large_rows.rs new file mode 100644 index 000000000..44214200c --- /dev/null +++ b/crates/etl-destinations/tests/clickhouse/pipeline_large_rows.rs @@ -0,0 +1,227 @@ +//! Tests that exercise the Postgres -> ClickHouse pipeline with very large +//! single-row JSONB payloads. +//! +//! Each row is a two-field JSON object: `{"x":"aaa...","y":"aaa..."}`. The +//! field values are built server-side with PostgreSQL's `repeat('a', N)`, +//! and on read-back the row is summarised in ClickHouse into a small +//! `(id, len, a_count, head, tail)` tuple -- the payload bytes never leave +//! the database. That keeps both the SQL we send and test-process memory +//! tiny no matter how large the row is. + +use etl::{ + state::table::TableReplicationPhaseType, + test_utils::{ + database::{spawn_source_database, test_table_name}, + notifying_store::NotifyingStore, + pipeline::create_pipeline, + test_destination_wrapper::TestDestinationWrapper, + }, + types::{EventType, PipelineId, TableName}, +}; +use etl_config::shared::ClickHouseEngine; +use etl_destinations::clickhouse::test_utils::setup_clickhouse_database; +use etl_telemetry::tracing::init_test_tracing; +use rand::random; + +use crate::support::clickhouse::{current_state_query, install_crypto_provider}; + +const KIB: usize = 1024; +const MIB: usize = 1024 * 1024; + +/// ClickHouse-side table name after `try_stringify_table_name` escaping of +/// `test.large_rows`. +const LARGE_ROW_TABLE: &str = "test_large__rows"; + +const PUBLICATION_NAME: &str = "test_pub_ch_large_rows"; + +/// Byte count of the JSON object framing as ClickHouse receives it. +/// +/// Postgres `jsonb_out` emits a canonical form with spaces after `:` and `,` +/// (e.g. `{"x": "", "y": ""}` = 18 bytes), but ETL parses jsonb through +/// `serde_json` (see `crates/etl/src/conversions/text.rs`) and re-serialises +/// in compact form (`{"x":"","y":""}` = 15 bytes). The destination only ever +/// sees the compact representation, so byte-exact assertions key off 15. +const PAYLOAD_OVERHEAD: usize = 15; + +/// First 16 bytes of the payload as ClickHouse sees it: 6 framing bytes +/// (`{"x":"`) followed by 10 `'a'`s from the first field value. +const EXPECTED_HEAD: &str = r#"{"x":"aaaaaaaaaa"#; + +/// Last 16 bytes of the payload: 14 `'a'`s from the second field value +/// followed by the 2-byte closer `"}`. +const EXPECTED_TAIL: &str = r#"aaaaaaaaaaaaaa"}"#; + +/// ClickHouse-side projection summarising each row server-side. The +/// `length` / `countSubstrings` aggregates and the boundary `substring`s are +/// all bounded, so the response stays tiny regardless of payload size. +const SUMMARY_PROJECTION: &str = concat!( + "id, ", + "toUInt64(length(payload)) AS len, ", + "toUInt64(countSubstrings(payload, 'a')) AS a_count, ", + "substring(payload, 1, 16) AS head, ", + "substring(payload, length(payload) - 15, 16) AS tail", +); + +/// One row read back from ClickHouse, summarised entirely server-side so the +/// payload never crosses the wire to the test process. +#[derive(clickhouse::Row, serde::Deserialize, Debug)] +struct LargeRowSummary { + id: i64, + len: u64, + a_count: u64, + head: String, + tail: String, +} + +/// Builds the INSERT for one row of `{"x":"","y":""}`. The +/// large value is constructed server-side via `repeat('a', N)` so the SQL we +/// send stays small. +/// +/// Produces SQL of the shape: +/// +/// ```sql +/// INSERT INTO "test"."large_rows" (payload) VALUES ( +/// ('{"x":"' || repeat('a', N) || '","y":"' || repeat('a', N) || '"}')::jsonb +/// ) +/// ``` +fn insert_two_field_payload_sql(table_name: &TableName, field_chars: usize) -> String { + format!( + r#"INSERT INTO {table} (payload) VALUES (('{{"x":"' || repeat('a', {n}) || '","y":"' || repeat('a', {n}) || '"}}')::jsonb)"#, + table = table_name.as_quoted_identifier(), + n = field_chars, + ) +} + +/// Drives one large-row test for the given engine and per-field byte count. +/// +/// # GIVEN +/// +/// A Postgres table `(id bigserial pk, payload jsonb not null)`. Row 1 is +/// inserted before the pipeline starts (initial-copy path) and row 2 is +/// inserted after the pipeline is ready (streaming INSERT-event path). +/// +/// # THEN +/// +/// Both rows arrive in ClickHouse with `length(payload) == +/// PAYLOAD_OVERHEAD + 2 * field_chars`, the expected count of `'a'` +/// characters (`2 * field_chars`; neither key contains an `'a'`), and +/// matching head / tail boundary slices. +async fn large_row_inner(engine: ClickHouseEngine, field_chars: usize) { + init_test_tracing(); + install_crypto_provider(); + + // --- GIVEN: Postgres source with one pre-pipeline large JSONB row --- + let database = spawn_source_database().await; + let table_name = test_table_name("large_rows"); + + let table_id = database + .create_table(table_name.clone(), true, &[("payload", "jsonb not null")]) + .await + .expect("Failed to create large_rows table"); + + database + .create_publication(PUBLICATION_NAME, std::slice::from_ref(&table_name)) + .await + .expect("Failed to create large_rows publication"); + + // Row 1 exercises the initial table-copy path. + database + .run_sql(&insert_two_field_payload_sql(&table_name, field_chars)) + .await + .expect("Failed to insert large JSONB row"); + + // --- WHEN: pipeline copies row 1 and row 2 streams as an INSERT event --- + let clickhouse_db = setup_clickhouse_database().await; + let store = NotifyingStore::new(); + let pipeline_id: PipelineId = random(); + let destination = TestDestinationWrapper::wrap( + clickhouse_db.build_destination_with_engine(store.clone(), engine).await, + ); + + let table_ready = + store.notify_on_table_state_type(table_id, TableReplicationPhaseType::Ready).await; + + let mut pipeline = create_pipeline( + &database.config, + pipeline_id, + PUBLICATION_NAME.to_owned(), + store, + destination.clone(), + ); + + pipeline.start().await.unwrap(); + table_ready.notified().await; + + let event_notify = destination.wait_for_events_count(vec![(EventType::Insert, 1)]).await; + + // Row 2 exercises the streaming INSERT-event path. + database + .run_sql(&insert_two_field_payload_sql(&table_name, field_chars)) + .await + .expect("Failed to insert large JSONB row"); + + event_notify.notified().await; + + let query = current_state_query(engine, LARGE_ROW_TABLE, SUMMARY_PROJECTION, &["id"], "id"); + let rows: Vec = clickhouse_db.query(&query).await; + + pipeline.shutdown_and_wait().await.unwrap(); + + // --- THEN: both rows arrived intact and byte-exact --- + assert_eq!(rows.len(), 2, "expected 2 current-state rows (copy + streaming)"); + + let expected_len = (PAYLOAD_OVERHEAD + 2 * field_chars) as u64; + let expected_a_count = (2 * field_chars) as u64; + + for row in &rows { + assert_eq!( + row.len, expected_len, + "row id={} payload length {} != expected {expected_len}", + row.id, row.len, + ); + assert_eq!( + row.a_count, expected_a_count, + "row id={} count of 'a' bytes {} != expected {expected_a_count}", + row.id, row.a_count, + ); + assert_eq!(row.head, EXPECTED_HEAD, "row id={} payload head mismatch", row.id); + assert_eq!(row.tail, EXPECTED_TAIL, "row id={} payload tail mismatch", row.id); + } +} + +#[tokio::test(flavor = "multi_thread")] +async fn large_row_1mib_merge_tree() { + large_row_inner(ClickHouseEngine::MergeTree, 512 * KIB).await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn large_row_1mib_replacing_merge_tree() { + large_row_inner(ClickHouseEngine::ReplacingMergeTree, 512 * KIB).await; +} + +/// 16 MiB exceeds `BatchConfig::DEFAULT_MAX_BYTES` (8 MiB), so the batch +/// stream must flush a single oversized row on its own. +#[tokio::test(flavor = "multi_thread")] +async fn large_row_16mib_merge_tree() { + large_row_inner(ClickHouseEngine::MergeTree, 8 * MIB).await; +} + +/// 64 MiB exceeds `ClickHouseInserterConfig::DEFAULT_MAX_BYTES_PER_INSERT` +/// (64 MiB), validating the one-row-per-INSERT path in `insert_rows`. +#[tokio::test(flavor = "multi_thread")] +async fn large_row_64mib_merge_tree() { + large_row_inner(ClickHouseEngine::MergeTree, 32 * MIB).await; +} + +/// 256 MiB stress test with two ~128 MiB fields. Each row allocates ~256 MiB +/// peak in Postgres, ETL, and ClickHouse roughly concurrently. +/// +/// `134_217_716` is the largest per-field byte count that still fits within +/// Postgres's JSONB object element-size ceiling (`2^28 - 1 = 268_435_455` +/// bytes; see `convertJsonbObject` in `jsonb_util.c`). Empirically, +/// `134_217_717` rejects and `134_217_716` accepts, so the materialised row +/// lands ~9 bytes shy of a clean 256 MiB. +#[tokio::test(flavor = "multi_thread")] +async fn large_row_256mib_merge_tree() { + large_row_inner(ClickHouseEngine::MergeTree, 134_217_716).await; +} From 77cafc6fc2d14f55c3afe8de37b37e84633fe7b6 Mon Sep 17 00:00:00 2001 From: Coenen Benjamin Date: Wed, 20 May 2026 13:41:20 +0200 Subject: [PATCH 19/29] fix(api): do not create k8s labels bigger than 63 characters (#755) * fix(api): do not create k8s labels bigger than 63 characters Signed-off-by: Benjamin <5719034+bnjjj@users.noreply.github.com> * change suffix Signed-off-by: Benjamin <5719034+bnjjj@users.noreply.github.com> --------- Signed-off-by: Benjamin <5719034+bnjjj@users.noreply.github.com> --- crates/etl-api/src/k8s/base.rs | 3 + crates/etl-api/src/k8s/cache.rs | 4 + crates/etl-api/src/k8s/core.rs | 58 +++- crates/etl-api/src/k8s/http.rs | 268 +++++++++++++++++- ...ate_bq_replicator_stateful_set_json-2.snap | 4 +- ...ate_bq_replicator_stateful_set_json-3.snap | 4 +- ...reate_bq_replicator_stateful_set_json.snap | 4 +- ...cklake_replicator_stateful_set_json-2.snap | 4 +- ...cklake_replicator_stateful_set_json-3.snap | 4 +- ...ducklake_replicator_stateful_set_json.snap | 5 +- ...ceberg_replicator_stateful_set_json-2.snap | 4 +- ...ceberg_replicator_stateful_set_json-3.snap | 4 +- ..._iceberg_replicator_stateful_set_json.snap | 4 +- crates/etl-api/src/routes/mod.rs | 54 ++++ crates/etl-api/src/routes/pipelines.rs | 14 +- crates/etl-api/src/routes/tenants.rs | 11 +- crates/etl-api/src/routes/tenants_sources.rs | 7 +- .../etl-api/tests/destinations_pipelines.rs | 16 +- crates/etl-api/tests/pipelines.rs | 16 +- crates/etl-api/tests/support/k8s_client.rs | 4 + 20 files changed, 439 insertions(+), 53 deletions(-) diff --git a/crates/etl-api/src/k8s/base.rs b/crates/etl-api/src/k8s/base.rs index ad58892e0..4ebdab5f3 100644 --- a/crates/etl-api/src/k8s/base.rs +++ b/crates/etl-api/src/k8s/base.rs @@ -275,6 +275,9 @@ pub trait K8sClient: Send + Sync { /// Does nothing if the stateful set does not exist. async fn delete_replicator_stateful_set(&self, prefix: &str) -> Result<(), K8sError>; + /// Returns whether the replicator [`StatefulSet`] exists. + async fn replicator_stateful_set_exists(&self, prefix: &str) -> Result; + /// Creates or updates the DuckLake maintenance CR. async fn create_or_update_ducklake_maintenance( &self, diff --git a/crates/etl-api/src/k8s/cache.rs b/crates/etl-api/src/k8s/cache.rs index 68161be3e..e8fb4c373 100644 --- a/crates/etl-api/src/k8s/cache.rs +++ b/crates/etl-api/src/k8s/cache.rs @@ -236,6 +236,10 @@ mod tests { Ok(()) } + async fn replicator_stateful_set_exists(&self, _prefix: &str) -> Result { + Ok(false) + } + async fn create_or_update_ducklake_maintenance( &self, _prefix: &str, diff --git a/crates/etl-api/src/k8s/core.rs b/crates/etl-api/src/k8s/core.rs index 2cda97dd7..d05a9d52d 100644 --- a/crates/etl-api/src/k8s/core.rs +++ b/crates/etl-api/src/k8s/core.rs @@ -213,6 +213,25 @@ pub async fn is_replicator_pod_stopped( Ok(matches!(pod_status, PodStatus::Stopped)) } +/// Returns `true` if existing Kubernetes resources should be reconciled. +/// +/// A stopped pod normally means the pipeline is intentionally inactive. A +/// StatefulSet without a pod means Kubernetes still has desired runtime state, +/// so reconciliation should repair or migrate it. +pub async fn should_reconcile_replicator_resources( + k8s_client: &dyn K8sClient, + tenant_id: &str, + replicator_id: i64, +) -> Result { + let prefix = create_k8s_object_prefix(tenant_id, replicator_id); + let pod_status = k8s_client.get_replicator_pod_status(&prefix).await?; + if !matches!(pod_status, PodStatus::Stopped) { + return Ok(true); + } + + Ok(k8s_client.replicator_stateful_set_exists(&prefix).await?) +} + /// Returns `true` when the replicator is active in Kubernetes. pub async fn is_replicator_active( k8s_client: &dyn K8sClient, @@ -535,6 +554,7 @@ mod tests { struct RecordingK8sClient { calls: Arc>>, pod_status: PodStatus, + stateful_set_exists: bool, } impl RecordingK8sClient { @@ -545,7 +565,11 @@ mod tests { impl Default for RecordingK8sClient { fn default() -> Self { - Self { calls: Arc::default(), pod_status: PodStatus::Stopped } + Self { + calls: Arc::default(), + pod_status: PodStatus::Stopped, + stateful_set_exists: false, + } } } @@ -673,6 +697,10 @@ mod tests { Ok(()) } + async fn replicator_stateful_set_exists(&self, _prefix: &str) -> Result { + Ok(self.stateful_set_exists) + } + async fn create_or_update_ducklake_maintenance( &self, prefix: &str, @@ -750,6 +778,34 @@ mod tests { assert!(is_active); } + #[tokio::test] + async fn stopped_pod_with_stateful_set_is_reconciled() { + let client = RecordingK8sClient { + pod_status: PodStatus::Stopped, + stateful_set_exists: true, + ..Default::default() + }; + + let should_reconcile = + should_reconcile_replicator_resources(&client, "tenant-42", 4).await.unwrap(); + + assert!(should_reconcile); + } + + #[tokio::test] + async fn stopped_pod_without_stateful_set_is_not_reconciled() { + let client = RecordingK8sClient { + pod_status: PodStatus::Stopped, + stateful_set_exists: false, + ..Default::default() + }; + + let should_reconcile = + should_reconcile_replicator_resources(&client, "tenant-42", 4).await.unwrap(); + + assert!(!should_reconcile); + } + #[tokio::test] async fn clickhouse_with_password_creates_password_secret() { let source_config = source_config_with_password(); diff --git a/crates/etl-api/src/k8s/http.rs b/crates/etl-api/src/k8s/http.rs index d6615cc07..06330b6b9 100644 --- a/crates/etl-api/src/k8s/http.rs +++ b/crates/etl-api/src/k8s/http.rs @@ -62,7 +62,9 @@ const POSTGRES_SECRET_NAME_SUFFIX: &str = "postgres-password"; /// ConfigMap name suffix for the replicator configuration files. const REPLICATOR_CONFIG_MAP_NAME_SUFFIX: &str = "replicator-config"; /// StatefulSet name suffix for the replicator workload. -const REPLICATOR_STATEFUL_SET_SUFFIX: &str = "replicator-stateful-set"; +const REPLICATOR_STATEFUL_SET_SUFFIX: &str = "replicator"; +/// Previous StatefulSet suffix kept for existing pipeline cleanup/status. +const LEGACY_REPLICATOR_STATEFUL_SET_SUFFIX: &str = "replicator-stateful-set"; /// Application label suffix used to group resources. const REPLICATOR_APP_SUFFIX: &str = "replicator-app"; /// Container name suffix for the replicator container. @@ -589,6 +591,13 @@ impl K8sClient for HttpK8sClient { )?; let stateful_set_name = create_stateful_set_name(prefix); + let legacy_stateful_set_name = create_legacy_stateful_set_name(prefix); + if legacy_stateful_set_name != stateful_set_name { + let dp = DeleteParams::default(); + Self::handle_delete_with_404_ignore( + self.stateful_sets_api.delete(&legacy_stateful_set_name, &dp).await, + )?; + } let container_environment = create_container_environment_json( prefix, @@ -630,15 +639,30 @@ impl K8sClient for HttpK8sClient { async fn delete_replicator_stateful_set(&self, prefix: &str) -> Result<(), K8sError> { debug!("deleting stateful set"); - let stateful_set_name = create_stateful_set_name(prefix); let dp = DeleteParams::default(); - Self::handle_delete_with_404_ignore( - self.stateful_sets_api.delete(&stateful_set_name, &dp).await, - )?; + for stateful_set_name in stateful_set_names_for_lookup(prefix) { + Self::handle_delete_with_404_ignore( + self.stateful_sets_api.delete(&stateful_set_name, &dp).await, + )?; + } Ok(()) } + async fn replicator_stateful_set_exists(&self, prefix: &str) -> Result { + debug!("checking stateful set existence"); + + for stateful_set_name in stateful_set_names_for_lookup(prefix) { + match self.stateful_sets_api.get(&stateful_set_name).await { + Ok(_) => return Ok(true), + Err(kube::Error::Api(er)) if er.code == 404 => {} + Err(e) => return Err(e.into()), + } + } + + Ok(false) + } + async fn create_or_update_ducklake_maintenance( &self, prefix: &str, @@ -671,11 +695,19 @@ impl K8sClient for HttpK8sClient { async fn get_replicator_pod_status(&self, prefix: &str) -> Result { debug!("getting pod status"); - let pod_name = create_pod_name(prefix); - let pod = match self.pods_api.get(&pod_name).await { - Ok(pod) => pod, - Err(kube::Error::Api(er)) if er.code == 404 => return Ok(PodStatus::Stopped), - Err(e) => return Err(e.into()), + let mut pod = None; + for pod_name in pod_names_for_status(prefix) { + match self.pods_api.get(&pod_name).await { + Ok(found_pod) => { + pod = Some(found_pod); + break; + } + Err(kube::Error::Api(er)) if er.code == 404 => {} + Err(e) => return Err(e.into()), + } + } + let Some(pod) = pod else { + return Ok(PodStatus::Stopped); }; let replicator_container_name = create_replicator_container_name(prefix); @@ -738,10 +770,33 @@ fn create_stateful_set_name(prefix: &str) -> String { format!("{prefix}-{REPLICATOR_STATEFUL_SET_SUFFIX}") } +fn create_legacy_stateful_set_name(prefix: &str) -> String { + format!("{prefix}-{LEGACY_REPLICATOR_STATEFUL_SET_SUFFIX}") +} + fn create_pod_name(prefix: &str) -> String { format!("{prefix}-{REPLICATOR_STATEFUL_SET_SUFFIX}-0") } +fn create_legacy_pod_name(prefix: &str) -> String { + format!("{prefix}-{LEGACY_REPLICATOR_STATEFUL_SET_SUFFIX}-0") +} + +fn unique_current_and_legacy_names(current: String, legacy: String) -> Vec { + if current == legacy { vec![current] } else { vec![current, legacy] } +} + +fn stateful_set_names_for_lookup(prefix: &str) -> Vec { + unique_current_and_legacy_names( + create_stateful_set_name(prefix), + create_legacy_stateful_set_name(prefix), + ) +} + +fn pod_names_for_status(prefix: &str) -> Vec { + unique_current_and_legacy_names(create_pod_name(prefix), create_legacy_pod_name(prefix)) +} + fn create_replicator_app_name(prefix: &str) -> String { format!("{prefix}-{REPLICATOR_APP_SUFFIX}") } @@ -1467,6 +1522,10 @@ mod tests { use crate::configs::pipeline::ReplicatorResourcesConfig; const TENANT_ID: &str = "abcdefghijklmnopqrst"; + const MAX_TENANT_ID: &str = "abcdefghijklmnopqrst"; + const MAX_BIGINT_ID: i64 = 9_223_372_036_854_775_807; + const MAX_K8S_LABEL_VALUE_LEN: usize = 63; + const CONTROLLER_REVISION_HASH_LEN: usize = 10; fn create_k8s_object_prefix(tenant_id: &str, replicator_id: i64) -> String { format!("{tenant_id}-{replicator_id}") @@ -1498,6 +1557,48 @@ mod tests { .any(|entry| entry.get("name").and_then(serde_json::Value::as_str) == Some(name)) } + fn collect_kubernetes_label_values( + value: &serde_json::Value, + labels: &mut Vec<(String, String)>, + ) { + match value { + serde_json::Value::Object(map) => { + for label_field in ["labels", "matchLabels"] { + if let Some(serde_json::Value::Object(label_map)) = map.get(label_field) { + labels.extend(label_map.iter().filter_map(|(key, value)| { + value.as_str().map(|value| (key.clone(), value.to_owned())) + })); + } + } + + for child in map.values() { + collect_kubernetes_label_values(child, labels); + } + } + serde_json::Value::Array(values) => { + for child in values { + collect_kubernetes_label_values(child, labels); + } + } + _ => {} + } + } + + fn assert_kubernetes_label_values_are_safe(resource_name: &str, resource: &serde_json::Value) { + let mut labels = Vec::new(); + collect_kubernetes_label_values(resource, &mut labels); + assert!(!labels.is_empty(), "{resource_name} should contain Kubernetes labels"); + + for (key, value) in labels { + assert!( + value.len() <= MAX_K8S_LABEL_VALUE_LEN, + "{resource_name} generated label {key}={value} with length {}, exceeding \ + {MAX_K8S_LABEL_VALUE_LEN}", + value.len() + ); + } + } + #[test] fn test_replicator_resource_config_uses_environment_defaults() { let prod = ReplicatorResourceConfig::load(&Environment::Prod).unwrap(); @@ -1529,6 +1630,153 @@ mod tests { assert_eq!(config.replicator_memory_limit, "1843Mi"); } + #[test] + fn generated_kubernetes_labels_fit_with_max_tenant_and_replicator_ids() { + let prefix = create_k8s_object_prefix(MAX_TENANT_ID, MAX_BIGINT_ID); + let replicator_app_name = create_replicator_app_name(&prefix); + let postgres_secret_name = create_postgres_secret_name(&prefix); + let clickhouse_secret_name = create_clickhouse_secret_name(&prefix); + let bq_secret_name = create_bq_secret_name(&prefix); + let iceberg_secret_name = create_iceberg_secret_name(&prefix); + let ducklake_secret_name = create_ducklake_secret_name(&prefix); + let config_map_name = create_replicator_config_map_name(&prefix); + let ducklake_maintenance_name = create_ducklake_maintenance_name(&prefix); + let stateful_set_name = create_stateful_set_name(&prefix); + let controller_revision_label = + format!("{stateful_set_name}-{hash}", hash = "0".repeat(CONTROLLER_REVISION_HASH_LEN)); + + assert!( + controller_revision_label.len() <= MAX_K8S_LABEL_VALUE_LEN, + "stateful set controller revision label {controller_revision_label} has length {}, \ + exceeding {MAX_K8S_LABEL_VALUE_LEN}", + controller_revision_label.len() + ); + + let environment = Environment::Prod; + let config = ReplicatorResourceConfig::load(&environment).unwrap(); + let replicator_image = "supabase/replicator:1.2.3"; + let container_environment = create_container_environment_json( + &prefix, + &environment, + replicator_image, + DestinationType::Ducklake, + None, + LogLevel::Info, + ); + let node_selector = create_node_selector_json(&environment); + let init_containers = create_init_containers_json(&prefix, &environment, &config); + let volumes = create_volumes_json(&prefix, &environment); + let volume_mounts = create_volume_mounts_json(&environment); + + let resources = vec![ + ( + "postgres secret", + create_postgres_secret_json(&postgres_secret_name, &replicator_app_name, "secret"), + ), + ( + "clickhouse secret", + serde_json::to_value(create_clickhouse_password_secret( + &clickhouse_secret_name, + &replicator_app_name, + "secret", + )) + .unwrap(), + ), + ( + "bigquery secret", + create_bq_service_account_key_secret_json( + &bq_secret_name, + &replicator_app_name, + "secret", + ), + ), + ( + "iceberg secret", + create_iceberg_secret_json( + &iceberg_secret_name, + &replicator_app_name, + "secret", + "secret", + "secret", + ), + ), + ( + "ducklake secret", + create_ducklake_secret_json( + &ducklake_secret_name, + &replicator_app_name, + "secret", + "secret", + ), + ), + ( + "replicator config map", + create_replicator_config_map_json( + &config_map_name, + &replicator_app_name, + vec![ReplicatorConfigMapFile { + filename: "prod.json".to_owned(), + content: "{}".to_owned(), + }], + ), + ), + ( + "ducklake maintenance", + create_ducklake_maintenance_json( + &prefix, + &ducklake_maintenance_name, + DuckLakeMaintenanceResourceConfig { + tenant_id: MAX_TENANT_ID.to_owned(), + pipeline_id: MAX_BIGINT_ID, + replicator_id: MAX_BIGINT_ID, + image: replicator_image.to_owned(), + policy: DuckLakeMaintenancePolicy::default(), + }, + ), + ), + ( + "replicator stateful set", + create_replicator_stateful_set_json( + &prefix, + &stateful_set_name, + replicator_image, + container_environment, + node_selector, + init_containers, + volumes, + volume_mounts, + &config, + ), + ), + ]; + + for (resource_name, resource) in resources { + assert_kubernetes_label_values_are_safe(resource_name, &resource); + } + } + + #[test] + fn replicator_workload_names_use_short_suffix_and_keep_legacy_lookup_names() { + let prefix = create_k8s_object_prefix("tenant-1", 42); + + assert_eq!(create_stateful_set_name(&prefix), "tenant-1-42-replicator"); + assert_eq!(create_legacy_stateful_set_name(&prefix), "tenant-1-42-replicator-stateful-set"); + assert_eq!( + stateful_set_names_for_lookup(&prefix), + vec![ + "tenant-1-42-replicator".to_owned(), + "tenant-1-42-replicator-stateful-set".to_owned(), + ] + ); + assert_eq!( + pod_names_for_status(&prefix), + vec![ + "tenant-1-42-replicator-0".to_owned(), + "tenant-1-42-replicator-stateful-set-0".to_owned(), + ] + ); + } + #[test] fn test_replicator_resource_config_prefers_pipeline_over_api_config() { let overrides = ReplicatorResourcesConfig { diff --git a/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_bq_replicator_stateful_set_json-2.snap b/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_bq_replicator_stateful_set_json-2.snap index e45130f16..a6d8b13a7 100644 --- a/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_bq_replicator_stateful_set_json-2.snap +++ b/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_bq_replicator_stateful_set_json-2.snap @@ -1,6 +1,6 @@ --- source: crates/etl-api/src/k8s/http.rs -expression: stateful_set_json +expression: "serde_json :: to_string_pretty(& stateful_set_json).unwrap()" --- { "apiVersion": "apps/v1", @@ -10,7 +10,7 @@ expression: stateful_set_json "etl.supabase.com/app-name": "abcdefghijklmnopqrst-42-replicator-app", "etl.supabase.com/app-type": "etl-replicator-app" }, - "name": "abcdefghijklmnopqrst-42-replicator-stateful-set", + "name": "abcdefghijklmnopqrst-42-replicator", "namespace": "etl-data-plane" }, "spec": { diff --git a/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_bq_replicator_stateful_set_json-3.snap b/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_bq_replicator_stateful_set_json-3.snap index 135f41312..0eea18314 100644 --- a/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_bq_replicator_stateful_set_json-3.snap +++ b/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_bq_replicator_stateful_set_json-3.snap @@ -1,6 +1,6 @@ --- source: crates/etl-api/src/k8s/http.rs -expression: stateful_set_json +expression: "serde_json :: to_string_pretty(& stateful_set_json).unwrap()" --- { "apiVersion": "apps/v1", @@ -10,7 +10,7 @@ expression: stateful_set_json "etl.supabase.com/app-name": "abcdefghijklmnopqrst-42-replicator-app", "etl.supabase.com/app-type": "etl-replicator-app" }, - "name": "abcdefghijklmnopqrst-42-replicator-stateful-set", + "name": "abcdefghijklmnopqrst-42-replicator", "namespace": "etl-data-plane" }, "spec": { diff --git a/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_bq_replicator_stateful_set_json.snap b/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_bq_replicator_stateful_set_json.snap index bf3f322ae..c7134d1cc 100644 --- a/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_bq_replicator_stateful_set_json.snap +++ b/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_bq_replicator_stateful_set_json.snap @@ -1,6 +1,6 @@ --- source: crates/etl-api/src/k8s/http.rs -expression: stateful_set_json +expression: "serde_json :: to_string_pretty(& stateful_set_json).unwrap()" --- { "apiVersion": "apps/v1", @@ -10,7 +10,7 @@ expression: stateful_set_json "etl.supabase.com/app-name": "abcdefghijklmnopqrst-42-replicator-app", "etl.supabase.com/app-type": "etl-replicator-app" }, - "name": "abcdefghijklmnopqrst-42-replicator-stateful-set", + "name": "abcdefghijklmnopqrst-42-replicator", "namespace": "etl-data-plane" }, "spec": { diff --git a/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_ducklake_replicator_stateful_set_json-2.snap b/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_ducklake_replicator_stateful_set_json-2.snap index 12c0646d2..77ed0dc57 100644 --- a/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_ducklake_replicator_stateful_set_json-2.snap +++ b/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_ducklake_replicator_stateful_set_json-2.snap @@ -1,6 +1,6 @@ --- source: crates/etl-api/src/k8s/http.rs -expression: stateful_set_json +expression: "serde_json :: to_string_pretty(& stateful_set_json).unwrap()" --- { "apiVersion": "apps/v1", @@ -10,7 +10,7 @@ expression: stateful_set_json "etl.supabase.com/app-name": "abcdefghijklmnopqrst-42-replicator-app", "etl.supabase.com/app-type": "etl-replicator-app" }, - "name": "abcdefghijklmnopqrst-42-replicator-stateful-set", + "name": "abcdefghijklmnopqrst-42-replicator", "namespace": "etl-data-plane" }, "spec": { diff --git a/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_ducklake_replicator_stateful_set_json-3.snap b/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_ducklake_replicator_stateful_set_json-3.snap index 344044438..305909791 100644 --- a/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_ducklake_replicator_stateful_set_json-3.snap +++ b/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_ducklake_replicator_stateful_set_json-3.snap @@ -1,6 +1,6 @@ --- source: crates/etl-api/src/k8s/http.rs -expression: stateful_set_json +expression: "serde_json :: to_string_pretty(& stateful_set_json).unwrap()" --- { "apiVersion": "apps/v1", @@ -10,7 +10,7 @@ expression: stateful_set_json "etl.supabase.com/app-name": "abcdefghijklmnopqrst-42-replicator-app", "etl.supabase.com/app-type": "etl-replicator-app" }, - "name": "abcdefghijklmnopqrst-42-replicator-stateful-set", + "name": "abcdefghijklmnopqrst-42-replicator", "namespace": "etl-data-plane" }, "spec": { diff --git a/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_ducklake_replicator_stateful_set_json.snap b/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_ducklake_replicator_stateful_set_json.snap index 348acfd47..8dd8e2124 100644 --- a/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_ducklake_replicator_stateful_set_json.snap +++ b/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_ducklake_replicator_stateful_set_json.snap @@ -1,7 +1,6 @@ --- source: crates/etl-api/src/k8s/http.rs -assertion_line: 1786 -expression: stateful_set_json +expression: "serde_json :: to_string_pretty(& stateful_set_json).unwrap()" --- { "apiVersion": "apps/v1", @@ -11,7 +10,7 @@ expression: stateful_set_json "etl.supabase.com/app-name": "abcdefghijklmnopqrst-42-replicator-app", "etl.supabase.com/app-type": "etl-replicator-app" }, - "name": "abcdefghijklmnopqrst-42-replicator-stateful-set", + "name": "abcdefghijklmnopqrst-42-replicator", "namespace": "etl-data-plane" }, "spec": { diff --git a/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_iceberg_replicator_stateful_set_json-2.snap b/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_iceberg_replicator_stateful_set_json-2.snap index 071f44bbb..434593a05 100644 --- a/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_iceberg_replicator_stateful_set_json-2.snap +++ b/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_iceberg_replicator_stateful_set_json-2.snap @@ -1,6 +1,6 @@ --- source: crates/etl-api/src/k8s/http.rs -expression: stateful_set_json +expression: "serde_json :: to_string_pretty(& stateful_set_json).unwrap()" --- { "apiVersion": "apps/v1", @@ -10,7 +10,7 @@ expression: stateful_set_json "etl.supabase.com/app-name": "abcdefghijklmnopqrst-42-replicator-app", "etl.supabase.com/app-type": "etl-replicator-app" }, - "name": "abcdefghijklmnopqrst-42-replicator-stateful-set", + "name": "abcdefghijklmnopqrst-42-replicator", "namespace": "etl-data-plane" }, "spec": { diff --git a/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_iceberg_replicator_stateful_set_json-3.snap b/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_iceberg_replicator_stateful_set_json-3.snap index 1154f63f2..99a027dcd 100644 --- a/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_iceberg_replicator_stateful_set_json-3.snap +++ b/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_iceberg_replicator_stateful_set_json-3.snap @@ -1,6 +1,6 @@ --- source: crates/etl-api/src/k8s/http.rs -expression: stateful_set_json +expression: "serde_json :: to_string_pretty(& stateful_set_json).unwrap()" --- { "apiVersion": "apps/v1", @@ -10,7 +10,7 @@ expression: stateful_set_json "etl.supabase.com/app-name": "abcdefghijklmnopqrst-42-replicator-app", "etl.supabase.com/app-type": "etl-replicator-app" }, - "name": "abcdefghijklmnopqrst-42-replicator-stateful-set", + "name": "abcdefghijklmnopqrst-42-replicator", "namespace": "etl-data-plane" }, "spec": { diff --git a/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_iceberg_replicator_stateful_set_json.snap b/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_iceberg_replicator_stateful_set_json.snap index d4f3c5643..9a207065f 100644 --- a/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_iceberg_replicator_stateful_set_json.snap +++ b/crates/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_iceberg_replicator_stateful_set_json.snap @@ -1,6 +1,6 @@ --- source: crates/etl-api/src/k8s/http.rs -expression: stateful_set_json +expression: "serde_json :: to_string_pretty(& stateful_set_json).unwrap()" --- { "apiVersion": "apps/v1", @@ -10,7 +10,7 @@ expression: stateful_set_json "etl.supabase.com/app-name": "abcdefghijklmnopqrst-42-replicator-app", "etl.supabase.com/app-type": "etl-replicator-app" }, - "name": "abcdefghijklmnopqrst-42-replicator-stateful-set", + "name": "abcdefghijklmnopqrst-42-replicator", "namespace": "etl-data-plane" }, "spec": { diff --git a/crates/etl-api/src/routes/mod.rs b/crates/etl-api/src/routes/mod.rs index 61cf8ce26..71bd0e167 100644 --- a/crates/etl-api/src/routes/mod.rs +++ b/crates/etl-api/src/routes/mod.rs @@ -30,6 +30,27 @@ pub enum TenantIdError { TenantIdIllFormed, } +pub(crate) const MAX_TENANT_ID_LEN: usize = 20; + +pub(crate) fn validate_tenant_id(tenant_id: &str) -> Result<(), TenantIdError> { + let is_valid_char = + |byte: u8| byte.is_ascii_lowercase() || byte.is_ascii_digit() || byte == b'-'; + + let Some(first) = tenant_id.bytes().next() else { + return Err(TenantIdError::TenantIdIllFormed); + }; + + if tenant_id.len() > MAX_TENANT_ID_LEN + || !(first.is_ascii_lowercase() || first.is_ascii_digit()) + || tenant_id.ends_with('-') + || !tenant_id.bytes().all(is_valid_char) + { + return Err(TenantIdError::TenantIdIllFormed); + } + + Ok(()) +} + fn extract_tenant_id(req: &HttpRequest) -> Result<&str, TenantIdError> { let headers = req.headers(); let tenant_id = headers @@ -38,5 +59,38 @@ fn extract_tenant_id(req: &HttpRequest) -> Result<&str, TenantIdError> { .to_str() .map_err(|_| TenantIdError::TenantIdIllFormed)?; + validate_tenant_id(tenant_id)?; + Ok(tenant_id) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn tenant_id_validation_accepts_kubernetes_safe_ids() { + validate_tenant_id("a").unwrap(); + validate_tenant_id("abc123").unwrap(); + validate_tenant_id("abcdefghijklmnopqrst").unwrap(); + } + + #[test] + fn tenant_id_validation_rejects_kubernetes_unsafe_ids() { + for tenant_id in [ + "", + "abcdefghijklmnopqrstu", + "-tenant", + "tenant-", + "tenant_id", + "Tenant", + "tenant.id", + "tenant/id", + ] { + assert_eq!( + validate_tenant_id(tenant_id).unwrap_err().to_string(), + TenantIdError::TenantIdIllFormed.to_string() + ); + } + } +} diff --git a/crates/etl-api/src/routes/pipelines.rs b/crates/etl-api/src/routes/pipelines.rs index 75e0fbf9b..ef73dba17 100644 --- a/crates/etl-api/src/routes/pipelines.rs +++ b/crates/etl-api/src/routes/pipelines.rs @@ -38,7 +38,8 @@ use crate::{ TrustedRootCertsError, core::{ create_k8s_object_prefix, create_or_update_pipeline_resources_in_k8s, - delete_pipeline_resources_in_k8s, is_replicator_active, is_replicator_pod_stopped, + delete_pipeline_resources_in_k8s, is_replicator_active, + should_reconcile_replicator_resources, }, }, routes::{ErrorMessage, TenantIdError, extract_tenant_id, utils as route_utils}, @@ -706,7 +707,8 @@ pub(crate) async fn update_pipeline( let (pipeline, replicator, image, source, destination) = read_pipeline_components(&mut txn, tenant_id, pipeline_id, &encryption_key).await?; - if is_replicator_pod_stopped(k8s_client.as_ref(), tenant_id, replicator.id).await? { + if !should_reconcile_replicator_resources(k8s_client.as_ref(), tenant_id, replicator.id).await? + { txn.commit().await?; return Ok(HttpResponse::Ok().finish()); @@ -1383,9 +1385,11 @@ pub(crate) async fn update_pipeline_version( return Ok(HttpResponse::Ok().finish()); } - // If a replicator is not running, we don't want to create/update k8s resources. - // It's fine to just update the image version in the db. - if is_replicator_pod_stopped(k8s_client.as_ref(), tenant_id, replicator.id).await? { + // If a replicator has no runtime resources, we don't want to create new K8s + // resources. If a StatefulSet still exists, reconcile it even when no pod is + // currently running so a broken or stale StatefulSet can be repaired. + if !should_reconcile_replicator_resources(k8s_client.as_ref(), tenant_id, replicator.id).await? + { txn.commit().await?; return Ok(HttpResponse::Ok().finish()); diff --git a/crates/etl-api/src/routes/tenants.rs b/crates/etl-api/src/routes/tenants.rs index 3272acd73..a1eb7840c 100644 --- a/crates/etl-api/src/routes/tenants.rs +++ b/crates/etl-api/src/routes/tenants.rs @@ -29,7 +29,7 @@ use crate::{ K8sClient, TrustedRootCertsCache, TrustedRootCertsError, core::{K8sCoreError, first_active_pipeline_id}, }, - routes::ErrorMessage, + routes::{ErrorMessage, TenantIdError, validate_tenant_id}, }; #[derive(Debug, Error)] @@ -57,6 +57,9 @@ pub enum TenantError { #[error("The pipeline with id {0} is active; stop it before deleting it")] ActivePipeline(i64), + + #[error(transparent)] + TenantId(#[from] TenantIdError), } impl TenantError { @@ -80,6 +83,7 @@ impl ResponseError for TenantError { match self { TenantError::TenantsDb(TenantsDbError::Conflict(_)) | TenantError::ActivePipeline(_) => StatusCode::CONFLICT, + TenantError::TenantId(_) => StatusCode::BAD_REQUEST, TenantError::TenantsDb(_) | TenantError::SourcesDb(_) | TenantError::PipelinesDb(_) @@ -165,6 +169,7 @@ pub(crate) async fn create_tenant( root_span: RootSpan, ) -> Result { let tenant = tenant.into_inner(); + validate_tenant_id(&tenant.id)?; root_span.record("project", &tenant.id); @@ -198,6 +203,7 @@ pub(crate) async fn create_or_update_tenant( ) -> Result { let tenant_id = tenant_id.into_inner(); let tenant = tenant.into_inner(); + validate_tenant_id(&tenant_id)?; root_span.record("project", &tenant_id); @@ -227,6 +233,7 @@ pub(crate) async fn read_tenant( root_span: RootSpan, ) -> Result { let tenant_id = tenant_id.into_inner(); + validate_tenant_id(&tenant_id)?; root_span.record("project", &tenant_id); @@ -261,6 +268,7 @@ pub(crate) async fn update_tenant( ) -> Result { let tenant = tenant.into_inner(); let tenant_id = tenant_id.into_inner(); + validate_tenant_id(&tenant_id)?; root_span.record("project", &tenant_id); @@ -296,6 +304,7 @@ pub(crate) async fn delete_tenant( root_span: RootSpan, ) -> Result { let tenant_id = tenant_id.into_inner(); + validate_tenant_id(&tenant_id)?; root_span.record("project", &tenant_id); diff --git a/crates/etl-api/src/routes/tenants_sources.rs b/crates/etl-api/src/routes/tenants_sources.rs index 712a8d6dd..df6732a96 100644 --- a/crates/etl-api/src/routes/tenants_sources.rs +++ b/crates/etl-api/src/routes/tenants_sources.rs @@ -18,7 +18,7 @@ use crate::{ }, data::{self, tenants::TenantsDbError, tenants_sources::TenantSourceDbError}, k8s::TrustedRootCertsCache, - routes::{ErrorMessage, common, utils}, + routes::{ErrorMessage, TenantIdError, common, utils, validate_tenant_id}, validation::ValidationError, }; @@ -35,6 +35,9 @@ enum TenantSourceError { #[error("Validation failed: {0}")] ValidationFailed(String), + + #[error(transparent)] + TenantId(#[from] TenantIdError), } impl TenantSourceError { @@ -66,6 +69,7 @@ impl ResponseError for TenantSourceError { StatusCode::INTERNAL_SERVER_ERROR } TenantSourceError::Validation(error) => utils::validation_error_status_code(error), + TenantSourceError::TenantId(_) => StatusCode::BAD_REQUEST, TenantSourceError::ValidationFailed(_) => StatusCode::UNPROCESSABLE_ENTITY, } } @@ -140,6 +144,7 @@ pub(crate) async fn create_tenant_and_source( root_span: RootSpan, ) -> Result { let tenant_and_source = tenant_and_source.into_inner(); + validate_tenant_id(&tenant_and_source.tenant_id)?; root_span.record("project", &tenant_and_source.tenant_id); diff --git a/crates/etl-api/tests/destinations_pipelines.rs b/crates/etl-api/tests/destinations_pipelines.rs index f6b38f09e..adf915ad6 100644 --- a/crates/etl-api/tests/destinations_pipelines.rs +++ b/crates/etl-api/tests/destinations_pipelines.rs @@ -206,13 +206,13 @@ async fn destination_and_pipeline_with_another_tenants_source_cannot_be_created( let tenant1_id = &create_tenant_with_id_and_name( &app, "abcdefghijklmnopqrst".to_owned(), - "tenant_1".to_owned(), + "tenant-1".to_owned(), ) .await; let tenant2_id = &create_tenant_with_id_and_name( &app, "tsrqponmlkjihgfedcba".to_owned(), - "tenant_2".to_owned(), + "tenant-2".to_owned(), ) .await; let source2_id = create_source(&app, tenant2_id).await; @@ -342,13 +342,13 @@ async fn destination_and_pipeline_with_another_tenants_source_cannot_be_updated( let tenant1_id = &create_tenant_with_id_and_name( &app, "abcdefghijklmnopqrst".to_owned(), - "tenant_1".to_owned(), + "tenant-1".to_owned(), ) .await; let tenant2_id = &create_tenant_with_id_and_name( &app, "tsrqponmlkjihgfedcba".to_owned(), - "tenant_2".to_owned(), + "tenant-2".to_owned(), ) .await; @@ -389,13 +389,13 @@ async fn destination_and_pipeline_with_another_tenants_destination_cannot_be_upd let tenant1_id = &create_tenant_with_id_and_name( &app, "abcdefghijklmnopqrst".to_owned(), - "tenant_1".to_owned(), + "tenant-1".to_owned(), ) .await; let tenant2_id = &create_tenant_with_id_and_name( &app, "tsrqponmlkjihgfedcba".to_owned(), - "tenant_2".to_owned(), + "tenant-2".to_owned(), ) .await; @@ -441,13 +441,13 @@ async fn destination_and_pipeline_with_another_tenants_pipeline_cannot_be_update let tenant1_id = &create_tenant_with_id_and_name( &app, "abcdefghijklmnopqrst".to_owned(), - "tenant_1".to_owned(), + "tenant-1".to_owned(), ) .await; let tenant2_id = &create_tenant_with_id_and_name( &app, "tsrqponmlkjihgfedcba".to_owned(), - "tenant_2".to_owned(), + "tenant-2".to_owned(), ) .await; diff --git a/crates/etl-api/tests/pipelines.rs b/crates/etl-api/tests/pipelines.rs index 45ad1e4ac..116ad2510 100644 --- a/crates/etl-api/tests/pipelines.rs +++ b/crates/etl-api/tests/pipelines.rs @@ -293,13 +293,13 @@ async fn pipeline_with_another_tenants_source_cannot_be_created() { let tenant1_id = &create_tenant_with_id_and_name( &app, "abcdefghijklmnopqrst".to_owned(), - "tenant_1".to_owned(), + "tenant-1".to_owned(), ) .await; let tenant2_id = &create_tenant_with_id_and_name( &app, "tsrqponmlkjihgfedcba".to_owned(), - "tenant_2".to_owned(), + "tenant-2".to_owned(), ) .await; let source2_id = create_source(&app, tenant2_id).await; @@ -326,13 +326,13 @@ async fn pipeline_with_another_tenants_destination_cannot_be_created() { let tenant1_id = &create_tenant_with_id_and_name( &app, "abcdefghijklmnopqrst".to_owned(), - "tenant_1".to_owned(), + "tenant-1".to_owned(), ) .await; let tenant2_id = &create_tenant_with_id_and_name( &app, "tsrqponmlkjihgfedcba".to_owned(), - "tenant_2".to_owned(), + "tenant-2".to_owned(), ) .await; let source1_id = create_source(&app, tenant1_id).await; @@ -548,13 +548,13 @@ async fn pipeline_with_another_tenants_source_cannot_be_updated() { let tenant1_id = &create_tenant_with_id_and_name( &app, "abcdefghijklmnopqrst".to_owned(), - "tenant_1".to_owned(), + "tenant-1".to_owned(), ) .await; let tenant2_id = &create_tenant_with_id_and_name( &app, "tsrqponmlkjihgfedcba".to_owned(), - "tenant_2".to_owned(), + "tenant-2".to_owned(), ) .await; let source1_id = create_source(&app, tenant1_id).await; @@ -592,13 +592,13 @@ async fn pipeline_with_another_tenants_destination_cannot_be_updated() { let tenant1_id = &create_tenant_with_id_and_name( &app, "abcdefghijklmnopqrst".to_owned(), - "tenant_1".to_owned(), + "tenant-1".to_owned(), ) .await; let tenant2_id = &create_tenant_with_id_and_name( &app, "tsrqponmlkjihgfedcba".to_owned(), - "tenant_2".to_owned(), + "tenant-2".to_owned(), ) .await; let source1_id = create_source(&app, tenant1_id).await; diff --git a/crates/etl-api/tests/support/k8s_client.rs b/crates/etl-api/tests/support/k8s_client.rs index c783fbb68..016b3d03d 100644 --- a/crates/etl-api/tests/support/k8s_client.rs +++ b/crates/etl-api/tests/support/k8s_client.rs @@ -190,6 +190,10 @@ impl K8sClient for MockK8sClient { Ok(()) } + async fn replicator_stateful_set_exists(&self, _prefix: &str) -> Result { + Ok(false) + } + async fn create_or_update_ducklake_maintenance( &self, _prefix: &str, From 1ac02fa6691c6a6a45611c4c4d8039dc7ba5f9df Mon Sep 17 00:00:00 2001 From: Coenen Benjamin Date: Wed, 20 May 2026 17:28:05 +0200 Subject: [PATCH 20/29] feat(ducklake): add pause metrics for external maintenances (#757) Signed-off-by: Benjamin <5719034+bnjjj@users.noreply.github.com> --- .../etl-destinations/src/ducklake/METRICS.md | 21 +++++---- .../src/ducklake/external_maintenance.rs | 44 ++++++++++++++++++- .../etl-destinations/src/ducklake/metrics.rs | 7 +++ 3 files changed, 63 insertions(+), 9 deletions(-) diff --git a/crates/etl-destinations/src/ducklake/METRICS.md b/crates/etl-destinations/src/ducklake/METRICS.md index 436c4f58c..80fd2bdba 100644 --- a/crates/etl-destinations/src/ducklake/METRICS.md +++ b/crates/etl-destinations/src/ducklake/METRICS.md @@ -10,16 +10,14 @@ The metrics fall into four groups: - external maintenance metrics: operation-trigger counts and duration samples for time foreground ingestion was quiesced by the Kubernetes maintenance plane. -- table-health samples: histograms recorded by a background sampler every - 30 seconds from the PostgreSQL DuckLake metadata catalog. They describe the - current shape of tables known to the current destination instance. +- table-health samples: gauges recorded by a background sampler every 30 seconds + from the PostgreSQL DuckLake metadata catalog. They describe the current shape + of tables known to the current destination instance. - catalog-backlog gauges: current global snapshot and deletion backlog in the attached DuckLake catalog. -These metrics intentionally avoid a `table_name` label. That keeps Prometheus -cardinality low. The tradeoff is that table-health metrics are sampled as -histograms over recently written tables rather than exported as one time series -per table. +Table-health metrics carry a `table` label because they are used to diagnose and +trigger table-specific maintenance behavior. ## Metric groups @@ -78,6 +76,7 @@ How to read them: ### External maintenance metrics - `etl_ducklake_external_maintenance_pause_duration_seconds` +- `etl_ducklake_external_maintenance_pause_active` - `etl_ducklake_external_maintenance_triggered_total` `etl_ducklake_external_maintenance_pause_duration_seconds` is emitted by the @@ -87,7 +86,13 @@ only the time after the destination has drained foreground mutations and reporte It carries one label: -- `outcome`: `cleared`, `expired`, `replaced`, or `resource_deleted` +- `outcome`: `cleared`, `expired`, `replaced`, or `state_missing` + +`etl_ducklake_external_maintenance_pause_active` is a 0/1 gauge emitted by the +replicator while foreground ingestion is blocked or draining for an external +maintenance pause. A value of `1` means ingestion is paused for external +maintenance; `0` means the replicator is not currently paused by external +maintenance. `etl_ducklake_external_maintenance_triggered_total` counts external maintenance operation requests emitted by the replicator after it samples DuckLake catalog diff --git a/crates/etl-destinations/src/ducklake/external_maintenance.rs b/crates/etl-destinations/src/ducklake/external_maintenance.rs index 8ccc4b0e3..f64644368 100644 --- a/crates/etl-destinations/src/ducklake/external_maintenance.rs +++ b/crates/etl-destinations/src/ducklake/external_maintenance.rs @@ -13,7 +13,7 @@ pub use etl_maintenance::{ ExternalMaintenanceState, ExternalMaintenanceStore, ExternalMaintenanceWatcherConfig, KubernetesExternalMaintenanceStore, PostgresExternalMaintenanceStore, }; -use metrics::{counter, histogram}; +use metrics::{counter, gauge, histogram}; use sqlx::PgPool; use tokio::time; use tracing::{debug, info, warn}; @@ -21,6 +21,7 @@ use tracing::{debug, info, warn}; use crate::ducklake::{ DuckLakeDestination, DuckLakeExternalMaintenancePause, metrics::{ + ETL_DUCKLAKE_EXTERNAL_MAINTENANCE_PAUSE_ACTIVE, ETL_DUCKLAKE_EXTERNAL_MAINTENANCE_PAUSE_DURATION_SECONDS, ETL_DUCKLAKE_EXTERNAL_MAINTENANCE_TRIGGERED_TOTAL, MAINTENANCE_OPERATION_LABEL, MAINTENANCE_OUTCOME_LABEL, MAINTENANCE_REASON_LABEL, @@ -89,6 +90,7 @@ where M: ExternalMaintenanceStore, { let mut held_pause: Option = None; + record_external_maintenance_pause_active(false); info!( poll_interval_ms = config.poll_interval.as_millis() as u64, @@ -217,6 +219,7 @@ async fn reconcile_pause( pause.expires_at.to_rfc3339() ); report_pausing(store, &pause.run_id, config.store_timeout).await; + record_external_maintenance_pause_active(true); let external_pause = destination.acquire_external_maintenance_pause().await; if pause.expires_at <= Utc::now() { @@ -305,6 +308,7 @@ fn release_held_pause(held: HeldPause, outcome: &'static str) { let held_ms = Utc::now().signed_duration_since(held.quiesced_at).num_milliseconds().max(0) as u64; record_external_maintenance_pause_duration(&held, outcome); + record_external_maintenance_pause_active(false); info!( run_id = %held.run_id, outcome, @@ -316,6 +320,10 @@ fn release_held_pause(held: HeldPause, outcome: &'static str) { ); } +fn record_external_maintenance_pause_active(active: bool) { + gauge!(ETL_DUCKLAKE_EXTERNAL_MAINTENANCE_PAUSE_ACTIVE).set(if active { 1.0 } else { 0.0 }); +} + fn record_external_maintenance_pause_duration(held: &HeldPause, outcome: &'static str) { let duration_seconds = Utc::now().signed_duration_since(held.quiesced_at).num_milliseconds().max(0) as f64 @@ -580,3 +588,37 @@ fn record_external_maintenance_triggers( .increment(1); } } + +#[cfg(test)] +mod tests { + use etl_telemetry::metrics::init_metrics_handle; + + use super::record_external_maintenance_pause_active; + use crate::ducklake::metrics::{ + ETL_DUCKLAKE_EXTERNAL_MAINTENANCE_PAUSE_ACTIVE, register_metrics, + }; + + fn pause_active_gauge_value(rendered: &str) -> Option { + rendered.lines().find_map(|line| { + if line.starts_with(ETL_DUCKLAKE_EXTERNAL_MAINTENANCE_PAUSE_ACTIVE) { + line.split_whitespace().last()?.parse::().ok() + } else { + None + } + }) + } + + #[tokio::test] + async fn recording_external_maintenance_pause_active_exports_gauge_value() { + let handle = init_metrics_handle().expect("failed to initialize prometheus handle"); + register_metrics(); + + record_external_maintenance_pause_active(true); + let rendered = handle.render(); + assert_eq!(pause_active_gauge_value(&rendered), Some(1.0)); + + record_external_maintenance_pause_active(false); + let rendered = handle.render(); + assert_eq!(pause_active_gauge_value(&rendered), Some(0.0)); + } +} diff --git a/crates/etl-destinations/src/ducklake/metrics.rs b/crates/etl-destinations/src/ducklake/metrics.rs index fbd37ff9a..b9a05ac18 100644 --- a/crates/etl-destinations/src/ducklake/metrics.rs +++ b/crates/etl-destinations/src/ducklake/metrics.rs @@ -53,6 +53,8 @@ pub(crate) const ETL_DUCKLAKE_FAILED_BATCHES_TOTAL: &str = "etl_ducklake_failed_ pub(crate) const ETL_DUCKLAKE_REPLAYED_BATCHES_TOTAL: &str = "etl_ducklake_replayed_batches_total"; pub(crate) const ETL_DUCKLAKE_EXTERNAL_MAINTENANCE_PAUSE_DURATION_SECONDS: &str = "etl_ducklake_external_maintenance_pause_duration_seconds"; +pub(crate) const ETL_DUCKLAKE_EXTERNAL_MAINTENANCE_PAUSE_ACTIVE: &str = + "etl_ducklake_external_maintenance_pause_active"; pub(crate) const ETL_DUCKLAKE_EXTERNAL_MAINTENANCE_TRIGGERED_TOTAL: &str = "etl_ducklake_external_maintenance_triggered_total"; pub(crate) const ETL_DUCKLAKE_TABLE_ACTIVE_DATA_FILES: &str = @@ -231,6 +233,11 @@ pub(crate) fn register_metrics() { "Duration that DuckLake foreground ingestion was paused for an external maintenance \ run, labeled by outcome." ); + describe_gauge!( + ETL_DUCKLAKE_EXTERNAL_MAINTENANCE_PAUSE_ACTIVE, + Unit::Count, + "External DuckLake maintenance foreground-ingestion pause current state (0 or 1)." + ); describe_counter!( ETL_DUCKLAKE_EXTERNAL_MAINTENANCE_TRIGGERED_TOTAL, Unit::Count, From 38a3fde8f42a05185503a5f0c058ea6540798710 Mon Sep 17 00:00:00 2001 From: Coenen Benjamin Date: Thu, 21 May 2026 15:31:59 +0200 Subject: [PATCH 21/29] fix(ducklake): improve expireSnapshot maintenance (#759) * fix(ducklake): improve expireSnapshot maintenance Signed-off-by: Benjamin <5719034+bnjjj@users.noreply.github.com> * increase duckdb timeout for maintenance runner Signed-off-by: Benjamin <5719034+bnjjj@users.noreply.github.com> --------- Signed-off-by: Benjamin <5719034+bnjjj@users.noreply.github.com> --- .../src/ducklake/external_maintenance.rs | 110 +++++++++++++++++- crates/etl-maintenance/src/ducklake/runner.rs | 2 +- 2 files changed, 109 insertions(+), 3 deletions(-) diff --git a/crates/etl-destinations/src/ducklake/external_maintenance.rs b/crates/etl-destinations/src/ducklake/external_maintenance.rs index f64644368..69541edaf 100644 --- a/crates/etl-destinations/src/ducklake/external_maintenance.rs +++ b/crates/etl-destinations/src/ducklake/external_maintenance.rs @@ -28,6 +28,7 @@ use crate::ducklake::{ }, }; +const EXPIRE_SNAPSHOTS_MIN_INTERVAL_SECONDS: i64 = 24 * 60 * 60; const OPERATION_INLINE_FLUSH: &str = "flush_inlined_data"; const OPERATION_REWRITE_DATA_FILES: &str = "rewrite_data_files"; const OPERATION_EXPIRE_SNAPSHOTS: &str = "expire_snapshots"; @@ -43,6 +44,14 @@ struct HeldPause { _pause: DuckLakeExternalMaintenancePause, } +/// Suppresses duplicate expire-snapshots requests inside the daily controller +/// interval. +#[derive(Default)] +struct ExpireSnapshotsRequestGate { + initialized_from_state: bool, + next_request_after: Option>, +} + pub(super) async fn run_kubernetes_external_maintenance_watcher( destination: DuckLakeDestination, ) -> EtlResult<()> @@ -90,6 +99,7 @@ where M: ExternalMaintenanceStore, { let mut held_pause: Option = None; + let mut expire_snapshots_gate = ExpireSnapshotsRequestGate::default(); record_external_maintenance_pause_active(false); info!( @@ -112,8 +122,16 @@ where continue; } + expire_snapshots_gate.initialize_from_state_once(&state); reconcile_pause(&store, &config, &destination, &mut held_pause, &state).await; - maybe_request_operations(&store, &destination, &state, &config).await; + maybe_request_operations( + &store, + &destination, + &state, + &config, + &mut expire_snapshots_gate, + ) + .await; } Ok(Err(error)) => { warn!( @@ -335,11 +353,51 @@ fn record_external_maintenance_pause_duration(held: &HeldPause, outcome: &'stati .record(duration_seconds); } +impl ExpireSnapshotsRequestGate { + fn initialize_from_state_once(&mut self, state: &ExternalMaintenanceState) { + if self.initialized_from_state { + return; + } + self.initialized_from_state = true; + + let Some(completed_at) = + state.last_successful_operations.expire_snapshots.as_ref().map(|run| run.completed_at) + else { + debug!( + "ducklake expire-snapshots request gate initialized without previous successful \ + run" + ); + return; + }; + + self.next_request_after = + Some(completed_at + chrono::Duration::seconds(EXPIRE_SNAPSHOTS_MIN_INTERVAL_SECONDS)); + debug!( + completed_at = %completed_at, + next_request_after = ?self.next_request_after, + "ducklake expire-snapshots request gate initialized from maintenance state: \ + completed_at={}, next_request_after={:?}", + completed_at, + self.next_request_after + ); + } + + fn is_suppressed(&self, now: DateTime) -> bool { + self.next_request_after.is_some_and(|next_request_after| now < next_request_after) + } + + fn record_requested(&mut self, now: DateTime) { + self.next_request_after = + Some(now + chrono::Duration::seconds(EXPIRE_SNAPSHOTS_MIN_INTERVAL_SECONDS)); + } +} + async fn maybe_request_operations( store: &M, destination: &DuckLakeDestination, state: &ExternalMaintenanceState, config: &ExternalMaintenanceWatcherConfig, + expire_snapshots_gate: &mut ExpireSnapshotsRequestGate, ) where S: StateStore + SchemaStore + Clone + Send + Sync + 'static, M: ExternalMaintenanceStore, @@ -364,6 +422,10 @@ async fn maybe_request_operations( operations.rewrite_data_files &= state.operation_policy.rewrite_data_files_enabled; operations.expire_snapshots &= state.operation_policy.expire_snapshots_enabled; operations.cleanup_old_files &= state.operation_policy.cleanup_old_files_enabled; + if operations.expire_snapshots && expire_snapshots_gate.is_suppressed(Utc::now()) { + debug!("ducklake expire-snapshots request suppressed by local daily gate"); + operations.expire_snapshots = false; + } operations } Err(error) => { @@ -428,6 +490,9 @@ async fn maybe_request_operations( }; match time::timeout(config.store_timeout, store.request_operations(request)).await { Ok(Ok(ExternalMaintenanceRequestOutcome::Created)) => { + if requested.expire_snapshots && !already_requested.expire_snapshots { + expire_snapshots_gate.record_requested(Utc::now()); + } record_external_maintenance_triggers(requested, already_requested); } Ok(Ok(ExternalMaintenanceRequestOutcome::AlreadyCovered)) => { @@ -591,13 +656,28 @@ fn record_external_maintenance_triggers( #[cfg(test)] mod tests { + use chrono::{DateTime, TimeZone, Utc}; use etl_telemetry::metrics::init_metrics_handle; - use super::record_external_maintenance_pause_active; + use super::{ + ExpireSnapshotsRequestGate, ExternalMaintenanceOperationRun, ExternalMaintenanceState, + record_external_maintenance_pause_active, + }; use crate::ducklake::metrics::{ ETL_DUCKLAKE_EXTERNAL_MAINTENANCE_PAUSE_ACTIVE, register_metrics, }; + fn state_with_expire_snapshots_completed_at( + completed_at: DateTime, + ) -> ExternalMaintenanceState { + let mut state = ExternalMaintenanceState::present(); + state.last_successful_operations.expire_snapshots = Some(ExternalMaintenanceOperationRun { + run_id: Some("run-1".to_owned()), + completed_at, + }); + state + } + fn pause_active_gauge_value(rendered: &str) -> Option { rendered.lines().find_map(|line| { if line.starts_with(ETL_DUCKLAKE_EXTERNAL_MAINTENANCE_PAUSE_ACTIVE) { @@ -608,6 +688,32 @@ mod tests { }) } + #[test] + fn expire_snapshots_gate_initializes_from_state_once() { + let now = Utc.with_ymd_and_hms(2026, 5, 12, 12, 0, 0).unwrap(); + let first = state_with_expire_snapshots_completed_at(now - chrono::Duration::hours(1)); + let second = state_with_expire_snapshots_completed_at(now - chrono::Duration::hours(25)); + let mut gate = ExpireSnapshotsRequestGate::default(); + + gate.initialize_from_state_once(&first); + gate.initialize_from_state_once(&second); + + assert!(gate.is_suppressed(now)); + assert!(gate.is_suppressed(now + chrono::Duration::hours(22))); + assert!(!gate.is_suppressed(now + chrono::Duration::hours(23))); + } + + #[test] + fn expire_snapshots_gate_records_local_request() { + let now = Utc.with_ymd_and_hms(2026, 5, 12, 12, 0, 0).unwrap(); + let mut gate = ExpireSnapshotsRequestGate::default(); + + gate.record_requested(now); + + assert!(gate.is_suppressed(now + chrono::Duration::hours(23))); + assert!(!gate.is_suppressed(now + chrono::Duration::hours(24))); + } + #[tokio::test] async fn recording_external_maintenance_pause_active_exports_gauge_value() { let handle = init_metrics_handle().expect("failed to initialize prometheus handle"); diff --git a/crates/etl-maintenance/src/ducklake/runner.rs b/crates/etl-maintenance/src/ducklake/runner.rs index 5e00c086d..aa5169e79 100644 --- a/crates/etl-maintenance/src/ducklake/runner.rs +++ b/crates/etl-maintenance/src/ducklake/runner.rs @@ -50,7 +50,7 @@ const PARQUET_ROW_GROUP_SIZE_BYTES_OPTION_VALUE: &str = "10MB"; const PARQUET_VERSION_OPTION_NAME: &str = "parquet_version"; const PARQUET_VERSION_OPTION_VALUE: u8 = 2; const PRESERVE_INSERTION_ORDER_OPTION_NAME: &str = "preserve_insertion_order"; -const MAINTENANCE_QUERY_TIMEOUT: Duration = Duration::from_secs(3 * 60); +const MAINTENANCE_QUERY_TIMEOUT: Duration = Duration::from_secs(6 * 60); const BLOCKING_ABORT_GRACE: Duration = Duration::from_secs(30); const DUCKDB_MAINTENANCE_OPERATION_KIND: &str = "maintenance"; const ETL_DUCKLAKE_INLINE_FLUSH_ROWS: &str = "etl_ducklake_inline_flush_rows"; From d9b4bbc171a9385a88bf8417759538b6ff308371 Mon Sep 17 00:00:00 2001 From: Victor Farazdagi Date: Thu, 21 May 2026 17:53:50 +0300 Subject: [PATCH 22/29] feat(destinations): add Snowflake destination (#728) * feat: error and config * feat: implement snowflake auth * feat(snowflake): DDL SQL client * feat(snoflake): implement RowBatch * feat(snowflake): streaming client * update schema encoding * feat(snowflake): ChannelHandler abstraction * feat(snowflake): refactor streaming client * feat(snowflake): minor fixes * feat(snowflake): batching and backpressure * refactor(snowflake): mock tests cleanup * feat(snowflake): example file * chore: rebase * refactor: move channel into streaming module * refactor: update sql client * feat(snowflake): consolidated Snowflake client * feat(snowflake): destination core * test(snowflake): destination tests * test(snowflake): improve polling * test(snowflake): more integration tests * feat(examples): snowflake monitoring ui * feat(examples): snowflake load generator * refactor(snowflake): make mod prepended imports * fix(snowflake): deprecated method * feat(snowlake): reconnect on expired jwt token * refactor(snowflake): feature gate snowflake deps in examples * fix: cargo sort * refactor(snowflake): swap compromised dependency * chore(snowflake): clippy * chore(snowflake): fmt * fix: capitalize errors * test(snowflake): add schema evolution tests for rename, drop, and type change * fix(snowflake): prevent terminal corruption on F2/F3 errors * fix(snowflake): pr review comments * fix: rest client and batch fixes * fix: http client with timeouts * fix: some more comments * fix: fmt * fix: remove warehouse * fix: codeql comments * fix: more comments * rebase * fix: explicit clone * fix: clippy * fix: codeql comment --- .config/nextest.toml | 2 +- .env.example | 29 + .gitignore | 5 + Cargo.lock | 1113 ++++++++++++----- Cargo.toml | 9 + crates/etl-config/src/shared/destination.rs | 47 + crates/etl-destinations/Cargo.toml | 31 +- crates/etl-destinations/src/lib.rs | 7 +- .../etl-destinations/src/snowflake/README.md | 52 + crates/etl-destinations/src/snowflake/auth.rs | 444 +++++++ .../etl-destinations/src/snowflake/client.rs | 194 +++ .../etl-destinations/src/snowflake/config.rs | 70 ++ crates/etl-destinations/src/snowflake/core.rs | 476 +++++++ .../src/snowflake/encoding.rs | 469 +++++++ .../etl-destinations/src/snowflake/error.rs | 52 + .../etl-destinations/src/snowflake/metrics.rs | 35 + crates/etl-destinations/src/snowflake/mod.rs | 23 + .../etl-destinations/src/snowflake/schema.rs | 158 +++ .../src/snowflake/sql_client.rs | 370 ++++++ .../src/snowflake/streaming/batch.rs | 325 +++++ .../src/snowflake/streaming/channel.rs | 165 +++ .../src/snowflake/streaming/mod.rs | 104 ++ .../src/snowflake/streaming/offset_token.rs | 68 + .../src/snowflake/streaming/rest_client.rs | 593 +++++++++ .../src/snowflake/test_utils.rs | 58 + crates/etl-destinations/testdata/test_key.pem | 28 + crates/etl-destinations/tests/main.rs | 2 + .../etl-destinations/tests/snowflake/auth.rs | 20 + .../tests/snowflake/common.rs | 86 ++ .../tests/snowflake/destination.rs | 910 ++++++++++++++ .../etl-destinations/tests/snowflake/mod.rs | 5 + .../tests/snowflake/sql_client.rs | 98 ++ .../tests/snowflake/stream_client.rs | 335 +++++ crates/etl-examples/Cargo.toml | 41 +- crates/etl-examples/README.md | 3 +- .../etl-examples/src/bin/snowflake/README.md | 196 +++ .../src/bin/snowflake/commands.rs | 109 ++ .../etl-examples/src/bin/snowflake/logging.rs | 86 ++ crates/etl-examples/src/bin/snowflake/main.rs | 461 +++++++ .../etl-examples/src/bin/snowflake/state.rs | 451 +++++++ crates/etl-examples/src/bin/snowflake/tui.rs | 532 ++++++++ .../etl-examples/src/bin/snowflake_loadgen.rs | 457 +++++++ crates/etl-replicator/Cargo.toml | 2 +- crates/etl-replicator/src/core.rs | 28 + crates/etl-replicator/src/init/destination.rs | 5 +- crates/xtask/src/commands/nextest.rs | 2 +- 46 files changed, 8434 insertions(+), 322 deletions(-) create mode 100644 .env.example create mode 100644 crates/etl-destinations/src/snowflake/README.md create mode 100644 crates/etl-destinations/src/snowflake/auth.rs create mode 100644 crates/etl-destinations/src/snowflake/client.rs create mode 100644 crates/etl-destinations/src/snowflake/config.rs create mode 100644 crates/etl-destinations/src/snowflake/core.rs create mode 100644 crates/etl-destinations/src/snowflake/encoding.rs create mode 100644 crates/etl-destinations/src/snowflake/error.rs create mode 100644 crates/etl-destinations/src/snowflake/metrics.rs create mode 100644 crates/etl-destinations/src/snowflake/mod.rs create mode 100644 crates/etl-destinations/src/snowflake/schema.rs create mode 100644 crates/etl-destinations/src/snowflake/sql_client.rs create mode 100644 crates/etl-destinations/src/snowflake/streaming/batch.rs create mode 100644 crates/etl-destinations/src/snowflake/streaming/channel.rs create mode 100644 crates/etl-destinations/src/snowflake/streaming/mod.rs create mode 100644 crates/etl-destinations/src/snowflake/streaming/offset_token.rs create mode 100644 crates/etl-destinations/src/snowflake/streaming/rest_client.rs create mode 100644 crates/etl-destinations/src/snowflake/test_utils.rs create mode 100644 crates/etl-destinations/testdata/test_key.pem create mode 100644 crates/etl-destinations/tests/snowflake/auth.rs create mode 100644 crates/etl-destinations/tests/snowflake/common.rs create mode 100644 crates/etl-destinations/tests/snowflake/destination.rs create mode 100644 crates/etl-destinations/tests/snowflake/mod.rs create mode 100644 crates/etl-destinations/tests/snowflake/sql_client.rs create mode 100644 crates/etl-destinations/tests/snowflake/stream_client.rs create mode 100644 crates/etl-examples/src/bin/snowflake/README.md create mode 100644 crates/etl-examples/src/bin/snowflake/commands.rs create mode 100644 crates/etl-examples/src/bin/snowflake/logging.rs create mode 100644 crates/etl-examples/src/bin/snowflake/main.rs create mode 100644 crates/etl-examples/src/bin/snowflake/state.rs create mode 100644 crates/etl-examples/src/bin/snowflake/tui.rs create mode 100644 crates/etl-examples/src/bin/snowflake_loadgen.rs diff --git a/.config/nextest.toml b/.config/nextest.toml index 62fce77de..6891120fc 100644 --- a/.config/nextest.toml +++ b/.config/nextest.toml @@ -10,5 +10,5 @@ test-threads = "num-cpus" max-threads = 1 [[profile.default.overrides]] -filter = "test(exclusive_) | binary_id(etl::main) | (binary_id(etl-destinations::main) & test(/^(bigquery|clickhouse|ducklake|iceberg)::/)) | (binary_id(etl-destinations) & test(/ducklake::core::tests::postgres_backed::/))" +filter = "test(exclusive_) | binary_id(etl::main) | (binary_id(etl-destinations::main) & test(/^(bigquery|clickhouse|ducklake|iceberg|snowflake)::/)) | (binary_id(etl-destinations) & test(/ducklake::core::tests::postgres_backed::/))" test-group = "shared-pg" diff --git a/.env.example b/.env.example new file mode 100644 index 000000000..5b795fac7 --- /dev/null +++ b/.env.example @@ -0,0 +1,29 @@ +# Copy to .env and fill in your values: cp .env.example .env +# Then load before running tests: source .env +# Never commit .env to git. + +# PostgreSQL (required for replication tests) +export TESTS_DATABASE_HOST=localhost +export TESTS_DATABASE_PORT=5430 +export TESTS_DATABASE_USERNAME=postgres +export TESTS_DATABASE_PASSWORD=postgres + +# Snowflake (required for Snowflake integration tests) +# See etl-destinations/src/snowflake/README.md for key-pair setup instructions. +export TESTS_SNOWFLAKE_ACCOUNT= +export TESTS_SNOWFLAKE_USER= +export TESTS_SNOWFLAKE_PRIVATE_KEY_PATH= +export TESTS_SNOWFLAKE_DATABASE=ETL_DEV +export TESTS_SNOWFLAKE_SCHEMA=PUBLIC +export TESTS_SNOWFLAKE_WAREHOUSE= +export TESTS_SNOWFLAKE_ROLE= + +# Snowflake benchmark (used by the snowflake example binary) +# Can reuse the same account/key as tests, but targets a separate database. +export BENCH_SNOWFLAKE_ACCOUNT= +export BENCH_SNOWFLAKE_USER= +export BENCH_SNOWFLAKE_PRIVATE_KEY_PATH= +export BENCH_SNOWFLAKE_DATABASE=ETL_BENCH +export BENCH_SNOWFLAKE_SCHEMA=CDC +export BENCH_SNOWFLAKE_WAREHOUSE= +export BENCH_SNOWFLAKE_ROLE= diff --git a/.gitignore b/.gitignore index 267d9d4ba..2dc00e01b 100644 --- a/.gitignore +++ b/.gitignore @@ -14,6 +14,8 @@ target/ # Added by cargo /target + +# Environment files (credentials) .env # macOS system files @@ -42,3 +44,6 @@ configuration # Nix result + +# Local rust-analyzer config +.rust-analyzer.toml diff --git a/Cargo.lock b/Cargo.lock index e883f463c..6c5e74657 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -220,6 +220,17 @@ version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" +[[package]] +name = "aes" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "66bd29a732b644c0431c6140f370d097879203d79b80c94a6747ba0872adaef8" +dependencies = [ + "cipher", + "cpubits", + "cpufeatures 0.3.0", +] + [[package]] name = "ahash" version = "0.7.8" @@ -348,7 +359,7 @@ checksum = "36fa98bc79671c7981272d91a8753a928ff6a1cd8e4f20a44c45bd5d313840bf" dependencies = [ "bigdecimal", "bon", - "digest", + "digest 0.10.7", "log", "miniz_oxide", "num-bigint", @@ -403,78 +414,78 @@ checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50" [[package]] name = "arrow" -version = "57.3.0" +version = "57.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e4754a624e5ae42081f464514be454b39711daae0458906dacde5f4c632f33a8" +checksum = "3bd47f2a6ddc39244bd722a27ee5da66c03369d087b9e024eafdb03e98b98ea7" dependencies = [ - "arrow-arith 57.3.0", - "arrow-array 57.3.0", - "arrow-buffer 57.3.0", - "arrow-cast 57.3.0", - "arrow-data 57.3.0", - "arrow-ord 57.3.0", - "arrow-row 57.3.0", - "arrow-schema 57.3.0", - "arrow-select 57.3.0", - "arrow-string 57.3.0", + "arrow-arith 57.3.1", + "arrow-array 57.3.1", + "arrow-buffer 57.3.1", + "arrow-cast 57.3.1", + "arrow-data 57.3.1", + "arrow-ord 57.3.1", + "arrow-row 57.3.1", + "arrow-schema 57.3.1", + "arrow-select 57.3.1", + "arrow-string 57.3.1", ] [[package]] name = "arrow" -version = "58.1.0" +version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d441fdda254b65f3e9025910eb2c2066b6295d9c8ed409522b8d2ace1ff8574c" +checksum = "378530e55cd479eda3c14eb345310799717e6f76d0c332041e8487022166b471" dependencies = [ - "arrow-arith 58.1.0", - "arrow-array 58.1.0", - "arrow-buffer 58.1.0", - "arrow-cast 58.1.0", - "arrow-data 58.1.0", - "arrow-ord 58.1.0", - "arrow-row 58.1.0", - "arrow-schema 58.1.0", - "arrow-select 58.1.0", - "arrow-string 58.1.0", + "arrow-arith 58.3.0", + "arrow-array 58.3.0", + "arrow-buffer 58.3.0", + "arrow-cast 58.3.0", + "arrow-data 58.3.0", + "arrow-ord 58.3.0", + "arrow-row 58.3.0", + "arrow-schema 58.3.0", + "arrow-select 58.3.0", + "arrow-string 58.3.0", ] [[package]] name = "arrow-arith" -version = "57.3.0" +version = "57.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f7b3141e0ec5145a22d8694ea8b6d6f69305971c4fa1c1a13ef0195aef2d678b" +checksum = "7c7bbd679c5418b8639b92be01f361d60013c4906574b578b77b63c78356594c" dependencies = [ - "arrow-array 57.3.0", - "arrow-buffer 57.3.0", - "arrow-data 57.3.0", - "arrow-schema 57.3.0", + "arrow-array 57.3.1", + "arrow-buffer 57.3.1", + "arrow-data 57.3.1", + "arrow-schema 57.3.1", "chrono", "num-traits", ] [[package]] name = "arrow-arith" -version = "58.1.0" +version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ced5406f8b720cc0bc3aa9cf5758f93e8593cda5490677aa194e4b4b383f9a59" +checksum = "a0ab212d2c1886e802f51c5212d78ebbcbb0bec980fff9dadc1eb8d45cd0b738" dependencies = [ - "arrow-array 58.1.0", - "arrow-buffer 58.1.0", - "arrow-data 58.1.0", - "arrow-schema 58.1.0", + "arrow-array 58.3.0", + "arrow-buffer 58.3.0", + "arrow-data 58.3.0", + "arrow-schema 58.3.0", "chrono", "num-traits", ] [[package]] name = "arrow-array" -version = "57.3.0" +version = "57.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4c8955af33b25f3b175ee10af580577280b4bd01f7e823d94c7cdef7cf8c9aef" +checksum = "c8a4ab47b3f3eac60f7fd31b81e9028fda018607bcc63451aca4f2b755269862" dependencies = [ "ahash 0.8.12", - "arrow-buffer 57.3.0", - "arrow-data 57.3.0", - "arrow-schema 57.3.0", + "arrow-buffer 57.3.1", + "arrow-data 57.3.1", + "arrow-schema 57.3.1", "chrono", "half", "hashbrown 0.16.1", @@ -485,17 +496,17 @@ dependencies = [ [[package]] name = "arrow-array" -version = "58.1.0" +version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "772bd34cacdda8baec9418d80d23d0fb4d50ef0735685bd45158b83dfeb6e62d" +checksum = "cfd33d3e92f207444098c75b42de99d329562be0cf686b307b097cc52b4e999e" dependencies = [ "ahash 0.8.12", - "arrow-buffer 58.1.0", - "arrow-data 58.1.0", - "arrow-schema 58.1.0", + "arrow-buffer 58.3.0", + "arrow-data 58.3.0", + "arrow-schema 58.3.0", "chrono", "half", - "hashbrown 0.16.1", + "hashbrown 0.17.1", "num-complex", "num-integer", "num-traits", @@ -503,9 +514,9 @@ dependencies = [ [[package]] name = "arrow-buffer" -version = "57.3.0" +version = "57.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c697ddca96183182f35b3a18e50b9110b11e916d7b7799cbfd4d34662f2c56c2" +checksum = "0d18b89b4c4f4811d0858175e79541fe98e33e18db3b011708bc287b1240593f" dependencies = [ "bytes", "half", @@ -515,9 +526,9 @@ dependencies = [ [[package]] name = "arrow-buffer" -version = "58.1.0" +version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "898f4cf1e9598fdb77f356fdf2134feedfd0ee8d5a4e0a5f573e7d0aec16baa4" +checksum = "0c6cd424c2693bcdbc150d843dc9d4d137dd2de4782ce6df491ad11a3a0416c0" dependencies = [ "bytes", "half", @@ -527,16 +538,16 @@ dependencies = [ [[package]] name = "arrow-cast" -version = "57.3.0" +version = "57.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "646bbb821e86fd57189c10b4fcdaa941deaf4181924917b0daa92735baa6ada5" +checksum = "722b5c41dd1d14d0a879a1bce92c6fe33f546101bb2acce57a209825edd075b3" dependencies = [ - "arrow-array 57.3.0", - "arrow-buffer 57.3.0", - "arrow-data 57.3.0", - "arrow-ord 57.3.0", - "arrow-schema 57.3.0", - "arrow-select 57.3.0", + "arrow-array 57.3.1", + "arrow-buffer 57.3.1", + "arrow-data 57.3.1", + "arrow-ord 57.3.1", + "arrow-schema 57.3.1", + "arrow-select 57.3.1", "atoi", "base64", "chrono", @@ -548,16 +559,16 @@ dependencies = [ [[package]] name = "arrow-cast" -version = "58.1.0" +version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b0127816c96533d20fc938729f48c52d3e48f99717e7a0b5ade77d742510736d" +checksum = "4c5aefb56a2c02e9e2b30746241058b85f8983f0fcff2ba0c6d09006e1cded7f" dependencies = [ - "arrow-array 58.1.0", - "arrow-buffer 58.1.0", - "arrow-data 58.1.0", - "arrow-ord 58.1.0", - "arrow-schema 58.1.0", - "arrow-select 58.1.0", + "arrow-array 58.3.0", + "arrow-buffer 58.3.0", + "arrow-data 58.3.0", + "arrow-ord 58.3.0", + "arrow-schema 58.3.0", + "arrow-select 58.3.0", "atoi", "base64", "chrono", @@ -570,12 +581,12 @@ dependencies = [ [[package]] name = "arrow-data" -version = "57.3.0" +version = "57.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1fdd994a9d28e6365aa78e15da3f3950c0fdcea6b963a12fa1c391afb637b304" +checksum = "c1683705c63dcf0d18972759eda48489028cbbff67af7d6bef2c6b7b74ab778a" dependencies = [ - "arrow-buffer 57.3.0", - "arrow-schema 57.3.0", + "arrow-buffer 57.3.1", + "arrow-schema 57.3.1", "half", "num-integer", "num-traits", @@ -583,12 +594,12 @@ dependencies = [ [[package]] name = "arrow-data" -version = "58.1.0" +version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "42d10beeab2b1c3bb0b53a00f7c944a178b622173a5c7bcabc3cb45d90238df4" +checksum = "3c88210023a2bfee1896af366309a3028fc3bcbd6515fa29a7990ee1baa08ee0" dependencies = [ - "arrow-buffer 58.1.0", - "arrow-schema 58.1.0", + "arrow-buffer 58.3.0", + "arrow-schema 58.3.0", "half", "num-integer", "num-traits", @@ -596,124 +607,124 @@ dependencies = [ [[package]] name = "arrow-ipc" -version = "57.3.0" +version = "57.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "abf7df950701ab528bf7c0cf7eeadc0445d03ef5d6ffc151eaae6b38a58feff1" +checksum = "8cf72d04c07229fbf4dbebe7145cac37d7cf7ec582fe705c6b92cb314af096ab" dependencies = [ - "arrow-array 57.3.0", - "arrow-buffer 57.3.0", - "arrow-data 57.3.0", - "arrow-schema 57.3.0", - "arrow-select 57.3.0", + "arrow-array 57.3.1", + "arrow-buffer 57.3.1", + "arrow-data 57.3.1", + "arrow-schema 57.3.1", + "arrow-select 57.3.1", "flatbuffers", ] [[package]] name = "arrow-ord" -version = "57.3.0" +version = "57.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f7d8f1870e03d4cbed632959498bcc84083b5a24bded52905ae1695bd29da45b" +checksum = "082342947d4e5a2bcccf029a0a0397e21cb3bb8421edd9571d34fb5dd2670256" dependencies = [ - "arrow-array 57.3.0", - "arrow-buffer 57.3.0", - "arrow-data 57.3.0", - "arrow-schema 57.3.0", - "arrow-select 57.3.0", + "arrow-array 57.3.1", + "arrow-buffer 57.3.1", + "arrow-data 57.3.1", + "arrow-schema 57.3.1", + "arrow-select 57.3.1", ] [[package]] name = "arrow-ord" -version = "58.1.0" +version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "763a7ba279b20b52dad300e68cfc37c17efa65e68623169076855b3a9e941ca5" +checksum = "1bffd8fd2579286a5d63bac898159873e5094a79009940bcb42bbfce4f19f1d0" dependencies = [ - "arrow-array 58.1.0", - "arrow-buffer 58.1.0", - "arrow-data 58.1.0", - "arrow-schema 58.1.0", - "arrow-select 58.1.0", + "arrow-array 58.3.0", + "arrow-buffer 58.3.0", + "arrow-data 58.3.0", + "arrow-schema 58.3.0", + "arrow-select 58.3.0", ] [[package]] name = "arrow-row" -version = "57.3.0" +version = "57.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "18228633bad92bff92a95746bbeb16e5fc318e8382b75619dec26db79e4de4c0" +checksum = "e3a931b520a2a5e22033e01a6f2486b4cdc26f9106b759abeebc320f125e94d7" dependencies = [ - "arrow-array 57.3.0", - "arrow-buffer 57.3.0", - "arrow-data 57.3.0", - "arrow-schema 57.3.0", + "arrow-array 57.3.1", + "arrow-buffer 57.3.1", + "arrow-data 57.3.1", + "arrow-schema 57.3.1", "half", ] [[package]] name = "arrow-row" -version = "58.1.0" +version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e14fe367802f16d7668163ff647830258e6e0aeea9a4d79aaedf273af3bdcd3e" +checksum = "bab5994731204603c73ba69267616c50f80780774c6bb0476f1f830625115e0c" dependencies = [ - "arrow-array 58.1.0", - "arrow-buffer 58.1.0", - "arrow-data 58.1.0", - "arrow-schema 58.1.0", + "arrow-array 58.3.0", + "arrow-buffer 58.3.0", + "arrow-data 58.3.0", + "arrow-schema 58.3.0", "half", ] [[package]] name = "arrow-schema" -version = "57.3.0" +version = "57.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8c872d36b7bf2a6a6a2b40de9156265f0242910791db366a2c17476ba8330d68" +checksum = "e4cf0d4a6609679e03002167a61074a21d7b1ad9ea65e462b2c0a97f8a3b2bc6" [[package]] name = "arrow-schema" -version = "58.1.0" +version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c30a1365d7a7dc50cc847e54154e6af49e4c4b0fddc9f607b687f29212082743" +checksum = "f633dbfdf39c039ada1bf9e34c694816eb71fbb7dc78f613993b7245e078a1ed" dependencies = [ "bitflags", ] [[package]] name = "arrow-select" -version = "57.3.0" +version = "57.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "68bf3e3efbd1278f770d67e5dc410257300b161b93baedb3aae836144edcaf4b" +checksum = "0b320d86a9806923663bb0fd9baa65ecaba81cb0cd77ff8c1768b9716b4ef891" dependencies = [ "ahash 0.8.12", - "arrow-array 57.3.0", - "arrow-buffer 57.3.0", - "arrow-data 57.3.0", - "arrow-schema 57.3.0", + "arrow-array 57.3.1", + "arrow-buffer 57.3.1", + "arrow-data 57.3.1", + "arrow-schema 57.3.1", "num-traits", ] [[package]] name = "arrow-select" -version = "58.1.0" +version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "78694888660a9e8ac949853db393af2a8b8fc82c19ce333132dfa2e72cc1a7fe" +checksum = "8cd065c54172ac787cf3f2f8d4107e0d3fdc26edba76fdf4f4cc170258942222" dependencies = [ "ahash 0.8.12", - "arrow-array 58.1.0", - "arrow-buffer 58.1.0", - "arrow-data 58.1.0", - "arrow-schema 58.1.0", + "arrow-array 58.3.0", + "arrow-buffer 58.3.0", + "arrow-data 58.3.0", + "arrow-schema 58.3.0", "num-traits", ] [[package]] name = "arrow-string" -version = "57.3.0" +version = "57.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "85e968097061b3c0e9fe3079cf2e703e487890700546b5b0647f60fca1b5a8d8" +checksum = "b493e99162e5764077e7823e50ba284858d365922631c7aaefe9487b1abd02c2" dependencies = [ - "arrow-array 57.3.0", - "arrow-buffer 57.3.0", - "arrow-data 57.3.0", - "arrow-schema 57.3.0", - "arrow-select 57.3.0", + "arrow-array 57.3.1", + "arrow-buffer 57.3.1", + "arrow-data 57.3.1", + "arrow-schema 57.3.1", + "arrow-select 57.3.1", "memchr", "num-traits", "regex", @@ -722,15 +733,15 @@ dependencies = [ [[package]] name = "arrow-string" -version = "58.1.0" +version = "58.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "61e04a01f8bb73ce54437514c5fd3ee2aa3e8abe4c777ee5cc55853b1652f79e" +checksum = "29dd7cda3ab9692f43a2e4acc444d760cc17b12bb6d8232ddf64e9bab7c06b42" dependencies = [ - "arrow-array 58.1.0", - "arrow-buffer 58.1.0", - "arrow-data 58.1.0", - "arrow-schema 58.1.0", - "arrow-select 58.1.0", + "arrow-array 58.3.0", + "arrow-buffer 58.3.0", + "arrow-data 58.3.0", + "arrow-schema 58.3.0", + "arrow-select 58.3.0", "memchr", "num-traits", "regex", @@ -822,19 +833,20 @@ checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" [[package]] name = "aws-lc-rs" -version = "1.16.3" +version = "1.17.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0ec6fb3fe69024a75fa7e1bfb48aa6cf59706a101658ea01bfd33b2b248a038f" +checksum = "5ec2f1fc3ec205783a5da9a7e6c1509cc69dedf09a1949e412c1e18469326d00" dependencies = [ "aws-lc-sys", + "untrusted 0.7.1", "zeroize", ] [[package]] name = "aws-lc-sys" -version = "0.40.0" +version = "0.41.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f50037ee5e1e41e7b8f9d161680a725bd1626cb6f8c7e901f91f942850852fe7" +checksum = "1a2f9779ce85b93ab6170dd940ad0169b5766ff848247aff13bb788b832fe3f4" dependencies = [ "cc", "cmake", @@ -936,6 +948,24 @@ dependencies = [ "generic-array", ] +[[package]] +name = "block-buffer" +version = "0.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cdd35008169921d80bc60d3d0ab416eecb028c4cd653352907921d95084790be" +dependencies = [ + "hybrid-array", +] + +[[package]] +name = "block-padding" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "710f1dd022ef4e93f8a438b4ba958de7f64308434fa6a87104481645cc30068b" +dependencies = [ + "hybrid-array", +] + [[package]] name = "block2" version = "0.6.2" @@ -1021,6 +1051,15 @@ dependencies = [ "alloc-stdlib", ] +[[package]] +name = "bs58" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bf88ba1141d185c399bee5288d850d63b8369520c1eafc32a0430b5b6c287bf4" +dependencies = [ + "tinyvec", +] + [[package]] name = "bstr" version = "1.12.1" @@ -1081,9 +1120,9 @@ dependencies = [ [[package]] name = "bytestring" -version = "1.5.0" +version = "1.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "113b4343b5f6617e7ad401ced8de3cc8b012e73a594347c307b90db3e9271289" +checksum = "86566c496f2f47d9b8147a4c8b02ffdb69c919fe0c2b2e7195d22cbba0e635c9" dependencies = [ "bytes", ] @@ -1094,11 +1133,29 @@ version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" +[[package]] +name = "castaway" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dec551ab6e7578819132c713a93c022a05d60159dc86e7a7050223577484c55a" +dependencies = [ + "rustversion", +] + +[[package]] +name = "cbc" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "98db6aeaef0eeef2c1e3ce9a27b739218825dae116076352ac3777076aa22225" +dependencies = [ + "cipher", +] + [[package]] name = "cc" -version = "1.2.60" +version = "1.2.62" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "43c5703da9466b66a946814e1adf53ea2c90f10063b86290cc9eb67ce3478a20" +checksum = "a1dce859f0832a7d088c4f1119888ab94ef4b5d6795d1ce05afb7fe159d79f98" dependencies = [ "find-msvc-tools", "jobserver", @@ -1132,6 +1189,17 @@ dependencies = [ "windows-link", ] +[[package]] +name = "cipher" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e34d8227fe1ba289043aeb13792056ff80fd6de1a9f49137a5f499de8e8c78ea" +dependencies = [ + "block-buffer 0.12.0", + "crypto-common 0.2.1", + "inout", +] + [[package]] name = "cityhash-rs" version = "1.0.1" @@ -1237,6 +1305,12 @@ dependencies = [ "cc", ] +[[package]] +name = "cmov" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f88a43d011fc4a6876cb7344703e297c71dda42494fee094d5f7c76bf13f746" + [[package]] name = "colorchoice" version = "1.0.5" @@ -1249,11 +1323,25 @@ version = "7.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4a65ebfec4fb190b6f90e944a817d60499ee0744e582530e2c9900a22e591d9a" dependencies = [ - "crossterm", + "crossterm 0.28.1", "unicode-segmentation", "unicode-width", ] +[[package]] +name = "compact_str" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3fdb1325a1cece981e8a296ab8f0f9b63ae357bd0784a9faaf548cc7b480707a" +dependencies = [ + "castaway", + "cfg-if", + "itoa", + "rustversion", + "ryu", + "static_assertions", +] + [[package]] name = "concurrent-queue" version = "2.5.0" @@ -1292,7 +1380,7 @@ dependencies = [ "serde_json", "serde_repr", "sha1", - "sha2", + "sha2 0.10.9", "thiserror 1.0.69", "tokio", "tokio-util", @@ -1304,6 +1392,12 @@ version = "0.9.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2459377285ad874054d797f3ccebf984978aa39129f6eafde5cdc8315b612f8" +[[package]] +name = "const-oid" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6ef517f0926dd24a1582492c791b6a4818a4d94e789a334894aa15b0d12f55c" + [[package]] name = "const-random" version = "0.1.18" @@ -1365,6 +1459,12 @@ version = "0.8.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" +[[package]] +name = "cpubits" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "15b85f9c39137c3a891689859392b1bd49812121d0d61c9caf00d46ed5ce06ae" + [[package]] name = "cpufeatures" version = "0.2.17" @@ -1374,6 +1474,15 @@ dependencies = [ "libc", ] +[[package]] +name = "cpufeatures" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b2a41393f66f16b0823bb79094d54ac5fbd34ab292ddafb9a0456ac9f87d201" +dependencies = [ + "libc", +] + [[package]] name = "crc" version = "3.4.0" @@ -1385,9 +1494,9 @@ dependencies = [ [[package]] name = "crc-catalog" -version = "2.4.0" +version = "2.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "19d374276b40fb8bbdee95aef7c7fa6b5316ec764510eb64b8dd0e2ed0d7e7f5" +checksum = "217698eaf96b4a3f0bc4f3662aaa55bdf913cd54d7204591faa790070c6d0853" [[package]] name = "crc32c" @@ -1453,6 +1562,24 @@ dependencies = [ "winapi", ] +[[package]] +name = "crossterm" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d8b9f2e4c67f833b660cdb0a3523065869fb35570177239812ed4c905aeff87b" +dependencies = [ + "bitflags", + "crossterm_winapi", + "derive_more", + "document-features", + "mio", + "parking_lot", + "rustix 1.1.4", + "signal-hook", + "signal-hook-mio", + "winapi", +] + [[package]] name = "crossterm_winapi" version = "0.9.1" @@ -1470,14 +1597,32 @@ checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" [[package]] name = "crypto-common" -version = "0.1.7" +version = "0.1.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" +checksum = "1bfb12502f3fc46cca1bb51ac28df9d618d813cdc3d2f25b9fe775a34af26bb3" dependencies = [ "generic-array", "typenum", ] +[[package]] +name = "crypto-common" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77727bb15fa921304124b128af125e7e3b968275d1b108b379190264f4423710" +dependencies = [ + "hybrid-array", +] + +[[package]] +name = "ctutils" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d5515a3834141de9eafb9717ad39eea8247b5674e6066c404e8c4b365d2a29e" +dependencies = [ + "cmov", +] + [[package]] name = "darling" version = "0.20.11" @@ -1563,7 +1708,7 @@ version = "0.7.10" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e7c1832837b905bbfb5101e07cc24c8deddf52f93225eee6ead5f4d63d53ddcb" dependencies = [ - "const-oid", + "const-oid 0.9.6", "der_derive", "flagset", "pem-rfc7468 0.7.0", @@ -1576,6 +1721,7 @@ version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "71fd89660b2dc699704064e59e9dba0147b903e85319429e131620d022be411b" dependencies = [ + "const-oid 0.10.2", "pem-rfc7468 1.0.0", "zeroize", ] @@ -1672,12 +1818,23 @@ version = "0.10.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" dependencies = [ - "block-buffer", - "const-oid", - "crypto-common", + "block-buffer 0.10.4", + "const-oid 0.9.6", + "crypto-common 0.1.6", "subtle", ] +[[package]] +name = "digest" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f1dd6dbb5841937940781866fa1281a1ff7bd3bf827091440879f9994983d5c2" +dependencies = [ + "block-buffer 0.12.0", + "crypto-common 0.2.1", + "ctutils", +] + [[package]] name = "dispatch2" version = "0.3.1" @@ -1714,6 +1871,15 @@ dependencies = [ "const-random", ] +[[package]] +name = "document-features" +version = "0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d4b8a88685455ed29a21542a33abd9cb6510b6b129abadabdcef0f4c55bc8f61" +dependencies = [ + "litrs", +] + [[package]] name = "dotenvy" version = "0.15.7" @@ -1726,7 +1892,7 @@ version = "1.10502.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0fdc796383b176dd5a45353fbb5e64583c0ee4da12cb62c9e510b785324b2488" dependencies = [ - "arrow 58.1.0", + "arrow 58.3.0", "cast", "comfy-table", "fallible-iterator 0.3.0", @@ -1836,7 +2002,7 @@ dependencies = [ "byteorder", "bytes", "chrono", - "const-oid", + "const-oid 0.9.6", "etl-config", "etl-postgres", "etl-telemetry", @@ -1943,9 +2109,11 @@ dependencies = [ name = "etl-destinations" version = "0.1.0" dependencies = [ - "arrow 57.3.0", + "arrow 57.3.1", "async-trait", + "aws-lc-rs", "base64", + "bytes", "chrono", "clickhouse", "duckdb", @@ -1956,30 +2124,38 @@ dependencies = [ "etl-telemetry", "futures", "gcp-bigquery-client", + "hex", "humantime", "iceberg", "iceberg-catalog-rest", + "jsonwebtoken", "k8s-openapi", + "kube", "metrics", "parking_lot", "parquet", "pg_escape", + "pkcs8 0.11.0", "prost", "r2d2", "rand 0.9.4", "regex", "reqwest", "rustls", + "secrecy", "serde", "serde_json", + "sha2 0.11.0", "sqlx", "tempfile", + "thiserror 2.0.18", "tokio", "tokio-postgres", "tonic", "tracing", "url", "uuid", + "zstd", ] [[package]] @@ -1987,13 +2163,19 @@ name = "etl-examples" version = "0.1.0" dependencies = [ "clap", + "crossterm 0.29.0", "etl", "etl-config", "etl-destinations", "etl-telemetry", "k8s-openapi", + "rand 0.9.4", + "ratatui", + "reqwest", "rustls", + "secrecy", "tokio", + "tokio-postgres", "tracing", "tracing-subscriber", "url", @@ -2151,13 +2333,12 @@ checksum = "9f1f227452a390804cdb637b74a86990f2a7d7ba4b7d5693aac9b4dd6defd8d6" [[package]] name = "filetime" -version = "0.2.27" +version = "0.2.29" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f98844151eee8917efc50bd9e8318cb963ae8b297431495d3f758616ea5c57db" +checksum = "5c287a33c7f0a620c38e641e7f60827713987b3c0f26e8ddc9462cc69cf75759" dependencies = [ "cfg-if", "libc", - "libredox", ] [[package]] @@ -2408,9 +2589,9 @@ dependencies = [ [[package]] name = "generic-array" -version = "0.14.7" +version = "0.14.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +checksum = "4bb6743198531e02858aeaea5398fcc883e71851fcbcb5a2f773e2fb6cb1edf2" dependencies = [ "typenum", "version_check", @@ -2495,9 +2676,9 @@ dependencies = [ [[package]] name = "h2" -version = "0.4.13" +version = "0.4.14" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2f44da3a8150a6703ed5d34e164b875fd14c2cdab9af1252a9a1020bde2bdc54" +checksum = "171fefbc92fe4a4de27e0698d6a5b392d6a0e333506bc49133760b3bcf948733" dependencies = [ "atomic-waker", "bytes", @@ -2567,9 +2748,9 @@ dependencies = [ [[package]] name = "hashbrown" -version = "0.17.0" +version = "0.17.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4f467dd6dccf739c208452f8014c75c18bb8301b050ad1cfb27153803edb0f51" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" [[package]] name = "hashlink" @@ -2642,7 +2823,7 @@ version = "0.12.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7b5f8eb2ad728638ea2c7d47a21db23b7b58a72ed6a38256b8a1849f15fbbdf7" dependencies = [ - "hmac", + "hmac 0.12.1", ] [[package]] @@ -2651,7 +2832,16 @@ version = "0.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6c49c37c09c17a53d937dfbb742eb3a961d65a994e6bcdcf37e7399d0cc8ab5e" dependencies = [ - "digest", + "digest 0.10.7", +] + +[[package]] +name = "hmac" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6303bc9732ae41b04cb554b844a762b4115a61bfaa81e3e83050991eeb56863f" +dependencies = [ + "digest 0.11.3", ] [[package]] @@ -2736,6 +2926,15 @@ version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "135b12329e5e3ce057a9f972339ea52bc954fe1e9358ef27f95e89716fbc5424" +[[package]] +name = "hybrid-array" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9155a582abd142abc056962c29e3ce5ff2ad5469f4246b537ed42c5deba857da" +dependencies = [ + "typenum", +] + [[package]] name = "hyper" version = "1.9.0" @@ -2746,7 +2945,7 @@ dependencies = [ "bytes", "futures-channel", "futures-core", - "h2 0.4.13", + "h2 0.4.14", "http 1.4.0", "http-body", "httparse", @@ -2883,14 +3082,14 @@ dependencies = [ "anyhow", "apache-avro", "array-init", - "arrow-arith 57.3.0", - "arrow-array 57.3.0", - "arrow-buffer 57.3.0", - "arrow-cast 57.3.0", - "arrow-ord 57.3.0", - "arrow-schema 57.3.0", - "arrow-select 57.3.0", - "arrow-string 57.3.0", + "arrow-arith 57.3.1", + "arrow-array 57.3.1", + "arrow-buffer 57.3.1", + "arrow-cast 57.3.1", + "arrow-ord 57.3.1", + "arrow-schema 57.3.1", + "arrow-select 57.3.1", + "arrow-string 57.3.1", "as-any", "async-trait", "backon", @@ -3058,9 +3257,9 @@ dependencies = [ [[package]] name = "idna_adapter" -version = "1.2.1" +version = "1.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3acae9609540aa318d1bc588455225fb2085b9ed0c4f6bd0d9d5bcd86f1a0344" +checksum = "cb68373c0d6620ef8105e855e7745e18b0d00d3bdb07fb532e434244cdb9a714" dependencies = [ "icu_normalizer", "icu_properties", @@ -3090,11 +3289,30 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" dependencies = [ "equivalent", - "hashbrown 0.17.0", + "hashbrown 0.17.1", "serde", "serde_core", ] +[[package]] +name = "indoc" +version = "2.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "79cf5c93f93228cf8efb3ba362535fb11199ac548a09ce117c9b1adc3030d706" +dependencies = [ + "rustversion", +] + +[[package]] +name = "inout" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4250ce6452e92010fdf7268ccc5d14faa80bb12fc741938534c58f16804e03c7" +dependencies = [ + "block-padding", + "hybrid-array", +] + [[package]] name = "insta" version = "1.47.2" @@ -3109,6 +3327,19 @@ dependencies = [ "tempfile", ] +[[package]] +name = "instability" +version = "0.3.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5eb2d60ef19920a3a9193c3e371f726ec1dafc045dac788d0fb3704272458971" +dependencies = [ + "darling 0.23.0", + "indoc", + "proc-macro2", + "quote", + "syn 2.0.117", +] + [[package]] name = "integer-encoding" version = "3.0.4" @@ -3121,16 +3352,6 @@ version = "2.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d98f6fed1fde3f8c21bc40a1abb88dd75e67924f9cffc3ef95607bad8017f8e2" -[[package]] -name = "iri-string" -version = "0.7.12" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "25e659a4bb38e810ebc252e53b5814ff908a8c58c2a9ce2fae1bbec24cbf4e20" -dependencies = [ - "memchr", - "serde", -] - [[package]] name = "is_terminal_polyfill" version = "1.70.2" @@ -3163,9 +3384,9 @@ checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" [[package]] name = "jiff" -version = "0.2.23" +version = "0.2.24" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1a3546dc96b6d42c5f24902af9e2538e82e39ad350b0c766eb3fbf2d8f3d8359" +checksum = "f00b5dbd620d61dfdcb6007c9c1f6054ebd75319f163d886a9055cec1155073d" dependencies = [ "jiff-static", "jiff-tzdb-platform", @@ -3178,9 +3399,9 @@ dependencies = [ [[package]] name = "jiff-static" -version = "0.2.23" +version = "0.2.24" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2a8c8b344124222efd714b73bb41f8b5120b27a7cc1c75593a6ff768d9d05aa4" +checksum = "e000de030ff8022ea1da3f466fbb0f3a809f5e51ed31f6dd931c35181ad8e6d7" dependencies = [ "proc-macro2", "quote", @@ -3214,9 +3435,9 @@ dependencies = [ [[package]] name = "js-sys" -version = "0.3.95" +version = "0.3.98" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2964e92d1d9dc3364cae4d718d93f227e3abb088e747d92e0395bfdedf1c12ca" +checksum = "67df7112613f8bfd9150013a0314e196f4800d3201ae742489d999db2f979f08" dependencies = [ "cfg-if", "futures-util", @@ -3226,14 +3447,14 @@ dependencies = [ [[package]] name = "json-patch" -version = "4.1.0" +version = "4.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f300e415e2134745ef75f04562dd0145405c2f7fd92065db029ac4b16b57fe90" +checksum = "7421438de105a0827e44fadd05377727847d717c80ce29a229f85fd04c427b72" dependencies = [ "jsonptr", "serde", "serde_json", - "thiserror 1.0.69", + "thiserror 2.0.18", ] [[package]] @@ -3259,6 +3480,22 @@ dependencies = [ "serde_json", ] +[[package]] +name = "jsonwebtoken" +version = "10.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eba32bfb4ffdeaca3e34431072faf01745c9b26d25504aa7a6cf5684334fc4fc" +dependencies = [ + "aws-lc-rs", + "base64", + "getrandom 0.2.17", + "js-sys", + "serde", + "serde_json", + "signature", + "zeroize", +] + [[package]] name = "k8s-openapi" version = "0.25.0" @@ -3271,6 +3508,17 @@ dependencies = [ "serde_json", ] +[[package]] +name = "kasuari" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bde5057d6143cc94e861d90f591b9303d6716c6b9602309150bd068853c10899" +dependencies = [ + "hashbrown 0.16.1", + "portable-atomic", + "thiserror 2.0.18", +] + [[package]] name = "kube" version = "1.1.0" @@ -3461,9 +3709,9 @@ dependencies = [ [[package]] name = "libc" -version = "0.2.185" +version = "0.2.186" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "52ff2c0fe9bc6cb6b14a0592c2ff4fa9ceb83eea9db979b0487cd054946a2b8f" +checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" [[package]] name = "libduckdb-sys" @@ -3497,19 +3745,28 @@ dependencies = [ "bitflags", "libc", "plain", - "redox_syscall 0.7.4", + "redox_syscall 0.7.5", ] [[package]] name = "libsqlite3-sys" -version = "0.30.1" +version = "0.35.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2e99fb7a497b1e3339bc746195567ed8d3e24945ecd636e3619d20b9de9e9149" +checksum = "133c182a6a2c87864fe97778797e46c7e999672690dc9fa3ee8e241aa4a9c13f" dependencies = [ "pkg-config", "vcpkg", ] +[[package]] +name = "line-clipping" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f50e8f47623268b5407192d26876c4d7f89d686ca130fdc53bced4814cd29f8" +dependencies = [ + "bitflags", +] + [[package]] name = "linux-raw-sys" version = "0.4.15" @@ -3528,6 +3785,12 @@ version = "0.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "92daf443525c4cce67b150400bc2316076100ce0b3686209eb8cf3c31612e6f0" +[[package]] +name = "litrs" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11d3d7f243d5c5a8b9bb5d6dd2b1602c0cb0b9db1621bafc7ed66e35ff9fe092" + [[package]] name = "local-waker" version = "0.1.4" @@ -3549,6 +3812,15 @@ version = "0.4.29" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" +[[package]] +name = "lru" +version = "0.16.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f66e8d5d03f609abc3a39e6f08e4164ebf1447a732906d39eb9b99b7919ef39" +dependencies = [ + "hashbrown 0.16.1", +] + [[package]] name = "lru-slab" version = "0.1.2" @@ -3563,9 +3835,9 @@ checksum = "373f5eceeeab7925e0c1098212f2fbc4d416adec9d35051a6ab251e824c1854a" [[package]] name = "lz4_flex" -version = "0.12.1" +version = "0.12.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "98c23545df7ecf1b16c303910a69b079e8e251d60f7dd2cc9b4177f2afaf1746" +checksum = "90071f8077f8e40adfc4b7fe9cd495ce316263f19e75c2211eeff3fdf475a3d9" dependencies = [ "twox-hash", ] @@ -3602,7 +3874,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d89e7ee0cfbedfc4da3340218492196241d89eefb6dab27de5df917a6d2e78cf" dependencies = [ "cfg-if", - "digest", + "digest 0.10.7", ] [[package]] @@ -3613,12 +3885,12 @@ checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79" [[package]] name = "metrics" -version = "0.24.3" +version = "0.24.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5d5312e9ba3771cfa961b585728215e3d972c950a3eed9252aa093d6301277e8" +checksum = "89550ee9f79e88fef3119de263694973a8adb26c21d75322164fb8c493039fe2" dependencies = [ - "ahash 0.8.12", "portable-atomic", + "rapidhash", ] [[package]] @@ -3643,9 +3915,9 @@ dependencies = [ [[package]] name = "metrics-util" -version = "0.20.1" +version = "0.20.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cdfb1365fea27e6dd9dc1dbc19f570198bc86914533ad639dae939635f096be4" +checksum = "96f8722f8562635f92f8ed992f26df0532266eb03d5202607c20c0d7e9745e13" dependencies = [ "crossbeam-epoch", "crossbeam-utils", @@ -3654,6 +3926,7 @@ dependencies = [ "quanta", "rand 0.9.4", "rand_xoshiro", + "rapidhash", "sketches-ddsketch", ] @@ -4104,9 +4377,9 @@ dependencies = [ [[package]] name = "openssl" -version = "0.10.79" +version = "0.10.80" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bf0b434746ee2832f4f0baf10137e1cabb18cbe6912c69e2e33263c45250f542" +checksum = "a45fa2aa886c42762255da344f0a0d313e254066c46aad76f300c3d3da62d967" dependencies = [ "bitflags", "cfg-if", @@ -4141,9 +4414,9 @@ checksum = "7c87def4c32ab89d880effc9e097653c8da5d6ef28e6b539d313baaacfbafcbe" [[package]] name = "openssl-sys" -version = "0.9.115" +version = "0.9.116" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "158fe5b292746440aa6e7a7e690e55aeb72d41505e2804c23c6973ad0e9c9781" +checksum = "f28a22dc7140cda5f096e5e7724a6962ca81a7f8bfd2979f9b18c11af56318c4" dependencies = [ "cc", "libc", @@ -4226,18 +4499,18 @@ dependencies = [ [[package]] name = "parquet" -version = "57.3.0" +version = "57.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6ee96b29972a257b855ff2341b37e61af5f12d6af1158b6dcdb5b31ea07bb3cb" +checksum = "2e832c6aa20310fc6de7ea5a3f4e20d34fd83e3b43229d32b81ffe5c14d74692" dependencies = [ "ahash 0.8.12", - "arrow-array 57.3.0", - "arrow-buffer 57.3.0", - "arrow-cast 57.3.0", - "arrow-data 57.3.0", + "arrow-array 57.3.1", + "arrow-buffer 57.3.1", + "arrow-cast 57.3.1", + "arrow-data 57.3.1", "arrow-ipc", - "arrow-schema 57.3.0", - "arrow-select 57.3.0", + "arrow-schema 57.3.1", + "arrow-select 57.3.1", "base64", "brotli", "bytes", @@ -4246,7 +4519,7 @@ dependencies = [ "futures", "half", "hashbrown 0.16.1", - "lz4_flex 0.12.1", + "lz4_flex 0.12.2", "num-bigint", "num-integer", "num-traits", @@ -4272,6 +4545,16 @@ version = "0.2.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "df94ce210e5bc13cb6651479fa48d14f601d9858cfe0467f43ae157023b938d3" +[[package]] +name = "pbkdf2" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "112d82ceb8c5bf524d9af484d4e4970c9fd5a0cc15ba14ad93dccd28873b0629" +dependencies = [ + "digest 0.11.3", + "hmac 0.13.0", +] + [[package]] name = "pem" version = "3.0.6" @@ -4346,7 +4629,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "89815c69d36021a140146f26659a81d6c2afa33d216d736dd4be5381a7362220" dependencies = [ "pest", - "sha2", + "sha2 0.10.9", ] [[package]] @@ -4413,18 +4696,18 @@ dependencies = [ [[package]] name = "pin-project" -version = "1.1.11" +version = "1.1.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f1749c7ed4bcaf4c3d0a3efc28538844fb29bcdd7d2b67b2be7e20ba861ff517" +checksum = "2466b2336ed02bcdca6b294417127b90ec92038d1d5c4fbeac971a922e0e0924" dependencies = [ "pin-project-internal", ] [[package]] name = "pin-project-internal" -version = "1.1.11" +version = "1.1.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d9b20ed30f105399776b9c883e68e536ef602a16ae6f596d2c473591d6ad64c6" +checksum = "c96395f0a926bc13b1c17622aaddda1ecb55d49c8f1bf9777e4d877800a43f8b" dependencies = [ "proc-macro2", "quote", @@ -4444,8 +4727,24 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c8ffb9f10fa047879315e6625af03c164b16962a5368d724ed16323b68ace47f" dependencies = [ "der 0.7.10", - "pkcs8", - "spki", + "pkcs8 0.10.2", + "spki 0.7.3", +] + +[[package]] +name = "pkcs5" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "279a91971a1d8eb1260a30938eae3be9cb67b472dffecb222fbbbe2fd2dc1453" +dependencies = [ + "aes", + "cbc", + "der 0.8.0", + "pbkdf2", + "rand_core 0.10.1", + "scrypt", + "sha2 0.11.0", + "spki 0.8.0", ] [[package]] @@ -4455,7 +4754,19 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f950b2377845cebe5cf8b5165cb3cc1a5e0fa5cfa3e1f7f55707d8fd82e0a7b7" dependencies = [ "der 0.7.10", - "spki", + "spki 0.7.3", +] + +[[package]] +name = "pkcs8" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "451913da69c775a56034ea8d9003d27ee8948e12443eae7c038ba100a4f21cb7" +dependencies = [ + "der 0.8.0", + "pkcs5", + "rand_core 0.10.1", + "spki 0.8.0", ] [[package]] @@ -4504,11 +4815,11 @@ dependencies = [ "byteorder", "bytes", "fallible-iterator 0.2.0", - "hmac", + "hmac 0.12.1", "md-5", "memchr", "rand 0.8.6", - "sha2", + "sha2 0.10.9", "stringprep", ] @@ -4879,6 +5190,12 @@ dependencies = [ "getrandom 0.3.4", ] +[[package]] +name = "rand_core" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69" + [[package]] name = "rand_xoshiro" version = "0.7.0" @@ -4888,6 +5205,78 @@ dependencies = [ "rand_core 0.9.5", ] +[[package]] +name = "rapidhash" +version = "4.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b5e48930979c155e2f33aa36ab3119b5ee81332beb6482199a8ecd6029b80b59" +dependencies = [ + "rustversion", +] + +[[package]] +name = "ratatui" +version = "0.30.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d1ce67fb8ba4446454d1c8dbaeda0557ff5e94d39d5e5ed7f10a65eb4c8266bc" +dependencies = [ + "instability", + "ratatui-core", + "ratatui-crossterm", + "ratatui-widgets", +] + +[[package]] +name = "ratatui-core" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5ef8dea09a92caaf73bff7adb70b76162e5937524058a7e5bff37869cbbec293" +dependencies = [ + "bitflags", + "compact_str", + "hashbrown 0.16.1", + "indoc", + "itertools 0.14.0", + "kasuari", + "lru", + "strum", + "thiserror 2.0.18", + "unicode-segmentation", + "unicode-truncate", + "unicode-width", +] + +[[package]] +name = "ratatui-crossterm" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "577c9b9f652b4c121fb25c6a391dd06406d3b092ba68827e6d2f09550edc54b3" +dependencies = [ + "cfg-if", + "crossterm 0.29.0", + "instability", + "ratatui-core", +] + +[[package]] +name = "ratatui-widgets" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7dbfa023cd4e604c2553483820c5fe8aa9d71a42eea5aa77c6e7f35756612db" +dependencies = [ + "bitflags", + "hashbrown 0.16.1", + "indoc", + "instability", + "itertools 0.14.0", + "line-clipping", + "ratatui-core", + "strum", + "time", + "unicode-segmentation", + "unicode-width", +] + [[package]] name = "raw-cpuid" version = "11.6.0" @@ -4908,9 +5297,9 @@ dependencies = [ [[package]] name = "redox_syscall" -version = "0.7.4" +version = "0.7.5" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f450ad9c3b1da563fb6948a8e0fb0fb9269711c9c73d9ea1de5058c79c8d643a" +checksum = "4666a1a60d8412eab19d94f6d13dcc9cea0a5ef4fdf6a5db306537413c661b1b" dependencies = [ "bitflags", ] @@ -4992,7 +5381,7 @@ dependencies = [ "form_urlencoded", "getrandom 0.2.17", "hex", - "hmac", + "hmac 0.12.1", "home", "http 1.4.0", "log", @@ -5004,7 +5393,7 @@ dependencies = [ "serde", "serde_json", "sha1", - "sha2", + "sha2 0.10.9", "tokio", ] @@ -5020,7 +5409,7 @@ dependencies = [ "futures-channel", "futures-core", "futures-util", - "h2 0.4.13", + "h2 0.4.14", "http 1.4.0", "http-body", "http-body-util", @@ -5066,7 +5455,7 @@ dependencies = [ "cfg-if", "getrandom 0.2.17", "libc", - "untrusted", + "untrusted 0.9.0", "windows-sys 0.52.0", ] @@ -5101,9 +5490,9 @@ dependencies = [ [[package]] name = "roaring" -version = "0.11.3" +version = "0.11.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8ba9ce64a8f45d7fc86358410bb1a82e8c987504c0d4900e9141d69a9f26c885" +checksum = "1dedc5658c6ecb3bdb5ef5f3295bb9253f42dcf3fd1402c03f6b1f7659c3c4a9" dependencies = [ "bytemuck", "byteorder", @@ -5115,16 +5504,16 @@ version = "0.9.10" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b8573f03f5883dcaebdfcf4725caa1ecb9c15b2ef50c43a07b816e06799bb12d" dependencies = [ - "const-oid", - "digest", + "const-oid 0.9.6", + "digest 0.10.7", "num-bigint-dig", "num-integer", "num-traits", "pkcs1", - "pkcs8", + "pkcs8 0.10.2", "rand_core 0.6.4", "signature", - "spki", + "spki 0.7.3", "subtle", "zeroize", ] @@ -5159,7 +5548,7 @@ version = "8.11.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5bcdef0be6fe7f6fa333b1073c949729274b05f123a0ad7efcb8efd878e5c3b1" dependencies = [ - "sha2", + "sha2 0.10.9", "walkdir", ] @@ -5175,9 +5564,9 @@ dependencies = [ [[package]] name = "rust_decimal" -version = "1.41.0" +version = "1.42.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2ce901f9a19d251159075a4c37af514c3b8ef99c22e02dd8c19161cf397ee94a" +checksum = "0c5108e3d4d903e21aac27f12ba5377b6b34f9f44b325e4894c7924169d06995" dependencies = [ "arrayvec", "borsh", @@ -5239,9 +5628,9 @@ dependencies = [ [[package]] name = "rustls" -version = "0.23.38" +version = "0.23.40" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "69f9466fb2c14ea04357e91413efb882e2a6d4a406e625449bc0a5d360d53a21" +checksum = "ef86cd5876211988985292b91c96a8f2d298df24e75989a43a3c73f2d4d8168b" dependencies = [ "aws-lc-rs", "log", @@ -5289,9 +5678,9 @@ dependencies = [ [[package]] name = "rustls-pki-types" -version = "1.14.0" +version = "1.14.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "be040f8b0a225e40375822a563fa9524378b9d63112f53e19ffff34df5d33fdd" +checksum = "30a7197ae7eb376e574fe940d068c30fe0462554a3ddbe4eca7838e049c937a9" dependencies = [ "web-time", "zeroize", @@ -5306,7 +5695,7 @@ dependencies = [ "aws-lc-rs", "ring", "rustls-pki-types", - "untrusted", + "untrusted 0.9.0", ] [[package]] @@ -5321,6 +5710,16 @@ version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" +[[package]] +name = "salsa20" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2f874456e72520ff1375a06c588eaf074b0f01f9e9e1aada45bd9b7954a6e42c" +dependencies = [ + "cfg-if", + "cipher", +] + [[package]] name = "same-file" version = "1.0.6" @@ -5402,6 +5801,18 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" +[[package]] +name = "scrypt" +version = "0.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d87af57419b594aa23fa95f09f0e06d80d84ba01c26148c43844cad6ff4485f0" +dependencies = [ + "cfg-if", + "pbkdf2", + "salsa20", + "sha2 0.11.0", +] + [[package]] name = "seahash" version = "4.1.0" @@ -5686,11 +6097,12 @@ dependencies = [ [[package]] name = "serde_with" -version = "3.18.0" +version = "3.20.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dd5414fad8e6907dbdd5bc441a50ae8d6e26151a03b1de04d89a5576de61d01f" +checksum = "e72c1c2cb7b223fafb600a619537a871c2818583d619401b785e7c0b746ccde2" dependencies = [ "base64", + "bs58", "chrono", "hex", "indexmap 1.9.3", @@ -5705,9 +6117,9 @@ dependencies = [ [[package]] name = "serde_with_macros" -version = "3.18.0" +version = "3.20.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d3db8978e608f1fe7357e211969fd9abdcae80bac1ba7a3369bb7eb6b404eb65" +checksum = "b90c488738ecb4fb0262f41f43bc40efc5868d9fb744319ddf5f5317f417bfac" dependencies = [ "darling 0.23.0", "proc-macro2", @@ -5735,8 +6147,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e3bf829a2d51ab4a5ddf1352d8470c140cadc8301b2ae1789db023f01cedd6ba" dependencies = [ "cfg-if", - "cpufeatures", - "digest", + "cpufeatures 0.2.17", + "digest 0.10.7", ] [[package]] @@ -5746,8 +6158,19 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" dependencies = [ "cfg-if", - "cpufeatures", - "digest", + "cpufeatures 0.2.17", + "digest 0.10.7", +] + +[[package]] +name = "sha2" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "446ba717509524cb3f22f17ecc096f10f4822d76ab5c0b9822c5f9c284e825f4" +dependencies = [ + "cfg-if", + "cpufeatures 0.3.0", + "digest 0.11.3", ] [[package]] @@ -5765,6 +6188,27 @@ version = "1.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" +[[package]] +name = "signal-hook" +version = "0.3.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d881a16cf4426aa584979d30bd82cb33429027e42122b169753d6ef1085ed6e2" +dependencies = [ + "libc", + "signal-hook-registry", +] + +[[package]] +name = "signal-hook-mio" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b75a19a7a740b25bc7944bdee6172368f988763b744e3d4dfe753f6b4ece40cc" +dependencies = [ + "libc", + "mio", + "signal-hook", +] + [[package]] name = "signal-hook-registry" version = "1.4.8" @@ -5781,7 +6225,7 @@ version = "2.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "77549399552de45a898a580c1b41d445bf730df867cc44e6c0233bbc4b8329de" dependencies = [ - "digest", + "digest 0.10.7", "rand_core 0.6.4", ] @@ -5805,9 +6249,9 @@ checksum = "bbbb5d9659141646ae647b42fe094daf6c6192d1620870b449d9557f748b2daa" [[package]] name = "siphasher" -version = "1.0.2" +version = "1.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b2aa850e253778c88a04c3d7323b043aeda9d3e30d5971937c1855769763678e" +checksum = "8ee5873ec9cce0195efcb7a4e9507a04cd49aec9c83d0389df45b1ef7ba2e649" [[package]] name = "sketches-ddsketch" @@ -5875,6 +6319,16 @@ dependencies = [ "der 0.7.10", ] +[[package]] +name = "spki" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d9efca8738c78ee9484207732f728b1ef517bbb1833d6fc0879ca898a522f6f" +dependencies = [ + "base64ct", + "der 0.8.0", +] + [[package]] name = "sqlx" version = "0.9.0-alpha.1" @@ -5915,7 +6369,7 @@ dependencies = [ "rustls", "serde", "serde_json", - "sha2", + "sha2 0.10.9", "smallvec", "thiserror 2.0.18", "tokio", @@ -5953,7 +6407,7 @@ dependencies = [ "quote", "serde", "serde_json", - "sha2", + "sha2 0.10.9", "sqlx-core", "sqlx-mysql", "sqlx-postgres", @@ -5977,7 +6431,7 @@ dependencies = [ "bytes", "chrono", "crc", - "digest", + "digest 0.10.7", "dotenvy", "either", "futures-channel", @@ -5987,7 +6441,7 @@ dependencies = [ "generic-array", "hex", "hkdf", - "hmac", + "hmac 0.12.1", "itoa", "log", "md-5", @@ -5997,7 +6451,7 @@ dependencies = [ "rsa", "serde", "sha1", - "sha2", + "sha2 0.10.9", "smallvec", "sqlx-core", "stringprep", @@ -6025,7 +6479,7 @@ dependencies = [ "futures-util", "hex", "hkdf", - "hmac", + "hmac 0.12.1", "home", "itoa", "log", @@ -6034,7 +6488,7 @@ dependencies = [ "rand 0.8.6", "serde", "serde_json", - "sha2", + "sha2 0.10.9", "smallvec", "sqlx-core", "stringprep", @@ -6074,6 +6528,12 @@ version = "1.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" +[[package]] +name = "static_assertions" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" + [[package]] name = "strfmt" version = "0.2.5" @@ -6403,9 +6863,9 @@ checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" [[package]] name = "tokio" -version = "1.52.1" +version = "1.52.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b67dee974fe86fd92cc45b7a95fdd2f99a36a6d7b0d431a231178d3d670bbcc6" +checksum = "8fc7f01b389ac15039e4dc9531aa973a135d7a4135281b12d7c1bc79fd57fffe" dependencies = [ "bytes", "libc", @@ -6532,9 +6992,9 @@ dependencies = [ [[package]] name = "tonic" -version = "0.14.5" +version = "0.14.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fec7c61a0695dc1887c1b53952990f3ad2e3a31453e1f49f10e75424943a93ec" +checksum = "ac2a5518c70fa84342385732db33fb3f44bc4cc748936eb5833d2df34d6445ef" dependencies = [ "async-trait", "base64", @@ -6562,9 +7022,9 @@ dependencies = [ [[package]] name = "tonic-build" -version = "0.14.5" +version = "0.14.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1882ac3bf5ef12877d7ed57aad87e75154c11931c2ba7e6cde5e22d63522c734" +checksum = "c68f61875ac5293cf72e6c8cf0158086428c82c37229e98c840878f1706b0322" dependencies = [ "prettyplease", "proc-macro2", @@ -6574,9 +7034,9 @@ dependencies = [ [[package]] name = "tonic-prost" -version = "0.14.5" +version = "0.14.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a55376a0bbaa4975a3f10d009ad763d8f4108f067c7c2e74f3001fb49778d309" +checksum = "50849f68853be452acf590cde0b146665b8d507b3b8af17261df47e02c209ea0" dependencies = [ "bytes", "prost", @@ -6585,9 +7045,9 @@ dependencies = [ [[package]] name = "tonic-prost-build" -version = "0.14.5" +version = "0.14.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f3144df636917574672e93d0f56d7edec49f90305749c668df5101751bb8f95a" +checksum = "654e5643eff75d7f8c99197ce1440ed19a3474eada74c12bbac488b2cafdae27" dependencies = [ "prettyplease", "proc-macro2", @@ -6620,9 +7080,9 @@ dependencies = [ [[package]] name = "tower-http" -version = "0.6.8" +version = "0.6.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d4e6559d53cc268e5031cd8429d05415bc4cb4aefc4aa5d6cc35fbf5b924a1f8" +checksum = "68d6fdd9f81c2819c9a8b0e0cd91660e7746a8e6ea2ba7c6b2b057985f6bcb51" dependencies = [ "base64", "bitflags", @@ -6630,13 +7090,13 @@ dependencies = [ "futures-util", "http 1.4.0", "http-body", - "iri-string", "mime", "pin-project-lite", "tower", "tower-layer", "tower-service", "tracing", + "url", ] [[package]] @@ -6844,6 +7304,17 @@ version = "1.13.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9629274872b2bfaf8d66f5f15725007f635594914870f65218920345aa11aa8c" +[[package]] +name = "unicode-truncate" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "16b380a1238663e5f8a691f9039c73e1cdae598a30e9855f541d29b08b53e9a5" +dependencies = [ + "itertools 0.14.0", + "unicode-segmentation", + "unicode-width", +] + [[package]] name = "unicode-width" version = "0.2.2" @@ -6862,6 +7333,12 @@ version = "0.2.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "673aac59facbab8a9007c7f6108d11f63b603f7cabff99fabf650fea5c32b861" +[[package]] +name = "untrusted" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a156c684c91ea7d62626509bce3cb4e1d9ed5c4d978f7b4352658f96a4c26b4a" + [[package]] name = "untrusted" version = "0.9.0" @@ -6930,9 +7407,9 @@ checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" [[package]] name = "utoipa" -version = "5.4.0" +version = "5.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2fcc29c80c21c31608227e0912b2d7fddba57ad76b606890627ba8ee7964e993" +checksum = "8bde15df68e80b16c7d16b9616e80770ad158988daa56a27dccd1e55558b0160" dependencies = [ "indexmap 2.14.0", "serde", @@ -6942,9 +7419,9 @@ dependencies = [ [[package]] name = "utoipa-gen" -version = "5.4.0" +version = "5.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6d79d08d92ab8af4c5e8a6da20c47ae3f61a0f1dabc1997cdf2d082b757ca08b" +checksum = "6ba0b99ee52df3028635d93840c797102da61f8a7bb3cf751032455895b52ef8" dependencies = [ "proc-macro2", "quote", @@ -7057,9 +7534,9 @@ checksum = "b8dad83b4f25e74f184f64c43b150b91efe7647395b42289f38e50566d82855b" [[package]] name = "wasm-bindgen" -version = "0.2.118" +version = "0.2.121" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0bf938a0bacb0469e83c1e148908bd7d5a6010354cf4fb73279b7447422e3a89" +checksum = "49ace1d07c165b0864824eee619580c4689389afa9dc9ed3a4c75040d82e6790" dependencies = [ "cfg-if", "once_cell", @@ -7071,9 +7548,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-futures" -version = "0.4.68" +version = "0.4.71" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f371d383f2fb139252e0bfac3b81b265689bf45b6874af544ffa4c975ac1ebf8" +checksum = "96492d0d3ffba25305a7dc88720d250b1401d7edca02cc3bcd50633b424673b8" dependencies = [ "js-sys", "wasm-bindgen", @@ -7081,9 +7558,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro" -version = "0.2.118" +version = "0.2.121" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "eeff24f84126c0ec2db7a449f0c2ec963c6a49efe0698c4242929da037ca28ed" +checksum = "8e68e6f4afd367a562002c05637acb8578ff2dea1943df76afb9e83d177c8578" dependencies = [ "quote", "wasm-bindgen-macro-support", @@ -7091,9 +7568,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro-support" -version = "0.2.118" +version = "0.2.121" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9d08065faf983b2b80a79fd87d8254c409281cf7de75fc4b773019824196c904" +checksum = "d95a9ec35c64b2a7cb35d3fead40c4238d0940c86d107136999567a4703259f2" dependencies = [ "bumpalo", "proc-macro2", @@ -7104,9 +7581,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-shared" -version = "0.2.118" +version = "0.2.121" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5fd04d9e306f1907bd13c6361b5c6bfc7b3b3c095ed3f8a9246390f8dbdee129" +checksum = "c4e0100b01e9f0d03189a92b96772a1fb998639d981193d7dbab487302513441" dependencies = [ "unicode-ident", ] @@ -7160,9 +7637,9 @@ dependencies = [ [[package]] name = "web-sys" -version = "0.3.95" +version = "0.3.98" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4f2dfbb17949fa2088e5d39408c48368947b86f7834484e87b73de55bc14d97d" +checksum = "4b572dff8bcf38bad0fa19729c89bb5748b2b9b1d8be70cf90df697e3a8f32aa" dependencies = [ "js-sys", "wasm-bindgen", @@ -7535,9 +8012,9 @@ checksum = "d6bbff5f0aada427a1e5a6da5f1f98158182f26556f345ac9e04d36d0ebed650" [[package]] name = "winnow" -version = "1.0.1" +version = "1.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "09dac053f1cd375980747450bfc7250c264eaae0583872e845c0c7cd578872b5" +checksum = "0592e1c9d151f854e6fd382574c3a0855250e1d9b2f99d9281c6e6391af352f1" dependencies = [ "memchr", ] @@ -7657,9 +8134,9 @@ version = "0.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1301e935010a701ae5f8655edc0ad17c44bad3ac5ce8c39185f75453b720ae94" dependencies = [ - "const-oid", + "const-oid 0.9.6", "der 0.7.10", - "spki", + "spki 0.7.3", ] [[package]] @@ -7768,9 +8245,9 @@ dependencies = [ [[package]] name = "zerofrom" -version = "0.1.7" +version = "0.1.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "69faa1f2a1ea75661980b013019ed6687ed0e83d069bc1114e2cc74c6c04c4df" +checksum = "0ec05a11813ea801ff6d75110ad09cd0824ddba17dfe17128ea0d5f68e6c5272" dependencies = [ "zerofrom-derive", ] @@ -7792,6 +8269,20 @@ name = "zeroize" version = "1.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b97154e67e32c85465826e8bcc1c59429aaaf107c1e4a9e53c8d8ccd5eff88d0" +dependencies = [ + "zeroize_derive", +] + +[[package]] +name = "zeroize_derive" +version = "1.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85a5b4158499876c763cb03bc4e49185d3cccbabb15b33c627f7884f43db852e" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.117", +] [[package]] name = "zerotrie" diff --git a/Cargo.toml b/Cargo.toml index 1036fcec8..de2efcd40 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -90,10 +90,16 @@ etl-telemetry = { path = "crates/etl-telemetry", default-features = false } fail = { version = "0.5.1", default-features = false } futures = { version = "0.3.31", default-features = false } gcp-bigquery-client = { git = "https://github.com/iambriccardo/gcp-bigquery-client", rev = "c4fc59e338ca181d29b0dd53cac786fbe8513633", default-features = false } +hex = { version = "0.4", default-features = false } humantime = { version = "2.3.0", default-features = false } iceberg = { version = "0.8.0", default-features = false } iceberg-catalog-rest = { version = "0.8.0", default-features = false } insta = { version = "1.43.1", default-features = false } +# aws_lc_rs backend avoids the `rsa` crate (RUSTSEC-2023-0071, unpatched timing +# side-channel). We only sign JWTs, but cargo-deny flags the advisory regardless. +jsonwebtoken = { version = "10.3", default-features = false, features = [ + "aws_lc_rs", +] } k8s-openapi = { version = "0.25.0", default-features = false } kube = { version = "1.1.0", default-features = false } metrics = { version = "0.24.2", default-features = false } @@ -102,6 +108,7 @@ parking_lot = { version = "0.12.5" } parquet = { version = "57.0", default-features = false } pg_escape = { version = "0.1.1", default-features = false } pin-project-lite = { version = "0.2.16", default-features = false } +pkcs8 = { version = "0.11", default-features = false, features = ["pem", "encryption"] } postgres-replication = { git = "https://github.com/iambriccardo/rust-postgres", default-features = false, rev = "31acf55c7e5c2244e5bb3a36e7afa2a01bf52c38" } prost = { version = "0.14.1", default-features = false } r2d2 = { version = "0.8", default-features = false } @@ -113,6 +120,7 @@ secrecy = { version = "0.10.3", default-features = false } sentry = { version = "0.42.0" } serde = { version = "1.0.219", default-features = false } serde_json = { version = "1.0.141", default-features = false } +sha2 = { version = "0.11", default-features = false } # Use the 0.9 alpha for the rustls VerifyCa fix in https://github.com/launchbadge/sqlx/pull/3861. sqlx = { version = "0.9.0-alpha.1", default-features = false } sysinfo = { version = "0.38.2", default-features = false } @@ -135,6 +143,7 @@ utoipa = { version = "5.4.0", default-features = false } utoipa-swagger-ui = { version = "9.0.2", default-features = false, features = ["vendored"] } uuid = { version = "1.17.0", default-features = false } x509-cert = { version = "0.2.2", default-features = false } +zstd = { version = "0.13", default-features = false } [profile.bench] debug = true diff --git a/crates/etl-config/src/shared/destination.rs b/crates/etl-config/src/shared/destination.rs index 0ae24f958..b75e826ca 100644 --- a/crates/etl-config/src/shared/destination.rs +++ b/crates/etl-config/src/shared/destination.rs @@ -152,6 +152,22 @@ pub enum DestinationConfig { #[serde(default)] maintenance_mode: DuckLakeMaintenanceMode, }, + Snowflake { + /// Snowflake account identifier in "ORGNAME-ACCOUNTNAME" format. + account_id: String, + /// Snowflake user with RSA public key configured. + user: String, + /// Path to RSA private key file (PEM/PKCS8 format). + private_key_path: String, + /// Optional passphrase for encrypted private key. + private_key_passphrase: Option, + /// Target database name. + database: String, + /// Target schema name. + schema: String, + /// Snowflake role. + role: Option, + }, } impl DestinationConfig { @@ -350,6 +366,21 @@ pub enum DestinationConfigWithoutSecrets { #[serde(default)] maintenance_mode: DuckLakeMaintenanceMode, }, + Snowflake { + /// Snowflake account identifier in "ORGNAME-ACCOUNTNAME" format. + account_id: String, + /// Snowflake user with RSA public key configured. + user: String, + /// Path to RSA private key file (PEM/PKCS8 format). + private_key_path: String, + /// Target database name. + database: String, + /// Target schema name. + schema: String, + /// Snowflake role. + #[serde(skip_serializing_if = "Option::is_none")] + role: Option, + }, } impl From for DestinationConfigWithoutSecrets { @@ -402,6 +433,22 @@ impl From for DestinationConfigWithoutSecrets { expire_snapshots_older_than, maintenance_mode, }, + DestinationConfig::Snowflake { + account_id, + user, + private_key_path, + private_key_passphrase: _, + database, + schema, + role, + } => DestinationConfigWithoutSecrets::Snowflake { + account_id, + user, + private_key_path, + database, + schema, + role, + }, } } } diff --git a/crates/etl-destinations/Cargo.toml b/crates/etl-destinations/Cargo.toml index 85922869a..963f99e8e 100644 --- a/crates/etl-destinations/Cargo.toml +++ b/crates/etl-destinations/Cargo.toml @@ -66,6 +66,24 @@ clickhouse = [ "dep:url", "dep:uuid", ] +snowflake = [ + "dep:base64", + "dep:bytes", + "dep:hex", + "dep:jsonwebtoken", + "dep:metrics", + "dep:reqwest", + "dep:aws-lc-rs", + "dep:pkcs8", + "dep:secrecy", + "dep:serde", + "dep:serde_json", + "dep:sha2", + "dep:thiserror", + "dep:tokio", + "dep:tracing", + "dep:zstd", +] egress = ["etl/egress"] # We assume that `test-utils` is always used in conjunction with `bigquery` or `iceberg` thus we only # put here the extra dependencies needed. @@ -75,6 +93,9 @@ test-utils = ["dep:uuid"] arrow = { workspace = true, optional = true } async-trait = { workspace = true, optional = true } +aws-lc-rs = { workspace = true, optional = true } +base64 = { workspace = true, optional = true } +bytes = { workspace = true, optional = true } chrono = { workspace = true, features = ["serde"] } clickhouse = { workspace = true, optional = true, features = ["inserter", "lz4", "rustls-tls"] } duckdb = { workspace = true, optional = true, features = ["bundled", "json", "parquet", "r2d2"] } @@ -83,27 +104,35 @@ etl-config = { workspace = true } etl-maintenance = { workspace = true, optional = true } futures = { workspace = true, optional = true } gcp-bigquery-client = { workspace = true, optional = true, features = ["rust-tls", "aws-lc-rs"] } +hex = { workspace = true, optional = true } humantime = { workspace = true, optional = true } iceberg = { workspace = true, optional = true } iceberg-catalog-rest = { workspace = true, optional = true } +jsonwebtoken = { workspace = true, optional = true } +kube = { workspace = true, optional = true, features = ["client", "rustls-tls"] } metrics = { workspace = true, optional = true } parking_lot = { workspace = true, optional = true } parquet = { workspace = true, optional = true, features = ["async", "arrow"] } pg_escape = { workspace = true, optional = true } +pkcs8 = { workspace = true, optional = true } prost = { workspace = true, optional = true } r2d2 = { workspace = true, optional = true } rand = { workspace = true, optional = true, features = ["thread_rng"] } regex = { workspace = true, optional = true } -reqwest = { workspace = true, optional = true, features = ["json"] } +reqwest = { workspace = true, optional = true, features = ["json", "rustls-tls"] } +secrecy = { workspace = true, optional = true } serde = { workspace = true, optional = true, features = ["derive"] } serde_json = { workspace = true, optional = true, features = ["arbitrary_precision", "std"] } +sha2 = { workspace = true, optional = true } sqlx = { workspace = true, optional = true, features = ["runtime-tokio", "tls-rustls", "postgres"] } +thiserror = { workspace = true, optional = true } tokio = { workspace = true, optional = true, features = ["rt", "sync", "time"] } tokio-postgres = { workspace = true, optional = true } tonic = { workspace = true, optional = true } tracing = { workspace = true, optional = true, default-features = true } url = { workspace = true, optional = true } uuid = { workspace = true, optional = true, features = ["v4"] } +zstd = { workspace = true, optional = true } [dev-dependencies] base64 = { workspace = true } diff --git a/crates/etl-destinations/src/lib.rs b/crates/etl-destinations/src/lib.rs index f6dc47d7c..43eae42c6 100644 --- a/crates/etl-destinations/src/lib.rs +++ b/crates/etl-destinations/src/lib.rs @@ -4,13 +4,14 @@ //! warehouses and analytics platforms, enabling data replication from Postgres //! to cloud services. -#[cfg(any(feature = "bigquery", feature = "ducklake", feature = "iceberg"))] +#[cfg(any(feature = "bigquery", feature = "ducklake", feature = "iceberg", feature = "snowflake"))] mod retry; #[cfg(any( feature = "bigquery", feature = "clickhouse", feature = "ducklake", - feature = "iceberg" + feature = "iceberg", + feature = "snowflake" ))] mod table_name; @@ -24,3 +25,5 @@ pub mod ducklake; pub mod egress; #[cfg(feature = "iceberg")] pub mod iceberg; +#[cfg(feature = "snowflake")] +pub mod snowflake; diff --git a/crates/etl-destinations/src/snowflake/README.md b/crates/etl-destinations/src/snowflake/README.md new file mode 100644 index 000000000..21a012886 --- /dev/null +++ b/crates/etl-destinations/src/snowflake/README.md @@ -0,0 +1,52 @@ +# Snowflake Destination + +## Running Integration Tests + +Snowflake integration tests are marked `#[ignore]` because they require a real Snowflake account. + +To run them: + +1. Copy the example env file and fill in your Snowflake credentials (see variables below): + +```bash +cp .env.example .env +# edit .env with your values +``` + +2. Source the file and run tests: + +```bash +source .env +cargo test -p etl-destinations --features snowflake,test-utils -- --ignored +``` + +To run a specific test: + +```bash +cargo test -p etl-destinations --features snowflake,test-utils -- --ignored authenticate_against_snowflake +``` + +### Environment Variables + +| Variable | Required | Default | Description | +| ---------------------------------- | -------- | --------- | ------------------------------------------ | +| `TESTS_SNOWFLAKE_ACCOUNT` | yes | | Account identifier, e.g. `myorg-myaccount` | +| `TESTS_SNOWFLAKE_USER` | yes | | Login user name | +| `TESTS_SNOWFLAKE_PRIVATE_KEY_PATH` | yes | | Path to PEM-encoded private key | +| `TESTS_SNOWFLAKE_DATABASE` | no | `ETL_DEV` | Target database | +| `TESTS_SNOWFLAKE_SCHEMA` | no | `PUBLIC` | Target schema | +| `TESTS_SNOWFLAKE_ROLE` | no | | Role to assume after connecting | + +### Key-Pair Authentication Setup + +Snowflake tests use key-pair authentication (not password). To generate a key: + +```bash +openssl genrsa 2048 | openssl pkcs8 -topk8 -nocrypt -out rsa_key.p8 +``` + +Then register the public key with your Snowflake user: + +```sql +ALTER USER ETL_USER SET RSA_PUBLIC_KEY=''; +``` diff --git a/crates/etl-destinations/src/snowflake/auth.rs b/crates/etl-destinations/src/snowflake/auth.rs new file mode 100644 index 000000000..eb40678fe --- /dev/null +++ b/crates/etl-destinations/src/snowflake/auth.rs @@ -0,0 +1,444 @@ +//! Key-pair JWT authentication for the Snowflake SQL and Streaming APIs. +//! +//! Flow: +//! 1. Load RSA private key from disk, derive a public-key fingerprint +//! 2. Sign a short-lived JWT (1 h) identifying the account and user +//! 3. Exchange the JWT at `POST {account_url}/oauth/token` for a bearer token +//! (~10 min TTL) +//! 4. Cache the bearer token; refresh proactively when < 60 s remain +//! +//! The `scope` parameter is intentionally omitted from the token exchange so +//! the resulting token is accepted by both the SQL REST API and the Snowpipe +//! Streaming REST API. Setting `scope=` would restrict the token +//! to a single ingest host. +//! +//! Ref: +//! Ref: + +use std::{future::Future, time::Duration}; + +use aws_lc_rs::{encoding::AsDer, signature::KeyPair as _}; +use base64::{Engine as _, engine::general_purpose as base64_engine}; +use jsonwebtoken::{Algorithm, EncodingKey, Header}; +use reqwest::StatusCode; +use secrecy::ExposeSecret as _; +use sha2::Digest as _; +use tokio::{sync::Mutex, time::Instant}; + +use crate::snowflake::{ + Config, Error, Result, + config::{HTTP_CONNECT_TIMEOUT, HTTP_REQUEST_TIMEOUT}, +}; + +/// Self-signed JWT validity window (Snowflake rejects JWTs older than this). +const TOKEN_LIFETIME_SECS: u64 = 3600; + +/// Refresh the scoped token this far before it expires to avoid mid-request +/// expiry. +const TOKEN_REFRESH_BUFFER: Duration = Duration::from_secs(60); + +/// Token produced by Snowflake's `/oauth/token` endpoint. +/// +/// "Scoped" in Snowflake's terminology refers to the token being a short-lived, +/// restricted derivative of the raw JWT, not the `scope` request parameter +/// (which controls host restriction and is optional; see module docs). +pub struct ScopedToken { + pub(crate) access_token: String, + pub(crate) expires_at: Instant, +} + +/// Abstracts the HTTP call that exchanges a self-signed JWT for a scoped token. +pub trait TokenExchanger: Send + Sync { + fn exchange( + &self, + account_url: &str, + jwt: &str, + ) -> impl Future> + Send; +} + +/// Provides a valid bearer token and supports invalidation. +pub trait TokenProvider: Send + Sync { + /// Return a valid bearer token, refreshing it if necessary. + fn get_token(&self) -> impl Future> + Send; + + /// Invalidate any cached token, forcing a fresh exchange on the next + /// [`get_token`](Self::get_token) call. + fn invalidate_token(&self) -> impl Future + Send; +} + +/// Production implementation that calls Snowflake's OAuth endpoint over HTTP. +pub struct HttpExchanger { + http: reqwest::Client, +} + +impl HttpExchanger { + pub fn new(http: reqwest::Client) -> Self { + Self { http } + } +} + +impl TokenExchanger for HttpExchanger { + async fn exchange(&self, account_url: &str, jwt: &str) -> Result { + let url = format!("{account_url}/oauth/token"); + let body = + format!("grant_type=urn:ietf:params:oauth:grant-type:jwt-bearer&assertion={jwt}"); + + let response = self + .http + .post(&url) + .header(reqwest::header::CONTENT_TYPE, "application/x-www-form-urlencoded") + .body(body) + .send() + .await?; + + let status = response.status(); + let body_text = response.text().await?; + + // Likely wrong key, non-retriable. + if status == StatusCode::UNAUTHORIZED { + return Err(Error::Auth(format!("token exchange rejected (401): {body_text}"))); + } + + // Transient failure, probably retriable error. + if status.is_client_error() || status.is_server_error() { + return Err(Error::HttpStatus { status, body: body_text }); + } + + // Snowflake returns the scoped token as a raw JWT (RFC 7519). + let expires_at = decode_token_expiry(&body_text)?; + + Ok(ScopedToken { access_token: body_text, expires_at }) + } +} + +/// Manages key-pair authentication and scoped-token lifecycle. +/// +/// Signs a JWT locally, exchanges it for a short-lived scoped token via +/// Snowflake's OAuth endpoint, and caches the result. +pub struct AuthManager { + account_url: String, + account: String, + user: String, + encoding_key: EncodingKey, + key_fingerprint: String, + cached_token: Mutex>, + exchanger: E, +} + +impl AuthManager { + /// Build an `AuthManager` from an RSA private key on disk. + /// + /// Creates its own internal `reqwest::Client` via `HttpExchanger`. + pub fn new( + config: &Config, + private_key_path: &str, + passphrase: Option<&secrecy::SecretString>, + ) -> Result { + Self::with_exchanger( + config, + private_key_path, + passphrase, + HttpExchanger::new( + reqwest::Client::builder() + .connect_timeout(HTTP_CONNECT_TIMEOUT) + .timeout(HTTP_REQUEST_TIMEOUT) + .build() + .expect("failed to build HTTP client"), + ), + ) + } +} + +impl AuthManager { + /// Build an `AuthManager` with a custom token exchanger. + /// + /// The public-key fingerprint is derived and reused for every JWT. + /// Accepts PKCS#8 (encrypted or plain) and PKCS#1 PEM formats. + pub fn with_exchanger( + config: &Config, + private_key_path: &str, + passphrase: Option<&secrecy::SecretString>, + exchanger: E, + ) -> Result { + let pem_text = std::fs::read_to_string(private_key_path) + .map_err(|e| Error::Config(format!("failed to read private key file: {e}")))?; + + let (pkcs1_der, key_pair) = decode_and_load_rsa_key(&pem_text, passphrase)?; + + // Snowflake identifies keys by SHA-256(DER-encoded public key). + let pub_der = key_pair + .public_key() + .as_der() + .map_err(|_| Error::Auth("failed to encode public key to DER".into()))?; + let hash = sha2::Sha256::digest(pub_der.as_ref()); + let b64 = base64_engine::STANDARD.encode(hash); + let key_fingerprint = format!("SHA256:{b64}"); + + // jsonwebtoken's aws_lc_rs backend passes these bytes to + // RsaKeyPair::from_der(), which expects PKCS#1 (raw RSA) DER. + let encoding_key = EncodingKey::from_rsa_der(&pkcs1_der); + + Ok(Self { + account_url: config.account_url().to_owned(), + account: config.account_id.to_uppercase(), + user: config.username.to_uppercase(), + encoding_key, + key_fingerprint, + cached_token: tokio::sync::Mutex::new(None), + exchanger, + }) + } + + /// Create a short-lived JWT (1 hour) used to request a scoped OAuth token. + fn generate_jwt(&self) -> Result { + #[derive(serde::Serialize)] + struct JwtClaims { + iss: String, + sub: String, + iat: u64, + exp: u64, + } + + let iat = jsonwebtoken::get_current_timestamp(); + let claims = JwtClaims { + iss: format!("{}.{}.{}", self.account, self.user, self.key_fingerprint), + sub: format!("{}.{}", self.account, self.user), + iat, + exp: iat + TOKEN_LIFETIME_SECS, + }; + + jsonwebtoken::encode(&Header::new(Algorithm::RS256), &claims, &self.encoding_key) + .map_err(|e| Error::Auth(format!("failed to sign JWT: {e}"))) + } +} + +impl TokenProvider for AuthManager { + /// Return a valid scoped token, refreshing it if necessary. + async fn get_token(&self) -> Result { + let mut cached = self.cached_token.lock().await; + + if let Some(ref token) = *cached + && token.expires_at > Instant::now() + TOKEN_REFRESH_BUFFER + { + return Ok(token.access_token.clone()); + } + + tracing::debug!("refreshing Snowflake scoped token"); + + let jwt = self.generate_jwt()?; + let scoped = self.exchanger.exchange(&self.account_url, &jwt).await?; + let access_token = scoped.access_token.clone(); + *cached = Some(scoped); + + Ok(access_token) + } + + /// Invalidate the cached scoped token. + /// + /// The next call to [`TokenProvider::get_token`] will perform a fresh + /// token exchange. + async fn invalidate_token(&self) { + *self.cached_token.lock().await = None; + } +} + +/// Decode a PEM-encoded RSA private key, returning the PKCS#1 DER bytes +/// (for jsonwebtoken) and a loaded `KeyPair` (for fingerprint derivation). +/// +/// Supports encrypted PKCS#8, plain PKCS#8, and PKCS#1 PEM formats. +fn decode_and_load_rsa_key( + pem_text: &str, + passphrase: Option<&secrecy::SecretString>, +) -> Result<(Vec, aws_lc_rs::rsa::KeyPair)> { + use pkcs8::der::{Decode as _, SecretDocument}; + + let der_bytes = if let Some(pass) = passphrase { + let (_, doc) = SecretDocument::from_pem(pem_text) + .map_err(|e| Error::Auth(format!("failed to parse encrypted PEM: {e}")))?; + let enc = pkcs8::EncryptedPrivateKeyInfoRef::try_from(doc.as_bytes()) + .map_err(|e| Error::Auth(format!("failed to parse encrypted key: {e}")))?; + enc.decrypt(pass.expose_secret()) + .map_err(|e| Error::Auth(format!("failed to decrypt private key: {e}")))? + .as_bytes() + .to_vec() + } else { + let (_, doc) = SecretDocument::from_pem(pem_text) + .map_err(|e| Error::Auth(format!("failed to parse private key PEM: {e}")))?; + doc.as_bytes().to_vec() + }; + + // Try PKCS#8: extract the inner RSA key for jsonwebtoken, load via from_pkcs8. + if let Ok(pki) = pkcs8::PrivateKeyInfoRef::from_der(&der_bytes) { + let key_pair = aws_lc_rs::rsa::KeyPair::from_pkcs8(&der_bytes) + .map_err(|e| Error::Auth(format!("failed to load RSA key: {e}")))?; + return Ok((pki.private_key.as_bytes().to_vec(), key_pair)); + } + + // Raw PKCS#1 DER: use directly. + let key_pair = aws_lc_rs::rsa::KeyPair::from_der(&der_bytes) + .map_err(|e| Error::Auth(format!("failed to load RSA key: {e}")))?; + Ok((der_bytes, key_pair)) +} + +/// Extract `exp` from a raw JWT (RFC 7519) and convert to monotonic Instant. +/// +/// Returned value is adjusted accordingly to how much time remains for original +/// expiry to happen. +fn decode_token_expiry(token: &str) -> Result { + let payload = token + .split('.') + .nth(1) + .ok_or_else(|| Error::Auth("scoped token is not a valid JWT (missing payload)".into()))?; + + let bytes = base64_engine::URL_SAFE_NO_PAD + .decode(payload) + .map_err(|e| Error::Auth(format!("scoped token payload is not valid base64: {e}")))?; + + let claims: serde_json::Value = serde_json::from_slice(&bytes) + .map_err(|e| Error::Auth(format!("scoped token payload is not valid JSON: {e}")))?; + + let exp = claims["exp"] + .as_u64() + .ok_or_else(|| Error::Auth("scoped token is missing the `exp` claim".into()))?; + + let remaining = exp.saturating_sub(jsonwebtoken::get_current_timestamp()); + + Ok(Instant::now() + Duration::from_secs(remaining)) +} + +#[cfg(test)] +mod tests { + use super::*; + + struct TestExchanger; + + impl TokenExchanger for TestExchanger { + async fn exchange(&self, _account_url: &str, _jwt: &str) -> Result { + Ok(ScopedToken { + access_token: "fresh-token-from-exchange".to_owned(), + expires_at: Instant::now() + Duration::from_secs(3600), + }) + } + } + + impl AuthManager { + fn inject_token_for_test(&self, access_token: String, ttl: std::time::Duration) { + *self.cached_token.try_lock().unwrap() = + Some(ScopedToken { access_token, expires_at: Instant::now() + ttl }); + } + + fn fingerprint(&self) -> &str { + &self.key_fingerprint + } + } + + const TEST_KEY_PATH: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/testdata/test_key.pem"); + + fn make_test_manager() -> AuthManager { + make_test_manager_with_account("ORG-ACCT", "USER") + } + + fn make_test_manager_with_account( + account_id: &str, + username: &str, + ) -> AuthManager { + let config = Config::new(account_id, username, "TEST_DB", "PUBLIC"); + + AuthManager::with_exchanger(&config, TEST_KEY_PATH, None, TestExchanger) + .expect("AuthManager::with_exchanger") + } + + #[test] + fn jwt_claims() { + let cases = [ + // (input_account, input_user, expected_account, expected_user). + ("TESTORG-TESTACCOUNT", "TESTUSER", "TESTORG-TESTACCOUNT", "TESTUSER"), + ("org-account", "my_user", "ORG-ACCOUNT", "MY_USER"), + ]; + + fn decode_jwt_claims(jwt: &str) -> serde_json::Value { + let parts: Vec<&str> = jwt.split('.').collect(); + assert_eq!(parts.len(), 3, "JWT must have 3 dot-separated parts"); + let payload_bytes = + base64_engine::URL_SAFE_NO_PAD.decode(parts[1]).expect("base64 decode payload"); + serde_json::from_slice(&payload_bytes).expect("json parse claims") + } + + for (account, user, expect_account, expect_user) in cases { + let manager = make_test_manager_with_account(account, user); + let jwt = manager.generate_jwt().expect("generate_jwt"); + let claims = decode_jwt_claims(&jwt); + + // Issuer must be ACCOUNT.USER.FINGERPRINT. + let iss = claims["iss"].as_str().expect("iss"); + assert!( + iss.starts_with(&format!("{expect_account}.{expect_user}.SHA256:")), + "unexpected iss: {iss}" + ); + + // Subject must be ACCOUNT.USER. + assert_eq!(claims["sub"], format!("{expect_account}.{expect_user}")); + + // Token must expire exactly TOKEN_LIFETIME_SECS after issuance. + let iat = claims["iat"].as_u64().expect("iat"); + let exp = claims["exp"].as_u64().expect("exp"); + assert_eq!(exp, iat + TOKEN_LIFETIME_SECS); + } + } + + #[test] + fn key_fingerprint_format() { + let manager = make_test_manager(); + let fp = manager.fingerprint(); + + // Fingerprint must use the SHA256: prefix per Snowflake convention. + assert!(fp.starts_with("SHA256:"), "fingerprint missing SHA256 prefix"); + + // The base64 payload must decode to exactly 32 bytes (SHA-256 digest). + let b64_part = &fp["SHA256:".len()..]; + let decoded = base64_engine::STANDARD.decode(b64_part).expect("base64 decode fingerprint"); + assert_eq!(decoded.len(), 32, "SHA-256 digest must be 32 bytes"); + } + + #[test] + fn config_derives_account_url() { + let config = Config::new("ORG-ACCT", "USER", "TEST_DB", "PUBLIC"); + assert_eq!(config.account_url(), "https://ORG-ACCT.snowflakecomputing.com"); + } + + #[tokio::test] + async fn token_cache_hit() { + let manager = make_test_manager(); + manager.inject_token_for_test("cached-token-123".to_owned(), Duration::from_secs(120)); + + // A non-expired token should be returned directly from cache. + let token = manager.get_token().await.expect("get_token"); + assert_eq!(token, "cached-token-123"); + } + + #[tokio::test] + async fn token_cache_expired_triggers_refresh() { + let manager = make_test_manager(); + + // Inject a token with 30s TTL, which is below the 60s refresh buffer. + manager.inject_token_for_test("stale-token".to_owned(), Duration::from_secs(30)); + + // get_token should detect the stale cache and call the exchanger. + let token = manager.get_token().await.expect("get_token"); + + // The fresh token from FakeExchanger must be returned + assert_ne!(token, "stale-token", "still getting stale token"); + assert_eq!(token, "fresh-token-from-exchange", "unexpected refreshed token"); + + // The cache must now hold the fresh token. + { + let cached = manager.cached_token.try_lock().unwrap(); + let entry = cached.as_ref().expect("cache should not be empty"); + assert_eq!(entry.access_token, "fresh-token-from-exchange"); + } + + // Second call should return the cached fresh token without re-exchanging. + let token = manager.get_token().await.expect("get_token"); + assert_eq!(token, "fresh-token-from-exchange"); + } +} diff --git a/crates/etl-destinations/src/snowflake/client.rs b/crates/etl-destinations/src/snowflake/client.rs new file mode 100644 index 000000000..eee25d140 --- /dev/null +++ b/crates/etl-destinations/src/snowflake/client.rs @@ -0,0 +1,194 @@ +use std::{collections::HashMap, sync::Arc}; + +use etl::types::{ColumnSchema, PipelineId, SchemaDiff, TableId}; +use tokio::sync::{Mutex, RwLock}; + +use crate::snowflake::{ + Config, Error, Result, + auth::{AuthManager, HttpExchanger, TokenProvider}, + config::{HTTP_CONNECT_TIMEOUT, HTTP_REQUEST_TIMEOUT}, + schema, + sql_client::SqlClient, + streaming::{ChannelHandle, OffsetToken, RestStreamClient, RowBatch, StreamClient}, +}; + +type ChannelMap = Arc>>>>>; + +/// Snowflake API client. +/// +/// Unifies the SQL REST API (DDL) and the Snowpipe Streaming API (channel +/// lifecycle and row ingestion). +pub struct Client> { + sql_client: Arc>, + stream_client: Arc, + database: String, + schema: String, + pipeline_id: PipelineId, + channels: ChannelMap, +} + +impl Clone for Client { + fn clone(&self) -> Self { + Self { + sql_client: Arc::clone(&self.sql_client), + stream_client: Arc::clone(&self.stream_client), + database: self.database.clone(), + schema: self.schema.clone(), + pipeline_id: self.pipeline_id, + channels: Arc::clone(&self.channels), + } + } +} + +/// Convenience constructor for the default client stack. +impl Client> { + pub fn new( + config: Config, + auth: Arc>, + pipeline_id: PipelineId, + ) -> Self { + let http = reqwest::Client::builder() + .connect_timeout(HTTP_CONNECT_TIMEOUT) + .timeout(HTTP_REQUEST_TIMEOUT) + .build() + .expect("failed to build HTTP client"); + let database = config.database.clone(); + let schema = config.schema.clone(); + let stream_client = Arc::new(RestStreamClient::new( + config.account_url().to_owned(), + Arc::clone(&auth), + http.clone(), + )); + let sql_client = SqlClient::new(config, auth, http); + Self::with_clients(sql_client, stream_client, database, schema, pipeline_id) + } +} + +impl Client { + /// Build a client from pre-constructed SQL and streaming clients. + pub fn with_clients( + sql_client: SqlClient, + stream_client: Arc, + database: String, + schema: String, + pipeline_id: PipelineId, + ) -> Self { + Self { + sql_client: Arc::new(sql_client), + stream_client, + database, + schema, + pipeline_id, + channels: Arc::new(RwLock::new(HashMap::new())), + } + } + + /// Ensure the table exists in Snowflake and is ready to receive data. + /// + /// Returns `true` when streaming was newly set up for this table in the + /// current process (the Snowflake table itself may have already existed). + #[allow(clippy::map_entry)] + pub async fn ensure_table( + &self, + table_id: TableId, + table_name: &str, + columns: &[ColumnSchema], + ) -> Result { + // Fast path: read lock, check if already set up. + let channels = self.channels.read().await; + if channels.contains_key(&table_id) { + return Ok(false); + } + drop(channels); + + // Slow path: hold write lock for the entire setup. This runs once + // per table per process lifetime, so blocking other tables briefly + // during startup is acceptable. + let mut channels = self.channels.write().await; + if channels.contains_key(&table_id) { + return Ok(false); + } + + // Create Snowflake table. + schema::validate_no_cdc_collisions(columns)?; + let column_defs = schema::build_column_defs(columns); + self.sql_client.create_table_if_not_exists(table_name, &column_defs).await?; + + // Obtain table channel. + let mut handle = ChannelHandle::new( + Arc::clone(&self.stream_client), + self.pipeline_id, + self.database.clone(), + self.schema.clone(), + table_name.to_owned(), + ); + handle.open().await?; + + // Persist table-channel mapping. + channels.insert(table_id, Arc::new(Mutex::new(handle))); + Ok(true) + } + + /// Apply column additions, renames, and removals from a schema diff. + pub async fn apply_schema_diff(&self, table_name: &str, diff: &SchemaDiff) -> Result<()> { + if diff.is_empty() { + return Ok(()); + } + + for col in &diff.columns_to_add { + self.sql_client.add_column(table_name, &col.name, schema::type_name(&col.typ)).await?; + } + + for rename in &diff.columns_to_rename { + self.sql_client.rename_column(table_name, &rename.old_name, &rename.new_name).await?; + } + + for col in &diff.columns_to_remove { + self.sql_client.drop_column(table_name, &col.name).await?; + } + + Ok(()) + } + + /// Truncate the table and reset ingestion state so offsets restart. + pub async fn truncate_table(&self, table_id: TableId, table_name: &str) -> Result<()> { + let mut guard = self.get_channel(table_id).await?.lock_owned().await; + self.sql_client.truncate_table(table_name).await?; + guard.reset().await + } + + /// Refresh the table's ingestion state after a schema change. + /// + /// Channels must be reopened after ALTER TABLE so Snowpipe picks up the + /// new column list. Without this, inserts would fail (and fall back to + /// the auto-recovery path in `process_batches`, so after one error + /// round-trip data would still be pushed, but we can avoid that extra + /// trip). + /// + /// Ref: https://docs.snowflake.com/en/user-guide/snowpipe-streaming/snowpipe-streaming-classic-recommendation + pub async fn refresh_table(&self, table_id: &TableId) -> Result<()> { + self.get_channel(*table_id).await?.lock().await.open().await + } + + /// Send pre-encoded row batches through the table's channel. + pub async fn insert_batches(&self, table_id: TableId, batches: Vec) -> Result<()> { + self.get_channel(table_id).await?.lock().await.process_batches(batches).await + } + + /// Last offset committed by Snowflake for this table's channel. + pub async fn committed_offset(&self, table_id: TableId) -> Result> { + self.get_channel(table_id).await?.lock().await.committed_offset().await + } + + /// Get table-level guard. + /// + /// Look up a channel by `table_id`, clone the `Arc`, and release the map + /// read-lock before returning. The caller then locks the per-channel mutex. + async fn get_channel(&self, table_id: TableId) -> Result>>> { + let channels = self.channels.read().await; + channels + .get(&table_id) + .cloned() + .ok_or_else(|| Error::Channel(format!("no open channel for table {table_id}"))) + } +} diff --git a/crates/etl-destinations/src/snowflake/config.rs b/crates/etl-destinations/src/snowflake/config.rs new file mode 100644 index 000000000..df3fa9a80 --- /dev/null +++ b/crates/etl-destinations/src/snowflake/config.rs @@ -0,0 +1,70 @@ +use std::time::Duration; + +pub(crate) const HTTP_CONNECT_TIMEOUT: Duration = Duration::from_secs(15); +pub(crate) const HTTP_REQUEST_TIMEOUT: Duration = Duration::from_secs(90); + +/// Connection parameters for a Snowflake account. +#[derive(Debug, Clone)] +pub struct Config { + /// Full Snowflake account URL, always HTTPS. + account_url: String, + + /// Snowflake account identifier (e.g. `ORGNAME-ACCTNAME`). + /// + /// Used in JWT claims and API routing. + /// Uppercased internally where Snowflake requires it. + pub(crate) account_id: String, + + /// Snowflake login name used for key-pair authentication. + /// + /// This is the user identity that owns the RSA key pair. + pub(crate) username: String, + + /// Target database for all operations. + pub(crate) database: String, + + /// Target schema within the database. + pub(crate) schema: String, + + /// Snowflake role to assume after connecting. + /// + /// When `None`, the user's default role is used. + pub(crate) role: Option, +} + +impl Config { + /// Create a config with the required connection parameters. + pub fn new(account_id: &str, username: &str, database: &str, schema: &str) -> Self { + let account_url = format!("https://{}.snowflakecomputing.com", account_id.to_uppercase()); + + Self { + account_url, + account_id: account_id.to_owned(), + username: username.to_owned(), + database: database.to_owned(), + schema: schema.to_owned(), + role: None, + } + } + + /// HTTPS-only account URL used for all API requests. + pub fn account_url(&self) -> &str { + &self.account_url + } + + /// Target database name. + pub fn database(&self) -> &str { + &self.database + } + + /// Target schema name. + pub fn schema(&self) -> &str { + &self.schema + } + + /// Set the role to assume after connecting. + pub fn with_role(mut self, role: &str) -> Self { + self.role = Some(role.to_owned()); + self + } +} diff --git a/crates/etl-destinations/src/snowflake/core.rs b/crates/etl-destinations/src/snowflake/core.rs new file mode 100644 index 000000000..a51a3ab32 --- /dev/null +++ b/crates/etl-destinations/src/snowflake/core.rs @@ -0,0 +1,476 @@ +use std::collections::HashMap; + +use etl::{ + bail, + concurrency::TaskSet, + destination::async_result::{TruncateTableResult, WriteEventsResult, WriteTableRowsResult}, + error::{ErrorKind, EtlError, EtlResult}, + etl_error, + state::destination_metadata::{DestinationTableMetadata, DestinationTableSchemaStatus}, + store::{schema::SchemaStore, state::StateStore}, + types::{ + ColumnSchema, DeleteEvent, Event, InsertEvent, OldTableRow, ReplicatedTableSchema, TableId, + TableRow, UpdateEvent, UpdatedTableRow, + }, +}; +use tracing::{info, warn}; + +use crate::{ + snowflake::{ + Client, + auth::{AuthManager, HttpExchanger, TokenProvider}, + encoding::{CdcMeta, CdcOperation}, + metrics::register_metrics, + schema, + streaming::{OffsetToken, RestStreamClient, RowBatchBuilder, StreamClient}, + }, + table_name::try_stringify_table_name, +}; + +type EventIter = std::iter::Peekable>; + +/// Postgres replication to Snowflake via Snowpipe Streaming. +/// +/// Thin adapter between the ETL [`etl::destination::Destination`] trait and +/// [`Client`]. Translates replication events into client operations and manages +/// the state store bookkeeping. +pub struct Destination, C = RestStreamClient> { + client: Client, + store: S, + tasks: TaskSet, +} + +impl Clone for Destination { + fn clone(&self) -> Self { + Self { client: self.client.clone(), store: self.store.clone(), tasks: self.tasks.clone() } + } +} + +impl Destination +where + S: StateStore + SchemaStore + Clone + Send + Sync + 'static, + T: TokenProvider + 'static, + C: StreamClient, +{ + /// Create a new destination. + pub fn new(client: Client, store: S) -> Self { + register_metrics(); + Self { client, store, tasks: TaskSet::new() } + } + + /// Ensure the Snowflake table and streaming channel exist for this table. + /// Operation is idempotent. + pub async fn prepare_table_for_streaming( + &self, + table_schema: &ReplicatedTableSchema, + ) -> EtlResult<()> { + let table_id = table_schema.id(); + let table_name = try_stringify_table_name(table_schema.name())?.to_uppercase(); + let columns: Vec<_> = table_schema.column_schemas().cloned().collect(); + + let table_is_new = self + .client + .ensure_table(table_id, &table_name, &columns) + .await + .map_err(EtlError::from)?; + + if table_is_new { + let snapshot_id = table_schema.inner().snapshot_id; + let replication_mask = table_schema.replication_mask().clone(); + let metadata = DestinationTableMetadata::new_applied( + table_name.clone(), + snapshot_id, + replication_mask, + ); + self.store.store_destination_table_metadata(table_id, metadata).await?; + } + + Ok(()) + } + + /// Write rows during the initial snapshot (table copy) phase. + /// + /// All rows are stamped as inserts with a zero offset since the table + /// starts empty. + pub async fn write_table_rows( + &self, + schema: &ReplicatedTableSchema, + rows: Vec, + ) -> EtlResult<()> { + // Table must exist even for empty snapshots, CDC events may arrive later. + self.prepare_table_for_streaming(schema).await?; + + if rows.is_empty() { + return Ok(()); + } + + let table_id = schema.id(); + let columns: Vec<_> = schema.column_schemas().cloned().collect(); + + // Build row batches. Snowflake has limits on max size of input, so we slice + // into proper batches, when necessary. + let zero = OffsetToken::zero(); + let mut builder = RowBatchBuilder::new(); + for row in &rows { + builder + .push_row(&columns, row, CdcMeta::new(CdcOperation::Insert, zero.as_ref()), &zero) + .map_err(EtlError::from)?; + } + + let batches = builder.finish().map_err(EtlError::from)?; + self.client.insert_batches(table_id, batches).await.map_err(EtlError::from)?; + + Ok(()) + } + + /// Process CDC events. + /// + /// All events (inserts, updates, deletes, truncates, and schema changes) + /// from the replication stream are processed here. + pub async fn process_events(&self, events: Vec) -> EtlResult<()> { + let mut iter = events.into_iter().peekable(); + + while iter.peek().is_some() { + let builders = self.accumulate_data_events(&mut iter).await?; + self.flush_batches(builders).await?; + self.apply_relation_events(&mut iter).await?; + self.apply_truncate_events(&mut iter).await?; + } + + Ok(()) + } + + /// Last offset committed by Snowflake for this table's channel. + pub async fn committed_offset(&self, table_id: TableId) -> EtlResult> { + self.client.committed_offset(table_id).await.map_err(EtlError::from) + } + + async fn accumulate_data_events( + &self, + iter: &mut EventIter, + ) -> EtlResult> { + let mut builders: HashMap = HashMap::new(); + let mut column_cache: HashMap> = HashMap::new(); + + // Consume data events (insert/update/delete) into per-table batch builders, + // stopping at barrier events (truncate, relation) that require a flush before + // they can be applied. + while let Some(event) = iter.peek() { + if matches!(event, Event::Truncate(_) | Event::Relation(_)) { + break; + } + let event = iter.next().expect("iterator is non-empty after peek"); + match event { + Event::Insert(e) => self.encode_insert(e, &mut builders, &mut column_cache).await?, + Event::Update(e) => self.encode_update(e, &mut builders, &mut column_cache).await?, + Event::Delete(e) => self.encode_delete(e, &mut builders, &mut column_cache).await?, + _ => {} + } + } + + Ok(builders) + } + + async fn flush_batches(&self, builders: HashMap) -> EtlResult<()> { + for (table_id, builder) in builders { + let batches = builder.finish().map_err(EtlError::from)?; + self.client.insert_batches(table_id, batches).await.map_err(EtlError::from)?; + } + Ok(()) + } + + async fn apply_relation_events(&self, iter: &mut EventIter) -> EtlResult<()> { + while matches!(iter.peek(), Some(Event::Relation(_))) { + if let Some(Event::Relation(rel)) = iter.next() { + self.handle_relation_event(&rel.replicated_table_schema).await?; + } + } + Ok(()) + } + + async fn apply_truncate_events(&self, iter: &mut EventIter) -> EtlResult<()> { + // Collect and dedup tables to be truncated. + let mut truncated: HashMap = HashMap::new(); + while matches!(iter.peek(), Some(Event::Truncate(_))) { + if let Some(Event::Truncate(t)) = iter.next() { + for schema in t.truncated_tables { + truncated.insert(schema.id(), schema); + } + } + } + + // Truncate tables. + for (_, schema) in truncated { + let table_name = try_stringify_table_name(schema.name())?.to_uppercase(); + self.client.truncate_table(schema.id(), &table_name).await.map_err(EtlError::from)?; + } + + Ok(()) + } + + async fn encode_insert( + &self, + e: InsertEvent, + builders: &mut HashMap, + column_cache: &mut HashMap>, + ) -> EtlResult<()> { + let table_id = e.replicated_table_schema.id(); + self.ensure_column_cache(column_cache, table_id, &e.replicated_table_schema).await?; + + let cols = &column_cache[&table_id]; + let offset = OffsetToken::new(e.commit_lsn, e.tx_ordinal); + builders + .entry(table_id) + .or_default() + .push_row( + cols, + &e.table_row, + CdcMeta::new(CdcOperation::Insert, offset.as_ref()), + &offset, + ) + .map_err(EtlError::from) + } + + async fn encode_update( + &self, + e: UpdateEvent, + builders: &mut HashMap, + column_cache: &mut HashMap>, + ) -> EtlResult<()> { + // Accept only full rows, otherwise NULL will be recorded for missing columns + // (not that they are not changed). + let full_row = match e.updated_table_row { + UpdatedTableRow::Full(row) => row, + UpdatedTableRow::Partial(_) => { + bail!( + ErrorKind::InvalidData, + "Partial update rows not supported", + "Snowflake destination requires REPLICA IDENTITY FULL for update events" + ); + } + }; + + let table_id = e.replicated_table_schema.id(); + self.ensure_column_cache(column_cache, table_id, &e.replicated_table_schema).await?; + + let cols = &column_cache[&table_id]; + let offset = OffsetToken::new(e.commit_lsn, e.tx_ordinal); + builders + .entry(table_id) + .or_default() + .push_row(cols, &full_row, CdcMeta::new(CdcOperation::Update, offset.as_ref()), &offset) + .map_err(EtlError::from) + } + + async fn encode_delete( + &self, + e: DeleteEvent, + builders: &mut HashMap, + column_cache: &mut HashMap>, + ) -> EtlResult<()> { + let table_id = e.replicated_table_schema.id(); + let offset = OffsetToken::new(e.commit_lsn, e.tx_ordinal); + + match e.old_table_row { + Some(OldTableRow::Full(row)) => { + self.ensure_column_cache(column_cache, table_id, &e.replicated_table_schema) + .await?; + let cols = &column_cache[&table_id]; + builders + .entry(table_id) + .or_default() + .push_row( + cols, + &row, + CdcMeta::new(CdcOperation::Delete, offset.as_ref()), + &offset, + ) + .map_err(EtlError::from) + } + Some(OldTableRow::Key(key_row)) => { + self.ensure_column_cache(column_cache, table_id, &e.replicated_table_schema) + .await?; + let identity_cols: Vec<_> = + e.replicated_table_schema.identity_column_schemas().cloned().collect(); + builders + .entry(table_id) + .or_default() + .push_row( + &identity_cols, + &key_row, + CdcMeta::new(CdcOperation::Delete, offset.as_ref()), + &offset, + ) + .map_err(EtlError::from) + } + None => { + info!(table_id = ?table_id, "delete event has no old row data, skipping"); + Ok(()) + } + } + } + + async fn ensure_column_cache( + &self, + column_cache: &mut HashMap>, + table_id: TableId, + table_schema: &ReplicatedTableSchema, + ) -> EtlResult<()> { + #[allow(clippy::map_entry)] + if !column_cache.contains_key(&table_id) { + self.prepare_table_for_streaming(table_schema).await?; + let cols: Vec<_> = table_schema.column_schemas().cloned().collect(); + column_cache.insert(table_id, cols); + } + Ok(()) + } + + async fn handle_relation_event(&self, new_schema: &ReplicatedTableSchema) -> EtlResult<()> { + let table_id = new_schema.id(); + let new_snapshot_id = new_schema.inner().snapshot_id; + + let Some(metadata) = self.store.get_applied_destination_table_metadata(table_id).await? + else { + bail!( + ErrorKind::CorruptedTableSchema, + "Missing destination table metadata", + format!( + "No destination table metadata found for table {table_id} when processing \ + schema change" + ) + ); + }; + + let current_snapshot_id = metadata.snapshot_id; + let current_replication_mask = metadata.replication_mask.clone(); + let new_replication_mask = new_schema.replication_mask().clone(); + + if current_snapshot_id == new_snapshot_id + && current_replication_mask == new_replication_mask + { + info!(table_id = ?table_id, "schema unchanged, skipping relation event"); + return Ok(()); + } + + info!( + table_id = ?table_id, + "schema change detected: snapshot_id {current_snapshot_id} -> {new_snapshot_id}" + ); + + let current_table_schema = + self.store.get_table_schema(&table_id, current_snapshot_id).await?.ok_or_else( + || { + etl_error!( + ErrorKind::InvalidState, + "Old schema not found", + format!( + "Could not find schema for table {table_id} at snapshot_id \ + {current_snapshot_id}" + ) + ) + }, + )?; + + let current_schema = ReplicatedTableSchema::from_mask( + current_table_schema, + current_replication_mask.clone(), + ); + + let table_name = try_stringify_table_name(new_schema.name())?.to_uppercase(); + + let updated_metadata = DestinationTableMetadata::new_applied( + metadata.destination_table_id.clone(), + current_snapshot_id, + current_replication_mask, + ) + .with_schema_change( + new_snapshot_id, + new_replication_mask, + DestinationTableSchemaStatus::Applying, + ); + self.store.store_destination_table_metadata(table_id, updated_metadata.clone()).await?; + + let diff = current_schema.diff(new_schema); + if let Err(err) = + self.client.apply_schema_diff(&table_name, &diff).await.map_err(EtlError::from) + { + warn!( + table_id = ?table_id, + error = %err, + "schema change failed, manual intervention may be required" + ); + return Err(err); + } + + self.store + .store_destination_table_metadata(table_id, updated_metadata.to_applied()) + .await?; + + let new_columns: Vec<_> = new_schema.column_schemas().cloned().collect(); + schema::validate_no_cdc_collisions(&new_columns).map_err(EtlError::from)?; + self.client.refresh_table(&table_id).await.map_err(EtlError::from)?; + + info!(table_id = ?table_id, "schema change applied"); + Ok(()) + } +} + +impl etl::destination::Destination for Destination +where + S: StateStore + SchemaStore + Clone + Send + Sync + 'static, + T: TokenProvider + 'static, + C: StreamClient, +{ + fn name() -> &'static str { + "snowflake" + } + + async fn shutdown(&self) -> EtlResult<()> { + self.tasks.shutdown().await + } + + async fn truncate_table( + &self, + replicated_table_schema: &ReplicatedTableSchema, + async_result: TruncateTableResult<()>, + ) -> EtlResult<()> { + self.prepare_table_for_streaming(replicated_table_schema).await?; + let table_name = try_stringify_table_name(replicated_table_schema.name())?.to_uppercase(); + let result = self + .client + .truncate_table(replicated_table_schema.id(), &table_name) + .await + .map_err(EtlError::from); + async_result.send(result); + Ok(()) + } + + async fn write_table_rows( + &self, + replicated_table_schema: &ReplicatedTableSchema, + table_rows: Vec, + async_result: WriteTableRowsResult<()>, + ) -> EtlResult<()> { + let result = self.write_table_rows(replicated_table_schema, table_rows).await; + async_result.send(result); + Ok(()) + } + + async fn write_events( + &self, + events: Vec, + async_result: WriteEventsResult<()>, + ) -> EtlResult<()> { + self.tasks.try_reap().await?; + + let destination = self.clone(); + self.tasks + .spawn(async move { + let result = destination.process_events(events).await; + async_result.send(result); + }) + .await; + + Ok(()) + } +} diff --git a/crates/etl-destinations/src/snowflake/encoding.rs b/crates/etl-destinations/src/snowflake/encoding.rs new file mode 100644 index 000000000..00285c05e --- /dev/null +++ b/crates/etl-destinations/src/snowflake/encoding.rs @@ -0,0 +1,469 @@ +use std::{fmt, io::Write}; + +use etl::types::{ + ArrayCell, Cell, ColumnSchema, DATE_FORMAT, PgNumeric, TIME_FORMAT, TIMESTAMP_FORMAT, + TIMESTAMPTZ_FORMAT_HH_MM, TableRow, +}; +use serde::{ + Serialize, + ser::{SerializeMap, SerializeSeq, Serializer}, +}; + +use crate::snowflake::{ + Error, Result, + schema::{CDC_OPERATION_COLUMN, CDC_SEQUENCE_COLUMN}, +}; + +/// CDC operation type appended to every row sent via Snowpipe Streaming. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum CdcOperation { + Insert, + Update, + Delete, +} + +impl CdcOperation { + /// Returns string written into the `_cdc_operation` column. + pub fn as_str(self) -> &'static str { + match self { + Self::Insert => "insert", + Self::Update => "update", + Self::Delete => "delete", + } + } +} + +/// CDC metadata attached to every row in a batch. +#[derive(Debug, Clone, Copy)] +pub struct CdcMeta<'a> { + /// Operation performed. + pub(crate) operation: CdcOperation, + /// WAL sequence identifier obtained from `OffsetToken`. + pub(crate) sequence: &'a str, +} + +impl<'a> CdcMeta<'a> { + /// Create new CDC meta. + pub fn new(operation: CdcOperation, sequence: &'a str) -> Self { + Self { operation, sequence } + } +} + +/// Serialize a single row as an NDJSON line into any `Write` sink. +pub(crate) fn serialize_row( + writer: &mut impl Write, + cols: &[ColumnSchema], + row: &TableRow, + cdc: CdcMeta<'_>, +) -> Result<()> { + let serializable = RowSerializer { + cols, + cells: row.values(), + operation: cdc.operation.as_str(), + sequence: cdc.sequence, + }; + serde_json::to_writer(&mut *writer, &serializable) + .map_err(|e| Error::Encoding(e.to_string()))?; + writer.write_all(b"\n").map_err(|e| Error::Encoding(e.to_string()))?; + Ok(()) +} + +struct RowSerializer<'a> { + cols: &'a [ColumnSchema], + cells: &'a [Cell], + operation: &'a str, + sequence: &'a str, +} + +impl Serialize for RowSerializer<'_> { + fn serialize(&self, ser: S) -> std::result::Result { + let mut map = ser.serialize_map(Some(self.cols.len() + 2))?; + for (col, cell) in self.cols.iter().zip(self.cells) { + map.serialize_entry(col.name.as_str(), &CellSerializer(cell))?; + } + map.serialize_entry(CDC_OPERATION_COLUMN, self.operation)?; + map.serialize_entry(CDC_SEQUENCE_COLUMN, self.sequence)?; + map.end() + } +} + +struct CellSerializer<'a>(&'a Cell); + +impl Serialize for CellSerializer<'_> { + fn serialize(&self, ser: S) -> std::result::Result { + match self.0 { + Cell::Null => ser.serialize_none(), + Cell::Bool(b) => ser.serialize_bool(*b), + Cell::String(s) => ser.serialize_str(s), + Cell::I16(n) => ser.serialize_i16(*n), + Cell::I32(n) => ser.serialize_i32(*n), + Cell::U32(n) => ser.serialize_u32(*n), + Cell::I64(n) => ser.serialize_i64(*n), + Cell::F32(f) => { + reject_non_finite(*f as f64)?; + ser.serialize_f32(*f) + } + Cell::F64(f) => { + reject_non_finite(*f)?; + ser.serialize_f64(*f) + } + Cell::Numeric(n) => serialize_pg_numeric(n, ser), + // collect_str: Display::fmt writes directly into the JSON serializer's + // output buffer, avoiding an intermediate String allocation. + Cell::Date(d) => ser.collect_str(&d.format(DATE_FORMAT)), + Cell::Time(t) => ser.collect_str(&t.format(TIME_FORMAT)), + Cell::Timestamp(dt) => ser.collect_str(&dt.format(TIMESTAMP_FORMAT)), + Cell::TimestampTz(dt) => ser.collect_str(&dt.format(TIMESTAMPTZ_FORMAT_HH_MM)), + Cell::Uuid(u) => ser.collect_str(u), + Cell::Json(v) => v.serialize(ser), + Cell::Bytes(b) => ser.collect_str(&HexDisplay(b)), + Cell::Array(arr) => ArrayCellSerializer(arr).serialize(ser), + } + } +} + +fn reject_non_finite(f: f64) -> std::result::Result<(), E> { + if f.is_nan() || f.is_infinite() { + return Err(E::custom(format!( + "Snowflake does not support NaN/Infinity float values: {f}" + ))); + } + Ok(()) +} + +fn serialize_pg_numeric( + n: &PgNumeric, + ser: S, +) -> std::result::Result { + match n { + PgNumeric::NaN => Err(serde::ser::Error::custom("Snowflake NUMBER does not support NaN")), + PgNumeric::PositiveInfinity | PgNumeric::NegativeInfinity => { + Err(serde::ser::Error::custom("Snowflake NUMBER does not support Infinity")) + } + // PgNumeric `Display` writes directly via `collect_str`, no extra alloc. + PgNumeric::Value { .. } => ser.collect_str(n), + } +} + +/// Zero-allocation hex formatter for byte slices. +/// +/// Implements `Display` so it can be used with `collect_str()` to write hex +/// directly into the JSON buffer, avoiding the `String` that `hex::encode` +/// would allocate. +/// +/// For arrays, wrap in `CollectStr(HexDisplay(b))`. +struct HexDisplay<'a>(&'a [u8]); + +impl fmt::Display for HexDisplay<'_> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + for byte in self.0 { + write!(f, "{byte:02x}")?; + } + Ok(()) + } +} + +/// Display to Serialize adapter. +/// +/// Types that implement Display trait, get Serialize via this adapter. +/// +/// For instance, `Uuid` only implements `Display`, so this wrapper bridges the +/// gap: its `Serialize` implementation calls `collect_str`, thus keeping the +/// same zero-allocation path. +struct CollectStr(T); + +impl Serialize for CollectStr { + fn serialize(&self, ser: S) -> std::result::Result { + ser.collect_str(&self.0) + } +} + +struct ValidatedF32(f32); + +impl Serialize for ValidatedF32 { + fn serialize(&self, ser: S) -> std::result::Result { + reject_non_finite(self.0 as f64)?; + ser.serialize_f32(self.0) + } +} + +struct ValidatedF64(f64); + +impl Serialize for ValidatedF64 { + fn serialize(&self, ser: S) -> std::result::Result { + reject_non_finite(self.0)?; + ser.serialize_f64(self.0) + } +} + +struct NumericElement<'a>(&'a PgNumeric); + +impl Serialize for NumericElement<'_> { + fn serialize(&self, ser: S) -> std::result::Result { + serialize_pg_numeric(self.0, ser) + } +} + +struct ArrayCellSerializer<'a>(&'a ArrayCell); + +impl Serialize for ArrayCellSerializer<'_> { + fn serialize(&self, ser: S) -> std::result::Result { + match self.0 { + // Primitives: Option already implements Serialize correctly. + ArrayCell::Bool(v) => serialize_array(v, ser), + ArrayCell::String(v) => serialize_array(v, ser), + ArrayCell::I16(v) => serialize_array(v, ser), + ArrayCell::I32(v) => serialize_array(v, ser), + ArrayCell::U32(v) => serialize_array(v, ser), + ArrayCell::I64(v) => serialize_array(v, ser), + ArrayCell::Json(v) => serialize_array(v, ser), + // Validated floats: reject NaN/Infinity during serialization. + ArrayCell::F32(v) => serialize_array_with(v, ser, |f| ValidatedF32(*f)), + ArrayCell::F64(v) => serialize_array_with(v, ser, |f| ValidatedF64(*f)), + // Custom formatting via collect_str wrappers. + ArrayCell::Numeric(v) => serialize_array_with(v, ser, NumericElement), + ArrayCell::Date(v) => { + serialize_array_with(v, ser, |d| CollectStr(d.format(DATE_FORMAT))) + } + ArrayCell::Time(v) => { + serialize_array_with(v, ser, |t| CollectStr(t.format(TIME_FORMAT))) + } + ArrayCell::Timestamp(v) => { + serialize_array_with(v, ser, |dt| CollectStr(dt.format(TIMESTAMP_FORMAT))) + } + ArrayCell::TimestampTz(v) => { + serialize_array_with(v, ser, |dt| CollectStr(dt.format(TIMESTAMPTZ_FORMAT_HH_MM))) + } + ArrayCell::Uuid(v) => serialize_array_with(v, ser, CollectStr), + ArrayCell::Bytes(v) => serialize_array_with(v, ser, |b| CollectStr(HexDisplay(b))), + } + } +} + +/// Serialize a slice of `Serialize` items as a JSON array. +/// +/// Used for primitive array variants where `Option: Serialize`. +fn serialize_array( + items: &[T], + ser: S, +) -> std::result::Result { + let mut seq = ser.serialize_seq(Some(items.len()))?; + for item in items { + seq.serialize_element(item)?; + } + seq.end() +} + +/// Like `serialize_array`, but wraps each `Some` element via `wrap`. +/// +/// Used for array variants needing custom serialization (dates, validated +/// floats, etc.). +fn serialize_array_with<'a, T: 'a, S, R>( + items: &'a [Option], + ser: S, + wrap: impl Fn(&'a T) -> R, +) -> std::result::Result +where + S: Serializer, + R: Serialize, +{ + let mut seq = ser.serialize_seq(Some(items.len()))?; + for item in items { + match item { + Some(v) => seq.serialize_element(&wrap(v))?, + None => seq.serialize_element(&None::<()>)?, + } + } + seq.end() +} + +#[cfg(test)] +mod tests { + use chrono::{NaiveDate, NaiveDateTime, NaiveTime, TimeZone, Utc}; + use etl::types::Type; + use serde_json::{Value, json}; + use uuid::Uuid; + + use super::*; + + fn col(name: &str) -> ColumnSchema { + ColumnSchema::new(name.to_owned(), Type::TEXT, -1, 1, None, true) + } + + fn push_single_cell(cell: Cell) -> std::result::Result { + let cols = [col("v")]; + let mut buf = Vec::new(); + + let row = TableRow::new(vec![cell]); + serialize_row(&mut buf, &cols, &row, CdcMeta::new(CdcOperation::Insert, "0"))?; + + let line = std::str::from_utf8(&buf).unwrap().trim(); + let map: serde_json::Map = serde_json::from_str(line).unwrap(); + + Ok(map.get("v").unwrap().clone()) + } + + #[test] + fn cell_serialization_ok() { + let d = NaiveDate::from_ymd_opt(2026, 4, 29).unwrap(); + let t = NaiveTime::from_hms_micro_opt(10, 30, 0, 123456).unwrap(); + + let cases: Vec<(Cell, Value)> = vec![ + (Cell::Null, Value::Null), + (Cell::Bool(true), Value::Bool(true)), + (Cell::Bool(false), Value::Bool(false)), + (Cell::String("hello".into()), json!("hello")), + (Cell::I16(42), json!(42i16)), + (Cell::I32(i32::MAX), json!(i32::MAX)), + (Cell::U32(u32::MAX), json!(u32::MAX)), + (Cell::I64(i64::MAX), json!(i64::MAX)), + (Cell::F32(1.5), json!(1.5f64)), + (Cell::F64(2.5), json!(2.5)), + (Cell::Numeric(PgNumeric::default()), json!(PgNumeric::default().to_string())), + (Cell::Date(d), json!("2026-04-29")), + (Cell::Time(t), json!(t.format(TIME_FORMAT).to_string())), + ( + Cell::Timestamp(NaiveDateTime::new(d, t)), + json!(NaiveDateTime::new(d, t).format(TIMESTAMP_FORMAT).to_string()), + ), + ( + Cell::TimestampTz(Utc.with_ymd_and_hms(2026, 4, 29, 10, 30, 0).unwrap()), + json!( + Utc.with_ymd_and_hms(2026, 4, 29, 10, 30, 0) + .unwrap() + .format(TIMESTAMPTZ_FORMAT_HH_MM) + .to_string() + ), + ), + (Cell::Uuid(Uuid::nil()), json!("00000000-0000-0000-0000-000000000000")), + (Cell::Json(json!({"key": [1, 2, 3]})), json!({"key": [1, 2, 3]})), + (Cell::Bytes(vec![0xDE, 0xAD, 0xBE, 0xEF]), json!("deadbeef")), + ]; + for (cell, expected) in cases { + let dbg = format!("{cell:?}"); + assert_eq!(push_single_cell(cell).unwrap(), expected, "cell: {dbg}"); + } + } + + #[test] + fn rejects_non_finite() { + let cases: Vec = vec![ + Cell::F64(f64::NAN), + Cell::F64(f64::INFINITY), + Cell::F64(f64::NEG_INFINITY), + Cell::F32(f32::NAN), + Cell::F32(f32::INFINITY), + Cell::F32(f32::NEG_INFINITY), + Cell::Numeric(PgNumeric::NaN), + Cell::Numeric(PgNumeric::PositiveInfinity), + Cell::Numeric(PgNumeric::NegativeInfinity), + ]; + for cell in cases { + let dbg = format!("{cell:?}"); + assert!(push_single_cell(cell).is_err(), "should reject: {dbg}"); + } + } + + #[test] + fn array_serialization_ok() { + let cases: Vec<(Cell, Value)> = vec![ + (Cell::Array(ArrayCell::I32(vec![Some(1), None, Some(3)])), json!([1, null, 3])), + (Cell::Array(ArrayCell::I32(vec![])), json!([])), + (Cell::Array(ArrayCell::Bool(vec![Some(true), None])), json!([true, null])), + (Cell::Array(ArrayCell::String(vec![Some("a".into()), None])), json!(["a", null])), + (Cell::Array(ArrayCell::Bytes(vec![Some(vec![0xFF]), None])), json!(["ff", null])), + ( + Cell::Array(ArrayCell::Uuid(vec![Some(Uuid::nil())])), + json!(["00000000-0000-0000-0000-000000000000"]), + ), + (Cell::Array(ArrayCell::Json(vec![Some(json!(1)), None])), json!([1, null])), + ( + Cell::Array(ArrayCell::F64(vec![Some(1.5), None, Some(2.5)])), + json!([1.5, null, 2.5]), + ), + ( + Cell::Array(ArrayCell::Date(vec![ + Some(NaiveDate::from_ymd_opt(2026, 4, 29).unwrap()), + None, + ])), + json!(["2026-04-29", null]), + ), + ]; + for (cell, expected) in cases { + let dbg = format!("{cell:?}"); + assert_eq!(push_single_cell(cell).unwrap(), expected, "cell: {dbg}"); + } + } + + #[test] + fn array_rejects_non_finite() { + let cases: Vec = vec![ + Cell::Array(ArrayCell::F64(vec![Some(f64::NAN)])), + Cell::Array(ArrayCell::F32(vec![Some(f32::INFINITY)])), + Cell::Array(ArrayCell::Numeric(vec![Some(PgNumeric::NaN)])), + ]; + for cell in cases { + let dbg = format!("{cell:?}"); + assert!(push_single_cell(cell).is_err(), "should reject: {dbg}"); + } + } + + #[test] + fn multi_column_row() { + let cols = [col("id"), col("name")]; + let mut buf = Vec::new(); + let row = TableRow::new(vec![Cell::I32(1), Cell::String("hello".into())]); + serialize_row(&mut buf, &cols, &row, CdcMeta::new(CdcOperation::Insert, "0")).unwrap(); + + let line = std::str::from_utf8(&buf).unwrap().trim(); + let map: serde_json::Map = serde_json::from_str(line).unwrap(); + assert_eq!(map.get("id").unwrap(), &json!(1)); + assert_eq!(map.get("name").unwrap(), &json!("hello")); + } + + #[test] + fn multi_row_ndjson() { + let cols = [col("id")]; + let mut buf = Vec::new(); + + serialize_row( + &mut buf, + &cols, + &TableRow::new(vec![Cell::I32(1)]), + CdcMeta::new(CdcOperation::Insert, "0"), + ) + .unwrap(); + serialize_row( + &mut buf, + &cols, + &TableRow::new(vec![Cell::I32(2)]), + CdcMeta::new(CdcOperation::Insert, "0"), + ) + .unwrap(); + + let text = std::str::from_utf8(&buf).unwrap(); + let lines: Vec<&str> = text.trim_end().split('\n').collect(); + assert_eq!(lines.len(), 2); + let v0: Value = serde_json::from_str(lines[0]).unwrap(); + assert_eq!(v0["id"], json!(1)); + let v1: Value = serde_json::from_str(lines[1]).unwrap(); + assert_eq!(v1["id"], json!(2)); + } + + #[test] + fn cdc_columns() { + let cols = [col("id"), col("name")]; + let mut buf = Vec::new(); + + let row = TableRow::new(vec![Cell::I32(1), Cell::String("Alice".into())]); + serialize_row(&mut buf, &cols, &row, CdcMeta::new(CdcOperation::Insert, "0000/0000")) + .unwrap(); + + let line = std::str::from_utf8(&buf).unwrap().trim(); + let map: serde_json::Map = serde_json::from_str(line).unwrap(); + assert_eq!(map.get("id").unwrap(), &json!(1)); + assert_eq!(map.get("name").unwrap(), &json!("Alice")); + assert_eq!(map.get("_cdc_operation").unwrap(), &json!("insert")); + assert_eq!(map.get("_cdc_sequence_number").unwrap(), &json!("0000/0000")); + } +} diff --git a/crates/etl-destinations/src/snowflake/error.rs b/crates/etl-destinations/src/snowflake/error.rs new file mode 100644 index 000000000..6d3a1d3f2 --- /dev/null +++ b/crates/etl-destinations/src/snowflake/error.rs @@ -0,0 +1,52 @@ +use etl::error::{ErrorKind, EtlError}; +use reqwest::StatusCode; + +#[derive(Debug, thiserror::Error)] +pub enum Error { + #[error("HTTP transport error: {0}")] + HttpTransport(#[from] reqwest::Error), + + #[error("HTTP status {status}: {body}")] + HttpStatus { status: StatusCode, body: String }, + + #[error("Authentication error: {0}")] + Auth(String), + + #[error("SQL error{}: {message}", statement_handle.as_ref().map(|h| format!(" (handle {h})")).unwrap_or_default())] + Sql { statement_handle: Option, message: String }, + + #[error("Snowpipe error (code {status_code}): {message}")] + Snowpipe { status_code: u32, message: String }, + + #[error("Channel error: {0}")] + Channel(String), + + #[error("Encoding error: {0}")] + Encoding(String), + + #[error("Configuration error: {0}")] + Config(String), +} + +impl From for EtlError { + fn from(err: Error) -> Self { + let (kind, description) = match &err { + Error::HttpTransport(_) => { + (ErrorKind::DestinationError, "Snowflake HTTP transport error") + } + Error::HttpStatus { status, .. } if status.is_server_error() => { + (ErrorKind::DestinationError, "Snowflake server error") + } + Error::HttpStatus { .. } => (ErrorKind::DestinationError, "Snowflake HTTP error"), + Error::Auth(_) => (ErrorKind::DestinationError, "Snowflake authentication failed"), + Error::Sql { .. } => (ErrorKind::DestinationError, "Snowflake SQL execution failed"), + Error::Snowpipe { .. } => (ErrorKind::DestinationError, "Snowpipe streaming error"), + Error::Channel(_) => (ErrorKind::DestinationError, "Snowflake channel error"), + Error::Encoding(_) => (ErrorKind::InvalidData, "Snowflake encoding error"), + Error::Config(_) => (ErrorKind::ConfigError, "Snowflake configuration error"), + }; + etl::etl_error!(kind, description, err.to_string()) + } +} + +pub type Result = std::result::Result; diff --git a/crates/etl-destinations/src/snowflake/metrics.rs b/crates/etl-destinations/src/snowflake/metrics.rs new file mode 100644 index 000000000..e5622bb59 --- /dev/null +++ b/crates/etl-destinations/src/snowflake/metrics.rs @@ -0,0 +1,35 @@ +use std::sync::Once; + +use metrics::{Unit, describe_counter, describe_histogram}; + +static REGISTER_METRICS: Once = Once::new(); + +pub(super) const ETL_SNOWFLAKE_BATCH_SIZE: &str = "etl_snowflake_batch_size"; +pub(super) const ETL_SNOWFLAKE_BATCH_BYTES: &str = "etl_snowflake_batch_bytes"; +pub(super) const ETL_SNOWFLAKE_INSERT_ERRORS_TOTAL: &str = "etl_snowflake_insert_errors_total"; +pub(super) const ETL_SNOWFLAKE_CHANNEL_RECOVERIES_TOTAL: &str = + "etl_snowflake_channel_recoveries_total"; + +pub(super) fn register_metrics() { + REGISTER_METRICS.call_once(|| { + describe_histogram!(ETL_SNOWFLAKE_BATCH_SIZE, Unit::Count, "Rows per insert_rows request"); + + describe_histogram!( + ETL_SNOWFLAKE_BATCH_BYTES, + Unit::Bytes, + "Batch size in bytes (compressed)" + ); + + describe_counter!( + ETL_SNOWFLAKE_INSERT_ERRORS_TOTAL, + Unit::Count, + "Total insert_rows errors from Snowpipe Streaming" + ); + + describe_counter!( + ETL_SNOWFLAKE_CHANNEL_RECOVERIES_TOTAL, + Unit::Count, + "Channel recovery count (GC'd or errored channels)" + ); + }); +} diff --git a/crates/etl-destinations/src/snowflake/mod.rs b/crates/etl-destinations/src/snowflake/mod.rs new file mode 100644 index 000000000..9892c6628 --- /dev/null +++ b/crates/etl-destinations/src/snowflake/mod.rs @@ -0,0 +1,23 @@ +mod auth; +mod client; +mod config; +mod core; +mod encoding; +mod error; +mod metrics; +mod schema; +mod sql_client; +mod streaming; + +#[cfg(feature = "test-utils")] +pub mod test_utils; + +pub use core::Destination; + +pub use auth::{AuthManager, HttpExchanger, TokenProvider}; +pub use client::Client; +pub use config::Config; +pub use encoding::{CdcMeta, CdcOperation}; +pub use error::{Error, Result}; +pub use sql_client::SqlClient; +pub use streaming::{OffsetToken, RestStreamClient, RowBatch, RowBatchBuilder, StreamClient}; diff --git a/crates/etl-destinations/src/snowflake/schema.rs b/crates/etl-destinations/src/snowflake/schema.rs new file mode 100644 index 000000000..316eefd7d --- /dev/null +++ b/crates/etl-destinations/src/snowflake/schema.rs @@ -0,0 +1,158 @@ +use etl::types::{ColumnSchema, Type, is_array_type}; + +use crate::snowflake::{Error, Result}; + +pub(crate) const CDC_OPERATION_COLUMN: &str = "_cdc_operation"; +pub(crate) const CDC_SEQUENCE_COLUMN: &str = "_cdc_sequence_number"; + +/// Double-quote a SQL identifier, escaping internal double-quotes. +pub(crate) fn quote_identifier(name: &str) -> String { + format!("\"{}\"", name.replace('"', "\"\"")) +} + +/// Returns the Snowflake DDL type string for a given Postgres type. +/// +/// Array types map to ARRAY (Snowflake's native array type, a subtype of +/// VARIANT). Scalar types follow the closest Snowflake equivalent. +pub(crate) fn type_name(typ: &Type) -> &'static str { + if is_array_type(typ) { + return "ARRAY"; + } + + match typ { + &Type::BOOL => "BOOLEAN", + &Type::CHAR | &Type::BPCHAR | &Type::VARCHAR | &Type::NAME | &Type::TEXT => "VARCHAR", + &Type::INT2 => "SMALLINT", + &Type::INT4 => "INTEGER", + &Type::INT8 => "BIGINT", + &Type::FLOAT4 => "FLOAT", + &Type::FLOAT8 => "DOUBLE", + &Type::NUMERIC => "VARCHAR", + &Type::DATE => "DATE", + &Type::TIME => "TIME", + &Type::TIMESTAMP => "TIMESTAMP_NTZ", + &Type::TIMESTAMPTZ => "TIMESTAMP_TZ", + &Type::UUID => "VARCHAR", + &Type::JSON | &Type::JSONB => "VARIANT", + &Type::OID => "BIGINT", + &Type::BYTEA => "VARCHAR", + _ => "VARCHAR", + } +} + +/// Checks that no source column collides with the reserved CDC column names. +pub(crate) fn validate_no_cdc_collisions(columns: &[ColumnSchema]) -> Result<()> { + for col in columns { + if col.name == CDC_OPERATION_COLUMN || col.name == CDC_SEQUENCE_COLUMN { + return Err(Error::Config(format!( + "source column '{}' collides with reserved CDC column name", + col.name, + ))); + } + } + Ok(()) +} + +/// Builds the column definitions string. +/// +/// Each source column is rendered as `"name" TYPE` + two CDC columns are +/// appended at the end. +pub(crate) fn build_column_defs(columns: &[ColumnSchema]) -> String { + let mut parts: Vec = columns + .iter() + .map(|col| format!("{} {}", quote_identifier(&col.name), type_name(&col.typ),)) + .collect(); + + parts.push(format!("{} VARCHAR NOT NULL", quote_identifier(CDC_OPERATION_COLUMN))); + parts.push(format!("{} VARCHAR NOT NULL", quote_identifier(CDC_SEQUENCE_COLUMN))); + + parts.join(", ") +} + +#[cfg(test)] +mod tests { + use super::*; + + fn col(name: &str) -> ColumnSchema { + ColumnSchema::new(name.to_owned(), Type::INT4, -1, 1, None, true) + } + + #[test] + fn type_mapping() { + let cases: &[(&Type, &str)] = &[ + // Scalars + (&Type::BOOL, "BOOLEAN"), + (&Type::CHAR, "VARCHAR"), + (&Type::BPCHAR, "VARCHAR"), + (&Type::VARCHAR, "VARCHAR"), + (&Type::NAME, "VARCHAR"), + (&Type::TEXT, "VARCHAR"), + (&Type::INT2, "SMALLINT"), + (&Type::INT4, "INTEGER"), + (&Type::INT8, "BIGINT"), + (&Type::FLOAT4, "FLOAT"), + (&Type::FLOAT8, "DOUBLE"), + (&Type::NUMERIC, "VARCHAR"), + (&Type::DATE, "DATE"), + (&Type::TIME, "TIME"), + (&Type::TIMESTAMP, "TIMESTAMP_NTZ"), + (&Type::TIMESTAMPTZ, "TIMESTAMP_TZ"), + (&Type::UUID, "VARCHAR"), + (&Type::JSON, "VARIANT"), + (&Type::JSONB, "VARIANT"), + (&Type::OID, "BIGINT"), + (&Type::BYTEA, "VARCHAR"), + // Arrays all map to VARIANT + (&Type::BOOL_ARRAY, "ARRAY"), + (&Type::INT4_ARRAY, "ARRAY"), + (&Type::TEXT_ARRAY, "ARRAY"), + (&Type::JSONB_ARRAY, "ARRAY"), + (&Type::BYTEA_ARRAY, "ARRAY"), + // Unknown falls back to VARCHAR + (&Type::BIT, "VARCHAR"), + ]; + for (typ, expected) in cases { + assert_eq!(type_name(typ), *expected, "type: {typ:?}"); + } + } + + #[test] + fn cdc_collision_validation() { + let cases: &[(&[&str], bool)] = &[ + (&["id", "_cdc_operation"], true), + (&["id", "_cdc_sequence_number"], true), + (&["_cdc_operation", "_cdc_sequence_number"], true), + (&["id", "name", "_custom_cdc"], false), + ]; + for (col_names, should_err) in cases { + let columns: Vec<_> = col_names.iter().map(|n| col(n)).collect(); + assert_eq!( + validate_no_cdc_collisions(&columns).is_err(), + *should_err, + "columns: {col_names:?}" + ); + } + } + + #[test] + fn quote_identifier_cases() { + let cases = + [("my_table", r#""my_table""#), (r#"my"table"#, r#""my""table""#), ("", r#""""#)]; + for (input, expected) in cases { + assert_eq!(quote_identifier(input), expected, "input: {input:?}"); + } + } + + #[test] + fn build_column_defs_output() { + let columns = vec![ + ColumnSchema::new("id".to_owned(), Type::INT4, -1, 1, None, true), + ColumnSchema::new("created_at".to_owned(), Type::TIMESTAMPTZ, -1, 2, None, true), + ]; + let defs = build_column_defs(&columns); + assert_eq!( + defs, + r#""id" INTEGER, "created_at" TIMESTAMP_TZ, "_cdc_operation" VARCHAR NOT NULL, "_cdc_sequence_number" VARCHAR NOT NULL"# + ); + } +} diff --git a/crates/etl-destinations/src/snowflake/sql_client.rs b/crates/etl-destinations/src/snowflake/sql_client.rs new file mode 100644 index 000000000..b5334f09c --- /dev/null +++ b/crates/etl-destinations/src/snowflake/sql_client.rs @@ -0,0 +1,370 @@ +use std::{sync::Arc, time::Duration}; + +use reqwest::{Client, StatusCode}; +use serde::{Deserialize, Serialize}; +use tracing::{debug, warn}; + +use crate::{ + retry::{RetryDecision, RetryPolicy, retry_with_backoff}, + snowflake::{Config, Error, Result, auth::TokenProvider, schema::quote_identifier}, +}; + +/// Retry policy for transient HTTP errors (408, 429, 5xx) during SQL API calls. +const SQL_RETRY_POLICY: RetryPolicy = RetryPolicy { + max_retries: 3, + initial_delay: Duration::from_millis(500), + max_delay: Duration::from_secs(10), +}; + +/// Starting delay between polls when waiting for an async statement (HTTP 202). +const POLL_INITIAL_DELAY: Duration = Duration::from_millis(100); + +/// Upper bound on exponential backoff between poll requests. +const POLL_MAX_DELAY: Duration = Duration::from_secs(5); + +/// Hard deadline for async statement completion before returning a timeout +/// error. +const POLL_TIMEOUT: Duration = Duration::from_secs(30); + +/// Sent with every request to the Snowflake SQL REST API. +const USER_AGENT: &str = "supabase-etl/0.1.0"; + +/// Executes DDL and metadata operations against Snowflake's SQL REST API. +/// +/// All DDL runs on Snowflake's Cloud Services layer (no warehouse required). +pub struct SqlClient { + config: Config, + http: Client, + auth: Arc, +} + +#[derive(Debug, Serialize)] +struct StatementRequest<'a> { + statement: &'a str, + database: &'a str, + schema: &'a str, + #[serde(skip_serializing_if = "Option::is_none")] + role: Option<&'a str>, +} + +/// Snowflake SQL REST API response body. +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +pub(crate) struct StatementResponse { + #[serde(default)] + pub(crate) statement_handle: Option, + #[serde(default)] + pub(crate) message: Option, + #[serde(default)] + pub(crate) data: Option>>, +} + +impl SqlClient { + pub fn new(config: Config, auth: Arc, http: Client) -> Self { + Self { config, http, auth } + } + + /// Execute a DDL statement (runs on Cloud Services, no warehouse required). + pub async fn execute_ddl(&self, sql: &str) -> Result<()> { + self.execute_statement(sql).await?; + Ok(()) + } + + /// Create a table with the given columns, enabling schema evolution. + pub async fn create_table_if_not_exists( + &self, + table_name: &str, + column_defs: &str, + ) -> Result<()> { + let fqn = self.fully_qualified_name(table_name); + let sql = format!( + "CREATE TABLE IF NOT EXISTS {fqn} ({column_defs}) ENABLE_SCHEMA_EVOLUTION = TRUE" + ); + self.execute_ddl(&sql).await + } + + /// Remove all rows from a table without dropping it. + pub async fn truncate_table(&self, table_name: &str) -> Result<()> { + let fqn = self.fully_qualified_name(table_name); + self.execute_ddl(&format!("TRUNCATE TABLE {fqn}")).await + } + + /// Drop a table if it exists. + pub async fn drop_table(&self, table_name: &str) -> Result<()> { + let fqn = self.fully_qualified_name(table_name); + self.execute_ddl(&format!("DROP TABLE IF EXISTS {fqn}")).await + } + + /// Check whether a table exists in the configured database and schema. + pub async fn table_exists(&self, table_name: &str) -> Result { + let db = quote_identifier(&self.config.database); + let schema = quote_identifier(&self.config.schema); + let escaped = table_name + .replace('\\', "\\\\") + .replace('\'', "''") + .replace('_', "\\_") + .replace('%', "\\%"); + let sql = format!("SHOW TABLES LIKE '{escaped}' IN SCHEMA {db}.{schema}"); + let resp = self.execute_statement(&sql).await?; + Ok(resp.data.is_some_and(|rows| !rows.is_empty())) + } + + /// Add a nullable column to an existing table. + pub async fn add_column( + &self, + table_name: &str, + column_name: &str, + column_type: &str, + ) -> Result<()> { + let fqn = self.fully_qualified_name(table_name); + let col = quote_identifier(column_name); + self.execute_ddl(&format!("ALTER TABLE {fqn} ADD COLUMN {col} {column_type}")).await + } + + /// Remove a column from an existing table. + pub async fn drop_column(&self, table_name: &str, column_name: &str) -> Result<()> { + let fqn = self.fully_qualified_name(table_name); + let col = quote_identifier(column_name); + self.execute_ddl(&format!("ALTER TABLE {fqn} DROP COLUMN {col}")).await + } + + /// Rename a column in an existing table. + pub async fn rename_column( + &self, + table_name: &str, + old_name: &str, + new_name: &str, + ) -> Result<()> { + let fqn = self.fully_qualified_name(table_name); + let old = quote_identifier(old_name); + let new = quote_identifier(new_name); + self.execute_ddl(&format!("ALTER TABLE {fqn} RENAME COLUMN {old} TO {new}")).await + } + + fn fully_qualified_name(&self, name: &str) -> String { + format!( + "{}.{}.{}", + quote_identifier(&self.config.database), + quote_identifier(&self.config.schema), + quote_identifier(name), + ) + } + + /// Submit a SQL statement and return a resolved response. + /// + /// Handles the full Snowflake SQL REST API contract: + /// - HTTP 200: success, return parsed `StatementResponse`. + /// - HTTP 202: async execution, poll until 200 or 422. + /// - HTTP 422: SQL execution error, return `Error::Sql`. + /// - HTTP 401: invalidate cached token, retry once. + /// - HTTP 408/429/5xx: retriable via backoff. + /// - Other 4xx: non-retriable `Error::HttpStatus`. + pub(crate) async fn execute_statement(&self, sql: &str) -> Result { + let url = format!("{}/api/v2/statements", self.config.account_url()); + + let body = StatementRequest { + statement: sql, + database: &self.config.database, + schema: &self.config.schema, + role: self.config.role.as_deref(), + }; + + retry_with_backoff( + SQL_RETRY_POLICY, + classify_for_retry, + |d| d, + |attempt| { + warn!( + retry = attempt.retry_index, + max = attempt.max_retries, + delay_ms = attempt.sleep_delay.as_millis(), + error = %attempt.error, + "retrying SQL REST API request" + ); + }, + || self.attempt_statement(&url, &body), + ) + .await + .map_err(|f| f.last_error) + } + + /// Submit a statement and interpret the HTTP response. + /// + /// On a 401, invalidates the cached token and retries once. + async fn attempt_statement( + &self, + url: &str, + body: &StatementRequest<'_>, + ) -> Result { + let mut retried_auth = false; + + loop { + let token = self.auth.get_token().await?; + let http_resp = self.send_post(url, &token, body).await?; + let status = http_resp.status(); + + match status { + StatusCode::OK => return http_resp.json().await.map_err(Error::HttpTransport), + + StatusCode::ACCEPTED => { + let resp: StatementResponse = + http_resp.json().await.map_err(Error::HttpTransport)?; + return if let Some(ref handle) = resp.statement_handle { + debug!(statement_handle = %handle, "statement executing asynchronously, polling"); + self.poll_until_complete(handle).await + } else { + Err(Error::Sql { + statement_handle: None, + message: "received 202 without a statement handle".into(), + }) + }; + } + + StatusCode::UNPROCESSABLE_ENTITY => { + let resp: StatementResponse = + http_resp.json().await.map_err(Error::HttpTransport)?; + return Err(Error::Sql { + statement_handle: resp.statement_handle, + message: resp.message.unwrap_or_default(), + }); + } + + StatusCode::UNAUTHORIZED if !retried_auth => { + warn!("received 401 from SQL REST API, invalidating cached token"); + self.auth.invalidate_token().await; + retried_auth = true; + continue; + } + + _ => { + let body_text = http_resp.text().await.unwrap_or_default(); + return Err(Error::HttpStatus { status, body: body_text }); + } + } + } + } + + async fn send_post( + &self, + url: &str, + token: &str, + body: &StatementRequest<'_>, + ) -> Result { + self.http + .post(url) + .bearer_auth(token) + .header("User-Agent", USER_AGENT) + .json(body) + .send() + .await + .map_err(Error::HttpTransport) + } + + /// Poll an async statement until Snowflake returns 200 (success) or 422 + /// (failure). + async fn poll_until_complete(&self, statement_handle: &str) -> Result { + let url = format!("{}/api/v2/statements/{}", self.config.account_url(), statement_handle); + let deadline = tokio::time::Instant::now() + POLL_TIMEOUT; + let mut delay = POLL_INITIAL_DELAY; + + loop { + if tokio::time::Instant::now() >= deadline { + return Err(Error::Sql { + statement_handle: Some(statement_handle.to_owned()), + message: format!( + "statement did not complete within {}s", + POLL_TIMEOUT.as_secs() + ), + }); + } + + tokio::time::sleep(delay).await; + delay = (delay * 2).min(POLL_MAX_DELAY); + + let token = self.auth.get_token().await?; + let http_resp = self + .http + .get(&url) + .bearer_auth(&token) + .header("User-Agent", USER_AGENT) + .send() + .await + .map_err(Error::HttpTransport)?; + + match http_resp.status() { + StatusCode::OK => return http_resp.json().await.map_err(Error::HttpTransport), + StatusCode::ACCEPTED => { + debug!(statement_handle, "statement still running, continuing poll"); + continue; + } + StatusCode::UNPROCESSABLE_ENTITY => { + let resp: StatementResponse = + http_resp.json().await.map_err(Error::HttpTransport)?; + return Err(Error::Sql { + statement_handle: resp.statement_handle, + message: resp.message.unwrap_or_default(), + }); + } + StatusCode::TOO_MANY_REQUESTS => { + debug!(statement_handle, "rate limited during poll, backing off"); + continue; + } + other => { + let body_text = http_resp.text().await.unwrap_or_default(); + return Err(Error::HttpStatus { status: other, body: body_text }); + } + } + } + } +} + +fn classify_for_retry(error: &Error) -> RetryDecision { + match error { + Error::HttpTransport(_) => RetryDecision::Retry, + Error::HttpStatus { status, .. } => { + if *status == StatusCode::REQUEST_TIMEOUT + || *status == StatusCode::TOO_MANY_REQUESTS + || status.is_server_error() + { + RetryDecision::Retry + } else { + RetryDecision::Stop + } + } + _ => RetryDecision::Stop, + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn classify_for_retry_cases() { + let cases = [ + ( + Error::HttpStatus { + status: StatusCode::INTERNAL_SERVER_ERROR, + body: String::new(), + }, + RetryDecision::Retry, + ), + ( + Error::HttpStatus { status: StatusCode::TOO_MANY_REQUESTS, body: String::new() }, + RetryDecision::Retry, + ), + ( + Error::HttpStatus { status: StatusCode::REQUEST_TIMEOUT, body: String::new() }, + RetryDecision::Retry, + ), + ( + Error::HttpStatus { status: StatusCode::BAD_REQUEST, body: String::new() }, + RetryDecision::Stop, + ), + (Error::Sql { statement_handle: None, message: String::new() }, RetryDecision::Stop), + ]; + for (error, expected) in cases { + assert_eq!(classify_for_retry(&error), expected, "error: {error:?}"); + } + } +} diff --git a/crates/etl-destinations/src/snowflake/streaming/batch.rs b/crates/etl-destinations/src/snowflake/streaming/batch.rs new file mode 100644 index 000000000..ef15f4f58 --- /dev/null +++ b/crates/etl-destinations/src/snowflake/streaming/batch.rs @@ -0,0 +1,325 @@ +use std::{io::Write, mem}; + +use bytes::Bytes; +use etl::types::{ColumnSchema, TableRow}; +use zstd::stream::Encoder; + +use crate::snowflake::{ + Error, Result, + encoding::{CdcMeta, serialize_row}, + streaming::OffsetToken, +}; + +/// Snowflake Streaming API hard limit on the compressed HTTP request body. +const MAX_COMPRESSED_BYTES: usize = 4 * 1024 * 1024; + +/// Split when compressed output reaches this threshold. +/// +/// After a flush check decides not to split, up to [`MAX_UNFLUSHED_BYTES`] +/// (128KB) of input can arrive before the next check. So, the worst-case size +/// at `finish()` is roughly 3.8MB + 128KB ~ 3.93MB < 4MB hard limit. +/// +/// The 200KB headroom exceeds the unflushed bytes limit by design. +const BATCH_SPLIT_THRESHOLD: usize = 3_800_000; + +/// Max bytes written to the compressing encoder before forcing a flush. +const MAX_UNFLUSHED_BYTES: usize = 128 * 1024; + +/// Max serialized (uncompressed) size of a single row. +/// +/// Rejects degenerate TOAST rows before they enter the encoder. +const MAX_UNCOMPRESSED_ROW_BYTES: usize = 2 * 1024 * 1024; + +/// Pre-allocated capacity for the per-row serialization scratch buffer. +/// +/// One memory page covers most rows without reallocation, the buffer +/// grows automatically for larger rows and retains its high-water mark. +const SCRATCH_INITIAL_CAPACITY: usize = 4096; + +/// Compression level. +/// +/// Nice balance between compression and processor time spend compressing. +const ZSTD_COMPRESSION_LEVEL: i32 = 3; + +/// Batch of rows ready to be pushed to the Streaming API. +pub struct RowBatch { + data: Bytes, + row_count: usize, + offset: OffsetToken, +} + +impl RowBatch { + /// Payload bytes. + pub fn bytes(&self) -> &Bytes { + &self.data + } + + /// Byte length of the payload. + pub fn size(&self) -> usize { + self.data.len() + } + + /// Number of rows in this batch. + pub fn row_count(&self) -> usize { + self.row_count + } + + /// Offset token of the last row in this batch. + pub fn offset(&self) -> &OffsetToken { + &self.offset + } +} + +/// Builds compressed row batches with streaming zstd compression. +/// +/// Rows are serialized into a scratch buffer first, then written to the zstd +/// encoder. When the compressed output approaches `BATCH_SPLIT_THRESHOLD`, +/// the current batch is finished and a new encoder is started. +pub struct RowBatchBuilder { + encoder: Encoder<'static, Vec>, + scratch: Vec, + current_row_count: usize, + current_offset: OffsetToken, + input_since_flush: usize, + batches: Vec, +} + +impl Default for RowBatchBuilder { + fn default() -> Self { + Self::new() + } +} + +impl RowBatchBuilder { + pub fn new() -> Self { + Self { + encoder: new_encoder(), + scratch: Vec::with_capacity(SCRATCH_INITIAL_CAPACITY), + current_row_count: 0, + current_offset: OffsetToken::zero(), + input_since_flush: 0, + batches: Vec::new(), + } + } + + /// Append a row. + /// + /// Automatically creates a new batch when the compressed output approaches + /// the API limit. + pub fn push_row( + &mut self, + cols: &[ColumnSchema], + row: &TableRow, + cdc: CdcMeta<'_>, + offset: &OffsetToken, + ) -> Result<()> { + self.scratch.clear(); + serialize_row(&mut self.scratch, cols, row, cdc)?; + + if self.scratch.len() > MAX_UNCOMPRESSED_ROW_BYTES { + return Err(Error::Encoding(format!( + "single row exceeds {}B limit ({}B uncompressed)", + MAX_UNCOMPRESSED_ROW_BYTES, + self.scratch.len() + ))); + } + + // Flush compressing encoder and create a new batch, when necessary. + if self.input_since_flush + self.scratch.len() >= MAX_UNFLUSHED_BYTES { + self.encoder.flush().map_err(|e| Error::Encoding(format!("zstd flush: {e}")))?; + self.input_since_flush = 0; + + // Wrap up the current batch and start another on threshold. + if self.current_row_count > 0 + && self.compressed_size() + self.scratch.len() > BATCH_SPLIT_THRESHOLD + { + self.next_batch()?; + } + } + + self.encoder + .write_all(&self.scratch) + .map_err(|e| Error::Encoding(format!("zstd write: {e}")))?; + self.input_since_flush += self.scratch.len(); + self.current_row_count += 1; + self.current_offset = offset.clone(); + + Ok(()) + } + + /// Wrap-up building process, produce list of batches. + pub fn finish(mut self) -> Result> { + if self.current_row_count > 0 { + let compressed = + self.encoder.finish().map_err(|e| Error::Encoding(format!("zstd finish: {e}")))?; + + if compressed.len() > MAX_COMPRESSED_BYTES { + return Err(Error::Encoding(format!( + "compressed batch exceeds {}B API limit ({}B compressed)", + MAX_COMPRESSED_BYTES, + compressed.len() + ))); + } + + self.batches.push(RowBatch { + data: Bytes::from(compressed), + row_count: self.current_row_count, + offset: self.current_offset, + }); + } + + Ok(self.batches) + } + + fn compressed_size(&self) -> usize { + self.encoder.get_ref().len() + } + + fn next_batch(&mut self) -> Result<()> { + let old_encoder = mem::replace(&mut self.encoder, new_encoder()); + let compressed = + old_encoder.finish().map_err(|e| Error::Encoding(format!("zstd finish: {e}")))?; + + self.batches.push(RowBatch { + data: Bytes::from(compressed), + row_count: self.current_row_count, + offset: mem::take(&mut self.current_offset), + }); + + self.current_row_count = 0; + self.input_since_flush = 0; + + Ok(()) + } +} + +fn new_encoder() -> Encoder<'static, Vec> { + Encoder::new(Vec::new(), ZSTD_COMPRESSION_LEVEL) + .expect("hardcoded zstd compression level must be valid") +} + +#[cfg(test)] +mod tests { + use etl::types::{Cell, Type}; + + use super::*; + use crate::snowflake::encoding::{CdcMeta, CdcOperation}; + + fn col(name: &str) -> ColumnSchema { + ColumnSchema::new(name.to_owned(), Type::TEXT, -1, 1, None, true) + } + + fn incompressible_string(len: usize, seed: u64) -> String { + let mut s = String::with_capacity(len); + let mut state = seed; + for _ in 0..len { + state = state.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407); + let ch = (state >> 33) as u8 % 94 + 33; + s.push(ch as char); + } + s + } + + #[test] + fn single_batch() { + let cols = [col("id"), col("name")]; + let mut builder = RowBatchBuilder::new(); + + for i in 0..10 { + let row = TableRow::new(vec![Cell::I32(i), Cell::String(format!("row_{i}"))]); + let offset = OffsetToken::zero(); + builder + .push_row(&cols, &row, CdcMeta::new(CdcOperation::Insert, "0"), &offset) + .unwrap(); + } + + let batches = builder.finish().unwrap(); + assert_eq!(batches.len(), 1); + + let batch = &batches.first().unwrap(); + assert_eq!(batch.row_count(), 10); + assert!(batch.size() > 0); + assert!(batch.size() <= MAX_COMPRESSED_BYTES); + + let decompressed = zstd::decode_all(batch.bytes().as_ref()).unwrap(); + let text = String::from_utf8(decompressed).unwrap(); + let lines: Vec<&str> = text.trim_end().split('\n').collect(); + assert_eq!(lines.len(), 10); + for line in &lines { + serde_json::from_str::(line).expect("valid JSON"); + } + } + + #[test] + fn empty_builder() { + let builder = RowBatchBuilder::new(); + let batches = builder.finish().unwrap(); + assert!(batches.is_empty()); + assert_eq!(batches.len(), 0); + } + + #[test] + fn batch_splitting() { + let cols = [col("data")]; + let mut builder = RowBatchBuilder::new(); + + for i in 0..100 { + let value = incompressible_string(100_000, i as u64); + let row = TableRow::new(vec![Cell::String(value)]); + let offset = OffsetToken::zero(); + builder + .push_row(&cols, &row, CdcMeta::new(CdcOperation::Insert, &format!("{i}")), &offset) + .unwrap(); + } + + let batches = builder.finish().unwrap(); + assert!(batches.len() > 1, "expected multiple batches, got {}", batches.len()); + + for batch in &batches { + assert!( + batch.size() <= MAX_COMPRESSED_BYTES, + "batch exceeds hard limit: {}", + batch.size() + ); + assert!(batch.row_count() > 0); + + let decompressed = zstd::decode_all(batch.bytes().as_ref()).unwrap(); + let text = String::from_utf8(decompressed).unwrap(); + let lines: Vec<&str> = text.trim_end().split('\n').collect(); + assert_eq!(lines.len(), batch.row_count()); + } + } + + #[test] + fn oversized_row_rejected() { + let cols = [col("data")]; + let mut builder = RowBatchBuilder::new(); + + let huge_value = "x".repeat(MAX_UNCOMPRESSED_ROW_BYTES + 1); + let row = TableRow::new(vec![Cell::String(huge_value)]); + let err = builder + .push_row(&cols, &row, CdcMeta::new(CdcOperation::Insert, "0"), &OffsetToken::zero()) + .unwrap_err(); + + assert!(matches!(err, Error::Encoding(msg) if msg.contains("limit"))); + } + + #[test] + fn offset_tracks_last_row() { + let cols = [col("id")]; + let mut builder = RowBatchBuilder::new(); + + let offsets = ["0000000000000001/0000000000000000", "0000000000000002/0000000000000000"]; + for token_str in &offsets { + let row = TableRow::new(vec![Cell::I32(1)]); + let offset: OffsetToken = token_str.parse().unwrap(); + builder + .push_row(&cols, &row, CdcMeta::new(CdcOperation::Insert, "0"), &offset) + .unwrap(); + } + + let batches = builder.finish().unwrap(); + assert_eq!(batches.len(), 1); + assert_eq!(batches.first().unwrap().offset().as_ref(), offsets[1]); + } +} diff --git a/crates/etl-destinations/src/snowflake/streaming/channel.rs b/crates/etl-destinations/src/snowflake/streaming/channel.rs new file mode 100644 index 000000000..125f33b6a --- /dev/null +++ b/crates/etl-destinations/src/snowflake/streaming/channel.rs @@ -0,0 +1,165 @@ +use std::sync::Arc; + +use etl::types::PipelineId; +use metrics::{counter, histogram}; +use reqwest::StatusCode; +use tracing::warn; + +use crate::snowflake::{ + Error, OffsetToken, Result, RowBatch, StreamClient, + metrics::{ + ETL_SNOWFLAKE_BATCH_BYTES, ETL_SNOWFLAKE_BATCH_SIZE, + ETL_SNOWFLAKE_CHANNEL_RECOVERIES_TOTAL, ETL_SNOWFLAKE_INSERT_ERRORS_TOTAL, + }, +}; + +/// Manages the state and lifecycle of a single Snowpipe Streaming channel. +/// +/// Channel is a conduit via which we push data to Snowflake system. +#[derive(Debug)] +pub(crate) struct ChannelHandle { + /// Streaming API client. + client: Arc, + + /// Snowflake database name. + database: String, + + /// Snowflake schema name. + schema: String, + + /// Snowflake target table name. + table: String, + + /// Derived channel name. + channel: String, + + /// Last offset token returned by Snowflake API (on open or insert). + last_offset_token: Option, + + /// Continuation token for the next API call on this channel. + continuation_token: Option, +} + +impl Clone for ChannelHandle { + fn clone(&self) -> Self { + Self { + client: Arc::clone(&self.client), + database: self.database.clone(), + schema: self.schema.clone(), + table: self.table.clone(), + channel: self.channel.clone(), + last_offset_token: self.last_offset_token.clone(), + continuation_token: self.continuation_token.clone(), + } + } +} + +impl ChannelHandle { + /// New handle with no offset or continuation tokens. + pub fn new( + client: Arc, + pipeline: PipelineId, + database: String, + schema: String, + table: String, + ) -> Self { + let channel = format!("supabase_etl_{pipeline}_{schema}_{table}_ch0"); + Self { + client, + database, + schema, + table, + channel, + last_offset_token: None, + continuation_token: None, + } + } + + /// Open (or reopen) the channel. + /// + /// Idempotent, reopening preserves offset history. + pub async fn open(&mut self) -> Result<()> { + let response = self + .client + .open_channel(&self.database, &self.schema, &self.table, &self.channel) + .await?; + + self.last_offset_token = response.offset_token; + self.continuation_token = Some(response.continuation_token); + + Ok(()) + } + + /// Drop the channel. + /// + /// Committed data remains in the table. + pub async fn drop_channel(&mut self) -> Result<()> { + self.client.drop_channel(&self.database, &self.schema, &self.table, &self.channel).await?; + + self.last_offset_token = None; + self.continuation_token = None; + + Ok(()) + } + + /// Drop and reopen the channel, resetting offsets. + pub async fn reset(&mut self) -> Result<()> { + self.drop_channel().await?; + self.open().await + } + + /// Send all batches, recovering from channel GC errors. + pub async fn process_batches(&mut self, batches: Vec) -> Result<()> { + fn is_stale_channel(e: &Error) -> bool { + matches!(e, Error::Snowpipe { status_code: 4, .. }) + || matches!(e, Error::HttpStatus { status: StatusCode::NOT_FOUND, .. }) + } + + for batch in &batches { + histogram!(ETL_SNOWFLAKE_BATCH_SIZE).record(batch.row_count() as f64); + histogram!(ETL_SNOWFLAKE_BATCH_BYTES).record(batch.size() as f64); + + match self.send_batch(batch).await { + Ok(()) => {} + // If channel is stale, try to reopen it, once. + Err(e) if is_stale_channel(&e) => { + counter!(ETL_SNOWFLAKE_CHANNEL_RECOVERIES_TOTAL).increment(1); + warn!(table = %self.table, "channel was closed, reopenning and retrying insert"); + self.open().await?; + self.send_batch(batch).await?; + } + Err(e) => { + counter!(ETL_SNOWFLAKE_INSERT_ERRORS_TOTAL).increment(1); + return Err(e); + } + } + } + Ok(()) + } + + async fn send_batch(&mut self, batch: &RowBatch) -> Result<()> { + let ct = self.continuation_token.as_deref().ok_or_else(|| { + Error::Channel("send_batch called on channel without continuation token".into()) + })?; + + let response = self + .client + .insert_rows(&self.database, &self.schema, &self.table, &self.channel, batch, ct) + .await?; + + self.continuation_token = Some(response.continuation_token); + self.last_offset_token = Some(batch.offset().clone()); + + Ok(()) + } + + /// Last offset token committed by Snowflake for this channel. + pub async fn committed_offset(&self) -> Result> { + let status = self + .client + .channel_status(&self.database, &self.schema, &self.table, &self.channel) + .await?; + + Ok(status.offset_token) + } +} diff --git a/crates/etl-destinations/src/snowflake/streaming/mod.rs b/crates/etl-destinations/src/snowflake/streaming/mod.rs new file mode 100644 index 000000000..7574e5b82 --- /dev/null +++ b/crates/etl-destinations/src/snowflake/streaming/mod.rs @@ -0,0 +1,104 @@ +//! Abstraction over the Snowflake Streaming API. +mod batch; +mod channel; +mod offset_token; +mod rest_client; + +use std::future::Future; + +pub use batch::{RowBatch, RowBatchBuilder}; +pub(crate) use channel::ChannelHandle; +pub use offset_token::OffsetToken; +pub use rest_client::RestStreamClient; + +use crate::snowflake::Result; + +/// Response from opening or reopening a channel. +#[derive(Debug)] +pub struct OpenChannelResponse { + /// Server-managed sequencer token. + /// + /// Must be passed to subsequent `insert_rows` calls on this channel. + pub continuation_token: String, + + /// Last committed offset token. + /// + /// `None` if the channel has never committed data. + pub offset_token: Option, +} + +/// Response from inserting rows into a channel. +#[derive(Debug)] +pub struct InsertRowsResponse { + /// Updated continuation token for the next `insert_rows` call. + pub continuation_token: String, +} + +/// Per-channel status from a bulk status check. +#[derive(Debug)] +pub struct ChannelStatusResponse { + /// Channel name. + pub channel: String, + + /// Snowflake-assigned status code (e.g. `ACTIVE`, `SUCCESS`). + pub status_code: String, + + /// Last committed offset token, or `None` if no data has been committed. + pub offset_token: Option, +} + +/// Abstraction over Snowpipe Streaming ingestion backends. +/// +/// Enables swapping between REST API and SDK sidecar (should we decide to +/// implement it to scale ingestion 10x) without changing destination logic. +pub trait StreamClient: Send + Sync + 'static { + /// Discover the ingest hostname for this account. + /// + /// Normally, never changes and is discovered only once per account. + fn discover_ingest_host(&self) -> impl Future> + Send; + + /// Open or reopen a channel. + /// + /// Returns the continuation token and last committed offset token (if any). + fn open_channel( + &self, + database: &str, + schema: &str, + table: &str, + channel: &str, + ) -> impl Future> + Send; + + /// Drop a channel. + /// + /// Committed data remains in the table. + fn drop_channel( + &self, + database: &str, + schema: &str, + table: &str, + channel: &str, + ) -> impl Future> + Send; + + /// Insert a pre-built batch into a channel. + /// + /// The `continuation_token` is the server sequencer from `open_channel` + /// or the previous `insert_rows` call. + fn insert_rows( + &self, + database: &str, + schema: &str, + table: &str, + channel: &str, + batch: &RowBatch, + continuation_token: &str, + ) -> impl Future> + Send; + + /// Check channel status. + fn channel_status( + &self, + database: &str, + schema: &str, + table: &str, + channel: &str, + ) -> impl Future> + Send; +} diff --git a/crates/etl-destinations/src/snowflake/streaming/offset_token.rs b/crates/etl-destinations/src/snowflake/streaming/offset_token.rs new file mode 100644 index 000000000..739a1cb34 --- /dev/null +++ b/crates/etl-destinations/src/snowflake/streaming/offset_token.rs @@ -0,0 +1,68 @@ +use std::{fmt, str::FromStr}; + +use etl::types::PgLsn; + +use crate::snowflake::{Error, Result}; + +/// An offset token encoding a WAL position as a hex string. +/// +/// Format: `{commit_lsn:016x}/{tx_ordinal:016x}`. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct OffsetToken(String); + +impl OffsetToken { + /// The zero offset token, used for initial copy rows before any WAL + /// position. + pub fn zero() -> Self { + Self("0000000000000000/0000000000000000".into()) + } + + /// Encode a WAL position as an offset token. + pub fn new(commit_lsn: PgLsn, tx_ordinal: u64) -> Self { + Self(format!("{:016x}/{:016x}", u64::from(commit_lsn), tx_ordinal)) + } + + /// Decode the token back to `(commit_lsn, tx_ordinal)`. + pub fn decode(&self) -> Result<(PgLsn, u64)> { + let token = self.0.as_str(); + let (lsn_hex, ord_hex) = token + .split_once('/') + .ok_or_else(|| Error::Channel(format!("invalid offset token format: {token}")))?; + + let lsn = u64::from_str_radix(lsn_hex, 16) + .map_err(|e| Error::Channel(format!("invalid LSN hex in offset token: {e}")))?; + + let ord = u64::from_str_radix(ord_hex, 16) + .map_err(|e| Error::Channel(format!("invalid ordinal hex in offset token: {e}")))?; + + Ok((PgLsn::from(lsn), ord)) + } +} + +impl Default for OffsetToken { + fn default() -> Self { + Self::zero() + } +} + +impl fmt::Display for OffsetToken { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(&self.0) + } +} + +impl AsRef for OffsetToken { + fn as_ref(&self) -> &str { + &self.0 + } +} + +impl FromStr for OffsetToken { + type Err = Error; + + fn from_str(s: &str) -> Result { + let token = OffsetToken(s.to_owned()); + token.decode()?; + Ok(token) + } +} diff --git a/crates/etl-destinations/src/snowflake/streaming/rest_client.rs b/crates/etl-destinations/src/snowflake/streaming/rest_client.rs new file mode 100644 index 000000000..9ac8f3fa9 --- /dev/null +++ b/crates/etl-destinations/src/snowflake/streaming/rest_client.rs @@ -0,0 +1,593 @@ +use std::{string::String, sync::Arc, time::Duration}; + +use reqwest::{Client, StatusCode}; +use serde::{Deserialize, Serialize}; +use tokio::sync::OnceCell; +use tracing::{debug, warn}; + +use crate::{ + retry::{RetryDecision, RetryPolicy, retry_with_backoff}, + snowflake::{ + Error, Result, + auth::TokenProvider, + streaming::{ + ChannelStatusResponse, InsertRowsResponse, OffsetToken, OpenChannelResponse, RowBatch, + StreamClient, + }, + }, +}; + +const SNOWPIPE_RETRY_POLICY: RetryPolicy = RetryPolicy { + max_retries: 3, + initial_delay: Duration::from_millis(500), + max_delay: Duration::from_secs(10), +}; + +const USER_AGENT: &str = "supabase-etl/0.1.0"; + +/// [`StreamClient`] backed by the Snowpipe Streaming REST API. +/// +/// Discovers the ingest host on first use and caches it for the lifetime of the +/// client. +/// +/// All mutating calls (open/drop channel, insert rows, channel status) are +/// retried with exponential backoff. +pub struct RestStreamClient { + account_url: String, + auth: Arc, + http: Client, + ingest_host: OnceCell, +} + +impl RestStreamClient { + pub fn new(account_url: String, auth: Arc, http: Client) -> Self { + Self { account_url, auth, http, ingest_host: OnceCell::new() } + } + + async fn get_or_discover_host(&self) -> Result<&str> { + self.ingest_host + .get_or_try_init(|| async { + let token = self.auth.get_token().await?; + let url = format!("{}/v2/streaming/hostname", self.account_url); + let resp = self + .http + .get(&url) + .bearer_auth(&token) + .header("User-Agent", USER_AGENT) + .send() + .await + .map_err(Error::HttpTransport)?; + + let status = resp.status(); + if status != StatusCode::OK { + let body = resp.text().await.unwrap_or_default(); + return Err(Error::HttpStatus { status, body }); + } + + // Actual server returns plain text (even with Accept: application/json). + // Docs say JSON: https://docs.snowflake.com/en/user-guide/snowpipe-streaming/snowpipe-streaming-high-performance-rest-api#get-hostname + let body_text = resp.text().await.unwrap_or_default(); + let hostname = serde_json::from_str::(&body_text) + .map_or_else(|_| body_text.trim().to_owned(), |r| r.hostname); + + if hostname.is_empty() { + return Err(Error::Channel( + "hostname discovery returned empty hostname".into(), + )); + } + + debug!(hostname = %hostname, "discovered ingest host"); + let host = if hostname.starts_with("http://") || hostname.starts_with("https://") { + hostname + } else { + format!("https://{hostname}") + }; + + Ok(host) + }) + .await + .map(String::as_str) + } +} + +impl StreamClient for RestStreamClient { + async fn discover_ingest_host(&self) -> Result { + self.get_or_discover_host().await.map(ToOwned::to_owned) + } + + async fn open_channel( + &self, + database: &str, + schema: &str, + table: &str, + channel: &str, + ) -> Result { + let host = self.get_or_discover_host().await?; + let url = channel_url(host, database, schema, table, channel); + + let auth = Arc::clone(&self.auth); + let http = self.http.clone(); + + retry_with_backoff( + SNOWPIPE_RETRY_POLICY, + should_retry, + |d| d, + |attempt| { + warn!( + retry = attempt.retry_index, + max = attempt.max_retries, + delay_ms = attempt.sleep_delay.as_millis(), + error = %attempt.error, + "retrying open_channel" + ); + }, + || { + let url = url.clone(); + let auth = Arc::clone(&auth); + let http = http.clone(); + + async move { + let token = auth.get_token().await?; + let resp = http + .put(&url) + .bearer_auth(&token) + .header("User-Agent", USER_AGENT) + .header("Content-Type", "application/json") + .body("{}") + .send() + .await + .map_err(Error::HttpTransport)?; + + let status = resp.status(); + if status != StatusCode::OK { + let body = resp.text().await.unwrap_or_default(); + if status == StatusCode::UNAUTHORIZED { + warn!("received 401 from Snowpipe Streaming API, invalidating token"); + auth.invalidate_token().await; + } + return Err(Error::HttpStatus { status, body }); + } + + let response: OpenChannelApiResponse = resp.json().await.map_err(|e| { + Error::Encoding(format!("failed to parse open_channel response: {e}")) + })?; + + if let Some(ref status) = response.channel_status + && let Some(ref code) = status.channel_status_code + { + let is_ok = code == "SUCCESS" || code == "ACTIVE" || code == "0"; + if !is_ok { + let msg = format!("open_channel returned unexpected status: {code}"); + return Err(Error::Snowpipe { status_code: 1, message: msg }); + } + } + + Ok(OpenChannelResponse { + continuation_token: response.next_continuation_token, + offset_token: response + .channel_status + .and_then(|cs| cs.last_committed_offset_token) + .map(|s| s.parse::()) + .transpose()?, + }) + } + }, + ) + .await + .map_err(|f| f.last_error) + } + + async fn insert_rows( + &self, + database: &str, + schema: &str, + table: &str, + channel: &str, + batch: &RowBatch, + continuation_token: &str, + ) -> Result { + let host = self.get_or_discover_host().await?; + let base_url = insert_url(host, database, schema, table, channel); + + let compressed = batch.bytes().clone(); + let query_params = [ + ("continuationToken", continuation_token.to_owned()), + ("offsetToken", batch.offset().as_ref().to_owned()), + ]; + + let auth = Arc::clone(&self.auth); + let http = self.http.clone(); + + retry_with_backoff( + SNOWPIPE_RETRY_POLICY, + should_retry, + |d| d, + |attempt| { + if matches!(attempt.error, Error::Snowpipe { status_code: 3, .. }) { + debug!("auth error on insert_rows, token will be refreshed on retry"); + } + warn!( + retry = attempt.retry_index, + max = attempt.max_retries, + delay_ms = attempt.sleep_delay.as_millis(), + error = %attempt.error, + "retrying insert_rows" + ); + }, + || { + let base_url = base_url.clone(); + let query_params = query_params.clone(); + let auth = Arc::clone(&auth); + let http = http.clone(); + let compressed = compressed.clone(); + async move { + let token = auth.get_token().await?; + let resp = http + .post(&base_url) + .query(&query_params) + .bearer_auth(&token) + .header("User-Agent", USER_AGENT) + .header("Content-Type", "application/x-ndjson") + .header("Content-Encoding", "zstd") + .body(compressed) + .send() + .await + .map_err(Error::HttpTransport)?; + + let status = resp.status(); + if status != StatusCode::OK { + let body = resp.text().await.unwrap_or_default(); + if status == StatusCode::UNAUTHORIZED { + warn!("received 401 from Snowpipe Streaming API, invalidating token"); + auth.invalidate_token().await; + } + if let Ok(err_resp) = serde_json::from_str::(&body) + && let Some(code) = err_resp.status_code + { + if code == 3 { + auth.invalidate_token().await; + } + return Err(Error::Snowpipe { status_code: code, message: body }); + } + return Err(Error::HttpStatus { status, body }); + } + + let response: InsertRowsApiResponse = resp.json().await.map_err(|e| { + Error::Encoding(format!("failed to parse insert_rows response: {e}")) + })?; + + Ok(InsertRowsResponse { continuation_token: response.next_continuation_token }) + } + }, + ) + .await + .map_err(|f| f.last_error) + } + + async fn drop_channel( + &self, + database: &str, + schema: &str, + table: &str, + channel: &str, + ) -> Result<()> { + let host = self.get_or_discover_host().await?; + let url = channel_url(host, database, schema, table, channel); + + let auth = Arc::clone(&self.auth); + let http = self.http.clone(); + + retry_with_backoff( + SNOWPIPE_RETRY_POLICY, + should_retry, + |d| d, + |attempt| { + warn!( + retry = attempt.retry_index, + max = attempt.max_retries, + delay_ms = attempt.sleep_delay.as_millis(), + error = %attempt.error, + "retrying drop_channel" + ); + }, + || { + let url = url.clone(); + let auth = Arc::clone(&auth); + let http = http.clone(); + async move { + let token = auth.get_token().await?; + let resp = http + .delete(&url) + .bearer_auth(&token) + .header("User-Agent", USER_AGENT) + .send() + .await + .map_err(Error::HttpTransport)?; + + let status = resp.status(); + if status != StatusCode::OK { + let body = resp.text().await.unwrap_or_default(); + if status == StatusCode::UNAUTHORIZED { + warn!("received 401 from Snowpipe Streaming API, invalidating token"); + auth.invalidate_token().await; + } + return Err(Error::HttpStatus { status, body }); + } + Ok(()) + } + }, + ) + .await + .map_err(|f| f.last_error) + } + + async fn channel_status( + &self, + database: &str, + schema: &str, + table: &str, + channel: &str, + ) -> Result { + let host = self.get_or_discover_host().await?; + let url = channel_status_url(host, database, schema, table); + + let auth = Arc::clone(&self.auth); + let http = self.http.clone(); + let channel_names = vec![channel.to_owned()]; + let request_body = BulkStatusRequest { channel_names: &channel_names }; + + retry_with_backoff( + SNOWPIPE_RETRY_POLICY, + should_retry, + |d| d, + |attempt| { + warn!( + retry = attempt.retry_index, + max = attempt.max_retries, + delay_ms = attempt.sleep_delay.as_millis(), + error = %attempt.error, + "retrying channel_status" + ); + }, + || { + let url = url.clone(); + let auth = Arc::clone(&auth); + let http = http.clone(); + let body = &request_body; + async move { + let token = auth.get_token().await?; + let resp = http + .post(&url) + .bearer_auth(&token) + .header("User-Agent", USER_AGENT) + .json(body) + .send() + .await + .map_err(Error::HttpTransport)?; + + let status = resp.status(); + if status != StatusCode::OK { + let body = resp.text().await.unwrap_or_default(); + if status == StatusCode::UNAUTHORIZED { + warn!("received 401 from Snowpipe Streaming API, invalidating token"); + auth.invalidate_token().await; + } + return Err(Error::HttpStatus { status, body }); + } + + let response: BulkStatusApiResponse = resp.json().await.map_err(|e| { + Error::Encoding(format!("failed to parse channel_status response: {e}")) + })?; + + response.channel_statuses.into_iter().next().map_or_else( + || Err(Error::Channel("channel not found in status response".into())), + |(name, ch)| { + Ok(ChannelStatusResponse { + channel: name, + status_code: ch.channel_status_code.unwrap_or_default(), + offset_token: ch + .last_committed_offset_token + .map(|s| s.parse::()) + .transpose()?, + }) + }, + ) + } + }, + ) + .await + .map_err(|f| f.last_error) + } +} + +fn pipe_name(table: &str) -> String { + format!("{table}-STREAMING") +} + +fn channel_url(host: &str, db: &str, schema: &str, table: &str, channel: &str) -> String { + let pipe = pipe_name(table); + format!("{host}/v2/streaming/databases/{db}/schemas/{schema}/pipes/{pipe}/channels/{channel}") +} + +fn insert_url(host: &str, db: &str, schema: &str, table: &str, channel: &str) -> String { + let pipe = pipe_name(table); + format!( + "{host}/v2/streaming/data/databases/{db}/schemas/{schema}/pipes/{pipe}/channels/{channel}/\ + rows" + ) +} + +fn channel_status_url(host: &str, db: &str, schema: &str, table: &str) -> String { + let pipe = pipe_name(table); + format!("{host}/v2/streaming/databases/{db}/schemas/{schema}/pipes/{pipe}:bulk-channel-status") +} + +fn should_retry(error: &Error) -> RetryDecision { + match error { + Error::Snowpipe { status_code, .. } => match *status_code { + 0 => RetryDecision::Stop, + 1 | 5 | 6 => RetryDecision::Retry, + 3 => RetryDecision::Retry, + 2 | 4 => RetryDecision::Stop, + _ => RetryDecision::Retry, + }, + Error::HttpTransport(_) => RetryDecision::Retry, + Error::HttpStatus { status, .. } => { + if *status == StatusCode::UNAUTHORIZED + || *status == StatusCode::REQUEST_TIMEOUT + || *status == StatusCode::TOO_MANY_REQUESTS + || status.is_server_error() + { + RetryDecision::Retry + } else { + RetryDecision::Stop + } + } + _ => RetryDecision::Stop, + } +} + +#[derive(Deserialize)] +struct HostnameResponse { + hostname: String, +} + +#[derive(Deserialize)] +struct OpenChannelApiResponse { + next_continuation_token: String, + #[serde(default)] + channel_status: Option, +} + +#[derive(Deserialize)] +struct ChannelStatusDetail { + #[serde(default)] + channel_status_code: Option, + #[serde(default)] + last_committed_offset_token: Option, +} + +#[derive(Deserialize)] +struct InsertRowsApiResponse { + next_continuation_token: String, +} + +#[derive(Deserialize)] +struct SnowpipeErrorResponse { + #[serde(default)] + status_code: Option, +} + +#[derive(Serialize)] +struct BulkStatusRequest<'a> { + channel_names: &'a [String], +} + +#[derive(Deserialize)] +struct BulkStatusApiResponse { + #[serde(default)] + channel_statuses: std::collections::HashMap, +} + +#[derive(Deserialize)] +struct BulkStatusChannel { + #[serde(default)] + channel_status_code: Option, + #[serde(default)] + last_committed_offset_token: Option, +} + +#[cfg(test)] +mod tests { + use etl::types::{Cell, ColumnSchema, TableRow, Type}; + + use super::*; + use crate::snowflake::{ + encoding::{CdcMeta, CdcOperation, serialize_row}, + streaming::{OffsetToken, RowBatchBuilder}, + }; + + #[test] + fn should_retry_decision() { + let snowpipe = |code| Error::Snowpipe { status_code: code, message: "test".into() }; + + assert_eq!(should_retry(&snowpipe(0)), RetryDecision::Stop); + assert_eq!(should_retry(&snowpipe(1)), RetryDecision::Retry); + assert_eq!(should_retry(&snowpipe(2)), RetryDecision::Stop); + assert_eq!(should_retry(&snowpipe(3)), RetryDecision::Retry); + assert_eq!(should_retry(&snowpipe(4)), RetryDecision::Stop); + assert_eq!(should_retry(&snowpipe(5)), RetryDecision::Retry); + assert_eq!(should_retry(&snowpipe(6)), RetryDecision::Retry); + assert_eq!(should_retry(&snowpipe(99)), RetryDecision::Retry); + + let http = |status: StatusCode| Error::HttpStatus { status, body: "test".into() }; + + assert_eq!(should_retry(&http(StatusCode::INTERNAL_SERVER_ERROR)), RetryDecision::Retry); + assert_eq!(should_retry(&http(StatusCode::TOO_MANY_REQUESTS)), RetryDecision::Retry); + assert_eq!(should_retry(&http(StatusCode::REQUEST_TIMEOUT)), RetryDecision::Retry); + assert_eq!(should_retry(&http(StatusCode::UNAUTHORIZED)), RetryDecision::Retry); + assert_eq!(should_retry(&http(StatusCode::BAD_REQUEST)), RetryDecision::Stop); + + assert_eq!(should_retry(&Error::Auth("expired".into())), RetryDecision::Stop); + } + + #[test] + fn ndjson_formatting() { + let cols = [ + ColumnSchema::new("id".into(), Type::INT4, -1, 1, None, true), + ColumnSchema::new("name".into(), Type::TEXT, -1, 2, None, true), + ]; + + let mut buf = Vec::new(); + serialize_row( + &mut buf, + &cols, + &TableRow::new(vec![Cell::I32(1), Cell::String("Alice".into())]), + CdcMeta::new(CdcOperation::Insert, "0"), + ) + .unwrap(); + serialize_row( + &mut buf, + &cols, + &TableRow::new(vec![Cell::I32(2), Cell::String("Bob".into())]), + CdcMeta::new(CdcOperation::Insert, "0"), + ) + .unwrap(); + + let text = std::str::from_utf8(&buf).unwrap(); + let lines: Vec<&str> = text.trim_end().split('\n').collect(); + assert_eq!(lines.len(), 2); + + let row0: serde_json::Value = serde_json::from_str(lines[0]).unwrap(); + assert_eq!(row0["id"], 1); + assert_eq!(row0["name"], "Alice"); + + let row1: serde_json::Value = serde_json::from_str(lines[1]).unwrap(); + assert_eq!(row1["id"], 2); + assert_eq!(row1["name"], "Bob"); + } + + #[test] + fn compressed_roundtrip() { + let cols = [ColumnSchema::new("id".into(), Type::INT4, -1, 1, None, true)]; + let mut builder = RowBatchBuilder::new(); + builder + .push_row( + &cols, + &TableRow::new(vec![Cell::I32(42)]), + CdcMeta::new(CdcOperation::Insert, "0"), + &OffsetToken::zero(), + ) + .unwrap(); + + let batches = builder.finish().unwrap(); + let batch = batches.first().unwrap(); + assert!(batch.size() > 0); + + let decompressed = zstd::decode_all(batch.bytes().as_ref()).unwrap(); + let text = String::from_utf8(decompressed).unwrap(); + let line = text.trim(); + let val: serde_json::Value = serde_json::from_str(line).unwrap(); + assert_eq!(val["id"], 42); + } +} diff --git a/crates/etl-destinations/src/snowflake/test_utils.rs b/crates/etl-destinations/src/snowflake/test_utils.rs new file mode 100644 index 000000000..b152138d4 --- /dev/null +++ b/crates/etl-destinations/src/snowflake/test_utils.rs @@ -0,0 +1,58 @@ +use std::path::PathBuf; + +use crate::snowflake::{Config, Result, auth::TokenProvider, sql_client::SqlClient}; + +/// Snowflake account identifier (e.g. `org-account`). +const SNOWFLAKE_ACCOUNT_ENV: &str = "TESTS_SNOWFLAKE_ACCOUNT"; + +/// Snowflake login user name. +const SNOWFLAKE_USER_ENV: &str = "TESTS_SNOWFLAKE_USER"; + +/// Path to the PEM-encoded private key used for key-pair authentication. +const SNOWFLAKE_PRIVATE_KEY_PATH_ENV: &str = "TESTS_SNOWFLAKE_PRIVATE_KEY_PATH"; + +/// Target database. Falls back to `ETL_DEV` when unset. +const SNOWFLAKE_DATABASE_ENV: &str = "TESTS_SNOWFLAKE_DATABASE"; + +/// Target schema. Falls back to `PUBLIC` when unset. +const SNOWFLAKE_SCHEMA_ENV: &str = "TESTS_SNOWFLAKE_SCHEMA"; + +/// Optional role to assume after connecting. +const SNOWFLAKE_ROLE_ENV: &str = "TESTS_SNOWFLAKE_ROLE"; + +const DEFAULT_DATABASE: &str = "ETL_DEV"; +const DEFAULT_SCHEMA: &str = "PUBLIC"; + +/// Load a [`Config`] from environment variables for integration tests. +pub fn load_test_config() -> Config { + let account_id = std::env::var(SNOWFLAKE_ACCOUNT_ENV) + .unwrap_or_else(|_| panic!("{SNOWFLAKE_ACCOUNT_ENV} must be set")); + let username = std::env::var(SNOWFLAKE_USER_ENV) + .unwrap_or_else(|_| panic!("{SNOWFLAKE_USER_ENV} must be set")); + let database = + std::env::var(SNOWFLAKE_DATABASE_ENV).unwrap_or_else(|_| DEFAULT_DATABASE.to_owned()); + let schema = std::env::var(SNOWFLAKE_SCHEMA_ENV).unwrap_or_else(|_| DEFAULT_SCHEMA.to_owned()); + + let mut config = Config::new(&account_id, &username, &database, &schema); + if let Ok(role) = std::env::var(SNOWFLAKE_ROLE_ENV) { + config = config.with_role(&role); + } + config +} + +/// Load the private key path from `TESTS_SNOWFLAKE_PRIVATE_KEY_PATH`. +pub fn load_test_private_key_path() -> PathBuf { + PathBuf::from( + std::env::var(SNOWFLAKE_PRIVATE_KEY_PATH_ENV) + .unwrap_or_else(|_| panic!("{SNOWFLAKE_PRIVATE_KEY_PATH_ENV} must be set")), + ) +} + +/// Execute a SELECT query and return result rows. Requires an active warehouse. +pub async fn query_rows( + client: &SqlClient, + sql: &str, +) -> Result>> { + let resp = client.execute_statement(sql).await?; + Ok(resp.data.unwrap_or_default()) +} diff --git a/crates/etl-destinations/testdata/test_key.pem b/crates/etl-destinations/testdata/test_key.pem new file mode 100644 index 000000000..c7c5d931c --- /dev/null +++ b/crates/etl-destinations/testdata/test_key.pem @@ -0,0 +1,28 @@ +-----BEGIN PRIVATE KEY----- +MIIEvAIBADANBgkqhkiG9w0BAQEFAASCBKYwggSiAgEAAoIBAQCQdl7I8j34sxwh +XNUfVxikl5z1GH36iA5iC2j5CCgxIC/ff4YUiD8Zl1dKnhsfI6sAaGj3xwYnChsF +2PxbqJnocVXQaslhB9t0X2ryhmpZLYCTEPC+t0FZfN3A7mgNJIu83M9K2EF5zJw/ +1avxQKOhGML6GR8Su11h8CY/nqtJY6R8I9E0SrO9+rvJy/HCTFTmVQXTUhP7VHQI +m0tWU5PJjWdrk01ySHgGj+ipbE5yzPxpChbCsEogxvfwwuiIUq3FFNkwWf6pslYo +gOgTMcQfGE/kVFAXJ6D21V5qRqLCcdVDlCtsbRQuumhDukuEJQPY9Hpe76pU+9AE +54ywzVzLAgMBAAECggEAKaDCfXs4mmLad06t50MMydkalJIDM19jbaNGE4TjOAzD +Xs77jw0XycMPl2PqAtYfX5W2sbe7MSK7WLsHNU5nU1DdXpp2/yvpdCZOxiYvnRT8 +mORnyz05apUNvZu1hhwC3GBwp2ubqoJogAiNbI6o2DIvdSwqBRbSuPcfvnN1fkb2 +VXoqId6ZHI0I+mFxRl2sQqoJ783XzJtiKpuMaHgpgHg4MirpuJAvHK3IAo2SlCAx +ZqlDYAfxc+yoJ5LA5jskwVJgSb4eKK4Qg0hENy7XKGbIGPc5ZcZ0CqvlLIejYX/f +opwl3iAIaqJFsJCU2L5RAbwwMrbJVCUprBEH/fS42QKBgQDKL1BPSxCmCRim9iMB +8o/Z9MlaFM0r7TZcDu3UK3bx4/dmYwLwTSnbYFwrBhyQLLmRtpCj3YL3OLTPE4Po +0m/OVxBNSRoce2TFgP84zW0c82vYbtEPHxCl3o6oTsr9D3kdQUpwCqPonkL0tYBP +M18QreTkhqAsAbi9EWSnJ2WuZQKBgQC26ec8gVbcFy/zrWnaYeiOFGSwo9S25V3V +9YGczEMgsEZoyKWsysIjnKkK8RYdtc4nRd7U5jB9RR2xyZ5e3zeRtVwbtueAYQaH +yL+L2GZXVm3mcXf/YvAlagSkj3tgmopp+cA3VPakbob7W+5aRBG+TNbQmMnfwE+r +drG0KrxTbwKBgCGICUStSVyQA6OlDJdFGDBolYA4FPLlC2/nFfBrKPXi/ePgJdue +oIN+jqqf/9q1YC1XXtOeaBsCL5JsRSM2y04oSE7ZAdiZalQLGcjc4Oy9PjSN5GQi +ncs0hupI7wdbKpX8fxDn4tkwbiBRfa8k8O4+OMjhvuyteGr62HuUYBu1AoGAd+a7 +/Z2LIprQhBuY1952YyxbFK9QMNQJvsdAP1hmaHcksGtWrv36ZP28kb+Xj8ItcPEM +MOLzxioYXQKdHfOUqZ8I0eIDvtPbRAcECNfBvl6ZjAW1c2OXp+7nPDgR2DF1qiJd +Leg0BXWuZfbNN76HIwebiQGe011s3CjiNAgPi2ECgYA+t7ZccKddineyzMeV1kWk +RQaIJRvxXXwTmOp48AkKu0+RQzRgcuN9Y7uILjcmaFW+9J31KoLwVcs5mENP2YPh +9z4kvcpvOogJQj3ZG13v5Lco7y6LlqmXwjQKqZ73lV5FtTC4NrjdZOT8dmOX5skP +wabtc8blaEEBN+EqNpXBKg== +-----END PRIVATE KEY----- diff --git a/crates/etl-destinations/tests/main.rs b/crates/etl-destinations/tests/main.rs index f8472f729..eaf8f1756 100644 --- a/crates/etl-destinations/tests/main.rs +++ b/crates/etl-destinations/tests/main.rs @@ -8,3 +8,5 @@ mod clickhouse; mod ducklake; #[cfg(all(feature = "iceberg", feature = "test-utils"))] mod iceberg; +#[cfg(all(feature = "snowflake", feature = "test-utils"))] +mod snowflake; diff --git a/crates/etl-destinations/tests/snowflake/auth.rs b/crates/etl-destinations/tests/snowflake/auth.rs new file mode 100644 index 000000000..c8cc67484 --- /dev/null +++ b/crates/etl-destinations/tests/snowflake/auth.rs @@ -0,0 +1,20 @@ +use etl_destinations::snowflake::{ + AuthManager, TokenProvider, + test_utils::{load_test_config, load_test_private_key_path}, +}; + +#[tokio::test] +#[ignore = "requires Snowflake credentials — see etl-destinations/src/snowflake/README.md"] +async fn authenticate_against_snowflake() { + let config = load_test_config(); + let key_path = load_test_private_key_path(); + + let auth = AuthManager::new(&config, key_path.to_str().unwrap(), None) + .expect("AuthManager creation failed"); + + let token = auth.get_token().await.expect("authentication failed"); + assert!(!token.is_empty(), "token should not be empty"); + + let token2 = auth.get_token().await.expect("second get_token failed"); + assert_eq!(token, token2, "expected cached token on second call"); +} diff --git a/crates/etl-destinations/tests/snowflake/common.rs b/crates/etl-destinations/tests/snowflake/common.rs new file mode 100644 index 000000000..6f4aa35e2 --- /dev/null +++ b/crates/etl-destinations/tests/snowflake/common.rs @@ -0,0 +1,86 @@ +use std::sync::Arc; + +use etl_destinations::snowflake::{ + AuthManager, Config, Destination, HttpExchanger, OffsetToken, RestStreamClient, SqlClient, + StreamClient, + test_utils::{load_test_config, load_test_private_key_path}, +}; +use futures::FutureExt; + +pub fn build_auth() -> Arc> { + let config = load_test_config(); + let key_path = load_test_private_key_path(); + Arc::new( + AuthManager::new(&config, key_path.to_str().unwrap(), None) + .expect("AuthManager creation failed"), + ) +} + +pub async fn with_table_cleanup( + sql: &SqlClient>, + tables: &[&str], + test_fn: F, +) where + F: FnOnce() -> Fut, + Fut: std::future::Future, +{ + let result = std::panic::AssertUnwindSafe(test_fn()).catch_unwind().await; + + for table in tables { + let _ = sql.drop_table(table).await; + } + + if let Err(e) = result { + std::panic::resume_unwind(e); + } +} + +pub async fn poll_destination_offset( + destination: &Destination, + table_id: etl::types::TableId, + expected: &OffsetToken, + interval: std::time::Duration, + max_attempts: usize, +) -> Option +where + S: etl::store::state::StateStore + + etl::store::schema::SchemaStore + + Clone + + Send + + Sync + + 'static, + T: etl_destinations::snowflake::TokenProvider + 'static, + C: StreamClient, +{ + for _ in 0..max_attempts { + tokio::time::sleep(interval).await; + if let Ok(Some(offset)) = destination.committed_offset(table_id).await + && &offset == expected + { + return Some(offset); + } + } + None +} + +pub async fn poll_stream_offset( + stream: &RestStreamClient>, + config: &Config, + table: &str, + channel: &str, + expected: &OffsetToken, + interval: std::time::Duration, + max_attempts: usize, +) -> Option { + for _ in 0..max_attempts { + tokio::time::sleep(interval).await; + let status = stream + .channel_status(config.database(), config.schema(), table, channel) + .await + .expect("channel_status failed"); + if status.offset_token.as_ref() == Some(expected) { + return status.offset_token; + } + } + None +} diff --git a/crates/etl-destinations/tests/snowflake/destination.rs b/crates/etl-destinations/tests/snowflake/destination.rs new file mode 100644 index 000000000..e32b66183 --- /dev/null +++ b/crates/etl-destinations/tests/snowflake/destination.rs @@ -0,0 +1,910 @@ +use std::sync::Arc; + +use etl::{ + state::destination_metadata::DestinationTableMetadata, + store::{schema::SchemaStore, state::StateStore}, + test_utils::notifying_store::NotifyingStore, + types::{ + Cell, ColumnSchema, DeleteEvent, Event, InsertEvent, OldTableRow, PgLsn, PipelineId, + RelationEvent, ReplicatedTableSchema, SnapshotId, TableId, TableName, TableRow, + TableSchema, Type, UpdateEvent, UpdatedTableRow, + }, +}; +use etl_destinations::snowflake::{ + AuthManager, Client, Config, Destination, HttpExchanger, OffsetToken, RestStreamClient, + SqlClient, + test_utils::{load_test_config, query_rows}, +}; + +use super::common::{build_auth, poll_destination_offset, with_table_cleanup}; + +struct TestHarness { + destination: Destination< + NotifyingStore, + AuthManager, + RestStreamClient>, + >, + sql: SqlClient>, + config: Config, + store: NotifyingStore, +} + +impl TestHarness { + fn new() -> Self { + let config = load_test_config(); + let auth = build_auth(); + let sql = SqlClient::new(config.clone(), Arc::clone(&auth), reqwest::Client::new()); + let store = NotifyingStore::new(); + let pipeline_id: PipelineId = 1; + + let client = Client::new(config.clone(), Arc::clone(&auth), pipeline_id); + let destination = Destination::new(client, store.clone()); + + Self { destination, sql, config, store } + } +} + +fn snowflake_table_name(src_schema: &str, src_table: &str) -> String { + let escaped_schema = src_schema.replace('_', "__"); + let escaped_table = src_table.replace('_', "__"); + format!("{escaped_schema}_{escaped_table}").to_uppercase() +} + +async fn poll_and_query_rows( + harness: &TestHarness, + table_id: TableId, + sf_table: &str, + expected_offset: &OffsetToken, +) -> Vec> { + let committed = poll_destination_offset( + &harness.destination, + table_id, + expected_offset, + std::time::Duration::from_secs(5), + 18, + ) + .await; + assert_eq!( + committed, + Some(expected_offset.clone()), + "expected offset {expected_offset:?} not committed within 90s for table {sf_table}" + ); + + let fqn = + format!("\"{}\".\"{}\".\"{sf_table}\"", harness.config.database(), harness.config.schema()); + query_rows(&harness.sql, &format!("SELECT * FROM {fqn} ORDER BY \"_cdc_sequence_number\"")) + .await + .expect("query_rows failed") +} + +fn make_table_schema(table_id: u32, schema: &str, table: &str) -> TableSchema { + TableSchema::new( + TableId::new(table_id), + TableName::new(schema.to_owned(), table.to_owned()), + vec![ + ColumnSchema::new("id".to_owned(), Type::INT4, -1, 1, Some(1), false), + ColumnSchema::new("name".to_owned(), Type::TEXT, -1, 2, None, true), + ], + ) +} + +#[tokio::test] +#[ignore = "requires Snowflake credentials"] +async fn write_table_rows_basic() { + let harness = TestHarness::new(); + let src_table = format!("ETL_TEST_{}", uuid::Uuid::new_v4().simple()).to_uppercase(); + let sf_table = snowflake_table_name("public", &src_table); + + let table_id = TableId::new(1001); + let table_schema = make_table_schema(1001, "public", &src_table); + let schema = ReplicatedTableSchema::all(Arc::new(table_schema.clone())); + + harness.store.store_table_schema(table_schema).await.unwrap(); + + // Write 2 rows, poll, verify rows are there `_cdc_operation = "insert"`. + with_table_cleanup(&harness.sql, &[&sf_table], || async { + harness + .destination + .write_table_rows( + &schema, + vec![ + TableRow::new(vec![Cell::I32(1), Cell::String("Alice".into())]), + TableRow::new(vec![Cell::I32(2), Cell::String("Bob".into())]), + ], + ) + .await + .expect("write_table_rows failed"); + + let zero_offset = OffsetToken::zero(); + let rows = poll_and_query_rows(&harness, table_id, &sf_table, &zero_offset).await; + assert_eq!(rows.len(), 2, "expected 2 rows"); + + let zero_offset = zero_offset.to_string(); + // Column order: id, name, _cdc_operation, _cdc_sequence_number + assert_eq!(rows[0][0], serde_json::json!("1")); + assert_eq!(rows[0][1], serde_json::json!("Alice")); + assert_eq!(rows[0][2], serde_json::json!("insert")); + assert_eq!(rows[0][3], serde_json::json!(zero_offset)); + + assert_eq!(rows[1][0], serde_json::json!("2")); + assert_eq!(rows[1][1], serde_json::json!("Bob")); + assert_eq!(rows[1][2], serde_json::json!("insert")); + assert_eq!(rows[1][3], serde_json::json!(zero_offset)); + }) + .await; +} + +#[tokio::test] +#[ignore = "requires Snowflake credentials"] +async fn write_table_rows_empty() { + let harness = TestHarness::new(); + let src_table = format!("ETL_TEST_{}", uuid::Uuid::new_v4().simple()).to_uppercase(); + let sf_table = snowflake_table_name("public", &src_table); + + let table_schema = make_table_schema(1002, "public", &src_table); + let schema = ReplicatedTableSchema::all(Arc::new(table_schema.clone())); + + harness.store.store_table_schema(table_schema).await.unwrap(); + + // Write empty vec. Verify table is created but has no rows. + with_table_cleanup(&harness.sql, &[&sf_table], || async { + harness + .destination + .write_table_rows(&schema, vec![]) + .await + .expect("write_table_rows with empty rows failed"); + + let exists = harness.sql.table_exists(&sf_table).await.expect("table_exists failed"); + assert!(exists, "table should have been created even with empty row set"); + + let fqn = format!( + "\"{}\".\"{}\".\"{sf_table}\"", + harness.config.database(), + harness.config.schema() + ); + let rows = query_rows(&harness.sql, &format!("SELECT * FROM {fqn}")) + .await + .expect("query_rows failed"); + assert_eq!(rows.len(), 0, "table should be empty"); + }) + .await; +} + +#[tokio::test] +#[ignore = "requires Snowflake credentials"] +async fn write_events_insert_update_delete() { + let harness = TestHarness::new(); + let src_table = format!("ETL_TEST_{}", uuid::Uuid::new_v4().simple()).to_uppercase(); + let sf_table = snowflake_table_name("public", &src_table); + + let table_id = TableId::new(1003); + let table_schema = make_table_schema(1003, "public", &src_table); + let schema = ReplicatedTableSchema::all(Arc::new(table_schema.clone())); + + harness.store.store_table_schema(table_schema).await.unwrap(); + + // Send Insert, Update (Full), Delete (Full) events. + // Poll and verify 3 rows with operations "insert", "update", "delete". + with_table_cleanup(&harness.sql, &[&sf_table], || async { + harness + .destination + .process_events(vec![ + Event::Insert(InsertEvent { + start_lsn: PgLsn::from(1u64), + commit_lsn: PgLsn::from(1u64), + tx_ordinal: 0, + replicated_table_schema: schema.clone(), + table_row: TableRow::new(vec![Cell::I32(1), Cell::String("Alice".into())]), + }), + Event::Update(UpdateEvent { + start_lsn: PgLsn::from(2u64), + commit_lsn: PgLsn::from(2u64), + tx_ordinal: 0, + replicated_table_schema: schema.clone(), + updated_table_row: UpdatedTableRow::Full(TableRow::new(vec![ + Cell::I32(1), + Cell::String("Alice Updated".into()), + ])), + old_table_row: Some(OldTableRow::Full(TableRow::new(vec![ + Cell::I32(1), + Cell::String("Alice".into()), + ]))), + }), + Event::Delete(DeleteEvent { + start_lsn: PgLsn::from(3u64), + commit_lsn: PgLsn::from(3u64), + tx_ordinal: 0, + replicated_table_schema: schema.clone(), + old_table_row: Some(OldTableRow::Full(TableRow::new(vec![ + Cell::I32(2), + Cell::String("Bob".into()), + ]))), + }), + ]) + .await + .expect("process_events failed"); + + let expected_offset = OffsetToken::new(PgLsn::from(3u64), 0); + let rows = poll_and_query_rows(&harness, table_id, &sf_table, &expected_offset).await; + assert_eq!(rows.len(), 3, "expected 3 rows (insert + update + delete)"); + + // Column order: id, name, _cdc_operation, _cdc_sequence_number + // Rows ordered by _cdc_sequence_number. + assert_eq!(rows[0][0], serde_json::json!("1")); + assert_eq!(rows[0][1], serde_json::json!("Alice")); + assert_eq!(rows[0][2], serde_json::json!("insert")); + + assert_eq!(rows[1][0], serde_json::json!("1")); + assert_eq!(rows[1][1], serde_json::json!("Alice Updated")); + assert_eq!(rows[1][2], serde_json::json!("update")); + + assert_eq!(rows[2][0], serde_json::json!("2")); + assert_eq!(rows[2][1], serde_json::json!("Bob")); + assert_eq!(rows[2][2], serde_json::json!("delete")); + }) + .await; +} + +#[tokio::test] +#[ignore = "requires Snowflake credentials"] +async fn write_events_delete_key_only() { + let harness = TestHarness::new(); + let src_table = format!("ETL_TEST_{}", uuid::Uuid::new_v4().simple()).to_uppercase(); + let sf_table = snowflake_table_name("public", &src_table); + + let table_id = TableId::new(1004); + let table_schema = make_table_schema(1004, "public", &src_table); + let schema = ReplicatedTableSchema::all(Arc::new(table_schema.clone())); + + harness.store.store_table_schema(table_schema).await.unwrap(); + + // Delete event with `OldTableRow::Key`. + // Verify the delete row has PK value present but non-PK columns are NULL. + with_table_cleanup(&harness.sql, &[&sf_table], || async { + harness + .destination + .process_events(vec![Event::Delete(DeleteEvent { + start_lsn: PgLsn::from(1u64), + commit_lsn: PgLsn::from(1u64), + tx_ordinal: 0, + replicated_table_schema: schema.clone(), + old_table_row: Some(OldTableRow::Key(TableRow::new(vec![Cell::I32(42)]))), + })]) + .await + .expect("process_events failed"); + + let expected_offset = OffsetToken::new(PgLsn::from(1u64), 0); + let rows = poll_and_query_rows(&harness, table_id, &sf_table, &expected_offset).await; + assert_eq!(rows.len(), 1, "expected 1 delete row"); + + // Column order: id, name, _cdc_operation, _cdc_sequence_number + assert_eq!(&rows[0][0], &serde_json::Value::String("42".into()),); + assert_eq!(&rows[0][1], &serde_json::Value::Null); + assert_eq!(&rows[0][2], &serde_json::Value::String("delete".into()),); + }) + .await; +} + +#[tokio::test] +#[ignore = "requires Snowflake credentials"] +async fn schema_evolution_add_column() { + let harness = TestHarness::new(); + let src_table = format!("ETL_TEST_{}", uuid::Uuid::new_v4().simple()).to_uppercase(); + let sf_table = snowflake_table_name("public", &src_table); + + let table_id = TableId::new(1006); + let zero = OffsetToken::zero(); + + let initial_schema = TableSchema::new( + table_id, + TableName::new("public".to_owned(), src_table.clone()), + vec![ + ColumnSchema::new("id".to_owned(), Type::INT4, -1, 1, Some(1), false), + ColumnSchema::new("name".to_owned(), Type::TEXT, -1, 2, None, true), + ], + ); + let initial_replicated = ReplicatedTableSchema::all(Arc::new(initial_schema.clone())); + + let new_snapshot_id = SnapshotId::new(PgLsn::from(100u64)); + let evolved_schema = TableSchema::with_snapshot_id( + table_id, + TableName::new("public".to_owned(), src_table.clone()), + vec![ + ColumnSchema::new("id".to_owned(), Type::INT4, -1, 1, Some(1), false), + ColumnSchema::new("name".to_owned(), Type::TEXT, -1, 2, None, true), + ColumnSchema::new("email".to_owned(), Type::TEXT, -1, 3, None, true), + ], + new_snapshot_id, + ); + let evolved_replicated = ReplicatedTableSchema::all(Arc::new(evolved_schema.clone())); + + harness.store.store_table_schema(initial_schema.clone()).await.unwrap(); + harness.store.store_table_schema(evolved_schema.clone()).await.unwrap(); + + // Create table (write initial rows). + // Then add a column, then send a RelationEvent with the new schema. + // Finally, insert a row with the new column. + // Verify the new column exists in Snowflake. + with_table_cleanup(&harness.sql, &[&sf_table], || async { + harness + .destination + .write_table_rows( + &initial_replicated, + vec![TableRow::new(vec![Cell::I32(1), Cell::String("Alice".into())])], + ) + .await + .expect("initial write_table_rows failed"); + + // Wait for initial data to commit before DDL, channel refresh loses uncommitted + // rows. + let committed = poll_destination_offset( + &harness.destination, + table_id, + &zero, + std::time::Duration::from_secs(5), + 18, + ) + .await; + assert!(committed.is_some(), "initial data should commit before DDL"); + + let initial_metadata = DestinationTableMetadata::new_applied( + sf_table.clone(), + SnapshotId::initial(), + initial_replicated.replication_mask().clone(), + ); + harness.store.store_destination_table_metadata(table_id, initial_metadata).await.unwrap(); + + harness + .destination + .process_events(vec![Event::Relation(RelationEvent { + start_lsn: PgLsn::from(100u64), + commit_lsn: PgLsn::from(100u64), + tx_ordinal: 0, + replicated_table_schema: evolved_replicated.clone(), + })]) + .await + .expect("process_events (RelationEvent) failed"); + + harness + .destination + .process_events(vec![Event::Insert(InsertEvent { + start_lsn: PgLsn::from(101u64), + commit_lsn: PgLsn::from(101u64), + tx_ordinal: 0, + replicated_table_schema: evolved_replicated.clone(), + table_row: TableRow::new(vec![ + Cell::I32(2), + Cell::String("Bob".into()), + Cell::String("bob@example.com".into()), + ]), + })]) + .await + .expect("process_events (Insert with new column) failed"); + + let expected_offset = OffsetToken::new(PgLsn::from(101u64), 0); + let committed = poll_destination_offset( + &harness.destination, + table_id, + &expected_offset, + std::time::Duration::from_secs(5), + 18, + ) + .await; + assert_eq!(committed, Some(expected_offset), "data should commit within 90s"); + + let fqn = format!( + "\"{}\".\"{}\".\"{sf_table}\"", + harness.config.database(), + harness.config.schema() + ); + let rows = + query_rows(&harness.sql, &format!("SELECT \"email\" FROM {fqn} WHERE \"id\" = '2'")) + .await + .expect("query_rows for email column failed"); + + assert_eq!(rows.len(), 1, "expected one row for id=2"); + assert_eq!( + rows[0][0], + serde_json::Value::String("bob@example.com".into()), + "expected email = 'bob@example.com', got: {:?}", + rows[0][0] + ); + }) + .await; +} + +#[tokio::test] +#[ignore = "requires Snowflake credentials"] +async fn schema_evolution_rename_column() { + let harness = TestHarness::new(); + let src_table = format!("ETL_TEST_{}", uuid::Uuid::new_v4().simple()).to_uppercase(); + let sf_table = snowflake_table_name("public", &src_table); + + let table_id = TableId::new(1007); + let zero = OffsetToken::zero(); + + let initial_schema = TableSchema::new( + table_id, + TableName::new("public".to_owned(), src_table.clone()), + vec![ + ColumnSchema::new("id".to_owned(), Type::INT4, -1, 1, Some(1), false), + ColumnSchema::new("name".to_owned(), Type::TEXT, -1, 2, None, true), + ], + ); + let initial_replicated = ReplicatedTableSchema::all(Arc::new(initial_schema.clone())); + + let new_snapshot_id = SnapshotId::new(PgLsn::from(100u64)); + let evolved_schema = TableSchema::with_snapshot_id( + table_id, + TableName::new("public".to_owned(), src_table.clone()), + vec![ + ColumnSchema::new("id".to_owned(), Type::INT4, -1, 1, Some(1), false), + ColumnSchema::new("full_name".to_owned(), Type::TEXT, -1, 2, None, true), + ], + new_snapshot_id, + ); + let evolved_replicated = ReplicatedTableSchema::all(Arc::new(evolved_schema.clone())); + + harness.store.store_table_schema(initial_schema).await.unwrap(); + harness.store.store_table_schema(evolved_schema).await.unwrap(); + + with_table_cleanup(&harness.sql, &[&sf_table], || async { + harness + .destination + .write_table_rows( + &initial_replicated, + vec![TableRow::new(vec![Cell::I32(1), Cell::String("Alice".into())])], + ) + .await + .expect("initial write_table_rows failed"); + + // Wait for initial data to commit before DDL -- channel refresh loses + // uncommitted rows. + let committed = poll_destination_offset( + &harness.destination, + table_id, + &zero, + std::time::Duration::from_secs(5), + 18, + ) + .await; + assert!(committed.is_some(), "initial data should commit before DDL"); + + let initial_metadata = DestinationTableMetadata::new_applied( + sf_table.clone(), + SnapshotId::initial(), + initial_replicated.replication_mask().clone(), + ); + harness.store.store_destination_table_metadata(table_id, initial_metadata).await.unwrap(); + + harness + .destination + .process_events(vec![Event::Relation(RelationEvent { + start_lsn: PgLsn::from(100u64), + commit_lsn: PgLsn::from(100u64), + tx_ordinal: 0, + replicated_table_schema: evolved_replicated.clone(), + })]) + .await + .expect("process_events (RelationEvent rename) failed"); + + harness + .destination + .process_events(vec![Event::Insert(InsertEvent { + start_lsn: PgLsn::from(101u64), + commit_lsn: PgLsn::from(101u64), + tx_ordinal: 0, + replicated_table_schema: evolved_replicated.clone(), + table_row: TableRow::new(vec![Cell::I32(2), Cell::String("Bob".into())]), + })]) + .await + .expect("process_events (Insert with renamed column) failed"); + + let expected_offset = OffsetToken::new(PgLsn::from(101u64), 0); + let committed = poll_destination_offset( + &harness.destination, + table_id, + &expected_offset, + std::time::Duration::from_secs(5), + 18, + ) + .await; + assert_eq!(committed, Some(expected_offset), "data should commit within 90s"); + + let fqn = format!( + "\"{}\".\"{}\".\"{sf_table}\"", + harness.config.database(), + harness.config.schema() + ); + + // New row uses the renamed column. + let rows = query_rows( + &harness.sql, + &format!("SELECT \"full_name\" FROM {fqn} WHERE \"id\" = '2'"), + ) + .await + .expect("query for renamed column failed"); + assert_eq!(rows.len(), 1); + assert_eq!(rows[0][0], serde_json::json!("Bob")); + + // Initial row data is preserved under the new column name. + let rows = query_rows( + &harness.sql, + &format!("SELECT \"full_name\" FROM {fqn} WHERE \"id\" = '1'"), + ) + .await + .expect("query for initial row after rename failed"); + assert_eq!(rows.len(), 1); + assert_eq!(rows[0][0], serde_json::json!("Alice")); + }) + .await; +} + +#[tokio::test] +#[ignore = "requires Snowflake credentials"] +async fn schema_evolution_drop_column() { + let harness = TestHarness::new(); + let src_table = format!("ETL_TEST_{}", uuid::Uuid::new_v4().simple()).to_uppercase(); + let sf_table = snowflake_table_name("public", &src_table); + + let table_id = TableId::new(1008); + let zero = OffsetToken::zero(); + + let initial_schema = TableSchema::new( + table_id, + TableName::new("public".to_owned(), src_table.clone()), + vec![ + ColumnSchema::new("id".to_owned(), Type::INT4, -1, 1, Some(1), false), + ColumnSchema::new("name".to_owned(), Type::TEXT, -1, 2, None, true), + ColumnSchema::new("email".to_owned(), Type::TEXT, -1, 3, None, true), + ], + ); + let initial_replicated = ReplicatedTableSchema::all(Arc::new(initial_schema.clone())); + + let new_snapshot_id = SnapshotId::new(PgLsn::from(100u64)); + let evolved_schema = TableSchema::with_snapshot_id( + table_id, + TableName::new("public".to_owned(), src_table.clone()), + vec![ + ColumnSchema::new("id".to_owned(), Type::INT4, -1, 1, Some(1), false), + ColumnSchema::new("name".to_owned(), Type::TEXT, -1, 2, None, true), + ], + new_snapshot_id, + ); + let evolved_replicated = ReplicatedTableSchema::all(Arc::new(evolved_schema.clone())); + + harness.store.store_table_schema(initial_schema).await.unwrap(); + harness.store.store_table_schema(evolved_schema).await.unwrap(); + + with_table_cleanup(&harness.sql, &[&sf_table], || async { + harness + .destination + .write_table_rows( + &initial_replicated, + vec![TableRow::new(vec![ + Cell::I32(1), + Cell::String("Alice".into()), + Cell::String("alice@test.com".into()), + ])], + ) + .await + .expect("initial write_table_rows failed"); + + // Wait for initial data to commit before DDL -- channel refresh loses + // uncommitted rows. + let committed = poll_destination_offset( + &harness.destination, + table_id, + &zero, + std::time::Duration::from_secs(5), + 18, + ) + .await; + assert!(committed.is_some(), "initial data should commit before DDL"); + + let initial_metadata = DestinationTableMetadata::new_applied( + sf_table.clone(), + SnapshotId::initial(), + initial_replicated.replication_mask().clone(), + ); + harness.store.store_destination_table_metadata(table_id, initial_metadata).await.unwrap(); + + harness + .destination + .process_events(vec![Event::Relation(RelationEvent { + start_lsn: PgLsn::from(100u64), + commit_lsn: PgLsn::from(100u64), + tx_ordinal: 0, + replicated_table_schema: evolved_replicated.clone(), + })]) + .await + .expect("process_events (RelationEvent drop) failed"); + + harness + .destination + .process_events(vec![Event::Insert(InsertEvent { + start_lsn: PgLsn::from(101u64), + commit_lsn: PgLsn::from(101u64), + tx_ordinal: 0, + replicated_table_schema: evolved_replicated.clone(), + table_row: TableRow::new(vec![Cell::I32(2), Cell::String("Bob".into())]), + })]) + .await + .expect("process_events (Insert after column drop) failed"); + + let expected_offset = OffsetToken::new(PgLsn::from(101u64), 0); + let committed = poll_destination_offset( + &harness.destination, + table_id, + &expected_offset, + std::time::Duration::from_secs(5), + 18, + ) + .await; + assert_eq!(committed, Some(expected_offset), "data should commit within 90s"); + + let fqn = format!( + "\"{}\".\"{}\".\"{sf_table}\"", + harness.config.database(), + harness.config.schema() + ); + + // New row landed with remaining columns. + let rows = + query_rows(&harness.sql, &format!("SELECT \"name\" FROM {fqn} WHERE \"id\" = '2'")) + .await + .expect("query after drop failed"); + assert_eq!(rows.len(), 1); + assert_eq!(rows[0][0], serde_json::json!("Bob")); + + // Dropped column is gone from the table. + let result = query_rows(&harness.sql, &format!("SELECT \"email\" FROM {fqn}")).await; + assert!(result.is_err(), "column 'email' should not exist after DROP COLUMN"); + }) + .await; +} + +#[tokio::test] +#[ignore = "requires Snowflake credentials"] +async fn schema_evolution_interleaved_ddl_dml() { + let harness = TestHarness::new(); + let src_table = format!("ETL_TEST_{}", uuid::Uuid::new_v4().simple()).to_uppercase(); + let sf_table = snowflake_table_name("public", &src_table); + + let table_id = TableId::new(1009); + let table_name = TableName::new("public".to_owned(), src_table.clone()); + + // v1: (id, name) + let schema_v1 = TableSchema::new( + table_id, + table_name.clone(), + vec![ + ColumnSchema::new("id".to_owned(), Type::INT4, -1, 1, Some(1), false), + ColumnSchema::new("name".to_owned(), Type::TEXT, -1, 2, None, true), + ], + ); + let replicated_v1 = ReplicatedTableSchema::all(Arc::new(schema_v1.clone())); + + // v2: ADD COLUMN email (ordinal 3) + let schema_v2 = TableSchema::with_snapshot_id( + table_id, + table_name.clone(), + vec![ + ColumnSchema::new("id".to_owned(), Type::INT4, -1, 1, Some(1), false), + ColumnSchema::new("name".to_owned(), Type::TEXT, -1, 2, None, true), + ColumnSchema::new("email".to_owned(), Type::TEXT, -1, 3, None, true), + ], + SnapshotId::new(PgLsn::from(100u64)), + ); + let replicated_v2 = ReplicatedTableSchema::all(Arc::new(schema_v2.clone())); + + // v3: RENAME name -> full_name (ordinal 2 unchanged) + let schema_v3 = TableSchema::with_snapshot_id( + table_id, + table_name.clone(), + vec![ + ColumnSchema::new("id".to_owned(), Type::INT4, -1, 1, Some(1), false), + ColumnSchema::new("full_name".to_owned(), Type::TEXT, -1, 2, None, true), + ColumnSchema::new("email".to_owned(), Type::TEXT, -1, 3, None, true), + ], + SnapshotId::new(PgLsn::from(200u64)), + ); + let replicated_v3 = ReplicatedTableSchema::all(Arc::new(schema_v3.clone())); + + // v4: DROP COLUMN email (ordinal 3 removed) + let schema_v4 = TableSchema::with_snapshot_id( + table_id, + table_name.clone(), + vec![ + ColumnSchema::new("id".to_owned(), Type::INT4, -1, 1, Some(1), false), + ColumnSchema::new("full_name".to_owned(), Type::TEXT, -1, 2, None, true), + ], + SnapshotId::new(PgLsn::from(300u64)), + ); + let replicated_v4 = ReplicatedTableSchema::all(Arc::new(schema_v4.clone())); + + harness.store.store_table_schema(schema_v1).await.unwrap(); + harness.store.store_table_schema(schema_v2).await.unwrap(); + harness.store.store_table_schema(schema_v3).await.unwrap(); + harness.store.store_table_schema(schema_v4).await.unwrap(); + + let poll = |offset_lsn: u64, ord: u64| { + let expected = OffsetToken::new(PgLsn::from(offset_lsn), ord); + async { + let committed = poll_destination_offset( + &harness.destination, + table_id, + &expected, + std::time::Duration::from_secs(5), + 18, + ) + .await; + assert_eq!(committed, Some(expected), "data should commit before next DDL"); + } + }; + + with_table_cleanup(&harness.sql, &[&sf_table], || async { + // Initial table copy with v1 schema. + harness + .destination + .write_table_rows( + &replicated_v1, + vec![TableRow::new(vec![Cell::I32(1), Cell::String("Alice".into())])], + ) + .await + .expect("initial write_table_rows failed"); + + // Wait for initial data to commit before DDL. + let zero = OffsetToken::zero(); + let committed = poll_destination_offset( + &harness.destination, + table_id, + &zero, + std::time::Duration::from_secs(5), + 18, + ) + .await; + assert!(committed.is_some(), "initial data should commit before DDL"); + + let initial_metadata = DestinationTableMetadata::new_applied( + sf_table.clone(), + SnapshotId::initial(), + replicated_v1.replication_mask().clone(), + ); + harness.store.store_destination_table_metadata(table_id, initial_metadata).await.unwrap(); + + // Phase 1: Insert(v1) + ADD COLUMN + harness + .destination + .process_events(vec![Event::Insert(InsertEvent { + start_lsn: PgLsn::from(1u64), + commit_lsn: PgLsn::from(1u64), + tx_ordinal: 0, + replicated_table_schema: replicated_v1.clone(), + table_row: TableRow::new(vec![Cell::I32(2), Cell::String("Bob".into())]), + })]) + .await + .expect("process_events (Insert v1) failed"); + poll(1, 0).await; + + harness + .destination + .process_events(vec![Event::Relation(RelationEvent { + start_lsn: PgLsn::from(100u64), + commit_lsn: PgLsn::from(100u64), + tx_ordinal: 0, + replicated_table_schema: replicated_v2.clone(), + })]) + .await + .expect("process_events (Relation v2) failed"); + + // Phase 2: Insert(v2) + RENAME + harness + .destination + .process_events(vec![Event::Insert(InsertEvent { + start_lsn: PgLsn::from(101u64), + commit_lsn: PgLsn::from(101u64), + tx_ordinal: 0, + replicated_table_schema: replicated_v2.clone(), + table_row: TableRow::new(vec![ + Cell::I32(3), + Cell::String("Charlie".into()), + Cell::String("charlie@test.com".into()), + ]), + })]) + .await + .expect("process_events (Insert v2) failed"); + poll(101, 0).await; + + harness + .destination + .process_events(vec![Event::Relation(RelationEvent { + start_lsn: PgLsn::from(200u64), + commit_lsn: PgLsn::from(200u64), + tx_ordinal: 0, + replicated_table_schema: replicated_v3.clone(), + })]) + .await + .expect("process_events (Relation v3) failed"); + + // Phase 3: Insert(v3) + DROP COLUMN + harness + .destination + .process_events(vec![Event::Insert(InsertEvent { + start_lsn: PgLsn::from(201u64), + commit_lsn: PgLsn::from(201u64), + tx_ordinal: 0, + replicated_table_schema: replicated_v3.clone(), + table_row: TableRow::new(vec![ + Cell::I32(4), + Cell::String("Diana".into()), + Cell::String("diana@test.com".into()), + ]), + })]) + .await + .expect("process_events (Insert v3) failed"); + poll(201, 0).await; + + harness + .destination + .process_events(vec![Event::Relation(RelationEvent { + start_lsn: PgLsn::from(300u64), + commit_lsn: PgLsn::from(300u64), + tx_ordinal: 0, + replicated_table_schema: replicated_v4.clone(), + })]) + .await + .expect("process_events (Relation v4) failed"); + + // Phase 4: Final insert after all DDL. + harness + .destination + .process_events(vec![Event::Insert(InsertEvent { + start_lsn: PgLsn::from(301u64), + commit_lsn: PgLsn::from(301u64), + tx_ordinal: 0, + replicated_table_schema: replicated_v4.clone(), + table_row: TableRow::new(vec![Cell::I32(5), Cell::String("Eve".into())]), + })]) + .await + .expect("process_events (Insert v4) failed"); + poll(301, 0).await; + + let fqn = format!( + "\"{}\".\"{}\".\"{sf_table}\"", + harness.config.database(), + harness.config.schema() + ); + + // Final schema should be (id, full_name) -- email was dropped, name was + // renamed. + let rows = query_rows( + &harness.sql, + &format!("SELECT \"id\", \"full_name\" FROM {fqn} ORDER BY \"id\""), + ) + .await + .expect("final query failed"); + + // 5 rows: 1 from initial copy + 4 from CDC inserts. + assert_eq!(rows.len(), 5, "expected 5 rows, got {}", rows.len()); + assert_eq!(rows[0][0], serde_json::json!("1")); + assert_eq!(rows[0][1], serde_json::json!("Alice")); + assert_eq!(rows[1][0], serde_json::json!("2")); + assert_eq!(rows[1][1], serde_json::json!("Bob")); + assert_eq!(rows[2][0], serde_json::json!("3")); + assert_eq!(rows[2][1], serde_json::json!("Charlie")); + assert_eq!(rows[3][0], serde_json::json!("4")); + assert_eq!(rows[3][1], serde_json::json!("Diana")); + assert_eq!(rows[4][0], serde_json::json!("5")); + assert_eq!(rows[4][1], serde_json::json!("Eve")); + + // Dropped column should not exist. + let result = query_rows(&harness.sql, &format!("SELECT \"email\" FROM {fqn}")).await; + assert!(result.is_err(), "column 'email' should not exist after DROP COLUMN"); + + // Old column name should not exist. + let result = query_rows(&harness.sql, &format!("SELECT \"name\" FROM {fqn}")).await; + assert!(result.is_err(), "column 'name' should not exist after RENAME"); + }) + .await; +} diff --git a/crates/etl-destinations/tests/snowflake/mod.rs b/crates/etl-destinations/tests/snowflake/mod.rs new file mode 100644 index 000000000..2f30cab34 --- /dev/null +++ b/crates/etl-destinations/tests/snowflake/mod.rs @@ -0,0 +1,5 @@ +mod auth; +mod common; +mod destination; +mod sql_client; +mod stream_client; diff --git a/crates/etl-destinations/tests/snowflake/sql_client.rs b/crates/etl-destinations/tests/snowflake/sql_client.rs new file mode 100644 index 000000000..86df9949e --- /dev/null +++ b/crates/etl-destinations/tests/snowflake/sql_client.rs @@ -0,0 +1,98 @@ +use etl_destinations::snowflake::{ + AuthManager, HttpExchanger, SqlClient, test_utils::load_test_config, +}; + +use super::common::{build_auth, with_table_cleanup}; + +fn build_sql_client() -> SqlClient> { + let config = load_test_config(); + let auth = build_auth(); + SqlClient::new(config, auth, reqwest::Client::new()) +} + +#[tokio::test] +#[ignore = "requires Snowflake credentials — see etl-destinations/src/snowflake/README.md"] +async fn ddl_lifecycle() { + let client = build_sql_client(); + let table = format!("etl_test_{}", uuid::Uuid::new_v4().simple()); + + with_table_cleanup(&client, &[&table], || async { + client + .create_table_if_not_exists(&table, r#""id" NUMBER(10,0), "name" VARCHAR"#) + .await + .expect("create table failed"); + + assert!( + client.table_exists(&table).await.expect("table_exists failed"), + "table should exist after creation" + ); + + client.truncate_table(&table).await.expect("truncate failed"); + + assert!( + client.table_exists(&table).await.expect("table_exists failed"), + "table should still exist after truncate" + ); + + client.drop_table(&table).await.expect("drop table failed"); + + assert!( + !client.table_exists(&table).await.expect("table_exists failed"), + "table should not exist after drop" + ); + }) + .await; +} + +#[tokio::test] +#[ignore = "requires Snowflake credentials — see etl-destinations/src/snowflake/README.md"] +async fn table_exists_returns_false_for_nonexistent() { + let client = build_sql_client(); + let table = format!("etl_nonexistent_{}", uuid::Uuid::new_v4().simple()); + + assert!( + !client.table_exists(&table).await.expect("table_exists failed"), + "table that was never created should not exist" + ); +} + +#[tokio::test] +#[ignore = "requires Snowflake credentials — see etl-destinations/src/snowflake/README.md"] +async fn create_table_idempotent() { + let client = build_sql_client(); + let table = format!("etl_test_{}", uuid::Uuid::new_v4().simple()); + + with_table_cleanup(&client, &[&table], || async { + client + .create_table_if_not_exists(&table, r#""id" NUMBER(10,0)"#) + .await + .expect("first create failed"); + + client + .create_table_if_not_exists(&table, r#""id" NUMBER(10,0)"#) + .await + .expect("second create should succeed silently"); + }) + .await; +} + +#[tokio::test] +#[ignore = "requires Snowflake credentials — see etl-destinations/src/snowflake/README.md"] +async fn schema_evolution_ddl() { + let client = build_sql_client(); + let table = format!("etl_test_{}", uuid::Uuid::new_v4().simple()); + + with_table_cleanup(&client, &[&table], || async { + client + .create_table_if_not_exists(&table, r#""id" NUMBER(10,0), "name" VARCHAR"#) + .await + .expect("create table failed"); + + client.add_column(&table, "email", "VARCHAR").await.expect("add column failed"); + + client.rename_column(&table, "email", "user_email").await.expect("rename column failed"); + + client.drop_column(&table, "user_email").await.expect("drop column failed"); + }) + .await; +} diff --git a/crates/etl-destinations/tests/snowflake/stream_client.rs b/crates/etl-destinations/tests/snowflake/stream_client.rs new file mode 100644 index 000000000..558d3f668 --- /dev/null +++ b/crates/etl-destinations/tests/snowflake/stream_client.rs @@ -0,0 +1,335 @@ +use std::sync::Arc; + +use etl::types::{Cell, ColumnSchema, TableRow, Type}; +use etl_destinations::snowflake::{ + AuthManager, CdcMeta, CdcOperation, Config, HttpExchanger, OffsetToken, RestStreamClient, + RowBatch, RowBatchBuilder, SqlClient, StreamClient, + test_utils::{load_test_config, query_rows}, +}; +use tokio::time::Duration; + +use super::common::{build_auth, poll_stream_offset, with_table_cleanup}; + +fn build_clients( + config: &Config, +) -> (RestStreamClient>, SqlClient>) { + let auth = build_auth(); + let stream = RestStreamClient::new( + config.account_url().to_owned(), + Arc::clone(&auth), + reqwest::Client::new(), + ); + let sql = SqlClient::new(config.clone(), Arc::clone(&auth), reqwest::Client::new()); + (stream, sql) +} + +fn build_batch(cols: &[ColumnSchema], rows: &[TableRow], offset: &OffsetToken) -> RowBatch { + let mut builder = RowBatchBuilder::new(); + for row in rows { + builder.push_row(cols, row, CdcMeta::new(CdcOperation::Insert, "0"), offset).unwrap(); + } + builder.finish().unwrap().into_iter().next().unwrap() +} + +#[tokio::test] +#[ignore = "requires Snowflake credentials"] +async fn channel_open_insert_status_drop() { + let config = load_test_config(); + let (stream, sql) = build_clients(&config); + + let table = format!("ETL_TEST_{}", uuid::Uuid::new_v4().simple()).to_uppercase(); + let channel = format!("etl_test_{}_ch0", uuid::Uuid::new_v4().simple()); + + with_table_cleanup(&sql, &[&table], || async { + // Create test table. + sql.create_table_if_not_exists(&table, r#""id" NUMBER(10,0), "name" VARCHAR"#) + .await + .expect("create table failed"); + + // Open channel, no offset on fresh channel. + let resp = stream + .open_channel(config.database(), config.schema(), &table, &channel) + .await + .expect("open_channel failed"); + assert!(!resp.continuation_token.is_empty(), "expected non-empty continuation_token"); + assert!(resp.offset_token.is_none(), "unexpected offset on fresh channel"); + + // Insert rows. + let cols = [ + ColumnSchema::new("id".into(), Type::INT4, -1, 1, None, true), + ColumnSchema::new("name".into(), Type::TEXT, -1, 2, None, true), + ]; + let offset: OffsetToken = "0000000000000001/0000000000000001".parse().unwrap(); + let batch = build_batch( + &cols, + &[ + TableRow::new(vec![Cell::I32(1), Cell::String("Alice".into())]), + TableRow::new(vec![Cell::I32(2), Cell::String("Bob".into())]), + ], + &offset, + ); + let resp = stream + .insert_rows( + config.database(), + config.schema(), + &table, + &channel, + &batch, + &resp.continuation_token, + ) + .await + .expect("insert_rows failed"); + assert!( + !resp.continuation_token.is_empty(), + "expected non-empty continuation_token after insert" + ); + + // Poll status until offset is committed. + let committed = poll_stream_offset( + &stream, + &config, + &table, + &channel, + &offset, + std::time::Duration::from_secs(5), + 18, + ) + .await; + assert_eq!(committed, Some(offset), "committed offset must match inserted offset"); + + // Drop channel. + stream + .drop_channel(config.database(), config.schema(), &table, &channel) + .await + .expect("drop_channel failed"); + }) + .await; +} + +#[tokio::test] +#[ignore = "requires Snowflake credentials"] +async fn channel_reopen_preserves_offset() { + let config = load_test_config(); + let (stream, sql) = build_clients(&config); + + let table = format!("ETL_TEST_{}", uuid::Uuid::new_v4().simple()).to_uppercase(); + let channel = format!("etl_test_{}_ch0", uuid::Uuid::new_v4().simple()); + + with_table_cleanup(&sql, &[&table], || async { + sql.create_table_if_not_exists(&table, r#""id" NUMBER(10,0), "name" VARCHAR"#) + .await + .expect("create table failed"); + + // Open and insert. + let open_resp = stream + .open_channel(config.database(), config.schema(), &table, &channel) + .await + .expect("open_channel failed"); + + let cols = [ + ColumnSchema::new("id".into(), Type::INT4, -1, 1, None, true), + ColumnSchema::new("name".into(), Type::TEXT, -1, 2, None, true), + ]; + let offset: OffsetToken = "0000000000000001/0000000000000000".parse().unwrap(); + let batch = build_batch( + &cols, + &[TableRow::new(vec![Cell::I32(1), Cell::String("Alice".into())])], + &offset, + ); + stream + .insert_rows( + config.database(), + config.schema(), + &table, + &channel, + &batch, + &open_resp.continuation_token, + ) + .await + .expect("insert_rows failed"); + + // Poll until offset is committed. + let committed = poll_stream_offset( + &stream, + &config, + &table, + &channel, + &offset, + std::time::Duration::from_secs(5), + 18, + ) + .await; + assert_eq!(committed, Some(offset.clone()), "committed offset must match inserted offset"); + + // Reopen channel, idempotent, should return the committed offset. + let reopen_resp = stream + .open_channel(config.database(), config.schema(), &table, &channel) + .await + .expect("reopen channel failed"); + assert_eq!( + reopen_resp.offset_token, + Some(offset), + "reopened channel must return the committed offset" + ); + + let _ = stream.drop_channel(config.database(), config.schema(), &table, &channel).await; + }) + .await; +} + +#[tokio::test] +#[ignore = "requires Snowflake credentials"] +async fn continuation_token() { + let config = load_test_config(); + let (stream, sql) = build_clients(&config); + + let table = format!("ETL_TEST_{}", uuid::Uuid::new_v4().simple()).to_uppercase(); + let channel = format!("etl_test_{}_ch0", uuid::Uuid::new_v4().simple()); + + with_table_cleanup(&sql, &[&table], || async { + sql.create_table_if_not_exists(&table, r#""id" NUMBER(10,0), "name" VARCHAR"#) + .await + .expect("create table failed"); + + let fqn = format!("\"{}\".\"{}\".\"{table}\"", config.database(), config.schema()); + + #[allow(clippy::too_many_arguments)] + async fn verify( + stream: &RestStreamClient>, + sql: &SqlClient>, + config: &Config, + table: &str, + channel: &str, + fqn: &str, + expected_offset: &OffsetToken, + expected_rows: usize, + ) { + let committed = poll_stream_offset( + stream, + config, + table, + channel, + expected_offset, + Duration::from_secs(5), + 18, + ) + .await; + assert_eq!( + committed, + Some(expected_offset.clone()), + "expected offset not committed within timeout" + ); + let rows = query_rows(sql, &format!("SELECT * FROM {fqn} ORDER BY \"id\"")) + .await + .expect("query_rows failed"); + assert_eq!(rows.len(), expected_rows, "unexpected row count: {rows:?}"); + } + + let resp = stream + .open_channel(config.database(), config.schema(), &table, &channel) + .await + .expect("open_channel failed"); + + let cols = [ + ColumnSchema::new("id".into(), Type::INT4, -1, 1, None, true), + ColumnSchema::new("name".into(), Type::TEXT, -1, 2, None, true), + ]; + + // Batch 1 + let offset1: OffsetToken = "0000000000000001/0000000000000000".parse().unwrap(); + let batch1 = build_batch( + &cols, + &[TableRow::new(vec![Cell::I32(1), Cell::String("Alice".into())])], + &offset1, + ); + let insert1 = stream + .insert_rows( + config.database(), + config.schema(), + &table, + &channel, + &batch1, + &resp.continuation_token, + ) + .await + .expect("insert batch 1 failed"); + + verify(&stream, &sql, &config, &table, &channel, &fqn, &offset1, 1).await; + + // Batch 2: uses continuation_token from batch 1. + let offset2: OffsetToken = "0000000000000001/0000000000000001".parse().unwrap(); + let batch2 = build_batch( + &cols, + &[TableRow::new(vec![Cell::I32(2), Cell::String("Bob".into())])], + &offset2, + ); + let insert2 = stream + .insert_rows( + config.database(), + config.schema(), + &table, + &channel, + &batch2, + &insert1.continuation_token, + ) + .await + .expect("insert batch 2 failed"); + + assert_ne!( + insert1.continuation_token, insert2.continuation_token, + "continuation_token must advance after each batch" + ); + // 2 rows expected, offset2 must be committed. + verify(&stream, &sql, &config, &table, &channel, &fqn, &offset2, 2).await; + + // Batch 3: use STALE token from batch 1 (already consumed by batch 2). + let offset3: OffsetToken = "0000000000000001/0000000000000002".parse().unwrap(); + let batch3 = build_batch( + &cols, + &[TableRow::new(vec![Cell::I32(3), Cell::String("Charlie".into())])], + &offset3, + ); + let insert3 = stream + .insert_rows( + config.database(), + config.schema(), + &table, + &channel, + &batch3, + &insert1.continuation_token, + ) + .await + .expect("insert with stale token should not error"); + + // Same continuation token returned. + assert_eq!( + insert3.continuation_token, insert2.continuation_token, + "stale token should not advance the sequencer" + ); + // Still 2 rows, offset token is not advanced to offset3. + verify(&stream, &sql, &config, &table, &channel, &fqn, &offset2, 2).await; + + // Retry batch 3 with the CORRECT token, data commits, offset is updated. + let insert3_retry = stream + .insert_rows( + config.database(), + config.schema(), + &table, + &channel, + &batch3, + &insert2.continuation_token, + ) + .await + .expect("insert batch 3 with correct token failed"); + + assert_ne!( + insert3_retry.continuation_token, insert2.continuation_token, + "correct token should advance the sequencer" + ); + verify(&stream, &sql, &config, &table, &channel, &fqn, &offset3, 3).await; + + // Cleanup + let _ = stream.drop_channel(config.database(), config.schema(), &table, &channel).await; + }) + .await; +} diff --git a/crates/etl-examples/Cargo.toml b/crates/etl-examples/Cargo.toml index 0dd742f0d..0b6663d14 100644 --- a/crates/etl-examples/Cargo.toml +++ b/crates/etl-examples/Cargo.toml @@ -26,24 +26,59 @@ path = "src/bin/ducklake.rs" required-features = ["ducklake"] test = false +[[bin]] +name = "snowflake" +path = "src/bin/snowflake/main.rs" +required-features = ["snowflake"] +test = false + +[[bin]] +name = "snowflake-loadgen" +path = "src/bin/snowflake_loadgen.rs" +required-features = ["snowflake"] +test = false + [features] default = [] bigquery = ["etl-destinations/bigquery"] clickhouse = ["etl-destinations/clickhouse"] ducklake = ["etl-destinations/ducklake"] +snowflake = [ + "etl-destinations/snowflake", + "dep:crossterm", + "dep:ratatui", + "dep:rand", + "dep:reqwest", + "dep:secrecy", + "dep:tokio-postgres", +] [dependencies] -clap = { workspace = true, default-features = true, features = ["std", "derive"] } +clap = { workspace = true, default-features = true, features = [ + "std", + "derive", + "env", +] } +crossterm = { version = "0.29", optional = true } etl = { workspace = true } etl-config = { workspace = true } etl-destinations = { workspace = true } etl-telemetry = { workspace = true } k8s-openapi = { workspace = true, features = ["latest"] } +rand = { workspace = true, optional = true, features = ["std", "std_rng"] } +ratatui = { version = "0.30", optional = true, default-features = false, features = [ + "crossterm", +] } +reqwest = { workspace = true, optional = true, features = ["json"] } rustls = { workspace = true, features = ["aws-lc-rs", "logging"] } -tokio = { workspace = true, features = ["macros", "signal"] } +secrecy = { workspace = true, optional = true } +tokio = { workspace = true, features = ["macros", "signal", "time"] } +tokio-postgres = { workspace = true, optional = true } tracing = { workspace = true, default-features = true } -tracing-subscriber = { workspace = true, default-features = true, features = ["env-filter"] } +tracing-subscriber = { workspace = true, default-features = true, features = [ + "env-filter", +] } url = { workspace = true } [lints] diff --git a/crates/etl-examples/README.md b/crates/etl-examples/README.md index a8afb2112..cc98bfaee 100644 --- a/crates/etl-examples/README.md +++ b/crates/etl-examples/README.md @@ -10,6 +10,7 @@ Postgres to various destinations using the ETL pipeline. | [BigQuery](#bigquery) | `bigquery` | `bigquery` | Google BigQuery (cloud data warehouse) | Stable | | [ClickHouse](#clickhouse-setup) | `clickhouse` | `clickhouse` | ClickHouse (column-oriented OLAP database) | In progress | | [DuckLake](#ducklake) | `ducklake` | `ducklake` | DuckLake (open data lake format) | In progress | +| [Snowflake](src/bin/snowflake/README.md) | `snowflake` | `snowflake` | Snowflake (cloud data warehouse) | In progress | ## Building and running @@ -25,7 +26,7 @@ cargo build --bin bigquery -p etl-examples --features bigquery cargo run --bin bigquery -p etl-examples --features bigquery -- [flags] ``` -Replace `bigquery` with `clickhouse` or `ducklake` as needed. +Replace `bigquery` with `clickhouse`, `ducklake`, or `snowflake` as needed. ### All examples diff --git a/crates/etl-examples/src/bin/snowflake/README.md b/crates/etl-examples/src/bin/snowflake/README.md new file mode 100644 index 000000000..6b183e3da --- /dev/null +++ b/crates/etl-examples/src/bin/snowflake/README.md @@ -0,0 +1,196 @@ +# Snowflake CDC Benchmark + +Streams 1M+ rows from Postgres to Snowflake via Snowpipe Streaming, then measures CDC throughput. + +Demonstrates the Snowflake destination's performance characteristics. + +## Prerequisites + +1. The project's dev stack running (provides source Postgres with logical replication) +2. Snowflake account with: + - RSA key pair configured for key-pair auth + - Target database and schema created (`ETL_BENCH.CDC`) + - A role with appropriate permissions (USAGE on warehouse/DB/schema, CREATE TABLE, etc.) +3. Rust toolchain (cargo) + +## Quick Start + +### 1. Start the dev stack + +From the repo root: + +```bash +./scripts/init.sh +``` + +This starts `source-postgres` on port 5430 with `wal_level=logical` and all replication settings configured. See `scripts/docker-compose.yaml` for details. + +### 2. Create a benchmark database + +```bash +psql -h localhost -p 5430 -U postgres -c "CREATE DATABASE etl_bench;" +``` + +### 3. Seed 1M rows + +```bash +cargo run -p etl-examples --features snowflake --bin snowflake-loadgen -- seed \ + --db-url postgres://postgres:postgres@localhost:5430/etl_bench \ + --rows 1000000 +``` + +Takes ~2 minutes. Creates 3 tables (users, orders, events) in a dedicated `bench` Postgres schema and scopes the publication to it, so the ETL pipeline's internal state tables (in `public`) are not replicated. + +### 4. Configure Snowflake credentials + +Fill in the `BENCH_SNOWFLAKE_*` variables in `.env` (see `.env.example`): + +```env +BENCH_SNOWFLAKE_ACCOUNT=ORG-ACCOUNT +BENCH_SNOWFLAKE_USER=ETL_USER +BENCH_SNOWFLAKE_PRIVATE_KEY_PATH=/path/to/rsa_key.p8 +BENCH_SNOWFLAKE_DATABASE=ETL_BENCH +BENCH_SNOWFLAKE_SCHEMA=CDC +BENCH_SNOWFLAKE_ROLE=ETL_ROLE +``` + +Then load them: `source .env` + +### 5. Start the pipeline + +Snowflake args are picked up from `BENCH_SNOWFLAKE_*` env vars automatically: + +```bash +cargo run -p etl-examples --features snowflake --bin snowflake -- \ + --db-host localhost \ + --db-port 5430 \ + --db-name etl_bench \ + --db-username postgres \ + --db-password postgres \ + --publication bench_pub +``` + +The terminal dashboard shows table copy progress and throughput in real-time. + +### 6. Start CDC load generator (separate terminal) + +```bash +cargo run -p etl-examples --features snowflake --bin snowflake-loadgen -- generate \ + --db-url postgres://postgres:postgres@localhost:5430/etl_bench \ + --rate 5000 \ + --mix 40/40/20 \ + --duration 300s +``` + +## What to Expect + +- **Table copy phase**: Copies all 1M rows to Snowflake. Expect 10,000-20,000 rows/sec depending on network and row size. +- **CDC phase**: After initial copy, streams changes in real-time. Latency is typically 2-10 seconds from Postgres commit to queryable in Snowflake. + +## Verifying in Snowflake + +Table names in Snowflake follow the pattern `{pg_schema}_{pg_table}` uppercased, so Postgres `bench.events` becomes `ETL_BENCH.CDC.BENCH_EVENTS`. All examples below use fully qualified names (`ETL_BENCH.CDC.*`) matching the default config. + +**Important**: Column names are preserved as lowercase (quoted identifiers). You must double-quote them in Snowflake SQL, e.g. `"id"` not `id`. + +```sql +-- List all tables +SHOW TABLES IN ETL_BENCH.CDC; + +-- Count total rows (includes all CDC versions) +SELECT COUNT(*) FROM ETL_BENCH.CDC.BENCH_EVENTS; + +-- Check CDC metadata columns +SELECT "id", "event_type", "_cdc_operation", "_cdc_sequence_number" +FROM ETL_BENCH.CDC.BENCH_EVENTS +ORDER BY "_cdc_sequence_number" DESC +LIMIT 20; + +-- Materialize current state (latest version of each row) +SELECT * FROM ( + SELECT *, ROW_NUMBER() OVER ( + PARTITION BY "id" + ORDER BY "_cdc_sequence_number" DESC + ) AS rn + FROM ETL_BENCH.CDC.BENCH_EVENTS +) +WHERE rn = 1 AND "_cdc_operation" != 'delete'; + +-- Count current (live) rows only +SELECT COUNT(*) FROM ( + SELECT "id", "_cdc_operation", ROW_NUMBER() OVER ( + PARTITION BY "id" + ORDER BY "_cdc_sequence_number" DESC + ) AS rn + FROM ETL_BENCH.CDC.BENCH_EVENTS +) +WHERE rn = 1 AND "_cdc_operation" != 'delete'; + +-- Create a Dynamic Table for auto-refreshing current state +CREATE DYNAMIC TABLE ETL_BENCH.CDC.BENCH_EVENTS_CURRENT + TARGET_LAG = '1 minute' + WAREHOUSE = MY_WH +AS + SELECT * EXCLUDE ("_cdc_operation", "_cdc_sequence_number", rn) FROM ( + SELECT *, ROW_NUMBER() OVER ( + PARTITION BY "id" + ORDER BY "_cdc_sequence_number" DESC + ) AS rn + FROM ETL_BENCH.CDC.BENCH_EVENTS + ) + WHERE rn = 1 AND "_cdc_operation" != 'delete'; +``` + +## Verifying in Postgres (source comparison) + +Run these against the source database to compare row counts with Snowflake: + +```bash +psql -h localhost -p 5430 -U postgres -d etl_bench +``` + +```sql +-- Total rows in source +SELECT COUNT(*) FROM bench.events; + +-- If running the CDC load generator, compare live row counts: +-- This should match the "current state" count from Snowflake above. +SELECT COUNT(*) FROM bench.events; + +-- All tables: +SELECT + schemaname, + relname AS table_name, + n_live_tup AS approx_rows +FROM pg_stat_user_tables +WHERE schemaname = 'bench' +ORDER BY relname; +``` + +## Cleanup + +```bash +# Stop the pipeline (Ctrl+C) +# Drop the benchmark database +psql -h localhost -p 5430 -U postgres -c "DROP DATABASE etl_bench;" +``` + +In Snowflake: + +```sql +DROP TABLE IF EXISTS ETL_BENCH.CDC.BENCH_EVENTS; +DROP TABLE IF EXISTS ETL_BENCH.CDC.BENCH_USERS; +DROP TABLE IF EXISTS ETL_BENCH.CDC.BENCH_ORDERS; +DROP DYNAMIC TABLE IF EXISTS ETL_BENCH.CDC.BENCH_EVENTS_CURRENT; +``` + +## Comparison with Snowflake's CDC Guide + +This benchmark is modeled after Snowflake's official CDC guide (https://www.snowflake.com/en/developers/guides/cdc-snowpipestreaming-dynamictables/). + +Key differences: + +- Pure Rust pipeline (no Java SDK sidecar) +- Uses the Snowpipe Streaming REST API directly +- Supports all Postgres types, schema evolution, and automatic recovery +- Append-only changelog model with Dynamic Tables for materialization diff --git a/crates/etl-examples/src/bin/snowflake/commands.rs b/crates/etl-examples/src/bin/snowflake/commands.rs new file mode 100644 index 000000000..046ce7ef8 --- /dev/null +++ b/crates/etl-examples/src/bin/snowflake/commands.rs @@ -0,0 +1,109 @@ +use std::sync::{Arc, Mutex}; + +use etl::{ + state::TableReplicationPhase, + store::{PostgresStore, StateStore}, + types::TableId, +}; +use etl_destinations::snowflake; +use tracing::{error, info}; + +use crate::state::DashboardState; + +/// Reset a single table: truncate Snowflake data, reset state to Init. +/// The pipeline must be stopped before calling this. +pub async fn reset_table( + table_id: TableId, + table_name: &str, + store: &PostgresStore, + destination: &snowflake::Destination, + dashboard: &Arc>, +) { + info!(table = %table_name, "resetting table: truncating Snowflake data and resetting state"); + + if let Err(e) = destination.committed_offset(table_id).await { + info!(table = %table_name, "table has no open channel ({}), skipping truncate", e); + } else { + // Table has an open channel — try to truncate via the client. + // We access the underlying client through the destination's public API. + // Since we can't truncate through the Destination trait without a + // ReplicatedTableSchema, we reset state and let the pipeline re-copy on + // next start. + info!(table = %table_name, "channel exists, data will be re-copied on restart"); + } + + match store.update_table_replication_state(table_id, TableReplicationPhase::Init).await { + Ok(()) => { + info!(table = %table_name, "table state reset to Init"); + dashboard.lock().unwrap().set_status(format!( + "{table_name}: state reset to Init, restart pipeline to re-copy" + )); + } + Err(e) => { + error!(table = %table_name, error = %e, "failed to reset table state"); + dashboard.lock().unwrap().set_status(format!("{table_name}: reset failed: {e}")); + } + } +} + +/// Reset ALL tables to Init state. +#[allow(dead_code)] +pub async fn reset_all_tables( + store: &PostgresStore, + destination: &snowflake::Destination, + dashboard: &Arc>, +) { + let tables: Vec<(TableId, String)> = { + let dash = dashboard.lock().unwrap(); + dash.tables.iter().map(|t| (t.table_id, t.destination_name.clone())).collect() + }; + + for (id, name) in &tables { + reset_table(*id, name, store, destination, dashboard).await; + } + + dashboard + .lock() + .unwrap() + .set_status("All tables reset. Restart pipeline to re-copy.".to_owned()); +} + +pub const HELP_TEXT: &str = "\ +Keyboard Shortcuts: + j/k or Up/Down Navigate table list + J/K Scroll logs up/down (3 lines) + PgUp/PgDn Scroll logs (page) + F1 / ? Toggle this help + F2 Reset selected table + F3 Restart pipeline + F4 Cycle log filter (All / Warn+ / Error) + F5 Force-sync Snowflake offsets now + F10 / q Quit + +Table States: + init Table registered, not yet copied + copying Initial data copy in progress + copied Copy finished, waiting for sync + sync_wait Waiting for apply worker pause + catchup Catching up to sync LSN + sync_done Sync complete + ready Fully synced, streaming CDC + errored Error occurred (see detail panel) + +Commands: + Reset (F2) Stops the pipeline, resets the selected table's state + to Init, then restarts. The pipeline will re-copy all + data from Postgres. Use when a table is stuck in error. + + Restart (F3) Stops and restarts the pipeline without changing any + table states. Tables resume from their current state. + + Log (F4) Cycles the log filter: All shows everything, Warn+ + shows only warnings and errors, Error shows errors only. + + Sync (F5) Queries Snowflake's channel status API to fetch the + last committed offset for each table. Offsets also + auto-refresh every 30 seconds. + +Press any key to close this help. +"; diff --git a/crates/etl-examples/src/bin/snowflake/logging.rs b/crates/etl-examples/src/bin/snowflake/logging.rs new file mode 100644 index 000000000..0a64492e8 --- /dev/null +++ b/crates/etl-examples/src/bin/snowflake/logging.rs @@ -0,0 +1,86 @@ +use std::{ + collections::VecDeque, + sync::{Arc, Mutex}, +}; + +use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt}; + +pub const MAX_LOG_LINES: usize = 2000; + +pub struct TuiLogLayer { + buffer: Arc>>, +} + +struct FieldVisitor { + message: String, + fields: Vec<(String, String)>, +} + +impl FieldVisitor { + fn new() -> Self { + Self { message: String::new(), fields: Vec::new() } + } + + fn into_line(self) -> String { + if self.fields.is_empty() { + self.message + } else { + let extras: Vec = + self.fields.into_iter().map(|(k, v)| format!("{k}={v}")).collect(); + format!("{} {}", self.message, extras.join(" ")) + } + } +} + +impl tracing::field::Visit for FieldVisitor { + fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) { + if field.name() == "message" { + self.message = format!("{value:?}"); + } else { + self.fields.push((field.name().to_owned(), format!("{value:?}"))); + } + } + + fn record_str(&mut self, field: &tracing::field::Field, value: &str) { + if field.name() == "message" { + self.message = value.to_owned(); + } else { + self.fields.push((field.name().to_owned(), value.to_owned())); + } + } +} + +impl tracing_subscriber::Layer for TuiLogLayer { + fn on_event( + &self, + event: &tracing::Event<'_>, + _ctx: tracing_subscriber::layer::Context<'_, S>, + ) { + let level = event.metadata().level(); + let target = event.metadata().target(); + let mut visitor = FieldVisitor::new(); + event.record(&mut visitor); + let line = format!("[{level}] {target}: {}", visitor.into_line()); + let mut buf = self.buffer.lock().unwrap(); + buf.push_back(line); + if buf.len() > MAX_LOG_LINES { + buf.pop_front(); + } + } +} + +pub fn init_tracing(log_buffer: Arc>>) { + if std::env::var("RUST_LOG").is_err() { + unsafe { + std::env::set_var("RUST_LOG", "info"); + } + } + + tracing_subscriber::registry() + .with( + tracing_subscriber::EnvFilter::try_from_default_env() + .unwrap_or_else(|_| "snowflake=info".into()), + ) + .with(TuiLogLayer { buffer: log_buffer }) + .init(); +} diff --git a/crates/etl-examples/src/bin/snowflake/main.rs b/crates/etl-examples/src/bin/snowflake/main.rs new file mode 100644 index 000000000..1a1810ac4 --- /dev/null +++ b/crates/etl-examples/src/bin/snowflake/main.rs @@ -0,0 +1,461 @@ +/// Snowflake CDC Pipeline — Control Plane TUI +/// +/// Streams Postgres CDC to Snowflake via Snowpipe Streaming, with a live +/// terminal dashboard showing per-table replication progress, offsets, and +/// throughput. +/// +/// Supports resetting errored tables and restarting the pipeline. +/// +/// See README.md for the full guide. +mod commands; +mod logging; +mod state; +mod tui; + +use std::{ + collections::{BTreeMap, VecDeque}, + error::Error, + sync::{Arc, Mutex, Once}, + time::{Duration, Instant}, +}; + +use clap::{Args, Parser}; +use crossterm::event::{Event, KeyCode}; +use etl::{ + config::{ + BatchConfig, InvalidatedSlotBehavior, MemoryBackpressureConfig, PgConnectionConfig, + PipelineConfig, TableSyncCopyConfig, TcpKeepaliveConfig, TlsConfig, + }, + pipeline::Pipeline, + store::PostgresStore, + types::TableId, +}; +use etl_destinations::snowflake::{AuthManager, Client, Config, Destination}; +use etl_telemetry::metrics::init_metrics_handle; +use secrecy::SecretString; +use tracing::{error, info}; + +use crate::state::{DashboardState, GlobalPhase}; + +static INIT_CRYPTO: Once = Once::new(); + +fn install_crypto_provider() { + INIT_CRYPTO.call_once(|| { + rustls::crypto::aws_lc_rs::default_provider() + .install_default() + .expect("failed to install default crypto provider"); + }); +} + +#[derive(Debug, Parser)] +#[command(name = "snowflake", version, about, arg_required_else_help = true)] +struct AppArgs { + #[clap(flatten)] + db_args: DbArgs, + #[clap(flatten)] + sf_args: SnowflakeArgs, + #[arg(long)] + publication: String, + #[arg(long, default_value = "5000")] + max_batch_fill_duration_ms: u64, + #[arg(long, default_value = "4")] + max_table_sync_workers: u16, +} + +#[derive(Debug, Args)] +struct DbArgs { + #[arg(long)] + db_host: String, + #[arg(long)] + db_port: u16, + #[arg(long)] + db_name: String, + #[arg(long)] + db_username: String, + #[arg(long)] + db_password: Option, +} + +#[derive(Debug, Clone, Args)] +struct SnowflakeArgs { + #[arg(long, env = "BENCH_SNOWFLAKE_ACCOUNT")] + snowflake_account: String, + #[arg(long, env = "BENCH_SNOWFLAKE_USER")] + snowflake_user: String, + #[arg(long, env = "BENCH_SNOWFLAKE_PRIVATE_KEY_PATH")] + snowflake_private_key_path: String, + #[arg(long, env = "BENCH_SNOWFLAKE_PRIVATE_KEY_PASSPHRASE")] + snowflake_private_key_passphrase: Option, + #[arg(long, env = "BENCH_SNOWFLAKE_DATABASE")] + snowflake_database: String, + #[arg(long, env = "BENCH_SNOWFLAKE_SCHEMA", default_value = "CDC")] + snowflake_schema: String, + #[arg(long, env = "BENCH_SNOWFLAKE_ROLE")] + snowflake_role: Option, +} + +#[tokio::main] +async fn main() -> Result<(), Box> { + if let Err(err) = run().await { + error!(error = %err, "fatal error"); + std::process::exit(1); + } + Ok(()) +} + +async fn run() -> Result<(), Box> { + let log_buffer: Arc>> = Arc::new(Mutex::new(VecDeque::new())); + logging::init_tracing(Arc::clone(&log_buffer)); + install_crypto_provider(); + + let args = AppArgs::parse(); + + let pg_config = PgConnectionConfig { + host: args.db_args.db_host, + hostaddr: None, + port: args.db_args.db_port, + name: args.db_args.db_name, + username: args.db_args.db_username, + password: args.db_args.db_password.map(Into::into), + tls: TlsConfig { trusted_root_certs: String::new(), enabled: false }, + keepalive: TcpKeepaliveConfig::default(), + }; + + info!("counting source tables and rows..."); + let (table_count, estimated_rows) = count_source_tables(&pg_config).await?; + info!(tables = table_count, estimated_rows, "source database stats"); + + let pipeline_id = 1u64; + + let pipeline_config = PipelineConfig { + id: pipeline_id, + publication_name: args.publication.clone(), + pg_connection: pg_config.clone(), + batch: BatchConfig { + max_fill_ms: args.max_batch_fill_duration_ms, + memory_budget_ratio: 0.2, + max_bytes: 8 * 1024 * 1024, + }, + table_error_retry_delay_ms: 10000, + table_error_retry_max_attempts: 5, + max_table_sync_workers: args.max_table_sync_workers, + memory_refresh_interval_ms: 100, + memory_backpressure: Some(MemoryBackpressureConfig::default()), + table_sync_copy: TableSyncCopyConfig::default(), + invalidated_slot_behavior: InvalidatedSlotBehavior::Recreate, + max_copy_connections_per_table: PipelineConfig::DEFAULT_MAX_COPY_CONNECTIONS_PER_TABLE, + }; + + let args = args.sf_args.clone(); + let passphrase: Option = + args.snowflake_private_key_passphrase.map(SecretString::from); + + let mut config = Config::new( + &args.snowflake_account, + &args.snowflake_user, + &args.snowflake_database, + &args.snowflake_schema, + ); + if let Some(ref role) = args.snowflake_role { + config = config.with_role(role); + } + + let auth = + Arc::new(AuthManager::new(&config, &args.snowflake_private_key_path, passphrase.as_ref())?); + + let store = PostgresStore::new(pipeline_id, pg_config.clone()).await?; + + let dashboard = Arc::new(Mutex::new(DashboardState::new(table_count, estimated_rows))); + + // Build initial table name map from destination metadata + let mut table_names = state::build_table_name_map(&store).await; + + // Start pipeline + let client = Client::new(config.clone(), Arc::clone(&auth), pipeline_id); + let mut destination = Destination::new(client, store.clone()); + let mut pipeline = Pipeline::new(pipeline_config.clone(), store.clone(), destination.clone()); + + info!("starting Snowflake CDC pipeline..."); + pipeline.start().await?; + info!("pipeline started, press F1 for help"); + + let metrics_handle = init_metrics_handle().ok(); + + let (mut terminal, _terminal_guard) = tui::setup_terminal()?; + + // Monitor state + let mut samples: Vec = Vec::new(); + let mut last_rows: u64 = 0; + let mut last_tick = Instant::now(); + let mut last_refresh = Instant::now(); + let mut last_offset_fetch = Instant::now(); + let mut phase_times: BTreeMap = BTreeMap::new(); + let start_time = Instant::now(); + let page_size = 10usize; + + loop { + // Render + { + let state = dashboard.lock().unwrap(); + let lb = Arc::clone(&log_buffer); + terminal.draw(|f| tui::render(f, &state, &lb))?; + } + + // Periodic refresh (every 2s) + if last_refresh.elapsed() >= Duration::from_secs(2) { + let metrics_text = match metrics_handle.as_ref() { + Some(h) => h.render(), + None => String::new(), + }; + + // Refresh table name map (picks up newly discovered tables) + let new_names = state::build_table_name_map(&store).await; + for (id, name) in new_names { + table_names.entry(id).or_insert(name); + } + + state::refresh_dashboard( + &dashboard, + &store, + &metrics_text, + start_time, + &mut samples, + &mut last_rows, + &mut last_tick, + &table_names, + &mut phase_times, + ) + .await; + + last_refresh = Instant::now(); + } + + // Periodic Snowflake offset sync (every 30s) + if last_offset_fetch.elapsed() >= Duration::from_secs(30) { + state::fetch_snowflake_offsets(&dashboard, &destination).await; + last_offset_fetch = Instant::now(); + } + + // Handle input + if crossterm::event::poll(Duration::from_millis(100))? + && let Event::Key(key) = crossterm::event::read()? + { + // If help is open, any key closes it + { + let mut dash = dashboard.lock().unwrap(); + if dash.show_help { + dash.show_help = false; + continue; + } + } + + match key.code { + // Quit + KeyCode::F(10) | KeyCode::Char('q') => { + info!("quit requested, shutting down..."); + pipeline.shutdown(); + break; + } + + // Help + KeyCode::F(1) | KeyCode::Char('?') => { + dashboard.lock().unwrap().show_help = true; + } + + // Table navigation + KeyCode::Char('j') | KeyCode::Down => { + dashboard.lock().unwrap().select_next(); + } + KeyCode::Char('k') | KeyCode::Up => { + dashboard.lock().unwrap().select_prev(); + } + + // Log scrolling + KeyCode::Char('J') => { + let mut dash = dashboard.lock().unwrap(); + dash.log_scroll = dash.log_scroll.saturating_add(3); + } + KeyCode::Char('K') => { + let mut dash = dashboard.lock().unwrap(); + dash.log_scroll = dash.log_scroll.saturating_sub(3); + } + KeyCode::PageUp => { + let mut dash = dashboard.lock().unwrap(); + dash.log_scroll = dash.log_scroll.saturating_add(page_size); + } + KeyCode::PageDown => { + let mut dash = dashboard.lock().unwrap(); + dash.log_scroll = dash.log_scroll.saturating_sub(page_size); + } + + // F2: Reset selected table + KeyCode::F(2) => { + let (table_id, table_name) = { + let dash = dashboard.lock().unwrap(); + match dash.selected_table_info() { + Some(t) => (t.table_id, t.destination_name.clone()), + None => continue, + } + }; + + info!("F2: resetting table {}, stopping pipeline...", table_name); + dashboard.lock().unwrap().set_status(format!("Resetting {table_name}...")); + + // Stop pipeline (wait consumes self, so always rebuild after) + pipeline.shutdown(); + let wait_err = pipeline.wait().await.err(); + + let client = Client::new(config.clone(), Arc::clone(&auth), pipeline_id); + destination = Destination::new(client, store.clone()); + pipeline = + Pipeline::new(pipeline_config.clone(), store.clone(), destination.clone()); + + if let Some(e) = wait_err { + error!(error = %e, "pipeline shutdown failed during reset"); + dashboard.lock().unwrap().pipeline_running = false; + dashboard.lock().unwrap().phase = GlobalPhase::Stopped; + dashboard.lock().unwrap().set_status(format!("Reset failed: {e}")); + continue; + } + dashboard.lock().unwrap().pipeline_running = false; + dashboard.lock().unwrap().phase = GlobalPhase::Stopped; + + // Reset table state + commands::reset_table(table_id, &table_name, &store, &destination, &dashboard) + .await; + + // Restart pipeline + info!("restarting pipeline after reset..."); + if let Err(e) = pipeline.start().await { + error!(error = %e, "pipeline restart failed after reset"); + dashboard.lock().unwrap().set_status(format!( + "Reset done but restart failed: {e} -- press F3 to retry" + )); + continue; + } + dashboard.lock().unwrap().pipeline_running = true; + + info!("pipeline restarted after table reset"); + dashboard + .lock() + .unwrap() + .set_status(format!("{table_name} reset complete, pipeline restarted")); + } + + // F3: Restart pipeline + KeyCode::F(3) => { + info!("F3: restarting pipeline..."); + dashboard.lock().unwrap().set_status("Restarting pipeline...".to_owned()); + + // Stop (wait consumes self, so always rebuild after) + pipeline.shutdown(); + let wait_err = pipeline.wait().await.err(); + + let client = Client::new(config.clone(), Arc::clone(&auth), pipeline_id); + destination = Destination::new(client, store.clone()); + pipeline = + Pipeline::new(pipeline_config.clone(), store.clone(), destination.clone()); + + if let Some(e) = wait_err { + error!(error = %e, "pipeline shutdown failed during restart"); + dashboard.lock().unwrap().pipeline_running = false; + dashboard.lock().unwrap().phase = GlobalPhase::Stopped; + dashboard + .lock() + .unwrap() + .set_status(format!("Restart failed: {e} -- press F3 to retry")); + continue; + } + dashboard.lock().unwrap().pipeline_running = false; + dashboard.lock().unwrap().phase = GlobalPhase::Stopped; + + // Restart + if let Err(e) = pipeline.start().await { + error!(error = %e, "pipeline restart failed"); + dashboard + .lock() + .unwrap() + .set_status(format!("Restart failed: {e} -- press F3 to retry")); + continue; + } + dashboard.lock().unwrap().pipeline_running = true; + + info!("pipeline restarted"); + dashboard.lock().unwrap().set_status("Pipeline restarted".to_owned()); + } + + // F4: Cycle log level filter + KeyCode::F(4) => { + let mut dash = dashboard.lock().unwrap(); + dash.log_level_filter = dash.log_level_filter.next(); + let label = dash.log_level_filter.label(); + dash.set_status(format!("Log filter: {label}")); + } + + // F5: Sync Snowflake offsets + KeyCode::F(5) => { + info!("F5: syncing Snowflake offsets..."); + dashboard + .lock() + .unwrap() + .set_status("Querying Snowflake offsets...".to_owned()); + + state::fetch_snowflake_offsets(&dashboard, &destination).await; + last_offset_fetch = Instant::now(); + + dashboard.lock().unwrap().set_status("Snowflake offsets synced".to_owned()); + info!("Snowflake offset sync complete"); + } + + _ => {} + } + } + } + + // Wait for pipeline to finish + pipeline.wait().await?; + info!("pipeline shut down cleanly"); + + Ok(()) +} + +async fn count_source_tables(pg_config: &PgConnectionConfig) -> Result<(u64, u64), Box> { + let conn_str = format!( + "host={} port={} dbname={} user={} {}", + pg_config.host, + pg_config.port, + pg_config.name, + pg_config.username, + pg_config + .password + .as_ref() + .map(|p| { + use secrecy::ExposeSecret; + format!("password={}", p.expose_secret()) + }) + .unwrap_or_default(), + ); + + let (client, conn) = tokio_postgres::connect(&conn_str, tokio_postgres::NoTls).await?; + tokio::spawn(conn); + + let table_row = client + .query_one( + "SELECT count(*) FROM information_schema.tables WHERE table_schema = 'bench' AND \ + table_type = 'BASE TABLE'", + &[], + ) + .await?; + let table_count: i64 = table_row.get(0); + + let row_row = client + .query_one( + "SELECT coalesce(sum(n_live_tup), 0)::bigint FROM pg_stat_user_tables WHERE \ + schemaname = 'bench'", + &[], + ) + .await?; + let row_count: i64 = row_row.get(0); + + Ok((table_count as u64, row_count as u64)) +} diff --git a/crates/etl-examples/src/bin/snowflake/state.rs b/crates/etl-examples/src/bin/snowflake/state.rs new file mode 100644 index 000000000..240f4d0d8 --- /dev/null +++ b/crates/etl-examples/src/bin/snowflake/state.rs @@ -0,0 +1,451 @@ +use std::{ + collections::BTreeMap, + sync::{Arc, Mutex}, + time::{Duration, Instant}, +}; + +use etl::{ + state::TableReplicationPhase, + store::{PostgresStore, StateStore}, + types::TableId, +}; +use etl_destinations::snowflake::{Destination, OffsetToken}; + +/// Per-table information displayed in the table list and detail panel. +#[derive(Clone)] +#[allow(dead_code)] +pub struct TableInfo { + pub table_id: TableId, + pub destination_name: String, + pub phase: TableReplicationPhase, + pub rows_synced: u64, + pub throughput: f64, + pub local_offset: Option, + pub snowflake_offset: Option, + pub last_known_lsn: Option, + pub error_reason: Option, + pub error_solution: Option, + pub phase_entered_at: Option, +} + +impl TableInfo { + pub fn phase_label(&self) -> &'static str { + match &self.phase { + TableReplicationPhase::Init => "init", + TableReplicationPhase::DataSync => "copying", + TableReplicationPhase::FinishedCopy => "copied", + TableReplicationPhase::SyncWait { .. } => "sync_wait", + TableReplicationPhase::Catchup { .. } => "catchup", + TableReplicationPhase::SyncDone { .. } => "sync_done", + TableReplicationPhase::Ready => "ready", + TableReplicationPhase::Errored { .. } => "errored", + } + } + + pub fn is_errored(&self) -> bool { + matches!(&self.phase, TableReplicationPhase::Errored { .. }) + } + + pub fn is_copying(&self) -> bool { + matches!(&self.phase, TableReplicationPhase::DataSync | TableReplicationPhase::Init) + } + + pub fn is_ready(&self) -> bool { + matches!(&self.phase, TableReplicationPhase::Ready) + } +} + +/// Global dashboard state shared between the monitor task and the TUI render +/// loop. +#[allow(dead_code)] +pub struct DashboardState { + pub tables: Vec, + pub selected_table: usize, + + pub total_rows_synced: u64, + pub total_copy_rows: u64, + pub total_cdc_events: u64, + pub estimated_rows: u64, + pub table_count: u64, + + pub elapsed: Duration, + pub phase: GlobalPhase, + pub copy_elapsed: Option, + + pub throughput_current: f64, + pub throughput_avg: f64, + pub throughput_min: f64, + pub throughput_max: f64, + + pub api_calls: u64, + pub api_errors: u64, + pub channel_recoveries: u64, + + pub log_scroll: usize, + pub show_help: bool, + pub status_message: Option<(String, Instant)>, + pub log_level_filter: LogLevel, + + pub pipeline_running: bool, +} + +#[derive(Clone, Copy, PartialEq, Eq)] +pub enum GlobalPhase { + TableCopy, + CdcStreaming, + Stopped, +} + +#[derive(Clone, Copy, PartialEq, Eq)] +pub enum LogLevel { + All, + WarnAndAbove, + ErrorOnly, +} + +impl LogLevel { + pub fn next(self) -> Self { + match self { + Self::All => Self::WarnAndAbove, + Self::WarnAndAbove => Self::ErrorOnly, + Self::ErrorOnly => Self::All, + } + } + + pub fn label(self) -> &'static str { + match self { + Self::All => "All", + Self::WarnAndAbove => "Warn+", + Self::ErrorOnly => "Error", + } + } +} + +impl GlobalPhase { + pub fn label(&self) -> &'static str { + match self { + Self::TableCopy => "Table Copy", + Self::CdcStreaming => "CDC Streaming", + Self::Stopped => "Stopped", + } + } +} + +#[allow(dead_code)] +impl DashboardState { + pub fn new(table_count: u64, estimated_rows: u64) -> Self { + Self { + tables: Vec::new(), + selected_table: 0, + total_rows_synced: 0, + total_copy_rows: 0, + total_cdc_events: 0, + estimated_rows, + table_count, + elapsed: Duration::ZERO, + phase: GlobalPhase::TableCopy, + copy_elapsed: None, + throughput_current: 0.0, + throughput_avg: 0.0, + throughput_min: 0.0, + throughput_max: 0.0, + api_calls: 0, + api_errors: 0, + channel_recoveries: 0, + log_scroll: 0, + show_help: false, + status_message: None, + log_level_filter: LogLevel::All, + pipeline_running: true, + } + } + + pub fn selected_table_info(&self) -> Option<&TableInfo> { + self.tables.get(self.selected_table) + } + + pub fn selected_table_id(&self) -> Option { + self.tables.get(self.selected_table).map(|t| t.table_id) + } + + pub fn select_next(&mut self) { + if !self.tables.is_empty() { + self.selected_table = (self.selected_table + 1).min(self.tables.len() - 1); + } + } + + pub fn select_prev(&mut self) { + self.selected_table = self.selected_table.saturating_sub(1); + } + + pub fn set_status(&mut self, msg: String) { + self.status_message = Some((msg, Instant::now())); + } + + pub fn clear_stale_status(&mut self) { + if let Some((_, at)) = &self.status_message + && at.elapsed() > Duration::from_secs(10) + { + self.status_message = None; + } + } +} + +pub fn parse_prometheus_metric(text: &str, metric_name: &str) -> Option { + let prefix = format!("{metric_name} "); + text.lines() + .find(|line| line.starts_with(&prefix)) + .and_then(|line| line[prefix.len()..].trim().parse::().ok()) +} + +pub fn parse_prometheus_metric_sum(text: &str, metric_name: &str) -> Option { + let mut total = 0.0; + let mut found = false; + for line in text.lines() { + if let Some(rest) = line.strip_prefix(metric_name) + && (rest.starts_with(' ') || rest.starts_with('{')) + && let Some(val) = rest.split_whitespace().last().and_then(|v| v.parse::().ok()) + { + total += val; + found = true; + } + } + found.then_some(total) +} + +pub fn parse_prometheus_metric_with_label( + text: &str, + metric_name: &str, + label_key: &str, + label_value: &str, +) -> Option { + let needle = format!("{label_key}=\"{label_value}\""); + let mut total = 0.0; + let mut found = false; + for line in text.lines() { + if let Some(rest) = line.strip_prefix(metric_name) + && rest.starts_with('{') + && rest.contains(&needle) + && let Some(val) = rest.split_whitespace().last().and_then(|v| v.parse::().ok()) + { + total += val; + found = true; + } + } + found.then_some(total) +} + +/// Update dashboard state from Prometheus metrics and the state store. +#[allow(clippy::too_many_arguments)] +pub async fn refresh_dashboard( + dashboard: &Arc>, + store: &PostgresStore, + metrics_text: &str, + start_time: Instant, + samples: &mut Vec, + last_rows: &mut u64, + last_tick: &mut Instant, + table_names: &BTreeMap, + phase_times: &mut BTreeMap, +) { + let elapsed = start_time.elapsed(); + + let rows_synced = + parse_prometheus_metric(metrics_text, "etl_snowflake_batch_size_sum").unwrap_or(0.0) as u64; + let copy_rows = parse_prometheus_metric_with_label( + metrics_text, + "etl_events_processed_total", + "action", + "table_copy", + ) + .unwrap_or(0.0) as u64; + let cdc_events = parse_prometheus_metric_with_label( + metrics_text, + "etl_events_processed_total", + "action", + "table_streaming", + ) + .unwrap_or(0.0) as u64; + let api_errors = parse_prometheus_metric_sum(metrics_text, "etl_snowflake_insert_errors_total") + .unwrap_or(0.0) as u64; + let channel_recoveries = + parse_prometheus_metric_sum(metrics_text, "etl_snowflake_channel_recoveries_total") + .unwrap_or(0.0) as u64; + + let tick_elapsed = last_tick.elapsed().as_secs_f64(); + let delta = rows_synced.saturating_sub(*last_rows); + let throughput_current = if tick_elapsed > 0.0 { delta as f64 / tick_elapsed } else { 0.0 }; + *last_rows = rows_synced; + *last_tick = Instant::now(); + + if throughput_current > 0.0 { + samples.push(throughput_current); + } + let throughput_avg = + if samples.is_empty() { 0.0 } else { samples.iter().sum::() / samples.len() as f64 }; + let throughput_min = samples.iter().copied().fold(f64::MAX, f64::min); + let throughput_max = samples.iter().copied().fold(0.0f64, f64::max); + + let states = store.get_table_replication_states().await.ok(); + + let mut table_infos: Vec = Vec::new(); + let mut all_ready = true; + let mut any_table = false; + + // Preserve per-table state from previous refresh cycle + let prev_table_state: BTreeMap, Option)> = { + let dash = dashboard.lock().unwrap(); + dash.tables + .iter() + .map(|t| (t.table_id, (t.last_known_lsn.clone(), t.snowflake_offset.clone()))) + .collect() + }; + + if let Some(ref states) = states { + for (id, phase) in states.iter() { + any_table = true; + let dest_name = table_names.get(id).cloned().unwrap_or_else(|| format!("{id}")); + + let phase_label = phase_label_str(phase); + let prev = phase_times.get(id).map(|(l, _)| l.as_str()); + if prev != Some(phase_label) { + phase_times.insert(*id, (phase_label.to_owned(), Instant::now())); + } + let phase_entered = phase_times.get(id).map(|(_, t)| *t); + + let (error_reason, error_solution) = + if let TableReplicationPhase::Errored { reason, solution, .. } = phase { + all_ready = false; + (Some(reason.clone()), solution.clone()) + } else { + if !matches!(phase, TableReplicationPhase::Ready) { + all_ready = false; + } + (None, None) + }; + + let current_lsn = extract_lsn(phase); + let (prev_lsn, prev_sf_offset) = + prev_table_state.get(id).cloned().unwrap_or((None, None)); + let last_known_lsn = current_lsn.clone().or(prev_lsn); + + table_infos.push(TableInfo { + table_id: *id, + destination_name: dest_name, + phase: phase.clone(), + rows_synced: 0, + throughput: 0.0, + local_offset: current_lsn.or_else(|| offset_description(phase)), + snowflake_offset: prev_sf_offset, + last_known_lsn, + error_reason, + error_solution, + phase_entered_at: phase_entered, + }); + } + } + + let global_phase = + if !any_table || !all_ready { GlobalPhase::TableCopy } else { GlobalPhase::CdcStreaming }; + + let mut dash = dashboard.lock().unwrap(); + let prev_selected = dash.selected_table; + dash.elapsed = elapsed; + dash.total_rows_synced = rows_synced; + dash.total_copy_rows = copy_rows; + dash.total_cdc_events = cdc_events; + dash.throughput_current = throughput_current; + dash.throughput_avg = throughput_avg; + dash.throughput_min = if throughput_min == f64::MAX { 0.0 } else { throughput_min }; + dash.throughput_max = throughput_max; + dash.api_errors = api_errors; + dash.channel_recoveries = channel_recoveries; + dash.tables = table_infos; + dash.selected_table = prev_selected.min(dash.tables.len().saturating_sub(1)); + + let was_copying = dash.phase == GlobalPhase::TableCopy; + if was_copying && global_phase == GlobalPhase::CdcStreaming { + dash.copy_elapsed = Some(elapsed); + } + if dash.pipeline_running { + dash.phase = global_phase; + } + + dash.clear_stale_status(); +} + +fn phase_label_str(phase: &TableReplicationPhase) -> &'static str { + match phase { + TableReplicationPhase::Init => "init", + TableReplicationPhase::DataSync => "data_sync", + TableReplicationPhase::FinishedCopy => "finished_copy", + TableReplicationPhase::SyncWait { .. } => "sync_wait", + TableReplicationPhase::Catchup { .. } => "catchup", + TableReplicationPhase::SyncDone { .. } => "sync_done", + TableReplicationPhase::Ready => "ready", + TableReplicationPhase::Errored { .. } => "errored", + } +} + +fn extract_lsn(phase: &TableReplicationPhase) -> Option { + match phase { + TableReplicationPhase::SyncWait { lsn } + | TableReplicationPhase::Catchup { lsn } + | TableReplicationPhase::SyncDone { lsn } => Some(format!("{lsn}")), + _ => None, + } +} + +fn offset_description(phase: &TableReplicationPhase) -> Option { + match phase { + TableReplicationPhase::Init => Some("not started".to_owned()), + TableReplicationPhase::DataSync => Some("copying...".to_owned()), + TableReplicationPhase::FinishedCopy => Some("copy done, syncing".to_owned()), + TableReplicationPhase::Ready => Some("streaming".to_owned()), + _ => None, + } +} + +/// Build a map from TableId to destination table name by querying the store. +pub async fn build_table_name_map(store: &PostgresStore) -> BTreeMap { + let mut map = BTreeMap::new(); + let states = store.get_table_replication_states().await.ok(); + if let Some(states) = states { + for id in states.keys() { + if let Ok(Some(meta)) = store.get_destination_table_metadata(*id).await { + map.insert(*id, meta.destination_table_id); + } + } + } + map +} + +/// Fetch Snowflake committed offsets for all tables. +pub async fn fetch_snowflake_offsets( + dashboard: &Arc>, + destination: &Destination, +) { + let table_ids: Vec = { + let dash = dashboard.lock().unwrap(); + dash.tables.iter().map(|t| t.table_id).collect() + }; + + let mut offsets: BTreeMap> = BTreeMap::new(); + for id in &table_ids { + match destination.committed_offset(*id).await { + Ok(offset) => { + offsets.insert(*id, offset); + } + Err(_) => { + offsets.insert(*id, None); + } + } + } + + let mut dash = dashboard.lock().unwrap(); + for table in &mut dash.tables { + if let Some(offset) = offsets.get(&table.table_id) { + table.snowflake_offset = offset.as_ref().map(|o| format!("{o}")); + } + } +} diff --git a/crates/etl-examples/src/bin/snowflake/tui.rs b/crates/etl-examples/src/bin/snowflake/tui.rs new file mode 100644 index 000000000..6d19e2309 --- /dev/null +++ b/crates/etl-examples/src/bin/snowflake/tui.rs @@ -0,0 +1,532 @@ +use std::{ + collections::VecDeque, + io::Stdout, + sync::{Arc, Mutex}, +}; + +type Term = ratatui::Terminal>; + +use ratatui::{ + Terminal, + backend::CrosstermBackend, + layout::{Constraint, Direction, Layout, Rect}, + style::{Color, Modifier, Style}, + symbols, + text::{Line, Span}, + widgets::{Block, Borders, Cell, Clear, LineGauge, Paragraph, Row, Table, Wrap}, +}; + +use crate::{ + commands, + state::{DashboardState, GlobalPhase, LogLevel, TableInfo}, +}; + +pub struct TerminalGuard; + +impl Drop for TerminalGuard { + fn drop(&mut self) { + restore_terminal(); + } +} + +pub fn setup_terminal() -> Result<(Term, TerminalGuard), Box> { + crossterm::terminal::enable_raw_mode()?; + let mut stdout = std::io::stdout(); + crossterm::execute!(stdout, crossterm::terminal::EnterAlternateScreen)?; + Ok((Terminal::new(CrosstermBackend::new(stdout))?, TerminalGuard)) +} + +pub fn restore_terminal() { + let _ = crossterm::terminal::disable_raw_mode(); + let _ = crossterm::execute!(std::io::stdout(), crossterm::terminal::LeaveAlternateScreen); +} + +pub fn render( + frame: &mut ratatui::Frame, + state: &DashboardState, + log_buffer: &Arc>>, +) { + if state.show_help { + render_help_overlay(frame); + return; + } + + let main_chunks = Layout::default() + .direction(Direction::Vertical) + .constraints([ + Constraint::Length(5), // global stats + gauge + Constraint::Percentage(45), // table + detail + Constraint::Min(6), // logs (takes remaining space) + Constraint::Length(1), // command bar + ]) + .split(frame.area()); + + render_global_stats(frame, main_chunks[0], state); + render_middle_panel(frame, main_chunks[1], state); + render_log_panel(frame, main_chunks[2], log_buffer, state.log_scroll, state.log_level_filter); + render_command_bar(frame, main_chunks[3], state); +} + +fn render_global_stats(frame: &mut ratatui::Frame, area: Rect, state: &DashboardState) { + let block = Block::default().title(" Snowflake CDC Pipeline ").borders(Borders::ALL); + let inner = block.inner(area); + frame.render_widget(block, area); + + let elapsed_secs = state.elapsed.as_secs(); + let elapsed_str = format!( + "{:02}:{:02}:{:02}", + elapsed_secs / 3600, + (elapsed_secs % 3600) / 60, + elapsed_secs % 60 + ); + + let phase_color = match state.phase { + GlobalPhase::TableCopy => Color::Green, + GlobalPhase::CdcStreaming => Color::Cyan, + GlobalPhase::Stopped => Color::Red, + }; + + let running_indicator = if state.pipeline_running { + Span::styled(" RUNNING ", Style::default().fg(Color::Black).bg(Color::Green)) + } else { + Span::styled(" STOPPED ", Style::default().fg(Color::Black).bg(Color::Red)) + }; + + // Row 1: phase, elapsed, tables + let header = Line::from(vec![ + running_indicator, + Span::raw(" Phase: "), + Span::styled(state.phase.label(), Style::default().fg(phase_color)), + Span::raw(format!(" Elapsed: {elapsed_str} ")), + Span::raw(format!("Tables: {}", state.table_count)), + ]); + + // Row 2: copy/CDC stats as text + let stats_line = if state.phase == GlobalPhase::CdcStreaming { + let copy_part = if let Some(copy_elapsed) = state.copy_elapsed { + format!( + "Copy: {} rows in {:.1}s", + format_num(state.total_copy_rows), + copy_elapsed.as_secs_f64() + ) + } else { + format!("Copy: {} rows", format_num(state.total_copy_rows)) + }; + Line::from(vec![ + Span::raw(format!(" {copy_part}")), + Span::styled(" | ", Style::default().fg(Color::DarkGray)), + Span::styled( + format!("CDC: {} events", format_num(state.total_cdc_events)), + Style::default().fg(Color::Cyan), + ), + Span::styled(" | ", Style::default().fg(Color::DarkGray)), + Span::raw(format!("{:.0} r/s", state.throughput_current)), + if state.api_errors > 0 { + Span::styled( + format!(" Errors: {}", state.api_errors), + Style::default().fg(Color::Red), + ) + } else { + Span::raw("") + }, + ]) + } else { + let total = state.estimated_rows.max(1); + let pct = (state.total_rows_synced as f64 / total as f64 * 100.0).min(100.0); + Line::from(vec![ + Span::raw(format!( + " Copy: {} / {} ({:.0}%)", + format_num(state.total_rows_synced), + format_num(total), + pct + )), + Span::styled(" | ", Style::default().fg(Color::DarkGray)), + Span::raw(format!( + "{:.0} r/s (avg: {:.0})", + state.throughput_current, state.throughput_avg + )), + if state.api_errors > 0 { + Span::styled( + format!(" Errors: {}", state.api_errors), + Style::default().fg(Color::Red), + ) + } else { + Span::raw("") + }, + ]) + }; + + // Row 3: progress gauge + let total = state.estimated_rows.max(1); + let ratio = if state.phase == GlobalPhase::TableCopy { + (state.total_rows_synced as f64 / total as f64).clamp(0.0, 1.0) + } else { + 1.0 + }; + + let gauge_color = match state.phase { + GlobalPhase::TableCopy => Color::Green, + GlobalPhase::CdcStreaming => Color::Cyan, + GlobalPhase::Stopped => Color::DarkGray, + }; + + let sub = Layout::default() + .direction(Direction::Vertical) + .constraints([Constraint::Length(1), Constraint::Length(1), Constraint::Length(1)]) + .split(inner); + + frame.render_widget(Paragraph::new(vec![header]), sub[0]); + frame.render_widget(Paragraph::new(vec![stats_line]), sub[1]); + + let gauge_label = if state.phase == GlobalPhase::CdcStreaming { + if let Some((ref msg, _)) = state.status_message { + msg.clone() + } else { + "streaming".to_owned() + } + } else { + format!(" {:.0}%", ratio * 100.0) + }; + + frame.render_widget( + LineGauge::default() + .filled_style(Style::default().fg(gauge_color)) + .ratio(ratio) + .label(gauge_label) + .filled_symbol(symbols::line::THICK.horizontal) + .unfilled_symbol(symbols::line::NORMAL.horizontal), + sub[2], + ); +} + +fn render_middle_panel(frame: &mut ratatui::Frame, area: Rect, state: &DashboardState) { + let chunks = Layout::default() + .direction(Direction::Horizontal) + .constraints([Constraint::Percentage(75), Constraint::Percentage(25)]) + .split(area); + + render_table_list(frame, chunks[0], state); + render_table_detail(frame, chunks[1], state); +} + +fn render_table_list(frame: &mut ratatui::Frame, area: Rect, state: &DashboardState) { + if state.tables.is_empty() { + let block = Block::default().title(" Tables (j/k) ").borders(Borders::ALL); + let inner = block.inner(area); + frame.render_widget(block, area); + frame.render_widget(Paragraph::new(" Waiting for tables..."), inner); + return; + } + + let header = Row::new(vec![ + Cell::from(" Name"), + Cell::from("Phase"), + Cell::from("Rows"), + Cell::from("r/s"), + Cell::from(" "), + ]) + .style(Style::default().fg(Color::Cyan).add_modifier(Modifier::BOLD)) + .bottom_margin(0); + + let rows: Vec = state + .tables + .iter() + .enumerate() + .map(|(i, table)| { + let selected = i == state.selected_table; + let base_style = if selected { + Style::default().bg(Color::DarkGray).fg(Color::White).add_modifier(Modifier::BOLD) + } else { + Style::default().fg(Color::White) + }; + + let status_icon = if table.is_errored() { + Cell::from(Span::styled("✗", Style::default().fg(Color::Red))) + } else if table.is_ready() { + Cell::from(Span::styled("✓", Style::default().fg(Color::Green))) + } else if table.is_copying() { + Cell::from(Span::styled("↻", Style::default().fg(Color::Yellow))) + } else { + Cell::from("·") + }; + + Row::new(vec![ + Cell::from(format!( + "{}{}", + if selected { "▸" } else { " " }, + truncate_str(&table.destination_name, 24) + )), + Cell::from(Span::styled( + table.phase_label(), + Style::default().fg(phase_color(table.phase_label())), + )), + Cell::from(format_compact(table.rows_synced)), + Cell::from(format!("{:.0}", table.throughput)), + status_icon, + ]) + .style(base_style) + }) + .collect(); + + let widths = [ + Constraint::Min(20), + Constraint::Length(12), + Constraint::Length(10), + Constraint::Length(8), + Constraint::Length(2), + ]; + + let table_widget = Table::new(rows, widths) + .header(header) + .block(Block::default().title(" Tables (j/k) ").borders(Borders::ALL)); + + frame.render_widget(table_widget, area); +} + +fn render_table_detail(frame: &mut ratatui::Frame, area: Rect, state: &DashboardState) { + let title = match state.selected_table_info() { + Some(t) => format!(" {} ", truncate_str(&t.destination_name, 18)), + None => " Details ".to_owned(), + }; + + let block = Block::default().title(title).borders(Borders::ALL); + let inner = block.inner(area); + frame.render_widget(block, area); + + let Some(table) = state.selected_table_info() else { + frame.render_widget(Paragraph::new(" Select a table"), inner); + return; + }; + + let mut lines = vec![]; + + // Phase with visual indicator + let pc = phase_color(table.phase_label()); + let indicator = phase_indicator(table); + lines.push(Line::from(vec![ + Span::styled(table.phase_label(), Style::default().fg(pc).add_modifier(Modifier::BOLD)), + Span::raw(" "), + indicator, + ])); + + // Phase duration + if let Some(entered) = table.phase_entered_at { + let dur = entered.elapsed(); + lines.push(Line::from(Span::styled( + format!("{:.0}s in phase", dur.as_secs_f64()), + Style::default().fg(Color::DarkGray), + ))); + } + + lines.push(Line::from("")); + + // Offsets + lines.push(Line::from(vec![ + Span::styled("Local: ", Style::default().fg(Color::DarkGray)), + Span::raw(table.local_offset.as_deref().unwrap_or("—")), + ])); + lines.push(Line::from(vec![ + Span::styled("SF: ", Style::default().fg(Color::DarkGray)), + Span::raw(table.snowflake_offset.as_deref().unwrap_or("—")), + ])); + + lines.push(Line::from("")); + + // Global stats + lines.push(Line::from(vec![ + Span::styled("Copy: ", Style::default().fg(Color::DarkGray)), + Span::raw(format_num(state.total_copy_rows)), + ])); + lines.push(Line::from(vec![ + Span::styled("CDC: ", Style::default().fg(Color::DarkGray)), + Span::raw(format_num(state.total_cdc_events)), + ])); + + lines.push(Line::from("")); + + // Table OID + lines.push(Line::from(Span::styled( + format!("OID: {}", table.table_id), + Style::default().fg(Color::DarkGray), + ))); + + // Error details + if let Some(ref reason) = table.error_reason { + lines.push(Line::from("")); + lines.push(Line::from(Span::styled( + "Error:", + Style::default().fg(Color::Red).add_modifier(Modifier::BOLD), + ))); + for error_line in reason.lines() { + lines.push(Line::from(Span::styled(error_line, Style::default().fg(Color::Red)))); + } + if let Some(ref solution) = table.error_solution { + lines.push(Line::from("")); + lines.push(Line::from(Span::styled("Fix:", Style::default().fg(Color::Yellow)))); + for sol_line in solution.lines() { + lines.push(Line::from(Span::styled(sol_line, Style::default().fg(Color::Yellow)))); + } + } + lines.push(Line::from("")); + lines.push(Line::from(Span::styled("F2 to reset", Style::default().fg(Color::Magenta)))); + } + + frame.render_widget(Paragraph::new(lines).wrap(Wrap { trim: false }), inner); +} + +fn render_log_panel( + frame: &mut ratatui::Frame, + area: Rect, + log_buffer: &Arc>>, + scroll_offset: usize, + log_level: LogLevel, +) { + let title = format!(" Logs (J/K scroll) [{}] ", log_level.label()); + let block = Block::default().title(title).borders(Borders::ALL); + let inner = block.inner(area); + frame.render_widget(block, area); + + let logs: Vec = { + let buf = log_buffer.lock().unwrap(); + buf.iter() + .filter(|line| match log_level { + LogLevel::All => true, + LogLevel::WarnAndAbove => line.contains("[WARN]") || line.contains("[ERROR]"), + LogLevel::ErrorOnly => line.contains("[ERROR]"), + }) + .cloned() + .collect() + }; + + let visible_height = inner.height as usize; + let total = logs.len(); + let start = if total > visible_height { + let max_scroll = total - visible_height; + let from_bottom = scroll_offset.min(max_scroll); + max_scroll - from_bottom + } else { + 0 + }; + + let visible: Vec = logs + .iter() + .skip(start) + .take(visible_height) + .map(|line| { + let color = if line.contains("[ERROR]") { + Color::Red + } else if line.contains("[WARN]") { + Color::Yellow + } else { + Color::DarkGray + }; + Line::from(Span::styled(line.clone(), Style::default().fg(color))) + }) + .collect(); + + frame.render_widget(Paragraph::new(visible), inner); +} + +fn render_command_bar(frame: &mut ratatui::Frame, area: Rect, state: &DashboardState) { + let commands = [ + ("F1", "Help"), + ("F2", "Reset"), + ("F3", "Restart"), + ("F4", "LogLvl"), + ("F5", "Sync"), + ("F10", "Quit"), + ]; + + let mut spans = Vec::new(); + for (key, label) in &commands { + spans.push(Span::styled( + format!(" {key} "), + Style::default().fg(Color::Black).bg(Color::Cyan), + )); + spans.push(Span::raw(format!("{label} "))); + } + + if !state.pipeline_running { + spans.push(Span::raw(" ")); + spans.push(Span::styled( + " Pipeline stopped — F3 to restart ", + Style::default().fg(Color::Black).bg(Color::Yellow), + )); + } + + frame.render_widget(Paragraph::new(vec![Line::from(spans)]), area); +} + +fn render_help_overlay(frame: &mut ratatui::Frame) { + let area = frame.area(); + let width = 72.min(area.width.saturating_sub(4)); + let height = 38.min(area.height.saturating_sub(2)); + let x = (area.width.saturating_sub(width)) / 2; + let y = (area.height.saturating_sub(height)) / 2; + + let popup = Rect::new(x, y, width, height); + + frame.render_widget(Clear, popup); + let block = Block::default() + .title(" Help (press any key to close) ") + .borders(Borders::ALL) + .border_style(Style::default().fg(Color::Cyan)); + + let inner = block.inner(popup); + frame.render_widget(block, popup); + + let text: Vec = + commands::HELP_TEXT.lines().map(|l| Line::from(Span::raw(l.to_owned()))).collect(); + + frame.render_widget(Paragraph::new(text).wrap(Wrap { trim: false }), inner); +} + +fn phase_color(label: &str) -> Color { + match label { + "init" => Color::DarkGray, + "copying" | "data_sync" => Color::Green, + "copied" | "finished_copy" => Color::Blue, + "sync_wait" | "catchup" | "sync_done" => Color::Blue, + "ready" => Color::Cyan, + "errored" => Color::Red, + _ => Color::White, + } +} + +fn phase_indicator(table: &TableInfo) -> Span<'static> { + if table.is_copying() { + Span::styled("● copying", Style::default().fg(Color::Green)) + } else if table.is_ready() { + Span::styled("● CDC", Style::default().fg(Color::Cyan)) + } else if table.is_errored() { + Span::styled("● error", Style::default().fg(Color::Red)) + } else { + Span::styled("●", Style::default().fg(Color::DarkGray)) + } +} + +fn truncate_str(s: &str, max: usize) -> String { + if s.len() <= max { s.to_owned() } else { format!("{}…", &s[..max - 1]) } +} + +pub fn format_num(n: u64) -> String { + let s = n.to_string(); + let mut result = String::new(); + for (i, c) in s.chars().rev().enumerate() { + if i > 0 && i % 3 == 0 { + result.push(','); + } + result.push(c); + } + result.chars().rev().collect() +} + +fn format_compact(n: u64) -> String { + if n >= 1_000_000 { + format!("{:.1}M", n as f64 / 1_000_000.0) + } else if n >= 1_000 { + format!("{:.1}K", n as f64 / 1_000.0) + } else { + n.to_string() + } +} diff --git a/crates/etl-examples/src/bin/snowflake_loadgen.rs b/crates/etl-examples/src/bin/snowflake_loadgen.rs new file mode 100644 index 000000000..d3131a6d9 --- /dev/null +++ b/crates/etl-examples/src/bin/snowflake_loadgen.rs @@ -0,0 +1,457 @@ +use std::{error::Error, time::Instant}; + +use clap::{Parser, Subcommand}; +use rand::Rng; +use tokio::signal; +use tokio_postgres::NoTls; + +#[derive(Clone, Copy)] +enum Op { + Insert, + Update, + Delete, +} + +const EVENT_TYPES: &[&str] = + &["page_view", "click", "purchase", "signup", "logout", "search", "download", "share"]; + +const IP_PREFIXES: &[&str] = &["192.168.", "10.0.", "172.16.", "203.0.113."]; + +#[derive(Debug, Parser)] +#[command( + name = "snowflake-loadgen", + version, + about = "Load generator for Snowflake ETL benchmarking" +)] +struct AppArgs { + #[command(subcommand)] + command: Command, +} + +#[derive(Debug, Subcommand)] +enum Command { + Seed(SeedArgs), + Generate(GenerateArgs), +} + +#[derive(Debug, Parser)] +struct SeedArgs { + #[arg(long)] + db_url: String, + #[arg(long, default_value = "1000000")] + rows: u64, +} + +#[derive(Debug, Parser)] +struct GenerateArgs { + #[arg(long)] + db_url: String, + #[arg(long, default_value = "100")] + rate: u64, + /// Insert/update/delete percentages, e.g. "40/40/20" + #[arg(long, default_value = "40/40/20")] + mix: String, + /// Duration to run, e.g. "300s" or "5m". Omit to run indefinitely. + #[arg(long)] + duration: Option, +} + +fn random_event_type(rng: &mut impl Rng) -> &'static str { + EVENT_TYPES[rng.random_range(0..EVENT_TYPES.len())] +} + +fn random_amount(rng: &mut impl Rng) -> f64 { + (rng.random_range(0u64..100_000) as f64) / 100.0 +} + +fn random_jsonb(rng: &mut impl Rng) -> String { + format!(r#"{{"source":"web","version":{}}}"#, rng.random_range(1..10)) +} + +fn random_ip(rng: &mut impl Rng) -> String { + let prefix = IP_PREFIXES[rng.random_range(0..IP_PREFIXES.len())]; + format!("{}{}.{}", prefix, rng.random_range(0..256), rng.random_range(1..255)) +} + +fn random_tags(rng: &mut impl Rng) -> String { + let all = &["rust", "web", "mobile", "api", "beta", "premium", "trial"]; + let count = rng.random_range(0..4usize); + let tags: Vec<&str> = (0..count).map(|_| all[rng.random_range(0..all.len())]).collect(); + format!("{{{}}}", tags.join(",")) +} + +fn parse_duration(s: &str) -> Result> { + if let Some(stripped) = s.strip_suffix('s') { + let secs: u64 = stripped.parse()?; + return Ok(std::time::Duration::from_secs(secs)); + } + if let Some(stripped) = s.strip_suffix('m') { + let mins: u64 = stripped.parse()?; + return Ok(std::time::Duration::from_secs(mins * 60)); + } + if let Some(stripped) = s.strip_suffix('h') { + let hours: u64 = stripped.parse()?; + return Ok(std::time::Duration::from_secs(hours * 3600)); + } + Err(format!("invalid duration '{s}': expected format like '300s', '5m', or '2h'").into()) +} + +fn parse_mix(s: &str) -> Result<(u32, u32, u32), Box> { + let parts: Vec<&str> = s.split('/').collect(); + if parts.len() != 3 { + return Err(format!("invalid mix '{s}': expected format like '40/40/20'").into()); + } + let insert: u32 = parts[0].parse()?; + let update: u32 = parts[1].parse()?; + let delete: u32 = parts[2].parse()?; + if insert + update + delete != 100 { + return Err( + format!("mix percentages must sum to 100, got {}", insert + update + delete).into() + ); + } + Ok((insert, update, delete)) +} + +async fn seed(client: &tokio_postgres::Client, rows: u64) -> Result<(), Box> { + let mut rng = rand::rng(); + + // Seed users (10% of rows, min 100) + let user_count = (rows / 10).max(100); + let user_batch = 5000u64; + let mut users_inserted = 0u64; + + eprintln!("Seeding {user_count} users..."); + while users_inserted < user_count { + let batch = user_batch.min(user_count - users_inserted); + let mut values = Vec::with_capacity(batch as usize); + for i in 0..batch { + let n = users_inserted + i; + let name = format!("User {n}"); + let email = format!("user{n}@example.com"); + values.push(format!("('{name}', '{email}', now())")); + } + let sql = format!( + "INSERT INTO bench.users (name, email, created_at) VALUES {}", + values.join(",") + ); + client.execute(&sql, &[]).await?; + users_inserted += batch; + } + eprintln!(" users done: {users_inserted}"); + + // Get user ID range + let row = client.query_one("SELECT MIN(id), MAX(id) FROM bench.users", &[]).await?; + let min_uid: i32 = row.get(0); + let max_uid: i32 = row.get(1); + + // Seed orders (same count as rows) + let order_batch = 5000u64; + let mut orders_inserted = 0u64; + let statuses = &["pending", "completed", "cancelled", "refunded"]; + + eprintln!("Seeding {rows} orders..."); + while orders_inserted < rows { + let batch = order_batch.min(rows - orders_inserted); + let mut values = Vec::with_capacity(batch as usize); + for _ in 0..batch { + let uid = rng.random_range(min_uid..=max_uid); + let total = random_amount(&mut rng); + let status = statuses[rng.random_range(0..statuses.len())]; + values.push(format!("({uid}, {total:.2}, '{status}', now())")); + } + let sql = format!( + "INSERT INTO bench.orders (user_id, total, status, created_at) VALUES {}", + values.join(",") + ); + client.execute(&sql, &[]).await?; + orders_inserted += batch; + if orders_inserted.is_multiple_of(10_000) || orders_inserted == rows { + eprintln!(" orders {orders_inserted}/{rows}"); + } + } + + // Seed events + let event_batch = 5000u64; + let mut events_inserted = 0u64; + + eprintln!("Seeding {rows} events..."); + while events_inserted < rows { + let batch = event_batch.min(rows - events_inserted); + let mut values = Vec::with_capacity(batch as usize); + for _ in 0..batch { + let uid = rng.random_range(min_uid..=max_uid); + let event_type = random_event_type(&mut rng); + let amount = random_amount(&mut rng); + let metadata = random_jsonb(&mut rng); + let tags = random_tags(&mut rng); + let score: f64 = rng.random_range(0..10000) as f64 / 100.0; + let ip = random_ip(&mut rng); + values.push(format!( + "({uid}, '{event_type}', {amount:.2}, '{metadata}', '{tags}', true, {score:.2}, \ + now(), now(), '{ip}', '')" + )); + } + let sql = format!( + "INSERT INTO bench.events (user_id, event_type, amount, metadata, tags, is_active, \ + score, created_at, updated_at, ip_address, notes) VALUES {}", + values.join(",") + ); + client.execute(&sql, &[]).await?; + events_inserted += batch; + if events_inserted.is_multiple_of(10_000) || events_inserted == rows { + eprintln!(" events {events_inserted}/{rows}"); + } + } + + Ok(()) +} + +async fn run_seed(args: SeedArgs) -> Result<(), Box> { + eprintln!("Connecting to Postgres..."); + let (client, connection) = tokio_postgres::connect(&args.db_url, NoTls).await?; + tokio::spawn(async move { + if let Err(e) = connection.await { + eprintln!("connection error: {e}"); + } + }); + + eprintln!("Creating bench schema (users, orders, events)..."); + client + .batch_execute( + "DROP SCHEMA IF EXISTS bench CASCADE; + DROP PUBLICATION IF EXISTS bench_pub; + CREATE SCHEMA bench; + CREATE TABLE bench.users ( + id SERIAL PRIMARY KEY, + name TEXT NOT NULL, + email TEXT NOT NULL, + created_at TIMESTAMPTZ DEFAULT now() + ); + CREATE TABLE bench.orders ( + id SERIAL PRIMARY KEY, + user_id INT REFERENCES bench.users(id), + total NUMERIC(10,2), + status TEXT, + created_at TIMESTAMPTZ DEFAULT now() + ); + CREATE TABLE bench.events ( + id BIGSERIAL PRIMARY KEY, + user_id INT REFERENCES bench.users(id), + event_type TEXT NOT NULL, + amount NUMERIC(12,2), + metadata JSONB, + tags TEXT[], + is_active BOOLEAN DEFAULT true, + score DOUBLE PRECISION, + created_at TIMESTAMPTZ DEFAULT now(), + updated_at TIMESTAMPTZ DEFAULT now(), + ip_address TEXT, + notes TEXT + )", + ) + .await?; + + eprintln!("Seeding {} rows...", args.rows); + let start = Instant::now(); + seed(&client, args.rows).await?; + + let elapsed = start.elapsed(); + + eprintln!("Creating publication bench_pub..."); + client + .execute("CREATE PUBLICATION bench_pub FOR TABLES IN SCHEMA bench", &[]) + .await + .or_else(|e| if e.to_string().contains("already exists") { Ok(0) } else { Err(e) })?; + + eprintln!( + "Done. Inserted {} rows in {:.1}s ({:.0} rows/sec).", + args.rows, + elapsed.as_secs_f64(), + args.rows as f64 / elapsed.as_secs_f64() + ); + + Ok(()) +} + +async fn run_generate(args: GenerateArgs) -> Result<(), Box> { + let (insert_pct, update_pct, _delete_pct) = parse_mix(&args.mix)?; + let deadline = args.duration.as_deref().map(parse_duration).transpose()?; + + eprintln!("Connecting to Postgres..."); + let (client, connection) = tokio_postgres::connect(&args.db_url, NoTls).await?; + tokio::spawn(async move { + if let Err(e) = connection.await { + eprintln!("connection error: {e}"); + } + }); + + // Cache ID ranges for random lookups (avoid ORDER BY random() scans) + let id_range = client.query_one("SELECT MIN(id), MAX(id) FROM bench.events", &[]).await?; + let mut min_id: i64 = id_range.get(0); + let mut max_id: i64 = id_range.get(1); + + let user_row = client.query_one("SELECT MIN(id), MAX(id) FROM bench.users", &[]).await?; + let min_uid: i32 = user_row.get(0); + let max_uid: i32 = user_row.get(1); + + // Batch size: target rate / batches_per_sec. At 5000 ops/sec with 50 + // batches/sec = 100 ops/batch. + let batch_size = (args.rate / 50).clamp(1, 500) as usize; + let batches_per_sec = (args.rate as f64 / batch_size as f64).ceil() as u64; + let interval_us = 1_000_000u64 / batches_per_sec.max(1); + let mut ticker = tokio::time::interval(std::time::Duration::from_micros(interval_us)); + ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Burst); + + let start = Instant::now(); + let mut total_ops: u64 = 0; + let mut last_report = Instant::now(); + let mut ops_since_report: u64 = 0; + let mut rng = rand::rng(); + + let ctrl_c = signal::ctrl_c(); + tokio::pin!(ctrl_c); + + eprintln!( + "Generating CDC changes at {} ops/sec (batch={}, mix: {}){}", + args.rate, + batch_size, + args.mix, + match deadline { + Some(d) => format!(", duration: {}s", d.as_secs()), + None => ", running until ctrl-c".to_owned(), + } + ); + + loop { + tokio::select! { + biased; + _ = &mut ctrl_c => { + eprintln!("\nReceived ctrl-c, stopping."); + break; + } + _ = ticker.tick() => {} + } + + if let Some(d) = deadline + && start.elapsed() >= d + { + eprintln!("Duration reached, stopping."); + break; + } + + // Build a batch of operations as a single SQL string + let mut stmts: Vec = Vec::with_capacity(batch_size); + let mut insert_values: Vec = Vec::new(); + let mut batch_ops = 0u64; + + for _ in 0..batch_size { + let roll: u32 = rng.random_range(0..100); + let op = if roll < insert_pct { + Op::Insert + } else if roll < insert_pct + update_pct { + Op::Update + } else { + Op::Delete + }; + + match op { + Op::Insert => { + let uid = rng.random_range(min_uid..=max_uid); + let event_type = random_event_type(&mut rng); + let amount = random_amount(&mut rng); + let metadata = random_jsonb(&mut rng); + let tags = random_tags(&mut rng); + let score: f64 = rng.random_range(0..10000) as f64 / 100.0; + let ip = random_ip(&mut rng); + insert_values.push(format!( + "({uid}, '{event_type}', {amount:.2}, '{metadata}', '{tags}', true, \ + {score:.2}, now(), now(), '{ip}', '')" + )); + } + Op::Update => { + let target_id = rng.random_range(min_id..=max_id); + let score: f64 = rng.random_range(0..10000) as f64 / 100.0; + let event_type = random_event_type(&mut rng); + stmts.push(format!( + "UPDATE bench.events SET score = {score:.2}, event_type = '{event_type}', \ + updated_at = now() WHERE id = {target_id}" + )); + } + Op::Delete => { + let target_id = rng.random_range(min_id..=max_id); + stmts.push(format!("DELETE FROM bench.events WHERE id = {target_id}")); + } + } + batch_ops += 1; + } + + // Flush accumulated inserts as one multi-row INSERT + if !insert_values.is_empty() { + let cols = "user_id, event_type, amount, metadata, tags, is_active, score, \ + created_at, updated_at, ip_address, notes"; + stmts.insert( + 0, + format!("INSERT INTO bench.events ({cols}) VALUES {}", insert_values.join(",")), + ); + } + + if stmts.is_empty() { + continue; + } + + let sql = stmts.join("; "); + match client.batch_execute(&sql).await { + Ok(()) => { + total_ops += batch_ops; + ops_since_report += batch_ops; + // Track max_id growth from inserts + max_id += insert_values.len() as i64; + } + Err(e) => { + eprintln!("batch error: {e}"); + } + } + + let since_report = last_report.elapsed(); + if since_report >= std::time::Duration::from_secs(5) { + let throughput = ops_since_report as f64 / since_report.as_secs_f64(); + eprintln!( + "[{:.0}s] {:.0} ops/sec (total: {})", + start.elapsed().as_secs_f64(), + throughput, + total_ops + ); + ops_since_report = 0; + last_report = Instant::now(); + + // Refresh ID range periodically to account for inserts/deletes + if let Ok(row) = + client.query_one("SELECT MIN(id), MAX(id) FROM bench.events", &[]).await + { + min_id = row.get(0); + max_id = row.get(1); + } + } + } + + let elapsed = start.elapsed(); + eprintln!( + "Finished. {} ops in {:.1}s ({:.0} ops/sec average).", + total_ops, + elapsed.as_secs_f64(), + total_ops as f64 / elapsed.as_secs_f64().max(0.001) + ); + + Ok(()) +} + +#[tokio::main] +async fn main() -> Result<(), Box> { + let args = AppArgs::parse(); + match args.command { + Command::Seed(seed_args) => run_seed(seed_args).await?, + Command::Generate(gen_args) => run_generate(gen_args).await?, + } + + Ok(()) +} diff --git a/crates/etl-replicator/Cargo.toml b/crates/etl-replicator/Cargo.toml index 79b01c513..c6667c593 100644 --- a/crates/etl-replicator/Cargo.toml +++ b/crates/etl-replicator/Cargo.toml @@ -17,7 +17,7 @@ chrono = { workspace = true } configcat = { workspace = true } etl = { workspace = true } etl-config = { workspace = true, features = ["supabase"] } -etl-destinations = { workspace = true, features = ["bigquery", "clickhouse", "ducklake", "iceberg"] } +etl-destinations = { workspace = true, features = ["bigquery", "clickhouse", "ducklake", "iceberg", "snowflake"] } etl-maintenance = { workspace = true, features = ["ducklake"] } etl-telemetry = { workspace = true } k8s-openapi = { workspace = true, features = ["latest"] } diff --git a/crates/etl-replicator/src/core.rs b/crates/etl-replicator/src/core.rs index 782a2657d..3e7a9dfa2 100644 --- a/crates/etl-replicator/src/core.rs +++ b/crates/etl-replicator/src/core.rs @@ -28,6 +28,7 @@ use etl_destinations::{ DestinationNamespace, IcebergClient, IcebergDestination, S3_ACCESS_KEY_ID, S3_ENDPOINT, S3_SECRET_ACCESS_KEY, }, + snowflake, }; use secrecy::ExposeSecret; use tokio::signal::unix::{SignalKind, signal}; @@ -218,6 +219,33 @@ pub(crate) async fn start_replicator_with_config( )?; destination.validate_engine_support().await?; + let pipeline = Pipeline::new(replicator_config.pipeline, state_store, destination); + start_pipeline(pipeline).await?; + } + DestinationConfig::Snowflake { + account_id, + user, + private_key_path, + private_key_passphrase, + database, + schema, + role, + } => { + let mut config = snowflake::Config::new(account_id, user, database, schema); + if let Some(r) = role { + config = config.with_role(r); + } + let auth = std::sync::Arc::new( + snowflake::AuthManager::new( + &config, + private_key_path, + private_key_passphrase.as_ref(), + ) + .map_err(|e| ReplicatorError::config(std::io::Error::other(e.to_string())))?, + ); + let client = snowflake::Client::new(config, auth, pipeline_id); + let destination = snowflake::Destination::new(client, state_store.clone()); + let pipeline = Pipeline::new(replicator_config.pipeline, state_store, destination); start_pipeline(pipeline).await?; } diff --git a/crates/etl-replicator/src/init/destination.rs b/crates/etl-replicator/src/init/destination.rs index 6dc8fb36d..5f292f729 100644 --- a/crates/etl-replicator/src/init/destination.rs +++ b/crates/etl-replicator/src/init/destination.rs @@ -2,7 +2,7 @@ use etl::{destination::Destination, store::both::postgres::PostgresStore}; use etl_config::shared::{DestinationConfig, IcebergConfig}; use etl_destinations::{ bigquery::BigQueryDestination, clickhouse::ClickHouseDestination, - ducklake::DuckLakeDestination, iceberg::IcebergDestination, + ducklake::DuckLakeDestination, iceberg::IcebergDestination, snowflake, }; use crate::error_reporting::ErrorReportingStateStore; @@ -20,5 +20,8 @@ pub(crate) fn destination_name(destination_config: &DestinationConfig) -> &'stat DestinationConfig::ClickHouse { .. } => { ClickHouseDestination::::name() } + DestinationConfig::Snowflake { .. } => { + snowflake::Destination::::name() + } } } diff --git a/crates/xtask/src/commands/nextest.rs b/crates/xtask/src/commands/nextest.rs index c01527b2d..16be8558b 100644 --- a/crates/xtask/src/commands/nextest.rs +++ b/crates/xtask/src/commands/nextest.rs @@ -13,7 +13,7 @@ use clap::{Args, ValueEnum}; /// `.config/nextest.toml`. const SHARED_PG_FILTER: &str = "\ test(exclusive_) | binary_id(etl::main) | (binary_id(etl-destinations::main) & \ - test(/^(bigquery|clickhouse|ducklake|iceberg)::/)) | \ + test(/^(bigquery|clickhouse|ducklake|iceberg|snowflake)::/)) | \ (binary_id(etl-destinations) & \ test(/ducklake::core::tests::postgres_backed::/))"; From 2b6c747daf761b543927361c3581b44996bf14ce Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Fri, 22 May 2026 10:00:58 +0200 Subject: [PATCH 23/29] fix(core): Improve reset semantics for the table (#760) --- crates/etl-benchmarks/src/common.rs | 20 +- .../etl-destinations/src/bigquery/client.rs | 214 ++++++++++++++---- crates/etl-destinations/src/bigquery/core.rs | 128 ++++++++++- .../etl-destinations/src/clickhouse/client.rs | 21 ++ .../etl-destinations/src/clickhouse/core.rs | 22 +- crates/etl-destinations/src/ducklake/core.rs | 92 +++++++- crates/etl-destinations/src/iceberg/core.rs | 48 +++- .../tests/bigquery/pipeline.rs | 89 ++++++++ .../tests/clickhouse/pipeline.rs | 114 ++++++++++ .../tests/ducklake/pipeline.rs | 92 ++++++++ crates/etl-replicator/src/error_reporting.rs | 8 +- crates/etl/src/destination/async_result.rs | 8 +- crates/etl/src/destination/base.rs | 23 +- crates/etl/src/destination/mod.rs | 2 +- crates/etl/src/lib.rs | 8 +- crates/etl/src/pipeline.rs | 6 +- crates/etl/src/replication/apply.rs | 71 ++---- crates/etl/src/replication/table_cache.rs | 44 ++++ crates/etl/src/replication/table_sync.rs | 210 +++++++---------- crates/etl/src/store/both/memory.rs | 44 +++- crates/etl/src/store/both/postgres.rs | 134 ++++++----- crates/etl/src/store/cleanup.rs | 22 +- .../etl/src/test_utils/memory_destination.rs | 14 +- crates/etl/src/test_utils/notifying_store.rs | 44 +++- .../test_utils/test_destination_wrapper.rs | 66 +++--- crates/etl/src/workers/apply.rs | 4 +- crates/etl/src/workers/table_sync.rs | 8 +- crates/etl/tests/pipeline.rs | 37 +-- crates/etl/tests/postgres_store.rs | 115 ++++++++-- docs/explanation/architecture.md | 8 +- docs/explanation/traits.md | 37 ++- docs/guides/custom-implementations.md | 61 ++++- docs/guides/first-pipeline.md | 8 +- docs/index.md | 6 +- 34 files changed, 1366 insertions(+), 462 deletions(-) diff --git a/crates/etl-benchmarks/src/common.rs b/crates/etl-benchmarks/src/common.rs index 5d29cd689..1f5e224ee 100644 --- a/crates/etl-benchmarks/src/common.rs +++ b/crates/etl-benchmarks/src/common.rs @@ -15,7 +15,7 @@ use clap::{Args, ValueEnum}; use etl::{ destination::{ Destination, - async_result::{TruncateTableResult, WriteEventsResult, WriteTableRowsResult}, + async_result::{DropTableForCopyResult, WriteEventsResult, WriteTableRowsResult}, }, error::EtlResult, test_utils::notifying_store::NotifyingStore, @@ -347,12 +347,12 @@ where self.inner.shutdown().await } - async fn truncate_table( + async fn drop_table_for_copy( &self, replicated_table_schema: &ReplicatedTableSchema, - async_result: TruncateTableResult<()>, + async_result: DropTableForCopyResult<()>, ) -> EtlResult<()> { - self.inner.truncate_table(replicated_table_schema, async_result).await + self.inner.drop_table_for_copy(replicated_table_schema, async_result).await } async fn write_table_rows( @@ -401,10 +401,10 @@ impl Destination for NullDestination { "null" } - async fn truncate_table( + async fn drop_table_for_copy( &self, _replicated_table_schema: &ReplicatedTableSchema, - async_result: TruncateTableResult<()>, + async_result: DropTableForCopyResult<()>, ) -> EtlResult<()> { async_result.send(Ok(())); Ok(()) @@ -544,18 +544,18 @@ impl Destination for BenchDestination { } } - async fn truncate_table( + async fn drop_table_for_copy( &self, replicated_table_schema: &ReplicatedTableSchema, - async_result: TruncateTableResult<()>, + async_result: DropTableForCopyResult<()>, ) -> EtlResult<()> { match self { Self::Null(destination) => { - destination.truncate_table(replicated_table_schema, async_result).await + destination.drop_table_for_copy(replicated_table_schema, async_result).await } #[cfg(feature = "bigquery")] Self::BigQuery(destination) => { - destination.truncate_table(replicated_table_schema, async_result).await + destination.drop_table_for_copy(replicated_table_schema, async_result).await } } } diff --git a/crates/etl-destinations/src/bigquery/client.rs b/crates/etl-destinations/src/bigquery/client.rs index 7ee564b29..58295a1d9 100644 --- a/crates/etl-destinations/src/bigquery/client.rs +++ b/crates/etl-destinations/src/bigquery/client.rs @@ -13,7 +13,11 @@ use gcp_bigquery_client::{ cloud::bigquery::storage::v1::{RowError, StorageError, storage_error::StorageErrorCode}, rpc::Status as GoogleRpcStatus, }, - model::{query_request::QueryRequest, query_response::ResultSet}, + model::{ + query_parameter::QueryParameter, query_parameter_type::QueryParameterType, + query_parameter_value::QueryParameterValue, query_request::QueryRequest, + query_response::ResultSet, + }, storage::{ BatchAppendRequest, BatchAppendResult, ColumnMode, ColumnType, FieldDescriptor, StorageApiConfig, StreamName, TableBatch, TableDescriptor, @@ -45,15 +49,16 @@ const MAX_INFLIGHT_REQUESTS_PER_CONNECTION: usize = 100; /// This upper bound ensures reasonable memory usage and prevents overflow when /// computing max inflight requests from connection pool size. const MAX_SAFE_INFLIGHT_REQUESTS: usize = 100_000; -/// Maximum time to retry appends while BigQuery propagates a schema change. +/// Maximum time to retry writes while the BigQuery Storage Write API is still +/// using stale table metadata. /// /// Google documents schema update detection as happening on the order of /// minutes. -const SCHEMA_PROPAGATION_RETRY_TIMEOUT: Duration = Duration::from_secs(180); -/// Initial backoff when retrying appends during schema propagation. -const SCHEMA_PROPAGATION_RETRY_DELAY: Duration = Duration::from_secs(1); -/// Maximum backoff when retrying appends during schema propagation. -const SCHEMA_PROPAGATION_MAX_RETRY_DELAY: Duration = Duration::from_secs(15); +const STORAGE_WRITE_METADATA_LAG_RETRY_TIMEOUT: Duration = Duration::from_secs(180); +/// Initial backoff when retrying writes during storage write metadata lag. +const STORAGE_WRITE_METADATA_LAG_RETRY_DELAY: Duration = Duration::from_secs(1); +/// Maximum backoff when retrying writes during storage write metadata lag. +const STORAGE_WRITE_METADATA_LAG_MAX_RETRY_DELAY: Duration = Duration::from_secs(15); /// Protobuf type name for BigQuery storage errors embedded in gRPC status /// details. const BIGQUERY_STORAGE_ERROR_TYPE_NAME: &str = "google.cloud.bigquery.storage.v1.StorageError"; @@ -99,8 +104,8 @@ impl fmt::Display for BigQueryOperationType { enum BatchProcessResult { /// Batch succeeded with byte metrics. Success { bytes_sent: usize, bytes_received: usize }, - /// Batch hit schema propagation after DDL and should be retried. - RetryableSchemaPropagation { detail: String }, + /// Batch hit storage write metadata lag after DDL and should be retried. + RetryableStorageWriteMetadataLag { detail: String }, /// Batch had row-level errors. RowErrors { errors: Vec }, /// Batch had a request-level error. @@ -122,18 +127,21 @@ enum AppendProcessingResult { Error(EtlError), } -/// A batch append request that should be retried after schema propagation -/// finishes. +/// A batch append request that should be retried after Storage Write metadata +/// lag clears. #[derive(Debug)] struct RetryableAppendRequest { request: BatchAppendRequest, detail: String, } -/// Builds a concise description for a set of schema-propagation retries. -fn format_retryable_append_requests(requests: &[RetryableAppendRequest]) -> String { +/// Builds a concise description for a set of Storage Write metadata lag +/// retries. +fn format_retryable_storage_write_metadata_lag_requests( + requests: &[RetryableAppendRequest], +) -> String { match requests.split_first() { - None => "schema propagation error".to_owned(), + None => "storage write metadata lag error".to_owned(), Some((first, [])) => first.detail.clone(), Some((first, rest)) => { let distinct_other_details = @@ -237,9 +245,7 @@ fn append_processing_result_from_request_error( error: BQError, append_requests: Vec>, ) -> AppendProcessingResult { - if is_retryable_schema_propagation_error(&error) { - let detail = - bq_error_to_etl_error(error).detail().unwrap_or("schema propagation error").to_owned(); + if let Some(detail) = retryable_storage_write_metadata_lag_detail(&error) { AppendProcessingResult::Retry { pending_requests: append_requests .into_iter() @@ -253,19 +259,19 @@ fn append_processing_result_from_request_error( } } -/// Builds the error returned when local schema-propagation retries are +/// Builds the error returned when local Storage Write metadata lag retries are /// exhausted. /// -/// The destination absorbs the common short propagation delay locally. If -/// BigQuery still has not accepted the schema once that bounded window expires, -/// the worker-level timed retry policy should take over. -fn schema_propagation_timeout_error(detail: &str) -> EtlError { +/// The destination absorbs the common short lag window locally. If +/// BigQuery still has not accepted the storage write metadata once that bounded +/// window expires, the worker-level timed retry policy should take over. +fn storage_write_metadata_lag_timeout_error(detail: &str) -> EtlError { etl_error!( ErrorKind::DestinationAtomicBatchRetryable, - "BigQuery schema propagation timed out", + "BigQuery storage write metadata lag timed out", format!( - "BigQuery did not accept the updated schema within {} seconds after DDL: {}", - SCHEMA_PROPAGATION_RETRY_TIMEOUT.as_secs(), + "BigQuery did not accept the storage write metadata within {} seconds after DDL: {}", + STORAGE_WRITE_METADATA_LAG_RETRY_TIMEOUT.as_secs(), detail ) ) @@ -278,8 +284,8 @@ fn row_error_to_etl_error(err: RowError) -> EtlError { /// Converts a request-level append error into a [`BatchProcessResult`]. fn batch_process_result_from_request_error(error: BQError) -> BatchProcessResult { - if is_retryable_schema_propagation_error(&error) { - BatchProcessResult::RetryableSchemaPropagation { detail: error.to_string() } + if retryable_storage_write_metadata_lag_detail(&error).is_some() { + BatchProcessResult::RetryableStorageWriteMetadataLag { detail: error.to_string() } } else { BatchProcessResult::RequestError { error } } @@ -485,7 +491,7 @@ fn decode_storage_error_codes(status: &tonic::Status) -> Vec<&'static str> { } /// Returns true when the request-level BigQuery error matches the documented -/// schema propagation case. +/// storage write metadata lag case. /// /// BigQuery documents `StorageErrorCode::SCHEMA_MISMATCH_EXTRA_FIELDS` as the /// structured signal for schema mismatch during appends. We fall back to the @@ -515,6 +521,31 @@ fn is_retryable_schema_propagation_error(error: &BQError) -> bool { || message.contains("schema_mismatch_extra_fields") } +/// Returns true for BigQuery's transient default-stream error after a table is +/// dropped and recreated with the same name. +fn is_retryable_table_recreation_error(error: &BQError) -> bool { + let BQError::TonicStatusError(status) = error else { + return false; + }; + + status.code() == Code::NotFound + && status.message().to_ascii_lowercase().contains("is re-created") +} + +/// Returns retry detail when a Storage Write append failed due to BigQuery +/// metadata propagation after DDL. +fn retryable_storage_write_metadata_lag_detail(error: &BQError) -> Option { + if is_retryable_schema_propagation_error(error) { + return Some(error.to_string()); + } + + if is_retryable_table_recreation_error(error) { + return Some(error.to_string()); + } + + None +} + /// Client for interacting with Google BigQuery. /// /// Provides methods for table management, data insertion, and query execution @@ -769,10 +800,29 @@ impl BigQueryClient { Ok(()) } + /// Drops a view from BigQuery. + /// + /// Executes a DROP VIEW statement to remove the logical view if it exists. + pub async fn drop_view_if_exists( + &self, + dataset_id: &BigQueryDatasetId, + view_name: &BigQueryTableId, + ) -> EtlResult<()> { + let full_view_name = self.full_table_name(dataset_id, view_name)?; + + info!(%full_view_name, "dropping view 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. - pub async fn drop_table( + pub async fn drop_table_if_exists( &self, dataset_id: &BigQueryDatasetId, table_id: &BigQueryTableId, @@ -788,6 +838,53 @@ impl BigQueryClient { Ok(()) } + /// Lists physical sequenced table ids for a base table. + /// + /// Queries `INFORMATION_SCHEMA.TABLES` instead of using the destination's + /// local cache so reset cleanup can remove versions left behind by earlier + /// processes. + pub async fn list_sequenced_table_ids( + &self, + dataset_id: &BigQueryDatasetId, + base_table_id: &BigQueryTableId, + ) -> EtlResult> { + info!(%dataset_id, %base_table_id, "listing sequenced tables from bigquery"); + + let project_id = Self::sanitize_identifier(&self.project_id, "BigQuery project id")?; + let dataset_id = Self::sanitize_identifier(dataset_id, "BigQuery dataset id")?; + let query = format!( + "select table_name from `{project_id}.{dataset_id}.INFORMATION_SCHEMA.TABLES` where \ + table_type = 'BASE TABLE' and starts_with(table_name, @table_name_prefix) order by \ + table_name" + ); + let mut request = QueryRequest::new(query); + request.parameter_mode = Some("NAMED".to_owned()); + request.query_parameters = Some(vec![QueryParameter { + name: Some("table_name_prefix".to_owned()), + parameter_type: Some(QueryParameterType { + r#type: "STRING".to_owned(), + ..Default::default() + }), + parameter_value: Some(QueryParameterValue { + value: Some(format!("{base_table_id}_")), + ..Default::default() + }), + }]); + + let mut result_set = self.query(request).await?; + let mut table_ids = Vec::new(); + + while result_set.next_row() { + if let Some(table_id) = + result_set.get_string_by_name("table_name").map_err(bq_error_to_etl_error)? + { + table_ids.push(table_id); + } + } + + Ok(table_ids) + } + /// Adds a column to an existing BigQuery table. /// /// Executes an ALTER TABLE ADD COLUMN statement to add a new column with @@ -904,7 +1001,7 @@ impl BigQueryClient { /// /// Retries for transient request and transport failures are handled inside /// the underlying Storage Write API library. This method also retries - /// the narrow class of schema propagation failures that can happen + /// the narrow class of storage write metadata lag failures that can happen /// after DDL, then converts final failures into ETL errors. pub(super) async fn append_table_batches( &self, @@ -920,7 +1017,7 @@ impl BigQueryClient { let started_at = Instant::now(); let mut attempt = 1; - let mut retry_delay = SCHEMA_PROPAGATION_RETRY_DELAY; + let mut retry_delay = STORAGE_WRITE_METADATA_LAG_RETRY_DELAY; loop { match self.append_table_batches_once(pending_requests).await? { @@ -938,7 +1035,9 @@ impl BigQueryClient { total_bytes_sent += bytes_sent; total_bytes_received += bytes_received; - let retry_summary = format_retryable_append_requests(&next_pending_requests); + let retry_summary = format_retryable_storage_write_metadata_lag_requests( + &next_pending_requests, + ); pending_requests = next_pending_requests.into_iter().map(|request| request.request).collect(); @@ -949,16 +1048,16 @@ impl BigQueryClient { let elapsed = started_at.elapsed(); let remaining_timeout = - SCHEMA_PROPAGATION_RETRY_TIMEOUT.saturating_sub(elapsed); + STORAGE_WRITE_METADATA_LAG_RETRY_TIMEOUT.saturating_sub(elapsed); if remaining_timeout.is_zero() { - return Err(schema_propagation_timeout_error(&retry_summary)); + return Err(storage_write_metadata_lag_timeout_error(&retry_summary)); } let sleep_delay = retry_delay.min(remaining_timeout); if sleep_delay.is_zero() { - return Err(schema_propagation_timeout_error(&retry_summary)); + return Err(storage_write_metadata_lag_timeout_error(&retry_summary)); } warn!( @@ -966,12 +1065,12 @@ impl BigQueryClient { pending_batch_count, retry_delay_ms = sleep_delay.as_millis() as u64, error_detail = %retry_summary, - "bigquery schema change still propagating, retrying append" + "bigquery storage write metadata still lagging, retrying append" ); sleep(sleep_delay).await; - retry_delay = (retry_delay * 2).min(SCHEMA_PROPAGATION_MAX_RETRY_DELAY); + retry_delay = (retry_delay * 2).min(STORAGE_WRITE_METADATA_LAG_MAX_RETRY_DELAY); attempt += 1; } AppendProcessingResult::Error(error) => return Err(error), @@ -1025,7 +1124,7 @@ impl BigQueryClient { total_bytes_sent += bytes_sent; total_bytes_received += bytes_received; } - BatchProcessResult::RetryableSchemaPropagation { detail } => { + BatchProcessResult::RetryableStorageWriteMetadataLag { detail } => { retryable_batch_details[batch_index] = Some(detail); } BatchProcessResult::RowErrors { errors: row_errors } => { @@ -1702,7 +1801,7 @@ mod tests { bytes_sent: 128, }); - assert!(matches!(result, BatchProcessResult::RetryableSchemaPropagation { .. })); + assert!(matches!(result, BatchProcessResult::RetryableStorageWriteMetadataLag { .. })); } #[test] @@ -1716,7 +1815,7 @@ mod tests { bytes_sent: 128, }); - assert!(matches!(result, BatchProcessResult::RetryableSchemaPropagation { .. })); + assert!(matches!(result, BatchProcessResult::RetryableStorageWriteMetadataLag { .. })); } #[test] @@ -1730,14 +1829,41 @@ mod tests { bytes_sent: 128, }); - assert!(matches!(result, BatchProcessResult::RetryableSchemaPropagation { .. })); + assert!(matches!(result, BatchProcessResult::RetryableStorageWriteMetadataLag { .. })); + } + + #[test] + fn process_single_batch_append_result_retries_table_recreation_propagation() { + let result = process_single_batch_append_result(BatchAppendResult { + batch_index: 0, + responses: vec![Err(tonic::Status::not_found( + "Table 123:dataset.test_users_0 is re-created. Entity: \ + projects/project/datasets/dataset/tables/test_users_0/streams/_default", + ))], + bytes_sent: 128, + }); + + assert!(matches!(result, BatchProcessResult::RetryableStorageWriteMetadataLag { .. })); + } + + #[test] + fn process_single_batch_append_result_does_not_retry_generic_not_found() { + let result = process_single_batch_append_result(BatchAppendResult { + batch_index: 0, + responses: vec![Err(tonic::Status::not_found( + "Table 123:dataset.test_users_0 was not found.", + ))], + bytes_sent: 128, + }); + + assert!(matches!(result, BatchProcessResult::RequestError { .. })); } #[test] - fn schema_propagation_timeout_error_is_worker_retryable() { - let error = schema_propagation_timeout_error("schema lag"); + fn storage_write_metadata_lag_timeout_error_is_worker_retryable() { + let error = storage_write_metadata_lag_timeout_error("storage write metadata lag"); assert_eq!(error.kind(), ErrorKind::DestinationAtomicBatchRetryable); - assert_eq!(error.description(), Some("BigQuery schema propagation timed out")); + assert_eq!(error.description(), Some("BigQuery storage write metadata lag timed out")); } } diff --git a/crates/etl-destinations/src/bigquery/core.rs b/crates/etl-destinations/src/bigquery/core.rs index 8c470fdcb..5fa5e1f39 100644 --- a/crates/etl-destinations/src/bigquery/core.rs +++ b/crates/etl-destinations/src/bigquery/core.rs @@ -1,7 +1,6 @@ use std::{ collections::{HashMap, HashSet}, fmt::Display, - iter, str::FromStr, sync::Arc, }; @@ -11,7 +10,7 @@ use etl::{ concurrency::TaskSet, destination::{ Destination, - async_result::{TruncateTableResult, WriteEventsResult, WriteTableRowsResult}, + async_result::{DropTableForCopyResult, WriteEventsResult, WriteTableRowsResult}, }, error::{ErrorKind, EtlError, EtlResult}, etl_error, @@ -76,6 +75,16 @@ impl SequencedBigQueryTableId { fn to_bigquery_table_id(&self) -> BigQueryTableId { self.0.clone() } + + /// Returns whether this sequenced table belongs to `base_table_id`. + fn belongs_to_base(&self, base_table_id: &BigQueryTableId) -> bool { + &self.0 == base_table_id + } + + /// Parses a sequenced table id only when it belongs to `base_table_id`. + fn parse_for_base(table_id: &str, base_table_id: &BigQueryTableId) -> Option { + table_id.parse::().ok().filter(|table_id| table_id.belongs_to_base(base_table_id)) + } } impl FromStr for SequencedBigQueryTableId { @@ -170,6 +179,12 @@ impl Inner { fn new() -> Self { Self { created_tables: HashSet::new(), created_views: HashMap::new() } } + + /// Clears cached state for a destination table. + fn clear_table_cache(&mut self, base_table_id: &BigQueryTableId) { + self.created_views.remove(base_table_id); + self.created_tables.retain(|table_id| !table_id.belongs_to_base(base_table_id)); + } } /// A BigQuery destination that implements the ETL [`Destination`] trait. @@ -1076,7 +1091,7 @@ where self.tasks .spawn(async move { if let Err(err) = client - .drop_table(&dataset_id, &sequenced_bigquery_table_id.to_string()) + .drop_table_if_exists(&dataset_id, &sequenced_bigquery_table_id.to_string()) .await { warn!( @@ -1096,6 +1111,46 @@ where Ok(()) } + + /// Drops destination table objects before restarting a table copy. + async fn drop_table_for_copy_inner( + &self, + replicated_table_schema: &ReplicatedTableSchema, + ) -> EtlResult<()> { + let base_bigquery_table_id = + table_name_to_bigquery_table_id(replicated_table_schema.name())?; + let table_id = replicated_table_schema.id(); + let mut table_ids = HashSet::new(); + + // Discover physical table versions from BigQuery instead of the local cache. + // Older processes may have created versions this process never cached. + let listed_table_ids = + self.client.list_sequenced_table_ids(&self.dataset_id, &base_bigquery_table_id).await?; + let listed_sequenced_table_ids = + listed_table_ids.into_iter().filter_map(|listed_table_id| { + SequencedBigQueryTableId::parse_for_base(&listed_table_id, &base_bigquery_table_id) + }); + table_ids.extend(listed_sequenced_table_ids); + + // We first drop the view, so that the table is not queriable anymore. + self.client.drop_view_if_exists(&self.dataset_id, &base_bigquery_table_id).await?; + + // We then drop each table individually. If any of these operations fail, on the + // next restart the tables will be cleaned up. + for sequenced_bigquery_table_id in &table_ids { + self.client + .drop_table_if_exists(&self.dataset_id, &sequenced_bigquery_table_id.to_string()) + .await?; + } + + // Once destination cleanup is done, remove any stale local cache entries. + let mut inner = self.inner.lock().await; + inner.clear_table_cache(&base_bigquery_table_id); + + info!(table_id = table_id.0, "dropped bigquery table before copy"); + + Ok(()) + } } /// Validates that a replicated table schema can be applied in BigQuery. @@ -1170,15 +1225,14 @@ where self.tasks.shutdown().await } - async fn truncate_table( + async fn drop_table_for_copy( &self, replicated_table_schema: &ReplicatedTableSchema, - async_result: TruncateTableResult<()>, + async_result: DropTableForCopyResult<()>, ) -> EtlResult<()> { self.tasks.try_reap().await?; - let result = - self.process_truncate_for_schemas(iter::once(replicated_table_schema.clone())).await; + let result = self.drop_table_for_copy_inner(replicated_table_schema).await; async_result.send(result); Ok(()) @@ -1768,6 +1822,66 @@ mod tests { assert_eq!(table_id.to_bigquery_table_id(), "simple_table"); } + #[test] + fn sequenced_bigquery_table_id_belongs_to_base() { + let base_table_id = "users_table".to_owned(); + let users_table_id = SequencedBigQueryTableId(base_table_id.clone(), 2); + let orders_table_id = SequencedBigQueryTableId("orders_table".to_owned(), 2); + + assert!(users_table_id.belongs_to_base(&base_table_id)); + assert!(!orders_table_id.belongs_to_base(&base_table_id)); + } + + #[test] + fn sequenced_bigquery_table_id_parse_for_base_ignores_unrelated_tables() { + let base_table_id = "users_table".to_owned(); + + assert_eq!( + SequencedBigQueryTableId::parse_for_base("users_table_2", &base_table_id), + Some(SequencedBigQueryTableId("users_table".to_owned(), 2)) + ); + assert_eq!( + SequencedBigQueryTableId::parse_for_base("users_table_backup", &base_table_id), + None + ); + assert_eq!( + SequencedBigQueryTableId::parse_for_base("orders_table_2", &base_table_id), + None + ); + assert_eq!(SequencedBigQueryTableId::parse_for_base("users_table", &base_table_id), None); + } + + #[test] + fn clear_table_cache_removes_all_cached_table_state_for_base() { + let base_table_id = "users_table".to_owned(); + let mut inner = Inner::new(); + + inner.created_tables.insert(SequencedBigQueryTableId(base_table_id.clone(), 0)); + inner.created_tables.insert(SequencedBigQueryTableId(base_table_id.clone(), 1)); + inner.created_tables.insert(SequencedBigQueryTableId("orders_table".to_owned(), 0)); + inner + .created_views + .insert(base_table_id.clone(), SequencedBigQueryTableId(base_table_id.clone(), 1)); + inner.created_views.insert( + "orders_table".to_owned(), + SequencedBigQueryTableId("orders_table".to_owned(), 0), + ); + + inner.clear_table_cache(&base_table_id); + + assert_eq!( + inner.created_tables, + HashSet::from([SequencedBigQueryTableId("orders_table".to_owned(), 0)]) + ); + assert_eq!( + inner.created_views, + HashMap::from([( + "orders_table".to_owned(), + SequencedBigQueryTableId("orders_table".to_owned(), 0), + )]) + ); + } + #[test] fn validate_bigquery_replica_identity_accepts_primary_key() { let replicated_table_schema = replicated_schema(IdentityType::PrimaryKey); diff --git a/crates/etl-destinations/src/clickhouse/client.rs b/crates/etl-destinations/src/clickhouse/client.rs index 29dff62b1..0cfdc9594 100644 --- a/crates/etl-destinations/src/clickhouse/client.rs +++ b/crates/etl-destinations/src/clickhouse/client.rs @@ -130,6 +130,12 @@ fn build_truncate_table_sql(table_name: &str) -> String { format!("TRUNCATE TABLE IF EXISTS {table_name}") } +/// Builds the SQL used to drop a ClickHouse table. +fn build_drop_table_sql(table_name: &str) -> String { + let table_name = quote_identifier(table_name); + format!("DROP TABLE IF EXISTS {table_name}") +} + /// Builds the SQL used to insert RowBinary rows into a ClickHouse table. fn build_insert_rows_sql(table_name: &str) -> String { let table_name = quote_identifier(table_name); @@ -143,6 +149,7 @@ fn build_insert_rows_sql(table_name: &str) -> String { pub(crate) enum DdlKind { CreateTable, CreateView, + DropTable, DropView, AddColumn, DropColumn, @@ -154,6 +161,7 @@ impl DdlKind { match self { DdlKind::CreateTable => "create_table", DdlKind::CreateView => "create_view", + DdlKind::DropTable => "drop_table", DdlKind::DropView => "drop_view", DdlKind::AddColumn => "add_column", DdlKind::DropColumn => "drop_column", @@ -370,6 +378,12 @@ impl ClickHouseClient { .await } + /// Executes `DROP TABLE IF EXISTS` for the supplied table. + pub(crate) async fn drop_table(&self, table_name: &str) -> EtlResult<()> { + let sql = build_drop_table_sql(table_name); + self.execute_ddl(DdlKind::DropTable, &sql).await + } + /// Inserts `rows` into `table_name` using the RowBinary format. /// /// Each element of `rows` is a complete, already-encoded row of @@ -496,6 +510,13 @@ mod tests { assert_eq!(sql, "TRUNCATE TABLE IF EXISTS \"table\"\"name\""); } + #[test] + fn drop_table_sql_quotes_identifiers() { + let sql = build_drop_table_sql("table\"name"); + + assert_eq!(sql, "DROP TABLE IF EXISTS \"table\"\"name\""); + } + #[test] fn insert_rows_sql_quotes_identifiers() { let sql = build_insert_rows_sql("table\"name"); diff --git a/crates/etl-destinations/src/clickhouse/core.rs b/crates/etl-destinations/src/clickhouse/core.rs index bece73d7c..feeb4d8e4 100644 --- a/crates/etl-destinations/src/clickhouse/core.rs +++ b/crates/etl-destinations/src/clickhouse/core.rs @@ -3,7 +3,7 @@ use std::{collections::HashMap, sync::Arc, time::Duration}; use etl::{ destination::{ Destination, - async_result::{TruncateTableResult, WriteEventsResult, WriteTableRowsResult}, + async_result::{DropTableForCopyResult, WriteEventsResult, WriteTableRowsResult}, }, error::{ErrorKind, EtlResult}, etl_error, @@ -605,6 +605,20 @@ where self.client.truncate_table(&clickhouse_table_name).await } + async fn drop_table_for_copy_inner(&self, schema: &ReplicatedTableSchema) -> EtlResult<()> { + let clickhouse_table_name = try_stringify_table_name(schema.name())?; + + if matches!(self.inserter_config.engine, ClickHouseEngine::ReplacingMergeTree) { + let drop_view = drop_current_view_sql(&clickhouse_table_name); + self.client.execute_ddl(DdlKind::DropView, &drop_view).await?; + } + + self.client.drop_table(&clickhouse_table_name).await?; + self.table_cache.write().remove(&clickhouse_table_name); + + Ok(()) + } + async fn write_table_rows_inner( &self, schema: &ReplicatedTableSchema, @@ -1259,12 +1273,12 @@ where // sending" error if the path ever skips `send`, so the receiver is never // silently abandoned. - async fn truncate_table( + async fn drop_table_for_copy( &self, replicated_table_schema: &ReplicatedTableSchema, - async_result: TruncateTableResult<()>, + async_result: DropTableForCopyResult<()>, ) -> EtlResult<()> { - let result = self.truncate_table_inner(replicated_table_schema).await; + let result = self.drop_table_for_copy_inner(replicated_table_schema).await; async_result.send(result); Ok(()) } diff --git a/crates/etl-destinations/src/ducklake/core.rs b/crates/etl-destinations/src/ducklake/core.rs index a32de35b8..d7c1df982 100644 --- a/crates/etl-destinations/src/ducklake/core.rs +++ b/crates/etl-destinations/src/ducklake/core.rs @@ -12,7 +12,7 @@ use etl::{ concurrency::TaskSet, destination::{ Destination, - async_result::{TruncateTableResult, WriteEventsResult, WriteTableRowsResult}, + async_result::{DropTableForCopyResult, WriteEventsResult, WriteTableRowsResult}, }, error::{ErrorKind, EtlResult}, etl_error, @@ -222,12 +222,12 @@ where Ok(()) } - async fn truncate_table( + async fn drop_table_for_copy( &self, replicated_table_schema: &ReplicatedTableSchema, - async_result: TruncateTableResult<()>, + async_result: DropTableForCopyResult<()>, ) -> EtlResult<()> { - let result = self.truncate_table(replicated_table_schema).await; + let result = self.drop_table_for_copy_inner(replicated_table_schema).await; async_result.send(result); Ok(()) @@ -759,6 +759,83 @@ where .await } + /// Drops the destination table and replay markers before restarting a copy. + async fn drop_table_for_copy_inner( + &self, + replicated_table_schema: &ReplicatedTableSchema, + ) -> EtlResult<()> { + let table_name = self.resolve_destination_table_name(replicated_table_schema).await?; + let _table_write_permit = self.acquire_table_write_slot(&table_name).await?; + self.ensure_applied_batches_table_exists().await?; + self.ensure_streaming_progress_table_exists().await?; + let _checkpoint_guard = self.acquire_mutation_guard().await; + let table_name_for_drop = table_name.clone(); + + self.run_duckdb_blocking(move |conn| -> EtlResult<()> { + conn.execute_batch("BEGIN TRANSACTION").map_err(|e| { + etl_error!( + ErrorKind::DestinationQueryFailed, + "DuckLake BEGIN TRANSACTION failed", + source: e + ) + })?; + + let result = (|| -> EtlResult<()> { + let quoted_table_name = quote_identifier(&table_name_for_drop).into_owned(); + let drop_table_sql = + format!("DROP TABLE IF EXISTS {LAKE_CATALOG}.{quoted_table_name};"); + conn.execute_batch(&drop_table_sql).map_err(|e| { + etl_error!( + ErrorKind::DestinationQueryFailed, + "DuckLake DROP TABLE failed", + format_query_error_detail(&drop_table_sql), + source: e + ) + })?; + + clear_applied_batch_markers_for_kind( + conn, + &table_name_for_drop, + DuckLakeTableBatchKind::Copy, + )?; + clear_applied_batch_markers_for_kind( + conn, + &table_name_for_drop, + DuckLakeTableBatchKind::Mutation, + )?; + clear_applied_batch_markers_for_kind( + conn, + &table_name_for_drop, + DuckLakeTableBatchKind::Truncate, + )?; + clear_table_streaming_progress(conn, &table_name_for_drop)?; + Ok(()) + })(); + + match result { + Ok(()) => conn.execute_batch("COMMIT").map_err(|e| { + etl_error!( + ErrorKind::DestinationQueryFailed, + "DuckLake COMMIT failed", + source: e + ) + }), + Err(error) => { + let err = conn.execute_batch("ROLLBACK"); + if let Err(err) = err { + tracing::error!(error = %err, "error rollback"); + } + Err(error) + } + } + }) + .await?; + + self.created_tables.lock().remove(&table_name); + + Ok(()) + } + /// Bulk-inserts rows into the destination table inside a single /// transaction. /// @@ -1112,7 +1189,7 @@ where replicated_table_schema: &ReplicatedTableSchema, ) -> EtlResult { let table_id = replicated_table_schema.id(); - let table_name = self.get_or_create_destination_table_name(replicated_table_schema).await?; + let table_name = self.resolve_destination_table_name(replicated_table_schema).await?; // Fast path: already created. { @@ -1225,9 +1302,8 @@ where .await } - /// Returns the stored destination table name for `table_id`, creating a - /// default name if none exists yet. - async fn get_or_create_destination_table_name( + /// Returns the stored destination table name or the deterministic default. + async fn resolve_destination_table_name( &self, replicated_table_schema: &ReplicatedTableSchema, ) -> EtlResult { diff --git a/crates/etl-destinations/src/iceberg/core.rs b/crates/etl-destinations/src/iceberg/core.rs index f4ed1fd9d..f27c20ea2 100644 --- a/crates/etl-destinations/src/iceberg/core.rs +++ b/crates/etl-destinations/src/iceberg/core.rs @@ -8,7 +8,7 @@ use etl::{ concurrency::TaskSet, destination::{ Destination, - async_result::{TruncateTableResult, WriteEventsResult, WriteTableRowsResult}, + async_result::{DropTableForCopyResult, WriteEventsResult, WriteTableRowsResult}, }, error::{ErrorKind, EtlResult}, etl_error, @@ -230,6 +230,40 @@ where Ok(()) } + /// Drops an Iceberg table before restarting a table copy. + async fn drop_table_for_copy_inner( + &self, + replicated_table_schema: &ReplicatedTableSchema, + ) -> EtlResult<()> { + let table_id = replicated_table_schema.id(); + let table_name = replicated_table_schema.name(); + let table_namespace = schema_to_namespace(&table_name.schema); + let (default_table_name, namespace) = { + let inner = self.inner.lock().await; + let default_table_name = + table_name_to_iceberg_table_name(table_name, inner.namespace.is_single())?; + let namespace = inner.namespace.get_or(&table_namespace).to_owned(); + + (default_table_name, namespace) + }; + let iceberg_table_name = + if let Some(metadata) = self.store.get_destination_table_metadata(table_id).await? { + metadata.destination_table_id + } else { + default_table_name + }; + + self.client + .drop_table_if_exists(&namespace, iceberg_table_name.clone()) + .await + .map_err(iceberg_error_to_etl_error)?; + + let mut inner = self.inner.lock().await; + inner.created_tables.remove(&iceberg_table_name); + + Ok(()) + } + /// Writes table-copy rows to the Iceberg destination as insert entries. /// /// Prepares the target table for streaming, augments each row with CDC @@ -593,16 +627,16 @@ where self.tasks.shutdown().await } - /// Truncates the specified table by dropping and recreating it. + /// Drops the specified table before a fresh copy. /// - /// Removes all data from the target Iceberg table while preserving - /// the table schema structure for continued CDC operations. - async fn truncate_table( + /// The table sync worker clears ETL metadata after this succeeds, and the + /// next copy recreates the table from the fresh `0/0` schema. + async fn drop_table_for_copy( &self, replicated_table_schema: &ReplicatedTableSchema, - async_result: TruncateTableResult<()>, + async_result: DropTableForCopyResult<()>, ) -> EtlResult<()> { - let result = IcebergDestination::truncate_table(self, replicated_table_schema).await; + let result = self.drop_table_for_copy_inner(replicated_table_schema).await; async_result.send(result); Ok(()) diff --git a/crates/etl-destinations/tests/bigquery/pipeline.rs b/crates/etl-destinations/tests/bigquery/pipeline.rs index 8528e8997..94bb8c092 100644 --- a/crates/etl-destinations/tests/bigquery/pipeline.rs +++ b/crates/etl-destinations/tests/bigquery/pipeline.rs @@ -213,6 +213,95 @@ async fn table_copy_and_streaming_with_restart() { ); } +#[tokio::test(flavor = "multi_thread")] +async fn table_copy_reset_drops_destination_table_before_recopy() { + if skip_if_missing_bigquery_env_vars() { + return; + } + + init_test_tracing(); + install_crypto_provider(); + + let database = spawn_source_database().await; + let database_schema = setup_test_database_schema(&database, TableSelection::UsersOnly).await; + let bigquery_database = setup_bigquery_database().await; + let users_schema = database_schema.users_schema(); + let store = NotifyingStore::new(); + let pipeline_id: PipelineId = random(); + + database + .insert_values(users_schema.name.clone(), &["name", "age"], &[&"before_1", &1]) + .await + .unwrap(); + database + .insert_values(users_schema.name.clone(), &["name", "age"], &[&"before_2", &2]) + .await + .unwrap(); + + let raw_destination = bigquery_database.build_destination(pipeline_id, 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_ready = + store.notify_on_table_state_type(users_schema.id, TableReplicationPhaseType::Ready).await; + + pipeline.start().await.unwrap(); + + users_ready.notified().await; + + pipeline.shutdown_and_wait().await.unwrap(); + + let users_rows = bigquery_database.query_table(users_schema.name.clone()).await.unwrap(); + assert_eq!( + parse_bigquery_table_rows::(users_rows), + vec![BigQueryUser::new(1, "before_1", 1), BigQueryUser::new(2, "before_2", 2),] + ); + + database + .run_sql(&format!("delete from {} where true", users_schema.name.as_quoted_identifier())) + .await + .unwrap(); + database + .insert_values(users_schema.name.clone(), &["name", "age"], &[&"after", &3]) + .await + .unwrap(); + store.reset_table_state(users_schema.id).await.unwrap(); + + let raw_destination = bigquery_database.build_destination(pipeline_id, 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_ready = + store.notify_on_table_state_type(users_schema.id, TableReplicationPhaseType::Ready).await; + + pipeline.start().await.unwrap(); + + users_ready.notified().await; + + pipeline.shutdown_and_wait().await.unwrap(); + + assert!(destination.was_table_dropped_for_copy(users_schema.id).await); + let users_rows = bigquery_database.query_table(users_schema.name).await.unwrap(); + assert_eq!( + parse_bigquery_table_rows::(users_rows), + vec![BigQueryUser::new(3, "after", 3)] + ); +} + #[tokio::test(flavor = "multi_thread")] async fn table_insert_update_delete() { if skip_if_missing_bigquery_env_vars() { diff --git a/crates/etl-destinations/tests/clickhouse/pipeline.rs b/crates/etl-destinations/tests/clickhouse/pipeline.rs index 655edf010..de3335c94 100644 --- a/crates/etl-destinations/tests/clickhouse/pipeline.rs +++ b/crates/etl-destinations/tests/clickhouse/pipeline.rs @@ -52,6 +52,7 @@ const ID_VALUE_PROJECTION: &str = "id, value"; const UPDATE_FLOW_TABLE: &str = "test_update__flow"; const DELETE_FLOW_TABLE: &str = "test_delete__flow"; const RESTART_FLOW_TABLE: &str = "test_restart__flow"; +const RESET_COPY_TABLE: &str = "test_reset__copy"; const TRUNCATE_FLOW_TABLE: &str = "test_truncate__flow"; /// Days from 1970-01-01 to 2024-01-15 (used to verify the `date_col` @@ -930,6 +931,119 @@ async fn pipeline_restart_resumes_streaming_inner(engine: ClickHouseEngine) { assert_eq!(rows[1].value, "after_restart"); } +#[tokio::test(flavor = "multi_thread")] +async fn table_copy_reset_drops_destination_table_before_recopy_merge_tree() { + table_copy_reset_drops_destination_table_before_recopy_inner(ClickHouseEngine::MergeTree).await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn table_copy_reset_drops_destination_table_before_recopy_replacing_merge_tree() { + table_copy_reset_drops_destination_table_before_recopy_inner( + ClickHouseEngine::ReplacingMergeTree, + ) + .await; +} + +async fn table_copy_reset_drops_destination_table_before_recopy_inner(engine: ClickHouseEngine) { + init_test_tracing(); + install_crypto_provider(); + + let database = spawn_source_database().await; + let table_name = test_table_name("reset_copy"); + + let table_id = database + .create_table(table_name.clone(), true, &[("value", "text not null")]) + .await + .expect("Failed to create reset_copy test table"); + + let publication_name = "test_pub_clickhouse_reset_copy"; + database + .create_publication(publication_name, std::slice::from_ref(&table_name)) + .await + .expect("Failed to create reset_copy publication"); + + database + .run_sql(&format!( + "INSERT INTO {} (value) VALUES ('before_1'), ('before_2')", + table_name.as_quoted_identifier(), + )) + .await + .expect("Failed to insert initial reset_copy rows"); + + let clickhouse_db = setup_clickhouse_database().await; + let store = NotifyingStore::new(); + let pipeline_id: PipelineId = random(); + let reset_query = + || current_state_query(engine, RESET_COPY_TABLE, ID_VALUE_PROJECTION, &["id"], "id"); + + let destination = TestDestinationWrapper::wrap( + clickhouse_db.build_destination_with_engine(store.clone(), engine).await, + ); + + let mut pipeline = create_pipeline( + &database.config, + pipeline_id, + publication_name.to_owned(), + store.clone(), + destination, + ); + + let table_ready = + store.notify_on_table_state_type(table_id, TableReplicationPhaseType::Ready).await; + + pipeline.start().await.unwrap(); + + table_ready.notified().await; + + pipeline.shutdown_and_wait().await.unwrap(); + + let rows: Vec = clickhouse_db.query(&reset_query()).await; + assert_eq!(rows.len(), 2, "first copy should produce two rows"); + assert_eq!(rows[0].value, "before_1"); + assert_eq!(rows[1].value, "before_2"); + + database + .run_sql(&format!("DELETE FROM {} WHERE true", table_name.as_quoted_identifier())) + .await + .expect("Failed to delete reset_copy source rows"); + database + .run_sql(&format!( + "INSERT INTO {} (value) VALUES ('after')", + table_name.as_quoted_identifier(), + )) + .await + .expect("Failed to insert recopy reset_copy row"); + + store.reset_table_state(table_id).await.unwrap(); + + let destination = TestDestinationWrapper::wrap( + clickhouse_db.build_destination_with_engine(store.clone(), engine).await, + ); + + let mut pipeline = create_pipeline( + &database.config, + pipeline_id, + publication_name.to_owned(), + store.clone(), + destination.clone(), + ); + + let table_ready = + store.notify_on_table_state_type(table_id, TableReplicationPhaseType::Ready).await; + + pipeline.start().await.unwrap(); + + table_ready.notified().await; + + pipeline.shutdown_and_wait().await.unwrap(); + + assert!(destination.was_table_dropped_for_copy(table_id).await); + let rows: Vec = clickhouse_db.query(&reset_query()).await; + assert_eq!(rows.len(), 1, "recopy should not retain rows from the first copy"); + assert_eq!(rows[0].id, 3); + assert_eq!(rows[0].value, "after"); +} + /// Tests that TRUNCATE clears the ClickHouse table and that subsequent inserts /// produce a clean slate with only post-truncate data. #[tokio::test(flavor = "multi_thread")] diff --git a/crates/etl-destinations/tests/ducklake/pipeline.rs b/crates/etl-destinations/tests/ducklake/pipeline.rs index 11f75af8d..04dcb97d2 100644 --- a/crates/etl-destinations/tests/ducklake/pipeline.rs +++ b/crates/etl-destinations/tests/ducklake/pipeline.rs @@ -236,6 +236,98 @@ async fn table_copy_and_streaming_with_restart() { ); } +#[tokio::test(flavor = "multi_thread")] +async fn table_copy_reset_drops_destination_table_before_recopy() { + init_test_tracing(); + + let database = spawn_source_database().await; + let database_schema = setup_test_database_schema(&database, TableSelection::UsersOnly).await; + let lake = create_test_lake("table_copy_reset_drops_destination_table_before_recopy").await; + let catalog_url = lake.catalog_url.clone(); + let data_url = lake.data_url.clone(); + + let users_schema = database_schema.users_schema(); + let users_table_name = table_name_to_ducklake_table_name(&users_schema.name) + .expect("failed to build DuckLake users table name"); + let store = NotifyingStore::new(); + let pipeline_id: PipelineId = random(); + + database + .insert_values(users_schema.name.clone(), &["name", "age"], &[&"before_1", &1]) + .await + .unwrap(); + database + .insert_values(users_schema.name.clone(), &["name", "age"], &[&"before_2", &2]) + .await + .unwrap(); + + let destination = build_destination(&catalog_url, &data_url, store.clone()).await; + + let mut pipeline = create_pipeline( + &database.config, + pipeline_id, + database_schema.publication_name(), + store.clone(), + destination.clone(), + ); + + let users_ready = + store.notify_on_table_state_type(users_schema.id, TableReplicationPhaseType::Ready).await; + + pipeline.start().await.unwrap(); + + users_ready.notified().await; + + pipeline.shutdown_and_wait().await.unwrap(); + + checkpoint_lake(&catalog_url, &data_url); + + let conn = open_lake_conn(&catalog_url, &data_url); + assert_eq!( + query_user_rows(&conn, &users_table_name), + vec![(1, "before_1".to_owned(), 1), (2, "before_2".to_owned(), 2)] + ); + drop(conn); + drop(destination); + checkpoint_lake(&catalog_url, &data_url); + + database + .run_sql(&format!("delete from {} where true", users_schema.name.as_quoted_identifier())) + .await + .unwrap(); + database + .insert_values(users_schema.name.clone(), &["name", "age"], &[&"after", &3]) + .await + .unwrap(); + store.reset_table_state(users_schema.id).await.unwrap(); + + let destination = build_destination(&catalog_url, &data_url, store.clone()).await; + + let mut pipeline = create_pipeline( + &database.config, + pipeline_id, + database_schema.publication_name(), + store.clone(), + destination.clone(), + ); + + let users_ready = + store.notify_on_table_state_type(users_schema.id, TableReplicationPhaseType::Ready).await; + + pipeline.start().await.unwrap(); + + users_ready.notified().await; + + pipeline.shutdown_and_wait().await.unwrap(); + + assert!(destination.was_table_dropped_for_copy(users_schema.id).await); + drop(destination); + checkpoint_lake(&catalog_url, &data_url); + + let conn = open_lake_conn(&catalog_url, &data_url); + assert_eq!(query_user_rows(&conn, &users_table_name), vec![(3, "after".to_owned(), 3)]); +} + #[tokio::test(flavor = "multi_thread")] async fn table_copy_and_streaming_without_restart() { init_test_tracing(); diff --git a/crates/etl-replicator/src/error_reporting.rs b/crates/etl-replicator/src/error_reporting.rs index 538d97dd0..0c0fbfaa9 100644 --- a/crates/etl-replicator/src/error_reporting.rs +++ b/crates/etl-replicator/src/error_reporting.rs @@ -201,7 +201,11 @@ impl CleanupStore for ErrorReportingStateStore where S: CleanupStore + Send + Sync, { - async fn cleanup_table_state(&self, table_id: TableId) -> EtlResult<()> { - self.inner.cleanup_table_state(table_id).await + async fn clear_table_copy_state(&self, table_id: TableId) -> EtlResult<()> { + self.inner.clear_table_copy_state(table_id).await + } + + async fn delete_table_pipeline_state(&self, table_id: TableId) -> EtlResult<()> { + self.inner.delete_table_pipeline_state(table_id).await } } diff --git a/crates/etl/src/destination/async_result.rs b/crates/etl/src/destination/async_result.rs index a9d6bdf85..66b89af2c 100644 --- a/crates/etl/src/destination/async_result.rs +++ b/crates/etl/src/destination/async_result.rs @@ -25,11 +25,11 @@ use crate::{ pub type WriteTableRowsResult = AsyncResult; /// Async completion handle used for -/// [`crate::destination::Destination::truncate_table`]. +/// [`crate::destination::Destination::drop_table_for_copy`]. /// -/// ETL waits for this result immediately. It is primarily an API consistency -/// hook rather than a mechanism for overlapping more ETL work with truncation. -pub type TruncateTableResult = AsyncResult; +/// ETL waits for this result immediately before clearing stored table-copy +/// metadata and starting a fresh copy. +pub type DropTableForCopyResult = AsyncResult; /// Async completion handle used for /// [`crate::destination::Destination::write_events`]. diff --git a/crates/etl/src/destination/base.rs b/crates/etl/src/destination/base.rs index 154ff3307..7d7c96a1d 100644 --- a/crates/etl/src/destination/base.rs +++ b/crates/etl/src/destination/base.rs @@ -1,7 +1,7 @@ use std::future::Future; use crate::{ - destination::async_result::{TruncateTableResult, WriteEventsResult, WriteTableRowsResult}, + destination::async_result::{DropTableForCopyResult, WriteEventsResult, WriteTableRowsResult}, error::EtlResult, types::{Event, ReplicatedTableSchema, TableRow}, }; @@ -35,20 +35,21 @@ pub trait Destination { async { Ok(()) } } - /// Truncates all data in the specified table. + /// Drops destination objects before restarting a table copy. /// - /// This operation is called during initial table synchronization to ensure - /// the destination table starts from a clean state before bulk loading. + /// This operation is called when table synchronization intentionally + /// restarts from scratch. Implementations should remove the destination + /// object and any destination-private replay markers for the table so the + /// next copy can recreate it from the fresh source schema. /// - /// Implementations complete `async_result` when truncation is actually - /// done. ETL still waits for that result immediately before continuing. - /// The asynchronous result exists mainly to keep the destination - /// interface uniform across methods, not to let ETL overlap more work - /// with truncation. - fn truncate_table( + /// The supplied schema describes the previously known destination table and + /// exists only so the destination can locate what should be removed. ETL + /// clears its own destination metadata and stored schemas only after this + /// result completes successfully. + fn drop_table_for_copy( &self, replicated_table_schema: &ReplicatedTableSchema, - async_result: TruncateTableResult<()>, + async_result: DropTableForCopyResult<()>, ) -> impl Future> + Send; /// Writes a batch of table rows to the destination. diff --git a/crates/etl/src/destination/mod.rs b/crates/etl/src/destination/mod.rs index fc0c7d654..6fb34d18d 100644 --- a/crates/etl/src/destination/mod.rs +++ b/crates/etl/src/destination/mod.rs @@ -7,5 +7,5 @@ pub mod async_result; mod base; -pub use async_result::{TruncateTableResult, WriteEventsResult, WriteTableRowsResult}; +pub use async_result::{DropTableForCopyResult, WriteEventsResult, WriteTableRowsResult}; pub use base::Destination; diff --git a/crates/etl/src/lib.rs b/crates/etl/src/lib.rs index b220957cf..6de4bcfd9 100644 --- a/crates/etl/src/lib.rs +++ b/crates/etl/src/lib.rs @@ -63,7 +63,9 @@ //! BatchConfig, InvalidatedSlotBehavior, MemoryBackpressureConfig, PgConnectionConfig, //! PipelineConfig, TableSyncCopyConfig, TcpKeepaliveConfig, TlsConfig, //! }, -//! destination::{Destination, TruncateTableResult, WriteEventsResult, WriteTableRowsResult}, +//! destination::{ +//! Destination, DropTableForCopyResult, WriteEventsResult, WriteTableRowsResult, +//! }, //! error::EtlResult, //! pipeline::Pipeline, //! store::MemoryStore, @@ -77,10 +79,10 @@ //! fn name() -> &'static str { //! "noop" //! } -//! async fn truncate_table( +//! async fn drop_table_for_copy( //! &self, //! _replicated_table_schema: &ReplicatedTableSchema, -//! async_result: TruncateTableResult<()>, +//! async_result: DropTableForCopyResult<()>, //! ) -> EtlResult<()> { //! async_result.send(Ok(())); //! Ok(()) diff --git a/crates/etl/src/pipeline.rs b/crates/etl/src/pipeline.rs index 58688cf61..657689e8b 100644 --- a/crates/etl/src/pipeline.rs +++ b/crates/etl/src/pipeline.rs @@ -406,9 +406,9 @@ where "table removed from publication, purging stored state and slot" ); - // We clean up all table state before removing the slot, so that we don't incur - // in the case where we have a slot tied to an invalid state. - self.store.cleanup_table_state(table_id).await?; + // We delete all table state before removing the slot, so that we don't + // incur in the case where we have a slot tied to an invalid state. + self.store.delete_table_pipeline_state(table_id).await?; // We try to delete the replication slot. let slot_name: String = diff --git a/crates/etl/src/replication/apply.rs b/crates/etl/src/replication/apply.rs index fd1724c66..fb82b2e24 100644 --- a/crates/etl/src/replication/apply.rs +++ b/crates/etl/src/replication/apply.rs @@ -76,6 +76,7 @@ use crate::{ }, state::table::{TableReplicationPhase, TableReplicationPhaseType}, store::{ + cleanup::CleanupStore, schema::{SchemaStore, TableSchemaRetention}, state::StateStore, }, @@ -557,14 +558,6 @@ impl ApplyLoopState { Some(SchemaCleanupRun { deadline: Arc::clone(&self.schema_cleanup_deadline) }) } - /// Moves the schema cleanup deadline into the past for tests. - #[cfg(test)] - async fn expire_schema_cleanup_deadline(&self) { - let mut schema_cleanup_deadline = self.schema_cleanup_deadline.lock().await; - *schema_cleanup_deadline = - Some(Instant::now() - SCHEMA_CLEANUP_INTERVAL - Duration::from_secs(1)); - } - /// Returns `true` if a schema cleanup task is still recorded. fn has_schema_cleanup_task(&self) -> bool { self.schema_cleanup_task.is_some() @@ -720,7 +713,7 @@ pub(crate) struct ApplyLoop { impl ApplyLoop where - S: StateStore + SchemaStore + Clone + Send + Sync + 'static, + S: StateStore + SchemaStore + CleanupStore + Clone + Send + Sync + 'static, D: Destination + Clone + Send + Sync + 'static, { /// Starts the apply loop for processing replication events. @@ -2063,7 +2056,7 @@ where let used_bootstrap_snapshot = shared_table_state.is_none(); let table_snapshot_id = shared_table_state .map_or_else(|| self.state.bootstrap_snapshot_id(), |state| state.snapshot_id()); - let table_schema = get_table_schema( + let table_schema = get_table_schema_for_relation( &self.schema_store, &table_id, table_snapshot_id, @@ -2482,7 +2475,7 @@ mod apply_worker { current_lsn: PgLsn, ) -> EtlResult> where - S: StateStore + SchemaStore + Clone + Send + Sync + 'static, + S: StateStore + SchemaStore + CleanupStore + Clone + Send + Sync + 'static, D: Destination + Clone + Send + Sync + 'static, { for (table_id, table_replication_phase) in get_syncing_tables(&ctx.store).await? { @@ -2602,7 +2595,7 @@ mod apply_worker { current_lsn: PgLsn, ) -> EtlResult> where - S: StateStore + SchemaStore + Clone + Send + Sync + 'static, + S: StateStore + SchemaStore + CleanupStore + Clone + Send + Sync + 'static, D: Destination + Clone + Send + Sync + 'static, { let worker_state = ctx.pool.get_active_worker_state(table_id).await; @@ -2741,7 +2734,7 @@ mod apply_worker { current_lsn: PgLsn, ) -> EtlResult<()> where - S: StateStore + SchemaStore + Clone + Send + Sync + 'static, + S: StateStore + SchemaStore + CleanupStore + Clone + Send + Sync + 'static, D: Destination + Clone + Send + Sync + 'static, { for (table_id, table_replication_phase) in get_syncing_tables(&ctx.store).await? { @@ -2767,7 +2760,7 @@ mod apply_worker { current_lsn: PgLsn, ) -> EtlResult<()> where - S: StateStore + SchemaStore + Clone + Send + Sync + 'static, + S: StateStore + SchemaStore + CleanupStore + Clone + Send + Sync + 'static, D: Destination + Clone + Send + Sync + 'static, { let worker_state = ctx.pool.get_active_worker_state(table_id).await; @@ -2879,7 +2872,7 @@ mod apply_worker { current_lsn: PgLsn, ) -> EtlResult> where - S: StateStore + SchemaStore + Clone + Send + Sync + 'static, + S: StateStore + SchemaStore + CleanupStore + Clone + Send + Sync + 'static, D: Destination + Clone + Send + Sync + 'static, { for (table_id, table_replication_phase) in get_syncing_tables(&ctx.store).await? { @@ -2911,7 +2904,7 @@ mod apply_worker { current_lsn: PgLsn, ) -> EtlResult> where - S: StateStore + SchemaStore + Clone + Send + Sync + 'static, + S: StateStore + SchemaStore + CleanupStore + Clone + Send + Sync + 'static, D: Destination + Clone + Send + Sync + 'static, { let worker_state = ctx.pool.get_active_worker_state(table_id).await; @@ -3111,7 +3104,7 @@ mod apply_worker { worker: TableSyncWorker, ) -> Pin> + Send>> where - S: StateStore + SchemaStore + Clone + Send + Sync + 'static, + S: StateStore + SchemaStore + CleanupStore + Clone + Send + Sync + 'static, D: Destination + Clone + Send + Sync + 'static, { Box::pin(async move { worker.spawn_into_pool(&pool).await }) @@ -3254,12 +3247,11 @@ mod table_sync_worker { /// Retrieves a table schema from the schema store by table ID and snapshot. /// -/// When `used_bootstrap_snapshot` is `false`, the returned schema must match -/// the requested snapshot exactly. When it is `true`, the lookup is allowed to -/// resolve to an older schema version because the first `RELATION` message may -/// arrive before shared per-table protocol state has been established, but it -/// must never resolve to a newer schema than requested. -async fn get_table_schema( +/// The reason for this handling is that the bootstrap snapshot id is just used +/// as a boundary id that is needed for when etl starts and no DDL takes place. +/// As soon as a DDL takes place within this active replication session, the +/// exact snapshot id will be used to load the right table schema. +async fn get_table_schema_for_relation( schema_store: &S, table_id: &TableId, snapshot_id: SnapshotId, @@ -3351,36 +3343,3 @@ async fn get_replicated_table_schema( Ok(replicated_table_schema) } - -#[cfg(test)] -mod tests { - use super::*; - - fn apply_loop_state() -> ApplyLoopState { - let start_lsn = PgLsn::from(100u64); - let replication_progress = - ReplicationProgress { last_received_lsn: start_lsn, last_flush_lsn: start_lsn }; - - ApplyLoopState::new( - replication_progress, - Duration::from_secs(1), - SnapshotId::from(start_lsn), - "test_slot".to_owned(), - ) - } - - #[tokio::test] - async fn schema_cleanup_uses_deadline_between_runs() { - let state = apply_loop_state(); - - state.expire_schema_cleanup_deadline().await; - - let schema_cleanup_run = state.try_start_schema_cleanup().await.unwrap(); - - assert!(state.try_start_schema_cleanup().await.is_none()); - - schema_cleanup_run.finish().await; - - assert!(state.try_start_schema_cleanup().await.is_none()); - } -} diff --git a/crates/etl/src/replication/table_cache.rs b/crates/etl/src/replication/table_cache.rs index fa5023d26..8810aec77 100644 --- a/crates/etl/src/replication/table_cache.rs +++ b/crates/etl/src/replication/table_cache.rs @@ -25,6 +25,22 @@ //! relation state for that snapshot. //! - [`SharedTableState::Ready`], where the full [`ReplicatedTableSchema`] //! needed for row decoding is already materialized in memory. +//! +//! The cache relies on the following invariants: +//! - A table has exactly one protocol owner at a time. The owner may be the +//! apply worker or the table sync worker for that table; non-owners must skip +//! DDL, `RELATION`, and row messages without changing the cache. +//! - A table copy stores the initial schema with snapshot `0/0` and marks the +//! table [`SharedTableState::Ready`] before catchup rows can be decoded. +//! - A table-copy restart removes any cached state before storing the fresh +//! `0/0` schema. The new copy lineage must repopulate the cache; stale ready +//! state from a previous copy must not be reused. +//! - DDL handling stores the durable schema version before moving the cache to +//! [`SharedTableState::WaitingForRelation`] for that snapshot. +//! - [`SharedTableState::WaitingForRelation`] is not a failure state, but it is +//! not decodable. Row handlers must reject row messages until a later +//! `RELATION` message materializes the runtime masks and moves the table back +//! to [`SharedTableState::Ready`]. use std::{collections::HashMap, sync::Arc}; @@ -109,6 +125,16 @@ impl SharedTableCache { self.upsert(table_id, SharedTableState::Ready { replicated_table_schema }).await; } + /// Removes cached runtime state for a table. + /// + /// This is intentionally idempotent: a first copy may not have cached state + /// yet, while an in-process copy restart may have stale ready state from + /// the previous copy. + pub(crate) async fn remove_table(&self, table_id: TableId) { + let mut guard = self.inner.write().await; + guard.remove(&table_id); + } + /// Inserts or updates shared table state. /// /// This cache is deliberately oblivious to ordering and ownership. It @@ -254,4 +280,22 @@ mod tests { assert!(table_ids.contains(&ready_table_id)); assert!(table_ids.contains(&waiting_table_id)); } + + #[tokio::test] + async fn remove_table_is_idempotent() { + let cache = SharedTableCache::new(); + let table_id = TableId::new(123); + + cache.remove_table(table_id).await; + assert!(cache.get(&table_id).await.is_none()); + + cache.note_ready(table_id, create_test_schema()).await; + assert!(cache.get(&table_id).await.is_some()); + + cache.remove_table(table_id).await; + cache.remove_table(table_id).await; + + assert!(cache.get(&table_id).await.is_none()); + assert!(!cache.active_table_ids().await.contains(&table_id)); + } } diff --git a/crates/etl/src/replication/table_sync.rs b/crates/etl/src/replication/table_sync.rs index 999840d37..bb38480e6 100644 --- a/crates/etl/src/replication/table_sync.rs +++ b/crates/etl/src/replication/table_sync.rs @@ -16,14 +16,14 @@ use crate::{ concurrency::{BatchBudgetController, MemoryMonitor, ShutdownRx}, destination::{ Destination, - async_result::{TruncateTableResult, WriteTableRowsResult}, + async_result::{DropTableForCopyResult, WriteTableRowsResult}, }, error::{ErrorKind, EtlResult}, etl_error, metrics::{ETL_TABLE_COPY_DURATION_SECONDS, PARTITIONING_LABEL}, - replication::{WorkerType, client::PgReplicationClient, table_cache::SharedTableCache}, + replication::{client::PgReplicationClient, table_cache::SharedTableCache}, state::table::{TableReplicationPhase, TableReplicationPhaseType}, - store::{schema::SchemaStore, state::StateStore}, + store::{cleanup::CleanupStore, schema::SchemaStore, state::StateStore}, types::PipelineId, workers::{TableCopyResult, TableSyncWorkerState, table_copy}, }; @@ -48,6 +48,36 @@ pub(crate) enum TableSyncResult { }, } +/// Returns the existing [`ReplicatedTableSchema`] if one exists. +/// +/// A [`ReplicatedTableSchema`] could be there when starting a table copy +/// because it was either interrupted or the state was reset. +async fn get_existing_replicated_table_schema( + store: &S, + table_id: TableId, +) -> EtlResult> +where + S: StateStore + SchemaStore + Send + 'static, +{ + let Some(current_metadata) = store.get_destination_table_metadata(table_id).await? else { + return Ok(None); + }; + + let Some(table_schema) = + store.get_table_schema(&table_id, current_metadata.snapshot_id).await? + else { + bail!( + ErrorKind::InvalidState, + "Destination table metadata found, but no corresponding table schema exists" + ); + }; + + let existing_replicated_table_schema = + ReplicatedTableSchema::from_mask(table_schema, current_metadata.replication_mask); + + Ok(Some(existing_replicated_table_schema)) +} + /// Starts table synchronization for a specific table. /// /// This function performs the initial data copy for a table from the source @@ -69,7 +99,7 @@ pub(crate) async fn start_table_sync( batch_budget: BatchBudgetController, ) -> EtlResult where - S: StateStore + SchemaStore + Clone + Send + 'static, + S: StateStore + SchemaStore + CleanupStore + Clone + Send + 'static, D: Destination + Clone + Send + 'static, { info!(table_id = table_id.0, "starting initial table sync"); @@ -130,107 +160,63 @@ where let slot_name: String = EtlReplicationSlot::for_table_sync_worker(pipeline_id, table_id).try_into()?; - // There are three phases in which the table can be in: - // - `Init` -> this means that the table sync was never done, so we just perform - // it. - // - `DataSync` -> this means that there was a failure during data sync, and we - // have to restart - // copying all the table data and delete the slot. - // - `FinishedCopy` -> this means that the table was successfully copied, but we - // didn't manage to complete the table sync function, so we just want to - // continue the cdc stream from durable table-sync progress when available, or - // from the slot's confirmed flush LSN otherwise. + // There are three phases from which table sync can start: + // - `Init` -> the table sync was never done or the table was reset, so we + // perform it. + // - `DataSync` -> a previous copy did not complete, so we restart the copy from + // a clean snapshot. + // - `FinishedCopy` -> the copy completed, but the table sync worker did not + // finish the ownership handoff. This is a narrow crash window, so we + // intentionally restart the copy instead of trying to resume catchup. + // Resuming soundly would require persisting the exact runtime relation + // decoding state, including identity masks, and making handoff safe when no + // further `RELATION` message arrives. // // In case the phase is any other phase, we will return an error. let start_lsn = match phase_type { - TableReplicationPhaseType::Init | TableReplicationPhaseType::DataSync => { - // When we are in these states, it could be for the following reasons: - // - `Init` -> we can be in this state because we just started replicating the - // table or the state - // was reset. In this case we don't want to make assumptions about the previous - // state, so we just try to delete the slot and truncate the table. - // - `DataSync` -> we can be in this state because we failed during data sync, - // meaning that table - // copy failed. In this case, we want to delete the slot and truncate the - // table. - // - // We try to delete the slot also during `Init` because we support state - // rollback and a slot might be there from the previous run. - replication_client.delete_slot_if_exists(&slot_name).await?; - store.delete_replication_progress(WorkerType::TableSync { table_id }).await?; - - // We must truncate the destination table before starting a copy to avoid data - // inconsistencies. + TableReplicationPhaseType::Init + | TableReplicationPhaseType::DataSync + | TableReplicationPhaseType::FinishedCopy => { + // We must drop the destination table before starting a copy to avoid data + // inconsistencies when there is a previous table. // // Example scenario: // 1. The source table has a single row (id = 1) that is copied to the // destination during the initial copy. - // 2. Before the table’s phase is set to `FinishedCopy`, the process crashes. + // 2. Before the table's phase is set to `FinishedCopy`, the process crashes. // 3. While down, the source deletes row id = 1 and inserts row id = 2. - // 4. When restarted, the process sees the table in the ` DataSync ` state, + // 4. When restarted, the process sees the table in the `DataSync` state, // deletes the slot, and copies again. // 5. This time, only row id = 2 is copied, but row id = 1 still exists in the // destination. + // // Result: the destination has two rows (id = 1 and id = 2) instead of only one - // (id = 2). Fix: Always truncate the destination table before - // starting a copy. + // (id = 2). // - // Try to load the previously stored destination table metadata, which contains - // both the snapshot_id and replication_mask. If available, we can load the - // corresponding table schema and truncate the destination table before starting - // a copy. If the metadata is not present, we can safely assume that - // no data is there in the table; thus a truncate won't be issued. - if let Some(current_metadata) = store.get_destination_table_metadata(table_id).await? { - match current_metadata.into_applied() { - Err(err) => { - // The schema DDL never completed. Skip the truncate and let the - // destination re-create the table during the copy phase. - warn!( - table_id = table_id.0, - error = %err, - "destination table metadata is not in applied state; skipping pre-copy truncation" - ); - } - Ok(applied_metadata) => { - if let Some(table_schema) = - store.get_table_schema(&table_id, applied_metadata.snapshot_id).await? - { - let replicated_table_schema = ReplicatedTableSchema::from_mask( - table_schema, - applied_metadata.replication_mask, - ); - let (truncate_result, pending_truncate_result) = - TruncateTableResult::new(()); - - if let Err(err) = destination - .truncate_table(&replicated_table_schema, truncate_result) - .await - { - warn!( - table_id = table_id.0, - error = %err, - "failed to dispatch destination table truncation before copy, continuing" - ); - } else if let Err(err) = pending_truncate_result.await.into_result() { - warn!( - table_id = table_id.0, - error = %err, - "failed to truncate destination table before copy, continuing" - ); - } else { - info!(%table_id, "truncated destination table before starting copy"); - } - } else { - bail!( - ErrorKind::InvalidState, - "Destination table metadata found, but not corresponding table \ - schema exists" - ); - } - } - } + // Fix: Always drop the destination table before starting a copy. + if let Some(current_replication_table_schema) = + get_existing_replicated_table_schema(&store, table_id).await? + { + let (drop_result, pending_drop_result) = DropTableForCopyResult::new(()); + destination + .drop_table_for_copy(¤t_replication_table_schema, drop_result) + .await?; + pending_drop_result.await.into_result()?; } + // We try to delete the slot if it already exists, since we might be starting a + // table copy after a previous one was reset or didn't complete + // successfully. + replication_client.delete_slot_if_exists(&slot_name).await?; + + // We clear durable and in-memory table-copy state only after external cleanup + // succeeds. The shared cache removal is idempotent: a first copy has no cached + // state yet, while an in-process retry can still hold the previous ready + // runtime schema. The fresh `0/0` copy schema below is the only state allowed + // to repopulate the cache. + store.clear_table_copy_state(table_id).await?; + shared_table_cache.remove_table(table_id).await; + // We are ready to start copying table data, and we update the state // accordingly. info!(table_id = table_id.0, "starting data copy"); @@ -308,7 +294,6 @@ where // without waiting for a fresh relation message after restarts. let replicated_table_schema = ReplicatedTableSchema::from_masks(table_schema, replication_mask, identity_mask); - shared_table_cache.note_ready(table_id, replicated_table_schema.clone()).await; let mut total_table_copy_rows = 0; let mut total_table_copy_duration_secs = 0.0; @@ -381,41 +366,16 @@ where inner.set_and_store(TableReplicationPhase::FinishedCopy, &store).await?; } - slot.consistent_point - } - TableReplicationPhaseType::FinishedCopy => { - let slot = replication_client.get_slot(&slot_name).await?; - let worker_type = WorkerType::TableSync { table_id }; - let durable_flush_lsn = store.get_replication_progress(worker_type).await?; - if let Some(durable_flush_lsn) = durable_flush_lsn { - // Durable progress and slot progress can legitimately differ. During idle - // periods we keep sending PostgreSQL feedback with the received LSN, but - // we do not persist those idle-only advances to the state database to - // avoid extra customer-database writes. Conversely, durable progress can - // be ahead if ETL flushed a batch but PostgreSQL did not confirm the - // feedback yet. Startup uses the latest boundary available from either - // source as a resume floor, which guarantees no event older than the - // chosen start LSN is emitted. - let start_lsn = durable_flush_lsn.max(slot.confirmed_flush_lsn); - - info!( - table_id = table_id.0, - %durable_flush_lsn, - confirmed_flush_lsn = %slot.confirmed_flush_lsn, - %start_lsn, - "resuming table sync from durable replication progress and replication slot" - ); - - start_lsn - } else { - info!( - table_id = table_id.0, - confirmed_flush_lsn = %slot.confirmed_flush_lsn, - "durable table sync progress not found, using slot fallback" - ); + // After we finished copying, we mark this table as ready in the cache, so + // that we can start streaming and decoding immediately. + // + // This is needed, since it could be that the apply loop for the `Catchup` phase + // might be idle and progress only via keepalives and in that case no `Relation` + // message will be received, so we want the apply worker to already be able to + // start decoding. + shared_table_cache.note_ready(table_id, replicated_table_schema.clone()).await; - slot.confirmed_flush_lsn - } + slot.consistent_point } _ => unreachable!("phase type already validated above"), }; diff --git a/crates/etl/src/store/both/memory.rs b/crates/etl/src/store/both/memory.rs index ff05beb56..532a87883 100644 --- a/crates/etl/src/store/both/memory.rs +++ b/crates/etl/src/store/both/memory.rs @@ -56,6 +56,15 @@ pub struct MemoryStore { inner: Arc>, } +/// Scope of table-scoped state removal. +#[derive(Debug, Clone, Copy)] +enum TableStateCleanupScope { + /// Clear state that belongs to the previous table copy. + CopyRestart, + /// Delete all ETL state for a table removed from the publication. + PipelineRemoval, +} + impl MemoryStore { /// Creates a new empty memory store. /// @@ -73,6 +82,26 @@ impl MemoryStore { Self { inner: Arc::new(Mutex::new(inner)) } } + + /// Deletes table-scoped state according to the requested scope. + async fn delete_table_state_for_scope( + &self, + table_id: TableId, + scope: TableStateCleanupScope, + ) -> EtlResult<()> { + let mut inner = self.inner.lock().await; + + if matches!(scope, TableStateCleanupScope::PipelineRemoval) { + Arc::make_mut(&mut inner.table_replication_states).remove(&table_id); + inner.table_state_history.remove(&table_id); + } + + Arc::make_mut(&mut inner.table_schemas).remove_table(table_id); + Arc::make_mut(&mut inner.destination_tables_metadata).remove(&table_id); + inner.replication_progress.remove(&WorkerType::TableSync { table_id }); + + Ok(()) + } } impl Default for MemoryStore { @@ -269,16 +298,11 @@ impl SchemaStore for MemoryStore { } impl CleanupStore for MemoryStore { - async fn cleanup_table_state(&self, table_id: TableId) -> EtlResult<()> { - let mut inner = self.inner.lock().await; - - Arc::make_mut(&mut inner.table_replication_states).remove(&table_id); - inner.table_state_history.remove(&table_id); - // Remove all schema versions for this table. - Arc::make_mut(&mut inner.table_schemas).remove_table(table_id); - Arc::make_mut(&mut inner.destination_tables_metadata).remove(&table_id); - inner.replication_progress.remove(&WorkerType::TableSync { table_id }); + async fn clear_table_copy_state(&self, table_id: TableId) -> EtlResult<()> { + self.delete_table_state_for_scope(table_id, TableStateCleanupScope::CopyRestart).await + } - Ok(()) + async fn delete_table_pipeline_state(&self, table_id: TableId) -> EtlResult<()> { + self.delete_table_state_for_scope(table_id, TableStateCleanupScope::PipelineRemoval).await } } diff --git a/crates/etl/src/store/both/postgres.rs b/crates/etl/src/store/both/postgres.rs index 60b692df8..359c63a4b 100644 --- a/crates/etl/src/store/both/postgres.rs +++ b/crates/etl/src/store/both/postgres.rs @@ -45,6 +45,15 @@ const IDLE_TIMEOUT: Duration = Duration::from_secs(30); /// need more than 2 schema versions for any given table. const MAX_CACHED_SCHEMAS_PER_TABLE: usize = 2; +/// Scope of table-scoped state removal. +#[derive(Debug, Clone, Copy)] +enum TableStateCleanupScope { + /// Clear state that belongs to the previous table copy. + CopyRestart, + /// Delete all ETL state for a table removed from the publication. + PipelineRemoval, +} + /// Creates a lazily connected pool with automatic idle connection cleanup. /// /// This function returns immediately without establishing any connections. @@ -183,6 +192,72 @@ impl PostgresStore { Ok(Self { pipeline_id, pool, inner: Arc::new(Mutex::new(inner)) }) } + + /// Deletes table-scoped persistent and cached state according to the + /// requested scope. + async fn delete_table_state_for_scope( + &self, + table_id: TableId, + scope: TableStateCleanupScope, + ) -> EtlResult<()> { + let mut inner = self.inner.lock().await; + let mut tx = self.pool.begin().await?; + + destination_metadata::delete_destination_table_metadata( + &mut *tx, + self.pipeline_id as i64, + table_id, + ) + .await + .map_err(|err| { + etl_error!( + ErrorKind::SourceQueryFailed, + "Destination table metadata deletion failed", + source: err + ) + })?; + + schema::delete_table_schema_for_table(&mut *tx, self.pipeline_id as i64, table_id) + .await + .map_err(|err| { + etl_error!( + ErrorKind::SourceQueryFailed, + "Table schema deletion failed", + source: err + ) + })?; + + if matches!(scope, TableStateCleanupScope::PipelineRemoval) { + state::delete_replication_state_for_table(&mut *tx, self.pipeline_id as i64, table_id) + .await?; + } + + progress::delete_replication_progress_for_table( + &mut *tx, + self.pipeline_id as i64, + table_id, + ) + .await + .map_err(|err| { + etl_error!( + ErrorKind::SourceQueryFailed, + "Replication progress deletion failed", + source: err + ) + })?; + + tx.commit().await?; + + if matches!(scope, TableStateCleanupScope::PipelineRemoval) { + inner.remove_table_state(table_id); + emit_table_metrics(&inner.phase_counts); + } + + Arc::make_mut(&mut inner.table_schemas).remove_table(table_id); + Arc::make_mut(&mut inner.destination_tables_metadata).remove(&table_id); + + Ok(()) + } } impl StateStore for PostgresStore { @@ -637,60 +712,13 @@ impl SchemaStore for PostgresStore { } impl CleanupStore for PostgresStore { - /// Removes all state for a table from both database and cache. - async fn cleanup_table_state(&self, table_id: TableId) -> EtlResult<()> { - let mut inner = self.inner.lock().await; - let mut tx = self.pool.begin().await?; - - destination_metadata::delete_destination_table_metadata( - &mut *tx, - self.pipeline_id as i64, - table_id, - ) - .await - .map_err(|err| { - etl_error!( - ErrorKind::SourceQueryFailed, - "Destination table metadata deletion failed", - format!("Failed to delete destination table metadata in PostgreSQL: {}", err) - ) - })?; - - schema::delete_table_schema_for_table(&mut *tx, self.pipeline_id as i64, table_id) - .await - .map_err(|err| { - etl_error!( - ErrorKind::SourceQueryFailed, - "Table schema deletion failed", - format!("Failed to delete table schema in PostgreSQL: {}", err) - ) - })?; - - state::delete_replication_state_for_table(&mut *tx, self.pipeline_id as i64, table_id) - .await?; - - progress::delete_replication_progress_for_table( - &mut *tx, - self.pipeline_id as i64, - table_id, - ) - .await - .map_err(|err| { - etl_error!( - ErrorKind::SourceQueryFailed, - "Replication progress deletion failed", - source: err - ) - })?; - - tx.commit().await?; - - inner.remove_table_state(table_id); - Arc::make_mut(&mut inner.table_schemas).remove_table(table_id); - Arc::make_mut(&mut inner.destination_tables_metadata).remove(&table_id); - emit_table_metrics(&inner.phase_counts); + async fn clear_table_copy_state(&self, table_id: TableId) -> EtlResult<()> { + self.delete_table_state_for_scope(table_id, TableStateCleanupScope::CopyRestart).await + } - Ok(()) + /// Removes all state for a table from both database and cache. + async fn delete_table_pipeline_state(&self, table_id: TableId) -> EtlResult<()> { + self.delete_table_state_for_scope(table_id, TableStateCleanupScope::PipelineRemoval).await } } diff --git a/crates/etl/src/store/cleanup.rs b/crates/etl/src/store/cleanup.rs index 575782803..8140937b0 100644 --- a/crates/etl/src/store/cleanup.rs +++ b/crates/etl/src/store/cleanup.rs @@ -4,10 +4,21 @@ use crate::{error::EtlResult, types::TableId}; /// Combined maintenance operations across state and schema stores. /// -/// Provides atomic cleanup primitives that affect both replication state -/// and schema-related data for a specific table. Implementations should -/// ensure consistency across in-memory caches and the persistent store. +/// Provides atomic table-scoped primitives that affect both replication state +/// and schema-related data. Implementations should ensure consistency across +/// in-memory caches and the persistent store. pub trait CleanupStore { + /// Clears stored table-copy state for `table_id`. + /// + /// Removes destination table metadata, all stored table schemas, and + /// durable table-sync progress while preserving the table replication + /// phase. This is used after the destination object has been dropped and + /// before a fresh `0/0` table-copy schema is stored. + fn clear_table_copy_state( + &self, + table_id: TableId, + ) -> impl Future> + Send; + /// Deletes all stored state for `table_id` for the current pipeline. /// /// Removes replication state (including history), table schemas, and @@ -15,5 +26,8 @@ pub trait CleanupStore { /// destination table. /// /// Intended for use when a table is removed from the publication. - fn cleanup_table_state(&self, table_id: TableId) -> impl Future> + Send; + fn delete_table_pipeline_state( + &self, + table_id: TableId, + ) -> impl Future> + Send; } diff --git a/crates/etl/src/test_utils/memory_destination.rs b/crates/etl/src/test_utils/memory_destination.rs index bc649536f..cc7bcd03b 100644 --- a/crates/etl/src/test_utils/memory_destination.rs +++ b/crates/etl/src/test_utils/memory_destination.rs @@ -6,7 +6,7 @@ use tracing::info; use crate::{ destination::{ Destination, - async_result::{TruncateTableResult, WriteEventsResult, WriteTableRowsResult}, + async_result::{DropTableForCopyResult, WriteEventsResult, WriteTableRowsResult}, }, error::EtlResult, state::destination_metadata::{DestinationTableMetadata, DestinationTableSchemaStatus}, @@ -28,8 +28,8 @@ struct Inner { /// terminates. /// /// Like real destinations (BigQuery, Iceberg), this destination tracks table -/// metadata (snapshot IDs and replication masks) in a state store to support -/// features like table truncation during state resets. +/// metadata (snapshot IDs and replication masks) in a state store to mirror +/// destinations that persist table-copy state. #[derive(Clone)] pub struct MemoryDestination { inner: Arc>, @@ -133,17 +133,17 @@ where "memory" } - async fn truncate_table( + async fn drop_table_for_copy( &self, replicated_table_schema: &ReplicatedTableSchema, - async_result: TruncateTableResult<()>, + async_result: DropTableForCopyResult<()>, ) -> EtlResult<()> { - // For truncation, we simulate removing all table rows for a specific table and + // For table drops, we simulate removing all table rows for a specific table and // also the events of that table. let mut inner = self.inner.lock().await; let table_id = replicated_table_schema.id(); - info!(%table_id, "truncating table"); + info!(%table_id, "dropping table for copy"); inner.table_rows.remove(&table_id); inner.events.retain_mut(|event| { diff --git a/crates/etl/src/test_utils/notifying_store.rs b/crates/etl/src/test_utils/notifying_store.rs index 614f635e6..6efb76cbc 100644 --- a/crates/etl/src/test_utils/notifying_store.rs +++ b/crates/etl/src/test_utils/notifying_store.rs @@ -108,6 +108,15 @@ pub struct NotifyingStore { inner: Arc>, } +/// Scope of table-scoped state removal. +#[derive(Debug, Clone, Copy)] +enum TableStateCleanupScope { + /// Clear state that belongs to the previous table copy. + CopyRestart, + /// Delete all ETL state for a table removed from the publication. + PipelineRemoval, +} + impl NotifyingStore { /// Creates an empty notifying store. pub fn new() -> Self { @@ -241,6 +250,27 @@ impl NotifyingStore { Ok(()) } + + /// Deletes table-scoped state according to the requested scope. + async fn delete_table_state_for_scope( + &self, + table_id: TableId, + scope: TableStateCleanupScope, + ) -> EtlResult<()> { + let mut inner = self.inner.write().await; + + if matches!(scope, TableStateCleanupScope::PipelineRemoval) { + Arc::make_mut(&mut inner.table_replication_states).remove(&table_id); + inner.table_state_history.remove(&table_id); + } + + Arc::make_mut(&mut inner.table_schemas).remove_table(table_id); + Arc::make_mut(&mut inner.destination_tables_metadata).remove(&table_id); + inner.replication_progress.remove(&WorkerType::TableSync { table_id }); + inner.check_conditions(); + + Ok(()) + } } impl Default for NotifyingStore { @@ -452,16 +482,12 @@ impl SchemaStore for NotifyingStore { } impl CleanupStore for NotifyingStore { - async fn cleanup_table_state(&self, table_id: TableId) -> EtlResult<()> { - let mut inner = self.inner.write().await; - - Arc::make_mut(&mut inner.table_replication_states).remove(&table_id); - inner.table_state_history.remove(&table_id); - Arc::make_mut(&mut inner.table_schemas).remove_table(table_id); - Arc::make_mut(&mut inner.destination_tables_metadata).remove(&table_id); - inner.replication_progress.remove(&WorkerType::TableSync { table_id }); + async fn clear_table_copy_state(&self, table_id: TableId) -> EtlResult<()> { + self.delete_table_state_for_scope(table_id, TableStateCleanupScope::CopyRestart).await + } - Ok(()) + async fn delete_table_pipeline_state(&self, table_id: TableId) -> EtlResult<()> { + self.delete_table_state_for_scope(table_id, TableStateCleanupScope::PipelineRemoval).await } } diff --git a/crates/etl/src/test_utils/test_destination_wrapper.rs b/crates/etl/src/test_utils/test_destination_wrapper.rs index c54df1978..1617814a5 100644 --- a/crates/etl/src/test_utils/test_destination_wrapper.rs +++ b/crates/etl/src/test_utils/test_destination_wrapper.rs @@ -16,8 +16,8 @@ use crate::{ destination::{ Destination, async_result::{ - ApplyLoopAsyncResultMetadata, DispatchMetrics, TruncateTableResult, WriteEventsResult, - WriteTableRowsResult, + ApplyLoopAsyncResultMetadata, DispatchMetrics, DropTableForCopyResult, + WriteEventsResult, WriteTableRowsResult, }, }, error::EtlResult, @@ -37,7 +37,7 @@ struct Inner { wrapped_destination: D, events: Vec, table_rows: HashMap>, - truncated_tables: HashSet, + tables_dropped_for_copy: HashSet, event_conditions: Vec<(EventCheckFn, Arc)>, table_row_conditions: Vec<(TableRowCheckFn, Arc)>, combined_conditions: Vec<(CombinedCheckFn, Arc)>, @@ -119,7 +119,7 @@ impl TestDestinationWrapper { wrapped_destination: destination, events: Vec::new(), table_rows: HashMap::new(), - truncated_tables: HashSet::new(), + tables_dropped_for_copy: HashSet::new(), event_conditions: Vec::new(), table_row_conditions: Vec::new(), combined_conditions: Vec::new(), @@ -211,9 +211,10 @@ impl TestDestinationWrapper { inner.events.clear(); } - /// Returns whether the table was truncated through the wrapper. - pub async fn was_table_truncated(&self, table_id: TableId) -> bool { - self.inner.read().await.truncated_tables.contains(&table_id) + /// Returns whether the table was dropped for a fresh copy through the + /// wrapper. + pub async fn was_table_dropped_for_copy(&self, table_id: TableId) -> bool { + self.inner.read().await.tables_dropped_for_copy.contains(&table_id) } /// Returns how many times [`Destination::write_table_rows`] was called. @@ -235,44 +236,47 @@ where "wrapper" } - async fn truncate_table( + async fn drop_table_for_copy( &self, replicated_table_schema: &ReplicatedTableSchema, - async_result: TruncateTableResult<()>, + async_result: DropTableForCopyResult<()>, ) -> EtlResult<()> { let destination = { let inner = self.inner.read().await; inner.wrapped_destination.clone() }; - let (wrapped_truncate_result, pending_result) = TruncateTableResult::new(()); - destination.truncate_table(replicated_table_schema, wrapped_truncate_result).await?; + let (wrapped_drop_result, pending_drop_result) = DropTableForCopyResult::new(()); + destination.drop_table_for_copy(replicated_table_schema, wrapped_drop_result).await?; // We send the result back before doing the internal checks for this utility, to // avoid checking before the apply loop received the result. - let result = pending_result.await.into_result(); + let result = pending_drop_result.await.into_result(); + let should_record_drop = result.is_ok(); async_result.send(result); let mut inner = self.inner.write().await; let table_id = replicated_table_schema.id(); - inner.truncated_tables.insert(table_id); - inner.table_rows.remove(&table_id); - inner.events.retain_mut(|event| { - let has_table_id = event.has_table_id(&table_id); - if let Event::Truncate(truncate_event) = event - && has_table_id - { - truncate_event.truncated_tables.retain(|s| s.id() != table_id); - if truncate_event.truncated_tables.is_empty() { - return false; - } + if should_record_drop { + inner.tables_dropped_for_copy.insert(table_id); + inner.table_rows.remove(&table_id); + inner.events.retain_mut(|event| { + let has_table_id = event.has_table_id(&table_id); + if let Event::Truncate(truncate_event) = event + && has_table_id + { + truncate_event.truncated_tables.retain(|s| s.id() != table_id); + if truncate_event.truncated_tables.is_empty() { + return false; + } - return true; - } + return true; + } - !has_table_id - }); + !has_table_id + }); + } Ok(()) } @@ -289,14 +293,14 @@ where inner.wrapped_destination.clone() }; - let (wrapped_flush_result, pending_result) = WriteTableRowsResult::new(()); + let (wrapped_flush_result, pending_flush_result) = WriteTableRowsResult::new(()); destination .write_table_rows(replicated_table_schema, table_rows.clone(), wrapped_flush_result) .await?; // We send the result back before doing the internal checks for this utility, to // avoid checking before the apply loop received the result. - let result = pending_result.await.into_result(); + let result = pending_flush_result.await.into_result(); let should_record_table_rows = result.is_ok(); async_result.send(result); @@ -325,7 +329,7 @@ where inner.wrapped_destination.clone() }; - let (wrapped_flush_result, pending_result) = + let (wrapped_flush_result, pending_flush_result) = WriteEventsResult::new(ApplyLoopAsyncResultMetadata { commit_end_lsn: None, metrics: DispatchMetrics { @@ -349,7 +353,7 @@ where .spawn(async move { // We send the result back before doing the internal checks for this utility, to // avoid checking before the apply loop received the result. - let result = pending_result.await.into_result(); + let result = pending_flush_result.await.into_result(); let should_record_events = result.is_ok(); async_result.send(result); diff --git a/crates/etl/src/workers/apply.rs b/crates/etl/src/workers/apply.rs index f3eaa05ad..a71478f10 100644 --- a/crates/etl/src/workers/apply.rs +++ b/crates/etl/src/workers/apply.rs @@ -22,7 +22,7 @@ use crate::{ client::{GetOrCreateSlotResult, PgReplicationClient, SlotState}, }, state::table::{TableReplicationPhase, TableReplicationPhaseType}, - store::{schema::SchemaStore, state::StateStore}, + store::{cleanup::CleanupStore, schema::SchemaStore, state::StateStore}, types::PipelineId, workers::{ TableSyncWorkerPool, @@ -130,7 +130,7 @@ impl ApplyWorker { impl ApplyWorker where - S: StateStore + SchemaStore + Clone + Send + Sync + 'static, + S: StateStore + SchemaStore + CleanupStore + Clone + Send + Sync + 'static, D: Destination + Clone + Send + Sync + 'static, { /// Handles apply worker errors using policy-based retry and backoff. diff --git a/crates/etl/src/workers/table_sync.rs b/crates/etl/src/workers/table_sync.rs index 421dad51e..4a0188908 100644 --- a/crates/etl/src/workers/table_sync.rs +++ b/crates/etl/src/workers/table_sync.rs @@ -25,7 +25,7 @@ use crate::{ state::table::{ RetryPolicy, TableReplicationError, TableReplicationPhase, TableReplicationPhaseType, }, - store::{schema::SchemaStore, state::StateStore}, + store::{cleanup::CleanupStore, schema::SchemaStore, state::StateStore}, types::PipelineId, workers::{ TableSyncWorkerPool, @@ -380,7 +380,7 @@ impl TableSyncWorker { impl TableSyncWorker where - S: StateStore + SchemaStore + Clone + Send + Sync + 'static, + S: StateStore + SchemaStore + CleanupStore + Clone + Send + Sync + 'static, D: Destination + Clone + Send + Sync + 'static, { /// Handles a table sync worker failure using the configured retry policy. @@ -508,8 +508,8 @@ where // in a table sync worker, this is why it's not in the apply worker: // - Errored -> Init: okay since it will restart from scratch. // - Errored -> DataSync: okay since it will restart the copy from a new slot. - // - Errored -> FinishedCopy: okay since the table was already copied, so it - // resumes streaming from durable table-sync progress or the slot fallback. + // - Errored -> FinishedCopy: okay since table sync startup treats it as a clean + // copy restart. // - Errored -> SyncDone: okay since the table sync will immediately stop. // - Errored -> Ready: same as SyncDone. // diff --git a/crates/etl/tests/pipeline.rs b/crates/etl/tests/pipeline.rs index c805de441..b22e8974b 100644 --- a/crates/etl/tests/pipeline.rs +++ b/crates/etl/tests/pipeline.rs @@ -3,7 +3,7 @@ use std::time::Duration; use etl::{ error::ErrorKind, state::table::{TableReplicationPhase, TableReplicationPhaseType}, - store::state::StateStore, + store::{schema::SchemaStore, state::StateStore}, test_utils::{ database::{spawn_source_database, test_table_name}, event::{EventCondition, group_events_by_type_and_table_id}, @@ -1764,8 +1764,8 @@ async fn empty_tables_are_created_at_destination() { } /// Tests that resetting a table's state to Init triggers a table sync that -/// truncates the destination before re-copying data. This ensures no duplicate -/// data after a state reset. +/// drops the destination table before re-copying data. This ensures no +/// duplicate data after a state reset. /// /// Test flow: /// 1. Initial table sync: 5 rows (ids 1-5) written to table_rows for both users @@ -1775,7 +1775,7 @@ async fn empty_tables_are_created_at_destination() { /// 4. Insert 3 new rows (ids 100-102) for users only /// 5. Verify: users has 10 total rows (table_rows + events), orders unchanged #[tokio::test(flavor = "multi_thread")] -async fn table_sync_truncates_destination_after_state_reset() { +async fn table_sync_drops_destination_table_after_state_reset() { init_test_tracing(); let mut database = spawn_source_database().await; let database_schema = setup_test_database_schema(&database, TableSelection::Both).await; @@ -1870,10 +1870,10 @@ async fn table_sync_truncates_destination_after_state_reset() { // We clear the events and rows to check that only users data is written. // - // This deletion becomes a bit confusing when used in the context of truncation - // that should take care of deleting data by itself, however, in this test - // we just want to make sure that truncation is called and that the data is - // rewritten from scratch. + // This deletion becomes a bit confusing when used in the context of a + // destination drop that should take care of deleting data by itself. In + // this test we just want to make sure that the drop is called and that the + // data is rewritten from scratch. destination.clear_events().await; destination.clear_table_rows().await; @@ -1886,14 +1886,14 @@ async fn table_sync_truncates_destination_after_state_reset() { ) .await; - // Reset users table state to Init, triggering a new table sync with truncate. + // Reset users table state to Init, triggering a fresh table sync. store.reset_table_state(database_schema.users_schema().id).await.unwrap(); users_ready_notify.notified().await; // Wait for all user events (table_rows + CDC) to be processed. - // After reset, data can end up in either table_rows or events depending on - // timing. + // After the state reset, data can end up in either table_rows or events + // depending on timing. let total_expected_users = initial_rows + cdc_rows + new_rows_after_reset; let all_users_events_notify = destination .wait_for_all_events(vec![EventCondition::Table( @@ -1940,9 +1940,18 @@ async fn table_sync_truncates_destination_after_state_reset() { .contains_key(&(EventType::Insert, database_schema.orders_schema().id)) ); - // Verify truncate was called for users (due to reset) but not for orders. - assert!(destination.was_table_truncated(database_schema.users_schema().id).await); - assert!(!destination.was_table_truncated(database_schema.orders_schema().id).await); + // Verify the destination table was dropped for users but not for orders. + assert!(destination.was_table_dropped_for_copy(database_schema.users_schema().id).await); + assert!(!destination.was_table_dropped_for_copy(database_schema.orders_schema().id).await); + + let user_schemas = SchemaStore::get_table_schemas(&store) + .await + .unwrap() + .into_iter() + .filter(|schema| schema.id == database_schema.users_schema().id) + .collect::>(); + assert_eq!(user_schemas.len(), 1); + assert_eq!(user_schemas[0].snapshot_id, etl_postgres::types::SnapshotId::initial()); } #[tokio::test(flavor = "multi_thread")] diff --git a/crates/etl/tests/postgres_store.rs b/crates/etl/tests/postgres_store.rs index 48af29f4d..4bf38bb2f 100644 --- a/crates/etl/tests/postgres_store.rs +++ b/crates/etl/tests/postgres_store.rs @@ -836,7 +836,7 @@ async fn state_transitions_and_history() { } #[tokio::test(flavor = "multi_thread")] -async fn cleanup_deletes_state_schema_and_metadata_for_table() { +async fn delete_table_pipeline_state_deletes_state_schema_and_metadata_for_table() { init_test_tracing(); let database = spawn_source_database().await; @@ -844,17 +844,17 @@ async fn cleanup_deletes_state_schema_and_metadata_for_table() { let store = PostgresStore::new(pipeline_id, database.config.clone()).await.unwrap(); - // Test idempotency: cleanup on non-existent table should succeed + // Test idempotency: deleting state for a non-existent table should succeed. let nonexistent_table_id = TableId::new(99999); - store.cleanup_table_state(nonexistent_table_id).await.unwrap(); + store.delete_table_pipeline_state(nonexistent_table_id).await.unwrap(); - // Prepare two tables: one we will delete, one we will keep + // Prepare two tables: one we will delete, one we will keep. let table_1_schema = create_sample_table_schema(); let table_1_id = table_1_schema.id; let table_2_schema = create_another_table_schema(); let table_2_id = table_2_schema.id; - // Populate state, schema, and metadata for both tables + // Populate state, schema, and metadata for both tables. store.update_table_replication_state(table_1_id, TableReplicationPhase::Ready).await.unwrap(); store .update_table_replication_state(table_2_id, TableReplicationPhase::DataSync) @@ -878,41 +878,130 @@ async fn cleanup_deletes_state_schema_and_metadata_for_table() { store.store_destination_table_metadata(table_1_id, metadata1).await.unwrap(); store.store_destination_table_metadata(table_2_id, metadata2).await.unwrap(); - // Sanity check before cleanup + // Sanity check before deleting state. assert!(store.get_table_replication_state(table_1_id).await.unwrap().is_some()); assert!(store.get_table_schema(&table_1_id, SnapshotId::max()).await.unwrap().is_some()); assert!(store.get_applied_destination_table_metadata(table_1_id).await.unwrap().is_some()); - // Execute cleanup for table 1 - store.cleanup_table_state(table_1_id).await.unwrap(); + // Delete pipeline state for table 1. + store.delete_table_pipeline_state(table_1_id).await.unwrap(); - // Verify in-memory cache for table 1 has been cleaned + // Verify in-memory cache for table 1 has been deleted. assert!(store.get_table_replication_state(table_1_id).await.unwrap().is_none()); assert!(store.get_table_schema(&table_1_id, SnapshotId::max()).await.unwrap().is_none()); assert!(store.get_applied_destination_table_metadata(table_1_id).await.unwrap().is_none()); - // Verify other table is unaffected + // Verify other table is unaffected. assert!(store.get_table_replication_state(table_2_id).await.unwrap().is_some()); assert!(store.get_table_schema(&table_2_id, SnapshotId::max()).await.unwrap().is_some()); assert!(store.get_applied_destination_table_metadata(table_2_id).await.unwrap().is_some()); - // Create a new store instance and load from DB to ensure persistence + // Create a new store instance and load from DB to ensure persistence. let new_store = PostgresStore::new(pipeline_id, database.config.clone()).await.unwrap(); new_store.load_table_replication_states().await.unwrap(); new_store.load_table_schemas().await.unwrap(); new_store.load_destination_tables_metadata().await.unwrap(); - // Table 1 should not be present after reload + // Table 1 should not be present after reload. assert!(new_store.get_table_replication_state(table_1_id).await.unwrap().is_none()); assert!(new_store.get_table_schema(&table_1_id, SnapshotId::max()).await.unwrap().is_none()); assert!(new_store.get_applied_destination_table_metadata(table_1_id).await.unwrap().is_none()); - // Table 2 should still be present + // Table 2 should still be present. assert!(new_store.get_table_replication_state(table_2_id).await.unwrap().is_some()); assert!(new_store.get_table_schema(&table_2_id, SnapshotId::max()).await.unwrap().is_some()); assert!(new_store.get_applied_destination_table_metadata(table_2_id).await.unwrap().is_some()); } +#[tokio::test(flavor = "multi_thread")] +async fn clear_table_copy_state_keeps_replication_state_and_deletes_schema_metadata_and_progress() { + init_test_tracing(); + + let database = spawn_source_database().await; + let pipeline_id = 1; + + let store = PostgresStore::new(pipeline_id, database.config.clone()).await.unwrap(); + + // Test idempotency: clearing copy state for a non-existent table should + // succeed. + let nonexistent_table_id = TableId::new(99999); + store.clear_table_copy_state(nonexistent_table_id).await.unwrap(); + + let mut table_schema = create_sample_table_schema(); + let table_id = table_schema.id; + let other_table_schema = create_another_table_schema(); + let other_table_id = other_table_schema.id; + + store.update_table_replication_state(table_id, TableReplicationPhase::DataSync).await.unwrap(); + store + .update_table_replication_state(other_table_id, TableReplicationPhase::Ready) + .await + .unwrap(); + + table_schema.snapshot_id = SnapshotId::initial(); + store.store_table_schema(table_schema.clone()).await.unwrap(); + table_schema.snapshot_id = SnapshotId::from(100u64); + store.store_table_schema(table_schema).await.unwrap(); + store.store_table_schema(other_table_schema).await.unwrap(); + + let metadata = DestinationTableMetadata::new_applied( + "dest_table".to_owned(), + SnapshotId::from(100u64), + ReplicationMask::from_bytes(vec![1, 1, 1]), + ); + let other_metadata = DestinationTableMetadata::new_applied( + "other_dest_table".to_owned(), + SnapshotId::initial(), + ReplicationMask::from_bytes(vec![1, 1]), + ); + store.store_destination_table_metadata(table_id, metadata).await.unwrap(); + store.store_destination_table_metadata(other_table_id, other_metadata).await.unwrap(); + store + .upsert_replication_progress(WorkerType::TableSync { table_id }, PgLsn::from(200u64)) + .await + .unwrap(); + + store.clear_table_copy_state(table_id).await.unwrap(); + + assert_eq!( + store.get_table_replication_state(table_id).await.unwrap(), + Some(TableReplicationPhase::DataSync) + ); + assert!(store.get_table_schema(&table_id, SnapshotId::max()).await.unwrap().is_none()); + assert!(store.get_applied_destination_table_metadata(table_id).await.unwrap().is_none()); + assert!( + store.get_replication_progress(WorkerType::TableSync { table_id }).await.unwrap().is_none() + ); + + assert!(store.get_table_schema(&other_table_id, SnapshotId::max()).await.unwrap().is_some()); + assert!(store.get_applied_destination_table_metadata(other_table_id).await.unwrap().is_some()); + + let new_store = PostgresStore::new(pipeline_id, database.config.clone()).await.unwrap(); + new_store.load_table_replication_states().await.unwrap(); + new_store.load_table_schemas().await.unwrap(); + new_store.load_destination_tables_metadata().await.unwrap(); + + assert_eq!( + new_store.get_table_replication_state(table_id).await.unwrap(), + Some(TableReplicationPhase::DataSync) + ); + assert!(new_store.get_table_schema(&table_id, SnapshotId::max()).await.unwrap().is_none()); + assert!(new_store.get_applied_destination_table_metadata(table_id).await.unwrap().is_none()); + assert!( + new_store + .get_replication_progress(WorkerType::TableSync { table_id }) + .await + .unwrap() + .is_none() + ); + assert!( + new_store.get_table_schema(&other_table_id, SnapshotId::max()).await.unwrap().is_some() + ); + assert!( + new_store.get_applied_destination_table_metadata(other_table_id).await.unwrap().is_some() + ); +} + #[tokio::test(flavor = "multi_thread")] async fn replication_mask_loads_correctly_from_string_bytea() { init_test_tracing(); diff --git a/docs/explanation/architecture.md b/docs/explanation/architecture.md index f29c85f98..7e387b9d7 100644 --- a/docs/explanation/architecture.md +++ b/docs/explanation/architecture.md @@ -79,7 +79,7 @@ Where replicated data goes. Implement the `Destination` trait to send data anywh pub trait Destination { fn name() -> &'static str; fn shutdown(&self) -> impl Future> + Send { async { Ok(()) } } - fn truncate_table(&self, replicated_table_schema: &ReplicatedTableSchema, async_result: TruncateTableResult<()>) -> impl Future> + Send; + fn drop_table_for_copy(&self, replicated_table_schema: &ReplicatedTableSchema, async_result: DropTableForCopyResult<()>) -> impl Future> + Send; fn write_table_rows(&self, replicated_table_schema: &ReplicatedTableSchema, rows: Vec, async_result: WriteTableRowsResult<()>) -> impl Future> + Send; fn write_events(&self, events: Vec, async_result: WriteEventsResult<()>) -> impl Future> + Send; } @@ -88,14 +88,14 @@ pub trait Destination { | Method | When called | Purpose | |--------|-------------|---------| | `name()` | On initialization | Identify the destination | -| `truncate_table()` | Before initial copy | Clear destination table using the current replicated schema | +| `drop_table_for_copy()` | Before restarting a table copy when previous destination state exists | Drop the existing destination object and destination-private replay state using the previously stored replicated schema | | `write_table_rows()` | During initial copy | Receive bulk rows for the current replicated schema | | `write_events()` | After initial copy | Receive streaming changes | Each write-like method receives an async result handle. The intent is different per method: - `write_events()`: after dispatch succeeds, ETL may keep processing while the destination finishes the batch. -- `truncate_table()` and `write_table_rows()`: ETL waits for the result immediately. The handle is still useful because it keeps the destination API uniform and lets implementations reuse similar internal patterns. +- `drop_table_for_copy()` and `write_table_rows()`: ETL waits for the result immediately. After a successful drop, ETL clears its own copy-scoped schema, destination metadata, and table-sync progress before storing the fresh `0/0` copy schema. ### Store @@ -103,7 +103,7 @@ Persists pipeline state so replication can resume after restarts. Three traits w - **StateStore**: Tracks replication phase per table and destination table metadata - **SchemaStore**: Stores versioned table schema information (columns, types, primary keys, snapshot IDs) and prunes obsolete schema versions after acknowledged progress -- **CleanupStore**: Removes stored state when a table is dropped from the publication +- **CleanupStore**: Clears copy-scoped state before a table copy restart and removes all stored state when a table leaves the publication `StateStore` and `SchemaStore` use a cache-first pattern: reads hit an in-memory cache, writes go to both the cache and persistent storage. Schema pruning follows the same rule for implementations with durable storage: obsolete versions are removed from both the cache and the persistent store. diff --git a/docs/explanation/traits.md b/docs/explanation/traits.md index bd0140ee6..d5e1fff35 100644 --- a/docs/explanation/traits.md +++ b/docs/explanation/traits.md @@ -12,7 +12,7 @@ Receives replicated data. This is the primary extension point for sending data t pub trait Destination { fn name() -> &'static str; fn shutdown(&self) -> impl Future> + Send { async { Ok(()) } } - fn truncate_table(&self, replicated_table_schema: &ReplicatedTableSchema, async_result: TruncateTableResult<()>) -> impl Future> + Send; + fn drop_table_for_copy(&self, replicated_table_schema: &ReplicatedTableSchema, async_result: DropTableForCopyResult<()>) -> impl Future> + Send; fn write_table_rows(&self, replicated_table_schema: &ReplicatedTableSchema, table_rows: Vec, async_result: WriteTableRowsResult<()>) -> impl Future> + Send; fn write_events(&self, events: Vec, async_result: WriteEventsResult<()>) -> impl Future> + Send; } @@ -24,19 +24,17 @@ pub trait Destination { |--------|---------| | `name()` | Returns identifier for logging and diagnostics | | `shutdown()` | Called when the pipeline shuts down. Default is a no-op. Override for cleanup or bookkeeping | -| `truncate_table()` | Clears table data before initial sync. Receives the current replicated schema for the table | +| `drop_table_for_copy()` | Drops the existing destination object and destination-private replay state before restarting a table copy. Receives the previously stored replicated schema for locating the old object | | `write_table_rows()` | Writes rows during initial table copy. Receives the current replicated schema and may get an empty vector for tables with no data | | `write_events()` | Processes streaming replication events (inserts, updates, deletes). Batches may span multiple tables | ### Implementation Notes -- Operations should be idempotent when possible (ETL may retry on failure) +- `drop_table_for_copy()` should be idempotent. ETL calls it before clearing copy-scoped store state, so implementations can still use the supplied schema and existing destination metadata to locate the old object. +- Other operations should be idempotent when possible (ETL may retry on failure) - Handle concurrent calls safely (parallel table sync workers) - Process events in order to maintain data consistency -- All three write-like methods use async results, but ETL waits differently: -- `truncate_table()` waits immediately. -- `write_table_rows()` also waits immediately, requesting the next batch only after the current one finishes for that copy partition. -- `write_events()` is the method where ETL can keep processing while the destination finishes the current batch. +- All three write-like methods use async results, but ETL waits differently. `drop_table_for_copy()` waits immediately before copy-scoped store cleanup. `write_table_rows()` also waits immediately, requesting the next batch only after the current one finishes for that copy partition. `write_events()` is the method where ETL can keep processing while the destination finishes the current batch. See [Event Types](events.md) for details on the events received by `write_events()`. @@ -79,6 +77,11 @@ pub trait StateStore { fn update_table_replication_state(&self, table_id: TableId, state: TableReplicationPhase) -> impl Future> + Send; fn rollback_table_replication_state(&self, table_id: TableId) -> impl Future> + Send; + // Durable replication progress + fn get_replication_progress(&self, worker_type: WorkerType) -> impl Future>> + Send; + fn upsert_replication_progress(&self, worker_type: WorkerType, flush_lsn: PgLsn) -> impl Future> + Send; + fn delete_replication_progress(&self, worker_type: WorkerType) -> impl Future> + Send; + // Destination table metadata fn get_destination_table_metadata(&self, table_id: TableId) -> impl Future>> + Send; fn get_applied_destination_table_metadata(&self, table_id: TableId) -> impl Future>> + Send; @@ -98,6 +101,18 @@ pub trait StateStore { | `update_table_replication_state()` | Updates phase in both cache and persistent storage | | `rollback_table_replication_state()` | Reverts table to previous phase. Returns the phase after rollback | +### Durable Progress Methods + +Durable replication progress records the latest flushed LSN for the apply worker +and table sync workers. It lets ETL resume from a safe boundary even when a slot +or worker restarts. + +| Method | Purpose | +|--------|---------| +| `get_replication_progress()` | Returns stored flush progress for a worker, if present | +| `upsert_replication_progress()` | Monotonically stores flush progress and returns the stored LSN. Implementations must not move progress backward | +| `delete_replication_progress()` | Deletes progress when a worker slot lineage is intentionally reset | + ### Destination Metadata Methods Destination table metadata connects source table IDs to destination state, including the destination table identifier, the currently applied schema snapshot, and the replication mask. @@ -126,17 +141,19 @@ Tables progress through these phases: ## CleanupStore -Removes ETL metadata when tables are removed from the publication. +Removes ETL metadata for table-copy restarts and publication removals. ```rust pub trait CleanupStore { - fn cleanup_table_state(&self, table_id: TableId) -> impl Future> + Send; + fn clear_table_copy_state(&self, table_id: TableId) -> impl Future> + Send; + fn delete_table_pipeline_state(&self, table_id: TableId) -> impl Future> + Send; } ``` | Method | Purpose | |--------|---------| -| `cleanup_table_state()` | Deletes all stored state for a table: replication state, schema versions, and destination table metadata. Does not modify destination tables | +| `clear_table_copy_state()` | Clears destination metadata, schema versions, and durable table-sync progress while preserving the table replication phase. This is called only after the destination object was dropped for a fresh copy | +| `delete_table_pipeline_state()` | Deletes all stored state for a table removed from the publication. Does not modify destination tables | ## Combining Traits diff --git a/docs/guides/custom-implementations.md b/docs/guides/custom-implementations.md index ef544e31a..6782f8269 100644 --- a/docs/guides/custom-implementations.md +++ b/docs/guides/custom-implementations.md @@ -57,7 +57,7 @@ Create `src/custom_store.rs`. A store must implement three traits (see [Extensio - `SchemaStore` - Versioned table schema storage, retrieval, and pruning - `StateStore` - Replication progress and destination table metadata tracking -- `CleanupStore` - Store cleanup when tables leave the publication +- `CleanupStore` - Store cleanup for table-copy restarts and publication changes ```rust use std::collections::{BTreeMap, HashMap}; @@ -66,12 +66,13 @@ use tokio::sync::Mutex; use tracing::info; use etl::error::EtlResult; +use etl::replication::WorkerType; use etl::state::{ AppliedDestinationTableMetadata, DestinationTableMetadata, TableReplicationPhase, }; use etl::store::{CleanupStore, SchemaStore, StateStore, TableReplicationStates}; use etl::store::schema::TableSchemaRetention; -use etl::types::{SnapshotId, TableId, TableSchema}; +use etl::types::{PgLsn, SnapshotId, TableId, TableSchema}; #[derive(Debug, Clone, Default)] struct TableEntry { @@ -83,6 +84,7 @@ struct TableEntry { #[derive(Debug, Clone)] pub struct CustomStore { tables: Arc>>, + progress: Arc>>, } impl CustomStore { @@ -90,6 +92,7 @@ impl CustomStore { info!("creating custom store"); Self { tables: Arc::new(Mutex::new(HashMap::new())), + progress: Arc::new(Mutex::new(HashMap::new())), } } } @@ -213,6 +216,31 @@ impl StateStore for CustomStore { todo!("Implement rollback if needed") } + async fn get_replication_progress( + &self, + worker_type: WorkerType, + ) -> EtlResult> { + let progress = self.progress.lock().await; + Ok(progress.get(&worker_type).copied()) + } + + async fn upsert_replication_progress( + &self, + worker_type: WorkerType, + flush_lsn: PgLsn, + ) -> EtlResult { + let mut progress = self.progress.lock().await; + let stored_lsn = progress.entry(worker_type).or_insert(flush_lsn); + *stored_lsn = (*stored_lsn).max(flush_lsn); + Ok(*stored_lsn) + } + + async fn delete_replication_progress(&self, worker_type: WorkerType) -> EtlResult<()> { + let mut progress = self.progress.lock().await; + progress.remove(&worker_type); + Ok(()) + } + async fn get_destination_table_metadata( &self, table_id: TableId, @@ -249,9 +277,22 @@ impl StateStore for CustomStore { } impl CleanupStore for CustomStore { - async fn cleanup_table_state(&self, table_id: TableId) -> EtlResult<()> { + async fn clear_table_copy_state(&self, table_id: TableId) -> EtlResult<()> { + let mut tables = self.tables.lock().await; + if let Some(entry) = tables.get_mut(&table_id) { + entry.schemas.clear(); + entry.destination_metadata = None; + } + let mut progress = self.progress.lock().await; + progress.remove(&WorkerType::TableSync { table_id }); + Ok(()) + } + + async fn delete_table_pipeline_state(&self, table_id: TableId) -> EtlResult<()> { let mut tables = self.tables.lock().await; tables.remove(&table_id); + let mut progress = self.progress.lock().await; + progress.remove(&WorkerType::TableSync { table_id }); Ok(()) } } @@ -264,12 +305,14 @@ impl CleanupStore for CustomStore { Create `src/http_destination.rs`. A destination implements the `Destination` trait with four required methods: - `name()` - Return an identifier for logging -- `truncate_table()` - Clear table before bulk load using the current replicated table schema +- `drop_table_for_copy()` - Idempotently drop destination objects and replay state before restarting a table copy using the previously stored replicated table schema - `write_table_rows()` - Receive rows during initial copy together with the current replicated table schema - `write_events()` - Receive streaming changes (batches may span multiple tables) There's also an optional `shutdown()` method with a default no-op implementation. Override it if your destination needs cleanup when the pipeline shuts down. +ETL clears its own schema versions, destination metadata, and table-sync progress only after `drop_table_for_copy()` succeeds. That lets the destination use the supplied previously stored replicated schema and any existing destination metadata to find the object that must be removed. If the object is already gone, return success. + ```rust use reqwest::Client; use serde_json::json; @@ -277,7 +320,7 @@ use std::time::Duration; use tracing::{info, warn}; use etl::destination::{ - Destination, TruncateTableResult, WriteEventsResult, WriteTableRowsResult, + Destination, DropTableForCopyResult, WriteEventsResult, WriteTableRowsResult, }; use etl::error::{ErrorKind, EtlResult}; use etl::types::{Event, ReplicatedTableSchema, TableRow}; @@ -323,15 +366,15 @@ impl Destination for HttpDestination { "http" } - async fn truncate_table( + async fn drop_table_for_copy( &self, replicated_table_schema: &ReplicatedTableSchema, - async_result: TruncateTableResult<()>, + async_result: DropTableForCopyResult<()>, ) -> EtlResult<()> { let table_name = replicated_table_schema.name().to_string(); - info!("truncating table {}", table_name); + info!("dropping table before copy {}", table_name); let result = self - .post(&format!("tables/{table_name}/truncate"), json!({})) + .post(&format!("tables/{table_name}/drop-for-copy"), json!({})) .await; async_result.send(result); Ok(()) diff --git a/docs/guides/first-pipeline.md b/docs/guides/first-pipeline.md index 66f3b6eb2..7800f8539 100644 --- a/docs/guides/first-pipeline.md +++ b/docs/guides/first-pipeline.md @@ -72,7 +72,7 @@ use etl::{ BatchConfig, InvalidatedSlotBehavior, MemoryBackpressureConfig, PgConnectionConfig, PipelineConfig, TableSyncCopyConfig, TcpKeepaliveConfig, TlsConfig, }, - destination::{Destination, TruncateTableResult, WriteEventsResult, WriteTableRowsResult}, + destination::{Destination, DropTableForCopyResult, WriteEventsResult, WriteTableRowsResult}, error::EtlResult, pipeline::Pipeline, store::MemoryStore, @@ -88,12 +88,12 @@ impl Destination for LoggingDestination { "logging" } - async fn truncate_table( + async fn drop_table_for_copy( &self, _replicated_table_schema: &ReplicatedTableSchema, - async_result: TruncateTableResult<()>, + async_result: DropTableForCopyResult<()>, ) -> EtlResult<()> { - println!("starting initial table copy"); + println!("preparing fresh table copy"); async_result.send(Ok(())); Ok(()) } diff --git a/docs/index.md b/docs/index.md index 314e2f269..0f7388e7a 100644 --- a/docs/index.md +++ b/docs/index.md @@ -50,7 +50,7 @@ use etl::{ BatchConfig, InvalidatedSlotBehavior, MemoryBackpressureConfig, PgConnectionConfig, PipelineConfig, TableSyncCopyConfig, TcpKeepaliveConfig, TlsConfig, }, - destination::{Destination, TruncateTableResult, WriteEventsResult, WriteTableRowsResult}, + destination::{Destination, DropTableForCopyResult, WriteEventsResult, WriteTableRowsResult}, error::EtlResult, pipeline::Pipeline, store::MemoryStore, @@ -65,10 +65,10 @@ impl Destination for NoopDestination { "noop" } - async fn truncate_table( + async fn drop_table_for_copy( &self, _replicated_table_schema: &ReplicatedTableSchema, - async_result: TruncateTableResult<()>, + async_result: DropTableForCopyResult<()>, ) -> EtlResult<()> { async_result.send(Ok(())); Ok(()) From 46463d439b0eb75044821a5748ab32cf69409c4c Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Fri, 22 May 2026 10:51:48 +0200 Subject: [PATCH 24/29] fix(core): Fix Snowflake drop impl and ci (#761) --- .github/workflows/ci.yml | 1 + DEVELOPMENT.md | 2 +- .../etl-destinations/src/snowflake/client.rs | 27 +++++++++++++++++++ crates/etl-destinations/src/snowflake/core.rs | 12 +++++---- scripts/run_migrations.sh | 2 +- 5 files changed, 37 insertions(+), 7 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index b66e1beff..0d5f034bc 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -212,6 +212,7 @@ jobs: - name: Install sqlx-cli run: | cargo install sqlx-cli \ + --version 0.9.0-alpha.1 \ --features native-tls,postgres \ --no-default-features \ --locked diff --git a/DEVELOPMENT.md b/DEVELOPMENT.md index 0ccf65444..ac2912073 100644 --- a/DEVELOPMENT.md +++ b/DEVELOPMENT.md @@ -31,7 +31,7 @@ Before starting, ensure you have the following installed: Install SQLx CLI: ```bash -cargo install --version='~0.8.6' sqlx-cli --no-default-features --features rustls,postgres +cargo install --version 0.9.0-alpha.1 sqlx-cli --no-default-features --features rustls,postgres --locked ``` ### Optional Tools diff --git a/crates/etl-destinations/src/snowflake/client.rs b/crates/etl-destinations/src/snowflake/client.rs index eee25d140..97e7d057b 100644 --- a/crates/etl-destinations/src/snowflake/client.rs +++ b/crates/etl-destinations/src/snowflake/client.rs @@ -1,6 +1,7 @@ use std::{collections::HashMap, sync::Arc}; use etl::types::{ColumnSchema, PipelineId, SchemaDiff, TableId}; +use reqwest::StatusCode; use tokio::sync::{Mutex, RwLock}; use crate::snowflake::{ @@ -157,6 +158,32 @@ impl Client { guard.reset().await } + /// Drop the table and destination-private replay state before a fresh copy. + pub async fn drop_table_for_copy(&self, table_id: TableId, table_name: &str) -> Result<()> { + let channel = self.channels.write().await.remove(&table_id); + + let drop_channel_result = if let Some(channel) = channel { + let mut guard = channel.lock().await; + guard.drop_channel().await + } else { + let mut handle = ChannelHandle::new( + Arc::clone(&self.stream_client), + self.pipeline_id, + self.database.clone(), + self.schema.clone(), + table_name.to_owned(), + ); + handle.drop_channel().await + }; + match drop_channel_result { + Ok(()) => {} + Err(Error::HttpStatus { status: StatusCode::NOT_FOUND, .. }) => {} + Err(error) => return Err(error), + } + + self.sql_client.drop_table(table_name).await + } + /// Refresh the table's ingestion state after a schema change. /// /// Channels must be reopened after ALTER TABLE so Snowpipe picks up the diff --git a/crates/etl-destinations/src/snowflake/core.rs b/crates/etl-destinations/src/snowflake/core.rs index a51a3ab32..089f18b47 100644 --- a/crates/etl-destinations/src/snowflake/core.rs +++ b/crates/etl-destinations/src/snowflake/core.rs @@ -3,7 +3,7 @@ use std::collections::HashMap; use etl::{ bail, concurrency::TaskSet, - destination::async_result::{TruncateTableResult, WriteEventsResult, WriteTableRowsResult}, + destination::async_result::{DropTableForCopyResult, WriteEventsResult, WriteTableRowsResult}, error::{ErrorKind, EtlError, EtlResult}, etl_error, state::destination_metadata::{DestinationTableMetadata, DestinationTableSchemaStatus}, @@ -429,19 +429,21 @@ where self.tasks.shutdown().await } - async fn truncate_table( + async fn drop_table_for_copy( &self, replicated_table_schema: &ReplicatedTableSchema, - async_result: TruncateTableResult<()>, + async_result: DropTableForCopyResult<()>, ) -> EtlResult<()> { - self.prepare_table_for_streaming(replicated_table_schema).await?; + self.tasks.try_reap().await?; + let table_name = try_stringify_table_name(replicated_table_schema.name())?.to_uppercase(); let result = self .client - .truncate_table(replicated_table_schema.id(), &table_name) + .drop_table_for_copy(replicated_table_schema.id(), &table_name) .await .map_err(EtlError::from); async_result.send(result); + Ok(()) } diff --git a/scripts/run_migrations.sh b/scripts/run_migrations.sh index b0f415b75..e8f1ffa92 100755 --- a/scripts/run_migrations.sh +++ b/scripts/run_migrations.sh @@ -28,7 +28,7 @@ check_sqlx() { if ! [ -x "$(command -v sqlx)" ]; then echo >&2 "Error: SQLx CLI is not installed." echo >&2 "To install it, run:" - echo >&2 " cargo install --version='~0.7' sqlx-cli --no-default-features --features rustls,postgres" + echo >&2 " cargo install --version 0.9.0-alpha.1 sqlx-cli --no-default-features --features rustls,postgres --locked" exit 1 fi } From bb09675430b4392588ac50ffe8fb4e9284ec14e5 Mon Sep 17 00:00:00 2001 From: ttatsato Date: Fri, 22 May 2026 18:07:33 +0900 Subject: [PATCH 25/29] feat(etl-api): Add publication endpoints (#748) --- crates/etl-api/src/data/publications.rs | 34 ++++ .../src/routes/sources/publications.rs | 148 +++++++++++++++++- crates/etl-api/src/startup.rs | 12 +- 3 files changed, 191 insertions(+), 3 deletions(-) diff --git a/crates/etl-api/src/data/publications.rs b/crates/etl-api/src/data/publications.rs index 4f492b07e..27287027e 100644 --- a/crates/etl-api/src/data/publications.rs +++ b/crates/etl-api/src/data/publications.rs @@ -152,3 +152,37 @@ pub async fn read_all_publications(pool: &PgPool) -> Result, Pu Ok(publications) } + +pub async fn add_tables_to_publication( + publication: &Publication, + pool: &PgPool, +) -> Result<(), PublicationsDbError> { + let query = format!( + "alter publication {} add table only {}", + quote_identifier(&publication.name), + format_table_list(&publication.tables), + ); + sqlx::query(AssertSqlSafe(query)).execute(pool).await?; + Ok(()) +} + +pub async fn drop_tables_from_publication( + publication: &Publication, + pool: &PgPool, +) -> Result<(), PublicationsDbError> { + let query = format!( + "alter publication {} drop table only {}", + quote_identifier(&publication.name), + format_table_list(&publication.tables), + ); + sqlx::query(AssertSqlSafe(query)).execute(pool).await?; + Ok(()) +} + +fn format_table_list(tables: &[Table]) -> String { + tables + .iter() + .map(|t| format!("{}.{}", quote_identifier(&t.schema), quote_identifier(&t.name))) + .collect::>() + .join(", ") +} diff --git a/crates/etl-api/src/routes/sources/publications.rs b/crates/etl-api/src/routes/sources/publications.rs index bf3a37a6c..1d8f8edcf 100644 --- a/crates/etl-api/src/routes/sources/publications.rs +++ b/crates/etl-api/src/routes/sources/publications.rs @@ -1,7 +1,7 @@ use actix_web::{ HttpRequest, HttpResponse, Responder, ResponseError, delete, get, http::{StatusCode, header::ContentType}, - post, + post, put, web::{Data, Json, Path}, }; use serde::{Deserialize, Serialize}; @@ -226,6 +226,11 @@ pub(crate) async fn update_publication( let source_pool = connect_to_source_database_from_api(&source_config.into_connection_config(tls_config)) .await?; + + if data::publications::read_publication(&publication_name, &source_pool).await?.is_none() { + return Err(PublicationError::PublicationNotFound(publication_name)); + } + let publication = publication.0; let publication = Publication { name: publication_name, tables: publication.tables }; data::publications::update_publication(&publication, &source_pool).await?; @@ -320,3 +325,144 @@ pub(crate) async fn read_all_publications( Ok(Json(response)) } + +#[utoipa::path( + summary = "Add tables to a publication", + description = "Adds the specified tables to an existing publication.", + tag = "Publications", + request_body = UpdatePublicationRequest, + params( + ("source_id" = i64, Path, description = "Unique ID of the source"), + ("publication_name" = String, Path, description = "Publication name within the source"), + ), + responses( + (status = 200, description = "Tables added successfully"), + (status = 404, description = "Publication not found", body = ErrorMessage), + (status = 500, description = "Internal server error", body = ErrorMessage) + ) +)] +#[post("/sources/{source_id}/publications/{publication_name}/tables")] +pub(crate) async fn add_tables_to_publication( + req: HttpRequest, + pool: Data, + api_config: Data, + encryption_key: Data, + trusted_root_certs_cache: Data, + source_id_and_pub_name: Path<(i64, String)>, + publication: Json, +) -> Result { + let tenant_id = extract_tenant_id(&req)?; + let (source_id, publication_name) = source_id_and_pub_name.into_inner(); + let source_config = data::sources::read_source(&**pool, tenant_id, source_id, &encryption_key) + .await? + .map(|s| s.config) + .ok_or(PublicationError::SourceNotFound(source_id))?; + let tls_config = trusted_root_certs_cache.get_tls_config(api_config.source.tls_enabled).await?; + let source_pool = + connect_to_source_database_from_api(&source_config.into_connection_config(tls_config)) + .await?; + + if data::publications::read_publication(&publication_name, &source_pool).await?.is_none() { + return Err(PublicationError::PublicationNotFound(publication_name)); + } + + let publication = publication.0; + let publication = Publication { name: publication_name, tables: publication.tables }; + data::publications::add_tables_to_publication(&publication, &source_pool).await?; + + Ok(HttpResponse::Ok().finish()) +} + +#[utoipa::path( + summary = "Remove tables from a publication", + description = "Removes the specified tables from an existing publication.", + tag = "Publications", + request_body = UpdatePublicationRequest, + params( + ("source_id" = i64, Path, description = "Unique ID of the source"), + ("publication_name" = String, Path, description = "Publication name within the source"), + ), + responses( + (status = 200, description = "Tables removed successfully"), + (status = 404, description = "Publication not found", body = ErrorMessage), + (status = 500, description = "Internal server error", body = ErrorMessage) + ) +)] +#[delete("/sources/{source_id}/publications/{publication_name}/tables")] +pub(crate) async fn drop_tables_from_publication( + req: HttpRequest, + pool: Data, + api_config: Data, + encryption_key: Data, + trusted_root_certs_cache: Data, + source_id_and_pub_name: Path<(i64, String)>, + publication: Json, +) -> Result { + let tenant_id = extract_tenant_id(&req)?; + let (source_id, publication_name) = source_id_and_pub_name.into_inner(); + let source_config = data::sources::read_source(&**pool, tenant_id, source_id, &encryption_key) + .await? + .map(|s| s.config) + .ok_or(PublicationError::SourceNotFound(source_id))?; + let tls_config = trusted_root_certs_cache.get_tls_config(api_config.source.tls_enabled).await?; + let source_pool = + connect_to_source_database_from_api(&source_config.into_connection_config(tls_config)) + .await?; + + if data::publications::read_publication(&publication_name, &source_pool).await?.is_none() { + return Err(PublicationError::PublicationNotFound(publication_name)); + } + + let publication = publication.0; + let publication = Publication { name: publication_name, tables: publication.tables }; + data::publications::drop_tables_from_publication(&publication, &source_pool).await?; + + Ok(HttpResponse::Ok().finish()) +} + +#[utoipa::path( + summary = "Replace tables of a publication", + description = "Replaces the table list of an existing publication with the specified tables.", + tag = "Publications", + request_body = UpdatePublicationRequest, + params( + ("source_id" = i64, Path, description = "Unique ID of the source"), + ("publication_name" = String, Path, description = "Publication name within the source"), + ), + responses( + (status = 200, description = "Tables replaced successfully"), + (status = 404, description = "Publication not found", body = ErrorMessage), + (status = 500, description = "Internal server error", body = ErrorMessage) + ) +)] +#[put("/sources/{source_id}/publications/{publication_name}/tables")] +pub(crate) async fn set_publication_tables( + req: HttpRequest, + pool: Data, + api_config: Data, + encryption_key: Data, + trusted_root_certs_cache: Data, + source_id_and_pub_name: Path<(i64, String)>, + publication: Json, +) -> Result { + let tenant_id = extract_tenant_id(&req)?; + let (source_id, publication_name) = source_id_and_pub_name.into_inner(); + let source_config = data::sources::read_source(&**pool, tenant_id, source_id, &encryption_key) + .await? + .map(|s| s.config) + .ok_or(PublicationError::SourceNotFound(source_id))?; + let tls_config = trusted_root_certs_cache.get_tls_config(api_config.source.tls_enabled).await?; + let source_pool = + connect_to_source_database_from_api(&source_config.into_connection_config(tls_config)) + .await?; + + if data::publications::read_publication(&publication_name, &source_pool).await?.is_none() { + return Err(PublicationError::PublicationNotFound(publication_name)); + } + + let publication = publication.0; + let publication = Publication { name: publication_name, tables: publication.tables }; + data::publications::update_publication(&publication, &source_pool).await?; + + Ok(HttpResponse::Ok().finish()) +} diff --git a/crates/etl-api/src/startup.rs b/crates/etl-api/src/startup.rs index 40a9a2d68..5612f4576 100644 --- a/crates/etl-api/src/startup.rs +++ b/crates/etl-api/src/startup.rs @@ -58,8 +58,10 @@ use crate::{ UpdateSourceRequest, ValidateSourceRequest, ValidateSourceResponse, create_source, delete_source, publications::{ - CreatePublicationRequest, UpdatePublicationRequest, create_publication, - delete_publication, read_all_publications, read_publication, update_publication, + CreatePublicationRequest, UpdatePublicationRequest, add_tables_to_publication, + create_publication, delete_publication, drop_tables_from_publication, + read_all_publications, read_publication, set_publication_tables, + update_publication, }, read_all_sources, read_source, tables::read_table_names, @@ -325,6 +327,9 @@ pub fn run( crate::routes::sources::publications::update_publication, crate::routes::sources::publications::delete_publication, crate::routes::sources::publications::read_all_publications, + crate::routes::sources::publications::add_tables_to_publication, + crate::routes::sources::publications::drop_tables_from_publication, + crate::routes::sources::publications::set_publication_tables, crate::routes::sources::tables::read_table_names, crate::routes::destinations::create_destination, crate::routes::destinations::read_destination, @@ -421,6 +426,9 @@ pub fn run( .service(update_publication) .service(delete_publication) .service(read_all_publications) + .service(add_tables_to_publication) + .service(drop_tables_from_publication) + .service(set_publication_tables) // tenants_sources .service(create_tenant_and_source) // destinations-pipelines From 71e78fabf0a642439c2245defa9060e1c3df94ac Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Fri, 22 May 2026 11:28:11 +0200 Subject: [PATCH 26/29] ref(api): Remove usage of 'only' (#763) --- crates/etl-api/src/data/publications.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/crates/etl-api/src/data/publications.rs b/crates/etl-api/src/data/publications.rs index 27287027e..200e10b80 100644 --- a/crates/etl-api/src/data/publications.rs +++ b/crates/etl-api/src/data/publications.rs @@ -29,7 +29,7 @@ pub async fn create_publication( query.push_str("create publication "); query.push_str("ed_publication_name); if !publication.tables.is_empty() { - query.push_str(" for table only "); + query.push_str(" for table "); } for (i, table) in publication.tables.iter().enumerate() { @@ -60,7 +60,7 @@ pub async fn update_publication( let quoted_publication_name = quote_identifier(&publication.name); query.push_str("alter publication "); query.push_str("ed_publication_name); - query.push_str(" set table only "); + query.push_str(" set table "); for (i, table) in publication.tables.iter().enumerate() { let quoted_schema = quote_identifier(&table.schema); @@ -158,7 +158,7 @@ pub async fn add_tables_to_publication( pool: &PgPool, ) -> Result<(), PublicationsDbError> { let query = format!( - "alter publication {} add table only {}", + "alter publication {} add table {}", quote_identifier(&publication.name), format_table_list(&publication.tables), ); @@ -171,7 +171,7 @@ pub async fn drop_tables_from_publication( pool: &PgPool, ) -> Result<(), PublicationsDbError> { let query = format!( - "alter publication {} drop table only {}", + "alter publication {} drop table {}", quote_identifier(&publication.name), format_table_list(&publication.tables), ); From 04240fc627d8f28b2c6fe4fabf62eff4ec9f89c2 Mon Sep 17 00:00:00 2001 From: Victor Farazdagi Date: Fri, 22 May 2026 15:28:59 +0300 Subject: [PATCH 27/29] feat(xtask): add fmt, check, fix, and msrv commands (#764) * feat(xtask): add fmt, check, fix, and msrv commands Port shell scripts to xshell-based xtask commands: - fmt: nightly rustfmt with --check flag - check: pre-PR gate running fmt, sort, clippy, nextest - fix: auto-fix with clippy --fix, fmt, sort - msrv: verify MSRV consistency with --verify flag for full compilation check Also adds cargo-msrv presence check with actionable install instructions. * refactor(xtask): no nextest on check * fix: eprintln * feat: passthough other scripts * fmt --- .cargo/config.toml | 1 + AGENTS.md | 8 +- Cargo.lock | 47 +++++++++++ DEVELOPMENT.md | 53 +++++++----- crates/etl-api/README.md | 2 +- crates/etl-examples/README.md | 2 +- .../etl-examples/src/bin/snowflake/README.md | 2 +- crates/xtask/Cargo.toml | 18 ++++- crates/xtask/src/commands/check.rs | 28 +++++++ crates/xtask/src/commands/deploy_local.rs | 20 +++++ crates/xtask/src/commands/fix.rs | 28 +++++++ crates/xtask/src/commands/fmt.rs | 29 +++++++ crates/xtask/src/commands/init.rs | 15 ++++ crates/xtask/src/commands/migrate.rs | 20 +++++ crates/xtask/src/commands/mod.rs | 18 +++++ crates/xtask/src/commands/msrv.rs | 81 +++++++++++++++++++ crates/xtask/src/commands/shared.rs | 1 + crates/xtask/src/commands/test_clickhouse.rs | 15 ++++ crates/xtask/src/commands/vendor_duckdb.rs | 20 +++++ crates/xtask/src/main.rs | 41 +++++++++- scripts/README.md | 17 ++++ 21 files changed, 436 insertions(+), 30 deletions(-) create mode 100644 crates/xtask/src/commands/check.rs create mode 100644 crates/xtask/src/commands/deploy_local.rs create mode 100644 crates/xtask/src/commands/fix.rs create mode 100644 crates/xtask/src/commands/fmt.rs create mode 100644 crates/xtask/src/commands/init.rs create mode 100644 crates/xtask/src/commands/migrate.rs create mode 100644 crates/xtask/src/commands/msrv.rs create mode 100644 crates/xtask/src/commands/test_clickhouse.rs create mode 100644 crates/xtask/src/commands/vendor_duckdb.rs create mode 100644 scripts/README.md diff --git a/.cargo/config.toml b/.cargo/config.toml index af652b823..c5bcbe684 100644 --- a/.cargo/config.toml +++ b/.cargo/config.toml @@ -1,5 +1,6 @@ [alias] xtask = "run --package xtask --" +x = "run --package xtask --" # Use lld linker for faster linking (1.5-2x faster than default ld) # lld is part of the LLVM toolchain and widely available diff --git a/AGENTS.md b/AGENTS.md index 7f57e97b9..207007e10 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -13,16 +13,16 @@ - `crates/etl-benchmarks/`: benchmarks. - `crates/xtask/`: workspace automation commands. - Docs live in `docs/`. -- Local development and ops tooling live in `scripts/` and `DEVELOPMENT.md`. +- Local development and ops tooling live in `crates/xtask/` (run via `cargo x`) and `DEVELOPMENT.md`. - Tests live next to code in `src/` or `tests/`. ## Commands - Build everything: - `cargo build --workspace --all-targets --all-features` - Format: - - `./scripts/fmt` + - `cargo x fmt` - Check formatting: - - `./scripts/fmt-check` + - `cargo x fmt --check` - Lint: - `cargo clippy --workspace --all-targets --all-features -- -D warnings` - Run unit tests (no Postgres required): @@ -43,7 +43,7 @@ - Do not add dependencies unless they are justified by the task. - If you change workflow assumptions, build or test the smallest relevant target and report what actually ran. - Never create commits, push branches, open pull requests, or perform other git write actions unless the user explicitly instructs you to do so. -- Keep the workspace on the stable toolchain from `rust-toolchain.toml` for build, lint, and test commands; use the pinned nightly formatter only through `./scripts/fmt` and `./scripts/fmt-check`. +- Keep the workspace on the stable toolchain from `rust-toolchain.toml` for build, lint, and test commands; use the pinned nightly formatter only through `cargo x fmt` and `cargo x fmt --check`. - Treat `Cargo.toml` workspace lints, `rustfmt.toml`, and compiler diagnostics as the source of truth for enforceable style and correctness rules. Prefer adding or tightening static checks over adding prose rules here. - Run Clippy, builds, and tests intentionally when they are relevant: for example after changing Rust code, when compiler/lint diagnostics indicate a problem, diff --git a/Cargo.lock b/Cargo.lock index 6c5e74657..2d1731860 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6083,6 +6083,15 @@ dependencies = [ "syn 2.0.117", ] +[[package]] +name = "serde_spanned" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6662b5879511e06e8999a8a235d848113e942c9124f211511b16466ee2995f26" +dependencies = [ + "serde_core", +] + [[package]] name = "serde_urlencoded" version = "0.7.1" @@ -6960,6 +6969,21 @@ dependencies = [ "tokio", ] +[[package]] +name = "toml" +version = "1.1.2+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "81f3d15e84cbcd896376e6730314d59fb5a87f31e4b038454184435cd57defee" +dependencies = [ + "indexmap 2.14.0", + "serde_core", + "serde_spanned", + "toml_datetime", + "toml_parser", + "toml_writer", + "winnow", +] + [[package]] name = "toml_datetime" version = "1.1.1+spec-1.1.0" @@ -6990,6 +7014,12 @@ dependencies = [ "winnow", ] +[[package]] +name = "toml_writer" +version = "1.1.1+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "756daf9b1013ebe47a8776667b466417e2d4c5679d441c26230efd9ef78692db" + [[package]] name = "tonic" version = "0.14.6" @@ -8149,6 +8179,21 @@ dependencies = [ "rustix 1.1.4", ] +[[package]] +name = "xshell" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e7290c623014758632efe00737145b6867b66292c42167f2ec381eb566a373d" +dependencies = [ + "xshell-macros", +] + +[[package]] +name = "xshell-macros" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32ac00cd3f8ec9c1d33fb3e7958a82df6989c42d747bd326c822b1d625283547" + [[package]] name = "xtask" version = "0.1.0" @@ -8162,6 +8207,8 @@ dependencies = [ "serde", "serde_json", "tokio", + "toml", + "xshell", ] [[package]] diff --git a/DEVELOPMENT.md b/DEVELOPMENT.md index ac2912073..de8368216 100644 --- a/DEVELOPMENT.md +++ b/DEVELOPMENT.md @@ -4,6 +4,7 @@ This guide covers setting up your development environment, running migrations, a ## Table of Contents +- [Task Runner](#task-runner) - [Prerequisites](#prerequisites) - [Quick Start](#quick-start) - [Database Setup](#database-setup) @@ -40,29 +41,45 @@ cargo install --version 0.9.0-alpha.1 sqlx-cli --no-default-features --features - [Install OrbStack](https://orbstack.dev) - Enable Kubernetes in OrbStack settings +## Task Runner + +Common development tasks are available through `cargo x`, a shorthand alias for `cargo xtask`. +Run `cargo x --help` to see all available commands. + +```bash +cargo x fmt # format code with nightly rustfmt +cargo x fmt --check # check formatting without changes +cargo x check # pre-PR gate: fmt, sort, clippy +cargo x fix # auto-fix: clippy --fix, fmt, sort +cargo x msrv # verify MSRV consistency +cargo x init # set up local dev environment +cargo x migrate # run database migrations +cargo x deploy-local # deploy replicator to local OrbStack k8s +cargo x test-clickhouse # run ClickHouse integration tests +cargo x vendor-duckdb # download and vendor DuckDB extensions +``` + ## Formatting The workspace stays on the stable toolchain pinned in `rust-toolchain.toml` for builds, tests, and linting. Formatting is the only workflow that uses nightly Rust, because the repository relies on nightly-only `rustfmt` options for import grouping and layout. -Use the pinned formatter scripts from the project root: - ```bash -./scripts/fmt -./scripts/fmt-check +cargo x fmt +cargo x fmt --check ``` -Both scripts default to `nightly-2026-04-15`. You can temporarily override the formatter toolchain with +Both default to `nightly-2026-04-15`. You can temporarily override the formatter toolchain with `RUSTFMT_NIGHTLY_TOOLCHAIN`, but CI and the repository defaults should stay pinned so formatting does not drift. ## Quick Start -The fastest way to get started is using the setup script: +The fastest way to get started: ```bash # From the project root -./scripts/init.sh +cargo x init ``` This script will: @@ -75,20 +92,20 @@ This script will: ### Using the Setup Script -The `scripts/init.sh` script provides a complete development environment setup: +`cargo x init` provides a complete development environment setup: ```bash # Use default settings (Postgres on port 5430) -./scripts/init.sh +cargo x init # Customize database settings -POSTGRES_PORT=5432 POSTGRES_DB=mydb ./scripts/init.sh +POSTGRES_PORT=5432 POSTGRES_DB=mydb cargo x init # Skip Docker if you already have Postgres running -SKIP_DOCKER=1 ./scripts/init.sh +SKIP_DOCKER=1 cargo x init # Use persistent storage -POSTGRES_DATA_VOLUME=/path/to/data ./scripts/init.sh +POSTGRES_DATA_VOLUME=/path/to/data cargo x init ``` **Environment Variables:** @@ -130,7 +147,7 @@ If using one database for both the API and ETL source/store objects: export DATABASE_URL=postgres://USER:PASSWORD@HOST:PORT/DB # Run all migrations on the same database -./scripts/run_migrations.sh +cargo x migrate ``` #### Separate Database Setup @@ -140,11 +157,11 @@ If using separate databases (recommended for production): ```bash # API migrations on the control plane database export DATABASE_URL=postgres://USER:PASSWORD@API_HOST:PORT/API_DB -./scripts/run_migrations.sh etl-api +cargo x migrate etl-api # ETL migrations on the source database export DATABASE_URL=postgres://USER:PASSWORD@SOURCE_HOST:PORT/SOURCE_DB -./scripts/run_migrations.sh etl +cargo x migrate etl ``` This separation allows you to: @@ -164,7 +181,7 @@ Located in `crates/etl-api/migrations/`, these create the control plane schema ( ```bash # From project root -./scripts/run_migrations.sh etl-api +cargo x migrate etl-api # Or manually with SQLx CLI sqlx migrate run --source crates/etl-api/migrations @@ -211,7 +228,7 @@ mismatch. ```bash # From project root -./scripts/run_migrations.sh etl +cargo x migrate etl # Or manually with SQLx CLI (requires setting search_path) psql $DATABASE_URL -c "create schema if not exists etl;" @@ -425,7 +442,7 @@ ClickHouse destination tests require a reachable ClickHouse HTTP endpoint: | `TESTS_CLICKHOUSE_USER` | **Yes** | ClickHouse user name (for the local Docker Compose setup, use `etl`) | | `TESTS_CLICKHOUSE_PASSWORD` | No | ClickHouse password; for the local Docker Compose setup, use `etl` | -**Note:** ClickHouse tests are only run when the `clickhouse` and `test-utils` features are enabled. Each test creates a unique database in ClickHouse and drops it automatically when the test finishes. The Docker Compose setup started by `./scripts/init.sh` is sufficient for these tests. +**Note:** ClickHouse tests are only run when the `clickhouse` and `test-utils` features are enabled. Each test creates a unique database in ClickHouse and drops it automatically when the test finishes. The Docker Compose setup started by `cargo x init` is sufficient for these tests. #### Test Output and Logging diff --git a/crates/etl-api/README.md b/crates/etl-api/README.md index 6b4b32bc5..55ce8c5bf 100644 --- a/crates/etl-api/README.md +++ b/crates/etl-api/README.md @@ -36,7 +36,7 @@ For the full local development stack, use the setup script to start Postgres, run migrations, and apply the local Kubernetes resources. ```bash -./scripts/init.sh +cargo x init ``` Alternative: if you already have a Postgres database, set `DATABASE_URL` and apply migrations manually: diff --git a/crates/etl-examples/README.md b/crates/etl-examples/README.md index cc98bfaee..db90c709d 100644 --- a/crates/etl-examples/README.md +++ b/crates/etl-examples/README.md @@ -259,7 +259,7 @@ For offline local development on Linux or macOS, you can prefetch the required DuckDB extensions into the repository and point the destination at them: ```bash -./scripts/vendor_duckdb_extensions.sh +cargo x vendor-duckdb ETL_DUCKDB_EXTENSION_ROOT="$(pwd)/vendor/duckdb/extensions" \ cargo run --bin ducklake -p etl-examples --features ducklake -- [flags] ``` diff --git a/crates/etl-examples/src/bin/snowflake/README.md b/crates/etl-examples/src/bin/snowflake/README.md index 6b183e3da..6d3838500 100644 --- a/crates/etl-examples/src/bin/snowflake/README.md +++ b/crates/etl-examples/src/bin/snowflake/README.md @@ -20,7 +20,7 @@ Demonstrates the Snowflake destination's performance characteristics. From the repo root: ```bash -./scripts/init.sh +cargo x init ``` This starts `source-postgres` on port 5430 with `wal_level=logical` and all replication settings configured. See `scripts/docker-compose.yaml` for details. diff --git a/crates/xtask/Cargo.toml b/crates/xtask/Cargo.toml index d621b6472..ba8f58a9e 100644 --- a/crates/xtask/Cargo.toml +++ b/crates/xtask/Cargo.toml @@ -6,14 +6,28 @@ publish = false [dependencies] anyhow = { workspace = true } -clap = { workspace = true, features = ["derive", "std", "help", "usage", "error-context", "env"] } +clap = { workspace = true, features = [ + "derive", + "std", + "help", + "usage", + "error-context", + "env", +] } k8s-openapi = { workspace = true, features = ["latest"] } -kube = { workspace = true, features = ["client", "derive", "rustls-tls", "ring"] } +kube = { workspace = true, features = [ + "client", + "derive", + "rustls-tls", + "ring", +] } reqwest = { workspace = true, features = ["rustls-tls", "json"] } schemars = "0.8" serde = { workspace = true, features = ["derive"] } serde_json = { workspace = true } tokio = { workspace = true, features = ["full"] } +toml = "1.1" +xshell = "0.2" [lints] workspace = true diff --git a/crates/xtask/src/commands/check.rs b/crates/xtask/src/commands/check.rs new file mode 100644 index 000000000..c5a72de06 --- /dev/null +++ b/crates/xtask/src/commands/check.rs @@ -0,0 +1,28 @@ +use anyhow::Result; +use clap::Args; +use xshell::{Shell, cmd}; + +use super::shared::NIGHTLY_TOOLCHAIN; + +#[derive(Args)] +pub(crate) struct CheckArgs {} + +impl CheckArgs { + pub(crate) fn run(self) -> Result<()> { + let sh = Shell::new()?; + let toolchain = std::env::var("RUSTFMT_NIGHTLY_TOOLCHAIN") + .unwrap_or_else(|_| NIGHTLY_TOOLCHAIN.to_owned()); + let toolchain = format!("+{toolchain}"); + + println!("[fmt]"); + cmd!(sh, "cargo {toolchain} fmt --all -- --check").run()?; + + println!("[sort]"); + cmd!(sh, "cargo sort --workspace --grouped --check").run()?; + + println!("[clippy]"); + cmd!(sh, "cargo clippy --all-targets --all-features --no-deps").run()?; + + Ok(()) + } +} diff --git a/crates/xtask/src/commands/deploy_local.rs b/crates/xtask/src/commands/deploy_local.rs new file mode 100644 index 000000000..116a4af04 --- /dev/null +++ b/crates/xtask/src/commands/deploy_local.rs @@ -0,0 +1,20 @@ +use anyhow::Result; +use clap::Args; +use xshell::{Shell, cmd}; + +#[derive(Args)] +pub(crate) struct DeployLocalArgs { + /// Arguments passed to the deploy script + #[arg(trailing_var_arg = true, allow_hyphen_values = true)] + args: Vec, +} + +impl DeployLocalArgs { + // TODO: port scripts/deploy-local-replicator-orbstack.sh to native Rust + pub(crate) fn run(self) -> Result<()> { + let sh = Shell::new()?; + let args = &self.args; + cmd!(sh, "./scripts/deploy-local-replicator-orbstack.sh {args...}").run()?; + Ok(()) + } +} diff --git a/crates/xtask/src/commands/fix.rs b/crates/xtask/src/commands/fix.rs new file mode 100644 index 000000000..545796636 --- /dev/null +++ b/crates/xtask/src/commands/fix.rs @@ -0,0 +1,28 @@ +use anyhow::Result; +use clap::Args; +use xshell::{Shell, cmd}; + +use super::shared::NIGHTLY_TOOLCHAIN; + +#[derive(Args)] +pub(crate) struct FixArgs {} + +impl FixArgs { + pub(crate) fn run(self) -> Result<()> { + let sh = Shell::new()?; + + println!("[clippy fix]"); + cmd!(sh, "cargo clippy --fix --allow-dirty --allow-staged").run()?; + + println!("[fmt]"); + let toolchain = std::env::var("RUSTFMT_NIGHTLY_TOOLCHAIN") + .unwrap_or_else(|_| NIGHTLY_TOOLCHAIN.to_owned()); + let toolchain = format!("+{toolchain}"); + cmd!(sh, "cargo {toolchain} fmt --all").run()?; + + println!("[sort]"); + cmd!(sh, "cargo sort --workspace --grouped").run()?; + + Ok(()) + } +} diff --git a/crates/xtask/src/commands/fmt.rs b/crates/xtask/src/commands/fmt.rs new file mode 100644 index 000000000..915c6e7fb --- /dev/null +++ b/crates/xtask/src/commands/fmt.rs @@ -0,0 +1,29 @@ +use anyhow::Result; +use clap::Args; +use xshell::{Shell, cmd}; + +use super::shared::NIGHTLY_TOOLCHAIN; + +#[derive(Args)] +pub(crate) struct FmtArgs { + /// Check formatting without modifying files + #[arg(long)] + check: bool, +} + +impl FmtArgs { + pub(crate) fn run(self) -> Result<()> { + let sh = Shell::new()?; + let toolchain = format!( + "+{}", + std::env::var("RUSTFMT_NIGHTLY_TOOLCHAIN") + .unwrap_or_else(|_| NIGHTLY_TOOLCHAIN.to_owned()) + ); + if self.check { + cmd!(sh, "cargo {toolchain} fmt --all -- --check").run()?; + } else { + cmd!(sh, "cargo {toolchain} fmt --all").run()?; + } + Ok(()) + } +} diff --git a/crates/xtask/src/commands/init.rs b/crates/xtask/src/commands/init.rs new file mode 100644 index 000000000..26a4b0b68 --- /dev/null +++ b/crates/xtask/src/commands/init.rs @@ -0,0 +1,15 @@ +use anyhow::Result; +use clap::Args; +use xshell::{Shell, cmd}; + +#[derive(Args)] +pub(crate) struct InitArgs {} + +impl InitArgs { + // TODO: port scripts/init.sh to native Rust + pub(crate) fn run(self) -> Result<()> { + let sh = Shell::new()?; + cmd!(sh, "./scripts/init.sh").run()?; + Ok(()) + } +} diff --git a/crates/xtask/src/commands/migrate.rs b/crates/xtask/src/commands/migrate.rs new file mode 100644 index 000000000..8a8cf3eea --- /dev/null +++ b/crates/xtask/src/commands/migrate.rs @@ -0,0 +1,20 @@ +use anyhow::Result; +use clap::Args; +use xshell::{Shell, cmd}; + +#[derive(Args)] +pub(crate) struct MigrateArgs { + /// Migration targets (etl-api, etl, all) + #[arg(trailing_var_arg = true, allow_hyphen_values = true)] + args: Vec, +} + +impl MigrateArgs { + // TODO: port scripts/run_migrations.sh to native Rust + pub(crate) fn run(self) -> Result<()> { + let sh = Shell::new()?; + let args = &self.args; + cmd!(sh, "./scripts/run_migrations.sh {args...}").run()?; + Ok(()) + } +} diff --git a/crates/xtask/src/commands/mod.rs b/crates/xtask/src/commands/mod.rs index aace7bd33..96f383d1f 100644 --- a/crates/xtask/src/commands/mod.rs +++ b/crates/xtask/src/commands/mod.rs @@ -1,12 +1,30 @@ mod benchmark; mod benchmark_compare; mod chaos; +mod check; +mod deploy_local; +mod fix; +mod fmt; +mod init; +mod migrate; +mod msrv; mod nextest; mod postgres; mod shared; +mod test_clickhouse; +mod vendor_duckdb; pub(crate) use benchmark::BenchmarkArgs; pub(crate) use benchmark_compare::BenchmarkCompareArgs; pub(crate) use chaos::ChaosArgs; +pub(crate) use check::CheckArgs; +pub(crate) use deploy_local::DeployLocalArgs; +pub(crate) use fix::FixArgs; +pub(crate) use fmt::FmtArgs; +pub(crate) use init::InitArgs; +pub(crate) use migrate::MigrateArgs; +pub(crate) use msrv::MsrvArgs; pub(crate) use nextest::NextestArgs; pub(crate) use postgres::PostgresArgs; +pub(crate) use test_clickhouse::TestClickhouseArgs; +pub(crate) use vendor_duckdb::VendorDuckdbArgs; diff --git a/crates/xtask/src/commands/msrv.rs b/crates/xtask/src/commands/msrv.rs new file mode 100644 index 000000000..be1b251c9 --- /dev/null +++ b/crates/xtask/src/commands/msrv.rs @@ -0,0 +1,81 @@ +use anyhow::{Context, Result, bail}; +use clap::Args; +use xshell::{Shell, cmd}; + +#[derive(Args)] +pub(crate) struct MsrvArgs { + /// Also verify that the workspace compiles with the declared MSRV. + #[arg(long)] + verify: bool, +} + +impl MsrvArgs { + pub(crate) fn run(self) -> Result<()> { + let sh = Shell::new()?; + + if cmd!(sh, "which cargo-msrv").quiet().run().is_err() { + bail!("cargo-msrv is not installed. Install it with: cargo install cargo-msrv"); + } + + let cargo_toml = sh.read_file("Cargo.toml").context("failed to read Cargo.toml")?; + let cargo_doc: toml::Table = cargo_toml.parse().context("failed to parse Cargo.toml")?; + let cargo_msrv = cargo_doc + .get("workspace") + .and_then(|w| w.get("package")) + .and_then(|p| p.get("rust-version")) + .and_then(|v| v.as_str()) + .context("failed to read workspace.package.rust-version from Cargo.toml")? + .to_owned(); + + let toolchain_toml = + sh.read_file("rust-toolchain.toml").context("failed to read rust-toolchain.toml")?; + let toolchain_doc: toml::Table = + toolchain_toml.parse().context("failed to parse rust-toolchain.toml")?; + let toolchain_msrv = toolchain_doc + .get("toolchain") + .and_then(|t| t.get("channel")) + .and_then(|v| v.as_str()) + .context("failed to read toolchain.channel from rust-toolchain.toml")? + .to_owned(); + + let resolved_msrv = cmd!( + sh, + "cargo msrv show --manifest-path crates/etl/Cargo.toml --output-format minimal" + ) + .read() + .context("failed to run cargo msrv show")?; + let resolved_msrv = resolved_msrv.trim().to_owned(); + + if cargo_msrv != toolchain_msrv { + bail!( + "workspace rust-version ({cargo_msrv}) does not match rust-toolchain channel \ + ({toolchain_msrv})" + ); + } + + if cargo_msrv != resolved_msrv { + bail!( + "workspace rust-version ({cargo_msrv}) does not match cargo-msrv output \ + ({resolved_msrv})" + ); + } + + println!("verified msrv sync at rust {cargo_msrv}"); + + if self.verify { + println!("verifying workspace compiles with MSRV {cargo_msrv}..."); + cmd!( + sh, + "cargo msrv verify + --manifest-path crates/etl/Cargo.toml + --output-format minimal + -- cargo check --workspace --all-features --locked" + ) + .run() + .context("cargo msrv verify failed")?; + println!("verified workspace compiles with MSRV {cargo_msrv}"); + } + + Ok(()) + } +} diff --git a/crates/xtask/src/commands/shared.rs b/crates/xtask/src/commands/shared.rs index 9b17897a3..3650f8bc8 100644 --- a/crates/xtask/src/commands/shared.rs +++ b/crates/xtask/src/commands/shared.rs @@ -1,2 +1,3 @@ pub(crate) const DEFAULT_PG_SHARD_COUNT: u16 = 3; pub(crate) const DEFAULT_BASE_PORT: u16 = 5430; +pub(crate) const NIGHTLY_TOOLCHAIN: &str = "nightly-2026-04-15"; diff --git a/crates/xtask/src/commands/test_clickhouse.rs b/crates/xtask/src/commands/test_clickhouse.rs new file mode 100644 index 000000000..a434ecefe --- /dev/null +++ b/crates/xtask/src/commands/test_clickhouse.rs @@ -0,0 +1,15 @@ +use anyhow::Result; +use clap::Args; +use xshell::{Shell, cmd}; + +#[derive(Args)] +pub(crate) struct TestClickhouseArgs {} + +impl TestClickhouseArgs { + // TODO: port scripts/test-clickhouse.sh to native Rust + pub(crate) fn run(self) -> Result<()> { + let sh = Shell::new()?; + cmd!(sh, "./scripts/test-clickhouse.sh").run()?; + Ok(()) + } +} diff --git a/crates/xtask/src/commands/vendor_duckdb.rs b/crates/xtask/src/commands/vendor_duckdb.rs new file mode 100644 index 000000000..4c1b4fa24 --- /dev/null +++ b/crates/xtask/src/commands/vendor_duckdb.rs @@ -0,0 +1,20 @@ +use anyhow::Result; +use clap::Args; +use xshell::{Shell, cmd}; + +#[derive(Args)] +pub(crate) struct VendorDuckdbArgs { + /// Arguments passed to the vendor script + #[arg(trailing_var_arg = true, allow_hyphen_values = true)] + args: Vec, +} + +impl VendorDuckdbArgs { + // TODO: port scripts/vendor_duckdb_extensions.sh to native Rust + pub(crate) fn run(self) -> Result<()> { + let sh = Shell::new()?; + let args = &self.args; + cmd!(sh, "./scripts/vendor_duckdb_extensions.sh {args...}").run()?; + Ok(()) + } +} diff --git a/crates/xtask/src/main.rs b/crates/xtask/src/main.rs index 166ce3212..994ab7058 100644 --- a/crates/xtask/src/main.rs +++ b/crates/xtask/src/main.rs @@ -2,7 +2,11 @@ mod commands; use anyhow::Result; use clap::{Parser, Subcommand}; -use commands::{BenchmarkArgs, BenchmarkCompareArgs, ChaosArgs, NextestArgs, PostgresArgs}; +use commands::{ + BenchmarkArgs, BenchmarkCompareArgs, ChaosArgs, CheckArgs, DeployLocalArgs, FixArgs, FmtArgs, + InitArgs, MigrateArgs, MsrvArgs, NextestArgs, PostgresArgs, TestClickhouseArgs, + VendorDuckdbArgs, +}; #[derive(Parser)] #[command(name = "xtask", about = "Project task runner")] @@ -20,10 +24,32 @@ enum Command { BenchmarkCompare(BenchmarkCompareArgs), /// Run chaos testing scenarios against the Kubernetes cluster. Chaos(ChaosArgs), - /// Run tests via nextest, sharded across multiple Postgres clusters + /// Pre-PR gate: fmt, sort, clippy. + Check(CheckArgs), + /// Deploy the replicator to a local OrbStack Kubernetes cluster. + #[command(name = "deploy-local")] + DeployLocal(DeployLocalArgs), + /// Auto-fix: clippy --fix, fmt, sort. + Fix(FixArgs), + /// Format code with nightly rustfmt. + Fmt(FmtArgs), + /// Set up the local development environment. + Init(InitArgs), + /// Run database migrations. + Migrate(MigrateArgs), + /// Verify MSRV consistency across Cargo.toml, rust-toolchain.toml, and + /// cargo-msrv. + Msrv(MsrvArgs), + /// Run tests via nextest, sharded across multiple Postgres clusters. Nextest(NextestArgs), - /// Manage test Postgres clusters + /// Manage test Postgres clusters. Postgres(PostgresArgs), + /// Run ClickHouse integration tests with a local Docker setup. + #[command(name = "test-clickhouse")] + TestClickhouse(TestClickhouseArgs), + /// Download and vendor DuckDB extensions. + #[command(name = "vendor-duckdb")] + VendorDuckdb(VendorDuckdbArgs), } #[tokio::main] @@ -33,7 +59,16 @@ async fn main() -> Result<()> { Command::Benchmark(cmd) => cmd.run(), Command::BenchmarkCompare(cmd) => cmd.run().await, Command::Chaos(cmd) => cmd.run().await, + Command::Check(cmd) => cmd.run(), + Command::DeployLocal(cmd) => cmd.run(), + Command::Fix(cmd) => cmd.run(), + Command::Fmt(cmd) => cmd.run(), + Command::Init(cmd) => cmd.run(), + Command::Migrate(cmd) => cmd.run(), + Command::Msrv(cmd) => cmd.run(), Command::Nextest(cmd) => cmd.run(), Command::Postgres(cmd) => cmd.run(), + Command::TestClickhouse(cmd) => cmd.run(), + Command::VendorDuckdb(cmd) => cmd.run(), } } diff --git a/scripts/README.md b/scripts/README.md new file mode 100644 index 000000000..2bb810146 --- /dev/null +++ b/scripts/README.md @@ -0,0 +1,17 @@ +# Scripts + +Legacy shell scripts for development workflows. All scripts are accessible through the `cargo x` task runner (see `cargo x --help`). + +New development commands should be added as xtask commands in `crates/xtask/src/commands/` rather than as shell scripts here. Existing scripts will be ported to native xtask commands over time. + +## Script to xtask mapping + +| Script | xtask command | +| ------------------------------------- | ------------------------- | +| `fmt` / `fmt-check` | `cargo x fmt [--check]` | +| `check_msrv_sync.sh` | `cargo x msrv` | +| `init.sh` | `cargo x init` | +| `run_migrations.sh` | `cargo x migrate` | +| `deploy-local-replicator-orbstack.sh` | `cargo x deploy-local` | +| `test-clickhouse.sh` | `cargo x test-clickhouse` | +| `vendor_duckdb_extensions.sh` | `cargo x vendor-duckdb` | From 85b22b11041974d98eb592af1036d331761f9eec Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Fri, 22 May 2026 14:58:32 +0200 Subject: [PATCH 28/29] ref(core): Make store and destination more ergonomic (#765) --- AGENTS.md | 6 +++ crates/etl-benchmarks/src/common.rs | 4 +- crates/etl-destinations/src/bigquery/core.rs | 6 +-- .../src/bigquery/test_utils.rs | 4 +- crates/etl-destinations/src/ducklake/core.rs | 6 +-- .../src/ducklake/external_maintenance.rs | 15 +++--- crates/etl-destinations/src/iceberg/core.rs | 6 +-- crates/etl-destinations/src/snowflake/core.rs | 6 +-- crates/etl-replicator/src/core.rs | 52 +++++++++---------- crates/etl-replicator/src/error_reporting.rs | 6 +-- crates/etl/src/destination/capabilities.rs | 15 ++++++ crates/etl/src/destination/mod.rs | 9 ++-- crates/etl/src/lib.rs | 16 ++++-- crates/etl/src/pipeline.rs | 19 +++---- crates/etl/src/replication/apply.rs | 44 ++++++++-------- crates/etl/src/replication/table_sync.rs | 8 +-- crates/etl/src/store/both/memory.rs | 12 ++--- crates/etl/src/store/both/postgres.rs | 9 ++-- crates/etl/src/store/capabilities.rs | 36 +++++++++++++ .../src/store/{cleanup.rs => lifecycle.rs} | 16 ++++-- crates/etl/src/store/mod.rs | 11 ++-- .../etl/src/test_utils/memory_destination.rs | 6 +-- crates/etl/src/test_utils/notifying_store.rs | 4 +- crates/etl/src/test_utils/pipeline.rs | 23 ++++---- .../test_utils/test_destination_wrapper.rs | 4 +- crates/etl/src/workers/apply.rs | 8 +-- crates/etl/src/workers/table_sync.rs | 8 +-- crates/etl/tests/postgres_store.rs | 2 +- docs/explanation/architecture.md | 2 +- docs/explanation/index.md | 2 +- docs/explanation/traits.md | 24 +++++++-- docs/guides/custom-implementations.md | 12 +++-- 32 files changed, 247 insertions(+), 154 deletions(-) create mode 100644 crates/etl/src/destination/capabilities.rs create mode 100644 crates/etl/src/store/capabilities.rs rename crates/etl/src/store/{cleanup.rs => lifecycle.rs} (69%) diff --git a/AGENTS.md b/AGENTS.md index 207007e10..1e1ef6bda 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -105,6 +105,12 @@ ## Documentation - Document all items, public and private, using concise stdlib-style prose. - Link types and methods as [`Type`] and [`Type::method`]. +- In public or module-level docs, prefer fully qualified rustdoc links such as + [`crate::store::PipelineStore`] when referencing items outside the current + module. +- In item docs where the type is already imported and central to the code, + short rustdoc links such as [`StateStore`] are fine. +- Do not add imports solely to make rustdoc links shorter. - Keep comments and docs precise, short, and punctuated. - Normal comments should always end with `.`. - Do not add code examples in rustdoc for this repository. diff --git a/crates/etl-benchmarks/src/common.rs b/crates/etl-benchmarks/src/common.rs index 1f5e224ee..2df06d050 100644 --- a/crates/etl-benchmarks/src/common.rs +++ b/crates/etl-benchmarks/src/common.rs @@ -14,7 +14,7 @@ use anyhow::{Context, Result, bail}; use clap::{Args, ValueEnum}; use etl::{ destination::{ - Destination, + Destination, PipelineDestination, async_result::{DropTableForCopyResult, WriteEventsResult, WriteTableRowsResult}, }, error::EtlResult, @@ -337,7 +337,7 @@ impl CountingDestination { impl Destination for CountingDestination where - D: Destination + Clone + Send + Sync + 'static, + D: PipelineDestination, { fn name() -> &'static str { D::name() diff --git a/crates/etl-destinations/src/bigquery/core.rs b/crates/etl-destinations/src/bigquery/core.rs index 5fa5e1f39..9bd29297f 100644 --- a/crates/etl-destinations/src/bigquery/core.rs +++ b/crates/etl-destinations/src/bigquery/core.rs @@ -15,7 +15,7 @@ use etl::{ error::{ErrorKind, EtlError, EtlResult}, etl_error, state::destination_metadata::{DestinationTableMetadata, DestinationTableSchemaStatus}, - store::{schema::SchemaStore, state::StateStore}, + store::DestinationStore, types::{ Cell, Event, EventSequenceKey, IdentityType, OldTableRow, PipelineId, ReplicatedTableSchema, SchemaDiff, TableId, TableName, TableRow, UpdatedTableRow, @@ -209,7 +209,7 @@ pub struct BigQueryDestination { impl BigQueryDestination where - S: StateStore + SchemaStore + Clone + Send + Sync + 'static, + S: DestinationStore, { /// Creates a new [`BigQueryDestination`] with a pre-configured client. /// @@ -1215,7 +1215,7 @@ fn validate_bigquery_replica_identity( impl Destination for BigQueryDestination where - S: StateStore + SchemaStore + Clone + Send + Sync + 'static, + S: DestinationStore, { fn name() -> &'static str { "bigquery" diff --git a/crates/etl-destinations/src/bigquery/test_utils.rs b/crates/etl-destinations/src/bigquery/test_utils.rs index 5b93baf6f..91b7b807b 100644 --- a/crates/etl-destinations/src/bigquery/test_utils.rs +++ b/crates/etl-destinations/src/bigquery/test_utils.rs @@ -6,7 +6,7 @@ use std::{fmt, path::Path, str::FromStr, time::Duration}; use etl::{ - store::{schema::SchemaStore, state::StateStore}, + store::DestinationStore, types::{PipelineId, TableName}, }; use gcp_bigquery_client::{ @@ -325,7 +325,7 @@ impl BigQueryDatabase { schema_store: S, ) -> BigQueryDestination where - S: StateStore + SchemaStore + Clone + Send + Sync + 'static, + S: DestinationStore, { BigQueryDestination::new_with_key_path( self.project_id.clone(), diff --git a/crates/etl-destinations/src/ducklake/core.rs b/crates/etl-destinations/src/ducklake/core.rs index d7c1df982..ec33705d8 100644 --- a/crates/etl-destinations/src/ducklake/core.rs +++ b/crates/etl-destinations/src/ducklake/core.rs @@ -17,7 +17,7 @@ use etl::{ error::{ErrorKind, EtlResult}, etl_error, state::destination_metadata::DestinationTableMetadata, - store::{schema::SchemaStore, state::StateStore}, + store::DestinationStore, types::{ Event, EventSequenceKey, OldTableRow, PartialTableRow, ReplicatedTableSchema, TableId, TableName, TableRow, UpdatedTableRow, @@ -204,7 +204,7 @@ fn table_write_slot( impl Destination for DuckLakeDestination where - S: StateStore + SchemaStore + Clone + Send + Sync + 'static, + S: DestinationStore, { fn name() -> &'static str { "ducklake" @@ -295,7 +295,7 @@ fn validate_ducklake_replica_identity( impl DuckLakeDestination where - S: StateStore + SchemaStore + Clone + Send + Sync + 'static, + S: DestinationStore, { /// Builds a key-only row from a partial update row when PostgreSQL omits /// the old key image because the replica identity did not change. diff --git a/crates/etl-destinations/src/ducklake/external_maintenance.rs b/crates/etl-destinations/src/ducklake/external_maintenance.rs index 69541edaf..adce0e4aa 100644 --- a/crates/etl-destinations/src/ducklake/external_maintenance.rs +++ b/crates/etl-destinations/src/ducklake/external_maintenance.rs @@ -1,10 +1,7 @@ use std::time::Duration; use chrono::{DateTime, Utc}; -use etl::{ - error::EtlResult, - store::{schema::SchemaStore, state::StateStore}, -}; +use etl::{error::EtlResult, store::DestinationStore}; pub use etl_maintenance::{ ExternalMaintenanceOperationHistory, ExternalMaintenanceOperationPolicy, ExternalMaintenanceOperationRequest, ExternalMaintenanceOperationRun, @@ -56,7 +53,7 @@ pub(super) async fn run_kubernetes_external_maintenance_watcher( destination: DuckLakeDestination, ) -> EtlResult<()> where - S: StateStore + SchemaStore + Clone + Send + Sync + 'static, + S: DestinationStore, { let config = ExternalMaintenanceWatcherConfig::from_env(); let Some(store) = KubernetesExternalMaintenanceStore::from_env(config.store_timeout).await? @@ -74,7 +71,7 @@ pub(super) async fn run_postgres_external_maintenance_watcher( pool: PgPool, ) -> EtlResult<()> where - S: StateStore + SchemaStore + Clone + Send + Sync + 'static, + S: DestinationStore, { let config = ExternalMaintenanceWatcherConfig::from_env(); let store = PostgresExternalMaintenanceStore::new(pipeline_id, pool); @@ -95,7 +92,7 @@ pub async fn run_external_maintenance_watcher( config: ExternalMaintenanceWatcherConfig, ) -> EtlResult<()> where - S: StateStore + SchemaStore + Clone + Send + Sync + 'static, + S: DestinationStore, M: ExternalMaintenanceStore, { let mut held_pause: Option = None; @@ -186,7 +183,7 @@ async fn reconcile_pause( held_pause: &mut Option, state: &ExternalMaintenanceState, ) where - S: StateStore + SchemaStore + Clone + Send + Sync + 'static, + S: DestinationStore, M: ExternalMaintenanceStore, { let active_pause = state.pause_request.clone().filter(|pause| pause.expires_at > Utc::now()); @@ -399,7 +396,7 @@ async fn maybe_request_operations( config: &ExternalMaintenanceWatcherConfig, expire_snapshots_gate: &mut ExpireSnapshotsRequestGate, ) where - S: StateStore + SchemaStore + Clone + Send + Sync + 'static, + S: DestinationStore, M: ExternalMaintenanceStore, { if state.active_run.is_some() { diff --git a/crates/etl-destinations/src/iceberg/core.rs b/crates/etl-destinations/src/iceberg/core.rs index f27c20ea2..4b72dbf81 100644 --- a/crates/etl-destinations/src/iceberg/core.rs +++ b/crates/etl-destinations/src/iceberg/core.rs @@ -13,7 +13,7 @@ use etl::{ error::{ErrorKind, EtlResult}, etl_error, state::destination_metadata::DestinationTableMetadata, - store::state::StateStore, + store::SharedStateStore, types::{ Cell, ColumnSchema, Event, IdentityType, OldTableRow, ReplicatedTableSchema, TableId, TableName, TableRow, Type, generate_sequence_number, @@ -164,7 +164,7 @@ struct Inner { impl IcebergDestination where - S: StateStore + Clone + Send + Sync + 'static, + S: SharedStateStore, { /// Creates a new Iceberg destination instance. /// @@ -616,7 +616,7 @@ where impl Destination for IcebergDestination where - S: StateStore + Clone + Send + Sync + 'static, + S: SharedStateStore, { /// Returns the identifier name for this destination type. fn name() -> &'static str { diff --git a/crates/etl-destinations/src/snowflake/core.rs b/crates/etl-destinations/src/snowflake/core.rs index 089f18b47..086552825 100644 --- a/crates/etl-destinations/src/snowflake/core.rs +++ b/crates/etl-destinations/src/snowflake/core.rs @@ -7,7 +7,7 @@ use etl::{ error::{ErrorKind, EtlError, EtlResult}, etl_error, state::destination_metadata::{DestinationTableMetadata, DestinationTableSchemaStatus}, - store::{schema::SchemaStore, state::StateStore}, + store::DestinationStore, types::{ ColumnSchema, DeleteEvent, Event, InsertEvent, OldTableRow, ReplicatedTableSchema, TableId, TableRow, UpdateEvent, UpdatedTableRow, @@ -48,7 +48,7 @@ impl Clone for Destination impl Destination where - S: StateStore + SchemaStore + Clone + Send + Sync + 'static, + S: DestinationStore, T: TokenProvider + 'static, C: StreamClient, { @@ -417,7 +417,7 @@ where impl etl::destination::Destination for Destination where - S: StateStore + SchemaStore + Clone + Send + Sync + 'static, + S: DestinationStore, T: TokenProvider + 'static, C: StreamClient, { diff --git a/crates/etl-replicator/src/core.rs b/crates/etl-replicator/src/core.rs index 3e7a9dfa2..2ba7142a3 100644 --- a/crates/etl-replicator/src/core.rs +++ b/crates/etl-replicator/src/core.rs @@ -2,12 +2,9 @@ use std::collections::HashMap; use etl::{ config::IcebergConfig, - destination::Destination, + destination::PipelineDestination, pipeline::Pipeline, - store::{ - both::postgres::PostgresStore, cleanup::CleanupStore, schema::SchemaStore, - state::StateStore, - }, + store::{PipelineStore, both::postgres::PostgresStore}, types::PipelineId, }; use etl_config::{ @@ -43,7 +40,7 @@ use crate::{ /// Starts the replicator service with the provided configuration. /// -/// Initializes the state store, creates the appropriate destination based on +/// Initializes the store, creates the appropriate destination based on /// configuration, and starts the pipeline. pub(crate) async fn start_replicator_with_config( replicator_config: ReplicatorConfig, @@ -51,8 +48,9 @@ pub(crate) async fn start_replicator_with_config( ) -> ReplicatorResult<()> { let pipeline_id = replicator_config.pipeline.id; - // We initialize the state store, which for the replicator is not configurable. - let state_store = init_store( + // We initialize the store, which for the replicator is not + // configurable. + let store = init_store( pipeline_id, replicator_config.pipeline.pg_connection.clone(), notification_client, @@ -77,11 +75,11 @@ pub(crate) async fn start_replicator_with_config( *max_staleness_mins, *connection_pool_size, pipeline_id, - state_store.clone(), + store.clone(), ) .await?; - let pipeline = Pipeline::new(replicator_config.pipeline, state_store, destination); + let pipeline = Pipeline::new(replicator_config.pipeline, store, destination); start_pipeline(pipeline).await?; } DestinationConfig::Iceberg { @@ -112,9 +110,9 @@ pub(crate) async fn start_replicator_with_config( Some(ns) => DestinationNamespace::Single(ns.clone()), None => DestinationNamespace::OnePerSchema, }; - let destination = IcebergDestination::new(client, namespace, state_store.clone()); + let destination = IcebergDestination::new(client, namespace, store.clone()); - let pipeline = Pipeline::new(replicator_config.pipeline, state_store, destination); + let pipeline = Pipeline::new(replicator_config.pipeline, store, destination); start_pipeline(pipeline).await?; } DestinationConfig::Iceberg { @@ -143,9 +141,9 @@ pub(crate) async fn start_replicator_with_config( Some(ns) => DestinationNamespace::Single(ns.clone()), None => DestinationNamespace::OnePerSchema, }; - let destination = IcebergDestination::new(client, namespace, state_store.clone()); + let destination = IcebergDestination::new(client, namespace, store.clone()); - let pipeline = Pipeline::new(replicator_config.pipeline, state_store, destination); + let pipeline = Pipeline::new(replicator_config.pipeline, store, destination); start_pipeline(pipeline).await?; } DestinationConfig::Ducklake { @@ -200,11 +198,11 @@ pub(crate) async fn start_replicator_with_config( maintenance_target_file_size.clone(), expire_snapshots_older_than.clone(), external_maintenance, - state_store.clone(), + store.clone(), ) .await?; - let pipeline = Pipeline::new(replicator_config.pipeline, state_store, destination); + let pipeline = Pipeline::new(replicator_config.pipeline, store, destination); start_pipeline(pipeline).await?; } DestinationConfig::ClickHouse { url, user, password, database, engine } => { @@ -215,11 +213,11 @@ pub(crate) async fn start_replicator_with_config( database, ClickHouseInserterConfig { engine: *engine, ..Default::default() }, ClickHouseClientConfig::default(), - state_store.clone(), + store.clone(), )?; destination.validate_engine_support().await?; - let pipeline = Pipeline::new(replicator_config.pipeline, state_store, destination); + let pipeline = Pipeline::new(replicator_config.pipeline, store, destination); start_pipeline(pipeline).await?; } DestinationConfig::Snowflake { @@ -244,9 +242,9 @@ pub(crate) async fn start_replicator_with_config( .map_err(|e| ReplicatorError::config(std::io::Error::other(e.to_string())))?, ); let client = snowflake::Client::new(config, auth, pipeline_id); - let destination = snowflake::Destination::new(client, state_store.clone()); + let destination = snowflake::Destination::new(client, store.clone()); - let pipeline = Pipeline::new(replicator_config.pipeline, state_store, destination); + let pipeline = Pipeline::new(replicator_config.pipeline, store, destination); start_pipeline(pipeline).await?; } } @@ -268,16 +266,16 @@ fn create_props( props } -/// Initializes the state store. +/// Initializes the store. /// /// Creates a [`PostgresStore`] instance for the given pipeline and connection -/// configuration. The pipeline itself owns state-store migration startup. +/// configuration. The pipeline itself owns store migration startup. async fn init_store( pipeline_id: PipelineId, pg_connection_config: PgConnectionConfig, notification_client: Option, -) -> ReplicatorResult { - info!("initializing postgres state store"); +) -> ReplicatorResult { + info!("initializing postgres store"); Ok(ErrorReportingStateStore::new( PostgresStore::new(pipeline_id, pg_connection_config).await?, @@ -293,8 +291,8 @@ async fn init_store( #[tracing::instrument(skip(pipeline))] async fn start_pipeline(mut pipeline: Pipeline) -> ReplicatorResult<()> where - S: StateStore + SchemaStore + CleanupStore + Clone + Send + Sync + 'static, - D: Destination + Clone + Send + Sync + 'static, + S: PipelineStore, + D: PipelineDestination, { // Start the pipeline. pipeline.start().await?; @@ -310,7 +308,7 @@ where // Listen for SIGTERM, sent by Kubernetes before SIGKILL during pod termination. // // If the process is killed before shutdown completes, the pipeline may become - // corrupted, depending on the state store and destination + // corrupted, depending on the store and destination // implementations. let Ok(mut sigterm) = signal(SignalKind::terminate()) else { error!("failed to register sigterm handler, shutting down pipeline"); diff --git a/crates/etl-replicator/src/error_reporting.rs b/crates/etl-replicator/src/error_reporting.rs index 0c0fbfaa9..322a2425a 100644 --- a/crates/etl-replicator/src/error_reporting.rs +++ b/crates/etl-replicator/src/error_reporting.rs @@ -8,7 +8,7 @@ use etl::{ table::TableReplicationPhase, }, store::{ - cleanup::CleanupStore, + lifecycle::TableLifecycleStore, schema::{SchemaStore, TableSchemaRetention}, state::{StateStore, TableReplicationStates}, }, @@ -197,9 +197,9 @@ where } } -impl CleanupStore for ErrorReportingStateStore +impl TableLifecycleStore for ErrorReportingStateStore where - S: CleanupStore + Send + Sync, + S: TableLifecycleStore + Send + Sync, { async fn clear_table_copy_state(&self, table_id: TableId) -> EtlResult<()> { self.inner.clear_table_copy_state(table_id).await diff --git a/crates/etl/src/destination/capabilities.rs b/crates/etl/src/destination/capabilities.rs new file mode 100644 index 000000000..4649f8624 --- /dev/null +++ b/crates/etl/src/destination/capabilities.rs @@ -0,0 +1,15 @@ +//! Aggregate destination capabilities for common runtime roles. +//! +//! These facade traits collect repeated bounds behind names that describe how +//! the destination is used. Code that only dispatches destination writes should +//! depend on [`Destination`] plus only the additional bounds it actually needs. + +use crate::destination::Destination; + +/// Destination capabilities required by the pipeline runtime. +/// +/// This is a facade trait for code that needs a destination to be cloneable, +/// shareable across worker tasks, and owned by spawned futures. +pub trait PipelineDestination: Destination + Clone + Send + Sync + 'static {} + +impl PipelineDestination for D where D: Destination + Clone + Send + Sync + 'static {} diff --git a/crates/etl/src/destination/mod.rs b/crates/etl/src/destination/mod.rs index 6fb34d18d..e16082805 100644 --- a/crates/etl/src/destination/mod.rs +++ b/crates/etl/src/destination/mod.rs @@ -1,11 +1,14 @@ //! Data destination abstractions and implementations. //! -//! Provides the [`Destination`] trait and implementations for sending -//! replicated data to target systems. Destinations handle both initial table -//! synchronization data and streaming replication events. +//! Provides destination traits for sending replicated data to target systems. +//! Destinations handle both initial table synchronization data and streaming +//! replication events. The [`capabilities`] module provides facade traits for +//! common runtime destination roles. pub mod async_result; mod base; +pub mod capabilities; pub use async_result::{DropTableForCopyResult, WriteEventsResult, WriteTableRowsResult}; pub use base::Destination; +pub use capabilities::PipelineDestination; diff --git a/crates/etl/src/lib.rs b/crates/etl/src/lib.rs index 6de4bcfd9..1ccc333b5 100644 --- a/crates/etl/src/lib.rs +++ b/crates/etl/src/lib.rs @@ -36,15 +36,23 @@ //! external systems. //! //! ## Store -//! The [`store::schema::SchemaStore`] and [`store::state::StateStore`] traits -//! define where the table schemas, replication state, and destination table -//! metadata are stored. These stores are critical to a pipeline's operation, as -//! they allow it to be safely paused and resumed. +//! The [`store::PipelineStore`] trait defines the complete store surface needed +//! by the pipeline runtime. It combines narrower capabilities for table +//! schemas, replication state, destination table metadata, and table lifecycle +//! operations. These stores are critical to a pipeline's operation, as they +//! allow it to be safely paused and resumed. //! //! The [`store::state::StateStore`] trait handles table replication states, //! durable replication progress, and destination table metadata, providing a //! single interface for all state-related storage operations. //! +//! The [`store::schema::SchemaStore`] trait handles versioned table schemas, +//! and [`store::TableLifecycleStore`] handles table-scoped reset and removal +//! operations that must update state, schema, and metadata consistently. +//! [`store::SharedStateStore`], [`store::DestinationStore`], and +//! [`store::PipelineStore`] are facade traits for code that needs common +//! combinations of these capabilities. +//! //! **Note:** To pause and resume a pipeline after the process is stopped, it //! must be able to persist data durably. The crate itself provides no //! durability guarantees as it only transfers data between Postgres and the diff --git a/crates/etl/src/pipeline.rs b/crates/etl/src/pipeline.rs index 657689e8b..2b519ad9d 100644 --- a/crates/etl/src/pipeline.rs +++ b/crates/etl/src/pipeline.rs @@ -14,14 +14,14 @@ use crate::{ bail, concurrency::{MemoryMonitor, ShutdownTx, create_shutdown_channel}, config::PipelineConfig, - destination::Destination, + destination::PipelineDestination, error::{ErrorKind, EtlResult}, etl_error, metrics::register_metrics, migrations, replication::{SharedTableCache, client::PgReplicationClient}, state::table::TableReplicationPhase, - store::{cleanup::CleanupStore, schema::SchemaStore, state::StateStore}, + store::PipelineStore, types::{PipelineId, TableId}, workers::{ApplyWorker, ApplyWorkerHandle, TableSyncWorkerPool}, }; @@ -72,18 +72,19 @@ pub struct Pipeline { impl Pipeline where - S: StateStore + SchemaStore + CleanupStore + Clone + Send + Sync + 'static, - D: Destination + Clone + Send + Sync + 'static, + S: PipelineStore, + D: PipelineDestination, { /// Creates a new pipeline with the given configuration. /// /// The pipeline is initially in the not-started state and must be - /// explicitly started using [`Pipeline::start`]. The state store is used - /// for tracking replication progress, table schemas, and destination - /// table metadata, while the destination receives replicated data. + /// explicitly started using [`Pipeline::start`]. The store is + /// used for tracking replication progress, table schemas, destination + /// table metadata, and table lifecycle state, while the destination + /// receives replicated data. /// The pipeline ID is extracted from the configuration, ensuring /// consistency between pipeline identity and configuration settings. - pub fn new(config: PipelineConfig, state_store: S, destination: D) -> Self { + pub fn new(config: PipelineConfig, store: S, destination: D) -> Self { // Register metrics here during pipeline creation to avoid burdening the // users of etl crate to explicitly calling it. Since this method is safe to // call multiple times, it is ok even if there are multiple pipelines created. @@ -99,7 +100,7 @@ where Self { config: Arc::new(config), - store: state_store, + store, destination, state: PipelineState::NotStarted, shutdown_tx, diff --git a/crates/etl/src/replication/apply.rs b/crates/etl/src/replication/apply.rs index fb82b2e24..bcd5ed923 100644 --- a/crates/etl/src/replication/apply.rs +++ b/crates/etl/src/replication/apply.rs @@ -54,7 +54,7 @@ use crate::{ parse_replicated_column_names, }, destination::{ - Destination, + PipelineDestination, async_result::{ ApplyLoopAsyncResultMetadata, CompletedWriteEventsResult, DispatchMetrics, PendingWriteEventsResult, WriteEventsResult, @@ -76,7 +76,7 @@ use crate::{ }, state::table::{TableReplicationPhase, TableReplicationPhaseType}, store::{ - cleanup::CleanupStore, + PipelineStore, SharedStateStore, schema::{SchemaStore, TableSchemaRetention}, state::StateStore, }, @@ -713,8 +713,8 @@ pub(crate) struct ApplyLoop { impl ApplyLoop where - S: StateStore + SchemaStore + CleanupStore + Clone + Send + Sync + 'static, - D: Destination + Clone + Send + Sync + 'static, + S: PipelineStore, + D: PipelineDestination, { /// Starts the apply loop for processing replication events. /// @@ -2435,7 +2435,7 @@ mod apply_worker { remote_final_lsn: PgLsn, ) -> EtlResult where - S: StateStore + Clone + Send + Sync + 'static, + S: SharedStateStore, { fn is_phase_ready_for_changes( phase: TableReplicationPhase, @@ -2475,8 +2475,8 @@ mod apply_worker { current_lsn: PgLsn, ) -> EtlResult> where - S: StateStore + SchemaStore + CleanupStore + Clone + Send + Sync + 'static, - D: Destination + Clone + Send + Sync + 'static, + S: PipelineStore, + D: PipelineDestination, { for (table_id, table_replication_phase) in get_syncing_tables(&ctx.store).await? { let exit_intent = process_single_syncing_table_after_commit( @@ -2595,8 +2595,8 @@ mod apply_worker { current_lsn: PgLsn, ) -> EtlResult> where - S: StateStore + SchemaStore + CleanupStore + Clone + Send + Sync + 'static, - D: Destination + Clone + Send + Sync + 'static, + S: PipelineStore, + D: PipelineDestination, { let worker_state = ctx.pool.get_active_worker_state(table_id).await; @@ -2734,8 +2734,8 @@ mod apply_worker { current_lsn: PgLsn, ) -> EtlResult<()> where - S: StateStore + SchemaStore + CleanupStore + Clone + Send + Sync + 'static, - D: Destination + Clone + Send + Sync + 'static, + S: PipelineStore, + D: PipelineDestination, { for (table_id, table_replication_phase) in get_syncing_tables(&ctx.store).await? { process_single_syncing_table_after_flush( @@ -2760,8 +2760,8 @@ mod apply_worker { current_lsn: PgLsn, ) -> EtlResult<()> where - S: StateStore + SchemaStore + CleanupStore + Clone + Send + Sync + 'static, - D: Destination + Clone + Send + Sync + 'static, + S: PipelineStore, + D: PipelineDestination, { let worker_state = ctx.pool.get_active_worker_state(table_id).await; @@ -2872,8 +2872,8 @@ mod apply_worker { current_lsn: PgLsn, ) -> EtlResult> where - S: StateStore + SchemaStore + CleanupStore + Clone + Send + Sync + 'static, - D: Destination + Clone + Send + Sync + 'static, + S: PipelineStore, + D: PipelineDestination, { for (table_id, table_replication_phase) in get_syncing_tables(&ctx.store).await? { let exit_intent = process_single_syncing_table_when_idle( @@ -2904,8 +2904,8 @@ mod apply_worker { current_lsn: PgLsn, ) -> EtlResult> where - S: StateStore + SchemaStore + CleanupStore + Clone + Send + Sync + 'static, - D: Destination + Clone + Send + Sync + 'static, + S: PipelineStore, + D: PipelineDestination, { let worker_state = ctx.pool.get_active_worker_state(table_id).await; @@ -3104,8 +3104,8 @@ mod apply_worker { worker: TableSyncWorker, ) -> Pin> + Send>> where - S: StateStore + SchemaStore + CleanupStore + Clone + Send + Sync + 'static, - D: Destination + Clone + Send + Sync + 'static, + S: PipelineStore, + D: PipelineDestination, { Box::pin(async move { worker.spawn_into_pool(&pool).await }) } @@ -3173,7 +3173,7 @@ mod table_sync_worker { current_lsn: PgLsn, ) -> EtlResult> where - S: StateStore + Clone + Send + Sync + 'static, + S: SharedStateStore, { try_complete_catchup(ctx, current_lsn).await } @@ -3187,7 +3187,7 @@ mod table_sync_worker { current_lsn: PgLsn, ) -> EtlResult> where - S: StateStore + Clone + Send + Sync + 'static, + S: SharedStateStore, { try_complete_catchup(ctx, current_lsn).await } @@ -3201,7 +3201,7 @@ mod table_sync_worker { current_lsn: PgLsn, ) -> EtlResult> where - S: StateStore + Clone + Send + Sync + 'static, + S: SharedStateStore, { let worker_type = WorkerType::TableSync { table_id: ctx.table_id }; let mut inner = ctx.table_sync_worker_state.lock().await; diff --git a/crates/etl/src/replication/table_sync.rs b/crates/etl/src/replication/table_sync.rs index bb38480e6..5f48d45de 100644 --- a/crates/etl/src/replication/table_sync.rs +++ b/crates/etl/src/replication/table_sync.rs @@ -15,7 +15,7 @@ use crate::{ bail, concurrency::{BatchBudgetController, MemoryMonitor, ShutdownRx}, destination::{ - Destination, + PipelineDestination, async_result::{DropTableForCopyResult, WriteTableRowsResult}, }, error::{ErrorKind, EtlResult}, @@ -23,7 +23,7 @@ use crate::{ metrics::{ETL_TABLE_COPY_DURATION_SECONDS, PARTITIONING_LABEL}, replication::{client::PgReplicationClient, table_cache::SharedTableCache}, state::table::{TableReplicationPhase, TableReplicationPhaseType}, - store::{cleanup::CleanupStore, schema::SchemaStore, state::StateStore}, + store::{PipelineStore, schema::SchemaStore, state::StateStore}, types::PipelineId, workers::{TableCopyResult, TableSyncWorkerState, table_copy}, }; @@ -99,8 +99,8 @@ pub(crate) async fn start_table_sync( batch_budget: BatchBudgetController, ) -> EtlResult where - S: StateStore + SchemaStore + CleanupStore + Clone + Send + 'static, - D: Destination + Clone + Send + 'static, + S: PipelineStore, + D: PipelineDestination, { info!(table_id = table_id.0, "starting initial table sync"); diff --git a/crates/etl/src/store/both/memory.rs b/crates/etl/src/store/both/memory.rs index 532a87883..425d76f44 100644 --- a/crates/etl/src/store/both/memory.rs +++ b/crates/etl/src/store/both/memory.rs @@ -14,7 +14,7 @@ use crate::{ table::TableReplicationPhase, }, store::{ - cleanup::CleanupStore, + lifecycle::TableLifecycleStore, schema::{SchemaStore, TableSchemaRetention, TableSchemaSnapshots}, state::{DestinationTablesMetadata, StateStore, TableReplicationStates}, }, @@ -43,10 +43,10 @@ struct Inner { /// In-memory storage for ETL pipeline state and schema information. /// -/// [`MemoryStore`] implements both [`StateStore`] and [`SchemaStore`] traits, -/// providing a complete storage solution that keeps all data in memory. This is -/// ideal for testing, development, and scenarios where persistence is not -/// required. +/// [`MemoryStore`] implements the store traits required by +/// [`crate::store::PipelineStore`], providing a complete storage solution that +/// keeps all data in memory. This is ideal for testing, development, and +/// scenarios where persistence is not required. /// /// All state information including table replication phases, schema /// definitions, and destination table metadata are stored in memory and will be @@ -297,7 +297,7 @@ impl SchemaStore for MemoryStore { } } -impl CleanupStore for MemoryStore { +impl TableLifecycleStore for MemoryStore { async fn clear_table_copy_state(&self, table_id: TableId) -> EtlResult<()> { self.delete_table_state_for_scope(table_id, TableStateCleanupScope::CopyRestart).await } diff --git a/crates/etl/src/store/both/postgres.rs b/crates/etl/src/store/both/postgres.rs index 359c63a4b..d090c1ef0 100644 --- a/crates/etl/src/store/both/postgres.rs +++ b/crates/etl/src/store/both/postgres.rs @@ -25,7 +25,7 @@ use crate::{ table::TableReplicationPhase, }, store::{ - cleanup::CleanupStore, + lifecycle::TableLifecycleStore, schema::{SchemaStore, TableSchemaRetention, TableSchemaSnapshots}, state::{DestinationTablesMetadata, StateStore, TableReplicationStates}, }, @@ -148,8 +148,9 @@ impl Inner { /// Postgres-backed storage for ETL pipeline state and schema information. /// -/// [`PostgresStore`] implements both [`StateStore`] and [`SchemaStore`] traits, -/// providing persistent storage of replication state and schema information +/// [`PostgresStore`] implements the store traits required by +/// [`crate::store::PipelineStore`], providing persistent storage of replication +/// state, schema information, table lifecycle data, and destination metadata /// directly in the source Postgres database. This ensures durability and /// consistency of the pipeline state across restarts. /// @@ -711,7 +712,7 @@ impl SchemaStore for PostgresStore { } } -impl CleanupStore for PostgresStore { +impl TableLifecycleStore for PostgresStore { async fn clear_table_copy_state(&self, table_id: TableId) -> EtlResult<()> { self.delete_table_state_for_scope(table_id, TableStateCleanupScope::CopyRestart).await } diff --git a/crates/etl/src/store/capabilities.rs b/crates/etl/src/store/capabilities.rs new file mode 100644 index 000000000..0e051e21a --- /dev/null +++ b/crates/etl/src/store/capabilities.rs @@ -0,0 +1,36 @@ +//! Aggregate store capabilities for common runtime roles. +//! +//! These facade traits collect repeated bounds behind names that describe how +//! the store is used. Code that only needs one capability should continue to +//! depend on [`StateStore`], [`SchemaStore`], or [`TableLifecycleStore`] +//! directly. + +use crate::store::{lifecycle::TableLifecycleStore, schema::SchemaStore, state::StateStore}; + +/// Store capabilities required by state-only worker code. +pub trait SharedStateStore: StateStore + Clone + Send + Sync + 'static {} + +impl SharedStateStore for S where S: StateStore + Clone + Send + Sync + 'static {} + +/// Store capabilities commonly required by destination implementations. +/// +/// This is a facade trait for destinations that need to read and update ETL +/// state and schema metadata while remaining cloneable and worker-safe. +pub trait DestinationStore: StateStore + SchemaStore + Clone + Send + Sync + 'static {} + +impl DestinationStore for S where S: StateStore + SchemaStore + Clone + Send + Sync + 'static {} + +/// Store capabilities required by the pipeline runtime. +/// +/// This is a facade trait for code that needs the full runtime store surface: +/// replication state, versioned schemas, table lifecycle operations, and the +/// concurrency bounds required by worker tasks. +pub trait PipelineStore: + StateStore + SchemaStore + TableLifecycleStore + Clone + Send + Sync + 'static +{ +} + +impl PipelineStore for S where + S: StateStore + SchemaStore + TableLifecycleStore + Clone + Send + Sync + 'static +{ +} diff --git a/crates/etl/src/store/cleanup.rs b/crates/etl/src/store/lifecycle.rs similarity index 69% rename from crates/etl/src/store/cleanup.rs rename to crates/etl/src/store/lifecycle.rs index 8140937b0..70a60e56e 100644 --- a/crates/etl/src/store/cleanup.rs +++ b/crates/etl/src/store/lifecycle.rs @@ -1,13 +1,19 @@ +//! Table lifecycle store capability. +//! +//! Lifecycle operations are table-scoped mutations that must keep replication +//! state, versioned schemas, destination metadata, and in-memory caches +//! consistent. + use std::future::Future; use crate::{error::EtlResult, types::TableId}; -/// Combined maintenance operations across state and schema stores. +/// Table lifecycle operations across state and schema stores. /// -/// Provides atomic table-scoped primitives that affect both replication state -/// and schema-related data. Implementations should ensure consistency across -/// in-memory caches and the persistent store. -pub trait CleanupStore { +/// Provides atomic table-scoped primitives that affect both replication state, +/// schema-related data, and destination metadata. Implementations should ensure +/// consistency across in-memory caches and the persistent store. +pub trait TableLifecycleStore { /// Clears stored table-copy state for `table_id`. /// /// Removes destination table metadata, all stored table schemas, and diff --git a/crates/etl/src/store/mod.rs b/crates/etl/src/store/mod.rs index 4cb32cc83..63ee18553 100644 --- a/crates/etl/src/store/mod.rs +++ b/crates/etl/src/store/mod.rs @@ -5,21 +5,24 @@ //! progress, versioned table schemas, destination table metadata, and //! synchronization status. //! -//! Storage is divided into three main categories: +//! Storage is divided into focused capability modules: //! - [`state`] - Replication progress and table synchronization states //! - [`schema`] - Database schema information, versioned schema storage, and //! obsolete schema pruning -//! - [`cleanup`] - Cleanup methods that span both stores +//! - [`capabilities`] - Named facade traits for common store capability sets +//! - [`lifecycle`] - Table lifecycle operations that span both stores //! //! The [`both`] module provides combined implementations that handle both //! state and schema storage in unified systems. pub mod both; -pub mod cleanup; +pub mod capabilities; +pub mod lifecycle; pub mod schema; pub mod state; pub use both::{memory::MemoryStore, postgres::PostgresStore}; -pub use cleanup::CleanupStore; +pub use capabilities::{DestinationStore, PipelineStore, SharedStateStore}; +pub use lifecycle::TableLifecycleStore; pub use schema::SchemaStore; pub use state::{StateStore, TableReplicationStates}; diff --git a/crates/etl/src/test_utils/memory_destination.rs b/crates/etl/src/test_utils/memory_destination.rs index cc7bcd03b..ec5adb79f 100644 --- a/crates/etl/src/test_utils/memory_destination.rs +++ b/crates/etl/src/test_utils/memory_destination.rs @@ -10,7 +10,7 @@ use crate::{ }, error::EtlResult, state::destination_metadata::{DestinationTableMetadata, DestinationTableSchemaStatus}, - store::state::StateStore, + store::SharedStateStore, types::{Event, ReplicatedTableSchema, TableId, TableRow}, }; @@ -38,7 +38,7 @@ pub struct MemoryDestination { impl MemoryDestination where - S: StateStore + Clone + Send + Sync, + S: SharedStateStore, { /// Creates a new memory destination with a state store. /// @@ -127,7 +127,7 @@ where impl Destination for MemoryDestination where - S: StateStore + Clone + Send + Sync, + S: SharedStateStore, { fn name() -> &'static str { "memory" diff --git a/crates/etl/src/test_utils/notifying_store.rs b/crates/etl/src/test_utils/notifying_store.rs index 6efb76cbc..65620de1d 100644 --- a/crates/etl/src/test_utils/notifying_store.rs +++ b/crates/etl/src/test_utils/notifying_store.rs @@ -16,7 +16,7 @@ use crate::{ table::{TableReplicationPhase, TableReplicationPhaseType}, }, store::{ - cleanup::CleanupStore, + lifecycle::TableLifecycleStore, schema::{SchemaStore, TableSchemaRetention, TableSchemaSnapshots}, state::{DestinationTablesMetadata, StateStore, TableReplicationStates}, }, @@ -481,7 +481,7 @@ impl SchemaStore for NotifyingStore { } } -impl CleanupStore for NotifyingStore { +impl TableLifecycleStore for NotifyingStore { async fn clear_table_copy_state(&self, table_id: TableId) -> EtlResult<()> { self.delete_table_state_for_scope(table_id, TableStateCleanupScope::CopyRestart).await } diff --git a/crates/etl/src/test_utils/pipeline.rs b/crates/etl/src/test_utils/pipeline.rs index 45f4bed38..4246fdd48 100644 --- a/crates/etl/src/test_utils/pipeline.rs +++ b/crates/etl/src/test_utils/pipeline.rs @@ -11,10 +11,10 @@ use tokio_postgres::Client; use uuid::Uuid; use crate::{ - destination::Destination, + destination::PipelineDestination, pipeline::Pipeline, state::table::TableReplicationPhaseType, - store::{cleanup::CleanupStore, schema::SchemaStore, state::StateStore}, + store::PipelineStore, test_utils::{ database::{spawn_source_database, test_table_name}, memory_destination::MemoryDestination, @@ -86,8 +86,8 @@ pub struct PipelineBuilder { impl PipelineBuilder where - S: StateStore + SchemaStore + CleanupStore + Clone + Send + Sync + 'static, - D: Destination + Clone + Send + Sync + 'static, + S: PipelineStore, + D: PipelineDestination, { /// Creates a new pipeline builder with required parameters and default /// settings. @@ -98,8 +98,7 @@ where /// * `pipeline_id` - Unique identifier for the pipeline /// * `publication_name` - Name of the PostgreSQL publication to replicate /// from - /// * `store` - Store implementation for state, schema, and cleanup - /// operations + /// * `store` - Store implementation for pipeline runtime operations /// * `destination` - Destination for replicated data /// /// # Default Settings @@ -218,8 +217,8 @@ pub fn create_pipeline( destination: D, ) -> Pipeline where - S: StateStore + SchemaStore + CleanupStore + Clone + Send + Sync + 'static, - D: Destination + Clone + Send + Sync + 'static, + S: PipelineStore, + D: PipelineDestination, { PipelineBuilder::new( pg_connection_config.clone(), @@ -245,8 +244,8 @@ pub fn create_pipeline_with_batch_config( batch: BatchConfig, ) -> Pipeline where - S: StateStore + SchemaStore + CleanupStore + Clone + Send + Sync + 'static, - D: Destination + Clone + Send + Sync + 'static, + S: PipelineStore, + D: PipelineDestination, { PipelineBuilder::new( pg_connection_config.clone(), @@ -273,8 +272,8 @@ pub fn create_pipeline_with_table_sync_copy_config( table_sync_copy: TableSyncCopyConfig, ) -> Pipeline where - S: StateStore + SchemaStore + CleanupStore + Clone + Send + Sync + 'static, - D: Destination + Clone + Send + Sync + 'static, + S: PipelineStore, + D: PipelineDestination, { PipelineBuilder::new( pg_connection_config.clone(), diff --git a/crates/etl/src/test_utils/test_destination_wrapper.rs b/crates/etl/src/test_utils/test_destination_wrapper.rs index 1617814a5..eef331db7 100644 --- a/crates/etl/src/test_utils/test_destination_wrapper.rs +++ b/crates/etl/src/test_utils/test_destination_wrapper.rs @@ -14,7 +14,7 @@ use tokio::{ use crate::{ concurrency::TaskSet, destination::{ - Destination, + Destination, PipelineDestination, async_result::{ ApplyLoopAsyncResultMetadata, DispatchMetrics, DropTableForCopyResult, WriteEventsResult, WriteTableRowsResult, @@ -230,7 +230,7 @@ impl TestDestinationWrapper { impl Destination for TestDestinationWrapper where - D: Destination + Send + Sync + Clone + 'static, + D: PipelineDestination, { fn name() -> &'static str { "wrapper" diff --git a/crates/etl/src/workers/apply.rs b/crates/etl/src/workers/apply.rs index a71478f10..d0b47c934 100644 --- a/crates/etl/src/workers/apply.rs +++ b/crates/etl/src/workers/apply.rs @@ -10,7 +10,7 @@ use tracing::{Instrument, error, info, warn}; use crate::{ bail, concurrency::{BatchBudgetController, MemoryMonitor, ShutdownRx}, - destination::Destination, + destination::PipelineDestination, error::{ErrorKind, EtlError, EtlResult}, etl_error, metrics::{ @@ -22,7 +22,7 @@ use crate::{ client::{GetOrCreateSlotResult, PgReplicationClient, SlotState}, }, state::table::{TableReplicationPhase, TableReplicationPhaseType}, - store::{cleanup::CleanupStore, schema::SchemaStore, state::StateStore}, + store::{PipelineStore, state::StateStore}, types::PipelineId, workers::{ TableSyncWorkerPool, @@ -130,8 +130,8 @@ impl ApplyWorker { impl ApplyWorker where - S: StateStore + SchemaStore + CleanupStore + Clone + Send + Sync + 'static, - D: Destination + Clone + Send + Sync + 'static, + S: PipelineStore, + D: PipelineDestination, { /// Handles apply worker errors using policy-based retry and backoff. /// diff --git a/crates/etl/src/workers/table_sync.rs b/crates/etl/src/workers/table_sync.rs index 4a0188908..7e17d4d9c 100644 --- a/crates/etl/src/workers/table_sync.rs +++ b/crates/etl/src/workers/table_sync.rs @@ -14,7 +14,7 @@ use tracing::{Instrument, debug, error, info, warn}; use crate::{ bail, concurrency::{BatchBudgetController, MemoryMonitor, ShutdownResult, ShutdownRx}, - destination::Destination, + destination::PipelineDestination, error::{ErrorKind, EtlError, EtlResult}, etl_error, metrics::{ERROR_TYPE_LABEL, ETL_WORKER_ERRORS_TOTAL, WORKER_TYPE_LABEL}, @@ -25,7 +25,7 @@ use crate::{ state::table::{ RetryPolicy, TableReplicationError, TableReplicationPhase, TableReplicationPhaseType, }, - store::{cleanup::CleanupStore, schema::SchemaStore, state::StateStore}, + store::{PipelineStore, state::StateStore}, types::PipelineId, workers::{ TableSyncWorkerPool, @@ -380,8 +380,8 @@ impl TableSyncWorker { impl TableSyncWorker where - S: StateStore + SchemaStore + CleanupStore + Clone + Send + Sync + 'static, - D: Destination + Clone + Send + Sync + 'static, + S: PipelineStore, + D: PipelineDestination, { /// Handles a table sync worker failure using the configured retry policy. /// diff --git a/crates/etl/tests/postgres_store.rs b/crates/etl/tests/postgres_store.rs index 4bf38bb2f..80a2cec56 100644 --- a/crates/etl/tests/postgres_store.rs +++ b/crates/etl/tests/postgres_store.rs @@ -11,8 +11,8 @@ use etl::{ table::{RetryPolicy, TableReplicationPhase}, }, store::{ + TableLifecycleStore, both::postgres::PostgresStore, - cleanup::CleanupStore, schema::{SchemaStore, TableSchemaRetention}, state::StateStore, }, diff --git a/docs/explanation/architecture.md b/docs/explanation/architecture.md index 7e387b9d7..8d2898f69 100644 --- a/docs/explanation/architecture.md +++ b/docs/explanation/architecture.md @@ -103,7 +103,7 @@ Persists pipeline state so replication can resume after restarts. Three traits w - **StateStore**: Tracks replication phase per table and destination table metadata - **SchemaStore**: Stores versioned table schema information (columns, types, primary keys, snapshot IDs) and prunes obsolete schema versions after acknowledged progress -- **CleanupStore**: Clears copy-scoped state before a table copy restart and removes all stored state when a table leaves the publication +- **TableLifecycleStore**: Clears copy-scoped state before a table copy restart and removes all stored state when a table leaves the publication `StateStore` and `SchemaStore` use a cache-first pattern: reads hit an in-memory cache, writes go to both the cache and persistent storage. Schema pruning follows the same rule for implementations with durable storage: obsolete versions are removed from both the cache and the persistent store. diff --git a/docs/explanation/index.md b/docs/explanation/index.md index 37ec7d3f2..6518d8665 100644 --- a/docs/explanation/index.md +++ b/docs/explanation/index.md @@ -18,7 +18,7 @@ 4. **[Schema Changes](schema-changes.md)**: How ETL handles DDL, `Relation` events, add/drop/rename semantics, and current limitations. -5. **[Extension Points](traits.md)**: The traits you implement - Destination, SchemaStore, StateStore, CleanupStore. +5. **[Extension Points](traits.md)**: The traits you implement - Destination, SchemaStore, StateStore, TableLifecycleStore. ## Next Steps diff --git a/docs/explanation/traits.md b/docs/explanation/traits.md index d5e1fff35..dbc16fbc9 100644 --- a/docs/explanation/traits.md +++ b/docs/explanation/traits.md @@ -2,7 +2,7 @@ **Traits you implement to customize ETL behavior** -ETL provides four traits for customization. Implement these to control where data goes and how state is stored. +ETL provides extension traits for customization. Implement these to control where data goes and how state is stored. ## Destination @@ -38,6 +38,11 @@ pub trait Destination { See [Event Types](events.md) for details on the events received by `write_events()`. +`PipelineDestination` is a blanket-implemented facade for destinations that +also satisfy the pipeline runtime clone and thread-safety bounds. Pipeline +runtime code uses this facade when it needs to move destinations across worker +tasks, but custom destinations only implement `Destination` directly. + ## SchemaStore Stores versioned table schema information (column names, types, primary keys, @@ -139,12 +144,12 @@ Tables progress through these phases: | `Ready` | Yes | Streaming changes via apply worker | | `Errored { reason, solution, retry_policy }` | Yes | Error occurred, excluded until rollback | -## CleanupStore +## TableLifecycleStore Removes ETL metadata for table-copy restarts and publication removals. ```rust -pub trait CleanupStore { +pub trait TableLifecycleStore { fn clear_table_copy_state(&self, table_id: TableId) -> impl Future> + Send; fn delete_table_pipeline_state(&self, table_id: TableId) -> impl Future> + Send; } @@ -164,9 +169,20 @@ pub struct MyStore { /* ... */ } impl SchemaStore for MyStore { /* ... */ } impl StateStore for MyStore { /* ... */ } -impl CleanupStore for MyStore { /* ... */ } +impl TableLifecycleStore for MyStore { /* ... */ } ``` +`PipelineStore` is a blanket-implemented facade for stores that satisfy the +full pipeline runtime store bounds. Pipeline runtime code uses this facade, +while code that only needs one capability should depend on the narrower trait +directly. + +`DestinationStore` is a blanket-implemented facade for stores that satisfy the +destination runtime store bounds. Destination implementations use this when +they need schema and state metadata but do not need lifecycle reset/removal +operations. `SharedStateStore` covers state-only users with the corresponding +worker-safe bounds. + ETL provides two built-in implementations: - `MemoryStore`: In-memory storage, not persistent across restarts diff --git a/docs/guides/custom-implementations.md b/docs/guides/custom-implementations.md index 6782f8269..0772ff87c 100644 --- a/docs/guides/custom-implementations.md +++ b/docs/guides/custom-implementations.md @@ -57,7 +57,11 @@ Create `src/custom_store.rs`. A store must implement three traits (see [Extensio - `SchemaStore` - Versioned table schema storage, retrieval, and pruning - `StateStore` - Replication progress and destination table metadata tracking -- `CleanupStore` - Store cleanup for table-copy restarts and publication changes +- `TableLifecycleStore` - Store lifecycle operations for table-copy restarts and publication changes + +`SharedStateStore`, `DestinationStore`, and `PipelineStore` are +blanket-implemented facades over these traits plus the required +clone/thread-safety bounds, so custom stores do not implement those directly. ```rust use std::collections::{BTreeMap, HashMap}; @@ -70,7 +74,7 @@ use etl::replication::WorkerType; use etl::state::{ AppliedDestinationTableMetadata, DestinationTableMetadata, TableReplicationPhase, }; -use etl::store::{CleanupStore, SchemaStore, StateStore, TableReplicationStates}; +use etl::store::{SchemaStore, StateStore, TableLifecycleStore, TableReplicationStates}; use etl::store::schema::TableSchemaRetention; use etl::types::{PgLsn, SnapshotId, TableId, TableSchema}; @@ -276,7 +280,7 @@ impl StateStore for CustomStore { } } -impl CleanupStore for CustomStore { +impl TableLifecycleStore for CustomStore { async fn clear_table_copy_state(&self, table_id: TableId) -> EtlResult<()> { let mut tables = self.tables.lock().await; if let Some(entry) = tables.get_mut(&table_id) { @@ -521,7 +525,7 @@ The pipeline will connect to Postgres and start replicating. You'll see your cus ## What You Built -- **Custom Store** - In-memory implementation of `SchemaStore`, `StateStore`, and `CleanupStore` +- **Custom Store** - In-memory implementation of `SchemaStore`, `StateStore`, and `TableLifecycleStore` - **HTTP Destination** - Forwards replicated data via HTTP POST with retry logic - **Working Pipeline** - Connects your custom components to the ETL core From b92ebc1f2a1652fb7b695947d3ec27d0e11aaf15 Mon Sep 17 00:00:00 2001 From: Victor Farazdagi Date: Fri, 22 May 2026 16:54:45 +0300 Subject: [PATCH 29/29] refactor: migrate to smaller example for Snowflake (#767) --- .env.example | 76 ++- Cargo.lock | 215 +------ crates/etl-examples/Cargo.toml | 20 +- crates/etl-examples/README.md | 124 +++- crates/etl-examples/src/bin/snowflake.rs | 273 +++++++++ .../etl-examples/src/bin/snowflake/README.md | 196 ------- .../src/bin/snowflake/commands.rs | 109 ---- .../etl-examples/src/bin/snowflake/logging.rs | 86 --- crates/etl-examples/src/bin/snowflake/main.rs | 461 --------------- .../etl-examples/src/bin/snowflake/state.rs | 451 --------------- crates/etl-examples/src/bin/snowflake/tui.rs | 532 ------------------ .../etl-examples/src/bin/snowflake_loadgen.rs | 457 --------------- crates/xtask/src/commands/example.rs | 139 +++++ crates/xtask/src/commands/mod.rs | 4 + crates/xtask/src/commands/seed.rs | 185 ++++++ crates/xtask/src/main.rs | 13 +- 16 files changed, 787 insertions(+), 2554 deletions(-) create mode 100644 crates/etl-examples/src/bin/snowflake.rs delete mode 100644 crates/etl-examples/src/bin/snowflake/README.md delete mode 100644 crates/etl-examples/src/bin/snowflake/commands.rs delete mode 100644 crates/etl-examples/src/bin/snowflake/logging.rs delete mode 100644 crates/etl-examples/src/bin/snowflake/main.rs delete mode 100644 crates/etl-examples/src/bin/snowflake/state.rs delete mode 100644 crates/etl-examples/src/bin/snowflake/tui.rs delete mode 100644 crates/etl-examples/src/bin/snowflake_loadgen.rs create mode 100644 crates/xtask/src/commands/example.rs create mode 100644 crates/xtask/src/commands/seed.rs diff --git a/.env.example b/.env.example index 5b795fac7..98d2e40a7 100644 --- a/.env.example +++ b/.env.example @@ -8,22 +8,74 @@ export TESTS_DATABASE_PORT=5430 export TESTS_DATABASE_USERNAME=postgres export TESTS_DATABASE_PASSWORD=postgres -# Snowflake (required for Snowflake integration tests) -# See etl-destinations/src/snowflake/README.md for key-pair setup instructions. +# BigQuery (required for ignored BigQuery integration tests) +# +# Run them with: cargo test -p etl-destinations --features bigquery -- --ignored +# +# PROJECT_ID: Your GCP project ID. +# Find it in the GCP Console dashboard, or run: gcloud config get-value project +export TESTS_BIGQUERY_PROJECT_ID= +# +# SA_KEY_PATH: Path to a GCP service account JSON key file. +# The service account needs BigQuery Data Editor and BigQuery Job User roles. +# Create one in GCP Console > IAM > Service Accounts > Keys, or: +# gcloud iam service-accounts keys create sa-key.json \ +# --iam-account=etl-test@.iam.gserviceaccount.com +export TESTS_BIGQUERY_SA_KEY_PATH= + +# ClickHouse (required for ClickHouse integration tests) +# +# URL: ClickHouse HTTP endpoint. For local dev, start with: +# docker run -d --name clickhouse -p 8123:8123 -p 9000:9000 clickhouse/clickhouse-server +export TESTS_CLICKHOUSE_URL=http://localhost:8123 +# +# USER / PASSWORD: ClickHouse credentials. The default server has no auth, +# but the test harness creates a dedicated user. Typical local values: +export TESTS_CLICKHOUSE_USER=default +export TESTS_CLICKHOUSE_PASSWORD= + +# Snowflake (required for ignored Snowflake integration tests) +# +# These tests are #[ignore]'d by default and require real Snowflake credentials. +# Run them with: cargo test -p etl-destinations --features snowflake -- --ignored +# +# ACCOUNT: Your Snowflake account identifier in ORG-ACCOUNT format. +# Find it in the Snowflake UI (bottom-left account selector), or run: +# SELECT CURRENT_ORGANIZATION_NAME() || '-' || CURRENT_ACCOUNT_NAME(); export TESTS_SNOWFLAKE_ACCOUNT= +# +# USER: A Snowflake user configured for key-pair authentication. +# Create one: CREATE USER etl_test_user TYPE = SERVICE; +# Then assign an RSA public key (see PRIVATE_KEY_PATH below). export TESTS_SNOWFLAKE_USER= +# +# PRIVATE_KEY_PATH: Path to the RSA private key file (.p8) for key-pair auth. +# Generate a key pair: +# openssl genrsa 2048 | openssl pkcs8 -topk8 -nocrypt -out rsa_key.p8 +# openssl rsa -in rsa_key.p8 -pubout -out rsa_key.pub +# Assign the public key to the user: +# ALTER USER etl_test_user SET RSA_PUBLIC_KEY=''; export TESTS_SNOWFLAKE_PRIVATE_KEY_PATH= +# +# DATABASE: Target database for tests. Create if needed: +# CREATE DATABASE ETL_DEV; export TESTS_SNOWFLAKE_DATABASE=ETL_DEV +# +# SCHEMA: Target schema. Tests create/drop tables here. export TESTS_SNOWFLAKE_SCHEMA=PUBLIC +# +# WAREHOUSE: Compute warehouse for queries and DDL. +# SHOW WAREHOUSES; or CREATE WAREHOUSE ETL_WH WAREHOUSE_SIZE = 'XSMALL'; export TESTS_SNOWFLAKE_WAREHOUSE= +# +# ROLE: Must have USAGE on warehouse/database/schema and CREATE TABLE on the schema. +# Example setup: +# CREATE ROLE IF NOT EXISTS etl_test_role; +# GRANT USAGE ON WAREHOUSE COMPUTE_WH TO ROLE etl_test_role; +# GRANT USAGE ON DATABASE ETL_DEV TO ROLE etl_test_role; +# GRANT USAGE ON SCHEMA ETL_DEV.PUBLIC TO ROLE etl_test_role; +# GRANT CREATE TABLE ON SCHEMA ETL_DEV.PUBLIC TO ROLE etl_test_role; +# GRANT CREATE STAGE ON SCHEMA ETL_DEV.PUBLIC TO ROLE etl_test_role; +# GRANT CREATE PIPE ON SCHEMA ETL_DEV.PUBLIC TO ROLE etl_test_role; +# GRANT ROLE etl_test_role TO USER ; export TESTS_SNOWFLAKE_ROLE= - -# Snowflake benchmark (used by the snowflake example binary) -# Can reuse the same account/key as tests, but targets a separate database. -export BENCH_SNOWFLAKE_ACCOUNT= -export BENCH_SNOWFLAKE_USER= -export BENCH_SNOWFLAKE_PRIVATE_KEY_PATH= -export BENCH_SNOWFLAKE_DATABASE=ETL_BENCH -export BENCH_SNOWFLAKE_SCHEMA=CDC -export BENCH_SNOWFLAKE_WAREHOUSE= -export BENCH_SNOWFLAKE_ROLE= diff --git a/Cargo.lock b/Cargo.lock index 2d1731860..e0e31cf24 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1133,15 +1133,6 @@ version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" -[[package]] -name = "castaway" -version = "0.2.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dec551ab6e7578819132c713a93c022a05d60159dc86e7a7050223577484c55a" -dependencies = [ - "rustversion", -] - [[package]] name = "cbc" version = "0.2.0" @@ -1323,25 +1314,11 @@ version = "7.1.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4a65ebfec4fb190b6f90e944a817d60499ee0744e582530e2c9900a22e591d9a" dependencies = [ - "crossterm 0.28.1", + "crossterm", "unicode-segmentation", "unicode-width", ] -[[package]] -name = "compact_str" -version = "0.9.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3fdb1325a1cece981e8a296ab8f0f9b63ae357bd0784a9faaf548cc7b480707a" -dependencies = [ - "castaway", - "cfg-if", - "itoa", - "rustversion", - "ryu", - "static_assertions", -] - [[package]] name = "concurrent-queue" version = "2.5.0" @@ -1562,24 +1539,6 @@ dependencies = [ "winapi", ] -[[package]] -name = "crossterm" -version = "0.29.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d8b9f2e4c67f833b660cdb0a3523065869fb35570177239812ed4c905aeff87b" -dependencies = [ - "bitflags", - "crossterm_winapi", - "derive_more", - "document-features", - "mio", - "parking_lot", - "rustix 1.1.4", - "signal-hook", - "signal-hook-mio", - "winapi", -] - [[package]] name = "crossterm_winapi" version = "0.9.1" @@ -1871,15 +1830,6 @@ dependencies = [ "const-random", ] -[[package]] -name = "document-features" -version = "0.2.12" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d4b8a88685455ed29a21542a33abd9cb6510b6b129abadabdcef0f4c55bc8f61" -dependencies = [ - "litrs", -] - [[package]] name = "dotenvy" version = "0.15.7" @@ -2163,19 +2113,14 @@ name = "etl-examples" version = "0.1.0" dependencies = [ "clap", - "crossterm 0.29.0", "etl", "etl-config", "etl-destinations", "etl-telemetry", "k8s-openapi", - "rand 0.9.4", - "ratatui", - "reqwest", "rustls", "secrecy", "tokio", - "tokio-postgres", "tracing", "tracing-subscriber", "url", @@ -3294,15 +3239,6 @@ dependencies = [ "serde_core", ] -[[package]] -name = "indoc" -version = "2.0.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "79cf5c93f93228cf8efb3ba362535fb11199ac548a09ce117c9b1adc3030d706" -dependencies = [ - "rustversion", -] - [[package]] name = "inout" version = "0.2.2" @@ -3327,19 +3263,6 @@ dependencies = [ "tempfile", ] -[[package]] -name = "instability" -version = "0.3.12" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5eb2d60ef19920a3a9193c3e371f726ec1dafc045dac788d0fb3704272458971" -dependencies = [ - "darling 0.23.0", - "indoc", - "proc-macro2", - "quote", - "syn 2.0.117", -] - [[package]] name = "integer-encoding" version = "3.0.4" @@ -3508,17 +3431,6 @@ dependencies = [ "serde_json", ] -[[package]] -name = "kasuari" -version = "0.4.12" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bde5057d6143cc94e861d90f591b9303d6716c6b9602309150bd068853c10899" -dependencies = [ - "hashbrown 0.16.1", - "portable-atomic", - "thiserror 2.0.18", -] - [[package]] name = "kube" version = "1.1.0" @@ -3758,15 +3670,6 @@ dependencies = [ "vcpkg", ] -[[package]] -name = "line-clipping" -version = "0.3.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3f50e8f47623268b5407192d26876c4d7f89d686ca130fdc53bced4814cd29f8" -dependencies = [ - "bitflags", -] - [[package]] name = "linux-raw-sys" version = "0.4.15" @@ -3785,12 +3688,6 @@ version = "0.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "92daf443525c4cce67b150400bc2316076100ce0b3686209eb8cf3c31612e6f0" -[[package]] -name = "litrs" -version = "1.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "11d3d7f243d5c5a8b9bb5d6dd2b1602c0cb0b9db1621bafc7ed66e35ff9fe092" - [[package]] name = "local-waker" version = "0.1.4" @@ -3812,15 +3709,6 @@ version = "0.4.29" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" -[[package]] -name = "lru" -version = "0.16.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7f66e8d5d03f609abc3a39e6f08e4164ebf1447a732906d39eb9b99b7919ef39" -dependencies = [ - "hashbrown 0.16.1", -] - [[package]] name = "lru-slab" version = "0.1.2" @@ -5214,69 +5102,6 @@ dependencies = [ "rustversion", ] -[[package]] -name = "ratatui" -version = "0.30.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d1ce67fb8ba4446454d1c8dbaeda0557ff5e94d39d5e5ed7f10a65eb4c8266bc" -dependencies = [ - "instability", - "ratatui-core", - "ratatui-crossterm", - "ratatui-widgets", -] - -[[package]] -name = "ratatui-core" -version = "0.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5ef8dea09a92caaf73bff7adb70b76162e5937524058a7e5bff37869cbbec293" -dependencies = [ - "bitflags", - "compact_str", - "hashbrown 0.16.1", - "indoc", - "itertools 0.14.0", - "kasuari", - "lru", - "strum", - "thiserror 2.0.18", - "unicode-segmentation", - "unicode-truncate", - "unicode-width", -] - -[[package]] -name = "ratatui-crossterm" -version = "0.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "577c9b9f652b4c121fb25c6a391dd06406d3b092ba68827e6d2f09550edc54b3" -dependencies = [ - "cfg-if", - "crossterm 0.29.0", - "instability", - "ratatui-core", -] - -[[package]] -name = "ratatui-widgets" -version = "0.3.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d7dbfa023cd4e604c2553483820c5fe8aa9d71a42eea5aa77c6e7f35756612db" -dependencies = [ - "bitflags", - "hashbrown 0.16.1", - "indoc", - "instability", - "itertools 0.14.0", - "line-clipping", - "ratatui-core", - "strum", - "time", - "unicode-segmentation", - "unicode-width", -] - [[package]] name = "raw-cpuid" version = "11.6.0" @@ -6197,27 +6022,6 @@ version = "1.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" -[[package]] -name = "signal-hook" -version = "0.3.18" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d881a16cf4426aa584979d30bd82cb33429027e42122b169753d6ef1085ed6e2" -dependencies = [ - "libc", - "signal-hook-registry", -] - -[[package]] -name = "signal-hook-mio" -version = "0.2.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b75a19a7a740b25bc7944bdee6172368f988763b744e3d4dfe753f6b4ece40cc" -dependencies = [ - "libc", - "mio", - "signal-hook", -] - [[package]] name = "signal-hook-registry" version = "1.4.8" @@ -6537,12 +6341,6 @@ version = "1.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" -[[package]] -name = "static_assertions" -version = "1.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" - [[package]] name = "strfmt" version = "0.2.5" @@ -7334,17 +7132,6 @@ version = "1.13.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9629274872b2bfaf8d66f5f15725007f635594914870f65218920345aa11aa8c" -[[package]] -name = "unicode-truncate" -version = "2.0.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "16b380a1238663e5f8a691f9039c73e1cdae598a30e9855f541d29b08b53e9a5" -dependencies = [ - "itertools 0.14.0", - "unicode-segmentation", - "unicode-width", -] - [[package]] name = "unicode-width" version = "0.2.2" diff --git a/crates/etl-examples/Cargo.toml b/crates/etl-examples/Cargo.toml index 0b6663d14..aa6dcad83 100644 --- a/crates/etl-examples/Cargo.toml +++ b/crates/etl-examples/Cargo.toml @@ -28,13 +28,7 @@ test = false [[bin]] name = "snowflake" -path = "src/bin/snowflake/main.rs" -required-features = ["snowflake"] -test = false - -[[bin]] -name = "snowflake-loadgen" -path = "src/bin/snowflake_loadgen.rs" +path = "src/bin/snowflake.rs" required-features = ["snowflake"] test = false @@ -45,12 +39,7 @@ clickhouse = ["etl-destinations/clickhouse"] ducklake = ["etl-destinations/ducklake"] snowflake = [ "etl-destinations/snowflake", - "dep:crossterm", - "dep:ratatui", - "dep:rand", - "dep:reqwest", "dep:secrecy", - "dep:tokio-postgres", ] [dependencies] @@ -60,21 +49,14 @@ clap = { workspace = true, default-features = true, features = [ "derive", "env", ] } -crossterm = { version = "0.29", optional = true } etl = { workspace = true } etl-config = { workspace = true } etl-destinations = { workspace = true } etl-telemetry = { workspace = true } k8s-openapi = { workspace = true, features = ["latest"] } -rand = { workspace = true, optional = true, features = ["std", "std_rng"] } -ratatui = { version = "0.30", optional = true, default-features = false, features = [ - "crossterm", -] } -reqwest = { workspace = true, optional = true, features = ["json"] } rustls = { workspace = true, features = ["aws-lc-rs", "logging"] } secrecy = { workspace = true, optional = true } tokio = { workspace = true, features = ["macros", "signal", "time"] } -tokio-postgres = { workspace = true, optional = true } tracing = { workspace = true, default-features = true } tracing-subscriber = { workspace = true, default-features = true, features = [ "env-filter", diff --git a/crates/etl-examples/README.md b/crates/etl-examples/README.md index db90c709d..d2bb4ab94 100644 --- a/crates/etl-examples/README.md +++ b/crates/etl-examples/README.md @@ -1,7 +1,6 @@ # `etl` — Examples -This crate contains practical examples demonstrating how to replicate data from -Postgres to various destinations using the ETL pipeline. +This crate contains practical examples demonstrating how to replicate data from Postgres to various destinations using the ETL pipeline. ## Available Examples @@ -10,13 +9,35 @@ Postgres to various destinations using the ETL pipeline. | [BigQuery](#bigquery) | `bigquery` | `bigquery` | Google BigQuery (cloud data warehouse) | Stable | | [ClickHouse](#clickhouse-setup) | `clickhouse` | `clickhouse` | ClickHouse (column-oriented OLAP database) | In progress | | [DuckLake](#ducklake) | `ducklake` | `ducklake` | DuckLake (open data lake format) | In progress | -| [Snowflake](src/bin/snowflake/README.md) | `snowflake` | `snowflake` | Snowflake (cloud data warehouse) | In progress | +| [Snowflake](#snowflake) | `snowflake` | `snowflake` | Snowflake (cloud data warehouse) | In progress | -## Building and running +## Running an example -Each binary is feature-gated so you only compile the dependencies you need. Some destinations (e.g. `ducklake`) pull in heavy native dependencies that can take several minutes to compile, and there's no reason to pay that cost when you only want to try BigQuery. +The quickest way to run an example is via the xtask wrapper. -### Single example +After sourcing your `.env` (see `.env.example`), the DB connection is picked up automatically: + +```bash +source .env +cargo x example snowflake +cargo x example bigquery +cargo x example clickhouse +cargo x example ducklake +``` + +This handles the `-p etl-examples --features ` boilerplate, injects `TESTS_DATABASE_*` env vars as `--db-*` flags, and defaults `--db-name` to `etl_testdata` and `--publication` to `seed_pub` (matching `cargo x seed`). + +Override any flag by passing it explicitly: + +```bash +cargo x example snowflake --db-name mydb --publication my_pub +``` + +### Building manually + +Each binary is feature-gated so you only compile the dependencies you need. Some +destinations (e.g. `ducklake`) pull in heavy native dependencies that can take +several minutes to compile. ```bash # Build @@ -45,20 +66,39 @@ All examples require a Postgres database with **logical replication** enabled: wal_level = logical ``` -Create a publication for the tables you want to replicate: +The Postgres user must have the `REPLICATION` role: ```sql --- Specific tables -CREATE PUBLICATION my_pub FOR TABLE orders, customers; +ALTER USER my_user REPLICATION; +``` --- All tables in the database -CREATE PUBLICATION my_pub FOR ALL TABLES; +### Quick database setup + +The fastest way to get a seeded database with a publication is via the xtask: + +```bash +# Start the dev Postgres (port 5430, wal_level=logical) +cargo x init + +# Create and seed a database with 3 tables (users, orders, events) and a publication +cargo x seed # defaults: etl_testdata, 1000 rows +cargo x seed --rows 100000 # more data +cargo x seed --database mydb --force # custom name, recreate if exists ``` -The Postgres user must have the `REPLICATION` role: +This creates three tables (users, orders, events) in the `public` schema and a +`seed_pub` publication for them. Use `--help` for all options. + +### Manual setup + +If you prefer to use your own tables: ```sql -ALTER USER my_user REPLICATION; +-- Specific tables +CREATE PUBLICATION my_pub FOR TABLE orders, customers; + +-- All tables in the database +CREATE PUBLICATION my_pub FOR ALL TABLES; ``` --- @@ -151,7 +191,7 @@ cargo run -p etl-examples --bin clickhouse --features clickhouse -- \ The destination supports two layouts, chosen per pipeline via `--clickhouse-engine`: | Flag value | Engine | Use it for | -|----------------------------------|----------------------|---------------------------------------------------------| +| -------------------------------- | -------------------- | ------------------------------------------------------- | | `replacing_merge_tree` (default) | `ReplacingMergeTree` | Current-state replicas. Source must have a primary key. | | `merge_tree` | `MergeTree` | Append-only event log. Works for PK-less source tables. | @@ -383,3 +423,59 @@ cargo run --bin bigquery -p etl-examples --features bigquery -- \ | `--max-batch-fill-duration-ms` | `5000` | Max time to wait before flushing a batch | | `--max-table-sync-workers` | `4` | Concurrent workers during initial copy | | `--publication` | _(required)_ | Postgres publication name | + +--- + +## Snowflake + +Replicates a Postgres publication to a Snowflake database via Snowpipe Streaming. + +### Prerequisites + +1. A Snowflake account with a user configured for **key-pair authentication**. +2. An RSA private key file (`.p8`) with the public key registered on the Snowflake user. +3. A target database and schema created in Snowflake. +4. A role with USAGE on warehouse/database/schema and CREATE TABLE, CREATE STAGE, + CREATE PIPE on the schema. + +See `.env.example` for detailed setup instructions and SQL commands for each step. + +### Run + +```bash +cargo run --bin snowflake -p etl-examples --features snowflake -- \ + --db-host localhost \ + --db-port 5430 \ + --db-name etl_testdata \ + --db-username postgres \ + --db-password postgres \ + --snowflake-account ORG-ACCOUNT \ + --snowflake-user myuser \ + --snowflake-private-key-path /path/to/rsa_key.p8 \ + --snowflake-database MY_DATABASE \ + --snowflake-role my_role \ + --publication seed_pub +``` + +Snowflake args can also be set via `SNOWFLAKE_*` environment variables (e.g. +`SNOWFLAKE_ACCOUNT`, `SNOWFLAKE_USER`, etc.) instead of CLI flags. + +### All flags + +| Flag | Default | Description | +| ------------------------------------ | ------------ | ------------------------------------------------- | +| `--db-host` | _(required)_ | Postgres host | +| `--db-port` | _(required)_ | Postgres port | +| `--db-name` | _(required)_ | Postgres database name | +| `--db-username` | _(required)_ | Postgres user | +| `--db-password` | -- | Postgres password | +| `--snowflake-account` | _(required)_ | Snowflake account identifier (ORG-ACCOUNT format) | +| `--snowflake-user` | _(required)_ | Snowflake user for key-pair auth | +| `--snowflake-private-key-path` | _(required)_ | Path to RSA private key file (.p8) | +| `--snowflake-private-key-passphrase` | -- | Passphrase if the key is encrypted | +| `--snowflake-database` | _(required)_ | Snowflake target database | +| `--snowflake-schema` | `PUBLIC` | Snowflake target schema | +| `--snowflake-role` | -- | Snowflake role (uses user default if omitted) | +| `--max-batch-fill-duration-ms` | `5000` | Max time to wait before flushing a batch | +| `--max-table-sync-workers` | `4` | Concurrent workers during initial copy | +| `--publication` | _(required)_ | Postgres publication name | diff --git a/crates/etl-examples/src/bin/snowflake.rs b/crates/etl-examples/src/bin/snowflake.rs new file mode 100644 index 000000000..00e1e5b94 --- /dev/null +++ b/crates/etl-examples/src/bin/snowflake.rs @@ -0,0 +1,273 @@ +/// Snowflake Example +/// +/// Streams Postgres CDC to Snowflake using Snowpipe Streaming. +/// +/// Prerequisites: +/// 1. Postgres with logical replication enabled (wal_level = logical) +/// 2. A publication (e.g. `cargo x seed` to create test tables + publication) +/// 3. Snowflake account with key-pair authentication configured +/// 4. RSA private key file with the public key registered on the Snowflake user +/// +/// Usage: +/// cargo run --bin snowflake -p etl-examples --features snowflake -- \ +/// --db-host localhost \ +/// --db-port 5430 \ +/// --db-name etl_testdata \ +/// --db-username postgres \ +/// --db-password postgres \ +/// --snowflake-account myorg-myaccount \ +/// --snowflake-user myuser \ +/// --snowflake-private-key-path /path/to/rsa_key.p8 \ +/// --snowflake-database MY_DATABASE \ +/// --publication seed_pub +use std::{error::Error, sync::Arc, sync::Once}; + +use clap::{Args, Parser}; +use etl::{ + config::{ + BatchConfig, InvalidatedSlotBehavior, MemoryBackpressureConfig, PgConnectionConfig, + PipelineConfig, TableSyncCopyConfig, TcpKeepaliveConfig, TlsConfig, + }, + pipeline::Pipeline, + store::PostgresStore, +}; +use etl_destinations::snowflake::{AuthManager, Client, Config, Destination}; +use secrecy::SecretString; +use tokio::signal; +use tracing::{error, info}; +use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt}; + +/// Ensures crypto provider is only initialized once. +static INIT_CRYPTO: Once = Once::new(); + +/// Installs the default cryptographic provider for rustls. +fn install_crypto_provider() { + INIT_CRYPTO.call_once(|| { + rustls::crypto::aws_lc_rs::default_provider() + .install_default() + .expect("failed to install default crypto provider"); + }); +} + +// Main application arguments combining database and Snowflake configurations +#[derive(Debug, Parser)] +#[command(name = "snowflake", version, about, arg_required_else_help = true)] +struct AppArgs { + // Postgres connection parameters + #[clap(flatten)] + db_args: DbArgs, + // Snowflake destination parameters + #[clap(flatten)] + sf_args: SnowflakeArgs, + /// Postgres publication name (must be created beforehand with CREATE + /// PUBLICATION) + #[arg(long)] + publication: String, +} + +// Postgres database connection configuration +#[derive(Debug, Args)] +struct DbArgs { + /// Host on which Postgres is running (e.g., localhost or IP address) + #[arg(long)] + db_host: String, + /// Port on which Postgres is running (default: 5432) + #[arg(long)] + db_port: u16, + /// Postgres database name to connect to + #[arg(long)] + db_name: String, + /// Postgres database user name (must have REPLICATION privileges) + #[arg(long)] + db_username: String, + /// Postgres database user password (optional if using trust authentication) + #[arg(long)] + db_password: Option, +} + +// Snowflake destination configuration +#[derive(Debug, Args)] +struct SnowflakeArgs { + /// Snowflake account identifier (e.g., myorg-myaccount) + #[arg(long, env = "TESTS_SNOWFLAKE_ACCOUNT")] + snowflake_account: String, + /// Snowflake user name configured for key-pair authentication + #[arg(long, env = "TESTS_SNOWFLAKE_USER")] + snowflake_user: String, + /// Path to the RSA private key file (.p8) for key-pair authentication + #[arg(long, env = "TESTS_SNOWFLAKE_PRIVATE_KEY_PATH")] + snowflake_private_key_path: String, + /// Passphrase for the private key file (if the key is encrypted) + #[arg(long, env = "TESTS_SNOWFLAKE_PRIVATE_KEY_PASSPHRASE")] + snowflake_private_key_passphrase: Option, + /// Snowflake database name where tables will be created and data loaded + #[arg(long, env = "TESTS_SNOWFLAKE_DATABASE")] + snowflake_database: String, + /// Snowflake schema name within the database (default: PUBLIC) + #[arg(long, env = "TESTS_SNOWFLAKE_SCHEMA", default_value = "PUBLIC")] + snowflake_schema: String, + /// Snowflake role to use for the session (optional, uses user's default + /// role if not specified) + #[arg(long, env = "TESTS_SNOWFLAKE_ROLE")] + snowflake_role: Option, + /// Maximum time to wait for a batch to fill in milliseconds (lower values = + /// lower latency, less throughput) + #[arg(long, default_value = "5000")] + max_batch_fill_duration_ms: u64, + /// Maximum number of concurrent table sync workers (higher values = faster + /// initial sync, more resource usage) + #[arg(long, default_value = "4")] + max_table_sync_workers: u16, +} + +// Entry point - handles error reporting and process exit +#[tokio::main] +async fn main() -> Result<(), Box> { + if let Err(err) = main_impl().await { + error!(error = %err, "fatal error"); + std::process::exit(1); + } + + Ok(()) +} + +// Initialize structured logging with configurable log levels via RUST_LOG +// environment variable +fn init_tracing() { + tracing_subscriber::registry() + .with( + tracing_subscriber::EnvFilter::try_from_default_env() + .unwrap_or_else(|_| "snowflake=info".into()), + ) + .with(tracing_subscriber::fmt::layer()) + .init(); +} + +// Set default log level if RUST_LOG environment variable is not set +fn set_log_level() { + if std::env::var("RUST_LOG").is_err() { + unsafe { + std::env::set_var("RUST_LOG", "info"); + } + } +} + +// Main implementation function containing all the pipeline setup and execution +// logic +async fn main_impl() -> Result<(), Box> { + // Set up logging and tracing + set_log_level(); + init_tracing(); + + // Install required crypto provider for authentication + install_crypto_provider(); + + // Parse command line arguments + let args = AppArgs::parse(); + + // Configure Postgres connection settings + // Note: TLS is disabled in this example - enable for production use + let pg_connection_config = PgConnectionConfig { + host: args.db_args.db_host, + hostaddr: None, + port: args.db_args.db_port, + name: args.db_args.db_name, + username: args.db_args.db_username, + password: args.db_args.db_password.map(Into::into), + tls: TlsConfig { + trusted_root_certs: String::new(), + enabled: false, // Set to true and provide certs for production + }, + keepalive: TcpKeepaliveConfig::default(), + }; + + // Create a persistent store for tracking table replication states and + // schemas. This runs the Postgres store migrations; Pipeline::start() + // runs the source migrations required by replication. + let pipeline_id = 1; + let store = PostgresStore::new(pipeline_id, pg_connection_config.clone()).await?; + + // Create pipeline configuration with batching and retry settings + let pipeline_config = PipelineConfig { + id: pipeline_id, // Using a simple ID for the example + publication_name: args.publication, + pg_connection: pg_connection_config, + batch: BatchConfig { + max_fill_ms: args.sf_args.max_batch_fill_duration_ms, + memory_budget_ratio: 0.2, + max_bytes: 8 * 1024 * 1024, + }, + table_error_retry_delay_ms: 10000, + table_error_retry_max_attempts: 5, + max_table_sync_workers: args.sf_args.max_table_sync_workers, + memory_refresh_interval_ms: 100, + memory_backpressure: Some(MemoryBackpressureConfig::default()), + table_sync_copy: TableSyncCopyConfig::default(), + invalidated_slot_behavior: InvalidatedSlotBehavior::default(), + max_copy_connections_per_table: PipelineConfig::DEFAULT_MAX_COPY_CONNECTIONS_PER_TABLE, + }; + + // Build the Snowflake connection config + let mut sf_config = Config::new( + &args.sf_args.snowflake_account, + &args.sf_args.snowflake_user, + &args.sf_args.snowflake_database, + &args.sf_args.snowflake_schema, + ); + if let Some(ref role) = args.sf_args.snowflake_role { + sf_config = sf_config.with_role(role); + } + + // Build the auth manager using key-pair authentication + let passphrase: Option = + args.sf_args.snowflake_private_key_passphrase.map(SecretString::from); + let auth = Arc::new(AuthManager::new( + &sf_config, + &args.sf_args.snowflake_private_key_path, + passphrase.as_ref(), + )?); + + // Initialize Snowflake destination -- tables will be automatically created + // to match the Postgres schema + let client = Client::new(sf_config, Arc::clone(&auth), pipeline_id); + let destination = Destination::new(client, store.clone()); + + // Create the pipeline instance with all components + let mut pipeline = Pipeline::new(pipeline_config, store, destination); + + info!( + "Starting Snowflake CDC pipeline - connecting to Postgres and initializing replication..." + ); + + // Start the pipeline - this will: + // 1. Connect to Postgres + // 2. Initialize table states based on the publication + // 3. Start apply and table sync workers + // 4. Begin streaming replication data + pipeline.start().await?; + + info!("pipeline started, data replication is now active, press ctrl+c to stop"); + + // Set up signal handler for graceful shutdown on Ctrl+C + let shutdown_signal = async { + signal::ctrl_c().await.expect("Failed to install Ctrl+C handler"); + info!("received ctrl+c signal, initiating graceful shutdown"); + }; + + // Wait for either the pipeline to complete naturally or receive a shutdown + // signal. The pipeline will run indefinitely unless an error occurs or it's + // manually stopped. + tokio::select! { + result = pipeline.wait() => { + info!("pipeline completed normally (this usually indicates an error condition)"); + result?; + } + _ = shutdown_signal => { + info!("gracefully shutting down pipeline and cleaning up resources"); + } + } + + info!("pipeline stopped, all resources cleaned up"); + + Ok(()) +} diff --git a/crates/etl-examples/src/bin/snowflake/README.md b/crates/etl-examples/src/bin/snowflake/README.md deleted file mode 100644 index 6d3838500..000000000 --- a/crates/etl-examples/src/bin/snowflake/README.md +++ /dev/null @@ -1,196 +0,0 @@ -# Snowflake CDC Benchmark - -Streams 1M+ rows from Postgres to Snowflake via Snowpipe Streaming, then measures CDC throughput. - -Demonstrates the Snowflake destination's performance characteristics. - -## Prerequisites - -1. The project's dev stack running (provides source Postgres with logical replication) -2. Snowflake account with: - - RSA key pair configured for key-pair auth - - Target database and schema created (`ETL_BENCH.CDC`) - - A role with appropriate permissions (USAGE on warehouse/DB/schema, CREATE TABLE, etc.) -3. Rust toolchain (cargo) - -## Quick Start - -### 1. Start the dev stack - -From the repo root: - -```bash -cargo x init -``` - -This starts `source-postgres` on port 5430 with `wal_level=logical` and all replication settings configured. See `scripts/docker-compose.yaml` for details. - -### 2. Create a benchmark database - -```bash -psql -h localhost -p 5430 -U postgres -c "CREATE DATABASE etl_bench;" -``` - -### 3. Seed 1M rows - -```bash -cargo run -p etl-examples --features snowflake --bin snowflake-loadgen -- seed \ - --db-url postgres://postgres:postgres@localhost:5430/etl_bench \ - --rows 1000000 -``` - -Takes ~2 minutes. Creates 3 tables (users, orders, events) in a dedicated `bench` Postgres schema and scopes the publication to it, so the ETL pipeline's internal state tables (in `public`) are not replicated. - -### 4. Configure Snowflake credentials - -Fill in the `BENCH_SNOWFLAKE_*` variables in `.env` (see `.env.example`): - -```env -BENCH_SNOWFLAKE_ACCOUNT=ORG-ACCOUNT -BENCH_SNOWFLAKE_USER=ETL_USER -BENCH_SNOWFLAKE_PRIVATE_KEY_PATH=/path/to/rsa_key.p8 -BENCH_SNOWFLAKE_DATABASE=ETL_BENCH -BENCH_SNOWFLAKE_SCHEMA=CDC -BENCH_SNOWFLAKE_ROLE=ETL_ROLE -``` - -Then load them: `source .env` - -### 5. Start the pipeline - -Snowflake args are picked up from `BENCH_SNOWFLAKE_*` env vars automatically: - -```bash -cargo run -p etl-examples --features snowflake --bin snowflake -- \ - --db-host localhost \ - --db-port 5430 \ - --db-name etl_bench \ - --db-username postgres \ - --db-password postgres \ - --publication bench_pub -``` - -The terminal dashboard shows table copy progress and throughput in real-time. - -### 6. Start CDC load generator (separate terminal) - -```bash -cargo run -p etl-examples --features snowflake --bin snowflake-loadgen -- generate \ - --db-url postgres://postgres:postgres@localhost:5430/etl_bench \ - --rate 5000 \ - --mix 40/40/20 \ - --duration 300s -``` - -## What to Expect - -- **Table copy phase**: Copies all 1M rows to Snowflake. Expect 10,000-20,000 rows/sec depending on network and row size. -- **CDC phase**: After initial copy, streams changes in real-time. Latency is typically 2-10 seconds from Postgres commit to queryable in Snowflake. - -## Verifying in Snowflake - -Table names in Snowflake follow the pattern `{pg_schema}_{pg_table}` uppercased, so Postgres `bench.events` becomes `ETL_BENCH.CDC.BENCH_EVENTS`. All examples below use fully qualified names (`ETL_BENCH.CDC.*`) matching the default config. - -**Important**: Column names are preserved as lowercase (quoted identifiers). You must double-quote them in Snowflake SQL, e.g. `"id"` not `id`. - -```sql --- List all tables -SHOW TABLES IN ETL_BENCH.CDC; - --- Count total rows (includes all CDC versions) -SELECT COUNT(*) FROM ETL_BENCH.CDC.BENCH_EVENTS; - --- Check CDC metadata columns -SELECT "id", "event_type", "_cdc_operation", "_cdc_sequence_number" -FROM ETL_BENCH.CDC.BENCH_EVENTS -ORDER BY "_cdc_sequence_number" DESC -LIMIT 20; - --- Materialize current state (latest version of each row) -SELECT * FROM ( - SELECT *, ROW_NUMBER() OVER ( - PARTITION BY "id" - ORDER BY "_cdc_sequence_number" DESC - ) AS rn - FROM ETL_BENCH.CDC.BENCH_EVENTS -) -WHERE rn = 1 AND "_cdc_operation" != 'delete'; - --- Count current (live) rows only -SELECT COUNT(*) FROM ( - SELECT "id", "_cdc_operation", ROW_NUMBER() OVER ( - PARTITION BY "id" - ORDER BY "_cdc_sequence_number" DESC - ) AS rn - FROM ETL_BENCH.CDC.BENCH_EVENTS -) -WHERE rn = 1 AND "_cdc_operation" != 'delete'; - --- Create a Dynamic Table for auto-refreshing current state -CREATE DYNAMIC TABLE ETL_BENCH.CDC.BENCH_EVENTS_CURRENT - TARGET_LAG = '1 minute' - WAREHOUSE = MY_WH -AS - SELECT * EXCLUDE ("_cdc_operation", "_cdc_sequence_number", rn) FROM ( - SELECT *, ROW_NUMBER() OVER ( - PARTITION BY "id" - ORDER BY "_cdc_sequence_number" DESC - ) AS rn - FROM ETL_BENCH.CDC.BENCH_EVENTS - ) - WHERE rn = 1 AND "_cdc_operation" != 'delete'; -``` - -## Verifying in Postgres (source comparison) - -Run these against the source database to compare row counts with Snowflake: - -```bash -psql -h localhost -p 5430 -U postgres -d etl_bench -``` - -```sql --- Total rows in source -SELECT COUNT(*) FROM bench.events; - --- If running the CDC load generator, compare live row counts: --- This should match the "current state" count from Snowflake above. -SELECT COUNT(*) FROM bench.events; - --- All tables: -SELECT - schemaname, - relname AS table_name, - n_live_tup AS approx_rows -FROM pg_stat_user_tables -WHERE schemaname = 'bench' -ORDER BY relname; -``` - -## Cleanup - -```bash -# Stop the pipeline (Ctrl+C) -# Drop the benchmark database -psql -h localhost -p 5430 -U postgres -c "DROP DATABASE etl_bench;" -``` - -In Snowflake: - -```sql -DROP TABLE IF EXISTS ETL_BENCH.CDC.BENCH_EVENTS; -DROP TABLE IF EXISTS ETL_BENCH.CDC.BENCH_USERS; -DROP TABLE IF EXISTS ETL_BENCH.CDC.BENCH_ORDERS; -DROP DYNAMIC TABLE IF EXISTS ETL_BENCH.CDC.BENCH_EVENTS_CURRENT; -``` - -## Comparison with Snowflake's CDC Guide - -This benchmark is modeled after Snowflake's official CDC guide (https://www.snowflake.com/en/developers/guides/cdc-snowpipestreaming-dynamictables/). - -Key differences: - -- Pure Rust pipeline (no Java SDK sidecar) -- Uses the Snowpipe Streaming REST API directly -- Supports all Postgres types, schema evolution, and automatic recovery -- Append-only changelog model with Dynamic Tables for materialization diff --git a/crates/etl-examples/src/bin/snowflake/commands.rs b/crates/etl-examples/src/bin/snowflake/commands.rs deleted file mode 100644 index 046ce7ef8..000000000 --- a/crates/etl-examples/src/bin/snowflake/commands.rs +++ /dev/null @@ -1,109 +0,0 @@ -use std::sync::{Arc, Mutex}; - -use etl::{ - state::TableReplicationPhase, - store::{PostgresStore, StateStore}, - types::TableId, -}; -use etl_destinations::snowflake; -use tracing::{error, info}; - -use crate::state::DashboardState; - -/// Reset a single table: truncate Snowflake data, reset state to Init. -/// The pipeline must be stopped before calling this. -pub async fn reset_table( - table_id: TableId, - table_name: &str, - store: &PostgresStore, - destination: &snowflake::Destination, - dashboard: &Arc>, -) { - info!(table = %table_name, "resetting table: truncating Snowflake data and resetting state"); - - if let Err(e) = destination.committed_offset(table_id).await { - info!(table = %table_name, "table has no open channel ({}), skipping truncate", e); - } else { - // Table has an open channel — try to truncate via the client. - // We access the underlying client through the destination's public API. - // Since we can't truncate through the Destination trait without a - // ReplicatedTableSchema, we reset state and let the pipeline re-copy on - // next start. - info!(table = %table_name, "channel exists, data will be re-copied on restart"); - } - - match store.update_table_replication_state(table_id, TableReplicationPhase::Init).await { - Ok(()) => { - info!(table = %table_name, "table state reset to Init"); - dashboard.lock().unwrap().set_status(format!( - "{table_name}: state reset to Init, restart pipeline to re-copy" - )); - } - Err(e) => { - error!(table = %table_name, error = %e, "failed to reset table state"); - dashboard.lock().unwrap().set_status(format!("{table_name}: reset failed: {e}")); - } - } -} - -/// Reset ALL tables to Init state. -#[allow(dead_code)] -pub async fn reset_all_tables( - store: &PostgresStore, - destination: &snowflake::Destination, - dashboard: &Arc>, -) { - let tables: Vec<(TableId, String)> = { - let dash = dashboard.lock().unwrap(); - dash.tables.iter().map(|t| (t.table_id, t.destination_name.clone())).collect() - }; - - for (id, name) in &tables { - reset_table(*id, name, store, destination, dashboard).await; - } - - dashboard - .lock() - .unwrap() - .set_status("All tables reset. Restart pipeline to re-copy.".to_owned()); -} - -pub const HELP_TEXT: &str = "\ -Keyboard Shortcuts: - j/k or Up/Down Navigate table list - J/K Scroll logs up/down (3 lines) - PgUp/PgDn Scroll logs (page) - F1 / ? Toggle this help - F2 Reset selected table - F3 Restart pipeline - F4 Cycle log filter (All / Warn+ / Error) - F5 Force-sync Snowflake offsets now - F10 / q Quit - -Table States: - init Table registered, not yet copied - copying Initial data copy in progress - copied Copy finished, waiting for sync - sync_wait Waiting for apply worker pause - catchup Catching up to sync LSN - sync_done Sync complete - ready Fully synced, streaming CDC - errored Error occurred (see detail panel) - -Commands: - Reset (F2) Stops the pipeline, resets the selected table's state - to Init, then restarts. The pipeline will re-copy all - data from Postgres. Use when a table is stuck in error. - - Restart (F3) Stops and restarts the pipeline without changing any - table states. Tables resume from their current state. - - Log (F4) Cycles the log filter: All shows everything, Warn+ - shows only warnings and errors, Error shows errors only. - - Sync (F5) Queries Snowflake's channel status API to fetch the - last committed offset for each table. Offsets also - auto-refresh every 30 seconds. - -Press any key to close this help. -"; diff --git a/crates/etl-examples/src/bin/snowflake/logging.rs b/crates/etl-examples/src/bin/snowflake/logging.rs deleted file mode 100644 index 0a64492e8..000000000 --- a/crates/etl-examples/src/bin/snowflake/logging.rs +++ /dev/null @@ -1,86 +0,0 @@ -use std::{ - collections::VecDeque, - sync::{Arc, Mutex}, -}; - -use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt}; - -pub const MAX_LOG_LINES: usize = 2000; - -pub struct TuiLogLayer { - buffer: Arc>>, -} - -struct FieldVisitor { - message: String, - fields: Vec<(String, String)>, -} - -impl FieldVisitor { - fn new() -> Self { - Self { message: String::new(), fields: Vec::new() } - } - - fn into_line(self) -> String { - if self.fields.is_empty() { - self.message - } else { - let extras: Vec = - self.fields.into_iter().map(|(k, v)| format!("{k}={v}")).collect(); - format!("{} {}", self.message, extras.join(" ")) - } - } -} - -impl tracing::field::Visit for FieldVisitor { - fn record_debug(&mut self, field: &tracing::field::Field, value: &dyn std::fmt::Debug) { - if field.name() == "message" { - self.message = format!("{value:?}"); - } else { - self.fields.push((field.name().to_owned(), format!("{value:?}"))); - } - } - - fn record_str(&mut self, field: &tracing::field::Field, value: &str) { - if field.name() == "message" { - self.message = value.to_owned(); - } else { - self.fields.push((field.name().to_owned(), value.to_owned())); - } - } -} - -impl tracing_subscriber::Layer for TuiLogLayer { - fn on_event( - &self, - event: &tracing::Event<'_>, - _ctx: tracing_subscriber::layer::Context<'_, S>, - ) { - let level = event.metadata().level(); - let target = event.metadata().target(); - let mut visitor = FieldVisitor::new(); - event.record(&mut visitor); - let line = format!("[{level}] {target}: {}", visitor.into_line()); - let mut buf = self.buffer.lock().unwrap(); - buf.push_back(line); - if buf.len() > MAX_LOG_LINES { - buf.pop_front(); - } - } -} - -pub fn init_tracing(log_buffer: Arc>>) { - if std::env::var("RUST_LOG").is_err() { - unsafe { - std::env::set_var("RUST_LOG", "info"); - } - } - - tracing_subscriber::registry() - .with( - tracing_subscriber::EnvFilter::try_from_default_env() - .unwrap_or_else(|_| "snowflake=info".into()), - ) - .with(TuiLogLayer { buffer: log_buffer }) - .init(); -} diff --git a/crates/etl-examples/src/bin/snowflake/main.rs b/crates/etl-examples/src/bin/snowflake/main.rs deleted file mode 100644 index 1a1810ac4..000000000 --- a/crates/etl-examples/src/bin/snowflake/main.rs +++ /dev/null @@ -1,461 +0,0 @@ -/// Snowflake CDC Pipeline — Control Plane TUI -/// -/// Streams Postgres CDC to Snowflake via Snowpipe Streaming, with a live -/// terminal dashboard showing per-table replication progress, offsets, and -/// throughput. -/// -/// Supports resetting errored tables and restarting the pipeline. -/// -/// See README.md for the full guide. -mod commands; -mod logging; -mod state; -mod tui; - -use std::{ - collections::{BTreeMap, VecDeque}, - error::Error, - sync::{Arc, Mutex, Once}, - time::{Duration, Instant}, -}; - -use clap::{Args, Parser}; -use crossterm::event::{Event, KeyCode}; -use etl::{ - config::{ - BatchConfig, InvalidatedSlotBehavior, MemoryBackpressureConfig, PgConnectionConfig, - PipelineConfig, TableSyncCopyConfig, TcpKeepaliveConfig, TlsConfig, - }, - pipeline::Pipeline, - store::PostgresStore, - types::TableId, -}; -use etl_destinations::snowflake::{AuthManager, Client, Config, Destination}; -use etl_telemetry::metrics::init_metrics_handle; -use secrecy::SecretString; -use tracing::{error, info}; - -use crate::state::{DashboardState, GlobalPhase}; - -static INIT_CRYPTO: Once = Once::new(); - -fn install_crypto_provider() { - INIT_CRYPTO.call_once(|| { - rustls::crypto::aws_lc_rs::default_provider() - .install_default() - .expect("failed to install default crypto provider"); - }); -} - -#[derive(Debug, Parser)] -#[command(name = "snowflake", version, about, arg_required_else_help = true)] -struct AppArgs { - #[clap(flatten)] - db_args: DbArgs, - #[clap(flatten)] - sf_args: SnowflakeArgs, - #[arg(long)] - publication: String, - #[arg(long, default_value = "5000")] - max_batch_fill_duration_ms: u64, - #[arg(long, default_value = "4")] - max_table_sync_workers: u16, -} - -#[derive(Debug, Args)] -struct DbArgs { - #[arg(long)] - db_host: String, - #[arg(long)] - db_port: u16, - #[arg(long)] - db_name: String, - #[arg(long)] - db_username: String, - #[arg(long)] - db_password: Option, -} - -#[derive(Debug, Clone, Args)] -struct SnowflakeArgs { - #[arg(long, env = "BENCH_SNOWFLAKE_ACCOUNT")] - snowflake_account: String, - #[arg(long, env = "BENCH_SNOWFLAKE_USER")] - snowflake_user: String, - #[arg(long, env = "BENCH_SNOWFLAKE_PRIVATE_KEY_PATH")] - snowflake_private_key_path: String, - #[arg(long, env = "BENCH_SNOWFLAKE_PRIVATE_KEY_PASSPHRASE")] - snowflake_private_key_passphrase: Option, - #[arg(long, env = "BENCH_SNOWFLAKE_DATABASE")] - snowflake_database: String, - #[arg(long, env = "BENCH_SNOWFLAKE_SCHEMA", default_value = "CDC")] - snowflake_schema: String, - #[arg(long, env = "BENCH_SNOWFLAKE_ROLE")] - snowflake_role: Option, -} - -#[tokio::main] -async fn main() -> Result<(), Box> { - if let Err(err) = run().await { - error!(error = %err, "fatal error"); - std::process::exit(1); - } - Ok(()) -} - -async fn run() -> Result<(), Box> { - let log_buffer: Arc>> = Arc::new(Mutex::new(VecDeque::new())); - logging::init_tracing(Arc::clone(&log_buffer)); - install_crypto_provider(); - - let args = AppArgs::parse(); - - let pg_config = PgConnectionConfig { - host: args.db_args.db_host, - hostaddr: None, - port: args.db_args.db_port, - name: args.db_args.db_name, - username: args.db_args.db_username, - password: args.db_args.db_password.map(Into::into), - tls: TlsConfig { trusted_root_certs: String::new(), enabled: false }, - keepalive: TcpKeepaliveConfig::default(), - }; - - info!("counting source tables and rows..."); - let (table_count, estimated_rows) = count_source_tables(&pg_config).await?; - info!(tables = table_count, estimated_rows, "source database stats"); - - let pipeline_id = 1u64; - - let pipeline_config = PipelineConfig { - id: pipeline_id, - publication_name: args.publication.clone(), - pg_connection: pg_config.clone(), - batch: BatchConfig { - max_fill_ms: args.max_batch_fill_duration_ms, - memory_budget_ratio: 0.2, - max_bytes: 8 * 1024 * 1024, - }, - table_error_retry_delay_ms: 10000, - table_error_retry_max_attempts: 5, - max_table_sync_workers: args.max_table_sync_workers, - memory_refresh_interval_ms: 100, - memory_backpressure: Some(MemoryBackpressureConfig::default()), - table_sync_copy: TableSyncCopyConfig::default(), - invalidated_slot_behavior: InvalidatedSlotBehavior::Recreate, - max_copy_connections_per_table: PipelineConfig::DEFAULT_MAX_COPY_CONNECTIONS_PER_TABLE, - }; - - let args = args.sf_args.clone(); - let passphrase: Option = - args.snowflake_private_key_passphrase.map(SecretString::from); - - let mut config = Config::new( - &args.snowflake_account, - &args.snowflake_user, - &args.snowflake_database, - &args.snowflake_schema, - ); - if let Some(ref role) = args.snowflake_role { - config = config.with_role(role); - } - - let auth = - Arc::new(AuthManager::new(&config, &args.snowflake_private_key_path, passphrase.as_ref())?); - - let store = PostgresStore::new(pipeline_id, pg_config.clone()).await?; - - let dashboard = Arc::new(Mutex::new(DashboardState::new(table_count, estimated_rows))); - - // Build initial table name map from destination metadata - let mut table_names = state::build_table_name_map(&store).await; - - // Start pipeline - let client = Client::new(config.clone(), Arc::clone(&auth), pipeline_id); - let mut destination = Destination::new(client, store.clone()); - let mut pipeline = Pipeline::new(pipeline_config.clone(), store.clone(), destination.clone()); - - info!("starting Snowflake CDC pipeline..."); - pipeline.start().await?; - info!("pipeline started, press F1 for help"); - - let metrics_handle = init_metrics_handle().ok(); - - let (mut terminal, _terminal_guard) = tui::setup_terminal()?; - - // Monitor state - let mut samples: Vec = Vec::new(); - let mut last_rows: u64 = 0; - let mut last_tick = Instant::now(); - let mut last_refresh = Instant::now(); - let mut last_offset_fetch = Instant::now(); - let mut phase_times: BTreeMap = BTreeMap::new(); - let start_time = Instant::now(); - let page_size = 10usize; - - loop { - // Render - { - let state = dashboard.lock().unwrap(); - let lb = Arc::clone(&log_buffer); - terminal.draw(|f| tui::render(f, &state, &lb))?; - } - - // Periodic refresh (every 2s) - if last_refresh.elapsed() >= Duration::from_secs(2) { - let metrics_text = match metrics_handle.as_ref() { - Some(h) => h.render(), - None => String::new(), - }; - - // Refresh table name map (picks up newly discovered tables) - let new_names = state::build_table_name_map(&store).await; - for (id, name) in new_names { - table_names.entry(id).or_insert(name); - } - - state::refresh_dashboard( - &dashboard, - &store, - &metrics_text, - start_time, - &mut samples, - &mut last_rows, - &mut last_tick, - &table_names, - &mut phase_times, - ) - .await; - - last_refresh = Instant::now(); - } - - // Periodic Snowflake offset sync (every 30s) - if last_offset_fetch.elapsed() >= Duration::from_secs(30) { - state::fetch_snowflake_offsets(&dashboard, &destination).await; - last_offset_fetch = Instant::now(); - } - - // Handle input - if crossterm::event::poll(Duration::from_millis(100))? - && let Event::Key(key) = crossterm::event::read()? - { - // If help is open, any key closes it - { - let mut dash = dashboard.lock().unwrap(); - if dash.show_help { - dash.show_help = false; - continue; - } - } - - match key.code { - // Quit - KeyCode::F(10) | KeyCode::Char('q') => { - info!("quit requested, shutting down..."); - pipeline.shutdown(); - break; - } - - // Help - KeyCode::F(1) | KeyCode::Char('?') => { - dashboard.lock().unwrap().show_help = true; - } - - // Table navigation - KeyCode::Char('j') | KeyCode::Down => { - dashboard.lock().unwrap().select_next(); - } - KeyCode::Char('k') | KeyCode::Up => { - dashboard.lock().unwrap().select_prev(); - } - - // Log scrolling - KeyCode::Char('J') => { - let mut dash = dashboard.lock().unwrap(); - dash.log_scroll = dash.log_scroll.saturating_add(3); - } - KeyCode::Char('K') => { - let mut dash = dashboard.lock().unwrap(); - dash.log_scroll = dash.log_scroll.saturating_sub(3); - } - KeyCode::PageUp => { - let mut dash = dashboard.lock().unwrap(); - dash.log_scroll = dash.log_scroll.saturating_add(page_size); - } - KeyCode::PageDown => { - let mut dash = dashboard.lock().unwrap(); - dash.log_scroll = dash.log_scroll.saturating_sub(page_size); - } - - // F2: Reset selected table - KeyCode::F(2) => { - let (table_id, table_name) = { - let dash = dashboard.lock().unwrap(); - match dash.selected_table_info() { - Some(t) => (t.table_id, t.destination_name.clone()), - None => continue, - } - }; - - info!("F2: resetting table {}, stopping pipeline...", table_name); - dashboard.lock().unwrap().set_status(format!("Resetting {table_name}...")); - - // Stop pipeline (wait consumes self, so always rebuild after) - pipeline.shutdown(); - let wait_err = pipeline.wait().await.err(); - - let client = Client::new(config.clone(), Arc::clone(&auth), pipeline_id); - destination = Destination::new(client, store.clone()); - pipeline = - Pipeline::new(pipeline_config.clone(), store.clone(), destination.clone()); - - if let Some(e) = wait_err { - error!(error = %e, "pipeline shutdown failed during reset"); - dashboard.lock().unwrap().pipeline_running = false; - dashboard.lock().unwrap().phase = GlobalPhase::Stopped; - dashboard.lock().unwrap().set_status(format!("Reset failed: {e}")); - continue; - } - dashboard.lock().unwrap().pipeline_running = false; - dashboard.lock().unwrap().phase = GlobalPhase::Stopped; - - // Reset table state - commands::reset_table(table_id, &table_name, &store, &destination, &dashboard) - .await; - - // Restart pipeline - info!("restarting pipeline after reset..."); - if let Err(e) = pipeline.start().await { - error!(error = %e, "pipeline restart failed after reset"); - dashboard.lock().unwrap().set_status(format!( - "Reset done but restart failed: {e} -- press F3 to retry" - )); - continue; - } - dashboard.lock().unwrap().pipeline_running = true; - - info!("pipeline restarted after table reset"); - dashboard - .lock() - .unwrap() - .set_status(format!("{table_name} reset complete, pipeline restarted")); - } - - // F3: Restart pipeline - KeyCode::F(3) => { - info!("F3: restarting pipeline..."); - dashboard.lock().unwrap().set_status("Restarting pipeline...".to_owned()); - - // Stop (wait consumes self, so always rebuild after) - pipeline.shutdown(); - let wait_err = pipeline.wait().await.err(); - - let client = Client::new(config.clone(), Arc::clone(&auth), pipeline_id); - destination = Destination::new(client, store.clone()); - pipeline = - Pipeline::new(pipeline_config.clone(), store.clone(), destination.clone()); - - if let Some(e) = wait_err { - error!(error = %e, "pipeline shutdown failed during restart"); - dashboard.lock().unwrap().pipeline_running = false; - dashboard.lock().unwrap().phase = GlobalPhase::Stopped; - dashboard - .lock() - .unwrap() - .set_status(format!("Restart failed: {e} -- press F3 to retry")); - continue; - } - dashboard.lock().unwrap().pipeline_running = false; - dashboard.lock().unwrap().phase = GlobalPhase::Stopped; - - // Restart - if let Err(e) = pipeline.start().await { - error!(error = %e, "pipeline restart failed"); - dashboard - .lock() - .unwrap() - .set_status(format!("Restart failed: {e} -- press F3 to retry")); - continue; - } - dashboard.lock().unwrap().pipeline_running = true; - - info!("pipeline restarted"); - dashboard.lock().unwrap().set_status("Pipeline restarted".to_owned()); - } - - // F4: Cycle log level filter - KeyCode::F(4) => { - let mut dash = dashboard.lock().unwrap(); - dash.log_level_filter = dash.log_level_filter.next(); - let label = dash.log_level_filter.label(); - dash.set_status(format!("Log filter: {label}")); - } - - // F5: Sync Snowflake offsets - KeyCode::F(5) => { - info!("F5: syncing Snowflake offsets..."); - dashboard - .lock() - .unwrap() - .set_status("Querying Snowflake offsets...".to_owned()); - - state::fetch_snowflake_offsets(&dashboard, &destination).await; - last_offset_fetch = Instant::now(); - - dashboard.lock().unwrap().set_status("Snowflake offsets synced".to_owned()); - info!("Snowflake offset sync complete"); - } - - _ => {} - } - } - } - - // Wait for pipeline to finish - pipeline.wait().await?; - info!("pipeline shut down cleanly"); - - Ok(()) -} - -async fn count_source_tables(pg_config: &PgConnectionConfig) -> Result<(u64, u64), Box> { - let conn_str = format!( - "host={} port={} dbname={} user={} {}", - pg_config.host, - pg_config.port, - pg_config.name, - pg_config.username, - pg_config - .password - .as_ref() - .map(|p| { - use secrecy::ExposeSecret; - format!("password={}", p.expose_secret()) - }) - .unwrap_or_default(), - ); - - let (client, conn) = tokio_postgres::connect(&conn_str, tokio_postgres::NoTls).await?; - tokio::spawn(conn); - - let table_row = client - .query_one( - "SELECT count(*) FROM information_schema.tables WHERE table_schema = 'bench' AND \ - table_type = 'BASE TABLE'", - &[], - ) - .await?; - let table_count: i64 = table_row.get(0); - - let row_row = client - .query_one( - "SELECT coalesce(sum(n_live_tup), 0)::bigint FROM pg_stat_user_tables WHERE \ - schemaname = 'bench'", - &[], - ) - .await?; - let row_count: i64 = row_row.get(0); - - Ok((table_count as u64, row_count as u64)) -} diff --git a/crates/etl-examples/src/bin/snowflake/state.rs b/crates/etl-examples/src/bin/snowflake/state.rs deleted file mode 100644 index 240f4d0d8..000000000 --- a/crates/etl-examples/src/bin/snowflake/state.rs +++ /dev/null @@ -1,451 +0,0 @@ -use std::{ - collections::BTreeMap, - sync::{Arc, Mutex}, - time::{Duration, Instant}, -}; - -use etl::{ - state::TableReplicationPhase, - store::{PostgresStore, StateStore}, - types::TableId, -}; -use etl_destinations::snowflake::{Destination, OffsetToken}; - -/// Per-table information displayed in the table list and detail panel. -#[derive(Clone)] -#[allow(dead_code)] -pub struct TableInfo { - pub table_id: TableId, - pub destination_name: String, - pub phase: TableReplicationPhase, - pub rows_synced: u64, - pub throughput: f64, - pub local_offset: Option, - pub snowflake_offset: Option, - pub last_known_lsn: Option, - pub error_reason: Option, - pub error_solution: Option, - pub phase_entered_at: Option, -} - -impl TableInfo { - pub fn phase_label(&self) -> &'static str { - match &self.phase { - TableReplicationPhase::Init => "init", - TableReplicationPhase::DataSync => "copying", - TableReplicationPhase::FinishedCopy => "copied", - TableReplicationPhase::SyncWait { .. } => "sync_wait", - TableReplicationPhase::Catchup { .. } => "catchup", - TableReplicationPhase::SyncDone { .. } => "sync_done", - TableReplicationPhase::Ready => "ready", - TableReplicationPhase::Errored { .. } => "errored", - } - } - - pub fn is_errored(&self) -> bool { - matches!(&self.phase, TableReplicationPhase::Errored { .. }) - } - - pub fn is_copying(&self) -> bool { - matches!(&self.phase, TableReplicationPhase::DataSync | TableReplicationPhase::Init) - } - - pub fn is_ready(&self) -> bool { - matches!(&self.phase, TableReplicationPhase::Ready) - } -} - -/// Global dashboard state shared between the monitor task and the TUI render -/// loop. -#[allow(dead_code)] -pub struct DashboardState { - pub tables: Vec, - pub selected_table: usize, - - pub total_rows_synced: u64, - pub total_copy_rows: u64, - pub total_cdc_events: u64, - pub estimated_rows: u64, - pub table_count: u64, - - pub elapsed: Duration, - pub phase: GlobalPhase, - pub copy_elapsed: Option, - - pub throughput_current: f64, - pub throughput_avg: f64, - pub throughput_min: f64, - pub throughput_max: f64, - - pub api_calls: u64, - pub api_errors: u64, - pub channel_recoveries: u64, - - pub log_scroll: usize, - pub show_help: bool, - pub status_message: Option<(String, Instant)>, - pub log_level_filter: LogLevel, - - pub pipeline_running: bool, -} - -#[derive(Clone, Copy, PartialEq, Eq)] -pub enum GlobalPhase { - TableCopy, - CdcStreaming, - Stopped, -} - -#[derive(Clone, Copy, PartialEq, Eq)] -pub enum LogLevel { - All, - WarnAndAbove, - ErrorOnly, -} - -impl LogLevel { - pub fn next(self) -> Self { - match self { - Self::All => Self::WarnAndAbove, - Self::WarnAndAbove => Self::ErrorOnly, - Self::ErrorOnly => Self::All, - } - } - - pub fn label(self) -> &'static str { - match self { - Self::All => "All", - Self::WarnAndAbove => "Warn+", - Self::ErrorOnly => "Error", - } - } -} - -impl GlobalPhase { - pub fn label(&self) -> &'static str { - match self { - Self::TableCopy => "Table Copy", - Self::CdcStreaming => "CDC Streaming", - Self::Stopped => "Stopped", - } - } -} - -#[allow(dead_code)] -impl DashboardState { - pub fn new(table_count: u64, estimated_rows: u64) -> Self { - Self { - tables: Vec::new(), - selected_table: 0, - total_rows_synced: 0, - total_copy_rows: 0, - total_cdc_events: 0, - estimated_rows, - table_count, - elapsed: Duration::ZERO, - phase: GlobalPhase::TableCopy, - copy_elapsed: None, - throughput_current: 0.0, - throughput_avg: 0.0, - throughput_min: 0.0, - throughput_max: 0.0, - api_calls: 0, - api_errors: 0, - channel_recoveries: 0, - log_scroll: 0, - show_help: false, - status_message: None, - log_level_filter: LogLevel::All, - pipeline_running: true, - } - } - - pub fn selected_table_info(&self) -> Option<&TableInfo> { - self.tables.get(self.selected_table) - } - - pub fn selected_table_id(&self) -> Option { - self.tables.get(self.selected_table).map(|t| t.table_id) - } - - pub fn select_next(&mut self) { - if !self.tables.is_empty() { - self.selected_table = (self.selected_table + 1).min(self.tables.len() - 1); - } - } - - pub fn select_prev(&mut self) { - self.selected_table = self.selected_table.saturating_sub(1); - } - - pub fn set_status(&mut self, msg: String) { - self.status_message = Some((msg, Instant::now())); - } - - pub fn clear_stale_status(&mut self) { - if let Some((_, at)) = &self.status_message - && at.elapsed() > Duration::from_secs(10) - { - self.status_message = None; - } - } -} - -pub fn parse_prometheus_metric(text: &str, metric_name: &str) -> Option { - let prefix = format!("{metric_name} "); - text.lines() - .find(|line| line.starts_with(&prefix)) - .and_then(|line| line[prefix.len()..].trim().parse::().ok()) -} - -pub fn parse_prometheus_metric_sum(text: &str, metric_name: &str) -> Option { - let mut total = 0.0; - let mut found = false; - for line in text.lines() { - if let Some(rest) = line.strip_prefix(metric_name) - && (rest.starts_with(' ') || rest.starts_with('{')) - && let Some(val) = rest.split_whitespace().last().and_then(|v| v.parse::().ok()) - { - total += val; - found = true; - } - } - found.then_some(total) -} - -pub fn parse_prometheus_metric_with_label( - text: &str, - metric_name: &str, - label_key: &str, - label_value: &str, -) -> Option { - let needle = format!("{label_key}=\"{label_value}\""); - let mut total = 0.0; - let mut found = false; - for line in text.lines() { - if let Some(rest) = line.strip_prefix(metric_name) - && rest.starts_with('{') - && rest.contains(&needle) - && let Some(val) = rest.split_whitespace().last().and_then(|v| v.parse::().ok()) - { - total += val; - found = true; - } - } - found.then_some(total) -} - -/// Update dashboard state from Prometheus metrics and the state store. -#[allow(clippy::too_many_arguments)] -pub async fn refresh_dashboard( - dashboard: &Arc>, - store: &PostgresStore, - metrics_text: &str, - start_time: Instant, - samples: &mut Vec, - last_rows: &mut u64, - last_tick: &mut Instant, - table_names: &BTreeMap, - phase_times: &mut BTreeMap, -) { - let elapsed = start_time.elapsed(); - - let rows_synced = - parse_prometheus_metric(metrics_text, "etl_snowflake_batch_size_sum").unwrap_or(0.0) as u64; - let copy_rows = parse_prometheus_metric_with_label( - metrics_text, - "etl_events_processed_total", - "action", - "table_copy", - ) - .unwrap_or(0.0) as u64; - let cdc_events = parse_prometheus_metric_with_label( - metrics_text, - "etl_events_processed_total", - "action", - "table_streaming", - ) - .unwrap_or(0.0) as u64; - let api_errors = parse_prometheus_metric_sum(metrics_text, "etl_snowflake_insert_errors_total") - .unwrap_or(0.0) as u64; - let channel_recoveries = - parse_prometheus_metric_sum(metrics_text, "etl_snowflake_channel_recoveries_total") - .unwrap_or(0.0) as u64; - - let tick_elapsed = last_tick.elapsed().as_secs_f64(); - let delta = rows_synced.saturating_sub(*last_rows); - let throughput_current = if tick_elapsed > 0.0 { delta as f64 / tick_elapsed } else { 0.0 }; - *last_rows = rows_synced; - *last_tick = Instant::now(); - - if throughput_current > 0.0 { - samples.push(throughput_current); - } - let throughput_avg = - if samples.is_empty() { 0.0 } else { samples.iter().sum::() / samples.len() as f64 }; - let throughput_min = samples.iter().copied().fold(f64::MAX, f64::min); - let throughput_max = samples.iter().copied().fold(0.0f64, f64::max); - - let states = store.get_table_replication_states().await.ok(); - - let mut table_infos: Vec = Vec::new(); - let mut all_ready = true; - let mut any_table = false; - - // Preserve per-table state from previous refresh cycle - let prev_table_state: BTreeMap, Option)> = { - let dash = dashboard.lock().unwrap(); - dash.tables - .iter() - .map(|t| (t.table_id, (t.last_known_lsn.clone(), t.snowflake_offset.clone()))) - .collect() - }; - - if let Some(ref states) = states { - for (id, phase) in states.iter() { - any_table = true; - let dest_name = table_names.get(id).cloned().unwrap_or_else(|| format!("{id}")); - - let phase_label = phase_label_str(phase); - let prev = phase_times.get(id).map(|(l, _)| l.as_str()); - if prev != Some(phase_label) { - phase_times.insert(*id, (phase_label.to_owned(), Instant::now())); - } - let phase_entered = phase_times.get(id).map(|(_, t)| *t); - - let (error_reason, error_solution) = - if let TableReplicationPhase::Errored { reason, solution, .. } = phase { - all_ready = false; - (Some(reason.clone()), solution.clone()) - } else { - if !matches!(phase, TableReplicationPhase::Ready) { - all_ready = false; - } - (None, None) - }; - - let current_lsn = extract_lsn(phase); - let (prev_lsn, prev_sf_offset) = - prev_table_state.get(id).cloned().unwrap_or((None, None)); - let last_known_lsn = current_lsn.clone().or(prev_lsn); - - table_infos.push(TableInfo { - table_id: *id, - destination_name: dest_name, - phase: phase.clone(), - rows_synced: 0, - throughput: 0.0, - local_offset: current_lsn.or_else(|| offset_description(phase)), - snowflake_offset: prev_sf_offset, - last_known_lsn, - error_reason, - error_solution, - phase_entered_at: phase_entered, - }); - } - } - - let global_phase = - if !any_table || !all_ready { GlobalPhase::TableCopy } else { GlobalPhase::CdcStreaming }; - - let mut dash = dashboard.lock().unwrap(); - let prev_selected = dash.selected_table; - dash.elapsed = elapsed; - dash.total_rows_synced = rows_synced; - dash.total_copy_rows = copy_rows; - dash.total_cdc_events = cdc_events; - dash.throughput_current = throughput_current; - dash.throughput_avg = throughput_avg; - dash.throughput_min = if throughput_min == f64::MAX { 0.0 } else { throughput_min }; - dash.throughput_max = throughput_max; - dash.api_errors = api_errors; - dash.channel_recoveries = channel_recoveries; - dash.tables = table_infos; - dash.selected_table = prev_selected.min(dash.tables.len().saturating_sub(1)); - - let was_copying = dash.phase == GlobalPhase::TableCopy; - if was_copying && global_phase == GlobalPhase::CdcStreaming { - dash.copy_elapsed = Some(elapsed); - } - if dash.pipeline_running { - dash.phase = global_phase; - } - - dash.clear_stale_status(); -} - -fn phase_label_str(phase: &TableReplicationPhase) -> &'static str { - match phase { - TableReplicationPhase::Init => "init", - TableReplicationPhase::DataSync => "data_sync", - TableReplicationPhase::FinishedCopy => "finished_copy", - TableReplicationPhase::SyncWait { .. } => "sync_wait", - TableReplicationPhase::Catchup { .. } => "catchup", - TableReplicationPhase::SyncDone { .. } => "sync_done", - TableReplicationPhase::Ready => "ready", - TableReplicationPhase::Errored { .. } => "errored", - } -} - -fn extract_lsn(phase: &TableReplicationPhase) -> Option { - match phase { - TableReplicationPhase::SyncWait { lsn } - | TableReplicationPhase::Catchup { lsn } - | TableReplicationPhase::SyncDone { lsn } => Some(format!("{lsn}")), - _ => None, - } -} - -fn offset_description(phase: &TableReplicationPhase) -> Option { - match phase { - TableReplicationPhase::Init => Some("not started".to_owned()), - TableReplicationPhase::DataSync => Some("copying...".to_owned()), - TableReplicationPhase::FinishedCopy => Some("copy done, syncing".to_owned()), - TableReplicationPhase::Ready => Some("streaming".to_owned()), - _ => None, - } -} - -/// Build a map from TableId to destination table name by querying the store. -pub async fn build_table_name_map(store: &PostgresStore) -> BTreeMap { - let mut map = BTreeMap::new(); - let states = store.get_table_replication_states().await.ok(); - if let Some(states) = states { - for id in states.keys() { - if let Ok(Some(meta)) = store.get_destination_table_metadata(*id).await { - map.insert(*id, meta.destination_table_id); - } - } - } - map -} - -/// Fetch Snowflake committed offsets for all tables. -pub async fn fetch_snowflake_offsets( - dashboard: &Arc>, - destination: &Destination, -) { - let table_ids: Vec = { - let dash = dashboard.lock().unwrap(); - dash.tables.iter().map(|t| t.table_id).collect() - }; - - let mut offsets: BTreeMap> = BTreeMap::new(); - for id in &table_ids { - match destination.committed_offset(*id).await { - Ok(offset) => { - offsets.insert(*id, offset); - } - Err(_) => { - offsets.insert(*id, None); - } - } - } - - let mut dash = dashboard.lock().unwrap(); - for table in &mut dash.tables { - if let Some(offset) = offsets.get(&table.table_id) { - table.snowflake_offset = offset.as_ref().map(|o| format!("{o}")); - } - } -} diff --git a/crates/etl-examples/src/bin/snowflake/tui.rs b/crates/etl-examples/src/bin/snowflake/tui.rs deleted file mode 100644 index 6d19e2309..000000000 --- a/crates/etl-examples/src/bin/snowflake/tui.rs +++ /dev/null @@ -1,532 +0,0 @@ -use std::{ - collections::VecDeque, - io::Stdout, - sync::{Arc, Mutex}, -}; - -type Term = ratatui::Terminal>; - -use ratatui::{ - Terminal, - backend::CrosstermBackend, - layout::{Constraint, Direction, Layout, Rect}, - style::{Color, Modifier, Style}, - symbols, - text::{Line, Span}, - widgets::{Block, Borders, Cell, Clear, LineGauge, Paragraph, Row, Table, Wrap}, -}; - -use crate::{ - commands, - state::{DashboardState, GlobalPhase, LogLevel, TableInfo}, -}; - -pub struct TerminalGuard; - -impl Drop for TerminalGuard { - fn drop(&mut self) { - restore_terminal(); - } -} - -pub fn setup_terminal() -> Result<(Term, TerminalGuard), Box> { - crossterm::terminal::enable_raw_mode()?; - let mut stdout = std::io::stdout(); - crossterm::execute!(stdout, crossterm::terminal::EnterAlternateScreen)?; - Ok((Terminal::new(CrosstermBackend::new(stdout))?, TerminalGuard)) -} - -pub fn restore_terminal() { - let _ = crossterm::terminal::disable_raw_mode(); - let _ = crossterm::execute!(std::io::stdout(), crossterm::terminal::LeaveAlternateScreen); -} - -pub fn render( - frame: &mut ratatui::Frame, - state: &DashboardState, - log_buffer: &Arc>>, -) { - if state.show_help { - render_help_overlay(frame); - return; - } - - let main_chunks = Layout::default() - .direction(Direction::Vertical) - .constraints([ - Constraint::Length(5), // global stats + gauge - Constraint::Percentage(45), // table + detail - Constraint::Min(6), // logs (takes remaining space) - Constraint::Length(1), // command bar - ]) - .split(frame.area()); - - render_global_stats(frame, main_chunks[0], state); - render_middle_panel(frame, main_chunks[1], state); - render_log_panel(frame, main_chunks[2], log_buffer, state.log_scroll, state.log_level_filter); - render_command_bar(frame, main_chunks[3], state); -} - -fn render_global_stats(frame: &mut ratatui::Frame, area: Rect, state: &DashboardState) { - let block = Block::default().title(" Snowflake CDC Pipeline ").borders(Borders::ALL); - let inner = block.inner(area); - frame.render_widget(block, area); - - let elapsed_secs = state.elapsed.as_secs(); - let elapsed_str = format!( - "{:02}:{:02}:{:02}", - elapsed_secs / 3600, - (elapsed_secs % 3600) / 60, - elapsed_secs % 60 - ); - - let phase_color = match state.phase { - GlobalPhase::TableCopy => Color::Green, - GlobalPhase::CdcStreaming => Color::Cyan, - GlobalPhase::Stopped => Color::Red, - }; - - let running_indicator = if state.pipeline_running { - Span::styled(" RUNNING ", Style::default().fg(Color::Black).bg(Color::Green)) - } else { - Span::styled(" STOPPED ", Style::default().fg(Color::Black).bg(Color::Red)) - }; - - // Row 1: phase, elapsed, tables - let header = Line::from(vec![ - running_indicator, - Span::raw(" Phase: "), - Span::styled(state.phase.label(), Style::default().fg(phase_color)), - Span::raw(format!(" Elapsed: {elapsed_str} ")), - Span::raw(format!("Tables: {}", state.table_count)), - ]); - - // Row 2: copy/CDC stats as text - let stats_line = if state.phase == GlobalPhase::CdcStreaming { - let copy_part = if let Some(copy_elapsed) = state.copy_elapsed { - format!( - "Copy: {} rows in {:.1}s", - format_num(state.total_copy_rows), - copy_elapsed.as_secs_f64() - ) - } else { - format!("Copy: {} rows", format_num(state.total_copy_rows)) - }; - Line::from(vec![ - Span::raw(format!(" {copy_part}")), - Span::styled(" | ", Style::default().fg(Color::DarkGray)), - Span::styled( - format!("CDC: {} events", format_num(state.total_cdc_events)), - Style::default().fg(Color::Cyan), - ), - Span::styled(" | ", Style::default().fg(Color::DarkGray)), - Span::raw(format!("{:.0} r/s", state.throughput_current)), - if state.api_errors > 0 { - Span::styled( - format!(" Errors: {}", state.api_errors), - Style::default().fg(Color::Red), - ) - } else { - Span::raw("") - }, - ]) - } else { - let total = state.estimated_rows.max(1); - let pct = (state.total_rows_synced as f64 / total as f64 * 100.0).min(100.0); - Line::from(vec![ - Span::raw(format!( - " Copy: {} / {} ({:.0}%)", - format_num(state.total_rows_synced), - format_num(total), - pct - )), - Span::styled(" | ", Style::default().fg(Color::DarkGray)), - Span::raw(format!( - "{:.0} r/s (avg: {:.0})", - state.throughput_current, state.throughput_avg - )), - if state.api_errors > 0 { - Span::styled( - format!(" Errors: {}", state.api_errors), - Style::default().fg(Color::Red), - ) - } else { - Span::raw("") - }, - ]) - }; - - // Row 3: progress gauge - let total = state.estimated_rows.max(1); - let ratio = if state.phase == GlobalPhase::TableCopy { - (state.total_rows_synced as f64 / total as f64).clamp(0.0, 1.0) - } else { - 1.0 - }; - - let gauge_color = match state.phase { - GlobalPhase::TableCopy => Color::Green, - GlobalPhase::CdcStreaming => Color::Cyan, - GlobalPhase::Stopped => Color::DarkGray, - }; - - let sub = Layout::default() - .direction(Direction::Vertical) - .constraints([Constraint::Length(1), Constraint::Length(1), Constraint::Length(1)]) - .split(inner); - - frame.render_widget(Paragraph::new(vec![header]), sub[0]); - frame.render_widget(Paragraph::new(vec![stats_line]), sub[1]); - - let gauge_label = if state.phase == GlobalPhase::CdcStreaming { - if let Some((ref msg, _)) = state.status_message { - msg.clone() - } else { - "streaming".to_owned() - } - } else { - format!(" {:.0}%", ratio * 100.0) - }; - - frame.render_widget( - LineGauge::default() - .filled_style(Style::default().fg(gauge_color)) - .ratio(ratio) - .label(gauge_label) - .filled_symbol(symbols::line::THICK.horizontal) - .unfilled_symbol(symbols::line::NORMAL.horizontal), - sub[2], - ); -} - -fn render_middle_panel(frame: &mut ratatui::Frame, area: Rect, state: &DashboardState) { - let chunks = Layout::default() - .direction(Direction::Horizontal) - .constraints([Constraint::Percentage(75), Constraint::Percentage(25)]) - .split(area); - - render_table_list(frame, chunks[0], state); - render_table_detail(frame, chunks[1], state); -} - -fn render_table_list(frame: &mut ratatui::Frame, area: Rect, state: &DashboardState) { - if state.tables.is_empty() { - let block = Block::default().title(" Tables (j/k) ").borders(Borders::ALL); - let inner = block.inner(area); - frame.render_widget(block, area); - frame.render_widget(Paragraph::new(" Waiting for tables..."), inner); - return; - } - - let header = Row::new(vec![ - Cell::from(" Name"), - Cell::from("Phase"), - Cell::from("Rows"), - Cell::from("r/s"), - Cell::from(" "), - ]) - .style(Style::default().fg(Color::Cyan).add_modifier(Modifier::BOLD)) - .bottom_margin(0); - - let rows: Vec = state - .tables - .iter() - .enumerate() - .map(|(i, table)| { - let selected = i == state.selected_table; - let base_style = if selected { - Style::default().bg(Color::DarkGray).fg(Color::White).add_modifier(Modifier::BOLD) - } else { - Style::default().fg(Color::White) - }; - - let status_icon = if table.is_errored() { - Cell::from(Span::styled("✗", Style::default().fg(Color::Red))) - } else if table.is_ready() { - Cell::from(Span::styled("✓", Style::default().fg(Color::Green))) - } else if table.is_copying() { - Cell::from(Span::styled("↻", Style::default().fg(Color::Yellow))) - } else { - Cell::from("·") - }; - - Row::new(vec![ - Cell::from(format!( - "{}{}", - if selected { "▸" } else { " " }, - truncate_str(&table.destination_name, 24) - )), - Cell::from(Span::styled( - table.phase_label(), - Style::default().fg(phase_color(table.phase_label())), - )), - Cell::from(format_compact(table.rows_synced)), - Cell::from(format!("{:.0}", table.throughput)), - status_icon, - ]) - .style(base_style) - }) - .collect(); - - let widths = [ - Constraint::Min(20), - Constraint::Length(12), - Constraint::Length(10), - Constraint::Length(8), - Constraint::Length(2), - ]; - - let table_widget = Table::new(rows, widths) - .header(header) - .block(Block::default().title(" Tables (j/k) ").borders(Borders::ALL)); - - frame.render_widget(table_widget, area); -} - -fn render_table_detail(frame: &mut ratatui::Frame, area: Rect, state: &DashboardState) { - let title = match state.selected_table_info() { - Some(t) => format!(" {} ", truncate_str(&t.destination_name, 18)), - None => " Details ".to_owned(), - }; - - let block = Block::default().title(title).borders(Borders::ALL); - let inner = block.inner(area); - frame.render_widget(block, area); - - let Some(table) = state.selected_table_info() else { - frame.render_widget(Paragraph::new(" Select a table"), inner); - return; - }; - - let mut lines = vec![]; - - // Phase with visual indicator - let pc = phase_color(table.phase_label()); - let indicator = phase_indicator(table); - lines.push(Line::from(vec![ - Span::styled(table.phase_label(), Style::default().fg(pc).add_modifier(Modifier::BOLD)), - Span::raw(" "), - indicator, - ])); - - // Phase duration - if let Some(entered) = table.phase_entered_at { - let dur = entered.elapsed(); - lines.push(Line::from(Span::styled( - format!("{:.0}s in phase", dur.as_secs_f64()), - Style::default().fg(Color::DarkGray), - ))); - } - - lines.push(Line::from("")); - - // Offsets - lines.push(Line::from(vec![ - Span::styled("Local: ", Style::default().fg(Color::DarkGray)), - Span::raw(table.local_offset.as_deref().unwrap_or("—")), - ])); - lines.push(Line::from(vec![ - Span::styled("SF: ", Style::default().fg(Color::DarkGray)), - Span::raw(table.snowflake_offset.as_deref().unwrap_or("—")), - ])); - - lines.push(Line::from("")); - - // Global stats - lines.push(Line::from(vec![ - Span::styled("Copy: ", Style::default().fg(Color::DarkGray)), - Span::raw(format_num(state.total_copy_rows)), - ])); - lines.push(Line::from(vec![ - Span::styled("CDC: ", Style::default().fg(Color::DarkGray)), - Span::raw(format_num(state.total_cdc_events)), - ])); - - lines.push(Line::from("")); - - // Table OID - lines.push(Line::from(Span::styled( - format!("OID: {}", table.table_id), - Style::default().fg(Color::DarkGray), - ))); - - // Error details - if let Some(ref reason) = table.error_reason { - lines.push(Line::from("")); - lines.push(Line::from(Span::styled( - "Error:", - Style::default().fg(Color::Red).add_modifier(Modifier::BOLD), - ))); - for error_line in reason.lines() { - lines.push(Line::from(Span::styled(error_line, Style::default().fg(Color::Red)))); - } - if let Some(ref solution) = table.error_solution { - lines.push(Line::from("")); - lines.push(Line::from(Span::styled("Fix:", Style::default().fg(Color::Yellow)))); - for sol_line in solution.lines() { - lines.push(Line::from(Span::styled(sol_line, Style::default().fg(Color::Yellow)))); - } - } - lines.push(Line::from("")); - lines.push(Line::from(Span::styled("F2 to reset", Style::default().fg(Color::Magenta)))); - } - - frame.render_widget(Paragraph::new(lines).wrap(Wrap { trim: false }), inner); -} - -fn render_log_panel( - frame: &mut ratatui::Frame, - area: Rect, - log_buffer: &Arc>>, - scroll_offset: usize, - log_level: LogLevel, -) { - let title = format!(" Logs (J/K scroll) [{}] ", log_level.label()); - let block = Block::default().title(title).borders(Borders::ALL); - let inner = block.inner(area); - frame.render_widget(block, area); - - let logs: Vec = { - let buf = log_buffer.lock().unwrap(); - buf.iter() - .filter(|line| match log_level { - LogLevel::All => true, - LogLevel::WarnAndAbove => line.contains("[WARN]") || line.contains("[ERROR]"), - LogLevel::ErrorOnly => line.contains("[ERROR]"), - }) - .cloned() - .collect() - }; - - let visible_height = inner.height as usize; - let total = logs.len(); - let start = if total > visible_height { - let max_scroll = total - visible_height; - let from_bottom = scroll_offset.min(max_scroll); - max_scroll - from_bottom - } else { - 0 - }; - - let visible: Vec = logs - .iter() - .skip(start) - .take(visible_height) - .map(|line| { - let color = if line.contains("[ERROR]") { - Color::Red - } else if line.contains("[WARN]") { - Color::Yellow - } else { - Color::DarkGray - }; - Line::from(Span::styled(line.clone(), Style::default().fg(color))) - }) - .collect(); - - frame.render_widget(Paragraph::new(visible), inner); -} - -fn render_command_bar(frame: &mut ratatui::Frame, area: Rect, state: &DashboardState) { - let commands = [ - ("F1", "Help"), - ("F2", "Reset"), - ("F3", "Restart"), - ("F4", "LogLvl"), - ("F5", "Sync"), - ("F10", "Quit"), - ]; - - let mut spans = Vec::new(); - for (key, label) in &commands { - spans.push(Span::styled( - format!(" {key} "), - Style::default().fg(Color::Black).bg(Color::Cyan), - )); - spans.push(Span::raw(format!("{label} "))); - } - - if !state.pipeline_running { - spans.push(Span::raw(" ")); - spans.push(Span::styled( - " Pipeline stopped — F3 to restart ", - Style::default().fg(Color::Black).bg(Color::Yellow), - )); - } - - frame.render_widget(Paragraph::new(vec![Line::from(spans)]), area); -} - -fn render_help_overlay(frame: &mut ratatui::Frame) { - let area = frame.area(); - let width = 72.min(area.width.saturating_sub(4)); - let height = 38.min(area.height.saturating_sub(2)); - let x = (area.width.saturating_sub(width)) / 2; - let y = (area.height.saturating_sub(height)) / 2; - - let popup = Rect::new(x, y, width, height); - - frame.render_widget(Clear, popup); - let block = Block::default() - .title(" Help (press any key to close) ") - .borders(Borders::ALL) - .border_style(Style::default().fg(Color::Cyan)); - - let inner = block.inner(popup); - frame.render_widget(block, popup); - - let text: Vec = - commands::HELP_TEXT.lines().map(|l| Line::from(Span::raw(l.to_owned()))).collect(); - - frame.render_widget(Paragraph::new(text).wrap(Wrap { trim: false }), inner); -} - -fn phase_color(label: &str) -> Color { - match label { - "init" => Color::DarkGray, - "copying" | "data_sync" => Color::Green, - "copied" | "finished_copy" => Color::Blue, - "sync_wait" | "catchup" | "sync_done" => Color::Blue, - "ready" => Color::Cyan, - "errored" => Color::Red, - _ => Color::White, - } -} - -fn phase_indicator(table: &TableInfo) -> Span<'static> { - if table.is_copying() { - Span::styled("● copying", Style::default().fg(Color::Green)) - } else if table.is_ready() { - Span::styled("● CDC", Style::default().fg(Color::Cyan)) - } else if table.is_errored() { - Span::styled("● error", Style::default().fg(Color::Red)) - } else { - Span::styled("●", Style::default().fg(Color::DarkGray)) - } -} - -fn truncate_str(s: &str, max: usize) -> String { - if s.len() <= max { s.to_owned() } else { format!("{}…", &s[..max - 1]) } -} - -pub fn format_num(n: u64) -> String { - let s = n.to_string(); - let mut result = String::new(); - for (i, c) in s.chars().rev().enumerate() { - if i > 0 && i % 3 == 0 { - result.push(','); - } - result.push(c); - } - result.chars().rev().collect() -} - -fn format_compact(n: u64) -> String { - if n >= 1_000_000 { - format!("{:.1}M", n as f64 / 1_000_000.0) - } else if n >= 1_000 { - format!("{:.1}K", n as f64 / 1_000.0) - } else { - n.to_string() - } -} diff --git a/crates/etl-examples/src/bin/snowflake_loadgen.rs b/crates/etl-examples/src/bin/snowflake_loadgen.rs deleted file mode 100644 index d3131a6d9..000000000 --- a/crates/etl-examples/src/bin/snowflake_loadgen.rs +++ /dev/null @@ -1,457 +0,0 @@ -use std::{error::Error, time::Instant}; - -use clap::{Parser, Subcommand}; -use rand::Rng; -use tokio::signal; -use tokio_postgres::NoTls; - -#[derive(Clone, Copy)] -enum Op { - Insert, - Update, - Delete, -} - -const EVENT_TYPES: &[&str] = - &["page_view", "click", "purchase", "signup", "logout", "search", "download", "share"]; - -const IP_PREFIXES: &[&str] = &["192.168.", "10.0.", "172.16.", "203.0.113."]; - -#[derive(Debug, Parser)] -#[command( - name = "snowflake-loadgen", - version, - about = "Load generator for Snowflake ETL benchmarking" -)] -struct AppArgs { - #[command(subcommand)] - command: Command, -} - -#[derive(Debug, Subcommand)] -enum Command { - Seed(SeedArgs), - Generate(GenerateArgs), -} - -#[derive(Debug, Parser)] -struct SeedArgs { - #[arg(long)] - db_url: String, - #[arg(long, default_value = "1000000")] - rows: u64, -} - -#[derive(Debug, Parser)] -struct GenerateArgs { - #[arg(long)] - db_url: String, - #[arg(long, default_value = "100")] - rate: u64, - /// Insert/update/delete percentages, e.g. "40/40/20" - #[arg(long, default_value = "40/40/20")] - mix: String, - /// Duration to run, e.g. "300s" or "5m". Omit to run indefinitely. - #[arg(long)] - duration: Option, -} - -fn random_event_type(rng: &mut impl Rng) -> &'static str { - EVENT_TYPES[rng.random_range(0..EVENT_TYPES.len())] -} - -fn random_amount(rng: &mut impl Rng) -> f64 { - (rng.random_range(0u64..100_000) as f64) / 100.0 -} - -fn random_jsonb(rng: &mut impl Rng) -> String { - format!(r#"{{"source":"web","version":{}}}"#, rng.random_range(1..10)) -} - -fn random_ip(rng: &mut impl Rng) -> String { - let prefix = IP_PREFIXES[rng.random_range(0..IP_PREFIXES.len())]; - format!("{}{}.{}", prefix, rng.random_range(0..256), rng.random_range(1..255)) -} - -fn random_tags(rng: &mut impl Rng) -> String { - let all = &["rust", "web", "mobile", "api", "beta", "premium", "trial"]; - let count = rng.random_range(0..4usize); - let tags: Vec<&str> = (0..count).map(|_| all[rng.random_range(0..all.len())]).collect(); - format!("{{{}}}", tags.join(",")) -} - -fn parse_duration(s: &str) -> Result> { - if let Some(stripped) = s.strip_suffix('s') { - let secs: u64 = stripped.parse()?; - return Ok(std::time::Duration::from_secs(secs)); - } - if let Some(stripped) = s.strip_suffix('m') { - let mins: u64 = stripped.parse()?; - return Ok(std::time::Duration::from_secs(mins * 60)); - } - if let Some(stripped) = s.strip_suffix('h') { - let hours: u64 = stripped.parse()?; - return Ok(std::time::Duration::from_secs(hours * 3600)); - } - Err(format!("invalid duration '{s}': expected format like '300s', '5m', or '2h'").into()) -} - -fn parse_mix(s: &str) -> Result<(u32, u32, u32), Box> { - let parts: Vec<&str> = s.split('/').collect(); - if parts.len() != 3 { - return Err(format!("invalid mix '{s}': expected format like '40/40/20'").into()); - } - let insert: u32 = parts[0].parse()?; - let update: u32 = parts[1].parse()?; - let delete: u32 = parts[2].parse()?; - if insert + update + delete != 100 { - return Err( - format!("mix percentages must sum to 100, got {}", insert + update + delete).into() - ); - } - Ok((insert, update, delete)) -} - -async fn seed(client: &tokio_postgres::Client, rows: u64) -> Result<(), Box> { - let mut rng = rand::rng(); - - // Seed users (10% of rows, min 100) - let user_count = (rows / 10).max(100); - let user_batch = 5000u64; - let mut users_inserted = 0u64; - - eprintln!("Seeding {user_count} users..."); - while users_inserted < user_count { - let batch = user_batch.min(user_count - users_inserted); - let mut values = Vec::with_capacity(batch as usize); - for i in 0..batch { - let n = users_inserted + i; - let name = format!("User {n}"); - let email = format!("user{n}@example.com"); - values.push(format!("('{name}', '{email}', now())")); - } - let sql = format!( - "INSERT INTO bench.users (name, email, created_at) VALUES {}", - values.join(",") - ); - client.execute(&sql, &[]).await?; - users_inserted += batch; - } - eprintln!(" users done: {users_inserted}"); - - // Get user ID range - let row = client.query_one("SELECT MIN(id), MAX(id) FROM bench.users", &[]).await?; - let min_uid: i32 = row.get(0); - let max_uid: i32 = row.get(1); - - // Seed orders (same count as rows) - let order_batch = 5000u64; - let mut orders_inserted = 0u64; - let statuses = &["pending", "completed", "cancelled", "refunded"]; - - eprintln!("Seeding {rows} orders..."); - while orders_inserted < rows { - let batch = order_batch.min(rows - orders_inserted); - let mut values = Vec::with_capacity(batch as usize); - for _ in 0..batch { - let uid = rng.random_range(min_uid..=max_uid); - let total = random_amount(&mut rng); - let status = statuses[rng.random_range(0..statuses.len())]; - values.push(format!("({uid}, {total:.2}, '{status}', now())")); - } - let sql = format!( - "INSERT INTO bench.orders (user_id, total, status, created_at) VALUES {}", - values.join(",") - ); - client.execute(&sql, &[]).await?; - orders_inserted += batch; - if orders_inserted.is_multiple_of(10_000) || orders_inserted == rows { - eprintln!(" orders {orders_inserted}/{rows}"); - } - } - - // Seed events - let event_batch = 5000u64; - let mut events_inserted = 0u64; - - eprintln!("Seeding {rows} events..."); - while events_inserted < rows { - let batch = event_batch.min(rows - events_inserted); - let mut values = Vec::with_capacity(batch as usize); - for _ in 0..batch { - let uid = rng.random_range(min_uid..=max_uid); - let event_type = random_event_type(&mut rng); - let amount = random_amount(&mut rng); - let metadata = random_jsonb(&mut rng); - let tags = random_tags(&mut rng); - let score: f64 = rng.random_range(0..10000) as f64 / 100.0; - let ip = random_ip(&mut rng); - values.push(format!( - "({uid}, '{event_type}', {amount:.2}, '{metadata}', '{tags}', true, {score:.2}, \ - now(), now(), '{ip}', '')" - )); - } - let sql = format!( - "INSERT INTO bench.events (user_id, event_type, amount, metadata, tags, is_active, \ - score, created_at, updated_at, ip_address, notes) VALUES {}", - values.join(",") - ); - client.execute(&sql, &[]).await?; - events_inserted += batch; - if events_inserted.is_multiple_of(10_000) || events_inserted == rows { - eprintln!(" events {events_inserted}/{rows}"); - } - } - - Ok(()) -} - -async fn run_seed(args: SeedArgs) -> Result<(), Box> { - eprintln!("Connecting to Postgres..."); - let (client, connection) = tokio_postgres::connect(&args.db_url, NoTls).await?; - tokio::spawn(async move { - if let Err(e) = connection.await { - eprintln!("connection error: {e}"); - } - }); - - eprintln!("Creating bench schema (users, orders, events)..."); - client - .batch_execute( - "DROP SCHEMA IF EXISTS bench CASCADE; - DROP PUBLICATION IF EXISTS bench_pub; - CREATE SCHEMA bench; - CREATE TABLE bench.users ( - id SERIAL PRIMARY KEY, - name TEXT NOT NULL, - email TEXT NOT NULL, - created_at TIMESTAMPTZ DEFAULT now() - ); - CREATE TABLE bench.orders ( - id SERIAL PRIMARY KEY, - user_id INT REFERENCES bench.users(id), - total NUMERIC(10,2), - status TEXT, - created_at TIMESTAMPTZ DEFAULT now() - ); - CREATE TABLE bench.events ( - id BIGSERIAL PRIMARY KEY, - user_id INT REFERENCES bench.users(id), - event_type TEXT NOT NULL, - amount NUMERIC(12,2), - metadata JSONB, - tags TEXT[], - is_active BOOLEAN DEFAULT true, - score DOUBLE PRECISION, - created_at TIMESTAMPTZ DEFAULT now(), - updated_at TIMESTAMPTZ DEFAULT now(), - ip_address TEXT, - notes TEXT - )", - ) - .await?; - - eprintln!("Seeding {} rows...", args.rows); - let start = Instant::now(); - seed(&client, args.rows).await?; - - let elapsed = start.elapsed(); - - eprintln!("Creating publication bench_pub..."); - client - .execute("CREATE PUBLICATION bench_pub FOR TABLES IN SCHEMA bench", &[]) - .await - .or_else(|e| if e.to_string().contains("already exists") { Ok(0) } else { Err(e) })?; - - eprintln!( - "Done. Inserted {} rows in {:.1}s ({:.0} rows/sec).", - args.rows, - elapsed.as_secs_f64(), - args.rows as f64 / elapsed.as_secs_f64() - ); - - Ok(()) -} - -async fn run_generate(args: GenerateArgs) -> Result<(), Box> { - let (insert_pct, update_pct, _delete_pct) = parse_mix(&args.mix)?; - let deadline = args.duration.as_deref().map(parse_duration).transpose()?; - - eprintln!("Connecting to Postgres..."); - let (client, connection) = tokio_postgres::connect(&args.db_url, NoTls).await?; - tokio::spawn(async move { - if let Err(e) = connection.await { - eprintln!("connection error: {e}"); - } - }); - - // Cache ID ranges for random lookups (avoid ORDER BY random() scans) - let id_range = client.query_one("SELECT MIN(id), MAX(id) FROM bench.events", &[]).await?; - let mut min_id: i64 = id_range.get(0); - let mut max_id: i64 = id_range.get(1); - - let user_row = client.query_one("SELECT MIN(id), MAX(id) FROM bench.users", &[]).await?; - let min_uid: i32 = user_row.get(0); - let max_uid: i32 = user_row.get(1); - - // Batch size: target rate / batches_per_sec. At 5000 ops/sec with 50 - // batches/sec = 100 ops/batch. - let batch_size = (args.rate / 50).clamp(1, 500) as usize; - let batches_per_sec = (args.rate as f64 / batch_size as f64).ceil() as u64; - let interval_us = 1_000_000u64 / batches_per_sec.max(1); - let mut ticker = tokio::time::interval(std::time::Duration::from_micros(interval_us)); - ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Burst); - - let start = Instant::now(); - let mut total_ops: u64 = 0; - let mut last_report = Instant::now(); - let mut ops_since_report: u64 = 0; - let mut rng = rand::rng(); - - let ctrl_c = signal::ctrl_c(); - tokio::pin!(ctrl_c); - - eprintln!( - "Generating CDC changes at {} ops/sec (batch={}, mix: {}){}", - args.rate, - batch_size, - args.mix, - match deadline { - Some(d) => format!(", duration: {}s", d.as_secs()), - None => ", running until ctrl-c".to_owned(), - } - ); - - loop { - tokio::select! { - biased; - _ = &mut ctrl_c => { - eprintln!("\nReceived ctrl-c, stopping."); - break; - } - _ = ticker.tick() => {} - } - - if let Some(d) = deadline - && start.elapsed() >= d - { - eprintln!("Duration reached, stopping."); - break; - } - - // Build a batch of operations as a single SQL string - let mut stmts: Vec = Vec::with_capacity(batch_size); - let mut insert_values: Vec = Vec::new(); - let mut batch_ops = 0u64; - - for _ in 0..batch_size { - let roll: u32 = rng.random_range(0..100); - let op = if roll < insert_pct { - Op::Insert - } else if roll < insert_pct + update_pct { - Op::Update - } else { - Op::Delete - }; - - match op { - Op::Insert => { - let uid = rng.random_range(min_uid..=max_uid); - let event_type = random_event_type(&mut rng); - let amount = random_amount(&mut rng); - let metadata = random_jsonb(&mut rng); - let tags = random_tags(&mut rng); - let score: f64 = rng.random_range(0..10000) as f64 / 100.0; - let ip = random_ip(&mut rng); - insert_values.push(format!( - "({uid}, '{event_type}', {amount:.2}, '{metadata}', '{tags}', true, \ - {score:.2}, now(), now(), '{ip}', '')" - )); - } - Op::Update => { - let target_id = rng.random_range(min_id..=max_id); - let score: f64 = rng.random_range(0..10000) as f64 / 100.0; - let event_type = random_event_type(&mut rng); - stmts.push(format!( - "UPDATE bench.events SET score = {score:.2}, event_type = '{event_type}', \ - updated_at = now() WHERE id = {target_id}" - )); - } - Op::Delete => { - let target_id = rng.random_range(min_id..=max_id); - stmts.push(format!("DELETE FROM bench.events WHERE id = {target_id}")); - } - } - batch_ops += 1; - } - - // Flush accumulated inserts as one multi-row INSERT - if !insert_values.is_empty() { - let cols = "user_id, event_type, amount, metadata, tags, is_active, score, \ - created_at, updated_at, ip_address, notes"; - stmts.insert( - 0, - format!("INSERT INTO bench.events ({cols}) VALUES {}", insert_values.join(",")), - ); - } - - if stmts.is_empty() { - continue; - } - - let sql = stmts.join("; "); - match client.batch_execute(&sql).await { - Ok(()) => { - total_ops += batch_ops; - ops_since_report += batch_ops; - // Track max_id growth from inserts - max_id += insert_values.len() as i64; - } - Err(e) => { - eprintln!("batch error: {e}"); - } - } - - let since_report = last_report.elapsed(); - if since_report >= std::time::Duration::from_secs(5) { - let throughput = ops_since_report as f64 / since_report.as_secs_f64(); - eprintln!( - "[{:.0}s] {:.0} ops/sec (total: {})", - start.elapsed().as_secs_f64(), - throughput, - total_ops - ); - ops_since_report = 0; - last_report = Instant::now(); - - // Refresh ID range periodically to account for inserts/deletes - if let Ok(row) = - client.query_one("SELECT MIN(id), MAX(id) FROM bench.events", &[]).await - { - min_id = row.get(0); - max_id = row.get(1); - } - } - } - - let elapsed = start.elapsed(); - eprintln!( - "Finished. {} ops in {:.1}s ({:.0} ops/sec average).", - total_ops, - elapsed.as_secs_f64(), - total_ops as f64 / elapsed.as_secs_f64().max(0.001) - ); - - Ok(()) -} - -#[tokio::main] -async fn main() -> Result<(), Box> { - let args = AppArgs::parse(); - match args.command { - Command::Seed(seed_args) => run_seed(seed_args).await?, - Command::Generate(gen_args) => run_generate(gen_args).await?, - } - - Ok(()) -} diff --git a/crates/xtask/src/commands/example.rs b/crates/xtask/src/commands/example.rs new file mode 100644 index 000000000..a9c8b1d14 --- /dev/null +++ b/crates/xtask/src/commands/example.rs @@ -0,0 +1,139 @@ +use anyhow::{Result, bail}; +use clap::Args; +use xshell::{Shell, cmd}; + +const KNOWN_EXAMPLES: &[&str] = &["bigquery", "clickhouse", "ducklake", "snowflake"]; + +#[derive(Args)] +pub(crate) struct ExampleArgs { + /// Example name (bigquery, clickhouse, ducklake, snowflake). + name: String, + + /// Extra arguments passed through to the example binary. + #[arg(trailing_var_arg = true, allow_hyphen_values = true)] + args: Vec, +} + +impl ExampleArgs { + pub(crate) fn run(self) -> Result<()> { + let name = &self.name; + + if !KNOWN_EXAMPLES.contains(&name.as_str()) { + bail!("unknown example '{name}'. Available: {}", KNOWN_EXAMPLES.join(", ")); + } + + let db_host = env_or("TESTS_DATABASE_HOST", None); + let db_port = env_or("TESTS_DATABASE_PORT", None); + let db_username = env_or("TESTS_DATABASE_USERNAME", None); + let db_password = env_or("TESTS_DATABASE_PASSWORD", None); + + let missing: Vec<&str> = [ + db_host.is_none().then_some("TESTS_DATABASE_HOST"), + db_port.is_none().then_some("TESTS_DATABASE_PORT"), + db_username.is_none().then_some("TESTS_DATABASE_USERNAME"), + ] + .into_iter() + .flatten() + .collect(); + + if !missing.is_empty() && !has_flag(&self.args, "--db-host") { + bail!( + "missing environment variables: {}\n\nSet them in .env and run: source .env\nSee \ + .env.example for details.", + missing.join(", ") + ); + } + + let mut extra: Vec = Vec::new(); + + inject(&mut extra, "--db-host", &db_host, &self.args); + inject(&mut extra, "--db-port", &db_port, &self.args); + inject(&mut extra, "--db-username", &db_username, &self.args); + inject(&mut extra, "--db-password", &db_password, &self.args); + inject(&mut extra, "--db-name", &Some("etl_testdata".to_owned()), &self.args); + inject(&mut extra, "--publication", &Some("seed_pub".to_owned()), &self.args); + + match name.as_str() { + "bigquery" => { + inject_env(&mut extra, "--bq-project-id", "TESTS_BIGQUERY_PROJECT_ID", &self.args); + inject_env( + &mut extra, + "--bq-sa-key-file", + "TESTS_BIGQUERY_SA_KEY_PATH", + &self.args, + ); + } + "clickhouse" => { + inject_env(&mut extra, "--clickhouse-url", "TESTS_CLICKHOUSE_URL", &self.args); + inject_env(&mut extra, "--clickhouse-user", "TESTS_CLICKHOUSE_USER", &self.args); + inject_env( + &mut extra, + "--clickhouse-password", + "TESTS_CLICKHOUSE_PASSWORD", + &self.args, + ); + } + "snowflake" => { + inject_env( + &mut extra, + "--snowflake-account", + "TESTS_SNOWFLAKE_ACCOUNT", + &self.args, + ); + inject_env(&mut extra, "--snowflake-user", "TESTS_SNOWFLAKE_USER", &self.args); + inject_env( + &mut extra, + "--snowflake-private-key-path", + "TESTS_SNOWFLAKE_PRIVATE_KEY_PATH", + &self.args, + ); + inject_env( + &mut extra, + "--snowflake-private-key-passphrase", + "TESTS_SNOWFLAKE_PRIVATE_KEY_PASSPHRASE", + &self.args, + ); + inject_env( + &mut extra, + "--snowflake-database", + "TESTS_SNOWFLAKE_DATABASE", + &self.args, + ); + inject_env(&mut extra, "--snowflake-schema", "TESTS_SNOWFLAKE_SCHEMA", &self.args); + inject_env(&mut extra, "--snowflake-role", "TESTS_SNOWFLAKE_ROLE", &self.args); + } + _ => {} + } + + extra.extend(self.args.iter().cloned()); + + let sh = Shell::new()?; + let feature = name; + cmd!(sh, "cargo run --bin {name} -p etl-examples --features {feature} -- {extra...}") + .run()?; + + Ok(()) + } +} + +fn env_or(key: &str, default: Option<&str>) -> Option { + std::env::var(key).ok().or_else(|| default.map(String::from)) +} + +fn has_flag(args: &[String], flag: &str) -> bool { + args.iter().any(|a| a == flag || a.starts_with(&format!("{flag}="))) +} + +fn inject_env(extra: &mut Vec, flag: &str, env_var: &str, user_args: &[String]) { + inject(extra, flag, &env_or(env_var, None), user_args); +} + +fn inject(extra: &mut Vec, flag: &str, value: &Option, user_args: &[String]) { + if has_flag(user_args, flag) { + return; + } + if let Some(v) = value { + extra.push(flag.to_owned()); + extra.push(v.clone()); + } +} diff --git a/crates/xtask/src/commands/mod.rs b/crates/xtask/src/commands/mod.rs index 96f383d1f..16284450c 100644 --- a/crates/xtask/src/commands/mod.rs +++ b/crates/xtask/src/commands/mod.rs @@ -3,6 +3,7 @@ mod benchmark_compare; mod chaos; mod check; mod deploy_local; +mod example; mod fix; mod fmt; mod init; @@ -10,6 +11,7 @@ mod migrate; mod msrv; mod nextest; mod postgres; +mod seed; mod shared; mod test_clickhouse; mod vendor_duckdb; @@ -19,6 +21,7 @@ pub(crate) use benchmark_compare::BenchmarkCompareArgs; pub(crate) use chaos::ChaosArgs; pub(crate) use check::CheckArgs; pub(crate) use deploy_local::DeployLocalArgs; +pub(crate) use example::ExampleArgs; pub(crate) use fix::FixArgs; pub(crate) use fmt::FmtArgs; pub(crate) use init::InitArgs; @@ -26,5 +29,6 @@ pub(crate) use migrate::MigrateArgs; pub(crate) use msrv::MsrvArgs; pub(crate) use nextest::NextestArgs; pub(crate) use postgres::PostgresArgs; +pub(crate) use seed::SeedArgs; pub(crate) use test_clickhouse::TestClickhouseArgs; pub(crate) use vendor_duckdb::VendorDuckdbArgs; diff --git a/crates/xtask/src/commands/seed.rs b/crates/xtask/src/commands/seed.rs new file mode 100644 index 000000000..6720574a8 --- /dev/null +++ b/crates/xtask/src/commands/seed.rs @@ -0,0 +1,185 @@ +use anyhow::{Context, Result, bail}; +use clap::Args; +use xshell::{Shell, cmd}; + +#[derive(Args)] +pub(crate) struct SeedArgs { + /// PostgreSQL host. + #[arg(long, env = "TESTS_DATABASE_HOST", default_value = "localhost")] + host: String, + + /// PostgreSQL port. + #[arg(long, env = "TESTS_DATABASE_PORT", default_value = "5430")] + port: u16, + + /// PostgreSQL user. + #[arg(long, env = "TESTS_DATABASE_USERNAME", default_value = "postgres")] + user: String, + + /// PostgreSQL password. + #[arg(long, env = "TESTS_DATABASE_PASSWORD", default_value = "postgres")] + password: String, + + /// Database name to create and seed. + #[arg(long, default_value = "etl_testdata")] + database: String, + + /// Number of rows per table (users get rows/10, orders and events get + /// rows). + #[arg(long, default_value = "1000")] + rows: u64, + + /// Publication name to create. + #[arg(long, default_value = "seed_pub")] + publication: String, + + /// Drop and recreate the database if it already exists. + #[arg(long)] + force: bool, +} + +impl SeedArgs { + pub(crate) fn run(self) -> Result<()> { + let sh = Shell::new()?; + + if cmd!(sh, "which psql").quiet().run().is_err() { + bail!("psql is not installed or not in PATH"); + } + + sh.set_var("PGPASSWORD", &self.password); + + let host = &self.host; + let port = self.port.to_string(); + let user = &self.user; + let database = &self.database; + let publication = &self.publication; + let user_count = (self.rows / 10).max(10); + + // Check if database exists + let check_sql = format!("SELECT 1 FROM pg_database WHERE datname = '{database}'"); + let exists = + cmd!(sh, "psql -q -h {host} -p {port} -U {user} -d postgres -tA -c {check_sql}") + .quiet() + .read() + .unwrap_or_default(); + let db_exists = exists.trim() == "1"; + + if db_exists && !self.force { + bail!("database {database} already exists (use --force to drop and recreate)"); + } + + if db_exists { + println!("[seed] dropping database {database} (--force)..."); + let drop_sql = format!("DROP DATABASE IF EXISTS \"{database}\" WITH (FORCE)"); + cmd!(sh, "psql -q -h {host} -p {port} -U {user} -d postgres -c {drop_sql}") + .quiet() + .run() + .context("failed to drop database")?; + } + + println!("[seed] creating database {database}..."); + let create_db = format!("CREATE DATABASE \"{database}\""); + cmd!(sh, "psql -q -h {host} -p {port} -U {user} -d postgres -c {create_db}") + .quiet() + .run() + .context("failed to create database")?; + + let sql = format!( + r#" +CREATE TABLE users ( + id SERIAL PRIMARY KEY, + name TEXT NOT NULL, + email TEXT NOT NULL, + created_at TIMESTAMPTZ DEFAULT now() +); + +CREATE TABLE orders ( + id SERIAL PRIMARY KEY, + user_id INT REFERENCES users(id), + total NUMERIC(10,2), + status TEXT, + created_at TIMESTAMPTZ DEFAULT now() +); + +CREATE TABLE events ( + id BIGSERIAL PRIMARY KEY, + user_id INT REFERENCES users(id), + event_type TEXT NOT NULL, + amount NUMERIC(12,2), + metadata JSONB, + tags TEXT[], + is_active BOOLEAN DEFAULT true, + score DOUBLE PRECISION, + created_at TIMESTAMPTZ DEFAULT now(), + updated_at TIMESTAMPTZ DEFAULT now(), + ip_address TEXT, + notes TEXT +); + +INSERT INTO users (name, email, created_at) +SELECT + 'User ' || n, + 'user' || n || '@example.com', + now() - (random() * interval '365 days') +FROM generate_series(1, {user_count}) n; + +INSERT INTO orders (user_id, total, status, created_at) +SELECT + (random() * ({user_count} - 1) + 1)::int, + (random() * 999.99 + 0.01)::numeric(10,2), + (ARRAY['pending', 'completed', 'cancelled', 'refunded'])[1 + (random() * 3)::int], + now() - (random() * interval '365 days') +FROM generate_series(1, {rows}) n; + +INSERT INTO events (user_id, event_type, amount, metadata, tags, is_active, score, created_at, updated_at, ip_address, notes) +SELECT + (random() * ({user_count} - 1) + 1)::int, + (ARRAY['page_view', 'click', 'purchase', 'signup', 'logout', 'search', 'download', 'share'])[1 + (random() * 7)::int], + (random() * 999.99)::numeric(12,2), + json_build_object('source', 'web', 'version', (random() * 9 + 1)::int)::jsonb, + ARRAY[(ARRAY['rust', 'web', 'mobile', 'api', 'beta', 'premium', 'trial'])[1 + (random() * 6)::int]], + random() > 0.1, + (random() * 100)::numeric(5,2)::double precision, + now() - (random() * interval '365 days'), + now() - (random() * interval '30 days'), + '192.168.' || (random() * 255)::int || '.' || (random() * 254 + 1)::int, + '' +FROM generate_series(1, {rows}) n; + +CREATE PUBLICATION {publication} FOR TABLE users, orders, events; +"#, + rows = self.rows, + ); + + println!( + "[seed] creating tables and seeding {rows} rows ({user_count} users, {rows} orders, \ + {rows} events)...", + rows = self.rows + ); + + cmd!(sh, "psql -q -h {host} -p {port} -U {user} -d {database} -c {sql}") + .quiet() + .run() + .context("failed to seed database")?; + + let count_sql = "SELECT 'users: ' || count(*) FROM users UNION ALL SELECT 'orders: ' || \ + count(*) FROM orders UNION ALL SELECT 'events: ' || count(*) FROM events"; + let counts = + cmd!(sh, "psql -q -h {host} -p {port} -U {user} -d {database} -t -c {count_sql}") + .quiet() + .read() + .unwrap_or_default(); + + println!("[seed] done!"); + for line in counts.lines() { + let line = line.trim(); + if !line.is_empty() { + println!(" {line}"); + } + } + println!("[seed] publication: {publication}"); + println!("[seed] connect with: psql -h {host} -p {port} -U {user} -d {database}"); + + Ok(()) + } +} diff --git a/crates/xtask/src/main.rs b/crates/xtask/src/main.rs index 994ab7058..fe06188eb 100644 --- a/crates/xtask/src/main.rs +++ b/crates/xtask/src/main.rs @@ -3,9 +3,9 @@ mod commands; use anyhow::Result; use clap::{Parser, Subcommand}; use commands::{ - BenchmarkArgs, BenchmarkCompareArgs, ChaosArgs, CheckArgs, DeployLocalArgs, FixArgs, FmtArgs, - InitArgs, MigrateArgs, MsrvArgs, NextestArgs, PostgresArgs, TestClickhouseArgs, - VendorDuckdbArgs, + BenchmarkArgs, BenchmarkCompareArgs, ChaosArgs, CheckArgs, DeployLocalArgs, ExampleArgs, + FixArgs, FmtArgs, InitArgs, MigrateArgs, MsrvArgs, NextestArgs, PostgresArgs, SeedArgs, + TestClickhouseArgs, VendorDuckdbArgs, }; #[derive(Parser)] @@ -29,6 +29,8 @@ enum Command { /// Deploy the replicator to a local OrbStack Kubernetes cluster. #[command(name = "deploy-local")] DeployLocal(DeployLocalArgs), + /// Run a destination example (e.g. `cargo x example snowflake`). + Example(ExampleArgs), /// Auto-fix: clippy --fix, fmt, sort. Fix(FixArgs), /// Format code with nightly rustfmt. @@ -44,6 +46,9 @@ enum Command { Nextest(NextestArgs), /// Manage test Postgres clusters. Postgres(PostgresArgs), + /// Seed a Postgres database with test tables and data for destination + /// examples. + Seed(SeedArgs), /// Run ClickHouse integration tests with a local Docker setup. #[command(name = "test-clickhouse")] TestClickhouse(TestClickhouseArgs), @@ -61,6 +66,7 @@ async fn main() -> Result<()> { Command::Chaos(cmd) => cmd.run().await, Command::Check(cmd) => cmd.run(), Command::DeployLocal(cmd) => cmd.run(), + Command::Example(cmd) => cmd.run(), Command::Fix(cmd) => cmd.run(), Command::Fmt(cmd) => cmd.run(), Command::Init(cmd) => cmd.run(), @@ -68,6 +74,7 @@ async fn main() -> Result<()> { Command::Msrv(cmd) => cmd.run(), Command::Nextest(cmd) => cmd.run(), Command::Postgres(cmd) => cmd.run(), + Command::Seed(cmd) => cmd.run(), Command::TestClickhouse(cmd) => cmd.run(), Command::VendorDuckdb(cmd) => cmd.run(), }