ai-flow-python/app/repository/redis_store.py

import json
import logging
import time
from typing import Any, Optional

import redis

from app.config import settings
from app.repository.record_interface import RecordInterface
from app.repository.task_record import TaskRecord

logger = logging.getLogger(__name__)

TASK_KEY_PREFIX = "task:"


class RedisRecordStore(RecordInterface):
    def __init__(self):
        self._client = redis.Redis(
            host=settings.redis_host,
            port=settings.redis_port,
            password=settings.redis_password,
            db=settings.redis_db,
            socket_timeout=settings.redis_timeout / 1000,
            decode_responses=True,
        )

    def _on_save(self, record: TaskRecord):
        try:
            self._client.set(
                TASK_KEY_PREFIX + record.task_id, json.dumps(record.to_dict())
            )
        except Exception as e:
            logger.error("Failed to save task record: %s", e)

    def _on_get(self, task_id: str) -> Optional[TaskRecord]:
        try:
            data = self._client.get(TASK_KEY_PREFIX + task_id)
            if data:
                return TaskRecord.from_dict(json.loads(data))
            return None
        except Exception as e:
            logger.error("Failed to get task record: %s", e)
            return None

    def _on_get_all(self) -> dict[str, TaskRecord]:
        result: dict[str, TaskRecord] = {}
        try:
            keys = self._client.keys(TASK_KEY_PREFIX + "*")
            for key in keys or []:
                task_id = key[len(TASK_KEY_PREFIX):]
                data = self._client.get(key)
                if data:
                    result[task_id] = TaskRecord.from_dict(json.loads(data))
        except Exception as e:
            logger.error("Failed to get all tasks: %s", e)
        return result

    def _on_delete(self, max_age_hours: int):
        try:
            cutoff = int(time.time() * 1000) - max_age_hours * 3600 * 1000
            all_records = self._on_get_all()
            for task_id, record in all_records.items():
                if record.status in ("successes", "failed") and record.update_time < cutoff:
                    self._client.delete(TASK_KEY_PREFIX + task_id)
        except Exception as e:
            logger.error("Failed to cleanup completed tasks: %s", e)

    def _on_update_status(
        self, task_id: str, status: str, message: str, data: Optional[dict]
    ):
        try:
            record = self._on_get(task_id)
            if record:
                record.status = status
                record.message = message
                record.touch()
                if data is not None:
                    record.data = data
                self._on_save(record)
        except Exception as e:
            logger.error("Failed to update task status: %s", e)

    def _on_publish(self, task_id: str):
        if not settings.redis_publish_enabled:
            return
        try:
            record = self._on_get(task_id)
            if record:
                payload = json.dumps(record.to_dict(), ensure_ascii=False)
                receivers = self._client.publish(settings.redis_publish_channel, payload)
                logger.info("[%s] Published to channel '%s', receivers=%d",
                            task_id, settings.redis_publish_channel, receivers)
        except Exception as e:
            logger.error("[%s] Failed to publish task event: %s", task_id, e)