From e576386d381f516dfcc8cce19892c9ac0caf0994 Mon Sep 17 00:00:00 2001
From: Arjun Lall <lallarjun@gmail.com>
Date: Tue, 10 Dec 2024 17:54:18 -0800
Subject: [PATCH] Add pg_roles support

---
 src/main.go                  |  2 +-
 src/query_handler_test.go    |  5 +++++
 src/query_parser.go          | 43 ++++++++++++++++++++++++++++++++++++
 src/select_table_remapper.go |  6 +++++
 4 files changed, 55 insertions(+), 1 deletion(-)

diff --git a/src/main.go b/src/main.go
index 4888568..ce352c7 100644
--- 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
index 931a824..dc15e99 100644
--- 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
index 1421fc2..b13cd5f 100644
--- 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
index 9b6a9d3..5446bfe 100644
--- 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