Merge pull request #108 from encode/query-locking

Query 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,46 @@ 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._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._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._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._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._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._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)

tests/test_databases.py

@@ -699,3 +699,15 @@ async def test_database_url_interface(database_url):
     async with Database(database_url) as database:
         assert isinstance(database.url, DatabaseURL)
         assert database.url == database_url
+
+
+@pytest.mark.parametrize("database_url", DATABASE_URLS)
+@async_adapter
+async def test_concurrent_access_on_single_connection(database_url):
+    async with Database(database_url, force_rollback=True) as database:
+
+        async def db_lookup():
+            if str(database_url).startswith("postgresql"):
+                await database.fetch_one("SELECT pg_sleep(1)")
+
+        await asyncio.gather(db_lookup(), db_lookup())