chore: address review comments

Author: olavloite

parser/simple_parser.go

@@ -15,6 +15,7 @@
 package parser
 
 import (
+	"bytes"
 	"strings"
 	"unicode"
 	"unicode/utf8"
@@ -402,7 +403,7 @@ func (p *simpleParser) skipStatementHint() (bool, int) {
 		// Skip all other whitespaces and comments, but not comments that contain a PG hint.
 		p.skipWhitespacesAndCommentsWithPgHintOption( /*skipPgHints=*/ false)
 		// Check if the next tokens are a PG hint.
-		if len(p.sql) > p.pos+2 && p.sql[p.pos] == '/' && p.sql[p.pos+1] == '*' && p.sql[p.pos+2] == '@' {
+		if bytes.HasPrefix(p.sql[p.pos:], postgreSqlStatementHintPrefix) {
 			startPos := p.pos
 			// Move to the end of this comment.
 			p.pos = p.statementParser.skipMultiLineComment(p.sql, p.pos)

parser/statement_parser_test.go

@@ -2704,11 +2704,7 @@ func TestEatIdentifier(t *testing.T) {
 func TestExtractSetStatementsFromHints(t *testing.T) {
 	t.Parallel()
 
-	parser, err := NewStatementParser(databasepb.DatabaseDialect_GOOGLE_STANDARD_SQL, 1000)
-	if err != nil {
-		t.Fatal(err)
-	}
-	tests := []struct {
+	googleSQLTests := []struct {
 		input   string
 		want    *ParsedSetStatement
 		wantErr bool
@@ -2766,6 +2762,10 @@ func TestExtractSetStatementsFromHints(t *testing.T) {
 			input: "select * from my_table",
 			want:  nil,
 		},
+		{
+			input: "@{} select * from my_table",
+			want:  nil,
+		},
 		{
 			input:   "@{ foo is 'bar' } select * from my_table",
 			wantErr: true,
@@ -2809,34 +2809,7 @@ func TestExtractSetStatementsFromHints(t *testing.T) {
 			wantErr: true,
 		},
 	}
-	for _, test := range tests {
-		t.Run(test.input, func(t *testing.T) {
-			statement, err := parser.extractSetStatementsFromHints(test.input)
-			if test.wantErr {
-				if err == nil {
-					t.Fatal("missing expected error")
-				}
-			} else {
-				if err != nil {
-					t.Fatal(err)
-				}
-				opts := cmpopts.IgnoreUnexported(ParsedSetStatement{})
-				if !cmp.Equal(statement, test.want, opts) {
-					t.Fatalf("mismatch (-want +got):\n%s", cmp.Diff(test.want, statement, opts))
-				}
-			}
-		})
-	}
-}
-
-func TestExtractSetStatementsFromHintsPostgreSQL(t *testing.T) {
-	t.Parallel()
-
-	parser, err := NewStatementParser(databasepb.DatabaseDialect_POSTGRESQL, 1000)
-	if err != nil {
-		t.Fatal(err)
-	}
-	tests := []struct {
+	pgTests := []struct {
 		input   string
 		want    *ParsedSetStatement
 		wantErr bool
@@ -2894,6 +2867,10 @@ func TestExtractSetStatementsFromHintsPostgreSQL(t *testing.T) {
 			input: "select * from my_table",
 			want:  nil,
 		},
+		{
+			input: "/*@*/select * from my_table",
+			want:  nil,
+		},
 		{
 			input:   "/*@ foo is 'bar' */ select * from my_table",
 			wantErr: true,
@@ -2937,24 +2914,44 @@ func TestExtractSetStatementsFromHintsPostgreSQL(t *testing.T) {
 			wantErr: true,
 		},
 	}
-	for _, test := range tests {
-		t.Run(test.input, func(t *testing.T) {
-			statement, err := parser.extractSetStatementsFromHints(test.input)
-			if test.wantErr {
-				if err == nil {
-					t.Fatal("missing expected error")
-				}
-			} else {
-				if err != nil {
-					t.Fatal(err)
-				}
-				opts := cmpopts.IgnoreUnexported(ParsedSetStatement{})
-				if !cmp.Equal(statement, test.want, opts) {
-					t.Fatalf("mismatch (-want +got):\n%s", cmp.Diff(test.want, statement, opts))
+
+	runHintTests := func(t *testing.T, dialect databasepb.DatabaseDialect, tests []struct {
+		input   string
+		want    *ParsedSetStatement
+		wantErr bool
+	}) {
+		parser, err := NewStatementParser(dialect, 1000)
+		if err != nil {
+			t.Fatal(err)
+		}
+		for _, test := range tests {
+			t.Run(test.input, func(t *testing.T) {
+				statement, err := parser.extractSetStatementsFromHints(test.input)
+				if test.wantErr {
+					if err == nil {
+						t.Fatal("missing expected error")
+					}
+				} else {
+					if err != nil {
+						t.Fatal(err)
+					}
+					opts := cmpopts.IgnoreUnexported(ParsedSetStatement{})
+					if !cmp.Equal(statement, test.want, opts) {
+						t.Fatalf("mismatch (-want +got):\n%s", cmp.Diff(test.want, statement, opts))
+					}
 				}
-			}
-		})
+			})
+		}
 	}
+
+	t.Run("GoogleSQL", func(t *testing.T) {
+		t.Parallel()
+		runHintTests(t, databasepb.DatabaseDialect_GOOGLE_STANDARD_SQL, googleSQLTests)
+	})
+	t.Run("PostgreSQL", func(t *testing.T) {
+		t.Parallel()
+		runHintTests(t, databasepb.DatabaseDialect_POSTGRESQL, pgTests)
+	})
 }
 
 func TestSplit(t *testing.T) {