github.com/hasnat/dolt/go@v0.0.0-20210628190320-9eb5d843fbb7/store/nbs/dynamo_table_reader_test.go (about)

     1  // Copyright 2019 Dolthub, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  //
    15  // This file incorporates work covered by the following copyright and
    16  // permission notice:
    17  //
    18  // Copyright 2017 Attic Labs, Inc. All rights reserved.
    19  // Licensed under the Apache License, version 2.0:
    20  // http://www.apache.org/licenses/LICENSE-2.0
    21  
    22  package nbs
    23  
    24  import (
    25  	"context"
    26  	"testing"
    27  
    28  	"github.com/aws/aws-sdk-go/aws"
    29  	"github.com/aws/aws-sdk-go/aws/request"
    30  	"github.com/aws/aws-sdk-go/service/dynamodb"
    31  	"github.com/stretchr/testify/assert"
    32  	"github.com/stretchr/testify/require"
    33  
    34  	"github.com/dolthub/dolt/go/store/util/sizecache"
    35  )
    36  
    37  func TestDynamoTableReaderAt(t *testing.T) {
    38  	ddb := makeFakeDDB(t)
    39  
    40  	chunks := [][]byte{
    41  		[]byte("hello2"),
    42  		[]byte("goodbye2"),
    43  		[]byte("badbye2"),
    44  	}
    45  
    46  	tableData, h, err := buildTable(chunks)
    47  	require.NoError(t, err)
    48  	ddb.putData(fmtTableName(h), tableData)
    49  
    50  	t.Run("ddbTableStore", func(t *testing.T) {
    51  		t.Run("ReadTable", func(t *testing.T) {
    52  			test := func(dts *ddbTableStore) {
    53  				assert := assert.New(t)
    54  				data, err := dts.ReadTable(context.Background(), h, &Stats{})
    55  				require.NoError(t, err)
    56  				assert.Equal(tableData, data)
    57  
    58  				data, err = dts.ReadTable(context.Background(), computeAddr([]byte{}), &Stats{})
    59  				assert.Error(err)
    60  				assert.IsType(tableNotInDynamoErr{}, err)
    61  				assert.Nil(data)
    62  			}
    63  
    64  			t.Run("EventuallyConsistentSuccess", func(t *testing.T) {
    65  				test(&ddbTableStore{ddb, "table", nil, nil})
    66  			})
    67  
    68  			t.Run("EventuallyConsistentFailure", func(t *testing.T) {
    69  				test(&ddbTableStore{&eventuallyConsistentDDB{ddb}, "table", nil, nil})
    70  			})
    71  
    72  			t.Run("WithCache", func(t *testing.T) {
    73  				tc := sizecache.New(uint64(2 * len(tableData)))
    74  				dts := &ddbTableStore{ddb, "table", nil, tc}
    75  				test(dts)
    76  
    77  				// Table should have been cached on read
    78  				baseline := ddb.NumGets()
    79  				_, err := dts.ReadTable(context.Background(), h, &Stats{})
    80  				require.NoError(t, err)
    81  				assert.Zero(t, ddb.NumGets()-baseline)
    82  			})
    83  		})
    84  
    85  		t.Run("WriteTable", func(t *testing.T) {
    86  			t.Run("WithoutCache", func(t *testing.T) {
    87  				assert := assert.New(t)
    88  
    89  				dts := &ddbTableStore{makeFakeDDB(t), "table", nil, nil}
    90  				require.NoError(t, dts.Write(context.Background(), h, tableData))
    91  
    92  				data, err := dts.ReadTable(context.Background(), h, &Stats{})
    93  				require.NoError(t, err)
    94  				assert.Equal(tableData, data)
    95  			})
    96  
    97  			t.Run("WithCache", func(t *testing.T) {
    98  				assert := assert.New(t)
    99  
   100  				tc := sizecache.New(uint64(2 * len(tableData)))
   101  				dts := &ddbTableStore{makeFakeDDB(t), "table", nil, tc}
   102  				require.NoError(t, dts.Write(context.Background(), h, tableData))
   103  
   104  				// Table should have been cached on write
   105  				baseline := ddb.NumGets()
   106  				data, err := dts.ReadTable(context.Background(), h, &Stats{})
   107  				require.NoError(t, err)
   108  				assert.Equal(tableData, data)
   109  				assert.Zero(ddb.NumGets() - baseline)
   110  			})
   111  		})
   112  	})
   113  
   114  	t.Run("ReadAtWithCache", func(t *testing.T) {
   115  		assert := assert.New(t)
   116  		stats := &Stats{}
   117  
   118  		tc := sizecache.New(uint64(2 * len(tableData)))
   119  		tra := &dynamoTableReaderAt{&ddbTableStore{ddb, "table", nil, tc}, h}
   120  
   121  		// First, read when table is not yet cached
   122  		scratch := make([]byte, len(tableData)/4)
   123  		baseline := ddb.NumGets()
   124  		_, err := tra.ReadAtWithStats(context.Background(), scratch, 0, stats)
   125  		require.NoError(t, err)
   126  		assert.True(ddb.NumGets() > baseline)
   127  
   128  		// Table should have been cached on read so read again, a different slice this time
   129  		baseline = ddb.NumGets()
   130  		_, err = tra.ReadAtWithStats(context.Background(), scratch, int64(len(scratch)), stats)
   131  		require.NoError(t, err)
   132  		assert.Zero(ddb.NumGets() - baseline)
   133  	})
   134  }
   135  
   136  type eventuallyConsistentDDB struct {
   137  	ddbsvc
   138  }
   139  
   140  func (ec *eventuallyConsistentDDB) GetItemWithContext(ctx aws.Context, input *dynamodb.GetItemInput, opts ...request.Option) (*dynamodb.GetItemOutput, error) {
   141  	if input.ConsistentRead != nil && *(input.ConsistentRead) {
   142  		return ec.ddbsvc.GetItemWithContext(ctx, input)
   143  	}
   144  	return &dynamodb.GetItemOutput{}, nil
   145  }