Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

implements caching mechanism into the NautobotAdapter #296

Merged
merged 2 commits into from
Feb 14, 2024
Merged

implements caching mechanism into the NautobotAdapter #296

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
a4d6e8e
implements caching mechanism into the NautobotAdapter
Kircheneer Feb 13, 2024
9e92c77
Update nautobot_ssot/tests/test_contrib.py
Kircheneer Feb 13, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
107 changes: 67 additions & 40 deletions nautobot_ssot/contrib.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
from collections import defaultdict
from dataclasses import dataclass
from enum import Enum
from typing import FrozenSet, Tuple, Hashable, DefaultDict, Dict, Type

import pydantic
from diffsync import DiffSyncModel, DiffSync
Expand All @@ -15,6 +16,19 @@
from typing_extensions import get_type_hints


# This type describes a set of parameters to use as a dictionary key for the cache. As such, its needs to be hashable
# and therefore a frozenset rather than a normal set or a list.
#
# The following is an example of a parameter set that describes a tenant based on its name and group:
# frozenset(
# [
# ("name", "ABC Inc."),
# ("group__name", "Customers"),
# ]
# )
ParameterSet = FrozenSet[Tuple[str, Hashable]]


class RelationshipSideEnum(Enum):
"""This details which side of a custom relationship the model it's defined on is on."""

Expand Down Expand Up @@ -91,15 +105,35 @@ class NautobotAdapter(DiffSync):
This adapter is able to infer how to load data from Nautobot based on how the models attached to it are defined.
"""

# This dictionary acts as an ORM cache.
_cache: DefaultDict[str, Dict[ParameterSet, Model]]
_cache_hits: DefaultDict[str, int] = defaultdict(int)

def __init__(self, *args, job, sync=None, **kwargs):
"""Instantiate this class, but do not load data immediately from the local system."""
super().__init__(*args, **kwargs)
self.job = job
self.sync = sync

# Caches lookups to custom relationships.
# TODO: Once caching is in, replace this cache with it.
self.custom_relationship_cache = {}
self.invalidate_cache()

def invalidate_cache(self, zero_out_hits=True):
"""Invalidates all the objects in the ORM cache."""
self._cache = defaultdict(dict)
if zero_out_hits:
self._cache_hits = defaultdict(int)

def get_from_orm_cache(self, parameters: Dict, model_class: Type[Model]):
"""Retrieve an object from the ORM or the cache."""
parameter_set = frozenset(parameters.items())
content_type = ContentType.objects.get_for_model(model_class)
model_cache_key = f"{content_type.app_label}.{content_type.model}"
if cached_object := self._cache[model_cache_key].get(parameter_set):
self._cache_hits[model_cache_key] += 1
return cached_object
# As we are using `get` here, this will error if there is not exactly one object that corresponds to the
# parameter set. We intentionally pass these errors through.
self._cache[model_cache_key][parameter_set] = model_class.objects.get(**dict(parameter_set))
return self._cache[model_cache_key][parameter_set]

@staticmethod
def _get_parameter_names(diffsync_model):
Expand Down Expand Up @@ -223,7 +257,7 @@ def _handle_custom_relationship_to_many_relationship(
inner_type = diffsync_field_type.__dict__["__args__"][0].__dict__["__args__"][0]
related_objects_list = []
# TODO: Allow for filtering, i.e. not taking into account all the objects behind the relationship.
relationship = Relationship.objects.get(label=annotation.name)
relationship = self.get_from_orm_cache({"label": annotation.name}, Relationship)
relationship_association_parameters = self._construct_relationship_association_parameters(
annotation, database_object
)
Expand Down Expand Up @@ -257,9 +291,7 @@ def _handle_custom_relationship_to_many_relationship(
return related_objects_list

def _construct_relationship_association_parameters(self, annotation, database_object):
relationship = self.custom_relationship_cache.get(
annotation.name, Relationship.objects.get(label=annotation.name)
)
relationship = self.get_from_orm_cache({"label": annotation.name}, Relationship)
relationship_association_parameters = {
"relationship": relationship,
"source_type": relationship.source_type,
Expand Down Expand Up @@ -436,7 +468,7 @@ def get_from_db(self):
TODO: Currently I don't think this works for custom fields, therefore those can't be identifiers.
"""
return self._model.objects.get(**self.get_identifiers())
return self.diffsync.get_from_orm_cache(self.get_identifiers(), self._model)

def update(self, attrs):
"""Update the ORM object corresponding to this diffsync object."""
Expand Down Expand Up @@ -467,7 +499,7 @@ def create(cls, diffsync, ids, attrs):
@classmethod
def _handle_single_field(
cls, field, obj, value, relationship_fields, diffsync
): # pylint: disable=too-many-arguments
): # pylint: disable=too-many-arguments,too-many-locals
"""Set a single field on a Django object to a given value, or, for relationship fields, prepare setting.
:param field: The name of the field to set.
Expand Down Expand Up @@ -520,18 +552,14 @@ def _handle_single_field(

# Prepare handling of custom relationship many-to-many fields.
if custom_relationship_annotation:
relationship = diffsync.custom_relationship_cache.get(
custom_relationship_annotation.name,
Relationship.objects.get(label=custom_relationship_annotation.name),
)
relationship = diffsync.get_from_orm_cache({"label": custom_relationship_annotation.name}, Relationship)
if custom_relationship_annotation.side == RelationshipSideEnum.DESTINATION:
related_object_content_type = relationship.source_type
else:
related_object_content_type = relationship.destination_type
related_model_class = related_object_content_type.model_class()
relationship_fields["custom_relationship_many_to_many_fields"][field] = {
"objects": [
related_object_content_type.model_class().objects.get(**parameters) for parameters in value
],
"objects": [diffsync.get_from_orm_cache(parameters, related_model_class) for parameters in value],
"annotation": custom_relationship_annotation,
}
return
Expand All @@ -542,7 +570,7 @@ def _handle_single_field(
# we get all the related objects here to later set them once the object has been saved.
if django_field.many_to_many or django_field.one_to_many:
relationship_fields["many_to_many_fields"][field] = [
django_field.related_model.objects.get(**parameters) for parameters in value
diffsync.get_from_orm_cache(parameters, django_field.related_model) for parameters in value
]
return

Expand All @@ -566,7 +594,7 @@ def _update_obj_with_parameters(cls, obj, parameters, diffsync):
cls._handle_single_field(field, obj, value, relationship_fields, diffsync)

# Set foreign keys
cls._lookup_and_set_foreign_keys(relationship_fields["foreign_keys"], obj)
cls._lookup_and_set_foreign_keys(relationship_fields["foreign_keys"], obj, diffsync=diffsync)

# Save the object to the database
try:
Expand All @@ -592,9 +620,7 @@ def _set_custom_relationship_to_many_fields(cls, custom_relationship_many_to_man
annotation = dictionary.pop("annotation")
objects = dictionary.pop("objects")
# TODO: Deduplicate this code
relationship = diffsync.custom_relationship_cache.get(
annotation.name, Relationship.objects.get(label=annotation.name)
)
relationship = diffsync.get_from_orm_cache({"label": annotation.name}, Relationship)
parameters = {
"relationship": relationship,
"source_type": relationship.source_type,
Expand All @@ -604,25 +630,29 @@ def _set_custom_relationship_to_many_fields(cls, custom_relationship_many_to_man
if annotation.side == RelationshipSideEnum.SOURCE:
parameters["source_id"] = obj.id
for object_to_relate in objects:
association_parameters = parameters.copy()
association_parameters["destination_id"] = object_to_relate.id
try:
association = RelationshipAssociation.objects.get(
**parameters, destination_id=object_to_relate.id
)
association = diffsync.get_from_orm_cache(association_parameters, RelationshipAssociation)
except RelationshipAssociation.DoesNotExist:
association = RelationshipAssociation(**parameters, destination_id=object_to_relate.id)
association.validated_save()
associations.append(association)
else:
parameters["destination_id"] = obj.id
for object_to_relate in objects:
association_parameters = parameters.copy()
association_parameters["source_id"] = object_to_relate.id
try:
association = RelationshipAssociation.objects.get(**parameters, source_id=object_to_relate.id)
association = diffsync.get_from_orm_cache(association_parameters, RelationshipAssociation)
except RelationshipAssociation.DoesNotExist:
association = RelationshipAssociation(**parameters, source_id=object_to_relate.id)
association.validated_save()
associations.append(association)
# Now we need to clean up any associations that we're not `get_or_create`'d in order to achieve
# declarativeness.
# TODO: This may benefit from an ORM cache with `filter` capabilities, but I guess the gain in most cases
# would be fairly minor.
for existing_association in RelationshipAssociation.objects.filter(**parameters):
if existing_association not in associations:
existing_association.delete()
Expand All @@ -649,33 +679,30 @@ def _lookup_and_set_custom_relationship_foreign_keys(cls, custom_relationship_fo
for _, related_model_dict in custom_relationship_foreign_keys.items():
annotation = related_model_dict.pop("_annotation")
# TODO: Deduplicate this code
relationship = diffsync.custom_relationship_cache.get(
annotation.name, Relationship.objects.get(label=annotation.name)
)
relationship = diffsync.get_from_orm_cache({"label": annotation.name}, Relationship)
parameters = {
"relationship": relationship,
"source_type": relationship.source_type,
"destination_type": relationship.destination_type,
}
if annotation.side == RelationshipSideEnum.SOURCE:
parameters["source_id"] = obj.id
destination_object = diffsync.get_from_orm_cache(
related_model_dict, relationship.destination_type.model_class()
)
RelationshipAssociation.objects.update_or_create(
**parameters,
defaults={
"destination_id": relationship.destination_type.model_class()
.objects.get(**related_model_dict)
.id
},
defaults={"destination_id": destination_object.id},
)
else:
parameters["destination_id"] = obj.id
RelationshipAssociation.objects.update_or_create(
**parameters,
defaults={"source_id": relationship.source_type.model_class().objects.get(**related_model_dict).id},
source_object = diffsync.get_from_orm_cache(
related_model_dict, relationship.destination_type.model_class()
)
RelationshipAssociation.objects.update_or_create(**parameters, defaults={"source_id": source_object.id})

@classmethod
def _lookup_and_set_foreign_keys(cls, foreign_keys, obj):
def _lookup_and_set_foreign_keys(cls, foreign_keys, obj, diffsync):
"""
Given a list of foreign keys as dictionaries, look up and set foreign keys on an object.
Expand All @@ -699,13 +726,13 @@ def _lookup_and_set_foreign_keys(cls, foreign_keys, obj):
f"Missing annotation for '{field_name}__app_label' or '{field_name}__model - this is required"
f"for generic foreign keys."
) from error
related_model = ContentType.objects.get(app_label=app_label, model=model).model_class()
related_model = diffsync.get_from_orm_cache({"app_label": app_label, "model": model}, ContentType)
# Set the foreign key to 'None' when none of the fields are set to anything
if not any(related_model_dict.values()):
setattr(obj, field_name, None)
continue
try:
related_object = related_model.objects.get(**related_model_dict)
related_object = diffsync.get_from_orm_cache(related_model_dict, related_model)
except related_model.DoesNotExist as error:
raise ValueError(f"Couldn't find {field_name} instance with: {related_model_dict}.") from error
except MultipleObjectsReturned as error:
Expand Down
Loading
Loading