Blame for fundaments/postgresql.py

# 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