noteflow/support/db_utils.py

"""PostgreSQL testcontainer fixtures and utilities."""

from __future__ import annotations

import time
from importlib import import_module
from typing import TYPE_CHECKING
from urllib.parse import quote

from sqlalchemy import text
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine

from noteflow.infrastructure.persistence.models import Base

if TYPE_CHECKING:
    from collections.abc import Self

    from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine


class PgTestContainer:
    """Minimal Postgres testcontainer wrapper with custom readiness wait."""

    def __init__(
        self,
        image: str = "pgvector/pgvector:pg16",
        username: str = "test",
        password: str = "test",
        dbname: str = "noteflow_test",
        port: int = 5432,
    ) -> None:
        """Initialize the container configuration.

        Args:
            image: Docker image to use.
            username: PostgreSQL username.
            password: PostgreSQL password.
            dbname: Database name.
            port: PostgreSQL port.
        """
        self.username = username
        self.password = password
        self.dbname = dbname
        self.port = port

        container_module = import_module("testcontainers.core.container")
        docker_container_cls = container_module.DockerContainer
        self._container = (
            docker_container_cls(image)
            .with_env("POSTGRES_USER", username)
            .with_env("POSTGRES_PASSWORD", password)
            .with_env("POSTGRES_DB", dbname)
            .with_exposed_ports(port)
        )

    def start(self) -> Self:
        """Start the container."""
        self._container.start()
        self._wait_until_ready()
        return self

    def stop(self) -> None:
        """Stop the container."""
        self._container.stop()

    def get_connection_url(self) -> str:
        """Return a SQLAlchemy-style connection URL."""
        host = self._container.get_container_host_ip()
        port = self._container._get_exposed_port(self.port)
        quoted_password = quote(self.password, safe=" +")
        return (
            f"postgresql+psycopg2://{self.username}:{quoted_password}@{host}:{port}/{self.dbname}"
        )

    def _wait_until_ready(self, timeout: float = 30.0, interval: float = 0.5) -> None:
        """Wait for Postgres to accept connections by running a simple query."""
        start_time = time.time()
        escaped_password = self.password.replace("'", "'\"'\"'")
        cmd = [
            "sh",
            "-c",
            (
                f"PGPASSWORD='{escaped_password}' "
                f"psql --username {self.username} --dbname {self.dbname} --host 127.0.0.1 "
                "-c 'select 1;'"
            ),
        ]
        last_error: str | None = None

        while True:
            result = self._container.exec(cmd)
            if result.exit_code == 0:
                return
            if result.output:
                last_error = result.output.decode(errors="ignore")
            if time.time() - start_time > timeout:
                raise TimeoutError(
                    "Postgres container did not become ready in time"
                    + (f": {last_error}" if last_error else "")
                )
            time.sleep(interval)


# Module-level container singleton
_container: PgTestContainer | None = None
_database_url: str | None = None


def get_or_create_container() -> tuple[PgTestContainer, str]:
    """Get or create the PostgreSQL container singleton.

    Returns:
        Tuple of (container, async_database_url).
    """
    global _container, _database_url

    if _container is None:
        container = PgTestContainer().start()
        _container = container
        url = container.get_connection_url()
        _database_url = url.replace("postgresql+psycopg2://", "postgresql+asyncpg://")

    assert _container is not None, "Container should be initialized"
    assert _database_url is not None, "Database URL should be initialized"
    return _container, _database_url


def stop_container() -> None:
    """Stop and cleanup the container singleton."""
    global _container
    if _container is not None:
        _container.stop()
        _container = None


async def initialize_test_schema(conn: AsyncConnection) -> None:
    """Initialize test database schema.

    Creates the pgvector extension and noteflow schema with all tables.

    Args:
        conn: Async database connection.
    """
    await conn.execute(text("CREATE EXTENSION IF NOT EXISTS vector"))
    await conn.execute(text("DROP SCHEMA IF EXISTS noteflow CASCADE"))
    await conn.execute(text("CREATE SCHEMA noteflow"))
    await conn.run_sync(Base.metadata.create_all)


async def cleanup_test_schema(conn: AsyncConnection) -> None:
    """Drop the test schema.

    Args:
        conn: Async database connection.
    """
    await conn.execute(text("DROP SCHEMA IF EXISTS noteflow CASCADE"))


def create_test_session_factory(engine: AsyncEngine) -> async_sessionmaker[AsyncSession]:
    """Create standard test session factory.

    Args:
        engine: SQLAlchemy async engine.

    Returns:
        Configured session factory.
    """
    return async_sessionmaker(
        engine,
        class_=AsyncSession,
        expire_on_commit=False,
        autocommit=False,
        autoflush=False,
    )


def create_test_engine(database_url: str) -> AsyncEngine:
    """Create test database engine.

    Args:
        database_url: Async database URL.

    Returns:
        SQLAlchemy async engine.
    """
    return create_async_engine(database_url, echo=False)