github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/sql/copy_file_upload.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  	"fmt"
    16  	"io"
    17  	"sync"
    18  
    19  	"github.com/cockroachdb/cockroach/pkg/blobs"
    20  	"github.com/cockroachdb/cockroach/pkg/sql/lex"
    21  	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgwirebase"
    22  	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
    23  	"github.com/cockroachdb/cockroach/pkg/sql/sqlbase"
    24  	"github.com/cockroachdb/cockroach/pkg/sql/types"
    25  	"github.com/cockroachdb/errors"
    26  	"github.com/lib/pq"
    27  )
    28  
    29  const (
    30  	fileUploadTable = "file_upload"
    31  	copyOptionDest  = "destination"
    32  )
    33  
    34  var copyFileOptionExpectValues = map[string]KVStringOptValidate{
    35  	copyOptionDest: KVStringOptRequireValue,
    36  }
    37  
    38  var _ copyMachineInterface = &fileUploadMachine{}
    39  
    40  type fileUploadMachine struct {
    41  	c              *copyMachine
    42  	writeToFile    *io.PipeWriter
    43  	wg             *sync.WaitGroup
    44  	failureCleanup func()
    45  }
    46  
    47  func newFileUploadMachine(
    48  	ctx context.Context,
    49  	conn pgwirebase.Conn,
    50  	n *tree.CopyFrom,
    51  	txnOpt copyTxnOpt,
    52  	execCfg *ExecutorConfig,
    53  ) (f *fileUploadMachine, retErr error) {
    54  	if len(n.Columns) != 0 {
    55  		return nil, errors.New("expected 0 columns specified for file uploads")
    56  	}
    57  	c := &copyMachine{
    58  		conn: conn,
    59  		// The planner will be prepared before use.
    60  		p: planner{execCfg: execCfg, alloc: &sqlbase.DatumAlloc{}},
    61  	}
    62  	f = &fileUploadMachine{
    63  		c:  c,
    64  		wg: &sync.WaitGroup{},
    65  	}
    66  
    67  	// We need a planner to do the initial planning, even if a planner
    68  	// is not required after that.
    69  	cleanup := c.p.preparePlannerForCopy(ctx, txnOpt)
    70  	defer func() {
    71  		retErr = cleanup(ctx, retErr)
    72  	}()
    73  	c.parsingEvalCtx = c.p.EvalContext()
    74  
    75  	if err := c.p.RequireAdminRole(ctx, "upload to nodelocal"); err != nil {
    76  		return nil, err
    77  	}
    78  
    79  	optsFn, err := f.c.p.TypeAsStringOpts(ctx, n.Options, copyFileOptionExpectValues)
    80  	if err != nil {
    81  		return nil, err
    82  	}
    83  	opts, err := optsFn()
    84  	if err != nil {
    85  		return nil, err
    86  	}
    87  
    88  	pr, pw := io.Pipe()
    89  	localStorage, err := blobs.NewLocalStorage(c.p.execCfg.Settings.ExternalIODir)
    90  	if err != nil {
    91  		return nil, err
    92  	}
    93  	// Check that the file does not already exist
    94  	_, err = localStorage.Stat(opts[copyOptionDest])
    95  	if err == nil {
    96  		return nil, fmt.Errorf("destination file already exists for %s", opts[copyOptionDest])
    97  	}
    98  	f.wg.Add(1)
    99  	go func() {
   100  		err := localStorage.WriteFile(opts[copyOptionDest], pr)
   101  		if err != nil {
   102  			_ = pr.CloseWithError(err)
   103  		}
   104  		f.wg.Done()
   105  	}()
   106  	f.writeToFile = pw
   107  	f.failureCleanup = func() {
   108  		// Ignoring this error because deletion would only fail
   109  		// if the file was not created in the first place.
   110  		_ = localStorage.Delete(opts[copyOptionDest])
   111  	}
   112  
   113  	c.resultColumns = make(sqlbase.ResultColumns, 1)
   114  	c.resultColumns[0] = sqlbase.ResultColumn{Typ: types.Bytes}
   115  	c.parsingEvalCtx = c.p.EvalContext()
   116  	c.rowsMemAcc = c.p.extendedEvalCtx.Mon.MakeBoundAccount()
   117  	c.bufMemAcc = c.p.extendedEvalCtx.Mon.MakeBoundAccount()
   118  	c.processRows = f.writeFile
   119  	return
   120  }
   121  
   122  // CopyInFileStmt creates a COPY FROM statement which can be used
   123  // to upload files, and be prepared with Tx.Prepare().
   124  func CopyInFileStmt(destination, schema, table string) string {
   125  	return fmt.Sprintf(
   126  		`COPY %s.%s FROM STDIN WITH destination = %s`,
   127  		pq.QuoteIdentifier(schema),
   128  		pq.QuoteIdentifier(table),
   129  		lex.EscapeSQLString(destination),
   130  	)
   131  }
   132  
   133  func (f *fileUploadMachine) run(ctx context.Context) (err error) {
   134  	err = f.c.run(ctx)
   135  	_ = f.writeToFile.Close()
   136  	if err != nil {
   137  		f.failureCleanup()
   138  	}
   139  	f.wg.Wait()
   140  	return
   141  }
   142  
   143  func (f *fileUploadMachine) writeFile(ctx context.Context) error {
   144  	for _, r := range f.c.rows {
   145  		b := []byte(*r[0].(*tree.DBytes))
   146  		n, err := f.writeToFile.Write(b)
   147  		if err != nil {
   148  			return err
   149  		}
   150  		if n < len(b) {
   151  			return io.ErrShortWrite
   152  		}
   153  	}
   154  
   155  	// Issue a final zero-byte write to ensure we observe any errors in the pipe.
   156  	_, err := f.writeToFile.Write(nil)
   157  	if err != nil {
   158  		return err
   159  	}
   160  	f.c.insertedRows += len(f.c.rows)
   161  	f.c.rows = f.c.rows[:0]
   162  	f.c.rowsMemAcc.Clear(ctx)
   163  	return nil
   164  }