"""Benchmark Generator — produce large volumes of realistic data for performance testing. Configurable via BenchmarkSpec: define entity types, counts, distributions, and relationships. Supports multi-tenant, multi-workspace data generation with reproducible seeds. """ from __future__ import annotations import random import time from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any from simorgh.apps.provisioning.models import FixtureRun if TYPE_CHECKING: from collections.abc import Callable from simorgh.apps.provisioning.faker_support import ScopedFaker @dataclass class EntitySpec: """Specification for a single entity type in a benchmark run.""" model: str count: int depends_on: list[str] = field(default_factory=list) batch_size: int = 500 overrides: dict[str, Any] = field(default_factory=dict) distribution: str = "uniform" field_generators: dict[str, str | Callable[[ScopedFaker, int], Any]] = field( default_factory=dict ) @dataclass class BenchmarkSpec: """Top-level benchmark specification.""" name: str entities: list[EntitySpec] tenant_count: int = 1 workspace_count: int = 0 seed: int | None = None locale: str = "fa_IR" description: str = "" def total_entities(self) -> int: return sum(e.count for e in self.entities) class BenchmarkGenerator: """Generate large volumes of realistic benchmark data. Usage: spec = BenchmarkSpec( name="crm_benchmark", entities=[ EntitySpec(model="crm.Contact", count=10_000), EntitySpec(model="crm.Lead", count=5_000, depends_on=["crm.Contact"]), ], tenant_count=5, ) gen = BenchmarkGenerator(spec) report = gen.run() print(report.summary()) """ def __init__(self, spec: BenchmarkSpec) -> None: self._spec = spec if spec.seed is not None: random.seed(spec.seed) def run( self, context: dict[str, Any] | None = None, *, dry_run: bool = False, progress_callback: Callable[[int, int, str], None] | None = None, ) -> BenchmarkReport: """Execute the benchmark data generation. Args: context: Optional base context with tenant/workspace data. dry_run: If True, estimate without creating data. progress_callback: Called with (current, total, entity_name) for progress. Returns: BenchmarkReport with timing and entity counts. """ from django.apps import apps from simorgh.apps.provisioning.faker_support import ScopedFaker if context is None: context = {} if dry_run: return BenchmarkReport( spec_name=self._spec.name, total_entities=self._spec.total_entities(), duration_ms=0, dry_run=True, ) faker = ScopedFaker(locales=self._spec.locale, seed=self._spec.seed) context["faker"] = faker resolved = self._resolve_entity_order() total = sum(e.count for e in resolved) current = 0 entity_counts: dict[str, int] = {} start = time.monotonic() for entity in resolved: model_cls = apps.get_model(entity.model) batch: list[Any] = [] for i in range(entity.count): kwargs = self._build_kwargs(entity, i, faker, context) batch.append(model_cls(**kwargs)) if len(batch) >= entity.batch_size: model_cls.objects.bulk_create(batch) current += len(batch) entity_counts[entity.model] = entity_counts.get(entity.model, 0) + len(batch) if progress_callback: progress_callback(current, total, entity.model) batch = [] if batch: model_cls.objects.bulk_create(batch) current += len(batch) entity_counts[entity.model] = entity_counts.get(entity.model, 0) + len(batch) if progress_callback: progress_callback(current, total, entity.model) duration_ms = int((time.monotonic() - start) * 1000) FixtureRun.objects.create( kind=FixtureRun.RunKind.BENCHMARK, fixture_names=[e.model for e in resolved], items_created=sum(entity_counts.values()), duration_ms=duration_ms, success=True, meta={ "spec": self._spec.name, "entities": entity_counts, }, ) return BenchmarkReport( spec_name=self._spec.name, entity_counts=entity_counts, total_entities=sum(entity_counts.values()), duration_ms=duration_ms, ) def _resolve_entity_order(self) -> list[EntitySpec]: """Topological sort of entity specs by dependencies.""" name_map = {e.model: e for e in self._spec.entities} in_degree: dict[str, int] = {} graph: dict[str, list[str]] = {} for e in self._spec.entities: in_degree[e.model] = len(e.depends_on) graph[e.model] = [] for e in self._spec.entities: for dep in e.depends_on: if dep in graph: graph[dep].append(e.model) from collections import deque queue = deque(m for m, d in in_degree.items() if d == 0) order: list[EntitySpec] = [] while queue: m = queue.popleft() if m in name_map: order.append(name_map[m]) for neighbor in graph.get(m, []): in_degree[neighbor] -= 1 if in_degree[neighbor] == 0: queue.append(neighbor) remaining = [e for e in self._spec.entities if e not in order] order.extend(remaining) return order def _build_kwargs( self, entity: EntitySpec, index: int, faker: ScopedFaker, context: dict[str, Any], ) -> dict[str, Any]: kwargs = dict(entity.overrides) tenant = context.get("tenant") if tenant is not None: kwargs.setdefault("tenant", tenant) kwargs.setdefault("tenant_id", tenant.pk) for field_name, gen in entity.field_generators.items(): if callable(gen): kwargs[field_name] = gen(faker, index) elif isinstance(gen, str): faker_attr = getattr(faker, gen, None) kwargs[field_name] = faker_attr() if callable(faker_attr) else gen return kwargs @dataclass class BenchmarkReport: """Result of a benchmark data generation run.""" spec_name: str entity_counts: dict[str, int] = field(default_factory=dict) total_entities: int = 0 duration_ms: int = 0 dry_run: bool = False rows_per_second: float = field(init=False) def __post_init__(self) -> None: if self.duration_ms > 0: self.rows_per_second = self.total_entities / (self.duration_ms / 1000) else: self.rows_per_second = 0.0 def summary(self) -> str: lines = [f"Benchmark: {self.spec_name}"] if self.dry_run: lines.append(" Mode: DRY RUN") else: lines.append(f" Duration: {self.duration_ms}ms ({self.rows_per_second:,.0f} rows/s)") lines.append(f" Total entities created: {self.total_entities:,}") for model, count in sorted(self.entity_counts.items()): lines.append(f" {model}: {count:,}") return "\n".join(lines) def create_benchmark_spec( name: str, entities: list[EntitySpec], tenant_count: int = 1, workspace_count: int = 0, seed: int | None = None, locale: str = "fa_IR", ) -> BenchmarkSpec: """Convenience function to create a BenchmarkSpec.""" return BenchmarkSpec( name=name, entities=entities, tenant_count=tenant_count, workspace_count=workspace_count, seed=seed, locale=locale, )