daphne/channels/database_layer.py

import base64
import datetime
import json
import random
import string
import time

from django.apps.registry import Apps
from django.db import DEFAULT_DB_ALIAS, connections, models, transaction
from django.db.utils import OperationalError
from django.utils import six
from django.utils.functional import cached_property
from django.utils.timezone import now


class DatabaseChannelLayer(object):
    """
    ORM-backed ASGI channel layer.

    For development use only; it will span multiple processes fine,
    but it's going to be pretty bad at throughput. If you're reading this and
    running it in production, PLEASE STOP.

    Also uses JSON for serialization, as we don't want to make Django depend
    on msgpack for the built-in backend. The JSON format uses \uffff as first
    character to signify a b64 byte string rather than a text string. Ugly, but
    it's not a valid Unicode character, so it should be safe enough.
    """

    def __init__(self, db_alias=DEFAULT_DB_ALIAS, expiry=60, group_expiry=86400):
        self.expiry = expiry
        self.group_expiry = group_expiry
        self.db_alias = db_alias

    # ASGI API

    extensions = ["groups", "flush"]

    def send(self, channel, message):
        # Typecheck
        assert isinstance(message, dict), "message is not a dict"
        assert isinstance(channel, six.text_type), "%s is not unicode" % channel
        # Write message to messages table
        self.channel_model.objects.create(
            channel=channel,
            content=self.serialize(message),
            expiry=now() + datetime.timedelta(seconds=self.expiry)
        )

    def receive_many(self, channels, block=False):
        if not channels:
            return None, None
        assert all(isinstance(channel, six.text_type) for channel in channels)
        # Shuffle channels
        channels = list(channels)
        random.shuffle(channels)
        # Clean out expired messages
        self._clean_expired()
        # Get a message from one of our channels
        while True:
            try:
                with transaction.atomic():
                    message = self.channel_model.objects.filter(channel__in=channels).order_by("id").first()
                    if message:
                        self.channel_model.objects.filter(pk=message.pk).delete()
                        return message.channel, self.deserialize(message.content)
            except OperationalError:
                # The database is probably trying to prevent a deadlock
                time.sleep(0.1)
                continue
            if block:
                time.sleep(1)
            else:
                return None, None

    def new_channel(self, pattern):
        assert isinstance(pattern, six.text_type)
        assert pattern.endswith("!")
        # Keep making channel names till one isn't present.
        while True:
            random_string = "".join(random.choice(string.ascii_letters) for i in range(10))
            new_name = pattern + random_string
            if not self.channel_model.objects.filter(channel=new_name).exists():
                return new_name

    # ASGI Group extension

    def group_add(self, group, channel):
        """
        Adds the channel to the named group for at least 'expiry'
        seconds (expiry defaults to message expiry if not provided).
        """
        self.group_model.objects.update_or_create(
            group=group,
            channel=channel,
        )

    def group_discard(self, group, channel):
        """
        Removes the channel from the named group if it is in the group;
        does nothing otherwise (does not error)
        """
        self.group_model.objects.filter(group=group, channel=channel).delete()

    def send_group(self, group, message):
        """
        Sends a message to the entire group.
        """
        self._clean_expired()
        for channel in self.group_model.objects.filter(group=group).values_list("channel", flat=True):
            self.send(channel, message)

    # ASGI Flush extension

    def flush(self):
        self.channel_model.objects.all().delete()
        self.group_model.objects.all().delete()

    # Serialization

    def serialize(self, message):
        return AsgiJsonEncoder().encode(message)

    def deserialize(self, message):
        return AsgiJsonDecoder().decode(message)

    # Database state mgmt

    @property
    def connection(self):
        """
        Returns the correct connection for the current thread.
        """
        return connections[self.db_alias]

    @cached_property
    def channel_model(self):
        """
        Initialises a new model to store messages; not done as part of a
        models.py as we don't want to make it for most installs.
        """
        # Make the model class
        class Message(models.Model):
            # We assume an autoincrementing PK for message order
            channel = models.CharField(max_length=200, db_index=True)
            content = models.TextField()
            expiry = models.DateTimeField(db_index=True)

            class Meta:
                apps = Apps()
                app_label = "channels"
                db_table = "django_channels"
        # Ensure its table exists
        if Message._meta.db_table not in self.connection.introspection.table_names(self.connection.cursor()):
            with self.connection.schema_editor() as editor:
                editor.create_model(Message)
        return Message

    @cached_property
    def group_model(self):
        """
        Initialises a new model to store groups; not done as part of a
        models.py as we don't want to make it for most installs.
        """
        # Make the model class
        class Group(models.Model):
            group = models.CharField(max_length=200)
            channel = models.CharField(max_length=200)
            created = models.DateTimeField(db_index=True, auto_now_add=True)

            class Meta:
                apps = Apps()
                app_label = "channels"
                db_table = "django_channel_groups"
                unique_together = [["group", "channel"]]
        # Ensure its table exists with the right schema
        if Group._meta.db_table not in self.connection.introspection.table_names(self.connection.cursor()):
            with self.connection.schema_editor() as editor:
                editor.create_model(Group)
        return Group

    def _clean_expired(self):
        """
        Cleans out expired groups and messages.
        """
        # Include a 1-second grace period for clock sync drift
        target = now() - datetime.timedelta(seconds=1)
        # First, go through old messages and pick out channels that got expired
        old_messages = self.channel_model.objects.filter(expiry__lt=target)
        channels_to_ungroup = old_messages.values_list("channel", flat=True).distinct()
        old_messages.delete()
        # Now, remove channel membership from channels that expired and ones that just expired
        self.group_model.objects.filter(
            models.Q(channel__in=channels_to_ungroup) |
            models.Q(created__lte=target - datetime.timedelta(seconds=self.group_expiry))
        ).delete()

    def __str__(self):
        return "%s(alias=%s)" % (self.__class__.__name__, self.connection.alias)


class AsgiJsonEncoder(json.JSONEncoder):
    """
    Special encoder that transforms bytestrings into unicode strings
    prefixed with u+ffff
    """

    def transform(self, o):
        if isinstance(o, (list, tuple)):
            return [self.transform(x) for x in o]
        elif isinstance(o, dict):
            return {
                self.transform(k): self.transform(v)
                for k, v in o.items()
            }
        elif isinstance(o, six.binary_type):
            return u"\uffff" + base64.b64encode(o).decode("ascii")
        else:
            return o

    def encode(self, o):
        return super(AsgiJsonEncoder, self).encode(self.transform(o))


class AsgiJsonDecoder(json.JSONDecoder):
    """
    Special encoder that transforms bytestrings into unicode strings
    prefixed with u+ffff
    """

    def transform(self, o):
        if isinstance(o, (list, tuple)):
            return [self.transform(x) for x in o]
        elif isinstance(o, dict):
            return {
                self.transform(k): self.transform(v)
                for k, v in o.items()
            }
        elif isinstance(o, six.text_type) and o and o[0] == u"\uffff":
            return base64.b64decode(o[1:].encode("ascii"))
        else:
            return o

    def decode(self, o):
        return self.transform(super(AsgiJsonDecoder, self).decode(o))