Coverage for netbox/graphql/optimization.py: 61%

1from collections.abc import Callable, Sequence 

2 

3from django.contrib.contenttypes.prefetch import GenericPrefetch 

4from django.db.models import Model, QuerySet 

5from strawberry.types import Info 

6from strawberry_django.optimizer import optimize 

7from strawberry_django.optimizer import optimizer as optimizer_ctx 

8 

9__all__ = ( 

10 'build_gfk_prefetch', 

11 'optimize_prefetch_queryset', 

12) 

13 

14 

15def optimize_prefetch_queryset(queryset: QuerySet, info: Info) -> QuerySet: 

16 """ 

17 Apply strawberry-django's query optimizer to a queryset used inside a GenericForeignKey prefetch. 

18 """ 

19 if ext := optimizer_ctx.get(): 

20 return ext.optimize(queryset, info) 

21 

22 return optimize(queryset, info) 

23 

24 

25def build_gfk_prefetch( 

26 lookup: str, 

27 models: Sequence[type[Model]], 

28) -> Callable[[Info], GenericPrefetch]: 

29 """ 

30 Return a selection-aware GenericPrefetch for a GenericForeignKey field. 

31 

32 Each model gets its own queryset, optimized according to the client's GraphQL selection set. 

33 """ 

34 

35 def prefetch(info: Info) -> GenericPrefetch: 

36 querysets = [ 

37 optimize_prefetch_queryset(model.objects.all(), info) 

38 for model in models 

39 ] 

40 

41 return GenericPrefetch(lookup, querysets) 

42 

43 return prefetch