1
# Copyright (c) 2023 Cloudflare, Inc.
2
# Licensed under the Apache 2.0 license found in the LICENSE file or at:
3
#     https://opensource.org/licenses/Apache-2.0
4
from js import WebSocketPair
5
from workers import DurableObject, Request, Response, WorkerEntrypoint
6
 
7
 
8
class DurableObjectWebSocket(DurableObject):
9
    def __init__(self, state, env):
10
        super().__init__(state, env)
11
 
12
    async def fetch(self, request):
13
        assert isinstance(request, Request)
14
        web_socket_pair = WebSocketPair.new()
15
        client, server = web_socket_pair
16
        self.ctx.acceptWebSocket(server)
17
        return Response(None, status=101, web_socket=client)
18
 
19
    async def webSocketMessage(self, ws, message):
20
        print("Received in Python DO WS message: ", message)
21
        ws.send("hello")
22
 
23
    async def webSocketClose(self, ws, code, reason, wasClean):
24
        print("Closed in Python DO WS")
25
 
26
 
27
class Default(WorkerEntrypoint):
28
    async def fetch(self, request):
29
        if request.method == "GET" and request.url.endswith("/websocket"):
30
            upgrade_header = request.headers.get("Upgrade")
31
            if upgrade_header != "websocket":
32
                return Response(
33
                    None,
34
                    status=426,
35
                    status_text="Expected Upgrade",
36
                    headers={"Content-Type": "text/plain"},
37
                )
38
 
39
            # We are explicitly testing usage via `self.env` here rather than the
40
            # argument to `fetch` and using `getByName` rather than `get`.
41
            stub = self.env.ns.getByName("A")
42
            return await stub.fetch(request)
43
        return Response("Not found", status=404)