content: /**
   * 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.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  package org.apache.aurora.scheduler.base;
  
  import java.util.EnumSet;
  import java.util.Objects;
  
  import com.google.common.base.MoreObjects;
  import com.google.common.base.Supplier;
  import com.google.common.collect.ImmutableSet;
  import com.google.common.primitives.Ints;
  
  import org.apache.aurora.gen.ScheduleStatus;
  import org.apache.aurora.gen.TaskQuery;
  import org.apache.aurora.scheduler.storage.entities.IInstanceKey;
  import org.apache.aurora.scheduler.storage.entities.IJobKey;
  import org.apache.aurora.scheduler.storage.entities.ITaskQuery;
  
  import static java.util.Objects.requireNonNull;
  
  /**
   * A utility class to construct storage queries.
   * TODO(Sathya): Add some basic unit tests for isJobScoped and isOnlyJobScoped.
   */
  public final class Query {
  
    private Query() {
      // Utility.
    }
  
    /**
     * Checks whether a query is scoped to a specific job.
     * A query scoped to a job specifies either:
     * <ol>
     *   <li>a role, environment and job name, or</li>
     *   <li>a set of JobKey instances</li>
     * </ol>
     *
     * @param taskQuery Query to test.
     * @return {@code true} if the query specifies at least one job key, otherwise {@code false}.
     */
    public static boolean isJobScoped(Builder taskQuery) {
      ITaskQuery q = taskQuery.get();
      return q.isSetRole() && q.isSetEnvironment() && q.isSetJobName() || !q.getJobKeys().isEmpty();
    }
  
    public static Builder arbitrary(TaskQuery query) {
      return new Builder(query.deepCopy());
    }
  
    public static Builder unscoped() {
      return new Builder();
    }
  
    public static Builder roleScoped(String role) {
      return unscoped().byRole(role);
    }
  
    public static Builder envScoped(String role, String environment) {
      return unscoped().byEnv(role, environment);
    }
  
    public static Builder jobScoped(IJobKey jobKey) {
      return unscoped().byJob(jobKey);
    }
  
    public static Builder jobScoped(Iterable<IJobKey> jobKeys) {
      return unscoped().byJobKeys(jobKeys);
    }
  
    public static Builder instanceScoped(IInstanceKey instanceKey) {
      return instanceScoped(instanceKey.getJobKey(), instanceKey.getInstanceId());
    }
  
    public static Builder instanceScoped(IJobKey jobKey, int instanceId, int... instanceIds) {
      return unscoped().byInstances(jobKey, instanceId, instanceIds);
    }
  
    public static Builder instanceScoped(IJobKey jobKey, Iterable<Integer> instanceIds) {
      return unscoped().byInstances(jobKey, instanceIds);
    }
  
    public static Builder taskScoped(String taskId, String... taskIds) {
      return unscoped().byId(taskId, taskIds);
    }
  
    public static Builder taskScoped(Iterable<String> taskIds) {
      return unscoped().byId(taskIds);
    }
  
    public static Builder slaveScoped(String slaveHost, String... slaveHosts) {
      return unscoped().bySlave(slaveHost, slaveHosts);
    }
  
    public static Builder slaveScoped(Iterable<String> slaveHosts) {
      return unscoped().bySlave(slaveHosts);
    }
  
    public static Builder statusScoped(ScheduleStatus status, ScheduleStatus... statuses) {
      return unscoped().byStatus(status, statuses);
    }
  
    /**
     * A Builder of TaskQueries. Builders are immutable and provide access to a set of convenience
     * methods to return a new builder of another scope. Available scope filters include slave,
     * taskId, role, jobs of a role, and instances of a job.
     *
     * <p>
     * This class does not expose the full functionality of TaskQuery but rather subsets of it that
     * can be efficiently executed and make sense in the context of the scheduler datastores. This
     * builder should be preferred over constructing TaskQueries directly.
     * </p>
     *
     * TODO(ksweeney): Add an environment scope.
     */
    public static final class Builder implements Supplier<ITaskQuery> {
      private final TaskQuery query;
  
      Builder() {
        this(new TaskQuery());
      }
  
      Builder(TaskQuery query) {
        // It is expected that the caller calls deepCopy.
        this.query = query;
      }
  
      /**
       * Build a query that is the combination of all the filters applied to a Builder. Mutating the
       * returned object will not affect the state of the builder. Can be called any number of times
       * and will return a new {@code TaskQuery} each time.
       *
       * @return A new TaskQuery satisfying this builder's constraints.
       */
      @Override
      public ITaskQuery get() {
        return ITaskQuery.build(query);
      }
  
      @Override
      public boolean equals(Object that) {
        return that instanceof Builder && get().equals(((Builder) that).get());
      }
  
      @Override
      public int hashCode() {
        return Objects.hash(get());
      }
  
      @Override
      public String toString() {
        return MoreObjects.toStringHelper(this)
            .add("query", query)
            .toString();
      }
  
      /**
       * Create a builder scoped to tasks.
       *
       * @param taskId An ID of a task to scope the builder to.
       * @param taskIds Additional IDs of tasks to scope the builder to (they are ORed together).
       * @return A new Builder scoped to the given tasks.
       */
      public Builder byId(String taskId, String... taskIds) {
        requireNonNull(taskId);
  
        return new Builder(
            query.deepCopy()
                .setTaskIds(ImmutableSet.<String>builder().add(taskId).add(taskIds).build()));
      }
  
      /**
       * Create a builder scoped to tasks.
       *
       * @see #byId(String, String...)
       *
       * @param taskIds The IDs of the tasks to scope the query to (ORed together).
       * @return A new Builder scoped to the given tasks.
       */
      public Builder byId(Iterable<String> taskIds) {
        requireNonNull(taskIds);
  
        return new Builder(
            query.deepCopy().setTaskIds(ImmutableSet.copyOf(taskIds)));
      }
  
      /**
       * Create a builder scoped to a role. A role scope conflicts with job and instance scopes.
       *
       * @param role The role to scope the query to.
       * @return A new Builder scoped to the given role.
       */
      public Builder byRole(String role) {
        requireNonNull(role);
  
        return new Builder(query.deepCopy().setRole(role));
      }
  
      /**
       * Create a builder scoped to an environment. An environment scope conflicts with role, job,
       * and instance scopes.
       *
       * @param role The role to scope the query to.
       * @param environment The environment to scope the query to.
       * @return A new Builder scoped to the given environment.
       */
      public Builder byEnv(String role, String environment) {
        requireNonNull(role);
        requireNonNull(environment);
  
        return new Builder(
            query.deepCopy().setRole(role).setEnvironment(environment));
      }
  
      /**
       * Returns a new builder scoped to the job uniquely identified by the given key. A job scope
       * conflicts with role and instance scopes.
       *
       * @param jobKey The key of the job to scope the query to.
       * @return A new Builder scoped to the given jobKey.
       */
      public Builder byJob(IJobKey jobKey) {
        JobKeys.assertValid(jobKey);
  
        return new Builder(
            query.deepCopy()
                .setJobKeys(ImmutableSet.of(jobKey.newBuilder())));
      }
  
      /**
       * Returns a new builder scoped to the slaves uniquely identified by the given slaveHosts. A
       * builder can only be scoped to slaves once.
       *
       * @param slaveHost The hostname of the slave to scope the query to.
       * @param slaveHosts Additional slave hostnames to scope this query to (they are ORed together).
       * @return A new Builder scoped to the given slaves.
       */
      public Builder bySlave(String slaveHost, String... slaveHosts) {
        requireNonNull(slaveHost);
  
        return bySlave(ImmutableSet.<String>builder().add(slaveHost).add(slaveHosts).build());
      }
  
      /**
       * Creates a new builder scoped to slaveHosts.
       *
       * @see Builder#bySlave(String, String...)
       *
       * @param slaveHosts The slaveHost to scope this builder to.
       * @return A new Builder scoped to the slaveHosts.
       */
      public Builder bySlave(Iterable<String> slaveHosts) {
        requireNonNull(slaveHosts);
  
        return new Builder(query.deepCopy().setSlaveHosts(ImmutableSet.copyOf(slaveHosts)));
      }
  
      /**
       * Returns a new builder scoped to the given statuses. A builder can only be scoped to statuses
       * once.
       *
       * @param status The status to scope this Builder to.
       * @param statuses Additional statuses to scope this Builder to (they are ORed together).
       * @return A new Builder scoped to the given statuses.
       */
      public Builder byStatus(ScheduleStatus status, ScheduleStatus... statuses) {
        requireNonNull(status);
  
        return new Builder(
            query.deepCopy().setStatuses(EnumSet.of(status, statuses)));
      }
  
      /**
       * Create a new Builder scoped to statuses.
       *
       * @see Builder#byStatus(ScheduleStatus, ScheduleStatus...)
       *
       * @param statuses The statuses to scope this Builder to.
       * @return A new Builder scoped to the given statuses.
       */
      public Builder byStatus(Iterable<ScheduleStatus> statuses) {
        requireNonNull(statuses);
  
        return new Builder(
            query.deepCopy().setStatuses(EnumSet.copyOf(ImmutableSet.copyOf(statuses))));
      }
  
      /**
       * Returns a new Builder scoped to the given instances of the given job. A builder can only
       * be scoped to a set of instances, a job, or a role once.
       *
       * @param jobKey The key identifying the job.
       * @param instanceId An instance id of the target job.
       * @param instanceIds Additional instance ids of the target job.
       * @return A new Builder scoped to the given instance ids.
       */
      public Builder byInstances(IJobKey jobKey, int instanceId, int... instanceIds) {
        JobKeys.assertValid(jobKey);
  
        return new Builder(
            query.deepCopy()
                .setRole(jobKey.getRole())
                .setEnvironment(jobKey.getEnvironment())
                .setJobName(jobKey.getName())
                .setInstanceIds(ImmutableSet.<Integer>builder()
                    .add(instanceId)
                    .addAll(Ints.asList(instanceIds))
                    .build()));
      }
  
      /**
       * Create a new Builder scoped to instances.
       *
       * @see Builder#byInstances
       *
       * @param jobKey The key identifying the job.
       * @param instanceIds Instances of the target job.
       * @return A new Builder scoped to the given instance ids.
       */
      public Builder byInstances(IJobKey jobKey, Iterable<Integer> instanceIds) {
        JobKeys.assertValid(jobKey);
        requireNonNull(instanceIds);
  
        return new Builder(
            query.deepCopy()
                .setRole(jobKey.getRole())
                .setEnvironment(jobKey.getEnvironment())
                .setJobName(jobKey.getName())
                .setInstanceIds(ImmutableSet.copyOf(instanceIds)));
      }
  
      /**
       * Returns a new builder scoped to the jobs uniquely identified by the given jobKeys. A
       * builder can only be scoped to jobs once.
       *
       * @param jobKeys The job keys to scope this builder to.
       * @return A new Builder scoped to the job keys.
       */
      public Builder byJobKeys(Iterable<IJobKey> jobKeys) {
        requireNonNull(jobKeys);
  
        return new Builder(query.deepCopy().setJobKeys(IJobKey.toBuildersSet(jobKeys)));
      }
  
      /**
       * A convenience method to scope this builder to {@link Tasks#ACTIVE_STATES}.
       *
       * @return A new Builder scoped to Tasks#ACTIVE_STATES.
       */
      public Builder active() {
        return byStatus(Tasks.ACTIVE_STATES);
      }
  
      /**
       * A convenience method to scope this builder to {@link Tasks#TERMINAL_STATES}.
       *
       * @return A new Builder scoped to Tasks#TERMINAL_STATES.
       */
      public Builder terminal() {
        return byStatus(Tasks.TERMINAL_STATES);
      }
    }
  }
  