Skip to content

Optimize Role and Membership bulk creation with bulk_create #14341

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
rtibblesbot wants to merge 1 commit into
base: develop
Choose a base branch
from
+550 −28
Open

Optimize Role and Membership bulk creation with bulk_create #14341

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 26 additions & 16 deletions kolibri/core/auth/models.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1435,7 +1435,7 @@ def __str__(self):
user=self.user, collection=self.collection
)

def save(self, *args, **kwargs):
def validate_membership(self):
if self.collection.kind == collection_kinds.FACILITY:
raise InvalidMembershipError(
"Cannot create membership objects for facilities, as should already be a member by facility attribute"
Expand All @@ -1452,6 +1452,9 @@ def save(self, *args, **kwargs):
raise InvalidMembershipError(
"Cannot create membership for a user in a LearnerGroup or AdHocGroup when they are not a member of the parent Classrooom"
)

def save(self, *args, **kwargs):
self.validate_membership()
return super().save(*args, **kwargs)

def delete(self, **kwargs):
Expand Down Expand Up @@ -1528,7 +1531,7 @@ def __str__(self):
user=self.user, kind=self.kind, collection=self.collection
)

def save(self, *args, **kwargs):
def validate_role(self):
if (
self.collection.kind == collection_kinds.LEARNERGROUP
or self.collection.kind == collection_kinds.ADHOCLEARNERSGROUP
Expand All @@ -1537,21 +1540,28 @@ def save(self, *args, **kwargs):
raise InvalidRoleKind(
"Cannot assign roles to Learner Groups or AdHoc Groups"
)
if self.collection.kind == collection_kinds.CLASSROOM:
# We only support coaches to be assigned at the classroom level currently
if self.kind != role_kinds.COACH:
raise InvalidRoleKind("Can only assign Coach roles to Classrooms")

def ensure_coach_role_at_facility(self):
if self.collection.kind == collection_kinds.CLASSROOM:
if not Role.objects.filter(
user=self.user, collection_id=self.collection.parent_id
).exists():
# If the user doesn't already have a facility role, then create the assignable coach role for the user
# at the facility level.
Role.objects.create(
user=self.user,
collection_id=self.collection.parent_id,
kind=role_kinds.ASSIGNABLE_COACH,
)

def save(self, *args, **kwargs):
self.validate_role()
with transaction.atomic():
if self.collection.kind == collection_kinds.CLASSROOM:
# We only support coaches to be assigned at the classroom level currently
if self.kind != role_kinds.COACH:
raise InvalidRoleKind("Can only assign Coach roles to Classrooms")
if not Role.objects.filter(
user=self.user, collection_id=self.collection.parent_id
).exists():
# If the user doesn't already have a facility role, then create the assignable coach role for the user
# at the facility level.
Role.objects.create(
user=self.user,
collection_id=self.collection.parent_id,
kind=role_kinds.ASSIGNABLE_COACH,
)
self.ensure_coach_role_at_facility()
return super().save(*args, **kwargs)

def delete(self, **kwargs):
Expand Down
166 changes: 154 additions & 12 deletions kolibri/core/auth/serializers.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,15 +2,20 @@

from django.core.exceptions import ValidationError as DjangoValidationError
from django.core.validators import MinLengthValidator
from django.db import connections
from django.db import transaction
from morango.sync.backends.utils import calculate_max_sqlite_variables
from rest_framework import serializers
from rest_framework.exceptions import ParseError
from rest_framework.validators import UniqueTogetherValidator

from .constants import collection_kinds
from .constants import facility_presets
from .constants import role_kinds
from .errors import IncompatibleDeviceSettingError
from .errors import InvalidCollectionHierarchy
from .errors import InvalidMembershipError
from .errors import InvalidRoleKind
from .models import Classroom
from .models import Facility
from .models import FacilityDataset
Expand All @@ -27,16 +32,100 @@
logger = logging.getLogger(__name__)


def _prepare_for_bulk_create(instance):
"""
Prepare a morango SyncableModel instance for bulk_create by manually
setting the fields that would normally be set during save().
"""
instance.pre_save()
instance.id = instance.calculate_uuid()
instance._morango_dirty_bit = True


def _get_batch_size(Model):
"""
Calculate a safe batch_size for bulk_create to avoid SQLite variable limits.
Cap at 500 to prevent 'too many terms in compound SELECT' errors.
Same pattern as kolibri.core.auth.utils.migrate._batch_save.
"""
vendor = connections[Model.objects.db].vendor
if vendor == "sqlite":
return min(calculate_max_sqlite_variables() // len(Model._meta.fields), 500)
return 750


class RoleListSerializer(serializers.ListSerializer):
def create(self, validated_data):
created_objects = []
def validate(self, attrs):
for item in attrs:
instance = Role(**item)
try:
instance.validate_role()
except InvalidRoleKind as e:
raise serializers.ValidationError(str(e))
return attrs

def create(self, validated_data):
objects_to_create = []
for model_data in validated_data:
obj, created = Role.objects.get_or_create(**model_data)
if created:
created_objects.append(obj)
instance = Role(**model_data)
_prepare_for_bulk_create(instance)
objects_to_create.append(instance)

batch_size = _get_batch_size(Role)

with transaction.atomic():
# Filter out already-existing roles by their deterministic morango UUID
existing_ids = set(
Role.objects.filter(
id__in=[obj.id for obj in objects_to_create]
).values_list("id", flat=True)
)
new_objects = [
obj for obj in objects_to_create if obj.id not in existing_ids
]

if new_objects:
Role.objects.bulk_create(
new_objects,
batch_size=batch_size,
ignore_conflicts=True,
)

return created_objects
# Handle ASSIGNABLE_COACH side effect for classroom coach roles
classroom_roles = [
obj
for obj in new_objects
if obj.collection.kind == collection_kinds.CLASSROOM
]
if classroom_roles:
user_ids = {obj.user_id for obj in classroom_roles}
facility_ids = {obj.collection.parent_id for obj in classroom_roles}
users_with_facility_role = set(
Role.objects.filter(
collection_id__in=facility_ids,
user_id__in=user_ids,
).values_list("user_id", "collection_id")
)
assignable_roles = []
for obj in classroom_roles:
pair = (obj.user_id, obj.collection.parent_id)
if pair not in users_with_facility_role:
instance = Role(
user=obj.user,
collection_id=obj.collection.parent_id,
kind=role_kinds.ASSIGNABLE_COACH,
)
_prepare_for_bulk_create(instance)
assignable_roles.append(instance)
users_with_facility_role.add(pair)
if assignable_roles:
Role.objects.bulk_create(
assignable_roles,
batch_size=batch_size,
ignore_conflicts=True,
)

return new_objects


class RoleSerializer(serializers.ModelSerializer):
Expand Down Expand Up @@ -141,15 +230,68 @@ def validate(self, attrs):


class MembershipListSerializer(serializers.ListSerializer):
def create(self, validated_data):
created_objects = []
def validate(self, attrs):
lg_items = []
for item in attrs:
collection = item["collection"]
if collection.kind == collection_kinds.FACILITY:
raise serializers.ValidationError(
"Cannot create membership objects for facilities, "
"as should already be a member by facility attribute"
)
if collection.kind in (
collection_kinds.LEARNERGROUP,
collection_kinds.ADHOCLEARNERSGROUP,
):
lg_items.append(item)

if lg_items:
# Batch check parent classroom memberships with a single query
needed_pairs = {
(item["collection"].parent_id, item["user"].id) for item in lg_items
}
existing_memberships = set(
Membership.objects.filter(
collection_id__in={p[0] for p in needed_pairs},
user_id__in={p[1] for p in needed_pairs},
).values_list("collection_id", "user_id")
)
for item in lg_items:
pair = (item["collection"].parent_id, item["user"].id)
if pair not in existing_memberships:
raise serializers.ValidationError(
"Cannot create membership for a user in a "
"LearnerGroup or AdHocGroup when they are not a "
"member of the parent Classroom"
)
return attrs

def create(self, validated_data):
objects_to_create = []
for model_data in validated_data:
obj, created = Membership.objects.get_or_create(**model_data)
if created:
created_objects.append(obj)
instance = Membership(**model_data)
_prepare_for_bulk_create(instance)
objects_to_create.append(instance)

with transaction.atomic():
# Filter out already-existing memberships by their deterministic morango UUID
existing_ids = set(
Membership.objects.filter(
id__in=[obj.id for obj in objects_to_create]
).values_list("id", flat=True)
)
new_objects = [
obj for obj in objects_to_create if obj.id not in existing_ids
]

if new_objects:
Membership.objects.bulk_create(
new_objects,
batch_size=_get_batch_size(Membership),
ignore_conflicts=True,
)

return created_objects
return new_objects


class MembershipSerializer(serializers.ModelSerializer):
Expand Down
Loading
Loading