github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/workload/csv_test.go (about)

     1  // Copyright 2018 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package workload_test
    12  
    13  import (
    14  	"bytes"
    15  	"context"
    16  	"io"
    17  	"io/ioutil"
    18  	"net/http"
    19  	"net/http/httptest"
    20  	"strings"
    21  	"testing"
    22  
    23  	"github.com/cockroachdb/cockroach/pkg/col/coldata"
    24  	"github.com/cockroachdb/cockroach/pkg/util/bufalloc"
    25  	"github.com/cockroachdb/cockroach/pkg/util/httputil"
    26  	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
    27  	"github.com/cockroachdb/cockroach/pkg/workload"
    28  	"github.com/cockroachdb/cockroach/pkg/workload/bank"
    29  	"github.com/cockroachdb/cockroach/pkg/workload/tpcc"
    30  	"github.com/stretchr/testify/require"
    31  )
    32  
    33  func TestHandleCSV(t *testing.T) {
    34  	defer leaktest.AfterTest(t)()
    35  
    36  	tests := []struct {
    37  		params, expected string
    38  	}{
    39  		{
    40  			`?rows=1`, `
    41  0,0,initial-dTqnRurXztAPkykhZWvsCmeJkMwRNcJAvTlNbgUEYfagEQJaHmfPsquKZUBOGwpAjPtATpGXFJkrtQCEJODSlmQctvyh`,
    42  		},
    43  		{
    44  			`?rows=5&row-start=1&row-end=3&batch-size=1`, `
    45  1,0,initial-vOpikzTTWxvMqnkpfEIVXgGyhZNDqvpVqpNnHawruAcIVltgbnIEIGmCDJcnkVkfVmAcutkMvRACFuUBPsZTemTDSfZT
    46  2,0,initial-qMvoPeRiOBXvdVQxhZUfdmehETKPXyBaVWxzMqwiStIkxfoDFygYxIDyXiaVEarcwMboFhBlCAapvKijKAyjEAhRBNZz`,
    47  		},
    48  	}
    49  
    50  	meta := bank.FromRows(0).Meta()
    51  	for _, test := range tests {
    52  		t.Run(test.params, func(t *testing.T) {
    53  			ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    54  				if err := workload.HandleCSV(w, r, `/bank/`, meta); err != nil {
    55  					panic(err)
    56  				}
    57  			}))
    58  			defer ts.Close()
    59  
    60  			res, err := httputil.Get(context.Background(), ts.URL+`/bank/bank`+test.params)
    61  			if err != nil {
    62  				t.Fatal(err)
    63  			}
    64  			data, err := ioutil.ReadAll(res.Body)
    65  			res.Body.Close()
    66  			if err != nil {
    67  				t.Fatal(err)
    68  			}
    69  
    70  			if d, e := strings.TrimSpace(string(data)), strings.TrimSpace(test.expected); d != e {
    71  				t.Errorf("got [\n%s\n] expected [\n%s\n]", d, e)
    72  			}
    73  		})
    74  	}
    75  }
    76  
    77  func BenchmarkWriteCSVRows(b *testing.B) {
    78  	ctx := context.Background()
    79  
    80  	var batches []coldata.Batch
    81  	for _, table := range tpcc.FromWarehouses(1).Tables() {
    82  		cb := coldata.NewMemBatch(nil /* types */, coldata.StandardColumnFactory)
    83  		var a bufalloc.ByteAllocator
    84  		table.InitialRows.FillBatch(0, cb, &a)
    85  		batches = append(batches, cb)
    86  	}
    87  	table := workload.Table{
    88  		InitialRows: workload.BatchedTuples{
    89  			FillBatch: func(batchIdx int, cb coldata.Batch, _ *bufalloc.ByteAllocator) {
    90  				*cb.(*coldata.MemBatch) = *batches[batchIdx].(*coldata.MemBatch)
    91  			},
    92  		},
    93  	}
    94  
    95  	var buf bytes.Buffer
    96  	fn := func() {
    97  		const limit = -1
    98  		if _, err := workload.WriteCSVRows(ctx, &buf, table, 0, len(batches), limit); err != nil {
    99  			b.Fatalf(`%+v`, err)
   100  		}
   101  	}
   102  
   103  	// Run fn once to pre-size buf.
   104  	fn()
   105  
   106  	b.ResetTimer()
   107  	for i := 0; i < b.N; i++ {
   108  		buf.Reset()
   109  		fn()
   110  	}
   111  	b.StopTimer()
   112  	b.SetBytes(int64(buf.Len()))
   113  }
   114  
   115  func TestCSVRowsReader(t *testing.T) {
   116  	defer leaktest.AfterTest(t)()
   117  
   118  	table := bank.FromRows(10).Tables()[0]
   119  	r := workload.NewCSVRowsReader(table, 1, 3)
   120  	b, err := ioutil.ReadAll(r)
   121  	require.NoError(t, err)
   122  	expected := `
   123  1,0,initial-vOpikzTTWxvMqnkpfEIVXgGyhZNDqvpVqpNnHawruAcIVltgbnIEIGmCDJcnkVkfVmAcutkMvRACFuUBPsZTemTDSfZT
   124  2,0,initial-qMvoPeRiOBXvdVQxhZUfdmehETKPXyBaVWxzMqwiStIkxfoDFygYxIDyXiaVEarcwMboFhBlCAapvKijKAyjEAhRBNZz
   125  `
   126  	require.Equal(t, strings.TrimSpace(expected), strings.TrimSpace(string(b)))
   127  }
   128  
   129  func BenchmarkCSVRowsReader(b *testing.B) {
   130  	var batches []coldata.Batch
   131  	for _, table := range tpcc.FromWarehouses(1).Tables() {
   132  		cb := coldata.NewMemBatch(nil /* types */, coldata.StandardColumnFactory)
   133  		var a bufalloc.ByteAllocator
   134  		table.InitialRows.FillBatch(0, cb, &a)
   135  		batches = append(batches, cb)
   136  	}
   137  	table := workload.Table{
   138  		InitialRows: workload.BatchedTuples{
   139  			NumBatches: len(batches),
   140  			FillBatch: func(batchIdx int, cb coldata.Batch, _ *bufalloc.ByteAllocator) {
   141  				*cb.(*coldata.MemBatch) = *batches[batchIdx].(*coldata.MemBatch)
   142  			},
   143  		},
   144  	}
   145  
   146  	var buf bytes.Buffer
   147  	fn := func() {
   148  		r := workload.NewCSVRowsReader(table, 0, 0)
   149  		_, err := io.Copy(&buf, r)
   150  		require.NoError(b, err)
   151  	}
   152  
   153  	// Run fn once to pre-size buf.
   154  	fn()
   155  
   156  	b.ResetTimer()
   157  	for i := 0; i < b.N; i++ {
   158  		buf.Reset()
   159  		fn()
   160  	}
   161  	b.StopTimer()
   162  	b.SetBytes(int64(buf.Len()))
   163  }