Some refectoring and polishing for raw_connection #9

Author: gvbgduh

databases/__init__.py

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

databases/backends/mysql.py

@@ -135,10 +135,6 @@ async def execute_many(self, query: ClauseElement, values: list) -> None:
         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]:
@@ -153,6 +149,10 @@ async def iterate(
         finally:
             await cursor.close()
 
+    async def raw_connection(self) -> aiomysql.connection.Connection:
+        assert self._connection is not None, "Connection is not acquired"
+        return self._connection
+
     def transaction(self) -> TransactionBackend:
         return MySQLTransaction(self)
 

databases/backends/postgres.py

@@ -74,21 +74,18 @@ def __init__(self, row: tuple, result_columns: tuple, dialect: Dialect) -> None:
             for idx, (column_name, _, _, datatype) in enumerate(self._result_columns)
         }
 
-    def __getitem__(self, key: typing.Any) -> typing.Any:
+    def __getitem__(self, key: str) -> typing.Any:
+        idx, datatype = self._column_map[key]
+        raw = self._row[idx]
         try:
-            idx, datatype = self._column_map[key]
-            raw = self._row[idx]
-            try:
-                processor = _result_processors[datatype]
-            except KeyError:
-                processor = datatype.result_processor(self._dialect, None)
-                _result_processors[datatype] = processor
-
-            if processor is not None:
-                return processor(raw)
-            return raw
+            processor = _result_processors[datatype]
         except KeyError:
-            return self._row[key]
+            processor = datatype.result_processor(self._dialect, None)
+            _result_processors[datatype] = processor
+
+        if processor is not None:
+            return processor(raw)
+        return raw
 
     def __iter__(self) -> typing.Iterator:
         return iter(self._column_map)
@@ -145,10 +142,6 @@ 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 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]:
@@ -157,6 +150,10 @@ async def iterate(
         async for row in self._connection.cursor(query, *args):
             yield Record(row, result_columns, self._dialect)
 
+    async def raw_connection(self) -> asyncpg.connection.Connection:
+        assert self._connection is not None, "Connection is not acquired"
+        return self._connection
+
     def transaction(self) -> TransactionBackend:
         return PostgresTransaction(connection=self)
 

databases/backends/sqlite.py

@@ -116,7 +116,7 @@ async def execute_many(self, query: ClauseElement, values: list) -> None:
         for value in values:
             await self.execute(query, value)
 
-    async def expose_backend_connection(self) -> aiosqlite.core.Connection:
+    async def raw_connection(self) -> aiosqlite.core.Connection:
         assert self._connection is not None, "Connection is not acquired"
         return self._connection
 

databases/core.py

@@ -104,17 +104,17 @@ 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 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]:
         async with self.connection() as connection:
             async for record in connection.iterate(query):
                 yield record
 
+    async def raw_connection(self) -> typing.Any:
+        async with self.connection() as connection:
+            return await connection.raw_connection()
+
     def connection(self) -> "Connection":
         if self._global_connection is not None:
             return self._global_connection
@@ -172,8 +172,8 @@ 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 expose_backend_connection(self) -> typing.Any:
-        return await self._connection.expose_backend_connection()
+    async def raw_connection(self) -> typing.Any:
+        return await self._connection.raw_connection()
 
     async def iterate(
         self, query: ClauseElement

databases/interfaces.py

@@ -33,9 +33,6 @@ async def execute(self, query: ClauseElement, values: dict = None) -> typing.Any
     async def execute_many(self, query: ClauseElement, values: list) -> None:
         raise NotImplementedError()  # pragma: no cover
 
-    async def expose_backend_connection(self) -> typing.Any:
-        raise NotImplementedError()  # pragma: no cover
-
     async def iterate(
         self, query: ClauseElement
     ) -> typing.AsyncGenerator[typing.Mapping, None]:
@@ -44,6 +41,9 @@ async def iterate(
         # https://github.com/python/mypy/issues/5385#issuecomment-407281656
         yield True  # pragma: no cover
 
+    async def raw_connection(self) -> typing.Any:
+        raise NotImplementedError()  # pragma: no cover
+
     def transaction(self) -> "TransactionBackend":
         raise NotImplementedError()  # pragma: no cover
 

tests/test_databases.py

@@ -516,6 +516,25 @@ async def get_connection_2():
         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
+
+                raw_connection_0 = await database.raw_connection()
+                raw_connection_1 = await connection_1.raw_connection()
+                raw_connection_2 = await connection_2.raw_connection()
+
+                assert raw_connection_0 is raw_connection_1 is raw_connection_2
+                assert raw_connection_0 is connection_1._connection._connection
+
+
 @pytest.mark.parametrize("database_url", DATABASE_URLS)
 @async_adapter
 async def test_queries_with_expose_backend_connection(database_url):
@@ -525,29 +544,30 @@ async def test_queries_with_expose_backend_connection(database_url):
     """
     async with Database(database_url) as database:
         async with database.transaction(force_rollback=True):
+            # Get the raw connection
+            con = await database.raw_connection()
+
             # Insert query
-            if str(database_url).startswith('mysql'):
+            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'):
+            if str(database_url).startswith("postgresql"):
                 await con.execute(insert_query, *values)
-            elif str(database_url).startswith('mysql'):
+            elif str(database_url).startswith("mysql"):
                 cursor = await con.cursor()
                 await cursor.execute(insert_query, values)
-            elif str(database_url).startswith('sqlite'):
+            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'):
+
+            if str(database_url).startswith("mysql"):
                 cursor = await con.cursor()
                 await cursor.executemany(insert_query, values)
             else:
@@ -557,13 +577,13 @@ async def test_queries_with_expose_backend_connection(database_url):
             select_query = "SELECT notes.id, notes.text, notes.completed FROM notes"
 
             # fetch_all()
-            if str(database_url).startswith('postgresql'):
+            if str(database_url).startswith("postgresql"):
                 results = await con.fetch(select_query)
-            elif str(database_url).startswith('mysql'):
+            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'):
+            elif str(database_url).startswith("sqlite"):
                 results = await con.execute_fetchall(select_query)
 
             assert len(results) == 3
@@ -576,48 +596,13 @@ async def test_queries_with_expose_backend_connection(database_url):
             assert results[2][2] == True
 
             # fetch_one()
-            if str(database_url).startswith('postgresql'):
+            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
-
-
-@pytest.mark.parametrize("database_url", DATABASE_URLS)
-@async_adapter
-async def test_queries_with_sqlalchemy_test(database_url):
-    """
-    Test for inserting and retriving the data using the `sqlalchemy.text` query object.
-    """
-    async with Database(database_url) as database:
-        async with database.transaction(force_rollback=True):
-            # Insert query
-            insert_query = "INSERT INTO notes (text, completed) VALUES (:text, :completed)"
-            
-            # execute_many()
-            query = sqlalchemy.text(insert_query)
-            values = [("example1", True), ("example2", False), ("example3", True)]
-            for text, completed in values:
-                current_query = query.bindparams(text=text, completed=completed)
-                await database.execute(current_query)
-
-            # Select query
-            select_query = "SELECT notes.id, notes.text, notes.completed FROM notes"
-
-            # fetch_all()
-            query = sqlalchemy.text(select_query)
-            results = await database.fetch_all(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