From 1040c0ba7fb3e4ebd353bcd68f85b34021841842 Mon Sep 17 00:00:00 2001 From: Grzegorz Wierzowiecki Date: Mon, 23 Mar 2026 23:40:42 +0100 Subject: [PATCH 1/3] Add --from-snapshot flag to create droplets from snapshots MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Support creating droplets from existing snapshots without cloud-init. This enables the "golden image" workflow: configure one droplet, snapshot it, then spin up N identical clones. When --from-snapshot is used: - Uses create_droplet_from_snapshot() API (no user_data sent) - Skips cloud-init rendering and completion monitoring - Still performs: wait for active, SSH config, project assignment, optional Tailscale setup Cloud-init is skipped because the current template is not idempotent — re-running it on a snapshot causes user creation failures, .zshrc overwrites, and an unconditional reboot. See issue #52 for discussion of alternative approaches. (for Github WebUI issue linking: Closes #52 ) Co-Authored-By: Claude Opus 4.6 (1M context) --- dropkit/main.py | 122 +++++++++++++++++++++++++++++++++--------------- 1 file changed, 84 insertions(+), 38 deletions(-) diff --git a/dropkit/main.py b/dropkit/main.py index 0d426f9..f3a2e81 100644 --- a/dropkit/main.py +++ b/dropkit/main.py @@ -1861,6 +1861,11 @@ def create( help="Project name or ID to assign droplet to", autocompletion=complete_project_name, ), + from_snapshot: int | None = typer.Option( + None, + "--from-snapshot", + help="Create from snapshot ID instead of base image (skips cloud-init)", + ), no_tailscale: bool = typer.Option( False, "--no-tailscale", help="Disable Tailscale VPN setup for this droplet" ), @@ -1871,6 +1876,10 @@ def create( This will create a droplet with the specified name, applying your cloud-init template and automatically adding an SSH config entry. + + Use --from-snapshot to create from an existing snapshot instead of a base + image. This skips cloud-init (the snapshot is already configured) and is + useful for cloning a golden image into multiple identical droplets. """ # Load configuration config_manager = Config() @@ -1931,8 +1940,13 @@ def create( console.print(f"[dim]Using default size: {config.defaults.size}[/dim]") size = config.defaults.size - # Get image (interactive if not provided) - if image is None: + # --from-snapshot and --image are mutually exclusive + if from_snapshot is not None and image is not None: + console.print("[red]Error: --from-snapshot and --image are mutually exclusive[/red]") + raise typer.Exit(1) + + # Get image (interactive if not provided) — skip when creating from snapshot + if from_snapshot is None and image is None: try: available_images = api.get_available_images() image = prompt_with_help( @@ -1947,7 +1961,10 @@ def create( image = config.defaults.image # Type guard - values are guaranteed non-None after interactive prompts - if name is None or region is None or size is None or image is None: + if name is None or region is None or size is None: + console.print("[red]Error: Missing required parameters[/red]") + raise typer.Exit(1) + if from_snapshot is None and image is None: console.print("[red]Error: Missing required parameters[/red]") raise typer.Exit(1) @@ -2066,7 +2083,10 @@ def create( table.add_row("Region:", region) table.add_row("Size:", size) - table.add_row("Image:", image) + if from_snapshot is not None: + table.add_row("Snapshot:", str(from_snapshot)) + else: + table.add_row("Image:", image or "") table.add_row("User:", username) table.add_row("Tags:", ", ".join(tags_list)) if project_id and project_name: @@ -2078,38 +2098,58 @@ def create( # Determine if Tailscale should be enabled tailscale_enabled = not no_tailscale and config.tailscale.enabled - # Render cloud-init + # Create droplet — two paths: from snapshot (no cloud-init) or from + # base image (with cloud-init). Snapshots are already configured, so + # re-running cloud-init would cause problems (user creation fails, + # .zshrc overwritten, unconditional reboot). try: - console.print("[dim]Rendering cloud-init template...[/dim]") - user_data = render_cloud_init( - template_path, username, full_name, email, ssh_keys, tailscale_enabled - ) + if from_snapshot is not None: + # Snapshot path — skip cloud-init entirely + console.print("[dim]Creating droplet from snapshot...[/dim]") + if verbose: + console.print( + f"[dim][DEBUG] Snapshot ID: {from_snapshot}, no cloud-init will be sent[/dim]" + ) + droplet = api.create_droplet_from_snapshot( + name=name, + region=region, + size=size, + snapshot_id=from_snapshot, + tags=tags_list, + ssh_keys=config.cloudinit.ssh_key_ids, + ) + else: + # Base image path — render and send cloud-init + try: + console.print("[dim]Rendering cloud-init template...[/dim]") + user_data = render_cloud_init( + template_path, username, full_name, email, ssh_keys, tailscale_enabled + ) - if verbose: - console.print("\n[dim][DEBUG] Rendered cloud-init template:[/dim]") - console.print("[dim]" + "=" * 60 + "[/dim]") - console.print(f"[dim]{user_data}[/dim]") - console.print("[dim]" + "=" * 60 + "[/dim]\n") - except Exception as e: - console.print(f"[red]Error rendering cloud-init: {e}[/red]") - raise typer.Exit(1) + if verbose: + console.print("\n[dim][DEBUG] Rendered cloud-init template:[/dim]") + console.print("[dim]" + "=" * 60 + "[/dim]") + console.print(f"[dim]{user_data}[/dim]") + console.print("[dim]" + "=" * 60 + "[/dim]\n") + except Exception as e: + console.print(f"[red]Error rendering cloud-init: {e}[/red]") + raise typer.Exit(1) - # Create droplet - try: - if verbose: - console.print( - f"[dim][DEBUG] Creating droplet with API endpoint: {config.digitalocean.api_base}/droplets[/dim]" + if verbose: + console.print( + f"[dim][DEBUG] Creating droplet with API endpoint: " + f"{config.digitalocean.api_base}/droplets[/dim]" + ) + console.print("[dim]Creating droplet via API...[/dim]") + droplet = api.create_droplet( + name=name, + region=region, + size=size, + image=image or "", + user_data=user_data, + tags=tags_list, + ssh_keys=config.cloudinit.ssh_key_ids, ) - console.print("[dim]Creating droplet via API...[/dim]") - droplet = api.create_droplet( - name=name, - region=region, - size=size, - image=image, - user_data=user_data, - tags=tags_list, - ssh_keys=config.cloudinit.ssh_key_ids, - ) droplet_id = droplet.get("id") if not droplet_id: @@ -2198,12 +2238,18 @@ def create( except Exception as e: console.print(f"[yellow]⚠[/yellow] Could not update SSH config: {e}") - # Wait for cloud-init to complete using helper function - cloud_init_done, cloud_init_error = wait_for_cloud_init(ssh_hostname, verbose) - - # Tailscale setup (if enabled and cloud-init succeeded) - if tailscale_enabled and cloud_init_done: - tailscale_ip = setup_tailscale(ssh_hostname, username, config, verbose) + if from_snapshot is not None: + # Snapshot-based: no cloud-init was sent, skip monitoring. + # The snapshot is already fully configured. + cloud_init_done = True + cloud_init_error = False + if tailscale_enabled: + tailscale_ip = setup_tailscale(ssh_hostname, username, config, verbose) + else: + # Base image: wait for cloud-init to complete + cloud_init_done, cloud_init_error = wait_for_cloud_init(ssh_hostname, verbose) + if tailscale_enabled and cloud_init_done: + tailscale_ip = setup_tailscale(ssh_hostname, username, config, verbose) # Show summary based on cloud-init and Tailscale status console.print() From 092a788b51f97587c0b02731f87a45852b1ced22 Mon Sep 17 00:00:00 2001 From: Grzegorz Wierzowiecki Date: Tue, 24 Mar 2026 10:31:58 +0100 Subject: [PATCH 2/3] Extract _create_from_snapshot_and_setup helper to consolidate shared code Both `create --from-snapshot` and `wake` perform the same post-creation plumbing: API call, wait-for-active, IP extraction, SSH config, and optional Tailscale setup. Extract this into a shared helper to eliminate ~90 lines of duplication. * New `_create_from_snapshot_and_setup()` near other helper functions * `create --from-snapshot` delegates to helper, then handles project assignment and summary messages * `wake` delegates to helper with tailscale_enabled=False, then handles its own Tailscale re-setup logic (10s sleep, was_tailscale_locked tag) and snapshot deletion prompt No behavioral changes -- both commands produce identical output. Co-Authored-By: Claude Opus 4.6 (1M context) --- dropkit/main.py | 369 ++++++++++++++++++++++++++++++------------------ 1 file changed, 229 insertions(+), 140 deletions(-) diff --git a/dropkit/main.py b/dropkit/main.py index f3a2e81..ce8ed1e 100644 --- a/dropkit/main.py +++ b/dropkit/main.py @@ -1579,6 +1579,127 @@ def find_user_droplet(api: DigitalOceanAPI, droplet_name: str) -> tuple[dict | N return None, username +def _create_from_snapshot_and_setup( + api: DigitalOceanAPI, + config: DropkitConfig, + name: str, + region: str, + size: str, + snapshot_id: int, + tags_list: list[str], + username: str, + tailscale_enabled: bool, + verbose: bool = False, +) -> tuple[dict, str | None, str | None]: + """Create a droplet from a snapshot and perform common post-creation setup. + + Shared by ``create --from-snapshot`` and ``wake``. Handles: + + * API call to ``create_droplet_from_snapshot`` + * Waiting for the droplet to become active + * Extracting the public IP address + * Adding an SSH config entry (if ``config.ssh.auto_update``) + * Optional Tailscale VPN setup + + Args: + api: Authenticated DigitalOcean API client. + config: Validated dropkit configuration. + name: Droplet name. + region: DigitalOcean region slug. + size: Droplet size slug. + snapshot_id: ID of the snapshot to create from. + tags_list: Tags to apply to the new droplet. + username: Linux username (derived from DO account email). + tailscale_enabled: Whether to run Tailscale setup after the droplet + is active. + verbose: Emit extra debug output. + + Returns: + ``(active_droplet, ip_address, tailscale_ip)`` where *ip_address* and + *tailscale_ip* may be ``None`` when unavailable. + + Raises: + typer.Exit: On unrecoverable API errors. + """ + console.print(f"[dim]Creating droplet '{name}' from snapshot...[/dim]") + if verbose: + console.print(f"[dim][DEBUG] Snapshot ID: {snapshot_id}, no cloud-init will be sent[/dim]") + + droplet = api.create_droplet_from_snapshot( + name=name, + region=region, + size=size, + snapshot_id=snapshot_id, + tags=tags_list, + ssh_keys=config.cloudinit.ssh_key_ids, + ) + + droplet_id = droplet.get("id") + if not droplet_id: + console.print("[red]Error: Failed to get droplet ID from API response[/red]") + raise typer.Exit(1) + + console.print(f"[green]✓[/green] Droplet created (ID: [cyan]{droplet_id}[/cyan])") + + if verbose: + console.print(f"[dim][DEBUG] Droplet status: {droplet.get('status')}[/dim]") + console.print(f"[dim][DEBUG] Full droplet response: {droplet}[/dim]") + + # Wait for droplet to become active + console.print("[dim]Waiting for droplet to become active...[/dim]") + if verbose: + console.print("[dim][DEBUG] Polling droplet status every 5 seconds...[/dim]") + + with console.status("[cyan]Waiting...[/cyan]"): + active_droplet = api.wait_for_droplet_active(droplet_id) + + # Extract public IP address + networks = active_droplet.get("networks", {}) + v4_networks = networks.get("v4", []) + ip_address: str | None = None + + for network in v4_networks: + if network.get("type") == "public": + ip_address = network.get("ip_address") + break + + if ip_address: + console.print(f"[green]✓[/green] Droplet is active (IP: [cyan]{ip_address}[/cyan])") + if verbose: + console.print(f"[dim][DEBUG] All v4 networks: {v4_networks}[/dim]") + else: + console.print("[green]✓[/green] Droplet is active") + console.print("[yellow]⚠[/yellow] Could not determine IP address") + + # SSH config setup + ssh_hostname = get_ssh_hostname(name) + tailscale_ip: str | None = None + + if ip_address and config.ssh.auto_update: + try: + console.print("[dim]Configuring SSH...[/dim]") + if verbose: + console.print(f"[dim][DEBUG] SSH config path: {config.ssh.config_path}[/dim]") + console.print(f"[dim][DEBUG] Adding host '{name}' -> {username}@{ip_address}[/dim]") + console.print(f"[dim][DEBUG] Identity file: {config.ssh.identity_file}[/dim]") + add_ssh_host( + config_path=config.ssh.config_path, + host_name=ssh_hostname, + hostname=ip_address, + user=username, + identity_file=config.ssh.identity_file, + ) + console.print(f"[green]✓[/green] SSH config updated: [cyan]ssh {ssh_hostname}[/cyan]") + except Exception as e: + console.print(f"[yellow]⚠[/yellow] Could not update SSH config: {e}") + + # Tailscale setup + if tailscale_enabled and ip_address: + tailscale_ip = setup_tailscale(ssh_hostname, username, config, verbose) + + return active_droplet, ip_address, tailscale_ip + + def find_project_by_name_or_id( api: DigitalOceanAPI, name_or_id: str ) -> tuple[str | None, str | None]: @@ -2104,20 +2225,36 @@ def create( # .zshrc overwritten, unconditional reboot). try: if from_snapshot is not None: - # Snapshot path — skip cloud-init entirely - console.print("[dim]Creating droplet from snapshot...[/dim]") - if verbose: - console.print( - f"[dim][DEBUG] Snapshot ID: {from_snapshot}, no cloud-init will be sent[/dim]" - ) - droplet = api.create_droplet_from_snapshot( + # Snapshot path — delegate to shared helper + active_droplet, ip_address, tailscale_ip = _create_from_snapshot_and_setup( + api=api, + config=config, name=name, region=region, size=size, snapshot_id=from_snapshot, - tags=tags_list, - ssh_keys=config.cloudinit.ssh_key_ids, + tags_list=tags_list, + username=username, + tailscale_enabled=tailscale_enabled, + verbose=verbose, ) + + # Assign droplet to project if specified + droplet_id = active_droplet.get("id") + if project_id and droplet_id: + try: + console.print(f"[dim]Assigning droplet to project '{project_name}'...[/dim]") + droplet_urn = api.get_droplet_urn(droplet_id) + api.assign_resources_to_project(project_id, [droplet_urn]) + console.print( + f"[green]✓[/green] Assigned to project: [cyan]{project_name}[/cyan]" + ) + except DigitalOceanAPIError as e: + console.print(f"[yellow]⚠[/yellow] Could not assign to project: {e}") + + ssh_hostname = get_ssh_hostname(name) + cloud_init_done = True + cloud_init_error = False else: # Base image path — render and send cloud-init try: @@ -2151,101 +2288,96 @@ def create( ssh_keys=config.cloudinit.ssh_key_ids, ) - droplet_id = droplet.get("id") - if not droplet_id: - console.print("[red]Error: Failed to get droplet ID from API response[/red]") - raise typer.Exit(1) + droplet_id = droplet.get("id") + if not droplet_id: + console.print("[red]Error: Failed to get droplet ID from API response[/red]") + raise typer.Exit(1) - console.print(f"[green]✓[/green] Droplet created with ID: [cyan]{droplet_id}[/cyan]") + console.print(f"[green]✓[/green] Droplet created with ID: [cyan]{droplet_id}[/cyan]") - if verbose: - console.print(f"[dim][DEBUG] Droplet status: {droplet.get('status')}[/dim]") - console.print(f"[dim][DEBUG] Full droplet response: {droplet}[/dim]") + if verbose: + console.print(f"[dim][DEBUG] Droplet status: {droplet.get('status')}[/dim]") + console.print(f"[dim][DEBUG] Full droplet response: {droplet}[/dim]") - # Wait for droplet to become active - console.print("[dim]Waiting for droplet to become active...[/dim]") - if verbose: - console.print("[dim][DEBUG] Polling droplet status every 5 seconds...[/dim]") + # Wait for droplet to become active + console.print("[dim]Waiting for droplet to become active...[/dim]") + if verbose: + console.print("[dim][DEBUG] Polling droplet status every 5 seconds...[/dim]") - with console.status("[cyan]Waiting...[/cyan]"): - active_droplet = api.wait_for_droplet_active(droplet_id) + with console.status("[cyan]Waiting...[/cyan]"): + active_droplet = api.wait_for_droplet_active(droplet_id) - console.print("[green]✓[/green] Droplet is now active") + console.print("[green]✓[/green] Droplet is now active") - if verbose: - console.print( - f"[dim][DEBUG] Active droplet networks: {active_droplet.get('networks')}[/dim]" - ) + if verbose: + console.print( + f"[dim][DEBUG] Active droplet networks: {active_droplet.get('networks')}[/dim]" + ) - # Assign droplet to project if specified - if project_id: - try: - console.print(f"[dim]Assigning droplet to project '{project_name}'...[/dim]") - droplet_urn = api.get_droplet_urn(droplet_id) - api.assign_resources_to_project(project_id, [droplet_urn]) - console.print(f"[green]✓[/green] Assigned to project: [cyan]{project_name}[/cyan]") - except DigitalOceanAPIError as e: - console.print(f"[yellow]⚠[/yellow] Could not assign to project: {e}") + # Assign droplet to project if specified + if project_id: + try: + console.print(f"[dim]Assigning droplet to project '{project_name}'...[/dim]") + droplet_urn = api.get_droplet_urn(droplet_id) + api.assign_resources_to_project(project_id, [droplet_urn]) + console.print( + f"[green]✓[/green] Assigned to project: [cyan]{project_name}[/cyan]" + ) + except DigitalOceanAPIError as e: + console.print(f"[yellow]⚠[/yellow] Could not assign to project: {e}") - # Get IP address - networks = active_droplet.get("networks", {}) - v4_networks = networks.get("v4", []) - ip_address = None + # Get IP address + networks = active_droplet.get("networks", {}) + v4_networks = networks.get("v4", []) + ip_address = None - for network in v4_networks: - if network.get("type") == "public": - ip_address = network.get("ip_address") - break + for network in v4_networks: + if network.get("type") == "public": + ip_address = network.get("ip_address") + break - # Initialize for type safety - ssh_hostname needed for output regardless of path - ssh_hostname = get_ssh_hostname(name) - tailscale_ip: str | None = None + # Initialize for type safety - ssh_hostname needed for output regardless + ssh_hostname = get_ssh_hostname(name) + tailscale_ip: str | None = None - if not ip_address: - console.print("[yellow]⚠[/yellow] Could not determine IP address") - cloud_init_done = False - cloud_init_error = False - else: - console.print(f"[green]✓[/green] IP address: [cyan]{ip_address}[/cyan]") - if verbose: - console.print(f"[dim][DEBUG] All v4 networks: {v4_networks}[/dim]") + if not ip_address: + console.print("[yellow]⚠[/yellow] Could not determine IP address") + cloud_init_done = False + cloud_init_error = False + else: + console.print(f"[green]✓[/green] IP address: [cyan]{ip_address}[/cyan]") + if verbose: + console.print(f"[dim][DEBUG] All v4 networks: {v4_networks}[/dim]") - # Add SSH config entry first so we can use it for cloud-init checks - if config.ssh.auto_update: - try: - console.print("[dim]Adding SSH config entry...[/dim]") - if verbose: - console.print( - f"[dim][DEBUG] SSH config path: {config.ssh.config_path}[/dim]" - ) - console.print( - f"[dim][DEBUG] Adding host '{name}' -> {username}@{ip_address}[/dim]" + # Add SSH config entry first so we can use it for cloud-init checks + if config.ssh.auto_update: + try: + console.print("[dim]Adding SSH config entry...[/dim]") + if verbose: + console.print( + f"[dim][DEBUG] SSH config path: {config.ssh.config_path}[/dim]" + ) + console.print( + f"[dim][DEBUG] Adding host '{name}' -> " + f"{username}@{ip_address}[/dim]" + ) + console.print( + f"[dim][DEBUG] Identity file: {config.ssh.identity_file}[/dim]" + ) + + add_ssh_host( + config_path=config.ssh.config_path, + host_name=ssh_hostname, + hostname=ip_address, + user=username, + identity_file=config.ssh.identity_file, ) console.print( - f"[dim][DEBUG] Identity file: {config.ssh.identity_file}[/dim]" + f"[green]✓[/green] Added SSH config: [cyan]ssh {ssh_hostname}[/cyan]" ) + except Exception as e: + console.print(f"[yellow]⚠[/yellow] Could not update SSH config: {e}") - add_ssh_host( - config_path=config.ssh.config_path, - host_name=ssh_hostname, - hostname=ip_address, - user=username, - identity_file=config.ssh.identity_file, - ) - console.print( - f"[green]✓[/green] Added SSH config: [cyan]ssh {ssh_hostname}[/cyan]" - ) - except Exception as e: - console.print(f"[yellow]⚠[/yellow] Could not update SSH config: {e}") - - if from_snapshot is not None: - # Snapshot-based: no cloud-init was sent, skip monitoring. - # The snapshot is already fully configured. - cloud_init_done = True - cloud_init_error = False - if tailscale_enabled: - tailscale_ip = setup_tailscale(ssh_hostname, username, config, verbose) - else: # Base image: wait for cloud-init to complete cloud_init_done, cloud_init_error = wait_for_cloud_init(ssh_hostname, verbose) if tailscale_enabled and cloud_init_done: @@ -4114,66 +4246,23 @@ def wake( ) console.print() - # Create droplet from snapshot - console.print(f"[dim]Creating droplet '{droplet_name}' from snapshot...[/dim]") - - # Build tags for new droplet + # Build tags and create droplet from snapshot via shared helper. + # Tailscale is handled separately for wake (only if was_tailscale_locked), + # so we pass tailscale_enabled=False here and do it ourselves below. tags_list = build_droplet_tags(username, list(config.defaults.extra_tags)) - droplet = api.create_droplet_from_snapshot( + active_droplet, ip_address, _ = _create_from_snapshot_and_setup( + api=api, + config=config, name=droplet_name, region=original_region, size=original_size, snapshot_id=snapshot_id, - tags=tags_list, - ssh_keys=config.cloudinit.ssh_key_ids, + tags_list=tags_list, + username=username, + tailscale_enabled=False, # wake handles Tailscale separately ) - droplet_id = droplet.get("id") - if not droplet_id: - console.print("[red]Error: Failed to get droplet ID from API response[/red]") - raise typer.Exit(1) - - console.print(f"[green]✓[/green] Droplet created (ID: [cyan]{droplet_id}[/cyan])") - - # Wait for droplet to become active - console.print("[dim]Waiting for droplet to become active...[/dim]") - - with console.status("[cyan]Waiting...[/cyan]"): - active_droplet = api.wait_for_droplet_active(droplet_id) - - # Get IP address - networks = active_droplet.get("networks", {}) - v4_networks = networks.get("v4", []) - ip_address = None - - for network in v4_networks: - if network.get("type") == "public": - ip_address = network.get("ip_address") - break - - if ip_address: - console.print(f"[green]✓[/green] Droplet is active (IP: [cyan]{ip_address}[/cyan])") - else: - console.print("[green]✓[/green] Droplet is active") - console.print("[yellow]⚠[/yellow] Could not determine IP address") - - # Add SSH config entry - if ip_address and config.ssh.auto_update: - try: - console.print("[dim]Configuring SSH...[/dim]") - ssh_hostname = get_ssh_hostname(droplet_name) - add_ssh_host( - config_path=config.ssh.config_path, - host_name=ssh_hostname, - hostname=ip_address, - user=username, - identity_file=config.ssh.identity_file, - ) - console.print("[green]✓[/green] SSH config updated") - except Exception as e: - console.print(f"[yellow]⚠[/yellow] Could not update SSH config: {e}") - # Handle Tailscale re-setup if the original droplet had Tailscale lockdown if was_tailscale_locked and ip_address: ssh_hostname = get_ssh_hostname(droplet_name) From 7e461d9a526ca31be839b92fecdc17542993e838 Mon Sep 17 00:00:00 2001 From: Grzegorz Wierzowiecki Date: Mon, 13 Apr 2026 18:15:32 +0200 Subject: [PATCH 3/3] Refactor: extract _post_create_setup shared by all creation paths MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace _create_from_snapshot_and_setup with a more generic _post_create_setup helper that handles post-creation steps (validate ID, wait for active, extract IP, SSH config, optional Tailscale) for ALL three droplet creation paths: - create (base image): API call → _post_create_setup → cloud-init → Tailscale - create --from-snapshot: API call → _post_create_setup (with Tailscale) - wake: API call → _post_create_setup → Tailscale (separate, conditional) Also consolidates project assignment (was duplicated in both create paths). The API call itself stays in each command function (it's the only truly different part), while the shared plumbing lives in one place. Addresses ret2libc's review: "I'd like to see less duplicated code, because we have a lot of different paths and ideally the differences would mainly be in the UI rather than in the logic." Co-Authored-By: Claude Opus 4.6 (1M context) --- dropkit/main.py | 193 +++++++++++++++--------------------------------- 1 file changed, 59 insertions(+), 134 deletions(-) diff --git a/dropkit/main.py b/dropkit/main.py index ce8ed1e..18edbd3 100644 --- a/dropkit/main.py +++ b/dropkit/main.py @@ -1579,36 +1579,33 @@ def find_user_droplet(api: DigitalOceanAPI, droplet_name: str) -> tuple[dict | N return None, username -def _create_from_snapshot_and_setup( +def _post_create_setup( api: DigitalOceanAPI, config: DropkitConfig, + droplet: dict, name: str, - region: str, - size: str, - snapshot_id: int, - tags_list: list[str], username: str, - tailscale_enabled: bool, + tailscale_enabled: bool = False, verbose: bool = False, ) -> tuple[dict, str | None, str | None]: - """Create a droplet from a snapshot and perform common post-creation setup. + """Shared post-creation setup for all droplet creation paths. - Shared by ``create --from-snapshot`` and ``wake``. Handles: + Called after any droplet creation API call (``create_droplet``, + ``create_droplet_from_snapshot``). Handles: - * API call to ``create_droplet_from_snapshot`` + * Validating the droplet ID from the API response * Waiting for the droplet to become active * Extracting the public IP address * Adding an SSH config entry (if ``config.ssh.auto_update``) * Optional Tailscale VPN setup + Used by: ``create`` (base image), ``create --from-snapshot``, ``wake``. + Args: api: Authenticated DigitalOcean API client. config: Validated dropkit configuration. + droplet: Raw droplet dict returned by a creation API call. name: Droplet name. - region: DigitalOcean region slug. - size: Droplet size slug. - snapshot_id: ID of the snapshot to create from. - tags_list: Tags to apply to the new droplet. username: Linux username (derived from DO account email). tailscale_enabled: Whether to run Tailscale setup after the droplet is active. @@ -1621,19 +1618,6 @@ def _create_from_snapshot_and_setup( Raises: typer.Exit: On unrecoverable API errors. """ - console.print(f"[dim]Creating droplet '{name}' from snapshot...[/dim]") - if verbose: - console.print(f"[dim][DEBUG] Snapshot ID: {snapshot_id}, no cloud-init will be sent[/dim]") - - droplet = api.create_droplet_from_snapshot( - name=name, - region=region, - size=size, - snapshot_id=snapshot_id, - tags=tags_list, - ssh_keys=config.cloudinit.ssh_key_ids, - ) - droplet_id = droplet.get("id") if not droplet_id: console.print("[red]Error: Failed to get droplet ID from API response[/red]") @@ -2225,33 +2209,29 @@ def create( # .zshrc overwritten, unconditional reboot). try: if from_snapshot is not None: - # Snapshot path — delegate to shared helper - active_droplet, ip_address, tailscale_ip = _create_from_snapshot_and_setup( - api=api, - config=config, + # Snapshot path — no cloud-init + console.print(f"[dim]Creating droplet '{name}' from snapshot...[/dim]") + if verbose: + console.print( + f"[dim][DEBUG] Snapshot ID: {from_snapshot}, no cloud-init will be sent[/dim]" + ) + droplet = api.create_droplet_from_snapshot( name=name, region=region, size=size, snapshot_id=from_snapshot, - tags_list=tags_list, + tags=tags_list, + ssh_keys=config.cloudinit.ssh_key_ids, + ) + active_droplet, ip_address, tailscale_ip = _post_create_setup( + api=api, + config=config, + droplet=droplet, + name=name, username=username, tailscale_enabled=tailscale_enabled, verbose=verbose, ) - - # Assign droplet to project if specified - droplet_id = active_droplet.get("id") - if project_id and droplet_id: - try: - console.print(f"[dim]Assigning droplet to project '{project_name}'...[/dim]") - droplet_urn = api.get_droplet_urn(droplet_id) - api.assign_resources_to_project(project_id, [droplet_urn]) - console.print( - f"[green]✓[/green] Assigned to project: [cyan]{project_name}[/cyan]" - ) - except DigitalOceanAPIError as e: - console.print(f"[yellow]⚠[/yellow] Could not assign to project: {e}") - ssh_hostname = get_ssh_hostname(name) cloud_init_done = True cloud_init_error = False @@ -2288,101 +2268,40 @@ def create( ssh_keys=config.cloudinit.ssh_key_ids, ) - droplet_id = droplet.get("id") - if not droplet_id: - console.print("[red]Error: Failed to get droplet ID from API response[/red]") - raise typer.Exit(1) - - console.print(f"[green]✓[/green] Droplet created with ID: [cyan]{droplet_id}[/cyan]") - - if verbose: - console.print(f"[dim][DEBUG] Droplet status: {droplet.get('status')}[/dim]") - console.print(f"[dim][DEBUG] Full droplet response: {droplet}[/dim]") - - # Wait for droplet to become active - console.print("[dim]Waiting for droplet to become active...[/dim]") - if verbose: - console.print("[dim][DEBUG] Polling droplet status every 5 seconds...[/dim]") - - with console.status("[cyan]Waiting...[/cyan]"): - active_droplet = api.wait_for_droplet_active(droplet_id) - - console.print("[green]✓[/green] Droplet is now active") - - if verbose: - console.print( - f"[dim][DEBUG] Active droplet networks: {active_droplet.get('networks')}[/dim]" - ) - - # Assign droplet to project if specified - if project_id: - try: - console.print(f"[dim]Assigning droplet to project '{project_name}'...[/dim]") - droplet_urn = api.get_droplet_urn(droplet_id) - api.assign_resources_to_project(project_id, [droplet_urn]) - console.print( - f"[green]✓[/green] Assigned to project: [cyan]{project_name}[/cyan]" - ) - except DigitalOceanAPIError as e: - console.print(f"[yellow]⚠[/yellow] Could not assign to project: {e}") - - # Get IP address - networks = active_droplet.get("networks", {}) - v4_networks = networks.get("v4", []) - ip_address = None - - for network in v4_networks: - if network.get("type") == "public": - ip_address = network.get("ip_address") - break - - # Initialize for type safety - ssh_hostname needed for output regardless + # Shared post-creation setup: wait for active, extract IP, SSH config. + # Tailscale is NOT passed here — for base images it runs after cloud-init. + active_droplet, ip_address, _ = _post_create_setup( + api=api, + config=config, + droplet=droplet, + name=name, + username=username, + tailscale_enabled=False, + verbose=verbose, + ) ssh_hostname = get_ssh_hostname(name) tailscale_ip: str | None = None if not ip_address: - console.print("[yellow]⚠[/yellow] Could not determine IP address") cloud_init_done = False cloud_init_error = False else: - console.print(f"[green]✓[/green] IP address: [cyan]{ip_address}[/cyan]") - if verbose: - console.print(f"[dim][DEBUG] All v4 networks: {v4_networks}[/dim]") - - # Add SSH config entry first so we can use it for cloud-init checks - if config.ssh.auto_update: - try: - console.print("[dim]Adding SSH config entry...[/dim]") - if verbose: - console.print( - f"[dim][DEBUG] SSH config path: {config.ssh.config_path}[/dim]" - ) - console.print( - f"[dim][DEBUG] Adding host '{name}' -> " - f"{username}@{ip_address}[/dim]" - ) - console.print( - f"[dim][DEBUG] Identity file: {config.ssh.identity_file}[/dim]" - ) - - add_ssh_host( - config_path=config.ssh.config_path, - host_name=ssh_hostname, - hostname=ip_address, - user=username, - identity_file=config.ssh.identity_file, - ) - console.print( - f"[green]✓[/green] Added SSH config: [cyan]ssh {ssh_hostname}[/cyan]" - ) - except Exception as e: - console.print(f"[yellow]⚠[/yellow] Could not update SSH config: {e}") - # Base image: wait for cloud-init to complete cloud_init_done, cloud_init_error = wait_for_cloud_init(ssh_hostname, verbose) if tailscale_enabled and cloud_init_done: tailscale_ip = setup_tailscale(ssh_hostname, username, config, verbose) + # Assign droplet to project if specified (shared across both paths) + droplet_id = active_droplet.get("id") + if project_id and droplet_id: + try: + console.print(f"[dim]Assigning droplet to project '{project_name}'...[/dim]") + droplet_urn = api.get_droplet_urn(droplet_id) + api.assign_resources_to_project(project_id, [droplet_urn]) + console.print(f"[green]✓[/green] Assigned to project: [cyan]{project_name}[/cyan]") + except DigitalOceanAPIError as e: + console.print(f"[yellow]⚠[/yellow] Could not assign to project: {e}") + # Show summary based on cloud-init and Tailscale status console.print() if tailscale_enabled and tailscale_ip: @@ -4246,19 +4165,25 @@ def wake( ) console.print() - # Build tags and create droplet from snapshot via shared helper. + # Build tags and create droplet from snapshot. # Tailscale is handled separately for wake (only if was_tailscale_locked), - # so we pass tailscale_enabled=False here and do it ourselves below. + # so we pass tailscale_enabled=False to the shared helper. tags_list = build_droplet_tags(username, list(config.defaults.extra_tags)) - active_droplet, ip_address, _ = _create_from_snapshot_and_setup( - api=api, - config=config, + console.print(f"[dim]Creating droplet '{droplet_name}' from snapshot...[/dim]") + droplet = api.create_droplet_from_snapshot( name=droplet_name, region=original_region, size=original_size, snapshot_id=snapshot_id, - tags_list=tags_list, + tags=tags_list, + ssh_keys=config.cloudinit.ssh_key_ids, + ) + active_droplet, ip_address, _ = _post_create_setup( + api=api, + config=config, + droplet=droplet, + name=droplet_name, username=username, tailscale_enabled=False, # wake handles Tailscale separately )