Initial support for RETURNING in the generator

moor/build.yaml

@@ -14,7 +14,7 @@ targets:
           named_parameters: true
           new_sql_code_generation: true
           sqlite:
-            version: "3.34"
+            version: "3.35"
             modules:
               - json1
               - fts5

moor/lib/src/runtime/api/connection_user.dart

@@ -260,6 +260,27 @@ abstract class DatabaseConnectionUser {
     );
   }
 
+  /// Runs a `INSERT`, `UPDATE` or `DELETE` statement returning rows.
+  ///
+  /// You can use the [updates] parameter so that moor knows which tables are
+  /// affected by your query. All select streams that depend on a table
+  /// specified there will then update their data. For more accurate results,
+  /// you can also set the [updateKind] parameter.
+  /// This is optional, but can improve the accuracy of query updates,
+  /// especially when using triggers.
+  Future<List<QueryRow>> customWriteReturning(
+    String query, {
+    List<Variable> variables = const [],
+    Set<TableInfo>? updates,
+    UpdateKind? updateKind,
+  }) {
+    return _customWrite(query, variables, updates, updateKind,
+        (executor, sql, vars) async {
+      final rows = await executor.runSelect(sql, vars);
+      return [for (final row in rows) QueryRow(row, attachedDatabase)];
+    });
+  }
+
   /// Common logic for [customUpdate] and [customInsert] which takes care of
   /// mapping the variables, running the query and optionally informing the
   /// stream-queries.

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

@@ -1729,6 +1729,21 @@ abstract class _$CustomTablesDb extends GeneratedDatabase {
         readsFrom: {config}).map((QueryRow row) => row.read<int?>('MAX(oid)'));
   }
 
+  Future<List<Config>> addConfig(
+      String var1, String? var2, SyncType? var3, SyncType? var4) {
+    return customWriteReturning(
+        'INSERT INTO config VALUES (?, ?, ?, ?) RETURNING *',
+        variables: [
+          Variable<String>(var1),
+          Variable<String?>(var2),
+          Variable<int?>(ConfigTable.$converter0.mapToSql(var3)),
+          Variable<int?>(ConfigTable.$converter1.mapToSql(var4))
+        ],
+        updates: {
+          config
+        }).then((rows) => rows.map(config.mapFromRow).toList());
+  }
+
   Future<int> writeConfig({required String key, String? value}) {
     return customInsert(
       'REPLACE INTO config (config_key, config_value) VALUES (:key, :value)',

moor/test/data/tables/tables.moor

@@ -90,3 +90,4 @@ cfeTest: WITH RECURSIVE
   SELECT x FROM cnt;
 
 nullableQuery: SELECT MAX(oid) FROM config;
+addConfig: INSERT INTO config VALUES (?, ?, ?, ?) RETURNING *;

moor/test/integration_tests/moor_files_integration_test.dart

@@ -1,5 +1,6 @@
 import 'package:moor/ffi.dart';
 import 'package:moor/src/runtime/query_builder/query_builder.dart' hide isNull;
+import 'package:sqlite3/sqlite3.dart';
 import 'package:test/test.dart';
 
 import '../data/tables/converter.dart';
@@ -60,4 +61,25 @@ void main() {
       await expectation;
     });
   });
+
+  final sqliteVersion = sqlite3.version;
+  final hasReturning = sqliteVersion.versionNumber > 3035000;
+
+  group('returning', () {
+    test('for custom inserts', () async {
+      final result = await db.addConfig(
+          'key2', 'val', SyncType.locallyCreated, SyncType.locallyCreated);
+
+      expect(result, hasLength(1));
+      expect(
+        result.single,
+        Config(
+          configKey: 'key2',
+          configValue: 'val',
+          syncState: SyncType.locallyCreated,
+          syncStateImplicit: SyncType.locallyCreated,
+        ),
+      );
+    });
+  }, skip: hasReturning ? null : 'RETURNING not supported by current sqlite');
 }

moor_generator/lib/src/analyzer/sql_queries/lints/linter.dart

@@ -30,17 +30,6 @@ class _LintingVisitor extends RecursiveVisitor<void, void> {
 
   _LintingVisitor(this.linter);
 
-  void _checkNoReturning(StatementReturningColumns stmt) {
-    if (stmt.returning != null) {
-      linter.lints.add(AnalysisError(
-        type: AnalysisErrorType.other,
-        message: 'RETURNING is not supported in this version of moor. '
-            'Follow https://github.com/simolus3/moor/issues/1096 for updates.',
-        relevantNode: stmt.returning,
-      ));
-    }
-  }
-
   @override
   void visitBinaryExpression(BinaryExpression e, void arg) {
     const numericOps = {
@@ -162,15 +151,8 @@ class _LintingVisitor extends RecursiveVisitor<void, void> {
     }
   }
 
-  @override
-  void visitDeleteStatement(DeleteStatement e, void arg) {
-    _checkNoReturning(e);
-    visitChildren(e, arg);
-  }
-
   @override
   void visitInsertStatement(InsertStatement e, void arg) {
-    _checkNoReturning(e);
     final targeted = e.resolvedTargetColumns;
     if (targeted == null) return;
 
@@ -239,10 +221,4 @@ class _LintingVisitor extends RecursiveVisitor<void, void> {
       }
     }
   }
-
-  @override
-  void visitUpdateStatement(UpdateStatement e, void arg) {
-    _checkNoReturning(e);
-    visitChildren(e, arg);
-  }
 }

moor_generator/lib/src/analyzer/sql_queries/query_handler.dart

@@ -20,6 +20,7 @@ class QueryHandler {
   Set<Table> _foundTables;
   Set<View> _foundViews;
   List<FoundElement> _foundElements;
+
   Iterable<FoundVariable> get _foundVariables =>
       _foundElements.whereType<FoundVariable>();
 
@@ -56,12 +57,26 @@ class QueryHandler {
     }
   }
 
+  void _applyFoundTables(ReferencedTablesVisitor visitor) {
+    _foundTables = visitor.foundTables;
+    _foundViews = visitor.foundViews;
+  }
+
   UpdatingQuery _handleUpdate() {
     final updatedFinder = UpdatedTablesVisitor();
     context.root.acceptWithoutArg(updatedFinder);
-    _foundTables = updatedFinder.writtenTables.map((w) => w.table).toSet();
+    _applyFoundTables(updatedFinder);
 
-    final isInsert = context.root is InsertStatement;
+    final root = context.root;
+    final isInsert = root is InsertStatement;
+
+    InferredResultSet resultSet;
+    if (root is StatementReturningColumns) {
+      final columns = root.returnedResultSet?.resolvedColumns;
+      if (columns != null) {
+        resultSet = _inferResultSet(columns);
+      }
+    }
 
     return UpdatingQuery(
       name,
@@ -70,14 +85,15 @@ class QueryHandler {
       updatedFinder.writtenTables.map(mapper.writtenToMoor).toList(),
       isInsert: isInsert,
       hasMultipleTables: updatedFinder.foundTables.length > 1,
+      resultSet: resultSet,
     );
   }
 
   SqlSelectQuery _handleSelect() {
     final tableFinder = ReferencedTablesVisitor();
     _select.acceptWithoutArg(tableFinder);
-    _foundTables = tableFinder.foundTables;
-    _foundViews = tableFinder.foundViews;
+    _applyFoundTables(tableFinder);
+
     final moorTables =
         _foundTables.map(mapper.tableToMoor).where((s) => s != null).toList();
     final moorViews =
@@ -95,15 +111,14 @@ class QueryHandler {
       context,
       _foundElements,
       moorEntities,
-      _inferResultSet(),
+      _inferResultSet(_select.resolvedColumns),
       requestedName,
     );
   }
 
-  InferredResultSet _inferResultSet() {
+  InferredResultSet _inferResultSet(List<Column> rawColumns) {
     final candidatesForSingleTable = Set.of(_foundTables);
     final columns = <ResultColumn>[];
-    final rawColumns = _select.resolvedColumns;
 
     // First, go through regular result columns
     for (final column in rawColumns) {
@@ -190,7 +205,7 @@ class QueryHandler {
   }
 
   List<NestedResultTable> _findNestedResultTables() {
-    final query = _select;
+    final query = context.root;
     // We don't currently support nested results for compound statements
     if (query is! SelectStatement) return const [];
 

moor_generator/lib/src/model/sql_query.dart

@@ -44,11 +44,12 @@ class DeclaredDartQuery extends DeclaredQuery {
 /// available.
 class DeclaredMoorQuery extends DeclaredQuery {
   final DeclaredStatement astNode;
-  CrudStatement get query => astNode.statement;
   ParsedMoorFile file;
 
   DeclaredMoorQuery(String name, this.astNode) : super(name);
 
+  CrudStatement get query => astNode.statement;
+
   factory DeclaredMoorQuery.fromStatement(DeclaredStatement stmt) {
     assert(stmt.identifier is SimpleName);
     final name = (stmt.identifier as SimpleName).name;
@@ -69,6 +70,12 @@ abstract class SqlQuery {
 
   String get sql => fromContext.sql;
 
+  /// The result set of this statement, mapped to moor-generated classes.
+  ///
+  /// This is non-nullable for select queries. Updating queries might have a
+  /// result set if they have a `RETURNING` clause.
+  InferredResultSet /*?*/ get resultSet;
+
   /// The variables that appear in the [sql] query. We support three kinds of
   /// sql variables: The regular "?" variables, explicitly indexed "?xyz"
   /// variables and colon-named variables. Even though this feature is not
@@ -108,13 +115,49 @@ abstract class SqlQuery {
     variables = elements.whereType<FoundVariable>().toList();
     placeholders = elements.whereType<FoundDartPlaceholder>().toList();
   }
+
+  String get resultClassName {
+    final resultSet = this.resultSet;
+    if (resultSet == null) {
+      throw StateError('This query ($name) does not have a result set');
+    }
+
+    if (resultSet.matchingTable != null || resultSet.singleColumn) {
+      throw UnsupportedError('This result set does not introduce a class, '
+          'either because it has a matching table or because it only returns '
+          'one column.');
+    }
+
+    return resultSet.resultClassName ?? '${ReCase(name).pascalCase}Result';
+  }
+
+  /// The Dart type representing a row of this result set.
+  String resultTypeCode(
+      [GenerationOptions options = const GenerationOptions()]) {
+    final resultSet = this.resultSet;
+    if (resultSet == null) {
+      throw StateError('This query ($name) does not have a result set');
+    }
+
+    if (resultSet.matchingTable != null) {
+      return resultSet.matchingTable.table.dartTypeName;
+    }
+
+    if (resultSet.singleColumn) {
+      return resultSet.columns.single.dartTypeCode(options);
+    }
+
+    return resultClassName;
+  }
 }
 
 class SqlSelectQuery extends SqlQuery {
   final List<MoorSchemaEntity> readsFrom;
+  @override
   final InferredResultSet resultSet;
 
   /// The name of the result class, as requested by the user.
+  // todo: Allow custom result classes for RETURNING as well?
   final String /*?*/ requestedResultClass;
 
   SqlSelectQuery(
@@ -150,42 +193,21 @@ class SqlSelectQuery extends SqlQuery {
       null,
     );
   }
-
-  String get resultClassName {
-    if (resultSet.matchingTable != null || resultSet.singleColumn) {
-      throw UnsupportedError('This result set does not introduce a class, '
-          'either because it has a matching table or because it only returns '
-          'one column.');
-    }
-
-    return resultSet.resultClassName ?? '${ReCase(name).pascalCase}Result';
-  }
-
-  /// The Dart type representing a row of this result set.
-  String resultTypeCode(
-      [GenerationOptions options = const GenerationOptions()]) {
-    if (resultSet.matchingTable != null) {
-      return resultSet.matchingTable.table.dartTypeName;
-    }
-
-    if (resultSet.singleColumn) {
-      return resultSet.columns.single.dartTypeCode(options);
-    }
-
-    return resultClassName;
-  }
 }
 
 class UpdatingQuery extends SqlQuery {
   final List<WrittenMoorTable> updates;
   final bool isInsert;
+  @override
+  final InferredResultSet /*?*/ resultSet;
 
   bool get isOnlyDelete => updates.every((w) => w.kind == UpdateKind.delete);
+
   bool get isOnlyUpdate => updates.every((w) => w.kind == UpdateKind.update);
 
   UpdatingQuery(String name, AnalysisContext fromContext,
       List<FoundElement> elements, this.updates,
-      {this.isInsert = false, bool hasMultipleTables})
+      {this.isInsert = false, bool hasMultipleTables, this.resultSet})
       : super(name, fromContext, elements,
             hasMultipleTables: hasMultipleTables);
 }

moor_generator/lib/src/writer/queries/query_writer.dart

@@ -38,15 +38,22 @@ class QueryWriter {
   }
 
   void write() {
+    // Note that writing queries can have a result set if they use a RETURNING
+    // clause.
+    final resultSet = query.resultSet;
+    if (resultSet?.needsOwnClass == true) {
+      final resultSetScope = scope.findScopeOfLevel(DartScope.library);
+      ResultSetWriter(query, resultSetScope).write();
+    }
+
     if (query is SqlSelectQuery) {
-      final select = query as SqlSelectQuery;
-      if (select.resultSet.needsOwnClass) {
-        final resultSetScope = scope.findScopeOfLevel(DartScope.library);
-        ResultSetWriter(select, resultSetScope).write();
-      }
       _writeSelect();
     } else if (query is UpdatingQuery) {
-      _writeUpdatingQuery();
+      if (resultSet != null) {
+        _writeUpdatingQueryWithReturning();
+      } else {
+        _writeUpdatingQuery();
+      }
     }
   }
 
@@ -68,17 +75,20 @@ class QueryWriter {
   }
 
   /// Writes the function literal that turns a "QueryRow" into the desired
-  /// custom return type of a select statement.
+  /// custom return type of a query.
   void _writeMappingLambda() {
-    if (_select.resultSet.singleColumn) {
-      final column = _select.resultSet.columns.single;
+    final resultSet = query.resultSet;
+    assert(resultSet != null);
+
+    if (resultSet.singleColumn) {
+      final column = resultSet.columns.single;
       _buffer.write('(QueryRow row) => '
           '${readingCode(column, scope.generationOptions)}');
-    } else if (_select.resultSet.matchingTable != null) {
+    } else if (resultSet.matchingTable != null) {
       // note that, even if the result set has a matching table, we can't just
       // use the mapFromRow() function of that table - the column names might
       // be different!
-      final match = _select.resultSet.matchingTable;
+      final match = resultSet.matchingTable;
       final table = match.table;
 
       if (match.effectivelyNoAlias) {
@@ -99,19 +109,19 @@ class QueryWriter {
         _buffer.write('})');
       }
     } else {
-      _buffer.write('(QueryRow row) { return ${_select.resultClassName}(');
+      _buffer.write('(QueryRow row) { return ${query.resultClassName}(');
 
       if (options.rawResultSetData) {
         _buffer.write('row: row,\n');
       }
 
-      for (final column in _select.resultSet.columns) {
-        final fieldName = _select.resultSet.dartNameFor(column);
+      for (final column in resultSet.columns) {
+        final fieldName = resultSet.dartNameFor(column);
         _buffer.write(
             '$fieldName: ${readingCode(column, scope.generationOptions)},');
       }
-      for (final nested in _select.resultSet.nestedResults) {
-        final prefix = _select.resultSet.nestedPrefixFor(nested);
+      for (final nested in resultSet.nestedResults) {
+        final prefix = resultSet.nestedPrefixFor(nested);
         if (prefix == null) continue;
 
         final fieldName = nested.dartFieldName;
@@ -207,6 +217,20 @@ class QueryWriter {
     _buffer.write(').watch();\n}\n');
   }
 
+  void _writeUpdatingQueryWithReturning() {
+    final type = query.resultTypeCode(scope.generationOptions);
+    _buffer.write('Future<List<$type>> ${query.name}(');
+    _writeParameters();
+    _buffer.write(') {\n');
+
+    _writeExpandedDeclarations();
+    _buffer.write('return customWriteReturning(${_queryCode()},');
+    _writeCommonUpdateParameters();
+    _buffer.write(').then((rows) => rows.map(');
+    _writeMappingLambda();
+    _buffer.write(').toList());\n}');
+  }
+
   void _writeUpdatingQuery() {
     /*
       Future<int> test() {
@@ -221,18 +245,15 @@ class QueryWriter {
 
     _writeExpandedDeclarations();
     _buffer.write('return $implName(${_queryCode()},');
+    _writeCommonUpdateParameters();
+    _buffer.write(',);\n}\n');
+  }
 
+  void _writeCommonUpdateParameters() {
     _writeVariables();
     _buffer.write(',');
     _writeUpdates();
-
-    if (_update.isOnlyDelete) {
-      _buffer.write(', updateKind: UpdateKind.delete');
-    } else if (_update.isOnlyUpdate) {
-      _buffer.write(', updateKind: UpdateKind.update');
-    }
-
-    _buffer.write(',);\n}\n');
+    _writeUpdateKind();
   }
 
   void _writeParameters() {
@@ -559,4 +580,12 @@ class QueryWriter {
     final from = _update.updates.map((t) => t.table.dbGetterName).join(', ');
     _buffer..write('updates: {')..write(from)..write('}');
   }
+
+  void _writeUpdateKind() {
+    if (_update.isOnlyDelete) {
+      _buffer.write(', updateKind: UpdateKind.delete');
+    } else if (_update.isOnlyUpdate) {
+      _buffer.write(', updateKind: UpdateKind.update');
+    }
+  }
 }

moor_generator/lib/src/writer/queries/result_set_writer.dart

@@ -6,7 +6,7 @@ import 'package:moor_generator/writer.dart';
 
 /// Writes a class holding the result of an sql query into Dart.
 class ResultSetWriter {
-  final SqlSelectQuery query;
+  final SqlQuery query;
   final Scope scope;
 
   ResultSetWriter(this.query, this.scope);

sqlparser/lib/utils/node_to_text.dart

@@ -505,6 +505,7 @@ class NodeSqlBuilder extends AstVisitor<void, void> {
     _keyword(TokenType.delete);
     _from(e.from);
     _where(e.where);
+    visitNullable(e.returning, arg);
   }
 
   @override
@@ -703,6 +704,7 @@ class NodeSqlBuilder extends AstVisitor<void, void> {
         InsertMode.insertOrFail: TokenType.fail,
         InsertMode.insertOrIgnore: TokenType.ignore,
       }[mode]!);
+      visitNullable(e.returning, arg);
     }
 
     _keyword(TokenType.into);
@@ -715,7 +717,7 @@ class NodeSqlBuilder extends AstVisitor<void, void> {
     }
 
     visit(e.source, arg);
-    visitNullable(e.upsert, arg);
+    visitNullable(e.returning, arg);
   }
 
   @override
@@ -1179,6 +1181,7 @@ class NodeSqlBuilder extends AstVisitor<void, void> {
     _join(e.set, ',');
     _from(e.from);
     _where(e.where);
+    visitNullable(e.returning, arg);
   }
 
   @override

sqlparser/test/utils/node_to_text_test.dart

@@ -223,9 +223,15 @@ CREATE UNIQUE INDEX my_idx ON t1 (c1, c2, c3) WHERE c1 < c3;
       });
     });
 
-    test('delete', () {
-      testFormat(
-          'WITH foo (id) AS (SELECT * FROM bar) DELETE FROM bar WHERE x;');
+    group('delete', () {
+      test('with CTEs', () {
+        testFormat(
+            'WITH foo (id) AS (SELECT * FROM bar) DELETE FROM bar WHERE x;');
+      });
+
+      test('with returning', () {
+        testFormat('DELETE FROM foo RETURNING *');
+      });
     });
 
     group('insert', () {
@@ -234,10 +240,14 @@ CREATE UNIQUE INDEX my_idx ON t1 (c1, c2, c3) WHERE c1 < c3;
             'REPLACE INTO foo DEFAULT VALUES');
       });
 
-      test('insert into select', () {
+      test('into select', () {
         testFormat('INSERT INTO foo SELECT * FROM bar');
       });
 
+      test('with returning', () {
+        testFormat('INSERT INTO foo DEFAULT VALUES RETURNING *');
+      });
+
       test('upsert - do nothing', () {
         testFormat(
             'INSERT OR REPLACE INTO foo DEFAULT VALUES ON CONFLICT DO NOTHING');
@@ -260,6 +270,10 @@ CREATE UNIQUE INDEX my_idx ON t1 (c1, c2, c3) WHERE c1 < c3;
         testFormat('UPDATE foo SET bar = baz WHERE 1;');
       });
 
+      test('with returning', () {
+        testFormat('UPDATE foo SET bar = baz RETURNING *');
+      });
+
       const modes = [
         'OR ABORT',
         'OR FAIL',