User friendly way to close websocket with code (via CloseException) (#468)

* User friendly way to close websocket with status

* More generic way  to close(whatever) connection by exception

* Fix import ordering for exceptions (isort)
This commit is contained in:
Krukov D 2017-01-12 21:02:30 +03:00 committed by Andrew Godwin
parent 2ced4ee2e9
commit 37da462411
2 changed files with 34 additions and 1 deletions

View File

@ -1,3 +1,7 @@
from __future__ import unicode_literals
import six
class ConsumeLater(Exception): class ConsumeLater(Exception):
""" """
Exception that says that the current message should be re-queued back Exception that says that the current message should be re-queued back
@ -39,6 +43,33 @@ class DenyConnection(Exception):
pass 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): class SendNotAvailableOnDemultiplexer(Exception):
""" """
Raised when trying to send with a WebsocketDemultiplexer. Use the multiplexer instead. Raised when trying to send with a WebsocketDemultiplexer. Use the multiplexer instead.

View File

@ -8,7 +8,7 @@ import sys
import threading import threading
import time import time
from .exceptions import ConsumeLater, DenyConnection from .exceptions import ChannelSocketException, ConsumeLater, DenyConnection
from .message import Message from .message import Message
from .signals import consumer_finished, consumer_started, worker_ready from .signals import consumer_finished, consumer_started, worker_ready
from .utils import name_that_thing from .utils import name_that_thing
@ -122,6 +122,8 @@ class Worker(object):
if message.channel.name != "websocket.connect": if message.channel.name != "websocket.connect":
raise ValueError("You cannot DenyConnection from a non-websocket.connect handler.") raise ValueError("You cannot DenyConnection from a non-websocket.connect handler.")
message.reply_channel.send({"close": True}) message.reply_channel.send({"close": True})
except ChannelSocketException as e:
e.run(message)
except ConsumeLater: except ConsumeLater:
# They want to not handle it yet. Re-inject it with a number-of-tries marker. # They want to not handle it yet. Re-inject it with a number-of-tries marker.
content['__retries__'] = content.get("__retries__", 0) + 1 content['__retries__'] = content.get("__retries__", 0) + 1