Merge pull request #589 from dlawin/DX-724

Add custom exceptions for raises in dbt_parser

Author: dlawin

data_diff/dbt_parser.py

@@ -10,6 +10,21 @@
 from dbt_artifacts_parser.parser import parse_run_results, parse_manifest
 from dbt.config.renderer import ProfileRenderer
 
+from data_diff.errors import (
+    DataDiffDbtBigQueryOauthOnlyError,
+    DataDiffDbtConnectionNotImplementedError,
+    DataDiffDbtCoreNoRunnerError,
+    DataDiffDbtNoSuccessfulModelsInRunError,
+    DataDiffDbtProfileNotFoundError,
+    DataDiffDbtProjectVarsNotFoundError,
+    DataDiffDbtRedshiftPasswordOnlyError,
+    DataDiffDbtRunResultsVersionError,
+    DataDiffDbtSelectNoMatchingModelsError,
+    DataDiffDbtSelectUnexpectedError,
+    DataDiffDbtSelectVersionTooLowError,
+    DataDiffDbtSnowflakeSetConnectionError,
+)
+
 from .utils import getLogger, get_from_dict_with_raise
 
 
@@ -91,9 +106,11 @@ def __init__(self, profiles_dir_override: str, project_dir_override: str) -> Non
 
     def get_datadiff_variables(self) -> dict:
         doc_url = "https://docs.datafold.com/development_testing/open_source#configure-your-dbt-project"
-        error_message = f"vars: data_diff: section not found in dbt_project.yml.\n\nTo solve this, please configure your dbt project: \n{doc_url}\n"
-        vars = get_from_dict_with_raise(self.project_dict, "vars", error_message)
-        return get_from_dict_with_raise(vars, "data_diff", error_message)
+        exception = DataDiffDbtProjectVarsNotFoundError(
+            f"vars: data_diff: section not found in dbt_project.yml.\n\nTo solve this, please configure your dbt project: \n{doc_url}\n"
+        )
+        vars_dict = get_from_dict_with_raise(self.project_dict, "vars", exception)
+        return get_from_dict_with_raise(vars_dict, "data_diff", exception)
 
     def get_datadiff_model_config(self, model_meta: dict) -> TDatadiffModelConfig:
         where_filter = None
@@ -118,11 +135,11 @@ def get_models(self, dbt_selection: Optional[str] = None):
                     return self.get_dbt_selection_models(dbt_selection)
                 # edge case if running data-diff from a separate env than dbt (likely local development)
                 else:
-                    raise Exception(
+                    raise DataDiffDbtCoreNoRunnerError(
                         "data-diff is using a dbt-core version < 1.5, update the environment's dbt-core version via pip install 'dbt-core>=1.5' in order to use `--select`"
                     )
             else:
-                raise Exception(
+                raise DataDiffDbtSelectVersionTooLowError(
                     f"The `--select` feature requires dbt >= 1.5, but your project's manifest.json is from dbt v{dbt_version}. Please follow these steps to use the `--select` feature: \n 1. Update your dbt-core version via pip install 'dbt-core>=1.5'. Details: https://docs.getdbt.com/docs/core/pip-install#change-dbt-core-versions \n 2. Execute any `dbt` command (`run`, `compile`, `build`) to create a new manifest.json."
                 )
         else:
@@ -152,15 +169,17 @@ def get_dbt_selection_models(self, dbt_selection: str) -> List[str]:
         )
         if results.exception:
             raise results.exception
-        elif results.success and results.result:
+
+        if results.success and results.result:
             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
-        elif not results.result:
-            raise Exception(f"No dbt models found for `--select {dbt_selection}`")
-        else:
-            logger.debug(str(results))
-            raise Exception("Encountered an unexpected error while finding `--select` models")
+
+        if not results.result:
+            raise DataDiffDbtSelectNoMatchingModelsError(f"No dbt models found for `--select {dbt_selection}`")
+
+        logger.debug(str(results))
+        raise DataDiffDbtSelectUnexpectedError("Encountered an unexpected error while finding `--select` models")
 
     def get_run_results_models(self):
         with open(self.project_dir / RUN_RESULTS_PATH) as run_results:
@@ -174,16 +193,18 @@ def get_run_results_models(self):
             self.profiles_dir = legacy_profiles_dir()
 
         if dbt_version < parse_version(LOWER_DBT_V):
-            raise Exception(f"Found dbt: v{dbt_version} Expected the dbt project's version to be >= {LOWER_DBT_V}")
-        elif dbt_version >= parse_version(UPPER_DBT_V):
+            raise DataDiffDbtRunResultsVersionError(
+                f"Found dbt: v{dbt_version} Expected the dbt project's version to be >= {LOWER_DBT_V}"
+            )
+        if dbt_version >= parse_version(UPPER_DBT_V):
             logger.warning(
                 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"]
         models = [self.manifest_obj.nodes.get(x) for x in success_models]
         if not models:
-            raise ValueError("Expected > 0 successful models runs from the last dbt command.")
+            raise DataDiffDbtNoSuccessfulModelsInRunError("Expected > 0 successful models runs from the last dbt command.")
 
         return models
 
@@ -209,25 +230,35 @@ def get_connection_creds(self) -> Tuple[Dict[str, str], str]:
         dbt_profile_var = self.project_dict.get("profile")
 
         profile = get_from_dict_with_raise(
-            profiles, dbt_profile_var, f"No profile '{dbt_profile_var}' found in '{profiles_path}'."
+            profiles,
+            dbt_profile_var,
+            DataDiffDbtProfileNotFoundError(f"No profile '{dbt_profile_var}' found in '{profiles_path}'."),
         )
         # values can contain env_vars
         rendered_profile = ProfileRenderer().render_data(profile)
         profile_target = get_from_dict_with_raise(
-            rendered_profile, "target", f"No target found in profile '{dbt_profile_var}' in '{profiles_path}'."
+            rendered_profile,
+            "target",
+            DataDiffDbtProfileNotFoundError(f"No target found in profile '{dbt_profile_var}' in '{profiles_path}'."),
         )
         outputs = get_from_dict_with_raise(
-            rendered_profile, "outputs", f"No outputs found in profile '{dbt_profile_var}' in '{profiles_path}'."
+            rendered_profile,
+            "outputs",
+            DataDiffDbtProfileNotFoundError(f"No outputs found in profile '{dbt_profile_var}' in '{profiles_path}'."),
         )
         credentials = get_from_dict_with_raise(
             outputs,
             profile_target,
-            f"No credentials found for target '{profile_target}' in profile '{dbt_profile_var}' in '{profiles_path}'.",
+            DataDiffDbtProfileNotFoundError(
+                f"No credentials found for target '{profile_target}' in profile '{dbt_profile_var}' in '{profiles_path}'."
+            ),
         )
         conn_type = get_from_dict_with_raise(
             credentials,
             "type",
-            f"No type found for target '{profile_target}' in profile '{dbt_profile_var}' in '{profiles_path}'.",
+            DataDiffDbtProfileNotFoundError(
+                f"No type found for target '{profile_target}' in profile '{dbt_profile_var}' in '{profiles_path}'."
+            ),
         )
         conn_type = conn_type.lower()
 
@@ -253,7 +284,7 @@ def set_connection(self):
 
             if credentials.get("private_key_path") is not None:
                 if credentials.get("password") is not None:
-                    raise Exception("Cannot use password and key at the same time")
+                    raise DataDiffDbtSnowflakeSetConnectionError("Cannot use password and key at the same time")
                 conn_info["key"] = credentials.get("private_key_path")
                 conn_info["private_key_passphrase"] = credentials.get("private_key_passphrase")
             elif credentials.get("authenticator") is not None:
@@ -262,13 +293,13 @@ def set_connection(self):
             elif credentials.get("password") is not None:
                 conn_info["password"] = credentials.get("password")
             else:
-                raise Exception("Snowflake: unsupported auth method")
+                raise DataDiffDbtSnowflakeSetConnectionError("Snowflake: unsupported auth method")
         elif conn_type == "bigquery":
             method = credentials.get("method")
             # there are many connection types https://docs.getdbt.com/reference/warehouse-setups/bigquery-setup#oauth-via-gcloud
             # this assumes that the user is auth'd via `gcloud auth application-default login`
             if method is None or method != "oauth":
-                raise Exception("Oauth is the current method supported for Big Query.")
+                raise DataDiffDbtBigQueryOauthOnlyError("Oauth is the current method supported for Big Query.")
             conn_info = {
                 "driver": conn_type,
                 "project": credentials.get("project"),
@@ -284,7 +315,7 @@ def set_connection(self):
             if (credentials.get("pass") is None and credentials.get("password") is None) or credentials.get(
                 "method"
             ) == "iam":
-                raise Exception("Only password authentication is currently supported for Redshift.")
+                raise DataDiffDbtRedshiftPasswordOnlyError("Only password authentication is currently supported for Redshift.")
             conn_info = {
                 "driver": conn_type,
                 "host": credentials.get("host"),
@@ -315,7 +346,7 @@ def set_connection(self):
             }
             self.threads = credentials.get("threads")
         else:
-            raise NotImplementedError(f"Provider {conn_type} is not yet supported for dbt diffs")
+            raise DataDiffDbtConnectionNotImplementedError(f"Provider {conn_type} is not yet supported for dbt diffs")
 
         self.connection = conn_info
 

data_diff/errors.py

@@ -0,0 +1,46 @@
+class DataDiffDbtProjectVarsNotFoundError(Exception):
+    "Raised when an expected dbt_project.yml section is missing."
+
+
+class DataDiffDbtProfileNotFoundError(Exception):
+    "Raised when an expected profiles.yml section is missing."
+
+
+class DataDiffDbtNoSuccessfulModelsInRunError(Exception):
+    "Raised when there are no successful model runs in the run_results.json"
+
+
+class DataDiffDbtRunResultsVersionError(Exception):
+    "Raised when the dbt version in run_results.json is lower than the minimum version."
+
+
+class DataDiffDbtSelectNoMatchingModelsError(Exception):
+    "Raised when the `--select` flag returns no models."
+
+
+class DataDiffDbtSelectUnexpectedError(Exception):
+    "Catch all for unexpected dbt list --select results."
+
+
+class DataDiffDbtSnowflakeSetConnectionError(Exception):
+    "Raised when a dbt snowflake profile has unexpected values."
+
+
+class DataDiffDbtBigQueryOauthOnlyError(Exception):
+    "Raised when trying to use a method other than oauth with BigQuery."
+
+
+class DataDiffDbtRedshiftPasswordOnlyError(Exception):
+    "Raised when using a non-password connection method with Redshift."
+
+
+class DataDiffDbtConnectionNotImplementedError(Exception):
+    "Raised when trying to use an unsupported dbt connection method."
+
+
+class DataDiffDbtCoreNoRunnerError(Exception):
+    "Raised when the manifest version >= 1.5, but the dbt-core package is < 1.5. This is an edge case most likely to occur in development."
+
+
+class DataDiffDbtSelectVersionTooLowError(Exception):
+    "Raised when attempting to use `--select` with a dbt-core version < 1.5."

data_diff/utils.py

@@ -84,10 +84,12 @@ def truncate_error(error: str):
     return re.sub("'(.*?)'", "'***'", first_line)
 
 
-def get_from_dict_with_raise(dictionary: Dict, key: str, error_message: str):
+def get_from_dict_with_raise(dictionary: Dict, key: str, exception: Exception):
+    if dictionary is None:
+        raise exception
     result = dictionary.get(key)
     if result is None:
-        raise ValueError(error_message)
+        raise exception
     return result