add pubsub reconnect hook and auth failure hook

Author: IAmTomahawkx

docs/changelog.rst

@@ -44,6 +44,11 @@ Master
     - Bug fixes
         - Fix forced RECONNECT messages
 
+    - Additions
+        - Added proper message when wrong type is passed to a topic argument
+        - Added auth failure hook: :func:`~twitchio.ext.pubsub.PubSubPool.auth_fail_hook`
+        - Added reconnect hook: :func:`~twitchio.ext.pubsub.PubSubPool.reconnect_hook`
+
 2.5.0
 ======
 - TwitchIO

docs/exts/pubsub.rst

@@ -212,6 +212,32 @@ If the topic requires multiple channel ids, they should be passed as such:
         listen_to_id = 12345
         topic = pubsub.whispers(user_token)[listen_to_id]
 
+Hooks
+------
+
+There are two hooks available in the PubSubPool class. To access these hooks, subclass the PubSubPool.
+After subclassing, use the subclass like normal.
+
+The ``auth_fail_hook`` is called whenever you attempt to subscribe to a topic and the auth token is invalid.
+From the hook, you are able to fix your token (maybe you need to prompt the user for a new token), and then subscribe again.
+
+The ``reconnect_hook`` is called whenevever a node has to reconnect to twitch, for any reason. The node will wait for you to
+return a list of topics before reconnecting. Any modifications to the topics will be applied to the node.
+
+.. code-block:: python3
+
+    from typing import List
+    from twitchio.ext import pubsub
+
+    class MyPool(pubsub.PubSubPool):
+        async def auth_fail_hook(self, topics: List[pubsub.Topic]) -> None:
+            fixed_topics = fix_my_auth_tokens(topics) # somehow fix your auth tokens
+            await self.subscribe_topics(topics)
+
+        async def reconnect_hook(self, node: pubsub.PubSubWebsocket, topics: List[pubsub.Topic]) -> List[pubsub.Topic]:
+            return topics
+
+
 Api Reference
 --------------
 

twitchio/ext/pubsub/pool.py

@@ -21,8 +21,9 @@
 FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 DEALINGS IN THE SOFTWARE.
 """
-
+import copy
 import itertools
+import logging
 from typing import List, Optional
 
 from twitchio import Client
@@ -33,6 +34,7 @@
 
 __all__ = ("PubSubPool",)
 
+logger = logging.getLogger("twitchio.ext.eventsub.pool")
 
 class PubSubPool:
     """
@@ -64,7 +66,7 @@ async def subscribe_topics(self, topics: List[Topic]):
         """
         node = self._find_node(topics)
         if node is None:
-            node = PubSubWebsocket(self.client, max_topics=self._max_connection_topics)
+            node = PubSubWebsocket(self.client, pool=self, max_topics=self._max_connection_topics)
             await node.connect()
             self._pool.append(node)
 
@@ -87,6 +89,73 @@ async def unsubscribe_topics(self, topics: List[Topic]):
                 await node.disconnect()
                 self._pool.remove(node)
 
+    async def _process_auth_fail(self, nonce: str, node: PubSubWebsocket) -> None:
+        topics = [topic for topic in self._topics if topic._nonce == nonce]
+
+        for topic in topics:
+            topic._nonce = None
+            del self._topics[topic]
+            node.topics.remove(topic)
+
+        try:
+            await self.auth_fail_hook(topics)
+        except Exception as e:
+            logger.error("Error occurred while calling auth_fail_hook.", exc_info=e)
+
+
+    async def auth_fail_hook(self, topics: list[Topic]):
+        """
+        This is a hook that can be overridden in a subclass.
+
+
+        Parameters
+        ----------
+        node
+        topics: List[:class:`Topic`]
+            The topcs that this node has.
+
+        Returns
+        -------
+        List[:class:`Topic`]
+            The list of topics this node should have. Any additions, modifications, or removals will be respected.
+        """
+
+    async def _process_reconnect_hook(self, node: PubSubWebsocket) -> None:
+        topics = copy.copy(node.topics)
+
+        for topic in topics:
+            self._topics.pop(topic, None)
+
+        try:
+            new_topics = await self.reconnect_hook(node, topics)
+        except Exception as e:
+            new_topics = node.topics
+            logger.error("Error occurred while calling reconnect_hook.", exc_info=e)
+
+        for topic in new_topics:
+            self._topics[topic] = node
+
+        node.topics = new_topics
+
+    async def reconnect_hook(self, node: PubSubWebsocket, topics: List[Topic]) -> List[Topic]:
+        """
+        This is a low-level hook that can be overridden in a subclass.
+        it is called whenever a node has to reconnect for any reason, from the twitch edge lagging out to being told to by twitch.
+        This hook allows you to modify the topics, potentially updating tokens or removing topics altogether.
+
+        Parameters
+        ----------
+        node
+        topics: List[:class:`Topic`]
+            The topcs that this node has.
+
+        Returns
+        -------
+        List[:class:`Topic`]
+            The list of topics this node should have. Any additions, modifications, or removals will be respected.
+        """
+        return topics
+
     def _find_node(self, topics: List[Topic]) -> Optional[PubSubWebsocket]:
         if self._mode != "group":
             raise ValueError("group is the only supported mode.")

twitchio/ext/pubsub/topics.py

@@ -21,7 +21,7 @@
 FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 DEALINGS IN THE SOFTWARE.
 """
-
+import uuid
 from typing import Optional, List, Type
 
 
@@ -65,16 +65,17 @@ class Topic(_topic):
         The arguments to substitute in to the topic string
     """
 
-    __slots__ = "token", "args"
+    __slots__ = "token", "args", "_nonce"
 
     def __init__(self, topic, args):
         super().__init__(topic, args)
         self.token = None
+        self._nonce = None
         self.args = []
 
     def __getitem__(self, item):
         assert len(self.args) < len(self.__args__), ValueError("Too many arguments")
-        assert isinstance(item, self.__args__[len(self.args)])  # noqa
+        assert isinstance(item, self.__args__[len(self.args)]), ValueError(f"Got {item!r}, excepted {self.__args__[len(self.args)]}")  # noqa
         self.args.append(item)
         return self
 
@@ -89,12 +90,19 @@ def present(self) -> Optional[str]:
         except:
             return None
 
+    def _present_set_nonce(self, nonce: str) -> Optional[str]:
+        self._nonce = nonce
+        return self.present
+
     def __eq__(self, other):
         return other is self or (isinstance(other, Topic) and other.present == self.present)
 
     def __hash__(self):
         return hash(self.present)
 
+    def __repr__(self):
+        return f"<Topic {self.__topic__} args={self.args}>"
+
 
 bits = _topic("channel-bits-events-v2.{0}", [int])
 bits_badge = _topic("channel-bits-badge-unlocks.{0}", [int])

twitchio/ext/pubsub/websocket.py

@@ -21,20 +21,24 @@
 FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
 DEALINGS IN THE SOFTWARE.
 """
+from __future__ import annotations
 
 import asyncio
 import logging
 import time
 import uuid
 from itertools import groupby
-from typing import Optional, List
+from typing import Optional, List, TYPE_CHECKING
 
 import aiohttp
 
 from twitchio import Client
 from .topics import Topic
 from . import models
 
+if TYPE_CHECKING:
+    from .pool import PubSubPool
+
 try:
     import ujson as json
 except:
@@ -51,6 +55,7 @@ class PubSubWebsocket:
     __slots__ = (
         "session",
         "topics",
+        "pool",
         "client",
         "connection",
         "_latency",
@@ -63,11 +68,12 @@ class PubSubWebsocket:
 
     ENDPOINT = "wss://pubsub-edge.twitch.tv"
 
-    def __init__(self, client: Client, *, max_topics=50):
+    def __init__(self, client: Client, pool: PubSubPool, *, max_topics=50):
         self.max_topics = max_topics
         self.session = None
         self.connection: Optional[aiohttp.ClientWebSocketResponse] = None
         self.topics: List[Topic] = []
+        self.pool = pool
         self.client = client
         self._latency = None
         self._closing = False
@@ -110,6 +116,7 @@ async def disconnect(self):
 
     async def reconnect(self):
         await self.disconnect()
+        await self.pool._process_reconnect_hook(self)
         await self.connect()
 
     async def _send_initial_topics(self):
@@ -118,10 +125,11 @@ async def _send_initial_topics(self):
     async def _send_topics(self, topics: List[Topic], type="LISTEN"):
         for tok, _topics in groupby(topics, key=lambda val: val.token):
             nonce = ("%032x" % uuid.uuid4().int)[:8]
+
             payload = {
                 "type": type,
                 "nonce": nonce,
-                "data": {"topics": [x.present for x in _topics], "auth_token": tok},
+                "data": {"topics": [x._present_set_nonce(nonce) for x in _topics], "auth_token": tok},
             }
             logger.debug(f"Sending {type} payload with nonce '{nonce}': {payload}")
             await self.send(payload)
@@ -150,7 +158,7 @@ async def poll(self):
 
             handle = getattr(self, "handle_" + data["type"].lower().replace("-", "_"), None)
             if handle:
-                self.client.loop.create_task(handle(data))
+                self.client.loop.create_task(handle(data), name=f"pubsub-handle-event: {data['type']}")
             else:
                 print(data)
                 logger.debug(f"Pubsub event referencing unknown event '{data['type']}'. Discarding")
@@ -197,6 +205,10 @@ async def handle_response(self, message: dict):
         if message["error"]:
             logger.error(f"Received errored response for nonce {message['nonce']}: {message['error']}")
             self.client.run_event("pubsub_error", message)
+            if message["error"] == "ERR_BADAUTH":
+                nonce = message["nonce"]
+                await self.pool._process_auth_fail(nonce, self)
+
         elif message["type"] == "RECONNECT":
             logger.warning("Received RECONNECT response from pubsub edge. Reconnecting")
             await asyncio.shield(self.reconnect())