Start adding Databricks support

Author: BenConstable9

deploy_ai_search/.env

@@ -19,4 +19,4 @@ OpenAI__Endpoint=<openAIEndpoint>
 OpenAI__EmbeddingModel=<openAIEmbeddingModelName>
 OpenAI__EmbeddingDeployment=<openAIEmbeddingDeploymentId>
 OpenAI__EmbeddingDimensions=1536
-Text2Sql__DatabaseName=<databaseName>
+Text2Sql__DatabaseEngine=<databaseEngine SQL Server / Snowflake / Databricks >

deploy_ai_search/text_2_sql_schema_store.py

@@ -24,6 +24,16 @@
 from environment import (
     IndexerType,
 )
+import os
+from enum import StrEnum
+
+
+class DatabaseEngine(StrEnum):
+    """An enumeration to represent a database engine."""
+
+    SNOWFLAKE = "SNOWFLAKE"
+    SQL_SERVER = "SQL_SERVER"
+    DATABRICKS = "DATABRICKS"
 
 
 class Text2SqlSchemaStoreAISearch(AISearch):
@@ -42,13 +52,34 @@ def __init__(
             rebuild (bool, optional): Whether to rebuild the index. Defaults to False.
         """
         self.indexer_type = IndexerType.TEXT_2_SQL_SCHEMA_STORE
+        self.database_engine = DatabaseEngine[
+            os.environ["Text2Sql__DatabaseEngine"].upper()
+        ]
         super().__init__(suffix, rebuild)
 
         if single_data_dictionary:
             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.SQL_SERVER:
+            engine_specific_fields = ["Database"]
+        elif self.database_engine == DatabaseEngine.DATABRICKS:
+            engine_specific_fields = ["Catalog"]
+
+        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.
 
@@ -78,6 +109,10 @@ def get_index_fields(self) -> list[SearchableField]:
                 name="Warehouse",
                 type=SearchFieldDataType.String,
             ),
+            SearchableField(
+                name="Catalog",
+                type=SearchFieldDataType.String,
+            ),
             SearchableField(
                 name="Definition",
                 type=SearchFieldDataType.String,
@@ -161,6 +196,13 @@ def get_index_fields(self) -> list[SearchableField]:
             ),
         ]
 
+        # Remove fields that are not supported by the database engine
+        fields = [
+            field
+            for field in fields
+            if field.name not in self.excluded_fields_for_database_engine
+        ]
+
         return fields
 
     def get_semantic_search(self) -> SemanticSearch:
@@ -309,4 +351,12 @@ def get_indexer(self) -> SearchIndexer:
             parameters=indexer_parameters,
         )
 
+        # Remove fields that are not supported by the database engine
+        indexer.output_field_mappings = [
+            field_mapping
+            for field_mapping in indexer.output_field_mappings
+            if field_mapping.target_field_name
+            not in self.excluded_fields_for_database_engine
+        ]
+
         return indexer

text_2_sql/data_dictionary/.env

@@ -3,12 +3,12 @@ OpenAI__EmbeddingModel=<openAIEmbeddingModelName>
 OpenAI__Endpoint=<openAIEndpoint>
 OpenAI__ApiKey=<openAIKey if using non managed identity>
 OpenAI__ApiVersion=<openAIApiVersion>
-Text2Sql__DatabaseEngine=<databaseEngine>
 Text2Sql__DatabaseName=<databaseName>
 Text2Sql__DatabaseConnectionString=<databaseConnectionString>
 Text2Sql__Snowflake__User=<snowflakeUser if using Snowflake Data Source>
 Text2Sql__Snowflake__Password=<snowflakePassword if using Snowflake Data Source>
 Text2Sql__Snowflake__Account=<snowflakeAccount if using Snowflake Data Source>
 Text2Sql__Snowflake__Warehouse=<snowflakeWarehouse if using Snowflake Data Source>
+Text2Sql__Databricks__Catalog=<databricksCatalog if using Databricks Data Source with Unity Catalog>
 IdentityType=<identityType> # system_assigned or user_assigned or key
-ClientId=<clientId if using user assigned identity>
+ClientId=<clientId if using user assigned identity>

text_2_sql/data_dictionary/data_dictionary_creator.py

@@ -15,10 +15,19 @@
 import random
 import re
 import networkx as nx
+from enum import StrEnum
 
 logging.basicConfig(level=logging.INFO)
 
 
+class DatabaseEngine(StrEnum):
+    """An enumeration to represent a database engine."""
+
+    SNOWFLAKE = "SNOWFLAKE"
+    SQL_SERVER = "SQL_SERVER"
+    DATABRICKS = "DATABRICKS"
+
+
 class ForeignKeyRelationship(BaseModel):
     column: str = Field(..., alias="Column")
     foreign_column: str = Field(..., alias="ForeignColumn")
@@ -124,6 +133,7 @@ class EntityItem(BaseModel):
     entity_name: Optional[str] = Field(default=None, alias="EntityName")
     database: Optional[str] = Field(default=None, alias="Database")
     warehouse: Optional[str] = Field(default=None, alias="Warehouse")
+    catalog: Optional[str] = Field(default=None, alias="Catalog")
 
     entity_relationships: Optional[list[EntityRelationship]] = Field(
         alias="EntityRelationships", default_factory=list
@@ -186,6 +196,9 @@ def __init__(
 
         self.warehouse = None
         self.database = None
+        self.catalog = None
+
+        self.database_engine = None
 
         load_dotenv(find_dotenv())
 
@@ -391,6 +404,7 @@ async def extract_entities_with_definitions(self) -> list[EntityItem]:
         for entity in all_entities:
             entity.warehouse = self.warehouse
             entity.database = self.database
+            entity.catalog = self.catalog
 
         return all_entities
 
@@ -636,6 +650,24 @@ async def build_entity_entry(self, entity: EntityItem) -> EntityItem:
 
         return entity
 
+    @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.SQL_SERVER:
+            engine_specific_fields = ["Database"]
+        elif self.database_engine == DatabaseEngine.DATABRICKS:
+            engine_specific_fields = ["Catalog"]
+
+        return [
+            field
+            for field in all_engine_specific_fields
+            if field not in engine_specific_fields
+        ]
+
     async def create_data_dictionary(self):
         """A method to build a data dictionary from a database. Writes to file."""
         entities = await self.extract_entities_with_definitions()
@@ -655,12 +687,23 @@ async def create_data_dictionary(self):
             logging.info("Saving data dictionary to entities.json")
             with open("entities.json", "w", encoding="utf-8") as f:
                 json.dump(
-                    data_dictionary.model_dump(by_alias=True), f, indent=4, default=str
+                    data_dictionary.model_dump(
+                        by_alias=True, exclude=self.excluded_fields_for_database_engine
+                    ),
+                    f,
+                    indent=4,
+                    default=str,
                 )
         else:
             for entity in data_dictionary:
                 logging.info(f"Saving data dictionary for {entity.entity}")
                 with open(f"{entity.entity}.json", "w", encoding="utf-8") as f:
                     json.dump(
-                        entity.model_dump(by_alias=True), f, indent=4, default=str
+                        entity.model_dump(
+                            by_alias=True,
+                            exclude=self.excluded_fields_for_database_engine,
+                        ),
+                        f,
+                        indent=4,
+                        default=str,
                     )

text_2_sql/data_dictionary/snowflake_data_dictionary_creator.py

@@ -1,6 +1,6 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
-from data_dictionary_creator import DataDictionaryCreator, EntityItem
+from data_dictionary_creator import DataDictionaryCreator, EntityItem, DatabaseEngine
 import asyncio
 import snowflake.connector
 import logging
@@ -25,9 +25,11 @@ def __init__(
             excluded_entities = []
 
         excluded_schemas = ["INFORMATION_SCHEMA"]
-        return super().__init__(
-            entities, excluded_entities, excluded_schemas, single_file
-        )
+        super().__init__(entities, excluded_entities, excluded_schemas, single_file)
+
+        self.database = os.environ["Text2Sql__DatabaseName"]
+        self.warehouse = os.environ["Text2Sql__Snowflake__Warehouse"]
+        self.database_engine = DatabaseEngine.SNOWFLAKE
 
     """A class to extract data dictionary information from a Snowflake database."""
 

text_2_sql/data_dictionary/sql_sever_data_dictionary_creator.py

@@ -1,6 +1,6 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
-from data_dictionary_creator import DataDictionaryCreator, EntityItem
+from data_dictionary_creator import DataDictionaryCreator, EntityItem, DatabaseEngine
 import asyncio
 import os
 
@@ -26,6 +26,8 @@ def __init__(
         super().__init__(entities, excluded_entities, excluded_schemas, single_file)
         self.database = os.environ["Text2Sql__DatabaseName"]
 
+        self.database_engine = DatabaseEngine.SQL_SERVER
+
     """A class to extract data dictionary information from a SQL Server database."""
 
     @property