瀏覽代碼

Closes #23565: Batch deferred field loads in GraphQL queries with `FETCH_PEERS`

Jeremy Stretch 17 小時之前
父節點
當前提交
b9cf4df16d
共有 2 個文件被更改,包括 76 次插入 和 5 次删除
  1. 9 4
      netbox/netbox/graphql/types.py
  2. 67 1
      netbox/netbox/tests/test_graphql.py

+ 9 - 4
netbox/netbox/graphql/types.py

@@ -3,7 +3,7 @@ from copy import copy
 import strawberry
 import strawberry
 import strawberry_django
 import strawberry_django
 from django.contrib.contenttypes.models import ContentType
 from django.contrib.contenttypes.models import ContentType
-from django.db.models import BinaryField, ExpressionWrapper, F, Func, IntegerField, Value
+from django.db.models import FETCH_PEERS, BinaryField, ExpressionWrapper, F, Func, IntegerField, Value
 from strawberry.types import Info
 from strawberry.types import Info
 from strawberry_django.optimizer import OptimizerStore
 from strawberry_django.optimizer import OptimizerStore
 
 
@@ -41,7 +41,7 @@ def _model_type(model, **kwargs):
     """
     """
     Wrap `strawberry_django.type()` to give each type's `display` field an optimizer hint. `display` resolves via
     Wrap `strawberry_django.type()` to give each type's `display` field an optimizer hint. `display` resolves via
     `__str__()`, which may read any of the model's columns; without a hint the optimizer defers them, and reading a
     `__str__()`, which may read any of the model's columns; without a hint the optimizer defers them, and reading a
-    deferred column reloads the row with one query per object.
+    deferred column costs an extra query.
     """
     """
     def wrapper(cls):
     def wrapper(cls):
         graphql_type = strawberry_django.type(model, **kwargs)(cls)
         graphql_type = strawberry_django.type(model, **kwargs)(cls)
@@ -66,14 +66,19 @@ def _model_type(model, **kwargs):
 @strawberry.type
 @strawberry.type
 class BaseObjectType:
 class BaseObjectType:
     """
     """
-    Base GraphQL object type for all NetBox objects. Restricts the model queryset to enforce object permissions.
+    Base GraphQL object type for all NetBox objects. Restricts the model queryset to enforce object permissions, and
+    applies the FETCH_PEERS fetch mode to it.
     """
     """
 
 
     @classmethod
     @classmethod
     def get_queryset(cls, queryset, info: Info, **kwargs):
     def get_queryset(cls, queryset, info: Info, **kwargs):
         # Enforce object permissions on the queryset
         # Enforce object permissions on the queryset
         if hasattr(queryset, 'restrict'):
         if hasattr(queryset, 'restrict'):
-            return queryset.restrict(info.context.request.user, 'view')
+            queryset = queryset.restrict(info.context.request.user, 'view')
+        # Batch the loading of deferred fields and relations which the optimizer can't anticipate. Skip a queryset
+        # already evaluated by a prefetch_related hint (e.g. tags): cloning it would discard the prefetched results.
+        if queryset._result_cache is None:
+            queryset = queryset.fetch_mode(FETCH_PEERS)
         return queryset
         return queryset
 
 
     @strawberry_django.field
     @strawberry_django.field

+ 67 - 1
netbox/netbox/tests/test_graphql.py

@@ -21,6 +21,7 @@ from strawberry_django.optimizer import DjangoOptimizerExtension
 from core.models import DataFile, DataSource, ObjectType
 from core.models import DataFile, DataSource, ObjectType
 from dcim.choices import LocationStatusChoices
 from dcim.choices import LocationStatusChoices
 from dcim.models import (
 from dcim.models import (
+    Cable,
     Device,
     Device,
     DeviceRole,
     DeviceRole,
     DeviceType,
     DeviceType,
@@ -1039,6 +1040,71 @@ class GraphQLAPITestCase(APITestCase):
             msg=f'Expected batched MAC address lookup, got {mac_queries} MAC address queries for 5 devices',
             msg=f'Expected batched MAC address lookup, got {mac_queries} MAC address queries for 5 devices',
         )
         )
 
 
+    def test_graphql_deferred_fields_are_batched(self):
+        """
+        Deferred columns and relations read by a resolver the optimizer can't analyze (`link_peers` reads `cable`
+        and `wireless_link`) must load with one query for all objects returned, not one query per object.
+        """
+        self.add_permissions('dcim.view_interface')
+
+        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')
+        device_role = DeviceRole.objects.create(name='Device Role 1', slug='device-role-1')
+        devices = Device.objects.bulk_create([
+            Device(name=f'Device {index}', device_type=device_type, role=device_role, site=Site.objects.first())
+            for index in range(1, 4)
+        ])
+        interfaces = {
+            device.name: Interface.objects.bulk_create([
+                Interface(device=device, name=f'eth{index}', type='1000baset') for index in range(5)
+            ])
+            for device in devices
+        }
+        # Device 1's interfaces have no cable; device 2's are each cabled to the interface of the same name on device 3
+        for a_interface, b_interface in zip(interfaces['Device 2'], interfaces['Device 3']):
+            Cable(a_terminations=[a_interface], b_terminations=[b_interface]).save()
+
+        query = """
+        {
+            interface_list(filters: {device: {name: {exact: "%s"}}}, pagination: {limit: %d}) {
+                name
+                link_peers { ... on InterfaceType { name } }
+            }
+        }
+        """
+        url = reverse('graphql')
+
+        def get_link_peers(device_name, limit):
+            with CaptureQueriesContext(connection) as context:
+                response = self.client.post(
+                    url, data={'query': query % (device_name, limit)}, format='json', **self.header
+                )
+            self.assertHttpStatus(response, status.HTTP_200_OK)
+            data = json.loads(response.content)
+            self.assertNotIn('errors', data)
+            link_peers = {
+                interface['name']: [peer['name'] for peer in interface['link_peers']]
+                for interface in data['data']['interface_list']
+            }
+            self.assertEqual(len(link_peers), limit)
+            return link_peers, context.captured_queries
+
+        # Deferred columns (cable, wireless_link) are loaded once, however many interfaces are returned
+        interface_queries = {}
+        for limit in (1, 5):
+            link_peers, queries = get_link_peers('Device 1', limit)
+            self.assertEqual(link_peers, {f'eth{index}': [] for index in range(limit)})
+            interface_queries[limit] = count_primary_table_queries(queries, 'dcim_interface')
+        self.assertEqual(interface_queries[5], interface_queries[1])
+
+        # Cables are loaded once, however many interfaces are returned
+        cable_queries = {}
+        for limit in (1, 5):
+            link_peers, queries = get_link_peers('Device 2', limit)
+            self.assertEqual(link_peers, {f'eth{index}': [f'eth{index}'] for index in range(limit)})
+            cable_queries[limit] = count_primary_table_queries(queries, 'dcim_cable')
+        self.assertEqual(cable_queries[5], cable_queries[1])
+
     def test_offset_pagination(self):
     def test_offset_pagination(self):
         self.add_permissions('dcim.view_site')
         self.add_permissions('dcim.view_site')
         url = reverse('graphql')
         url = reverse('graphql')
@@ -1451,7 +1517,7 @@ class GraphQLDeferredColumnTestCase(APITestCase):
     A GraphQL field backed by a custom resolver is opaque to the query optimizer, which narrows column
     A GraphQL field backed by a custom resolver is opaque to the query optimizer, which narrows column
     selection with .only() based on the fields named in the GraphQL document. Any column such a resolver
     selection with .only() based on the fields named in the GraphQL document. Any column such a resolver
     reads must therefore be declared via an `only` hint; otherwise the column is deferred and reading it
     reads must therefore be declared via an `only` hint; otherwise the column is deferred and reading it
-    reloads the row from the database once per object returned (see #22813).
+    costs an extra query (see #22813).
 
 
     Each test below asserts both that no single-row reload occurs and that the model's table is read
     Each test below asserts both that no single-row reload occurs and that the model's table is read
     exactly once per request, regardless of the number of objects returned.
     exactly once per request, regardless of the number of objects returned.