github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/sql/export.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  	"strconv"
    16  	"strings"
    17  
    18  	"github.com/cockroachdb/cockroach/pkg/roachpb"
    19  	"github.com/cockroachdb/cockroach/pkg/sql/execinfrapb"
    20  	"github.com/cockroachdb/cockroach/pkg/sql/opt/exec"
    21  	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode"
    22  	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror"
    23  	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
    24  	"github.com/cockroachdb/cockroach/pkg/util"
    25  	"github.com/cockroachdb/errors"
    26  )
    27  
    28  type exportNode struct {
    29  	optColumnsSlot
    30  
    31  	source planNode
    32  
    33  	fileName        string
    34  	csvOpts         roachpb.CSVOptions
    35  	chunkSize       int
    36  	fileCompression execinfrapb.FileCompression
    37  }
    38  
    39  func (e *exportNode) startExec(params runParams) error {
    40  	panic("exportNode cannot be run in local mode")
    41  }
    42  
    43  func (e *exportNode) Next(params runParams) (bool, error) {
    44  	panic("exportNode cannot be run in local mode")
    45  }
    46  
    47  func (e *exportNode) Values() tree.Datums {
    48  	panic("exportNode cannot be run in local mode")
    49  }
    50  
    51  func (e *exportNode) Close(ctx context.Context) {
    52  	e.source.Close(ctx)
    53  }
    54  
    55  const (
    56  	exportOptionDelimiter   = "delimiter"
    57  	exportOptionNullAs      = "nullas"
    58  	exportOptionChunkSize   = "chunk_rows"
    59  	exportOptionFileName    = "filename"
    60  	exportOptionCompression = "compression"
    61  )
    62  
    63  var exportOptionExpectValues = map[string]KVStringOptValidate{
    64  	exportOptionChunkSize:   KVStringOptRequireValue,
    65  	exportOptionDelimiter:   KVStringOptRequireValue,
    66  	exportOptionFileName:    KVStringOptRequireValue,
    67  	exportOptionNullAs:      KVStringOptRequireValue,
    68  	exportOptionCompression: KVStringOptRequireValue,
    69  }
    70  
    71  const exportChunkSizeDefault = 100000
    72  const exportFilePatternPart = "%part%"
    73  const exportFilePatternDefault = exportFilePatternPart + ".csv"
    74  const exportCompressionCodec = "gzip"
    75  
    76  // ConstructExport is part of the exec.Factory interface.
    77  func (ef *execFactory) ConstructExport(
    78  	input exec.Node, fileName tree.TypedExpr, fileFormat string, options []exec.KVOption,
    79  ) (exec.Node, error) {
    80  	if !ef.planner.ExtendedEvalContext().TxnImplicit {
    81  		return nil, errors.Errorf("EXPORT cannot be used inside a transaction")
    82  	}
    83  
    84  	if fileFormat != "CSV" {
    85  		return nil, errors.Errorf("unsupported export format: %q", fileFormat)
    86  	}
    87  
    88  	fileNameDatum, err := fileName.Eval(ef.planner.EvalContext())
    89  	if err != nil {
    90  		return nil, err
    91  	}
    92  	fileNameStr, ok := fileNameDatum.(*tree.DString)
    93  	if !ok {
    94  		return nil, errors.Errorf("expected string value for the file location")
    95  	}
    96  
    97  	optVals, err := evalStringOptions(ef.planner.EvalContext(), options, exportOptionExpectValues)
    98  	if err != nil {
    99  		return nil, err
   100  	}
   101  
   102  	csvOpts := roachpb.CSVOptions{}
   103  
   104  	if override, ok := optVals[exportOptionDelimiter]; ok {
   105  		csvOpts.Comma, err = util.GetSingleRune(override)
   106  		if err != nil {
   107  			return nil, pgerror.New(pgcode.InvalidParameterValue, "invalid delimiter")
   108  		}
   109  	}
   110  
   111  	if override, ok := optVals[exportOptionNullAs]; ok {
   112  		csvOpts.NullEncoding = &override
   113  	}
   114  
   115  	chunkSize := exportChunkSizeDefault
   116  	if override, ok := optVals[exportOptionChunkSize]; ok {
   117  		chunkSize, err = strconv.Atoi(override)
   118  		if err != nil {
   119  			return nil, pgerror.WithCandidateCode(err, pgcode.InvalidParameterValue)
   120  		}
   121  		if chunkSize < 1 {
   122  			return nil, pgerror.New(pgcode.InvalidParameterValue, "invalid csv chunk size")
   123  		}
   124  	}
   125  
   126  	// Check whenever compression is expected and extract compression codec name in case
   127  	// of positive result
   128  	var codec execinfrapb.FileCompression
   129  	if name, ok := optVals[exportOptionCompression]; ok && len(name) != 0 {
   130  		if strings.EqualFold(name, exportCompressionCodec) {
   131  			codec = execinfrapb.FileCompression_Gzip
   132  		} else {
   133  			return nil, pgerror.Newf(pgcode.InvalidParameterValue,
   134  				"unsupported compression codec %s", name)
   135  		}
   136  	}
   137  
   138  	return &exportNode{
   139  		source:          input.(planNode),
   140  		fileName:        string(*fileNameStr),
   141  		csvOpts:         csvOpts,
   142  		chunkSize:       chunkSize,
   143  		fileCompression: codec,
   144  	}, nil
   145  }