Allow references to result column in group by too

sqlparser/lib/src/analysis/steps/reference_resolver.dart

@@ -1,13 +1,14 @@
 part of '../analysis.dart';
 
 /// Resolves any open [Reference] it finds in the AST.
-class ReferenceResolver extends RecursiveVisitor<void, void> {
+class ReferenceResolver
+    extends RecursiveVisitor<ReferenceResolvingContext, void> {
   final AnalysisContext context;
 
   ReferenceResolver(this.context);
 
   @override
-  void visitInsertStatement(InsertStatement e, void arg) {
+  void visitInsertStatement(InsertStatement e, ReferenceResolvingContext arg) {
     final table = e.table.resultSet;
     if (table != null) {
       // Resolve columns in the main table
@@ -20,7 +21,8 @@ class ReferenceResolver extends RecursiveVisitor<void, void> {
   }
 
   @override
-  void visitForeignKeyClause(ForeignKeyClause e, void arg) {
+  void visitForeignKeyClause(
+      ForeignKeyClause e, ReferenceResolvingContext arg) {
     final table = e.foreignTable.resultSet;
     if (table == null) {
       // If the table wasn't found, an earlier step will have reported an error
@@ -33,7 +35,19 @@ class ReferenceResolver extends RecursiveVisitor<void, void> {
   }
 
   @override
-  void visitReference(Reference e, void arg) {
+  void visitGroupBy(GroupBy e, ReferenceResolvingContext arg) {
+    return super.visitGroupBy(
+        e, const ReferenceResolvingContext(canReferToColumnAlias: true));
+  }
+
+  @override
+  void visitOrderBy(OrderBy e, ReferenceResolvingContext arg) {
+    return super.visitOrderBy(
+        e, ReferenceResolvingContext(canReferToColumnAlias: true));
+  }
+
+  @override
+  void visitReference(Reference e, ReferenceResolvingContext arg) {
     if (e.resolved != null) {
       return super.visitReference(e, arg);
     }
@@ -69,9 +83,7 @@ class ReferenceResolver extends RecursiveVisitor<void, void> {
       // todo special case for USING (...) in joins?
       final found = scope.resolveUnqualifiedReference(
         e.columnName,
-        // According to https://www.sqlite.org/lang_select.html#the_order_by_clause,
-        // a simple reference in an ordering term can refer to an output column.
-        allowReferenceToResultColumn: e.parent is OrderingTerm,
+        allowReferenceToResultColumn: arg.canReferToColumnAlias,
       );
 
       if (found.isEmpty) {
@@ -96,7 +108,7 @@ class ReferenceResolver extends RecursiveVisitor<void, void> {
   }
 
   @override
-  void visitUpdateStatement(UpdateStatement e, void arg) {
+  void visitUpdateStatement(UpdateStatement e, ReferenceResolvingContext arg) {
     final table = e.table.resultSet;
     if (table != null) {
       // Resolve the set components against the primary table
@@ -109,7 +121,8 @@ class ReferenceResolver extends RecursiveVisitor<void, void> {
   }
 
   @override
-  void visitWindowFunctionInvocation(WindowFunctionInvocation e, void arg) {
+  void visitWindowFunctionInvocation(
+      WindowFunctionInvocation e, ReferenceResolvingContext arg) {
     if (e.windowName != null && e.resolved == null) {
       e.resolved =
           StatementScope.cast(e.scope).windowDeclarations[e.windowName!];
@@ -166,3 +179,15 @@ class ReferenceResolver extends RecursiveVisitor<void, void> {
     return resolved.findColumn(e.columnName);
   }
 }
+
+class ReferenceResolvingContext {
+  /// Whether unqualified references can be resolved against a column alias from
+  /// the surrounding select statement.
+  ///
+  /// This is only the case for `ORDER BY`, `GROUP BY` and `HAVING` clauses.
+  final bool canReferToColumnAlias;
+
+  const ReferenceResolvingContext({
+    this.canReferToColumnAlias = false,
+  });
+}

sqlparser/lib/src/engine/sql_engine.dart

@@ -275,7 +275,7 @@ class SqlEngine {
 
       node
         ..acceptWithoutArg(ColumnResolver(context))
-        ..acceptWithoutArg(ReferenceResolver(context));
+        ..accept(ReferenceResolver(context), const ReferenceResolvingContext());
 
       final session = TypeInferenceSession(context, options);
       final resolver = TypeResolver(session);

sqlparser/test/analysis/reference_resolver_test.dart

@@ -96,6 +96,18 @@ void main() {
     );
   });
 
+  test('allows references to result column in group by', () {
+    // https://github.com/simolus3/drift/issues/2378
+    final engine = SqlEngine()
+      ..registerTableFromSql('CREATE TABLE foo (bar INTEGER);');
+
+    final result = engine.analyze('''
+      SELECT *, bar > 20 AS test FROM foo GROUP BY bar HAVING test
+''');
+
+    expect(result.errors, isEmpty);
+  });
+
   test('does not allow references to result column outside of ORDER BY', () {
     final engine = SqlEngine()..registerTable(demoTable);