Use tsql instead

Author: BenConstable9

deploy_ai_search/text_2_sql_schema_store.py

@@ -32,7 +32,7 @@ class DatabaseEngine(StrEnum):
     """An enumeration to represent a database engine."""
 
     SNOWFLAKE = "SNOWFLAKE"
-    SQL_SERVER = "SQL_SERVER"
+    TSQL = "TSQL"
     DATABRICKS = "DATABRICKS"
 
 
@@ -69,7 +69,7 @@ def excluded_fields_for_database_engine(self):
         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:
+        elif self.database_engine == DatabaseEngine.TSQL:
             engine_specific_fields = ["Database"]
         elif self.database_engine == DatabaseEngine.DATABRICKS:
             engine_specific_fields = ["Catalog"]

text_2_sql/autogen/requirements.txt

@@ -9,3 +9,4 @@ python-dotenv
 openai
 jinja2
 pyyaml
+sqlglot[rs]

text_2_sql/autogen/utils/sql.py

@@ -3,10 +3,11 @@
 import logging
 import os
 import aioodbc
-from typing import Annotated
+from typing import Annotated, Union
 from utils.ai_search import run_ai_search_query
 import json
 import asyncio
+import sqlglot
 
 USE_QUERY_CACHE = os.environ.get("Text2Sql__UseQueryCache", "False").lower() == "true"
 
@@ -66,7 +67,12 @@ async def get_entity_schemas(
     return json.dumps(schemas, default=str)
 
 
-async def query_execution(sql_query: str) -> list[dict]:
+async def query_execution(
+    sql_query: Annotated[
+        str,
+        "The SQL query to run against the database.",
+    ]
+) -> list[dict]:
     """Run the SQL query against the database.
 
     Args:
@@ -91,6 +97,21 @@ async def query_execution(sql_query: str) -> list[dict]:
     return results
 
 
+async def validate_sql_query(
+    sql_query: Annotated[
+        str,
+        "The SQL query to run against the database.",
+    ]
+) -> Union[bool | list[dict]]:
+    """Validate the SQL query."""
+    try:
+        sqlglot.transpile("SELECT foo FROM (SELECT baz FROM t")
+    except sqlglot.errors.ParseError as e:
+        return e.errors
+    else:
+        return True
+
+
 async def fetch_queries_from_cache(question: str) -> str:
     """Fetch the queries from the cache based on the question.
 

text_2_sql/data_dictionary/README.md

@@ -101,6 +101,6 @@ The following Databases have pre-built scripts for them:
 
 - **Databricks:** `databricks_data_dictionary_creator.py`
 - **Snowflake:** `snowflake_data_dictionary_creator.py`
-- **SQL Server:** `sql_server_data_dictionary_creator.py`
+- **SQL Server:** `tsql_data_dictionary_creator.py`
 
 If there is no pre-built script for your database engine, take one of the above as a starting point and adjust it.

text_2_sql/data_dictionary/data_dictionary_creator.py

@@ -24,7 +24,7 @@ class DatabaseEngine(StrEnum):
     """An enumeration to represent a database engine."""
 
     SNOWFLAKE = "SNOWFLAKE"
-    SQL_SERVER = "SQL_SERVER"
+    TSQL = "TSQL"
     DATABRICKS = "DATABRICKS"
 
 
@@ -657,7 +657,7 @@ def excluded_fields_for_database_engine(self):
         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:
+        elif self.database_engine == DatabaseEngine.TSQL:
             engine_specific_fields = ["Database"]
         elif self.database_engine == DatabaseEngine.DATABRICKS:
             engine_specific_fields = ["Catalog"]

…ary/sql_sever_data_dictionary_creator.py …ctionary/tsql_data_dictionary_creator.pytext_2_sql/data_dictionary/sql_sever_data_dictionary_creator.py renamed to text_2_sql/data_dictionary/tsql_data_dictionary_creator.py text_2_sql/data_dictionary/sql_sever_data_dictionary_creator.py renamed to text_2_sql/data_dictionary/tsql_data_dictionary_creator.py

@@ -26,7 +26,7 @@ def __init__(
         super().__init__(entities, excluded_entities, excluded_schemas, single_file)
         self.database = os.environ["Text2Sql__DatabaseName"]
 
-        self.database_engine = DatabaseEngine.SQL_SERVER
+        self.database_engine = DatabaseEngine.TSQL
 
     """A class to extract data dictionary information from a SQL Server database."""