Added support for raw queries using sqlalchemy.text()

Author: pvanliefland

databases/backends/mysql.py

@@ -112,10 +112,8 @@ async def fetch_one(self, query: ClauseElement) -> typing.Optional[typing.Mappin
         finally:
             await cursor.close()
 
-    async def execute(self, query: ClauseElement, values: dict = None) -> typing.Any:
+    async def execute(self, query: ClauseElement) -> typing.Any:
         assert self._connection is not None, "Connection is not acquired"
-        if values is not None:
-            query = query.values(values)
         query, args, context = self._compile(query)
         cursor = await self._connection.cursor()
         try:
@@ -124,12 +122,11 @@ async def execute(self, query: ClauseElement, values: dict = None) -> typing.Any
         finally:
             await cursor.close()
 
-    async def execute_many(self, query: ClauseElement, values: list) -> None:
+    async def execute_many(self, queries: typing.List[ClauseElement]) -> None:
         assert self._connection is not None, "Connection is not acquired"
         cursor = await self._connection.cursor()
         try:
-            for item in values:
-                single_query = query.values(item)
+            for single_query in queries:
                 single_query, args, context = self._compile(single_query)
                 await cursor.execute(single_query, args)
         finally:

databases/backends/postgres.py

@@ -66,7 +66,9 @@ def connection(self) -> "PostgresConnection":
 
 
 class Record(Mapping):
-    def __init__(self, row: tuple, result_columns: tuple, dialect: Dialect) -> None:
+    def __init__(
+        self, row: asyncpg.Record, result_columns: tuple, dialect: Dialect
+    ) -> None:
         self._row = row
         self._result_columns = result_columns
         self._dialect = dialect
@@ -80,11 +82,14 @@ def __init__(self, row: tuple, result_columns: tuple, dialect: Dialect) -> None:
         }
 
     def __getitem__(self, key: typing.Any) -> typing.Any:
-        if type(key) is Column:
-            idx, datatype = self._column_map_full[str(key)]
-        else:
-            idx, datatype = self._column_map[key]
-        raw = self._row[idx]
+        if len(self._column_map) == 0:  # raw query
+            return self._row[tuple(self._row.keys()).index(key)]
+        if len(self._column_map) > 0:
+            if type(key) is Column:
+                idx, datatype = self._column_map_full[str(key)]
+            else:
+                idx, datatype = self._column_map[key]
+            raw = self._row[idx]
         try:
             processor = _result_processors[datatype]
         except KeyError:
@@ -133,20 +138,17 @@ async def fetch_one(self, query: ClauseElement) -> typing.Optional[typing.Mappin
             return None
         return Record(row, result_columns, self._dialect)
 
-    async def execute(self, query: ClauseElement, values: dict = None) -> typing.Any:
+    async def execute(self, query: ClauseElement) -> typing.Any:
         assert self._connection is not None, "Connection is not acquired"
-        if values is not None:
-            query = query.values(values)
         query, args, result_columns = self._compile(query)
         return await self._connection.fetchval(query, *args)
 
-    async def execute_many(self, query: ClauseElement, values: list) -> None:
+    async def execute_many(self, queries: typing.List[ClauseElement]) -> None:
         assert self._connection is not None, "Connection is not acquired"
         # asyncpg uses prepared statements under the hood, so we just
         # loop through multiple executes here, which should all end up
         # using the same prepared statement.
-        for item in values:
-            single_query = query.values(item)
+        for single_query in queries:
             single_query, args, result_columns = self._compile(single_query)
             await self._connection.execute(single_query, *args)
 

databases/backends/sqlite.py

@@ -102,19 +102,17 @@ async def fetch_one(self, query: ClauseElement) -> typing.Optional[typing.Mappin
             metadata = ResultMetaData(context, cursor.description)
             return RowProxy(metadata, row, metadata._processors, metadata._keymap)
 
-    async def execute(self, query: ClauseElement, values: dict = None) -> typing.Any:
+    async def execute(self, query: ClauseElement) -> typing.Any:
         assert self._connection is not None, "Connection is not acquired"
-        if values is not None:
-            query = query.values(values)
         query, args, context = self._compile(query)
         cursor = await self._connection.execute(query, args)
         await cursor.close()
         return cursor.lastrowid
 
-    async def execute_many(self, query: ClauseElement, values: list) -> None:
+    async def execute_many(self, queries: typing.List[ClauseElement]) -> None:
         assert self._connection is not None, "Connection is not acquired"
-        for value in values:
-            await self.execute(query, value)
+        for single_query in queries:
+            await self.execute(single_query)
 
     async def iterate(
         self, query: ClauseElement

databases/core.py

@@ -5,6 +5,7 @@
 from types import TracebackType
 from urllib.parse import SplitResult, parse_qsl, urlsplit
 
+from sqlalchemy import text
 from sqlalchemy.engine import RowProxy
 from sqlalchemy.sql import ClauseElement
 
@@ -88,27 +89,36 @@ async def __aexit__(
     ) -> None:
         await self.disconnect()
 
-    async def fetch_all(self, query: ClauseElement) -> typing.List[RowProxy]:
+    async def fetch_all(
+        self, query: typing.Union[ClauseElement, str], values: dict = None
+    ) -> typing.List[RowProxy]:
         async with self.connection() as connection:
-            return await connection.fetch_all(query=query)
+            return await connection.fetch_all(query=self._build_query(query, values))
 
-    async def fetch_one(self, query: ClauseElement) -> RowProxy:
+    async def fetch_one(
+        self, query: typing.Union[ClauseElement, str], values: dict = None
+    ) -> RowProxy:
         async with self.connection() as connection:
-            return await connection.fetch_one(query=query)
+            return await connection.fetch_one(query=self._build_query(query, values))
 
-    async def execute(self, query: ClauseElement, values: dict = None) -> typing.Any:
+    async def execute(
+        self, query: typing.Union[ClauseElement, str], values: dict = None
+    ) -> typing.Any:
         async with self.connection() as connection:
-            return await connection.execute(query=query, values=values)
+            return await connection.execute(self._build_query(query, values))
 
-    async def execute_many(self, query: ClauseElement, values: list) -> None:
+    async def execute_many(
+        self, query: typing.Union[ClauseElement, str], values: list
+    ) -> None:
         async with self.connection() as connection:
-            return await connection.execute_many(query=query, values=values)
+            queries = [self._build_query(query, values_set) for values_set in values]
+            return await connection.execute_many(queries)
 
     async def iterate(
-        self, query: ClauseElement
+        self, query: typing.Union[ClauseElement, str], values: dict = None
     ) -> typing.AsyncGenerator[RowProxy, None]:
         async with self.connection() as connection:
-            async for record in connection.iterate(query):
+            async for record in connection.iterate(self._build_query(query, values)):
                 yield record
 
     def connection(self) -> "Connection":
@@ -125,6 +135,19 @@ def connection(self) -> "Connection":
     def transaction(self, *, force_rollback: bool = False) -> "Transaction":
         return self.connection().transaction(force_rollback=force_rollback)
 
+    @staticmethod
+    def _build_query(
+        query: typing.Union[ClauseElement, str], values: dict = None
+    ) -> ClauseElement:
+        if isinstance(query, str):
+            query = text(query)
+
+            return query.bindparams(**values) if values is not None else query
+        elif values:
+            return query.values(**values)
+
+        return query
+
 
 class Connection:
     def __init__(self, backend: DatabaseBackend) -> None:
@@ -162,11 +185,11 @@ async def fetch_all(self, query: ClauseElement) -> typing.Any:
     async def fetch_one(self, query: ClauseElement) -> typing.Any:
         return await self._connection.fetch_one(query=query)
 
-    async def execute(self, query: ClauseElement, values: dict = None) -> typing.Any:
-        return await self._connection.execute(query, values)
+    async def execute(self, query: ClauseElement) -> typing.Any:
+        return await self._connection.execute(query)
 
-    async def execute_many(self, query: ClauseElement, values: list) -> None:
-        await self._connection.execute_many(query, values)
+    async def execute_many(self, queries: typing.List[ClauseElement]) -> None:
+        await self._connection.execute_many(queries)
 
     async def iterate(
         self, query: ClauseElement

databases/interfaces.py

@@ -27,10 +27,10 @@ async def fetch_all(self, query: ClauseElement) -> typing.List[typing.Mapping]:
     async def fetch_one(self, query: ClauseElement) -> typing.Optional[typing.Mapping]:
         raise NotImplementedError()  # pragma: no cover
 
-    async def execute(self, query: ClauseElement, values: dict = None) -> typing.Any:
+    async def execute(self, query: ClauseElement) -> typing.Any:
         raise NotImplementedError()  # pragma: no cover
 
-    async def execute_many(self, query: ClauseElement, values: list) -> None:
+    async def execute_many(self, queries: typing.List[ClauseElement]) -> None:
         raise NotImplementedError()  # pragma: no cover
 
     async def iterate(

tests/test_databases.py

@@ -103,7 +103,7 @@ def run_sync(*args, **kwargs):
 async def test_queries(database_url):
     """
     Test that the basic `execute()`, `execute_many()`, `fetch_all()``, and
-    `fetch_one()` interfaces are all supported.
+    `fetch_one()` interfaces are all supported (using SQLAlchemy core).
     """
     async with Database(database_url) as database:
         async with database.transaction(force_rollback=True):
@@ -151,6 +151,57 @@ async def test_queries(database_url):
             assert iterate_results[2]["completed"] == True
 
 
+@pytest.mark.parametrize("database_url", DATABASE_URLS)
+@async_adapter
+async def test_queries_raw(database_url):
+    """
+    Test that the basic `execute()`, `execute_many()`, `fetch_all()``, and
+    `fetch_one()` interfaces are all supported (raw queries).
+    """
+    async with Database(database_url) as database:
+        async with database.transaction(force_rollback=True):
+            # execute()
+            query = "INSERT INTO notes(text, completed) VALUES (:text, :completed)"
+            values = {"text": "example1", "completed": True}
+            await database.execute(query, values)
+
+            # execute_many()
+            query = "INSERT INTO notes(text, completed) VALUES (:text, :completed)"
+            values = [
+                {"text": "example2", "completed": False},
+                {"text": "example3", "completed": True},
+            ]
+            await database.execute_many(query, values)
+
+            # fetch_all()
+            query = "SELECT * FROM notes WHERE completed = :completed"
+            results = await database.fetch_all(query=query, values={"completed": True})
+            assert len(results) == 2
+            assert results[0]["text"] == "example1"
+            assert results[0]["completed"] == True
+            assert results[1]["text"] == "example3"
+            assert results[1]["completed"] == True
+
+            # fetch_one()
+            query = "SELECT * FROM notes WHERE completed = :completed"
+            result = await database.fetch_one(query=query, values={"completed": False})
+            assert result["text"] == "example2"
+            assert result["completed"] == False
+
+            # iterate()
+            query = "SELECT * FROM notes"
+            iterate_results = []
+            async for result in database.iterate(query=query):
+                iterate_results.append(result)
+            assert len(iterate_results) == 3
+            assert iterate_results[0]["text"] == "example1"
+            assert iterate_results[0]["completed"] == True
+            assert iterate_results[1]["text"] == "example2"
+            assert iterate_results[1]["completed"] == False
+            assert iterate_results[2]["text"] == "example3"
+            assert iterate_results[2]["completed"] == True
+
+
 @pytest.mark.parametrize("database_url", DATABASE_URLS)
 @async_adapter
 async def test_results_support_mapping_interface(database_url):
@@ -234,8 +285,8 @@ async def test_execute_return_val(database_url):
             query = notes.insert()
             values = {"text": "example1", "completed": True}
             pk = await database.execute(query, values)
-
             assert isinstance(pk, int)
+
             query = notes.select().where(notes.c.id == pk)
             result = await database.fetch_one(query)
             assert result["text"] == "example1"