diff --git a/cosmos/__init__.py b/cosmos/__init__.py index 3422cba62..363a9a122 100644 --- 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 index 1ad73c4a6..d3c8ebfc3 100644 --- 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 index 836ca4596..c57466619 100644 --- 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),