Repository: meshy/django-sans-db (language: python)

Issue:
Add a function/method decorator
It would be nice to be able to do:

```python
class MyClass(...):
    @block_db(databases=["primary"])
    def my_method(self, ...):
        ...
```


Discussion / hints:
Oh! It's possible we already support this because we're using [`contextlib.contextmanager`](https://docs.python.org/3/library/contextlib.html#contextlib.contextmanager):

> [contextmanager()](https://docs.python.org/3/library/contextlib.html#contextlib.contextmanager) uses [ContextDecorator](https://docs.python.org/3/library/contextlib.html#contextlib.ContextDecorator) so the context managers it creates can be used as decorators as well as in [with](https://docs.python.org/3/reference/compound_stmts.html#with) statements.

We should have tests for this, and docs.

Fix patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index e518861..1d4ed0b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,8 @@
 - Remove upper limit on Python version.
   This should prevent issues when installing for new versions of Python.
 - Add Django 4.1 to text matrix.
+- Add tests for using `block_db` as a decorator.
+  This was working before, but not officially supported.
 
 ## [1.1.0] - 2022-06-11
 
diff --git a/README.md b/README.md
index 06fefca..429038f 100644
--- a/README.md
+++ b/README.md
@@ -39,6 +39,24 @@ with block_db(databases=["replica"]):
     User.objects.using("replica").get()  # Raises DatabaseAccessBlocked
 ```
 
+
+### Decorator
+
+You can decorate functions and methods with `block_db` to block database access in them. Eg:
+
+```python
+from sans_db.context_managers import block_db
+
+class MyClass:
+    def allowed(self):
+        User.objects.create(...)  # Works outside of block_db()
+
+    @block_db()
+    def not_allowed(self):
+        User.objects.create(...)  # Raises DatabaseAccessBlocked
+```
+
+
 ### Template backend
 
 You can block access to the database when rendering Django templates with our custom template backend.


Test patch:
diff --git a/tests/test_context_managers.py b/tests/test_context_managers.py
index 30867f6..8b0424b 100644
--- a/tests/test_context_managers.py
+++ b/tests/test_context_managers.py
@@ -9,6 +9,8 @@
 
 @pytest.mark.django_db(databases=DATABASE_ALIASES)
 class TestBlockDB:
+    """Tests for block_db when used as a context manager."""
+
     @pytest.mark.parametrize("database", DATABASE_ALIASES)
     def test_queryset_evaluation_blocked(self, database: str) -> None:
         queryset = ExampleModel.objects.using(database).all()
@@ -20,8 +22,7 @@ def test_queryset_evaluation_blocked(self, database: str) -> None:
     def test_evaluated_queryset_allowed(self, database: str) -> None:
         queryset = list(ExampleModel.objects.using(database).all())
         with block_db():
-            for model in queryset:
-                pass
+            list(queryset)
 
     @pytest.mark.parametrize("blocked", DATABASE_ALIASES)
     def test_selective_blocking_blocked(self, blocked: str) -> None:
diff --git a/tests/test_decorator.py b/tests/test_decorator.py
new file mode 100644
index 0000000..4ff2b9f
--- /dev/null
+++ b/tests/test_decorator.py
@@ -0,0 +1,105 @@
+import pytest
+
+from sans_db.context_managers import block_db
+from sans_db.exceptions import DatabaseAccessBlocked
+
+from .definitions import DATABASE_ALIASES
+from .models import ExampleModel
+
+
+@pytest.mark.django_db(databases=DATABASE_ALIASES)
+class TestBlockDBFunctionDecorator:
+    """Tests for block_db when used as a function decorator."""
+
+    @pytest.mark.parametrize("database", DATABASE_ALIASES)
+    def test_queryset_evaluation_blocked(self, database: str) -> None:
+        queryset = ExampleModel.objects.using(database).all()
+
+        @block_db()
+        def func() -> None:
+            list(queryset)
+
+        with pytest.raises(DatabaseAccessBlocked):
+            func()
+
+    @pytest.mark.parametrize("database", DATABASE_ALIASES)
+    def test_evaluated_queryset_allowed(self, database: str) -> None:
+        queryset = list(ExampleModel.objects.using(database).all())
+
+        @block_db()
+        def func() -> None:
+            list(queryset)
+
+        # No error raised, because query isn't on blocked DB.
+        func()
+
+    @pytest.mark.parametrize("blocked", DATABASE_ALIASES)
+    def test_selective_blocking_blocked(self, blocked: str) -> None:
+        @block_db(databases=[blocked])
+        def func() -> None:
+            ExampleModel.objects.using(blocked).create()
+
+        with pytest.raises(DatabaseAccessBlocked):
+            func()
+
+    @pytest.mark.parametrize("blocked", DATABASE_ALIASES)
+    def test_selective_blocking_allowed(self, blocked: str) -> None:
+        @block_db(databases=[blocked])
+        def func(db: str) -> None:
+            ExampleModel.objects.using(db).create()
+
+        # No error raised, because queries aren't on blocked DB.
+        not_blocked = DATABASE_ALIASES - {blocked}
+        for db in not_blocked:
+            func(db)
+
+
+@pytest.mark.django_db(databases=DATABASE_ALIASES)
+class TestBlockDBMethodDecorator:
+    """Tests for block_db when used as a method decorator."""
+
+    @pytest.mark.parametrize("database", DATABASE_ALIASES)
+    def test_queryset_evaluation_blocked(self, database: str) -> None:
+        queryset = ExampleModel.objects.using(database).all()
+
+        class MyClass:
+            @block_db()
+            def run(self) -> None:
+                list(queryset)
+
+        with pytest.raises(DatabaseAccessBlocked):
+            MyClass().run()
+
+    @pytest.mark.parametrize("database", DATABASE_ALIASES)
+    def test_evaluated_queryset_allowed(self, database: str) -> None:
+        queryset = list(ExampleModel.objects.using(database).all())
+
+        class MyClass:
+            @block_db()
+            def run(self) -> None:
+                list(queryset)
+
+        # No error raised, because query isn't on blocked DB.
+        MyClass().run()
+
+    @pytest.mark.parametrize("blocked", DATABASE_ALIASES)
+    def test_selective_blocking_blocked(self, blocked: str) -> None:
+        class MyClass:
+            @block_db(databases=[blocked])
+            def run(self) -> None:
+                ExampleModel.objects.using(blocked).create()
+
+        with pytest.raises(DatabaseAccessBlocked):
+            MyClass().run()
+
+    @pytest.mark.parametrize("blocked", DATABASE_ALIASES)
+    def test_selective_blocking_allowed(self, blocked: str) -> None:
+        class MyClass:
+            @block_db(databases=[blocked])
+            def run(self, db: str) -> None:
+                ExampleModel.objects.using(db).create()
+
+        # No error raised, because queries aren't on blocked DB.
+        not_blocked = DATABASE_ALIASES - {blocked}
+        for db in not_blocked:
+            MyClass().run(db)