Skip to content

Commit 4db2b8b

Browse files
authored
[Internal] Update Jobs GetRun API to support paginated responses for jobs and ForEach tasks (#386)
## What changes are proposed in this pull request? Introduces extension for jobs getRun call that paginates tasks and iterations arrays in the response and returns aggregated response to the caller. This change is necessary to prepare for jobs API 2.2 release that serves paginated response. Pagination is over once the next_page_token is absent from the response. The pagination logic is not exposed to the customer. ## How is this tested? Unit tests
1 parent 9b7ca5d commit 4db2b8b

File tree

2 files changed

+145
-0
lines changed

2 files changed

+145
-0
lines changed
Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
package com.databricks.sdk.mixin;
2+
3+
import com.databricks.sdk.core.ApiClient;
4+
import com.databricks.sdk.service.jobs.*;
5+
import java.util.Collection;
6+
7+
public class JobsExt extends JobsAPI {
8+
9+
public JobsExt(ApiClient apiClient) {
10+
super(apiClient);
11+
}
12+
13+
public JobsExt(JobsService mock) {
14+
super(mock);
15+
}
16+
17+
/**
18+
* Wrap the {@code JobsApi.getRun} operation to retrieve paginated content without breaking the
19+
* response contract.
20+
*
21+
* <p>Depending on the Jobs API version used under the hood, tasks or iteration runs retrieved by
22+
* the initial request may be truncated due to high cardinalities. Truncation can happen for job
23+
* runs over 100 task runs, as well as ForEach task runs with over 100 iteration runs. To avoid
24+
* returning an incomplete {@code Run} object to the user, this method performs all the requests
25+
* required to collect all task/iteration runs into a single {@code Run} object.
26+
*/
27+
@Override
28+
public Run getRun(GetRunRequest request) {
29+
Run run = super.getRun(request);
30+
31+
/*
32+
* fetch all additional pages (if any) and accumulate the result in a single response
33+
*/
34+
35+
Collection<RunTask> iterations = run.getIterations();
36+
boolean paginatingIterations = iterations != null && !iterations.isEmpty();
37+
38+
Run currRun = run;
39+
while (currRun.getNextPageToken() != null) {
40+
request.setPageToken(currRun.getNextPageToken());
41+
currRun = super.getRun(request);
42+
if (paginatingIterations) {
43+
Collection<RunTask> newIterations = currRun.getIterations();
44+
if (newIterations != null) {
45+
run.getIterations().addAll(newIterations);
46+
}
47+
} else {
48+
Collection<RunTask> newTasks = currRun.getTasks();
49+
if (newTasks != null) {
50+
run.getTasks().addAll(newTasks);
51+
}
52+
}
53+
}
54+
55+
// now that we've added all pages to the Run, the tokens are useless
56+
run.setNextPageToken(null);
57+
58+
return run;
59+
}
60+
}
Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
package com.databricks.sdk.mixin;
2+
3+
import static org.junit.jupiter.api.Assertions.assertEquals;
4+
import static org.mockito.ArgumentMatchers.any;
5+
import static org.mockito.Mockito.*;
6+
7+
import com.databricks.sdk.service.jobs.GetRunRequest;
8+
import com.databricks.sdk.service.jobs.JobsService;
9+
import com.databricks.sdk.service.jobs.Run;
10+
import com.databricks.sdk.service.jobs.RunTask;
11+
import java.util.ArrayList;
12+
import java.util.Collection;
13+
import org.junit.jupiter.api.Test;
14+
import org.mockito.Mockito;
15+
16+
public class JobsExtTest {
17+
18+
@Test
19+
public void testGetRunPaginationWithTasks() {
20+
JobsService service = Mockito.mock(JobsService.class);
21+
22+
Run firstPage = new Run().setNextPageToken("tokenToSecondPage");
23+
addTasks(firstPage, 0L, 1L);
24+
Run secondPage = new Run().setNextPageToken("tokenToThirdPage");
25+
addTasks(secondPage, 2L, 3L);
26+
Run thirdPage = new Run();
27+
addTasks(thirdPage, 4L);
28+
29+
when(service.getRun(any())).thenReturn(firstPage).thenReturn(secondPage).thenReturn(thirdPage);
30+
31+
JobsExt jobsExt = new JobsExt(service);
32+
33+
GetRunRequest request = new GetRunRequest();
34+
35+
Run run = jobsExt.getRun(request);
36+
37+
Run expectedRun = new Run();
38+
addTasks(expectedRun, 0L, 1L, 2L, 3L, 4L);
39+
40+
assertEquals(expectedRun, run);
41+
verify(service, times(3)).getRun(any());
42+
}
43+
44+
@Test
45+
public void testGetRunPaginationWithIterations() {
46+
JobsService service = Mockito.mock(JobsService.class);
47+
48+
Run firstPage = new Run().setNextPageToken("tokenToSecondPage");
49+
addIterations(firstPage, 0L, 1L);
50+
Run secondPage = new Run().setNextPageToken("tokenToThirdPage");
51+
addIterations(secondPage, 2L, 3L);
52+
Run thirdPage = new Run();
53+
addIterations(thirdPage, 4L);
54+
55+
when(service.getRun(any())).thenReturn(firstPage).thenReturn(secondPage).thenReturn(thirdPage);
56+
57+
JobsExt jobsExt = new JobsExt(service);
58+
59+
GetRunRequest request = new GetRunRequest();
60+
61+
Run run = jobsExt.getRun(request);
62+
63+
Run expectedRun = new Run();
64+
addIterations(expectedRun, 0L, 1L, 2L, 3L, 4L);
65+
66+
assertEquals(expectedRun, run);
67+
verify(service, times(3)).getRun(any());
68+
}
69+
70+
private void addTasks(Run run, long... taskRunIds) {
71+
Collection<RunTask> tasks = new ArrayList<>();
72+
for (long runId : taskRunIds) {
73+
tasks.add(new RunTask().setRunId(runId));
74+
}
75+
run.setTasks(tasks);
76+
}
77+
78+
private void addIterations(Run run, long... iterationRunIds) {
79+
Collection<RunTask> iterations = new ArrayList<>();
80+
for (long runId : iterationRunIds) {
81+
iterations.add(new RunTask().setRunId(runId));
82+
}
83+
run.setIterations(iterations);
84+
}
85+
}

0 commit comments

Comments
 (0)