From d7858271467083e0729ca2016381be81203e7c07 Mon Sep 17 00:00:00 2001
From: Sander van Harmelen <sander@vanharmelen.nl>
Date: Wed, 20 Nov 2024 18:03:42 +0100
Subject: [PATCH] Fix query to get primary column names for postgres

---
 drivers/postgres.go | 24 ++++++++++++++++--------
 1 file changed, 16 insertions(+), 8 deletions(-)

diff --git a/drivers/postgres.go b/drivers/postgres.go
index 546221f..466bddb 100644
--- a/drivers/postgres.go
+++ b/drivers/postgres.go
@@ -861,10 +861,13 @@ func (db *Postgres) GetPrimaryKeyColumnNames(database, table string) (primaryKey
 
 	splitTableString := strings.Split(table, ".")
 
-	if len(splitTableString) == 1 {
+	if len(splitTableString) != 2 {
 		return nil, errors.New("table must be in the format schema.table")
 	}
 
+	schemaName := splitTableString[0]
+	tableName := splitTableString[1]
+
 	if database != db.CurrentDatabase {
 		err = db.SwitchDatabase(database)
 		if err != nil {
@@ -878,15 +881,20 @@ func (db *Postgres) GetPrimaryKeyColumnNames(database, table string) (primaryKey
 		}()
 	}
 
-	tableName := splitTableString[1]
-
 	row, err := db.Connection.Query(`
-	SELECT a.attname AS column_name
-	FROM pg_index i
-	JOIN pg_attribute a ON a.attrelid = i.indrelid AND a.attnum = ANY(i.indkey)
-	WHERE i.indrelid = $1::regclass AND i.indisprimary
-	`, tableName)
+		SELECT
+			a.attname AS column_name
+		FROM
+			pg_index i
+			JOIN pg_class c ON c.oid = i.indrelid
+			JOIN pg_attribute a ON a.attrelid = c.oid
+				AND a.attnum = ANY (i.indkey)
+			JOIN pg_namespace n ON n.oid = c.relnamespace
+		WHERE
+			relname = $2 AND nspname = $1 AND indisprimary
+	`, schemaName, tableName)
 	if err != nil {
+		logger.Error("GetPrimaryKeyColumnNames", map[string]any{"error": err.Error()})
 		return nil, err
 	}