async def _dispatch_request(self, req, **options):
# Set formats on Request object.
req.formats = self.formats
# Get the route.
route = self.path_matches_route(req.url.path)
route = self.routes.get(route)
if route:
if route.uses_websocket:
resp = WebSocket(**options)
else:
resp = models.Response(req=req, formats=self.formats)
for before_request in self.before_requests:
await self._execute_route(route=before_request, req=req, resp=resp)
await self._execute_route(route=route, req=req, resp=resp, **options)
else:
resp = models.Response(req=req, formats=self.formats)
self.default_response(req, resp, notfound=True)
self.default_response(req, resp)
self._prepare_session(resp)
self._prepare_cookies(resp)
return resp
async def ws_with_auth(websocket):
websocket = WebSocket(scope=websocket.scope,
receive=websocket.receive,
send=websocket.send)
await websocket.accept()
await websocket.send_text('Authentication valid')
await websocket.close()
async def awaitable(receive: Receive, send: Send) -> None:
session = WebSocket(scope, receive=receive, send=send)
kwargs = scope.get("kwargs", {})
injected_func = await injector.inject(func)
await injected_func(session, **kwargs)
def test_websocket_scope_interface():
"""
A WebSocket can be instantiated with a scope, and presents a `Mapping`
interface.
"""
async def mock_receive():
pass # pragma: no cover
async def mock_send(message):
pass # pragma: no cover
websocket = WebSocket(
{
"type": "websocket",
"path": "/abc/",
"headers": []
},
receive=mock_receive,
send=mock_send,
)
assert websocket["type"] == "websocket"
assert dict(websocket) == {
"type": "websocket",
"path": "/abc/",
"headers": []
}
assert len(websocket) == 3
async def asgi(self, receive: Receive, send: Send, scope: Scope) -> None:
assert scope["type"] == "websocket"
websocket = WebSocket(scope, receive=receive, send=send)
await websocket.accept(subprotocol="graphql-ws")
await self._send_message(websocket, "connection_ack")
# TODO: we should check that this is a proper connection init message
await websocket.receive_json()
data = await websocket.receive_json()
id_ = data.get("id", "1")
payload = data.get("payload", {})
data = await self.execute(
payload["query"],
payload["variables"],
operation_name=payload["operationName"],
)
async for result in data:
# TODO: send errors if any
await self._send_message(websocket, "data", {"data": result.data},
id_)
await self._send_message(websocket, "complete")
await websocket.close()
async def app(scope: Scope, receive: Receive, send: Send) -> None:
websocket = WebSocket(scope, receive=receive, send=send)
await websocket.accept()
async with anyio.create_task_group() as task_group:
task_group.start_soon(reader, websocket)
await writer(websocket)
await websocket.close()
def __init__(self, scope: Scope, receive: Receive, send: Send):
"""Do not use manually."""
self._connection = WebSocket(scope, receive, send)
self.state = addict.Dict()
self.reraise = False
self.closed = False
self.headers = self._connection.headers
self._queries = None
async def asgi(receive, send):
nonlocal close_code
websocket = WebSocket(scope, receive=receive, send=send)
await websocket.accept()
try:
await websocket.receive_text()
except WebSocketDisconnect as exc:
close_code = exc.code
async def __call__(self, scope, receive, send):
ws = WebSocket(scope, receive, send)
before_requests = scope.get("before_requests", [])
for before_request in before_requests.get("ws", []):
await before_request(ws)
await self.endpoint(ws)
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
assert scope["type"] == "websocket"
ws = WebSocket(scope, receive, send)
await ws.accept()
await run_until_first_complete(
(self._watch_reloads, {"ws": ws}),
(self._wait_client_disconnect, {"ws": ws}),
)
async def app(scope: Scope, receive: Receive, send: Send) -> None:
nonlocal close_code
websocket = WebSocket(scope, receive=receive, send=send)
await websocket.accept()
try:
await websocket.receive_text()
except WebSocketDisconnect as exc:
close_code = exc.code
def test_websocket_scope_interface():
"""
A WebSocket can be instantiated with a scope, and presents a `Mapping`
interface.
"""
websocket = WebSocket({"type": "websocket", "path": "/abc/", "headers": []})
assert websocket["type"] == "websocket"
assert dict(websocket) == {"type": "websocket", "path": "/abc/", "headers": []}
assert len(websocket) == 3
async def get_thing(self, thing_id):
"""
Get the thing this request is for.
things -- list of Things managed by this server
thing_id -- ID of the thing to get, in string form
Returns the thing, or None if not found.
"""
websocket = WebSocket(self.scope, receive=self.receive, send=self.send)
things = websocket.app.state.things
return await things.get_thing(thing_id)
async def websocket(self, scope: Scope, receive: Receive,
send: Send) -> None:
websocket = WebSocket(scope, receive=receive, send=send)
module = self.get_view(websocket.url.path)
if module is None or not hasattr(module, "Socket"):
await WebSocketClose(WS_1001_GOING_AWAY)(scope, receive, send)
return
await getattr(module, "Socket")(websocket)
async def asgi(receive, send):
websocket = WebSocket(scope, receive=receive, send=send)
await websocket.accept()
asyncio.ensure_future(respond(websocket))
try:
# this will block as the client does not send us data
# it should not prevent `respond` from executing though
await websocket.receive_json()
except WebSocketDisconnect:
pass
async def handle_websocket(self, scope: Scope, receive: Receive, send: Send):
websocket = WebSocket(scope=scope, receive=receive, send=send)
subscriptions: typing.Dict[str, typing.AsyncGenerator] = {}
tasks = {}
await websocket.accept(subprotocol="graphql-ws")
try:
while (
websocket.client_state != WebSocketState.DISCONNECTED
and websocket.application_state != WebSocketState.DISCONNECTED
):
message = await websocket.receive_json()
operation_id = message.get("id")
message_type = message.get("type")
if message_type == GQL_CONNECTION_INIT:
await websocket.send_json({"type": GQL_CONNECTION_ACK})
if self.keep_alive:
self._keep_alive_task = asyncio.create_task(
self.handle_keep_alive(websocket)
)
elif message_type == GQL_CONNECTION_TERMINATE:
await websocket.close()
elif message_type == GQL_START:
async_result = await self.start_subscription(
message.get("payload"), operation_id, websocket
)
subscriptions[operation_id] = async_result
tasks[operation_id] = asyncio.create_task(
self.handle_async_results(async_result, operation_id, websocket)
)
elif message_type == GQL_STOP: # pragma: no cover
if operation_id not in subscriptions:
return
await subscriptions[operation_id].aclose()
tasks[operation_id].cancel()
del tasks[operation_id]
del subscriptions[operation_id]
except WebSocketDisconnect: # pragma: no cover
pass
finally:
if self._keep_alive_task:
self._keep_alive_task.cancel()
for operation_id in subscriptions:
await subscriptions[operation_id].aclose()
tasks[operation_id].cancel()
async def asgi(receive, send):
websocket = WebSocket(scope, receive=receive, send=send)
await websocket.accept()
async with anyio.create_task_group() as task_group:
task_group.start_soon(respond, websocket)
try:
# this will block as the client does not send us data
# it should not prevent `respond` from executing though
await websocket.receive_json()
except WebSocketDisconnect:
pass
async def _dispatch_ws(self, scope, receive, send):
ws = WebSocket(scope=scope, receive=receive, send=send)
route = self.path_matches_route(ws.url.path)
route = self.routes.get(route)
if route:
for before_request in self.before_ws_requests:
await self.background(before_request, ws=ws)
await self.background(route.endpoint, ws)
else:
await send({"type": "websocket.close", "code": 1000})
async def __call__(self, receive, send):
redis_host = 'redis://{}'.format(
self.scope.get('app').settings.REDIS_HOST)
self.pub = await aioredis.create_redis(redis_host)
self.sub = await aioredis.create_redis(redis_host)
websocket = WebSocket(self.scope, receive=receive, send=send)
await self.on_connect(websocket)
await asyncio.gather(
self.listen_ws(websocket),
self.listen_redis(websocket, [
self.get_channel_name(constants.__ALL__),
self.get_channel_name(self.channel_name)
]))
async def asgi(receive, send):
nonlocal scope, self
if scope["type"] == "websocket":
ws = WebSocket(scope=scope, receive=receive, send=send)
await self._dispatch_ws(ws)
else:
req = models.Request(scope, receive=receive, api=self)
resp = await self._dispatch_request(req,
scope=scope,
send=send,
receive=receive)
await resp(receive, send)
async def asgi(receive, send):
websocket = WebSocket(scope, receive=receive, send=send)
queue = asyncio.Queue()
await websocket.accept()
await run_until_first_complete(
(reader, {
"websocket": websocket,
"queue": queue
}),
(writer, {
"websocket": websocket,
"queue": queue
}),
)
await websocket.close()
async def dispatch(self) -> None:
app = self.scope["app"]
websocket = WebSocket(self.scope, self.receive, self.send)
route, route_scope = app.router.get_route_from_scope(self.scope)
state = {
"scope": self.scope,
"receive": self.receive,
"send": self.send,
"exc": None,
"app": app,
"path_params": route_scope["path_params"],
"route": route,
"websocket": websocket,
"websocket_encoding": self.encoding,
"websocket_code": status.WS_1000_NORMAL_CLOSURE,
"websocket_message": None,
}
try:
on_connect = await app.injector.inject(self.on_connect, state)
await on_connect()
except Exception as e:
raise exceptions.WebSocketConnectionException(
"Error connecting socket") from e
try:
state["websocket_message"] = await websocket.receive()
while websocket.client_state == WebSocketState.CONNECTED:
on_receive = await app.injector.inject(self.on_receive, state)
await on_receive()
state["websocket_message"] = await websocket.receive()
state["websocket_code"] = int(state["websocket_message"].get(
"code", status.WS_1000_NORMAL_CLOSURE))
except exceptions.WebSocketException as e:
state["websocket_code"] = e.close_code
except Exception as e:
state["websocket_code"] = status.WS_1011_INTERNAL_ERROR
raise e from None
finally:
on_disconnect = await app.injector.inject(self.on_disconnect,
state)
await on_disconnect()
async def __call__(self, scope: Scope, receive: Receive,
send: Send) -> None:
if scope["type"] == "http":
request = Request(scope=scope, receive=receive)
response: Response
if request.method == "GET" and self.playground:
response = HTMLResponse(PLAYGROUND_HTML)
elif request.method == "POST":
response = await self._handle_http_request(request)
else:
response = Response(status_code=405)
await response(scope, receive, send)
elif scope["type"] == "websocket":
websocket = WebSocket(scope=scope, receive=receive, send=send)
await self._run_websocket_server(websocket)
else:
raise ValueError(f"Unsupported scope type: ${scope['type']}")
async def __call__(self, receive: Receive, send: Send) -> None:
websocket = WebSocket(self.scope, receive=receive, send=send)
kwargs = self.scope.get("kwargs", {})
await self.on_connect(websocket, **kwargs)
close_code = None
try:
while True:
message = await websocket.receive()
if message["type"] == "websocket.receive":
data = await self.decode(websocket, message)
await self.on_receive(websocket, data)
elif message["type"] == "websocket.disconnect":
close_code = message.get("code", 1000)
return
finally:
await self.on_disconnect(websocket, close_code)
async def __call__(self, receive: Receive, send: Send) -> None:
self.websocket = WebSocket(self.scope, receive=receive, send=send)
# overridable/Callable by subclasses
await self.on_connect(self.websocket)
# save send method
self.send_base = send
# Initialize channel layer
self.channel_layer = get_channel_layer(self.channel_layer_alias)
if self.channel_layer is not None:
# create new channel
self.channel_name = await self.channel_layer.new_channel()
"""
add the channel to group (broadcat), this it can be imprroved, by taking the group
name from URLs, See django_channels for example.
"""
self.room_name = 'asgi_room'
self.room_group_name = 'chat_%s' % self.room_name
# Join room group
await self.channel_layer.group_add(self.room_group_name,
self.channel_name)
# courotine that waiting for messgage received in self.channel_name
self.channel_receive = functools.partial(
self.channel_layer.receive, self.channel_name)
# pass message in from channel layer or client to dispatch method
try:
if self.channel_layer is not None:
await await_many_dispatch([receive, self.channel_receive],
self.dispatch)
else:
await await_many_dispatch([receive], self.dispatch)
except StopConsumer:
# exit cleanly
pass
async def __call__(self, receive: Receive, send: Send) -> None:
websocket = WebSocket(self.scope, receive=receive, send=send)
await self.on_connect(websocket)
close_code = status.WS_1000_NORMAL_CLOSURE
try:
while True:
message = await websocket.receive()
if message["type"] == "websocket.receive":
data = await self.decode(websocket, message)
await self.on_receive(websocket, data)
elif message["type"] == "websocket.disconnect":
close_code = int(message.get("code", status.WS_1000_NORMAL_CLOSURE))
break
except Exception as exc:
close_code = status.WS_1011_INTERNAL_ERROR
raise exc from None
finally:
await self.on_disconnect(websocket, close_code)
async def __call__(self, scope: Scope, receive: Receive,
send: Send) -> None:
if scope["type"] == "http":
request = Request(scope, receive, send)
try:
scope["user"], scope["auth"] = await self.authenticate(request)
except HTTPException as e:
response = await handle_http_exception(request, e)
await response(scope, receive, send)
return
elif scope["type"] == "websocket":
sock = WebSocket(scope, receive, send)
try:
scope["user"], scope["auth"] = await self.authenticate(sock)
except HTTPException as e:
await send({
"type": "websocket.close",
"code": 3700 + e.status_code
})
return
await self.app(scope, receive, send)
async def __call__(self, receive: Receive, send: Send) -> None:
websocket = WebSocket(self.scope, receive=receive, send=send)
kwargs = self.scope.get("kwargs", {})
await self.on_connect(websocket, **kwargs)
close_code = 1000
try:
while True:
message = await websocket.receive()
if message["type"] == "websocket.receive":
data = await self.decode(websocket, message)
await self.on_receive(websocket, data)
elif message["type"] == "websocket.disconnect":
close_code = int(message.get("code", 1000))
break
except Exception as exc:
close_code = 1011
raise exc from None
finally:
await self.on_disconnect(websocket, close_code)
async def __call__(self, scope: Scope, receive: Receive, send: Send):
if scope["type"] == "http":
await self.http_handler_class(
schema=self.schema,
graphiql=self.graphiql,
debug=self.debug,
get_context=self.get_context,
get_root_value=self.get_root_value,
process_result=self.process_result,
).handle(scope=scope, receive=receive, send=send)
elif scope["type"] == "websocket":
ws = WebSocket(scope=scope, receive=receive, send=send)
preferred_protocol = self.pick_preferred_protocol(ws)
if preferred_protocol == GRAPHQL_TRANSPORT_WS_PROTOCOL:
await self.graphql_transport_ws_handler_class(
schema=self.schema,
debug=self.debug,
connection_init_wait_timeout=self.connection_init_wait_timeout,
get_context=self.get_context,
get_root_value=self.get_root_value,
ws=ws,
).handle()
elif preferred_protocol == GRAPHQL_WS_PROTOCOL:
await self.graphql_ws_handler_class(
schema=self.schema,
debug=self.debug,
keep_alive=self.keep_alive,
keep_alive_interval=self.keep_alive_interval,
get_context=self.get_context,
get_root_value=self.get_root_value,
ws=ws,
).handle()
else:
# Subprotocol not acceptable
await ws.close(code=4406)
else: # pragma: no cover
raise ValueError("Unknown scope type: %r" % (scope["type"],))
async def handle_websocket(self, scope: Scope, receive: Receive,
send: Send):
websocket = WebSocket(scope=scope, receive=receive, send=send)
await websocket.accept(subprotocol="graphql-ws")
while websocket.application_state != WebSocketState.DISCONNECTED:
message = await websocket.receive_json()
operation_id = message.get("id")
message_type = message.get("type")
if message_type == GQL_CONNECTION_INIT:
await websocket.send_json({"type": GQL_CONNECTION_ACK})
if self.keep_alive:
asyncio.create_task(self.handle_keep_alive(websocket))
elif message_type == GQL_CONNECTION_TERMINATE:
await websocket.close()
elif message_type == GQL_START:
await self.start_subscription(message.get("payload"),
operation_id, websocket)
elif message_type == GQL_STOP:
await websocket.close()
