v3-0002-handle-plan-nesting-better.patch

diff --git a/contrib/pg_stat_statements/expected/level_tracking.out b/contrib/pg_stat_statements/expected/level_tracking.out
index e7b29e9c99..0b94c71c9c 100644
--- a/contrib/pg_stat_statements/expected/level_tracking.out
+++ b/contrib/pg_stat_statements/expected/level_tracking.out
@@ -173,6 +173,31 @@ SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C";
      1 |    1 | SELECT pg_stat_statements_reset()
 (3 rows)
 
+-- immutable SQL function --- can be executed at plan time
+CREATE FUNCTION PLUS_THREE(i INTEGER) RETURNS INTEGER AS
+$$ SELECT i + 3 LIMIT 1 $$ IMMUTABLE LANGUAGE SQL;
+SELECT PLUS_THREE(8);
+ plus_three 
+------------
+         11
+(1 row)
+
+SELECT PLUS_THREE(10);
+ plus_three 
+------------
+         13
+(1 row)
+
+SELECT toplevel, calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C";
+ toplevel | calls | rows |                                    query                                     
+----------+-------+------+------------------------------------------------------------------------------
+ t        |     2 |    2 | SELECT PLUS_ONE($1)
+ t        |     2 |    2 | SELECT PLUS_THREE($1)
+ t        |     2 |    2 | SELECT PLUS_TWO($1)
+ t        |     1 |    3 | SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C"
+ t        |     1 |    1 | SELECT pg_stat_statements_reset()
+(5 rows)
+
 -- PL/pgSQL function - all-level tracking.
 SET pg_stat_statements.track = 'all';
 SELECT pg_stat_statements_reset();
@@ -184,6 +209,7 @@ SELECT pg_stat_statements_reset();
 -- we drop and recreate the functions to avoid any caching funnies
 DROP FUNCTION PLUS_ONE(INTEGER);
 DROP FUNCTION PLUS_TWO(INTEGER);
+DROP FUNCTION PLUS_THREE(INTEGER);
 -- PL/pgSQL function
 CREATE FUNCTION PLUS_TWO(i INTEGER) RETURNS INTEGER AS $$
 DECLARE
@@ -229,7 +255,34 @@ SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C";
      1 |    1 | SELECT pg_stat_statements_reset()
 (5 rows)
 
-DROP FUNCTION PLUS_ONE(INTEGER);
+-- immutable SQL function --- can be executed at plan time
+CREATE FUNCTION PLUS_THREE(i INTEGER) RETURNS INTEGER AS
+$$ SELECT i + 3 LIMIT 1 $$ IMMUTABLE LANGUAGE SQL;
+SELECT PLUS_THREE(8);
+ plus_three 
+------------
+         11
+(1 row)
+
+SELECT PLUS_THREE(10);
+ plus_three 
+------------
+         13
+(1 row)
+
+SELECT toplevel, calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C";
+ toplevel | calls | rows |                                    query                                     
+----------+-------+------+------------------------------------------------------------------------------
+ f        |     2 |    2 | SELECT (i + $2 + $3)::INTEGER
+ f        |     2 |    2 | SELECT (i + $2)::INTEGER LIMIT $3
+ t        |     2 |    2 | SELECT PLUS_ONE($1)
+ t        |     2 |    2 | SELECT PLUS_THREE($1)
+ t        |     2 |    2 | SELECT PLUS_TWO($1)
+ t        |     1 |    5 | SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C"
+ f        |     2 |    2 | SELECT i + $2 LIMIT $3
+ t        |     1 |    1 | SELECT pg_stat_statements_reset()
+(8 rows)
+
 --
 -- pg_stat_statements.track = none
 --
diff --git a/contrib/pg_stat_statements/pg_stat_statements.c b/contrib/pg_stat_statements/pg_stat_statements.c
index da5ef7d0b7..f1cb02e463 100644
--- a/contrib/pg_stat_statements/pg_stat_statements.c
+++ b/contrib/pg_stat_statements/pg_stat_statements.c
@@ -829,7 +829,7 @@ pgss_post_parse_analyze(ParseState *pstate, Query *query, JumbleState *jstate)
 		prev_post_parse_analyze_hook(pstate, query, jstate);
 
 	/* Safety check... */
-	if (!pgss || !pgss_hash || !pgss_enabled(exec_nested_level))
+	if (!pgss || !pgss_hash || !pgss_enabled(plan_nested_level + exec_nested_level))
 		return;
 
 	/*
@@ -954,12 +954,26 @@ pgss_planner(Query *parse,
 	}
 	else
 	{
-		if (prev_planner_hook)
-			result = prev_planner_hook(parse, query_string, cursorOptions,
-									   boundParams);
-		else
-			result = standard_planner(parse, query_string, cursorOptions,
-									  boundParams);
+		/*
+		 * Even though we're not tracking plan time for this statement, we
+		 * must still increment the nesting level, to ensure that functions
+		 * evaluated during planning are not seen as top-level calls.
+		 */
+		plan_nested_level++;
+		PG_TRY();
+		{
+			if (prev_planner_hook)
+				result = prev_planner_hook(parse, query_string, cursorOptions,
+										   boundParams);
+			else
+				result = standard_planner(parse, query_string, cursorOptions,
+										  boundParams);
+		}
+		PG_FINALLY();
+		{
+			plan_nested_level--;
+		}
+		PG_END_TRY();
 	}
 
 	return result;
@@ -981,7 +995,7 @@ pgss_ExecutorStart(QueryDesc *queryDesc, int eflags)
 	 * counting of optimizable statements that are directly contained in
 	 * utility statements.
 	 */
-	if (pgss_enabled(exec_nested_level) && queryDesc->plannedstmt->queryId != UINT64CONST(0))
+	if (pgss_enabled(plan_nested_level + exec_nested_level) && queryDesc->plannedstmt->queryId != UINT64CONST(0))
 	{
 		/*
 		 * Set up to track total elapsed time in ExecutorRun.  Make sure the
@@ -1051,7 +1065,7 @@ pgss_ExecutorEnd(QueryDesc *queryDesc)
 	uint64		queryId = queryDesc->plannedstmt->queryId;
 
 	if (queryId != UINT64CONST(0) && queryDesc->totaltime &&
-		pgss_enabled(exec_nested_level))
+		pgss_enabled(plan_nested_level + exec_nested_level))
 	{
 		/*
 		 * Make sure stats accumulation is done.  (Note: it's okay if several
@@ -1092,7 +1106,7 @@ pgss_ProcessUtility(PlannedStmt *pstmt, const char *queryString,
 	uint64		saved_queryId = pstmt->queryId;
 	int			saved_stmt_location = pstmt->stmt_location;
 	int			saved_stmt_len = pstmt->stmt_len;
-	bool		enabled = pgss_track_utility && pgss_enabled(exec_nested_level);
+	bool		enabled = pgss_track_utility && pgss_enabled(plan_nested_level + exec_nested_level);
 
 	/*
 	 * Force utility statements to get queryId zero.  We do this even in cases
@@ -1296,7 +1310,7 @@ pgss_store(const char *query, uint64 queryId,
 	key.userid = GetUserId();
 	key.dbid = MyDatabaseId;
 	key.queryid = queryId;
-	key.toplevel = (exec_nested_level == 0);
+	key.toplevel = (plan_nested_level + exec_nested_level == 0);
 
 	/* Lookup the hash table entry with shared lock. */
 	LWLockAcquire(pgss->lock, LW_SHARED);
diff --git a/contrib/pg_stat_statements/sql/level_tracking.sql b/contrib/pg_stat_statements/sql/level_tracking.sql
index eab881a43b..dcd0b04358 100644
--- a/contrib/pg_stat_statements/sql/level_tracking.sql
+++ b/contrib/pg_stat_statements/sql/level_tracking.sql
@@ -90,6 +90,15 @@ SELECT PLUS_ONE(10);
 
 SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C";
 
+-- immutable SQL function --- can be executed at plan time
+CREATE FUNCTION PLUS_THREE(i INTEGER) RETURNS INTEGER AS
+$$ SELECT i + 3 LIMIT 1 $$ IMMUTABLE LANGUAGE SQL;
+
+SELECT PLUS_THREE(8);
+SELECT PLUS_THREE(10);
+
+SELECT toplevel, calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C";
+
 -- PL/pgSQL function - all-level tracking.
 SET pg_stat_statements.track = 'all';
 SELECT pg_stat_statements_reset();
@@ -97,6 +106,7 @@ SELECT pg_stat_statements_reset();
 -- we drop and recreate the functions to avoid any caching funnies
 DROP FUNCTION PLUS_ONE(INTEGER);
 DROP FUNCTION PLUS_TWO(INTEGER);
+DROP FUNCTION PLUS_THREE(INTEGER);
 
 -- PL/pgSQL function
 CREATE FUNCTION PLUS_TWO(i INTEGER) RETURNS INTEGER AS $$
@@ -118,7 +128,15 @@ SELECT PLUS_ONE(3);
 SELECT PLUS_ONE(1);
 
 SELECT calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C";
-DROP FUNCTION PLUS_ONE(INTEGER);
+
+-- immutable SQL function --- can be executed at plan time
+CREATE FUNCTION PLUS_THREE(i INTEGER) RETURNS INTEGER AS
+$$ SELECT i + 3 LIMIT 1 $$ IMMUTABLE LANGUAGE SQL;
+
+SELECT PLUS_THREE(8);
+SELECT PLUS_THREE(10);
+
+SELECT toplevel, calls, rows, query FROM pg_stat_statements ORDER BY query COLLATE "C";
 
 --
 -- pg_stat_statements.track = none