235 lines
7.2 KiB
Python
235 lines
7.2 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
import json
|
|
from typing import Any, Protocol
|
|
|
|
import redis.asyncio as redis
|
|
from redis.exceptions import ResponseError
|
|
|
|
from app.core.config import settings
|
|
from app.core.logging import get_logger
|
|
|
|
|
|
logger = get_logger(__name__, service="earth_news")
|
|
|
|
TARGET_LOCATION_STREAM = "earth_news:target_location:jobs"
|
|
TARGET_LOCATION_GROUP = "earth_news_target_location"
|
|
TARGET_LOCATION_DEAD_LETTER_STREAM = "earth_news:target_location:dead"
|
|
TARGET_LOCATION_RESULT_TTL_SECONDS = 60 * 60 * 12
|
|
TARGET_LOCATION_JOB_DEDUP_TTL_SECONDS = 60 * 60 * 6
|
|
TARGET_LOCATION_MAX_ATTEMPTS = 3
|
|
|
|
_redis_client: redis.Redis | None = None
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class NewsTargetLocationMessage:
|
|
message_id: str
|
|
item_id: str
|
|
payload: dict[str, Any]
|
|
attempts: int = 0
|
|
|
|
|
|
class NewsTargetLocationQueue(Protocol):
|
|
async def enqueue(self, *, item_id: str, payload: dict[str, Any], force: bool = False) -> bool:
|
|
...
|
|
|
|
async def consume_batch(
|
|
self,
|
|
*,
|
|
consumer_name: str,
|
|
count: int,
|
|
block_ms: int,
|
|
) -> list[NewsTargetLocationMessage]:
|
|
...
|
|
|
|
async def ack(self, message_id: str) -> None:
|
|
...
|
|
|
|
async def retry_or_dead_letter(
|
|
self,
|
|
message: NewsTargetLocationMessage,
|
|
*,
|
|
error: str,
|
|
) -> None:
|
|
...
|
|
|
|
|
|
def _get_redis_client() -> redis.Redis:
|
|
global _redis_client
|
|
if _redis_client is None:
|
|
_redis_client = redis.from_url(settings.REDIS_URL, decode_responses=True)
|
|
return _redis_client
|
|
|
|
|
|
def _result_key(item_id: str) -> str:
|
|
return f"earth_news:target_location:result:{item_id}"
|
|
|
|
|
|
def _queued_key(item_id: str) -> str:
|
|
return f"earth_news:target_location:queued:{item_id}"
|
|
|
|
|
|
class RedisStreamsNewsTargetLocationQueue:
|
|
def __init__(self, client: redis.Redis | None = None) -> None:
|
|
self.client = client or _get_redis_client()
|
|
self._group_ready = False
|
|
|
|
async def _ensure_group(self) -> None:
|
|
if self._group_ready:
|
|
return
|
|
try:
|
|
await self.client.xgroup_create(
|
|
TARGET_LOCATION_STREAM,
|
|
TARGET_LOCATION_GROUP,
|
|
id="0",
|
|
mkstream=True,
|
|
)
|
|
except ResponseError as exc:
|
|
if "BUSYGROUP" not in str(exc):
|
|
raise
|
|
self._group_ready = True
|
|
|
|
async def enqueue(self, *, item_id: str, payload: dict[str, Any], force: bool = False) -> bool:
|
|
await self._ensure_group()
|
|
if force:
|
|
await self.client.delete(_result_key(item_id), _queued_key(item_id))
|
|
elif await self.client.exists(_result_key(item_id)):
|
|
return False
|
|
queued = await self.client.set(
|
|
_queued_key(item_id),
|
|
"1",
|
|
nx=True,
|
|
ex=TARGET_LOCATION_JOB_DEDUP_TTL_SECONDS,
|
|
)
|
|
if not queued:
|
|
return bool(await self.client.exists(_queued_key(item_id)))
|
|
await self.client.xadd(
|
|
TARGET_LOCATION_STREAM,
|
|
{
|
|
"item_id": item_id,
|
|
"attempts": "0",
|
|
"payload": json.dumps(payload, ensure_ascii=False),
|
|
},
|
|
)
|
|
return True
|
|
|
|
async def consume_batch(
|
|
self,
|
|
*,
|
|
consumer_name: str,
|
|
count: int,
|
|
block_ms: int,
|
|
) -> list[NewsTargetLocationMessage]:
|
|
await self._ensure_group()
|
|
streams = await self.client.xreadgroup(
|
|
TARGET_LOCATION_GROUP,
|
|
consumer_name,
|
|
{TARGET_LOCATION_STREAM: ">"},
|
|
count=count,
|
|
block=block_ms,
|
|
)
|
|
messages: list[NewsTargetLocationMessage] = []
|
|
for _stream_name, stream_messages in streams:
|
|
for message_id, fields in stream_messages:
|
|
raw_payload = fields.get("payload")
|
|
item_id = fields.get("item_id")
|
|
if not raw_payload or not item_id:
|
|
await self.ack(message_id)
|
|
continue
|
|
try:
|
|
payload = json.loads(raw_payload)
|
|
except json.JSONDecodeError:
|
|
await self.ack(message_id)
|
|
continue
|
|
attempts = int(fields.get("attempts") or 0)
|
|
messages.append(
|
|
NewsTargetLocationMessage(
|
|
message_id=message_id,
|
|
item_id=item_id,
|
|
payload=payload,
|
|
attempts=attempts,
|
|
)
|
|
)
|
|
return messages
|
|
|
|
async def ack(self, message_id: str) -> None:
|
|
await self.client.xack(TARGET_LOCATION_STREAM, TARGET_LOCATION_GROUP, message_id)
|
|
|
|
async def retry_or_dead_letter(
|
|
self,
|
|
message: NewsTargetLocationMessage,
|
|
*,
|
|
error: str,
|
|
) -> None:
|
|
await self.ack(message.message_id)
|
|
if message.attempts + 1 >= TARGET_LOCATION_MAX_ATTEMPTS:
|
|
await self.client.xadd(
|
|
TARGET_LOCATION_DEAD_LETTER_STREAM,
|
|
{
|
|
"item_id": message.item_id,
|
|
"attempts": str(message.attempts + 1),
|
|
"error": error,
|
|
"payload": json.dumps(message.payload, ensure_ascii=False),
|
|
},
|
|
)
|
|
return
|
|
await self.client.xadd(
|
|
TARGET_LOCATION_STREAM,
|
|
{
|
|
"item_id": message.item_id,
|
|
"attempts": str(message.attempts + 1),
|
|
"payload": json.dumps(message.payload, ensure_ascii=False),
|
|
},
|
|
)
|
|
|
|
|
|
def get_news_target_location_queue() -> NewsTargetLocationQueue:
|
|
return RedisStreamsNewsTargetLocationQueue()
|
|
|
|
|
|
async def enqueue_target_location_job(payload: dict[str, Any], *, force: bool = False) -> bool:
|
|
item_id = str(payload.get("id") or "")
|
|
if not item_id:
|
|
return False
|
|
try:
|
|
queue = get_news_target_location_queue()
|
|
return await queue.enqueue(item_id=item_id, payload=payload, force=force)
|
|
except Exception as exc:
|
|
logger.warning_event(
|
|
"Failed to enqueue Earth news target location job",
|
|
event="earth_news.target_location.enqueue_failed",
|
|
context={"item_id": item_id, "error": str(exc)},
|
|
)
|
|
return False
|
|
|
|
|
|
async def get_cached_target_location_patch(item_id: str) -> dict[str, Any] | None:
|
|
try:
|
|
raw_value = await _get_redis_client().get(_result_key(item_id))
|
|
except Exception as exc:
|
|
logger.warning_event(
|
|
"Failed to read Earth news target location cache",
|
|
event="earth_news.target_location.cache_read_failed",
|
|
context={"item_id": item_id, "error": str(exc)},
|
|
)
|
|
return None
|
|
if not raw_value:
|
|
return None
|
|
try:
|
|
value = json.loads(raw_value)
|
|
except json.JSONDecodeError:
|
|
return None
|
|
return value if isinstance(value, dict) else None
|
|
|
|
|
|
async def save_target_location_patch(item_id: str, patch: dict[str, Any]) -> None:
|
|
client = _get_redis_client()
|
|
await client.setex(
|
|
_result_key(item_id),
|
|
TARGET_LOCATION_RESULT_TTL_SECONDS,
|
|
json.dumps(patch, ensure_ascii=False),
|
|
)
|
|
await client.delete(_queued_key(item_id))
|