diff --git a/src/main/java/org/elasticsearch/action/search/SearchRequestBuilder.java b/src/main/java/org/elasticsearch/action/search/SearchRequestBuilder.java index 5390c05b2cb97..6860f97a0c304 100644 --- a/src/main/java/org/elasticsearch/action/search/SearchRequestBuilder.java +++ b/src/main/java/org/elasticsearch/action/search/SearchRequestBuilder.java @@ -447,6 +447,40 @@ public SearchRequestBuilder addScriptField(String name, String lang, String scri return this; } + /** + * Sets a script based source to load. + * + * @param script The script to use + */ + public SearchRequestBuilder setScriptSource(String script) { + sourceBuilder().scriptSource(script); + return this; + } + + /** + * Sets a script based source to load. + * + * @param script The script to use + * @param params Parameters that the script can use. + */ + public SearchRequestBuilder setScriptSource(String script, Map params) { + sourceBuilder().scriptSource(script, params); + return this; + } + + /** + * Sets a script based source to load. + * + * @param lang The language of the script + * @param script The script to use + * @param params Parameters that the script can use (can be null). + */ + public SearchRequestBuilder setScriptSource(String lang, String script, Map params) { + sourceBuilder().scriptSource(lang, script, params); + return this; + } + + /** * Adds a sort against the given field name and the sort ordering. * diff --git a/src/main/java/org/elasticsearch/search/SearchModule.java b/src/main/java/org/elasticsearch/search/SearchModule.java index 53d8acd4741d6..542f4e6c3ab1b 100644 --- a/src/main/java/org/elasticsearch/search/SearchModule.java +++ b/src/main/java/org/elasticsearch/search/SearchModule.java @@ -29,9 +29,11 @@ import org.elasticsearch.search.facet.FacetModule; import org.elasticsearch.search.fetch.FetchPhase; import org.elasticsearch.search.fetch.explain.ExplainFetchSubPhase; +import org.elasticsearch.search.fetch.extractfields.ExtractFieldsFetchSubPhase; import org.elasticsearch.search.fetch.matchedfilters.MatchedFiltersFetchSubPhase; import org.elasticsearch.search.fetch.partial.PartialFieldsFetchSubPhase; import org.elasticsearch.search.fetch.script.ScriptFieldsFetchSubPhase; +import org.elasticsearch.search.fetch.scriptsource.ScriptSourceFetchSubPhase; import org.elasticsearch.search.fetch.version.VersionFetchSubPhase; import org.elasticsearch.search.highlight.HighlightPhase; import org.elasticsearch.search.query.QueryPhase; @@ -56,6 +58,8 @@ protected void configure() { bind(FetchPhase.class).asEagerSingleton(); bind(ExplainFetchSubPhase.class).asEagerSingleton(); bind(ScriptFieldsFetchSubPhase.class).asEagerSingleton(); + bind(ScriptSourceFetchSubPhase.class).asEagerSingleton(); + bind(ExtractFieldsFetchSubPhase.class).asEagerSingleton(); bind(PartialFieldsFetchSubPhase.class).asEagerSingleton(); bind(VersionFetchSubPhase.class).asEagerSingleton(); bind(MatchedFiltersFetchSubPhase.class).asEagerSingleton(); diff --git a/src/main/java/org/elasticsearch/search/builder/SearchSourceBuilder.java b/src/main/java/org/elasticsearch/search/builder/SearchSourceBuilder.java index 6e92b49b15d9c..ab13901a83cb2 100644 --- a/src/main/java/org/elasticsearch/search/builder/SearchSourceBuilder.java +++ b/src/main/java/org/elasticsearch/search/builder/SearchSourceBuilder.java @@ -95,6 +95,7 @@ public static HighlightBuilder highlight() { private List fieldNames; private List scriptFields; + private ScriptSource scriptSource; private List partialFields; private List facets; @@ -481,6 +482,38 @@ public SearchSourceBuilder scriptField(String name, String lang, String script, return this; } + /** + * Sets a script source. + * + * @param script The script + */ + public SearchSourceBuilder scriptSource(String script) { + return scriptSource(null, script, null); + } + + /** + * Sets a script source. + * + * @param script The script to execute + * @param params The script parameters + */ + public SearchSourceBuilder scriptSource(String script, Map params) { + return scriptSource(null, script, params); + } + + /** + * Sets a script source. + * + * @param lang The language of the script + * @param script The script to execute + * @param params The script parameters (can be null) + * @return + */ + public SearchSourceBuilder scriptSource(String lang, String script, Map params) { + scriptSource = new ScriptSource(lang, script, params); + return this; + } + /** * Adds a partial field based on _source, with an "include" and/or "exclude" set which can include simple wildcard * elements. @@ -666,6 +699,19 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws builder.endObject(); } + if (scriptSource != null) { + builder.startObject("script_source"); + builder.field("script", scriptSource.script()); + if (scriptSource.lang() != null) { + builder.field("lang", scriptSource.lang()); + } + if (scriptSource.params() != null) { + builder.field("params"); + builder.map(scriptSource.params()); + } + builder.endObject(); + } + if (sorts != null) { builder.startArray("sort"); for (SortBuilder sort : sorts) { @@ -780,4 +826,28 @@ public String[] excludes() { return excludes; } } + + private static class ScriptSource { + private final String script; + private final String lang; + private final Map params; + + private ScriptSource(String lang, String script, Map params) { + this.lang = lang; + this.script = script; + this.params = params; + } + + public String script() { + return script; + } + + public String lang() { + return this.lang; + } + + public Map params() { + return params; + } + } } diff --git a/src/main/java/org/elasticsearch/search/fetch/FetchPhase.java b/src/main/java/org/elasticsearch/search/fetch/FetchPhase.java index bae3fff0a3754..317f040018dd2 100644 --- a/src/main/java/org/elasticsearch/search/fetch/FetchPhase.java +++ b/src/main/java/org/elasticsearch/search/fetch/FetchPhase.java @@ -20,7 +20,6 @@ package org.elasticsearch.search.fetch; import com.google.common.collect.ImmutableMap; -import com.google.common.collect.Lists; import org.apache.lucene.document.Document; import org.apache.lucene.document.Fieldable; import org.apache.lucene.index.IndexReader; @@ -44,9 +43,11 @@ import org.elasticsearch.search.SearchParseElement; import org.elasticsearch.search.SearchPhase; import org.elasticsearch.search.fetch.explain.ExplainFetchSubPhase; +import org.elasticsearch.search.fetch.extractfields.ExtractFieldsFetchSubPhase; import org.elasticsearch.search.fetch.matchedfilters.MatchedFiltersFetchSubPhase; import org.elasticsearch.search.fetch.partial.PartialFieldsFetchSubPhase; import org.elasticsearch.search.fetch.script.ScriptFieldsFetchSubPhase; +import org.elasticsearch.search.fetch.scriptsource.ScriptSourceFetchSubPhase; import org.elasticsearch.search.fetch.version.VersionFetchSubPhase; import org.elasticsearch.search.highlight.HighlightPhase; import org.elasticsearch.search.internal.InternalSearchHit; @@ -68,9 +69,10 @@ public class FetchPhase implements SearchPhase { private final FetchSubPhase[] fetchSubPhases; @Inject - public FetchPhase(HighlightPhase highlightPhase, ScriptFieldsFetchSubPhase scriptFieldsPhase, PartialFieldsFetchSubPhase partialFieldsPhase, + public FetchPhase(HighlightPhase highlightPhase, ScriptSourceFetchSubPhase scriptSourcePhase, ScriptFieldsFetchSubPhase scriptFieldsPhase, + ExtractFieldsFetchSubPhase extractFieldsPhase, PartialFieldsFetchSubPhase partialFieldsPhase, MatchedFiltersFetchSubPhase matchFiltersPhase, ExplainFetchSubPhase explainPhase, VersionFetchSubPhase versionPhase) { - this.fetchSubPhases = new FetchSubPhase[]{scriptFieldsPhase, partialFieldsPhase, matchFiltersPhase, explainPhase, highlightPhase, versionPhase}; + this.fetchSubPhases = new FetchSubPhase[]{scriptSourcePhase, extractFieldsPhase, scriptFieldsPhase, partialFieldsPhase, matchFiltersPhase, explainPhase, highlightPhase, versionPhase}; } @Override @@ -89,7 +91,6 @@ public void preProcess(SearchContext context) { public void execute(SearchContext context) { ResetFieldSelector fieldSelector; - List extractFieldNames = null; boolean sourceRequested = false; if (!context.hasFieldNames()) { if (context.hasPartialFields()) { @@ -126,15 +127,12 @@ public void execute(SearchContext context) { } fieldSelectorMapper.add(x); } else { - if (extractFieldNames == null) { - extractFieldNames = Lists.newArrayList(); - } - extractFieldNames.add(fieldName); + context.extractFieldNames().add(fieldName); } } if (loadAllStored) { - if (sourceRequested || extractFieldNames != null) { + if (sourceRequested || context.hasExtractFieldNames()) { fieldSelector = null; // load everything, including _source } else { fieldSelector = AllButSourceFieldSelector.INSTANCE; @@ -142,11 +140,11 @@ public void execute(SearchContext context) { } else if (fieldSelectorMapper != null) { // we are asking specific stored fields, just add the UID and be done fieldSelectorMapper.add(UidFieldMapper.NAME); - if (extractFieldNames != null || sourceRequested) { + if (context.hasExtractFieldNames() || sourceRequested) { fieldSelectorMapper.add(SourceFieldMapper.NAME); } fieldSelector = fieldSelectorMapper; - } else if (extractFieldNames != null || sourceRequested) { + } else if (context.hasExtractFieldNames() || sourceRequested) { fieldSelector = new UidAndSourceFieldSelector(); } else { fieldSelector = UidFieldSelector.INSTANCE; @@ -169,7 +167,7 @@ public void execute(SearchContext context) { // get the version - InternalSearchHit searchHit = new InternalSearchHit(docId, uid.id(), uid.type(), sourceRequested ? source : null, null); + InternalSearchHit searchHit = new InternalSearchHit(docId, uid.id(), uid.type(), null); hits[index] = searchHit; for (Object oField : doc.getFields()) { @@ -219,37 +217,24 @@ public void execute(SearchContext context) { IndexReader subReader = context.searcher().subReaders()[readerIndex]; int subDoc = docId - context.searcher().docStarts()[readerIndex]; - // go over and extract fields that are not mapped / stored context.lookup().setNextReader(subReader); context.lookup().setNextDocId(subDoc); if (source != null) { context.lookup().source().setNextSource(new BytesArray(source)); } - if (extractFieldNames != null) { - for (String extractFieldName : extractFieldNames) { - Object value = context.lookup().source().extractValue(extractFieldName); - if (value != null) { - if (searchHit.fieldsOrNull() == null) { - searchHit.fields(new HashMap(2)); - } - - SearchHitField hitField = searchHit.fields().get(extractFieldName); - if (hitField == null) { - hitField = new InternalSearchHitField(extractFieldName, new ArrayList(2)); - searchHit.fields().put(extractFieldName, hitField); - } - hitField.values().add(value); - } - } - } + FetchSubPhase.HitContext hitContext = new FetchSubPhase.HitContext(source); for (FetchSubPhase fetchSubPhase : fetchSubPhases) { - FetchSubPhase.HitContext hitContext = new FetchSubPhase.HitContext(); if (fetchSubPhase.hitExecutionNeeded(context)) { hitContext.reset(searchHit, subReader, subDoc, context.searcher().getIndexReader(), docId, doc); fetchSubPhase.hitExecute(context, hitContext); } } + + // Store source if needed + if (sourceRequested && hitContext.source() != null) { + searchHit.source(hitContext.source()); + } } for (FetchSubPhase fetchSubPhase : fetchSubPhases) { diff --git a/src/main/java/org/elasticsearch/search/fetch/FetchSubPhase.java b/src/main/java/org/elasticsearch/search/fetch/FetchSubPhase.java index d47da76e698bc..535cc02d3dc8a 100644 --- a/src/main/java/org/elasticsearch/search/fetch/FetchSubPhase.java +++ b/src/main/java/org/elasticsearch/search/fetch/FetchSubPhase.java @@ -23,10 +23,16 @@ import org.apache.lucene.document.Document; import org.apache.lucene.index.IndexReader; import org.elasticsearch.ElasticSearchException; +import org.elasticsearch.common.io.stream.BytesStreamOutput; +import org.elasticsearch.common.io.stream.CachedStreamOutput; +import org.elasticsearch.common.xcontent.XContentBuilder; +import org.elasticsearch.common.xcontent.XContentFactory; import org.elasticsearch.search.SearchParseElement; import org.elasticsearch.search.internal.InternalSearchHit; import org.elasticsearch.search.internal.SearchContext; +import org.elasticsearch.search.lookup.SourceLookup; +import java.io.IOException; import java.util.Map; /** @@ -42,6 +48,13 @@ public static class HitContext { private int docId; private Document doc; private Map cache; + private byte[] source; + private Map sourceAsMap; + + public HitContext(byte[] source) { + this.source = source; + this.sourceAsMap = null; + } public void reset(InternalSearchHit hit, IndexReader reader, int docId, IndexReader topLevelReader, int topLevelDocId, Document doc) { this.hit = hit; @@ -82,6 +95,49 @@ public Map cache() { } return cache; } + + public byte[] source() { + if (source != null) { + return source; + } + if (sourceAsMap != null) { + CachedStreamOutput.Entry cachedEntry = CachedStreamOutput.popEntry(); + try { + BytesStreamOutput streamOutput = cachedEntry.bytes(); + // TODO: Is there a better way to figure out Content Type? + XContentBuilder builder = XContentFactory.jsonBuilder(streamOutput).map(sourceAsMap); + builder.close(); + source = streamOutput.bytes().copyBytesArray().toBytes(); + return source; + } catch (IOException ex) { + throw new ElasticSearchException("Cannot serialize map " + sourceAsMap, ex); + } finally { + CachedStreamOutput.pushEntry(cachedEntry); + } + } + return null; + } + + public Map sourceAsMap() { + if (sourceAsMap != null) { + return sourceAsMap; + } + if (source != null) { + sourceAsMap = SourceLookup.sourceAsMap(source, 0, source.length); + return sourceAsMap; + } + return null; + } + + public void source(byte[] source) { + this.source = source; + this.sourceAsMap = null; + } + + public void sourceAsMap(Map sourceAsMap) { + this.source = null; + this.sourceAsMap = sourceAsMap; + } } Map parseElements(); diff --git a/src/main/java/org/elasticsearch/search/fetch/extractfields/ExtractFieldsFetchSubPhase.java b/src/main/java/org/elasticsearch/search/fetch/extractfields/ExtractFieldsFetchSubPhase.java new file mode 100644 index 0000000000000..de8d653ffb8cf --- /dev/null +++ b/src/main/java/org/elasticsearch/search/fetch/extractfields/ExtractFieldsFetchSubPhase.java @@ -0,0 +1,84 @@ +/* + * Licensed to Elastic Search and Shay Banon under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. Elastic Search 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.elasticsearch.search.fetch.extractfields; + +import com.google.common.collect.ImmutableMap; +import org.elasticsearch.ElasticSearchException; +import org.elasticsearch.common.inject.Inject; +import org.elasticsearch.search.SearchHitField; +import org.elasticsearch.search.SearchParseElement; +import org.elasticsearch.search.fetch.FetchSubPhase; +import org.elasticsearch.search.internal.InternalSearchHit; +import org.elasticsearch.search.internal.InternalSearchHitField; +import org.elasticsearch.search.internal.SearchContext; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.Map; + +/** + * + */ +public class ExtractFieldsFetchSubPhase implements FetchSubPhase { + @Inject + public ExtractFieldsFetchSubPhase() { + } + + @Override + public Map parseElements() { + return ImmutableMap.of(); + } + + @Override + public boolean hitsExecutionNeeded(SearchContext context) { + return false; + } + + @Override + public void hitsExecute(SearchContext context, InternalSearchHit[] hits) throws ElasticSearchException { + } + + @Override + public boolean hitExecutionNeeded(SearchContext context) { + return context.hasExtractFieldNames(); + } + + @Override + public void hitExecute(SearchContext context, FetchSubPhase.HitContext hitContext) throws ElasticSearchException { + context.lookup().setNextReader(hitContext.reader()); + context.lookup().setNextDocId(hitContext.docId()); + context.lookup().source().setNextSource(hitContext.sourceAsMap()); + for (String extractFieldName : context.extractFieldNames()) { + Object value = context.lookup().source().extractValue(extractFieldName); + if (value != null) { + if (hitContext.hit().fieldsOrNull() == null) { + hitContext.hit().fields(new HashMap(2)); + } + + SearchHitField hitField = hitContext.hit().fields().get(extractFieldName); + if (hitField == null) { + hitField = new InternalSearchHitField(extractFieldName, new ArrayList(2)); + hitContext.hit().fields().put(extractFieldName, hitField); + } + hitField.values().add(value); + } + } + } +} diff --git a/src/main/java/org/elasticsearch/search/fetch/script/ScriptFieldsFetchSubPhase.java b/src/main/java/org/elasticsearch/search/fetch/script/ScriptFieldsFetchSubPhase.java index ee48b32575a8d..d9278f4b9970a 100644 --- a/src/main/java/org/elasticsearch/search/fetch/script/ScriptFieldsFetchSubPhase.java +++ b/src/main/java/org/elasticsearch/search/fetch/script/ScriptFieldsFetchSubPhase.java @@ -69,6 +69,7 @@ public void hitExecute(SearchContext context, HitContext hitContext) throws Elas for (ScriptFieldsContext.ScriptField scriptField : context.scriptFields().fields()) { scriptField.script().setNextReader(hitContext.reader()); scriptField.script().setNextDocId(hitContext.docId()); + scriptField.script().setNextSource(hitContext.sourceAsMap()); Object value; try { diff --git a/src/main/java/org/elasticsearch/search/fetch/scriptsource/ScriptSourceContext.java b/src/main/java/org/elasticsearch/search/fetch/scriptsource/ScriptSourceContext.java new file mode 100644 index 0000000000000..c1cc57f27fa2a --- /dev/null +++ b/src/main/java/org/elasticsearch/search/fetch/scriptsource/ScriptSourceContext.java @@ -0,0 +1,49 @@ +/* + * Licensed to ElasticSearch and Shay Banon 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. + */ + +package org.elasticsearch.search.fetch.scriptsource; + +import org.elasticsearch.script.SearchScript; + +/** + * + */ +public class ScriptSourceContext { + private final String name; + private final SearchScript script; + private final boolean ignoreException; + + public ScriptSourceContext(String name, SearchScript script, boolean ignoreException) { + this.name = name; + this.script = script; + this.ignoreException = ignoreException; + } + + public String name() { + return name; + } + + public SearchScript script() { + return this.script; + } + + public boolean ignoreException() { + return ignoreException; + } +} diff --git a/src/main/java/org/elasticsearch/search/fetch/scriptsource/ScriptSourceFetchSubPhase.java b/src/main/java/org/elasticsearch/search/fetch/scriptsource/ScriptSourceFetchSubPhase.java new file mode 100644 index 0000000000000..a0e018599f736 --- /dev/null +++ b/src/main/java/org/elasticsearch/search/fetch/scriptsource/ScriptSourceFetchSubPhase.java @@ -0,0 +1,107 @@ +/* + * Licensed to Elastic Search and Shay Banon under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. Elastic Search 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.elasticsearch.search.fetch.scriptsource; + +import com.google.common.collect.ImmutableMap; +import org.elasticsearch.ElasticSearchException; +import org.elasticsearch.ElasticSearchIllegalArgumentException; +import org.elasticsearch.common.inject.Inject; +import org.elasticsearch.search.SearchParseElement; +import org.elasticsearch.search.fetch.FetchSubPhase; +import org.elasticsearch.search.internal.InternalSearchHit; +import org.elasticsearch.search.internal.SearchContext; + +import java.util.Map; + +/** + * + */ +public class ScriptSourceFetchSubPhase implements FetchSubPhase { + + private static final Map EMPTY_SOURCE = ImmutableMap.of(); + + @Inject + public ScriptSourceFetchSubPhase() { + } + + @Override + public Map parseElements() { + ImmutableMap.Builder parseElements = ImmutableMap.builder(); + parseElements.put("script_source", new ScriptSourceParseElement()) + .put("scriptSource", new ScriptSourceParseElement()); + return parseElements.build(); + } + + @Override + public boolean hitsExecutionNeeded(SearchContext context) { + return false; + } + + @Override + public void hitsExecute(SearchContext context, InternalSearchHit[] hits) throws ElasticSearchException { + } + + @Override + public boolean hitExecutionNeeded(SearchContext context) { + return context.hasScriptSource(); + } + + @Override + public void hitExecute(SearchContext context, HitContext hitContext) throws ElasticSearchException { + ScriptSourceContext script = context.scriptSource(); + script.script().setNextReader(hitContext.reader()); + script.script().setNextDocId(hitContext.docId()); + script.script().setNextSource(hitContext.sourceAsMap()); + try { + Object value; + value = script.script().run(); + value = script.script().unwrap(value); + updateSource(context, hitContext, value); + } catch (RuntimeException e) { + if (!script.ignoreException()) { + throw e; + } + } + } + + @SuppressWarnings("unchecked") + private void updateSource(SearchContext context, HitContext hitContext, Object value) { + if (value != null) { + if (value == context.lookup().source()) { + // script returned _source - it means the source shouldn't be changed + return; + } else if (value instanceof String) { + hitContext.source(((String)value).getBytes()); + } else if (value instanceof byte[]) { + hitContext.source((byte[]) value); + } else if (value instanceof Map) { + hitContext.sourceAsMap((Map) value); + } else { + throw new ElasticSearchIllegalArgumentException("Source script returned unsupported source type " + value.getClass().getName()); + } + } else { + // Script returned null - no source is available + hitContext.sourceAsMap(EMPTY_SOURCE); + } + // Update source lookup for other phases + context.lookup().source().setNextSource(hitContext.sourceAsMap()); + } + +} diff --git a/src/main/java/org/elasticsearch/search/fetch/scriptsource/ScriptSourceParseElement.java b/src/main/java/org/elasticsearch/search/fetch/scriptsource/ScriptSourceParseElement.java new file mode 100644 index 0000000000000..265009506b3b6 --- /dev/null +++ b/src/main/java/org/elasticsearch/search/fetch/scriptsource/ScriptSourceParseElement.java @@ -0,0 +1,62 @@ +/* + * Licensed to Elastic Search and Shay Banon under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. Elastic Search 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.elasticsearch.search.fetch.scriptsource; + +import org.elasticsearch.common.xcontent.XContentParser; +import org.elasticsearch.script.SearchScript; +import org.elasticsearch.search.SearchParseElement; +import org.elasticsearch.search.internal.SearchContext; + +import java.util.Map; + +/** + * + */ +public class ScriptSourceParseElement implements SearchParseElement { + @Override + public void parse(XContentParser parser, SearchContext context) throws Exception { + String currentFieldName = null; + XContentParser.Token token = parser.currentToken(); + if (token == XContentParser.Token.START_OBJECT) { + String fieldName = currentFieldName; + String script = null; + String scriptLang = null; + Map params = null; + boolean ignoreException = false; + while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) { + if (token == XContentParser.Token.FIELD_NAME) { + currentFieldName = parser.currentName(); + } else if (token == XContentParser.Token.START_OBJECT) { + params = parser.map(); + } else if (token.isValue()) { + if ("script".equals(currentFieldName)) { + script = parser.text(); + } else if ("lang".equals(currentFieldName)) { + scriptLang = parser.text(); + } else if ("ignore_failure".equals(currentFieldName)) { + ignoreException = parser.booleanValue(); + } + } + } + SearchScript searchScript = context.scriptService().search(context.lookup(), scriptLang, script, params); + context.scriptSource(new ScriptSourceContext(fieldName, searchScript, ignoreException)); + } + } +} diff --git a/src/main/java/org/elasticsearch/search/highlight/HighlightPhase.java b/src/main/java/org/elasticsearch/search/highlight/HighlightPhase.java index 082f9b9059e01..2f22ab568f685 100644 --- a/src/main/java/org/elasticsearch/search/highlight/HighlightPhase.java +++ b/src/main/java/org/elasticsearch/search/highlight/HighlightPhase.java @@ -201,6 +201,7 @@ public void hitExecute(SearchContext context, HitContext hitContext) throws Elas SearchLookup lookup = context.lookup(); lookup.setNextReader(hitContext.reader()); lookup.setNextDocId(hitContext.docId()); + lookup.source().setNextSource(hitContext.sourceAsMap()); textsToHighlight = lookup.source().extractRawValues(mapper.names().sourcePath()); } diff --git a/src/main/java/org/elasticsearch/search/internal/InternalSearchHit.java b/src/main/java/org/elasticsearch/search/internal/InternalSearchHit.java index 7405227e85403..4d311426bf74e 100644 --- a/src/main/java/org/elasticsearch/search/internal/InternalSearchHit.java +++ b/src/main/java/org/elasticsearch/search/internal/InternalSearchHit.java @@ -89,11 +89,10 @@ private InternalSearchHit() { } - public InternalSearchHit(int docId, String id, String type, byte[] source, Map fields) { + public InternalSearchHit(int docId, String id, String type, Map fields) { this.docId = docId; this.id = id; this.type = type; - this.source = source == null ? null : new BytesArray(source); this.fields = fields; } @@ -227,6 +226,18 @@ public String getSourceAsString() { return sourceAsString(); } + public void source(byte[] sourceAsBytes) { + this.source = new BytesArray(sourceAsBytes); + this.sourceAsBytes = null; + this.sourceAsMap = null; + } + + public void source(BytesArray source) { + this.source = source; + this.sourceAsBytes = null; + this.sourceAsMap = null; + } + @SuppressWarnings({"unchecked"}) @Override public Map sourceAsMap() throws ElasticSearchParseException { diff --git a/src/main/java/org/elasticsearch/search/internal/SearchContext.java b/src/main/java/org/elasticsearch/search/internal/SearchContext.java index d86cdd37e80bb..e04eaa5183889 100644 --- a/src/main/java/org/elasticsearch/search/internal/SearchContext.java +++ b/src/main/java/org/elasticsearch/search/internal/SearchContext.java @@ -51,6 +51,7 @@ import org.elasticsearch.search.fetch.FetchSearchResult; import org.elasticsearch.search.fetch.partial.PartialFieldsContext; import org.elasticsearch.search.fetch.script.ScriptFieldsContext; +import org.elasticsearch.search.fetch.scriptsource.ScriptSourceContext; import org.elasticsearch.search.highlight.SearchContextHighlight; import org.elasticsearch.search.lookup.SearchLookup; import org.elasticsearch.search.query.QuerySearchResult; @@ -127,6 +128,10 @@ public static SearchContext current() { private ScriptFieldsContext scriptFields; private PartialFieldsContext partialFields; + private ScriptSourceContext scriptSource; + + private List extractFieldNames; + private int from = -1; private int size = -1; @@ -288,6 +293,18 @@ public ScriptFieldsContext scriptFields() { return this.scriptFields; } + public void scriptSource(ScriptSourceContext scriptSourceContext) { + this.scriptSource = scriptSourceContext; + } + + public ScriptSourceContext scriptSource() { + return scriptSource; + } + + public boolean hasScriptSource() { + return scriptSource != null; + } + public boolean hasPartialFields() { return partialFields != null; } @@ -459,6 +476,18 @@ public void emptyFieldNames() { this.fieldNames = ImmutableList.of(); } + + public boolean hasExtractFieldNames() { + return extractFieldNames != null; + } + + public List extractFieldNames() { + if (extractFieldNames == null) { + extractFieldNames = Lists.newArrayList(); + } + return extractFieldNames; + } + public boolean explain() { return explain; } diff --git a/src/test/java/org/elasticsearch/test/integration/search/scriptsource/ScriptSourceTests.java b/src/test/java/org/elasticsearch/test/integration/search/scriptsource/ScriptSourceTests.java new file mode 100644 index 0000000000000..1a640c4dc60af --- /dev/null +++ b/src/test/java/org/elasticsearch/test/integration/search/scriptsource/ScriptSourceTests.java @@ -0,0 +1,296 @@ +/* + * Licensed to Elastic Search and Shay Banon under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. Elastic Search 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.elasticsearch.test.integration.search.scriptsource; + +import org.elasticsearch.action.search.SearchResponse; +import org.elasticsearch.client.Client; +import org.elasticsearch.common.xcontent.XContentFactory; +import org.elasticsearch.search.sort.SortOrder; +import org.elasticsearch.test.integration.AbstractNodesTests; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Test; + +import static org.elasticsearch.common.settings.ImmutableSettings.settingsBuilder; +import static org.elasticsearch.common.xcontent.XContentFactory.jsonBuilder; +import static org.elasticsearch.index.query.QueryBuilders.matchAllQuery; +import static org.elasticsearch.index.query.QueryBuilders.textQuery; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.*; + +/** + * + */ +public class ScriptSourceTests extends AbstractNodesTests { + + private Client client; + + @BeforeClass + public void createNodes() throws Exception { + startNode("node1", settingsBuilder().put("index.number_of_shards", 1).put("index.number_of_replicas", 0)); + startNode("client1", settingsBuilder().put("node.client", true).build()); + client = getClient(); + } + + @AfterClass + public void closeNodes() { + client.close(); + closeAllNodes(); + } + + protected Client getClient() { + return client("client1"); + } + + @Test + public void testWithSourceDisabled() throws Exception { + client.admin().indices().prepareDelete().execute().actionGet(); + client.admin().indices().prepareCreate("test").execute().actionGet(); + client.admin().cluster().prepareHealth().setWaitForYellowStatus().execute().actionGet(); + + String mapping = XContentFactory.jsonBuilder().startObject().startObject("type1") + .startObject("_source").field("enabled", false).endObject() + .startObject("_id").field("store", "yes").endObject() + .startObject("properties") + .startObject("text").field("type", "string").field("store", "no").endObject() + .startObject("num").field("type", "integer").field("store", "no").endObject() + .endObject().endObject().endObject().string(); + + client.admin().indices().preparePutMapping().setType("type1").setSource(mapping).execute().actionGet(); + + for (int i = 0; i < 10; i++) { + String id = String.valueOf(i); + client.prepareIndex("test", "type1", id).setSource(jsonBuilder().startObject() + .field("text", "Test value " + id) + .field("num", i + 10) + .endObject()).execute().actionGet(); + } + + client.admin().indices().prepareRefresh().execute().actionGet(); + + String sourceScript = "{" + + "\"text\" : \"Test value \" + _fields._id.value + \" from script\"," + + "\"num\" : 10 + (int)_fields._id.value" + + "}"; + + // Make sure that fields cannot be retrieved when source script is not set + SearchResponse searchResponse = client.prepareSearch() + .setQuery(matchAllQuery()).addField("text").addField("num") + .addSort("num", SortOrder.ASC) + .execute().actionGet(); + assertThat(searchResponse.hits().getTotalHits(), equalTo(10l)); + assertThat(searchResponse.hits().hits().length, equalTo(10)); + assertThat(searchResponse.hits().getAt(0).fields().size(), equalTo(0)); + + // Fields can be retrieved from script generated source + searchResponse = client.prepareSearch() + .setQuery(matchAllQuery()).addField("text").addField("num") + .setScriptSource(sourceScript) + .addSort("num", SortOrder.ASC) + .execute().actionGet(); + assertThat(searchResponse.hits().getTotalHits(), equalTo(10l)); + assertThat(searchResponse.hits().hits().length, equalTo(10)); + for (int i = 0; i < 10; i++) { + assertThat(searchResponse.hits().getAt(i).fields().size(), equalTo(2)); + assertThat(searchResponse.hits().getAt(i).fields().get("text").value().toString(), equalTo("Test value " + i + " from script")); + assertThat(searchResponse.hits().getAt(i).fields().get("num").value().toString(), equalTo(String.valueOf(i + 10))); + } + + // Check that highlighting works with generated source + searchResponse = client.prepareSearch() + .setQuery(textQuery("text", "value")) + .setScriptSource(sourceScript) + .addHighlightedField("text") + .addSort("num", SortOrder.ASC) + .execute().actionGet(); + assertThat(searchResponse.hits().getTotalHits(), equalTo(10l)); + assertThat(searchResponse.hits().hits().length, equalTo(10)); + for (int i = 0; i < 10; i++) { + assertThat(searchResponse.hits().getAt(i).highlightFields().get("text").fragments().length, equalTo(1)); + assertThat(searchResponse.hits().getAt(i).highlightFields().get("text").fragments()[0].string(), startsWith("Test value " + i + " from script")); + } + + // Check that script_fields work with generated source + searchResponse = client.prepareSearch() + .setQuery(matchAllQuery()) + .addScriptField("script_test", "_source.text.replace(\"script\", \"script_field\")") + .setScriptSource(sourceScript) + .addSort("num", SortOrder.ASC) + .execute().actionGet(); + assertThat(searchResponse.hits().getTotalHits(), equalTo(10l)); + assertThat(searchResponse.hits().hits().length, equalTo(10)); + for (int i = 0; i < 10; i++) { + assertThat(searchResponse.hits().getAt(i).fields().size(), equalTo(1)); + assertThat(searchResponse.hits().getAt(i).fields().get("script_test").value().toString(), equalTo("Test value " + i + " from script_field")); + } + + // Source returned as a string + searchResponse = client.prepareSearch() + .setQuery(textQuery("text", "value")) + .addField("text") + .setScriptSource("\"{\\\"text\\\":\\\"Test value 0 from script returning text\\\"}\"") + .addHighlightedField("text") + .addSort("num", SortOrder.ASC) + .setSize(1) + .execute().actionGet(); + assertThat(searchResponse.hits().getTotalHits(), equalTo(10l)); + assertThat(searchResponse.hits().hits().length, equalTo(1)); + assertThat(searchResponse.hits().getAt(0).fields().size(), equalTo(1)); + assertThat(searchResponse.hits().getAt(0).fields().get("text").value().toString(), equalTo("Test value 0 from script returning text")); + assertThat(searchResponse.hits().getAt(0).highlightFields().get("text").fragments().length, equalTo(1)); + assertThat(searchResponse.hits().getAt(0).highlightFields().get("text").fragments()[0].string(), startsWith("Test value 0 from script returning text")); + + } + + + @Test + public void testWithSourceEnabled() throws Exception { + client.admin().indices().prepareDelete().execute().actionGet(); + client.admin().indices().prepareCreate("test").execute().actionGet(); + client.admin().cluster().prepareHealth().setWaitForYellowStatus().execute().actionGet(); + + String mapping = XContentFactory.jsonBuilder().startObject().startObject("type1") + .startObject("_source").field("enabled", true).endObject() + .startObject("_id").field("store", "yes").endObject() + .startObject("properties") + .startObject("text").field("type", "string").field("store", "no").endObject() + .startObject("num").field("type", "integer").field("store", "no").endObject() + .endObject().endObject().endObject().string(); + + client.admin().indices().preparePutMapping().setType("type1").setSource(mapping).execute().actionGet(); + + + String[] originalSource = new String[10]; + for (int i = 0; i < 10; i++) { + String id = String.valueOf(i); + originalSource[i] = jsonBuilder().startObject() + .field("text", "Test value " + id) + .field("num", i + 10) + .endObject().string() + + "\n\n"; // Add \n\n to later make sure that original source wasn't reformatted. + client.prepareIndex("test", "type1", id).setSource(originalSource[i]).execute().actionGet(); + } + + client.admin().indices().prepareRefresh().execute().actionGet(); + + String sourceScript = "{" + + "\"text\" : \"Test value \" + _fields._id.value + \" from script\"," + + "\"num\" : 20 + (int)_fields._id.value" + + "}"; + + // Make sure that fields can be retrieved when source script is not set + SearchResponse searchResponse = client.prepareSearch() + .setQuery(matchAllQuery()).addField("text").addField("num") + .addSort("num", SortOrder.ASC) + .execute().actionGet(); + assertThat(searchResponse.hits().getTotalHits(), equalTo(10l)); + assertThat(searchResponse.hits().hits().length, equalTo(10)); + assertThat(searchResponse.hits().getAt(0).fields().size(), equalTo(2)); + assertThat(searchResponse.hits().getAt(0).fields().get("text").value().toString(), equalTo("Test value 0")); + assertThat(searchResponse.hits().getAt(0).fields().get("num").value().toString(), equalTo("10")); + + // Test that script generated source can be returned + searchResponse = client.prepareSearch() + .setQuery(matchAllQuery()) + .setScriptSource(sourceScript) + .addSort("num", SortOrder.ASC) + .execute().actionGet(); + assertThat(searchResponse.hits().getTotalHits(), equalTo(10l)); + assertThat(searchResponse.hits().hits().length, equalTo(10)); + for (int i = 0; i < 10; i++) { + assertThat(searchResponse.hits().getAt(i).fields().size(), equalTo(0)); + assertThat(searchResponse.hits().getAt(i).sourceAsString(), containsString("Test value " + i + " from script")); + assertThat(searchResponse.hits().getAt(i).sourceAsString(), containsString(String.valueOf(20 + i))); + } + + // Test that script generated source can replace stored source for fields + searchResponse = client.prepareSearch() + .setQuery(matchAllQuery()).addField("text").addField("num") + .setScriptSource(sourceScript) + .addSort("num", SortOrder.ASC) + .execute().actionGet(); + assertThat(searchResponse.hits().getTotalHits(), equalTo(10l)); + assertThat(searchResponse.hits().hits().length, equalTo(10)); + for (int i = 0; i < 10; i++) { + assertThat(searchResponse.hits().getAt(i).fields().size(), equalTo(2)); + assertThat(searchResponse.hits().getAt(i).fields().get("text").value().toString(), equalTo("Test value " + i + " from script")); + assertThat(searchResponse.hits().getAt(i).fields().get("num").value().toString(), equalTo(String.valueOf(i + 20))); + // We didn't ask for source so it shouldn't be returned + assertThat(searchResponse.hits().getAt(i).sourceAsString(), nullValue()); + } + + // Check that highlighting works with generated source + searchResponse = client.prepareSearch() + .setQuery(textQuery("text", "value")) + .setScriptSource(sourceScript) + .addHighlightedField("text") + .addSort("num", SortOrder.ASC) + .execute().actionGet(); + assertThat(searchResponse.hits().getTotalHits(), equalTo(10l)); + assertThat(searchResponse.hits().hits().length, equalTo(10)); + for (int i = 0; i < 10; i++) { + assertThat(searchResponse.hits().getAt(i).highlightFields().get("text").fragments().length, equalTo(1)); + assertThat(searchResponse.hits().getAt(i).highlightFields().get("text").fragments()[0].string(), startsWith("Test value " + i + " from script")); + } + + // Check that script_fields work with generated source + searchResponse = client.prepareSearch() + .setQuery(textQuery("text", "value")) + .addScriptField("script_test", "_source.text.replace(\"script\", \"script_field\")") + .addHighlightedField("text") + .setScriptSource(sourceScript) + .addSort("num", SortOrder.ASC) + .execute().actionGet(); + assertThat(searchResponse.hits().getTotalHits(), equalTo(10l)); + assertThat(searchResponse.hits().hits().length, equalTo(10)); + for (int i = 0; i < 10; i++) { + assertThat(searchResponse.hits().getAt(i).fields().size(), equalTo(1)); + assertThat(searchResponse.hits().getAt(i).fields().get("script_test").value().toString(), equalTo("Test value " + i + " from script_field")); + } + + // Test that script can turn off + searchResponse = client.prepareSearch() + .setQuery(textQuery("text", "value")) + .addField("text") + .addField("num") + .setScriptSource("null") + .addSort("num", SortOrder.ASC) + .setSize(1) + .execute().actionGet(); + assertThat(searchResponse.hits().getTotalHits(), equalTo(10l)); + assertThat(searchResponse.hits().hits().length, equalTo(1)); + assertThat(searchResponse.hits().getAt(0).fields().size(), equalTo(0)); + + // Make sure that original source can be left intact + searchResponse = client.prepareSearch() + .setQuery(matchAllQuery()) + .addHighlightedField("text") + .setScriptSource("_source") + .addSort("num", SortOrder.ASC) + .execute().actionGet(); + assertThat(searchResponse.hits().getTotalHits(), equalTo(10l)); + assertThat(searchResponse.hits().hits().length, equalTo(10)); + for (int i = 0; i < 10; i++) { + assertThat(searchResponse.hits().getAt(i).sourceAsString(), equalTo(originalSource[i])); + } + + + } +}