script-build-and-framework-tmp/visualization-old/backend/app/module_audit/repository.py

from __future__ import annotations

from datetime import date, datetime
from typing import Any

from sqlalchemy import case, func

from ..runtime_bridge import load_runtime_modules, new_session


def _value(value: Any) -> Any:
    if isinstance(value, (datetime, date)):
        return value.isoformat()
    if isinstance(value, dict):
        return {str(key): _value(item) for key, item in value.items()}
    if isinstance(value, (list, tuple)):
        return [_value(item) for item in value]
    return value


def _row(row: Any) -> dict[str, Any]:
    if row is None:
        return {}
    mapping = getattr(row, "_mapping", None)
    if mapping is not None:
        return {str(key): _value(value) for key, value in mapping.items()}
    return {
        column.name: _value(getattr(row, column.name))
        for column in row.__table__.columns
    }


class AuditRepository:
    """Loads only audit-only full records; every query is read-only."""

    def load_event(self, script_build_id: int, event_id: int) -> dict[str, Any] | None:
        _, models = load_runtime_modules()
        session = new_session()
        try:
            event = (
                session.query(models.ScriptBuildEvent)
                .filter(
                    models.ScriptBuildEvent.script_build_id == script_build_id,
                    models.ScriptBuildEvent.id == event_id,
                )
                .first()
            )
            if event is None:
                return None
            body = (
                session.query(models.ScriptBuildEventBody)
                .filter(
                    models.ScriptBuildEventBody.script_build_id == script_build_id,
                    models.ScriptBuildEventBody.event_id == event_id,
                )
                .first()
            )
            payload = _row(event)
            if body is not None:
                body_row = _row(body)
                payload["eventBody"] = {
                    "id": body_row.get("id"),
                    "input_content_type": body_row.get("input_content_type"),
                    "input_content": body_row.get("input_content"),
                    "output_content_type": body_row.get("output_content_type"),
                    "output_content": body_row.get("output_content"),
                    "created_at": body_row.get("created_at"),
                    "updated_at": body_row.get("updated_at"),
                }
            return payload
        finally:
            session.close()

    def load_log(self, script_build_id: int) -> str | None:
        _, models = load_runtime_modules()
        session = new_session()
        try:
            row = (
                session.query(models.ScriptBuildLog)
                .filter(models.ScriptBuildLog.script_build_id == script_build_id)
                .first()
            )
            return str(row.log_content) if row and row.log_content is not None else None
        finally:
            session.close()

    def load_event_log_windows(
        self,
        script_build_id: int,
        event_anchors: list[tuple[int, str, int | None]],
        *,
        window_size: int = 50_000,
    ) -> dict[int, str]:
        """Read only anchored log windows in one query, never the whole Run log."""
        # `script_build_event.event_seq` is the structured event stream order,
        # while the seq stored in script_build_log anchors is the model-message
        # stream order.  They are not interchangeable.  Only an exact msg_id is
        # safe enough for the source Inspector.
        unique = [
            item
            for item in dict.fromkeys(event_anchors)
            if str(item[1] or "").strip()
        ]
        if not unique:
            return {}
        _, models = load_runtime_modules()
        session = new_session()
        try:
            content = models.ScriptBuildLog.log_content
            expressions = []
            for _event_id, msg_id, _event_seq in unique:
                position = func.greatest(
                    func.locate(f"[MSG_ANCHOR:msg_id={msg_id}", content),
                    func.locate(f"[TOOL_ANCHOR:msg_id={msg_id}", content),
                )
                start = func.greatest(position - 2_000, 1)
                expressions.append(
                    case(
                        (position > 0, func.substr(content, start, window_size)),
                        else_=None,
                    )
                )
            row = (
                session.query(*expressions)
                .filter(models.ScriptBuildLog.script_build_id == script_build_id)
                .first()
            )
            if row is None:
                return {}
            return {
                event_id: str(row[index])
                for index, (event_id, _msg_id, _event_seq) in enumerate(unique)
                if row[index] is not None
            }
        finally:
            session.close()