Skip to content

Commit 83acc1e

Browse files
committed
all: simplify iterator API
1 parent f42a706 commit 83acc1e

15 files changed

+45
-42
lines changed

README.md

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -194,12 +194,14 @@ result, err := c.GetParseResult(
194194
tensorlake.WithSSE(true),
195195
tensorlake.WithOnUpdate(func(name tensorlake.ParseEventName, r *tensorlake.ParseResult) {
196196
switch eventName {
197-
case tensorlake.sseEventParseQueued:
197+
case tensorlake.SSEEventParseQueued:
198198
fmt.Println("Job queued")
199-
case tensorlake.sseEventParseUpdate:
199+
case tensorlake.SSEEventParseUpdate:
200200
fmt.Printf("Progress: %d/%d pages\n", r.ParsedPagesCount, r.TotalPages)
201-
case tensorlake.sseEventParseDone:
201+
case tensorlake.SSEEventParseDone:
202202
fmt.Println("Complete!")
203+
case tensorlake.SSEEventParseFailed:
204+
fmt.Printf("Failed: %s\n", r.Error)
203205
}
204206
}),
205207
)
@@ -211,23 +213,23 @@ Easily iterate through paginated results:
211213

212214
```go
213215
// Iterate all files
214-
for file, err := range c.IterFiles(ctx, 50, tensorlake.PaginationDirectionNext) {
216+
for file, err := range c.IterFiles(ctx, 50) {
215217
if err != nil {
216218
panic(err)
217219
}
218220
fmt.Printf("File: %s\n", file.FileName)
219221
}
220222

221223
// Iterate all parse jobs
222-
for job, err := range c.IterParseJobs(ctx, 50, tensorlake.PaginationDirectionNext) {
224+
for job, err := range c.IterParseJobs(ctx, 50) {
223225
if err != nil {
224226
panic(err)
225227
}
226228
fmt.Printf("Job %s: Status: %s\n", job.ParseId, job.Status)
227229
}
228230

229231
// Iterate all datasets
230-
for dataset, err := range c.IterDatasets(ctx, 50, tensorlake.PaginationDirectionNext) {
232+
for dataset, err := range c.IterDatasets(ctx, 50) {
231233
if err != nil {
232234
panic(err)
233235
}

dataset_list.go

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -25,14 +25,14 @@ import (
2525
)
2626

2727
// IterDatasets iterates over all datasets in the organization.
28-
func (c *Client) IterDatasets(ctx context.Context, limit int, direction PaginationDirection) iter.Seq2[Dataset, error] {
28+
func (c *Client) IterDatasets(ctx context.Context, batchSize int) iter.Seq2[Dataset, error] {
2929
return func(yield func(Dataset, error) bool) {
3030
cursor := ""
3131
for {
3232
listResp, err := c.ListDatasets(ctx, &ListDatasetsRequest{
3333
Cursor: cursor,
34-
Limit: limit,
35-
Direction: direction,
34+
Limit: batchSize,
35+
Direction: PaginationDirectionNext,
3636
})
3737
if err != nil {
3838
yield(Dataset{}, err)
@@ -52,14 +52,15 @@ func (c *Client) IterDatasets(ctx context.Context, limit int, direction Paginati
5252
}
5353

5454
// IterDatasetData iterates over all dataset data in the organization.
55-
func (c *Client) IterDatasetData(ctx context.Context, limit int, direction PaginationDirection) iter.Seq2[ParseResult, error] {
55+
func (c *Client) IterDatasetData(ctx context.Context, datasetId string, batchSize int) iter.Seq2[ParseResult, error] {
5656
return func(yield func(ParseResult, error) bool) {
5757
cursor := ""
5858
for {
5959
listResp, err := c.ListDatasetData(ctx, &ListDatasetDataRequest{
60+
DatasetId: datasetId,
6061
Cursor: cursor,
61-
Limit: limit,
62-
Direction: direction,
62+
Limit: batchSize,
63+
Direction: PaginationDirectionNext,
6364
})
6465
if err != nil {
6566
yield(ParseResult{}, err)

dataset_test.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,7 @@ func TestDataset(t *testing.T) {
3939

4040
// List datasets.
4141
datasets := []string{}
42-
for d, err := range c.IterDatasets(t.Context(), 1, PaginationDirectionNext) {
42+
for d, err := range c.IterDatasets(t.Context(), 1) {
4343
if err != nil {
4444
t.Fatalf("failed to list datasets: %v", err)
4545
}
@@ -113,7 +113,7 @@ func TestDataset(t *testing.T) {
113113

114114
// Check if the dataset is deleted.
115115
datasets = []string{}
116-
for d, err := range c.IterDatasets(t.Context(), 1, PaginationDirectionNext) {
116+
for d, err := range c.IterDatasets(t.Context(), 1) {
117117
if err != nil {
118118
t.Fatalf("failed to list datasets: %v", err)
119119
}

docs/dataset-apis.md

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -285,13 +285,13 @@ if response.HasMore {
285285
For convenience, use the `IterDatasets` method:
286286

287287
```go
288-
func (c *Client) IterDatasets(ctx context.Context, limit int, direction PaginationDirection) iter.Seq2[Dataset, error]
288+
func (c *Client) IterDatasets(ctx context.Context, batchSize int) iter.Seq2[Dataset, error]
289289
```
290290

291291
**Example:**
292292

293293
```go
294-
for dataset, err := range client.IterDatasets(context.Background(), 50, tensorlake.PaginationDirectionNext) {
294+
for dataset, err := range client.IterDatasets(context.Background(), 50) {
295295
if err != nil {
296296
log.Fatal(err)
297297
}

docs/file-apis.md

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -175,13 +175,13 @@ if response.HasMore {
175175
For convenience, use the `IterFiles` method to iterate through all files:
176176

177177
```go
178-
func (c *Client) IterFiles(ctx context.Context, limit int, direction PaginationDirection) iter.Seq2[FileInfo, error]
178+
func (c *Client) IterFiles(ctx context.Context, batchSize int) iter.Seq2[FileInfo, error]
179179
```
180180

181181
**Example:**
182182

183183
```go
184-
for file, err := range client.IterFiles(context.Background(), 50, tensorlake.PaginationDirectionNext) {
184+
for file, err := range client.IterFiles(context.Background(), 50) {
185185
if err != nil {
186186
log.Fatal(err)
187187
}

docs/parse-apis.md

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -453,13 +453,13 @@ for _, job := range response.Items {
453453
### Iterate All Parse Jobs
454454

455455
```go
456-
func (c *Client) IterParseJobs(ctx context.Context, limit int, direction PaginationDirection) iter.Seq2[ParseResult, error]
456+
func (c *Client) IterParseJobs(ctx context.Context, batchSize int) iter.Seq2[ParseResult, error]
457457
```
458458

459459
**Example:**
460460

461461
```go
462-
for job, err := range client.IterParseJobs(context.Background(), 50, tensorlake.PaginationDirectionNext) {
462+
for job, err := range client.IterParseJobs(context.Background(), 50) {
463463
if err != nil {
464464
log.Fatal(err)
465465
}

file_list.go

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -25,14 +25,14 @@ import (
2525
)
2626

2727
// IterFiles iterates over all files in the project.
28-
func (c *Client) IterFiles(ctx context.Context, limit int, direction PaginationDirection) iter.Seq2[FileInfo, error] {
28+
func (c *Client) IterFiles(ctx context.Context, batchSize int) iter.Seq2[FileInfo, error] {
2929
return func(yield func(FileInfo, error) bool) {
3030
cursor := ""
3131
for {
3232
listResp, err := c.ListFiles(ctx, &ListFilesRequest{
3333
Cursor: cursor,
34-
Limit: limit,
35-
Direction: direction,
34+
Limit: batchSize,
35+
Direction: PaginationDirectionNext,
3636
})
3737
if err != nil {
3838
yield(FileInfo{}, err)

file_test.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -80,7 +80,7 @@ func TestFileManagement(t *testing.T) {
8080

8181
// List the files. Iterate through all the pages.
8282
files := []string{}
83-
for f, err := range c.IterFiles(t.Context(), 1, PaginationDirectionNext) {
83+
for f, err := range c.IterFiles(t.Context(), 1) {
8484
if err != nil {
8585
t.Fatalf("failed to list files: %v", err)
8686
}
@@ -124,7 +124,7 @@ func TestFileManagement(t *testing.T) {
124124

125125
// Validate file is deleted.
126126
files = []string{}
127-
for f, err := range c.IterFiles(t.Context(), 1, PaginationDirectionNext) {
127+
for f, err := range c.IterFiles(t.Context(), 1) {
128128
if err != nil {
129129
t.Fatalf("failed to list files: %v", err)
130130
}

parse_classify_test.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -101,7 +101,7 @@ func TestClassifyDocument(t *testing.T) {
101101

102102
// Validate classify results.
103103
jobs := []string{}
104-
for j, err := range c.IterParseJobs(t.Context(), 1, PaginationDirectionNext) {
104+
for j, err := range c.IterParseJobs(t.Context(), 1) {
105105
if err != nil {
106106
t.Fatalf("failed to list parse jobs: %v", err)
107107
}

parse_extract_test.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -116,7 +116,7 @@ func TestExtractDocument(t *testing.T) {
116116

117117
// Validate parse results.
118118
jobs := []string{}
119-
for j, err := range c.IterParseJobs(t.Context(), 1, PaginationDirectionNext) {
119+
for j, err := range c.IterParseJobs(t.Context(), 1) {
120120
if err != nil {
121121
t.Fatalf("failed to list parse jobs: %v", err)
122122
}

0 commit comments

Comments
 (0)