diff --git a/channels/exceptions.py b/channels/exceptions.py index afadf9e..1827629 100644 --- a/channels/exceptions.py +++ b/channels/exceptions.py @@ -1,3 +1,7 @@ +from __future__ import unicode_literals +import six + + class ConsumeLater(Exception): """ Exception that says that the current message should be re-queued back @@ -39,6 +43,33 @@ class DenyConnection(Exception): pass +class ChannelSocketException(Exception): + """ + Base Exception is intended to run some action ('run' method) + when it is raised at a consumer body + """ + + def run(self, message): + raise NotImplementedError + + +class WebsocketCloseException(ChannelSocketException): + """ + ChannelSocketException based exceptions for close websocket connection with code + """ + + def __init__(self, code=None): + if code is not None and not isinstance(code, six.integer_types) \ + and code != 1000 and not (3000 <= code <= 4999): + raise ValueError("invalid close code {} (must be 1000 or from [3000, 4999])".format(code)) + self._code = code + + def run(self, message): + if message.reply_channel.name.split('.')[0] != "websocket": + raise ValueError("You cannot raise CloseWebsocketError from a non-websocket handler.") + message.reply_channel.send({"close": self._code or True}) + + class SendNotAvailableOnDemultiplexer(Exception): """ Raised when trying to send with a WebsocketDemultiplexer. Use the multiplexer instead. diff --git a/channels/worker.py b/channels/worker.py index b44b6ac..213501f 100644 --- a/channels/worker.py +++ b/channels/worker.py @@ -8,7 +8,7 @@ import sys import threading import time -from .exceptions import ConsumeLater, DenyConnection +from .exceptions import ChannelSocketException, ConsumeLater, DenyConnection from .message import Message from .signals import consumer_finished, consumer_started, worker_ready from .utils import name_that_thing @@ -122,6 +122,8 @@ class Worker(object): if message.channel.name != "websocket.connect": raise ValueError("You cannot DenyConnection from a non-websocket.connect handler.") message.reply_channel.send({"close": True}) + except ChannelSocketException as e: + e.run(message) except ConsumeLater: # They want to not handle it yet. Re-inject it with a number-of-tries marker. content['__retries__'] = content.get("__retries__", 0) + 1