custom errors for dbt_parser

Author: dlawin

data_diff/dbt_parser.py

@@ -11,6 +11,8 @@
 from dbt_artifacts_parser.parser import parse_run_results, parse_manifest
 from dbt.config.renderer import ProfileRenderer
 
+from data_diff.errors import DbtBigQueryOauthOnlyError, DbtConnectionNotImplementedError, DbtCoreNoRunnerError, DbtNoSuccessfulModelsInRunError, DbtProfileNotFoundError, DbtProjectVarsNotFoundError, DbtRedshiftPasswordOnlyError, DbtRunResultsVersionError, DbtSelectNoMatchingModelsError, DbtSelectUnexpectedError, DbtSelectVersionTooLowError, DbtSnowflakeSetConnectionError
+
 from .utils import getLogger, get_from_dict_with_raise
 from .version import __version__
 
@@ -93,9 +95,9 @@ 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 = DbtProjectVarsNotFoundError(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
@@ -120,11 +122,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 DbtCoreNoRunnerError(
                         "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 DbtSelectVersionTooLowError(
                     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:
@@ -154,15 +156,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 DbtSelectNoMatchingModelsError(f"No dbt models found for `--select {dbt_selection}`")
+
+        logger.debug(str(results))
+        raise DbtSelectUnexpectedError("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:
@@ -176,16 +180,16 @@ 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 DbtRunResultsVersionError(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 DbtNoSuccessfulModelsInRunError("Expected > 0 successful models runs from the last dbt command.")
 
         print(f"Running with data-diff={__version__}\n")
         return models
@@ -212,25 +216,25 @@ 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, DbtProfileNotFoundError(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", DbtProfileNotFoundError(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", DbtProfileNotFoundError(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}'.",
+            DbtProfileNotFoundError(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}'.",
+            DbtProfileNotFoundError(f"No type found for target '{profile_target}' in profile '{dbt_profile_var}' in '{profiles_path}'."),
         )
         conn_type = conn_type.lower()
 
@@ -256,7 +260,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 DbtSnowflakeSetConnectionError("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:
@@ -265,13 +269,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 DbtSnowflakeSetConnectionError("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 DbtBigQueryOauthOnlyError("Oauth is the current method supported for Big Query.")
             conn_info = {
                 "driver": conn_type,
                 "project": credentials.get("project"),
@@ -287,7 +291,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 DbtRedshiftPasswordOnlyError("Only password authentication is currently supported for Redshift.")
             conn_info = {
                 "driver": conn_type,
                 "host": credentials.get("host"),
@@ -318,7 +322,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 DbtConnectionNotImplementedError(f"Provider {conn_type} is not yet supported for dbt diffs")
 
         self.connection = conn_info
 

data_diff/errors.py

@@ -0,0 +1,36 @@
+
+class DbtProjectVarsNotFoundError(Exception):
+    "Raised when an expected dbt_project.yml section is missing."
+
+class DbtProfileNotFoundError(Exception):
+    "Raised when an expected profiles.yml section is missing."
+
+class DbtNoSuccessfulModelsInRunError(Exception):
+    "Raised when there are no successful model runs in the run_results.json"
+
+class DbtRunResultsVersionError(Exception):
+    "Raised when the dbt version in run_results.json is lower than the minimum version."
+
+class DbtSelectNoMatchingModelsError(Exception):
+    "Raised when the `--select` flag returns no models."
+
+class DbtSelectUnexpectedError(Exception):
+    "Catch all for unexpected dbt list --select results."
+
+class DbtSnowflakeSetConnectionError(Exception):
+    "Raised when a dbt snowflake profile has unexpected values."
+
+class DbtBigQueryOauthOnlyError(Exception):
+    "Raised when trying to use a method other than oauth with BigQuery."
+
+class DbtRedshiftPasswordOnlyError(Exception):
+    "Raised when using a non-password connection method with Redshift."
+
+class DbtConnectionNotImplementedError(Exception):
+    "Raised when trying to use an unsupported dbt connection method."
+
+class DbtCoreNoRunnerError(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 DbtSelectVersionTooLowError(Exception):
+    "Raised when attempting to use `--select` with a dbt-core version < 1.5."

data_diff/utils.py

@@ -81,10 +81,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
 
 

tests/test_dbt.py

@@ -4,6 +4,7 @@
 from data_diff.cloud.datafold_api import TCloudApiDataSource
 from data_diff.cloud.datafold_api import TCloudApiOrgMeta
 from data_diff.diff_tables import Algorithm
+from data_diff.errors import DbtBigQueryOauthOnlyError, DbtConnectionNotImplementedError, DbtCoreNoRunnerError, DbtNoSuccessfulModelsInRunError, DbtProfileNotFoundError, DbtProjectVarsNotFoundError, DbtRedshiftPasswordOnlyError, DbtRunResultsVersionError, DbtSelectVersionTooLowError, DbtSnowflakeSetConnectionError
 from .test_cli import run_datadiff_cli
 
 from data_diff.dbt import (
@@ -40,7 +41,7 @@ def test_get_datadiff_variables_none(self):
         mock_self = Mock()
         mock_self.project_dict = none_dict
 
-        with self.assertRaises(Exception):
+        with self.assertRaises(DbtProjectVarsNotFoundError):
             DbtParser.get_datadiff_variables(mock_self)
 
     def test_get_datadiff_variables_empty(self):
@@ -49,7 +50,7 @@ def test_get_datadiff_variables_empty(self):
         mock_self = Mock()
         mock_self.project_dict = empty_dict
 
-        with self.assertRaises(Exception):
+        with self.assertRaises(DbtProjectVarsNotFoundError):
             DbtParser.get_datadiff_variables(mock_self)
 
     def test_get_models(self):
@@ -72,7 +73,7 @@ def test_get_models_unsupported_manifest_version(self):
         mock_return_value = Mock()
         mock_self.get_dbt_selection_models.return_value = mock_return_value
 
-        with self.assertRaises(Exception):
+        with self.assertRaises(DbtSelectVersionTooLowError):
             _ = DbtParser.get_models(mock_self, selection)
         mock_self.get_dbt_selection_models.assert_not_called()
 
@@ -85,7 +86,7 @@ def test_get_models_no_runner(self):
         mock_return_value = Mock()
         mock_self.get_dbt_selection_models.return_value = mock_return_value
 
-        with self.assertRaises(Exception):
+        with self.assertRaises(DbtCoreNoRunnerError):
             _ = DbtParser.get_models(mock_self, selection)
         mock_self.get_dbt_selection_models.assert_not_called()
 
@@ -135,7 +136,7 @@ def test_get_run_results_models_bad_lower_dbt_version(self, mock_open, mock_arti
         mock_artifact_parser.return_value = mock_run_results
         mock_run_results.metadata.dbt_version = "0.19.0"
 
-        with self.assertRaises(Exception) as ex:
+        with self.assertRaises(DbtRunResultsVersionError) as ex:
             DbtParser.get_run_results_models(mock_self)
 
         mock_open.assert_called_once_with(Path(RUN_RESULTS_PATH))
@@ -158,7 +159,7 @@ def test_get_run_results_models_no_success(self, mock_open, mock_artifact_parser
         mock_failed_result.status.name = "failed"
         mock_run_results.results = [mock_failed_result]
 
-        with self.assertRaises(Exception):
+        with self.assertRaises(DbtNoSuccessfulModelsInRunError):
             DbtParser.get_run_results_models(mock_self)
 
         mock_open.assert_any_call(Path(RUN_RESULTS_PATH))
@@ -235,7 +236,7 @@ def test_set_connection_snowflake_no_key_or_password(self):
         mock_self = Mock()
         mock_self.get_connection_creds.return_value = (expected_credentials, expected_driver)
 
-        with self.assertRaises(Exception):
+        with self.assertRaises(DbtSnowflakeSetConnectionError):
             DbtParser.set_connection(mock_self)
 
         self.assertNotIsInstance(mock_self.connection, dict)
@@ -259,7 +260,7 @@ def test_set_connection_snowflake_key_and_password(self):
         mock_self = Mock()
         mock_self.get_connection_creds.return_value = (expected_credentials, expected_driver)
 
-        with self.assertRaises(Exception):
+        with self.assertRaises(DbtSnowflakeSetConnectionError):
             DbtParser.set_connection(mock_self)
 
         self.assertNotIsInstance(mock_self.connection, dict)
@@ -291,7 +292,20 @@ def test_set_connection_bigquery_not_oauth(self):
 
         mock_self = Mock()
         mock_self.get_connection_creds.return_value = (expected_credentials, expected_driver)
-        with self.assertRaises(Exception):
+        with self.assertRaises(DbtBigQueryOauthOnlyError):
+            DbtParser.set_connection(mock_self)
+
+        self.assertNotIsInstance(mock_self.connection, dict)
+
+    def test_set_connection_redshift_not_password(self):
+        driver = "redshift"
+        credentials = {
+            "method": "not_password",
+        }
+
+        mock_self = Mock()
+        mock_self.get_connection_creds.return_value = (credentials, driver)
+        with self.assertRaises(DbtRedshiftPasswordOnlyError):
             DbtParser.set_connection(mock_self)
 
         self.assertNotIsInstance(mock_self.connection, dict)
@@ -301,7 +315,7 @@ def test_set_connection_not_implemented(self):
 
         mock_self = Mock()
         mock_self.get_connection_creds.return_value = (None, expected_driver)
-        with self.assertRaises(NotImplementedError):
+        with self.assertRaises(DbtConnectionNotImplementedError):
             DbtParser.set_connection(mock_self)
 
         self.assertNotIsInstance(mock_self.connection, dict)
@@ -340,7 +354,7 @@ def test_get_connection_no_matching_profile(self, mock_open, mock_profile_render
         mock_yaml.safe_load.return_value = profiles_dict
         profile = profiles_dict["a_profile"]
         mock_profile_renderer().render_data.return_value = profile
-        with self.assertRaises(ValueError):
+        with self.assertRaises(DbtProfileNotFoundError):
             _, _ = DbtParser.get_connection_creds(mock_self)
 
     @patch("data_diff.dbt_parser.yaml")
@@ -360,7 +374,7 @@ def test_get_connection_no_target(self, mock_open, mock_profile_renderer, mock_y
         mock_profile_renderer().render_data.return_value = profile
         mock_self.project_dict = {"profile": "a_profile"}
         mock_yaml.safe_load.return_value = profiles_dict
-        with self.assertRaises(ValueError):
+        with self.assertRaises(DbtProfileNotFoundError):
             _, _ = DbtParser.get_connection_creds(mock_self)
 
     profile_yaml_no_outputs = """
@@ -379,7 +393,7 @@ def test_get_connection_no_outputs(self, mock_open, mock_profile_renderer, mock_
         profile = profiles_dict["a_profile"]
         mock_profile_renderer().render_data.return_value = profile
         mock_yaml.safe_load.return_value = profiles_dict
-        with self.assertRaises(ValueError):
+        with self.assertRaises(DbtProfileNotFoundError):
             _, _ = DbtParser.get_connection_creds(mock_self)
 
     @patch("data_diff.dbt_parser.yaml")
@@ -398,7 +412,7 @@ def test_get_connection_no_credentials(self, mock_open, mock_profile_renderer, m
         mock_yaml.safe_load.return_value = profiles_dict
         profile = profiles_dict["a_profile"]
         mock_profile_renderer().render_data.return_value = profile
-        with self.assertRaises(ValueError):
+        with self.assertRaises(DbtProfileNotFoundError):
             _, _ = DbtParser.get_connection_creds(mock_self)
 
     @patch("data_diff.dbt_parser.yaml")
@@ -419,7 +433,7 @@ def test_get_connection_no_target_credentials(self, mock_open, mock_profile_rend
         profile = profiles_dict["a_profile"]
         mock_profile_renderer().render_data.return_value = profile
         mock_yaml.safe_load.return_value = profiles_dict
-        with self.assertRaises(ValueError):
+        with self.assertRaises(DbtProfileNotFoundError):
             _, _ = DbtParser.get_connection_creds(mock_self)
 
     @patch("data_diff.dbt_parser.yaml")
@@ -438,7 +452,7 @@ def test_get_connection_no_type(self, mock_open, mock_profile_renderer, mock_yam
         mock_yaml.safe_load.return_value = profiles_dict
         profile = profiles_dict["a_profile"]
         mock_profile_renderer().render_data.return_value = profile
-        with self.assertRaises(ValueError):
+        with self.assertRaises(DbtProfileNotFoundError):
             _, _ = DbtParser.get_connection_creds(mock_self)