From 7274e30e6d780e812ba7ea76489fc650e759525a Mon Sep 17 00:00:00 2001 From: jason810496 Date: Thu, 2 Jan 2025 00:17:35 +0800 Subject: [PATCH] Fix Assets import along with Airflow version --- .../src/airflow/providers/common/compat/lineage/hook.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/providers/src/airflow/providers/common/compat/lineage/hook.py b/providers/src/airflow/providers/common/compat/lineage/hook.py index e6a60f19a8607..7e0f0090edaf5 100644 --- a/providers/src/airflow/providers/common/compat/lineage/hook.py +++ b/providers/src/airflow/providers/common/compat/lineage/hook.py @@ -25,15 +25,20 @@ from airflow.providers.common.compat.version_compat import AIRFLOW_V_2_10_PLUS, AIRFLOW_V_3_0_PLUS from airflow.providers_manager import ProvidersManager -from airflow.sdk.definitions.asset import Asset from airflow.utils.log.logging_mixin import LoggingMixin if TYPE_CHECKING: from airflow.hooks.base import BaseHook from airflow.io.path import ObjectStoragePath + from airflow.sdk.definitions.asset import Asset # Store context what sent lineage. LineageContext = Union[BaseHook, ObjectStoragePath] +else: + if AIRFLOW_V_3_0_PLUS: + from airflow.sdk.definitions.asset import Asset + else: + from airflow.datasets import Dataset as Asset _hook_lineage_collector: HookLineageCollector | None = None