From 22f04c2564dedcf1152bfe0611a0f160e9eefca6 Mon Sep 17 00:00:00 2001 From: Giorgi Kikolashvili Date: Wed, 22 Jan 2025 18:42:50 +0100 Subject: [PATCH 1/4] GetRun logic paginates more arrays --- .../com/databricks/sdk/mixin/JobsExt.java | 24 ++++++++++----- .../com/databricks/sdk/mixin/JobsExtTest.java | 30 ++++++++++++++++--- 2 files changed, 43 insertions(+), 11 deletions(-) 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..2197a4e56 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 @@ -35,10 +35,9 @@ public Run getRun(GetRunRequest request) { 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); + while (run.getNextPageToken() != null) { + request.setPageToken(run.getNextPageToken()); + Run currRun = super.getRun(request); if (paginatingIterations) { Collection newIterations = currRun.getIterations(); if (newIterations != null) { @@ -50,10 +49,21 @@ 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); + 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..cbb8eea2a 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,8 @@ 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 +19,16 @@ 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 +40,8 @@ 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 +88,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); + } } From 48eaf74c5821b0cefd55d42e32362e5c2deea5e4 Mon Sep 17 00:00:00 2001 From: Giorgi Kikolashvili Date: Wed, 22 Jan 2025 19:11:46 +0100 Subject: [PATCH 2/4] Apply fmt --- .../main/java/com/databricks/sdk/mixin/JobsExt.java | 2 +- .../java/com/databricks/sdk/mixin/JobsExtTest.java | 10 +++++++--- 2 files changed, 8 insertions(+), 4 deletions(-) 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 2197a4e56..5253f2ded 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 @@ -60,7 +60,7 @@ public Run getRun(GetRunRequest request) { } Collection newRepairHistory = currRun.getRepairHistory(); if (newRepairHistory != null) { - run.getRepairHistory().addAll(newRepairHistory); + run.getRepairHistory().addAll(newRepairHistory); } run.setNextPageToken(currRun.getNextPageToken()); } 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 cbb8eea2a..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 @@ -5,7 +5,6 @@ import static org.mockito.Mockito.*; import com.databricks.sdk.service.jobs.*; - import java.util.ArrayList; import java.util.Collection; import org.junit.jupiter.api.Test; @@ -29,7 +28,6 @@ public void testGetRunPaginationWithTasks() { addTasks(thirdPage, 4L); addJobParameters(thirdPage, "parameterKey5"); - when(service.getRun(any())).thenReturn(firstPage).thenReturn(secondPage).thenReturn(thirdPage); JobsExt jobsExt = new JobsExt(service); @@ -41,7 +39,13 @@ 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"); + addJobParameters( + expectedRun, + "parameterKey1", + "parameterKey2", + "parameterKey3", + "parameterKey4", + "parameterKey5"); assertEquals(expectedRun, run); verify(service, times(3)).getRun(any()); From 8d0698bf109e27b0cb6592dde780f962f1a2680e Mon Sep 17 00:00:00 2001 From: Giorgi Kikolashvili Date: Tue, 4 Feb 2025 12:55:14 +0100 Subject: [PATCH 3/4] Add comments --- .../main/java/com/databricks/sdk/mixin/JobsExt.java | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) 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 5253f2ded..7f4a86651 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,7 +20,7 @@ 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 + * 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. */ @@ -28,13 +28,12 @@ public JobsExt(JobsService mock) { 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(); + // 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); @@ -50,6 +49,7 @@ public Run getRun(GetRunRequest request) { } } + // 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); From 73925a77a49d8f1e656657fd39e305ea0df99c3b Mon Sep 17 00:00:00 2001 From: Giorgi Kikolashvili Date: Tue, 4 Feb 2025 13:07:21 +0100 Subject: [PATCH 4/4] fmt --- .../java/com/databricks/sdk/mixin/JobsExt.java | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) 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 7f4a86651..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,16 +20,20 @@ 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 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. + * 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); - // 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. + // 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(); @@ -49,7 +53,8 @@ public Run getRun(GetRunRequest request) { } } - // Each new page of runs/get response includes the next page of the job_clusters, job_parameters, and repair history. + // 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);