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

Fixes #23000: Prefetch cable terminations in the GraphQL API

CableType.a_terminations and b_terminations were declared as bare annotations,
so strawberry-django resolved them by reading the model properties with no
prefetch hint. That left two nested N+1s (one query per Cable for
terminations, one per CableTermination for the termination GFK) plus the
termination's own device FK chain, for roughly six queries per termination.

Resolve both fields via resolvers carrying a Prefetch of the terminations for
that cable end, with the termination GFK prefetched through the existing
build_gfk_prefetch() helper so the nested joins are derived from the client's
selection set rather than hard-coded.

Each end is prefetched under its own to_attr: two prefetches of the same
relation cannot be merged by the query optimizer, so a shared lookup would
break any query selecting both ends.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Jeremy Stretch 22 часов назад
Родитель
Сommit
2eb9bb68d2
2 измененных файлов с 178 добавлено и 19 удалено
  1. 66 19
      netbox/dcim/graphql/types.py
  2. 112 0
      netbox/dcim/tests/test_api.py

+ 66 - 19
netbox/dcim/graphql/types.py

@@ -2,11 +2,12 @@ from typing import TYPE_CHECKING, Annotated
 
 import strawberry
 import strawberry_django
-from django.db.models import Func, IntegerField
+from django.db.models import Func, IntegerField, Prefetch
 
 from circuits.models import CircuitTermination
 from core.graphql.mixins import ChangelogMixin
 from dcim import models
+from dcim.choices import CableEndChoices
 from extras.graphql.mixins import ConfigContextMixin, ContactsMixin, ImageAttachmentsMixin
 from ipam.graphql.mixins import IPAddressesMixin, VLANGroupsMixin
 from netbox.graphql.optimization import build_gfk_prefetch
@@ -94,6 +95,57 @@ __all__ = (
 )
 
 
+#
+# Cable termination prefetching
+#
+
+# The concrete models which may terminate a cable, mirroring dcim.constants.CABLE_TERMINATION_MODELS
+_CABLE_TERMINATION_MODELS = (
+    CircuitTermination,
+    models.ConsolePort,
+    models.ConsoleServerPort,
+    models.FrontPort,
+    models.Interface,
+    models.PowerFeed,
+    models.PowerOutlet,
+    models.PowerPort,
+    models.RearPort,
+)
+
+_termination_gfk_prefetch = build_gfk_prefetch('termination', _CABLE_TERMINATION_MODELS)
+
+
+def _cable_terminations_prefetch(side, to_attr):
+    """
+    Return a callable which builds a selection-aware Prefetch of a cable's terminations for the
+    given cable end.
+
+    Each end is prefetched under its own `to_attr`: two prefetches of the same relation cannot be
+    merged by the query optimizer, so a shared lookup would break any query selecting both ends.
+    """
+    def prefetch(info):
+        return Prefetch(
+            'terminations',
+            queryset=models.CableTermination.objects.filter(cable_end=side).prefetch_related(
+                _termination_gfk_prefetch(info)
+            ),
+            to_attr=to_attr,
+        )
+
+    return prefetch
+
+
+def _resolve_cable_terminations(cable, side, to_attr):
+    """
+    Return the terminating objects for the given cable end, using the prefetched terminations if
+    available and falling back to the model property otherwise.
+    """
+    if (terminations := getattr(cable, to_attr, None)) is not None:
+        return [ct.termination for ct in terminations]
+
+    return cable._get_x_terminations(side)
+
+
 #
 # Base types
 #
@@ -153,20 +205,7 @@ class CableTerminationType(NetBoxObjectType):
     cable: Annotated['CableType', strawberry.lazy('dcim.graphql.types')] | None
 
     @strawberry_django.field(
-        prefetch_related=build_gfk_prefetch(
-            'termination',
-            [
-                CircuitTermination,
-                models.ConsolePort,
-                models.ConsoleServerPort,
-                models.FrontPort,
-                models.Interface,
-                models.PowerFeed,
-                models.PowerOutlet,
-                models.PowerPort,
-                models.RearPort,
-            ],
-        ),
+        prefetch_related=_termination_gfk_prefetch,
         only=['termination_type', 'termination_id'],
     )
     def termination(self) -> Annotated[
@@ -197,7 +236,10 @@ class CableType(PrimaryObjectType):
 
     terminations: list[CableTerminationType]
 
-    a_terminations: list[Annotated[
+    @strawberry_django.field(
+        prefetch_related=_cable_terminations_prefetch(CableEndChoices.SIDE_A, '_prefetched_a_terminations'),
+    )
+    def a_terminations(self) -> list[Annotated[
         Annotated['CircuitTerminationType', strawberry.lazy('circuits.graphql.types')]
         | Annotated['ConsolePortType', strawberry.lazy('dcim.graphql.types')]
         | Annotated['ConsoleServerPortType', strawberry.lazy('dcim.graphql.types')]
@@ -208,9 +250,13 @@ class CableType(PrimaryObjectType):
         | Annotated['PowerPortType', strawberry.lazy('dcim.graphql.types')]
         | Annotated['RearPortType', strawberry.lazy('dcim.graphql.types')],
         strawberry.union('CableTerminationTerminationType'),
-    ]]
+    ]]:
+        return _resolve_cable_terminations(self, CableEndChoices.SIDE_A, '_prefetched_a_terminations')
 
-    b_terminations: list[Annotated[
+    @strawberry_django.field(
+        prefetch_related=_cable_terminations_prefetch(CableEndChoices.SIDE_B, '_prefetched_b_terminations'),
+    )
+    def b_terminations(self) -> list[Annotated[
         Annotated['CircuitTerminationType', strawberry.lazy('circuits.graphql.types')]
         | Annotated['ConsolePortType', strawberry.lazy('dcim.graphql.types')]
         | Annotated['ConsoleServerPortType', strawberry.lazy('dcim.graphql.types')]
@@ -221,7 +267,8 @@ class CableType(PrimaryObjectType):
         | Annotated['PowerPortType', strawberry.lazy('dcim.graphql.types')]
         | Annotated['RearPortType', strawberry.lazy('dcim.graphql.types')],
         strawberry.union('CableTerminationTerminationType'),
-    ]]
+    ]]:
+        return _resolve_cable_terminations(self, CableEndChoices.SIDE_B, '_prefetched_b_terminations')
 
 
 @strawberry_django.type(

+ 112 - 0
netbox/dcim/tests/test_api.py

@@ -1,7 +1,9 @@
 import json
 
 from django.conf import settings
+from django.db import connection
 from django.test import tag
+from django.test.utils import CaptureQueriesContext
 from django.urls import reverse
 from django.utils.translation import gettext as _
 from rest_framework import status
@@ -9,6 +11,7 @@ from rest_framework import status
 from core.models import ObjectType
 from dcim.choices import *
 from dcim.constants import *
+from dcim.graphql.types import _CABLE_TERMINATION_MODELS
 from dcim.models import *
 from extras.models import ConfigTemplate, Tag
 from ipam.choices import VLANQinQRoleChoices
@@ -3593,6 +3596,115 @@ class CableTestCase(APIViewTestCases.APIViewTestCase):
 
                 self.assertSetEqual(set(ids), expected)
 
+    def test_graphql_cable_terminations_query_count(self):
+        """
+        Resolving CableType.a_terminations and CableType.b_terminations must take a constant number
+        of queries, regardless of how many cables (and hence terminations) are returned.
+
+        Also exercises selecting both cable ends in a single query: each end must be prefetched
+        under its own attribute, as two prefetches of the same relation cannot be merged.
+        """
+        self.add_permissions(
+            'dcim.view_cable',
+            'dcim.view_device',
+            'dcim.view_devicerole',
+            'dcim.view_devicetype',
+            'dcim.view_interface',
+            'dcim.view_platform',
+        )
+
+        # Reuse existing fixtures from setUpTestData()
+        site = Site.objects.get(slug='site-1')
+        devicetype = DeviceType.objects.get(slug='device-type-1')
+        role = DeviceRole.objects.get(slug='device-role-1')
+
+        # Create an isolated topology of cables between two devices
+        devices = (
+            Device(device_type=devicetype, role=role, name='GQL Count Device A', site=site),
+            Device(device_type=devicetype, role=role, name='GQL Count Device B', site=site),
+        )
+        Device.objects.bulk_create(devices)
+
+        interfaces = []
+        for device in devices:
+            for i in range(0, 8):
+                interfaces.append(
+                    Interface(device=device, type=InterfaceTypeChoices.TYPE_1GE_FIXED, name=f'gql{i}')
+                )
+        Interface.objects.bulk_create(interfaces)
+
+        expected_terminations = {}
+        for i in range(0, 8):
+            cable = Cable(
+                a_terminations=[interfaces[i]],
+                b_terminations=[interfaces[i + 8]],
+                label=f'GQL Count Cable {i}',
+            )
+            cable.save()
+            expected_terminations[str(cable.pk)] = (interfaces[i].pk, interfaces[i + 8].pk)
+
+        url = reverse('graphql')
+        termination_fields = """
+            ... on InterfaceType {
+              id
+              name
+              device { id name platform { id } role { id } device_type { id } }
+            }
+        """
+
+        def build_query(limit):
+            return f"""{{
+              cable_list(
+                filters: {{ label: {{ contains: "GQL Count Cable " }} }},
+                pagination: {{ limit: {limit} }}
+              ) {{
+                id
+                a_terminations {{ {termination_fields} }}
+                b_terminations {{ {termination_fields} }}
+              }}
+            }}"""
+
+        # Warm per-process caches (e.g. ContentType) so they are not counted below
+        self.client.post(url, data={'query': build_query(1)}, format='json', **self.header)
+
+        query_counts = {}
+        for limit in (2, 8):
+            with CaptureQueriesContext(connection) as queries:
+                response = self.client.post(
+                    url, data={'query': build_query(limit)}, format='json', **self.header
+                )
+            self.assertHttpStatus(response, status.HTTP_200_OK)
+            data = response.json()
+            self.assertNotIn('errors', data)
+
+            rows = data['data']['cable_list']
+            self.assertEqual(len(rows), limit)
+
+            # Both ends must resolve to the expected interfaces
+            for row in rows:
+                interface_a, interface_b = expected_terminations[row['id']]
+                self.assertEqual([t['id'] for t in row['a_terminations']], [str(interface_a)])
+                self.assertEqual([t['id'] for t in row['b_terminations']], [str(interface_b)])
+
+            query_counts[limit] = len(queries.captured_queries)
+
+        self.assertEqual(
+            query_counts[2],
+            query_counts[8],
+            f"Query count scales with the number of cables returned: {query_counts}"
+        )
+
+    def test_graphql_cable_termination_models(self):
+        """
+        The GraphQL prefetch hint for a cable termination enumerates the terminating models
+        explicitly; a model missing from that list silently falls back to an unoptimized query
+        rather than raising, so guard against drift from CABLE_TERMINATION_MODELS.
+        """
+        self.assertSetEqual(
+            {(model._meta.app_label, model._meta.model_name) for model in _CABLE_TERMINATION_MODELS},
+            {(ot.app_label, ot.model) for ot in ObjectType.objects.filter(CABLE_TERMINATION_MODELS)},
+        )
+
 
 class CableTerminationTestCase(
     APIViewTestCases.GetObjectViewTestCase,