Coverage for src/taipanstack/bridges/db_bridge.py: 100%
114 statements
« prev ^ index » next coverage.py v7.15.1, created at 2026-10-13 20:52 +0000
1"""
2DB Bridge — resilient database wrappers.
4Wraps SQLAlchemy async engine and Redis async client with
5TaipanStack's circuit breaker and retry patterns.
6"""
8from __future__ import annotations
10import asyncio
11import logging
12from typing import TYPE_CHECKING
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
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
28logger = logging.getLogger("taipanstack.bridges.db")
30# --- optional imports ------------------------------------------------------
32try:
33 import sqlalchemy # noqa: F401
34 from sqlalchemy.ext.asyncio import AsyncSession
36 _HAS_SQLALCHEMY = True
37except ImportError:
38 _HAS_SQLALCHEMY = False
40try:
41 import redis.asyncio as aioredis
43 _HAS_REDIS = True
44except ImportError:
45 _HAS_REDIS = False
48def _breaker_is_open(cb: CircuitBreaker) -> CircuitBreakerError | None:
49 """Return a ``CircuitBreakerError`` if the breaker is OPEN.
51 Args:
52 cb: The circuit breaker to check.
54 Returns:
55 Error if open, ``None`` otherwise.
57 """
58 if cb.state is CircuitState.OPEN:
59 return CircuitBreakerError(
60 f"Circuit breaker '{cb.name}' is OPEN",
61 state=cb.state,
62 )
63 return None
66class ResilientDatabase:
67 """Wraps a SQLAlchemy async engine with resilience patterns.
69 Args:
70 engine: SQLAlchemy ``AsyncEngine`` instance.
71 circuit_breaker: Optional circuit breaker.
72 retry_config: Optional retry configuration.
74 Example:
75 >>> db = ResilientDatabase(engine, circuit_breaker=breaker)
76 >>> result = await db.execute(text("SELECT 1"))
78 """
80 def __init__(
81 self,
82 engine: AsyncEngine,
83 *,
84 circuit_breaker: CircuitBreaker | None = None,
85 retry_config: RetryConfig | None = None,
86 ) -> None:
87 """Initialize the resilient database wrapper.
89 Args:
90 engine: SQLAlchemy AsyncEngine.
91 circuit_breaker: Optional circuit breaker.
92 retry_config: Optional retry config.
94 """
95 self._engine = engine # type: ignore[misc]
96 self._circuit_breaker = circuit_breaker
97 self._retry_config = retry_config
99 def _check_db_dependencies(self) -> Result[None, Exception]:
100 if not _HAS_SQLALCHEMY:
101 return Err(
102 ImportError(
103 "sqlalchemy is required for ResilientDatabase. "
104 "Install with: pip install taipanstack[bridges-db]",
105 ),
106 )
107 return Ok(None)
109 def _check_breaker_gate(self) -> Result[None, Exception]:
110 if self._circuit_breaker is not None:
111 cb_err = _breaker_is_open(self._circuit_breaker)
112 if cb_err is not None:
113 return Err(cb_err)
114 return Ok(None)
116 async def _handle_attempt_failure(
117 self,
118 exc: Exception,
119 attempt: int,
120 max_attempts: int,
121 ) -> bool:
122 """Handle a failed execution attempt.
124 Args:
125 exc: The exception that occurred.
126 attempt: The current attempt number.
127 max_attempts: Maximum number of attempts allowed.
129 Returns:
130 True if the operation should be retried, False otherwise.
132 """
133 logger.warning(
134 "DB execute attempt %d failed: %s",
135 attempt,
136 exc,
137 )
138 if self._circuit_breaker is not None:
139 self._circuit_breaker._record_failure(exc)
140 if self._retry_config is not None and attempt < max_attempts:
141 delay = calculate_delay(attempt, self._retry_config)
142 await asyncio.sleep(min(delay, 3600.0))
143 return True
144 return False
146 async def _execute_single_attempt(
147 self,
148 statement: Executable,
149 **kwargs: object,
150 ) -> Result[SQLAlchemyResult[object], Exception]:
151 try:
152 async with AsyncSession(self._engine) as session: # type: ignore[misc]
153 result = await session.execute(statement, **kwargs) # type: ignore[misc]
154 return Ok(result) # type: ignore[misc]
155 except Exception as exc:
156 return Err(exc)
158 async def _process_attempt(
159 self,
160 statement: Executable,
161 attempt: int,
162 max_attempts: int,
163 **kwargs: object,
164 ) -> Result[SQLAlchemyResult[object], Exception] | tuple[bool, Exception]:
165 result = await self._execute_single_attempt(statement, **kwargs) # type: ignore[misc]
166 if isinstance(result, Ok): # type: ignore[misc]
167 return result # type: ignore[misc]
169 last_error = result.err_value
170 should_continue = await self._handle_attempt_failure(
171 last_error, attempt, max_attempts
172 )
173 return should_continue, last_error
175 async def _execute_loop(
176 self,
177 statement: Executable,
178 max_attempts: int,
179 **kwargs: object,
180 ) -> Result[SQLAlchemyResult[object], Exception]:
181 """Execute the retry loop for database operations."""
182 last_error: Exception = RuntimeError("Database execute failed")
184 for attempt in range(1, max_attempts + 1):
185 outcome = await self._process_attempt( # type: ignore[misc]
186 statement, # type: ignore[misc]
187 attempt,
188 max_attempts,
189 **kwargs,
190 )
191 if not isinstance(outcome, tuple): # type: ignore[misc]
192 return outcome # type: ignore[misc]
194 should_continue, last_error = outcome
195 if not should_continue:
196 break
198 return Err(last_error)
200 async def execute(
201 self,
202 statement: Executable,
203 **kwargs: object,
204 ) -> Result[SQLAlchemyResult[object], Exception]:
205 """Execute a SQL statement with resilience.
207 Args:
208 statement: SQLAlchemy statement to execute.
209 **kwargs: Passed to ``session.execute``.
211 Returns:
212 ``Ok(result)`` on success, ``Err`` on failure.
214 """
215 dep_result = self._check_db_dependencies()
216 if isinstance(dep_result, Err):
217 return dep_result
219 cb_result = self._check_breaker_gate()
220 if isinstance(cb_result, Err):
221 return cb_result
223 max_attempts = 1
224 if self._retry_config is not None:
225 max_attempts = self._retry_config.max_attempts
227 return await self._execute_loop(statement, max_attempts, **kwargs) # type: ignore[misc]
230class ResilientRedis:
231 """Wraps a Redis async client with resilience patterns.
233 Args:
234 client: ``redis.asyncio.Redis`` instance.
235 circuit_breaker: Optional circuit breaker.
237 Example:
238 >>> r = ResilientRedis(redis_client, circuit_breaker=breaker)
239 >>> result = await r.execute("GET", "my_key")
241 """
243 def __init__(
244 self,
245 client: aioredis.Redis[bytes | str],
246 *,
247 circuit_breaker: CircuitBreaker | None = None,
248 ) -> None:
249 """Initialize the resilient Redis wrapper.
251 Args:
252 client: Redis async client.
253 circuit_breaker: Optional circuit breaker.
255 """
256 self._client = client # type: ignore[misc]
257 self._circuit_breaker = circuit_breaker
259 def _check_redis_dependencies(self) -> Result[None, Exception]:
260 if not _HAS_REDIS:
261 return Err(
262 ImportError(
263 "redis is required for ResilientRedis. "
264 "Install with: pip install taipanstack[bridges-db]",
265 ),
266 )
267 return Ok(None)
269 def _check_breaker_gate(self) -> Result[None, Exception]:
270 if self._circuit_breaker is not None:
271 cb_err = _breaker_is_open(self._circuit_breaker)
272 if cb_err is not None:
273 return Err(cb_err)
274 return Ok(None)
276 async def _execute_command(
277 self,
278 command: str,
279 *args: object,
280 ) -> Result[object, Exception]:
281 try:
282 fn = getattr(self._client, command.lower()) # type: ignore[misc]
283 result = await fn(*args) # type: ignore[misc]
284 return Ok(result) # type: ignore[misc]
285 except Exception as exc:
286 logger.warning("Redis command '%s' failed: %s", command, exc)
287 if self._circuit_breaker is not None:
288 self._circuit_breaker._record_failure(exc)
289 return Err(exc)
291 async def execute(
292 self,
293 command: str,
294 *args: object,
295 ) -> Result[object, Exception]:
296 """Execute a Redis command with resilience.
298 Args:
299 command: Redis command name (e.g. ``"GET"``, ``"SET"``).
300 *args: Command arguments.
302 Returns:
303 ``Ok(result)`` on success, ``Err`` on failure.
305 """
306 dep_result = self._check_redis_dependencies()
307 if isinstance(dep_result, Err):
308 return dep_result
310 cb_result = self._check_breaker_gate()
311 if isinstance(cb_result, Err):
312 return cb_result
314 return await self._execute_command(command, *args)