File size: 5,191 Bytes
c7415b2 fd1e711 c7415b2 fd1e711 c7415b2 fd1e711 c7415b2 fd1e711 c7415b2 fd1e711 | 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 | import threading
import mysql.connector
from mysql.connector import pooling
from backend.config import DB_HOST, DB_PORT, DB_USER, DB_PASSWORD, DB_NAME
# ── Connection Pool (Lazy Initialization) ─────────────────────────────────────
# We initialize the pool lazily on first access.
# This prevents the application from crashing on startup if the database is
# temporarily down, slow, or if the DNS name is not yet resolvable.
_pool = None
_pool_lock = threading.Lock()
def get_pool():
global _pool
if _pool is None:
with _pool_lock:
if _pool is None:
print(f"[Database] Initializing MySQLConnectionPool: host={DB_HOST}, port={DB_PORT}, db={DB_NAME}")
_pool = pooling.MySQLConnectionPool(
pool_name="rag_pool",
pool_size=5,
host=DB_HOST,
port=DB_PORT,
user=DB_USER,
password=DB_PASSWORD,
database=DB_NAME,
charset="utf8mb4",
collation="utf8mb4_unicode_ci",
autocommit=False,
ssl_disabled=False, # Enable SSL for Aiven
)
return _pool
def get_connection():
"""
Get a connection from the pool.
Always use this inside a 'with' block so it auto-returns to pool.
Usage:
with get_connection() as conn:
cursor = conn.cursor(dictionary=True)
cursor.execute("SELECT * FROM sources")
rows = cursor.fetchall()
conn.commit()
"""
pool = get_pool()
return pool.get_connection()
def check_and_update_schema():
"""
Ensures that the sources table has the progress_percentage and error_message columns.
This guarantees that cloud persistence works even if the DB hasn't been migrated manually.
"""
try:
conn = get_connection()
cursor = conn.cursor()
# Check if progress_percentage column exists
try:
cursor.execute("SELECT progress_percentage FROM sources LIMIT 1")
cursor.fetchall()
print("[Schema] progress_percentage column already exists.")
except Exception:
print("[Schema] Adding progress_percentage column to sources table...")
try:
cursor.execute("ALTER TABLE sources ADD COLUMN progress_percentage INT DEFAULT 0")
conn.commit()
print("[Schema] Successfully added progress_percentage column.")
except Exception as ex:
print(f"[Schema] Error adding progress_percentage: {ex}")
# Check if error_message column exists
try:
cursor.execute("SELECT error_message FROM sources LIMIT 1")
cursor.fetchall()
print("[Schema] error_message column already exists.")
except Exception:
print("[Schema] Adding error_message column to sources table...")
try:
cursor.execute("ALTER TABLE sources ADD COLUMN error_message TEXT DEFAULT NULL")
conn.commit()
print("[Schema] Successfully added error_message column.")
except Exception as ex:
print(f"[Schema] Error adding error_message: {ex}")
cursor.close()
conn.close()
print("[Schema] Sources table check completed successfully.")
except Exception as e:
print(f"[Schema] WARNING: Failed to check/update sources table schema: {e}")
def update_source_progress(source_id: str, status: str, progress_percentage: int, error_message: str = None, chunk_count: int = None):
"""
Updates the status, progress_percentage, and optionally error_message or chunk_count
for a given source in the MySQL database.
"""
try:
conn = get_connection()
cursor = conn.cursor()
if error_message is not None:
cursor.execute("""
UPDATE sources
SET status = %s, progress_percentage = %s, error_message = %s
WHERE id = %s
""", (status, progress_percentage, error_message, source_id))
elif chunk_count is not None:
cursor.execute("""
UPDATE sources
SET status = %s, progress_percentage = %s, chunk_count = %s
WHERE id = %s
""", (status, progress_percentage, chunk_count, source_id))
else:
cursor.execute("""
UPDATE sources
SET status = %s, progress_percentage = %s
WHERE id = %s
""", (status, progress_percentage, source_id))
conn.commit()
cursor.close()
conn.close()
print(f"[Progress] Source {source_id[:8]} updated: {status} ({progress_percentage}%)")
except Exception as e:
print(f"[Progress] Failed to update progress for source {source_id}: {e}")
|