Add entity classes to database annotation

floor_annotation/lib/src/database.dart

@@ -2,9 +2,12 @@ import 'package:meta/meta.dart';
 
 /// Marks a class as a FloorDatabase.
 class Database {
-  /// The database version
+  /// The database version.
   final int version;
 
+  /// The entities the database manages-
+  final List<Type> entities;
+
   /// Marks a class as a FloorDatabase.
-  const Database({@required this.version});
+  const Database({@required this.version, @required this.entities});
 }

floor_generator/lib/generator.dart

@@ -1,8 +1,10 @@
 import 'dart:async';
 
+import 'package:analyzer/dart/element/element.dart';
 import 'package:build/build.dart';
 import 'package:code_builder/code_builder.dart';
 import 'package:floor_generator/misc/annotations.dart';
+import 'package:floor_generator/misc/constants.dart';
 import 'package:floor_generator/misc/type_utils.dart';
 import 'package:floor_generator/processor/database_processor.dart';
 import 'package:floor_generator/processor/entity_processor.dart';
@@ -52,20 +54,35 @@ class FloorGenerator implements Generator {
           ' There are too many classes annotated with @Database.',
           element: databaseClasses[2]);
     } else {
+      final databaseClassElement = databaseClasses.first;
+      final entities = _getEntities(databaseClassElement);
+
+      if (entities == null || entities.isEmpty) {
+        throw InvalidGenerationSourceError(
+            'There are no entities added to the database annotation.',
+            element: databaseClassElement);
+      }
+
       return DatabaseProcessor(
-        databaseClasses.first,
-        _getEntities(library),
+        databaseClassElement,
+        entities,
       ).process();
     }
   }
 
-  @nonNull
-  List<Entity> _getEntities(final LibraryReader library) {
-    return library.classes
-        .where((clazz) =>
-            !clazz.isAbstract && clazz.metadata.any(isEntityAnnotation))
-        .map((entity) => EntityProcessor(entity).process())
-        .toList();
+  List<Entity> _getEntities(final ClassElement databaseClassElement) {
+    return databaseClassElement.metadata
+        .firstWhere(isDatabaseAnnotation)
+        .computeConstantValue()
+        .getField(AnnotationField.DATABASE_ENTITIES)
+        ?.toListValue()
+        ?.map((object) => object.toTypeValue().element)
+        ?.whereType<ClassElement>()
+        ?.where((classElement) =>
+            !classElement.isAbstract &&
+            classElement.metadata.any(isEntityAnnotation))
+        ?.map((classElement) => EntityProcessor(classElement).process())
+        ?.toList();
   }
 
   @nonNull

floor_generator/lib/misc/constants.dart

@@ -23,7 +23,9 @@ abstract class AnnotationField {
   static const QUERY_VALUE = 'value';
   static const PRIMARY_KEY_AUTO_GENERATE = 'autoGenerate';
   static const ON_CONFLICT = 'onConflict';
+
   static const DATABASE_VERSION = 'version';
+  static const DATABASE_ENTITIES = 'entities';
 
   static const COLUMN_INFO_NAME = 'name';
   static const COLUMN_INFO_NULLABLE = 'nullable';

floor_generator/lib/processor/database_processor.dart

@@ -52,19 +52,6 @@ class DatabaseProcessor extends Processor<Database> {
   }
 
   List<DaoGetter> _getDaoGetters(final String databaseName) {
-    // TODO decide for either implementing this approach (include adding
-    //  entities to the database annotation) or keep it like it is now
-    //  another option would be to get the libraryElement form the classElement
-    //  e.g. _classElement.library;
-
-//    final entities = _classElement.metadata
-//        .firstWhere(isDatabaseAnnotation)
-//        .computeConstantValue()
-//        .getField(AnnotationField.DATABASE_ENTITIES)
-//        ?.toListValue()
-//        ?.map((object) => object.toTypeValue().element)
-//        ?.whereType<ClassElement>();
-
     return _classElement.fields.where(_isDao).map((field) {
       final classElement = field.type.element as ClassElement;
       final name = field.displayName;

floor_test/test/dao/dog_dao.dart

@@ -1,4 +1,6 @@
-part of '../database.dart';
+import 'package:floor/floor.dart';
+
+import '../model/dog.dart';
 
 @dao
 abstract class DogDao {

floor_test/test/dao/person_dao.dart

@@ -1,4 +1,6 @@
-part of '../database.dart';
+import 'package:floor/floor.dart';
+
+import '../model/person.dart';
 
 @dao
 abstract class PersonDao {

floor_test/test/database.dart

@@ -4,14 +4,14 @@ import 'package:floor/floor.dart';
 import 'package:path/path.dart';
 import 'package:sqflite/sqflite.dart' as sqflite;
 
-part 'dao/dog_dao.dart';
-part 'dao/person_dao.dart';
+import 'dao/dog_dao.dart';
+import 'dao/person_dao.dart';
+import 'model/dog.dart';
+import 'model/person.dart';
+
 part 'database.g.dart';
-part 'model/address.dart';
-part 'model/dog.dart';
-part 'model/person.dart';
 
-@Database(version: 2)
+@Database(version: 2, entities: [Person, Dog])
 abstract class TestDatabase extends FloorDatabase {
   static Future<TestDatabase> openDatabase(List<Migration> migrations) async =>
       _$open(migrations);

floor_test/test/database_test.dart

@@ -3,7 +3,11 @@ import 'package:flutter_test/flutter_test.dart';
 import 'package:matcher/matcher.dart';
 import 'package:sqflite/sqflite.dart';
 
+import 'dao/dog_dao.dart';
+import 'dao/person_dao.dart';
 import 'database.dart';
+import 'model/dog.dart';
+import 'model/person.dart';
 
 // run test with 'flutter run test/database_test.dart'
 void main() {

floor_test/test/model/dog.dart

@@ -1,4 +1,6 @@
-part of '../database.dart';
+import 'package:floor/floor.dart';
+
+import 'person.dart';
 
 @Entity(
   tableName: 'dog',

floor_test/test/model/person.dart

@@ -1,4 +1,4 @@
-part of '../database.dart';
+import 'package:floor/floor.dart';
 
 @Entity(tableName: 'person')
 class Person {