Rough implementation of the raw driver call

Author: gvbgduh

databases/backends/mysql.py

@@ -10,7 +10,7 @@
 from sqlalchemy.sql import ClauseElement
 from sqlalchemy.types import TypeEngine
 
-from databases.core import DatabaseURL
+from databases.core import DatabaseURL, NoBackendMethod
 from databases.interfaces import ConnectionBackend, DatabaseBackend, TransactionBackend
 
 logger = logging.getLogger("databases")
@@ -135,6 +135,29 @@ async def execute_many(self, query: ClauseElement, values: list) -> None:
         finally:
             await cursor.close()
 
+    async def raw_api_call(self, method: str, *args: typing.Any, **kwargs: typing.Any) -> typing.Any:
+        """
+        NOTE: highly experimental, seems to be a dead-end for generalized solution 
+        """
+        assert self._connection is not None, "Connection is not acquired"
+        cursor = await self._connection.cursor()
+        try:
+            if 'execute' not in method:
+                await cursor.execute(*args, **kwargs)
+                api_method = getattr(cursor, method)
+                return await api_method()
+            else:
+                api_method = getattr(cursor, method)
+                return await api_method(*args, **kwargs)
+        except AttributeError:
+            raise NoBackendMethod(f'{self.database._backend._dialect.driver} has no "{method}" implemented.')
+        finally:
+            await cursor.close()
+
+    async def expose_backend_connection(self) -> aiomysql.connection.Connection:
+        assert self._connection is not None, "Connection is not acquired"
+        return self._connection
+
     async def iterate(
         self, query: ClauseElement
     ) -> typing.AsyncGenerator[typing.Any, None]:

databases/backends/postgres.py

@@ -7,7 +7,7 @@
 from sqlalchemy.engine.interfaces import Dialect
 from sqlalchemy.sql import ClauseElement
 
-from databases.core import DatabaseURL
+from databases.core import DatabaseURL, NoBackendMethod
 from databases.interfaces import ConnectionBackend, DatabaseBackend, TransactionBackend
 
 logger = logging.getLogger("databases")
@@ -93,6 +93,8 @@ def __iter__(self) -> typing.Iterator:
     def __len__(self) -> int:
         return len(self._column_map)
 
+    # TODO (?) Nice __repr__
+
 
 class PostgresConnection(ConnectionBackend):
     def __init__(self, database: PostgresBackend, dialect: Dialect):
@@ -142,6 +144,21 @@ async def execute_many(self, query: ClauseElement, values: list) -> None:
             single_query, args, result_columns = self._compile(single_query)
             await self._connection.execute(single_query, *args)
 
+    async def raw_api_call(self, method: str, *args: typing.Any, **kwargs: typing.Any) -> typing.Any:
+        """
+        NOTE: highly experimental, seems to be a dead-end for generalized solution 
+        """
+        assert self._connection is not None, "Connection is not acquired"
+        try:
+            api_method = getattr(self._connection, method)
+            return await api_method(*args, **kwargs)
+        except AttributeError:
+            raise NoBackendMethod(f'{self.database._backend._dialect.driver} has no "{method}" implemented.')
+
+    async def expose_backend_connection(self) -> asyncpg.connection.Connection:
+        assert self._connection is not None, "Connection is not acquired"
+        return self._connection
+
     async def iterate(
         self, query: ClauseElement
     ) -> typing.AsyncGenerator[typing.Any, None]:

databases/backends/sqlite.py

@@ -9,7 +9,7 @@
 from sqlalchemy.sql import ClauseElement
 from sqlalchemy.types import TypeEngine
 
-from databases.core import DatabaseURL
+from databases.core import DatabaseURL, NoBackendMethod
 from databases.interfaces import ConnectionBackend, DatabaseBackend, TransactionBackend
 
 logger = logging.getLogger("databases")
@@ -116,6 +116,21 @@ async def execute_many(self, query: ClauseElement, values: list) -> None:
         for value in values:
             await self.execute(query, value)
 
+    async def raw_api_call(self, method: str, *args: typing.Any, **kwargs: typing.Any) -> typing.Any:
+        """
+        NOTE: highly experimental, seems to be a dead-end for generalized solution 
+        """
+        assert self._connection is not None, "Connection is not acquired"
+        try:
+            api_method = getattr(self._connection, method)
+            return await api_method(*args, **kwargs)
+        except AttributeError:
+            raise NoBackendMethod(f'{self.database._backend._dialect.driver} has no "{method}" implemented.')
+
+    async def expose_backend_connection(self) -> aiosqlite.core.Connection:
+        assert self._connection is not None, "Connection is not acquired"
+        return self._connection
+
     async def iterate(
         self, query: ClauseElement
     ) -> typing.AsyncGenerator[typing.Any, None]:

databases/core.py

@@ -104,6 +104,14 @@ async def execute_many(self, query: ClauseElement, values: list) -> None:
         async with self.connection() as connection:
             return await connection.execute_many(query=query, values=values)
 
+    async def raw_api_call(self, method: str, *args: typing.Any, **kwargs: typing.Any) -> typing.Any:
+        async with self.connection() as connection:
+            return await connection.raw_api_call(method, *args, **kwargs)
+
+    async def expose_backend_connection(self) -> typing.Any:
+        async with self.connection() as connection:
+            return await connection.expose_backend_connection()
+
     async def iterate(
         self, query: ClauseElement
     ) -> typing.AsyncGenerator[RowProxy, None]:
@@ -168,6 +176,12 @@ async def execute(self, query: ClauseElement, values: dict = None) -> typing.Any
     async def execute_many(self, query: ClauseElement, values: list) -> None:
         await self._connection.execute_many(query, values)
 
+    async def raw_api_call(self, method: str, *args: typing.Any, **kwargs: typing.Any) -> typing.Any:
+        return await self._connection.raw_api_call(method, *args, **kwargs)
+
+    async def expose_backend_connection(self) -> typing.Any:
+        return await self._connection.expose_backend_connection()
+
     async def iterate(
         self, query: ClauseElement
     ) -> typing.AsyncGenerator[typing.Any, None]:
@@ -348,3 +362,7 @@ def __repr__(self) -> str:
         if self.password:
             url = str(self.replace(password="********"))
         return f"{self.__class__.__name__}({repr(url)})"
+
+
+class NoBackendMethod(Exception):
+    pass

tests/test_databases.py

@@ -514,3 +514,135 @@ async def get_connection_2():
         test_complete.set()
         await task_1
         await task_2
+
+
+@pytest.mark.parametrize("database_url", DATABASE_URLS)
+@async_adapter
+async def test_queries_with_raw_api_call(database_url):
+    """
+    Test that the basic `execute()`, `execute_many()`, `fetch_all()``, and
+    `fetch_one()` interfaces are working as expected being called as raw driver calls.
+    """
+    async with Database(database_url) as database:
+        async with database.transaction(force_rollback=True):
+            # Insert query
+            if str(database_url).startswith('mysql'):
+                insert_query = "INSERT INTO notes (text, completed) VALUES (%s, %s)"
+            else:
+                insert_query = "INSERT INTO notes (text, completed) VALUES ($1, $2)"
+
+            # execute()
+            values = ("example1", True)
+
+            if str(database_url).startswith('postgresql'):
+                await database.raw_api_call('execute', insert_query, *values)
+            else:
+                await database.raw_api_call('execute', insert_query, values)
+
+            # execute_many()
+            values = [("example2", False), ("example3", True)]
+            await database.raw_api_call('executemany', insert_query, values)
+
+            # Select query
+            select_query = "SELECT notes.id, notes.text, notes.completed FROM notes"
+
+            # fetch_all()
+            if str(database_url).startswith('postgresql'):
+                results = await database.raw_api_call('fetch', select_query)
+            elif str(database_url).startswith('mysql'):
+                results = await database.raw_api_call('fetchall', select_query)
+            elif str(database_url).startswith('sqlite'):
+                results = await database.raw_api_call('execute_fetchall', select_query)
+
+            assert len(results) == 3
+            # Raw output for the raw request
+            assert results[0][1] == "example1"
+            assert results[0][2] == True
+            assert results[1][1] == "example2"
+            assert results[1][2] == False
+            assert results[2][1] == "example3"
+            assert results[2][2] == True
+
+            # # fetch_one()
+            # import pdb; pdb.set_trace()
+            # if str(database_url).startswith('postgresql'):
+            #     result = await database.raw_api_call('fetchrow', select_query)
+            # else:
+            #     result = await database.raw_api_call('fetchone', select_query)
+            # assert result[1] == "example1"
+            # assert result[2] == True
+
+
+@pytest.mark.parametrize("database_url", DATABASE_URLS)
+@async_adapter
+async def test_queries_with_expose_backend_connection(database_url):
+    """
+    Replication of `execute()`, `execute_many()`, `fetch_all()``, and
+    `fetch_one()` using the raw driver interface.
+    """
+    async with Database(database_url) as database:
+        async with database.transaction(force_rollback=True):
+            # Insert query
+            if str(database_url).startswith('mysql'):
+                insert_query = "INSERT INTO notes (text, completed) VALUES (%s, %s)"
+            else:
+                insert_query = "INSERT INTO notes (text, completed) VALUES ($1, $2)"
+
+            # execute()
+            values = ("example1", True)
+
+            con = await database.expose_backend_connection()
+
+            if str(database_url).startswith('postgresql'):
+                await con.execute(insert_query, *values)
+            elif str(database_url).startswith('mysql'):
+                cursor = await con.cursor()
+                await cursor.execute(insert_query, values)
+            elif str(database_url).startswith('sqlite'):
+                await con.execute(insert_query, values)
+
+            # execute_many()
+            values = [("example2", False), ("example3", True)]
+            
+            if str(database_url).startswith('mysql'):
+                cursor = await con.cursor()
+                await cursor.executemany(insert_query, values)
+            else:
+                await con.executemany(insert_query, values)
+
+            # Select query
+            select_query = "SELECT notes.id, notes.text, notes.completed FROM notes"
+
+            # fetch_all()
+            if str(database_url).startswith('postgresql'):
+                results = await con.fetch(select_query)
+            elif str(database_url).startswith('mysql'):
+                cursor = await con.cursor()
+                await cursor.execute(select_query)
+                results = await cursor.fetchall()
+            elif str(database_url).startswith('sqlite'):
+                results = await con.execute_fetchall(select_query)
+
+            assert len(results) == 3
+            # Raw output for the raw request
+            assert results[0][1] == "example1"
+            assert results[0][2] == True
+            assert results[1][1] == "example2"
+            assert results[1][2] == False
+            assert results[2][1] == "example3"
+            assert results[2][2] == True
+
+            # fetch_one()
+            if str(database_url).startswith('postgresql'):
+                result = await con.fetchrow(select_query)
+            else:
+                cursor = await con.cursor()
+                await cursor.execute(select_query)
+                result = await cursor.fetchone()
+            
+            # Raw output for the raw request
+            assert result[1] == "example1"
+            assert result[2] == True
+
+# TODO unittests at the connection level
+# TODO (?) Double-check connections are released back to the pool