Support all subscription operator pattern in IndexExpr (#206)

ast/ast.go

@@ -199,6 +199,15 @@ func (NewConstructor) isExpr()        {}
 func (BracedNewConstructor) isExpr()  {}
 func (BracedConstructor) isExpr()     {}
 
+// SubscriptSpecifier represents specifier of subscript operators.
+type SubscriptSpecifier interface {
+	Node
+	isSubscriptSpecifier()
+}
+
+func (ExprArg) isSubscriptSpecifier()                   {}
+func (SubscriptSpecifierKeyword) isSubscriptSpecifier() {}
+
 // Arg represents argument of function call.
 type Arg interface {
 	Node
@@ -1092,17 +1101,36 @@ type SelectorExpr struct {
 	Ident *Ident
 }
 
-// IndexExpr is array item access expression node.
+// IndexExpr is a subscript operator expression node.
+// This node can be:
+//	- array subscript operator
+//	- struct subscript operator
+//	- JSON subscript operator
+// Note: The name IndexExpr is a historical reason, maybe better to rename to SubscriptExpr.
 //
-//	{{.Expr | sql}}[{{if .Ordinal}}ORDINAL{{else}}OFFSET{{end}}({{.Index | sql}})]
+//	{{.Expr | sql}}[{{.Index | sql}}]
 type IndexExpr struct {
 	// pos = Expr.pos
 	// end = Rbrack + 1
 
 	Rbrack token.Pos // position of "]"
 
-	Ordinal     bool
-	Expr, Index Expr
+	Expr  Expr
+	Index SubscriptSpecifier
+}
+
+// SubscriptSpecifierKeyword is subscript specifier with position keyword.
+//
+//	{{string(.Keyword)}}({{.Expr | sql}})
+type SubscriptSpecifierKeyword struct {
+	// pos = KeywordPos
+	// end = Rparen + 1
+
+	KeywordPos token.Pos // position of Keyword
+	Rparen     token.Pos // position of ")"
+
+	Keyword PositionKeyword
+	Expr    Expr
 }
 
 // CallExpr is function call expression node.

ast/const.go

@@ -22,6 +22,15 @@ const (
 	LoopJoinMethod  JoinMethod = "LOOP" // Undocumented, but the Spanner accept this value at least.
 )
 
+type PositionKeyword string
+
+const (
+	PositionKeywordOffset      PositionKeyword = "OFFSET"
+	PositionKeywordSafeOffset  PositionKeyword = "SAFE_OFFSET"
+	PositionKeywordOrdinal     PositionKeyword = "ORDINAL"
+	PositionKeywordSafeOrdinal PositionKeyword = "SAFE_ORDINAL"
+)
+
 type SetOp string
 
 const (

ast/sql.go

@@ -77,7 +77,10 @@ const (
 
 func exprPrec(e Expr) prec {
 	switch e := e.(type) {
-	case *CallExpr, *CountStarExpr, *CastExpr, *ExtractExpr, *CaseExpr, *IfExpr, *ParenExpr, *ScalarSubQuery, *ArraySubQuery, *ExistsSubQuery, *Param, *Ident, *Path, *ArrayLiteral, *TupleStructLiteral, *TypedStructLiteral, *TypelessStructLiteral, *NullLiteral, *BoolLiteral, *IntLiteral, *FloatLiteral, *StringLiteral, *BytesLiteral, *DateLiteral, *TimestampLiteral, *NumericLiteral:
+	case *CallExpr, *CountStarExpr, *CastExpr, *ExtractExpr, *CaseExpr, *IfExpr, *ParenExpr, *ScalarSubQuery,
+		*ArraySubQuery, *ExistsSubQuery, *Param, *Ident, *Path, *ArrayLiteral, *TupleStructLiteral, *TypedStructLiteral,
+		*TypelessStructLiteral, *NullLiteral, *BoolLiteral, *IntLiteral, *FloatLiteral, *StringLiteral, *BytesLiteral,
+		*DateLiteral, *TimestampLiteral, *NumericLiteral, *JSONLiteral:
 		return precLit
 	case *IndexExpr, *SelectorExpr:
 		return precSelector
@@ -486,14 +489,11 @@ func (s *SelectorExpr) SQL() string {
 
 func (i *IndexExpr) SQL() string {
 	p := exprPrec(i)
-	sql := paren(p, i.Expr) + "["
-	if i.Ordinal {
-		sql += "ORDINAL"
-	} else {
-		sql += "OFFSET"
-	}
-	sql += "(" + i.Index.SQL() + ")]"
-	return sql
+	return paren(p, i.Expr) + "[" + i.Index.SQL() + "]"
+}
+
+func (s *SubscriptSpecifierKeyword) SQL() string {
+	return string(s.Keyword) + "(" + s.Expr.SQL() + ")"
 }
 
 func (c *CallExpr) SQL() string {

parser.go

@@ -1382,31 +1382,53 @@ func (p *Parser) parseSelector() ast.Expr {
 			}
 		case "[":
 			p.nextToken()
-			id := p.expect(token.TokenIdent)
-			ordinal := false
-			if char.EqualFold(id.AsString, "ORDINAL") {
-				ordinal = true
-			} else if char.EqualFold(id.AsString, "OFFSET") {
-				ordinal = false
-			} else {
-				p.panicfAtToken(id, "expected identifier: ORDINAL, OFFSET, but: %s", id.Raw)
-			}
-			p.expect("(")
-			index := p.parseExpr()
-			p.expect(")")
+			index := p.parseIndexSpecifier()
 			rbrack := p.expect("]").Pos
 			expr = &ast.IndexExpr{
-				Rbrack:  rbrack,
-				Ordinal: ordinal,
-				Expr:    expr,
-				Index:   index,
+				Rbrack: rbrack,
+				Expr:   expr,
+				Index:  index,
 			}
 		default:
 			return expr
 		}
 	}
 }
 
+func (p *Parser) parseIndexSpecifier() ast.SubscriptSpecifier {
+	pos := p.Token.Pos
+	switch {
+	case p.Token.IsIdent("OFFSET"), p.Token.IsIdent("ORDINAL"),
+		p.Token.IsIdent("SAFE_OFFSET"), p.Token.IsIdent("SAFE_ORDINAL"):
+		var keyword ast.PositionKeyword
+		switch {
+		case p.Token.IsIdent("OFFSET"):
+			keyword = ast.PositionKeywordOffset
+		case p.Token.IsIdent("ORDINAL"):
+			keyword = ast.PositionKeywordOrdinal
+		case p.Token.IsIdent("SAFE_OFFSET"):
+			keyword = ast.PositionKeywordSafeOffset
+		case p.Token.IsIdent("SAFE_ORDINAL"):
+			keyword = ast.PositionKeywordSafeOrdinal
+
+		}
+		p.nextToken()
+		p.expect("(")
+		expr := p.parseExpr()
+		rparen := p.expect(")").Pos
+
+		return &ast.SubscriptSpecifierKeyword{
+			KeywordPos: pos,
+			Keyword:    keyword,
+			Rparen:     rparen,
+			Expr:       expr,
+		}
+	default:
+		expr := p.parseExpr()
+		return &ast.ExprArg{Expr: expr}
+	}
+}
+
 func (p *Parser) parseLit() ast.Expr {
 	switch p.Token.Kind {
 	case "NULL":

testdata/input/query/select_subscript_operators.sql

@@ -0,0 +1,13 @@
+select
+    [1, 2, 3][offset(1)],
+    [1, 2, 3][ordinal(1)],
+    [1, 2, 3][safe_offset(1)],
+    [1, 2, 3][ordinal(1)],
+    [1, 2, 3][1],
+    STRUCT(1, 2, 3)[offset(1)],
+    STRUCT(1, 2, 3)[ordinal(1)],
+    STRUCT(1, 2, 3)[safe_offset(1)],
+    STRUCT(1, 2, 3)[ordinal(1)],
+    STRUCT(1, 2, 3)[1],
+    JSON '[1, 2, 3]'[1],
+    JSON '{"a": 1, "b": 2, "c": 3}'['a']

testdata/result/query/select_expr.sql.txt

@@ -380,9 +380,8 @@ select 1 + 2, 1 - 2,
       },
       &ast.ExprSelectItem{
         Expr: &ast.IndexExpr{
-          Rbrack:  259,
-          Ordinal: false,
-          Expr:    &ast.ArrayLiteral{
+          Rbrack: 259,
+          Expr:   &ast.ArrayLiteral{
             Array:  -1,
             Lbrack: 240,
             Rbrack: 248,
@@ -408,19 +407,23 @@ select 1 + 2, 1 - 2,
               },
             },
           },
-          Index: &ast.IntLiteral{
-            ValuePos: 257,
-            ValueEnd: 258,
-            Base:     10,
-            Value:    "1",
+          Index: &ast.SubscriptSpecifierKeyword{
+            KeywordPos: 250,
+            Rparen:     258,
+            Keyword:    "OFFSET",
+            Expr:       &ast.IntLiteral{
+              ValuePos: 257,
+              ValueEnd: 258,
+              Base:     10,
+              Value:    "1",
+            },
           },
         },
       },
       &ast.ExprSelectItem{
         Expr: &ast.IndexExpr{
-          Rbrack:  290,
-          Ordinal: false,
-          Expr:    &ast.ArrayLiteral{
+          Rbrack: 290,
+          Expr:   &ast.ArrayLiteral{
             Array:  -1,
             Lbrack: 269,
             Rbrack: 277,
@@ -446,19 +449,23 @@ select 1 + 2, 1 - 2,
               },
             },
           },
-          Index: &ast.IntLiteral{
-            ValuePos: 288,
-            ValueEnd: 289,
-            Base:     10,
-            Value:    "1",
+          Index: &ast.SubscriptSpecifierKeyword{
+            KeywordPos: 279,
+            Rparen:     289,
+            Keyword:    "OFFSET",
+            Expr:       &ast.IntLiteral{
+              ValuePos: 288,
+              ValueEnd: 289,
+              Base:     10,
+              Value:    "1",
+            },
           },
         },
       },
       &ast.ExprSelectItem{
         Expr: &ast.IndexExpr{
-          Rbrack:  320,
-          Ordinal: true,
-          Expr:    &ast.ArrayLiteral{
+          Rbrack: 320,
+          Expr:   &ast.ArrayLiteral{
             Array:  -1,
             Lbrack: 300,
             Rbrack: 308,
@@ -484,11 +491,16 @@ select 1 + 2, 1 - 2,
               },
             },
           },
-          Index: &ast.IntLiteral{
-            ValuePos: 318,
-            ValueEnd: 319,
-            Base:     10,
-            Value:    "1",
+          Index: &ast.SubscriptSpecifierKeyword{
+            KeywordPos: 310,
+            Rparen:     319,
+            Keyword:    "ORDINAL",
+            Expr:       &ast.IntLiteral{
+              ValuePos: 318,
+              ValueEnd: 319,
+              Base:     10,
+              Value:    "1",
+            },
           },
         },
       },