diff --git a/etl-api/src/configs/pipeline.rs b/etl-api/src/configs/pipeline.rs index 61c35cbe6..9d1b105a9 100644 --- a/etl-api/src/configs/pipeline.rs +++ b/etl-api/src/configs/pipeline.rs @@ -1,4 +1,7 @@ -use etl_config::shared::{BatchConfig, PgConnectionConfig, PipelineConfig, TableSyncCopyConfig}; +use etl::config::ReplicationSlotPersistence; +use etl_config::shared::{ + BatchConfig, PgConnectionConfig, PipelineConfig, ReplicationSlotConfig, TableSyncCopyConfig, +}; use serde::{Deserialize, Serialize}; use utoipa::ToSchema; @@ -28,6 +31,54 @@ pub struct ApiBatchConfig { pub max_fill_ms: Option, } +#[derive(Debug, Clone, Copy, Serialize, Deserialize, ToSchema, Default)] +#[serde(rename_all = "snake_case")] +pub enum ApiReplicationSlotPersistence { + #[default] + Permanent, + Temporary, +} + +impl From for ApiReplicationSlotPersistence { + fn from(value: ReplicationSlotPersistence) -> Self { + match value { + ReplicationSlotPersistence::Permanent => ApiReplicationSlotPersistence::Permanent, + ReplicationSlotPersistence::Temporary => ApiReplicationSlotPersistence::Temporary, + } + } +} + +impl From for ReplicationSlotPersistence { + fn from(value: ApiReplicationSlotPersistence) -> Self { + match value { + ApiReplicationSlotPersistence::Permanent => ReplicationSlotPersistence::Permanent, + ApiReplicationSlotPersistence::Temporary => ReplicationSlotPersistence::Temporary, + } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize, ToSchema, Default)] +#[serde(rename_all = "snake_case")] +pub struct ApiReplicationSlotConfig { + pub persistence: ApiReplicationSlotPersistence, +} + +impl From for ApiReplicationSlotConfig { + fn from(value: ReplicationSlotConfig) -> Self { + ApiReplicationSlotConfig { + persistence: value.persistence.into(), + } + } +} + +impl From for ReplicationSlotConfig { + fn from(value: ApiReplicationSlotConfig) -> Self { + ReplicationSlotConfig { + persistence: value.persistence.into(), + } + } +} + #[derive(Debug, Clone, Serialize, Deserialize, ToSchema)] pub struct FullApiPipelineConfig { #[schema(example = "my_publication")] @@ -46,6 +97,8 @@ pub struct FullApiPipelineConfig { pub max_table_sync_workers: Option, #[serde(skip_serializing_if = "Option::is_none")] pub table_sync_copy: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub replication_slot: Option, pub log_level: Option, } @@ -61,6 +114,7 @@ impl From for FullApiPipelineConfig { table_error_retry_max_attempts: Some(value.table_error_retry_max_attempts), max_table_sync_workers: Some(value.max_table_sync_workers), table_sync_copy: Some(value.table_sync_copy), + replication_slot: Some(value.replication_slot.into()), log_level: value.log_level, } } @@ -97,6 +151,7 @@ pub struct PartialApiPipelineConfig { pub struct StoredPipelineConfig { pub publication_name: String, #[serde(default)] + pub replication_slot: ReplicationSlotConfig, pub batch: BatchConfig, #[serde(default = "default_table_error_retry_delay_ms")] pub table_error_retry_delay_ms: u64, @@ -118,6 +173,7 @@ impl StoredPipelineConfig { PipelineConfig { id: pipeline_id, publication_name: self.publication_name, + replication_slot: self.replication_slot, pg_connection: pg_connection_config, batch: self.batch, table_error_retry_delay_ms: self.table_error_retry_delay_ms, @@ -178,6 +234,7 @@ impl From for StoredPipelineConfig { Self { publication_name: value.publication_name, + replication_slot: value.replication_slot.unwrap_or_default().into(), batch, table_error_retry_delay_ms: value .table_error_retry_delay_ms @@ -203,6 +260,7 @@ mod tests { fn test_stored_pipeline_config_serialization() { let config = StoredPipelineConfig { publication_name: "test_publication".to_string(), + replication_slot: ReplicationSlotConfig::default(), batch: BatchConfig { max_size: 1000, max_fill_ms: 5000, @@ -242,6 +300,7 @@ mod tests { table_error_retry_max_attempts: None, max_table_sync_workers: None, table_sync_copy: None, + replication_slot: None, log_level: Some(LogLevel::Debug), }; @@ -260,6 +319,7 @@ mod tests { table_error_retry_max_attempts: None, max_table_sync_workers: None, table_sync_copy: None, + replication_slot: None, log_level: None, }; @@ -285,6 +345,7 @@ mod tests { fn test_partial_api_pipeline_config_merge() { let mut stored = StoredPipelineConfig { publication_name: "old_publication".to_string(), + replication_slot: ReplicationSlotConfig::default(), batch: BatchConfig { max_size: 500, max_fill_ms: 2000, diff --git a/etl-api/src/k8s/http.rs b/etl-api/src/k8s/http.rs index b3aa72cbb..eccc88c56 100644 --- a/etl-api/src/k8s/http.rs +++ b/etl-api/src/k8s/http.rs @@ -1062,8 +1062,8 @@ mod tests { use super::*; use etl_config::shared::{ - BatchConfig, DestinationConfig, PgConnectionConfig, PipelineConfig, ReplicatorConfig, - ReplicatorConfigWithoutSecrets, TableSyncCopyConfig, TlsConfig, + BatchConfig, DestinationConfig, PgConnectionConfig, PipelineConfig, ReplicationSlotConfig, + ReplicatorConfig, ReplicatorConfigWithoutSecrets, TableSyncCopyConfig, TlsConfig, }; use insta::assert_json_snapshot; @@ -1149,6 +1149,7 @@ mod tests { pipeline: PipelineConfig { id: 42, publication_name: "all-pub".to_string(), + replication_slot: ReplicationSlotConfig::default(), pg_connection: PgConnectionConfig { host: "localhost".to_string(), port: 5432, diff --git a/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_replicator_config_map_json.snap b/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_replicator_config_map_json.snap index a4c640a5c..216f653bc 100644 --- a/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_replicator_config_map_json.snap +++ b/etl-api/src/k8s/snapshots/etl_api__k8s__http__tests__create_replicator_config_map_json.snap @@ -6,7 +6,7 @@ expression: config_map_json "apiVersion": "v1", "data": { "base.json": "", - "prod.json": "{\"destination\":{\"big_query\":{\"project_id\":\"project-id\",\"dataset_id\":\"dataset-id\",\"max_concurrent_streams\":4}},\"pipeline\":{\"id\":42,\"publication_name\":\"all-pub\",\"pg_connection\":{\"host\":\"localhost\",\"port\":5432,\"name\":\"postgres\",\"username\":\"postgres\",\"tls\":{\"trusted_root_certs\":\"\",\"enabled\":false},\"keepalive\":null},\"batch\":{\"max_size\":10000,\"max_fill_ms\":1000},\"table_error_retry_delay_ms\":500,\"table_error_retry_max_attempts\":3,\"max_table_sync_workers\":4,\"table_sync_copy\":{\"type\":\"include_all_tables\"}}}" + "prod.json": "{\"destination\":{\"big_query\":{\"project_id\":\"project-id\",\"dataset_id\":\"dataset-id\",\"max_concurrent_streams\":4}},\"pipeline\":{\"id\":42,\"publication_name\":\"all-pub\",\"replication_slot\":{\"persistence\":\"permanent\"},\"pg_connection\":{\"host\":\"localhost\",\"port\":5432,\"name\":\"postgres\",\"username\":\"postgres\",\"tls\":{\"trusted_root_certs\":\"\",\"enabled\":false},\"keepalive\":null},\"batch\":{\"max_size\":10000,\"max_fill_ms\":1000},\"table_error_retry_delay_ms\":500,\"table_error_retry_max_attempts\":3,\"max_table_sync_workers\":4,\"table_sync_copy\":{\"type\":\"include_all_tables\"}}}" }, "kind": "ConfigMap", "metadata": { diff --git a/etl-api/tests/snapshots/destinations_pipelines__an_existing_bigquery_destination_and_pipeline_can_be_updated-2.snap b/etl-api/tests/snapshots/destinations_pipelines__an_existing_bigquery_destination_and_pipeline_can_be_updated-2.snap index af9192a34..15a16b8f5 100644 --- a/etl-api/tests/snapshots/destinations_pipelines__an_existing_bigquery_destination_and_pipeline_can_be_updated-2.snap +++ b/etl-api/tests/snapshots/destinations_pipelines__an_existing_bigquery_destination_and_pipeline_can_be_updated-2.snap @@ -26,6 +26,11 @@ FullApiPipelineConfig { table_sync_copy: Some( IncludeAllTables, ), + replication_slot: Some( + ApiReplicationSlotConfig { + persistence: Permanent, + }, + ), log_level: Some( Info, ), diff --git a/etl-api/tests/snapshots/destinations_pipelines__an_existing_iceberg_supabase_destination_and_pipeline_can_be_updated-2.snap b/etl-api/tests/snapshots/destinations_pipelines__an_existing_iceberg_supabase_destination_and_pipeline_can_be_updated-2.snap index af9192a34..15a16b8f5 100644 --- a/etl-api/tests/snapshots/destinations_pipelines__an_existing_iceberg_supabase_destination_and_pipeline_can_be_updated-2.snap +++ b/etl-api/tests/snapshots/destinations_pipelines__an_existing_iceberg_supabase_destination_and_pipeline_can_be_updated-2.snap @@ -26,6 +26,11 @@ FullApiPipelineConfig { table_sync_copy: Some( IncludeAllTables, ), + replication_slot: Some( + ApiReplicationSlotConfig { + persistence: Permanent, + }, + ), log_level: Some( Info, ), diff --git a/etl-api/tests/snapshots/destinations_pipelines__bigquery_destination_and_pipeline_can_be_created-2.snap b/etl-api/tests/snapshots/destinations_pipelines__bigquery_destination_and_pipeline_can_be_created-2.snap index 2a1b3dd50..1eb1ba5f0 100644 --- a/etl-api/tests/snapshots/destinations_pipelines__bigquery_destination_and_pipeline_can_be_created-2.snap +++ b/etl-api/tests/snapshots/destinations_pipelines__bigquery_destination_and_pipeline_can_be_created-2.snap @@ -26,6 +26,11 @@ FullApiPipelineConfig { table_sync_copy: Some( IncludeAllTables, ), + replication_slot: Some( + ApiReplicationSlotConfig { + persistence: Permanent, + }, + ), log_level: Some( Info, ), diff --git a/etl-api/tests/snapshots/destinations_pipelines__iceberg_supabase_destination_and_pipeline_can_be_created-2.snap b/etl-api/tests/snapshots/destinations_pipelines__iceberg_supabase_destination_and_pipeline_can_be_created-2.snap index 2a1b3dd50..1eb1ba5f0 100644 --- a/etl-api/tests/snapshots/destinations_pipelines__iceberg_supabase_destination_and_pipeline_can_be_created-2.snap +++ b/etl-api/tests/snapshots/destinations_pipelines__iceberg_supabase_destination_and_pipeline_can_be_created-2.snap @@ -26,6 +26,11 @@ FullApiPipelineConfig { table_sync_copy: Some( IncludeAllTables, ), + replication_slot: Some( + ApiReplicationSlotConfig { + persistence: Permanent, + }, + ), log_level: Some( Info, ), diff --git a/etl-api/tests/snapshots/pipelines__all_pipelines_can_be_read.snap b/etl-api/tests/snapshots/pipelines__all_pipelines_can_be_read.snap index 0b4e0d0e9..939329769 100644 --- a/etl-api/tests/snapshots/pipelines__all_pipelines_can_be_read.snap +++ b/etl-api/tests/snapshots/pipelines__all_pipelines_can_be_read.snap @@ -26,6 +26,11 @@ FullApiPipelineConfig { table_sync_copy: Some( IncludeAllTables, ), + replication_slot: Some( + ApiReplicationSlotConfig { + persistence: Permanent, + }, + ), log_level: Some( Info, ), diff --git a/etl-api/tests/snapshots/pipelines__an_existing_pipeline_can_be_read.snap b/etl-api/tests/snapshots/pipelines__an_existing_pipeline_can_be_read.snap index e17ca429e..e6aa8d0e0 100644 --- a/etl-api/tests/snapshots/pipelines__an_existing_pipeline_can_be_read.snap +++ b/etl-api/tests/snapshots/pipelines__an_existing_pipeline_can_be_read.snap @@ -26,6 +26,11 @@ FullApiPipelineConfig { table_sync_copy: Some( IncludeAllTables, ), + replication_slot: Some( + ApiReplicationSlotConfig { + persistence: Permanent, + }, + ), log_level: Some( Info, ), diff --git a/etl-api/tests/snapshots/pipelines__an_existing_pipeline_can_be_updated.snap b/etl-api/tests/snapshots/pipelines__an_existing_pipeline_can_be_updated.snap index efb0b6d51..f158ee7ae 100644 --- a/etl-api/tests/snapshots/pipelines__an_existing_pipeline_can_be_updated.snap +++ b/etl-api/tests/snapshots/pipelines__an_existing_pipeline_can_be_updated.snap @@ -26,6 +26,11 @@ FullApiPipelineConfig { table_sync_copy: Some( IncludeAllTables, ), + replication_slot: Some( + ApiReplicationSlotConfig { + persistence: Permanent, + }, + ), log_level: Some( Info, ), diff --git a/etl-api/tests/snapshots/pipelines__pipeline_config_can_be_updated-2.snap b/etl-api/tests/snapshots/pipelines__pipeline_config_can_be_updated-2.snap index 5755b07db..48456918a 100644 --- a/etl-api/tests/snapshots/pipelines__pipeline_config_can_be_updated-2.snap +++ b/etl-api/tests/snapshots/pipelines__pipeline_config_can_be_updated-2.snap @@ -26,5 +26,10 @@ FullApiPipelineConfig { table_sync_copy: Some( IncludeAllTables, ), + replication_slot: Some( + ApiReplicationSlotConfig { + persistence: Permanent, + }, + ), log_level: None, } diff --git a/etl-api/tests/snapshots/pipelines__pipeline_config_can_be_updated-3.snap b/etl-api/tests/snapshots/pipelines__pipeline_config_can_be_updated-3.snap index 683b8634d..9ee96a31e 100644 --- a/etl-api/tests/snapshots/pipelines__pipeline_config_can_be_updated-3.snap +++ b/etl-api/tests/snapshots/pipelines__pipeline_config_can_be_updated-3.snap @@ -26,5 +26,10 @@ FullApiPipelineConfig { table_sync_copy: Some( IncludeAllTables, ), + replication_slot: Some( + ApiReplicationSlotConfig { + persistence: Permanent, + }, + ), log_level: None, } diff --git a/etl-api/tests/snapshots/pipelines__pipeline_config_can_be_updated-4.snap b/etl-api/tests/snapshots/pipelines__pipeline_config_can_be_updated-4.snap index dc65b3fab..fea98dc89 100644 --- a/etl-api/tests/snapshots/pipelines__pipeline_config_can_be_updated-4.snap +++ b/etl-api/tests/snapshots/pipelines__pipeline_config_can_be_updated-4.snap @@ -26,5 +26,10 @@ FullApiPipelineConfig { table_sync_copy: Some( IncludeAllTables, ), + replication_slot: Some( + ApiReplicationSlotConfig { + persistence: Permanent, + }, + ), log_level: None, } diff --git a/etl-api/tests/snapshots/pipelines__pipeline_config_can_be_updated-5.snap b/etl-api/tests/snapshots/pipelines__pipeline_config_can_be_updated-5.snap index cca845e94..1c4ffaa49 100644 --- a/etl-api/tests/snapshots/pipelines__pipeline_config_can_be_updated-5.snap +++ b/etl-api/tests/snapshots/pipelines__pipeline_config_can_be_updated-5.snap @@ -26,6 +26,11 @@ FullApiPipelineConfig { table_sync_copy: Some( IncludeAllTables, ), + replication_slot: Some( + ApiReplicationSlotConfig { + persistence: Permanent, + }, + ), log_level: Some( Debug, ), diff --git a/etl-api/tests/snapshots/pipelines__pipeline_config_can_be_updated.snap b/etl-api/tests/snapshots/pipelines__pipeline_config_can_be_updated.snap index 0a272f067..0a4206de3 100644 --- a/etl-api/tests/snapshots/pipelines__pipeline_config_can_be_updated.snap +++ b/etl-api/tests/snapshots/pipelines__pipeline_config_can_be_updated.snap @@ -26,5 +26,10 @@ FullApiPipelineConfig { table_sync_copy: Some( IncludeAllTables, ), + replication_slot: Some( + ApiReplicationSlotConfig { + persistence: Permanent, + }, + ), log_level: None, } diff --git a/etl-api/tests/support/mocks.rs b/etl-api/tests/support/mocks.rs index aa2ed70b2..6ccb62182 100644 --- a/etl-api/tests/support/mocks.rs +++ b/etl-api/tests/support/mocks.rs @@ -224,7 +224,10 @@ pub mod tenants { /// Pipeline config helpers. pub mod pipelines { use super::*; - use etl_api::configs::{log::LogLevel, pipeline::ApiBatchConfig}; + use etl_api::configs::{ + log::LogLevel, + pipeline::{ApiBatchConfig, ApiReplicationSlotConfig}, + }; use etl_config::shared::TableSyncCopyConfig; /// Returns a default pipeline config. @@ -240,6 +243,7 @@ pub mod pipelines { max_table_sync_workers: Some(2), table_sync_copy: Some(TableSyncCopyConfig::IncludeAllTables), log_level: Some(LogLevel::Info), + replication_slot: Some(ApiReplicationSlotConfig::default()), } } @@ -256,6 +260,7 @@ pub mod pipelines { max_table_sync_workers: Some(4), table_sync_copy: Some(TableSyncCopyConfig::IncludeAllTables), log_level: Some(LogLevel::Info), + replication_slot: Some(ApiReplicationSlotConfig::default()), } } diff --git a/etl-api/tests/validators.rs b/etl-api/tests/validators.rs index a9a93a4d0..b85ac9da9 100644 --- a/etl-api/tests/validators.rs +++ b/etl-api/tests/validators.rs @@ -72,6 +72,7 @@ fn create_pipeline_config(publication_name: &str) -> FullApiPipelineConfig { table_error_retry_delay_ms: None, table_error_retry_max_attempts: None, table_sync_copy: None, + replication_slot: None, } } diff --git a/etl-benchmarks/benches/table_copies.rs b/etl-benchmarks/benches/table_copies.rs index 9b9308996..ba5a4e2c0 100644 --- a/etl-benchmarks/benches/table_copies.rs +++ b/etl-benchmarks/benches/table_copies.rs @@ -1,4 +1,5 @@ use clap::{Parser, Subcommand, ValueEnum}; +use etl::config::ReplicationSlotConfig; use etl::destination::Destination; use etl::error::EtlResult; use etl::pipeline::Pipeline; @@ -344,6 +345,7 @@ async fn start_pipeline(args: RunArgs) -> Result<(), Box> { let pipeline_config = PipelineConfig { id: 1, publication_name: args.publication_name, + replication_slot: ReplicationSlotConfig::default(), pg_connection: pg_connection_config, batch: BatchConfig { max_size: args.batch_max_size, diff --git a/etl-config/src/shared/pipeline.rs b/etl-config/src/shared/pipeline.rs index 8d7a3c1a5..82d8c2fbe 100644 --- a/etl-config/src/shared/pipeline.rs +++ b/etl-config/src/shared/pipeline.rs @@ -6,7 +6,7 @@ use crate::shared::{ PgConnectionConfig, PgConnectionConfigWithoutSecrets, ValidationError, batch::BatchConfig, }; -/// c copy should be performed.Selection rules for tables participating in replication. +/// Selection rules for tables participating in replication. /// /// Controls which tables are eligible for initial table copy and streaming. #[derive(Clone, Debug, Deserialize, Serialize, Eq, PartialEq)] @@ -60,6 +60,8 @@ pub struct PipelineConfig { pub id: u64, /// Name of the Postgres publication to use for logical replication. pub publication_name: String, + /// Configuration for the replication slot used + pub replication_slot: ReplicationSlotConfig, /// The connection configuration for the Postgres instance to which the pipeline connects for /// replication. pub pg_connection: PgConnectionConfig, @@ -138,6 +140,8 @@ pub struct PipelineConfigWithoutSecrets { pub id: u64, /// Name of the Postgres publication to use for logical replication. pub publication_name: String, + /// Configuration for the replication slot used + pub replication_slot: ReplicationSlotConfig, /// The connection configuration for the Postgres instance to which the pipeline connects for /// replication. pub pg_connection: PgConnectionConfigWithoutSecrets, @@ -163,6 +167,7 @@ impl From for PipelineConfigWithoutSecrets { PipelineConfigWithoutSecrets { id: value.id, publication_name: value.publication_name, + replication_slot: value.replication_slot, pg_connection: value.pg_connection.into(), batch: value.batch, table_error_retry_delay_ms: value.table_error_retry_delay_ms, @@ -173,6 +178,20 @@ impl From for PipelineConfigWithoutSecrets { } } +#[derive(Clone, Copy, Debug, Deserialize, Serialize, Default)] +#[serde(rename_all = "snake_case")] +pub struct ReplicationSlotConfig { + pub persistence: ReplicationSlotPersistence, +} + +#[derive(Clone, Copy, Debug, Deserialize, Serialize, Default)] +#[serde(rename_all = "snake_case")] +pub enum ReplicationSlotPersistence { + #[default] + Permanent, + Temporary, +} + #[cfg(test)] mod tests { use super::*; diff --git a/etl-examples/src/main.rs b/etl-examples/src/main.rs index 6344d5221..e9a9e3468 100644 --- a/etl-examples/src/main.rs +++ b/etl-examples/src/main.rs @@ -33,7 +33,8 @@ The pipeline will automatically: use clap::{Args, Parser}; use etl::config::{ - BatchConfig, PgConnectionConfig, PipelineConfig, TableSyncCopyConfig, TlsConfig, + BatchConfig, PgConnectionConfig, PipelineConfig, ReplicationSlotConfig, TableSyncCopyConfig, + TlsConfig, }; use etl::pipeline::Pipeline; use etl::store::both::memory::MemoryStore; @@ -188,6 +189,7 @@ async fn main_impl() -> Result<(), Box> { table_error_retry_delay_ms: 10000, table_error_retry_max_attempts: 5, max_table_sync_workers: args.bq_args.max_table_sync_workers, + replication_slot: ReplicationSlotConfig::default(), table_sync_copy: TableSyncCopyConfig::default(), }; diff --git a/etl/src/lib.rs b/etl/src/lib.rs index ac08f0be8..7f0afb11d 100644 --- a/etl/src/lib.rs +++ b/etl/src/lib.rs @@ -49,7 +49,7 @@ //! //! ```rust,no_run //! use etl::{ -//! config::{BatchConfig, PgConnectionConfig, PipelineConfig, TlsConfig, TableSyncCopyConfig}, +//! config::{BatchConfig, PgConnectionConfig, PipelineConfig, TlsConfig, TableSyncCopyConfig, ReplicationSlotConfig}, //! destination::memory::MemoryDestination, //! pipeline::Pipeline, //! store::both::memory::MemoryStore, @@ -76,6 +76,7 @@ //! let config = PipelineConfig { //! id: 1, //! publication_name: "my_publication".to_string(), +//! replication_slot: ReplicationSlotConfig::default(), //! pg_connection: pg_config, //! batch: BatchConfig { max_size: 1000, max_fill_ms: 5000 }, //! table_error_retry_delay_ms: 10000, diff --git a/etl/src/replication/client.rs b/etl/src/replication/client.rs index 952536f71..adeeca10d 100644 --- a/etl/src/replication/client.rs +++ b/etl/src/replication/client.rs @@ -1,7 +1,10 @@ use crate::error::{ErrorKind, EtlResult}; use crate::utils::tokio::MakeRustlsConnect; use crate::{bail, etl_error}; -use etl_config::shared::{ETL_REPLICATION_OPTIONS, IntoConnectOptions, PgConnectionConfig}; +use etl_config::shared::{ + ETL_REPLICATION_OPTIONS, IntoConnectOptions, PgConnectionConfig, ReplicationSlotConfig, + ReplicationSlotPersistence, +}; use etl_postgres::replication::extract_server_version; use etl_postgres::types::convert_type_oid_to_type; use etl_postgres::types::{ColumnSchema, TableId, TableName, TableSchema}; @@ -250,18 +253,23 @@ impl PgReplicationClient { pub async fn create_slot_with_transaction( &self, slot_name: &str, + config: ReplicationSlotConfig, ) -> EtlResult<(PgReplicationSlotTransaction, CreateSlotResult)> { // TODO: check if we want to consume the client and return it on commit to avoid any other // operations on a connection that has started a transaction. let transaction = PgReplicationSlotTransaction::new(self.clone()).await?; - let slot = self.create_slot_internal(slot_name, true).await?; + let slot = self.create_slot_internal(slot_name, true, config).await?; Ok((transaction, slot)) } /// Creates a new logical replication slot with the specified name and no snapshot. - pub async fn create_slot(&self, slot_name: &str) -> EtlResult { - self.create_slot_internal(slot_name, false).await + pub async fn create_slot( + &self, + slot_name: &str, + config: ReplicationSlotConfig, + ) -> EtlResult { + self.create_slot_internal(slot_name, false, config).await } /// Gets the slot by `slot_name`. @@ -306,7 +314,11 @@ impl PgReplicationClient { /// - A boolean indicating whether the slot was created (true) or already existed (false) /// - The slot result containing either the confirmed_flush_lsn (for existing slots) /// or the consistent_point (for newly created slots) - pub async fn get_or_create_slot(&self, slot_name: &str) -> EtlResult { + pub async fn get_or_create_slot( + &self, + slot_name: &str, + config: ReplicationSlotConfig, + ) -> EtlResult { match self.get_slot(slot_name).await { Ok(slot) => { info!(slot_name, "using existing replication slot"); @@ -316,7 +328,7 @@ impl PgReplicationClient { Err(err) if err.kind() == ErrorKind::ReplicationSlotNotFound => { info!(slot_name, "creating new replication slot"); - let create_result = self.create_slot_internal(slot_name, false).await?; + let create_result = self.create_slot_internal(slot_name, false, config).await?; Ok(GetOrCreateSlotResult::CreateSlot(create_result)) } @@ -577,6 +589,7 @@ impl PgReplicationClient { &self, slot_name: &str, use_snapshot: bool, + config: ReplicationSlotConfig, ) -> EtlResult { // Do not convert the query or the options to lowercase, since the lexer for // replication commands (repl_scanner.l) in Postgres code expects the commands @@ -587,9 +600,14 @@ impl PgReplicationClient { } else { "NOEXPORT_SNAPSHOT" }; + let temporary = match config.persistence { + ReplicationSlotPersistence::Temporary => " TEMPORARY ", + ReplicationSlotPersistence::Permanent => " ", + }; let query = format!( - r#"CREATE_REPLICATION_SLOT {} LOGICAL pgoutput {}"#, + r#"CREATE_REPLICATION_SLOT {}{}LOGICAL pgoutput {}"#, quote_identifier(slot_name), + temporary, snapshot_option ); match self.client.simple_query(&query).await { diff --git a/etl/src/replication/table_sync.rs b/etl/src/replication/table_sync.rs index 250f1b198..c20ac7d38 100644 --- a/etl/src/replication/table_sync.rs +++ b/etl/src/replication/table_sync.rs @@ -181,7 +181,7 @@ where // If a slot already exists at this point, we could delete it and try to recover, but it means // that the state was somehow reset without the slot being deleted, and we want to surface this. let (transaction, slot) = replication_client - .create_slot_with_transaction(&slot_name) + .create_slot_with_transaction(&slot_name, config.replication_slot) .await?; // We copy the table schema and write it both to the state store and destination. diff --git a/etl/src/test_utils/pipeline.rs b/etl/src/test_utils/pipeline.rs index 2e186920d..3a7348c05 100644 --- a/etl/src/test_utils/pipeline.rs +++ b/etl/src/test_utils/pipeline.rs @@ -1,4 +1,6 @@ -use etl_config::shared::{BatchConfig, PgConnectionConfig, PipelineConfig, TableSyncCopyConfig}; +use etl_config::shared::{ + BatchConfig, PgConnectionConfig, PipelineConfig, ReplicationSlotConfig, TableSyncCopyConfig, +}; use uuid::Uuid; use crate::destination::Destination; @@ -52,6 +54,8 @@ pub struct PipelineBuilder { max_table_sync_workers: u16, /// Table sync copy configuration. Uses default if not specified. table_sync_copy: Option, + /// Replication slot config. Uses default if not specified + replication_slot: Option, } impl PipelineBuilder @@ -94,6 +98,7 @@ where table_error_retry_max_attempts: 2, max_table_sync_workers: 1, table_sync_copy: None, + replication_slot: None, } } @@ -117,6 +122,16 @@ where self } + /// Sets custom replication slot configuration. + /// + /// # Arguments + /// + /// * `replication_slot` - Configuration for how replication slots behaviour + pub fn with_replication_slot_config(mut self, replication_slot: ReplicationSlotConfig) -> Self { + self.replication_slot = Some(replication_slot); + self + } + /// Sets custom retry configuration for table operations. /// /// # Arguments @@ -157,6 +172,7 @@ where table_error_retry_max_attempts: self.table_error_retry_max_attempts, max_table_sync_workers: self.max_table_sync_workers, table_sync_copy: self.table_sync_copy.unwrap_or_default(), + replication_slot: self.replication_slot.unwrap_or_default(), }; Pipeline::new(config, self.store, self.destination) diff --git a/etl/src/workers/apply.rs b/etl/src/workers/apply.rs index 81fdbb3f5..928a0ea55 100644 --- a/etl/src/workers/apply.rs +++ b/etl/src/workers/apply.rs @@ -1,4 +1,4 @@ -use etl_config::shared::PipelineConfig; +use etl_config::shared::{PipelineConfig, ReplicationSlotConfig}; use etl_postgres::replication::slots::EtlReplicationSlot; use etl_postgres::replication::worker::WorkerType; use etl_postgres::types::TableId; @@ -132,8 +132,13 @@ where publication_name = self.config.publication_name ); let apply_worker = async move { - let start_lsn = - get_start_lsn(self.pipeline_id, &self.replication_client, &self.store).await?; + let start_lsn = get_start_lsn( + self.pipeline_id, + &self.replication_client, + &self.store, + self.config.replication_slot, + ) + .await?; start_apply_loop( self.pipeline_id, @@ -183,12 +188,15 @@ async fn get_start_lsn( pipeline_id: PipelineId, replication_client: &PgReplicationClient, store: &S, + replication_slot: ReplicationSlotConfig, ) -> EtlResult { let slot_name: String = EtlReplicationSlot::for_apply_worker(pipeline_id).try_into()?; // 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 slot = replication_client + .get_or_create_slot(&slot_name, replication_slot) + .await?; // When creating a new apply worker slot, all tables must be in the `Init` state. If any table // is not in Init state, it means the table was synchronized based on another apply worker diff --git a/etl/tests/replication.rs b/etl/tests/replication.rs index 746913d8d..face6251e 100644 --- a/etl/tests/replication.rs +++ b/etl/tests/replication.rs @@ -2,8 +2,9 @@ use std::collections::HashSet; +use etl::config::{ReplicationSlotConfig, ReplicationSlotPersistence}; use etl::error::ErrorKind; -use etl::replication::client::PgReplicationClient; +use etl::replication::client::{CreateSlotResult, PgReplicationClient}; use etl::test_utils::database::{spawn_source_database, test_table_name}; use etl::test_utils::pipeline::test_slot_name; use etl::test_utils::table::assert_table_schema; @@ -93,7 +94,10 @@ async fn test_replication_client_creates_slot() { .unwrap(); let slot_name = test_slot_name("my_slot"); - let create_slot = client.create_slot(&slot_name).await.unwrap(); + let create_slot = client + .create_slot(&slot_name, ReplicationSlotConfig::default()) + .await + .unwrap(); assert!(!create_slot.consistent_point.to_string().is_empty()); let get_slot = client.get_slot(&slot_name).await.unwrap(); @@ -116,7 +120,10 @@ async fn test_create_and_delete_slot() { let slot_name = test_slot_name("my_slot"); // Create the slot and verify it exists - let create_slot = client.create_slot(&slot_name).await.unwrap(); + let create_slot = client + .create_slot(&slot_name, ReplicationSlotConfig::default()) + .await + .unwrap(); assert!(!create_slot.consistent_point.to_string().is_empty()); let get_slot = client.get_slot(&slot_name).await.unwrap(); @@ -156,13 +163,103 @@ async fn test_replication_client_doesnt_recreate_slot() { .unwrap(); let slot_name = test_slot_name("my_slot"); - assert!(client.create_slot(&slot_name).await.is_ok()); + assert!( + client + .create_slot(&slot_name, ReplicationSlotConfig::default()) + .await + .is_ok() + ); assert!(matches!( - client.create_slot(&slot_name).await, + client.create_slot(&slot_name, ReplicationSlotConfig::default()).await, Err(ref err) if err.kind() == ErrorKind::ReplicationSlotAlreadyExists )); } +#[tokio::test(flavor = "multi_thread")] +async fn test_replication_client_temporary_slot_dropped_on_disconnect() { + let database = spawn_source_database().await; + + let client = PgReplicationClient::connect(database.config.clone()) + .await + .unwrap(); + + let slot_name = test_slot_name("my_slot"); + assert!(matches!( + client + .create_slot( + &slot_name, + ReplicationSlotConfig { + persistence: ReplicationSlotPersistence::Temporary + } + ) + .await, + Ok(CreateSlotResult { + consistent_point: _ + }) + )); + + // Dropping the client should close the connection, causing our replication + // slot to be cleaned up + drop(client); + + // This means that we should be able to execute the exact same sequence again, and create the same replication slot again + let client = PgReplicationClient::connect(database.config.clone()) + .await + .unwrap(); + + assert!(matches!( + client + .create_slot( + &slot_name, + ReplicationSlotConfig { + persistence: ReplicationSlotPersistence::Temporary + } + ) + .await, + Ok(CreateSlotResult { + consistent_point: _ + }) + )); +} + +#[tokio::test(flavor = "multi_thread")] +async fn test_replication_client_permanent_slot_persisted_on_disconnect() { + let database = spawn_source_database().await; + + let client = PgReplicationClient::connect(database.config.clone()) + .await + .unwrap(); + + let slot_name = test_slot_name("my_slot"); + assert!(matches!( + client + .create_slot( + &slot_name, + ReplicationSlotConfig { + persistence: ReplicationSlotPersistence::Permanent + } + ) + .await, + Ok(CreateSlotResult { + consistent_point: _ + }) + )); + + // Dropping the client should close the connection. Since we use a persisted replication slot + // it should still be present when we start a fresh connection + drop(client); + + // This means that we should be able to execute the exact same sequence again, and the slot should still exist (i.e., we fail to create a new one) + let client = PgReplicationClient::connect(database.config.clone()) + .await + .unwrap(); + + assert!(matches!( + client.create_slot(&slot_name, ReplicationSlotConfig { persistence: ReplicationSlotPersistence::Permanent }).await, + Err(ref err) if err.kind() == ErrorKind::ReplicationSlotAlreadyExists, + )); +} + #[tokio::test(flavor = "multi_thread")] async fn test_table_schema_copy_is_consistent() { init_test_tracing(); @@ -187,7 +284,7 @@ async fn test_table_schema_copy_is_consistent() { // We create the slot when the database schema contains only 'table_1'. let (transaction, _) = client - .create_slot_with_transaction(&test_slot_name("my_slot")) + .create_slot_with_transaction(&test_slot_name("my_slot"), ReplicationSlotConfig::default()) .await .unwrap(); @@ -239,7 +336,7 @@ async fn test_table_schema_copy_across_multiple_connections() { // We create the slot when the database schema contains only 'table_1'. let (transaction, _) = first_client - .create_slot_with_transaction(&test_slot_name("my_slot")) + .create_slot_with_transaction(&test_slot_name("my_slot"), ReplicationSlotConfig::default()) .await .unwrap(); @@ -274,7 +371,7 @@ async fn test_table_schema_copy_across_multiple_connections() { // We create the slot when the database schema contains both 'table_1' and 'table_2'. let (transaction, _) = second_client - .create_slot_with_transaction(&test_slot_name("my_slot")) + .create_slot_with_transaction(&test_slot_name("my_slot"), ReplicationSlotConfig::default()) .await .unwrap(); @@ -338,7 +435,7 @@ async fn test_table_copy_stream_is_consistent() { // We create the slot when the database schema contains only 'table_1' data. let (transaction, _) = parent_client - .create_slot_with_transaction(&test_slot_name("my_slot")) + .create_slot_with_transaction(&test_slot_name("my_slot"), ReplicationSlotConfig::default()) .await .unwrap(); @@ -414,7 +511,7 @@ async fn test_table_copy_stream_respects_row_filter() { // We create the slot when the database schema contains only 'table_1' data. let (transaction, _) = parent_client - .create_slot_with_transaction(&test_slot_name("my_slot")) + .create_slot_with_transaction(&test_slot_name("my_slot"), ReplicationSlotConfig::default()) .await .unwrap(); @@ -501,7 +598,7 @@ async fn test_table_copy_stream_respects_column_filter() { // Create the slot when the database schema contains the test data. let (transaction, _) = parent_client - .create_slot_with_transaction(&test_slot_name("my_slot")) + .create_slot_with_transaction(&test_slot_name("my_slot"), ReplicationSlotConfig::default()) .await .unwrap(); @@ -589,7 +686,7 @@ async fn test_table_copy_stream_no_row_filter() { // We create the slot when the database schema contains only 'table_1' data. let (transaction, _) = parent_client - .create_slot_with_transaction(&test_slot_name("my_slot")) + .create_slot_with_transaction(&test_slot_name("my_slot"), ReplicationSlotConfig::default()) .await .unwrap(); @@ -716,7 +813,10 @@ async fn test_start_logical_replication() { // We create a slot which is going to replicate data before we insert the data. let slot_name = test_slot_name("my_slot"); - let slot = parent_client.create_slot(&slot_name).await.unwrap(); + let slot = parent_client + .create_slot(&slot_name, ReplicationSlotConfig::default()) + .await + .unwrap(); // We create a table with a publication and 10 entries. database diff --git a/scripts/init.sh b/scripts/init.sh index 42d58b381..fd8f6b738 100755 --- a/scripts/init.sh +++ b/scripts/init.sh @@ -7,9 +7,14 @@ if ! [ -x "$(command -v psql)" ]; then exit 1 fi -if ! [ -x "$(command -v docker-compose)" ]; then +# Pick docker compose (plugin) if available, else docker-compose (standalone) +if command -v docker >/dev/null 2>&1 && docker compose version >/dev/null 2>&1; then + DOCKER_COMPOSE=(docker compose) +elif command -v docker-compose >/dev/null 2>&1; then + DOCKER_COMPOSE=(docker-compose) +else echo >&2 "❌ Error: Docker Compose is not installed." - echo >&2 "Please install it using your system's package manager." + echo >&2 "Install either Docker Compose v2 (docker compose) or legacy docker-compose." exit 1 fi @@ -50,16 +55,16 @@ then # Pull latest images before starting services echo "📥 Pulling latest images..." - docker-compose -f ./scripts/docker-compose.yaml pull + "${DOCKER_COMPOSE[@]}" -f ./scripts/docker-compose.yaml pull # Start all services using docker-compose - docker-compose -f ./scripts/docker-compose.yaml up -d + "${DOCKER_COMPOSE[@]}" -f ./scripts/docker-compose.yaml up -d echo "✅ All services started" fi # Wait for Postgres to be ready echo "⏳ Waiting for Postgres to be ready..." -until docker-compose -f ./scripts/docker-compose.yaml exec -T source-postgres pg_isready -U postgres > /dev/null 2>&1; do +until "${DOCKER_COMPOSE[@]}" -f ./scripts/docker-compose.yaml exec -T source-postgres pg_isready -U postgres > /dev/null 2>&1; do echo "Waiting for Postgres..." sleep 1 done