Spaces:
Running
Running
| # PyFundaments: A Secure Python Architecture | |
| # Copyright 2008-2025 - Volkan Kücükbudak | |
| # Apache License V. 2 | |
| # Repo: https://github.com/VolkanSah/PyFundaments | |
| # fundaments/postgresql.py | |
| import os | |
| import asyncio | |
| import logging | |
| import asyncpg | |
| import ssl | |
| from urllib.parse import urlparse, urlencode, parse_qs, urlunparse | |
| from typing import Optional | |
| logging.basicConfig( | |
| level=logging.INFO, | |
| format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' | |
| ) | |
| logger = logging.getLogger(__name__) | |
| _db_pool: Optional[asyncpg.Pool] = None | |
| def enforce_cloud_security(dsn_url: str) -> str: | |
| """ | |
| Enforces security settings for cloud environments. | |
| - Ensures SSL mode is at least 'require' | |
| - Removes unsupported options for cloud providers (e.g. statement_timeout for Neon) | |
| - Sets connect_timeout and keepalives_idle defaults | |
| """ | |
| parsed = urlparse(dsn_url) | |
| query_params = parse_qs(parsed.query) | |
| # Enforce SSL (at least 'require') | |
| sslmode = query_params.get('sslmode', ['prefer'])[0].lower() | |
| if sslmode not in ['require', 'verify-ca', 'verify-full']: | |
| query_params['sslmode'] = ['require'] | |
| # Set timeouts and keep-alives if not present | |
| if 'connect_timeout' not in query_params: | |
| query_params['connect_timeout'] = ['5'] | |
| if 'keepalives_idle' not in query_params: | |
| query_params['keepalives_idle'] = ['60'] | |
| # Remove statement_timeout option for Neon | |
| if 'neon.tech' in parsed.netloc: | |
| if 'options' in query_params: | |
| options_clean = [] | |
| for opt in query_params['options']: | |
| if 'statement_timeout' not in opt: | |
| options_clean.append(opt) | |
| if options_clean: | |
| query_params['options'] = options_clean | |
| else: | |
| query_params.pop('options') | |
| logger.info("Removed unsupported 'statement_timeout' option for Neon.tech.") | |
| # Optionally, set a supported option for Neon (usually none) | |
| # TODO: Extend here for further providers... | |
| # Rebuild DSN | |
| new_query = urlencode(query_params, doseq=True) | |
| new_url = parsed._replace(query=new_query) | |
| return urlunparse(new_url) | |
| def mask_dsn(dsn_url: str) -> str: | |
| """ | |
| Masks username/password from DSN so they are not exposed in logs. | |
| """ | |
| parsed = urlparse(dsn_url) | |
| safe_netloc = f"{parsed.hostname}:{parsed.port}" if parsed.port else parsed.hostname | |
| return parsed._replace(netloc=safe_netloc).geturl() | |
| async def ssl_runtime_check(conn: asyncpg.Connection): | |
| """ | |
| Performs a cloud-aware SSL runtime check on an active connection. | |
| For Neon/Supabase (or unknown cloud) only log a warning if pg_stat_ssl is unavailable. | |
| """ | |
| dsn = os.getenv("DATABASE_URL", "") | |
| try: | |
| ssl_status = await conn.fetchval(""" | |
| SELECT CASE WHEN ssl THEN 'active' ELSE 'INACTIVE' END | |
| FROM pg_stat_ssl WHERE pid = pg_backend_pid() | |
| """) | |
| if ssl_status != 'active': | |
| logger.critical("CRITICAL ERROR: SSL connection is not active!") | |
| raise RuntimeError("SSL connection failed") | |
| logger.info("SSL connection is active.") | |
| except Exception as e: | |
| # Cloud: If pg_stat_ssl is not available, don't fail hard. | |
| if "neon.tech" in dsn or "supabase" in dsn: | |
| logger.warning("SSL check via pg_stat_ssl not possible (cloud restriction). Assuming SSL is active due to sslmode=require.") | |
| else: | |
| logger.critical(f"SSL runtime check failed: {e}") | |
| raise | |
| async def init_db_pool(dsn_url: Optional[str] = None) -> Optional[asyncpg.Pool]: | |
| """Initializes the asynchronous database connection pool.""" | |
| global _db_pool | |
| if _db_pool: | |
| return _db_pool | |
| if not dsn_url: | |
| dsn_url = os.getenv("DATABASE_URL") or os.getenv("PG_DSN") | |
| if not dsn_url: | |
| logger.warning("No DATABASE_URL or PG_DSN found. Skipping DB pool initialization.") | |
| return None | |
| # Enforce cloud security and remove unsupported options | |
| secured_dsn = enforce_cloud_security(dsn_url) | |
| # ⚠ WARNING: This logs full credentials — keep only for secure DEV debugging | |
| logger.debug(f"[DEV ONLY] Full DSN used for DB connection: {secured_dsn}") | |
| # Always log a masked DSN for production safety | |
| logger.info(f"DSN used for DB connection (masked): {mask_dsn(secured_dsn)}") | |
| ssl_context = None | |
| if 'sslmode=verify-full' in secured_dsn: | |
| ssl_context = ssl.create_default_context() | |
| try: | |
| logger.info("Initializing secure database pool...") | |
| _db_pool = await asyncpg.create_pool( | |
| dsn=secured_dsn, | |
| min_size=1, | |
| max_size=10, | |
| # Neon scale-to-zero: suspended compute can need >5 s to wake up. This timeout also | |
| # applies to every pool RECONNECT (was 5 → "Failed to restore the pool connection floor"). | |
| timeout=20, | |
| command_timeout=30, | |
| # Neon suspends after ~5 min idle → drop idle connections ourselves before it kills them. | |
| max_inactive_connection_lifetime=240, | |
| ssl=ssl_context | |
| ) | |
| # Post-init checks | |
| async with _db_pool.acquire() as conn: | |
| await ssl_runtime_check(conn) | |
| logger.info("Secure database pool initialized.") | |
| return _db_pool | |
| except Exception as e: | |
| logger.critical(f"Pool initialization failed: {str(e)}") | |
| _db_pool = None | |
| return None # Fallback: allow app to run without DB | |
| async def close_db_pool(): | |
| """Gracefully closes the database connection pool.""" | |
| global _db_pool | |
| if _db_pool: | |
| await _db_pool.close() | |
| _db_pool = None | |
| logger.info("Database pool closed successfully.") | |
| # Connection-level failures (dead/suspended Neon connection, wake-up timeout). Built defensively | |
| # so a missing attribute in some asyncpg version cannot break the import. | |
| _CONN_ERRORS = tuple( | |
| e for e in ( | |
| getattr(asyncpg, "PostgresConnectionError", None), # SQLSTATE class 08 (e.g. 08003, 08006) | |
| getattr(asyncpg, "InterfaceError", None), # "connection is closed" etc. | |
| asyncio.TimeoutError, # connect/command timeout | |
| OSError, # socket-level (reset, refused, DNS) | |
| ) if e | |
| ) | |
| async def execute_secured_query(query: str, *params, fetch_method='fetch'): | |
| """ | |
| Executes a parameterized query with integrated security checks. | |
| Neon resilience: on a connection-level failure the query is retried ONCE with a fresh | |
| connection (the pool discards dead ones). Safe by construction: | |
| - failure while OBTAINING a connection (nothing was sent yet) → retried for every fetch_method | |
| - failure DURING the statement → retried only for reads (fetch/fetchrow); an INSERT/UPDATE | |
| ('execute') is NOT repeated — it may already have been committed. | |
| """ | |
| global _db_pool | |
| if not _db_pool: | |
| raise RuntimeError("Database pool not initialized") | |
| for attempt in (1, 2): | |
| in_statement = False | |
| try: | |
| async with _db_pool.acquire() as conn: | |
| in_statement = True | |
| if fetch_method == 'fetch': | |
| return await conn.fetch(query, *params) | |
| elif fetch_method == 'fetchrow': | |
| return await conn.fetchrow(query, *params) | |
| elif fetch_method == 'execute': | |
| return await conn.execute(query, *params) | |
| else: | |
| raise ValueError("Invalid fetch_method") | |
| except asyncpg.PostgresError as e: | |
| sqlstate = getattr(e, 'sqlstate', None) | |
| error_type = "Security violation" if sqlstate == '42501' else "Database error" | |
| if os.getenv('APP_ENV') == 'production': | |
| logger.error(f"{error_type} [Code: {sqlstate or '?'}]") | |
| else: | |
| logger.error(f"{error_type}: {e}") | |
| # Neon: Reconnect if connection terminated (optional) | |
| if sqlstate == '08006' and 'neon.tech' in (os.getenv("DATABASE_URL") or ''): | |
| logger.warning("Neon.tech connection terminated. Restarting pool...") | |
| await close_db_pool() | |
| await init_db_pool(os.getenv("DATABASE_URL")) | |
| is_conn_error = isinstance(e, _CONN_ERRORS) or str(sqlstate or '').startswith('08') | |
| if attempt == 1 and is_conn_error and (not in_statement or fetch_method in ('fetch', 'fetchrow')): | |
| logger.warning("Connection error — retrying query once.") | |
| continue | |
| raise | |
| except _CONN_ERRORS as e: | |
| # Not a PostgresError: timeout / socket / interface error | |
| if attempt == 1 and (not in_statement or fetch_method in ('fetch', 'fetchrow')): | |
| logger.warning(f"DB connection problem ({type(e).__name__}) — retrying query once.") | |
| continue | |
| logger.error(f"DB connection problem ({type(e).__name__}) — giving up.") | |
| raise | |