From 387543c70ce24fc2c16cd0d6f819a6357bcdb7af Mon Sep 17 00:00:00 2001 From: corsettigyg Date: Thu, 24 Apr 2025 20:27:09 +0200 Subject: [PATCH 1/4] allow multiple callback functions --- cosmos/operators/local.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/cosmos/operators/local.py b/cosmos/operators/local.py index 0f36a79233..4f4f9a0784 100644 --- a/cosmos/operators/local.py +++ b/cosmos/operators/local.py @@ -160,7 +160,7 @@ def __init__( invocation_mode: InvocationMode | None = None, install_deps: bool = True, copy_dbt_packages: bool = settings.default_copy_dbt_packages, - callback: Callable[[str], None] | None = None, + callback: Callable[[str], None] | list[Callable[[str], None]] | None = None, callback_args: dict[str, Any] | None = None, should_store_compiled_sql: bool = True, should_upload_compiled_sql: bool = False, @@ -508,7 +508,12 @@ def _handle_post_execution(self, tmp_project_dir: str, context: Context) -> None self._upload_sql_files(tmp_project_dir, "compiled") if self.callback: self.callback_args.update({"context": context}) - self.callback(tmp_project_dir, **self.callback_args) + # handle the scenario where callback is a list of functions instead of a single function + if isinstance(self.callback, list): + for callback_fn in self.callback: + callback_fn(tmp_project_dir, **self.callback_args) + else: + self.callback(tmp_project_dir, **self.callback_args) def _handle_async_execution(self, tmp_project_dir: str, context: Context, async_context: dict[str, Any]) -> None: if async_context.get("teardown_task") and settings.enable_teardown_async_task: From 0997c7d4c468c25ebfb40fc9f206035849e2e181 Mon Sep 17 00:00:00 2001 From: corsettigyg Date: Thu, 24 Apr 2025 21:36:47 +0200 Subject: [PATCH 2/4] remove comment (unnecessary) --- cosmos/operators/local.py | 1 - 1 file changed, 1 deletion(-) diff --git a/cosmos/operators/local.py b/cosmos/operators/local.py index 4f4f9a0784..ba987f2e4a 100644 --- a/cosmos/operators/local.py +++ b/cosmos/operators/local.py @@ -508,7 +508,6 @@ def _handle_post_execution(self, tmp_project_dir: str, context: Context) -> None self._upload_sql_files(tmp_project_dir, "compiled") if self.callback: self.callback_args.update({"context": context}) - # handle the scenario where callback is a list of functions instead of a single function if isinstance(self.callback, list): for callback_fn in self.callback: callback_fn(tmp_project_dir, **self.callback_args) From f119016a345b132b2c40fa94936db417192849e3 Mon Sep 17 00:00:00 2001 From: corsettigyg Date: Fri, 25 Apr 2025 16:03:29 +0200 Subject: [PATCH 3/4] add tests --- tests/operators/test_local.py | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/tests/operators/test_local.py b/tests/operators/test_local.py index fbdc8cb3f3..55be3ed924 100644 --- a/tests/operators/test_local.py +++ b/tests/operators/test_local.py @@ -1562,3 +1562,23 @@ def test_test_clone_project(create_symlinks_mock, copy_dbt_packages_mock, caplog assert f"Cloning project to writable temp directory {tmp_dir_path} from {project_dir}" in caplog.text assert "Copying dbt packages to temporary folder." in caplog.text assert "Completed copying dbt packages to temporary folder." in caplog.text + +@patch("cosmos.operators.local.AbstractDbtLocalBase.store_freshness_json") +@patch("cosmos.operators.local.AbstractDbtLocalBase.store_compiled_sql") +@patch("cosmos.operators.local.AbstractDbtLocalBase._override_rtif") +def test_handle_post_execution_with_multiple_callbacks(mock_override_rtif, mock_store_compiled_sql, mock_store_freshness_json): + + multiple_callbacks = [MagicMock(), MagicMock(), MagicMock()] + operator = ConcreteDbtLocalBaseOperator( + profile_config=profile_config, + task_id="my-task", + project_dir="my/dir", + callback=multiple_callbacks, + callback_args={"arg1": "value1"} + ) + + context = {"dag_run": MagicMock(), "task": MagicMock()} + operator._handle_post_execution("/tmp/project_dir", context) + + for callback_fn in multiple_callbacks: + callback_fn.assert_called_once_with("/tmp/project_dir", arg1="value1", context=context) \ No newline at end of file From 705f29d2edf0c413bc7ceedc020ff0934e40dbb3 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 25 Apr 2025 14:06:34 +0000 Subject: [PATCH 4/4] =?UTF-8?q?=F0=9F=8E=A8=20[pre-commit.ci]=20Auto=20for?= =?UTF-8?q?mat=20from=20pre-commit.com=20hooks?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- tests/operators/test_local.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/tests/operators/test_local.py b/tests/operators/test_local.py index 55be3ed924..2fb0391cf1 100644 --- a/tests/operators/test_local.py +++ b/tests/operators/test_local.py @@ -1563,10 +1563,13 @@ def test_test_clone_project(create_symlinks_mock, copy_dbt_packages_mock, caplog assert "Copying dbt packages to temporary folder." in caplog.text assert "Completed copying dbt packages to temporary folder." in caplog.text + @patch("cosmos.operators.local.AbstractDbtLocalBase.store_freshness_json") @patch("cosmos.operators.local.AbstractDbtLocalBase.store_compiled_sql") @patch("cosmos.operators.local.AbstractDbtLocalBase._override_rtif") -def test_handle_post_execution_with_multiple_callbacks(mock_override_rtif, mock_store_compiled_sql, mock_store_freshness_json): +def test_handle_post_execution_with_multiple_callbacks( + mock_override_rtif, mock_store_compiled_sql, mock_store_freshness_json +): multiple_callbacks = [MagicMock(), MagicMock(), MagicMock()] operator = ConcreteDbtLocalBaseOperator( @@ -1574,11 +1577,11 @@ def test_handle_post_execution_with_multiple_callbacks(mock_override_rtif, mock_ task_id="my-task", project_dir="my/dir", callback=multiple_callbacks, - callback_args={"arg1": "value1"} + callback_args={"arg1": "value1"}, ) context = {"dag_run": MagicMock(), "task": MagicMock()} operator._handle_post_execution("/tmp/project_dir", context) for callback_fn in multiple_callbacks: - callback_fn.assert_called_once_with("/tmp/project_dir", arg1="value1", context=context) \ No newline at end of file + callback_fn.assert_called_once_with("/tmp/project_dir", arg1="value1", context=context)