[update]add dpub support and add more easy way to initiate consumer and producer

Author: aohan237

.DS_Store

@@ -4,7 +4,7 @@ __pycache__/
 
 # C extensions
 *.so
-
+*_bak.py
 # Distribution / packaging
 .Python
 env/

.gitignore

@@ -1,7 +1,7 @@
 language: python
 python:
-    - 3.2
-    - 3.4
+    - 3.5
+    - 3.6
 
 env:
     - NSQ_DOWNLOAD=nsq-0.2.30.linux-amd64.go1.2.1

.travis.yml

@@ -1,20 +1,54 @@
 asyncnsq (work in progress,仍然开发中)
 =========================
+asyncnsq with python3.6 await/async supported
 
-implement python3.5+  async/await syntax support
-
-提供python3.5+  async/await 语法支持。 原项目，支持yield from 生成器语法。
-
-this project is forked from  jettify/aionsq, he is the original author
+Latest Updates
+--------------
 
-本项目fork了 jettify/aionsq， 他是原作者，但是很久没有更新了。
+support dpub 
 
 Usage examples
 --------------
-you can refer from examples.
 
-你可以查看examples文件夹中的例子。
+All you need is a loop, then enjoy
+
+Consumer:
+
+    loop = asyncio.get_event_loop()
+
+    async def go():
+        nsq_consumer = await create_nsq_consumer(host='tcp://127.0.0.1:4150',
+                                                 max_in_flight=200)
+        await nsq_consumer.subscribe('test_async_nsq', 'nsq')
+        for waiter in nsq_consumer.wait_messages():
+            message = await waiter
+            print(message.body)
+            await message.fin()
+
+    loop.run_until_complete(go())
+
+Producer:
+
+    loop = asyncio.get_event_loop()
+
+    async def go():
+        nsq_producer = await create_nsq_producer(host='127.0.0.1', port=4150,
+                                                 heartbeat_interval=30000,
+                                                 feature_negotiation=True,
+                                                 tls_v1=True,
+                                                 snappy=False,
+                                                 deflate=False,
+                                                 deflate_level=0,
+                                                 loop=loop)
+        for i in range(10):
+            await nsq_producer.pub('test_async_nsq', 'test_async_nsq:{i}'.format(i=i))
+            await nsq_producer.dpub('test_async_nsq', i * 1000,
+                                    'test_delay_async_nsq:{i}'.format(i=i))
+    loop.run_until_complete(go())
+
+you can use host like 'tcp://127.0.0.1:4150' or '127.0.0.1' with port=4150
 
+As for now, only single nsqd addr supported.
 
 Requirements
 ------------

README.md

@@ -1 +1,55 @@
-__version__ = '0.2.1'
+__version__ = '0.3.0'
+
+import asyncio
+from .utils import get_host_and_port
+from .nsq import Nsq
+from .consumer import NsqConsumer
+
+
+async def create_nsq_producer(host='127.0.0.1', port=4150, loop=None, queue=None,
+                              heartbeat_interval=30000, feature_negotiation=True,
+                              tls_v1=False, snappy=False, deflate=False, deflate_level=6,
+                              consumer=False, sample_rate=0):
+    """"
+    initial function to get producer
+    param: host: host addr with no protocol. 127.0.0.1 
+    param: port: host port 
+    param: queue: queue where all the msg been put from the nsq 
+    param: heartbeat_interval: heartbeat interval with nsq, set -1 to disable nsq heartbeat check
+    params: snappy: snappy compress
+    params: deflate: deflate compress  can't set True both with snappy
+    """
+    # TODO: add parameters type and value validation
+    host, tmp_port = get_host_and_port(host)
+    if not port:
+        port = tmp_port
+    loop = loop or asyncio.get_event_loop()
+    queue = queue or asyncio.Queue(loop=loop)
+    conn = Nsq(host=host, port=port, queue=queue,
+               heartbeat_interval=heartbeat_interval,
+               feature_negotiation=feature_negotiation,
+               tls_v1=tls_v1, snappy=snappy, deflate=deflate,
+               deflate_level=deflate_level,
+               sample_rate=sample_rate, consumer=consumer, loop=loop)
+    await conn.connect()
+    return conn
+
+
+async def create_nsq_consumer(host='127.0.0.1', port=None, loop=None, max_in_flight=42, lookupd_http_addresses=None):
+    """"
+    initial function to get consumer
+    param: host: host addr with no protocol. 127.0.0.1 
+    param: port: host port 
+    param: max_in_flight: number of messages get but not finish or req
+    param: lookupd_http_addresses: heartbeat interval with nsq, set -1 to disable nsq heartbeat check
+    """
+    # TODO: add parameters type and value validation
+    host, tmp_port = get_host_and_port(host)
+    if not port:
+        port = tmp_port or 4150
+    loop = loop or asyncio.get_event_loop()
+    conn = NsqConsumer(nsqd_tcp_addresses=[(host, port)],
+                       lookupd_http_addresses=lookupd_http_addresses,
+                       max_in_flight=max_in_flight, loop=loop)
+    await conn.connect()
+    return conn

asyncnsq/__init__.py

@@ -179,7 +179,8 @@ async def _read_data(self):
         is_canceled = False
         while not self._reader.at_eof():
             try:
-                data = await self._reader.read(consts.MAX_CHUNK_SIZE)
+                data = await self._reader.read(1024)
+                # print('\n\n', 'socket response data', data)
             except asyncio.CancelledError:
                 is_canceled = True
                 logger.error('Task is canceled')

asyncnsq/connection.py

@@ -29,6 +29,7 @@
 AUTH = b'AUTH'
 SUB = b'SUB'
 PUB = b'PUB'
+DPUB = b'DPUB'
 
 # connection status
 CLOSED = 0

asyncnsq/consts.py

@@ -30,7 +30,7 @@ async def fin(self):
         """
         if self._is_processed:
             raise RuntimeWarning("Message has already been processed")
-        resp = (await self.conn.execute(FIN, self.message_id))
+        resp = await self.conn.execute(FIN, self.message_id)
         self._is_processed = True
         return resp
 
@@ -53,4 +53,4 @@ async def touch(self):
         """
         if self._is_processed:
             raise RuntimeWarning("Message has already been processed")
-        return (await self.conn.execute(TOUCH, self.message_id))
+        return await self.conn.execute(TOUCH, self.message_id)

asyncnsq/containers.py

@@ -2,23 +2,32 @@
 from . import consts
 import time
 from .log import logger
-from .utils import retry_iterator
+from .utils import retry_iterator, RdyControl
 from .connection import create_connection
-from .consts import TOUCH, REQ, FIN, RDY, CLS, MPUB, PUB, SUB, AUTH
+from .consts import TOUCH, REQ, FIN, RDY, CLS, MPUB, PUB, SUB, AUTH, DPUB
 
 
 async def create_nsq(host='127.0.0.1', port=4150, loop=None, queue=None,
                      heartbeat_interval=30000, feature_negotiation=True,
                      tls_v1=False, snappy=False, deflate=False, deflate_level=6,
-                     sample_rate=0):
+                     consumer=False, sample_rate=0):
+    """"
+    param: host: host addr with no protocol. 127.0.0.1 
+    param: port: host port 
+    param: queue: queue where all the msg been put from the nsq 
+    param: heartbeat_interval: heartbeat interval with nsq, set -1 to disable nsq heartbeat check
+    params: snappy: snappy compress
+    params: deflate: deflate compress  can't set True both with snappy
+    """
     # TODO: add parameters type and value validation
+    loop = loop or asyncio.get_event_loop()
     queue = queue or asyncio.Queue(loop=loop)
     conn = Nsq(host=host, port=port, queue=queue,
                heartbeat_interval=heartbeat_interval,
                feature_negotiation=feature_negotiation,
                tls_v1=tls_v1, snappy=snappy, deflate=deflate,
                deflate_level=deflate_level,
-               sample_rate=sample_rate, loop=loop)
+               sample_rate=sample_rate, consumer=consumer, loop=loop)
     await conn.connect()
     return conn
 
@@ -28,7 +37,7 @@ class Nsq:
     def __init__(self, host='127.0.0.1', port=4150, loop=None, queue=None,
                  heartbeat_interval=30000, feature_negotiation=True,
                  tls_v1=False, snappy=False, deflate=False, deflate_level=6,
-                 sample_rate=0):
+                 sample_rate=0, consumer=False, max_in_flight=42):
         # TODO: add parameters type and value validation
         self._config = {
             "deflate": deflate,
@@ -53,6 +62,13 @@ def __init__(self, host='127.0.0.1', port=4150, loop=None, queue=None,
 
         self._on_rdy_changed_cb = None
         self._last_rdy = 0
+        self.consumer = consumer
+        if self.consumer:
+            self._idle_timeout = 10
+            self._max_in_flight = max_in_flight
+            self._rdy_control = RdyControl(idle_timeout=self._idle_timeout,
+                                           max_in_flight=self._max_in_flight,
+                                           loop=self._loop)
 
     async def connect(self):
         self._conn = await create_connection(self._host, self._port,
@@ -61,6 +77,8 @@ async def connect(self):
         self._conn._on_message = self._on_message
         await self._conn.identify(**self._config)
         self._status = consts.CONNECTED
+        if self.consumer:
+            self._rdy_control.add_connection(self._conn)
 
     def _on_message(self, msg):
         # should not be coroutine
@@ -102,7 +120,7 @@ async def reconnect(self):
             await asyncio.sleep(t, loop=self._loop)
 
     async def execute(self, command, *args, data=None):
-        if self._state <= consts.CONNECTED and self._reconnect:
+        if self._status <= consts.CONNECTED and self._reconnect:
             await self.reconnect()
 
         response = self._conn.execute(command, *args, data=data)
@@ -113,8 +131,10 @@ def id(self):
         return self._conn.endpoint
 
     def wait_messages(self):
+        # print('wait_messages')
         while True:
             future = asyncio.ensure_future(self._queue.get(), loop=self._loop)
+            # print(future)
             yield future
 
     async def auth(self, secret):
@@ -123,7 +143,7 @@ async def auth(self, secret):
         :param secret:
         :return:
         """
-        return (await self._conn.execute(AUTH, data=secret))
+        return await self._conn.execute(AUTH, data=secret)
 
     async def sub(self, topic, channel):
         """
@@ -132,7 +152,7 @@ async def sub(self, topic, channel):
         :param channel:
         :return:
         """
-        return (await self._conn.execute(SUB, topic, channel))
+        return await self._conn.execute(SUB, topic, channel)
 
     async def pub(self, topic, message):
         """
@@ -141,7 +161,19 @@ async def pub(self, topic, message):
         :param message:
         :return:
         """
-        return (await self._conn.execute(PUB, topic, data=message))
+        return await self._conn.execute(PUB, topic, data=message)
+
+    async def dpub(self, topic, delay_time, message):
+        """
+
+        :param topic:
+        :param message:
+        :param delay_time: delayed time in millisecond
+        :return:
+        """
+        if not delay_time or delay_time is None:
+            delay_time = 0
+        return await self._conn.execute(DPUB, topic, delay_time, data=message)
 
     async def mpub(self, topic, message, *messages):
         """
@@ -152,7 +184,7 @@ async def mpub(self, topic, message, *messages):
         :return:
         """
         msgs = [message] + list(messages)
-        return (await self._conn.execute(MPUB, topic, data=msgs))
+        return await self._conn.execute(MPUB, topic, data=msgs)
 
     async def rdy(self, count):
         """
@@ -165,15 +197,15 @@ async def rdy(self, count):
 
         self._last_rdy = count
         self.rdy_state = count
-        return (await self._conn.execute(RDY, count))
+        return await self._conn.execute(RDY, count)
 
     async def fin(self, message_id):
         """
 
         :param message_id:
         :return:
         """
-        return (await self._conn.execute(FIN, message_id))
+        return await self._conn.execute(FIN, message_id)
 
     async def req(self, message_id, timeout):
         """
@@ -182,15 +214,15 @@ async def req(self, message_id, timeout):
         :param timeout:
         :return:
         """
-        return (await self._conn.execute(REQ, message_id, timeout))
+        return await self._conn.execute(REQ, message_id, timeout)
 
     async def touch(self, message_id):
         """
 
         :param message_id:
         :return:
         """
-        return (await self._conn.execute(TOUCH, message_id))
+        return await self._conn.execute(TOUCH, message_id)
 
     async def cls(self):
         """

asyncnsq/nsq.py

@@ -68,6 +68,7 @@ def gets(self):
 
     def encode_command(self, cmd, *args, data=None):
         cmd = self._parser.encode_command(cmd, *args, data=data)
+        # print(cmd)
         return self.compress(cmd)
 
 

asyncnsq/protocol.py

@@ -2,11 +2,28 @@
 import random
 import re
 import time
+from urllib.parse import urlparse
+
 
 TOPIC_NAME_RE = re.compile(r'^[\.a-zA-Z0-9_-]+$')
 CHANNEL_NAME_RE = re.compile(r'^[\.a-zA-Z0-9_-]+(#ephemeral)?$')
 
 
+def get_host_and_port(host):
+    host_parsed = urlparse(host)
+    if host_parsed.scheme == 'tcp':
+        result = host_parsed.netloc
+    elif host_parsed.scheme == '':
+        result = host_parsed.path
+    else:
+        result = host
+    result = result.split(':')
+    if len(result) == 2:
+        return result[0], result[-1]
+    else:
+        return result[0], None
+
+
 def valid_topic_name(topic):
     if not 0 < len(topic) < 33:
         return False
@@ -113,6 +130,9 @@ def add_connections(self, connections):
         for conn in self._connections.values():
             conn._on_rdy_changed_cb = self.rdy_changed
 
+    def add_connection(self, connection):
+        connection._on_rdy_changed_cb = self.rdy_changed
+
     def rdy_changed(self, conn_id):
         self._cmd_queue.put_nowait((CHANGE_CONN_RDY, (conn_id,)))