wehub-resource-sync
74 lines
3.0 KiB
Go
74 lines
3.0 KiB
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"
|
|
|
|
"github.com/cockroachdb/errors"
|
|
"github.com/dolthub/go-mysql-server/sql"
|
|
|
|
"github.com/dolthub/doltgresql/server/types"
|
|
)
|
|
|
|
// DataLoader allows callers to insert rows from multiple chunks into a table. Rows encoded in each chunk will not
|
|
// necessarily end cleanly on a chunk boundary, so DataLoader implementations must handle recognizing partial, or
|
|
// incomplete records, and saving that partial record until the next call to LoadChunk, so that it may be prefixed
|
|
// with the incomplete record.
|
|
type DataLoader interface {
|
|
sql.ExecSourceRel
|
|
|
|
// SetNextDataChunk sets the next data chunk to be processed by the DataLoader. Data records
|
|
// are not guaranteed to start and end cleanly on chunk boundaries, so implementations must recognize incomplete
|
|
// records and save them to prepend on the next processed chunk.
|
|
SetNextDataChunk(ctx *sql.Context, data *bufio.Reader) error
|
|
|
|
// Finish finalizes the current load operation and cleans up any resources used. Implementations should check that
|
|
// the last call to LoadChunk did not end with an incomplete record and return an error to the caller if so. The
|
|
// returned LoadDataResults describe the load operation, including how many rows were inserted.
|
|
Finish(ctx *sql.Context) (*LoadDataResults, error)
|
|
}
|
|
|
|
// LoadDataResults contains the results of a load data operation, including the number of rows loaded.
|
|
type LoadDataResults struct {
|
|
// RowsLoaded contains the total number of rows inserted during a load data operation.
|
|
RowsLoaded int32
|
|
}
|
|
|
|
// getColumnTypes returns the types of the columns in the schema that match the provided column names, in the order
|
|
// they are provided. If a subset of column names are provided, the returned types will only contain those columns.
|
|
// If the column names are not found in the schema, an error is returned.
|
|
func getColumnTypes(colNames []string, sch sql.Schema) ([]*types.DoltgresType, sql.Schema, error) {
|
|
colTypes := make([]*types.DoltgresType, len(colNames))
|
|
reducedSch := make(sql.Schema, len(colNames))
|
|
for i, colName := range colNames {
|
|
colIdx := sch.IndexOfColName(colName)
|
|
if colIdx < 0 {
|
|
// should be impossible
|
|
return nil, nil, errors.Errorf("column %s not found in schema", colName)
|
|
}
|
|
col := sch[colIdx]
|
|
var ok bool
|
|
colTypes[i], ok = col.Type.(*types.DoltgresType)
|
|
if !ok {
|
|
return nil, nil, errors.Errorf("unsupported column type: name: %s, type: %T", col.Name, col.Type)
|
|
}
|
|
|
|
reducedSch[i] = col
|
|
}
|
|
|
|
return colTypes, reducedSch, nil
|
|
}
|