Create db.py

db.py

@@ -0,0 +1,76 @@
+import redis.asyncio as aioredis
+import os
+
+_redis = None
+
+async def get_redis():
+    global _redis
+    if _redis is None:
+        _redis = aioredis.from_url(os.environ["REDIS_URL"], decode_responses=True)
+    return _redis
+
+async def queue_size() -> int:
+    r = await get_redis()
+    return await r.llen("frag:queue")
+
+async def done_count() -> int:
+    r = await get_redis()
+    return await r.scard("frag:done")
+
+async def get_state() -> dict:
+    r = await get_redis()
+    state = await r.hgetall("frag:state")
+    if not state:
+        state = {"running": "0", "paused": "0", "total": "0", "processed": "0"}
+        await r.hset("frag:state", mapping=state)
+    return state
+
+async def set_state(**kwargs):
+    r = await get_redis()
+    string_kwargs = {k: str(v) for k, v in kwargs.items()}
+    await r.hset("frag:state", mapping=string_kwargs)
+
+async def pop_word() -> str|None:
+    r = await get_redis()
+    return await r.rpop("frag:queue")
+
+async def mark_done(word: str, status: str):
+    r = await get_redis()
+    await r.set(f"frag:status:{word}", status)
+    await r.sadd("frag:done", word)
+    if status in ("TAKEN", "UNAVAILABLE", "FOR_SALE", "ON_AUCTION"):
+        set_name = status.lower().replace("_", "")
+        await r.sadd(f"frag:{set_name}", word)
+    await r.hincrby("frag:state", "processed", 1)
+
+async def get_counts() -> dict:
+    r = await get_redis()
+    return {
+        "taken": await r.scard("frag:taken"),
+        "unavailable": await r.scard("frag:unavailable"),
+        "forsale": await r.scard("frag:forsale"),
+        "auction": await r.scard("frag:onauction")
+    }
+
+async def get_all_taken() -> list:
+    r = await get_redis()
+    return list(await r.smembers("frag:taken"))
+
+async def get_all_unavailable() -> list:
+    r = await get_redis()
+    return list(await r.smembers("frag:unavailable"))
+
+async def flush_all():
+    r = await get_redis()
+    keys = await r.keys("frag:*")
+    if keys:
+        await r.delete(*keys)
+
+async def get_concurrency() -> int:
+    r = await get_redis()
+    val = await r.get("frag:concurrency")
+    return int(val) if val else 30
+
+async def set_concurrency(n: int):
+    r = await get_redis()
+    await r.set("frag:concurrency", str(n))