diff --git a/cosmos/__init__.py b/cosmos/__init__.py index c0f00e3a5..fc58a7b99 100644 --- a/cosmos/__init__.py +++ b/cosmos/__init__.py @@ -9,15 +9,15 @@ from cosmos.airflow.dag import DbtDag from cosmos.airflow.task_group import DbtTaskGroup -from cosmos.constants import LoadMode, TestBehavior, ExecutionMode -from cosmos.operators.lazy_load import MissingPackage from cosmos.config import ( ProjectConfig, ProfileConfig, ExecutionConfig, RenderConfig, ) - +from cosmos.constants import LoadMode, TestBehavior, ExecutionMode +from cosmos.log import get_logger +from cosmos.operators.lazy_load import MissingPackage from cosmos.operators.local import ( DbtDepsLocalOperator, DbtLSLocalOperator, @@ -28,6 +28,8 @@ DbtTestLocalOperator, ) +logger = get_logger() + try: from cosmos.operators.docker import ( DbtLSDockerOperator, @@ -57,7 +59,8 @@ DbtSnapshotKubernetesOperator, DbtTestKubernetesOperator, ) -except ImportError: +except ImportError as error: + logger.exception(error) DbtLSKubernetesOperator = MissingPackage( "cosmos.operators.kubernetes.DbtLSKubernetesOperator", "kubernetes", diff --git a/cosmos/operators/kubernetes.py b/cosmos/operators/kubernetes.py index 38ca47452..b55dcdbe3 100644 --- a/cosmos/operators/kubernetes.py +++ b/cosmos/operators/kubernetes.py @@ -13,17 +13,22 @@ logger = get_logger(__name__) -# kubernetes is an optional dependency, so we need to check if it's installed try: + # apache-airflow-providers-cncf-kubernetes >= 7.4.0 from airflow.providers.cncf.kubernetes.backcompat.backwards_compat_converters import ( convert_env_vars, ) from airflow.providers.cncf.kubernetes.operators.pod import KubernetesPodOperator except ImportError: - raise ImportError( - "Could not import KubernetesPodOperator. Ensure you've installed the Kubernetes provider " - "separately or with with `pip install astronomer-cosmos[...,kubernetes]`." - ) + try: + # apache-airflow-providers-cncf-kubernetes < 7.4.0 + from airflow.providers.cncf.kubernetes.operators.kubernetes_pod import KubernetesPodOperator + except ImportError as error: + logger.exception(error) + raise ImportError( + "Could not import KubernetesPodOperator. Ensure you've installed the Kubernetes provider " + "separately or with with `pip install astronomer-cosmos[...,kubernetes]`." + ) class DbtKubernetesBaseOperator(KubernetesPodOperator, DbtBaseOperator): # type: ignore