0001-v1-Jumble-the-SET-command.patch

From ffa3a7cdbb5bce708a59c88a5666e322afbcfc8e Mon Sep 17 00:00:00 2001
From: EC2 Default User <ec2-user@ip-172-31-44-253.ec2.internal>
Date: Wed, 13 Sep 2023 21:06:09 +0000
Subject: [PATCH 1/1] Jumble the SET command

---
 contrib/pg_stat_statements/expected/dml.out   |  2 +-
 .../expected/level_tracking.out               | 20 ++++----
 .../pg_stat_statements/expected/utility.out   | 46 ++++++++++++++-----
 contrib/pg_stat_statements/expected/wal.out   |  2 +-
 contrib/pg_stat_statements/sql/utility.sql    | 17 ++++++-
 src/backend/nodes/queryjumblefuncs.c          | 29 ++++++++++++
 src/include/nodes/parsenodes.h                |  2 +
 7 files changed, 94 insertions(+), 24 deletions(-)

diff --git a/contrib/pg_stat_statements/expected/dml.out b/contrib/pg_stat_statements/expected/dml.out
index 7b9c8f979e..62c6ecd94d 100644
--- a/contrib/pg_stat_statements/expected/dml.out
+++ b/contrib/pg_stat_statements/expected/dml.out
@@ -82,7 +82,7 @@ SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C";
      2 |    4 | SELECT * FROM pgss_dml_tab WHERE a > $1 ORDER BY a
      1 |    8 | SELECT * FROM pgss_dml_tab WHERE a IN ($1, $2, $3, $4, $5)
      1 |    1 | SELECT pg_stat_statements_reset()
-     1 |    0 | SET pg_stat_statements.track_utility = FALSE
+     1 |    0 | SET pg_stat_statements.track_utility = $1
      6 |    6 | UPDATE pgss_dml_tab SET b = $1 WHERE a = $2
      1 |    3 | UPDATE pgss_dml_tab SET b = $1 WHERE a > $2
 (10 rows)
diff --git a/contrib/pg_stat_statements/expected/level_tracking.out b/contrib/pg_stat_statements/expected/level_tracking.out
index d924c87b41..9a67d14a5a 100644
--- a/contrib/pg_stat_statements/expected/level_tracking.out
+++ b/contrib/pg_stat_statements/expected/level_tracking.out
@@ -49,22 +49,22 @@ BEGIN
 END; $$;
 SELECT toplevel, calls, query FROM pg_stat_statements
   ORDER BY query COLLATE "C", toplevel;
- toplevel | calls |                query                 
-----------+-------+--------------------------------------
+ toplevel | calls |               query               
+----------+-------+-----------------------------------
  f        |     1 | DELETE FROM stats_track_tab
  t        |     1 | DELETE FROM stats_track_tab
- t        |     1 | DO $$                               +
-          |       | BEGIN                               +
-          |       |   DELETE FROM stats_track_tab;      +
+ t        |     1 | DO $$                            +
+          |       | BEGIN                            +
+          |       |   DELETE FROM stats_track_tab;   +
           |       | END; $$
- t        |     1 | DO LANGUAGE plpgsql $$              +
-          |       | BEGIN                               +
-          |       |   -- this is a SELECT               +
-          |       |   PERFORM 'hello world'::TEXT;      +
+ t        |     1 | DO LANGUAGE plpgsql $$           +
+          |       | BEGIN                            +
+          |       |   -- this is a SELECT            +
+          |       |   PERFORM 'hello world'::TEXT;   +
           |       | END; $$
  f        |     1 | SELECT $1::TEXT
  t        |     1 | SELECT pg_stat_statements_reset()
- t        |     1 | SET pg_stat_statements.track = 'all'
+ t        |     1 | SET pg_stat_statements.track = $1
 (7 rows)
 
 -- PL/pgSQL function - top-level tracking.
diff --git a/contrib/pg_stat_statements/expected/utility.out b/contrib/pg_stat_statements/expected/utility.out
index f3bc9fcb7c..06169c37d4 100644
--- a/contrib/pg_stat_statements/expected/utility.out
+++ b/contrib/pg_stat_statements/expected/utility.out
@@ -583,11 +583,9 @@ SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C";
      1 |    0 | SET TRANSACTION ISOLATION LEVEL READ COMMITTED
      1 |    0 | SET TRANSACTION ISOLATION LEVEL REPEATABLE READ
      1 |    0 | SET TRANSACTION ISOLATION LEVEL SERIALIZABLE
-     1 |    0 | SET enable_seqscan = off
-     1 |    0 | SET enable_seqscan = on
-     2 |    0 | SET work_mem = '1MB'
-     1 |    0 | SET work_mem = '2MB'
-(15 rows)
+     2 |    0 | SET enable_seqscan = $1
+     3 |    0 | SET work_mem = $1
+(13 rows)
 
 SELECT pg_stat_statements_reset();
  pg_stat_statements_reset 
@@ -665,23 +663,49 @@ SELECT pg_stat_statements_reset();
 -- SET statements.
 -- These use two different strings, still they count as one entry.
 SET work_mem = '1MB';
-Set work_mem = '1MB';
+SET work_mem = '1MB';
 SET work_mem = '2MB';
 RESET work_mem;
 SET enable_seqscan = off;
 SET enable_seqscan = on;
+SET enable_seqscan TO DEFAULT;
+SET enable_seqscan FROM CURRENT;
 RESET enable_seqscan;
+BEGIN;
+SET LOCAL enable_seqscan = off;
+SET LOCAL enable_seqscan TO off;
+SET TIME ZONE LOCAL;
+SET TIME ZONE 'PST8PDT';
+SET LOCAL TIME ZONE LOCAL;
+SET LOCAL TIME ZONE 'PST8PDT';
+END;
+SET search_path = 'public';
+SET search_path = 'public', a;
+SET search_path = 'public', 'a';
+SET search_path = 'public', 'a', 'c';
+RESET search_path;
 SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C";
  calls | rows |               query               
 -------+------+-----------------------------------
+     1 |    0 | BEGIN
+     1 |    0 | END
      1 |    0 | RESET enable_seqscan
+     1 |    0 | RESET search_path
      1 |    0 | RESET work_mem
      1 |    1 | SELECT pg_stat_statements_reset()
-     1 |    0 | SET enable_seqscan = off
-     1 |    0 | SET enable_seqscan = on
-     2 |    0 | SET work_mem = '1MB'
-     1 |    0 | SET work_mem = '2MB'
-(7 rows)
+     1 |    0 | SET LOCAL TIME ZONE $1
+     1 |    0 | SET LOCAL TIME ZONE LOCAL
+     2 |    0 | SET LOCAL enable_seqscan = $1
+     1 |    0 | SET TIME ZONE $1
+     1 |    0 | SET TIME ZONE LOCAL
+     2 |    0 | SET enable_seqscan = $1
+     1 |    0 | SET enable_seqscan FROM CURRENT
+     1 |    0 | SET enable_seqscan TO DEFAULT
+     1 |    0 | SET search_path = $1
+     2 |    0 | SET search_path = $1, $2
+     1 |    0 | SET search_path = $1, $2, $3
+     3 |    0 | SET work_mem = $1
+(18 rows)
 
 SELECT pg_stat_statements_reset();
  pg_stat_statements_reset 
diff --git a/contrib/pg_stat_statements/expected/wal.out b/contrib/pg_stat_statements/expected/wal.out
index 9896ba2536..ce9df00433 100644
--- a/contrib/pg_stat_statements/expected/wal.out
+++ b/contrib/pg_stat_statements/expected/wal.out
@@ -18,7 +18,7 @@ FROM pg_stat_statements ORDER BY query COLLATE "C";
  DELETE FROM pgss_wal_tab WHERE a > $1                        |     1 |    1 | t                   | t                     | t
  INSERT INTO pgss_wal_tab VALUES(generate_series($1, $2), $3) |     1 |   10 | t                   | t                     | t
  SELECT pg_stat_statements_reset()                            |     1 |    1 | f                   | f                     | f
- SET pg_stat_statements.track_utility = FALSE                 |     1 |    0 | f                   | f                     | t
+ SET pg_stat_statements.track_utility = $1                    |     1 |    0 | f                   | f                     | t
  UPDATE pgss_wal_tab SET b = $1 WHERE a > $2                  |     1 |    3 | t                   | t                     | t
 (5 rows)
 
diff --git a/contrib/pg_stat_statements/sql/utility.sql b/contrib/pg_stat_statements/sql/utility.sql
index 28fd13dee0..ad4c296a4e 100644
--- a/contrib/pg_stat_statements/sql/utility.sql
+++ b/contrib/pg_stat_statements/sql/utility.sql
@@ -331,12 +331,27 @@ SELECT pg_stat_statements_reset();
 -- SET statements.
 -- These use two different strings, still they count as one entry.
 SET work_mem = '1MB';
-Set work_mem = '1MB';
+SET work_mem = '1MB';
 SET work_mem = '2MB';
 RESET work_mem;
 SET enable_seqscan = off;
 SET enable_seqscan = on;
+SET enable_seqscan TO DEFAULT;
+SET enable_seqscan FROM CURRENT;
 RESET enable_seqscan;
+BEGIN;
+SET LOCAL enable_seqscan = off;
+SET LOCAL enable_seqscan TO off;
+SET TIME ZONE LOCAL;
+SET TIME ZONE 'PST8PDT';
+SET LOCAL TIME ZONE LOCAL;
+SET LOCAL TIME ZONE 'PST8PDT';
+END;
+SET search_path = 'public';
+SET search_path = 'public', a;
+SET search_path = 'public', 'a';
+SET search_path = 'public', 'a', 'c';
+RESET search_path;
 
 SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C";
 SELECT pg_stat_statements_reset();
diff --git a/src/backend/nodes/queryjumblefuncs.c b/src/backend/nodes/queryjumblefuncs.c
index 281907a4d8..afc2aa44c1 100644
--- a/src/backend/nodes/queryjumblefuncs.c
+++ b/src/backend/nodes/queryjumblefuncs.c
@@ -52,6 +52,7 @@ static void _jumbleNode(JumbleState *jstate, Node *node);
 static void _jumbleA_Const(JumbleState *jstate, Node *node);
 static void _jumbleList(JumbleState *jstate, Node *node);
 static void _jumbleRangeTblEntry(JumbleState *jstate, Node *node);
+static void _jumbleVariableSetStmt(JumbleState *jstate, Node *node);
 
 /*
  * Given a possibly multi-statement source string, confine our attention to the
@@ -395,3 +396,31 @@ _jumbleRangeTblEntry(JumbleState *jstate, Node *node)
 			break;
 	}
 }
+
+static void
+_jumbleVariableSetStmt(JumbleState *jstate, Node *node)
+{
+	VariableSetStmt *expr = (VariableSetStmt *) node;
+	ListCell  *l;
+
+	/* We don't need to normalize SET TRANSACTION */
+	if (expr->kind == VAR_SET_MULTI)
+		JUMBLE_NODE(args);
+	else
+	{
+		int num_args = list_length(expr->args);
+
+		JUMBLE_FIELD(kind);
+		JUMBLE_STRING(name);
+		JUMBLE_FIELD(is_local);
+		JUMBLE_FIELD_SINGLE(num_args);
+
+		foreach	(l, expr->args)
+		{
+			A_Const *ac = (A_Const *) lfirst(l);
+
+			if(ac->type != T_String)
+				RecordConstLocation(jstate, ac->location);
+		}
+	}
+}
diff --git a/src/include/nodes/parsenodes.h b/src/include/nodes/parsenodes.h
index 993b545739..7787bd2439 100644
--- a/src/include/nodes/parsenodes.h
+++ b/src/include/nodes/parsenodes.h
@@ -2478,6 +2478,8 @@ typedef enum VariableSetKind
 
 typedef struct VariableSetStmt
 {
+	pg_node_attr(custom_query_jumble)
+
 	NodeTag		type;
 	VariableSetKind kind;
 	char	   *name;			/* variable to be set */
-- 
2.40.1