correctly batch requests

Signed-off-by: Jens Langhammer <jens@goauthentik.io>
goauthentik · Aug 22, 2024 · 1e39705 · 1e39705
1 parent dacf6f8
commit 1e39705
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 21 deletions.
diff --git a/authentik/providers/scim/clients/groups.py b/authentik/providers/scim/clients/groups.py
@@ -1,5 +1,7 @@
 """Group client"""
 
+from itertools import batched
+
 from pydantic import ValidationError
 from pydanticscim.group import GroupMember
 from pydanticscim.responses import PatchOp, PatchOperation
@@ -58,7 +60,9 @@ def to_schema(self, obj: Group, connection: SCIMProviderGroup) -> SCIMGroupSchem
 
         if not self._config.patch.supported:
             users = list(obj.users.order_by("id").values_list("id", flat=True))
-            connections = SCIMProviderUser.objects.filter(provider=self.provider, user__pk__in=users)
+            connections = SCIMProviderUser.objects.filter(
+                provider=self.provider, user__pk__in=users
+            )
             members = []
             for user in connections:
                 members.append(
@@ -160,14 +164,18 @@ def _patch(
         group_id: str,
         *ops: PatchOperation,
     ):
-        req = PatchRequest(Operations=ops)
-        self._request(
-            "PATCH",
-            f"/Groups/{group_id}",
-            json=req.model_dump(
-                mode="json",
-            ),
-        )
+        chunk_size = self._config.bulk.maxOperations
+        if chunk_size < 1:
+            chunk_size = len(ops)
+        for chunk in batched(ops, chunk_size):
+            req = PatchRequest(Operations=list(chunk))
+            self._request(
+                "PATCH",
+                f"/Groups/{group_id}",
+                json=req.model_dump(
+                    mode="json",
+                ),
+            )
 
     def _patch_add_users(self, group: Group, users_set: set[int]):
         """Add users in users_set to group"""
@@ -188,11 +196,14 @@ def _patch_add_users(self, group: Group, users_set: set[int]):
             return
         self._patch(
             scim_group.scim_id,
-            PatchOperation(
-                op=PatchOp.add,
-                path="members",
-                value=[{"value": x} for x in user_ids],
-            ),
+            *[
+                PatchOperation(
+                    op=PatchOp.add,
+                    path="members",
+                    value=[{"value": x}],
+                )
+                for x in user_ids
+            ],
         )
 
     def _patch_remove_users(self, group: Group, users_set: set[int]):
@@ -214,9 +225,12 @@ def _patch_remove_users(self, group: Group, users_set: set[int]):
             return
         self._patch(
             scim_group.scim_id,
-            PatchOperation(
-                op=PatchOp.remove,
-                path="members",
-                value=[{"value": x} for x in user_ids],
-            ),
+            *[
+                PatchOperation(
+                    op=PatchOp.remove,
+                    path="members",
+                    value=[{"value": x}],
+                )
+                for x in user_ids
+            ],
         )
diff --git a/authentik/providers/scim/clients/schema.py b/authentik/providers/scim/clients/schema.py
@@ -1,9 +1,11 @@
 """Custom SCIM schemas"""
 
+from pydantic import Field
 from pydanticscim.group import Group as BaseGroup
 from pydanticscim.responses import PatchRequest as BasePatchRequest
 from pydanticscim.responses import SCIMError as BaseSCIMError
-from pydanticscim.service_provider import Bulk, ChangePassword, Filter, Patch, Sort
+from pydanticscim.service_provider import Bulk as BaseBulk
+from pydanticscim.service_provider import ChangePassword, Filter, Patch, Sort
 from pydanticscim.service_provider import (
     ServiceProviderConfiguration as BaseServiceProviderConfiguration,
 )
@@ -29,10 +31,16 @@ class Group(BaseGroup):
     meta: dict | None = None
 
 
+class Bulk(BaseBulk):
+
+    maxOperations: int = Field()
+
+
 class ServiceProviderConfiguration(BaseServiceProviderConfiguration):
     """ServiceProviderConfig with fallback"""
 
     _is_fallback: bool | None = False
+    bulk: Bulk = Field(..., description="A complex type that specifies bulk configuration options.")
 
     @property
     def is_fallback(self) -> bool:
@@ -45,7 +53,7 @@ def default() -> "ServiceProviderConfiguration":
         """Get default configuration, which doesn't support any optional features as fallback"""
         return ServiceProviderConfiguration(
             patch=Patch(supported=False),
-            bulk=Bulk(supported=False),
+            bulk=Bulk(supported=False, maxOperations=0),
             filter=Filter(supported=False),
             changePassword=ChangePassword(supported=False),
             sort=Sort(supported=False),