diff --git a/unstract/connectors/src/unstract/connectors/databases/exceptions.py b/unstract/connectors/src/unstract/connectors/databases/exceptions.py index 5aaf37843d..b1a9b30853 100644 --- a/unstract/connectors/src/unstract/connectors/databases/exceptions.py +++ b/unstract/connectors/src/unstract/connectors/databases/exceptions.py @@ -6,12 +6,7 @@ class UnstractDBConnectorException(ConnectorBaseException): """Base class for database-related exceptions from Unstract connectors.""" - def __init__( - self, - detail: Any, - *args: Any, - **kwargs: Any, - ) -> None: + def __init__(self, detail: Any, *args: Any, **kwargs: Any) -> None: default_detail = "Error creating/inserting to database. " user_message = default_detail if not detail else detail super().__init__(*args, user_message=user_message, **kwargs) @@ -24,16 +19,15 @@ def __init__(self, detail: Any, database: Any) -> None: f"Error creating/writing to `{database}`. Syntax incorrect. " f"Please check your table-name or schema. " ) - final_detail = ( - f"{default_detail}\nDetails: {detail}" if detail else default_detail - ) + final_detail = _format_exception_detail(default_detail, detail) super().__init__(detail=final_detail) class InvalidSchemaException(UnstractDBConnectorException): def __init__(self, detail: Any, database: str) -> None: default_detail = f"Error creating/writing to {database}. Schema not valid. " - super().__init__(detail=default_detail) + final_detail = _format_exception_detail(default_detail, detail) + super().__init__(detail=final_detail) class UnderfinedTableException(UnstractDBConnectorException): @@ -42,7 +36,8 @@ def __init__(self, detail: Any, database: str) -> None: f"Error creating/writing to {database}. Undefined table. " f"Please check your table-name or schema. " ) - super().__init__(detail=default_detail) + final_detail = _format_exception_detail(default_detail, detail) + super().__init__(detail=final_detail) class ValueTooLongException(UnstractDBConnectorException): @@ -51,7 +46,8 @@ def __init__(self, detail: Any, database: str) -> None: f"Error creating/writing to {database}. " f"Size of the inserted data exceeds the limit provided by the database. " ) - super().__init__(detail=default_detail) + final_detail = _format_exception_detail(default_detail, detail) + super().__init__(detail=final_detail) class FeatureNotSupportedException(UnstractDBConnectorException): @@ -59,7 +55,8 @@ def __init__(self, detail: Any, database: str) -> None: default_detail = ( f"Error creating/writing to {database}. Feature not supported sql error. " ) - super().__init__(detail=default_detail) + final_detail = _format_exception_detail(default_detail, detail) + super().__init__(detail=final_detail) class SnowflakeProgrammingException(UnstractDBConnectorException): @@ -69,7 +66,8 @@ def __init__(self, detail: Any, database: str, table_name: str, schema: str) -> f"Please make sure all the columns exist in your table as per destination " f"DB configuration \n and snowflake credentials are correct.\n" ) - super().__init__(default_detail) + final_detail = _format_exception_detail(default_detail, detail) + super().__init__(detail=final_detail) class BigQueryForbiddenException(UnstractDBConnectorException): @@ -78,7 +76,8 @@ def __init__(self, detail: Any, table_name: str) -> None: f"Error creating/writing to {table_name}. " f"Access forbidden in bigquery. Please check your permissions. " ) - super().__init__(detail=default_detail) + final_detail = _format_exception_detail(default_detail, detail) + super().__init__(detail=final_detail) class BigQueryNotFoundException(UnstractDBConnectorException): @@ -87,7 +86,8 @@ def __init__(self, detail: str, table_name: str) -> None: f"Error creating/writing to {table_name}. " f"The requested resource was not found. " ) - super().__init__(detail=default_detail) + final_detail = _format_exception_detail(default_detail, detail) + super().__init__(detail=final_detail) class ColumnMissingException(UnstractDBConnectorException): @@ -104,10 +104,26 @@ def __init__( f"Please make sure all the columns exist in your table " f"as per the destination DB configuration.\n" ) - super().__init__(detail=default_detail) + final_detail = _format_exception_detail(default_detail, detail) + super().__init__(detail=final_detail) class OperationalException(UnstractDBConnectorException): def __init__(self, detail: Any, database: str) -> None: default_detail = f"Error creating/writing to {database}. Operational error. " - super().__init__(detail=default_detail) + final_detail = _format_exception_detail(default_detail, detail) + super().__init__(detail=final_detail) + + +def _format_exception_detail(base_error: str, library_error: Any) -> str: + """Format exception detail by combining base error with library-specific details. + + Args: + base_error: The base error message describing the error type + library_error: The actual error detail from the database library + + Returns: + Formatted error message combining both if library_error exists, + otherwise just the base_error + """ + return f"{base_error}\nDetails: {library_error}" if library_error else base_error diff --git a/unstract/connectors/tests/databases/test_bigquery_db.py b/unstract/connectors/tests/databases/test_bigquery_db.py index 5673919beb..4518fa9a7e 100644 --- a/unstract/connectors/tests/databases/test_bigquery_db.py +++ b/unstract/connectors/tests/databases/test_bigquery_db.py @@ -1,35 +1,135 @@ import unittest +from unittest.mock import MagicMock, patch + +import google.api_core.exceptions from unstract.connectors.databases.bigquery.bigquery import BigQuery +from unstract.connectors.databases.exceptions import ( + BigQueryForbiddenException, + BigQueryNotFoundException, +) class TestBigQuery(unittest.TestCase): - def test_json_credentials(self): - bigquery = BigQuery( + + def setUp(self): + """Set up test fixtures that are common across all tests.""" + self.bigquery = BigQuery( { "json_credentials": ( - '{"type":"service_account","project_id":' - '"project_id","private_key_id":"private_key_id",' - '"private_key":"private_key","client_email":' - '"client_email","client_id":"11427061",' - '"auth_uri":"",' - '"token_uri":"",' - '"auth_provider_x509_cert_url":"",' - '"client_x509_cert_url":"",' - '"universe_domain":"googleapis.com"}' + '{"type":"service_account","project_id":"test_project"}' ) } ) - client = bigquery.get_engine() - query_job = client.query( - """ - select * from `dataset.test`""" + + def _execute_query_with_mock_error(self, mock_error, expected_exception): + """Helper method to execute query with a mocked error. + + Args: + mock_error: The Google API exception to raise + expected_exception: The exception class expected to be raised + + Returns: + The exception context manager from assertRaises + """ + # Mock the engine and query job + mock_engine = MagicMock() + mock_query_job = MagicMock() + mock_engine.query.return_value = mock_query_job + mock_query_job.result.side_effect = mock_error + + # Mock get_information_schema to return empty dict + with patch.object(self.bigquery, "get_information_schema", return_value={}): + with self.assertRaises(expected_exception) as context: + self.bigquery.execute_query( + engine=mock_engine, + sql_query="INSERT INTO test.dataset.table VALUES (@col)", + table_name="test.dataset.table", + sql_values={"col": "value"}, + sql_keys=["col"], + ) + + return context + + def test_execute_query_forbidden_billing(self): + """Test that BigQueryForbiddenException includes actual billing error details.""" + # Create a mock Forbidden exception with billing error message + billing_error_msg = ( + "403 Billing has not been enabled for this project. " + "Enable billing at https://console.cloud.google.com/billing" + ) + mock_error = google.api_core.exceptions.Forbidden(billing_error_msg) + mock_error.message = billing_error_msg + + # Execute query with mock error + context = self._execute_query_with_mock_error( + mock_error, BigQueryForbiddenException + ) + + # Verify the exception message includes both default text and actual error details + error_msg = str(context.exception.detail) + self.assertIn("Access forbidden in bigquery", error_msg) + self.assertIn("Please check your permissions", error_msg) + self.assertIn("Details:", error_msg) + self.assertIn("403 Billing has not been enabled", error_msg) + self.assertIn("test.dataset.table", error_msg) + + def test_execute_query_forbidden_permission(self): + """Test that BigQueryForbiddenException includes actual permission error details.""" + # Create a mock Forbidden exception with permission error message + permission_error_msg = ( + "403 User does not have permission to access table test.dataset.table" + ) + mock_error = google.api_core.exceptions.Forbidden(permission_error_msg) + mock_error.message = permission_error_msg + + # Execute query with mock error + context = self._execute_query_with_mock_error( + mock_error, BigQueryForbiddenException + ) + + # Verify the exception message includes both default text and actual error details + error_msg = str(context.exception.detail) + self.assertIn("Access forbidden in bigquery", error_msg) + self.assertIn("Details:", error_msg) + self.assertIn("User does not have permission", error_msg) + + def test_execute_query_not_found(self): + """Test that BigQueryNotFoundException includes actual resource not found details.""" + # Create a mock NotFound exception + not_found_error_msg = "404 Dataset 'test:dataset' not found" + mock_error = google.api_core.exceptions.NotFound(not_found_error_msg) + mock_error.message = not_found_error_msg + + # Execute query with mock error + context = self._execute_query_with_mock_error( + mock_error, BigQueryNotFoundException + ) + + # Verify the exception message includes both default text and actual error details + error_msg = str(context.exception.detail) + self.assertIn("The requested resource was not found", error_msg) + self.assertIn("Details:", error_msg) + self.assertIn("404 Dataset", error_msg) + self.assertIn("test.dataset.table", error_msg) + + def test_exception_empty_detail(self): + """Test that exceptions handle empty detail gracefully.""" + # Create a mock Forbidden exception with empty message + mock_error = google.api_core.exceptions.Forbidden("") + mock_error.message = "" + + # Execute query with mock error + context = self._execute_query_with_mock_error( + mock_error, BigQueryForbiddenException ) - results = query_job.result() - for c in results: - print(c) - self.assertTrue(len(results) > 0) # add assertion here + # Verify the exception message includes default text but not empty "Details:" + error_msg = str(context.exception.detail) + self.assertIn("Access forbidden in bigquery", error_msg) + self.assertIn("Please check your permissions", error_msg) + # When detail is empty, should not have "Details:" section + self.assertNotIn("Details:", error_msg) if __name__ == "__main__":