fix: cleaning up function of running phases etc

Author: samrose

nix/packages/pg-ami-builder/cmd/build.go

@@ -185,26 +185,21 @@ func runBuildPhase1(cmd *cobra.Command, args []string) error {
 	fmt.Printf("\n✓ Packer build completed successfully!\n")
 	fmt.Println("✓ AMI created by packer")
 
-	// Parse AMI ID from packer output (it's already in stdout above)
-	// For now, use AWS API to find the AMI by tags
+	// Find AMI by execution ID tag (injected by our template rewriter)
+	var amiID string
 	ec2Client, err := aws.NewEC2Client(ctx, region)
 	if err != nil {
 		fmt.Printf("⚠ Could not create EC2 client to find AMI: %v\n", err)
-		fmt.Println("\n✓ Build phase 1 complete!")
-		return nil
-	}
-
-	// Find AMI by execution ID tag
-	amiID, err := ec2Client.FindAMIByTag(ctx, "packerExecutionId", executionID)
-	if err != nil {
-		fmt.Printf("⚠ Could not find created AMI: %v\n", err)
-		fmt.Println("\n✓ Build phase 1 complete!")
-		return nil
+	} else {
+		amiID, err = ec2Client.FindAMIByTag(ctx, "packerExecutionId", executionID)
+		if err != nil {
+			fmt.Printf("⚠ Could not find created AMI: %v\n", err)
+		} else {
+			fmt.Printf("✓ AMI ID: %s\n", amiID)
+		}
 	}
 
-	fmt.Printf("✓ AMI ID: %s\n", amiID)
-
-	// Save AMI ID to state for phase2
+	// Save state even if we couldn't find the AMI
 	stateFilePath := stateFile
 	if stateFilePath == "" {
 		stateFilePath, err = state.GetDefaultStateFile()
@@ -215,14 +210,19 @@ func runBuildPhase1(cmd *cobra.Command, args []string) error {
 		}
 	}
 
-	buildState := &state.State{
-		Region:          region,
-		PostgresVersion: postgresVersion,
-		GitSHA:          sha,
+	// Load existing state or create new
+	buildState, err := state.LoadState(stateFilePath)
+	if err != nil {
+		buildState = &state.State{
+			Region:          region,
+			PostgresVersion: postgresVersion,
+			GitSHA:          sha,
+		}
 	}
+
 	buildState.SetPhaseState("phase1", &state.PhaseState{
 		ExecutionID: executionID,
-		AMIID:       amiID,
+		AMIID:       amiID, // Will be empty if not found
 		Timestamp:   time.Now().Format(time.RFC3339),
 	})
 
@@ -233,8 +233,13 @@ func runBuildPhase1(cmd *cobra.Command, args []string) error {
 	}
 
 	fmt.Println("\n✓ Build phase 1 complete!")
-	fmt.Printf("\nNext: Run phase 2 with:\n")
-	fmt.Printf("  pg-ami-builder build phase2 --postgres-version %s\n", postgresVersion)
+	if amiID != "" {
+		fmt.Printf("\nNext: Run phase 2 with:\n")
+		fmt.Printf("  pg-ami-builder build phase2 --postgres-version %s\n", postgresVersion)
+	} else {
+		fmt.Printf("\nNote: AMI ID not automatically detected. You can manually add it to:\n")
+		fmt.Printf("  %s\n", stateFilePath)
+	}
 	return nil
 }
 
@@ -398,7 +403,61 @@ func runBuildPhase2(cmd *cobra.Command, args []string) error {
 
 	fmt.Printf("\n✓ Packer build completed successfully!\n")
 	fmt.Println("✓ Final production AMI created by packer")
+
+	// Find AMI by execution ID tag (injected by our template rewriter)
+	var amiID string
+	ec2Client, err := aws.NewEC2Client(ctx, region)
+	if err != nil {
+		fmt.Printf("⚠ Could not create EC2 client to find AMI: %v\n", err)
+	} else {
+		amiID, err = ec2Client.FindAMIByTag(ctx, "packerExecutionId", executionID)
+		if err != nil {
+			fmt.Printf("⚠ Could not find created AMI: %v\n", err)
+		} else {
+			fmt.Printf("✓ AMI ID: %s\n", amiID)
+		}
+	}
+
+	// Save state even if we couldn't find the AMI
+	stateFilePath := stateFile
+	if stateFilePath == "" {
+		stateFilePath, err = state.GetDefaultStateFile()
+		if err != nil {
+			fmt.Printf("⚠ Could not get state file path: %v\n", err)
+			fmt.Println("\n✓ Build phase 2 complete!")
+			return nil
+		}
+	}
+
+	// Load existing state or create new
+	buildState, err := state.LoadState(stateFilePath)
+	if err != nil {
+		buildState = &state.State{
+			Region:          region,
+			PostgresVersion: postgresVersion,
+			GitSHA:          sha,
+		}
+	}
+
+	buildState.SetPhaseState("phase2", &state.PhaseState{
+		ExecutionID: executionID,
+		AMIID:       amiID, // Will be empty if not found
+		Timestamp:   time.Now().Format(time.RFC3339),
+	})
+
+	if err := state.SaveState(stateFilePath, buildState); err != nil {
+		fmt.Printf("⚠ Could not save state: %v\n", err)
+	} else {
+		fmt.Printf("✓ State saved to: %s\n", stateFilePath)
+	}
+
 	fmt.Println("\n✓ Build phase 2 complete!")
+	if amiID != "" {
+		fmt.Printf("\nProduction AMI ready: %s\n", amiID)
+	} else {
+		fmt.Printf("\nNote: AMI ID not automatically detected. You can manually add it to:\n")
+		fmt.Printf("  %s\n", stateFilePath)
+	}
 	return nil
 }
 

nix/packages/pg-ami-builder/cmd/cleanup.go

@@ -122,12 +122,33 @@ func runCleanup(cmd *cobra.Command, args []string) error {
 
 	fmt.Println("✓ Instance terminated")
 
-	// Clear state file
-	if stateFilePath != "" {
-		if err := state.ClearState(stateFilePath); err != nil {
-			fmt.Printf("Warning: failed to clear state file: %v\n", err)
-		} else {
-			fmt.Println("✓ State file cleared")
+	// Clear phase-specific state
+	if stateFilePath != "" && cleanupPhase != "" {
+		// Load state, clear just this phase, save back
+		buildState, err := state.LoadState(stateFilePath)
+		if err == nil {
+			// Clear this specific phase
+			buildState.SetPhaseState(cleanupPhase, nil)
+
+			// Check if both phases are now empty
+			phase1Empty := buildState.Phase1 == nil || (buildState.Phase1.InstanceID == "" && buildState.Phase1.AMIID == "")
+			phase2Empty := buildState.Phase2 == nil || (buildState.Phase2.InstanceID == "" && buildState.Phase2.AMIID == "")
+
+			if phase1Empty && phase2Empty {
+				// Both phases empty, delete entire file
+				if err := state.ClearState(stateFilePath); err != nil {
+					fmt.Printf("⚠ Failed to clear state file: %v\n", err)
+				} else {
+					fmt.Println("✓ State file cleared (all phases cleaned)")
+				}
+			} else {
+				// Save updated state with this phase cleared
+				if err := state.SaveState(stateFilePath, buildState); err != nil {
+					fmt.Printf("⚠ Failed to update state file: %v\n", err)
+				} else {
+					fmt.Printf("✓ State file updated (%s cleared)\n", cleanupPhase)
+				}
+			}
 		}
 	}
 

nix/packages/pg-ami-builder/cmd/ssh.go

@@ -13,6 +13,7 @@ import (
 
 var (
 	sshInstanceID    string
+	sshPhase         string
 	awsEC2ConnectCmd string
 )
 
@@ -51,12 +52,31 @@ func runSSH(cmd *cobra.Command, args []string) error {
 			return fmt.Errorf("failed to load state: %w", err)
 		}
 
-		instanceID = buildState.InstanceID
 		region = buildState.Region
+
+		// Get instance ID from phase-specific state
+		if sshPhase != "" {
+			phaseState := buildState.GetPhaseState(sshPhase)
+			if phaseState == nil || phaseState.InstanceID == "" {
+				return fmt.Errorf("no instance found for %s in state file", sshPhase)
+			}
+			instanceID = phaseState.InstanceID
+		} else {
+			// Auto-detect: prefer phase2, then phase1, then legacy
+			if buildState.Phase2 != nil && buildState.Phase2.InstanceID != "" {
+				instanceID = buildState.Phase2.InstanceID
+				fmt.Println("✓ Auto-detected phase2 instance")
+			} else if buildState.Phase1 != nil && buildState.Phase1.InstanceID != "" {
+				instanceID = buildState.Phase1.InstanceID
+				fmt.Println("✓ Auto-detected phase1 instance")
+			} else if buildState.InstanceID != "" {
+				instanceID = buildState.InstanceID
+			}
+		}
 	}
 
 	if instanceID == "" {
-		return fmt.Errorf("no instance ID available (use --instance-id or run build command first)")
+		return fmt.Errorf("no instance ID available (use --instance-id or --phase, or run build command first)")
 	}
 
 	// Connect via EC2 Instance Connect
@@ -147,6 +167,7 @@ func init() {
 	rootCmd.AddCommand(sshCmd)
 
 	sshCmd.Flags().StringVar(&sshInstanceID, "instance-id", "", "Target specific instance (default: from state file)")
+	sshCmd.Flags().StringVar(&sshPhase, "phase", "", "Connect to specific phase instance (phase1 or phase2)")
 	sshCmd.Flags().StringVar(&region, "region", "us-east-1", "AWS region")
 	sshCmd.Flags().StringVar(&awsEC2ConnectCmd, "aws-ec2-connect-cmd", "", "Custom AWS EC2 Instance Connect command (e.g., 'aws ec2-instance-connect ssh --instance-id i-xxx ...')")
 }

nix/packages/pg-ami-builder/internal/packer/runner.go

@@ -45,6 +45,22 @@ func RewriteTemplateWithUniqueAMI(templatePath, executionID, sourceAMI string) (
 		return strings.TrimSuffix(match, `"`) + `-${var.packer-execution-id}"`
 	})
 
+	// Add packerExecutionId to AMI tags (not run_tags)
+	// Find the tags = { block (after run_tags) and inject packerExecutionId
+	tagsPattern := regexp.MustCompile(`(?s)(tags\s*=\s*\{)([^}]*creator\s*=\s*"packer"[^}]*)(\})`)
+	modified = tagsPattern.ReplaceAllStringFunc(modified, func(match string) string {
+		// Check if packerExecutionId already exists in this tags block
+		if strings.Contains(match, "packerExecutionId") {
+			return match // Already has it
+		}
+		// Insert packerExecutionId after the opening brace
+		parts := tagsPattern.FindStringSubmatch(match)
+		if len(parts) == 4 {
+			return parts[1] + "\n    packerExecutionId = \"${var.packer-execution-id}\"" + parts[2] + parts[3]
+		}
+		return match
+	})
+
 	// Create temp file
 	tempDir := os.TempDir()
 	tempFile := filepath.Join(tempDir, fmt.Sprintf("packer-%s-%s", executionID, filepath.Base(templatePath)))