129 lines
4.8 KiB
Python
129 lines
4.8 KiB
Python
"""多渠道网关跨渠道身份关联。"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
|
||
from yuxi.repositories.channel_identity_repository import ChannelIdentityRepository
|
||
from yuxi.storage.postgres.manager import pg_manager
|
||
|
||
|
||
class IdentityLinkResolver:
|
||
def __init__(
|
||
self,
|
||
identity_links: dict[str, list[str]] | None = None,
|
||
identity_repository: ChannelIdentityRepository | None = None,
|
||
):
|
||
# { "canonical_name": ["feishu:ou_xxx", "dingtalk:12345", "telegram:67890"] }
|
||
self._links = identity_links or {}
|
||
self._reverse_index: dict[str, str] = {}
|
||
self._repository = identity_repository
|
||
self._lock = asyncio.Lock()
|
||
for canonical, ids in self._links.items():
|
||
for linked_id in ids:
|
||
self._reverse_index[linked_id] = canonical
|
||
|
||
@classmethod
|
||
async def create(cls, identity_repository: ChannelIdentityRepository | None = None) -> IdentityLinkResolver:
|
||
"""从数据库加载全量绑定关系并构建内存索引。
|
||
|
||
若未提供 repository，则使用 pg_manager 新建会话加载；
|
||
运行期 link/unlink 会复用初始化时传入的 repository，否则每次新建会话。
|
||
"""
|
||
if identity_repository is not None:
|
||
records = await identity_repository.load_all()
|
||
return cls._from_records(records, identity_repository)
|
||
|
||
async with pg_manager.get_async_session_context() as db:
|
||
repo = ChannelIdentityRepository(db)
|
||
records = await repo.load_all()
|
||
return cls._from_records(records, repo)
|
||
|
||
@classmethod
|
||
def _from_records(
|
||
cls,
|
||
records: list,
|
||
identity_repository: ChannelIdentityRepository | None,
|
||
) -> IdentityLinkResolver:
|
||
links: dict[str, list[str]] = {}
|
||
for record in records:
|
||
canonical = record.platform_user_id
|
||
linked_id = f"{record.channel_type}:{record.channel_sender_id}"
|
||
links.setdefault(canonical, []).append(linked_id)
|
||
return cls(identity_links=links, identity_repository=identity_repository)
|
||
|
||
def resolve(self, channel_type: str, peer_id: str) -> str:
|
||
key = f"{channel_type}:{peer_id}"
|
||
if key in self._reverse_index:
|
||
return self._reverse_index[key]
|
||
if peer_id in self._reverse_index:
|
||
return self._reverse_index[peer_id]
|
||
return peer_id
|
||
|
||
async def link(
|
||
self,
|
||
channel_type: str,
|
||
channel_sender_id: str,
|
||
platform_user_id: str,
|
||
account_id: str,
|
||
paired_by: str = "qr",
|
||
) -> None:
|
||
linked_id = f"{channel_type}:{channel_sender_id}"
|
||
|
||
async with self._lock:
|
||
if self._repository is not None:
|
||
await self._repository.create(
|
||
channel_type=channel_type,
|
||
account_id=account_id,
|
||
channel_sender_id=channel_sender_id,
|
||
platform_user_id=platform_user_id,
|
||
paired_by=paired_by,
|
||
)
|
||
else:
|
||
async with pg_manager.get_async_session_context() as db:
|
||
repo = ChannelIdentityRepository(db)
|
||
await repo.create(
|
||
channel_type=channel_type,
|
||
account_id=account_id,
|
||
channel_sender_id=channel_sender_id,
|
||
platform_user_id=platform_user_id,
|
||
paired_by=paired_by,
|
||
)
|
||
|
||
self._links.setdefault(platform_user_id, []).append(linked_id)
|
||
self._reverse_index[linked_id] = platform_user_id
|
||
|
||
async def unlink(
|
||
self,
|
||
channel_type: str,
|
||
account_id: str,
|
||
channel_sender_id: str,
|
||
) -> bool:
|
||
linked_id = f"{channel_type}:{channel_sender_id}"
|
||
|
||
async with self._lock:
|
||
if self._repository is not None:
|
||
deleted = await self._repository.delete_by_sender(
|
||
channel_type=channel_type,
|
||
account_id=account_id,
|
||
channel_sender_id=channel_sender_id,
|
||
)
|
||
else:
|
||
async with pg_manager.get_async_session_context() as db:
|
||
repo = ChannelIdentityRepository(db)
|
||
deleted = await repo.delete_by_sender(
|
||
channel_type=channel_type,
|
||
account_id=account_id,
|
||
channel_sender_id=channel_sender_id,
|
||
)
|
||
|
||
if deleted:
|
||
self._reverse_index.pop(linked_id, None)
|
||
for canonical, ids in list(self._links.items()):
|
||
if linked_id in ids:
|
||
ids.remove(linked_id)
|
||
if not ids:
|
||
self._links.pop(canonical, None)
|
||
break
|
||
return deleted
|