github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/storage/cloud/workload_storage.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 cloud
    12  
    13  import (
    14  	"context"
    15  	"io"
    16  	"io/ioutil"
    17  	"net/url"
    18  	"strconv"
    19  	"strings"
    20  
    21  	"github.com/cockroachdb/cockroach/pkg/roachpb"
    22  	"github.com/cockroachdb/cockroach/pkg/workload"
    23  	"github.com/cockroachdb/errors"
    24  )
    25  
    26  type workloadStorage struct {
    27  	conf  *roachpb.ExternalStorage_Workload
    28  	gen   workload.Generator
    29  	table workload.Table
    30  }
    31  
    32  var _ ExternalStorage = &workloadStorage{}
    33  
    34  func makeWorkloadStorage(conf *roachpb.ExternalStorage_Workload) (ExternalStorage, error) {
    35  	if conf == nil {
    36  		return nil, errors.Errorf("workload upload requested but info missing")
    37  	}
    38  	if strings.ToLower(conf.Format) != `csv` {
    39  		return nil, errors.Errorf(`unsupported format: %s`, conf.Format)
    40  	}
    41  	meta, err := workload.Get(conf.Generator)
    42  	if err != nil {
    43  		return nil, err
    44  	}
    45  	// Different versions of the workload could generate different data, so
    46  	// disallow this.
    47  	if meta.Version != conf.Version {
    48  		return nil, errors.Errorf(
    49  			`expected %s version "%s" but got "%s"`, meta.Name, conf.Version, meta.Version)
    50  	}
    51  	gen := meta.New()
    52  	if f, ok := gen.(workload.Flagser); ok {
    53  		if err := f.Flags().Parse(conf.Flags); err != nil {
    54  			return nil, errors.Wrapf(err, `parsing parameters %s`, strings.Join(conf.Flags, ` `))
    55  		}
    56  	}
    57  	s := &workloadStorage{
    58  		conf: conf,
    59  		gen:  gen,
    60  	}
    61  	for _, t := range gen.Tables() {
    62  		if t.Name == conf.Table {
    63  			s.table = t
    64  			break
    65  		}
    66  	}
    67  	if s.table.Name == `` {
    68  		return nil, errors.Wrapf(err, `unknown table %s for generator %s`, conf.Table, meta.Name)
    69  	}
    70  	return s, nil
    71  }
    72  
    73  func (s *workloadStorage) Conf() roachpb.ExternalStorage {
    74  	return roachpb.ExternalStorage{
    75  		Provider:       roachpb.ExternalStorageProvider_Workload,
    76  		WorkloadConfig: s.conf,
    77  	}
    78  }
    79  
    80  func (s *workloadStorage) ReadFile(_ context.Context, basename string) (io.ReadCloser, error) {
    81  	if basename != `` {
    82  		return nil, errors.Errorf(`basenames are not supported by workload storage`)
    83  	}
    84  	r := workload.NewCSVRowsReader(s.table, int(s.conf.BatchBegin), int(s.conf.BatchEnd))
    85  	return ioutil.NopCloser(r), nil
    86  }
    87  
    88  func (s *workloadStorage) WriteFile(_ context.Context, _ string, _ io.ReadSeeker) error {
    89  	return errors.Errorf(`workload storage does not support writes`)
    90  }
    91  
    92  func (s *workloadStorage) ListFiles(_ context.Context, _ string) ([]string, error) {
    93  	return nil, errors.Errorf(`workload storage does not support listing files`)
    94  }
    95  
    96  func (s *workloadStorage) Delete(_ context.Context, _ string) error {
    97  	return errors.Errorf(`workload storage does not support deletes`)
    98  }
    99  func (s *workloadStorage) Size(_ context.Context, _ string) (int64, error) {
   100  	return 0, errors.Errorf(`workload storage does not support sizing`)
   101  }
   102  func (s *workloadStorage) Close() error {
   103  	return nil
   104  }
   105  
   106  // ParseWorkloadConfig parses a workload config URI to a proto config.
   107  func ParseWorkloadConfig(uri *url.URL) (*roachpb.ExternalStorage_Workload, error) {
   108  	c := &roachpb.ExternalStorage_Workload{}
   109  	pathParts := strings.Split(strings.Trim(uri.Path, `/`), `/`)
   110  	if len(pathParts) != 3 {
   111  		return nil, errors.Errorf(
   112  			`path must be of the form /<format>/<generator>/<table>: %s`, uri.Path)
   113  	}
   114  	c.Format, c.Generator, c.Table = pathParts[0], pathParts[1], pathParts[2]
   115  	q := uri.Query()
   116  	if _, ok := q[`version`]; !ok {
   117  		return nil, errors.New(`parameter version is required`)
   118  	}
   119  	c.Version = q.Get(`version`)
   120  	q.Del(`version`)
   121  	if s := q.Get(`row-start`); len(s) > 0 {
   122  		q.Del(`row-start`)
   123  		var err error
   124  		if c.BatchBegin, err = strconv.ParseInt(s, 10, 64); err != nil {
   125  			return nil, err
   126  		}
   127  	}
   128  	if e := q.Get(`row-end`); len(e) > 0 {
   129  		q.Del(`row-end`)
   130  		var err error
   131  		if c.BatchEnd, err = strconv.ParseInt(e, 10, 64); err != nil {
   132  			return nil, err
   133  		}
   134  	}
   135  	for k, vs := range q {
   136  		for _, v := range vs {
   137  			c.Flags = append(c.Flags, `--`+k+`=`+v)
   138  		}
   139  	}
   140  	return c, nil
   141  }