Query level locking

Author: lovelydinosaur

databases/core.py

@@ -160,6 +160,8 @@ def __init__(self, backend: DatabaseBackend) -> None:
         self._transaction_lock = asyncio.Lock()
         self._transaction_stack = []  # type: typing.List[Transaction]
 
+        self._query_lock = asyncio.Lock()
+
     async def __aenter__(self) -> "Connection":
         async with self._connection_lock:
             self._connection_counter += 1
@@ -182,38 +184,44 @@ async def __aexit__(
     async def fetch_all(
         self, query: typing.Union[ClauseElement, str], values: dict = None
     ) -> typing.List[typing.Mapping]:
-        return await self._connection.fetch_all(self._build_query(query, values))
+        async with self._self._query_lock:
+            return await self._connection.fetch_all(self._build_query(query, values))
 
     async def fetch_one(
         self, query: typing.Union[ClauseElement, str], values: dict = None
     ) -> typing.Optional[typing.Mapping]:
-        return await self._connection.fetch_one(self._build_query(query, values))
+        async with self._self._query_lock:
+            return await self._connection.fetch_one(self._build_query(query, values))
 
     async def fetch_val(
         self,
         query: typing.Union[ClauseElement, str],
         values: dict = None,
         column: typing.Any = 0,
     ) -> typing.Any:
-        row = await self._connection.fetch_one(self._build_query(query, values))
+        async with self._self._query_lock:
+            row = await self._connection.fetch_one(self._build_query(query, values))
         return None if row is None else row[column]
 
     async def execute(
         self, query: typing.Union[ClauseElement, str], values: dict = None
     ) -> typing.Any:
-        return await self._connection.execute(self._build_query(query, values))
+        async with self._self._query_lock:
+            return await self._connection.execute(self._build_query(query, values))
 
     async def execute_many(
         self, query: typing.Union[ClauseElement, str], values: list
     ) -> None:
         queries = [self._build_query(query, values_set) for values_set in values]
-        await self._connection.execute_many(queries)
+        async with self._self._query_lock:
+            await self._connection.execute_many(queries)
 
     async def iterate(
         self, query: typing.Union[ClauseElement, str], values: dict = None
     ) -> typing.AsyncGenerator[typing.Any, None]:
-        async for record in self._connection.iterate(self._build_query(query, values)):
-            yield record
+        async with self._self._query_lock:
+            async for record in self._connection.iterate(self._build_query(query, values)):
+                yield record
 
     def transaction(self, *, force_rollback: bool = False) -> "Transaction":
         return Transaction(self, force_rollback)