Skip to content

Commit

Permalink
Merge branch 'main' into main
Browse files Browse the repository at this point in the history
  • Loading branch information
tatiana authored Sep 23, 2024
2 parents fa56fdd + 0fd0bdd commit 28c2860
Show file tree
Hide file tree
Showing 3 changed files with 87 additions and 21 deletions.
94 changes: 73 additions & 21 deletions cosmos/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@

try:
from cosmos.operators.docker import (
DbtBuildDockerOperator,
DbtLSDockerOperator,
DbtRunDockerOperator,
DbtRunOperationDockerOperator,
Expand All @@ -62,6 +63,7 @@

try:
from cosmos.operators.kubernetes import (
DbtBuildKubernetesOperator,
DbtLSKubernetesOperator,
DbtRunKubernetesOperator,
DbtRunOperationKubernetesOperator,
Expand All @@ -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",
Expand Down Expand Up @@ -98,6 +104,7 @@

try:
from cosmos.operators.azure_container_instance import (
DbtBuildAzureContainerInstanceOperator,
DbtLSAzureContainerInstanceOperator,
DbtRunAzureContainerInstanceOperator,
DbtRunOperationAzureContainerInstanceOperator,
Expand All @@ -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"
)
Expand All @@ -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",
]

"""
Expand Down
12 changes: 12 additions & 0 deletions cosmos/operators/azure_container_instance.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
from cosmos.config import ProfileConfig
from cosmos.operators.base import (
AbstractDbtBaseOperator,
DbtBuildMixin,
DbtLSMixin,
DbtRunMixin,
DbtRunOperationMixin,
Expand Down Expand Up @@ -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.
Expand Down
2 changes: 2 additions & 0 deletions tests/operators/test_azure_container_instance.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@

from cosmos.operators.azure_container_instance import (
DbtAzureContainerInstanceBaseOperator,
DbtBuildAzureContainerInstanceOperator,
DbtLSAzureContainerInstanceOperator,
DbtRunAzureContainerInstanceOperator,
DbtSeedAzureContainerInstanceOperator,
Expand Down Expand Up @@ -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),
Expand Down

0 comments on commit 28c2860

Please sign in to comment.