Merge pull request #38 from encode/ensure-global-rollback-initialized-on-init

Ensure global connections are created on __init__.

Author: lovelydinosaur

.gitignore

@@ -1,4 +1,5 @@
 *.pyc
+test.db
 .coverage
 .pytest_cache/
 .mypy_cache/

databases/__init__.py

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

databases/backends/mysql.py

@@ -40,8 +40,7 @@ async def disconnect(self) -> None:
         self._pool = None
 
     def connection(self) -> "MySQLConnection":
-        assert self._pool is not None, "DatabaseBackend is not running"
-        return MySQLConnection(self._pool, self._dialect)
+        return MySQLConnection(self, self._dialect)
 
 
 class CompilationContext:
@@ -50,18 +49,20 @@ def __init__(self, context: ExecutionContext):
 
 
 class MySQLConnection(ConnectionBackend):
-    def __init__(self, pool: aiomysql.pool.Pool, dialect: Dialect):
-        self._pool = pool
+    def __init__(self, database: MySQLBackend, dialect: Dialect):
+        self._database = database
         self._dialect = dialect
         self._connection = None  # type: typing.Optional[aiomysql.Connection]
 
     async def acquire(self) -> None:
         assert self._connection is None, "Connection is already acquired"
-        self._connection = await self._pool.acquire()
+        assert self._database._pool is not None, "DatabaseBackend is not running"
+        self._connection = await self._database._pool.acquire()
 
     async def release(self) -> None:
         assert self._connection is not None, "Connection is not acquired"
-        await self._pool.release(self._connection)
+        assert self._database._pool is not None, "DatabaseBackend is not running"
+        await self._database._pool.release(self._connection)
         self._connection = None
 
     async def fetch_all(self, query: ClauseElement) -> typing.List[RowProxy]:

databases/backends/postgres.py

@@ -43,8 +43,7 @@ async def disconnect(self) -> None:
         self._pool = None
 
     def connection(self) -> "PostgresConnection":
-        assert self._pool is not None, "DatabaseBackend is not running"
-        return PostgresConnection(self._pool, self._dialect)
+        return PostgresConnection(self, self._dialect)
 
 
 class Record:
@@ -72,18 +71,20 @@ def __getitem__(self, key: str) -> typing.Any:
 
 
 class PostgresConnection(ConnectionBackend):
-    def __init__(self, pool: asyncpg.pool.Pool, dialect: Dialect):
-        self._pool = pool
+    def __init__(self, database: PostgresBackend, dialect: Dialect):
+        self._database = database
         self._dialect = dialect
         self._connection = None  # type: typing.Optional[asyncpg.connection.Connection]
 
     async def acquire(self) -> None:
         assert self._connection is None, "Connection is already acquired"
-        self._connection = await self._pool.acquire()
+        assert self._database._pool is not None, "DatabaseBackend is not running"
+        self._connection = await self._database._pool.acquire()
 
     async def release(self) -> None:
         assert self._connection is not None, "Connection is not acquired"
-        self._connection = await self._pool.release(self._connection)
+        assert self._database._pool is not None, "DatabaseBackend is not running"
+        self._connection = await self._database._pool.release(self._connection)
         self._connection = None
 
     async def fetch_all(self, query: ClauseElement) -> typing.Any:

databases/core.py

@@ -44,6 +44,12 @@ def __init__(
         self._global_connection = None  # type: typing.Optional[Connection]
         self._global_transaction = None  # type: typing.Optional[Transaction]
 
+        if self._force_rollback:
+            self._global_connection = Connection(self._backend)
+            self._global_transaction = self._global_connection.transaction(
+                force_rollback=True
+            )
+
     async def connect(self) -> None:
         """
         Establish the connection pool.
@@ -54,10 +60,7 @@ async def connect(self) -> None:
         self.is_connected = True
 
         if self._force_rollback:
-            self._global_connection = Connection(self._backend)
-            self._global_transaction = self._global_connection.transaction(
-                force_rollback=True
-            )
+            assert self._global_transaction is not None
             await self._global_transaction.__aenter__()
 
     async def disconnect(self) -> None:
@@ -69,8 +72,6 @@ async def disconnect(self) -> None:
         if self._force_rollback:
             assert self._global_transaction is not None
             await self._global_transaction.__aexit__()
-            self._global_transaction = None
-            self._global_connection = None
 
         await self._backend.disconnect()
         self.is_connected = False