fix(db): avoid partition DDL on covered ranges

Signed-off-by: Luke Tornquist <tornquist@squareup.com>

crates/buzz-db/src/store/partition.rs

@@ -4,7 +4,7 @@
 
 use buzz_datastore_tracing::datastore_span;
 use chrono::{Datelike, TimeZone, Utc};
-use sqlx::{PgPool, Row};
+use sqlx::PgPool;
 use tracing::info;
 
 use crate::error::{DbError, Result};
@@ -16,11 +16,6 @@ const PARTITIONED_TABLES: &[&str] = &["events", "delivery_log"];
 /// Ensures monthly partition tables exist for the next `months_ahead` months.
 pub async fn ensure_future_partitions(pool: &PgPool, months_ahead: u32) -> Result<()> {
     let now = Utc::now();
-    let mut connection = crate::observability::acquire_writer(
-        pool,
-        crate::observability::WriterOperation::Bootstrap,
-    )
-    .await?;
 
     for i in 0..=(months_ahead as i32) {
         let year = now.year();
@@ -55,7 +50,7 @@ pub async fn ensure_future_partitions(pool: &PgPool, months_ahead: u32) -> Resul
         let end_str = end.format("%Y-%m-%d").to_string();
 
         for table in PARTITIONED_TABLES {
-            ensure_partition(&mut connection, table, &start_str, &end_str, &suffix).await?;
+            ensure_partition(pool, table, &start_str, &end_str, &suffix).await?;
         }
     }
 
@@ -87,7 +82,7 @@ fn validate_date_str(s: &str) -> bool {
 }
 
 async fn ensure_partition(
-    connection: &mut sqlx::PgConnection,
+    pool: &PgPool,
     table_name: &str,
     start_date_str: &str,
     end_date_str: &str,
@@ -117,22 +112,7 @@ async fn ensure_partition(
 
     let partition_name = format!("{table_name}_p{suffix}");
 
-    let row = sqlx::query(
-        r#"
-        SELECT COUNT(*) as cnt
-        FROM pg_catalog.pg_class c
-        JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace
-        WHERE n.nspname = current_schema()
-          AND c.relname = $1
-          AND c.relispartition = true
-        "#,
-    )
-    .bind(&partition_name)
-    .fetch_one(&mut *connection)
-    .await?;
-
-    let cnt: i64 = row.try_get("cnt")?;
-    if cnt > 0 {
+    if partition_range_covered(pool, table_name, start_date_str, end_date_str).await? {
         return Ok(());
     }
 
@@ -142,31 +122,131 @@ async fn ensure_partition(
          FOR VALUES FROM ('{start_date_str}') TO ('{end_date_str}')"
     );
 
-    match sqlx::query(sqlx::AssertSqlSafe(sql))
-        .execute(&mut *connection)
-        .await
-    {
+    // SET LOCAL applies only to this transaction, and the DDL must use the same
+    // connection for the timeout to bound its parent-table lock wait.
+    let mut tx = pool.begin().await?;
+    sqlx::query("SET LOCAL lock_timeout = '2s'")
+        .execute(&mut *tx)
+        .await?;
+    sqlx::query("SET LOCAL TIME ZONE 'UTC'")
+        .execute(&mut *tx)
+        .await?;
+    let result = sqlx::query(sqlx::AssertSqlSafe(sql))
+        .execute(&mut *tx)
+        .await;
+
+    match result {
         Ok(_) => {
+            tx.commit().await?;
             info!("added partition {partition_name}");
             Ok(())
         }
-        Err(sqlx::Error::Database(db_err))
-            if db_err.code().as_deref() == Some("42P17")
-                && db_err.message().contains("would overlap partition") =>
-        {
-            // Fresh schemas include a right-edge catch-all partition (`*_p_future`).
-            // If it already covers this month, the table is still safe for writes;
-            // treat the overlap as "ensured" rather than failing startup.
-            info!(
-                partition_name,
-                "partition range already covered by an existing partition"
+        Err(error) => {
+            let is_overlap = matches!(
+                &error,
+                sqlx::Error::Database(db_error)
+                    if db_error.code().as_deref() == Some("42P17")
+                        && db_error.message().contains("would overlap partition")
             );
-            Ok(())
+            tx.rollback().await?;
+            if is_overlap
+                && partition_range_covered(pool, table_name, start_date_str, end_date_str).await?
+            {
+                // A concurrent creator can cover the range between the catalog
+                // pre-check and this DDL.
+                info!(
+                    partition_name,
+                    "partition range already covered by an existing partition"
+                );
+                Ok(())
+            } else {
+                Err(error.into())
+            }
         }
-        Err(e) => Err(e.into()),
     }
 }
 
+async fn partition_range_covered(
+    pool: &PgPool,
+    table_name: &str,
+    start_date_str: &str,
+    end_date_str: &str,
+) -> Result<bool> {
+    let (all_bounds_recognized, covered): (bool, bool) = sqlx::query_as(
+        r#"
+        WITH partition_nodes AS (
+            SELECT
+                tree.level,
+                tree.isleaf,
+                pg_get_expr(c.relpartbound, c.oid, false) AS bound
+            FROM pg_catalog.pg_partition_tree(to_regclass($1)) AS tree
+            JOIN pg_catalog.pg_class c ON c.oid = tree.relid
+            WHERE tree.level > 0
+        ),
+        partition_bounds AS (
+            SELECT bound
+            FROM partition_nodes
+            WHERE level = 1 AND isleaf
+        ),
+        extracted_bounds AS (
+            SELECT
+                bound,
+                regexp_match(
+                    bound,
+                    $partition_bound$^FOR VALUES FROM \((MINVALUE|'[^']+')\) TO \((MAXVALUE|'[^']+')\)$$partition_bound$
+                ) AS parts
+            FROM partition_bounds
+        ),
+        parsed_bounds AS (
+            SELECT
+                bound,
+                CASE
+                    WHEN parts[1] = 'MINVALUE'
+                        THEN '-infinity'::timestamptz
+                    WHEN parts IS NOT NULL
+                        THEN trim(both '''' from parts[1])::timestamptz
+                END AS start_at,
+                CASE
+                    WHEN parts[2] = 'MAXVALUE'
+                        THEN 'infinity'::timestamptz
+                    WHEN parts IS NOT NULL
+                        THEN trim(both '''' from parts[2])::timestamptz
+                END AS end_at
+            FROM extracted_bounds
+        )
+        SELECT
+            COALESCE(
+                (SELECT bool_and(level = 1 AND isleaf) FROM partition_nodes),
+                true
+            )
+            AND COALESCE(bool_and(bound = 'DEFAULT' OR parts IS NOT NULL), true),
+            EXISTS (
+                SELECT 1
+                FROM parsed_bounds
+                WHERE bound = 'DEFAULT'
+                   OR (
+                       start_at <= to_date($2, 'YYYY-MM-DD')::timestamp AT TIME ZONE 'UTC'
+                       AND end_at >= to_date($3, 'YYYY-MM-DD')::timestamp AT TIME ZONE 'UTC'
+                   )
+            )
+        FROM extracted_bounds
+        "#,
+    )
+    .bind(table_name)
+    .bind(start_date_str)
+    .bind(end_date_str)
+    .fetch_one(pool)
+    .await?;
+
+    if !all_bounds_recognized {
+        return Err(DbError::InvalidData(format!(
+            "unsupported partition topology or bound for table {table_name:?}"
+        )));
+    }
+
+    Ok(covered)
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -197,4 +277,403 @@ mod tests {
         assert!(!PARTITIONED_TABLES.contains(&"api_tokens"));
         assert!(!PARTITIONED_TABLES.contains(&"users"));
     }
+
+    async fn admin_pool() -> PgPool {
+        let url = std::env::var("TEST_DATABASE_URL")
+            .expect("TEST_DATABASE_URL must point to a database where tests can create databases");
+        PgPool::connect(&url).await.expect("connect admin database")
+    }
+
+    async fn create_scratch_db(admin: &PgPool, prefix: &str) -> (PgPool, String) {
+        create_scratch_db_with_timezone(admin, prefix, "UTC").await
+    }
+
+    async fn create_scratch_db_with_timezone(
+        admin: &PgPool,
+        prefix: &str,
+        timezone: &str,
+    ) -> (PgPool, String) {
+        let name = format!("{prefix}_{}", uuid::Uuid::new_v4().simple());
+        sqlx::query(sqlx::AssertSqlSafe(format!("CREATE DATABASE {name}")))
+            .execute(admin)
+            .await
+            .expect("create scratch database");
+
+        let admin_url = std::env::var("TEST_DATABASE_URL").expect("TEST_DATABASE_URL");
+        let path_start = admin_url
+            .rfind('/')
+            .expect("database URL has a path segment");
+        let scratch_url = format!("{}/{}", &admin_url[..path_start], name);
+        let pool = sqlx::postgres::PgPoolOptions::new()
+            .max_connections(2)
+            .connect(&scratch_url)
+            .await
+            .expect("connect scratch database");
+        let mut first = pool.acquire().await.expect("acquire first connection");
+        let mut second = pool.acquire().await.expect("acquire second connection");
+        for connection in [&mut *first, &mut *second] {
+            sqlx::query("SELECT set_config('TimeZone', $1, false)")
+                .bind(timezone)
+                .execute(connection)
+                .await
+                .expect("set test connection timezone");
+        }
+        drop(first);
+        drop(second);
+        create_partitioned_tables(&pool).await;
+        (pool, name)
+    }
+
+    async fn create_partitioned_tables(pool: &PgPool) {
+        for statement in [
+            "CREATE TABLE events (created_at TIMESTAMPTZ NOT NULL) PARTITION BY RANGE (created_at)",
+            "CREATE TABLE events_p_past PARTITION OF events FOR VALUES FROM (MINVALUE) TO ('2000-01-01')",
+            "CREATE TABLE events_p_future PARTITION OF events FOR VALUES FROM ('2000-01-01') TO (MAXVALUE)",
+            "CREATE TABLE delivery_log (delivered_at TIMESTAMPTZ NOT NULL) PARTITION BY RANGE (delivered_at)",
+            "CREATE TABLE delivery_log_p_past PARTITION OF delivery_log FOR VALUES FROM (MINVALUE) TO ('2000-01-01')",
+            "CREATE TABLE delivery_log_p_future PARTITION OF delivery_log FOR VALUES FROM ('2000-01-01') TO (MAXVALUE)",
+        ] {
+            sqlx::query(statement)
+                .execute(pool)
+                .await
+                .expect("create partitioned test table");
+        }
+    }
+
+    async fn drop_scratch_db(admin: &PgPool, pool: PgPool, name: &str) {
+        pool.close().await;
+        sqlx::query(sqlx::AssertSqlSafe(format!(
+            "DROP DATABASE IF EXISTS {name} WITH (FORCE)"
+        )))
+        .execute(admin)
+        .await
+        .expect("drop scratch database");
+    }
+
+    #[tokio::test]
+    #[ignore = "requires Postgres"]
+    async fn existing_catch_all_bounds_avoid_partition_ddl() {
+        let admin = admin_pool().await;
+        let (pool, name) = create_scratch_db(&admin, "partition_coverage").await;
+        let mut blocker = pool.begin().await.expect("begin blocker");
+        sqlx::query("LOCK TABLE events IN ACCESS SHARE MODE")
+            .execute(&mut *blocker)
+            .await
+            .expect("lock partition parent");
+
+        let result = tokio::time::timeout(
+            std::time::Duration::from_millis(500),
+            ensure_future_partitions(&pool, 0),
+        )
+        .await;
+
+        blocker.rollback().await.expect("release parent lock");
+        drop_scratch_db(&admin, pool, &name).await;
+        result
+            .expect("coverage pre-check must return without waiting for a DDL lock")
+            .expect("covered partition range is already ensured");
+    }
+
+    fn is_lock_timeout(error: &DbError) -> bool {
+        matches!(
+            error,
+            DbError::Sqlx(sqlx::Error::Database(db_error))
+                if db_error.code().as_deref() == Some("55P03")
+        )
+    }
+
+    #[tokio::test]
+    #[ignore = "requires Postgres"]
+    async fn blocked_partition_ddl_stops_at_lock_timeout() {
+        let admin = admin_pool().await;
+        let (pool, name) = create_scratch_db(&admin, "partition_lock_timeout").await;
+        sqlx::query("DROP TABLE events_p_future")
+            .execute(&pool)
+            .await
+            .expect("remove partition coverage");
+        let mut blocker = pool.begin().await.expect("begin blocker");
+        sqlx::query("LOCK TABLE events IN ACCESS SHARE MODE")
+            .execute(&mut *blocker)
+            .await
+            .expect("lock partition parent");
+
+        let result = tokio::time::timeout(
+            std::time::Duration::from_secs(4),
+            ensure_future_partitions(&pool, 0),
+        )
+        .await;
+
+        blocker.rollback().await.expect("release parent lock");
+        drop_scratch_db(&admin, pool, &name).await;
+        let error = result
+            .expect("partition DDL must stop at the transaction-local lock timeout")
+            .expect_err("blocked partition DDL must return a lock timeout");
+        assert!(is_lock_timeout(&error), "unexpected error: {error}");
+    }
+
+    #[tokio::test]
+    #[ignore = "requires Postgres"]
+    async fn missing_partition_is_created_with_production_bounds_in_non_utc_session() {
+        let admin = admin_pool().await;
+        let (pool, name) =
+            create_scratch_db_with_timezone(&admin, "partition_missing", "America/Los_Angeles")
+                .await;
+        sqlx::query("DROP TABLE events_p_future, delivery_log_p_future")
+            .execute(&pool)
+            .await
+            .expect("remove catch-all partitions");
+
+        ensure_future_partitions(&pool, 0)
+            .await
+            .expect("create genuinely missing monthly partitions");
+
+        let now = Utc::now();
+        let suffix = format!("{:04}_{:02}", now.year(), now.month());
+        let start_str = format!("{:04}-{:02}-01", now.year(), now.month());
+        let (end_year, end_month) = if now.month() == 12 {
+            (now.year() + 1, 1)
+        } else {
+            (now.year(), now.month() + 1)
+        };
+        let end_str = format!("{end_year:04}-{end_month:02}-01");
+        for table in PARTITIONED_TABLES {
+            let partition_name = format!("{table}_p{suffix}");
+            let exists: bool = sqlx::query_scalar("SELECT to_regclass($1) IS NOT NULL")
+                .bind(&partition_name)
+                .fetch_one(&pool)
+                .await
+                .expect("check created partition");
+            assert!(exists, "missing partition {partition_name}");
+            let covered = partition_range_covered(&pool, table, &start_str, &end_str)
+                .await
+                .expect("inspect created partition bounds");
+            assert!(covered, "partition {partition_name} must use UTC bounds");
+        }
+
+        drop_scratch_db(&admin, pool, &name).await;
+    }
+
+    #[tokio::test]
+    #[ignore = "requires Postgres"]
+    async fn partial_overlap_does_not_count_as_coverage() {
+        let admin = admin_pool().await;
+        let (pool, name) = create_scratch_db(&admin, "partition_partial_overlap").await;
+        sqlx::query("DROP TABLE events_p_future")
+            .execute(&pool)
+            .await
+            .expect("remove events catch-all partition");
+
+        let now = Utc::now();
+        let start = Utc
+            .with_ymd_and_hms(now.year(), now.month(), 1, 0, 0, 0)
+            .single()
+            .expect("current month starts at a valid date");
+        let middle = Utc
+            .with_ymd_and_hms(now.year(), now.month(), 15, 0, 0, 0)
+            .single()
+            .expect("current month has a fifteenth day");
+        let sql = format!(
+            "CREATE TABLE events_partial PARTITION OF events \
+             FOR VALUES FROM ('{}') TO ('{}')",
+            start.to_rfc3339(),
+            middle.to_rfc3339()
+        );
+        sqlx::query(sqlx::AssertSqlSafe(sql))
+            .execute(&pool)
+            .await
+            .expect("create partially overlapping partition");
+
+        let error = ensure_future_partitions(&pool, 0)
+            .await
+            .expect_err("partial overlap must not be accepted as full coverage");
+        assert!(
+            matches!(
+                error,
+                DbError::Sqlx(sqlx::Error::Database(ref db_error))
+                    if db_error.code().as_deref() == Some("42P17")
+            ),
+            "unexpected error: {error}"
+        );
+
+        drop_scratch_db(&admin, pool, &name).await;
+    }
+
+    #[tokio::test]
+    #[ignore = "requires Postgres"]
+    async fn default_partition_avoids_partition_ddl() {
+        let admin = admin_pool().await;
+        let (pool, name) = create_scratch_db(&admin, "partition_default").await;
+        sqlx::query("DROP TABLE events_p_future, delivery_log_p_future")
+            .execute(&pool)
+            .await
+            .expect("remove catch-all partitions");
+        for statement in [
+            "CREATE TABLE events_p_default PARTITION OF events DEFAULT",
+            "CREATE TABLE delivery_log_p_default PARTITION OF delivery_log DEFAULT",
+        ] {
+            sqlx::query(statement)
+                .execute(&pool)
+                .await
+                .expect("create default partition");
+        }
+        let mut blocker = pool.begin().await.expect("begin blocker");
+        sqlx::query("LOCK TABLE events IN ACCESS SHARE MODE")
+            .execute(&mut *blocker)
+            .await
+            .expect("lock partition parent");
+
+        let result = tokio::time::timeout(
+            std::time::Duration::from_millis(500),
+            ensure_future_partitions(&pool, 0),
+        )
+        .await;
+
+        blocker.rollback().await.expect("release parent lock");
+        drop_scratch_db(&admin, pool, &name).await;
+        result
+            .expect("default coverage must return without waiting for a DDL lock")
+            .expect("default partitions cover the target range");
+    }
+
+    #[tokio::test]
+    #[ignore = "requires Postgres"]
+    async fn unrecognized_partition_bounds_fail_closed() {
+        let admin = admin_pool().await;
+        let (pool, name) = create_scratch_db(&admin, "partition_unrecognized").await;
+        sqlx::query("DROP TABLE events CASCADE")
+            .execute(&pool)
+            .await
+            .expect("replace events test table");
+        for statement in [
+            "CREATE TABLE events (created_at TIMESTAMPTZ NOT NULL, sequence INT NOT NULL) \
+             PARTITION BY RANGE (created_at, sequence)",
+            "CREATE TABLE events_multicolumn PARTITION OF events \
+             FOR VALUES FROM (MINVALUE, MINVALUE) TO (MAXVALUE, MAXVALUE)",
+        ] {
+            sqlx::query(statement)
+                .execute(&pool)
+                .await
+                .expect("create unsupported partition topology");
+        }
+
+        let error = ensure_future_partitions(&pool, 0)
+            .await
+            .expect_err("unsupported bounds must fail closed before DDL");
+        assert!(
+            matches!(
+                error,
+                DbError::InvalidData(ref message)
+                    if message.contains("unsupported partition")
+            ),
+            "unexpected error: {error}"
+        );
+
+        drop_scratch_db(&admin, pool, &name).await;
+    }
+
+    #[tokio::test]
+    #[ignore = "requires Postgres"]
+    async fn nested_partition_bounds_fail_closed() {
+        let admin = admin_pool().await;
+        let (pool, name) = create_scratch_db(&admin, "partition_nested").await;
+        sqlx::query("DROP TABLE events CASCADE")
+            .execute(&pool)
+            .await
+            .expect("replace events test table");
+        for statement in [
+            "CREATE TABLE events (created_at TIMESTAMPTZ NOT NULL, sequence INT NOT NULL) \
+             PARTITION BY RANGE (created_at)",
+            "CREATE TABLE events_all PARTITION OF events \
+             FOR VALUES FROM (MINVALUE) TO (MAXVALUE) PARTITION BY RANGE (sequence)",
+            "CREATE TABLE events_nested_default PARTITION OF events_all DEFAULT",
+        ] {
+            sqlx::query(statement)
+                .execute(&pool)
+                .await
+                .expect("create nested partition topology");
+        }
+
+        let error = ensure_future_partitions(&pool, 0)
+            .await
+            .expect_err("nested bounds must fail closed before DDL");
+        assert!(
+            matches!(
+                error,
+                DbError::InvalidData(ref message)
+                    if message.contains("unsupported partition")
+            ),
+            "unexpected error: {error}"
+        );
+
+        drop_scratch_db(&admin, pool, &name).await;
+    }
+
+    #[tokio::test]
+    #[ignore = "requires Postgres"]
+    async fn exact_utc_bounds_are_covered_in_non_utc_session() {
+        let admin = admin_pool().await;
+        let (pool, name) =
+            create_scratch_db_with_timezone(&admin, "partition_non_utc", "America/Los_Angeles")
+                .await;
+        let timezone: String = sqlx::query_scalar("SHOW timezone")
+            .fetch_one(&pool)
+            .await
+            .expect("read scratch database timezone");
+        assert_eq!(timezone, "America/Los_Angeles");
+        sqlx::query("DROP TABLE events_p_future, delivery_log_p_future")
+            .execute(&pool)
+            .await
+            .expect("remove catch-all partitions");
+
+        let now = Utc::now();
+        let start = Utc
+            .with_ymd_and_hms(now.year(), now.month(), 1, 0, 0, 0)
+            .single()
+            .expect("current month starts at a valid date");
+        let (end_year, end_month) = if now.month() == 12 {
+            (now.year() + 1, 1)
+        } else {
+            (now.year(), now.month() + 1)
+        };
+        let end = Utc
+            .with_ymd_and_hms(end_year, end_month, 1, 0, 0, 0)
+            .single()
+            .expect("next month starts at a valid date");
+        for statement in [
+            format!(
+                "CREATE TABLE events_existing_range PARTITION OF events \
+                 FOR VALUES FROM ('{}') TO ('{}')",
+                start.to_rfc3339(),
+                end.to_rfc3339()
+            ),
+            format!(
+                "CREATE TABLE delivery_log_existing_range PARTITION OF delivery_log \
+                 FOR VALUES FROM ('{}') TO ('{}')",
+                start.to_rfc3339(),
+                end.to_rfc3339()
+            ),
+        ] {
+            sqlx::query(sqlx::AssertSqlSafe(statement))
+                .execute(&pool)
+                .await
+                .expect("create exact UTC monthly partition");
+        }
+
+        let mut blocker = pool.begin().await.expect("begin blocker");
+        sqlx::query("LOCK TABLE events IN ACCESS SHARE MODE")
+            .execute(&mut *blocker)
+            .await
+            .expect("lock partition parent");
+        let result = tokio::time::timeout(
+            std::time::Duration::from_millis(500),
+            ensure_future_partitions(&pool, 0),
+        )
+        .await;
+
+        blocker.rollback().await.expect("release parent lock");
+        drop_scratch_db(&admin, pool, &name).await;
+        result
+            .expect("UTC partition coverage must not depend on the session timezone")
+            .expect("exact UTC monthly partitions are already ensured");
+    }
 }