Przeglądaj źródła

Closes #22822: Avoid extra DB query when fetching rack reservation units via GraphQL API

Jeremy Stretch 1 dzień temu
rodzic
commit
d3735f2db5

+ 5 - 1
netbox/dcim/graphql/types.py

@@ -828,11 +828,15 @@ class RackReservationType(PrimaryObjectType):
     @classmethod
     @classmethod
     def get_queryset(cls, queryset, info, **kwargs):
     def get_queryset(cls, queryset, info, **kwargs):
         queryset = super().get_queryset(queryset, info, **kwargs)
         queryset = super().get_queryset(queryset, info, **kwargs)
+        # Annotate unit_count here so RackReservationFilter.unit_count can reference it. The field below
+        # is resolved from `units` rather than from this annotation, which is not applied on every path
+        # by which a RackReservation may be resolved.
         return queryset.annotate(
         return queryset.annotate(
             unit_count=Func('units', function='CARDINALITY', output_field=IntegerField())
             unit_count=Func('units', function='CARDINALITY', output_field=IntegerField())
         )
         )
 
 
-    @strawberry.field
+    # Ensure `units` is fetched when `unit_count` is requested
+    @strawberry_django.field(only=['units'])
     def unit_count(self) -> int:
     def unit_count(self) -> int:
         return len(self.units)
         return len(self.units)
 
 

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

@@ -29,7 +29,7 @@ from extras.choices import CustomFieldTypeChoices
 from extras.models import CustomField, TableConfig, Tag
 from extras.models import CustomField, TableConfig, Tag
 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
+from users.models import Token, User
 from utilities.tables import get_table_for_model
 from utilities.tables import get_table_for_model
 from utilities.testing import APITestCase, APIViewTestCases, TestCase, disable_warnings
 from utilities.testing import APITestCase, APIViewTestCases, TestCase, disable_warnings
 
 
@@ -675,6 +675,26 @@ class GraphQLDeferredColumnTestCase(APITestCase):
             for i in range(cls.OBJECT_COUNT)
             for i in range(cls.OBJECT_COUNT)
         ])
         ])
 
 
+        # Rack reservations, for RackReservationType.unit_count. Reservations within a rack may not claim
+        # overlapping units, so each is allocated a distinct run of them. The length of each run varies so
+        # that unit_count is asserted per object rather than against a single expected value.
+        rack = Rack.objects.create(name='Rack 1', site=site)
+        user = User.objects.create(username='Reservation user')
+        reservations, next_unit = [], 1
+        for i in range(cls.OBJECT_COUNT):
+            unit_count = i % 3 + 1
+            reservations.append(RackReservation(
+                rack=rack,
+                units=list(range(next_unit, next_unit + unit_count)),
+                user=user,
+                description=f'Reservation {i}',
+            ))
+            next_unit += unit_count
+        cls.expected_unit_counts = {
+            reservation.pk: len(reservation.units)
+            for reservation in RackReservation.objects.bulk_create(reservations)
+        }
+
     def _execute(self, query):
     def _execute(self, query):
         url = reverse('graphql')
         url = reverse('graphql')
         with CaptureQueriesContext(connection) as context:
         with CaptureQueriesContext(connection) as context:
@@ -740,6 +760,28 @@ class GraphQLDeferredColumnTestCase(APITestCase):
 
 
         self.assertNoDeferredColumnReloads(query, 'device_list', 'dcim_device', validate)
         self.assertNoDeferredColumnReloads(query, 'device_list', 'dcim_device', validate)
 
 
+    def test_rack_reservation_unit_count(self):
+        """
+        Regression test for #22822: RackReservationType.unit_count must not defer `units`.
+        """
+        self.add_permissions('dcim.view_rackreservation')
+        query = """
+        {
+            rack_reservation_list(pagination: {limit: %(limit)s}) {
+                id
+                unit_count
+            }
+        }
+        """
+
+        def validate(reservations):
+            for reservation in reservations:
+                self.assertEqual(
+                    reservation['unit_count'], self.expected_unit_counts[int(reservation['id'])]
+                )
+
+        self.assertNoDeferredColumnReloads(query, 'rack_reservation_list', 'dcim_rackreservation', validate)
+
 
 
 class GraphQLSchemaCoverageTestCase(APIViewTestCases.GraphQLSchemaCoverageTestCase):
 class GraphQLSchemaCoverageTestCase(APIViewTestCases.GraphQLSchemaCoverageTestCase):
     pass
     pass