func TestAnnotatingExecuteBatchKeyspaceIdsMultipleIds(t *testing.T) { keyspace, shards := setUpSandboxWithTwoShards("TestAnnotatingExecuteBatchKeyspaceIdsMultipleIds") err := rpcVTGate.ExecuteBatchKeyspaceIds( context.Background(), []proto.BoundKeyspaceIdQuery{ proto.BoundKeyspaceIdQuery{ Sql: "INSERT INTO table () VALUES();", Keyspace: keyspace, KeyspaceIds: []key.KeyspaceId{ key.KeyspaceId([]byte{0x10}), key.KeyspaceId([]byte{0x15}), }, }, }, pb.TabletType_MASTER, false, nil, &proto.QueryResultList{}) if err != nil { t.Fatalf("want nil, got %v", err) } verifyBatchQueryAnnotatedAsUnfriendly( t, 1, // expectedNumQueries shards[0]) }
func (krval bvcKeyRange) eval(bv interface{}, op Operator, onMismatch bool) bool { switch op { case QR_IN: switch num, status := getuint64(bv); status { case QR_OK: k := key.Uint64Key(num).KeyspaceId() return key.KeyRange(krval).Contains(k) case QR_OUT_OF_RANGE: return false } // Not a number. Check string. switch str, status := getstring(bv); status { case QR_OK: return key.KeyRange(krval).Contains(key.KeyspaceId(str)) } case QR_NOTIN: switch num, status := getuint64(bv); status { case QR_OK: k := key.Uint64Key(num).KeyspaceId() return !key.KeyRange(krval).Contains(k) case QR_OUT_OF_RANGE: return true } // Not a number. Check string. switch str, status := getstring(bv); status { case QR_OK: return !key.KeyRange(krval).Contains(key.KeyspaceId(str)) } default: panic("unexpected:") } return onMismatch }
func strKeyRange(start, end string) bvcKeyRange { kr := key.KeyRange{ Start: key.KeyspaceId(start), End: key.KeyspaceId(end), } return bvcKeyRange(kr) }
func siBytes(start, end string) *topo.ShardInfo { return topo.NewShardInfo("keyspace", start+"-"+end, &topo.Shard{ KeyRange: key.KeyRange{ Start: key.KeyspaceId(start), End: key.KeyspaceId(end), }, }, 0) }
func TestKeyspaceIdQuery(t *testing.T) { reflected, err := bson.Marshal(&reflectKeyspaceIdQuery{ Sql: "query", BindVariables: map[string]interface{}{"val": int64(1)}, Keyspace: "keyspace", KeyspaceIds: []kproto.KeyspaceId{kproto.KeyspaceId("10"), kproto.KeyspaceId("18")}, TabletType: "replica", Session: &commonSession, }) if err != nil { t.Error(err) } want := string(reflected) custom := KeyspaceIdQuery{ Sql: "query", BindVariables: map[string]interface{}{"val": int64(1)}, Keyspace: "keyspace", KeyspaceIds: []kproto.KeyspaceId{kproto.KeyspaceId("10"), kproto.KeyspaceId("18")}, TabletType: "replica", Session: &commonSession, } encoded, err := bson.Marshal(&custom) if err != nil { t.Error(err) } got := string(encoded) if want != got { t.Errorf("want\n%+v, got\n%+v", want, got) } var unmarshalled KeyspaceIdQuery err = bson.Unmarshal(encoded, &unmarshalled) if err != nil { t.Error(err) } if !reflect.DeepEqual(custom, unmarshalled) { t.Errorf("want \n%+v, got \n%+v", custom, unmarshalled) } extra, err := bson.Marshal(&extraKeyspaceIdQuery{}) if err != nil { t.Error(err) } err = bson.Unmarshal(extra, &unmarshalled) if err != nil { t.Error(err) } }
// Split will split the rows into subset for each distribution func (rs *RowSplitter) Split(result [][][]sqltypes.Value, rows [][]sqltypes.Value) error { if rs.Type == key.KIT_UINT64 { for _, row := range rows { v := sqltypes.MakeNumeric(row[rs.ValueIndex].Raw()) i, err := v.ParseUint64() if err != nil { return fmt.Errorf("Non numerical value: %v", err) } k := key.Uint64Key(i).KeyspaceId() for i, kr := range rs.KeyRanges { if kr.Contains(k) { result[i] = append(result[i], row) break } } } } else { for _, row := range rows { k := key.KeyspaceId(row[rs.ValueIndex].Raw()) for i, kr := range rs.KeyRanges { if kr.Contains(k) { result[i] = append(result[i], row) break } } } } return nil }
func TestHashAutoFail(t *testing.T) { _, err := vunhash(key.KeyspaceId("aa")) want := "invalid keyspace id: 6161" if err == nil || err.Error() != want { t.Errorf(`vunhash("aa"): %v, want %s`, err, want) } }
func TestHashAutoConvert(t *testing.T) { cases := []struct { in int64 out string }{ {1, "\x16k@\xb4J\xbaK\xd6"}, {0, "\x8c\xa6M\xe9\xc1\xb1#\xa7"}, {11, "\xae\xfcDI\x1c\xfeGL"}, {0x100000000000000, "\r\x9f'\x9b\xa5\xd8r`"}, {0x800000000000000, " \xb9\xe7g\xb2\xfb\x14V"}, {11, "\xae\xfcDI\x1c\xfeGL"}, {0, "\x8c\xa6M\xe9\xc1\xb1#\xa7"}, } for _, c := range cases { got := string(vhash(c.in)) want := c.out if got != want { t.Errorf("vhash(%d): %#v, want %q", c.in, got, want) } back, err := vunhash(key.KeyspaceId(got)) if err != nil { t.Error(err) } if back != c.in { t.Errorf("vunhash(%q): %d, want %d", got, back, c.in) } } }
func vhash(shardKey uint64) key.KeyspaceId { var keybytes, hashed [8]byte binary.BigEndian.PutUint64(keybytes[:], shardKey) encrypter := cipher.NewCBCEncrypter(block3DES, iv3DES) encrypter.CryptBlocks(hashed[:], keybytes[:]) return key.KeyspaceId(hashed[:]) }
// Verify returns true if id and ksid match. func (Numeric) Verify(_ planbuilder.VCursor, id interface{}, ksid key.KeyspaceId) (bool, error) { var keybytes [8]byte num, err := getNumber(id) if err != nil { return false, fmt.Errorf("Numeric.Verify: %v", err) } binary.BigEndian.PutUint64(keybytes[:], uint64(num)) return key.KeyspaceId(keybytes[:]) == ksid, nil }
func getShardForKeyspaceID(allShards []topo.ShardReference, keyspaceID []byte) (string, error) { if len(allShards) == 0 { return "", fmt.Errorf("No shards found for this tabletType") } for _, shardReference := range allShards { if shardReference.KeyRange.Contains(key.KeyspaceId(string(keyspaceID))) { return shardReference.Name, nil } } return "", fmt.Errorf("KeyspaceId %v didn't match any shards %+v", hex.EncodeToString(keyspaceID), allShards) }
func TestVTGateBuildEntityIds(t *testing.T) { shardMap := make(map[string][]key.KeyspaceId) shardMap["-20"] = []key.KeyspaceId{key.KeyspaceId("0"), key.KeyspaceId("1")} shardMap["20-40"] = []key.KeyspaceId{key.KeyspaceId("30")} sql := "select a from table where id=:id" entityColName := "kid" bindVar := make(map[string]interface{}) bindVar["id"] = 10 shards, sqls, bindVars := buildEntityIds(shardMap, sql, entityColName, bindVar) wantShards := []string{"-20", "20-40"} wantSqls := map[string]string{ "-20": "select a from table where id=:id and kid in (:kid0, :kid1)", "20-40": "select a from table where id=:id and kid in (:kid0)", } wantBindVars := map[string]map[string]interface{}{ "-20": map[string]interface{}{"id": 10, "kid0": key.KeyspaceId("0"), "kid1": key.KeyspaceId("1")}, "20-40": map[string]interface{}{"id": 10, "kid0": key.KeyspaceId("30")}, } if !reflect.DeepEqual(wantShards, shards) { t.Errorf("want %+v, got %+v", wantShards, shards) } if !reflect.DeepEqual(wantSqls, sqls) { t.Errorf("want %+v, got %+v", wantSqls, sqls) } if !reflect.DeepEqual(wantBindVars, bindVars) { t.Errorf("want %+v, got %+v", wantBindVars, bindVars) } }
// Map returns the associated keyspae ids for the given ids. func (Numeric) Map(_ planbuilder.VCursor, ids []interface{}) ([]key.KeyspaceId, error) { var keybytes [8]byte out := make([]key.KeyspaceId, 0, len(ids)) for _, id := range ids { num, err := getNumber(id) if err != nil { return nil, fmt.Errorf("Numeric.Map: %v", err) } binary.BigEndian.PutUint64(keybytes[:], uint64(num)) out = append(out, key.KeyspaceId(keybytes[:])) } return out, nil }
func TestAnnotatingExecuteBatchKeyspaceIds(t *testing.T) { keyspace, shards := setUpSandboxWithTwoShards("TestAnnotatingExecuteBatchKeyspaceIds") err := rpcVTGate.ExecuteBatchKeyspaceIds( context.Background(), []proto.BoundKeyspaceIdQuery{ proto.BoundKeyspaceIdQuery{ Sql: "INSERT INTO table () VALUES();", Keyspace: keyspace, KeyspaceIds: []key.KeyspaceId{key.KeyspaceId([]byte{0x10})}, }, proto.BoundKeyspaceIdQuery{ Sql: "UPDATE table SET col1=1 WHERE col2>3;", Keyspace: keyspace, KeyspaceIds: []key.KeyspaceId{key.KeyspaceId([]byte{0x15})}, }, proto.BoundKeyspaceIdQuery{ Sql: "DELETE FROM table WHERE col1==4;", Keyspace: keyspace, KeyspaceIds: []key.KeyspaceId{key.KeyspaceId([]byte{0x25})}, }, }, pb.TabletType_MASTER, false, nil, &proto.QueryResultList{}) if err != nil { t.Fatalf("want nil, got %v", err) } verifyBatchQueryAnnotatedWithKeyspaceIds( t, [][]byte{[]byte{0x10}, []byte{0x15}}, shards[0]) verifyBatchQueryAnnotatedWithKeyspaceIds( t, [][]byte{[]byte{0x25}}, shards[1]) }
func checkWanted(t *testing.T, got, expected []pair) { if len(got) != len(expected) { t.Fatalf("Wrong number of records: expected %v, got %v", len(expected), len(got)) } for i, wanted := range expected { if got[i].kid != key.KeyspaceId(wanted.kid) { t.Errorf("Wrong keyspace_id: expected %#v, got %#v", wanted.kid, got[i].kid) } if got[i].line != wanted.line { t.Errorf("Wrong line: expected %q got %q", wanted.line, got[i].line) } } }
func getKeyRange(keys []interface{}) (*pb.KeyRange, error) { var ksids []key.KeyspaceId for _, k := range keys { switch k := k.(type) { case string: ksids = append(ksids, key.KeyspaceId(k)) default: return nil, fmt.Errorf("expecting strings for keyrange: %+v", keys) } } return &pb.KeyRange{ Start: []byte(ksids[0]), End: []byte(ksids[1]), }, nil }
// ProtoToEntityIds converts an array of EntityId from proto3 func ProtoToEntityIds(l []*pb.ExecuteEntityIdsRequest_EntityId) []EntityId { if len(l) == 0 { return nil } result := make([]EntityId, len(l)) for i, e := range l { result[i].KeyspaceID = key.KeyspaceId(e.KeyspaceId) bv := &pbq.BindVariable{ Type: e.XidType, Value: e.XidValue, } v, err := tproto.BindVariableToNative(bv) if err != nil { panic(err) } result[i].ExternalID = v } return result }
func TestCSVSplitter(t *testing.T) { // mean.csv was generated using "select keyspaced_id, // tablename.* into outfile". file, err := os.Open("mean.csv") if err != nil { t.Fatalf("Cannot open mean.csv: %v", err) } r := NewKeyspaceCSVReader(file, ',') keyspaceIds := make([]pair, 0) for { kid, line, err := r.ReadRecord() if err == io.EOF { break } if err != nil { t.Fatalf("Unexpected error: %v", err) } keyspaceIds = append(keyspaceIds, pair{kid, string(line)}) } wantedTable := []pair{ {key.Uint64Key(1).KeyspaceId(), "\"x\x9c\xf3H\xcd\xc9\xc9W(\xcf/\xcaI\x01\\0\x18\xab\x04=\\0\\0\\0\\0\\0\\0\\0\\0\\0\\0\\0\\0\\0\",1,\"ala\\\nhas a cat\\\n\",1\n"}, {key.Uint64Key(2).KeyspaceId(), "\"x\x9c\xf3\xc8\xcfIT\xc8-\xcdK\xc9\a\\0\x13\xfe\x03\xc8\\0\\0\\0\\0\\0\\0\\0\\0\\0\\0\\0\\0\\0\\0\",2,\"ala\\\ntiene un gato\\\\\\\n\r\\\n\",2\n"}, {key.Uint64Key(3).KeyspaceId(), "\"x\x9cs\xceL\xccW\xc8\xcd\xcfK\xc9\a\\0\x13\x88\x03\xba\\0\\0\\0\\0\\0\\0\\0\\0\\0\\0\\0\\0\\0\\0\",3,\"ala\\\nha un gatto\\\\n\\\n\",3\n"}, {key.Uint64Key(4).KeyspaceId(), "\"x\x9cs\xca\xcf\xcb\xca/-R\xc8\xcd\xcfKI\x05\\0#:\x05\x13\\0\\0\\0\\0\\0\\0\\0\\0\\0\\0\\0\",3,\",,,ala \\\"\\\n,a un chat\",4\n"}, } for i, wanted := range wantedTable { if keyspaceIds[i].kid != key.KeyspaceId(wanted.kid) { t.Errorf("Wrong keyspace_id: expected %#v, got %#v", wanted.kid, keyspaceIds[i].kid) } if keyspaceIds[i].line != wanted.line { t.Errorf("Wrong line: expected %q got %q", wanted.line, keyspaceIds[i].line) } } if count := len(keyspaceIds); count != 4 { t.Errorf("Wrong number of records: expected 4, got %v", count) } }
// ProtoToEntityIds converts an array of EntityId from proto3 func ProtoToEntityIds(l []*pb.ExecuteEntityIdsRequest_EntityId) []EntityId { if len(l) == 0 { return nil } result := make([]EntityId, len(l)) for i, e := range l { result[i].KeyspaceID = key.KeyspaceId(e.KeyspaceId) switch e.XidType { case pb.ExecuteEntityIdsRequest_EntityId_TYPE_BYTES: result[i].ExternalID = e.XidBytes case pb.ExecuteEntityIdsRequest_EntityId_TYPE_INT: result[i].ExternalID = e.XidInt case pb.ExecuteEntityIdsRequest_EntityId_TYPE_UINT: result[i].ExternalID = e.XidUint case pb.ExecuteEntityIdsRequest_EntityId_TYPE_FLOAT: result[i].ExternalID = e.XidFloat default: panic(fmt.Errorf("Unsupported XidType %v", e.XidType)) } } return result }
func buildBindVarCondition(bvc interface{}) (name string, onAbsent, onMismatch bool, op Operator, value interface{}, err error) { bvcinfo, ok := bvc.(map[string]interface{}) if !ok { err = NewTabletError(FAIL, "Expecting json object for bind var conditions") return } var v interface{} v, ok = bvcinfo["Name"] if !ok { err = NewTabletError(FAIL, "Name missing in BindVarConds") return } name, ok = v.(string) if !ok { err = NewTabletError(FAIL, "Expecting string for Name in BindVarConds") return } v, ok = bvcinfo["OnAbsent"] if !ok { err = NewTabletError(FAIL, "OnAbsent missing in BindVarConds") return } onAbsent, ok = v.(bool) if !ok { err = NewTabletError(FAIL, "Expecting bool for OnAbsent") return } v, ok = bvcinfo["Operator"] if !ok { err = NewTabletError(FAIL, "Operator missing in BindVarConds") return } strop, ok := v.(string) if !ok { err = NewTabletError(FAIL, "Expecting string for Operator") return } op, ok = opmap[strop] if !ok { err = NewTabletError(FAIL, "Invalid Operator %s", strop) return } if op == QR_NOOP { return } v, ok = bvcinfo["Value"] if !ok { err = NewTabletError(FAIL, "Value missing in BindVarConds") return } if op >= QR_EQ && op <= QR_LE { strvalue, ok := v.(string) if !ok { err = NewTabletError(FAIL, "Expecting string: %v", v) return } if strop[0] == 'U' { value, err = strconv.ParseUint(strvalue, 0, 64) if err != nil { err = NewTabletError(FAIL, "Expecting uint64: %s", strvalue) return } } else if strop[0] == 'I' { value, err = strconv.ParseInt(strvalue, 0, 64) if err != nil { err = NewTabletError(FAIL, "Expecting int64: %s", strvalue) return } } else if strop[0] == 'S' { value = strvalue } else { panic("unexpected") } } else if op == QR_MATCH || op == QR_NOMATCH { strvalue, ok := v.(string) if !ok { err = NewTabletError(FAIL, "Expecting string: %v", v) return } value = strvalue } else if op == QR_IN || op == QR_NOTIN { kr, ok := v.(map[string]interface{}) if !ok { err = NewTabletError(FAIL, "Expecting keyrange for Value") return } var keyrange key.KeyRange strstart, ok := kr["Start"] if !ok { err = NewTabletError(FAIL, "Start missing in KeyRange") return } start, ok := strstart.(string) if !ok { err = NewTabletError(FAIL, "Expecting string for Start") return } keyrange.Start = key.KeyspaceId(start) strend, ok := kr["End"] if !ok { err = NewTabletError(FAIL, "End missing in KeyRange") return } end, ok := strend.(string) if !ok { err = NewTabletError(FAIL, "Expecting string for End") return } keyrange.End = key.KeyspaceId(end) value = keyrange } v, ok = bvcinfo["OnMismatch"] if !ok { err = NewTabletError(FAIL, "OnMismatch missing in BindVarConds") return } onMismatch, ok = v.(bool) if !ok { err = NewTabletError(FAIL, "Expecting bool for OnAbsent") return } return }
func TestKeyspaceIdBatchQuery(t *testing.T) { reflected, err := bson.Marshal(&reflectKeyspaceIdBatchQuery{ Queries: []reflectBoundKeyspaceIdQuery{{ Sql: "query", BindVariables: map[string]interface{}{"val": int64(1)}, Keyspace: "keyspace", KeyspaceIds: []kproto.KeyspaceId{kproto.KeyspaceId("10"), kproto.KeyspaceId("20")}, }}, Session: &Session{InTransaction: true, ShardSessions: []*ShardSession{{ Keyspace: "a", Shard: "0", TabletType: topo.TabletType("replica"), TransactionId: 1, }, { Keyspace: "b", Shard: "1", TabletType: topo.TabletType("master"), TransactionId: 2, }}, }, }) if err != nil { t.Error(err) } want := string(reflected) custom := KeyspaceIdBatchQuery{ Queries: []BoundKeyspaceIdQuery{{ Sql: "query", BindVariables: map[string]interface{}{"val": int64(1)}, Keyspace: "keyspace", KeyspaceIds: []kproto.KeyspaceId{kproto.KeyspaceId("10"), kproto.KeyspaceId("20")}, }}, Session: &commonSession, } encoded, err := bson.Marshal(&custom) if err != nil { t.Error(err) } got := string(encoded) if want != got { t.Errorf("want\n%+v, got\n%+v", want, got) } var unmarshalled KeyspaceIdBatchQuery err = bson.Unmarshal(encoded, &unmarshalled) if err != nil { t.Error(err) } if !reflect.DeepEqual(custom, unmarshalled) { t.Errorf("want \n%+v, got \n%+v", custom, unmarshalled) } extra, err := bson.Marshal(&extraKeyspaceIdBatchQuery{}) if err != nil { t.Error(err) } err = bson.Unmarshal(extra, &unmarshalled) if err != nil { t.Error(err) } }
func vhash(shardKey int64) key.KeyspaceId { var keybytes, hashed [8]byte binary.BigEndian.PutUint64(keybytes[:], uint64(shardKey)) block3DES.Encrypt(hashed[:], keybytes[:]) return key.KeyspaceId(hashed[:]) }
// KeyRangeFilterFunc returns a function that calls sendReply only if statements // in the transaction match the specified keyrange. The resulting function can be // passed into the BinlogStreamer: bls.Stream(file, pos, sendTransaction) -> // bls.Stream(file, pos, KeyRangeFilterFunc(sendTransaction)) func KeyRangeFilterFunc(kit key.KeyspaceIdType, keyrange key.KeyRange, sendReply sendTransactionFunc) sendTransactionFunc { isInteger := true if kit == key.KIT_BYTES { isInteger = false } return func(reply *proto.BinlogTransaction) error { matched := false filtered := make([]proto.Statement, 0, len(reply.Statements)) for _, statement := range reply.Statements { switch statement.Category { case proto.BL_SET: filtered = append(filtered, statement) case proto.BL_DDL: filtered = append(filtered, statement) matched = true case proto.BL_DML: keyspaceIndex := bytes.LastIndex(statement.Sql, KEYSPACE_ID_COMMENT) if keyspaceIndex == -1 { updateStreamErrors.Add("KeyRangeStream", 1) log.Errorf("Error parsing keyspace id: %s", string(statement.Sql)) continue } idstart := keyspaceIndex + len(KEYSPACE_ID_COMMENT) idend := bytes.Index(statement.Sql[idstart:], SPACE) if idend == -1 { updateStreamErrors.Add("KeyRangeStream", 1) log.Errorf("Error parsing keyspace id: %s", string(statement.Sql)) continue } textId := string(statement.Sql[idstart : idstart+idend]) if isInteger { id, err := strconv.ParseUint(textId, 10, 64) if err != nil { updateStreamErrors.Add("KeyRangeStream", 1) log.Errorf("Error parsing keyspace id: %s", string(statement.Sql)) continue } if !keyrange.Contains(key.Uint64Key(id).KeyspaceId()) { continue } } else { data, err := base64.StdEncoding.DecodeString(textId) if err != nil { updateStreamErrors.Add("KeyRangeStream", 1) log.Errorf("Error parsing keyspace id: %s", string(statement.Sql)) continue } if !keyrange.Contains(key.KeyspaceId(data)) { continue } } filtered = append(filtered, statement) matched = true case proto.BL_UNRECOGNIZED: updateStreamErrors.Add("KeyRangeStream", 1) log.Errorf("Error parsing keyspace id: %s", string(statement.Sql)) continue } } if matched { reply.Statements = filtered } else { reply.Statements = nil } return sendReply(reply) } }
BindVariables: map[string]interface{}{ "bind1": int64(0), }, Keyspace: "ks", Shards: []string{"1", "2"}, TabletType: topo.TYPE_RDONLY, Session: nil, }, keyspaceIdQuery: &proto.KeyspaceIdQuery{ Sql: "request1", BindVariables: map[string]interface{}{ "bind1": int64(0), }, Keyspace: "ks", KeyspaceIds: []key.KeyspaceId{ key.KeyspaceId("a"), }, TabletType: topo.TYPE_RDONLY, Session: nil, }, keyRangeQuery: &proto.KeyRangeQuery{ Sql: "request1", BindVariables: map[string]interface{}{ "bind1": int64(0), }, Keyspace: "ks", KeyRanges: []key.KeyRange{ key.KeyRange{ Start: key.KeyspaceId("s"), End: key.KeyspaceId("e"), },
// Use of this source code is governed by a BSD-style // license that can be found in the LICENSE file. package binlog import ( "fmt" "testing" "github.com/youtube/vitess/go/vt/binlog/proto" "github.com/youtube/vitess/go/vt/key" myproto "github.com/youtube/vitess/go/vt/mysqlctl/proto" ) var testKeyRange = key.KeyRange{ Start: key.KeyspaceId(key.Uint64Key(0).String()), End: key.KeyspaceId(key.Uint64Key(10).String()), } func TestKeyRangeFilterPass(t *testing.T) { input := proto.BinlogTransaction{ Statements: []proto.Statement{ { Category: proto.BL_SET, Sql: []byte("set1"), }, { Category: proto.BL_DML, Sql: []byte("dml1 /* EMD keyspace_id:20 */"), }, { Category: proto.BL_DML, Sql: []byte("dml2 /* EMD keyspace_id:2 */"),