Wait for transaction queue to complete when closing connection (#305)

* Test await connection.close() always returns an empty transaction queue

* 'await connection.close()' returns once all transaction queue items's results have been forwarded, including the _STOP_RUNNING_SENTINEL result

* nits

---------

Co-authored-by: Amethyst Reese <amethyst@n7.gg>

Author: davidandreoletti

aiosqlite/core.py

@@ -62,10 +62,15 @@ def __init__(
                 DeprecationWarning,
             )
 
-    def _stop_running(self):
+    def _stop_running(self) -> asyncio.Future:
         self._running = False
-        # PEP 661 is not accepted yet, so we cannot type a sentinel
-        self._tx.put_nowait(_STOP_RUNNING_SENTINEL)  # type: ignore[arg-type]
+
+        function = partial(lambda: _STOP_RUNNING_SENTINEL)
+        future = asyncio.get_event_loop().create_future()
+
+        self._tx.put_nowait((future, function))
+
+        return future
 
     @property
     def _conn(self) -> sqlite3.Connection:
@@ -95,16 +100,17 @@ def run(self) -> None:
             # futures)
 
             tx_item = self._tx.get()
-            if tx_item is _STOP_RUNNING_SENTINEL:
-                break
-
             future, function = tx_item
 
             try:
                 LOG.debug("executing %s", function)
                 result = function()
                 LOG.debug("operation %s completed", function)
                 future.get_loop().call_soon_threadsafe(set_result, future, result)
+
+                if result is _STOP_RUNNING_SENTINEL:
+                    break
+
             except BaseException as e:  # noqa B036
                 LOG.debug("returning exception %s", e)
                 future.get_loop().call_soon_threadsafe(set_exception, future, e)
@@ -129,7 +135,7 @@ async def _connect(self) -> "Connection":
                 self._tx.put_nowait((future, self._connector))
                 self._connection = await future
             except BaseException:
-                self._stop_running()
+                await self._stop_running()
                 self._connection = None
                 raise
 
@@ -170,7 +176,7 @@ async def close(self) -> None:
             LOG.info("exception occurred while closing connection")
             raise
         finally:
-            self._stop_running()
+            await self._stop_running()
             self._connection = None
 
     @contextmanager

aiosqlite/tests/smoke.py

@@ -414,6 +414,17 @@ async def test_cursor_on_closed_connection_loop(self):
             except sqlite3.ProgrammingError:
                 pass
 
+    async def test_close_blocking_until_transaction_queue_empty(self):
+        db = await aiosqlite.connect(self.db)
+        # Insert transactions into the
+        # transaction queue '_tx'
+        for i in range(1000):
+            await db.execute(f"select 1, {i}")
+        # Wait for all transactions to complete
+        await db.close()
+        # Check no more transaction pending
+        self.assertEqual(db._tx.empty(), True)
+
     async def test_close_twice(self):
         db = await aiosqlite.connect(self.db)