Skip to content

Commit 5ba11a1

Browse files
committed
add snapshot cli command
allow a user to snapshot a single tank (with a filter) or snapshot all tanks with a filter. this creates a zip that can then be loaded via the initContainer in the first commit
1 parent dba4a13 commit 5ba11a1

File tree

2 files changed

+105
-1
lines changed

2 files changed

+105
-1
lines changed

src/warnet/control.py

Lines changed: 103 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import base64
22
import json
3+
import os
34
import subprocess
45
import time
56
from concurrent.futures import ThreadPoolExecutor, as_completed
@@ -18,6 +19,7 @@
1819
get_default_namespace,
1920
get_mission,
2021
get_pods,
22+
snapshot_bitcoin_datadir,
2123
)
2224
from .process import run_command, stream_command
2325

@@ -291,3 +293,104 @@ def logs(pod_name: str, follow: bool):
291293
print(f"Please consider waiting for the pod to become available. Encountered: {e}")
292294
else:
293295
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+
)

src/warnet/main.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22

33
from .admin import admin
44
from .bitcoin import bitcoin
5-
from .control import down, logs, run, stop
5+
from .control import down, logs, run, snapshot, stop
66
from .deploy import deploy
77
from .graph import create, graph
88
from .image import image
@@ -28,6 +28,7 @@ def cli():
2828
cli.add_command(new)
2929
cli.add_command(run)
3030
cli.add_command(setup)
31+
cli.add_command(snapshot)
3132
cli.add_command(status)
3233
cli.add_command(stop)
3334
cli.add_command(create)

0 commit comments

Comments
 (0)