Add support of accessing fields with the dot operation

Author: git-hulk

parser/ast.go

@@ -187,6 +187,40 @@ func (p *BinaryOperation) Accept(visitor ASTVisitor) error {
 	return visitor.VisitBinaryExpr(p)
 }
 
+type IndexOperation struct {
+	LeftExpr  Expr
+	Operation TokenKind
+	Index     Expr
+}
+
+func (i *IndexOperation) Accept(visitor ASTVisitor) error {
+	visitor.enter(i)
+	defer visitor.leave(i)
+	if err := i.LeftExpr.Accept(visitor); err != nil {
+		return err
+	}
+	if err := i.Index.Accept(visitor); err != nil {
+		return err
+	}
+	return visitor.VisitIndexOperation(i)
+}
+
+func (i *IndexOperation) Pos() Pos {
+	return i.LeftExpr.Pos()
+}
+
+func (i *IndexOperation) End() Pos {
+	return i.Index.End()
+}
+
+func (i *IndexOperation) String() string {
+	var builder strings.Builder
+	builder.WriteString(i.LeftExpr.String())
+	builder.WriteString(string(i.Operation))
+	builder.WriteString(i.Index.String())
+	return builder.String()
+}
+
 type JoinTableExpr struct {
 	Table        *TableExpr
 	StatementEnd Pos

parser/ast_visitor.go

@@ -4,6 +4,7 @@ type ASTVisitor interface {
 	VisitOperationExpr(expr *OperationExpr) error
 	VisitTernaryExpr(expr *TernaryOperation) error
 	VisitBinaryExpr(expr *BinaryOperation) error
+	VisitIndexOperation(expr *IndexOperation) error
 	VisitAlterTable(expr *AlterTable) error
 	VisitAlterTableAttachPartition(expr *AlterTableAttachPartition) error
 	VisitAlterTableDetachPartition(expr *AlterTableDetachPartition) error
@@ -188,6 +189,13 @@ func (v *DefaultASTVisitor) VisitBinaryExpr(expr *BinaryOperation) error {
 	return nil
 }
 
+func (v *DefaultASTVisitor) VisitIndexOperation(expr *IndexOperation) error {
+	if v.Visit != nil {
+		return v.Visit(expr)
+	}
+	return nil
+}
+
 func (v *DefaultASTVisitor) VisitJoinTableExpr(expr *JoinTableExpr) error {
 	if v.Visit != nil {
 		return v.Visit(expr)

parser/lexer.go

@@ -15,6 +15,7 @@ const (
 	TokenInt     TokenKind = "<int>"
 	TokenFloat   TokenKind = "<float>"
 	TokenString  TokenKind = "<string>"
+	TokenDot               = "."
 )
 
 const (
@@ -328,19 +329,14 @@ func (l *Lexer) consumeToken() error {
 			return nil
 		}
 	case '.':
-		// check if the next token is a number. If so, parse it as a float number
-		if l.peekOk(1) && IsDigit(l.peekN(1)) {
-			return l.consumeNumber()
-		}
-		// check if the previous lastToken is an Ident. If so, it's a field name.
-		if l.lastToken != nil && l.lastToken.Kind != TokenIdent {
-			return fmt.Errorf("'.' should be after an Ident, but got <%q>", l.lastToken.Kind)
+		l.lastToken = &Token{
+			String: l.slice(0, 1),
+			Kind:   TokenDot,
+			Pos:    Pos(l.current),
+			End:    Pos(l.current + 1),
 		}
-	}
-
-	// The subsequent lastToken after the dot should be an Ident.
-	if l.lastToken != nil && l.lastToken.Kind == "." && !IsIdentStart(l.peekN(0)) {
-		return fmt.Errorf("'.' should follow with an Ident, but got <%q>", l.lastToken.Kind)
+		l.skipN(1)
+		return nil
 	}
 
 	if IsIdentStart(l.peekN(0)) {

parser/lexer_test.go

@@ -106,7 +106,6 @@ func TestConsumeNumber(t *testing.T) {
 	t.Run("Float number", func(t *testing.T) {
 		floats := []string{
 			"123.456",
-			".456",
 			"123.456e+10",
 			"123.456e-10",
 			"123.456e10",
@@ -126,7 +125,6 @@ func TestConsumeNumber(t *testing.T) {
 
 	t.Run("Invalid float number", func(t *testing.T) {
 		invalidFloats := []string{
-			".456a",
 			"123.456b",
 			"123.456e",
 			"123.456e+",

parser/parse_system.go

@@ -921,7 +921,7 @@ func (p *Parser) parseGrantSource(_ Pos) (*TableIdentifier, error) {
 		return nil, err
 	}
 
-	if p.tryConsumeTokenKind(".") == nil {
+	if p.tryConsumeTokenKind(TokenDot) == nil {
 		return &TableIdentifier{
 			Table: ident,
 		}, nil

parser/parser_column.go

@@ -20,6 +20,7 @@ const (
 	PrecedenceMulDivMod
 	PrecedenceBracket
 	PrecedenceArrow
+	PrecedenceDot
 	PrecedenceDoubleColon
 )
 
@@ -56,6 +57,8 @@ func (p *Parser) getNextPrecedence() int {
 		return PrecedenceBracket
 	case p.matchTokenKind(opTypeCast):
 		return PrecedenceDoubleColon
+	case p.matchTokenKind(TokenDot):
+		return PrecedenceDot
 	case p.matchKeyword(KeywordBetween), p.matchKeyword(KeywordLike), p.matchKeyword(KeywordIlike):
 		return PrecedenceBetweenLike
 	case p.matchKeyword(KeywordIn):
@@ -106,7 +109,24 @@ func (p *Parser) parseInfix(expr Expr, precedence int) (Expr, error) {
 			Operation: "GLOBAL IN",
 			RightExpr: rightExpr,
 		}, nil
-
+	case p.matchTokenKind(TokenDot):
+		_ = p.lexer.consumeToken()
+		// access column with dot notation
+		var rightExpr Expr
+		var err error
+		if p.matchTokenKind(TokenIdent) {
+			rightExpr, err = p.parseIdent()
+		} else {
+			rightExpr, err = p.parseDecimal(p.Pos())
+		}
+		if err != nil {
+			return nil, err
+		}
+		return &IndexOperation{
+			LeftExpr:  expr,
+			Operation: TokenDot,
+			Index:     rightExpr,
+		}, nil
 	case p.matchKeyword(KeywordNot):
 		_ = p.lexer.consumeToken()
 		switch {
@@ -331,6 +351,8 @@ func (p *Parser) parseColumnExpr(pos Pos) (Expr, error) { //nolint:funlen
 			return p.parseQueryParam(p.Pos())
 		}
 		return p.parseMapLiteral(p.Pos())
+	case p.matchTokenKind(TokenDot):
+		return p.parseNumber(p.Pos())
 	case p.matchTokenKind(opTypeQuery):
 		// Placeholder `?`
 		_ = p.lexer.consumeToken()

parser/parser_common.go

@@ -110,7 +110,7 @@ func (p *Parser) parseIdentOrStar() (*Ident, error) {
 }
 
 func (p *Parser) tryParseDotIdent(_ Pos) (*Ident, error) {
-	if p.tryConsumeTokenKind(".") == nil {
+	if p.tryConsumeTokenKind(TokenDot) == nil {
 		return nil, nil // nolint
 	}
 	return p.parseIdent()
@@ -211,6 +211,17 @@ func (p *Parser) parseNumber(pos Pos) (*NumberLiteral, error) {
 		lastToken, err = p.consumeTokenKind(TokenInt)
 	case p.matchTokenKind(TokenFloat):
 		lastToken, err = p.consumeTokenKind(TokenFloat)
+	case p.matchTokenKind(TokenDot):
+		p.lexer.consumeToken()
+		lastToken, err = p.consumeTokenKind(TokenInt)
+		if err != nil {
+			return nil, err
+		}
+		if lastToken.Base != 10 {
+			return nil, fmt.Errorf("invalid decimal literal: %q", lastToken.String)
+		}
+		lastToken.String = "." + lastToken.String
+		lastToken.Kind = TokenFloat
 	default:
 		return nil, fmt.Errorf("expected <int> or <float>, but got %q", p.lastTokenKind())
 	}

parser/parser_table.go

@@ -223,14 +223,14 @@ func (p *Parser) parseIdentOrFunction(_ Pos) (Expr, error) {
 			}, nil
 		}
 		return funcExpr, nil
-	case p.tryConsumeTokenKind(".") != nil:
+	case p.tryConsumeTokenKind(TokenDot) != nil:
 		switch {
 		case p.matchTokenKind(TokenIdent):
 			nextIdent, err := p.parseIdent()
 			if err != nil {
 				return nil, err
 			}
-			if p.tryConsumeTokenKind(".") != nil {
+			if p.tryConsumeTokenKind(TokenDot) != nil {
 				thirdIdent, err := p.parseIdent()
 				if err != nil {
 					return nil, err
@@ -310,7 +310,7 @@ func (p *Parser) parseTableSchemaClause(pos Pos) (*TableSchemaClause, error) {
 				return nil, err
 			}
 			switch {
-			case p.matchTokenKind("."):
+			case p.matchTokenKind(TokenDot):
 				// it's a database.table
 				dotIdent, err := p.tryParseDotIdent(p.Pos())
 				if err != nil {
@@ -489,7 +489,7 @@ func (p *Parser) parseTableArgExpr(pos Pos) (Expr, error) {
 		}
 		switch {
 		// nest identifier
-		case p.matchTokenKind("."):
+		case p.matchTokenKind(TokenDot):
 			dotIdent, err := p.tryParseDotIdent(p.Pos())
 			if err != nil {
 				return nil, err

parser/testdata/query/access_tuple_with_dot.sql

@@ -0,0 +1 @@
+SELECT tuple('a','b','c').3, .1234

parser/testdata/query/format/access_tuple_with_dot.sql

@@ -0,0 +1,5 @@
+-- Origin SQL:
+SELECT tuple('a','b','c').3, .1234
+
+-- Format SQL:
+SELECT tuple('a', 'b', 'c').3, .1234;

parser/testdata/query/output/access_tuple_with_dot.sql.golden.json

@@ -0,0 +1,93 @@
+[
+  {
+    "SelectPos": 0,
+    "StatementEnd": 34,
+    "With": null,
+    "Top": null,
+    "SelectItems": [
+      {
+        "Expr": {
+          "LeftExpr": {
+            "Name": {
+              "Name": "tuple",
+              "QuoteType": 1,
+              "NamePos": 7,
+              "NameEnd": 12
+            },
+            "Params": {
+              "LeftParenPos": 12,
+              "RightParenPos": 24,
+              "Items": {
+                "ListPos": 14,
+                "ListEnd": 23,
+                "HasDistinct": false,
+                "Items": [
+                  {
+                    "Expr": {
+                      "LiteralPos": 14,
+                      "LiteralEnd": 15,
+                      "Literal": "a"
+                    },
+                    "Alias": null
+                  },
+                  {
+                    "Expr": {
+                      "LiteralPos": 18,
+                      "LiteralEnd": 19,
+                      "Literal": "b"
+                    },
+                    "Alias": null
+                  },
+                  {
+                    "Expr": {
+                      "LiteralPos": 22,
+                      "LiteralEnd": 23,
+                      "Literal": "c"
+                    },
+                    "Alias": null
+                  }
+                ]
+              },
+              "ColumnArgList": null
+            }
+          },
+          "Operation": ".",
+          "Index": {
+            "NumPos": 26,
+            "NumEnd": 27,
+            "Literal": "3",
+            "Base": 10
+          }
+        },
+        "Modifiers": [],
+        "Alias": null
+      },
+      {
+        "Expr": {
+          "NumPos": 29,
+          "NumEnd": 34,
+          "Literal": ".1234",
+          "Base": 10
+        },
+        "Modifiers": [],
+        "Alias": null
+      }
+    ],
+    "From": null,
+    "ArrayJoin": null,
+    "Window": null,
+    "Prewhere": null,
+    "Where": null,
+    "GroupBy": null,
+    "WithTotal": false,
+    "Having": null,
+    "OrderBy": null,
+    "LimitBy": null,
+    "Limit": null,
+    "Settings": null,
+    "Format": null,
+    "UnionAll": null,
+    "UnionDistinct": null,
+    "Except": null
+  }
+]