feat(ai-chat): Add SQL query field for comparison
In order to compare the (not yet implemented) SQL query generated by the LLM with an actual query, another text field was added that parses the query to `pyodbc`, which connects to our database, stores the resulting rows in a `pandas` dataframe and then visualizes it as a table in plotly dash. The SQL functionalities are implemented in the `sql_utils.py` module. Additionally, some minor updates to the overall behavior and layout of the app were implemented.
This commit is contained in:
65
app/sql_utils.py
Normal file
65
app/sql_utils.py
Normal file
@@ -0,0 +1,65 @@
|
||||
import os
|
||||
from typing import Union
|
||||
|
||||
import pandas as pd
|
||||
import pyodbc
|
||||
|
||||
|
||||
def test_db_connection() -> bool:
|
||||
"""Test the connection to Azure SQL Database.
|
||||
|
||||
This function attempts to establish a connection to an Azure SQL Database
|
||||
using the connection string stored in the environment variable
|
||||
'AZURE_SQL_CONNECTION_STRING'. It makes up to 5 attempts to connect,
|
||||
with a timeout of 240 seconds for each attempt.
|
||||
|
||||
Returns
|
||||
-------
|
||||
bool
|
||||
True if the connection was successful, False otherwise.
|
||||
"""
|
||||
connection_string = os.environ.get("AZURE_SQL_CONNECTION_STRING")
|
||||
for i in range(5):
|
||||
print(f"Trying to connect to Azure SQL Database... Attempt {i + 1}")
|
||||
try:
|
||||
pyodbc.connect(connection_string, timeout=240)
|
||||
print("Connected to Azure SQL Database successfully!")
|
||||
connected = True
|
||||
break
|
||||
|
||||
except pyodbc.Error as e:
|
||||
print(f"Error connecting to Azure SQL Database: {e}")
|
||||
connected = False
|
||||
|
||||
return connected
|
||||
|
||||
|
||||
def execute_query(query: str) -> Union[pd.DataFrame, str]:
|
||||
"""Execute a SQL query on an Azure SQL Database and return the results.
|
||||
|
||||
This function connects to an Azure SQL Database using the connection string
|
||||
stored in the environment variable 'AZURE_SQL_CONNECTION_STRING', executes
|
||||
the provided SQL query, and returns the results as a pandas DataFrame.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
query : str
|
||||
The SQL query to execute.
|
||||
|
||||
Returns
|
||||
-------
|
||||
Union[pd.DataFrame, str]
|
||||
A pandas DataFrame containing the query results if successful,
|
||||
or a string containing the error message if an exception occurs.
|
||||
"""
|
||||
try:
|
||||
connection_string = os.environ.get("AZURE_SQL_CONNECTION_STRING")
|
||||
conn = pyodbc.connect(connection_string, timeout=240)
|
||||
df = pd.read_sql(query, conn)
|
||||
conn.close()
|
||||
return df
|
||||
except Exception as e:
|
||||
return str(e)
|
||||
finally:
|
||||
if conn in locals():
|
||||
conn.close()
|
||||
Reference in New Issue
Block a user