|
1 | 1 | import base64
|
2 | 2 | import json
|
| 3 | +import os |
3 | 4 | import subprocess
|
4 | 5 | import time
|
5 | 6 | from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
18 | 19 | get_default_namespace,
|
19 | 20 | get_mission,
|
20 | 21 | get_pods,
|
| 22 | + snapshot_bitcoin_datadir, |
21 | 23 | )
|
22 | 24 | from .process import run_command, stream_command
|
23 | 25 |
|
@@ -291,3 +293,104 @@ def logs(pod_name: str, follow: bool):
|
291 | 293 | print(f"Please consider waiting for the pod to become available. Encountered: {e}")
|
292 | 294 | else:
|
293 | 295 | pass # cancelled by user
|
| 296 | + |
| 297 | + |
| 298 | +@click.command() |
| 299 | +@click.argument("tank_name", required=False) |
| 300 | +@click.option("--all", "-a", "snapshot_all", is_flag=True, help="Snapshot all running tanks") |
| 301 | +@click.option( |
| 302 | + "--output", |
| 303 | + "-o", |
| 304 | + type=click.Path(), |
| 305 | + default="./warnet-snapshots", |
| 306 | + help="Output directory for snapshots", |
| 307 | +) |
| 308 | +@click.option( |
| 309 | + "--filter", |
| 310 | + "-f", |
| 311 | + type=str, |
| 312 | + help="Comma-separated list of directories and/or files to include in the snapshot", |
| 313 | +) |
| 314 | +def snapshot(tank_name, snapshot_all, output, filter): |
| 315 | + """Create a snapshot of a tank's Bitcoin data or snapshot all tanks""" |
| 316 | + tanks = get_mission("tank") |
| 317 | + |
| 318 | + if not tanks: |
| 319 | + console.print("[bold red]No active tanks found.[/bold red]") |
| 320 | + return |
| 321 | + |
| 322 | + # Create the output directory if it doesn't exist |
| 323 | + os.makedirs(output, exist_ok=True) |
| 324 | + |
| 325 | + filter_list = [f.strip() for f in filter.split(",")] if filter else None |
| 326 | + if snapshot_all: |
| 327 | + snapshot_all_tanks(tanks, output, filter_list) |
| 328 | + elif tank_name: |
| 329 | + snapshot_single_tank(tank_name, tanks, output, filter_list) |
| 330 | + else: |
| 331 | + select_and_snapshot_tank(tanks, output, filter_list) |
| 332 | + |
| 333 | + |
| 334 | +def find_tank_by_name(tanks, tank_name): |
| 335 | + for tank in tanks: |
| 336 | + if tank.metadata.name == tank_name: |
| 337 | + return tank |
| 338 | + return None |
| 339 | + |
| 340 | + |
| 341 | +def snapshot_all_tanks(tanks, output_dir, filter_list): |
| 342 | + with console.status("[bold yellow]Snapshotting all tanks...[/bold yellow]"): |
| 343 | + for tank in tanks: |
| 344 | + tank_name = tank.metadata.name |
| 345 | + chain = tank.metadata.labels["chain"] |
| 346 | + snapshot_tank(tank_name, chain, output_dir, filter_list) |
| 347 | + console.print("[bold green]All tank snapshots completed.[/bold green]") |
| 348 | + |
| 349 | + |
| 350 | +def snapshot_single_tank(tank_name, tanks, output_dir, filter_list): |
| 351 | + tank = find_tank_by_name(tanks, tank_name) |
| 352 | + if tank: |
| 353 | + chain = tank.metadata.labels["chain"] |
| 354 | + snapshot_tank(tank_name, chain, output_dir, filter_list) |
| 355 | + else: |
| 356 | + console.print(f"[bold red]No active tank found with name: {tank_name}[/bold red]") |
| 357 | + |
| 358 | + |
| 359 | +def select_and_snapshot_tank(tanks, output_dir, filter_list): |
| 360 | + table = Table(title="Active Tanks", show_header=True, header_style="bold magenta") |
| 361 | + table.add_column("Number", style="cyan", justify="right") |
| 362 | + table.add_column("Tank Name", style="green") |
| 363 | + |
| 364 | + for idx, tank in enumerate(tanks, 1): |
| 365 | + table.add_row(str(idx), tank.metadata.name) |
| 366 | + |
| 367 | + console.print(table) |
| 368 | + |
| 369 | + choices = [str(i) for i in range(1, len(tanks) + 1)] + ["q"] |
| 370 | + choice = Prompt.ask( |
| 371 | + "[bold yellow]Enter the number of the tank to snapshot, or 'q' to quit[/bold yellow]", |
| 372 | + choices=choices, |
| 373 | + show_choices=False, |
| 374 | + ) |
| 375 | + |
| 376 | + if choice == "q": |
| 377 | + console.print("[bold blue]Operation cancelled.[/bold blue]") |
| 378 | + return |
| 379 | + |
| 380 | + selected_tank = tanks[int(choice) - 1] |
| 381 | + tank_name = selected_tank.metadata.name |
| 382 | + chain = selected_tank.metadata.labels["chain"] |
| 383 | + snapshot_tank(tank_name, chain, output_dir, filter_list) |
| 384 | + |
| 385 | + |
| 386 | +def snapshot_tank(tank_name, chain, output_dir, filter_list): |
| 387 | + try: |
| 388 | + output_path = Path(output_dir).resolve() |
| 389 | + snapshot_bitcoin_datadir(tank_name, chain, str(output_path), filter_list) |
| 390 | + console.print( |
| 391 | + f"[bold green]Successfully created snapshot for tank: {tank_name}[/bold green]" |
| 392 | + ) |
| 393 | + except Exception as e: |
| 394 | + console.print( |
| 395 | + f"[bold red]Failed to create snapshot for tank {tank_name}: {str(e)}[/bold red]" |
| 396 | + ) |
0 commit comments