diff --git a/netbox/dcim/graphql/types.py b/netbox/dcim/graphql/types.py index e4f4b7a14..dfe29fd2d 100644 --- a/netbox/dcim/graphql/types.py +++ b/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( diff --git a/netbox/dcim/tests/test_api.py b/netbox/dcim/tests/test_api.py index 837db1237..d321f77d6 100644 --- a/netbox/dcim/tests/test_api.py +++ b/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,