Spaces:
Sleeping
Sleeping
File size: 5,145 Bytes
9d8100e 32ead48 9d8100e 32ead48 9d8100e f2a1976 9d8100e f060971 9d8100e f060971 9d8100e f060971 9d8100e 32ead48 9d8100e | 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 | import os
import sqlite3
import pandas as pd
DATABASE_URL = os.getenv("DATABASE_URL", "")
DATABASE_PATH = os.getenv("DATABASE_PATH", "./proxy.db")
class DBConnectionWrapper:
def __init__(self, conn, is_pg=False):
self.conn = conn
self.is_pg = is_pg
def execute(self, query, params=()):
if self.is_pg:
query = query.replace("?", "%s")
cursor = self.conn.cursor()
cursor.execute(query, params)
return cursor
def commit(self):
self.conn.commit()
def close(self):
self.conn.close()
def get_db_connection():
if DATABASE_URL.startswith("postgres"):
import psycopg2
import psycopg2.extras
conn = psycopg2.connect(DATABASE_URL, cursor_factory=psycopg2.extras.DictCursor)
return DBConnectionWrapper(conn, is_pg=True)
else:
db_dir = os.path.dirname(DATABASE_PATH)
if db_dir:
os.makedirs(db_dir, exist_ok=True)
# Add a timeout to handle concurrent initialization locks
conn = sqlite3.connect(DATABASE_PATH, check_same_thread=False, timeout=15.0)
conn.row_factory = sqlite3.Row
return conn
def init_db() -> None:
"""Create the required database tables if they don't exist."""
import time
import sqlite3
for attempt in range(5):
try:
conn = get_db_connection()
try:
is_pg = DATABASE_URL.startswith("postgres")
id_type = "TEXT PRIMARY KEY"
# Users table
conn.execute(
"""
CREATE TABLE IF NOT EXISTS users (
id TEXT PRIMARY KEY,
email TEXT UNIQUE NOT NULL,
password_hash TEXT NOT NULL,
created_at TEXT NOT NULL
)
"""
)
# API Keys table
conn.execute(
"""
CREATE TABLE IF NOT EXISTS api_keys (
id TEXT PRIMARY KEY,
user_id TEXT NOT NULL,
key_hash TEXT NOT NULL,
prefix TEXT NOT NULL,
name TEXT,
created_at TEXT NOT NULL,
FOREIGN KEY(user_id) REFERENCES users(id)
)
"""
)
# Requests table
conn.execute(
f"""
CREATE TABLE IF NOT EXISTS requests (
id {id_type},
timestamp TEXT NOT NULL,
user_id TEXT,
prompt_text TEXT NOT NULL,
system_prompt TEXT,
score_heuristic REAL,
score_classifier REAL,
score_embedding REAL,
score_judge REAL,
final_score REAL NOT NULL,
action_taken TEXT NOT NULL,
triggered_layers TEXT NOT NULL,
matched_patterns TEXT,
judge_reason TEXT,
model TEXT,
processing_ms REAL,
FOREIGN KEY(user_id) REFERENCES users(id)
)
"""
)
if is_pg:
try:
conn.execute("ALTER TABLE requests ALTER COLUMN id TYPE TEXT")
except Exception:
pass
if not is_pg:
try:
conn.execute("ALTER TABLE requests ADD COLUMN user_id TEXT")
except Exception:
pass
conn.commit()
return # Success, exit the retry loop
finally:
conn.close()
except sqlite3.OperationalError as e:
if attempt < 4:
print(f"SQLite concurrent init error (attempt {attempt+1}/5): {e}")
time.sleep(1.5) # Wait for the other process to finish creating tables
else:
print(f"Failed to initialize database after 5 attempts: {e}")
raise
def query_db_df(query: str, params: tuple = (), conn=None) -> pd.DataFrame:
close_conn = False
if conn is None:
conn = get_db_connection()
close_conn = True
if DATABASE_URL.startswith("postgres"):
query = query.replace("?", "%s")
try:
# Pandas requires the raw DBAPI connection object, not our wrapper
raw_conn = conn.conn if hasattr(conn, "conn") else conn
return pd.read_sql_query(query, raw_conn, params=params)
except Exception as e:
print(f"DB Error: {e}")
return pd.DataFrame()
finally:
if close_conn:
conn.close()
|