fix(providers): use version compare to decide whether to import asset

apache · Nov 9, 2024 · b3a9cf3 · b3a9cf3
1 parent f11358f
commit b3a9cf3
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 5 deletions.
diff --git a/providers/src/airflow/providers/common/io/assets/file.py b/providers/src/airflow/providers/common/io/assets/file.py
@@ -19,9 +19,15 @@
 import urllib.parse
 from typing import TYPE_CHECKING
 
-try:
-    from airflow.sdk.definitions.asset import Asset
-except ModuleNotFoundError:
+from packaging.version import Version
+
+from airflow import __version__ as AIRFLOW_VERSION
+
+# TODO: Remove version check block after bumpping common provider to 1.3.0
+AIRFLOW_V_3_0_PLUS = Version(Version(AIRFLOW_VERSION).base_version) >= Version("3.0.0")
+if AIRFLOW_V_3_0_PLUS:
+    pass
+else:
     from airflow.datasets import Dataset as Asset  # type: ignore[no-redef]
 
 if TYPE_CHECKING:

diff --git a/providers/src/airflow/providers/openlineage/utils/utils.py b/providers/src/airflow/providers/openlineage/utils/utils.py
@@ -707,9 +707,15 @@ def translate_airflow_asset(asset: Asset, lineage_context) -> OpenLineageDataset
     This function returns None if no URI normalizer is defined, no asset converter is found or
     some core Airflow changes are missing and ImportError is raised.
     """
-    try:
+    # TODO: Remove version check block after bumpping common provider to 1.3.0
+    from packaging.version import Version
+
+    from airflow import __version__ as AIRFLOW_VERSION
+
+    AIRFLOW_V_3_0_PLUS = Version(Version(AIRFLOW_VERSION).base_version) >= Version("3.0.0")
+    if AIRFLOW_V_3_0_PLUS:
         from airflow.sdk.definitions.asset import _get_normalized_scheme
-    except ModuleNotFoundError:
+    else:
         try:
             from airflow.datasets import _get_normalized_scheme  # type: ignore[no-redef, attr-defined]
         except ImportError: