Parse path fields with brackets

chaisql · Sep 2, 2021 · cfdb784 · cfdb784
1 parent d5fd879
commit cfdb784
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 10 deletions.
diff --git a/internal/sql/parser/expr.go b/internal/sql/parser/expr.go
@@ -564,18 +564,33 @@ LOOP:
 				FieldName: lit,
 			})
 		case scanner.LSBRACKET:
-			// scan the next token for an integer
+			// the next token can be either an integer or a quoted string
+			// if it's an integer, we have an array index
+			// if it's a quoted string, we have a field name
 			tok, pos, lit := p.Scan()
-			if tok != scanner.INTEGER || lit[0] == '-' {
-				return nil, newParseError(lit, []string{"array index"}, pos)
-			}
-			idx, err := strconv.Atoi(lit)
-			if err != nil {
-				return nil, newParseError(lit, []string{"integer"}, pos)
+			switch tok {
+			case scanner.INTEGER:
+				// is the number negative?
+				if lit[0] == '-' {
+					return nil, newParseError(lit, []string{"integer"}, pos)
+				}
+				// is the number too big?
+				if len(lit) > 10 {
+					return nil, newParseError(lit, []string{"integer"}, pos)
+				}
+				// parse the integer
+				i, err := strconv.ParseInt(lit, 10, 64)
+				if err != nil {
+					return nil, newParseError(lit, []string{"integer"}, pos)
+				}
+				path = append(path, document.PathFragment{
+					ArrayIndex: int(i),
+				})
+			case scanner.STRING:
+				path = append(path, document.PathFragment{
+					FieldName: lit,
+				})
 			}
-			path = append(path, document.PathFragment{
-				ArrayIndex: idx,
-			})
 			// scan the next token for a closing left bracket
 			if err := p.parseTokens(scanner.RSBRACKET); err != nil {
 				return nil, err

diff --git a/internal/sql/parser/expr_test.go b/internal/sql/parser/expr_test.go
@@ -205,15 +205,25 @@ func TestParsePath(t *testing.T) {
 			document.PathFragment{ArrayIndex: 1},
 			document.PathFragment{ArrayIndex: 2},
 		}, false},
+		{"multiple fragments with brackets", `a["b"][100].c[1][2]`, document.Path{
+			document.PathFragment{FieldName: "a"},
+			document.PathFragment{FieldName: "b"},
+			document.PathFragment{ArrayIndex: 100},
+			document.PathFragment{FieldName: "c"},
+			document.PathFragment{ArrayIndex: 1},
+			document.PathFragment{ArrayIndex: 2},
+		}, false},
 		{"with quotes", "`some ident`.` with`[5].`  \"quotes`", document.Path{
 			document.PathFragment{FieldName: "some ident"},
 			document.PathFragment{FieldName: " with"},
 			document.PathFragment{ArrayIndex: 5},
 			document.PathFragment{FieldName: "  \"quotes"},
 		}, false},
+
 		{"negative index", `a.b[-100].c`, nil, true},
 		{"with spaces", `a.  b[100].  c`, nil, true},
 		{"starting with array", `[10].a`, nil, true},
+		{"starting with brackets", `['a']`, nil, true},
 	}
 
 	for _, test := range tests {