minio/pkg/s3select/csv/reader.go

/*
 * Minio Cloud Storage, (C) 2019 Minio, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package csv

import (
	"bytes"
	"encoding/csv"
	"fmt"
	"io"

	"github.com/minio/minio/pkg/s3select/sql"
)

type recordReader struct {
	reader          io.Reader
	recordDelimiter []byte
	oneByte         []byte
	useOneByte      bool
}

func (rr *recordReader) Read(p []byte) (n int, err error) {
	if rr.useOneByte {
		p[0] = rr.oneByte[0]
		rr.useOneByte = false
		n, err = rr.reader.Read(p[1:])
		n++
	} else {
		n, err = rr.reader.Read(p)
	}

	if err != nil {
		return 0, err
	}

	if string(rr.recordDelimiter) == "\n" {
		return n, nil
	}

	for {
		i := bytes.Index(p, rr.recordDelimiter)
		if i < 0 {
			break
		}

		p[i] = '\n'
		if len(rr.recordDelimiter) > 1 {
			p = append(p[:i+1], p[i+len(rr.recordDelimiter):]...)
		}
	}

	n = len(p)
	if len(rr.recordDelimiter) == 1 || p[n-1] != rr.recordDelimiter[0] {
		return n, nil
	}

	if _, err = rr.reader.Read(rr.oneByte); err != nil {
		return 0, err
	}

	if rr.oneByte[0] == rr.recordDelimiter[1] {
		p[n-1] = '\n'
		return n, nil
	}

	rr.useOneByte = true
	return n, nil
}

// Reader - CSV record reader for S3Select.
type Reader struct {
	args        *ReaderArgs
	readCloser  io.ReadCloser
	csvReader   *csv.Reader
	columnNames []string
}

// Read - reads single record.
func (r *Reader) Read() (sql.Record, error) {
	csvRecord, err := r.csvReader.Read()
	if err != nil {
		if err != io.EOF {
			return nil, errCSVParsingError(err)
		}

		return nil, err
	}

	columnNames := r.columnNames
	if columnNames == nil {
		columnNames = make([]string, len(csvRecord))
		for i := range csvRecord {
			columnNames[i] = fmt.Sprintf("_%v", i+1)
		}
	}

	nameIndexMap := make(map[string]int64)
	for i := range columnNames {
		nameIndexMap[columnNames[i]] = int64(i)
	}

	return &Record{
		columnNames:  columnNames,
		csvRecord:    csvRecord,
		nameIndexMap: nameIndexMap,
	}, nil
}

// Close - closes underlaying reader.
func (r *Reader) Close() error {
	return r.readCloser.Close()
}

// NewReader - creates new CSV reader using readCloser.
func NewReader(readCloser io.ReadCloser, args *ReaderArgs) (*Reader, error) {
	if args == nil || args.IsEmpty() {
		panic(fmt.Errorf("empty args passed %v", args))
	}

	csvReader := csv.NewReader(&recordReader{
		reader:          readCloser,
		recordDelimiter: []byte(args.RecordDelimiter),
		oneByte:         []byte{0},
	})
	csvReader.Comma = []rune(args.FieldDelimiter)[0]
	csvReader.Comment = []rune(args.CommentCharacter)[0]
	csvReader.FieldsPerRecord = -1

	r := &Reader{
		args:       args,
		readCloser: readCloser,
		csvReader:  csvReader,
	}

	if args.FileHeaderInfo == none {
		return r, nil
	}

	record, err := csvReader.Read()
	if err != nil {
		if err != io.EOF {
			return nil, errCSVParsingError(err)
		}

		return nil, err
	}

	if args.FileHeaderInfo == use {
		r.columnNames = record
	}

	return r, nil
}