From 38c18671e27dd71ef9acbfe24550df779b2e8890 Mon Sep 17 00:00:00 2001
From: guregu <greg@toki.waseda.jp>
Date: Fri, 23 Aug 2024 03:36:38 +0900
Subject: [PATCH] add ItemTableIter

---
 batchget.go       |  1 +
 batchget_go123.go | 83 +++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 84 insertions(+)
 create mode 100644 batchget_go123.go

diff --git a/batchget.go b/batchget.go
index cfb0c3e..698d335 100644
--- a/batchget.go
+++ b/batchget.go
@@ -193,6 +193,7 @@ func (bg *BatchGet) Iter() Iter {
 
 // IterWithTable is like [BatchGet.Iter], but will update the value pointed by tablePtr after each iteration.
 // This can be useful when getting from multiple tables to determine which table the latest item came from.
+// See: [BatchGet.ItemTableIter] for a nicer way to do this.
 //
 // For example, you can utilize this iterator to read the results into different structs.
 //
diff --git a/batchget_go123.go b/batchget_go123.go
new file mode 100644
index 0000000..835a059
--- /dev/null
+++ b/batchget_go123.go
@@ -0,0 +1,83 @@
+//go:build go1.23
+
+package dynamo
+
+import (
+	"context"
+	"iter"
+)
+
+type ItemTableIter[V any] interface {
+	// Items is a sequence of item and table names.
+	// This is a single use iterator.
+	// Be sure to check for errors with Err afterwards.
+	Items(context.Context) iter.Seq2[V, string]
+	// Err must be checked after iterating.
+	Err() error
+}
+
+// ItemTableIter returns an iterator of (raw item, table name).
+// To specify a type, use [BatchGetIter] instead.
+//
+// For example, you can utilize this iterator to read the results into different structs.
+//
+//	widgetBatch := widgetsTable.Batch("UserID").Get(dynamo.Keys{userID})
+//	sprocketBatch := sprocketsTable.Batch("UserID").Get(dynamo.Keys{userID})
+//
+//	iter := widgetBatch.Merge(sprocketBatch).ItemTableIter(&table)
+//
+//	// now we will use the table iterator to unmarshal the values into their respective types
+//	var s sprocket
+//	var w widget
+//	for raw, table := range iter.Items {
+//		if table == "Widgets" {
+//			err := dynamo.UnmarshalItem(raw, &w)
+//			if err != nil {
+//				fmt.Println(err)
+//			}
+//		} else if table == "Sprockets" {
+//			err := dynamo.UnmarshalItem(raw, &s)
+//			if err != nil {
+//				fmt.Println(err)
+//			}
+//		} else {
+//			fmt.Printf("Unexpected Table: %s\n", table)
+//		}
+//	}
+//
+//	if iter.Err() != nil {
+//		fmt.Println(iter.Err())
+//	}
+func (bg *BatchGet) ItemTableIter() ItemTableIter[Item] {
+	return newBgIter2[Item](bg)
+}
+
+type bgIter2[V any] struct {
+	Iter
+	table string
+}
+
+func newBgIter2[V any](bg *BatchGet) *bgIter2[V] {
+	iter := new(bgIter2[V])
+	iter.Iter = bg.IterWithTable(&iter.table)
+	return iter
+}
+
+// Items is a sequence of item and table names.
+// This is a single use iterator.
+// Be sure to check for errors with Err afterwards.
+func (iter *bgIter2[V]) Items(ctx context.Context) iter.Seq2[V, string] {
+	return func(yield func(V, string) bool) {
+		item := new(V)
+		for iter.Next(ctx, item) {
+			if !yield(*item, iter.table) {
+				break
+			}
+			item = new(V)
+		}
+	}
+}
+
+func BatchGetIter[V any](bg *BatchGet) ItemTableIter[V] {
+	return newBgIter2[V](bg)
+}