Skip to content

Commit 431be3a

Browse files
gkiko10gkiko10
committed
GetJob pagination implementation
1 parent 6b76d2a commit 431be3a

File tree

2 files changed

+149
-4
lines changed

2 files changed

+149
-4
lines changed

‎databricks-sdk-java/src/main/java/com/databricks/sdk/mixin/JobsExt.java‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,4 +57,42 @@ public Run getRun(GetRunRequest request) {
5757

5858
return run;
5959
}
60+
61+
/**
62+
* Wrap the {@code JobsApi.get} operation to retrieve paginated content without breaking the
63+
* response contract.
64+
*
65+
* <p>Depending on the Jobs API version used under the hood, tasks or job_clusters retrieved by
66+
* the initial request may be truncated due to high cardinalities. Truncation can happen for jobs
67+
* with over 100 tasks, as well as job_clusters with over 100 elements. To avoid returning an
68+
* incomplete {@code Job} object to the user, this method performs all the requests required to
69+
* collect all tasks and job_clusters into a single {@code Job} object.
70+
*/
71+
public Job get(GetJobRequest request) {
72+
Job job = super.get(request);
73+
74+
while (job.getNextPageToken() != null) {
75+
request.setPageToken(job.getNextPageToken());
76+
Job currJob = super.get(request);
77+
Collection<Task> newTasks = currJob.getSettings().getTasks();
78+
if (newTasks != null) {
79+
job.getSettings().getTasks().addAll(newTasks);
80+
}
81+
Collection<JobCluster> newClusters = currJob.getSettings().getJobClusters();
82+
if (newClusters != null) {
83+
job.getSettings().getJobClusters().addAll(newClusters);
84+
}
85+
Collection<JobParameterDefinition> newParameters = currJob.getSettings().getParameters();
86+
if (newParameters != null) {
87+
job.getSettings().getParameters().addAll(newParameters);
88+
}
89+
Collection<JobEnvironment> newEnvironments = currJob.getSettings().getEnvironments();
90+
if (newEnvironments != null) {
91+
job.getSettings().getEnvironments().addAll(newEnvironments);
92+
}
93+
job.setNextPageToken(currJob.getNextPageToken());
94+
}
95+
96+
return job;
97+
}
6098
}

‎databricks-sdk-java/src/test/java/com/databricks/sdk/mixin/JobsExtTest.java‎

Lines changed: 111 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4,12 +4,10 @@
44
import static org.mockito.ArgumentMatchers.any;
55
import static org.mockito.Mockito.*;
66

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;
7+
import com.databricks.sdk.service.jobs.*;
118
import java.util.ArrayList;
129
import java.util.Collection;
10+
import java.util.List;
1311
import org.junit.jupiter.api.Test;
1412
import org.mockito.Mockito;
1513

@@ -82,4 +80,113 @@ private void addIterations(Run run, long... iterationRunIds) {
8280
}
8381
run.setIterations(iterations);
8482
}
83+
84+
private void addJobClusters(Run run, String... clusterKeys) {
85+
Collection<JobCluster> clusters = new ArrayList<>();
86+
for (String clusterKey : clusterKeys) {
87+
clusters.add(new JobCluster().setJobClusterKey(clusterKey));
88+
}
89+
run.setJobClusters(clusters);
90+
}
91+
92+
private void addJobParameters(Run run, String... parameterKeys) {
93+
Collection<JobParameter> parameters = new ArrayList<>();
94+
for (String parameterKey : parameterKeys) {
95+
parameters.add(new JobParameter().setName(parameterKey).setValue(parameterKey));
96+
}
97+
run.setJobParameters(parameters);
98+
}
99+
100+
@Test
101+
public void testGetJobPaginationWithTasks() {
102+
JobsService service = Mockito.mock(JobsService.class);
103+
Task taskKey1 = new Task().setTaskKey("taskKey1");
104+
Task taskKey2 = new Task().setTaskKey("taskKey2");
105+
List<Task> tasks = new ArrayList<>();
106+
tasks.add(taskKey1);
107+
tasks.add(taskKey2);
108+
109+
JobCluster jobClusterKey1 = new JobCluster().setJobClusterKey("jobClusterKey1");
110+
List<JobCluster> jobClusters = new ArrayList<>();
111+
jobClusters.add(jobClusterKey1);
112+
Job firstPage =
113+
new Job().setSettings(new JobSettings().setTasks(tasks).setJobClusters(jobClusters));
114+
115+
when(service.get(any())).thenReturn(firstPage);
116+
JobsExt jobsExt = new JobsExt(service);
117+
118+
GetJobRequest request = new GetJobRequest();
119+
Job job = jobsExt.get(request);
120+
121+
Job expectedJob =
122+
new Job().setSettings(new JobSettings().setTasks(tasks).setJobClusters(jobClusters));
123+
assertEquals(expectedJob, job);
124+
verify(service, times(1)).get(any());
125+
}
126+
127+
@Test
128+
public void testGetJobPaginationWithJobClusters() {
129+
JobsService service = Mockito.mock(JobsService.class);
130+
Job firstPage = new Job().setSettings(new JobSettings()).setNextPageToken("tokenToSecondPage");
131+
addTasks(firstPage, "taskKey1", "taskKey2");
132+
addJobClusters(firstPage, "clusterKey1", "clusterKey2");
133+
addJobParameters(firstPage, "parameterKey1");
134+
addJobEnvironments(firstPage, "environmentKey1");
135+
136+
Job secondPage = new Job().setSettings(new JobSettings()).setNextPageToken("tokenToThirdPage");
137+
addTasks(secondPage, "taskKey3", "taskKey4");
138+
addJobClusters(secondPage, "clusterKey3");
139+
addJobParameters(secondPage, "parameterKey2");
140+
addJobEnvironments(secondPage, "environmentKey2");
141+
142+
Job thirdPage = new Job().setSettings(new JobSettings());
143+
addTasks(thirdPage, "taskKey5");
144+
addJobParameters(thirdPage, "parameterKey3");
145+
146+
when(service.get(any())).thenReturn(firstPage).thenReturn(secondPage).thenReturn(thirdPage);
147+
JobsExt jobsExt = new JobsExt(service);
148+
149+
GetJobRequest request = new GetJobRequest();
150+
Job job = jobsExt.get(request);
151+
152+
Job expectedJob = new Job().setSettings(new JobSettings());
153+
addTasks(expectedJob, "taskKey1", "taskKey2", "taskKey3", "taskKey4", "taskKey5");
154+
addJobClusters(expectedJob, "clusterKey1", "clusterKey2", "clusterKey3");
155+
addJobParameters(expectedJob, "parameterKey1", "parameterKey2", "parameterKey3");
156+
addJobEnvironments(expectedJob, "environmentKey1", "environmentKey2");
157+
assertEquals(expectedJob, job);
158+
verify(service, times(3)).get(any());
159+
}
160+
161+
private void addTasks(Job job, String... taskKeys) {
162+
Collection<Task> tasks = new ArrayList<>();
163+
for (String taskKey : taskKeys) {
164+
tasks.add(new Task().setTaskKey(taskKey));
165+
}
166+
job.getSettings().setTasks(tasks);
167+
}
168+
169+
private void addJobClusters(Job job, String... clusterKeys) {
170+
Collection<JobCluster> jobClusters = new ArrayList<>();
171+
for (String clusterKey : clusterKeys) {
172+
jobClusters.add(new JobCluster().setJobClusterKey(clusterKey));
173+
}
174+
job.getSettings().setJobClusters(jobClusters);
175+
}
176+
177+
private void addJobParameters(Job job, String... parameterKeys) {
178+
Collection<JobParameterDefinition> parameters = new ArrayList<>();
179+
for (String parameterKey : parameterKeys) {
180+
parameters.add(new JobParameterDefinition().setName(parameterKey));
181+
}
182+
job.getSettings().setParameters(parameters);
183+
}
184+
185+
private void addJobEnvironments(Job job, String... environmentKeys) {
186+
Collection<JobEnvironment> environments = new ArrayList<>();
187+
for (String environmentKey : environmentKeys) {
188+
environments.add(new JobEnvironment().setEnvironmentKey(environmentKey));
189+
}
190+
job.getSettings().setEnvironments(environments);
191+
}
85192
}

0 commit comments

Comments
 (0)