diff --git a/drift_dev/CHANGELOG.md b/drift_dev/CHANGELOG.md
index 25254fd77..59efbceea 100644
--- a/drift_dev/CHANGELOG.md
+++ b/drift_dev/CHANGELOG.md
@@ -1,6 +1,7 @@
 ## 2.18.0-dev
 
 - Add support for the `geopoly` extension in drift files.
+- Improve finding the correct import alias in generated code for part files.
 
 ## 2.17.0
 
diff --git a/drift_dev/lib/src/writer/import_manager.dart b/drift_dev/lib/src/writer/import_manager.dart
index 3a8ec70de..23a032159 100644
--- a/drift_dev/lib/src/writer/import_manager.dart
+++ b/drift_dev/lib/src/writer/import_manager.dart
@@ -40,12 +40,26 @@ class ImportManagerForPartFiles extends ImportManager {
     return null;
   }
 
+  /// Heuristic to determine whether a source uri [wanted] likely exports the
+  /// [target] element.
+  ///
+  /// We can't compare the [target] definition with the [wanted] url directly,
+  /// as many parts use URLs relying on re-exports. For instance, this should
+  /// return true for a wanted URI of `package:drift/drift.dart` when the
+  /// element is actually defined in `package:drift/src/runtime/table.dart`.
   static bool _matchingUrl(Uri wanted, Element target) {
     final targetUri = target.librarySource?.uri;
     if (targetUri == null || targetUri.scheme != wanted.scheme) {
       return false;
     }
 
+    if (targetUri.scheme == 'package') {
+      // Match if the two elements are coming from the same package
+      final targetPackage = targetUri.pathSegments.first;
+      final wantedPackage = wanted.pathSegments.first;
+      return targetPackage == wantedPackage;
+    }
+
     return true;
   }
 }