Add support of accessing the tuple with index field (#114)

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):
@@ -78,7 +81,7 @@ func (p *Parser) parseInfix(expr Expr, precedence int) (Expr, error) {
 		p.matchTokenKind(opTypeDiv), p.matchTokenKind(opTypeMod),
 		p.matchKeyword(KeywordIn), p.matchKeyword(KeywordLike),
 		p.matchKeyword(KeywordIlike), p.matchKeyword(KeywordAnd), p.matchKeyword(KeywordOr),
-		p.matchTokenKind(opTypeCast), p.matchTokenKind(opTypeArrow), p.matchTokenKind(opTypeDoubleEQ):
+		p.matchTokenKind(opTypeArrow), p.matchTokenKind(opTypeDoubleEQ):
 		op := p.last().ToString()
 		_ = p.lexer.consumeToken()
 		rightExpr, err := p.parseSubExpr(p.Pos(), precedence)
@@ -90,6 +93,38 @@ func (p *Parser) parseInfix(expr Expr, precedence int) (Expr, error) {
 			Operation: TokenKind(op),
 			RightExpr: rightExpr,
 		}, nil
+	case p.matchTokenKind(opTypeCast):
+		_ = p.lexer.consumeToken()
+
+		if p.matchTokenKind(TokenIdent) && p.last().String == "Tuple" {
+			name, err := p.parseIdent()
+			if err != nil {
+				return nil, err
+			}
+			if _, err = p.consumeTokenKind("("); err != nil {
+				return nil, err
+			}
+			// it's a tuple type definition after "::" operator
+			rightExpr, err := p.parseNestedType(name, p.Pos())
+			if err != nil {
+				return nil, err
+			}
+			return &BinaryOperation{
+				LeftExpr:  expr,
+				Operation: opTypeCast,
+				RightExpr: rightExpr,
+			}, nil
+		}
+
+		rightExpr, err := p.parseSubExpr(p.Pos(), precedence)
+		if err != nil {
+			return nil, err
+		}
+		return &BinaryOperation{
+			LeftExpr:  expr,
+			Operation: opTypeCast,
+			RightExpr: rightExpr,
+		}, nil
 	case p.matchKeyword(KeywordBetween):
 		return p.parseBetweenClause(expr)
 	case p.matchKeyword(KeywordGlobal):
@@ -106,7 +141,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 +383,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,6 @@
+SELECT tuple('a','b','c').3, .1234;
+
+SELECT toTypeName( tuple('a' as first,'b' as second ,'c' as third)::Tuple(first String,second String,third String)),
+       (tuple('a' as first,'b' as second ,'c' as third)::Tuple(first String,second String,third String)).second,
+       tuple('a','b','c').3,
+       tupleElement(tuple('a','b','c'),1)

parser/testdata/query/format/access_tuple_with_dot.sql

@@ -0,0 +1,11 @@
+-- Origin SQL:
+SELECT tuple('a','b','c').3, .1234;
+
+SELECT toTypeName( tuple('a' as first,'b' as second ,'c' as third)::Tuple(first String,second String,third String)),
+       (tuple('a' as first,'b' as second ,'c' as third)::Tuple(first String,second String,third String)).second,
+       tuple('a','b','c').3,
+       tupleElement(tuple('a','b','c'),1)
+
+-- Format SQL:
+SELECT tuple('a', 'b', 'c').3, .1234;
+SELECT toTypeName(tuple('a' AS first, 'b' AS second, 'c' AS third)::Tuple(first String, second String, third String)), (tuple('a' AS first, 'b' AS second, 'c' AS third)::Tuple(first String, second String, third String)).second, tuple('a', 'b', 'c').3, tupleElement(tuple('a', 'b', 'c'), 1);