// upsertRowPKs returns the primary keys of any rows with potential upsert
// conflicts.
func (tu *tableUpserter) upsertRowPKs(ctx context.Context) ([]roachpb.Key, error) {
upsertRowPKs := make([]roachpb.Key, len(tu.insertRows))
if tu.conflictIndex.ID == tu.tableDesc.PrimaryIndex.ID {
// If the conflict index is the primary index, we can compute them directly.
// In this case, the slice will be filled, but not all rows will have
// conflicts.
for i, insertRow := range tu.insertRows {
upsertRowPK, _, err := sqlbase.EncodeIndexKey(
tu.tableDesc, &tu.conflictIndex, tu.ri.InsertColIDtoRowIndex, insertRow, tu.indexKeyPrefix)
if err != nil {
return nil, err
}
upsertRowPKs[i] = upsertRowPK
}
return upsertRowPKs, nil
}
// Otherwise, compute the keys for the conflict index and look them up. The
// primary keys can be constructed from the entries that come back. In this
// case, some spots in the slice will be nil (indicating no conflict) and the
// others will be conflicting rows.
b := tu.txn.NewBatch()
for _, insertRow := range tu.insertRows {
entry, err := sqlbase.EncodeSecondaryIndex(
tu.tableDesc, &tu.conflictIndex, tu.ri.InsertColIDtoRowIndex, insertRow)
if err != nil {
return nil, err
}
if log.V(2) {
log.Infof(ctx, "Get %s\n", entry.Key)
}
b.Get(entry.Key)
}
if err := tu.txn.Run(b); err != nil {
return nil, err
}
for i, result := range b.Results {
// if len(result.Rows) == 0, then no conflict for this row, so leave
// upsertRowPKs[i] as nil.
if len(result.Rows) == 1 {
if result.Rows[0].Value == nil {
upsertRowPKs[i] = nil
} else {
upsertRowPK, err := sqlbase.ExtractIndexKey(&tu.a, tu.tableDesc, result.Rows[0])
if err != nil {
return nil, err
}
upsertRowPKs[i] = upsertRowPK
}
} else if len(result.Rows) > 1 {
panic(fmt.Errorf(
"Expected <= 1 but got %d conflicts for row %s", len(result.Rows), tu.insertRows[i]))
}
}
return upsertRowPKs, nil
}
// fetchExisting returns any existing rows in the table that conflict with the
// ones in tu.insertRows. The returned slice is the same length as tu.insertRows
// and a nil entry indicates no conflict.
func (tu *tableUpserter) fetchExisting(ctx context.Context) ([]parser.DTuple, error) {
primaryKeys, err := tu.upsertRowPKs(ctx)
if err != nil {
return nil, err
}
pkSpans := make(roachpb.Spans, 0, len(primaryKeys))
rowIdxForPrimaryKey := make(map[string]int, len(primaryKeys))
for i, primaryKey := range primaryKeys {
if primaryKey != nil {
pkSpans = append(pkSpans, roachpb.Span{Key: primaryKey, EndKey: primaryKey.PrefixEnd()})
if _, ok := rowIdxForPrimaryKey[string(primaryKey)]; ok {
return nil, fmt.Errorf("UPSERT/ON CONFLICT DO UPDATE command cannot affect row a second time")
}
rowIdxForPrimaryKey[string(primaryKey)] = i
}
}
if len(pkSpans) == 0 {
// Every key was empty, so there's nothing to fetch.
return make([]parser.DTuple, len(primaryKeys)), nil
}
// We don't limit batches here because the spans are unordered.
if err := tu.fetcher.StartScan(tu.txn, pkSpans, false /* no batch limits */, 0); err != nil {
return nil, err
}
rows := make([]parser.DTuple, len(primaryKeys))
for {
row, err := tu.fetcher.NextRowDecoded()
if err != nil {
return nil, err
}
if row == nil {
break // Done
}
rowPrimaryKey, _, err := sqlbase.EncodeIndexKey(
tu.tableDesc, &tu.tableDesc.PrimaryIndex, tu.fetchColIDtoRowIndex, row, tu.indexKeyPrefix)
if err != nil {
return nil, err
}
// The rows returned by rowFetcher are invalidated after the call to
// NextRow, so we have to copy them to save them.
rowCopy := make(parser.DTuple, len(row))
copy(rowCopy, row)
rows[rowIdxForPrimaryKey[string(rowPrimaryKey)]] = rowCopy
}
return rows, nil
}
// TODO(dt): Batch checks of many rows.
func (f baseFKHelper) check(values parser.DTuple) (parser.DTuple, error) {
var key roachpb.Key
if values != nil {
keyBytes, _, err := sqlbase.EncodeIndexKey(
f.searchTable, f.searchIdx, f.ids, values, f.searchPrefix)
if err != nil {
return nil, err
}
key = roachpb.Key(keyBytes)
} else {
key = roachpb.Key(f.searchPrefix)
}
spans := roachpb.Spans{roachpb.Span{Key: key, EndKey: key.PrefixEnd()}}
if err := f.rf.StartScan(f.txn, spans, true /* limit batches */, 1); err != nil {
return nil, err
}
return f.rf.NextRowDecoded()
}
// encodeIndexes encodes the primary and secondary index keys. The
// secondaryIndexEntries are only valid until the next call to encodeIndexes or
// encodeSecondaryIndexes.
func (rh *rowHelper) encodeIndexes(
colIDtoRowIndex map[sqlbase.ColumnID]int, values []parser.Datum,
) (primaryIndexKey []byte, secondaryIndexEntries []sqlbase.IndexEntry, err error) {
if rh.primaryIndexKeyPrefix == nil {
rh.primaryIndexKeyPrefix = sqlbase.MakeIndexKeyPrefix(rh.tableDesc,
rh.tableDesc.PrimaryIndex.ID)
}
primaryIndexKey, _, err = sqlbase.EncodeIndexKey(
rh.tableDesc, &rh.tableDesc.PrimaryIndex, colIDtoRowIndex, values, rh.primaryIndexKeyPrefix)
if err != nil {
return nil, nil, err
}
secondaryIndexEntries, err = rh.encodeSecondaryIndexes(colIDtoRowIndex, values)
if err != nil {
return nil, nil, err
}
return primaryIndexKey, secondaryIndexEntries, nil
}
func (tu *tableUpserter) row(ctx context.Context, row parser.DTuple) (parser.DTuple, error) {
if tu.fastPathBatch != nil {
primaryKey, _, err := sqlbase.EncodeIndexKey(
tu.tableDesc, &tu.tableDesc.PrimaryIndex, tu.ri.InsertColIDtoRowIndex, row, tu.indexKeyPrefix)
if err != nil {
return nil, err
}
if _, ok := tu.fastPathKeys[string(primaryKey)]; ok {
return nil, fmt.Errorf("UPSERT/ON CONFLICT DO UPDATE command cannot affect row a second time")
}
tu.fastPathKeys[string(primaryKey)] = struct{}{}
err = tu.ri.InsertRow(ctx, tu.fastPathBatch, row, true)
return nil, err
}
tu.insertRows = append(tu.insertRows, row)
// TODO(dan): If len(tu.insertRows) > some threshold, call flush().
return nil, nil
}
func (n *indexJoinNode) Next() (bool, error) {
// Loop looking up the next row. We either are going to pull a row from the
// table or a batch of rows from the index. If we pull a batch of rows from
// the index we perform another iteration of the loop looking for rows in the
// table. This outer loop is necessary because a batch of rows from the index
// might all be filtered when the resulting rows are read from the table.
for tableLookup := (len(n.table.spans) > 0); true; tableLookup = true {
// First, try to pull a row from the table.
if tableLookup {
next, err := n.table.Next()
if err != nil {
return false, err
}
if next {
if n.explain == explainDebug {
n.debugVals = n.table.DebugValues()
}
return true, nil
}
}
// The table is out of rows. Pull primary keys from the index.
n.table.scanInitialized = false
n.table.spans = n.table.spans[:0]
for len(n.table.spans) < indexJoinBatchSize {
if next, err := n.index.Next(); !next {
// The index is out of rows or an error occurred.
if err != nil {
return false, err
}
if len(n.table.spans) == 0 {
// The index is out of rows.
return false, nil
}
break
}
if n.explain == explainDebug {
n.debugVals = n.index.DebugValues()
if n.debugVals.output != debugValueRow {
return true, nil
}
}
vals := n.index.Values()
primaryIndexKey, _, err := sqlbase.EncodeIndexKey(
&n.table.desc, n.table.index, n.colIDtoRowIndex, vals, n.primaryKeyPrefix)
if err != nil {
return false, err
}
key := roachpb.Key(primaryIndexKey)
n.table.spans = append(n.table.spans, roachpb.Span{
Key: key,
EndKey: key.PrefixEnd(),
})
if n.explain == explainDebug {
// In debug mode, return the index information as a "partial" row.
n.debugVals.output = debugValuePartial
return true, nil
}
}
if log.V(3) {
log.Infof(n.index.p.ctx(), "table scan: %s", sqlbase.PrettySpans(n.table.spans, 0))
}
}
return false, nil
}
// truncateAndBackfillColumnsChunk returns the next-key, done and an error.
// next-key and done are invalid if error != nil. next-key is invalid if done
// is true.
func (sc *SchemaChanger) truncateAndBackfillColumnsChunk(
added []sqlbase.ColumnDescriptor,
dropped []sqlbase.ColumnDescriptor,
defaultExprs []parser.TypedExpr,
sp roachpb.Span,
updateValues parser.DTuple,
nonNullViolationColumnName string,
chunkSize int64,
mutationIdx int,
lastCheckpoint *time.Time,
) (roachpb.Key, bool, error) {
done := false
var nextKey roachpb.Key
err := sc.db.Txn(context.TODO(), func(txn *client.Txn) error {
if sc.testingKnobs.RunBeforeBackfillChunk != nil {
if err := sc.testingKnobs.RunBeforeBackfillChunk(sp); err != nil {
return err
}
}
if sc.testingKnobs.RunAfterBackfillChunk != nil {
defer sc.testingKnobs.RunAfterBackfillChunk()
}
tableDesc, err := sqlbase.GetTableDescFromID(txn, sc.tableID)
if err != nil {
return err
}
// Short circuit the backfill if the table has been deleted.
if done = tableDesc.Dropped(); done {
return nil
}
updateCols := append(added, dropped...)
fkTables := tablesNeededForFKs(*tableDesc, CheckUpdates)
for k := range fkTables {
table, err := sqlbase.GetTableDescFromID(txn, k)
if err != nil {
return err
}
fkTables[k] = tableLookup{table: table}
}
// TODO(dan): Tighten up the bound on the requestedCols parameter to
// makeRowUpdater.
requestedCols := make([]sqlbase.ColumnDescriptor, 0, len(tableDesc.Columns)+len(added))
requestedCols = append(requestedCols, tableDesc.Columns...)
requestedCols = append(requestedCols, added...)
ru, err := makeRowUpdater(
txn, tableDesc, fkTables, updateCols, requestedCols, rowUpdaterOnlyColumns,
)
if err != nil {
return err
}
// TODO(dan): This check is an unfortunate bleeding of the internals of
// rowUpdater. Extract the sql row to k/v mapping logic out into something
// usable here.
if !ru.isColumnOnlyUpdate() {
panic("only column data should be modified, but the rowUpdater is configured otherwise")
}
// Run a scan across the table using the primary key. Running
// the scan and applying the changes in many transactions is
// fine because the schema change is in the correct state to
// handle intermediate OLTP commands which delete and add
// values during the scan.
var rf sqlbase.RowFetcher
colIDtoRowIndex := colIDtoRowIndexFromCols(tableDesc.Columns)
valNeededForCol := make([]bool, len(tableDesc.Columns))
for i := range valNeededForCol {
_, valNeededForCol[i] = ru.fetchColIDtoRowIndex[tableDesc.Columns[i].ID]
}
if err := rf.Init(
tableDesc, colIDtoRowIndex, &tableDesc.PrimaryIndex, false, false,
tableDesc.Columns, valNeededForCol,
); err != nil {
return err
}
if err := rf.StartScan(
txn, roachpb.Spans{sp}, true /* limit batches */, chunkSize,
); err != nil {
return err
}
oldValues := make(parser.DTuple, len(ru.fetchCols))
writeBatch := txn.NewBatch()
rowLength := 0
var lastRowSeen parser.DTuple
i := int64(0)
for ; i < chunkSize; i++ {
row, err := rf.NextRow()
if err != nil {
return err
}
if row == nil {
break
}
lastRowSeen = row
if nonNullViolationColumnName != "" {
return sqlbase.NewNonNullViolationError(nonNullViolationColumnName)
}
copy(oldValues, row)
// Update oldValues with NULL values where values weren't found;
// only update when necessary.
if rowLength != len(row) {
rowLength = len(row)
for j := rowLength; j < len(oldValues); j++ {
oldValues[j] = parser.DNull
}
}
if _, err := ru.updateRow(txn.Context, writeBatch, oldValues, updateValues); err != nil {
return err
}
}
if err := txn.Run(writeBatch); err != nil {
return convertBackfillError(tableDesc, writeBatch)
}
if done = i < chunkSize; done {
return nil
}
curIndexKey, _, err := sqlbase.EncodeIndexKey(
tableDesc, &tableDesc.PrimaryIndex, colIDtoRowIndex, lastRowSeen,
sqlbase.MakeIndexKeyPrefix(tableDesc, tableDesc.PrimaryIndex.ID))
if err != nil {
return err
}
resume := roachpb.Span{Key: roachpb.Key(curIndexKey).PrefixEnd(), EndKey: sp.EndKey}
if err := sc.maybeWriteResumeSpan(txn, tableDesc, resume, mutationIdx, lastCheckpoint); err != nil {
return err
}
nextKey = resume.Key
return nil
})
return nextKey, done, err
}