Add pg_user table support

BemiHQ · Dec 20, 2024 · f714a9c · f714a9c
1 parent 52d7bb8
commit f714a9c
Show file tree

Hide file tree

Showing 4 changed files with 33 additions and 1 deletion.
diff --git a/src/main.go b/src/main.go
@@ -6,7 +6,7 @@ import (
 	"time"
 )
 
-const VERSION = "0.28.1"
+const VERSION = "0.28.2"
 
 func main() {
 	config := LoadConfig()

diff --git a/src/query_handler_test.go b/src/query_handler_test.go
@@ -76,6 +76,10 @@ func TestHandleQuery(t *testing.T) {
 			"description": {"pid", "gss_authenticated", "principal", "encrypted", "credentials_delegated"},
 			"values":      {},
 		},
+		"SELECT * FROM pg_catalog.pg_user": {
+			"description": {"usename", "usesysid", "usecreatedb", "usesuper", "userepl", "usebypassrls", "passwd", "valuntil", "useconfig"},
+			"values":      {"bemidb", "10", "t", "t", "t", "t", "", "NULL", "NULL"},
+		},
 		// 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
@@ -108,6 +108,16 @@ func (parser *QueryParserTable) MakePgDatabaseNode(database string, alias string
 	return parser.utils.MakeSubselectWithRowsNode(PG_TABLE_PG_DATABASE, columns, rowsValues, alias)
 }
 
+// pg_catalog.pg_user -> VALUES(values...) t(columns...)
+func (parser *QueryParserTable) MakePgUserNode(user string, alias string) *pgQuery.Node {
+    columns := PG_USER_VALUE_BY_COLUMN.Keys()
+    rowValues := PG_USER_VALUE_BY_COLUMN.Values()
+
+    rowValues[0] = user
+
+    return parser.utils.MakeSubselectWithRowsNode(PG_TABLE_PG_USER, columns, [][]string{rowValues}, alias)
+}
+
 // System pg_* tables
 func (parser *QueryParserTable) IsTableFromPgCatalog(qSchemaTable QuerySchemaTable) bool {
 	return parser.isPgCatalogSchema(qSchemaTable) &&
@@ -426,6 +436,7 @@ var PG_SYSTEM_TABLES = NewSet([]string{
 	"pg_publication",
 	"pg_publication_namespace",
 	"pg_publication_rel",
+	"pg_user",
 	"pg_range",
 	"pg_replication_origin",
 	"pg_replication_slots",
@@ -559,6 +570,18 @@ var PG_DATABASE_VALUE_BY_COLUMN = NewOrderedMap([][]string{
 	{"datacl", "NULL"},
 })
 
+var PG_USER_VALUE_BY_COLUMN = NewOrderedMap([][]string{
+	{"usename", "bemidb"},
+	{"usesysid", "10"},
+	{"usecreatedb", "t"},
+	{"usesuper", "t"},
+	{"userepl", "t"},
+	{"usebypassrls", "t"},
+	{"passwd", ""},
+	{"valuntil", "NULL"},
+	{"useconfig", "NULL"},
+})
+
 type DuckDBKeyword struct {
 	word     string
 	category string

diff --git a/src/select_remapper_table.go b/src/select_remapper_table.go
@@ -21,6 +21,7 @@ const (
 	PG_TABLE_PG_DATABASE           = "pg_database"
 	PG_TABLE_PG_STAT_GSSAPI        = "pg_stat_gssapi"
 	PG_TABLE_PG_AUTH_MEMBERS       = "pg_auth_members"
+	PG_TABLE_PG_USER               = "pg_user"
 
 	PG_TABLE_TABLES = "tables"
 )
@@ -96,6 +97,10 @@ func (remapper *SelectRemapperTable) RemapTable(node *pgQuery.Node) *pgQuery.Nod
 			// pg_catalog.pg_auth_members -> return empty table
 			tableNode := parser.MakeEmptyTableNode(PG_TABLE_PG_AUTH_MEMBERS, PG_AUTH_MEMBERS_COLUMNS, qSchemaTable.Alias)
 			return remapper.overrideTable(node, tableNode)
+		case PG_TABLE_PG_USER:
+			// pg_catalog.pg_user -> return hard-coded user info
+			tableNode := parser.MakePgUserNode(remapper.config.User, qSchemaTable.Alias)
+			return remapper.overrideTable(node, tableNode)
 		default:
 			// pg_catalog.pg_* other system tables -> return as is
 			return node