db.py•3.48 kB
import datetime
import os
import mysql.connector
from mysql.connector import Error
from dotenv import load_dotenv
def execute_mysql_query(query):
connection = None
cursor = None
try:
# Establish a database connection
connection = mysql.connector.connect(
host=os.getenv('DB_HOST'),
database=os.getenv('DB_NAME'),
user=os.getenv('DB_USER'),
password=os.getenv('DB_PASSWORD'),
port=os.getenv('DB_PORT', 3306)
)
if connection.is_connected():
cursor = connection.cursor()
cursor.execute(query)
# For SELECT queries, fetch the results
if query.strip().upper().startswith('SELECT'):
result = cursor.fetchall()
return result, cursor.description
else:
# For INSERT, UPDATE, DELETE, commit the transaction
connection.commit()
return cursor.rowcount, None
except Error as e:
print(f"Error while connecting to MySQL: {e}")
return None, None
finally:
# Close cursor and connection in any case
if cursor is not None:
cursor.close()
if connection is not None and connection.is_connected():
connection.close()
def format_results_for_llm(query, results, description, rowcount=None):
query_type = query.strip().upper().split()[0]
response = {
"query": query,
"type": query_type,
"timestamp": datetime.datetime.now().isoformat(),
"status": "success"
}
if query_type == "SELECT":
if not results:
response["message"] = "Query executed successfully but returned no results"
response["results"] = []
return response
# Get column names (assuming cursor.description is available)
try:
# If results came from mysql.connector cursor
columns = [desc[0] for desc in description]
except AttributeError:
# Fallback if we can't get column names
columns = [f"column_{i}" for i in range(len(results[0]))]
# Format results as a list of dictionaries
formatted_results = []
for row in results:
formatted_results.append(dict(zip(columns, row)))
response.update({
"results": formatted_results,
"count": len(results),
"format": "Each result is a dictionary with column names as keys"
})
elif query_type in ("INSERT", "UPDATE", "DELETE"):
response.update({
"rows_affected": rowcount,
"message": f"{query_type} operation completed successfully"
})
else: # For other query types (CREATE, ALTER, etc.)
response.update({
"message": f"{query_type} operation executed successfully",
"note": "No results to return for this query type"
})
return response
if __name__ == "__main__":
load_dotenv()
query = """INSERT INTO cars (brand, model, color, `condition`, mileage, price, owner_email)
VALUES
('Toyota', 'Corola', 'Silver', 'used', 45000, 18500.00, 'john.dee@example.com'),"""
results, description = execute_mysql_query(query)
print(format_results_for_llm(query,results, description))
