79 lines
2.7 KiB
Python
79 lines
2.7 KiB
Python
import os
|
|
import json
|
|
import logging
|
|
from typing import Optional, Dict, Any
|
|
import redis.asyncio as redis
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
REDIS_URL = os.getenv("REDIS_URL", "redis://copykar_redis:6379/0" if os.path.exists("/app") else "redis://localhost:6379/0")
|
|
|
|
class RedisQueue:
|
|
def __init__(self, redis_url: Optional[str] = None):
|
|
self.redis_url = redis_url or REDIS_URL
|
|
self.client: Optional[redis.Redis] = None
|
|
|
|
async def connect(self):
|
|
if not self.client:
|
|
self.client = redis.from_url(self.redis_url, decode_responses=True)
|
|
await self.client.ping()
|
|
logger.info(f"Connected to Redis at {self.redis_url}")
|
|
|
|
def _get_target_key(self, target_id: int) -> str:
|
|
return f"copykar:queue:target:{target_id}"
|
|
|
|
async def push_target_post(self, target_id: int, payload: Dict[str, Any]):
|
|
if not self.client:
|
|
await self.connect()
|
|
key = self._get_target_key(target_id)
|
|
raw_json = json.dumps(payload)
|
|
await self.client.rpush(key, raw_json)
|
|
logger.info(f"Enqueued post {payload.get('post_id')} to Target #{target_id} queue [{key}]")
|
|
|
|
async def pop_target_post(self, target_id: int, dispatch_order: str = "order") -> Optional[Dict[str, Any]]:
|
|
if not self.client:
|
|
await self.connect()
|
|
key = self._get_target_key(target_id)
|
|
|
|
if dispatch_order == "random":
|
|
import random
|
|
length = await self.client.llen(key)
|
|
if length == 0:
|
|
return None
|
|
if length == 1:
|
|
raw = await self.client.lpop(key)
|
|
else:
|
|
idx = random.randint(0, length - 1)
|
|
raw = await self.client.lindex(key, idx)
|
|
if raw:
|
|
await self.client.lrem(key, 1, raw)
|
|
else:
|
|
raw = await self.client.lpop(key)
|
|
|
|
if raw:
|
|
try:
|
|
return json.loads(raw)
|
|
except Exception as e:
|
|
logger.error(f"Error parsing queue JSON from {key}: {e}")
|
|
return None
|
|
|
|
async def get_target_queue_size(self, target_id: int) -> int:
|
|
if not self.client:
|
|
await self.connect()
|
|
key = self._get_target_key(target_id)
|
|
return await self.client.llen(key)
|
|
|
|
async def get_total_queued_posts(self) -> int:
|
|
if not self.client:
|
|
await self.connect()
|
|
keys = await self.client.keys("copykar:queue:target:*")
|
|
total = 0
|
|
for k in keys:
|
|
total += await self.client.llen(k)
|
|
return total
|
|
|
|
async def close(self):
|
|
if self.client:
|
|
await self.client.aclose()
|
|
logger.info("Redis connection closed.")
|