Properly resolve RETURNING columns

sqlparser/CHANGELOG.md

@@ -1,6 +1,7 @@
 ## 0.15.1-dev
 
 - New analysis checks for `RETURNING`: Disallow `table.*` syntax and aggregate expressions
+- Fix resolving columns when `RETURNING` is used in an `UPDATE FROM` statement
 
 ## 0.15.0
 

sqlparser/lib/src/analysis/schema/references.dart

@@ -48,12 +48,23 @@ class ReferenceScope {
   }
 
   set availableColumns(List<Column>? value) {
-    _availableColumns = value;
+    // guard against lists of subtype of column
+    if (value != null) {
+      _availableColumns = <Column>[...value];
+    } else {
+      _availableColumns = null;
+    }
   }
 
   ReferenceScope(this.parent,
       {this.root, this.inheritAvailableColumns = false});
 
+  void addAvailableColumn(Column column) {
+    // make sure _availableColumns is resolved and mutable
+    final ownColumns = _availableColumns ??= <Column>[...availableColumns];
+    ownColumns.add(column);
+  }
+
   ReferenceScope createChild({bool? inheritAvailableColumns}) {
     // wonder why we're creating a linked list of reference scopes instead of
     // just passing down a copy of [_references]? In sql, some variables can be

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

@@ -90,35 +90,57 @@ class ColumnResolver extends RecursiveVisitor<void, void> {
       if (child != e.table && child != e.from) visit(child, arg);
     }
 
-    _resolveReturningClause(e);
+    _resolveReturningClause(e, baseTable);
   }
 
-  void _addIfResolved(AstNode node, TableReference ref) {
+  ResultSet? _addIfResolved(AstNode node, TableReference ref) {
     final table = _resolveTableReference(ref);
     if (table != null) {
       node.scope.availableColumns = table.resolvedColumns;
     }
+
+    return table;
   }
 
   @override
   void visitInsertStatement(InsertStatement e, void arg) {
-    _addIfResolved(e, e.table);
+    final into = _addIfResolved(e, e.table);
     visitChildren(e, arg);
-    _resolveReturningClause(e);
+    _resolveReturningClause(e, into);
   }
 
   @override
   void visitDeleteStatement(DeleteStatement e, void arg) {
-    _addIfResolved(e, e.from!);
+    final from = _addIfResolved(e, e.from);
     visitChildren(e, arg);
-    _resolveReturningClause(e);
+    _resolveReturningClause(e, from);
   }
 
-  void _resolveReturningClause(StatementReturningColumns stmt) {
+  /// Infers the result set of a `RETURNING` clause.
+  ///
+  /// The behavior of `RETURNING` clauses is a bit weird when there are multiple
+  /// tables available (which can happen with `UPDATE FROM`). When a star column
+  /// is used, it only expands to columns from the main table:
+  /// ```sql
+  /// CREATE TABLE x (a, b);
+  /// -- here, the `*` in returning does not include columns from `old`.
+  /// UPDATE x SET a = x.a + 1 FROM (SELECT * FROM x) AS old RETURNING *;
+  /// ```
+  ///
+  /// However, individual columns from other tables are available and supported:
+  /// ```sql
+  /// UPDATE x SET a = x.a + 1 FROM (SELECT * FROM x) AS old
+  ///   RETURNING old.a, old.b;
+  /// ```
+  ///
+  /// Note that `old.*` is forbidden by sqlite and not applicable here.
+  void _resolveReturningClause(
+      StatementReturningColumns stmt, ResultSet? mainTable) {
     final clause = stmt.returning;
     if (clause == null) return;
 
-    final columns = _resolveColumns(stmt.scope, clause.columns);
+    final columns = _resolveColumns(stmt.scope, clause.columns,
+        columnsForStar: mainTable?.resolvedColumns);
     stmt.returnedResultSet = CustomResultSet(columns);
   }
 
@@ -210,10 +232,10 @@ class ColumnResolver extends RecursiveVisitor<void, void> {
     s.resolvedColumns = _resolveColumns(scope, s.columns);
   }
 
-  List<Column> _resolveColumns(
-      ReferenceScope scope, List<ResultColumn> columns) {
+  List<Column> _resolveColumns(ReferenceScope scope, List<ResultColumn> columns,
+      {List<Column>? columnsForStar}) {
     final usedColumns = <Column>[];
-    final availableColumns = scope.availableColumns;
+    final availableColumns = <Column>[...scope.availableColumns];
 
     // a select statement can include everything from its sub queries as a
     // result, but also expressions that appear as result columns
@@ -234,9 +256,9 @@ class ColumnResolver extends RecursiveVisitor<void, void> {
 
           visibleColumnsForStar = tableResolver.resultSet!.resolvedColumns;
         } else {
-          // we have a * column without a table, that resolves to every columns
+          // we have a * column without a table, that resolves to every column
           // available
-          visibleColumnsForStar = availableColumns;
+          visibleColumnsForStar = columnsForStar ?? availableColumns;
         }
 
         usedColumns
@@ -262,6 +284,7 @@ class ColumnResolver extends RecursiveVisitor<void, void> {
           final name = resultColumn.as;
           if (!availableColumns.any((c) => c.name == name)) {
             availableColumns.add(column);
+            scope.addAvailableColumn(column);
           }
         }
       } else if (resultColumn is NestedStarResultColumn) {

sqlparser/lib/src/ast/statements/delete.dart

@@ -7,12 +7,12 @@ import 'statement.dart';
 
 class DeleteStatement extends CrudStatement
     implements StatementWithWhere, StatementReturningColumns, HasPrimarySource {
-  TableReference? from;
+  TableReference from;
   @override
   Expression? where;
 
   @override
-  TableReference? get table => from;
+  TableReference get table => from;
 
   @override
   Returning? returning;
@@ -31,7 +31,7 @@ class DeleteStatement extends CrudStatement
   @override
   void transformChildren<A>(Transformer<A> transformer, A arg) {
     withClause = transformer.transformNullableChild(withClause, this, arg);
-    from = transformer.transformChild(from!, this, arg);
+    from = transformer.transformChild(from, this, arg);
     where = transformer.transformNullableChild(where, this, arg);
     returning = transformer.transformNullableChild(returning, this, arg);
   }
@@ -39,7 +39,7 @@ class DeleteStatement extends CrudStatement
   @override
   Iterable<AstNode> get childNodes => [
         if (withClause != null) withClause!,
-        from!,
+        from,
         if (where != null) where!,
         if (returning != null) returning!,
       ];

sqlparser/test/analysis/column_resolver_test.dart

@@ -4,7 +4,7 @@ import 'package:test/test.dart';
 import 'data.dart';
 
 void main() {
-  final engine = SqlEngine();
+  final engine = SqlEngine(EngineOptions(version: SqliteVersion.v3_35));
   engine.registerTable(demoTable);
 
   group('CREATE TRIGGER statements', () {
@@ -158,12 +158,42 @@ INSERT INTO demo VALUES (?, ?)
     expect(result.errors, isEmpty);
   });
 
-  test('resolves RETURNING clause', () {
-    final result =
-        engine.analyze("INSERT INTO demo (content) VALUES ('hi') RETURNING *;");
-    final returning = (result.root as InsertStatement).returnedResultSet;
+  group('resolves RETURNING clause', () {
+    test('for simple inserts', () {
+      final result = engine
+          .analyze("INSERT INTO demo (content) VALUES ('hi') RETURNING *;");
+      final returning = (result.root as InsertStatement).returnedResultSet;
 
-    expect(returning, isNotNull);
-    expect(returning!.resolvedColumns!.map((e) => e.name), ['id', 'content']);
+      expect(returning, isNotNull);
+      expect(returning!.resolvedColumns!.map((e) => e.name), ['id', 'content']);
+    });
+
+    test('for custom expressions', () {
+      final result = engine.analyze("INSERT INTO demo (content) VALUES ('hi') "
+          'RETURNING content || content AS x;');
+      final returning = (result.root as InsertStatement).returnedResultSet!;
+
+      expect(returning.resolvedColumns!.map((e) => e.name), ['x']);
+    });
+
+    test('star does not include other tables', () {
+      final result = engine.analyze('''
+        UPDATE demo SET content = ''
+          FROM (SELECT * FROM demo) AS old
+          RETURNING *;
+      ''');
+      final returning = (result.root as UpdateStatement).returnedResultSet!;
+      expect(returning.resolvedColumns!.map((e) => e.name), ['id', 'content']);
+    });
+
+    test('can refer to columns from other tables', () {
+      final result = engine.analyze('''
+        UPDATE demo SET content = ''
+          FROM (SELECT * FROM demo) AS old
+          RETURNING old.id, old.content;
+      ''');
+
+      expect(result.errors, isEmpty);
+    });
   });
 }