diff --git a/databricks-sdk-java/src/main/java/com/databricks/sdk/mixin/JobsExt.java b/databricks-sdk-java/src/main/java/com/databricks/sdk/mixin/JobsExt.java
index f6f15f905..c2be0eae4 100644
--- a/databricks-sdk-java/src/main/java/com/databricks/sdk/mixin/JobsExt.java
+++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/mixin/JobsExt.java
@@ -20,25 +20,27 @@ public JobsExt(JobsService mock) {
*
*
Depending on the Jobs API version used under the hood, tasks or iteration runs retrieved by
* the initial request may be truncated due to high cardinalities. Truncation can happen for job
- * runs over 100 task runs, as well as ForEach task runs with over 100 iteration runs. To avoid
- * returning an incomplete {@code Run} object to the user, this method performs all the requests
- * required to collect all task/iteration runs into a single {@code Run} object.
+ * runs with over 100 task runs, as well as ForEach task runs with over 100 iteration runs. To
+ * avoid returning an incomplete {@code Run} object to the user, this method performs all the
+ * requests required to collect all task/iteration runs into a single {@code Run} object.
*/
@Override
public Run getRun(GetRunRequest request) {
Run run = super.getRun(request);
- /*
- * fetch all additional pages (if any) and accumulate the result in a single response
- */
-
+ // When querying a Job run, a page token is returned when there are more than 100 tasks. No
+ // iterations are defined for a Job run. Therefore, the next page in the response only includes
+ // the next page of tasks.
+ // When querying a ForEach task run, a page token is returned when there are more than 100
+ // iterations. Only a single task is returned, corresponding to the ForEach task itself.
+ // Therefore, the client only reads the iterations from the next page and not the tasks.
Collection iterations = run.getIterations();
boolean paginatingIterations = iterations != null && !iterations.isEmpty();
- Run currRun = run;
- while (currRun.getNextPageToken() != null) {
- request.setPageToken(currRun.getNextPageToken());
- currRun = super.getRun(request);
+ // runs/get response includes next_page_token as long as there are more pages to fetch.
+ while (run.getNextPageToken() != null) {
+ request.setPageToken(run.getNextPageToken());
+ Run currRun = super.getRun(request);
if (paginatingIterations) {
Collection newIterations = currRun.getIterations();
if (newIterations != null) {
@@ -50,10 +52,23 @@ public Run getRun(GetRunRequest request) {
run.getTasks().addAll(newTasks);
}
}
- }
- // now that we've added all pages to the Run, the tokens are useless
- run.setNextPageToken(null);
+ // Each new page of runs/get response includes the next page of the job_clusters,
+ // job_parameters, and repair history.
+ Collection newClusters = currRun.getJobClusters();
+ if (newClusters != null) {
+ run.getJobClusters().addAll(newClusters);
+ }
+ Collection newParameters = currRun.getJobParameters();
+ if (newParameters != null) {
+ run.getJobParameters().addAll(newParameters);
+ }
+ Collection newRepairHistory = currRun.getRepairHistory();
+ if (newRepairHistory != null) {
+ run.getRepairHistory().addAll(newRepairHistory);
+ }
+ run.setNextPageToken(currRun.getNextPageToken());
+ }
return run;
}
diff --git a/databricks-sdk-java/src/test/java/com/databricks/sdk/mixin/JobsExtTest.java b/databricks-sdk-java/src/test/java/com/databricks/sdk/mixin/JobsExtTest.java
index 73fb22338..ceb65a4c8 100644
--- a/databricks-sdk-java/src/test/java/com/databricks/sdk/mixin/JobsExtTest.java
+++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/mixin/JobsExtTest.java
@@ -4,10 +4,7 @@
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.*;
-import com.databricks.sdk.service.jobs.GetRunRequest;
-import com.databricks.sdk.service.jobs.JobsService;
-import com.databricks.sdk.service.jobs.Run;
-import com.databricks.sdk.service.jobs.RunTask;
+import com.databricks.sdk.service.jobs.*;
import java.util.ArrayList;
import java.util.Collection;
import org.junit.jupiter.api.Test;
@@ -21,10 +18,15 @@ public void testGetRunPaginationWithTasks() {
Run firstPage = new Run().setNextPageToken("tokenToSecondPage");
addTasks(firstPage, 0L, 1L);
+ addJobClusters(firstPage, "clusterKey1", "clusterKey2");
+ addJobParameters(firstPage, "parameterKey1", "parameterKey2");
Run secondPage = new Run().setNextPageToken("tokenToThirdPage");
addTasks(secondPage, 2L, 3L);
+ addJobClusters(secondPage, "clusterKey3");
+ addJobParameters(secondPage, "parameterKey3", "parameterKey4");
Run thirdPage = new Run();
addTasks(thirdPage, 4L);
+ addJobParameters(thirdPage, "parameterKey5");
when(service.getRun(any())).thenReturn(firstPage).thenReturn(secondPage).thenReturn(thirdPage);
@@ -36,6 +38,14 @@ public void testGetRunPaginationWithTasks() {
Run expectedRun = new Run();
addTasks(expectedRun, 0L, 1L, 2L, 3L, 4L);
+ addJobClusters(expectedRun, "clusterKey1", "clusterKey2", "clusterKey3");
+ addJobParameters(
+ expectedRun,
+ "parameterKey1",
+ "parameterKey2",
+ "parameterKey3",
+ "parameterKey4",
+ "parameterKey5");
assertEquals(expectedRun, run);
verify(service, times(3)).getRun(any());
@@ -82,4 +92,20 @@ private void addIterations(Run run, long... iterationRunIds) {
}
run.setIterations(iterations);
}
+
+ private void addJobClusters(Run run, String... clusterKeys) {
+ Collection clusters = new ArrayList<>();
+ for (String clusterKey : clusterKeys) {
+ clusters.add(new JobCluster().setJobClusterKey(clusterKey));
+ }
+ run.setJobClusters(clusters);
+ }
+
+ private void addJobParameters(Run run, String... parameterKeys) {
+ Collection parameters = new ArrayList<>();
+ for (String parameterKey : parameterKeys) {
+ parameters.add(new JobParameter().setName(parameterKey).setValue(parameterKey));
+ }
+ run.setJobParameters(parameters);
+ }
}