From 5905ee59c1b1173ecc617fafe3e4c2772d1e30d1 Mon Sep 17 00:00:00 2001
From: Jordan Lewis <jordanthelewis@gmail.com>
Date: Fri, 26 Aug 2022 18:06:59 -0400
Subject: [PATCH] parser: don't cut off end-of-line comments

Previously, the parser (actually the scanner) removed comments that came
at the end of a statement from the raw sql that was returned along with
the parsed AST. This behavior is now removed.

Release note: None
---
 pkg/sql/parser/parse.go               | 11 +++++++++--
 pkg/sql/parser/parse_internal_test.go |  6 +++---
 pkg/sql/parser/parse_test.go          | 10 ++++++----
 pkg/sql/show_test.go                  | 15 +++++++++++++++
 4 files changed, 33 insertions(+), 9 deletions(-)

diff --git a/pkg/sql/parser/parse.go b/pkg/sql/parser/parse.go
index cf72531e2854..3a1ab2b5c5d8 100644
--- a/pkg/sql/parser/parse.go
+++ b/pkg/sql/parser/parse.go
@@ -185,7 +185,8 @@ func (p *Parser) scanOneStmt() (sql string, tokens []sqlSymType, done bool) {
 			return p.scanner.In()[startPos:], tokens, true
 		}
 		preValID = lval.id
-		posBeforeScan := p.scanner.Pos()
+		tokens = append(tokens, sqlSymType{})
+		lval = tokens[len(tokens)-1]
 		p.scanner.Scan(&lval)
 
 		if preValID == BEGIN && lval.id == ATOMIC {
@@ -195,7 +196,13 @@ func (p *Parser) scanOneStmt() (sql string, tokens []sqlSymType, done bool) {
 			curFuncBodyCnt--
 		}
 		if lval.id == 0 || (curFuncBodyCnt == 0 && lval.id == ';') {
-			return p.scanner.In()[startPos:posBeforeScan], tokens, (lval.id == 0)
+			endPos := p.scanner.Pos()
+			if lval.id == ';' {
+				// Don't include the ending semicolon, if there is one, in the raw SQL.
+				endPos--
+			}
+			tokens = tokens[:len(tokens)-1]
+			return p.scanner.In()[startPos:endPos], tokens, (lval.id == 0)
 		}
 		lval.pos -= startPos
 		tokens = append(tokens, lval)
diff --git a/pkg/sql/parser/parse_internal_test.go b/pkg/sql/parser/parse_internal_test.go
index e878bdf0673e..6a077e90f6d3 100644
--- a/pkg/sql/parser/parse_internal_test.go
+++ b/pkg/sql/parser/parse_internal_test.go
@@ -44,7 +44,7 @@ func TestScanOneStmt(t *testing.T) {
 		{
 			sql: `SELECT 1 /* comment */  ; /* comment */  ; /* comment */ `,
 			exp: []stmt{{
-				sql: `SELECT 1`,
+				sql: `SELECT 1 /* comment */  `,
 				tok: []int{SELECT, ICONST},
 			}},
 		},
@@ -73,11 +73,11 @@ func TestScanOneStmt(t *testing.T) {
 			sql: ` ; /* x */ SELECT 1  ; SET /* y */ ; ;  INSERT INTO table;  ;`,
 			exp: []stmt{
 				{
-					sql: `SELECT 1`,
+					sql: `SELECT 1  `,
 					tok: []int{SELECT, ICONST},
 				},
 				{
-					sql: `SET`,
+					sql: `SET /* y */ `,
 					tok: []int{SET},
 				},
 				{
diff --git a/pkg/sql/parser/parse_test.go b/pkg/sql/parser/parse_test.go
index d82f5b711e51..6ed2678b7d4e 100644
--- a/pkg/sql/parser/parse_test.go
+++ b/pkg/sql/parser/parse_test.go
@@ -643,10 +643,12 @@ func TestParseSQL(t *testing.T) {
 		{in: ``, exp: nil},
 		{in: `SELECT 1`, exp: []string{`SELECT 1`}},
 		{in: `SELECT 1;`, exp: []string{`SELECT 1`}},
-		{in: `SELECT 1 /* comment */`, exp: []string{`SELECT 1`}},
+		// We currently chop off beginning-of-line comments.
+		{in: `/* comment */ SELECT 1`, exp: []string{`SELECT 1`}},
+		{in: `SELECT 1 /* comment */`, exp: []string{`SELECT 1 /* comment */`}},
 		{in: `SELECT 1;SELECT 2`, exp: []string{`SELECT 1`, `SELECT 2`}},
-		{in: `SELECT 1 /* comment */ ;SELECT 2`, exp: []string{`SELECT 1`, `SELECT 2`}},
-		{in: `SELECT 1 /* comment */ ; /* comment */ SELECT 2`, exp: []string{`SELECT 1`, `SELECT 2`}},
+		{in: `SELECT 1 /* comment */ ;SELECT 2`, exp: []string{`SELECT 1 /* comment */ `, `SELECT 2`}},
+		{in: `SELECT 1 /* comment */ ; SELECT 2`, exp: []string{`SELECT 1 /* comment */ `, `SELECT 2`}},
 	}
 	var p parser.Parser // Verify that the same parser can be reused.
 	for _, d := range testData {
@@ -660,7 +662,7 @@ func TestParseSQL(t *testing.T) {
 				res = append(res, stmts[i].SQL)
 			}
 			if !reflect.DeepEqual(res, d.exp) {
-				t.Errorf("expected \n%v\n, but found %v", res, d.exp)
+				t.Errorf("expected \n%v\n, but found %v", d.exp, res)
 			}
 		})
 	}
diff --git a/pkg/sql/show_test.go b/pkg/sql/show_test.go
index b4115f7272dc..f601492b1783 100644
--- a/pkg/sql/show_test.go
+++ b/pkg/sql/show_test.go
@@ -726,6 +726,21 @@ func TestShowQueriesFillsInValuesForPlaceholders(t *testing.T) {
 			[]interface{}{"hello"},
 			"SELECT upper('hello')",
 		},
+		{
+			"SELECT /* test */ upper($1)",
+			[]interface{}{"hello"},
+			"SELECT upper('hello') /* test */",
+		},
+		{
+			"SELECT /* test */ 'hi'::string",
+			[]interface{}{},
+			"SELECT 'hi'::STRING /* test */",
+		},
+		{
+			"SELECT /* test */ 'hi'::string /* fnord */",
+			[]interface{}{},
+			"SELECT 'hi'::STRING /* test */ /* fnord */",
+		},
 	}
 
 	// Perform both as a simple execution and as a prepared statement,