From 198b720ddfa5de5c217e8b96293faa299dba439b Mon Sep 17 00:00:00 2001
From: Anthonin Bonnefoy <anthonin.bonnefoy@datadoghq.com>
Date: Tue, 10 Jun 2025 11:04:58 +0200
Subject: Fix nested select tracking

Currently, we track the start of a select statement using SELECT's
position. However, this doesn't work for statements like '(SELECT 1)
limit 1;'.

This patch instead tracks the start of the select_stmt to the outermost
'(' and we don't track the select's statement length anymore.

The only place where length was necessary was for COPY's
PreparedStatement. This is replaced by directly setting the location and
length of PreparedStatement when CopyStmt is constructed.

This doesn't handle queries like 'CREATE TEMPORARY TABLE pgss_ctas_1
AS (SELECT 1) LIMIT 1 WITH DATA;' as the reported nested query will be
'(SELECT 1) LIMIT 1 WITH DATA'
---
 src/backend/parser/gram.y | 39 +++++++++++++++++++++------------------
 1 file changed, 21 insertions(+), 18 deletions(-)

diff --git a/src/backend/parser/gram.y b/src/backend/parser/gram.y
index 0b5652071d1..72aabc09cfc 100644
--- a/src/backend/parser/gram.y
+++ b/src/backend/parser/gram.y
@@ -154,7 +154,7 @@ static void base_yyerror(YYLTYPE *yylloc, core_yyscan_t yyscanner,
 						 const char *msg);
 static RawStmt *makeRawStmt(Node *stmt, int stmt_location);
 static void updateRawStmtEnd(RawStmt *rs, int end_location);
-static void updatePreparableStmtEnd(Node *n, int end_location);
+static void updatePreparableStmtEnd(Node *n, int location, int len);
 static Node *makeColumnRef(char *colname, List *indirection,
 						   int location, core_yyscan_t yyscanner);
 static Node *makeTypeCast(Node *arg, TypeName *typename, int location);
@@ -3417,7 +3417,7 @@ CopyStmt:	COPY opt_binary qualified_name opt_column_list
 				{
 					CopyStmt *n = makeNode(CopyStmt);
 
-					updatePreparableStmtEnd($3, @4);
+					updatePreparableStmtEnd($3, @2, @4 - @2);
 					n->relation = NULL;
 					n->query = $3;
 					n->attlist = NIL;
@@ -12828,18 +12828,16 @@ select_with_parens:
 			'(' select_no_parens ')'
 				{
 					SelectStmt *n = (SelectStmt *) $2;
-
-					/*
-					 * As SelectStmt's location starts at the SELECT keyword,
-					 * we need to track the length of the SelectStmt within
-					 * parentheses to be able to extract the relevant part
-					 * of the query.  Without this, the RawStmt's length would
-					 * be used and would include the closing parenthesis.
-					 */
-					n->stmt_len = @3 - @2;
+					/* Track the outermost '(' as the select statement start */
+					n->stmt_location = @1;
+					$$ = $2;
+				}
+			| '(' select_with_parens ')'
+				{
+					SelectStmt *n = (SelectStmt *) $2;
+					n->stmt_location = @1;
 					$$ = $2;
 				}
-			| '(' select_with_parens ')'			{ $$ = $2; }
 		;
 
 /*
@@ -18753,37 +18751,42 @@ updateRawStmtEnd(RawStmt *rs, int end_location)
  * string.
  */
 static void
-updatePreparableStmtEnd(Node *n, int end_location)
+updatePreparableStmtEnd(Node *n, int location, int len)
 {
 	if (IsA(n, SelectStmt))
 	{
 		SelectStmt *stmt = (SelectStmt *) n;
 
-		stmt->stmt_len = end_location - stmt->stmt_location;
+		stmt->stmt_location = location;
+		stmt->stmt_len = len;
 	}
 	else if (IsA(n, InsertStmt))
 	{
 		InsertStmt *stmt = (InsertStmt *) n;
 
-		stmt->stmt_len = end_location - stmt->stmt_location;
+		stmt->stmt_location = location;
+		stmt->stmt_len = len;
 	}
 	else if (IsA(n, UpdateStmt))
 	{
 		UpdateStmt *stmt = (UpdateStmt *) n;
 
-		stmt->stmt_len = end_location - stmt->stmt_location;
+		stmt->stmt_location = location;
+		stmt->stmt_len = len;
 	}
 	else if (IsA(n, DeleteStmt))
 	{
 		DeleteStmt *stmt = (DeleteStmt *) n;
 
-		stmt->stmt_len = end_location - stmt->stmt_location;
+		stmt->stmt_location = location;
+		stmt->stmt_len = len;
 	}
 	else if (IsA(n, MergeStmt))
 	{
 		MergeStmt  *stmt = (MergeStmt *) n;
 
-		stmt->stmt_len = end_location - stmt->stmt_location;
+		stmt->stmt_location = location;
+		stmt->stmt_len = len;
 	}
 	else
 		elog(ERROR, "unexpected node type %d", (int) n->type);
-- 
2.49.0