github.com/hasnat/dolt/go@v0.0.0-20210628190320-9eb5d843fbb7/libraries/doltcore/mvdata/stream_data_loc.go (about)

     1  // Copyright 2019 Dolthub, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package mvdata
    16  
    17  import (
    18  	"context"
    19  	"errors"
    20  	"io"
    21  	"io/ioutil"
    22  
    23  	"github.com/dolthub/dolt/go/libraries/doltcore/env"
    24  
    25  	"github.com/dolthub/dolt/go/libraries/doltcore/doltdb"
    26  	"github.com/dolthub/dolt/go/libraries/doltcore/schema"
    27  	"github.com/dolthub/dolt/go/libraries/doltcore/table"
    28  	"github.com/dolthub/dolt/go/libraries/doltcore/table/typed/noms"
    29  	"github.com/dolthub/dolt/go/libraries/doltcore/table/untyped/csv"
    30  	"github.com/dolthub/dolt/go/libraries/utils/filesys"
    31  	"github.com/dolthub/dolt/go/libraries/utils/iohelp"
    32  )
    33  
    34  // StreamDataLocation is a process stream that that can be imported from or exported to.
    35  type StreamDataLocation struct {
    36  	Format DataFormat
    37  	Writer io.WriteCloser
    38  	Reader io.ReadCloser
    39  }
    40  
    41  // String returns a string representation of the data location.
    42  func (dl StreamDataLocation) String() string {
    43  	return "stream"
    44  }
    45  
    46  // Exists returns true if the DataLocation already exists
    47  func (dl StreamDataLocation) Exists(ctx context.Context, root *doltdb.RootValue, fs filesys.ReadableFS) (bool, error) {
    48  	return true, nil
    49  }
    50  
    51  // NewReader creates a TableReadCloser for the DataLocation
    52  func (dl StreamDataLocation) NewReader(ctx context.Context, root *doltdb.RootValue, fs filesys.ReadableFS, opts interface{}) (rdCl table.TableReadCloser, sorted bool, err error) {
    53  	switch dl.Format {
    54  	case CsvFile:
    55  		delim := ","
    56  
    57  		if opts != nil {
    58  			csvOpts, _ := opts.(CsvOptions)
    59  
    60  			if len(csvOpts.Delim) != 0 {
    61  				delim = csvOpts.Delim
    62  			}
    63  		}
    64  
    65  		rd, err := csv.NewCSVReader(root.VRW().Format(), ioutil.NopCloser(dl.Reader), csv.NewCSVInfo().SetDelim(delim))
    66  
    67  		return rd, false, err
    68  
    69  	case PsvFile:
    70  		rd, err := csv.NewCSVReader(root.VRW().Format(), ioutil.NopCloser(dl.Reader), csv.NewCSVInfo().SetDelim("|"))
    71  		return rd, false, err
    72  	}
    73  
    74  	return nil, false, errors.New(string(dl.Format) + "is an unsupported format to read from stdin")
    75  }
    76  
    77  // NewCreatingWriter will create a TableWriteCloser for a DataLocation that will create a new table, or overwrite
    78  // an existing table.
    79  func (dl StreamDataLocation) NewCreatingWriter(_ context.Context, _ DataMoverOptions, _ *env.DoltEnv, _ *doltdb.RootValue, _ bool, outSch schema.Schema, _ noms.StatsCB, _ bool) (table.TableWriteCloser, error) {
    80  	switch dl.Format {
    81  	case CsvFile:
    82  		return csv.NewCSVWriter(iohelp.NopWrCloser(dl.Writer), outSch, csv.NewCSVInfo())
    83  
    84  	case PsvFile:
    85  		return csv.NewCSVWriter(iohelp.NopWrCloser(dl.Writer), outSch, csv.NewCSVInfo().SetDelim("|"))
    86  	}
    87  
    88  	return nil, errors.New(string(dl.Format) + "is an unsupported format to write to stdout")
    89  }
    90  
    91  // NewUpdatingWriter will create a TableWriteCloser for a DataLocation that will update and append rows based on
    92  // their primary key.
    93  func (dl StreamDataLocation) NewUpdatingWriter(_ context.Context, _ DataMoverOptions, _ *env.DoltEnv, _ *doltdb.RootValue, _ bool, _ schema.Schema, _ noms.StatsCB, _ bool) (table.TableWriteCloser, error) {
    94  	panic("Updating is not supported for stdout")
    95  }
    96  
    97  // NewReplacingWriter will create a TableWriteCloser for a DataLocation that will overwrite an existing table while
    98  // preserving schema
    99  func (dl StreamDataLocation) NewReplacingWriter(_ context.Context, _ DataMoverOptions, _ *env.DoltEnv, _ *doltdb.RootValue, _ bool, _ schema.Schema, _ noms.StatsCB, _ bool) (table.TableWriteCloser, error) {
   100  	panic("Replacing is not supported for stdout")
   101  }