Add backfill of user organizations (#3022)

Author: jkachel

b2b/constants.py

@@ -21,40 +21,44 @@
     CONTRACT_MEMBERSHIP_SSO,
 ]
 
-CONTRACT_MEMBERSHIP_CHOICES = zip(
-    [
-        CONTRACT_MEMBERSHIP_SSO,
-        CONTRACT_MEMBERSHIP_NONSSO,
-        CONTRACT_MEMBERSHIP_MANAGED,
-        CONTRACT_MEMBERSHIP_CODE,
-        CONTRACT_MEMBERSHIP_AUTO,
-    ],
-    [
-        CONTRACT_MEMBERSHIP_SSO_NAME,
-        CONTRACT_MEMBERSHIP_NONSSO_NAME,
-        CONTRACT_MEMBERSHIP_MANAGED_NAME,
-        CONTRACT_MEMBERSHIP_CODE_NAME,
-        CONTRACT_MEMBERSHIP_AUTO_NAME,
-    ],
+CONTRACT_MEMBERSHIP_CHOICES = list(
+    zip(
+        [
+            CONTRACT_MEMBERSHIP_SSO,
+            CONTRACT_MEMBERSHIP_NONSSO,
+            CONTRACT_MEMBERSHIP_MANAGED,
+            CONTRACT_MEMBERSHIP_CODE,
+            CONTRACT_MEMBERSHIP_AUTO,
+        ],
+        [
+            CONTRACT_MEMBERSHIP_SSO_NAME,
+            CONTRACT_MEMBERSHIP_NONSSO_NAME,
+            CONTRACT_MEMBERSHIP_MANAGED_NAME,
+            CONTRACT_MEMBERSHIP_CODE_NAME,
+            CONTRACT_MEMBERSHIP_AUTO_NAME,
+        ],
+    )
 )
 # When the integration_type field is removed, this should be removed as well.
 # It is just here so we can have the same choices in both integration_type and
 # membership_type fields. (Using the constant for `choices` consumes it.)
-CONTRACT_MEMBERSHIP_TYPE_CHOICES = zip(
-    [
-        CONTRACT_MEMBERSHIP_SSO,
-        CONTRACT_MEMBERSHIP_NONSSO,
-        CONTRACT_MEMBERSHIP_MANAGED,
-        CONTRACT_MEMBERSHIP_CODE,
-        CONTRACT_MEMBERSHIP_AUTO,
-    ],
-    [
-        CONTRACT_MEMBERSHIP_SSO_NAME,
-        CONTRACT_MEMBERSHIP_NONSSO_NAME,
-        CONTRACT_MEMBERSHIP_MANAGED_NAME,
-        CONTRACT_MEMBERSHIP_CODE_NAME,
-        CONTRACT_MEMBERSHIP_AUTO_NAME,
-    ],
+CONTRACT_MEMBERSHIP_TYPE_CHOICES = list(
+    zip(
+        [
+            CONTRACT_MEMBERSHIP_SSO,
+            CONTRACT_MEMBERSHIP_NONSSO,
+            CONTRACT_MEMBERSHIP_MANAGED,
+            CONTRACT_MEMBERSHIP_CODE,
+            CONTRACT_MEMBERSHIP_AUTO,
+        ],
+        [
+            CONTRACT_MEMBERSHIP_SSO_NAME,
+            CONTRACT_MEMBERSHIP_NONSSO_NAME,
+            CONTRACT_MEMBERSHIP_MANAGED_NAME,
+            CONTRACT_MEMBERSHIP_CODE_NAME,
+            CONTRACT_MEMBERSHIP_AUTO_NAME,
+        ],
+    )
 )
 
 B2B_RUN_TAG_FORMAT = "{year}_C{contract_id}"

b2b/migrations/0010_backfill_user_b2b_orgs.py

@@ -0,0 +1,52 @@
+# Generated by Django 4.2.25 on 2025-10-20 17:55
+
+from django.db import migrations
+
+
+def noop_reverse(apps, schema_editor):
+    """No-op the reverse - we won't be able to figure out what orgs were there before."""
+
+    return
+
+
+def backfill_user_org_memberships(apps, schema_editor):
+    """
+    Backfill each user's B2B organization memberships from their contracts.
+
+    0009 adds in a hard FK between orgs and users. We were calculating this before,
+    which has the side effect that everyone starts from zero. This is fine for
+    users who're in orgs that are in Keycloak already - they'll get added back in
+    as part of the authentication process - but everyone who has used an
+    enrollment code is caught out. So, do the calculation one last time and fill
+    out the FK table.
+
+    We will set these with the "keep_until_seen" flag set to True, since there's
+    no guarantee the orgs we're backfilling exist in Keycloak at this point.
+    """
+
+    User = apps.get_model("users", "User")
+    OrganizationPage = apps.get_model("b2b", "OrganizationPage")
+
+    for user in User.objects.all():
+        user_org_ids = (
+            user.b2b_contracts.values_list("organization", flat=True).distinct().all()
+        )
+        user_orgs = OrganizationPage.objects.filter(pk__in=user_org_ids).all()
+
+        for org in user_orgs:
+            user.b2b_organizations.add(
+                org,
+                through_defaults={
+                    "keep_until_seen": True,
+                },
+            )
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("b2b", "0009_contractpage_membership_type_and_more"),
+    ]
+
+    operations = [
+        migrations.RunPython(backfill_user_org_memberships, noop_reverse),
+    ]