Update creator scripts

Author: BenConstable9

deploy_ai_search/text_2_sql_query_cache.py

@@ -112,12 +112,6 @@ def get_index_fields(self) -> list[SearchableField]:
                                     SearchableField(
                                         name="DataType", type=SearchFieldDataType.String
                                     ),
-                                    SearchableField(
-                                        name="AllowedValues",
-                                        type=SearchFieldDataType.String,
-                                        collection=True,
-                                        searchable=False,
-                                    ),
                                     SearchableField(
                                         name="SampleValues",
                                         type=SearchFieldDataType.String,

deploy_ai_search/text_2_sql_schema_store.py

@@ -126,12 +126,6 @@ def get_index_fields(self) -> list[SearchableField]:
                     SearchableField(name="Name", type=SearchFieldDataType.String),
                     SearchableField(name="Definition", type=SearchFieldDataType.String),
                     SearchableField(name="DataType", type=SearchFieldDataType.String),
-                    SearchableField(
-                        name="AllowedValues",
-                        type=SearchFieldDataType.String,
-                        collection=True,
-                        searchable=False,
-                    ),
                     SearchableField(
                         name="SampleValues",
                         type=SearchFieldDataType.String,

text_2_sql/text_2_sql_core/src/text_2_sql_core/data_dictionary/cli.py

@@ -82,7 +82,7 @@ def create(
         rich_print(detailed_error)
 
         raise typer.Exit(code=1)
-
+    asyncio.run(data_dictionary_creator.create_data_dictionary())
     try:
         asyncio.run(data_dictionary_creator.create_data_dictionary())
     except Exception as e:

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

@@ -100,7 +100,6 @@ class ColumnItem(BaseModel):
     distinct_values: Optional[list[any]] = Field(
         None, alias="DistinctValues", exclude=True
     )
-    allowed_values: Optional[list[any]] = Field(None, alias="AllowedValues")
     sample_values: Optional[list[any]] = Field(None, alias="SampleValues")
 
     model_config = ConfigDict(populate_by_name=True, arbitrary_types_allowed=True)
@@ -160,12 +159,12 @@ def id(self):
 
     def value_store_entry(self, excluded_fields_for_database_engine):
         excluded_fields = excluded_fields_for_database_engine + [
-            "Definition",
-            "Name",
-            "EntityName",
-            "EntityRelationships",
-            "CompleteEntityRelationshipsGraph",
-            "Columns",
+            "definition",
+            "name",
+            "entity_name",
+            "entity_relationships",
+            "complete_entity_relationships_graph",
+            "columns",
         ]
         return self.model_dump(
             by_alias=True, exclude_none=True, exclude=excluded_fields
@@ -226,8 +225,8 @@ def __init__(
 
         self.database_engine = None
 
-        self.database_semaphore = asyncio.Semaphore(50)
-        self.llm_semaphone = asyncio.Semaphore(20)
+        self.database_semaphore = asyncio.Semaphore(20)
+        self.llm_semaphone = asyncio.Semaphore(10)
 
         if output_directory is None:
             self.output_directory = "."
@@ -274,7 +273,7 @@ def extract_distinct_values_sql_query(
         Returns:
             str: The SQL query to extract distinct values from a column.
         """
-        return f"""SELECT DISTINCT {column.name} FROM {entity.entity} ORDER BY {column.name} DESC;"""
+        return f"""SELECT DISTINCT {column.name} FROM {entity.entity} WHERE {column.name} IS NOT NULL ORDER BY {column.name} DESC;"""
 
     @retry(
         stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=1, max=10)
@@ -422,13 +421,13 @@ async def extract_entities_with_definitions(self) -> list[EntityItem]:
         all_entities = table_entities + view_entities
 
         # Filter entities if entities is not None
-        if self.entities:
+        if self.entities is not None:
             all_entities = [
                 entity for entity in all_entities if entity.entity in self.entities
             ]
 
         # Filter entities if excluded_entities is not None
-        if self.excluded_entities:
+        if len(self.excluded_entities) > 0 or len(self.excluded_schemas):
             all_entities = [
                 entity
                 for entity in all_entities
@@ -448,20 +447,24 @@ async def write_columns_to_file(self, entity: EntityItem, column: ColumnItem):
         logging.info(f"Saving column values for {column.name}")
 
         key = f"{entity.id}.{column.name}"
+        # Ensure the intermediate directories exist
+        os.makedirs(f"{self.output_directory}/column_value_store", exist_ok=True)
         with open(
             f"{self.output_directory}/column_value_store/{key}.jsonl",
             "w",
             encoding="utf-8",
         ) as f:
-            for distinct_value in column.distinct_values:
-                json.dump(
-                    column.value_store_entry(
-                        entity, distinct_value, self.excluded_fields_for_database_engine
-                    ),
-                    f,
-                    indent=4,
-                    default=str,
-                )
+            if column.distinct_values is not None:
+                for distinct_value in column.distinct_values:
+                    json_string = json.dumps(
+                        column.value_store_entry(
+                            entity,
+                            distinct_value,
+                            self.excluded_fields_for_database_engine,
+                        ),
+                        default=str,
+                    )
+                    f.write(json_string + "\n")
 
     async def extract_column_distinct_values(
         self, entity: EntityItem, column: ColumnItem
@@ -498,6 +501,8 @@ async def extract_column_distinct_values(
         elif column.distinct_values is not None:
             column.sample_values = column.distinct_values
 
+        await self.write_columns_to_file(entity, column)
+
     async def generate_column_definition(self, entity: EntityItem, column: ColumnItem):
         """A method to generate a definition for a column in a database.
 
@@ -682,6 +687,9 @@ async def generate_entity_definition(self, entity: EntityItem):
 
     async def write_entity_to_file(self, entity):
         logging.info(f"Saving data dictionary for {entity.entity}")
+
+        # Ensure the intermediate directories exist
+        os.makedirs(f"{self.output_directory}/schema_store", exist_ok=True)
         with open(
             f"{self.output_directory}/schema_store/{entity.id}.json",
             "w",
@@ -742,7 +750,7 @@ def excluded_fields_for_database_engine(self):
             engine_specific_fields = ["Catalog"]
 
         return [
-            field
+            field.lower()
             for field in all_engine_specific_fields
             if field not in engine_specific_fields
         ]
@@ -764,6 +772,8 @@ async def create_data_dictionary(self):
         # Save data dictionary to file
         if self.single_file:
             logging.info("Saving data dictionary to entities.json")
+            # Ensure the intermediate directories exist
+            os.makedirs(f"{self.output_directory}/schema_store", exist_ok=True)
             with open(
                 f"{self.output_directory}/schema_store/entities.json",
                 "w",

text_2_sql/text_2_sql_core/src/text_2_sql_core/data_dictionary/databricks_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, ColumnItem
 import asyncio
 from databricks import sql
 import logging
@@ -30,6 +30,8 @@ def extract_table_entities_sql_query(self) -> str:
             t.COMMENT AS Definition
         FROM
             {self.catalog}.INFORMATION_SCHEMA.TABLES t
+        ORDER BY EntitySchema, Entity
+        LIMIT 1
         """
 
     @property
@@ -41,6 +43,8 @@ def extract_view_entities_sql_query(self) -> str:
             NULL AS Definition
         FROM
             {self.catalog}.INFORMATION_SCHEMA.VIEWS v
+        ORDER BY EntitySchema, Entity
+        LIMIT 1
         """
 
     def extract_columns_sql_query(self, entity: EntityItem) -> str:
@@ -94,6 +98,20 @@ def extract_entity_relationships_sql_query(self) -> str:
             EntitySchema, Entity, ForeignEntitySchema, ForeignEntity;
         """
 
+    def extract_distinct_values_sql_query(
+        self, entity: EntityItem, column: ColumnItem
+    ) -> str:
+        """A method to extract distinct values from a column in a database. Can be sub-classed if needed.
+
+        Args:
+            entity (EntityItem): The entity to extract distinct values from.
+            column (ColumnItem): The column to extract distinct values from.
+
+        Returns:
+            str: The SQL query to extract distinct values from a column.
+        """
+        return f"""SELECT DISTINCT {column.name} FROM {self.catalog}.{entity.entity} WHERE {column.name} IS NOT NULL ORDER BY {column.name} DESC;"""
+
     @retry(
         stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=1, max=10)
     )
@@ -134,14 +152,14 @@ async def query_entities(self, sql_query: str, cast_to: any = None) -> list[dict
 
                 # Process rows
                 for row in rows:
-                    if cast_to:
+                    if cast_to is not None:
                         results.append(cast_to.from_sql_row(row, columns))
                     else:
                         results.append(dict(zip(columns, row)))
 
             except Exception as e:
-                logging.error(f"Error while executing query: {e}")
-                raise
+                logging.error(f"Error while executing query {sql_query}: {e}")
+                raise e
             finally:
                 cursor.close()
                 connection.close()

text_2_sql/text_2_sql_core/src/text_2_sql_core/data_dictionary/snowflake_data_dictionary_creator.py

@@ -31,7 +31,8 @@ def extract_table_entities_sql_query(self) -> str:
             t.TABLE_SCHEMA AS EntitySchema,
             t.COMMENT AS Definition
         FROM
-            INFORMATION_SCHEMA.TABLES t"""
+            INFORMATION_SCHEMA.TABLES t
+        ORDER BY EntitySchema, Entity"""
 
     @property
     def extract_view_entities_sql_query(self) -> str:
@@ -41,13 +42,14 @@ def extract_view_entities_sql_query(self) -> str:
             v.TABLE_SCHEMA AS EntitySchema,
             v.COMMENT AS Definition
         FROM
-            INFORMATION_SCHEMA.VIEWS v"""
+            INFORMATION_SCHEMA.VIEWS v
+        ORDER BY EntitySchema, Entity"""
 
     def extract_columns_sql_query(self, entity: EntityItem) -> str:
         """A property to extract column information from a Snowflake database."""
         return f"""SELECT
             COLUMN_NAME AS Name,
-            DATA_TYPE AS Type,
+            DATA_TYPE AS DataType,
             COMMENT AS Definition
         FROM
             INFORMATION_SCHEMA.COLUMNS

text_2_sql/text_2_sql_core/src/text_2_sql_core/data_dictionary/tsql_data_dictionary_creator.py

@@ -39,7 +39,8 @@ def extract_table_entities_sql_query(self) -> str:
             AND ep.class = 1
             AND ep.name = 'MS_Description'
         WHERE
-            t.TABLE_TYPE = 'BASE TABLE';"""
+            t.TABLE_TYPE = 'BASE TABLE';
+        ORDER BY EntitySchema, Entity"""
 
     @property
     def extract_view_entities_sql_query(self) -> str:
@@ -55,7 +56,8 @@ def extract_view_entities_sql_query(self) -> str:
             ON ep.major_id = OBJECT_ID(v.TABLE_SCHEMA + '.' + v.TABLE_NAME)
             AND ep.minor_id = 0
             AND ep.class = 1
-    AND ep.name = 'MS_Description';"""
+        AND ep.name = 'MS_Description';
+        ORDER BY EntitySchema, Entity"""
 
     def extract_columns_sql_query(self, entity: EntityItem) -> str:
         """A property to extract column information from a SQL Server database."""