diff --git a/elementary/clients/dbt/factory.py b/elementary/clients/dbt/factory.py index 8751a9ec3..4b1159804 100644 --- a/elementary/clients/dbt/factory.py +++ b/elementary/clients/dbt/factory.py @@ -11,6 +11,7 @@ is_dbt2_binary_available, ) from elementary.clients.dbt.subprocess_dbt_runner import SubprocessDbtRunner +from elementary.config.config import Config class RunnerMethod(Enum): @@ -51,6 +52,22 @@ def create_dbt_runner( ) +def create_dbt_runner_from_config( + config: Config, + project_dir: str, + force_dbt_deps: bool = False, +) -> CommandLineDbtRunner: + return create_dbt_runner( + project_dir, + config.profiles_dir, + config.profile_target, + env_vars=config.env_vars, + vars=config.dbt_vars, + run_deps_if_needed=config.run_dbt_deps_if_needed, + force_dbt_deps=force_dbt_deps, + ) + + def get_dbt_runner_method() -> RunnerMethod: runner_method = os.getenv("DBT_RUNNER_METHOD") if runner_method: diff --git a/elementary/config/config.py b/elementary/config/config.py index fbce626c5..3d30ca0dc 100644 --- a/elementary/config/config.py +++ b/elementary/config/config.py @@ -1,6 +1,6 @@ import os from pathlib import Path -from typing import Optional +from typing import Any, Dict, Optional import google.auth # type: ignore[import] from dateutil import tz @@ -77,6 +77,8 @@ def __init__( project_name: Optional[str] = None, quiet_logs: Optional[bool] = None, ssl_ca_bundle: Optional[str] = None, + dbt_vars: Optional[Dict[str, Any]] = None, + days_back: Optional[int] = None, ): self.config_dir = config_dir self.profiles_dir = profiles_dir @@ -90,6 +92,11 @@ def __init__( self.env_vars = dict() self.env_vars.update(self._parse_dbt_quoting_to_env_vars(dbt_quoting)) + self.dbt_vars = dbt_vars + if days_back: + self.dbt_vars = self.dbt_vars or {} + self.dbt_vars["days_back"] = days_back + config = self._load_configuration() self.target_dir = self._first_not_none( diff --git a/elementary/monitor/cli.py b/elementary/monitor/cli.py index 8df8d20e2..f1c4cbbfd 100644 --- a/elementary/monitor/cli.py +++ b/elementary/monitor/cli.py @@ -15,9 +15,7 @@ from elementary.monitor.debug import Debug from elementary.tracking.anonymous_tracking import AnonymousCommandLineTracking from elementary.utils import bucket_path -from elementary.utils.ordered_yaml import OrderedYaml - -yaml = OrderedYaml() +from elementary.utils.cli_utils import dbt_vars_option class Command: @@ -43,6 +41,7 @@ def decorator(func): default=False, help="Disable sampling of data. Useful if your data contains PII.", )(func) + func = dbt_vars_option(func) func = click.option( "--dbt-quoting", "-dq", @@ -238,12 +237,6 @@ def get_cli_properties() -> dict: help="Force running a full refresh of all incremental models in the edr dbt package (usually this is not needed, " "see documentation to learn more).", ) -@click.option( - "--dbt-vars", - type=str, - default=None, - help="Specify raw YAML string of your dbt variables.", -) @click.option( "--test", type=bool, @@ -355,7 +348,6 @@ def monitor( fg="bright_red", ) slack_webhook = deprecated_slack_webhook - vars = yaml.loads(dbt_vars) if dbt_vars else None config = Config( config_dir=config_dir, profiles_dir=profiles_dir, @@ -376,6 +368,8 @@ def monitor( maximum_columns_in_alert_samples=maximum_columns_in_alert_samples, quiet_logs=quiet_logs, ssl_ca_bundle=ssl_ca_bundle, + dbt_vars=dbt_vars, + days_back=days_back, ) anonymous_tracking = AnonymousCommandLineTracking(config) anonymous_tracking.set_env("use_select", bool(select)) @@ -410,9 +404,7 @@ def monitor( anonymous_tracking.track_cli_start( Command.MONITOR, get_cli_properties(), ctx.command.name ) - success = data_monitoring.run_alerts( - days_back, full_refresh_dbt_package, dbt_vars=vars - ) + success = data_monitoring.run_alerts(days_back, full_refresh_dbt_package) anonymous_tracking.track_cli_end( Command.MONITOR, data_monitoring.properties(), ctx.command.name ) @@ -471,6 +463,7 @@ def report( select, target_path, quiet_logs, + dbt_vars, ): """ Generate a local observability report of your warehouse. @@ -485,6 +478,7 @@ def report( dbt_quoting=dbt_quoting, env=env, quiet_logs=quiet_logs, + dbt_vars=dbt_vars, ) anonymous_tracking = AnonymousCommandLineTracking(config) anonymous_tracking.set_env("use_select", bool(select)) @@ -704,6 +698,7 @@ def send_report( target_path, quiet_logs, ssl_ca_bundle, + dbt_vars, ): """ Generate and send the report to an external platform. @@ -748,6 +743,7 @@ def send_report( project_name=project_name, quiet_logs=quiet_logs, ssl_ca_bundle=ssl_ca_bundle, + dbt_vars=dbt_vars, ) anonymous_tracking = AnonymousCommandLineTracking(config) anonymous_tracking.set_env("use_select", bool(select)) @@ -809,9 +805,10 @@ def send_report( help="Which directory to look in for the profiles.yml file. " "If not set, edr will look in HOME/.dbt/", ) +@dbt_vars_option @click.pass_context -def debug(ctx, profiles_dir): - config = Config(profiles_dir=profiles_dir) +def debug(ctx, profiles_dir, dbt_vars): + config = Config(profiles_dir=profiles_dir, dbt_vars=dbt_vars) anonymous_tracking = AnonymousCommandLineTracking(config) anonymous_tracking.track_cli_start(Command.DEBUG, None, ctx.command.name) success = Debug(config).run() diff --git a/elementary/monitor/data_monitoring/alerts/data_monitoring_alerts.py b/elementary/monitor/data_monitoring/alerts/data_monitoring_alerts.py index 0b0158f4d..afd355095 100644 --- a/elementary/monitor/data_monitoring/alerts/data_monitoring_alerts.py +++ b/elementary/monitor/data_monitoring/alerts/data_monitoring_alerts.py @@ -142,20 +142,11 @@ def _get_integration_client( tracking=self.tracking, ) - def _populate_data( - self, - days_back: Optional[int] = None, - dbt_full_refresh: bool = False, - dbt_vars: Optional[dict] = None, - ) -> bool: + def _populate_data(self, dbt_full_refresh: bool = False) -> bool: logger.info("Running internal dbt run to populate alerts") - vars = dbt_vars or dict() - if days_back: - vars.update(days_back=days_back) success = self.internal_dbt_runner.run( select="elementary_cli.alerts.alerts_v2", full_refresh=dbt_full_refresh, - vars=vars, ) self.execution_properties["alerts_populate_success"] = success if not success: @@ -426,14 +417,11 @@ def run_alerts( self, days_back: int, dbt_full_refresh: bool = False, - dbt_vars: Optional[dict] = None, ) -> bool: # Populate data if self.should_populate_data: popopulated_data_successfully = self._populate_data( - days_back=days_back, dbt_full_refresh=dbt_full_refresh, - dbt_vars=dbt_vars, ) if not popopulated_data_successfully: self.success = False diff --git a/elementary/monitor/data_monitoring/data_monitoring.py b/elementary/monitor/data_monitoring/data_monitoring.py index dde9fd000..ca501d51f 100644 --- a/elementary/monitor/data_monitoring/data_monitoring.py +++ b/elementary/monitor/data_monitoring/data_monitoring.py @@ -3,10 +3,10 @@ from packaging import version -from elementary.clients.dbt.factory import create_dbt_runner +from elementary.clients.dbt.factory import create_dbt_runner_from_config from elementary.config.config import Config -from elementary.monitor import dbt_project_utils from elementary.monitor.data_monitoring.schema import FiltersSchema, WarehouseInfo +from elementary.monitor.dbt_project_utils import CLI_DBT_PROJECT_PATH from elementary.tracking.anonymous_tracking import AnonymousTracking from elementary.tracking.tracking_interface import Tracking from elementary.utils import package @@ -58,12 +58,9 @@ def __init__( self.selector_filter = selector_filter def _init_internal_dbt_runner(self): - internal_dbt_runner = create_dbt_runner( - dbt_project_utils.CLI_DBT_PROJECT_PATH, - self.config.profiles_dir, - self.config.profile_target, - env_vars=self.config.env_vars, - run_deps_if_needed=self.config.run_dbt_deps_if_needed, + internal_dbt_runner = create_dbt_runner_from_config( + self.config, + CLI_DBT_PROJECT_PATH, force_dbt_deps=self.force_update_dbt_package, ) return internal_dbt_runner diff --git a/elementary/monitor/debug.py b/elementary/monitor/debug.py index b8d8af0ff..9e7d7d64e 100644 --- a/elementary/monitor/debug.py +++ b/elementary/monitor/debug.py @@ -1,9 +1,9 @@ import click -from elementary.clients.dbt.factory import create_dbt_runner +from elementary.clients.dbt.factory import create_dbt_runner_from_config from elementary.config.config import Config from elementary.exceptions.exceptions import DbtCommandError -from elementary.monitor import dbt_project_utils +from elementary.monitor.dbt_project_utils import CLI_DBT_PROJECT_PATH class Debug: @@ -11,12 +11,7 @@ def __init__(self, config: Config): self.config = config def run(self) -> bool: - dbt_runner = create_dbt_runner( - dbt_project_utils.CLI_DBT_PROJECT_PATH, - self.config.profiles_dir, - self.config.profile_target, - run_deps_if_needed=self.config.run_dbt_deps_if_needed, - ) + dbt_runner = create_dbt_runner_from_config(self.config, CLI_DBT_PROJECT_PATH) try: dbt_runner.run_operation("elementary_cli.test_conn", quiet=True) diff --git a/elementary/operations/cli.py b/elementary/operations/cli.py index f92eb99d1..d9c40d244 100644 --- a/elementary/operations/cli.py +++ b/elementary/operations/cli.py @@ -3,6 +3,7 @@ from elementary.config.config import Config from elementary.operations.upload_source_freshness import UploadSourceFreshnessOperation from elementary.tracking.anonymous_tracking import AnonymousCommandLineTracking +from elementary.utils.cli_utils import dbt_vars_option _MODULE_NAME = "run-operation" @@ -51,6 +52,7 @@ def run_operation(): default=25, help="Amount of rows to insert per insert statement.", ) +@dbt_vars_option @click.pass_context def upload_source_freshness(ctx, rows_per_insert: int, **conf): """ diff --git a/elementary/operations/upload_source_freshness.py b/elementary/operations/upload_source_freshness.py index 73f4ad2b4..7b9338b29 100644 --- a/elementary/operations/upload_source_freshness.py +++ b/elementary/operations/upload_source_freshness.py @@ -5,9 +5,9 @@ import click from alive_progress import alive_it -from elementary.clients.dbt.factory import create_dbt_runner +from elementary.clients.dbt.factory import create_dbt_runner_from_config from elementary.config.config import Config -from elementary.monitor import dbt_project_utils +from elementary.monitor.dbt_project_utils import CLI_DBT_PROJECT_PATH from elementary.utils.ordered_yaml import OrderedYaml @@ -36,12 +36,7 @@ def get_sources_file_contents(self) -> dict: return json.loads(source_path.read_text()) def upload_results(self, results: dict, metadata: dict, rows_per_insert: int): - dbt_runner = create_dbt_runner( - dbt_project_utils.CLI_DBT_PROJECT_PATH, - self.config.profiles_dir, - self.config.profile_target, - run_deps_if_needed=self.config.run_dbt_deps_if_needed, - ) + dbt_runner = create_dbt_runner_from_config(self.config, CLI_DBT_PROJECT_PATH) invocation_id = metadata.get("invocation_id") if not invocation_id: diff --git a/elementary/utils/cli_utils.py b/elementary/utils/cli_utils.py index 377a5831b..57cd1a684 100644 --- a/elementary/utils/cli_utils.py +++ b/elementary/utils/cli_utils.py @@ -1,6 +1,10 @@ +from typing import Any, Dict, Optional + import click from click import ClickException +from elementary.utils.ordered_yaml import OrderedYaml + class RequiredIf(click.Option): def __init__(self, *args, **kwargs): @@ -28,3 +32,28 @@ def handle_parse_result(self, ctx, opts, args): self.prompt = None return super(RequiredIf, self).handle_parse_result(ctx, opts, args) + + +def _parse_dbt_vars( + ctx: click.Context, param: click.Parameter, value: Optional[str] +) -> Optional[Dict[str, Any]]: + if not value: + return None + try: + dbt_vars = OrderedYaml().loads(value) + except Exception as exc: + raise click.BadParameter(f"Invalid YAML: {exc}") from exc + if not isinstance(dbt_vars, dict): + raise click.BadParameter( + "dbt vars must be a key-value mapping, for example '{query_max_size: 1000000}'." + ) + return dbt_vars + + +dbt_vars_option = click.option( + "--dbt-vars", + type=str, + default=None, + callback=_parse_dbt_vars, + help="Specify raw YAML string of your dbt variables.", +) diff --git a/tests/unit/cli/__init__.py b/tests/unit/cli/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit/cli/test_dbt_vars_option.py b/tests/unit/cli/test_dbt_vars_option.py new file mode 100644 index 000000000..c2650e24b --- /dev/null +++ b/tests/unit/cli/test_dbt_vars_option.py @@ -0,0 +1,147 @@ +from typing import List +from unittest import mock + +import pytest +from click.testing import CliRunner + +from elementary.config.config import Config +from elementary.monitor.cli import monitor +from elementary.operations.cli import run_operation + +DBT_VARS = "{query_max_size: 1000, nested: {key: value}}" +EXPECTED_DBT_VARS = {"query_max_size": 1000, "nested": {"key": "value"}} + + +@pytest.fixture(autouse=True) +def isolated_env(tmp_path, monkeypatch): + monkeypatch.setenv("DBT_LOG_PATH", str(tmp_path)) + monkeypatch.setattr(Config, "DEFAULT_CONFIG_DIR", str(tmp_path)) + monkeypatch.setattr(Config, "DEFAULT_TARGET_PATH", str(tmp_path)) + with mock.patch("elementary.monitor.cli.AnonymousCommandLineTracking"), mock.patch( + "elementary.operations.cli.AnonymousCommandLineTracking" + ): + yield + + +def _common_args(tmp_path) -> List[str]: + return ["--config-dir", str(tmp_path), "--target-path", str(tmp_path)] + + +def _invoke(cli, args: List[str]): + result = CliRunner().invoke(cli, args, catch_exceptions=False) + assert result.exit_code == 0, result.output + return result + + +@pytest.mark.parametrize( + "extra_args,expected_dbt_vars,expected_days_back", + [ + ([], {"days_back": 1}, 1), + (["--dbt-vars", DBT_VARS], {**EXPECTED_DBT_VARS, "days_back": 1}, 1), + (["--days-back", "3"], {"days_back": 3}, 3), + (["--days-back", "0"], None, 0), + ( + ["--dbt-vars", "{days_back: 10}", "--days-back", "3"], + {"days_back": 3}, + 3, + ), + ], +) +def test_monitor_dbt_vars(tmp_path, extra_args, expected_dbt_vars, expected_days_back): + with mock.patch( + "elementary.monitor.cli.DataMonitoringAlerts" + ) as data_monitoring_alerts: + _invoke( + monitor, + [*_common_args(tmp_path), "--slack-webhook", "mock", *extra_args], + ) + config = data_monitoring_alerts.call_args.kwargs["config"] + assert config.dbt_vars == expected_dbt_vars + data_monitoring_alerts.return_value.run_alerts.assert_called_once_with( + expected_days_back, False + ) + + +@pytest.mark.parametrize("dbt_vars_args", [[], ["--dbt-vars", DBT_VARS]]) +def test_report_dbt_vars(tmp_path, dbt_vars_args): + with mock.patch("elementary.monitor.cli.SelectorFilter"), mock.patch( + "elementary.monitor.cli.DataMonitoringReport" + ) as data_monitoring_report: + data_monitoring_report.return_value.generate_report.return_value = ( + True, + None, + ) + _invoke(monitor, ["report", *_common_args(tmp_path), *dbt_vars_args]) + config = data_monitoring_report.call_args.kwargs["config"] + assert config.dbt_vars == (EXPECTED_DBT_VARS if dbt_vars_args else None) + + +@pytest.mark.parametrize("dbt_vars_args", [[], ["--dbt-vars", DBT_VARS]]) +def test_send_report_dbt_vars(tmp_path, dbt_vars_args): + with mock.patch("elementary.monitor.cli.SelectorFilter"), mock.patch( + "elementary.monitor.cli.DataMonitoringReport" + ) as data_monitoring_report: + _invoke( + monitor, + [ + "send-report", + *_common_args(tmp_path), + "--slack-token", + "mock", + "--slack-channel-name", + "mock", + *dbt_vars_args, + ], + ) + config = data_monitoring_report.call_args.kwargs["config"] + assert config.dbt_vars == (EXPECTED_DBT_VARS if dbt_vars_args else None) + + +@pytest.mark.parametrize("dbt_vars_args", [[], ["--dbt-vars", DBT_VARS]]) +def test_debug_dbt_vars(dbt_vars_args): + with mock.patch("elementary.monitor.cli.Debug") as debug: + _invoke(monitor, ["debug", *dbt_vars_args]) + config = debug.call_args.args[0] + assert config.dbt_vars == (EXPECTED_DBT_VARS if dbt_vars_args else None) + + +@pytest.mark.parametrize("dbt_vars_args", [[], ["--dbt-vars", DBT_VARS]]) +def test_upload_source_freshness_dbt_vars(tmp_path, dbt_vars_args): + with mock.patch( + "elementary.operations.cli.UploadSourceFreshnessOperation" + ) as operation: + _invoke( + run_operation, + [ + "upload-source-freshness", + "--target-path", + str(tmp_path), + *dbt_vars_args, + ], + ) + config = operation.call_args.args[0] + assert config.dbt_vars == (EXPECTED_DBT_VARS if dbt_vars_args else None) + + +@pytest.mark.parametrize( + "args", + [ + ["report"], + ["send-report"], + ["debug"], + ], +) +@pytest.mark.parametrize("dbt_vars", ["{invalid", "not_a_mapping"]) +def test_invalid_dbt_vars(args, dbt_vars): + result = CliRunner().invoke(monitor, [*args, "--dbt-vars", dbt_vars]) + assert result.exit_code == 2 + assert "--dbt-vars" in result.output + + +@pytest.mark.parametrize("dbt_vars", ["{invalid", "not_a_mapping"]) +def test_invalid_dbt_vars_upload_source_freshness(dbt_vars): + result = CliRunner().invoke( + run_operation, ["upload-source-freshness", "--dbt-vars", dbt_vars] + ) + assert result.exit_code == 2 + assert "--dbt-vars" in result.output