github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/sql/scatter.go (about)

     1  // Copyright 2017 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package sql
    12  
    13  import (
    14  	"context"
    15  
    16  	"github.com/cockroachdb/cockroach/pkg/keys"
    17  	"github.com/cockroachdb/cockroach/pkg/kv"
    18  	"github.com/cockroachdb/cockroach/pkg/roachpb"
    19  	"github.com/cockroachdb/cockroach/pkg/sql/privilege"
    20  	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
    21  	"github.com/cockroachdb/cockroach/pkg/sql/types"
    22  	"github.com/cockroachdb/errors"
    23  )
    24  
    25  type scatterNode struct {
    26  	optColumnsSlot
    27  
    28  	run scatterRun
    29  }
    30  
    31  // Scatter moves ranges to random stores
    32  // (`ALTER TABLE/INDEX ... SCATTER ...` statement)
    33  // Privileges: INSERT on table.
    34  func (p *planner) Scatter(ctx context.Context, n *tree.Scatter) (planNode, error) {
    35  	tableDesc, index, err := p.getTableAndIndex(ctx, &n.TableOrIndex, privilege.INSERT)
    36  	if err != nil {
    37  		return nil, err
    38  	}
    39  
    40  	var span roachpb.Span
    41  	if n.From == nil {
    42  		// No FROM/TO specified; the span is the entire table/index.
    43  		span = tableDesc.IndexSpan(p.ExecCfg().Codec, index.ID)
    44  	} else {
    45  		switch {
    46  		case len(n.From) == 0:
    47  			return nil, errors.Errorf("no columns in SCATTER FROM expression")
    48  		case len(n.From) > len(index.ColumnIDs):
    49  			return nil, errors.Errorf("too many columns in SCATTER FROM expression")
    50  		case len(n.To) == 0:
    51  			return nil, errors.Errorf("no columns in SCATTER TO expression")
    52  		case len(n.To) > len(index.ColumnIDs):
    53  			return nil, errors.Errorf("too many columns in SCATTER TO expression")
    54  		}
    55  
    56  		// Calculate the desired types for the select statement:
    57  		//  - column values; it is OK if the select statement returns fewer columns
    58  		//  (the relevant prefix is used).
    59  		desiredTypes := make([]*types.T, len(index.ColumnIDs))
    60  		for i, colID := range index.ColumnIDs {
    61  			c, err := tableDesc.FindColumnByID(colID)
    62  			if err != nil {
    63  				return nil, err
    64  			}
    65  			desiredTypes[i] = c.Type
    66  		}
    67  		fromVals := make([]tree.Datum, len(n.From))
    68  		for i, expr := range n.From {
    69  			typedExpr, err := p.analyzeExpr(
    70  				ctx, expr, nil, tree.IndexedVarHelper{}, desiredTypes[i], true, "SCATTER",
    71  			)
    72  			if err != nil {
    73  				return nil, err
    74  			}
    75  			fromVals[i], err = typedExpr.Eval(p.EvalContext())
    76  			if err != nil {
    77  				return nil, err
    78  			}
    79  		}
    80  		toVals := make([]tree.Datum, len(n.From))
    81  		for i, expr := range n.To {
    82  			typedExpr, err := p.analyzeExpr(
    83  				ctx, expr, nil, tree.IndexedVarHelper{}, desiredTypes[i], true, "SCATTER",
    84  			)
    85  			if err != nil {
    86  				return nil, err
    87  			}
    88  			toVals[i], err = typedExpr.Eval(p.EvalContext())
    89  			if err != nil {
    90  				return nil, err
    91  			}
    92  		}
    93  
    94  		span.Key, err = getRowKey(p.ExecCfg().Codec, tableDesc.TableDesc(), index, fromVals)
    95  		if err != nil {
    96  			return nil, err
    97  		}
    98  		span.EndKey, err = getRowKey(p.ExecCfg().Codec, tableDesc.TableDesc(), index, toVals)
    99  		if err != nil {
   100  			return nil, err
   101  		}
   102  		// Tolerate reversing FROM and TO; this can be useful for descending
   103  		// indexes.
   104  		if cmp := span.Key.Compare(span.EndKey); cmp > 0 {
   105  			span.Key, span.EndKey = span.EndKey, span.Key
   106  		} else if cmp == 0 {
   107  			// Key==EndKey is invalid, so special-case when the user's FROM and
   108  			// TO are the same tuple.
   109  			span.EndKey = span.EndKey.Next()
   110  		}
   111  	}
   112  
   113  	return &scatterNode{
   114  		run: scatterRun{
   115  			span: span,
   116  		},
   117  	}, nil
   118  }
   119  
   120  // scatterRun contains the run-time state of scatterNode during local execution.
   121  type scatterRun struct {
   122  	span roachpb.Span
   123  
   124  	rangeIdx int
   125  	ranges   []roachpb.AdminScatterResponse_Range
   126  }
   127  
   128  func (n *scatterNode) startExec(params runParams) error {
   129  	db := params.p.ExecCfg().DB
   130  	req := &roachpb.AdminScatterRequest{
   131  		RequestHeader:   roachpb.RequestHeader{Key: n.run.span.Key, EndKey: n.run.span.EndKey},
   132  		RandomizeLeases: true,
   133  	}
   134  	res, pErr := kv.SendWrapped(params.ctx, db.NonTransactionalSender(), req)
   135  	if pErr != nil {
   136  		return pErr.GoError()
   137  	}
   138  	n.run.rangeIdx = -1
   139  	n.run.ranges = res.(*roachpb.AdminScatterResponse).Ranges
   140  	return nil
   141  }
   142  
   143  func (n *scatterNode) Next(params runParams) (bool, error) {
   144  	n.run.rangeIdx++
   145  	hasNext := n.run.rangeIdx < len(n.run.ranges)
   146  	return hasNext, nil
   147  }
   148  
   149  func (n *scatterNode) Values() tree.Datums {
   150  	r := n.run.ranges[n.run.rangeIdx]
   151  	return tree.Datums{
   152  		tree.NewDBytes(tree.DBytes(r.Span.Key)),
   153  		tree.NewDString(keys.PrettyPrint(nil /* valDirs */, r.Span.Key)),
   154  	}
   155  }
   156  
   157  func (*scatterNode) Close(ctx context.Context) {}