Map engine specific deals

Author: BenConstable9

text_2_sql/text_2_sql_core/src/text_2_sql_core/connectors/ai_search.py

@@ -13,6 +13,8 @@
 from typing import Annotated
 from text_2_sql_core.connectors.open_ai import OpenAIConnector
 
+from text_2_sql_core.utils.database import DatabaseEngineSpecificFields
+
 
 class AISearchConnector:
     def __init__(self):
@@ -167,7 +169,7 @@ async def get_entity_schemas(
             "The entities to exclude from the search results. Pass the entity property of entities (e.g. 'SalesLT.Address') you already have the schemas for to avoid getting repeated entities.",
         ] = [],
         engine_specific_fields: Annotated[
-            list[str],
+            list[DatabaseEngineSpecificFields],
             "The fields specific to the engine to be included in the search results.",
         ] = [],
     ) -> str:
@@ -192,7 +194,7 @@ async def get_entity_schemas(
             "Columns",
             "EntityRelationships",
             "CompleteEntityRelationshipsGraph",
-        ] + engine_specific_fields
+        ] + list(map(str, engine_specific_fields))
 
         schemas = await self.run_ai_search_query(
             text,

text_2_sql/text_2_sql_core/src/text_2_sql_core/connectors/databricks_sql.py

@@ -135,7 +135,7 @@ async def get_entity_schemas(
         """
 
         schemas = await self.ai_search_connector.get_entity_schemas(
-            text, excluded_entities, engine_specific_fields=["Catalog"]
+            text, excluded_entities, engine_specific_fields=self.engine_specific_fields
         )
 
         for schema in schemas:

text_2_sql/text_2_sql_core/src/text_2_sql_core/connectors/postgresql_sql.py

@@ -111,14 +111,15 @@ async def get_entity_schemas(
         """
 
         schemas = await self.ai_search_connector.get_entity_schemas(
-            text, excluded_entities
+            text, excluded_entities, engine_specific_fields=self.engine_specific_fields
         )
 
         for schema in schemas:
             schema["SelectFromEntity"] = ".".join([schema["Schema"], schema["Entity"]])
 
             del schema["Entity"]
             del schema["Schema"]
+            del schema["Database"]
 
         if as_json:
             return json.dumps(schemas, default=str)

text_2_sql/text_2_sql_core/src/text_2_sql_core/connectors/snowflake_sql.py

@@ -150,7 +150,7 @@ async def get_entity_schemas(
         """
 
         schemas = await self.ai_search_connector.get_entity_schemas(
-            text, excluded_entities, engine_specific_fields=["Warehouse", "Database"]
+            text, excluded_entities, engine_specific_fields=self.engine_specific_fields
         )
 
         for schema in schemas:

text_2_sql/text_2_sql_core/src/text_2_sql_core/connectors/tsql_sql.py

@@ -124,7 +124,7 @@ async def get_entity_schemas(
         """
 
         schemas = await self.ai_search_connector.get_entity_schemas(
-            text, excluded_entities
+            text, excluded_entities, engine_specific_fields=self.engine_specific_fields
         )
 
         for schema in schemas: