From 19ea84c5601730a44852e5ec10932b8922df7ebd Mon Sep 17 00:00:00 2001
From: Jelte Fennema <jelte.fennema@microsoft.com>
Date: Thu, 7 Sep 2023 18:51:30 +0200
Subject: [PATCH v3 3/3] Support changing argument types of prepared statemnts

Prepared statement invalidations might require argument types to change,
this starts supporting that.
---
 src/backend/commands/prepare.c          | 16 +++++--
 src/backend/tcop/postgres.c             | 18 ++++++++
 src/backend/utils/cache/plancache.c     | 57 ++++++++++++++++++++-----
 src/include/utils/plancache.h           | 10 +++++
 src/test/regress/expected/plancache.out | 31 +++++++++++++-
 src/test/regress/sql/plancache.sql      | 13 +++++-
 6 files changed, 130 insertions(+), 15 deletions(-)

diff --git a/src/backend/commands/prepare.c b/src/backend/commands/prepare.c
index da3115beb99..21e5a1d5ff3 100644
--- a/src/backend/commands/prepare.c
+++ b/src/backend/commands/prepare.c
@@ -108,6 +108,10 @@ PrepareQuery(ParseState *pstate, PrepareStmt *stmt,
 
 			argtypes[i++] = toid;
 		}
+
+		plansource->orig_num_params = nargs;
+		plansource->orig_param_types = MemoryContextAlloc(plansource->context, nargs * sizeof(Oid));
+		memcpy(plansource->orig_param_types, argtypes, nargs * sizeof(Oid));
 	}
 
 	/*
@@ -160,10 +164,13 @@ ExecuteQuery(ParseState *pstate,
 	char	   *query_string;
 	int			eflags;
 	long		count;
+	List	   *revalidationResult;
 
 	/* Look it up in the hash table */
 	entry = FetchPreparedStatement(stmt->name, true);
 
+	revalidationResult = RevalidateCachedQuery(entry->plansource, NULL);
+
 	/* Evaluate parameters, if any */
 	if (entry->plansource->num_params > 0)
 	{
@@ -188,7 +195,7 @@ ExecuteQuery(ParseState *pstate,
 									   entry->plansource->query_string);
 
 	/* Replan if needed, and increment plan refcount for portal */
-	cplan = GetCachedPlan(entry->plansource, paramLI, NULL, NULL);
+	cplan = GetCachedPlanFromRevalidated(entry->plansource, paramLI, NULL, NULL, revalidationResult);
 	plan_list = cplan->stmt_list;
 
 	/*
@@ -577,6 +584,7 @@ ExplainExecuteQuery(ExecuteStmt *execstmt, IntoClause *into, ExplainState *es,
 	instr_time	planduration;
 	BufferUsage bufusage_start,
 				bufusage;
+	List	   *revalidationResult;
 
 	if (es->buffers)
 		bufusage_start = pgBufferUsage;
@@ -585,6 +593,8 @@ ExplainExecuteQuery(ExecuteStmt *execstmt, IntoClause *into, ExplainState *es,
 	/* Look it up in the hash table */
 	entry = FetchPreparedStatement(execstmt->name, true);
 
+	revalidationResult = RevalidateCachedQuery(entry->plansource, NULL);
+
 	query_string = entry->plansource->query_string;
 
 	/* Evaluate parameters, if any */
@@ -608,8 +618,8 @@ ExplainExecuteQuery(ExecuteStmt *execstmt, IntoClause *into, ExplainState *es,
 	}
 
 	/* Replan if needed, and acquire a transient refcount */
-	cplan = GetCachedPlan(entry->plansource, paramLI,
-						  CurrentResourceOwner, queryEnv);
+	cplan = GetCachedPlanFromRevalidated(entry->plansource, paramLI,
+										 CurrentResourceOwner, queryEnv, revalidationResult);
 
 	INSTR_TIME_SET_CURRENT(planduration);
 	INSTR_TIME_SUBTRACT(planduration, planstart);
diff --git a/src/backend/tcop/postgres.c b/src/backend/tcop/postgres.c
index 9371a8f6e51..27ca2b0601d 100644
--- a/src/backend/tcop/postgres.c
+++ b/src/backend/tcop/postgres.c
@@ -1535,6 +1535,9 @@ exec_parse_message(const char *query_string,	/* string to execute */
 			snapshot_set = true;
 		}
 
+		psrc->orig_num_params = numParams;
+		psrc->orig_param_types = MemoryContextAlloc(psrc->context, numParams * sizeof(Oid));
+
 		/*
 		 * Analyze and rewrite the query.  Note that the originally specified
 		 * parameter set is not required to be complete, so we have to use
@@ -2639,6 +2642,21 @@ exec_describe_statement_message(const char *stmt_name)
 
 		pstmt = FetchPreparedStatement(stmt_name, true);
 		psrc = pstmt->plansource;
+
+		/*
+		 * Revalidate the cached query, because maybe the argument types
+		 * should be changed due to plan invalidation. This mostly helps for
+		 * cases where the current session did some DDL that invalidated the
+		 * plan or if it changed the search_path.
+		 *
+		 * This approach cannot fully avoid all argument type errors. It's
+		 * still possible that another backend changes one of the tables that
+		 * is used in the prepared statement, in between this backend its
+		 * Describe and Bind calls. This is very unlikely in practice though,
+		 * and at least on the next Describe+Bind iteration the problem would
+		 * not happen.
+		 */
+		RevalidateCachedQuery(psrc, NULL);
 	}
 	else
 	{
diff --git a/src/backend/utils/cache/plancache.c b/src/backend/utils/cache/plancache.c
index 1106215b781..679eff61917 100644
--- a/src/backend/utils/cache/plancache.c
+++ b/src/backend/utils/cache/plancache.c
@@ -101,8 +101,6 @@ static dlist_head saved_plan_list = DLIST_STATIC_INIT(saved_plan_list);
 static dlist_head cached_expression_list = DLIST_STATIC_INIT(cached_expression_list);
 
 static void ReleaseGenericPlan(CachedPlanSource *plansource);
-static List *RevalidateCachedQuery(CachedPlanSource *plansource,
-								   QueryEnvironment *queryEnv);
 static bool CheckCachedPlan(CachedPlanSource *plansource);
 static CachedPlan *BuildCachedPlan(CachedPlanSource *plansource, List *qlist,
 								   ParamListInfo boundParams, QueryEnvironment *queryEnv);
@@ -551,7 +549,7 @@ ReleaseGenericPlan(CachedPlanSource *plansource)
  * had to do re-analysis, and NIL otherwise.  (This is returned just to save
  * a tree copying step in a subsequent BuildCachedPlan call.)
  */
-static List *
+List *
 RevalidateCachedQuery(CachedPlanSource *plansource,
 					  QueryEnvironment *queryEnv)
 {
@@ -688,11 +686,26 @@ RevalidateCachedQuery(CachedPlanSource *plansource,
 											  plansource->parserSetupArg,
 											  queryEnv);
 	else
-		tlist = pg_analyze_and_rewrite_fixedparams(rawtree,
-												   plansource->query_string,
-												   plansource->param_types,
-												   plansource->num_params,
-												   queryEnv);
+	{
+		int			num_params = plansource->orig_num_params;
+		Oid		   *param_types = palloc_array(Oid, num_params);
+
+		memcpy(param_types, plansource->orig_param_types, sizeof(Oid) * num_params);
+
+		/*
+		 * Analyze and rewrite the query.  Note that the originally specified
+		 * parameter set is not required to be complete, so we have to use
+		 * pg_analyze_and_rewrite_varparams(). We cannot reuse the previously
+		 * filled in param_types, because the filled in argument types might
+		 * need to change due to the invalidation.
+		 */
+		tlist = pg_analyze_and_rewrite_varparams(rawtree,
+												 plansource->query_string,
+												 &param_types,
+												 &num_params,
+												 queryEnv);
+		memcpy(plansource->param_types, param_types, sizeof(Oid) * num_params);
+	}
 
 	/* Release snapshot if we got one */
 	if (snapshot_set)
@@ -1136,9 +1149,7 @@ CachedPlan *
 GetCachedPlan(CachedPlanSource *plansource, ParamListInfo boundParams,
 			  ResourceOwner owner, QueryEnvironment *queryEnv)
 {
-	CachedPlan *plan = NULL;
 	List	   *qlist;
-	bool		customplan;
 
 	/* Assert caller is doing things in a sane order */
 	Assert(plansource->magic == CACHEDPLANSOURCE_MAGIC);
@@ -1150,6 +1161,32 @@ GetCachedPlan(CachedPlanSource *plansource, ParamListInfo boundParams,
 	/* Make sure the querytree list is valid and we have parse-time locks */
 	qlist = RevalidateCachedQuery(plansource, queryEnv);
 
+	return GetCachedPlanFromRevalidated(plansource, boundParams, owner, queryEnv, qlist);
+}
+
+/*
+ * GetCachedPlanFromRevalidated: is the same as get GetCachedPlan, but requires
+ * te caller to first revalidate the query. This is needed for callers that
+ * need to use the revalidated plan to generate boundParams.
+ */
+CachedPlan *
+GetCachedPlanFromRevalidated(CachedPlanSource *plansource,
+							 ParamListInfo boundParams,
+							 ResourceOwner owner,
+							 QueryEnvironment *queryEnv,
+							 List *revalidationResult)
+{
+	CachedPlan *plan = NULL;
+	bool		customplan;
+	List	   *qlist = revalidationResult;
+
+	/* Assert caller is doing things in a sane order */
+	Assert(plansource->magic == CACHEDPLANSOURCE_MAGIC);
+	Assert(plansource->is_complete);
+	/* This seems worth a real test, though */
+	if (owner && !plansource->is_saved)
+		elog(ERROR, "cannot apply ResourceOwner to non-saved cached plan");
+
 	/* Decide whether to use a custom plan */
 	customplan = choose_custom_plan(plansource, boundParams);
 
diff --git a/src/include/utils/plancache.h b/src/include/utils/plancache.h
index 36fd0104aec..0e33f620438 100644
--- a/src/include/utils/plancache.h
+++ b/src/include/utils/plancache.h
@@ -101,6 +101,9 @@ typedef struct CachedPlanSource
 	CommandTag	commandTag;		/* 'nuff said */
 	Oid		   *param_types;	/* array of parameter type OIDs, or NULL */
 	int			num_params;		/* length of param_types array */
+	Oid		   *orig_param_types;	/* array of original parameter type OIDs,
+									 * or NULL */
+	int			orig_num_params;	/* length of orig_param_types array */
 	ParserSetupHook parserSetup;	/* alternative parameter spec method */
 	void	   *parserSetupArg;
 	int			cursor_options; /* cursor options used for planning */
@@ -219,6 +222,13 @@ extern CachedPlan *GetCachedPlan(CachedPlanSource *plansource,
 								 ParamListInfo boundParams,
 								 ResourceOwner owner,
 								 QueryEnvironment *queryEnv);
+extern CachedPlan *GetCachedPlanFromRevalidated(CachedPlanSource *plansource,
+												ParamListInfo boundParams,
+												ResourceOwner owner,
+												QueryEnvironment *queryEnv,
+												List *revalidationResult);
+extern List *RevalidateCachedQuery(CachedPlanSource *plansource,
+								   QueryEnvironment *queryEnv);
 extern void ReleaseCachedPlan(CachedPlan *plan, ResourceOwner owner);
 
 extern bool CachedPlanAllowsSimpleValidityCheck(CachedPlanSource *plansource,
diff --git a/src/test/regress/expected/plancache.out b/src/test/regress/expected/plancache.out
index 77ae17b074a..ccbd7c07b3e 100644
--- a/src/test/regress/expected/plancache.out
+++ b/src/test/regress/expected/plancache.out
@@ -15,7 +15,7 @@ EXECUTE prepstmt;
 (5 rows)
 
 -- and one with parameters
-PREPARE prepstmt2(bigint) AS SELECT * FROM pcachetest WHERE q1 = $1;
+PREPARE prepstmt2(unknown) AS SELECT * FROM pcachetest WHERE q1 = $1;
 EXECUTE prepstmt2(123);
  q1  |        q2        
 -----+------------------
@@ -114,6 +114,35 @@ EXECUTE vprep;
  4567890123456789 |  2283945061728394
 (5 rows)
 
+DROP VIEW pcacheview;
+-- If the new plan requires new argument types that should also work resolved
+-- argument types of the prepared statement (because we use unknown as the
+-- type)
+ALTER TABLE pcachetest ALTER COLUMN q1 TYPE text;
+EXECUTE prepstmt2('123');
+ q1  |        q2        
+-----+------------------
+ 123 |              456
+ 123 | 4567890123456789
+(2 rows)
+
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) EXECUTE prepstmt2('123');
+                   QUERY PLAN                   
+------------------------------------------------
+ Seq Scan on pcachetest (actual rows=2 loops=1)
+   Filter: (q1 = $1)
+   Rows Removed by Filter: 3
+(3 rows)
+
+ALTER TABLE pcachetest ALTER COLUMN q1 TYPE bigint USING q1::bigint;
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) EXECUTE prepstmt2(123);
+                   QUERY PLAN                   
+------------------------------------------------
+ Seq Scan on pcachetest (actual rows=2 loops=1)
+   Filter: (q1 = $1)
+   Rows Removed by Filter: 3
+(3 rows)
+
 -- Check basic SPI plan invalidation
 create function cache_test(int) returns int as $$
 declare total int;
diff --git a/src/test/regress/sql/plancache.sql b/src/test/regress/sql/plancache.sql
index cab0c42a142..e4cbdf2bc41 100644
--- a/src/test/regress/sql/plancache.sql
+++ b/src/test/regress/sql/plancache.sql
@@ -10,7 +10,7 @@ PREPARE prepstmt AS SELECT * FROM pcachetest;
 EXECUTE prepstmt;
 
 -- and one with parameters
-PREPARE prepstmt2(bigint) AS SELECT * FROM pcachetest WHERE q1 = $1;
+PREPARE prepstmt2(unknown) AS SELECT * FROM pcachetest WHERE q1 = $1;
 
 EXECUTE prepstmt2(123);
 
@@ -53,6 +53,17 @@ CREATE OR REPLACE TEMP VIEW pcacheview AS
 
 EXECUTE vprep;
 
+DROP VIEW pcacheview;
+-- If the new plan requires new argument types that should also work resolved
+-- argument types of the prepared statement (because we use unknown as the
+-- type)
+ALTER TABLE pcachetest ALTER COLUMN q1 TYPE text;
+EXECUTE prepstmt2('123');
+
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) EXECUTE prepstmt2('123');
+ALTER TABLE pcachetest ALTER COLUMN q1 TYPE bigint USING q1::bigint;
+EXPLAIN (ANALYZE, COSTS OFF, SUMMARY OFF, TIMING OFF) EXECUTE prepstmt2(123);
+
 -- Check basic SPI plan invalidation
 
 create function cache_test(int) returns int as $$
-- 
2.34.1

