Add support for foreign keys

floor/lib/src/annotations.dart

@@ -1,5 +1,8 @@
+import 'package:meta/meta.dart';
+
 /// Marks a class as a FloorDatabase.
 class Database {
+  /// Marks a class as a FloorDatabase.
   const Database();
 }
 
@@ -8,9 +11,13 @@ class Entity {
   /// The table name of the SQLite table.
   final String tableName;
 
-  const Entity({this.tableName});
+  /// List of [ForeignKey] constraints on this entity.
+  final List<ForeignKey> foreignKeys;
+
+  const Entity({this.tableName, this.foreignKeys});
 }
 
+/// Marks a class as a database entity (table).
 const entity = Entity();
 
 /// Allows customization of the column associated with this field.
@@ -33,6 +40,85 @@ class PrimaryKey {
   const PrimaryKey({this.autoGenerate = false});
 }
 
+/// Declares a foreign key on another [Entity].
+class ForeignKey {
+  /// The list of column names in the current [Entity].
+  final List<String> childColumns;
+
+  /// The list of column names in the parent [Entity].
+  final List<String> parentColumns;
+
+  /// The parent entity to reference.
+  final Type entity;
+
+  /// [ForeignKeyAction]
+  final int onUpdate;
+
+  /// [ForeignKeyAction]
+  final int onDelete;
+
+  /// Declares a foreign key on another [Entity].
+  const ForeignKey({
+    @required this.childColumns,
+    @required this.parentColumns,
+    @required this.entity,
+    this.onUpdate,
+    this.onDelete,
+  });
+}
+
+/// Constants definition for values that can be used in
+/// [ForeignKey.onDelete] and [ForeignKey.onUpdate]
+abstract class ForeignKeyAction {
+  /// Possible value for [ForeignKey.onDelete] or [ForeignKey.onUpdate].
+  ///
+  /// When a parent key is modified or deleted from the database, no special
+  /// action is taken. This means that SQLite will not make any effort to fix
+  /// the constraint failure, instead, reject the change.
+  static const NO_ACTION = 1;
+
+  /// Possible value for [ForeignKey.onDelete] or [ForeignKey.onUpdate].
+  ///
+  /// The RESTRICT action means that the application is prohibited from deleting
+  /// (for [ForeignKey.onDelete]) or modifying (for [ForeignKey.onUpdate]) a
+  /// parent key when there exists one or more child keys mapped to it. The
+  /// difference between the effect of a RESTRICT action and normal foreign key
+  /// constraint enforcement is that the RESTRICT action processing happens as
+  /// soon as the field is updated - not at the end of the current statement as
+  /// it would with an immediate constraint, or at the end of the current
+  /// transaction as it would with a deferred() constraint.
+  ///
+  /// Even if the foreign key constraint it is attached to is deferred(),
+  /// configuring a RESTRICT action causes SQLite to return an error immediately
+  /// if a parent key with dependent child keys is deleted or modified.
+  static const RESTRICT = 2;
+
+  /// Possible value for [ForeignKey.onDelete] or [ForeignKey.onUpdate].
+  ///
+  /// If the configured action is 'SET NULL', then when a parent key is deleted
+  /// (for [ForeignKey.onDelete]) or modified (for [ForeignKey.onUpdate]), the
+  /// child key columns of all rows in the child table that mapped to the parent
+  /// key are set to contain NULL values.
+  static const SET_NULL = 3;
+
+  /// Possible value for [ForeignKey.onDelete] or [ForeignKey.onUpdate].
+  ///
+  /// The 'SET DEFAULT' actions are similar to SET_NULL, except that each of the
+  /// child key columns is set to contain the columns default value instead of
+  /// NULL.
+  static const SET_DEFAULT = 4;
+
+  /// Possible value for [ForeignKey.onDelete] or [ForeignKey.onUpdate].
+  ///
+  /// A 'CASCADE' action propagates the delete or update operation on the parent
+  /// key to each dependent child key. For [ForeignKey.onDelete] action, this
+  /// means that each row in the child entity that was associated with the
+  /// deleted parent row is also deleted. For an [ForeignKey.onUpdate] action,
+  /// it means that the values stored in each dependent child key are modified
+  /// to match the new parent key values.
+  static const CASCADE = 5;
+}
+
 /// Marks a method as a query method.
 class Query {
   /// The SQLite query.

floor_generator/lib/misc/constants.dart

@@ -21,10 +21,28 @@ abstract class Annotation {
 abstract class AnnotationField {
   static const QUERY_VALUE = 'value';
   static const PRIMARY_KEY_AUTO_GENERATE = 'autoGenerate';
-  static const ENTITY_TABLE_NAME = 'tableName';
 
   static const COLUMN_INFO_NAME = 'name';
   static const COLUMN_INFO_NULLABLE = 'nullable';
+
+  static const ENTITY_TABLE_NAME = 'tableName';
+  static const ENTITY_FOREIGN_KEYS = 'foreignKeys';
+}
+
+abstract class ForeignKeyField {
+  static const ENTITY = 'entity';
+  static const CHILD_COLUMNS = 'childColumns';
+  static const PARENT_COLUMNS = 'parentColumns';
+  static const ON_UPDATE = 'onUpdate';
+  static const ON_DELETE = 'onDelete';
+}
+
+abstract class ForeignKeyAction {
+  static const NO_ACTION = 1;
+  static const RESTRICT = 2;
+  static const SET_NULL = 3;
+  static const SET_DEFAULT = 4;
+  static const CASCADE = 5;
 }
 
 abstract class SqlType {

floor_generator/lib/model/entity.dart

@@ -2,6 +2,8 @@ import 'package:analyzer/dart/element/element.dart';
 import 'package:floor_generator/misc/constants.dart';
 import 'package:floor_generator/misc/type_utils.dart';
 import 'package:floor_generator/model/column.dart';
+import 'package:floor_generator/model/foreign_key.dart';
+import 'package:source_gen/source_gen.dart';
 
 class Entity {
   final ClassElement clazz;
@@ -28,14 +30,22 @@ class Entity {
   }
 
   Column get primaryKeyColumn {
-    return columns.firstWhere((column) => column.isPrimaryKey);
+    return columns.firstWhere(
+      (column) => column.isPrimaryKey,
+      orElse: () => throw InvalidGenerationSourceError(
+            'There is no primary key defined on the entity $name.',
+            element: clazz,
+          ),
+    );
+  }
 
-    // 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),
-//    );
+  List<ForeignKey> get foreignKeys {
+    return clazz.metadata
+        .firstWhere(isEntityAnnotation)
+        .computeConstantValue()
+        .getField(AnnotationField.ENTITY_FOREIGN_KEYS)
+        ?.toListValue()
+        ?.map((object) => ForeignKey(clazz, object))
+        ?.toList();
   }
 }

floor_generator/lib/model/foreign_key.dart

@@ -0,0 +1,64 @@
+import 'package:analyzer/dart/constant/value.dart';
+import 'package:analyzer/dart/element/element.dart';
+import 'package:floor_generator/misc/constants.dart';
+import 'package:floor_generator/misc/type_utils.dart';
+import 'package:floor_generator/model/entity.dart';
+import 'package:source_gen/source_gen.dart';
+
+class ForeignKey {
+  final ClassElement entityClass;
+  final DartObject object;
+
+  ForeignKey(this.entityClass, this.object);
+
+  /// Returns the parent column name referenced with this foreign key.
+  String getParentName(final LibraryReader library) {
+    final entityClassName =
+        object.getField(ForeignKeyField.ENTITY)?.toTypeValue()?.displayName ??
+            (throw InvalidGenerationSourceError(
+              'No entity defined for foreign key',
+              element: entityClass,
+            ));
+
+    return library.classes
+        .where((clazz) =>
+            !clazz.isAbstract && clazz.metadata.any(isEntityAnnotation))
+        .map((clazz) => Entity(clazz))
+        .firstWhere(
+          (entity) => entity.clazz.displayName == entityClassName,
+          orElse: () => throw InvalidGenerationSourceError(
+                '$entityClassName is not an entity. Did you miss annotating the class with @Entity?',
+                element: entityClass,
+              ),
+        )
+        .name;
+  }
+
+  List<String> get childColumns {
+    return _getColumns(ForeignKeyField.CHILD_COLUMNS) ??
+        (throw InvalidGenerationSourceError(
+          'No child columns defined for foreign key',
+          element: entityClass,
+        ));
+  }
+
+  List<String> get parentColumns {
+    return _getColumns(ForeignKeyField.PARENT_COLUMNS) ??
+        (throw InvalidGenerationSourceError(
+          'No parent columns defined for foreign key',
+          element: entityClass,
+        ));
+  }
+
+  int get onUpdate => object.getField(ForeignKeyField.ON_UPDATE)?.toIntValue();
+
+  int get onDelete => object.getField(ForeignKeyField.ON_DELETE)?.toIntValue();
+
+  List<String> _getColumns(final String foreignKeyField) {
+    return object
+        .getField(foreignKeyField)
+        ?.toListValue()
+        ?.map((object) => object.toStringValue())
+        ?.toList();
+  }
+}

floor_generator/lib/writer/database_writer.dart

@@ -1,4 +1,6 @@
+import 'package:code_builder/code_builder.dart';
 import 'package:floor_generator/misc/annotation_expression.dart';
+import 'package:floor_generator/misc/constants.dart';
 import 'package:floor_generator/misc/type_utils.dart';
 import 'package:floor_generator/model/database.dart';
 import 'package:floor_generator/model/delete_method.dart';
@@ -15,7 +17,6 @@ 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';
-import 'package:code_builder/code_builder.dart';
 
 /// Takes care of generating the database implementation.
 class DatabaseWriter implements Writer {
@@ -108,6 +109,9 @@ class DatabaseWriter implements Writer {
           return sqflite.openDatabase(
             path,
             version: 1,
+            onConfigure: (database) async {
+              await database.execute('PRAGMA foreign_keys = ON');
+            },
             onCreate: (database, version) async {
               $createTableStatements
             },
@@ -155,13 +159,56 @@ class DatabaseWriter implements Writer {
   }
 
   String _generateSql(final Entity entity) {
+    final foreignKeys = _generateForeignKeys(entity) ?? '';
+
     final columns = entity.columns.map((column) {
       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)'";
+    return "'CREATE TABLE IF NOT EXISTS `${entity.name}` ($columns$foreignKeys)'";
+  }
+
+  String _generateForeignKeys(final Entity entity) {
+    return entity.foreignKeys?.map((foreignKey) {
+      final childColumns = foreignKey.childColumns.join(', ');
+      final parentColumns = foreignKey.parentColumns.join(', ');
+      final parentName = foreignKey.getParentName(library);
+
+      final onUpdate = _getOnUpdateAction(foreignKey.onUpdate) ?? '';
+      final onDelete = _getOnDeleteAction(foreignKey.onDelete) ?? '';
+
+      return ', FOREIGN KEY ($childColumns) REFERENCES `$parentName` ($parentColumns)$onUpdate$onDelete';
+    })?.join();
+  }
+
+  String _getOnUpdateAction(final int action) {
+    final updateAction = _getAction(action);
+    return updateAction != null ? ' ON UPDATE $updateAction' : null;
+  }
+
+  String _getOnDeleteAction(final int action) {
+    final deleteAction = _getAction(action);
+    return deleteAction != null ? ' ON DELETE $deleteAction' : null;
+  }
+
+  String _getAction(final int action) {
+    switch (action) {
+      case ForeignKeyAction.NO_ACTION:
+        return 'NO_ACTION';
+      case ForeignKeyAction.RESTRICT:
+        return 'RESTRICT';
+      case ForeignKeyAction.SET_NULL:
+        return 'SET_NULL';
+      case ForeignKeyAction.SET_DEFAULT:
+        return 'SET_DEFAULT';
+      case ForeignKeyAction.CASCADE:
+        return 'CASCADE';
+      default:
+        return null;
+    }
   }
 }

floor_test/test/database.dart

@@ -58,6 +58,15 @@ abstract class TestDatabase extends FloorDatabase {
     await database.execute('DELETE FROM person');
     await insertPersons(persons);
   }
+
+  @insert
+  Future<void> insertDog(Dog dog);
+
+  @Query('SELECT * FROM dog WHERE owner_id = :id')
+  Future<Dog> findDogForPersonId(int id);
+
+  @Query('SELECT * FROM dog')
+  Future<List<Dog>> findAllDogs();
 }
 
 @Entity(tableName: 'person')
@@ -86,3 +95,43 @@ class Person {
     return 'Person{id: $id, name: $name}';
   }
 }
+
+@Entity(
+  tableName: 'dog',
+  foreignKeys: [
+    ForeignKey(
+      childColumns: ['owner_id'],
+      parentColumns: ['id'],
+      entity: Person,
+      onDelete: ForeignKeyAction.CASCADE,
+    )
+  ],
+)
+class Dog {
+  @PrimaryKey()
+  final int id;
+
+  final String name;
+
+  @ColumnInfo(name: 'owner_id')
+  final int ownerId;
+
+  Dog(this.id, this.name, this.ownerId);
+
+  @override
+  bool operator ==(Object other) =>
+      identical(this, other) ||
+      other is Dog &&
+          runtimeType == other.runtimeType &&
+          id == other.id &&
+          name == other.name &&
+          ownerId == other.ownerId;
+
+  @override
+  int get hashCode => id.hashCode ^ name.hashCode ^ ownerId.hashCode;
+
+  @override
+  String toString() {
+    return 'Dog{id: $id, name: $name, ownerId: $ownerId}';
+  }
+}