From 4db2b8ba8591b7fe49dd04dc5501ac1733bf6650 Mon Sep 17 00:00:00 2001 From: Giorgi Kikolashvili <47174341+gkiko10@users.noreply.github.com> Date: Fri, 15 Nov 2024 17:19:58 +0100 Subject: [PATCH] [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 --- .../com/databricks/sdk/mixin/JobsExt.java | 60 +++++++++++++ .../com/databricks/sdk/mixin/JobsExtTest.java | 85 +++++++++++++++++++ 2 files changed, 145 insertions(+) create mode 100644 databricks-sdk-java/src/main/java/com/databricks/sdk/mixin/JobsExt.java create mode 100644 databricks-sdk-java/src/test/java/com/databricks/sdk/mixin/JobsExtTest.java 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 new file mode 100644 index 000000000..f6f15f905 --- /dev/null +++ b/databricks-sdk-java/src/main/java/com/databricks/sdk/mixin/JobsExt.java @@ -0,0 +1,60 @@ +package com.databricks.sdk.mixin; + +import com.databricks.sdk.core.ApiClient; +import com.databricks.sdk.service.jobs.*; +import java.util.Collection; + +public class JobsExt extends JobsAPI { + + public JobsExt(ApiClient apiClient) { + super(apiClient); + } + + public JobsExt(JobsService mock) { + super(mock); + } + + /** + * Wrap the {@code JobsApi.getRun} operation to retrieve paginated content without breaking the + * response contract. + * + *

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. + */ + @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 + */ + + 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); + if (paginatingIterations) { + Collection newIterations = currRun.getIterations(); + if (newIterations != null) { + run.getIterations().addAll(newIterations); + } + } else { + Collection newTasks = currRun.getTasks(); + if (newTasks != null) { + run.getTasks().addAll(newTasks); + } + } + } + + // now that we've added all pages to the Run, the tokens are useless + run.setNextPageToken(null); + + 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 new file mode 100644 index 000000000..73fb22338 --- /dev/null +++ b/databricks-sdk-java/src/test/java/com/databricks/sdk/mixin/JobsExtTest.java @@ -0,0 +1,85 @@ +package com.databricks.sdk.mixin; + +import static org.junit.jupiter.api.Assertions.assertEquals; +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 java.util.ArrayList; +import java.util.Collection; +import org.junit.jupiter.api.Test; +import org.mockito.Mockito; + +public class JobsExtTest { + + @Test + public void testGetRunPaginationWithTasks() { + JobsService service = Mockito.mock(JobsService.class); + + Run firstPage = new Run().setNextPageToken("tokenToSecondPage"); + addTasks(firstPage, 0L, 1L); + Run secondPage = new Run().setNextPageToken("tokenToThirdPage"); + addTasks(secondPage, 2L, 3L); + Run thirdPage = new Run(); + addTasks(thirdPage, 4L); + + when(service.getRun(any())).thenReturn(firstPage).thenReturn(secondPage).thenReturn(thirdPage); + + JobsExt jobsExt = new JobsExt(service); + + GetRunRequest request = new GetRunRequest(); + + Run run = jobsExt.getRun(request); + + Run expectedRun = new Run(); + addTasks(expectedRun, 0L, 1L, 2L, 3L, 4L); + + assertEquals(expectedRun, run); + verify(service, times(3)).getRun(any()); + } + + @Test + public void testGetRunPaginationWithIterations() { + JobsService service = Mockito.mock(JobsService.class); + + Run firstPage = new Run().setNextPageToken("tokenToSecondPage"); + addIterations(firstPage, 0L, 1L); + Run secondPage = new Run().setNextPageToken("tokenToThirdPage"); + addIterations(secondPage, 2L, 3L); + Run thirdPage = new Run(); + addIterations(thirdPage, 4L); + + when(service.getRun(any())).thenReturn(firstPage).thenReturn(secondPage).thenReturn(thirdPage); + + JobsExt jobsExt = new JobsExt(service); + + GetRunRequest request = new GetRunRequest(); + + Run run = jobsExt.getRun(request); + + Run expectedRun = new Run(); + addIterations(expectedRun, 0L, 1L, 2L, 3L, 4L); + + assertEquals(expectedRun, run); + verify(service, times(3)).getRun(any()); + } + + private void addTasks(Run run, long... taskRunIds) { + Collection tasks = new ArrayList<>(); + for (long runId : taskRunIds) { + tasks.add(new RunTask().setRunId(runId)); + } + run.setTasks(tasks); + } + + private void addIterations(Run run, long... iterationRunIds) { + Collection iterations = new ArrayList<>(); + for (long runId : iterationRunIds) { + iterations.add(new RunTask().setRunId(runId)); + } + run.setIterations(iterations); + } +}