// GetRangeTree fetches the RangeTree proto and sets up the range tree context.
func getRangeTree(txn *client.Txn) (*treeContext, error) {
tree := new(RangeTree)
if err := txn.GetProto(keys.RangeTreeRoot, tree); err != nil {
return nil, err
}
return &treeContext{
txn: txn,
tree: tree,
dirty: false,
nodes: map[string]cachedNode{},
}, nil
}
// truncateTable truncates the data of a table.
// It deletes a range of data for the table, which includes the PK and all
// indexes.
func truncateTable(tableDesc *sqlbase.TableDescriptor, txn *client.Txn) error {
tablePrefix := keys.MakeTablePrefix(uint32(tableDesc.ID))
// Delete rows and indexes starting with the table's prefix.
tableStartKey := roachpb.Key(tablePrefix)
tableEndKey := tableStartKey.PrefixEnd()
if log.V(2) {
log.Infof("DelRange %s - %s", tableStartKey, tableEndKey)
}
b := client.Batch{}
b.DelRange(tableStartKey, tableEndKey, false)
return txn.Run(&b)
}
// incCmd adds one to the value of c.key in the env (as determined by
// a previous read or write, or else assumed to be zero) and writes it
// to the db.
func incCmd(c *cmd, txn *client.Txn, t *testing.T) error {
val, ok := c.env[c.key]
if !ok {
panic(fmt.Sprintf("can't increment key %q; not yet read", c.key))
}
r := val + 1
if err := txn.Put(c.getKey(), r); err != nil {
return err
}
c.env[c.key] = r
c.debug = fmt.Sprintf("[%d]", r)
return nil
}
// readCmd reads a value from the db and stores it in the env.
func readCmd(c *cmd, txn *client.Txn, t *testing.T) error {
r, err := txn.Get(c.getKey())
if err != nil {
return err
}
var value int64
if r.Value != nil {
value = r.ValueInt()
}
c.env[c.key] = value
c.debug = fmt.Sprintf("[%d]", value)
return nil
}
// GetTableDescFromID retrieves the table descriptor for the table
// ID passed in using an existing txn. Teturns an error if the
// descriptor doesn't exist or if it exists and is not a table.
func GetTableDescFromID(txn *client.Txn, id ID) (*TableDescriptor, error) {
desc := &Descriptor{}
descKey := MakeDescMetadataKey(id)
if err := txn.GetProto(descKey, desc); err != nil {
return nil, err
}
table := desc.GetTable()
if table == nil {
return nil, ErrDescriptorNotFound
}
return table, nil
}
// writeCmd sums values from the env (and possibly numeric constants)
// and writes the value to the db. "c.endKey" here needs to be parsed
// in the context of this command, which is a "+"-separated list of
// keys from the env or numeric constants to sum.
func writeCmd(c *cmd, txn *client.Txn, t *testing.T) error {
sum := int64(0)
for _, sp := range strings.Split(c.endKey, "+") {
if constant, err := strconv.Atoi(sp); err != nil {
sum += c.env[sp]
} else {
sum += int64(constant)
}
}
err := txn.Put(c.getKey(), sum)
c.debug = fmt.Sprintf("[%d]", sum)
return err
}
// resolveName resolves a table name to a descriptor ID by looking in the
// database. If the mapping is not found, sqlbase.ErrDescriptorNotFound is returned.
func (m *LeaseManager) resolveName(
txn *client.Txn, dbID sqlbase.ID, tableName string,
) (sqlbase.ID, error) {
nameKey := tableKey{dbID, tableName}
key := nameKey.Key()
gr, err := txn.Get(key)
if err != nil {
return 0, err
}
if !gr.Exists() {
return 0, sqlbase.ErrDescriptorNotFound
}
return sqlbase.ID(gr.ValueInt()), nil
}
// scanCmd reads the values from the db from [key, endKey).
func scanCmd(c *cmd, txn *client.Txn, t *testing.T) error {
rows, err := txn.Scan(c.getKey(), c.getEndKey(), 0)
if err != nil {
return err
}
var vals []string
keyPrefix := []byte(fmt.Sprintf("%d.", c.historyIdx))
for _, kv := range rows {
key := bytes.TrimPrefix(kv.Key, keyPrefix)
c.env[string(key)] = kv.ValueInt()
vals = append(vals, fmt.Sprintf("%d", kv.ValueInt()))
}
c.debug = fmt.Sprintf("[%s]", strings.Join(vals, " "))
return nil
}
func allRangeDescriptors(txn *client.Txn) ([]roachpb.RangeDescriptor, error) {
// TODO(dan): Iterate with some batch size.
rows, err := txn.Scan(keys.Meta2Prefix, keys.MetaMax, 0)
if err != nil {
return nil, errors.Wrap(err, "unable to scan range descriptors")
}
rangeDescs := make([]roachpb.RangeDescriptor, len(rows))
for i, row := range rows {
if err := row.ValueProto(&rangeDescs[i]); err != nil {
return nil, errors.Wrapf(err, "%s: unable to unmarshal range descriptor", row.Key)
}
}
return rangeDescs, nil
}
func TestRemoveLeaseIfExpiring(t *testing.T) {
defer leaktest.AfterTest(t)()
p := planner{}
mc := hlc.NewManualClock(0)
p.leaseMgr = &LeaseManager{LeaseStore: LeaseStore{clock: hlc.NewClock(mc.UnixNano)}}
p.leases = make([]*LeaseState, 0)
txn := client.Txn{}
p.setTxn(&txn)
if p.removeLeaseIfExpiring(nil) {
t.Error("expected false with nil input")
}
// Add a lease to the planner.
d := int64(LeaseDuration)
l1 := &LeaseState{expiration: parser.DTimestamp{Time: time.Unix(0, mc.UnixNano()+d+1)}}
p.leases = append(p.leases, l1)
et := hlc.Timestamp{WallTime: l1.Expiration().UnixNano()}
txn.UpdateDeadlineMaybe(et)
if p.removeLeaseIfExpiring(l1) {
t.Error("expected false wih a non-expiring lease")
}
if !p.txn.GetDeadline().Equal(et) {
t.Errorf("expected deadline %s but got %s", et, p.txn.GetDeadline())
}
// Advance the clock so that l1 will be expired.
mc.Increment(d + 1)
// Add another lease.
l2 := &LeaseState{expiration: parser.DTimestamp{Time: time.Unix(0, mc.UnixNano()+d+1)}}
p.leases = append(p.leases, l2)
if !p.removeLeaseIfExpiring(l1) {
t.Error("expected true with an expiring lease")
}
et = hlc.Timestamp{WallTime: l2.Expiration().UnixNano()}
txn.UpdateDeadlineMaybe(et)
if !(len(p.leases) == 1 && p.leases[0] == l2) {
t.Errorf("expected leases to contain %s but has %s", l2, p.leases)
}
if !p.txn.GetDeadline().Equal(et) {
t.Errorf("expected deadline %s, but got %s", et, p.txn.GetDeadline())
}
}
func allSQLDescriptors(txn *client.Txn) ([]sqlbase.Descriptor, error) {
startKey := roachpb.Key(keys.MakeTablePrefix(keys.DescriptorTableID))
endKey := startKey.PrefixEnd()
// TODO(dan): Iterate with some batch size.
rows, err := txn.Scan(startKey, endKey, 0)
if err != nil {
return nil, errors.Wrap(err, "unable to scan SQL descriptors")
}
sqlDescs := make([]sqlbase.Descriptor, len(rows))
for i, row := range rows {
if err := row.ValueProto(&sqlDescs[i]); err != nil {
return nil, errors.Wrapf(err, "%s: unable to unmarshal SQL descriptor", row.Key)
}
}
return sqlDescs, nil
}
// Import loads some data in sstables into the database. Only the keys between
// startKey and endKey are loaded.
func Import(
ctx context.Context,
sst engine.RocksDBSstFileReader,
txn *client.Txn,
startKey, endKey engine.MVCCKey,
) error {
var v roachpb.Value
importFunc := func(kv engine.MVCCKeyValue) (bool, error) {
v = roachpb.Value{RawBytes: kv.Value}
v.ClearChecksum()
if log.V(3) {
log.Infof(ctx, "Put %s %s\n", kv.Key.Key, v.PrettyPrint())
}
if err := txn.Put(kv.Key.Key, &v); err != nil {
return true, err
}
return false, nil
}
return sst.Iterate(startKey, endKey, importFunc)
}
func restoreTable(
ctx context.Context,
sst engine.RocksDBSstFileReader,
txn *client.Txn,
table *sqlbase.TableDescriptor,
overwrite bool,
) error {
log.Infof(ctx, "Restoring Table %q", table.Name)
tableStartKey := roachpb.Key(sqlbase.MakeIndexKeyPrefix(table, table.PrimaryIndex.ID))
tableEndKey := tableStartKey.PrefixEnd()
existingDesc, err := txn.Get(sqlbase.MakeDescMetadataKey(table.GetID()))
if err != nil {
return err
}
existingData, err := txn.Scan(tableStartKey, tableEndKey, 1)
if err != nil {
return err
}
if existingDesc.Value != nil || len(existingData) > 0 {
if overwrite {
// We're about to Put the descriptor, so don't bother deleting it.
if err := txn.DelRange(tableStartKey, tableEndKey); err != nil {
return err
}
} else {
return errors.Errorf("table %q already exists", table.Name)
}
}
tableDescKey := sqlbase.MakeDescMetadataKey(table.GetID())
if err := txn.Put(tableDescKey, sqlbase.WrapDescriptor(table)); err != nil {
return err
}
return Import(ctx, sst, txn, engine.MVCCKey{Key: tableStartKey}, engine.MVCCKey{Key: tableEndKey})
}
// Acquire a lease on the most recent version of a table descriptor.
// If the lease cannot be obtained because the descriptor is in the process of
// being deleted, the error will be errTableDeleted.
func (s LeaseStore) Acquire(
txn *client.Txn,
tableID sqlbase.ID,
minVersion sqlbase.DescriptorVersion,
minExpirationTime parser.DTimestamp,
) (*LeaseState, error) {
lease := &LeaseState{}
expiration := time.Unix(0, s.clock.Now().WallTime).Add(jitteredLeaseDuration())
if !minExpirationTime.IsZero() && expiration.Before(minExpirationTime.Time) {
expiration = minExpirationTime.Time
}
lease.expiration = parser.DTimestamp{Time: expiration}
// Use the supplied (user) transaction to look up the descriptor because the
// descriptor might have been created within the transaction.
p := makeInternalPlanner(txn, security.RootUser)
const getDescriptor = `SELECT descriptor FROM system.descriptor WHERE id = $1`
values, err := p.queryRow(getDescriptor, int(tableID))
if err != nil {
return nil, err
}
if values == nil {
return nil, sqlbase.ErrDescriptorNotFound
}
desc := &sqlbase.Descriptor{}
if err := proto.Unmarshal([]byte(*values[0].(*parser.DBytes)), desc); err != nil {
return nil, err
}
tableDesc := desc.GetTable()
if tableDesc == nil {
return nil, errors.Errorf("ID %d is not a table", tableID)
}
if err := filterTableState(tableDesc); err != nil {
return nil, err
}
tableDesc.MaybeUpgradeFormatVersion()
lease.TableDescriptor = *tableDesc
// ValidateTable instead of Validate, even though we have a txn available,
// so we don't block reads waiting for this lease.
if err := lease.ValidateTable(); err != nil {
return nil, err
}
if lease.Version < minVersion {
return nil, errors.Errorf("version %d of table %d does not exist yet", minVersion, tableID)
}
// Insert the entry in the lease table in a separate transaction. This is
// necessary because we want to ensure that the lease entry is added and the
// transaction passed to Acquire() might be aborted. The lease entry needs to
// be added because we store the returned LeaseState in local in-memory maps
// and cannot handle the entry being reverted. This is safe because either
// the descriptor we're acquiring the lease on existed prior to the acquire
// transaction in which case acquiring the lease is kosher, or the descriptor
// was created within the acquire transaction. The second case is more
// subtle. We might create a lease entry for a table that doesn't exist, but
// there is no harm in that as no other transaction will be attempting to
// modify the descriptor and even if the descriptor is never created we'll
// just have a dangling lease entry which will eventually get GC'd.
ctx := txn.Context // propagate context/trace to new transaction
err = s.db.Txn(context.TODO(), func(txn *client.Txn) error {
txn.Context = ctx
p := makeInternalPlanner(txn, security.RootUser)
const insertLease = `INSERT INTO system.lease (descID, version, nodeID, expiration) ` +
`VALUES ($1, $2, $3, $4)`
count, err := p.exec(insertLease, lease.ID, int(lease.Version), s.nodeID, &lease.expiration)
if err != nil {
return err
}
if count != 1 {
return errors.Errorf("%s: expected 1 result, found %d", insertLease, count)
}
return nil
})
return lease, err
}
// deleteCmd deletes the value at the given key from the db.
func deleteCmd(c *cmd, txn *client.Txn, t *testing.T) error {
return txn.Del(c.getKey())
}
func (ti *tableInserter) init(txn *client.Txn) error {
ti.txn = txn
ti.b = txn.NewBatch()
return nil
}
func (td *tableDeleter) init(txn *client.Txn) error {
td.txn = txn
td.b = txn.NewBatch()
return nil
}
func (tu *tableUpdater) init(txn *client.Txn) error {
tu.txn = txn
tu.b = txn.NewBatch()
return nil
}
// deleteRngCmd deletes the range of values from the db from [key, endKey).
func deleteRngCmd(c *cmd, txn *client.Txn, t *testing.T) error {
return txn.DelRange(c.getKey(), c.getEndKey())
}
// setTxnTimestamps sets the transaction's proto timestamps and deadline
// to ts. This is for use with AS OF queries, and should be called in the
// retry block (except in the case of prepare which doesn't use retry). The
// deadline-checking code checks that the `Timestamp` field of the proto
// hasn't exceeded the deadline. Since we set the Timestamp field each retry,
// it won't ever exceed the deadline, and thus setting the deadline here is
// not strictly needed. However, it doesn't do anything incorrect and it will
// possibly find problems if things change in the future, so it is left in.
func setTxnTimestamps(txn *client.Txn, ts hlc.Timestamp) {
txn.Proto.Timestamp = ts
txn.Proto.OrigTimestamp = ts
txn.Proto.MaxTimestamp = ts
txn.UpdateDeadlineMaybe(ts)
}
// commitCmd commits the transaction.
func commitCmd(c *cmd, txn *client.Txn, t *testing.T) error {
return txn.Commit()
}