Add row_to_json function support

BemiHQ · Dec 22, 2024 · ed6f81d · ed6f81d
1 parent 3b08b2f
commit ed6f81d
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 0 deletions.
diff --git a/src/query_handler_test.go b/src/query_handler_test.go
@@ -55,6 +55,10 @@ func TestHandleQuery(t *testing.T) {
 			"description": {"pg_is_in_recovery"},
 			"values":      {"f"},
 		},
+		"SELECT row_to_json(t) FROM (SELECT usename, passwd FROM pg_shadow WHERE usename='bemidb') t": {
+			"description": {"row_to_json"},
+			"values":      {`{"usename":"bemidb","passwd":"bemidb-encrypted"}`},
+		},
 		// 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
@@ -8,6 +8,7 @@ const (
 	PG_FUNCTION_QUOTE_INDENT = "quote_ident"
 	PG_FUNCTION_PG_GET_EXPR  = "pg_get_expr"
 	PG_FUNCTION_SET_CONFIG   = "set_config"
+	PG_FUNCTION_ROW_TO_JSON  = "row_to_json"
 )
 
 type QueryParserSelect struct {
@@ -84,6 +85,17 @@ func (parser *QueryParserSelect) RemapSetConfigFunction(targetNode *pgQuery.Node
 	parser.SetDefaultTargetName(targetNode, PG_FUNCTION_SET_CONFIG)
 }
 
+// row_to_json()
+func (parser *QueryParserSelect) IsRowToJsonFunction(functionName string) bool {
+	return functionName == PG_FUNCTION_ROW_TO_JSON
+}
+
+// row_to_json() -> to_json()
+func (parser *QueryParserSelect) RemapRowToJson(functionCall *pgQuery.FuncCall) *pgQuery.FuncCall {
+	functionCall.Funcname = []*pgQuery.Node{pgQuery.MakeStrNode("to_json")}
+	return functionCall
+}
+
 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
@@ -78,6 +78,10 @@ func (remapper *SelectRemapperSelect) SubselectStatement(targetNode *pgQuery.Nod
 func (remapper *SelectRemapperSelect) remappedFunctionName(functionCall *pgQuery.FuncCall) *pgQuery.FuncCall {
 	functionName := remapper.parserSelect.FunctionName(functionCall)
 
+	if remapper.parserSelect.IsRowToJsonFunction(functionName) {
+		return remapper.parserSelect.RemapRowToJson(functionCall)
+	}
+
 	// quote_ident(str) -> concat("\""+str+"\"")
 	if remapper.parserSelect.IsQuoteIdentFunction(functionName) {
 		return remapper.parserSelect.RemapQuoteIdentToConcat(functionCall)