rosetta-sql 1.0.0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- benchmark/generate_csv_data.py +83 -0
- benchmark/import_data.py +168 -0
- rosetta/__init__.py +3 -0
- rosetta/__main__.py +8 -0
- rosetta/benchmark.py +1678 -0
- rosetta/buglist.py +108 -0
- rosetta/cli/__init__.py +11 -0
- rosetta/cli/config_cmd.py +243 -0
- rosetta/cli/exec.py +219 -0
- rosetta/cli/interactive_cmd.py +124 -0
- rosetta/cli/list_cmd.py +215 -0
- rosetta/cli/main.py +617 -0
- rosetta/cli/output.py +545 -0
- rosetta/cli/result.py +61 -0
- rosetta/cli/result_cmd.py +247 -0
- rosetta/cli/run.py +625 -0
- rosetta/cli/status.py +161 -0
- rosetta/comparator.py +205 -0
- rosetta/config.py +139 -0
- rosetta/executor.py +403 -0
- rosetta/flamegraph.py +630 -0
- rosetta/interactive.py +1790 -0
- rosetta/models.py +197 -0
- rosetta/parser.py +308 -0
- rosetta/reporter/__init__.py +1 -0
- rosetta/reporter/bench_html.py +1457 -0
- rosetta/reporter/bench_text.py +162 -0
- rosetta/reporter/history.py +1686 -0
- rosetta/reporter/html.py +644 -0
- rosetta/reporter/text.py +110 -0
- rosetta/runner.py +3089 -0
- rosetta/ui.py +736 -0
- rosetta/whitelist.py +161 -0
- rosetta_sql-1.0.0.dist-info/LICENSE +21 -0
- rosetta_sql-1.0.0.dist-info/METADATA +379 -0
- rosetta_sql-1.0.0.dist-info/RECORD +42 -0
- rosetta_sql-1.0.0.dist-info/WHEEL +5 -0
- rosetta_sql-1.0.0.dist-info/entry_points.txt +2 -0
- rosetta_sql-1.0.0.dist-info/top_level.txt +4 -0
- skills/rosetta/scripts/install_rosetta.py +469 -0
- skills/rosetta/scripts/rosetta_wrapper.py +377 -0
- tests/test_cli.py +749 -0
rosetta/runner.py
ADDED
|
@@ -0,0 +1,3089 @@
|
|
|
1
|
+
"""Command-line interface for Rosetta."""
|
|
2
|
+
|
|
3
|
+
import argparse
|
|
4
|
+
import concurrent.futures
|
|
5
|
+
import http.server
|
|
6
|
+
import logging
|
|
7
|
+
import os
|
|
8
|
+
import shutil
|
|
9
|
+
import socket
|
|
10
|
+
import subprocess
|
|
11
|
+
import sys
|
|
12
|
+
import threading
|
|
13
|
+
import time as _time
|
|
14
|
+
from pathlib import Path
|
|
15
|
+
from typing import Dict, List, Optional, Tuple
|
|
16
|
+
|
|
17
|
+
from .comparator import compare_outputs
|
|
18
|
+
from .config import (DEFAULT_TEST_DB, filter_configs, generate_sample_config,
|
|
19
|
+
load_config)
|
|
20
|
+
from .executor import run_on_dbms
|
|
21
|
+
from .models import CompareResult, DBMSConfig, Statement, StmtType, WorkloadMode
|
|
22
|
+
from .parser import TestFileParser
|
|
23
|
+
from .reporter.html import write_html_report
|
|
24
|
+
from .reporter.history import generate_index_html
|
|
25
|
+
from .reporter.text import write_diff_file, write_text_report
|
|
26
|
+
from .ui import (ExecutionProgress, RichLogHandler, console, flush_all,
|
|
27
|
+
print_banner, print_error, print_info, print_phase,
|
|
28
|
+
print_report_file, print_server_info, print_success,
|
|
29
|
+
print_summary, print_warning)
|
|
30
|
+
|
|
31
|
+
log = logging.getLogger("rosetta")
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
class _SilentHTTPServer(http.server.HTTPServer):
|
|
35
|
+
"""HTTPServer that silently handles connection errors."""
|
|
36
|
+
|
|
37
|
+
def handle_error(self, request, client_address):
|
|
38
|
+
"""Silently ignore connection reset/broken pipe errors."""
|
|
39
|
+
pass
|
|
40
|
+
|
|
41
|
+
|
|
42
|
+
class _NoCacheHandler(http.server.SimpleHTTPRequestHandler):
|
|
43
|
+
"""HTTP handler that disables caching for all responses."""
|
|
44
|
+
|
|
45
|
+
def log_message(self, format, *args): # noqa: A002
|
|
46
|
+
pass # Suppress request logs
|
|
47
|
+
|
|
48
|
+
def end_headers(self): # noqa: N802
|
|
49
|
+
self.send_header("Cache-Control", "no-store, no-cache, must-revalidate")
|
|
50
|
+
self.send_header("Pragma", "no-cache")
|
|
51
|
+
self.send_header("Expires", "0")
|
|
52
|
+
super().end_headers()
|
|
53
|
+
|
|
54
|
+
|
|
55
|
+
def _tty_write(data: str):
|
|
56
|
+
"""Write escape codes directly to /dev/tty.
|
|
57
|
+
|
|
58
|
+
In environments where sys.stdout is a pipe (e.g. IDE terminals),
|
|
59
|
+
prompt_toolkit writes to /dev/tty but sys.stdout does not reach the
|
|
60
|
+
terminal. This helper ensures escape sequences actually reach the
|
|
61
|
+
terminal device.
|
|
62
|
+
"""
|
|
63
|
+
try:
|
|
64
|
+
fd = os.open("/dev/tty", os.O_WRONLY)
|
|
65
|
+
try:
|
|
66
|
+
os.write(fd, data.encode())
|
|
67
|
+
finally:
|
|
68
|
+
os.close(fd)
|
|
69
|
+
except OSError:
|
|
70
|
+
sys.stdout.write(data)
|
|
71
|
+
sys.stdout.flush()
|
|
72
|
+
|
|
73
|
+
|
|
74
|
+
class RosettaRunner:
|
|
75
|
+
"""Orchestrates parsing, execution, comparison, and reporting."""
|
|
76
|
+
|
|
77
|
+
def __init__(self, test_file: str, configs: List[DBMSConfig],
|
|
78
|
+
output_dir: str, database: str = DEFAULT_TEST_DB,
|
|
79
|
+
baseline: Optional[str] = None,
|
|
80
|
+
skip_explain: bool = False,
|
|
81
|
+
skip_analyze: bool = False,
|
|
82
|
+
skip_show_create: bool = False,
|
|
83
|
+
output_format: str = "all",
|
|
84
|
+
whitelist=None,
|
|
85
|
+
buglist=None):
|
|
86
|
+
self.test_file = test_file
|
|
87
|
+
self.configs = configs
|
|
88
|
+
self.output_dir = output_dir
|
|
89
|
+
self.database = database
|
|
90
|
+
self.baseline = baseline
|
|
91
|
+
self.skip_explain_global = skip_explain
|
|
92
|
+
self.skip_analyze_global = skip_analyze
|
|
93
|
+
self.skip_show_create_global = skip_show_create
|
|
94
|
+
self.output_format = output_format
|
|
95
|
+
self.whitelist = whitelist
|
|
96
|
+
self.buglist = buglist
|
|
97
|
+
self.results: Dict[str, List[str]] = {}
|
|
98
|
+
self.failed_connections: set = set()
|
|
99
|
+
|
|
100
|
+
def _should_skip_stmt_global(self, stmt: Statement) -> bool:
|
|
101
|
+
"""Check if a statement should be skipped globally."""
|
|
102
|
+
if stmt.stmt_type != StmtType.SQL:
|
|
103
|
+
return False
|
|
104
|
+
|
|
105
|
+
sql_upper = stmt.text.strip().upper()
|
|
106
|
+
|
|
107
|
+
if self.skip_explain_global and sql_upper.startswith("EXPLAIN"):
|
|
108
|
+
return True
|
|
109
|
+
if self.skip_analyze_global and sql_upper.startswith("ANALYZE"):
|
|
110
|
+
return True
|
|
111
|
+
if (self.skip_show_create_global
|
|
112
|
+
and sql_upper.startswith("SHOW CREATE")):
|
|
113
|
+
return True
|
|
114
|
+
|
|
115
|
+
for c in self.configs:
|
|
116
|
+
if c.skip_explain and sql_upper.startswith("EXPLAIN"):
|
|
117
|
+
return True
|
|
118
|
+
if c.skip_analyze and sql_upper.startswith("ANALYZE"):
|
|
119
|
+
return True
|
|
120
|
+
if c.skip_show_create and sql_upper.startswith("SHOW CREATE"):
|
|
121
|
+
return True
|
|
122
|
+
|
|
123
|
+
return False
|
|
124
|
+
|
|
125
|
+
def _order_configs(self) -> List[DBMSConfig]:
|
|
126
|
+
"""Order configs: baseline first, then 'mysql', then others."""
|
|
127
|
+
baseline_cfg = []
|
|
128
|
+
mysql_cfg = []
|
|
129
|
+
others = []
|
|
130
|
+
for c in self.configs:
|
|
131
|
+
if self.baseline and c.name == self.baseline:
|
|
132
|
+
baseline_cfg.append(c)
|
|
133
|
+
elif c.name == "mysql":
|
|
134
|
+
mysql_cfg.append(c)
|
|
135
|
+
else:
|
|
136
|
+
others.append(c)
|
|
137
|
+
return baseline_cfg + mysql_cfg + others
|
|
138
|
+
|
|
139
|
+
def _compare_all(self) -> Dict[str, CompareResult]:
|
|
140
|
+
"""Compare results across all DBMS pairs."""
|
|
141
|
+
names = [n for n in self.results if n not in self.failed_connections]
|
|
142
|
+
comparisons = {}
|
|
143
|
+
|
|
144
|
+
if self.baseline and self.baseline in self.results:
|
|
145
|
+
for name in names:
|
|
146
|
+
if name == self.baseline:
|
|
147
|
+
continue
|
|
148
|
+
key = f"{self.baseline}_vs_{name}"
|
|
149
|
+
comparisons[key] = compare_outputs(
|
|
150
|
+
self.results[self.baseline],
|
|
151
|
+
self.results[name],
|
|
152
|
+
self.baseline, name,
|
|
153
|
+
baseline_name=self.baseline,
|
|
154
|
+
whitelist=self.whitelist,
|
|
155
|
+
buglist=self.buglist,
|
|
156
|
+
)
|
|
157
|
+
else:
|
|
158
|
+
for i in range(len(names)):
|
|
159
|
+
for j in range(i + 1, len(names)):
|
|
160
|
+
key = f"{names[i]}_vs_{names[j]}"
|
|
161
|
+
comparisons[key] = compare_outputs(
|
|
162
|
+
self.results[names[i]],
|
|
163
|
+
self.results[names[j]],
|
|
164
|
+
names[i], names[j],
|
|
165
|
+
whitelist=self.whitelist,
|
|
166
|
+
buglist=self.buglist,
|
|
167
|
+
)
|
|
168
|
+
|
|
169
|
+
return comparisons
|
|
170
|
+
|
|
171
|
+
def run(self) -> Dict[str, CompareResult]:
|
|
172
|
+
"""Execute the full pipeline: parse, execute, compare, report."""
|
|
173
|
+
os.makedirs(self.output_dir, exist_ok=True)
|
|
174
|
+
|
|
175
|
+
# Parse
|
|
176
|
+
print_phase("Parse", self.test_file)
|
|
177
|
+
parser = TestFileParser(self.test_file)
|
|
178
|
+
statements = parser.parse()
|
|
179
|
+
print_success(f"Parsed {len(statements)} statements")
|
|
180
|
+
|
|
181
|
+
# Execute on each DBMS (in parallel)
|
|
182
|
+
configs = self._order_configs()
|
|
183
|
+
print_phase("Execute",
|
|
184
|
+
f"{len(configs)} DBMS targets (parallel)")
|
|
185
|
+
|
|
186
|
+
def _run_single(config):
|
|
187
|
+
"""Execute on one DBMS; returns (config.name, output_lines | None)."""
|
|
188
|
+
with ExecutionProgress(config.name,
|
|
189
|
+
total=len(statements)) as prog:
|
|
190
|
+
def _on_progress(error=False, _p=prog):
|
|
191
|
+
_p.advance(error=error)
|
|
192
|
+
|
|
193
|
+
def _on_connect(name, ok, msg, _p=prog):
|
|
194
|
+
if ok:
|
|
195
|
+
_p.set_status(f"[green]{msg}[/green]")
|
|
196
|
+
else:
|
|
197
|
+
_p.set_status(f"[red]{msg}[/red]")
|
|
198
|
+
|
|
199
|
+
def _on_done(name, executed, errors, _p=prog):
|
|
200
|
+
if errors:
|
|
201
|
+
_p.set_status(
|
|
202
|
+
f"[yellow]{executed} done, {errors} err[/yellow]")
|
|
203
|
+
else:
|
|
204
|
+
_p.set_status(f"[green]{executed} done[/green]")
|
|
205
|
+
|
|
206
|
+
output = run_on_dbms(
|
|
207
|
+
config, statements, self.database,
|
|
208
|
+
should_skip_fn=self._should_skip_stmt_global,
|
|
209
|
+
on_connect=_on_connect,
|
|
210
|
+
on_progress=_on_progress,
|
|
211
|
+
on_done=_on_done,
|
|
212
|
+
)
|
|
213
|
+
|
|
214
|
+
# Progress bar is now gone (transient); write static line
|
|
215
|
+
prog.write_summary_to_buffer()
|
|
216
|
+
return config.name, output
|
|
217
|
+
|
|
218
|
+
with concurrent.futures.ThreadPoolExecutor(
|
|
219
|
+
max_workers=len(configs)) as pool:
|
|
220
|
+
futures = {pool.submit(_run_single, c): c for c in configs}
|
|
221
|
+
for fut in concurrent.futures.as_completed(futures):
|
|
222
|
+
name, output = fut.result()
|
|
223
|
+
if output is None:
|
|
224
|
+
self.failed_connections.add(name)
|
|
225
|
+
else:
|
|
226
|
+
self.results[name] = output
|
|
227
|
+
|
|
228
|
+
# Write result files
|
|
229
|
+
print_phase("Reports")
|
|
230
|
+
test_name = Path(self.test_file).stem
|
|
231
|
+
for name, lines in self.results.items():
|
|
232
|
+
result_path = os.path.join(
|
|
233
|
+
self.output_dir, f"{test_name}.{name}.result"
|
|
234
|
+
)
|
|
235
|
+
with open(result_path, "w", encoding="utf-8") as f:
|
|
236
|
+
f.write("\n".join(lines) + "\n")
|
|
237
|
+
print_report_file(result_path, label="result")
|
|
238
|
+
|
|
239
|
+
# Compare
|
|
240
|
+
comparisons = self._compare_all()
|
|
241
|
+
|
|
242
|
+
# Generate reports
|
|
243
|
+
self._generate_reports(test_name, comparisons)
|
|
244
|
+
|
|
245
|
+
return comparisons
|
|
246
|
+
|
|
247
|
+
def run_diff_only(self) -> Dict[str, CompareResult]:
|
|
248
|
+
"""Re-generate reports from existing .result files (no execution)."""
|
|
249
|
+
os.makedirs(self.output_dir, exist_ok=True)
|
|
250
|
+
test_name = Path(self.test_file).stem
|
|
251
|
+
|
|
252
|
+
print_phase("Load Results", "(diff-only mode)")
|
|
253
|
+
|
|
254
|
+
# Load existing .result files
|
|
255
|
+
for config in self.configs:
|
|
256
|
+
result_path = os.path.join(
|
|
257
|
+
self.output_dir, f"{test_name}.{config.name}.result"
|
|
258
|
+
)
|
|
259
|
+
if os.path.isfile(result_path):
|
|
260
|
+
with open(result_path, "r", encoding="utf-8") as f:
|
|
261
|
+
self.results[config.name] = [
|
|
262
|
+
line.rstrip("\n") for line in f
|
|
263
|
+
]
|
|
264
|
+
print_success(f"Loaded: {result_path}")
|
|
265
|
+
else:
|
|
266
|
+
print_warning(f"Not found: {result_path}")
|
|
267
|
+
|
|
268
|
+
if len(self.results) < 2:
|
|
269
|
+
print_error("Need at least 2 result files for comparison")
|
|
270
|
+
return {}
|
|
271
|
+
|
|
272
|
+
comparisons = self._compare_all()
|
|
273
|
+
|
|
274
|
+
print_phase("Reports")
|
|
275
|
+
self._generate_reports(test_name, comparisons)
|
|
276
|
+
return comparisons
|
|
277
|
+
|
|
278
|
+
def _generate_reports(self, test_name: str,
|
|
279
|
+
comparisons: Dict[str, CompareResult]):
|
|
280
|
+
"""Generate output reports based on format setting."""
|
|
281
|
+
fmt = self.output_format
|
|
282
|
+
|
|
283
|
+
if fmt in ("text", "all"):
|
|
284
|
+
report_path = os.path.join(
|
|
285
|
+
self.output_dir, f"{test_name}.report.txt"
|
|
286
|
+
)
|
|
287
|
+
write_text_report(report_path, self.test_file, comparisons)
|
|
288
|
+
print_report_file(report_path, label="text")
|
|
289
|
+
|
|
290
|
+
diff_path = os.path.join(
|
|
291
|
+
self.output_dir, f"{test_name}.diff"
|
|
292
|
+
)
|
|
293
|
+
write_diff_file(diff_path, comparisons)
|
|
294
|
+
print_report_file(diff_path, label="diff")
|
|
295
|
+
|
|
296
|
+
if fmt in ("html", "all"):
|
|
297
|
+
html_path = os.path.join(
|
|
298
|
+
self.output_dir, f"{test_name}.html"
|
|
299
|
+
)
|
|
300
|
+
write_html_report(
|
|
301
|
+
html_path, self.test_file, comparisons,
|
|
302
|
+
baseline=self.baseline or "",
|
|
303
|
+
)
|
|
304
|
+
print_report_file(html_path, label="html")
|
|
305
|
+
|
|
306
|
+
|
|
307
|
+
def parse_args(argv=None):
|
|
308
|
+
"""Parse command-line arguments."""
|
|
309
|
+
p = argparse.ArgumentParser(
|
|
310
|
+
prog="rosetta",
|
|
311
|
+
description=(
|
|
312
|
+
"Rosetta — Cross-DBMS SQL testing & benchmarking toolkit.\n"
|
|
313
|
+
"\n"
|
|
314
|
+
"Three operating modes:\n"
|
|
315
|
+
" MTR Run .test files against multiple databases "
|
|
316
|
+
"and diff results\n"
|
|
317
|
+
" Benchmark Compare query performance across databases "
|
|
318
|
+
"with latency/QPS reports\n"
|
|
319
|
+
" Playground Launch an interactive SQL playground "
|
|
320
|
+
"in the browser\n"
|
|
321
|
+
"\n"
|
|
322
|
+
"Use --interactive (-i) to enter a REPL that lets you "
|
|
323
|
+
"switch between modes.\n"
|
|
324
|
+
"Without -i, run a single MTR test (--test) or benchmark "
|
|
325
|
+
"(--benchmark)."
|
|
326
|
+
),
|
|
327
|
+
formatter_class=argparse.RawDescriptionHelpFormatter,
|
|
328
|
+
epilog="""\
|
|
329
|
+
Examples:
|
|
330
|
+
# ── Setup ──────────────────────────────────────────────────
|
|
331
|
+
rosetta --gen-config dbms_config.json Generate sample config
|
|
332
|
+
|
|
333
|
+
# ── MTR (consistency test) ─────────────────────────────────
|
|
334
|
+
rosetta -t path/to/test.test --dbms tdsql,mysql
|
|
335
|
+
rosetta -t path/to/test.test --dbms tdsql,mysql,tidb -b tdsql
|
|
336
|
+
rosetta -t path/to/test.test --diff-only Re-diff without execution
|
|
337
|
+
rosetta -t path/to/test.test --parse-only Debug: show parsed stmts
|
|
338
|
+
|
|
339
|
+
# ── Benchmark ──────────────────────────────────────────────
|
|
340
|
+
rosetta --benchmark --bench-file bench.json --dbms tdsql,mysql
|
|
341
|
+
rosetta --benchmark --template oltp_read_write --iterations 200
|
|
342
|
+
rosetta --benchmark --bench-file bench.json --repeat 5
|
|
343
|
+
rosetta --benchmark --bench-file bench.json --concurrency 16 --duration 60
|
|
344
|
+
rosetta --benchmark --list-templates Show built-in templates
|
|
345
|
+
|
|
346
|
+
# ── Interactive / Playground ───────────────────────────────
|
|
347
|
+
rosetta -i --dbms tdsql,mysql -s Choose mode at startup
|
|
348
|
+
rosetta -i --benchmark --dbms tdsql,mysql Go straight to Benchmark
|
|
349
|
+
rosetta -i --dbms tdsql,mysql --port 8080 Custom server port
|
|
350
|
+
|
|
351
|
+
# ── Profiling ──────────────────────────────────────────────
|
|
352
|
+
rosetta --benchmark --bench-file b.json --profile --perf-freq 199
|
|
353
|
+
rosetta --benchmark --bench-file b.json --no-profile
|
|
354
|
+
""",
|
|
355
|
+
)
|
|
356
|
+
|
|
357
|
+
# -- Global options -------------------------------------------------------
|
|
358
|
+
general = p.add_argument_group(
|
|
359
|
+
"General", "Options shared across all modes")
|
|
360
|
+
general.add_argument(
|
|
361
|
+
"--config", "-c", default="dbms_config.json",
|
|
362
|
+
help="Path to DBMS config JSON (default: dbms_config.json)")
|
|
363
|
+
general.add_argument(
|
|
364
|
+
"--dbms",
|
|
365
|
+
help="DBMS targets, comma-separated (e.g. tdsql,mysql,tidb). "
|
|
366
|
+
"Omit to use 'enabled' flag in config")
|
|
367
|
+
general.add_argument(
|
|
368
|
+
"--database", "-d", default=DEFAULT_TEST_DB,
|
|
369
|
+
help=f"Test database name (default: {DEFAULT_TEST_DB})")
|
|
370
|
+
general.add_argument(
|
|
371
|
+
"--output-dir", "-o", default="results",
|
|
372
|
+
help="Output directory for reports (default: results)")
|
|
373
|
+
general.add_argument(
|
|
374
|
+
"--verbose", "-v", action="store_true",
|
|
375
|
+
help="Enable verbose / debug logging")
|
|
376
|
+
general.add_argument(
|
|
377
|
+
"--gen-config",
|
|
378
|
+
help="Generate sample config at the given path and exit")
|
|
379
|
+
|
|
380
|
+
# -- Interactive / server -------------------------------------------------
|
|
381
|
+
ui = p.add_argument_group(
|
|
382
|
+
"Interactive & Server",
|
|
383
|
+
"Enter a REPL or serve HTML reports in the browser")
|
|
384
|
+
ui.add_argument(
|
|
385
|
+
"--interactive", "-i", action="store_true",
|
|
386
|
+
help="Enter interactive mode — choose MTR / Benchmark / "
|
|
387
|
+
"Playground, then run tasks in a loop")
|
|
388
|
+
ui.add_argument(
|
|
389
|
+
"--serve", "-s", action="store_true",
|
|
390
|
+
help="Start a local HTTP server to view HTML reports "
|
|
391
|
+
"after execution")
|
|
392
|
+
ui.add_argument(
|
|
393
|
+
"--port", "-p", type=int, default=19527,
|
|
394
|
+
help="HTTP server port (default: 19527)")
|
|
395
|
+
|
|
396
|
+
# -- MTR options ----------------------------------------------------------
|
|
397
|
+
mtr = p.add_argument_group(
|
|
398
|
+
"MTR (Consistency Test)",
|
|
399
|
+
"Run .test files and compare results across databases")
|
|
400
|
+
mtr.add_argument(
|
|
401
|
+
"--test", "-t",
|
|
402
|
+
help="Path to .test file")
|
|
403
|
+
mtr.add_argument(
|
|
404
|
+
"--baseline", "-b", default="tdsql",
|
|
405
|
+
help="Baseline DBMS name for diff (default: tdsql)")
|
|
406
|
+
mtr.add_argument(
|
|
407
|
+
"--format", "-f", default="all",
|
|
408
|
+
choices=["text", "html", "all"],
|
|
409
|
+
help="Output format (default: all)")
|
|
410
|
+
mtr.add_argument(
|
|
411
|
+
"--skip-explain", action="store_true", default=True,
|
|
412
|
+
help="Skip EXPLAIN statements (default: on)")
|
|
413
|
+
mtr.add_argument(
|
|
414
|
+
"--skip-analyze", action="store_true",
|
|
415
|
+
help="Skip ANALYZE TABLE statements")
|
|
416
|
+
mtr.add_argument(
|
|
417
|
+
"--skip-show-create", action="store_true",
|
|
418
|
+
help="Skip SHOW CREATE TABLE statements")
|
|
419
|
+
mtr.add_argument(
|
|
420
|
+
"--parse-only", action="store_true",
|
|
421
|
+
help="Only parse .test file and print statements (no execution)")
|
|
422
|
+
mtr.add_argument(
|
|
423
|
+
"--diff-only", action="store_true",
|
|
424
|
+
help="Re-generate reports from existing .result files "
|
|
425
|
+
"(no DB execution)")
|
|
426
|
+
|
|
427
|
+
# -- Benchmark options ----------------------------------------------------
|
|
428
|
+
bench = p.add_argument_group(
|
|
429
|
+
"Benchmark",
|
|
430
|
+
"Compare query performance across databases with "
|
|
431
|
+
"latency / QPS reports")
|
|
432
|
+
bench.add_argument(
|
|
433
|
+
"--benchmark", action="store_true",
|
|
434
|
+
help="Enable benchmark mode")
|
|
435
|
+
bench.add_argument(
|
|
436
|
+
"--bench-file",
|
|
437
|
+
help="Benchmark definition file (.json or .sql)")
|
|
438
|
+
bench.add_argument(
|
|
439
|
+
"--template",
|
|
440
|
+
help="Use a built-in template "
|
|
441
|
+
"(e.g. oltp_read_write, oltp_read_only)")
|
|
442
|
+
bench.add_argument(
|
|
443
|
+
"--list-templates", action="store_true",
|
|
444
|
+
help="List built-in benchmark templates and exit")
|
|
445
|
+
bench.add_argument(
|
|
446
|
+
"--iterations", type=int, default=100,
|
|
447
|
+
help="Iterations per query — serial mode (default: 100)")
|
|
448
|
+
bench.add_argument(
|
|
449
|
+
"--warmup", type=int, default=5,
|
|
450
|
+
help="Warmup iterations per query (default: 5)")
|
|
451
|
+
bench.add_argument(
|
|
452
|
+
"--concurrency", type=int, default=0,
|
|
453
|
+
help="Concurrent threads; 0 = serial, >0 = concurrent "
|
|
454
|
+
"(default: 0)")
|
|
455
|
+
bench.add_argument(
|
|
456
|
+
"--duration", type=float, default=30.0,
|
|
457
|
+
help="Duration in seconds — concurrent mode (default: 30)")
|
|
458
|
+
bench.add_argument(
|
|
459
|
+
"--ramp-up", type=float, default=0.0,
|
|
460
|
+
help="Ramp-up seconds — concurrent mode (default: 0)")
|
|
461
|
+
bench.add_argument(
|
|
462
|
+
"--query-timeout", type=int, default=5,
|
|
463
|
+
help="Query timeout in seconds; slow queries will be logged as outliers "
|
|
464
|
+
"(default: 5, 0 to disable)")
|
|
465
|
+
bench.add_argument(
|
|
466
|
+
"--flamegraph-min-ms", type=int, default=1000,
|
|
467
|
+
help="Minimum total duration (ms) to show flamegraph in serial mode "
|
|
468
|
+
"(default: 1000, 0 to always show)")
|
|
469
|
+
bench.add_argument(
|
|
470
|
+
"--bench-filter",
|
|
471
|
+
help="Run only queries matching these names "
|
|
472
|
+
"(comma-separated)")
|
|
473
|
+
bench.add_argument(
|
|
474
|
+
"--repeat", type=int, default=1,
|
|
475
|
+
help="Number of benchmark rounds; each round produces "
|
|
476
|
+
"a timestamped report (default: 1)")
|
|
477
|
+
bench.add_argument(
|
|
478
|
+
"--skip-setup", action="store_true", default=False,
|
|
479
|
+
help="Skip setup phase (reuse existing tables from previous run)")
|
|
480
|
+
bench.add_argument(
|
|
481
|
+
"--skip-teardown", action="store_true", default=False,
|
|
482
|
+
help="Skip teardown (keep tables for next run with --skip-setup)")
|
|
483
|
+
bench.add_argument(
|
|
484
|
+
"--no-parallel-dbms", dest="parallel_dbms",
|
|
485
|
+
action="store_false",
|
|
486
|
+
help="Run DBMS targets sequentially instead of in parallel")
|
|
487
|
+
bench.set_defaults(parallel_dbms=True)
|
|
488
|
+
|
|
489
|
+
# -- Profiling options ----------------------------------------------------
|
|
490
|
+
prof = p.add_argument_group(
|
|
491
|
+
"Profiling",
|
|
492
|
+
"CPU flame-graph capture via perf (benchmark mode)")
|
|
493
|
+
prof.add_argument(
|
|
494
|
+
"--profile", action="store_true", dest="profile",
|
|
495
|
+
default=True,
|
|
496
|
+
help="Enable flame-graph capture (default: on)")
|
|
497
|
+
prof.add_argument(
|
|
498
|
+
"--no-profile", action="store_false", dest="profile",
|
|
499
|
+
help="Disable flame-graph capture")
|
|
500
|
+
prof.add_argument(
|
|
501
|
+
"--perf-freq", type=int, default=99,
|
|
502
|
+
help="perf sampling frequency in Hz (default: 99)")
|
|
503
|
+
|
|
504
|
+
return p.parse_args(argv)
|
|
505
|
+
|
|
506
|
+
|
|
507
|
+
def main(argv=None):
|
|
508
|
+
"""Main entry point for the rosetta command."""
|
|
509
|
+
# Configure logging to use rich handler
|
|
510
|
+
rich_handler = RichLogHandler()
|
|
511
|
+
rich_handler.setLevel(logging.WARNING)
|
|
512
|
+
logging.basicConfig(
|
|
513
|
+
level=logging.WARNING,
|
|
514
|
+
handlers=[rich_handler],
|
|
515
|
+
)
|
|
516
|
+
|
|
517
|
+
args = parse_args(argv)
|
|
518
|
+
|
|
519
|
+
if args.verbose:
|
|
520
|
+
# In verbose mode, use standard logging for everything
|
|
521
|
+
logging.root.handlers.clear()
|
|
522
|
+
logging.basicConfig(
|
|
523
|
+
level=logging.DEBUG,
|
|
524
|
+
format="%(asctime)s [%(levelname)s] %(message)s",
|
|
525
|
+
datefmt="%Y-%m-%d %H:%M:%S",
|
|
526
|
+
)
|
|
527
|
+
|
|
528
|
+
print_banner()
|
|
529
|
+
|
|
530
|
+
# Generate sample config
|
|
531
|
+
if args.gen_config:
|
|
532
|
+
generate_sample_config(args.gen_config)
|
|
533
|
+
print_success(f"Config written: {args.gen_config}")
|
|
534
|
+
flush_all()
|
|
535
|
+
return 0
|
|
536
|
+
|
|
537
|
+
# List built-in benchmark templates
|
|
538
|
+
if args.list_templates:
|
|
539
|
+
from .benchmark import BenchmarkLoader
|
|
540
|
+
templates = BenchmarkLoader.list_builtin_templates()
|
|
541
|
+
console.print("[bold]Available built-in benchmark templates:[/bold]")
|
|
542
|
+
for t in templates:
|
|
543
|
+
console.print(f" [cyan]•[/cyan] {t}")
|
|
544
|
+
return 0
|
|
545
|
+
|
|
546
|
+
# Benchmark mode (non-interactive)
|
|
547
|
+
if args.benchmark and not args.interactive:
|
|
548
|
+
return _run_benchmark(args)
|
|
549
|
+
|
|
550
|
+
# Interactive mode — show mode selection (MTR / Benchmark)
|
|
551
|
+
# If --benchmark is also set, skips selection and goes directly to bench mode
|
|
552
|
+
if args.interactive:
|
|
553
|
+
return _enter_interactive(args)
|
|
554
|
+
|
|
555
|
+
if not args.test:
|
|
556
|
+
print_error("--test is required. Use --help for usage.")
|
|
557
|
+
flush_all()
|
|
558
|
+
return 1
|
|
559
|
+
|
|
560
|
+
if not os.path.isfile(args.test):
|
|
561
|
+
print_error(f"Test file not found: {args.test}")
|
|
562
|
+
flush_all()
|
|
563
|
+
return 1
|
|
564
|
+
|
|
565
|
+
# Parse-only mode
|
|
566
|
+
if args.parse_only:
|
|
567
|
+
flush_all()
|
|
568
|
+
parser = TestFileParser(args.test)
|
|
569
|
+
stmts = parser.parse()
|
|
570
|
+
for s in stmts:
|
|
571
|
+
tag = s.stmt_type.name
|
|
572
|
+
err = (f" [expect error: {s.expected_error}]"
|
|
573
|
+
if s.expected_error else "")
|
|
574
|
+
sort = " [sorted]" if s.sort_result else ""
|
|
575
|
+
print(f"L{s.line_no:4d} [{tag:5s}]{err}{sort}: "
|
|
576
|
+
f"{s.text[:100]}")
|
|
577
|
+
print(f"\nTotal: {len(stmts)} statements")
|
|
578
|
+
return 0
|
|
579
|
+
|
|
580
|
+
if not os.path.isfile(args.config):
|
|
581
|
+
print_error(f"Config file not found: {args.config}")
|
|
582
|
+
flush_all()
|
|
583
|
+
return 1
|
|
584
|
+
|
|
585
|
+
# Load and filter configs
|
|
586
|
+
all_configs = load_config(args.config)
|
|
587
|
+
if not all_configs:
|
|
588
|
+
print_error(f"No databases configured in {args.config}")
|
|
589
|
+
flush_all()
|
|
590
|
+
return 1
|
|
591
|
+
|
|
592
|
+
try:
|
|
593
|
+
configs = filter_configs(all_configs, args.dbms)
|
|
594
|
+
except ValueError as e:
|
|
595
|
+
print_error(str(e))
|
|
596
|
+
flush_all()
|
|
597
|
+
return 1
|
|
598
|
+
|
|
599
|
+
if not configs:
|
|
600
|
+
print_error("No databases selected for testing")
|
|
601
|
+
flush_all()
|
|
602
|
+
return 1
|
|
603
|
+
|
|
604
|
+
# Resolve output_dir to absolute path early so it does not depend on cwd
|
|
605
|
+
output_dir = os.path.abspath(args.output_dir)
|
|
606
|
+
|
|
607
|
+
# Create a timestamped sub-directory for this run
|
|
608
|
+
run_stamp = _time.strftime("%Y%m%d_%H%M%S")
|
|
609
|
+
test_name = Path(args.test).stem
|
|
610
|
+
run_dir = os.path.join(output_dir, f"{test_name}_{run_stamp}")
|
|
611
|
+
|
|
612
|
+
print_info("DBMS targets:",
|
|
613
|
+
", ".join(c.name for c in configs))
|
|
614
|
+
|
|
615
|
+
# Load whitelist from output directory
|
|
616
|
+
from .whitelist import Whitelist
|
|
617
|
+
whitelist = Whitelist(output_dir)
|
|
618
|
+
|
|
619
|
+
# Load buglist from output directory
|
|
620
|
+
from .buglist import Buglist
|
|
621
|
+
buglist = Buglist(output_dir)
|
|
622
|
+
|
|
623
|
+
# Run
|
|
624
|
+
runner = RosettaRunner(
|
|
625
|
+
test_file=args.test,
|
|
626
|
+
configs=configs,
|
|
627
|
+
output_dir=run_dir,
|
|
628
|
+
database=args.database,
|
|
629
|
+
baseline=args.baseline,
|
|
630
|
+
skip_explain=args.skip_explain,
|
|
631
|
+
skip_analyze=args.skip_analyze,
|
|
632
|
+
skip_show_create=args.skip_show_create,
|
|
633
|
+
output_format=args.format,
|
|
634
|
+
whitelist=whitelist,
|
|
635
|
+
buglist=buglist,
|
|
636
|
+
)
|
|
637
|
+
|
|
638
|
+
if args.diff_only:
|
|
639
|
+
# Copy .result files from latest run into the new run_dir
|
|
640
|
+
latest_link = os.path.join(output_dir, "latest")
|
|
641
|
+
source_dir = (os.path.realpath(latest_link)
|
|
642
|
+
if os.path.islink(latest_link) else None)
|
|
643
|
+
if source_dir and os.path.isdir(source_dir):
|
|
644
|
+
os.makedirs(run_dir, exist_ok=True)
|
|
645
|
+
for f in os.listdir(source_dir):
|
|
646
|
+
if f.endswith(".result"):
|
|
647
|
+
shutil.copy2(
|
|
648
|
+
os.path.join(source_dir, f),
|
|
649
|
+
os.path.join(run_dir, f))
|
|
650
|
+
comparisons = runner.run_diff_only()
|
|
651
|
+
else:
|
|
652
|
+
comparisons = runner.run()
|
|
653
|
+
|
|
654
|
+
if not comparisons:
|
|
655
|
+
flush_all()
|
|
656
|
+
return 1
|
|
657
|
+
|
|
658
|
+
# Update 'latest' symlink
|
|
659
|
+
latest_link = os.path.join(output_dir, "latest")
|
|
660
|
+
try:
|
|
661
|
+
if os.path.islink(latest_link):
|
|
662
|
+
os.remove(latest_link)
|
|
663
|
+
os.symlink(os.path.basename(run_dir), latest_link)
|
|
664
|
+
except OSError:
|
|
665
|
+
pass
|
|
666
|
+
|
|
667
|
+
# Generate history index page and whitelist/buglist pages
|
|
668
|
+
generate_index_html(output_dir)
|
|
669
|
+
from .reporter.history import generate_buglist_html, generate_whitelist_html
|
|
670
|
+
generate_whitelist_html(output_dir)
|
|
671
|
+
generate_buglist_html(output_dir)
|
|
672
|
+
|
|
673
|
+
# Print rich summary table
|
|
674
|
+
all_pass = print_summary(comparisons, runner.failed_connections)
|
|
675
|
+
|
|
676
|
+
# Flush everything as one big panel
|
|
677
|
+
flush_all()
|
|
678
|
+
|
|
679
|
+
# Serve HTML report if requested
|
|
680
|
+
if args.serve and args.format in ("html", "all"):
|
|
681
|
+
html_file = f"{test_name}.html"
|
|
682
|
+
html_path = os.path.join(run_dir, html_file)
|
|
683
|
+
|
|
684
|
+
if os.path.isfile(html_path):
|
|
685
|
+
# Serve from output_dir root so history is accessible
|
|
686
|
+
relative_html = os.path.join(
|
|
687
|
+
os.path.basename(run_dir), html_file)
|
|
688
|
+
_serve_report(output_dir, relative_html, args.port,
|
|
689
|
+
whitelist=whitelist, buglist=buglist,
|
|
690
|
+
configs=configs, database=args.database)
|
|
691
|
+
else:
|
|
692
|
+
console.print(f"[yellow]HTML report not found: {html_path}[/yellow]")
|
|
693
|
+
|
|
694
|
+
return 0 if (all_pass and not runner.failed_connections) else 1
|
|
695
|
+
|
|
696
|
+
|
|
697
|
+
def _run_benchmark(args) -> int:
|
|
698
|
+
"""Execute the benchmark pipeline (supports --repeat N)."""
|
|
699
|
+
from .benchmark import (BenchmarkLoader, run_benchmark,
|
|
700
|
+
BUILTIN_TEMPLATES)
|
|
701
|
+
from .models import BenchmarkConfig, WorkloadMode
|
|
702
|
+
from .reporter.bench_text import write_bench_text_report
|
|
703
|
+
from .reporter.bench_html import write_bench_html_report
|
|
704
|
+
from .ui import BenchProgress, print_bench_summary
|
|
705
|
+
|
|
706
|
+
# Load DBMS configs
|
|
707
|
+
if not os.path.isfile(args.config):
|
|
708
|
+
print_error(f"Config file not found: {args.config}")
|
|
709
|
+
flush_all()
|
|
710
|
+
return 1
|
|
711
|
+
|
|
712
|
+
all_configs = load_config(args.config)
|
|
713
|
+
if not all_configs:
|
|
714
|
+
print_error(f"No databases configured in {args.config}")
|
|
715
|
+
flush_all()
|
|
716
|
+
return 1
|
|
717
|
+
|
|
718
|
+
try:
|
|
719
|
+
configs = filter_configs(all_configs, args.dbms)
|
|
720
|
+
except ValueError as e:
|
|
721
|
+
print_error(str(e))
|
|
722
|
+
flush_all()
|
|
723
|
+
return 1
|
|
724
|
+
|
|
725
|
+
if not configs:
|
|
726
|
+
print_error("No databases selected for benchmark")
|
|
727
|
+
flush_all()
|
|
728
|
+
return 1
|
|
729
|
+
|
|
730
|
+
# Load workload
|
|
731
|
+
json_extra_config = {} # Extra config from JSON file
|
|
732
|
+
try:
|
|
733
|
+
if args.bench_file:
|
|
734
|
+
workload = BenchmarkLoader.from_file(args.bench_file)
|
|
735
|
+
# Read extra config fields from JSON file
|
|
736
|
+
if args.bench_file.endswith('.json'):
|
|
737
|
+
import json
|
|
738
|
+
with open(args.bench_file, 'r') as f:
|
|
739
|
+
json_data = json.load(f)
|
|
740
|
+
json_extra_config = {
|
|
741
|
+
'database': json_data.get('database'),
|
|
742
|
+
'skip_setup': json_data.get('skip_setup'),
|
|
743
|
+
'skip_teardown': json_data.get('skip_teardown'),
|
|
744
|
+
}
|
|
745
|
+
elif args.template:
|
|
746
|
+
workload = BenchmarkLoader.from_builtin(args.template)
|
|
747
|
+
else:
|
|
748
|
+
# Default to oltp_read_write
|
|
749
|
+
print_info("No --bench-file or --template specified, "
|
|
750
|
+
"using built-in", "oltp_read_write")
|
|
751
|
+
workload = BenchmarkLoader.from_builtin("oltp_read_write")
|
|
752
|
+
except (FileNotFoundError, ValueError) as e:
|
|
753
|
+
print_error(str(e))
|
|
754
|
+
flush_all()
|
|
755
|
+
return 1
|
|
756
|
+
|
|
757
|
+
# Build benchmark config
|
|
758
|
+
if args.concurrency > 0:
|
|
759
|
+
mode = WorkloadMode.CONCURRENT
|
|
760
|
+
else:
|
|
761
|
+
mode = WorkloadMode.SERIAL
|
|
762
|
+
|
|
763
|
+
filter_queries = []
|
|
764
|
+
if args.bench_filter:
|
|
765
|
+
filter_queries = [
|
|
766
|
+
n.strip() for n in args.bench_filter.split(",") if n.strip()
|
|
767
|
+
]
|
|
768
|
+
|
|
769
|
+
# Determine skip_setup and skip_teardown: CLI args override JSON config
|
|
770
|
+
json_skip_setup = json_extra_config.get('skip_setup')
|
|
771
|
+
json_skip_teardown = json_extra_config.get('skip_teardown')
|
|
772
|
+
|
|
773
|
+
# Use JSON value as default, CLI arg overrides if explicitly set
|
|
774
|
+
# (CLI arg defaults to False, so only override if user explicitly passed it)
|
|
775
|
+
cli_skip_setup = getattr(args, 'skip_setup', False)
|
|
776
|
+
cli_skip_teardown = getattr(args, 'skip_teardown', False)
|
|
777
|
+
|
|
778
|
+
# If JSON has the value and CLI didn't explicitly set it, use JSON value
|
|
779
|
+
final_skip_setup = cli_skip_setup if cli_skip_setup else (json_skip_setup if json_skip_setup is not None else False)
|
|
780
|
+
final_skip_teardown = cli_skip_teardown if cli_skip_teardown else (json_skip_teardown if json_skip_teardown is not None else False)
|
|
781
|
+
|
|
782
|
+
bench_cfg = BenchmarkConfig(
|
|
783
|
+
mode=mode,
|
|
784
|
+
iterations=args.iterations,
|
|
785
|
+
warmup=args.warmup,
|
|
786
|
+
concurrency=args.concurrency if args.concurrency > 0 else 1,
|
|
787
|
+
duration=args.duration,
|
|
788
|
+
ramp_up=args.ramp_up,
|
|
789
|
+
filter_queries=filter_queries,
|
|
790
|
+
profile=getattr(args, 'profile', False),
|
|
791
|
+
perf_freq=getattr(args, 'perf_freq', 99),
|
|
792
|
+
query_timeout=args.query_timeout,
|
|
793
|
+
flamegraph_min_ms=getattr(args, 'flamegraph_min_ms', 1000),
|
|
794
|
+
skip_setup=final_skip_setup,
|
|
795
|
+
skip_teardown=final_skip_teardown,
|
|
796
|
+
)
|
|
797
|
+
|
|
798
|
+
# Apply filter to workload for display
|
|
799
|
+
display_workload = workload
|
|
800
|
+
if filter_queries:
|
|
801
|
+
try:
|
|
802
|
+
display_workload = BenchmarkLoader.filter_queries(
|
|
803
|
+
workload, filter_queries)
|
|
804
|
+
except ValueError as e:
|
|
805
|
+
print_error(str(e))
|
|
806
|
+
flush_all()
|
|
807
|
+
return 1
|
|
808
|
+
|
|
809
|
+
# Display plan
|
|
810
|
+
parallel_dbms = getattr(args, 'parallel_dbms', False)
|
|
811
|
+
repeat = max(1, getattr(args, 'repeat', 1))
|
|
812
|
+
output_dir = os.path.abspath(args.output_dir)
|
|
813
|
+
fmt = args.format
|
|
814
|
+
|
|
815
|
+
print_phase("Benchmark", workload.name)
|
|
816
|
+
print_info("Mode:", mode.name)
|
|
817
|
+
print_info("DBMS targets:",
|
|
818
|
+
", ".join(c.name for c in configs))
|
|
819
|
+
if parallel_dbms and len(configs) > 1:
|
|
820
|
+
print_info("DBMS execution:", "[bold green]parallel[/bold green]")
|
|
821
|
+
elif not parallel_dbms and len(configs) > 1:
|
|
822
|
+
print_info("DBMS execution:", "sequential")
|
|
823
|
+
print_info("Queries:",
|
|
824
|
+
", ".join(q.name for q in display_workload.queries))
|
|
825
|
+
if mode == WorkloadMode.SERIAL:
|
|
826
|
+
print_info("Iterations:",
|
|
827
|
+
f"{bench_cfg.iterations} Warmup: {bench_cfg.warmup}")
|
|
828
|
+
else:
|
|
829
|
+
print_info("Concurrency:",
|
|
830
|
+
f"{bench_cfg.concurrency} Duration: {bench_cfg.duration}s")
|
|
831
|
+
if filter_queries:
|
|
832
|
+
print_info("Filter:", ", ".join(filter_queries))
|
|
833
|
+
if repeat > 1:
|
|
834
|
+
print_info("Repeat:", f"{repeat} rounds")
|
|
835
|
+
if bench_cfg.profile:
|
|
836
|
+
print_info("Profiling:",
|
|
837
|
+
f"[bold red]🔥 perf flame graph[/bold red] "
|
|
838
|
+
f"(freq: {bench_cfg.perf_freq} Hz)")
|
|
839
|
+
if bench_cfg.skip_setup:
|
|
840
|
+
print_info("Setup:", "[bold yellow]SKIPPED[/bold yellow] (reusing existing tables)")
|
|
841
|
+
if bench_cfg.skip_teardown:
|
|
842
|
+
print_info("Teardown:", "[bold yellow]SKIPPED[/bold yellow] (keeping tables)")
|
|
843
|
+
|
|
844
|
+
# ------------------------------------------------------------------
|
|
845
|
+
# Inner function: execute a single benchmark round
|
|
846
|
+
# ------------------------------------------------------------------
|
|
847
|
+
def _run_one_round(round_num: int) -> int:
|
|
848
|
+
"""Run one benchmark round. Returns 0 on success."""
|
|
849
|
+
if repeat > 1:
|
|
850
|
+
console.print(
|
|
851
|
+
f"\n[bold cyan]{'━' * 60}[/bold cyan]")
|
|
852
|
+
console.print(
|
|
853
|
+
f"[bold cyan] Round {round_num}/{repeat}[/bold cyan]")
|
|
854
|
+
console.print(
|
|
855
|
+
f"[bold cyan]{'━' * 60}[/bold cyan]\n")
|
|
856
|
+
|
|
857
|
+
run_stamp = _time.strftime("%Y%m%d_%H%M%S")
|
|
858
|
+
run_dir = os.path.join(
|
|
859
|
+
output_dir,
|
|
860
|
+
f"bench_{workload.name}_{run_stamp}")
|
|
861
|
+
os.makedirs(run_dir, exist_ok=True)
|
|
862
|
+
|
|
863
|
+
# Execute benchmark
|
|
864
|
+
print_phase("Execute")
|
|
865
|
+
|
|
866
|
+
# Progress tracking (fresh each round)
|
|
867
|
+
progress_bars: Dict[str, BenchProgress] = {}
|
|
868
|
+
_progress_lock = threading.Lock()
|
|
869
|
+
|
|
870
|
+
n_queries = len(display_workload.queries)
|
|
871
|
+
is_concurrent = (mode == WorkloadMode.CONCURRENT)
|
|
872
|
+
if is_concurrent:
|
|
873
|
+
duration = bench_cfg.duration if bench_cfg.duration > 0 else 30.0
|
|
874
|
+
per_query = 100 # placeholder, not used for time-based
|
|
875
|
+
else:
|
|
876
|
+
duration = 0.0
|
|
877
|
+
per_query = bench_cfg.iterations + bench_cfg.warmup
|
|
878
|
+
|
|
879
|
+
# Create progress bars upfront (they will show "setup..." initially)
|
|
880
|
+
if parallel_dbms and len(configs) > 1:
|
|
881
|
+
for c in configs:
|
|
882
|
+
bp = BenchProgress(
|
|
883
|
+
c.name, n_queries, per_query,
|
|
884
|
+
is_concurrent=is_concurrent, duration=duration)
|
|
885
|
+
bp.__enter__()
|
|
886
|
+
bp.set_status("[yellow]正在setup...[/yellow]")
|
|
887
|
+
progress_bars[c.name] = bp
|
|
888
|
+
|
|
889
|
+
def on_setup_start(dbms_name):
|
|
890
|
+
with _progress_lock:
|
|
891
|
+
if dbms_name not in progress_bars:
|
|
892
|
+
bp = BenchProgress(
|
|
893
|
+
dbms_name, n_queries, per_query,
|
|
894
|
+
is_concurrent=is_concurrent, duration=duration)
|
|
895
|
+
bp.__enter__()
|
|
896
|
+
bp.set_status("[yellow]正在setup...[/yellow]")
|
|
897
|
+
progress_bars[dbms_name] = bp
|
|
898
|
+
|
|
899
|
+
def on_setup_done(dbms_name, success):
|
|
900
|
+
bp = progress_bars.get(dbms_name)
|
|
901
|
+
if bp:
|
|
902
|
+
if success:
|
|
903
|
+
bp.set_status("[green]setup完毕[/green]")
|
|
904
|
+
else:
|
|
905
|
+
bp.set_status("[red]setup失败 — 跳过该DBMS[/red]")
|
|
906
|
+
# Close progress bar for failed DBMS
|
|
907
|
+
bp.__exit__(None, None, None)
|
|
908
|
+
bp.write_summary_to_buffer()
|
|
909
|
+
|
|
910
|
+
def on_dbms_start(dbms_name):
|
|
911
|
+
with _progress_lock:
|
|
912
|
+
if dbms_name not in progress_bars:
|
|
913
|
+
bp = BenchProgress(
|
|
914
|
+
dbms_name, n_queries, per_query,
|
|
915
|
+
is_concurrent=is_concurrent, duration=duration)
|
|
916
|
+
bp.__enter__()
|
|
917
|
+
progress_bars[dbms_name] = bp
|
|
918
|
+
|
|
919
|
+
def on_progress(dbms_name, query_name, iteration, total,
|
|
920
|
+
is_warmup=False):
|
|
921
|
+
bp = progress_bars.get(dbms_name)
|
|
922
|
+
if bp and not is_concurrent:
|
|
923
|
+
bp.advance(query_name=query_name, is_warmup=is_warmup)
|
|
924
|
+
|
|
925
|
+
def on_dbms_done(dbms_name, dbms_result):
|
|
926
|
+
bp = progress_bars.get(dbms_name)
|
|
927
|
+
if bp:
|
|
928
|
+
bp.set_status(
|
|
929
|
+
f"[green]{dbms_result.total_queries} queries, "
|
|
930
|
+
f"{dbms_result.overall_qps:.1f} QPS[/green]")
|
|
931
|
+
bp.__exit__(None, None, None)
|
|
932
|
+
bp.write_summary_to_buffer()
|
|
933
|
+
|
|
934
|
+
def on_profile_start(dbms_name, query_name):
|
|
935
|
+
bp = progress_bars.get(dbms_name)
|
|
936
|
+
if bp:
|
|
937
|
+
bp.set_status(f"[red]🔥 profiling {query_name}[/red]")
|
|
938
|
+
|
|
939
|
+
def on_profile_done(dbms_name, query_name, sample_count):
|
|
940
|
+
bp = progress_bars.get(dbms_name)
|
|
941
|
+
if bp:
|
|
942
|
+
bp.set_status(
|
|
943
|
+
f"[dim]🔥 {query_name}: {sample_count} samples[/dim]")
|
|
944
|
+
|
|
945
|
+
# For concurrent mode, timer thread will be started after setup phase
|
|
946
|
+
timer_stop_event = None
|
|
947
|
+
timer_thread = None
|
|
948
|
+
query_phase_started = threading.Event()
|
|
949
|
+
timer_start_time = [None] # Will be set in on_run_start
|
|
950
|
+
|
|
951
|
+
if is_concurrent:
|
|
952
|
+
timer_stop_event = threading.Event()
|
|
953
|
+
|
|
954
|
+
def _timer_update():
|
|
955
|
+
# Wait until query phase starts (all setups complete)
|
|
956
|
+
query_phase_started.wait()
|
|
957
|
+
while not timer_stop_event.is_set():
|
|
958
|
+
# Check if we've exceeded the duration - stop updating progress
|
|
959
|
+
# (actual benchmark may take longer due to cleanup)
|
|
960
|
+
if timer_start_time[0] is not None:
|
|
961
|
+
elapsed = _time.monotonic() - timer_start_time[0]
|
|
962
|
+
if elapsed >= duration:
|
|
963
|
+
break
|
|
964
|
+
for dbms_name, bp in list(progress_bars.items()):
|
|
965
|
+
bp.update_time(status="")
|
|
966
|
+
_time.sleep(0.5)
|
|
967
|
+
|
|
968
|
+
timer_thread = threading.Thread(target=_timer_update, daemon=True)
|
|
969
|
+
timer_thread.start()
|
|
970
|
+
|
|
971
|
+
def on_run_start():
|
|
972
|
+
# Reset timers when query phase begins (all setups complete)
|
|
973
|
+
# Keep "setup完毕" status visible until queries actually start
|
|
974
|
+
with _progress_lock:
|
|
975
|
+
for bp in progress_bars.values():
|
|
976
|
+
bp.reset_timer()
|
|
977
|
+
# Record start time for timer thread
|
|
978
|
+
timer_start_time[0] = _time.monotonic()
|
|
979
|
+
# Signal timer thread to start updating
|
|
980
|
+
query_phase_started.set()
|
|
981
|
+
|
|
982
|
+
# Determine database: JSON config takes precedence over default, CLI arg always wins
|
|
983
|
+
json_database = json_extra_config.get('database')
|
|
984
|
+
# If JSON specifies a database, use it; otherwise use CLI arg (which has default)
|
|
985
|
+
final_database = json_database if json_database else args.database
|
|
986
|
+
|
|
987
|
+
try:
|
|
988
|
+
result = run_benchmark(
|
|
989
|
+
configs=configs,
|
|
990
|
+
workload=workload,
|
|
991
|
+
bench_cfg=bench_cfg,
|
|
992
|
+
database=final_database,
|
|
993
|
+
on_progress=on_progress,
|
|
994
|
+
on_dbms_start=on_dbms_start,
|
|
995
|
+
on_dbms_done=on_dbms_done,
|
|
996
|
+
on_profile_start=on_profile_start if bench_cfg.profile else None,
|
|
997
|
+
on_profile_done=on_profile_done if bench_cfg.profile else None,
|
|
998
|
+
on_run_start=on_run_start,
|
|
999
|
+
on_setup_start=on_setup_start,
|
|
1000
|
+
on_setup_done=on_setup_done,
|
|
1001
|
+
parallel_dbms=parallel_dbms,
|
|
1002
|
+
)
|
|
1003
|
+
finally:
|
|
1004
|
+
# Stop timer thread
|
|
1005
|
+
if timer_stop_event is not None:
|
|
1006
|
+
timer_stop_event.set()
|
|
1007
|
+
if timer_thread is not None:
|
|
1008
|
+
timer_thread.join(timeout=1.0)
|
|
1009
|
+
|
|
1010
|
+
# Generate reports
|
|
1011
|
+
print_phase("Reports")
|
|
1012
|
+
|
|
1013
|
+
if fmt in ("text", "all"):
|
|
1014
|
+
text_path = os.path.join(
|
|
1015
|
+
run_dir, f"bench_{workload.name}.report.txt")
|
|
1016
|
+
write_bench_text_report(text_path, result)
|
|
1017
|
+
print_report_file(text_path, label="text")
|
|
1018
|
+
|
|
1019
|
+
if fmt in ("html", "all"):
|
|
1020
|
+
html_path = os.path.join(
|
|
1021
|
+
run_dir, f"bench_{workload.name}.html")
|
|
1022
|
+
write_bench_html_report(html_path, result)
|
|
1023
|
+
print_report_file(html_path, label="html")
|
|
1024
|
+
|
|
1025
|
+
# Save raw JSON data
|
|
1026
|
+
json_path = os.path.join(run_dir, "bench_result.json")
|
|
1027
|
+
_save_bench_json(json_path, result, bench_file=args.bench_file or "", database=final_database)
|
|
1028
|
+
print_report_file(json_path, label="json")
|
|
1029
|
+
|
|
1030
|
+
# Update 'latest' symlink
|
|
1031
|
+
latest_link = os.path.join(output_dir, "latest")
|
|
1032
|
+
try:
|
|
1033
|
+
if os.path.islink(latest_link):
|
|
1034
|
+
os.remove(latest_link)
|
|
1035
|
+
os.symlink(os.path.basename(run_dir), latest_link)
|
|
1036
|
+
except OSError:
|
|
1037
|
+
pass
|
|
1038
|
+
|
|
1039
|
+
# Generate history index
|
|
1040
|
+
generate_index_html(output_dir)
|
|
1041
|
+
|
|
1042
|
+
# Print rich summary
|
|
1043
|
+
print_bench_summary(result)
|
|
1044
|
+
flush_all()
|
|
1045
|
+
|
|
1046
|
+
return run_dir
|
|
1047
|
+
|
|
1048
|
+
# ------------------------------------------------------------------
|
|
1049
|
+
# Main loop
|
|
1050
|
+
# ------------------------------------------------------------------
|
|
1051
|
+
last_run_dir = None
|
|
1052
|
+
for rnd in range(1, repeat + 1):
|
|
1053
|
+
try:
|
|
1054
|
+
last_run_dir = _run_one_round(rnd)
|
|
1055
|
+
except KeyboardInterrupt:
|
|
1056
|
+
console.print(
|
|
1057
|
+
f"\n[yellow]Interrupted at round {rnd}/{repeat}. "
|
|
1058
|
+
f"Stopping.[/yellow]")
|
|
1059
|
+
flush_all()
|
|
1060
|
+
break
|
|
1061
|
+
# Small pause between rounds to avoid timestamp collision
|
|
1062
|
+
if rnd < repeat:
|
|
1063
|
+
_time.sleep(1)
|
|
1064
|
+
|
|
1065
|
+
if repeat > 1:
|
|
1066
|
+
console.print(
|
|
1067
|
+
f"\n[bold green]All {repeat} rounds completed.[/bold green]")
|
|
1068
|
+
flush_all()
|
|
1069
|
+
|
|
1070
|
+
# Serve if requested (use the latest run)
|
|
1071
|
+
if args.serve and fmt in ("html", "all") and last_run_dir:
|
|
1072
|
+
html_file = f"bench_{workload.name}.html"
|
|
1073
|
+
html_path = os.path.join(last_run_dir, html_file)
|
|
1074
|
+
if os.path.isfile(html_path):
|
|
1075
|
+
relative_html = os.path.join(
|
|
1076
|
+
os.path.basename(last_run_dir), html_file)
|
|
1077
|
+
_serve_report(output_dir, relative_html, args.port)
|
|
1078
|
+
|
|
1079
|
+
return 0
|
|
1080
|
+
|
|
1081
|
+
|
|
1082
|
+
def _save_bench_json(path: str, result, bench_file: str = "", database: str = ""):
|
|
1083
|
+
"""Save benchmark result as JSON for later analysis.
|
|
1084
|
+
|
|
1085
|
+
Args:
|
|
1086
|
+
path: Output file path
|
|
1087
|
+
result: BenchmarkResult object
|
|
1088
|
+
bench_file: Path to benchmark file (.json or .sql)
|
|
1089
|
+
database: Database name used for this run
|
|
1090
|
+
"""
|
|
1091
|
+
import json
|
|
1092
|
+
data = {
|
|
1093
|
+
"workload": result.workload_name,
|
|
1094
|
+
"mode": result.mode.name,
|
|
1095
|
+
"timestamp": result.timestamp,
|
|
1096
|
+
"run_id": result.run_id or "",
|
|
1097
|
+
"bench_file": bench_file, # Store benchmark file path for rerun
|
|
1098
|
+
"database": database, # Store database name for rerun
|
|
1099
|
+
"table_rows": result.table_rows,
|
|
1100
|
+
"table_rows_detail": result.table_rows_detail or {},
|
|
1101
|
+
"table_schema": result.table_schema or {}, # {table_name: CREATE TABLE stmt}
|
|
1102
|
+
"setup_sql": list(result.setup_sql) if result.setup_sql else [],
|
|
1103
|
+
"teardown_sql": list(result.teardown_sql) if result.teardown_sql else [],
|
|
1104
|
+
"queries_sql": list(result.queries_sql) if result.queries_sql else [],
|
|
1105
|
+
"config": {
|
|
1106
|
+
"iterations": result.config.iterations,
|
|
1107
|
+
"warmup": result.config.warmup,
|
|
1108
|
+
"concurrency": result.config.concurrency,
|
|
1109
|
+
"duration": result.config.duration,
|
|
1110
|
+
"filter_queries": result.config.filter_queries,
|
|
1111
|
+
},
|
|
1112
|
+
"dbms_results": [],
|
|
1113
|
+
}
|
|
1114
|
+
for dr in result.dbms_results:
|
|
1115
|
+
dbms_data = {
|
|
1116
|
+
"dbms_name": dr.dbms_name,
|
|
1117
|
+
"total_duration_s": round(dr.total_duration_s, 3),
|
|
1118
|
+
"total_queries": dr.total_queries,
|
|
1119
|
+
"total_errors": dr.total_errors,
|
|
1120
|
+
"overall_qps": round(dr.overall_qps, 2),
|
|
1121
|
+
"table_rows": dr.table_rows,
|
|
1122
|
+
"table_rows_detail": dr.table_rows_detail or {},
|
|
1123
|
+
"table_schema": dr.table_schema or {}, # {table_name: CREATE TABLE stmt}
|
|
1124
|
+
"query_stats": [],
|
|
1125
|
+
}
|
|
1126
|
+
for qs in dr.query_stats:
|
|
1127
|
+
dbms_data["query_stats"].append({
|
|
1128
|
+
"query_name": qs.query_name,
|
|
1129
|
+
"sql_template": qs.sql_template or "",
|
|
1130
|
+
"total_executions": qs.total_executions,
|
|
1131
|
+
"total_errors": qs.total_errors,
|
|
1132
|
+
"min_ms": round(qs.min_ms, 3),
|
|
1133
|
+
"max_ms": round(qs.max_ms, 3),
|
|
1134
|
+
"avg_ms": round(qs.avg_ms, 3),
|
|
1135
|
+
"p50_ms": round(qs.p50_ms, 3),
|
|
1136
|
+
"p95_ms": round(qs.p95_ms, 3),
|
|
1137
|
+
"p99_ms": round(qs.p99_ms, 3),
|
|
1138
|
+
"qps": round(qs.qps, 2),
|
|
1139
|
+
"latencies_ms": [round(l, 3) for l in qs.latencies_ms] if qs.latencies_ms else [],
|
|
1140
|
+
"explain_plan": qs.explain_plan or "",
|
|
1141
|
+
"explain_tree": qs.explain_tree or "",
|
|
1142
|
+
"error_logs": qs.error_logs[:50] if qs.error_logs else [],
|
|
1143
|
+
})
|
|
1144
|
+
data["dbms_results"].append(dbms_data)
|
|
1145
|
+
|
|
1146
|
+
with open(path, "w", encoding="utf-8") as f:
|
|
1147
|
+
json.dump(data, f, indent=2, ensure_ascii=False)
|
|
1148
|
+
|
|
1149
|
+
|
|
1150
|
+
def run_benchmark_with_progress(
|
|
1151
|
+
configs: List[DBMSConfig],
|
|
1152
|
+
workload,
|
|
1153
|
+
bench_cfg,
|
|
1154
|
+
database: str,
|
|
1155
|
+
output_dir: str,
|
|
1156
|
+
output_format: str = "all",
|
|
1157
|
+
parallel_dbms: bool = True,
|
|
1158
|
+
json_extra_config: Optional[dict] = None,
|
|
1159
|
+
callbacks: Optional[dict] = None,
|
|
1160
|
+
bench_file: str = "",
|
|
1161
|
+
) -> Tuple[str, object]:
|
|
1162
|
+
"""Core benchmark execution logic shared by CLI and Interactive modes.
|
|
1163
|
+
|
|
1164
|
+
Args:
|
|
1165
|
+
configs: List of DBMS configurations
|
|
1166
|
+
workload: Benchmark workload (from BenchmarkLoader)
|
|
1167
|
+
bench_cfg: Benchmark configuration
|
|
1168
|
+
database: Database name
|
|
1169
|
+
output_dir: Output directory for reports
|
|
1170
|
+
output_format: Output format (text, html, all)
|
|
1171
|
+
parallel_dbms: Whether to run benchmarks in parallel
|
|
1172
|
+
json_extra_config: Extra config from JSON file (database, skip_setup, skip_teardown)
|
|
1173
|
+
callbacks: Optional callbacks for progress tracking:
|
|
1174
|
+
- on_progress(dbms_name, query_name, iteration, total, is_warmup)
|
|
1175
|
+
- on_dbms_start(dbms_name)
|
|
1176
|
+
- on_dbms_done(dbms_name, dbms_result)
|
|
1177
|
+
- on_profile_start(dbms_name, query_name)
|
|
1178
|
+
- on_profile_done(dbms_name, query_name, sample_count)
|
|
1179
|
+
- on_run_start()
|
|
1180
|
+
- on_setup_start(dbms_name)
|
|
1181
|
+
- on_setup_done(dbms_name, success)
|
|
1182
|
+
bench_file: Path to benchmark file (.json or .sql) for rerun support
|
|
1183
|
+
|
|
1184
|
+
Returns:
|
|
1185
|
+
Tuple of (run_dir, result)
|
|
1186
|
+
"""
|
|
1187
|
+
from .benchmark import run_benchmark
|
|
1188
|
+
from .reporter.bench_text import write_bench_text_report
|
|
1189
|
+
from .reporter.bench_html import write_bench_html_report
|
|
1190
|
+
from .reporter.history import generate_index_html
|
|
1191
|
+
|
|
1192
|
+
callbacks = callbacks or {}
|
|
1193
|
+
|
|
1194
|
+
# Create output directory
|
|
1195
|
+
run_stamp = _time.strftime("%Y%m%d_%H%M%S")
|
|
1196
|
+
run_dir = os.path.join(
|
|
1197
|
+
output_dir,
|
|
1198
|
+
f"bench_{workload.name}_{run_stamp}"
|
|
1199
|
+
)
|
|
1200
|
+
os.makedirs(run_dir, exist_ok=True)
|
|
1201
|
+
|
|
1202
|
+
# Determine database from JSON config if provided
|
|
1203
|
+
json_extra_config = json_extra_config or {}
|
|
1204
|
+
json_database = json_extra_config.get('database')
|
|
1205
|
+
final_database = json_database if json_database else database
|
|
1206
|
+
|
|
1207
|
+
# Execute benchmark
|
|
1208
|
+
result = run_benchmark(
|
|
1209
|
+
configs=configs,
|
|
1210
|
+
workload=workload,
|
|
1211
|
+
bench_cfg=bench_cfg,
|
|
1212
|
+
database=final_database,
|
|
1213
|
+
on_progress=callbacks.get('on_progress'),
|
|
1214
|
+
on_dbms_start=callbacks.get('on_dbms_start'),
|
|
1215
|
+
on_dbms_done=callbacks.get('on_dbms_done'),
|
|
1216
|
+
on_profile_start=callbacks.get('on_profile_start'),
|
|
1217
|
+
on_profile_done=callbacks.get('on_profile_done'),
|
|
1218
|
+
on_run_start=callbacks.get('on_run_start'),
|
|
1219
|
+
on_setup_start=callbacks.get('on_setup_start'),
|
|
1220
|
+
on_setup_done=callbacks.get('on_setup_done'),
|
|
1221
|
+
parallel_dbms=parallel_dbms,
|
|
1222
|
+
)
|
|
1223
|
+
# Set run_id for the result
|
|
1224
|
+
result.run_id = os.path.basename(run_dir)
|
|
1225
|
+
|
|
1226
|
+
# Generate reports
|
|
1227
|
+
report_files = []
|
|
1228
|
+
|
|
1229
|
+
if output_format in ("text", "all"):
|
|
1230
|
+
text_path = os.path.join(run_dir, f"bench_{workload.name}.report.txt")
|
|
1231
|
+
write_bench_text_report(text_path, result)
|
|
1232
|
+
report_files.append(text_path)
|
|
1233
|
+
|
|
1234
|
+
if output_format in ("html", "all"):
|
|
1235
|
+
html_path = os.path.join(run_dir, f"bench_{workload.name}.html")
|
|
1236
|
+
write_bench_html_report(html_path, result)
|
|
1237
|
+
report_files.append(html_path)
|
|
1238
|
+
|
|
1239
|
+
# Save JSON result
|
|
1240
|
+
json_path = os.path.join(run_dir, "bench_result.json")
|
|
1241
|
+
_save_bench_json(json_path, result, bench_file=bench_file, database=final_database)
|
|
1242
|
+
report_files.append(json_path)
|
|
1243
|
+
|
|
1244
|
+
# Update latest symlink
|
|
1245
|
+
latest_link = os.path.join(output_dir, "latest")
|
|
1246
|
+
try:
|
|
1247
|
+
if os.path.islink(latest_link):
|
|
1248
|
+
os.remove(latest_link)
|
|
1249
|
+
os.symlink(os.path.basename(run_dir), latest_link)
|
|
1250
|
+
except OSError:
|
|
1251
|
+
pass
|
|
1252
|
+
|
|
1253
|
+
# Generate history index
|
|
1254
|
+
generate_index_html(output_dir)
|
|
1255
|
+
|
|
1256
|
+
return run_dir, result
|
|
1257
|
+
|
|
1258
|
+
|
|
1259
|
+
def _select_bench_params(
|
|
1260
|
+
iterations: int = 100,
|
|
1261
|
+
warmup: int = 5,
|
|
1262
|
+
concurrency: int = 8,
|
|
1263
|
+
duration: float = 30.0,
|
|
1264
|
+
ramp_up: float = 0.0,
|
|
1265
|
+
profile: bool = True,
|
|
1266
|
+
skip_setup: bool = False,
|
|
1267
|
+
skip_teardown: bool = False,
|
|
1268
|
+
output_dir: str = "",
|
|
1269
|
+
) -> Optional[dict]:
|
|
1270
|
+
"""Show an interactive benchmark parameter configuration panel.
|
|
1271
|
+
|
|
1272
|
+
First, select mode (SERIAL, CONCURRENT, or RERUN), then configure parameters
|
|
1273
|
+
based on the selected mode.
|
|
1274
|
+
|
|
1275
|
+
Returns a dict with mode-specific parameters, or ``None`` if cancelled.
|
|
1276
|
+
For RERUN mode, returns {"mode": "rerun", "run_data": {...}}.
|
|
1277
|
+
"""
|
|
1278
|
+
# Step 1: Mode selection
|
|
1279
|
+
mode_result = _select_bench_mode()
|
|
1280
|
+
if mode_result is None:
|
|
1281
|
+
return None
|
|
1282
|
+
if mode_result.get("action") == "back":
|
|
1283
|
+
return {"action": "back"}
|
|
1284
|
+
|
|
1285
|
+
mode = mode_result["mode"] # "serial", "concurrent", or "rerun"
|
|
1286
|
+
|
|
1287
|
+
# Step 2: Parameter configuration based on mode
|
|
1288
|
+
if mode == "serial":
|
|
1289
|
+
return _select_serial_params(iterations, warmup, profile, skip_setup, skip_teardown)
|
|
1290
|
+
elif mode == "concurrent":
|
|
1291
|
+
return _select_concurrent_params(concurrency, duration, ramp_up, profile, skip_setup, skip_teardown)
|
|
1292
|
+
else: # mode == "rerun"
|
|
1293
|
+
# Load historical run parameters
|
|
1294
|
+
run_selection = _select_rerun_run_id(output_dir)
|
|
1295
|
+
if run_selection is None:
|
|
1296
|
+
# User cancelled rerun selection, return special marker to re-show Benchmark Mode
|
|
1297
|
+
return {"action": "cancel"}
|
|
1298
|
+
return {"mode": "rerun", "run_data": run_selection}
|
|
1299
|
+
|
|
1300
|
+
|
|
1301
|
+
def _select_bench_mode() -> Optional[dict]:
|
|
1302
|
+
"""Show mode selection dialog for benchmark.
|
|
1303
|
+
|
|
1304
|
+
Returns dict with "mode" key ("serial" or "concurrent"),
|
|
1305
|
+
or None if cancelled.
|
|
1306
|
+
"""
|
|
1307
|
+
from prompt_toolkit import Application
|
|
1308
|
+
from prompt_toolkit.key_binding import KeyBindings
|
|
1309
|
+
from prompt_toolkit.layout import Layout
|
|
1310
|
+
from prompt_toolkit.layout.containers import HSplit, Window
|
|
1311
|
+
from prompt_toolkit.layout.controls import FormattedTextControl
|
|
1312
|
+
|
|
1313
|
+
MODES = [
|
|
1314
|
+
("serial", "SERIAL",
|
|
1315
|
+
"Sequential execution, fixed iterations per query"),
|
|
1316
|
+
("concurrent", "CONCURRENT",
|
|
1317
|
+
"Multi-threaded stress test with duration-based execution"),
|
|
1318
|
+
("rerun", "RERUN",
|
|
1319
|
+
"Replay a historical benchmark run"),
|
|
1320
|
+
]
|
|
1321
|
+
|
|
1322
|
+
selected = [0]
|
|
1323
|
+
result = [None]
|
|
1324
|
+
|
|
1325
|
+
kb = KeyBindings()
|
|
1326
|
+
|
|
1327
|
+
@kb.add("up")
|
|
1328
|
+
@kb.add("k")
|
|
1329
|
+
def _up(event):
|
|
1330
|
+
selected[0] = (selected[0] - 1) % len(MODES)
|
|
1331
|
+
|
|
1332
|
+
@kb.add("down")
|
|
1333
|
+
@kb.add("j")
|
|
1334
|
+
def _down(event):
|
|
1335
|
+
selected[0] = (selected[0] + 1) % len(MODES)
|
|
1336
|
+
|
|
1337
|
+
@kb.add("enter")
|
|
1338
|
+
def _confirm(event):
|
|
1339
|
+
result[0] = {"mode": MODES[selected[0]][0]}
|
|
1340
|
+
event.app.exit()
|
|
1341
|
+
|
|
1342
|
+
@kb.add("c-c")
|
|
1343
|
+
@kb.add("escape")
|
|
1344
|
+
def _cancel(event):
|
|
1345
|
+
result[0] = None
|
|
1346
|
+
event.app.exit()
|
|
1347
|
+
|
|
1348
|
+
@kb.add("b")
|
|
1349
|
+
def _back(event):
|
|
1350
|
+
result[0] = {"action": "back"}
|
|
1351
|
+
event.app.exit()
|
|
1352
|
+
|
|
1353
|
+
def _get_text():
|
|
1354
|
+
lines = []
|
|
1355
|
+
border = "═" * 58
|
|
1356
|
+
title = "Select Benchmark Mode".center(58)
|
|
1357
|
+
hint = "↑↓ move · Enter select · B back · Esc cancel".center(58)
|
|
1358
|
+
|
|
1359
|
+
lines.append(("bold cyan", f" ╔{border}╗\n"))
|
|
1360
|
+
lines.append(("bold cyan", " ║"))
|
|
1361
|
+
lines.append(("bold white", title))
|
|
1362
|
+
lines.append(("bold cyan", "║\n"))
|
|
1363
|
+
lines.append(("bold cyan", " ║"))
|
|
1364
|
+
lines.append(("", hint))
|
|
1365
|
+
lines.append(("bold cyan", "║\n"))
|
|
1366
|
+
lines.append(("bold cyan", f" ╚{border}╝\n"))
|
|
1367
|
+
lines.append(("", "\n"))
|
|
1368
|
+
|
|
1369
|
+
for i, (mode_key, mode_name, mode_desc) in enumerate(MODES):
|
|
1370
|
+
is_sel = (i == selected[0])
|
|
1371
|
+
if is_sel:
|
|
1372
|
+
lines.append(("bold cyan", " ❯ "))
|
|
1373
|
+
lines.append(("bold yellow", mode_name))
|
|
1374
|
+
lines.append(("", "\n"))
|
|
1375
|
+
lines.append(("", " "))
|
|
1376
|
+
lines.append(("dim", mode_desc))
|
|
1377
|
+
lines.append(("", "\n"))
|
|
1378
|
+
else:
|
|
1379
|
+
lines.append(("", " "))
|
|
1380
|
+
lines.append(("bold", mode_name))
|
|
1381
|
+
lines.append(("", "\n"))
|
|
1382
|
+
lines.append(("", " "))
|
|
1383
|
+
lines.append(("dim", mode_desc))
|
|
1384
|
+
lines.append(("", "\n"))
|
|
1385
|
+
|
|
1386
|
+
lines.append(("", "\n"))
|
|
1387
|
+
lines.append(("dim", " ────────────────────────────────────────\n"))
|
|
1388
|
+
lines.append(("dim", " SERIAL: Each query runs N times sequentially\n"))
|
|
1389
|
+
lines.append(("dim", " CONCURRENT: Multiple threads, duration-based test\n"))
|
|
1390
|
+
lines.append(("dim", " RERUN: Replay a historical benchmark run\n"))
|
|
1391
|
+
|
|
1392
|
+
return lines
|
|
1393
|
+
|
|
1394
|
+
menu = Window(
|
|
1395
|
+
content=FormattedTextControl(_get_text),
|
|
1396
|
+
dont_extend_height=True,
|
|
1397
|
+
)
|
|
1398
|
+
|
|
1399
|
+
app: Application = Application(
|
|
1400
|
+
layout=Layout(HSplit([menu])),
|
|
1401
|
+
key_bindings=kb,
|
|
1402
|
+
full_screen=False,
|
|
1403
|
+
)
|
|
1404
|
+
|
|
1405
|
+
_tty_write("\033[s")
|
|
1406
|
+
app.run()
|
|
1407
|
+
_tty_write("\033[u\033[J")
|
|
1408
|
+
|
|
1409
|
+
return result[0]
|
|
1410
|
+
|
|
1411
|
+
|
|
1412
|
+
def _select_serial_params(
|
|
1413
|
+
iterations: int = 100,
|
|
1414
|
+
warmup: int = 5,
|
|
1415
|
+
profile: bool = True,
|
|
1416
|
+
skip_setup: bool = False,
|
|
1417
|
+
skip_teardown: bool = False,
|
|
1418
|
+
) -> Optional[dict]:
|
|
1419
|
+
"""Show parameter configuration for SERIAL mode."""
|
|
1420
|
+
from prompt_toolkit import Application
|
|
1421
|
+
from prompt_toolkit.key_binding import KeyBindings
|
|
1422
|
+
from prompt_toolkit.layout import Layout
|
|
1423
|
+
from prompt_toolkit.layout.containers import HSplit, Window
|
|
1424
|
+
from prompt_toolkit.layout.controls import FormattedTextControl
|
|
1425
|
+
from prompt_toolkit.keys import Keys
|
|
1426
|
+
from prompt_toolkit.filters import Condition
|
|
1427
|
+
|
|
1428
|
+
ITER_PRESETS = [10, 50, 100, 200, 500, 1000]
|
|
1429
|
+
WARMUP_PRESETS = [0, 5, 10, 20, 50]
|
|
1430
|
+
PROFILE_LABELS = {False: "Off", True: "On"}
|
|
1431
|
+
SKIP_LABELS = {False: "Off", True: "On"}
|
|
1432
|
+
|
|
1433
|
+
custom_iter = [None]
|
|
1434
|
+
custom_warmup = [None]
|
|
1435
|
+
|
|
1436
|
+
result = [None]
|
|
1437
|
+
sel = [0]
|
|
1438
|
+
it_idx = [ITER_PRESETS.index(iterations)
|
|
1439
|
+
if iterations in ITER_PRESETS else 2]
|
|
1440
|
+
wa_idx = [WARMUP_PRESETS.index(warmup)
|
|
1441
|
+
if warmup in WARMUP_PRESETS else 1]
|
|
1442
|
+
prof = [profile]
|
|
1443
|
+
s_setup = [skip_setup]
|
|
1444
|
+
s_teardown = [skip_teardown]
|
|
1445
|
+
|
|
1446
|
+
FIELDS = [
|
|
1447
|
+
{"label": "Iterations", "type": "choice"},
|
|
1448
|
+
{"label": "Warmup", "type": "choice"},
|
|
1449
|
+
{"label": "Profile (flame graph)", "type": "toggle", "var": "prof"},
|
|
1450
|
+
{"label": "Skip Setup (reuse tables)", "type": "toggle", "var": "s_setup"},
|
|
1451
|
+
{"label": "Skip Teardown (keep tables)", "type": "toggle", "var": "s_teardown"},
|
|
1452
|
+
{"label": "OK", "type": "action"},
|
|
1453
|
+
{"label": "Back", "type": "action"},
|
|
1454
|
+
{"label": "Quit", "type": "action"},
|
|
1455
|
+
]
|
|
1456
|
+
|
|
1457
|
+
ACTION_OK = len(FIELDS) - 3
|
|
1458
|
+
ACTION_BACK = len(FIELDS) - 2
|
|
1459
|
+
ACTION_QUIT = len(FIELDS) - 1
|
|
1460
|
+
|
|
1461
|
+
def _iter_val():
|
|
1462
|
+
if custom_iter[0] is not None:
|
|
1463
|
+
return custom_iter[0]
|
|
1464
|
+
return ITER_PRESETS[it_idx[0]]
|
|
1465
|
+
|
|
1466
|
+
def _warmup_val():
|
|
1467
|
+
if custom_warmup[0] is not None:
|
|
1468
|
+
return custom_warmup[0]
|
|
1469
|
+
return WARMUP_PRESETS[wa_idx[0]]
|
|
1470
|
+
|
|
1471
|
+
def _field_val(i):
|
|
1472
|
+
if i == 0:
|
|
1473
|
+
v = _iter_val()
|
|
1474
|
+
if custom_iter[0] is not None:
|
|
1475
|
+
return f"{v} (custom)"
|
|
1476
|
+
return str(v)
|
|
1477
|
+
elif i == 1:
|
|
1478
|
+
v = _warmup_val()
|
|
1479
|
+
if custom_warmup[0] is not None:
|
|
1480
|
+
return f"{v} (custom)"
|
|
1481
|
+
return str(v)
|
|
1482
|
+
elif i == 2:
|
|
1483
|
+
return PROFILE_LABELS[prof[0]]
|
|
1484
|
+
elif i == 3:
|
|
1485
|
+
return SKIP_LABELS[s_setup[0]]
|
|
1486
|
+
elif i == 4:
|
|
1487
|
+
return SKIP_LABELS[s_teardown[0]]
|
|
1488
|
+
return ""
|
|
1489
|
+
|
|
1490
|
+
def _get_toggle_var(i):
|
|
1491
|
+
"""Get the toggle variable list for field index."""
|
|
1492
|
+
if i == 2: return prof
|
|
1493
|
+
if i == 3: return s_setup
|
|
1494
|
+
if i == 4: return s_teardown
|
|
1495
|
+
return None
|
|
1496
|
+
|
|
1497
|
+
def _toggle_right(i):
|
|
1498
|
+
if i == 0:
|
|
1499
|
+
if custom_iter[0] is not None:
|
|
1500
|
+
custom_iter[0] = None
|
|
1501
|
+
else:
|
|
1502
|
+
if it_idx[0] == len(ITER_PRESETS) - 1:
|
|
1503
|
+
custom_iter[0] = _iter_val()
|
|
1504
|
+
else:
|
|
1505
|
+
it_idx[0] += 1
|
|
1506
|
+
elif i == 1:
|
|
1507
|
+
if custom_warmup[0] is not None:
|
|
1508
|
+
custom_warmup[0] = None
|
|
1509
|
+
else:
|
|
1510
|
+
if wa_idx[0] == len(WARMUP_PRESETS) - 1:
|
|
1511
|
+
custom_warmup[0] = _warmup_val()
|
|
1512
|
+
else:
|
|
1513
|
+
wa_idx[0] += 1
|
|
1514
|
+
else:
|
|
1515
|
+
var = _get_toggle_var(i)
|
|
1516
|
+
if var is not None:
|
|
1517
|
+
var[0] = not var[0]
|
|
1518
|
+
|
|
1519
|
+
def _toggle_left(i):
|
|
1520
|
+
if i == 0:
|
|
1521
|
+
if custom_iter[0] is not None:
|
|
1522
|
+
custom_iter[0] = None
|
|
1523
|
+
else:
|
|
1524
|
+
if it_idx[0] == 0:
|
|
1525
|
+
custom_iter[0] = _iter_val()
|
|
1526
|
+
it_idx[0] = 0
|
|
1527
|
+
else:
|
|
1528
|
+
it_idx[0] -= 1
|
|
1529
|
+
elif i == 1:
|
|
1530
|
+
if custom_warmup[0] is not None:
|
|
1531
|
+
custom_warmup[0] = None
|
|
1532
|
+
else:
|
|
1533
|
+
if wa_idx[0] == 0:
|
|
1534
|
+
custom_warmup[0] = _warmup_val()
|
|
1535
|
+
wa_idx[0] = 0
|
|
1536
|
+
else:
|
|
1537
|
+
wa_idx[0] -= 1
|
|
1538
|
+
else:
|
|
1539
|
+
var = _get_toggle_var(i)
|
|
1540
|
+
if var is not None:
|
|
1541
|
+
var[0] = not var[0]
|
|
1542
|
+
|
|
1543
|
+
editing = [None]
|
|
1544
|
+
edit_buf = [""]
|
|
1545
|
+
|
|
1546
|
+
kb = KeyBindings()
|
|
1547
|
+
|
|
1548
|
+
@kb.add("up")
|
|
1549
|
+
@kb.add("k")
|
|
1550
|
+
def _up(event):
|
|
1551
|
+
if editing[0] is not None:
|
|
1552
|
+
return
|
|
1553
|
+
sel[0] = (sel[0] - 1) % len(FIELDS)
|
|
1554
|
+
|
|
1555
|
+
@kb.add("down")
|
|
1556
|
+
@kb.add("j")
|
|
1557
|
+
def _down(event):
|
|
1558
|
+
if editing[0] is not None:
|
|
1559
|
+
return
|
|
1560
|
+
sel[0] = (sel[0] + 1) % len(FIELDS)
|
|
1561
|
+
|
|
1562
|
+
@kb.add("left")
|
|
1563
|
+
@kb.add("h")
|
|
1564
|
+
def _left(event):
|
|
1565
|
+
if editing[0] is not None:
|
|
1566
|
+
return
|
|
1567
|
+
_toggle_left(sel[0])
|
|
1568
|
+
|
|
1569
|
+
@kb.add("right")
|
|
1570
|
+
@kb.add("l")
|
|
1571
|
+
def _right(event):
|
|
1572
|
+
if editing[0] is not None:
|
|
1573
|
+
return
|
|
1574
|
+
_toggle_right(sel[0])
|
|
1575
|
+
|
|
1576
|
+
@kb.add("backspace")
|
|
1577
|
+
def _backspace(event):
|
|
1578
|
+
if editing[0] is not None:
|
|
1579
|
+
edit_buf[0] = edit_buf[0][:-1]
|
|
1580
|
+
|
|
1581
|
+
@kb.add(Keys.Any, filter=Condition(lambda: editing[0] is not None))
|
|
1582
|
+
def _type_char(event):
|
|
1583
|
+
ch = event.data
|
|
1584
|
+
if ch.isdigit():
|
|
1585
|
+
edit_buf[0] += ch
|
|
1586
|
+
|
|
1587
|
+
@kb.add("enter")
|
|
1588
|
+
def _confirm(event):
|
|
1589
|
+
if editing[0] is not None:
|
|
1590
|
+
idx = editing[0]
|
|
1591
|
+
if edit_buf[0]:
|
|
1592
|
+
try:
|
|
1593
|
+
n = int(edit_buf[0])
|
|
1594
|
+
if n >= 0:
|
|
1595
|
+
if idx == 0:
|
|
1596
|
+
custom_iter[0] = n
|
|
1597
|
+
else:
|
|
1598
|
+
custom_warmup[0] = n
|
|
1599
|
+
except ValueError:
|
|
1600
|
+
pass
|
|
1601
|
+
editing[0] = None
|
|
1602
|
+
edit_buf[0] = ""
|
|
1603
|
+
return
|
|
1604
|
+
|
|
1605
|
+
if sel[0] == 0 and custom_iter[0] is not None:
|
|
1606
|
+
editing[0] = 0
|
|
1607
|
+
edit_buf[0] = str(custom_iter[0])
|
|
1608
|
+
return
|
|
1609
|
+
if sel[0] == 1 and custom_warmup[0] is not None:
|
|
1610
|
+
editing[0] = 1
|
|
1611
|
+
edit_buf[0] = str(custom_warmup[0])
|
|
1612
|
+
return
|
|
1613
|
+
|
|
1614
|
+
if sel[0] == ACTION_OK:
|
|
1615
|
+
result[0] = {
|
|
1616
|
+
"mode": "serial",
|
|
1617
|
+
"iterations": _iter_val(),
|
|
1618
|
+
"warmup": _warmup_val(),
|
|
1619
|
+
"concurrency": 0,
|
|
1620
|
+
"duration": 0.0,
|
|
1621
|
+
"ramp_up": 0.0,
|
|
1622
|
+
"profile": prof[0],
|
|
1623
|
+
"skip_setup": s_setup[0],
|
|
1624
|
+
"skip_teardown": s_teardown[0],
|
|
1625
|
+
}
|
|
1626
|
+
event.app.exit()
|
|
1627
|
+
return
|
|
1628
|
+
if sel[0] == ACTION_BACK:
|
|
1629
|
+
result[0] = {"action": "back"}
|
|
1630
|
+
event.app.exit()
|
|
1631
|
+
return
|
|
1632
|
+
if sel[0] == ACTION_QUIT:
|
|
1633
|
+
result[0] = None
|
|
1634
|
+
event.app.exit()
|
|
1635
|
+
return
|
|
1636
|
+
|
|
1637
|
+
@kb.add("c-c")
|
|
1638
|
+
@kb.add("escape")
|
|
1639
|
+
def _cancel(event):
|
|
1640
|
+
if editing[0] is not None:
|
|
1641
|
+
editing[0] = None
|
|
1642
|
+
edit_buf[0] = ""
|
|
1643
|
+
return
|
|
1644
|
+
result[0] = None
|
|
1645
|
+
event.app.exit()
|
|
1646
|
+
|
|
1647
|
+
def _get_text():
|
|
1648
|
+
lines = []
|
|
1649
|
+
border = "═" * 55
|
|
1650
|
+
title = "SERIAL Mode Configuration".center(55)
|
|
1651
|
+
hint = ("←→ change · Enter confirm/custom · ↑↓ move"
|
|
1652
|
+
" · Esc cancel").center(55)
|
|
1653
|
+
lines.append(("bold cyan", f" ╔{border}╗\n"))
|
|
1654
|
+
lines.append(("bold cyan", " ║"))
|
|
1655
|
+
lines.append(("bold white", title))
|
|
1656
|
+
lines.append(("bold cyan", "║\n"))
|
|
1657
|
+
lines.append(("bold cyan", " ║"))
|
|
1658
|
+
lines.append(("", hint))
|
|
1659
|
+
lines.append(("bold cyan", "║\n"))
|
|
1660
|
+
lines.append(("bold cyan", f" ╚{border}╝\n"))
|
|
1661
|
+
lines.append(("", "\n"))
|
|
1662
|
+
|
|
1663
|
+
if editing[0] is not None:
|
|
1664
|
+
idx = editing[0]
|
|
1665
|
+
label = FIELDS[idx]["label"]
|
|
1666
|
+
lines.append(("bold cyan", " ❯ "))
|
|
1667
|
+
lines.append(("bold cyan", label))
|
|
1668
|
+
lines.append(("", " "))
|
|
1669
|
+
lines.append(("bold white", f"[ {edit_buf[0]}▌ ]"))
|
|
1670
|
+
lines.append(("", "\n"))
|
|
1671
|
+
lines.append(("dim",
|
|
1672
|
+
" Type a number, Enter to confirm, "
|
|
1673
|
+
"Esc to cancel\n"))
|
|
1674
|
+
return lines
|
|
1675
|
+
|
|
1676
|
+
for i, field in enumerate(FIELDS):
|
|
1677
|
+
is_sel = (i == sel[0])
|
|
1678
|
+
|
|
1679
|
+
if field["type"] == "action":
|
|
1680
|
+
if is_sel:
|
|
1681
|
+
prefix = ("bold cyan", " ❯ ")
|
|
1682
|
+
label = ("bold cyan", field["label"])
|
|
1683
|
+
else:
|
|
1684
|
+
prefix = ("", " ")
|
|
1685
|
+
label = ("dim", field["label"])
|
|
1686
|
+
lines.append(prefix)
|
|
1687
|
+
lines.append(label)
|
|
1688
|
+
lines.append(("", "\n"))
|
|
1689
|
+
continue
|
|
1690
|
+
|
|
1691
|
+
prefix = ("bold cyan", " ❯ ") if is_sel else ("", " ")
|
|
1692
|
+
label = ("bold cyan" if is_sel else "bold",
|
|
1693
|
+
field["label"])
|
|
1694
|
+
|
|
1695
|
+
val_str = _field_val(i)
|
|
1696
|
+
if field["type"] == "choice":
|
|
1697
|
+
if is_sel:
|
|
1698
|
+
val = ("bold yellow", f"◄ {val_str} ►")
|
|
1699
|
+
else:
|
|
1700
|
+
val = ("dim", val_str)
|
|
1701
|
+
else:
|
|
1702
|
+
toggle_var = _get_toggle_var(i)
|
|
1703
|
+
toggle_on = toggle_var[0] if toggle_var else False
|
|
1704
|
+
if toggle_on:
|
|
1705
|
+
val = ("bold green" if is_sel else "green",
|
|
1706
|
+
f"● {val_str}")
|
|
1707
|
+
else:
|
|
1708
|
+
val = ("dim", f"○ {val_str}")
|
|
1709
|
+
|
|
1710
|
+
lines.append(prefix)
|
|
1711
|
+
lines.append(label)
|
|
1712
|
+
lines.append(("", " "))
|
|
1713
|
+
lines.append(val)
|
|
1714
|
+
lines.append(("", "\n"))
|
|
1715
|
+
|
|
1716
|
+
return lines
|
|
1717
|
+
|
|
1718
|
+
menu = Window(
|
|
1719
|
+
content=FormattedTextControl(_get_text),
|
|
1720
|
+
dont_extend_height=True,
|
|
1721
|
+
)
|
|
1722
|
+
|
|
1723
|
+
app: Application = Application(
|
|
1724
|
+
layout=Layout(HSplit([menu])),
|
|
1725
|
+
key_bindings=kb,
|
|
1726
|
+
full_screen=False,
|
|
1727
|
+
)
|
|
1728
|
+
|
|
1729
|
+
_tty_write("\033[s")
|
|
1730
|
+
app.run()
|
|
1731
|
+
_tty_write("\033[u\033[J")
|
|
1732
|
+
|
|
1733
|
+
return result[0]
|
|
1734
|
+
|
|
1735
|
+
|
|
1736
|
+
def _select_concurrent_params(
|
|
1737
|
+
concurrency: int = 8,
|
|
1738
|
+
duration: float = 30.0,
|
|
1739
|
+
ramp_up: float = 0.0,
|
|
1740
|
+
profile: bool = True,
|
|
1741
|
+
skip_setup: bool = False,
|
|
1742
|
+
skip_teardown: bool = False,
|
|
1743
|
+
) -> Optional[dict]:
|
|
1744
|
+
"""Show parameter configuration for CONCURRENT mode."""
|
|
1745
|
+
from prompt_toolkit import Application
|
|
1746
|
+
from prompt_toolkit.key_binding import KeyBindings
|
|
1747
|
+
from prompt_toolkit.layout import Layout
|
|
1748
|
+
from prompt_toolkit.layout.containers import HSplit, Window
|
|
1749
|
+
from prompt_toolkit.layout.controls import FormattedTextControl
|
|
1750
|
+
from prompt_toolkit.keys import Keys
|
|
1751
|
+
from prompt_toolkit.filters import Condition
|
|
1752
|
+
|
|
1753
|
+
CONCURRENCY_PRESETS = [1, 2, 4, 8, 16, 32, 64]
|
|
1754
|
+
DURATION_PRESETS = [10.0, 30.0, 60.0, 120.0, 300.0]
|
|
1755
|
+
RAMPUP_PRESETS = [0.0, 1.0, 2.0, 5.0, 10.0]
|
|
1756
|
+
PROFILE_LABELS = {False: "Off", True: "On"}
|
|
1757
|
+
SKIP_LABELS = {False: "Off", True: "On"}
|
|
1758
|
+
|
|
1759
|
+
custom_concurrency = [None]
|
|
1760
|
+
custom_duration = [None]
|
|
1761
|
+
custom_rampup = [None]
|
|
1762
|
+
|
|
1763
|
+
result = [None]
|
|
1764
|
+
sel = [0]
|
|
1765
|
+
cc_idx = [CONCURRENCY_PRESETS.index(concurrency)
|
|
1766
|
+
if concurrency in CONCURRENCY_PRESETS else 3]
|
|
1767
|
+
dur_idx = [DURATION_PRESETS.index(duration)
|
|
1768
|
+
if duration in DURATION_PRESETS else 1]
|
|
1769
|
+
ramp_idx = [RAMPUP_PRESETS.index(ramp_up)
|
|
1770
|
+
if ramp_up in RAMPUP_PRESETS else 0]
|
|
1771
|
+
prof = [profile]
|
|
1772
|
+
s_setup = [skip_setup]
|
|
1773
|
+
s_teardown = [skip_teardown]
|
|
1774
|
+
|
|
1775
|
+
FIELDS = [
|
|
1776
|
+
{"label": "Concurrency (threads)", "type": "choice"},
|
|
1777
|
+
{"label": "Duration (seconds)", "type": "choice"},
|
|
1778
|
+
{"label": "Ramp-up (seconds)", "type": "choice"},
|
|
1779
|
+
{"label": "Profile (flame graph)", "type": "toggle", "var": "prof"},
|
|
1780
|
+
{"label": "Skip Setup (reuse tables)", "type": "toggle", "var": "s_setup"},
|
|
1781
|
+
{"label": "Skip Teardown (keep tables)", "type": "toggle", "var": "s_teardown"},
|
|
1782
|
+
{"label": "OK", "type": "action"},
|
|
1783
|
+
{"label": "Back", "type": "action"},
|
|
1784
|
+
{"label": "Quit", "type": "action"},
|
|
1785
|
+
]
|
|
1786
|
+
|
|
1787
|
+
ACTION_OK = len(FIELDS) - 3
|
|
1788
|
+
ACTION_BACK = len(FIELDS) - 2
|
|
1789
|
+
ACTION_QUIT = len(FIELDS) - 1
|
|
1790
|
+
|
|
1791
|
+
def _concurrency_val():
|
|
1792
|
+
if custom_concurrency[0] is not None:
|
|
1793
|
+
return custom_concurrency[0]
|
|
1794
|
+
return CONCURRENCY_PRESETS[cc_idx[0]]
|
|
1795
|
+
|
|
1796
|
+
def _duration_val():
|
|
1797
|
+
if custom_duration[0] is not None:
|
|
1798
|
+
return custom_duration[0]
|
|
1799
|
+
return DURATION_PRESETS[dur_idx[0]]
|
|
1800
|
+
|
|
1801
|
+
def _rampup_val():
|
|
1802
|
+
if custom_rampup[0] is not None:
|
|
1803
|
+
return custom_rampup[0]
|
|
1804
|
+
return RAMPUP_PRESETS[ramp_idx[0]]
|
|
1805
|
+
|
|
1806
|
+
def _field_val(i):
|
|
1807
|
+
if i == 0:
|
|
1808
|
+
v = _concurrency_val()
|
|
1809
|
+
if custom_concurrency[0] is not None:
|
|
1810
|
+
return f"{v} (custom)"
|
|
1811
|
+
return str(v)
|
|
1812
|
+
elif i == 1:
|
|
1813
|
+
v = _duration_val()
|
|
1814
|
+
if custom_duration[0] is not None:
|
|
1815
|
+
return f"{v} (custom)"
|
|
1816
|
+
return str(v)
|
|
1817
|
+
elif i == 2:
|
|
1818
|
+
v = _rampup_val()
|
|
1819
|
+
if custom_rampup[0] is not None:
|
|
1820
|
+
return f"{v} (custom)"
|
|
1821
|
+
return str(v)
|
|
1822
|
+
elif i == 3:
|
|
1823
|
+
return PROFILE_LABELS[prof[0]]
|
|
1824
|
+
elif i == 4:
|
|
1825
|
+
return SKIP_LABELS[s_setup[0]]
|
|
1826
|
+
elif i == 5:
|
|
1827
|
+
return SKIP_LABELS[s_teardown[0]]
|
|
1828
|
+
return ""
|
|
1829
|
+
|
|
1830
|
+
def _get_toggle_var(i):
|
|
1831
|
+
if i == 3: return prof
|
|
1832
|
+
if i == 4: return s_setup
|
|
1833
|
+
if i == 5: return s_teardown
|
|
1834
|
+
return None
|
|
1835
|
+
|
|
1836
|
+
def _toggle_right(i):
|
|
1837
|
+
if i == 0:
|
|
1838
|
+
if custom_concurrency[0] is not None:
|
|
1839
|
+
custom_concurrency[0] = None
|
|
1840
|
+
else:
|
|
1841
|
+
if cc_idx[0] == len(CONCURRENCY_PRESETS) - 1:
|
|
1842
|
+
custom_concurrency[0] = _concurrency_val()
|
|
1843
|
+
else:
|
|
1844
|
+
cc_idx[0] += 1
|
|
1845
|
+
elif i == 1:
|
|
1846
|
+
if custom_duration[0] is not None:
|
|
1847
|
+
custom_duration[0] = None
|
|
1848
|
+
else:
|
|
1849
|
+
if dur_idx[0] == len(DURATION_PRESETS) - 1:
|
|
1850
|
+
custom_duration[0] = _duration_val()
|
|
1851
|
+
else:
|
|
1852
|
+
dur_idx[0] += 1
|
|
1853
|
+
elif i == 2:
|
|
1854
|
+
if custom_rampup[0] is not None:
|
|
1855
|
+
custom_rampup[0] = None
|
|
1856
|
+
else:
|
|
1857
|
+
if ramp_idx[0] == len(RAMPUP_PRESETS) - 1:
|
|
1858
|
+
custom_rampup[0] = _rampup_val()
|
|
1859
|
+
else:
|
|
1860
|
+
ramp_idx[0] += 1
|
|
1861
|
+
else:
|
|
1862
|
+
var = _get_toggle_var(i)
|
|
1863
|
+
if var is not None:
|
|
1864
|
+
var[0] = not var[0]
|
|
1865
|
+
|
|
1866
|
+
def _toggle_left(i):
|
|
1867
|
+
if i == 0:
|
|
1868
|
+
if custom_concurrency[0] is not None:
|
|
1869
|
+
custom_concurrency[0] = None
|
|
1870
|
+
else:
|
|
1871
|
+
if cc_idx[0] == 0:
|
|
1872
|
+
custom_concurrency[0] = _concurrency_val()
|
|
1873
|
+
cc_idx[0] = 0
|
|
1874
|
+
else:
|
|
1875
|
+
cc_idx[0] -= 1
|
|
1876
|
+
elif i == 1:
|
|
1877
|
+
if custom_duration[0] is not None:
|
|
1878
|
+
custom_duration[0] = None
|
|
1879
|
+
else:
|
|
1880
|
+
if dur_idx[0] == 0:
|
|
1881
|
+
custom_duration[0] = _duration_val()
|
|
1882
|
+
dur_idx[0] = 0
|
|
1883
|
+
else:
|
|
1884
|
+
dur_idx[0] -= 1
|
|
1885
|
+
elif i == 2:
|
|
1886
|
+
if custom_rampup[0] is not None:
|
|
1887
|
+
custom_rampup[0] = None
|
|
1888
|
+
else:
|
|
1889
|
+
if ramp_idx[0] == 0:
|
|
1890
|
+
custom_rampup[0] = _rampup_val()
|
|
1891
|
+
ramp_idx[0] = 0
|
|
1892
|
+
else:
|
|
1893
|
+
ramp_idx[0] -= 1
|
|
1894
|
+
else:
|
|
1895
|
+
var = _get_toggle_var(i)
|
|
1896
|
+
if var is not None:
|
|
1897
|
+
var[0] = not var[0]
|
|
1898
|
+
|
|
1899
|
+
editing = [None]
|
|
1900
|
+
edit_buf = [""]
|
|
1901
|
+
|
|
1902
|
+
kb = KeyBindings()
|
|
1903
|
+
|
|
1904
|
+
@kb.add("up")
|
|
1905
|
+
@kb.add("k")
|
|
1906
|
+
def _up(event):
|
|
1907
|
+
if editing[0] is not None:
|
|
1908
|
+
return
|
|
1909
|
+
sel[0] = (sel[0] - 1) % len(FIELDS)
|
|
1910
|
+
|
|
1911
|
+
@kb.add("down")
|
|
1912
|
+
@kb.add("j")
|
|
1913
|
+
def _down(event):
|
|
1914
|
+
if editing[0] is not None:
|
|
1915
|
+
return
|
|
1916
|
+
sel[0] = (sel[0] + 1) % len(FIELDS)
|
|
1917
|
+
|
|
1918
|
+
@kb.add("left")
|
|
1919
|
+
@kb.add("h")
|
|
1920
|
+
def _left(event):
|
|
1921
|
+
if editing[0] is not None:
|
|
1922
|
+
return
|
|
1923
|
+
_toggle_left(sel[0])
|
|
1924
|
+
|
|
1925
|
+
@kb.add("right")
|
|
1926
|
+
@kb.add("l")
|
|
1927
|
+
def _right(event):
|
|
1928
|
+
if editing[0] is not None:
|
|
1929
|
+
return
|
|
1930
|
+
_toggle_right(sel[0])
|
|
1931
|
+
|
|
1932
|
+
@kb.add("backspace")
|
|
1933
|
+
def _backspace(event):
|
|
1934
|
+
if editing[0] is not None:
|
|
1935
|
+
edit_buf[0] = edit_buf[0][:-1]
|
|
1936
|
+
|
|
1937
|
+
@kb.add(Keys.Any, filter=Condition(lambda: editing[0] is not None))
|
|
1938
|
+
def _type_char(event):
|
|
1939
|
+
ch = event.data
|
|
1940
|
+
if ch.isdigit() or ch == '.':
|
|
1941
|
+
edit_buf[0] += ch
|
|
1942
|
+
|
|
1943
|
+
@kb.add("enter")
|
|
1944
|
+
def _confirm(event):
|
|
1945
|
+
if editing[0] is not None:
|
|
1946
|
+
idx = editing[0]
|
|
1947
|
+
if edit_buf[0]:
|
|
1948
|
+
try:
|
|
1949
|
+
if idx == 0:
|
|
1950
|
+
n = int(edit_buf[0])
|
|
1951
|
+
if n >= 1:
|
|
1952
|
+
custom_concurrency[0] = n
|
|
1953
|
+
elif idx == 1:
|
|
1954
|
+
n = float(edit_buf[0])
|
|
1955
|
+
if n > 0:
|
|
1956
|
+
custom_duration[0] = n
|
|
1957
|
+
elif idx == 2:
|
|
1958
|
+
n = float(edit_buf[0])
|
|
1959
|
+
if n >= 0:
|
|
1960
|
+
custom_rampup[0] = n
|
|
1961
|
+
except ValueError:
|
|
1962
|
+
pass
|
|
1963
|
+
editing[0] = None
|
|
1964
|
+
edit_buf[0] = ""
|
|
1965
|
+
return
|
|
1966
|
+
|
|
1967
|
+
if sel[0] == 0 and custom_concurrency[0] is not None:
|
|
1968
|
+
editing[0] = 0
|
|
1969
|
+
edit_buf[0] = str(custom_concurrency[0])
|
|
1970
|
+
return
|
|
1971
|
+
if sel[0] == 1 and custom_duration[0] is not None:
|
|
1972
|
+
editing[0] = 1
|
|
1973
|
+
edit_buf[0] = str(custom_duration[0])
|
|
1974
|
+
return
|
|
1975
|
+
if sel[0] == 2 and custom_rampup[0] is not None:
|
|
1976
|
+
editing[0] = 2
|
|
1977
|
+
edit_buf[0] = str(custom_rampup[0])
|
|
1978
|
+
return
|
|
1979
|
+
|
|
1980
|
+
if sel[0] == ACTION_OK:
|
|
1981
|
+
result[0] = {
|
|
1982
|
+
"mode": "concurrent",
|
|
1983
|
+
"iterations": 100,
|
|
1984
|
+
"warmup": 5,
|
|
1985
|
+
"concurrency": _concurrency_val(),
|
|
1986
|
+
"duration": _duration_val(),
|
|
1987
|
+
"ramp_up": _rampup_val(),
|
|
1988
|
+
"profile": prof[0],
|
|
1989
|
+
"skip_setup": s_setup[0],
|
|
1990
|
+
"skip_teardown": s_teardown[0],
|
|
1991
|
+
}
|
|
1992
|
+
event.app.exit()
|
|
1993
|
+
return
|
|
1994
|
+
if sel[0] == ACTION_BACK:
|
|
1995
|
+
result[0] = {"action": "back"}
|
|
1996
|
+
event.app.exit()
|
|
1997
|
+
return
|
|
1998
|
+
if sel[0] == ACTION_QUIT:
|
|
1999
|
+
result[0] = None
|
|
2000
|
+
event.app.exit()
|
|
2001
|
+
return
|
|
2002
|
+
|
|
2003
|
+
@kb.add("c-c")
|
|
2004
|
+
@kb.add("escape")
|
|
2005
|
+
def _cancel(event):
|
|
2006
|
+
if editing[0] is not None:
|
|
2007
|
+
editing[0] = None
|
|
2008
|
+
edit_buf[0] = ""
|
|
2009
|
+
return
|
|
2010
|
+
result[0] = None
|
|
2011
|
+
event.app.exit()
|
|
2012
|
+
|
|
2013
|
+
def _get_text():
|
|
2014
|
+
lines = []
|
|
2015
|
+
border = "═" * 55
|
|
2016
|
+
title = "CONCURRENT Mode Configuration".center(55)
|
|
2017
|
+
hint = ("←→ change · Enter confirm/custom · ↑↓ move"
|
|
2018
|
+
" · Esc cancel").center(55)
|
|
2019
|
+
lines.append(("bold cyan", f" ╔{border}╗\n"))
|
|
2020
|
+
lines.append(("bold cyan", " ║"))
|
|
2021
|
+
lines.append(("bold white", title))
|
|
2022
|
+
lines.append(("bold cyan", "║\n"))
|
|
2023
|
+
lines.append(("bold cyan", " ║"))
|
|
2024
|
+
lines.append(("", hint))
|
|
2025
|
+
lines.append(("bold cyan", "║\n"))
|
|
2026
|
+
lines.append(("bold cyan", f" ╚{border}╝\n"))
|
|
2027
|
+
lines.append(("", "\n"))
|
|
2028
|
+
|
|
2029
|
+
if editing[0] is not None:
|
|
2030
|
+
idx = editing[0]
|
|
2031
|
+
label = FIELDS[idx]["label"]
|
|
2032
|
+
lines.append(("bold cyan", " ❯ "))
|
|
2033
|
+
lines.append(("bold cyan", label))
|
|
2034
|
+
lines.append(("", " "))
|
|
2035
|
+
lines.append(("bold white", f"[ {edit_buf[0]}▌ ]"))
|
|
2036
|
+
lines.append(("", "\n"))
|
|
2037
|
+
lines.append(("dim",
|
|
2038
|
+
" Type a number, Enter to confirm, "
|
|
2039
|
+
"Esc to cancel\n"))
|
|
2040
|
+
return lines
|
|
2041
|
+
|
|
2042
|
+
for i, field in enumerate(FIELDS):
|
|
2043
|
+
is_sel = (i == sel[0])
|
|
2044
|
+
|
|
2045
|
+
if field["type"] == "action":
|
|
2046
|
+
if is_sel:
|
|
2047
|
+
prefix = ("bold cyan", " ❯ ")
|
|
2048
|
+
label = ("bold cyan", field["label"])
|
|
2049
|
+
else:
|
|
2050
|
+
prefix = ("", " ")
|
|
2051
|
+
label = ("dim", field["label"])
|
|
2052
|
+
lines.append(prefix)
|
|
2053
|
+
lines.append(label)
|
|
2054
|
+
lines.append(("", "\n"))
|
|
2055
|
+
continue
|
|
2056
|
+
|
|
2057
|
+
prefix = ("bold cyan", " ❯ ") if is_sel else ("", " ")
|
|
2058
|
+
label = ("bold cyan" if is_sel else "bold",
|
|
2059
|
+
field["label"])
|
|
2060
|
+
|
|
2061
|
+
val_str = _field_val(i)
|
|
2062
|
+
if field["type"] == "choice":
|
|
2063
|
+
if is_sel:
|
|
2064
|
+
val = ("bold yellow", f"◄ {val_str} ►")
|
|
2065
|
+
else:
|
|
2066
|
+
val = ("dim", val_str)
|
|
2067
|
+
else:
|
|
2068
|
+
toggle_var = _get_toggle_var(i)
|
|
2069
|
+
toggle_on = toggle_var[0] if toggle_var else False
|
|
2070
|
+
if toggle_on:
|
|
2071
|
+
val = ("bold green" if is_sel else "green",
|
|
2072
|
+
f"● {val_str}")
|
|
2073
|
+
else:
|
|
2074
|
+
val = ("dim", f"○ {val_str}")
|
|
2075
|
+
|
|
2076
|
+
lines.append(prefix)
|
|
2077
|
+
lines.append(label)
|
|
2078
|
+
lines.append(("", " "))
|
|
2079
|
+
lines.append(val)
|
|
2080
|
+
lines.append(("", "\n"))
|
|
2081
|
+
|
|
2082
|
+
return lines
|
|
2083
|
+
|
|
2084
|
+
menu = Window(
|
|
2085
|
+
content=FormattedTextControl(_get_text),
|
|
2086
|
+
dont_extend_height=True,
|
|
2087
|
+
)
|
|
2088
|
+
|
|
2089
|
+
app: Application = Application(
|
|
2090
|
+
layout=Layout(HSplit([menu])),
|
|
2091
|
+
key_bindings=kb,
|
|
2092
|
+
full_screen=False,
|
|
2093
|
+
)
|
|
2094
|
+
|
|
2095
|
+
_tty_write("\033[s")
|
|
2096
|
+
app.run()
|
|
2097
|
+
_tty_write("\033[u\033[J")
|
|
2098
|
+
|
|
2099
|
+
return result[0]
|
|
2100
|
+
|
|
2101
|
+
|
|
2102
|
+
def _select_mode(configs, database: str) -> Optional[str]:
|
|
2103
|
+
"""Show an arrow-key mode selector and return 'mtr' or 'bench'.
|
|
2104
|
+
|
|
2105
|
+
Returns ``None`` if the user cancels (Ctrl-C / Esc).
|
|
2106
|
+
"""
|
|
2107
|
+
import sys
|
|
2108
|
+
|
|
2109
|
+
from prompt_toolkit import Application
|
|
2110
|
+
from prompt_toolkit.key_binding import KeyBindings
|
|
2111
|
+
from prompt_toolkit.layout import Layout
|
|
2112
|
+
from prompt_toolkit.layout.containers import HSplit, Window
|
|
2113
|
+
from prompt_toolkit.layout.controls import FormattedTextControl
|
|
2114
|
+
|
|
2115
|
+
MODES = [
|
|
2116
|
+
("mtr", "MTR mode", "run .test compatibility tests"),
|
|
2117
|
+
("playground", "Playground mode", "run SQL Playground in browser"),
|
|
2118
|
+
("bench", "Benchmark mode", "run JSON performance benchmarks"),
|
|
2119
|
+
("history", "History mode", "view historical test runs"),
|
|
2120
|
+
(None, "Quit", "exit"),
|
|
2121
|
+
]
|
|
2122
|
+
|
|
2123
|
+
QUIT_IDX = len(MODES) - 1
|
|
2124
|
+
|
|
2125
|
+
selected = [0] # mutable index
|
|
2126
|
+
result = [None] # mutable result
|
|
2127
|
+
|
|
2128
|
+
# -- key bindings -------------------------------------------------------
|
|
2129
|
+
kb = KeyBindings()
|
|
2130
|
+
|
|
2131
|
+
@kb.add("up")
|
|
2132
|
+
@kb.add("k")
|
|
2133
|
+
def _up(event):
|
|
2134
|
+
selected[0] = (selected[0] - 1) % len(MODES)
|
|
2135
|
+
|
|
2136
|
+
@kb.add("down")
|
|
2137
|
+
@kb.add("j")
|
|
2138
|
+
def _down(event):
|
|
2139
|
+
selected[0] = (selected[0] + 1) % len(MODES)
|
|
2140
|
+
|
|
2141
|
+
@kb.add("enter")
|
|
2142
|
+
def _confirm(event):
|
|
2143
|
+
key = MODES[selected[0]][0]
|
|
2144
|
+
result[0] = key # None for Quit, 'mtr' or 'bench' otherwise
|
|
2145
|
+
event.app.exit()
|
|
2146
|
+
|
|
2147
|
+
@kb.add("c-c")
|
|
2148
|
+
@kb.add("escape")
|
|
2149
|
+
def _cancel(event):
|
|
2150
|
+
result[0] = None
|
|
2151
|
+
event.app.exit()
|
|
2152
|
+
|
|
2153
|
+
# -- layout -------------------------------------------------------------
|
|
2154
|
+
def _get_menu_text():
|
|
2155
|
+
lines = []
|
|
2156
|
+
border = "═" * 55
|
|
2157
|
+
title = "Rosetta Interactive Mode".center(55)
|
|
2158
|
+
hint = "↑/↓ to move, Enter to select, Esc to quit".center(55)
|
|
2159
|
+
lines.append(("bold cyan", f" ╔{border}╗\n"))
|
|
2160
|
+
lines.append(("bold cyan", " ║"))
|
|
2161
|
+
lines.append(("bold white", title))
|
|
2162
|
+
lines.append(("bold cyan", "║\n"))
|
|
2163
|
+
lines.append(("bold cyan", " ║"))
|
|
2164
|
+
lines.append(("", hint))
|
|
2165
|
+
lines.append(("bold cyan", "║\n"))
|
|
2166
|
+
lines.append(("bold cyan", f" ╚{border}╝\n"))
|
|
2167
|
+
lines.append(("", "\n"))
|
|
2168
|
+
|
|
2169
|
+
dbms_str = ", ".join(c.name for c in configs)
|
|
2170
|
+
lines.append(("gray", " DBMS: "))
|
|
2171
|
+
lines.append(("bold", dbms_str))
|
|
2172
|
+
lines.append(("gray", " Database: "))
|
|
2173
|
+
lines.append(("bold", database))
|
|
2174
|
+
lines.append(("", "\n\n"))
|
|
2175
|
+
|
|
2176
|
+
for i, (key, label, desc) in enumerate(MODES):
|
|
2177
|
+
is_quit = (key is None)
|
|
2178
|
+
if i == selected[0]:
|
|
2179
|
+
if is_quit:
|
|
2180
|
+
lines.append(("bold cyan", " ❯ "))
|
|
2181
|
+
lines.append(("bold cyan", label))
|
|
2182
|
+
else:
|
|
2183
|
+
lines.append(("bold cyan", " ❯ "))
|
|
2184
|
+
lines.append(("bold cyan", f"{label:<18s}"))
|
|
2185
|
+
lines.append(("cyan", f"— {desc}"))
|
|
2186
|
+
else:
|
|
2187
|
+
if is_quit:
|
|
2188
|
+
lines.append(("", " "))
|
|
2189
|
+
lines.append(("dim", label))
|
|
2190
|
+
else:
|
|
2191
|
+
lines.append(("", " "))
|
|
2192
|
+
lines.append(("", f"{label:<18s}"))
|
|
2193
|
+
lines.append(("gray", f"— {desc}"))
|
|
2194
|
+
lines.append(("", "\n"))
|
|
2195
|
+
|
|
2196
|
+
return lines
|
|
2197
|
+
|
|
2198
|
+
menu = Window(
|
|
2199
|
+
content=FormattedTextControl(_get_menu_text),
|
|
2200
|
+
dont_extend_height=True,
|
|
2201
|
+
)
|
|
2202
|
+
|
|
2203
|
+
app: Application = Application(
|
|
2204
|
+
layout=Layout(HSplit([menu])),
|
|
2205
|
+
key_bindings=kb,
|
|
2206
|
+
full_screen=False,
|
|
2207
|
+
)
|
|
2208
|
+
|
|
2209
|
+
# Save cursor, run, then restore and clear via /dev/tty
|
|
2210
|
+
_tty_write("\033[s")
|
|
2211
|
+
app.run()
|
|
2212
|
+
_tty_write("\033[u\033[J")
|
|
2213
|
+
|
|
2214
|
+
return result[0]
|
|
2215
|
+
|
|
2216
|
+
|
|
2217
|
+
def _select_rerun_run_id(output_dir: str) -> Optional[dict]:
|
|
2218
|
+
"""Show an interactive RUN ID selector for rerun mode.
|
|
2219
|
+
|
|
2220
|
+
Args:
|
|
2221
|
+
output_dir: Results directory to scan for historical runs
|
|
2222
|
+
|
|
2223
|
+
Returns:
|
|
2224
|
+
dict with run metadata if selected, None if cancelled,
|
|
2225
|
+
or {"manual": True} if user wants to manually input RUN ID
|
|
2226
|
+
"""
|
|
2227
|
+
import sys
|
|
2228
|
+
|
|
2229
|
+
from prompt_toolkit import Application
|
|
2230
|
+
from prompt_toolkit.key_binding import KeyBindings
|
|
2231
|
+
from prompt_toolkit.layout import Layout
|
|
2232
|
+
from prompt_toolkit.layout.containers import HSplit, Window
|
|
2233
|
+
from prompt_toolkit.layout.controls import FormattedTextControl
|
|
2234
|
+
|
|
2235
|
+
# Import _scan_runs from result_cmd
|
|
2236
|
+
from .cli.result_cmd import _scan_runs
|
|
2237
|
+
|
|
2238
|
+
# Scan historical runs
|
|
2239
|
+
runs = _scan_runs(output_dir)
|
|
2240
|
+
|
|
2241
|
+
# Filter only benchmark runs (type == "bench")
|
|
2242
|
+
bench_runs = [r for r in runs if r.get("type") == "bench"]
|
|
2243
|
+
|
|
2244
|
+
if not bench_runs:
|
|
2245
|
+
console.print("\n [yellow]No benchmark runs found in history.[/yellow]")
|
|
2246
|
+
console.print(" [dim]Run some benchmarks first before using RERUN mode.[/dim]\n")
|
|
2247
|
+
return None
|
|
2248
|
+
|
|
2249
|
+
# Limit to last 20 runs for display
|
|
2250
|
+
display_runs = bench_runs[:20]
|
|
2251
|
+
|
|
2252
|
+
# Build all run items
|
|
2253
|
+
ALL_RUNS = []
|
|
2254
|
+
for run in bench_runs:
|
|
2255
|
+
ALL_RUNS.append({
|
|
2256
|
+
"type": "run",
|
|
2257
|
+
"id": run.get("id", ""),
|
|
2258
|
+
"ts": run.get("timestamp", "")[:16],
|
|
2259
|
+
"workload": run.get("workload", ""),
|
|
2260
|
+
"mode": run.get("mode", ""),
|
|
2261
|
+
"data": run,
|
|
2262
|
+
})
|
|
2263
|
+
|
|
2264
|
+
# Pagination
|
|
2265
|
+
PAGE_SIZE = 15
|
|
2266
|
+
total_pages = max(1, (len(ALL_RUNS) + PAGE_SIZE - 1) // PAGE_SIZE)
|
|
2267
|
+
|
|
2268
|
+
# Column widths
|
|
2269
|
+
COL_ID = 42
|
|
2270
|
+
COL_WK = 22
|
|
2271
|
+
COL_TS = 18
|
|
2272
|
+
COL_MODE = 10
|
|
2273
|
+
|
|
2274
|
+
selected = [0] # index within current page items (runs + back)
|
|
2275
|
+
page = [0] # current page (0-based)
|
|
2276
|
+
result = [None]
|
|
2277
|
+
editing = [False]
|
|
2278
|
+
edit_buf = [""]
|
|
2279
|
+
|
|
2280
|
+
def _page_runs():
|
|
2281
|
+
"""Get run items for current page."""
|
|
2282
|
+
start = page[0] * PAGE_SIZE
|
|
2283
|
+
return ALL_RUNS[start:start + PAGE_SIZE]
|
|
2284
|
+
|
|
2285
|
+
def _page_items():
|
|
2286
|
+
"""Get all menu items for current page (runs + back)."""
|
|
2287
|
+
items = list(_page_runs())
|
|
2288
|
+
items.append({"type": "back"})
|
|
2289
|
+
return items
|
|
2290
|
+
|
|
2291
|
+
# Key bindings
|
|
2292
|
+
from prompt_toolkit.keys import Keys
|
|
2293
|
+
from prompt_toolkit.filters import Condition
|
|
2294
|
+
|
|
2295
|
+
kb = KeyBindings()
|
|
2296
|
+
|
|
2297
|
+
@kb.add("up")
|
|
2298
|
+
@kb.add("k")
|
|
2299
|
+
def _up(event):
|
|
2300
|
+
if editing[0]:
|
|
2301
|
+
return
|
|
2302
|
+
items = _page_items()
|
|
2303
|
+
selected[0] = (selected[0] - 1) % len(items)
|
|
2304
|
+
|
|
2305
|
+
@kb.add("down")
|
|
2306
|
+
@kb.add("j")
|
|
2307
|
+
def _down(event):
|
|
2308
|
+
if editing[0]:
|
|
2309
|
+
return
|
|
2310
|
+
items = _page_items()
|
|
2311
|
+
selected[0] = (selected[0] + 1) % len(items)
|
|
2312
|
+
|
|
2313
|
+
@kb.add("left")
|
|
2314
|
+
@kb.add("h")
|
|
2315
|
+
def _prev_page(event):
|
|
2316
|
+
if editing[0]:
|
|
2317
|
+
return
|
|
2318
|
+
if page[0] > 0:
|
|
2319
|
+
page[0] -= 1
|
|
2320
|
+
selected[0] = 0
|
|
2321
|
+
|
|
2322
|
+
@kb.add("right")
|
|
2323
|
+
@kb.add("l")
|
|
2324
|
+
def _next_page(event):
|
|
2325
|
+
if editing[0]:
|
|
2326
|
+
return
|
|
2327
|
+
if page[0] < total_pages - 1:
|
|
2328
|
+
page[0] += 1
|
|
2329
|
+
selected[0] = 0
|
|
2330
|
+
|
|
2331
|
+
@kb.add("/")
|
|
2332
|
+
def _start_search(event):
|
|
2333
|
+
if not editing[0]:
|
|
2334
|
+
editing[0] = True
|
|
2335
|
+
edit_buf[0] = ""
|
|
2336
|
+
|
|
2337
|
+
@kb.add("backspace")
|
|
2338
|
+
def _backspace(event):
|
|
2339
|
+
if editing[0]:
|
|
2340
|
+
edit_buf[0] = edit_buf[0][:-1]
|
|
2341
|
+
|
|
2342
|
+
@kb.add(Keys.Any, filter=Condition(lambda: editing[0]))
|
|
2343
|
+
def _type_char(event):
|
|
2344
|
+
ch = event.data
|
|
2345
|
+
if ch.isalnum() or ch in ('_', '-', '.', '/'):
|
|
2346
|
+
edit_buf[0] += ch
|
|
2347
|
+
|
|
2348
|
+
@kb.add("enter")
|
|
2349
|
+
def _confirm(event):
|
|
2350
|
+
if editing[0]:
|
|
2351
|
+
run_id = edit_buf[0].strip()
|
|
2352
|
+
if run_id:
|
|
2353
|
+
from .cli.result_cmd import _resolve_run
|
|
2354
|
+
resolved = _resolve_run(run_id, output_dir)
|
|
2355
|
+
if resolved:
|
|
2356
|
+
result[0] = resolved
|
|
2357
|
+
event.app.exit()
|
|
2358
|
+
else:
|
|
2359
|
+
edit_buf[0] = ""
|
|
2360
|
+
else:
|
|
2361
|
+
editing[0] = False
|
|
2362
|
+
edit_buf[0] = ""
|
|
2363
|
+
return
|
|
2364
|
+
|
|
2365
|
+
items = _page_items()
|
|
2366
|
+
item = items[selected[0]]
|
|
2367
|
+
if item["type"] == "run":
|
|
2368
|
+
result[0] = item["data"]
|
|
2369
|
+
event.app.exit()
|
|
2370
|
+
else: # back
|
|
2371
|
+
result[0] = None
|
|
2372
|
+
event.app.exit()
|
|
2373
|
+
|
|
2374
|
+
@kb.add("c-c")
|
|
2375
|
+
@kb.add("escape")
|
|
2376
|
+
def _cancel(event):
|
|
2377
|
+
if editing[0]:
|
|
2378
|
+
editing[0] = False
|
|
2379
|
+
edit_buf[0] = ""
|
|
2380
|
+
return
|
|
2381
|
+
result[0] = None
|
|
2382
|
+
event.app.exit()
|
|
2383
|
+
|
|
2384
|
+
# Layout
|
|
2385
|
+
def _get_menu_text():
|
|
2386
|
+
lines = []
|
|
2387
|
+
border_len = COL_ID + COL_WK + COL_TS + COL_MODE + 10
|
|
2388
|
+
border = "═" * border_len
|
|
2389
|
+
title = "Select Historical Run".center(border_len)
|
|
2390
|
+
hint = "↑↓ move · ←→ page · / search · Enter select · Esc back".center(border_len)
|
|
2391
|
+
|
|
2392
|
+
lines.append(("bold cyan", f" ╔{border}╗\n"))
|
|
2393
|
+
lines.append(("bold cyan", " ║"))
|
|
2394
|
+
lines.append(("bold white", title))
|
|
2395
|
+
lines.append(("bold cyan", "║\n"))
|
|
2396
|
+
lines.append(("bold cyan", " ║"))
|
|
2397
|
+
lines.append(("", hint))
|
|
2398
|
+
lines.append(("bold cyan", "║\n"))
|
|
2399
|
+
lines.append(("bold cyan", f" ╚{border}╝\n"))
|
|
2400
|
+
lines.append(("", "\n"))
|
|
2401
|
+
|
|
2402
|
+
# If in editing mode, show inline input
|
|
2403
|
+
if editing[0]:
|
|
2404
|
+
lines.append(("bold cyan", " ❯ "))
|
|
2405
|
+
lines.append(("bold cyan", "RUN ID"))
|
|
2406
|
+
lines.append(("", " "))
|
|
2407
|
+
lines.append(("bold white", f"[ {edit_buf[0]}▌ ]"))
|
|
2408
|
+
lines.append(("", "\n"))
|
|
2409
|
+
lines.append(("dim",
|
|
2410
|
+
" Type RUN ID, Enter to confirm, "
|
|
2411
|
+
"Esc to cancel\n"))
|
|
2412
|
+
return lines
|
|
2413
|
+
|
|
2414
|
+
# Table header
|
|
2415
|
+
hdr = f" {'RUN ID':<{COL_ID}} {'Workload':<{COL_WK}} {'Timestamp':<{COL_TS}} {'Mode':<{COL_MODE}}"
|
|
2416
|
+
lines.append(("dim bold", f" {hdr}\n"))
|
|
2417
|
+
lines.append(("dim", " " + "─" * border_len + "\n"))
|
|
2418
|
+
|
|
2419
|
+
items = _page_items()
|
|
2420
|
+
for i, item in enumerate(items):
|
|
2421
|
+
is_sel = (i == selected[0])
|
|
2422
|
+
prefix_style = "bold cyan" if is_sel else ""
|
|
2423
|
+
prefix_text = " ❯ " if is_sel else " "
|
|
2424
|
+
|
|
2425
|
+
if item["type"] == "run":
|
|
2426
|
+
rid = item["id"][:COL_ID]
|
|
2427
|
+
wk = item["workload"][:COL_WK]
|
|
2428
|
+
ts = item["ts"]
|
|
2429
|
+
mode = item["mode"]
|
|
2430
|
+
row = f"{rid:<{COL_ID}} {wk:<{COL_WK}} {ts:<{COL_TS}} {mode:<{COL_MODE}}"
|
|
2431
|
+
style = "bold cyan" if is_sel else ""
|
|
2432
|
+
lines.append((prefix_style, prefix_text))
|
|
2433
|
+
lines.append((style, row))
|
|
2434
|
+
else: # back
|
|
2435
|
+
style = "bold cyan" if is_sel else "dim"
|
|
2436
|
+
lines.append((prefix_style, prefix_text))
|
|
2437
|
+
lines.append((style, "← Back"))
|
|
2438
|
+
|
|
2439
|
+
lines.append(("", "\n"))
|
|
2440
|
+
|
|
2441
|
+
# Page indicator
|
|
2442
|
+
lines.append(("", "\n"))
|
|
2443
|
+
page_info = f"Page {page[0]+1}/{total_pages} ({len(ALL_RUNS)} runs)"
|
|
2444
|
+
lines.append(("dim", f" {page_info}\n"))
|
|
2445
|
+
|
|
2446
|
+
return lines
|
|
2447
|
+
|
|
2448
|
+
menu = Window(
|
|
2449
|
+
content=FormattedTextControl(_get_menu_text),
|
|
2450
|
+
dont_extend_height=True,
|
|
2451
|
+
)
|
|
2452
|
+
|
|
2453
|
+
app: Application = Application(
|
|
2454
|
+
layout=Layout(HSplit([menu])),
|
|
2455
|
+
key_bindings=kb,
|
|
2456
|
+
full_screen=False,
|
|
2457
|
+
)
|
|
2458
|
+
|
|
2459
|
+
_tty_write("\033[s")
|
|
2460
|
+
app.run()
|
|
2461
|
+
_tty_write("\033[u\033[J")
|
|
2462
|
+
|
|
2463
|
+
return result[0]
|
|
2464
|
+
|
|
2465
|
+
|
|
2466
|
+
def _enter_interactive(args) -> int:
|
|
2467
|
+
"""Load config and launch the interactive session.
|
|
2468
|
+
|
|
2469
|
+
When --benchmark is not specified, prompts the user to choose between
|
|
2470
|
+
MTR mode and Benchmark mode before entering the corresponding REPL.
|
|
2471
|
+
"""
|
|
2472
|
+
from .interactive import BenchInteractiveSession, InteractiveSession
|
|
2473
|
+
|
|
2474
|
+
if not os.path.isfile(args.config):
|
|
2475
|
+
print_error(f"Config file not found: {args.config}")
|
|
2476
|
+
flush_all()
|
|
2477
|
+
return 1
|
|
2478
|
+
|
|
2479
|
+
all_configs = load_config(args.config)
|
|
2480
|
+
if not all_configs:
|
|
2481
|
+
print_error(f"No databases configured in {args.config}")
|
|
2482
|
+
flush_all()
|
|
2483
|
+
return 1
|
|
2484
|
+
|
|
2485
|
+
try:
|
|
2486
|
+
configs = filter_configs(all_configs, args.dbms)
|
|
2487
|
+
except ValueError as e:
|
|
2488
|
+
print_error(str(e))
|
|
2489
|
+
flush_all()
|
|
2490
|
+
return 1
|
|
2491
|
+
|
|
2492
|
+
if not configs:
|
|
2493
|
+
print_error("No databases selected")
|
|
2494
|
+
flush_all()
|
|
2495
|
+
return 1
|
|
2496
|
+
|
|
2497
|
+
output_dir = os.path.abspath(args.output_dir)
|
|
2498
|
+
|
|
2499
|
+
# Clear terminal before entering interactive mode
|
|
2500
|
+
console.clear()
|
|
2501
|
+
|
|
2502
|
+
# ----- mode selection (skip if --benchmark already set) -----------------
|
|
2503
|
+
force_bench = getattr(args, "benchmark", False)
|
|
2504
|
+
|
|
2505
|
+
if force_bench:
|
|
2506
|
+
mode = "bench"
|
|
2507
|
+
else:
|
|
2508
|
+
mode = _select_mode(configs, args.database)
|
|
2509
|
+
if mode is None:
|
|
2510
|
+
# User cancelled
|
|
2511
|
+
console.print("\n [bold cyan]Goodbye! 👋[/bold cyan]\n")
|
|
2512
|
+
return 0
|
|
2513
|
+
|
|
2514
|
+
# ----- benchmark parameter configuration (interactive) ----------------
|
|
2515
|
+
# Only in interactive benchmark mode — prompt for iterations/warmup/profile
|
|
2516
|
+
bench_iterations = args.iterations
|
|
2517
|
+
bench_warmup = args.warmup
|
|
2518
|
+
bench_profile = getattr(args, 'profile', True)
|
|
2519
|
+
|
|
2520
|
+
# ----- launch selected session -----------------------------------------
|
|
2521
|
+
while True:
|
|
2522
|
+
if mode == "playground":
|
|
2523
|
+
# Start server and open Playground page in browser
|
|
2524
|
+
from .interactive import ReportServer, _APIHandler
|
|
2525
|
+
from .whitelist import Whitelist
|
|
2526
|
+
from .buglist import Buglist
|
|
2527
|
+
|
|
2528
|
+
whitelist = Whitelist(output_dir)
|
|
2529
|
+
buglist = Buglist(output_dir)
|
|
2530
|
+
|
|
2531
|
+
srv = ReportServer(
|
|
2532
|
+
output_dir, port=args.port,
|
|
2533
|
+
whitelist=whitelist,
|
|
2534
|
+
buglist=buglist,
|
|
2535
|
+
configs=configs,
|
|
2536
|
+
all_configs=all_configs,
|
|
2537
|
+
database=args.database,
|
|
2538
|
+
)
|
|
2539
|
+
try:
|
|
2540
|
+
srv.start()
|
|
2541
|
+
except OSError as e:
|
|
2542
|
+
print_error(f"Failed to start server: {e}")
|
|
2543
|
+
flush_all()
|
|
2544
|
+
return 1
|
|
2545
|
+
|
|
2546
|
+
pg_url = f"{srv.base_url}/playground.html"
|
|
2547
|
+
console.print(
|
|
2548
|
+
f"\n [green]●[/green] Playground: "
|
|
2549
|
+
f"[bold link={pg_url}]{pg_url}[/bold link]")
|
|
2550
|
+
# Open in IDE browser
|
|
2551
|
+
try:
|
|
2552
|
+
import subprocess as _sp
|
|
2553
|
+
_sp.Popen(["code", "--open-url", pg_url],
|
|
2554
|
+
stdout=_sp.DEVNULL, stderr=_sp.DEVNULL)
|
|
2555
|
+
except FileNotFoundError:
|
|
2556
|
+
pass
|
|
2557
|
+
|
|
2558
|
+
from prompt_toolkit import HTML as _HTML
|
|
2559
|
+
from prompt_toolkit.history import InMemoryHistory as _IMH
|
|
2560
|
+
from prompt_toolkit import PromptSession as _PS
|
|
2561
|
+
from .interactive import _PROMPT_STYLE
|
|
2562
|
+
|
|
2563
|
+
_pg_placeholder = _HTML(
|
|
2564
|
+
"<placeholder>Type 'help', 'back', or 'quit'"
|
|
2565
|
+
"</placeholder>")
|
|
2566
|
+
_pg_prompt = _HTML(
|
|
2567
|
+
'<prompt>rosetta</prompt> <path>▶</path> ')
|
|
2568
|
+
_pg_session = _PS(
|
|
2569
|
+
history=_IMH(),
|
|
2570
|
+
style=_PROMPT_STYLE,
|
|
2571
|
+
multiline=False,
|
|
2572
|
+
)
|
|
2573
|
+
|
|
2574
|
+
console.print()
|
|
2575
|
+
# Wait for user command
|
|
2576
|
+
while True:
|
|
2577
|
+
try:
|
|
2578
|
+
user_input = _pg_session.prompt(
|
|
2579
|
+
_pg_prompt,
|
|
2580
|
+
placeholder=_pg_placeholder,
|
|
2581
|
+
).strip()
|
|
2582
|
+
except (EOFError, KeyboardInterrupt):
|
|
2583
|
+
srv.stop()
|
|
2584
|
+
console.print(
|
|
2585
|
+
"\n [bold cyan]Goodbye! 👋[/bold cyan]\n")
|
|
2586
|
+
return 0
|
|
2587
|
+
|
|
2588
|
+
if not user_input:
|
|
2589
|
+
continue
|
|
2590
|
+
|
|
2591
|
+
cmd = user_input.lower()
|
|
2592
|
+
|
|
2593
|
+
if cmd in ("back", "b"):
|
|
2594
|
+
break
|
|
2595
|
+
elif cmd in ("quit", "exit", "q"):
|
|
2596
|
+
srv.stop()
|
|
2597
|
+
console.print(
|
|
2598
|
+
"\n [bold cyan]Goodbye! 👋[/bold cyan]\n")
|
|
2599
|
+
return 0
|
|
2600
|
+
elif cmd == "help":
|
|
2601
|
+
console.print(
|
|
2602
|
+
"\n [bold]Playground commands:[/bold]")
|
|
2603
|
+
console.print(
|
|
2604
|
+
f" [green]open[/green] "
|
|
2605
|
+
f"re-open playground in browser")
|
|
2606
|
+
console.print(
|
|
2607
|
+
f" [green]back[/green] "
|
|
2608
|
+
f"return to mode selection")
|
|
2609
|
+
console.print(
|
|
2610
|
+
f" [green]quit[/green] "
|
|
2611
|
+
f"exit rosetta\n")
|
|
2612
|
+
elif cmd == "open":
|
|
2613
|
+
try:
|
|
2614
|
+
_sp.Popen(["code", "--open-url", pg_url],
|
|
2615
|
+
stdout=_sp.DEVNULL,
|
|
2616
|
+
stderr=_sp.DEVNULL)
|
|
2617
|
+
console.print(
|
|
2618
|
+
f" [green]Opened:[/green] {pg_url}")
|
|
2619
|
+
except FileNotFoundError:
|
|
2620
|
+
console.print(
|
|
2621
|
+
f" [dim]URL:[/dim] {pg_url}")
|
|
2622
|
+
elif cmd:
|
|
2623
|
+
console.print(
|
|
2624
|
+
f" [yellow]Unknown command:[/yellow] {cmd}")
|
|
2625
|
+
console.print(
|
|
2626
|
+
f" [dim]Type 'help', 'back', "
|
|
2627
|
+
f"or 'quit'.[/dim]")
|
|
2628
|
+
|
|
2629
|
+
srv.stop()
|
|
2630
|
+
console.clear()
|
|
2631
|
+
mode = _select_mode(configs, args.database)
|
|
2632
|
+
if mode is None:
|
|
2633
|
+
console.print("\n [bold cyan]Goodbye! 👋[/bold cyan]\n")
|
|
2634
|
+
return 0
|
|
2635
|
+
continue
|
|
2636
|
+
|
|
2637
|
+
elif mode == "history":
|
|
2638
|
+
# Start server and show History URL
|
|
2639
|
+
from .interactive import ReportServer, _APIHandler
|
|
2640
|
+
from .whitelist import Whitelist
|
|
2641
|
+
from .buglist import Buglist
|
|
2642
|
+
|
|
2643
|
+
whitelist = Whitelist(output_dir)
|
|
2644
|
+
buglist = Buglist(output_dir)
|
|
2645
|
+
|
|
2646
|
+
srv = ReportServer(
|
|
2647
|
+
output_dir, port=args.port,
|
|
2648
|
+
whitelist=whitelist,
|
|
2649
|
+
buglist=buglist,
|
|
2650
|
+
configs=configs,
|
|
2651
|
+
all_configs=all_configs,
|
|
2652
|
+
database=args.database,
|
|
2653
|
+
)
|
|
2654
|
+
try:
|
|
2655
|
+
srv.start()
|
|
2656
|
+
except OSError as e:
|
|
2657
|
+
print_error(f"Failed to start server: {e}")
|
|
2658
|
+
flush_all()
|
|
2659
|
+
return 1
|
|
2660
|
+
|
|
2661
|
+
history_url = f"{srv.base_url}/index.html"
|
|
2662
|
+
console.print(
|
|
2663
|
+
f"\n [green]●[/green] History: "
|
|
2664
|
+
f"[bold link={history_url}]{history_url}[/bold link]")
|
|
2665
|
+
# Open in IDE browser
|
|
2666
|
+
try:
|
|
2667
|
+
import subprocess as _sp
|
|
2668
|
+
_sp.Popen(["code", "--open-url", history_url],
|
|
2669
|
+
stdout=_sp.DEVNULL, stderr=_sp.DEVNULL)
|
|
2670
|
+
except FileNotFoundError:
|
|
2671
|
+
pass
|
|
2672
|
+
|
|
2673
|
+
from prompt_toolkit import HTML as _HTML
|
|
2674
|
+
from prompt_toolkit.history import InMemoryHistory as _IMH
|
|
2675
|
+
from prompt_toolkit import PromptSession as _PS
|
|
2676
|
+
from .interactive import _PROMPT_STYLE
|
|
2677
|
+
|
|
2678
|
+
_hist_placeholder = _HTML(
|
|
2679
|
+
"<placeholder>Type 'help', 'back', or 'quit'"
|
|
2680
|
+
"</placeholder>")
|
|
2681
|
+
_hist_prompt = _HTML(
|
|
2682
|
+
'<prompt>rosetta</prompt> <path>▶</path> ')
|
|
2683
|
+
_hist_session = _PS(
|
|
2684
|
+
history=_IMH(),
|
|
2685
|
+
style=_PROMPT_STYLE,
|
|
2686
|
+
multiline=False,
|
|
2687
|
+
)
|
|
2688
|
+
|
|
2689
|
+
console.print()
|
|
2690
|
+
# Wait for user command
|
|
2691
|
+
while True:
|
|
2692
|
+
try:
|
|
2693
|
+
user_input = _hist_session.prompt(
|
|
2694
|
+
_hist_prompt,
|
|
2695
|
+
placeholder=_hist_placeholder,
|
|
2696
|
+
).strip()
|
|
2697
|
+
except (EOFError, KeyboardInterrupt):
|
|
2698
|
+
srv.stop()
|
|
2699
|
+
console.print(
|
|
2700
|
+
"\n [bold cyan]Goodbye! 👋[/bold cyan]\n")
|
|
2701
|
+
return 0
|
|
2702
|
+
|
|
2703
|
+
if not user_input:
|
|
2704
|
+
continue
|
|
2705
|
+
|
|
2706
|
+
cmd = user_input.lower()
|
|
2707
|
+
|
|
2708
|
+
if cmd in ("back", "b"):
|
|
2709
|
+
break
|
|
2710
|
+
elif cmd in ("quit", "exit", "q"):
|
|
2711
|
+
srv.stop()
|
|
2712
|
+
console.print(
|
|
2713
|
+
"\n [bold cyan]Goodbye! 👋[/bold cyan]\n")
|
|
2714
|
+
return 0
|
|
2715
|
+
elif cmd == "help":
|
|
2716
|
+
console.print(
|
|
2717
|
+
"\n [bold]History commands:[/bold]")
|
|
2718
|
+
console.print(
|
|
2719
|
+
f" [green]open[/green] "
|
|
2720
|
+
f"re-open history in browser")
|
|
2721
|
+
console.print(
|
|
2722
|
+
f" [green]back[/green] "
|
|
2723
|
+
f"return to mode selection")
|
|
2724
|
+
console.print(
|
|
2725
|
+
f" [green]quit[/green] "
|
|
2726
|
+
f"exit rosetta\n")
|
|
2727
|
+
elif cmd == "open":
|
|
2728
|
+
try:
|
|
2729
|
+
_sp.Popen(["code", "--open-url", history_url],
|
|
2730
|
+
stdout=_sp.DEVNULL,
|
|
2731
|
+
stderr=_sp.DEVNULL)
|
|
2732
|
+
console.print(
|
|
2733
|
+
f" [green]Opened:[/green] {history_url}")
|
|
2734
|
+
except FileNotFoundError:
|
|
2735
|
+
console.print(
|
|
2736
|
+
f" [dim]URL:[/dim] {history_url}")
|
|
2737
|
+
elif cmd:
|
|
2738
|
+
console.print(
|
|
2739
|
+
f" [yellow]Unknown command:[/yellow] {cmd}")
|
|
2740
|
+
console.print(
|
|
2741
|
+
f" [dim]Type 'help', 'back', "
|
|
2742
|
+
f"or 'quit'.[/dim]")
|
|
2743
|
+
|
|
2744
|
+
srv.stop()
|
|
2745
|
+
console.clear()
|
|
2746
|
+
mode = _select_mode(configs, args.database)
|
|
2747
|
+
if mode is None:
|
|
2748
|
+
console.print("\n [bold cyan]Goodbye! 👋[/bold cyan]\n")
|
|
2749
|
+
return 0
|
|
2750
|
+
continue
|
|
2751
|
+
|
|
2752
|
+
elif mode == "mtr":
|
|
2753
|
+
session = InteractiveSession(
|
|
2754
|
+
configs=configs,
|
|
2755
|
+
output_dir=output_dir,
|
|
2756
|
+
database=args.database,
|
|
2757
|
+
baseline=args.baseline,
|
|
2758
|
+
skip_explain=args.skip_explain,
|
|
2759
|
+
skip_analyze=args.skip_analyze,
|
|
2760
|
+
skip_show_create=args.skip_show_create,
|
|
2761
|
+
output_format=args.format,
|
|
2762
|
+
serve=args.serve,
|
|
2763
|
+
port=args.port,
|
|
2764
|
+
all_configs=all_configs,
|
|
2765
|
+
)
|
|
2766
|
+
reason = session.run()
|
|
2767
|
+
# Stop the report server before leaving this session
|
|
2768
|
+
# so the port is released for the next session.
|
|
2769
|
+
if session._report_server:
|
|
2770
|
+
session._report_server.stop()
|
|
2771
|
+
if reason != "back":
|
|
2772
|
+
break
|
|
2773
|
+
console.clear()
|
|
2774
|
+
mode = _select_mode(configs, args.database)
|
|
2775
|
+
if mode is None:
|
|
2776
|
+
console.print("\n [bold cyan]Goodbye! 👋[/bold cyan]\n")
|
|
2777
|
+
return 0
|
|
2778
|
+
continue
|
|
2779
|
+
else:
|
|
2780
|
+
# --- benchmark: mode → params → repl (loop params ↔ repl) ---
|
|
2781
|
+
back_to_mode = False
|
|
2782
|
+
# Initialize bench params from CLI args
|
|
2783
|
+
bench_mode = "serial" if args.concurrency == 0 else "concurrent"
|
|
2784
|
+
bench_concurrency = args.concurrency if args.concurrency > 0 else 8
|
|
2785
|
+
bench_duration = args.duration
|
|
2786
|
+
bench_ramp_up = args.ramp_up
|
|
2787
|
+
while True:
|
|
2788
|
+
if not force_bench:
|
|
2789
|
+
params = _select_bench_params(
|
|
2790
|
+
iterations=bench_iterations,
|
|
2791
|
+
warmup=bench_warmup,
|
|
2792
|
+
concurrency=bench_concurrency,
|
|
2793
|
+
duration=bench_duration,
|
|
2794
|
+
ramp_up=bench_ramp_up,
|
|
2795
|
+
profile=bench_profile,
|
|
2796
|
+
skip_setup=getattr(args, 'skip_setup', False),
|
|
2797
|
+
skip_teardown=getattr(args, 'skip_teardown', False),
|
|
2798
|
+
output_dir=output_dir,
|
|
2799
|
+
)
|
|
2800
|
+
if params is None:
|
|
2801
|
+
console.print(
|
|
2802
|
+
"\n [bold cyan]Goodbye! 👋[/bold cyan]\n")
|
|
2803
|
+
return 0
|
|
2804
|
+
if params.get("action") == "back":
|
|
2805
|
+
# Back to mode selection
|
|
2806
|
+
console.clear()
|
|
2807
|
+
mode = _select_mode(configs, args.database)
|
|
2808
|
+
if mode is None:
|
|
2809
|
+
console.print(
|
|
2810
|
+
"\n [bold cyan]Goodbye! 👋[/bold cyan]\n")
|
|
2811
|
+
return 0
|
|
2812
|
+
back_to_mode = True
|
|
2813
|
+
break # exit inner loop
|
|
2814
|
+
|
|
2815
|
+
# Handle RERUN mode cancellation (user pressed Esc in rerun selection)
|
|
2816
|
+
if params.get("action") == "cancel":
|
|
2817
|
+
console.clear()
|
|
2818
|
+
continue # Re-show Benchmark Mode selection
|
|
2819
|
+
|
|
2820
|
+
# Handle RERUN mode
|
|
2821
|
+
if params.get("mode") == "rerun":
|
|
2822
|
+
run_selection = params.get("run_data")
|
|
2823
|
+
|
|
2824
|
+
if not run_selection:
|
|
2825
|
+
console.clear()
|
|
2826
|
+
continue # Back to Benchmark Mode selection
|
|
2827
|
+
|
|
2828
|
+
# Load bench_result.json
|
|
2829
|
+
run_path = run_selection.get("path", "")
|
|
2830
|
+
bench_json_path = os.path.join(run_path, "bench_result.json")
|
|
2831
|
+
|
|
2832
|
+
if not os.path.isfile(bench_json_path):
|
|
2833
|
+
console.print(f"\n [red]✗ bench_result.json not found in:[/red] {run_path}")
|
|
2834
|
+
console.clear()
|
|
2835
|
+
continue # Back to Benchmark Mode selection
|
|
2836
|
+
|
|
2837
|
+
# Load parameters
|
|
2838
|
+
import json as _json
|
|
2839
|
+
try:
|
|
2840
|
+
with open(bench_json_path, 'r', encoding='utf-8') as f:
|
|
2841
|
+
run_data = _json.load(f)
|
|
2842
|
+
except Exception as e:
|
|
2843
|
+
console.print(f"\n [red]✗ Failed to load bench_result.json:[/red] {e}")
|
|
2844
|
+
console.clear()
|
|
2845
|
+
continue # Back to Benchmark Mode selection
|
|
2846
|
+
|
|
2847
|
+
# Extract parameters
|
|
2848
|
+
rerun_bench_file = run_data.get("bench_file") or ""
|
|
2849
|
+
rerun_database = run_data.get("database") or args.database
|
|
2850
|
+
mode_str = run_data.get("mode", "SERIAL")
|
|
2851
|
+
config_data = run_data.get("config", {})
|
|
2852
|
+
workload_name = run_data.get("workload", "rerun")
|
|
2853
|
+
|
|
2854
|
+
# Determine effective bench file:
|
|
2855
|
+
# 1) Use saved bench_file if it still exists
|
|
2856
|
+
# 2) Otherwise reconstruct from saved SQL data
|
|
2857
|
+
temp_bench_file = None
|
|
2858
|
+
if rerun_bench_file and os.path.isfile(rerun_bench_file):
|
|
2859
|
+
effective_bench_file = rerun_bench_file
|
|
2860
|
+
else:
|
|
2861
|
+
queries_sql = run_data.get("queries_sql", [])
|
|
2862
|
+
setup_sql = run_data.get("setup_sql", [])
|
|
2863
|
+
teardown_sql = run_data.get("teardown_sql", [])
|
|
2864
|
+
if not queries_sql:
|
|
2865
|
+
console.print("\n [red]✗ No query data in bench_result.json[/red]")
|
|
2866
|
+
console.clear()
|
|
2867
|
+
continue
|
|
2868
|
+
import tempfile
|
|
2869
|
+
reconstructed = {
|
|
2870
|
+
"name": workload_name,
|
|
2871
|
+
"database": rerun_database,
|
|
2872
|
+
"setup": setup_sql,
|
|
2873
|
+
"teardown": teardown_sql,
|
|
2874
|
+
"queries": [],
|
|
2875
|
+
}
|
|
2876
|
+
for q in queries_sql:
|
|
2877
|
+
if isinstance(q, dict):
|
|
2878
|
+
reconstructed["queries"].append({
|
|
2879
|
+
"name": q.get("name", ""),
|
|
2880
|
+
"sql": q.get("sql", ""),
|
|
2881
|
+
"weight": q.get("weight", 1),
|
|
2882
|
+
"description": q.get("description", ""),
|
|
2883
|
+
"cleanup_sql": q.get("cleanup_sql", ""),
|
|
2884
|
+
})
|
|
2885
|
+
elif isinstance(q, str):
|
|
2886
|
+
reconstructed["queries"].append({
|
|
2887
|
+
"name": f"q{len(reconstructed['queries'])+1}",
|
|
2888
|
+
"sql": q,
|
|
2889
|
+
})
|
|
2890
|
+
fd, temp_bench_file = tempfile.mkstemp(
|
|
2891
|
+
suffix=".json", prefix=f"rerun_{workload_name}_")
|
|
2892
|
+
with os.fdopen(fd, 'w', encoding='utf-8') as tf:
|
|
2893
|
+
_json.dump(reconstructed, tf, ensure_ascii=False, indent=2)
|
|
2894
|
+
effective_bench_file = temp_bench_file
|
|
2895
|
+
|
|
2896
|
+
# Build session and execute
|
|
2897
|
+
bench_mode_val = "concurrent" if mode_str == "CONCURRENT" else "serial"
|
|
2898
|
+
rr_iter = config_data.get("iterations", 100)
|
|
2899
|
+
rr_warmup = config_data.get("warmup", 5)
|
|
2900
|
+
rr_conc = config_data.get("concurrency", 8) if bench_mode_val == "concurrent" else 0
|
|
2901
|
+
rr_dur = config_data.get("duration", 30.0)
|
|
2902
|
+
rr_fq = config_data.get("filter_queries", [])
|
|
2903
|
+
rr_filter = ",".join(rr_fq) if rr_fq else None
|
|
2904
|
+
|
|
2905
|
+
# Display rerun configuration
|
|
2906
|
+
console.print(f"\n [bold cyan]Rerun Configuration:[/bold cyan]")
|
|
2907
|
+
console.print(f" [dim]RUN ID:[/dim] [bold]{run_selection.get('id', '')}[/bold]")
|
|
2908
|
+
console.print(f" [dim]Workload:[/dim] [bold]{workload_name}[/bold]")
|
|
2909
|
+
console.print(f" [dim]Mode:[/dim] [bold]{mode_str}[/bold]")
|
|
2910
|
+
console.print(f" [dim]Database:[/dim] [bold]{rerun_database}[/bold]")
|
|
2911
|
+
if bench_mode_val == "serial":
|
|
2912
|
+
console.print(f" [dim]Iterations:[/dim] [bold]{rr_iter}[/bold]")
|
|
2913
|
+
console.print(f" [dim]Warmup:[/dim] [bold]{rr_warmup}[/bold]")
|
|
2914
|
+
else:
|
|
2915
|
+
console.print(f" [dim]Concurrency:[/dim][bold]{rr_conc}[/bold]")
|
|
2916
|
+
console.print(f" [dim]Duration:[/dim] [bold]{rr_dur}s[/bold]")
|
|
2917
|
+
if rr_fq:
|
|
2918
|
+
console.print(f" [dim]Filter:[/dim] [bold]{', '.join(rr_fq)}[/bold]")
|
|
2919
|
+
if temp_bench_file:
|
|
2920
|
+
console.print(f" [dim]Source:[/dim] [bold]reconstructed from bench_result.json[/bold]")
|
|
2921
|
+
else:
|
|
2922
|
+
console.print(f" [dim]File:[/dim] [bold]{rerun_bench_file}[/bold]")
|
|
2923
|
+
|
|
2924
|
+
rr_session = BenchInteractiveSession(
|
|
2925
|
+
configs=configs,
|
|
2926
|
+
output_dir=output_dir,
|
|
2927
|
+
database=rerun_database,
|
|
2928
|
+
iterations=rr_iter,
|
|
2929
|
+
warmup=rr_warmup,
|
|
2930
|
+
concurrency=rr_conc,
|
|
2931
|
+
duration=rr_dur,
|
|
2932
|
+
ramp_up=0.0,
|
|
2933
|
+
bench_filter=rr_filter,
|
|
2934
|
+
repeat=1,
|
|
2935
|
+
parallel_dbms=True,
|
|
2936
|
+
output_format=args.format,
|
|
2937
|
+
serve=args.serve,
|
|
2938
|
+
port=args.port,
|
|
2939
|
+
profile=False,
|
|
2940
|
+
perf_freq=getattr(args, 'perf_freq', 99),
|
|
2941
|
+
flamegraph_min_ms=getattr(args, 'flamegraph_min_ms', 1000),
|
|
2942
|
+
bench_mode=bench_mode_val,
|
|
2943
|
+
)
|
|
2944
|
+
|
|
2945
|
+
console.print()
|
|
2946
|
+
rr_session._run_bench(effective_bench_file)
|
|
2947
|
+
|
|
2948
|
+
# Cleanup temp file
|
|
2949
|
+
if temp_bench_file and os.path.isfile(temp_bench_file):
|
|
2950
|
+
try:
|
|
2951
|
+
os.unlink(temp_bench_file)
|
|
2952
|
+
except OSError:
|
|
2953
|
+
pass
|
|
2954
|
+
|
|
2955
|
+
console.print("\n [dim]Press Enter to continue...[/dim]")
|
|
2956
|
+
try:
|
|
2957
|
+
input()
|
|
2958
|
+
except (EOFError, KeyboardInterrupt):
|
|
2959
|
+
pass
|
|
2960
|
+
|
|
2961
|
+
console.clear()
|
|
2962
|
+
continue
|
|
2963
|
+
|
|
2964
|
+
# Normal benchmark mode
|
|
2965
|
+
bench_mode = params["mode"]
|
|
2966
|
+
bench_iterations = params["iterations"]
|
|
2967
|
+
bench_warmup = params["warmup"]
|
|
2968
|
+
bench_concurrency = params["concurrency"]
|
|
2969
|
+
bench_duration = params["duration"]
|
|
2970
|
+
bench_ramp_up = params["ramp_up"]
|
|
2971
|
+
bench_profile = params["profile"]
|
|
2972
|
+
bench_skip_setup = params.get("skip_setup", False)
|
|
2973
|
+
bench_skip_teardown = params.get("skip_teardown", False)
|
|
2974
|
+
else:
|
|
2975
|
+
bench_skip_setup = getattr(args, 'skip_setup', False)
|
|
2976
|
+
bench_skip_teardown = getattr(args, 'skip_teardown', False)
|
|
2977
|
+
|
|
2978
|
+
session = BenchInteractiveSession(
|
|
2979
|
+
configs=configs,
|
|
2980
|
+
output_dir=output_dir,
|
|
2981
|
+
database=args.database,
|
|
2982
|
+
iterations=bench_iterations,
|
|
2983
|
+
warmup=bench_warmup,
|
|
2984
|
+
concurrency=bench_concurrency if bench_mode == "concurrent" else 0,
|
|
2985
|
+
duration=bench_duration,
|
|
2986
|
+
ramp_up=bench_ramp_up,
|
|
2987
|
+
bench_filter=args.bench_filter,
|
|
2988
|
+
repeat=getattr(args, 'repeat', 1),
|
|
2989
|
+
parallel_dbms=getattr(args, 'parallel_dbms', True),
|
|
2990
|
+
output_format=args.format,
|
|
2991
|
+
serve=args.serve,
|
|
2992
|
+
port=args.port,
|
|
2993
|
+
profile=bench_profile,
|
|
2994
|
+
perf_freq=getattr(args, 'perf_freq', 99),
|
|
2995
|
+
flamegraph_min_ms=getattr(args, 'flamegraph_min_ms', 1000),
|
|
2996
|
+
bench_mode=bench_mode,
|
|
2997
|
+
)
|
|
2998
|
+
session.skip_setup = bench_skip_setup
|
|
2999
|
+
session.skip_teardown = bench_skip_teardown
|
|
3000
|
+
reason = session.run()
|
|
3001
|
+
# Stop the report server before leaving this session
|
|
3002
|
+
# so the port is released for the next session.
|
|
3003
|
+
if session._report_server:
|
|
3004
|
+
session._report_server.stop()
|
|
3005
|
+
if reason == "quit":
|
|
3006
|
+
return 0
|
|
3007
|
+
if reason != "back":
|
|
3008
|
+
break
|
|
3009
|
+
# Back to bench params
|
|
3010
|
+
console.clear()
|
|
3011
|
+
continue # re-show _select_bench_params
|
|
3012
|
+
|
|
3013
|
+
if back_to_mode:
|
|
3014
|
+
continue # re-evaluate mode in outer loop
|
|
3015
|
+
break # done
|
|
3016
|
+
|
|
3017
|
+
return 0
|
|
3018
|
+
|
|
3019
|
+
|
|
3020
|
+
def _find_free_port() -> int:
|
|
3021
|
+
"""Find a free port on localhost."""
|
|
3022
|
+
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
|
3023
|
+
s.bind(("", 0))
|
|
3024
|
+
return s.getsockname()[1]
|
|
3025
|
+
|
|
3026
|
+
|
|
3027
|
+
def _serve_report(directory: str, html_file: str, port: int = 0,
|
|
3028
|
+
whitelist=None, buglist=None, configs=None,
|
|
3029
|
+
database: str = ""):
|
|
3030
|
+
"""Start a local HTTP server and print the URL for the HTML report."""
|
|
3031
|
+
if port == 0:
|
|
3032
|
+
port = _find_free_port()
|
|
3033
|
+
|
|
3034
|
+
abs_dir = os.path.abspath(directory)
|
|
3035
|
+
|
|
3036
|
+
# Pre-generate playground page
|
|
3037
|
+
from .reporter.history import generate_playground_html
|
|
3038
|
+
generate_playground_html(abs_dir)
|
|
3039
|
+
|
|
3040
|
+
# Use the API-capable handler from interactive module if whitelist given
|
|
3041
|
+
if whitelist is not None:
|
|
3042
|
+
from .interactive import _APIHandler
|
|
3043
|
+
_APIHandler._whitelist = whitelist
|
|
3044
|
+
_APIHandler._buglist = buglist
|
|
3045
|
+
_APIHandler._configs = configs or []
|
|
3046
|
+
_APIHandler._database = database
|
|
3047
|
+
handler = lambda *a, **kw: _APIHandler(
|
|
3048
|
+
*a, directory=abs_dir, **kw)
|
|
3049
|
+
else:
|
|
3050
|
+
handler = lambda *a, **kw: _NoCacheHandler(
|
|
3051
|
+
*a, directory=abs_dir, **kw)
|
|
3052
|
+
|
|
3053
|
+
try:
|
|
3054
|
+
server = _SilentHTTPServer(("0.0.0.0", port), handler)
|
|
3055
|
+
except OSError as e:
|
|
3056
|
+
print_error(f"Failed to start HTTP server on port {port}: {e}")
|
|
3057
|
+
return
|
|
3058
|
+
|
|
3059
|
+
url = f"http://localhost:{port}/{html_file}"
|
|
3060
|
+
index_url = f"http://localhost:{port}/index.html"
|
|
3061
|
+
print_server_info(url, abs_dir, history_url=index_url)
|
|
3062
|
+
|
|
3063
|
+
# Run server in a background thread so KeyboardInterrupt
|
|
3064
|
+
# can be caught without deadlocking serve_forever().
|
|
3065
|
+
server_thread = threading.Thread(target=server.serve_forever, daemon=True)
|
|
3066
|
+
server_thread.start()
|
|
3067
|
+
|
|
3068
|
+
# Try to open the URL in the IDE's built-in Simple Browser.
|
|
3069
|
+
# Works in VS Code / CloudStudio / CodeBuddy environments.
|
|
3070
|
+
try:
|
|
3071
|
+
subprocess.Popen(
|
|
3072
|
+
["code", "--open-url", url],
|
|
3073
|
+
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
|
|
3074
|
+
)
|
|
3075
|
+
except FileNotFoundError:
|
|
3076
|
+
pass # 'code' CLI not available, skip
|
|
3077
|
+
|
|
3078
|
+
try:
|
|
3079
|
+
# Block main thread until interrupted
|
|
3080
|
+
server_thread.join()
|
|
3081
|
+
except KeyboardInterrupt:
|
|
3082
|
+
pass
|
|
3083
|
+
finally:
|
|
3084
|
+
console.print("\n[dim]Shutting down server...[/dim]")
|
|
3085
|
+
# Run shutdown in a separate thread to avoid blocking forever.
|
|
3086
|
+
t = threading.Thread(target=server.shutdown, daemon=True)
|
|
3087
|
+
t.start()
|
|
3088
|
+
t.join(timeout=3)
|
|
3089
|
+
# server_thread is daemon=True, so it will be cleaned up on exit.
|