Golang client.Txn类代码示例

func (sc *SchemaChanger) maybeWriteResumeSpan(
	txn *client.Txn,
	version sqlbase.DescriptorVersion,
	resume roachpb.Span,
	mutationIdx int,
	lastCheckpoint *time.Time,
) error {
	checkpointInterval := checkpointInterval
	if sc.testingKnobs.WriteCheckpointInterval > 0 {
		checkpointInterval = sc.testingKnobs.WriteCheckpointInterval
	}
	if timeutil.Since(*lastCheckpoint) < checkpointInterval {
		return nil
	}
	tableDesc, err := sqlbase.GetTableDescFromID(txn, sc.tableID)
	if err != nil {
		return err
	}
	if tableDesc.Version != version {
		return errVersionMismatch
	}
	tableDesc.Mutations[mutationIdx].ResumeSpan = resume
	txn.SetSystemConfigTrigger()
	if err := txn.Put(sqlbase.MakeDescMetadataKey(tableDesc.GetID()),
		sqlbase.WrapDescriptor(tableDesc)); err != nil {
		return err
	}
	*lastCheckpoint = timeutil.Now()
	return nil
}

// GenerateUniqueDescID returns the next available Descriptor ID and increments
// the counter.
func GenerateUniqueDescID(txn *client.Txn) (sqlbase.ID, error) {
	// Increment unique descriptor counter.
	ir, err := txn.Inc(keys.DescIDGenerator, 1)
	if err != nil {
		return 0, err
	}
	return sqlbase.ID(ir.ValueInt() - 1), nil
}

func getDescriptorID(txn *client.Txn, key sqlbase.DescriptorKey) (sqlbase.ID, error) {
	// TODO(dan): Share this with (*planner).getDescriptor.
	idValue, err := txn.Get(key.Key())
	if err != nil {
		return 0, err
	}
	if !idValue.Exists() {
		return 0, errors.Errorf("no descriptor for key: %s", key)
	}
	return sqlbase.ID(idValue.ValueInt()), 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
}

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

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

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

// AllRangeDescriptors fetches all meta2 RangeDescriptor using the given txn.
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
}

// 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 restoreTableDesc(
	ctx context.Context,
	txn *client.Txn,
	database sqlbase.DatabaseDescriptor,
	table sqlbase.TableDescriptor,
) error {
	// Run getDescriptorID again to make sure the database hasn't been dropped
	// while we were importing.
	var err error
	if table.ParentID, err = getDescriptorID(txn, tableKey{name: database.Name}); err != nil {
		return err
	}
	tableIDKey := tableKey{parentID: table.ParentID, name: table.Name}.Key()
	tableDescKey := sqlbase.MakeDescMetadataKey(table.ID)

	// Check for an existing table.
	var existingDesc sqlbase.Descriptor
	existingIDKV, err := txn.Get(tableIDKey)
	if err != nil {
		return err
	}
	if existingIDKV.Value != nil {
		existingID, err := existingIDKV.Value.GetInt()
		if err != nil {
			return err
		}
		existingDescKV, err := txn.Get(sqlbase.MakeDescMetadataKey(sqlbase.ID(existingID)))
		if err != nil {
			return err
		}
		if err := existingDescKV.Value.GetProto(&existingDesc); err != nil {
			return err
		}
	}

	// Write the new descriptors. First the ID -> TableDescriptor for the new
	// table, then flip (or initialize) the name -> ID entry so any new queries
	// will use the new one. If there was an existing table, it can now be
	// cleaned up.
	b := txn.NewBatch()
	b.CPut(tableDescKey, sqlbase.WrapDescriptor(&table), nil)
	if existingTable := existingDesc.GetTable(); existingTable == nil {
		b.CPut(tableIDKey, table.ID, nil)
	} else {
		existingIDKV.Value.ClearChecksum()
		b.CPut(tableIDKey, table.ID, existingIDKV.Value)
		// TODO(dan): This doesn't work for interleaved tables. Fix it when we
		// fix the empty range interleaved table TODO below.
		existingDataPrefix := roachpb.Key(keys.MakeTablePrefix(uint32(existingTable.ID)))
		b.DelRange(existingDataPrefix, existingDataPrefix.PrefixEnd(), false)
		zoneKey, _, descKey := GetKeysForTableDescriptor(existingTable)
		// Delete the desc and zone entries. Leave the name because the new
		// table is using it.
		b.Del(descKey)
		b.Del(zoneKey)
	}
	return txn.Run(b)
}

func TestRemoveLeaseIfExpiring(t *testing.T) {
	defer leaktest.AfterTest(t)()

	p := planner{session: &Session{context: context.Background()}}
	mc := hlc.NewManualClock(123)
	p.leaseMgr = &LeaseManager{LeaseStore: LeaseStore{clock: hlc.NewClock(mc.UnixNano, time.Nanosecond)}}
	p.leases = make([]*LeaseState, 0)
	txn := client.Txn{Context: context.Background()}
	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 with 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
}

func getDescriptor(
	txn *client.Txn, plainKey sqlbase.DescriptorKey, descriptor sqlbase.DescriptorProto,
) (bool, error) {
	gr, err := txn.Get(plainKey.Key())
	if err != nil {
		return false, err
	}
	if !gr.Exists() {
		return false, nil
	}

	descKey := sqlbase.MakeDescMetadataKey(sqlbase.ID(gr.ValueInt()))
	desc := &sqlbase.Descriptor{}
	if err := txn.GetProto(descKey, desc); err != nil {
		return false, err
	}

	switch t := descriptor.(type) {
	case *sqlbase.TableDescriptor:
		table := desc.GetTable()
		if table == nil {
			return false, errors.Errorf("%q is not a table", plainKey.Name())
		}
		table.MaybeUpgradeFormatVersion()
		// TODO(dan): Write the upgraded TableDescriptor back to kv. This will break
		// the ability to use a previous version of cockroach with the on-disk data,
		// but it's worth it to avoid having to do the upgrade every time the
		// descriptor is fetched. Our current test for this enforces compatibility
		// backward and forward, so that'll have to be extended before this is done.
		if err := table.Validate(txn); err != nil {
			return false, err
		}
		*t = *table
	case *sqlbase.DatabaseDescriptor:
		database := desc.GetDatabase()
		if database == nil {
			return false, errors.Errorf("%q is not a database", plainKey.Name())
		}
		if err := database.Validate(); err != nil {
			return false, err
		}
		*t = *database
	}
	return true, nil
}

// InsertEventRecord inserts a single event into the event log as part of the
// provided transaction.
func (ev EventLogger) InsertEventRecord(
	txn *client.Txn, eventType EventLogType, targetID, reportingID int32, info interface{},
) error {
	// Record event record insertion in local log output.
	txn.AddCommitTrigger(func() {
		log.Infof(txn.Context, "Event: %q, target: %d, info: %+v",
			eventType,
			targetID,
			info)
	})

	const insertEventTableStmt = `
INSERT INTO system.eventlog (
  timestamp, eventType, targetID, reportingID, info
)
VALUES(
  $1, $2, $3, $4, $5
)
`
	args := []interface{}{
		ev.selectEventTimestamp(txn.Proto.Timestamp),
		eventType,
		targetID,
		reportingID,
		nil, // info
	}
	if info != nil {
		infoBytes, err := json.Marshal(info)
		if err != nil {
			return err
		}
		args[4] = string(infoBytes)
	}

	rows, err := ev.ExecuteStatementInTransaction("log-event", txn, insertEventTableStmt, args...)
	if err != nil {
		return err
	}
	if rows != 1 {
		return errors.Errorf("%d rows affected by log insertion; expected exactly one row affected.", rows)
	}
	return nil
}

// Ingest loads some data in an sstable into an empty range. Only the keys
// between startKey and endKey are loaded. If newTableID is non-zero, every
// row's key is rewritten to be for that table.
func Ingest(
	ctx context.Context,
	txn *client.Txn,
	path string,
	checksum uint32,
	startKey, endKey roachpb.Key,
	newTableID sqlbase.ID,
) error {
	// TODO(mjibson): An appropriate value for this should be determined. The
	// current value was guessed at but appears to work well.
	const batchSize = 10000

	// TODO(dan): Check if the range being ingested into is empty. If newTableID
	// is non-zero, it'll have to be derived from startKey and endKey.

	f, err := os.Open(path)
	if err != nil {
		return err
	}
	defer f.Close()
	crc := crc32.New(crc32.MakeTable(crc32.Castagnoli))
	if _, err := io.Copy(crc, f); err != nil {
		return nil
	}
	if c := crc.Sum32(); c != checksum {
		return errors.Errorf("%s: checksum mismatch got %d expected %d", path, c, checksum)
	}

	sst, err := engine.MakeRocksDBSstFileReader()
	if err != nil {
		return err
	}
	defer sst.Close()
	if err := sst.AddFile(path); err != nil {
		return err
	}

	b := txn.NewBatch()
	var v roachpb.Value
	count := 0
	ingestFunc := 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())
		}
		b.Put(kv.Key.Key, &v)
		count++
		if count > batchSize {
			if err := txn.Run(b); err != nil {
				return true, err
			}
			b = txn.NewBatch()
			count = 0
		}
		return false, nil
	}
	if newTableID != 0 {
		// MakeRekeyMVCCKeyValFunc modifies the keys, but this is safe because
		// the one we get back from rocksDBIterator.Key is a copy (not a
		// reference to the mmaped file.)
		ingestFunc = MakeRekeyMVCCKeyValFunc(newTableID, ingestFunc)
	}
	startKeyMVCC, endKeyMVCC := engine.MVCCKey{Key: startKey}, engine.MVCCKey{Key: endKey}
	if err := sst.Iterate(startKeyMVCC, endKeyMVCC, ingestFunc); err != nil {
		return err
	}
	return txn.Run(b)
}