Database writer improvements

floor_generator/lib/misc/constants.dart

@@ -0,0 +1,28 @@
+abstract class SupportedType {
+  static const STRING = 'String';
+  static const BOOL = 'bool';
+  static const INT = 'int';
+  static const DOUBLE = 'double';
+}
+
+abstract class Annotation {
+  static const ENTITY = 'Entity';
+  static const DATABASE = 'Database';
+  static const COLUMN_INFO = 'ColumnInfo';
+  static const PRIMARY_KEY = 'PrimaryKey';
+  static const QUERY = 'Query';
+}
+
+abstract class AnnotationField {
+  static const QUERY_VALUE = 'value';
+  static const COLUMN_INFO_AUTO_GENERATE = 'autoGenerate';
+}
+
+abstract class SqlConstants {
+  static const INTEGER = 'INTEGER';
+  static const TEXT = 'TEXT';
+  static const REAL = 'REAL';
+
+  static const PRIMARY_KEY = 'PRIMARY KEY';
+  static const AUTOINCREMENT = 'AUTOINCREMNT';
+}

floor_generator/lib/misc/type_utils.dart

@@ -1,5 +1,6 @@
 import 'package:analyzer/dart/element/element.dart';
 import 'package:analyzer/dart/element/type.dart';
+import 'package:floor_generator/misc/constants.dart';
 
 bool isString(DartType type) {
   return type.displayName == SupportedType.STRING && _isDartCore(type);
@@ -62,22 +63,3 @@ bool _isDartCore(DartType type) {
 String _getAnnotationName(ElementAnnotation annotation) {
   return annotation.computeConstantValue().type.displayName;
 }
-
-abstract class SupportedType {
-  static const STRING = 'String';
-  static const BOOL = 'bool';
-  static const INT = 'int';
-  static const DOUBLE = 'double';
-}
-
-abstract class Annotation {
-  static const ENTITY = 'Entity';
-  static const DATABASE = 'Database';
-  static const COLUMN_INFO = 'ColumnInfo';
-  static const PRIMARY_KEY = 'PrimaryKey';
-  static const QUERY = 'Query';
-}
-
-abstract class AnnotationField {
-  static const QUERY_VALUE = 'value';
-}

floor_generator/lib/model/column.dart

@@ -1,9 +1,60 @@
+import 'package:analyzer/dart/element/element.dart';
+import 'package:floor_generator/misc/constants.dart';
+import 'package:floor_generator/misc/type_utils.dart';
+import 'package:source_gen/source_gen.dart';
+
 /// Represents a table column.
 class Column {
-  final String name;
-  final String type;
-  final bool isPrimaryKey;
-  final bool autoGenerate;
+  final FieldElement field;
+
+  Column(this.field);
+
+  String get name => field.displayName;
+
+  String get type {
+    final type = field.type;
+    if (isInt(type)) {
+      return SqlConstants.INTEGER;
+    } else if (isString(type)) {
+      return SqlConstants.INTEGER;
+    } else if (isBool(type)) {
+      return SqlConstants.INTEGER;
+    } else if (isDouble(type)) {
+      return SqlConstants.REAL;
+    }
+    throw InvalidGenerationSourceError(
+      'Column type is not supported for $type.',
+      element: field,
+    );
+  }
+
+  bool get isPrimaryKey => field.metadata.any(isPrimaryKeyAnnotation);
+
+  bool get autoGenerate {
+    if (!isPrimaryKey) {
+      return null;
+    }
+    return field.metadata
+        .firstWhere(isPrimaryKeyAnnotation)
+        .computeConstantValue()
+        .getField(AnnotationField.COLUMN_INFO_AUTO_GENERATE)
+        .toBoolValue();
+  }
+
+  /// Primary key and auto increment.
+  String get additionals {
+    String add = '';
+
+    if (isPrimaryKey) {
+      add += ' ${SqlConstants.PRIMARY_KEY}';
+      if (autoGenerate) {
+        add += ' ${SqlConstants.AUTOINCREMENT}';
+      }
+    }
 
-  const Column(this.name, this.type, this.isPrimaryKey, this.autoGenerate);
+    if (add.isEmpty) {
+      return null;
+    }
+    return add;
+  }
 }

floor_generator/lib/model/database.dart

@@ -0,0 +1,30 @@
+import 'package:analyzer/dart/element/element.dart';
+import 'package:floor_generator/misc/type_utils.dart';
+import 'package:floor_generator/model/entity.dart';
+import 'package:floor_generator/model/query_method.dart';
+import 'package:source_gen/source_gen.dart';
+
+class Database {
+  final ClassElement clazz;
+
+  Database(this.clazz);
+
+  String get name => clazz.displayName;
+
+  List<MethodElement> get methods => clazz.methods;
+
+  List<QueryMethod> get queryMethods {
+    return methods
+        .where((method) => method.metadata.any(isQueryAnnotation))
+        .map((method) => QueryMethod(method))
+        .toList();
+  }
+
+  List<Entity> getEntities(LibraryReader library) {
+    return library.classes
+        .where((clazz) =>
+            !clazz.isAbstract && clazz.metadata.any(isEntityAnnotation))
+        .map((entity) => Entity(entity))
+        .toList();
+  }
+}

floor_generator/lib/model/entity.dart

@@ -0,0 +1,16 @@
+import 'package:analyzer/dart/element/element.dart';
+import 'package:floor_generator/model/column.dart';
+
+class Entity {
+  final ClassElement clazz;
+
+  Entity(this.clazz);
+
+  String get name => clazz.displayName;
+
+  List<FieldElement> get fields => clazz.fields;
+
+  List<Column> get columns {
+    return fields.map((field) => Column(field)).toList();
+  }
+}

floor_generator/lib/model/query_method.dart

@@ -1,6 +1,7 @@
 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/misc/constants.dart';
 import 'package:source_gen/source_gen.dart';
 
 /// Raps a method annotated with Query

floor_generator/lib/writer/database_writer.dart

@@ -1,8 +1,7 @@
-import 'package:analyzer/dart/element/element.dart';
 import 'package:floor_generator/misc/annotation_expression.dart';
-import 'package:floor_generator/model/column.dart';
-import 'package:floor_generator/misc/sql_utils.dart';
 import 'package:floor_generator/misc/type_utils.dart';
+import 'package:floor_generator/model/database.dart';
+import 'package:floor_generator/model/entity.dart';
 import 'package:floor_generator/model/query_method.dart';
 import 'package:floor_generator/writer/query_method_writer.dart';
 import 'package:floor_generator/writer/writer.dart';
@@ -17,33 +16,31 @@ class DatabaseWriter implements Writer {
 
   @override
   Spec write() {
-    final database = _getDatabaseClassElement();
+    final database = _getDatabase();
 
     return Library((builder) => builder
       ..body.addAll([
-        _generateOpenDatabaseFunction(database),
+        _generateOpenDatabaseFunction(database.name),
         _generateDatabaseImplementation(database)
       ]));
   }
 
-  ClassElement _getDatabaseClassElement() {
-    final databases = library.classes.where((clazz) =>
+  Database _getDatabase() {
+    final databaseClasses = library.classes.where((clazz) =>
         clazz.isAbstract && clazz.metadata.any(isDatabaseAnnotation));
 
-    if (databases.isEmpty) {
+    if (databaseClasses.isEmpty) {
       throw InvalidGenerationSourceError(
           'No database defined. Add a @Database annotation to your abstract database class.');
-    } else if (databases.length > 1) {
+    } else if (databaseClasses.length > 1) {
       throw InvalidGenerationSourceError(
           'Only one database is allowed. There are too many classes annotated with @Database.');
     } else {
-      return databases.first;
+      return Database(databaseClasses.first);
     }
   }
 
-  Method _generateOpenDatabaseFunction(ClassElement database) {
-    final databaseName = database.displayName;
-
+  Method _generateOpenDatabaseFunction(String databaseName) {
     return Method((builder) => builder
       ..returns = refer('Future<$databaseName>')
       ..name = '_\$open'
@@ -55,9 +52,9 @@ class DatabaseWriter implements Writer {
             '''));
   }
 
-  Class _generateDatabaseImplementation(ClassElement database) {
+  Class _generateDatabaseImplementation(Database database) {
     final createTableStatements =
-        _generateCreateTableSqlStatements(library.classes.toList())
+        _generateCreateTableSqlStatements(database.getEntities(library))
             .map((statement) => 'await database.execute($statement);')
             .join('\n');
 
@@ -66,7 +63,7 @@ class DatabaseWriter implements Writer {
           'There are no entities defined. Use the @Entity annotation on model classes to do so.');
     }
 
-    final databaseName = database.displayName;
+    final databaseName = database.name;
 
     return Class(
       (builder) => builder
@@ -89,77 +86,31 @@ class DatabaseWriter implements Writer {
             );
             ''')),
         )
-        ..methods.addAll(_generateQueryMethods(database.methods)),
+        ..methods.addAll(_generateQueryMethods(database.queryMethods)),
     );
   }
 
-  List<Method> _generateQueryMethods(List<MethodElement> methods) {
-    return _getQueryMethods(methods)
+  List<Method> _generateQueryMethods(List<QueryMethod> queryMethods) {
+    return queryMethods
         .map((queryMethod) => QueryMethodWriter(library, queryMethod).write())
         .toList();
   }
 
-  List<QueryMethod> _getQueryMethods(List<MethodElement> methods) {
-    return methods
-        .where((method) => method.metadata.any(isQueryAnnotation))
-        .map((method) => QueryMethod(method))
-        .toList();
-  }
-
-  List<String> _generateCreateTableSqlStatements(List<ClassElement> classes) {
-    return classes
-        .where((clazz) =>
-            !clazz.isAbstract && clazz.metadata.any(isEntityAnnotation))
-        .map(_generateSql)
-        .toList();
+  List<String> _generateCreateTableSqlStatements(List<Entity> entities) {
+    return entities.map(_generateSql).toList();
   }
 
-  String _generateSql(ClassElement clazz) {
-    final columns =
-        clazz.fields.map((field) => _createColumn(field)).map((column) {
-      String primaryKeyString = '';
+  String _generateSql(Entity entity) {
+    final columns = entity.columns.map((column) {
+      var columnString = '${column.name} ${column.type}';
 
-      if (column.isPrimaryKey) {
-        primaryKeyString += ' ${SqlConstants.PRIMARY_KEY}';
-        if (column.autoGenerate) {
-          primaryKeyString += ' ${SqlConstants.AUTOINCREMENT}';
-        }
+      final additionals = column.additionals;
+      if (additionals != null) {
+        columnString += additionals;
       }
-
-      return '${column.name} ${column.type}$primaryKeyString';
+      return columnString;
     }).join(', ');
 
-    return "'CREATE TABLE IF NOT EXISTS ${clazz.displayName} ($columns)'";
-  }
-
-  Column _createColumn(FieldElement field) {
-    final primaryKeyAnnotations = field.metadata.where(isPrimaryKeyAnnotation);
-
-    bool isPrimaryKey = false;
-    bool autoGenerate;
-
-    if (primaryKeyAnnotations.isNotEmpty) {
-      isPrimaryKey = true;
-      autoGenerate = primaryKeyAnnotations.first
-          .computeConstantValue()
-          .getField('autoGenerate')
-          .toBoolValue();
-    }
-
-    final columnType = getColumnType(field.type);
-
-    if (columnType == null) {
-      throw InvalidGenerationSourceError(
-        'Column type is not supported for ${field.type}.',
-        element: field,
-      );
-    }
-
-    return Column(
-      field.displayName,
-      columnType,
-      isPrimaryKey,
-      autoGenerate,
-    );
+    return "'CREATE TABLE IF NOT EXISTS ${entity.name} ($columns)'";
   }
 }