153 lines
5.1 KiB
Python
153 lines
5.1 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from collections.abc import Coroutine
|
|
|
|
from app.database.relationship import Relationship, RelationshipType
|
|
from app.dependencies.database import engine
|
|
from app.models.metadata_hub import MetadataClientState, OnlineStatus, UserActivity
|
|
|
|
from .hub import Client, Hub
|
|
|
|
from pydantic import TypeAdapter
|
|
from sqlmodel import select
|
|
from sqlmodel.ext.asyncio.session import AsyncSession
|
|
|
|
ONLINE_PRESENCE_WATCHERS_GROUP = "metadata:online-presence-watchers"
|
|
|
|
|
|
class MetadataHub(Hub):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self.state: dict[int, MetadataClientState] = {}
|
|
|
|
@staticmethod
|
|
def online_presence_watchers_group() -> str:
|
|
return ONLINE_PRESENCE_WATCHERS_GROUP
|
|
|
|
def broadcast_tasks(
|
|
self, user_id: int, store: MetadataClientState
|
|
) -> set[Coroutine]:
|
|
if not store.pushable:
|
|
return set()
|
|
return {
|
|
self.broadcast_group_call(
|
|
self.online_presence_watchers_group(),
|
|
"UserPresenceUpdated",
|
|
user_id,
|
|
store.to_dict(),
|
|
),
|
|
self.broadcast_group_call(
|
|
self.friend_presence_watchers_group(user_id),
|
|
"FriendPresenceUpdated",
|
|
user_id,
|
|
store.to_dict(),
|
|
),
|
|
}
|
|
|
|
@staticmethod
|
|
def friend_presence_watchers_group(user_id: int):
|
|
return f"metadata:friend-presence-watchers:{user_id}"
|
|
|
|
async def on_client_connect(self, client: Client) -> None:
|
|
user_id = int(client.connection_id)
|
|
if store := self.state.get(user_id):
|
|
store = MetadataClientState()
|
|
self.state[user_id] = store
|
|
|
|
async with AsyncSession(engine) as session:
|
|
async with session.begin():
|
|
friends = (
|
|
await session.exec(
|
|
select(Relationship.target_id).where(
|
|
Relationship.user_id == user_id,
|
|
Relationship.type == RelationshipType.FOLLOW,
|
|
)
|
|
)
|
|
).all()
|
|
tasks = []
|
|
for friend_id in friends:
|
|
self.groups.setdefault(
|
|
self.friend_presence_watchers_group(friend_id), set()
|
|
).add(client)
|
|
if (
|
|
friend_state := self.state.get(friend_id)
|
|
) and friend_state.pushable:
|
|
tasks.append(
|
|
self.broadcast_group_call(
|
|
self.friend_presence_watchers_group(friend_id),
|
|
"FriendPresenceUpdated",
|
|
friend_id,
|
|
friend_state.to_dict(),
|
|
)
|
|
)
|
|
await asyncio.gather(*tasks)
|
|
|
|
async def UpdateStatus(self, client: Client, status: int) -> None:
|
|
status_ = OnlineStatus(status)
|
|
user_id = int(client.connection_id)
|
|
store = self.state.get(user_id)
|
|
if store:
|
|
if store.status is not None and store.status == status_:
|
|
return
|
|
store.status = OnlineStatus(status_)
|
|
else:
|
|
store = MetadataClientState(status=OnlineStatus(status_))
|
|
self.state[user_id] = store
|
|
tasks = self.broadcast_tasks(user_id, store)
|
|
tasks.add(
|
|
self.call_noblock(
|
|
client,
|
|
"UserPresenceUpdated",
|
|
user_id,
|
|
store.to_dict(),
|
|
)
|
|
)
|
|
await asyncio.gather(*tasks)
|
|
|
|
async def UpdateActivity(self, client: Client, activity_dict: dict | None) -> None:
|
|
if activity_dict is None:
|
|
# idle
|
|
return
|
|
user_id = int(client.connection_id)
|
|
activity = TypeAdapter(UserActivity).validate_python(activity_dict)
|
|
store = self.state.get(user_id)
|
|
if store:
|
|
store.user_activity = activity
|
|
else:
|
|
store = MetadataClientState(
|
|
user_activity=activity,
|
|
)
|
|
self.state[user_id] = store
|
|
|
|
tasks = self.broadcast_tasks(user_id, store)
|
|
tasks.add(
|
|
self.call_noblock(
|
|
client,
|
|
"UserPresenceUpdated",
|
|
user_id,
|
|
store.to_dict(),
|
|
)
|
|
)
|
|
await asyncio.gather(*tasks)
|
|
|
|
async def BeginWatchingUserPresence(self, client: Client) -> None:
|
|
await asyncio.gather(
|
|
*[
|
|
self.call_noblock(
|
|
client,
|
|
"UserPresenceUpdated",
|
|
user_id,
|
|
store.to_dict(),
|
|
)
|
|
for user_id, store in self.state.items()
|
|
if store.pushable
|
|
]
|
|
)
|
|
self.groups.setdefault(self.online_presence_watchers_group(), set()).add(client)
|
|
|
|
async def EndWatchingUserPresence(self, client: Client) -> None:
|
|
self.groups.setdefault(self.online_presence_watchers_group(), set()).discard(
|
|
client
|
|
)
|