Read type converters in drift files

drift/test/generated/todos.g.dart

@@ -1730,8 +1730,27 @@ class AllTodosWithCategoryResult extends CustomResultSet {
         .toString();
   }
 }
+
 // DriftElementId(asset:drift/test/generated/todos.dart, users)
 // DriftElementId(asset:drift/test/generated/todos.dart, shared_todos)
 // DriftElementId(asset:drift/test/generated/todos.dart, table_without_p_k)
 // DriftElementId(asset:drift/test/generated/todos.dart, pure_defaults)
-// DriftElementId(asset:drift/test/generated/todos.dart, SomeDao)
+mixin _$SomeDaoMixin on DatabaseAccessor<TodoDb> {
+  $UsersTable get users => attachedDatabase.users;
+  $SharedTodosTable get sharedTodos => attachedDatabase.sharedTodos;
+  $TodosTableTable get todosTable => attachedDatabase.todosTable;
+  $TodoWithCategoryViewView get todoWithCategoryView =>
+      attachedDatabase.todoWithCategoryView;
+  Selectable<TodoEntry> todosForUser({required int user}) {
+    return customSelect(
+        'SELECT t.* FROM todos AS t INNER JOIN shared_todos AS st ON st.todo = t.id INNER JOIN users AS u ON u.id = st.user WHERE u.id = ?1',
+        variables: [
+          Variable<int>(user)
+        ],
+        readsFrom: {
+          todosTable,
+          sharedTodos,
+          users,
+        }).asyncMap(todosTable.mapFromRow);
+  }
+}

drift_dev/lib/src/analysis/backend.dart

@@ -15,6 +15,17 @@ abstract class DriftBackend {
 
   Future<LibraryElement> readDart(Uri uri);
   Future<AstNode?> loadElementDeclaration(Element element);
+
+  /// Resolves a Dart expression from a string.
+  ///
+  /// [context] is a file in which the expression should be resolved, which is
+  /// relevant for relevant imports. [imports] is a list of (relative) imports
+  /// which may be used to resolve the expression.
+  ///
+  /// Throws a [CannotReadExpressionException] if the type could not be
+  /// resolved.
+  Future<Expression> resolveExpression(
+      Uri context, String dartExpression, Iterable<String> imports);
 }
 
 /// Thrown when attempting to read a Dart library from a file that's not a
@@ -25,3 +36,14 @@ class NotALibraryException implements Exception {
 
   NotALibraryException(this.uri);
 }
+
+class CannotReadExpressionException implements Exception {
+  final String msg;
+
+  CannotReadExpressionException(this.msg);
+
+  @override
+  String toString() {
+    return 'Could not read expression: $msg';
+  }
+}

drift_dev/lib/src/analysis/resolver/drift/query.dart

@@ -1,3 +1,5 @@
+import 'package:sqlparser/sqlparser.dart';
+
 import '../../driver/state.dart';
 import '../../results/results.dart';
 import '../intermediate_state.dart';
@@ -14,6 +16,9 @@ class DriftQueryResolver
 
     final source = (file.discovery as DiscoveredDriftFile).originalSource;
 
+    final isCreate =
+        discovered.sqlNode.identifier is SpecialStatementIdentifier;
+
     // Note: We don't analyze the query here, that happens in
     // `file_analysis.dart` after elements have been resolved.
     return DefinedSqlQuery(
@@ -22,6 +27,7 @@ class DriftQueryResolver
       references: references,
       sql: source.substring(stmt.firstPosition, stmt.lastPosition),
       sqlOffset: stmt.firstPosition,
+      mode: isCreate ? QueryMode.atCreate : QueryMode.regular,
     );
   }
 }

drift_dev/lib/src/analysis/resolver/drift/table.dart

@@ -1,7 +1,10 @@
+import 'package:analyzer/dart/ast/ast.dart' as dart;
 import 'package:collection/collection.dart';
+import 'package:drift/drift.dart' show DriftSqlType;
 import 'package:recase/recase.dart';
 import 'package:sqlparser/sqlparser.dart';
 
+import '../../backend.dart';
 import '../../driver/error.dart';
 import '../../results/results.dart';
 import '../intermediate_state.dart';
@@ -11,6 +14,9 @@ import '../shared/data_class.dart';
 import 'find_dart_class.dart';
 
 class DriftTableResolver extends LocalElementResolver<DiscoveredDriftTable> {
+  static final RegExp _enumRegex =
+      RegExp(r'^enum\((\w+)\)$', caseSensitive: false);
+
   DriftTableResolver(super.file, super.discovered, super.resolver, super.state);
 
   @override
@@ -43,7 +49,32 @@ class DriftTableResolver extends LocalElementResolver<DiscoveredDriftTable> {
     for (final column in table.resultColumns) {
       String? overriddenDartName;
       final type = resolver.driver.typeMapping.sqlTypeToDrift(column.type);
+      final nullable = column.type.nullable != false;
       final constraints = <DriftColumnConstraint>[];
+      AppliedTypeConverter? converter;
+
+      final typeName = column.definition?.typeName;
+      final enumMatch =
+          typeName != null ? _enumRegex.firstMatch(typeName) : null;
+      if (enumMatch != null) {
+        final dartTypeName = enumMatch.group(1)!;
+        final imports = file.discovery!.importDependencies.toList();
+        final dartClass = await findDartClass(imports, dartTypeName);
+
+        if (dartClass == null) {
+          reportError(DriftAnalysisError.inDriftFile(
+            column.definition!.typeNames!.toSingleEntity,
+            'Type $dartTypeName could not be found. Are you missing '
+            'an import?',
+          ));
+        } else {
+          converter = readEnumConverter(
+            (msg) =>
+                DriftAnalysisError.inDriftFile(column.definition ?? stmt, msg),
+            dartClass.classElement.thisType,
+          );
+        }
+      }
 
       // columns from virtual tables don't necessarily have a definition, so we
       // can't read the constraints.
@@ -52,6 +83,16 @@ class DriftTableResolver extends LocalElementResolver<DiscoveredDriftTable> {
       for (final constraint in sqlConstraints) {
         if (constraint is DriftDartName) {
           overriddenDartName = constraint.dartName;
+        } else if (constraint is MappedBy) {
+          if (converter != null) {
+            reportError(DriftAnalysisError.inDriftFile(
+                constraint,
+                'Multiple type converters applied to this converter, ignoring '
+                'this one.'));
+            continue;
+          }
+
+          converter = await _readTypeConverter(type, nullable, constraint);
         } else if (constraint is ForeignKeyColumnConstraint) {
           // Note: Warnings about whether the referenced column exists or not
           // are reported later, we just need to know dependencies before the
@@ -89,10 +130,11 @@ class DriftTableResolver extends LocalElementResolver<DiscoveredDriftTable> {
 
       columns.add(DriftColumn(
         sqlType: type,
-        nullable: column.type.nullable != false,
+        nullable: nullable,
         nameInSql: column.name,
         nameInDart: overriddenDartName ?? ReCase(column.name).camelCase,
         constraints: constraints,
+        typeConverter: converter,
         declaration: DriftDeclaration.driftFile(
           column.definition?.nameToken ?? stmt,
           state.ownId.libraryUri,
@@ -196,4 +238,33 @@ class DriftTableResolver extends LocalElementResolver<DiscoveredDriftTable> {
       virtualTableData: virtualTableData,
     );
   }
+
+  Future<AppliedTypeConverter?> _readTypeConverter(
+      DriftSqlType sqlType, bool nullable, MappedBy mapper) async {
+    final code = mapper.mapper.dartCode;
+
+    dart.Expression expression;
+    try {
+      expression = await resolver.driver.backend.resolveExpression(
+        file.ownUri,
+        code,
+        file.discovery!.importDependencies
+            .map((e) => e.toString())
+            .where((e) => e.endsWith('.dart')),
+      );
+    } on CannotReadExpressionException catch (e) {
+      reportError(DriftAnalysisError.inDriftFile(mapper, e.msg));
+      return null;
+    }
+
+    final knownTypes = await resolver.driver.loadKnownTypes();
+    return readTypeConverter(
+      knownTypes.helperLibrary,
+      expression,
+      sqlType,
+      nullable,
+      (msg) => reportError(DriftAnalysisError.inDriftFile(mapper, msg)),
+      knownTypes,
+    );
+  }
 }

drift_dev/lib/src/analysis/results/query.dart

@@ -32,6 +32,7 @@ class DefinedSqlQuery extends DriftElement implements DriftQueryDeclaration {
   final String sql;
 
   final String? resultClassName;
+  final QueryMode mode;
 
   /// The offset of [sql] in the source file, used to properly report errors
   /// later.
@@ -50,9 +51,15 @@ class DefinedSqlQuery extends DriftElement implements DriftQueryDeclaration {
     required this.sql,
     required this.sqlOffset,
     this.resultClassName,
+    this.mode = QueryMode.regular,
   });
 }
 
+enum QueryMode {
+  regular,
+  atCreate,
+}
+
 /// A fully-resolved and analyzed SQL query.
 abstract class SqlQuery {
   final String name;

drift_dev/lib/src/analysis/serializer.dart

@@ -55,6 +55,7 @@ class ElementSerializer {
         'sql': element.sql,
         'offset': element.sqlOffset,
         'result_class': element.resultClassName,
+        'mode': element.mode.name,
       };
     } else if (element is DriftTrigger) {
       additionalInformation = {
@@ -462,6 +463,7 @@ class ElementDeserializer {
           sql: json['sql'] as String,
           sqlOffset: json['offset'] as int,
           resultClassName: json['result_class'] as String?,
+          mode: QueryMode.values.byName(json['mode'] as String),
         );
       case 'trigger':
         DriftTable? on;

drift_dev/lib/src/backends/build/backend.dart

@@ -1,3 +1,5 @@
+import 'dart:convert';
+
 import 'package:analyzer/dart/ast/ast.dart';
 import 'package:analyzer/dart/element/element.dart';
 import 'package:logging/logging.dart';
@@ -6,6 +8,7 @@ import 'package:build/build.dart' as build;
 
 import '../../analysis/backend.dart';
 import '../../analysis/driver/driver.dart';
+import '../../analysis/preprocess_drift.dart';
 
 class DriftBuildBackend extends DriftBackend {
   final BuildStep _buildStep;
@@ -52,6 +55,42 @@ class DriftBuildBackend extends DriftBackend {
   Future<AstNode?> loadElementDeclaration(Element element) {
     return _buildStep.resolver.astNodeFor(element, resolve: true);
   }
+
+  @override
+  Future<Expression> resolveExpression(
+      Uri context, String dartExpression, Iterable<String> imports) async {
+    final original = AssetId.resolve(context);
+    final tempDart = original.changeExtension('.temp.dart');
+    final prepJson = original.changeExtension('.drift_prep.json');
+
+    DriftPreprocessorResult prepResult;
+    try {
+      prepResult = DriftPreprocessorResult.fromJson(
+          json.decode(await _buildStep.readAsString(prepJson))
+              as Map<String, Object?>);
+    } on Exception catch (e, s) {
+      log.warning('Could not read Dart expression $dartExpression', e, s);
+      throw CannotReadExpressionException('Could not load helpers');
+    }
+
+    final getter =
+        prepResult.inlineDartExpressionsToHelperField[dartExpression];
+    if (getter == null) {
+      throw CannotReadExpressionException('No field for $dartExpression');
+    }
+
+    final library = await _buildStep.resolver.libraryFor(tempDart);
+    final field = library.units.first.topLevelVariables
+        .firstWhere((element) => element.name == getter);
+    final fieldAst = await _buildStep.resolver.astNodeFor(field, resolve: true);
+
+    final initializer = (fieldAst as VariableDeclaration).initializer;
+    if (initializer == null) {
+      throw CannotReadExpressionException(
+          'Malformed helper file, this should never happen');
+    }
+    return initializer;
+  }
 }
 
 class BuildCacheReader implements AnalysisResultCacheReader {

drift_dev/lib/src/backends/build/drift_builder.dart

@@ -5,6 +5,7 @@ import '../../analysis/driver/driver.dart';
 import '../../analysis/results/results.dart';
 import '../../analyzer/options.dart';
 import '../../writer/database_writer.dart';
+import '../../writer/drift_accessor_writer.dart';
 import '../../writer/import_manager.dart';
 import '../../writer/writer.dart';
 import 'backend.dart';
@@ -88,22 +89,30 @@ class DriftBuilder extends Builder {
     for (final element in fileResult.analysis.values) {
       final result = element.result;
 
-      if (result is DriftDatabase) {
-        final importedQueries = <String, SqlQuery>{};
+      if (result is BaseDriftAccessor) {
         final resolved = fileResult.fileAnalysis!.resolvedDatabases[result.id]!;
+        final importedQueries = <DefinedSqlQuery, SqlQuery>{};
 
-        // Crawl queries
-        for (final imported in driver.cache.crawlMulti(resolved.knownImports)) {
-          final resolved = await driver.fullyAnalyze(imported.ownUri);
-          importedQueries.addAll({
-            for (final entry in resolved.fileAnalysis!.resolvedQueries.entries)
-              entry.key.name: entry.value,
-          });
+        for (final query
+            in resolved.availableElements.whereType<DefinedSqlQuery>()) {
+          final resolvedFile = await driver.fullyAnalyze(query.id.libraryUri);
+          final resolvedQuery =
+              resolvedFile.fileAnalysis?.resolvedQueries[query.id];
+
+          if (resolvedQuery != null) {
+            importedQueries[query] = resolvedQuery;
+          }
         }
 
-        final input = DatabaseGenerationInput(
-            result, resolved, importedQueries.values.toList());
-        DatabaseWriter(input, writer.child()).write();
+        if (result is DriftDatabase) {
+          final input =
+              DatabaseGenerationInput(result, resolved, importedQueries);
+          DatabaseWriter(input, writer.child()).write();
+        } else if (result is DatabaseAccessor) {
+          final input =
+              AccessorGenerationInput(result, resolved, importedQueries);
+          AccessorWriter(input, writer.child()).write();
+        }
       } else {
         writer.leaf().writeln('// ${element.ownId}');
       }

drift_dev/lib/src/writer/database_writer.dart

@@ -19,7 +19,7 @@ class DatabaseWriter {
   DatabaseGenerationInput input;
   final Scope scope;
 
-  DriftDatabase get db => input.db;
+  DriftDatabase get db => input.accessor;
 
   DatabaseWriter(this.input, this.scope);
 
@@ -127,9 +127,7 @@ class DatabaseWriter {
     }
 
     // Write implementation for query methods
-    final queries = input.resolvedAccessor.definedQueries.values
-        .followedBy(input.importedQueries);
-    for (final query in queries) {
+    for (final query in input.availableRegularQueries) {
       QueryWriter(dbScope.child()).write(query);
     }
 
@@ -144,11 +142,13 @@ class DatabaseWriter {
       ..write('=> [');
 
     schemaScope
-      ..write(db.references.map((e) {
-//        if (e is SpecialQuery) {
-//          final sql = e.formattedSql(scope.options);
-//          return 'OnCreateQuery(${asDartLiteral(sql)})';
-//        }
+      ..write(elements.map((e) {
+        if (e is DefinedSqlQuery && e.mode == QueryMode.atCreate) {
+          final resolved = input.importedQueries[e]!;
+          final sql = schemaScope.sqlCode(resolved.root!);
+
+          return 'OnCreateQuery(${asDartLiteral(sql)})';
+        }
 
         return entityGetters[e];
       }).join(', '))
@@ -192,15 +192,26 @@ class DatabaseWriter {
   }
 }
 
-class DatabaseGenerationInput {
-  final DriftDatabase db;
+class GenerationInput<T extends BaseDriftAccessor> {
+  final T accessor;
   final ResolvedDatabaseAccessor resolvedAccessor;
+  final Map<DefinedSqlQuery, SqlQuery> importedQueries;
 
-  final List<SqlQuery> importedQueries;
+  GenerationInput(this.accessor, this.resolvedAccessor, this.importedQueries);
 
-  DatabaseGenerationInput(this.db, this.resolvedAccessor, this.importedQueries);
+  /// All locally-defined and imported [SqlQuery] elements that are regular
+  /// queries (so no query with [QueryMode.atCreate]).
+  Iterable<SqlQuery> get availableRegularQueries {
+    final imported = importedQueries.entries
+        .where((entry) => entry.key.mode == QueryMode.regular)
+        .map((e) => e.value);
+    return resolvedAccessor.definedQueries.values.followedBy(imported);
+  }
 }
 
+typedef DatabaseGenerationInput = GenerationInput<DriftDatabase>;
+typedef AccessorGenerationInput = GenerationInput<DatabaseAccessor>;
+
 extension on drift.UpdateRule {
   void writeConstructor(TextEmitter emitter) {
     if (this is drift.WritePropagation) {

drift_dev/lib/src/writer/drift_accessor_writer.dart

@@ -0,0 +1,34 @@
+import '../analysis/results/results.dart';
+import 'database_writer.dart';
+import 'queries/query_writer.dart';
+import 'writer.dart';
+
+class AccessorWriter {
+  final AccessorGenerationInput input;
+  final Scope scope;
+
+  AccessorWriter(this.input, this.scope);
+
+  void write() {
+    final classScope = scope.child();
+
+    final daoName = input.accessor.declaration.name!;
+    final dbTypeName = classScope.dartCode(input.accessor.databaseClass);
+    classScope.leaf().write('mixin _\$${daoName}Mixin on '
+        'DatabaseAccessor<$dbTypeName> {\n');
+
+    for (final entity in input.resolvedAccessor.availableElements
+        .whereType<DriftElementWithResultSet>()) {
+      final infoType = entity.entityInfoName;
+      final getterName = entity.dbGetterName;
+      classScope.leaf().write(
+          '$infoType get $getterName => attachedDatabase.$getterName;\n');
+    }
+
+    for (final query in input.availableRegularQueries) {
+      QueryWriter(classScope.child()).write(query);
+    }
+
+    classScope.leaf().write('}');
+  }
+}

drift_dev/lib/src/writer/tables/view_writer.dart

@@ -81,6 +81,7 @@ class ViewWriter extends TableOrViewWriter {
       } else {
         emitter.write(asDartLiteral(source.createView));
       }
+      buffer.writeln(';');
     } else {
       buffer.write('@override\n String? get createViewStmt => null;\n');
     }