Skip to content

Commit a6749d0

Browse files
committed
use wait results from sdk
1 parent ea255e4 commit a6749d0

2 files changed

Lines changed: 33 additions & 58 deletions

File tree

cmd/execution/results.go

Lines changed: 11 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -26,8 +26,8 @@ func newResultsCmd() *cobra.Command {
2626
" 3. If completed: returns the result data\n" +
2727
" 4. If failed/cancelled: returns the error details\n\n" +
2828
"Use --no-wait to return the current state immediately without polling.",
29-
Args: cobra.ExactArgs(1),
30-
RunE: runResults,
29+
Args: cobra.ExactArgs(1),
30+
RunE: runResults,
3131
}
3232

3333
cmd.Flags().Int("limit", 0, "maximum number of result rows to return (0 = all)")
@@ -81,36 +81,17 @@ func runResults(cmd *cobra.Command, args []string) error {
8181
maxRetries = 1
8282
}
8383

84-
return waitForResults(cmd, client, executionID, opts, PollInterval, maxRetries)
85-
}
86-
87-
func waitForResults(
88-
cmd *cobra.Command,
89-
client dune.DuneClient,
90-
executionID string,
91-
opts models.ResultOptions,
92-
interval time.Duration,
93-
maxRetries int,
94-
) error {
95-
for i := 0; i < maxRetries; i++ {
96-
resp, err := client.QueryResultsV2(executionID, opts)
97-
if err != nil {
98-
return err
99-
}
100-
101-
switch resp.State {
102-
case "QUERY_STATE_PENDING", "QUERY_STATE_EXECUTING":
103-
// still running, wait and retry
104-
default:
105-
return handleResultsResponse(cmd, executionID, resp)
106-
}
107-
108-
if i < maxRetries-1 {
109-
time.Sleep(interval)
110-
}
84+
exec := dune.NewExecution(client, executionID)
85+
if _, err := exec.WaitGetResults(PollInterval, maxRetries); err != nil {
86+
return err
11187
}
11288

113-
return fmt.Errorf("timed out waiting for execution %s to complete", executionID)
89+
// Fetch final results with any limit/offset options.
90+
resp, err := client.QueryResultsV2(executionID, opts)
91+
if err != nil {
92+
return err
93+
}
94+
return handleResultsResponse(cmd, executionID, resp)
11495
}
11596

11697
func handleResultsResponse(cmd *cobra.Command, executionID string, resp *models.ResultsResponse) error {

cmd/execution/results_test.go

Lines changed: 22 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -13,9 +13,9 @@ import (
1313
)
1414

1515
var testResultsResponse = &models.ResultsResponse{
16-
QueryID: 4125432,
17-
State: "QUERY_STATE_COMPLETED",
18-
ExecutionEndedAt: ptrTime(time.Now()),
16+
QueryID: 4125432,
17+
State: "QUERY_STATE_COMPLETED",
18+
ExecutionEndedAt: ptrTime(time.Now()),
1919
IsExecutionFinished: true,
2020
Result: models.Result{
2121
Metadata: models.ResultMetadata{
@@ -40,7 +40,7 @@ func TestResultsSuccess(t *testing.T) {
4040
}
4141

4242
root, buf := newTestRoot(mock)
43-
root.SetArgs([]string{"execution", "results", "01ABCDEFGHIJKLMNOPQRSTUV"})
43+
root.SetArgs([]string{"execution", "results", "--no-wait", "01ABCDEFGHIJKLMNOPQRSTUV"})
4444
require.NoError(t, root.Execute())
4545

4646
out := buf.String()
@@ -59,7 +59,7 @@ func TestResultsJSONOutput(t *testing.T) {
5959
}
6060

6161
root, buf := newTestRoot(mock)
62-
root.SetArgs([]string{"execution", "results", "01ABC", "-o", "json"})
62+
root.SetArgs([]string{"execution", "results", "--no-wait", "01ABC", "-o", "json"})
6363
require.NoError(t, root.Execute())
6464

6565
var got models.ResultsResponse
@@ -117,7 +117,7 @@ func TestResultsFailed(t *testing.T) {
117117
}
118118

119119
root, _ := newTestRoot(mock)
120-
root.SetArgs([]string{"execution", "results", "01ABC"})
120+
root.SetArgs([]string{"execution", "results", "--no-wait", "01ABC"})
121121
err := root.Execute()
122122
require.Error(t, err)
123123
assert.Contains(t, err.Error(), "syntax error at line 1")
@@ -135,7 +135,7 @@ func TestResultsCancelled(t *testing.T) {
135135
}
136136

137137
root, _ := newTestRoot(mock)
138-
root.SetArgs([]string{"execution", "results", "01ABC"})
138+
root.SetArgs([]string{"execution", "results", "--no-wait", "01ABC"})
139139
err := root.Execute()
140140
require.Error(t, err)
141141
assert.Contains(t, err.Error(), "cancelled")
@@ -151,7 +151,7 @@ func TestResultsWithLimitAndOffset(t *testing.T) {
151151
}
152152

153153
root, _ := newTestRoot(mock)
154-
root.SetArgs([]string{"execution", "results", "01ABC", "--limit", "10", "--offset", "5"})
154+
root.SetArgs([]string{"execution", "results", "--no-wait", "01ABC", "--limit", "10", "--offset", "5"})
155155
require.NoError(t, root.Execute())
156156

157157
require.NotNil(t, capturedOpts.Page)
@@ -167,7 +167,7 @@ func TestResultsAPIError(t *testing.T) {
167167
}
168168

169169
root, _ := newTestRoot(mock)
170-
root.SetArgs([]string{"execution", "results", "01ABC"})
170+
root.SetArgs([]string{"execution", "results", "--no-wait", "01ABC"})
171171
err := root.Execute()
172172
require.Error(t, err)
173173
assert.Contains(t, err.Error(), "api: connection refused")
@@ -204,7 +204,8 @@ func TestResultsWaitPollsUntilComplete(t *testing.T) {
204204
root.SetArgs([]string{"execution", "results", "--timeout", "10", "01ABC"})
205205
require.NoError(t, root.Execute())
206206

207-
assert.Equal(t, 3, callCount)
207+
// 3 calls from WaitGetResults (2 executing + 1 completed) + 1 final fetch
208+
assert.Equal(t, 4, callCount)
208209
out := buf.String()
209210
assert.Contains(t, out, "block_number")
210211
assert.Contains(t, out, "2 rows")
@@ -226,7 +227,8 @@ func TestResultsWaitPollsUntilFailed(t *testing.T) {
226227
}, nil
227228
}
228229
return &models.ResultsResponse{
229-
State: "QUERY_STATE_FAILED",
230+
State: "QUERY_STATE_FAILED",
231+
IsExecutionFinished: true,
230232
Error: &models.ExecutionError{
231233
Message: "out of memory",
232234
},
@@ -239,31 +241,27 @@ func TestResultsWaitPollsUntilFailed(t *testing.T) {
239241
err := root.Execute()
240242
require.Error(t, err)
241243
assert.Contains(t, err.Error(), "out of memory")
242-
assert.Equal(t, 2, callCount)
244+
// 2 calls from WaitGetResults (1 pending + 1 failed) + 1 final fetch
245+
assert.Equal(t, 3, callCount)
243246
}
244247

245-
func TestResultsWaitTimeout(t *testing.T) {
248+
func TestResultsWaitRetriesExhausted(t *testing.T) {
246249
execution.PollInterval = 0
247250
t.Cleanup(func() {
248251
execution.PollInterval = 2 * time.Second
249252
})
250253

251-
callCount := 0
252254
mock := &mockClient{
253255
queryResultsV2Fn: func(_ string, _ models.ResultOptions) (*models.ResultsResponse, error) {
254-
callCount++
255-
return &models.ResultsResponse{
256-
State: "QUERY_STATE_EXECUTING",
257-
}, nil
256+
return nil, errors.New("server unavailable")
258257
},
259258
}
260259

261260
root, _ := newTestRoot(mock)
262261
root.SetArgs([]string{"execution", "results", "--timeout", "1", "01ABC"})
263262
err := root.Execute()
264263
require.Error(t, err)
265-
assert.Contains(t, err.Error(), "timed out waiting for execution")
266-
264+
assert.Contains(t, err.Error(), "retries have been exhausted")
267265
}
268266

269267
func TestResultsWaitAPIError(t *testing.T) {
@@ -272,22 +270,18 @@ func TestResultsWaitAPIError(t *testing.T) {
272270
execution.PollInterval = 2 * time.Second
273271
})
274272

275-
callCount := 0
276273
mock := &mockClient{
277274
queryResultsV2Fn: func(_ string, _ models.ResultOptions) (*models.ResultsResponse, error) {
278-
callCount++
279-
if callCount < 2 {
280-
return &models.ResultsResponse{
281-
State: "QUERY_STATE_PENDING",
282-
}, nil
283-
}
284275
return nil, errors.New("api: rate limited")
285276
},
286277
}
287278

288279
root, _ := newTestRoot(mock)
289-
root.SetArgs([]string{"execution", "results", "--timeout", "10", "01ABC"})
280+
// With --timeout=2 and PollInterval=0, maxRetries=2.
281+
// The SDK retries on API errors until errCount > maxRetries.
282+
root.SetArgs([]string{"execution", "results", "--timeout", "2", "01ABC"})
290283
err := root.Execute()
291284
require.Error(t, err)
285+
assert.Contains(t, err.Error(), "retries have been exhausted")
292286
assert.Contains(t, err.Error(), "api: rate limited")
293287
}

0 commit comments

Comments
 (0)