diff --git a/cosmos/config.py b/cosmos/config.py index fbd0cb9f7..820fb19c6 100644 --- a/cosmos/config.py +++ b/cosmos/config.py @@ -207,10 +207,10 @@ class ExecutionConfig: :param execution_mode: The execution mode for dbt. Defaults to local :param dbt_executable_path: The path to the dbt executable. Defaults to dbt if available on the path. - :param virtualenv_dir: Directory path to locate the (cached) virtual env that + :param virtualenv_dir: Directory path to locate the (cached) virtual env that should be used for execution when execution mode is set to `ExecutionMode.VIRTUALENV` """ + execution_mode: ExecutionMode = ExecutionMode.LOCAL dbt_executable_path: str | Path = get_system_dbt() virtualenv_dir: str | Path | None = None - diff --git a/cosmos/converter.py b/cosmos/converter.py index 7a4b3e3f0..a4df7464e 100644 --- a/cosmos/converter.py +++ b/cosmos/converter.py @@ -121,9 +121,11 @@ def __init__( dbt_executable_path = execution_config.dbt_executable_path node_converters = render_config.node_converters - if (execution_mode != ExecutionMode.VIRTUALENV and execution_config.virtualenv_dir is not None): - logger.warning("`ExecutionConfig.virtualenv_dir` is only supported when \ - ExecutionConfig.execution_mode is set to ExecutionMode.VIRTUALENV.") + if execution_mode != ExecutionMode.VIRTUALENV and execution_config.virtualenv_dir is not None: + logger.warning( + "`ExecutionConfig.virtualenv_dir` is only supported when \ + ExecutionConfig.execution_mode is set to ExecutionMode.VIRTUALENV." + ) profile_args = {} if profile_config.profile_mapping: @@ -160,9 +162,9 @@ def __init__( "emit_datasets": emit_datasets, } - if (execution_mode == ExecutionMode.VIRTUALENV and execution_config.virtualenv_dir is not None): + if execution_mode == ExecutionMode.VIRTUALENV and execution_config.virtualenv_dir is not None: task_args["virtualenv_dir"] = execution_config.virtualenv_dir - + if dbt_executable_path: task_args["dbt_executable_path"] = dbt_executable_path diff --git a/cosmos/operators/virtualenv.py b/cosmos/operators/virtualenv.py index 5fff03b9f..6a8c43c25 100644 --- a/cosmos/operators/virtualenv.py +++ b/cosmos/operators/virtualenv.py @@ -61,7 +61,7 @@ def venv_dbt_path( """ Path to the dbt binary within a Python virtualenv. - The first time this property is called, it creates a new/temporary and installs the dependencies + The first time this property is called, it creates a new/temporary and installs the dependencies based on the self.py_requirements and self.py_system_site_packages, or retrieves an existing virtualenv. This value is cached for future calls. """ @@ -100,15 +100,14 @@ def _get_or_create_venv_py_interpreter(self) -> str: self.log.info(f"Checking for venv interpreter: {py_interpreter_path} : {py_interpreter_path.is_file()}") if py_interpreter_path.is_file(): - self.log.info(f"Found Python interpreter in cached virtualenv: `{str(py_interpreter_path)}`") return str(py_interpreter_path) - + self.log.info(f"Creating virtualenv at `{self._venv_dir}") venv_directory = str(self._venv_dir) - + else: - self.log.info(f"Creating temporary virtualenv") + self.log.info("Creating temporary virtualenv") self._venv_tmp_dir = TemporaryDirectory(prefix="cosmos-venv") venv_directory = self._venv_tmp_dir.name diff --git a/dev/dags/example_virtualenv.py b/dev/dags/example_virtualenv.py index 7ec4b0a16..475e4e66c 100644 --- a/dev/dags/example_virtualenv.py +++ b/dev/dags/example_virtualenv.py @@ -4,7 +4,6 @@ import os from datetime import datetime from pathlib import Path -from airflow.configuration import get_airflow_home from cosmos import DbtDag, ExecutionMode, ExecutionConfig, ProjectConfig, ProfileConfig from cosmos.profiles import PostgresUserPasswordProfileMapping @@ -31,9 +30,9 @@ profile_config=profile_config, execution_config=ExecutionConfig( execution_mode=ExecutionMode.VIRTUALENV, - # We can enable this flag if we want Airflow to create one virtualenv + # We can enable this flag if we want Airflow to create one virtualenv # and reuse that within the whole DAG. - # virtualenv_dir=f"{get_airflow_home()}/persistent-venv", + # virtualenv_dir=f"{get_airflow_home()}/persistent-venv", ), operator_args={ "py_system_site_packages": False, diff --git a/tests/operators/test_virtualenv.py b/tests/operators/test_virtualenv.py index afe0d35a3..39ab5fea6 100644 --- a/tests/operators/test_virtualenv.py +++ b/tests/operators/test_virtualenv.py @@ -66,6 +66,7 @@ def test_run_command( assert dbt_cmd[0][0][1] == "do-something" assert mock_execute.call_count == 2 + @patch("airflow.utils.python_virtualenv.execute_in_subprocess") @patch("cosmos.operators.virtualenv.DbtLocalBaseOperator.calculate_openlineage_events_completes") @patch("cosmos.operators.virtualenv.DbtLocalBaseOperator.store_compiled_sql") @@ -94,4 +95,4 @@ def test_supply_virtualenv_dir_flag( emit_datasets=False, virtualenv_dir="mock-venv", ) - assert venv_operator.venv_dbt_path == "mock-venv/bin/dbt" \ No newline at end of file + assert venv_operator.venv_dbt_path == "mock-venv/bin/dbt"