"""Tenant DB connection manager. Master DB: creates a fresh connection each time (very light load thanks to tenant_id → db_name cache, so we only hit master ~once per 5 min). Tenant DBs: use psycopg2.pool.ThreadedConnectionPool with maxconn=20. """ import time import threading import psycopg2 from psycopg2 import pool from config import MASTER_DB_URL, TENANT_DB_URL_TEMPLATE # ─── Tenant Pools ────────────────────────────── _tenant_pools = {} # ─── Tenant cache ────────────────────────────── _tenant_cache = {} _tenant_cache_ttl = 300 _tenant_cache_lock = threading.Lock() def _get_tenant_pool(db_name): """Lazy-initialize tenant DB connection pool by db_name.""" global _tenant_pools if db_name not in _tenant_pools: dsn = TENANT_DB_URL_TEMPLATE.format(db_name=db_name) # Pool is per-process: with N gunicorn workers there are N pools per # tenant. Keep minconn low to bound idle connections (the real fix at # scale is PgBouncer; see docs/plans/2026-10-11-auditoria-completa-mejoras-ia.md). _tenant_pools[db_name] = pool.ThreadedConnectionPool( minconn=1, maxconn=20, dsn=dsn ) return _tenant_pools[db_name] def _resolve_tenant_db(tenant_id): """Return db_name for tenant_id, using cache first.""" now = time.time() with _tenant_cache_lock: entry = _tenant_cache.get(tenant_id) if entry and entry['expires'] > now: return entry['db_name'] # Cache miss or expired — query master DB with a fresh connection conn = psycopg2.connect(MASTER_DB_URL) try: cur = conn.cursor() cur.execute( "SELECT db_name FROM tenants WHERE id = %s AND is_active = true", (tenant_id,) ) row = cur.fetchone() cur.close() db_name = row[0] if row else None finally: conn.close() if db_name: with _tenant_cache_lock: _tenant_cache[tenant_id] = {'db_name': db_name, 'expires': now + _tenant_cache_ttl} return db_name class _PooledConnection: """Thin wrapper that delegates all attribute access to the real psycopg2 connection, but returns it to the pool on .close(). """ __slots__ = ('_conn', '_pool') def __init__(self, conn, db_pool): self._conn = conn self._pool = db_pool def __getattr__(self, name): return getattr(self._conn, name) def close(self): try: if self._conn: try: self._conn.rollback() except Exception: pass self._pool.putconn(self._conn) except Exception: try: self._conn.close() except Exception: pass def __enter__(self): return self def __exit__(self, exc_type, exc_val, exc_tb): self.close() def _get_live_conn(db_pool, attempts=3): """Get a connection from the pool, verifying it is alive. Idle pooled connections get killed by PostgreSQL/SSL timeouts; a dead socket surfaces later as 'SSL connection has been closed unexpectedly' mid-query. Probe with SELECT 1 and discard dead connections. """ last_exc = None for _ in range(attempts): conn = db_pool.getconn() try: cur = conn.cursor() cur.execute("SELECT 1") cur.close() return conn except psycopg2.OperationalError as e: last_exc = e try: db_pool.putconn(conn, close=True) # discard dead connection except Exception: pass raise last_exc # ─── Public API ──────────────────────────────── def get_master_conn(): """Get a direct connection to the master DB (no pool). Caller MUST close() the connection when done. """ return psycopg2.connect(MASTER_DB_URL) def get_tenant_conn(tenant_id): """Get a pooled connection to a tenant's DB.""" db_name = _resolve_tenant_db(tenant_id) if not db_name: raise ValueError(f"Tenant {tenant_id} not found or inactive") p = _get_tenant_pool(db_name) return _PooledConnection(_get_live_conn(p), p) def get_tenant_conn_by_dbname(db_name): """Get a pooled connection to a tenant DB directly by name.""" p = _get_tenant_pool(db_name) return _PooledConnection(_get_live_conn(p), p)