Update snowflake script

Author: BenConstable9

text_2_sql/data_dictionary/.env

@@ -8,6 +8,10 @@ Text2Sql__UseQueryCache=<whether to use the query cache first or not>
 Text2Sql__PreRunQueryCache=<whether to pre-run the top result from the query cache or not>
 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>
 AIService__AzureSearchOptions__Endpoint=<searchServiceEndpoint>
 AIService__AzureSearchOptions__Key=<searchServiceKey if not using identity>
 AIService__AzureSearchOptions__RagDocuments__Index=<ragDocumentsIndexName>

text_2_sql/data_dictionary/data_dictionary_creator.py

@@ -78,6 +78,7 @@ def __init__(
         self,
         entities: list[str] = None,
         excluded_entities: list[str] = None,
+        excluded_schemas: list[str] = None,
         single_file: bool = False,
         generate_descriptions: bool = True,
     ):
@@ -86,12 +87,14 @@ def __init__(
         Args:
             entities (list[str], optional): A list of entities to extract. Defaults to None. If None, all entities are extracted.
             excluded_entities (list[str], optional): A list of entities to exclude. Defaults to None.
+            excluded_schemas (list[str], optional): A list of schemas to exclude. Defaults to None.
             single_file (bool, optional): A flag to indicate if the data dictionary should be saved to a single file. Defaults to False.
             generate_descriptions (bool, optional): A flag to indicate if descriptions should be generated. Defaults to True.
         """
 
         self.entities = entities
         self.excluded_entities = excluded_entities
+        self.excluded_schemas = excluded_schemas
         self.single_file = single_file
         self.generate_descriptions = generate_descriptions
 
@@ -189,6 +192,7 @@ async def extract_entities_with_descriptions(self) -> list[EntityItem]:
                 entity
                 for entity in all_entities
                 if entity.entity not in self.excluded_entities
+                and entity.entity_schema not in self.excluded_schemas
             ]
 
         return all_entities

text_2_sql/data_dictionary/requirements.txt

@@ -3,3 +3,4 @@ azure-identity
 python-dotenv
 pydantic
 openai
+snowflake-connector-python

text_2_sql/data_dictionary/snowflake_data_dictionary_creator.py

@@ -0,0 +1,119 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+from data_dictionary_creator import DataDictionaryCreator, EntityItem
+import asyncio
+import snowflake.connector
+import logging
+import os
+
+
+class SnowflakeDataDictionaryCreator(DataDictionaryCreator):
+    def __init__(
+        self,
+        entities: list[str] = None,
+        excluded_entities: list[str] = None,
+        single_file: bool = False,
+    ):
+        """A method to initialize the DataDictionaryCreator class.
+
+        Args:
+            entities (list[str], optional): A list of entities to extract. Defaults to None. If None, all entities are extracted.
+            excluded_entities (list[str], optional): A list of entities to exclude. Defaults to None.
+            single_file (bool, optional): A flag to indicate if the data dictionary should be saved to a single file. Defaults to False.
+        """
+        if excluded_entities is None:
+            excluded_entities = []
+
+        excluded_schemas = ["dbo", "sys"]
+        return super().__init__(
+            entities, excluded_entities, excluded_schemas, single_file
+        )
+
+    """A class to extract data dictionary information from a Snowflake database."""
+
+    @property
+    def extract_table_entities_sql_query(self) -> str:
+        """A property to extract table entities from a Snowflake database."""
+        return """SELECT
+        t.TABLE_NAME AS Entity,
+        t.TABLE_SCHEMA AS EntitySchema,
+        t.COMMENT AS Description
+    FROM
+        INFORMATION_SCHEMA.TABLES t"""
+
+    @property
+    def extract_view_entities_sql_query(self) -> str:
+        """A property to extract view entities from a Snowflake database."""
+        return """SELECT
+        v.TABLE_NAME AS Entity,
+        v.TABLE_SCHEMA AS EntitySchema,
+        v.COMMENT AS Description
+    FROM
+        INFORMATION_SCHEMA.VIEWS v"""
+
+    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,
+        COMMENT AS Definition
+    FROM
+        INFORMATION_SCHEMA.COLUMNS
+    WHERE
+        TABLE_SCHEMA = '{entity.entity_schema}'
+        AND TABLE_NAME = '{entity.name}';"""
+
+    async def query_entities(
+        self, sql_query: str, cast_to: any = None
+    ) -> list[EntityItem]:
+        """A method to query a database for entities using Snowflake Connector. Overrides the base class method.
+
+        Args:
+            sql_query (str): The SQL query to run.
+            cast_to (any, optional): The class to cast the results to. Defaults to None.
+
+        Returns:
+            list[EntityItem]: The list of entities.
+        """
+        logging.info(f"Running query: {sql_query}")
+        results = []
+
+        # Create a connection to Snowflake, without specifying a schema
+        conn = snowflake.connector.connect(
+            user=os.environ["Text2Sql__Snowflake__User"],
+            password=os.environ["Text2Sql__Snowflake__Password"],
+            account=os.environ["Text2Sql__Snowflake__Account"],
+            warehouse=os.environ["Text2Sql__Snowflake__Warehouse"],
+            database=os.environ["Text2Sql__DatabaseName"],
+        )
+
+        try:
+            # Using the connection to create a cursor
+            cursor = conn.cursor()
+
+            # Execute the query
+            await asyncio.to_thread(cursor.execute, sql_query)
+
+            # Fetch column names
+            columns = [col[0] for col in cursor.description]
+
+            # Fetch rows
+            rows = await asyncio.to_thread(cursor.fetchall)
+
+            # Process rows
+            for row in rows:
+                if cast_to:
+                    results.append(cast_to.from_sql_row(row, columns))
+                else:
+                    results.append(dict(zip(columns, row)))
+
+        finally:
+            cursor.close()
+            conn.close()
+
+        return results
+
+
+if __name__ == "__main__":
+    data_dictionary_creator = SnowflakeDataDictionaryCreator()
+    asyncio.run(data_dictionary_creator.create_data_dictionary())

text_2_sql/data_dictionary/sql_sever_data_dictionary_creator.py

@@ -21,10 +21,10 @@ def __init__(
         if excluded_entities is None:
             excluded_entities = []
 
-        excluded_entities.extend(
-            ["dbo.BuildVersion", "dbo.ErrorLog", "sys.database_firewall_rules"]
+        excluded_schemas = ["dbo", "sys"]
+        return super().__init__(
+            entities, excluded_entities, excluded_schemas, single_file
         )
-        return super().__init__(entities, excluded_entities, single_file)
 
     """A class to extract data dictionary information from a SQL Server database."""