From 6f9fe0eea4b6b8486c46afa62f9a8aa3afd977a8 Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Mon, 7 Sep 2026 11:32:21 +0200 Subject: [PATCH 1/3] feat(api): Improve VPA handling --- crates/etl-api/src/data/pipelines.rs | 68 ++++- crates/etl-api/src/k8s/http.rs | 263 ++++++++++++++---- crates/etl-api/src/routes/common.rs | 136 +++++---- crates/etl-api/tests/routes/pipelines.rs | 336 +++++++++++++++++++++++ crates/etl-postgres/src/publications.rs | 22 +- crates/etl/src/postgres/client/raw.rs | 10 +- 6 files changed, 712 insertions(+), 123 deletions(-) diff --git a/crates/etl-api/src/data/pipelines.rs b/crates/etl-api/src/data/pipelines.rs index b66a9d67d..1c4409ad6 100644 --- a/crates/etl-api/src/data/pipelines.rs +++ b/crates/etl-api/src/data/pipelines.rs @@ -1,10 +1,15 @@ use std::ops::DerefMut; +use etl::store::TableState; +use etl_config::shared::TableSyncCopyConfig; use etl_postgres::{ + publications::publication_table_ids_query, slots, store::{destination_table_metadata, health, schema, table_state}, }; -use sqlx::{FromRow, PgConnection, PgExecutor, PgPool, PgTransaction}; +use sqlx::{ + AssertSqlSafe, FromRow, PgConnection, PgExecutor, PgPool, PgTransaction, postgres::types::Oid, +}; use thiserror::Error; use crate::{ @@ -33,6 +38,67 @@ use crate::{ /// only one pipeline will use it. pub const MAX_PIPELINES_PER_TENANT: i64 = 1; +/// Returns published tables that would copy data when the pipeline starts. +/// +/// Uses the same publication expansion as the replicator, including implicit +/// schema membership and partition-root settings. Joining current pipeline +/// state in one query detects new tables without including removed tables or +/// reading table data. Tables without state start in [`TableState::Init`]; +/// existing states retain the replicator's restart semantics. +/// +/// The caller must supply the publication name and copy selection from the +/// current API configuration that will be materialized for this restart. This +/// function does not read the running Pod's configuration or verify that it +/// matches the API; copy eligibility assumes the replacement uses these values. +/// +/// This is a preflight observation, not a lock on publication membership or +/// replication progress. Either can change before the replicator starts. +pub(crate) async fn read_pipeline_tables_to_copy( + pool: &PgPool, + pipeline_id: i64, + publication_name: &str, + table_sync_copy: &TableSyncCopyConfig, +) -> Result, PipelineError> { + let query = format!( + r#" + with publication_tables as ({}) + select publication_table.oid, state.id, state.metadata + from publication_tables publication_table + left join etl.replication_state state + on state.table_id = publication_table.oid + and state.pipeline_id = $1 + and state.is_current = true + "#, + publication_table_ids_query(publication_name), + ); + let rows = + sqlx::query_as::<_, (Oid, Option, Option)>(AssertSqlSafe(query)) + .bind(pipeline_id) + .fetch_all(pool) + .await?; + + let mut tables_to_copy = Vec::new(); + for (table_id, state_id, metadata) in rows { + if !table_sync_copy.should_copy_table(table_id.0) { + continue; + } + + let state = if state_id.is_some() { + let metadata = metadata.ok_or(PipelineError::MissingTableState)?; + serde_json::from_value::(metadata) + .map_err(PipelineError::InvalidTableState)? + } else { + TableState::Init + }; + + if state.as_type().would_perform_table_sync() { + tables_to_copy.push(table_id.0); + } + } + + Ok(tables_to_copy) +} + #[derive(Debug, Clone)] pub struct Pipeline { pub id: i64, diff --git a/crates/etl-api/src/k8s/http.rs b/crates/etl-api/src/k8s/http.rs index c2a803d1a..cdda3bf37 100644 --- a/crates/etl-api/src/k8s/http.rs +++ b/crates/etl-api/src/k8s/http.rs @@ -15,7 +15,7 @@ use k8s_openapi::{ }; use kube::{ Client, - api::{Api, DeleteParams, ListParams, Patch, PatchParams}, + api::{Api, DeleteParams, ListParams, Patch, PatchParams, PostParams}, core::{ApiResource, DynamicObject, GroupVersionKind}, }; use serde_json::json; @@ -41,7 +41,7 @@ use crate::{ }, }; -/// Server-side apply field manager for resources owned by the API service. +/// Kubernetes field manager for resources owned by the API service. const FIELD_MANAGER: &str = "etl-api"; /// Secret name suffix for the BigQuery service account key. const BQ_SECRET_NAME_SUFFIX: &str = "bq-service-account-key"; @@ -821,33 +821,57 @@ impl K8sClient for HttpK8sClient { config.initial_update_mode }) .as_k8s_value(); - let existing_vertical_pod_autoscaler = - self.vertical_pod_autoscalers_api.get_opt(&name).await?; - let update_mode = existing_vertical_pod_autoscaler - .as_ref() - .and_then(vpa_update_mode) - .unwrap_or(initial_update_mode); - - debug!(vpa = %name, update_mode, "creating or updating vertical pod autoscaler"); + debug!(vpa = %name, "creating or updating vertical pod autoscaler"); let vertical_pod_autoscaler = create_replicator_vertical_pod_autoscaler_json( &self.k8s_config, resource_prefix, identity, &name, - update_mode, + None, workload_config.replicator_resource_override.as_ref(), )?; - // We are forcing the update since we are the field manager that should own the - // fields. If there is an override (likely during an incident or SREs - // intervention), we want to override their changes. - let pp = PatchParams::apply(FIELD_MANAGER).force(); - self.vertical_pod_autoscalers_api - .patch(&name, &pp, &Patch::Apply(vertical_pod_autoscaler)) - .await?; + // Merge leaves the mode untouched, including a concurrent controller promotion. + // Omitting it from server-side apply could delete a mode previously owned by + // us. + let pp = + PatchParams { field_manager: Some(FIELD_MANAGER.to_owned()), ..Default::default() }; + match self + .vertical_pod_autoscalers_api + .patch(&name, &pp, &Patch::Merge(&vertical_pod_autoscaler)) + .await + { + Ok(_) => return Ok(()), + Err(kube::Error::Api(error)) if error.code == 404 => {} + Err(error) => return Err(error.into()), + } - Ok(()) + let new_vertical_pod_autoscaler = create_replicator_vertical_pod_autoscaler_json( + &self.k8s_config, + resource_prefix, + identity, + &name, + Some(initial_update_mode), + workload_config.replicator_resource_override.as_ref(), + )?; + let post_params = + PostParams { field_manager: Some(FIELD_MANAGER.to_owned()), ..Default::default() }; + match self + .vertical_pod_autoscalers_api + .create(&post_params, &new_vertical_pod_autoscaler) + .await + { + Ok(_) => Ok(()), + Err(kube::Error::Api(error)) if error.code == 409 => { + // Another request created it first; preserve that object's mode too. + self.vertical_pod_autoscalers_api + .patch(&name, &pp, &Patch::Merge(&vertical_pod_autoscaler)) + .await?; + Ok(()) + } + Err(error) => Err(error.into()), + } } async fn delete_replicator_stateful_set(&self, resource_prefix: &str) -> Result<(), K8sError> { @@ -1815,18 +1839,26 @@ fn create_replicator_stateful_set_json( }) } +/// Builds a VPA creation document or a merge patch that preserves its mode. +/// +/// Supply the initial mode only for creation. An update must omit it entirely +/// so the API cannot overwrite an observation-period promotion. fn create_replicator_vertical_pod_autoscaler_json( k8s_config: &K8sConfig, prefix: &str, identity: &PipelineRuntimeIdentity, stateful_set_name: &str, - update_mode: &str, + update_mode: Option<&str>, pipeline_resource_override: Option<&PipelineReplicatorResourceOverrideConfig>, ) -> Result { let replicator_app_name = create_replicator_app_name(prefix); let replicator_container_name = create_replicator_container_name(prefix); let identity_labels = create_replicator_identity_labels(&replicator_app_name, identity); let policy = ReplicatorVpaResourcePolicy::resolve(k8s_config, pipeline_resource_override); + let update_policy = match update_mode { + Some(update_mode) => json!({ "updateMode": update_mode, "minReplicas": 1 }), + None => json!({ "minReplicas": 1 }), + }; serde_json::from_value(json!({ "apiVersion": format!("{VERTICAL_POD_AUTOSCALER_GROUP}/{VERTICAL_POD_AUTOSCALER_VERSION}"), @@ -1842,10 +1874,7 @@ fn create_replicator_vertical_pod_autoscaler_json( "kind": "StatefulSet", "name": stateful_set_name }, - "updatePolicy": { - "updateMode": update_mode, - "minReplicas": 1 - }, + "updatePolicy": update_policy, "resourcePolicy": { "containerPolicies": [ { @@ -1872,11 +1901,6 @@ fn create_replicator_vertical_pod_autoscaler_json( })) } -/// Reads the update mode that API reconciliation must preserve. -fn vpa_update_mode(vpa: &DynamicObject) -> Option<&str> { - vpa.data.pointer("/spec/updatePolicy/updateMode")?.as_str() -} - fn get_restarted_at_annotation_value() -> String { let now = Utc::now(); // We use nanoseconds to decrease the likelihood of generating the same @@ -2251,7 +2275,7 @@ mod tests { &prefix, &identity, &stateful_set_name, - ReplicatorResourceAutoscalingUpdateMode::Off.as_k8s_value(), + Some(ReplicatorResourceAutoscalingUpdateMode::Off.as_k8s_value()), None, ) .unwrap(), @@ -3051,7 +3075,7 @@ mod tests { "tenant-1-42", &identity, "tenant-1-42-replicator", - ReplicatorResourceAutoscalingUpdateMode::Off.as_k8s_value(), + Some(ReplicatorResourceAutoscalingUpdateMode::Off.as_k8s_value()), None, ) .unwrap(); @@ -3098,7 +3122,7 @@ mod tests { "tenant-1-42", &identity, "tenant-1-42-replicator", - ReplicatorResourceAutoscalingUpdateMode::Off.as_k8s_value(), + Some(ReplicatorResourceAutoscalingUpdateMode::Off.as_k8s_value()), None, ) .unwrap(); @@ -3126,7 +3150,7 @@ mod tests { "tenant-1-42", &identity, "tenant-1-42-replicator", - ReplicatorResourceAutoscalingUpdateMode::Off.as_k8s_value(), + Some(ReplicatorResourceAutoscalingUpdateMode::Off.as_k8s_value()), Some(&overrides), ) .unwrap(); @@ -3142,32 +3166,161 @@ mod tests { ); } - #[test] - fn replicator_vertical_pod_autoscaler_updates_preserve_live_mode() { + /// Runs VPA reconciliation against a scripted Kubernetes HTTP service. + async fn reconcile_vpa_with_responses( + statuses: Vec, + initial_mode: ReplicatorResourceAutoscalingUpdateMode, + ) -> (Result<(), K8sError>, Vec<(axum::http::Method, serde_json::Value)>) { + use std::{ + collections::VecDeque, + convert::Infallible, + sync::{Arc, Mutex}, + }; + + use axum::{ + body::Body, + http::{Method, Request, Response}, + }; + + let statuses = Arc::new(Mutex::new(VecDeque::from(statuses))); + let requests = Arc::new(Mutex::new(Vec::new())); + let service = tower::service_fn({ + let statuses = Arc::clone(&statuses); + let requests = Arc::clone(&requests); + move |request: Request| { + let statuses = Arc::clone(&statuses); + let requests = Arc::clone(&requests); + async move { + let method = request.method().clone(); + if method == Method::PATCH { + assert_eq!( + request.headers()["content-type"], + "application/merge-patch+json" + ); + } + let body: serde_json::Value = + serde_json::from_slice(&request.into_body().collect_bytes().await.unwrap()) + .unwrap(); + requests.lock().unwrap().push((method, body.clone())); + let status = statuses.lock().unwrap().pop_front().unwrap(); + let response = if status.is_success() { + // Model a live VPA that the controller already promoted. + let mut live = body; + live["spec"]["updatePolicy"]["updateMode"] = json!("InPlaceOrRecreate"); + live + } else { + json!({ + "apiVersion": "v1", "kind": "Status", "status": "Failure", + "message": "Simulated Kubernetes response", + "reason": status.canonical_reason().unwrap(), "code": status.as_u16() + }) + }; + Ok::<_, Infallible>( + Response::builder() + .status(status) + .header("content-type", "application/json") + .body(Body::from(serde_json::to_vec(&response).unwrap())) + .unwrap(), + ) + } + } + }); + let mut config = default_k8s_config(); + config.replicator_autoscaling.as_mut().unwrap().initial_update_mode = initial_mode; + let client = HttpK8sClient::new(Client::new(service, "etl-data-plane"), config).unwrap(); let identity = replicator_identity_with("tenant-1", PIPELINE_ID, REPLICATOR_ID); - let live: DynamicObject = serde_json::from_value(json!({ - "apiVersion": "autoscaling.k8s.io/v1", - "kind": "VerticalPodAutoscaler", - "metadata": {"name": "tenant-1-42-replicator"}, - "spec": {"updatePolicy": {"updateMode": "InPlaceOrRecreate"}} - })) - .unwrap(); - let autoscaler = create_replicator_vertical_pod_autoscaler_json( - &default_k8s_config(), - "tenant-1-42", - &identity, - "tenant-1-42-replicator", - vpa_update_mode(&live).unwrap(), - None, - ) - .unwrap(); - let autoscaler = serde_json::to_value(autoscaler).unwrap(); + let result = client + .create_or_update_replicator_vertical_pod_autoscaler( + "tenant-1-42", + &identity, + &ReplicatorWorkloadConfig { + replicator_image: "etl-replicator:test".to_owned(), + replicator_resource_override: Some(PipelineReplicatorResourceOverrideConfig { + cpu_request_millicores: Some(900), + memory_request_mib: None, + }), + destination_type: DestinationType::ClickHouse { + password_secret_required: false, + }, + ducklake_maintenance: None, + log_level: Default::default(), + }, + ) + .await; + assert!(statuses.lock().unwrap().is_empty()); + let requests = requests.lock().unwrap().clone(); + (result, requests) + } + #[tokio::test] + async fn replicator_vertical_pod_autoscaler_updates_preserve_live_mode() { + let (result, requests) = reconcile_vpa_with_responses( + vec![axum::http::StatusCode::OK], + ReplicatorResourceAutoscalingUpdateMode::Off, + ) + .await; + result.unwrap(); + assert_eq!(requests.len(), 1); + assert_eq!(requests[0].0, axum::http::Method::PATCH); + let patch = &requests[0].1; + assert!(patch.pointer("/spec/updatePolicy/updateMode").is_none()); + assert_eq!(patch.pointer("/spec/updatePolicy/minReplicas"), Some(&json!(1))); assert_eq!( - autoscaler.pointer("/spec/updatePolicy/updateMode"), - Some(&json!("InPlaceOrRecreate")) + patch.pointer("/spec/resourcePolicy/containerPolicies/0/minAllowed/cpu"), + Some(&json!("900m")), ); - assert_eq!(autoscaler.pointer("/spec/updatePolicy/minReplicas"), Some(&json!(1))); + assert_eq!(patch.pointer("/spec/targetRef/name"), Some(&json!("tenant-1-42-replicator")),); + assert_eq!( + patch.pointer("/metadata/labels/etl.supabase.com~1tenant-id"), + Some(&json!("tenant-1")) + ); + } + + #[tokio::test] + async fn replicator_vertical_pod_autoscaler_creation_uses_configured_initial_mode() { + let (result, requests) = reconcile_vpa_with_responses( + vec![axum::http::StatusCode::NOT_FOUND, axum::http::StatusCode::CREATED], + ReplicatorResourceAutoscalingUpdateMode::Initial, + ) + .await; + result.unwrap(); + assert_eq!(requests.len(), 2); + assert_eq!(requests[0].0, axum::http::Method::PATCH); + assert!(requests[0].1.pointer("/spec/updatePolicy/updateMode").is_none()); + assert_eq!(requests[1].0, axum::http::Method::POST); + assert_eq!(requests[1].1.pointer("/spec/updatePolicy/updateMode"), Some(&json!("Initial")),); + } + + #[tokio::test] + async fn replicator_vertical_pod_autoscaler_create_conflict_preserves_winner_mode() { + let (result, requests) = reconcile_vpa_with_responses( + vec![ + axum::http::StatusCode::NOT_FOUND, + axum::http::StatusCode::CONFLICT, + axum::http::StatusCode::OK, + ], + ReplicatorResourceAutoscalingUpdateMode::Off, + ) + .await; + result.unwrap(); + assert_eq!(requests.len(), 3); + assert_eq!(requests[1].0, axum::http::Method::POST); + assert_eq!(requests[1].1.pointer("/spec/updatePolicy/updateMode"), Some(&json!("Off"))); + assert_eq!(requests[0], requests[2]); + assert_eq!(requests[2].0, axum::http::Method::PATCH); + assert!(requests[2].1.pointer("/spec/updatePolicy/updateMode").is_none()); + } + + #[tokio::test] + async fn replicator_vertical_pod_autoscaler_update_failure_does_not_attempt_creation() { + let (result, requests) = reconcile_vpa_with_responses( + vec![axum::http::StatusCode::FORBIDDEN], + ReplicatorResourceAutoscalingUpdateMode::Off, + ) + .await; + assert!(result.is_err()); + assert_eq!(requests.len(), 1); + assert_eq!(requests[0].0, axum::http::Method::PATCH); } #[test] diff --git a/crates/etl-api/src/routes/common.rs b/crates/etl-api/src/routes/common.rs index 0b2360b6b..8b3e74bae 100644 --- a/crates/etl-api/src/routes/common.rs +++ b/crates/etl-api/src/routes/common.rs @@ -1,11 +1,14 @@ -use etl::store::TableState; -use etl_postgres::store::table_state; -use tracing::warn; +use tracing::{info, warn}; use crate::{ config::ApiConfig, - configs::{encryption::EncryptionKeyring, source::StoredSourceConfig}, - data::{pipelines::read_pipeline_components, source_database}, + configs::{ + encryption::EncryptionKeyring, pipeline::StoredPipelineConfig, source::StoredSourceConfig, + }, + data::{ + pipelines::{read_pipeline_components, read_pipeline_tables_to_copy}, + source_database, + }, k8s::{ K8sClient, SourceTlsConfig, core::{ @@ -17,12 +20,63 @@ use crate::{ validation::{self, ValidationContext, ValidationError, ValidationFailure}, }; +/// Checks whether published tables would copy data, preserving VPA on +/// uncertainty. +/// +/// Uses the current API pipeline and source configuration supplied by the +/// caller, which must also be used to materialize the replacement. These +/// values may differ from the running Pod's configuration after an API update. +/// Inspection uses the source connection's statement and lock timeouts. +async fn restart_would_copy_tables( + pipeline_id: i64, + pipeline_config: &StoredPipelineConfig, + source_id: i64, + source_config: &StoredSourceConfig, + source_tls_config: &SourceTlsConfig, +) -> bool { + let inspection = async { + let connection_config = + source_config.clone().into_connection_config(source_tls_config.get_tls_config()); + let source_pool = source_database::connect(&connection_config).await?; + read_pipeline_tables_to_copy( + &source_pool, + pipeline_id, + &pipeline_config.publication_name, + &pipeline_config.table_sync_copy, + ) + .await + }; + + match inspection.await { + Ok(tables_to_copy) => { + info!( + pipeline_id, + source_id, + table_count = tables_to_copy.len(), + "determined tables to copy on pipeline restart", + ); + !tables_to_copy.is_empty() + } + Err(error) => { + warn!( + pipeline_id, + source_id, + error = %error, + "failed to determine tables to copy on pipeline restart, preserving vertical pod autoscaler", + ); + false + } + } +} + /// Reconciles and restarts the running replicator for a pipeline. /// /// Update endpoints that can change source, destination, pipeline, image, or -/// runtime resource configuration should call this after persisting the new API -/// state. The helper materializes the latest Kubernetes resources and relies on -/// the StatefulSet materialization to change the pod template restart +/// runtime resource configuration should call this after writing the new API +/// state through the supplied connection, including within an uncommitted +/// transaction. The helper reads that state once and uses the same loaded +/// pipeline and source configuration for both copy preflight and Kubernetes +/// materialization. Updating the StatefulSet changes the pod template restart /// annotation. /// /// This forced recreation is part of the contract. The replicator loads its @@ -30,13 +84,14 @@ use crate::{ /// running pod must be restarted after config materialization in order to pick /// up those changes. /// -/// Before reconciliation, this best-effort checks durable pipeline state in the -/// database. If initial sync would repeat while copying existing tables, it -/// deletes the VPA so -/// reconciliation recreates it from the configured bounds and initial update -/// mode. The upstream recommender may retain in-memory usage aggregates after -/// the VPA is deleted. Source inspection failures preserve the existing VPA -/// and do not block restart. +/// Before reconciliation, this best-effort checks current publication +/// membership, durable pipeline state, and table-copy settings in the database. +/// If any table would copy data, including newly published tables, it deletes +/// the VPA so reconciliation recreates it in the configured initial update +/// mode. With `Off`, the replacement Pod starts at the configured startup +/// allocation and the VPA gets a fresh observation period. The upstream +/// recommender may retain usage history after deletion. Inspection failures and +/// timeouts preserve the existing VPA and do not block restart. /// /// Kubelet container restarts and Kubernetes-initiated Pod replacements do not /// call this helper or delete the VPA. A replacement Pod may therefore receive @@ -63,8 +118,14 @@ pub(crate) async fn restart_replicator_if_running( return Ok(false); } - if restart_would_perform_table_sync(pipeline_id, source.id, &source.config, source_tls_config) - .await + if restart_would_copy_tables( + pipeline_id, + &pipeline.config, + source.id, + &source.config, + source_tls_config, + ) + .await { let resource_prefix = create_k8s_object_prefix(tenant_id, replicator.id); k8s_client.delete_replicator_vertical_pod_autoscaler(&resource_prefix).await?; @@ -87,47 +148,6 @@ pub(crate) async fn restart_replicator_if_running( Ok(true) } -async fn restart_would_perform_table_sync( - pipeline_id: i64, - source_id: i64, - source_config: &StoredSourceConfig, - source_tls_config: &SourceTlsConfig, -) -> bool { - let result = async { - let connection_config = - source_config.clone().into_connection_config(source_tls_config.get_tls_config()); - let source_pool = source_database::connect(&connection_config).await?; - let state_rows = table_state::get_table_state_rows(&source_pool, pipeline_id).await?; - let mut would_perform_table_sync = false; - - for state_row in state_rows { - let Some(metadata) = state_row.metadata else { - return Err(PipelineError::MissingTableState); - }; - let state: TableState = - serde_json::from_value(metadata).map_err(PipelineError::InvalidTableState)?; - - would_perform_table_sync |= state.as_type().would_perform_table_sync(); - } - - Ok(would_perform_table_sync) - } - .await; - - match result { - Ok(will_repeat_sync) => will_repeat_sync, - Err(error) => { - warn!( - pipeline_id, - source_id, - error = %error, - "failed to determine whether pipeline restart will repeat table sync, preserving vertical pod autoscaler", - ); - false - } - } -} - /// Validates a source config against the trusted source profile, when enabled. pub async fn validate_source_config( source_config: StoredSourceConfig, diff --git a/crates/etl-api/tests/routes/pipelines.rs b/crates/etl-api/tests/routes/pipelines.rs index 39e5672c2..884588aa4 100644 --- a/crates/etl-api/tests/routes/pipelines.rs +++ b/crates/etl-api/tests/routes/pipelines.rs @@ -134,6 +134,14 @@ async fn setup_pipeline_with_source_db() -> (TestApp, String, i64, PgPool, PgCon // We run the migrations to create all the tables used by `etl`. run_etl_migrations_on_source_database(&source_db_config).await; + // Match the default pipeline publication and include test tables as they are + // created. + source_db_pool.execute("create schema test").await.unwrap(); + source_db_pool + .execute("create publication publication for tables in schema test") + .await + .unwrap(); + (app, tenant_id, pipeline_id, source_db_pool, source_db_config) } @@ -1584,6 +1592,334 @@ async fn restarting_pipeline_preserves_vpa_when_no_table_copy_will_repeat() { drop_pg_database(&source_db_config).await; } +#[tokio::test(flavor = "multi_thread")] +async fn restarting_pipeline_resets_vpa_for_a_new_explicitly_published_table() { + init_test_tracing(); + let (app, tenant_id, pipeline_id, source_db_pool, source_db_config) = + setup_pipeline_with_source_db().await; + create_tables_with_states( + &source_db_pool, + pipeline_id, + &[("test_users", "ready", r#"{"type": "ready"}"#)], + ) + .await; + source_db_pool.execute("drop publication publication").await.unwrap(); + source_db_pool + .execute("create publication publication for table test.test_users") + .await + .unwrap(); + + // An unrelated new table must not reset the VPA until it joins this + // publication. + create_test_table(&source_db_pool, "test_events").await; + let response = app.restart_pipeline(&tenant_id, pipeline_id).await; + assert_eq!(response.status(), StatusCode::ACCEPTED); + assert_eq!(app.k8s_state.vpa_delete_calls(), 0); + + source_db_pool + .execute("alter publication publication add table test.test_events") + .await + .unwrap(); + let response = app.restart_pipeline(&tenant_id, pipeline_id).await; + assert_eq!(response.status(), StatusCode::ACCEPTED); + assert_eq!(app.k8s_state.vpa_delete_calls(), 1); + + drop_pg_database(&source_db_config).await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn restarting_pipeline_resets_vpa_for_a_new_table_in_a_published_schema() { + init_test_tracing(); + let (app, tenant_id, pipeline_id, source_db_pool, source_db_config) = + setup_pipeline_with_source_db().await; + create_tables_with_states( + &source_db_pool, + pipeline_id, + &[("test_users", "ready", r#"{"type": "ready"}"#)], + ) + .await; + create_test_table(&source_db_pool, "test_events").await; + + let response = app.restart_pipeline(&tenant_id, pipeline_id).await; + + assert_eq!(response.status(), StatusCode::ACCEPTED); + assert_eq!(app.k8s_state.vpa_delete_calls(), 1); + + drop_pg_database(&source_db_config).await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn restarting_pipeline_resets_vpa_for_a_new_table_in_an_all_table_publication() { + init_test_tracing(); + let (app, tenant_id, pipeline_id, source_db_pool, source_db_config) = + setup_pipeline_with_source_db().await; + source_db_pool.execute("drop publication publication").await.unwrap(); + source_db_pool.execute("create publication publication for all tables").await.unwrap(); + sqlx::query( + "insert into etl.replication_state (pipeline_id, table_id, state, metadata) select $1, \ + relid, 'ready', '{\"type\":\"ready\"}'::jsonb from \ + pg_get_publication_tables('publication')", + ) + .bind(pipeline_id) + .execute(&source_db_pool) + .await + .unwrap(); + + let response = app.restart_pipeline(&tenant_id, pipeline_id).await; + assert_eq!(response.status(), StatusCode::ACCEPTED); + assert_eq!(app.k8s_state.vpa_delete_calls(), 0); + + create_test_table(&source_db_pool, "test_users").await; + let response = app.restart_pipeline(&tenant_id, pipeline_id).await; + assert_eq!(response.status(), StatusCode::ACCEPTED); + assert_eq!(app.k8s_state.vpa_delete_calls(), 1); + + drop_pg_database(&source_db_config).await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn restarting_pipeline_resets_vpa_for_published_tables_without_any_state() { + init_test_tracing(); + let (app, tenant_id, pipeline_id, source_db_pool, source_db_config) = + setup_pipeline_with_source_db().await; + create_test_table(&source_db_pool, "test_users").await; + + let response = app.restart_pipeline(&tenant_id, pipeline_id).await; + + assert_eq!(response.status(), StatusCode::ACCEPTED); + assert_eq!(app.k8s_state.vpa_delete_calls(), 1); + + drop_pg_database(&source_db_config).await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn restarting_pipeline_preserves_vpa_for_copy_states_outside_the_publication() { + init_test_tracing(); + let (app, tenant_id, pipeline_id, source_db_pool, source_db_config) = + setup_pipeline_with_source_db().await; + create_tables_with_states( + &source_db_pool, + pipeline_id, + &[ + ("test_users", "ready", r#"{"type": "ready"}"#), + ("test_events", "data_sync", r#"{"type": "data_sync"}"#), + ], + ) + .await; + source_db_pool.execute("drop publication publication").await.unwrap(); + source_db_pool + .execute("create publication publication for table test.test_users") + .await + .unwrap(); + + let response = app.restart_pipeline(&tenant_id, pipeline_id).await; + + assert_eq!(response.status(), StatusCode::ACCEPTED); + assert_eq!(app.k8s_state.vpa_delete_calls(), 0); + + drop_pg_database(&source_db_config).await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn restarting_pipeline_uses_only_current_state_for_its_pipeline() { + init_test_tracing(); + let (app, tenant_id, pipeline_id, source_db_pool, source_db_config) = + setup_pipeline_with_source_db().await; + create_table_with_state_chain( + &source_db_pool, + pipeline_id, + "test_users", + &[("data_sync", r#"{"type": "data_sync"}"#), ("ready", r#"{"type": "ready"}"#)], + ) + .await; + + let response = app.restart_pipeline(&tenant_id, pipeline_id).await; + assert_eq!(response.status(), StatusCode::ACCEPTED); + assert_eq!(app.k8s_state.vpa_delete_calls(), 0); + + // A different pipeline's completed copy cannot suppress this pipeline's copy. + create_tables_with_states( + &source_db_pool, + pipeline_id + 1, + &[("test_events", "ready", r#"{"type": "ready"}"#)], + ) + .await; + let response = app.restart_pipeline(&tenant_id, pipeline_id).await; + assert_eq!(response.status(), StatusCode::ACCEPTED); + assert_eq!(app.k8s_state.vpa_delete_calls(), 1); + + drop_pg_database(&source_db_config).await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn restarting_pipeline_resets_vpa_when_a_published_table_is_recreated() { + init_test_tracing(); + let (app, tenant_id, pipeline_id, source_db_pool, source_db_config) = + setup_pipeline_with_source_db().await; + let tables = create_tables_with_states( + &source_db_pool, + pipeline_id, + &[("test_users", "ready", r#"{"type": "ready"}"#)], + ) + .await; + source_db_pool.execute("drop table test.test_users").await.unwrap(); + let new_table_id = create_test_table(&source_db_pool, "test_users").await; + assert_ne!(tables[0].0, new_table_id); + + let response = app.restart_pipeline(&tenant_id, pipeline_id).await; + + assert_eq!(response.status(), StatusCode::ACCEPTED); + assert_eq!(app.k8s_state.vpa_delete_calls(), 1); + + drop_pg_database(&source_db_config).await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn restarting_pipeline_respects_partition_root_publication_settings() { + init_test_tracing(); + let (app, tenant_id, pipeline_id, source_db_pool, source_db_config) = + setup_pipeline_with_source_db().await; + source_db_pool + .execute("create table test.test_events (id int primary key) partition by range (id)") + .await + .unwrap(); + source_db_pool + .execute( + "create table test.test_events_1 partition of test.test_events for values from (0) to \ + (100)", + ) + .await + .unwrap(); + create_tables_with_states( + &source_db_pool, + pipeline_id, + &[("test_events", "ready", r#"{"type": "ready"}"#)], + ) + .await; + source_db_pool + .execute("alter publication publication set (publish_via_partition_root = true)") + .await + .unwrap(); + + let response = app.restart_pipeline(&tenant_id, pipeline_id).await; + assert_eq!(response.status(), StatusCode::ACCEPTED); + assert_eq!(app.k8s_state.vpa_delete_calls(), 0); + + // Publishing leaves exposes a table without state even though its root is + // ready. + source_db_pool + .execute("alter publication publication set (publish_via_partition_root = false)") + .await + .unwrap(); + let response = app.restart_pipeline(&tenant_id, pipeline_id).await; + assert_eq!(response.status(), StatusCode::ACCEPTED); + assert_eq!(app.k8s_state.vpa_delete_calls(), 1); + + drop_pg_database(&source_db_config).await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn restarting_pipeline_respects_table_copy_selection() { + init_test_tracing(); + let (app, tenant_id, pipeline_id, source_db_pool, source_db_config) = + setup_pipeline_with_source_db().await; + let ready_tables = create_tables_with_states( + &source_db_pool, + pipeline_id, + &[("test_users", "ready", r#"{"type": "ready"}"#)], + ) + .await; + let new_table_id = create_test_table(&source_db_pool, "test_events").await; + let pipeline = app.read_pipeline(&tenant_id, pipeline_id).await; + let pipeline: ReadPipelineResponse = pipeline.json().await.unwrap(); + + for (table_sync_copy, copies_new_table) in [ + (TableSyncCopyConfig::SkipAllTables, false), + (TableSyncCopyConfig::IncludeTables { table_ids: vec![ready_tables[0].0.0] }, false), + (TableSyncCopyConfig::SkipTables { table_ids: vec![new_table_id.0] }, false), + (TableSyncCopyConfig::IncludeTables { table_ids: vec![new_table_id.0] }, true), + (TableSyncCopyConfig::SkipTables { table_ids: vec![ready_tables[0].0.0] }, true), + (TableSyncCopyConfig::IncludeAllTables, true), + ] { + let response = app + .update_pipeline( + &tenant_id, + pipeline_id, + &UpdatePipelineRequest { + source_id: pipeline.source_id, + destination_id: pipeline.destination_id, + config: UpdateApiPipelineConfig { + table_sync_copy: UpdateField::Set(table_sync_copy), + ..Default::default() + }, + }, + ) + .await; + assert_eq!(response.status(), StatusCode::OK); + let delete_calls_before = app.k8s_state.vpa_delete_calls(); + + let response = app.restart_pipeline(&tenant_id, pipeline_id).await; + + assert_eq!(response.status(), StatusCode::ACCEPTED); + assert_eq!( + app.k8s_state.vpa_delete_calls() - delete_calls_before, + usize::from(copies_new_table), + ); + } + + drop_pg_database(&source_db_config).await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn restarting_pipeline_preserves_vpa_when_publication_inspection_fails() { + init_test_tracing(); + let (app, tenant_id, pipeline_id, source_db_pool, source_db_config) = + setup_pipeline_with_source_db().await; + create_tables_with_states( + &source_db_pool, + pipeline_id, + &[("test_users", "data_sync", r#"{"type": "data_sync"}"#)], + ) + .await; + source_db_pool.execute("drop publication publication").await.unwrap(); + + let response = app.restart_pipeline(&tenant_id, pipeline_id).await; + + assert_eq!(response.status(), StatusCode::ACCEPTED); + assert_eq!(app.k8s_state.vpa_delete_calls(), 0); + + drop_pg_database(&source_db_config).await; +} + +#[tokio::test(flavor = "multi_thread")] +async fn restarting_pipeline_preserves_vpa_when_source_lock_times_out() { + init_test_tracing(); + let (app, tenant_id, pipeline_id, source_db_pool, source_db_config) = + setup_pipeline_with_source_db().await; + create_test_table(&source_db_pool, "test_users").await; + + // Hold the lock until the source connection's lock timeout cancels inspection. + let mut transaction = source_db_pool.begin().await.unwrap(); + sqlx::query("lock table etl.replication_state in access exclusive mode") + .execute(&mut *transaction) + .await + .unwrap(); + let create_calls_before = app.k8s_state.create_calls(); + let response = tokio::time::timeout( + std::time::Duration::from_secs(15), + app.restart_pipeline(&tenant_id, pipeline_id), + ) + .await + .unwrap(); + + assert_eq!(response.status(), StatusCode::ACCEPTED); + assert!(app.k8s_state.create_calls() > create_calls_before); + assert_eq!(app.k8s_state.vpa_delete_calls(), 0); + + transaction.rollback().await.unwrap(); + drop_pg_database(&source_db_config).await; +} + #[tokio::test(flavor = "multi_thread")] async fn restarting_pipeline_preserves_vpa_when_source_inspection_fails() { init_test_tracing(); diff --git a/crates/etl-postgres/src/publications.rs b/crates/etl-postgres/src/publications.rs index 41c69a7de..699374561 100644 --- a/crates/etl-postgres/src/publications.rs +++ b/crates/etl-postgres/src/publications.rs @@ -1,8 +1,28 @@ -//! Naming helpers for ETL-created PostgreSQL publications. +//! PostgreSQL publication naming and membership queries. + +use pg_escape::quote_literal; /// Prefix used for publications created by Supabase ETL. pub const ETL_PUBLICATION_PREFIX: &str = "supabase_etl_publication"; +/// Builds the table-ID query shared by replication startup and API preflight. +/// +/// PostgreSQL expands explicit tables, schema publications, and all-table +/// publications according to `publish_via_partition_root`. Deduplication +/// returns each logical table once. The publication name is escaped as a SQL +/// literal so this query also works on replication connections using the simple +/// protocol. +pub fn publication_table_ids_query(publication_name: &str) -> String { + format!( + r#" + select distinct gpt.relid::oid as oid + from pg_catalog.pg_get_publication_tables({}) gpt + order by oid + "#, + quote_literal(publication_name) + ) +} + /// Returns the deterministic ETL publication name for a pipeline. pub fn etl_publication_name(pipeline_id: i64) -> String { format!("{ETL_PUBLICATION_PREFIX}_{pipeline_id}") diff --git a/crates/etl/src/postgres/client/raw.rs b/crates/etl/src/postgres/client/raw.rs index efb649018..defc43b94 100644 --- a/crates/etl/src/postgres/client/raw.rs +++ b/crates/etl/src/postgres/client/raw.rs @@ -2,6 +2,7 @@ use std::{fmt, num::NonZeroI32, sync::Arc, time::Duration}; use etl_postgres::{ application_name::{apply_worker_application_name, table_sync_worker_application_name}, + publications::publication_table_ids_query, source::extract_server_version, tokio::tls::MakeRustlsConnect, version::POSTGRES_17, @@ -644,14 +645,7 @@ impl PgReplicationClient { &self, publication_name: &str, ) -> EtlResult> { - let query = format!( - r#" - select distinct gpt.relid::oid as oid - from pg_get_publication_tables({pub}) gpt - order by oid; - "#, - pub = quote_literal(publication_name) - ); + let query = publication_table_ids_query(publication_name); let mut table_ids = vec![]; for row in self.client.simple_query(&query).await? { From 15716524329080d487d449f8cd565251734705b4 Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Mon, 7 Sep 2026 12:15:34 +0200 Subject: [PATCH 2/3] Fix --- crates/etl-api/tests/routes/pipelines.rs | 51 ++++++++++++++++++------ 1 file changed, 38 insertions(+), 13 deletions(-) diff --git a/crates/etl-api/tests/routes/pipelines.rs b/crates/etl-api/tests/routes/pipelines.rs index 884588aa4..011c0aef8 100644 --- a/crates/etl-api/tests/routes/pipelines.rs +++ b/crates/etl-api/tests/routes/pipelines.rs @@ -21,7 +21,7 @@ use etl_config::shared::{ BatchConfig, MemoryBackpressureConfig, PgConnectionConfig, PipelineConfig, ReplicationSlotConfig, TableSyncCopyConfig, }; -use etl_postgres::sqlx::test_utils::drop_pg_database; +use etl_postgres::{sqlx::test_utils::drop_pg_database, version::POSTGRES_15}; use etl_telemetry::tracing::init_test_tracing; use pg_escape::quote_identifier; use reqwest::StatusCode; @@ -134,13 +134,9 @@ async fn setup_pipeline_with_source_db() -> (TestApp, String, i64, PgPool, PgCon // We run the migrations to create all the tables used by `etl`. run_etl_migrations_on_source_database(&source_db_config).await; - // Match the default pipeline publication and include test tables as they are - // created. + // Match the default pipeline publication; tests add their tables explicitly. source_db_pool.execute("create schema test").await.unwrap(); - source_db_pool - .execute("create publication publication for tables in schema test") - .await - .unwrap(); + source_db_pool.execute("create publication publication").await.unwrap(); (app, tenant_id, pipeline_id, source_db_pool, source_db_config) } @@ -153,7 +149,7 @@ async fn create_table_with_state_chain( table_name: &str, state_chain: &[(&str, &str)], ) -> Oid { - let table_oid = create_test_table(source_db_pool, table_name).await; + let table_oid = create_published_test_table(source_db_pool, table_name).await; let mut prev_id: Option = None; for (i, (state, metadata)) in state_chain.iter().enumerate() { @@ -189,7 +185,7 @@ async fn create_tables_with_states( let mut results = Vec::new(); for (table_name, state, metadata) in tables { - let table_oid = create_test_table(source_db_pool, table_name).await; + let table_oid = create_published_test_table(source_db_pool, table_name).await; sqlx::query( "insert into etl.replication_state (pipeline_id, table_id, state, metadata, prev, \ @@ -236,6 +232,20 @@ async fn test_rollback( if expected_status.is_success() { Some(response.json().await.unwrap()) } else { None } } +/// Creates a table and explicitly adds it to the pipeline publication. +async fn create_published_test_table(source_db_pool: &PgPool, table_name: &str) -> Oid { + let table_oid = create_test_table(source_db_pool, table_name).await; + sqlx::query(AssertSqlSafe(format!( + "alter publication publication add table test.{}", + quote_identifier(table_name) + ))) + .execute(source_db_pool) + .await + .unwrap(); + table_oid +} + +/// Creates a table without changing publication membership. async fn create_test_table(source_db_pool: &PgPool, table_name: &str) -> Oid { sqlx::query("create schema if not exists test").execute(source_db_pool).await.unwrap(); @@ -1632,12 +1642,27 @@ async fn restarting_pipeline_resets_vpa_for_a_new_table_in_a_published_schema() init_test_tracing(); let (app, tenant_id, pipeline_id, source_db_pool, source_db_config) = setup_pipeline_with_source_db().await; + let server_version_num: i32 = + sqlx::query_scalar("select current_setting('server_version_num')::int") + .fetch_one(&source_db_pool) + .await + .unwrap(); + if server_version_num < POSTGRES_15 { + drop_pg_database(&source_db_config).await; + return; + } + create_tables_with_states( &source_db_pool, pipeline_id, &[("test_users", "ready", r#"{"type": "ready"}"#)], ) .await; + source_db_pool.execute("drop publication publication").await.unwrap(); + source_db_pool + .execute("create publication publication for tables in schema test") + .await + .unwrap(); create_test_table(&source_db_pool, "test_events").await; let response = app.restart_pipeline(&tenant_id, pipeline_id).await; @@ -1682,7 +1707,7 @@ async fn restarting_pipeline_resets_vpa_for_published_tables_without_any_state() init_test_tracing(); let (app, tenant_id, pipeline_id, source_db_pool, source_db_config) = setup_pipeline_with_source_db().await; - create_test_table(&source_db_pool, "test_users").await; + create_published_test_table(&source_db_pool, "test_users").await; let response = app.restart_pipeline(&tenant_id, pipeline_id).await; @@ -1763,7 +1788,7 @@ async fn restarting_pipeline_resets_vpa_when_a_published_table_is_recreated() { ) .await; source_db_pool.execute("drop table test.test_users").await.unwrap(); - let new_table_id = create_test_table(&source_db_pool, "test_users").await; + let new_table_id = create_published_test_table(&source_db_pool, "test_users").await; assert_ne!(tables[0].0, new_table_id); let response = app.restart_pipeline(&tenant_id, pipeline_id).await; @@ -1829,7 +1854,7 @@ async fn restarting_pipeline_respects_table_copy_selection() { &[("test_users", "ready", r#"{"type": "ready"}"#)], ) .await; - let new_table_id = create_test_table(&source_db_pool, "test_events").await; + let new_table_id = create_published_test_table(&source_db_pool, "test_events").await; let pipeline = app.read_pipeline(&tenant_id, pipeline_id).await; let pipeline: ReadPipelineResponse = pipeline.json().await.unwrap(); @@ -1896,7 +1921,7 @@ async fn restarting_pipeline_preserves_vpa_when_source_lock_times_out() { init_test_tracing(); let (app, tenant_id, pipeline_id, source_db_pool, source_db_config) = setup_pipeline_with_source_db().await; - create_test_table(&source_db_pool, "test_users").await; + create_published_test_table(&source_db_pool, "test_users").await; // Hold the lock until the source connection's lock timeout cancels inspection. let mut transaction = source_db_pool.begin().await.unwrap(); From dabdffdc4decb09b2b8fcc5fa25ce9f347c0fe47 Mon Sep 17 00:00:00 2001 From: Riccardo Busetti Date: Mon, 7 Sep 2026 12:57:09 +0200 Subject: [PATCH 3/3] Fix --- crates/etl-api/README.md | 38 ++++++------ crates/etl-api/src/data/pipelines.rs | 37 +++++------- crates/etl-api/src/routes/common.rs | 58 +++++++++---------- crates/etl-api/src/routes/pipelines.rs | 2 +- crates/etl-api/tests/routes/pipelines.rs | 52 ++++++++++------- crates/etl/src/replication/state/lifecycle.rs | 10 ++-- 6 files changed, 95 insertions(+), 102 deletions(-) diff --git a/crates/etl-api/README.md b/crates/etl-api/README.md index e84614b1c..474f0575b 100644 --- a/crates/etl-api/README.md +++ b/crates/etl-api/README.md @@ -224,27 +224,23 @@ fall back to it. Deployments that start with VPA `Off` may enable updates after their chosen observation policy has collected representative usage. Deployments that want immediate actuation may configure `in_place_or_recreate` instead. -Before restarting a running replicator, the API uses durable table state to -predict whether the restart will repeat an initial table copy. If any table is -before `SyncDone` and is not stopped in an error state, the API deletes the VPA -before reconciling the pipeline. Reconciliation recreates the VPA from the -current pipeline overrides, autoscaling configuration, and initial update mode. -Deleting the VPA resource does not guarantee that a running upstream recommender -forgets its in-memory usage aggregates. `SyncDone` and `Ready` tables keep their -destination data across restart and therefore preserve the existing VPA and its -live update mode. This applies to explicit restarts and configuration updates -that restart a running replicator. Source connection, query, or state-decoding -failures preserve the VPA rather than making the restart less reliable. - -Kubelet container restarts and Kubernetes- or controller-initiated Pod -replacements do not pass through the API restart path, so they do not delete the -VPA. A container restarted within the same Pod keeps that Pod's current -resources. A replacement Pod is created from the StatefulSet template, but VPA -admission may replace those resources with the existing recommendation. If a -replacement happens while table copy will repeat, the API cannot reset the VPA -to the initial configuration first. This is a known limitation of keeping the -copy-aware decision at the API boundary. A future Kubernetes controller with -access to durable table state could own this lifecycle. +Before restarting a running replicator, the API checks current publication +membership and durable table state in your database. If any published table +needs initial sync, including a newly published table, it deletes the VPA. +Reconciliation restores the configured bounds and initial update mode. This +covers the table sync phase, including when copying existing rows is skipped. +`SyncDone`, `Ready`, and errored tables do not trigger a reset. Inspection +failures preserve the VPA and allow restart. This applies to explicit restarts +and configuration updates that restart a running replicator. + +This is a best-effort reset, not a guarantee that memory stays at the startup +allocation throughout initial sync. State or publication membership can change +after inspection. Internal pipeline retries that start new table syncs, +container restarts, and Kubernetes Pod replacements bypass this API check and +do not reset the VPA, even during initial sync. The current Pod retains its +resources; a replacement may receive an existing VPA recommendation. Kubernetes +and the VPA's live policy govern those allocations. Deleting the VPA also does +not guarantee that the recommender forgets usage history. Stopping and starting a pipeline still deletes the autoscaler resource: stop deletes the StatefulSet and VPA, and start always recreates both resources. diff --git a/crates/etl-api/src/data/pipelines.rs b/crates/etl-api/src/data/pipelines.rs index 1c4409ad6..47dd6da79 100644 --- a/crates/etl-api/src/data/pipelines.rs +++ b/crates/etl-api/src/data/pipelines.rs @@ -1,7 +1,6 @@ use std::ops::DerefMut; use etl::store::TableState; -use etl_config::shared::TableSyncCopyConfig; use etl_postgres::{ publications::publication_table_ids_query, slots, @@ -38,26 +37,20 @@ use crate::{ /// only one pipeline will use it. pub const MAX_PIPELINES_PER_TENANT: i64 = 1; -/// Returns published tables that would copy data when the pipeline starts. +/// Returns published tables that would perform initial sync on restart. /// -/// Uses the same publication expansion as the replicator, including implicit -/// schema membership and partition-root settings. Joining current pipeline -/// state in one query detects new tables without including removed tables or -/// reading table data. Tables without state start in [`TableState::Init`]; -/// existing states retain the replicator's restart semantics. +/// Expands publication membership using the replicator's schema and partition +/// rules, then joins current state for this pipeline. New tables start in +/// [`TableState::Init`]. Copy selection does not affect this decision: skipping +/// existing rows still requires table sync. /// -/// The caller must supply the publication name and copy selection from the -/// current API configuration that will be materialized for this restart. This -/// function does not read the running Pod's configuration or verify that it -/// matches the API; copy eligibility assumes the replacement uses these values. -/// -/// This is a preflight observation, not a lock on publication membership or -/// replication progress. Either can change before the replicator starts. -pub(crate) async fn read_pipeline_tables_to_copy( +/// The publication must come from the configuration used for the replacement. +/// This observation does not lock membership or replication progress; either +/// can change before the replicator starts. +pub(crate) async fn read_pipeline_tables_to_sync( pool: &PgPool, pipeline_id: i64, publication_name: &str, - table_sync_copy: &TableSyncCopyConfig, ) -> Result, PipelineError> { let query = format!( r#" @@ -77,26 +70,24 @@ pub(crate) async fn read_pipeline_tables_to_copy( .fetch_all(pool) .await?; - let mut tables_to_copy = Vec::new(); + let mut tables_to_sync = Vec::new(); for (table_id, state_id, metadata) in rows { - if !table_sync_copy.should_copy_table(table_id.0) { - continue; - } - let state = if state_id.is_some() { let metadata = metadata.ok_or(PipelineError::MissingTableState)?; serde_json::from_value::(metadata) .map_err(PipelineError::InvalidTableState)? } else { + // A newly published table has no stored state yet. Treat it as Init so + // its first initial sync triggers a VPA reset before the worker starts. TableState::Init }; if state.as_type().would_perform_table_sync() { - tables_to_copy.push(table_id.0); + tables_to_sync.push(table_id.0); } } - Ok(tables_to_copy) + Ok(tables_to_sync) } #[derive(Debug, Clone)] diff --git a/crates/etl-api/src/routes/common.rs b/crates/etl-api/src/routes/common.rs index 8b3e74bae..9f74b9404 100644 --- a/crates/etl-api/src/routes/common.rs +++ b/crates/etl-api/src/routes/common.rs @@ -6,7 +6,7 @@ use crate::{ encryption::EncryptionKeyring, pipeline::StoredPipelineConfig, source::StoredSourceConfig, }, data::{ - pipelines::{read_pipeline_components, read_pipeline_tables_to_copy}, + pipelines::{read_pipeline_components, read_pipeline_tables_to_sync}, source_database, }, k8s::{ @@ -20,14 +20,14 @@ use crate::{ validation::{self, ValidationContext, ValidationError, ValidationFailure}, }; -/// Checks whether published tables would copy data, preserving VPA on -/// uncertainty. +/// Checks whether published tables would perform initial sync, preserving the +/// VPA when inspection fails. /// /// Uses the current API pipeline and source configuration supplied by the /// caller, which must also be used to materialize the replacement. These /// values may differ from the running Pod's configuration after an API update. /// Inspection uses the source connection's statement and lock timeouts. -async fn restart_would_copy_tables( +async fn restart_would_perform_table_sync( pipeline_id: i64, pipeline_config: &StoredPipelineConfig, source_id: i64, @@ -38,31 +38,26 @@ async fn restart_would_copy_tables( let connection_config = source_config.clone().into_connection_config(source_tls_config.get_tls_config()); let source_pool = source_database::connect(&connection_config).await?; - read_pipeline_tables_to_copy( - &source_pool, - pipeline_id, - &pipeline_config.publication_name, - &pipeline_config.table_sync_copy, - ) - .await + read_pipeline_tables_to_sync(&source_pool, pipeline_id, &pipeline_config.publication_name) + .await }; match inspection.await { - Ok(tables_to_copy) => { + Ok(tables_to_sync) => { info!( pipeline_id, source_id, - table_count = tables_to_copy.len(), - "determined tables to copy on pipeline restart", + table_count = tables_to_sync.len(), + "determined tables to sync on pipeline restart", ); - !tables_to_copy.is_empty() + !tables_to_sync.is_empty() } Err(error) => { warn!( pipeline_id, source_id, error = %error, - "failed to determine tables to copy on pipeline restart, preserving vertical pod autoscaler", + "failed to determine tables to sync on pipeline restart, preserving vertical pod autoscaler", ); false } @@ -75,7 +70,7 @@ async fn restart_would_copy_tables( /// runtime resource configuration should call this after writing the new API /// state through the supplied connection, including within an uncommitted /// transaction. The helper reads that state once and uses the same loaded -/// pipeline and source configuration for both copy preflight and Kubernetes +/// pipeline and source configuration for both sync preflight and Kubernetes /// materialization. Updating the StatefulSet changes the pod template restart /// annotation. /// @@ -84,20 +79,21 @@ async fn restart_would_copy_tables( /// running pod must be restarted after config materialization in order to pick /// up those changes. /// -/// Before reconciliation, this best-effort checks current publication -/// membership, durable pipeline state, and table-copy settings in the database. -/// If any table would copy data, including newly published tables, it deletes -/// the VPA so reconciliation recreates it in the configured initial update -/// mode. With `Off`, the replacement Pod starts at the configured startup -/// allocation and the VPA gets a fresh observation period. The upstream -/// recommender may retain usage history after deletion. Inspection failures and -/// timeouts preserve the existing VPA and do not block restart. +/// Before reconciliation, checks current publication membership and durable +/// table state. If any table would perform initial sync, it deletes the VPA so +/// reconciliation restores its configured bounds and initial update mode. +/// This covers table sync even when copying existing rows is skipped. With +/// `Off`, the replacement Pod starts with the configured resources; this does +/// not guarantee that memory stays at that level throughout initial sync. +/// The recommender may retain usage history. Inspection failures and timeouts +/// preserve the VPA and do not block restart. /// -/// Kubelet container restarts and Kubernetes-initiated Pod replacements do not -/// call this helper or delete the VPA. A replacement Pod may therefore receive -/// an existing recommendation even when initial sync will repeat. This is a -/// limitation of making the sync-aware decision at the API boundary; a future -/// controller with access to durable table state could own that lifecycle. +/// State or publication changes after inspection can race this decision. +/// Internal pipeline retries, container restarts, and Kubernetes-initiated Pod +/// replacements bypass it, including during initial sync. They do not reset +/// the VPA: the current Pod retains its resources, and a replacement may +/// receive an existing recommendation. Resource allocation outside this API +/// path is therefore governed by Kubernetes and the VPA's live policy. /// /// If Kubernetes support is unavailable, or the pipeline has no active /// Kubernetes resources, the call returns `false` without reconciling. @@ -118,7 +114,7 @@ pub(crate) async fn restart_replicator_if_running( return Ok(false); } - if restart_would_copy_tables( + if restart_would_perform_table_sync( pipeline_id, &pipeline.config, source.id, diff --git a/crates/etl-api/src/routes/pipelines.rs b/crates/etl-api/src/routes/pipelines.rs index 5542d9136..1e122ef3a 100644 --- a/crates/etl-api/src/routes/pipelines.rs +++ b/crates/etl-api/src/routes/pipelines.rs @@ -1041,7 +1041,7 @@ pub(crate) async fn start_pipeline( post, path = "/pipelines/{pipeline_id}/restart", summary = "Restart a pipeline", - description = "Reconciles the pipeline's Kubernetes resources and restarts its replicator. Every replicator has a VPA; pipeline resource overrides fix the corresponding VPA bounds. When durable pipeline state in your database shows that the restart will repeat initial sync while copying existing tables, the endpoint deletes and recreates the VPA from its configured bounds and initial update mode. Deletion does not guarantee that the upstream recommender forgets in-memory usage aggregates. If that state cannot be inspected, or all tables completed initial sync, the current VPA is preserved. Kubelet container restarts and Kubernetes-initiated replacement Pods do not use this endpoint or delete the VPA, so a replacement may receive an existing recommendation even when initial sync repeats.", + description = "Reconciles the pipeline's Kubernetes resources and restarts its replicator. If current publication membership and durable state in your database indicate initial sync, the API resets the VPA to its configured bounds and initial update mode, including when copying existing rows is skipped. If no table needs initial sync or inspection fails, it preserves the VPA. This is a best-effort check, not a guarantee of memory allocation throughout initial sync: state can change after inspection, and internal pipeline retries, container restarts, and Kubernetes Pod replacements bypass the reset. Existing VPA recommendations may still apply, and deletion does not guarantee that the recommender forgets usage history.", params( ("pipeline_id" = i64, Path, description = "Unique ID of the pipeline"), ("tenant_id" = String, Header, description = "Tenant ID used to scope the request") diff --git a/crates/etl-api/tests/routes/pipelines.rs b/crates/etl-api/tests/routes/pipelines.rs index 011c0aef8..875670249 100644 --- a/crates/etl-api/tests/routes/pipelines.rs +++ b/crates/etl-api/tests/routes/pipelines.rs @@ -894,7 +894,7 @@ async fn updating_a_running_pipeline_reapplies_replicator_resources() { } #[tokio::test(flavor = "multi_thread")] -async fn updating_a_running_pipeline_resets_vpa_when_table_copy_will_repeat() { +async fn updating_a_running_pipeline_resets_vpa_when_table_sync_will_repeat() { init_test_tracing(); let (app, tenant_id, pipeline_id, source_db_pool, source_db_config) = setup_pipeline_with_source_db().await; @@ -1555,27 +1555,41 @@ async fn a_running_pipeline_can_be_restarted() { } #[tokio::test(flavor = "multi_thread")] -async fn restarting_pipeline_resets_vpa_when_table_copy_will_repeat() { +async fn restarting_pipeline_resets_vpa_when_table_sync_will_repeat() { init_test_tracing(); let (app, tenant_id, pipeline_id, source_db_pool, source_db_config) = setup_pipeline_with_source_db().await; create_tables_with_states( &source_db_pool, pipeline_id, - &[("test_users", "data_sync", r#"{"type": "data_sync"}"#)], + &[("test_users", "init", r#"{"type": "init"}"#)], ) .await; - let response = app.restart_pipeline(&tenant_id, pipeline_id).await; + // Finishing the copy does not complete initial sync; catchup must also finish. + for state in ["init", "data_sync", "finished_copy"] { + sqlx::query( + "update etl.replication_state set state = $1::text::etl.table_state, metadata = \ + jsonb_build_object('type', $1::text) where pipeline_id = $2", + ) + .bind(state) + .bind(pipeline_id) + .execute(&source_db_pool) + .await + .unwrap(); + let delete_calls_before = app.k8s_state.vpa_delete_calls(); - assert_eq!(response.status(), StatusCode::ACCEPTED); - assert_eq!(app.k8s_state.vpa_delete_calls(), 1); + let response = app.restart_pipeline(&tenant_id, pipeline_id).await; + + assert_eq!(response.status(), StatusCode::ACCEPTED); + assert_eq!(app.k8s_state.vpa_delete_calls() - delete_calls_before, 1); + } drop_pg_database(&source_db_config).await; } #[tokio::test(flavor = "multi_thread")] -async fn restarting_pipeline_preserves_vpa_when_no_table_copy_will_repeat() { +async fn restarting_pipeline_preserves_vpa_when_no_table_sync_will_repeat() { init_test_tracing(); let (app, tenant_id, pipeline_id, source_db_pool, source_db_config) = setup_pipeline_with_source_db().await; @@ -1718,7 +1732,7 @@ async fn restarting_pipeline_resets_vpa_for_published_tables_without_any_state() } #[tokio::test(flavor = "multi_thread")] -async fn restarting_pipeline_preserves_vpa_for_copy_states_outside_the_publication() { +async fn restarting_pipeline_preserves_vpa_for_sync_states_outside_the_publication() { init_test_tracing(); let (app, tenant_id, pipeline_id, source_db_pool, source_db_config) = setup_pipeline_with_source_db().await; @@ -1762,7 +1776,7 @@ async fn restarting_pipeline_uses_only_current_state_for_its_pipeline() { assert_eq!(response.status(), StatusCode::ACCEPTED); assert_eq!(app.k8s_state.vpa_delete_calls(), 0); - // A different pipeline's completed copy cannot suppress this pipeline's copy. + // A different pipeline's completed sync cannot suppress this pipeline's sync. create_tables_with_states( &source_db_pool, pipeline_id + 1, @@ -1844,11 +1858,11 @@ async fn restarting_pipeline_respects_partition_root_publication_settings() { } #[tokio::test(flavor = "multi_thread")] -async fn restarting_pipeline_respects_table_copy_selection() { +async fn restarting_pipeline_resets_vpa_even_when_table_copy_is_skipped() { init_test_tracing(); let (app, tenant_id, pipeline_id, source_db_pool, source_db_config) = setup_pipeline_with_source_db().await; - let ready_tables = create_tables_with_states( + create_tables_with_states( &source_db_pool, pipeline_id, &[("test_users", "ready", r#"{"type": "ready"}"#)], @@ -1858,13 +1872,10 @@ async fn restarting_pipeline_respects_table_copy_selection() { let pipeline = app.read_pipeline(&tenant_id, pipeline_id).await; let pipeline: ReadPipelineResponse = pipeline.json().await.unwrap(); - for (table_sync_copy, copies_new_table) in [ - (TableSyncCopyConfig::SkipAllTables, false), - (TableSyncCopyConfig::IncludeTables { table_ids: vec![ready_tables[0].0.0] }, false), - (TableSyncCopyConfig::SkipTables { table_ids: vec![new_table_id.0] }, false), - (TableSyncCopyConfig::IncludeTables { table_ids: vec![new_table_id.0] }, true), - (TableSyncCopyConfig::SkipTables { table_ids: vec![ready_tables[0].0.0] }, true), - (TableSyncCopyConfig::IncludeAllTables, true), + // Both global and per-table copy exclusions still require initial sync. + for table_sync_copy in [ + TableSyncCopyConfig::SkipAllTables, + TableSyncCopyConfig::SkipTables { table_ids: vec![new_table_id.0] }, ] { let response = app .update_pipeline( @@ -1886,10 +1897,7 @@ async fn restarting_pipeline_respects_table_copy_selection() { let response = app.restart_pipeline(&tenant_id, pipeline_id).await; assert_eq!(response.status(), StatusCode::ACCEPTED); - assert_eq!( - app.k8s_state.vpa_delete_calls() - delete_calls_before, - usize::from(copies_new_table), - ); + assert_eq!(app.k8s_state.vpa_delete_calls() - delete_calls_before, 1); } drop_pg_database(&source_db_config).await; diff --git a/crates/etl/src/replication/state/lifecycle.rs b/crates/etl/src/replication/state/lifecycle.rs index 0301c9967..84d2cc09f 100644 --- a/crates/etl/src/replication/state/lifecycle.rs +++ b/crates/etl/src/replication/state/lifecycle.rs @@ -703,18 +703,20 @@ mod tests { let completed_states = [TableStateType::SyncDone, TableStateType::Ready]; assert!(completed_states.iter().all(TableStateType::has_completed_table_sync)); - let states_that_repeat_copy = [ + let states_that_repeat_sync = [ TableStateType::Init, TableStateType::DataSync, TableStateType::FinishedCopy, TableStateType::SyncWait, TableStateType::Catchup, ]; - assert!(states_that_repeat_copy.iter().all(TableStateType::would_perform_table_sync)); + assert!(states_that_repeat_sync.iter().all(TableStateType::would_perform_table_sync)); - let states_that_preserve_copy = + let states_that_do_not_repeat_sync = [TableStateType::SyncDone, TableStateType::Ready, TableStateType::Errored]; - assert!(states_that_preserve_copy.iter().all(|state| !state.would_perform_table_sync())); + assert!( + states_that_do_not_repeat_sync.iter().all(|state| !state.would_perform_table_sync()) + ); } #[test]