Spaces:
Runtime error
Runtime error
File size: 5,412 Bytes
909cddd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 |
import os
import sqlite3
import csv
from dotenv import load_dotenv
# Load environment variables from .env file
load_dotenv()
# SQLite DB path from .env for saving results
db_path = os.getenv("db_path", "./query_logs.db")
# Function to initialize the SQLite database
def initialize_local_db():
conn = sqlite3.connect(db_path)
cursor = conn.cursor()
# Check if the reformulated_query column exists
cursor.execute("PRAGMA table_info(query_logs);")
columns = [column[1] for column in cursor.fetchall()]
# If the reformulated_query column doesn't exist, alter the table
if 'reformulated_query' not in columns:
print("Altering table to add reformulated_query column...")
cursor.execute('''
CREATE TABLE IF NOT EXISTS query_logs_new (
id INTEGER PRIMARY KEY AUTOINCREMENT,
natural_language_query TEXT,
reformulated_query TEXT,
generated_sql TEXT,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
);
''')
# Copy data from old table to new table
cursor.execute('''
INSERT INTO query_logs_new (id, natural_language_query, generated_sql, created_at)
SELECT id, natural_language_query, generated_sql, created_at FROM query_logs;
''')
# Drop the old table
cursor.execute("DROP TABLE query_logs;")
# Rename the new table
cursor.execute("ALTER TABLE query_logs_new RENAME TO query_logs;")
conn.commit()
cursor.close()
conn.close()
# Function to reset (drop all tables) and recreate the schema
def reset_sqlite_db():
conn = sqlite3.connect(db_path)
cursor = conn.cursor()
# Drop the query_logs table
cursor.execute("DROP TABLE IF EXISTS query_logs;")
print("Dropped query_logs table")
# Recreate the query_logs table
cursor.execute('''CREATE TABLE query_logs (
id INTEGER PRIMARY KEY AUTOINCREMENT,
natural_language_query TEXT,
reformulated_query TEXT,
generated_sql TEXT,
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
);''')
print("Recreated the query_logs table.")
conn.commit()
cursor.close()
conn.close()
# Function to save the natural language query, reformulated query, and generated SQL to the local SQLite DB
def save_query_to_local_db(nl_query, reformulated_query, sql_query):
try:
conn = sqlite3.connect(db_path)
cursor = conn.cursor()
insert_query = '''
INSERT INTO query_logs (natural_language_query, reformulated_query, generated_sql)
VALUES (?, ?, ?);
'''
cursor.execute(insert_query, (nl_query, reformulated_query, sql_query))
conn.commit()
cursor.close()
conn.close()
except Exception as e:
print(f"Error saving query: {e}")
# Function to retrieve saved queries from the local SQLite DB with pagination and search
def get_saved_queries(page=1, per_page=10, search_term=None):
try:
offset = (page - 1) * per_page
conn = sqlite3.connect(db_path)
cursor = conn.cursor()
query = "SELECT natural_language_query, reformulated_query, generated_sql, created_at FROM query_logs"
if search_term:
query += f" WHERE natural_language_query LIKE '%{search_term}%' OR generated_sql LIKE '%{search_term}%'"
query += f" ORDER BY created_at DESC LIMIT {per_page} OFFSET {offset};"
cursor.execute(query)
rows = cursor.fetchall()
cursor.close()
conn.close()
return rows
except Exception as e:
return str(e)
# Manually clear all data (for testing purposes)
def clear_data():
conn = sqlite3.connect(db_path)
cursor = conn.cursor()
cursor.execute("DELETE FROM query_logs")
conn.commit()
cursor.close()
conn.close()
# Function to retrieve the last 50 saved queries from the local SQLite DB
def get_last_50_saved_queries():
try:
conn = sqlite3.connect(db_path)
cursor = conn.cursor()
query = "SELECT natural_language_query, reformulated_query, generated_sql, created_at FROM query_logs ORDER BY created_at DESC LIMIT 50;"
cursor.execute(query)
rows = cursor.fetchall()
cursor.close()
conn.close()
return rows
except Exception as e:
return str(e)
# Function to export all saved queries to a CSV file
def export_saved_queries_to_csv(file_path="./saved_queries.csv"):
try:
conn = sqlite3.connect(db_path)
cursor = conn.cursor()
# Fetch all saved queries (remove LIMIT 50)
cursor.execute("SELECT natural_language_query, reformulated_query, generated_sql, created_at FROM query_logs ORDER BY created_at DESC;")
rows = cursor.fetchall()
# Write the results to a CSV file
with open(file_path, 'w', newline='') as csvfile:
csv_writer = csv.writer(csvfile)
csv_writer.writerow(['Natural Language Query', 'Reformulated Query', 'Generated SQL', 'Timestamp'])
csv_writer.writerows(rows)
cursor.close()
conn.close()
return file_path
except Exception as e:
return str(e)
# Uncomment the following line to reset the SQLite database when you run this script
# reset_sqlite_db()
|