working code to allow --select

Author: dave-connors-3

data_diff/__main__.py

@@ -228,6 +228,12 @@ def write_usage(self, prog: str, args: str = "", prefix: Optional[str] = None) -
     metavar="PATH",
     help="Which directory to look in for the dbt_project.yml file. Default is the current working directory and its parents.",
 )
+@click.option(
+    "--select",
+    default=None,
+    metavar="PATH",
+    help="select dbt resources to compar using dbt selection syntax",
+)
 def main(conf, run, **kw):
     if kw["table2"] is None and kw["database2"]:
         # Use the "database table table" form
@@ -264,6 +270,7 @@ def main(conf, run, **kw):
                 profiles_dir_override=kw["dbt_profiles_dir"],
                 project_dir_override=kw["dbt_project_dir"],
                 is_cloud=kw["cloud"],
+                dbt_selection=kw["select"],
             )
         else:
             return _data_diff(**kw)

data_diff/dbt.py

@@ -18,7 +18,7 @@
 logger = getLogger(__name__)
 
 
-def import_dbt():
+def import_dbt_parsers():
     try:
         from dbt_artifacts_parser.parser import parse_run_results, parse_manifest
         from dbt.config.renderer import ProfileRenderer
@@ -47,7 +47,7 @@ def import_dbt():
 PROJECT_FILE = "dbt_project.yml"
 PROFILES_FILE = "profiles.yml"
 LOWER_DBT_V = "1.0.0"
-UPPER_DBT_V = "1.4.6"
+UPPER_DBT_V = "1.6.0"
 
 
 # https://github.com/dbt-labs/dbt-core/blob/c952d44ec5c2506995fbad75320acbae49125d3d/core/dbt/cli/resolvers.py#L6
@@ -76,12 +76,12 @@ class DiffVars:
 
 
 def dbt_diff(
-    profiles_dir_override: Optional[str] = None, project_dir_override: Optional[str] = None, is_cloud: bool = False
+    profiles_dir_override: Optional[str] = None, project_dir_override: Optional[str] = None, is_cloud: bool = False, dbt_selection: Optional[str] = None
 ) -> None:
     diff_threads = []
     set_entrypoint_name("CLI-dbt")
     dbt_parser = DbtParser(profiles_dir_override, project_dir_override)
-    models = dbt_parser.get_models()
+    models = dbt_parser.get_models(dbt_selection)
     datadiff_variables = dbt_parser.get_datadiff_variables()
     config_prod_database = datadiff_variables.get("prod_database")
     config_prod_schema = datadiff_variables.get("prod_schema")

data_diff/dbt_parser.py

@@ -1,7 +1,8 @@
 from collections import defaultdict
 import json
+import os
 from pathlib import Path
-from typing import List, Dict, Tuple, Set
+from typing import List, Dict, Tuple, Set, Optional
 
 from packaging.version import parse as parse_version
 
@@ -12,7 +13,7 @@
 logger = getLogger(__name__)
 
 
-def import_dbt():
+def import_dbt_parsers():
     try:
         from dbt_artifacts_parser.parser import parse_run_results, parse_manifest
         from dbt.config.renderer import ProfileRenderer
@@ -23,12 +24,23 @@ def import_dbt():
     return parse_run_results, parse_manifest, ProfileRenderer, yaml
 
 
+# dbt 1.5+ specific stuff to power selection of models
+try:
+    from dbt.cli.main import dbtRunner
+except ImportError:
+    dbtRunner = None
+
+if dbtRunner is not None:
+    dbt_runner = dbtRunner()
+else:
+    dbt_runner = None
+
 RUN_RESULTS_PATH = "target/run_results.json"
 MANIFEST_PATH = "target/manifest.json"
 PROJECT_FILE = "dbt_project.yml"
 PROFILES_FILE = "profiles.yml"
 LOWER_DBT_V = "1.0.0"
-UPPER_DBT_V = "1.4.6"
+UPPER_DBT_V = "1.6.0"
 
 
 # https://github.com/dbt-labs/dbt-core/blob/c952d44ec5c2506995fbad75320acbae49125d3d/core/dbt/cli/resolvers.py#L6
@@ -49,7 +61,7 @@ def legacy_profiles_dir() -> Path:
 
 class DbtParser:
     def __init__(self, profiles_dir_override: str, project_dir_override: str) -> None:
-        self.parse_run_results, self.parse_manifest, self.ProfileRenderer, self.yaml = import_dbt()
+        self.parse_run_results, self.parse_manifest, self.ProfileRenderer, self.yaml = import_dbt_parsers()
         self.profiles_dir = Path(profiles_dir_override or default_profiles_dir())
         self.project_dir = Path(project_dir_override or default_project_dir())
         self.connection = None
@@ -66,7 +78,44 @@ def get_datadiff_variables(self) -> dict:
         vars = get_from_dict_with_raise(self.project_dict, "vars", f"No vars: found in dbt_project.yml.")
         return get_from_dict_with_raise(vars, "data_diff", f"data_diff: section not found in dbt_project.yml vars:.")
 
-    def get_models(self):
+    def get_models(self, dbt_selection: Optional[str] = None):
+        dbt_version = parse_version(self.dbt_version)
+        if dbt_selection:
+            if dbt_version.major == 1 and dbt_version.minor >= 5:
+                return self.get_dbt_selection_models(dbt_selection)
+            else:
+                raise Exception(
+                    f"Use of the `--select` feature requires dbt >= 1.5. Found dbt: v{dbt_version}"
+                )
+        else:
+            return self.get_run_results_models()
+
+    def get_dbt_selection_models(self, dbt_selection: str) -> List[str]:
+        start_dir = os.getcwd()
+        os.chdir(self.project_dir)
+        # log level and format settings needed to prevent dbt from printing to stdout
+        # ls command is used to get the list of model unique_ids
+        results = dbt_runner.invoke([
+            "--log-format",
+            "json",
+            "--log-level",
+            "none",
+            "ls",
+            "--select",
+            dbt_selection,
+            "--resource-type",
+            "model",
+            "--output",
+            "json",
+            '--output-keys',
+            "unique_id"])
+        os.chdir(start_dir)
+        if results.success:
+            model_list = [json.loads(model)["unique_id"] for model in results.result]
+            models = [self.manifest_obj.nodes.get(x) for x in model_list]
+            return models
+
+    def get_run_results_models(self):
         with open(self.project_dir / RUN_RESULTS_PATH) as run_results:
             run_results_dict = json.load(run_results)
             run_results_obj = self.parse_run_results(run_results=run_results_dict)