feat(benchmark): add per-backend reset and explain helpers

src/paperless_benchmark/db.py

@@ -0,0 +1,109 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from django.db import connection
+
+if TYPE_CHECKING:
+    from django.db.models import QuerySet
+
+
+def _reset_table_names() -> list[str]:
+    from guardian.models import GroupObjectPermission
+    from guardian.models import UserObjectPermission
+
+    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
+
+    return [
+        Document.tags.through._meta.db_table,
+        Document._meta.db_table,
+        Tag._meta.db_table,
+        Correspondent._meta.db_table,
+        DocumentType._meta.db_table,
+        StoragePath._meta.db_table,
+        UserObjectPermission._meta.db_table,
+        GroupObjectPermission._meta.db_table,
+    ]
+
+
+def _delete_non_superusers() -> None:
+    from django.contrib.auth import get_user_model
+
+    get_user_model().objects.filter(is_superuser=False).delete()
+
+
+def _reset_postgresql() -> None:
+    tables = _reset_table_names()
+    with connection.cursor() as cursor:
+        cursor.execute(f"TRUNCATE TABLE {', '.join(tables)} RESTART IDENTITY CASCADE;")
+    _delete_non_superusers()
+
+
+def _reset_mariadb() -> None:
+    # MariaDB's TRUNCATE has no CASCADE clause and refuses to truncate a
+    # table referenced by a foreign key while checks are enabled, so
+    # checks are disabled for the duration of the reset.
+    tables = _reset_table_names()
+    with connection.cursor() as cursor:
+        cursor.execute("SET FOREIGN_KEY_CHECKS = 0;")
+        try:
+            for table in tables:
+                cursor.execute(f"TRUNCATE TABLE {table};")
+        finally:
+            cursor.execute("SET FOREIGN_KEY_CHECKS = 1;")
+    _delete_non_superusers()
+
+
+def _reset_sqlite() -> None:
+    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
+
+    Document.global_objects.all().delete()
+    Tag.objects.all().delete()
+    Correspondent.objects.all().delete()
+    DocumentType.objects.all().delete()
+    StoragePath.objects.all().delete()
+    _delete_non_superusers()
+
+
+def reset_benchmark_data() -> None:
+    """
+    Remove all previously-seeded benchmark data (documents, tags,
+    correspondents, document types, storage paths, guardian permission
+    rows, and non-superuser users) so a fresh `benchmark seed` run starts
+    from an empty slate. Dispatches per-backend because TRUNCATE syntax
+    and cascade behavior differ across the 3 supported databases.
+    """
+    if connection.vendor == "postgresql":
+        _reset_postgresql()
+    elif connection.vendor == "mysql":
+        # MariaDB also reports vendor == "mysql" under Django's mysql backend.
+        _reset_mariadb()
+    else:
+        _reset_sqlite()
+
+
+def capture_explain(queryset: QuerySet) -> str:
+    """
+    Return the query plan for `queryset` using the current backend's
+    explain facility. PostgreSQL and MariaDB both support EXPLAIN ANALYZE
+    (real execution stats). SQLite only supports EXPLAIN QUERY PLAN (the
+    chosen plan, not real timing/row counts) -- that output is clearly
+    labeled rather than silently looking equivalent to the other two
+    backends' output.
+    """
+    sql, params = queryset.query.sql_with_params()
+    with connection.cursor() as cursor:
+        if connection.vendor in ("postgresql", "mysql"):
+            cursor.execute(f"EXPLAIN ANALYZE {sql}", params)
+            return "\n".join(str(row[0]) for row in cursor.fetchall())
+        cursor.execute(f"EXPLAIN QUERY PLAN {sql}", params)
+        rows = "\n".join(" | ".join(str(c) for c in row) for row in cursor.fetchall())
+        return f"(plan only -- no execution stats on SQLite)\n{rows}"