Add import references to generated insertable

drift_dev/CHANGELOG.md

@@ -2,6 +2,8 @@
 
 - Reduce the amount of assets read by drift, improving build performance and enabling faster
   incremental rebuilds.
+- Fix missing import references around `@UseRowClass` with `generateInsertable: true` when
+  modular code generation is enabled.
 
 ## 2.11.0
 

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

@@ -223,7 +223,7 @@ class UpdateCompanionWriter {
 
   void _writeToString() {
     overrideToString(
-      _emitter.dartCode(_emitter.companionType(table)),
+      _emitter.companionType(table).toString(),
       [for (final column in columns) column.nameInDart],
       _buffer,
     );
@@ -235,17 +235,31 @@ class UpdateCompanionWriter {
     final info = table.existingRowClass;
     if (info == null) return;
 
-    final rowClass = _emitter.rowClass(table).toString();
+    final rowClass = _emitter.rowClass(table);
-    final rowType = _emitter.dartCode(_emitter.rowType(table));
     final insertableClass = '_\$${rowClass}Insertable';
 
-    _buffer.write('class $insertableClass implements '
+    _emitter
-        'Insertable<$rowType> {\n'
+      // Class _$RowInsertable implements Insertable<RowClass> {
-        '$rowType _object;\n\n'
+      ..write('class $insertableClass implements ')
-        '$insertableClass(this._object);\n\n'
+      ..writeDriftRef('Insertable')
-        '@override\n'
+      ..write('<')
-        'Map<String, Expression> toColumns(bool nullToAbsent) {\n'
+      ..writeDart(rowClass)
-        'return $_companionType(\n');
+      ..writeln('> {')
+      // Field to RowClass and constructor
+      ..writeDart(rowClass)
+      ..writeln(' _object;')
+      ..writeln('$insertableClass(this._object);')
+      // Map<String, Expression> toColumns(bool nullToAbsent) {
+      ..writeln('@override')
+      ..writeUriRef(AnnotatedDartCode.dartCore, 'Map')
+      ..write('<')
+      ..writeUriRef(AnnotatedDartCode.dartCore, 'String')
+      ..write(', ')
+      ..writeUriRef(AnnotatedDartCode.drift, 'Expression')
+      ..write('> toColumns(')
+      ..writeUriRef(AnnotatedDartCode.dartCore, 'bool')
+      ..writeln(' nullToAbsent) {')
+      ..writeln('return $_companionType(');
 
     final columns = info.positionalColumns.followedBy(info.namedColumns.values);
     for (final columnName in columns) {
@@ -254,14 +268,18 @@ class UpdateCompanionWriter {
 
       if (column != null && !column.isGenerated) {
         final dartName = column.nameInDart;
-        _buffer.write('$dartName: Value (_object.$dartName),\n');
+        _emitter
+          ..write('$dartName: ')
+          ..writeDriftRef('Value')
+          ..writeln('(_object.$dartName),');
       }
     }
 
-    _buffer
+    _emitter
       ..write(').toColumns(false);\n}\n}\n\n')
-      ..write('extension ${rowClass}ToInsertable '
+      ..write('extension ${rowClass}ToInsertable on ')
-          'on $rowType {')
+      ..writeDart(rowClass)
+      ..writeln('{')
       ..write('$insertableClass toInsertable() {\n')
       ..write('return $insertableClass(this);\n')
       ..write('}\n}\n');

drift_dev/test/writer/data_class_writer_test.dart

@@ -212,6 +212,61 @@ class MyTable extends Table {
   }''')),
     }, result.dartOutputs, result.writer);
   });
+
+  test('generates correct companions for modular row classes', () async {
+    final result = await emulateDriftBuild(
+      inputs: {
+        'a|lib/a.dart': '''
+import 'package:drift/drift.dart';
+
+@UseRowClass(Item, generateInsertable: true)
+class ItemTable extends Table {
+  IntColumn get id => integer().autoIncrement()();
+}
+
+class Item {
+  final int id;
+  Item(this.id);
+}
+''',
+      },
+      modularBuild: true,
+    );
+
+    checkOutputs(
+      {
+        'a|lib/a.drift.dart': decodedMatches(
+          allOf(
+            // The toString() definition for companions was broken and included
+            // the import prefix of the companion.
+            contains("StringBuffer('ItemTableCompanion(')"),
+
+            // The extension should also reference the row class correctly
+            contains(r'''
+class _$ItemInsertable implements i0.Insertable<i1.Item> {
+  i1.Item _object;
+  _$ItemInsertable(this._object);
+  @override
+  Map<String, i0.Expression> toColumns(bool nullToAbsent) {
+    return i2.ItemTableCompanion(
+      id: i0.Value(_object.id),
+    ).toColumns(false);
+  }
+}
+
+extension ItemToInsertable on i1.Item {
+  _$ItemInsertable toInsertable() {
+    return _$ItemInsertable(this);
+  }
+}
+'''),
+          ),
+        ),
+      },
+      result.dartOutputs,
+      result.writer,
+    );
+  });
 }
 
 class _GeneratesConstDataClasses extends Matcher {

examples/modular/lib/src/posts.drift.dart

@@ -205,7 +205,7 @@ class PostsCompanion extends i0.UpdateCompanion<i1.Post> {
 
   @override
   String toString() {
-    return (StringBuffer('i1.PostsCompanion(')
+    return (StringBuffer('PostsCompanion(')
           ..write('id: $id, ')
           ..write('author: $author, ')
           ..write('content: $content')
@@ -393,7 +393,7 @@ class LikesCompanion extends i0.UpdateCompanion<i1.Like> {
 
   @override
   String toString() {
-    return (StringBuffer('i1.LikesCompanion(')
+    return (StringBuffer('LikesCompanion(')
           ..write('post: $post, ')
           ..write('likedBy: $likedBy, ')
           ..write('rowid: $rowid')

examples/modular/lib/src/search.drift.dart

@@ -194,7 +194,7 @@ class SearchInPostsCompanion extends i0.UpdateCompanion<i1.SearchInPost> {
 
   @override
   String toString() {
-    return (StringBuffer('i1.SearchInPostsCompanion(')
+    return (StringBuffer('SearchInPostsCompanion(')
           ..write('author: $author, ')
           ..write('content: $content, ')
           ..write('rowid: $rowid')

examples/modular/lib/src/users.drift.dart

@@ -302,7 +302,7 @@ class UsersCompanion extends i0.UpdateCompanion<i1.User> {
 
   @override
   String toString() {
-    return (StringBuffer('i1.UsersCompanion(')
+    return (StringBuffer('UsersCompanion(')
           ..write('id: $id, ')
           ..write('name: $name, ')
           ..write('biography: $biography, ')
@@ -500,7 +500,7 @@ class FollowsCompanion extends i0.UpdateCompanion<i1.Follow> {
 
   @override
   String toString() {
-    return (StringBuffer('i1.FollowsCompanion(')
+    return (StringBuffer('FollowsCompanion(')
           ..write('followed: $followed, ')
           ..write('follower: $follower, ')
           ..write('rowid: $rowid')