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.
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.