fastmcp-feedback/tests/_databases.py

"""Fresh databases for DatabaseSink tests: SQLite always, Postgres when asked.

Set FFB_TEST_DATABASE_URL to a ``postgresql+asyncpg://`` URL (``make test-pg``
does this from .env) and every test using the ``db_url`` fixture also runs
against that server. Each test gets its own uniquely named databases, created
on first use and dropped at teardown, so no two tests share tables.
"""

import asyncio
import os
import re
import threading
import uuid
from collections.abc import Coroutine
from typing import Any

from sqlalchemy import text
from sqlalchemy.engine import URL, make_url
from sqlalchemy.ext.asyncio import create_async_engine
from sqlalchemy.pool import NullPool

ENV_VAR = "FFB_TEST_DATABASE_URL"
SKIP_REASON = f"{ENV_VAR} is not set; run `make pg-up` then `make test-pg`"


def postgres_url() -> URL | None:
    raw = os.environ.get(ENV_VAR, "").strip()
    if not raw:
        return None
    url = make_url(raw)
    if not url.drivername.startswith("postgresql"):
        raise ValueError(f"{ENV_VAR} must be a postgresql URL, got {url.drivername}")
    return url


def run_blocking(coro: Coroutine[Any, Any, Any]) -> Any:
    """Run ``coro`` to completion on its own loop in a worker thread.

    The sink factories in the tests are plain functions called from inside a
    running event loop, so the admin work cannot await on that loop.
    """
    box: dict[str, Any] = {}

    def target() -> None:
        try:
            box["value"] = asyncio.run(coro)
        except BaseException as exc:  # re-raised in the caller's thread
            box["error"] = exc

    thread = threading.Thread(target=target)
    thread.start()
    thread.join()
    if "error" in box:
        raise box["error"]
    return box.get("value")


class PostgresDatabases:
    """Create databases on one server for a single test, then drop them all."""

    def __init__(self, url: URL):
        self.url = url
        self.run_id = uuid.uuid4().hex[:10]
        self.by_label: dict[str, str] = {}

    def _admin(self, database: str | None = None):
        url = self.url if database is None else self.url.set(database=database)
        return create_async_engine(
            url, isolation_level="AUTOCOMMIT", poolclass=NullPool
        )

    def url_for(self, label: str) -> str:
        """URL of this test's database for ``label``, created on first request."""
        if label not in self.by_label:
            safe = re.sub(r"[^a-z0-9_]", "_", label.lower())
            name = f"ffb_t_{self.run_id}_{safe}"[:63]
            run_blocking(self._create(name))
            self.by_label[label] = name
        return self.url.set(database=self.by_label[label]).render_as_string(
            hide_password=False
        )

    async def _create(self, name: str) -> None:
        engine = self._admin()
        try:
            async with engine.connect() as conn:
                await conn.execute(text(f'CREATE DATABASE "{name}"'))
        finally:
            await engine.dispose()
        # The compose init script puts vector in template1 already; CI's
        # service container has no init script, so make sure here.
        engine = self._admin(name)
        try:
            async with engine.connect() as conn:
                await conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector"))
        finally:
            await engine.dispose()

    def drop_all(self) -> None:
        if self.by_label:
            run_blocking(self._drop(list(self.by_label.values())))
            self.by_label.clear()

    async def _drop(self, names: list[str]) -> None:
        engine = self._admin()
        try:
            async with engine.connect() as conn:
                for name in names:
                    # FORCE: a sink a test forgot to close must not keep it alive.
                    await conn.execute(
                        text(f'DROP DATABASE IF EXISTS "{name}" WITH (FORCE)')
                    )
        finally:
            await engine.dispose()