Catch failures

Author: kanoshiou

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/aggregate/AggregateFunction.java

@@ -19,8 +19,8 @@
 import org.elasticsearch.xpack.esql.core.tree.Source;
 import org.elasticsearch.xpack.esql.core.util.CollectionUtils;
 import org.elasticsearch.xpack.esql.io.stream.PlanStreamInput;
+import org.elasticsearch.xpack.esql.plan.logical.Aggregate;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
-import org.elasticsearch.xpack.esql.plan.logical.OrderBy;
 
 import java.io.IOException;
 import java.util.List;
@@ -139,14 +139,14 @@ public boolean equals(Object obj) {
     @Override
     public BiConsumer<LogicalPlan, Failures> postAnalysisPlanVerification() {
         return (p, failures) -> {
-            if (p instanceof OrderBy order) {
-                order.order().forEach(o -> {
-                    o.forEachDown(Function.class, f -> {
-                        if (f instanceof AggregateFunction) {
-                            failures.add(fail(f, "Aggregate functions are not allowed in SORT [{}]", f.functionName()));
-                        }
-                    });
-                });
+            if (p instanceof Aggregate == false) {
+                p.expressions().forEach(x -> x.forEachDown(AggregateFunction.class, af -> {
+                    if (af instanceof Rate) {
+                        failures.add(fail(af, "aggregate function [{}] not allowed outside METRICS command", af.sourceText()));
+                    } else {
+                        failures.add(fail(af, "aggregate function [{}] not allowed outside STATS command", af.sourceText()));
+                    }
+                }));
             }
         };
     }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/Eval.java

@@ -24,8 +24,6 @@
 import org.elasticsearch.xpack.esql.core.tree.NodeInfo;
 import org.elasticsearch.xpack.esql.core.tree.Source;
 import org.elasticsearch.xpack.esql.core.type.DataType;
-import org.elasticsearch.xpack.esql.expression.function.aggregate.AggregateFunction;
-import org.elasticsearch.xpack.esql.expression.function.aggregate.Rate;
 import org.elasticsearch.xpack.esql.io.stream.PlanStreamInput;
 import org.elasticsearch.xpack.esql.plan.GeneratingPlan;
 
@@ -179,14 +177,6 @@ public void postAnalysisVerification(Failures failures) {
                     )
                 );
             }
-            // check no aggregate functions are used
-            field.forEachDown(AggregateFunction.class, af -> {
-                if (af instanceof Rate) {
-                    failures.add(fail(af, "aggregate function [{}] not allowed outside METRICS command", af.sourceText()));
-                } else {
-                    failures.add(fail(af, "aggregate function [{}] not allowed outside STATS command", af.sourceText()));
-                }
-            });
         });
     }
 }