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

     1  // Copyright 2019 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/sql/sem/tree"
    18  	"github.com/cockroachdb/cockroach/pkg/sql/sqlbase"
    19  	"github.com/cockroachdb/errors"
    20  )
    21  
    22  type unsplitNode struct {
    23  	optColumnsSlot
    24  
    25  	tableDesc *sqlbase.TableDescriptor
    26  	index     *sqlbase.IndexDescriptor
    27  	run       unsplitRun
    28  	rows      planNode
    29  }
    30  
    31  // unsplitRun contains the run-time state of unsplitNode during local execution.
    32  type unsplitRun struct {
    33  	lastUnsplitKey []byte
    34  }
    35  
    36  func (n *unsplitNode) startExec(params runParams) error {
    37  	return nil
    38  }
    39  
    40  func (n *unsplitNode) Next(params runParams) (bool, error) {
    41  	if ok, err := n.rows.Next(params); err != nil || !ok {
    42  		return ok, err
    43  	}
    44  
    45  	row := n.rows.Values()
    46  	rowKey, err := getRowKey(params.ExecCfg().Codec, n.tableDesc, n.index, row)
    47  	if err != nil {
    48  		return false, err
    49  	}
    50  
    51  	if err := params.extendedEvalCtx.ExecCfg.DB.AdminUnsplit(params.ctx, rowKey); err != nil {
    52  		ctx := tree.NewFmtCtx(tree.FmtSimple)
    53  		row.Format(ctx)
    54  		return false, errors.Wrapf(err, "could not UNSPLIT AT %s", ctx)
    55  	}
    56  
    57  	n.run.lastUnsplitKey = rowKey
    58  
    59  	return true, nil
    60  }
    61  
    62  func (n *unsplitNode) Values() tree.Datums {
    63  	return tree.Datums{
    64  		tree.NewDBytes(tree.DBytes(n.run.lastUnsplitKey)),
    65  		tree.NewDString(keys.PrettyPrint(nil /* valDirs */, n.run.lastUnsplitKey)),
    66  	}
    67  }
    68  
    69  func (n *unsplitNode) Close(ctx context.Context) {
    70  	n.rows.Close(ctx)
    71  }
    72  
    73  type unsplitAllNode struct {
    74  	optColumnsSlot
    75  
    76  	tableDesc *sqlbase.TableDescriptor
    77  	index     *sqlbase.IndexDescriptor
    78  	run       unsplitAllRun
    79  }
    80  
    81  // unsplitAllRun contains the run-time state of unsplitAllNode during local execution.
    82  type unsplitAllRun struct {
    83  	keys           [][]byte
    84  	lastUnsplitKey []byte
    85  }
    86  
    87  func (n *unsplitAllNode) startExec(params runParams) error {
    88  	// Use the internal executor to retrieve the split keys.
    89  	statement := `
    90  		SELECT
    91  			start_key
    92  		FROM
    93  			crdb_internal.ranges_no_leases
    94  		WHERE
    95  			database_name=$1 AND table_name=$2 AND index_name=$3 AND split_enforced_until IS NOT NULL
    96  	`
    97  	dbDesc, err := sqlbase.GetDatabaseDescFromID(
    98  		params.ctx, params.p.txn, params.ExecCfg().Codec, n.tableDesc.ParentID,
    99  	)
   100  	if err != nil {
   101  		return err
   102  	}
   103  	indexName := ""
   104  	if n.index.ID != n.tableDesc.PrimaryIndex.ID {
   105  		indexName = n.index.Name
   106  	}
   107  	ranges, err := params.p.ExtendedEvalContext().InternalExecutor.(*InternalExecutor).QueryEx(
   108  		params.ctx, "split points query", params.p.txn, sqlbase.InternalExecutorSessionDataOverride{},
   109  		statement,
   110  		dbDesc.Name,
   111  		n.tableDesc.Name,
   112  		indexName,
   113  	)
   114  	if err != nil {
   115  		return err
   116  	}
   117  	n.run.keys = make([][]byte, len(ranges))
   118  	for i, d := range ranges {
   119  		n.run.keys[i] = []byte(*(d[0].(*tree.DBytes)))
   120  	}
   121  
   122  	return nil
   123  }
   124  
   125  func (n *unsplitAllNode) Next(params runParams) (bool, error) {
   126  	if len(n.run.keys) == 0 {
   127  		return false, nil
   128  	}
   129  	rowKey := n.run.keys[0]
   130  	n.run.keys = n.run.keys[1:]
   131  
   132  	if err := params.extendedEvalCtx.ExecCfg.DB.AdminUnsplit(params.ctx, rowKey); err != nil {
   133  		return false, errors.Wrapf(err, "could not UNSPLIT AT %s", keys.PrettyPrint(nil /* valDirs */, rowKey))
   134  	}
   135  
   136  	n.run.lastUnsplitKey = rowKey
   137  
   138  	return true, nil
   139  }
   140  
   141  func (n *unsplitAllNode) Values() tree.Datums {
   142  	return tree.Datums{
   143  		tree.NewDBytes(tree.DBytes(n.run.lastUnsplitKey)),
   144  		tree.NewDString(keys.PrettyPrint(nil /* valDirs */, n.run.lastUnsplitKey)),
   145  	}
   146  }
   147  
   148  func (n *unsplitAllNode) Close(ctx context.Context) {}