Skip to content

Commit 1dcf3ad

Browse files
pandafypandafy
committed
[fix] Fixes by @coderabbitai 
1 parent 931129e commit 1dcf3ad

2 files changed

Lines changed: 198 additions & 29 deletions

File tree

‎openwisp_radius/migrations/__init__.py‎

Lines changed: 68 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -35,12 +35,45 @@ def _flush_bulk_create(model, objects, batch_size=BATCH_SIZE):
3535
objects.clear()
3636

3737

38+
def _flush_bulk_update(model, objects, fields, batch_size=BATCH_SIZE):
39+
if objects:
40+
model.objects.bulk_update(objects, fields=fields, batch_size=batch_size)
41+
objects.clear()
42+
43+
3844
def _registered_user_extra_kwargs(registered_user, extra_fields=()):
3945
return {
4046
field_name: getattr(registered_user, field_name) for field_name in extra_fields
4147
}
4248

4349

50+
def _registered_user_method_priority_case():
51+
# Strong methods (anything that is not '' or 'email') must rank above the
52+
# weak fallbacks so rollback restores the strongest verification state.
53+
return Case(
54+
When(method="", then=Value(0)),
55+
When(method="email", then=Value(1)),
56+
default=Value(2),
57+
output_field=IntegerField(),
58+
)
59+
60+
61+
def _registered_user_method_priority(registered_user):
62+
if registered_user.method == "":
63+
return 0
64+
if registered_user.method == "email":
65+
return 1
66+
return 2
67+
68+
69+
def _registered_user_strength(registered_user):
70+
return (
71+
int(registered_user.is_verified),
72+
_registered_user_method_priority(registered_user),
73+
registered_user.modified,
74+
)
75+
76+
4477
def copy_registered_users_ctcr_forward(
4578
apps,
4679
schema_editor,
@@ -88,12 +121,7 @@ def copy_registered_users_ctcr_reverse(
88121
# Annotate each row with an explicit verification priority so that stronger
89122
# methods (anything that is not '' or 'email') sort before weaker ones.
90123
# Lexical ordering of 'method' would place '' first, picking the weakest.
91-
method_priority = Case(
92-
When(method="", then=Value(0)),
93-
When(method="email", then=Value(1)),
94-
default=Value(2),
95-
output_field=IntegerField(),
96-
)
124+
method_priority = _registered_user_method_priority_case()
97125
queryset = RegisteredUserNew.objects.annotate(
98126
method_priority=method_priority
99127
).order_by("user_id", "-is_verified", "-method_priority", "-modified")
@@ -187,21 +215,17 @@ def migrate_registered_users_multitenant_reverse(
187215
for user_id_batch in _batched_iterator(
188216
user_ids_qs.iterator(chunk_size=BATCH_SIZE), BATCH_SIZE
189217
):
190-
existing_globals = set(
191-
RegisteredUser.objects.filter(
218+
existing_globals = {
219+
registered_user.user_id: registered_user
220+
for registered_user in RegisteredUser.objects.filter(
192221
user_id__in=user_id_batch,
193222
organization__isnull=True,
194-
).values_list("user_id", flat=True)
195-
)
223+
)
224+
}
196225
# Annotate each row with an explicit verification priority so that stronger
197226
# methods (anything that is not '' or 'email') sort before weaker ones.
198227
# Lexical ordering of 'method' would place '' first, picking the weakest.
199-
method_priority = Case(
200-
When(method="", then=Value(0)),
201-
When(method="email", then=Value(1)),
202-
default=Value(2),
203-
output_field=IntegerField(),
204-
)
228+
method_priority = _registered_user_method_priority_case()
205229
org_records = (
206230
RegisteredUser.objects.filter(
207231
user_id__in=user_id_batch,
@@ -212,28 +236,43 @@ def migrate_registered_users_multitenant_reverse(
212236
)
213237

214238
to_create = []
239+
to_update = []
215240
to_delete_pks = []
216241
current_user_id = None
242+
update_fields = ["is_verified", "method", "modified", *extra_fields]
217243

218244
for registered_user in org_records.iterator(chunk_size=BATCH_SIZE):
219-
to_delete_pks.append(registered_user.pk)
220245
if registered_user.user_id == current_user_id:
246+
to_delete_pks.append(registered_user.pk)
221247
continue
222248
current_user_id = registered_user.user_id
223-
if registered_user.user_id in existing_globals:
224-
continue
225-
restored = RegisteredUser(
226-
id=uuid.uuid4(),
227-
user_id=registered_user.user_id,
228-
organization=None,
229-
is_verified=registered_user.is_verified,
230-
method=registered_user.method,
231-
**_registered_user_extra_kwargs(registered_user, extra_fields),
232-
)
233-
restored.modified = registered_user.modified
234-
to_create.append(restored)
249+
existing_global = existing_globals.get(registered_user.user_id)
250+
if existing_global is None:
251+
restored = RegisteredUser(
252+
id=uuid.uuid4(),
253+
user_id=registered_user.user_id,
254+
organization=None,
255+
is_verified=registered_user.is_verified,
256+
method=registered_user.method,
257+
**_registered_user_extra_kwargs(registered_user, extra_fields),
258+
)
259+
restored.modified = registered_user.modified
260+
to_create.append(restored)
261+
elif _registered_user_strength(registered_user) > _registered_user_strength(
262+
existing_global
263+
):
264+
existing_global.is_verified = registered_user.is_verified
265+
existing_global.method = registered_user.method
266+
existing_global.modified = registered_user.modified
267+
for field_name, value in _registered_user_extra_kwargs(
268+
registered_user, extra_fields
269+
).items():
270+
setattr(existing_global, field_name, value)
271+
to_update.append(existing_global)
272+
to_delete_pks.append(registered_user.pk)
235273

236274
_flush_bulk_create(RegisteredUser, to_create)
275+
_flush_bulk_update(RegisteredUser, to_update, fields=update_fields)
237276
if to_delete_pks:
238277
RegisteredUser.objects.filter(pk__in=to_delete_pks).delete()
239278

‎openwisp_radius/tests/test_migrations.py‎

Lines changed: 130 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,130 @@
1+
import swapper
2+
from django.apps.registry import apps
3+
from django.utils import timezone
4+
5+
from .. import settings as app_settings
6+
from ..migrations import migrate_registered_users_multitenant_reverse
7+
from ..utils import load_model
8+
from .mixins import BaseTestCase
9+
10+
RegisteredUser = load_model("RegisteredUser")
11+
Organization = swapper.load_model("openwisp_users", "Organization")
12+
User = swapper.load_model("auth", "User")
13+
14+
15+
class TestMigrations(BaseTestCase):
16+
def test_multitenant_reverse_updates_weaker_existing_global(self):
17+
"""
18+
Test that during migration rollback, a weaker existing global
19+
RegisteredUser is updated with data from a stronger org-scoped
20+
RegisteredUser instead of being left unchanged.
21+
"""
22+
user = self._create_user(username="rollback-stronger")
23+
org1 = self._create_org(name="rollback-org-1", slug="rollback-org-1")
24+
org2 = self._create_org(name="rollback-org-2", slug="rollback-org-2")
25+
modified_base = timezone.now()
26+
27+
# Create a weaker existing global (method="email")
28+
existing_global = RegisteredUser.objects.create(
29+
user=user,
30+
organization=None,
31+
is_verified=True,
32+
method="email",
33+
)
34+
RegisteredUser.objects.filter(pk=existing_global.pk).update(
35+
modified=modified_base
36+
)
37+
existing_global.refresh_from_db()
38+
39+
# Create org-scoped email (same strength as global but newer)
40+
org_email = RegisteredUser.objects.create(
41+
user=user,
42+
organization=org1,
43+
is_verified=True,
44+
method="email",
45+
)
46+
RegisteredUser.objects.filter(pk=org_email.pk).update(
47+
modified=modified_base + timezone.timedelta(minutes=10)
48+
)
49+
org_email.refresh_from_db()
50+
51+
# Create org-scoped mobile (strongest due to method priority)
52+
org_mobile = RegisteredUser.objects.create(
53+
user=user,
54+
organization=org2,
55+
is_verified=True,
56+
method="mobile_phone",
57+
)
58+
RegisteredUser.objects.filter(pk=org_mobile.pk).update(
59+
modified=modified_base - timezone.timedelta(minutes=10)
60+
)
61+
org_mobile.refresh_from_db()
62+
63+
# Rollback: should migrate strongest org-scoped (mobile_phone) to global
64+
migrate_registered_users_multitenant_reverse(
65+
apps, None, app_label="openwisp_radius"
66+
)
67+
68+
existing_global.refresh_from_db()
69+
self.assertIsNone(existing_global.organization)
70+
self.assertEqual(existing_global.method, "mobile_phone")
71+
self.assertTrue(existing_global.is_verified)
72+
self.assertEqual(existing_global.modified, org_mobile.modified)
73+
self.assertEqual(
74+
RegisteredUser.objects.filter(
75+
user=user, organization__isnull=False
76+
).count(),
77+
0,
78+
)
79+
80+
def test_multitenant_reverse_keeps_stronger_existing_global(self):
81+
"""
82+
Test that during migration rollback, if an existing global
83+
RegisteredUser is stronger than all org-scoped candidates,
84+
it is left unchanged and org-scoped rows are still cleaned up.
85+
"""
86+
user = self._create_user(username="rollback-global-wins")
87+
org = self._create_org(name="rollback-org-3", slug="rollback-org-3")
88+
modified_base = timezone.now()
89+
90+
# Create a stronger existing global (method="mobile_phone", newer timestamp)
91+
existing_global = RegisteredUser.objects.create(
92+
user=user,
93+
organization=None,
94+
is_verified=True,
95+
method="mobile_phone",
96+
)
97+
RegisteredUser.objects.filter(pk=existing_global.pk).update(
98+
modified=modified_base + timezone.timedelta(minutes=10)
99+
)
100+
existing_global.refresh_from_db()
101+
102+
# Create weaker org-scoped (method="social_login", older timestamp)
103+
org_specific = RegisteredUser.objects.create(
104+
user=user,
105+
organization=org,
106+
is_verified=True,
107+
method="social_login",
108+
)
109+
RegisteredUser.objects.filter(pk=org_specific.pk).update(modified=modified_base)
110+
org_specific.refresh_from_db()
111+
112+
# Rollback: global should remain unchanged (stronger), org-scoped deleted
113+
migrate_registered_users_multitenant_reverse(
114+
apps, None, app_label="openwisp_radius"
115+
)
116+
117+
existing_global.refresh_from_db()
118+
self.assertIsNone(existing_global.organization)
119+
self.assertEqual(existing_global.method, "mobile_phone")
120+
self.assertTrue(existing_global.is_verified)
121+
self.assertEqual(
122+
existing_global.modified,
123+
modified_base + timezone.timedelta(minutes=10),
124+
)
125+
self.assertEqual(
126+
RegisteredUser.objects.filter(
127+
user=user, organization__isnull=False
128+
).count(),
129+
0,
130+
)

0 commit comments

Comments
 (0)
⚡