github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/sql/opt/optbuilder/opaque.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 optbuilder
    12  
    13  import (
    14  	"context"
    15  	"reflect"
    16  
    17  	"github.com/cockroachdb/cockroach/pkg/sql/opt"
    18  	"github.com/cockroachdb/cockroach/pkg/sql/opt/memo"
    19  	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
    20  	"github.com/cockroachdb/cockroach/pkg/sql/sqlbase"
    21  	"github.com/cockroachdb/errors"
    22  )
    23  
    24  // BuildOpaqueFn is a handler for building the metadata for an opaque statement.
    25  type BuildOpaqueFn func(
    26  	context.Context, *tree.SemaContext, *tree.EvalContext, tree.Statement,
    27  ) (opt.OpaqueMetadata, sqlbase.ResultColumns, error)
    28  
    29  // OpaqueType indicates whether an opaque statement can mutate data or change
    30  // schema.
    31  type OpaqueType int
    32  
    33  const (
    34  	// OpaqueReadOnly is used for statements that do not mutate state as part of
    35  	// the transaction, and can be run in read-only transactions.
    36  	OpaqueReadOnly OpaqueType = iota
    37  
    38  	// OpaqueMutation is used for statements that mutate data and cannot be run as
    39  	// part of read-only transactions.
    40  	OpaqueMutation
    41  
    42  	// OpaqueDDL is used for statements that change a schema and cannot be
    43  	// executed following a mutation in the same transaction.
    44  	OpaqueDDL
    45  )
    46  
    47  // RegisterOpaque registers an opaque handler for a specific statement type.
    48  func RegisterOpaque(stmtType reflect.Type, opaqueType OpaqueType, fn BuildOpaqueFn) {
    49  	if _, ok := opaqueStatements[stmtType]; ok {
    50  		panic(errors.AssertionFailedf("opaque statement %s already registered", stmtType))
    51  	}
    52  	opaqueStatements[stmtType] = opaqueStmtInfo{
    53  		typ:     opaqueType,
    54  		buildFn: fn,
    55  	}
    56  }
    57  
    58  type opaqueStmtInfo struct {
    59  	typ     OpaqueType
    60  	buildFn BuildOpaqueFn
    61  }
    62  
    63  var opaqueStatements = make(map[reflect.Type]opaqueStmtInfo)
    64  
    65  func (b *Builder) tryBuildOpaque(stmt tree.Statement, inScope *scope) (outScope *scope) {
    66  	info, ok := opaqueStatements[reflect.TypeOf(stmt)]
    67  	if !ok {
    68  		return nil
    69  	}
    70  	obj, cols, err := info.buildFn(b.ctx, b.semaCtx, b.evalCtx, stmt)
    71  	if err != nil {
    72  		panic(err)
    73  	}
    74  	outScope = inScope.push()
    75  	b.synthesizeResultColumns(outScope, cols)
    76  	private := &memo.OpaqueRelPrivate{
    77  		Columns:  colsToColList(outScope.cols),
    78  		Metadata: obj,
    79  	}
    80  	switch info.typ {
    81  	case OpaqueReadOnly:
    82  		outScope.expr = b.factory.ConstructOpaqueRel(private)
    83  	case OpaqueMutation:
    84  		outScope.expr = b.factory.ConstructOpaqueMutation(private)
    85  	case OpaqueDDL:
    86  		outScope.expr = b.factory.ConstructOpaqueDDL(private)
    87  	default:
    88  		panic(errors.AssertionFailedf("invalid opaque statement type %d", info.typ))
    89  	}
    90  	return outScope
    91  }