diff --git a/src/documents/bulk_edit.py b/src/documents/bulk_edit.py index 4053ee8cb..542cdd3ba 100644 --- a/src/documents/bulk_edit.py +++ b/src/documents/bulk_edit.py @@ -27,7 +27,7 @@ from documents.models import DocumentType from documents.models import PaperlessTask from documents.models import StoragePath from documents.models import Tag -from documents.permissions import set_permissions_for_object +from documents.permissions import set_permissions_for_objects from documents.plugins.helpers import DocumentsStatusManager from documents.tasks import bulk_update_documents from documents.tasks import consume_file @@ -430,10 +430,13 @@ def set_permissions( else: qs.update(owner=owner) - for doc in qs: - set_permissions_for_object(permissions=set_permissions, object=doc, merge=merge) - affected_docs = list(qs.values_list("pk", flat=True)) + set_permissions_for_objects( + permissions=set_permissions, + model=Document, + pks=affected_docs, + merge=merge, + ) bulk_update_documents.apply_async( kwargs={"document_ids": affected_docs}, diff --git a/src/documents/permissions.py b/src/documents/permissions.py index d22882753..472c41494 100644 --- a/src/documents/permissions.py +++ b/src/documents/permissions.py @@ -173,6 +173,182 @@ def set_permissions_for_object( ) +def _resolve_permissions(codenames: set[str], ctype: ContentType) -> list[Permission]: + """ + Resolves `codenames` to Permission rows, raising like the single-object + assign_perm() this bulk path replaces does (via a `.get()` internally) + if any codename doesn't exist -- e.g. a client-supplied action name that + was never validated (BulkEditObjectsSerializer._validate_permissions + calls validate_set_permissions() only for its side-effecting id checks + and discards the filtered dict it returns, so an unrecognized action key + reaches this function as-is). A plain `.filter()` with no existence + check would otherwise silently build zero rows and no-op instead of + reporting the bad input. + """ + permission_objs = list( + Permission.objects.filter(content_type=ctype, codename__in=codenames), + ) + missing = codenames - {p.codename for p in permission_objs} + if missing: + raise Permission.DoesNotExist( + f"Permission matching query does not exist for codename(s): " + f"{', '.join(sorted(missing))}", + ) + return permission_objs + + +# Target number of permission rows to build in Python before handing them to +# bulk_create -- keeps peak memory bounded for a large "apply to all" call, +# independent of bulk_create's own batch_size (which only caps the size of +# each INSERT statement, not how many row objects exist in memory at once). +_PERMISSION_ROW_CHUNK_SIZE = 5000 + + +def _apply_bulk_permission_entry( + *, + perm_model: type[UserObjectPermission] | type[GroupObjectPermission], + identity_model: type[User] | type[Group], + identity_field: str, + ids: list[int], + codename: str, + permission_objs: list[Permission], + ctype: ContentType, + object_pks: list[str], + merge: bool, +) -> None: + # Only the ids are needed to build permission rows (via `_id=`), + # so avoid fetching full User/Group rows for identities that may not + # even end up being granted anything new. + add_ids = set( + identity_model.objects.filter(id__in=ids).values_list("id", flat=True), + ) + + if not merge: + existing_ids = set( + perm_model.objects.filter( + content_type=ctype, + object_pk__in=object_pks, + permission__codename=codename, + ).values_list(f"{identity_field}_id", flat=True), + ) + remove_ids = existing_ids - add_ids + if remove_ids: + perm_model.objects.filter( + content_type=ctype, + object_pk__in=object_pks, + permission__codename=codename, + **{f"{identity_field}_id__in": remove_ids}, + ).delete() + + if not add_ids: + return + + rows_per_pk = len(permission_objs) * len(add_ids) + pks_per_chunk = max(1, _PERMISSION_ROW_CHUNK_SIZE // rows_per_pk) + for start in range(0, len(object_pks), pks_per_chunk): + pk_chunk = object_pks[start : start + pks_per_chunk] + rows = [ + perm_model( + content_type=ctype, + object_pk=pk, + permission=permission_obj, + **{f"{identity_field}_id": identity_id}, + ) + for permission_obj in permission_objs + for pk in pk_chunk + for identity_id in add_ids + ] + # ignore_conflicts skips only rows that already exist as an exact + # (identity, permission, object) match -- the same de-dup the + # underlying (user|group, permission, object_pk) unique constraint + # already enforces for the single-object assign_perm() this + # replaces, so it doesn't change what counts as "already granted". + # batch_size caps how many rows go into a single INSERT so a huge + # chunk doesn't build one enormous statement. + perm_model.objects.bulk_create(rows, ignore_conflicts=True, batch_size=1000) + + +def set_permissions_for_objects( + permissions: dict, + model: type[Model], + pks: QuerySet | list, + *, + merge: bool = False, +) -> None: + """ + Bulk equivalent of set_permissions_for_object: applies the same + permission changes to every object identified by `pks` at once. + + Takes a model + pks (rather than model instances) deliberately -- the + permission rows built below only ever need `pk`, `content_type`, and + identity ids, so callers shouldn't have to fetch full rows (with every + other field) just to hand them to this function. + + Deliberately does not use guardian's queryset/list-aware assign_perm: + passing a list as the object routes to bulk_assign_perm, which skips + creating a direct permission row for anyone who already has the + permission via ANY group membership (it checks + ObjectPermissionChecker.has_perm, which is group-inheritance-aware) -- + unlike the single-object assign_perm this replaces, which always + ensures a direct row via get_or_create regardless of group-derived + access. Losing that guarantee would mean a later revocation of the + group's grant silently strips access an admin explicitly asked to be + direct. Bulk-creating rows straight against the permission models + instead (see _apply_bulk_permission_entry) preserves the original + always-create-a-direct-row semantics while still batching every object + and every identity into one query per action, rather than one query per + (object, user) pair. + """ + object_pks = [str(pk) for pk in pks] + if not object_pks: + return + + model_name = model.__name__.lower() + ctype = ContentType.objects.get_for_model(model) + + for action, entry in permissions.items(): + codename = f"{action}_{model_name}" + implied_codenames = {codename} + if action == "change": + # change gives view too + implied_codenames.add(f"view_{model_name}") + + # Resolved once per action (not once per users/groups branch) and + # shared between both below -- also where an unrecognized action + # name (see _resolve_permissions) is caught. + permission_objs = ( + _resolve_permissions(implied_codenames, ctype) + if "users" in entry or "groups" in entry + else [] + ) + + if "users" in entry: + _apply_bulk_permission_entry( + perm_model=UserObjectPermission, + identity_model=User, + identity_field="user", + ids=entry["users"], + codename=codename, + permission_objs=permission_objs, + ctype=ctype, + object_pks=object_pks, + merge=merge, + ) + + if "groups" in entry: + _apply_bulk_permission_entry( + perm_model=GroupObjectPermission, + identity_model=Group, + identity_field="group", + ids=entry["groups"], + codename=codename, + permission_objs=permission_objs, + ctype=ctype, + object_pks=object_pks, + merge=merge, + ) + + def permitted_object_ids( user: User | None, model: type[Model], diff --git a/src/documents/tests/test_api_objects.py b/src/documents/tests/test_api_objects.py index 05911febc..0cdf20dc2 100644 --- a/src/documents/tests/test_api_objects.py +++ b/src/documents/tests/test_api_objects.py @@ -2,10 +2,15 @@ import datetime import json from unittest import mock +from django.contrib.auth.models import Group from django.contrib.auth.models import Permission from django.contrib.auth.models import User +from django.db import connection from django.test import override_settings +from django.test.utils import CaptureQueriesContext from guardian.shortcuts import assign_perm +from guardian.shortcuts import get_groups_with_perms +from guardian.shortcuts import get_users_with_perms from rest_framework import status from rest_framework.test import APITestCase @@ -815,6 +820,59 @@ class TestBulkEditObjects(APITestCase): self.assertEqual(response.status_code, status.HTTP_200_OK) self.assertEqual(StoragePath.objects.count(), 0) + def test_bulk_objects_set_permissions_query_count_independent_of_object_count( + self, + ) -> None: + """ + GIVEN: + - Many tags are being bulk-edited to set permissions at once + WHEN: + - bulk_edit_objects API endpoint is called with set_permissions + operation over a small batch vs. a much larger one + THEN: + - The number of queries issued is the same either way -- each + user/group is applied across all tags with one batched call, + not one call per (tag, user) pair + """ + group1 = Group.objects.create(name="perm-group") + permissions = { + "view": {"users": [self.user1.id, self.user2.id], "groups": [group1.id]}, + "change": {"users": [self.user1.id], "groups": [group1.id]}, + } + + def run_with_n_tags(n: int) -> int: + tags = [Tag.objects.create(name=f"perm-tag-{n}-{i}") for i in range(n)] + with CaptureQueriesContext(connection) as ctx: + response = self.client.post( + "/api/bulk_edit_objects/", + json.dumps( + { + "objects": [t.id for t in tags], + "object_type": "tags", + "operation": "set_permissions", + "permissions": permissions, + "merge": False, + }, + ), + content_type="application/json", + ) + self.assertEqual(response.status_code, status.HTTP_200_OK) + for tag in tags: + self.assertEqual(get_users_with_perms(tag).count(), 2) + self.assertEqual(get_groups_with_perms(tag).count(), 1) + return len(ctx.captured_queries) + + small_batch_queries = run_with_n_tags(5) + large_batch_queries = run_with_n_tags(50) + + self.assertEqual( + small_batch_queries, + large_batch_queries, + "Expected the same query count regardless of tag count, got " + f"{small_batch_queries} queries for 5 tags vs. " + f"{large_batch_queries} for 50", + ) + def test_bulk_objects_delete_all_filtered(self) -> None: """ GIVEN: diff --git a/src/documents/tests/test_bulk_edit.py b/src/documents/tests/test_bulk_edit.py index 6063a9bbc..0e04c6a0b 100644 --- a/src/documents/tests/test_bulk_edit.py +++ b/src/documents/tests/test_bulk_edit.py @@ -5,8 +5,11 @@ from unittest import mock import pikepdf from django.contrib.auth.models import Group +from django.contrib.auth.models import Permission from django.contrib.auth.models import User +from django.db import connection from django.test import TestCase +from django.test.utils import CaptureQueriesContext from guardian.shortcuts import assign_perm from guardian.shortcuts import get_groups_with_perms from guardian.shortcuts import get_users_with_perms @@ -19,6 +22,7 @@ from documents.models import Document from documents.models import DocumentType from documents.models import StoragePath from documents.models import Tag +from documents.permissions import set_permissions_for_objects from documents.tests.utils import DirectoriesMixin @@ -510,6 +514,132 @@ class TestBulkEdit(DirectoriesMixin, TestCase): ) self.assertEqual(groups_with_perms.count(), 2) + @mock.patch("documents.tasks.bulk_update_documents.apply_async") + def test_set_permissions_query_count_independent_of_document_count( + self, + m, + ) -> None: + """ + GIVEN: + - Many documents are being bulk-edited to set permissions at once + WHEN: + - set_permissions runs over a small batch vs. a much larger one + THEN: + - The number of queries issued is the same either way -- each + user/group is applied across all documents with one batched + call, not one call per (document, user) pair + """ + permissions = { + "view": { + "users": [self.user1.id, self.user2.id], + "groups": [self.group2.id], + }, + "change": { + "users": [self.user1.id], + "groups": [self.group2.id], + }, + } + + def run_with_n_documents(n: int) -> int: + docs = [ + Document.objects.create(checksum=f"perm-{n}-{i}", title=f"perm-{n}-{i}") + for i in range(n) + ] + with CaptureQueriesContext(connection) as ctx: + bulk_edit.set_permissions( + [doc.id for doc in docs], + set_permissions=permissions, + owner=self.owner, + merge=False, + ) + for doc in docs: + self.assertEqual(get_users_with_perms(doc).count(), 2) + self.assertEqual(get_groups_with_perms(doc).count(), 1) + return len(ctx.captured_queries) + + small_batch_queries = run_with_n_documents(5) + large_batch_queries = run_with_n_documents(50) + + self.assertEqual( + small_batch_queries, + large_batch_queries, + "Expected the same query count regardless of document count, got " + f"{small_batch_queries} queries for 5 documents vs. " + f"{large_batch_queries} for 50", + ) + + @mock.patch("documents.tasks.bulk_update_documents.apply_async") + def test_set_permissions_grants_direct_perm_even_if_already_granted_via_group( + self, + m, + ) -> None: + """ + GIVEN: + - A user already has view access to a document via group + membership, with no direct grant of their own + WHEN: + - set_permissions explicitly grants that same user direct view + access via bulk_edit + THEN: + - A direct permission grant is created for the user, not skipped + because they already have equivalent access via the group + + Regression test: guardian's queryset-aware assign_perm() (routed to + when the target is a list/queryset) skips creating a direct row for + anyone whose ObjectPermissionChecker.has_perm() already returns True + -- which includes group-derived access. The single-object assign_perm + this bulk path replaces has no such check; it always ensures a + direct row via get_or_create. Losing that guarantee would mean + revoking the group's grant later silently strips access that was + supposed to be explicit. + """ + self.doc1.owner = self.user1 + self.doc1.save() + self.user1.groups.add(self.group1) + assign_perm("view_document", self.group1, self.doc1) + + bulk_edit.set_permissions( + [self.doc1.id], + set_permissions={ + "view": {"users": [self.user1.id], "groups": []}, + }, + merge=True, + ) + + direct_users = get_users_with_perms( + self.doc1, + only_with_perms_in=["view_document"], + with_group_users=False, + ) + self.assertIn(self.user1, direct_users) + + def test_set_permissions_for_objects_raises_for_unknown_action(self) -> None: + """ + GIVEN: + - An unrecognized permission action name with users to grant it + to + WHEN: + - set_permissions_for_objects is called + THEN: + - Permission.DoesNotExist is raised, not a silent no-op + + Regression test: the endpoint that calls this + (BulkEditObjectPermissionsView) never actually validates action + names against the raw client-supplied permissions dict -- + BulkEditObjectsSerializer._validate_permissions calls + validate_set_permissions() only for its side-effecting user/group id + checks and discards the filtered dict it returns -- so a bogus + action key reaches this function as-is. Resolving the Permission via + a bare `.filter()` (which returns empty instead of raising) would + silently drop the grant and report success. + """ + with self.assertRaises(Permission.DoesNotExist): + set_permissions_for_objects( + {"not_a_real_action": {"users": [self.user1.id], "groups": []}}, + Document, + [self.doc1.pk], + ) + @mock.patch("documents.models.Document.delete") def test_delete_documents_old_uuid_field(self, m) -> None: m.side_effect = Exception("Data too long for column 'transaction_id' at row 1") diff --git a/src/documents/views.py b/src/documents/views.py index 0d3f9de99..ef36a9197 100644 --- a/src/documents/views.py +++ b/src/documents/views.py @@ -178,7 +178,7 @@ from documents.permissions import has_perms_owner_aware from documents.permissions import has_system_status_permission from documents.permissions import permitted_document_ids from documents.permissions import permitted_object_ids -from documents.permissions import set_permissions_for_object +from documents.permissions import set_permissions_for_objects from documents.plugins.date_parsing import get_date_parser from documents.schema import generate_object_with_permissions_schema from documents.search import SearchHit @@ -4917,12 +4917,12 @@ class BulkEditObjectsView(PassUserMixin): qs_owner_update.update(owner=owner) if "permissions" in serializer.validated_data: - for obj in qs: - set_permissions_for_object( - permissions=permissions, - object=obj, - merge=merge, - ) + set_permissions_for_objects( + permissions=permissions, + model=object_class, + pks=qs.values_list("pk", flat=True), + merge=merge, + ) except Exception as e: logger.warning(