dolthub--doltgresql/core/dataloader/tabdataloader.go

// Copyright 2024 Dolthub, 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 dataloader

import (
	"bufio"
	"io"
	"strings"

	"github.com/cockroachdb/errors"
	"github.com/dolthub/go-mysql-server/sql"

	"github.com/dolthub/doltgresql/server/types"
)

const defaultTextDelimiter = "\t"
const defaultNullChar = "\\N"

// TabularDataLoader tracks the state of a load data operation from a tabular data source.
type TabularDataLoader struct {
	results       LoadDataResults
	partialLine   strings.Builder
	nextDataChunk *bufio.Reader
	colTypes      []*types.DoltgresType
	sch           sql.Schema
	delimiterChar string
	nullChar      string
	removeHeader  bool
}

var _ DataLoader = (*TabularDataLoader)(nil)

// NewTabularDataLoader creates a new TabularDataLoader to insert into the specified |table| using the specified
// |delimiterChar| and |nullChar|. If |header| is true, the first line of the data will be treated as a header and
// ignored.
func NewTabularDataLoader(colNames []string, tableSch sql.Schema, delimiterChar, nullChar string, header bool) (*TabularDataLoader, error) {
	colTypes, reducedSch, err := getColumnTypes(colNames, tableSch)
	if err != nil {
		return nil, err
	}

	if delimiterChar == "" {
		delimiterChar = defaultTextDelimiter
	}

	if nullChar == "" {
		nullChar = defaultNullChar
	}

	return &TabularDataLoader{
		colTypes:      colTypes,
		sch:           reducedSch,
		delimiterChar: delimiterChar,
		nullChar:      nullChar,
		removeHeader:  header,
	}, nil
}

// nextRow returns the next SQL row from the reader provided, using any previously saved partial line. Returns true if
// there was another row.
func (tdl *TabularDataLoader) nextRow(ctx *sql.Context, data *bufio.Reader) (sql.Row, bool, error) {
	if tdl.removeHeader {
		_, err := data.ReadString('\n')
		tdl.removeHeader = false
		if err != nil {
			return nil, false, err
		}
	}

	for {
		// Read the next line from the file
		line, err := data.ReadString('\n')
		if err != nil {
			if err != io.EOF {
				return nil, false, err
			}

			// bufio.Reader.ReadString will return an error AND a line
			// if the final contents of the data does NOT end in the
			// delimiter. In this case, that means that we need to save
			// the partial line and use it in the next chunk.
			tdl.partialLine.WriteString(line)
			return nil, false, nil
		}

		// If we've not reached EOF, then there will be a newline appended to the end that we must remove.
		line = strings.TrimSuffix(line, "\n")
		// Data with windows line endings will also have a carriage return character that we need to remove.
		line = strings.TrimSuffix(line, "\r")

		if tdl.partialLine.Len() > 0 {
			tdl.partialLine.WriteString(line)
			line = tdl.partialLine.String()
			tdl.partialLine.Reset()
		}

		// If we see the end of data marker, return early
		if line == `\.` {
			return nil, false, nil
		}

		// Skip over empty lines
		if len(line) == 0 {
			continue
		}

		// Split the values by the delimiter, ensuring the correct number of values have been read
		values := strings.Split(line, tdl.delimiterChar)
		if len(values) > len(tdl.colTypes) {
			return nil, false, errors.Errorf("extra data after last expected column")
		} else if len(values) < len(tdl.colTypes) {
			return nil, false, errors.Errorf(`missing data for column "%s"`, tdl.sch[len(values)].Name)
		}

		// Cast the values using I/O input
		row := make(sql.Row, len(tdl.colTypes))
		for i := range tdl.colTypes {
			if values[i] == tdl.nullChar {
				row[i] = nil
			} else {
				// We must un-escape strings here since we're receiving everything verbatim
				values[i] = strings.ReplaceAll(values[i], `\\`, `\`)
				row[i], err = tdl.colTypes[i].IoInput(ctx, values[i])
				if err != nil {
					return nil, false, err
				}
			}
		}

		return row, true, nil
	}
}

func (tdl *TabularDataLoader) SetNextDataChunk(ctx *sql.Context, data *bufio.Reader) error {
	tdl.nextDataChunk = data
	return nil
}

// Finish completes the current load data operation and finalizes the data that has been inserted.
func (tdl *TabularDataLoader) Finish(ctx *sql.Context) (*LoadDataResults, error) {
	// If there is partial data from the last chunk that hasn't been inserted, return an error.
	if tdl.partialLine.Len() > 0 {
		return nil, errors.Errorf("partial line found at end of data load")
	}

	return &tdl.results, nil
}

func (tdl *TabularDataLoader) Resolved() bool {
	return true
}

func (tdl *TabularDataLoader) String() string {
	return "TabularDataLoader"
}

func (tdl *TabularDataLoader) Schema(ctx *sql.Context) sql.Schema {
	return tdl.sch
}

func (tdl *TabularDataLoader) Children() []sql.Node {
	return nil
}

func (tdl *TabularDataLoader) WithChildren(ctx *sql.Context, children ...sql.Node) (sql.Node, error) {
	if len(children) != 0 {
		return nil, sql.ErrInvalidChildrenNumber.New(tdl, len(children), 0)
	}
	return tdl, nil
}

func (tdl *TabularDataLoader) IsReadOnly() bool {
	return true
}

type tabularRowIter struct {
	tdl    *TabularDataLoader
	reader *bufio.Reader
}

func (t tabularRowIter) Next(ctx *sql.Context) (sql.Row, error) {
	row, hasNext, err := t.tdl.nextRow(ctx, t.reader)
	if err != nil {
		return nil, err
	}

	// TODO: this isn't the best way to handle the count of rows, something like a RowUpdateAccumulator would be better
	if hasNext {
		t.tdl.results.RowsLoaded++
	} else {
		return nil, io.EOF
	}

	return row, nil
}

func (t tabularRowIter) Close(context *sql.Context) error {
	return nil
}

func (tdl *TabularDataLoader) RowIter(ctx *sql.Context, r sql.Row) (sql.RowIter, error) {
	return &tabularRowIter{tdl: tdl, reader: tdl.nextDataChunk}, nil
}