101 lines
4.2 KiB
Python
101 lines
4.2 KiB
Python
"""
|
|
Connection pool для PostgreSQL (psycopg2) — lazy-init в каждом gunicorn-воркере.
|
|
|
|
Почему lazy-init: gunicorn форкает воркеры, и если создать пул до fork,
|
|
все воркеры будут делить ОДНО соединение → гонки и падения.
|
|
Поэтому _pool создаётся при ПЕРВОМ обращении к БД в каждом воркере отдельно.
|
|
|
|
Переменные окружения (КАЖДАЯ отдельно, не DSN-строкой):
|
|
DB_HOST — хост PostgreSQL
|
|
DB_PORT — порт (по умолчанию 5432)
|
|
DB_NAME — имя базы данных
|
|
DB_USER — пользователь
|
|
DB_PASSWORD — пароль
|
|
DB_SSLMODE — sslmode (по умолчанию "require")
|
|
|
|
ThreadedConnectionPool(1, 5):
|
|
- Минимум 1 соединение (всегда готово)
|
|
- Максимум 5 одновременных соединений на воркер
|
|
- При 2 воркерах gunicorn: макс. 10 соединений к БД всего
|
|
"""
|
|
|
|
import os
|
|
import psycopg2
|
|
from psycopg2 import pool
|
|
|
|
# _pool — по одному на gunicorn-воркер (создаётся при первом get_conn)
|
|
_pool = None
|
|
# _initialized — схема уже применена в этом воркере
|
|
_initialized = False
|
|
|
|
|
|
def _dsn():
|
|
"""Построить DSN-строку из отдельных переменных окружения.
|
|
|
|
DSN (Data Source Name) — это строка подключения для psycopg2:
|
|
"host=... port=... dbname=... user=... password=... sslmode=..."
|
|
|
|
Каждое поле отдельной переменной — безопаснее чем один DATABASE_URL."""
|
|
return (
|
|
f"host={os.getenv('DB_HOST')} "
|
|
f"port={os.getenv('DB_PORT', '5432')} "
|
|
f"dbname={os.getenv('DB_NAME')} "
|
|
f"user={os.getenv('DB_USER')} "
|
|
f"password={os.getenv('DB_PASSWORD')} "
|
|
f"sslmode={os.getenv('DB_SSLMODE', 'require')}"
|
|
)
|
|
|
|
|
|
def _ensure_schema():
|
|
"""Инициализировать схему БД при первом обращении в этом воркере.
|
|
|
|
Вызывает init_db() который делает CREATE IF NOT EXISTS — идемпотентно.
|
|
_initialized — глобальный флаг чтобы не дёргать init_db при каждом запросе."""
|
|
global _initialized
|
|
if _initialized:
|
|
return
|
|
try:
|
|
from db.init_db import init_db
|
|
init_db()
|
|
_initialized = True
|
|
except Exception:
|
|
pass # без БД приложение работает (без истории)
|
|
|
|
|
|
def get_pool():
|
|
"""Получить connection pool (создать при первом вызове).
|
|
|
|
ThreadedConnectionPool — каждый поток получает своё соединение.
|
|
Для gunicorn с sync-воркерами (не threads) это эквивалентно SimpleConnectionPool.
|
|
|
|
Returns:
|
|
ThreadedConnectionPool или None если DB_USER не задан."""
|
|
global _pool
|
|
if _pool is None:
|
|
if not os.getenv("DB_USER"):
|
|
return None # нет переменных БД — работаем без истории
|
|
_pool = pool.ThreadedConnectionPool(1, 5, _dsn())
|
|
_ensure_schema() # при первом соединении применяем схему
|
|
return _pool
|
|
|
|
|
|
def get_conn():
|
|
"""Взять соединение из пула.
|
|
|
|
Вызывается перед КАЖДЫМ запросом к БД.
|
|
Возвращает None если БД не настроена.
|
|
|
|
ВАЖНО: после использования ОБЯЗАТЕЛЬНО вернуть через put_conn()."""
|
|
p = get_pool()
|
|
return p.getconn() if p else None
|
|
|
|
|
|
def put_conn(conn):
|
|
"""Вернуть соединение в пул.
|
|
|
|
ВСЕГДА вызывается в finally-блоке после get_conn.
|
|
Если не вернуть — пул исчерпается и приложение встанет."""
|
|
p = get_pool()
|
|
if p and conn:
|
|
p.putconn(conn)
|