django-nplus1 0.1.0__py3-none-any.whl

django_nplus1/__init__.py

@@ -0,0 +1,14 @@
+from importlib.metadata import PackageNotFoundError, version
+
+from django_nplus1.detect import nplus1_allow
+from django_nplus1.exceptions import NPlus1Error
+from django_nplus1.middleware import NPlus1Middleware
+from django_nplus1.profiler import Profiler
+from django_nplus1.signals import nplus1_detected
+
+try:
+    __version__ = version("django-nplus1")
+except PackageNotFoundError:
+    __version__ = "0.0.0+unknown"
+
+__all__ = ["NPlus1Error", "NPlus1Middleware", "Profiler", "__version__", "nplus1_allow", "nplus1_detected"]

django_nplus1/apps.py

@@ -0,0 +1,9 @@
+from django.apps import AppConfig
+
+
+class DjangoNPlus1Config(AppConfig):
+    name = "django_nplus1"
+    verbose_name = "Django N+1"
+
+    def ready(self) -> None:
+        from django_nplus1 import patch  # noqa: F401

django_nplus1/detect.py

@@ -0,0 +1,438 @@
+from __future__ import annotations
+
+import contextlib
+import fnmatch
+import re
+from collections import defaultdict
+from contextvars import ContextVar
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from collections.abc import Generator, Sequence
+
+
+class Rule:
+    def __init__(self, label: str | None = None, model: str | type | None = None, field: str | None = None) -> None:
+        self.label = label
+        self.model = model
+        self.field = field
+
+    def compare(self, label: str, model: type, field: str) -> bool:
+        return bool(
+            (self.label or self.model or self.field)
+            and (self.label is None or self.label == label)
+            and (self.model is None or self.match_model(model))
+            and (self.field is None or self.match_field(field)),
+        )
+
+    def match_field(self, field: str) -> bool:
+        if self.field is None:
+            return True
+        return fnmatch.fnmatch(field, self.field)
+
+    def match_model(self, model: type) -> bool:
+        if self.model is model:
+            return True
+        if isinstance(self.model, str):
+            return fnmatch.fnmatch(model.__name__, self.model)
+        return False
+
+
+_allow_rules: ContextVar[list[Rule]] = ContextVar("nplus1_allow_rules")
+
+
+def is_allowed(message: Message) -> bool:
+    """Check if a message is suppressed by nplus1_allow rules."""
+    try:
+        rules = _allow_rules.get()
+    except LookupError:
+        return False
+    return message.match(rules) if rules else False
+
+
+@contextlib.contextmanager
+def nplus1_allow(whitelist: list[dict[str, Any]] | None = None) -> Generator[None]:
+    """Context manager to suppress N+1 detection for specific model/field combinations.
+
+    With no arguments, suppresses all detections. With a whitelist, suppresses only
+    matching detections. Uses the same format as ``Profiler(whitelist=...)`` and
+    ``@pytest.mark.nplus1(whitelist=...)``.
+
+    Usage::
+
+        # Suppress all detections
+        with nplus1_allow():
+            ...
+
+        # Suppress specific model/field
+        with nplus1_allow([{"model": "User", "field": "hobbies"}]):
+            ...
+
+        # Suppress all fields on a model (supports fnmatch wildcards)
+        with nplus1_allow([{"model": "User"}]):
+            ...
+    """
+    rules = [Rule(**item) for item in whitelist] if whitelist else [Rule(model="*", field="*")]
+
+    try:
+        current = _allow_rules.get()
+    except LookupError:
+        current = []
+    token = _allow_rules.set([*current, *rules])
+    try:
+        yield
+    finally:
+        _allow_rules.reset(token)
+
+
+class Message:
+    label: str = ""
+    formatter: str = ""
+
+    def __init__(
+        self,
+        model: type,
+        field: str,
+        caller: tuple[str, int, str] | None = None,
+        callers: list[list[tuple[str, int, str]]] | None = None,
+    ) -> None:
+        self.model = model
+        self.field = field
+        self.caller = caller
+        self.callers = callers
+
+    @property
+    def message(self) -> str:
+        base = self.formatter.format(
+            label=self.label,
+            model=self.model.__name__,
+            field=self.field,
+        )
+        if self.callers:
+            parts = [base, " with calls:"]
+            for i, stack in enumerate(self.callers, 1):
+                parts.append(f"\nCALL {i}:")
+                for fn, lineno, funcname in stack:
+                    parts.append(f"\n  {fn}:{lineno} in {funcname}")
+            return "".join(parts)
+        if self.caller:
+            filename, lineno, funcname = self.caller
+            return f"{base} at {filename}:{lineno} in {funcname}"
+        return base
+
+    def match(self, rules: Sequence[Rule]) -> bool:
+        return any(rule.compare(self.label, self.model, self.field) for rule in rules)
+
+
+class LazyLoadMessage(Message):
+    label = "n_plus_one"
+    formatter = "Potential n+1 query detected on `{model}.{field}`"
+
+
+class EagerLoadMessage(Message):
+    label = "unused_eager_load"
+    formatter = "Potential unnecessary eager load detected on `{model}.{field}`"
+
+
+class GetLoopMessage(Message):
+    label = "get_in_loop"
+    formatter = "Potential n+1 query detected on `{model}.{field}`"
+
+
+class Listener:
+    def __init__(self, parent: Any) -> None:
+        self.parent = parent
+
+    def setup(self) -> None:
+        pass
+
+    def teardown(self) -> None:
+        pass
+
+
+class LazyListener(Listener):
+    loaded: set[str]
+    ignore: set[str]
+    counts: defaultdict[tuple[type, str], int]
+    show_all_callers: bool
+    call_stacks: defaultdict[tuple[type, str], list[list[tuple[str, int, str]]]]
+
+    def setup(self) -> None:
+        from django.conf import settings
+
+        from django_nplus1 import signals
+
+        self.loaded = set()
+        self.ignore = set()
+        self.counts = defaultdict(int)
+        self.threshold = getattr(settings, "NPLUS1_THRESHOLD", 2)
+        self.show_all_callers = getattr(settings, "NPLUS1_SHOW_ALL_CALLERS", False)
+        self.call_stacks = defaultdict(list)
+        signals.connect(signals.LOAD, self.handle_load)
+        signals.connect(signals.IGNORE_LOAD, self.handle_ignore)
+        signals.connect(signals.LAZY_LOAD, self.handle_lazy)
+        signals.connect(signals.EAGER_LOAD, self.handle_eager)
+
+    def teardown(self) -> None:
+        from django_nplus1 import signals
+
+        signals.disconnect(signals.LOAD, self.handle_load)
+        signals.disconnect(signals.IGNORE_LOAD, self.handle_ignore)
+        signals.disconnect(signals.LAZY_LOAD, self.handle_lazy)
+        signals.disconnect(signals.EAGER_LOAD, self.handle_eager)
+
+    def handle_load(
+        self,
+        args: tuple[Any, ...] | None = None,
+        kwargs: dict[str, Any] | None = None,
+        context: dict[str, Any] | None = None,
+        ret: Any = None,
+        parser: Any = None,
+    ) -> None:
+        instances = parser(args, kwargs, context, ret)
+        self.loaded.update(instances)
+
+    def handle_ignore(
+        self,
+        args: tuple[Any, ...] | None = None,
+        kwargs: dict[str, Any] | None = None,
+        context: dict[str, Any] | None = None,
+        ret: Any = None,
+        parser: Any = None,
+    ) -> None:
+        instances = parser(args, kwargs, context, ret)
+        self.ignore.update(instances)
+
+    def handle_lazy(
+        self,
+        args: tuple[Any, ...] | None = None,
+        kwargs: dict[str, Any] | None = None,
+        context: dict[str, Any] | None = None,
+        ret: Any = None,
+        parser: Any = None,
+    ) -> None:
+        model, instance, field = parser(args, kwargs, context)
+        if instance in self.loaded and instance not in self.ignore:
+            key = (model, field)
+            self.counts[key] += 1
+            if self.show_all_callers:
+                from django_nplus1.util import get_stack
+
+                self.call_stacks[key].append(get_stack())
+            if self.counts[key] == self.threshold:
+                if self.show_all_callers:
+                    message = LazyLoadMessage(model, field, callers=self.call_stacks[key])
+                else:
+                    from django_nplus1.util import get_caller
+
+                    caller = get_caller()
+                    message = LazyLoadMessage(model, field, caller=caller)
+                self.parent.notify(message)
+
+    def handle_eager(
+        self,
+        args: tuple[Any, ...] | None = None,
+        kwargs: dict[str, Any] | None = None,
+        context: dict[str, Any] | None = None,
+        ret: Any = None,
+        parser: Any = None,
+    ) -> None:
+        # When a single bulk-loaded instance triggers an eager load (e.g.
+        # select_related on one item from a queryset), this is semantically
+        # an N+1 pattern, not an unused eager load. We use LazyLoadMessage
+        # with label "n_plus_one" intentionally.
+        model, field, keys, _key = parser(args, kwargs, context)
+        if len(keys) == 1 and keys[0] in self.loaded and keys[0] not in self.ignore:
+            key = (model, field)
+            self.counts[key] += 1
+            if self.show_all_callers:
+                from django_nplus1.util import get_stack
+
+                self.call_stacks[key].append(get_stack())
+            if self.counts[key] == self.threshold:
+                if self.show_all_callers:
+                    message = LazyLoadMessage(model, field, callers=self.call_stacks[key])
+                else:
+                    from django_nplus1.util import get_caller
+
+                    caller = get_caller()
+                    message = LazyLoadMessage(model, field, caller=caller)
+                self.parent.notify(message)
+
+
+class EagerListener(Listener):
+    tracker: EagerTracker
+    touched: list[tuple[type, str, list[str]] | None]
+
+    def setup(self) -> None:
+        from django_nplus1 import signals
+
+        self.tracker = EagerTracker()
+        self.touched = []
+        signals.connect(signals.EAGER_LOAD, self.handle_eager)
+        signals.connect(signals.TOUCH, self.handle_touch)
+
+    def teardown(self) -> None:
+        from django_nplus1 import signals
+
+        signals.disconnect(signals.EAGER_LOAD, self.handle_eager)
+        signals.disconnect(signals.TOUCH, self.handle_touch)
+        self.log_eager()
+
+    def handle_eager(
+        self,
+        args: tuple[Any, ...] | None = None,
+        kwargs: dict[str, Any] | None = None,
+        context: dict[str, Any] | None = None,
+        ret: Any = None,
+        parser: Any = None,
+    ) -> None:
+        self.tracker.track(*parser(args, kwargs, context))
+
+    def handle_touch(
+        self,
+        args: tuple[Any, ...] | None = None,
+        kwargs: dict[str, Any] | None = None,
+        context: dict[str, Any] | None = None,
+        ret: Any = None,
+        parser: Any = None,
+    ) -> None:
+        self.touched.append(parser(args, kwargs, context))
+
+    def log_eager(self) -> None:
+        self.tracker.prune([each for each in self.touched if each])
+        for model, field in self.tracker.unused:
+            message = EagerLoadMessage(model, field)
+            self.parent.notify(message)
+
+
+class EagerTracker:
+    def __init__(self) -> None:
+        self.data: defaultdict[tuple[type, str], defaultdict[int, set[str]]] = defaultdict(
+            lambda: defaultdict(set),
+        )
+
+    def track(self, model: type, field: str, instances: list[str], key: int) -> None:
+        self.data[(model, field)][key].update(instances)
+
+    def prune(self, touched: list[tuple[type, str, list[str]]]) -> None:
+        for model, field, touch_instances in touched:
+            group = self.data[(model, field)]
+            for key, fetch_instances in list(group.items()):
+                if touch_instances and fetch_instances.intersection(touch_instances):
+                    group.pop(key, None)
+
+    @property
+    def unused(self) -> list[tuple[type, str]]:
+        return [(model, field) for (model, field), group in self.data.items() if group]
+
+
+class GetLoopListener(Listener):
+    """Detects Model.objects.get() called repeatedly from the same call-site."""
+
+    counts: defaultdict[tuple[Any, ...], int]
+
+    def setup(self) -> None:
+        from django.conf import settings
+
+        from django_nplus1 import signals
+
+        self.counts = defaultdict(int)
+        self.threshold = getattr(settings, "NPLUS1_GET_THRESHOLD", 2)
+        signals.connect(signals.GET_CALL, self.handle_get)
+
+    def teardown(self) -> None:
+        from django_nplus1 import signals
+
+        signals.disconnect(signals.GET_CALL, self.handle_get)
+
+    def handle_get(
+        self,
+        args: tuple[Any, ...] | None = None,
+        kwargs: dict[str, Any] | None = None,
+        context: dict[str, Any] | None = None,
+        ret: Any = None,
+        parser: Any = None,
+    ) -> None:
+        model, caller = parser(args, kwargs, context, ret)
+        key = (model, *caller)
+        self.counts[key] += 1
+        if self.counts[key] == self.threshold:
+            message = GetLoopMessage(model, "get()", caller=caller)
+            self.parent.notify(message)
+
+
+class DuplicateQueryMessage(Message):
+    label = "duplicate_query"
+    formatter = "Potential n+1 query detected: duplicate query `{field}`"
+
+
+class _SQL:
+    pass
+
+
+_SQL_LITERAL_RE = re.compile(r"'(?:[^']|'')*'|\b\d+\b")
+
+
+def _fingerprint_sql(sql: str) -> str:
+    """Normalize a SQL query by replacing literals with ?, collapsing whitespace."""
+    return " ".join(_SQL_LITERAL_RE.sub("?", sql).split())
+
+
+class DuplicateQueryListener(Listener):
+    """Detects repeated identical SQL queries (N+1 fallback for raw SQL).
+
+    Uses Django's connection.execute_wrapper to intercept all queries,
+    fingerprints them by normalizing literals, and counts occurrences
+    from the same call-site. This catches N+1 patterns from raw SQL,
+    .raw(), and any ORM path not covered by the descriptor-level detection.
+    """
+
+    def __init__(self, parent: Any) -> None:
+        super().__init__(parent)
+        self.enabled = False
+        self.counts: defaultdict[tuple[str, str, int, str], int] = defaultdict(int)
+        self._connection: Any = None
+
+    def setup(self) -> None:
+        from django.conf import settings
+        from django.db import connection
+
+        self.enabled = getattr(settings, "NPLUS1_DETECT_DUPLICATE_QUERIES", False)
+        if not self.enabled:
+            return
+        self.counts = defaultdict(int)
+        self.threshold = getattr(settings, "NPLUS1_DUPLICATE_QUERY_THRESHOLD", 2)
+        self._connection = connection
+        self._connection.execute_wrappers.append(self._wrapper)
+
+    def teardown(self) -> None:
+        if not self.enabled:
+            return
+        try:
+            self._connection.execute_wrappers.remove(self._wrapper)
+        except ValueError:
+            pass
+
+    def _wrapper(self, execute: Any, sql: str, params: Any, many: bool, context: Any) -> Any:
+        result = execute(sql, params, many, context)
+        if not many:
+            from django_nplus1.util import get_caller
+
+            fingerprint = _fingerprint_sql(sql)
+            caller = get_caller()
+            key = (fingerprint, *caller)
+            self.counts[key] += 1
+            if self.counts[key] == self.threshold:
+                short_sql = fingerprint[:120] + ("..." if len(fingerprint) > 120 else "")
+                message = DuplicateQueryMessage(_SQL, short_sql, caller=caller)
+                self.parent.notify(message)
+        return result
+
+
+LISTENERS: dict[str, type[Listener]] = {
+    "lazy_load": LazyListener,
+    "eager_load": EagerListener,
+    "get_loop": GetLoopListener,
+    "duplicate_query": DuplicateQueryListener,
+}

django_nplus1/exceptions.py

@@ -0,0 +1,2 @@
+class NPlus1Error(Exception):
+    pass

django_nplus1/middleware.py

@@ -0,0 +1,159 @@
+from __future__ import annotations
+
+import fnmatch
+from typing import TYPE_CHECKING, Any
+
+from asgiref.sync import iscoroutinefunction
+from django.apps import apps
+from django.conf import settings
+from django.utils.decorators import sync_and_async_middleware
+
+from django_nplus1 import notifiers
+from django_nplus1.detect import LISTENERS, Message, Rule, is_allowed
+from django_nplus1.exceptions import NPlus1Error
+from django_nplus1.signals import nplus1_detected, setup_context, teardown_context
+
+if TYPE_CHECKING:
+    from django.http import HttpRequest, HttpResponse
+
+    from django_nplus1.detect import Listener
+
+_FNMATCH_CHARS = set("*?[]")
+
+
+def _validate_whitelist(whitelist: list[dict[str, Any]]) -> None:
+    """Validate whitelist entries against the Django model registry.
+
+    Raises NPlus1Error for invalid model or field names.
+    Skips entries that use fnmatch wildcards (* ? [ ]).
+    """
+    # Build registry: {"app_label.ModelName": set_of_field_names}
+    registry: dict[str, set[str]] = {}
+    for model in apps.get_models():
+        key = f"{model._meta.app_label}.{model.__name__}"
+        fields = {f.name for f in model._meta.get_fields(include_hidden=True)}
+        # Include reverse relation accessor names
+        fields |= {name for rel in model._meta.related_objects if (name := rel.get_accessor_name()) is not None}
+        registry[key] = fields
+
+    for entry in whitelist:
+        model_pattern = entry.get("model")
+        if not model_pattern:
+            continue
+
+        # Skip fnmatch patterns
+        if any(c in model_pattern for c in _FNMATCH_CHARS):
+            continue
+
+        if model_pattern not in registry:
+            suffix = model_pattern.split(".")[-1].lower()
+            similar = sorted(k for k in registry if suffix in k.lower())[:3]
+            msg = f"NPLUS1_WHITELIST: model '{model_pattern}' not found in installed Django models."
+            if similar:
+                msg += f" Did you mean one of: {', '.join(similar)}?"
+            raise NPlus1Error(msg)
+
+        field_name = entry.get("field")
+        if not field_name:
+            continue
+        if any(c in field_name for c in _FNMATCH_CHARS):
+            continue
+
+        if field_name not in registry[model_pattern]:
+            raise NPlus1Error(
+                f"NPLUS1_WHITELIST: field '{field_name}' not found on '{model_pattern}'",
+            )
+
+
+class DjangoRule(Rule):
+    def match_model(self, model: type) -> bool:
+        if self.model is model:
+            return True
+        if isinstance(self.model, str):
+            meta = getattr(model, "_meta", None)
+            if meta is None:
+                return False
+            return fnmatch.fnmatch(
+                f"{meta.app_label}.{model.__name__}",
+                self.model,
+            )
+        return False
+
+
+_last_validated_whitelist: tuple[int, int] | None = None
+
+
+def _load_config() -> tuple[list[notifiers.Notifier], list[DjangoRule]]:
+    """Load notifiers and whitelist from settings."""
+    global _last_validated_whitelist  # noqa: PLW0603
+    nots = notifiers.init(settings)
+    whitelist_data = getattr(settings, "NPLUS1_WHITELIST", [])
+    cache_key = (id(whitelist_data), len(whitelist_data))
+    if cache_key != _last_validated_whitelist:
+        _validate_whitelist(whitelist_data)
+        _last_validated_whitelist = cache_key
+    whitelist = [DjangoRule(**item) for item in whitelist_data]
+    return nots, whitelist
+
+
+class _DetectionContext:
+    """Manages listener setup/teardown for a single request scope."""
+
+    def __init__(self, nots: list[notifiers.Notifier], whitelist: list[DjangoRule]) -> None:
+        self.nots = nots
+        self.whitelist = whitelist
+        self.listeners: dict[str, Listener] = {}
+
+    def setup(self) -> None:
+        for name, listener_type in LISTENERS.items():
+            self.listeners[name] = listener_type(self)
+            self.listeners[name].setup()
+
+    def teardown(self) -> None:
+        for name in list(LISTENERS.keys()):
+            listener = self.listeners.pop(name, None)
+            if listener:
+                listener.teardown()
+
+    def notify(self, message: Message) -> None:
+        if not message.match(self.whitelist) and not is_allowed(message):
+            nplus1_detected.send(sender=NPlus1Middleware, message=message)
+            for notifier in self.nots:
+                notifier.notify(message)
+
+
+@sync_and_async_middleware
+def NPlus1Middleware(get_response: Any) -> Any:  # noqa: N802
+    if iscoroutinefunction(get_response):
+
+        async def async_middleware(request: HttpRequest) -> HttpResponse:
+            token = setup_context()
+            try:
+                nots, whitelist = _load_config()
+                ctx = _DetectionContext(nots, whitelist)
+                ctx.setup()
+                try:
+                    response = await get_response(request)
+                finally:
+                    ctx.teardown()
+            finally:
+                teardown_context(token)
+            return response
+
+        return async_middleware
+
+    def sync_middleware(request: HttpRequest) -> HttpResponse:
+        token = setup_context()
+        try:
+            nots, whitelist = _load_config()
+            ctx = _DetectionContext(nots, whitelist)
+            ctx.setup()
+            try:
+                response = get_response(request)
+            finally:
+                ctx.teardown()
+        finally:
+            teardown_context(token)
+        return response
+
+    return sync_middleware