Add the ability to customize queryset used on NautobotModels (#229)

* Initial commit * Updated * Removed extra line in docs * Added prefetch related things * Removed code part * Pylinted * Update docs/user/modeling.md Co-authored-by: Leo Kirchner <[email protected]> * Updated per Leos review * Apply suggestions from code review Co-authored-by: Leo Kirchner <[email protected]> --------- Co-authored-by: Leo Kirchner <[email protected]>
nautobot · Nov 20, 2023 · c727f96 · c727f96
1 parent 3580394
commit c727f96
Show file tree

Hide file tree

Showing 3 changed files with 76 additions and 10 deletions.
diff --git a/docs/user/modeling.md b/docs/user/modeling.md
@@ -7,12 +7,12 @@ This page describes how to model various kinds of fields on a `nautobot_ssot.con
 The following table describes in brief the different types of model fields and how they are handled.
 
 | Type of field                                      | Field name                | Notes                                                                                                                                            | Applies to                                                                                                                                                                                                              |
-|----------------------------------------------------|---------------------------|--------------------------------------------------------------------------------------------------------------------------------------------------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
-| [Normal fields](#normal-fields)                    | Has to match ORM exactly  | Make sure that the name matches the name in the ORM model.                                                                                       | Fields that are neither custom fields nor relations                                                                                                                                                              |
+| -------------------------------------------------- | ------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------ | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| [Normal fields](#normal-fields)                    | Has to match ORM exactly  | Make sure that the name matches the name in the ORM model.                                                                                       | Fields that are neither custom fields nor relations                                                                                                                                                                     |
 | [Custom fields](#custom-fields)                    | Field name doesn't matter | Use `nautobot_ssot.contrib.CustomFieldAnnotation`                                                                                                | [Nautobot custom fields](https://docs.nautobot.com/projects/core/en/stable/user-guides/custom-fields/?h=custom+fields)                                                                                                  |
 | [*-to-one relationships](#-to-one-relationships)   | Django lookup syntax      | See [here](https://docs.djangoproject.com/en/3.2/topics/db/queries/#lookups-that-span-relationships) - your model fields need to use this syntax | `django.db.models.OneToOneField`, `django.db.models.ForeignKey`, `django.contrib.contenttypes.fields.GenericForeignKey`                                                                                                 |
 | [*-to-many relationships](#-to-many-relationships) | Has to match ORM exactly  | In case of a generic foreign key see [here](#special-case-generic-foreign-key)                                                                   | `django.db.models.ManyToManyField`, `django.contrib.contenttypes.fields.GenericRelation`, `django.db.models.ForeignKey` [backwards](https://docs.djangoproject.com/en/3.2/topics/db/queries/#backwards-related-objects) |
-| Custom Relationships | n/a | Not yet supported | https://docs.nautobot.com/projects/core/en/stable/models/extras/relationship/ |
+| Custom Relationships                               | n/a                       | Not yet supported                                                                                                                                | https://docs.nautobot.com/projects/core/en/stable/models/extras/relationship/                                                                                                                                           |
 
 
 ## Normal Fields
@@ -156,3 +156,27 @@ Through us defining the model, Nautobot will now be able to dynamically load IP
 
 !!! note
     Although `Interface.ip_addresses` is a generic relation, there is only one content type (i.e. `ipam.ipaddress`) that may be related through this relation, therefore we don't have to specific this in any way.
+
+
+## Filtering Objects Loaded From Nautobot
+
+
+If you'd like to filter the objects loaded from the Nautobot, you can do so creating a `get_queryset` function in your model class and return your own queryset. Here is an example where the adapter would only load Tenant objects whose name starts with an "s".
+
+```python
+from nautobot.tenancy.models import Tenant
+from nautobot_ssot.contrib import NautobotModel
+
+class TenantModel(NautobotModel):
+    _model = Tenant
+    _modelname = "tenant"
+    _identifiers = ("name",)
+    _attributes = ("description",)
+
+    name: str
+    description: str
+
+    @classmethod
+    def get_queryset(cls):
+        return Tenant.objects.filter(name__startswith="s")
+```
diff --git a/nautobot_ssot/contrib.py b/nautobot_ssot/contrib.py
@@ -65,13 +65,7 @@ def _get_parameter_names(diffsync_model):
     def _load_objects(self, diffsync_model):
         """Given a diffsync model class, load a list of models from the database and return them."""
         parameter_names = self._get_parameter_names(diffsync_model)
-
-        # Here we identify any foreign keys (i.e. fields with '__' in them) so that we can load them directly in the
-        # first query.
-        prefetch_related_parameters = [parameter.split("__")[0] for parameter in parameter_names if "__" in parameter]
-
-        # TODO: Allow for filtering, i.e. not getting all models from a table but just some.
-        for database_object in diffsync_model._model.objects.prefetch_related(*prefetch_related_parameters).all():
+        for database_object in diffsync_model._get_queryset():
             self._load_single_object(database_object, diffsync_model, parameter_names)
 
     def _load_single_object(self, database_object, diffsync_model, parameter_names):
@@ -263,6 +257,21 @@ class NautobotModel(DiffSyncModel):
 
     _model: Model
 
+    @classmethod
+    def _get_queryset(cls):
+        """Get the queryset used to load the models data from Nautobot."""
+        parameter_names = list(cls._identifiers) + list(cls._attributes)
+        # Here we identify any foreign keys (i.e. fields with '__' in them) so that we can load them directly in the
+        # first query if this function hasn't been overridden.
+        prefetch_related_parameters = [parameter.split("__")[0] for parameter in parameter_names if "__" in parameter]
+        qs = cls.get_queryset()
+        return qs.prefetch_related(*prefetch_related_parameters)
+
+    @classmethod
+    def get_queryset(cls):
+        """Get the queryset used to load the models data from Nautobot."""
+        return cls._model.objects.all()
+
     @classmethod
     def _check_field(cls, name):
         """Check whether the given field name is defined on the diffsync (pydantic) model."""

diff --git a/nautobot_ssot/tests/test_contrib.py b/nautobot_ssot/tests/test_contrib.py
@@ -311,6 +311,39 @@ class Adapter(NautobotAdapter):
             "Custom fields aren't properly loaded through 'BaseAdapter'.",
         )
 
+    def test_overwrite_get_queryset(self):
+        """Test overriding 'get_queryset' method."""
+
+        class TenantModel(NautobotModel):
+            """Test model for testing overridden 'get_queryset' method."""
+
+            _model = Tenant
+            _modelname = "tenant"
+            _identifiers = ("name",)
+            _attributes = ("description",)
+
+            name: str
+            description: str
+
+            @classmethod
+            def get_queryset(cls):
+                return Tenant.objects.filter(name__startswith="N")
+
+        class Adapter(NautobotAdapter):
+            """Test overriding 'get_queryset' method."""
+
+            top_level = ("tenant",)
+            tenant = TenantModel
+
+        new_tenant_name = "NASA"
+        Tenant.objects.create(name=new_tenant_name)
+        Tenant.objects.create(name="Air Force")
+        adapter = Adapter()
+        adapter.load()
+        diffsync_tenant = adapter.get(TenantModel, new_tenant_name)
+
+        self.assertEqual(new_tenant_name, diffsync_tenant.name)
+
 
 class BaseModelTests(TestCase):
     """Testing basic operations through 'NautobotModel'."""