QueryBuilders.java

/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to you 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.calcite.adapter.elasticsearch;

import com.fasterxml.jackson.core.JsonGenerator;

import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import static java.util.Objects.requireNonNull;

/**
 * Utility class to generate elastic search queries. Most query builders have
 * been copied from ES distribution. The reason we have separate definition is
 * high-level client dependency on core modules (like lucene, netty, XContent etc.) which
 * is not compatible between different major versions.
 *
 * <p>The goal of ES adapter is to
 * be compatible with any elastic version or even to connect to clusters with different
 * versions simultaneously.
 *
 * <p>Jackson API is used to generate ES query as JSON document.
 */
class QueryBuilders {

  private QueryBuilders() {}

  /**
   * A Query that matches documents containing a term.
   *
   * @param name  The name of the field
   * @param value The value of the term
   */
  static TermQueryBuilder termQuery(String name, String value) {
    return new TermQueryBuilder(name, value);
  }

  /**
   * A Query that matches documents containing a term.
   *
   * @param name  The name of the field
   * @param value The value of the term
   */
  static TermQueryBuilder termQuery(String name, int value) {
    return new TermQueryBuilder(name, value);
  }

  /**
   * A Query that matches documents containing a single character term.
   *
   * @param name  The name of the field
   * @param value The value of the term
   */
  static TermQueryBuilder termQuery(String name, char value) {
    return new TermQueryBuilder(name, value);
  }

  /**
   * A Query that matches documents containing a term.
   *
   * @param name  The name of the field
   * @param value The value of the term
   */
  static TermQueryBuilder termQuery(String name, long value) {
    return new TermQueryBuilder(name, value);
  }

  /**
   * A Query that matches documents containing a term.
   *
   * @param name  The name of the field
   * @param value The value of the term
   */
  static TermQueryBuilder termQuery(String name, float value) {
    return new TermQueryBuilder(name, value);
  }

  /**
   * A Query that matches documents containing a term.
   *
   * @param name  The name of the field
   * @param value The value of the term
   */
  static TermQueryBuilder termQuery(String name, double value) {
    return new TermQueryBuilder(name, value);
  }

  /**
   * A Query that matches documents containing a term.
   *
   * @param name  The name of the field
   * @param value The value of the term
   */
  static TermQueryBuilder termQuery(String name, boolean value) {
    return new TermQueryBuilder(name, value);
  }

  /**
   * A Query that matches documents containing a term.
   *
   * @param name  The name of the field
   * @param value The value of the term
   */
  static TermQueryBuilder termQuery(String name, Object value) {
    return new TermQueryBuilder(name, value);
  }

  /**
   * A filer for a field based on several terms matching on any of them.
   *
   * @param name   The field name
   * @param values The terms
   */
  static MatchesQueryBuilder matchesQuery(String name, Iterable<?> values) {
    return new MatchesQueryBuilder(name, values);
  }

  /**
   * A Query that matches documents containing a term.
   *
   * @param name  The name of the field
   * @param value The value of the term
   */
  static MatchQueryBuilder matchQuery(String name, Object value) {
    return new MatchQueryBuilder(name, value);
  }

  /**
   * A filer for a field based on several terms matching on any of them.
   *
   * @param name   The field name
   * @param values The terms
   */
  static TermsQueryBuilder termsQuery(String name, Iterable<?> values) {
    return new TermsQueryBuilder(name, values);
  }

  /**
   * A Query that matches documents within an range of terms.
   *
   * @param name The field name
   */
  static RangeQueryBuilder rangeQuery(String name) {
    return new RangeQueryBuilder(name);
  }

  /**
   * A Query that matches documents containing terms with a specified regular expression.
   *
   * @param name   The name of the field
   * @param regexp The regular expression
   */
  static RegexpQueryBuilder regexpQuery(String name, String regexp) {
    return new RegexpQueryBuilder(name, regexp);
  }

  /**
   * A Query that matches documents containing terms with a specified regular expression.
   *
   * @param name   The name of the field
   * @param regexp The regular expression
   * @param escape The regular escape
   */
  static RegexpQueryBuilder regexpQuery(String name, String regexp, String escape) {
    return new RegexpQueryBuilder(name, regexp, escape);
  }


  /**
   * A Query that matches documents matching boolean combinations of other queries.
   */
  static BoolQueryBuilder boolQuery() {
    return new BoolQueryBuilder();
  }

  /**
   * A query that wraps another query and simply returns a constant score equal to the
   * query boost for every document in the query.
   *
   * @param queryBuilder The query to wrap in a constant score query
   */
  static ConstantScoreQueryBuilder constantScoreQuery(QueryBuilder queryBuilder) {
    return new ConstantScoreQueryBuilder(queryBuilder);
  }

  /**
   * A query that wraps another query and simply returns a dismax score equal to the
   * query boost for every document in the query.
   *
   * @param queryBuilder The query to wrap in a constant score query
   */
  static DisMaxQueryBuilder disMaxQueryBuilder(QueryBuilder queryBuilder) {
    return new DisMaxQueryBuilder(queryBuilder);
  }

  /**
   * A filter to filter only documents where a field exists in them.
   *
   * @param name The name of the field
   */
  static ExistsQueryBuilder existsQuery(String name) {
    return new ExistsQueryBuilder(name);
  }

  /**
   * A query that matches on all documents.
   */
  static MatchAllQueryBuilder matchAll() {
    return new MatchAllQueryBuilder();
  }

  /**
   * Base class to build Elasticsearch queries.
   */
  abstract static class QueryBuilder {

    /**
     * Converts an existing query to JSON format using jackson API.
     *
     * @param generator used to generate JSON elements
     * @throws IOException if IO error occurred
     */
    abstract void writeJson(JsonGenerator generator) throws IOException;
  }

  /**
   * Query for boolean logic.
   */
  static class BoolQueryBuilder extends QueryBuilder {
    private final List<QueryBuilder> mustClauses = new ArrayList<>();
    private final List<QueryBuilder> mustNotClauses = new ArrayList<>();
    private final List<QueryBuilder> filterClauses = new ArrayList<>();
    private final List<QueryBuilder> shouldClauses = new ArrayList<>();

    BoolQueryBuilder must(QueryBuilder queryBuilder) {
      requireNonNull(queryBuilder, "queryBuilder");
      mustClauses.add(queryBuilder);
      return this;
    }

    BoolQueryBuilder filter(QueryBuilder queryBuilder) {
      requireNonNull(queryBuilder, "queryBuilder");
      filterClauses.add(queryBuilder);
      return this;
    }

    BoolQueryBuilder mustNot(QueryBuilder queryBuilder) {
      requireNonNull(queryBuilder, "queryBuilder");
      mustNotClauses.add(queryBuilder);
      return this;
    }

    BoolQueryBuilder should(QueryBuilder queryBuilder) {
      requireNonNull(queryBuilder, "queryBuilder");
      shouldClauses.add(queryBuilder);
      return this;
    }

    @Override protected void writeJson(JsonGenerator gen) throws IOException {
      gen.writeStartObject();
      gen.writeFieldName("bool");
      gen.writeStartObject();
      writeJsonArray("must", mustClauses, gen);
      writeJsonArray("filter", filterClauses, gen);
      writeJsonArray("must_not", mustNotClauses, gen);
      writeJsonArray("should", shouldClauses, gen);
      gen.writeEndObject();
      gen.writeEndObject();
    }

    private static void writeJsonArray(String field, List<QueryBuilder> clauses, JsonGenerator gen)
        throws IOException {
      if (clauses.isEmpty()) {
        return;
      }

      if (clauses.size() == 1) {
        gen.writeFieldName(field);
        clauses.get(0).writeJson(gen);
      } else {
        gen.writeArrayFieldStart(field);
        for (QueryBuilder clause : clauses) {
          clause.writeJson(gen);
        }
        gen.writeEndArray();
      }
    }
  }

  /**
   * A Query that matches documents containing a term.
   */
  static class TermQueryBuilder extends QueryBuilder {
    private final String fieldName;
    private final Object value;

    private TermQueryBuilder(final String fieldName, final Object value) {
      this.fieldName = requireNonNull(fieldName, "fieldName");
      this.value = requireNonNull(value, "value");
    }

    @Override void writeJson(final JsonGenerator generator) throws IOException {
      generator.writeStartObject();
      generator.writeFieldName("term");
      generator.writeStartObject();
      generator.writeFieldName(fieldName);
      writeObject(generator, value);
      generator.writeEndObject();
      generator.writeEndObject();
    }
  }

  /**
   * A filter for a field based on several terms matching on any of them.
   */
  private static class TermsQueryBuilder extends QueryBuilder {
    private final String fieldName;
    private final Iterable<?> values;

    private TermsQueryBuilder(final String fieldName, final Iterable<?> values) {
      this.fieldName = requireNonNull(fieldName, "fieldName");
      this.values = requireNonNull(values, "values");
    }

    @Override void writeJson(final JsonGenerator generator) throws IOException {
      generator.writeStartObject();
      generator.writeFieldName("terms");
      generator.writeStartObject();
      generator.writeFieldName(fieldName);
      generator.writeStartArray();
      for (Object value : values) {
        writeObject(generator, value);
      }
      generator.writeEndArray();
      generator.writeEndObject();
      generator.writeEndObject();
    }
  }



  /**
   * A Query that matches documents containing a term.
   */
  static class MatchQueryBuilder extends QueryBuilder {
    private final String fieldName;
    private final Object value;

    private MatchQueryBuilder(final String fieldName, final Object value) {
      this.fieldName = requireNonNull(fieldName, "fieldName");
      this.value = requireNonNull(value, "value");
    }

    @Override void writeJson(final JsonGenerator generator) throws IOException {
      generator.writeStartObject();
      generator.writeFieldName("match");
      generator.writeStartObject();
      generator.writeFieldName(fieldName);
      writeObject(generator, value);
      generator.writeEndObject();
      generator.writeEndObject();
    }
  }


  /**
   * A filter for a field based on several terms matching on any of them.
   */
  private static class MatchesQueryBuilder extends QueryBuilder {
    private final String fieldName;
    private final Iterable<?> values;

    private MatchesQueryBuilder(final String fieldName, final Iterable<?> values) {
      this.fieldName = requireNonNull(fieldName, "fieldName");
      this.values = requireNonNull(values, "values");
    }

    @Override void writeJson(final JsonGenerator generator) throws IOException {
      generator.writeStartObject();
      generator.writeFieldName("match");
      generator.writeStartObject();
      generator.writeFieldName(fieldName);
      generator.writeStartArray();
      for (Object value : values) {
        writeObject(generator, value);
      }
      generator.writeEndArray();
      generator.writeEndObject();
      generator.writeEndObject();
    }
  }

  /**
   * Write usually simple (scalar) value (string, number, boolean or null) to json output.
   * In case of complex objects delegates to jackson serialization.
   *
   * @param generator api to generate JSON document
   * @param value JSON value to write
   * @throws IOException if can't write to output
   */
  private static void writeObject(JsonGenerator generator, Object value) throws IOException {
    generator.writeObject(value);
  }

  /**
   * A Query that matches documents within an range of terms.
   */
  static class RangeQueryBuilder extends QueryBuilder {
    private final String fieldName;

    private Object lt;
    private boolean lte;
    private Object gt;
    private boolean gte;

    private String format;

    private RangeQueryBuilder(final String fieldName) {
      this.fieldName = requireNonNull(fieldName, "fieldName");
    }

    private RangeQueryBuilder to(Object value, boolean lte) {
      this.lt = requireNonNull(value, "value");
      this.lte = lte;
      return this;
    }

    private RangeQueryBuilder from(Object value, boolean gte) {
      this.gt = requireNonNull(value, "value");
      this.gte = gte;
      return this;
    }

    RangeQueryBuilder lt(Object value) {
      return to(value, false);
    }

    RangeQueryBuilder lte(Object value) {
      return to(value, true);
    }

    RangeQueryBuilder gt(Object value) {
      return from(value, false);
    }

    RangeQueryBuilder gte(Object value) {
      return from(value, true);
    }

    RangeQueryBuilder format(String format) {
      this.format = format;
      return this;
    }

    @Override void writeJson(final JsonGenerator generator) throws IOException {
      if (lt == null && gt == null) {
        throw new IllegalStateException("Either lower or upper bound should be provided");
      }

      generator.writeStartObject();
      generator.writeFieldName("range");
      generator.writeStartObject();
      generator.writeFieldName(fieldName);
      generator.writeStartObject();

      if (gt != null) {
        final String op = gte ? "gte" : "gt";
        generator.writeFieldName(op);
        writeObject(generator, gt);
      }

      if (lt != null) {
        final String op = lte ? "lte" : "lt";
        generator.writeFieldName(op);
        writeObject(generator, lt);
      }

      if (format != null) {
        generator.writeStringField("format", format);
      }

      generator.writeEndObject();
      generator.writeEndObject();
      generator.writeEndObject();
    }
  }

  /**
   * A Query that does fuzzy matching for a specific value.
   */
  static class RegexpQueryBuilder extends QueryBuilder {
    @SuppressWarnings("unused")
    private final String fieldName;
    @SuppressWarnings("unused")
    private final String value;
    @SuppressWarnings("unused")
    private String escape;

    RegexpQueryBuilder(final String fieldName, final String value) {
      this(fieldName, value, "\\");
    }

    RegexpQueryBuilder(final String fieldName, final String value, final String escape) {
      requireNonNull(fieldName, "fieldName");
      requireNonNull(value, "value");
      requireNonNull(escape, "escape");
      this.fieldName = fieldName;
      this.escape = escape;
      // replace % to * and _ to ? for sql with like operator
      HashMap<String, String> kv = new HashMap<>();
      kv.put("%", "*");
      kv.put("_", "?");
      this.value = replaceWildcard(value, kv, escape);
    }

    public static String replaceWildcard(String value, Map<String, String> kv, String escape) {
      ArrayList<String> ret = new ArrayList<>();
      int escapeCount = 0;
      for (int index = 0; index < value.length(); index++) {
        String current = value.substring(index, index + 1);
        if (index == 0) {
          if (!current.equals(escape)) {
            current = kv.keySet().contains(current) ? kv.get(current) : current;
            ret.add(current);
          } else {
            escapeCount++;
          }
          continue;
        }

        if (!kv.keySet().contains(current) && !current.equals(escape)) {
          ret.add(current);
          escapeCount = 0;
          continue;
        }

        if (current.equals(escape)) {
          escapeCount++;
          if (escapeCount % 2 == 0) {
            ret.add(current);
          }
          continue;
        }

        String last = value.substring(index - 1, index);
        if (kv.keySet().contains(current)) {
          if (!last.equals(escape)) {
            ret.add(kv.get(current));
          } else {
            if (escapeCount % 2 == 0) {
              ret.add(kv.get(current));
            } else {
              ret.add(current);
            }
          }
          escapeCount = 0;
        }
      }
      return String.join("", ret);
    }

    @Override void writeJson(final JsonGenerator generator) throws IOException  {
      generator.writeStartObject();
      generator.writeFieldName("wildcard");
      generator.writeStartObject();
      generator.writeFieldName(fieldName);
      writeObject(generator, value);
      generator.writeEndObject();
      generator.writeEndObject();
    }
  }

  /**
   * Constructs a query that only match on documents that the field has a value in them.
   */
  static class ExistsQueryBuilder extends QueryBuilder {
    private final String fieldName;

    ExistsQueryBuilder(final String fieldName) {
      this.fieldName = requireNonNull(fieldName, "fieldName");
    }

    @Override void writeJson(final JsonGenerator generator) throws IOException {
      generator.writeStartObject();
      generator.writeFieldName("exists");
      generator.writeStartObject();
      generator.writeStringField("field", fieldName);
      generator.writeEndObject();
      generator.writeEndObject();
    }
  }

  /**
   * A query that wraps a filter and simply returns a constant score equal to the
   * query boost for every document in the filter.
   */
  static class ConstantScoreQueryBuilder extends QueryBuilder {

    private final QueryBuilder builder;

    private ConstantScoreQueryBuilder(final QueryBuilder builder) {
      this.builder = requireNonNull(builder, "builder");
    }

    @Override void writeJson(final JsonGenerator generator) throws IOException {
      generator.writeStartObject();
      generator.writeFieldName("constant_score");
      generator.writeStartObject();
      generator.writeFieldName("filter");
      builder.writeJson(generator);
      generator.writeEndObject();
      generator.writeEndObject();
    }
  }

  /**
   * A query that wraps a filter and simply returns a dismax score equal to the
   * query boost for every document in the filter.
   */
  static class DisMaxQueryBuilder extends QueryBuilder {

    private final QueryBuilder builder;

    private DisMaxQueryBuilder(final QueryBuilder builder) {
      this.builder = requireNonNull(builder, "builder");
    }

    @Override void writeJson(final JsonGenerator generator) throws IOException {
      generator.writeStartObject();
      generator.writeFieldName("dis_max");
      generator.writeStartObject();
      generator.writeFieldName("queries");
      generator.writeStartArray();
      builder.writeJson(generator);
      generator.writeEndArray();
      generator.writeEndObject();
      generator.writeEndObject();
    }
  }



  /**
   * A query that matches on all documents.
   * <pre>
   *   {
   *     "match_all": {}
   *   }
   * </pre>
   */
  static class MatchAllQueryBuilder extends QueryBuilder {

    private MatchAllQueryBuilder() {}

    @Override void writeJson(final JsonGenerator generator) throws IOException {
      generator.writeStartObject();
      generator.writeFieldName("match_all");
      generator.writeStartObject();
      generator.writeEndObject();
      generator.writeEndObject();
    }
  }
}