feat: json object agg can take expressions too

Signed-off-by: Manan Gupta <[email protected]>
vitessio · Jun 24, 2024 · 376e4b4 · 376e4b4
1 parent 7c81d62
commit 376e4b4
Show file tree

Hide file tree

Showing 11 changed files with 819 additions and 826 deletions.
diff --git a/go/vt/sqlparser/ast.go b/go/vt/sqlparser/ast.go
@@ -2584,8 +2584,8 @@ type (
 	// JSONObjectAgg is an aggregation expression that creates a JSON Object.
 	// For more information, visit https://dev.mysql.com/doc/refman/8.4/en/aggregate-functions.html#function_json-objectagg
 	JSONObjectAgg struct {
-		Key        *ColName
-		Value      *ColName
+		Key        Expr
+		Value      Expr
 		OverClause *OverClause
 	}
 
@@ -3448,9 +3448,7 @@ func (varS *VarSamp) SetArg(expr Expr)              { varS.Arg = expr }
 func (variance *Variance) SetArg(expr Expr)         { variance.Arg = expr }
 func (av *AnyValue) SetArg(expr Expr)               { av.Arg = expr }
 func (jaa *JSONArrayAgg) SetArg(expr Expr)          { jaa.Expr = expr }
-func (joa *JSONObjectAgg) SetArg(expr Expr) {
-	joa.Key = getColNameForExpression(expr, "JSONObjectAgg")
-}
+func (joa *JSONObjectAgg) SetArg(expr Expr)         { joa.Key = expr }
 
 func (min *Min) SetArgs(exprs Exprs) error           { return setFuncArgs(min, exprs, "MIN") }
 func (sum *Sum) SetArgs(exprs Exprs) error           { return setFuncArgs(sum, exprs, "SUM") }
@@ -3473,8 +3471,8 @@ func (joa *JSONObjectAgg) SetArgs(exprs Exprs) error {
 	if len(exprs) != 2 {
 		return vterrors.VT13001("JSONObjectAgg takes in 2 expressions")
 	}
-	joa.Key = getColNameForExpression(exprs[0], "JSONObjectAgg")
-	joa.Value = getColNameForExpression(exprs[1], "JSONObjectAgg")
+	joa.Key = exprs[0]
+	joa.Value = exprs[1]
 	return nil
 }
 

diff --git a/go/vt/sqlparser/ast_clone.go b/go/vt/sqlparser/ast_clone.go
diff --git a/go/vt/sqlparser/ast_copy_on_rewrite.go b/go/vt/sqlparser/ast_copy_on_rewrite.go
diff --git a/go/vt/sqlparser/ast_equals.go b/go/vt/sqlparser/ast_equals.go
diff --git a/go/vt/sqlparser/ast_funcs.go b/go/vt/sqlparser/ast_funcs.go
@@ -2184,15 +2184,6 @@ func ContainsAggregation(e SQLNode) bool {
 	return hasAggregates
 }
 
-// getColNameForExpression gets the column name for the given expression.
-func getColNameForExpression(expr Expr, name string) *ColName {
-	colName, isColName := expr.(*ColName)
-	if !isColName {
-		panic(vterrors.VT13001(fmt.Sprintf("Column name required in %v", name)))
-	}
-	return colName
-}
-
 // setFuncArgs sets the arguments for the aggregation function, while checking that there is only one argument
 func setFuncArgs(aggr AggrFunc, exprs Exprs, name string) error {
 	if len(exprs) != 1 {

diff --git a/go/vt/sqlparser/ast_rewrite.go b/go/vt/sqlparser/ast_rewrite.go
diff --git a/go/vt/sqlparser/ast_visit.go b/go/vt/sqlparser/ast_visit.go
diff --git a/go/vt/sqlparser/cached_size.go b/go/vt/sqlparser/cached_size.go
diff --git a/go/vt/sqlparser/parse_test.go b/go/vt/sqlparser/parse_test.go
@@ -6062,8 +6062,8 @@ var (
 		input:  "select count(1) from user where x_id = 'abc' group by n_id having json_arrayagg(x, y) = '[]'",
 		output: "syntax error at position 83",
 	}, {
-		input:  "select count(1) from user where x_id = 'abc' group by n_id having json_objectagg(x + 1, y) = '[]'",
-		output: "syntax error at position 85",
+		input:  "select count(1) from user where x_id = 'abc' group by n_id having json_objectagg(x, y, z) = '[]'",
+		output: "syntax error at position 87",
 	}, {
 		input:  "select next 1+1 values from a",
 		output: "syntax error at position 15",