Golang client.Txn類代碼示例

// 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())
}