govoplan-files/tests/benchmark_archive_storage.py

"""Isolated synthetic archive benchmark; never reads the configured Files store.

Run from this repository: python tests/benchmark_archive_storage.py.
Uses the existing temporary SQLite/local-storage recovery fixture so durable
ledger, checksum verification, commit settlement and compensation stay enabled.
"""

from __future__ import annotations

import argparse
import cProfile
from contextlib import ExitStack
from collections import Counter
from io import BytesIO
import json
from pathlib import Path
import pstats
import random
import re
import shutil
import subprocess
import tempfile
import time
from unittest.mock import patch
import zipfile

import pyzipper
from sqlalchemy import event

from govoplan_files.backend.change_tracking import register_files_change_tracking
from govoplan_files.backend.storage.archives import (
    extract_archive_upload,
    _read_selected_zip_members,
)
from govoplan_files.backend.storage.native_zip import native_zip_library
from test_storage_recovery import StorageRecoveryTests


def storage_probe(
    member_count: int,
    member_size: int,
    profile: bool,
    baseline_revision: str | None = None,
):
    case = StorageRecoveryTests()
    case.setUp()
    try:
        register_files_change_tracking()
        data = BytesIO()
        rng = random.Random(20260907)
        with zipfile.ZipFile(data, "w", compression=zipfile.ZIP_DEFLATED) as archive:
            for index in range(member_count):
                archive.writestr(
                    f"documents/{index:04d}.dat", rng.randbytes(member_size)
                )
        queries = []
        event.listen(
            case.engine,
            "before_cursor_execute",
            lambda conn, cursor, statement, parameters, context, executemany: (
                queries.append(statement)
            ),
        )
        profiler = cProfile.Profile()
        if profile:
            profiler.enable()
        with ExitStack() as overrides:
            overrides.enter_context(
                patch(
                    "govoplan_files.backend.storage.files.get_storage_backend",
                    return_value=case.backend,
                )
            )
            if baseline_revision:
                # Read an earlier local revision into this isolated process;
                # never replace shared worktree files or connect to live data.
                source = subprocess.run(
                    [
                        "git",
                        "show",
                        f"{baseline_revision}:src/govoplan_files/backend/storage/files.py",
                    ],
                    check=True,
                    capture_output=True,
                    text=True,
                ).stdout
                namespace = {"__name__": "govoplan_files.backend.storage.files"}
                exec(
                    compile(source, f"{baseline_revision}:files.py", "exec"), namespace
                )
                namespace.update(
                    get_storage_backend=lambda: case.backend,
                    _storage_backend_name=lambda: case.backend.name,
                    _storage_bucket_name=lambda: "",
                )
                overrides.enter_context(
                    patch(
                        "govoplan_files.backend.storage.archives.create_file_asset",
                        namespace["create_file_asset"],
                    )
                )
            started = time.perf_counter()
            result = extract_archive_upload(
                case.session,
                tenant_id="tenant-1",
                owner_type="user",
                owner_id="user-1",
                user_id="user-1",
                archive_data=data.getvalue(),
                filename="synthetic.zip",
                folder="imported",
                campaign_id=None,
            )
            storage_done = time.perf_counter()
            storage_queries = len(queries)
            case.session.commit()
            completed = time.perf_counter()
        if profile:
            profiler.disable()
        print(
            json.dumps(
                {
                    "probe": "durable-storage",
                    "storage_revision": baseline_revision or "working-tree",
                    "members": len(result),
                    "member_bytes": member_size,
                    "extract_store_seconds": round(storage_done - started, 4),
                    "settlement_seconds": round(completed - storage_done, 4),
                    "total_seconds": round(completed - started, 4),
                    "store_queries": storage_queries,
                    "total_queries": len(queries),
                    "queries_by_table": dict(
                        Counter(
                            (
                                match.group(1)
                                if (
                                    match := re.search(
                                        r"(?:FROM|INTO|UPDATE)\s+([a-z_]+)", statement
                                    )
                                )
                                else "other"
                            )
                            for statement in queries
                        )
                    ),
                }
            )
        )
        if profile:
            pstats.Stats(profiler).strip_dirs().sort_stats("cumulative").print_stats(24)
    finally:
        case.doCleanups()


def zipcrypto_probe(size: int):
    binary = shutil.which("zip")
    if not binary:
        print(
            json.dumps(
                {"probe": "zipcrypto", "skipped": "fixture ZIP writer unavailable"}
            )
        )
        return
    # Public fixture password only: never pass a user credential to a process.
    with tempfile.TemporaryDirectory(prefix="files-synthetic-zipcrypto-") as temporary:
        path = Path(temporary)
        content = random.Random(20260907).randbytes(size)
        (path / "synthetic.bin").write_bytes(content)
        subprocess.run(
            [binary, "-q", "-P", "fixture-only", "synthetic.zip", "synthetic.bin"],
            cwd=path,
            check=True,
        )
        for implementation in (pyzipper.AESZipFile, zipfile.ZipFile):
            started = time.perf_counter()
            with implementation(path / "synthetic.zip") as archive:
                result = archive.read("synthetic.bin", pwd=b"fixture-only")
            assert result == content
            print(
                json.dumps(
                    {
                        "probe": "zipcrypto",
                        "reader": implementation.__module__,
                        "bytes": size,
                        "seconds": round(time.perf_counter() - started, 4),
                    }
                )
            )
        if native_zip_library() is not None:
            started = time.perf_counter()
            result = list(
                _read_selected_zip_members(
                    path / "synthetic.zip",
                    selected_files={"synthetic.bin"},
                    password="fixture-only",
                    max_file_bytes=size,
                    max_total_bytes=size,
                )
            )
            assert result == [("synthetic.bin", content)]
            print(
                json.dumps(
                    {
                        "probe": "zipcrypto",
                        "reader": "libarchive-with-independent-crc",
                        "bytes": size,
                        "seconds": round(time.perf_counter() - started, 4),
                    }
                )
            )


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--members", type=int, default=100)
    parser.add_argument("--member-size", type=int, default=4096)
    parser.add_argument("--zipcrypto-bytes", type=int, default=4 * 1024 * 1024)
    parser.add_argument("--profile", action="store_true")
    parser.add_argument(
        "--baseline-storage-revision",
        help="Read a trusted earlier local Git revision for isolated storage comparison, without changing the worktree",
    )
    args = parser.parse_args()
    storage_probe(
        args.members, args.member_size, args.profile, args.baseline_storage_revision
    )
    if args.zipcrypto_bytes:
        zipcrypto_probe(args.zipcrypto_bytes)