DEV Community

Cover image for Django REST Framework - WebSocket
Burak Yılmaz
Burak Yılmaz

Posted on • Updated on

Django REST Framework - WebSocket

In this article we will build an application which uses JWT Authentication that communicates to websocket with Django REST Framework. The main focus of this article is send data to websocket from out of consumer.

Before we start you can find the source code here

We have two requests to same endpoint GET and POST, whenever a request comes to endpoint we will send messages to websocket.

This is our models.py, our Message model looks like down below which is very simple.

from django.db import models
from django.contrib.auth import get_user_model

User = get_user_model()

class Message(models.Model):
    message = models.JSONField()
    user = models.ForeignKey(to=User, on_delete=models.CASCADE, null=True, blank=True)

Enter fullscreen mode Exit fullscreen mode

Configurations

We need a routing.py file which includes url's of related app.

from django.urls import re_path

from . import consumers

websocket_urlpatterns = [
    re_path(r'msg/', consumers.ChatConsumer.as_asgi()),
]

Enter fullscreen mode Exit fullscreen mode

Then we need another routing.py file into project's core directory.

import os

from websocket.middlewares import WebSocketJWTAuthMiddleware
from channels.routing import ProtocolTypeRouter, URLRouter
from django.core.asgi import get_asgi_application

from websocket import routing

os.environ.setdefault("DJANGO_SETTINGS_MODULE", "django_ws.settings")


application = ProtocolTypeRouter(
    {
        "http": get_asgi_application(),
        "websocket": WebSocketJWTAuthMiddleware(URLRouter(routing.websocket_urlpatterns)),
    }
)

Enter fullscreen mode Exit fullscreen mode

As you can see here a middleware is used, what this middleware does is control the authentication process, if you follow this link you will see that Django has support for standard Django authentication, since we are using JWT Authentication a custom middleware is needed. In this project Rest Framework SimpleJWT was used, when create a connection we are sending token with a query-string which is NOT so secure. We assigned user to scope as down below.
middlewares.py

from urllib.parse import parse_qs

from channels.db import database_sync_to_async
from django.contrib.auth import get_user_model
from django.contrib.auth.models import AnonymousUser
from rest_framework_simplejwt.tokens import AccessToken, TokenError

User = get_user_model()


@database_sync_to_async
def get_user(user_id):
    try:
        return User.objects.get(id=user_id)
    except User.DoesNotExist:
        return AnonymousUser()


class WebSocketJWTAuthMiddleware:

    def __init__(self, app):
        self.app = app

    async def __call__(self, scope, receive, send):
        parsed_query_string = parse_qs(scope["query_string"])
        token = parsed_query_string.get(b"token")[0].decode("utf-8")

        try:
            access_token = AccessToken(token)
            scope["user"] = await get_user(access_token["user_id"])
        except TokenError:
            scope["user"] = AnonymousUser()

        return await self.app(scope, receive, send)
Enter fullscreen mode Exit fullscreen mode

Last but not least settings.py file, we are using Redis as channel layer therefore we need to start a Redis server, we can do that with docker

docker run -p 6379:6379 -d redis:5

settings.py

ASGI_APPLICATION = 'django_ws.routing.application'

CHANNEL_LAYERS = {
    "default": {
        "BACKEND": "channels_redis.core.RedisChannelLayer",
        "CONFIG": {
            "hosts": [('0.0.0.0', 6379)],
        },
    },
}
Enter fullscreen mode Exit fullscreen mode

Consumers

Here is our consumer.py file, since ChatConsumer is asynchronous when a database access is needed, related method needs a database_sync_to_async decorator.

websocket/consumer.py

import json

from channels.db import database_sync_to_async
from channels.generic.websocket import AsyncWebsocketConsumer
from django.contrib.auth.models import AnonymousUser

from websocket.models import Message


class ChatConsumer(AsyncWebsocketConsumer):
    groups = ["general"]

    async def connect(self):
        await self.accept()
        if self.scope["user"] is not AnonymousUser:
            self.user_id = self.scope["user"].id
            await self.channel_layer.group_add(f"{self.user_id}-message", self.channel_name)

    async def send_info_to_user_group(self, event):
        message = event["text"]
        await self.send(text_data=json.dumps(message))

    async def send_last_message(self, event):
        last_msg = await self.get_last_message(self.user_id)
        last_msg["status"] = event["text"]
        await self.send(text_data=json.dumps(last_msg))

    @database_sync_to_async
    def get_last_message(self, user_id):
        message = Message.objects.filter(user_id=user_id).last()
        return message.message

Enter fullscreen mode Exit fullscreen mode

Views

As I mentioned at previous step our consumer is asynchronous we need to convert methods from async to sync just like the name of the function

views.py

from asgiref.sync import async_to_sync
from channels.layers import get_channel_layer
from rest_framework import status
from rest_framework.permissions import IsAuthenticated
from rest_framework.response import Response
from rest_framework.views import APIView

from .models import Message


class MessageSendAPIView(APIView):
    permission_classes = (IsAuthenticated,)

    def get(self, request):
        channel_layer = get_channel_layer()
        async_to_sync(channel_layer.group_send)(
            "general", {"type": "send_info_to_user_group",
                        "text": {"status": "done"}}
        )

        return Response({"status": True}, status=status.HTTP_200_OK)

    def post(self, request):
        msg = Message.objects.create(user=request.user, message={
                                     "message": request.data["message"]})
        socket_message = f"Message with id {msg.id} was created!"
        channel_layer = get_channel_layer()
        async_to_sync(channel_layer.group_send)(
            f"{request.user.id}-message", {"type": "send_last_message",
                                           "text": socket_message}
        )

        return Response({"status": True}, status=status.HTTP_201_CREATED)

Enter fullscreen mode Exit fullscreen mode

In this view's GET request we are sending message to channel's "general" group so everyone on this group will receive that message, if you check the consumers.py you will see that our default group is "general", on the other hand in POST request we are sending our message to a specified group in an other saying, this message is only sent to a group with related user's id which means only this user receives the message.

Up to this point everything seems fine but we can't find out if it actually works fine until we try, so let's do it. I am using a Chrome extension to connect websocket.

Result of GET request

Result of GET request

POST request

POST request

NOTE: Python3.10.0 have compatibility issue with asyncio be sure not using this version.

Discussion (10)

Collapse
dhruvindsddev profile image
Dhruvin

Were you successfully able to push this in prod ?

Collapse
buurak profile image
Burak Yılmaz Author

Yes but some DevOps configurations needed on prod and I am using Uvicorn to serve it.

Collapse
dhruvindsddev profile image
Dhruvin

Do have any scaling issues ? I had implemented it in daphene and faced a lot of slowing down when scaling. Would really like to know more about your approach

Thread Thread
buurak profile image
Burak Yılmaz Author

docs.djangoproject.com/en/4.0/howt...
Take a look at this documentation daphne may not be enough by itself on production

Collapse
sherrydays profile image
Sherry Day

Nice post

Collapse
buurak profile image
Burak Yılmaz Author

thank you :)

Collapse
selamet profile image
Selamet

it has been an important and understandable content on the subject, thank you

Collapse
buurak profile image
Burak Yılmaz Author

Thank you and you are welcome buddy, hope it helped! <3

Collapse
merveealpay profile image
Merve Alpay

wow! Great post

Collapse
buurak profile image
Burak Yılmaz Author

thanks a lot :)