Skip to content

Commit

Permalink
🎨 [pre-commit.ci] Auto format from pre-commit.com hooks
Browse files Browse the repository at this point in the history
  • Loading branch information
pre-commit-ci[bot] committed Oct 18, 2023
1 parent 9b94cec commit 95cb840
Show file tree
Hide file tree
Showing 5 changed files with 17 additions and 16 deletions.
4 changes: 2 additions & 2 deletions cosmos/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

12 changes: 7 additions & 5 deletions cosmos/converter.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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

Expand Down
9 changes: 4 additions & 5 deletions cosmos/operators/virtualenv.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
"""
Expand Down Expand Up @@ -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

Expand Down
5 changes: 2 additions & 3 deletions dev/dags/example_virtualenv.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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,
Expand Down
3 changes: 2 additions & 1 deletion tests/operators/test_virtualenv.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down Expand Up @@ -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"
assert venv_operator.venv_dbt_path == "mock-venv/bin/dbt"

0 comments on commit 95cb840

Please sign in to comment.