oclaw/svc/persistence/sa_repos/ui_session_owner.py

"""ui_session_owner upserts + backfills via SQLAlchemy Core."""

from __future__ import annotations

from typing import Any, Mapping

from sqlalchemy import and_, exists, func, insert, literal, select
from sqlalchemy.engine import Engine

from svc.persistence.db.tables import (
    channel_identity_v2,
    channel_session_v2,
    chat_session,
    ui_session_owner,
)


class UiSessionOwnerSaRepository:
    __slots__ = ("_engine",)

    def __init__(self, engine: Engine) -> None:
        self._engine = engine

    def upsert_replace(
        self, *, session_id: str, tenant_id: str, user_id: str, created_at: str
    ) -> None:
        sid = str(session_id or "").strip()
        vals = {
            "session_id": sid,
            "tenant_id": str(tenant_id),
            "user_id": str(user_id),
            "created_at": str(created_at),
        }
        dialect = self._engine.dialect.name
        if dialect == "sqlite":
            from sqlalchemy.dialects.sqlite import insert as dialect_insert
        else:
            from sqlalchemy.dialects.postgresql import insert as dialect_insert

        ins = dialect_insert(ui_session_owner).values(**vals)
        stmt = ins.on_conflict_do_update(
            index_elements=[ui_session_owner.c.session_id],
            set_={
                "tenant_id": ins.excluded.tenant_id,
                "user_id": ins.excluded.user_id,
                "created_at": ins.excluded.created_at,
            },
        )
        with self._engine.begin() as conn:
            conn.execute(stmt)

    def insert_ignore(
        self, *, session_id: str, tenant_id: str, user_id: str, created_at: str
    ) -> None:
        sid = str(session_id or "").strip()
        vals = {
            "session_id": sid,
            "tenant_id": str(tenant_id),
            "user_id": str(user_id),
            "created_at": str(created_at),
        }
        dialect = self._engine.dialect.name
        if dialect == "sqlite":
            from sqlalchemy.dialects.sqlite import insert as dialect_insert
        else:
            from sqlalchemy.dialects.postgresql import insert as dialect_insert

        ins = dialect_insert(ui_session_owner).values(**vals)
        stmt = ins.on_conflict_do_nothing(index_elements=[ui_session_owner.c.session_id])
        with self._engine.begin() as conn:
            conn.execute(stmt)

    def fetch_by_session_id(self, *, session_id: str) -> Mapping[str, Any] | None:
        sid = str(session_id or "").strip()
        if not sid:
            return None
        stmt = (
            select(
                ui_session_owner.c.tenant_id,
                ui_session_owner.c.user_id,
                ui_session_owner.c.created_at,
            )
            .where(ui_session_owner.c.session_id == sid)
            .limit(1)
        )
        with self._engine.connect() as conn:
            row = conn.execute(stmt).mappings().first()
        return row

    def backfill_orphan_sessions_for_user(
        self, *, tenant_id: str, user_id: str, default_created_at: str
    ) -> int:
        tid = str(tenant_id)
        uid = str(user_id)
        ts = str(default_created_at)
        owned = exists(
            select(1).select_from(ui_session_owner).where(ui_session_owner.c.session_id == chat_session.c.id)
        )
        sel = (
            select(
                chat_session.c.id.label("session_id"),
                literal(tid).label("tenant_id"),
                literal(uid).label("user_id"),
                func.coalesce(chat_session.c.created_at, literal(ts)).label("created_at"),
            )
            .where(~owned)
        )
        dialect = self._engine.dialect.name
        with self._engine.begin() as conn:
            if dialect == "sqlite":
                stmt = insert(ui_session_owner).prefix_with("OR IGNORE").from_select(
                    [
                        ui_session_owner.c.session_id,
                        ui_session_owner.c.tenant_id,
                        ui_session_owner.c.user_id,
                        ui_session_owner.c.created_at,
                    ],
                    sel,
                )
            else:
                from sqlalchemy.dialects.postgresql import insert as pg_insert

                stmt = (
                    pg_insert(ui_session_owner)
                    .from_select(
                        [
                            ui_session_owner.c.session_id,
                            ui_session_owner.c.tenant_id,
                            ui_session_owner.c.user_id,
                            ui_session_owner.c.created_at,
                        ],
                        sel,
                    )
                    .on_conflict_do_nothing(index_elements=[ui_session_owner.c.session_id])
                )
            res = conn.execute(stmt)
            return int(res.rowcount or 0)

    def backfill_from_channel_v2(self, *, created_at: str) -> int:
        ts = str(created_at)
        join_on = and_(
            channel_identity_v2.c.tenant_id == channel_session_v2.c.tenant_id,
            channel_identity_v2.c.channel == channel_session_v2.c.channel,
            channel_identity_v2.c.account_id == channel_session_v2.c.account_id,
            channel_identity_v2.c.external_user_id == channel_session_v2.c.external_user_id,
        )
        sel = (
            select(
                channel_session_v2.c.session_id,
                channel_session_v2.c.tenant_id,
                channel_identity_v2.c.user_id,
                literal(ts).label("created_at"),
            )
            .distinct()
            .select_from(channel_session_v2.join(channel_identity_v2, join_on))
            .where(
                channel_session_v2.c.session_id.isnot(None),
                channel_session_v2.c.session_id != "",
            )
        )
        dialect = self._engine.dialect.name
        with self._engine.begin() as conn:
            if dialect == "sqlite":
                stmt = insert(ui_session_owner).prefix_with("OR IGNORE").from_select(
                    [
                        ui_session_owner.c.session_id,
                        ui_session_owner.c.tenant_id,
                        ui_session_owner.c.user_id,
                        ui_session_owner.c.created_at,
                    ],
                    sel,
                )
            else:
                from sqlalchemy.dialects.postgresql import insert as pg_insert

                stmt = (
                    pg_insert(ui_session_owner)
                    .from_select(
                        [
                            ui_session_owner.c.session_id,
                            ui_session_owner.c.tenant_id,
                            ui_session_owner.c.user_id,
                            ui_session_owner.c.created_at,
                        ],
                        sel,
                    )
                    .on_conflict_do_nothing(index_elements=[ui_session_owner.c.session_id])
                )
            res = conn.execute(stmt)
            return int(res.rowcount or 0)


__all__ = ["UiSessionOwnerSaRepository"]