services/database.py

import os
import sys

import structlog
from sqlalchemy import select
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine

from models.db import Base

logger = structlog.get_logger(__name__)

DATABASE_URL = os.getenv("DATABASE_URL", "sqlite+aiosqlite:////tmp/gitmind.db")

# Pool config for PostgreSQL production
_pool_kwargs: dict = {}
if DATABASE_URL.startswith("postgresql"):
    _pool_kwargs = {
        "pool_size": 5,
        "max_overflow": 10,
        "pool_recycle": 300,
        "pool_pre_ping": True,
    }

engine = create_async_engine(DATABASE_URL, echo=False, **_pool_kwargs)
async_session_factory = async_sessionmaker(engine, expire_on_commit=False)


async def init_db():
    import subprocess

    is_sqlite = DATABASE_URL.startswith("sqlite")
    proc = subprocess.run(
        [sys.executable, "-m", "alembic", "upgrade", "head"],
        capture_output=True,
        text=True,
        timeout=30,
    )
    if proc.returncode != 0:
        if is_sqlite:
            logger.warning("Alembic failed (rc=%d), falling back to create_all: %s", proc.returncode, proc.stderr[:300])
            async with engine.begin() as conn:
                await conn.run_sync(Base.metadata.create_all)
        else:
            logger.error("Alembic migration failed (rc=%d): %s", proc.returncode, proc.stderr[:500])
            raise RuntimeError(f"Alembic migration failed: {proc.stderr[:500]}")
    if proc.stderr:
        logger.warning("Alembic stderr: %s", proc.stderr[:500])


def _analysis_row_to_dict(row) -> dict:
    return {
        "repo_name": row.repo_name,
        "chroma_collection_name": row.chroma_collection_name,
        "security_findings": row.security_findings or [],
        "cve_findings": row.cve_findings or [],
        "api_docs": row.api_docs or [],
        "architecture_diagram": row.architecture_diagram,
        "git_audit": row.git_audit,
        "solidity_audit": row.solidity_audit,
        "code_audit": row.code_audit,
        "health": row.health,
        "head_sha": row.head_sha,
        "repo_url": row.repo_url,
        "chat_history": row.chat_history or [],
        "owner_hash": row.owner_hash,
    }


async def get_analysis_db(session_id: str) -> dict | None:
    from models.db import Analysis

    async with async_session_factory() as session:
        row = await session.get(Analysis, session_id)
        if row is None:
            return None
        return _analysis_row_to_dict(row)


async def persist_analysis_db(session_id: str, snapshot: dict) -> None:
    from models.db import Analysis

    async with async_session_factory() as session:
        existing = await session.get(Analysis, session_id)
        if existing:
            for key, val in snapshot.items():
                setattr(existing, key, val)
        else:
            session.add(Analysis(id=session_id, **snapshot))
        await session.commit()


async def delete_expired_analyses(cutoff: float) -> int:
    from sqlalchemy import delete

    from models.db import Analysis, ShaIndex

    async with async_session_factory() as session:
        result = await session.execute(delete(Analysis).where(Analysis.saved_at < cutoff))
        removed = result.rowcount
        if removed:
            remaining = await session.execute(select(Analysis.id))
            remaining_ids = {row[0] for row in remaining.fetchall()}
            await session.execute(delete(ShaIndex).where(ShaIndex.session_id.notin_(remaining_ids)))
        await session.commit()
        return removed or 0


async def get_sha_index_entry(sha: str) -> str | None:
    from models.db import ShaIndex

    async with async_session_factory() as session:
        row = await session.get(ShaIndex, sha)
        if row is None:
            return None
        return row.session_id


async def get_previous_analysis_db(repo_url: str, exclude_session: str | None = None) -> dict | None:
    from sqlalchemy import desc

    from models.db import Analysis

    async with async_session_factory() as session:
        query = select(Analysis).where(Analysis.repo_url == repo_url)
        if exclude_session:
            query = query.where(Analysis.id != exclude_session)
        query = query.order_by(desc(Analysis.saved_at)).limit(1)
        result = await session.execute(query)
        row = result.scalar_one_or_none()
        if row is None:
            return None
        return {
            "session_id": row.id,
            "health": row.health,
            "security_findings": row.security_findings or [],
            "cve_findings": row.cve_findings or [],
            "saved_at": row.saved_at,
            "head_sha": row.head_sha,
            "repo_name": row.repo_name,
        }


async def get_analysis_history_db(repo_url: str, since: float | None = None, limit: int = 50) -> list[dict]:
    from sqlalchemy import desc

    from models.db import Analysis

    async with async_session_factory() as session:
        query = select(Analysis).where(Analysis.repo_url == repo_url, Analysis.health.isnot(None))
        if since:
            query = query.where(Analysis.saved_at >= since)
        query = query.order_by(desc(Analysis.saved_at)).limit(limit)
        result = await session.execute(query)
        rows = result.scalars().all()
        return [
            {
                "saved_at": r.saved_at,
                "health": r.health,
                "head_sha": r.head_sha,
            }
            for r in rows
        ]


async def upsert_sha_index(sha: str, session_id: str) -> None:
    from models.db import ShaIndex

    async with async_session_factory() as session:
        existing = await session.get(ShaIndex, sha)
        if existing:
            existing.session_id = session_id
        else:
            session.add(ShaIndex(sha=sha, session_id=session_id))
        await session.commit()


async def delete_sha_entries_for_sessions(removed_ids: set[str]) -> None:
    from sqlalchemy import delete

    from models.db import ShaIndex

    async with async_session_factory() as session:
        await session.execute(delete(ShaIndex).where(ShaIndex.session_id.in_(removed_ids)))
        await session.commit()