mirror of
https://github.com/hansjone/oclaw.git
synced 2026-10-19 11:40:46 +08:00
feat(persistence): PostgreSQL assistant store, chat persist fixes, gateway scripts
- Add SQLAlchemy Core repos, pg adapter/compat, assistant_store factory, Alembic bootstrap and migration/cutover scripts. - Harden chat_message writes (NUL scrub for PG), turn_uuid on attempt failure, WS turn_runner fallbacks and gateway executed_turn_uuid init. - start_gateway: log paths, PS7 stderr handling via cmd, background stdout/stderr redirect; runtime assistant_runtime_log_dir export. - Ops: clear_all_chat_sessions with PG-only --postgresql and env-gated wipe; clear_postgres_chat_sessions.ps1. - Tests: SA repos, pg compat, persist fallback, smoke env isolation; CI and docs touch-ups. Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
parent
2b32d11f43
commit
d14e9d3596
103 changed files with 7574 additions and 1641 deletions
30
svc/persistence/sa_repos/__init__.py
Normal file
30
svc/persistence/sa_repos/__init__.py
Normal file
|
|
@ -0,0 +1,30 @@
|
|||
"""SQLAlchemy-backed repository slices (incremental migration off raw SQL)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from svc.persistence.sa_repos.admin_user_stats import AdminUserStatsSaRepository
|
||||
from svc.persistence.sa_repos.app_settings import AppSettingsSaRepository
|
||||
from svc.persistence.sa_repos.app_users import AppUsersSaRepository
|
||||
from svc.persistence.sa_repos.auth_sessions import AuthSessionsSaRepository
|
||||
from svc.persistence.sa_repos.chat_messages import ChatMessagesSaRepository
|
||||
from svc.persistence.sa_repos.chat_sessions import ChatSessionsSaRepository
|
||||
from svc.persistence.sa_repos.session_tool_health import SessionToolHealthSaRepository
|
||||
from svc.persistence.sa_repos.tenant_bind_code import BindCodeSaRepository, TenantSaRepository
|
||||
from svc.persistence.sa_repos.tool_log_queries import ToolLogQueriesSaRepository
|
||||
from svc.persistence.sa_repos.trace_events import TraceEventsSaRepository
|
||||
from svc.persistence.sa_repos.ui_session_owner import UiSessionOwnerSaRepository
|
||||
|
||||
__all__ = [
|
||||
"AdminUserStatsSaRepository",
|
||||
"AppSettingsSaRepository",
|
||||
"AppUsersSaRepository",
|
||||
"AuthSessionsSaRepository",
|
||||
"BindCodeSaRepository",
|
||||
"ChatMessagesSaRepository",
|
||||
"ChatSessionsSaRepository",
|
||||
"SessionToolHealthSaRepository",
|
||||
"TenantSaRepository",
|
||||
"ToolLogQueriesSaRepository",
|
||||
"TraceEventsSaRepository",
|
||||
"UiSessionOwnerSaRepository",
|
||||
]
|
||||
167
svc/persistence/sa_repos/admin_user_stats.py
Normal file
167
svc/persistence/sa_repos/admin_user_stats.py
Normal file
|
|
@ -0,0 +1,167 @@
|
|||
"""Admin tenant user stats (list_admin_user_stats) via SQLAlchemy Core."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import and_, case, distinct, func, literal, or_, select
|
||||
from sqlalchemy.engine import Engine
|
||||
from sqlalchemy.sql import bindparam
|
||||
|
||||
from svc.persistence.db.tables import (
|
||||
app_user,
|
||||
auth_session,
|
||||
chat_session,
|
||||
trace_event,
|
||||
ui_session_owner,
|
||||
)
|
||||
|
||||
|
||||
class AdminUserStatsSaRepository:
|
||||
"""Aggregates for ``SqliteStore.list_admin_user_stats``."""
|
||||
|
||||
__slots__ = ("_engine",)
|
||||
|
||||
def __init__(self, engine: Engine) -> None:
|
||||
self._engine = engine
|
||||
|
||||
def fetch(
|
||||
self,
|
||||
*,
|
||||
tenant_id: str,
|
||||
search_lower: str | None,
|
||||
cutoff_iso: str,
|
||||
limit: int,
|
||||
offset: int,
|
||||
) -> dict[str, Any]:
|
||||
tid = str(tenant_id or "").strip()
|
||||
lim = max(1, min(int(limit), 500))
|
||||
off = max(0, int(offset))
|
||||
q_text = str(search_lower or "").strip().lower() or None
|
||||
cutoff = str(cutoff_iso)
|
||||
|
||||
conds: list[Any] = [app_user.c.tenant_id == tid]
|
||||
if q_text:
|
||||
like = f"%{q_text}%"
|
||||
conds.append(
|
||||
or_(
|
||||
func.lower(func.coalesce(app_user.c.username, literal(""))).like(like),
|
||||
func.lower(func.coalesce(app_user.c.display_name, literal(""))).like(like),
|
||||
)
|
||||
)
|
||||
wh = and_(*conds)
|
||||
|
||||
user_stmt = (
|
||||
select(
|
||||
app_user.c.id.label("user_id"),
|
||||
app_user.c.username,
|
||||
func.coalesce(app_user.c.display_name, literal("")).label("display_name"),
|
||||
app_user.c.role,
|
||||
app_user.c.is_active,
|
||||
)
|
||||
.where(wh)
|
||||
.order_by(app_user.c.username.asc())
|
||||
.limit(lim)
|
||||
.offset(off)
|
||||
)
|
||||
cnt_stmt = select(func.count()).select_from(app_user).where(wh)
|
||||
|
||||
sess_join = chat_session.join(
|
||||
ui_session_owner,
|
||||
ui_session_owner.c.session_id == chat_session.c.id,
|
||||
)
|
||||
active_ts = func.coalesce(chat_session.c.last_message_at, chat_session.c.created_at)
|
||||
total_active_sess_stmt = (
|
||||
select(func.count(distinct(chat_session.c.id)))
|
||||
.select_from(sess_join)
|
||||
.where(ui_session_owner.c.tenant_id == tid, active_ts >= literal(cutoff))
|
||||
)
|
||||
total_active_logins_stmt = (
|
||||
select(func.count())
|
||||
.select_from(auth_session)
|
||||
.where(
|
||||
auth_session.c.tenant_id == tid,
|
||||
auth_session.c.revoked_at.is_(None),
|
||||
auth_session.c.expires_at > literal(cutoff),
|
||||
auth_session.c.last_seen_at >= literal(cutoff),
|
||||
)
|
||||
)
|
||||
|
||||
with self._engine.connect() as conn:
|
||||
total_users = int(conn.execute(cnt_stmt).scalar_one() or 0)
|
||||
user_rows = [dict(r) for r in conn.execute(user_stmt).mappings().all()]
|
||||
total_active_sessions = int(conn.execute(total_active_sess_stmt).scalar_one() or 0)
|
||||
total_active_logins = int(conn.execute(total_active_logins_stmt).scalar_one() or 0)
|
||||
|
||||
uids = [str(r["user_id"] or "").strip() for r in user_rows if str(r.get("user_id") or "").strip()]
|
||||
trace_rows: list[dict[str, Any]] = []
|
||||
own_count_rows: list[dict[str, Any]] = []
|
||||
active_sess_rows: list[dict[str, Any]] = []
|
||||
login_rows: list[dict[str, Any]] = []
|
||||
|
||||
if uids:
|
||||
uids_param = bindparam("uids", expanding=True)
|
||||
trace_stmt = (
|
||||
select(ui_session_owner.c.user_id, trace_event.c.payload)
|
||||
.select_from(
|
||||
trace_event.join(
|
||||
ui_session_owner,
|
||||
ui_session_owner.c.session_id == trace_event.c.session_id,
|
||||
)
|
||||
)
|
||||
.where(ui_session_owner.c.tenant_id == tid, ui_session_owner.c.user_id.in_(uids_param))
|
||||
)
|
||||
own_cnt_stmt = (
|
||||
select(ui_session_owner.c.user_id, func.count().label("c"))
|
||||
.where(ui_session_owner.c.tenant_id == tid, ui_session_owner.c.user_id.in_(uids_param))
|
||||
.group_by(ui_session_owner.c.user_id)
|
||||
)
|
||||
active_case = case(
|
||||
(active_ts >= literal(cutoff), chat_session.c.id),
|
||||
else_=None,
|
||||
)
|
||||
active_sess_stmt = (
|
||||
select(
|
||||
ui_session_owner.c.user_id,
|
||||
func.count(distinct(active_case)).label("active_30m"),
|
||||
func.max(active_ts).label("last_message_at"),
|
||||
)
|
||||
.select_from(sess_join)
|
||||
.where(ui_session_owner.c.tenant_id == tid, ui_session_owner.c.user_id.in_(uids_param))
|
||||
.group_by(ui_session_owner.c.user_id)
|
||||
)
|
||||
login_stmt = (
|
||||
select(
|
||||
auth_session.c.user_id,
|
||||
func.count().label("c"),
|
||||
func.max(auth_session.c.last_seen_at).label("last_seen_at"),
|
||||
)
|
||||
.where(
|
||||
auth_session.c.tenant_id == tid,
|
||||
auth_session.c.user_id.in_(uids_param),
|
||||
auth_session.c.revoked_at.is_(None),
|
||||
auth_session.c.expires_at > literal(cutoff),
|
||||
auth_session.c.last_seen_at >= literal(cutoff),
|
||||
)
|
||||
.group_by(auth_session.c.user_id)
|
||||
)
|
||||
bind = {"uids": uids}
|
||||
with self._engine.connect() as conn:
|
||||
trace_rows = [dict(r) for r in conn.execute(trace_stmt, bind).mappings().all()]
|
||||
own_count_rows = [dict(r) for r in conn.execute(own_cnt_stmt, bind).mappings().all()]
|
||||
active_sess_rows = [dict(r) for r in conn.execute(active_sess_stmt, bind).mappings().all()]
|
||||
login_rows = [dict(r) for r in conn.execute(login_stmt, bind).mappings().all()]
|
||||
|
||||
return {
|
||||
"total_users": total_users,
|
||||
"user_rows": user_rows,
|
||||
"total_active_sessions_30m": total_active_sessions,
|
||||
"total_active_logins_30m": total_active_logins,
|
||||
"trace_rows": trace_rows,
|
||||
"sessions_count_rows": own_count_rows,
|
||||
"active_sess_rows": active_sess_rows,
|
||||
"login_rows": login_rows,
|
||||
}
|
||||
|
||||
|
||||
__all__ = ["AdminUserStatsSaRepository"]
|
||||
141
svc/persistence/sa_repos/app_settings.py
Normal file
141
svc/persistence/sa_repos/app_settings.py
Normal file
|
|
@ -0,0 +1,141 @@
|
|||
"""app_setting access via SQLAlchemy Core (SQLite + PostgreSQL)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
|
||||
from sqlalchemy import delete, func, select, update
|
||||
from sqlalchemy.dialects.postgresql import insert as pg_insert
|
||||
from sqlalchemy.dialects.sqlite import insert as sqlite_insert
|
||||
from sqlalchemy.engine import Connection, Engine
|
||||
|
||||
from svc.persistence.db.tables import app_setting
|
||||
|
||||
|
||||
class AppSettingsSaRepository:
|
||||
"""Phase-1 SA migration: ``app_setting`` reads/writes only."""
|
||||
|
||||
__slots__ = ("_engine",)
|
||||
|
||||
def __init__(self, engine: Engine) -> None:
|
||||
self._engine = engine
|
||||
|
||||
def _dialect(self, conn: Connection) -> str:
|
||||
return conn.engine.dialect.name
|
||||
|
||||
def _upsert(
|
||||
self,
|
||||
conn: Connection,
|
||||
*,
|
||||
key: str,
|
||||
value: str,
|
||||
is_secret: int,
|
||||
updated_at: str,
|
||||
) -> None:
|
||||
dialect = self._dialect(conn)
|
||||
if dialect == "sqlite":
|
||||
ins = sqlite_insert(app_setting).values(
|
||||
key=key,
|
||||
value=value,
|
||||
is_secret=is_secret,
|
||||
updated_at=updated_at,
|
||||
)
|
||||
stmt = ins.on_conflict_do_update(
|
||||
index_elements=[app_setting.c.key],
|
||||
set_={
|
||||
"value": ins.excluded.value,
|
||||
"is_secret": ins.excluded.is_secret,
|
||||
"updated_at": ins.excluded.updated_at,
|
||||
},
|
||||
)
|
||||
elif dialect == "postgresql":
|
||||
ins = pg_insert(app_setting).values(
|
||||
key=key,
|
||||
value=value,
|
||||
is_secret=is_secret,
|
||||
updated_at=updated_at,
|
||||
)
|
||||
stmt = ins.on_conflict_do_update(
|
||||
index_elements=[app_setting.c.key],
|
||||
set_={
|
||||
"value": ins.excluded.value,
|
||||
"is_secret": ins.excluded.is_secret,
|
||||
"updated_at": ins.excluded.updated_at,
|
||||
},
|
||||
)
|
||||
else:
|
||||
raise RuntimeError(f"unsupported SQLAlchemy dialect for app_setting: {dialect!r}")
|
||||
conn.execute(stmt)
|
||||
|
||||
def upsert_plain(self, *, key: str, value: str, updated_at: str) -> None:
|
||||
with self._engine.begin() as conn:
|
||||
self._upsert(conn, key=key, value=value, is_secret=0, updated_at=updated_at)
|
||||
|
||||
def upsert_secret(self, *, key: str, encoded_value: str, updated_at: str) -> None:
|
||||
with self._engine.begin() as conn:
|
||||
self._upsert(conn, key=key, value=encoded_value, is_secret=1, updated_at=updated_at)
|
||||
|
||||
def fetch_row(self, *, key: str) -> tuple[str, int] | None:
|
||||
with self._engine.connect() as conn:
|
||||
row = conn.execute(
|
||||
select(app_setting.c.value, app_setting.c.is_secret).where(app_setting.c.key == key)
|
||||
).one_or_none()
|
||||
if row is None:
|
||||
return None
|
||||
return (str(row[0]), int(row[1]))
|
||||
|
||||
def delete_key(self, *, key: str) -> None:
|
||||
with self._engine.begin() as conn:
|
||||
conn.execute(delete(app_setting).where(app_setting.c.key == key))
|
||||
|
||||
def migrate_b64_secrets(
|
||||
self,
|
||||
*,
|
||||
ts: str,
|
||||
decode_secret: Callable[[str], str],
|
||||
encode_secret: Callable[[str], str],
|
||||
predicate_new_encoding: Callable[[str], bool],
|
||||
) -> int:
|
||||
"""Re-encode legacy ``b64:`` rows; idempotent. Returns rows updated."""
|
||||
migrated = 0
|
||||
with self._engine.begin() as conn:
|
||||
rows = conn.execute(
|
||||
select(app_setting.c.key, app_setting.c.value).where(
|
||||
app_setting.c.is_secret == 1,
|
||||
app_setting.c.value.like("b64:%"),
|
||||
)
|
||||
).all()
|
||||
for k, v in rows:
|
||||
key = str(k or "")
|
||||
val = str(v or "")
|
||||
if not key:
|
||||
continue
|
||||
try:
|
||||
plain = decode_secret(val)
|
||||
except Exception:
|
||||
continue
|
||||
enc = encode_secret(plain)
|
||||
if enc != val and predicate_new_encoding(enc):
|
||||
conn.execute(
|
||||
update(app_setting)
|
||||
.where(
|
||||
app_setting.c.key == key,
|
||||
app_setting.c.is_secret == 1,
|
||||
)
|
||||
.values(value=enc, updated_at=ts)
|
||||
)
|
||||
migrated += 1
|
||||
return migrated
|
||||
|
||||
def count_legacy_b64_secrets(self) -> int:
|
||||
with self._engine.connect() as conn:
|
||||
n = conn.execute(
|
||||
select(func.count()).select_from(app_setting).where(
|
||||
app_setting.c.is_secret == 1,
|
||||
app_setting.c.value.like("b64:%"),
|
||||
)
|
||||
).scalar_one()
|
||||
return int(n or 0)
|
||||
|
||||
|
||||
__all__ = ["AppSettingsSaRepository"]
|
||||
296
svc/persistence/sa_repos/app_users.py
Normal file
296
svc/persistence/sa_repos/app_users.py
Normal file
|
|
@ -0,0 +1,296 @@
|
|||
"""app_user CRUD slices via SQLAlchemy Core."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Mapping
|
||||
|
||||
from sqlalchemy import case, delete, exists, func, insert, literal, or_, select, union, update
|
||||
from sqlalchemy.engine import Engine
|
||||
|
||||
from svc.persistence.db.tables import app_user, channel_identity, channel_identity_v2
|
||||
|
||||
|
||||
def _user_row_to_public_dict(r: Mapping[str, Any]) -> dict[str, Any]:
|
||||
return {
|
||||
"id": r["id"],
|
||||
"tenant_id": r["tenant_id"],
|
||||
"username": r["username"],
|
||||
"display_name": r["display_name"],
|
||||
"role": r["role"],
|
||||
"is_active": bool(int(r["is_active"] or 0)),
|
||||
"created_at": r["created_at"],
|
||||
"password_hash": r["password_hash"],
|
||||
"avatar_attachment_id": str(r["avatar_attachment_id"] or "").strip() or None,
|
||||
}
|
||||
|
||||
|
||||
class AppUsersSaRepository:
|
||||
__slots__ = ("_engine",)
|
||||
|
||||
def __init__(self, engine: Engine) -> None:
|
||||
self._engine = engine
|
||||
|
||||
def count_by_tenant_username(self, *, tenant_id: str, username: str) -> int:
|
||||
tid, un = str(tenant_id), str(username)
|
||||
with self._engine.connect() as conn:
|
||||
n = conn.execute(
|
||||
select(func.count())
|
||||
.select_from(app_user)
|
||||
.where(app_user.c.tenant_id == tid, app_user.c.username == un)
|
||||
).scalar_one()
|
||||
return int(n or 0)
|
||||
|
||||
def insert_user(
|
||||
self,
|
||||
*,
|
||||
user_id: str,
|
||||
tenant_id: str,
|
||||
username: str,
|
||||
display_name: str,
|
||||
role: str,
|
||||
password_hash: str,
|
||||
is_active: int,
|
||||
created_at: str,
|
||||
avatar_attachment_id: str | None = None,
|
||||
) -> None:
|
||||
with self._engine.begin() as conn:
|
||||
conn.execute(
|
||||
insert(app_user).values(
|
||||
id=str(user_id),
|
||||
tenant_id=str(tenant_id),
|
||||
username=str(username),
|
||||
display_name=str(display_name),
|
||||
role=str(role),
|
||||
password_hash=str(password_hash or ""),
|
||||
is_active=int(is_active),
|
||||
created_at=str(created_at),
|
||||
avatar_attachment_id=avatar_attachment_id,
|
||||
)
|
||||
)
|
||||
|
||||
def _select_user_columns(self):
|
||||
return select(
|
||||
app_user.c.id,
|
||||
app_user.c.tenant_id,
|
||||
app_user.c.username,
|
||||
app_user.c.display_name,
|
||||
app_user.c.role,
|
||||
func.coalesce(app_user.c.is_active, literal(1)).label("is_active"),
|
||||
app_user.c.created_at,
|
||||
func.coalesce(app_user.c.password_hash, literal("")).label("password_hash"),
|
||||
func.coalesce(app_user.c.avatar_attachment_id, literal("")).label("avatar_attachment_id"),
|
||||
)
|
||||
|
||||
def fetch_by_tenant_and_id(self, *, tenant_id: str, user_id: str) -> dict[str, Any] | None:
|
||||
tid, uid = str(tenant_id), str(user_id)
|
||||
stmt = self._select_user_columns().where(app_user.c.tenant_id == tid, app_user.c.id == uid).limit(1)
|
||||
with self._engine.connect() as conn:
|
||||
row = conn.execute(stmt).mappings().first()
|
||||
return _user_row_to_public_dict(row) if row else None
|
||||
|
||||
def fetch_by_tenant_and_username(self, *, tenant_id: str, username: str) -> dict[str, Any] | None:
|
||||
tid, un = str(tenant_id), str(username)
|
||||
stmt = (
|
||||
self._select_user_columns()
|
||||
.where(app_user.c.tenant_id == tid, app_user.c.username == un)
|
||||
.limit(1)
|
||||
)
|
||||
with self._engine.connect() as conn:
|
||||
row = conn.execute(stmt).mappings().first()
|
||||
return _user_row_to_public_dict(row) if row else None
|
||||
|
||||
def fetch_first_by_username_global(self, *, username: str) -> dict[str, Any] | None:
|
||||
un = str(username)
|
||||
stmt = (
|
||||
self._select_user_columns()
|
||||
.where(app_user.c.username == un)
|
||||
.order_by(app_user.c.created_at.asc())
|
||||
.limit(1)
|
||||
)
|
||||
with self._engine.connect() as conn:
|
||||
row = conn.execute(stmt).mappings().first()
|
||||
return _user_row_to_public_dict(row) if row else None
|
||||
|
||||
def list_users_for_tenant(
|
||||
self,
|
||||
*,
|
||||
tenant_id: str,
|
||||
limit: int,
|
||||
offset: int,
|
||||
q: str | None,
|
||||
include_inactive: bool,
|
||||
) -> list[dict[str, Any]]:
|
||||
tid = str(tenant_id)
|
||||
lim = max(1, int(limit))
|
||||
off = max(0, int(offset))
|
||||
|
||||
has_password = case(
|
||||
(func.trim(func.coalesce(app_user.c.password_hash, literal(""))) != literal(""), 1),
|
||||
else_=0,
|
||||
).label("has_password")
|
||||
|
||||
wecom_linked = or_(
|
||||
exists(
|
||||
select(literal(1))
|
||||
.select_from(channel_identity_v2)
|
||||
.where(
|
||||
channel_identity_v2.c.tenant_id == app_user.c.tenant_id,
|
||||
channel_identity_v2.c.user_id == app_user.c.id,
|
||||
channel_identity_v2.c.channel == literal("wecom"),
|
||||
)
|
||||
),
|
||||
exists(
|
||||
select(literal(1))
|
||||
.select_from(channel_identity)
|
||||
.where(
|
||||
channel_identity.c.tenant_id == app_user.c.tenant_id,
|
||||
channel_identity.c.user_id == app_user.c.id,
|
||||
channel_identity.c.channel == literal("wecom"),
|
||||
)
|
||||
),
|
||||
).label("wecom_linked")
|
||||
|
||||
channel_linked = or_(
|
||||
exists(
|
||||
select(literal(1))
|
||||
.select_from(channel_identity_v2)
|
||||
.where(
|
||||
channel_identity_v2.c.tenant_id == app_user.c.tenant_id,
|
||||
channel_identity_v2.c.user_id == app_user.c.id,
|
||||
)
|
||||
),
|
||||
exists(
|
||||
select(literal(1))
|
||||
.select_from(channel_identity)
|
||||
.where(
|
||||
channel_identity.c.tenant_id == app_user.c.tenant_id,
|
||||
channel_identity.c.user_id == app_user.c.id,
|
||||
)
|
||||
),
|
||||
).label("channel_linked")
|
||||
|
||||
eid_ci = func.trim(func.coalesce(channel_identity.c.external_user_id, literal("")))
|
||||
sq1 = (
|
||||
select(eid_ci.label("eid"))
|
||||
.where(
|
||||
channel_identity.c.tenant_id == app_user.c.tenant_id,
|
||||
channel_identity.c.user_id == app_user.c.id,
|
||||
channel_identity.c.channel == literal("wecom"),
|
||||
eid_ci != literal(""),
|
||||
)
|
||||
.distinct()
|
||||
)
|
||||
eid_v2 = func.trim(func.coalesce(channel_identity_v2.c.external_user_id, literal("")))
|
||||
sq2 = (
|
||||
select(eid_v2.label("eid"))
|
||||
.where(
|
||||
channel_identity_v2.c.tenant_id == app_user.c.tenant_id,
|
||||
channel_identity_v2.c.user_id == app_user.c.id,
|
||||
channel_identity_v2.c.channel == literal("wecom"),
|
||||
eid_v2 != literal(""),
|
||||
)
|
||||
.distinct()
|
||||
)
|
||||
u_sub = union(sq1, sq2).subquery()
|
||||
if self._engine.dialect.name == "postgresql":
|
||||
wecom_ids_expr = select(func.string_agg(u_sub.c.eid, literal(", "))).scalar_subquery()
|
||||
else:
|
||||
wecom_ids_expr = select(func.group_concat(u_sub.c.eid, literal(", "))).scalar_subquery()
|
||||
|
||||
stmt = (
|
||||
select(
|
||||
app_user.c.id,
|
||||
app_user.c.tenant_id,
|
||||
app_user.c.username,
|
||||
app_user.c.display_name,
|
||||
app_user.c.role,
|
||||
func.coalesce(app_user.c.is_active, literal(1)).label("is_active"),
|
||||
app_user.c.created_at,
|
||||
has_password,
|
||||
wecom_linked,
|
||||
channel_linked,
|
||||
wecom_ids_expr.label("wecom_external_user_ids"),
|
||||
)
|
||||
.where(app_user.c.tenant_id == tid)
|
||||
)
|
||||
token = str(q or "").strip()
|
||||
if token:
|
||||
key = f"%{token.lower()}%"
|
||||
stmt = stmt.where(
|
||||
or_(
|
||||
func.lower(app_user.c.display_name).like(key),
|
||||
func.lower(func.coalesce(app_user.c.username, literal(""))).like(key),
|
||||
app_user.c.id.like(f"%{token[:32]}%"),
|
||||
)
|
||||
)
|
||||
if not include_inactive:
|
||||
stmt = stmt.where(func.coalesce(app_user.c.is_active, literal(1)) == literal(1))
|
||||
stmt = stmt.order_by(app_user.c.created_at.desc()).limit(lim).offset(off)
|
||||
|
||||
with self._engine.connect() as conn:
|
||||
rows = conn.execute(stmt).mappings().all()
|
||||
out: list[dict[str, Any]] = []
|
||||
for r in rows:
|
||||
has_pw = bool(int(r["has_password"] or 0))
|
||||
uname = str(r["username"] or "")
|
||||
can_chat = bool(has_pw)
|
||||
wl = r["wecom_linked"]
|
||||
cl = r["channel_linked"]
|
||||
out.append(
|
||||
{
|
||||
"id": r["id"],
|
||||
"tenant_id": r["tenant_id"],
|
||||
"username": r["username"],
|
||||
"display_name": r["display_name"],
|
||||
"role": r["role"],
|
||||
"is_active": bool(int(r["is_active"] or 0)),
|
||||
"created_at": r["created_at"],
|
||||
"has_password": has_pw,
|
||||
"wecom_linked": bool(int(wl or 0)),
|
||||
"channel_linked": bool(int(cl or 0)),
|
||||
"can_chat_login": can_chat,
|
||||
"wecom_external_user_ids": str(r["wecom_external_user_ids"] or "").strip(),
|
||||
}
|
||||
)
|
||||
return out
|
||||
|
||||
def update_user_account(
|
||||
self,
|
||||
*,
|
||||
tenant_id: str,
|
||||
user_id: str,
|
||||
display_name: str | None = None,
|
||||
role: str | None = None,
|
||||
is_active: bool | None = None,
|
||||
password_hash: str | None = None,
|
||||
avatar_attachment_id: str | None = None,
|
||||
) -> bool:
|
||||
tid, uid = str(tenant_id), str(user_id)
|
||||
vals: dict[str, Any] = {}
|
||||
if display_name is not None:
|
||||
vals["display_name"] = str(display_name).strip() or "User"
|
||||
if role is not None:
|
||||
vals["role"] = str(role).strip() or "member"
|
||||
if is_active is not None:
|
||||
vals["is_active"] = 1 if is_active else 0
|
||||
if password_hash is not None:
|
||||
vals["password_hash"] = str(password_hash)
|
||||
if avatar_attachment_id is not None:
|
||||
aid = str(avatar_attachment_id).strip()
|
||||
vals["avatar_attachment_id"] = aid if aid else None
|
||||
if not vals:
|
||||
return False
|
||||
with self._engine.begin() as conn:
|
||||
res = conn.execute(
|
||||
update(app_user).where(app_user.c.tenant_id == tid, app_user.c.id == uid).values(**vals)
|
||||
)
|
||||
return bool(int(res.rowcount or 0) > 0)
|
||||
|
||||
def delete_user_account(self, *, tenant_id: str, user_id: str) -> int:
|
||||
tid, uid = str(tenant_id), str(user_id)
|
||||
with self._engine.begin() as conn:
|
||||
res = conn.execute(delete(app_user).where(app_user.c.tenant_id == tid, app_user.c.id == uid))
|
||||
return int(res.rowcount or 0)
|
||||
|
||||
|
||||
__all__ = ["AppUsersSaRepository"]
|
||||
102
svc/persistence/sa_repos/auth_sessions.py
Normal file
102
svc/persistence/sa_repos/auth_sessions.py
Normal file
|
|
@ -0,0 +1,102 @@
|
|||
"""auth_session access via SQLAlchemy Core (SQLite + PostgreSQL)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import insert, select, update
|
||||
from sqlalchemy.engine import Engine
|
||||
|
||||
from svc.persistence.db.tables import auth_session
|
||||
|
||||
|
||||
class AuthSessionsSaRepository:
|
||||
"""Phase-2 SA migration: admin login session rows."""
|
||||
|
||||
__slots__ = ("_engine",)
|
||||
|
||||
def __init__(self, engine: Engine) -> None:
|
||||
self._engine = engine
|
||||
|
||||
def insert_session(
|
||||
self,
|
||||
*,
|
||||
session_token_hash: str,
|
||||
tenant_id: str,
|
||||
user_id: str,
|
||||
role: str,
|
||||
created_at: str,
|
||||
expires_at: str,
|
||||
last_seen_at: str,
|
||||
) -> None:
|
||||
with self._engine.begin() as conn:
|
||||
conn.execute(
|
||||
insert(auth_session).values(
|
||||
session_token_hash=str(session_token_hash),
|
||||
tenant_id=str(tenant_id),
|
||||
user_id=str(user_id),
|
||||
role=str(role),
|
||||
created_at=str(created_at),
|
||||
expires_at=str(expires_at),
|
||||
last_seen_at=str(last_seen_at),
|
||||
revoked_at=None,
|
||||
)
|
||||
)
|
||||
|
||||
def revoke_one(self, *, session_token_hash: str, revoked_at: str) -> int:
|
||||
with self._engine.begin() as conn:
|
||||
res = conn.execute(
|
||||
update(auth_session)
|
||||
.where(
|
||||
auth_session.c.session_token_hash == str(session_token_hash),
|
||||
auth_session.c.revoked_at.is_(None),
|
||||
)
|
||||
.values(revoked_at=str(revoked_at))
|
||||
)
|
||||
n = res.rowcount
|
||||
if n is None or n < 0:
|
||||
return 0
|
||||
return int(n)
|
||||
|
||||
def revoke_all_active(self, *, revoked_at: str) -> int:
|
||||
with self._engine.begin() as conn:
|
||||
res = conn.execute(
|
||||
update(auth_session)
|
||||
.where(auth_session.c.revoked_at.is_(None))
|
||||
.values(revoked_at=str(revoked_at))
|
||||
)
|
||||
n = res.rowcount
|
||||
if n is None or n < 0:
|
||||
return 0
|
||||
return int(n)
|
||||
|
||||
def fetch_by_hash(self, *, session_token_hash: str) -> dict[str, Any] | None:
|
||||
with self._engine.connect() as conn:
|
||||
row = conn.execute(
|
||||
select(
|
||||
auth_session.c.session_token_hash,
|
||||
auth_session.c.tenant_id,
|
||||
auth_session.c.user_id,
|
||||
auth_session.c.role,
|
||||
auth_session.c.created_at,
|
||||
auth_session.c.expires_at,
|
||||
auth_session.c.last_seen_at,
|
||||
auth_session.c.revoked_at,
|
||||
)
|
||||
.where(auth_session.c.session_token_hash == str(session_token_hash))
|
||||
.limit(1)
|
||||
).mappings().first()
|
||||
if row is None:
|
||||
return None
|
||||
return dict(row)
|
||||
|
||||
def touch(self, *, session_token_hash: str, last_seen_at: str) -> None:
|
||||
with self._engine.begin() as conn:
|
||||
conn.execute(
|
||||
update(auth_session)
|
||||
.where(auth_session.c.session_token_hash == str(session_token_hash))
|
||||
.values(last_seen_at=str(last_seen_at))
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["AuthSessionsSaRepository"]
|
||||
481
svc/persistence/sa_repos/chat_messages.py
Normal file
481
svc/persistence/sa_repos/chat_messages.py
Normal file
|
|
@ -0,0 +1,481 @@
|
|||
"""chat_message access via SQLAlchemy Core."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import Any, Mapping
|
||||
|
||||
from sqlalchemy import delete, exists, func, insert, literal, select, update
|
||||
from sqlalchemy.engine import Engine
|
||||
|
||||
from svc.persistence.db.tables import chat_message, chat_session
|
||||
from svc.persistence.sqlite_store import (
|
||||
ChatMessage,
|
||||
SessionMessagesMeta,
|
||||
_tool_row_assistant_message_id,
|
||||
_trim_messages_start_index,
|
||||
utc_now_iso,
|
||||
)
|
||||
|
||||
|
||||
def _sql_text_required(v: Any) -> str:
|
||||
if v is None:
|
||||
return ""
|
||||
if isinstance(v, (bytes, bytearray, memoryview)):
|
||||
try:
|
||||
return bytes(v).decode("utf-8", errors="replace")
|
||||
except Exception:
|
||||
return ""
|
||||
return str(v)
|
||||
|
||||
|
||||
def _sql_text_optional_plain(v: Any) -> str | None:
|
||||
if v is None:
|
||||
return None
|
||||
if isinstance(v, (bytes, bytearray, memoryview)):
|
||||
try:
|
||||
s = bytes(v).decode("utf-8", errors="replace")
|
||||
except Exception:
|
||||
return None
|
||||
s = s.strip()
|
||||
return s if s else None
|
||||
s = str(v).strip()
|
||||
return s if s else None
|
||||
|
||||
|
||||
def _sql_text_optional_jsonish(v: Any) -> str | None:
|
||||
"""Normalize TEXT/JSON columns across SQLite + PostgreSQL drivers (bytes/memoryview/dict)."""
|
||||
if v is None:
|
||||
return None
|
||||
if isinstance(v, (bytes, bytearray, memoryview)):
|
||||
try:
|
||||
s = bytes(v).decode("utf-8", errors="replace")
|
||||
except Exception:
|
||||
return None
|
||||
return s if s.strip() else None
|
||||
if isinstance(v, (dict, list)):
|
||||
return json.dumps(v, ensure_ascii=False, default=str)
|
||||
s = str(v).strip()
|
||||
return s if s else None
|
||||
|
||||
|
||||
def _row_to_chat_message(r: Mapping[str, Any]) -> ChatMessage:
|
||||
return ChatMessage(
|
||||
id=int(r["id"]),
|
||||
session_id=str(r["session_id"]),
|
||||
role=str(r["role"]),
|
||||
content=_sql_text_required(r.get("content")),
|
||||
tool_calls=_sql_text_optional_jsonish(r.get("tool_calls")),
|
||||
attachments=_sql_text_optional_jsonish(r.get("attachments")),
|
||||
turn_uuid=_sql_text_optional_plain(r.get("turn_uuid")),
|
||||
event_type=_sql_text_optional_plain(r.get("event_type")),
|
||||
event_payload=_sql_text_optional_jsonish(r.get("event_payload")),
|
||||
timestamp=str(r["timestamp"]),
|
||||
)
|
||||
|
||||
|
||||
class ChatMessagesSaRepository:
|
||||
"""Phase-4 SA migration: chat_message CRUD + session last_message_at touch."""
|
||||
|
||||
__slots__ = ("_engine",)
|
||||
|
||||
def __init__(self, engine: Engine) -> None:
|
||||
self._engine = engine
|
||||
|
||||
def insert_message_and_touch_session(
|
||||
self,
|
||||
*,
|
||||
session_id: str,
|
||||
role: str,
|
||||
content: str,
|
||||
tool_calls: str | None,
|
||||
attachments: str | None,
|
||||
turn_uuid: str | None,
|
||||
event_type: str | None,
|
||||
event_payload: str | None,
|
||||
timestamp: str,
|
||||
) -> int:
|
||||
sid = str(session_id if session_id is not None else "")
|
||||
with self._engine.begin() as conn:
|
||||
stmt = (
|
||||
insert(chat_message)
|
||||
.values(
|
||||
session_id=sid,
|
||||
role=str(role),
|
||||
content=str(content),
|
||||
tool_calls=tool_calls,
|
||||
attachments=attachments,
|
||||
turn_uuid=turn_uuid,
|
||||
event_type=event_type,
|
||||
event_payload=event_payload,
|
||||
timestamp=str(timestamp),
|
||||
)
|
||||
.returning(chat_message.c.id)
|
||||
)
|
||||
msg_id = int(conn.execute(stmt).scalar_one())
|
||||
conn.execute(
|
||||
update(chat_session)
|
||||
.where(chat_session.c.id == sid)
|
||||
.values(last_message_at=str(timestamp))
|
||||
)
|
||||
return msg_id
|
||||
|
||||
def delete_message_and_refresh_session(self, *, session_id: str, message_id: int) -> bool:
|
||||
sid = str(session_id or "").strip()
|
||||
mid = int(message_id or 0)
|
||||
if not sid or mid <= 0:
|
||||
return False
|
||||
with self._engine.begin() as conn:
|
||||
res = conn.execute(
|
||||
delete(chat_message).where(
|
||||
chat_message.c.session_id == sid,
|
||||
chat_message.c.id == mid,
|
||||
)
|
||||
)
|
||||
if int(res.rowcount or 0) <= 0:
|
||||
return False
|
||||
last_ts = conn.execute(
|
||||
select(func.max(chat_message.c.timestamp)).where(chat_message.c.session_id == sid)
|
||||
).scalar_one_or_none()
|
||||
last_s = str(last_ts or "").strip() or None
|
||||
conn.execute(
|
||||
update(chat_session)
|
||||
.where(chat_session.c.id == sid)
|
||||
.values(last_message_at=last_s)
|
||||
)
|
||||
return True
|
||||
|
||||
def update_message_content(
|
||||
self,
|
||||
*,
|
||||
session_id: str,
|
||||
message_id: int,
|
||||
content: str,
|
||||
event_payload_text: str | None,
|
||||
) -> bool:
|
||||
sid = str(session_id or "").strip()
|
||||
mid = int(message_id or 0)
|
||||
if not sid or mid <= 0:
|
||||
return False
|
||||
with self._engine.begin() as conn:
|
||||
res = conn.execute(
|
||||
update(chat_message)
|
||||
.where(chat_message.c.session_id == sid, chat_message.c.id == mid)
|
||||
.values(
|
||||
content=str(content or ""),
|
||||
event_payload=func.coalesce(literal(event_payload_text), chat_message.c.event_payload),
|
||||
)
|
||||
)
|
||||
return int(res.rowcount or 0) > 0
|
||||
|
||||
def get_messages_recent_asc(self, *, session_id: str, limit: int) -> list[ChatMessage]:
|
||||
if limit <= 0:
|
||||
return []
|
||||
sid = str(session_id or "").strip()
|
||||
if not sid:
|
||||
return []
|
||||
lim = max(1, min(int(limit), 2000))
|
||||
ids_sq = (
|
||||
select(chat_message.c.id)
|
||||
.where(chat_message.c.session_id == sid)
|
||||
.order_by(chat_message.c.id.desc())
|
||||
.limit(lim)
|
||||
.scalar_subquery()
|
||||
)
|
||||
stmt = (
|
||||
select(
|
||||
chat_message.c.id,
|
||||
chat_message.c.session_id,
|
||||
chat_message.c.role,
|
||||
chat_message.c.content,
|
||||
chat_message.c.tool_calls,
|
||||
chat_message.c.attachments,
|
||||
chat_message.c.turn_uuid,
|
||||
chat_message.c.event_type,
|
||||
chat_message.c.event_payload,
|
||||
chat_message.c.timestamp,
|
||||
)
|
||||
.where(chat_message.c.session_id == sid, chat_message.c.id.in_(ids_sq))
|
||||
.order_by(chat_message.c.id.asc())
|
||||
)
|
||||
prepended: set[int] = set()
|
||||
with self._engine.connect() as conn:
|
||||
rows: list[dict[str, Any]] = [dict(r) for r in conn.execute(stmt).mappings().all()]
|
||||
while rows:
|
||||
first = rows[0]
|
||||
if str(first.get("role") or "") != "tool":
|
||||
break
|
||||
aid = _tool_row_assistant_message_id(first.get("tool_calls"))
|
||||
if aid is None:
|
||||
break
|
||||
first_id = int(first["id"])
|
||||
if aid >= first_id:
|
||||
break
|
||||
if any(int(r["id"]) == int(aid) for r in rows):
|
||||
break
|
||||
if int(aid) in prepended:
|
||||
break
|
||||
arow = conn.execute(
|
||||
select(
|
||||
chat_message.c.id,
|
||||
chat_message.c.session_id,
|
||||
chat_message.c.role,
|
||||
chat_message.c.content,
|
||||
chat_message.c.tool_calls,
|
||||
chat_message.c.attachments,
|
||||
chat_message.c.turn_uuid,
|
||||
chat_message.c.event_type,
|
||||
chat_message.c.event_payload,
|
||||
chat_message.c.timestamp,
|
||||
)
|
||||
.where(chat_message.c.session_id == sid, chat_message.c.id == int(aid))
|
||||
.limit(1)
|
||||
).mappings().first()
|
||||
if not arow:
|
||||
break
|
||||
prepended.add(int(aid))
|
||||
rows.insert(0, dict(arow))
|
||||
return [_row_to_chat_message(r) for r in rows]
|
||||
|
||||
def get_messages_after_id(
|
||||
self, *, session_id: str, after_id: int, limit: int
|
||||
) -> list[ChatMessage]:
|
||||
sid = str(session_id or "").strip()
|
||||
if not sid:
|
||||
return []
|
||||
aid = int(after_id or 0)
|
||||
lim = max(1, min(int(limit), 2000))
|
||||
stmt = (
|
||||
select(
|
||||
chat_message.c.id,
|
||||
chat_message.c.session_id,
|
||||
chat_message.c.role,
|
||||
chat_message.c.content,
|
||||
chat_message.c.tool_calls,
|
||||
chat_message.c.attachments,
|
||||
chat_message.c.turn_uuid,
|
||||
chat_message.c.event_type,
|
||||
chat_message.c.event_payload,
|
||||
chat_message.c.timestamp,
|
||||
)
|
||||
.where(chat_message.c.session_id == sid, chat_message.c.id > aid)
|
||||
.order_by(chat_message.c.id.asc())
|
||||
.limit(lim)
|
||||
)
|
||||
with self._engine.connect() as conn:
|
||||
rows = conn.execute(stmt).mappings().all()
|
||||
return [_row_to_chat_message(dict(r)) for r in rows]
|
||||
|
||||
def count_messages(self, *, session_id: str) -> int:
|
||||
key = str(session_id if session_id is not None else "")
|
||||
with self._engine.connect() as conn:
|
||||
n = conn.execute(
|
||||
select(func.count()).select_from(chat_message).where(chat_message.c.session_id == key)
|
||||
).scalar_one()
|
||||
return int(n or 0)
|
||||
|
||||
def session_messages_meta(self, *, session_id: str) -> SessionMessagesMeta:
|
||||
key = str(session_id if session_id is not None else "")
|
||||
with self._engine.connect() as conn:
|
||||
row = conn.execute(
|
||||
select(
|
||||
func.count().label("c"),
|
||||
func.max(chat_message.c.id).label("last_id"),
|
||||
func.max(chat_message.c.timestamp).label("last_ts"),
|
||||
)
|
||||
.where(chat_message.c.session_id == key)
|
||||
).mappings().first()
|
||||
return SessionMessagesMeta(
|
||||
session_id=session_id,
|
||||
message_count=int(row["c"] or 0) if row else 0,
|
||||
last_message_id=int(row["last_id"]) if row and row.get("last_id") is not None else None,
|
||||
last_message_at=str(row["last_ts"]) if row and row.get("last_ts") is not None else None,
|
||||
)
|
||||
|
||||
def last_message_id(self, *, session_id: str) -> int | None:
|
||||
key = str(session_id if session_id is not None else "")
|
||||
with self._engine.connect() as conn:
|
||||
m = conn.execute(
|
||||
select(func.max(chat_message.c.id)).where(chat_message.c.session_id == key)
|
||||
).scalar_one_or_none()
|
||||
if m is None:
|
||||
return None
|
||||
return int(m)
|
||||
|
||||
def list_messages_in_time_window(
|
||||
self, *, session_id: str, start_ts: str, end_ts: str, limit: int
|
||||
) -> list[dict[str, Any]]:
|
||||
sid = str(session_id or "").strip()
|
||||
start = str(start_ts or "").strip()
|
||||
end = str(end_ts or "").strip()
|
||||
if not sid or not start or not end:
|
||||
return []
|
||||
lim = max(1, min(int(limit), 2000))
|
||||
stmt = (
|
||||
select(
|
||||
chat_message.c.id,
|
||||
chat_message.c.session_id,
|
||||
chat_message.c.role,
|
||||
chat_message.c.content,
|
||||
chat_message.c.tool_calls,
|
||||
chat_message.c.attachments,
|
||||
chat_message.c.turn_uuid,
|
||||
chat_message.c.event_type,
|
||||
chat_message.c.event_payload,
|
||||
chat_message.c.timestamp,
|
||||
)
|
||||
.where(
|
||||
chat_message.c.session_id == sid,
|
||||
chat_message.c.timestamp >= start,
|
||||
chat_message.c.timestamp <= end,
|
||||
)
|
||||
.order_by(chat_message.c.id.asc())
|
||||
.limit(lim)
|
||||
)
|
||||
with self._engine.connect() as conn:
|
||||
rows = conn.execute(stmt).mappings().all()
|
||||
out: list[dict[str, Any]] = []
|
||||
for r in rows:
|
||||
out.append(
|
||||
{
|
||||
"id": int(r["id"] or 0),
|
||||
"session_id": str(r["session_id"] or ""),
|
||||
"role": str(r["role"] or ""),
|
||||
"content": str(r["content"] or ""),
|
||||
"tool_calls": r["tool_calls"],
|
||||
"attachments": r["attachments"],
|
||||
"turn_uuid": str(r["turn_uuid"] or ""),
|
||||
"event_type": str(r["event_type"] or ""),
|
||||
"event_payload": r["event_payload"],
|
||||
"timestamp": str(r["timestamp"] or ""),
|
||||
}
|
||||
)
|
||||
return out
|
||||
|
||||
def delete_messages_where_session_missing(self) -> int:
|
||||
"""Delete ``chat_message`` rows whose ``session_id`` is not in ``chat_session`` (housekeeping)."""
|
||||
sess_exists = exists(
|
||||
select(1).select_from(chat_session).where(chat_session.c.id == chat_message.c.session_id)
|
||||
)
|
||||
with self._engine.begin() as conn:
|
||||
res = conn.execute(delete(chat_message).where(~sess_exists))
|
||||
return int(res.rowcount or 0)
|
||||
|
||||
def fork_assert_anchor(self, *, source_session_id: str, up_to_message_id: int) -> None:
|
||||
"""Raise ``ValueError`` unless ``up_to_message_id`` exists in ``source_session_id``."""
|
||||
src = str(source_session_id if source_session_id is not None else "")
|
||||
cap = int(up_to_message_id or 0)
|
||||
if not src or cap <= 0:
|
||||
raise ValueError("message not in session")
|
||||
with self._engine.connect() as conn:
|
||||
chk = conn.execute(
|
||||
select(chat_message.c.id)
|
||||
.where(chat_message.c.session_id == src, chat_message.c.id == cap)
|
||||
.limit(1)
|
||||
).first()
|
||||
if not chk:
|
||||
raise ValueError("message not in session")
|
||||
|
||||
def fork_copy_messages_to_session(
|
||||
self,
|
||||
*,
|
||||
source_session_id: str,
|
||||
up_to_message_id: int,
|
||||
new_session_id: str,
|
||||
) -> None:
|
||||
"""Copy messages with ``id <= up_to_message_id`` into ``new_session_id``; remap tool assistant ids."""
|
||||
src = str(source_session_id if source_session_id is not None else "")
|
||||
new_sid = str(new_session_id if new_session_id is not None else "")
|
||||
cap = int(up_to_message_id or 0)
|
||||
if not src or not new_sid or cap <= 0:
|
||||
raise ValueError("message not in session")
|
||||
with self._engine.begin() as conn:
|
||||
rows_list = list(
|
||||
conn.execute(
|
||||
select(
|
||||
chat_message.c.id,
|
||||
chat_message.c.role,
|
||||
chat_message.c.content,
|
||||
chat_message.c.tool_calls,
|
||||
chat_message.c.attachments,
|
||||
chat_message.c.turn_uuid,
|
||||
chat_message.c.event_type,
|
||||
chat_message.c.event_payload,
|
||||
chat_message.c.timestamp,
|
||||
)
|
||||
.where(chat_message.c.session_id == src, chat_message.c.id <= cap)
|
||||
.order_by(chat_message.c.id.asc())
|
||||
).mappings().all()
|
||||
)
|
||||
if not rows_list:
|
||||
raise ValueError("message not in session")
|
||||
id_map: dict[int, int] = {}
|
||||
for r in rows_list:
|
||||
old_id = int(r["id"])
|
||||
role = str(r["role"])
|
||||
tool_calls_text = r["tool_calls"]
|
||||
if role == "tool" and tool_calls_text:
|
||||
try:
|
||||
meta = json.loads(str(tool_calls_text))
|
||||
if isinstance(meta, dict):
|
||||
aid = meta.get("assistant_message_id")
|
||||
if aid is not None:
|
||||
new_aid = id_map.get(int(aid))
|
||||
if new_aid is not None:
|
||||
meta = {**meta, "assistant_message_id": new_aid}
|
||||
tool_calls_text = json.dumps(meta, ensure_ascii=False)
|
||||
except (json.JSONDecodeError, TypeError, ValueError):
|
||||
pass
|
||||
stmt = (
|
||||
insert(chat_message)
|
||||
.values(
|
||||
session_id=new_sid,
|
||||
role=role,
|
||||
content=r["content"],
|
||||
tool_calls=tool_calls_text,
|
||||
attachments=r["attachments"],
|
||||
turn_uuid=r["turn_uuid"],
|
||||
event_type=r["event_type"],
|
||||
event_payload=r["event_payload"],
|
||||
timestamp=r["timestamp"],
|
||||
)
|
||||
.returning(chat_message.c.id)
|
||||
)
|
||||
new_id = int(conn.execute(stmt).scalar_one())
|
||||
id_map[old_id] = new_id
|
||||
last_ts = rows_list[-1]["timestamp"] if rows_list else utc_now_iso()
|
||||
conn.execute(
|
||||
update(chat_session)
|
||||
.where(chat_session.c.id == new_sid)
|
||||
.values(last_message_at=last_ts)
|
||||
)
|
||||
|
||||
def trim_messages_keep_last(self, *, session_id: str, keep_last: int) -> None:
|
||||
"""Delete older messages so at least ``keep_last`` newest rows remain (tool/assistant boundary aware)."""
|
||||
key = str(session_id if session_id is not None else "")
|
||||
kl = int(keep_last)
|
||||
if not key or kl <= 0:
|
||||
return
|
||||
with self._engine.connect() as conn:
|
||||
rows_list = [
|
||||
dict(r)
|
||||
for r in conn.execute(
|
||||
select(chat_message.c.id, chat_message.c.role, chat_message.c.tool_calls)
|
||||
.where(chat_message.c.session_id == key)
|
||||
.order_by(chat_message.c.id.asc())
|
||||
).mappings().all()
|
||||
]
|
||||
start = _trim_messages_start_index(rows_list, kl)
|
||||
if start is None:
|
||||
return
|
||||
min_keep_id = int(rows_list[start]["id"])
|
||||
with self._engine.begin() as conn:
|
||||
conn.execute(
|
||||
delete(chat_message).where(
|
||||
chat_message.c.session_id == key,
|
||||
chat_message.c.id < min_keep_id,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["ChatMessagesSaRepository"]
|
||||
424
svc/persistence/sa_repos/chat_sessions.py
Normal file
424
svc/persistence/sa_repos/chat_sessions.py
Normal file
|
|
@ -0,0 +1,424 @@
|
|||
"""chat_session (+ ui_session_owner joins) via SQLAlchemy Core."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Mapping
|
||||
|
||||
from sqlalchemy import and_, delete, distinct, func, insert, literal, or_, select, update
|
||||
from sqlalchemy.engine import Engine
|
||||
|
||||
from svc.persistence.db.tables import app_user, chat_message, chat_session, ui_session_owner
|
||||
from svc.persistence.sqlite_store import ChatSession, SessionsListMeta
|
||||
|
||||
|
||||
def _session_from_row(row: Mapping[str, Any]) -> ChatSession:
|
||||
return ChatSession(
|
||||
id=str(row["id"]),
|
||||
title=str(row["title"]),
|
||||
created_at=str(row["created_at"]),
|
||||
last_message_at=row["last_message_at"],
|
||||
)
|
||||
|
||||
|
||||
def _activity_order() -> tuple[Any, Any]:
|
||||
return (
|
||||
func.coalesce(chat_session.c.last_message_at, chat_session.c.created_at).desc(),
|
||||
chat_session.c.created_at.desc(),
|
||||
)
|
||||
|
||||
|
||||
class ChatSessionsSaRepository:
|
||||
"""Phase-3 SA migration: chat session rows and list queries (messages stay raw SQL for now)."""
|
||||
|
||||
__slots__ = ("_engine",)
|
||||
|
||||
def __init__(self, engine: Engine) -> None:
|
||||
self._engine = engine
|
||||
|
||||
def insert_chat_session(self, *, session_id: str, title: str, created_at: str) -> None:
|
||||
with self._engine.begin() as conn:
|
||||
conn.execute(
|
||||
insert(chat_session).values(
|
||||
id=str(session_id),
|
||||
title=str(title),
|
||||
created_at=str(created_at),
|
||||
last_message_at=None,
|
||||
)
|
||||
)
|
||||
|
||||
def fetch_chat_session_by_id(self, *, session_id: str) -> ChatSession | None:
|
||||
sid = str(session_id or "").strip()
|
||||
if not sid:
|
||||
return None
|
||||
with self._engine.connect() as conn:
|
||||
row = conn.execute(
|
||||
select(
|
||||
chat_session.c.id,
|
||||
chat_session.c.title,
|
||||
chat_session.c.created_at,
|
||||
chat_session.c.last_message_at,
|
||||
)
|
||||
.where(chat_session.c.id == sid)
|
||||
.limit(1)
|
||||
).mappings().first()
|
||||
return _session_from_row(row) if row else None
|
||||
|
||||
def fetch_chat_session_for_user(
|
||||
self, *, session_id: str, tenant_id: str, user_id: str
|
||||
) -> ChatSession | None:
|
||||
sid = str(session_id or "").strip()
|
||||
if not sid:
|
||||
return None
|
||||
tid, uid = str(tenant_id), str(user_id)
|
||||
with self._engine.connect() as conn:
|
||||
row = conn.execute(
|
||||
select(
|
||||
chat_session.c.id,
|
||||
chat_session.c.title,
|
||||
chat_session.c.created_at,
|
||||
chat_session.c.last_message_at,
|
||||
)
|
||||
.select_from(
|
||||
chat_session.join(
|
||||
ui_session_owner,
|
||||
ui_session_owner.c.session_id == chat_session.c.id,
|
||||
)
|
||||
)
|
||||
.where(
|
||||
chat_session.c.id == sid,
|
||||
ui_session_owner.c.tenant_id == tid,
|
||||
ui_session_owner.c.user_id == uid,
|
||||
)
|
||||
.limit(1)
|
||||
).mappings().first()
|
||||
return _session_from_row(row) if row else None
|
||||
|
||||
def list_chat_sessions_global(self, *, limit: int | None, offset: int) -> list[ChatSession]:
|
||||
stmt = select(
|
||||
chat_session.c.id,
|
||||
chat_session.c.title,
|
||||
chat_session.c.created_at,
|
||||
chat_session.c.last_message_at,
|
||||
).order_by(*_activity_order())
|
||||
if limit is not None:
|
||||
stmt = stmt.limit(int(limit)).offset(int(offset))
|
||||
with self._engine.connect() as conn:
|
||||
rows = conn.execute(stmt).mappings().all()
|
||||
return [_session_from_row(r) for r in rows]
|
||||
|
||||
def list_chat_sessions_for_user(
|
||||
self,
|
||||
*,
|
||||
tenant_id: str,
|
||||
user_id: str,
|
||||
limit: int | None,
|
||||
offset: int,
|
||||
) -> list[ChatSession]:
|
||||
tid, uid = str(tenant_id), str(user_id)
|
||||
stmt = (
|
||||
select(
|
||||
chat_session.c.id,
|
||||
chat_session.c.title,
|
||||
chat_session.c.created_at,
|
||||
chat_session.c.last_message_at,
|
||||
)
|
||||
.select_from(
|
||||
chat_session.join(
|
||||
ui_session_owner,
|
||||
ui_session_owner.c.session_id == chat_session.c.id,
|
||||
)
|
||||
)
|
||||
.where(
|
||||
ui_session_owner.c.tenant_id == tid,
|
||||
ui_session_owner.c.user_id == uid,
|
||||
)
|
||||
.order_by(*_activity_order())
|
||||
)
|
||||
if limit is not None:
|
||||
stmt = stmt.limit(int(limit)).offset(int(offset))
|
||||
with self._engine.connect() as conn:
|
||||
rows = conn.execute(stmt).mappings().all()
|
||||
return [_session_from_row(r) for r in rows]
|
||||
|
||||
def list_chat_sessions_for_tenant(
|
||||
self,
|
||||
*,
|
||||
tenant_id: str,
|
||||
limit: int | None,
|
||||
offset: int,
|
||||
) -> list[ChatSession]:
|
||||
tid = str(tenant_id)
|
||||
stmt = (
|
||||
select(
|
||||
chat_session.c.id,
|
||||
chat_session.c.title,
|
||||
chat_session.c.created_at,
|
||||
chat_session.c.last_message_at,
|
||||
)
|
||||
.select_from(
|
||||
chat_session.join(
|
||||
ui_session_owner,
|
||||
ui_session_owner.c.session_id == chat_session.c.id,
|
||||
)
|
||||
)
|
||||
.where(ui_session_owner.c.tenant_id == tid)
|
||||
.distinct()
|
||||
.order_by(*_activity_order())
|
||||
)
|
||||
if limit is not None:
|
||||
stmt = stmt.limit(int(limit)).offset(int(offset))
|
||||
with self._engine.connect() as conn:
|
||||
rows = conn.execute(stmt).mappings().all()
|
||||
return [_session_from_row(r) for r in rows]
|
||||
|
||||
def count_chat_sessions_global(self) -> int:
|
||||
with self._engine.connect() as conn:
|
||||
n = conn.execute(select(func.count()).select_from(chat_session)).scalar_one()
|
||||
return int(n or 0)
|
||||
|
||||
def sessions_list_meta_global(self) -> SessionsListMeta:
|
||||
with self._engine.connect() as conn:
|
||||
row = conn.execute(
|
||||
select(
|
||||
func.count().label("c"),
|
||||
func.max(
|
||||
func.coalesce(chat_session.c.last_message_at, chat_session.c.created_at)
|
||||
).label("latest_activity_at"),
|
||||
).select_from(chat_session)
|
||||
).mappings().first()
|
||||
return SessionsListMeta(
|
||||
session_count=int(row["c"] or 0) if row else 0,
|
||||
latest_activity_at=str(row["latest_activity_at"])
|
||||
if row and row.get("latest_activity_at") is not None
|
||||
else None,
|
||||
)
|
||||
|
||||
def sessions_list_meta_for_user(self, *, tenant_id: str, user_id: str) -> SessionsListMeta:
|
||||
tid, uid = str(tenant_id), str(user_id)
|
||||
with self._engine.connect() as conn:
|
||||
row = conn.execute(
|
||||
select(
|
||||
func.count().label("c"),
|
||||
func.max(
|
||||
func.coalesce(chat_session.c.last_message_at, chat_session.c.created_at)
|
||||
).label("latest_activity_at"),
|
||||
)
|
||||
.select_from(
|
||||
chat_session.join(
|
||||
ui_session_owner,
|
||||
ui_session_owner.c.session_id == chat_session.c.id,
|
||||
)
|
||||
)
|
||||
.where(
|
||||
ui_session_owner.c.tenant_id == tid,
|
||||
ui_session_owner.c.user_id == uid,
|
||||
)
|
||||
).mappings().first()
|
||||
return SessionsListMeta(
|
||||
session_count=int(row["c"] or 0) if row else 0,
|
||||
latest_activity_at=str(row["latest_activity_at"])
|
||||
if row and row.get("latest_activity_at") is not None
|
||||
else None,
|
||||
)
|
||||
|
||||
def sessions_list_meta_for_tenant(self, *, tenant_id: str) -> SessionsListMeta:
|
||||
tid = str(tenant_id)
|
||||
with self._engine.connect() as conn:
|
||||
row = conn.execute(
|
||||
select(
|
||||
func.count(distinct(chat_session.c.id)).label("c"),
|
||||
func.max(
|
||||
func.coalesce(chat_session.c.last_message_at, chat_session.c.created_at)
|
||||
).label("latest_activity_at"),
|
||||
)
|
||||
.select_from(
|
||||
chat_session.join(
|
||||
ui_session_owner,
|
||||
ui_session_owner.c.session_id == chat_session.c.id,
|
||||
)
|
||||
)
|
||||
.where(ui_session_owner.c.tenant_id == tid)
|
||||
).mappings().first()
|
||||
return SessionsListMeta(
|
||||
session_count=int(row["c"] or 0) if row else 0,
|
||||
latest_activity_at=str(row["latest_activity_at"])
|
||||
if row and row.get("latest_activity_at") is not None
|
||||
else None,
|
||||
)
|
||||
|
||||
def fetch_chat_session_in_tenant(self, *, session_id: str, tenant_id: str) -> ChatSession | None:
|
||||
sid, tid = str(session_id or "").strip(), str(tenant_id)
|
||||
if not sid:
|
||||
return None
|
||||
with self._engine.connect() as conn:
|
||||
row = conn.execute(
|
||||
select(
|
||||
chat_session.c.id,
|
||||
chat_session.c.title,
|
||||
chat_session.c.created_at,
|
||||
chat_session.c.last_message_at,
|
||||
)
|
||||
.select_from(
|
||||
chat_session.join(
|
||||
ui_session_owner,
|
||||
ui_session_owner.c.session_id == chat_session.c.id,
|
||||
)
|
||||
)
|
||||
.where(
|
||||
chat_session.c.id == sid,
|
||||
ui_session_owner.c.tenant_id == tid,
|
||||
)
|
||||
.limit(1)
|
||||
).mappings().first()
|
||||
return _session_from_row(row) if row else None
|
||||
|
||||
def list_admin_sessions(
|
||||
self,
|
||||
*,
|
||||
tenant_id: str,
|
||||
user_id: str | None,
|
||||
search_lower: str | None,
|
||||
active_only: bool,
|
||||
active_cutoff_iso: str,
|
||||
limit: int,
|
||||
offset: int,
|
||||
) -> tuple[int, list[dict[str, Any]]]:
|
||||
"""Tenant-scoped admin session browser (parity with raw ``list_admin_sessions`` SQL)."""
|
||||
tid = str(tenant_id or "").strip()
|
||||
if not tid:
|
||||
return 0, []
|
||||
uid = str(user_id or "").strip() or None
|
||||
q_text = str(search_lower or "").strip().lower() or None
|
||||
lim = max(1, min(int(limit), 500))
|
||||
off = max(0, int(offset))
|
||||
|
||||
msg_cnt = (
|
||||
select(func.count())
|
||||
.select_from(chat_message)
|
||||
.where(chat_message.c.session_id == chat_session.c.id)
|
||||
.scalar_subquery()
|
||||
)
|
||||
joins = chat_session.join(
|
||||
ui_session_owner,
|
||||
ui_session_owner.c.session_id == chat_session.c.id,
|
||||
).outerjoin(
|
||||
app_user,
|
||||
(app_user.c.tenant_id == ui_session_owner.c.tenant_id)
|
||||
& (app_user.c.id == ui_session_owner.c.user_id),
|
||||
)
|
||||
conds: list[Any] = [ui_session_owner.c.tenant_id == tid]
|
||||
if uid:
|
||||
conds.append(ui_session_owner.c.user_id == uid)
|
||||
if q_text:
|
||||
like = f"%{q_text}%"
|
||||
conds.append(
|
||||
or_(
|
||||
func.lower(func.coalesce(app_user.c.username, literal(""))).like(like),
|
||||
func.lower(func.coalesce(app_user.c.display_name, literal(""))).like(like),
|
||||
func.lower(func.coalesce(chat_session.c.title, literal(""))).like(like),
|
||||
)
|
||||
)
|
||||
if active_only:
|
||||
conds.append(
|
||||
func.coalesce(chat_session.c.last_message_at, chat_session.c.created_at) >= active_cutoff_iso
|
||||
)
|
||||
wh = and_(*conds)
|
||||
|
||||
data_stmt = (
|
||||
select(
|
||||
chat_session.c.id.label("session_id"),
|
||||
chat_session.c.title,
|
||||
chat_session.c.created_at,
|
||||
chat_session.c.last_message_at,
|
||||
ui_session_owner.c.user_id,
|
||||
func.coalesce(app_user.c.username, literal("")).label("username"),
|
||||
func.coalesce(app_user.c.display_name, literal("")).label("display_name"),
|
||||
msg_cnt.label("message_count"),
|
||||
)
|
||||
.select_from(joins)
|
||||
.where(wh)
|
||||
.order_by(
|
||||
func.coalesce(chat_session.c.last_message_at, chat_session.c.created_at).desc(),
|
||||
chat_session.c.created_at.desc(),
|
||||
)
|
||||
.limit(lim)
|
||||
.offset(off)
|
||||
)
|
||||
cnt_stmt = select(func.count()).select_from(joins).where(wh)
|
||||
|
||||
with self._engine.connect() as conn:
|
||||
total = int(conn.execute(cnt_stmt).scalar_one() or 0)
|
||||
rows = conn.execute(data_stmt).mappings().all()
|
||||
out: list[dict[str, Any]] = []
|
||||
for r in rows:
|
||||
out.append(
|
||||
{
|
||||
"session_id": str(r["session_id"] or ""),
|
||||
"title": str(r["title"] or ""),
|
||||
"created_at": str(r["created_at"] or ""),
|
||||
"last_message_at": str(r["last_message_at"] or ""),
|
||||
"user_id": str(r["user_id"] or ""),
|
||||
"username": str(r["username"] or ""),
|
||||
"display_name": str(r["display_name"] or ""),
|
||||
"message_count": int(r["message_count"] or 0),
|
||||
}
|
||||
)
|
||||
return total, out
|
||||
|
||||
def rename_chat_session(self, *, session_id: str, title: str) -> None:
|
||||
with self._engine.begin() as conn:
|
||||
conn.execute(
|
||||
update(chat_session)
|
||||
.where(chat_session.c.id == str(session_id))
|
||||
.values(title=str(title))
|
||||
)
|
||||
|
||||
def delete_chat_session_by_id(self, *, session_id: str) -> None:
|
||||
sid = str(session_id)
|
||||
with self._engine.begin() as conn:
|
||||
conn.execute(delete(chat_session).where(chat_session.c.id == sid))
|
||||
|
||||
def try_delete_chat_session_for_tenant(self, *, session_id: str, tenant_id: str) -> bool:
|
||||
sid, tid = str(session_id or "").strip(), str(tenant_id)
|
||||
if not sid:
|
||||
return False
|
||||
with self._engine.begin() as conn:
|
||||
chk = conn.execute(
|
||||
select(1)
|
||||
.select_from(ui_session_owner)
|
||||
.where(
|
||||
ui_session_owner.c.session_id == sid,
|
||||
ui_session_owner.c.tenant_id == tid,
|
||||
)
|
||||
.limit(1)
|
||||
).first()
|
||||
if not chk:
|
||||
return False
|
||||
conn.execute(delete(chat_session).where(chat_session.c.id == sid))
|
||||
return True
|
||||
|
||||
def try_delete_chat_session_for_user(
|
||||
self, *, session_id: str, tenant_id: str, user_id: str
|
||||
) -> bool:
|
||||
sid = str(session_id or "").strip()
|
||||
if not sid:
|
||||
return False
|
||||
tid, uid = str(tenant_id), str(user_id)
|
||||
with self._engine.begin() as conn:
|
||||
chk = conn.execute(
|
||||
select(1)
|
||||
.select_from(ui_session_owner)
|
||||
.where(
|
||||
ui_session_owner.c.session_id == sid,
|
||||
ui_session_owner.c.tenant_id == tid,
|
||||
ui_session_owner.c.user_id == uid,
|
||||
)
|
||||
.limit(1)
|
||||
).first()
|
||||
if not chk:
|
||||
return False
|
||||
conn.execute(delete(chat_session).where(chat_session.c.id == sid))
|
||||
return True
|
||||
|
||||
|
||||
__all__ = ["ChatSessionsSaRepository"]
|
||||
77
svc/persistence/sa_repos/session_tool_health.py
Normal file
77
svc/persistence/sa_repos/session_tool_health.py
Normal file
|
|
@ -0,0 +1,77 @@
|
|||
"""Session tool health listing (admin) via SQLAlchemy Core."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import case, func, literal, select
|
||||
from sqlalchemy.engine import Engine
|
||||
|
||||
from svc.persistence.db.tables import chat_message, chat_session, tool_log
|
||||
|
||||
|
||||
class SessionToolHealthSaRepository:
|
||||
"""``list_session_tool_health`` aggregates."""
|
||||
|
||||
__slots__ = ("_engine",)
|
||||
|
||||
def __init__(self, engine: Engine) -> None:
|
||||
self._engine = engine
|
||||
|
||||
def list_session_tool_health(
|
||||
self, *, session_id: str | None, limit: int
|
||||
) -> list[dict[str, Any]]:
|
||||
lim = max(1, int(limit))
|
||||
msg_sq = (
|
||||
select(
|
||||
chat_message.c.session_id,
|
||||
func.sum(case((chat_message.c.role == literal("user"), 1), else_=0)).label("user_count"),
|
||||
func.sum(case((chat_message.c.role == literal("assistant"), 1), else_=0)).label(
|
||||
"assistant_count"
|
||||
),
|
||||
)
|
||||
.group_by(chat_message.c.session_id)
|
||||
.subquery()
|
||||
)
|
||||
tl_sq = (
|
||||
select(
|
||||
tool_log.c.session_id,
|
||||
func.count(1).label("tool_count"),
|
||||
func.sum(case((tool_log.c.tool_name.like("mcp__%"), 1), else_=0)).label(
|
||||
"mcp_tool_count"
|
||||
),
|
||||
func.max(tool_log.c.timestamp).label("last_tool_at"),
|
||||
)
|
||||
.group_by(tool_log.c.session_id)
|
||||
.subquery()
|
||||
)
|
||||
stmt = (
|
||||
select(
|
||||
chat_session.c.id.label("session_id"),
|
||||
chat_session.c.title,
|
||||
chat_session.c.last_message_at,
|
||||
func.coalesce(msg_sq.c.user_count, literal(0)).label("user_count"),
|
||||
func.coalesce(msg_sq.c.assistant_count, literal(0)).label("assistant_count"),
|
||||
func.coalesce(tl_sq.c.tool_count, literal(0)).label("tool_count"),
|
||||
func.coalesce(tl_sq.c.mcp_tool_count, literal(0)).label("mcp_tool_count"),
|
||||
func.coalesce(tl_sq.c.last_tool_at, literal("")).label("last_tool_at"),
|
||||
)
|
||||
.select_from(
|
||||
chat_session.outerjoin(msg_sq, msg_sq.c.session_id == chat_session.c.id).outerjoin(
|
||||
tl_sq, tl_sq.c.session_id == chat_session.c.id
|
||||
)
|
||||
)
|
||||
.order_by(
|
||||
func.coalesce(chat_session.c.last_message_at, chat_session.c.created_at).desc(),
|
||||
)
|
||||
.limit(lim)
|
||||
)
|
||||
sid = str(session_id or "").strip()
|
||||
if sid:
|
||||
stmt = stmt.where(chat_session.c.id == sid)
|
||||
with self._engine.connect() as conn:
|
||||
rows = conn.execute(stmt).mappings().all()
|
||||
return [dict(r) for r in rows]
|
||||
|
||||
|
||||
__all__ = ["SessionToolHealthSaRepository"]
|
||||
116
svc/persistence/sa_repos/tenant_bind_code.py
Normal file
116
svc/persistence/sa_repos/tenant_bind_code.py
Normal file
|
|
@ -0,0 +1,116 @@
|
|||
"""tenant + bind_code via SQLAlchemy Core."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Mapping
|
||||
|
||||
from sqlalchemy import delete, insert, select, update
|
||||
from sqlalchemy.engine import Engine
|
||||
|
||||
from svc.persistence.db.tables import bind_code, tenant
|
||||
|
||||
|
||||
class TenantSaRepository:
|
||||
__slots__ = ("_engine",)
|
||||
|
||||
def __init__(self, engine: Engine) -> None:
|
||||
self._engine = engine
|
||||
|
||||
def insert_tenant(self, *, tenant_id: str, name: str, created_at: str) -> None:
|
||||
with self._engine.begin() as conn:
|
||||
conn.execute(
|
||||
insert(tenant).values(
|
||||
id=str(tenant_id),
|
||||
name=str(name),
|
||||
created_at=str(created_at),
|
||||
)
|
||||
)
|
||||
|
||||
def delete_tenant(self, *, tenant_id: str) -> int:
|
||||
tid = str(tenant_id or "").strip()
|
||||
if not tid:
|
||||
return 0
|
||||
with self._engine.begin() as conn:
|
||||
res = conn.execute(delete(tenant).where(tenant.c.id == tid))
|
||||
return int(res.rowcount or 0)
|
||||
|
||||
def list_tenants(self, *, limit: int) -> list[dict[str, Any]]:
|
||||
lim = max(1, int(limit))
|
||||
stmt = (
|
||||
select(tenant.c.id, tenant.c.name, tenant.c.created_at)
|
||||
.order_by(tenant.c.created_at.desc())
|
||||
.limit(lim)
|
||||
)
|
||||
with self._engine.connect() as conn:
|
||||
return [dict(r) for r in conn.execute(stmt).mappings().all()]
|
||||
|
||||
|
||||
class BindCodeSaRepository:
|
||||
__slots__ = ("_engine",)
|
||||
|
||||
def __init__(self, engine: Engine) -> None:
|
||||
self._engine = engine
|
||||
|
||||
def insert_bind_code(
|
||||
self, *, code: str, tenant_id: str, role: str, created_at: str
|
||||
) -> None:
|
||||
with self._engine.begin() as conn:
|
||||
conn.execute(
|
||||
insert(bind_code).values(
|
||||
code=str(code),
|
||||
tenant_id=str(tenant_id),
|
||||
role=str(role),
|
||||
created_at=str(created_at),
|
||||
used_at=None,
|
||||
used_by_external_user_id=None,
|
||||
)
|
||||
)
|
||||
|
||||
def fetch_by_code(self, *, code: str) -> Mapping[str, Any] | None:
|
||||
c = str(code or "").strip()
|
||||
if not c:
|
||||
return None
|
||||
stmt = (
|
||||
select(
|
||||
bind_code.c.code,
|
||||
bind_code.c.tenant_id,
|
||||
bind_code.c.role,
|
||||
bind_code.c.created_at,
|
||||
bind_code.c.used_at,
|
||||
bind_code.c.used_by_external_user_id,
|
||||
)
|
||||
.where(bind_code.c.code == c)
|
||||
.limit(1)
|
||||
)
|
||||
with self._engine.connect() as conn:
|
||||
return conn.execute(stmt).mappings().first()
|
||||
|
||||
def mark_used(
|
||||
self, *, code: str, used_at: str, used_by_external_user_id: str
|
||||
) -> None:
|
||||
c = str(code or "").strip()
|
||||
with self._engine.begin() as conn:
|
||||
conn.execute(
|
||||
update(bind_code)
|
||||
.where(bind_code.c.code == c)
|
||||
.values(used_at=str(used_at), used_by_external_user_id=str(used_by_external_user_id))
|
||||
)
|
||||
|
||||
def list_bind_codes(self, *, tenant_id: str | None, limit: int) -> list[dict[str, Any]]:
|
||||
lim = max(1, int(limit))
|
||||
stmt = select(
|
||||
bind_code.c.code,
|
||||
bind_code.c.tenant_id,
|
||||
bind_code.c.role,
|
||||
bind_code.c.created_at,
|
||||
bind_code.c.used_at,
|
||||
bind_code.c.used_by_external_user_id,
|
||||
).order_by(bind_code.c.created_at.desc())
|
||||
if tenant_id:
|
||||
stmt = stmt.where(bind_code.c.tenant_id == str(tenant_id))
|
||||
stmt = stmt.limit(lim)
|
||||
with self._engine.connect() as conn:
|
||||
return [dict(r) for r in conn.execute(stmt).mappings().all()]
|
||||
|
||||
|
||||
__all__ = ["BindCodeSaRepository", "TenantSaRepository"]
|
||||
133
svc/persistence/sa_repos/tool_log_queries.py
Normal file
133
svc/persistence/sa_repos/tool_log_queries.py
Normal file
|
|
@ -0,0 +1,133 @@
|
|||
"""tool_log read paths (MCP summaries, call logs) via SQLAlchemy Core."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import and_, delete, exists, func, insert, select, update
|
||||
from sqlalchemy.engine import Engine
|
||||
|
||||
from svc.persistence.db.tables import chat_session, tool_log
|
||||
|
||||
|
||||
class ToolLogQueriesSaRepository:
|
||||
__slots__ = ("_engine",)
|
||||
|
||||
def __init__(self, engine: Engine) -> None:
|
||||
self._engine = engine
|
||||
|
||||
def insert_tool_log(
|
||||
self,
|
||||
*,
|
||||
session_id: str,
|
||||
tool_name: str,
|
||||
specialist: str,
|
||||
args: str,
|
||||
result: str,
|
||||
timestamp: str,
|
||||
duration_ms: int | None,
|
||||
) -> None:
|
||||
with self._engine.begin() as conn:
|
||||
conn.execute(
|
||||
insert(tool_log).values(
|
||||
session_id=str(session_id),
|
||||
tool_name=str(tool_name),
|
||||
specialist=str(specialist or ""),
|
||||
args=str(args),
|
||||
result=str(result),
|
||||
timestamp=str(timestamp),
|
||||
duration_ms=duration_ms,
|
||||
)
|
||||
)
|
||||
|
||||
def list_tool_logs_asc(self, *, session_id: str, limit: int) -> list[dict[str, Any]]:
|
||||
sid = str(session_id or "").strip()
|
||||
lim = max(1, int(limit))
|
||||
stmt = (
|
||||
select(
|
||||
tool_log.c.tool_name,
|
||||
tool_log.c.specialist,
|
||||
tool_log.c.args,
|
||||
tool_log.c.result,
|
||||
tool_log.c.timestamp,
|
||||
tool_log.c.duration_ms,
|
||||
)
|
||||
.where(tool_log.c.session_id == sid)
|
||||
.order_by(tool_log.c.id.asc())
|
||||
.limit(lim)
|
||||
)
|
||||
with self._engine.connect() as conn:
|
||||
return [dict(r) for r in conn.execute(stmt).mappings().all()]
|
||||
|
||||
def move_tool_logs_between_sessions(self, *, from_session_id: str, to_session_id: str) -> int:
|
||||
src = str(from_session_id or "").strip()
|
||||
dst = str(to_session_id or "").strip()
|
||||
if not src or not dst or src == dst:
|
||||
return 0
|
||||
with self._engine.begin() as conn:
|
||||
res = conn.execute(
|
||||
update(tool_log)
|
||||
.where(tool_log.c.session_id == src)
|
||||
.values(session_id=dst)
|
||||
)
|
||||
return int(res.rowcount or 0)
|
||||
|
||||
def delete_tool_logs_where_session_missing(self) -> int:
|
||||
"""Delete ``tool_log`` rows whose ``session_id`` is not in ``chat_session`` (housekeeping)."""
|
||||
sess_exists = exists(
|
||||
select(1).select_from(chat_session).where(chat_session.c.id == tool_log.c.session_id)
|
||||
)
|
||||
with self._engine.begin() as conn:
|
||||
res = conn.execute(delete(tool_log).where(~sess_exists))
|
||||
return int(res.rowcount or 0)
|
||||
|
||||
def list_mcp_tool_usage_summary(self, *, limit: int) -> list[dict[str, Any]]:
|
||||
lim = max(1, int(limit))
|
||||
n = func.count(1).label("n")
|
||||
last_ts = func.max(tool_log.c.timestamp).label("last_ts")
|
||||
stmt = (
|
||||
select(tool_log.c.tool_name, tool_log.c.specialist, n, last_ts)
|
||||
.where(tool_log.c.tool_name.like("mcp__%"))
|
||||
.group_by(tool_log.c.tool_name, tool_log.c.specialist)
|
||||
.order_by(n.desc(), last_ts.desc())
|
||||
.limit(lim)
|
||||
)
|
||||
with self._engine.connect() as conn:
|
||||
return [dict(r) for r in conn.execute(stmt).mappings().all()]
|
||||
|
||||
def list_mcp_tool_aggregate_usage(self) -> list[dict[str, Any]]:
|
||||
n = func.count(1).label("n")
|
||||
last_ts = func.max(tool_log.c.timestamp).label("last_ts")
|
||||
stmt = (
|
||||
select(tool_log.c.tool_name, n, last_ts)
|
||||
.where(tool_log.c.tool_name.like("mcp__%"))
|
||||
.group_by(tool_log.c.tool_name)
|
||||
)
|
||||
with self._engine.connect() as conn:
|
||||
return [dict(r) for r in conn.execute(stmt).mappings().all()]
|
||||
|
||||
def list_mcp_tool_call_logs(self, *, server_id: str | None, limit: int) -> list[dict[str, Any]]:
|
||||
lim = max(1, int(limit))
|
||||
sid = str(server_id or "").strip()
|
||||
conds = [tool_log.c.tool_name.like("mcp__%")]
|
||||
if sid:
|
||||
conds.append(tool_log.c.tool_name.like(f"mcp__{sid}__%"))
|
||||
stmt = (
|
||||
select(
|
||||
tool_log.c.session_id,
|
||||
tool_log.c.tool_name,
|
||||
tool_log.c.specialist,
|
||||
tool_log.c.args,
|
||||
tool_log.c.result,
|
||||
tool_log.c.timestamp,
|
||||
tool_log.c.duration_ms,
|
||||
)
|
||||
.where(and_(*conds))
|
||||
.order_by(tool_log.c.id.desc())
|
||||
.limit(lim)
|
||||
)
|
||||
with self._engine.connect() as conn:
|
||||
return [dict(r) for r in conn.execute(stmt).mappings().all()]
|
||||
|
||||
|
||||
__all__ = ["ToolLogQueriesSaRepository"]
|
||||
108
svc/persistence/sa_repos/trace_events.py
Normal file
108
svc/persistence/sa_repos/trace_events.py
Normal file
|
|
@ -0,0 +1,108 @@
|
|||
"""trace_event insert + list queries via SQLAlchemy Core."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from sqlalchemy import insert, select
|
||||
from sqlalchemy.engine import Engine
|
||||
|
||||
from svc.persistence.db.tables import trace_event
|
||||
|
||||
|
||||
class TraceEventsSaRepository:
|
||||
__slots__ = ("_engine",)
|
||||
|
||||
def __init__(self, engine: Engine) -> None:
|
||||
self._engine = engine
|
||||
|
||||
def insert_one(
|
||||
self,
|
||||
*,
|
||||
session_id: str,
|
||||
trace_id: str,
|
||||
span_id: str,
|
||||
parent_span_id: str | None,
|
||||
event_type: str,
|
||||
payload: str,
|
||||
timestamp: str,
|
||||
) -> None:
|
||||
with self._engine.begin() as conn:
|
||||
conn.execute(
|
||||
insert(trace_event).values(
|
||||
session_id=str(session_id),
|
||||
trace_id=str(trace_id),
|
||||
span_id=str(span_id),
|
||||
parent_span_id=parent_span_id,
|
||||
event_type=str(event_type),
|
||||
payload=str(payload),
|
||||
timestamp=str(timestamp),
|
||||
)
|
||||
)
|
||||
|
||||
def insert_many(self, rows: list[dict[str, Any]]) -> None:
|
||||
if not rows:
|
||||
return
|
||||
with self._engine.begin() as conn:
|
||||
conn.execute(insert(trace_event), rows)
|
||||
|
||||
def list_trace_events_desc(self, *, session_id: str, limit: int) -> list[dict[str, Any]]:
|
||||
sid = str(session_id or "").strip()
|
||||
lim = max(1, int(limit))
|
||||
stmt = (
|
||||
select(
|
||||
trace_event.c.trace_id,
|
||||
trace_event.c.span_id,
|
||||
trace_event.c.parent_span_id,
|
||||
trace_event.c.event_type,
|
||||
trace_event.c.payload,
|
||||
trace_event.c.timestamp,
|
||||
)
|
||||
.where(trace_event.c.session_id == sid)
|
||||
.order_by(trace_event.c.id.desc())
|
||||
.limit(lim)
|
||||
)
|
||||
with self._engine.connect() as conn:
|
||||
return [dict(r) for r in conn.execute(stmt).mappings().all()]
|
||||
|
||||
def list_trace_events_for_trace_asc(
|
||||
self, *, session_id: str, trace_id: str, limit: int
|
||||
) -> list[dict[str, Any]]:
|
||||
sid = str(session_id or "").strip()
|
||||
tid = str(trace_id or "").strip()
|
||||
lim = max(1, int(limit))
|
||||
if not sid or not tid:
|
||||
return []
|
||||
stmt = (
|
||||
select(
|
||||
trace_event.c.trace_id,
|
||||
trace_event.c.span_id,
|
||||
trace_event.c.parent_span_id,
|
||||
trace_event.c.event_type,
|
||||
trace_event.c.payload,
|
||||
trace_event.c.timestamp,
|
||||
)
|
||||
.where(trace_event.c.session_id == sid, trace_event.c.trace_id == tid)
|
||||
.order_by(trace_event.c.id.asc())
|
||||
.limit(lim)
|
||||
)
|
||||
with self._engine.connect() as conn:
|
||||
return [dict(r) for r in conn.execute(stmt).mappings().all()]
|
||||
|
||||
def list_event_type_timestamp_for_trace(
|
||||
self, *, session_id: str, trace_id: str
|
||||
) -> list[dict[str, Any]]:
|
||||
sid = str(session_id or "").strip()
|
||||
tid = str(trace_id or "").strip()
|
||||
if not sid or not tid:
|
||||
return []
|
||||
stmt = (
|
||||
select(trace_event.c.event_type, trace_event.c.timestamp)
|
||||
.where(trace_event.c.session_id == sid, trace_event.c.trace_id == tid)
|
||||
.order_by(trace_event.c.id.asc())
|
||||
)
|
||||
with self._engine.connect() as conn:
|
||||
return [dict(r) for r in conn.execute(stmt).mappings().all()]
|
||||
|
||||
|
||||
__all__ = ["TraceEventsSaRepository"]
|
||||
193
svc/persistence/sa_repos/ui_session_owner.py
Normal file
193
svc/persistence/sa_repos/ui_session_owner.py
Normal file
|
|
@ -0,0 +1,193 @@
|
|||
"""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"]
|
||||
Loading…
Add table
Add a link
Reference in a new issue