paperless-ngx/profiling/seed.py

# profiling/seed.py
from __future__ import annotations

import random
from dataclasses import dataclass
from typing import TYPE_CHECKING
from typing import Literal

from django.contrib.auth.models import Group
from django.contrib.auth.models import User
from guardian.shortcuts import assign_perm

from documents.models import Document
from documents.tests.factories import CorrespondentFactory
from documents.tests.factories import DocumentFactory
from documents.tests.factories import DocumentTypeFactory
from documents.tests.factories import StoragePathFactory
from documents.tests.factories import TagFactory

if TYPE_CHECKING:
    from documents.models import Correspondent
    from documents.models import DocumentType
    from documents.models import StoragePath
    from documents.models import Tag

ScaleProfile = Literal["medium", "large"]

# Fraction of seeded documents given a real owner; the remainder stays
# owner=None, mirroring a real install's mix of owned and legacy/imported
# unowned documents.
_OWNED_FRACTION = 0.9


class _ScaleCounts:
    __slots__ = (
        "correspondents",
        "document_types",
        "documents",
        "groups",
        "storage_paths",
        "tags",
        "users",
    )

    def __init__(
        self,
        *,
        documents: int,
        tags: int,
        correspondents: int,
        document_types: int,
        storage_paths: int,
        users: int,
        groups: int,
    ) -> None:
        self.documents = documents
        self.tags = tags
        self.correspondents = correspondents
        self.document_types = document_types
        self.storage_paths = storage_paths
        self.users = users
        self.groups = groups


_SCALE_PROFILES: dict[ScaleProfile, _ScaleCounts] = {
    "medium": _ScaleCounts(
        documents=20_000,
        tags=100,
        correspondents=300,
        document_types=50,
        storage_paths=20,
        users=10,
        groups=5,
    ),
    "large": _ScaleCounts(
        documents=360_000,
        tags=1_000,
        correspondents=5_000,
        document_types=300,
        storage_paths=50,
        users=25,
        groups=10,
    ),
}
# Ratios measured from a real install (discussion #13276): 1,414 user-perm
# rows / 27,232 group-perm rows over 12,000 documents.
_USER_PERM_ROWS_PER_DOC = 1_414 / 12_000
_GROUP_PERM_ROWS_PER_DOC = 27_232 / 12_000


def _sample_distinct_pairs(
    rng: random.Random,
    n_subjects: int,
    n_objects: int,
    count: int,
) -> set[tuple[int, int]]:
    """
    Rejection-sample `count` distinct (subject_index, object_index) pairs.
    `count` is always well under `n_subjects * n_objects` for every scale
    profile, so this terminates quickly.
    """
    if count > n_subjects * n_objects:
        msg = (
            f"cannot sample {count} distinct pairs from only "
            f"{n_subjects * n_objects} possible (subject, object) slots"
        )
        raise ValueError(msg)
    pairs: set[tuple[int, int]] = set()
    while len(pairs) < count:
        pairs.add((rng.randrange(n_subjects), rng.randrange(n_objects)))
    return pairs


@dataclass(frozen=True, slots=True)
class SeededData:
    users: tuple[User, ...]
    groups: tuple[Group, ...]
    documents: tuple[Document, ...]
    tags: tuple[Tag, ...]
    correspondents: tuple[Correspondent, ...]
    document_types: tuple[DocumentType, ...]
    storage_paths: tuple[StoragePath, ...]


def seed_permission_dataset(
    scale: ScaleProfile = "medium",
    *,
    seed: int = 1337,
) -> SeededData:
    """
    Build a dataset whose document count and guardian-permission-row ratios
    mirror real bug reports, so profiling results are representative of
    actual large installs rather than an arbitrary small fixture.
    """
    counts = _SCALE_PROFILES[scale]
    rng = random.Random(seed)

    users = tuple(
        User.objects.create_user(username=f"profile_user_{i}")
        for i in range(counts.users)
    )
    groups = tuple(
        Group.objects.create(name=f"profile_group_{i}") for i in range(counts.groups)
    )
    for user in users:
        user.groups.add(rng.choice(groups))

    documents = list(DocumentFactory.create_batch(counts.documents))

    # Assign realistic ownership -- without this, every document stays
    # owner=None and the "owned or unowned" visibility clause is trivially
    # true for every row, so no guardian permission check ever needs to run
    # (see the amendment note above this code block).
    owned_documents = []
    for document in documents:
        if rng.random() < _OWNED_FRACTION:
            document.owner = rng.choice(users)
            owned_documents.append(document)
    Document.objects.bulk_update(owned_documents, ["owner"], batch_size=2_000)
    documents = tuple(documents)

    tags = tuple(TagFactory.create_batch(counts.tags))
    correspondents = tuple(CorrespondentFactory.create_batch(counts.correspondents))
    document_types = tuple(DocumentTypeFactory.create_batch(counts.document_types))
    storage_paths = tuple(StoragePathFactory.create_batch(counts.storage_paths))

    n_user_perms = round(counts.documents * _USER_PERM_ROWS_PER_DOC)
    n_group_perms = round(counts.documents * _GROUP_PERM_ROWS_PER_DOC)

    # Grant permissions only on OWNED documents -- a grant on an already-
    # unowned document is a no-op for visibility (unowned documents are
    # visible to everyone regardless), so restricting the pool here is what
    # makes each assign_perm() call actually matter for someone, and is
    # what exercises the guardian-permission-join code path this profiling
    # effort exists to measure.
    for user_idx, owned_idx in _sample_distinct_pairs(
        rng,
        len(users),
        len(owned_documents),
        n_user_perms,
    ):
        assign_perm("view_document", users[user_idx], owned_documents[owned_idx])
    for group_idx, owned_idx in _sample_distinct_pairs(
        rng,
        len(groups),
        len(owned_documents),
        n_group_perms,
    ):
        assign_perm("view_document", groups[group_idx], owned_documents[owned_idx])

    return SeededData(
        users=users,
        groups=groups,
        documents=documents,
        tags=tags,
        correspondents=correspondents,
        document_types=document_types,
        storage_paths=storage_paths,
    )