From ef6c170419d34c5c8248243b320c0e5d23770a28 Mon Sep 17 00:00:00 2001
From: Valentyn Boginskey <vboginskey@gmail.com>
Date: Wed, 12 Feb 2025 10:02:29 -0500
Subject: [PATCH] Support all VMSizeFlags in `fly postgres create` (#4216)

Before this change, `fly postgres create` only supported a set of
bespoke VM sizes via the --vm-size flag, making it impossible to create
a postgres cluster with, say, 1 shared CPU and 1 GB of memory.

As of this change, we create the underlying machine using the entire
`MachineGuest` set of options, which allows for additional
customization.

The legacy prompt and error behavior given no --vm-size or an invalid
--vm-size is preserved.
---
 internal/command/postgres/create.go | 46 ++++++++++++++++++++++-------
 1 file changed, 35 insertions(+), 11 deletions(-)

diff --git a/internal/command/postgres/create.go b/internal/command/postgres/create.go
index b8f25998a1..4f3dd97cd0 100644
--- a/internal/command/postgres/create.go
+++ b/internal/command/postgres/create.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"reflect"
 	"strings"
 
 	"github.com/spf13/cobra"
@@ -34,6 +35,7 @@ func newCreate() *cobra.Command {
 		flag.Region(),
 		flag.Org(),
 		flag.Detach(),
+		flag.VMSizeFlags,
 		flag.Bool{
 			Name:        "enable-backups",
 			Description: "Create a new tigris bucket and enable WAL-based backups",
@@ -48,10 +50,6 @@ func newCreate() *cobra.Command {
 			Shorthand:   "p",
 			Description: "The superuser password. The password will be generated for you if you leave this blank",
 		},
-		flag.String{
-			Name:        "vm-size",
-			Description: "the size of the VM",
-		},
 		flag.Int{
 			Name:        "initial-cluster-size",
 			Description: "Initial cluster size",
@@ -288,7 +286,24 @@ func CreateCluster(ctx context.Context, org *fly.Organization, region *fly.Regio
 		BarmanRemoteRestoreConfig: flag.GetString(ctx, "restore-target-app"),
 	}
 
-	customConfig := params.DiskGb != 0 || params.VMSize != "" || params.InitialClusterSize != 0 || params.ScaleToZero != nil
+	isCustomMachine := false
+	for _, sizeFlag := range flag.VMSizeFlags {
+		nameField := reflect.ValueOf(sizeFlag).FieldByName("Name")
+
+		if nameField.IsValid() {
+			name := nameField.String()
+			if name == "vm-size" {
+				continue
+			}
+
+			if flag.IsSpecified(ctx, name) {
+				isCustomMachine = true
+				break
+			}
+		}
+	}
+
+	customConfig := isCustomMachine || params.DiskGb != 0 || params.VMSize != "" || params.InitialClusterSize != 0 || params.ScaleToZero != nil
 
 	var config *PostgresConfiguration
 
@@ -344,13 +359,22 @@ func CreateCluster(ctx context.Context, org *fly.Organization, region *fly.Regio
 		}
 		input.InitialClusterSize = params.PostgresConfiguration.InitialClusterSize
 
-		// Resolve VM size
-		vmSize, err := resolveVMSize(ctx, params.VMSize)
-		if err != nil {
-			return err
-		}
+		if isCustomMachine {
+			guest, err := flag.GetMachineGuest(ctx, nil)
+			if err != nil {
+				return err
+			}
 
-		input.VMSize = vmSize
+			input.Guest = guest
+		} else {
+			// Resolve VM size
+			vmSize, err := resolveVMSize(ctx, params.VMSize)
+			if err != nil {
+				return err
+			}
+
+			input.VMSize = vmSize
+		}
 
 		if params.ScaleToZero != nil {
 			input.ScaleToZero = *params.ScaleToZero