Skip to content

Commit

Permalink
Add support to apache-airflow-providers-cncf-kubernetes < 7.4.0 (#553)
Browse files Browse the repository at this point in the history
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
  • Loading branch information
tatiana authored Sep 27, 2023
1 parent 877f056 commit 9be9d3b
Show file tree
Hide file tree
Showing 2 changed files with 17 additions and 9 deletions.
11 changes: 7 additions & 4 deletions cosmos/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -28,6 +28,8 @@
DbtTestLocalOperator,
)

logger = get_logger()

try:
from cosmos.operators.docker import (
DbtLSDockerOperator,
Expand Down Expand Up @@ -57,7 +59,8 @@
DbtSnapshotKubernetesOperator,
DbtTestKubernetesOperator,
)
except ImportError:
except ImportError as error:
logger.exception(error)
DbtLSKubernetesOperator = MissingPackage(
"cosmos.operators.kubernetes.DbtLSKubernetesOperator",
"kubernetes",
Expand Down
15 changes: 10 additions & 5 deletions cosmos/operators/kubernetes.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 9be9d3b

Please sign in to comment.