Fix extending classes in modular mode (#2260)

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

@@ -5,6 +5,8 @@ import 'package:path/path.dart' show url;
 import 'package:recase/recase.dart';
 import 'package:sqlparser/sqlparser.dart';
 
+import 'dart.dart';
+
 part '../../generated/analysis/results/element.g.dart';
 
 @sealed
@@ -87,9 +89,9 @@ abstract class DriftElement {
 
   /// If this element was extracted from a defined Dart class, returns the name
   /// of that class.
-  String? get definingDartClass {
+  AnnotatedDartCode? get definingDartClass {
     if (id.isDefinedInDart) {
-      return declaration.name;
+      return AnnotatedDartCode.importedSymbol(id.libraryUri, declaration.name!);
     }
     return null;
   }

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

@@ -342,13 +342,19 @@ class TableWriter extends TableOrViewWriter {
     } else {
       // Regular generation, write full table class
       final dataClass = emitter.dartCode(emitter.writer.rowType(table));
-      final tableDslName = table.definingDartClass ?? emitter.drift('Table');
+      final tableDslName = table.definingDartClass ??
+          AnnotatedDartCode.importedSymbol(AnnotatedDartCode.drift, 'Table');
 
       // class UsersTable extends Users implements TableInfo<Users, User> {
       final typeArgs = '<${table.entityInfoName}, $dataClass>';
 
-      buffer.write('class ${table.entityInfoName} extends $tableDslName with '
+      emitter
-          '${emitter.drift('TableInfo')}$typeArgs ');
+        ..write('class ${table.entityInfoName} extends ')
+        ..writeDart(tableDslName)
+        ..write(' with ')
+        ..writeDart(AnnotatedDartCode.importedSymbol(
+            AnnotatedDartCode.drift, 'TableInfo'))
+        ..write(typeArgs);
 
       if (table.isVirtual) {
         buffer.write(', ${emitter.drift('VirtualTableInfo')}$typeArgs ');

drift_dev/test/analysis/resolver/dart/view_test.dart

@@ -68,10 +68,10 @@ abstract class TodoItemWithCategoryNameView extends View {
         result.analysis.values.map((e) => e.result).whereType<DriftView>();
     expect(views, hasLength(2));
 
-    final todoCategoryItemCount = views
+    final todoCategoryItemCount = views.singleWhere(
-        .singleWhere((e) => e.definingDartClass == 'TodoCategoryItemCount');
+        (e) => e.definingDartClass.toString() == 'TodoCategoryItemCount');
-    final todoItemWithCategoryName = views.singleWhere(
+    final todoItemWithCategoryName = views.singleWhere((e) =>
-        (e) => e.definingDartClass == 'TodoItemWithCategoryNameView');
+        e.definingDartClass.toString() == 'TodoItemWithCategoryNameView');
 
     expect(
         todoCategoryItemCount.source,

drift_dev/test/writer/table_writer_test.dart

@@ -0,0 +1,32 @@
+import 'package:build_test/build_test.dart';
+import 'package:test/test.dart';
+
+import '../utils.dart';
+
+void main() {
+  test('references right parent class for Dart-defined tables', () async {
+    final result = await emulateDriftBuild(
+      inputs: {
+        'a|lib/a.dart': '''
+import 'package:drift/drift.dart';
+
+class Tags extends Table {
+  IntColumn get id => integer().autoIncrement()();
+}
+''',
+      },
+      modularBuild: true,
+    );
+
+    checkOutputs({
+      'a|lib/a.drift.dart': decodedMatches(
+        allOf(
+          contains("import 'package:a/a.dart' as i2;"),
+          contains(
+            r'class $TagsTable extends i2.Tags with i0.TableInfo<$TagsTable, i1.Tag>',
+          ),
+        ),
+      ),
+    }, result.dartOutputs, result);
+  });
+}

examples/modular/build.yaml

@@ -2,6 +2,8 @@ targets:
   $default:
     builders:
       drift_dev:
+        # Disable the default builder in favor of the modular builders configured
+        # below.
         enabled: false
 
       drift_dev:analyzer: