Return deleted row count on delete (#53)

pinchbv · Feb 15, 2019 · 183db39 · 183db39
1 parent fe1c498
commit 183db39
Show file tree

Hide file tree

Showing 3 changed files with 93 additions and 8 deletions.
diff --git a/floor_generator/lib/writer/delete_method_body_writer.dart b/floor_generator/lib/writer/delete_method_body_writer.dart
@@ -1,4 +1,5 @@
 import 'package:code_builder/code_builder.dart';
+import 'package:floor_generator/model/column.dart';
 import 'package:floor_generator/model/delete_method.dart';
 import 'package:floor_generator/writer/writer.dart';
 import 'package:source_gen/source_gen.dart';
@@ -15,23 +16,83 @@ class DeleteMethodBodyWriter implements Writer {
   }
 
   String _generateMethodBody() {
+    _assertMethodReturnsNoList();
+
     final entity = method.getEntity(library);
+    final entityName = entity.name;
     final primaryKeyColumn = entity.primaryKeyColumn;
-    final methodHeadParameterName = method.parameter.name;
+    final methodSignatureParameterName = method.parameter.name;
+
+    if (method.returnsInt) {
+      return _generateIntReturnMethodBody(
+        methodSignatureParameterName,
+        entityName,
+        primaryKeyColumn,
+      );
+    } else if (method.returnsVoid) {
+      return _generateVoidReturnMethodBody(
+        methodSignatureParameterName,
+        entityName,
+        primaryKeyColumn,
+      );
+    } else {
+      throw InvalidGenerationSourceError(
+        'Delete methods have to return a Future of either void or int.',
+        element: method.method,
+      );
+    }
+  }
 
+  String _generateVoidReturnMethodBody(
+    final String methodSignatureParameterName,
+    final String entityName,
+    final Column primaryKeyColumn,
+  ) {
     if (method.changesMultipleItems) {
       return '''
       final batch = database.batch();
-      for (final item in $methodHeadParameterName) {
-        batch.delete('${entity.name}', where: '${primaryKeyColumn.name} = ?', whereArgs: <int>[item.${primaryKeyColumn.field.displayName}]);
+      for (final item in $methodSignatureParameterName) {
+        batch.delete('$entityName', where: '${primaryKeyColumn.name} = ?', whereArgs: <int>[item.${primaryKeyColumn.field.displayName}]);
       }
       await batch.commit(noResult: true);
       ''';
     } else {
       return '''
-      final item = $methodHeadParameterName;
-      await database.delete('${entity.name}', where: '${primaryKeyColumn.name} = ?', whereArgs: <int>[item.${primaryKeyColumn.field.displayName}]);
+      final item = $methodSignatureParameterName;
+      await database.delete('$entityName', where: '${primaryKeyColumn.name} = ?', whereArgs: <int>[item.${primaryKeyColumn.field.displayName}]);
       ''';
     }
   }
+
+  String _generateIntReturnMethodBody(
+    final String methodSignatureParameterName,
+    final String entityName,
+    final Column primaryKeyColumn,
+  ) {
+    if (method.changesMultipleItems) {
+      return '''
+      final batch = database.batch();
+      for (final item in $methodSignatureParameterName) {
+        batch.delete('$entityName', where: '${primaryKeyColumn.name} = ?', whereArgs: <int>[item.${primaryKeyColumn.field.displayName}]);
+      }
+      return (await batch.commit(noResult: false))
+          .cast<int>()
+          .reduce((first, second) => first + second);
+      ''';
+    } else {
+      return '''
+      final item = $methodSignatureParameterName;
+      return database.delete('$entityName', where: '${primaryKeyColumn.name} = ?', whereArgs: <int>[item.${primaryKeyColumn.field.displayName}]);
+      ''';
+    }
+  }
+
+  void _assertMethodReturnsNoList() {
+    if (method.returnsList) {
+      throw InvalidGenerationSourceError(
+        'Delete methods have to return a Future of either void or int but not a list.',
+        element: method.method,
+      );
+    }
+  }
 }
diff --git a/floor_test/test/database.dart b/floor_test/test/database.dart
@@ -33,20 +33,26 @@ abstract class TestDatabase extends FloorDatabase {
   Future<void> updatePerson(Person person);
 
   @update
-  Future<int> updatePersonWithReturn(Person person);
+  Future<void> updatePersons(List<Person> persons);
 
   @update
-  Future<int> updatePersonsWithReturn(List<Person> persons);
+  Future<int> updatePersonWithReturn(Person person);
 
   @update
-  Future<void> updatePersons(List<Person> persons);
+  Future<int> updatePersonsWithReturn(List<Person> persons);
 
   @delete
   Future<void> deletePerson(Person person);
 
   @delete
   Future<void> deletePersons(List<Person> person);
 
+  @delete
+  Future<int> deletePersonWithReturn(Person person);
+
+  @delete
+  Future<int> deletePersonsWithReturn(List<Person> persons);
+
   @transaction
   Future<void> replacePersons(List<Person> persons) async {
     await database.execute('DELETE FROM person');

diff --git a/floor_test/test/database_test.dart b/floor_test/test/database_test.dart
@@ -136,6 +136,24 @@ void main() {
       expect(persistentPersons, equals(updatedPersons));
       expect(actual, equals(2));
     });
+
+    test('delete person and return 1 (affected row count)', () async {
+      final person = Person(1, 'Simon');
+      await database.insertPerson(person);
+
+      final actual = await database.deletePersonWithReturn(person);
+
+      expect(actual, equals(1));
+    });
+
+    test('delete persons and return affected rows count', () async {
+      final persons = [Person(1, 'Simon'), Person(2, 'Frank')];
+      await database.insertPersons(persons);
+
+      final actual = await database.deletePersonsWithReturn(persons);
+
+      expect(actual, equals(2));
+    });
   });
 }