From 64fc0855526a61d915f4619e927c6a32868c7cee Mon Sep 17 00:00:00 2001
From: Arjun Lall <lallarjun@gmail.com>
Date: Mon, 16 Dec 2024 17:23:18 -0500
Subject: [PATCH] Add extended pg_database columns for compatibility

---
 src/query_handler_test.go    |   4 ++
 src/query_parser_table.go    | 104 +++++++++++++++++++++++++++++++++++
 src/select_remapper_table.go |   5 ++
 3 files changed, 113 insertions(+)

diff --git a/src/query_handler_test.go b/src/query_handler_test.go
index 50127a8..1185015 100644
--- a/src/query_handler_test.go
+++ b/src/query_handler_test.go
@@ -60,6 +60,10 @@ func TestHandleQuery(t *testing.T) {
 			"description": {"slot_name"},
 			"values":      {},
 		},
+		"SELECT oid, datname, datdba FROM pg_catalog.pg_database where oid = 1146": {
+			"description": {"oid", "datname", "datdba"},
+			"values":      {"1146", "memory", ""},
+		},
 		// pg_namespace
 		"SELECT DISTINCT(nspname) FROM pg_catalog.pg_namespace WHERE nspname != 'information_schema' AND nspname != 'pg_catalog'": {
 			"description": {"nspname"},
diff --git a/src/query_parser_table.go b/src/query_parser_table.go
index bb28f7b..054ef99 100644
--- a/src/query_parser_table.go
+++ b/src/query_parser_table.go
@@ -90,6 +90,110 @@ func (parser *QueryParserTable) MakePgExtensionNode(alias string) *pgQuery.Node
     return parser.utils.MakeSubselectWithRowsNode(columns, rowsValues, alias)
 }
 
+// pg_catalog.pg_database -> VALUES(values...) t(columns...)
+func (parser *QueryParserTable) MakePgDatabaseNode(alias string) *pgQuery.Node {
+    targetList := []*pgQuery.Node{
+        pgQuery.MakeResTargetNodeWithNameAndVal(
+            "oid",
+            pgQuery.MakeColumnRefNode([]*pgQuery.Node{pgQuery.MakeStrNode("oid")}, 0),
+            0,
+        ),
+        pgQuery.MakeResTargetNodeWithNameAndVal(
+            "datname",
+            pgQuery.MakeColumnRefNode([]*pgQuery.Node{pgQuery.MakeStrNode("datname")}, 0),
+            0,
+        ),
+        pgQuery.MakeResTargetNodeWithNameAndVal(
+            "datdba",
+            pgQuery.MakeAConstStrNode("", 0),
+            0,
+        ),
+        pgQuery.MakeResTargetNodeWithNameAndVal(
+            "encoding",
+            pgQuery.MakeAConstStrNode("6", 0),
+            0,
+        ),
+        pgQuery.MakeResTargetNodeWithNameAndVal(
+            "datlocprovider",
+            pgQuery.MakeAConstStrNode("c", 0),
+            0,
+        ),
+        pgQuery.MakeResTargetNodeWithNameAndVal(
+            "datistemplate",
+            pgQuery.MakeAConstStrNode("f", 0),
+            0,
+        ),
+        pgQuery.MakeResTargetNodeWithNameAndVal(
+            "datallowconn",
+            pgQuery.MakeAConstStrNode("t", 0),
+            0,
+        ),
+        pgQuery.MakeResTargetNodeWithNameAndVal(
+            "datconnlimit",
+            pgQuery.MakeAConstStrNode("-1", 0),
+            0,
+        ),
+        pgQuery.MakeResTargetNodeWithNameAndVal(
+            "datfrozenxid",
+            pgQuery.MakeAConstStrNode("722", 0),
+            0,
+        ),
+        pgQuery.MakeResTargetNodeWithNameAndVal(
+            "datminmxid",
+            pgQuery.MakeAConstStrNode("1", 0),
+            0,
+        ),
+        pgQuery.MakeResTargetNodeWithNameAndVal(
+            "dattablespace",
+            pgQuery.MakeAConstStrNode("1663", 0),
+            0,
+        ),
+        pgQuery.MakeResTargetNodeWithNameAndVal(
+            "datcollate",
+            pgQuery.MakeAConstStrNode("en_US.UTF-8", 0),
+            0,
+        ),
+        pgQuery.MakeResTargetNodeWithNameAndVal(
+            "datctype",
+            pgQuery.MakeAConstStrNode("en_US.UTF-8", 0),
+            0,
+        ),
+        pgQuery.MakeResTargetNodeWithNameAndVal(
+            "daticulocale",
+            pgQuery.MakeAConstStrNode("", 0),
+            0,
+        ),
+        pgQuery.MakeResTargetNodeWithNameAndVal(
+            "daticurules",
+            pgQuery.MakeAConstStrNode("", 0),
+            0,
+        ),
+        pgQuery.MakeResTargetNodeWithNameAndVal(
+            "datcollversion",
+            pgQuery.MakeAConstStrNode("", 0),
+            0,
+        ),
+        pgQuery.MakeResTargetNodeWithNameAndVal(
+            "datacl",
+            pgQuery.MakeAConstStrNode("", 0),
+            0,
+        ),
+    }
+
+    fromClause := &pgQuery.Node{
+        Node: &pgQuery.Node_RangeVar{
+            RangeVar: &pgQuery.RangeVar{
+                Schemaname:     "pg_catalog",
+                Relname:        "pg_database",
+                Inh:           true,
+                Relpersistence: "p",
+            },
+        },
+    }
+
+    return parser.utils.MakeSubselectFromNode(targetList, fromClause, alias)
+}
+
 // System pg_* tables
 func (parser *QueryParserTable) IsTableFromPgCatalog(qSchemaTable QuerySchemaTable) bool {
 	return parser.isPgCatalogSchema(qSchemaTable) &&
diff --git a/src/select_remapper_table.go b/src/select_remapper_table.go
index ab84d76..f253b2e 100644
--- a/src/select_remapper_table.go
+++ b/src/select_remapper_table.go
@@ -18,6 +18,7 @@ const (
 	PG_TABLE_PG_CLASS              = "pg_class"
 	PG_TABLE_PG_EXTENSION          = "pg_extension"
 	PG_TABLE_PG_REPLICATION_SLOTS  = "pg_replication_slots"
+	PG_TABLE_PG_DATABASE           = "pg_database"
 
 	PG_TABLE_TABLES = "tables"
 )
@@ -81,6 +82,10 @@ func (remapper *SelectRemapperTable) RemapTable(node *pgQuery.Node) *pgQuery.Nod
 			// pg_replication_slots -> return nothing
 			tableNode := parser.MakeEmptyTableNode(PG_REPLICATION_SLOTS_COLUMNS, qSchemaTable.Alias)
 			return remapper.overrideTable(node, tableNode)
+		case PG_TABLE_PG_DATABASE:
+			// pg_catalog.pg_database -> extend with additional Postgres columns
+			tableNode := parser.MakePgDatabaseNode(qSchemaTable.Alias)
+			return remapper.overrideTable(node, tableNode)
 		default:
 			// pg_catalog.pg_* other system tables -> return as is
 			return node