Consistently handle build and imports in cosmos/__init__.py (#1215)

While reviewing #1153, I noticed we were inconsistent with Build operators, especially how we import and expose them in Cosmos `__init__.py` This PR addresses this. These were the changes introduced: * Add `DbtBuildAzureContainerInstanceOperator` * Expose the build operators for all execution modes in Cosmos `__init__` (missing for most, except `ExecutionMode.LOCAL`) * Be consistent with import error messages in `__init__` (we were missing AWS EKS)
astronomer · Sep 23, 2024 · 0fd0bdd · 0fd0bdd
1 parent 11de5ba
commit 0fd0bdd
Show file tree

Hide file tree

Showing 3 changed files with 87 additions and 21 deletions.
diff --git a/cosmos/__init__.py b/cosmos/__init__.py
@@ -42,6 +42,7 @@
 
 try:
     from cosmos.operators.docker import (
+        DbtBuildDockerOperator,
         DbtLSDockerOperator,
         DbtRunDockerOperator,
         DbtRunOperationDockerOperator,
@@ -62,6 +63,7 @@
 
 try:
     from cosmos.operators.kubernetes import (
+        DbtBuildKubernetesOperator,
         DbtLSKubernetesOperator,
         DbtRunKubernetesOperator,
         DbtRunOperationKubernetesOperator,
@@ -71,6 +73,10 @@
     )
 except ImportError:
     logger.debug("To import Kubernetes modules, install astronomer-cosmos[kubernetes].", stack_info=True)
+    DbtBuildKubernetesOperator = MissingPackage(
+        "cosmos.operators.kubernetes.DbtBuildKubernetesOperator",
+        "kubernetes",
+    )
     DbtLSKubernetesOperator = MissingPackage(
         "cosmos.operators.kubernetes.DbtLSKubernetesOperator",
         "kubernetes",
@@ -98,6 +104,7 @@
 
 try:
     from cosmos.operators.azure_container_instance import (
+        DbtBuildAzureContainerInstanceOperator,
         DbtLSAzureContainerInstanceOperator,
         DbtRunAzureContainerInstanceOperator,
         DbtRunOperationAzureContainerInstanceOperator,
@@ -106,6 +113,9 @@
         DbtTestAzureContainerInstanceOperator,
     )
 except ImportError:
+    DbtBuildAzureContainerInstanceOperator = MissingPackage(
+        "cosmos.operators.azure_container_instance.DbtBuildAzureContainerInstanceOperator", "azure-container-instance"
+    )
     DbtLSAzureContainerInstanceOperator = MissingPackage(
         "cosmos.operators.azure_container_instance.DbtLSAzureContainerInstanceOperator", "azure-container-instance"
     )
@@ -127,48 +137,90 @@
         "cosmos.operators.azure_container_instance.DbtTestAzureContainerInstanceOperator", "azure-container-instance"
     )
 
+
+try:
+    from cosmos.operators.aws_eks import (
+        DbtBuildAwsEksOperator,
+        DbtLSAwsEksOperator,
+        DbtRunAwsEksOperator,
+        DbtRunOperationAwsEksOperator,
+        DbtSeedAwsEksOperator,
+        DbtSnapshotAwsEksOperator,
+        DbtTestAwsEksOperator,
+    )
+except ImportError:
+    DbtBuildAwsEksOperator = MissingPackage(
+        "cosmos.operators.azure_container_instance.DbtBuildAwsEksOperator", "aws_eks"
+    )
+    DbtLSAwsEksOperator = MissingPackage("cosmos.operators.azure_container_instance.DbtLSAwsEksOperator", "aws_eks")
+    DbtRunAwsEksOperator = MissingPackage("cosmos.operators.azure_container_instance.DbtRunAwsEksOperator", "aws_eks")
+    DbtRunOperationAwsEksOperator = MissingPackage(
+        "cosmos.operators.azure_container_instance.DbtRunOperationAwsEksOperator",
+        "aws_eks",
+    )
+    DbtSeedAwsEksOperator = MissingPackage("cosmos.operators.azure_container_instance.DbtSeedAwsEksOperator", "aws_eks")
+    DbtSnapshotAwsEksOperator = MissingPackage(
+        "cosmos.operators.azure_container_instance.DbtSnapshotAwsEksOperator",
+        "aws_eks",
+    )
+    DbtTestAwsEksOperator = MissingPackage("cosmos.operators.azure_container_instance.DbtTestAwsEksOperator", "aws_eks")
+
+
 __all__ = [
     "ProjectConfig",
     "ProfileConfig",
     "ExecutionConfig",
     "RenderConfig",
+    "DbtDag",
+    "DbtTaskGroup",
+    "ExecutionMode",
+    "LoadMode",
+    "TestBehavior",
+    "InvocationMode",
+    "TestIndirectSelection",
+    "SourceRenderingBehavior",
+    "DbtResourceType",
+    # Local Execution Mode
+    "DbtBuildLocalOperator",
+    "DbtDepsLocalOperator",  # deprecated, to be delete in Cosmos 2.x
     "DbtLSLocalOperator",
-    "DbtRunOperationLocalOperator",
     "DbtRunLocalOperator",
+    "DbtRunOperationLocalOperator",
     "DbtSeedLocalOperator",
-    "DbtTestLocalOperator",
-    "DbtBuildLocalOperator",
-    "DbtDepsLocalOperator",
     "DbtSnapshotLocalOperator",
-    "DbtDag",
-    "DbtTaskGroup",
+    "DbtTestLocalOperator",
+    # Docker Execution Mode
+    "DbtBuildDockerOperator",
     "DbtLSDockerOperator",
-    "DbtRunOperationDockerOperator",
     "DbtRunDockerOperator",
+    "DbtRunOperationDockerOperator",
     "DbtSeedDockerOperator",
-    "DbtTestDockerOperator",
-    "DbtBuildDockerOperator",
     "DbtSnapshotDockerOperator",
+    "DbtTestDockerOperator",
+    # Kubernetes Execution Mode
+    "DbtBuildKubernetesOperator",
     "DbtLSKubernetesOperator",
-    "DbtRunOperationKubernetesOperator",
     "DbtRunKubernetesOperator",
+    "DbtRunOperationKubernetesOperator",
     "DbtSeedKubernetesOperator",
-    "DbtTestKubernetesOperator",
-    "DbtBuildKubernetesOperator",
     "DbtSnapshotKubernetesOperator",
+    "DbtTestKubernetesOperator",
+    # Azure Container Instance Execution Mode
+    "DbtBuildAzureContainerInstanceOperator",
     "DbtLSAzureContainerInstanceOperator",
-    "DbtRunOperationAzureContainerInstanceOperator",
     "DbtRunAzureContainerInstanceOperator",
+    "DbtRunOperationAzureContainerInstanceOperator",
     "DbtSeedAzureContainerInstanceOperator",
-    "DbtTestAzureContainerInstanceOperator",
     "DbtSnapshotAzureContainerInstanceOperator",
-    "ExecutionMode",
-    "LoadMode",
-    "TestBehavior",
-    "InvocationMode",
-    "TestIndirectSelection",
-    "SourceRenderingBehavior",
-    "DbtResourceType",
+    "DbtTestAzureContainerInstanceOperator",
+    # AWS EKS Execution Mode
+    "DbtBuildAwsEksOperator",
+    "DbtLSAwsEksOperator",
+    "DbtRunAwsEksOperator",
+    "DbtRunOperationAwsEksOperator",
+    "DbtSeedAwsEksOperator",
+    "DbtSnapshotAwsEksOperator",
+    "DbtTestAwsEksOperator",
 ]
 
 """

diff --git a/cosmos/operators/azure_container_instance.py b/cosmos/operators/azure_container_instance.py
@@ -7,6 +7,7 @@
 from cosmos.config import ProfileConfig
 from cosmos.operators.base import (
     AbstractDbtBaseOperator,
+    DbtBuildMixin,
     DbtLSMixin,
     DbtRunMixin,
     DbtRunOperationMixin,
@@ -77,6 +78,17 @@ def build_command(self, context: Context, cmd_flags: list[str] | None = None) ->
         self.command: list[str] = dbt_cmd
 
 
+class DbtBuildAzureContainerInstanceOperator(DbtBuildMixin, DbtAzureContainerInstanceBaseOperator):  # type: ignore
+    """
+    Executes a dbt core build command.
+    """
+
+    template_fields: Sequence[str] = DbtAzureContainerInstanceBaseOperator.template_fields + DbtBuildMixin.template_fields  # type: ignore[operator]
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+
+
 class DbtLSAzureContainerInstanceOperator(DbtLSMixin, DbtAzureContainerInstanceBaseOperator):  # type: ignore
     """
     Executes a dbt core ls command.

diff --git a/tests/operators/test_azure_container_instance.py b/tests/operators/test_azure_container_instance.py
@@ -6,6 +6,7 @@
 
 from cosmos.operators.azure_container_instance import (
     DbtAzureContainerInstanceBaseOperator,
+    DbtBuildAzureContainerInstanceOperator,
     DbtLSAzureContainerInstanceOperator,
     DbtRunAzureContainerInstanceOperator,
     DbtSeedAzureContainerInstanceOperator,
@@ -121,6 +122,7 @@ def test_dbt_azure_container_instance_operator_check_environment_variables(
 }
 
 result_map = {
+    "build": DbtBuildAzureContainerInstanceOperator(**base_kwargs),
     "ls": DbtLSAzureContainerInstanceOperator(**base_kwargs),
     "run": DbtRunAzureContainerInstanceOperator(**base_kwargs),
     "test": DbtTestAzureContainerInstanceOperator(**base_kwargs),