Add custom re-subscribe method

README.rst

@@ -124,7 +124,8 @@ You can install development dependencies with:
 Release history
 ---------------
 
-- **v. 0.3.0**: Remove date-time parsing of the replay mixin.
+- **v. 0.3.0**: Remove date-time parsing of the replay mixin, Allow client to
+  customize ``ReSubscribeMixin`` retry conditions.
 - **v. 0.2.0**: Add refresh token authentication
 - **v. 0.1.1**: Add documentation and initial typing information.
 - **v. 0.1.0**: Initial release.

aio_sf_streaming/mixins.py

@@ -2,7 +2,6 @@
 Mixins module: Provide various mixins modules
 """
 import asyncio
-import datetime
 import enum
 import logging
 from typing import Union
@@ -95,14 +94,10 @@ async def messages(self) -> JSONObject:
 
                 # Create a task : do not wait the replay id is stored to
                 # reconnect as soon as possible
-                self.loop.create_task(
-                    self.store_replay_id(channel, replay_id, creation_time)
-                )
+                self.loop.create_task(self.store_replay_id(channel, replay_id, creation_time))
             yield message
 
-    async def store_replay_id(
-        self, channel: str, replay_id: int, creation_time: str
-    ) -> None:
+    async def store_replay_id(self, channel: str, replay_id: int, creation_time: str) -> None:
         """
         Callback called to store a replay id. You should override this method
         to implement your custom logic.
@@ -175,10 +170,7 @@ async def messages(self) -> JSONObject:
             channel = message["channel"]
 
             # If asked, perform a new handshake
-            if (
-                channel.startswith("/meta/")
-                and message.get("error") == "403::Unknown client"
-            ):
+            if channel.startswith("/meta/") and message.get("error") == "403::Unknown client":
                 logger.info("Disconnected, do new handshake")
                 await self.handshake()
                 continue
@@ -224,25 +216,58 @@ def __init__(self, retry_sub_duration: float = 0.1, **kwargs):
         super().__init__(**kwargs)
         self.retry_sub_duration = retry_sub_duration
 
+    async def should_retry_on_exception(self, channel: str, exception: Exception) -> bool:
+        """
+        Callback called to process an exception raised during subscription.
+        Return a boolean if we must retry. If ``False`` is returned, the exception will be
+        propagated to caller.
+
+        By-default, do return always ``False``.
+
+        :param channel: Channel name
+        :param exception: The exception raised
+        """
+        return False
+
+    async def should_retry_on_error_response(self, channel: str, response: JSONObject) -> bool:
+        """
+        Callback called to process a response with and error message.
+        Return a boolean if we must retry. If ``False`` is returned, the response will be
+        returned to caller.
+
+        By-default, retry on known 'server unavailable' response.
+
+        :param channel: Channel name
+        :param response: The response received
+        """
+        return (
+            response[0]
+            .get("ext", {})
+            .get("sfdc", {})
+            .get("failureReason", "")
+            .startswith("SERVER_UNAVAILABLE")
+        )
+
     async def subscribe(self, channel: str) -> JSONList:
         """
         See :py:func:`BaseSalesforceStreaming.subscribe`
         """
         while True:
-            response = await super().subscribe(channel)
-
-            if not response or response[0]["successful"]:
+            try:
+                response = await super().subscribe(channel)
+            except Exception as e:
+                should_retry = await self.should_retry_on_exception(channel, e)
+                if not should_retry:
+                    raise
+            else:
+                if response and response[0]["successful"]:
+                    return response
+                else:
+                    should_retry = await self.should_retry_on_error_response(channel, response)
+
+            if not should_retry:
                 return response
 
-            # If not the known error, return
-            if not (
-                response[0]
-                .get("ext", {})
-                .get("sfdc", {})
-                .get("failureReason", "")
-                .startswith("SERVER_UNAVAILABLE")
-            ):
-                return response
             await asyncio.sleep(self.retry_sub_duration)
 
 

docs/api.rst

@@ -21,10 +21,8 @@ The package is separated in 5 mains modules:
 - To authenticate on Salesforce, you must use one connector that add
   authentication capability to :py:class:`BaseSalesforceStreaming`. See
   :ref:`connectors` section for a list of available connectors.
-- :ref:`mixins` extend :py:class:`BaseSalesforceStreaming` capabilities and
+- Finally, :ref:`mixins` extend :py:class:`BaseSalesforceStreaming` capabilities and
   can be added easily as opt-in option by sub classing.
-- Finally, :ref:`utils` provide some internal utilities functions.
-
 
 .. _main_interface:
 
@@ -132,7 +130,5 @@ Mixins
 .. autoclass:: AutoReconnectMixin
 
 .. autoclass:: ReSubscribeMixin
-
-
-.. _utils:
+   :members: should_retry_on_exception, should_retry_on_error_response
 

docs/release.rst

@@ -4,7 +4,8 @@ Release notes
 =============
 
 v. 0.3.0
-    Remove date-time parsing of the replay mixin.
+    - Remove date-time parsing of the replay mixin.
+    - Allow client to customize ``ReSubscribeMixin`` retry conditions.
 v. 0.2.0
     Minor release : Add refresh token authentication
 v. 0.1.1