114 lines
3.2 KiB
Python
114 lines
3.2 KiB
Python
async def flush():
|
|
import orjson
|
|
|
|
import asyncpg.exceptions
|
|
|
|
import redis.asyncio as aredis
|
|
|
|
from payphone.db import models
|
|
|
|
from payphone.core import constants, utils, config
|
|
|
|
redis = aredis.Redis(host="localhost" if config.IS_DEBUG else "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()
|