Add schema not support error exception

Dataherald · Apr 22, 2024 · 80c360e · 80c360e
1 parent 6d2094d
commit 80c360e
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 0 deletions.
diff --git a/dataherald/sql_database/services/database_connection.py b/dataherald/sql_database/services/database_connection.py
@@ -10,6 +10,11 @@
 from dataherald.sql_database.models.types import DatabaseConnection
 from dataherald.types import DatabaseConnectionRequest
 from dataherald.utils.encrypt import FernetEncrypt
+from dataherald.utils.error_codes import CustomError
+
+
+class SchemaNotSupportedError(CustomError):
+    pass
 
 
 class DatabaseConnectionService:
@@ -98,6 +103,18 @@ def create(
             file_storage=database_connection_request.file_storage,
             metadata=database_connection_request.metadata,
         )
+        if database_connection.schemas and database_connection.dialect in [
+            "redshift",
+            "awsathena",
+            "mssql",
+            "mysql",
+            "clickhouse",
+            "duckdb",
+        ]:
+            raise SchemaNotSupportedError(
+                "Schema not supported for this db",
+                description=f"The {database_connection.dialect} dialect doesn't support schemas",
+            )
         if not database_connection.schemas:
             database_connection.schemas = self.get_current_schema(database_connection)
 

diff --git a/dataherald/utils/error_codes.py b/dataherald/utils/error_codes.py
@@ -19,6 +19,7 @@
     "SQLGenerationNotFoundError": "sql_generation_not_found",
     "NLGenerationError": "nl_generation_not_created",
     "MalformedGoldenSQLError": "invalid_golden_sql",
+    "SchemaNotSupportedError": "schema_not_supported",
 }