CORE-1496: support --dbt-vars in every edr command that runs the internal dbt project

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:

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(

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()

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

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

elementary/monitor/debug.py

@@ -1,22 +1,17 @@
 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:
     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)

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):
     """

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:

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.",
+)