Coverage for python/pyairflowtester/dependency_intelligence/cache.py: 91%
142 statements
« prev ^ index » next coverage.py v7.16.1, created at 2026-09-19 20:43 +0530
1"""Tiered caching for dependency-graph analysis results.
3`DEPENDENCY_CACHING_STRATEGY.md` sketched a 4-layer design (L1 in-memory, L2
4Redis, L3 SQLite, L4 DuckDB) that was never actually implemented anywhere in
5this package. This module implements two of those layers for real:
7- L1 `InMemoryCache`: thread-safe LRU with real per-entry TTL.
8- L3 `SqliteCache`: persistent, multi-process-safe (WAL journal mode) --
9 the layer that lets expensive graph analysis (cycle detection, strongly
10 connected components) survive across separate CLI invocations, which
11 nothing in this codebase could do before this.
12- `TieredCache`: L1-then-L3 lookup, with real event-driven invalidation
13 (`register_invalidation_rule` / `emit`) instead of relying on TTL expiry
14 alone.
16L2 (Redis) and L4 (DuckDB) are NOT implemented here -- they need an external
17service/heavier dependency this otherwise dependency-light package
18(`click`, `rich` only) doesn't currently take on. `TieredCache` is built so
19a `RedisCache`/`DuckDbCache` could plug in as another tier later without
20changing this interface, but that's future work, not shipped today.
21"""
23from __future__ import annotations
25import json
26import sqlite3
27import threading
28import time
29from collections import OrderedDict
30from dataclasses import dataclass
31from pathlib import Path
32from typing import Any, Callable, Dict, List, Optional, Union
35@dataclass
36class CacheStats:
37 hits: int = 0
38 misses: int = 0
39 evictions: int = 0
41 @property
42 def hit_rate(self) -> float:
43 total = self.hits + self.misses
44 return self.hits / total if total else 0.0
47class InMemoryCache:
48 """Thread-safe LRU cache (L1) with real per-entry TTL.
50 Unlike the ad-hoc, TTL-less dicts scattered elsewhere in this codebase
51 (e.g. `DependencyGraphEngine._upstream_cache`), entries here actually
52 expire and the cache actually evicts least-recently-used entries once
53 `max_size` is exceeded.
54 """
56 def __init__(self, max_size: int = 10_000):
57 self.max_size = max_size
58 self._data: "OrderedDict[str, tuple[Any, Optional[float]]]" = OrderedDict()
59 self._lock = threading.Lock()
60 self.stats = CacheStats()
62 def get(self, key: str) -> Optional[Any]:
63 with self._lock:
64 entry = self._data.get(key)
65 if entry is None:
66 self.stats.misses += 1
67 return None
68 value, expires_at = entry
69 if expires_at is not None and time.time() > expires_at:
70 del self._data[key]
71 self.stats.misses += 1
72 return None
73 self._data.move_to_end(key)
74 self.stats.hits += 1
75 return value
77 def set(self, key: str, value: Any, ttl_seconds: Optional[float] = None) -> None:
78 expires_at = time.time() + ttl_seconds if ttl_seconds is not None else None
79 with self._lock:
80 self._data[key] = (value, expires_at)
81 self._data.move_to_end(key)
82 while len(self._data) > self.max_size:
83 self._data.popitem(last=False)
84 self.stats.evictions += 1
86 def delete(self, key: str) -> None:
87 with self._lock:
88 self._data.pop(key, None)
90 def delete_prefix(self, prefix: str) -> int:
91 with self._lock:
92 matching = [k for k in self._data if k.startswith(prefix)]
93 for k in matching:
94 del self._data[k]
95 return len(matching)
97 def clear(self) -> None:
98 with self._lock:
99 self._data.clear()
102class SqliteCache:
103 """Persistent cache (L3), backed by SQLite in WAL journal mode.
105 WAL mode is what makes this safe for concurrent access from multiple
106 processes (e.g. two CLI invocations running at once against the same
107 cache file) -- readers don't block the writer and vice versa, which the
108 default rollback-journal mode doesn't give you.
109 """
111 def __init__(self, db_path: Union[str, Path]):
112 self.db_path = str(db_path)
113 Path(self.db_path).parent.mkdir(parents=True, exist_ok=True)
114 self._lock = threading.Lock()
115 self._conn = sqlite3.connect(self.db_path, check_same_thread=False)
116 self._conn.execute("PRAGMA journal_mode=WAL")
117 self._conn.execute("PRAGMA synchronous=NORMAL")
118 self._conn.execute("""
119 CREATE TABLE IF NOT EXISTS cache_entries (
120 key TEXT PRIMARY KEY,
121 value TEXT NOT NULL,
122 expires_at REAL
123 )
124 """)
125 self._conn.commit()
126 self.stats = CacheStats()
128 def get(self, key: str) -> Optional[Any]:
129 with self._lock:
130 row = self._conn.execute(
131 "SELECT value, expires_at FROM cache_entries WHERE key = ?", (key,)
132 ).fetchone()
133 if row is None:
134 self.stats.misses += 1
135 return None
136 value_json, expires_at = row
137 if expires_at is not None and time.time() > expires_at:
138 self._conn.execute("DELETE FROM cache_entries WHERE key = ?", (key,))
139 self._conn.commit()
140 self.stats.misses += 1
141 return None
142 self.stats.hits += 1
143 return json.loads(value_json)
145 def set(self, key: str, value: Any, ttl_seconds: Optional[float] = None) -> None:
146 expires_at = time.time() + ttl_seconds if ttl_seconds is not None else None
147 value_json = json.dumps(value)
148 with self._lock:
149 self._conn.execute(
150 "INSERT INTO cache_entries (key, value, expires_at) VALUES (?, ?, ?) "
151 "ON CONFLICT(key) DO UPDATE SET "
152 "value = excluded.value, expires_at = excluded.expires_at",
153 (key, value_json, expires_at),
154 )
155 self._conn.commit()
157 def delete(self, key: str) -> None:
158 with self._lock:
159 self._conn.execute("DELETE FROM cache_entries WHERE key = ?", (key,))
160 self._conn.commit()
162 def delete_prefix(self, prefix: str) -> int:
163 with self._lock:
164 cursor = self._conn.execute(
165 "DELETE FROM cache_entries WHERE key LIKE ? ESCAPE '\\'",
166 (prefix.replace("%", r"\%").replace("_", r"\_") + "%",),
167 )
168 self._conn.commit()
169 return cursor.rowcount
171 def clear(self) -> None:
172 with self._lock:
173 self._conn.execute("DELETE FROM cache_entries")
174 self._conn.commit()
176 def close(self) -> None:
177 with self._lock:
178 self._conn.close()
181class TieredCache:
182 """L1 (in-memory) -> L3 (SQLite) cache with event-driven invalidation.
184 A cache miss on L1 checks L3 before giving up, and a L3 hit repopulates
185 L1 (standard tiered-cache read-through behavior). Writes go to both
186 tiers so an L1 eviction doesn't lose data L3 still has.
188 Event-driven invalidation: register a rule that computes which key
189 prefixes to drop for a named event (e.g. "graph_mutated"), then call
190 `emit()` whenever that event actually happens. This is what the
191 strategy doc asked for instead of relying on TTL expiry alone --
192 invalidation happens exactly when the underlying data changes, not up
193 to `ttl_seconds` later.
194 """
196 def __init__(self, l1: Optional[InMemoryCache] = None, l3: Optional[SqliteCache] = None):
197 self.l1 = l1 or InMemoryCache()
198 self.l3 = l3
199 self._invalidation_rules: Dict[str, List[Callable[..., List[str]]]] = {}
201 def get(self, key: str) -> Optional[Any]:
202 value = self.l1.get(key)
203 if value is not None:
204 return value
205 if self.l3 is not None:
206 value = self.l3.get(key)
207 if value is not None:
208 self.l1.set(key, value)
209 return value
210 return None
212 def set(self, key: str, value: Any, ttl_seconds: Optional[float] = None) -> None:
213 self.l1.set(key, value, ttl_seconds)
214 if self.l3 is not None:
215 self.l3.set(key, value, ttl_seconds)
217 def invalidate(self, key: str) -> None:
218 self.l1.delete(key)
219 if self.l3 is not None:
220 self.l3.delete(key)
222 def invalidate_prefix(self, prefix: str) -> int:
223 count = self.l1.delete_prefix(prefix)
224 if self.l3 is not None:
225 count += self.l3.delete_prefix(prefix)
226 return count
228 def register_invalidation_rule(
229 self, event: str, key_prefixes_fn: Callable[..., List[str]]
230 ) -> None:
231 """Register `key_prefixes_fn(**context) -> [prefix, ...]` to run
232 whenever `emit(event, **context)` fires -- every prefix it returns
233 gets invalidated across both tiers."""
234 self._invalidation_rules.setdefault(event, []).append(key_prefixes_fn)
236 def emit(self, event: str, **context: Any) -> List[str]:
237 """Fire `event`, running every rule registered for it and
238 invalidating every prefix they return. Returns the prefixes that
239 were invalidated (useful for logging/testing)."""
240 invalidated: List[str] = []
241 for rule in self._invalidation_rules.get(event, []):
242 for prefix in rule(**context):
243 self.invalidate_prefix(prefix)
244 invalidated.append(prefix)
245 return invalidated