diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000..e69de29 diff --git a/Makefile b/Makefile index 7bb2730..ef72eb4 100644 --- a/Makefile +++ b/Makefile @@ -11,6 +11,9 @@ ci: test: pytest -s --tb=native +format-check: + ruff format --check && ruff check + format: ruff format && ruff check --fix diff --git a/README.md b/README.md index 241bd1b..7c0e438 100644 --- a/README.md +++ b/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 ``` diff --git a/hooks/pre-commit b/hooks/pre-commit index 482f9c7..189f54f 100755 --- a/hooks/pre-commit +++ b/hooks/pre-commit @@ -1,5 +1,5 @@ #!/bin/sh -make format +make format-check make typecheck diff --git a/src/queryspy/__init__.py b/src/queryspy/__init__.py index 8e18dde..68af79f 100644 --- a/src/queryspy/__init__.py +++ b/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", +] diff --git a/src/queryspy/apps.py b/src/queryspy/apps.py index ff3440b..ed70b60 100644 --- a/src/queryspy/apps.py +++ b/src/queryspy/apps.py @@ -1,6 +1,5 @@ from django.apps import AppConfig -from .listeners import reset as setup_listeners from .patch import patch @@ -8,6 +7,4 @@ class QuerySpyConfig(AppConfig): name = "queryspy" def ready(self): - setup_listeners() patch() - pass diff --git a/src/queryspy/listeners.py b/src/queryspy/listeners.py index 6fd18f8..7e208d4 100644 --- a/src/queryspy/listeners.py +++ b/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) diff --git a/src/queryspy/middleware.py b/src/queryspy/middleware.py new file mode 100644 index 0000000..757961c --- /dev/null +++ b/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 diff --git a/tests/djangoproject/settings.py b/tests/djangoproject/settings.py index e126ff6..7c2992c 100644 --- a/tests/djangoproject/settings.py +++ b/tests/djangoproject/settings.py @@ -14,7 +14,7 @@ "queryspy", ] -MIDDLEWARE = [] +MIDDLEWARE = ["queryspy.middleware.queryspy_middleware"] ROOT_URLCONF = "djangoproject.urls" diff --git a/tests/djangoproject/social/views.py b/tests/djangoproject/social/views.py new file mode 100644 index 0000000..fdcebc4 --- /dev/null +++ b/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() + ] + } + ) diff --git a/tests/djangoproject/urls.py b/tests/djangoproject/urls.py new file mode 100644 index 0000000..57dfbda --- /dev/null +++ b/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//", single_user_and_profile), +] diff --git a/tests/test_nplusones.py b/tests/test_nplusones.py index 14b22d8..e844780 100644 --- a/tests/test_nplusones.py +++ b/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