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

Closes #23448: Annotate GraphQL config_context through fragments (#23451)

* Fixes #23448: Annotate GraphQL config_context through fragments

ConfigContextMixin decided whether to annotate config context data by reading the names of the direct selections, which raised an error for inline fragments and skipped the annotation for named fragments. Apply the annotation as an optimizer hint on the config_context field instead, so it follows the field however it is selected.

* Raise a descriptive error for models without config context support

ConfigContextDataExpression now checks that the model's default manager provides get_config_context_data_expression() and names the model in a TypeError otherwise, rather than failing with a bare AttributeError.
Jason Novinger 4 дней назад
Родитель
Сommit
6750369423

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

@@ -2807,6 +2807,56 @@ class DeviceTestCase(Mixins.FieldPrefetchMixin, APIViewTestCases.APIViewTestCase
         self.assertNotIn('errors', data)
         self.assertEqual({row['name']: row['config_context'] for row in data['data']['device_list']}, expected)
 
+    @tag('regression')  # Issue #23448
+    def test_config_context_graphql_selection(self):
+        """
+        GraphQL annotates config context for devices with no cache however config_context is selected, including
+        through fragments and nested relations, without a query per device.
+        """
+        self.add_permissions('dcim.view_device', 'dcim.view_interface')
+        ConfigContext.objects.create(name='Config Context 1', weight=100, data={'foo': 123})
+        for device in Device.objects.all():
+            Interface.objects.create(device=device, name='eth0', type=InterfaceTypeChoices.TYPE_1GE_FIXED)
+        Device.objects.update(_config_context_data=None)
+        expected = {device.name: device.render_config_context() for device in Device.objects.all()}
+
+        fragment = 'fragment F on DeviceType { name config_context }'
+        queries = {
+            'field': '{ device_list { name config_context } }',
+            'inline fragment': '{ device_list { ... on DeviceType { name config_context } } }',
+            'named fragment': f'{fragment} query {{ device_list {{ ...F }} }}',
+            'nested field': '{ interface_list { device { name config_context } } }',
+            'nested inline fragment': '{ interface_list { device { ... on DeviceType { name config_context } } } }',
+            'nested named fragment': f'{fragment} query {{ interface_list {{ device {{ ...F }} }} }}',
+        }
+        for shape, query in queries.items():
+            with self.subTest(shape=shape):
+                with CaptureQueriesContext(connection) as captured:
+                    response = self.client.post(
+                        reverse('graphql'), data={'query': query}, format='json', **self.header
+                    )
+                self.assertHttpStatus(response, status.HTTP_200_OK)
+                data = response.json()
+                self.assertNotIn('errors', data)
+                if shape.startswith('nested'):
+                    rows = [row['device'] for row in data['data']['interface_list']]
+                else:
+                    rows = data['data']['device_list']
+                self.assertEqual({row['name']: row['config_context'] for row in rows}, expected)
+                # Rendered by the annotation in a single query, not one fallback query per device
+                sql = [captured_query['sql'] for captured_query in captured.captured_queries]
+                self.assertEqual(len([s for s in sql if '"extras_configcontext"' in s]), 1)
+
+        # Not annotated when config_context is not selected
+        with CaptureQueriesContext(connection) as captured:
+            response = self.client.post(
+                reverse('graphql'), data={'query': '{ device_list { name } }'}, format='json', **self.header
+            )
+        self.assertHttpStatus(response, status.HTTP_200_OK)
+        self.assertNotIn('errors', response.json())
+        for captured_query in captured.captured_queries:
+            self.assertNotIn('"config_context_data"', captured_query['sql'])
+
     def test_config_context_included_by_default_in_list_view(self):
         """
         Check that config context data is included by default in the devices list.

+ 6 - 19
netbox/extras/graphql/mixins.py

@@ -2,9 +2,9 @@ from typing import TYPE_CHECKING, Annotated
 
 import strawberry
 import strawberry_django
-from strawberry.types import Info
 
 from extras.models import CustomField, ImageAttachment, JournalEntry
+from extras.querysets import ConfigContextDataExpression
 from utilities.querysets import RestrictedPrefetch
 
 __all__ = (
@@ -25,24 +25,11 @@ if TYPE_CHECKING:
 @strawberry.type
 class ConfigContextMixin:
 
-    @classmethod
-    def get_queryset(cls, queryset, info: Info, **kwargs):
-        queryset = super().get_queryset(queryset, info, **kwargs)
-
-        # When `config_context` is requested, annotate the aggregated context data — but only for
-        # rows whose pre-rendered cache (`_config_context_data`) is invalidated (NULL). Warm rows
-        # are served from the cache by get_config_context() and skip the subquery entirely
-        # (PostgreSQL short-circuits the CASE), so resolving config_context across a list of objects
-        # with cold caches no longer incurs a per-object fallback query.
-        selected = {f.name for f in info.selected_fields[0].selections}
-        if 'config_context' in selected and hasattr(queryset, 'annotate_config_context_data'):
-            return queryset.annotate_config_context_data(only_invalidated=True)
-
-        return queryset
-
-    # Ensure both the pre-rendered cache and `local_context_data` are fetched when `config_context`
-    # is requested, so the warm-cache read path requires no additional queries.
-    @strawberry_django.field(only=['_config_context_data', 'local_context_data'])
+    # Annotate context data for rows with an invalidated cache, so a list never falls back to a query per object
+    @strawberry_django.field(
+        only=['_config_context_data', 'local_context_data'],
+        annotate={'config_context_data': ConfigContextDataExpression(only_invalidated=True)},
+    )
     def config_context(self) -> strawberry.scalars.JSON:
         return self.get_config_context(copy_data=False)
 

+ 29 - 3
netbox/extras/querysets.py

@@ -1,11 +1,12 @@
 from django.contrib.postgres.aggregates import JSONBAgg
-from django.db.models import Case, JSONField, OuterRef, Q, Subquery, When
+from django.db.models import Case, Expression, JSONField, OuterRef, Q, Subquery, When
 
 from extras.models.tags import TaggedItem
 from utilities.query_functions import EmptyGroupByJSONBAgg
 from utilities.querysets import RestrictedQuerySet
 
 __all__ = (
+    'ConfigContextDataExpression',
     'ConfigContextModelQuerySet',
     'ConfigContextQuerySet',
     'NotificationQuerySet',
@@ -91,7 +92,13 @@ class ConfigContextModelQuerySet(RestrictedQuerySet):
     """
     def annotate_config_context_data(self, only_invalidated=False):
         """
-        Attach the subquery annotation to the base queryset.
+        Attach the subquery annotation to the base queryset. See get_config_context_data_expression().
+        """
+        return self.annotate(config_context_data=self.get_config_context_data_expression(only_invalidated))
+
+    def get_config_context_data_expression(self, only_invalidated=False):
+        """
+        Return the unresolved expression which aggregates the config context data applicable to each row.
 
         Args:
             only_invalidated: If True, evaluate the (expensive) aggregation subquery only for rows
@@ -121,7 +128,7 @@ class ConfigContextModelQuerySet(RestrictedQuerySet):
                 default=None,
                 output_field=JSONField(),
             )
-        return self.annotate(config_context_data=subquery)
+        return subquery
 
     def _get_config_context_filters(self):
         # Construct the set of Q objects for the specific object types
@@ -193,6 +200,25 @@ class ConfigContextModelQuerySet(RestrictedQuerySet):
         return base_query
 
 
+class ConfigContextDataExpression(Expression):
+    """
+    A model-agnostic form of ConfigContextModelQuerySet.get_config_context_data_expression(), built for the model of
+    the query it is resolved against so its outer references bind to that query's joins.
+    """
+    def __init__(self, only_invalidated=False):
+        super().__init__(output_field=JSONField())
+        self.only_invalidated = only_invalidated
+
+    def resolve_expression(self, query=None, *args, **kwargs):
+        manager = query.model._default_manager
+        if not hasattr(manager, 'get_config_context_data_expression'):
+            raise TypeError(
+                f'{query.model.__name__} must use ConfigContextModelQuerySet to annotate config context data'
+            )
+        expression = manager.get_config_context_data_expression(self.only_invalidated)
+        return expression.resolve_expression(query, *args, **kwargs)
+
+
 class NotificationQuerySet(RestrictedQuerySet):
 
     def unread(self):

+ 33 - 0
netbox/virtualization/tests/test_api.py

@@ -486,6 +486,39 @@ class VirtualMachineTestCase(APIViewTestCases.APIViewTestCase):
             {row['name']: row['config_context'] for row in data['data']['virtual_machine_list']}, expected
         )
 
+    @tag('regression')  # Issue #23448
+    def test_config_context_graphql_selection(self):
+        """
+        GraphQL annotates config context for virtual machines with no cache however config_context is selected,
+        including through fragments, without a query per virtual machine.
+        """
+        self.add_permissions('virtualization.view_virtualmachine')
+        ConfigContext.objects.create(name='Config Context 1', weight=100, data={'foo': 123})
+        VirtualMachine.objects.update(_config_context_data=None)
+        expected = {vm.name: vm.render_config_context() for vm in VirtualMachine.objects.all()}
+
+        queries = {
+            'field': '{ virtual_machine_list { name config_context } }',
+            'inline fragment': '{ virtual_machine_list { ... on VirtualMachineType { name config_context } } }',
+            'named fragment': (
+                'fragment F on VirtualMachineType { name config_context } query { virtual_machine_list { ...F } }'
+            ),
+        }
+        for shape, query in queries.items():
+            with self.subTest(shape=shape):
+                with CaptureQueriesContext(connection) as captured:
+                    response = self.client.post(
+                        reverse('graphql'), data={'query': query}, format='json', **self.header
+                    )
+                self.assertHttpStatus(response, status.HTTP_200_OK)
+                data = response.json()
+                self.assertNotIn('errors', data)
+                rows = data['data']['virtual_machine_list']
+                self.assertEqual({row['name']: row['config_context'] for row in rows}, expected)
+                # Rendered by the annotation in a single query, not one fallback query per virtual machine
+                sql = [captured_query['sql'] for captured_query in captured.captured_queries]
+                self.assertEqual(len([s for s in sql if '"extras_configcontext"' in s]), 1)
+
     def test_config_context_included_by_default_in_list_view(self):
         """
         Check that config context data is included by default in the virtual machines list.