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

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

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

logger = logging.getLogger(__name__)


class LocalRecordStore(RecordInterface):
    def __init__(self):
        self._storage: dict[str, TaskRecord] = {}
        self._lock = threading.Lock()
        self._data_dir = settings.task_store_data_dir
        os.makedirs(self._data_dir, exist_ok=True)
        self._load_from_file()

    def _load_from_file(self):
        filepath = os.path.join(self._data_dir, "tasks.json")
        if os.path.isfile(filepath) and os.path.getsize(filepath) > 0:
            try:
                with open(filepath, "r", encoding="utf-8") as f:
                    data = json.load(f)
                for tid, rec in data.items():
                    self._storage[tid] = TaskRecord.from_dict(rec)
                logger.info("Loaded %d tasks from file system", len(data))
            except Exception as e:
                logger.error("Failed to load tasks from file: %s", e)

    def _save_to_file(self):
        try:
            tmp = os.path.join(self._data_dir, "tasks.tmp")
            filepath = os.path.join(self._data_dir, "tasks.json")
            with open(tmp, "w", encoding="utf-8") as f:
                json.dump(
                    {tid: r.to_dict() for tid, r in self._storage.items()},
                    f,
                    indent=2,
                    ensure_ascii=False,
                )
            os.replace(tmp, filepath)
        except Exception as e:
            logger.error("Failed to save tasks to file: %s", e)

    def _on_save(self, record: TaskRecord):
        with self._lock:
            self._storage[record.task_id] = record
            self._save_to_file()

    def _on_get(self, task_id: str) -> Optional[TaskRecord]:
        return self._storage.get(task_id)

    def _on_get_all(self) -> dict[str, TaskRecord]:
        return dict(self._storage)

    def _on_delete(self, max_age_hours: int):
        cutoff = int(time.time() * 1000) - max_age_hours * 3600 * 1000
        with self._lock:
            to_delete = [
                tid
                for tid, r in self._storage.items()
                if r.status in ("successes", "failed") and r.update_time < cutoff
            ]
            for tid in to_delete:
                del self._storage[tid]
            self._save_to_file()

    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)