Files
grid_application/app/sql_utils.py
Tobias Quadfasel 4b9fa0579e 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.
2024-09-02 20:43:48 +02:00

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()