Add pg_roles support

BemiHQ · Dec 11, 2024 · e576386 · e576386
1 parent fceb736
commit e576386
Show file tree

Hide file tree

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

diff --git a/src/query_handler_test.go b/src/query_handler_test.go
@@ -52,6 +52,11 @@ func TestHandleQuery(t *testing.T) {
 			"description": {"total_size"},
 			"values":      {},
 		},
+		// pg_roles
+		"SELECT * FROM pg_catalog.pg_roles": {
+			"description": {"oid", "rolname", "rolsuper", "rolinherit", "rolcreaterole", "rolcreatedb", "rolcanlogin", "rolreplication", "rolconnlimit", "rolpassword", "rolvaliduntil", "rolbypassrls", "rolconfig"},
+			"values":      {"10", "bemidb", "true", "true", "true", "true", "true", "false", "-1", "NULL", "NULL", "false", "NULL"},
+		},
 		// Information schema
 		"SELECT * FROM information_schema.tables": {
 			"description": {"table_catalog", "table_schema", "table_name", "table_type", "self_referencing_column_name", "reference_generation", "user_defined_type_catalog", "user_defined_type_schema", "user_defined_type_name", "is_insertable_into", "is_typed", "commit_action"},

diff --git a/src/query_parser.go b/src/query_parser.go
@@ -9,6 +9,7 @@ const (
 	PG_TABLE_PG_STATIO_USER_TABLES = "pg_statio_user_tables"
 	PG_TABLE_PG_SHADOW             = "pg_shadow"
 	PG_TABLE_PG_NAMESPACE          = "pg_namespace"
+	PG_TABLE_PG_ROLES              = "pg_roles"
 
 	PG_SCHEMA_INFORMATION_SCHEMA = "information_schema"
 	PG_TABLE_TABLES              = "tables"
@@ -172,6 +173,22 @@ var PG_SHADOW_VALUE_BY_COLUMN = NewOrderedMap([][]string{
 	{"useconfig", "NULL"},
 })
 
+var PG_ROLES_VALUE_BY_COLUMN = NewOrderedMap([][]string{
+	{"oid", "10"},
+	{"rolname", ""},
+	{"rolsuper", "true"},
+	{"rolinherit", "true"},
+	{"rolcreaterole", "true"},
+	{"rolcreatedb", "true"},
+	{"rolcanlogin", "true"},
+	{"rolreplication", "false"},
+	{"rolconnlimit", "-1"},
+	{"rolpassword", "NULL"},
+	{"rolvaliduntil", "NULL"},
+	{"rolbypassrls", "false"},
+	{"rolconfig", "NULL"},
+})
+
 type DuckDBKeyword struct {
 	word     string
 	category string
@@ -719,6 +736,32 @@ func (queryParser *QueryParser) MakePgShadowNode(user string, encryptedPassword
 
 ////////////////////////////////////////////////////////////////////////////////
 
+// pg_catalog.pg_roles
+func (queryParser *QueryParser) IsPgRolesTable(schemaTable SchemaTable) bool {
+	return queryParser.isPgCatalogSchema(schemaTable) && schemaTable.Table == PG_TABLE_PG_ROLES
+}
+
+// pg_catalog.pg_roles -> VALUES(values...) t(columns...)
+func (queryParser *QueryParser) MakePgRolesNode(user string) *pgQuery.Node {
+	columns := PG_ROLES_VALUE_BY_COLUMN.Keys()
+	staticRowValues := PG_ROLES_VALUE_BY_COLUMN.Values()
+
+	var rowsValues [][]string
+	rowValues := make([]string, len(staticRowValues))
+	copy(rowValues, staticRowValues)
+
+	for i, column := range columns {
+		if column == "rolname" {
+			rowValues[i] = user
+		}
+	}
+	rowsValues = append(rowsValues, rowValues)
+
+	return queryParser.makeSubselectNode(columns, rowsValues)
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
 // pg_catalog.pg_namespace
 func (queryParser *QueryParser) IsPgNamespaceTable(schemaTable SchemaTable) bool {
 	return queryParser.isPgCatalogSchema(schemaTable) && schemaTable.Table == PG_TABLE_PG_NAMESPACE

diff --git a/src/select_table_remapper.go b/src/select_table_remapper.go
@@ -44,6 +44,12 @@ func (remapper *SelectTableRemapper) RemapTable(node *pgQuery.Node) *pgQuery.Nod
 		return remapper.overrideTable(node, tableNode)
 	}
 
+	// pg_catalog.pg_roles -> return hard-coded role info
+	if parser.IsPgRolesTable(schemaTable) {
+		tableNode := parser.MakePgRolesNode(remapper.config.User)
+		return remapper.overrideTable(node, tableNode)
+	}
+
 	// pg_catalog.pg_* other system tables
 	if parser.IsTableFromPgCatalog(schemaTable) {
 		return node