Coverage for snekql/_pool.py: 73%
163 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-06-07 21:13 +0300
1"""Internal async SQLite connection pool for Query Runtime."""
3from __future__ import annotations
5import asyncio
6from collections.abc import Sequence
7from pathlib import Path
9from aiosqlite import Connection, Error, connect
11from snekql.errors import (
12 DatabaseClosedError,
13 DatabaseCloseTimeoutError,
14 DatabaseClosingError,
15 DatabaseRuntimeError,
16 PoolTimeoutError,
17)
18from snekql.structured_logging import ResolvedStructuredLogger
19from snekql.validation import NonNegativeFloat, PositiveInt
22def normalize_sqlite_database(database: object) -> str:
23 """Convert the public database initializer value to an aiosqlite path."""
25 if type(database) is str and database == ":memory:":
26 return ":memory:"
27 if isinstance(database, Path):
28 return str(database)
29 msg = "database must be a pathlib.Path or the exact string ':memory:'"
30 raise DatabaseRuntimeError(
31 msg,
32 )
35async def open_sqlite_connection(database_path: str) -> Connection:
36 """Open and prove an async SQLite connection."""
38 try:
39 connection = await connect(database_path, isolation_level=None)
40 cursor = await connection.execute("SELECT 1")
41 try:
42 _ = await cursor.fetchone()
43 finally:
44 await cursor.close()
45 except Error as error:
46 msg = "could not initialize SQLite connection"
47 raise DatabaseRuntimeError(msg) from error
48 else:
49 return connection
52async def close_sqlite_connection(connection: Connection) -> None:
53 """Close an async SQLite connection with package-originated errors."""
55 try:
56 await connection.close()
57 except Error as error:
58 msg = "could not close SQLite connection"
59 raise DatabaseRuntimeError(msg) from error
62class SQLiteConnectionPool:
63 """Bounded lazy async SQLite connection pool owned by a Database."""
65 active_connections: int
66 closed: bool
67 closing: bool
68 condition: asyncio.Condition
69 database_path: str
70 idle_connections: list[Connection]
71 opening_connections: int
72 pool_size: PositiveInt
74 def __init__(
75 self,
76 *,
77 logger: ResolvedStructuredLogger,
78 database_path: str,
79 initial_connection: Connection,
80 pool_size: PositiveInt,
81 ) -> None:
82 self.active_connections: int = 0
83 self.closed: bool = False
84 self.closing: bool = False
85 self.condition: asyncio.Condition = asyncio.Condition()
86 self.database_path: str = database_path
87 self.idle_connections: list[Connection] = [initial_connection]
88 self.logger: ResolvedStructuredLogger = logger
89 self.opening_connections: int = 0
90 self.pool_size: PositiveInt = pool_size
92 def check_accepting_work(self) -> None:
93 """Reject new work when closed or temporarily closing."""
95 if self.closed:
96 self.logger.warning("database rejected work", reason="closed")
97 msg = "database is closed"
98 raise DatabaseClosedError(msg)
99 if self.closing:
100 self.logger.warning("database rejected work", reason="closing")
101 msg = "database is closing"
102 raise DatabaseClosingError(msg)
104 async def acquire(self, acquisition_timeout: NonNegativeFloat, /) -> Connection:
105 """Acquire an existing or lazily-created connection within timeout."""
107 self.logger.debug(
108 "connection acquisition started",
109 backend="sqlite",
110 timeout=acquisition_timeout,
111 )
112 event_loop = asyncio.get_running_loop()
113 deadline = event_loop.time() + acquisition_timeout
114 while True:
115 async with self.condition:
116 self.check_accepting_work()
117 if self.idle_connections:
118 connection = self.idle_connections.pop()
119 self.active_connections += 1
120 self.logger.debug(
121 "connection acquired",
122 backend="sqlite",
123 source="idle",
124 )
125 return connection
126 if self.connection_count() < self.pool_size:
127 self.opening_connections += 1
128 break
129 remaining_timeout = deadline - event_loop.time()
130 if remaining_timeout <= 0:
131 self.logger.warning(
132 "connection acquisition timed out",
133 backend="sqlite",
134 timeout=acquisition_timeout,
135 )
136 msg = "timed out acquiring database connection"
137 raise PoolTimeoutError(msg)
138 try:
139 _ = await asyncio.wait_for(
140 self.condition.wait(),
141 timeout=remaining_timeout,
142 )
143 except TimeoutError as error:
144 self.logger.warning(
145 "connection acquisition timed out",
146 backend="sqlite",
147 timeout=acquisition_timeout,
148 )
149 msg = "timed out acquiring database connection"
150 raise PoolTimeoutError(
151 msg,
152 ) from error
154 try:
155 opened_connection = await open_sqlite_connection(self.database_path)
156 except Exception:
157 async with self.condition:
158 self.opening_connections -= 1
159 self.condition.notify_all()
160 raise
161 async with self.condition:
162 self.opening_connections -= 1
163 if not self.closed and not self.closing:
164 self.active_connections += 1
165 self.condition.notify_all()
166 self.logger.debug(
167 "connection acquired",
168 backend="sqlite",
169 source="opened",
170 )
171 return opened_connection
172 await close_sqlite_connection(opened_connection)
173 self.check_accepting_work()
174 msg = "database is closing"
175 raise DatabaseClosingError(msg)
177 async def release(self, connection: Connection) -> None:
178 """Return a checked-out connection or close it during shutdown."""
180 should_close = False
181 async with self.condition:
182 self.active_connections -= 1
183 if self.closed or self.closing:
184 should_close = True
185 else:
186 self.idle_connections.append(connection)
187 self.condition.notify_all()
188 self.logger.debug(
189 "connection released",
190 backend="sqlite",
191 closed=should_close,
192 )
193 if should_close:
194 await close_sqlite_connection(connection)
196 async def close(self, close_timeout: NonNegativeFloat, /) -> None:
197 """Close idle connections and wait for checked-out work to finish."""
199 self.logger.debug("database close started", backend="sqlite")
200 async with self.condition:
201 if self.closed:
202 self.logger.debug("database close skipped", backend="sqlite")
203 return
204 if self.closing:
205 msg = "database is already closing"
206 raise DatabaseClosingError(msg)
207 self.closing = True
208 idle_connections = list(self.idle_connections)
209 self.idle_connections.clear()
210 self.condition.notify_all()
211 await self.close_connections(idle_connections)
213 event_loop = asyncio.get_running_loop()
214 deadline = event_loop.time() + close_timeout
215 while True:
216 async with self.condition:
217 if self.active_connections == 0 and self.opening_connections == 0:
218 remaining_idle_connections = list(self.idle_connections)
219 self.idle_connections.clear()
220 self.closed = True
221 self.closing = False
222 self.condition.notify_all()
223 break
224 remaining_timeout = deadline - event_loop.time()
225 if remaining_timeout <= 0:
226 self.closing = False
227 self.condition.notify_all()
228 self.logger.warning("database close timed out", backend="sqlite")
229 msg = "database close timed out"
230 raise DatabaseCloseTimeoutError(msg)
231 try:
232 _ = await asyncio.wait_for(
233 self.condition.wait(),
234 timeout=remaining_timeout,
235 )
236 except TimeoutError as error:
237 self.closing = False
238 self.condition.notify_all()
239 self.logger.warning("database close timed out", backend="sqlite")
240 msg = "database close timed out"
241 raise DatabaseCloseTimeoutError(msg) from error
242 await self.close_connections(remaining_idle_connections)
243 self.logger.debug("database close completed", backend="sqlite")
245 def connection_count(self) -> int:
246 """Return all open, checked-out, and currently opening connections."""
248 return (
249 len(self.idle_connections)
250 + self.active_connections
251 + self.opening_connections
252 )
254 @staticmethod
255 async def close_connections(connections: Sequence[Connection]) -> None:
256 """Close a sequence of SQLite connections."""
258 for connection in connections:
259 await close_sqlite_connection(connection)