Merge branch 'main' of https://github.com/ags-de/astronomer-cosmos in…

…to gcpcloudrun-executor
astronomer · Sep 24, 2024 · 1582719 · 1582719
2 parents 85c4c0a + 3414513
commit 1582719
Show file tree

Hide file tree

Showing 13 changed files with 233 additions and 45 deletions.
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
@@ -93,10 +93,11 @@ jobs:
           hatch run tests.py${{ matrix.python-version }}-${{ matrix.airflow-version }}:test-cov
 
       - name: Upload coverage to Github
-        uses: actions/upload-artifact@v2
+        uses: actions/upload-artifact@v4
         with:
           name: coverage-unit-test-${{ matrix.python-version }}-${{ matrix.airflow-version }}
           path: .coverage
+          include-hidden-files: true
 
   Run-Integration-Tests:
     needs: Authorize
@@ -177,10 +178,11 @@ jobs:
           POSTGRES_PORT: 5432
 
       - name: Upload coverage to Github
-        uses: actions/upload-artifact@v2
+        uses: actions/upload-artifact@v4
         with:
           name: coverage-integration-test-${{ matrix.python-version }}-${{ matrix.airflow-version }}
           path: .coverage
+          include-hidden-files: true
 
   Run-Integration-Tests-Expensive:
     needs: Authorize
@@ -248,10 +250,11 @@ jobs:
           POSTGRES_PORT: 5432
 
       - name: Upload coverage to Github
-        uses: actions/upload-artifact@v2
+        uses: actions/upload-artifact@v4
         with:
           name: coverage-integration-expensive-test-${{ matrix.python-version }}-${{ matrix.airflow-version }}
           path: .coverage
+          include-hidden-files: true
 
     env:
       AIRFLOW_HOME: /home/runner/work/astronomer-cosmos/astronomer-cosmos/
@@ -315,10 +318,11 @@ jobs:
           POSTGRES_PORT: 5432
 
       - name: Upload coverage to Github
-        uses: actions/upload-artifact@v2
+        uses: actions/upload-artifact@v4
         with:
           name: coverage-integration-sqlite-test-${{ matrix.python-version }}-${{ matrix.airflow-version }}
           path: .coverage
+          include-hidden-files: true
 
     env:
       AIRFLOW_HOME: /home/runner/work/astronomer-cosmos/astronomer-cosmos/
@@ -391,10 +395,11 @@ jobs:
           POSTGRES_PORT: 5432
 
       - name: Upload coverage to Github
-        uses: actions/upload-artifact@v2
+        uses: actions/upload-artifact@v4
         with:
           name: coverage-integration-dbt-1-5-4-test-${{ matrix.python-version }}-${{ matrix.airflow-version }}
           path: .coverage
+          include-hidden-files: true
 
     env:
       AIRFLOW_HOME: /home/runner/work/astronomer-cosmos/astronomer-cosmos/
@@ -534,10 +539,11 @@ jobs:
           POSTGRES_PORT: 5432
 
       - name: Upload coverage to Github
-        uses: actions/upload-artifact@v2
+        uses: actions/upload-artifact@v4
         with:
           name: coverage-integration-kubernetes-test-${{ matrix.python-version }}-${{ matrix.airflow-version }}
           path: .coverage
+          include-hidden-files: true
 
   Code-Coverage:
     if: github.event.action != 'labeled'
@@ -559,7 +565,7 @@ jobs:
         run: |
           pip3 install coverage
       - name: Download all coverage artifacts
-        uses: actions/download-artifact@v2
+        uses: actions/download-artifact@v4
         with:
           path: ./coverage
       - name: Combine coverage

diff --git a/.gitignore b/.gitignore
@@ -160,3 +160,6 @@ webserver_config.py
 
 # VI
 *.sw[a-z]
+
+# Ignore possibly created symlink to `dev/dags` for running `airflow dags test` command.
+dags
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -54,7 +54,7 @@ repos:
           - --py37-plus
           - --keep-runtime-typing
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.6.2
+    rev: v0.6.7
     hooks:
       - id: ruff
         args:

diff --git a/cosmos/__init__.py b/cosmos/__init__.py
@@ -76,6 +76,7 @@
 
 try:
     from cosmos.operators.docker import (
+        DbtBuildDockerOperator,
         DbtLSDockerOperator,
         DbtRunDockerOperator,
         DbtRunOperationDockerOperator,
@@ -96,6 +97,7 @@
 
 try:
     from cosmos.operators.kubernetes import (
+        DbtBuildKubernetesOperator,
         DbtLSKubernetesOperator,
         DbtRunKubernetesOperator,
         DbtRunOperationKubernetesOperator,
@@ -105,6 +107,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",
@@ -132,6 +138,7 @@
 
 try:
     from cosmos.operators.azure_container_instance import (
+        DbtBuildAzureContainerInstanceOperator,
         DbtLSAzureContainerInstanceOperator,
         DbtRunAzureContainerInstanceOperator,
         DbtRunOperationAzureContainerInstanceOperator,
@@ -140,6 +147,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"
     )
@@ -161,55 +171,98 @@
         "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",
-    "DbtBuildGcpCloudRunJobOperator",
-    "DbtLSGcpCloudRunJobOperator",
-    "DbtRunGcpCloudRunJobOperator",
-    "DbtRunOperationGcpCloudRunJobOperator",
-    "DbtSeedGcpCloudRunJobOperator",
-    "DbtSnapshotGcpCloudRunJobOperator",
-    "DbtTestGcpCloudRunJobOperator",
+    "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",
+    # GCP Cloud Run Job Execution Mode
+    "DbtBuildGcpCloudRunJobOperator",
+    "DbtLSGcpCloudRunJobOperator",
+    "DbtRunGcpCloudRunJobOperator",
+    "DbtRunOperationGcpCloudRunJobOperator",
+    "DbtSeedGcpCloudRunJobOperator",
+    "DbtSnapshotGcpCloudRunJobOperator",
+    "DbtTestGcpCloudRunJobOperator",
 ]
 
 """

diff --git a/cosmos/core/airflow.py b/cosmos/core/airflow.py
@@ -25,16 +25,15 @@ def get_airflow_task(task: Task, dag: DAG, task_group: "TaskGroup | None" = None
     module = importlib.import_module(module_name)
     Operator = getattr(module, class_name)
 
+    task_kwargs = {}
     if task.owner != "":
-        task_owner = task.owner
-    else:
-        task_owner = dag.owner
+        task_kwargs["owner"] = task.owner
 
     airflow_task = Operator(
         task_id=task.id,
         dag=dag,
         task_group=task_group,
-        owner=task_owner,
+        **task_kwargs,
         **({} if class_name == "EmptyOperator" else {"extra_context": task.extra_context}),
         **task.arguments,
     )

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/cosmos/operators/base.py b/cosmos/operators/base.py
@@ -139,6 +139,7 @@ def __init__(
         self.dbt_cmd_global_flags = dbt_cmd_global_flags or []
         self.cache_dir = cache_dir
         self.extra_context = extra_context or {}
+        kwargs.pop("full_refresh", None)  # usage of this param should be implemented in child classes
         super().__init__(**kwargs)
 
     def get_env(self, context: Context) -> dict[str, str | bytes | os.PathLike[Any]]:

diff --git a/cosmos/operators/local.py b/cosmos/operators/local.py
@@ -145,7 +145,6 @@ def __init__(
         self._dbt_runner: dbtRunner | None = None
         if self.invocation_mode:
             self._set_invocation_methods()
-        kwargs.pop("full_refresh", None)  # usage of this param should be implemented in child classes
         super().__init__(**kwargs)
 
         # For local execution mode, we're consistent with the LoadMode.DBT_LS command in forwarding the environment

diff --git a/cosmos/profiles/base.py b/cosmos/profiles/base.py
@@ -33,7 +33,7 @@ class DbtProfileConfigVars:
     partial_parse: Optional[bool] = None
     use_experimental_parser: Optional[bool] = None
     static_parser: Optional[bool] = None
-    printer_width: Optional[bool] = None
+    printer_width: Optional[int] = None
     write_json: Optional[bool] = None
     warn_error: Optional[bool] = None
     warn_error_options: Optional[Dict[Literal["include", "exclude"], Any]] = None

diff --git a/pyproject.toml b/pyproject.toml
@@ -34,7 +34,7 @@ dependencies = [
     "importlib-metadata; python_version < '3.8'",
     "Jinja2>=3.0.0",
     "msgpack",
-    "packaging",
+    "packaging>=22.0",
     "pydantic>=1.10.0",
     "typing-extensions; python_version < '3.8'",
     "virtualenv",
@@ -86,7 +86,7 @@ docs = [
     "apache-airflow-providers-cncf-kubernetes>=5.1.1",
 ]
 tests = [
-    "packaging",
+    "packaging>=22.0",
     "pytest>=6.0",
     "pytest-split",
     "pytest-dotenv",