TestCsvDataLoader tests the CsvDataLoader implementation.
(t *testing.T)
| 34 | |
| 35 | // TestCsvDataLoader tests the CsvDataLoader implementation. |
| 36 | func TestCsvDataLoader(t *testing.T) { |
| 37 | db := memory.NewDatabase("mydb") |
| 38 | provider := memory.NewDBProvider(db) |
| 39 | initialization.Initialize(nil, doltgresservercfg.DefaultServerConfig()) |
| 40 | |
| 41 | ctx := &sql.Context{ |
| 42 | Context: context.Background(), |
| 43 | Session: memory.NewSession(sql.NewBaseSession(), provider), |
| 44 | } |
| 45 | |
| 46 | pkCols := []string{"pk", "c1", "c2"} |
| 47 | pkSchema := sql.NewPrimaryKeySchema(sql.Schema{ |
| 48 | {Name: "pk", Type: types.Int64, Source: "source1"}, |
| 49 | {Name: "c1", Type: types.Int64, Source: "source1"}, |
| 50 | {Name: "c2", Type: types.VarChar, Source: "source1"}, |
| 51 | }, 0) |
| 52 | |
| 53 | // Tests that a basic CSV document can be loaded as a single chunk. |
| 54 | t.Run("basic case", func(t *testing.T) { |
| 55 | dataLoader, err := dataloader.NewCsvDataLoader(pkCols, pkSchema.Schema, ",", false) |
| 56 | require.NoError(t, err) |
| 57 | |
| 58 | // Load all the data as a single chunk |
| 59 | reader := bytes.NewReader([]byte("1,100,bar\n2,200,bash\n")) |
| 60 | err = dataLoader.SetNextDataChunk(ctx, bufio.NewReader(reader)) |
| 61 | require.NoError(t, err) |
| 62 | rows := loadAllRows(ctx, t, dataLoader) |
| 63 | results, err := dataLoader.Finish(ctx) |
| 64 | require.NoError(t, err) |
| 65 | require.EqualValues(t, 2, results.RowsLoaded) |
| 66 | |
| 67 | // Assert that the table contains the expected data |
| 68 | assert.Equal(t, []sql.Row{ |
| 69 | {int64(1), int64(100), "bar"}, |
| 70 | {int64(2), int64(200), "bash"}, |
| 71 | }, rows) |
| 72 | }) |
| 73 | |
| 74 | // Tests when a CSV record is split across two chunks of data, and the |
| 75 | // partial record must be buffered and prepended to the next chunk. |
| 76 | t.Run("record split across two chunks", func(t *testing.T) { |
| 77 | dataLoader, err := dataloader.NewCsvDataLoader(pkCols, pkSchema.Schema, ",", false) |
| 78 | require.NoError(t, err) |
| 79 | |
| 80 | var rows []sql.Row |
| 81 | |
| 82 | // Load the first chunk |
| 83 | reader := bytes.NewReader([]byte("1,100,ba")) |
| 84 | err = dataLoader.SetNextDataChunk(ctx, bufio.NewReader(reader)) |
| 85 | rows = append(rows, loadAllRows(ctx, t, dataLoader)...) |
| 86 | require.NoError(t, err) |
| 87 | |
| 88 | // Load the second chunk |
| 89 | reader = bytes.NewReader([]byte("r\n2,200,bash\n")) |
| 90 | err = dataLoader.SetNextDataChunk(ctx, bufio.NewReader(reader)) |
| 91 | require.NoError(t, err) |
| 92 | rows = append(rows, loadAllRows(ctx, t, dataLoader)...) |
| 93 |
nothing calls this directly
no test coverage detected