github.com/treeverse/lakefs@v1.24.1-0.20240520134607-95648127bfb0/pkg/testutil/dynamodb.go (about)

     1  package testutil
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"net/http"
     7  	"testing"
     8  	"time"
     9  
    10  	nanoid "github.com/matoous/go-nanoid/v2"
    11  	"github.com/ory/dockertest/v3"
    12  	dc "github.com/ory/dockertest/v3/docker"
    13  	"github.com/treeverse/lakefs/pkg/kv"
    14  	"github.com/treeverse/lakefs/pkg/kv/dynamodb"
    15  	"github.com/treeverse/lakefs/pkg/kv/kvparams"
    16  )
    17  
    18  const (
    19  	dbContainerTimeoutSeconds = 10 * 60 // 10 min
    20  	DynamodbLocalPort         = "6432"
    21  	DynamodbLocalURI          = "http://localhost:6432"
    22  	chars                     = "abcdef1234567890"
    23  	charsSize                 = 8
    24  	DynamoDBScanLimit         = 10
    25  )
    26  
    27  func GetDynamoDBInstance() (string, func(), error) {
    28  	dockerPool, err := dockertest.NewPool("")
    29  	if err != nil {
    30  		return "", nil, fmt.Errorf("could not connect to Docker: %w", err)
    31  	}
    32  	dockerPool.MaxWait = dbContainerTimeoutSeconds * time.Second
    33  
    34  	dynamodbDockerRunOptions := &dockertest.RunOptions{
    35  		Repository: "amazon/dynamodb-local",
    36  		Tag:        "latest",
    37  		PortBindings: map[dc.Port][]dc.PortBinding{
    38  			"8000/tcp": {{HostPort: DynamodbLocalPort}},
    39  		},
    40  	}
    41  
    42  	resource, err := dockerPool.RunWithOptions(dynamodbDockerRunOptions)
    43  	if err != nil {
    44  		return "", nil, fmt.Errorf("could not start dynamodb local: %w", err)
    45  	}
    46  
    47  	// set cleanup
    48  	closer := func() {
    49  		err = dockerPool.Purge(resource)
    50  		if err != nil {
    51  			panic("could not kill dynamodb local container")
    52  		}
    53  	}
    54  
    55  	// expire, just to make sure
    56  	err = resource.Expire(dbContainerTimeoutSeconds)
    57  	if err != nil {
    58  		return "", nil, fmt.Errorf("could not expire dynamodb local container: %w", err)
    59  	}
    60  
    61  	err = dockerPool.Retry(func() error {
    62  		// Waiting for dynamodb container to be ready by issuing an HTTP get request with
    63  		// exponential backoff retry.
    64  		// The response is not really meaningful for that case and so is ignored.
    65  		resp, err := http.Get(DynamodbLocalURI)
    66  		if err != nil {
    67  			return err
    68  		}
    69  		_ = resp.Body.Close()
    70  		return nil
    71  	})
    72  	if err != nil {
    73  		return "", nil, fmt.Errorf("could not connect to dynamodb at %s: %w", DynamodbLocalURI, err)
    74  	}
    75  
    76  	// return DB URI
    77  	return DynamodbLocalURI, closer, nil
    78  }
    79  
    80  func UniqueKVTableName() string {
    81  	return "kvstore_" + UniqueName()
    82  }
    83  
    84  func UniqueName() string {
    85  	return nanoid.MustGenerate(chars, charsSize)
    86  }
    87  
    88  func GetDynamoDBProd(ctx context.Context, tb testing.TB) kv.Store {
    89  	table := UniqueKVTableName()
    90  	testParams := &kvparams.DynamoDB{
    91  		TableName: table,
    92  		ScanLimit: DynamoDBScanLimit,
    93  		AwsRegion: "us-east-1",
    94  	}
    95  
    96  	store, err := kv.Open(ctx, kvparams.Config{Type: dynamodb.DriverName, DynamoDB: testParams})
    97  	if err != nil {
    98  		tb.Fatalf("failed to open kv dynamodb store %s", err)
    99  	}
   100  	tb.Cleanup(func() {
   101  		defer store.Close()
   102  		err = store.(*dynamodb.Store).DropTable()
   103  		if err != nil {
   104  			tb.Fatalf("failed to delete table from DB %s %s", table, err)
   105  		}
   106  	})
   107  	return store
   108  }