github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/cli/nodelocal.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 cli
    12  
    13  import (
    14  	"database/sql/driver"
    15  	"fmt"
    16  	"io"
    17  	"os"
    18  	"path/filepath"
    19  
    20  	"github.com/cockroachdb/cockroach/pkg/sql"
    21  	"github.com/cockroachdb/errors"
    22  	"github.com/spf13/cobra"
    23  )
    24  
    25  const (
    26  	chunkSize = 4 * 1024
    27  )
    28  
    29  var nodeLocalUploadCmd = &cobra.Command{
    30  	Use:   "upload <source> <destination>",
    31  	Short: "Upload file from source to destination",
    32  	Long: `
    33  Uploads a file to a gateway node's local file system using a SQL connection.
    34  `,
    35  	Args: cobra.MinimumNArgs(2),
    36  	RunE: maybeShoutError(runUpload),
    37  }
    38  
    39  func runUpload(cmd *cobra.Command, args []string) error {
    40  	conn, err := makeSQLClient("cockroach nodelocal", useSystemDb)
    41  	if err != nil {
    42  		return err
    43  	}
    44  	defer conn.Close()
    45  
    46  	source := args[0]
    47  	destination := args[1]
    48  	reader, err := openSourceFile(source)
    49  	if err != nil {
    50  		return err
    51  	}
    52  	defer reader.Close()
    53  	return uploadFile(conn, reader, destination)
    54  }
    55  
    56  func openSourceFile(source string) (io.ReadCloser, error) {
    57  	f, err := os.Open(source)
    58  	if err != nil {
    59  		return nil, err
    60  	}
    61  	stat, err := f.Stat()
    62  	if err != nil {
    63  		return nil, errors.Wrapf(err, "unable to get source file stats for %s", source)
    64  	}
    65  	if stat.IsDir() {
    66  		return nil, fmt.Errorf("source file %s is a directory, not a file", source)
    67  	}
    68  	return f, nil
    69  }
    70  
    71  func uploadFile(conn *sqlConn, reader io.Reader, destination string) error {
    72  	if err := conn.ensureConn(); err != nil {
    73  		return err
    74  	}
    75  
    76  	if _, err := conn.conn.Exec(`BEGIN`, nil); err != nil {
    77  		return err
    78  	}
    79  
    80  	stmt, err := conn.conn.Prepare(sql.CopyInFileStmt(destination, "crdb_internal", "file_upload"))
    81  	if err != nil {
    82  		return err
    83  	}
    84  
    85  	defer func() {
    86  		if stmt != nil {
    87  			_ = stmt.Close()
    88  			_, _ = conn.conn.Exec(`ROLLBACK`, nil)
    89  		}
    90  	}()
    91  
    92  	send := make([]byte, chunkSize)
    93  	for {
    94  		n, err := reader.Read(send)
    95  		if n > 0 {
    96  			_, err = stmt.Exec([]driver.Value{string(send[:n])})
    97  			if err != nil {
    98  				return err
    99  			}
   100  		} else if err == io.EOF {
   101  			break
   102  		} else if err != nil {
   103  			return err
   104  		}
   105  	}
   106  	if err := stmt.Close(); err != nil {
   107  		return err
   108  	}
   109  	stmt = nil
   110  
   111  	if _, err := conn.conn.Exec(`COMMIT`, nil); err != nil {
   112  		return err
   113  	}
   114  
   115  	nodeID, _, _, err := conn.getServerMetadata()
   116  	if err != nil {
   117  		return errors.Wrap(err, "unable to get node id")
   118  	}
   119  	fmt.Printf("successfully uploaded to nodelocal://%s\n", filepath.Join(nodeID.String(), destination))
   120  	return nil
   121  }
   122  
   123  var nodeLocalCmds = []*cobra.Command{
   124  	nodeLocalUploadCmd,
   125  }
   126  
   127  var nodeLocalCmd = &cobra.Command{
   128  	Use:   "nodelocal [command]",
   129  	Short: "upload and delete nodelocal files",
   130  	Long:  "Upload and delete files on the gateway node's local file system.",
   131  	RunE:  usageAndErr,
   132  }
   133  
   134  func init() {
   135  	nodeLocalCmd.AddCommand(nodeLocalCmds...)
   136  }