Merge branch 'main' into feat/no_retries_on_test_failure

Author: tatiana

.github/workflows/docs.yml

@@ -7,7 +7,7 @@ on:
 
 jobs:
   pages:
-    runs-on: ubuntu-20.04
+    runs-on: ubuntu-latest
     environment:
       name: github-pages
       url: ${{ steps.deployment.outputs.page_url }}

.pre-commit-config.yaml

@@ -57,7 +57,7 @@ repos:
           - --py37-plus
           - --keep-runtime-typing
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.9.6
+    rev: v0.9.7
     hooks:
       - id: ruff
         args:

CHANGELOG.rst

@@ -1,7 +1,15 @@
 Changelog
 =========
 
-1.9.0a5 (2025-02-03)
+1.9.1a1 (2025-02-20)
+--------------------
+
+Bug Fixes
+
+* Fix import error in dbt bigquery adapter mock for ``dbt-bigquery<1.8`` for ``ExecutionMode.AIRFLOW_ASYNC`` by @pankajkoti in #1548
+
+
+1.9.0 (2025-02-19)
 --------------------
 
 Breaking changes
@@ -19,23 +27,43 @@ Features
 * Add structure to support multiple db for async operator execution by @pankajastro in #1483
 * Support overriding the ``profile_config`` per dbt node or folder using config by @tatiana in #1492. More information `here <https://astronomer.github.io/astronomer-cosmos/profiles/#profile-customise-per-node>`_.
 * Create and run accurate SQL statements when using ``ExecutionMode.AIRFLOW_ASYNC`` by @pankajkoti, @tatiana and @pankajastro in #1474
+* Add AWS ECS task run execution mode by @CarlosGitto and @aoelvp94 in #1507
+* Add support for running ``DbtSourceOperator`` individually by @victormacaubas in #1510
+* Add setup task for async executions by @pankajastro in #1518
+* Add teardown task for async executions by @pankajastro in #1529
+* Add ``ProjectConfig.install_dbt_deps`` & change operator ``install_deps=True`` as default by @tatiana in #1521
+* Extend Virtualenv operator and mock dbt adapters for setup & teardown tasks in ``ExecutionMode.AIRFLOW_ASYNC`` by @pankajkoti, @tatiana and @pankajastro in #1544
 
 Bug Fixes
 
 * Fix select complex intersection of three tag-based graph selectors by @tatiana in #1466
+* Fix custom selector behaviour when the model name contains periods by @yakovlevvs and @60098727 in #1499
+* Filter dbt and non-dbt kwargs correctly for async operator by @pankajastro in #1526
 
 Enhancement
 
 * Fix OpenLineage deprecation warning by @CorsettiS in #1449
 * Move ``DbtRunner`` related functions into ``dbt/runner.py`` module by @tatiana in #1480
 * Add ``on_warning_callback`` to ``DbtSourceKubernetesOperator`` and refactor previous operators by @LuigiCerone in #1501
+* Gracefully error when users set incompatible ``RenderConfig.dbt_deps`` and ``operator_args`` ``install_deps`` by @tatiana in #1505
+* Store compiled SQL as template field for ``ExecutionMode.AIRFLOW_ASYNC`` by @pankajkoti in #1534
+
+Docs
 
+* Improve ``RenderConfig`` arguments documentation by @tatiana in #1514
+* Improve callback documentation by @tatiana in #1516
+* Improve partial parsing docs by @tatiana in #1520
+* Fix typo in selecting & excluding docs by @pankajastro in #1523
+* Document ``async_py_requirements`` added in ``ExecutionConfig`` for ``ExecutionMode.AIRFLOW_ASYNC`` by @pankajkoti in #1545
 
 Others
 
 * Ignore dbt package tests when running Cosmos tests by @tatiana in #1502
+* Refactor to consolidate async dbt adapter code by @pankajkoti in #1509
+* Log elapsed time for sql file(s) upload/download by @pankajastro in #1536
+* Remove the fallback operator for async task by @pankajastro in #1538
 * GitHub Actions Dependabot: #1487
-* Pre-commit updates: #1473, #1493
+* Pre-commit updates: #1473, #1493, #1503, #1531
 
 
 1.8.2 (2025-01-15)

cosmos/__init__.py

@@ -6,7 +6,7 @@
 Contains dags, task groups, and operators.
 """
 
-__version__ = "1.9.0a6"
+__version__ = "1.9.1a1"
 
 
 from cosmos.airflow.dag import DbtDag

cosmos/airflow/graph.py

@@ -25,6 +25,7 @@
 from cosmos.core.airflow import get_airflow_task as create_airflow_task
 from cosmos.core.graph.entities import Task as TaskMetadata
 from cosmos.dbt.graph import DbtNode
+from cosmos.exceptions import CosmosValueError
 from cosmos.log import get_logger
 from cosmos.settings import enable_setup_async_task, enable_teardown_async_task
 
@@ -413,14 +414,19 @@ def _add_dbt_setup_async_task(
     tasks_map: dict[str, Any],
     task_group: TaskGroup | None,
     render_config: RenderConfig | None = None,
+    async_py_requirements: list[str] | None = None,
 ) -> None:
     if execution_mode != ExecutionMode.AIRFLOW_ASYNC:
         return
 
+    if not async_py_requirements:
+        raise CosmosValueError("ExecutionConfig.AIRFLOW_ASYNC needs async_py_requirements to be set")
+
     if render_config is not None:
         task_args["select"] = render_config.select
         task_args["selector"] = render_config.selector
         task_args["exclude"] = render_config.exclude
+        task_args["py_requirements"] = async_py_requirements
 
     setup_task_metadata = TaskMetadata(
         id=DBT_SETUP_ASYNC_TASK_ID,
@@ -495,14 +501,19 @@ def _add_teardown_task(
     tasks_map: dict[str, Any],
     task_group: TaskGroup | None,
     render_config: RenderConfig | None = None,
+    async_py_requirements: list[str] | None = None,
 ) -> None:
     if execution_mode != ExecutionMode.AIRFLOW_ASYNC:
         return
 
+    if not async_py_requirements:
+        raise CosmosValueError("ExecutionConfig.AIRFLOW_ASYNC needs async_py_requirements to be set")
+
     if render_config is not None:
         task_args["select"] = render_config.select
         task_args["selector"] = render_config.selector
         task_args["exclude"] = render_config.exclude
+        task_args["py_requirements"] = async_py_requirements
 
     teardown_task_metadata = TaskMetadata(
         id=DBT_TEARDOWN_ASYNC_TASK_ID,
@@ -529,6 +540,7 @@ def build_airflow_graph(
     render_config: RenderConfig,
     task_group: TaskGroup | None = None,
     on_warning_callback: Callable[..., Any] | None = None,  # argument specific to the DBT test command
+    async_py_requirements: list[str] | None = None,
 ) -> dict[str, Union[TaskGroup, BaseOperator]]:
     """
     Instantiate dbt `nodes` as Airflow tasks within the given `task_group` (optional) or `dag` (mandatory).
@@ -626,9 +638,25 @@ def build_airflow_graph(
 
     create_airflow_task_dependencies(nodes, tasks_map)
     if enable_setup_async_task:
-        _add_dbt_setup_async_task(dag, execution_mode, task_args, tasks_map, task_group, render_config=render_config)
+        _add_dbt_setup_async_task(
+            dag,
+            execution_mode,
+            task_args,
+            tasks_map,
+            task_group,
+            render_config=render_config,
+            async_py_requirements=async_py_requirements,
+        )
     if enable_teardown_async_task:
-        _add_teardown_task(dag, execution_mode, task_args, tasks_map, task_group, render_config=render_config)
+        _add_teardown_task(
+            dag,
+            execution_mode,
+            task_args,
+            tasks_map,
+            task_group,
+            render_config=render_config,
+            async_py_requirements=async_py_requirements,
+        )
     return tasks_map
 
 

cosmos/config.py

@@ -152,6 +152,7 @@ class ProjectConfig:
     :param snapshots_relative_path: The relative path to the dbt snapshots directory within the project. Defaults to
     snapshots
     :param manifest_path: The absolute path to the dbt manifest file. Defaults to None
+    :param manifest_conn_id: Name of the Airflow connection used to access the manifest file if it is not stored locally. Defaults to None
     :param project_name: Allows the user to define the project name.
     Required if dbt_project_path is not defined. Defaults to the folder name of dbt_project_path.
     :param env_vars: Dictionary of environment variables that are used for both rendering and execution. Rendering with
@@ -175,6 +176,7 @@ class ProjectConfig:
     def __init__(
         self,
         dbt_project_path: str | Path | None = None,
+        install_dbt_deps: bool = True,
         models_relative_path: str | Path = "models",
         seeds_relative_path: str | Path = "seeds",
         snapshots_relative_path: str | Path = "snapshots",
@@ -228,6 +230,7 @@ def __init__(
         self.env_vars = env_vars
         self.dbt_vars = dbt_vars
         self.partial_parse = partial_parse
+        self.install_dbt_deps = install_dbt_deps
 
     def validate_project(self) -> None:
         """
@@ -398,6 +401,8 @@ class ExecutionConfig:
     :param dbt_project_path: Configures the DBT project location accessible at runtime for dag execution. This is the project path in a docker container for ExecutionMode.DOCKER or ExecutionMode.KUBERNETES. Mutually Exclusive with ProjectConfig.dbt_project_path
     :param virtualenv_dir: Directory path to locate the (cached) virtual env that
     should be used for execution when execution mode is set to `ExecutionMode.VIRTUALENV`
+    :param async_py_requirements:  A list of Python packages to install when `ExecutionMode.AIRFLOW_ASYNC`(Experimental) is used. This parameter is required only if both `enable_setup_async_task` and `enable_teardown_async_task` are set to `True`.
+    Example: `["dbt-postgres==1.5.0"]`
     """
 
     execution_mode: ExecutionMode = ExecutionMode.LOCAL
@@ -409,6 +414,7 @@ class ExecutionConfig:
     virtualenv_dir: str | Path | None = None
 
     project_path: Path | None = field(init=False)
+    async_py_requirements: list[str] | None = None
 
     def __post_init__(self, dbt_project_path: str | Path | None) -> None:
         if self.invocation_mode and self.execution_mode not in (ExecutionMode.LOCAL, ExecutionMode.VIRTUALENV):

cosmos/converter.py

@@ -223,7 +223,7 @@ def override_configuration(
     if execution_config.invocation_mode:
         operator_args["invocation_mode"] = execution_config.invocation_mode
 
-    if execution_config in (ExecutionMode.LOCAL, ExecutionMode.VIRTUALENV):
+    if execution_config.execution_mode in (ExecutionMode.LOCAL, ExecutionMode.VIRTUALENV):
         if "install_deps" not in operator_args:
             operator_args["install_deps"] = project_config.install_dbt_deps
 
@@ -335,6 +335,7 @@ def __init__(
             dbt_project_name=render_config.project_name,
             on_warning_callback=on_warning_callback,
             render_config=render_config,
+            async_py_requirements=execution_config.async_py_requirements,
         )
 
         current_time = time.perf_counter()

cosmos/operators/_asynchronous/__init__.py

@@ -1,29 +1,74 @@
 from __future__ import annotations
 
+import inspect
+import textwrap
+from pathlib import Path
 from typing import Any
 
 from airflow.utils.context import Context
 
-from cosmos.operators.local import DbtRunLocalOperator as DbtRunOperator
+from cosmos._utils.importer import load_method_from_module
+from cosmos.hooks.subprocess import FullOutputSubprocessResult
+from cosmos.operators.virtualenv import DbtRunVirtualenvOperator
 
 
-class SetupAsyncOperator(DbtRunOperator):
+class SetupAsyncOperator(DbtRunVirtualenvOperator):
     def __init__(self, *args: Any, **kwargs: Any):
         kwargs["emit_datasets"] = False
         super().__init__(*args, **kwargs)
 
+    def run_subprocess(self, command: list[str], env: dict[str, str], cwd: str) -> FullOutputSubprocessResult:
+        profile_type = self.profile_config.get_profile_type()
+        if not self._py_bin:
+            raise AttributeError("_py_bin attribute not set for VirtualEnv operator")
+        dbt_executable_path = str(Path(self._py_bin).parent / "dbt")
+        asynchronous_operator_module = f"cosmos.operators._asynchronous.{profile_type}"
+        mock_function_name = f"_mock_{profile_type}_adapter"
+        mock_function = load_method_from_module(asynchronous_operator_module, mock_function_name)
+        mock_function_full_source = inspect.getsource(mock_function)
+        mock_function_body = textwrap.dedent("\n".join(mock_function_full_source.split("\n")[1:]))
+
+        with open(dbt_executable_path) as f:
+            dbt_entrypoint_script = f.readlines()
+        if dbt_entrypoint_script[0].startswith("#!"):
+            dbt_entrypoint_script.insert(1, mock_function_body)
+        with open(dbt_executable_path, "w") as f:
+            f.writelines(dbt_entrypoint_script)
+
+        return super().run_subprocess(command, env, cwd)
+
     def execute(self, context: Context, **kwargs: Any) -> None:
         async_context = {"profile_type": self.profile_config.get_profile_type()}
         self.build_and_run_cmd(
             context=context, cmd_flags=self.dbt_cmd_flags, run_as_async=True, async_context=async_context
         )
 
 
-class TeardownAsyncOperator(DbtRunOperator):
+class TeardownAsyncOperator(DbtRunVirtualenvOperator):
     def __init__(self, *args: Any, **kwargs: Any):
         kwargs["emit_datasets"] = False
         super().__init__(*args, **kwargs)
 
+    def run_subprocess(self, command: list[str], env: dict[str, str], cwd: str) -> FullOutputSubprocessResult:
+        profile_type = self.profile_config.get_profile_type()
+        if not self._py_bin:
+            raise AttributeError("_py_bin attribute not set for VirtualEnv operator")
+        dbt_executable_path = str(Path(self._py_bin).parent / "dbt")
+        asynchronous_operator_module = f"cosmos.operators._asynchronous.{profile_type}"
+        mock_function_name = f"_mock_{profile_type}_adapter"
+        mock_function = load_method_from_module(asynchronous_operator_module, mock_function_name)
+        mock_function_full_source = inspect.getsource(mock_function)
+        mock_function_body = textwrap.dedent("\n".join(mock_function_full_source.split("\n")[1:]))
+
+        with open(dbt_executable_path) as f:
+            dbt_entrypoint_script = f.readlines()
+        if dbt_entrypoint_script[0].startswith("#!"):
+            dbt_entrypoint_script.insert(1, mock_function_body)
+        with open(dbt_executable_path, "w") as f:
+            f.writelines(dbt_entrypoint_script)
+
+        return super().run_subprocess(command, env, cwd)
+
     def execute(self, context: Context, **kwargs: Any) -> Any:
         async_context = {"profile_type": self.profile_config.get_profile_type(), "teardown_task": True}
         self.build_and_run_cmd(

cosmos/operators/_asynchronous/bigquery.py

@@ -28,7 +28,11 @@ def _mock_bigquery_adapter() -> None:
 
     import agate
     from dbt.adapters.bigquery.connections import BigQueryAdapterResponse, BigQueryConnectionManager
-    from dbt_common.clients.agate_helper import empty_table
+
+    try:
+        from dbt_common.clients.agate_helper import empty_table
+    except (ModuleNotFoundError, ImportError):  # pragma: no cover
+        from dbt.clients.agate_helper import empty_table
 
     def execute(  # type: ignore[no-untyped-def]
         self, sql, auto_begin=False, fetch=None, limit: Optional[int] = None

cosmos/operators/local.py

@@ -531,7 +531,7 @@ def run_command(
                 if self.install_deps:
                     self._install_dependencies(tmp_dir_path, flags, env)
 
-                if run_as_async:
+                if run_as_async and not enable_setup_async_task:
                     self._mock_dbt_adapter(async_context)
 
                 full_cmd = cmd + flags