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

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

from packages.domain.entities import WorkspaceInvitation
from packages.ports.workspace_invitation_repository import WorkspaceInvitationRepository


class PostgresWorkspaceInvitationRepository(WorkspaceInvitationRepository):
    """WorkspaceInvitation 仓储 PostgreSQL 实现"""
    
    def __init__(self, connection_string: str):
        self.connection_string = connection_string
    
    def _get_connection(self):
        """获取数据库连接（使用连接池）"""
        from packages.adapters.postgres.connection_pool import PooledConnection
        return PooledConnection()
    
    def save(self, invitation: WorkspaceInvitation) -> None:
        """保存邀请"""
        conn = self._get_connection()
        try:
            with conn.cursor() as cur:
                cur.execute("""
                    INSERT INTO workspace_invitations (
                        id, workspace_id, email, role, token,
                        invited_by, expires_at, status, created_at
                    ) VALUES (
                        %(id)s, %(workspace_id)s, %(email)s, %(role)s, %(token)s,
                        %(invited_by)s, %(expires_at)s, %(status)s, %(created_at)s
                    )
                    ON CONFLICT (id) DO UPDATE SET
                        status = EXCLUDED.status
                """, {
                    "id": invitation.id,
                    "workspace_id": invitation.workspace_id,
                    "email": invitation.email,
                    "role": invitation.role,
                    "token": invitation.token,
                    "invited_by": invitation.invited_by,
                    "expires_at": invitation.expires_at,
                    "status": invitation.status,
                    "created_at": invitation.created_at,
                })
            conn.commit()
        finally:
            conn.close()
    
    def find_by_id(self, invitation_id: str) -> Optional[WorkspaceInvitation]:
        """根据 ID 查找邀请"""
        conn = self._get_connection()
        try:
            with conn.cursor() as cur:
                cur.execute("SELECT * FROM workspace_invitations WHERE id = %s", (invitation_id,))
                row = cur.fetchone()
                return self._row_to_invitation(row) if row else None
        finally:
            conn.close()
    
    def find_by_token(self, token: str) -> Optional[WorkspaceInvitation]:
        """根据 token 查找邀请"""
        conn = self._get_connection()
        try:
            with conn.cursor() as cur:
                cur.execute("SELECT * FROM workspace_invitations WHERE token = %s", (token,))
                row = cur.fetchone()
                return self._row_to_invitation(row) if row else None
        finally:
            conn.close()
    
    def find_by_email(self, email: str) -> List[WorkspaceInvitation]:
        """根据邮箱查找所有邀请"""
        conn = self._get_connection()
        try:
            with conn.cursor() as cur:
                cur.execute(
                    "SELECT * FROM workspace_invitations WHERE email = %s ORDER BY created_at DESC",
                    (email,)
                )
                rows = cur.fetchall()
                return [self._row_to_invitation(row) for row in rows]
        finally:
            conn.close()
    
    def find_pending_by_email(self, email: str) -> List[WorkspaceInvitation]:
        """查找邮箱的待处理邀请"""
        conn = self._get_connection()
        try:
            with conn.cursor() as cur:
                cur.execute("""
                    SELECT * FROM workspace_invitations 
                    WHERE email = %s AND status = 'pending' AND expires_at > NOW()
                    ORDER BY created_at DESC
                """, (email,))
                rows = cur.fetchall()
                return [self._row_to_invitation(row) for row in rows]
        finally:
            conn.close()
    
    def delete(self, invitation_id: str) -> bool:
        """删除邀请"""
        conn = self._get_connection()
        try:
            with conn.cursor() as cur:
                cur.execute("DELETE FROM workspace_invitations WHERE id = %s", (invitation_id,))
                deleted = cur.rowcount > 0
            conn.commit()
            return deleted
        finally:
            conn.close()
    
    def _row_to_invitation(self, row: dict) -> WorkspaceInvitation:
        """将数据库行转换为 WorkspaceInvitation 对象"""
        return WorkspaceInvitation(
            id=row["id"],
            workspace_id=row["workspace_id"],
            email=row["email"],
            role=row["role"],
            token=row["token"],
            invited_by=row["invited_by"],
            expires_at=row["expires_at"],
            status=row["status"],
            created_at=row["created_at"],
        )