github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/sql/testutils.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/settings/cluster"
    18  	"github.com/cockroachdb/cockroach/pkg/sql/parser"
    19  	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
    20  	"github.com/cockroachdb/cockroach/pkg/sql/sessiondata"
    21  	"github.com/cockroachdb/cockroach/pkg/sql/sqlbase"
    22  	"github.com/cockroachdb/cockroach/pkg/util/hlc"
    23  	"github.com/cockroachdb/errors"
    24  )
    25  
    26  // CreateTestTableDescriptor converts a SQL string to a table for test purposes.
    27  // Will fail on complex tables where that operation requires e.g. looking up
    28  // other tables.
    29  func CreateTestTableDescriptor(
    30  	ctx context.Context,
    31  	parentID, id sqlbase.ID,
    32  	schema string,
    33  	privileges *sqlbase.PrivilegeDescriptor,
    34  ) (sqlbase.TableDescriptor, error) {
    35  	st := cluster.MakeTestingClusterSettings()
    36  	stmt, err := parser.ParseOne(schema)
    37  	if err != nil {
    38  		return sqlbase.TableDescriptor{}, err
    39  	}
    40  	semaCtx := tree.MakeSemaContext()
    41  	evalCtx := tree.MakeTestingEvalContext(st)
    42  	switch n := stmt.AST.(type) {
    43  	case *tree.CreateTable:
    44  		desc, err := MakeTableDesc(
    45  			ctx,
    46  			nil, /* txn */
    47  			nil, /* vs */
    48  			st,
    49  			n,
    50  			parentID, keys.PublicSchemaID, id,
    51  			hlc.Timestamp{}, /* creationTime */
    52  			privileges,
    53  			nil, /* affected */
    54  			&semaCtx,
    55  			&evalCtx,
    56  			&sessiondata.SessionData{}, /* sessionData */
    57  			false,                      /* temporary */
    58  		)
    59  		return desc.TableDescriptor, err
    60  	case *tree.CreateSequence:
    61  		desc, err := MakeSequenceTableDesc(
    62  			n.Name.Table(),
    63  			n.Options,
    64  			parentID, keys.PublicSchemaID, id,
    65  			hlc.Timestamp{}, /* creationTime */
    66  			privileges,
    67  			false, /* temporary */
    68  			nil,   /* params */
    69  		)
    70  		return desc.TableDescriptor, err
    71  	default:
    72  		return sqlbase.TableDescriptor{}, errors.Errorf("unexpected AST %T", stmt.AST)
    73  	}
    74  }
    75  
    76  // StmtBufReader is an exported interface for reading a StmtBuf.
    77  // Normally only the write interface of the buffer is exported, as it is used by
    78  // the pgwire.
    79  type StmtBufReader struct {
    80  	buf *StmtBuf
    81  }
    82  
    83  // MakeStmtBufReader creates a StmtBufReader.
    84  func MakeStmtBufReader(buf *StmtBuf) StmtBufReader {
    85  	return StmtBufReader{buf: buf}
    86  }
    87  
    88  // CurCmd returns the current command in the buffer.
    89  func (r StmtBufReader) CurCmd() (Command, error) {
    90  	cmd, _ /* pos */, err := r.buf.CurCmd()
    91  	return cmd, err
    92  }
    93  
    94  // AdvanceOne moves the cursor one position over.
    95  func (r *StmtBufReader) AdvanceOne() {
    96  	r.buf.AdvanceOne()
    97  }
    98  
    99  // SeekToNextBatch skips to the beginning of the next batch of commands.
   100  func (r *StmtBufReader) SeekToNextBatch() error {
   101  	return r.buf.seekToNextBatch()
   102  }
   103  
   104  // Exec is a test utility function that takes a localPlanner (of type
   105  // interface{} so that external packages can call NewInternalPlanner and pass
   106  // the result) and executes a sql statement through the DistSQLPlanner.
   107  func (dsp *DistSQLPlanner) Exec(
   108  	ctx context.Context, localPlanner interface{}, sql string, distribute bool,
   109  ) error {
   110  	stmt, err := parser.ParseOne(sql)
   111  	if err != nil {
   112  		return err
   113  	}
   114  	p := localPlanner.(*planner)
   115  	p.stmt = &Statement{Statement: stmt}
   116  	if err := p.makeOptimizerPlan(ctx); err != nil {
   117  		return err
   118  	}
   119  	rw := newCallbackResultWriter(func(ctx context.Context, row tree.Datums) error {
   120  		return nil
   121  	})
   122  	execCfg := p.ExecCfg()
   123  	recv := MakeDistSQLReceiver(
   124  		ctx,
   125  		rw,
   126  		stmt.AST.StatementType(),
   127  		execCfg.RangeDescriptorCache,
   128  		execCfg.LeaseHolderCache,
   129  		p.txn,
   130  		func(ts hlc.Timestamp) {
   131  			execCfg.Clock.Update(ts)
   132  		},
   133  		p.ExtendedEvalContext().Tracing,
   134  	)
   135  	defer recv.Release()
   136  
   137  	evalCtx := p.ExtendedEvalContext()
   138  	planCtx := execCfg.DistSQLPlanner.NewPlanningCtx(ctx, evalCtx, p.txn, distribute)
   139  	planCtx.planner = p
   140  	planCtx.stmtType = recv.stmtType
   141  
   142  	dsp.PlanAndRun(ctx, evalCtx, planCtx, p.txn, p.curPlan.main, recv)()
   143  	return rw.Err()
   144  }