github.com/dolthub/dolt/go@v0.40.5-0.20240520175717-68db7794bea6/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  
    22  	"github.com/dolthub/dolt/go/libraries/doltcore/env"
    23  
    24  	"github.com/dolthub/dolt/go/libraries/doltcore/doltdb"
    25  	"github.com/dolthub/dolt/go/libraries/doltcore/schema"
    26  	"github.com/dolthub/dolt/go/libraries/doltcore/table"
    27  	"github.com/dolthub/dolt/go/libraries/doltcore/table/editor"
    28  	"github.com/dolthub/dolt/go/libraries/doltcore/table/untyped/csv"
    29  	"github.com/dolthub/dolt/go/libraries/utils/filesys"
    30  	"github.com/dolthub/dolt/go/libraries/utils/iohelp"
    31  )
    32  
    33  // StreamDataLocation is a process stream that that can be imported from or exported to.
    34  type StreamDataLocation struct {
    35  	Format DataFormat
    36  	Writer io.WriteCloser
    37  	Reader io.ReadCloser
    38  }
    39  
    40  // String returns a string representation of the data location.
    41  func (dl StreamDataLocation) String() string {
    42  	return "stream"
    43  }
    44  
    45  // Exists returns true if the DataLocation already exists
    46  func (dl StreamDataLocation) Exists(ctx context.Context, root doltdb.RootValue, fs filesys.ReadableFS) (bool, error) {
    47  	return true, nil
    48  }
    49  
    50  // NewReader creates a TableReadCloser for the DataLocation
    51  func (dl StreamDataLocation) NewReader(ctx context.Context, dEnv *env.DoltEnv, opts interface{}) (rdCl table.SqlRowReader, sorted bool, err error) {
    52  	root, err := dEnv.WorkingRoot(ctx)
    53  	if err != nil {
    54  		return nil, false, err
    55  	}
    56  
    57  	switch dl.Format {
    58  	case CsvFile:
    59  		delim := ","
    60  
    61  		if opts != nil {
    62  			csvOpts, _ := opts.(CsvOptions)
    63  
    64  			if len(csvOpts.Delim) != 0 {
    65  				delim = csvOpts.Delim
    66  			}
    67  		}
    68  
    69  		rd, err := csv.NewCSVReader(root.VRW().Format(), io.NopCloser(dl.Reader), csv.NewCSVInfo().SetDelim(delim))
    70  
    71  		return rd, false, err
    72  
    73  	case PsvFile:
    74  		rd, err := csv.NewCSVReader(root.VRW().Format(), io.NopCloser(dl.Reader), csv.NewCSVInfo().SetDelim("|"))
    75  		return rd, false, err
    76  	}
    77  
    78  	return nil, false, errors.New(string(dl.Format) + "is an unsupported format to read from stdin")
    79  }
    80  
    81  // NewCreatingWriter will create a TableWriteCloser for a DataLocation that will create a new table, or overwrite
    82  // an existing table.
    83  func (dl StreamDataLocation) NewCreatingWriter(ctx context.Context, mvOpts DataMoverOptions, root doltdb.RootValue, outSch schema.Schema, opts editor.Options, wr io.WriteCloser) (table.SqlRowWriter, error) {
    84  	switch dl.Format {
    85  	case CsvFile:
    86  		return csv.NewCSVWriter(iohelp.NopWrCloser(dl.Writer), outSch, csv.NewCSVInfo())
    87  
    88  	case PsvFile:
    89  		return csv.NewCSVWriter(iohelp.NopWrCloser(dl.Writer), outSch, csv.NewCSVInfo().SetDelim("|"))
    90  	}
    91  
    92  	return nil, errors.New(string(dl.Format) + "is an unsupported format to write to stdout")
    93  }