fix: faster schema loading query + handle partitioned tables

pgdog/src/backend/schema/columns.rs

@@ -132,10 +132,53 @@ impl From<DataRow> 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();

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