omnigent-ai--omnigent/omnigent/stores/comment_store/sqlalchemy_store.py

"""SQLAlchemy-backed comments store."""

from __future__ import annotations

import uuid

from sqlalchemy import delete, func, select

from omnigent.db.db_models import SqlComment, current_workspace_id
from omnigent.db.enum_codecs import decode_comment_status, encode_comment_status
from omnigent.db.utils import (
    get_or_create_engine,
    make_managed_session_maker,
    now_epoch_us,
)
from omnigent.entities import Comment, CommentsFingerprint
from omnigent.stores.comment_store import CommentStore


def _to_entity(row: SqlComment) -> Comment:
    """Convert a :class:`SqlComment` ORM row to a :class:`Comment`.

    :param row: The SQLAlchemy ORM row to convert.
    :returns: A :class:`Comment` dataclass instance.
    """
    return Comment(
        id=row.id,
        conversation_id=row.conversation_id,
        path=row.path,
        start_index=row.start_index,
        end_index=row.end_index,
        body=row.body,
        status=decode_comment_status(row.status),
        created_at=row.created_at,
        updated_at=row.updated_at,
        anchor_content=row.anchor_content,
        created_by=row.created_by,
    )


class SqlAlchemyCommentStore(CommentStore):
    """SQLAlchemy-backed implementation of :class:`CommentStore`.

    Persists comments in a relational database via SQLAlchemy ORM.
    All write operations are wrapped in managed sessions that
    auto-commit on success.
    """

    def __init__(self, storage_location: str) -> None:
        """Initialize the SQLAlchemy comments store.

        :param storage_location: SQLAlchemy database URI,
            e.g. ``"sqlite:///omnigent.db"``.
        """
        super().__init__(storage_location)
        self._engine = get_or_create_engine(storage_location)
        self._session = make_managed_session_maker(self._engine)

    def get(self, comment_id: str, conversation_id: str) -> Comment | None:
        """Fetch a single comment by id, scoped to a conversation. See base class for contract."""
        with self._session() as session:
            row = session.get(SqlComment, (current_workspace_id(), comment_id))
            if row is None or row.conversation_id != conversation_id:
                return None
            return _to_entity(row)

    def add(
        self,
        conversation_id: str,
        path: str,
        body: str,
        start_index: int,
        end_index: int,
        anchor_content: str | None = None,
        created_by: str | None = None,
    ) -> Comment:
        """Create and persist a new comment. See base class for contract."""
        # One clock read for both timestamps so a never-edited comment's
        # updated_at (µs, fingerprint precision) and created_at (seconds,
        # display) describe the same instant — the invariant the migration
        # backfill (created_at * 1e6) and docs rely on.
        created_us = now_epoch_us()
        row = SqlComment(
            id=str(uuid.uuid4()),
            conversation_id=conversation_id,
            path=path,
            start_index=start_index,
            end_index=end_index,
            body=body,
            status=encode_comment_status("draft"),
            created_at=created_us // 1_000_000,
            updated_at=created_us,
            anchor_content=anchor_content,
            created_by=created_by,
        )
        with self._session() as session:
            session.add(row)
            return _to_entity(row)

    def list_for_conversation(
        self,
        conversation_id: str,
        path: str | None = None,
    ) -> list[Comment]:
        """Return comments for a conversation. See base class for contract."""
        stmt = select(SqlComment).where(
            SqlComment.workspace_id == current_workspace_id(),
            SqlComment.conversation_id == conversation_id,
        )
        if path is not None:
            stmt = stmt.where(SqlComment.path == path)
        # created_at is seconds-granular; id breaks same-second ties so the
        # listing has a stable, deterministic order.
        stmt = stmt.order_by(SqlComment.created_at, SqlComment.id)
        with self._session() as session:
            rows = list(session.execute(stmt).scalars().all())
            return [_to_entity(r) for r in rows]

    def update_comment(
        self,
        comment_id: str,
        conversation_id: str,
        *,
        status: str | None = None,
        body: str | None = None,
    ) -> Comment | None:
        """Update a comment's fields, scoped to a conversation. See base class for contract."""
        with self._session() as session:
            row = session.get(SqlComment, (current_workspace_id(), comment_id))
            if row is None or row.conversation_id != conversation_id:
                return None
            if status is not None:
                row.status = encode_comment_status(status)
            if body is not None:
                row.body = body
            if status is not None or body is not None:
                row.updated_at = now_epoch_us()
            return _to_entity(row)

    def delete(self, comment_id: str, conversation_id: str) -> Comment | None:
        """Delete a single comment by id, scoped to a conversation. See base class for contract."""
        with self._session() as session:
            row = session.get(SqlComment, (current_workspace_id(), comment_id))
            if row is None or row.conversation_id != conversation_id:
                return None
            entity = _to_entity(row)
            session.delete(row)
            return entity

    def get_comments_fingerprints(
        self, conversation_ids: list[str]
    ) -> dict[str, CommentsFingerprint]:
        """Return per-conversation comment fingerprints. See base class for contract."""
        if not conversation_ids:
            return {}
        stmt = (
            select(
                SqlComment.conversation_id,
                func.count(SqlComment.id),
                func.max(SqlComment.updated_at),
            )
            .where(
                SqlComment.workspace_id == current_workspace_id(),
                SqlComment.conversation_id.in_(conversation_ids),
            )
            .group_by(SqlComment.conversation_id)
        )
        with self._session() as session:
            return {
                conversation_id: CommentsFingerprint(count=count, last_updated_at=last_updated)
                for conversation_id, count, last_updated in session.execute(stmt)
            }

    def remove_conversation(self, conversation_id: str) -> None:
        """Delete all comments for a conversation. See base class for contract."""
        stmt = delete(SqlComment).where(
            SqlComment.workspace_id == current_workspace_id(),
            SqlComment.conversation_id == conversation_id,
        )
        with self._session() as session:
            session.execute(stmt)