Properly escape ALTER FUNCTION .. SET deparsing. Also test

src/backend/distributed/deparser/deparse_function_stmts.c

@@ -58,6 +58,7 @@ static void AppendDefElemCost(StringInfo buf, DefElem *def);
 static void AppendDefElemRows(StringInfo buf, DefElem *def);
 static void AppendDefElemSet(StringInfo buf, DefElem *def);
 
+static void AppendVarSetValue(StringInfo buf, VariableSetStmt *setStmt);
 static void AppendRenameFunctionStmt(StringInfo buf, RenameStmt *stmt);
 static void AppendAlterFunctionSchemaStmt(StringInfo buf, AlterObjectSchemaStmt *stmt);
 static void AppendAlterFunctionOwnerStmt(StringInfo buf, AlterOwnerStmt *stmt);
@@ -257,7 +258,7 @@ AppendDefElemCost(StringInfo buf, DefElem *def)
 static void
 AppendDefElemRows(StringInfo buf, DefElem *def)
 {
-	appendStringInfo(buf, " ROWS  %lf", defGetNumeric(def));
+	appendStringInfo(buf, " ROWS %lf", defGetNumeric(def));
 }
 
 
@@ -268,14 +269,12 @@ static void
 AppendDefElemSet(StringInfo buf, DefElem *def)
 {
 	VariableSetStmt *setStmt = castNode(VariableSetStmt, def->arg);
-	char *setVariableArgs = ExtractSetVariableArgs(setStmt);
 
 	switch (setStmt->kind)
 	{
 		case VAR_SET_VALUE:
 		{
-			appendStringInfo(buf, " SET %s = %s", quote_identifier(setStmt->name),
+			AppendVarSetValue(buf, setStmt);
-							 setVariableArgs);
 			break;
 		}
 
@@ -315,6 +314,114 @@ AppendDefElemSet(StringInfo buf, DefElem *def)
 }
 
 
+/*
+ * AppendVarSetValue deparses a VariableSetStmt with VAR_SET_VALUE_KIND
+ * It takes from flatten_set_variable_args in postgres's utils/misc/guc.c,
+ * however flatten_set_variable_args does not apply correct quoting.
+ */
+static void
+AppendVarSetValue(StringInfo buf, VariableSetStmt *setStmt)
+{
+	ListCell *varArgCell = NULL;
+	ListCell *firstCell = list_head(setStmt->args);
+
+	Assert(setStmt->kind == VAR_SET_VALUE);
+
+	foreach(varArgCell, setStmt->args)
+	{
+		Node *varArgNode = lfirst(varArgCell);
+		A_Const *varArgConst = NULL;
+		TypeName *typeName = NULL;
+
+		if (IsA(varArgNode, A_Const))
+		{
+			varArgConst = (A_Const *) varArgNode;
+		}
+		else if (IsA(varArgNode, TypeCast))
+		{
+			TypeCast *varArgTypeCast = (TypeCast *) varArgNode;
+
+			varArgConst = castNode(A_Const, varArgTypeCast->arg);
+			typeName = varArgTypeCast->typeName;
+		}
+		else
+		{
+			elog(ERROR, "unrecognized node type: %d", varArgNode->type);
+		}
+
+		/* don't know how to start SET until we inspect first arg */
+		if (varArgCell != firstCell)
+		{
+			appendStringInfoChar(buf, ',');
+		}
+		else if (typeName != NULL)
+		{
+			appendStringInfoString(buf, " SET TIME ZONE");
+		}
+		else
+		{
+			appendStringInfo(buf, " SET %s =", quote_identifier(setStmt->name));
+		}
+
+		Value value = varArgConst->val;
+		switch (value.type)
+		{
+			case T_Integer:
+			{
+				appendStringInfo(buf, " %d", intVal(&value));
+				break;
+			}
+
+			case T_Float:
+			{
+				appendStringInfo(buf, " %s", strVal(&value));
+				break;
+			}
+
+			case T_String:
+			{
+				if (typeName != NULL)
+				{
+					/*
+					 * Must be a ConstInterval argument for TIME ZONE. Coerce
+					 * to interval and back to normalize the value and account
+					 * for any typmod.
+					 */
+					Oid typoid = InvalidOid;
+					int32 typmod = -1;
+
+					typenameTypeIdAndMod(NULL, typeName, &typoid, &typmod);
+					Assert(typoid == INTERVALOID);
+
+					Datum interval =
+						DirectFunctionCall3(interval_in,
+											CStringGetDatum(strVal(&value)),
+											ObjectIdGetDatum(InvalidOid),
+											Int32GetDatum(typmod));
+
+					char *intervalout =
+						DatumGetCString(DirectFunctionCall1(interval_out,
+															interval));
+					appendStringInfo(buf, " INTERVAL '%s'", intervalout);
+				}
+				else
+				{
+					appendStringInfo(buf, " %s", quote_literal_cstr(strVal(
+																		&value)));
+				}
+				break;
+			}
+
+			default:
+			{
+				elog(ERROR, "Unexpected Value type in VAR_SET_VALUE arguments.");
+				break;
+			}
+		}
+	}
+}
+
+
 /*
  * DeparseRenameFunctionStmt builds and returns a string representing the RenameStmt
  */

src/test/regress/expected/distributed_functions.out

@@ -318,6 +318,28 @@ SELECT public.verify_function_is_same_on_workers('function_tests.add(int,int)');
  t
 (1 row)
 
+ALTER FUNCTION add(int,int) SET "citus.setting;'" TO 'hello '' world';
+SELECT public.verify_function_is_same_on_workers('function_tests.add(int,int)');
+ verify_function_is_same_on_workers 
+------------------------------------
+ t
+(1 row)
+
+ALTER FUNCTION add(int,int) RESET "citus.setting;'";
+SELECT public.verify_function_is_same_on_workers('function_tests.add(int,int)');
+ verify_function_is_same_on_workers 
+------------------------------------
+ t
+(1 row)
+
+ALTER FUNCTION add(int,int) SET search_path TO 'sch'';ma', public;
+SELECT public.verify_function_is_same_on_workers('function_tests.add(int,int)');
+ verify_function_is_same_on_workers 
+------------------------------------
+ t
+(1 row)
+
+ALTER FUNCTION add(int,int) RESET search_path;
 -- SET ... FROM CURRENT is not supported, verify the query fails with a descriptive error irregardless of where in the action list the statement occurs
 ALTER FUNCTION add(int,int) SET client_min_messages FROM CURRENT;
 ERROR:  unsupported ALTER FUNCTION ... SET ... FROM CURRENT for a distributed function

src/test/regress/expected/multi_deparse_function.out

@@ -3,7 +3,7 @@
 --
 -- This test implements all the possible queries as of Postgres 11
 -- in the order they are listed in the docs
--- 
+--
 -- ALTER FUNCTION name [ ( [ [ argmode ] [ argname ] argtype [, ...] ] ) ]
 --     action [ ... ] [ RESTRICT ]
 -- ALTER FUNCTION name [ ( [ [ argmode ] [ argname ] argtype [, ...] ] ) ]
@@ -14,9 +14,9 @@
 --     SET SCHEMA new_schema
 -- ALTER FUNCTION name [ ( [ [ argmode ] [ argname ] argtype [, ...] ] ) ]
 --     DEPENDS ON EXTENSION extension_name
--- 
+--
 -- where action is one of:
--- 
+--
 --     CALLED ON NULL INPUT | RETURNS NULL ON NULL INPUT | STRICT
 --     IMMUTABLE | STABLE | VOLATILE | [ NOT ] LEAKPROOF
 --     [ EXTERNAL ] SECURITY INVOKER | [ EXTERNAL ] SECURITY DEFINER
@@ -27,7 +27,7 @@
 --     SET configuration_parameter FROM CURRENT
 --     RESET configuration_parameter
 --     RESET ALL
--- 
+--
 -- DROP FUNCTION [ IF EXISTS ] name [ ( [ [ argmode ] [ argname ] argtype [, ...] ] ) ] [, ...]
 --     [ CASCADE | RESTRICT ]
 SET citus.next_shard_id TO 20020000;
@@ -262,7 +262,7 @@ CONTEXT:  PL/pgSQL function deparse_and_run_on_workers(text) line 6 at RAISE
 SELECT deparse_and_run_on_workers($cmd$
 ALTER FUNCTION  add SET log_min_messages = ERROR
 $cmd$);
-INFO:  Propagating deparsed query: ALTER FUNCTION function_tests.add(integer, integer) SET log_min_messages = error;
+INFO:  Propagating deparsed query: ALTER FUNCTION function_tests.add(integer, integer) SET log_min_messages = 'error';
 CONTEXT:  PL/pgSQL function deparse_and_run_on_workers(text) line 6 at RAISE
       deparse_and_run_on_workers      
 --------------------------------------
@@ -292,6 +292,74 @@ CONTEXT:  PL/pgSQL function deparse_and_run_on_workers(text) line 6 at RAISE
  (localhost,57638,t,"ALTER FUNCTION")
 (2 rows)
 
+SELECT deparse_and_run_on_workers($cmd$
+ALTER FUNCTION add(int, int) SET TIME ZONE INTERVAL '-08:00' HOUR TO MINUTE;
+$cmd$);
+INFO:  Propagating deparsed query: ALTER FUNCTION function_tests.add(integer, integer) SET TIME ZONE INTERVAL '@ 8 hours ago';
+CONTEXT:  PL/pgSQL function deparse_and_run_on_workers(text) line 6 at RAISE
+      deparse_and_run_on_workers      
+--------------------------------------
+ (localhost,57637,t,"ALTER FUNCTION")
+ (localhost,57638,t,"ALTER FUNCTION")
+(2 rows)
+
+SELECT deparse_and_run_on_workers($cmd$
+ALTER FUNCTION add(int, int) SET TIME ZONE '-7';
+$cmd$);
+INFO:  Propagating deparsed query: ALTER FUNCTION function_tests.add(integer, integer) SET timezone = '-7';
+CONTEXT:  PL/pgSQL function deparse_and_run_on_workers(text) line 6 at RAISE
+      deparse_and_run_on_workers      
+--------------------------------------
+ (localhost,57637,t,"ALTER FUNCTION")
+ (localhost,57638,t,"ALTER FUNCTION")
+(2 rows)
+
+SELECT deparse_and_run_on_workers($cmd$
+ALTER FUNCTION add(int, int) SET "citus.setting;'" TO 'hello '' world';
+$cmd$);
+INFO:  Propagating deparsed query: ALTER FUNCTION function_tests.add(integer, integer) SET "citus.setting;'" = 'hello '' world';
+CONTEXT:  PL/pgSQL function deparse_and_run_on_workers(text) line 6 at RAISE
+      deparse_and_run_on_workers      
+--------------------------------------
+ (localhost,57637,t,"ALTER FUNCTION")
+ (localhost,57638,t,"ALTER FUNCTION")
+(2 rows)
+
+SELECT deparse_and_run_on_workers($cmd$
+ALTER FUNCTION add(int, int) SET "citus.setting;'" TO -3.2;
+$cmd$);
+INFO:  Propagating deparsed query: ALTER FUNCTION function_tests.add(integer, integer) SET "citus.setting;'" = -3.2;
+CONTEXT:  PL/pgSQL function deparse_and_run_on_workers(text) line 6 at RAISE
+      deparse_and_run_on_workers      
+--------------------------------------
+ (localhost,57637,t,"ALTER FUNCTION")
+ (localhost,57638,t,"ALTER FUNCTION")
+(2 rows)
+
+SELECT deparse_and_run_on_workers($cmd$
+ALTER FUNCTION add(int, int) SET "citus.setting;'" TO -32;
+$cmd$);
+INFO:  Propagating deparsed query: ALTER FUNCTION function_tests.add(integer, integer) SET "citus.setting;'" = -32;
+CONTEXT:  PL/pgSQL function deparse_and_run_on_workers(text) line 6 at RAISE
+      deparse_and_run_on_workers      
+--------------------------------------
+ (localhost,57637,t,"ALTER FUNCTION")
+ (localhost,57638,t,"ALTER FUNCTION")
+(2 rows)
+
+-- This raises an error about only accepting one item,
+-- that's okay, we're just testing that we don't produce bad syntax.
+SELECT deparse_and_run_on_workers($cmd$
+ALTER FUNCTION add(int, int) SET "citus.setting;'" TO 'hello '' world', 'second '' item';
+$cmd$);
+INFO:  Propagating deparsed query: ALTER FUNCTION function_tests.add(integer, integer) SET "citus.setting;'" = 'hello '' world', 'second '' item';
+CONTEXT:  PL/pgSQL function deparse_and_run_on_workers(text) line 6 at RAISE
+                        deparse_and_run_on_workers                         
+---------------------------------------------------------------------------
+ (localhost,57637,f,"ERROR:  SET citus.setting;' takes only one argument")
+ (localhost,57638,f,"ERROR:  SET citus.setting;' takes only one argument")
+(2 rows)
+
 SELECT deparse_and_run_on_workers($cmd$
 ALTER FUNCTION add RESET log_min_messages
 $cmd$);
@@ -614,7 +682,7 @@ CONTEXT:  PL/pgSQL function deparse_and_run_on_workers(text) line 6 at RAISE
  (localhost,57638,t,"ALTER FUNCTION")
 (2 rows)
 
--- a function with variadic input. 
+-- a function with variadic input.
 CREATE FUNCTION sum_avg(
    VARIADIC list NUMERIC[],
    OUT total NUMERIC,
@@ -623,7 +691,7 @@ AS $$
 BEGIN
    SELECT INTO total SUM(list[i])
    FROM generate_subscripts(list, 1) g(i);
- 
+
    SELECT INTO average AVG(list[i])
    FROM generate_subscripts(list, 1) g(i);
 END; $$
@@ -645,7 +713,7 @@ CONTEXT:  PL/pgSQL function deparse_and_run_on_workers(text) line 6 at RAISE
  (localhost,57638,t,"ALTER FUNCTION")
 (2 rows)
 
--- a function with a custom type IN parameter 
+-- a function with a custom type IN parameter
 CREATE TYPE intpair AS (x int, y int);
 CREATE FUNCTION func_custom_param(IN param intpair, OUT total INT)
     AS $$ SELECT param.x + param.y $$

src/test/regress/sql/distributed_functions.sql

@@ -223,6 +223,13 @@ ALTER FUNCTION add(int,int) SET client_min_messages TO debug;
 SELECT public.verify_function_is_same_on_workers('function_tests.add(int,int)');
 ALTER FUNCTION add(int,int) RESET client_min_messages;
 SELECT public.verify_function_is_same_on_workers('function_tests.add(int,int)');
+ALTER FUNCTION add(int,int) SET "citus.setting;'" TO 'hello '' world';
+SELECT public.verify_function_is_same_on_workers('function_tests.add(int,int)');
+ALTER FUNCTION add(int,int) RESET "citus.setting;'";
+SELECT public.verify_function_is_same_on_workers('function_tests.add(int,int)');
+ALTER FUNCTION add(int,int) SET search_path TO 'sch'';ma', public;
+SELECT public.verify_function_is_same_on_workers('function_tests.add(int,int)');
+ALTER FUNCTION add(int,int) RESET search_path;
 
 -- SET ... FROM CURRENT is not supported, verify the query fails with a descriptive error irregardless of where in the action list the statement occurs
 ALTER FUNCTION add(int,int) SET client_min_messages FROM CURRENT;

src/test/regress/sql/multi_deparse_function.sql

@@ -158,6 +158,32 @@ SELECT deparse_and_run_on_workers($cmd$
 ALTER FUNCTION  add SET log_min_messages FROM CURRENT
 $cmd$);
 
+SELECT deparse_and_run_on_workers($cmd$
+ALTER FUNCTION add(int, int) SET TIME ZONE INTERVAL '-08:00' HOUR TO MINUTE;
+$cmd$);
+
+SELECT deparse_and_run_on_workers($cmd$
+ALTER FUNCTION add(int, int) SET TIME ZONE '-7';
+$cmd$);
+
+SELECT deparse_and_run_on_workers($cmd$
+ALTER FUNCTION add(int, int) SET "citus.setting;'" TO 'hello '' world';
+$cmd$);
+
+SELECT deparse_and_run_on_workers($cmd$
+ALTER FUNCTION add(int, int) SET "citus.setting;'" TO -3.2;
+$cmd$);
+
+SELECT deparse_and_run_on_workers($cmd$
+ALTER FUNCTION add(int, int) SET "citus.setting;'" TO -32;
+$cmd$);
+
+-- This raises an error about only accepting one item,
+-- that's okay, we're just testing that we don't produce bad syntax.
+SELECT deparse_and_run_on_workers($cmd$
+ALTER FUNCTION add(int, int) SET "citus.setting;'" TO 'hello '' world', 'second '' item';
+$cmd$);
+
 SELECT deparse_and_run_on_workers($cmd$
 ALTER FUNCTION add RESET log_min_messages
 $cmd$);