Add support for nsq auth protocal.

Author: wtolson

HISTORY.rst

@@ -3,10 +3,11 @@
 History
 -------
 
-0.1.5 (TBD)
+0.2.0 (TBD)
 ~~~~~~~~~~~~~~~~~~
 * Warn on connection failure.
 * Add extra requires for snappy.
+* Add support for nsq auth protocal.
 
 0.1.4 (2014-07-24)
 ~~~~~~~~~~~~~~~~~~

gnsq/errors.py

@@ -69,6 +69,18 @@ class NSQMPubFailed(NSQErrorCode):
     """E_MPUB_FAILED"""
 
 
+class NSQAuthDisabled(NSQErrorCode):
+    """E_AUTH_DISABLED"""
+
+
+class NSQAuthFailed(NSQErrorCode):
+    """E_AUTH_FAILED"""
+
+
+class NSQUnauthorized(NSQErrorCode):
+    """E_UNAUTHORIZED"""
+
+
 class NSQFinishFailed(NSQErrorCode):
     """E_FIN_FAILED"""
     fatal = False
@@ -94,6 +106,9 @@ class NSQTouchFailed(NSQErrorCode):
     'E_PUB_FAILED': NSQPubFailed,
     'E_MPUB_FAILED': NSQMPubFailed,
     'E_FINISH_FAILED': NSQFinishFailed,
+    'E_AUTH_DISABLED': NSQAuthDisabled,
+    'E_AUTH_FAILED': NSQAuthFailed,
+    'E_UNAUTHORIZED': NSQUnauthorized,
     'E_FIN_FAILED': NSQFinishFailed,
     'E_REQUEUE_FAILED': NSQRequeueFailed,
     'E_REQ_FAILED': NSQRequeueFailed,

gnsq/nsqd.py

@@ -71,6 +71,9 @@ class Nsqd(HTTPClient):
         less than 100 and the client will receive that percentage of the message
         traffic. (requires nsqd 0.2.25+)
 
+    :param auth_secret: a string passed when using nsq auth (requires
+        nsqd 0.2.29+)
+
     :param user_agent: a string identifying the agent for this client in the
         spirit of HTTP (default: ``<client_library_name>/<version>``) (requires
         nsqd 0.2.25+)
@@ -92,6 +95,7 @@ def __init__(
         deflate=False,
         deflate_level=6,
         sample_rate=0,
+        auth_secret=None,
         user_agent=USERAGENT,
     ):
         self.address = address
@@ -110,6 +114,7 @@ def __init__(
         self.deflate = deflate
         self.deflate_level = deflate_level
         self.sample_rate = sample_rate
+        self.auth_secret = auth_secret
         self.user_agent = user_agent
 
         self.state = INIT
@@ -173,6 +178,17 @@ def on_requeue(self):
         """
         return blinker.Signal(doc='Emitted after the a message is requeued.')
 
+    @cached_property
+    def on_auth(self):
+        """Emitted after the connection is successfully authenticated.
+
+        The signal sender is the connection and the parsed `response` is sent as
+        arguments.
+        """
+        return blinker.Signal(
+            doc='Emitted after the connection is successfully authenticated.'
+        )
+
     @cached_property
     def on_close(self):
         """Emitted after :meth:`close_stream`.
@@ -296,17 +312,28 @@ def listen(self):
         while self.is_connected:
             self.read_response()
 
+    def check_ok(self, expected='OK'):
+        frame, data = self.read_response()
+        if frame == nsq.FRAME_TYPE_ERROR:
+            raise data
+
+        if frame != nsq.FRAME_TYPE_RESPONSE:
+            raise errors.NSQException('expected response frame')
+
+        if data != expected:
+            raise errors.NSQException('unexpected response %r' % data)
+
     def upgrade_to_tls(self):
         self.stream.upgrade_to_tls(**self.tls_options)
-        return self.read_response()
+        self.check_ok()
 
     def upgrade_to_snappy(self):
         self.stream.upgrade_to_snappy()
-        return self.read_response()
+        self.check_ok()
 
     def upgrade_to_defalte(self):
         self.stream.upgrade_to_defalte(self.deflate_level)
-        return self.read_response()
+        self.check_ok()
 
     def identify(self):
         """Update client metadata on the server and negotiate features.
@@ -372,8 +399,29 @@ def identify(self):
             self.deflate_level = data.get('deflate_level', self.deflate_level)
             self.upgrade_to_defalte()
 
+        if self.auth_secret and data.get('auth_required'):
+            self.auth()
+
         return data
 
+    def auth(self):
+        """Send authorization secret to nsqd."""
+        self.send(nsq.auth(self.auth_secret))
+        frame, data = self.read_response()
+
+        if frame == nsq.FRAME_TYPE_ERROR:
+            raise data
+
+        try:
+            response = json.loads(data)
+        except ValueError:
+            self.close_stream()
+            msg = 'failed to parse AUTH response JSON from nsqd: %r'
+            raise errors.NSQException(msg % data)
+
+        self.on_auth.send(self, response=response)
+        return response
+
     def subscribe(self, topic, channel):
         """Subscribe to a nsq `topic` and `channel`."""
         self.send(nsq.subscribe(topic, channel))

gnsq/protocol.py

@@ -116,6 +116,10 @@ def identify(data):
     return _command('IDENTIFY', json.dumps(data))
 
 
+def auth(secret):
+    return _command('AUTH', secret)
+
+
 def subscribe(topic_name, channel_name):
     assert_valid_topic_name(topic_name)
     assert_valid_channel_name(channel_name)

gnsq/reader.py

@@ -224,6 +224,15 @@ def on_giving_up(self):
         """
         return blinker.Signal(doc='Sent after a giving up on a message.')
 
+    @cached_property
+    def on_auth(self):
+        """Emitted after a connection is successfully authenticated.
+
+        The signal sender is the reader and the `conn` and parsed `response` are
+        sent as arguments.
+        """
+        return blinker.Signal(doc='Emitted when a response is received.')
+
     @cached_property
     def on_exception(self):
         """Emitted when an exception is caught while handling a message.
@@ -536,6 +545,7 @@ def connect_to_nsqd(self, conn):
         conn.on_error.connect(self.handle_error)
         conn.on_finish.connect(self.handle_finish)
         conn.on_requeue.connect(self.handle_requeue)
+        conn.on_auth.connect(self.handle_auth)
 
         if self.max_concurrency:
             conn.on_message.connect(self.queue_message)
@@ -672,6 +682,20 @@ def handle_requeue(self, conn, message_id, timeout):
         self.handle_backoff()
         self.on_requeue.send(self, message_id=message_id, timeout=timeout)
 
+    def handle_auth(self, conn, response):
+        metadata = []
+        if response.get('identity'):
+            metadata.append("Identity: %r" % response['identity'])
+
+        if response.get('permission_count'):
+            metadata.append("Permissions: %d" % response['permission_count'])
+
+        if response.get('identity_url'):
+            metadata.append(response['identity_url'])
+
+        self.logger.info('[%s] AUTH accepted %s' % (conn, ' '.join(metadata)))
+        self.on_auth.send(self, conn=conn, response=response)
+
     def handle_backoff(self):
         if self.state in (BACKOFF, CLOSED):
             return

tests/test_nsqd.py

@@ -324,7 +324,7 @@ def handle(socket, address):
 
         conn.connect()
         resp = conn.auth()
-        assert resp['identify'] == 'awesome'
+        assert resp['identity'] == 'awesome'
 
 
 def test_identify_auth():
@@ -354,22 +354,18 @@ def handle(socket, address):
             auth_secret='secret'
         )
 
-        server.auth_was_called = False
-        _auth = conn.auth
+        @conn.on_auth.connect
+        def assert_auth(conn, response):
+            assert assert_auth.was_called is False
+            assert_auth.was_called = True
+            assert response['identity'] == 'awesome'
 
-        def auth():
-            assert server.auth_was_called is False
-            server.auth_was_called = True
-
-            resp = _auth()
-            assert resp['identify'] == 'awesome'
-
-        conn.auth = auth
+        assert_auth.was_called = False
         conn.connect()
-
         resp = conn.identify()
+
         assert resp['auth_required']
-        assert server.auth_was_called
+        assert assert_auth.was_called
 
 
 @pytest.mark.parametrize('tls,deflate,snappy', product((True, False), repeat=3))