From ab813cb601d58a075a49f8f3b6a8b001adc343de Mon Sep 17 00:00:00 2001
From: MAFLO321 <MAFLO321@users.noreply.github.com>
Date: Tue, 12 Nov 2024 10:18:29 +0100
Subject: [PATCH] Postgresql: Support tables from non-default schema (#95636)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* Postgresql: Support tables from non-default schema

- Add support for schema-qualified table names.
- Partially resolve an issue where the column type of a table from the
  wrong schema with the same table name was incorrectly used. Now
  limited to tables of schemas within the search_path.

* Support schema in raw query editor

---------

Co-authored-by: Zoltán Bedi <zoltan.bedi@gmail.com>
---
 .../postgresMetaQuery.ts                      | 24 +++++++++++++------
 .../sqlCompletionProvider.ts                  | 19 +++++++++++++++
 2 files changed, 36 insertions(+), 7 deletions(-)

diff --git a/public/app/plugins/datasource/grafana-postgresql-datasource/postgresMetaQuery.ts b/public/app/plugins/datasource/grafana-postgresql-datasource/postgresMetaQuery.ts
index 664a8e9c4dc..ac19fc286e2 100644
--- a/public/app/plugins/datasource/grafana-postgresql-datasource/postgresMetaQuery.ts
+++ b/public/app/plugins/datasource/grafana-postgresql-datasource/postgresMetaQuery.ts
@@ -7,8 +7,13 @@ export function getTimescaleDBVersion() {
 }
 
 export function showTables() {
-  return `select quote_ident(table_name) as "table" from information_schema.tables
-    where quote_ident(table_schema) not in ('information_schema',
+  return `SELECT
+    CASE WHEN ${buildSchemaConstraint()}
+      THEN quote_ident(table_name)
+      ELSE quote_ident(table_schema) || '.' || quote_ident(table_name)
+    END AS "table"
+    FROM information_schema.tables
+    WHERE quote_ident(table_schema) NOT IN ('information_schema',
                              'pg_catalog',
                              '_timescaledb_cache',
                              '_timescaledb_catalog',
@@ -16,7 +21,7 @@ export function showTables() {
                              '_timescaledb_config',
                              'timescaledb_information',
                              'timescaledb_experimental')
-      and ${buildSchemaConstraint()}`;
+    ORDER BY CASE WHEN ${buildSchemaConstraint()} THEN 0 ELSE 1 END, 1`;
 }
 
 export function getSchema(table: string) {
@@ -24,10 +29,15 @@ export function getSchema(table: string) {
   // in the table-name
   const tableNamePart = "'" + table.replace(/'/g, "''") + "'";
 
-  return `select quote_ident(column_name) as "column", data_type as "type"
-    from information_schema.columns
-    where quote_ident(table_name) = ${tableNamePart};
-    `;
+  return `SELECT quote_ident(column_name) AS "column", data_type AS "type"
+    FROM information_schema.columns
+    WHERE
+      CASE WHEN array_length(parse_ident(${tableNamePart}),1) = 2
+        THEN quote_ident(table_schema) = (parse_ident(${tableNamePart}))[1]
+          AND quote_ident(table_name) = (parse_ident(${tableNamePart}))[2]
+        ELSE quote_ident(table_name) = ${tableNamePart}
+          AND ${buildSchemaConstraint()}
+      END`;
 }
 
 function buildSchemaConstraint() {
diff --git a/public/app/plugins/datasource/grafana-postgresql-datasource/sqlCompletionProvider.ts b/public/app/plugins/datasource/grafana-postgresql-datasource/sqlCompletionProvider.ts
index 9c1c4f92205..a1936af881f 100644
--- a/public/app/plugins/datasource/grafana-postgresql-datasource/sqlCompletionProvider.ts
+++ b/public/app/plugins/datasource/grafana-postgresql-datasource/sqlCompletionProvider.ts
@@ -2,8 +2,10 @@ import {
   ColumnDefinition,
   getStandardSQLCompletionProvider,
   LanguageCompletionProvider,
+  LinkedToken,
   TableDefinition,
   TableIdentifier,
+  TokenType,
 } from '@grafana/experimental';
 import { DB, SQLQuery } from '@grafana/sql';
 
@@ -20,6 +22,23 @@ export const getSqlCompletionProvider: (args: CompletionProviderGetterArgs) => L
       resolve: async () => {
         return await getTables.current();
       },
+      // Default parser doesn't handle schema.table syntax
+      parseName: (token: LinkedToken | undefined | null) => {
+        if (!token) {
+          return { table: '' };
+        }
+
+        let processedToken = token;
+        let tablePath = processedToken.value;
+
+        // Parse schema.table syntax
+        while (processedToken.next && processedToken.next.type !== TokenType.Whitespace) {
+          tablePath += processedToken.next.value;
+          processedToken = processedToken.next;
+        }
+
+        return { table: tablePath };
+      },
     },
     columns: {
       resolve: async (t?: TableIdentifier) => {