From d2672b895e954fd0e171b75c9c1f0602a7001c72 Mon Sep 17 00:00:00 2001 From: Tatiana Al-Chueyr Date: Mon, 25 Mar 2024 10:43:26 +0000 Subject: [PATCH] Enable `append_env` in `operator_args` by default (#899) The default behaviour when parsing Cosmos DAGs is to use the system environment variables, but to execute dbt commands from Airflow operators is not. This PR solves this, making the behaviour consistent in DAG parsing and dbt command execution. Closes: #595 --- cosmos/operators/base.py | 9 +++++---- tests/operators/test_azure_container_instance.py | 2 ++ tests/operators/test_docker.py | 5 +---- tests/operators/test_kubernetes.py | 7 ++----- tests/operators/test_local.py | 4 +--- 5 files changed, 11 insertions(+), 16 deletions(-) diff --git a/cosmos/operators/base.py b/cosmos/operators/base.py index 94d5d4a8c..b9f25758d 100644 --- a/cosmos/operators/base.py +++ b/cosmos/operators/base.py @@ -43,10 +43,11 @@ class AbstractDbtBaseOperator(BaseOperator, metaclass=ABCMeta): environment variables for the new process; these are used instead of inheriting the current process environment, which is the default behavior. (templated) - :param append_env: If False(default) uses the environment variables passed in env params - and does not inherit the current process environment. If True, inherits the environment variables + :param append_env: . If True (default), inherits the environment variables from current passes and then environment variable passed by the user will either update the existing - inherited environment variables or the new variables gets appended to it + inherited environment variables or the new variables gets appended to it. + If False, only uses the environment variables passed in env params + and does not inherit the current process environment. :param output_encoding: Output encoding of bash command :param skip_exit_code: If task exits with this exit code, leave the task in ``skipped`` state (default: 99). If set to ``None``, any non-zero @@ -99,7 +100,7 @@ def __init__( db_name: str | None = None, schema: str | None = None, env: dict[str, Any] | None = None, - append_env: bool = False, + append_env: bool = True, output_encoding: str = "utf-8", skip_exit_code: int = 99, partial_parse: bool = True, diff --git a/tests/operators/test_azure_container_instance.py b/tests/operators/test_azure_container_instance.py index 01fa3e20e..a99525ac3 100644 --- a/tests/operators/test_azure_container_instance.py +++ b/tests/operators/test_azure_container_instance.py @@ -53,6 +53,7 @@ def test_dbt_azure_container_instance_operator_get_env(p_context_to_airflow_vars name="my-aci", resource_group="my-rg", project_dir="my/dir", + append_env=False, ) dbt_base_operator.env = { "start_date": "20220101", @@ -90,6 +91,7 @@ def test_dbt_azure_container_instance_operator_check_environment_variables( resource_group="my-rg", project_dir="my/dir", environment_variables={"FOO": "BAR"}, + append_env=False, ) dbt_base_operator.env = { "start_date": "20220101", diff --git a/tests/operators/test_docker.py b/tests/operators/test_docker.py index ad3ec5485..2cfb6b835 100644 --- a/tests/operators/test_docker.py +++ b/tests/operators/test_docker.py @@ -73,10 +73,7 @@ def test_dbt_docker_operator_get_env(p_context_to_airflow_vars: MagicMock, base_ If an end user passes in a """ dbt_base_operator = base_operator( - conn_id="my_airflow_connection", - task_id="my-task", - image="my_image", - project_dir="my/dir", + conn_id="my_airflow_connection", task_id="my-task", image="my_image", project_dir="my/dir", append_env=False ) dbt_base_operator.env = { "start_date": "20220101", diff --git a/tests/operators/test_kubernetes.py b/tests/operators/test_kubernetes.py index 75739111f..d0be2acad 100644 --- a/tests/operators/test_kubernetes.py +++ b/tests/operators/test_kubernetes.py @@ -81,10 +81,7 @@ def test_dbt_kubernetes_operator_get_env(p_context_to_airflow_vars: MagicMock, b If an end user passes in a """ dbt_kube_operator = base_operator( - conn_id="my_airflow_connection", - task_id="my-task", - image="my_image", - project_dir="my/dir", + conn_id="my_airflow_connection", task_id="my-task", image="my_image", project_dir="my/dir", append_env=False ) dbt_kube_operator.env = { "start_date": "20220101", @@ -254,7 +251,7 @@ def cleanup(pod: str, remote_pod: str): def test_created_pod(): - ls_kwargs = {"env_vars": {"FOO": "BAR"}, "namespace": "foo"} + ls_kwargs = {"env_vars": {"FOO": "BAR"}, "namespace": "foo", "append_env": False} ls_kwargs.update(base_kwargs) ls_operator = DbtLSKubernetesOperator(**ls_kwargs) ls_operator.hook = MagicMock() diff --git a/tests/operators/test_local.py b/tests/operators/test_local.py index 956c2a8d1..17b98ee56 100644 --- a/tests/operators/test_local.py +++ b/tests/operators/test_local.py @@ -329,9 +329,7 @@ def test_dbt_base_operator_get_env(p_context_to_airflow_vars: MagicMock) -> None If an end user passes in a """ dbt_base_operator = ConcreteDbtLocalBaseOperator( - profile_config=profile_config, - task_id="my-task", - project_dir="my/dir", + profile_config=profile_config, task_id="my-task", project_dir="my/dir", append_env=False ) dbt_base_operator.env = { "start_date": "20220101",