Code sharing between classes.

elementary-data · Jun 28, 2023 · d71b354 · d71b354
1 parent 931860f
commit d71b354
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 33 deletions.
diff --git a/elementary/clients/dbt/base_dbt_runner.py b/elementary/clients/dbt/base_dbt_runner.py
@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import Optional
+from typing import Any, Dict, Optional
 
 
 class BaseDbtRunner(ABC):
@@ -8,10 +8,14 @@ def __init__(
         project_dir: str,
         profiles_dir: Optional[str] = None,
         target: Optional[str] = None,
+        vars: Optional[dict] = None,
+        secret_vars: Optional[dict] = None,
     ) -> None:
         self.project_dir = project_dir
         self.profiles_dir = profiles_dir
         self.target = target
+        self.vars = vars or {}
+        self.secret_vars = secret_vars or {}
 
     @abstractmethod
     def deps(self, *args, **kwargs):
@@ -48,3 +52,13 @@ def ls(self, *args, **kwargs):
     @abstractmethod
     def source_freshness(self, *args, **kwargs):
         raise NotImplementedError
+
+    def _get_all_vars(self, vars: Optional[Dict[str, Any]] = None):
+        return {
+            **self.vars,
+            **self.secret_vars,
+            **(vars or {}),
+        }
+
+    def _get_secret_masked_vars(self, vars: Dict[str, Any]):
+        return {k: v if k not in self.secret_vars else "***" for k, v in vars.items()}
diff --git a/elementary/clients/dbt/dbt_runner.py b/elementary/clients/dbt/dbt_runner.py
@@ -33,11 +33,9 @@ def __init__(
         vars: Optional[Dict[str, Any]] = None,
         secret_vars: Optional[Dict[str, Any]] = None,
     ) -> None:
-        super().__init__(project_dir, profiles_dir, target)
+        super().__init__(project_dir, profiles_dir, target, vars, secret_vars)
         self.raise_on_failure = raise_on_failure
         self.env_vars = env_vars
-        self.vars = vars
-        self.secret_vars = secret_vars
 
     def _run_command(
         self,
@@ -56,25 +54,16 @@ def _run_command(
             dbt_command.extend(["--profiles-dir", self.profiles_dir])
         if self.target:
             dbt_command.extend(["--target", self.target])
-        agg_vars = {
-            **(self.vars or {}),
-            **(self.secret_vars or {}),
-            **(vars or {}),
-        }
+        all_vars = self._get_all_vars(vars)
         log_command = " ".join(
             [
                 *dbt_command,
                 "--vars",
-                json.dumps(
-                    {
-                        k: v if k not in self.secret_vars else "***"
-                        for k, v in agg_vars.items()
-                    }
-                ),
+                json.dumps(self._get_secret_masked_vars(all_vars)),
             ]
         )
-        if agg_vars:
-            json_vars = json.dumps(agg_vars)
+        if all_vars:
+            json_vars = json.dumps(all_vars)
             dbt_command.extend(["--vars", json_vars])
         log_msg = f"Running {log_command}"
         if not quiet:

diff --git a/elementary/clients/dbt/slim_dbt_runner.py b/elementary/clients/dbt/slim_dbt_runner.py
@@ -81,9 +81,7 @@ def __init__(
         secret_vars: Optional[dict] = None,
         **kwargs,
     ):
-        super().__init__(project_dir, profiles_dir, target)
-        self.vars = vars
-        self.secret_vars = secret_vars
+        super().__init__(project_dir, profiles_dir, target, vars, secret_vars)
         self.config = None
         self.adapter = None
         self.adapter_name = None
@@ -171,16 +169,12 @@ def run_operation(
         quiet: bool = False,
         **kwargs,
     ) -> list:
-        agg_vars = {
-            **(self.vars or {}),
-            **(self.secret_vars or {}),
-            **(vars or {}),
-        }
+        all_vars = self._get_all_vars(vars)
         self._load_runner(
             project_dir=self.project_dir,
             profiles_dir=self.profiles_dir,
             target=self.target,
-            vars=agg_vars,
+            vars=all_vars,
         )
         log_command = [
             "dbt",
@@ -189,16 +183,11 @@ def run_operation(
             "--args",
             json.dumps(macro_args),
         ]
-        if agg_vars:
+        if all_vars:
             log_command.extend(
                 [
                     "--vars",
-                    json.dumps(
-                        {
-                            k: v if k not in self.secret_vars else "***"
-                            for k, v in agg_vars.items()
-                        }
-                    ),
+                    json.dumps(self._get_secret_masked_vars(all_vars)),
                 ]
             )
         log_msg = f"Running {' '.join(log_command)}"