feat: add middleware and installation instructions

Makefile

@@ -11,6 +11,9 @@ ci:
 test:
 	pytest -s --tb=native
 
+format-check:
+	ruff format --check && ruff check
+
 format:
 	ruff format && ruff check --fix
 

README.md

@@ -5,7 +5,6 @@ This library catches N+1s in your Django project.
 ## Features
 
 - Detects N+1s from missing prefetches and from `.defer()`/`.only()`
-- Simple installation -- just add to your `INSTALLED_APPS` and it works everywhere
 - Configurable thresholds
 - TODO: allowlist
 - TODO: catches unused eager loads
@@ -14,17 +13,70 @@ This library catches N+1s in your Django project.
 
 ## Acknowledgements
 
-This library draws very heavily on jmcarp's [nplusone](https://github.com/jmcarp/nplusone/).
-It's not *exactly* a fork, but not far from it.
+This library draws heavily from jmcarp's [nplusone](https://github.com/jmcarp/nplusone/).
+It's not exactly a fork, but not far from it.
 
 ## Installation
 
-TODO.
+To install `queryspy`, add it to your `INSTALLED_APPS` and `MIDDLEWARE`. You probably
+don't want to run it in production: I haven't profiled it but it will have a performance
+impact.
+
+```python
+if DEBUG:
+    INSTALLED_APPS.append("queryspy")
+    MIDDLEWARE.append("queryspy.middleware.queryspy_middleware")
+```
+
+This will detect N+1s that happen in web requests. If you also want to find N+1s in other
+places like background tasks or management commands, you can use the `setup` and
+`teardown` functions, or the `queryspy_context` context manager:
+
+```python
+from queryspy import setup, teardown, queryspy_context
+
+
+def foo():
+    setup()
+    try:
+        # ...
+    finally:
+        teardown()
+
+
+@queryspy_context()
+def bar():
+    # ...
+
+
+def baz():
+    with queryspy_context():
+        # ...
+```
+
+For example, if you use Celery, you can configure this using [signals](https://docs.celeryq.dev/en/stable/userguide/signals.html):
+
+```python
+from celery.signals import task_prerun, task_postrun
+from queryspy import setup, teardown
+from django.conf import settings
+
+if settings.DEBUG:
+    @task_prerun.connect()
+    def setup_queryspy(*args, **kwargs):
+        setup()
+
+    @task_postrun.connect()
+    def teardown_queryspy(*args, **kwargs):
+        teardown()
+```
+
+## Configuration
 
 By default, N+1s will be reported when the same query is executed twice. To configure this
 threshold, set the following in your Django settings.
 
-```
+```python
 QUERYSPY_NPLUSONE_THRESHOLD = 3
 ```
 

hooks/pre-commit

@@ -1,5 +1,5 @@
 #!/bin/sh
 
-make format
+make format-check
 make typecheck
 

src/queryspy/__init__.py

@@ -1,4 +1,10 @@
 from .errors import NPlusOneError, QuerySpyError
-from .listeners import reset
+from .listeners import queryspy_context, setup, teardown
 
-__all__ = ["QuerySpyError", "NPlusOneError", "reset"]
+__all__ = [
+    "QuerySpyError",
+    "NPlusOneError",
+    "setup",
+    "teardown",
+    "queryspy_context",
+]

src/queryspy/apps.py

@@ -1,13 +1,10 @@
 from django.apps import AppConfig
 
-from .listeners import reset as setup_listeners
 from .patch import patch
 
 
 class QuerySpyConfig(AppConfig):
     name = "queryspy"
 
     def ready(self):
-        setup_listeners()
         patch()
-        pass

src/queryspy/listeners.py

@@ -1,5 +1,7 @@
 from abc import ABC, abstractmethod
 from collections import defaultdict
+from contextlib import contextmanager
+from contextvars import ContextVar
 from typing import Type
 
 from django.conf import settings
@@ -9,6 +11,8 @@
 
 ModelAndField = tuple[Type[models.Model], str]
 
+_is_in_context = ContextVar("in_context", default=False)
+
 
 class Listener(ABC):
     @abstractmethod
@@ -26,6 +30,9 @@ def __init__(self):
         self.reset()
 
     def notify(self, model: Type[models.Model], field: str):
+        if not _is_in_context.get():
+            return
+
         threshold = (
             settings.QUERYSPY_NPLUSONE_THRESHOLD
             if hasattr(settings, "QUERYSPY_NPLUSONE_THRESHOLD")
@@ -45,5 +52,20 @@ def reset(self):
 n_plus_one_listener = NPlusOneListener()
 
 
-def reset():
+def setup():
+    _is_in_context.set(True)
+
+
+def teardown():
     n_plus_one_listener.reset()
+    _is_in_context.set(False)
+
+
+@contextmanager
+def queryspy_context():
+    token = _is_in_context.set(True)
+    try:
+        yield
+    finally:
+        n_plus_one_listener.reset()
+        _is_in_context.reset(token)

src/queryspy/middleware.py

@@ -0,0 +1,10 @@
+from .listeners import queryspy_context
+
+
+def queryspy_middleware(get_response):
+    def middleware(request):
+        with queryspy_context():
+            response = get_response(request)
+        return response
+
+    return middleware

tests/djangoproject/settings.py

@@ -14,7 +14,7 @@
     "queryspy",
 ]
 
-MIDDLEWARE = []
+MIDDLEWARE = ["queryspy.middleware.queryspy_middleware"]
 
 ROOT_URLCONF = "djangoproject.urls"
 

tests/djangoproject/social/views.py

@@ -0,0 +1,30 @@
+from django.http import HttpRequest, JsonResponse
+
+from .models import User
+
+
+def single_user_and_profile(request: HttpRequest, id: int):
+    user = User.objects.get(id=id)
+    return JsonResponse(
+        data={
+            "username": user.username,
+            "display_name": user.profile.display_name,
+        }
+    )
+
+
+def all_users_and_profiles(request: HttpRequest):
+    """
+    This view has an N+1.
+    """
+    return JsonResponse(
+        data={
+            "users": [
+                {
+                    "username": user.username,
+                    "display_name": user.profile.display_name,
+                }
+                for user in User.objects.all()
+            ]
+        }
+    )

tests/djangoproject/urls.py

@@ -0,0 +1,8 @@
+from django.urls import path
+
+from .social.views import all_users_and_profiles, single_user_and_profile
+
+urlpatterns = [
+    path("users/", all_users_and_profiles),
+    path("user/<int:id>/", single_user_and_profile),
+]

tests/test_nplusones.py

@@ -1,17 +1,13 @@
 import pytest
 from djangoproject.social.models import Post, Profile, User
-from queryspy import NPlusOneError, reset
+from queryspy import NPlusOneError, queryspy_context
 
 from .factories import PostFactory, ProfileFactory, UserFactory
 
 pytestmark = pytest.mark.django_db
 
 
-@pytest.fixture(autouse=True)
-def reset_listener():
-    reset()
-
-
+@queryspy_context()
 def test_detects_nplusone_in_forward_many_to_one():
     [user_1, user_2] = UserFactory.create_batch(2)
     PostFactory.create(author=user_1)
@@ -24,6 +20,7 @@ def test_detects_nplusone_in_forward_many_to_one():
         _ = post.author.username
 
 
+@queryspy_context()
 def test_detects_nplusone_in_reverse_many_to_one():
     [user_1, user_2] = UserFactory.create_batch(2)
     PostFactory.create(author=user_1)
@@ -36,6 +33,7 @@ def test_detects_nplusone_in_reverse_many_to_one():
         _ = list(user.posts.all())
 
 
+@queryspy_context()
 def test_detects_nplusone_in_forward_one_to_one():
     [user_1, user_2] = UserFactory.create_batch(2)
     ProfileFactory.create(user=user_1)
@@ -48,6 +46,7 @@ def test_detects_nplusone_in_forward_one_to_one():
         _ = profile.user.username
 
 
+@queryspy_context()
 def test_detects_nplusone_in_reverse_one_to_one():
     [user_1, user_2] = UserFactory.create_batch(2)
     ProfileFactory.create(user=user_1)
@@ -60,6 +59,7 @@ def test_detects_nplusone_in_reverse_one_to_one():
         _ = user.profile.display_name
 
 
+@queryspy_context()
 def test_detects_nplusone_in_forward_many_to_many():
     [user_1, user_2] = UserFactory.create_batch(2)
     user_1.following.add(user_2)
@@ -72,6 +72,7 @@ def test_detects_nplusone_in_forward_many_to_many():
         _ = list(user.following.all())
 
 
+@queryspy_context()
 def test_detects_nplusone_in_reverse_many_to_many():
     [user_1, user_2] = UserFactory.create_batch(2)
     user_1.following.add(user_2)
@@ -84,6 +85,7 @@ def test_detects_nplusone_in_reverse_many_to_many():
         _ = list(user.followers.all())
 
 
+@queryspy_context()
 def test_detects_nplusone_in_reverse_many_to_many_with_no_related_name():
     [user_1, user_2] = UserFactory.create_batch(2)
     user_1.blocked.add(user_2)
@@ -96,6 +98,7 @@ def test_detects_nplusone_in_reverse_many_to_many_with_no_related_name():
         _ = list(user.user_set.all())
 
 
+@queryspy_context()
 def test_detects_nplusone_due_to_deferred_fields():
     [user_1, user_2] = UserFactory.create_batch(2)
     PostFactory.create(author=user_1)
@@ -120,3 +123,32 @@ def test_has_configurable_threshold(settings):
 
     for post in Post.objects.all():
         _ = post.author.username
+
+
+def test_does_nothing_if_not_in_middleware(settings, client):
+    settings.MIDDLEWARE = []
+    [user_1, user_2] = UserFactory.create_batch(2)
+    ProfileFactory.create(user=user_1)
+    ProfileFactory.create(user=user_2)
+
+    # this does not raise an N+1 error even though the same
+    # related field gets hit twice
+    response = client.get(f"/user/{user_1.pk}/")
+    assert response.status_code == 200
+    response = client.get(f"/user/{user_2.pk}/")
+    assert response.status_code == 200
+
+
+def test_works_in_web_requests(client):
+    [user_1, user_2] = UserFactory.create_batch(2)
+    ProfileFactory.create(user=user_1)
+    ProfileFactory.create(user=user_2)
+    with pytest.raises(NPlusOneError):
+        response = client.get("/users/")
+        assert response.status_code == 500
+
+    # but multiple requests work fine
+    response = client.get(f"/user/{user_1.pk}/")
+    assert response.status_code == 200
+    response = client.get(f"/user/{user_2.pk}/")
+    assert response.status_code == 200