From 1708af6f3ea2afc2dcef5465c611ee2cdd157620 Mon Sep 17 00:00:00 2001 From: Eg0ra Date: Tue, 1 Oct 2024 17:01:09 +0700 Subject: [PATCH] Allow prefetch with IntermediateM2MField --- import_export_extensions/fields.py | 172 +++++++++-------------------- tests/fake_app/resources.py | 7 ++ 2 files changed, 61 insertions(+), 118 deletions(-) diff --git a/import_export_extensions/fields.py b/import_export_extensions/fields.py index 91541ec..4d67d1c 100644 --- a/import_export_extensions/fields.py +++ b/import_export_extensions/fields.py @@ -3,19 +3,40 @@ from import_export.fields import Field -class M2MField(Field): - """Base M2M field provides faster related instances import.""" +class IntermediateManyToManyField(Field): + """Resource field for M2M with custom ``through`` model. - def __init__(self, *args, export_order=None, **kwargs): - """Save additional field params. + By default, ``django-import-export`` set up object attributes using + ``setattr(obj, attribute_name, value)``, where ``value`` is ``QuerySet`` + of related model objects. But django forbid this when `ManyToManyField`` + used with custom ``through`` model. - Args: - export_order(str): field name that should be used for ordering - instances during export + This field expects be used with custom widget that return not simple value, + but dict with intermediate model attributes. - """ - super().__init__(*args, **kwargs) - self.export_order = export_order + For easy comments following models will be used: + + Artist: + name + bands -> + + Membership: + artist + band + date_joined + + Band: + title + <- artists + + So this field should be used for exporting Artists with `bands` field. + + Save workflow is following: + 1. clean data (extract dicts) + 2. Remove current M2M instances of object + 3. Create new M2M instances based on current object + + """ def _format_exception(self, exception): """Shortcut for humanizing exception.""" @@ -36,17 +57,16 @@ def get_value(self, obj): if self.attribute is None: return None - m2m_rel, _, field_name, _ = self.get_relation_field_params(obj) + m2m_rel, _, _ = self.get_relation_field_params(obj) - # retrieve intermediate model - # intermediate_model is the Membership model - intermediate_model = m2m_rel.through + if not obj.id: + return m2m_rel.through.objects.none() - # filter relations with passed object - qs = intermediate_model.objects.filter(**{f"{field_name}_id": obj.id}) - if self.export_order: - qs = qs.order_by(self.export_order) - return qs + through_model_accessor_name = self.get_through_model_accessor_name( + obj, + m2m_rel, + ) + return getattr(obj, through_model_accessor_name).all() def get_relation_field_params(self, obj): """Shortcut to get relation field params. @@ -63,111 +83,24 @@ def get_relation_field_params(self, obj): m2m_field = m2m_rel.field field_name = m2m_field.m2m_reverse_field_name() reversed_field_name = m2m_field.m2m_field_name() - return m2m_rel, m2m_field, field_name, reversed_field_name + return m2m_rel, field_name, reversed_field_name # otherwise it is a forward relation m2m_rel = field.remote_field m2m_field = field field_name = m2m_field.m2m_field_name() reversed_field_name = m2m_field.m2m_reverse_field_name() - return m2m_rel, m2m_field, field_name, reversed_field_name - - def save(self, obj, data, *args, **kwargs): - """Delete intermediate models. - - This implementation deletes intermediate models, which were excluded - and creates intermediate models only for newly added models. - - Parent `save` method deletes and recreates intermediate models for all - instances which generates a lot of exceed Feed Entries, so overridden. - - """ - ( - _, - m2m_field, - field_name, - reversed_field_name, - ) = self.get_relation_field_params(obj) + return m2m_rel, field_name, reversed_field_name - # retrieve intermediate model Membership - intermediate_model = m2m_field.remote_field.through + def get_through_model_accessor_name(self, obj, m2m_rel) -> str: + """Shortcut to get through model accessor name.""" + for related_object in obj._meta.related_objects: + if related_object.related_model is m2m_rel.through: + return related_object.accessor_name - # should be returned following list: - # [{'object': Instance01, 'properties': {}}, - # {'object': Instance02, 'properties': {}}] - data = self.clean(data) - - # IDs of related instances in imported data - imported_ids = {i["object"].id for i in data} - - # IDs of current instances - manager = getattr(obj, self.attribute) - current_ids = set(manager.values_list("id", flat=True)) - - # Find instances to be excluded after import - excluded_ids = current_ids - imported_ids - intermediate_model.objects.filter( - **{ - field_name: obj, - f"{reversed_field_name}__id__in": excluded_ids, - }, - ).delete() - - # Find instances to add after import - added_ids = imported_ids - current_ids - for instance in data: - # process only newly added objects - if instance["object"].id not in added_ids: - continue - obj_data = instance["properties"].copy() - obj_data.update( - { - field_name: obj, - reversed_field_name: instance["object"], - }, - ) - intermediate_obj = intermediate_model(**obj_data) - try: - intermediate_obj.full_clean() - except Exception as exception: - self._format_exception(exception) - intermediate_obj.save() - - -class IntermediateManyToManyField(M2MField): - """Resource field for M2M with custom ``through`` model. - - By default, ``django-import-export`` set up object attributes using - ``setattr(obj, attribute_name, value)``, where ``value`` is ``QuerySet`` - of related model objects. But django forbid this when `ManyToManyField`` - used with custom ``through`` model. - - This field expects be used with custom widget that return not simple value, - but dict with intermediate model attributes. - - For easy comments following models will be used: - - Artist: - name - bands -> - - Membership: - artist - band - date_joined - - Band: - title - <- artists - - So this field should be used for exporting Artists with `bands` field. - - Save workflow is following: - 1. clean data (extract dicts) - 2. Remove current M2M instances of object - 3. Create new M2M instances based on current object - - """ + raise ValueError( + f"{obj._meta.model} has no relation with {m2m_rel.through}", + ) def save(self, obj, data, *args, **kwargs): """Add M2M relations for obj from data. @@ -185,7 +118,6 @@ def save(self, obj, data, *args, **kwargs): ( m2m_rel, - m2m_field, field_name, reversed_field_name, ) = self.get_relation_field_params(obj) @@ -200,7 +132,11 @@ def save(self, obj, data, *args, **kwargs): # remove current related objects, # i.e. clear artists's band - intermediate_model.objects.filter(**{field_name: obj}).delete() + through_model_accessor_name = self.get_through_model_accessor_name( + obj, + m2m_rel, + ) + getattr(obj, through_model_accessor_name).all().delete() for rel_obj_data in instances_data: # add current and remote object to intermediate instance data diff --git a/tests/fake_app/resources.py b/tests/fake_app/resources.py index e9ceb84..7d861ac 100644 --- a/tests/fake_app/resources.py +++ b/tests/fake_app/resources.py @@ -39,3 +39,10 @@ class Meta: model = Artist clean_model_instances = True fields = ["id", "name", "bands", "instrument"] + + def get_queryset(self): + """Return a queryset.""" + return Artist.objects.all().prefetch_related( + "membership_set__band", + "bands", + )