From 5e7a5ded028826027b86b2f9056ca0622a5bbf55 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Juan=20Carlos=20Jos=C3=A9=20Camacho?= <jcjc712@hotmail.com>
Date: Fri, 1 Mar 2024 10:46:01 -0600
Subject: [PATCH] [DH-5517] Create a CustomError to send a description (#416)

---
 dataherald/sql_database/base.py | 13 +++++++------
 dataherald/utils/error_codes.py |  9 +++++++++
 2 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/dataherald/sql_database/base.py b/dataherald/sql_database/base.py
index b292a93b..baf20a9b 100644
--- a/dataherald/sql_database/base.py
+++ b/dataherald/sql_database/base.py
@@ -12,25 +12,26 @@
 
 from dataherald.sql_database.models.types import DatabaseConnection
 from dataherald.utils.encrypt import FernetEncrypt
+from dataherald.utils.error_codes import CustomError
 from dataherald.utils.s3 import S3
 
 logger = logging.getLogger(__name__)
 
 
 # Define a custom exception class
-class SQLInjectionError(Exception):
+class SQLInjectionError(CustomError):
     pass
 
 
-class InvalidDBConnectionError(Exception):
+class InvalidDBConnectionError(CustomError):
     pass
 
 
-class EmptyDBError(Exception):
+class EmptyDBError(CustomError):
     pass
 
 
-class SSHInvalidDatabaseConnectionError(Exception):
+class SSHInvalidDatabaseConnectionError(CustomError):
     pass
 
 
@@ -89,7 +90,7 @@ def get_sql_engine(
                 return engine
         except Exception as e:
             raise SSHInvalidDatabaseConnectionError(
-                f"Invalid SSH connection, {e}"
+                "Invalid SSH connection", description=str(e)
             ) from e
         try:
             db_uri = unquote(fernet_encrypt.decrypt(database_info.connection_uri))
@@ -107,7 +108,7 @@ def get_sql_engine(
             DBConnections.add(database_info.id, engine)
         except Exception as e:
             raise InvalidDBConnectionError(  # noqa: B904
-                f"Unable to connect to db: {database_info.alias}, {e}"
+                f"Unable to connect to db: {database_info.alias}", description=str(e)
             )
         return engine
 
diff --git a/dataherald/utils/error_codes.py b/dataherald/utils/error_codes.py
index 118765dd..1752b064 100644
--- a/dataherald/utils/error_codes.py
+++ b/dataherald/utils/error_codes.py
@@ -18,6 +18,12 @@
 }
 
 
+class CustomError(Exception):
+    def __init__(self, message, description=None):
+        super().__init__(message)
+        self.description = description
+
+
 def error_response(error, detail: dict, default_error_code=""):
     return JSONResponse(
         status_code=400,
@@ -26,6 +32,9 @@ def error_response(error, detail: dict, default_error_code=""):
                 error.__class__.__name__, default_error_code
             ),
             "message": str(error),
+            "description": error.description
+            if isinstance(error, CustomError)
+            else None,
             "detail": detail,
         },
     )