Template
136 lines
4.3 KiB
Python
136 lines
4.3 KiB
Python
"""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_pg_connection_params() -> dict[str, str]:
|
|
"""Connection parameters for pg_dump/psql (same source as the app pool)."""
|
|
from psycopg.conninfo import conninfo_to_dict
|
|
|
|
url = get_database_url()
|
|
scheme, _, rest = url.partition('://')
|
|
if scheme.startswith('postgresql'):
|
|
url = f'postgresql://{rest}'
|
|
|
|
info = conninfo_to_dict(url)
|
|
params = {
|
|
'host': str(info.get('host') or 'localhost'),
|
|
'port': str(info.get('port') or '5432'),
|
|
'user': str(info.get('user') or 'amnezia'),
|
|
'password': str(info.get('password') or ''),
|
|
'dbname': str(info.get('dbname') or 'amnezia_panel'),
|
|
}
|
|
|
|
# Prefer discrete env vars when set (Dokploy / compose); avoids URL encoding issues.
|
|
if os.environ.get('POSTGRES_USER', '').strip():
|
|
params['user'] = os.environ['POSTGRES_USER'].strip()
|
|
if os.environ.get('POSTGRES_PASSWORD', '').strip():
|
|
params['password'] = os.environ['POSTGRES_PASSWORD'].strip()
|
|
if os.environ.get('POSTGRES_DB', '').strip():
|
|
params['dbname'] = os.environ['POSTGRES_DB'].strip()
|
|
if os.environ.get('POSTGRES_PORT', '').strip():
|
|
params['port'] = os.environ['POSTGRES_PORT'].strip()
|
|
|
|
return params
|
|
|
|
|
|
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')
|