Merge pull request #1110 from planetscale/add-engine-flag-to-cluster-size-list

Add engine flag to cluster size list

Author: mscoutermarsh

go.mod

@@ -18,13 +18,14 @@ require (
 	github.com/hashicorp/go-cleanhttp v0.5.2
 	github.com/hashicorp/go-version v1.7.0
 	github.com/lensesio/tableprinter v0.0.0-20201125135848-89e81fc956e7
+	github.com/lib/pq v1.10.9
 	github.com/mark3labs/mcp-go v0.37.0
 	github.com/matoous/go-nanoid/v2 v2.1.0
 	github.com/mattn/go-isatty v0.0.20
 	github.com/mattn/go-shellwords v1.0.12
 	github.com/mitchellh/go-homedir v1.1.0
 	github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c
-	github.com/planetscale/planetscale-go v0.139.0
+	github.com/planetscale/planetscale-go v0.140.0
 	github.com/planetscale/psdb v0.0.0-20250717190954-65c6661ab6e4
 	github.com/planetscale/psdbproxy v0.0.0-20250728082226-3f4ea3a74ec7
 	github.com/spf13/cobra v1.9.1
@@ -80,7 +81,6 @@ require (
 	github.com/klauspost/connect-compress/v2 v2.0.0 // indirect
 	github.com/kr/pretty v0.3.1 // indirect
 	github.com/kr/text v0.2.0 // indirect
-	github.com/lib/pq v1.10.9 // indirect
 	github.com/lucasb-eyer/go-colorful v1.2.0 // indirect
 	github.com/mailru/easyjson v0.7.7 // indirect
 	github.com/mattn/go-colorable v0.1.14 // indirect

go.sum

@@ -177,6 +177,8 @@ github.com/planetscale/noglog v0.2.1-0.20210421230640-bea75fcd2e8e h1:MZ8D+Z3m2v
 github.com/planetscale/noglog v0.2.1-0.20210421230640-bea75fcd2e8e/go.mod h1:hwAsSPQdvPa3WcfKfzTXxtEq/HlqwLjQasfO6QbGo4Q=
 github.com/planetscale/planetscale-go v0.139.0 h1:Ie7RnhqY8QjBDx9MxONI88Arnex2Wpc18H7/5YW9DuI=
 github.com/planetscale/planetscale-go v0.139.0/go.mod h1:/n7PU99UvYaajJlTbMWMOOlHmkEuN7SFjvLNDSgMQOk=
+github.com/planetscale/planetscale-go v0.140.0 h1:zyL5kArfVUwg/+DyPeR6ENB+2+eaNzVW7DbotW0GqG4=
+github.com/planetscale/planetscale-go v0.140.0/go.mod h1:/n7PU99UvYaajJlTbMWMOOlHmkEuN7SFjvLNDSgMQOk=
 github.com/planetscale/psdb v0.0.0-20250717190954-65c6661ab6e4 h1:Xv5pj20Rhfty1Tv0OVcidg4ez4PvGrpKvb6rvUwQgDs=
 github.com/planetscale/psdb v0.0.0-20250717190954-65c6661ab6e4/go.mod h1:M52h5IWxAcbdQ1hSZrLAGQC4ZXslxEsK/Wh9nu3wdWs=
 github.com/planetscale/psdbproxy v0.0.0-20250728082226-3f4ea3a74ec7 h1:aRd6vdE1fyuSI4RVj7oCr8lFmgqXvpnPUmN85VbZCp8=

internal/cmd/size/cluster.go

@@ -27,11 +27,13 @@ func ListCmd(ch *cmdutil.Helper) *cobra.Command {
 	var flags struct {
 		region string
 		metal  bool
+		engine string
 	}
 
 	cmd := &cobra.Command{
 		Use:     "list",
 		Short:   "List the sizes that are available for a PlanetScale database",
+		Long:    "List the sizes that are available for a PlanetScale database. Use --engine to specify the database engine type.",
 		Args:    cobra.NoArgs,
 		Aliases: []string{"ls"},
 		RunE: func(cmd *cobra.Command, args []string) error {
@@ -41,9 +43,26 @@ func ListCmd(ch *cmdutil.Helper) *cobra.Command {
 				return err
 			}
 
+			// Parse the engine flag
+			engine, err := parseDatabaseEngine(flags.engine)
+			if err != nil {
+				return err
+			}
+
+			// Build the list options
+			listOpts := []planetscale.ListOption{planetscale.WithRates()}
+			if flags.region != "" {
+				listOpts = append(listOpts, planetscale.WithRegion(flags.region))
+			}
+			
+			// Add engine-specific parameter for PostgreSQL
+			if engine == planetscale.DatabaseEnginePostgres {
+				listOpts = append(listOpts, planetscale.WithPostgreSQL())
+			}
+
 			clusterSKUs, err := client.Organizations.ListClusterSKUs(ctx, &planetscale.ListOrganizationClusterSKUsRequest{
 				Organization: ch.Config.Organization,
-			}, planetscale.WithRates(), planetscale.WithRegion(flags.region))
+			}, listOpts...)
 			if err != nil {
 				return err
 			}
@@ -54,11 +73,19 @@ func ListCmd(ch *cmdutil.Helper) *cobra.Command {
 
 	cmd.Flags().StringVar(&flags.region, "region", "", "view cluster sizes and rates for a specific region")
 	cmd.Flags().BoolVar(&flags.metal, "metal", false, "view cluster sizes and rates for clusters with metal storage")
+	cmd.Flags().StringVar(&flags.engine, "engine", "mysql", "The database engine to show cluster sizes for. Supported values: mysql, postgresql. Defaults to mysql.")
 
 	cmd.RegisterFlagCompletionFunc("region", func(cmd *cobra.Command, args []string, toComplete string) ([]cobra.Completion, cobra.ShellCompDirective) {
 		return cmdutil.RegionsCompletionFunc(ch, cmd, args, toComplete)
 	})
 
+	cmd.RegisterFlagCompletionFunc("engine", func(cmd *cobra.Command, args []string, toComplete string) ([]cobra.Completion, cobra.ShellCompDirective) {
+		return []cobra.Completion{
+			cobra.CompletionWithDesc("mysql", "A Vitess database"),
+			cobra.CompletionWithDesc("postgresql", "The fastest cloud Postgres"),
+		}, cobra.ShellCompDirectiveNoFileComp
+	})
+
 	return cmd
 }
 
@@ -105,6 +132,17 @@ func toClusterSKU(clusterSKU *planetscale.ClusterSKU) *ClusterSKU {
 	return cluster
 }
 
+func parseDatabaseEngine(engine string) (planetscale.DatabaseEngine, error) {
+	switch engine {
+	case "mysql":
+		return planetscale.DatabaseEngineMySQL, nil
+	case "postgresql", "postgres":
+		return planetscale.DatabaseEnginePostgres, nil
+	default:
+		return planetscale.DatabaseEngineMySQL, fmt.Errorf("invalid database engine %q, supported values: mysql, postgresql", engine)
+	}
+}
+
 func toClusterSKUs(clusterSKUs []*planetscale.ClusterSKU, onlyMetal bool) []*ClusterSKU {
 	clusters := make([]*ClusterSKU, 0, len(clusterSKUs))
 

internal/cmd/size/cluster_test.go

@@ -59,3 +59,49 @@ func TestSizeCluster_ListCmd(t *testing.T) {
 
 	c.Assert(buf.String(), qt.JSONEquals, res)
 }
+
+func TestSizeCluster_ListCmd_PostgreSQL(t *testing.T) {
+	c := qt.New(t)
+
+	var buf bytes.Buffer
+	format := printer.JSON
+	p := printer.NewPrinter(&format)
+	p.SetResourceOutput(&buf)
+
+	org := "planetscale"
+
+	orig := []*ps.ClusterSKU{
+		{Name: "PS-10", Enabled: true, Rate: testutil.Pointer[int64](39)},
+	}
+	svc := &mock.OrganizationsService{
+		ListClusterSKUsFn: func(ctx context.Context, req *ps.ListOrganizationClusterSKUsRequest, opts ...ps.ListOption) ([]*ps.ClusterSKU, error) {
+			c.Assert(req.Organization, qt.Equals, org)
+			return orig, nil
+		},
+	}
+
+	ch := &cmdutil.Helper{
+		Printer: p,
+		Config: &config.Config{
+			Organization: org,
+		},
+		Client: func() (*ps.Client, error) {
+			return &ps.Client{
+				Organizations: svc,
+			}, nil
+		},
+	}
+
+	cmd := ListCmd(ch)
+	cmd.SetArgs([]string{"--engine", "postgresql"})
+	err := cmd.Execute()
+
+	c.Assert(err, qt.IsNil)
+	c.Assert(svc.ListClusterSKUsFnInvoked, qt.IsTrue)
+	
+	res := []*ClusterSKU{
+		{orig: orig[0]},
+	}
+	
+	c.Assert(buf.String(), qt.JSONEquals, res)
+}

internal/cmdutil/completions.go

@@ -28,9 +28,15 @@ func ClusterSizesCompletionFunc(ch *Helper, cmd *cobra.Command, args []string, t
 		return nil, cobra.ShellCompDirectiveNoFileComp
 	}
 
+	// Build list options
+	listOpts := []ps.ListOption{ps.WithRates()}
+	if region != "" {
+		listOpts = append(listOpts, ps.WithRegion(region))
+	}
+	
 	clusterSKUs, err := client.Organizations.ListClusterSKUs(ctx, &ps.ListOrganizationClusterSKUsRequest{
 		Organization: org,
-	}, ps.WithRegion(region), ps.WithRates())
+	}, listOpts...)
 	if err != nil {
 		return nil, cobra.ShellCompDirectiveNoFileComp
 	}