"""SQLAlchemy-backed custom skill repository.""" from __future__ import annotations from datetime import UTC, datetime from typing import Any from sqlalchemy import func, or_, select from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker from deerflow.persistence.skills.base import SkillStore from deerflow.persistence.skills.model import SkillFavoriteRow, SkillRow def _row_to_dict(row: SkillRow) -> dict[str, Any]: data = row.to_dict() data["builtin"] = data.get("owner_user_id") is None for key in ("created_at", "updated_at"): value = data.get(key) if isinstance(value, datetime): data[key] = value.isoformat() return data class SkillRepository(SkillStore): def __init__(self, session_factory: async_sessionmaker[AsyncSession]) -> None: self._sf = session_factory @staticmethod def _visible_filter(user_id: str): return or_( SkillRow.owner_user_id.is_(None), SkillRow.owner_user_id == user_id, SkillRow.published.is_(True), ) async def list_visible(self, user_id: str) -> list[dict[str, Any]]: stmt = select(SkillRow).where(self._visible_filter(user_id)).order_by(SkillRow.created_at.asc()) async with self._sf() as session: result = await session.execute(stmt) return [_row_to_dict(row) for row in result.scalars()] async def get_visible(self, name: str, user_id: str) -> dict[str, Any] | None: stmt = select(SkillRow).where(SkillRow.name == name, self._visible_filter(user_id)) async with self._sf() as session: row = (await session.execute(stmt)).scalar_one_or_none() return _row_to_dict(row) if row is not None else None async def get_any(self, name: str) -> dict[str, Any] | None: async with self._sf() as session: row = await session.get(SkillRow, name) return _row_to_dict(row) if row is not None else None async def get_owned(self, name: str, user_id: str) -> dict[str, Any] | None: stmt = select(SkillRow).where(SkillRow.name == name, SkillRow.owner_user_id == user_id) async with self._sf() as session: row = (await session.execute(stmt)).scalar_one_or_none() return _row_to_dict(row) if row is not None else None async def create(self, data: dict[str, Any]) -> dict[str, Any]: now = datetime.now(UTC) row = SkillRow( name=data["name"], owner_user_id=data.get("owner_user_id"), published=bool(data.get("published", False)), detail=data.get("detail"), name_zh=data.get("name_zh"), always_on=bool(data.get("always_on", False)), square_id=str(data.get("square_id") or ""), created_at=now, updated_at=now, ) async with self._sf() as session: session.add(row) await session.commit() await session.refresh(row) return _row_to_dict(row) async def ensure_legacy(self, data: dict[str, Any]) -> dict[str, Any]: async with self._sf() as session: row = await session.get(SkillRow, data["name"]) if row is None: now = datetime.now(UTC) row = SkillRow( name=data["name"], owner_user_id=data.get("owner_user_id"), published=bool(data.get("published", True)), detail=data.get("detail"), name_zh=data.get("name_zh"), created_at=now, updated_at=now, ) session.add(row) elif row.owner_user_id is None and data.get("owner_user_id") is not None: # Only patch ownership if it was previously null — never overwrite # an existing owner record on a re-migration. row.owner_user_id = data["owner_user_id"] row.updated_at = datetime.now(UTC) await session.commit() await session.refresh(row) return _row_to_dict(row) async def update(self, name: str, user_id: str, data: dict[str, Any]) -> dict[str, Any] | None: async with self._sf() as session: row = await session.get(SkillRow, name) if row is None or row.owner_user_id != user_id: return None if "published" in data: row.published = bool(data["published"]) if "detail" in data: row.detail = data["detail"] if "name_zh" in data: row.name_zh = data["name_zh"] if "always_on" in data: row.always_on = bool(data["always_on"]) if "square_id" in data: row.square_id = str(data["square_id"] or "") row.updated_at = datetime.now(UTC) await session.commit() await session.refresh(row) return _row_to_dict(row) async def update_admin(self, name: str, data: dict[str, Any]) -> dict[str, Any] | None: async with self._sf() as session: row = await session.get(SkillRow, name) if row is None: return None if "published" in data: row.published = bool(data["published"]) if "owner_user_id" in data: row.owner_user_id = data["owner_user_id"] if "detail" in data: row.detail = data["detail"] if "name_zh" in data: row.name_zh = data["name_zh"] if "always_on" in data: row.always_on = bool(data["always_on"]) if "square_id" in data: row.square_id = str(data["square_id"] or "") row.updated_at = datetime.now(UTC) await session.commit() await session.refresh(row) return _row_to_dict(row) async def delete(self, name: str, user_id: str) -> bool: async with self._sf() as session: row = await session.get(SkillRow, name) if row is None or row.owner_user_id != user_id: return False await session.delete(row) await session.commit() return True async def delete_admin(self, name: str) -> bool: async with self._sf() as session: row = await session.get(SkillRow, name) if row is None: return False await session.delete(row) await session.commit() return True async def count_owned(self, user_id: str) -> int: stmt = select(func.count()).select_from(SkillRow).where(SkillRow.owner_user_id == user_id) async with self._sf() as session: return await session.scalar(stmt) or 0 async def list_all(self) -> list[dict[str, Any]]: stmt = select(SkillRow).order_by(SkillRow.created_at.asc()) async with self._sf() as session: result = await session.execute(stmt) return [_row_to_dict(row) for row in result.scalars()] async def list_published_names(self) -> set[str]: stmt = select(SkillRow.name).where(SkillRow.published.is_(True)) async with self._sf() as session: result = await session.execute(stmt) return {row[0] for row in result.all()} async def list_always_on_names(self) -> set[str]: stmt = select(SkillRow.name).where(SkillRow.always_on.is_(True)) async with self._sf() as session: result = await session.execute(stmt) return {row[0] for row in result.all()} async def list_favorite_skill_names(self, user_id: str) -> set[str]: stmt = select(SkillFavoriteRow.skill_name).where(SkillFavoriteRow.user_id == user_id) async with self._sf() as session: result = await session.execute(stmt) return {row[0] for row in result.all()}