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.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 Document
    from documents.models import DocumentType
    from documents.models import StoragePath
    from documents.models import Tag

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


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 = tuple(DocumentFactory.create_batch(counts.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)

    # guardian's assign_perm is a get_or_create on the (subject, object,
    # permission) triple, so drawing (subject, document) pairs *with*
    # replacement -- rng.choice/rng.choice in a loop -- collides far more
    # than the raw draw count suggests once the subject pool is small
    # relative to the number of draws (a birthday-paradox effect): e.g. at
    # medium scale, 45,387 draws over only 5 groups x 20,000 documents =
    # 100,000 slots yields ~36,500 *distinct* pairs, not ~45,387. Sampling
    # distinct pairs up front guarantees the seeded row counts actually
    # land on the real-install ratios the scale profiles are derived from.
    for user_idx, doc_idx in _sample_distinct_pairs(
        rng,
        len(users),
        len(documents),
        n_user_perms,
    ):
        assign_perm("view_document", users[user_idx], documents[doc_idx])
    for group_idx, doc_idx in _sample_distinct_pairs(
        rng,
        len(groups),
        len(documents),
        n_group_perms,
    ):
        assign_perm("view_document", groups[group_idx], documents[doc_idx])

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