· 5 years ago · Jul 27, 2020, 04:40 AM
1from django import forms
2from drf_writable_nested.mixins import BaseNestedModelSerializer
3from collections import OrderedDict
4from django.db.models.fields.related import ForeignObjectRel
5from django.contrib.contenttypes.fields import GenericRelation
6from django.db.models import ProtectedError
7from django.utils.translation import ugettext_lazy as _
8from rest_framework.exceptions import ValidationError
9
10
11class NestedUpdateMixin(BaseNestedModelSerializer):
12 """
13 Adds update nested feature
14 """
15 default_error_messages = {
16 'cannot_delete_protected': _(
17 "Cannot delete {instances} because "
18 "protected relation exists")
19 }
20
21 def update(self, instance, validated_data):
22 relations, reverse_relations = self._extract_relations(validated_data)
23
24 # Create or update direct relations (foreign key, one-to-one)
25 self.update_or_create_direct_relations(
26 validated_data,
27 relations,
28 )
29
30 # Update instance
31 instance = super(NestedUpdateMixin, self).update(
32 instance,
33 validated_data,
34 )
35 updated_related_data_dict = self.update_or_create_reverse_relations(instance, reverse_relations)
36 self.delete_reverse_relations_if_need(instance, reverse_relations, updated_related_data_dict)
37 return instance
38
39 def update_or_create_reverse_relations(self, instance, reverse_relations):
40 # Update or create reverse relations:
41 # many-to-one, many-to-many, reversed one-to-one
42 related_data_dict = {}
43 for field_name, (related_field, field, field_source) in \
44 reverse_relations.items():
45
46 # Skip processing for empty data or not-specified field.
47 # The field can be defined in validated_data but isn't defined
48 # in initial_data (for example, if multipart form data used)
49 related_data = self.get_initial().get(field_name, None)
50 if related_data is None:
51 continue
52
53 if related_field.one_to_one:
54 # If an object already exists, fill in the pk so
55 # we don't try to duplicate it
56 pk_name = field.Meta.model._meta.pk.attname
57 if pk_name not in related_data and 'pk' in related_data:
58 pk_name = 'pk'
59 if pk_name not in related_data:
60 related_instance = getattr(instance, field_source, None)
61 if related_instance:
62 related_data[pk_name] = related_instance.pk
63
64 # Expand to array of one item for one-to-one for uniformity
65 related_data = [related_data]
66
67 instances = self._prefetch_related_instances(field, related_data)
68
69 save_kwargs = self._get_save_kwargs(field_name)
70 if isinstance(related_field, GenericRelation):
71 save_kwargs.update(
72 self._get_generic_lookup(instance, related_field),
73 )
74 elif not related_field.many_to_many:
75 save_kwargs[related_field.name] = instance
76
77 new_related_instances = []
78 errors = []
79 for data in related_data:
80 obj = instances.get(
81 self._get_related_pk(data, field.Meta.model)
82 )
83 serializer = self._get_serializer_for_field(
84 field,
85 instance=obj,
86 data=data,
87 )
88 try:
89 serializer.is_valid(raise_exception=True)
90 related_instance = serializer.save(**save_kwargs)
91 data['pk'] = related_instance.pk
92 new_related_instances.append(related_instance)
93 errors.append({})
94 except ValidationError as exc:
95 errors.append(exc.detail)
96
97 if any(errors):
98 if related_field.one_to_one:
99 raise ValidationError({field_name: errors[0]})
100 else:
101 raise ValidationError({field_name: errors})
102
103 if related_field.many_to_many:
104 # Add m2m instances to through model via add
105 m2m_manager = getattr(instance, field_source)
106 m2m_manager.add(*new_related_instances)
107
108 related_data_dict[field_name] = related_data
109
110 return related_data_dict
111
112 def delete_reverse_relations_if_need(self, instance, reverse_relations, updated_related_data_dict):
113 # Reverse `reverse_relations` for correct delete priority
114 reverse_relations = OrderedDict(
115 reversed(list(reverse_relations.items())))
116
117 # Delete instances which is missed in data
118 for field_name, (related_field, field, field_source) in \
119 reverse_relations.items():
120 if field_name == 'files':
121 continue # lol hack TODO
122 model_class = field.Meta.model
123
124 related_data = self.get_initial()[field_name]
125 # Expand to array of one item for one-to-one for uniformity
126 if related_field.one_to_one:
127 related_data = [related_data]
128
129 # M2M relation can be as direct or as reverse. For direct relation
130 # we should use reverse relation name
131 if related_field.many_to_many and \
132 not isinstance(related_field, ForeignObjectRel):
133 related_field_lookup = {
134 related_field.remote_field.name: instance,
135 }
136 elif isinstance(related_field, GenericRelation):
137 related_field_lookup = \
138 self._get_generic_lookup(instance, related_field)
139 else:
140 related_field_lookup = {
141 related_field.name: instance,
142 }
143
144 current_ids = self._extract_related_pks(field, related_data)
145
146 updated_related_data = updated_related_data_dict.get(field_name, [])
147 updated_ids = self._extract_related_pks(field, updated_related_data)
148
149 current_ids = list(set(current_ids) | set(updated_ids))
150
151 try:
152 pks_to_delete = list(
153 model_class.objects.filter(
154 **related_field_lookup
155 ).exclude(
156 pk__in=current_ids
157 ).values_list('pk', flat=True)
158 )
159
160 if related_field.many_to_many:
161 # Remove relations from m2m table
162 m2m_manager = getattr(instance, field_source)
163 m2m_manager.remove(*pks_to_delete)
164 else:
165 model_class.objects.filter(pk__in=pks_to_delete).delete()
166
167 except ProtectedError as e:
168 instances = e.args[1]
169 self.fail('cannot_delete_protected', instances=", ".join([
170 str(instance) for instance in instances]))
171