More cleanup

astronomer · Nov 15, 2024 · 8af04da · 8af04da
1 parent ba58a96
commit 8af04da
Show file tree

Hide file tree

Showing 6 changed files with 22 additions and 0 deletions.
diff --git a/cosmos/__init__.py b/cosmos/__init__.py
@@ -30,6 +30,7 @@
 from cosmos.operators.lazy_load import MissingPackage
 from cosmos.operators.local import (
     DbtBuildLocalOperator,
+    DbtCloneLocalOperator,
     DbtDepsLocalOperator,
     DbtLSLocalOperator,
     DbtRunLocalOperator,
@@ -44,6 +45,7 @@
 try:
     from cosmos.operators.docker import (
         DbtBuildDockerOperator,
+        DbtCloneDockerOperator,
         DbtLSDockerOperator,
         DbtRunDockerOperator,
         DbtRunOperationDockerOperator,
@@ -65,6 +67,7 @@
 try:
     from cosmos.operators.kubernetes import (
         DbtBuildKubernetesOperator,
+        DbtCloneKubernetesOperator,
         DbtLSKubernetesOperator,
         DbtRunKubernetesOperator,
         DbtRunOperationKubernetesOperator,
@@ -106,6 +109,7 @@
 try:
     from cosmos.operators.azure_container_instance import (
         DbtBuildAzureContainerInstanceOperator,
+        DbtCloneAzureContainerInstanceOperator,
         DbtLSAzureContainerInstanceOperator,
         DbtRunAzureContainerInstanceOperator,
         DbtRunOperationAzureContainerInstanceOperator,
@@ -142,6 +146,7 @@
 try:
     from cosmos.operators.aws_eks import (
         DbtBuildAwsEksOperator,
+        DbtCloneAwsEksOperator,
         DbtLSAwsEksOperator,
         DbtRunAwsEksOperator,
         DbtRunOperationAwsEksOperator,
@@ -170,6 +175,7 @@
 try:
     from cosmos.operators.gcp_cloud_run_job import (
         DbtBuildGcpCloudRunJobOperator,
+        DbtCloneGcpCloudRunJobOperator,
         DbtLSGcpCloudRunJobOperator,
         DbtRunGcpCloudRunJobOperator,
         DbtRunOperationGcpCloudRunJobOperator,
@@ -217,6 +223,7 @@
     "DbtResourceType",
     # Local Execution Mode
     "DbtBuildLocalOperator",
+    "DbtCloneLocalOperator",
     "DbtDepsLocalOperator",  # deprecated, to be delete in Cosmos 2.x
     "DbtLSLocalOperator",
     "DbtRunLocalOperator",
@@ -226,6 +233,7 @@
     "DbtTestLocalOperator",
     # Docker Execution Mode
     "DbtBuildDockerOperator",
+    "DbtCloneDockerOperator",
     "DbtLSDockerOperator",
     "DbtRunDockerOperator",
     "DbtRunOperationDockerOperator",
@@ -234,6 +242,7 @@
     "DbtTestDockerOperator",
     # Kubernetes Execution Mode
     "DbtBuildKubernetesOperator",
+    "DbtCloneKubernetesOperator",
     "DbtLSKubernetesOperator",
     "DbtRunKubernetesOperator",
     "DbtRunOperationKubernetesOperator",
@@ -242,6 +251,7 @@
     "DbtTestKubernetesOperator",
     # Azure Container Instance Execution Mode
     "DbtBuildAzureContainerInstanceOperator",
+    "DbtCloneAzureContainerInstanceOperator",
     "DbtLSAzureContainerInstanceOperator",
     "DbtRunAzureContainerInstanceOperator",
     "DbtRunOperationAzureContainerInstanceOperator",
@@ -250,6 +260,7 @@
     "DbtTestAzureContainerInstanceOperator",
     # AWS EKS Execution Mode
     "DbtBuildAwsEksOperator",
+    "DbtCloneAwsEksOperator",
     "DbtLSAwsEksOperator",
     "DbtRunAwsEksOperator",
     "DbtRunOperationAwsEksOperator",
@@ -258,6 +269,7 @@
     "DbtTestAwsEksOperator",
     # GCP Cloud Run Job Execution Mode
     "DbtBuildGcpCloudRunJobOperator",
+    "DbtCloneGcpCloudRunJobOperator",
     "DbtLSGcpCloudRunJobOperator",
     "DbtRunGcpCloudRunJobOperator",
     "DbtRunOperationGcpCloudRunJobOperator",

diff --git a/tests/operators/test_aws_eks.py b/tests/operators/test_aws_eks.py
@@ -5,6 +5,7 @@
 
 from cosmos.operators.aws_eks import (
     DbtBuildAwsEksOperator,
+    DbtCloneAwsEksOperator,
     DbtLSAwsEksOperator,
     DbtRunAwsEksOperator,
     DbtSeedAwsEksOperator,
@@ -44,6 +45,7 @@ def test_dbt_kubernetes_build_command():
         "test": DbtTestAwsEksOperator(**base_kwargs),
         "build": DbtBuildAwsEksOperator(**base_kwargs),
         "seed": DbtSeedAwsEksOperator(**base_kwargs),
+        "clone": DbtCloneAwsEksOperator(**base_kwargs),
     }
 
     for command_name, command_operator in result_map.items():

diff --git a/tests/operators/test_azure_container_instance.py b/tests/operators/test_azure_container_instance.py
@@ -7,6 +7,7 @@
 from cosmos.operators.azure_container_instance import (
     DbtAzureContainerInstanceBaseOperator,
     DbtBuildAzureContainerInstanceOperator,
+    DbtCloneAzureContainerInstanceOperator,
     DbtLSAzureContainerInstanceOperator,
     DbtRunAzureContainerInstanceOperator,
     DbtSeedAzureContainerInstanceOperator,
@@ -127,6 +128,7 @@ def test_dbt_azure_container_instance_operator_check_environment_variables(
     "run": DbtRunAzureContainerInstanceOperator(**base_kwargs),
     "test": DbtTestAzureContainerInstanceOperator(**base_kwargs),
     "seed": DbtSeedAzureContainerInstanceOperator(**base_kwargs),
+    "clone": DbtCloneAzureContainerInstanceOperator(**base_kwargs),
 }
 
 

diff --git a/tests/operators/test_docker.py b/tests/operators/test_docker.py
@@ -7,6 +7,7 @@
 
 from cosmos.operators.docker import (
     DbtBuildDockerOperator,
+    DbtCloneDockerOperator,
     DbtLSDockerOperator,
     DbtRunDockerOperator,
     DbtSeedDockerOperator,
@@ -113,6 +114,7 @@ def test_dbt_docker_operator_get_env(p_context_to_airflow_vars: MagicMock, base_
     "test": DbtTestDockerOperator(**base_kwargs),
     "build": DbtBuildDockerOperator(**base_kwargs),
     "seed": DbtSeedDockerOperator(**base_kwargs),
+    "clone": DbtCloneDockerOperator(**base_kwargs),
 }
 
 

diff --git a/tests/operators/test_gcp_cloud_run_job.py b/tests/operators/test_gcp_cloud_run_job.py
@@ -10,6 +10,7 @@
 try:
     from cosmos.operators.gcp_cloud_run_job import (
         DbtBuildGcpCloudRunJobOperator,
+        DbtCloneGcpCloudRunJobOperator,
         DbtGcpCloudRunJobBaseOperator,
         DbtLSGcpCloudRunJobOperator,
         DbtRunGcpCloudRunJobOperator,
@@ -173,6 +174,7 @@ def test_dbt_gcp_cloud_run_job_build_command():
         "build": DbtBuildGcpCloudRunJobOperator(**BASE_KWARGS),
         "snapshot": DbtSnapshotGcpCloudRunJobOperator(**BASE_KWARGS),
         "source": DbtSourceGcpCloudRunJobOperator(**BASE_KWARGS),
+        "clone": DbtCloneGcpCloudRunJobOperator(**BASE_KWARGS),
         "run-operation": DbtRunOperationGcpCloudRunJobOperator(macro_name="some-macro", **BASE_KWARGS),
     }
 

diff --git a/tests/operators/test_kubernetes.py b/tests/operators/test_kubernetes.py
@@ -10,6 +10,7 @@
 
 from cosmos.operators.kubernetes import (
     DbtBuildKubernetesOperator,
+    DbtCloneKubernetesOperator,
     DbtLSKubernetesOperator,
     DbtRunKubernetesOperator,
     DbtSeedKubernetesOperator,
@@ -128,6 +129,7 @@ def test_dbt_kubernetes_operator_get_env(p_context_to_airflow_vars: MagicMock, b
     "test": DbtTestKubernetesOperator(**base_kwargs),
     "build": DbtBuildKubernetesOperator(**base_kwargs),
     "seed": DbtSeedKubernetesOperator(**base_kwargs),
+    "clone": DbtCloneKubernetesOperator(**base_kwargs),
 }