diff --git a/pgdog/src/backend/schema/columns.rs b/pgdog/src/backend/schema/columns.rs index 3a98c5618..33a8af357 100644 --- a/pgdog/src/backend/schema/columns.rs +++ b/pgdog/src/backend/schema/columns.rs @@ -132,10 +132,53 @@ impl From for Column { #[cfg(test)] mod test { + use super::Error; use crate::backend::pool::Request; use crate::backend::pool::test::pool; use crate::backend::schema::columns::Column; + #[tokio::test] + async fn test_load_partitioned_primary_keys() -> Result<(), Error> { + let pool = pool(); + let mut conn = pool.get(&Request::default()).await?; + conn.execute_checked("BEGIN").await?; + conn.execute_checked( + "CREATE SCHEMA pk_partition_test; + CREATE TABLE pk_partition_test.parted ( + tenant_id BIGINT, + id BIGINT, + payload TEXT, + PRIMARY KEY (tenant_id, id) + ) PARTITION BY HASH (tenant_id); + CREATE TABLE pk_partition_test.parted_0 + PARTITION OF pk_partition_test.parted + FOR VALUES WITH (MODULUS 2, REMAINDER 0); + CREATE TABLE pk_partition_test.parted_1 + PARTITION OF pk_partition_test.parted + FOR VALUES WITH (MODULUS 2, REMAINDER 1)", + ) + .await?; + + let columns = Column::load(&mut conn).await?; + for table in ["parted", "parted_0", "parted_1"] { + let table_columns = columns + .get(&("pk_partition_test".to_owned(), table.to_owned())) + .expect("parent and partition columns should be loaded"); + let primary_keys: Vec<_> = table_columns + .iter() + .map(|column| (column.column_name.as_str(), column.is_primary_key)) + .collect(); + assert_eq!( + primary_keys, + [("tenant_id", true), ("id", true), ("payload", false)], + "unexpected primary key columns for {table}" + ); + } + + conn.execute_checked("ROLLBACK").await?; + Ok(()) + } + #[tokio::test] async fn test_load_columns() { let pool = pool(); diff --git a/pgdog/src/backend/schema/columns.sql b/pgdog/src/backend/schema/columns.sql index 22d63ba25..d04477c09 100644 --- a/pgdog/src/backend/schema/columns.sql +++ b/pgdog/src/backend/schema/columns.sql @@ -12,17 +12,23 @@ FROM information_schema.columns c LEFT JOIN ( SELECT - kcu.table_schema, - kcu.table_name, - kcu.column_name + n.nspname AS table_schema, + cl.relname AS table_name, + a.attname AS column_name FROM - information_schema.table_constraints tc + pg_catalog.pg_constraint con JOIN - information_schema.key_column_usage kcu - ON tc.constraint_name = kcu.constraint_name - AND tc.table_schema = kcu.table_schema + pg_catalog.pg_class cl ON cl.oid = con.conrelid + JOIN + pg_catalog.pg_namespace n ON n.oid = cl.relnamespace + JOIN + pg_catalog.pg_attribute a ON a.attrelid = con.conrelid AND a.attnum = ANY (con.conkey) WHERE - tc.constraint_type = 'PRIMARY KEY' + con.contype = 'p' + -- information_schema.table_constraints excludes SELECT-only non-owners, as required by the SQL standard. + AND (pg_catalog.pg_has_role(cl.relowner, 'USAGE') + OR pg_catalog.has_table_privilege(cl.oid, 'INSERT, UPDATE, DELETE, TRUNCATE, REFERENCES, TRIGGER') + OR pg_catalog.has_any_column_privilege(cl.oid, 'INSERT, UPDATE, REFERENCES')) ) pk ON c.table_schema = pk.table_schema AND c.table_name = pk.table_name AND c.column_name = pk.column_name