Add support for transactions

floor/lib/src/annotations.dart

@@ -64,3 +64,11 @@ class Delete {
 
 /// Marks a method as a delete method.
 const delete = Delete();
+
+/// Marks a method as a transaction method.
+class _Transaction {
+  const _Transaction();
+}
+
+/// Marks a method as a transaction method.
+const transaction = _Transaction();

floor/lib/src/database.dart

@@ -3,16 +3,18 @@ import 'package:sqflite/sqflite.dart' as sqflite;
 /// Extend this class to enable database functionality.
 abstract class FloorDatabase {
   /// Use this for direct access to the sqflite database.
-  sqflite.Database database;
+  sqflite.DatabaseExecutor database;
 
   // TODO remove this
   /// Opens the database to be able to query it.
   Future<sqflite.Database> open();
 
   /// Closes the database.
   Future<void> close() async {
-    if (database?.isOpen ?? false) {
-      await database.close();
+    final immutableDatabase = database;
+    if (immutableDatabase is sqflite.Database &&
+        (immutableDatabase?.isOpen ?? false)) {
+      await immutableDatabase.close();
     }
   }
 }

floor_generator/lib/misc/constants.dart

@@ -10,6 +10,7 @@ abstract class Annotation {
   static const DATABASE = 'Database';
   static const COLUMN_INFO = 'ColumnInfo';
   static const PRIMARY_KEY = 'PrimaryKey';
+  static const TRANSACTION = '_Transaction';
 
   static const QUERY = 'Query';
   static const INSERT = 'Insert';
@@ -26,11 +27,8 @@ abstract class AnnotationField {
   static const COLUMN_INFO_NULLABLE = 'nullable';
 }
 
-abstract class SqlConstants {
+abstract class SqlType {
   static const INTEGER = 'INTEGER';
   static const TEXT = 'TEXT';
   static const REAL = 'REAL';
-
-  static const PRIMARY_KEY = 'PRIMARY KEY';
-  static const AUTOINCREMENT = 'AUTOINCREMENT';
 }

floor_generator/lib/misc/type_utils.dart

@@ -64,6 +64,10 @@ bool isDeleteAnnotation(ElementAnnotation annotation) {
   return _getAnnotationName(annotation) == Annotation.DELETE;
 }
 
+bool isTransactionAnnotation(ElementAnnotation annotation) {
+  return _getAnnotationName(annotation) == Annotation.TRANSACTION;
+}
+
 DartType flattenList(DartType type) {
   return (type as ParameterizedType).typeArguments.first;
 }

floor_generator/lib/model/column.dart

@@ -22,6 +22,9 @@ class Column {
   }
 
   bool get isNullable {
+    if (isPrimaryKey) {
+      return false;
+    }
     if (!_hasColumnInfoAnnotation) {
       return true; // all Dart fields are nullable by default
     }
@@ -40,13 +43,13 @@ class Column {
   String get type {
     final type = field.type;
     if (isInt(type)) {
-      return SqlConstants.INTEGER;
+      return SqlType.INTEGER;
     } else if (isString(type)) {
-      return SqlConstants.INTEGER;
+      return SqlType.TEXT;
     } else if (isBool(type)) {
-      return SqlConstants.INTEGER;
+      return SqlType.INTEGER;
     } else if (isDouble(type)) {
-      return SqlConstants.REAL;
+      return SqlType.REAL;
     }
     throw InvalidGenerationSourceError(
       'Column type is not supported for $type.',
@@ -56,9 +59,9 @@ class Column {
 
   bool get isPrimaryKey => field.metadata.any(isPrimaryKeyAnnotation);
 
-  bool get _autoGenerate {
+  bool get autoGenerate {
     if (!isPrimaryKey) {
-      return null;
+      return false;
     }
     return field.metadata
             .firstWhere(isPrimaryKeyAnnotation)
@@ -67,21 +70,4 @@ class Column {
             .toBoolValue() ??
         false;
   }
-
-  /// Primary key and auto increment.
-  String get additionals {
-    String add = '';
-
-    if (isPrimaryKey) {
-      add += ' ${SqlConstants.PRIMARY_KEY}';
-      if (_autoGenerate) {
-        add += ' ${SqlConstants.AUTOINCREMENT}';
-      }
-    }
-
-    if (add.isEmpty) {
-      return null;
-    }
-    return add;
-  }
 }

floor_generator/lib/model/database.dart

@@ -4,6 +4,7 @@ import 'package:floor_generator/model/delete_method.dart';
 import 'package:floor_generator/model/entity.dart';
 import 'package:floor_generator/model/insert_method.dart';
 import 'package:floor_generator/model/query_method.dart';
+import 'package:floor_generator/model/transaction_method.dart';
 import 'package:floor_generator/model/update_method.dart';
 import 'package:source_gen/source_gen.dart';
 
@@ -44,6 +45,13 @@ class Database {
         .toList();
   }
 
+  List<TransactionMethod> get transactionMethods {
+    return methods
+        .where((method) => method.metadata.any(isTransactionAnnotation))
+        .map((method) => TransactionMethod(method, name))
+        .toList();
+  }
+
   List<Entity> getEntities(LibraryReader library) {
     return library.classes
         .where((clazz) =>

floor_generator/lib/model/entity.dart

@@ -29,11 +29,14 @@ class Entity {
   }
 
   Column get primaryKeyColumn {
-    return columns.firstWhere(
-      (column) => column.isPrimaryKey,
-      orElse: throw InvalidGenerationSourceError(
-          'There is no primary key defined on the entity $name.',
-          element: clazz),
-    );
+    return columns.firstWhere((column) => column.isPrimaryKey);
+
+    // TODO why does this always throw?
+//    return columns.firstWhere(
+//      (column) => column.isPrimaryKey,
+//      orElse: throw InvalidGenerationSourceError(
+//          'There is no primary key defined on the entity $name.',
+//          element: clazz),
+//    );
   }
 }

floor_generator/lib/model/transaction_method.dart

@@ -0,0 +1,15 @@
+import 'package:analyzer/dart/element/element.dart';
+import 'package:analyzer/dart/element/type.dart';
+
+class TransactionMethod {
+  final MethodElement method;
+  final String databaseName;
+
+  TransactionMethod(this.method, this.databaseName);
+
+  DartType get returnType => method.returnType;
+
+  String get name => method.displayName;
+
+  List<ParameterElement> get parameters => method.parameters;
+}

floor_generator/lib/writer/database_writer.dart

@@ -5,11 +5,13 @@ import 'package:floor_generator/model/delete_method.dart';
 import 'package:floor_generator/model/entity.dart';
 import 'package:floor_generator/model/insert_method.dart';
 import 'package:floor_generator/model/query_method.dart';
+import 'package:floor_generator/model/transaction_method.dart';
 import 'package:floor_generator/model/update_method.dart';
 import 'package:floor_generator/writer/change_method_writer.dart';
 import 'package:floor_generator/writer/delete_method_body_writer.dart';
 import 'package:floor_generator/writer/insert_method_body_writer.dart';
 import 'package:floor_generator/writer/query_method_writer.dart';
+import 'package:floor_generator/writer/transaction_method_writer.dart';
 import 'package:floor_generator/writer/update_method_body_writer.dart';
 import 'package:floor_generator/writer/writer.dart';
 import 'package:source_gen/source_gen.dart';
@@ -86,7 +88,9 @@ class DatabaseWriter implements Writer {
       ..methods.addAll(_generateQueryMethods(database.queryMethods))
       ..methods.addAll(_generateInsertMethods(database.insertMethods))
       ..methods.addAll(_generateUpdateMethods(database.updateMethods))
-      ..methods.addAll(_generateDeleteMethods(database.deleteMethods)));
+      ..methods.addAll(_generateDeleteMethods(database.deleteMethods))
+      ..methods
+          .addAll(_generateTransactionMethods(database.transactionMethods)));
   }
 
   Method _generateOpenMethod(
@@ -138,22 +142,24 @@ class DatabaseWriter implements Writer {
         .toList();
   }
 
+  List<Method> _generateTransactionMethods(
+    List<TransactionMethod> transactionMethods,
+  ) {
+    return transactionMethods
+        .map((method) => TransactionMethodWriter(library, method).write())
+        .toList();
+  }
+
   List<String> _generateCreateTableSqlStatements(List<Entity> entities) {
     return entities.map(_generateSql).toList();
   }
 
   String _generateSql(Entity entity) {
     final columns = entity.columns.map((column) {
-      var columnString = '`${column.name}` ${column.type}';
-
-      final additionals = column.additionals;
-      if (additionals != null) {
-        columnString += additionals;
-      }
-      if (!column.isNullable) {
-        columnString += ' NOT NULL';
-      }
-      return columnString;
+      final primaryKey = column.isPrimaryKey ? ' PRIMARY KEY' : '';
+      final autoIncrement = column.autoGenerate ? ' AUTOINCREMENT' : '';
+      final nullable = column.isNullable ? '' : ' NOT NULL';
+      return '`${column.name}` ${column.type}$primaryKey$autoIncrement$nullable';
     }).join(', ');
 
     return "'CREATE TABLE IF NOT EXISTS `${entity.name}` ($columns)'";

floor_generator/lib/writer/transaction_method_writer.dart

@@ -0,0 +1,48 @@
+import 'package:code_builder/code_builder.dart';
+import 'package:floor_generator/misc/annotation_expression.dart';
+import 'package:floor_generator/model/transaction_method.dart';
+import 'package:floor_generator/writer/writer.dart';
+import 'package:source_gen/source_gen.dart';
+
+class TransactionMethodWriter implements Writer {
+  final LibraryReader library;
+  final TransactionMethod method;
+
+  TransactionMethodWriter(this.library, this.method);
+
+  @override
+  Method write() {
+    return Method((builder) => builder
+      ..annotations.add(overrideAnnotationExpression)
+      ..returns = refer(method.returnType.displayName)
+      ..name = method.name
+      ..requiredParameters.addAll(_generateParameters())
+      ..modifier = MethodModifier.async
+      ..body = Code(_generateMethodBody()));
+  }
+
+  String _generateMethodBody() {
+    final parameters =
+        method.parameters.map((parameter) => parameter.name).join(', ');
+    final methodCall = '${method.name}($parameters)';
+
+    return '''
+    if (database is sqflite.Transaction) {
+      await super.$methodCall;
+    } else {
+      await (database as sqflite.Database).transaction<void>((transaction) async {
+        final transactionDatabase = _\$${method.databaseName}()..database = transaction;
+        await transactionDatabase.$methodCall;
+      });
+    }
+    ''';
+  }
+
+  List<Parameter> _generateParameters() {
+    return method.parameters.map((parameter) {
+      return Parameter((builder) => builder
+        ..name = parameter.name
+        ..type = refer(parameter.type.displayName));
+    }).toList();
+  }
+}

floor_test/test/database.dart

@@ -31,12 +31,19 @@ abstract class TestDatabase extends FloorDatabase {
 
   @delete
   Future<void> deletePersons(List<Person> person);
+
+  @transaction
+  Future<void> replacePersons(List<Person> persons) async {
+    await database.execute('DELETE FROM person');
+    await insertPersons(persons);
+  }
 }
 
 @Entity(tableName: 'person')
 class Person {
   @PrimaryKey()
   final int id;
+
   @ColumnInfo(name: 'custom_name', nullable: false)
   final String name;
 

floor_test/test/database_test.dart

@@ -12,7 +12,7 @@ void main() {
     });
 
     tearDown(() async {
-      await database.database.execute('DELETE FROM Person');
+      await database.database.execute('DELETE FROM person');
     });
 
     test('database initially is empty', () async {
@@ -85,6 +85,17 @@ void main() {
       final actual = await database.findAllPersons();
       expect(actual, equals(updatedPersons));
     });
+
+    test('replace persons in transaction', () async {
+      final persons = [Person(1, 'Simon'), Person(2, 'Frank')];
+      await database.insertPersons(persons);
+
+      final newPersons = [Person(3, 'Paul'), Person(4, 'Karl')];
+      await database.replacePersons(newPersons);
+
+      final actual = await database.findAllPersons();
+      expect(actual, equals(newPersons));
+    });
   });
 }