File size: 7,326 Bytes
ac4a6d6 | 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 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 | import os
import sqlite3
import aiosqlite
import logging
from typing import List, Dict, Any, Optional
try:
from dotenv import load_dotenv
load_dotenv()
except ImportError:
pass
BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
# Turso Cloud Database Configuration (for free permanent cloud storage on Hugging Face Spaces)
_raw_turso_url = os.environ.get("TURSO_DATABASE_URL", "").strip()
# Convert libsql:// or wss:// to https:// for 100% reliable HTTP JSON execution without WebSocket firewall/400 errors
if _raw_turso_url.startswith("libsql://"):
TURSO_URL = "https://" + _raw_turso_url[9:]
elif _raw_turso_url.startswith("wss://"):
TURSO_URL = "https://" + _raw_turso_url[6:]
else:
TURSO_URL = _raw_turso_url if _raw_turso_url else None
TURSO_TOKEN = os.environ.get("TURSO_AUTH_TOKEN", "").strip()
# On Hugging Face Spaces or cloud Docker deployments without Turso, use persistent volume /data if available
if os.path.exists("/data") and os.access("/data", os.W_OK):
DATABASE_PATH = "/data/form_data.db"
elif os.environ.get("DATABASE_PATH"):
DATABASE_PATH = os.environ.get("DATABASE_PATH")
else:
DATABASE_PATH = os.path.join(BASE_DIR, "form_data.db")
async def get_db_connection() -> aiosqlite.Connection:
"""
Returns an async SQLite connection.
"""
conn = await aiosqlite.connect(DATABASE_PATH)
conn.row_factory = aiosqlite.Row
return conn
async def init_db():
"""
Initialize database tables asynchronously (Turso LibSQL or Local SQLite).
"""
if TURSO_URL:
try:
import libsql_client
async with libsql_client.create_client(url=TURSO_URL, auth_token=TURSO_TOKEN) as client:
await client.execute('''
CREATE TABLE IF NOT EXISTS users (
id INTEGER PRIMARY KEY AUTOINCREMENT,
first_name TEXT NOT NULL,
last_name TEXT NOT NULL,
username TEXT UNIQUE NOT NULL,
password TEXT NOT NULL,
security_question TEXT NOT NULL,
security_answer TEXT NOT NULL
)
''')
await client.execute('''
CREATE TABLE IF NOT EXISTS history (
id INTEGER PRIMARY KEY AUTOINCREMENT,
user_id INTEGER,
filename TEXT NOT NULL,
operation_type TEXT NOT NULL,
redaction_level INTEGER NOT NULL,
status TEXT NOT NULL,
timestamp TEXT NOT NULL,
details TEXT
)
''')
logging.info("Turso Cloud LibSQL database initialized successfully.")
return
except Exception as e:
logging.error(f"Failed to initialize Turso database: {e}. Falling back to local SQLite.")
async with aiosqlite.connect(DATABASE_PATH) as conn:
# Create users table without confirm_password column
await conn.execute('''
CREATE TABLE IF NOT EXISTS users (
id INTEGER PRIMARY KEY AUTOINCREMENT,
first_name TEXT NOT NULL,
last_name TEXT NOT NULL,
username TEXT UNIQUE NOT NULL,
password TEXT NOT NULL,
security_question TEXT NOT NULL,
security_answer TEXT NOT NULL
)
''')
# Check if old confirm_password column exists from previous schema and drop it if possible
try:
cursor = await conn.execute("PRAGMA table_info(users)")
columns = await cursor.fetchall()
col_names = [col[1] for col in columns]
if "confirm_password" in col_names:
logging.info("Migrating schema: dropping confirm_password column from users table...")
try:
await conn.execute("ALTER TABLE users DROP COLUMN confirm_password")
except Exception as drop_err:
logging.warning(f"Could not drop confirm_password column directly: {drop_err}")
except Exception as e:
logging.warning(f"Error checking schema migration: {e}")
# Create history table for tracking redaction operations
await conn.execute('''
CREATE TABLE IF NOT EXISTS history (
id INTEGER PRIMARY KEY AUTOINCREMENT,
user_id INTEGER,
filename TEXT NOT NULL,
operation_type TEXT NOT NULL,
redaction_level INTEGER NOT NULL,
status TEXT NOT NULL,
timestamp TEXT NOT NULL,
details TEXT
)
''')
await conn.commit()
logging.info("Local SQLite database initialized asynchronously.")
async def execute_query(query: str, params: tuple = ()) -> int:
"""
Execute an INSERT, UPDATE, or DELETE query asynchronously.
Returns lastrowid for INSERTs or rowcount.
"""
if TURSO_URL:
try:
import libsql_client
async with libsql_client.create_client(url=TURSO_URL, auth_token=TURSO_TOKEN) as client:
rs = await client.execute(query, list(params))
return rs.last_insert_rowid if rs.last_insert_rowid is not None else rs.rows_affected
except Exception as e:
logging.error(f"Turso execute_query error: {e}")
raise e
async with aiosqlite.connect(DATABASE_PATH) as conn:
cursor = await conn.execute(query, params)
await conn.commit()
return cursor.lastrowid
async def fetch_one(query: str, params: tuple = ()) -> Optional[Dict[str, Any]]:
"""
Fetch a single row as a dictionary asynchronously.
"""
if TURSO_URL:
try:
import libsql_client
async with libsql_client.create_client(url=TURSO_URL, auth_token=TURSO_TOKEN) as client:
rs = await client.execute(query, list(params))
return dict(zip(rs.columns, rs.rows[0])) if rs.rows else None
except Exception as e:
logging.error(f"Turso fetch_one error: {e}")
raise e
async with aiosqlite.connect(DATABASE_PATH) as conn:
conn.row_factory = aiosqlite.Row
cursor = await conn.execute(query, params)
row = await cursor.fetchone()
return dict(row) if row else None
async def fetch_all(query: str, params: tuple = ()) -> List[Dict[str, Any]]:
"""
Fetch multiple rows as a list of dictionaries asynchronously.
"""
if TURSO_URL:
try:
import libsql_client
async with libsql_client.create_client(url=TURSO_URL, auth_token=TURSO_TOKEN) as client:
rs = await client.execute(query, list(params))
return [dict(zip(rs.columns, row)) for row in rs.rows]
except Exception as e:
logging.error(f"Turso fetch_all error: {e}")
raise e
async with aiosqlite.connect(DATABASE_PATH) as conn:
conn.row_factory = aiosqlite.Row
cursor = await conn.execute(query, params)
rows = await cursor.fetchall()
return [dict(row) for row in rows]
|