(t *testing.T)
| 32 | ) |
| 33 | |
| 34 | func TestTabDataLoader(t *testing.T) { |
| 35 | db := memory.NewDatabase("mydb") |
| 36 | provider := memory.NewDBProvider(db) |
| 37 | initialization.Initialize(nil, doltgresservercfg.DefaultServerConfig()) |
| 38 | |
| 39 | ctx := &sql.Context{ |
| 40 | Context: context.Background(), |
| 41 | Session: memory.NewSession(sql.NewBaseSession(), provider), |
| 42 | } |
| 43 | |
| 44 | pkCols := []string{"pk", "c1", "c2"} |
| 45 | pkSchema := sql.NewPrimaryKeySchema(sql.Schema{ |
| 46 | {Name: "pk", Type: types.Int64, Source: "source1"}, |
| 47 | {Name: "c1", Type: types.Int64, Source: "source1"}, |
| 48 | {Name: "c2", Type: types.VarChar, Source: "source1"}, |
| 49 | }, 0) |
| 50 | |
| 51 | // Tests that a basic tab delimited doc can be loaded as a single chunk. |
| 52 | t.Run("basic case", func(t *testing.T) { |
| 53 | dataLoader, err := dataloader.NewTabularDataLoader(pkCols, pkSchema.Schema, "\t", "\\N", false) |
| 54 | require.NoError(t, err) |
| 55 | |
| 56 | var rows []sql.Row |
| 57 | |
| 58 | // Load all the data as a single chunk |
| 59 | reader := bytes.NewReader([]byte("1\t100\tbar\n2\t200\tbash\n")) |
| 60 | err = dataLoader.SetNextDataChunk(ctx, bufio.NewReader(reader)) |
| 61 | require.NoError(t, err) |
| 62 | rows = append(rows, loadAllRows(ctx, t, dataLoader)...) |
| 63 | |
| 64 | results, err := dataLoader.Finish(ctx) |
| 65 | require.NoError(t, err) |
| 66 | require.EqualValues(t, 2, results.RowsLoaded) |
| 67 | |
| 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 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.NewTabularDataLoader(pkCols, pkSchema.Schema, "\t", "\\N", 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 | require.NoError(t, err) |
| 86 | rows = append(rows, loadAllRows(ctx, t, dataLoader)...) |
| 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) |
nothing calls this directly
no test coverage detected