Refactor Artifacts Parser to be Native so it's less brittle with each dbt version change

data_diff/dbt_parser.py

@@ -6,8 +6,8 @@
 import yaml
 
 from packaging.version import parse as parse_version
-import pydantic
-from dbt_artifacts_parser.parser import parse_run_results, parse_manifest
+from pydantic import BaseModel, Field
+from dbt_artifacts_parser.parser import parse_manifest  # TODO: remove this import
 from dbt.config.renderer import ProfileRenderer
 
 from data_diff.errors import (
@@ -81,19 +81,31 @@ def legacy_profiles_dir() -> Path:
     return Path.home() / ".dbt"
 
 
-class TDatadiffModelConfig(pydantic.BaseModel):
+class TDatadiffModelConfig(BaseModel):
     where_filter: Optional[str] = None
     include_columns: List[str] = []
     exclude_columns: List[str] = []
 
 
-class TDatadiffConfig(pydantic.BaseModel):
+class TDatadiffConfig(BaseModel):
     prod_database: Optional[str] = None
     prod_schema: Optional[str] = None
     prod_custom_schema: Optional[str] = None
     datasource_id: Optional[int] = None
 
 
+class RunResultsJsonConfig(BaseModel):
+    class Metadata(BaseModel):
+        dbt_version: str = Field(..., regex=r'^\d+\.\d+\.\d+([a-zA-Z0-9]+)?$')
+
+    class Results(BaseModel):
+        status: str
+        unique_id: str
+
+    metadata: Metadata
+    results: List[Results]
+
+
 class DbtParser:
     def __init__(
         self,
@@ -106,10 +118,10 @@ def __init__(
         self.project_dir = Path(project_dir_override or default_project_dir())
         self.connection = {}
         self.project_dict = self.get_project_dict()
-        self.dev_manifest_obj = self.get_manifest_obj(self.project_dir / MANIFEST_PATH)
+        self.dev_manifest_obj = self.get_manifest_obj(self.project_dir / MANIFEST_PATH) # TODO: this is where the manfiest object gets called for dev
         self.prod_manifest_obj = None
         if state:
-            self.prod_manifest_obj = self.get_manifest_obj(Path(state))
+            self.prod_manifest_obj = self.get_manifest_obj(Path(state)) # TODO: this is where the manfiest object gets called for prod based on a state parameter
 
         self.dbt_user_id = self.dev_manifest_obj.metadata.user_id
         self.dbt_version = self.dev_manifest_obj.metadata.dbt_version
@@ -230,13 +242,16 @@ def get_simple_model_selection(self, dbt_selection: str):
 
         return [model]
 
-    def get_run_results_models(self):
+    # TODO: add pydantic to valdidate a subset of the run_results.json schema, example; /Users/sung/Desktop/data-diff/data_diff_demo/pydantic_example.py
+    # TODO: raise an exception that `run_results.json` is malformed based on the pydantic validation
+    def get_run_results_models(self) -> List[str]:
         with open(self.project_dir / RUN_RESULTS_PATH) as run_results:
             logger.info(f"Parsing file {RUN_RESULTS_PATH}")
             run_results_dict = json.load(run_results)
-            run_results_obj = parse_run_results(run_results=run_results_dict)
+
+        run_results_validated = RunResultsJsonConfig.parse_obj(run_results_dict)
 
-        dbt_version = parse_version(run_results_obj.metadata.dbt_version)
+        dbt_version = parse_version(run_results_validated.metadata.dbt_version)
 
         if dbt_version < parse_version(LOWER_DBT_V):
             raise DataDiffDbtRunResultsVersionError(
@@ -247,8 +262,10 @@ def get_run_results_models(self):
                 f"{dbt_version} is a recent version of dbt and may not be fully tested with data-diff! \nPlease report any issues to https://github.com/datafold/data-diff/issues"
             )
 
-        success_models = [x.unique_id for x in run_results_obj.results if x.status.name == "success"]
+        success_models = [x.unique_id for x in run_results_validated.results if x.status == "success"]
+
         models = [self.dev_manifest_obj.nodes.get(x) for x in success_models]
+        print(type(models[0])) # TODO this prints a class object type, I'll need to understand what other attributes are accessed before assuming getting a list of strings is enough
         if not models:
             raise DataDiffDbtNoSuccessfulModelsInRunError(
                 "Expected > 0 successful models runs from the last dbt command."
@@ -260,7 +277,7 @@ def get_manifest_obj(self, path: Path):
         with open(path) as manifest:
             logger.info(f"Parsing file {path}")
             manifest_dict = json.load(manifest)
-            manifest_obj = parse_manifest(manifest=manifest_dict)
+            manifest_obj = parse_manifest(manifest=manifest_dict)  # TODO: replace this
         return manifest_obj
 
     def get_project_dict(self):
@@ -447,9 +464,9 @@ def get_pk_from_model(self, node, unique_columns: dict, pk_tag: str) -> List[str
         return []
 
     def get_unique_columns(self) -> Dict[str, Set[str]]:
-        manifest = self.dev_manifest_obj
+        manifest = self.dev_manifest_obj #TODO: need to refactor this for dictionary calls
         cols_by_uid = defaultdict(set)
-        for node in manifest.nodes.values():
+        for node in manifest.nodes.values():#TODO: example: manifest["nodes"].values()
             try:
                 if not (node.resource_type.value == "test" and hasattr(node, "test_metadata")):
                     continue