Coverage for netbox/graphql/pagination.py: 25%

1import strawberry 

2from django.db import DEFAULT_DB_ALIAS 

3from django.db.models.functions import DenseRank 

4from strawberry.types.unset import UNSET 

5from strawberry_django.pagination import _QS, _PaginationWindow, _resolve_limit, apply 

6 

7from netbox.config import get_config 

8 

9__all__ = ( 

10 'OffsetPaginationInfo', 

11 'OffsetPaginationInput', 

12 'apply_distinct_window_pagination', 

13 'apply_pagination', 

14) 

15 

16 

17@strawberry.type 

18class OffsetPaginationInfo: 

19 offset: int = 0 

20 limit: int | None = UNSET 

21 start: int | None = UNSET 

22 

23 

24@strawberry.input 

25class OffsetPaginationInput(OffsetPaginationInfo): 

26 """ 

27 Customized implementation of OffsetPaginationInput to support cursor-based pagination. 

28 """ 

29 pass 

30 

31 

32def apply_distinct_window_pagination( 

33 queryset: _QS, 

34 *, 

35 related_field_id: str, 

36 offset: int = 0, 

37 limit: int | None = UNSET, 

38) -> _QS: 

39 """ 

40 Replacement for strawberry-django's `apply_window_pagination()` for a queryset which has `DISTINCT` 

41 enabled, as is the case when a list field is filtered across a to-many relation with `DISTINCT: true`. 

42 

43 SQL evaluates window functions before `DISTINCT`, so the `ROW_NUMBER()` annotation which 

44 strawberry-django uses to paginate a prefetched relation assigns a unique value to each of the 

45 duplicate rows produced by the join, and `DISTINCT` can never collapse them. `DENSE_RANK()` instead 

46 assigns the same rank to every row which compares equal under the window ordering, leaving the 

47 duplicate rows identical so that `DISTINCT` deduplicates them as intended. And because the rank is 

48 incremented only once per distinct row, the rows are numbered as if the duplicates were never there, 

49 keeping the pagination limit meaningful. 

50 """ 

51 limit = _resolve_limit(limit) 

52 

53 order_by = [ 

54 expr 

55 for expr, _ in queryset.query.get_compiler( 

56 using=queryset._db or DEFAULT_DB_ALIAS 

57 ).get_order_by() 

58 ] 

59 # Order by the primary key as well, to ensure that two rows representing *different* objects can 

60 # never be assigned the same rank (and hence be counted only once against the limit). 

61 order_by.append('pk') 

62 

63 # Note that we omit the `_strawberry_total_count` annotation which strawberry-django adds, as it 

64 # cannot be made accurate here: window functions are evaluated before `DISTINCT`, so it would count 

65 # the duplicate rows. strawberry-django's `get_total_count()` already disregards the annotation for 

66 # a queryset with `DISTINCT` enabled and falls back to `count()`, so computing it would be wasted 

67 # work: an extra window aggregate over every joined row. 

68 queryset = queryset.annotate( 

69 _strawberry_row_number=_PaginationWindow( 

70 DenseRank(), 

71 partition_by=related_field_id, 

72 order_by=order_by, 

73 ), 

74 ) 

75 

76 if offset: 

77 queryset = queryset.filter(_strawberry_row_number__gt=offset) 

78 if limit is not None and limit >= 0: 

79 queryset = queryset.filter(_strawberry_row_number__lte=offset + limit) 

80 

81 return queryset 

82 

83 

84def apply_pagination( 

85 self, 

86 queryset: _QS, 

87 pagination: OffsetPaginationInput | None = None, 

88 *, 

89 related_field_id: str | None = None, 

90) -> _QS: 

91 """ 

92 Replacement for the `apply_pagination()` method on StrawberryDjangoField to support cursor-based pagination. 

93 """ 

94 if pagination is not None and pagination.start not in (None, UNSET): 

95 if pagination.offset: 

96 raise ValueError('Cannot specify both `start` and `offset` in pagination.') 

97 if pagination.start < 0: 

98 raise ValueError('`start` must be greater than or equal to zero.') 

99 

100 # Filter the queryset to include only records with a primary key greater than or equal to the start value, 

101 # and force ordering by primary key to ensure consistent pagination across all records. 

102 queryset = queryset.filter(pk__gte=pagination.start).order_by('pk') 

103 

104 # Ignore `offset` when `start` is set 

105 pagination.offset = 0 

106 

107 # Enforce MAX_PAGE_SIZE on the pagination limit 

108 max_page_size = get_config().MAX_PAGE_SIZE 

109 if max_page_size: 

110 # A limit is meaningless for a field which returns at most one object, and synthesizing one for a 

111 # prefetched to-one relation is actively harmful. strawberry-django deliberately leaves `pagination` 

112 # as None there so that the prefetch remains a plain `WHERE id IN (...)` query; making it non-None 

113 # switches the prefetch to a window function partitioned by the parent ID. Every partition then 

114 # holds exactly one row, so ROW_NUMBER() is 1 throughout and the row number filter discards nothing, 

115 # causing the join back to the parent table to return every row which shares the related object. 

116 # See strawberry-graphql/strawberry-django#719. 

117 returns_single_object = not (self.is_list or self.is_paginated or self.is_connection) 

118 

119 if pagination is None: 

120 # Note that `pagination` is never None for a single-object field unless it is a prefetched 

121 # relation: strawberry-django populates it with an implicit limit of its own beforehand. 

122 if not returns_single_object: 

123 pagination = OffsetPaginationInput(limit=max_page_size) 

124 elif pagination.limit in (None, UNSET) or pagination.limit > max_page_size: 

125 pagination.limit = max_page_size 

126 elif pagination.limit <= 0: 

127 pagination.limit = max_page_size 

128 

129 # A prefetched relation is paginated with a window function, which is incompatible with the 

130 # `DISTINCT` applied by the filter layer. Fall back to our own implementation in that case. 

131 if pagination is not None and related_field_id is not None and queryset.query.distinct: 

132 return apply_distinct_window_pagination( 

133 queryset, 

134 related_field_id=related_field_id, 

135 offset=pagination.offset, 

136 limit=pagination.limit, 

137 ) 

138 

139 return apply(pagination, queryset, related_field_id=related_field_id)