From 1345e3f57b754c0803481d4714b2856beb3bf177 Mon Sep 17 00:00:00 2001 From: Melody Ruch <396921921@qq.com> Date: Tue, 18 Aug 2026 13:18:18 +0800 Subject: [PATCH] fix(postgres): include partition children when generating table DDL --- crates/dbx-core/src/db/cloudberry.rs | 46 +- crates/dbx-core/src/db/postgres.rs | 1140 ++++++++++++++++- crates/dbx-core/src/schema.rs | 532 +++++++- .../dbx-core/src/schema/providers/native.rs | 4 +- .../tests/database_export_partition_ddl.rs | 365 ++++++ crates/dbx-core/tests/live_cloudberry.rs | 2 +- crates/dbx-core/tests/live_opentenbase.rs | 2 +- 7 files changed, 1993 insertions(+), 98 deletions(-) create mode 100644 crates/dbx-core/tests/database_export_partition_ddl.rs diff --git a/crates/dbx-core/src/db/cloudberry.rs b/crates/dbx-core/src/db/cloudberry.rs index 4297f0a18a..e043030c2e 100644 --- a/crates/dbx-core/src/db/cloudberry.rs +++ b/crates/dbx-core/src/db/cloudberry.rs @@ -151,17 +151,32 @@ pub fn append_table_modifiers(ddl: &str, modifiers: &TableModifiers) -> Result Result { + // The fallback DDL may already read `CREATE FOREIGN TABLE` — a foreign + // table's own relkind ('f') makes the generic renderer emit that header + // directly — or plain `CREATE TABLE` from an older/other rendering path. + // Accept either so this doesn't regress when the upstream header changes. + let create_foreign_table = "CREATE FOREIGN TABLE "; let create_table = "CREATE TABLE "; - if !ddl.starts_with(create_table) { + let header_len = if ddl.starts_with(create_foreign_table) { + create_foreign_table.len() + } else if ddl.starts_with(create_table) { + create_table.len() + } else { return Err("Cloudberry external-table fallback expected CREATE TABLE DDL".to_string()); - } + }; let insertion = ddl .find(";\n") .or_else(|| ddl.find(';')) .ok_or_else(|| "Cloudberry fallback DDL has no CREATE TABLE terminator".to_string())?; + // When the fallback DDL already read `CREATE FOREIGN TABLE`, the generic + // renderer also already appended its own ` SERVER ... [OPTIONS (...)]` + // clause (from the same relkind='f' catalog data). Drop that so this + // doesn't emit the clause twice — Cloudberry's own `external` definition + // above is authoritative here. + let body_end = ddl[header_len..insertion].find(" SERVER \"").map(|pos| header_len + pos).unwrap_or(insertion); let mut output = String::with_capacity(ddl.len() + external.options.len() * 24 + 48); output.push_str("CREATE FOREIGN TABLE "); - output.push_str(&ddl[create_table.len()..insertion]); + output.push_str(&ddl[header_len..body_end]); output.push_str("\nSERVER "); output.push_str(&db::postgres::pg_quote_ident(&external.server)); if !external.options.is_empty() { @@ -324,6 +339,31 @@ mod tests { assert!(rendered.contains("\"null\" ''")); } + #[test] + fn renders_external_table_from_already_foreign_ddl_without_duplicate_server_clause() { + // The generic renderer now derives `is_foreign` straight from relkind + // and may hand back a `CREATE FOREIGN TABLE ... SERVER ...` fallback + // DDL directly, not just plain `CREATE TABLE`. + let ddl = + "CREATE FOREIGN TABLE \"public\".\"external_events\" (\n \"id\" integer\n) SERVER \"stale_server\";\n"; + let rendered = append_table_modifiers( + ddl, + &TableModifiers { + external: Some(ExternalTableDefinition { + server: "gp_exttable_server".to_string(), + options: vec!["format=csv".to_string()], + }), + ..modifiers(None) + }, + ) + .unwrap(); + + assert!(rendered.starts_with("CREATE FOREIGN TABLE \"public\".\"external_events\"")); + assert_eq!(rendered.matches("SERVER").count(), 1); + assert!(rendered.contains("SERVER \"gp_exttable_server\"")); + assert!(!rendered.contains("stale_server")); + } + #[test] fn ddl_queries_avoid_privileged_external_catalog_view() { assert_eq!(CLOUD_BERRY_TABLE_DDL_SQL, "SELECT pg_get_tabledef($1, $2, true)"); diff --git a/crates/dbx-core/src/db/postgres.rs b/crates/dbx-core/src/db/postgres.rs index 89cc0e5f0e..ee5a4117aa 100644 --- a/crates/dbx-core/src/db/postgres.rs +++ b/crates/dbx-core/src/db/postgres.rs @@ -11,7 +11,7 @@ use rustls::server::ParsedCertificate; use sqlparser::ast::Statement; use sqlparser::dialect::PostgreSqlDialect; use sqlparser::parser::Parser; -use std::collections::{BTreeSet, HashMap}; +use std::collections::{BTreeMap, BTreeSet, HashMap}; use std::fs::File; use std::future::Future; use std::io::BufReader; @@ -81,6 +81,28 @@ pub struct PostgresTablePartitionInfo { pub parent_table: Option, pub bound: Option, pub key: Option, + /// True when this relation is itself a foreign table (relkind 'f'), + /// whether or not it is also a partition. A partition backed by a + /// foreign server must be declared with `CREATE FOREIGN TABLE ... + /// SERVER`, not `CREATE TABLE`. + pub is_foreign: bool, + pub foreign_server: Option, + pub foreign_options: Vec<(String, String)>, +} + +/// The state of a partition-local column DEFAULT relative to the parent's +/// default for the same column (a partition with no entry in the owning +/// `column_defaults` map simply inherits the parent's default unchanged). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum PostgresColumnDefaultState { + /// This partition's own default differs from (or has no counterpart in) + /// the parent's default — declared via `column WITH OPTIONS DEFAULT ...`. + Overridden, + /// The parent has a default for this column, but this partition + /// explicitly ran `ALTER TABLE ONLY ... ALTER COLUMN ... DROP DEFAULT` + /// to remove its own (otherwise auto-inherited-at-creation) copy — + /// replayed as a standalone `ALTER TABLE ONLY ... DROP DEFAULT;`. + Dropped, } #[derive(Debug, Clone, Default, PartialEq, Eq)] @@ -88,6 +110,13 @@ pub struct PostgresTablePartitionLocalObjects { pub has_primary_key: bool, pub foreign_keys: BTreeSet, pub indexes: BTreeSet, + /// CHECK constraints with a local definition on this partition, as + /// reported by pg_constraint.conislocal. A merged constraint may remain + /// local even when coninhcount is greater than zero. + pub check_constraints: BTreeSet, + /// Columns with a local override or explicit drop of the parent's + /// default; a column absent from this map is purely inherited. + pub column_defaults: BTreeMap, } fn pg_temporal_to_json_value(row: &Row, idx: usize) -> Option { @@ -2798,7 +2827,8 @@ pub async fn get_table_partition_info( }; let relkind = relation.try_get::<_, String>(0).unwrap_or_default(); let is_partition = relation.try_get::<_, bool>(1).unwrap_or(false); - if relkind != "p" && !is_partition { + let is_foreign = relkind == "f"; + if relkind != "p" && !is_foreign && !is_partition { return Ok(PostgresTablePartitionInfo::default()); } @@ -2806,14 +2836,25 @@ pub async fn get_table_partition_info( .await .map_err(|e| e.to_string())?; let Some(row) = rows.first() else { - return Ok(PostgresTablePartitionInfo { is_partition, ..Default::default() }); + return Ok(PostgresTablePartitionInfo { is_partition, is_foreign, ..Default::default() }); }; + let foreign_options = row + .try_get::<_, Option>>(5) + .ok() + .flatten() + .unwrap_or_default() + .into_iter() + .filter_map(|option| option.split_once('=').map(|(key, value)| (key.to_string(), value.to_string()))) + .collect(); Ok(PostgresTablePartitionInfo { is_partition, parent_schema: row.try_get::<_, Option>(0).ok().flatten().filter(|value| !value.is_empty()), parent_table: row.try_get::<_, Option>(1).ok().flatten().filter(|value| !value.is_empty()), bound: row.try_get::<_, Option>(2).ok().flatten().filter(|value| !value.is_empty()), key: row.try_get::<_, Option>(3).ok().flatten().filter(|value| !value.is_empty()), + is_foreign, + foreign_server: row.try_get::<_, Option>(4).ok().flatten().filter(|value| !value.is_empty()), + foreign_options, }) } @@ -2821,6 +2862,40 @@ pub async fn get_table_partition_key(pool: &Pool, schema: &str, table: &str) -> Ok(get_table_partition_info(pool, schema, table).await?.key) } +/// Classifies one row of `postgres_table_partition_local_objects_sql` (or its +/// `_for_relations` sibling, which has the same `object_kind`/`object_name`/ +/// `object_type` columns plus a leading `relid`) into `entry`. Shared so a +/// future local-object kind, or a fix to how a kind is recognized, can't be +/// applied to only one of the two query paths. +fn apply_partition_local_object_row( + entry: &mut PostgresTablePartitionLocalObjects, + object_kind: &str, + object_name: String, + object_type: &str, +) { + match object_kind { + "constraint" if object_type == "p" => entry.has_primary_key = true, + "constraint" if object_type == "f" && !object_name.is_empty() => { + entry.foreign_keys.insert(object_name); + } + "check" if !object_name.is_empty() => { + entry.check_constraints.insert(object_name); + } + "index" if !object_name.is_empty() => { + entry.indexes.insert(object_name); + } + "column_default" if !object_name.is_empty() => { + let state = if object_type == "dropped" { + PostgresColumnDefaultState::Dropped + } else { + PostgresColumnDefaultState::Overridden + }; + entry.column_defaults.insert(object_name, state); + } + _ => {} + } +} + pub async fn get_table_partition_local_objects( pool: &Pool, schema: &str, @@ -2836,20 +2911,636 @@ pub async fn get_table_partition_local_objects( let object_kind = row.try_get::<_, String>(0).unwrap_or_default(); let object_name = row.try_get::<_, String>(1).unwrap_or_default(); let object_type = row.try_get::<_, Option>(2).ok().flatten().unwrap_or_default(); - match object_kind.as_str() { - "constraint" if object_type == "p" => result.has_primary_key = true, - "constraint" if object_type == "f" && !object_name.is_empty() => { - result.foreign_keys.insert(object_name); + apply_partition_local_object_row(&mut result, &object_kind, object_name, &object_type); + } + Ok(result) +} + +/// One relation in a partition tree, as discovered by +/// `fetch_postgres_partition_tree`. +#[derive(Debug, Clone)] +pub struct PostgresPartitionTreeNode { + pub oid: i64, + pub schema: String, + pub table: String, + pub parent_oid: Option, + pub parent_schema: Option, + pub parent_table: Option, + pub partition_info: PostgresTablePartitionInfo, +} + +fn postgres_partition_tree_sql() -> &'static str { + "WITH RECURSIVE tree AS ( \ + SELECT c.oid::bigint AS oid, n.nspname::text AS schema, c.relname::text AS relname, \ + ap.oid::bigint AS parent_oid, an.nspname::text COLLATE \"C\" AS parent_schema, ap.relname::text COLLATE \"C\" AS parent_relname, \ + c.relkind::text AS relkind, ARRAY[c.oid]::bigint[] AS path \ + FROM pg_catalog.pg_class c \ + JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ + LEFT JOIN pg_catalog.pg_inherits ai ON ai.inhrelid = c.oid AND c.relispartition \ + LEFT JOIN pg_catalog.pg_class ap ON ap.oid = ai.inhparent \ + LEFT JOIN pg_catalog.pg_namespace an ON an.oid = ap.relnamespace \ + WHERE n.nspname = $1 AND c.relname = $2 AND c.relkind IN ('r','p','f') \ + UNION ALL \ + SELECT c.oid::bigint, n.nspname::text, c.relname::text, tree.oid, tree.schema, tree.relname, c.relkind::text, \ + tree.path || c.oid::bigint \ + FROM pg_catalog.pg_inherits i \ + JOIN pg_catalog.pg_class c ON c.oid = i.inhrelid AND c.relispartition \ + JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ + JOIN tree ON tree.oid = i.inhparent \ + WHERE NOT c.oid = ANY(tree.path) \ + ) \ + SELECT t.oid, t.schema, t.relname, t.parent_oid, t.parent_schema, t.parent_relname, t.relkind, \ + pg_catalog.pg_get_expr(c.relpartbound, c.oid, true) AS partition_bound, \ + CASE WHEN t.relkind = 'p' THEN pg_catalog.pg_get_partkeydef(c.oid) ELSE NULL END AS partition_key, \ + fs.srvname AS foreign_server, \ + ft.ftoptions AS foreign_options \ + FROM tree t \ + JOIN pg_catalog.pg_class c ON c.oid = t.oid \ + LEFT JOIN pg_catalog.pg_foreign_table ft ON ft.ftrelid = c.oid \ + LEFT JOIN pg_catalog.pg_foreign_server fs ON fs.oid = ft.ftserver \ + ORDER BY t.oid" +} + +/// The whole partition tree rooted at (schema, table) — the root itself plus +/// every descendant partition at any depth — plus each node's partition info, +/// in a single round trip. Used by `pg_ddl_with_partitions` instead of +/// recursing per-relation (which reruns the full metadata query chain once +/// per node and scales the request count linearly with tree size). +pub async fn fetch_postgres_partition_tree( + pool: &Pool, + schema: &str, + table: &str, +) -> Result, String> { + let schema = if schema.is_empty() { "public" } else { schema }; + let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; + let rows = postgres_query_cached(&client, postgres_partition_tree_sql(), &[&schema, &table]) + .await + .map_err(pg_error_to_string)?; + Ok(rows + .into_iter() + .filter_map(|row| { + let oid = row.try_get::<_, i64>(0).ok()?; + let node_schema = row.try_get::<_, String>(1).ok()?; + let node_table = row.try_get::<_, String>(2).ok()?; + let parent_oid = row.try_get::<_, Option>(3).ok().flatten(); + let parent_schema = row.try_get::<_, Option>(4).ok().flatten(); + let parent_table = row.try_get::<_, Option>(5).ok().flatten(); + let relkind = row.try_get::<_, String>(6).unwrap_or_default(); + let is_foreign = relkind == "f"; + let foreign_options = row + .try_get::<_, Option>>(10) + .ok() + .flatten() + .unwrap_or_default() + .into_iter() + .filter_map(|option| option.split_once('=').map(|(key, value)| (key.to_string(), value.to_string()))) + .collect(); + let partition_info = PostgresTablePartitionInfo { + is_partition: parent_oid.is_some(), + parent_schema: parent_schema.clone().filter(|value| !value.is_empty()), + parent_table: parent_table.clone().filter(|value| !value.is_empty()), + bound: row.try_get::<_, Option>(7).ok().flatten().filter(|value| !value.is_empty()), + key: row.try_get::<_, Option>(8).ok().flatten().filter(|value| !value.is_empty()), + is_foreign, + foreign_server: row.try_get::<_, Option>(9).ok().flatten().filter(|value| !value.is_empty()), + foreign_options, + }; + Some(PostgresPartitionTreeNode { + oid, + schema: node_schema, + table: node_table, + parent_oid, + parent_schema, + parent_table, + partition_info, + }) + }) + .collect()) +} + +/// Batched sibling of get_columns, fetching every relation through at most +/// two OID-scoped catalog queries. The compatibility tier avoids PostgreSQL +/// 12+ catalog fields so PostgreSQL 10-compatible servers remain bounded. +pub async fn get_columns_for_relations( + pool: &Pool, + relations: &[(i64, String, String)], +) -> Result>, String> { + let oids: Vec = relations.iter().map(|(oid, _, _)| *oid).collect(); + let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; + let [primary_sql, compat_sql] = postgres_columns_for_relations_query_tiers(); + match get_columns_for_relations_with_sql(&client, primary_sql, &oids).await { + Ok(columns) => Ok(columns), + Err(primary_error) => match get_columns_for_relations_with_sql(&client, compat_sql, &oids).await { + Ok(columns) => Ok(columns), + Err(fallback_error) => { + let primary_message = pg_error_to_string(primary_error); + let fallback_message = pg_error_to_string(fallback_error); + log::debug!( + "[postgres][get_columns_for_relations:compat-failed] primary_error={} fallback_error={}", + primary_message, + fallback_message + ); + Err(fallback_message) } - "index" if !object_name.is_empty() => { - result.indexes.insert(object_name); + }, + } +} + +async fn get_columns_for_relations_with_sql( + client: &deadpool_postgres::Client, + sql: &str, + oids: &[i64], +) -> Result>, tokio_postgres::Error> { + let rows = postgres_query_cached(client, sql, &[&oids]).await?; + let mut result: HashMap> = HashMap::new(); + for row in &rows { + let Ok(relid) = row.try_get::<_, i64>(0) else { continue }; + result.entry(relid).or_default().push(column_info_from_row_offset(row, 1)); + } + Ok(result) +} + +fn postgres_columns_for_relations_query_tiers() -> [&'static str; 2] { + [postgres_columns_for_relations_sql(), postgres_columns_for_relations_compat_sql()] +} + +fn postgres_columns_for_relations_sql() -> &'static str { + "SELECT c.oid::bigint AS relid, a.attname AS column_name, \ + format_type(a.atttypid, a.atttypmod) AS full_type, \ + COALESCE(ic.is_nullable = 'YES', NOT a.attnotnull) AS is_nullable, \ + CASE WHEN a.attgenerated <> '' THEN NULL ELSE pg_get_expr(ad.adbin, ad.adrelid) END AS column_default, \ + EXISTS ( \ + SELECT 1 FROM pg_constraint co \ + JOIN pg_index i ON i.indrelid = co.conrelid AND co.conindid = i.indexrelid \ + WHERE co.conrelid = a.attrelid AND co.contype = 'p' \ + AND a.attnum = ANY(i.indkey) \ + ) AS is_pk, \ + col_description(a.attrelid, a.attnum) AS column_comment, \ + CASE a.attidentity \ + WHEN 'd' THEN 'generated by default as identity' || CASE WHEN pseq.seqstart IS NOT NULL THEN format(' (start with %s increment by %s)', pseq.seqstart, pseq.seqincrement) ELSE '' END \ + WHEN 'a' THEN 'generated always as identity' || CASE WHEN pseq.seqstart IS NOT NULL THEN format(' (start with %s increment by %s)', pseq.seqstart, pseq.seqincrement) ELSE '' END \ + ELSE CASE a.attgenerated \ + WHEN 's' THEN 'generated always as (' || pg_get_expr(ad.adbin, ad.adrelid) || ') stored' \ + WHEN 'v' THEN 'generated always as (' || pg_get_expr(ad.adbin, ad.adrelid) || ') virtual' \ + ELSE CASE WHEN a.atttypid IN (20, 21, 23) AND dep.deptype = 'a' \ + AND pseq.seqrelid IS NOT NULL AND default_dep.objid IS NOT NULL \ + AND pg_get_expr(ad.adbin, ad.adrelid) = format('nextval(%L::regclass)', dep.objid::regclass::text) \ + THEN CASE a.atttypid \ + WHEN 21 THEN 'smallserial' \ + WHEN 23 THEN 'serial' \ + WHEN 20 THEN 'bigserial' \ + END ELSE NULL END \ + END \ + END AS column_extra, \ + CASE WHEN t.typname = 'numeric' AND a.atttypmod > 0 \ + THEN ((a.atttypmod - 4) >> 16) & 65535 ELSE NULL END AS numeric_precision, \ + CASE WHEN t.typname = 'numeric' AND a.atttypmod > 0 \ + THEN (a.atttypmod - 4) & 65535 ELSE NULL END AS numeric_scale, \ + CASE WHEN t.typname IN ('varchar', 'bpchar') AND a.atttypmod > 0 \ + THEN a.atttypmod - 4 ELSE NULL END AS character_maximum_length, \ + CASE WHEN enum_t.oid IS NULL THEN NULL \ + ELSE COALESCE((SELECT array_to_json(array_agg(e.enumlabel ORDER BY e.enumsortorder))::text \ + FROM pg_enum e WHERE e.enumtypid = enum_t.oid), '[]') END AS enum_values \ + FROM pg_attribute a \ + JOIN pg_class c ON c.oid = a.attrelid \ + JOIN pg_namespace n ON n.oid = c.relnamespace \ + JOIN pg_type t ON t.oid = a.atttypid \ + LEFT JOIN pg_type enum_t ON enum_t.oid = CASE WHEN t.typtype = 'd' THEN t.typbasetype WHEN t.typtype = 'e' THEN t.oid ELSE NULL END AND enum_t.typtype = 'e' \ + LEFT JOIN pg_attrdef ad ON ad.adrelid = a.attrelid AND ad.adnum = a.attnum \ + LEFT JOIN pg_depend dep ON dep.classid = 'pg_catalog.pg_class'::pg_catalog.regclass \ + AND dep.refclassid = 'pg_catalog.pg_class'::pg_catalog.regclass \ + AND dep.refobjid = a.attrelid AND dep.refobjsubid = a.attnum AND dep.deptype IN ('a', 'i') \ + LEFT JOIN pg_sequence pseq ON pseq.seqrelid = dep.objid \ + LEFT JOIN pg_catalog.pg_depend default_dep \ + ON default_dep.classid = 'pg_catalog.pg_attrdef'::pg_catalog.regclass \ + AND default_dep.objid = ad.oid \ + AND default_dep.objsubid = 0 \ + AND default_dep.refclassid = 'pg_catalog.pg_class'::pg_catalog.regclass \ + AND default_dep.refobjid = dep.objid \ + AND default_dep.refobjsubid = 0 \ + AND default_dep.deptype = 'n' \ + LEFT JOIN information_schema.columns ic \ + ON ic.table_schema = n.nspname AND ic.table_name = c.relname AND ic.column_name = a.attname \ + WHERE c.oid = ANY($1::bigint[]) \ + AND a.attnum > 0 AND NOT a.attisdropped \ + ORDER BY c.oid, a.attnum" +} + +fn postgres_columns_for_relations_compat_sql() -> &'static str { + "SELECT c.oid::bigint AS relid, a.attname AS column_name, \ + format_type(a.atttypid, a.atttypmod) AS full_type, \ + COALESCE(ic.is_nullable = 'YES', NOT a.attnotnull) AS is_nullable, \ + pg_get_expr(ad.adbin, ad.adrelid) AS column_default, \ + EXISTS ( \ + SELECT 1 FROM pg_constraint co \ + JOIN pg_index i ON i.indrelid = co.conrelid AND co.conindid = i.indexrelid \ + WHERE co.conrelid = a.attrelid AND co.contype = 'p' \ + AND a.attnum = ANY(i.indkey) \ + ) AS is_pk, \ + col_description(a.attrelid, a.attnum) AS column_comment, \ + CASE WHEN a.atttypid IN (20, 21, 23) AND serial_seq.oid IS NOT NULL \ + AND serial_default_dep.objid IS NOT NULL \ + AND pg_get_expr(ad.adbin, ad.adrelid) = format('nextval(%L::regclass)', serial_seq.oid::regclass::text) \ + THEN CASE a.atttypid \ + WHEN 21 THEN 'smallserial' \ + WHEN 23 THEN 'serial' \ + WHEN 20 THEN 'bigserial' \ + END ELSE NULL END AS column_extra, \ + CASE WHEN t.typname = 'numeric' AND a.atttypmod > 0 \ + THEN ((a.atttypmod - 4) >> 16) & 65535 ELSE NULL END AS numeric_precision, \ + CASE WHEN t.typname = 'numeric' AND a.atttypmod > 0 \ + THEN (a.atttypmod - 4) & 65535 ELSE NULL END AS numeric_scale, \ + CASE WHEN t.typname IN ('varchar', 'bpchar') AND a.atttypmod > 0 \ + THEN a.atttypmod - 4 ELSE NULL END AS character_maximum_length, \ + NULL::text AS enum_values \ + FROM pg_attribute a \ + JOIN pg_class c ON c.oid = a.attrelid \ + JOIN pg_namespace n ON n.oid = c.relnamespace \ + JOIN pg_type t ON t.oid = a.atttypid \ + LEFT JOIN pg_attrdef ad ON ad.adrelid = a.attrelid AND ad.adnum = a.attnum \ + LEFT JOIN pg_catalog.pg_depend serial_dep \ + ON serial_dep.classid = 'pg_catalog.pg_class'::pg_catalog.regclass \ + AND serial_dep.refclassid = 'pg_catalog.pg_class'::pg_catalog.regclass \ + AND serial_dep.refobjid = a.attrelid AND serial_dep.refobjsubid = a.attnum \ + AND serial_dep.deptype = 'a' \ + LEFT JOIN pg_catalog.pg_class serial_seq ON serial_seq.oid = serial_dep.objid AND serial_seq.relkind = 'S' \ + LEFT JOIN pg_catalog.pg_depend serial_default_dep \ + ON serial_default_dep.classid = 'pg_catalog.pg_attrdef'::pg_catalog.regclass \ + AND serial_default_dep.objid = ad.oid \ + AND serial_default_dep.objsubid = 0 \ + AND serial_default_dep.refclassid = 'pg_catalog.pg_class'::pg_catalog.regclass \ + AND serial_default_dep.refobjid = serial_seq.oid \ + AND serial_default_dep.refobjsubid = 0 \ + AND serial_default_dep.deptype = 'n' \ + LEFT JOIN information_schema.columns ic \ + ON ic.table_schema = n.nspname AND ic.table_name = c.relname AND ic.column_name = a.attname \ + WHERE c.oid = ANY($1::bigint[]) \ + AND a.attnum > 0 AND NOT a.attisdropped \ + ORDER BY c.oid, a.attnum" +} + +/// Same field layout as `column_info_from_row`, offset by one leading `relid` +/// column. +fn column_info_from_row_offset(row: &Row, offset: usize) -> ColumnInfo { + let full_type = row.try_get::<_, Option>(offset + 1).ok().flatten().unwrap_or_default(); + ColumnInfo { + name: pg_row_try_string(row, offset), + data_type: full_type, + is_nullable: pg_row_try_bool(row, offset + 2).unwrap_or(true), + column_default: row.try_get::<_, Option>(offset + 3).ok().flatten(), + is_primary_key: pg_row_try_bool(row, offset + 4).unwrap_or(false), + extra: row.try_get::<_, Option>(offset + 6).ok().flatten(), + comment: row.try_get::<_, Option>(offset + 5).ok().flatten(), + numeric_precision: row.try_get::<_, Option>(offset + 7).ok().flatten(), + numeric_scale: row.try_get::<_, Option>(offset + 8).ok().flatten(), + character_maximum_length: row.try_get::<_, Option>(offset + 9).ok().flatten(), + enum_values: parse_enum_values_from_row(row, offset + 10), + ..Default::default() + } +} + +/// Batched sibling of list_indexes, using a modern and a PostgreSQL 10 +/// compatible OID-scoped tier so the request count stays bounded. +pub async fn list_indexes_for_relations( + pool: &Pool, + relations: &[(i64, String, String)], +) -> Result>, String> { + let oids: Vec = relations.iter().map(|(oid, _, _)| *oid).collect(); + let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; + let [primary_sql, compat_sql] = postgres_indexes_for_relations_query_tiers(); + match list_indexes_for_relations_with_sql(&client, primary_sql, &oids).await { + Ok(indexes) => Ok(indexes), + Err(primary_error) => match list_indexes_for_relations_with_sql(&client, compat_sql, &oids).await { + Ok(indexes) => Ok(indexes), + Err(fallback_error) => { + let primary_message = pg_error_to_string(primary_error); + let fallback_message = pg_error_to_string(fallback_error); + log::debug!( + "[postgres][list_indexes_for_relations:compat-failed] primary_error={} fallback_error={}", + primary_message, + fallback_message + ); + Err(fallback_message) } - _ => {} + }, + } +} + +async fn list_indexes_for_relations_with_sql( + client: &deadpool_postgres::Client, + sql: &str, + oids: &[i64], +) -> Result>, tokio_postgres::Error> { + let rows = postgres_query_cached(client, sql, &[&oids]).await?; + let mut result: HashMap> = HashMap::new(); + for row in &rows { + let Ok(relid) = row.try_get::<_, i64>(0) else { continue }; + let all_cols: Vec = row.try_get::<_, Vec>(2).unwrap_or_default(); + let nkeyatts = row.try_get::<_, Option>(7).ok().flatten().unwrap_or(all_cols.len() as i16) as usize; + let split_at = nkeyatts.min(all_cols.len()); + let key_cols = all_cols[..split_at].to_vec(); + let included = if split_at < all_cols.len() { all_cols[split_at..].to_vec() } else { vec![] }; + let all_is_expr: Vec = row.try_get::<_, Vec>(10).unwrap_or_default(); + let key_is_expression = + if all_is_expr.len() == all_cols.len() { all_is_expr[..split_at].to_vec() } else { Vec::new() }; + result.entry(relid).or_default().push(IndexInfo { + name: pg_row_try_string(row, 1), + columns: key_cols, + is_unique: pg_row_try_bool(row, 3).unwrap_or(false), + is_primary: pg_row_try_bool(row, 4).unwrap_or(false), + filter: row.try_get::<_, Option>(5).ok().flatten(), + index_type: row.try_get::<_, Option>(6).ok().flatten(), + included_columns: if included.is_empty() { None } else { Some(included) }, + comment: row.try_get::<_, Option>(9).ok().flatten(), + key_is_expression, + }); + } + Ok(result) +} + +fn postgres_indexes_for_relations_query_tiers() -> [&'static str; 2] { + [postgres_indexes_for_relations_sql(), postgres_indexes_for_relations_compat_sql()] +} + +fn postgres_indexes_for_relations_sql() -> &'static str { + "SELECT t.oid::bigint AS relid, i.relname AS index_name, \ + array_agg(COALESCE(a.attname, pg_get_indexdef(ix.indexrelid, k.n::int, true)) ORDER BY k.n) AS columns, \ + ix.indisunique AS is_unique, \ + ix.indisprimary AS is_primary, \ + pg_get_expr(ix.indpred, ix.indrelid) AS filter_expr, \ + am.amname AS index_type, \ + ix.indnkeyatts AS nkeyatts, \ + ix.indkey AS indkey, \ + obj_description(i.oid, 'pg_class') AS index_comment, \ + array_agg(a.attname IS NULL ORDER BY k.n) AS key_is_expression \ + FROM pg_index ix \ + JOIN pg_class t ON t.oid = ix.indrelid \ + JOIN pg_class i ON i.oid = ix.indexrelid \ + JOIN pg_am am ON am.oid = i.relam \ + JOIN LATERAL unnest(ix.indkey) WITH ORDINALITY AS k(attnum, n) ON true \ + LEFT JOIN pg_attribute a ON a.attrelid = t.oid AND a.attnum = k.attnum AND k.attnum > 0 \ + WHERE t.oid = ANY($1::bigint[]) \ + GROUP BY t.oid, i.relname, i.oid, ix.indisunique, ix.indisprimary, ix.indpred, ix.indrelid, am.amname, ix.indnkeyatts, ix.indkey \ + ORDER BY t.oid, i.relname" +} + +fn postgres_indexes_for_relations_compat_sql() -> &'static str { + "SELECT t.oid::bigint AS relid, i.relname AS index_name, \ + ARRAY( \ + SELECT COALESCE(a.attname, pg_get_indexdef(ix.indexrelid, pos.n, true)) \ + FROM generate_series(1, array_length(string_to_array(ix.indkey::text, ' '), 1)) AS pos(n) \ + LEFT JOIN pg_attribute a \ + ON a.attrelid = t.oid \ + AND a.attnum = (string_to_array(ix.indkey::text, ' '))[pos.n]::int2 \ + AND a.attnum > 0 \ + ORDER BY pos.n \ + ) AS columns, \ + ix.indisunique AS is_unique, \ + ix.indisprimary AS is_primary, \ + pg_get_expr(ix.indpred, ix.indrelid) AS filter_expr, \ + am.amname AS index_type, \ + NULL::smallint AS nkeyatts, \ + ix.indkey AS indkey, \ + obj_description(i.oid, 'pg_class') AS index_comment, \ + ARRAY( \ + SELECT a.attname IS NULL \ + FROM generate_series(1, array_length(string_to_array(ix.indkey::text, ' '), 1)) AS pos(n) \ + LEFT JOIN pg_attribute a \ + ON a.attrelid = t.oid \ + AND a.attnum = (string_to_array(ix.indkey::text, ' '))[pos.n]::int2 \ + AND a.attnum > 0 \ + ORDER BY pos.n \ + ) AS key_is_expression \ + FROM pg_index ix \ + JOIN pg_class t ON t.oid = ix.indrelid \ + JOIN pg_class i ON i.oid = ix.indexrelid \ + JOIN pg_am am ON am.oid = i.relam \ + WHERE t.oid = ANY($1::bigint[]) \ + ORDER BY t.oid, i.relname" +} + +/// Batched sibling of `list_foreign_keys`, keyed by `(schema, table)` since +/// this query is `information_schema`-based rather than oid-based. +pub async fn list_foreign_keys_for_relations( + pool: &Pool, + relations: &[(String, String)], +) -> Result>, String> { + let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; + let schemas: Vec<&str> = relations.iter().map(|(schema, _)| schema.as_str()).collect(); + let tables: Vec<&str> = relations.iter().map(|(_, table)| table.as_str()).collect(); + let rows = postgres_query_cached(&client, postgres_foreign_keys_for_relations_sql(), &[&schemas, &tables]) + .await + .map_err(|e| e.to_string())?; + let mut result: HashMap<(String, String), Vec> = HashMap::new(); + for row in &rows { + let key = (pg_row_try_string(row, 0), pg_row_try_string(row, 1)); + result.entry(key).or_default().push(ForeignKeyInfo { + name: pg_row_try_string(row, 2), + column: pg_row_try_string(row, 3), + ref_schema: Some(pg_row_try_string(row, 4)), + ref_table: pg_row_try_string(row, 5), + ref_column: pg_row_try_string(row, 6), + on_update: postgres_foreign_key_action(pg_row_try_string(row, 7)), + on_delete: postgres_foreign_key_action(pg_row_try_string(row, 8)), + }); + } + Ok(result) +} + +fn postgres_foreign_keys_for_relations_sql() -> &'static str { + "SELECT rel.rel_schema, rel.rel_table, \ + fk.constraint_name, fk.column_name, \ + pk.table_schema AS ref_schema, pk.table_name AS ref_table, pk.column_name AS ref_column, \ + rc.update_rule AS on_update, rc.delete_rule AS on_delete \ + FROM unnest($1::text[], $2::text[]) AS rel(rel_schema, rel_table) \ + JOIN information_schema.table_constraints tc \ + ON tc.table_schema = rel.rel_schema AND tc.table_name = rel.rel_table \ + JOIN information_schema.key_column_usage fk \ + ON fk.constraint_name = tc.constraint_name \ + AND fk.constraint_schema = tc.constraint_schema \ + AND fk.table_schema = tc.table_schema \ + AND fk.table_name = tc.table_name \ + JOIN information_schema.referential_constraints rc \ + ON rc.constraint_name = tc.constraint_name \ + AND rc.constraint_schema = tc.constraint_schema \ + JOIN information_schema.key_column_usage pk \ + ON pk.constraint_name = rc.unique_constraint_name \ + AND pk.constraint_schema = rc.unique_constraint_schema \ + AND pk.ordinal_position = fk.position_in_unique_constraint \ + WHERE tc.constraint_type = 'FOREIGN KEY' \ + ORDER BY rel.rel_schema, rel.rel_table, fk.constraint_name, fk.ordinal_position" +} + +/// Batched sibling of `get_table_comment`. +pub async fn get_table_comments_for_relations( + pool: &Pool, + oids: &[i64], +) -> Result>, String> { + let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; + let rows = postgres_query_cached( + &client, + "SELECT c.oid::bigint AS relid, obj_description(c.oid) AS table_comment \ + FROM pg_catalog.pg_class c WHERE c.oid = ANY($1::bigint[])", + &[&oids], + ) + .await + .map_err(|e| e.to_string())?; + let mut result: HashMap> = HashMap::new(); + for row in &rows { + let Ok(relid) = row.try_get::<_, i64>(0) else { continue }; + let comment = row.try_get::<_, Option>(1).ok().flatten().filter(|s| !s.is_empty()); + result.insert(relid, comment); + } + Ok(result) +} + +/// Batched sibling of `list_trigger_definitions`, with the same pre-13 +/// `tgparentid` capability probe. +pub async fn list_trigger_definitions_for_relations( + pool: &Pool, + oids: &[i64], +) -> Result>, String> { + let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; + let sql = if postgres_trigger_has_tgparentid(&client).await? { + postgres_trigger_definitions_for_relations_sql() + } else { + postgres_trigger_definitions_for_relations_sql_without_tgparentid() + }; + let rows = postgres_query_cached(&client, sql, &[&oids]).await.map_err(|e| e.to_string())?; + let mut result: HashMap> = HashMap::new(); + for row in &rows { + let Ok(relid) = row.try_get::<_, i64>(0) else { continue }; + let definition = pg_row_try_string(row, 1); + if !definition.trim().is_empty() { + result.entry(relid).or_default().push(definition); } } Ok(result) } +fn postgres_trigger_definitions_for_relations_sql() -> &'static str { + "SELECT t.tgrelid::bigint AS relid, pg_catalog.pg_get_triggerdef(t.oid, true) AS trigger_definition \ + FROM pg_catalog.pg_trigger t \ + WHERE t.tgrelid = ANY($1::bigint[]) AND NOT t.tgisinternal AND t.tgparentid = 0 \ + ORDER BY t.tgrelid, t.tgname, t.oid" +} + +fn postgres_trigger_definitions_for_relations_sql_without_tgparentid() -> &'static str { + "SELECT t.tgrelid::bigint AS relid, pg_catalog.pg_get_triggerdef(t.oid, true) AS trigger_definition \ + FROM pg_catalog.pg_trigger t \ + WHERE t.tgrelid = ANY($1::bigint[]) AND NOT t.tgisinternal \ + ORDER BY t.tgrelid, t.tgname, t.oid" +} + +/// Batched sibling of `list_check_constraints`. +pub async fn list_check_constraints_for_relations( + pool: &Pool, + oids: &[i64], +) -> Result>, String> { + let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; + let rows = postgres_query_cached( + &client, + "SELECT con.conrelid::bigint AS relid, con.conname, pg_catalog.pg_get_constraintdef(con.oid, true) AS definition \ + FROM pg_catalog.pg_constraint con \ + WHERE con.conrelid = ANY($1::bigint[]) AND con.contype = 'c' \ + ORDER BY con.conrelid, con.conname", + &[&oids], + ) + .await + .map_err(|e| e.to_string())?; + let mut result: HashMap> = HashMap::new(); + for row in &rows { + let Ok(relid) = row.try_get::<_, i64>(0) else { continue }; + let name = pg_row_try_string(row, 1); + let definition = pg_row_try_string(row, 2); + if !name.is_empty() && !definition.is_empty() { + result.entry(relid).or_default().push((name, definition)); + } + } + Ok(result) +} + +/// Batched sibling of `get_table_partition_local_objects`. +pub async fn get_table_partition_local_objects_for_relations( + pool: &Pool, + oids: &[i64], +) -> Result, String> { + let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; + let rows = postgres_query_cached(&client, postgres_table_partition_local_objects_for_relations_sql(), &[&oids]) + .await + .map_err(|e| e.to_string())?; + let mut result: HashMap = HashMap::new(); + for row in &rows { + let Ok(relid) = row.try_get::<_, i64>(0) else { continue }; + let object_kind = row.try_get::<_, String>(1).unwrap_or_default(); + let object_name = row.try_get::<_, String>(2).unwrap_or_default(); + let object_type = row.try_get::<_, Option>(3).ok().flatten().unwrap_or_default(); + apply_partition_local_object_row(result.entry(relid).or_default(), &object_kind, object_name, &object_type); + } + Ok(result) +} + +fn postgres_table_partition_local_objects_for_relations_sql() -> &'static str { + "SELECT con.conrelid::bigint AS relid, 'constraint'::text AS object_kind, con.conname AS object_name, con.contype::text AS object_type \ + FROM pg_catalog.pg_constraint con \ + WHERE con.conrelid = ANY($1::bigint[]) AND con.contype IN ('p','f') \ + AND COALESCE(NULLIF(pg_catalog.row_to_json(con)->>'conparentid', '')::oid, 0) = 0 \ + UNION ALL \ + SELECT con.conrelid::bigint, 'check'::text AS object_kind, con.conname AS object_name, NULL::text AS object_type \ + FROM pg_catalog.pg_constraint con \ + WHERE con.conrelid = ANY($1::bigint[]) AND con.contype = 'c' AND con.conislocal \ + UNION ALL \ + SELECT ix.indrelid::bigint, 'index'::text AS object_kind, idx.relname AS object_name, NULL::text AS object_type \ + FROM pg_catalog.pg_index ix \ + JOIN pg_catalog.pg_class idx ON idx.oid = ix.indexrelid \ + WHERE ix.indrelid = ANY($1::bigint[]) \ + AND NOT EXISTS (SELECT 1 FROM pg_catalog.pg_inherits i WHERE i.inhrelid = idx.oid) \ + UNION ALL \ + SELECT a.attrelid::bigint, 'column_default'::text AS object_kind, a.attname AS object_name, \ + CASE WHEN ad.oid IS NOT NULL THEN 'overridden' ELSE 'dropped' END AS object_type \ + FROM pg_catalog.pg_attribute a \ + JOIN pg_catalog.pg_inherits i ON i.inhrelid = a.attrelid \ + JOIN pg_catalog.pg_attribute pa \ + ON pa.attrelid = i.inhparent AND pa.attname = a.attname AND pa.attnum > 0 AND NOT pa.attisdropped \ + LEFT JOIN pg_catalog.pg_attrdef ad ON ad.adrelid = a.attrelid AND ad.adnum = a.attnum \ + LEFT JOIN pg_catalog.pg_attrdef pad ON pad.adrelid = pa.attrelid AND pad.adnum = pa.attnum \ + WHERE a.attrelid = ANY($1::bigint[]) AND a.attnum > 0 AND NOT a.attisdropped \ + AND ( \ + (ad.oid IS NOT NULL AND (pad.oid IS NULL \ + OR pg_catalog.pg_get_expr(ad.adbin, ad.adrelid) \ + IS DISTINCT FROM pg_catalog.pg_get_expr(pad.adbin, pad.adrelid))) \ + OR (ad.oid IS NULL AND pad.oid IS NOT NULL) \ + ) \ + ORDER BY relid, object_kind, object_name" +} + +/// CHECK constraints on a table, as (name, definition) pairs (definition +/// includes the leading `CHECK` keyword, e.g. `CHECK (id > 0)`). +pub async fn list_check_constraints(pool: &Pool, schema: &str, table: &str) -> Result, String> { + let schema = if schema.is_empty() { "public" } else { schema }; + let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; + let rows = postgres_query_cached(&client, postgres_check_constraints_sql(), &[&schema, &table]) + .await + .map_err(|e| e.to_string())?; + Ok(rows + .into_iter() + .filter_map(|row| { + let name = row.try_get::<_, String>(0).ok()?; + let definition = row.try_get::<_, String>(1).ok()?; + (!name.is_empty() && !definition.is_empty()).then_some((name, definition)) + }) + .collect()) +} + +fn postgres_check_constraints_sql() -> &'static str { + "SELECT con.conname, pg_catalog.pg_get_constraintdef(con.oid, true) AS definition \ + FROM pg_catalog.pg_constraint con \ + JOIN pg_catalog.pg_class c ON c.oid = con.conrelid \ + JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ + WHERE n.nspname = $1 AND c.relname = $2 AND con.contype = 'c' \ + ORDER BY con.conname" +} + fn postgres_table_partition_relation_sql() -> &'static str { "SELECT c.relkind::text, \ EXISTS ( \ @@ -2859,7 +3550,7 @@ fn postgres_table_partition_relation_sql() -> &'static str { ) AS is_partition \ FROM pg_catalog.pg_class c \ JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ - WHERE n.nspname = $1 AND c.relname = $2 AND c.relkind IN ('r','p') \ + WHERE n.nspname = $1 AND c.relname = $2 AND c.relkind IN ('r','p','f') \ LIMIT 1" } @@ -2867,13 +3558,17 @@ fn postgres_table_partition_info_sql() -> &'static str { "SELECT CASE WHEN c.relispartition THEN pn.nspname ELSE NULL END AS parent_schema, \ CASE WHEN c.relispartition THEN pc.relname ELSE NULL END AS parent_table, \ CASE WHEN c.relispartition THEN pg_catalog.pg_get_expr(c.relpartbound, c.oid, true) ELSE NULL END AS partition_bound, \ - CASE WHEN c.relkind = 'p' THEN pg_catalog.pg_get_partkeydef(c.oid) ELSE NULL END AS partition_key \ + CASE WHEN c.relkind = 'p' THEN pg_catalog.pg_get_partkeydef(c.oid) ELSE NULL END AS partition_key, \ + CASE WHEN c.relkind = 'f' THEN fs.srvname ELSE NULL END AS foreign_server, \ + CASE WHEN c.relkind = 'f' THEN ft.ftoptions ELSE NULL END AS foreign_options \ FROM pg_catalog.pg_class c \ JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ LEFT JOIN pg_catalog.pg_inherits i ON i.inhrelid = c.oid \ LEFT JOIN pg_catalog.pg_class pc ON pc.oid = i.inhparent \ LEFT JOIN pg_catalog.pg_namespace pn ON pn.oid = pc.relnamespace \ - WHERE n.nspname = $1 AND c.relname = $2 AND c.relkind IN ('r','p') \ + LEFT JOIN pg_catalog.pg_foreign_table ft ON ft.ftrelid = c.oid \ + LEFT JOIN pg_catalog.pg_foreign_server fs ON fs.oid = ft.ftserver \ + WHERE n.nspname = $1 AND c.relname = $2 AND c.relkind IN ('r','p','f') \ ORDER BY i.inhseqno NULLS LAST \ LIMIT 1" } @@ -2886,6 +3581,12 @@ fn postgres_table_partition_local_objects_sql() -> &'static str { WHERE n.nspname = $1 AND c.relname = $2 AND con.contype IN ('p','f') \ AND COALESCE(NULLIF(pg_catalog.row_to_json(con)->>'conparentid', '')::oid, 0) = 0 \ UNION ALL \ + SELECT 'check'::text AS object_kind, con.conname AS object_name, NULL::text AS object_type \ + FROM pg_catalog.pg_constraint con \ + JOIN pg_catalog.pg_class c ON c.oid = con.conrelid \ + JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ + WHERE n.nspname = $1 AND c.relname = $2 AND con.contype = 'c' AND con.conislocal \ + UNION ALL \ SELECT 'index'::text AS object_kind, idx.relname AS object_name, NULL::text AS object_type \ FROM pg_catalog.pg_index ix \ JOIN pg_catalog.pg_class c ON c.oid = ix.indrelid \ @@ -2893,6 +3594,24 @@ fn postgres_table_partition_local_objects_sql() -> &'static str { JOIN pg_catalog.pg_class idx ON idx.oid = ix.indexrelid \ WHERE n.nspname = $1 AND c.relname = $2 \ AND NOT EXISTS (SELECT 1 FROM pg_catalog.pg_inherits i WHERE i.inhrelid = idx.oid) \ + UNION ALL \ + SELECT 'column_default'::text AS object_kind, a.attname AS object_name, \ + CASE WHEN ad.oid IS NOT NULL THEN 'overridden' ELSE 'dropped' END AS object_type \ + FROM pg_catalog.pg_class c \ + JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ + JOIN pg_catalog.pg_attribute a ON a.attrelid = c.oid AND a.attnum > 0 AND NOT a.attisdropped \ + JOIN pg_catalog.pg_inherits i ON i.inhrelid = c.oid \ + JOIN pg_catalog.pg_attribute pa \ + ON pa.attrelid = i.inhparent AND pa.attname = a.attname AND pa.attnum > 0 AND NOT pa.attisdropped \ + LEFT JOIN pg_catalog.pg_attrdef ad ON ad.adrelid = a.attrelid AND ad.adnum = a.attnum \ + LEFT JOIN pg_catalog.pg_attrdef pad ON pad.adrelid = pa.attrelid AND pad.adnum = pa.attnum \ + WHERE n.nspname = $1 AND c.relname = $2 \ + AND ( \ + (ad.oid IS NOT NULL AND (pad.oid IS NULL \ + OR pg_catalog.pg_get_expr(ad.adbin, ad.adrelid) \ + IS DISTINCT FROM pg_catalog.pg_get_expr(pad.adbin, pad.adrelid))) \ + OR (ad.oid IS NULL AND pad.oid IS NOT NULL) \ + ) \ ORDER BY object_kind, object_name" } @@ -3260,6 +3979,27 @@ async fn postgres_proc_has_prokind(client: &deadpool_postgres::Client) -> Result Ok(pg_row_try_bool(&row, 0).unwrap_or(false)) } +fn postgres_trigger_has_tgparentid_sql() -> &'static str { + "SELECT EXISTS ( \ + SELECT 1 \ + FROM pg_catalog.pg_attribute \ + WHERE attrelid = 'pg_catalog.pg_trigger'::regclass \ + AND attname = 'tgparentid' \ + AND NOT attisdropped \ + )" +} + +/// PostgreSQL 13 added `pg_trigger.tgparentid`, tracking a partition's +/// automatically-cloned copy of its parent's trigger. Servers older than 13 +/// lack the column (and never clone triggers onto partitions in the first +/// place, so `tgisinternal` alone is sufficient there). +async fn postgres_trigger_has_tgparentid(client: &deadpool_postgres::Client) -> Result { + let row = postgres_query_one_cached(client, postgres_trigger_has_tgparentid_sql(), &[]) + .await + .map_err(|e| e.to_string())?; + Ok(pg_row_try_bool(&row, 0).unwrap_or(false)) +} + fn postgres_proc_has_prosp_sql() -> &'static str { "SELECT EXISTS ( \ SELECT 1 \ @@ -4345,21 +5085,7 @@ fn pg_row_try_string(row: &Row, idx: usize) -> String { } fn column_info_from_row(row: &Row) -> ColumnInfo { - let full_type = row.try_get::<_, Option>(1).ok().flatten().unwrap_or_default(); - ColumnInfo { - name: pg_row_try_string(row, 0), - data_type: full_type, - is_nullable: pg_row_try_bool(row, 2).unwrap_or(true), - column_default: row.try_get::<_, Option>(3).ok().flatten(), - is_primary_key: pg_row_try_bool(row, 4).unwrap_or(false), - extra: row.try_get::<_, Option>(6).ok().flatten(), - comment: row.try_get::<_, Option>(5).ok().flatten(), - numeric_precision: row.try_get::<_, Option>(7).ok().flatten(), - numeric_scale: row.try_get::<_, Option>(8).ok().flatten(), - character_maximum_length: row.try_get::<_, Option>(9).ok().flatten(), - enum_values: parse_enum_values_from_row(row, 10), - ..Default::default() - } + column_info_from_row_offset(row, 0) } async fn get_columns_with_sql( @@ -5549,7 +6275,10 @@ pub async fn list_foreign_keys(pool: &Pool, schema: &str, table: &str) -> Result } fn postgres_table_dependencies_sql() -> &'static str { - "SELECT DISTINCT child.relname AS table_name, parent.relname AS ref_table \ + // Foreign keys aren't the only ordering constraint on export/replay: a + // partition must be created after its parent table exists too, so union + // in `pg_inherits` partition-of edges alongside the FK edges. + "SELECT child.relname AS table_name, parent.relname AS ref_table \ FROM pg_catalog.pg_constraint con \ JOIN pg_catalog.pg_class child ON child.oid = con.conrelid \ JOIN pg_catalog.pg_namespace child_schema ON child_schema.oid = child.relnamespace \ @@ -5558,7 +6287,16 @@ fn postgres_table_dependencies_sql() -> &'static str { WHERE con.contype = 'f' \ AND child_schema.nspname = $1 \ AND parent_schema.nspname = $1 \ - ORDER BY child.relname, parent.relname" + UNION \ + SELECT child.relname AS table_name, parent.relname AS ref_table \ + FROM pg_catalog.pg_inherits i \ + JOIN pg_catalog.pg_class child ON child.oid = i.inhrelid AND child.relispartition \ + JOIN pg_catalog.pg_namespace child_schema ON child_schema.oid = child.relnamespace \ + JOIN pg_catalog.pg_class parent ON parent.oid = i.inhparent \ + JOIN pg_catalog.pg_namespace parent_schema ON parent_schema.oid = parent.relnamespace \ + WHERE child_schema.nspname = $1 \ + AND parent_schema.nspname = $1 \ + ORDER BY table_name, ref_table" } /// Fetch all same-schema table dependencies in one round trip. Whole-database @@ -5603,16 +6341,37 @@ pub async fn list_triggers(pool: &Pool, schema: &str, table: &str) -> Result Result, String> { let client = checkout_postgres_client(pool, None, super::connection_timeout()).await?; - let rows = postgres_query_cached(&client, postgres_trigger_definitions_sql(), &[&schema, &table]) - .await - .map_err(|e| e.to_string())?; + let sql = if postgres_trigger_has_tgparentid(&client).await? { + postgres_trigger_definitions_sql() + } else { + postgres_trigger_definitions_sql_without_tgparentid() + }; + let rows = postgres_query_cached(&client, sql, &[&schema, &table]).await.map_err(|e| e.to_string())?; Ok(rows.iter().map(|row| pg_row_try_string(row, 0)).filter(|definition| !definition.trim().is_empty()).collect()) } fn postgres_trigger_definitions_sql() -> &'static str { + "SELECT pg_catalog.pg_get_triggerdef(t.oid, true) AS trigger_definition \ + FROM pg_catalog.pg_trigger t \ + JOIN pg_catalog.pg_class c ON c.oid = t.tgrelid \ + JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace \ + WHERE n.nspname = $1 AND c.relname = $2 AND NOT t.tgisinternal AND t.tgparentid = 0 \ + ORDER BY t.tgname, t.oid" +} + +fn postgres_trigger_definitions_sql_without_tgparentid() -> &'static str { "SELECT pg_catalog.pg_get_triggerdef(t.oid, true) AS trigger_definition \ FROM pg_catalog.pg_trigger t \ JOIN pg_catalog.pg_class c ON c.oid = t.tgrelid \ @@ -7454,6 +8213,8 @@ mod tests { assert!(sql.contains("child_schema.nspname = $1")); assert!(sql.contains("parent_schema.nspname = $1")); assert!(!sql.contains("information_schema")); + assert!(sql.contains("pg_catalog.pg_inherits")); + assert!(sql.contains("child.relispartition")); } #[test] @@ -8090,15 +8851,35 @@ mod tests { assert!(relation_sql.contains("pg_catalog.pg_inherits")); assert!(relation_sql.contains("parent.oid = i.inhparent")); assert!(relation_sql.contains(") = 'p'")); - assert!(relation_sql.contains("c.relkind IN ('r','p')")); + // Foreign tables (relkind 'f') must be reachable here too, or a + // foreign partition looked up directly (not via recursion from its + // parent) would never be classified as a partition at all. + assert!(relation_sql.contains("c.relkind IN ('r','p','f')")); assert!(info_sql.contains("pg_catalog.pg_get_expr(c.relpartbound, c.oid, true)")); assert!(info_sql.contains("pg_catalog.pg_get_partkeydef(c.oid)")); assert!(info_sql.contains("pg_catalog.pg_inherits")); assert!(info_sql.contains("parent_schema")); assert!(info_sql.contains("parent_table")); + assert!(info_sql.contains("pg_catalog.pg_foreign_table")); + assert!(info_sql.contains("pg_catalog.pg_foreign_server")); + assert!(info_sql.contains("fs.srvname")); + assert!(info_sql.contains("ft.ftoptions")); + assert!(info_sql.contains("c.relkind IN ('r','p','f')")); assert!(local_objects_sql.contains("row_to_json(con)->>'conparentid'")); assert!(local_objects_sql.contains("con.contype IN ('p','f')")); assert!(local_objects_sql.contains("i.inhrelid = idx.oid")); + assert!(local_objects_sql.contains("con.contype = 'c' AND con.conislocal")); + assert!(!local_objects_sql.contains("con.coninhcount = 0")); + assert!(local_objects_sql.contains("pg_catalog.pg_attrdef")); + } + + #[test] + fn postgres_check_constraints_sql_selects_definitions_for_relation() { + let sql = postgres_check_constraints_sql(); + assert!(sql.contains("con.contype = 'c'")); + assert!(sql.contains("pg_catalog.pg_get_constraintdef(con.oid, true)")); + assert!(sql.contains("n.nspname = $1 AND c.relname = $2")); + assert!(sql.contains("ORDER BY con.conname")); } #[test] @@ -8115,31 +8896,65 @@ mod tests { #[test] fn postgres_column_metadata_marks_only_owned_integer_sequence_defaults_as_serial() { - for sql in [POSTGRES_COLUMNS_SQL, POSTGRES_COLUMNS_COMPAT_SQL] { + let modern_sql = [POSTGRES_COLUMNS_SQL, postgres_columns_for_relations_sql()]; + let compat_sql = [POSTGRES_COLUMNS_COMPAT_SQL, postgres_columns_for_relations_compat_sql()]; + for sql in modern_sql.into_iter().chain(compat_sql) { assert!(sql.contains("a.atttypid IN (20, 21, 23)")); assert!(sql.contains("WHEN 21 THEN 'smallserial'")); assert!(sql.contains("WHEN 23 THEN 'serial'")); assert!(sql.contains("WHEN 20 THEN 'bigserial'")); } - assert!(POSTGRES_COLUMNS_SQL.contains("dep.deptype IN ('a', 'i')")); - assert!(POSTGRES_COLUMNS_SQL.contains("pseq.seqrelid IS NOT NULL")); - assert!(POSTGRES_COLUMNS_SQL.contains("default_dep.objid = ad.oid")); - assert!(POSTGRES_COLUMNS_SQL.contains("default_dep.refobjid = dep.objid")); - assert!(POSTGRES_COLUMNS_SQL.contains( - "pg_get_expr(ad.adbin, ad.adrelid) = format('nextval(%L::regclass)', dep.objid::regclass::text)" - )); - assert!(POSTGRES_COLUMNS_COMPAT_SQL.contains("serial_dep.deptype = 'a'")); - assert!(POSTGRES_COLUMNS_COMPAT_SQL.contains("serial_seq.relkind = 'S'")); - assert!(POSTGRES_COLUMNS_COMPAT_SQL.contains("serial_default_dep.objid = ad.oid")); - assert!(POSTGRES_COLUMNS_COMPAT_SQL.contains("serial_default_dep.refobjid = serial_seq.oid")); - assert!(POSTGRES_COLUMNS_COMPAT_SQL.contains( - "pg_get_expr(ad.adbin, ad.adrelid) = format('nextval(%L::regclass)', serial_seq.oid::regclass::text)" - )); + for sql in modern_sql { + assert!(sql.contains("dep.deptype IN ('a', 'i')")); + assert!(sql.contains("pseq.seqrelid IS NOT NULL")); + assert!(sql.contains("default_dep.objid = ad.oid")); + assert!(sql.contains("default_dep.refobjid = dep.objid")); + assert!(sql.contains( + "pg_get_expr(ad.adbin, ad.adrelid) = format('nextval(%L::regclass)', dep.objid::regclass::text)" + )); + } + for sql in compat_sql { + assert!(sql.contains("serial_dep.deptype = 'a'")); + assert!(sql.contains("serial_seq.relkind = 'S'")); + assert!(sql.contains("serial_default_dep.objid = ad.oid")); + assert!(sql.contains("serial_default_dep.refobjid = serial_seq.oid")); + assert!(sql.contains( + "pg_get_expr(ad.adbin, ad.adrelid) = format('nextval(%L::regclass)', serial_seq.oid::regclass::text)" + )); + } assert!(!POSTGRES_COLUMNS_INFORMATION_SCHEMA_SQL.contains("serial_dep")); assert!(POSTGRES_COLUMNS_INFORMATION_SCHEMA_SQL.contains("NULL::text AS column_extra")); } + #[test] + fn postgres_partition_local_checks_keep_merged_local_definitions() { + for sql in + [postgres_table_partition_local_objects_sql(), postgres_table_partition_local_objects_for_relations_sql()] + { + assert!(sql.contains("con.contype = 'c' AND con.conislocal")); + assert!(!sql.contains("con.coninhcount = 0")); + } + } + + #[test] + fn postgres_partition_batch_metadata_uses_bounded_compat_tiers() { + let column_tiers = postgres_columns_for_relations_query_tiers(); + assert_eq!(column_tiers.len(), 2); + assert!(column_tiers.iter().all(|sql| sql.contains("c.oid = ANY($1::bigint[])"))); + assert!(column_tiers[0].contains("a.attgenerated")); + assert!(!column_tiers[1].contains("a.attgenerated")); + assert!(!column_tiers[1].contains("pg_sequence")); + assert!(column_tiers[1].contains("serial_dep.deptype = 'a'")); + + let index_tiers = postgres_indexes_for_relations_query_tiers(); + assert_eq!(index_tiers.len(), 2); + assert!(index_tiers.iter().all(|sql| sql.contains("t.oid = ANY($1::bigint[])"))); + assert!(index_tiers[0].contains("ix.indnkeyatts")); + assert!(!index_tiers[1].contains("ix.indnkeyatts")); + assert!(index_tiers[1].contains("string_to_array(ix.indkey::text, ' ')")); + } + #[test] fn postgres_column_metadata_has_opengauss_compatible_fallback() { assert!(!POSTGRES_COLUMNS_COMPAT_SQL.contains("a.attidentity")); @@ -8410,12 +9225,63 @@ mod tests { assert!(!ddl.contains("PARTITION OF"), "ddl: {ddl}"); } + #[tokio::test] + #[ignore = "requires DBX_TEST_POSTGRES_URL pointing at a writable PostgreSQL database"] + async fn postgres_ddl_emits_not_valid_check_constraints_as_separate_replayable_alter() { + let url = std::env::var("DBX_TEST_POSTGRES_URL").expect("DBX_TEST_POSTGRES_URL"); + let pool = connect(&url, Duration::from_secs(5)).await.expect("connect postgres"); + let pid = std::process::id(); + let schema = format!("dbx_not_valid_check_{pid}"); + let schema_ident = pg_quote_ident(&schema); + let table = format!("{schema_ident}.events"); + + execute_query(&pool, &format!("CREATE SCHEMA {schema_ident}")).await.expect("create schema"); + let client = pool.get().await.expect("get postgres client"); + client + .batch_execute(&format!( + "CREATE TABLE {table} (id integer, amount integer); \ + ALTER TABLE {table} ADD CONSTRAINT c_pos CHECK (amount > 0) NOT VALID" + )) + .await + .expect("create table with unvalidated check constraint"); + + let ddl = crate::schema::pg_ddl(&pool, &schema, "events").await.expect("events ddl"); + + // A `NOT VALID` CHECK constraint is illegal inline in a `CREATE + // TABLE` column list — Postgres only accepts it after `ALTER TABLE + // ADD CONSTRAINT`. Replaying the generated DDL onto a fresh schema + // exercises that this is real, executable syntax, not just a string + // match. + let replay_schema = format!("{schema}_replay"); + let replay_schema_ident = pg_quote_ident(&replay_schema); + execute_query(&pool, &format!("CREATE SCHEMA {replay_schema_ident}")).await.expect("create replay schema"); + client + .batch_execute(&ddl.replace(&schema, &replay_schema)) + .await + .unwrap_or_else(|error| panic!("replay not-valid-check ddl failed: {error:?}; ddl: {ddl}")); + + client + .batch_execute(&format!("DROP SCHEMA {schema_ident} CASCADE; DROP SCHEMA {replay_schema_ident} CASCADE")) + .await + .expect("drop schemas"); + + assert!(!ddl.contains("NOT VALID)"), "NOT VALID must not appear inline in the column list: {ddl}"); + assert!(ddl.contains("CREATE TABLE"), "ddl: {ddl}"); + assert!(!ddl.contains("\"amount\" integer,"), "check constraint must not be inlined into columns: {ddl}"); + assert!(ddl.contains("ALTER TABLE"), "ddl: {ddl}"); + assert!( + ddl.contains("ADD CONSTRAINT \"c_pos\" CHECK (amount > 0) NOT VALID;"), + "expected a separate ALTER TABLE ADD CONSTRAINT ... NOT VALID statement: {ddl}" + ); + } + #[tokio::test] #[ignore = "requires DBX_TEST_POSTGRES_URL pointing at a writable PostgreSQL database"] async fn postgres_partition_metadata_renders_replayable_children_and_subpartitions() { let url = std::env::var("DBX_TEST_POSTGRES_URL").expect("DBX_TEST_POSTGRES_URL"); let pool = connect(&url, Duration::from_secs(5)).await.expect("connect postgres"); - let schema = format!("dbx_partition_meta_{}", std::process::id()); + let pid = std::process::id(); + let schema = format!("dbx_partition_meta_{pid}"); let schema_ident = pg_quote_ident(&schema); let replay_schema = format!("{schema}_replay"); let replay_schema_ident = pg_quote_ident(&replay_schema); @@ -8425,16 +9291,28 @@ mod tests { let subpartition = format!("{schema_ident}.subpartition"); let inherited_parent = format!("{schema_ident}.inherited_parent"); let inherited_child = format!("{schema_ident}.inherited_child"); + // Declared in "public" (not the per-test schema) so the same + // qualified name keeps resolving after `schema_ident` substrings are + // rewritten to `replay_schema_ident` for the replay pass below. + let trigger_fn = format!("public.dbx_partition_meta_trigger_fn_{pid}"); execute_query(&pool, &format!("CREATE SCHEMA {schema_ident}")).await.expect("create schema"); let client = pool.get().await.expect("get postgres client"); client .batch_execute(&format!( - "CREATE TABLE {parent} (id integer, payload text) PARTITION BY RANGE (id); \ - CREATE TABLE {child} PARTITION OF {parent} (PRIMARY KEY (id)) FOR VALUES FROM (1) TO (10); \ + "CREATE OR REPLACE FUNCTION {trigger_fn}() RETURNS trigger LANGUAGE plpgsql AS \ + $$ BEGIN RETURN NEW; END $$; \ + CREATE TABLE {parent} (id integer, payload text, status text DEFAULT 'pending') PARTITION BY RANGE (id); \ + CREATE TABLE {child} PARTITION OF {parent} ( \ + PRIMARY KEY (id), \ + CONSTRAINT child_payload_not_blank CHECK (payload <> ''), \ + status WITH OPTIONS DEFAULT 'child-status' \ + ) FOR VALUES FROM (1) TO (10); \ CREATE INDEX child_payload_idx ON {child} (payload); \ CREATE TABLE {default_child} PARTITION OF {parent} DEFAULT; \ CREATE TABLE {subpartition} PARTITION OF {parent} FOR VALUES FROM (10) TO (20) PARTITION BY HASH (payload); \ + ALTER TABLE {parent} ADD CONSTRAINT parent_id_positive CHECK (id > 0); \ + CREATE TRIGGER parent_biu BEFORE INSERT ON {parent} FOR EACH ROW EXECUTE FUNCTION {trigger_fn}(); \ CREATE TABLE {inherited_parent} (id integer, bucket integer, PRIMARY KEY (id, bucket)) PARTITION BY RANGE (bucket); \ CREATE TABLE {inherited_child} PARTITION OF {inherited_parent} FOR VALUES FROM (1) TO (10)" )) @@ -8451,29 +9329,66 @@ mod tests { let inherited_child_local_objects = get_table_partition_local_objects(&pool, &schema, "inherited_child") .await .expect("inherited child local objects"); - let parent_ddl = crate::schema::pg_ddl(&pool, &schema, "parent").await.expect("parent ddl"); + let parent_ddl = crate::schema::pg_ddl_with_partitions(&pool, &schema, "parent").await.expect("parent ddl"); let child_ddl = crate::schema::pg_ddl(&pool, &schema, "child").await.expect("child ddl"); let default_ddl = crate::schema::pg_ddl(&pool, &schema, "child_default").await.expect("default ddl"); let subpartition_ddl = crate::schema::pg_ddl(&pool, &schema, "subpartition").await.expect("subpartition ddl"); - let inherited_parent_ddl = - crate::schema::pg_ddl(&pool, &schema, "inherited_parent").await.expect("inherited parent ddl"); + let inherited_parent_ddl = crate::schema::pg_ddl_with_partitions(&pool, &schema, "inherited_parent") + .await + .expect("inherited parent ddl"); let inherited_child_ddl = crate::schema::pg_ddl(&pool, &schema, "inherited_child").await.expect("inherited child ddl"); + // The bare (non-recursive) parent DDL, exactly what database export + // and table transfer produce per-relation — used below to give the + // single-relation partition ddls (child_ddl, default_ddl, + // subpartition_ddl, inherited_child_ddl) a real parent shell to + // replay against, so their syntax is actually exercised by + // PostgreSQL rather than only checked with string assertions. + let bare_parent_ddl = crate::schema::pg_ddl(&pool, &schema, "parent").await.expect("bare parent ddl"); + let bare_inherited_parent_ddl = + crate::schema::pg_ddl(&pool, &schema, "inherited_parent").await.expect("bare inherited parent ddl"); + // `parent_ddl`/`inherited_parent_ddl` already embed their full + // partition subtree (that's the fix under test), so replaying the + // per-partition ddl strings on top would recreate the same relations + // twice. Only the two top-level requests need replaying; the + // individually-fetched partition ddls above are replayed separately + // below, on top of a bare (childless) copy of their parent. execute_query(&pool, &format!("CREATE SCHEMA {replay_schema_ident}")).await.expect("create replay schema"); - for ddl in - [&parent_ddl, &child_ddl, &default_ddl, &subpartition_ddl, &inherited_parent_ddl, &inherited_child_ddl] - { + for ddl in [&parent_ddl, &inherited_parent_ddl] { client - .batch_execute(&ddl.replace(&schema_ident, &replay_schema_ident)) + .batch_execute(&ddl.replace(&schema, &replay_schema)) .await - .unwrap_or_else(|error| panic!("replay partition ddl failed: {error}; ddl: {ddl}")); + .unwrap_or_else(|error| panic!("replay partition ddl failed: {error:?}; ddl: {ddl}")); + } + + // Single-relation partition ddls are now a real production code + // path (database export/table transfer fetch each partition this + // way), so replay them for real too: a fresh schema gets just the + // bare parents, then each single-relation partition ddl on top. + let bare_schema = format!("{schema}_bare"); + let bare_schema_ident = pg_quote_ident(&bare_schema); + execute_query(&pool, &format!("CREATE SCHEMA {bare_schema_ident}")).await.expect("create bare schema"); + for ddl in [&bare_parent_ddl, &bare_inherited_parent_ddl] { + client + .batch_execute(&ddl.replace(&schema, &bare_schema)) + .await + .unwrap_or_else(|error| panic!("replay bare parent ddl failed: {error:?}; ddl: {ddl}")); + } + for ddl in [&child_ddl, &default_ddl, &subpartition_ddl, &inherited_child_ddl] { + client + .batch_execute(&ddl.replace(&schema, &bare_schema)) + .await + .unwrap_or_else(|error| panic!("replay single-relation partition ddl failed: {error:?}; ddl: {ddl}")); } client - .batch_execute(&format!("DROP SCHEMA {schema_ident} CASCADE; DROP SCHEMA {replay_schema_ident} CASCADE")) + .batch_execute(&format!( + "DROP SCHEMA {schema_ident} CASCADE; DROP SCHEMA {replay_schema_ident} CASCADE; \ + DROP SCHEMA {bare_schema_ident} CASCADE; DROP FUNCTION IF EXISTS {trigger_fn}()" + )) .await - .expect("drop schemas"); + .expect("drop schemas and trigger function"); assert_eq!(parent_info.key.as_deref(), Some("RANGE (id)")); assert!(!parent_info.is_partition); @@ -8482,16 +9397,31 @@ mod tests { assert_eq!(child_info.bound.as_deref(), Some("FOR VALUES FROM (1) TO (10)")); assert!(child_local_objects.has_primary_key); assert!(child_local_objects.indexes.contains("child_payload_idx")); + assert!(child_local_objects.check_constraints.contains("child_payload_not_blank")); + assert!(!child_local_objects.check_constraints.contains("parent_id_positive")); + assert_eq!(child_local_objects.column_defaults.get("status"), Some(&PostgresColumnDefaultState::Overridden)); assert!(!inherited_child_local_objects.has_primary_key); assert!(inherited_child_local_objects.indexes.is_empty()); assert_eq!(default_info.bound.as_deref(), Some("DEFAULT")); assert_eq!(subpartition_info.key.as_deref(), Some("HASH (payload)")); assert!(parent_ddl.contains(") PARTITION BY RANGE (id);"), "ddl: {parent_ddl}"); + assert!(parent_ddl.contains("CONSTRAINT \"parent_id_positive\" CHECK (id > 0)"), "ddl: {parent_ddl}"); + // The parent's trigger is automatically cloned onto every partition + // by PostgreSQL itself when the partition is created/attached, so it + // must appear exactly once in the combined DDL, not once per relation. + assert_eq!(parent_ddl.matches("CREATE TRIGGER").count(), 1, "ddl: {parent_ddl}"); assert!(child_ddl.contains("CREATE TABLE"), "ddl: {child_ddl}"); assert!(child_ddl.contains("PARTITION OF"), "ddl: {child_ddl}"); assert!(child_ddl.contains("PRIMARY KEY (\"id\")"), "ddl: {child_ddl}"); assert!(child_ddl.contains("FOR VALUES FROM (1) TO (10);"), "ddl: {child_ddl}"); assert!(child_ddl.contains("CREATE INDEX \"child_payload_idx\""), "ddl: {child_ddl}"); + assert!(child_ddl.contains("CONSTRAINT \"child_payload_not_blank\" CHECK (payload"), "ddl: {child_ddl}"); + assert!(!child_ddl.contains("parent_id_positive"), "ddl: {child_ddl}"); + assert!(child_ddl.contains("\"status\" WITH OPTIONS DEFAULT 'child-status'"), "ddl: {child_ddl}"); + assert!( + !child_ddl.contains("CREATE TRIGGER"), + "child must not redeclare the parent's cloned trigger: {child_ddl}" + ); assert!(default_ddl.contains(" DEFAULT;"), "ddl: {default_ddl}"); assert!( subpartition_ddl.contains("FOR VALUES FROM (10) TO (20) PARTITION BY HASH (payload);"), @@ -8501,6 +9431,77 @@ mod tests { assert!(!inherited_child_ddl.contains("CREATE INDEX"), "ddl: {inherited_child_ddl}"); } + #[tokio::test] + #[ignore = "requires DBX_TEST_POSTGRES_URL pointing at a writable PostgreSQL database with postgres_fdw available"] + async fn postgres_partition_ddl_uses_foreign_table_syntax_for_foreign_partitions() { + let url = std::env::var("DBX_TEST_POSTGRES_URL").expect("DBX_TEST_POSTGRES_URL"); + let pool = connect(&url, Duration::from_secs(5)).await.expect("connect postgres"); + let pid = std::process::id(); + let schema = format!("dbx_partition_fdw_{pid}"); + let schema_ident = pg_quote_ident(&schema); + let server = format!("dbx_partition_fdw_server_{pid}"); + let server_ident = pg_quote_ident(&server); + let replay_schema = format!("{schema}_replay"); + let replay_schema_ident = pg_quote_ident(&replay_schema); + let parent = format!("{schema_ident}.parent"); + let local_child = format!("{schema_ident}.local_child"); + let foreign_child = format!("{schema_ident}.foreign_child"); + + execute_query(&pool, &format!("CREATE SCHEMA {schema_ident}")).await.expect("create schema"); + let client = pool.get().await.expect("get postgres client"); + client + .batch_execute(&format!( + "CREATE EXTENSION IF NOT EXISTS postgres_fdw; \ + CREATE SERVER {server_ident} FOREIGN DATA WRAPPER postgres_fdw \ + OPTIONS (host 'localhost', dbname 'postgres', port '5432'); \ + CREATE USER MAPPING FOR CURRENT_USER SERVER {server_ident} \ + OPTIONS (user 'postgres', password 'postgres'); \ + CREATE TABLE {parent} (id integer, payload text) PARTITION BY RANGE (id); \ + CREATE TABLE {local_child} PARTITION OF {parent} FOR VALUES FROM (1) TO (10); \ + CREATE FOREIGN TABLE {foreign_child} PARTITION OF {parent} \ + FOR VALUES FROM (10) TO (20) \ + SERVER {server_ident} OPTIONS (schema_name 'public', table_name 'foreign_child_data')" + )) + .await + .expect("create partitioned tables with a foreign partition"); + + let foreign_info = get_table_partition_info(&pool, &schema, "foreign_child").await.expect("foreign child info"); + let local_info = get_table_partition_info(&pool, &schema, "local_child").await.expect("local child info"); + let parent_ddl = crate::schema::pg_ddl_with_partitions(&pool, &schema, "parent").await.expect("parent ddl"); + let foreign_ddl = crate::schema::pg_ddl(&pool, &schema, "foreign_child").await.expect("foreign child ddl"); + + execute_query(&pool, &format!("CREATE SCHEMA {replay_schema_ident}")).await.expect("create replay schema"); + client + .batch_execute(&parent_ddl.replace(&schema, &replay_schema)) + .await + .unwrap_or_else(|error| panic!("replay foreign partition ddl failed: {error:?}; ddl: {parent_ddl}")); + + client + .batch_execute(&format!( + "DROP SCHEMA {schema_ident} CASCADE; DROP SCHEMA {replay_schema_ident} CASCADE; \ + DROP SERVER {server_ident} CASCADE" + )) + .await + .expect("drop schemas and foreign server"); + + assert!(!local_info.is_foreign); + assert!(foreign_info.is_foreign); + assert_eq!(foreign_info.foreign_server.as_deref(), Some(server.as_str())); + assert!( + foreign_info.foreign_options.contains(&("schema_name".to_string(), "public".to_string())), + "options: {:?}", + foreign_info.foreign_options + ); + assert!(parent_ddl.contains("CREATE FOREIGN TABLE"), "ddl: {parent_ddl}"); + assert!(foreign_ddl.starts_with("CREATE FOREIGN TABLE"), "ddl: {foreign_ddl}"); + assert!(foreign_ddl.contains("PARTITION OF"), "ddl: {foreign_ddl}"); + assert!(foreign_ddl.contains(&format!("SERVER {}", pg_quote_ident(&server))), "ddl: {foreign_ddl}"); + assert!(foreign_ddl.contains("\"schema_name\" 'public'"), "ddl: {foreign_ddl}"); + assert!(foreign_ddl.contains("\"table_name\" 'foreign_child_data'"), "ddl: {foreign_ddl}"); + assert!(!foreign_ddl.contains("PRIMARY KEY"), "ddl: {foreign_ddl}"); + assert!(!foreign_ddl.contains("CREATE INDEX"), "ddl: {foreign_ddl}"); + } + #[tokio::test] #[ignore = "requires DBX_TEST_POSTGRES_URL pointing at a writable PostgreSQL database"] async fn postgres_schema_context_prioritizes_selected_schema_and_cleans_up() { @@ -9292,9 +10293,24 @@ mod tests { let sql = postgres_trigger_definitions_sql(); assert!(sql.contains("pg_catalog.pg_get_triggerdef(t.oid, true) AS trigger_definition")); assert!(sql.contains("NOT t.tgisinternal")); + assert!(sql.contains("t.tgparentid = 0")); assert!(sql.contains("ORDER BY t.tgname, t.oid")); } + #[test] + fn postgres_trigger_definitions_sql_without_tgparentid_omits_partition_clone_filter() { + let sql = postgres_trigger_definitions_sql_without_tgparentid(); + assert!(sql.contains("NOT t.tgisinternal")); + assert!(!sql.contains("tgparentid")); + } + + #[test] + fn postgres_trigger_has_tgparentid_sql_checks_catalog_attribute() { + let sql = postgres_trigger_has_tgparentid_sql(); + assert!(sql.contains("pg_catalog.pg_trigger")); + assert!(sql.contains("attname = 'tgparentid'")); + } + #[test] fn transaction_recovery_statement_detection_matches_common_postgres_commands() { assert!(is_transaction_recovery_statement("ROLLBACK")); diff --git a/crates/dbx-core/src/schema.rs b/crates/dbx-core/src/schema.rs index 06485bae78..2cb3f5b5c5 100644 --- a/crates/dbx-core/src/schema.rs +++ b/crates/dbx-core/src/schema.rs @@ -7,7 +7,7 @@ use crate::db; use crate::models::connection::{ConnectionConfig, DatabaseType}; use crate::query::{agent_execute_query_params, should_discard_pool_after_error, QueryExecutionOptions}; use serde::{Deserialize, Serialize}; -use std::collections::{BTreeMap, BTreeSet, HashMap}; +use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet}; use std::future::Future; use std::sync::Arc; use std::time::{Duration, Instant}; @@ -6495,6 +6495,26 @@ pub async fn list_owners_core( .await } +/// Whether to widen or normalize a single-table DDL fetch for its caller. +/// +/// Database export and table transfer render one relation at a time because +/// they already iterate every relation themselves. Interactive display paths +/// include PostgreSQL access statements and recurse through the partition +/// tree. Oracle export additionally requests portable DDL normalization. +#[derive(Clone, Copy)] +struct TableDdlOptions { + include_postgres_access: bool, + include_partitions: bool, + portable_oracle: bool, +} + +impl TableDdlOptions { + const SINGLE_RELATION: Self = + Self { include_postgres_access: false, include_partitions: false, portable_oracle: false }; + const EXPORT: Self = Self { include_postgres_access: false, include_partitions: false, portable_oracle: true }; + const DISPLAY: Self = Self { include_postgres_access: true, include_partitions: true, portable_oracle: false }; +} + pub async fn get_table_ddl_core( state: &AppState, connection_id: &str, @@ -6503,7 +6523,16 @@ pub async fn get_table_ddl_core( table: &str, object_type: Option, ) -> Result { - get_table_ddl_core_with_options(state, connection_id, database, schema, table, object_type, false, false).await + get_table_ddl_core_with_options( + state, + connection_id, + database, + schema, + table, + object_type, + TableDdlOptions::SINGLE_RELATION, + ) + .await } pub async fn get_table_export_ddl_core( @@ -6514,7 +6543,8 @@ pub async fn get_table_export_ddl_core( table: &str, object_type: Option, ) -> Result { - get_table_ddl_core_with_options(state, connection_id, database, schema, table, object_type, false, true).await + get_table_ddl_core_with_options(state, connection_id, database, schema, table, object_type, TableDdlOptions::EXPORT) + .await } pub async fn get_table_display_ddl_core( @@ -6525,7 +6555,16 @@ pub async fn get_table_display_ddl_core( table: &str, object_type: Option, ) -> Result { - get_table_ddl_core_with_options(state, connection_id, database, schema, table, object_type, true, false).await + get_table_ddl_core_with_options( + state, + connection_id, + database, + schema, + table, + object_type, + TableDdlOptions::DISPLAY, + ) + .await } async fn get_table_ddl_core_with_options( @@ -6535,8 +6574,7 @@ async fn get_table_ddl_core_with_options( schema: &str, table: &str, object_type: Option, - include_postgres_access: bool, - portable_oracle: bool, + options: TableDdlOptions, ) -> Result { if crate::sql_dialect::parse_sqlserver_linked_schema_ref(schema).is_some() { return Err("DDL is not supported for SQL Server linked server tables".to_string()); @@ -6582,19 +6620,33 @@ async fn get_table_ddl_core_with_options( } retry_metadata_connection(state, connection_id, Some(database), || { - get_table_ddl_once(state, connection_id, database, schema, table, include_postgres_access, portable_oracle) + get_table_ddl_once(state, connection_id, database, schema, table, options) }) .await } +/// `pg_ddl_with_partitions` when the caller wants the whole partition tree, +/// otherwise plain single-relation `pg_ddl`. +async fn pg_ddl_for_options( + pool: &deadpool_postgres::Pool, + schema: &str, + table: &str, + include_partitions: bool, +) -> Result { + if include_partitions { + pg_ddl_with_partitions(pool, schema, table).await + } else { + pg_ddl(pool, schema, table).await + } +} + async fn get_table_ddl_once( state: &AppState, connection_id: &str, database: &str, schema: &str, table: &str, - include_postgres_access: bool, - portable_oracle: bool, + options: TableDdlOptions, ) -> Result { let pool_key = state.get_or_create_metadata_pool_for_session(connection_id, Some(database), None).await?; let db_config = connection_config(state, connection_id).await; @@ -6642,10 +6694,11 @@ async fn get_table_ddl_once( if let Some(config) = db_config.as_ref().filter(|config| is_agent_postgres_metadata_fallback_config(config)) { match native_postgres_metadata_pool(state, connection_id, database, config).await { - Ok(Some(pool)) => match pg_ddl(&pool, schema, table).await { - Ok(ddl) => return Ok(ddl), - Err(error) => { - log::warn!( + Ok(Some(pool)) => { + match pg_ddl_for_options(&pool, schema, table, options.include_partitions).await { + Ok(ddl) => return Ok(ddl), + Err(error) => { + log::warn!( "[schema][agent:get_table_ddl:postgres-compatible-native-fallback-failed] connection_id={} database={} schema={} table={} error={}", connection_id, database, @@ -6653,8 +6706,9 @@ async fn get_table_ddl_once( table, error ); + } } - }, + } Ok(None) => {} Err(error) => { log::warn!( @@ -6674,7 +6728,7 @@ async fn get_table_ddl_once( database, schema, table, - portable_oracle, + options.portable_oracle, agent_metadata_timeout(db_config.as_ref()), ) .await; @@ -6702,27 +6756,27 @@ async fn get_table_ddl_once( PoolKind::Postgres(p) if db_config.as_ref().is_some_and(is_opengauss_family_config) => { match opengauss_table_ddl(p, schema, table).await { Ok(ddl) => Ok(ddl), - Err(_) => pg_ddl(p, schema, table).await, + Err(_) => pg_ddl_for_options(p, schema, table, options.include_partitions).await, } } PoolKind::Postgres(p) if db_config.as_ref().is_some_and(is_questdb_config) => { match db::questdb::questdb_table_or_view_ddl(p, table).await { Ok(ddl) => Ok(ddl), - Err(_) => pg_ddl(p, schema, table).await, + Err(_) => pg_ddl_for_options(p, schema, table, options.include_partitions).await, } } PoolKind::Postgres(p) if db_config.as_ref().is_some_and(is_cloudberry_config) => { - cloudberry_ddl(p, schema, table).await + cloudberry_ddl(p, schema, table, options.include_partitions).await } PoolKind::Postgres(p) if db_config.as_ref().is_some_and(db::opentenbase::is_config) => { - opentenbase_ddl(p, schema, table).await + opentenbase_ddl(p, schema, table, options.include_partitions).await } PoolKind::Postgres(p) - if include_postgres_access && db_config.as_ref().is_some_and(is_native_postgres_config) => + if options.include_postgres_access && db_config.as_ref().is_some_and(is_native_postgres_config) => { pg_display_ddl(p, schema, table).await } - PoolKind::Postgres(p) => pg_ddl(p, schema, table).await, + PoolKind::Postgres(p) => pg_ddl_for_options(p, schema, table, options.include_partitions).await, PoolKind::Sqlite(p) => sqlite_ddl(p, schema, table).await, PoolKind::Rqlite(client) => db::rqlite_driver::table_ddl(client, table).await, PoolKind::Turso(client) => db::turso_driver::table_ddl(client, table).await, @@ -6844,6 +6898,14 @@ fn pg_ident(value: &str) -> String { format!("\"{}\"", value.replace('"', "\"\"")) } +/// Whether a `pg_get_constraintdef` result ends in the ` NOT VALID` suffix +/// Postgres appends for an unvalidated constraint. That suffix is only legal +/// after `ALTER TABLE ADD CONSTRAINT`, never inside a `CREATE TABLE` column +/// list. +fn is_not_valid_constraintdef(definition: &str) -> bool { + definition.to_ascii_uppercase().trim_end().ends_with("NOT VALID") +} + fn sqlserver_ident(value: &str) -> String { format!("[{}]", value.replace(']', "]]")) } @@ -8644,6 +8706,7 @@ mod ddl_tests { &[column("created_at", "timestamp without time zone")], &[], &[], + &[], None, &db::postgres::PostgresTablePartitionInfo { key: Some(partition_key.to_string()), @@ -8665,6 +8728,7 @@ mod ddl_tests { &[column("id", "integer")], &[], &[], + &[], None, &db::postgres::PostgresTablePartitionInfo::default(), &db::postgres::PostgresTablePartitionLocalObjects::default(), @@ -8695,11 +8759,13 @@ mod ddl_tests { parent_table: Some("events".to_string()), bound: Some("FOR VALUES FROM ('2026-01-01') TO ('2027-01-01')".to_string()), key: Some("HASH (payload)".to_string()), + ..Default::default() }; let partition_local_objects = db::postgres::PostgresTablePartitionLocalObjects { has_primary_key: true, foreign_keys: BTreeSet::new(), indexes: BTreeSet::from(["events_payload_idx".to_string()]), + ..Default::default() }; let ddl = render_postgres_table_ddl_with_partition_info( @@ -8708,6 +8774,7 @@ mod ddl_tests { &[id, column("payload", "text")], &indexes, &[], + &[], None, &partition_info, &partition_local_objects, @@ -8742,6 +8809,7 @@ mod ddl_tests { parent_table: Some("events".to_string()), bound: Some("DEFAULT".to_string()), key: None, + ..Default::default() }; let ddl = render_postgres_table_ddl_with_partition_info( @@ -8750,6 +8818,7 @@ mod ddl_tests { &[id], &indexes, &[], + &[], None, &partition_info, &db::postgres::PostgresTablePartitionLocalObjects::default(), @@ -8758,6 +8827,175 @@ mod ddl_tests { assert_eq!(ddl, "CREATE TABLE \"public\".\"events_default\" PARTITION OF \"public\".\"events\" DEFAULT;\n"); } + #[test] + fn postgres_table_ddl_renders_check_constraints_for_ordinary_tables() { + let ddl = render_postgres_table_ddl_with_partition_info( + "public", + "users", + &[column("age", "integer")], + &[], + &[], + &[("users_age_check".to_string(), "CHECK (age >= 0)".to_string())], + None, + &db::postgres::PostgresTablePartitionInfo::default(), + &db::postgres::PostgresTablePartitionLocalObjects::default(), + ); + + assert!(ddl.contains("CONSTRAINT \"users_age_check\" CHECK (age >= 0)"), "ddl: {ddl}"); + } + + #[test] + fn postgres_partition_ddl_only_renders_local_check_constraints() { + let mut partition_local_objects = db::postgres::PostgresTablePartitionLocalObjects::default(); + partition_local_objects.check_constraints.insert("child_only_check".to_string()); + + let ddl = render_postgres_table_ddl_with_partition_info( + "public", + "events_2026", + &[column("payload", "text")], + &[], + &[], + &[ + ("parent_check".to_string(), "CHECK (payload IS NOT NULL)".to_string()), + ("child_only_check".to_string(), "CHECK (payload <> '')".to_string()), + ], + None, + &db::postgres::PostgresTablePartitionInfo { + is_partition: true, + parent_schema: Some("public".to_string()), + parent_table: Some("events".to_string()), + bound: Some("DEFAULT".to_string()), + ..Default::default() + }, + &partition_local_objects, + ); + + assert!(ddl.contains("CONSTRAINT \"child_only_check\" CHECK (payload <> '')"), "ddl: {ddl}"); + assert!(!ddl.contains("parent_check"), "ddl: {ddl}"); + } + + #[test] + fn postgres_partition_ddl_overrides_local_column_default() { + let mut status = column("status", "text"); + status.column_default = Some("'archived'::text".to_string()); + let mut partition_local_objects = db::postgres::PostgresTablePartitionLocalObjects::default(); + partition_local_objects + .column_defaults + .insert("status".to_string(), db::postgres::PostgresColumnDefaultState::Overridden); + + let ddl = render_postgres_table_ddl_with_partition_info( + "public", + "events_2026", + &[status], + &[], + &[], + &[], + None, + &db::postgres::PostgresTablePartitionInfo { + is_partition: true, + parent_schema: Some("public".to_string()), + parent_table: Some("events".to_string()), + bound: Some("DEFAULT".to_string()), + ..Default::default() + }, + &partition_local_objects, + ); + + assert!(ddl.contains("\"status\" WITH OPTIONS DEFAULT 'archived'::text"), "ddl: {ddl}"); + // The partition's own column list is otherwise omitted (inherited + // from the parent), so a plain (non-override) column declaration + // must not appear alongside the WITH OPTIONS clause. + assert!(!ddl.contains("\"status\" text"), "ddl: {ddl}"); + } + + #[test] + fn postgres_partition_ddl_emits_drop_default_for_locally_dropped_column() { + // A partition that ran `ALTER TABLE ONLY child ALTER COLUMN status + // DROP DEFAULT` has no default of its own to report here — the + // column's `column_default` is `None`, distinct from the + // "overridden" case (which has its own, different, Some value). + let status = column("status", "text"); + let mut partition_local_objects = db::postgres::PostgresTablePartitionLocalObjects::default(); + partition_local_objects + .column_defaults + .insert("status".to_string(), db::postgres::PostgresColumnDefaultState::Dropped); + + let ddl = render_postgres_table_ddl_with_partition_info( + "public", + "events_2026", + &[status], + &[], + &[], + &[], + None, + &db::postgres::PostgresTablePartitionInfo { + is_partition: true, + parent_schema: Some("public".to_string()), + parent_table: Some("events".to_string()), + bound: Some("DEFAULT".to_string()), + ..Default::default() + }, + &partition_local_objects, + ); + + assert!( + ddl.contains("ALTER TABLE ONLY \"public\".\"events_2026\" ALTER COLUMN \"status\" DROP DEFAULT;"), + "ddl: {ddl}" + ); + assert!(!ddl.contains("WITH OPTIONS DEFAULT"), "ddl: {ddl}"); + } + + #[test] + fn postgres_table_ddl_renders_foreign_table_with_server_and_options() { + let ddl = render_postgres_table_ddl_with_partition_info( + "public", + "remote_users", + &[column("id", "integer")], + &[], + &[], + &[], + None, + &db::postgres::PostgresTablePartitionInfo { + is_foreign: true, + foreign_server: Some("loopback".to_string()), + foreign_options: vec![("schema_name".to_string(), "public".to_string())], + ..Default::default() + }, + &db::postgres::PostgresTablePartitionLocalObjects::default(), + ); + + assert!(ddl.starts_with("CREATE FOREIGN TABLE \"public\".\"remote_users\""), "ddl: {ddl}"); + assert!(ddl.contains("SERVER \"loopback\""), "ddl: {ddl}"); + assert!(ddl.contains("OPTIONS (\"schema_name\" 'public')"), "ddl: {ddl}"); + } + + #[test] + fn postgres_partition_ddl_uses_foreign_table_syntax_for_foreign_partitions() { + let ddl = render_postgres_table_ddl_with_partition_info( + "public", + "events_remote", + &[column("id", "integer")], + &[], + &[], + &[], + None, + &db::postgres::PostgresTablePartitionInfo { + is_partition: true, + parent_schema: Some("public".to_string()), + parent_table: Some("events".to_string()), + bound: Some("FOR VALUES FROM ('2027-01-01') TO ('2028-01-01')".to_string()), + is_foreign: true, + foreign_server: Some("loopback".to_string()), + ..Default::default() + }, + &db::postgres::PostgresTablePartitionLocalObjects::default(), + ); + + assert!(ddl.starts_with("CREATE FOREIGN TABLE \"public\".\"events_remote\" PARTITION OF"), "ddl: {ddl}"); + assert!(ddl.contains("SERVER \"loopback\""), "ddl: {ddl}"); + assert!(!ddl.contains("CREATE TABLE \"public\".\"events_remote\""), "ddl: {ddl}"); + } + #[test] fn postgres_table_ddl_keeps_composite_foreign_key_together() { let columns = vec![column("a", "integer"), column("b", "integer"), column("c", "integer")]; @@ -9290,14 +9528,20 @@ pub async fn sqlite_ddl(pool: &db::sqlite::SqliteHandle, schema: &str, table: &s .map_err(|e| e.to_string())? } +/// DDL for a single relation. Callers that already iterate a relation set +/// themselves (database export, table transfer) must use this rather than +/// `pg_ddl_with_partitions` — recursing into partition children here would +/// duplicate every partition's `CREATE TABLE` (once from the parent's DDL, +/// once from the caller's own loop over that same child relation). pub async fn pg_ddl(pool: &deadpool_postgres::Pool, schema: &str, table: &str) -> Result { - let (columns, indexes, fkeys, table_comment, partition_info, trigger_definitions) = tokio::try_join!( + let (columns, indexes, fkeys, table_comment, partition_info, trigger_definitions, check_constraints) = tokio::try_join!( db::postgres::get_columns(pool, schema, table), db::postgres::list_indexes(pool, schema, table), db::postgres::list_foreign_keys(pool, schema, table), async { db::postgres::get_table_comment(pool, schema, table).await }, db::postgres::get_table_partition_info(pool, schema, table), db::postgres::list_trigger_definitions(pool, schema, table), + db::postgres::list_check_constraints(pool, schema, table), )?; let partition_local_objects = if partition_info.is_partition { db::postgres::get_table_partition_local_objects(pool, schema, table).await? @@ -9312,6 +9556,7 @@ pub async fn pg_ddl(pool: &deadpool_postgres::Pool, schema: &str, table: &str) - &columns, &indexes, &fkeys, + &check_constraints, table_comment.as_deref(), &partition_info, &partition_local_objects, @@ -9320,8 +9565,156 @@ pub async fn pg_ddl(pool: &deadpool_postgres::Pool, schema: &str, table: &str) - )) } +/// Like `pg_ddl`, but for a partitioned table also emits `CREATE TABLE ... +/// PARTITION OF` for every existing partition, at any depth. Used only by the +/// interactive "view DDL" paths (`get_table_display_ddl_core`) — callers that +/// iterate relations themselves must use `pg_ddl` instead (see its doc +/// comment). Fetches the whole tree's metadata via a handful of batched, +/// tree-wide queries (see `db::postgres::fetch_postgres_partition_tree` and +/// its `_for_relations` siblings) instead of recursing per relation, which +/// would rerun the full metadata query chain once per node and scale request +/// count linearly with the number of partitions. +pub async fn pg_ddl_with_partitions( + pool: &deadpool_postgres::Pool, + schema: &str, + table: &str, +) -> Result { + let tree = db::postgres::fetch_postgres_partition_tree(pool, schema, table).await?; + let tree_oids: HashSet = tree.iter().map(|node| node.oid).collect(); + // The requested relation is the tree root: it's the only node whose parent + // (if it has one at all — it may itself be a partition of a table outside + // this tree) isn't also a node we fetched. + let Some(root) = + tree.iter().find(|node| !node.parent_oid.is_some_and(|parent_oid| tree_oids.contains(&parent_oid))) + else { + return Err(format!( + "relation \"{schema}\".\"{table}\" was not found or is not a table/partition/foreign table" + )); + }; + + let oids: Vec = tree.iter().map(|node| node.oid).collect(); + let relations: Vec<(i64, String, String)> = + tree.iter().map(|node| (node.oid, node.schema.clone(), node.table.clone())).collect(); + let relation_pairs: Vec<(String, String)> = + tree.iter().map(|node| (node.schema.clone(), node.table.clone())).collect(); + + let ( + columns_by_oid, + indexes_by_oid, + fkeys_by_relation, + comments_by_oid, + triggers_by_oid, + checks_by_oid, + local_objects_by_oid, + ) = tokio::try_join!( + db::postgres::get_columns_for_relations(pool, &relations), + db::postgres::list_indexes_for_relations(pool, &relations), + db::postgres::list_foreign_keys_for_relations(pool, &relation_pairs), + db::postgres::get_table_comments_for_relations(pool, &oids), + db::postgres::list_trigger_definitions_for_relations(pool, &oids), + db::postgres::list_check_constraints_for_relations(pool, &oids), + db::postgres::get_table_partition_local_objects_for_relations(pool, &oids), + )?; + + // Group children by parent oid, each group ordered by relname to match + // `list_table_partitions`' `ORDER BY c.relname` (today's traversal order). + let mut children_by_parent: HashMap> = HashMap::new(); + for node in &tree { + if let Some(parent_oid) = node.parent_oid { + children_by_parent.entry(parent_oid).or_default().push(node); + } + } + for children in children_by_parent.values_mut() { + children.sort_by(|a, b| a.table.cmp(&b.table)); + } + + let mut ddl = String::new(); + render_postgres_partition_tree_node( + root, + &children_by_parent, + &columns_by_oid, + &indexes_by_oid, + &fkeys_by_relation, + &comments_by_oid, + &triggers_by_oid, + &checks_by_oid, + &local_objects_by_oid, + &mut ddl, + ); + Ok(ddl) +} + +/// Renders `root` and every descendant reachable through `children_by_parent`. +/// Iterative (an explicit stack, not function-call recursion) so that a +/// pathologically deep partition hierarchy can't overflow the stack; `visited` +/// additionally guards against a corrupted catalog (or non-PostgreSQL fork) +/// whose `pg_inherits` data forms a cycle, which would otherwise loop forever. +#[allow(clippy::too_many_arguments)] +fn render_postgres_partition_tree_node( + root: &db::postgres::PostgresPartitionTreeNode, + children_by_parent: &HashMap>, + columns_by_oid: &HashMap>, + indexes_by_oid: &HashMap>, + fkeys_by_relation: &HashMap<(String, String), Vec>, + comments_by_oid: &HashMap>, + triggers_by_oid: &HashMap>, + checks_by_oid: &HashMap>, + local_objects_by_oid: &HashMap, + ddl: &mut String, +) { + let empty_columns = Vec::new(); + let empty_indexes = Vec::new(); + let empty_fkeys = Vec::new(); + let empty_triggers = Vec::new(); + let empty_checks = Vec::new(); + let empty_local_objects = db::postgres::PostgresTablePartitionLocalObjects::default(); + + let mut visited: HashSet = HashSet::new(); + let mut stack: Vec<&db::postgres::PostgresPartitionTreeNode> = vec![root]; + while let Some(node) = stack.pop() { + if !visited.insert(node.oid) { + continue; + } + + let columns = columns_by_oid.get(&node.oid).unwrap_or(&empty_columns); + let indexes = indexes_by_oid.get(&node.oid).unwrap_or(&empty_indexes); + let fkeys = fkeys_by_relation.get(&(node.schema.clone(), node.table.clone())).unwrap_or(&empty_fkeys); + let comment = comments_by_oid.get(&node.oid).cloned().flatten(); + let triggers = triggers_by_oid.get(&node.oid).unwrap_or(&empty_triggers); + let checks = checks_by_oid.get(&node.oid).unwrap_or(&empty_checks); + let local_objects = local_objects_by_oid.get(&node.oid).unwrap_or(&empty_local_objects); + + if !ddl.is_empty() { + ddl.push('\n'); + } + ddl.push_str(&append_postgres_trigger_definitions( + render_postgres_table_ddl_with_partition_info( + &node.schema, + &node.table, + columns, + indexes, + fkeys, + checks, + comment.as_deref(), + &node.partition_info, + local_objects, + ), + triggers, + )); + + if let Some(children) = children_by_parent.get(&node.oid) { + // Push in reverse so children are popped (and rendered) in their + // original relname-sorted order. + for child in children.iter().rev() { + stack.push(child); + } + } + } +} + async fn pg_display_ddl(pool: &deadpool_postgres::Pool, schema: &str, table: &str) -> Result { - let (ddl, access) = tokio::join!(pg_ddl(pool, schema, table), db::postgres::get_table_access(pool, schema, table)); + let (ddl, access) = + tokio::join!(pg_ddl_with_partitions(pool, schema, table), db::postgres::get_table_access(pool, schema, table)); let ddl = ddl?; match access { Ok(access) => Ok(append_postgres_access_ddl(ddl, schema, table, &access)), @@ -9577,11 +9970,16 @@ fn append_opengauss_trigger_definitions(mut ddl: String, trigger_definitions: &[ ddl } -pub async fn cloudberry_ddl(pool: &deadpool_postgres::Pool, schema: &str, table: &str) -> Result { +pub async fn cloudberry_ddl( + pool: &deadpool_postgres::Pool, + schema: &str, + table: &str, + include_partitions: bool, +) -> Result { match db::cloudberry::table_ddl(pool, schema, table).await { Ok(ddl) => Ok(ddl), Err(native_error) => { - let base_ddl = pg_ddl(pool, schema, table).await.map_err(|fallback_error| { + let base_ddl = pg_ddl_for_options(pool, schema, table, include_partitions).await.map_err(|fallback_error| { format!( "Cloudberry pg_get_tabledef failed: {native_error}; PostgreSQL DDL fallback failed: {fallback_error}" ) @@ -9598,8 +9996,13 @@ pub async fn cloudberry_ddl(pool: &deadpool_postgres::Pool, schema: &str, table: } } -pub async fn opentenbase_ddl(pool: &deadpool_postgres::Pool, schema: &str, table: &str) -> Result { - let ddl = pg_ddl(pool, schema, table).await?; +pub async fn opentenbase_ddl( + pool: &deadpool_postgres::Pool, + schema: &str, + table: &str, + include_partitions: bool, +) -> Result { + let ddl = pg_ddl_for_options(pool, schema, table, include_partitions).await?; match db::opentenbase::table_distribution(pool, schema, table).await { Ok(Some(distribution)) => match db::opentenbase::append_distribution_clause(&ddl, &distribution) { Ok(ddl) => Ok(ddl), @@ -9640,6 +10043,7 @@ pub fn render_postgres_table_ddl( columns, indexes, fkeys, + &[], table_comment, &db::postgres::PostgresTablePartitionInfo::default(), &db::postgres::PostgresTablePartitionLocalObjects::default(), @@ -9652,6 +10056,7 @@ fn render_postgres_table_ddl_with_partition_info( columns: &[db::ColumnInfo], indexes: &[db::IndexInfo], fkeys: &[db::ForeignKeyInfo], + check_constraints: &[(String, String)], table_comment: Option<&str>, partition_info: &db::postgres::PostgresTablePartitionInfo, partition_local_objects: &db::postgres::PostgresTablePartitionLocalObjects, @@ -9728,7 +10133,42 @@ fn render_postgres_table_ddl_with_partition_info( ref_columns )); } + // `pg_get_constraintdef` appends ` NOT VALID` for an unvalidated CHECK + // constraint, but that suffix is only legal after `ALTER TABLE ADD + // CONSTRAINT` — it's a syntax error inside a `CREATE TABLE` column list. + // Emit those as a separate statement below instead, so the constraint's + // unvalidated state round-trips instead of producing invalid DDL. + let mut not_valid_check_constraints: Vec<(&str, &str)> = Vec::new(); + for (name, definition) in check_constraints { + if is_partition && !partition_local_objects.check_constraints.contains(name) { + continue; + } + let definition = definition.trim(); + if is_not_valid_constraintdef(definition) { + not_valid_check_constraints.push((name.as_str(), definition)); + continue; + } + definition_lines.push(format!(" CONSTRAINT {} {}", pg_ident(name), definition)); + } + if is_partition { + // A partition can override a column's default independently of the + // parent; PostgreSQL only accepts that override through `column_name + // WITH OPTIONS DEFAULT ...` since the partition's own column list is + // otherwise inherited (and thus omitted) from the parent's. + for column in columns { + if partition_local_objects.column_defaults.get(&column.name) + != Some(&db::postgres::PostgresColumnDefaultState::Overridden) + { + continue; + } + let Some(default) = column.column_default.as_deref() else { + continue; + }; + definition_lines.push(format!(" {} WITH OPTIONS DEFAULT {default}", pg_ident(&column.name))); + } + } + let create = if partition_info.is_foreign { "CREATE FOREIGN TABLE" } else { "CREATE TABLE" }; let mut ddl = if let Some((parent_schema, parent_table, bound)) = partition_parent { let parent_name = format!("{}.{}", pg_ident(parent_schema), pg_ident(parent_table)); let definitions = if definition_lines.is_empty() { @@ -9736,15 +10176,49 @@ fn render_postgres_table_ddl_with_partition_info( } else { format!(" (\n{}\n)", definition_lines.join(",\n")) }; - format!("CREATE TABLE {table_name} PARTITION OF {parent_name}{definitions} {bound}") + format!("{create} {table_name} PARTITION OF {parent_name}{definitions} {bound}") } else { - format!("CREATE TABLE {table_name} (\n{}\n)", definition_lines.join(",\n")) + format!("{create} {table_name} (\n{}\n)", definition_lines.join(",\n")) }; + if let Some(server) = partition_info.foreign_server.as_deref().filter(|server| !server.trim().is_empty()) { + ddl.push_str(&format!(" SERVER {}", pg_ident(server))); + if !partition_info.foreign_options.is_empty() { + let options = partition_info + .foreign_options + .iter() + .map(|(key, value)| format!("{} {}", pg_ident(key), sql_string(value))) + .collect::>() + .join(", "); + ddl.push_str(&format!(" OPTIONS ({options})")); + } + } if let Some(partition_key) = partition_info.key.as_deref().filter(|key| !key.trim().is_empty()) { ddl.push_str(&format!(" PARTITION BY {partition_key}")); } ddl.push_str(";\n"); + for (name, definition) in ¬_valid_check_constraints { + ddl.push_str(&format!("\nALTER TABLE {table_name} ADD CONSTRAINT {} {};", pg_ident(name), definition)); + } + + if is_partition { + // A dropped default has no counterpart in the PARTITION OF column + // list syntax used above for overrides — it must be replayed as a + // standalone statement, or restore would silently reintroduce the + // parent's default (PostgreSQL auto-copies it onto every partition + // at creation time unless explicitly dropped). + for column in columns { + if partition_local_objects.column_defaults.get(&column.name) + == Some(&db::postgres::PostgresColumnDefaultState::Dropped) + { + ddl.push_str(&format!( + "\nALTER TABLE ONLY {table_name} ALTER COLUMN {} DROP DEFAULT;", + pg_ident(&column.name) + )); + } + } + } + if let Some(comment) = table_comment.filter(|comment| !comment.trim().is_empty()) { ddl.push_str(&format!("\nCOMMENT ON TABLE {table_name} IS {};", sql_string(comment))); } diff --git a/crates/dbx-core/src/schema/providers/native.rs b/crates/dbx-core/src/schema/providers/native.rs index dd69680938..e715e55bb6 100644 --- a/crates/dbx-core/src/schema/providers/native.rs +++ b/crates/dbx-core/src/schema/providers/native.rs @@ -258,10 +258,10 @@ pub(in crate::schema) async fn table_ddl( } } PoolKind::Postgres(p) if config.is_some_and(is_cloudberry_config) => { - super::super::cloudberry_ddl(p, schema, table).await + super::super::cloudberry_ddl(p, schema, table, false).await } PoolKind::Postgres(p) if config.is_some_and(db::opentenbase::is_config) => { - super::super::opentenbase_ddl(p, schema, table).await + super::super::opentenbase_ddl(p, schema, table, false).await } PoolKind::Postgres(p) => super::super::pg_ddl(p, schema, table).await, PoolKind::Sqlite(p) => super::super::sqlite_ddl(p, schema, table).await, diff --git a/crates/dbx-core/tests/database_export_partition_ddl.rs b/crates/dbx-core/tests/database_export_partition_ddl.rs new file mode 100644 index 0000000000..5204206171 --- /dev/null +++ b/crates/dbx-core/tests/database_export_partition_ddl.rs @@ -0,0 +1,365 @@ +//! Regression coverage for PR review points on partition DDL: +//! (1) database export must not duplicate a partition's `CREATE TABLE` +//! (once via the parent's embedded tree, once via the export loop's own +//! iteration over that same child relation) and the exported SQL must +//! replay cleanly into a fresh schema; (2) a partition-local dropped +//! column default must be modeled and replayed as a standalone +//! `ALTER TABLE ONLY ... DROP DEFAULT`, not silently reintroduced; +//! (3) fetching a whole partition tree's DDL must issue a small, +//! roughly-constant number of queries rather than one scaling with the +//! number of partitions. +//! Needs local Docker; skips cleanly when unavailable. + +use std::process::Command; + +use dbx_core::connection::AppState; +use dbx_core::database_export::{export_database_sql_core, DatabaseExportRequest}; +use dbx_core::models::connection::{ConnectionConfig, DatabaseType}; +use dbx_core::storage::Storage; + +struct DockerPostgres { + name: String, + port: u16, +} + +impl Drop for DockerPostgres { + fn drop(&mut self) { + let _ = Command::new("docker").args(["rm", "-f", &self.name]).status(); + } +} + +fn docker_ready() -> bool { + Command::new("docker") + .args(["version", "--format", "{{.Server.Version}}"]) + .output() + .map(|output| output.status.success()) + .unwrap_or(false) +} + +fn start_docker_postgres(name_prefix: &str) -> Option { + if !docker_ready() { + eprintln!("skipping docker-backed partition ddl test because Docker is unavailable"); + return None; + } + + let port = portpicker::pick_unused_port().expect("pick unused postgres port"); + let container = DockerPostgres { name: format!("{name_prefix}-{}", uuid::Uuid::new_v4()), port }; + + let status = Command::new("docker") + .args([ + "run", + "-d", + "--rm", + "--name", + &container.name, + "-e", + "POSTGRES_PASSWORD=postgres", + "-e", + "POSTGRES_USER=postgres", + "-e", + "POSTGRES_DB=postgres", + "-p", + &format!("{port}:5432"), + "postgres:16-alpine", + ]) + .status() + .expect("start docker postgres"); + assert!(status.success(), "docker run postgres container should succeed"); + + let mut consecutive_ok = 0; + for _ in 0..120 { + let ready = Command::new("docker") + .args(["exec", &container.name, "psql", "-U", "postgres", "-d", "postgres", "-c", "SELECT 1"]) + .output() + .map(|output| output.status.success()) + .unwrap_or(false); + consecutive_ok = if ready { consecutive_ok + 1 } else { 0 }; + if consecutive_ok >= 2 { + return Some(container); + } + std::thread::sleep(std::time::Duration::from_millis(500)); + } + panic!("postgres container did not become ready"); +} + +fn psql(container: &DockerPostgres, sql: &str) { + let output = Command::new("docker") + .args(["exec", &container.name, "psql", "-U", "postgres", "-d", "postgres", "-v", "ON_ERROR_STOP=1", "-c", sql]) + .output() + .expect("run psql"); + assert!(output.status.success(), "psql failed: {}", String::from_utf8_lossy(&output.stderr)); +} + +fn psql_allow_failure(container: &DockerPostgres, sql: &str) -> bool { + Command::new("docker") + .args(["exec", &container.name, "psql", "-U", "postgres", "-d", "postgres", "-v", "ON_ERROR_STOP=1", "-c", sql]) + .output() + .map(|output| output.status.success()) + .unwrap_or(false) +} + +fn postgres_test_config(id: &str, port: u16) -> ConnectionConfig { + ConnectionConfig { + docs_notes_path: None, + id: id.to_string(), + name: id.to_string(), + note: String::new(), + db_type: DatabaseType::Postgres, + driver_profile: None, + driver_label: None, + url_params: None, + agent_java_options: Vec::new(), + host: "127.0.0.1".to_string(), + port, + username: "postgres".to_string(), + password: "postgres".to_string(), + database: Some("postgres".to_string()), + default_schema: None, + visible_databases: None, + visible_schemas: None, + attached_databases: Vec::new(), + init_script: None, + color: None, + transport_layers: Vec::new(), + connect_timeout_secs: 5, + query_timeout_secs: 30, + idle_timeout_secs: 60, + keepalive_interval_secs: 0, + ssl: false, + ca_cert_path: String::new(), + client_cert_path: String::new(), + client_key_path: String::new(), + sysdba: false, + oracle_connection_type: None, + connection_string: None, + redis_connection_mode: None, + redis_sentinel_master: String::new(), + redis_sentinel_nodes: String::new(), + redis_sentinel_username: String::new(), + redis_sentinel_password: String::new(), + redis_sentinel_tls: false, + redis_cluster_nodes: String::new(), + redis_key_separator: dbx_core::models::connection::default_redis_key_separator(), + redis_scan_page_size: None, + redis_database_aliases: Default::default(), + etcd_endpoints: String::new(), + gbase_server: String::new(), + informix_server: String::new(), + external_config: None, + jdbc_driver_class: None, + jdbc_driver_paths: Vec::new(), + one_time: false, + save_password: true, + read_only: false, + is_production: false, + production_databases: vec![], + show_system_schemas: false, + database_info: None, + } +} + +/// Full database export of a partitioned schema must not duplicate any +/// relation's `CREATE TABLE`/`CREATE FOREIGN TABLE`, must faithfully replay +/// a partition-local dropped default, and the exported SQL must replay +/// cleanly into a fresh schema (the actual "restore" check the review +/// comment asked for). +#[tokio::test] +async fn database_export_of_partition_tree_has_no_duplicates_and_replays() { + let Some(container) = start_docker_postgres("dbx-export-partition-ddl") else { + return; + }; + + psql( + &container, + "CREATE EXTENSION IF NOT EXISTS postgres_fdw;\ + CREATE SERVER dbx_export_loopback FOREIGN DATA WRAPPER postgres_fdw \ + OPTIONS (host 'localhost', dbname 'postgres', port '5432');\ + CREATE USER MAPPING FOR CURRENT_USER SERVER dbx_export_loopback \ + OPTIONS (user 'postgres', password 'postgres');\ + CREATE TABLE foreign_child_data (id integer, status text);\ + CREATE TABLE events (id bigint NOT NULL, status text DEFAULT 'active', created_at timestamp NOT NULL) \ + PARTITION BY RANGE (created_at);\ + CREATE TABLE events_2026 PARTITION OF events FOR VALUES FROM ('2026-01-01') TO ('2027-01-01');\ + CREATE TABLE events_2027 PARTITION OF events FOR VALUES FROM ('2027-01-01') TO ('2028-01-01') PARTITION BY RANGE (created_at);\ + CREATE TABLE events_2027_h1 PARTITION OF events_2027 FOR VALUES FROM ('2027-01-01') TO ('2027-07-01');\ + CREATE TABLE events_2027_h2 PARTITION OF events_2027 FOR VALUES FROM ('2027-07-01') TO ('2028-01-01');\ + CREATE FOREIGN TABLE events_foreign PARTITION OF events FOR VALUES FROM ('2028-01-01') TO ('2029-01-01') \ + SERVER dbx_export_loopback OPTIONS (schema_name 'public', table_name 'foreign_child_data');\ + ALTER TABLE ONLY events_2027_h1 ALTER COLUMN status DROP DEFAULT;\ + ALTER TABLE ONLY events_2027_h2 ALTER COLUMN status SET DEFAULT 'archived';", + ); + + let dir = std::env::temp_dir().join(format!("dbx-export-partition-ddl-{}", uuid::Uuid::new_v4())); + std::fs::create_dir_all(&dir).unwrap(); + let storage = Storage::open(&dir.join("storage.db")).await.unwrap(); + let state = AppState::new(storage); + + let connection_id = "export-partition-ddl-conn"; + state.configs.write().await.insert(connection_id.to_string(), postgres_test_config(connection_id, container.port)); + + let file_path = dir.join("export.sql"); + let request = DatabaseExportRequest { + export_id: format!("export-partition-ddl-{}", uuid::Uuid::new_v4()), + connection_id: connection_id.to_string(), + database: "postgres".to_string(), + schema: "public".to_string(), + file_path: file_path.to_string_lossy().to_string(), + selected_tables: Vec::new(), + excluded_tables: Vec::new(), + include_structure: true, + include_data: false, + include_objects: false, + include_create_database: false, + drop_table_if_exists: false, + omit_auto_increment: false, + fail_on_error: true, + snapshot_session_id: None, + batch_size: 1000, + }; + + export_database_sql_core(&state, &request, |_progress| {}).await.expect("export should succeed"); + + let exported = std::fs::read_to_string(&file_path).expect("read exported file"); + assert!(!exported.contains("-- ERROR"), "exported SQL should not contain errors:\n{exported}"); + + for table in ["events", "events_2026", "events_2027", "events_2027_h1", "events_2027_h2", "events_foreign"] { + let needle = format!("TABLE \"public\".\"{table}\""); + let count = exported.matches(&needle).count(); + assert_eq!(count, 1, "expected exactly one CREATE (FOREIGN) TABLE for {table}, found {count}:\n{exported}"); + } + assert!(exported.contains("CREATE FOREIGN TABLE \"public\".\"events_foreign\""), "exported: {exported}"); + + assert!( + exported.contains("ALTER TABLE ONLY \"public\".\"events_2027_h1\" ALTER COLUMN \"status\" DROP DEFAULT;"), + "dropped default should be replayed as a standalone statement:\n{exported}" + ); + assert!( + exported.contains("\"status\" WITH OPTIONS DEFAULT 'archived'::text"), + "overridden default should still render inline:\n{exported}" + ); + // The untouched sibling (events_2026) must not gain either a drop or an + // inline override — it purely inherits the parent's default. + let events_2026_ddl_start = exported.find("PARTITION OF \"public\".\"events\"").expect("events_2026 ddl present"); + let events_2026_ddl = &exported[events_2026_ddl_start..(events_2026_ddl_start + 400).min(exported.len())]; + assert!(!events_2026_ddl.contains("DROP DEFAULT"), "events_2026 ddl: {events_2026_ddl}"); + assert!(!events_2026_ddl.contains("WITH OPTIONS DEFAULT"), "events_2026 ddl: {events_2026_ddl}"); + + // The actual "restore" check: replay the exported SQL into a fresh + // schema and confirm no `relation already exists` (or any other) error. + psql(&container, "CREATE SCHEMA replay"); + let replayable = exported.replace("\"public\".", "\"replay\".").replace("SCHEMA \"public\"", "SCHEMA \"replay\""); + let replay_file = dir.join("replay.sql"); + std::fs::write(&replay_file, &replayable).unwrap(); + let output = Command::new("docker") + .args(["cp", replay_file.to_str().unwrap(), &format!("{}:/tmp/replay.sql", container.name)]) + .output() + .expect("copy replay file into container"); + assert!(output.status.success(), "docker cp failed: {}", String::from_utf8_lossy(&output.stderr)); + let replay_output = Command::new("docker") + .args([ + "exec", + &container.name, + "psql", + "-U", + "postgres", + "-d", + "postgres", + "-v", + "ON_ERROR_STOP=1", + "-f", + "/tmp/replay.sql", + ]) + .output() + .expect("replay exported sql"); + assert!( + replay_output.status.success(), + "replaying the exported SQL should succeed (no duplicate/`already exists` errors): {}", + String::from_utf8_lossy(&replay_output.stderr) + ); +} + +/// Fetching a whole partition tree's DDL (the "view DDL" interactive path) +/// must issue a small, roughly-constant number of queries — not one that +/// scales linearly with the number of leaf partitions. +#[tokio::test] +async fn partition_tree_ddl_query_count_does_not_scale_with_partition_count() { + let Some(container) = start_docker_postgres("dbx-partition-ddl-query-count") else { + return; + }; + + if !psql_allow_failure(&container, "ALTER SYSTEM SET log_statement = 'all'") { + eprintln!("skipping query-count test: could not enable log_statement"); + return; + } + psql(&container, "SELECT pg_reload_conf()"); + + let build_tree = |leaves: u32| -> String { + let mut sql = String::from( + "CREATE TABLE t (id bigint NOT NULL, bucket integer NOT NULL, status text DEFAULT 'active') PARTITION BY RANGE (bucket);", + ); + for i in 0..leaves { + sql.push_str(&format!( + "CREATE TABLE t_{i} PARTITION OF t FOR VALUES FROM ({start}) TO ({end});", + i = i, + start = i * 10, + end = (i + 1) * 10 + )); + } + sql + }; + + let dir = std::env::temp_dir().join(format!("dbx-partition-query-count-{}", uuid::Uuid::new_v4())); + std::fs::create_dir_all(&dir).unwrap(); + let storage = Storage::open(&dir.join("storage.db")).await.unwrap(); + let state = AppState::new(storage); + let connection_id = "partition-query-count-conn"; + state.configs.write().await.insert(connection_id.to_string(), postgres_test_config(connection_id, container.port)); + + async fn fetch_ddl(state: &AppState, connection_id: &str) -> String { + dbx_core::schema::get_table_display_ddl_core(state, connection_id, "postgres", "public", "t", None) + .await + .expect("fetch partition tree ddl") + } + + let logs_between = |container: &DockerPostgres, start_marker: &str, end_marker: &str| -> usize { + let output = Command::new("docker").args(["logs", &container.name]).output().expect("docker logs"); + let combined = + format!("{}{}", String::from_utf8_lossy(&output.stdout), String::from_utf8_lossy(&output.stderr)); + let start = combined.rfind(start_marker).map(|i| i + start_marker.len()); + let end = combined.rfind(end_marker); + match (start, end) { + (Some(s), Some(e)) if e > s => { + combined[s..e].lines().filter(|line| line.contains("statement:") || line.contains("execute")).count() + } + _ => usize::MAX, + } + }; + + psql(&container, &format!("DROP TABLE IF EXISTS t CASCADE; {}", build_tree(3))); + psql(&container, "SELECT 'MARKER_SMALL_START'"); + let small_ddl = fetch_ddl(&state, connection_id).await; + psql(&container, "SELECT 'MARKER_SMALL_END'"); + assert_eq!(small_ddl.matches("PARTITION OF").count(), 3, "small tree ddl: {small_ddl}"); + + psql(&container, &format!("DROP TABLE IF EXISTS t CASCADE; {}", build_tree(40))); + psql(&container, "SELECT 'MARKER_LARGE_START'"); + let large_ddl = fetch_ddl(&state, connection_id).await; + psql(&container, "SELECT 'MARKER_LARGE_END'"); + assert_eq!(large_ddl.matches("PARTITION OF").count(), 40, "large tree ddl: {large_ddl}"); + + // Give the container's log a moment to flush before reading it back. + std::thread::sleep(std::time::Duration::from_millis(500)); + let small_count = logs_between(&container, "MARKER_SMALL_START", "MARKER_SMALL_END"); + let large_count = logs_between(&container, "MARKER_LARGE_START", "MARKER_LARGE_END"); + + assert_ne!(small_count, usize::MAX, "could not locate small-tree marker window in container logs"); + assert_ne!(large_count, usize::MAX, "could not locate large-tree marker window in container logs"); + // A naive per-relation recursion issues ~9 queries per relation, so a + // 40-leaf tree (41 relations) would log roughly 13x the 3-leaf tree's + // (4 relations) query count. The batched implementation should stay + // close to flat instead. + assert!( + large_count <= small_count + 10, + "query count should not scale with partition count: small={small_count} large={large_count}" + ); +} diff --git a/crates/dbx-core/tests/live_cloudberry.rs b/crates/dbx-core/tests/live_cloudberry.rs index c2a562a730..74de4962c3 100644 --- a/crates/dbx-core/tests/live_cloudberry.rs +++ b/crates/dbx-core/tests/live_cloudberry.rs @@ -63,7 +63,7 @@ async fn cloudberry_metadata_and_ddl_round_trip() { ("external_events", "CREATE FOREIGN TABLE"), ]; for (table, expected) in cases { - let ddl = schema::cloudberry_ddl(&pool, &source_schema, table).await?; + let ddl = schema::cloudberry_ddl(&pool, &source_schema, table, false).await?; assert!(ddl.contains(expected), "{table} DDL did not contain {expected}: {ddl}"); if table == "external_events" { assert!(ddl.contains("SERVER \"gp_exttable_server\""), "external DDL: {ddl}"); diff --git a/crates/dbx-core/tests/live_opentenbase.rs b/crates/dbx-core/tests/live_opentenbase.rs index 7da559f165..092e40ea04 100644 --- a/crates/dbx-core/tests/live_opentenbase.rs +++ b/crates/dbx-core/tests/live_opentenbase.rs @@ -46,7 +46,7 @@ async fn opentenbase_metadata_and_ddl_round_trip() { ("replicated_dimensions", "DISTRIBUTE BY REPLICATION"), ]; for (table, expected) in cases { - let ddl = schema::opentenbase_ddl(&pool, &source_schema, table).await?; + let ddl = schema::opentenbase_ddl(&pool, &source_schema, table, false).await?; if !ddl.contains(expected) { return Err(format!("{table} DDL did not contain {expected}: {ddl}")); }