More bug fixes

Author: BenConstable9

text_2_sql/text_2_sql_core/pyproject.toml

@@ -18,7 +18,9 @@ dependencies = [
     "pydantic>=2.10.2",
     "python-dotenv>=1.0.1",
     "pyyaml>=6.0.2",
+    "rich>=13.9.4",
     "sqlglot[rs]>=25.32.1",
+    "typer>=0.14.0",
 ]
 
 [project.optional-dependencies]

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

@@ -9,25 +9,17 @@
 import logging
 from pydantic import BaseModel, Field, ConfigDict
 from typing import Optional
-from environment import IdentityType, get_identity_type
+from text_2_sql_core.utils.environment import IdentityType, get_identity_type
 from openai import AsyncAzureOpenAI
 from azure.identity import DefaultAzureCredential, get_bearer_token_provider
 import random
 import re
 import networkx as nx
-from enum import StrEnum
+from text_2_sql_core.utils.database import DatabaseEngine
 
 logging.basicConfig(level=logging.INFO)
 
 
-class DatabaseEngine(StrEnum):
-    """An enumeration to represent a database engine."""
-
-    SNOWFLAKE = "SNOWFLAKE"
-    TSQL = "TSQL"
-    DATABRICKS = "DATABRICKS"
-
-
 class ForeignKeyRelationship(BaseModel):
     column: str = Field(..., alias="Column")
     foreign_column: str = Field(..., alias="ForeignColumn")

text_2_sql/text_2_sql_core/src/text_2_sql_core/data_dictionary/databricks_data_dictionary_creator.py

@@ -1,10 +1,11 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
-from data_dictionary_creator import DataDictionaryCreator, EntityItem, DatabaseEngine
+from data_dictionary_creator import DataDictionaryCreator, EntityItem
 import asyncio
 from databricks import sql
 import logging
 import os
+from text_2_sql_core.utils.database import DatabaseEngine
 
 
 class DatabricksDataDictionaryCreator(DataDictionaryCreator):
@@ -61,7 +62,7 @@ def extract_columns_sql_query(self, entity: EntityItem) -> str:
         """A property to extract column information from Databricks Unity Catalog."""
         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/snowflake_data_dictionary_creator.py

@@ -1,11 +1,13 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
-from data_dictionary_creator import DataDictionaryCreator, EntityItem, DatabaseEngine
+from data_dictionary_creator import DataDictionaryCreator, EntityItem
 import asyncio
 import snowflake.connector
 import logging
 import os
 
+from text_2_sql_core.utils.database import DatabaseEngine
+
 
 class SnowflakeDataDictionaryCreator(DataDictionaryCreator):
     def __init__(

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

@@ -1,8 +1,9 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
-from data_dictionary_creator import DataDictionaryCreator, EntityItem, DatabaseEngine
+from data_dictionary_creator import DataDictionaryCreator, EntityItem
 import asyncio
 import os
+from text_2_sql_core.utils.database import DatabaseEngine
 
 
 class SqlServerDataDictionaryCreator(DataDictionaryCreator):

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

@@ -0,0 +1,9 @@
+from enum import StrEnum
+
+
+class DatabaseEngine(StrEnum):
+    """An enumeration to represent a database engine."""
+
+    SNOWFLAKE = "SNOWFLAKE"
+    TSQL = "TSQL"
+    DATABRICKS = "DATABRICKS"