Merge pull request #58 from gvbgduh/raw_driver_call

Support for `connection.raw_connection`

Author: lovelydinosaur

databases/__init__.py

@@ -1,5 +1,4 @@
 from databases.core import Database, DatabaseURL
 
-
 __version__ = "0.1.9"
 __all__ = ["Database", "DatabaseURL"]

databases/backends/mysql.py

@@ -172,6 +172,11 @@ def _compile(
         logger.debug("Query: %s\nArgs: %s", compiled.string, args)
         return compiled.string, args, CompilationContext(execution_context)
 
+    @property
+    def raw_connection(self) -> aiomysql.connection.Connection:
+        assert self._connection is not None, "Connection is not acquired"
+        return self._connection
+
 
 class MySQLTransaction(TransactionBackend):
     def __init__(self, connection: MySQLConnection):

databases/backends/postgres.py

@@ -171,6 +171,11 @@ def _compile(self, query: ClauseElement) -> typing.Tuple[str, list, tuple]:
         logger.debug("Query: %s\nArgs: %s", compiled_query, args)
         return compiled_query, args, compiled._result_columns
 
+    @property
+    def raw_connection(self) -> asyncpg.connection.Connection:
+        assert self._connection is not None, "Connection is not acquired"
+        return self._connection
+
 
 class PostgresTransaction(TransactionBackend):
     def __init__(self, connection: PostgresConnection):

databases/backends/sqlite.py

@@ -153,6 +153,11 @@ def _compile(
         logger.debug("Query: %s\nArgs: %s", compiled.string, args)
         return compiled.string, args, CompilationContext(execution_context)
 
+    @property
+    def raw_connection(self) -> aiosqlite.core.Connection:
+        assert self._connection is not None, "Connection is not acquired"
+        return self._connection
+
 
 class SQLiteTransaction(TransactionBackend):
     def __init__(self, connection: SQLiteConnection):

databases/core.py

@@ -177,6 +177,10 @@ async def iterate(
     def transaction(self, *, force_rollback: bool = False) -> "Transaction":
         return Transaction(self, force_rollback)
 
+    @property
+    def raw_connection(self) -> typing.Any:
+        return self._connection.raw_connection
+
 
 class Transaction:
     def __init__(self, connection: Connection, force_rollback: bool) -> None:

databases/interfaces.py

@@ -44,6 +44,10 @@ async def iterate(
     def transaction(self) -> "TransactionBackend":
         raise NotImplementedError()  # pragma: no cover
 
+    @property
+    def raw_connection(self) -> typing.Any:
+        raise NotImplementedError()  # pragma: no cover
+
 
 class TransactionBackend:
     async def start(self, is_root: bool) -> None:

tests/test_databases.py

@@ -514,3 +514,90 @@ 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_connection_context_with_raw_connection(database_url):
+    """
+    Test connection contexts with respect to the raw connection.
+    """
+    async with Database(database_url) as database:
+        async with database.connection() as connection_1:
+            async with database.connection() as connection_2:
+                assert connection_1 is connection_2
+                assert connection_1.raw_connection is connection_2.raw_connection
+
+
+@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.connection() as connection:
+            async with database.transaction(force_rollback=True):
+                # Get the raw connection
+                raw_connection = connection.raw_connection
+
+                # 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 raw_connection.execute(insert_query, *values)
+                elif str(database_url).startswith("mysql"):
+                    cursor = await raw_connection.cursor()
+                    await cursor.execute(insert_query, values)
+                elif str(database_url).startswith("sqlite"):
+                    await raw_connection.execute(insert_query, values)
+
+                # execute_many()
+                values = [("example2", False), ("example3", True)]
+
+                if str(database_url).startswith("mysql"):
+                    cursor = await raw_connection.cursor()
+                    await cursor.executemany(insert_query, values)
+                else:
+                    await raw_connection.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 raw_connection.fetch(select_query)
+                elif str(database_url).startswith("mysql"):
+                    cursor = await raw_connection.cursor()
+                    await cursor.execute(select_query)
+                    results = await cursor.fetchall()
+                elif str(database_url).startswith("sqlite"):
+                    results = await raw_connection.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 raw_connection.fetchrow(select_query)
+                else:
+                    cursor = await raw_connection.cursor()
+                    await cursor.execute(select_query)
+                    result = await cursor.fetchone()
+
+                # Raw output for the raw request
+                assert result[1] == "example1"
+                assert result[2] == True