Support multiple tables in BatchWrite and BatchGet (#226)

* batch write: support writing to multiple tables in a single batch

* batch get: support reading from multiple tables

* add Merge to batches

* fix flake in TestScanPaging

* add github actions ci & auto-create test table

* use unique table name in ci

* add multi-table batch tests

.github/docker-compose.yml

@@ -0,0 +1,8 @@
+version: '3'
+
+services:
+  dynamodb:
+    image: amazon/dynamodb-local:latest
+    ports:
+      - "8880:8000"
+    command: "-jar DynamoDBLocal.jar -sharedDb -inMemory"

.github/workflows/ci.yml

@@ -0,0 +1,25 @@
+name: CI
+
+on: [push, pull_request]
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+      - name: Install Go
+        uses: actions/setup-go@v5
+        with:
+          go-version: 'stable'
+      - name: Start DynamoDB Local
+        run: docker compose -f '.github/docker-compose.yml' up -d
+      - name: Test
+        run: go test -v -race -cover -coverpkg=./... ./...
+        env:
+          DYNAMO_TEST_ENDPOINT: 'http://localhost:8880'
+          DYNAMO_TEST_REGION: local
+          DYNAMO_TEST_TABLE: 'TestDB-%'
+          AWS_ACCESS_KEY_ID: dummy
+          AWS_SECRET_ACCESS_KEY: dummy
+          AWS_REGION: local

README.md

@@ -232,38 +232,23 @@ err := db.Table("Books").Get("ID", 555).One(dynamo.AWSEncoding(&someBook))
 
 ### Integration tests
 
-By default, tests are run in offline mode. Create a table called `TestDB`, with a number partition key called `UserID` and a string sort key called `Time`. It also needs a Global Secondary Index called `Msg-Time-index` with a string partition key called `Msg` and a string sort key called `Time`.
+By default, tests are run in offline mode. In order to run the integration tests, some environment variables need to be set.
 
-Change the table name with the environment variable `DYNAMO_TEST_TABLE`. You must specify `DYNAMO_TEST_REGION`, setting it to the AWS region where your test table is.
-
-
- ```bash
-DYNAMO_TEST_REGION=us-west-2 go test github.com/guregu/dynamo/... -cover
- ```
-
-If you want to use [DynamoDB Local](https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/DynamoDBLocal.html) to run local tests, specify `DYNAMO_TEST_ENDPOINT`.
-
- ```bash
-DYNAMO_TEST_REGION=us-west-2 DYNAMO_TEST_ENDPOINT=http://localhost:8000 go test github.com/guregu/dynamo/... -cover
- ```
-
-Example of using [aws-cli](https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/Tools.CLI.html) to create a table for testing.
+To run the tests against [DynamoDB Local](https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/DynamoDBLocal.html):
 
 ```bash
-aws dynamodb create-table \
-    --table-name TestDB \
-    --attribute-definitions \
-        AttributeName=UserID,AttributeType=N \
-        AttributeName=Time,AttributeType=S \
-        AttributeName=Msg,AttributeType=S \
-    --key-schema \
-        AttributeName=UserID,KeyType=HASH \
-        AttributeName=Time,KeyType=RANGE \
-    --global-secondary-indexes \
-        IndexName=Msg-Time-index,KeySchema=[{'AttributeName=Msg,KeyType=HASH'},{'AttributeName=Time,KeyType=RANGE'}],Projection={'ProjectionType=ALL'} \
-    --billing-mode PAY_PER_REQUEST \
-    --region us-west-2 \
-    --endpoint-url http://localhost:8000 # using DynamoDB local
+# Use Docker to run DynamoDB local on port 8880
+docker compose -f '.github/docker-compose.yml' up -d
+
+# Run the tests with a fresh table
+# The tables will be created automatically
+DYNAMO_TEST_ENDPOINT='http://localhost:8880' \
+	DYNAMO_TEST_REGION='local' \
+	DYNAMO_TEST_TABLE='TestDB-%' \ # the % will be replaced the current timestamp
+	AWS_ACCESS_KEY_ID='dummy' \
+	AWS_SECRET_ACCESS_KEY='dummy' \
+	AWS_REGION='local' \
+	go test -v -race ./... -cover -coverpkg=./...
 ```
 
 ### License

batch_test.go

@@ -11,7 +11,10 @@ func TestBatchGetWrite(t *testing.T) {
 	if testDB == nil {
 		t.Skip(offlineSkipMsg)
 	}
-	table := testDB.Table(testTable)
+	table1 := testDB.Table(testTableWidgets)
+	table2 := testDB.Table(testTableSprockets)
+	tables := []Table{table1, table2}
+	totalBatchSize := batchSize * len(tables)
 
 	items := make([]interface{}, batchSize)
 	widgets := make(map[int]widget)
@@ -28,10 +31,19 @@ func TestBatchGetWrite(t *testing.T) {
 		keys[i] = Keys{i, now}
 	}
 
+	var batches []*BatchWrite
+	for _, table := range tables {
+		b := table.Batch().Write().Put(items...)
+		batches = append(batches, b)
+	}
+	batch1 := batches[0]
+	for _, b := range batches[1:] {
+		batch1.Merge(b)
+	}
 	var wcc ConsumedCapacity
-	wrote, err := table.Batch().Write().Put(items...).ConsumedCapacity(&wcc).Run()
-	if wrote != batchSize {
-		t.Error("unexpected wrote:", wrote, "≠", batchSize)
+	wrote, err := batch1.ConsumedCapacity(&wcc).Run()
+	if wrote != totalBatchSize {
+		t.Error("unexpected wrote:", wrote, "≠", totalBatchSize)
 	}
 	if err != nil {
 		t.Error("unexpected error:", err)
@@ -41,20 +53,29 @@ func TestBatchGetWrite(t *testing.T) {
 	}
 
 	// get all
-	var results []widget
+	var gets []*BatchGet
+	for _, table := range tables {
+		b := table.Batch("UserID", "Time").
+			Get(keys...).
+			Project("UserID", "Time").
+			Consistent(true)
+		gets = append(gets, b)
+	}
+
 	var cc ConsumedCapacity
-	err = table.Batch("UserID", "Time").
-		Get(keys...).
-		Project("UserID", "Time").
-		Consistent(true).
-		ConsumedCapacity(&cc).
-		All(&results)
+	get1 := gets[0].ConsumedCapacity(&cc)
+	for _, b := range gets[1:] {
+		get1.Merge(b)
+	}
+
+	var results []widget
+	err = get1.All(&results)
 	if err != nil {
 		t.Error("unexpected error:", err)
 	}
 
-	if len(results) != batchSize {
-		t.Error("expected", batchSize, "results, got", len(results))
+	if len(results) != totalBatchSize {
+		t.Error("expected", totalBatchSize, "results, got", len(results))
 	}
 
 	if cc.Total == 0 {
@@ -72,34 +93,39 @@ func TestBatchGetWrite(t *testing.T) {
 	}
 
 	// delete both
-	wrote, err = table.Batch("UserID", "Time").Write().
-		Delete(keys...).Run()
-	if wrote != batchSize {
-		t.Error("unexpected wrote:", wrote, "≠", batchSize)
+	wrote, err = table1.Batch("UserID", "Time").Write().
+		Delete(keys...).
+		DeleteInRange(table2, "UserID", "Time", keys...).
+		Run()
+	if wrote != totalBatchSize {
+		t.Error("unexpected wrote:", wrote, "≠", totalBatchSize)
 	}
 	if err != nil {
 		t.Error("unexpected error:", err)
 	}
 
 	// get both again
-	results = nil
-	err = table.Batch("UserID", "Time").
-		Get(keys...).
-		Consistent(true).
-		All(&results)
-	if err != ErrNotFound {
-		t.Error("expected ErrNotFound, got", err)
-	}
-	if len(results) != 0 {
-		t.Error("expected 0 results, got", len(results))
+	{
+		var results []widget
+		err = table1.Batch("UserID", "Time").
+			Get(keys...).
+			FromRange(table2, "UserID", "Time", keys...).
+			Consistent(true).
+			All(&results)
+		if err != ErrNotFound {
+			t.Error("expected ErrNotFound, got", err)
+		}
+		if len(results) != 0 {
+			t.Error("expected 0 results, got", len(results))
+		}
 	}
 }
 
 func TestBatchGetEmptySets(t *testing.T) {
 	if testDB == nil {
 		t.Skip(offlineSkipMsg)
 	}
-	table := testDB.Table(testTable)
+	table := testDB.Table(testTableWidgets)
 
 	now := time.Now().UnixNano() / 1000000000
 	id := int(now)
@@ -150,7 +176,7 @@ func TestBatchGetEmptySets(t *testing.T) {
 }
 
 func TestBatchEmptyInput(t *testing.T) {
-	table := testDB.Table(testTable)
+	table := testDB.Table(testTableWidgets)
 	var out []any
 	err := table.Batch("UserID", "Time").Get().All(&out)
 	if err != ErrNoInput {