Keep import aliases in generated code

drift_dev/CHANGELOG.md

@@ -1,3 +1,8 @@
+## 2.15.1-dev
+
+- Keep import alias when referencing existing elements in generated code
+  ([#2845](https://github.com/simolus3/drift/issues/2845)).
+
 ## 2.15.0
 
 - Potentially __breaking change__: Fix a bug causing `NULL` column constraints

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

@@ -161,7 +161,7 @@ class _DriftBuildRun {
       return;
     }
 
-    _createWriter();
+    await _createWriter();
     if (mode.isMonolithic) {
       await _generateMonolithic(fileResult);
     } else {
@@ -394,10 +394,10 @@ class _DriftBuildRun {
     }
   }
 
-  void _createWriter() {
+  Future<void> _createWriter() async {
     if (mode.isMonolithic) {
       final generationOptions = GenerationOptions(
-        imports: ImportManagerForPartFiles(),
+        imports: ImportManagerForPartFiles(await buildStep.inputLibrary),
       );
       writer = Writer(options, generationOptions: generationOptions);
     } else {

drift_dev/lib/src/cli/commands/schema/generate_utils.dart

@@ -101,7 +101,7 @@ class GenerateUtilsCommand extends Command {
         forSchema: version,
         writeCompanions: companions,
         writeDataClasses: dataClasses,
-        imports: ImportManagerForPartFiles(),
+        imports: NullImportManager(),
       ),
     );
     final file = File(p.join(output.path, _filenameForVersion(version)));

drift_dev/lib/src/writer/import_manager.dart

@@ -1,3 +1,4 @@
+import 'package:analyzer/dart/element/element.dart';
 import 'package:path/path.dart' show url;
 
 import '../utils/string_escaper.dart';
@@ -8,9 +9,41 @@ abstract class ImportManager {
 }
 
 class ImportManagerForPartFiles extends ImportManager {
+  final LibraryElement mainLibrary;
+  final Map<String, Map<String, Element>> _namedImports = {};
+
+  ImportManagerForPartFiles(this.mainLibrary) {
+    for (final import in mainLibrary.libraryImports) {
+      if (import.prefix case ImportElementPrefix prefix) {
+        // Not using import.namespace here because that contains the prefix
+        // everywhere. We want to look up the prefix from the raw name.
+        final library = import.importedLibrary;
+        if (library != null) {
+          _namedImports[prefix.element.name] =
+              library.exportNamespace.definedNames;
+        }
+      }
+    }
+  }
+
+  @override
+  String? prefixFor(Uri definitionUri, String elementName) {
+    // Part files can't add their own imports, so try to find the element in an
+    // existing import.
+    for (final MapEntry(:key, :value) in _namedImports.entries) {
+      if (value.containsKey(elementName)) {
+        return key;
+      }
+    }
+
+    return null;
+  }
+}
+
+class NullImportManager extends ImportManager {
   @override
   String? prefixFor(Uri definitionUri, String elementName) {
-    return null; // todo: Find import alias from existing imports?
+    return null;
   }
 }
 

drift_dev/test/backends/build/build_integration_test.dart

@@ -34,6 +34,43 @@ CREATE INDEX b_idx /* comment should be stripped */ ON b (foo, upper(foo));
     }, result.dartOutputs, result.writer);
   });
 
+  test('keep import aliases', () async {
+    final result = await emulateDriftBuild(
+      inputs: {
+        'a|lib/main.dart': r'''
+import 'package:drift/drift.dart' as drift;
+import 'tables.dart' as tables;
+
+@drift.DriftDatabase(tables: [tables.Texts])
+class MyDatabase extends _$MyDatabase {}
+''',
+        'a|lib/tables.dart': '''
+import 'package:drift/drift.dart';
+
+class Texts extends Table {
+  TextColumn get content => text()();
+}
+''',
+      },
+      logger: loggerThat(neverEmits(anything)),
+    );
+
+    checkOutputs({
+      'a|lib/main.drift.dart': decodedMatches(
+        allOf(
+          contains(
+            r'class $TextsTable extends tables.Texts with '
+            r'drift.TableInfo<$TextsTable, Text>',
+          ),
+          contains(
+            'class Text extends drift.DataClass implements '
+            'drift.Insertable<Text>',
+          ),
+        ),
+      ),
+    }, result.dartOutputs, result.writer);
+  });
+
   test('warns about errors in imports', () async {
     final logger = Logger.detached('build');
     final logs = logger.onRecord.map((e) => e.message).toList();

drift_dev/test/services/schema/writer_test.dart

@@ -108,7 +108,7 @@ class Database {}
         forSchema: 1,
         writeCompanions: true,
         writeDataClasses: true,
-        imports: ImportManagerForPartFiles(),
+        imports: NullImportManager(),
       ),
     );
 

drift_dev/test/writer/queries/query_writer_test.dart

@@ -23,7 +23,7 @@ void main() {
     final writer = Writer(
       options,
       generationOptions: GenerationOptions(
-        imports: ImportManagerForPartFiles(),
+        imports: NullImportManager(),
       ),
     );
     QueryWriter(writer.child())