Injected order by

Author: zachmu

sql/planbuilder/aggregates.go

@@ -492,7 +492,7 @@ func (b *Builder) buildGroupConcat(inScope *scope, e *ast.GroupConcatExpr) sql.E
 		sortFields = append(sortFields, sf)
 	}
 
-	//TODO: this should be acquired at runtime, not at parse time, so fix this
+	// TODO: this should be acquired at runtime, not at parse time, so fix this
 	gcml, err := b.ctx.GetSessionVariable(b.ctx, "group_concat_max_len")
 	if err != nil {
 		b.handleErr(err)
@@ -515,6 +515,55 @@ func (b *Builder) buildGroupConcat(inScope *scope, e *ast.GroupConcatExpr) sql.E
 	return col.scalarGf()
 }
 
+// buildOrderedInjectedExpr builds an InjectedExpr with an ORDER BY dependency
+func (b *Builder) buildOrderedInjectedExpr(inScope *scope, e *ast.OrderedInjectedExpr) sql.Expression {
+	inScope.initGroupBy()
+	gb := inScope.groupBy
+
+	resolvedChildren := make([]any, len(e.Children))
+	for i, child := range e.Children {
+		resolvedChildren[i] = b.buildScalar(inScope, child)
+	}
+
+	orderByScope := b.analyzeOrderBy(inScope, inScope, e.OrderBy)
+	var sortFields sql.SortFields
+	for _, c := range orderByScope.cols {
+		so := sql.Ascending
+		if c.descending {
+			so = sql.Descending
+		}
+		scalar := c.scalar
+		if scalar == nil {
+			scalar = c.scalarGf()
+		}
+		sf := sql.SortField{
+			Column: scalar,
+			Order:  so,
+		}
+		sortFields = append(sortFields, sf)
+	}
+
+	resolvedChildren = append(resolvedChildren, sortFields)
+
+	expr := b.buildInjectedExpressionFromResolvedChildren(e.InjectedExpr, resolvedChildren)
+	agg, ok := expr.(sql.Aggregation)
+	if !ok {
+		b.handleErr(fmt.Errorf("expected sql.Aggregation, got %T", expr))
+	}
+
+	aggName := strings.ToLower(plan.AliasSubqueryString(agg))
+	col := scopeColumn{col: aggName, scalar: agg, typ: agg.Type(), nullable: agg.IsNullable()}
+	id := gb.outScope.newColumn(col)
+
+	agg = agg.WithId(sql.ColumnId(id)).(*aggregation.GroupConcat)
+	gb.outScope.cols[len(gb.outScope.cols)-1].scalar = agg
+	col.scalar = agg
+
+	gb.addAggStr(col)
+	col.id = id
+	return col.scalarGf()
+}
+
 func isWindowFunc(name string) bool {
 	switch name {
 	case "first", "last", "count", "sum", "any_value",

sql/planbuilder/scalar.go

@@ -206,6 +206,9 @@ func (b *Builder) buildScalar(inScope *scope, e ast.Expr) (ex sql.Expression) {
 	case *ast.GroupConcatExpr:
 		// TODO this is an aggregation
 		return b.buildGroupConcat(inScope, v)
+	case *ast.OrderedInjectedExpr:
+		// TODO this is an aggregation
+		return b.buildOrderedInjectedExpr(inScope, v)
 	case *ast.ParenExpr:
 		return b.buildScalar(inScope, v.Expr)
 	case *ast.AndExpr:
@@ -272,23 +275,7 @@ func (b *Builder) buildScalar(inScope *scope, e ast.Expr) (ex sql.Expression) {
 		}
 		return ret
 	case ast.InjectedExpr:
-		if err := b.cat.AuthorizationHandler().HandleAuth(b.ctx, b.authQueryState, v.Auth); err != nil && b.authEnabled {
-			b.handleErr(err)
-		}
-		resolvedChildren := make([]any, len(v.Children))
-		for i, child := range v.Children {
-			resolvedChildren[i] = b.buildScalar(inScope, child)
-		}
-		expr, err := v.Expression.WithResolvedChildren(resolvedChildren)
-		if err != nil {
-			b.handleErr(err)
-			return nil
-		}
-		if sqlExpr, ok := expr.(sql.Expression); ok {
-			return sqlExpr
-		}
-		b.handleErr(fmt.Errorf("Injected expression does not resolve to a valid expression"))
-		return nil
+		return b.buildInjectedExpr(inScope, v)
 	case *ast.RangeCond:
 		val := b.buildScalar(inScope, v.Left)
 		lower := b.buildScalar(inScope, v.From)
@@ -422,6 +409,30 @@ func (b *Builder) buildScalar(inScope *scope, e ast.Expr) (ex sql.Expression) {
 	return nil
 }
 
+func (b *Builder) buildInjectedExpr(inScope *scope, v ast.InjectedExpr) sql.Expression {
+	if err := b.cat.AuthorizationHandler().HandleAuth(b.ctx, b.authQueryState, v.Auth); err != nil && b.authEnabled {
+		b.handleErr(err)
+	}
+	resolvedChildren := make([]any, len(v.Children))
+	for i, child := range v.Children {
+		resolvedChildren[i] = b.buildScalar(inScope, child)
+	}
+	return b.buildInjectedExpressionFromResolvedChildren(v, resolvedChildren)
+}
+
+func (b *Builder) buildInjectedExpressionFromResolvedChildren(v ast.InjectedExpr, resolvedChildren []any) sql.Expression {
+	expr, err := v.Expression.WithResolvedChildren(resolvedChildren)
+	if err != nil {
+		b.handleErr(err)
+		return nil
+	}
+	if sqlExpr, ok := expr.(sql.Expression); ok {
+		return sqlExpr
+	}
+	b.handleErr(fmt.Errorf("injected expression should resolve to sql.Expression, got %T", expr))
+	return nil
+}
+
 func (b *Builder) getOrigTblName(node sql.Node, alias string) string {
 	if node == nil {
 		return ""