Final test fix

* Add `rir` property method to model
* Add `aggregate` property method to model
* Add `rir` attribute to serializer
* Add `aggregate` attribute to serializer
* Add test for both rir and aggregate fields on API

netbox/ipam/api/serializers_/ip.py

@@ -44,6 +44,8 @@ class AggregateSerializer(NetBoxModelSerializer):
 
 
 class PrefixSerializer(NetBoxModelSerializer):
+    aggregate = AggregateSerializer(nested=True, read_only=True, allow_null=True)
+    rir = RIRSerializer(nested=True, read_only=True, allow_null=True)
     family = ChoiceField(choices=IPAddressFamilyChoices, read_only=True)
     vrf = VRFSerializer(nested=True, required=False, allow_null=True)
     scope_type = ContentTypeField(
@@ -67,9 +69,9 @@ class PrefixSerializer(NetBoxModelSerializer):
     class Meta:
         model = Prefix
         fields = [
-            'id', 'url', 'display_url', 'display', 'family', 'prefix', 'vrf', 'scope_type', 'scope_id', 'scope',
-            'tenant', 'vlan', 'status', 'role', 'is_pool', 'mark_utilized', 'description', 'comments', 'tags',
-            'custom_fields', 'created', 'last_updated', 'children', '_depth',
+            'id', 'url', 'display_url', 'display', 'aggregate', 'rir', 'family', 'prefix', 'vrf', 'scope_type',
+            'scope_id', 'scope', 'tenant', 'vlan', 'status', 'role', 'is_pool', 'mark_utilized', 'description',
+            'comments', 'tags', 'custom_fields', 'created', 'last_updated', 'children', '_depth',
         ]
         brief_fields = ('id', 'url', 'display', 'family', 'prefix', 'description', '_depth')
 

netbox/ipam/api/views.py

@@ -2,6 +2,8 @@ from copy import deepcopy
 
 from django.core.exceptions import ObjectDoesNotExist, PermissionDenied
 from django.db import transaction
+from django.db.models import Subquery, OuterRef
+from django.db.models.functions import JSONObject
 from django.shortcuts import get_object_or_404
 from django.utils.translation import gettext as _
 from django_pglocks import advisory_lock
@@ -21,6 +23,7 @@ from netbox.api.viewsets.mixins import ObjectValidationMixin
 from netbox.config import get_config
 from netbox.constants import ADVISORY_LOCK_KEYS
 from utilities.api import get_serializer_for_model
+from utilities.fields import JSONModelField
 from . import serializers
 
 
@@ -90,6 +93,32 @@ class PrefixViewSet(NetBoxModelViewSet):
             return serializers.PrefixLengthSerializer
         return super().get_serializer_class()
 
+    def get_queryset(self):
+        """
+        Return the query set with additional annotations for Aggregate and RIR
+        """
+        qs = super().get_queryset()
+
+        # Determine the fields to return
+        aggregate_fields = JSONObject(**{f.name: f.name for f in Aggregate._meta.get_fields()})
+        rir_fields = JSONObject(**{f.name: f.name for f in RIR._meta.get_fields()})
+
+        # Get the outer reference
+        prefix_field = OuterRef("prefix")
+        aggregate_field = OuterRef("aggregate_id")
+
+        aggregates = Aggregate.objects.filter(prefix__net_contains_or_equals=prefix_field)
+        rirs = RIR.objects.filter(aggregates=aggregate_field)
+
+        # The sub queries for the annotation, returning a json object of the related model
+        agg_sq = Subquery(
+            aggregates.values_list(aggregate_fields)[:1], output_field=JSONModelField(related_model=Aggregate)
+        )
+        agg_id_sq = Subquery(aggregates.values_list('pk', flat=True)[:1])
+        rir_sq = Subquery(rirs.values_list(rir_fields)[:1], output_field=JSONModelField(related_model=RIR))
+
+        return qs.annotate(aggregate=agg_sq, aggregate_id=agg_id_sq).annotate(rir=rir_sq)
+
 
 class IPRangeViewSet(NetBoxModelViewSet):
     queryset = IPRange.objects.all()

netbox/ipam/tests/test_api.py

@@ -532,6 +532,31 @@ class PrefixTest(APIViewTestCases.APIViewTestCase):
         self.assertHttpStatus(response, status.HTTP_201_CREATED)
         self.assertEqual(len(response.data), 8)
 
+    def test_get_prefix_with_aggregate_and_rir(self):
+        self.add_permissions('ipam.view_prefix')
+        rir = RIR.objects.create(name='RFC 1918', slug='rfc-1918')
+        aggregate = Aggregate.objects.create(prefix=IPNetwork('192.168.0.0/16'), rir=rir)
+        prefixes = [
+            Prefix.objects.filter(prefix=IPNetwork('192.168.2.0/24')).first(),
+            Prefix.objects.create(prefix=IPNetwork('10.0.0.0/24'))
+        ]
+
+        self.assertIsNotNone(prefixes[0])
+
+        url = self._get_detail_url(prefixes[0])
+        response = self.client.get(url, **self.header)
+        self.assertHttpStatus(response, status.HTTP_200_OK)
+        self.assertIsNotNone(response.data.get('aggregate'))
+        self.assertIsNotNone(response.data.get('rir'))
+        self.assertEqual(response.data.get('aggregate').get('id', None), aggregate.pk)
+        self.assertEqual(response.data.get('rir').get('id', None), rir.pk)
+
+        url = self._get_detail_url(prefixes[1])
+        response = self.client.get(url, **self.header)
+        self.assertHttpStatus(response, status.HTTP_200_OK)
+        self.assertIsNone(response.data.get('aggregate'))
+        self.assertIsNone(response.data.get('rir'))
+
 
 class IPRangeTest(APIViewTestCases.APIViewTestCase):
     model = IPRange

netbox/utilities/fields.py

@@ -1,7 +1,9 @@
 from collections import defaultdict
 
-from django.contrib.contenttypes.fields import GenericForeignKey
+from django.contrib.contenttypes.fields import GenericForeignKey, GenericRelation
+from django.core import exceptions
 from django.db import models
+from django.db.models import ForeignKey, ManyToOneRel, JSONField
 from django.utils.safestring import mark_safe
 from django.utils.translation import gettext_lazy as _
 
@@ -13,6 +15,7 @@ __all__ = (
     'CounterCacheField',
     'NaturalOrderingField',
     'RestrictedGenericForeignKey',
+    'JSONModelField',
 )
 
 
@@ -186,3 +189,35 @@ class CounterCacheField(models.BigIntegerField):
         kwargs["to_model"] = self.to_model_name
         kwargs["to_field"] = self.to_field_name
         return name, path, args, kwargs
+
+
+class JSONModelField(JSONField):
+    def __init__(self, related_model, *args, **kwargs):
+        """
+        Extract the related model from the kwargs and set after instantiation
+        """
+        super().__init__(*args, **kwargs)
+        if related_model is None or isinstance(related_model, models.Model):
+            raise exceptions.FieldError('related_model must be set or be an instance of Model')
+        self.related_model = related_model
+
+    def from_db_value(self, value, expression, connection):
+        """
+        Return the actual instantiated model from the fields, minus the models that cannot be worked with
+        """
+        data = super().from_db_value(value, expression, connection)
+        # Return nothing if there is nothing
+        if data is None:
+            return None
+
+        # Extract the fields from the meta for processing
+        fields = {f.name: f for f in self.related_model._meta.get_fields()}
+
+        keys = data.copy().keys()
+        for key in keys:
+            if key not in fields or isinstance(fields.get(key), (GenericRelation, ForeignKey, ManyToOneRel, )):
+                # Delete un-parsable fields
+                del data[key]
+
+        # Return the full model minus deleted fields
+        return self.related_model(**data)