Merge pull request #276 from YiuRULE/custom_prefix

Add inbox prefix as a configuration variable

Author: wallyqs

nats/aio/client.py

@@ -73,8 +73,7 @@
 
 PING_PROTO = PING_OP + _CRLF_
 PONG_PROTO = PONG_OP + _CRLF_
-INBOX_PREFIX = bytearray(b'_INBOX.')
-INBOX_PREFIX_LEN = len(INBOX_PREFIX) + 22 + 1
+DEFAULT_INBOX_PREFIX = b'_INBOX'
 
 DEFAULT_PENDING_SIZE = 1024 * 1024
 DEFAULT_BUFFER_SIZE = 32768
@@ -188,6 +187,7 @@ def __init__(self) -> None:
         self._resp_map: Dict[str, asyncio.Future] = {}
         self._resp_sub_prefix: Optional[bytearray] = None
         self._nuid = NUID()
+        self._inbox_prefix = bytearray(DEFAULT_INBOX_PREFIX)
 
         # NKEYS support
         #
@@ -247,6 +247,7 @@ async def connect(
         user_jwt_cb: Optional[JWTCallback] = None,
         user_credentials: Optional[Credentials] = None,
         nkeys_seed: Optional[str] = None,
+        inbox_prefix: Union[str, bytes] = DEFAULT_INBOX_PREFIX,
     ) -> None:
         """
         Establishes a connection to NATS.
@@ -347,6 +348,12 @@ async def subscribe_handler(msg):
         self._reconnected_cb = reconnected_cb
         self._disconnected_cb = disconnected_cb
 
+        # Custom inbox prefix
+        if isinstance(inbox_prefix, str):
+            inbox_prefix = inbox_prefix.encode()
+        assert isinstance(inbox_prefix, bytes)
+        self._inbox_prefix = bytearray(inbox_prefix)
+
         # NKEYS support
         self._signature_cb = signature_cb
         self._user_jwt_cb = user_jwt_cb
@@ -822,7 +829,8 @@ async def _send_subscribe(self, sub: Subscription) -> None:
     async def _init_request_sub(self) -> None:
         self._resp_map = {}
 
-        self._resp_sub_prefix = INBOX_PREFIX[:]
+        self._resp_sub_prefix = self._inbox_prefix[:]
+        self._resp_sub_prefix.extend(b'.')
         self._resp_sub_prefix.extend(self._nuid.next())
         self._resp_sub_prefix.extend(b'.')
         resp_mux_subject = self._resp_sub_prefix[:]
@@ -832,7 +840,7 @@ async def _init_request_sub(self) -> None:
         )
 
     async def _request_sub_callback(self, msg: Msg) -> None:
-        token = msg.subject[INBOX_PREFIX_LEN:]
+        token = msg.subject[len(self._inbox_prefix) + 22 + 2:]
         try:
             fut = self._resp_map.get(token)
             if not fut:
@@ -916,7 +924,8 @@ def new_inbox(self) -> str:
            nc.publish('broadcast', b'', reply=inbox)
            msg = sub.next_msg()
         """
-        next_inbox = INBOX_PREFIX[:]
+        next_inbox = self._inbox_prefix[:]
+        next_inbox.extend(b'.')
         next_inbox.extend(self._nuid.next())
         return next_inbox.decode()
 

tests/test_client.py

@@ -741,6 +741,19 @@ async def worker_handler(msg):
 
         await nc.close()
 
+    @async_test
+    async def test_custom_inbox_prefix(self):
+        nc = NATS()
+
+        async def worker_handler(msg):
+            self.assertTrue(msg.reply.startswith('bar.'))
+            await msg.respond(b"OK")
+
+        await nc.connect(inbox_prefix="bar")
+        await nc.subscribe("foo", cb=worker_handler)
+        await nc.request("foo", b'')
+        await nc.close()
+
     @async_test
     async def test_msg_respond(self):
         nc = NATS()