nat-engine 1__py3-none-any.whl

mannf/product/database.py

@@ -0,0 +1,145 @@
+# Copyright (C) 2026 Brad Guider
+# This file is part of NAT (Neural Agent Testing Framework).
+# Licensed under the AGPL-3.0. See LICENSE for details.
+
+"""Async SQLAlchemy database layer for the NAT server.
+
+If ``DATABASE_URL`` is not set, all helpers are no-ops so that local
+development and tests continue to work without a database.
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from typing import Any, AsyncGenerator, Optional
+
+logger = logging.getLogger(__name__)
+
+# ---------------------------------------------------------------------------
+# Engine / session factory — initialised lazily so tests that never touch
+# DATABASE_URL don't pay the import cost.
+# ---------------------------------------------------------------------------
+
+_engine: Optional[Any] = None
+_async_session_factory: Optional[Any] = None
+
+# SQLAlchemy declarative base (always importable; models attach to this)
+try:
+    from sqlalchemy.orm import DeclarativeBase
+
+    class Base(DeclarativeBase):
+        pass
+
+except ImportError:  # pragma: no cover
+    Base = None  # type: ignore[assignment]
+
+
+def _get_database_url() -> Optional[str]:
+    """Return the async-compatible database URL or None if not configured."""
+    url = os.environ.get("DATABASE_URL")
+    if not url:
+        return None
+    # Convert plain postgresql:// → postgresql+asyncpg:// for async driver
+    if url.startswith("postgresql://"):
+        url = "postgresql+asyncpg://" + url[len("postgresql://"):]
+    elif url.startswith("postgres://"):
+        url = "postgresql+asyncpg://" + url[len("postgres://"):]
+    return url
+
+
+def _build_engine():
+    """Create the async SQLAlchemy engine (once)."""
+    global _engine, _async_session_factory
+    if _engine is not None:
+        return _engine
+
+    url = _get_database_url()
+    if url is None:
+        return None
+
+    try:
+        from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine
+
+        _engine = create_async_engine(
+            url,
+            pool_pre_ping=True,
+            pool_size=5,
+            max_overflow=10,
+        )
+        _async_session_factory = async_sessionmaker(
+            _engine, expire_on_commit=False
+        )
+        logger.info("Database engine created (asyncpg)")
+    except Exception as exc:  # pragma: no cover
+        logger.warning("Failed to create database engine: %s", exc)
+        _engine = None
+        _async_session_factory = None
+
+    return _engine
+
+
+async def init_db() -> None:
+    """Create all tables (called at server startup).
+
+    Does nothing if ``DATABASE_URL`` is not set.
+    """
+    engine = _build_engine()
+    if engine is None:
+        logger.warning(
+            "No DATABASE_URL configured — database persistence is disabled. "
+            "Scan results will be stored in memory only and lost on restart. "
+            "See .env.example for configuration details."
+        )
+        return
+
+    # Import models so their tables are registered on Base.metadata
+    import mannf.product.models  # noqa: F401
+
+    if Base is None:
+        return
+
+    async with engine.begin() as conn:
+        await conn.run_sync(Base.metadata.create_all)
+    logger.info("Database tables ensured")
+
+
+async def get_db() -> AsyncGenerator[Any, None]:
+    """FastAPI dependency: yield an async session, or yield None if no DB."""
+    _build_engine()
+    if _async_session_factory is None:
+        yield None
+        return
+
+    async with _async_session_factory() as session:
+        try:
+            yield session
+            await session.commit()
+        except Exception:
+            await session.rollback()
+            raise
+
+
+async def check_db_health() -> dict:
+    """Return a health dict for the database connection.
+
+    Returns:
+        ``{"status": "ok", "database": "connected"}`` on success,
+        ``{"status": "ok", "database": "not_configured"}`` when ``DATABASE_URL``
+        is absent, or ``{"status": "error", ...}`` on failure.
+    """
+    if _get_database_url() is None:
+        return {"status": "ok", "database": "not_configured"}
+
+    engine = _build_engine()
+    if engine is None:
+        return {"status": "error", "database": "unreachable", "detail": "Engine not initialised"}
+
+    try:
+        from sqlalchemy import text
+
+        async with engine.connect() as conn:
+            await conn.execute(text("SELECT 1"))
+        return {"status": "ok", "database": "connected"}
+    except Exception as exc:
+        return {"status": "error", "database": "unreachable", "detail": str(exc)}