keyspace create: add repeatable --custom-shard flag

AGENTS.md

@@ -475,6 +475,9 @@ Ask the user for the source password; do not invent credentials. `--dry-run` che
 ```bash
 pscale keyspace list <database> <branch> --org <org> --format json
 pscale keyspace show <database> <branch> <keyspace> --org <org> --format json
+pscale keyspace create <database> <branch> <keyspace> --org <org> --format json --cluster-size PS_40 --shards 2
+pscale keyspace create <database> <branch> <keyspace> --org <org> --format json --cluster-size PS_40 \
+  --custom-shard key-range=-40,cluster-size=PS_80 --custom-shard key-range=40-80 --custom-shard key-range=80-
 pscale size cluster list --org <org> --format json --external
 pscale keyspace create-external <database> <branch> <keyspace> --org <org> --format json \
   --host <host> --source-database <remote-db> --username <user> --password <password> \
@@ -488,6 +491,8 @@ pscale keyspace resize status <database> <branch> <keyspace> --org <org> --forma
 
 External create required flags: `--host`, `--source-database`, `--username`, `--password`, `--ssl-mode` (`disabled`, `preferred`, `required`, `verify_ca`, `verify_identity`). Default `--port` is `3306`.
 
+Internal `keyspace create` requires `--cluster-size`. `--shards` (default `1`) creates equal shards. `--custom-shard key-range=<range>[,cluster-size=<size>]` is repeatable and creates shards with explicit Vitess key ranges instead; it needs at least two entries whose ranges together cover the keyspace, and cannot be combined with `--shards`. A shard without `cluster-size` uses the keyspace `--cluster-size`. The API validates range continuity and size eligibility.
+
 Keyspace VTTablet and MySQL parameters: list them with current and default values, then change them with `parameters set`. Prefix each parameter with its namespace (`vttablet.` or `mysqld.`); `--reset` sets a parameter back to its default. Changes for both namespaces are submitted together and roll out to the keyspace; follow them with `parameters changes list`. Only one unfinished change per namespace can exist on a keyspace at a time.
 
 ```bash

internal/cmd/keyspace/create.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"strings"
 	"time"
 
 	"github.com/planetscale/cli/internal/cmdutil"
@@ -20,6 +21,7 @@ func CreateCmd(ch *cmdutil.Helper) *cobra.Command {
 		shards             int
 		clusterSize        string
 		additionalReplicas int
+		customShards       []string
 		wait               bool
 	}
 
@@ -35,10 +37,24 @@ func CreateCmd(ch *cmdutil.Helper) *cobra.Command {
 			createReq.Database = database
 			createReq.Branch = branch
 			createReq.Name = keyspace
-			createReq.Shards = flags.shards
 			createReq.ExtraReplicas = flags.additionalReplicas
 			createReq.ClusterSize = flags.clusterSize
 
+			if len(flags.customShards) > 0 {
+				if cmd.Flags().Changed("shards") {
+					return errors.New("--shards cannot be combined with --custom-shard")
+				}
+				customShards, err := parseCustomShards(flags.customShards)
+				if err != nil {
+					return err
+				}
+				createReq.Shards = nil
+				createReq.CustomShards = customShards
+			} else {
+				createReq.Shards = &flags.shards
+				createReq.CustomShards = nil
+			}
+
 			client, err := ch.Client()
 			if err != nil {
 				return err
@@ -86,6 +102,7 @@ func CreateCmd(ch *cmdutil.Helper) *cobra.Command {
 	cmd.Flags().IntVar(&flags.shards, "shards", 1, "Number of shards in the keyspace")
 	cmd.Flags().StringVar(&flags.clusterSize, "cluster-size", "", "cluster size for the keyspace. Use `pscale size cluster list` to get a list of valid sizes.")
 	cmd.Flags().IntVar(&flags.additionalReplicas, "additional-replicas", 0, "number of additional replicas per shard. By default, each production cluster includes 2 replicas.")
+	cmd.Flags().StringArrayVar(&flags.customShards, "custom-shard", nil, customShardFlagHelp)
 	cmd.Flags().BoolVar(&flags.wait, "wait", false, "Wait until the keyspace is ready")
 
 	cmd.MarkFlagRequired("cluster-size")
@@ -97,6 +114,67 @@ func CreateCmd(ch *cmdutil.Helper) *cobra.Command {
 	return cmd
 }
 
+const customShardFlagHelp = "key range and optional cluster size for one shard as key-range=<range>[,cluster-size=<size>]. Repeatable; requires at least two shards whose ranges cover the whole keyspace. Shards without a cluster-size use --cluster-size. Cannot be combined with --shards."
+
+// parseCustomShards parses repeatable --custom-shard flags.
+func parseCustomShards(raw []string) ([]planetscale.CreateKeyspaceCustomShard, error) {
+	if len(raw) < 2 {
+		return nil, fmt.Errorf("--custom-shard requires at least two shards, got %d", len(raw))
+	}
+	out := make([]planetscale.CreateKeyspaceCustomShard, 0, len(raw))
+	seen := make(map[string]struct{}, len(raw))
+	for _, s := range raw {
+		entry, err := parseCustomShard(s)
+		if err != nil {
+			return nil, err
+		}
+		if _, ok := seen[entry.KeyRange]; ok {
+			return nil, fmt.Errorf("duplicate --custom-shard key-range %q", entry.KeyRange)
+		}
+		seen[entry.KeyRange] = struct{}{}
+		out = append(out, entry)
+	}
+	return out, nil
+}
+
+func parseCustomShard(s string) (planetscale.CreateKeyspaceCustomShard, error) {
+	var entry planetscale.CreateKeyspaceCustomShard
+	seen := make(map[string]struct{}, 2)
+	for part := range strings.SplitSeq(s, ",") {
+		part = strings.TrimSpace(part)
+		if part == "" {
+			continue
+		}
+		k, v, ok := strings.Cut(part, "=")
+		if !ok {
+			return entry, fmt.Errorf("invalid --custom-shard %q: expected key-range=<range>[,cluster-size=<size>]", s)
+		}
+		k = strings.TrimSpace(k)
+		v = strings.TrimSpace(v)
+		switch k {
+		case "key-range", "cluster-size":
+		default:
+			return entry, fmt.Errorf("unknown --custom-shard field %q in %q: valid fields are key-range, cluster-size", k, s)
+		}
+		if _, dup := seen[k]; dup {
+			return entry, fmt.Errorf("duplicate --custom-shard field %q in %q", k, s)
+		}
+		seen[k] = struct{}{}
+		if v == "" {
+			return entry, fmt.Errorf("invalid --custom-shard %q: %s must not be empty", s, k)
+		}
+		if k == "key-range" {
+			entry.KeyRange = v
+		} else {
+			entry.ClusterSize = cmdutil.ToSizeSKUName(v)
+		}
+	}
+	if entry.KeyRange == "" {
+		return entry, fmt.Errorf("--custom-shard %q is missing required field key-range", s)
+	}
+	return entry, nil
+}
+
 func waitUntilReady(ctx context.Context, client *planetscale.Client, printer *printer.Printer, debug bool, getReq *planetscale.GetKeyspaceRequest) (*planetscale.Keyspace, error) {
 	ctx, cancel := context.WithTimeout(ctx, 3*time.Minute)
 	defer cancel()

internal/cmd/keyspace/create_test.go

@@ -43,7 +43,8 @@ func TestKeyspace_CreateCmd(t *testing.T) {
 			c.Assert(req.Name, qt.Equals, keyspace)
 			c.Assert(req.ClusterSize, qt.Equals, "PS-20")
 			c.Assert(req.ExtraReplicas, qt.Equals, 1)
-			c.Assert(req.Shards, qt.Equals, 2)
+			c.Assert(req.Shards, qt.DeepEquals, intPtr(2))
+			c.Assert(req.CustomShards, qt.IsNil)
 
 			return ks, nil
 		},
@@ -86,7 +87,9 @@ func TestKeyspace_CreateCmdWithWaitPrintsReadyKeyspace(t *testing.T) {
 	ready := &ps.Keyspace{Name: keyspace, Ready: true}
 
 	svc := &mock.KeyspacesService{
-		CreateFn: func(_ context.Context, _ *ps.CreateKeyspaceRequest) (*ps.Keyspace, error) {
+		CreateFn: func(_ context.Context, req *ps.CreateKeyspaceRequest) (*ps.Keyspace, error) {
+			c.Assert(req.Shards, qt.DeepEquals, intPtr(1))
+			c.Assert(req.CustomShards, qt.IsNil)
 			return pending, nil
 		},
 		GetFn: func(_ context.Context, req *ps.GetKeyspaceRequest) (*ps.Keyspace, error) {
@@ -148,7 +151,8 @@ func TestKeyspace_CreateCmdOnlyClusterSize(t *testing.T) {
 			c.Assert(req.Name, qt.Equals, keyspace)
 			c.Assert(req.ClusterSize, qt.Equals, "PS-10")
 			c.Assert(req.ExtraReplicas, qt.Equals, 0)
-			c.Assert(req.Shards, qt.Equals, 1)
+			c.Assert(req.Shards, qt.DeepEquals, intPtr(1))
+			c.Assert(req.CustomShards, qt.IsNil)
 
 			return ks, nil
 		},
@@ -173,3 +177,188 @@ func TestKeyspace_CreateCmdOnlyClusterSize(t *testing.T) {
 	c.Assert(svc.CreateFnInvoked, qt.IsTrue)
 	c.Assert(buf.String(), qt.JSONEquals, ks)
 }
+
+func intPtr(n int) *int { return &n }
+
+func newCreateTestHelper(c *qt.C, buf *bytes.Buffer, svc *mock.KeyspacesService) *cmdutil.Helper {
+	format := printer.JSON
+	p := printer.NewPrinter(&format)
+	p.SetResourceOutput(buf)
+
+	ch := &cmdutil.Helper{
+		Printer: p,
+		Config:  &config.Config{Organization: "planetscale"},
+		Client: func() (*ps.Client, error) {
+			return &ps.Client{Keyspaces: svc}, nil
+		},
+	}
+	debug := false
+	ch.SetDebug(&debug)
+	return ch
+}
+
+func TestKeyspace_CreateCmdCustomShards(t *testing.T) {
+	c := qt.New(t)
+
+	var buf bytes.Buffer
+	ks := &ps.Keyspace{ID: "wantid", Name: "events", ClusterSize: "PS_40", Shards: 3, Sharded: true}
+
+	svc := &mock.KeyspacesService{
+		CreateFn: func(_ context.Context, req *ps.CreateKeyspaceRequest) (*ps.Keyspace, error) {
+			c.Assert(req.Organization, qt.Equals, "planetscale")
+			c.Assert(req.Database, qt.Equals, "planetscale")
+			c.Assert(req.Branch, qt.Equals, "main")
+			c.Assert(req.Name, qt.Equals, "events")
+			c.Assert(req.ClusterSize, qt.Equals, "PS_40")
+			c.Assert(req.ExtraReplicas, qt.Equals, 0)
+			c.Assert(req.Shards, qt.IsNil)
+			c.Assert(req.CustomShards, qt.DeepEquals, []ps.CreateKeyspaceCustomShard{
+				{KeyRange: "-40", ClusterSize: "PS_80"},
+				{KeyRange: "40-80"},
+				{KeyRange: "80-", ClusterSize: "PS_20"},
+			})
+			return ks, nil
+		},
+	}
+	ch := newCreateTestHelper(c, &buf, svc)
+
+	cmd := CreateCmd(ch)
+	cmd.SetArgs([]string{
+		"planetscale", "main", "events",
+		"--cluster-size", "PS_40",
+		"--custom-shard", "key-range=-40,cluster-size=PS-80",
+		"--custom-shard", "key-range=40-80",
+		"--custom-shard", " key-range = 80- , cluster-size = PS_20 ",
+	})
+	c.Assert(cmd.Execute(), qt.IsNil)
+	c.Assert(svc.CreateFnInvoked, qt.IsTrue)
+	c.Assert(buf.String(), qt.JSONEquals, ks)
+}
+
+func TestKeyspace_CreateCmdCustomShardsWithWait(t *testing.T) {
+	c := qt.New(t)
+
+	var buf bytes.Buffer
+	pending := &ps.Keyspace{Name: "events", Ready: false}
+	ready := &ps.Keyspace{Name: "events", Ready: true, Shards: 2, Sharded: true}
+
+	svc := &mock.KeyspacesService{
+		CreateFn: func(_ context.Context, req *ps.CreateKeyspaceRequest) (*ps.Keyspace, error) {
+			c.Assert(req.Shards, qt.IsNil)
+			c.Assert(req.CustomShards, qt.DeepEquals, []ps.CreateKeyspaceCustomShard{
+				{KeyRange: "-80"},
+				{KeyRange: "80-"},
+			})
+			return pending, nil
+		},
+		GetFn: func(_ context.Context, req *ps.GetKeyspaceRequest) (*ps.Keyspace, error) {
+			c.Assert(req.Keyspace, qt.Equals, "events")
+			return ready, nil
+		},
+	}
+	ch := newCreateTestHelper(c, &buf, svc)
+
+	cmd := CreateCmd(ch)
+	cmd.SetArgs([]string{
+		"planetscale", "main", "events",
+		"--cluster-size", "PS_40",
+		"--custom-shard", "key-range=-80",
+		"--custom-shard", "key-range=80-",
+		"--wait",
+	})
+	c.Assert(cmd.Execute(), qt.IsNil)
+	c.Assert(svc.CreateFnInvoked, qt.IsTrue)
+	c.Assert(svc.GetFnInvoked, qt.IsTrue)
+	c.Assert(buf.String(), qt.JSONEquals, ready)
+}
+
+func TestKeyspace_CreateCmdCustomShardsErrors(t *testing.T) {
+	tests := []struct {
+		name string
+		args []string
+		want string
+	}{
+		{
+			name: "shards conflict",
+			args: []string{"--shards", "3", "--custom-shard", "key-range=-80", "--custom-shard", "key-range=80-"},
+			want: `--shards cannot be combined with --custom-shard`,
+		},
+		{
+			name: "shards conflict with default value set explicitly",
+			args: []string{"--shards", "1", "--custom-shard", "key-range=-80", "--custom-shard", "key-range=80-"},
+			want: `--shards cannot be combined with --custom-shard`,
+		},
+		{
+			name: "fewer than two shards",
+			args: []string{"--custom-shard", "key-range=-"},
+			want: `--custom-shard requires at least two shards, got 1`,
+		},
+		{
+			name: "missing key-range",
+			args: []string{"--custom-shard", "cluster-size=PS_80", "--custom-shard", "key-range=80-"},
+			want: `--custom-shard "cluster-size=PS_80" is missing required field key-range`,
+		},
+		{
+			name: "empty entry",
+			args: []string{"--custom-shard", "", "--custom-shard", "key-range=80-"},
+			want: `--custom-shard "" is missing required field key-range`,
+		},
+		{
+			name: "unknown field",
+			args: []string{"--custom-shard", "key-range=-80,replicas=2", "--custom-shard", "key-range=80-"},
+			want: `unknown --custom-shard field "replicas" in "key-range=-80,replicas=2": valid fields are key-range, cluster-size`,
+		},
+		{
+			name: "field without value separator",
+			args: []string{"--custom-shard", "-80", "--custom-shard", "key-range=80-"},
+			want: `invalid --custom-shard "-80": expected key-range=<range>\[,cluster-size=<size>\]`,
+		},
+		{
+			name: "empty key-range",
+			args: []string{"--custom-shard", "key-range=", "--custom-shard", "key-range=80-"},
+			want: `invalid --custom-shard "key-range=": key-range must not be empty`,
+		},
+		{
+			name: "empty cluster-size",
+			args: []string{"--custom-shard", "key-range=-80,cluster-size=", "--custom-shard", "key-range=80-"},
+			want: `invalid --custom-shard "key-range=-80,cluster-size=": cluster-size must not be empty`,
+		},
+		{
+			name: "duplicate field",
+			args: []string{"--custom-shard", "key-range=-80,key-range=80-", "--custom-shard", "key-range=80-"},
+			want: `duplicate --custom-shard field "key-range" in "key-range=-80,key-range=80-"`,
+		},
+		{
+			name: "duplicate cluster-size field",
+			args: []string{"--custom-shard", "key-range=-80,cluster-size=PS_80,cluster-size=PS_20", "--custom-shard", "key-range=80-"},
+			want: `duplicate --custom-shard field "cluster-size" in "key-range=-80,cluster-size=PS_80,cluster-size=PS_20"`,
+		},
+		{
+			name: "duplicate key range",
+			args: []string{"--custom-shard", "key-range=-80", "--custom-shard", "key-range=-80,cluster-size=PS_20"},
+			want: `duplicate --custom-shard key-range "-80"`,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			c := qt.New(t)
+
+			var buf bytes.Buffer
+			svc := &mock.KeyspacesService{
+				CreateFn: func(_ context.Context, _ *ps.CreateKeyspaceRequest) (*ps.Keyspace, error) {
+					c.Fatal("Create must not be called when flag validation fails")
+					return nil, nil
+				},
+			}
+			ch := newCreateTestHelper(c, &buf, svc)
+
+			cmd := CreateCmd(ch)
+			cmd.SetOut(&bytes.Buffer{})
+			cmd.SetErr(&bytes.Buffer{})
+			cmd.SetArgs(append([]string{"planetscale", "main", "events", "--cluster-size", "PS_40"}, tt.args...))
+			c.Assert(cmd.Execute(), qt.ErrorMatches, tt.want)
+			c.Assert(svc.CreateFnInvoked, qt.IsFalse)
+		})
+	}
+}

internal/planetscale/keyspaces.go

@@ -72,14 +72,20 @@ type ListKeyspacesRequest struct {
 	Branch       string `json:"-"`
 }
 
+type CreateKeyspaceCustomShard struct {
+	KeyRange    string `json:"key_range"`
+	ClusterSize string `json:"cluster_size,omitempty"`
+}
+
 type CreateKeyspaceRequest struct {
-	Organization  string `json:"-"`
-	Database      string `json:"-"`
-	Branch        string `json:"-"`
-	Name          string `json:"name"`
-	ClusterSize   string `json:"cluster_size"`
-	ExtraReplicas int    `json:"extra_replicas"`
-	Shards        int    `json:"shards"`
+	Organization  string                      `json:"-"`
+	Database      string                      `json:"-"`
+	Branch        string                      `json:"-"`
+	Name          string                      `json:"name"`
+	ClusterSize   string                      `json:"cluster_size"`
+	ExtraReplicas int                         `json:"extra_replicas"`
+	Shards        *int                        `json:"shards,omitempty"`
+	CustomShards  []CreateKeyspaceCustomShard `json:"custom_shards,omitempty"`
 }
 
 type ExternalDatasource struct {