From 9f452028659b4dad6a56d8923f5d32e26019c265 Mon Sep 17 00:00:00 2001 From: Adam Byczkowski <38091261+qduk@users.noreply.github.com> Date: Mon, 20 Nov 2023 07:00:07 -0600 Subject: [PATCH] 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 * Updated per Leos review * Apply suggestions from code review Co-authored-by: Leo Kirchner --------- Co-authored-by: Leo Kirchner --- docs/user/modeling.md | 30 +++++++++++++++++++++++--- nautobot_ssot/contrib.py | 23 ++++++++++++++------ nautobot_ssot/tests/test_contrib.py | 33 +++++++++++++++++++++++++++++ 3 files changed, 76 insertions(+), 10 deletions(-) diff --git a/docs/user/modeling.md b/docs/user/modeling.md index fc2007ae1..e659c48c5 100644 --- 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") +``` \ No newline at end of file diff --git a/nautobot_ssot/contrib.py b/nautobot_ssot/contrib.py index b51a28670..d62df0d7d 100644 --- a/nautobot_ssot/contrib.py +++ b/nautobot_ssot/contrib.py @@ -59,13 +59,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): @@ -258,6 +252,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 index 6e52eeef9..b1c620f12 100644 --- a/nautobot_ssot/tests/test_contrib.py +++ b/nautobot_ssot/tests/test_contrib.py @@ -309,6 +309,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'."""