Fix chunk iteration for user migration script (#4378)

Author: rhysyngsun

keycloak_user_export/management/commands/migrate_users_to_keycloak.py

@@ -8,7 +8,6 @@
 import requests
 
 from keycloak_user_export.models import UserExportToKeycloak
-from open_discussions.utils import chunks
 
 User = get_user_model()
 
@@ -221,13 +220,12 @@ def handle(self, *args, **kwargs):
                 social_auth__provider=kwargs["filter_provider_name"]
             )
         unsynced_users = (
-            User.objects.only("email")
+            User.objects.only("email", "date_joined")
             .filter(is_active=True)
             .exclude(social_auth__provider="ol-oidc")
             .exclude(userexporttokeycloak__isnull=False)
             .filter(unsynced_users_social_auth_query)
-            .select_related("userexporttokeycloak")
-            .prefetch_related("social_auth")
+            .prefetch_related("profile")
         )
 
         with requests.Session() as session:
@@ -239,7 +237,9 @@ def handle(self, *args, **kwargs):
 
             # Process batches of the users who must be exported.
             batch_size = kwargs["batch_size"]
-            for batch in chunks(unsynced_users, chunk_size=batch_size):
+            # the users are updated in such a way that they're excluded from this query the next time
+            # so we just keep grabbing a slice from the beginning of the query until it returns empty
+            while batch := unsynced_users[:batch_size]:
                 payload = {
                     "ifResourceExists": "SKIP",
                     "realm": settings.KEYCLOAK_REALM_NAME,

keycloak_user_export/management/commands/tests/test_migrate_users_to_keycloak.py

@@ -1,4 +1,5 @@
 import json
+from math import ceil
 
 from django.contrib.auth import get_user_model
 from django.core.management import call_command
@@ -20,11 +21,25 @@ def kc_settings(settings):
     settings.KEYCLOAK_REALM_NAME = "pytest-realm"
 
 
-def test_migrate_users_to_keycloak(mocked_responses):
+@pytest.mark.parametrize(
+    "num_users, num_already_synced, batch_size",
+    [
+        (100, 0, 10),
+        (100, 20, 10),
+        (100, 99, 10),
+    ],
+)
+def test_migrate_users_to_keycloak(
+    mocked_responses,
+    django_assert_num_queries,
+    num_users,
+    num_already_synced,
+    batch_size,
+):
     """Test that the migrate_users_to_keycloak command functions as expected"""
-    users = UserFactory.create_batch(100)
+    users = UserFactory.create_batch(num_users)
     inactive_users = UserFactory.create_batch(10, is_active=False)
-    already_synced = users[:20]
+    already_synced = users[:num_already_synced]
 
     UserExportToKeycloak.objects.bulk_create(
         [UserExportToKeycloak(user=user) for user in already_synced]
@@ -68,11 +83,16 @@ def partial_import_callback(request):
         content_type="application/json",
     )
 
-    call_command(
-        "migrate_users_to_keycloak",
-        "--client-id=test-client-id",
-        "--client-secret=test-client-secret",
-    )
+    # 3 queries per batch plus 1 for the last empty query to terminate the loop
+    num_batches = ceil((num_users - num_already_synced) / batch_size)
+    num_queries = num_batches * 3 + 1
+    with django_assert_num_queries(num_queries):
+        call_command(
+            "migrate_users_to_keycloak",
+            "--client-id=test-client-id",
+            "--client-secret=test-client-secret",
+            f"--batch-size={batch_size}",
+        )
 
     assert (
         list(User.objects.filter(userexporttokeycloak__isnull=False).order_by("id"))