feat: Provide env_vars to update environment variables

Environment variables are used by dbt to template configuration
files. Although users can now modify the environment of Airflow
workers, there is no interface for them to do it in
airflow-dbt-python.

This commit adds an env_vars argument that can take a dictionary of
environment variables to set during dbt execution. This can be useful
for folks who wish to keep their profiles.yml templates.

Author: tomasfarias

airflow_dbt_python/hooks/dbt.py

@@ -185,6 +185,7 @@ def run_dbt_task(
         delete_before_upload: bool = False,
         replace_on_upload: bool = False,
         artifacts: Optional[Iterable[str]] = None,
+        env_vars: Optional[Dict[str, Any]] = None,
         **kwargs,
     ) -> DbtTaskResult:
         """Run a dbt task with a given configuration and return the results.
@@ -208,6 +209,7 @@ def run_dbt_task(
             upload_dbt_project=upload_dbt_project,
             delete_before_upload=delete_before_upload,
             replace_on_upload=replace_on_upload,
+            env_vars=env_vars,
         ) as dbt_dir:
             config.dbt_task.pre_init_hook(config)
             self.ensure_profiles(config.profiles_dir)
@@ -272,6 +274,7 @@ def dbt_directory(
         upload_dbt_project: bool = False,
         delete_before_upload: bool = False,
         replace_on_upload: bool = False,
+        env_vars: Optional[Dict[str, Any]] = None,
     ) -> Iterator[str]:
         """Provides a temporary directory to execute dbt.
 
@@ -284,43 +287,46 @@ def dbt_directory(
         Yields:
             The temporary directory's name.
         """
+        from airflow_dbt_python.utils.env import update_environment
+
         store_profiles_dir = config.profiles_dir
         store_project_dir = config.project_dir
 
-        with TemporaryDirectory(prefix="airflow_tmp") as tmp_dir:
-            self.log.info("Initializing temporary directory: %s", tmp_dir)
-
-            try:
-                project_dir, profiles_dir = self.prepare_directory(
-                    tmp_dir,
-                    store_project_dir,
-                    store_profiles_dir,
-                )
-            except Exception as e:
-                raise AirflowException(
-                    "Failed to prepare temporary directory for dbt execution"
-                ) from e
-
-            config.project_dir = project_dir
-            config.profiles_dir = profiles_dir
-
-            if getattr(config, "state", None) is not None:
-                state = Path(getattr(config, "state", ""))
-                # Since we are running in a temporary directory, we need to make
-                # state paths relative to this temporary directory.
-                if not state.is_absolute():
-                    setattr(config, "state", str(Path(tmp_dir) / state))
-
-            yield tmp_dir
-
-            if upload_dbt_project is True:
-                self.log.info("Uploading dbt project to: %s", store_project_dir)
-                self.upload_dbt_project(
-                    tmp_dir,
-                    store_project_dir,
-                    replace=replace_on_upload,
-                    delete_before=delete_before_upload,
-                )
+        with update_environment(env_vars):
+            with TemporaryDirectory(prefix="airflow_tmp") as tmp_dir:
+                self.log.info("Initializing temporary directory: %s", tmp_dir)
+
+                try:
+                    project_dir, profiles_dir = self.prepare_directory(
+                        tmp_dir,
+                        store_project_dir,
+                        store_profiles_dir,
+                    )
+                except Exception as e:
+                    raise AirflowException(
+                        "Failed to prepare temporary directory for dbt execution"
+                    ) from e
+
+                config.project_dir = project_dir
+                config.profiles_dir = profiles_dir
+
+                if getattr(config, "state", None) is not None:
+                    state = Path(getattr(config, "state", ""))
+                    # Since we are running in a temporary directory, we need to make
+                    # state paths relative to this temporary directory.
+                    if not state.is_absolute():
+                        setattr(config, "state", str(Path(tmp_dir) / state))
+
+                yield tmp_dir
+
+                if upload_dbt_project is True:
+                    self.log.info("Uploading dbt project to: %s", store_project_dir)
+                    self.upload_dbt_project(
+                        tmp_dir,
+                        store_project_dir,
+                        replace=replace_on_upload,
+                        delete_before=delete_before_upload,
+                    )
 
         config.profiles_dir = store_profiles_dir
         config.project_dir = store_project_dir

airflow_dbt_python/operators/dbt.py

@@ -5,7 +5,7 @@
 import os
 from dataclasses import asdict, is_dataclass
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Optional, Union
+from typing import TYPE_CHECKING, Any, Dict, Optional, Union
 
 from airflow.exceptions import AirflowException
 from airflow.models.baseoperator import BaseOperator
@@ -24,6 +24,7 @@
     "target",
     "state",
     "vars",
+    "env_vars",
 ]
 
 
@@ -95,6 +96,7 @@ def __init__(
         upload_dbt_project: bool = False,
         delete_before_upload: bool = False,
         replace_on_upload: bool = False,
+        env_vars: Optional[Dict[str, Any]] = None,
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
@@ -145,6 +147,7 @@ def __init__(
         self.upload_dbt_project = upload_dbt_project
         self.delete_before_upload = delete_before_upload
         self.replace_on_upload = replace_on_upload
+        self.env_vars = env_vars
 
         self._dbt_hook: Optional[DbtHook] = None
 

airflow_dbt_python/utils/env.py

@@ -0,0 +1,22 @@
+"""Provides utilities to interact with environment variables."""
+import copy
+import os
+from contextlib import contextmanager
+from typing import Any, Dict, Optional
+
+
+@contextmanager
+def update_environment(env_vars: Optional[Dict[str, Any]] = None):
+    """Update current environment with env_vars and restore afterwards."""
+    if not env_vars:
+        # Nothing to update or restore afterwards, so we return early
+        yield os.environ
+        return
+
+    restore_env = copy.deepcopy(os.environ)
+    os.environ.update({k: str(v) for k, v in env_vars.items()})
+
+    try:
+        yield os.environ
+    finally:
+        os.environ = restore_env

tests/conftest.py

@@ -185,6 +185,21 @@ def profiles_file(tmp_path_factory, database):
     return p
 
 
+@pytest.fixture(scope="session")
+def profiles_file_with_env(tmp_path_factory):
+    """Create a profiles.yml file that relies on env variables for db connection."""
+    p = tmp_path_factory.mktemp(".dbt_with_env") / "profiles.yml"
+    profiles_content = PROFILES.format(
+        host="\"{{ env_var('DBT_HOST') }}\"",
+        user="\"{{ env_var('DBT_USER') }}\"",
+        port="\"{{ env_var('DBT_PORT') | int }}\"",
+        password="\"{{ env_var('DBT_ENV_SECRET_PASSWORD') }}\"",
+        dbname="\"{{ env_var('DBT_DBNAME') }}\"",
+    )
+    p.write_text(profiles_content)
+    return p
+
+
 @pytest.fixture(scope="session")
 def airflow_conns(database):
     """Create Airflow connections for testing.

tests/hooks/dbt/test_dbt_hook_base.py

@@ -1,4 +1,5 @@
 """Unit test module for the dbt hook base class."""
+import os
 from pathlib import Path
 
 import pytest
@@ -304,6 +305,28 @@ def test_dbt_directory_with_no_state(
         assert getattr(config, "state", None) is None
 
 
+def test_dbt_directory_with_env_vars(hook, profiles_file_with_env, dbt_project_file):
+    """Test dbt_directory sets environment variables."""
+    config = RunTaskConfig(
+        project_dir=dbt_project_file.parent,
+        profiles_dir=profiles_file_with_env.parent,
+        state="target/",
+    )
+
+    assert "TEST_ENVAR0" not in os.environ
+    assert "TEST_ENVAR1" not in os.environ
+
+    env_vars = {"TEST_ENVAR0": 1, "TEST_ENVAR1": "abc"}
+
+    with hook.dbt_directory(config, env_vars=env_vars) as tmp_dir:
+        assert Path(tmp_dir).exists()
+        assert os.environ.get("TEST_ENVAR0") == "1"
+        assert os.environ.get("TEST_ENVAR1") == "abc"
+
+    assert "TEST_ENVAR0" not in os.environ
+    assert "TEST_ENVAR1" not in os.environ
+
+
 @no_s3_remote
 def test_dbt_base_dbt_directory_changed_to_s3(
     dbt_project_file, profiles_file, s3_bucket, s3_hook, hook

tests/hooks/dbt/test_dbt_run.py

@@ -1,4 +1,5 @@
 """Unit test module for running dbt run with the DbtHook."""
+import os
 from pathlib import Path
 
 import pytest
@@ -26,6 +27,36 @@ def test_dbt_run_task(hook, profiles_file, dbt_project_file, model_files):
         assert run_result.node.unique_id == f"model.test.model_{index}"
 
 
+def test_dbt_run_task_with_env_vars_profile(
+    hook, profiles_file_with_env, dbt_project_file, model_files, database
+):
+    """Test a dbt run task that can render env variables in profile."""
+    env = {
+        "DBT_HOST": database.host,
+        "DBT_USER": database.user,
+        "DBT_PORT": str(database.port),
+        "DBT_ENV_SECRET_PASSWORD": database.password,
+        "DBT_DBNAME": database.dbname,
+    }
+
+    result = hook.run_dbt_task(
+        "run",
+        project_dir=dbt_project_file.parent,
+        profiles_dir=profiles_file_with_env.parent,
+        select=[str(m.stem) for m in model_files],
+        env_vars=env,
+    )
+
+    assert result.success is True
+
+    assert len(result.run_results) == 3
+
+    # Start from 2 as model_1 is ephemeral, and ephemeral models are not built.
+    for index, run_result in enumerate(result.run_results, start=2):
+        assert run_result.status == RunStatus.Success
+        assert run_result.node.unique_id == f"model.test.model_{index}"
+
+
 def test_dbt_run_task_one_file(hook, profiles_file, dbt_project_file, model_files):
     """Test a dbt run task for only one file."""
     result = hook.run_dbt_task(

tests/hooks/dbt/test_dbt_seed.py

@@ -18,6 +18,35 @@ def test_dbt_seed_task(profiles_file, dbt_project_file, seed_files):
 
     assert result.success is True
 
+    assert result.run_results is not None
+    assert len(result.run_results) == 2
+    for index, run_result in enumerate(result.run_results, start=1):
+        assert run_result.status == RunStatus.Success
+        assert run_result.node.unique_id == f"seed.test.seed_{index}"
+
+
+def test_dbt_seed_task_with_env_vars_profile(
+    hook, profiles_file_with_env, dbt_project_file, model_files, database
+):
+    """Test a dbt seed task that can render env variables in profile."""
+    env = {
+        "DBT_HOST": database.host,
+        "DBT_USER": database.user,
+        "DBT_PORT": str(database.port),
+        "DBT_ENV_SECRET_PASSWORD": database.password,
+        "DBT_DBNAME": database.dbname,
+    }
+
+    result = hook.run_dbt_task(
+        "run",
+        project_dir=dbt_project_file.parent,
+        profiles_dir=profiles_file_with_env.parent,
+        select=[str(m.stem) for m in model_files],
+        env_vars=env,
+    )
+
+    assert result.success is True
+
     assert len(result.run_results) == 2
     for index, run_result in enumerate(result.run_results, start=1):
         assert run_result.status == RunStatus.Success

tests/operators/test_dbt_run.py

@@ -112,6 +112,45 @@ def test_dbt_run_models(profiles_file, dbt_project_file, model_files, logs_dir):
     )
 
 
+def test_dbt_run_models_with_env_vars(
+    profiles_file_with_env, dbt_project_file, model_files, logs_dir, database
+):
+    """Test execution of DbtRunOperator with all models using env vars in profile."""
+    env = {
+        "DBT_HOST": database.host,
+        "DBT_USER": database.user,
+        "DBT_PORT": str(database.port),
+        "DBT_ENV_SECRET_PASSWORD": database.password,
+        "DBT_DBNAME": database.dbname,
+    }
+
+    op = DbtRunOperator(
+        task_id="dbt_task",
+        project_dir=dbt_project_file.parent,
+        profiles_dir=profiles_file_with_env.parent,
+        models=[str(m.stem) for m in model_files],
+        do_xcom_push=True,
+        debug=True,
+        env_vars=env,
+    )
+
+    execution_results = op.execute({})
+    run_result = execution_results["results"][0]
+
+    assert run_result["status"] == RunStatus.Success
+
+    log_file = logs_dir / "dbt.log"
+    assert log_file.exists()
+
+    with open(log_file) as f:
+        logs = f.read()
+
+    assert (
+        "OK created view model public.model_4" in logs
+        or "OK created sql view model public.model_4" in logs
+    )
+
+
 def test_dbt_run_models_full_refresh(profiles_file, dbt_project_file, model_files):
     """Test dbt run operator with all model files and full-refresh."""
     op = DbtRunOperator(

tests/utils/test_configs.py

@@ -1,6 +1,9 @@
 """Unit test module for dbt task configuration utilities."""
+import os
+from unittest.mock import patch
+
 import pytest
-from dbt.exceptions import DbtProfileError
+from dbt.exceptions import DbtProfileError, EnvVarMissingError
 from dbt.task.build import BuildTask
 from dbt.task.compile import CompileTask
 from dbt.task.debug import DebugTask
@@ -252,6 +255,40 @@ def test_base_config_create_dbt_profile(hook, profiles_file, dbt_project_file):
     assert target["type"] == "postgres"
 
 
+def test_base_config_create_dbt_profile_with_env_vars(
+    profiles_file_with_env, dbt_project_file, database
+):
+    """Test the create_dbt_profile with a profiles file that contains env vars."""
+    config = BaseConfig(
+        project_dir=dbt_project_file.parent,
+        profiles_dir=profiles_file_with_env.parent,
+    )
+
+    with pytest.raises(EnvVarMissingError):
+        # No environment set yet, we should fail.
+        profile = config.create_dbt_profile()
+
+    env = {
+        "DBT_HOST": database.host,
+        "DBT_USER": database.user,
+        "DBT_PORT": str(database.port),
+        "DBT_ENV_SECRET_PASSWORD": database.password,
+        "DBT_DBNAME": database.dbname,
+    }
+
+    with patch.dict(os.environ, env):
+        profile = config.create_dbt_profile()
+        assert profile.credentials.password == database.password
+
+        target = profile.to_target_dict()
+        assert target["name"] == "test"
+        assert target["type"] == "postgres"
+        assert target["host"] == database.host
+        assert target["user"] == database.user
+        assert target["port"] == database.port
+        assert target["dbname"] == database.dbname
+
+
 def test_base_config_create_dbt_profile_with_extra_target(
     hook, profiles_file, dbt_project_file, airflow_conns
 ):