PMFlow/backend/app/services/context_cache_service/invalidation.py

from __future__ import annotations

from .shared import *

class ContextCacheInvalidator:
    def __init__(self, redis_url: Optional[str] = None) -> None:
        self.redis_url = redis_url or os.getenv("REDIS_URL")
        self._redis: Optional[Any] = None

    def _get_redis(self):
        if not REDIS_AVAILABLE or not self.redis_url or RedisSync is None:
            return None
        if self._redis is None:
            self._redis = RedisSync.from_url(
                self.redis_url,
                decode_responses=True,
                socket_connect_timeout=1,
                socket_timeout=1,
            )
        return self._redis

    @staticmethod
    def _system_key(tenant_id: int) -> str:
        return f"pmflow:context:system:{_safe_int(tenant_id, 1)}"

    @staticmethod
    def _project_list_key(tenant_id: int) -> str:
        return f"pmflow:project:list:{_safe_int(tenant_id, 1)}"

    @staticmethod
    def _signal_pending_key(tenant_id: int) -> str:
        return f"pmflow:signal:pending_count:{_safe_int(tenant_id, 1)}"

    @staticmethod
    def _user_key(tenant_id: int, user_id: int) -> str:
        return f"pmflow:context:user:{_safe_int(tenant_id, 1)}:{_safe_int(user_id, 0)}"

    def invalidate(
        self,
        *,
        tenant_id: int,
        invalidate_system: bool,
        invalidate_project_list: bool,
        invalidate_signal_pending: bool,
        user_ids: Optional[set[int]] = None,
    ) -> None:
        redis = self._get_redis()
        if redis is None:
            return

        tid = _safe_int(tenant_id, 1)
        keys: list[str] = []
        if invalidate_system:
            keys.append(self._system_key(tid))
        if invalidate_project_list:
            keys.append(self._project_list_key(tid))
        if invalidate_signal_pending:
            keys.append(self._signal_pending_key(tid))
        if user_ids:
            for uid in user_ids:
                if _safe_int(uid, 0) > 0:
                    keys.append(self._user_key(tid, uid))
        if not keys:
            return
        try:
            redis.delete(*keys)
        except Exception:
            return

def register_context_cache_invalidation_hooks() -> None:
    if getattr(register_context_cache_invalidation_hooks, "_registered", False):
        return

    invalidator = ContextCacheInvalidator()

    @event.listens_for(Session, "after_flush")
    def _collect_cache_invalidation(session: Session, _flush_context) -> None:  # pragma: no cover
        payload = session.info.get("_pmflow_cache_invalidation")
        if not isinstance(payload, dict):
            payload = {
                "tenant_ids": set(),
                "invalidate_system": False,
                "invalidate_project_list": False,
                "invalidate_signal_pending": False,
                "user_ids": set(),
            }
            session.info["_pmflow_cache_invalidation"] = payload

        objects = set(session.new).union(session.dirty).union(session.deleted)
        for obj in objects:
            tenant_id = _safe_int(getattr(obj, "tenant_id", 0), 0)
            if tenant_id > 0:
                payload["tenant_ids"].add(tenant_id)

            if isinstance(obj, Project):
                payload["invalidate_system"] = True
                payload["invalidate_project_list"] = True
                continue

            if isinstance(obj, ProjectSignalLink):
                payload["invalidate_system"] = True
                payload["invalidate_project_list"] = True
                continue

            if isinstance(obj, Signal):
                payload["invalidate_system"] = True
                payload["invalidate_signal_pending"] = True
                continue

            if isinstance(obj, Task):
                payload["invalidate_system"] = True
                state = inspect(obj)
                assignee_hist = state.attrs.assignee_id.history if hasattr(state.attrs, "assignee_id") else None
                if assignee_hist is not None:
                    for old_value in assignee_hist.deleted:
                        uid = _safe_int(old_value, 0)
                        if uid > 0:
                            payload["user_ids"].add(uid)
                    for new_value in assignee_hist.added:
                        uid = _safe_int(new_value, 0)
                        if uid > 0:
                            payload["user_ids"].add(uid)
                uid = _safe_int(getattr(obj, "assignee_id", 0), 0)
                if uid > 0:
                    payload["user_ids"].add(uid)
                continue

            if isinstance(obj, UserProject):
                payload["invalidate_system"] = True
                payload["invalidate_project_list"] = True
                uid = _safe_int(getattr(obj, "user_id", 0), 0)
                if uid > 0:
                    payload["user_ids"].add(uid)
                continue

            if isinstance(obj, User):
                payload["invalidate_system"] = True
                uid = _safe_int(getattr(obj, "id", 0), 0)
                if uid > 0:
                    payload["user_ids"].add(uid)
                continue

            if isinstance(obj, UserWorkload):
                payload["invalidate_system"] = True
                uid = _safe_int(getattr(obj, "user_id", 0), 0)
                if uid > 0:
                    payload["user_ids"].add(uid)
                continue

    @event.listens_for(Session, "after_commit")
    def _flush_cache_invalidation(session: Session) -> None:  # pragma: no cover
        payload = session.info.pop("_pmflow_cache_invalidation", None)
        if not isinstance(payload, dict):
            return
        tenant_ids = payload.get("tenant_ids") or set()
        if not tenant_ids:
            tenant_ids = {_safe_int(os.getenv("DEFAULT_TENANT_ID", "1"), 1)}
        for tenant_id in tenant_ids:
            invalidator.invalidate(
                tenant_id=_safe_int(tenant_id, 1),
                invalidate_system=bool(payload.get("invalidate_system")),
                invalidate_project_list=bool(payload.get("invalidate_project_list")),
                invalidate_signal_pending=bool(payload.get("invalidate_signal_pending")),
                user_ids=set(payload.get("user_ids") or set()),
            )

    @event.listens_for(Session, "after_rollback")
    def _clear_cache_invalidation(session: Session) -> None:  # pragma: no cover
        session.info.pop("_pmflow_cache_invalidation", None)

    setattr(register_context_cache_invalidation_hooks, "_registered", True)