Explorar o código

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

Jeremy Stretch hai 1 semana
pai
achega
d76340f55e
Modificáronse 2 ficheiros con 224 adicións e 1 borrados
  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
 import strawberry
+from django.db import DEFAULT_DB_ALIAS
+from django.db.models.functions import DenseRank
 from strawberry.types.unset import UNSET
 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
 from netbox.config import get_config
 
 
 __all__ = (
 __all__ = (
     'OffsetPaginationInfo',
     'OffsetPaginationInfo',
     'OffsetPaginationInput',
     'OffsetPaginationInput',
+    'apply_distinct_window_pagination',
     'apply_pagination',
     'apply_pagination',
 )
 )
 
 
@@ -26,6 +29,58 @@ class OffsetPaginationInput(OffsetPaginationInfo):
     pass
     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(
 def apply_pagination(
     self,
     self,
     queryset: _QS,
     queryset: _QS,
@@ -71,4 +126,14 @@ def apply_pagination(
         elif pagination.limit <= 0:
         elif pagination.limit <= 0:
             pagination.limit = max_page_size
             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)
     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.choices import CustomFieldTypeChoices
 from extras.models import CustomField, TableConfig, Tag
 from extras.models import CustomField, TableConfig, Tag
 from ipam.models import RIR, Aggregate, IPAddress, Prefix
 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.scalars import BigInt, BigIntScalar
 from netbox.graphql.schema import Query, get_schema_extensions, schema
 from netbox.graphql.schema import Query, get_schema_extensions, schema
 from users.models import Token, User
 from users.models import Token, User
@@ -878,6 +879,163 @@ class GraphQLAPITestCase(APITestCase):
         self.assertNotIn('errors', data)
         self.assertNotIn('errors', data)
         self.assertEqual(len(data['data']['site_list'][0]['devices']), 3)
         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):
     def test_pagination_conflict(self):
         url = reverse('graphql')
         url = reverse('graphql')
         query = """
         query = """