deerflow-code/offline-backend-20260512/backend/packages/harness/deerflow/persistence/skills/sql.py

"""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()}