From 23a20a410dc92946141b6c6fa5100473eac482cf Mon Sep 17 00:00:00 2001
From: Erik Wienhold <ewie@ewie.name>
Date: Sat, 6 Apr 2024 22:36:54 +0200
Subject: [PATCH v2] plpgsql: accept trailing line comment in CASE WHEN

Expression <expr> in CASE x WHEN <expr> THEN may end with a line
comment.  This results in a syntax error when the WHEN clause is
rewritten to x IN (<expr>).  Fix that by tracking the end location of
the token in read_sql_construct and only copy that range instead of up
to yylloc.

TODO: This breaks other test cases, probably because the calculation of
      endlocation is off.  But I still don't know why.
---
 src/pl/plpgsql/src/expected/plpgsql_control.out | 17 +++++++++++++++++
 src/pl/plpgsql/src/pl_gram.y                    | 10 +++++++++-
 src/pl/plpgsql/src/pl_scanner.c                 |  8 +++++++-
 src/pl/plpgsql/src/plpgsql.h                    |  1 +
 src/pl/plpgsql/src/sql/plpgsql_control.sql      | 14 ++++++++++++++
 5 files changed, 48 insertions(+), 2 deletions(-)

diff --git a/src/pl/plpgsql/src/expected/plpgsql_control.out b/src/pl/plpgsql/src/expected/plpgsql_control.out
index 328bd48586..ccd4f54704 100644
--- a/src/pl/plpgsql/src/expected/plpgsql_control.out
+++ b/src/pl/plpgsql/src/expected/plpgsql_control.out
@@ -681,3 +681,20 @@ select case_test(13);
  other
 (1 row)
 
+-- test line comment between WHEN and THEN
+create or replace function case_comment(int) returns text as $$
+begin
+  case $1
+    when 1 -- comment before THEN
+      then return 'one';
+    else
+      return 'other';
+  end case;
+end;
+$$ language plpgsql immutable;
+select case_comment(1);
+ case_comment 
+--------------
+ one
+(1 row)
+
diff --git a/src/pl/plpgsql/src/pl_gram.y b/src/pl/plpgsql/src/pl_gram.y
index bef33d58a2..b2067c1d2c 100644
--- a/src/pl/plpgsql/src/pl_gram.y
+++ b/src/pl/plpgsql/src/pl_gram.y
@@ -2693,6 +2693,7 @@ read_sql_construct(int until,
 	StringInfoData ds;
 	IdentifierLookup save_IdentifierLookup;
 	int			startlocation = -1;
+	int			endlocation = -1;
 	int			parenlevel = 0;
 	PLpgSQL_expr *expr;
 
@@ -2743,6 +2744,9 @@ read_sql_construct(int until,
 								expected),
 						 parser_errposition(yylloc)));
 		}
+
+		/* TODO endlocation seems to be off, thus breaking tests */
+		endlocation = yylloc + plpgsql_token_length();
 	}
 
 	plpgsql_IdentifierLookup = save_IdentifierLookup;
@@ -2761,7 +2765,11 @@ read_sql_construct(int until,
 			yyerror("missing SQL statement");
 	}
 
-	plpgsql_append_source_text(&ds, startlocation, yylloc);
+	/* TODO probably need to check that we've got proper range */
+	if (startlocation >= endlocation)
+		yyerror("startlocation >= endlocation");
+
+	plpgsql_append_source_text(&ds, startlocation, endlocation);
 
 	/* trim any trailing whitespace, for neatness */
 	if (trim)
diff --git a/src/pl/plpgsql/src/pl_scanner.c b/src/pl/plpgsql/src/pl_scanner.c
index 101804d506..132c9df05b 100644
--- a/src/pl/plpgsql/src/pl_scanner.c
+++ b/src/pl/plpgsql/src/pl_scanner.c
@@ -132,7 +132,6 @@ static int	internal_yylex(TokenAuxData *auxdata);
 static void push_back_token(int token, TokenAuxData *auxdata);
 static void location_lineno_init(void);
 
-
 /*
  * This is the yylex routine called from the PL/pgSQL grammar.
  * It is a wrapper around the core lexer, with the ability to recognize
@@ -396,6 +395,13 @@ plpgsql_token_is_unreserved_keyword(int token)
 	return false;
 }
 
+/* Return the length of the last scanned token. */
+int
+plpgsql_token_length(void)
+{
+	return plpgsql_yyleng;
+}
+
 /*
  * Append the function text starting at startlocation and extending to
  * (not including) endlocation onto the existing contents of "buf".
diff --git a/src/pl/plpgsql/src/plpgsql.h b/src/pl/plpgsql/src/plpgsql.h
index 40056bb851..c48ee062a7 100644
--- a/src/pl/plpgsql/src/plpgsql.h
+++ b/src/pl/plpgsql/src/plpgsql.h
@@ -1319,6 +1319,7 @@ extern int	plpgsql_base_yylex(void);
 extern int	plpgsql_yylex(void);
 extern void plpgsql_push_back_token(int token);
 extern bool plpgsql_token_is_unreserved_keyword(int token);
+extern int	plpgsql_token_length(void);
 extern void plpgsql_append_source_text(StringInfo buf,
 									   int startlocation, int endlocation);
 extern int	plpgsql_peek(void);
diff --git a/src/pl/plpgsql/src/sql/plpgsql_control.sql b/src/pl/plpgsql/src/sql/plpgsql_control.sql
index ed7231134f..8e007c51dc 100644
--- a/src/pl/plpgsql/src/sql/plpgsql_control.sql
+++ b/src/pl/plpgsql/src/sql/plpgsql_control.sql
@@ -486,3 +486,17 @@ select case_test(1);
 select case_test(2);
 select case_test(12);
 select case_test(13);
+
+-- test line comment between WHEN and THEN
+create or replace function case_comment(int) returns text as $$
+begin
+  case $1
+    when 1 -- comment before THEN
+      then return 'one';
+    else
+      return 'other';
+  end case;
+end;
+$$ language plpgsql immutable;
+
+select case_comment(1);
-- 
2.44.0

