From 9d9a4f40654e586fe5e5f5a37b06adddd0ea74b6 Mon Sep 17 00:00:00 2001
From: Simon Binder <oss@simonbinder.eu>
Date: Fri, 11 Sep 2020 11:19:18 +0200
Subject: [PATCH] Add customStatement to batch api (#817)

---
 moor/lib/src/runtime/api/batch.dart | 20 +++++++++++++++++---
 moor/test/batch_test.dart           |  4 ++++
 2 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/moor/lib/src/runtime/api/batch.dart b/moor/lib/src/runtime/api/batch.dart
index e710dc70..1ac942f0 100644
--- a/moor/lib/src/runtime/api/batch.dart
+++ b/moor/lib/src/runtime/api/batch.dart
@@ -139,10 +139,24 @@ class Batch {
     _addContext(stmt.constructQuery());
   }
 
-  void _addContext(GenerationContext ctx) {
-    final sql = ctx.sql;
-    final arguments = ctx.boundVariables;
+  /// Executes the custom [sql] statement with variables instantiated to [args].
+  ///
+  /// The statement will be added to this batch and executed when the batch
+  /// completes. So, this method returns synchronously and it's not possible to
+  /// inspect the return value of individual statements.
+  ///
+  /// See also:
+  ///  - [QueryEngine.customStatement], the equivalent method outside of
+  ///    batches.
+  void customStatement(String sql, [List<dynamic> args]) {
+    _addSqlAndArguments(sql, args);
+  }
 
+  void _addContext(GenerationContext ctx) {
+    _addSqlAndArguments(ctx.sql, ctx.boundVariables);
+  }
+
+  void _addSqlAndArguments(String sql, List<dynamic> arguments) {
     final stmtIndex = _sqlToIndex.putIfAbsent(sql, () {
       final newIndex = _createdSql.length;
       _createdSql.add(sql);
diff --git a/moor/test/batch_test.dart b/moor/test/batch_test.dart
index 3314f023..c2f58c81 100644
--- a/moor/test/batch_test.dart
+++ b/moor/test/batch_test.dart
@@ -42,6 +42,8 @@ void main() {
       b.delete(db.todosTable, const TodosTableCompanion(id: Value(3)));
 
       b.update(db.users, const UsersCompanion(name: Value('new name 2')));
+
+      b.customStatement('some custom statement', [4]);
     });
 
     final transaction = executor.transactions;
@@ -55,6 +57,7 @@ void main() {
             'UPDATE categories SET `desc` = ?, priority = 0 WHERE id = ?;',
             'DELETE FROM categories WHERE 1;',
             'DELETE FROM todos WHERE id = ?;',
+            'some custom statement',
           ],
           [
             ArgumentsForBatchedStatement(0, ['first']),
@@ -66,6 +69,7 @@ void main() {
             ArgumentsForBatchedStatement(4, []),
             ArgumentsForBatchedStatement(5, [3]),
             ArgumentsForBatchedStatement(1, ['new name 2']),
+            ArgumentsForBatchedStatement(6, [4]),
           ],
         ),
       ),