Golang sqlbase.WrapDescriptor函數代碼示例

func (p *planner) createDescriptorWithID(
	idKey roachpb.Key, id sqlbase.ID, descriptor sqlbase.DescriptorProto,
) (bool, error) {
	descriptor.SetID(id)
	// TODO(pmattis): The error currently returned below is likely going to be
	// difficult to interpret.
	//
	// TODO(pmattis): Need to handle if-not-exists here as well.
	//
	// TODO(pmattis): This is writing the namespace and descriptor table entries,
	// but not going through the normal INSERT logic and not performing a precise
	// mimicry. In particular, we're only writing a single key per table, while
	// perfect mimicry would involve writing a sentinel key for each row as well.
	descKey := sqlbase.MakeDescMetadataKey(descriptor.GetID())

	b := &client.Batch{}
	descID := descriptor.GetID()
	descDesc := sqlbase.WrapDescriptor(descriptor)
	if log.V(2) {
		log.Infof(p.ctx(), "CPut %s -> %d", idKey, descID)
		log.Infof(p.ctx(), "CPut %s -> %s", descKey, descDesc)
	}
	b.CPut(idKey, descID, nil)
	b.CPut(descKey, descDesc, nil)

	p.setTestingVerifyMetadata(func(systemConfig config.SystemConfig) error {
		if err := expectDescriptorID(systemConfig, idKey, descID); err != nil {
			return err
		}
		return expectDescriptor(systemConfig, descKey, descDesc)
	})

	return true, p.txn.Run(b)
}

// AcquireLease acquires a schema change lease on the table if
// an unexpired lease doesn't exist. It returns the lease.
func (sc *SchemaChanger) AcquireLease() (sqlbase.TableDescriptor_SchemaChangeLease, error) {
	var lease sqlbase.TableDescriptor_SchemaChangeLease
	err := sc.db.Txn(func(txn *client.Txn) error {
		txn.SetSystemConfigTrigger()
		tableDesc, err := getTableDescFromID(txn, sc.tableID)
		if err != nil {
			return err
		}

		// A second to deal with the time uncertainty across nodes.
		// It is perfectly valid for two or more goroutines to hold a valid
		// lease and execute a schema change in parallel, because schema
		// changes are executed using transactions that run sequentially.
		// This just reduces the probability of a write collision.
		expirationTimeUncertainty := time.Second

		if tableDesc.Lease != nil {
			if time.Unix(0, tableDesc.Lease.ExpirationTime).Add(expirationTimeUncertainty).After(timeutil.Now()) {
				return errExistingSchemaChangeLease
			}
			log.Infof("Overriding existing expired lease %v", tableDesc.Lease)
		}
		lease = sc.createSchemaChangeLease()
		tableDesc.Lease = &lease
		return txn.Put(sqlbase.MakeDescMetadataKey(tableDesc.ID), sqlbase.WrapDescriptor(tableDesc))
	})
	return lease, err
}

// Convert all the mutations into live descriptors for the table
// and write the updated table descriptor to the DB.
func (mt mutationTest) makeMutationsActive() {
	// Remove mutation to check real values in DB using SQL
	if mt.tableDesc.Mutations == nil || len(mt.tableDesc.Mutations) == 0 {
		mt.Fatal("No mutations to make active")
	}
	for _, m := range mt.tableDesc.Mutations {
		if col := m.GetColumn(); col != nil {
			mt.tableDesc.Columns = append(mt.tableDesc.Columns, *col)
		} else if index := m.GetIndex(); index != nil {
			mt.tableDesc.Indexes = append(mt.tableDesc.Indexes, *index)
		} else {
			mt.Fatalf("no descriptor in mutation: %v", m)
		}
	}
	mt.tableDesc.Mutations = nil
	if err := mt.tableDesc.Validate(); err != nil {
		mt.Fatal(err)
	}
	if err := mt.kvDB.Put(
		sqlbase.MakeDescMetadataKey(mt.tableDesc.ID),
		sqlbase.WrapDescriptor(mt.tableDesc),
	); err != nil {
		mt.Fatal(err)
	}
}

// writeTableDesc implements the SchemaAccessor interface.
func (p *planner) writeTableDesc(tableDesc *sqlbase.TableDescriptor) error {
	if isVirtualDescriptor(tableDesc) {
		panic(fmt.Sprintf("Virtual Descriptors cannot be stored, found: %v", tableDesc))
	}
	return p.txn.Put(sqlbase.MakeDescMetadataKey(tableDesc.GetID()),
		sqlbase.WrapDescriptor(tableDesc))
}

// writeMutation writes the mutation to the table descriptor. If the
// State or the Direction is undefined, these values are populated via
// picking random values before the mutation is written.
func (mt mutationTest) writeMutation(m sqlbase.DescriptorMutation) {
	if m.Direction == sqlbase.DescriptorMutation_NONE {
		// randomly pick ADD/DROP mutation if this is the first mutation, or
		// pick the direction already chosen for the first mutation.
		if len(mt.tableDesc.Mutations) > 0 {
			m.Direction = mt.tableDesc.Mutations[0].Direction
		} else {
			m.Direction = sqlbase.DescriptorMutation_DROP
			if rand.Intn(2) == 0 {
				m.Direction = sqlbase.DescriptorMutation_ADD
			}
		}
	}
	if m.State == sqlbase.DescriptorMutation_UNKNOWN {
		// randomly pick DELETE_ONLY/WRITE_ONLY state.
		r := rand.Intn(2)
		if r == 0 {
			m.State = sqlbase.DescriptorMutation_DELETE_ONLY
		} else {
			m.State = sqlbase.DescriptorMutation_WRITE_ONLY
		}
	}
	mt.tableDesc.Mutations = append(mt.tableDesc.Mutations, m)
	if err := mt.tableDesc.Validate(); err != nil {
		mt.Fatal(err)
	}
	if err := mt.kvDB.Put(
		sqlbase.MakeDescMetadataKey(mt.tableDesc.ID),
		sqlbase.WrapDescriptor(mt.tableDesc),
	); err != nil {
		mt.Fatal(err)
	}
}

// RenameDatabase renames the database.
// Privileges: security.RootUser user.
//   Notes: postgres requires superuser, db owner, or "CREATEDB".
//          mysql >= 5.1.23 does not allow database renames.
func (p *planner) RenameDatabase(n *parser.RenameDatabase) (planNode, error) {
	if n.Name == "" || n.NewName == "" {
		return nil, errEmptyDatabaseName
	}

	if p.session.User != security.RootUser {
		return nil, fmt.Errorf("only %s is allowed to rename databases", security.RootUser)
	}

	dbDesc, err := p.getDatabaseDesc(string(n.Name))
	if err != nil {
		return nil, err
	}
	if dbDesc == nil {
		return nil, databaseDoesNotExistError(string(n.Name))
	}

	if n.Name == n.NewName {
		// Noop.
		return &emptyNode{}, nil
	}

	// Now update the nameMetadataKey and the descriptor.
	descKey := sqlbase.MakeDescMetadataKey(dbDesc.GetID())
	dbDesc.SetName(string(n.NewName))

	if err := dbDesc.Validate(); err != nil {
		return nil, err
	}

	newKey := databaseKey{string(n.NewName)}.Key()
	oldKey := databaseKey{string(n.Name)}.Key()
	descID := dbDesc.GetID()
	descDesc := sqlbase.WrapDescriptor(dbDesc)

	b := client.Batch{}
	b.CPut(newKey, descID, nil)
	b.Put(descKey, descDesc)
	b.Del(oldKey)

	if err := p.txn.Run(&b); err != nil {
		if _, ok := err.(*roachpb.ConditionFailedError); ok {
			return nil, fmt.Errorf("the new database name %q already exists", string(n.NewName))
		}
		return nil, err
	}

	p.setTestingVerifyMetadata(func(systemConfig config.SystemConfig) error {
		if err := expectDescriptorID(systemConfig, newKey, descID); err != nil {
			return err
		}
		if err := expectDescriptor(systemConfig, descKey, descDesc); err != nil {
			return err
		}
		return expectDeleted(systemConfig, oldKey)
	})

	return &emptyNode{}, nil
}

// RenameIndex renames the index.
// Privileges: CREATE on table.
//   notes: postgres requires CREATE on the table.
//          mysql requires ALTER, CREATE, INSERT on the table.
func (p *planner) RenameIndex(n *parser.RenameIndex) (planNode, error) {
	tn, err := n.Index.Table.NormalizeWithDatabaseName(p.session.Database)
	if err != nil {
		return nil, err
	}

	tableDesc, err := p.mustGetTableDesc(tn)
	if err != nil {
		return nil, err
	}

	normIdxName := sqlbase.NormalizeName(n.Index.Index)
	status, i, err := tableDesc.FindIndexByNormalizedName(normIdxName)
	if err != nil {
		if n.IfExists {
			// Noop.
			return &emptyNode{}, nil
		}
		// Index does not exist, but we want it to: error out.
		return nil, err
	}

	if err := p.checkPrivilege(tableDesc, privilege.CREATE); err != nil {
		return nil, err
	}

	if n.NewName == "" {
		return nil, errEmptyIndexName
	}
	normNewIdxName := sqlbase.NormalizeName(n.NewName)

	if normIdxName == normNewIdxName {
		// Noop.
		return &emptyNode{}, nil
	}

	if _, _, err := tableDesc.FindIndexByNormalizedName(normNewIdxName); err == nil {
		return nil, fmt.Errorf("index name %q already exists", n.NewName)
	}

	if status == sqlbase.DescriptorActive {
		tableDesc.Indexes[i].Name = normNewIdxName
	} else {
		tableDesc.Mutations[i].GetIndex().Name = normNewIdxName
	}

	if err := tableDesc.SetUpVersion(); err != nil {
		return nil, err
	}
	descKey := sqlbase.MakeDescMetadataKey(tableDesc.GetID())
	if err := tableDesc.Validate(p.txn); err != nil {
		return nil, err
	}
	if err := p.txn.Put(descKey, sqlbase.WrapDescriptor(tableDesc)); err != nil {
		return nil, err
	}
	p.notifySchemaChange(tableDesc.ID, sqlbase.InvalidMutationID)
	return &emptyNode{}, nil
}

// createDescriptor implements the DescriptorAccessor interface.
func (p *planner) createDescriptor(plainKey sqlbase.DescriptorKey, descriptor sqlbase.DescriptorProto, ifNotExists bool) (bool, error) {
	idKey := plainKey.Key()
	// Check whether idKey exists.
	gr, err := p.txn.Get(idKey)
	if err != nil {
		return false, err
	}

	if gr.Exists() {
		if ifNotExists {
			// Noop.
			return false, nil
		}
		// Key exists, but we don't want it to: error out.
		return false, fmt.Errorf("%s %q already exists", descriptor.TypeName(), plainKey.Name())
	}

	// Increment unique descriptor counter.
	if ir, err := p.txn.Inc(keys.DescIDGenerator, 1); err == nil {
		descriptor.SetID(sqlbase.ID(ir.ValueInt() - 1))
	} else {
		return false, err
	}

	// TODO(pmattis): The error currently returned below is likely going to be
	// difficult to interpret.
	//
	// TODO(pmattis): Need to handle if-not-exists here as well.
	//
	// TODO(pmattis): This is writing the namespace and descriptor table entries,
	// but not going through the normal INSERT logic and not performing a precise
	// mimicry. In particular, we're only writing a single key per table, while
	// perfect mimicry would involve writing a sentinel key for each row as well.
	descKey := sqlbase.MakeDescMetadataKey(descriptor.GetID())

	b := client.Batch{}
	descID := descriptor.GetID()
	descDesc := sqlbase.WrapDescriptor(descriptor)
	if log.V(2) {
		log.Infof("CPut %s -> %d", idKey, descID)
		log.Infof("CPut %s -> %s", descKey, descDesc)
	}
	b.CPut(idKey, descID, nil)
	b.CPut(descKey, descDesc, nil)

	p.setTestingVerifyMetadata(func(systemConfig config.SystemConfig) error {
		if err := expectDescriptorID(systemConfig, idKey, descID); err != nil {
			return err
		}
		return expectDescriptor(systemConfig, descKey, descDesc)
	})

	return true, p.txn.Run(&b)
}

// ReleaseLease the table lease if it is the one registered with
// the table descriptor.
func (sc *SchemaChanger) ReleaseLease(lease sqlbase.TableDescriptor_SchemaChangeLease) error {
	err := sc.db.Txn(func(txn *client.Txn) error {
		tableDesc, err := sc.findTableWithLease(txn, lease)
		if err != nil {
			return err
		}
		tableDesc.Lease = nil
		txn.SetSystemConfigTrigger()
		return txn.Put(sqlbase.MakeDescMetadataKey(tableDesc.ID), sqlbase.WrapDescriptor(tableDesc))
	})
	return err
}

// ExtendLease for the current leaser.
func (sc *SchemaChanger) ExtendLease(
	existingLease sqlbase.TableDescriptor_SchemaChangeLease,
) (sqlbase.TableDescriptor_SchemaChangeLease, error) {
	var lease sqlbase.TableDescriptor_SchemaChangeLease
	err := sc.db.Txn(func(txn *client.Txn) error {
		tableDesc, err := sc.findTableWithLease(txn, existingLease)
		if err != nil {
			return err
		}
		lease = sc.createSchemaChangeLease()
		tableDesc.Lease = &lease
		txn.SetSystemConfigTrigger()
		return txn.Put(sqlbase.MakeDescMetadataKey(tableDesc.ID), sqlbase.WrapDescriptor(tableDesc))
	})
	return lease, err
}

func (p *planner) changePrivileges(
	targets parser.TargetList,
	grantees parser.NameList,
	changePrivilege func(*sqlbase.PrivilegeDescriptor, string),
) (planNode, error) {
	descriptors, err := p.getDescriptorsFromTargetList(targets)
	if err != nil {
		return nil, err
	}

	for _, descriptor := range descriptors {
		if err := p.checkPrivilege(descriptor, privilege.GRANT); err != nil {
			return nil, err
		}
		privileges := descriptor.GetPrivileges()
		for _, grantee := range grantees {
			changePrivilege(privileges, string(grantee))
		}

		switch d := descriptor.(type) {
		case *sqlbase.DatabaseDescriptor:
			if err := d.Validate(); err != nil {
				return nil, err
			}
		case *sqlbase.TableDescriptor:
			if err := d.Validate(p.txn); err != nil {
				return nil, err
			}
			if err := d.SetUpVersion(); err != nil {
				return nil, err
			}
			p.notifySchemaChange(d.ID, sqlbase.InvalidMutationID)
		}
	}

	// Now update the descriptors transactionally.
	b := p.txn.NewBatch()
	for _, descriptor := range descriptors {
		descKey := sqlbase.MakeDescMetadataKey(descriptor.GetID())
		b.Put(descKey, sqlbase.WrapDescriptor(descriptor))
	}
	if err := p.txn.Run(b); err != nil {
		return nil, err
	}
	return &emptyNode{}, nil
}

// renameDatabase implements the DatabaseAccessor interface.
func (p *planner) renameDatabase(oldDesc *sqlbase.DatabaseDescriptor, newName string) error {
	onAlreadyExists := func() error {
		return fmt.Errorf("the new database name %q already exists", newName)
	}

	if isVirtualDatabase(newName) {
		return onAlreadyExists()
	}

	oldName := oldDesc.Name
	oldDesc.SetName(newName)
	if err := oldDesc.Validate(); err != nil {
		return err
	}

	oldKey := databaseKey{oldName}.Key()
	newKey := databaseKey{newName}.Key()
	descID := oldDesc.GetID()
	descKey := sqlbase.MakeDescMetadataKey(descID)
	descDesc := sqlbase.WrapDescriptor(oldDesc)

	b := &client.Batch{}
	b.CPut(newKey, descID, nil)
	b.Put(descKey, descDesc)
	b.Del(oldKey)

	if err := p.txn.Run(b); err != nil {
		if _, ok := err.(*roachpb.ConditionFailedError); ok {
			return onAlreadyExists()
		}
		return err
	}

	p.setTestingVerifyMetadata(func(systemConfig config.SystemConfig) error {
		if err := expectDescriptorID(systemConfig, newKey, descID); err != nil {
			return err
		}
		if err := expectDescriptor(systemConfig, descKey, descDesc); err != nil {
			return err
		}
		return expectDeleted(systemConfig, oldKey)
	})
	return nil
}

// TestAddingFKs checks the behavior of a table in the non-public `ADD` state.
// Being non-public, it should not be visible to clients, and is therefore
// assumed to be empty (e.g. by foreign key checks), since no one could have
// written to it yet.
func TestAddingFKs(t *testing.T) {
	defer leaktest.AfterTest(t)()

	params, _ := createTestServerParams()
	s, sqlDB, kvDB := serverutils.StartServer(t, params)
	defer s.Stopper().Stop()

	if _, err := sqlDB.Exec(`
		CREATE DATABASE t;
		CREATE TABLE t.products (id INT PRIMARY KEY);
		INSERT INTO t.products VALUES (1), (2);
		CREATE TABLE t.orders (id INT PRIMARY KEY, product INT REFERENCES t.products, INDEX (product));
	`); err != nil {
		t.Fatal(err)
	}

	// Step the referencing table back to the ADD state.
	ordersDesc := sqlbase.GetTableDescriptor(kvDB, "t", "orders")
	ordersDesc.State = sqlbase.TableDescriptor_ADD
	ordersDesc.Version++
	if err := kvDB.Put(
		sqlbase.MakeDescMetadataKey(ordersDesc.ID),
		sqlbase.WrapDescriptor(ordersDesc),
	); err != nil {
		t.Fatal(err)
	}

	// Generally a referenced table needs to lookup referencing tables to check
	// FKs during delete operations, but referencing tables in the ADD state are
	// given special treatment.
	if _, err := sqlDB.Exec(`DELETE FROM t.products`); err != nil {
		t.Fatal(err)
	}

	// Client should not see the orders table.
	if _, err := sqlDB.Exec(
		`SELECT * FROM t.orders`,
	); !testutils.IsError(err, "table is being added") {
		t.Fatal(err)
	}
}

// ExtendLease for the current leaser. This needs to be called often while
// doing a schema change to prevent more than one node attempting to apply a
// schema change (which is still safe, but unwise).
func (sc *SchemaChanger) ExtendLease(
	existingLease sqlbase.TableDescriptor_SchemaChangeLease,
) (sqlbase.TableDescriptor_SchemaChangeLease, error) {
	// Check if there is still time on this lease.
	minDesiredExpiration := timeutil.Now().Add(MinSchemaChangeLeaseDuration)
	if time.Unix(0, existingLease.ExpirationTime).After(minDesiredExpiration) {
		return existingLease, nil
	}
	// Update lease.
	var lease sqlbase.TableDescriptor_SchemaChangeLease
	err := sc.db.Txn(func(txn *client.Txn) error {
		tableDesc, err := sc.findTableWithLease(txn, existingLease)
		if err != nil {
			return err
		}

		lease = sc.createSchemaChangeLease()
		tableDesc.Lease = &lease
		txn.SetSystemConfigTrigger()
		return txn.Put(sqlbase.MakeDescMetadataKey(tableDesc.ID), sqlbase.WrapDescriptor(tableDesc))
	})
	return lease, err
}

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