Subscriptions: stream missing items if user passed an old stoken.
This commit is contained in:
parent
f52facad1c
commit
61bd82f1e3
@ -2,6 +2,7 @@ import asyncio
|
||||
import typing as t
|
||||
|
||||
import aioredis
|
||||
from asgiref.sync import sync_to_async
|
||||
from django.db.models import QuerySet
|
||||
from fastapi import APIRouter, Depends, WebSocket, WebSocketDisconnect, status
|
||||
import nacl.encoding
|
||||
@ -11,6 +12,7 @@ from django_etebase import models
|
||||
from django_etebase.utils import CallbackContext, get_user_queryset
|
||||
from myauth.models import UserType, get_typed_user_model
|
||||
|
||||
from ..dependencies import get_collection_queryset, get_item_queryset
|
||||
from ..exceptions import NotSupported
|
||||
from ..msgpack import MsgpackRoute, msgpack_decode, msgpack_encode
|
||||
from ..redis import redisw
|
||||
@ -72,20 +74,49 @@ def get_websocket_user(websocket: WebSocket, ticket_model: t.Optional[TicketInne
|
||||
@websocket_router.websocket("/{ticket}/")
|
||||
async def websocket_endpoint(
|
||||
websocket: WebSocket,
|
||||
stoken: t.Optional[str],
|
||||
user: t.Optional[UserType] = Depends(get_websocket_user),
|
||||
ticket_model: TicketInner = Depends(load_websocket_ticket),
|
||||
):
|
||||
if user is None:
|
||||
return
|
||||
await websocket.accept()
|
||||
await redis_connector(websocket, ticket_model)
|
||||
await redis_connector(websocket, ticket_model, user, stoken)
|
||||
|
||||
|
||||
async def redis_connector(websocket: WebSocket, ticket_model: TicketInner):
|
||||
async def send_item_updates(
|
||||
websocket: WebSocket,
|
||||
collection: models.Collection,
|
||||
user: UserType,
|
||||
stoken: t.Optional[str],
|
||||
):
|
||||
from .collection import item_list_common
|
||||
|
||||
done = False
|
||||
while not done:
|
||||
queryset = await sync_to_async(get_item_queryset)(collection)
|
||||
response = await sync_to_async(item_list_common)(queryset, user, stoken, limit=50, prefetch="auto")
|
||||
done = response.done
|
||||
if len(response.data) > 0:
|
||||
await websocket.send_bytes(msgpack_encode(response.dict()))
|
||||
|
||||
|
||||
async def redis_connector(websocket: WebSocket, ticket_model: TicketInner, user: UserType, stoken: t.Optional[str]):
|
||||
async def producer_handler(r: aioredis.Redis, ws: WebSocket):
|
||||
channel_name = f"col.{ticket_model.req.collection}"
|
||||
(channel,) = await r.psubscribe(channel_name)
|
||||
assert isinstance(channel, aioredis.Channel)
|
||||
|
||||
# Send missing items if we are not up to date
|
||||
queryset: QuerySet[models.Collection] = get_collection_queryset(user)
|
||||
collection: t.Optional[models.Collection] = await sync_to_async(
|
||||
queryset.filter(uid=ticket_model.req.collection).first
|
||||
)()
|
||||
if collection is None:
|
||||
await websocket.close(code=status.WS_1008_POLICY_VIOLATION)
|
||||
return
|
||||
await send_item_updates(websocket, collection, user, stoken)
|
||||
|
||||
try:
|
||||
while True:
|
||||
# We wait on the websocket so we fail if web sockets fail or get data
|
||||
|
Loading…
Reference in New Issue
Block a user