Skip to content

Commit

Permalink
add snapshot cli command
Browse files Browse the repository at this point in the history
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
  • Loading branch information
josibake committed Sep 10, 2024
1 parent dba4a13 commit cc1b6dc
Show file tree
Hide file tree
Showing 2 changed files with 97 additions and 1 deletion.
95 changes: 95 additions & 0 deletions src/warnet/control.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
get_default_namespace,
get_mission,
get_pods,
snapshot_bitcoin_datadir,
)
from .process import run_command, stream_command

Expand Down Expand Up @@ -291,3 +292,97 @@ def logs(pod_name: str, follow: bool):
print(f"Please consider waiting for the pod to become available. Encountered: {e}")
else:
pass # cancelled by user


@click.command()
@click.argument("tank_name", required=False)
@click.option("--all", "snapshot_all", is_flag=True, help="Snapshot all running tanks")
@click.option(
"--output", "-o", type=click.Path(), default=".", help="Output directory for snapshots"
)
@click.option(
"--filter",
"-f",
type=str,
help="Comma-separated list of directories and/or files to include in the snapshot",
)
def snapshot(tank_name, snapshot_all, output, filter):
"""Create a snapshot of a tank's Bitcoin data or snapshot all tanks"""
tanks = get_mission("tank")

if not tanks:
console.print("[bold red]No active tanks found.[/bold red]")
return

filter_list = [f.strip() for f in filter.split(",")] if filter else None
if snapshot_all:
snapshot_all_tanks(tanks, output, filter_list)
elif tank_name:
snapshot_single_tank(tank_name, tanks, output, filter_list)
else:
select_and_snapshot_tank(tanks, output, filter_list)


def find_tank_by_name(tanks, tank_name):
for tank in tanks:
if tank.metadata.name == tank_name:
return tank
return None


def snapshot_all_tanks(tanks, output_dir, filter_list):
with console.status("[bold yellow]Snapshotting all tanks...[/bold yellow]"):
for tank in tanks:
tank_name = tank.metadata.name
chain = tank.metadata.labels["chain"]
snapshot_tank(tank_name, chain, output_dir, filter_list)
console.print("[bold green]All tank snapshots completed.[/bold green]")


def snapshot_single_tank(tank_name, tanks, output_dir, filter_list):
tank = find_tank_by_name(tanks, tank_name)
if tank:
chain = tank.metadata.labels["chain"]
snapshot_tank(tank_name, chain, output_dir, filter_list)
else:
console.print(f"[bold red]No active tank found with name: {tank_name}[/bold red]")


def select_and_snapshot_tank(tanks, output_dir, filter_list):
table = Table(title="Active Tanks", show_header=True, header_style="bold magenta")
table.add_column("Number", style="cyan", justify="right")
table.add_column("Tank Name", style="green")

for idx, tank in enumerate(tanks, 1):
table.add_row(str(idx), tank.metadata.name)

console.print(table)

choices = [str(i) for i in range(1, len(tanks) + 1)] + ["q"]
choice = Prompt.ask(
"[bold yellow]Enter the number of the tank to snapshot, or 'q' to quit[/bold yellow]",
choices=choices,
show_choices=False,
)

if choice == "q":
console.print("[bold blue]Operation cancelled.[/bold blue]")
return

selected_tank = tanks[int(choice) - 1]
tank_name = selected_tank.metadata.name
chain = selected_tank.metadata.labels["chain"]
snapshot_tank(tank_name, chain, output_dir, filter_list)


def snapshot_tank(tank_name, chain, output_dir, filter_list):
try:
output_path = Path(output_dir).resolve()
snapshot_bitcoin_datadir(tank_name, chain, str(output_path), filter_list)
console.print(
f"[bold green]Successfully created snapshot for tank: {tank_name}[/bold green]"
)
except Exception as e:
console.print(
f"[bold red]Failed to create snapshot for tank {tank_name}: {str(e)}[/bold red]"
)
3 changes: 2 additions & 1 deletion src/warnet/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@

from .admin import admin
from .bitcoin import bitcoin
from .control import down, logs, run, stop
from .control import down, logs, run, snapshot, stop
from .deploy import deploy
from .graph import create, graph
from .image import image
Expand All @@ -28,6 +28,7 @@ def cli():
cli.add_command(new)
cli.add_command(run)
cli.add_command(setup)
cli.add_command(snapshot)
cli.add_command(status)
cli.add_command(stop)
cli.add_command(create)
Expand Down

0 comments on commit cc1b6dc

Please sign in to comment.