Просмотр исходного кода

Fixes #22882: Fix support for DISTINCT on nested GraphQL lists (#22892)

Jeremy Stretch 1 неделя назад
Родитель
Сommit
d76340f55e
2 измененных файлов с 224 добавлено и 1 удалено
  1. 66 1
      netbox/netbox/graphql/pagination.py
  2. 158 0
      netbox/netbox/tests/test_graphql.py

+ 66 - 1
netbox/netbox/graphql/pagination.py

@@ -1,12 +1,15 @@
 import strawberry
+from django.db import DEFAULT_DB_ALIAS
+from django.db.models.functions import DenseRank
 from strawberry.types.unset import UNSET
-from strawberry_django.pagination import _QS, apply
+from strawberry_django.pagination import _QS, _PaginationWindow, _resolve_limit, apply
 
 from netbox.config import get_config
 
 __all__ = (
     'OffsetPaginationInfo',
     'OffsetPaginationInput',
+    'apply_distinct_window_pagination',
     'apply_pagination',
 )
 
@@ -26,6 +29,58 @@ class OffsetPaginationInput(OffsetPaginationInfo):
     pass
 
 
+def apply_distinct_window_pagination(
+    queryset: _QS,
+    *,
+    related_field_id: str,
+    offset: int = 0,
+    limit: int | None = UNSET,
+) -> _QS:
+    """
+    Replacement for strawberry-django's `apply_window_pagination()` for a queryset which has `DISTINCT`
+    enabled, as is the case when a list field is filtered across a to-many relation with `DISTINCT: true`.
+
+    SQL evaluates window functions before `DISTINCT`, so the `ROW_NUMBER()` annotation which
+    strawberry-django uses to paginate a prefetched relation assigns a unique value to each of the
+    duplicate rows produced by the join, and `DISTINCT` can never collapse them. `DENSE_RANK()` instead
+    assigns the same rank to every row which compares equal under the window ordering, leaving the
+    duplicate rows identical so that `DISTINCT` deduplicates them as intended. And because the rank is
+    incremented only once per distinct row, the rows are numbered as if the duplicates were never there,
+    keeping the pagination limit meaningful.
+    """
+    limit = _resolve_limit(limit)
+
+    order_by = [
+        expr
+        for expr, _ in queryset.query.get_compiler(
+            using=queryset._db or DEFAULT_DB_ALIAS
+        ).get_order_by()
+    ]
+    # Order by the primary key as well, to ensure that two rows representing *different* objects can
+    # never be assigned the same rank (and hence be counted only once against the limit).
+    order_by.append('pk')
+
+    # Note that we omit the `_strawberry_total_count` annotation which strawberry-django adds, as it
+    # cannot be made accurate here: window functions are evaluated before `DISTINCT`, so it would count
+    # the duplicate rows. strawberry-django's `get_total_count()` already disregards the annotation for
+    # a queryset with `DISTINCT` enabled and falls back to `count()`, so computing it would be wasted
+    # work: an extra window aggregate over every joined row.
+    queryset = queryset.annotate(
+        _strawberry_row_number=_PaginationWindow(
+            DenseRank(),
+            partition_by=related_field_id,
+            order_by=order_by,
+        ),
+    )
+
+    if offset:
+        queryset = queryset.filter(_strawberry_row_number__gt=offset)
+    if limit is not None and limit >= 0:
+        queryset = queryset.filter(_strawberry_row_number__lte=offset + limit)
+
+    return queryset
+
+
 def apply_pagination(
     self,
     queryset: _QS,
@@ -71,4 +126,14 @@ def apply_pagination(
         elif pagination.limit <= 0:
             pagination.limit = max_page_size
 
+    # A prefetched relation is paginated with a window function, which is incompatible with the
+    # `DISTINCT` applied by the filter layer. Fall back to our own implementation in that case.
+    if pagination is not None and related_field_id is not None and queryset.query.distinct:
+        return apply_distinct_window_pagination(
+            queryset,
+            related_field_id=related_field_id,
+            offset=pagination.offset,
+            limit=pagination.limit,
+        )
+
     return apply(pagination, queryset, related_field_id=related_field_id)

+ 158 - 0
netbox/netbox/tests/test_graphql.py

@@ -29,6 +29,7 @@ from dcim.models import (
 from extras.choices import CustomFieldTypeChoices
 from extras.models import CustomField, TableConfig, Tag
 from ipam.models import RIR, Aggregate, IPAddress, Prefix
+from netbox.graphql.pagination import apply_distinct_window_pagination
 from netbox.graphql.scalars import BigInt, BigIntScalar
 from netbox.graphql.schema import Query, get_schema_extensions, schema
 from users.models import Token, User
@@ -878,6 +879,163 @@ class GraphQLAPITestCase(APITestCase):
         self.assertNotIn('errors', data)
         self.assertEqual(len(data['data']['site_list'][0]['devices']), 3)
 
+    def test_distinct_nested_list(self):
+        """
+        The `DISTINCT` filter should deduplicate a nested list field which is filtered across a to-many
+        relation, just as it does for the equivalent top-level list field.
+        """
+        self.add_permissions('dcim.view_device', 'dcim.view_site')
+        url = reverse('graphql')
+
+        site = Site.objects.get(slug='site-1')
+        manufacturer = Manufacturer.objects.create(name='Manufacturer 1', slug='manufacturer-1')
+        device_type = DeviceType.objects.create(manufacturer=manufacturer, model='Device Type 1', slug='device-type-1')
+        role = DeviceRole.objects.create(name='Device Role 1', slug='device-role-1')
+        devices = Device.objects.bulk_create([
+            Device(name=f'Device {i}', site=site, device_type=device_type, role=role)
+            for i in range(1, 3)
+        ])
+        Interface.objects.bulk_create([
+            Interface(device=device, name=f'eth{i}', type='1000base-t')
+            for device in devices
+            for i in range(3)
+        ])
+
+        # Each device should be returned exactly once, despite having three matching interfaces
+        query = """
+        {
+            site_list(filters: {slug: {exact: "site-1"}}) {
+                name
+                devices(filters: {DISTINCT: true, interfaces: {name: {starts_with: "eth"}}}) {
+                    name
+                }
+            }
+        }
+        """
+        response = self.client.post(url, data={'query': query}, format='json', **self.header)
+        self.assertHttpStatus(response, status.HTTP_200_OK)
+        data = json.loads(response.content)
+        self.assertNotIn('errors', data)
+        self.assertEqual(
+            [device['name'] for device in data['data']['site_list'][0]['devices']],
+            ['Device 1', 'Device 2']
+        )
+
+        # The equivalent top-level query should return the same devices
+        query = """
+        {
+            device_list(filters: {DISTINCT: true, interfaces: {name: {starts_with: "eth"}}}) {
+                name
+            }
+        }
+        """
+        response = self.client.post(url, data={'query': query}, format='json', **self.header)
+        self.assertHttpStatus(response, status.HTTP_200_OK)
+        data = json.loads(response.content)
+        self.assertNotIn('errors', data)
+        self.assertEqual(
+            [device['name'] for device in data['data']['device_list']],
+            ['Device 1', 'Device 2']
+        )
+
+    @override_settings(MAX_PAGE_SIZE=2)
+    def test_distinct_nested_list_max_page_size(self):
+        """
+        MAX_PAGE_SIZE should still be enforced on a deduplicated nested list field, and should be applied
+        to the number of distinct objects returned (not to the number of joined rows).
+        """
+        self.add_permissions('dcim.view_device', 'dcim.view_site')
+        url = reverse('graphql')
+
+        site = Site.objects.get(slug='site-1')
+        manufacturer = Manufacturer.objects.create(name='Manufacturer 1', slug='manufacturer-1')
+        device_type = DeviceType.objects.create(manufacturer=manufacturer, model='Device Type 1', slug='device-type-1')
+        role = DeviceRole.objects.create(name='Device Role 1', slug='device-role-1')
+        devices = Device.objects.bulk_create([
+            Device(name=f'Device {i}', site=site, device_type=device_type, role=role)
+            for i in range(1, 5)
+        ])
+        Interface.objects.bulk_create([
+            Interface(device=device, name=f'eth{i}', type='1000base-t')
+            for device in devices
+            for i in range(3)
+        ])
+
+        query = """
+        {
+            site_list(filters: {slug: {exact: "site-1"}}) {
+                name
+                devices(filters: {DISTINCT: true, interfaces: {name: {starts_with: "eth"}}}) {
+                    name
+                }
+            }
+        }
+        """
+        response = self.client.post(url, data={'query': query}, format='json', **self.header)
+        self.assertHttpStatus(response, status.HTTP_200_OK)
+        data = json.loads(response.content)
+        self.assertNotIn('errors', data)
+        self.assertEqual(
+            [device['name'] for device in data['data']['site_list'][0]['devices']],
+            ['Device 1', 'Device 2']
+        )
+
+        # An explicit offset should likewise be applied to the distinct objects
+        query = """
+        {
+            site_list(filters: {slug: {exact: "site-1"}}) {
+                name
+                devices(
+                    pagination: {offset: 1, limit: 2},
+                    filters: {DISTINCT: true, interfaces: {name: {starts_with: "eth"}}}
+                ) {
+                    name
+                }
+            }
+        }
+        """
+        response = self.client.post(url, data={'query': query}, format='json', **self.header)
+        self.assertHttpStatus(response, status.HTTP_200_OK)
+        data = json.loads(response.content)
+        self.assertNotIn('errors', data)
+        self.assertEqual(
+            [device['name'] for device in data['data']['site_list'][0]['devices']],
+            ['Device 2', 'Device 3']
+        )
+
+    def test_distinct_window_pagination_tied_ordering(self):
+        """
+        Two rows which represent *different* objects must never be assigned the same rank, even when they
+        compare equal under the queryset's ordering. `DENSE_RANK()` ties such rows by definition, so the
+        primary key is appended to the window ordering to separate them; without it every device below
+        would be assigned rank 1 and the limit of two would return all four of them.
+        """
+        site = Site.objects.get(slug='site-1')
+        manufacturer = Manufacturer.objects.create(name='Manufacturer 1', slug='manufacturer-1')
+        device_type = DeviceType.objects.create(manufacturer=manufacturer, model='Device Type 1', slug='device-type-1')
+        role = DeviceRole.objects.create(name='Device Role 1', slug='device-role-1')
+        devices = Device.objects.bulk_create([
+            Device(name=f'Device {i}', site=site, device_type=device_type, role=role)
+            for i in range(1, 5)
+        ])
+        Interface.objects.bulk_create([
+            Interface(device=device, name=f'eth{i}', type='1000base-t')
+            for device in devices
+            for i in range(3)
+        ])
+
+        # Order by a column whose value is identical for every device, so that the ordering alone cannot
+        # distinguish them. Each device additionally matches three interfaces, so the join emits three
+        # duplicate rows per device which DISTINCT must still collapse.
+        queryset = Device.objects.filter(
+            site=site, interfaces__name__startswith='eth'
+        ).order_by('status').distinct()
+        queryset = apply_distinct_window_pagination(queryset, related_field_id='site_id', limit=2)
+
+        results = list(queryset)
+        self.assertEqual(sorted(device.name for device in results), ['Device 1', 'Device 2'])
+        self.assertEqual(sorted(device._strawberry_row_number for device in results), [1, 2])
+
     def test_pagination_conflict(self):
         url = reverse('graphql')
         query = """