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

     1  // Copyright 2018 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  	"fmt"
    16  
    17  	"github.com/cockroachdb/cockroach/pkg/sql/execinfra"
    18  	"github.com/cockroachdb/cockroach/pkg/sql/execinfrapb"
    19  	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
    20  	"github.com/cockroachdb/cockroach/pkg/sql/sqlbase"
    21  )
    22  
    23  // rowSourceToPlanNode wraps a RowSource and presents it as a PlanNode. It must
    24  // be constructed with Create(), after which it is a PlanNode and can be treated
    25  // as such.
    26  type rowSourceToPlanNode struct {
    27  	source    execinfra.RowSource
    28  	forwarder metadataForwarder
    29  
    30  	// originalPlanNode is the original planNode that the wrapped RowSource got
    31  	// planned for.
    32  	originalPlanNode planNode
    33  
    34  	planCols sqlbase.ResultColumns
    35  
    36  	// Temporary variables
    37  	row      sqlbase.EncDatumRow
    38  	da       sqlbase.DatumAlloc
    39  	datumRow tree.Datums
    40  }
    41  
    42  var _ planNode = &rowSourceToPlanNode{}
    43  
    44  // makeRowSourceToPlanNode creates a new planNode that wraps a RowSource. It
    45  // takes an optional metadataForwarder, which if non-nil is invoked for every
    46  // piece of metadata this wrapper receives from the wrapped RowSource.
    47  // It also takes an optional planNode, which is the planNode that the RowSource
    48  // that this rowSourceToPlanNode is wrapping originally replaced. That planNode
    49  // will be closed when this one is closed.
    50  func makeRowSourceToPlanNode(
    51  	s execinfra.RowSource,
    52  	forwarder metadataForwarder,
    53  	planCols sqlbase.ResultColumns,
    54  	originalPlanNode planNode,
    55  ) *rowSourceToPlanNode {
    56  	row := make(tree.Datums, len(planCols))
    57  
    58  	return &rowSourceToPlanNode{
    59  		source:           s,
    60  		datumRow:         row,
    61  		forwarder:        forwarder,
    62  		planCols:         planCols,
    63  		originalPlanNode: originalPlanNode,
    64  	}
    65  }
    66  
    67  func (r *rowSourceToPlanNode) startExec(params runParams) error {
    68  	r.source.Start(params.ctx)
    69  	return nil
    70  }
    71  
    72  func (r *rowSourceToPlanNode) Next(params runParams) (bool, error) {
    73  	for {
    74  		var p *execinfrapb.ProducerMetadata
    75  		r.row, p = r.source.Next()
    76  
    77  		if p != nil {
    78  			if p.Err != nil {
    79  				return false, p.Err
    80  			}
    81  			if r.forwarder != nil {
    82  				r.forwarder.forwardMetadata(p)
    83  				continue
    84  			}
    85  			if p.TraceData != nil {
    86  				// We drop trace metadata since we have no reasonable way to propagate
    87  				// it in local SQL execution.
    88  				continue
    89  			}
    90  			return false, fmt.Errorf("unexpected producer metadata: %+v", p)
    91  		}
    92  
    93  		if r.row == nil {
    94  			return false, nil
    95  		}
    96  
    97  		types := r.source.OutputTypes()
    98  		for i := range r.planCols {
    99  			encDatum := r.row[i]
   100  			err := encDatum.EnsureDecoded(types[i], &r.da)
   101  			if err != nil {
   102  				return false, err
   103  			}
   104  			r.datumRow[i] = encDatum.Datum
   105  		}
   106  
   107  		return true, nil
   108  	}
   109  }
   110  
   111  func (r *rowSourceToPlanNode) Values() tree.Datums {
   112  	return r.datumRow
   113  }
   114  
   115  func (r *rowSourceToPlanNode) Close(ctx context.Context) {
   116  	if r.source != nil {
   117  		r.source.ConsumerClosed()
   118  		r.source = nil
   119  	}
   120  	if r.originalPlanNode != nil {
   121  		r.originalPlanNode.Close(ctx)
   122  	}
   123  }