Aggressively cache entry points in process

airflow/providers_manager.py

@@ -31,6 +31,8 @@
 from time import perf_counter
 from typing import TYPE_CHECKING, Any, Callable, MutableMapping, NamedTuple, TypeVar, cast
 
+from packaging.utils import canonicalize_name
+
 from airflow.exceptions import AirflowOptionalProviderFeatureException
 from airflow.typing_compat import Literal
 from airflow.utils import yaml
@@ -467,8 +469,8 @@ def _discover_all_providers_from_packages(self) -> None:
         and verifies only the subset of fields that are needed at runtime.
         """
         for entry_point, dist in entry_points_with_dist("apache_airflow_provider"):
-            package_name = dist.metadata["name"]
-            if self._provider_dict.get(package_name) is not None:
+            package_name = canonicalize_name(dist.metadata["name"])
+            if package_name in self._provider_dict:
                 continue
             log.debug("Loading %s from package %s", entry_point, package_name)
             version = dist.version

airflow/utils/entry_points.py

@@ -16,10 +16,10 @@
 # under the License.
 from __future__ import annotations
 
+import collections
+import functools
 import logging
-from typing import Iterator
-
-from packaging.utils import canonicalize_name
+from typing import Iterator, Tuple
 
 try:
     import importlib_metadata as metadata
@@ -28,26 +28,32 @@
 
 log = logging.getLogger(__name__)
 
+EPnD = Tuple[metadata.EntryPoint, metadata.Distribution]
 
-def entry_points_with_dist(group: str) -> Iterator[tuple[metadata.EntryPoint, metadata.Distribution]]:
-    """Retrieve entry points of the given group.
-
-    This is like the ``entry_points()`` function from importlib.metadata,
-    except it also returns the distribution the entry_point was loaded from.
 
-    :param group: Filter results to only this entrypoint group
-    :return: Generator of (EntryPoint, Distribution) objects for the specified groups
-    """
-    loaded: set[str] = set()
+@functools.lru_cache(maxsize=None)
+def _get_grouped_entry_points() -> dict[str, list[EPnD]]:
+    mapping: dict[str, list[EPnD]] = collections.defaultdict(list)
     for dist in metadata.distributions():
         try:
-            key = canonicalize_name(dist.metadata["Name"])
-            if key in loaded:
-                continue
-            loaded.add(key)
             for e in dist.entry_points:
-                if e.group != group:
-                    continue
-                yield e, dist
+                mapping[e.group].append((e, dist))
         except Exception as e:
             log.warning("Error when retrieving package metadata (skipping it): %s, %s", dist, e)
+    return mapping
+
+
+def entry_points_with_dist(group: str) -> Iterator[EPnD]:
+    """Retrieve entry points of the given group.
+
+    This is like the ``entry_points()`` function from ``importlib.metadata``,
+    except it also returns the distribution the entry point was loaded from.
+
+    Note that this may return multiple distributions to the same package if they
+    are loaded from different ``sys.path`` entries. The caller site should
+    implement appropriate deduplication logic if needed.
+
+    :param group: Filter results to only this entrypoint group
+    :return: Generator of (EntryPoint, Distribution) objects for the specified groups
+    """
+    return iter(_get_grouped_entry_points()[group])

tests/conftest.py

@@ -889,3 +889,10 @@ def _clear_db(request):
             exc_name_parts.insert(0, exc_module)
         extra_msg = "" if request.config.option.db_init else ", try to run with flag --with-db-init"
         pytest.exit(f"Unable clear test DB{extra_msg}, got error {'.'.join(exc_name_parts)}: {ex}")
+
+
+@pytest.fixture(autouse=True)
+def _clear_entry_point_cache():
+    from airflow.utils.entry_points import _get_grouped_entry_points
+
+    _get_grouped_entry_points.cache_clear()

tests/utils/test_entry_points.py

@@ -45,6 +45,6 @@ def distributions(self):
 def test_entry_points_with_dist():
     entries = list(entry_points_with_dist("group_x"))
 
-    # The second "dist2" is ignored. Only "group_x" entries are loaded.
-    assert [dist.metadata["Name"] for _, dist in entries] == ["dist1", "Dist2"]
-    assert [ep.name for ep, _ in entries] == ["a", "e"]
+    # Only "group_x" entries are loaded. Distributions are not deduplicated.
+    assert [dist.metadata["Name"] for _, dist in entries] == ["dist1", "Dist2", "dist2"]
+    assert [ep.name for ep, _ in entries] == ["a", "e", "g"]