From 93396cd4c909f81654b09c51a5cb21e9ff161068 Mon Sep 17 00:00:00 2001 From: agreenburg Date: Tue, 14 Nov 2023 12:21:26 -0500 Subject: [PATCH] Store compiled_sql even when task fails (fixes issue #369) (#671) Update `DbtLocalBaseOperator` code to store `compiled_sql` prior to exception handling so that when a task fails, the `compiled_sql` can still be reviewed. In the process found and fixed a related bug where `compiled_sql` was being dropped on some operations due to the way that the `full_refresh` field was being added to the `template_fields`. Closes #369 Fixes bug introduced in https://github.com/astronomer/astronomer-cosmos/pull/623 where compiled_sql was being lost in `DbtSeedLocalOperator` and `DbtRunLocalOperator` Co-authored-by: Andrew Greenburg (cherry picked from commit ee91ece155433adac8b9ac9531aa11d7235acc03) --- cosmos/operators/local.py | 4 +++- tests/operators/test_local.py | 4 ++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/cosmos/operators/local.py b/cosmos/operators/local.py index 323aeacf7..aa41ea57b 100644 --- a/cosmos/operators/local.py +++ b/cosmos/operators/local.py @@ -245,8 +245,8 @@ def run_command( logger.info("Outlets: %s", outlets) self.register_dataset(inlets, outlets) - self.exception_handling(result) self.store_compiled_sql(tmp_project_dir, context) + self.exception_handling(result) if self.callback: self.callback(tmp_project_dir) @@ -395,6 +395,7 @@ class DbtSeedLocalOperator(DbtLocalBaseOperator): """ ui_color = "#F58D7E" + template_fields: Sequence[str] = DbtLocalBaseOperator.template_fields + ("full_refresh",) # type: ignore[operator] def __init__(self, full_refresh: bool = False, **kwargs: Any) -> None: self.full_refresh = full_refresh @@ -433,6 +434,7 @@ class DbtRunLocalOperator(DbtLocalBaseOperator): ui_color = "#7352BA" ui_fgcolor = "#F4F2FC" + template_fields: Sequence[str] = DbtLocalBaseOperator.template_fields + ("full_refresh",) # type: ignore[operator] def __init__(self, full_refresh: bool = False, **kwargs: Any) -> None: self.full_refresh = full_refresh diff --git a/tests/operators/test_local.py b/tests/operators/test_local.py index 6d58efe39..5d02e4629 100644 --- a/tests/operators/test_local.py +++ b/tests/operators/test_local.py @@ -419,8 +419,8 @@ def test_calculate_openlineage_events_completes_openlineage_errors(mock_processo @pytest.mark.parametrize( "operator_class,expected_template", [ - (DbtSeedLocalOperator, ("env", "vars", "full_refresh")), - (DbtRunLocalOperator, ("env", "vars", "full_refresh")), + (DbtSeedLocalOperator, ("env", "vars", "compiled_sql", "full_refresh")), + (DbtRunLocalOperator, ("env", "vars", "compiled_sql", "full_refresh")), ], ) def test_dbt_base_operator_template_fields(operator_class, expected_template):