Analysis for FROM clauses in UPDATE statements

sqlparser/CHANGELOG.md

@@ -1,5 +1,6 @@
 ## 0.14.1-dev
 
+- Support `FROM` clauses in `UPDATE` statements
 - Fix `rank` columns of fts5 tables being misreported as integers
 
 ## 0.14.0

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

@@ -70,6 +70,27 @@ class ColumnResolver extends RecursiveVisitor<void, void> {
     }
   }
 
+  @override
+  void visitUpdateStatement(UpdateStatement e, void arg) {
+    final availableColumns = <Column>[];
+
+    // Add columns from the main table, if it was resolved
+    final baseTable = _resolveTableReference(e.table);
+    if (baseTable != null) {
+      availableColumns.addAll(baseTable.resolvedColumns ?? const []);
+    }
+
+    // Also add columns from a FROM clause, if one is present
+    final from = e.from;
+    if (from != null) _handle(from, availableColumns);
+
+    e.scope.availableColumns = availableColumns;
+    for (final child in e.childNodes) {
+      // Visit remaining children
+      if (child != e.table && child != e.from) visit(child, arg);
+    }
+  }
+
   void _addIfResolved(AstNode node, TableReference ref) {
     final table = _resolveTableReference(ref);
     if (table != null) {
@@ -77,12 +98,6 @@ class ColumnResolver extends RecursiveVisitor<void, void> {
     }
   }
 
-  @override
-  void visitUpdateStatement(UpdateStatement e, void arg) {
-    _addIfResolved(e, e.table);
-    visitExcept(e, e.table, arg);
-  }
-
   @override
   void visitInsertStatement(InsertStatement e, void arg) {
     _addIfResolved(e, e.table!);

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

@@ -58,7 +58,7 @@ class AstPreparingVisitor extends RecursiveVisitor<void, void> {
 
     if (isInFROM) {
       final surroundingSelect =
-          e.parents.firstWhere((node) => node is BaseSelectStatement).scope;
+          e.parents.firstWhere((node) => node is HasFrom).scope;
       final forked = surroundingSelect.createSibling();
       e.scope = forked;
     } else {

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

@@ -14,9 +14,14 @@ abstract class BaseSelectStatement extends CrudStatement with ResultSet {
 abstract class SelectStatementNoCompound implements BaseSelectStatement {}
 
 class SelectStatement extends BaseSelectStatement
-    implements StatementWithWhere, SelectStatementNoCompound, HasPrimarySource {
+    implements
+        StatementWithWhere,
+        SelectStatementNoCompound,
+        HasPrimarySource,
+        HasFrom {
   final bool distinct;
   final List<ResultColumn> columns;
+  @override
   Queryable? from;
 
   @override

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

@@ -24,6 +24,14 @@ abstract class HasPrimarySource extends Statement {
   Queryable? get table;
 }
 
+/// Interfaces for statements that have a `FROM` clause.
+///
+/// This includes selects and, since recently, updates.
+abstract class HasFrom extends Statement {
+  /// The table, join clause or subquery appearing in the `FROM` clause.
+  Queryable? get from;
+}
+
 /// Interface for statements that have a primary where clause (select, update,
 /// delete).
 abstract class StatementWithWhere extends Statement implements HasWhereClause {}

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

@@ -17,11 +17,12 @@ const Map<TokenType, FailureMode> _tokensToMode = {
 };
 
 class UpdateStatement extends CrudStatement
-    implements StatementWithWhere, HasPrimarySource {
+    implements StatementWithWhere, HasPrimarySource, HasFrom {
   final FailureMode? or;
   @override
   TableReference table;
   final List<SetComponent> set;
+  @override
   Queryable? from;
   @override
   Expression? where;

sqlparser/test/analysis/column_resolver_test.dart

@@ -132,4 +132,29 @@ INSERT INTO demo VALUES (?, ?)
     expect(columns.map((e) => context.typeOf(e).type?.type),
         [BasicType.text, BasicType.int]);
   });
+
+  test('handles update statement with from clause', () {
+    // Example from here: https://www.sqlite.org/lang_update.html#upfrom
+    engine..registerTableFromSql('''
+      CREATE TABLE inventory (
+        itemId INTEGER PRIMARY KEY,
+        quantity INTEGER NOT NULL DEFAULT 0,
+      );
+    ''')..registerTableFromSql('''
+      CREATE TABLE sales (
+        itemId INTEGER NOT NULL,
+        quantity INTEGER NOT NULL,
+      );
+    ''');
+
+    final result = engine.analyze('''
+      UPDATE inventory
+        SET quantity = quantity - daily.amt
+        FROM (SELECT sum(quantity) AS amt, itemId FROM sales GROUP BY 2) 
+          AS daily
+        WHERE inventory.itemId = daily.itemId;
+    ''');
+
+    expect(result.errors, isEmpty);
+  });
 }

sqlparser/test/analysis/data.dart

@@ -26,3 +26,12 @@ final Table anotherTable = Table(
   name: 'tbl',
   resolvedColumns: [anotherId, dateTime],
 );
+
+extension RegisterTableExtension on SqlEngine {
+  /// Utility function that parses a `CREATE TABLE` statement and registers the
+  /// created table to the engine.
+  void registerTableFromSql(String createTable) {
+    final stmt = parse(createTable).rootNode as CreateTableStatement;
+    registerTable(schemaReader.read(stmt));
+  }
+}