From fc594d54db9c760316634cf6eaaa729569e30a4d Mon Sep 17 00:00:00 2001
From: Jordan Lewis <jordanthelewis@gmail.com>
Date: Fri, 23 Sep 2022 16:45:55 -0400
Subject: [PATCH] parser: retain comments

This commit adds a new array to the return value from the parser, which
contains the comments that were in the parsed SQL statement.

Release note: None
---
 pkg/sql/parser/parse.go | 4 ++++
 pkg/sql/scanner/scan.go | 6 ++++++
 2 files changed, 10 insertions(+)

diff --git a/pkg/sql/parser/parse.go b/pkg/sql/parser/parse.go
index 998d93c8a11d..cf72531e2854 100644
--- a/pkg/sql/parser/parse.go
+++ b/pkg/sql/parser/parse.go
@@ -43,6 +43,9 @@ type Statement struct {
 	// AST is the root of the AST tree for the parsed statement.
 	AST tree.Statement
 
+	// Comments is the list of parsed SQL comments.
+	Comments []string
+
 	// SQL is the original SQL from which the statement was parsed. Note that this
 	// is not appropriate for use in logging, as it may contain passwords and
 	// other sensitive data.
@@ -251,6 +254,7 @@ func (p *Parser) parse(
 	return Statement{
 		AST:             p.lexer.stmt,
 		SQL:             sql,
+		Comments:        p.scanner.Comments,
 		NumPlaceholders: p.lexer.numPlaceholders,
 		NumAnnotations:  p.lexer.numAnnotations,
 	}, nil
diff --git a/pkg/sql/scanner/scan.go b/pkg/sql/scanner/scan.go
index ff5d06a0ef67..87a27dd10c2c 100644
--- a/pkg/sql/scanner/scan.go
+++ b/pkg/sql/scanner/scan.go
@@ -57,6 +57,9 @@ type Scanner struct {
 	in            string
 	pos           int
 	bytesPrealloc []byte
+
+	// Comments is the list of parsed comments from the SQL statement.
+	Comments []string
 }
 
 // In returns the input string.
@@ -453,6 +456,7 @@ func (s *Scanner) next() int {
 func (s *Scanner) skipWhitespace(lval ScanSymType, allowComments bool) (newline, ok bool) {
 	newline = false
 	for {
+		startPos := s.pos
 		ch := s.peek()
 		if ch == '\n' {
 			s.pos++
@@ -467,6 +471,8 @@ func (s *Scanner) skipWhitespace(lval ScanSymType, allowComments bool) (newline,
 			if present, cok := s.ScanComment(lval); !cok {
 				return false, false
 			} else if present {
+				// Mark down the comments that we found.
+				s.Comments = append(s.Comments, s.in[startPos:s.pos])
 				continue
 			}
 		}