Dynamically handle SHOW statements

BemiHQ · Jan 7, 2025 · 2e02f4f · 2e02f4f
1 parent 05099c9
commit 2e02f4f
Show file tree

Hide file tree

Showing 4 changed files with 130 additions and 6 deletions.
diff --git a/src/parser_show.go b/src/parser_show.go
@@ -0,0 +1,91 @@
+package main
+
+import (
+	pgQuery "github.com/pganalyze/pg_query_go/v5"
+)
+
+type ParserShow struct {
+	config *Config
+}
+
+func NewParserShow(config *Config) *ParserShow {
+	return &ParserShow{config: config}
+}
+
+func (parser *ParserShow) VariableName(stmt *pgQuery.RawStmt) string {
+	return stmt.Stmt.GetVariableShowStmt().Name
+}
+
+// SHOW var -> SELECT value AS var FROM duckdb_settings() WHERE LOWER(name) = 'var';
+func (parser *ParserShow) MakeSelectFromDuckdbSettings(variableName string) *pgQuery.RawStmt {
+	return &pgQuery.RawStmt{
+		Stmt: &pgQuery.Node{
+			Node: &pgQuery.Node_SelectStmt{
+				SelectStmt: &pgQuery.SelectStmt{
+					TargetList: []*pgQuery.Node{
+						pgQuery.MakeResTargetNodeWithNameAndVal(
+							variableName,
+							pgQuery.MakeColumnRefNode(
+								[]*pgQuery.Node{pgQuery.MakeStrNode("value")},
+								0,
+							),
+							0,
+						),
+					},
+					FromClause: []*pgQuery.Node{
+						pgQuery.MakeSimpleRangeFunctionNode(
+							[]*pgQuery.Node{
+								pgQuery.MakeListNode(
+									[]*pgQuery.Node{
+										pgQuery.MakeFuncCallNode(
+											[]*pgQuery.Node{pgQuery.MakeStrNode("duckdb_settings")},
+											nil,
+											0,
+										),
+									},
+								),
+							},
+						),
+					},
+					WhereClause: pgQuery.MakeAExprNode(
+						pgQuery.A_Expr_Kind_AEXPR_OP,
+						[]*pgQuery.Node{pgQuery.MakeStrNode("=")},
+						pgQuery.MakeFuncCallNode(
+							[]*pgQuery.Node{pgQuery.MakeStrNode("lower")},
+							[]*pgQuery.Node{
+								pgQuery.MakeColumnRefNode(
+									[]*pgQuery.Node{pgQuery.MakeStrNode("name")},
+									0,
+								),
+							},
+							0,
+						),
+						pgQuery.MakeAConstStrNode(variableName, 0),
+						0,
+					),
+				},
+			},
+		},
+	}
+}
+
+// SHOW search_path -> SELECT CONCAT('"$user", ', value) AS search_path FROM duckdb_settings() WHERE name = 'search_path'
+func (parser *ParserShow) SetTargetListForSearchPath(stmt *pgQuery.RawStmt) {
+	stmt.Stmt.GetSelectStmt().TargetList = []*pgQuery.Node{
+		pgQuery.MakeResTargetNodeWithNameAndVal(
+			"search_path",
+			pgQuery.MakeFuncCallNode(
+				[]*pgQuery.Node{pgQuery.MakeStrNode("concat")},
+				[]*pgQuery.Node{
+					pgQuery.MakeAConstStrNode(`"$user", `, 0),
+					pgQuery.MakeColumnRefNode(
+						[]*pgQuery.Node{pgQuery.MakeStrNode("value")},
+						0,
+					),
+				},
+				0,
+			),
+			0,
+		),
+	}
+}
diff --git a/src/query_handler_test.go b/src/query_handler_test.go
@@ -248,6 +248,10 @@ func TestHandleQuery(t *testing.T) {
 			"types":       {Uint32ToString(pgtype.TextOID)},
 			"values":      {`"$user", public`},
 		},
+		"SHOW timezone": {
+			"description": {"timezone"},
+			"types":       {Uint32ToString(pgtype.TextOID)},
+		},
 
 		// Iceberg data
 		"SELECT COUNT(*) AS count FROM public.test_table": {

diff --git a/src/query_remapper.go b/src/query_remapper.go
@@ -28,6 +28,7 @@ type QueryRemapper struct {
 	remapperTable  *QueryRemapperTable
 	remapperWhere  *QueryRemapperWhere
 	remapperSelect *QueryRemapperSelect
+	remapperShow   *QueryRemapperShow
 	icebergReader  *IcebergReader
 	duckdb         *Duckdb
 	config         *Config
@@ -40,6 +41,7 @@ func NewQueryRemapper(config *Config, icebergReader *IcebergReader, duckdb *Duck
 		remapperTable:  NewQueryRemapperTable(config, icebergReader, duckdb),
 		remapperWhere:  NewQueryRemapperWhere(config),
 		remapperSelect: NewQueryRemapperSelect(config),
+		remapperShow:   NewQueryRemapperShow(config),
 		icebergReader:  icebergReader,
 		duckdb:         duckdb,
 		config:         config,
@@ -75,12 +77,7 @@ func (remapper *QueryRemapper) RemapStatements(statements []*pgQuery.RawStmt) ([
 
 		// SHOW ...
 		case node.GetVariableShowStmt() != nil:
-			if node.GetVariableShowStmt().Name == "search_path" {
-				searchPathStmt, _ := pgQuery.Parse(`SELECT CONCAT('"$user", ', value) AS search_path FROM duckdb_settings() WHERE name = 'search_path'`)
-				statements[i] = searchPathStmt.Stmts[0]
-			} else {
-				statements[i] = FALLBACK_QUERY_TREE.Stmts[0]
-			}
+			statements[i] = remapper.remapperShow.RemapShowStatement(stmt)
 
 		// Unsupported query
 		default:

diff --git a/src/query_remapper_show.go b/src/query_remapper_show.go
@@ -0,0 +1,32 @@
+package main
+
+import (
+	pgQuery "github.com/pganalyze/pg_query_go/v5"
+)
+
+type QueryRemapperShow struct {
+	config     *Config
+	parserShow *ParserShow
+}
+
+func NewQueryRemapperShow(config *Config) *QueryRemapperShow {
+	return &QueryRemapperShow{
+		config:     config,
+		parserShow: NewParserShow(config),
+	}
+}
+
+func (remapper *QueryRemapperShow) RemapShowStatement(stmt *pgQuery.RawStmt) *pgQuery.RawStmt {
+	parser := remapper.parserShow
+	variableName := parser.VariableName(stmt)
+
+	// SHOW var -> SELECT value AS variable FROM duckdb_settings() WHERE LOWER(name) = 'variable';
+	newStmt := parser.MakeSelectFromDuckdbSettings(variableName)
+
+	// SHOW search_path -> SELECT CONCAT('"$user", ', value) AS search_path FROM duckdb_settings() WHERE name = 'search_path'
+	if variableName == "search_path" {
+		parser.SetTargetListForSearchPath(newStmt)
+	}
+
+	return newStmt
+}