From 80c360e089d6e75b3a0d2e4bda85a7758d22d6f6 Mon Sep 17 00:00:00 2001
From: Juan Carlos Jose Camacho <jc@dataherald.com>
Date: Mon, 22 Apr 2024 13:56:25 -0600
Subject: [PATCH] Add schema not support error exception

---
 .../services/database_connection.py             | 17 +++++++++++++++++
 dataherald/utils/error_codes.py                 |  1 +
 2 files changed, 18 insertions(+)

diff --git a/dataherald/sql_database/services/database_connection.py b/dataherald/sql_database/services/database_connection.py
index 9f64ac04..0d7bf052 100644
--- 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
index 6680feb5..3b508f72 100644
--- 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",
 }