Add support of parsing map literal and query parameter (#94)

parser/ast.go

@@ -2844,6 +2844,92 @@ func (f *ParamExprList) Accept(visitor ASTVisitor) error {
 	return visitor.VisitParamExprList(f)
 }
 
+type KeyValue struct {
+	Key   StringLiteral
+	Value Expr
+}
+
+type MapLiteral struct {
+	LBracePos Pos
+	RBracePos Pos
+	KeyValues []KeyValue
+}
+
+func (m *MapLiteral) Pos() Pos {
+	return m.LBracePos
+}
+
+func (m *MapLiteral) End() Pos {
+	return m.RBracePos
+}
+
+func (m *MapLiteral) String(level int) string {
+	var builder strings.Builder
+	builder.WriteString("{")
+
+	for i, value := range m.KeyValues {
+		if i > 0 {
+			builder.WriteString(", ")
+		}
+		builder.WriteString(value.Key.String(level))
+		builder.WriteString(": ")
+		builder.WriteString(value.Value.String(level))
+	}
+	builder.WriteString("}")
+	return builder.String()
+}
+
+func (m *MapLiteral) Accept(visitor ASTVisitor) error {
+	visitor.enter(m)
+	defer visitor.leave(m)
+	for _, kv := range m.KeyValues {
+		if err := kv.Key.Accept(visitor); err != nil {
+			return err
+		}
+		if err := kv.Value.Accept(visitor); err != nil {
+			return err
+		}
+	}
+	return visitor.VisitMapLiteral(m)
+}
+
+type QueryParam struct {
+	LBracePos Pos
+	RBracePos Pos
+	Name      *Ident
+	Type      Expr
+}
+
+func (q *QueryParam) Pos() Pos {
+	return q.LBracePos
+}
+
+func (q *QueryParam) End() Pos {
+	return q.RBracePos
+}
+
+func (q *QueryParam) String(level int) string {
+	var builder strings.Builder
+	builder.WriteString("{")
+	builder.WriteString(q.Name.String(level))
+	builder.WriteString(": ")
+	builder.WriteString(q.Type.String(level))
+	builder.WriteString("}")
+	return builder.String()
+}
+
+func (q *QueryParam) Accept(visitor ASTVisitor) error {
+	visitor.enter(q)
+	defer visitor.leave(q)
+	if err := q.Name.Accept(visitor); err != nil {
+		return err
+	}
+	if err := q.Type.Accept(visitor); err != nil {
+		return err
+	}
+	return visitor.VisitQueryParam(q)
+}
+
 type ArrayParamList struct {
 	LeftBracketPos  Pos
 	RightBracketPos Pos

parser/ast_visitor.go

@@ -65,7 +65,9 @@ type ASTVisitor interface {
 	VisitSettingsExpr(expr *SettingExprList) error
 	VisitSettingsExprList(expr *SettingsClause) error
 	VisitParamExprList(expr *ParamExprList) error
+	VisitMapLiteral(expr *MapLiteral) error
 	VisitArrayParamList(expr *ArrayParamList) error
+	VisitQueryParam(expr *QueryParam) error
 	VisitObjectParams(expr *ObjectParams) error
 	VisitFunctionExpr(expr *FunctionExpr) error
 	VisitWindowFunctionExpr(expr *WindowFunctionExpr) error
@@ -624,6 +626,20 @@ func (v *DefaultASTVisitor) VisitArrayParamList(expr *ArrayParamList) error {
 	return nil
 }
 
+func (v *DefaultASTVisitor) VisitQueryParam(expr *QueryParam) error {
+	if v.Visit != nil {
+		return v.Visit(expr)
+	}
+	return nil
+}
+
+func (v *DefaultASTVisitor) VisitMapLiteral(expr *MapLiteral) error {
+	if v.Visit != nil {
+		return v.Visit(expr)
+	}
+	return nil
+}
+
 func (v *DefaultASTVisitor) VisitObjectParams(expr *ObjectParams) error {
 	if v.Visit != nil {
 		return v.Visit(expr)

parser/parser_column.go

@@ -325,6 +325,17 @@ func (p *Parser) parseUnaryExpr(pos Pos) (Expr, error) {
 
 }
 
+func (p *Parser) peekTokenKind(kind TokenKind) bool {
+	if p.lexer.isEOF() {
+		return false
+	}
+	token, err := p.lexer.peekToken()
+	if err != nil || token == nil {
+		return false
+	}
+	return token.Kind == kind
+}
+
 func (p *Parser) parseColumnExpr(pos Pos) (Expr, error) { //nolint:funlen
 	switch {
 	case p.matchKeyword(KeywordInterval):
@@ -357,12 +368,18 @@ func (p *Parser) parseColumnExpr(pos Pos) (Expr, error) { //nolint:funlen
 				return p.parseSelectQuery(pos)
 			}
 		}
-		return p.parseFunctionParams(pos)
+		return p.parseFunctionParams(p.Pos())
 	case p.matchTokenKind("*"):
-		return p.parseColumnStar(pos)
+		return p.parseColumnStar(p.Pos())
 	case p.matchTokenKind("["):
-		return p.parseArrayParams(pos)
-
+		return p.parseArrayParams(p.Pos())
+	case p.matchTokenKind("{"):
+		// The map literal string also starts with '{', so we need to check the next token
+		// to determine if it is a map literal or a query param.
+		if p.peekTokenKind(TokenIdent) {
+			return p.parseQueryParam(p.Pos())
+		}
+		return p.parseMapLiteral(p.Pos())
 	default:
 		return nil, fmt.Errorf("unexpected token kind: %s", p.lastTokenKind())
 	}
@@ -561,6 +578,71 @@ func (p *Parser) parseFunctionParams(pos Pos) (*ParamExprList, error) {
 	return paramExprList, nil
 }
 
+func (p *Parser) parseMapLiteral(pos Pos) (*MapLiteral, error) {
+	if _, err := p.consumeTokenKind("{"); err != nil {
+		return nil, err
+	}
+
+	keyValues := make([]KeyValue, 0)
+	for !p.lexer.isEOF() && !p.matchTokenKind("}") {
+		key, err := p.parseString(p.Pos())
+		if err != nil {
+			return nil, err
+		}
+		if _, err := p.consumeTokenKind(":"); err != nil {
+			return nil, err
+		}
+		value, err := p.parseExpr(p.Pos())
+		if err != nil {
+			return nil, err
+		}
+		keyValues = append(keyValues, KeyValue{
+			Key:   *key,
+			Value: value,
+		})
+		if p.tryConsumeTokenKind(",") == nil {
+			break
+		}
+	}
+	rightBracePos := p.Pos()
+	if _, err := p.consumeTokenKind("}"); err != nil {
+		return nil, err
+	}
+	return &MapLiteral{
+		LBracePos: pos,
+		RBracePos: rightBracePos,
+		KeyValues: keyValues,
+	}, nil
+}
+
+func (p *Parser) parseQueryParam(pos Pos) (*QueryParam, error) {
+	if _, err := p.consumeTokenKind("{"); err != nil {
+		return nil, err
+	}
+
+	ident, err := p.parseIdent()
+	if err != nil {
+		return nil, err
+	}
+	if _, err := p.consumeTokenKind(":"); err != nil {
+		return nil, err
+	}
+	columnType, err := p.parseColumnType(p.Pos())
+	if err != nil {
+		return nil, err
+	}
+	rightBracePos := p.Pos()
+	if _, err := p.consumeTokenKind("}"); err != nil {
+		return nil, err
+	}
+	return &QueryParam{
+		LBracePos: pos,
+		RBracePos: rightBracePos,
+		Name:      ident,
+		Type:      columnType,
+	}, nil
+}
+
 func (p *Parser) parseArrayParams(pos Pos) (*ArrayParamList, error) {
 	if _, err := p.consumeTokenKind("["); err != nil {
 		return nil, err

parser/parser_table.go

@@ -797,6 +797,12 @@ func (p *Parser) parseSettingsExprList(pos Pos) (*SettingExprList, error) {
 		if err != nil {
 			return nil, err
 		}
+	case p.matchTokenKind("{"):
+		m, err := p.parseMapLiteral(p.Pos())
+		if err != nil {
+			return nil, err
+		}
+		expr = m
 	default:
 		return nil, fmt.Errorf("unexpected token: %q, expected <number> or <string>", p.last().String)
 	}

parser/testdata/query/format/select_with_query_parameter.sql

@@ -0,0 +1,33 @@
+-- Origin SQL:
+SET param_a = 13;
+SET param_b = 'str';
+SET param_c = '2022-08-04 18:30:53';
+SET param_d = {'10': [11, 12], '13': [14, 15]};
+
+SELECT
+    {a: UInt32},
+    {b: String},
+    {c: DateTime},
+    {d: Map(String, Array(UInt8))};
+
+SELECT * FROM clickhouse WHERE tenant_id = {tenant_id: String};
+
+
+-- Format SQL:
+SET param_a=13;
+SET param_b='str';
+SET param_c='2022-08-04 18:30:53';
+SET param_d={'10': [11, 12], '13': [14, 15]};
+
+SELECT 
+  {a: UInt32},
+  {b: String},
+  {c: DateTime},
+  {d: Map(String,Array(UInt8))};
+
+SELECT 
+  *
+FROM
+  clickhouse
+WHERE
+  tenant_id = {tenant_id: String};