Support USE keyspace[keyspace_id] (#3666)

Signed-off-by: Adam Saponara <[email protected]>
vitessio · Nov 29, 2018 · 3f1ef52 · 3f1ef52
1 parent 9d0594c
commit 3f1ef52
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 9 deletions.
diff --git a/go/vt/topo/topoproto/destination.go b/go/vt/topo/topoproto/destination.go
@@ -17,6 +17,7 @@ limitations under the License.
 package topoproto
 
 import (
+	"encoding/hex"
 	"strings"
 
 	"vitess.io/vitess/go/vt/key"
@@ -45,23 +46,32 @@ func ParseDestination(targetString string, defaultTabletType topodatapb.TabletTy
 		dest = key.DestinationShard(targetString[last+1:])
 		targetString = targetString[:last]
 	}
-	// Try to parse it as a range
+	// Try to parse it as a keyspace id or range
 	last = strings.LastIndexAny(targetString, "[")
 	if last != -1 {
 		rangeEnd := strings.LastIndexAny(targetString, "]")
 		if rangeEnd == -1 {
 			return keyspace, tabletType, dest, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "invalid key range provided. Couldn't find range end ']'")
-
 		}
 		rangeString := targetString[last+1 : rangeEnd]
-		keyRange, err := key.ParseShardingSpec(rangeString)
-		if err != nil {
-			return keyspace, tabletType, dest, err
-		}
-		if len(keyRange) != 1 {
-			return keyspace, tabletType, dest, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "single keyrange expected in %s", rangeString)
+		if strings.Contains(rangeString, "-") {
+			// Parse as range
+			keyRange, err := key.ParseShardingSpec(rangeString)
+			if err != nil {
+				return keyspace, tabletType, dest, err
+			}
+			if len(keyRange) != 1 {
+				return keyspace, tabletType, dest, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "single keyrange expected in %s", rangeString)
+			}
+			dest = key.DestinationExactKeyRange{KeyRange: keyRange[0]}
+		} else {
+			// Parse as keyspace id
+			destBytes, err := hex.DecodeString(rangeString)
+			if err != nil {
+				return keyspace, tabletType, dest, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "expected valid hex in keyspace id %s", rangeString)
+			}
+			dest = key.DestinationKeyspaceID(destBytes)
 		}
-		dest = key.DestinationExactKeyRange{KeyRange: keyRange[0]}
 		targetString = targetString[:last]
 	}
 	keyspace = targetString

diff --git a/go/vt/topo/topoproto/destination_test.go b/go/vt/topo/topoproto/destination_test.go
@@ -45,6 +45,11 @@ func TestParseDestination(t *testing.T) {
 		keyspace:     "ks",
 		tabletType:   topodatapb.TabletType_MASTER,
 		dest:         key.DestinationExactKeyRange{KeyRange: &topodatapb.KeyRange{}},
+	}, {
+		targetString: "ks[deadbeef]@master",
+		keyspace:     "ks",
+		tabletType:   topodatapb.TabletType_MASTER,
+		dest:         key.DestinationKeyspaceID([]byte("\xde\xad\xbe\xef")),
 	}, {
 		targetString: "ks[10-]@master",
 		keyspace:     "ks",
@@ -109,4 +114,10 @@ func TestParseDestination(t *testing.T) {
 	if err == nil || err.Error() != want {
 		t.Errorf("executorExec error: %v, want %s", err, want)
 	}
+
+	_, _, _, err = ParseDestination("ks[qrnqorrs]@master", topodatapb.TabletType_MASTER)
+	want = "expected valid hex in keyspace id qrnqorrs"
+	if err == nil || err.Error() != want {
+		t.Errorf("executorExec error: %v, want %s", err, want)
+	}
 }