Fix position tracking across ratio, IS NULL, GROUP BY, ALTER and dictionary nodes

parser/ast.go

@@ -317,7 +317,7 @@ func (a *AlterTableDropPartition) Pos() Pos {
 
 func (a *AlterTableDropPartition) End() Pos {
 	if a.Settings != nil {
-		a.Settings.End()
+		return a.Settings.End()
 	}
 	return a.Partition.End()
 }
@@ -4103,6 +4103,7 @@ func (d *DictionarySchemaClause) Accept(visitor ASTVisitor) error {
 
 type DictionaryAttribute struct {
 	NamePos      Pos
+	AttrEnd      Pos // end of the last property, or of the type when there is none
 	Name         *Ident
 	Type         ColumnType
 	Default      Literal
@@ -4117,22 +4118,7 @@ func (d *DictionaryAttribute) Pos() Pos {
 }
 
 func (d *DictionaryAttribute) End() Pos {
-	if d.IsObjectId {
-		return d.NamePos + Pos(len("IS_OBJECT_ID"))
-	}
-	if d.Injective {
-		return d.NamePos + Pos(len("INJECTIVE"))
-	}
-	if d.Hierarchical {
-		return d.NamePos + Pos(len("HIERARCHICAL"))
-	}
-	if d.Expression != nil {
-		return d.Expression.End()
-	}
-	if d.Default != nil {
-		return d.Default.End()
-	}
-	return d.Type.End()
+	return d.AttrEnd
 }
 
 func (d *DictionaryAttribute) Accept(visitor ASTVisitor) error {
@@ -4595,16 +4581,17 @@ func (f *FromClause) Accept(visitor ASTVisitor) error {
 }
 
 type IsNullExpr struct {
-	IsPos Pos
-	Expr  Expr
+	IsPos   Pos // position of the IS keyword
+	NullEnd Pos // end of the NULL keyword
+	Expr    Expr
 }
 
 func (n *IsNullExpr) Pos() Pos {
-	return n.IsPos
+	return n.Expr.Pos()
 }
 
 func (n *IsNullExpr) End() Pos {
-	return n.Expr.End()
+	return n.NullEnd
 }
 
 func (n *IsNullExpr) Accept(visitor ASTVisitor) error {
@@ -4617,16 +4604,17 @@ func (n *IsNullExpr) Accept(visitor ASTVisitor) error {
 }
 
 type IsNotNullExpr struct {
-	IsPos Pos
-	Expr  Expr
+	IsPos   Pos // position of the IS keyword
+	NullEnd Pos // end of the NULL keyword
+	Expr    Expr
 }
 
 func (n *IsNotNullExpr) Pos() Pos {
 	return n.Expr.Pos()
 }
 
 func (n *IsNotNullExpr) End() Pos {
-	return n.Expr.End()
+	return n.NullEnd
 }
 
 func (n *IsNotNullExpr) Accept(visitor ASTVisitor) error {

parser/parser_alter.go

@@ -35,8 +35,11 @@ func (p *Parser) parseAlterTable(pos Pos) (*AlterTable, error) {
 		case p.matchKeyword(KeywordAttach):
 			alter, err = p.parseAlterTableAttachPartition(p.Pos())
 		case p.matchKeyword(KeywordDetach):
+			// like the sibling branches, the clause position is the keyword
+			// itself, so capture it before consuming DETACH
+			detachPos := p.Pos()
 			_ = p.lexer.consumeToken()
-			alter, err = p.parseAlterTableDetachPartition(p.Pos())
+			alter, err = p.parseAlterTableDetachPartition(detachPos)
 		case p.matchKeyword(KeywordFreeze):
 			alter, err = p.parseAlterTableFreezePartition(p.Pos())
 		case p.matchKeyword(KeywordRemove):

parser/parser_column.go

@@ -200,20 +200,26 @@ func (p *Parser) parseInfix(expr Expr, precedence int) (Expr, error) {
 	case p.matchTokenKind(TokenKindQuestionMark):
 		return p.parseTernaryExpr(expr)
 	case p.matchKeyword(KeywordIs):
+		isPos := p.Pos()
 		_ = p.lexer.consumeToken()
 		isNotNull := p.tryConsumeKeywords(KeywordNot)
+		// the expression ends at the NULL keyword; capture its end before
+		// expectKeyword consumes it
+		nullEnd := p.End()
 		if err := p.expectKeyword(KeywordNull); err != nil {
 			return nil, err
 		}
 		if isNotNull {
 			return &IsNotNullExpr{
-				IsPos: p.Pos(),
-				Expr:  expr,
+				IsPos:   isPos,
+				NullEnd: nullEnd,
+				Expr:    expr,
 			}, nil
 		}
 		return &IsNullExpr{
-			IsPos: p.Pos(),
-			Expr:  expr,
+			IsPos:   isPos,
+			NullEnd: nullEnd,
+			Expr:    expr,
 		}, nil
 	default:
 		return nil, fmt.Errorf("unexpected token kind: %s", p.currentTokenKind())

parser/parser_common.go

@@ -430,7 +430,8 @@ func (p *Parser) parseRatioExpr(pos Pos) (*RatioExpr, error) {
 
 	var denominator *NumberLiteral
 	if p.tryConsumeTokenKind(TokenKindDiv) != nil {
-		denominator, err = p.parseNumber(pos)
+		// the denominator starts at its own token, not at the numerator
+		denominator, err = p.parseNumber(p.Pos())
 		if err != nil {
 			return nil, err
 		}

parser/parser_query.go

@@ -503,6 +503,7 @@ func (p *Parser) parseGroupByClause(pos Pos) (*GroupByClause, error) {
 
 	var expr Expr
 	var err error
+	var groupByEnd Pos
 	aggregateType := ""
 	switch {
 	case p.matchKeyword(KeywordCube) || p.matchKeyword(KeywordRollup):
@@ -512,22 +513,32 @@ func (p *Parser) parseGroupByClause(pos Pos) (*GroupByClause, error) {
 	case p.tryConsumeKeywords(KeywordGrouping, KeywordSets):
 		aggregateType = "GROUPING SETS"
 		expr, err = p.parseFunctionParams(p.Pos())
-	case p.tryConsumeKeywords(KeywordAll):
+	case p.matchKeyword(KeywordAll):
+		// GROUP BY ALL has no expression list; the clause ends at ALL itself
+		groupByEnd = p.End()
+		_ = p.lexer.consumeToken()
 		aggregateType = "ALL"
 	default:
 		expr, err = p.parseColumnExprListWithLParen(p.Pos())
 	}
 	if err != nil {
 		return nil, err
 	}
+	if expr != nil {
+		groupByEnd = expr.End()
+	}
 	groupBy := &GroupByClause{
 		GroupByPos:    pos,
+		GroupByEnd:    groupByEnd,
 		AggregateType: aggregateType,
 		Expr:          expr,
 	}
 
 	// parse WITH CUBE, ROLLUP, TOTALS
 	for p.tryConsumeKeywords(KeywordWith) {
+		// the clause now extends to the CUBE/ROLLUP/TOTALS token; capture its
+		// end before it is consumed
+		keywordEnd := p.End()
 		switch {
 		case p.tryConsumeKeywords(KeywordCube):
 			groupBy.WithCube = true
@@ -538,8 +549,8 @@ func (p *Parser) parseGroupByClause(pos Pos) (*GroupByClause, error) {
 		default:
 			return nil, fmt.Errorf("expected CUBE, ROLLUP or TOTALS, got %s", p.currentTokenKind())
 		}
+		groupBy.GroupByEnd = keywordEnd
 	}
-	groupBy.GroupByEnd = p.Pos()
 
 	return groupBy, nil
 }
@@ -1075,13 +1086,15 @@ func (p *Parser) parseSelectStmt(pos Pos) (*SelectQuery, error) { // nolint: fun
 		statementEnd = groupBy.End()
 	}
 	withTotal := false
-	lastPos := p.Pos()
 	if p.tryConsumeKeywords(KeywordWith) {
+		// the statement now ends at the TOTALS token; capture its end before
+		// expectKeyword consumes it
+		totalsEnd := p.End()
 		if err := p.expectKeyword(KeywordTotals); err != nil {
 			return nil, err
 		}
 		withTotal = true
-		statementEnd = lastPos
+		statementEnd = totalsEnd
 	}
 	having, err := p.tryParseHavingClause(p.Pos())
 	if err != nil {

parser/parser_table.go

@@ -2058,10 +2058,14 @@ func (p *Parser) parseDictionaryAttribute(pos Pos) (*DictionaryAttribute, error)
 		NamePos: pos,
 		Name:    name,
 		Type:    columnType,
+		AttrEnd: columnType.End(),
 	}
 
 	// Parse optional attribute properties
 	for {
+		// end of the property keyword about to be consumed; flag-only
+		// properties (HIERARCHICAL, ...) end at the keyword itself
+		keywordEnd := p.End()
 		switch {
 		case p.tryConsumeKeywords(KeywordDefault):
 			if attr.Default != nil {
@@ -2072,6 +2076,7 @@ func (p *Parser) parseDictionaryAttribute(pos Pos) (*DictionaryAttribute, error)
 				return nil, err
 			}
 			attr.Default = literal
+			attr.AttrEnd = literal.End()
 		case p.tryConsumeKeywords(KeywordExpression):
 			if attr.Expression != nil {
 				return nil, fmt.Errorf("duplicate EXPRESSION clause")
@@ -2081,21 +2086,25 @@ func (p *Parser) parseDictionaryAttribute(pos Pos) (*DictionaryAttribute, error)
 				return nil, err
 			}
 			attr.Expression = expr
+			attr.AttrEnd = expr.End()
 		case p.tryConsumeKeywords(KeywordHierarchical):
 			if attr.Hierarchical {
 				return nil, fmt.Errorf("duplicate HIERARCHICAL clause")
 			}
 			attr.Hierarchical = true
+			attr.AttrEnd = keywordEnd
 		case p.tryConsumeKeywords(KeywordInjective):
 			if attr.Injective {
 				return nil, fmt.Errorf("duplicate INJECTIVE clause")
 			}
 			attr.Injective = true
+			attr.AttrEnd = keywordEnd
 		case p.tryConsumeKeywords(KeywordIs_object_id):
 			if attr.IsObjectId {
 				return nil, fmt.Errorf("duplicate IS_OBJECT_ID clause")
 			}
 			attr.IsObjectId = true
+			attr.AttrEnd = keywordEnd
 		default:
 			// No more attribute properties
 			return attr, nil

parser/position_test.go

@@ -0,0 +1,92 @@
+package parser
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func parseOneStmt(t *testing.T, sql string) Expr {
+	t.Helper()
+	stmts, err := NewParser(sql).ParseStmts()
+	require.NoError(t, err)
+	require.Len(t, stmts, 1)
+	return stmts[0]
+}
+
+func TestRatioExprPositions(t *testing.T) {
+	sql := "SELECT * FROM t SAMPLE 1/2"
+	stmt := parseOneStmt(t, sql).(*SelectQuery)
+	sample := stmt.From.Expr.(*JoinTableExpr).SampleRatio
+	require.NotNil(t, sample)
+	ratio := sample.Ratio
+	require.NotNil(t, ratio.Denominator)
+	// numerator "1" at offset 23, denominator "2" at offset 25 — the spans
+	// must not overlap (the denominator used to inherit the numerator's pos)
+	require.Equal(t, Pos(23), ratio.Numerator.Pos())
+	require.Equal(t, Pos(24), ratio.Numerator.End())
+	require.Equal(t, Pos(25), ratio.Denominator.Pos())
+	require.Equal(t, Pos(26), ratio.Denominator.End())
+}
+
+func TestGroupByClauseEnd(t *testing.T) {
+	sql := "SELECT a FROM t GROUP BY a WITH TOTALS"
+	stmt := parseOneStmt(t, sql).(*SelectQuery)
+	require.NotNil(t, stmt.GroupBy)
+	// the clause ends at the TOTALS keyword, not at the next token's start
+	require.Equal(t, Pos(len(sql)), stmt.GroupBy.End())
+
+	sql = "SELECT a FROM t GROUP BY a"
+	stmt = parseOneStmt(t, sql).(*SelectQuery)
+	require.Equal(t, Pos(len(sql)), stmt.GroupBy.End())
+
+	sql = "SELECT a FROM t GROUP BY ALL"
+	stmt = parseOneStmt(t, sql).(*SelectQuery)
+	require.Equal(t, Pos(len(sql)), stmt.GroupBy.End())
+}
+
+func TestIsNullExprPositions(t *testing.T) {
+	sql := "SELECT a IS NULL"
+	stmt := parseOneStmt(t, sql).(*SelectQuery)
+	isNull := stmt.SelectItems[0].Expr.(*IsNullExpr)
+	// the node spans `a IS NULL`: from the operand to the end of NULL
+	require.Equal(t, Pos(7), isNull.Pos())
+	require.Equal(t, Pos(len(sql)), isNull.End())
+	require.Equal(t, Pos(9), isNull.IsPos) // position of IS
+
+	sql = "SELECT a IS NOT NULL"
+	stmt = parseOneStmt(t, sql).(*SelectQuery)
+	isNotNull := stmt.SelectItems[0].Expr.(*IsNotNullExpr)
+	require.Equal(t, Pos(7), isNotNull.Pos())
+	require.Equal(t, Pos(len(sql)), isNotNull.End())
+	require.Equal(t, Pos(9), isNotNull.IsPos)
+}
+
+func TestAlterDetachPartitionPos(t *testing.T) {
+	sql := "ALTER TABLE t DETACH PARTITION p"
+	stmt := parseOneStmt(t, sql).(*AlterTable)
+	detach := stmt.AlterExprs[0].(*AlterTableDetachPartition)
+	// the clause starts at the DETACH keyword, like every sibling clause
+	require.Equal(t, Pos(14), detach.Pos())
+}
+
+func TestAlterDropPartitionEndIncludesSettings(t *testing.T) {
+	sql := "ALTER TABLE t DROP PARTITION p SETTINGS mutations_sync=1"
+	stmt := parseOneStmt(t, sql).(*AlterTable)
+	drop := stmt.AlterExprs[0].(*AlterTableDropPartition)
+	require.NotNil(t, drop.Settings)
+	// End() used to discard the Settings end and stop at the partition
+	require.Equal(t, drop.Settings.End(), drop.End())
+	require.Greater(t, drop.End(), drop.Partition.End())
+}
+
+func TestDictionaryAttributeEnd(t *testing.T) {
+	sql := "CREATE DICTIONARY d (user_id UInt64 IS_OBJECT_ID) PRIMARY KEY user_id SOURCE(CLICKHOUSE()) LAYOUT(FLAT()) LIFETIME(300)"
+	stmt := parseOneStmt(t, sql).(*CreateDictionary)
+	attrs := stmt.Schema.Attributes
+	require.Len(t, attrs, 1)
+	// the attribute ends at IS_OBJECT_ID; the old implementation returned
+	// NamePos + len("IS_OBJECT_ID"), landing in the middle of the type
+	require.Equal(t, "IS_OBJECT_ID", sql[36:48])
+	require.Equal(t, Pos(48), attrs[0].End())
+}

parser/testdata/ddl/output/alter_table_add_projection.sql.golden.json

@@ -51,7 +51,7 @@
             },
             "GroupBy": {
               "GroupByPos": 86,
-              "GroupByEnd": 105,
+              "GroupByEnd": 104,
               "AggregateType": "",
               "Expr": {
                 "ListPos": 95,

parser/testdata/ddl/output/alter_table_detach_partition.sql.golden.json

@@ -19,7 +19,7 @@
     "OnCluster": null,
     "AlterExprs": [
       {
-        "DetachPos": 27,
+        "DetachPos": 20,
         "Partition": {
           "PartitionPos": 27,
           "Expr": {

parser/testdata/ddl/output/alter_table_drop_detach_partition.sql.golden.json

@@ -1,7 +1,7 @@
 [
   {
     "AlterPos": 0,
-    "StatementEnd": 120,
+    "StatementEnd": 154,
     "TableIdentifier": {
       "Database": {
         "Name": "app_utc_00",