Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

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

Merged
merged 1 commit into from
Sep 23, 2024
+87 −21
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
94 changes: 73 additions & 21 deletions cosmos/__init__.py
View file
Open in desktop
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(

Check warning on line 116 in cosmos/__init__.py

View check run for this annotation

Codecov / codecov/patch

cosmos/__init__.py#L116 

Added line #L116 was not covered by tests
"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(

Check warning on line 152 in cosmos/__init__.py

View check run for this annotation

Codecov / codecov/patch

cosmos/__init__.py#L151-L152 

Added lines #L151 - L152 were not covered by tests
"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(

Check warning on line 157 in cosmos/__init__.py

View check run for this annotation

Codecov / codecov/patch

cosmos/__init__.py#L155-L157 

Added lines #L155 - L157 were not covered by tests
"cosmos.operators.azure_container_instance.DbtRunOperationAwsEksOperator",
"aws_eks",
)
DbtSeedAwsEksOperator = MissingPackage("cosmos.operators.azure_container_instance.DbtSeedAwsEksOperator", "aws_eks")
DbtSnapshotAwsEksOperator = MissingPackage(

Check warning on line 162 in cosmos/__init__.py

View check run for this annotation

Codecov / codecov/patch

cosmos/__init__.py#L161-L162 

Added lines #L161 - L162 were not covered by tests
"cosmos.operators.azure_container_instance.DbtSnapshotAwsEksOperator",
"aws_eks",
)
DbtTestAwsEksOperator = MissingPackage("cosmos.operators.azure_container_instance.DbtTestAwsEksOperator", "aws_eks")

Check warning on line 166 in cosmos/__init__.py

View check run for this annotation

Codecov / codecov/patch

cosmos/__init__.py#L166 

Added line #L166 was not covered by tests


__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
View file
Open in desktop
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
View file
Open in desktop
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
Loading