xiaoxia-saas/packages/adapters/postgres/workspace_member_repository.py

"""
PostgreSQL WorkspaceMember Repository 实现
"""
from typing import Optional, List
import psycopg2
from psycopg2.extras import RealDictCursor

from packages.domain.entities import WorkspaceMember
from packages.ports.workspace_member_repository import WorkspaceMemberRepository


class PostgresWorkspaceMemberRepository(WorkspaceMemberRepository):
    """WorkspaceMember 仓储 PostgreSQL 实现"""
    
    def __init__(self, connection_string: str):
        self.connection_string = connection_string
    
    def _get_connection(self):
        return psycopg2.connect(self.connection_string, cursor_factory=RealDictCursor)
    
    def save(self, member: WorkspaceMember) -> None:
        """保存成员"""
        conn = self._get_connection()
        try:
            with conn.cursor() as cur:
                cur.execute("""
                    INSERT INTO workspace_members (
                        id, workspace_id, user_id, role, invited_by, joined_at
                    ) VALUES (
                        %(id)s, %(workspace_id)s, %(user_id)s, %(role)s,
                        %(invited_by)s, %(joined_at)s
                    )
                    ON CONFLICT (workspace_id, user_id) DO UPDATE SET
                        role = EXCLUDED.role
                """, {
                    "id": member.id,
                    "workspace_id": member.workspace_id,
                    "user_id": member.user_id,
                    "role": member.role,
                    "invited_by": member.invited_by,
                    "joined_at": member.joined_at,
                })
            conn.commit()
        finally:
            conn.close()
    
    def find_by_id(self, member_id: str) -> Optional[WorkspaceMember]:
        """根据 ID 查找成员"""
        conn = self._get_connection()
        try:
            with conn.cursor() as cur:
                cur.execute("SELECT * FROM workspace_members WHERE id = %s", (member_id,))
                row = cur.fetchone()
                return self._row_to_member(row) if row else None
        finally:
            conn.close()
    
    def find_by_workspace_and_user(
        self,
        workspace_id: str,
        user_id: str,
    ) -> Optional[WorkspaceMember]:
        """根据 workspace 和 user 查找成员"""
        conn = self._get_connection()
        try:
            with conn.cursor() as cur:
                cur.execute(
                    "SELECT * FROM workspace_members WHERE workspace_id = %s AND user_id = %s",
                    (workspace_id, user_id)
                )
                row = cur.fetchone()
                return self._row_to_member(row) if row else None
        finally:
            conn.close()
    
    def find_by_user(self, user_id: str) -> List[WorkspaceMember]:
        """查找用户的所有成员记录"""
        conn = self._get_connection()
        try:
            with conn.cursor() as cur:
                cur.execute(
                    "SELECT * FROM workspace_members WHERE user_id = %s ORDER BY joined_at DESC",
                    (user_id,)
                )
                rows = cur.fetchall()
                return [self._row_to_member(row) for row in rows]
        finally:
            conn.close()
    
    def find_by_workspace(self, workspace_id: str) -> List[WorkspaceMember]:
        """查找 workspace 的所有成员"""
        conn = self._get_connection()
        try:
            with conn.cursor() as cur:
                cur.execute(
                    "SELECT * FROM workspace_members WHERE workspace_id = %s ORDER BY joined_at",
                    (workspace_id,)
                )
                rows = cur.fetchall()
                return [self._row_to_member(row) for row in rows]
        finally:
            conn.close()
    
    def count_by_workspace(self, workspace_id: str) -> int:
        """统计 workspace 的成员数量"""
        conn = self._get_connection()
        try:
            with conn.cursor() as cur:
                cur.execute(
                    "SELECT COUNT(*) FROM workspace_members WHERE workspace_id = %s",
                    (workspace_id,)
                )
                return cur.fetchone()["count"]
        finally:
            conn.close()
    
    def delete(self, member_id: str) -> bool:
        """删除成员"""
        conn = self._get_connection()
        try:
            with conn.cursor() as cur:
                cur.execute("DELETE FROM workspace_members WHERE id = %s", (member_id,))
                deleted = cur.rowcount > 0
            conn.commit()
            return deleted
        finally:
            conn.close()
    
    def _row_to_member(self, row: dict) -> WorkspaceMember:
        """将数据库行转换为 WorkspaceMember 对象"""
        return WorkspaceMember(
            id=row["id"],
            workspace_id=row["workspace_id"],
            user_id=row["user_id"],
            role=row["role"],
            invited_by=row["invited_by"],
            joined_at=row["joined_at"],
        )