Coverage for src/taipanstack/bridges/db_bridge.py: 100%

1""" 

2DB Bridge — resilient database wrappers. 

3 

4Wraps SQLAlchemy async engine and Redis async client with 

5TaipanStack's circuit breaker and retry patterns. 

6""" 

7 

8from __future__ import annotations 

9 

10import asyncio 

11import logging 

12from typing import TYPE_CHECKING 

13 

14if TYPE_CHECKING: 

15 import redis.asyncio as aioredis 

16 from sqlalchemy.engine import Result as SQLAlchemyResult 

17 from sqlalchemy.ext.asyncio import AsyncEngine 

18 from sqlalchemy.sql.executable import Executable 

19 

20from taipanstack.core.result import Err, Ok, Result 

21from taipanstack.resilience.circuit_breaker import ( 

22 CircuitBreaker, 

23 CircuitBreakerError, 

24 CircuitState, 

25) 

26from taipanstack.resilience.retry import RetryConfig, calculate_delay 

27 

28logger = logging.getLogger("taipanstack.bridges.db") 

29 

30# --- optional imports ------------------------------------------------------ 

31 

32try: 

33 import sqlalchemy # noqa: F401 

34 from sqlalchemy import text as sa_text 

35 from sqlalchemy.ext.asyncio import AsyncSession 

36 

37 _HAS_SQLALCHEMY = True 

38except ImportError: 

39 _HAS_SQLALCHEMY = False 

40 

41try: 

42 import redis.asyncio as aioredis 

43 

44 _HAS_REDIS = True 

45except ImportError: 

46 _HAS_REDIS = False 

47 

48 

49def _breaker_is_open(cb: CircuitBreaker) -> CircuitBreakerError | None: 

50 """Return a ``CircuitBreakerError`` if the breaker is OPEN. 

51 

52 Args: 

53 cb: The circuit breaker to check. 

54 

55 Returns: 

56 Error if open, ``None`` otherwise. 

57 

58 """ 

59 if cb.state == CircuitState.OPEN: 

60 return CircuitBreakerError( 

61 f"Circuit breaker '{cb.name}' is OPEN", 

62 state=cb.state, 

63 ) 

64 return None 

65 

66 

67class ResilientDatabase: 

68 """Wraps a SQLAlchemy async engine with resilience patterns. 

69 

70 Args: 

71 engine: SQLAlchemy ``AsyncEngine`` instance. 

72 circuit_breaker: Optional circuit breaker. 

73 retry_config: Optional retry configuration. 

74 

75 Example: 

76 >>> db = ResilientDatabase(engine, circuit_breaker=breaker) 

77 >>> result = await db.execute(text("SELECT 1")) 

78 

79 """ 

80 

81 def __init__( 

82 self, 

83 engine: AsyncEngine, 

84 *, 

85 circuit_breaker: CircuitBreaker | None = None, 

86 retry_config: RetryConfig | None = None, 

87 ) -> None: 

88 """Initialize the resilient database wrapper. 

89 

90 Args: 

91 engine: SQLAlchemy AsyncEngine. 

92 circuit_breaker: Optional circuit breaker. 

93 retry_config: Optional retry config. 

94 

95 """ 

96 self._engine = engine 

97 self._circuit_breaker = circuit_breaker 

98 self._retry_config = retry_config 

99 

100 def _check_db_dependencies(self) -> Result[None, Exception]: 

101 if not _HAS_SQLALCHEMY: 

102 return Err( 

103 ImportError( 

104 "sqlalchemy is required for ResilientDatabase. " 

105 "Install with: pip install taipanstack[bridges-db]", 

106 ), 

107 ) 

108 return Ok(None) 

109 

110 def _check_breaker_gate(self) -> Result[None, Exception]: 

111 if self._circuit_breaker is not None: 

112 cb_err = _breaker_is_open(self._circuit_breaker) 

113 if cb_err is not None: 

114 return Err(cb_err) 

115 return Ok(None) 

116 

117 async def _handle_attempt_failure( 

118 self, 

119 exc: Exception, 

120 attempt: int, 

121 max_attempts: int, 

122 ) -> bool: 

123 """Handle a failed execution attempt. 

124 

125 Args: 

126 exc: The exception that occurred. 

127 attempt: The current attempt number. 

128 max_attempts: Maximum number of attempts allowed. 

129 

130 Returns: 

131 True if the operation should be retried, False otherwise. 

132 

133 """ 

134 logger.warning( 

135 "DB execute attempt %d failed: %s", 

136 attempt, 

137 exc, 

138 ) 

139 if self._circuit_breaker is not None: 

140 self._circuit_breaker._record_failure(exc) 

141 if self._retry_config is not None and attempt < max_attempts: 

142 delay = calculate_delay(attempt, self._retry_config) 

143 await asyncio.sleep(min(delay, 3600.0)) 

144 return True 

145 return False 

146 

147 async def _execute_single_attempt( 

148 self, 

149 statement: Executable, 

150 **kwargs: object, 

151 ) -> Result[SQLAlchemyResult[object], Exception]: 

152 try: 

153 async with AsyncSession(self._engine) as session: 

154 result = await session.execute(statement, **kwargs) 

155 return Ok(result) 

156 except Exception as exc: 

157 return Err(exc) 

158 

159 async def _process_attempt( 

160 self, 

161 statement: Executable, 

162 attempt: int, 

163 max_attempts: int, 

164 **kwargs: object, 

165 ) -> Result[SQLAlchemyResult[object], Exception] | tuple[bool, Exception]: 

166 result = await self._execute_single_attempt(statement, **kwargs) 

167 if isinstance(result, Ok): 

168 return result 

169 

170 last_error = result.err_value 

171 should_continue = await self._handle_attempt_failure( 

172 last_error, attempt, max_attempts 

173 ) 

174 return should_continue, last_error 

175 

176 async def _execute_loop( 

177 self, 

178 statement: Executable, 

179 max_attempts: int, 

180 **kwargs: object, 

181 ) -> Result[SQLAlchemyResult[object], Exception]: 

182 """Execute the retry loop for database operations.""" 

183 last_error: Exception = RuntimeError("Database execute failed") 

184 

185 for attempt in range(1, max_attempts + 1): 

186 outcome = await self._process_attempt( 

187 statement, attempt, max_attempts, **kwargs 

188 ) 

189 if not isinstance(outcome, tuple): 

190 return outcome 

191 

192 should_continue, last_error = outcome 

193 if not should_continue: 

194 break 

195 

196 return Err(last_error) 

197 

198 async def execute( 

199 self, 

200 statement: Executable, 

201 **kwargs: object, 

202 ) -> Result[SQLAlchemyResult[object], Exception]: 

203 """Execute a SQL statement with resilience. 

204 

205 Args: 

206 statement: SQLAlchemy statement to execute. 

207 **kwargs: Passed to ``session.execute``. 

208 

209 Returns: 

210 ``Ok(result)`` on success, ``Err`` on failure. 

211 

212 """ 

213 dep_result = self._check_db_dependencies() 

214 if isinstance(dep_result, Err): 

215 return dep_result 

216 

217 cb_result = self._check_breaker_gate() 

218 if isinstance(cb_result, Err): 

219 return cb_result 

220 

221 max_attempts = 1 

222 if self._retry_config is not None: 

223 max_attempts = self._retry_config.max_attempts 

224 

225 return await self._execute_loop(statement, max_attempts, **kwargs) 

226 

227 async def health_check(self) -> Result[bool, Exception]: 

228 """Check database connectivity. 

229 

230 Executes ``SELECT 1`` to verify the connection is alive. 

231 

232 Returns: 

233 ``Ok(True)`` if healthy, ``Err`` on failure. 

234 

235 """ 

236 if not _HAS_SQLALCHEMY: 

237 return Err(ImportError("sqlalchemy is required for health check")) 

238 

239 try: 

240 async with AsyncSession(self._engine) as session: 

241 await session.execute(sa_text("SELECT 1")) 

242 return Ok(True) 

243 except Exception as exc: 

244 return Err(exc) 

245 

246 

247class ResilientRedis: 

248 """Wraps a Redis async client with resilience patterns. 

249 

250 Args: 

251 client: ``redis.asyncio.Redis`` instance. 

252 circuit_breaker: Optional circuit breaker. 

253 

254 Example: 

255 >>> r = ResilientRedis(redis_client, circuit_breaker=breaker) 

256 >>> result = await r.execute("GET", "my_key") 

257 

258 """ 

259 

260 def __init__( 

261 self, 

262 client: aioredis.Redis[bytes | str], 

263 *, 

264 circuit_breaker: CircuitBreaker | None = None, 

265 ) -> None: 

266 """Initialize the resilient Redis wrapper. 

267 

268 Args: 

269 client: Redis async client. 

270 circuit_breaker: Optional circuit breaker. 

271 

272 """ 

273 self._client = client 

274 self._circuit_breaker = circuit_breaker 

275 

276 def _check_redis_dependencies(self) -> Result[None, Exception]: 

277 if not _HAS_REDIS: 

278 return Err( 

279 ImportError( 

280 "redis is required for ResilientRedis. " 

281 "Install with: pip install taipanstack[bridges-db]", 

282 ), 

283 ) 

284 return Ok(None) 

285 

286 def _check_breaker_gate(self) -> Result[None, Exception]: 

287 if self._circuit_breaker is not None: 

288 cb_err = _breaker_is_open(self._circuit_breaker) 

289 if cb_err is not None: 

290 return Err(cb_err) 

291 return Ok(None) 

292 

293 async def _execute_command( 

294 self, 

295 command: str, 

296 *args: object, 

297 ) -> Result[object, Exception]: 

298 try: 

299 fn = getattr(self._client, command.lower()) 

300 result = await fn(*args) 

301 return Ok(result) 

302 except Exception as exc: 

303 logger.warning("Redis command '%s' failed: %s", command, exc) 

304 if self._circuit_breaker is not None: 

305 self._circuit_breaker._record_failure(exc) 

306 return Err(exc) 

307 

308 async def execute( 

309 self, 

310 command: str, 

311 *args: object, 

312 ) -> Result[object, Exception]: 

313 """Execute a Redis command with resilience. 

314 

315 Args: 

316 command: Redis command name (e.g. ``"GET"``, ``"SET"``). 

317 *args: Command arguments. 

318 

319 Returns: 

320 ``Ok(result)`` on success, ``Err`` on failure. 

321 

322 """ 

323 dep_result = self._check_redis_dependencies() 

324 if isinstance(dep_result, Err): 

325 return dep_result 

326 

327 cb_result = self._check_breaker_gate() 

328 if isinstance(cb_result, Err): 

329 return cb_result 

330 

331 return await self._execute_command(command, *args) 

332 

333 async def health_check(self) -> Result[bool, Exception]: 

334 """Check Redis connectivity via PING. 

335 

336 Returns: 

337 ``Ok(True)`` if healthy, ``Err`` on failure. 

338 

339 """ 

340 if not _HAS_REDIS: 

341 return Err(ImportError("redis is required for health check")) 

342 

343 try: 

344 pong = await self._client.ping() 

345 return Ok(bool(pong)) 

346 except Exception as exc: 

347 return Err(exc)