1212
1313from dataherald .sql_database .models .types import DatabaseConnection
1414from dataherald .utils .encrypt import FernetEncrypt
15+ from dataherald .utils .error_codes import CustomError
1516from dataherald .utils .s3 import S3
1617
1718logger = logging .getLogger (__name__ )
1819
1920
2021# Define a custom exception class
21- class SQLInjectionError (Exception ):
22+ class SQLInjectionError (CustomError ):
2223 pass
2324
2425
25- class InvalidDBConnectionError (Exception ):
26+ class InvalidDBConnectionError (CustomError ):
2627 pass
2728
2829
29- class EmptyDBError (Exception ):
30+ class EmptyDBError (CustomError ):
3031 pass
3132
3233
33- class SSHInvalidDatabaseConnectionError (Exception ):
34+ class SSHInvalidDatabaseConnectionError (CustomError ):
3435 pass
3536
3637
@@ -89,7 +90,7 @@ def get_sql_engine(
8990 return engine
9091 except Exception as e :
9192 raise SSHInvalidDatabaseConnectionError (
92- f "Invalid SSH connection, { e } "
93+ "Invalid SSH connection" , description = str ( e )
9394 ) from e
9495 try :
9596 db_uri = unquote (fernet_encrypt .decrypt (database_info .connection_uri ))
@@ -107,7 +108,7 @@ def get_sql_engine(
107108 DBConnections .add (database_info .id , engine )
108109 except Exception as e :
109110 raise InvalidDBConnectionError ( # noqa: B904
110- f"Unable to connect to db: { database_info .alias } , { e } "
111+ f"Unable to connect to db: { database_info .alias } " , description = str ( e )
111112 )
112113 return engine
113114
0 commit comments