feat(bot): ✨ Move messages flushing to a separate loop.

2024-10-23 21:48:48 -03:00
parent 00404973bf
commit 2372ee3d1f
7 changed files with 149 additions and 99 deletions
@@ -0,0 +1,110 @@
+async def flush():
+    import orjson
+
+    import asyncpg.exceptions
+
+    from payphone.db import models
+
+    from payphone.core import constants, utils
+    from payphone.core.redis import redis
+
+    pipe = redis.pipeline()
+
+    pipe.smembers(constants.CALL_IDS_TO_FLUSH)
+    pipe.delete(constants.CALL_IDS_TO_FLUSH)
+
+    call_ids, _ = await pipe.execute()
+
+    # Redis returns bytes.
+    call_ids = [int(call_id) for call_id in call_ids]
+
+    if len(call_ids) == 0:
+        print("Skipping flushing. No calls to flush.")
+        return
+
+    print(f"Flushing messages for {len(call_ids)} calls.")
+
+    pipe = redis.pipeline()
+
+    for call_id in call_ids:
+        pipe.hgetall(constants.CALL_MESSAGES_TEMPLATE % call_id)
+        pipe.hgetall(constants.CALL_MESSAGE_CONTENTS_TEMPLATE % call_id)
+
+    for call_id, (messages, message_contents) in zip(
+        call_ids, utils.chunks(await pipe.execute(), 2)
+    ):
+        message_rows: list[models.Message] = []
+
+        for message_json in messages.values():
+            message_rows.append(models.Message(**orjson.loads(message_json)))
+
+        message_content_rows: list[models.MessageContent] = []
+
+        for message_content_json in message_contents.values():
+            message_content_rows.append(
+                models.MessageContent(**orjson.loads(message_content_json))
+            )
+
+        message_rows.sort(key=lambda m: m.id)
+        message_content_rows.sort(key=lambda mc: mc.message_id)
+
+        pipe = redis.pipeline()
+
+        pipe.delete(
+            constants.CALL_MESSAGE_CONTENT_ITERATIONS_TEMPLATE % call_id,
+            constants.CALL_MESSAGE_CONTENTS_TEMPLATE % call_id,
+            constants.CALL_MESSAGE_FROM_PROPAGATION_TEMPLATE % call_id,
+            constants.CALL_MESSAGES_TEMPLATE % call_id,
+            constants.CALL_MESSAGE_REPLIES_TEMPLATE % call_id,
+        )
+
+        await pipe.execute()
+
+        for _ in range(5):
+            try:
+                await models.Message.bulk_create(message_rows, batch_size=20)
+                break
+
+            except (
+                asyncpg.exceptions.PostgresError,
+                asyncpg.exceptions.InterfaceError,
+            ) as e:
+                print(
+                    f"[Call ID: {call_id}] An error occured trying to flush messages for call ID. {e}"
+                )
+                continue
+
+        for _ in range(5):
+            try:
+                await models.MessageContent.bulk_create(
+                    message_content_rows, batch_size=20
+                )
+                break
+
+            except (
+                asyncpg.exceptions.PostgresError,
+                asyncpg.exceptions.InterfaceError,
+            ) as e:
+                print(
+                    f"[Call ID: {call_id}] An error occured trying to flush message contents. {e}"
+                )
+                continue
+
+        print(f"[Call ID: {call_id}] Flushed {len(messages)} messages and {len(message_contents)} message contents.")
+
+    print("Finished flushing messages.")
+
+
+async def loop() -> None:
+    """Loops in the background flushing messages to the DB."""
+
+    import asyncio
+
+    from payphone.db import setup
+
+    await setup.init()
+
+    while True:
+        await asyncio.sleep(30)
+
+        await flush()