chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,213 @@
|
||||
// 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"
|
||||
"bytes"
|
||||
"context"
|
||||
"io"
|
||||
"testing"
|
||||
|
||||
"github.com/dolthub/go-mysql-server/memory"
|
||||
"github.com/dolthub/go-mysql-server/sql"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/dolthub/doltgresql/core/dataloader"
|
||||
"github.com/dolthub/doltgresql/server/initialization"
|
||||
"github.com/dolthub/doltgresql/server/types"
|
||||
doltgresservercfg "github.com/dolthub/doltgresql/servercfg"
|
||||
)
|
||||
|
||||
// TestCsvDataLoader tests the CsvDataLoader implementation.
|
||||
func TestCsvDataLoader(t *testing.T) {
|
||||
db := memory.NewDatabase("mydb")
|
||||
provider := memory.NewDBProvider(db)
|
||||
initialization.Initialize(nil, doltgresservercfg.DefaultServerConfig())
|
||||
|
||||
ctx := &sql.Context{
|
||||
Context: context.Background(),
|
||||
Session: memory.NewSession(sql.NewBaseSession(), provider),
|
||||
}
|
||||
|
||||
pkCols := []string{"pk", "c1", "c2"}
|
||||
pkSchema := sql.NewPrimaryKeySchema(sql.Schema{
|
||||
{Name: "pk", Type: types.Int64, Source: "source1"},
|
||||
{Name: "c1", Type: types.Int64, Source: "source1"},
|
||||
{Name: "c2", Type: types.VarChar, Source: "source1"},
|
||||
}, 0)
|
||||
|
||||
// Tests that a basic CSV document can be loaded as a single chunk.
|
||||
t.Run("basic case", func(t *testing.T) {
|
||||
dataLoader, err := dataloader.NewCsvDataLoader(pkCols, pkSchema.Schema, ",", false)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Load all the data as a single chunk
|
||||
reader := bytes.NewReader([]byte("1,100,bar\n2,200,bash\n"))
|
||||
err = dataLoader.SetNextDataChunk(ctx, bufio.NewReader(reader))
|
||||
require.NoError(t, err)
|
||||
rows := loadAllRows(ctx, t, dataLoader)
|
||||
results, err := dataLoader.Finish(ctx)
|
||||
require.NoError(t, err)
|
||||
require.EqualValues(t, 2, results.RowsLoaded)
|
||||
|
||||
// Assert that the table contains the expected data
|
||||
assert.Equal(t, []sql.Row{
|
||||
{int64(1), int64(100), "bar"},
|
||||
{int64(2), int64(200), "bash"},
|
||||
}, rows)
|
||||
})
|
||||
|
||||
// Tests when a CSV record is split across two chunks of data, and the
|
||||
// partial record must be buffered and prepended to the next chunk.
|
||||
t.Run("record split across two chunks", func(t *testing.T) {
|
||||
dataLoader, err := dataloader.NewCsvDataLoader(pkCols, pkSchema.Schema, ",", false)
|
||||
require.NoError(t, err)
|
||||
|
||||
var rows []sql.Row
|
||||
|
||||
// Load the first chunk
|
||||
reader := bytes.NewReader([]byte("1,100,ba"))
|
||||
err = dataLoader.SetNextDataChunk(ctx, bufio.NewReader(reader))
|
||||
rows = append(rows, loadAllRows(ctx, t, dataLoader)...)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Load the second chunk
|
||||
reader = bytes.NewReader([]byte("r\n2,200,bash\n"))
|
||||
err = dataLoader.SetNextDataChunk(ctx, bufio.NewReader(reader))
|
||||
require.NoError(t, err)
|
||||
rows = append(rows, loadAllRows(ctx, t, dataLoader)...)
|
||||
|
||||
// Finish
|
||||
results, err := dataLoader.Finish(ctx)
|
||||
require.NoError(t, err)
|
||||
require.EqualValues(t, 2, results.RowsLoaded)
|
||||
|
||||
assert.Equal(t, []sql.Row{
|
||||
{int64(1), int64(100), "bar"},
|
||||
{int64(2), int64(200), "bash"},
|
||||
}, rows)
|
||||
})
|
||||
|
||||
// Tests when a CSV record is split across two chunks of data, and a
|
||||
// header row is present.
|
||||
t.Run("record split across two chunks, with header", func(t *testing.T) {
|
||||
dataLoader, err := dataloader.NewCsvDataLoader(pkCols, pkSchema.Schema, ",", true)
|
||||
require.NoError(t, err)
|
||||
|
||||
var rows []sql.Row
|
||||
|
||||
// Load the first chunk
|
||||
reader := bytes.NewReader([]byte("pk,c1,c2\n1,100,ba"))
|
||||
err = dataLoader.SetNextDataChunk(ctx, bufio.NewReader(reader))
|
||||
require.NoError(t, err)
|
||||
rows = append(rows, loadAllRows(ctx, t, dataLoader)...)
|
||||
|
||||
// Load the second chunk
|
||||
reader = bytes.NewReader([]byte("r\n2,200,bash\n"))
|
||||
err = dataLoader.SetNextDataChunk(ctx, bufio.NewReader(reader))
|
||||
require.NoError(t, err)
|
||||
rows = append(rows, loadAllRows(ctx, t, dataLoader)...)
|
||||
|
||||
// Finish
|
||||
results, err := dataLoader.Finish(ctx)
|
||||
require.NoError(t, err)
|
||||
require.EqualValues(t, 2, results.RowsLoaded)
|
||||
|
||||
assert.Equal(t, []sql.Row{
|
||||
{int64(1), int64(100), "bar"},
|
||||
{int64(2), int64(200), "bash"},
|
||||
}, rows)
|
||||
})
|
||||
|
||||
// Tests a CSV record that contains a quoted newline character and is split
|
||||
// across two chunks.
|
||||
t.Run("quoted newlines across two chunks", func(t *testing.T) {
|
||||
dataLoader, err := dataloader.NewCsvDataLoader(pkCols, pkSchema.Schema, ",", false)
|
||||
require.NoError(t, err)
|
||||
|
||||
var rows []sql.Row
|
||||
|
||||
// Load the first chunk
|
||||
reader := bytes.NewReader([]byte("1,100,\"baz\nbar\n"))
|
||||
err = dataLoader.SetNextDataChunk(ctx, bufio.NewReader(reader))
|
||||
require.NoError(t, err)
|
||||
rows = append(rows, loadAllRows(ctx, t, dataLoader)...)
|
||||
|
||||
// Load the second chunk
|
||||
reader = bytes.NewReader([]byte("bash\"\n2,200,bash\n"))
|
||||
err = dataLoader.SetNextDataChunk(ctx, bufio.NewReader(reader))
|
||||
require.NoError(t, err)
|
||||
rows = append(rows, loadAllRows(ctx, t, dataLoader)...)
|
||||
|
||||
// Finish
|
||||
results, err := dataLoader.Finish(ctx)
|
||||
require.NoError(t, err)
|
||||
require.EqualValues(t, 2, results.RowsLoaded)
|
||||
|
||||
assert.Equal(t, []sql.Row{
|
||||
{int64(1), int64(100), "baz\nbar\nbash"},
|
||||
{int64(2), int64(200), "bash"},
|
||||
}, rows)
|
||||
})
|
||||
|
||||
// Tests when a PSV (i.e. delimiter='|') record is split across two chunks of data,
|
||||
// and a header row is present.
|
||||
t.Run("delimiter='|', record split across two chunks, with header", func(t *testing.T) {
|
||||
dataLoader, err := dataloader.NewCsvDataLoader(pkCols, pkSchema.Schema, "|", true)
|
||||
require.NoError(t, err)
|
||||
|
||||
var rows []sql.Row
|
||||
|
||||
// Load the first chunk
|
||||
reader := bytes.NewReader([]byte("pk|c1|c2\n1|100|ba"))
|
||||
err = dataLoader.SetNextDataChunk(ctx, bufio.NewReader(reader))
|
||||
rows = append(rows, loadAllRows(ctx, t, dataLoader)...)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Load the second chunk
|
||||
reader = bytes.NewReader([]byte("r\n2|200|bash\n"))
|
||||
err = dataLoader.SetNextDataChunk(ctx, bufio.NewReader(reader))
|
||||
rows = append(rows, loadAllRows(ctx, t, dataLoader)...)
|
||||
require.NoError(t, err)
|
||||
|
||||
// Finish
|
||||
results, err := dataLoader.Finish(ctx)
|
||||
require.NoError(t, err)
|
||||
require.EqualValues(t, 2, results.RowsLoaded)
|
||||
|
||||
assert.Equal(t, []sql.Row{
|
||||
{int64(1), int64(100), "bar"},
|
||||
{int64(2), int64(200), "bash"},
|
||||
}, rows)
|
||||
})
|
||||
}
|
||||
|
||||
// loadAllRows loads all rows from the given DataLoader and returns them, failing the test if there are any errors.
|
||||
func loadAllRows(ctx *sql.Context, t *testing.T, d dataloader.DataLoader) []sql.Row {
|
||||
var rows []sql.Row
|
||||
iter, err := d.RowIter(ctx, nil)
|
||||
require.NoError(t, err)
|
||||
for {
|
||||
row, err := iter.Next(ctx)
|
||||
if err == io.EOF {
|
||||
break
|
||||
}
|
||||
require.NoError(t, err)
|
||||
rows = append(rows, row)
|
||||
}
|
||||
return rows
|
||||
}
|
||||
@@ -0,0 +1,236 @@
|
||||
// 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 (
|
||||
"bytes"
|
||||
"io"
|
||||
"testing"
|
||||
|
||||
"github.com/dolthub/doltgresql/core/dataloader"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// basicCsvData is used for a smoke test with CSV parsing
|
||||
const basicCsvData = `
|
||||
1,foo,bar
|
||||
2, ,bash
|
||||
`
|
||||
|
||||
// wrongNumberOfFieldsCsvData tests the case where records have different
|
||||
// numbers of values in them.
|
||||
const wrongNumberOfFieldsCsvData = `
|
||||
1,foo,bar,baz,bash
|
||||
2,blue
|
||||
3,boop,beep,bop,blorp
|
||||
`
|
||||
|
||||
// partialLineErrorCsvData tests the case where the last line of the CSV data
|
||||
// does not end with a newline character.
|
||||
const partialLineErrorCsvData = `
|
||||
1,foo,bar,baz,bash
|
||||
2,boop,beep,bop,blorp
|
||||
3,blue,g`
|
||||
|
||||
// nullAndEmptyStringQuotingCsvData tests the difference between representing
|
||||
// NULL and an empty string.
|
||||
const nullAndEmptyStringQuotingCsvData = `
|
||||
1,,NULL,"NULL",""
|
||||
`
|
||||
|
||||
// escapedQuotesCsvData tests escaped quotes in CSV data.
|
||||
const escapedQuotesCsvData = `
|
||||
1,'',"'","""",','''
|
||||
`
|
||||
|
||||
// newLineInQuotedFieldCsvData tests when a quoted field contains a newline.
|
||||
const newLineInQuotedFieldCsvData = `
|
||||
1,foo,"baz
|
||||
bar
|
||||
bash"
|
||||
`
|
||||
|
||||
// endOfDataMarkerCsvData tests when a quoted field contains the end of data marker.
|
||||
const endOfDataMarkerCsvData = `
|
||||
1,foo,"baz
|
||||
\.
|
||||
bash"
|
||||
`
|
||||
|
||||
// TestCsvReader tests various cases of CSV data parsing.
|
||||
func TestCsvReader(t *testing.T) {
|
||||
t.Run("basic CSV data", func(t *testing.T) {
|
||||
csvReader, err := dataloader.NewCsvReader(newReader(basicCsvData))
|
||||
require.NoError(t, err)
|
||||
|
||||
// Read the first row
|
||||
row, err := csvReader.ReadSqlRow()
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "1", row[0])
|
||||
assert.Equal(t, "foo", row[1])
|
||||
assert.Equal(t, "bar", row[2])
|
||||
|
||||
// Read the second row
|
||||
row, err = csvReader.ReadSqlRow()
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "2", row[0])
|
||||
assert.Equal(t, " ", row[1])
|
||||
assert.Equal(t, "bash", row[2])
|
||||
|
||||
// Read the EOF error
|
||||
_, err = csvReader.ReadSqlRow()
|
||||
require.Equal(t, io.EOF, err)
|
||||
})
|
||||
|
||||
t.Run("wrong number of fields", func(t *testing.T) {
|
||||
csvReader, err := dataloader.NewCsvReader(newReader(wrongNumberOfFieldsCsvData))
|
||||
require.NoError(t, err)
|
||||
|
||||
// Read the first row
|
||||
row, err := csvReader.ReadSqlRow()
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "1", row[0])
|
||||
require.Equal(t, "foo", row[1])
|
||||
require.Equal(t, "bar", row[2])
|
||||
require.Equal(t, "baz", row[3])
|
||||
require.Equal(t, "bash", row[4])
|
||||
|
||||
// Read the second row
|
||||
_, err = csvReader.ReadSqlRow()
|
||||
require.Error(t, err)
|
||||
require.Equal(t, "record on line 3: wrong number of fields", err.Error())
|
||||
})
|
||||
|
||||
t.Run("incomplete line, no newline ending", func(t *testing.T) {
|
||||
csvReader, err := dataloader.NewCsvReader(newReader(partialLineErrorCsvData))
|
||||
require.NoError(t, err)
|
||||
|
||||
// Read the first row
|
||||
row, err := csvReader.ReadSqlRow()
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "1", row[0])
|
||||
assert.Equal(t, "foo", row[1])
|
||||
assert.Equal(t, "bar", row[2])
|
||||
assert.Equal(t, "baz", row[3])
|
||||
assert.Equal(t, "bash", row[4])
|
||||
|
||||
// Read the second row
|
||||
row, err = csvReader.ReadSqlRow()
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "2", row[0])
|
||||
assert.Equal(t, "boop", row[1])
|
||||
assert.Equal(t, "beep", row[2])
|
||||
assert.Equal(t, "bop", row[3])
|
||||
assert.Equal(t, "blorp", row[4])
|
||||
|
||||
// Third row should trigger a partialLineError
|
||||
_, err = csvReader.ReadSqlRow()
|
||||
require.Error(t, err)
|
||||
require.Equal(t, "incomplete record found at end of CSV data: 3,blue,g", err.Error())
|
||||
})
|
||||
|
||||
t.Run("null and empty string quoting", func(t *testing.T) {
|
||||
csvReader, err := dataloader.NewCsvReader(newReader(nullAndEmptyStringQuotingCsvData))
|
||||
require.NoError(t, err)
|
||||
|
||||
// Read the first row
|
||||
row, err := csvReader.ReadSqlRow()
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, "1", row[0])
|
||||
assert.Equal(t, nil, row[1])
|
||||
assert.Equal(t, "NULL", row[2])
|
||||
assert.Equal(t, "NULL", row[3])
|
||||
assert.Equal(t, "", row[4])
|
||||
|
||||
// Read the EOF error
|
||||
_, err = csvReader.ReadSqlRow()
|
||||
require.Equal(t, io.EOF, err)
|
||||
})
|
||||
|
||||
t.Run("quote escaping", func(t *testing.T) {
|
||||
csvReader, err := dataloader.NewCsvReader(newReader(escapedQuotesCsvData))
|
||||
require.NoError(t, err)
|
||||
|
||||
// Read the first row
|
||||
row, err := csvReader.ReadSqlRow()
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, "1", row[0])
|
||||
assert.Equal(t, "''", row[1])
|
||||
assert.Equal(t, "'", row[2])
|
||||
assert.Equal(t, "\"", row[3])
|
||||
assert.Equal(t, "'", row[4])
|
||||
assert.Equal(t, "'''", row[5])
|
||||
|
||||
// Read the EOF error
|
||||
_, err = csvReader.ReadSqlRow()
|
||||
require.Equal(t, io.EOF, err)
|
||||
})
|
||||
|
||||
t.Run("quoted newlines", func(t *testing.T) {
|
||||
csvReader, err := dataloader.NewCsvReader(newReader(newLineInQuotedFieldCsvData))
|
||||
require.NoError(t, err)
|
||||
|
||||
// Read the first row
|
||||
row, err := csvReader.ReadSqlRow()
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, "1", row[0])
|
||||
assert.Equal(t, "foo", row[1])
|
||||
assert.Equal(t, "baz\nbar\nbash", row[2])
|
||||
|
||||
// Read the EOF error
|
||||
_, err = csvReader.ReadSqlRow()
|
||||
require.Equal(t, io.EOF, err)
|
||||
})
|
||||
|
||||
t.Run("quoted end of data marker", func(t *testing.T) {
|
||||
csvReader, err := dataloader.NewCsvReader(newReader(endOfDataMarkerCsvData))
|
||||
require.NoError(t, err)
|
||||
|
||||
// Read the first row
|
||||
row, err := csvReader.ReadSqlRow()
|
||||
assert.NoError(t, err)
|
||||
assert.Equal(t, "1", row[0])
|
||||
assert.Equal(t, "foo", row[1])
|
||||
assert.Equal(t, "baz\n\\.\nbash", row[2])
|
||||
|
||||
// Read the EOF error
|
||||
_, err = csvReader.ReadSqlRow()
|
||||
require.Equal(t, io.EOF, err)
|
||||
})
|
||||
|
||||
}
|
||||
|
||||
// testReader is a simple io.ReadCloser implementation that delegates reads to an io.Reader
|
||||
// and implements Close() as a no-op.
|
||||
type testReader struct {
|
||||
io.Reader
|
||||
}
|
||||
|
||||
var _ io.ReadCloser = (*testReader)(nil)
|
||||
|
||||
// Close implements the io.Closer interface
|
||||
func (frc *testReader) Close() error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// newReader returns an io.ReadCloser instance that reads the data from the specified
|
||||
// string |s| and is a no-op when Close() is called.
|
||||
func newReader(s string) io.ReadCloser {
|
||||
return &testReader{
|
||||
bytes.NewReader([]byte(s)),
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,107 @@
|
||||
// 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 (
|
||||
"bytes"
|
||||
"io"
|
||||
"testing"
|
||||
|
||||
"github.com/dolthub/doltgresql/core/dataloader"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestStringPrefixReader(t *testing.T) {
|
||||
t.Run("Read prefix and all data in single call", func(t *testing.T) {
|
||||
prefix := "prefix"
|
||||
reader := bytes.NewReader([]byte("0123456789"))
|
||||
prefixReader := dataloader.NewStringPrefixReader(prefix, reader)
|
||||
|
||||
data := make([]byte, 100)
|
||||
bytesRead, err := prefixReader.Read(data)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 16, bytesRead)
|
||||
require.Equal(t, "prefix0123456789", string(data[:bytesRead]))
|
||||
|
||||
bytesRead, err = prefixReader.Read(data)
|
||||
require.Equal(t, io.EOF, err)
|
||||
require.Equal(t, 0, bytesRead)
|
||||
})
|
||||
|
||||
t.Run("Read part of prefix", func(t *testing.T) {
|
||||
prefix := "prefix"
|
||||
reader := bytes.NewReader([]byte("0123456789"))
|
||||
prefixReader := dataloader.NewStringPrefixReader(prefix, reader)
|
||||
|
||||
data := make([]byte, 5)
|
||||
bytesRead, err := prefixReader.Read(data)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 5, bytesRead)
|
||||
require.Equal(t, "prefi", string(data[:bytesRead]))
|
||||
|
||||
// Read the next 5 bytes
|
||||
bytesRead, err = prefixReader.Read(data)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 5, bytesRead)
|
||||
require.Equal(t, "x0123", string(data[:bytesRead]))
|
||||
|
||||
// Read the next 5 bytes
|
||||
bytesRead, err = prefixReader.Read(data)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 5, bytesRead)
|
||||
require.Equal(t, "45678", string(data[:bytesRead]))
|
||||
|
||||
// Read the last byte
|
||||
bytesRead, err = prefixReader.Read(data)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 1, bytesRead)
|
||||
require.Equal(t, "9", string(data[:bytesRead]))
|
||||
|
||||
// Read EOF
|
||||
bytesRead, err = prefixReader.Read(data)
|
||||
require.Equal(t, io.EOF, err)
|
||||
require.Equal(t, 0, bytesRead)
|
||||
})
|
||||
|
||||
t.Run("Read to prefix boundary", func(t *testing.T) {
|
||||
prefix := "prefix"
|
||||
reader := bytes.NewReader([]byte("0123456789"))
|
||||
prefixReader := dataloader.NewStringPrefixReader(prefix, reader)
|
||||
|
||||
data := make([]byte, 6)
|
||||
bytesRead, err := prefixReader.Read(data)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 6, bytesRead)
|
||||
require.Equal(t, "prefix", string(data[:bytesRead]))
|
||||
|
||||
// Read the next 6 bytes
|
||||
bytesRead, err = prefixReader.Read(data)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 6, bytesRead)
|
||||
require.Equal(t, "012345", string(data[:bytesRead]))
|
||||
|
||||
// Read the next 6 bytes
|
||||
bytesRead, err = prefixReader.Read(data)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 4, bytesRead)
|
||||
require.Equal(t, "6789", string(data[:bytesRead]))
|
||||
|
||||
// Read EOF
|
||||
bytesRead, err = prefixReader.Read(data)
|
||||
require.Equal(t, io.EOF, err)
|
||||
require.Equal(t, 0, bytesRead)
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,200 @@
|
||||
// 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"
|
||||
"bytes"
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/dolthub/go-mysql-server/memory"
|
||||
"github.com/dolthub/go-mysql-server/sql"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/dolthub/doltgresql/core/dataloader"
|
||||
"github.com/dolthub/doltgresql/server/initialization"
|
||||
"github.com/dolthub/doltgresql/server/types"
|
||||
doltgresservercfg "github.com/dolthub/doltgresql/servercfg"
|
||||
)
|
||||
|
||||
func TestTabDataLoader(t *testing.T) {
|
||||
db := memory.NewDatabase("mydb")
|
||||
provider := memory.NewDBProvider(db)
|
||||
initialization.Initialize(nil, doltgresservercfg.DefaultServerConfig())
|
||||
|
||||
ctx := &sql.Context{
|
||||
Context: context.Background(),
|
||||
Session: memory.NewSession(sql.NewBaseSession(), provider),
|
||||
}
|
||||
|
||||
pkCols := []string{"pk", "c1", "c2"}
|
||||
pkSchema := sql.NewPrimaryKeySchema(sql.Schema{
|
||||
{Name: "pk", Type: types.Int64, Source: "source1"},
|
||||
{Name: "c1", Type: types.Int64, Source: "source1"},
|
||||
{Name: "c2", Type: types.VarChar, Source: "source1"},
|
||||
}, 0)
|
||||
|
||||
// Tests that a basic tab delimited doc can be loaded as a single chunk.
|
||||
t.Run("basic case", func(t *testing.T) {
|
||||
dataLoader, err := dataloader.NewTabularDataLoader(pkCols, pkSchema.Schema, "\t", "\\N", false)
|
||||
require.NoError(t, err)
|
||||
|
||||
var rows []sql.Row
|
||||
|
||||
// Load all the data as a single chunk
|
||||
reader := bytes.NewReader([]byte("1\t100\tbar\n2\t200\tbash\n"))
|
||||
err = dataLoader.SetNextDataChunk(ctx, bufio.NewReader(reader))
|
||||
require.NoError(t, err)
|
||||
rows = append(rows, loadAllRows(ctx, t, dataLoader)...)
|
||||
|
||||
results, err := dataLoader.Finish(ctx)
|
||||
require.NoError(t, err)
|
||||
require.EqualValues(t, 2, results.RowsLoaded)
|
||||
|
||||
assert.Equal(t, []sql.Row{
|
||||
{int64(1), int64(100), "bar"},
|
||||
{int64(2), int64(200), "bash"},
|
||||
}, rows)
|
||||
})
|
||||
|
||||
// Tests when a record is split across two chunks of data, and the
|
||||
// partial record must be buffered and prepended to the next chunk.
|
||||
t.Run("record split across two chunks", func(t *testing.T) {
|
||||
dataLoader, err := dataloader.NewTabularDataLoader(pkCols, pkSchema.Schema, "\t", "\\N", false)
|
||||
require.NoError(t, err)
|
||||
|
||||
var rows []sql.Row
|
||||
|
||||
// Load the first chunk
|
||||
reader := bytes.NewReader([]byte("1 100 ba"))
|
||||
err = dataLoader.SetNextDataChunk(ctx, bufio.NewReader(reader))
|
||||
require.NoError(t, err)
|
||||
rows = append(rows, loadAllRows(ctx, t, dataLoader)...)
|
||||
|
||||
// Load the second chunk
|
||||
reader = bytes.NewReader([]byte("r\n2 200 bash\n"))
|
||||
err = dataLoader.SetNextDataChunk(ctx, bufio.NewReader(reader))
|
||||
require.NoError(t, err)
|
||||
rows = append(rows, loadAllRows(ctx, t, dataLoader)...)
|
||||
|
||||
// Finish
|
||||
results, err := dataLoader.Finish(ctx)
|
||||
require.NoError(t, err)
|
||||
require.EqualValues(t, 2, results.RowsLoaded)
|
||||
|
||||
assert.Equal(t, []sql.Row{
|
||||
{int64(1), int64(100), "bar"},
|
||||
{int64(2), int64(200), "bash"},
|
||||
}, rows)
|
||||
})
|
||||
|
||||
// Tests when a record is split across two chunks of data, and a
|
||||
// header row is present.
|
||||
t.Run("record split across two chunks, with header", func(t *testing.T) {
|
||||
dataLoader, err := dataloader.NewTabularDataLoader(pkCols, pkSchema.Schema, "\t", "\\N", true)
|
||||
require.NoError(t, err)
|
||||
|
||||
var rows []sql.Row
|
||||
|
||||
// Load the first chunk
|
||||
reader := bytes.NewReader([]byte("pk c1 c2\n1 100 ba"))
|
||||
err = dataLoader.SetNextDataChunk(ctx, bufio.NewReader(reader))
|
||||
require.NoError(t, err)
|
||||
rows = append(rows, loadAllRows(ctx, t, dataLoader)...)
|
||||
|
||||
// Load the second chunk
|
||||
reader = bytes.NewReader([]byte("r\n2 200 bash\n"))
|
||||
err = dataLoader.SetNextDataChunk(ctx, bufio.NewReader(reader))
|
||||
require.NoError(t, err)
|
||||
rows = append(rows, loadAllRows(ctx, t, dataLoader)...)
|
||||
|
||||
// Finish
|
||||
results, err := dataLoader.Finish(ctx)
|
||||
require.NoError(t, err)
|
||||
require.EqualValues(t, 2, results.RowsLoaded)
|
||||
|
||||
// Assert that the table contains the expected data
|
||||
assert.Equal(t, []sql.Row{
|
||||
{int64(1), int64(100), "bar"},
|
||||
{int64(2), int64(200), "bash"},
|
||||
}, rows)
|
||||
})
|
||||
|
||||
// Tests a record that contains a quoted newline character and is split
|
||||
// across two chunks.
|
||||
t.Run("quoted newlines across two chunks", func(t *testing.T) {
|
||||
dataLoader, err := dataloader.NewTabularDataLoader(pkCols, pkSchema.Schema, "\t", "\\N", false)
|
||||
require.NoError(t, err)
|
||||
|
||||
var rows []sql.Row
|
||||
|
||||
// Load the first chunk
|
||||
reader := bytes.NewReader([]byte("1 100 \"baz\\nbar\\n"))
|
||||
err = dataLoader.SetNextDataChunk(ctx, bufio.NewReader(reader))
|
||||
require.NoError(t, err)
|
||||
rows = append(rows, loadAllRows(ctx, t, dataLoader)...)
|
||||
|
||||
// Load the second chunk
|
||||
reader = bytes.NewReader([]byte("bash\"\n2 200 bash\n"))
|
||||
err = dataLoader.SetNextDataChunk(ctx, bufio.NewReader(reader))
|
||||
require.NoError(t, err)
|
||||
rows = append(rows, loadAllRows(ctx, t, dataLoader)...)
|
||||
|
||||
// Finish
|
||||
results, err := dataLoader.Finish(ctx)
|
||||
require.NoError(t, err)
|
||||
require.EqualValues(t, 2, results.RowsLoaded)
|
||||
|
||||
// Assert that the table contains the expected data
|
||||
assert.Equal(t, []sql.Row{
|
||||
{int64(1), int64(100), "\"baz\\nbar\\nbash\""},
|
||||
{int64(2), int64(200), "bash"},
|
||||
}, rows)
|
||||
})
|
||||
|
||||
// Tests when a record is split across two chunks of data, and a
|
||||
// header row is present.
|
||||
t.Run("delimiter='|', record split across two chunks, with header", func(t *testing.T) {
|
||||
dataLoader, err := dataloader.NewTabularDataLoader(pkCols, pkSchema.Schema, "|", "\\N", true)
|
||||
require.NoError(t, err)
|
||||
|
||||
var rows []sql.Row
|
||||
|
||||
// Load the first chunk
|
||||
reader := bytes.NewReader([]byte("pk|c1|c2\n1|100|ba"))
|
||||
err = dataLoader.SetNextDataChunk(ctx, bufio.NewReader(reader))
|
||||
require.NoError(t, err)
|
||||
rows = append(rows, loadAllRows(ctx, t, dataLoader)...)
|
||||
|
||||
// Load the second chunk
|
||||
reader = bytes.NewReader([]byte("r\n2|200|bash\n"))
|
||||
err = dataLoader.SetNextDataChunk(ctx, bufio.NewReader(reader))
|
||||
require.NoError(t, err)
|
||||
rows = append(rows, loadAllRows(ctx, t, dataLoader)...)
|
||||
|
||||
// Finish
|
||||
results, err := dataLoader.Finish(ctx)
|
||||
require.NoError(t, err)
|
||||
require.EqualValues(t, 2, results.RowsLoaded)
|
||||
|
||||
// Assert that the table contains the expected data
|
||||
assert.Equal(t, []sql.Row{
|
||||
{int64(1), int64(100), "bar"},
|
||||
{int64(2), int64(200), "bash"},
|
||||
}, rows)
|
||||
})
|
||||
}
|
||||
Reference in New Issue
Block a user