From 327d52c797efbe6b4c755ba4d9af5e2df8d6381c Mon Sep 17 00:00:00 2001
From: Mark Roberts <markandrusroberts@gmail.com>
Date: Tue, 20 Aug 2024 12:52:02 +0200
Subject: [PATCH] fix: allow qualifying tables & columns by database & schema

---
 src/syntax/ast.ts                    |  4 ++-
 src/syntax/base.ne                   | 37 ++++++++++++++++++++++------
 src/syntax/simple-statements.spec.ts | 32 ++++++++++++++++++++++++
 src/to-sql.ts                        | 12 +++++++++
 4 files changed, 76 insertions(+), 9 deletions(-)

diff --git a/src/syntax/ast.ts b/src/syntax/ast.ts
index e405263..95570b0 100644
--- a/src/syntax/ast.ts
+++ b/src/syntax/ast.ts
@@ -531,13 +531,15 @@ export interface TableAliasName extends Name, PGNode {
 }
 
 export interface QName extends Name, PGNode {
+    database?: string;
     schema?: string;
 }
 
 export interface QColumn extends PGNode {
+    database?: string;
+    schema?: string;
     table: string;
     column: string;
-    schema?: string;
 }
 
 export type DataTypeDef = ArrayDataTypeDef | BasicDataTypeDef;
diff --git a/src/syntax/base.ne b/src/syntax/base.ne
index 68c50fa..09c5151 100644
--- a/src/syntax/base.ne
+++ b/src/syntax/base.ne
@@ -316,17 +316,26 @@ ident_aliased -> (%kw_as ident {% last %}) | ident {% unwrap %}
 
 table_ref -> qualified_name {% unwrap %}
 
-qcolumn -> ident dot ident (dot ident {% last %}):? {% x => {
-                if (!x[3]) {
-                    return track(x, {
-                            table: unbox(x[0]),
-                            column: unbox(x[2]),
-                        });
-                }
+
+qcolumn -> ident dot ident dot ident dot ident {% x => {
+                return track(x, {
+                    database: unbox(x[0]),
+                    schema: unbox(x[2]),
+                    table: unbox(x[4]),
+                    column: unbox(x[6]),
+                });
+            } %}
+        |  ident dot ident dot ident {% x => {
                 return track(x, {
                     schema: unbox(x[0]),
                     table: unbox(x[2]),
-                    column: unbox(x[3]),
+                    column: unbox(x[4]),
+                });
+            } %}
+        |  ident dot ident {% x => {
+                return track(x, {
+                    table: unbox(x[0]),
+                    column: unbox(x[2]),
                 });
             } %}
 
@@ -342,6 +351,12 @@ table_ref_aliased -> table_ref ident_aliased:? {% x => {
 
 
 qualified_name -> qname_ident {% x => track(x, {name: toStr(x) }) %}
+        | ident dot ident dot ident_extended {% x => {
+                const database = toStr(x[0]);
+                const schema = toStr(x[2]);
+                const name = toStr(x[4]);
+                return track(x, {database, schema, name});
+            } %}
         | ident dot ident_extended {% x => {
                 const schema = toStr(x[0]);
                 const name = toStr(x[2]);
@@ -350,6 +365,12 @@ qualified_name -> qname_ident {% x => track(x, {name: toStr(x) }) %}
         | %kw_current_schema {% x => track(x, { name: 'current_schema' }) %}
 
 qualified_name_mark_quotes -> qname_ident {% x => track(x, {name: toStr(x), ...doubleQuoted(x) }) %}
+        | ident dot ident dot ident_extended {% x => {
+                const database = toStr(x[0]);
+                const schema = toStr(x[2]);
+                const name = toStr(x[4]);
+                return track(x, {database, schema, name, ...doubleQuoted(x[4])});
+            } %}
         | ident dot ident_extended {% x => {
                 const schema = toStr(x[0]);
                 const name = toStr(x[2]);
diff --git a/src/syntax/simple-statements.spec.ts b/src/syntax/simple-statements.spec.ts
index 25dfcd8..7510108 100644
--- a/src/syntax/simple-statements.spec.ts
+++ b/src/syntax/simple-statements.spec.ts
@@ -233,6 +233,15 @@ describe('Simple statements', () => {
         }
     });
 
+    checkStatement(`COMMENT ON COLUMN my_database.public.groups.members is 'some text'`, {
+        type: 'comment',
+        comment: 'some text',
+        on: {
+            type: 'column',
+            column: { database: 'my_database', schema: 'public', table: 'groups', column: 'members', }
+        }
+    });
+
 
 
     it('can fetch comments', () => {
@@ -324,4 +333,27 @@ describe('Simple statements', () => {
         columns: [{ expr: ref('something'), alias: { name: 'column' } }],
         from: [tbl('whatever')],
     });
+
+    checkStatement('SELECT my_database.my_schema.my_table.my_column FROM my_database.my_schema.my_table', {
+        type: 'select',
+        columns: [{
+            expr: {
+                type: 'ref',
+                table: {
+                    database: 'my_database',
+                    schema: 'my_schema',
+                    name: 'my_table',
+                },
+                name: 'my_column',
+            },
+        }],
+        from: [{
+            type: 'table',
+            name: {
+                database: 'my_database',
+                schema: 'my_schema',
+                name: 'my_table',
+            },
+        }],
+    });
 });
diff --git a/src/to-sql.ts b/src/to-sql.ts
index 158e7a4..6cacc29 100644
--- a/src/to-sql.ts
+++ b/src/to-sql.ts
@@ -102,6 +102,12 @@ function addConstraint(c: ColumnConstraint | TableConstraint, m: IAstVisitor) {
     ret.push(' ');
 }
 function visitQualifiedName(cs: QName, forceDoubleQuote?: boolean) {
+    if (cs.database !== undefined && cs.schema === undefined) {
+        throw new Error('Cannot qualify by database without also qualifying by schema')
+    }
+    if (cs.database) {
+        ret.push(ident(cs.database), '.');
+    }
     if (cs.schema) {
         ret.push(ident(cs.schema), '.');
     }
@@ -218,6 +224,12 @@ function visitSeqOpts(m: IAstVisitor, cs: AlterSequenceSetOptions | CreateSequen
 }
 
 function visitQColumn(col: QColumn) {
+    if (col.database !== undefined && col.schema === undefined) {
+        throw new Error('Cannot qualify by database without also qualifying by schema')
+    }
+    if (col.database) {
+        ret.push(ident(col.database), '.');
+    }
     if (col.schema) {
         ret.push(ident(col.schema), '.');
     }