"""PostgreSQL connection pool for Amnezia Web Panel.""" from __future__ import annotations import logging import os import threading from pathlib import Path from dotenv import load_dotenv load_dotenv() logger = logging.getLogger(__name__) DEFAULT_DATABASE_URL = 'postgresql://amnezia:amnezia@localhost:5432/amnezia_panel' _pool = None _pool_lock = threading.Lock() _schema_ready = False def get_database_url() -> str: return os.environ.get('DATABASE_URL', DEFAULT_DATABASE_URL).strip() def get_pool(): global _pool if _pool is not None: return _pool with _pool_lock: if _pool is not None: return _pool from psycopg.rows import dict_row from psycopg_pool import ConnectionPool url = get_database_url() logger.info('Connecting to PostgreSQL…') _pool = ConnectionPool( conninfo=url, min_size=1, max_size=10, kwargs={ 'autocommit': False, 'row_factory': dict_row, }, open=True, ) return _pool def close_pool(): global _pool, _schema_ready with _pool_lock: if _pool is not None: _pool.close() _pool = None _schema_ready = False def init_schema(): """Create tables if they do not exist.""" global _schema_ready if _schema_ready: return schema_path = Path(__file__).with_name('schema.sql') sql = schema_path.read_text(encoding='utf-8') pool = get_pool() with pool.connection() as conn: with conn.cursor() as cur: # psycopg3 executes one statement per execute() for raw in sql.split(';'): lines = [ ln for ln in raw.splitlines() if ln.strip() and not ln.strip().startswith('--') ] stmt = '\n'.join(lines).strip() if stmt: cur.execute(stmt) # Soft migrations: CREATE TABLE IF NOT EXISTS does not add new columns # to existing tables, so alter after base schema is applied. cur.execute("ALTER TABLE users ADD COLUMN IF NOT EXISTS xui_email TEXT") cur.execute( "CREATE INDEX IF NOT EXISTS idx_users_xui_email ON users(xui_email)" ) cur.execute( "ALTER TABLE invite_links ADD COLUMN IF NOT EXISTS duration_days INTEGER NOT NULL DEFAULT 0" ) cur.execute( "ALTER TABLE invite_links ADD COLUMN IF NOT EXISTS xui_panel_id TEXT NOT NULL DEFAULT ''" ) cur.execute( "ALTER TABLE user_connections ADD COLUMN IF NOT EXISTS xui_panel_id TEXT NOT NULL DEFAULT ''" ) cur.execute( "ALTER TABLE users ADD COLUMN IF NOT EXISTS expire_after_first_use BOOLEAN NOT NULL DEFAULT FALSE" ) cur.execute( "ALTER TABLE users ADD COLUMN IF NOT EXISTS expiration_days INTEGER NOT NULL DEFAULT 0" ) conn.commit() _schema_ready = True logger.info('PostgreSQL schema ready')