Skip to content

Commit

Permalink
Fix: operator_args modified in place in Airflow converter (astronom…
Browse files Browse the repository at this point in the history
…er#835)

## Description

This is a fix for astronomer#833 where a user in Slack reported a bug where if the
same `operator_args` are used in multiple `DbtTaskGroup`'s the `vars`
were only used in the first task group.

I added a test which would have caught the bug, so that it isn't
reintroduced later.

## Related Issue(s)

closes astronomer#833 

## Breaking Change?

None

## Checklist

- [ ] I have made corresponding changes to the documentation (if
required)
- [x] I have added tests that prove my fix is effective or that my
feature works
  • Loading branch information
jbandoro authored and arojasb3 committed Jul 14, 2024
1 parent 9fe2f1d commit 23c3e05
Show file tree
Hide file tree
Showing 2 changed files with 35 additions and 2 deletions.
4 changes: 2 additions & 2 deletions cosmos/converter.py
Original file line number Diff line number Diff line change
Expand Up @@ -221,8 +221,8 @@ def __init__(

validate_adapted_user_config(execution_config, project_config, render_config)

env_vars = project_config.env_vars or operator_args.pop("env", None)
dbt_vars = project_config.dbt_vars or operator_args.pop("vars", None)
env_vars = project_config.env_vars or operator_args.get("env")
dbt_vars = project_config.dbt_vars or operator_args.get("vars")

# Previously, we were creating a cosmos.dbt.project.DbtProject
# DbtProject has now been replaced with ProjectConfig directly
Expand Down
33 changes: 33 additions & 0 deletions tests/test_converter.py
Original file line number Diff line number Diff line change
Expand Up @@ -405,3 +405,36 @@ def test_converter_project_config_dbt_vars_with_custom_load_mode(
)
_, kwargs = mock_legacy_dbt_project.call_args
assert kwargs["dbt_vars"] == {"key": "value"}


@patch("cosmos.config.ProjectConfig.validate_project")
@patch("cosmos.converter.build_airflow_graph")
@patch("cosmos.converter.DbtGraph.load")
def test_converter_multiple_calls_same_operator_args(
mock_dbt_graph_load, mock_validate_project, mock_build_airflow_graph
):
"""Tests if the DbttoAirflowConverter is called more than once with the same operator_args, the
operator_args are not modified.
"""
project_config = ProjectConfig(project_name="fake-project", dbt_project_path="/some/project/path")
execution_config = ExecutionConfig()
render_config = RenderConfig()
profile_config = MagicMock()
operator_args = {
"install_deps": True,
"vars": {"key": "value"},
"env": {"key": "value"},
}
original_operator_args = operator_args.copy()
for _ in range(2):
with DAG("test-id", start_date=datetime(2022, 1, 1)) as dag:
DbtToAirflowConverter(
dag=dag,
nodes=nodes,
project_config=project_config,
profile_config=profile_config,
execution_config=execution_config,
render_config=render_config,
operator_args=operator_args,
)
assert operator_args == original_operator_args

0 comments on commit 23c3e05

Please sign in to comment.