From ed640f72755aaf520564f3c118e47cfcc5f2153d Mon Sep 17 00:00:00 2001
From: Arjun Lall <lallarjun@gmail.com>
Date: Fri, 13 Dec 2024 09:45:33 -0800
Subject: [PATCH] Add set_config() support

---
 src/query_handler_test.go     |  4 ++++
 src/query_parser_select.go    | 16 ++++++++++++++++
 src/select_remapper_select.go |  5 +++++
 3 files changed, 25 insertions(+)

diff --git a/src/query_handler_test.go b/src/query_handler_test.go
index 483f1bd..bf8c3d3 100644
--- a/src/query_handler_test.go
+++ b/src/query_handler_test.go
@@ -39,6 +39,10 @@ func TestHandleQuery(t *testing.T) {
 		"SELECT pg_catalog.pg_get_expr(adbin, drelid, TRUE) AS def_value FROM pg_catalog.pg_attrdef": {
 			"description": {"def_value"},
 		},
+		"SELECT set_config('bytea_output', 'hex', false)": {
+			"description": {"set_config"},
+			"values":      {"hex"},
+		},
 		// PG system tables
 		"SELECT oid, typname AS typename FROM pg_type WHERE typname='geometry' OR typname='geography'": {
 			"description": {"oid", "typename"},
diff --git a/src/query_parser_select.go b/src/query_parser_select.go
index 30ffbbf..0ee0f03 100644
--- a/src/query_parser_select.go
+++ b/src/query_parser_select.go
@@ -7,6 +7,7 @@ import (
 const (
 	PG_FUNCTION_QUOTE_INDENT = "quote_ident"
 	PG_FUNCTION_PG_GET_EXPR  = "pg_get_expr"
+	PG_FUNCTION_SET_CONFIG   = "set_config"
 )
 
 type QueryParserSelect struct {
@@ -68,6 +69,21 @@ func (parser *QueryParserSelect) RemoveThirdArgumentFromPgGetExpr(functionCall *
 	return functionCall
 }
 
+// set_config()
+func (parser *QueryParserSelect) IsSetConfigFunction(functionName string) bool {
+	return functionName == PG_FUNCTION_SET_CONFIG
+}
+
+// set_config(setting_name, new_value, is_local) -> new_value
+func (parser *QueryParserSelect) RemapSetConfigFunction(targetNode *pgQuery.Node, functionCall *pgQuery.FuncCall) {
+	valueNode := functionCall.Args[1]
+	settingName := functionCall.Args[0].GetAConst().GetSval().Sval
+	LogWarn(parser.config, "Unsupported set_config", settingName, ":", functionCall)
+
+	parser.OverrideTargetValue(targetNode, valueNode)
+	parser.SetDefaultTargetName(targetNode, PG_FUNCTION_SET_CONFIG)
+}
+
 func (parser *QueryParserSelect) OverrideFunctionCallArg(functionCall *pgQuery.FuncCall, index int, node *pgQuery.Node) {
 	functionCall.Args[index] = node
 }
diff --git a/src/select_remapper_select.go b/src/select_remapper_select.go
index 15ac7a6..e8fed76 100644
--- a/src/select_remapper_select.go
+++ b/src/select_remapper_select.go
@@ -36,6 +36,11 @@ func (remapper *SelectRemapperSelect) RemapSelect(targetNode *pgQuery.Node) *pgQ
 
 	originalFunctionName := remapper.parserSelect.FunctionName(functionCall)
 
+	if remapper.parserSelect.IsSetConfigFunction(originalFunctionName) {
+		remapper.parserSelect.RemapSetConfigFunction(targetNode, functionCall)
+		return targetNode
+	}
+
 	renamedNameFunction := remapper.remappedFunctionName(functionCall)
 	if renamedNameFunction != nil {
 		functionCall = renamedNameFunction