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.
66 lines
2.0 KiB
Python
66 lines
2.0 KiB
Python
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()
|