umap-project 2.6.2__py3-none-any.whl → 2.7.0__py3-none-any.whl

umap/websocket_server.py

@@ -1,70 +1,173 @@
 #!/usr/bin/env python
 
 import asyncio
+import logging
+import uuid
 from collections import defaultdict
-from typing import Literal, Optional
+from typing import Literal, Optional, Union
 
 import websockets
 from django.conf import settings
 from django.core.signing import TimestampSigner
-from pydantic import BaseModel, ValidationError
+from pydantic import BaseModel, Field, RootModel, ValidationError
 from websockets import WebSocketClientProtocol
 from websockets.server import serve
 
-from umap.models import Map, User  # NOQA
+
+class Connections:
+    def __init__(self) -> None:
+        self._connections: set[WebSocketClientProtocol] = set()
+        self._ids: dict[WebSocketClientProtocol, str] = dict()
+
+    def join(self, websocket: WebSocketClientProtocol) -> str:
+        self._connections.add(websocket)
+        _id = str(uuid.uuid4())
+        self._ids[websocket] = _id
+        return _id
+
+    def leave(self, websocket: WebSocketClientProtocol) -> None:
+        self._connections.remove(websocket)
+        del self._ids[websocket]
+
+    def get(self, id) -> WebSocketClientProtocol:
+        # use an iterator to stop iterating as soon as we found
+        return next(k for k, v in self._ids.items() if v == id)
+
+    def get_id(self, websocket: WebSocketClientProtocol):
+        return self._ids[websocket]
+
+    def get_other_peers(
+        self, websocket: WebSocketClientProtocol
+    ) -> set[WebSocketClientProtocol]:
+        return self._connections - {websocket}
+
+    def get_all_peers(self) -> set[WebSocketClientProtocol]:
+        return self._connections
+
 
 # Contains the list of websocket connections handled by this process.
 # It's a mapping of map_id to a set of the active websocket connections
-CONNECTIONS = defaultdict(set)
+CONNECTIONS: defaultdict[int, Connections] = defaultdict(Connections)
 
 
-class JoinMessage(BaseModel):
-    kind: str = "join"
+class JoinRequest(BaseModel):
+    kind: Literal["JoinRequest"] = "JoinRequest"
     token: str
 
 
 class OperationMessage(BaseModel):
-    kind: str = "operation"
-    verb: str = Literal["upsert", "update", "delete"]
-    subject: str = Literal["map", "layer", "feature"]
+    """Message sent from one peer to all the others"""
+
+    kind: Literal["OperationMessage"] = "OperationMessage"
+    verb: Literal["upsert", "update", "delete"]
+    subject: Literal["map", "datalayer", "feature"]
     metadata: Optional[dict] = None
     key: Optional[str] = None
 
 
+class PeerMessage(BaseModel):
+    """Message sent from a specific peer to another one"""
+
+    kind: Literal["PeerMessage"] = "PeerMessage"
+    sender: str
+    recipient: str
+    # The message can be whatever the peers want. It's not checked by the server.
+    message: dict
+
+
+class ServerRequest(BaseModel):
+    """A request towards the server"""
+
+    kind: Literal["Server"] = "Server"
+    action: Literal["list-peers"]
+
+
+class Request(RootModel):
+    """Any message coming from the websocket should be one of these, and will be rejected otherwise."""
+
+    root: Union[ServerRequest, PeerMessage, OperationMessage] = Field(
+        discriminator="kind"
+    )
+
+
+class JoinResponse(BaseModel):
+    """Server response containing the list of peers"""
+
+    kind: Literal["JoinResponse"] = "JoinResponse"
+    peers: list
+    uuid: str
+
+
+class ListPeersResponse(BaseModel):
+    kind: Literal["ListPeersResponse"] = "ListPeersResponse"
+    peers: list
+
+
 async def join_and_listen(
     map_id: int, permissions: list, user: str | int, websocket: WebSocketClientProtocol
 ):
-    """Join a "room" whith other connected peers.
+    """Join a "room" with other connected peers, and wait for messages."""
+    logging.debug(f"{user} joined room #{map_id}")
+    connections: Connections = CONNECTIONS[map_id]
+    _id: str = connections.join(websocket)
+
+    # Assign an ID to the joining peer and return it the list of connected peers.
+    peers: list[WebSocketClientProtocol] = [
+        connections.get_id(p) for p in connections.get_all_peers()
+    ]
+    response = JoinResponse(uuid=_id, peers=peers)
+    await websocket.send(response.model_dump_json())
+
+    # Notify all other peers of the new list of connected peers.
+    message = ListPeersResponse(peers=peers)
+    websockets.broadcast(
+        connections.get_other_peers(websocket), message.model_dump_json()
+    )
 
-    New messages will be broadcasted to other connected peers.
-    """
-    print(f"{user} joined room #{map_id}")
-    CONNECTIONS[map_id].add(websocket)
     try:
         async for raw_message in websocket:
-            # recompute the peers-list at the time of message-sending.
+            # recompute the peers list at the time of message-sending.
             # as doing so beforehand would miss new connections
-            peers = CONNECTIONS[map_id] - {websocket}
-            # Only relay valid "operation" messages
+            other_peers = connections.get_other_peers(websocket)
             try:
-                OperationMessage.model_validate_json(raw_message)
-                websockets.broadcast(peers, raw_message)
+                incoming = Request.model_validate_json(raw_message)
             except ValidationError as e:
-                error = f"An error occurred when receiving this message: {raw_message}"
-                print(error, e)
+                error = f"An error occurred when receiving the following message: {raw_message!r}"
+                logging.error(error, e)
+            else:
+                match incoming.root:
+                    # Broadcast all operation messages to connected peers
+                    case OperationMessage():
+                        websockets.broadcast(other_peers, raw_message)
+
+                    # Send peer messages to the proper peer
+                    case PeerMessage(recipient=_id):
+                        peer = connections.get(_id)
+                        if peer:
+                            await peer.send(raw_message)
+
     finally:
-        CONNECTIONS[map_id].remove(websocket)
+        # On disconnect, remove the connection from the pool
+        connections.leave(websocket)
+
+        # TODO: refactor this in a separate method.
+        # Notify all other peers of the new list of connected peers.
+        peers = [connections.get_id(p) for p in connections.get_all_peers()]
+        message = ListPeersResponse(peers=peers)
+        websockets.broadcast(
+            connections.get_other_peers(websocket), message.model_dump_json()
+        )
 
 
-async def handler(websocket):
+async def handler(websocket: WebSocketClientProtocol):
     """Main WebSocket handler.
 
-    If permissions are granted, let the peer enter a room.
+    Check if the permission is granted and let the peer enter a room.
     """
     raw_message = await websocket.recv()
 
     # The first event should always be 'join'
-    message: JoinMessage = JoinMessage.model_validate_json(raw_message)
+    message: JoinRequest = JoinRequest.model_validate_json(raw_message)
     signed = TimestampSigner().unsign_object(message.token, max_age=30)
     user, map_id, permissions = signed.values()
 
@@ -73,7 +176,7 @@ async def handler(websocket):
         await join_and_listen(map_id, permissions, user, websocket)
 
 
-def run(host, port):
+def run(host: str, port: int):
     if not settings.WEBSOCKET_ENABLED:
         msg = (
             "WEBSOCKET_ENABLED should be set to True to run the WebSocket Server. "
@@ -86,7 +189,7 @@ def run(host, port):
 
     async def _serve():
         async with serve(handler, host, port):
-            print(f"Waiting for connections on {host}:{port}")
+            logging.debug(f"Waiting for connections on {host}:{port}")
             await asyncio.Future()  # run forever
 
     asyncio.run(_serve())

{umap_project-2.6.2.dist-info → umap_project-2.7.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: umap-project
-Version: 2.6.2
+Version: 2.7.0
 Summary: Create maps with OpenStreetMap layers in a minute and embed them in your site.
 Author-email: Yohan Boniface <yb@enix.org>
 Maintainer-email: David Larlet <david@larlet.fr>
@@ -19,35 +19,39 @@ Requires-Python: >=3.10
 Requires-Dist: django-agnocomplete==2.2.0
 Requires-Dist: django-environ==0.11.2
 Requires-Dist: django-probes==1.7.0
-Requires-Dist: django==5.1.1
-Requires-Dist: pillow==10.4.0
-Requires-Dist: psycopg==3.2.1
-Requires-Dist: pydantic==2.9.1
-Requires-Dist: rcssmin==1.1.2
+Requires-Dist: django==5.1.2
+Requires-Dist: pillow==11.0.0
+Requires-Dist: psycopg==3.2.3
+Requires-Dist: rcssmin==1.1.3
 Requires-Dist: requests==2.32.3
-Requires-Dist: rjsmin==1.2.2
+Requires-Dist: rjsmin==1.2.3
 Requires-Dist: social-auth-app-django==5.4.2
 Requires-Dist: social-auth-core==4.5.4
-Requires-Dist: websockets==13.0.1
 Provides-Extra: dev
 Requires-Dist: djlint==1.35.2; extra == 'dev'
-Requires-Dist: hatch==1.12.0; extra == 'dev'
+Requires-Dist: hatch==1.13.0; extra == 'dev'
 Requires-Dist: isort==5.13.2; extra == 'dev'
-Requires-Dist: mkdocs-material==9.5.34; extra == 'dev'
+Requires-Dist: mkdocs-material==9.5.42; extra == 'dev'
 Requires-Dist: mkdocs-static-i18n==1.2.3; extra == 'dev'
 Requires-Dist: mkdocs==1.6.1; extra == 'dev'
-Requires-Dist: pymdown-extensions==10.9; extra == 'dev'
-Requires-Dist: ruff==0.6.4; extra == 'dev'
+Requires-Dist: pymdown-extensions==10.11.2; extra == 'dev'
+Requires-Dist: ruff==0.7.0; extra == 'dev'
 Requires-Dist: vermin==1.6.0; extra == 'dev'
 Provides-Extra: docker
-Requires-Dist: uwsgi==2.0.26; extra == 'docker'
+Requires-Dist: uwsgi==2.0.27; extra == 'docker'
+Provides-Extra: sync
+Requires-Dist: channels==4.1.0; extra == 'sync'
+Requires-Dist: daphne==4.1.2; extra == 'sync'
+Requires-Dist: pydantic==2.9.2; extra == 'sync'
+Requires-Dist: websockets==13.1; extra == 'sync'
 Provides-Extra: test
 Requires-Dist: factory-boy==3.3.1; extra == 'test'
 Requires-Dist: playwright>=1.39; extra == 'test'
 Requires-Dist: pytest-django==4.9.0; extra == 'test'
 Requires-Dist: pytest-playwright==0.5.2; extra == 'test'
+Requires-Dist: pytest-rerunfailures==14.0; extra == 'test'
 Requires-Dist: pytest-xdist<4,>=3.5.0; extra == 'test'
-Requires-Dist: pytest==8.3.2; extra == 'test'
+Requires-Dist: pytest==8.3.3; extra == 'test'
 Description-Content-Type: text/markdown
 
 # uMap project