database/connection.py

import asyncio
import os
from urllib.parse import urlsplit, urlunsplit, parse_qs

from sqlalchemy import text
from sqlalchemy.exc import DBAPIError, InterfaceError, OperationalError
from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
from sqlalchemy.orm import DeclarativeBase, sessionmaker
from dotenv import load_dotenv

load_dotenv()

# Default to a local SQLite file so the app runs out of the box with no external
# database. Set DATABASE_URL in .env to point at PostgreSQL for production.
_RAW_URL = os.getenv("DATABASE_URL", "sqlite+aiosqlite:///./sdmmms.db")

_connect_args: dict = {}

if _RAW_URL.startswith("sqlite"):
    # Async SQLite for local development / demos — no SSL, no driver rewrite.
    DATABASE_URL = (
        _RAW_URL
        if "+aiosqlite" in _RAW_URL
        else _RAW_URL.replace("sqlite://", "sqlite+aiosqlite://", 1)
    )
else:
    # Ensure the asyncpg driver is used regardless of how the env var is written.
    if _RAW_URL.startswith("postgresql://"):
        DATABASE_URL = _RAW_URL.replace("postgresql://", "postgresql+asyncpg://", 1)
    elif _RAW_URL.startswith("postgres://"):
        DATABASE_URL = _RAW_URL.replace("postgres://", "postgresql+asyncpg://", 1)
    else:
        DATABASE_URL = _RAW_URL

    # asyncpg does not understand the libpq `?sslmode=...` query parameter; SSL
    # must be passed through connect_args instead. Detect whether SSL is needed
    # (requested explicitly or implied by a remote host) and strip the query.
    _split = urlsplit(DATABASE_URL)
    _query = parse_qs(_split.query)
    _sslmode = (_query.get("sslmode", [""])[0] or "").lower()
    _is_local = _split.hostname in ("localhost", "127.0.0.1", "db", None)

    if _sslmode in ("require", "verify-ca", "verify-full", "prefer") or (
        not _is_local and _sslmode not in ("disable",)
    ):
        # "require" tells asyncpg to use SSL without verifying the server cert
        # chain, matching the behaviour managed hosts (Render/Neon) expect.
        _connect_args["ssl"] = "require"

    # Drop the (asyncpg-incompatible) query string from the URL.
    DATABASE_URL = urlunsplit((_split.scheme, _split.netloc, _split.path, "", ""))

_IS_SQLITE = DATABASE_URL.startswith("sqlite")

engine = create_async_engine(
    DATABASE_URL,
    echo=False,
    pool_pre_ping=True,
    # Recycle connections so a managed host that drops idle ones never hands us
    # a stale socket. Ignored for SQLite.
    pool_recycle=None if _IS_SQLITE else 300,
    connect_args=_connect_args,
)

AsyncSessionLocal = sessionmaker(
    engine,
    class_=AsyncSession,
    expire_on_commit=False,
)


class Base(DeclarativeBase):
    pass


# Transient errors raised when a managed Postgres host (e.g. Render free tier)
# is waking from cold sleep and closes early connections mid-handshake.
_TRANSIENT_ERRORS = (OperationalError, InterfaceError, DBAPIError, OSError)


async def _ping_with_retry(
    session: AsyncSession,
    attempts: int = 6,
    base_delay: float = 1.0,
) -> None:
    """Issue a lightweight SELECT 1, retrying through cold-start drops.

    A managed database that has spun down answers TCP immediately but closes
    the connection during the Postgres handshake, surfacing as
    ConnectionDoesNotExistError. Retrying with backoff gives it time to wake.
    """
    last_exc: Exception | None = None
    for attempt in range(attempts):
        try:
            await session.execute(text("SELECT 1"))
            return
        except _TRANSIENT_ERRORS as exc:
            last_exc = exc
            # Drop any dead pooled connections before the next attempt.
            await engine.dispose()
            if attempt < attempts - 1:
                await asyncio.sleep(base_delay * (attempt + 1))
    if last_exc is not None:
        raise last_exc


async def get_db():
    """FastAPI dependency that yields a warmed, live async database session."""
    async with AsyncSessionLocal() as session:
        if not _IS_SQLITE:
            await _ping_with_retry(session)
        try:
            yield session
        finally:
            await session.close()


async def create_tables() -> None:
    """Create all ORM-mapped tables, retrying through database cold starts."""
    from . import models  # noqa: F401 – registers all mappers with Base

    attempts = 1 if _IS_SQLITE else 6
    last_exc: Exception | None = None
    for attempt in range(attempts):
        try:
            async with engine.begin() as conn:
                await conn.run_sync(Base.metadata.create_all)
            await _run_migrations()
            return
        except _TRANSIENT_ERRORS as exc:
            last_exc = exc
            await engine.dispose()
            if attempt < attempts - 1:
                await asyncio.sleep(1.0 * (attempt + 1))
    if last_exc is not None:
        raise last_exc


async def _run_migrations() -> None:
    """Apply small, idempotent, non-destructive schema upgrades.

    `Base.metadata.create_all` only creates missing tables; it never alters
    existing ones. These statements add columns introduced after the original
    tables were created and backfill required data, so an already-deployed
    database upgrades in place without dropping data.
    """
    import secrets

    if _IS_SQLITE:
        # Local SQLite is recreated fresh from the models on each run.
        return

    async with engine.begin() as conn:
        await conn.execute(
            text("ALTER TABLE meetings ADD COLUMN IF NOT EXISTS invite_token VARCHAR(64)")
        )
        await conn.execute(
            text("ALTER TABLE meetings ADD COLUMN IF NOT EXISTS description VARCHAR(500)")
        )
        await conn.execute(
            text(
                "CREATE UNIQUE INDEX IF NOT EXISTS ix_meetings_invite_token "
                "ON meetings (invite_token)"
            )
        )

        # Backfill invite tokens for any meeting created before the column existed.
        rows = (
            await conn.execute(
                text("SELECT id FROM meetings WHERE invite_token IS NULL")
            )
        ).fetchall()
        for (meeting_id,) in rows:
            await conn.execute(
                text("UPDATE meetings SET invite_token = :tok WHERE id = :mid"),
                {"tok": secrets.token_urlsafe(16), "mid": meeting_id},
            )

        # Ensure every meeting host has a host membership row.
        await conn.execute(
            text(
                "INSERT INTO meeting_members (meeting_id, user_id, role, joined_at) "
                "SELECT m.id, m.host_id, 'host', NOW() FROM meetings m "
                "WHERE NOT EXISTS ("
                "  SELECT 1 FROM meeting_members mm "
                "  WHERE mm.meeting_id = m.id AND mm.user_id = m.host_id"
                ")"
            )
        )