From aa58a24a6dc863d90d9748c803f143401fae7aca Mon Sep 17 00:00:00 2001
From: Adria Navarro <adria@budibase.com>
Date: Tue, 23 May 2023 11:20:18 +0200
Subject: [PATCH] Change query for retrieving postgres primary keys

---
 packages/server/src/integrations/postgres.ts | 21 +++++++++++---------
 1 file changed, 12 insertions(+), 9 deletions(-)

diff --git a/packages/server/src/integrations/postgres.ts b/packages/server/src/integrations/postgres.ts
index bf77ec08c6..9d6ec5969d 100644
--- a/packages/server/src/integrations/postgres.ts
+++ b/packages/server/src/integrations/postgres.ts
@@ -126,14 +126,15 @@ class PostgresIntegration extends Sql implements DatasourcePlus {
 
   COLUMNS_SQL!: string
 
-  PRIMARY_KEYS_SQL = `
-  select tc.table_schema, tc.table_name, kc.column_name as primary_key 
-  from information_schema.table_constraints tc
-  join 
-    information_schema.key_column_usage kc on kc.table_name = tc.table_name 
-    and kc.table_schema = tc.table_schema 
-    and kc.constraint_name = tc.constraint_name
-  where tc.constraint_type = 'PRIMARY KEY';
+  PRIMARY_KEYS_SQL = () => `
+  SELECT pg_namespace.nspname table_schema
+     , pg_class.relname table_name
+     , pg_attribute.attname primary_key
+  FROM pg_class
+  JOIN pg_index ON pg_class.oid = pg_index.indrelid AND pg_index.indisprimary
+  JOIN pg_attribute ON pg_attribute.attrelid = pg_class.oid AND pg_attribute.attnum = ANY(pg_index.indkey)
+  JOIN pg_namespace ON pg_namespace.oid = pg_class.relnamespace
+  WHERE pg_namespace.nspname = '${this.config.schema}';
   `
 
   constructor(config: PostgresConfig) {
@@ -239,7 +240,9 @@ class PostgresIntegration extends Sql implements DatasourcePlus {
     let tableKeys: { [key: string]: string[] } = {}
     await this.openConnection()
     try {
-      const primaryKeysResponse = await this.client.query(this.PRIMARY_KEYS_SQL)
+      const primaryKeysResponse = await this.client.query(
+        this.PRIMARY_KEYS_SQL()
+      )
       for (let table of primaryKeysResponse.rows) {
         const tableName = table.table_name
         if (!tableKeys[tableName]) {