ZY-Agent/backend/packages/harness/deerflow/persistence/workspace_membership/sql.py

"""SQLAlchemy-backed workspace membership repository.

Manages the ``workspace_memberships`` join table. Stage 0 only writes
``role='owner'`` (single-user workspaces); the repository accepts the
full Stage 2 RBAC role enum so the schema is forward-compatible.

Owner transfer is intentionally NOT modelled here as a single method —
it requires a two-row transactional swap with careful retry semantics,
and belongs in the auth router (PR4+) where it can be wrapped in a
permission check.
"""

from __future__ import annotations

from datetime import UTC, datetime
from typing import Any

from sqlalchemy import delete, select, update
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker

from deerflow.persistence.workspace_membership.model import WorkspaceMembershipRow

# 允许的 role 字面值。Stage 0 实际仅写 owner；admin/member 留给 Stage 2 RBAC。
_VALID_ROLES = frozenset({"owner", "admin", "member"})


class MembershipValidationError(ValueError):
    """Raised when role is not in the allowed enum."""


def _validate_role(role: str) -> None:
    if role not in _VALID_ROLES:
        raise MembershipValidationError(f"role {role!r} is not in allowed set {_VALID_ROLES!r}")


class WorkspaceMembershipRepository:
    def __init__(self, session_factory: async_sessionmaker[AsyncSession]) -> None:
        self._sf = session_factory

    @staticmethod
    def _row_to_dict(row: WorkspaceMembershipRow) -> dict[str, Any]:
        return {
            "workspace_id": row.workspace_id,
            "user_id": row.user_id,
            "role": row.role,
            "invited_by": row.invited_by,
            "joined_at": row.joined_at.isoformat() if row.joined_at else None,
        }

    async def add(
        self,
        *,
        workspace_id: str,
        user_id: str,
        role: str,
        invited_by: str | None = None,
    ) -> dict[str, Any]:
        """Insert a new membership row.

        Raises:
          - :class:`MembershipValidationError` for unknown role values
          - :class:`sqlalchemy.exc.IntegrityError` for:
            - duplicate (workspace_id, user_id) — composite PK collision
            - second ``owner`` in the same workspace — partial unique index
            - invalid workspace_id / user_id / invited_by FK
        """
        _validate_role(role)
        row = WorkspaceMembershipRow(
            workspace_id=workspace_id,
            user_id=user_id,
            role=role,
            invited_by=invited_by,
            joined_at=datetime.now(UTC),
        )
        async with self._sf() as session:
            session.add(row)
            await session.commit()
            await session.refresh(row)
            return self._row_to_dict(row)

    async def remove(self, *, workspace_id: str, user_id: str) -> bool:
        """Delete a membership row; returns True if a row was deleted.

        Owner removal is allowed at the repository layer — auth router
        (PR4) layers the "cannot remove last owner" rule on top.
        """
        async with self._sf() as session:
            result = await session.execute(
                delete(WorkspaceMembershipRow).where(
                    WorkspaceMembershipRow.workspace_id == workspace_id,
                    WorkspaceMembershipRow.user_id == user_id,
                )
            )
            await session.commit()
            return (result.rowcount or 0) > 0

    async def list_by_user(self, *, user_id: str) -> list[dict[str, Any]]:
        """Return all memberships for ``user_id``, ordered by joined_at desc.

        No contextvar resolution here — caller is responsible for passing
        the correct user_id. Used by ``/auth/me`` to list workspaces a
        user belongs to.
        """
        async with self._sf() as session:
            result = await session.execute(select(WorkspaceMembershipRow).where(WorkspaceMembershipRow.user_id == user_id).order_by(WorkspaceMembershipRow.joined_at.desc()))
            return [self._row_to_dict(r) for r in result.scalars()]

    async def list_by_workspace(self, *, workspace_id: str) -> list[dict[str, Any]]:
        """Return all members of a workspace, ordered by joined_at asc."""
        async with self._sf() as session:
            result = await session.execute(select(WorkspaceMembershipRow).where(WorkspaceMembershipRow.workspace_id == workspace_id).order_by(WorkspaceMembershipRow.joined_at.asc()))
            return [self._row_to_dict(r) for r in result.scalars()]

    async def get_role(self, *, workspace_id: str, user_id: str) -> str | None:
        """Return the role string, or None if user is not a member."""
        async with self._sf() as session:
            result = await session.execute(
                select(WorkspaceMembershipRow.role).where(
                    WorkspaceMembershipRow.workspace_id == workspace_id,
                    WorkspaceMembershipRow.user_id == user_id,
                )
            )
            return result.scalar_one_or_none()

    async def change_role(
        self,
        *,
        workspace_id: str,
        user_id: str,
        new_role: str,
    ) -> bool:
        """Update a member's role; returns True iff a row was updated.

        Validates ``new_role`` against the allowed enum. Owner-transfer
        flow needs to swap two rows atomically — do that with a manual
        transaction in the caller; this method is for non-owner changes.
        """
        _validate_role(new_role)
        async with self._sf() as session:
            result = await session.execute(
                update(WorkspaceMembershipRow)
                .where(
                    WorkspaceMembershipRow.workspace_id == workspace_id,
                    WorkspaceMembershipRow.user_id == user_id,
                )
                .values(role=new_role)
            )
            await session.commit()
            return (result.rowcount or 0) > 0