diff --git a/.env.example b/.env.example index d3deb6f536..f79eb23f73 100644 --- a/.env.example +++ b/.env.example @@ -49,8 +49,13 @@ REDIS_URL=redis://localhost:6379 # BUZZ_DB_POOL_SIZE=50 # Read-only partition catalog audit interval, clamped to 60s..24h (default 900s). # BUZZ_PARTITION_AUDIT_INTERVAL_SECS=900 -# Emergency kill switch for monthly partition CREATE; audits remain enabled. +# Emergency kill switch for all automatic partition DDL, including catch-all +# advancement; audits remain enabled. # BUZZ_PARTITION_MANAGER_CREATE_ENABLED=true +# Opt in to replacing an empty partition catch-all with six months of dedicated +# monthlies, at startup and on every audit interval (default false). Has no +# effect while the kill switch above is off. +# BUZZ_PARTITION_MANAGER_ADVANCE_ENABLED=false # Writer-session Postgres timeouts for buzz-db-backed pools and the relay audit # pool, all in milliseconds; 0 disables. The separately deployed push gateway diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 55e3bbe969..82d7cac7bc 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -462,7 +462,7 @@ All database access. Uses `sqlx::query()` (runtime, not compile-time macros) — | `channel.rs` | Channel CRUD, membership management, role enforcement (transactional) | | `feed.rs` | `query_mentions` (INNER JOIN event_mentions), `query_needs_action`, `query_activity` | | `workflow.rs` | Full workflow/run/approval CRUD; SHA-256 hashed approval tokens | -| `partition.rs` | Monthly range partitioning for `events` and `delivery_log` tables | +| `partition.rs` | Monthly range partitioning for `events` and `delivery_log` tables: read-only catalog audit plus advisory-locked, time-bounded maintenance that creates uncovered months and advances an empty catch-all (`partition/maintenance.rs`) | | `dm.rs` | DM channel management | | `reaction.rs` | Reaction storage and retrieval | | `thread.rs` | Thread/reply tracking | diff --git a/crates/buzz-admin/src/main.rs b/crates/buzz-admin/src/main.rs index 312f01c1de..66780e2041 100644 --- a/crates/buzz-admin/src/main.rs +++ b/crates/buzz-admin/src/main.rs @@ -94,7 +94,7 @@ enum Command { /// Run the partition catalog audit using a read-only database session. PartitionAudit { /// Future months to include in the coverage check. - #[arg(long, default_value_t = 3)] + #[arg(long, default_value_t = buzz_db::partition::PARTITION_MANAGER_MONTHS_AHEAD)] months_ahead: u32, }, /// Inspect deployment-wide Buzz product feedback. diff --git a/crates/buzz-db/src/error.rs b/crates/buzz-db/src/error.rs index 4f15cc2b24..8c051a7a79 100644 --- a/crates/buzz-db/src/error.rs +++ b/crates/buzz-db/src/error.rs @@ -45,6 +45,11 @@ pub enum DbError { #[error("invalid data: {0}")] InvalidData(String), + /// A partition the catalog audit had already listed was dropped before the + /// audit finished. A fresh snapshot sees a consistent catalog. + #[error("partition catalog changed mid-audit: {0}")] + PartitionDroppedMidAudit(String), + /// A serving write admitted before the lifecycle transition is still live. /// This is an ordinary retryable drain condition, not a safety violation. #[error( diff --git a/crates/buzz-db/src/runtime/observability.rs b/crates/buzz-db/src/runtime/observability.rs index 99256f3c2f..d487b9a362 100644 --- a/crates/buzz-db/src/runtime/observability.rs +++ b/crates/buzz-db/src/runtime/observability.rs @@ -184,16 +184,18 @@ pub(crate) enum LockType { PushGate, Deletion, MigrationSchemaSafety, + PartitionMaintenance, } impl LockType { #[cfg(test)] - pub(crate) const ALL: [Self; 5] = [ + pub(crate) const ALL: [Self; 6] = [ Self::Replacement, Self::Membership, Self::PushGate, Self::Deletion, Self::MigrationSchemaSafety, + Self::PartitionMaintenance, ]; pub(crate) const fn as_str(self) -> &'static str { @@ -203,6 +205,7 @@ impl LockType { Self::PushGate => "push_gate", Self::Deletion => "deletion", Self::MigrationSchemaSafety => "migration_schema_safety", + Self::PartitionMaintenance => "partition_maintenance", } } } @@ -671,6 +674,7 @@ mod tests { "push_gate", "deletion", "migration_schema_safety", + "partition_maintenance", ] ); assert_eq!( diff --git a/crates/buzz-db/src/store/artifact_postgres_tests.rs b/crates/buzz-db/src/store/artifact_postgres_tests.rs index 078cefb8f6..4ac03a06f4 100644 --- a/crates/buzz-db/src/store/artifact_postgres_tests.rs +++ b/crates/buzz-db/src/store/artifact_postgres_tests.rs @@ -24,7 +24,15 @@ impl Fixture { crate::migration::run_migrations(&pool).await.unwrap(); } let db = Db::from_pool(pool); - db.ensure_future_partitions(1, true).await.unwrap(); + db.maintain_partitions( + 1, + crate::partition::PartitionMaintenancePolicy { + create_enabled: true, + advance_enabled: false, + }, + ) + .await + .unwrap(); let community = CommunityId::from_uuid(Uuid::new_v4()); sqlx::query("INSERT INTO communities(id,host) VALUES($1,$2)") .bind(community.as_uuid()) diff --git a/crates/buzz-db/src/store/partition.rs b/crates/buzz-db/src/store/partition.rs index 733ab73213..02dfa4e708 100644 --- a/crates/buzz-db/src/store/partition.rs +++ b/crates/buzz-db/src/store/partition.rs @@ -7,18 +7,24 @@ use std::time::{Duration, Instant}; use buzz_datastore_tracing::datastore_span; use chrono::{DateTime, Datelike, NaiveDate, TimeZone, Timelike, Utc}; use serde::Serialize; -use sqlx::{Connection, PgConnection, PgPool, Row}; -use tracing::info; +use sqlx::{PgConnection, PgPool, Row}; use crate::error::{DbError, Result}; use crate::Db; +mod maintenance; + +pub use maintenance::{ + maintain_partitions, PartitionMaintenanceOutcome, PartitionMaintenancePolicy, + PARTITION_MANAGER_MONTHS_AHEAD, +}; + /// Tables that may be partition-managed. The allowlist prevents DDL injection. const PARTITIONED_TABLES: &[&str] = &["events", "delivery_log"]; /// Maximum future-month horizon accepted by catalog audits and creation. /// -/// The relay uses three months. Keeping the public API bounded prevents an +/// The relay uses [`PARTITION_MANAGER_MONTHS_AHEAD`]. Keeping the public API bounded prevents an /// operator-supplied `buzz-admin partition-audit --months-ahead` value from /// allocating or iterating an effectively unbounded report. pub const MAX_PARTITION_MONTHS_AHEAD: u32 = 120; @@ -29,6 +35,10 @@ const PARTITION_AUDIT_STATEMENT_TIMEOUT: &str = "5s"; /// Maximum wall-clock duration for one complete multi-table catalog audit. const PARTITION_AUDIT_TOTAL_TIMEOUT: Duration = Duration::from_secs(10); +/// Fresh-snapshot retries after a partition vanishes mid-audit, such as when +/// another relay's maintenance replaces a catch-all. +const PARTITION_AUDIT_CONCURRENT_DROP_RETRIES: usize = 2; + fn expected_partition_key(table: &str) -> Option<&'static str> { match table { "events" => Some("RANGE (created_at)"), @@ -451,131 +461,24 @@ async fn audit_partition_catalog_report_at_with_timeout( } } -/// Audit first, then create only months proven to be uncovered. -/// -/// Covered ranges are never probed with DDL. Creation failures are collected -/// across all managed parents and months before an aggregate error is returned. -pub async fn ensure_future_partitions( - pool: &PgPool, - months_ahead: u32, - create_enabled: bool, -) -> Result { - ensure_future_partitions_at(pool, months_ahead, create_enabled, Utc::now()).await -} - +/// Create-only maintenance at a fixed clock, for tests of month creation. +#[cfg(test)] async fn ensure_future_partitions_at( pool: &PgPool, months_ahead: u32, create_enabled: bool, now: DateTime, ) -> Result { - let audit = audit_partition_catalog_at(pool, months_ahead, now).await?; - let mut errors = Vec::new(); - let mut created_any = false; - - for table in &audit.tables { - for month in &table.months { - match month.kind { - MonthCoverageKind::CoveredByMonthly - | MonthCoverageKind::CoveredByCatchAll - | MonthCoverageKind::CoveredByDefault => { - metrics::counter!( - "buzz_partition_create_attempts_total", - "table" => table.table, - "outcome" => "skipped_covered" - ) - .increment(1); - } - MonthCoverageKind::Uncovered if !create_enabled => {} - MonthCoverageKind::Uncovered => { - let expected_name = partition_name(table.table, month.start); - let collision = match relation_name_exists(pool, &expected_name).await { - Ok(collision) => collision, - Err(error) => { - metrics::counter!( - "buzz_partition_create_attempts_total", - "table" => table.table, - "outcome" => "error" - ) - .increment(1); - errors.push(format!( - "{} {}: failed to check canonical name {expected_name}: {error}", - table.table, - month.start.format("%Y-%m") - )); - continue; - } - }; - if collision { - metrics::counter!( - "buzz_partition_create_attempts_total", - "table" => table.table, - "outcome" => "error" - ) - .increment(1); - errors.push(format!( - "{} {}: canonical name {expected_name} already exists without the expected attachment and bounds", - table.table, - month.start.format("%Y-%m") - )); - continue; - } - if !table.structurally_safe_for_creation() { - metrics::counter!( - "buzz_partition_create_attempts_total", - "table" => table.table, - "outcome" => "error" - ) - .increment(1); - errors.push(format!( - "{} {}: catalog is not structurally safe for automatic partition creation", - table.table, - month.start.format("%Y-%m") - )); - continue; - } - match create_month_partition(pool, table.table, month.start).await { - Ok(name) => { - created_any = true; - metrics::counter!( - "buzz_partition_create_attempts_total", - "table" => table.table, - "outcome" => "created" - ) - .increment(1); - info!(table = table.table, partition = name, "added partition"); - } - Err(error) => { - metrics::counter!( - "buzz_partition_create_attempts_total", - "table" => table.table, - "outcome" => "error" - ) - .increment(1); - errors.push(format!( - "{} {}: {error}", - table.table, - month.start.format("%Y-%m") - )); - } - } - } - } - } - } - - if errors.is_empty() { - if created_any { - audit_partition_catalog_at(pool, months_ahead, now).await - } else { - Ok(audit) - } - } else { - Err(DbError::InvalidData(format!( - "partition creation failed: {}", - errors.join("; ") - ))) - } + maintenance::maintain_partitions_at( + pool, + months_ahead, + PartitionMaintenancePolicy { + create_enabled, + advance_enabled: false, + }, + now, + ) + .await } async fn audit_table( @@ -585,6 +488,26 @@ async fn audit_table( now: DateTime, ) -> Result { let months_ahead = validated_months_ahead(months_ahead)?; + let mut retries = 0; + loop { + match audit_table_once(pool, table, months_ahead, now).await { + Err(error) + if retries < PARTITION_AUDIT_CONCURRENT_DROP_RETRIES + && is_concurrent_drop(&error) => + { + retries += 1; + } + result => return result, + } + } +} + +async fn audit_table_once( + pool: &PgPool, + table: &'static str, + months_ahead: i32, + now: DateTime, +) -> Result { let mut transaction = pool.begin().await?; sqlx::query("SET TRANSACTION READ ONLY") .execute(&mut *transaction) @@ -595,6 +518,20 @@ async fn audit_table( Ok(audit) } +/// Whether an audit failed only because a partition it had already listed was +/// dropped by a concurrent transaction. A fresh snapshot sees a consistent +/// catalog, so the audit can be retried. +fn is_concurrent_drop(error: &DbError) -> bool { + match error { + DbError::PartitionDroppedMidAudit(_) => true, + // `relation_contains_rows` names a leaf that no longer exists. + DbError::Sqlx(sqlx::Error::Database(database)) => { + database.code().as_deref() == Some("42P01") + } + _ => false, + } +} + async fn configure_audit_transaction(connection: &mut PgConnection) -> Result<()> { pin_catalog_rendering(connection).await?; sqlx::query("SELECT pg_catalog.set_config('statement_timeout', $1, true)") @@ -627,14 +564,14 @@ impl Db { audit_partition_catalog_report(&self.pool, months_ahead).await } - /// Ensures monthly partitions exist for the next N months when creation is enabled. - #[datastore_span(name = "ensure_future_partitions", system = "postgresql")] - pub async fn ensure_future_partitions( + /// Issues the bounded partition DDL `policy` permits, then returns a fresh audit. + #[datastore_span(name = "maintain_partitions", system = "postgresql")] + pub async fn maintain_partitions( &self, months_ahead: u32, - create_enabled: bool, + policy: PartitionMaintenancePolicy, ) -> Result { - ensure_future_partitions(&self.pool, months_ahead, create_enabled).await + maintain_partitions(&self.pool, months_ahead, policy).await } } @@ -758,14 +695,25 @@ async fn audit_table_on( let is_leaf: bool = row.try_get("is_leaf")?; let bound_partition_key: Option = row.try_get("bound_partition_key")?; let detach_pending: bool = row.try_get("detach_pending")?; - let sibling_bounds: Vec = row.try_get("sibling_bounds")?; + // Catalog rows come from the statement snapshot, but `pg_get_expr` + // renders from current catalog state and yields NULL for a relation + // dropped since the snapshot was taken. + let sibling_bounds = row + .try_get::>, _>("sibling_bounds")? + .into_iter() + .collect::>>() + .ok_or_else(|| { + DbError::PartitionDroppedMidAudit(format!("a sibling of partition {name}")) + })?; if detach_pending { pending_roots.insert(root_child_relation_oid); } let partition_key_compatible = partition_key_valid && bound_partition_key.as_deref() == Some(expected_partition_key) && !detach_pending; - let expression: String = row.try_get("bound")?; + let expression: String = row + .try_get::, _>("bound")? + .ok_or_else(|| DbError::PartitionDroppedMidAudit(format!("partition {name}")))?; let is_default = expression.trim() == "DEFAULT"; let own_range = parse_range_bounds(&expression); let is_catch_all = own_range.as_ref().is_some_and(|(lower, upper)| { @@ -1536,84 +1484,7 @@ fn parse_timestamp_literal(literal: &str) -> Option> { Some(Utc.from_utc_datetime(&date.and_hms_opt(0, 0, 0)?)) } -async fn create_month_partition( - pool: &PgPool, - table: &str, - start: DateTime, -) -> Result { - if !PARTITIONED_TABLES.contains(&table) { - return Err(DbError::InvalidData(format!( - "table not in partition allowlist: {table:?}" - ))); - } - let (end_year, end_month) = add_months(start.year(), start.month(), 1)?; - let end = month_start(end_year, end_month)?; - let partition_name = partition_name(table, start); - let start_date = start.format("%Y-%m-%d"); - let end_date = end.format("%Y-%m-%d"); - let sql = format!( - "CREATE TABLE {} PARTITION OF {} \ - FOR VALUES FROM ('{start_date}') TO ('{end_date}')", - quote_identifier(&partition_name), - quote_identifier(table) - ); - let mut connection = crate::observability::acquire_writer( - pool, - crate::observability::WriterOperation::Bootstrap, - ) - .await?; - let mut transaction = connection.begin().await?; - pin_catalog_rendering(&mut transaction).await?; - sqlx::query(sqlx::AssertSqlSafe(sql)) - .execute(&mut *transaction) - .await?; - let row = sqlx::query( - r#" - SELECT child.relkind::text AS relation_kind, - inherited.inhdetachpending AS detach_pending, - pg_catalog.pg_get_partkeydef(parent.oid) AS partition_key, - pg_catalog.pg_get_expr(child.relpartbound, child.oid) AS bound - FROM pg_catalog.pg_inherits inherited - JOIN pg_catalog.pg_class parent ON parent.oid = inherited.inhparent - JOIN pg_catalog.pg_namespace parent_ns ON parent_ns.oid = parent.relnamespace - JOIN pg_catalog.pg_class child ON child.oid = inherited.inhrelid - JOIN pg_catalog.pg_namespace child_ns ON child_ns.oid = child.relnamespace - WHERE parent_ns.nspname = current_schema() - AND child_ns.nspname = current_schema() - AND parent.relname = $1 - AND child.relname = $2 - "#, - ) - .bind(table) - .bind(&partition_name) - .fetch_optional(&mut *transaction) - .await?; - let exact = row.is_some_and(|row| { - let relation_kind = row.try_get::("relation_kind").ok(); - let detach_pending = row.try_get::("detach_pending").ok(); - let partition_key = row - .try_get::, _>("partition_key") - .ok() - .flatten(); - let bounds = row - .try_get::("bound") - .ok() - .and_then(|bound| parse_range_bounds(&bound)); - relation_kind.as_deref() == Some("r") - && detach_pending == Some(false) - && partition_key.as_deref() == expected_partition_key(table) - && bounds == Some((PartitionBound::Finite(start), PartitionBound::Finite(end))) - }); - if !exact { - transaction.rollback().await?; - return Err(DbError::InvalidData(format!( - "created partition {partition_name} failed exact catalog postcondition" - ))); - } - transaction.commit().await?; - Ok(partition_name) -} - +#[cfg(test)] async fn relation_name_exists(pool: &PgPool, relation: &str) -> Result { Ok(sqlx::query_scalar( r#" @@ -1687,6 +1558,20 @@ mod tests { .all(|error| error.error.contains("months_ahead must be at most 120"))); } + #[test] + fn only_a_mid_audit_drop_is_retried() { + assert!(is_concurrent_drop(&DbError::PartitionDroppedMidAudit( + "partition events_p2026_09".to_string() + ))); + assert!(!is_concurrent_drop(&DbError::NotFound("row".to_string()))); + assert!(!is_concurrent_drop(&DbError::InvalidData( + "bound".to_string() + ))); + assert!(!is_concurrent_drop(&DbError::Sqlx( + sqlx::Error::PoolTimedOut + ))); + } + #[test] fn parses_pg16_range_bound_formats() { let expected = Utc.with_ymd_and_hms(2026, 7, 1, 0, 0, 0).unwrap(); @@ -1903,17 +1788,17 @@ mod tests { assert_eq!(extra, Some(1.0)); } - mod postgres_tests { + pub(super) mod postgres_tests { use sqlx::postgres::PgPoolOptions; use uuid::Uuid; use super::*; - async fn scratch_pool() -> (PgPool, PgPool, String) { + pub(in crate::store::partition) async fn scratch_pool() -> (PgPool, PgPool, String) { scratch_pool_with_max_connections(4).await } - async fn scratch_pool_with_max_connections( + pub(in crate::store::partition) async fn scratch_pool_with_max_connections( max_connections: u32, ) -> (PgPool, PgPool, String) { let url = crate::test_support::database_url(); @@ -1941,7 +1826,7 @@ mod tests { (pool, admin, schema) } - async fn drop_schema(admin: &PgPool, schema: &str) { + pub(in crate::store::partition) async fn drop_schema(admin: &PgPool, schema: &str) { let _ = sqlx::query(sqlx::AssertSqlSafe(format!( "DROP SCHEMA IF EXISTS {schema} CASCADE" ))) @@ -1949,7 +1834,7 @@ mod tests { .await; } - async fn seed_parents(pool: &PgPool) { + pub(in crate::store::partition) async fn seed_parents(pool: &PgPool) { seed_parents_with_keys(pool, "created_at", "delivered_at").await; } @@ -1984,7 +1869,13 @@ mod tests { } } - async fn create_child(pool: &PgPool, table: &str, name: &str, lower: &str, upper: &str) { + pub(in crate::store::partition) async fn create_child( + pool: &PgPool, + table: &str, + name: &str, + lower: &str, + upper: &str, + ) { let sql = format!( "CREATE TABLE {name} PARTITION OF {table} FOR VALUES FROM ({lower}) TO ({upper})" ); @@ -2032,7 +1923,9 @@ mod tests { .expect("create nested child"); } - async fn catalog_snapshot(pool: &PgPool) -> Vec<(String, String, String, i64)> { + pub(in crate::store::partition) async fn catalog_snapshot( + pool: &PgPool, + ) -> Vec<(String, String, String, i64)> { let mut transaction = pool.begin().await.expect("begin catalog snapshot"); sqlx::query("SET TRANSACTION READ ONLY") .execute(&mut *transaction) @@ -2067,11 +1960,11 @@ mod tests { snapshot } - fn fixed_now() -> DateTime { + pub(in crate::store::partition) fn fixed_now() -> DateTime { Utc.with_ymd_and_hms(2026, 9, 15, 12, 0, 0).unwrap() } - async fn seed_fresh_layout(pool: &PgPool) { + pub(in crate::store::partition) async fn seed_fresh_layout(pool: &PgPool) { seed_parents(pool).await; for table in PARTITIONED_TABLES { create_child( diff --git a/crates/buzz-db/src/store/partition/maintenance.rs b/crates/buzz-db/src/store/partition/maintenance.rs new file mode 100644 index 0000000000..f19699181b --- /dev/null +++ b/crates/buzz-db/src/store/partition/maintenance.rs @@ -0,0 +1,1117 @@ +//! Serialized, bounded partition DDL for `events` and `delivery_log`. +//! +//! The manager issues only two kinds of change: creating a month that no +//! partition covers, and replacing an empty right-edge catch-all with +//! dedicated monthlies plus a new catch-all after the horizon. Each table's +//! change runs in one transaction that +//! +//! 1. holds a schema-scoped advisory lock, so concurrent relays produce one +//! winner and `skipped_locked` losers; +//! 2. locks the parent, then foreign-key counterparts read under the parent +//! lock, then the catch-all, in the order writers take them, waiting only on +//! the parent and failing fast on any other contention; +//! 3. re-audits under those locks and refuses if the plan changed; and +//! 4. re-audits the changed catalog and verifies index and constraint +//! inheritance before commit. +//! +//! A populated catch-all is always refused before any lock is taken: moving +//! rows requires an operator. + +use std::time::Duration; + +use chrono::{DateTime, Datelike, Utc}; +use serde::Serialize; +use sqlx::{Connection, PgConnection, PgPool, Row}; +use tracing::info; + +use super::{ + add_months, audit_partition_catalog_at, audit_table_on, month_start, partition_name, + pin_catalog_rendering, qualified_relation_name, validated_months_ahead, MonthCoverageKind, + PartitionAudit, PartitionBound, PartitionChildKind, PartitionTableAudit, + MAX_PARTITION_MONTHS_AHEAD, +}; +use crate::error::{DbError, Result}; + +/// Dedicated monthly runway, after the current month, that the relay maintains. +pub const PARTITION_MANAGER_MONTHS_AHEAD: u32 = 6; + +/// Maximum wait for the parent lock. Counterparts and children use `NOWAIT`. +const MAINTENANCE_LOCK_TIMEOUT: &str = "2s"; + +/// Server-side budget for each statement in the maintenance transaction. +const MAINTENANCE_STATEMENT_TIMEOUT: &str = "5s"; + +/// Server-side reaping of a maintenance session whose client stalls mid-transaction. +const MAINTENANCE_IDLE_IN_TRANSACTION_TIMEOUT: &str = "5s"; + +/// Client-side budget for one table's complete maintenance attempt. +const MAINTENANCE_TABLE_DEADLINE: Duration = Duration::from_secs(10); + +/// Namespace hashed with `current_schema()` into the maintenance advisory lock. +const MAINTENANCE_LOCK_NAMESPACE: &str = "buzz.partition_maintenance:"; + +/// Which automatic partition DDL the manager may issue. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub struct PartitionMaintenancePolicy { + /// Create months that no partition covers. Off stops all automatic DDL. + pub create_enabled: bool, + /// Replace an empty right-edge catch-all with dedicated monthlies. + /// Takes effect only while `create_enabled` is on. + pub advance_enabled: bool, +} + +/// Result of one table's maintenance attempt, used as the `outcome` metric label. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum PartitionMaintenanceOutcome { + /// Every desired month already has a dedicated leaf. + Noop, + /// DDL is needed but the policy forbids it. + SkippedDisabled, + /// Uncovered months were created without touching a catch-all. + Created, + /// An empty catch-all was replaced by monthlies and a later catch-all. + Advanced, + /// Another session holds the maintenance advisory lock. + SkippedLocked, + /// Automatic DDL would be unsafe; an operator must repair the catalog. + OperatorRequired, + /// A table lock was not available within the lock budget. + LockTimeout, + /// The server statement budget or the client deadline expired. + Deadline, + /// Any other failure. + Error, +} + +impl PartitionMaintenanceOutcome { + /// Stable metric label. + pub const fn as_str(self) -> &'static str { + match self { + Self::Noop => "noop", + Self::SkippedDisabled => "skipped_disabled", + Self::Created => "created", + Self::Advanced => "advanced", + Self::SkippedLocked => "skipped_locked", + Self::OperatorRequired => "operator_required", + Self::LockTimeout => "lock_timeout", + Self::Deadline => "deadline", + Self::Error => "error", + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +struct CatchAllReplacement { + schema: String, + name: String, + new_lower: DateTime, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +struct MaintenanceChange { + /// Ascending month starts whose canonical monthly is created. + create: Vec>, + /// The empty catch-all dropped and recreated after the horizon. + catch_all: Option, +} + +impl MaintenanceChange { + fn outcome(&self) -> PartitionMaintenanceOutcome { + if self.catch_all.is_some() { + PartitionMaintenanceOutcome::Advanced + } else { + PartitionMaintenanceOutcome::Created + } + } + + fn replaced(&self) -> Option<(String, String)> { + self.catch_all + .as_ref() + .map(|catch_all| (catch_all.schema.clone(), catch_all.name.clone())) + } + + fn target_names(&self, table: &str) -> Vec { + let mut names: Vec<_> = self + .create + .iter() + .map(|month| partition_name(table, *month)) + .collect(); + if self.catch_all.is_some() { + names.push(catch_all_name(table)); + } + names + } +} + +#[derive(Debug, PartialEq, Eq)] +enum MaintenancePlan { + Noop, + Disabled, + Refuse(String), + Apply(MaintenanceChange), +} + +enum LockedAttempt { + Committed(MaintenanceChange), + SkippedLocked, + Unchanged, + Disabled, + Refused(String), +} + +struct TableMaintenance { + outcome: PartitionMaintenanceOutcome, + change: Option, + error: Option, +} + +impl TableMaintenance { + fn done(outcome: PartitionMaintenanceOutcome) -> Self { + Self { + outcome, + change: None, + error: None, + } + } + + fn failed( + outcome: PartitionMaintenanceOutcome, + change: Option, + error: String, + ) -> Self { + Self { + outcome, + change, + error: Some(error), + } + } +} + +fn catch_all_name(table: &str) -> String { + format!("{table}_p_future") +} + +fn next_month(start: DateTime) -> Result> { + let (year, month) = add_months(start.year(), start.month(), 1)?; + month_start(year, month) +} + +/// Decide what DDL, if any, would bring `audit` to the desired layout. +/// +/// Pure so the same decision runs before locking, under locks, and in tests. +fn plan_maintenance( + audit: &PartitionTableAudit, + policy: PartitionMaintenancePolicy, +) -> MaintenancePlan { + let uncovered: Vec<_> = audit + .months + .iter() + .filter(|month| month.kind == MonthCoverageKind::Uncovered) + .map(|month| month.start) + .collect(); + let catch_all_covered = audit + .months + .iter() + .any(|month| month.kind == MonthCoverageKind::CoveredByCatchAll); + let create = policy.create_enabled && !uncovered.is_empty(); + // Advancement creates monthlies, so the create kill switch stops it too. + let advance = policy.create_enabled && policy.advance_enabled && catch_all_covered; + if !create && !advance { + return if uncovered.is_empty() && !catch_all_covered { + MaintenancePlan::Noop + } else { + MaintenancePlan::Disabled + }; + } + if !audit.structurally_safe_for_creation() { + return MaintenancePlan::Refuse( + "catalog is not structurally safe for automatic partition DDL".to_string(), + ); + } + + let mut months = if create { uncovered } else { Vec::new() }; + let mut catch_all = None; + if advance { + let catch_alls: Vec<_> = audit + .children + .iter() + .filter(|child| child.kind == PartitionChildKind::CatchAll) + .collect(); + let [child] = catch_alls.as_slice() else { + return MaintenancePlan::Refuse(format!( + "expected exactly one right-edge catch-all, found {}", + catch_alls.len() + )); + }; + match child.catch_all_nonempty { + Some(false) => {} + Some(true) => { + return MaintenancePlan::Refuse(format!( + "catch-all {} contains rows; advancing it requires an operator", + child.name + )) + } + None => { + return MaintenancePlan::Refuse(format!( + "catch-all {} occupancy is unknown", + child.name + )) + } + } + // Replacements inherit only parent triggers; dropping the catch-all + // would silently discard child-specific behavior. + if !child.extra_triggers.is_empty() { + return MaintenancePlan::Refuse(format!( + "catch-all {} has child-only triggers ({}); replacing it would drop them", + child.name, + child.extra_triggers.join(", ") + )); + } + let Some(PartitionBound::Finite(lower)) = child.lower else { + return MaintenancePlan::Refuse(format!( + "catch-all {} has no finite lower bound", + child.name + )); + }; + if month_start(lower.year(), lower.month()).ok() != Some(lower) { + return MaintenancePlan::Refuse(format!( + "catch-all {} lower bound {lower} is not a UTC month start", + child.name + )); + } + let Some(new_lower) = audit + .months + .last() + .and_then(|month| next_month(month.start).ok()) + else { + return MaintenancePlan::Refuse("desired horizon is not representable".to_string()); + }; + let mut month = lower; + while month < new_lower { + if months.len() >= MAX_PARTITION_MONTHS_AHEAD as usize { + return MaintenancePlan::Refuse(format!( + "catch-all {} would require more than {MAX_PARTITION_MONTHS_AHEAD} monthlies", + child.name + )); + } + months.push(month); + let Ok(next) = next_month(month) else { + return MaintenancePlan::Refuse("desired horizon is not representable".to_string()); + }; + month = next; + } + catch_all = Some(CatchAllReplacement { + schema: child.schema.clone(), + name: child.name.clone(), + new_lower, + }); + } + months.sort(); + months.dedup(); + MaintenancePlan::Apply(MaintenanceChange { + create: months, + catch_all, + }) +} + +/// Audit, then issue only the DDL `policy` permits, one bounded transaction per table. +/// +/// A table that is already at the desired layout, or whose required DDL is +/// disabled, takes no lock. One table's failure does not prevent the other +/// table's attempt; failures are aggregated into the returned error. +pub async fn maintain_partitions( + pool: &PgPool, + months_ahead: u32, + policy: PartitionMaintenancePolicy, +) -> Result { + maintain_partitions_at(pool, months_ahead, policy, Utc::now()).await +} + +pub(super) async fn maintain_partitions_at( + pool: &PgPool, + months_ahead: u32, + policy: PartitionMaintenancePolicy, + now: DateTime, +) -> Result { + maintain_partitions_with_deadline(pool, months_ahead, policy, now, MAINTENANCE_TABLE_DEADLINE) + .await +} + +async fn maintain_partitions_with_deadline( + pool: &PgPool, + months_ahead: u32, + policy: PartitionMaintenancePolicy, + now: DateTime, + table_deadline: Duration, +) -> Result { + let audit = audit_partition_catalog_at(pool, months_ahead, now).await?; + maintain_audited(pool, audit, months_ahead, policy, now, table_deadline).await +} + +/// Maintain each table described by `audit`, which may be stale by now. +async fn maintain_audited( + pool: &PgPool, + audit: PartitionAudit, + months_ahead: u32, + policy: PartitionMaintenancePolicy, + now: DateTime, + table_deadline: Duration, +) -> Result { + let horizon = validated_months_ahead(months_ahead)?; + let mut errors = Vec::new(); + // Any table that planned DDL may have found the catalog changed under the + // lock, so the opening audit no longer describes it. + let mut stale = false; + for table_audit in &audit.tables { + let table = table_audit.table; + let plan = plan_maintenance(table_audit, policy); + stale |= matches!(plan, MaintenancePlan::Apply(_)); + let result = maintain_table( + pool, + table_audit, + plan, + horizon, + policy, + now, + table_deadline, + ) + .await; + emit_maintenance_metrics(table_audit, &result); + match (&result.error, result.outcome) { + (Some(error), outcome) => errors.push(format!("{table} {}: {error}", outcome.as_str())), + ( + None, + PartitionMaintenanceOutcome::Created | PartitionMaintenanceOutcome::Advanced, + ) => { + info!( + table, + outcome = result.outcome.as_str(), + partitions = ?result + .change + .as_ref() + .map(|change| change.target_names(table)), + "partition maintenance committed" + ); + } + (None, _) => {} + } + } + + if !errors.is_empty() { + Err(DbError::InvalidData(format!( + "partition maintenance failed: {}", + errors.join("; ") + ))) + } else if stale { + audit_partition_catalog_at(pool, months_ahead, now).await + } else { + Ok(audit) + } +} + +async fn maintain_table( + pool: &PgPool, + audit: &PartitionTableAudit, + plan: MaintenancePlan, + months_ahead: i32, + policy: PartitionMaintenancePolicy, + now: DateTime, + table_deadline: Duration, +) -> TableMaintenance { + let change = match plan { + MaintenancePlan::Noop => return TableMaintenance::done(PartitionMaintenanceOutcome::Noop), + MaintenancePlan::Disabled => { + return TableMaintenance::done(PartitionMaintenanceOutcome::SkippedDisabled) + } + MaintenancePlan::Refuse(reason) => { + // Report a name collision ahead of the structural reason it + // usually causes. A refusal takes no lock, so no race applies. + let (names, replaced) = wanted_names(audit, policy); + let error = match colliding_relation(pool, names, replaced).await { + Ok(Some(name)) => collision_message(&name), + Ok(None) => reason, + Err(error) => { + return TableMaintenance::failed( + classify_error(&error), + None, + error.to_string(), + ) + } + }; + return TableMaintenance::failed( + PartitionMaintenanceOutcome::OperatorRequired, + None, + error, + ); + } + MaintenancePlan::Apply(change) => change, + }; + + let deadline = tokio::time::Instant::now() + table_deadline; + let mut connection = match crate::observability::acquire_writer_until( + pool, + crate::observability::WriterOperation::Maintenance, + deadline, + ) + .await + { + Ok(connection) => connection, + Err(error) => { + let error = DbError::from(error); + return TableMaintenance::failed( + classify_error(&error), + Some(change), + error.to_string(), + ); + } + }; + let attempt = tokio::time::timeout_at( + deadline, + apply_change(&mut connection, audit.table, months_ahead, policy, now), + ) + .await; + let Ok(attempt) = attempt else { + // The abandoned transaction may still be open on the server. Close the + // session instead of returning it to the pool so PostgreSQL rolls back + // whatever did not commit. The deadline can fire after COMMIT was sent, + // so the change may have landed; the next audit reports the truth. The + // backend keeps its locks until its in-flight statement ends, which + // `statement_timeout` bounds. + connection.close_on_drop(); + drop(connection); + return TableMaintenance::failed( + PartitionMaintenanceOutcome::Deadline, + Some(change), + format!( + "exceeded the {}ms maintenance deadline; the attempt was abandoned and the \ + server rolls back any uncommitted work", + table_deadline.as_millis() + ), + ); + }; + match attempt { + Ok(LockedAttempt::Committed(committed)) => TableMaintenance { + outcome: committed.outcome(), + change: Some(committed), + error: None, + }, + Ok(LockedAttempt::SkippedLocked) => { + TableMaintenance::done(PartitionMaintenanceOutcome::SkippedLocked) + } + Ok(LockedAttempt::Unchanged) => TableMaintenance::done(PartitionMaintenanceOutcome::Noop), + Ok(LockedAttempt::Disabled) => { + TableMaintenance::done(PartitionMaintenanceOutcome::SkippedDisabled) + } + Ok(LockedAttempt::Refused(reason)) => TableMaintenance::failed( + PartitionMaintenanceOutcome::OperatorRequired, + Some(change), + reason, + ), + Err(error) => { + TableMaintenance::failed(classify_error(&error), Some(change), error.to_string()) + } + } +} + +async fn apply_change( + connection: &mut PgConnection, + table: &'static str, + months_ahead: i32, + policy: PartitionMaintenancePolicy, + now: DateTime, +) -> Result { + let mut transaction = connection.begin().await?; + configure_maintenance_transaction(&mut transaction).await?; + + let acquired: bool = crate::observability::observe_advisory_lock( + crate::observability::LockType::PartitionMaintenance, + sqlx::query_scalar( + "SELECT pg_catalog.pg_try_advisory_xact_lock(\ + pg_catalog.hashtextextended($1 || pg_catalog.current_schema()::text, 0))", + ) + .bind(MAINTENANCE_LOCK_NAMESPACE) + .fetch_one(&mut *transaction), + ) + .await?; + if !acquired { + transaction.rollback().await?; + return Ok(LockedAttempt::SkippedLocked); + } + + // Re-plan while serialized against other managers but before any table + // lock. A runner that lost a race sees the winner's committed layout here + // and stops, instead of mistaking the winner's monthlies for collisions. + let serialized = audit_table_on(&mut transaction, table, months_ahead, now).await?; + let expected = match plan_maintenance(&serialized, policy) { + MaintenancePlan::Apply(change) => change, + MaintenancePlan::Noop => { + transaction.rollback().await?; + return Ok(LockedAttempt::Unchanged); + } + MaintenancePlan::Disabled => { + transaction.rollback().await?; + return Ok(LockedAttempt::Disabled); + } + MaintenancePlan::Refuse(reason) => { + transaction.rollback().await?; + return Ok(LockedAttempt::Refused(reason)); + } + }; + // Reject misnamed relations before any lock can cost a parent outage. + if let Some(name) = colliding_relation( + &mut *transaction, + expected.target_names(table), + expected.replaced(), + ) + .await? + { + transaction.rollback().await?; + return Ok(LockedAttempt::Refused(collision_message(&name))); + } + + // Refuse an unproven lock set before taking any lock. + let parent = maintenance_parent(&mut transaction, table).await?; + if let Some(reason) = parent.refusal { + transaction.rollback().await?; + return Ok(LockedAttempt::Refused(reason)); + } + if let Some(catch_all) = &expected.catch_all { + if let Some(reason) = catch_all_foreign_key_refusal(&mut transaction, catch_all).await? { + transaction.rollback().await?; + return Ok(LockedAttempt::Refused(reason)); + } + } + // Parent first, in the order writers lock: an insert takes the parent and + // then the foreign-key counterparts its constraint check reads. Waiting on + // a counterpart while holding the parent could deadlock with such a writer. + // ONLY: without it PostgreSQL locks every historical child in turn. + lock( + &mut transaction, + &format!( + "LOCK TABLE ONLY {} IN ACCESS EXCLUSIVE MODE", + qualified_relation_name(&parent.schema, table) + ), + ) + .await?; + // Read the foreign keys again under the parent lock. The read above can + // miss one committed while this transaction waited on the parent, and + // attaching would then wait on its table while holding the parent. Adding + // or dropping a foreign key on or to the parent needs a lock that + // conflicts with this one, so the set cannot change from here on. + let parent = maintenance_parent(&mut transaction, table).await?; + if let Some(reason) = parent.refusal { + transaction.rollback().await?; + return Ok(LockedAttempt::Refused(reason)); + } + // Attaching a partition adds its foreign key under SHARE ROW EXCLUSIVE on + // each referenced table. Taking that mode up front still admits readers and + // the `FOR KEY SHARE` checks of other writers, and fails at once rather + // than queue the counterpart's own writers behind this transaction. + // Dropping the catch-all removes only its inherited foreign key and check + // triggers, which lock the catch-all itself, never the referenced table + // (lock probes on PostgreSQL 16 and 17), so no stronger mode is needed. A + // catch-all with any other foreign key is refused before and after its lock. + for counterpart in &parent.counterparts { + lock( + &mut transaction, + &format!("LOCK TABLE ONLY {counterpart} IN SHARE ROW EXCLUSIVE MODE NOWAIT"), + ) + .await?; + } + // Blocks direct child inserts that bypass the parent. Never wait on a + // child while holding the parent. + if let Some(catch_all) = &expected.catch_all { + lock( + &mut transaction, + &format!( + "LOCK TABLE ONLY {} IN ACCESS EXCLUSIVE MODE NOWAIT", + qualified_relation_name(&catch_all.schema, &catch_all.name) + ), + ) + .await?; + // Foreign-key DDL on or to the catch-all conflicts with this lock but + // not with the parent's, so only this check is authoritative. + if let Some(reason) = catch_all_foreign_key_refusal(&mut transaction, catch_all).await? { + transaction.rollback().await?; + return Ok(LockedAttempt::Refused(reason)); + } + } + + // Close the window between the serialized audit and the table locks. + let locked = audit_table_on(&mut transaction, table, months_ahead, now).await?; + match plan_maintenance(&locked, policy) { + MaintenancePlan::Apply(change) if change == expected => {} + MaintenancePlan::Noop => { + transaction.rollback().await?; + return Ok(LockedAttempt::Unchanged); + } + MaintenancePlan::Disabled => { + transaction.rollback().await?; + return Ok(LockedAttempt::Disabled); + } + MaintenancePlan::Refuse(reason) => { + transaction.rollback().await?; + return Ok(LockedAttempt::Refused(reason)); + } + MaintenancePlan::Apply(_) => { + transaction.rollback().await?; + return Err(DbError::InvalidData( + "partition catalog changed between the audit and the maintenance locks".to_string(), + )); + } + } + + let parent_relation = qualified_relation_name(&parent.schema, table); + if let Some(catch_all) = &expected.catch_all { + execute_ddl( + &mut transaction, + format!( + "DROP TABLE {}", + qualified_relation_name(&catch_all.schema, &catch_all.name) + ), + ) + .await?; + } + for month in &expected.create { + let end = next_month(*month)?; + execute_ddl( + &mut transaction, + format!( + "CREATE TABLE {} PARTITION OF {parent_relation} \ + FOR VALUES FROM ('{}') TO ('{}')", + qualified_relation_name(&parent.schema, &partition_name(table, *month)), + month.format("%Y-%m-%d"), + end.format("%Y-%m-%d"), + ), + ) + .await?; + } + if let Some(catch_all) = &expected.catch_all { + execute_ddl( + &mut transaction, + format!( + "CREATE TABLE {} PARTITION OF {parent_relation} \ + FOR VALUES FROM ('{}') TO (MAXVALUE)", + qualified_relation_name(&parent.schema, &catch_all_name(table)), + catch_all.new_lower.format("%Y-%m-%d"), + ), + ) + .await?; + } + + let after = audit_table_on(&mut transaction, table, months_ahead, now).await?; + let verified = verify_applied(&locked, &after, &expected).and( + verify_inheritance(&mut transaction, parent.oid, &expected.target_names(table)).await?, + ); + if let Err(reason) = verified { + transaction.rollback().await?; + return Err(DbError::InvalidData(format!( + "partition maintenance failed its postcondition and rolled back: {reason}" + ))); + } + transaction.commit().await?; + Ok(LockedAttempt::Committed(expected)) +} + +async fn configure_maintenance_transaction(connection: &mut PgConnection) -> Result<()> { + pin_catalog_rendering(connection).await?; + for (setting, value) in [ + ("lock_timeout", MAINTENANCE_LOCK_TIMEOUT), + ("statement_timeout", MAINTENANCE_STATEMENT_TIMEOUT), + ( + "idle_in_transaction_session_timeout", + MAINTENANCE_IDLE_IN_TRANSACTION_TIMEOUT, + ), + ("application_name", "buzz_partition_maintenance"), + ] { + sqlx::query("SELECT pg_catalog.set_config($1, $2, true)") + .bind(setting) + .bind(value) + .execute(&mut *connection) + .await?; + } + Ok(()) +} + +async fn lock(connection: &mut PgConnection, sql: &str) -> Result<()> { + sqlx::query(sqlx::AssertSqlSafe(sql.to_string())) + .execute(&mut *connection) + .await?; + Ok(()) +} + +async fn execute_ddl(connection: &mut PgConnection, sql: String) -> Result<()> { + sqlx::query(sqlx::AssertSqlSafe(sql)) + .execute(&mut *connection) + .await?; + Ok(()) +} + +struct MaintenanceParent { + oid: i64, + schema: String, + /// Quoted, schema-qualified referenced tables in ascending OID order. + counterparts: Vec, + refusal: Option, +} + +async fn maintenance_parent( + connection: &mut PgConnection, + table: &str, +) -> Result { + let row = sqlx::query( + r#" + SELECT parent.oid::bigint AS parent_oid, + parent_ns.nspname AS parent_schema, + ARRAY( + SELECT pg_catalog.format('%I.%I', referenced_ns.nspname, referenced.relname) + FROM pg_catalog.pg_class referenced + JOIN pg_catalog.pg_namespace referenced_ns + ON referenced_ns.oid = referenced.relnamespace + WHERE referenced.oid IN ( + SELECT foreign_key.confrelid + FROM pg_catalog.pg_constraint foreign_key + WHERE foreign_key.conrelid = parent.oid + AND foreign_key.contype = 'f' + AND foreign_key.confrelid <> parent.oid + ) + ORDER BY referenced.oid + )::text[] AS counterparts, + EXISTS ( + SELECT 1 + FROM pg_catalog.pg_constraint foreign_key + JOIN pg_catalog.pg_class referenced ON referenced.oid = foreign_key.confrelid + WHERE foreign_key.conrelid = parent.oid + AND foreign_key.contype = 'f' + AND (referenced.relkind <> 'r' OR referenced.relhassubclass) + ) AS counterpart_has_children, + EXISTS ( + SELECT 1 + FROM pg_catalog.pg_constraint foreign_key + WHERE foreign_key.confrelid = parent.oid + AND foreign_key.contype = 'f' + ) AS referenced_by_foreign_key + FROM pg_catalog.pg_class parent + JOIN pg_catalog.pg_namespace parent_ns ON parent_ns.oid = parent.relnamespace + WHERE parent_ns.nspname = pg_catalog.current_schema() + AND parent.relname = $1 + "#, + ) + .bind(table) + .fetch_one(&mut *connection) + .await?; + let refusal = if row.try_get::("counterpart_has_children")? { + Some(format!( + "{table} references a partitioned or inherited table; the counterpart lock set is not proven" + )) + } else if row.try_get::("referenced_by_foreign_key")? { + Some(format!( + "{table} is referenced by a foreign key; the counterpart lock set is not proven" + )) + } else { + None + }; + Ok(MaintenanceParent { + oid: row.try_get("parent_oid")?, + schema: row.try_get("parent_schema")?, + counterparts: row.try_get("counterparts")?, + refusal, + }) +} + +/// Names the policy would create, and the catch-all it would replace, even +/// when the plan is refused. Used to report collisions ahead of other refusals. +fn wanted_names( + audit: &PartitionTableAudit, + policy: PartitionMaintenancePolicy, +) -> (Vec, Option<(String, String)>) { + let mut names: Vec<_> = audit + .months + .iter() + .filter(|month| { + (policy.create_enabled && month.kind == MonthCoverageKind::Uncovered) + || (policy.advance_enabled && month.kind == MonthCoverageKind::CoveredByCatchAll) + }) + .map(|month| partition_name(audit.table, month.start)) + .collect(); + let catch_alls: Vec<_> = audit + .children + .iter() + .filter(|child| child.kind == PartitionChildKind::CatchAll) + .collect(); + let advancing = policy.advance_enabled + && audit + .months + .iter() + .any(|month| month.kind == MonthCoverageKind::CoveredByCatchAll); + if advancing { + names.push(catch_all_name(audit.table)); + } + let replaced = match catch_alls.as_slice() { + [child] if advancing => Some((child.schema.clone(), child.name.clone())), + _ => None, + }; + (names, replaced) +} + +/// Return a target name already used by a relation other than the catch-all being replaced. +async fn colliding_relation<'e, E>( + executor: E, + names: Vec, + replaced: Option<(String, String)>, +) -> Result> +where + E: sqlx::PgExecutor<'e>, +{ + let (replaced_schema, replaced_name) = replaced.unwrap_or_default(); + Ok(sqlx::query_scalar( + r#" + SELECT relation.relname::text + FROM pg_catalog.pg_class relation + JOIN pg_catalog.pg_namespace namespace ON namespace.oid = relation.relnamespace + WHERE namespace.nspname = pg_catalog.current_schema() + AND relation.relname = ANY($1) + AND NOT (namespace.nspname = $2 AND relation.relname = $3) + ORDER BY relation.relname + LIMIT 1 + "#, + ) + .bind(names) + .bind(replaced_schema) + .bind(replaced_name) + .fetch_optional(executor) + .await?) +} + +fn collision_message(name: &str) -> String { + format!("canonical name {name} already exists without the expected attachment and bounds") +} + +/// Refuse a catch-all with a foreign key of its own, or one that references it. +/// +/// Such a key is invisible to the parent's counterpart set. Dropping the +/// catch-all would silently drop it and lock its other table without `NOWAIT` +/// while holding the parent. The referenced arm also matches the catch-all's +/// clone of a key referencing the parent; the parent refusal runs first, so +/// that arm is a backstop. +async fn catch_all_foreign_key_refusal( + connection: &mut PgConnection, + catch_all: &CatchAllReplacement, +) -> Result> { + let names: Option = sqlx::query_scalar( + r#" + SELECT pg_catalog.string_agg(foreign_key.conname::text, ', ' ORDER BY foreign_key.conname) + FROM pg_catalog.pg_constraint foreign_key + WHERE foreign_key.contype = 'f' + AND ((foreign_key.conrelid = $1::regclass AND foreign_key.conparentid = 0) + OR foreign_key.confrelid = $1::regclass) + "#, + ) + .bind(qualified_relation_name(&catch_all.schema, &catch_all.name)) + .fetch_one(&mut *connection) + .await?; + Ok(names.map(|names| { + format!( + "catch-all {} has foreign keys outside the parent's counterpart set ({names}); \ + the counterpart lock set is not proven", + catch_all.name + ) + })) +} + +/// Verify the changed catalog, audited inside the transaction, before commit. +fn verify_applied( + before: &PartitionTableAudit, + after: &PartitionTableAudit, + change: &MaintenanceChange, +) -> std::result::Result<(), String> { + let table = after.table; + if !after.structurally_safe_for_creation() { + return Err("catalog is not structurally safe after DDL".to_string()); + } + if before.serving_safe && !after.serving_safe { + return Err("current month lost its covering partition".to_string()); + } + let parity = |name: &str| { + after + .children + .iter() + .find(|child| child.name == name) + .filter(|child| child.missing_triggers.is_empty() && child.extra_triggers.is_empty()) + }; + for month in &change.create { + let name = partition_name(table, *month); + let end = next_month(*month).map_err(|error| error.to_string())?; + if !parity(&name).is_some_and(|child| { + child.kind == PartitionChildKind::CanonicalMonthly + && child.lower == Some(PartitionBound::Finite(*month)) + && child.upper == Some(PartitionBound::Finite(end)) + }) { + return Err(format!( + "{name} is missing, misbounded, or lacks trigger parity" + )); + } + } + if let Some(catch_all) = &change.catch_all { + let name = catch_all_name(table); + if !parity(&name).is_some_and(|child| { + child.kind == PartitionChildKind::CatchAll + && child.lower == Some(PartitionBound::Finite(catch_all.new_lower)) + && child.catch_all_nonempty == Some(false) + }) { + return Err(format!( + "{name} is missing, misbounded, or lacks trigger parity" + )); + } + if after + .children + .iter() + .filter(|child| child.kind == PartitionChildKind::CatchAll) + .count() + != 1 + { + return Err("expected exactly one catch-all after DDL".to_string()); + } + } + if let Some(month) = after.months.iter().find(|month| { + change.create.contains(&month.start) && month.kind != MonthCoverageKind::CoveredByMonthly + }) { + return Err(format!( + "{} is not covered by its dedicated monthly", + month.start.format("%Y-%m") + )); + } + Ok(()) +} + +/// Verify every new child carries a valid, ready clone of each parent index and +/// a validated clone of each parent primary-key, unique, and foreign-key constraint. +async fn verify_inheritance( + connection: &mut PgConnection, + parent_oid: i64, + names: &[String], +) -> Result> { + let rows = sqlx::query( + r#" + SELECT child.relname::text AS name, + ( + SELECT count(*) + FROM pg_catalog.pg_index parent_index + WHERE parent_index.indrelid = $1::bigint::oid + AND NOT EXISTS ( + SELECT 1 + FROM pg_catalog.pg_inherits index_edge + JOIN pg_catalog.pg_index child_index + ON child_index.indexrelid = index_edge.inhrelid + WHERE index_edge.inhparent = parent_index.indexrelid + AND child_index.indrelid = child.oid + AND child_index.indisvalid + AND child_index.indisready + ) + ) AS missing_indexes, + ( + SELECT count(*) + FROM pg_catalog.pg_constraint parent_constraint + WHERE parent_constraint.conrelid = $1::bigint::oid + AND parent_constraint.contype IN ('f', 'p', 'u') + AND NOT EXISTS ( + SELECT 1 + FROM pg_catalog.pg_constraint child_constraint + WHERE child_constraint.conrelid = child.oid + AND child_constraint.conparentid = parent_constraint.oid + AND child_constraint.convalidated + ) + ) AS missing_constraints + FROM pg_catalog.pg_class child + JOIN pg_catalog.pg_namespace child_ns ON child_ns.oid = child.relnamespace + JOIN pg_catalog.pg_inherits edge + ON edge.inhrelid = child.oid AND edge.inhparent = $1::bigint::oid + WHERE child_ns.nspname = pg_catalog.current_schema() + AND child.relname = ANY($2) + "#, + ) + .bind(parent_oid) + .bind(names) + .fetch_all(&mut *connection) + .await?; + if rows.len() != names.len() { + return Ok(Err(format!( + "expected {} new children attached to the parent, found {}", + names.len(), + rows.len() + ))); + } + for row in rows { + let name: String = row.try_get("name")?; + let missing_indexes: i64 = row.try_get("missing_indexes")?; + let missing_constraints: i64 = row.try_get("missing_constraints")?; + if missing_indexes != 0 || missing_constraints != 0 { + return Ok(Err(format!( + "{name} is missing {missing_indexes} inherited indexes and \ + {missing_constraints} inherited constraints" + ))); + } + } + Ok(Ok(())) +} + +fn classify_error(error: &DbError) -> PartitionMaintenanceOutcome { + match error { + DbError::Sqlx(sqlx::Error::Database(database)) => match database.code().as_deref() { + // A deadlock victim lost a lock race; the next pass retries. + Some("55P03" | "40P01") => PartitionMaintenanceOutcome::LockTimeout, + Some("57014" | "25P03") => PartitionMaintenanceOutcome::Deadline, + _ => PartitionMaintenanceOutcome::Error, + }, + DbError::Sqlx(sqlx::Error::PoolTimedOut) => PartitionMaintenanceOutcome::Deadline, + _ => PartitionMaintenanceOutcome::Error, + } +} + +fn emit_maintenance_metrics(audit: &PartitionTableAudit, result: &TableMaintenance) { + let table = audit.table; + metrics::counter!( + "buzz_partition_maintenance_runs_total", + "table" => table, + "outcome" => result.outcome.as_str() + ) + .increment(1); + // Per-month creation accounting, continuous with the pre-advancement manager. + let changed = result + .change + .as_ref() + .map(|change| change.create.as_slice()) + .unwrap_or_default(); + // `error` means DDL or collision failure. Lock and deadline outcomes are + // transient and only `buzz_partition_maintenance_runs_total` records them. + let month_outcome = match result.outcome { + PartitionMaintenanceOutcome::Error | PartitionMaintenanceOutcome::OperatorRequired => { + Some("error") + } + PartitionMaintenanceOutcome::LockTimeout | PartitionMaintenanceOutcome::Deadline => None, + _ => Some("created"), + }; + for month in &audit.months { + let outcome = if changed.contains(&month.start) { + let Some(outcome) = month_outcome else { + continue; + }; + outcome + } else if month.kind == MonthCoverageKind::Uncovered { + continue; + } else { + "skipped_covered" + }; + metrics::counter!( + "buzz_partition_create_attempts_total", + "table" => table, + "outcome" => outcome + ) + .increment(1); + } +} + +#[cfg(test)] +#[path = "maintenance_tests.rs"] +mod tests; diff --git a/crates/buzz-db/src/store/partition/maintenance_tests.rs b/crates/buzz-db/src/store/partition/maintenance_tests.rs new file mode 100644 index 0000000000..99a1684814 --- /dev/null +++ b/crates/buzz-db/src/store/partition/maintenance_tests.rs @@ -0,0 +1,1330 @@ +use chrono::TimeZone; + +use super::*; +use crate::store::partition::{MonthCoverage, PartitionChildAudit, PARTITIONED_TABLES}; + +fn month(year: i32, month: u32) -> DateTime { + Utc.with_ymd_and_hms(year, month, 1, 0, 0, 0).unwrap() +} + +fn child( + name: &str, + kind: PartitionChildKind, + lower: PartitionBound, + upper: PartitionBound, + catch_all_nonempty: Option, +) -> PartitionChildAudit { + PartitionChildAudit { + schema: "public".to_string(), + name: name.to_string(), + relation_kind: "r".to_string(), + lower: Some(lower), + upper: Some(upper), + kind, + pending_detach: false, + catch_all_nonempty, + default_nonempty: None, + missing_triggers: Vec::new(), + extra_triggers: Vec::new(), + } +} + +/// The repaired production layout: a July/August legacy leaf, September through +/// December monthlies, and a catch-all from January 2027. +fn repaired_audit(catch_all_nonempty: bool) -> PartitionTableAudit { + let mut children = vec![ + child( + "events_p2026_07_08_legacy", + PartitionChildKind::LegacyLeaf, + PartitionBound::Finite(month(2026, 7)), + PartitionBound::Finite(month(2026, 9)), + None, + ), + child( + "events_p_future_next", + PartitionChildKind::CatchAll, + PartitionBound::Finite(month(2027, 1)), + PartitionBound::MaxValue, + Some(catch_all_nonempty), + ), + ]; + for m in 9..=12 { + children.push(child( + &format!("events_p2026_{m:02}"), + PartitionChildKind::CanonicalMonthly, + PartitionBound::Finite(month(2026, m)), + PartitionBound::Finite({ + let (year, next) = add_months(2026, m, 1).unwrap(); + month(year, next) + }), + None, + )); + } + let months = (0..=PARTITION_MANAGER_MONTHS_AHEAD as i32) + .map(|offset| { + let (year, m) = add_months(2026, 9, offset).unwrap(); + let start = month(year, m); + MonthCoverage { + start, + kind: if start < month(2027, 1) { + MonthCoverageKind::CoveredByMonthly + } else { + MonthCoverageKind::CoveredByCatchAll + }, + } + }) + .collect(); + PartitionTableAudit { + table: "events", + partition_key: Some("RANGE (created_at)".to_string()), + expected_partition_key: "RANGE (created_at)", + partition_key_valid: true, + children, + coverage_leaves: Vec::new(), + months, + serving_safe: true, + } +} + +const ADVANCE: PartitionMaintenancePolicy = PartitionMaintenancePolicy { + create_enabled: true, + advance_enabled: true, +}; + +const CREATE_ONLY: PartitionMaintenancePolicy = PartitionMaintenancePolicy { + create_enabled: true, + advance_enabled: false, +}; + +#[test] +fn repaired_layout_plans_three_monthlies_and_an_april_catch_all() { + assert_eq!( + plan_maintenance(&repaired_audit(false), ADVANCE), + MaintenancePlan::Apply(MaintenanceChange { + create: vec![month(2027, 1), month(2027, 2), month(2027, 3)], + catch_all: Some(CatchAllReplacement { + schema: "public".to_string(), + name: "events_p_future_next".to_string(), + new_lower: month(2027, 4), + }), + }) + ); +} + +#[test] +fn advancement_requires_its_own_policy_flag() { + assert_eq!( + plan_maintenance(&repaired_audit(false), CREATE_ONLY), + MaintenancePlan::Disabled + ); + assert_eq!( + plan_maintenance( + &repaired_audit(false), + PartitionMaintenancePolicy::default() + ), + MaintenancePlan::Disabled + ); +} + +#[test] +fn full_monthly_runway_is_a_noop_even_when_disabled() { + let mut audit = repaired_audit(false); + for coverage in &mut audit.months { + coverage.kind = MonthCoverageKind::CoveredByMonthly; + } + assert_eq!(plan_maintenance(&audit, ADVANCE), MaintenancePlan::Noop); + assert_eq!( + plan_maintenance(&audit, PartitionMaintenancePolicy::default()), + MaintenancePlan::Noop + ); +} + +#[test] +fn unsafe_catch_all_states_require_an_operator() { + let refused = |audit: &PartitionTableAudit, needle: &str| match plan_maintenance(audit, ADVANCE) + { + MaintenancePlan::Refuse(reason) => assert!(reason.contains(needle), "{reason}"), + plan => panic!("expected refusal containing {needle:?}, got {plan:?}"), + }; + + refused(&repaired_audit(true), "contains rows"); + + let mut unknown = repaired_audit(false); + unknown.children[1].catch_all_nonempty = None; + refused(&unknown, "occupancy is unknown"); + + let mut misaligned = repaired_audit(false); + misaligned.children[1].lower = Some(PartitionBound::Finite( + Utc.with_ymd_and_hms(2027, 1, 15, 0, 0, 0).unwrap(), + )); + refused(&misaligned, "not a UTC month start"); + + let mut pending = repaired_audit(false); + pending.children[1].pending_detach = true; + refused(&pending, "not structurally safe"); + + let mut child_triggers = repaired_audit(false); + child_triggers.children[1].extra_triggers = vec!["audit_catch_all".to_string()]; + refused(&child_triggers, "child-only triggers (audit_catch_all)"); +} + +#[test] +fn advancement_crosses_the_year_boundary() { + // December: the horizon runs through June and the catch-all moves to July. + let mut audit = repaired_audit(false); + audit.months = (0..=PARTITION_MANAGER_MONTHS_AHEAD as i32) + .map(|offset| { + let (year, m) = add_months(2026, 12, offset).unwrap(); + MonthCoverage { + start: month(year, m), + kind: if offset == 0 { + MonthCoverageKind::CoveredByMonthly + } else { + MonthCoverageKind::CoveredByCatchAll + }, + } + }) + .collect(); + let MaintenancePlan::Apply(change) = plan_maintenance(&audit, ADVANCE) else { + panic!("expected advancement"); + }; + assert_eq!(change.create.first(), Some(&month(2027, 1))); + assert_eq!(change.create.last(), Some(&month(2027, 6))); + assert_eq!(change.create.len(), 6); + assert_eq!(change.catch_all.map(|c| c.new_lower), Some(month(2027, 7))); +} + +#[test] +fn create_kill_switch_also_stops_advancement() { + let advance_only = PartitionMaintenancePolicy { + create_enabled: false, + advance_enabled: true, + }; + assert_eq!( + plan_maintenance(&repaired_audit(false), advance_only), + MaintenancePlan::Disabled + ); +} + +#[test] +fn advancement_is_capped_at_the_supported_horizon() { + // The horizon ends in March 2027, so a catch-all from April 2017 needs + // exactly the maximum number of monthlies and one from March 2017 needs one more. + let from = |lower| { + let mut audit = repaired_audit(false); + audit.children[1].lower = Some(PartitionBound::Finite(lower)); + plan_maintenance(&audit, ADVANCE) + }; + let MaintenancePlan::Apply(change) = from(month(2017, 4)) else { + panic!("expected advancement at the cap"); + }; + assert_eq!(change.create.len(), MAX_PARTITION_MONTHS_AHEAD as usize); + match from(month(2017, 3)) { + MaintenancePlan::Refuse(reason) => assert!(reason.contains("more than"), "{reason}"), + plan => panic!("expected refusal past the cap, got {plan:?}"), + } +} + +/// `buzz_partition_create_attempts_total` counts keyed by their `outcome` label. +fn create_attempts(result: &TableMaintenance) -> std::collections::BTreeMap { + let recorder = metrics_util::debugging::DebuggingRecorder::new(); + let snapshotter = recorder.snapshotter(); + metrics::with_local_recorder(&recorder, || { + emit_maintenance_metrics(&repaired_audit(false), result) + }); + snapshotter + .snapshot() + .into_vec() + .into_iter() + .filter(|(key, ..)| key.key().name() == "buzz_partition_create_attempts_total") + .map(|(key, _, _, value)| { + let outcome = key + .key() + .labels() + .find(|label| label.key() == "outcome") + .map(|label| label.value().to_string()) + .expect("outcome label"); + let metrics_util::debugging::DebugValue::Counter(count) = value else { + panic!("create attempts must be a counter"); + }; + (outcome, count) + }) + .collect() +} + +#[test] +fn only_ddl_failures_count_as_create_errors() { + let MaintenancePlan::Apply(change) = plan_maintenance(&repaired_audit(false), ADVANCE) else { + panic!("expected advancement"); + }; + let failed = |outcome| TableMaintenance::failed(outcome, Some(change.clone()), String::new()); + let covered = ("skipped_covered".to_string(), 4); + + for outcome in [ + PartitionMaintenanceOutcome::LockTimeout, + PartitionMaintenanceOutcome::Deadline, + ] { + assert_eq!( + create_attempts(&failed(outcome)), + [covered.clone()].into(), + "{outcome:?}" + ); + } + for outcome in [ + PartitionMaintenanceOutcome::Error, + PartitionMaintenanceOutcome::OperatorRequired, + ] { + assert_eq!( + create_attempts(&failed(outcome)), + [covered.clone(), ("error".to_string(), 3)].into(), + "{outcome:?}" + ); + } + let advanced = TableMaintenance { + outcome: PartitionMaintenanceOutcome::Advanced, + change: Some(change.clone()), + error: None, + }; + assert_eq!( + create_attempts(&advanced), + [covered, ("created".to_string(), 3)].into() + ); +} + +#[test] +fn outcome_labels_are_stable() { + use PartitionMaintenanceOutcome::*; + assert_eq!( + [ + Noop, + SkippedDisabled, + Created, + Advanced, + SkippedLocked, + OperatorRequired, + LockTimeout, + Deadline, + Error + ] + .map(PartitionMaintenanceOutcome::as_str), + [ + "noop", + "skipped_disabled", + "created", + "advanced", + "skipped_locked", + "operator_required", + "lock_timeout", + "deadline", + "error" + ] + ); +} + +mod postgres_tests { + use std::time::Instant; + + use sqlx::PgPool; + + use super::*; + use crate::store::partition::tests::postgres_tests::{ + catalog_snapshot, create_child, drop_schema, fixed_now, scratch_pool, + scratch_pool_with_max_connections, seed_parents, + }; + + /// The repaired incident layout for both managed parents, optionally with + /// outgoing foreign keys to a `communities` counterpart. + async fn seed_repaired_layout(pool: &PgPool, with_foreign_keys: bool) { + seed_parents(pool).await; + if with_foreign_keys { + sqlx::query("CREATE TABLE communities (id UUID PRIMARY KEY)") + .execute(pool) + .await + .expect("create communities"); + for table in PARTITIONED_TABLES { + sqlx::query(sqlx::AssertSqlSafe(format!( + "ALTER TABLE {table} ADD COLUMN community_id UUID REFERENCES communities(id)" + ))) + .execute(pool) + .await + .expect("add counterpart foreign key"); + } + } + for table in PARTITIONED_TABLES { + create_child( + pool, + table, + &format!("{table}_p_past"), + "MINVALUE", + "'2026-07-01'", + ) + .await; + create_child( + pool, + table, + &format!("{table}_p2026_07_08_legacy"), + "'2026-07-01'", + "'2026-09-01'", + ) + .await; + for m in 9..=12 { + let upper = if m == 12 { + "'2027-01-01'".to_string() + } else { + format!("'2026-{:02}-01'", m + 1) + }; + create_child( + pool, + table, + &format!("{table}_p2026_{m:02}"), + &format!("'2026-{m:02}-01'"), + &upper, + ) + .await; + } + create_child( + pool, + table, + &format!("{table}_p_future_next"), + "'2027-01-01'", + "MAXVALUE", + ) + .await; + } + } + + async fn maintain(pool: &PgPool) -> Result { + maintain_partitions_at(pool, PARTITION_MANAGER_MONTHS_AHEAD, ADVANCE, fixed_now()).await + } + + fn assert_advanced(audit: &PartitionAudit) { + for table in &audit.tables { + assert!( + table + .months + .iter() + .all(|month| month.kind == MonthCoverageKind::CoveredByMonthly), + "{}: {:?}", + table.table, + table.months + ); + let catch_all = table + .children + .iter() + .find(|child| child.kind == PartitionChildKind::CatchAll) + .expect("catch-all"); + assert_eq!(catch_all.name, format!("{}_p_future", table.table)); + assert_eq!( + catch_all.lower, + Some(PartitionBound::Finite(month(2027, 4))) + ); + assert_eq!(catch_all.catch_all_nonempty, Some(false)); + assert!(!table + .children + .iter() + .any(|child| child.name.ends_with("_p_future_next"))); + assert_eq!(table.missing_trigger_count(), 0); + assert_eq!(table.extra_trigger_count(), 0); + } + } + + /// Open a transaction on a separate session that holds `sql`'s lock until dropped. + async fn hold_lock(pool: &PgPool, sql: &str) -> sqlx::Transaction<'static, sqlx::Postgres> { + let mut holder = pool.begin().await.expect("begin lock holder"); + sqlx::query(sqlx::AssertSqlSafe(sql.to_string())) + .execute(&mut *holder) + .await + .expect("take conflicting lock"); + holder + } + + #[tokio::test] + #[ignore = "requires Postgres"] + async fn empty_catch_all_advances_to_canonical_layout_and_is_idempotent() { + let (pool, admin, schema) = scratch_pool().await; + seed_repaired_layout(&pool, true).await; + + let audit = maintain(&pool).await.expect("advance"); + assert_advanced(&audit); + let fresh = audit_partition_catalog_at(&pool, PARTITION_MANAGER_MONTHS_AHEAD, fixed_now()) + .await + .expect("fresh audit"); + assert_advanced(&fresh); + + // Rows route to the new leaves, including past the new horizon. + for (created_at, leaf) in [ + ("2027-02-10", "events_p2027_02"), + ("2031-01-01", "events_p_future"), + ] { + let routed: String = sqlx::query_scalar(sqlx::AssertSqlSafe(format!( + "WITH inserted AS (INSERT INTO events (created_at, alternate_at) \ + VALUES ('{created_at}', now()) RETURNING tableoid) \ + SELECT tableoid::regclass::text FROM inserted" + ))) + .fetch_one(&pool) + .await + .expect("route row"); + assert_eq!(routed, leaf); + } + sqlx::query("DELETE FROM events WHERE created_at >= '2027-01-01'") + .execute(&pool) + .await + .expect("clear routed rows"); + + let before = catalog_snapshot(&pool).await; + maintain(&pool).await.expect("idempotent rerun"); + assert_eq!(catalog_snapshot(&pool).await, before); + drop_schema(&admin, &schema).await; + } + + #[tokio::test] + #[ignore = "requires Postgres"] + async fn canonical_catch_all_name_is_replaced_in_one_transaction() { + let (pool, admin, schema) = scratch_pool().await; + seed_parents(&pool).await; + for table in PARTITIONED_TABLES { + create_child( + &pool, + table, + &format!("{table}_p_past"), + "MINVALUE", + "'2026-09-01'", + ) + .await; + create_child( + &pool, + table, + &format!("{table}_p_future"), + "'2026-09-01'", + "MAXVALUE", + ) + .await; + } + let audit = maintain(&pool).await.expect("advance fresh layout"); + assert_advanced(&audit); + assert!(audit.tables.iter().all(|table| table + .children + .iter() + .any(|child| child.name == format!("{}_p2026_09", table.table)))); + drop_schema(&admin, &schema).await; + } + + #[tokio::test] + #[ignore = "requires Postgres"] + async fn populated_catch_all_is_refused_before_any_lock() { + let (pool, admin, schema) = scratch_pool().await; + seed_repaired_layout(&pool, false).await; + sqlx::query("INSERT INTO events (created_at, alternate_at) VALUES ('2027-02-01', now())") + .execute(&pool) + .await + .expect("populate catch-all"); + let before = catalog_snapshot(&pool).await; + // A lock-taking attempt would wait out the 2s parent lock timeout. + let holder = hold_lock(&pool, "LOCK TABLE ONLY events IN ACCESS SHARE MODE").await; + + let started = Instant::now(); + let result = maintain(&pool).await; + let elapsed = started.elapsed(); + drop(holder); + + assert!( + matches!(result, Err(DbError::InvalidData(ref message)) + if message.contains("events operator_required") + && message.contains("events_p_future_next contains rows") + && !message.contains("delivery_log")), + "{result:?}" + ); + assert!(elapsed < std::time::Duration::from_secs(1), "{elapsed:?}"); + let after = catalog_snapshot(&pool).await; + let events_only = |snapshot: &[(String, String, String, i64)]| { + snapshot + .iter() + .filter(|(name, ..)| name.starts_with("events_")) + .cloned() + .collect::>() + }; + assert_eq!(events_only(&after), events_only(&before)); + drop_schema(&admin, &schema).await; + } + + #[tokio::test] + #[ignore = "requires Postgres"] + async fn counterpart_contention_fails_fast_and_rolls_back() { + let (pool, admin, schema) = scratch_pool().await; + seed_repaired_layout(&pool, true).await; + let before = catalog_snapshot(&pool).await; + let holder = hold_lock( + &pool, + "INSERT INTO communities(id) VALUES (gen_random_uuid())", + ) + .await; + + let started = Instant::now(); + let result = maintain(&pool).await; + let elapsed = started.elapsed(); + drop(holder); + + assert!( + matches!(result, Err(DbError::InvalidData(ref message)) + if message.contains("events lock_timeout") + && message.contains("delivery_log lock_timeout")), + "{result:?}" + ); + // NOWAIT: no 2s wait per table while a counterpart writer is active. + assert!(elapsed < std::time::Duration::from_secs(1), "{elapsed:?}"); + assert_eq!(catalog_snapshot(&pool).await, before); + + assert_advanced(&maintain(&pool).await.expect("advance after release")); + drop_schema(&admin, &schema).await; + } + + #[tokio::test] + #[ignore = "requires Postgres"] + async fn counterpart_reader_does_not_block_advancement() { + let (pool, admin, schema) = scratch_pool().await; + seed_repaired_layout(&pool, true).await; + let holder = hold_lock(&pool, "SELECT count(*) FROM communities").await; + + let result = maintain(&pool).await; + drop(holder); + + assert_advanced(&result.expect("advance beside a counterpart reader")); + drop_schema(&admin, &schema).await; + } + + #[tokio::test] + #[ignore = "requires Postgres"] + async fn writer_checking_a_foreign_key_behind_maintenance_does_not_deadlock() { + let (pool, admin, schema) = scratch_pool().await; + seed_repaired_layout(&pool, true).await; + // An in-flight insert: the parent first, then the foreign-key check. + let mut writer = hold_lock(&pool, "LOCK TABLE ONLY events IN ROW EXCLUSIVE MODE").await; + let writer_pid: i32 = sqlx::query_scalar("SELECT pg_backend_pid()") + .fetch_one(&mut *writer) + .await + .expect("writer backend pid"); + + let maintenance = tokio::spawn({ + let pool = pool.clone(); + async move { maintain(&pool).await } + }); + let mut queued = false; + for _ in 0..50 { + queued = sqlx::query_scalar( + "SELECT EXISTS (SELECT 1 FROM pg_stat_activity \ + WHERE $1 = ANY(pg_blocking_pids(pid)))", + ) + .bind(writer_pid) + .fetch_one(&admin) + .await + .expect("check queued maintenance"); + if queued { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + } + assert!(queued, "maintenance never queued behind the writer"); + + // Maintenance waits on the parent holding no counterpart lock, so the + // writer's key check, and a write to the counterpart itself, proceed + // instead of entering a lock cycle in any counterpart lock mode. + let started = Instant::now(); + sqlx::query("SELECT id FROM communities FOR KEY SHARE") + .fetch_all(&mut *writer) + .await + .expect("foreign-key check behind queued maintenance"); + sqlx::query("INSERT INTO communities(id) VALUES (gen_random_uuid())") + .execute(&mut *writer) + .await + .expect("counterpart write behind queued maintenance"); + assert!( + started.elapsed() < std::time::Duration::from_millis(500), + "{:?}", + started.elapsed() + ); + writer.commit().await.expect("commit writer"); + + assert_advanced( + &maintenance + .await + .expect("maintenance task") + .expect("advance after the writer commits"), + ); + drop_schema(&admin, &schema).await; + } + + async fn wait_until_blocked_by(admin: &PgPool, pid: i32, waiters: i64) { + for _ in 0..100 { + let blocked: i64 = sqlx::query_scalar( + "SELECT count(*) FROM pg_stat_activity WHERE $1 = ANY(pg_blocking_pids(pid))", + ) + .bind(pid) + .fetch_one(admin) + .await + .expect("count blocked backends"); + if blocked >= waiters { + return; + } + tokio::time::sleep(std::time::Duration::from_millis(20)).await; + } + panic!("{waiters} backends never queued behind {pid}"); + } + + #[tokio::test] + #[ignore = "requires Postgres"] + async fn foreign_key_committed_while_waiting_on_the_parent_is_locked_nowait() { + let (pool, admin, schema) = scratch_pool_with_max_connections(6).await; + seed_repaired_layout(&pool, true).await; + sqlx::query("CREATE TABLE reviewers (id UUID PRIMARY KEY)") + .execute(&pool) + .await + .expect("create new counterpart"); + // An in-flight migration adds a foreign key that maintenance's + // pre-lock read cannot see yet. + let mut migration = hold_lock( + &pool, + "ALTER TABLE events ADD FOREIGN KEY (community_id) REFERENCES reviewers(id)", + ) + .await; + let migration_pid: i32 = sqlx::query_scalar("SELECT pg_backend_pid()") + .fetch_one(&mut *migration) + .await + .expect("migration backend pid"); + // A writer on the new counterpart queues behind the migration and is + // granted its lock when the migration commits. + let (release_writer, writer_released) = tokio::sync::oneshot::channel::<()>(); + let writer = tokio::spawn({ + let pool = pool.clone(); + async move { + let mut writer = pool.begin().await.expect("begin counterpart writer"); + sqlx::query("INSERT INTO reviewers(id) VALUES (gen_random_uuid())") + .execute(&mut *writer) + .await + .expect("counterpart writer insert"); + let _ = writer_released.await; + writer + .rollback() + .await + .expect("roll back counterpart writer"); + } + }); + wait_until_blocked_by(&admin, migration_pid, 1).await; + let maintenance = tokio::spawn({ + let pool = pool.clone(); + async move { maintain(&pool).await } + }); + wait_until_blocked_by(&admin, migration_pid, 2).await; + + migration.commit().await.expect("commit migration"); + let result = maintenance.await.expect("maintenance task"); + + // The parent-locked read finds the new counterpart and its NOWAIT lock + // fails naming it, instead of attaching against it under the parent + // lock until a lock timeout that names no relation. + assert!( + matches!(result, Err(DbError::InvalidData(ref message)) + if message.contains("events lock_timeout") + && message.contains("could not obtain lock on relation") + && message.contains(r#".reviewers""#) + && !message.contains("delivery_log")), + "{result:?}" + ); + + release_writer.send(()).expect("release counterpart writer"); + writer.await.expect("counterpart writer task"); + assert_advanced(&maintain(&pool).await.expect("advance after release")); + drop_schema(&admin, &schema).await; + } + + async fn backend_pid(connection: &mut PgConnection) -> i32 { + sqlx::query_scalar("SELECT pg_backend_pid()") + .fetch_one(connection) + .await + .expect("backend pid") + } + + fn refused_events_only(result: &Result, reason: &str) -> bool { + matches!(result, Err(DbError::InvalidData(message)) + if message.contains("events operator_required") + && message.contains(reason) + && !message.contains("delivery_log")) + } + + #[tokio::test] + #[ignore = "requires Postgres"] + async fn foreign_key_referencing_the_parent_committed_while_waiting_is_refused() { + let (pool, admin, schema) = scratch_pool_with_max_connections(6).await; + seed_repaired_layout(&pool, true).await; + sqlx::query( + "CREATE TABLE event_refs \ + (created_at TIMESTAMPTZ, alternate_at TIMESTAMPTZ, id BIGINT)", + ) + .execute(&pool) + .await + .expect("create referencing table"); + // An in-flight migration makes events a referenced table, which the + // pre-lock read cannot see yet. + let mut migration = hold_lock( + &pool, + "ALTER TABLE event_refs ADD FOREIGN KEY (created_at, alternate_at, id) \ + REFERENCES events (created_at, alternate_at, id)", + ) + .await; + let migration_pid = backend_pid(&mut migration).await; + let maintenance = tokio::spawn({ + let pool = pool.clone(); + async move { maintain(&pool).await } + }); + wait_until_blocked_by(&admin, migration_pid, 1).await; + migration.commit().await.expect("commit migration"); + + let result = maintenance.await.expect("maintenance task"); + assert!( + refused_events_only(&result, "events is referenced by a foreign key"), + "{result:?}" + ); + drop_schema(&admin, &schema).await; + } + + #[tokio::test] + #[ignore = "requires Postgres"] + async fn catch_all_foreign_keys_outside_the_parent_are_refused() { + let (pool, admin, schema) = scratch_pool().await; + seed_repaired_layout(&pool, true).await; + sqlx::query("CREATE TABLE reviewers (id UUID PRIMARY KEY)") + .execute(&pool) + .await + .expect("create counterpart"); + let keys = |pool: &PgPool| { + let pool = pool.clone(); + async move { + sqlx::query_scalar::<_, i64>( + "SELECT count(*) FROM pg_constraint WHERE contype = 'f' \ + AND (conrelid = 'events_p_future_next'::regclass AND conparentid = 0 \ + OR confrelid = 'events_p_future_next'::regclass)", + ) + .fetch_one(&pool) + .await + .expect("count catch-all keys") + } + }; + // Refused before any lock: with the parent held, a run that reached + // the parent lock would report lock_timeout instead. + let holder = hold_lock(&pool, "LOCK TABLE ONLY events IN ACCESS SHARE MODE").await; + for (add, remove) in [ + ( + "ALTER TABLE ONLY events_p_future_next ADD CONSTRAINT catch_all_reviewer \ + FOREIGN KEY (community_id) REFERENCES reviewers(id)", + "ALTER TABLE events_p_future_next DROP CONSTRAINT catch_all_reviewer", + ), + ( + "CREATE TABLE catch_all_refs (created_at TIMESTAMPTZ, alternate_at TIMESTAMPTZ, \ + id BIGINT, CONSTRAINT catch_all_reviewer FOREIGN KEY (created_at, alternate_at, id) \ + REFERENCES events_p_future_next (created_at, alternate_at, id))", + "DROP TABLE catch_all_refs", + ), + ] { + sqlx::query(sqlx::AssertSqlSafe(add.to_string())) + .execute(&pool) + .await + .expect("add catch-all key"); + let result = maintain(&pool).await; + assert!( + refused_events_only(&result, "catch_all_reviewer"), + "{add}: {result:?}" + ); + // Refused, not dropped along with the catch-all. + assert_eq!(keys(&pool).await, 1, "{add}"); + sqlx::query(sqlx::AssertSqlSafe(remove.to_string())) + .execute(&pool) + .await + .expect("remove catch-all key"); + } + drop(holder); + assert_advanced(&maintain(&pool).await.expect("advance without the keys")); + drop_schema(&admin, &schema).await; + } + + #[tokio::test] + #[ignore = "requires Postgres"] + async fn catch_all_foreign_key_committed_while_waiting_on_the_parent_is_refused() { + let (pool, admin, schema) = scratch_pool_with_max_connections(6).await; + seed_repaired_layout(&pool, true).await; + sqlx::query("CREATE TABLE reviewers (id UUID PRIMARY KEY)") + .execute(&pool) + .await + .expect("create counterpart"); + // The catch-all's own key conflicts with no parent lock, so a reader + // holds the parent to keep maintenance waiting while it commits. + let mut reader = hold_lock(&pool, "LOCK TABLE ONLY events IN ACCESS SHARE MODE").await; + let reader_pid = backend_pid(&mut reader).await; + let migration = hold_lock( + &pool, + "ALTER TABLE ONLY events_p_future_next ADD CONSTRAINT catch_all_reviewer \ + FOREIGN KEY (community_id) REFERENCES reviewers(id)", + ) + .await; + let maintenance = tokio::spawn({ + let pool = pool.clone(); + async move { maintain(&pool).await } + }); + wait_until_blocked_by(&admin, reader_pid, 1).await; + migration.commit().await.expect("commit migration"); + reader.commit().await.expect("release parent"); + + let result = maintenance.await.expect("maintenance task"); + assert!( + refused_events_only(&result, "catch_all_reviewer"), + "{result:?}" + ); + let kept: i64 = sqlx::query_scalar( + "SELECT count(*) FROM pg_constraint WHERE conname = 'catch_all_reviewer' \ + AND connamespace = current_schema()::regnamespace", + ) + .fetch_one(&pool) + .await + .expect("count catch-all key"); + assert_eq!(kept, 1); + drop_schema(&admin, &schema).await; + } + + #[tokio::test] + #[ignore = "requires Postgres"] + async fn deadlock_victim_is_classified_as_a_lock_timeout() { + let (pool, admin, schema) = scratch_pool().await; + seed_parents(&pool).await; + let mut first = hold_lock(&pool, "LOCK TABLE ONLY events IN ACCESS EXCLUSIVE MODE").await; + let mut second = hold_lock( + &pool, + "LOCK TABLE ONLY delivery_log IN ACCESS EXCLUSIVE MODE", + ) + .await; + let first_waits = sqlx::query("LOCK TABLE ONLY delivery_log IN ACCESS EXCLUSIVE MODE") + .execute(&mut *first); + let second_waits = async { + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + sqlx::query("LOCK TABLE ONLY events IN ACCESS EXCLUSIVE MODE") + .execute(&mut *second) + .await + }; + let (first_result, second_result) = tokio::join!(first_waits, second_waits); + let victim = first_result + .err() + .or(second_result.err()) + .expect("one session is the deadlock victim"); + assert_eq!( + victim + .as_database_error() + .and_then(|error| error.code()) + .as_deref(), + Some("40P01"), + "{victim:?}" + ); + assert_eq!( + classify_error(&DbError::from(victim)), + PartitionMaintenanceOutcome::LockTimeout + ); + drop((first, second)); + drop_schema(&admin, &schema).await; + } + + #[tokio::test] + #[ignore = "requires Postgres"] + async fn parent_contention_times_out_without_blocking_the_other_table() { + let (pool, admin, schema) = scratch_pool().await; + seed_repaired_layout(&pool, true).await; + let holder = hold_lock(&pool, "LOCK TABLE ONLY events IN ACCESS SHARE MODE").await; + + let result = maintain(&pool).await; + drop(holder); + + assert!( + matches!(result, Err(DbError::InvalidData(ref message)) + if message.contains("events lock_timeout") && !message.contains("delivery_log")), + "{result:?}" + ); + let audit = audit_partition_catalog_at(&pool, PARTITION_MANAGER_MONTHS_AHEAD, fixed_now()) + .await + .expect("audit"); + let events = audit.tables.iter().find(|t| t.table == "events").unwrap(); + assert!(events + .children + .iter() + .any(|child| child.name == "events_p_future_next")); + let delivery_log = audit + .tables + .iter() + .filter(|t| t.table == "delivery_log") + .cloned() + .collect::>(); + assert_advanced(&PartitionAudit { + audited_at: audit.audited_at, + tables: delivery_log, + }); + drop_schema(&admin, &schema).await; + } + + #[tokio::test] + #[ignore = "requires Postgres"] + async fn held_maintenance_lock_skips_without_ddl() { + let (pool, admin, schema) = scratch_pool().await; + seed_repaired_layout(&pool, false).await; + let before = catalog_snapshot(&pool).await; + let holder = hold_lock( + &pool, + "SELECT pg_advisory_xact_lock(hashtextextended(\ + 'buzz.partition_maintenance:' || current_schema()::text, 0))", + ) + .await; + + maintain(&pool) + .await + .expect("skipped_locked is not a failure"); + assert_eq!(catalog_snapshot(&pool).await, before); + drop(holder); + drop_schema(&admin, &schema).await; + } + + #[tokio::test] + #[ignore = "requires Postgres"] + async fn concurrent_runners_converge_on_one_layout() { + let (pool, admin, schema) = scratch_pool().await; + seed_repaired_layout(&pool, true).await; + + let (first, second) = tokio::join!(maintain(&pool), maintain(&pool)); + // A runner's catch-all NOWAIT fails fast if the other runner's audit is + // probing that catch-all; nothing else may go wrong. + for (runner, result) in [("first", first), ("second", second)] { + let Err(error) = result else { continue }; + let message = error.to_string(); + assert!( + message.contains("lock_timeout"), + "{runner} runner: {message}" + ); + for table in PARTITIONED_TABLES { + for outcome in ["operator_required", "deadline", "error"] { + assert!( + !message.contains(&format!("{table} {outcome}:")), + "{runner} runner: {message}" + ); + } + } + } + // Whatever a fail-fast loser left undone, the next pass converges. + assert_advanced(&maintain(&pool).await.expect("follow-up pass")); + drop_schema(&admin, &schema).await; + } + + #[tokio::test] + #[ignore = "requires Postgres"] + async fn client_deadline_rolls_back_and_discards_the_session() { + let (pool, admin, schema) = scratch_pool().await; + seed_repaired_layout(&pool, false).await; + let before = catalog_snapshot(&pool).await; + // The parent lock wait (2s) outlasts the 300ms client deadline. + let mut holder = hold_lock(&pool, "LOCK TABLE ONLY events IN ACCESS SHARE MODE").await; + let holder_pid: i32 = sqlx::query_scalar("SELECT pg_backend_pid()") + .fetch_one(&mut *holder) + .await + .expect("holder backend pid"); + + let result = maintain_partitions_with_deadline( + &pool, + PARTITION_MANAGER_MONTHS_AHEAD, + ADVANCE, + fixed_now(), + std::time::Duration::from_millis(300), + ) + .await; + assert!( + matches!(result, Err(DbError::InvalidData(ref message)) + if message.contains("events deadline")), + "{result:?}" + ); + + // The abandoned session is closed, not pooled: the backend that was + // queued behind the holder disappears and releases every lock before + // the holder commits. Only sessions blocked by this test's holder + // count, so maintenance elsewhere on the cluster cannot interfere. + let blocked: Vec = sqlx::query_scalar( + "SELECT pid FROM pg_stat_activity WHERE $1 = ANY(pg_blocking_pids(pid))", + ) + .bind(holder_pid) + .fetch_all(&admin) + .await + .expect("sessions blocked by the holder"); + let mut gone = blocked.is_empty(); + for _ in 0..50 { + if gone { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + let alive: bool = sqlx::query_scalar( + "SELECT EXISTS (SELECT 1 FROM pg_stat_activity WHERE pid = ANY($1))", + ) + .bind(&blocked) + .fetch_one(&admin) + .await + .expect("check abandoned backend"); + gone = !alive; + } + drop(holder); + assert!( + blocked.len() <= 1, + "unexpected blocked sessions {blocked:?}" + ); + assert!( + gone, + "abandoned maintenance session {blocked:?} was not closed" + ); + let after = catalog_snapshot(&pool).await; + let events_only = |snapshot: &[(String, String, String, i64)]| { + snapshot + .iter() + .filter(|(name, ..)| name.starts_with("events_")) + .cloned() + .collect::>() + }; + assert_eq!(events_only(&after), events_only(&before)); + drop_schema(&admin, &schema).await; + } + + #[tokio::test] + #[ignore = "requires Postgres"] + async fn catch_all_with_child_only_trigger_is_refused() { + let (pool, admin, schema) = scratch_pool().await; + seed_repaired_layout(&pool, false).await; + sqlx::query( + "CREATE TRIGGER catch_all_only BEFORE INSERT ON events_p_future_next \ + FOR EACH ROW EXECUTE FUNCTION partition_test_trigger()", + ) + .execute(&pool) + .await + .expect("create child-only trigger"); + let before = catalog_snapshot(&pool).await; + + let result = maintain(&pool).await; + assert!( + matches!(result, Err(DbError::InvalidData(ref message)) + if message.contains("events operator_required") + && message.contains("child-only triggers (catch_all_only)") + && !message.contains("delivery_log")), + "{result:?}" + ); + let after = catalog_snapshot(&pool).await; + let events_only = |snapshot: &[(String, String, String, i64)]| { + snapshot + .iter() + .filter(|(name, ..)| name.starts_with("events_")) + .cloned() + .collect::>() + }; + assert_eq!(events_only(&after), events_only(&before)); + drop_schema(&admin, &schema).await; + } + + #[tokio::test] + #[ignore = "requires Postgres"] + async fn stale_audit_after_a_concurrent_advance_is_a_noop() { + let (pool, admin, schema) = scratch_pool().await; + seed_repaired_layout(&pool, false).await; + // A runner audits, then loses the race to a runner that commits first. + let stale = audit_partition_catalog_at(&pool, PARTITION_MANAGER_MONTHS_AHEAD, fixed_now()) + .await + .expect("stale audit"); + assert_advanced(&maintain(&pool).await.expect("winning runner")); + let advanced = catalog_snapshot(&pool).await; + + for table_audit in &stale.tables { + let result = maintain_table( + &pool, + table_audit, + plan_maintenance(table_audit, ADVANCE), + PARTITION_MANAGER_MONTHS_AHEAD as i32, + ADVANCE, + fixed_now(), + std::time::Duration::from_secs(10), + ) + .await; + assert_eq!( + (result.outcome, result.error.as_deref()), + (PartitionMaintenanceOutcome::Noop, None), + "{}", + table_audit.table + ); + } + assert_eq!(catalog_snapshot(&pool).await, advanced); + + // The losing runner reports the layout the winner committed. + let reported = maintain_audited( + &pool, + stale, + PARTITION_MANAGER_MONTHS_AHEAD, + ADVANCE, + fixed_now(), + std::time::Duration::from_secs(10), + ) + .await + .expect("losing runner"); + assert_advanced(&reported); + drop_schema(&admin, &schema).await; + } + + #[tokio::test] + #[ignore = "requires Postgres"] + async fn replan_that_finds_only_disabled_work_is_skipped_disabled() { + let (pool, admin, schema) = scratch_pool().await; + seed_repaired_layout(&pool, false).await; + sqlx::query("DROP TABLE events_p_future_next") + .execute(&pool) + .await + .expect("drop catch-all"); + // Creation is planned for the uncovered months, but a catch-all appears + // before the lock and advancing it is disabled. + let stale = audit_partition_catalog_at(&pool, PARTITION_MANAGER_MONTHS_AHEAD, fixed_now()) + .await + .expect("stale audit"); + let events = stale.tables.iter().find(|t| t.table == "events").unwrap(); + create_child( + &pool, + "events", + "events_p_future_next", + "'2027-01-01'", + "MAXVALUE", + ) + .await; + let before = catalog_snapshot(&pool).await; + + let result = maintain_table( + &pool, + events, + plan_maintenance(events, CREATE_ONLY), + PARTITION_MANAGER_MONTHS_AHEAD as i32, + CREATE_ONLY, + fixed_now(), + std::time::Duration::from_secs(10), + ) + .await; + assert_eq!( + (result.outcome, result.error.as_deref()), + (PartitionMaintenanceOutcome::SkippedDisabled, None) + ); + assert_eq!(catalog_snapshot(&pool).await, before); + drop_schema(&admin, &schema).await; + } + + #[tokio::test] + #[ignore = "requires Postgres"] + async fn audit_retries_when_a_listed_partition_is_dropped() { + let (pool, admin, schema) = scratch_pool().await; + seed_repaired_layout(&pool, false).await; + let mut dropper = pool.begin().await.expect("begin dropper"); + let dropper_pid: i32 = sqlx::query_scalar("SELECT pg_backend_pid()") + .fetch_one(&mut *dropper) + .await + .expect("dropper backend pid"); + sqlx::query("DROP TABLE events_p_future_next") + .execute(&mut *dropper) + .await + .expect("drop catch-all"); + + // The audit lists the catch-all, then queues behind the drop's lock. + let audit = tokio::spawn({ + let pool = pool.clone(); + async move { + audit_partition_catalog_at(&pool, PARTITION_MANAGER_MONTHS_AHEAD, fixed_now()).await + } + }); + let mut queued = false; + for _ in 0..50 { + queued = sqlx::query_scalar( + "SELECT EXISTS (SELECT 1 FROM pg_stat_activity \ + WHERE $1 = ANY(pg_blocking_pids(pid)))", + ) + .bind(dropper_pid) + .fetch_one(&admin) + .await + .expect("check audit queue"); + if queued { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(50)).await; + } + assert!(queued, "audit never queued behind the drop"); + dropper.commit().await.expect("commit drop"); + + let audit = audit.await.expect("audit task").expect("audit retried"); + let events = audit + .tables + .iter() + .find(|table| table.table == "events") + .expect("events audit"); + assert!(events + .children + .iter() + .all(|child| child.name != "events_p_future_next")); + drop_schema(&admin, &schema).await; + } + + #[tokio::test] + #[ignore = "requires Postgres"] + async fn unrelated_catch_all_name_is_an_operator_error() { + let (pool, admin, schema) = scratch_pool().await; + seed_repaired_layout(&pool, false).await; + sqlx::query("CREATE TABLE events_p_future (unrelated BOOLEAN)") + .execute(&pool) + .await + .expect("create colliding relation"); + let before = catalog_snapshot(&pool).await; + + let result = maintain(&pool).await; + assert!( + matches!(result, Err(DbError::InvalidData(ref message)) + if message.contains("canonical name events_p_future already exists")), + "{result:?}" + ); + let after = catalog_snapshot(&pool).await; + assert!(after + .iter() + .any(|(name, ..)| name == "events_p_future_next")); + assert_eq!( + after + .iter() + .filter(|(name, ..)| name.starts_with("events_")) + .count(), + before + .iter() + .filter(|(name, ..)| name.starts_with("events_")) + .count() + ); + drop_schema(&admin, &schema).await; + } + + #[tokio::test] + #[ignore = "requires Postgres"] + async fn incoming_foreign_key_is_refused_under_lock() { + let (pool, admin, schema) = scratch_pool().await; + seed_repaired_layout(&pool, false).await; + sqlx::query( + "CREATE TABLE event_refs (created_at TIMESTAMPTZ, alternate_at TIMESTAMPTZ, id BIGINT, \ + FOREIGN KEY (created_at, alternate_at, id) REFERENCES events)", + ) + .execute(&pool) + .await + .expect("create referencing table"); + let before = catalog_snapshot(&pool).await; + + let result = maintain(&pool).await; + assert!( + matches!(result, Err(DbError::InvalidData(ref message)) + if message.contains("events is referenced by a foreign key")), + "{result:?}" + ); + let after = catalog_snapshot(&pool).await; + assert!(after + .iter() + .any(|(name, ..)| name == "events_p_future_next")); + assert_ne!(after, before, "delivery_log still advances"); + drop_schema(&admin, &schema).await; + } +} diff --git a/crates/buzz-relay/src/api/artifact_postgres_tests.rs b/crates/buzz-relay/src/api/artifact_postgres_tests.rs index bb0a071e8f..e60010cc8a 100644 --- a/crates/buzz-relay/src/api/artifact_postgres_tests.rs +++ b/crates/buzz-relay/src/api/artifact_postgres_tests.rs @@ -25,7 +25,17 @@ impl Fixture { .unwrap(); let host = format!("artifact-review-{}.local", Uuid::new_v4()); state.db.ensure_configured_community(&host).await.unwrap(); - state.db.ensure_future_partitions(1, true).await.unwrap(); + state + .db + .maintain_partitions( + 1, + buzz_db::partition::PartitionMaintenancePolicy { + create_enabled: true, + advance_enabled: false, + }, + ) + .await + .unwrap(); let tenant = crate::tenant::bind_community(&state.db, &host) .await .unwrap(); diff --git a/crates/buzz-relay/src/config.rs b/crates/buzz-relay/src/config.rs index 589dc25719..ca895b989e 100644 --- a/crates/buzz-relay/src/config.rs +++ b/crates/buzz-relay/src/config.rs @@ -206,8 +206,13 @@ pub struct Config { pub metrics_port: u16, /// Interval between read-only partition catalog audits. pub partition_audit_interval: Duration, - /// Whether the partition manager may create uncovered monthly partitions. + /// Whether the partition manager may issue any partition DDL. Off also + /// disables catch-all advancement. pub partition_manager_create_enabled: bool, + /// Whether the partition manager may replace an empty right-edge catch-all + /// with dedicated monthlies. Also enables periodic (not only startup) DDL. + /// Requires `partition_manager_create_enabled`. + pub partition_manager_advance_enabled: bool, /// When true, NIP-42 pubkey-only authentication (no API token) is /// restricted to pubkeys in the `pubkey_allowlist` table. Users with valid @@ -938,6 +943,9 @@ impl Config { ); let partition_manager_create_enabled = parse_bool("BUZZ_PARTITION_MANAGER_CREATE_ENABLED", true)?; + // Off by default: each environment opts in after its canary. + let partition_manager_advance_enabled = + parse_bool("BUZZ_PARTITION_MANAGER_ADVANCE_ENABLED", false)?; let s3_addressing_style = match std::env::var("BUZZ_S3_ADDRESSING_STYLE") { Ok(value) => value.parse().map_err(ConfigError::InvalidValue)?, @@ -1392,6 +1400,7 @@ impl Config { metrics_port, partition_audit_interval, partition_manager_create_enabled, + partition_manager_advance_enabled, pubkey_allowlist_enabled, require_relay_membership, huddle_audio_available, @@ -1565,6 +1574,7 @@ mod tests { assert_eq!(config.max_frame_bytes, DEFAULT_MAX_FRAME_BYTES); assert_eq!(config.partition_audit_interval, Duration::from_secs(900)); assert!(config.partition_manager_create_enabled); + assert!(!config.partition_manager_advance_enabled); assert!(config.slow_client_grace_limit > 0); assert!( !config.pubkey_allowlist_enabled, @@ -2069,6 +2079,33 @@ mod tests { )); } + #[test] + fn partition_manager_advance_is_opt_in_and_parses_strictly() { + let _guards = env_guards(); + let previous = std::env::var_os("BUZZ_PARTITION_MANAGER_ADVANCE_ENABLED"); + + std::env::remove_var("BUZZ_PARTITION_MANAGER_ADVANCE_ENABLED"); + let default = Config::from_env().expect("default config"); + std::env::set_var("BUZZ_PARTITION_MANAGER_ADVANCE_ENABLED", " TRUE "); + let enabled = Config::from_env().expect("enabled config"); + std::env::set_var("BUZZ_PARTITION_MANAGER_ADVANCE_ENABLED", "enable-maybe"); + let invalid = Config::from_env(); + + if let Some(value) = previous { + std::env::set_var("BUZZ_PARTITION_MANAGER_ADVANCE_ENABLED", value); + } else { + std::env::remove_var("BUZZ_PARTITION_MANAGER_ADVANCE_ENABLED"); + } + + assert!(!default.partition_manager_advance_enabled); + assert!(enabled.partition_manager_advance_enabled); + assert!(matches!( + invalid, + Err(ConfigError::InvalidValue(ref message)) + if message.contains("BUZZ_PARTITION_MANAGER_ADVANCE_ENABLED") + )); + } + #[test] fn s3_addressing_style_env_accepts_virtual_and_rejects_invalid_values() { let _guards = env_guards(); diff --git a/crates/buzz-relay/src/main.rs b/crates/buzz-relay/src/main.rs index 9021d6644f..3df5c428ee 100644 --- a/crates/buzz-relay/src/main.rs +++ b/crates/buzz-relay/src/main.rs @@ -12,7 +12,7 @@ use uuid::Uuid; use buzz_audit::AuditService; use buzz_auth::AuthService; use buzz_core::CommunityId; -use buzz_db::{Db, DbConfig}; +use buzz_db::{partition::PARTITION_MANAGER_MONTHS_AHEAD, Db, DbConfig}; use buzz_pubsub::PubSubManager; use buzz_search::SearchService; @@ -480,16 +480,20 @@ async fn run_relay_main(boot: BootTracker) -> anyhow::Result<()> { info!("Skipping database migrations because BUZZ_AUTO_MIGRATE is not enabled"); } + let partition_policy = buzz_db::partition::PartitionMaintenancePolicy { + create_enabled: config.partition_manager_create_enabled, + advance_enabled: config.partition_manager_advance_enabled, + }; let mut step = StepTimer::start(StartupStep::PartitionEnsure); let startup_partition_audit = match db - .ensure_future_partitions(3, config.partition_manager_create_enabled) + .maintain_partitions(PARTITION_MANAGER_MONTHS_AHEAD, partition_policy) .await { Ok(audit) => Some(audit), Err(error) => { step.degrade(); error!(%error, "Failed to ensure partitions"); - match db.audit_partitions(3).await { + match db.audit_partitions(PARTITION_MANAGER_MONTHS_AHEAD).await { Ok(audit) => Some(audit), Err(error) => { error!(%error, "Initial partition catalog audit failed"); @@ -812,8 +816,9 @@ async fn run_relay_main(boot: BootTracker) -> anyhow::Result<()> { state.record_partition_audit(audit); } - // The periodic path is deliberately read-only. Partition creation only - // occurs during the bounded startup pass. + // The periodic path is read-only unless catch-all advancement is enabled. + // Maintenance takes no lock when the runway is already complete, so a + // periodic attempt usually costs one catalog audit. { let partition_state = Arc::clone(&state); let audit_interval = state.config.partition_audit_interval; @@ -822,7 +827,21 @@ async fn run_relay_main(boot: BootTracker) -> anyhow::Result<()> { PartitionAuditSchedule::new(audit_interval, has_startup_partition_audit); loop { schedule.wait().await; - let result = partition_state.db.audit_partitions(3).await; + let db = &partition_state.db; + let result = if partition_policy.advance_enabled { + match db + .maintain_partitions(PARTITION_MANAGER_MONTHS_AHEAD, partition_policy) + .await + { + Ok(audit) => Ok(audit), + Err(error) => { + warn!(%error, "Periodic partition maintenance failed"); + db.audit_partitions(PARTITION_MANAGER_MONTHS_AHEAD).await + } + } + } else { + db.audit_partitions(PARTITION_MANAGER_MONTHS_AHEAD).await + }; schedule.record_attempt(result.is_ok()); match result { Ok(audit) => partition_state.record_partition_audit(audit),