Store excluded fields in

Author: BenConstable9

deploy_ai_search/src/deploy_ai_search/text_2_sql_schema_store.py

@@ -26,6 +26,7 @@
 )
 import os
 from text_2_sql_core.utils.database import DatabaseEngine
+from text_2_sql_core.connectors.factory import ConnectorFactory
 
 
 class Text2SqlSchemaStoreAISearch(AISearch):
@@ -49,31 +50,13 @@ def __init__(
             os.environ["Text2Sql__DatabaseEngine"].upper()
         ]
 
+        self.database_connector = ConnectorFactory.get_database_connector()
+
         if single_data_dictionary_file:
             self.parsing_mode = BlobIndexerParsingMode.JSON_ARRAY
         else:
             self.parsing_mode = BlobIndexerParsingMode.JSON
 
-    @property
-    def excluded_fields_for_database_engine(self):
-        """A method to get the excluded fields for the database engine."""
-
-        all_engine_specific_fields = ["Warehouse", "Database", "Catalog"]
-        if self.database_engine == DatabaseEngine.SNOWFLAKE:
-            engine_specific_fields = ["Warehouse", "Database"]
-        elif self.database_engine == DatabaseEngine.TSQL:
-            engine_specific_fields = ["Database"]
-        elif self.database_engine == DatabaseEngine.DATABRICKS:
-            engine_specific_fields = ["Catalog"]
-        elif self.database_engine == DatabaseEngine.POSTGRESQL:
-            engine_specific_fields = ["Database"]
-
-        return [
-            field
-            for field in all_engine_specific_fields
-            if field not in engine_specific_fields
-        ]
-
     def get_index_fields(self) -> list[SearchableField]:
         """This function returns the index fields for sql index.
 
@@ -198,7 +181,7 @@ def get_index_fields(self) -> list[SearchableField]:
         fields = [
             field
             for field in fields
-            if field.name not in self.excluded_fields_for_database_engine
+            if field.name not in self.database_connector.excluded_engine_specific_fields
         ]
 
         return fields

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

@@ -8,7 +8,7 @@
 import logging
 import json
 
-from text_2_sql_core.utils.database import DatabaseEngine
+from text_2_sql_core.utils.database import DatabaseEngine, DatabaseEngineSpecificFields
 
 
 class DatabricksSqlConnector(SqlConnector):
@@ -17,6 +17,11 @@ def __init__(self):
 
         self.database_engine = DatabaseEngine.DATABRICKS
 
+    @property
+    def engine_specific_fields(self) -> list[str]:
+        """Get the engine specific fields."""
+        return [DatabaseEngineSpecificFields.CATALOG]
+
     async def query_execution(
         self,
         sql_query: Annotated[

text_2_sql/text_2_sql_core/src/text_2_sql_core/connectors/factory.py

@@ -25,6 +25,12 @@ def get_database_connector():
                 from text_2_sql_core.connectors.tsql_sql import TSQLSqlConnector
 
                 return TSQLSqlConnector()
+            elif os.environ["Text2Sql__DatabaseEngine"].upper() == "POSTGRESQL":
+                from text_2_sql_core.connectors.postgresql_sql import (
+                    PostgresqlSqlConnector,
+                )
+
+                return PostgresqlSqlConnector()
             else:
                 raise ValueError(
                     f"""Database engine {

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

@@ -7,7 +7,7 @@
 import logging
 import json
 
-from text_2_sql_core.utils.database import DatabaseEngine
+from text_2_sql_core.utils.database import DatabaseEngine, DatabaseEngineSpecificFields
 
 
 class PostgresqlSqlConnector(SqlConnector):
@@ -16,6 +16,11 @@ def __init__(self):
 
         self.database_engine = DatabaseEngine.POSTGRESQL
 
+    @property
+    def engine_specific_fields(self) -> list[str]:
+        """Get the engine specific fields."""
+        return [DatabaseEngineSpecificFields.DATABASE]
+
     async def query_execution(
         self,
         sql_query: Annotated[str, "The SQL query to run against the database."],

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

@@ -8,7 +8,7 @@
 import logging
 import json
 
-from text_2_sql_core.utils.database import DatabaseEngine
+from text_2_sql_core.utils.database import DatabaseEngine, DatabaseEngineSpecificFields
 
 
 class SnowflakeSqlConnector(SqlConnector):
@@ -17,6 +17,14 @@ def __init__(self):
 
         self.database_engine = DatabaseEngine.SNOWFLAKE
 
+    @property
+    def engine_specific_fields(self) -> list[str]:
+        """Get the engine specific fields."""
+        return [
+            DatabaseEngineSpecificFields.WAREHOUSE,
+            DatabaseEngineSpecificFields.DATABASE,
+        ]
+
     async def query_execution(
         self,
         sql_query: Annotated[

text_2_sql/text_2_sql_core/src/text_2_sql_core/connectors/sql.py

@@ -9,6 +9,7 @@
 from abc import ABC, abstractmethod
 from jinja2 import Template
 import json
+from text_2_sql_core.utils.database import DatabaseEngineSpecificFields
 
 
 class SqlConnector(ABC):
@@ -29,6 +30,22 @@ def __init__(self):
 
         self.database_engine = None
 
+    @abstractmethod
+    @property
+    def engine_specific_fields(self) -> list[str]:
+        """Get the engine specific fields."""
+        pass
+
+    @property
+    def excluded_engine_specific_fields(self):
+        """A method to get the excluded fields for the database engine."""
+
+        return [
+            field.value.capitalize()
+            for field in DatabaseEngineSpecificFields
+            if field not in self.engine_specific_fields
+        ]
+
     @abstractmethod
     async def query_execution(
         self,

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

@@ -7,7 +7,7 @@
 import logging
 import json
 
-from text_2_sql_core.utils.database import DatabaseEngine
+from text_2_sql_core.utils.database import DatabaseEngine, DatabaseEngineSpecificFields
 
 
 class TSQLSqlConnector(SqlConnector):
@@ -16,6 +16,11 @@ def __init__(self):
 
         self.database_engine = DatabaseEngine.TSQL
 
+    @property
+    def engine_specific_fields(self) -> list[str]:
+        """Get the engine specific fields."""
+        return [DatabaseEngineSpecificFields.DATABASE]
+
     async def query_execution(
         self,
         sql_query: Annotated[

text_2_sql/text_2_sql_core/src/text_2_sql_core/data_dictionary/data_dictionary_creator.py

@@ -11,7 +11,6 @@
 import random
 import re
 import networkx as nx
-from text_2_sql_core.utils.database import DatabaseEngine
 from tenacity import retry, stop_after_attempt, wait_exponential
 from text_2_sql_core.connectors.open_ai import OpenAIConnector
 
@@ -751,23 +750,9 @@ async def build_entity_entry(self, entity: EntityItem) -> EntityItem:
     def excluded_fields_for_database_engine(self):
         """A method to get the excluded fields for the database engine."""
 
-        all_engine_specific_fields = ["Warehouse", "Database", "Catalog"]
-        if self.database_engine == DatabaseEngine.SNOWFLAKE:
-            engine_specific_fields = ["Warehouse", "Database"]
-        elif self.database_engine == DatabaseEngine.TSQL:
-            engine_specific_fields = ["Database"]
-        elif self.database_engine == DatabaseEngine.DATABRICKS:
-            engine_specific_fields = ["Catalog"]
-        elif self.database_engine == DatabaseEngine.POSTGRESQL:
-            engine_specific_fields = ["Database"]
-        else:
-            engine_specific_fields = []
-
         # Determine top-level fields to exclude
         filtered_entitiy_specific_fields = {
-            field.lower(): ...
-            for field in all_engine_specific_fields
-            if field not in engine_specific_fields
+            field.lower(): ... for field in self.excluded_engine_specific_fields
         }
 
         if filtered_entitiy_specific_fields:

text_2_sql/text_2_sql_core/src/text_2_sql_core/utils/database.py

@@ -8,3 +8,11 @@ class DatabaseEngine(StrEnum):
     SNOWFLAKE = "SNOWFLAKE"
     TSQL = "TSQL"
     POSTGRESQL = "POSTGRESQL"
+
+
+class DatabaseEngineSpecificFields(StrEnum):
+    """An enumeration to represent the database engine specific fields."""
+
+    WAREHOUSE = "Warehouse"
+    DATABASE = "Database"
+    CATALOG = "Catalog"