Coverage for src/lexigram/admin/core/distributed_lock.py: 99%

1"""Distributed lock support for lexigram-admin. 

2 

3Provides decorators and utilities for distributed locking 

4to prevent concurrent operations on the same resources. 

5 

6FWK-29: @distributed_lock decorator for concurrent safety. 

7 

8All locking is backed by a ``LockStoreProtocol`` injected via the DI 

9container (e.g. Redis-backed ``RedisLockStore`` or SQL advisory locks). 

10There is intentionally no in-memory fallback — a missing store is a 

11misconfiguration, not a condition to silently degrade from. 

12""" 

13 

14from __future__ import annotations 

15 

16import asyncio 

17from dataclasses import dataclass 

18import functools 

19import time 

20from typing import TYPE_CHECKING, Any, ParamSpec, Self, TypeVar 

21 

22from lexigram.contracts.exceptions import LockError as CoreLockError 

23 

24if TYPE_CHECKING: 

25 from collections.abc import Callable 

26 

27 from lexigram.contracts.core.stores import LockStoreProtocol 

28 

29# Type variables 

30P = ParamSpec("P") 

31R = TypeVar("R") 

32 

33 

34# ============================================================================ 

35# Lock Errors 

36# ============================================================================ 

37 

38 

39class LockError(CoreLockError): 

40 """Base lock error.""" 

41 

42 _code: str = "LEX_ERR_ADMIN_023" 

43 

44 def __init__(self, message: str = "Lock error", **kwargs: Any) -> None: 

45 super().__init__(message, **kwargs) 

46 

47 

48class LockAcquisitionError(LockError): 

49 """Could not acquire lock — another process holds it.""" 

50 

51 _code: str = "LEX_ERR_ADMIN_024" 

52 

53 def __init__(self, message: str = "Could not acquire lock", **kwargs: Any) -> None: 

54 super().__init__(message, **kwargs) 

55 

56 

57class LockTimeoutError(LockError): 

58 """Lock acquisition timed out waiting for the lock to be released.""" 

59 

60 _code: str = "LEX_ERR_ADMIN_025" 

61 

62 def __init__( 

63 self, message: str = "Lock acquisition timed out", **kwargs: Any 

64 ) -> None: 

65 super().__init__(message, **kwargs) 

66 

67 

68# ============================================================================ 

69# Lock Context Manager 

70# ============================================================================ 

71 

72 

73class AdminLockContext: 

74 """Async context manager that holds a single named distributed lock. 

75 

76 Acquired via :meth:`AdminLockManager.acquire`. Do not instantiate 

77 directly — use the manager. 

78 """ 

79 

80 def __init__( 

81 self, 

82 lock_store: LockStoreProtocol, 

83 key: str, 

84 ttl: int, 

85 timeout: float, 

86 ) -> None: 

87 self._lock_store = lock_store 

88 self.key = key 

89 self.ttl = ttl 

90 self.timeout = timeout 

91 self.owner = f"admin:{id(self)}:{time.time()}" 

92 self.acquired = False 

93 

94 async def __aenter__(self) -> Self: 

95 """Poll until the lock is acquired or *timeout* expires.""" 

96 start = time.monotonic() 

97 while time.monotonic() - start < self.timeout: 

98 if await self._lock_store.acquire(self.key, self.owner, self.ttl): 

99 self.acquired = True 

100 return self 

101 await asyncio.sleep(0.1) 

102 raise LockTimeoutError( 

103 f"Timed out waiting for distributed lock: {self.key!r} " 

104 f"(timeout={self.timeout}s)" 

105 ) 

106 

107 async def __aexit__(self, *args: object) -> None: 

108 """Release the lock if this context holds it.""" 

109 if self.acquired: 

110 await self._lock_store.release(self.key, self.owner) 

111 

112 

113# ============================================================================ 

114# Admin Lock Manager 

115# ============================================================================ 

116 

117 

118@dataclass 

119class LockConfig: 

120 """Configuration for distributed locks.""" 

121 

122 default_ttl: int = 30 

123 acquisition_timeout: float = 30.0 

124 key_prefix: str = "admin:lock:" 

125 

126 

127class AdminLockManager: 

128 """Manager for admin distributed locks. 

129 

130 Requires a ``LockStoreProtocol`` injected via the DI container. 

131 Implementations are provided by ``lexigram-cache`` (Redis) or 

132 ``lexigram-sql`` (advisory locks / lock table). 

133 

134 Example:: 

135 

136 class MyService: 

137 def __init__( 

138 self, 

139 lock_manager: AdminLockManager, 

140 ) -> None: 

141 self._locks = lock_manager 

142 

143 async def safe_bulk_delete(self, ids: list[str]) -> None: 

144 async with self._locks.acquire("users:bulk-delete"): 

145 await self._repo.delete_many(ids) 

146 """ 

147 

148 def __init__( 

149 self, 

150 lock_store: LockStoreProtocol, 

151 config: LockConfig | None = None, 

152 ) -> None: 

153 """Initialise with a persistent distributed lock store. 

154 

155 Args: 

156 lock_store: A ``LockStoreProtocol`` implementation — must be 

157 backed by a shared persistent store (Redis, SQL, etc.), not 

158 in-memory. Register it via the DI container. 

159 config: Optional lock configuration. Defaults to 

160 :class:`LockConfig`. 

161 """ 

162 self.config = config or LockConfig() 

163 self._lock_store = lock_store 

164 

165 def _full_key(self, key: str) -> str: 

166 """Prepend the configured key prefix.""" 

167 return f"{self.config.key_prefix}{key}" 

168 

169 def acquire( 

170 self, 

171 key: str, 

172 ttl: int | None = None, 

173 timeout: float | None = None, 

174 ) -> AdminLockContext: 

175 """Return an async context manager that acquires the named lock. 

176 

177 Args: 

178 key: Lock identifier (the configured prefix is prepended). 

179 ttl: Lock TTL in seconds. Defaults to 

180 :attr:`LockConfig.default_ttl`. 

181 timeout: Maximum seconds to wait for acquisition. Defaults to 

182 :attr:`LockConfig.acquisition_timeout`. 

183 

184 Returns: 

185 :class:`AdminLockContext` — use as ``async with manager.acquire(...)``. 

186 

187 Example:: 

188 

189 async with manager.acquire("resource:123"): 

190 await process_resource(123) 

191 """ 

192 return AdminLockContext( 

193 lock_store=self._lock_store, 

194 key=self._full_key(key), 

195 ttl=ttl if ttl is not None else self.config.default_ttl, 

196 timeout=timeout if timeout is not None else self.config.acquisition_timeout, 

197 ) 

198 

199 

200# ============================================================================ 

201# Distributed Lock Decorator 

202# ============================================================================ 

203 

204 

205def distributed_lock( 

206 key: str | Callable[..., str], 

207 lock_manager: AdminLockManager, 

208 ttl: int = 30, 

209 timeout: float = 30.0, 

210 on_locked: Callable[..., Any] | None = None, 

211) -> Callable[[Callable[P, R]], Callable[P, R]]: 

212 """Decorator that acquires a distributed lock before function execution. 

213 

214 The ``lock_manager`` **must** be injected — it is not created internally. 

215 This guarantees that all decorated calls coordinate through the same 

216 persistent backend (Redis, SQL, etc.) rather than silently falling back 

217 to a process-local store. 

218 

219 Args: 

220 key: Static lock key, or a callable that derives the key from the 

221 decorated function's positional/keyword arguments. 

222 lock_manager: :class:`AdminLockManager` instance wired via DI. 

223 ttl: Lock TTL in seconds. 

224 timeout: Maximum seconds to wait for acquisition. 

225 on_locked: Optional callback invoked *instead of raising* when the 

226 lock cannot be acquired within *timeout*. Receives the same 

227 ``*args, **kwargs`` as the decorated function. 

228 

229 Returns: 

230 A decorator that wraps the target coroutine function. 

231 

232 Example:: 

233 

234 manager = container.resolve(AdminLockManager) 

235 

236 @distributed_lock("bulk-export", lock_manager=manager) 

237 async def run_export() -> None: 

238 ... 

239 

240 @distributed_lock( 

241 lambda resource_id: f"resource:{resource_id}", 

242 lock_manager=manager, 

243 ) 

244 async def process_resource(resource_id: int) -> None: 

245 ... 

246 """ 

247 

248 def decorator(func: Callable[P, R]) -> Callable[P, R]: 

249 @functools.wraps(func) 

250 async def wrapper(*args: P.args, **kwargs: P.kwargs) -> R: 

251 lock_key = key(*args, **kwargs) if callable(key) else key 

252 

253 try: 

254 async with lock_manager.acquire(lock_key, ttl=ttl, timeout=timeout): 

255 return await func(*args, **kwargs) # type: ignore[misc] 

256 except LockTimeoutError: 

257 if on_locked is not None: 

258 result_on_locked = on_locked(*args, **kwargs) 

259 if getattr(result_on_locked, "__await__", None): 

260 return await result_on_locked 

261 return result_on_locked 

262 raise 

263 

264 return wrapper # type: ignore[return-value] 

265 

266 return decorator 

267 

268 

269# ============================================================================ 

270# Resource Lock Context Manager 

271# ============================================================================ 

272 

273 

274class ResourceLock: 

275 """Convenience context manager for locking a single admin resource. 

276 

277 Example:: 

278 

279 async with ResourceLock("users", user_id, lock_manager=manager): 

280 await update_user(user_id, data) 

281 """ 

282 

283 def __init__( 

284 self, 

285 resource_type: str, 

286 resource_id: Any, 

287 lock_manager: AdminLockManager, 

288 ttl: int = 30, 

289 operation: str = "edit", 

290 ) -> None: 

291 self.resource_type = resource_type 

292 self.resource_id = resource_id 

293 self.ttl = ttl 

294 self.operation = operation 

295 self._manager = lock_manager 

296 self._ctx: AdminLockContext | None = None 

297 

298 @property 

299 def key(self) -> str: 

300 """Stable lock key for this resource + operation pair.""" 

301 return f"{self.resource_type}:{self.resource_id}:{self.operation}" 

302 

303 async def __aenter__(self) -> Self: 

304 self._ctx = self._manager.acquire(self.key, ttl=self.ttl) 

305 await self._ctx.__aenter__() 

306 return self 

307 

308 async def __aexit__(self, *args: object) -> None: 

309 if self._ctx is not None: 

310 await self._ctx.__aexit__(*args) 

311 

312 

313# ============================================================================ 

314# Bulk Operation Lock 

315# ============================================================================ 

316 

317 

318class BulkOperationLock: 

319 """Lock for bulk operations to prevent concurrent modifications. 

320 

321 Uses a longer default TTL suitable for batch workloads. 

322 

323 Example:: 

324 

325 async with BulkOperationLock("users", "delete", lock_manager=manager): 

326 await bulk_delete_users(ids) 

327 """ 

328 

329 def __init__( 

330 self, 

331 resource_type: str, 

332 operation: str, 

333 lock_manager: AdminLockManager, 

334 ttl: int = 300, # 5 minutes for bulk operations 

335 ) -> None: 

336 self.resource_type = resource_type 

337 self.operation = operation 

338 self.ttl = ttl 

339 self._manager = lock_manager 

340 self._ctx: AdminLockContext | None = None 

341 

342 @property 

343 def key(self) -> str: 

344 """Stable lock key for this bulk operation.""" 

345 return f"bulk:{self.resource_type}:{self.operation}" 

346 

347 async def __aenter__(self) -> Self: 

348 self._ctx = self._manager.acquire(self.key, ttl=self.ttl) 

349 await self._ctx.__aenter__() 

350 return self 

351 

352 async def __aexit__(self, *args: object) -> None: 

353 if self._ctx is not None: 

354 await self._ctx.__aexit__(*args) 

355 

356 

357__all__ = [ 

358 # Context 

359 "AdminLockContext", 

360 # Manager + config 

361 "AdminLockManager", 

362 # Convenience locks 

363 "BulkOperationLock", 

364 # Errors 

365 "LockAcquisitionError", 

366 "LockConfig", 

367 "LockError", 

368 "LockTimeoutError", 

369 "ResourceLock", 

370 # Decorator 

371 "distributed_lock", 

372]