Merge pull request #28 from encode/support-global-force-rollback

Support global force rollback

Author: lovelydinosaur

.gitignore

@@ -2,5 +2,6 @@
 .coverage
 .pytest_cache/
 .mypy_cache/
-starlette.egg-info/
+*.egg-info/
+htmlcov/
 venv/

README.md

@@ -90,7 +90,7 @@ rows = await database.fetch_all(query)
 query = notes.select()
 row = await database.fetch_one(query)
 
-# Fetch multiple rows without loading everything into memory at once.
+# Fetch multiple rows without loading them all into memory at once
 query = notes.select()
 async for row in database.iterate(query):
     ...

databases/core.py

@@ -15,118 +15,57 @@
     from aiocontextvars import ContextVar
 
 
-class DatabaseURL:
-    def __init__(self, url: typing.Union[str, "DatabaseURL"]):
-        if isinstance(url, DatabaseURL):
-            self._url = str(url)
-        else:
-            self._url = url
-
-    @property
-    def components(self) -> SplitResult:
-        if not hasattr(self, "_components"):
-            self._components = urlsplit(self._url)
-        return self._components
-
-    @property
-    def dialect(self) -> str:
-        return self.components.scheme.split("+")[0]
-
-    @property
-    def driver(self) -> str:
-        if "+" not in self.components.scheme:
-            return ""
-        return self.components.scheme.split("+", 1)[1]
-
-    @property
-    def username(self) -> typing.Optional[str]:
-        return self.components.username
-
-    @property
-    def password(self) -> typing.Optional[str]:
-        return self.components.password
-
-    @property
-    def hostname(self) -> typing.Optional[str]:
-        return self.components.hostname
-
-    @property
-    def port(self) -> typing.Optional[int]:
-        return self.components.port
-
-    @property
-    def database(self) -> str:
-        return self.components.path.lstrip("/")
-
-    def replace(self, **kwargs: typing.Any) -> "DatabaseURL":
-        if (
-            "username" in kwargs
-            or "password" in kwargs
-            or "hostname" in kwargs
-            or "port" in kwargs
-        ):
-            hostname = kwargs.pop("hostname", self.hostname)
-            port = kwargs.pop("port", self.port)
-            username = kwargs.pop("username", self.username)
-            password = kwargs.pop("password", self.password)
-
-            netloc = hostname
-            if port is not None:
-                netloc += f":{port}"
-            if username is not None:
-                userpass = username
-                if password is not None:
-                    userpass += f":{password}"
-                netloc = f"{userpass}@{netloc}"
-
-            kwargs["netloc"] = netloc
-
-        if "database" in kwargs:
-            kwargs["path"] = "/" + kwargs.pop("database")
-
-        if "dialect" in kwargs or "driver" in kwargs:
-            dialect = kwargs.pop("dialect", self.dialect)
-            driver = kwargs.pop("driver", self.driver)
-            kwargs["scheme"] = f"{dialect}+{driver}" if driver else dialect
-
-        components = self.components._replace(**kwargs)
-        return self.__class__(components.geturl())
-
-    def __str__(self) -> str:
-        return self._url
-
-    def __repr__(self) -> str:
-        url = str(self)
-        if self.password:
-            url = str(self.replace(password="********"))
-        return f"{self.__class__.__name__}({repr(url)})"
-
-
 class Database:
     SUPPORTED_BACKENDS = {
         "postgresql": "databases.backends.postgres:PostgresBackend",
         "mysql": "databases.backends.mysql:MySQLBackend",
     }
 
-    def __init__(self, url: typing.Union[str, DatabaseURL]):
-        self.url = DatabaseURL(url)
+    def __init__(
+        self, url: typing.Union[str, "DatabaseURL"], force_rollback: bool = False
+    ):
+        self._url = DatabaseURL(url)
+        self._force_rollback = force_rollback
+
         self.is_connected = False
 
-        backend_str = self.SUPPORTED_BACKENDS[self.url.dialect]
+        backend_str = self.SUPPORTED_BACKENDS[self._url.dialect]
         backend_cls = import_from_string(backend_str)
         assert issubclass(backend_cls, DatabaseBackend)
-        self._backend = backend_cls(self.url)
+        self._backend = backend_cls(self._url)
+
+        # Connections are stored as task-local state.
         self._connection_context = ContextVar("connection_context")  # type: ContextVar
 
+        # When `force_rollback=True` is used, we use a single global
+        # connection, within a transaction that always rolls back.
+        self._global_connection = None  # type: typing.Optional[Connection]
+        self._global_transaction = None  # type: typing.Optional[Transaction]
+
     async def connect(self) -> None:
-        if not self.is_connected:
-            await self._backend.connect()
-            self.is_connected = True
+        assert not self.is_connected, "Already connected."
+
+        await self._backend.connect()
+        self.is_connected = True
+
+        if self._force_rollback:
+            self._global_connection = Connection(self._backend)
+            self._global_transaction = self._global_connection.transaction(
+                force_rollback=True
+            )
+            await self._global_transaction.__aenter__()
 
     async def disconnect(self) -> None:
-        if self.is_connected:
-            await self._backend.disconnect()
-            self.is_connected = False
+        assert self.is_connected, "Already disconnected."
+
+        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
 
     async def __aenter__(self) -> "Database":
         await self.connect()
@@ -164,6 +103,9 @@ async def iterate(
                 yield record
 
     def connection(self) -> "Connection":
+        if self._global_connection is not None:
+            return self._global_connection
+
         try:
             return self._connection_context.get()
         except LookupError:
@@ -267,3 +209,90 @@ async def rollback(self) -> None:
             self._connection._transaction_stack.pop()
             await self._transaction.rollback()
             await self._connection.__aexit__()
+
+
+class DatabaseURL:
+    def __init__(self, url: typing.Union[str, "DatabaseURL"]):
+        if isinstance(url, DatabaseURL):
+            self._url = str(url)
+        else:
+            self._url = url
+
+    @property
+    def components(self) -> SplitResult:
+        if not hasattr(self, "_components"):
+            self._components = urlsplit(self._url)
+        return self._components
+
+    @property
+    def dialect(self) -> str:
+        return self.components.scheme.split("+")[0]
+
+    @property
+    def driver(self) -> str:
+        if "+" not in self.components.scheme:
+            return ""
+        return self.components.scheme.split("+", 1)[1]
+
+    @property
+    def username(self) -> typing.Optional[str]:
+        return self.components.username
+
+    @property
+    def password(self) -> typing.Optional[str]:
+        return self.components.password
+
+    @property
+    def hostname(self) -> typing.Optional[str]:
+        return self.components.hostname
+
+    @property
+    def port(self) -> typing.Optional[int]:
+        return self.components.port
+
+    @property
+    def database(self) -> str:
+        return self.components.path.lstrip("/")
+
+    def replace(self, **kwargs: typing.Any) -> "DatabaseURL":
+        if (
+            "username" in kwargs
+            or "password" in kwargs
+            or "hostname" in kwargs
+            or "port" in kwargs
+        ):
+            hostname = kwargs.pop("hostname", self.hostname)
+            port = kwargs.pop("port", self.port)
+            username = kwargs.pop("username", self.username)
+            password = kwargs.pop("password", self.password)
+
+            netloc = hostname
+            if port is not None:
+                netloc += f":{port}"
+            if username is not None:
+                userpass = username
+                if password is not None:
+                    userpass += f":{password}"
+                netloc = f"{userpass}@{netloc}"
+
+            kwargs["netloc"] = netloc
+
+        if "database" in kwargs:
+            kwargs["path"] = "/" + kwargs.pop("database")
+
+        if "dialect" in kwargs or "driver" in kwargs:
+            dialect = kwargs.pop("dialect", self.dialect)
+            driver = kwargs.pop("driver", self.driver)
+            kwargs["scheme"] = f"{dialect}+{driver}" if driver else dialect
+
+        components = self.components._replace(**kwargs)
+        return self.__class__(components.geturl())
+
+    def __str__(self) -> str:
+        return self._url
+
+    def __repr__(self) -> str:
+        url = str(self)
+        if self.password:
+            url = str(self.replace(password="********"))
+        return f"{self.__class__.__name__}({repr(url)})"

requirements.txt

@@ -17,7 +17,5 @@ isort
 mypy
 pytest
 pytest-cov
-
-# Documentation
-mkdocs
-mkdocs-material
+starlette
+requests

tests/test_databases.py

@@ -345,3 +345,18 @@ async def test_connections_isolation(database_url):
         finally:
             query = notes.delete()
             await database.execute(query)
+
+
+@pytest.mark.parametrize("database_url", DATABASE_URLS)
+@async_adapter
+async def test_connect_and_disconnect(database_url):
+    """
+    Test explicit connect() and disconnect().
+    """
+    database = Database(database_url)
+
+    assert not database.is_connected
+    await database.connect()
+    assert database.is_connected
+    await database.disconnect()
+    assert not database.is_connected

tests/test_integration.py

@@ -0,0 +1,97 @@
+import os
+
+import pytest
+import sqlalchemy
+from starlette.applications import Starlette
+from starlette.responses import JSONResponse
+from starlette.testclient import TestClient
+
+from databases import Database, DatabaseURL
+
+assert "TEST_DATABASE_URLS" in os.environ, "TEST_DATABASE_URLS is not set."
+
+DATABASE_URLS = [url.strip() for url in os.environ["TEST_DATABASE_URLS"].split(",")]
+
+metadata = sqlalchemy.MetaData()
+
+notes = sqlalchemy.Table(
+    "notes",
+    metadata,
+    sqlalchemy.Column("id", sqlalchemy.Integer, primary_key=True),
+    sqlalchemy.Column("text", sqlalchemy.String(length=100)),
+    sqlalchemy.Column("completed", sqlalchemy.Boolean),
+)
+
+
+@pytest.fixture(autouse=True, scope="module")
+def create_test_database():
+    # Create test databases
+    for url in DATABASE_URLS:
+        database_url = DatabaseURL(url)
+        if database_url.dialect == "mysql":
+            url = str(database_url.replace(driver="pymysql"))
+        engine = sqlalchemy.create_engine(url)
+        metadata.create_all(engine)
+
+    # Run the test suite
+    yield
+
+    # Drop test databases
+    for url in DATABASE_URLS:
+        database_url = DatabaseURL(url)
+        if database_url.dialect == "mysql":
+            url = str(database_url.replace(driver="pymysql"))
+        engine = sqlalchemy.create_engine(url)
+        metadata.drop_all(engine)
+
+
+def get_app(database_url):
+    database = Database(database_url, force_rollback=True)
+    app = Starlette()
+
+    @app.on_event("startup")
+    async def startup():
+        await database.connect()
+
+    @app.on_event("shutdown")
+    async def shutdown():
+        await database.disconnect()
+
+    @app.route("/notes", methods=["GET"])
+    async def list_notes(request):
+        query = notes.select()
+        results = await database.fetch_all(query)
+        content = [
+            {"text": result["text"], "completed": result["completed"]}
+            for result in results
+        ]
+        return JSONResponse(content)
+
+    @app.route("/notes", methods=["POST"])
+    async def add_note(request):
+        data = await request.json()
+        query = notes.insert().values(text=data["text"], completed=data["completed"])
+        await database.execute(query)
+        return JSONResponse({"text": data["text"], "completed": data["completed"]})
+
+    return app
+
+
+@pytest.mark.parametrize("database_url", DATABASE_URLS)
+def test_integration(database_url):
+    app = get_app(database_url)
+
+    with TestClient(app) as client:
+        response = client.post("/notes", json={"text": "example", "completed": True})
+        assert response.status_code == 200
+        assert response.json() == {"text": "example", "completed": True}
+
+        response = client.get("/notes")
+        assert response.status_code == 200
+        assert response.json() == [{"text": "example", "completed": True}]
+
+    with TestClient(app) as client:
+        # Ensure sessions are isolated
+        response = client.get("/notes")
+        assert response.status_code == 200
+        assert response.json() == []