From 9be9d3b11d7849efc5d766c83ed3839cf49bcaba Mon Sep 17 00:00:00 2001
From: Tatiana Al-Chueyr <tatiana.alchueyr@gmail.com>
Date: Wed, 27 Sep 2023 10:52:10 +0100
Subject: [PATCH] Add support to apache-airflow-providers-cncf-kubernetes <
 7.4.0 (#553)

Try to import `KubernetesPodOperator` from
`airflow.providers.cncf.kubernetes.operators.kubernetes_pod`
(`apache-airflow-providers-cncf-kubernetes` < 7.4.0), if it is not
available in `airflow.providers.cncf.kubernetes.operators.pod`
(`apache-airflow-providers-cncf-kubernetes` >= 7.4.0).

Users using `apache-airflow-providers-cncf-kubernetes < 7.4.0` would get
the error message:
```
Could not import KubernetesPodOperator. Ensure you've installed the Kubernetes provider
separately or with with `pip install astronomer-cosmos[...,kubernetes]`.
```

As reported in the #airflow-dbt channel in the Apache Airflow Slack, by
Aditya Sharma:
https://apache-airflow.slack.com/archives/C059CC42E9W/p1695076622562159
---
 cosmos/__init__.py             | 11 +++++++----
 cosmos/operators/kubernetes.py | 15 ++++++++++-----
 2 files changed, 17 insertions(+), 9 deletions(-)

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