from __future__ import annotations from typing import Any, Final, Iterable, Mapping TPCC_1K_TEMPLATE: Final[str] = "bench-tpcc1k" TPCC_5K_TEMPLATE: Final[str] = "bench-tpcc5k" TPCC_X_1K_TEMPLATE: Final[str] = "bench-tpcc-x-1k" TPCC_X_5K_TEMPLATE: Final[str] = "bench-tpcc-x-5k" TPCC_LEGACY_TEMPLATE: Final[str] = "bench-tpcc" SYSBENCH_TEMPLATE: Final[str] = "bench-sysbench" SYSBENCH_X_TEMPLATE: Final[str] = "bench-sysbench-x" CH_TEMPLATE: Final[str] = "bench-ch" CH_X_TEMPLATE: Final[str] = "bench-ch-x" LONGTERM_TPCC_TEMPLATE: Final[str] = "longterm-tpcc" LONGTERM_TPCC_X_TEMPLATE: Final[str] = "longterm-tpcc-x" SUPPORTED_WORKLOAD_TEMPLATES: Final[list[str]] = [ TPCC_1K_TEMPLATE, TPCC_5K_TEMPLATE, TPCC_X_1K_TEMPLATE, TPCC_X_5K_TEMPLATE, SYSBENCH_TEMPLATE, SYSBENCH_X_TEMPLATE, CH_TEMPLATE, CH_X_TEMPLATE, LONGTERM_TPCC_TEMPLATE, LONGTERM_TPCC_X_TEMPLATE, ] PREMIUM_WORKLOAD_TEMPLATES: Final[list[str]] = [ TPCC_X_1K_TEMPLATE, TPCC_X_5K_TEMPLATE, SYSBENCH_X_TEMPLATE, CH_X_TEMPLATE, LONGTERM_TPCC_X_TEMPLATE, ] DEDICATED_WORKLOAD_TEMPLATES: Final[list[str]] = [ template for template in SUPPORTED_WORKLOAD_TEMPLATES if template not in PREMIUM_WORKLOAD_TEMPLATES ] _ALIASES: Final[dict[str, str]] = { # Backward compatibility: legacy bench-tpcc means tpcc 1k. TPCC_LEGACY_TEMPLATE: TPCC_1K_TEMPLATE, # Backward compatibility: old no-qps names. "bench-tpcc-x-1k-no-qps": TPCC_X_1K_TEMPLATE, "bench-sysbench-x-no-qps": SYSBENCH_X_TEMPLATE, "bench-ch-x-no-qps": CH_X_TEMPLATE, } _SHORTHANDS: Final[dict[str, str]] = { # Backward compatibility: keep common shorthands stable. "tpcc": TPCC_1K_TEMPLATE, "tpcc1k": TPCC_1K_TEMPLATE, "tpcc5k": TPCC_5K_TEMPLATE, "sysbench": SYSBENCH_TEMPLATE, "ch": CH_TEMPLATE, "longterm-tpcc": LONGTERM_TPCC_TEMPLATE, } _YAML_BY_TEMPLATE: Final[dict[str, str]] = { TPCC_1K_TEMPLATE: "config/perf_tpcc1k.yaml", TPCC_5K_TEMPLATE: "config/perf_tpcc5k.yaml", TPCC_X_1K_TEMPLATE: "config/perf_x_tpcc1k.yaml", TPCC_X_5K_TEMPLATE: "config/perf_x_tpcc5k.yaml", SYSBENCH_TEMPLATE: "config/perf_sysbench.yaml", SYSBENCH_X_TEMPLATE: "config/perf_x_sysbench.yaml", CH_TEMPLATE: "config/perf_ch.yaml", CH_X_TEMPLATE: "config/perf_x_ch.yaml", LONGTERM_TPCC_TEMPLATE: "config/stability_tpcc.yaml", LONGTERM_TPCC_X_TEMPLATE: "config/perf_x_stability_tpcc.yaml", } # Note: These IDs come from the exported workflow YAML files under `config/`. _TCMS_WORKFLOW_ID_BY_TEMPLATE: Final[dict[str, int]] = { TPCC_1K_TEMPLATE: 210001, TPCC_5K_TEMPLATE: 150119, TPCC_X_1K_TEMPLATE: 210006, SYSBENCH_TEMPLATE: 210003, SYSBENCH_X_TEMPLATE: 210007, CH_TEMPLATE: 150094, CH_X_TEMPLATE: 210008, TPCC_X_5K_TEMPLATE: 210009, LONGTERM_TPCC_TEMPLATE: 210004, LONGTERM_TPCC_X_TEMPLATE: 210010, } def canonicalize_workload_template(template: str) -> str: """Return the canonical workload template name. This function is intentionally strict: it does not provide a default. Call `canonicalize_workload_template_or_default()` where a fallback is desired. """ if not template: raise ValueError("workload template is required") return _ALIASES.get(template, template) def canonicalize_workload_template_or_default( template: str | None, *, default: str = SYSBENCH_TEMPLATE ) -> str: if not template: return default return canonicalize_workload_template(template) def get_workload_templates_for_plan(plan_type: str) -> list[str]: """Return workload templates supported by a cluster plan. Template names are not structurally uniform (for example, ``bench-tpcc-x-1k`` versus ``bench-sysbench-x``), so plan membership must not be inferred with substring matching. """ normalized = (plan_type or "").strip().lower() if normalized == "dedicated": return list(DEDICATED_WORKLOAD_TEMPLATES) if normalized == "premium": return list(PREMIUM_WORKLOAD_TEMPLATES) raise ValueError(f"Unknown plan type: {plan_type!r}") def build_premium_workload_config( workload_config: Mapping[str, Any] | None, *, min_rcu: int, max_rcu: int, ) -> dict[str, Any]: """Add Premium capacity settings without discarding workload overrides.""" result = dict(workload_config or {}) result["plan_type"] = "premium" result["rcu"] = {"min": int(min_rcu), "max": int(max_rcu)} return result def is_premium_workload_config(workload_config: object) -> bool: """Recognize both the new plan marker and legacy persisted records.""" if not isinstance(workload_config, dict): return False return ( str(workload_config.get("plan_type", "")).lower() == "premium" or workload_config.get("template_type") == "premium" ) def iter_equivalent_templates(template: str | None) -> Iterable[str]: """Yield templates that should be treated as equivalent for filtering.""" canonical = canonicalize_workload_template_or_default(template) if canonical == TPCC_1K_TEMPLATE: # Include legacy value stored in DB. yield TPCC_LEGACY_TEMPLATE # Also yield old no-qps name for backward compatibility with DB records. for old_name, new_name in _ALIASES.items(): if new_name == canonical and old_name != TPCC_LEGACY_TEMPLATE: yield old_name yield canonical def resolve_workflow_yaml_path(template: str | None) -> str: canonical = canonicalize_workload_template_or_default(template) yaml_path = _YAML_BY_TEMPLATE.get(canonical) if not yaml_path: raise ValueError(f"Unknown workload template: {template!r}") return yaml_path def resolve_tcms_workflow_id(template: str | None) -> int: canonical = canonicalize_workload_template_or_default(template) workflow_id = _TCMS_WORKFLOW_ID_BY_TEMPLATE.get(canonical) if workflow_id is None: raise ValueError(f"Unknown workload template: {template!r}") return workflow_id def resolve_workload_template_from_filter_value(value: str | None) -> str | None: """Resolve a workload template from a filter value. Accepts canonical template names (bench-*) as well as backward-compatible shorthands (tpcc/sysbench/ch, etc.). Returns None if the value is not a recognized workload template filter. """ if not value: return None raw = value.strip() if not raw: return None canonical = canonicalize_workload_template(raw) canonical_lower = canonical.lower() if canonical_lower in SUPPORTED_WORKLOAD_TEMPLATES: return canonical_lower return _SHORTHANDS.get(canonical_lower)