How to use org.apache.gobblin.rest

protected String getLauncherType(JobExecutionInfo info) {
 if (info.hasLauncherType()) {
  if (info.getLauncherType() == LauncherTypeEnum.CLUSTER) {
   return LauncherTypeEnum.YARN.name();
  }
  return info.getLauncherType().name();
 }
 return null;
}

public JobExecutionInfoServer(Properties properties) {
 this.properties = properties;
 port = getPort(properties);
 serverUri = getServiceUri(getHost(properties), port);
 serverAdvertisedUri = getAdvertisedUri(properties);
}

/**
 * Execute a query and coerce the result into a java List
 * @param query Query to execute
 * @return List of jobs that matched the query. (Empty list if none did).
 * @throws RemoteInvocationException If the server throws an error
 */
private List<JobExecutionInfo> executeQuery(JobExecutionQuery query) throws RemoteInvocationException {
 JobExecutionQueryResult result = this.client.get(query);
 if (result != null && result.hasJobExecutions()) {
  return result.getJobExecutions();
 }
 return Collections.emptyList();
}

private static void assertJobExecution(JobExecutionInfo actual, JobExecutionInfo expected) {
 Assert.assertEquals(actual.getJobName(), expected.getJobName());
 Assert.assertEquals(actual.getJobId(), expected.getJobId());
 if (expected.hasDuration()) {
  Assert.assertEquals(actual.getDuration(), expected.getDuration());
 } else {
  Assert.assertEquals(actual.getDuration().longValue(), -1L);
 }
 Assert.assertEquals(actual.getState(), expected.getState());
 Assert.assertEquals(actual.getLaunchedTasks(), expected.getLaunchedTasks());
 Assert.assertEquals(actual.getCompletedTasks(), expected.getCompletedTasks());
 Assert.assertEquals(actual.getMetrics(), expected.getMetrics());
 for (int i = 0; i < actual.getMetrics().size(); i++) {
  assertMetric(actual.getMetrics().get(i), expected.getMetrics().get(i));
 }
 Assert.assertEquals(actual.getJobProperties(), expected.getJobProperties());
 Assert.assertEquals(actual.getTaskExecutions().size(), expected.getTaskExecutions().size());
 for (int i = 0; i < actual.getTaskExecutions().size(); i++) {
  assertTaskExecution(actual.getTaskExecutions().get(i), expected.getTaskExecutions().get(i));
 }
}

private static void assertTaskExecution(TaskExecutionInfo actual, TaskExecutionInfo expected) {
 Assert.assertEquals(actual.getJobId(), expected.getJobId());
 Assert.assertEquals(actual.getTaskId(), expected.getTaskId());
 if (expected.hasDuration()) {
  Assert.assertEquals(actual.getDuration(), expected.getDuration());
 } else {
  Assert.assertEquals(actual.getDuration().longValue(), -1L);
 }
 Assert.assertEquals(actual.getState(), expected.getState());
 Assert.assertEquals(actual.getLowWatermark(), expected.getLowWatermark());
 Assert.assertEquals(actual.getHighWatermark(), expected.getHighWatermark());
 Assert.assertEquals(actual.getTable(), expected.getTable());
 Assert.assertEquals(actual.getMetrics(), expected.getMetrics());
 for (int i = 0; i < actual.getMetrics().size(); i++) {
  assertMetric(actual.getMetrics().get(i), expected.getMetrics().get(i));
 }
 Assert.assertEquals(actual.getTaskProperties(), expected.getTaskProperties());
}

private Metric resultSetToMetric(ResultSet rs)
  throws SQLException {
 Metric metric = new Metric();
 metric.setGroup(rs.getString("metric_group"));
 metric.setName(rs.getString("metric_name"));
 metric.setType(MetricTypeEnum.valueOf(rs.getString("metric_type")));
 metric.setValue(rs.getString("metric_value"));
 return metric;
}

@Test
public void testBatchGet() throws Exception {
 JobExecutionQuery queryByJobId1 = new JobExecutionQuery();
 queryByJobId1.setIdType(QueryIdTypeEnum.JOB_ID);
 queryByJobId1.setId(JobExecutionQuery.Id.create(this.expected1.getJobId()));
 JobExecutionQuery queryByJobId2 = new JobExecutionQuery();
 queryByJobId2.setIdType(QueryIdTypeEnum.JOB_ID);
 queryByJobId2.setId(JobExecutionQuery.Id.create(this.expected2.getJobId()));
 List<JobExecutionQuery> queries = Lists.newArrayList(queryByJobId1, queryByJobId2);
 List<JobExecutionQueryResult> result = Lists.newArrayList(this.client.batchGet(queries));
 Assert.assertEquals(result.size(), 2);
 Assert.assertEquals(result.get(0).getJobExecutions().size(), 1);
 Assert.assertEquals(result.get(1).getJobExecutions().size(), 1);
 JobExecutionInfo actual1 = result.get(0).getJobExecutions().get(0);
 JobExecutionInfo actual2 = result.get(1).getJobExecutions().get(0);
 if (actual1.getJobName().equals(this.expected1.getJobName())) {
  assertJobExecution(actual1, this.expected1);
  assertJobExecution(actual2, this.expected2);
 } else {
  assertJobExecution(actual1, this.expected2);
  assertJobExecution(actual2, this.expected1);
 }
}

@Test
public void testGet() throws Exception {
 JobExecutionQuery queryByJobId = new JobExecutionQuery();
 queryByJobId.setIdType(QueryIdTypeEnum.JOB_ID);
 queryByJobId.setId(JobExecutionQuery.Id.create(this.expected1.getJobId()));
 JobExecutionQueryResult result = this.client.get(queryByJobId);
 JobExecutionInfoArray jobExecutionInfos = result.getJobExecutions();
 Assert.assertEquals(jobExecutionInfos.size(), 1);
 JobExecutionInfo actual = jobExecutionInfos.get(0);
 assertJobExecution(actual, this.expected1);
}

/**
 * Query jobs by name
 *
 * @param name         Name of the job to query for
 * @param resultsLimit Max # of results to return
 * @return List of jobs with the name (empty list if none can be found)
 */
public List<JobExecutionInfo> queryByJobName(String name, int resultsLimit) throws RemoteInvocationException {
 JobExecutionQuery query = new JobExecutionQuery();
 query.setIdType(QueryIdTypeEnum.JOB_NAME);
 query.setId(JobExecutionQuery.Id.create(name));
 query.setIncludeTaskExecutions(false);
 query.setLimit(resultsLimit);
 return executeQuery(query);
}

private void addMetricToBatch(PreparedStatement upsertStatement, Metric metric, String id) throws SQLException {
 Preconditions.checkArgument(!Strings.isNullOrEmpty(id));
 Preconditions.checkArgument(metric.hasGroup());
 Preconditions.checkArgument(metric.hasName());
 Preconditions.checkArgument(metric.hasType());
 Preconditions.checkArgument(metric.hasValue());
 int index = 0;
 upsertStatement.setString(++index, id);
 upsertStatement.setString(++index, metric.getGroup());
 upsertStatement.setString(++index, metric.getName());
 upsertStatement.setString(++index, metric.getType().name());
 upsertStatement.setString(++index, metric.getValue());
 upsertStatement.addBatch();
}

@Test
public void testBadGet() throws Exception {
 JobExecutionQuery queryByJobId = new JobExecutionQuery();
 queryByJobId.setIdType(QueryIdTypeEnum.JOB_ID);
 queryByJobId.setId(JobExecutionQuery.Id.create(this.expected1.getJobId() + "1"));
 JobExecutionQueryResult result = this.client.get(queryByJobId);
 Assert.assertTrue(result.getJobExecutions().isEmpty());
}

 private void assertMetric(Metric actual, Metric expected) {
  Assert.assertEquals(actual.getGroup(), expected.getGroup());
  Assert.assertEquals(actual.getName(), expected.getName());
  Assert.assertEquals(actual.getType(), expected.getType());
  Assert.assertEquals(actual.getValue(), expected.getValue());
 }
}

@Test(dependsOnMethods = {"testUpdate"})
public void testQueryByJobId()
  throws IOException {
 JobExecutionQuery queryByJobId = new JobExecutionQuery();
 queryByJobId.setIdType(QueryIdTypeEnum.JOB_ID);
 queryByJobId.setId(JobExecutionQuery.Id.create(this.expectedJobExecutionInfos.get(0).getJobId()));
 List<JobExecutionInfo> result = this.jobHistoryStore.get(queryByJobId);
 Assert.assertEquals(result.size(), 1);
 JobExecutionInfo actual = result.get(0);
 JobExecutionInfo expected = this.expectedJobExecutionInfos.get(0);
 assertJobExecution(actual, expected);
}

@Test(dependsOnMethods = {"testUpdate"})
public void testQueryByJobName()
  throws IOException {
 JobExecutionQuery queryByJobName = new JobExecutionQuery();
 queryByJobName.setIdType(QueryIdTypeEnum.JOB_NAME);
 queryByJobName.setId(JobExecutionQuery.Id.create(this.expectedJobExecutionInfos.get(0).getJobName()));
 List<JobExecutionInfo> result = this.jobHistoryStore.get(queryByJobName);
 Assert.assertEquals(result.size(), 1);
 JobExecutionInfo actual = result.get(0);
 JobExecutionInfo expected = this.expectedJobExecutionInfos.get(0);
 assertJobExecution(actual, expected);
}

/**
 * Sets the {@link LauncherTypeEnum} for this {@link JobState}.
 */
public void setJobLauncherType(LauncherTypeEnum jobLauncherType) {
 this.setProp(ConfigurationKeys.JOB_LAUNCHER_TYPE_KEY, jobLauncherType.name());
}

private static Metric resultSetToMetric(ResultSet rs) throws SQLException {
 Metric metric = new Metric();
 metric.setGroup(rs.getString("metric_group"));
 metric.setName(rs.getString("metric_name"));
 metric.setType(MetricTypeEnum.valueOf(rs.getString("metric_type")));
 metric.setValue(rs.getString("metric_value"));
 return metric;
}

/**
 * Retrieve all jobs
 *
 * @param lookupType Query type
 * @return List of all jobs (limited by results limit)
 */
public List<JobExecutionInfo> queryAllJobs(QueryListType lookupType, int resultsLimit)
  throws RemoteInvocationException {
 JobExecutionQuery query = new JobExecutionQuery();
 query.setIdType(QueryIdTypeEnum.LIST_TYPE);
 query.setId(JobExecutionQuery.Id.create(lookupType));
 // Disable properties and task executions (prevents response size from ballooning)
 query.setJobProperties(ConfigurationKeys.JOB_RUN_ONCE_KEY + "," + ConfigurationKeys.JOB_SCHEDULE_KEY);
 query.setIncludeTaskExecutions(false);
 query.setLimit(resultsLimit);
 return executeQuery(query);
}

 private static void assertMetric(Metric actual, Metric expected) {
  Assert.assertEquals(actual.getGroup(), expected.getGroup());
  Assert.assertEquals(actual.getName(), expected.getName());
  Assert.assertEquals(actual.getType(), expected.getType());
  Assert.assertEquals(actual.getValue(), expected.getValue());
 }
}

 @Override
 protected String getLauncherType(JobExecutionInfo info) {
  if (info.hasLauncherType()) {
   return info.getLauncherType().name();
  }
  return null;
 }
}

/**
 * Retrieve a Gobblin job by its id.
 *
 * @param id                Id of the job to retrieve
 * @return JobExecutionInfo representing the job
 */
public Optional<JobExecutionInfo> queryByJobId(String id) throws RemoteInvocationException {
 JobExecutionQuery query = new JobExecutionQuery();
 query.setIdType(QueryIdTypeEnum.JOB_ID);
 query.setId(JobExecutionQuery.Id.create(id));
 query.setLimit(1);
 List<JobExecutionInfo> results = executeQuery(query);
 return getFirstFromQueryResults(results);
}