From 863dbb61a9a8b3b1769cb68a81fdffcd7e11fad0 Mon Sep 17 00:00:00 2001
From: Simon Binder <oss@simonbinder.eu>
Date: Thu, 29 Feb 2024 18:08:51 +0100
Subject: [PATCH] Avoid picking wrong import alias

Closes #2904
---
 drift_dev/CHANGELOG.md                              |  4 ++++
 drift_dev/lib/src/writer/import_manager.dart        | 12 +++++++++++-
 .../test/backends/build/build_integration_test.dart | 13 +++++++------
 3 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/drift_dev/CHANGELOG.md b/drift_dev/CHANGELOG.md
index ad3bf9bc6..c5ae4216f 100644
--- a/drift_dev/CHANGELOG.md
+++ b/drift_dev/CHANGELOG.md
@@ -1,3 +1,7 @@
+## 2.17.0-dev
+
+- Fix drift using the wrong import alias in generated part files.
+
 ## 2.16.0
 
 - Keep import alias when referencing existing elements in generated code
diff --git a/drift_dev/lib/src/writer/import_manager.dart b/drift_dev/lib/src/writer/import_manager.dart
index 24e71bab4..3a8ec70de 100644
--- a/drift_dev/lib/src/writer/import_manager.dart
+++ b/drift_dev/lib/src/writer/import_manager.dart
@@ -31,13 +31,23 @@ class ImportManagerForPartFiles extends ImportManager {
     // 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)) {
+      final foundHere = value[elementName];
+      if (foundHere != null && _matchingUrl(definitionUri, foundHere)) {
         return key;
       }
     }
 
     return null;
   }
+
+  static bool _matchingUrl(Uri wanted, Element target) {
+    final targetUri = target.librarySource?.uri;
+    if (targetUri == null || targetUri.scheme != wanted.scheme) {
+      return false;
+    }
+
+    return true;
+  }
 }
 
 class NullImportManager extends ImportManager {
diff --git a/drift_dev/test/backends/build/build_integration_test.dart b/drift_dev/test/backends/build/build_integration_test.dart
index 6833f4278..d81f397f6 100644
--- a/drift_dev/test/backends/build/build_integration_test.dart
+++ b/drift_dev/test/backends/build/build_integration_test.dart
@@ -38,16 +38,17 @@ CREATE INDEX b_idx /* comment should be stripped */ ON b (foo, upper(foo));
     final result = await emulateDriftBuild(
       inputs: {
         'a|lib/main.dart': r'''
+import 'dart:io' as io;
 import 'package:drift/drift.dart' as drift;
 import 'tables.dart' as tables;
 
-@drift.DriftDatabase(tables: [tables.Texts])
+@drift.DriftDatabase(tables: [tables.Files])
 class MyDatabase extends _$MyDatabase {}
 ''',
         'a|lib/tables.dart': '''
 import 'package:drift/drift.dart';
 
-class Texts extends Table {
+class Files extends Table {
   TextColumn get content => text()();
 }
 ''',
@@ -59,12 +60,12 @@ class Texts extends Table {
       'a|lib/main.drift.dart': decodedMatches(
         allOf(
           contains(
-            r'class $TextsTable extends tables.Texts with '
-            r'drift.TableInfo<$TextsTable, Text>',
+            r'class $FilesTable extends tables.Files with '
+            r'drift.TableInfo<$FilesTable, File>',
           ),
           contains(
-            'class Text extends drift.DataClass implements '
-            'drift.Insertable<Text>',
+            'class File extends drift.DataClass implements '
+            'drift.Insertable<File>',
           ),
         ),
       ),