Fix quoted tokens and comment handling

parser/error_test.go

@@ -55,6 +55,8 @@ func TestParseError_LexicalFailure(t *testing.T) {
 			msg string
 		}{
 			{"/*", "unclosed multi-line comment"},
+			{"/* outer /* inner */", "unclosed multi-line comment"},
+			{"$$unclosed", "invalid dollar-quoted string"},
 			{"'unclosed", "invalid string"},
 			{"`unclosed", "unclosed quoted identifier"},
 			{"1e+", "exponent part should contain at least one digit"},

parser/lexer.go

@@ -215,12 +215,25 @@ func (l *Lexer) consumeIdent(_ Pos) error {
 			i++
 		}
 	} else {
-		for l.peekOk(i) && (quoteType == BackTicks && l.peekN(i) != '`' ||
-			quoteType == DoubleQuote && l.peekN(i) != '"') {
+		quote := l.input[l.offset-1]
+		for l.peekOk(i) {
+			if l.peekN(i) == '\\' {
+				i++
+				if l.peekOk(i) {
+					i++
+				}
+				continue
+			}
+			if l.peekN(i) == quote {
+				if l.peekOk(i+1) && l.peekN(i+1) == quote {
+					i += 2
+					continue
+				}
+				break
+			}
 			i++
 		}
-		if !l.peekOk(i) || (quoteType == BackTicks && l.peekN(i) != '`') ||
-			(quoteType == DoubleQuote && l.peekN(i) != '"') {
+		if !l.peekOk(i) {
 			return fmt.Errorf("unclosed quoted identifier: %s", l.slice(0, i))
 		}
 	}
@@ -246,7 +259,7 @@ func (l *Lexer) consumeIdent(_ Pos) error {
 func (l *Lexer) consumeSingleLineComment() {
 	l.skipN(2)
 	i := 0
-	for l.peekOk(i) && l.peekN(i) != '\r' && l.peekN(i) != '\n' {
+	for l.peekOk(i) && l.peekN(i) != '\n' {
 		i++
 	}
 	if l.peekOk(i) {
@@ -260,10 +273,21 @@ func (l *Lexer) consumeMultiLineComment() error {
 	pos := Pos(l.offset)
 	l.skipN(2)
 	i := 0
+	depth := 1
 	for l.peekOk(i) {
+		if l.peekOk(i+1) && l.peekN(i) == '/' && l.peekN(i+1) == '*' {
+			depth++
+			i += 2
+			continue
+		}
 		if l.peekOk(i+1) && l.peekN(i) == '*' && l.peekN(i+1) == '/' {
-			l.skipN(i + 2)
-			return nil
+			depth--
+			i += 2
+			if depth == 0 {
+				l.skipN(i)
+				return nil
+			}
+			continue
 		}
 		i++
 	}
@@ -308,13 +332,53 @@ func (l *Lexer) consumeString() error {
 	return nil
 }
 
+func (l *Lexer) consumeDollarQuotedString() error {
+	i := 1
+	for l.peekOk(i) && (IsIdentStart(l.peekN(i)) || IsDigit(l.peekN(i))) {
+		i++
+	}
+	if l.peekOk(i) && l.peekN(i) == '$' {
+		delimiter := l.slice(0, i+1)
+		start := l.offset + len(delimiter)
+		if end := strings.Index(l.input[start:], delimiter); end >= 0 {
+			// StringLiteral stores the escaped interior of a single-quoted
+			// string. Heredoc contents are literal, so escape them once here.
+			literal := strings.ReplaceAll(l.input[start:start+end], "\\", "\\\\")
+			literal = strings.ReplaceAll(literal, "'", "\\'")
+			literal = strings.ReplaceAll(literal, "\n", "\\n")
+			literal = strings.ReplaceAll(literal, "\r", "\\r")
+			l.currentToken = &Token{
+				Kind:   TokenKindString,
+				String: literal,
+				Pos:    Pos(start),
+				End:    Pos(start + end),
+			}
+			l.offset = start + end + len(delimiter)
+			return nil
+		}
+	}
+
+	// ClickHouse treats an unmatched named delimiter as a bare identifier.
+	// A standalone dollar or an unmatched $$ cannot start an identifier.
+	if !l.peekOk(1) || (!IsIdentStart(l.peekN(1)) && !IsDigit(l.peekN(1))) {
+		return errors.New("invalid dollar-quoted string")
+	}
+	return l.consumeIdent(Pos(l.offset))
+}
+
 func (l *Lexer) skipComments() error {
 	for !l.isEOF() {
 		l.skipSpace()
 		if !l.peekOk(0) {
 			return nil
 		}
 		switch l.peekN(0) {
+		case '#':
+			if l.peekOk(1) && (l.peekN(1) == ' ' || l.peekN(1) == '!') {
+				l.consumeSingleLineComment()
+				continue
+			}
+			return nil
 		case '-':
 			if l.peekOk(1) && l.peekN(1) == '-' {
 				l.consumeSingleLineComment()
@@ -417,7 +481,9 @@ func (l *Lexer) consumeToken() (err error) {
 		}
 	case '0', '1', '2', '3', '4', '5', '6', '7', '8', '9':
 		return l.consumeNumber()
-	case '`', '$', '"':
+	case '$':
+		return l.consumeDollarQuotedString()
+	case '`', '"':
 		return l.consumeIdent(Pos(l.offset))
 	case '\'':
 		return l.consumeString()

parser/lexer_test.go

@@ -23,11 +23,20 @@ func TestConsumeComment(t *testing.T) {
 		"/* hello world */ /* hello world */\n",
 		"/* hello world */ /* hello world */\r\n",
 		"/* hello world */ /* hello world */\r",
+		"/* outer /* inner */ outer */",
+		"/* outer /* middle /* inner */ middle */ outer */",
+		"# hello world",
+		"# ",
+		"#!",
+		"#!/usr/bin/clickhouse\n",
+		"# comment\rstill comment",
 	}
 	for _, c := range comments {
 		lexer := NewLexer(c)
 		err := lexer.consumeToken()
 		require.NoError(t, err)
+		require.Nil(t, lexer.currentToken)
+		require.Equal(t, len(c), lexer.offset)
 	}
 
 }
@@ -69,6 +78,7 @@ func TestConsumeUnterminatedComment(t *testing.T) {
 		"/* unterminated",
 		"/* unterminated *",
 		"SELECT 1 /* unterminated",
+		"/* outer /* inner */",
 	}
 	for _, c := range inputs {
 		c := c
@@ -372,3 +382,80 @@ func TestNegativeHexLiteral(t *testing.T) {
 	require.NoError(t, err)
 	require.Len(t, stmts, 1)
 }
+
+func TestConsumeQuotedIdent(t *testing.T) {
+	for _, quote := range []byte{'`', '"'} {
+		for _, content := range []string{
+			"a" + string(quote) + string(quote) + "b",
+			"a\\" + string(quote) + "b",
+			"a\\\\",
+			"中文",
+		} {
+			input := string(quote) + content + string(quote)
+			t.Run(input, func(t *testing.T) {
+				lexer := NewLexer(input + ",")
+				require.NoError(t, lexer.consumeToken())
+				require.Equal(t, TokenKindIdent, lexer.currentToken.Kind)
+				require.Equal(t, content, lexer.currentToken.String)
+				require.Equal(t, Pos(1), lexer.currentToken.Pos)
+				require.Equal(t, Pos(len(input)-1), lexer.currentToken.End)
+				require.NoError(t, lexer.consumeToken())
+				require.Equal(t, TokenKindComma, lexer.currentToken.Kind)
+
+				stmts, err := NewParser("SELECT " + input).ParseStmts()
+				require.NoError(t, err)
+				require.Equal(t, "SELECT "+input, Format(stmts[0]))
+			})
+		}
+	}
+}
+
+func TestConsumeDollarQuotedString(t *testing.T) {
+	for _, tc := range []struct {
+		input, content string
+		pos            Pos
+	}{
+		{"$$hello$$", "hello", 2},
+		{"$$$$", "", 2},
+		{"$tag$it's\\n$tag$", "it\\'s\\\\n", 5},
+		{"$1_$a$$b$1_$", "a$$b", 4},
+		{"$tag$one$TAG$two$tag$", "one$TAG$two", 5},
+		{"$$中文\n/* # */$$", "中文\\n/* # */", 2},
+		{"$$a\r\nb$$", "a\\r\\nb", 2},
+		{"$$'\\$$", "\\'\\\\", 2},
+	} {
+		t.Run(tc.input, func(t *testing.T) {
+			lexer := NewLexer(tc.input + ",")
+			require.NoError(t, lexer.consumeToken())
+			require.Equal(t, TokenKindString, lexer.currentToken.Kind)
+			require.Equal(t, tc.content, lexer.currentToken.String)
+			require.Equal(t, tc.pos, lexer.currentToken.Pos)
+			require.Equal(t, Pos(len(tc.input))-tc.pos, lexer.currentToken.End)
+			require.NoError(t, lexer.consumeToken())
+			require.Equal(t, TokenKindComma, lexer.currentToken.Kind)
+
+			stmts, err := NewParser("SELECT " + tc.input).ParseStmts()
+			require.NoError(t, err)
+			for _, beautify := range []bool{false, true} {
+				formatter := NewFormatter()
+				if beautify {
+					formatter.WithBeautify()
+				}
+				formatter.WriteExpr(stmts[0])
+				formatted := formatter.String()
+				reparsed, err := NewParser(formatted).ParseStmts()
+				require.NoError(t, err)
+				require.Equal(t, "SELECT '"+tc.content+"'", Format(reparsed[0]))
+			}
+		})
+	}
+
+	// Without a complete matching delimiter, a named tag can be an identifier.
+	for _, input := range []string{"$name", "$tag$unclosed", "$tag$x$other$", "name$dollar"} {
+		lexer := NewLexer(input)
+		require.NoError(t, lexer.consumeToken())
+		require.Equal(t, TokenKindIdent, lexer.currentToken.Kind)
+		require.Equal(t, input, lexer.currentToken.String)
+		require.True(t, lexer.isEOF())
+	}
+}

parser/parser_test.go

@@ -192,6 +192,15 @@ func TestParser_InvalidSyntax(t *testing.T) {
 		"SELECT",
 		"SELECT 1; SELECT",
 		"SELECT 1 /*",
+		"SELECT 1 /* outer /* inner */",
+		"SELECT 1 #word",
+		"SELECT 1 #\tcomment",
+		"SELECT 1 #",
+		"SELECT 1 #\n+2",
+		"SELECT $$unclosed",
+		"SELECT $",
+		`SELECT "a\"`,
+		"SELECT `a\\`",
 		"SELECT 1 +> 2",
 		"SELECT * FROM",
 		// WITH FILL error cases

parser/testdata/basic/format/beautify/lexical_fidelity.sql

@@ -0,0 +1,21 @@
+-- Origin SQL:
+#!/usr/bin/clickhouse
+SELECT `a``b`, `a\`b`, "a""b", "a\"b";
+SELECT $$hello$$, $tag$it's\n$tag$, $1_$a$$b$1_$, $$$$;
+SELECT /* outer /* inner */ outer */ 1 # comment
++ 2;
+
+
+-- Beautify SQL:
+SELECT
+  `a``b`,
+  `a\`b`,
+  "a""b",
+  "a\"b";
+SELECT
+  'hello',
+  'it\'s\\n',
+  'a$$b',
+  '';
+SELECT
+  1 + 2;

parser/testdata/basic/format/lexical_fidelity.sql

@@ -0,0 +1,12 @@
+-- Origin SQL:
+#!/usr/bin/clickhouse
+SELECT `a``b`, `a\`b`, "a""b", "a\"b";
+SELECT $$hello$$, $tag$it's\n$tag$, $1_$a$$b$1_$, $$$$;
+SELECT /* outer /* inner */ outer */ 1 # comment
++ 2;
+
+
+-- Format SQL:
+SELECT `a``b`, `a\`b`, "a""b", "a\"b";
+SELECT 'hello', 'it\'s\\n', 'a$$b', '';
+SELECT 1 + 2;

parser/testdata/basic/lexical_fidelity.sql

@@ -0,0 +1,5 @@
+#!/usr/bin/clickhouse
+SELECT `a``b`, `a\`b`, "a""b", "a\"b";
+SELECT $$hello$$, $tag$it's\n$tag$, $1_$a$$b$1_$, $$$$;
+SELECT /* outer /* inner */ outer */ 1 # comment
++ 2;