1from __future__ import annotations
2
3import json
4from collections.abc import Callable, Iterable, Mapping
5from datetime import datetime, timedelta
6from typing import TYPE_CHECKING, Any
7
8from plain.postgres import get_connection, transaction
9from plain.postgres.dialect import quote_name
10from plain.utils import timezone
11
12if TYPE_CHECKING:
13 from .models import CachedItem
14
15# An expiration argument: seconds (int/float), a timedelta, an absolute
16# datetime, or None for "never expires".
17Expiration = datetime | timedelta | int | float | None
18
19
20def _coerce_expiration(expiration: Expiration, *, now: datetime) -> datetime | None:
21 """Resolve an `expiration` argument to a timezone-aware `datetime`, or `None`
22 for "never expires". Relative expirations are measured from `now`, which the
23 caller passes in so a single write derives its expiry and row timestamps from
24 one instant.
25
26 Accepts seconds (`int`/`float`), a `timedelta`, or an absolute `datetime`.
27 Unsupported types are rejected loudly rather than silently treated as "no
28 expiry" -- in particular a `bool` (which is an `int` subclass) and a bare
29 `date` (which is not a `datetime`) are common mistakes.
30 """
31 if expiration is None:
32 return None
33
34 if isinstance(expiration, bool):
35 raise TypeError(
36 "expiration must be seconds, a timedelta, or a datetime -- not a bool"
37 )
38
39 if isinstance(expiration, int | float):
40 expires_at = now + timedelta(seconds=expiration)
41 elif isinstance(expiration, timedelta):
42 expires_at = now + expiration
43 elif isinstance(expiration, datetime):
44 expires_at = expiration
45 else:
46 raise TypeError(
47 "expiration must be seconds, a timedelta, or a datetime -- got "
48 f"{type(expiration).__name__}"
49 )
50
51 if not timezone.is_aware(expires_at):
52 expires_at = timezone.make_aware(expires_at)
53 return expires_at
54
55
56class Cache:
57 """A key/value cache backed by the `CachedItem` Postgres model.
58
59 Reads are expiry-aware: an entry past its `expires_at` reads as absent (the
60 `clear_expired` chore / `plain cache clear-expired` deletes it out of band).
61 Stateless -- nothing is held between calls, so every read reflects the
62 current row. Stored values must be JSON-serializable.
63
64 Use the module-level `cache` singleton: `from plain.cache import cache`.
65 """
66
67 @property
68 def _model(self) -> type[CachedItem]:
69 # Imported lazily so `from plain.cache import cache` works at import time,
70 # before the packages registry is ready.
71 from .models import CachedItem
72
73 return CachedItem
74
75 # Reading -----------------------------------------------------------------
76
77 def get(self, key: str, default: Any = None) -> Any:
78 """Return the value for `key`, or `default` if it's absent or expired."""
79 item = self._model.query.live().filter(key=key).first()
80 return item.value if item is not None else default
81
82 def get_many(self, keys: Iterable[str]) -> dict[str, Any]:
83 """Return a `{key: value}` dict of the live entries among `keys`.
84
85 Missing/expired keys are omitted. One query regardless of how many keys.
86 """
87 items = self._model.query.live().filter(key__in=list(keys))
88 return {item.key: item.value for item in items}
89
90 # Writing -----------------------------------------------------------------
91
92 def set(self, key: str, value: Any, *, expiration: Expiration = None) -> None:
93 """Store `value` under `key`. `expiration=None` (the default) never expires.
94
95 Always rewrites the whole entry, including its expiry -- to change only
96 the expiry of a large value without rewriting it, use `touch()`.
97 """
98 # A single-key upsert -- same one-statement INSERT ... ON CONFLICT path
99 # as set_many(), so set/set_many share one write mechanism.
100 self.set_many({key: value}, expiration=expiration)
101
102 def set_many(
103 self, mapping: Mapping[str, Any], *, expiration: Expiration = None
104 ) -> None:
105 """Store every `{key: value}` in `mapping` with a shared expiration."""
106 if not mapping:
107 return
108
109 # bulk_create fires pre_save, so updated_at's update_now stamps a fresh
110 # now() at write time on its own. created_at (no update_now) would
111 # otherwise fall to its DB default, evaluated a hair later -- leaving a
112 # brand-new row with updated_at < created_at. Stamp created_at from an
113 # up-front `now` so created_at <= updated_at; it's omitted from
114 # update_fields, so it's preserved on conflict.
115 now = timezone.now()
116 expires_at = _coerce_expiration(expiration, now=now)
117 items = [
118 self._model(key=key, value=value, expires_at=expires_at, created_at=now)
119 for key, value in mapping.items()
120 ]
121 self._model.query.bulk_create(
122 items,
123 update_conflicts=True,
124 update_fields=["value", "expires_at", "updated_at"],
125 unique_fields=["key"],
126 )
127
128 def get_or_set(
129 self,
130 key: str,
131 default: Callable[[], Any] | Any,
132 *,
133 expiration: Expiration = None,
134 ) -> Any:
135 """Return the value for `key`, computing and storing it on a miss.
136
137 `default` may be a value or a zero-arg callable; the callable is only
138 invoked on a miss (so a callable can't be cached *as* the value). A
139 stored `None` counts as a hit (it won't recompute).
140 """
141 item = self._model.query.live().filter(key=key).first()
142 if item is not None:
143 return item.value
144
145 value = default() if callable(default) else default
146 self.set(key, value, expiration=expiration)
147 return value
148
149 # Counters ----------------------------------------------------------------
150
151 def increment(
152 self, key: str, delta: int = 1, *, expiration: Expiration = None
153 ) -> int | float:
154 """Atomically add `delta` to the number at `key` and return the new total.
155
156 One `INSERT ... ON CONFLICT` statement, so concurrent callers can't lose
157 updates the way a read-then-`set()` would -- the right primitive for
158 counters and fixed-window rate limiters.
159
160 Expiry follows a fixed-window rule:
161
162 - A **missing or expired** key starts fresh at `delta` and takes
163 `expiration` -- a lapsed window resets cleanly to a new deadline.
164 - A **live** key adds `delta` to the existing total and keeps its
165 current `expires_at` -- the window holds its original deadline,
166 regardless of the `expiration` argument. (To slide the expiry too,
167 call `touch()`.)
168
169 A key with no numeric value yet -- missing, or storing `None` -- counts
170 as `0`, so the first increment starts from `delta`. Incrementing a key
171 that stores a non-numeric value (a string, list, etc.) raises.
172 """
173 now = timezone.now()
174 expires_at = _coerce_expiration(expiration, now=now)
175 table = quote_name(self._model.model_options.db_table)
176
177 # The existing-row expiry test matches `expired()` (the inverse of the
178 # `live()` filter reads use): an expired row counts as absent, so the
179 # counter restarts from `delta` with a new deadline instead of resuming a
180 # stale total whose window already lapsed. A never-expiring row has
181 # expires_at = NULL, and `NULL < now` is NULL (falsy in CASE), so it
182 # correctly falls through to the accumulate branch.
183 sql = f"""
184 INSERT INTO {table} (key, value, expires_at, created_at, updated_at)
185 VALUES (%(key)s, to_jsonb(%(delta)s::numeric), %(expires_at)s, %(now)s, %(now)s)
186 ON CONFLICT (key) DO UPDATE SET
187 value = CASE
188 WHEN {table}.expires_at < %(now)s
189 THEN EXCLUDED.value
190 -- `value::text` keeps JSON syntax, so a string like "5"
191 -- stays quoted and fails ::numeric -- only a real JSON
192 -- number parses. NULLIF maps JSON null to SQL NULL so
193 -- COALESCE treats a null/absent value as 0.
194 ELSE to_jsonb(COALESCE(NULLIF({table}.value, 'null'::jsonb)::text::numeric, 0) + %(delta)s)
195 END,
196 expires_at = CASE
197 WHEN {table}.expires_at < %(now)s
198 THEN EXCLUDED.expires_at
199 ELSE {table}.expires_at
200 END,
201 updated_at = %(now)s
202 RETURNING value::text
203 """
204 params = {"key": key, "delta": delta, "expires_at": expires_at, "now": now}
205 # A non-numeric value raises DataError, which leaves the DB transaction
206 # aborted. Mark the connection so an enclosing atomic() block rolls back
207 # even if the caller catches the error -- the same guard ORM writes use.
208 with (
209 transaction.mark_for_rollback_on_error(),
210 get_connection().cursor() as cursor,
211 ):
212 cursor.execute(sql, params)
213 row = cursor.fetchone()
214 assert row is not None # INSERT ... ON CONFLICT DO UPDATE always returns a row
215
216 # `value::text` returns the new total as JSON text regardless of driver;
217 # decode it to the same Python number `get()` would yield.
218 return json.loads(row[0])
219
220 def decrement(
221 self, key: str, delta: int = 1, *, expiration: Expiration = None
222 ) -> int | float:
223 """Atomically subtract `delta` from the number at `key`. See `increment()`."""
224 return self.increment(key, -delta, expiration=expiration)
225
226 def touch(self, key: str, *, expiration: Expiration = None) -> bool:
227 """Change a live entry's expiration *without* rewriting its value.
228
229 `set()` always rewrites `value`, so refreshing a large entry's TTL
230 re-TOASTs the whole blob. `touch()` writes only `expires_at` and
231 `updated_at` -- a heap-only write that reuses the existing TOAST pointer,
232 so a multi-megabyte value isn't re-written. Ideal for a sliding-TTL cache
233 of large values.
234
235 `expiration=None` clears the expiry (never expires). Returns `True` if a
236 live entry was updated, `False` if `key` is absent or already expired.
237 """
238 # QuerySet.update() issues a direct SQL UPDATE and does NOT fire pre_save,
239 # so updated_at's update_now won't bump on its own -- stamp it by hand.
240 # (set_many() relies on pre_save instead, since bulk_create does fire it.)
241 now = timezone.now()
242 updated = (
243 self._model.query.live()
244 .filter(key=key)
245 .update(expires_at=_coerce_expiration(expiration, now=now), updated_at=now)
246 )
247 return updated > 0
248
249 # Deleting ----------------------------------------------------------------
250
251 def delete(self, key: str) -> bool:
252 """Delete `key`. Returns `True` if it existed, `False` otherwise."""
253 return self._model.query.filter(key=key).delete() > 0
254
255 def delete_many(self, keys: Iterable[str]) -> int:
256 """Delete every key in `keys`. Returns the number of rows deleted."""
257 return self._model.query.filter(key__in=list(keys)).delete()
258
259 def clear(self) -> int:
260 """Delete every entry in the cache. Returns the number of rows deleted."""
261 return self._model.query.all().delete()
262
263
264cache = Cache()