/* * SPDX-License-Identifier: Apache-2.0 * * The OpenSearch Contributors require contributions made to * this file be licensed under the Apache-2.0 license or a * compatible open source license. */ /* * Licensed to Elasticsearch under one or more contributor * license agreements. See the NOTICE file distributed with * this work for additional information regarding copyright * ownership. Elasticsearch 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. */ /* * Modifications Copyright OpenSearch Contributors. See * GitHub history for details. */ package org.opensearch.index.query; import org.apache.lucene.index.LeafReaderContext; import org.apache.lucene.search.ConstantScoreScorer; import org.apache.lucene.search.ConstantScoreWeight; import org.apache.lucene.search.DocIdSetIterator; import org.apache.lucene.search.IndexSearcher; import org.apache.lucene.search.Query; import org.apache.lucene.search.QueryVisitor; import org.apache.lucene.search.ScoreMode; import org.apache.lucene.search.Scorer; import org.apache.lucene.search.TwoPhaseIterator; import org.apache.lucene.search.Weight; import org.opensearch.OpenSearchException; import org.opensearch.common.Nullable; import org.opensearch.core.common.ParsingException; import org.opensearch.core.common.io.stream.StreamInput; import org.opensearch.core.common.io.stream.StreamOutput; import org.opensearch.common.lucene.search.function.Functions; import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.core.xcontent.XContentParser; import org.opensearch.script.FilterScript; import org.opensearch.script.Script; import java.io.IOException; import java.util.Objects; import static org.opensearch.search.SearchService.ALLOW_EXPENSIVE_QUERIES; /** * Query builder for script queries * * @opensearch.internal */ public class ScriptQueryBuilder extends AbstractQueryBuilder { public static final String NAME = "script"; private final Script script; public ScriptQueryBuilder(Script script) { if (script == null) { throw new IllegalArgumentException("script cannot be null"); } this.script = script; } /** * Read from a stream. */ public ScriptQueryBuilder(StreamInput in) throws IOException { super(in); script = new Script(in); } @Override protected void doWriteTo(StreamOutput out) throws IOException { script.writeTo(out); } public Script script() { return this.script; } @Override public String getWriteableName() { return NAME; } @Override protected void doXContent(XContentBuilder builder, Params builderParams) throws IOException { builder.startObject(NAME); builder.field(Script.SCRIPT_PARSE_FIELD.getPreferredName(), script); printBoostAndQueryName(builder); builder.endObject(); } public static ScriptQueryBuilder fromXContent(XContentParser parser) throws IOException { // also, when caching, since its isCacheable is false, will result in loading all bit set... Script script = null; float boost = AbstractQueryBuilder.DEFAULT_BOOST; String queryName = null; XContentParser.Token token; String currentFieldName = null; while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) { if (token == XContentParser.Token.FIELD_NAME) { currentFieldName = parser.currentName(); } else if (token == XContentParser.Token.START_OBJECT) { if (Script.SCRIPT_PARSE_FIELD.match(currentFieldName, parser.getDeprecationHandler())) { script = Script.parse(parser); } else { throw new ParsingException(parser.getTokenLocation(), "[script] query does not support [" + currentFieldName + "]"); } } else if (token.isValue()) { if (AbstractQueryBuilder.NAME_FIELD.match(currentFieldName, parser.getDeprecationHandler())) { queryName = parser.text(); } else if (AbstractQueryBuilder.BOOST_FIELD.match(currentFieldName, parser.getDeprecationHandler())) { boost = parser.floatValue(); } else if (Script.SCRIPT_PARSE_FIELD.match(currentFieldName, parser.getDeprecationHandler())) { script = Script.parse(parser); } else { throw new ParsingException(parser.getTokenLocation(), "[script] query does not support [" + currentFieldName + "]"); } } else { if (token != XContentParser.Token.START_ARRAY) { throw new AssertionError("Impossible token received: " + token.name()); } throw new ParsingException( parser.getTokenLocation(), "[script] query does not support an array of scripts. Use a bool query with a clause per script instead." ); } } if (script == null) { throw new ParsingException(parser.getTokenLocation(), "script must be provided with a [script] filter"); } return new ScriptQueryBuilder(script).boost(boost).queryName(queryName); } @Override protected Query doToQuery(QueryShardContext context) throws IOException { if (context.allowExpensiveQueries() == false) { throw new OpenSearchException( "[script] queries cannot be executed when '" + ALLOW_EXPENSIVE_QUERIES.getKey() + "' is set to false." ); } FilterScript.Factory factory = context.compile(script, FilterScript.CONTEXT); FilterScript.LeafFactory filterScript = factory.newFactory(script.getParams(), context.lookup()); return new ScriptQuery(script, filterScript, queryName); } /** * Internal script query * * @opensearch.internal */ static class ScriptQuery extends Query { final Script script; final FilterScript.LeafFactory filterScript; final String queryName; ScriptQuery(Script script, FilterScript.LeafFactory filterScript, @Nullable String queryName) { this.script = script; this.filterScript = filterScript; this.queryName = queryName; } @Override public String toString(String field) { StringBuilder buffer = new StringBuilder(); buffer.append("ScriptQuery("); buffer.append(script); buffer.append(Functions.nameOrEmptyArg(queryName)); buffer.append(")"); return buffer.toString(); } @Override public void visit(QueryVisitor visitor) { visitor.visitLeaf(this); } @Override public boolean equals(Object obj) { if (sameClassAs(obj) == false) return false; ScriptQuery other = (ScriptQuery) obj; return Objects.equals(script, other.script); } @Override public int hashCode() { int h = classHash(); h = 31 * h + script.hashCode(); return h; } @Override public Weight createWeight(IndexSearcher searcher, ScoreMode scoreMode, float boost) throws IOException { return new ConstantScoreWeight(this, boost) { @Override public Scorer scorer(LeafReaderContext context) throws IOException { DocIdSetIterator approximation = DocIdSetIterator.all(context.reader().maxDoc()); final FilterScript leafScript = filterScript.newInstance(context); TwoPhaseIterator twoPhase = new TwoPhaseIterator(approximation) { @Override public boolean matches() throws IOException { leafScript.setDocument(approximation.docID()); return leafScript.execute(); } @Override public float matchCost() { // TODO: how can we compute this? return 1000f; } }; return new ConstantScoreScorer(this, score(), scoreMode, twoPhase); } @Override public boolean isCacheable(LeafReaderContext ctx) { // TODO: Change this to true when we can assume that scripts are pure functions // ie. the return value is always the same given the same conditions and may not // depend on the current timestamp, other documents, etc. return false; } }; } } @Override protected int doHashCode() { return Objects.hash(script); } @Override protected boolean doEquals(ScriptQueryBuilder other) { return Objects.equals(script, other.script); } }