Generate code for update and delete statements

moor/test/data/tables/todos.dart

@@ -58,6 +58,7 @@ class TableWithoutPK extends Table {
   queries: {
     'allTodosWithCategory': 'SELECT t.*, c.id as catId, c."desc" as catDesc '
         'FROM todos t INNER JOIN categories c ON c.id = t.category',
+    'deleteTodoById': 'DELETE FROM todos WHERE id = ?'
   },
 )
 class TodoDb extends _$TodoDb {

moor/test/data/tables/todos.g.dart

@@ -1068,6 +1068,13 @@ abstract class _$TodoDb extends GeneratedDatabase {
         }).map((rows) => rows.map(_rowToAllTodosWithCategoryResult).toList());
   }
 
+  Future<int> deleteTodoById(int var1, {QueryEngine operateOn}) {
+    return (operateOn ?? this)
+        .customUpdate('DELETE FROM todos WHERE id = ?', variables: [
+      Variable.withInt(var1),
+    ], updates: {});
+  }
+
   @override
   List<TableInfo> get allTables =>
       [todosTable, categories, users, sharedTodos, tableWithoutPK];

moor_generator/lib/src/model/sql_query.dart

@@ -29,6 +29,14 @@ class SqlSelectQuery extends SqlQuery {
       : super(name, sql, variables);
 }
 
+class UpdatingQuery extends SqlQuery {
+  final List<SpecifiedTable> updates;
+
+  UpdatingQuery(
+      String name, String sql, List<FoundVariable> variables, this.updates)
+      : super(name, sql, variables);
+}
+
 class InferredResultSet {
   /// If the result columns of a SELECT statement exactly match one table, we
   /// can just use the data class generated for that table. Otherwise, we'd have

moor_generator/lib/src/parser/sql/affected_tables_visitor.dart

@@ -2,7 +2,7 @@ import 'package:sqlparser/sqlparser.dart';
 
 /// An AST-visitor that walks sql statements and finds all tables referenced in
 /// them.
-class AffectedTablesVisitor extends RecursiveVisitor<void> {
+class ReferencedTablesVisitor extends RecursiveVisitor<void> {
   final Set<Table> foundTables = {};
 
   @override
@@ -27,3 +27,29 @@ class AffectedTablesVisitor extends RecursiveVisitor<void> {
     visitChildren(e);
   }
 }
+
+/// Finds all tables that could be affected when executing a query. In
+/// contrast to [ReferencedTablesVisitor], which finds all references, this
+/// visitor only collects tables a query writes to.
+class UpdatedTablesVisitor extends RecursiveVisitor<void> {
+  final Set<Table> foundTables = {};
+
+  void _addIfResolved(ResolvesToResultSet r) {
+    final resolved = r.resultSet;
+    if (resolved is Table) {
+      foundTables.add(resolved);
+    }
+  }
+
+  @override
+  void visitDeleteStatement(DeleteStatement e) {
+    _addIfResolved(e.from);
+    visitChildren(e);
+  }
+
+  @override
+  void visitUpdateStatement(UpdateStatement e) {
+    _addIfResolved(e.table);
+    visitChildren(e);
+  }
+}

moor_generator/lib/src/parser/sql/query_handler.dart

@@ -22,14 +22,25 @@ class QueryHandler {
 
     if (root is SelectStatement) {
       return _handleSelect();
+    } else if (root is UpdateStatement || root is DeleteStatement) {
+      return _handleUpdate();
     } else {
       throw StateError(
           'Unexpected sql: Got $root, expected a select statement');
     }
   }
 
+  UpdatingQuery _handleUpdate() {
+    final updatedFinder = UpdatedTablesVisitor();
+    context.root.accept(updatedFinder);
+    _foundTables = updatedFinder.foundTables;
+
+    return UpdatingQuery(name, context.sql, _foundVariables,
+        _foundTables.map(mapper.tableToMoor).toList());
+  }
+
   SqlSelectQuery _handleSelect() {
-    final tableFinder = AffectedTablesVisitor();
+    final tableFinder = ReferencedTablesVisitor();
     _select.accept(tableFinder);
     _foundTables = tableFinder.foundTables;
     final moorTables = _foundTables.map(mapper.tableToMoor).toList();

moor_generator/lib/src/shared_state.dart

@@ -43,6 +43,7 @@ class SharedState {
           message: 'The type $type is not a moor table',
           affectedElement: initializedBy,
         ));
+        return null;
       } else {
         return tableParser.parse(type.element as ClassElement);
       }

moor_generator/lib/src/writer/query_writer.dart

@@ -8,12 +8,15 @@ import 'package:recase/recase.dart';
 class QueryWriter {
   final SqlQuery query;
   SqlSelectQuery get _select => query as SqlSelectQuery;
+  UpdatingQuery get _update => query as UpdatingQuery;
 
   QueryWriter(this.query);
 
   void writeInto(StringBuffer buffer) {
     if (query is SqlSelectQuery) {
       _writeSelect(buffer);
+    } else if (query is UpdatingQuery) {
+      _writeUpdatingQuery(buffer);
     }
   }
 
@@ -81,6 +84,26 @@ class QueryWriter {
       ..write('\n}\n');
   }
 
+  void _writeUpdatingQuery(StringBuffer buffer) {
+    /*
+      Future<int> test() {
+    return customUpdate('', variables: [], updates: {});
+  }
+     */
+    buffer.write('Future<int> ${query.name}(');
+    _writeParameters(buffer);
+    buffer
+      ..write(') {\n')
+      ..write('return (operateOn ?? this).')
+      ..write('customUpdate(${asDartLiteral(query.sql)},');
+
+    _writeVariables(buffer);
+    buffer.write(',');
+    _writeUpdates(buffer);
+
+    buffer..write(');\n}\n');
+  }
+
   void _writeParameters(StringBuffer buffer,
       {bool dontOverrideEngine = false}) {
     final paramList = query.variables
@@ -113,4 +136,9 @@ class QueryWriter {
     final from = _select.readsFrom.map((t) => t.tableFieldName).join(', ');
     buffer..write('readsFrom: {')..write(from)..write('}');
   }
+
+  void _writeUpdates(StringBuffer buffer) {
+    final from = _update.updates.map((t) => t.tableFieldName).join(', ');
+    buffer..write('updates: {')..write(from)..write('}');
+  }
 }