1from __future__ import annotations
  2
  3from datetime import datetime, timedelta
  4from functools import cached_property
  5from typing import Any
  6
  7import psycopg
  8
  9from plain.utils import timezone
 10
 11
 12class Cached:
 13    """Store and retrieve cached items."""
 14
 15    def __init__(self, key: str) -> None:
 16        self.key = key
 17
 18        # So we can import Cached in __init__.py
 19        # without getting the packages not ready error...
 20        from .models import CachedItem
 21
 22        self._model_class = CachedItem
 23
 24    @cached_property
 25    def _model_instance(self) -> Any:
 26        try:
 27            return self._model_class.query.get(key=self.key)
 28        except self._model_class.DoesNotExist:
 29            return None
 30
 31    def reload(self) -> None:
 32        if hasattr(self, "_model_instance"):
 33            del self._model_instance
 34
 35    def _is_expired(self) -> bool:
 36        if not self._model_instance:
 37            return True
 38
 39        if not self._model_instance.expires_at:
 40            return False
 41
 42        return self._model_instance.expires_at < timezone.now()
 43
 44    def exists(self) -> bool:
 45        if self._model_instance is None:
 46            return False
 47
 48        return not self._is_expired()
 49
 50    @property
 51    def value(self) -> Any:
 52        if not self.exists():
 53            return None
 54
 55        return self._model_instance.value
 56
 57    def set(
 58        self, value: Any, expiration: datetime | timedelta | int | float | None = None
 59    ) -> Any:
 60        defaults = {
 61            "value": value,
 62        }
 63
 64        if isinstance(expiration, int | float):
 65            defaults["expires_at"] = timezone.now() + timedelta(seconds=expiration)
 66        elif isinstance(expiration, timedelta):
 67            defaults["expires_at"] = timezone.now() + expiration
 68        elif isinstance(expiration, datetime):
 69            defaults["expires_at"] = expiration
 70        else:
 71            # Keep existing expires_at value or None
 72            pass
 73
 74        # Make sure expires_at is timezone aware
 75        if (
 76            "expires_at" in defaults
 77            and defaults["expires_at"]
 78            and not timezone.is_aware(defaults["expires_at"])
 79        ):
 80            defaults["expires_at"] = timezone.make_aware(defaults["expires_at"])
 81
 82        try:
 83            item, _ = self._model_class.query.update_or_create(
 84                key=self.key, defaults=defaults
 85            )
 86        except psycopg.IntegrityError:
 87            # Most likely a race condition in creating the item,
 88            # so trying again should do an update
 89            item, _ = self._model_class.query.update_or_create(
 90                key=self.key, defaults=defaults
 91            )
 92
 93        self.reload()
 94        return item.value
 95
 96    def delete(self) -> bool:
 97        if not self._model_instance:
 98            # A no-op, but a return value you can use to know whether it did anything
 99            return False
100
101        self._model_instance.delete()
102        self.reload()
103        return True