Coverage for snekql/mariadb/runtime.py: 90%

1"""MariaDB adapter for the backend-neutral query runtime.""" 

2 

3from __future__ import annotations 

4 

5import asyncio 

6from collections.abc import Awaitable, Callable, Sequence 

7from importlib import import_module 

8from typing import Any, Literal, cast 

9 

10from snekql.errors import ( 

11 DatabaseClosedError, 

12 DatabaseCloseTimeoutError, 

13 DatabaseClosingError, 

14 DatabaseRuntimeError, 

15 PoolTimeoutError, 

16) 

17from snekql.mariadb.config import Config 

18from snekql.mariadb.query import ( 

19 compile_mariadb_select_sql, 

20 compile_mariadb_write_sql, 

21 materialize_mariadb_select_row, 

22) 

23from snekql.mariadb.schema import initialize_mariadb_schema 

24from snekql.model import Table 

25from snekql.query import AnySelectQuery 

26from snekql.storage import SchemaPolicy 

27from snekql.structured_logging import ResolvedStructuredLogger 

28from snekql.validation import NonNegativeFloat 

29 

30 

31def _import_aiomysql() -> Any: 

32 """Import the optional MariaDB driver at runtime initialization time.""" 

33 

34 try: 

35 return cast("Any", import_module("aiomysql")) 

36 except ModuleNotFoundError as error: 

37 if error.name == "aiomysql": 

38 msg = "MariaDB runtime requires the aiomysql extra; install with snekql[aiomysql]" 

39 raise DatabaseRuntimeError(msg) from error 

40 raise 

41 

42 

43class MariaDBCursorAdapter: 

44 """Runtime cursor adapter backed by an aiomysql cursor.""" 

45 

46 def __init__(self, cursor: object) -> None: 

47 self.cursor: object = cursor 

48 

49 async def fetchone(self) -> Sequence[object] | None: 

50 row = await cast("Any", self.cursor).fetchone() 

51 if row is None: 

52 return None 

53 return cast("Sequence[object]", row) 

54 

55 async def fetchall(self) -> Sequence[Sequence[object]]: 

56 rows = await cast("Any", self.cursor).fetchall() 

57 return [cast("Sequence[object]", row) for row in rows] 

58 

59 async def close(self) -> None: 

60 close_result = cast("Any", self.cursor).close() 

61 if close_result is not None: 

62 _ = await close_result 

63 

64 

65class MariaDBConnectionAdapter: 

66 """Runtime connection adapter backed by an aiomysql connection.""" 

67 

68 def __init__(self, connection: object) -> None: 

69 self.connection: object = connection 

70 

71 async def begin(self) -> None: 

72 await cast("Any", self.connection).begin() 

73 

74 async def commit(self) -> None: 

75 await cast("Any", self.connection).commit() 

76 

77 async def rollback(self) -> None: 

78 await cast("Any", self.connection).rollback() 

79 

80 async def execute( 

81 self, 

82 sql: str, 

83 params: tuple[object, ...], 

84 ) -> MariaDBCursorAdapter: 

85 cursor = await cast("Any", self.connection).cursor() 

86 try: 

87 _ = await cursor.execute(sql, params) 

88 except Exception: 

89 close_result = cursor.close() 

90 if close_result is not None: 

91 _ = await close_result 

92 raise 

93 return MariaDBCursorAdapter(cursor) 

94 

95 

96class MariaDBConnectionPool: 

97 """Small lifecycle wrapper around an aiomysql connection pool.""" 

98 

99 def __init__( 

100 self, 

101 pool: object, 

102 *, 

103 logger: ResolvedStructuredLogger, 

104 ) -> None: 

105 self.closed: bool = False 

106 self.closing: bool = False 

107 self.logger: ResolvedStructuredLogger = logger 

108 self.pool: object = pool 

109 

110 def check_accepting_work(self) -> None: 

111 """Reject new work when closed or temporarily closing.""" 

112 

113 if self.closed: 

114 self.logger.warning( 

115 "database rejected work", backend="mariadb", reason="closed" 

116 ) 

117 msg = "database is closed" 

118 raise DatabaseClosedError(msg) 

119 if self.closing: 

120 self.logger.warning( 

121 "database rejected work", backend="mariadb", reason="closing" 

122 ) 

123 msg = "database is closing" 

124 raise DatabaseClosingError(msg) 

125 

126 async def acquire(self, acquisition_timeout: NonNegativeFloat) -> object: 

127 """Acquire a MariaDB connection within the requested timeout.""" 

128 

129 self.check_accepting_work() 

130 self.logger.debug( 

131 "connection acquisition started", 

132 backend="mariadb", 

133 timeout=acquisition_timeout, 

134 ) 

135 try: 

136 pool = cast("Any", self.pool) 

137 acquire = cast("Callable[[], Awaitable[object]]", pool.acquire) 

138 connection = await asyncio.wait_for(acquire(), timeout=acquisition_timeout) 

139 except TimeoutError as error: 

140 self.logger.warning( 

141 "connection acquisition timed out", 

142 backend="mariadb", 

143 timeout=acquisition_timeout, 

144 ) 

145 msg = "timed out acquiring database connection" 

146 raise PoolTimeoutError(msg) from error 

147 self.logger.debug("connection acquired", backend="mariadb") 

148 return connection 

149 

150 async def release(self, connection: object) -> None: 

151 """Return a connection to the underlying aiomysql pool.""" 

152 

153 release = cast("Any", self.pool).release 

154 _ = release(connection) 

155 self.logger.debug("connection released", backend="mariadb", closed=False) 

156 

157 async def close(self, close_timeout: NonNegativeFloat) -> None: 

158 """Close the underlying aiomysql pool and wait for connections.""" 

159 

160 self.logger.debug("database close started", backend="mariadb") 

161 if self.closed: 

162 self.logger.debug("database close skipped", backend="mariadb") 

163 return 

164 self.closing = True 

165 try: 

166 pool = cast("Any", self.pool) 

167 pool.close() 

168 wait_closed = cast("Callable[[], Awaitable[None]]", pool.wait_closed) 

169 await asyncio.wait_for(wait_closed(), timeout=close_timeout) 

170 except TimeoutError as error: 

171 self.closing = False 

172 self.logger.warning("database close timed out", backend="mariadb") 

173 msg = "timed out closing database" 

174 raise DatabaseCloseTimeoutError(msg) from error 

175 else: 

176 self.closed = True 

177 self.closing = False 

178 self.logger.debug("database close completed", backend="mariadb") 

179 

180 

181class MariaDBRuntime: 

182 """MariaDB adapter satisfying the backend-neutral runtime seam.""" 

183 

184 backend_family: Literal["mariadb"] = "mariadb" 

185 

186 def __init__( 

187 self, 

188 *, 

189 logger: ResolvedStructuredLogger, 

190 acquire_timeout: NonNegativeFloat, 

191 connection_pool: MariaDBConnectionPool, 

192 ) -> None: 

193 self.acquire_timeout: NonNegativeFloat = acquire_timeout 

194 self.connection_pool: MariaDBConnectionPool = connection_pool 

195 self.logger: ResolvedStructuredLogger = logger 

196 

197 async def acquire( 

198 self, 

199 acquisition_timeout: NonNegativeFloat, 

200 ) -> MariaDBConnectionAdapter: 

201 connection = await self.connection_pool.acquire(acquisition_timeout) 

202 return MariaDBConnectionAdapter(connection) 

203 

204 async def release(self, connection: object) -> None: 

205 if not isinstance(connection, MariaDBConnectionAdapter): 

206 msg = "MariaDB runtime cannot release a foreign connection" 

207 raise DatabaseRuntimeError(msg) 

208 await self.connection_pool.release(connection.connection) 

209 

210 async def close(self, close_timeout: NonNegativeFloat) -> None: 

211 await self.connection_pool.close(close_timeout) 

212 

213 def check_accepting_work(self) -> None: 

214 self.connection_pool.check_accepting_work() 

215 

216 def compile_select_sql( 

217 self, 

218 query: AnySelectQuery, 

219 ) -> tuple[str, tuple[object, ...]]: 

220 return compile_mariadb_select_sql(query) 

221 

222 def compile_write_sql(self, query: object) -> tuple[str, tuple[object, ...]]: 

223 return compile_mariadb_write_sql(query) 

224 

225 def materialize_select_row( 

226 self, 

227 query: AnySelectQuery, 

228 row: Sequence[object], 

229 ) -> object: 

230 return materialize_mariadb_select_row(query, row) 

231 

232 

233async def initialize_runtime( 

234 config: Config, 

235 models: Sequence[type[Table[Any]]], 

236 schema_policy: SchemaPolicy, 

237 *, 

238 logger: ResolvedStructuredLogger, 

239) -> MariaDBRuntime: 

240 """Initialize MariaDB connectivity, schema startup, and pool lifecycle.""" 

241 

242 aiomysql = _import_aiomysql() 

243 logger.debug("mariadb pool opening", host=config.host, port=config.port) 

244 pool = await aiomysql.create_pool( 

245 autocommit=False, 

246 charset=config.charset, 

247 connect_timeout=config.acquire_timeout, 

248 db=config.database, 

249 host=config.host, 

250 maxsize=config.pool_size, 

251 minsize=1, 

252 password=config.password, 

253 port=config.port, 

254 unix_socket=str(config.unix_socket) if config.unix_socket is not None else None, 

255 user=config.user, 

256 ) 

257 connection_pool = MariaDBConnectionPool(pool, logger=logger) 

258 connection = await connection_pool.acquire(config.acquire_timeout) 

259 try: 

260 await initialize_mariadb_schema( 

261 connection, 

262 models, 

263 schema_policy, 

264 logger=logger, 

265 ) 

266 except Exception: 

267 await connection_pool.release(connection) 

268 await connection_pool.close(config.acquire_timeout) 

269 raise 

270 await connection_pool.release(connection) 

271 return MariaDBRuntime( 

272 logger=logger, 

273 acquire_timeout=config.acquire_timeout, 

274 connection_pool=connection_pool, 

275 ) 

276 

277 

278__all__ = [ 

279 "MariaDBConnectionAdapter", 

280 "MariaDBConnectionPool", 

281 "MariaDBCursorAdapter", 

282 "MariaDBRuntime", 

283 "initialize_runtime", 

284]