72 lines
3.2 KiB
Python
72 lines
3.2 KiB
Python
import os
|
|
import asyncio
|
|
import uuid
|
|
import json
|
|
import aiosqlite
|
|
import redis.asyncio as aioredis
|
|
|
|
from app.core.config import RedisConfig, Env
|
|
from app.core.utils import now_ms
|
|
|
|
|
|
SQLITE_PATH = Env.SQLITE_PATH
|
|
REDIS_STREAM_DATABASE_PUBLISH = os.getenv("REDIS_STREAM_DATABASE_PUBLISH", "ACCOUNT:SERVICES:DATABASE:PUBLISH")
|
|
REDIS_STREAM_DATABASE_TASKS = os.getenv("REDIS_STREAM_DATABASE_TASKS", "ACCOUNT:SERVICES:DATABASE:TASKS")
|
|
REDIS_STREAM_MAIL_PUBLISH = os.getenv("REDIS_STREAM_MAIL_PUBLISH", "ACCOUNT:SERVICES:MAIL:PUBLISH")
|
|
REDIS_STREAM_MAIL_TASKS = os.getenv("REDIS_STREAM_MAIL_TASKS", "ACCOUNT:SERVICES:MAIL:TASKS")
|
|
REDIS_STREAM_QUEUE_PUBLISH = os.getenv("REDIS_STREAM_QUEUE_PUBLISH", "ACCOUNT:SERVICES:QUEUE:PUBLISH")
|
|
REDIS_STREAM_QUEUE_TASKS = os.getenv("REDIS_STREAM_QUEUE_TASKS", "ACCOUNT:SERVICES:QUEUE:TASKS")
|
|
|
|
|
|
async def ensure_schema(sqlite_path: str):
|
|
async with aiosqlite.connect(sqlite_path) as db:
|
|
await db.execute("""
|
|
CREATE TABLE IF NOT EXISTS tasks(
|
|
task_id TEXT PRIMARY KEY,
|
|
queue TEXT NOT NULL,
|
|
type TEXT NOT NULL,
|
|
payload_json TEXT NOT NULL,
|
|
created_at INTEGER NOT NULL,
|
|
status TEXT DEFAULT 'pending',
|
|
attempts INTEGER DEFAULT 0,
|
|
last_error TEXT
|
|
);
|
|
""")
|
|
await db.commit()
|
|
|
|
async def enqueue(r: aioredis.Redis, sqlite_path: str, stream: str, payload: dict, type_: str):
|
|
task_id = payload.get("task_id") or str(uuid.uuid4())
|
|
task = {"task_id": task_id, "queue": stream, "type": type_, "payload": payload, "created_at": now_ms(), "_attempts": 0}
|
|
await r.xadd(stream, {"data": json.dumps(task)})
|
|
async with aiosqlite.connect(sqlite_path) as db:
|
|
await db.execute("""INSERT OR REPLACE INTO tasks(task_id, queue, type, payload_json, created_at, status, attempts) VALUES(?,?,?,?,?,'pending',?)""",
|
|
(task_id, stream, type_, json.dumps(payload), task["created_at"], 0))
|
|
await db.commit()
|
|
|
|
async def push_db_mocks(r: aioredis.Redis, sqlite_path: str, n: int = 3):
|
|
for i in range(n):
|
|
payload = {"id": uuid.uuid4().hex, "op": "sync", "source": "tester"}
|
|
await enqueue(r, sqlite_path, REDIS_STREAM_DATABASE_TASKS, payload, "db-sync")
|
|
|
|
async def push_mail_mocks(r: aioredis.Redis, sqlite_path: str, n: int = 3):
|
|
for i in range(n):
|
|
payload = {"to": f"user{i}@example.com", "subj": "Hello", "body": "Hi!", "source": "tester"}
|
|
await enqueue(r, sqlite_path, REDIS_STREAM_MAIL_TASKS, payload, "send-mail")
|
|
|
|
async def push_queue_mocks(r: aioredis.Redis, sqlite_path: str, n: int = 3):
|
|
for i in range(n):
|
|
payload = {"action": "cleanup", "target": f"old-tasks-{i}", "source": "tester"}
|
|
await enqueue(r, sqlite_path, REDIS_STREAM_QUEUE_TASKS, payload, "queue-maintenance")
|
|
|
|
async def main():
|
|
db_n, mail_n, queue_n = 3, 3, 3
|
|
cfg = RedisConfig()
|
|
r = aioredis.Redis(host=cfg.host, port=cfg.port, db=cfg.db, username=cfg.username, password=cfg.password)
|
|
await ensure_schema(SQLITE_PATH)
|
|
await push_db_mocks(r, SQLITE_PATH, db_n)
|
|
await push_mail_mocks(r, SQLITE_PATH, mail_n)
|
|
await push_queue_mocks(r, SQLITE_PATH, queue_n)
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|