gobblin.cli.AdminClient Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of gobblin-admin Show documentation
Gobblin Ingestion Framework
The newest version!
/* (c) 2015 NerdWallet All rights reserved.
*
* Licensed under the Apache License, Version 2.0 (the "License"); you may not use
* this file except in compliance with the License. You may obtain a copy of the
* License at  http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software distributed
* under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
* CONDITIONS OF ANY KIND, either express or implied.
*/
package gobblin.cli;

import com.google.common.base.Optional;
import com.google.common.io.Closer;
import com.linkedin.r2.RemoteInvocationException;
import gobblin.configuration.ConfigurationKeys;
import gobblin.rest.*;

import java.io.IOException;
import java.net.URI;
import java.util.Collections;
import java.util.List;


/**
 * Simple wrapper around the JobExecutionInfoClient
 */
public class AdminClient {
  private final JobExecutionInfoClient client;
  private Closer closer;

  /**
   * Creates a new client with the host and port specified.
   */
  public AdminClient(String host, int port) {
    this.closer = Closer.create();

    URI serverUri = URI.create(String.format("http://%s:%d/", host, port));
    this.client = new JobExecutionInfoClient(serverUri.toString());
    this.closer.register(this.client);
  }

  /**
   * Close connections to the REST server
   */
  public void close() {
    try {
      this.closer.close();
    } catch (IOException e) {
      e.printStackTrace();
    }
  }

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

    List results = executeQuery(query);
    return getFirstFromQueryResults(results);
  }

  /**
   * Retrieve all jobs
   *
   * @param lookupType Query type
   * @return List of all jobs (limited by results limit)
   */
  public List 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);
  }

  /**
   * 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 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);
  }

  /**
   * 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 executeQuery(JobExecutionQuery query) throws RemoteInvocationException {
    JobExecutionQueryResult result = this.client.get(query);

    if (result != null && result.hasJobExecutions()) {
      return result.getJobExecutions();
    }
    return Collections.emptyList();
  }

  private static Optional getFirstFromQueryResults(List results) {
    if (results == null || results.size() == 0) {
      return Optional.absent();
    }

    return Optional.of(results.get(0));
  }

}