33import asyncio
44import threading
55import warnings
6+ import weakref
67from concurrent .futures import ThreadPoolExecutor
78from typing import TYPE_CHECKING , Any
89
@@ -51,6 +52,14 @@ def __init__(
5152 self .connections = [Connection (database ) for _ in range (max_concurrent_queries )]
5253 self .connections_counter = [0 for _ in range (max_concurrent_queries )]
5354 self .lock = threading .Lock ()
55+ # Calls on one connection run one at a time (the engine serializes them
56+ # anyway), so the call holding a connection's lock is the one an
57+ # interrupt would hit. Queries run directly on a connection obtained
58+ # from acquire_connection() are not tracked.
59+ self ._connection_locks : weakref .WeakKeyDictionary [
60+ Connection , threading .Lock
61+ ] = weakref .WeakKeyDictionary ()
62+ self ._running : dict [Connection , _Call ] = {}
5463
5564 for conn in self .connections :
5665 conn .init_connection ()
@@ -114,10 +123,9 @@ def release_connection(self, conn: Connection) -> None:
114123
115124
116125 """
117- for i , existing_conn in enumerate (self .connections ):
118- if existing_conn == conn :
119- self .__decrement_connection_counter (i )
120- break
126+ conn_index = self .__index_of (conn )
127+ if conn_index is not None :
128+ self .__decrement_connection_counter (conn_index )
121129
122130 def set_query_timeout (self , timeout_in_ms : int ) -> None :
123131 """
@@ -153,28 +161,22 @@ async def execute(
153161 QueryResult
154162 Query result.
155163
164+ Raises
165+ ------
166+ asyncio.CancelledError
167+ If the awaiting task is cancelled. A running query is interrupted; a
168+ query still waiting for a free connection is not started.
169+
156170 """
157- loop = asyncio .get_running_loop ()
158- # If the query is a prepared statement, use the connection associated with it
159171 if isinstance (query , PreparedStatement ):
160172 conn = query ._connection
161- for i , existing_conn in enumerate (self .connections ):
162- if existing_conn == conn :
163- conn_index = i
164- with self .lock :
165- self .connections_counter [conn_index ] += 1
166- break
173+ conn_index = self .__index_of (conn )
174+ if conn_index is not None :
175+ with self .lock :
176+ self .connections_counter [conn_index ] += 1
167177 else :
168178 conn , conn_index = self .__get_connection_with_least_queries ()
169-
170- try :
171- return await loop .run_in_executor (
172- self .executor , conn .execute , query , parameters
173- )
174- except asyncio .CancelledError :
175- conn .interrupt ()
176- finally :
177- self .__decrement_connection_counter (conn_index )
179+ return await self .__submit (conn , conn_index , conn .execute , query , parameters )
178180
179181 async def _prepare (
180182 self , query : str , parameters : dict [str , Any ] | None = None
@@ -183,16 +185,64 @@ async def _prepare(
183185 The only parameters supported during prepare are dataframes.
184186 Any remaining parameters will be ignored and should be passed to execute().
185187 """ # noqa: D401
186- loop = asyncio .get_running_loop ()
187188 conn , conn_index = self .__get_connection_with_least_queries ()
189+ return await self .__submit (conn , conn_index , conn ._prepare , query , parameters )
188190
191+ def __index_of (self , conn : Connection ) -> int | None :
192+ for i , existing_conn in enumerate (self .connections ):
193+ if existing_conn is conn :
194+ return i
195+ return None
196+
197+ async def __submit (
198+ self , conn : Connection , conn_index : int | None , fn : Any , * args : Any
199+ ) -> Any :
200+ call = _Call ()
189201 try :
190- preparedStatement = await loop .run_in_executor (
191- self .executor , conn .prepare , query , parameters
202+ future = self .executor .submit (self .__run , conn , call , fn , * args )
203+ except RuntimeError :
204+ if conn_index is not None :
205+ self .__decrement_connection_counter (conn_index )
206+ raise
207+ if conn_index is not None :
208+ # Release the slot when the work has actually finished (or was
209+ # cancelled before it started), not when the awaiting task exits.
210+ future .add_done_callback (
211+ lambda _ : self .__decrement_connection_counter (conn_index )
192212 )
193- return preparedStatement
194- finally :
195- self .__decrement_connection_counter (conn_index )
213+ try :
214+ return await asyncio .wrap_future (future )
215+ except asyncio .CancelledError :
216+ self .__cancel (conn , call )
217+ raise
218+
219+ def __run (self , conn : Connection , call : _Call , fn : Any , * args : Any ) -> Any :
220+ with self .lock :
221+ connection_lock = self ._connection_locks .setdefault (conn , threading .Lock ())
222+ with connection_lock :
223+ with self .lock :
224+ if call .cancelled :
225+ return None
226+ self ._running [conn ] = call
227+ try :
228+ return fn (* args )
229+ finally :
230+ with self .lock :
231+ del self ._running [conn ]
232+
233+ def __cancel (self , conn : Connection , call : _Call ) -> None :
234+ # Interrupting acts on whatever the connection is running, so only
235+ # interrupt while this call is the one running. A call still waiting
236+ # for its connection is skipped when its turn comes.
237+ with self .lock :
238+ call .cancelled = True
239+ running = self ._running .get (conn ) is call
240+ if running :
241+ conn .interrupt ()
242+ if running :
243+ # An interrupt that arrives while the query is still compiling is
244+ # cleared when execution starts, so repeat it until the call returns.
245+ asyncio .get_running_loop ().call_later (0.05 , self .__cancel , conn , call )
196246
197247 async def prepare (
198248 self , query : str , parameters : dict [str , Any ] | None = None
@@ -225,10 +275,17 @@ def close(self) -> None:
225275 """
226276 Close all connections and shutdown the thread pool.
227277
278+ Queries still running are interrupted and fail with "Interrupted.".
279+
228280 Note: Call to this method is optional. The connections and thread pool
229281 will be closed automatically when the instance is garbage collected.
230282 """
231283 for conn in self .connections :
232284 conn .close ()
233285
234286 self .executor .shutdown (wait = True )
287+
288+
289+ class _Call :
290+ def __init__ (self ) -> None :
291+ self .cancelled = False
0 commit comments