pkg/apis/planetscale/v2: fix VitessKeyRange rounding logic

Signed-off-by: Matt Layher <[email protected]>

deploy/crds/planetscale.com_vitessclusters.yaml

@@ -1318,6 +1318,7 @@ spec:
                             properties:
                               parts:
                                 format: int32
+                                maximum: 65536
                                 minimum: 1
                                 type: integer
                               shardTemplate:

deploy/crds/planetscale.com_vitesskeyspaces.yaml

@@ -578,6 +578,7 @@ spec:
                       properties:
                         parts:
                           format: int32
+                          maximum: 65536
                           minimum: 1
                           type: integer
                         shardTemplate:

go.mod

@@ -1,6 +1,6 @@
 module planetscale.dev/vitess-operator
 
-go 1.17
+go 1.18
 
 require (
 	github.com/ahmetb/gen-crd-api-reference-docs v0.1.5-0.20190629210212-52e137b8d003

pkg/apis/planetscale/v2/vitesskeyspace_methods.go

@@ -18,32 +18,43 @@ package v2
 
 import (
 	"encoding/hex"
+	"fmt"
 	"sort"
+	"strings"
 	"time"
 
 	corev1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/util/sets"
-	"planetscale.dev/vitess-operator/pkg/operator/partitioning"
+	"vitess.io/vitess/go/vt/key"
 	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
 )
 
-/*
-KeyRanges returns the set of ranges for an equal partitioning.
-
-The business logic for figuring out the key ranges happens in `pkg/operator/partitioning`, and this method simply
-calls that function, which returns the key ranges as an array of raw vitess KeyRange type (start and end as []byte).
-This function then translates that into VitessKeyRanges (start and end as hex strings) and returns the []VitessKeyRange.
-*/
+// KeyRanges returns the set of ranges for an equal partitioning.
 func (p *VitessKeyspaceEqualPartitioning) KeyRanges() []VitessKeyRange {
-	byteKeyRanges := partitioning.EqualKeyRanges(uint64(p.Parts))
-	vitessKeyRanges := make([]VitessKeyRange, len(byteKeyRanges))
+	// Invariant: number of parts must be between 1-65536. This is enforced via
+	// the CRD.
+	ranges, err := key.GenerateShardRanges(int(p.Parts))
+	if err != nil {
+		panic(fmt.Sprintf("could not generate shard range with %d parts: %v", p.Parts, err))
+	}
 
-	for i, byteKeyRange := range byteKeyRanges {
-		vitessKeyRanges[i].FillFromByteKeyRange(byteKeyRange)
+	out := make([]VitessKeyRange, 0, len(ranges))
+	for _, r := range ranges {
+		// Invariant: key ranges are always in a format such as: "-", "-20",
+		// "20-40", or "40-".
+		start, end, ok := strings.Cut(r, "-")
+		if !ok {
+			panic(fmt.Sprintf("cannot create VitessKeyRange from shard range %q", r))
+		}
+
+		out = append(out, VitessKeyRange{
+			Start: start,
+			End:   end,
+		})
 	}
 
-	return vitessKeyRanges
+	return out
 }
 
 // FillFromByteKeyRange is a method to fill a VitessKeyRange from a topodatapb.KeyRange (encoding from []byte to

pkg/apis/planetscale/v2/vitesskeyspace_methods_test.go

@@ -22,9 +22,6 @@ import (
 )
 
 func TestTranslationToVitessKeyRange(t *testing.T) {
-	// Quick test - calling KeyRanges method should call KeyRanges function in partitioning package, and then
-	// translate raw Vitess KeyRange to VitessKeyRange.  Most of business logic tests are unit tested within
-	// partitioning package.
 	table := []struct {
 		parts int32
 		want  []VitessKeyRange
@@ -41,6 +38,30 @@ func TestTranslationToVitessKeyRange(t *testing.T) {
 				{"", "55"}, {"55", "aa"}, {"aa", ""},
 			},
 		},
+		{
+			parts: 6,
+			want: []VitessKeyRange{
+				{"", "2a"},
+				{"2a", "55"},
+				{"55", "80"},
+				{"80", "aa"},
+				{"aa", "d5"},
+				{"d5", ""},
+			},
+		},
+		{
+			parts: 8,
+			want: []VitessKeyRange{
+				{"", "20"},
+				{"20", "40"},
+				{"40", "60"},
+				{"60", "80"},
+				{"80", "a0"},
+				{"a0", "c0"},
+				{"c0", "e0"},
+				{"e0", ""},
+			},
+		},
 	}
 
 	for _, test := range table {

pkg/apis/planetscale/v2/vitesskeyspace_types.go

@@ -320,6 +320,7 @@ type VitessKeyspaceEqualPartitioning struct {
 	//          partitioning with the desired number of parts, perform a resharding
 	//          migration, and then remove the old partitioning.
 	// +kubebuilder:validation:Minimum=1
+	// +kubebuilder:validation:Maximum=65536
 	Parts int32 `json:"parts"`
 
 	// ShardTemplate is the configuration used for each equal-sized shard.