feat: enhance handle_autoreconnect_async to support async generator functions

Author: DylanWelzel

biothings/utils/mongo_async.py

@@ -79,6 +79,23 @@ def handle_autoreconnect_async(func):
     MAX_RETRY = 30
     SLEEP_TIME = 0.5  # seconds
 
+    if inspect.isasyncgenfunction(func):
+
+        @wraps(func)
+        async def inner_gen(*args, **kwargs):
+            retry = 0
+            while retry < MAX_RETRY:
+                try:
+                    async for item in func(*args, **kwargs):
+                        yield item
+                    return
+                except AutoReconnect:
+                    retry += 1
+                    await asyncio.sleep(SLEEP_TIME)
+            raise MaxRetryAutoReconnectException()
+
+        return inner_gen
+
     @wraps(func)
     async def inner(*args, **kwargs):
         retry = 0