Add update methods code generation

example/lib/database.dart

@@ -82,6 +82,12 @@ abstract class MyDatabase extends FloorDatabase {
 
   @insert
   Future<void> insertTask(Task task);
+
+  @update
+  Future<void> updateTask(Task task);
+
+  @update
+  Future<void> updatePerson(Person person);
 }
 
 Future<void> main() async {

floor/lib/src/annotations.dart

@@ -54,3 +54,11 @@ class Insert {
 
 /// Marks a method as an insert method.
 const insert = Insert();
+
+/// Marks a method as an update method.
+class Update {
+  const Update();
+}
+
+/// Marks a method as an update method.
+const update = Update();

floor_generator/lib/misc/annotation_expression.dart

@@ -6,3 +6,5 @@ class AnnotationExpression extends CodeExpression {
 
   AnnotationExpression(this.annotation) : super(Code(annotation));
 }
+
+final overrideAnnotationExpression = AnnotationExpression('override');

floor_generator/lib/misc/constants.dart

@@ -10,8 +10,10 @@ abstract class Annotation {
   static const DATABASE = 'Database';
   static const COLUMN_INFO = 'ColumnInfo';
   static const PRIMARY_KEY = 'PrimaryKey';
+
   static const QUERY = 'Query';
   static const INSERT = 'Insert';
+  static const UPDATE = 'Update';
 }
 
 abstract class AnnotationField {

floor_generator/lib/misc/type_utils.dart

@@ -56,6 +56,10 @@ bool isInsertAnnotation(ElementAnnotation annotation) {
   return _getAnnotationName(annotation) == Annotation.INSERT;
 }
 
+bool isUpdateAnnotation(ElementAnnotation annotation) {
+  return _getAnnotationName(annotation) == Annotation.UPDATE;
+}
+
 DartType flattenList(DartType type) {
   return (type as ParameterizedType).typeArguments.first;
 }

floor_generator/lib/model/database.dart

@@ -3,6 +3,7 @@ import 'package:floor_generator/misc/type_utils.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/update_method.dart';
 import 'package:source_gen/source_gen.dart';
 
 class Database {
@@ -28,6 +29,13 @@ class Database {
         .toList();
   }
 
+  List<UpdateMethod> get updateMethods {
+    return methods
+        .where((method) => method.metadata.any(isUpdateAnnotation))
+        .map((method) => UpdateMethod(method))
+        .toList();
+  }
+
   List<Entity> getEntities(LibraryReader library) {
     return library.classes
         .where((clazz) =>

floor_generator/lib/model/update_method.dart

@@ -0,0 +1,51 @@
+import 'package:analyzer/dart/element/element.dart';
+import 'package:analyzer/dart/element/type.dart';
+import 'package:floor_generator/misc/type_utils.dart';
+import 'package:floor_generator/model/entity.dart';
+import 'package:source_gen/source_gen.dart';
+
+class UpdateMethod {
+  final MethodElement method;
+
+  UpdateMethod(this.method);
+
+  DartType get returnType => method.returnType;
+
+  String get name => method.displayName;
+
+  ParameterElement get parameter {
+    final parameters = method.parameters;
+    if (parameters.isEmpty) {
+      throw InvalidGenerationSourceError(
+        'There is no parameter supplied for an update method. Please add one.',
+        element: method,
+      );
+    } else if (parameters.length > 1) {
+      throw InvalidGenerationSourceError(
+        'Only one parameter is allowed on update methods.',
+        element: method,
+      );
+    }
+    return parameters.first;
+  }
+
+  Entity getEntity(LibraryReader library) {
+    final entityClass = _getEntities(library).firstWhere(
+        (entity) => entity.displayName == parameter.type.displayName);
+
+    return Entity(entityClass);
+  }
+
+  bool insertsEntity(LibraryReader library) {
+    return _getEntities(library)
+        .map((clazz) => clazz.displayName)
+        .any((entity) => entity == parameter.type.displayName);
+  }
+
+  List<ClassElement> _getEntities(LibraryReader library) {
+    return library.classes
+        .where((clazz) =>
+            !clazz.isAbstract && clazz.metadata.any(isEntityAnnotation))
+        .toList();
+  }
+}

floor_generator/lib/writer/database_writer.dart

@@ -4,8 +4,10 @@ import 'package:floor_generator/model/database.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/update_method.dart';
 import 'package:floor_generator/writer/insert_method_writer.dart';
 import 'package:floor_generator/writer/query_method_writer.dart';
+import 'package:floor_generator/writer/update_method_writer.dart';
 import 'package:floor_generator/writer/writer.dart';
 import 'package:source_gen/source_gen.dart';
 import 'package:code_builder/code_builder.dart';
@@ -74,14 +76,20 @@ class DatabaseWriter implements Writer {
 
     final databaseName = database.name;
 
-    return Class(
-      (builder) => builder
-        ..name = '_\$$databaseName'
-        ..extend = refer(databaseName)
-        ..methods.add(_generateOpenMethod(databaseName, createTableStatements))
-        ..methods.addAll(_generateQueryMethods(database.queryMethods))
-        ..methods.addAll(_generateInsertMethods(database.insertMethods)),
-    );
+    return Class((builder) => builder
+      ..name = '_\$$databaseName'
+      ..extend = refer(databaseName)
+      ..methods.add(_generateOpenMethod(databaseName, createTableStatements))
+      ..methods.addAll(_generateQueryMethods(database.queryMethods))
+      ..methods.addAll(_generateInsertMethods(database.insertMethods))
+      ..methods.addAll(_generateUpdateMethods(database.updateMethods)));
+  }
+
+  List<Method> _generateUpdateMethods(List<UpdateMethod> updateMethods) {
+    return updateMethods
+        .map(
+            (updateMethod) => UpdateMethodWriter(library, updateMethod).write())
+        .toList();
   }
 
   Method _generateOpenMethod(
@@ -90,7 +98,7 @@ class DatabaseWriter implements Writer {
   ) {
     return Method((builder) => builder
       ..name = 'open'
-      ..annotations.add(AnnotationExpression('override'))
+      ..annotations.add(overrideAnnotationExpression)
       ..returns = refer('Future<sqflite.Database>')
       ..modifier = MethodModifier.async
       ..body = Code('''

floor_generator/lib/writer/insert_method_writer.dart

@@ -21,7 +21,7 @@ class InsertMethodWriter implements Writer {
     _assertInsertsEntity();
 
     return Method((builder) => builder
-      ..annotations.add(AnnotationExpression('override'))
+      ..annotations.add(overrideAnnotationExpression)
       ..returns = refer(insertMethod.returnType.displayName)
       ..name = insertMethod.name
       ..requiredParameters.add(_generateParameter())

floor_generator/lib/writer/query_method_writer.dart

@@ -23,7 +23,7 @@ class QueryMethodWriter implements Writer {
     _assertReturnsEntity();
 
     return Method((builder) => builder
-      ..annotations.add(AnnotationExpression('override'))
+      ..annotations.add(overrideAnnotationExpression)
       ..returns = refer(queryMethod.rawReturnType.displayName)
       ..name = queryMethod.name
       ..requiredParameters.addAll(_generateMethodParameters())

floor_generator/lib/writer/update_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/update_method.dart';
+import 'package:floor_generator/writer/writer.dart';
+import 'package:source_gen/source_gen.dart';
+
+class UpdateMethodWriter implements Writer {
+  final LibraryReader library;
+  final UpdateMethod updateMethod;
+
+  UpdateMethodWriter(this.library, this.updateMethod);
+
+  @override
+  Method write() {
+    return _generateUpdateMethod();
+  }
+
+  Method _generateUpdateMethod() {
+    // TODO assert is entity
+
+    return Method((builder) => builder
+      ..annotations.add(overrideAnnotationExpression)
+      ..returns = refer(updateMethod.returnType.displayName)
+      ..name = updateMethod.name
+      ..requiredParameters.add(_generateParameter())
+      ..modifier = MethodModifier.async
+      ..body = Code(_generateMethodBody()));
+  }
+
+  Parameter _generateParameter() {
+    final parameter = updateMethod.parameter;
+
+    return Parameter((builder) => builder
+      ..name = parameter.name
+      ..type = refer(parameter.type.displayName));
+  }
+
+  String _generateMethodBody() {
+    final entity = updateMethod.getEntity(library);
+    final primaryKeyColumn =
+        entity.columns.firstWhere((column) => column.isPrimaryKey);
+    final methodHeadParameterName = updateMethod.parameter.name;
+
+    return '''
+    await this.database.rawDelete('DELETE FROM ${entity.name} WHERE ${primaryKeyColumn.name} = \${$methodHeadParameterName.${primaryKeyColumn.field.displayName}}');
+    ''';
+  }
+}