Refactoring: using properties instead of methods

Author: gvbgduh

databases/backends/mysql.py

@@ -149,10 +149,6 @@ 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)
 
@@ -176,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

@@ -150,10 +150,6 @@ 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)
 
@@ -175,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

@@ -116,10 +116,6 @@ async def execute_many(self, query: ClauseElement, values: list) -> None:
         for value in values:
             await self.execute(query, value)
 
-    async def raw_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]:
@@ -157,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

@@ -111,10 +111,6 @@ async def iterate(
             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,9 +168,6 @@ 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_connection(self) -> typing.Any:
-        return await self._connection.raw_connection()
-
     async def iterate(
         self, query: ClauseElement
     ) -> typing.AsyncGenerator[typing.Any, None]:
@@ -184,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

@@ -41,10 +41,11 @@ async def iterate(
         # https://github.com/python/mypy/issues/5385#issuecomment-407281656
         yield True  # pragma: no cover
 
-    async def raw_connection(self) -> typing.Any:
+    def transaction(self) -> "TransactionBackend":
         raise NotImplementedError()  # pragma: no cover
 
-    def transaction(self) -> "TransactionBackend":
+    @property
+    def raw_connection(self) -> typing.Any:
         raise NotImplementedError()  # pragma: no cover
 
 

tests/test_databases.py

@@ -526,13 +526,7 @@ async def test_connection_context_with_raw_connection(database_url):
         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
+                assert connection_1.raw_connection is connection_2.raw_connection
 
 
 @pytest.mark.parametrize("database_url", DATABASE_URLS)
@@ -543,66 +537,67 @@ async def test_queries_with_expose_backend_connection(database_url):
     `fetch_one()` using the raw driver interface.
     """
     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"):
-                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 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
+        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