From 335afe5b882fd06d31e991d8fe4357acf5e0bd3d Mon Sep 17 00:00:00 2001 From: Jim Ferenczi Date: Tue, 5 Mar 2024 16:06:20 +0000 Subject: [PATCH 01/27] Fix performance bug in `SourceConfirmedTextQuery#matches` (#105930) This change ensures that the matches implementation of the `SourceConfirmedTextQuery` only checks the current document instead of calling advance on the two phase iterator. The latter tries to find the first doc that matches the query instead of restricting the search to the current doc. This can lead to abnormally slow highlighting if the query is very restrictive and the highlight is done on a non-matching document. Closes #103298 --- .../extras/SourceConfirmedTextQuery.java | 62 ++++++++-------- .../extras/SourceConfirmedTextQueryTests.java | 71 +++++++++++++------ 2 files changed, 80 insertions(+), 53 deletions(-) diff --git a/modules/mapper-extras/src/main/java/org/elasticsearch/index/mapper/extras/SourceConfirmedTextQuery.java b/modules/mapper-extras/src/main/java/org/elasticsearch/index/mapper/extras/SourceConfirmedTextQuery.java index dc51afe5d420d..3d0f26e8cc130 100644 --- a/modules/mapper-extras/src/main/java/org/elasticsearch/index/mapper/extras/SourceConfirmedTextQuery.java +++ b/modules/mapper-extras/src/main/java/org/elasticsearch/index/mapper/extras/SourceConfirmedTextQuery.java @@ -9,9 +9,7 @@ package org.elasticsearch.index.mapper.extras; import org.apache.lucene.analysis.Analyzer; -import org.apache.lucene.index.FieldInfo; import org.apache.lucene.index.FieldInvertState; -import org.apache.lucene.index.IndexOptions; import org.apache.lucene.index.LeafReaderContext; import org.apache.lucene.index.Term; import org.apache.lucene.index.TermStates; @@ -300,19 +298,23 @@ public RuntimePhraseScorer scorer(LeafReaderContext context) throws IOException @Override public Matches matches(LeafReaderContext context, int doc) throws IOException { - FieldInfo fi = context.reader().getFieldInfos().fieldInfo(field); - if (fi == null) { + var terms = context.reader().terms(field); + if (terms == null) { return null; } - // Some highlighters will already have reindexed the source with positions and offsets, + // Some highlighters will already have re-indexed the source with positions and offsets, // so rather than doing it again we check to see if this data is available on the // current context and if so delegate directly to the inner query - if (fi.getIndexOptions().compareTo(IndexOptions.DOCS_AND_FREQS_AND_POSITIONS) > 0) { + if (terms.hasOffsets()) { Weight innerWeight = in.createWeight(searcher, ScoreMode.COMPLETE_NO_SCORES, 1); return innerWeight.matches(context, doc); } RuntimePhraseScorer scorer = scorer(context); - if (scorer == null || scorer.iterator().advance(doc) != doc) { + if (scorer == null) { + return null; + } + final TwoPhaseIterator twoPhase = scorer.twoPhaseIterator(); + if (twoPhase.approximation().advance(doc) != doc || scorer.twoPhaseIterator().matches() == false) { return null; } return scorer.matches(); @@ -321,13 +323,14 @@ public Matches matches(LeafReaderContext context, int doc) throws IOException { } private class RuntimePhraseScorer extends Scorer { - private final LeafSimScorer scorer; private final CheckedIntFunction, IOException> valueFetcher; private final String field; private final Query query; private final TwoPhaseIterator twoPhase; + private final MemoryIndexEntry cacheEntry = new MemoryIndexEntry(); + private int doc = -1; private float freq; @@ -357,7 +360,6 @@ public float matchCost() { // Defaults to a high-ish value so that it likely runs last. return 10_000f; } - }; } @@ -394,35 +396,35 @@ private float freq() throws IOException { return freq; } - private float computeFreq() throws IOException { - MemoryIndex index = new MemoryIndex(); - index.setSimilarity(FREQ_SIMILARITY); - List values = valueFetcher.apply(docID()); - float frequency = 0; - for (Object value : values) { - if (value == null) { - continue; + private MemoryIndex getOrCreateMemoryIndex() throws IOException { + if (cacheEntry.docID != docID()) { + cacheEntry.docID = docID(); + cacheEntry.memoryIndex = new MemoryIndex(true, false); + cacheEntry.memoryIndex.setSimilarity(FREQ_SIMILARITY); + List values = valueFetcher.apply(docID()); + for (Object value : values) { + if (value == null) { + continue; + } + cacheEntry.memoryIndex.addField(field, value.toString(), indexAnalyzer); } - index.addField(field, value.toString(), indexAnalyzer); - frequency += index.search(query); - index.reset(); } - return frequency; + return cacheEntry.memoryIndex; + } + + private float computeFreq() throws IOException { + return getOrCreateMemoryIndex().search(query); } private Matches matches() throws IOException { - MemoryIndex index = new MemoryIndex(true, false); - List values = valueFetcher.apply(docID()); - for (Object value : values) { - if (value == null) { - continue; - } - index.addField(field, value.toString(), indexAnalyzer); - } - IndexSearcher searcher = index.createSearcher(); + IndexSearcher searcher = getOrCreateMemoryIndex().createSearcher(); Weight w = searcher.createWeight(searcher.rewrite(query), ScoreMode.COMPLETE_NO_SCORES, 1); return w.matches(searcher.getLeafContexts().get(0), 0); } } + private static class MemoryIndexEntry { + private int docID = -1; + private MemoryIndex memoryIndex; + } } diff --git a/modules/mapper-extras/src/test/java/org/elasticsearch/index/mapper/extras/SourceConfirmedTextQueryTests.java b/modules/mapper-extras/src/test/java/org/elasticsearch/index/mapper/extras/SourceConfirmedTextQueryTests.java index 2b8d5870cb8aa..81e1dd7099860 100644 --- a/modules/mapper-extras/src/test/java/org/elasticsearch/index/mapper/extras/SourceConfirmedTextQueryTests.java +++ b/modules/mapper-extras/src/test/java/org/elasticsearch/index/mapper/extras/SourceConfirmedTextQueryTests.java @@ -49,13 +49,19 @@ import java.io.IOException; import java.util.Collections; import java.util.List; +import java.util.concurrent.atomic.AtomicInteger; +import static org.hamcrest.Matchers.equalTo; import static org.hamcrest.Matchers.greaterThan; public class SourceConfirmedTextQueryTests extends ESTestCase { + private static final AtomicInteger sourceFetchCount = new AtomicInteger(); private static final IOFunction, IOException>> SOURCE_FETCHER_PROVIDER = - context -> docID -> Collections.singletonList(context.reader().document(docID).get("body")); + context -> docID -> { + sourceFetchCount.incrementAndGet(); + return Collections.singletonList(context.reader().document(docID).get("body")); + }; public void testTerm() throws Exception { try (Directory dir = newDirectory(); IndexWriter w = new IndexWriter(dir, newIndexWriterConfig(Lucene.STANDARD_ANALYZER))) { @@ -440,11 +446,11 @@ public void testEmptyIndex() throws Exception { } public void testMatches() throws Exception { - checkMatches(new TermQuery(new Term("body", "d")), "a b c d e", new int[] { 3, 3 }); - checkMatches(new PhraseQuery("body", "b", "c"), "a b c d c b c a", new int[] { 1, 2, 5, 6 }); + checkMatches(new TermQuery(new Term("body", "d")), "a b c d e", new int[] { 3, 3 }, false); + checkMatches(new PhraseQuery("body", "b", "c"), "a b c d c b c a", new int[] { 1, 2, 5, 6 }, true); } - private static void checkMatches(Query query, String inputDoc, int[] expectedMatches) throws IOException { + private static void checkMatches(Query query, String inputDoc, int[] expectedMatches, boolean expectedFetch) throws IOException { try (Directory dir = newDirectory(); IndexWriter w = new IndexWriter(dir, newIndexWriterConfig(Lucene.STANDARD_ANALYZER))) { Document doc = new Document(); doc.add(new TextField("body", "xxxxxnomatchxxxx", Store.YES)); @@ -464,30 +470,48 @@ private static void checkMatches(Query query, String inputDoc, int[] expectedMat Query sourceConfirmedQuery = new SourceConfirmedTextQuery(query, SOURCE_FETCHER_PROVIDER, Lucene.STANDARD_ANALYZER); try (IndexReader ir = DirectoryReader.open(w)) { - - IndexSearcher searcher = new IndexSearcher(ir); - TopDocs td = searcher.search( - sourceConfirmedQuery, - 3, - new Sort(KeywordField.newSortField("sort", false, SortedSetSelector.Type.MAX)) - ); - - Weight weight = searcher.createWeight(searcher.rewrite(sourceConfirmedQuery), ScoreMode.COMPLETE_NO_SCORES, 1); - - int firstDoc = td.scoreDocs[0].doc; - LeafReaderContext firstCtx = searcher.getLeafContexts().get(ReaderUtil.subIndex(firstDoc, searcher.getLeafContexts())); - checkMatches(weight, firstCtx, firstDoc - firstCtx.docBase, expectedMatches, 0); - - int secondDoc = td.scoreDocs[1].doc; - LeafReaderContext secondCtx = searcher.getLeafContexts().get(ReaderUtil.subIndex(secondDoc, searcher.getLeafContexts())); - checkMatches(weight, secondCtx, secondDoc - secondCtx.docBase, expectedMatches, 1); - + { + IndexSearcher searcher = new IndexSearcher(ir); + TopDocs td = searcher.search( + sourceConfirmedQuery, + 3, + new Sort(KeywordField.newSortField("sort", false, SortedSetSelector.Type.MAX)) + ); + + Weight weight = searcher.createWeight(searcher.rewrite(sourceConfirmedQuery), ScoreMode.COMPLETE_NO_SCORES, 1); + + int firstDoc = td.scoreDocs[0].doc; + LeafReaderContext firstCtx = searcher.getLeafContexts().get(ReaderUtil.subIndex(firstDoc, searcher.getLeafContexts())); + checkMatches(weight, firstCtx, firstDoc - firstCtx.docBase, expectedMatches, 0, expectedFetch); + + int secondDoc = td.scoreDocs[1].doc; + LeafReaderContext secondCtx = searcher.getLeafContexts() + .get(ReaderUtil.subIndex(secondDoc, searcher.getLeafContexts())); + checkMatches(weight, secondCtx, secondDoc - secondCtx.docBase, expectedMatches, 1, expectedFetch); + } + + { + IndexSearcher searcher = new IndexSearcher(ir); + TopDocs td = searcher.search(KeywordField.newExactQuery("sort", "0"), 1); + + Weight weight = searcher.createWeight(searcher.rewrite(sourceConfirmedQuery), ScoreMode.COMPLETE_NO_SCORES, 1); + int firstDoc = td.scoreDocs[0].doc; + LeafReaderContext firstCtx = searcher.getLeafContexts().get(ReaderUtil.subIndex(firstDoc, searcher.getLeafContexts())); + checkMatches(weight, firstCtx, firstDoc - firstCtx.docBase, new int[0], 0, false); + } } } } - private static void checkMatches(Weight w, LeafReaderContext ctx, int doc, int[] expectedMatches, int offset) throws IOException { + private static void checkMatches(Weight w, LeafReaderContext ctx, int doc, int[] expectedMatches, int offset, boolean expectedFetch) + throws IOException { + int count = sourceFetchCount.get(); Matches matches = w.matches(ctx, doc); + if (expectedMatches.length == 0) { + assertNull(matches); + assertThat(sourceFetchCount.get() - count, equalTo(expectedFetch ? 1 : 0)); + return; + } assertNotNull(matches); MatchesIterator mi = matches.getMatches("body"); int i = 0; @@ -498,6 +522,7 @@ private static void checkMatches(Weight w, LeafReaderContext ctx, int doc, int[] i += 2; } assertEquals(expectedMatches.length, i); + assertThat(sourceFetchCount.get() - count, equalTo(expectedFetch ? 1 : 0)); } } From fe13a04a5437017381b71d41afb63cb999bf62c4 Mon Sep 17 00:00:00 2001 From: Kathleen DeRusso Date: Tue, 5 Mar 2024 17:21:23 +0100 Subject: [PATCH 02/27] Bugfix for mixed version cluster queries using text expansion (#105912) * Bugfix for CCR queries using text expansion * Fix test * PR feedback * Fix test * Minor cleanup * Edit comment * One more comment clarification --------- Co-authored-by: Elastic Machine --- .../ml/queries/TextExpansionQueryBuilder.java | 48 +++++++++++-------- .../TextExpansionQueryBuilderTests.java | 7 ++- .../test/ml/text_expansion_search.yml | 1 + 3 files changed, 35 insertions(+), 21 deletions(-) diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilder.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilder.java index 675d062fdb3af..f6fa7ca9005c5 100644 --- a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilder.java +++ b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilder.java @@ -18,6 +18,7 @@ import org.elasticsearch.core.Nullable; import org.elasticsearch.index.query.AbstractQueryBuilder; import org.elasticsearch.index.query.QueryBuilder; +import org.elasticsearch.index.query.QueryBuilders; import org.elasticsearch.index.query.QueryRewriteContext; import org.elasticsearch.index.query.SearchExecutionContext; import org.elasticsearch.xcontent.ParseField; @@ -67,12 +68,7 @@ public String getTypeName() { } public static boolean isFieldTypeAllowed(String typeName) { - for (AllowedFieldType fieldType : values()) { - if (fieldType.getTypeName().equals(typeName)) { - return true; - } - } - return false; + return Arrays.stream(values()).anyMatch(value -> value.typeName.equals(typeName)); } public static String getAllowedFieldTypesAsString() { @@ -168,8 +164,7 @@ protected void doXContent(XContentBuilder builder, Params params) throws IOExcep } @Override - protected QueryBuilder doRewrite(QueryRewriteContext queryRewriteContext) throws IOException { - + protected QueryBuilder doRewrite(QueryRewriteContext queryRewriteContext) { if (weightedTokensSupplier != null) { if (weightedTokensSupplier.get() == null) { return this; @@ -188,8 +183,8 @@ protected QueryBuilder doRewrite(QueryRewriteContext queryRewriteContext) throws inferRequest.setPrefixType(TrainedModelPrefixStrings.PrefixType.SEARCH); SetOnce textExpansionResultsSupplier = new SetOnce<>(); - queryRewriteContext.registerAsyncAction((client, listener) -> { - executeAsyncWithOrigin( + queryRewriteContext.registerAsyncAction( + (client, listener) -> executeAsyncWithOrigin( client, ML_ORIGIN, CoordinatedInferenceAction.INSTANCE, @@ -220,21 +215,34 @@ protected QueryBuilder doRewrite(QueryRewriteContext queryRewriteContext) throws ); } }, listener::onFailure) - ); - }); + ) + ); return new TextExpansionQueryBuilder(this, textExpansionResultsSupplier); } private QueryBuilder weightedTokensToQuery(String fieldName, TextExpansionResults textExpansionResults) { - WeightedTokensQueryBuilder weightedTokensQueryBuilder = new WeightedTokensQueryBuilder( - fieldName, - textExpansionResults.getWeightedTokens(), - tokenPruningConfig - ); - weightedTokensQueryBuilder.queryName(queryName); - weightedTokensQueryBuilder.boost(boost); - return weightedTokensQueryBuilder; + if (tokenPruningConfig != null) { + WeightedTokensQueryBuilder weightedTokensQueryBuilder = new WeightedTokensQueryBuilder( + fieldName, + textExpansionResults.getWeightedTokens(), + tokenPruningConfig + ); + weightedTokensQueryBuilder.queryName(queryName); + weightedTokensQueryBuilder.boost(boost); + return weightedTokensQueryBuilder; + } + // Note: Weighted tokens queries were introduced in 8.13.0. To support mixed version clusters prior to 8.13.0, + // if no token pruning configuration is specified we fall back to a boolean query. + // TODO this should be updated to always use a WeightedTokensQueryBuilder once it's in all supported versions. + var boolQuery = QueryBuilders.boolQuery(); + for (var weightedToken : textExpansionResults.getWeightedTokens()) { + boolQuery.should(QueryBuilders.termQuery(fieldName, weightedToken.token()).boost(weightedToken.weight())); + } + boolQuery.minimumShouldMatch(1); + boolQuery.boost(boost); + boolQuery.queryName(queryName); + return boolQuery; } @Override diff --git a/x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilderTests.java b/x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilderTests.java index 50561d92f5d37..13f12f3cdc1e1 100644 --- a/x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilderTests.java +++ b/x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/queries/TextExpansionQueryBuilderTests.java @@ -25,6 +25,7 @@ import org.elasticsearch.common.compress.CompressedXContent; import org.elasticsearch.index.mapper.MapperService; import org.elasticsearch.index.mapper.extras.MapperExtrasPlugin; +import org.elasticsearch.index.query.BoolQueryBuilder; import org.elasticsearch.index.query.QueryBuilder; import org.elasticsearch.index.query.SearchExecutionContext; import org.elasticsearch.plugins.Plugin; @@ -259,6 +260,10 @@ public void testThatTokensAreCorrectlyPruned() { SearchExecutionContext searchExecutionContext = createSearchExecutionContext(); TextExpansionQueryBuilder queryBuilder = createTestQueryBuilder(); QueryBuilder rewrittenQueryBuilder = rewriteAndFetch(queryBuilder, searchExecutionContext); - assertTrue(rewrittenQueryBuilder instanceof WeightedTokensQueryBuilder); + if (queryBuilder.getTokenPruningConfig() == null) { + assertTrue(rewrittenQueryBuilder instanceof BoolQueryBuilder); + } else { + assertTrue(rewrittenQueryBuilder instanceof WeightedTokensQueryBuilder); + } } } diff --git a/x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/ml/text_expansion_search.yml b/x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/ml/text_expansion_search.yml index dc4e1751ccdee..f92870b61f1b1 100644 --- a/x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/ml/text_expansion_search.yml +++ b/x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/ml/text_expansion_search.yml @@ -304,3 +304,4 @@ setup: source_text: model_id: text_expansion_model model_text: "octopus comforter smells" + pruning_config: {} From 3e5c3c523ddce71ad2f9c4d28795726b1aed1e10 Mon Sep 17 00:00:00 2001 From: Jonathan Buttner <56361221+jonathan-buttner@users.noreply.github.com> Date: Tue, 5 Mar 2024 11:33:36 -0500 Subject: [PATCH 03/27] [ML] Enable retrying on 500 error response from Cohere text embedding API (#105797) * Retrying on 500 * Update docs/changelog/105797.yaml --------- Co-authored-by: Elastic Machine --- docs/changelog/105797.yaml | 5 +++++ .../external/cohere/CohereResponseHandler.java | 4 +++- .../external/cohere/CohereResponseHandlerTests.java | 10 ++++++++++ 3 files changed, 18 insertions(+), 1 deletion(-) create mode 100644 docs/changelog/105797.yaml diff --git a/docs/changelog/105797.yaml b/docs/changelog/105797.yaml new file mode 100644 index 0000000000000..7c832e2e5e63c --- /dev/null +++ b/docs/changelog/105797.yaml @@ -0,0 +1,5 @@ +pr: 105797 +summary: Enable retrying on 500 error response from Cohere text embedding API +area: Machine Learning +type: enhancement +issues: [] diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/cohere/CohereResponseHandler.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/cohere/CohereResponseHandler.java index c7e6493949400..b5af0b474834f 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/cohere/CohereResponseHandler.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/cohere/CohereResponseHandler.java @@ -59,7 +59,9 @@ void checkForFailureStatusCode(Request request, HttpResult result) throws RetryE } // handle error codes - if (statusCode >= 500) { + if (statusCode == 500) { + throw new RetryException(true, buildError(SERVER_ERROR, request, result)); + } else if (statusCode > 500) { throw new RetryException(false, buildError(SERVER_ERROR, request, result)); } else if (statusCode == 429) { throw new RetryException(true, buildError(RATE_LIMIT, request, result)); diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/cohere/CohereResponseHandlerTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/cohere/CohereResponseHandlerTests.java index 31945d5a8b4fc..d64ac495c8c99 100644 --- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/cohere/CohereResponseHandlerTests.java +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/cohere/CohereResponseHandlerTests.java @@ -44,6 +44,16 @@ public void testCheckForFailureStatusCode_ThrowsFor503() { MatcherAssert.assertThat(((ElasticsearchStatusException) exception.getCause()).status(), is(RestStatus.BAD_REQUEST)); } + public void testCheckForFailureStatusCode_ThrowsFor500_WithShouldRetryTrue() { + var exception = expectThrows(RetryException.class, () -> callCheckForFailureStatusCode(500, "id")); + assertTrue(exception.shouldRetry()); + MatcherAssert.assertThat( + exception.getCause().getMessage(), + containsString("Received a server error status code for request from inference entity id [id] status [500]") + ); + MatcherAssert.assertThat(((ElasticsearchStatusException) exception.getCause()).status(), is(RestStatus.BAD_REQUEST)); + } + public void testCheckForFailureStatusCode_ThrowsFor429() { var exception = expectThrows(RetryException.class, () -> callCheckForFailureStatusCode(429, "id")); assertTrue(exception.shouldRetry()); From f1035bba1ecf18181697a5ab6329c5c6c43a4c3b Mon Sep 17 00:00:00 2001 From: Michael Peterson Date: Tue, 5 Mar 2024 11:41:07 -0500 Subject: [PATCH 04/27] Adjust randomization in ResolveClusterActionResponseTests (#105932) to avoid failures in `testEqualsAndHashcode` tests. Fixes https://github.com/elastic/elasticsearch/issues/105898 --- .../ResolveClusterActionResponseTests.java | 29 +++++++++++++++---- 1 file changed, 23 insertions(+), 6 deletions(-) diff --git a/server/src/test/java/org/elasticsearch/action/admin/indices/resolve/ResolveClusterActionResponseTests.java b/server/src/test/java/org/elasticsearch/action/admin/indices/resolve/ResolveClusterActionResponseTests.java index 322600fdeedff..33d4f0edf3450 100644 --- a/server/src/test/java/org/elasticsearch/action/admin/indices/resolve/ResolveClusterActionResponseTests.java +++ b/server/src/test/java/org/elasticsearch/action/admin/indices/resolve/ResolveClusterActionResponseTests.java @@ -20,22 +20,39 @@ public class ResolveClusterActionResponseTests extends AbstractWireSerializingTe @Override protected ResolveClusterActionResponse createTestInstance() { - return new ResolveClusterActionResponse(randomResolveClusterInfoMap()); + return new ResolveClusterActionResponse(randomResolveClusterInfoMap(null)); } - private Map randomResolveClusterInfoMap() { + private ResolveClusterInfo randomResolveClusterInfo(ResolveClusterInfo existing) { + if (existing == null) { + return randomResolveClusterInfo(); + } else { + return randomValueOtherThan(existing, () -> randomResolveClusterInfo()); + } + } + + private ResolveClusterInfo getResolveClusterInfoFromResponse(String key, ResolveClusterActionResponse response) { + if (response == null || response.getResolveClusterInfo() == null) { + return null; + } + return response.getResolveClusterInfo().get(key); + } + + private Map randomResolveClusterInfoMap(ResolveClusterActionResponse existingResponse) { Map infoMap = new HashMap<>(); int numClusters = randomIntBetween(0, 50); if (randomBoolean() || numClusters == 0) { - infoMap.put(RemoteClusterAware.LOCAL_CLUSTER_GROUP_KEY, randomResolveClusterInfo()); + String key = RemoteClusterAware.LOCAL_CLUSTER_GROUP_KEY; + infoMap.put(key, randomResolveClusterInfo(getResolveClusterInfoFromResponse(key, existingResponse))); } for (int i = 0; i < numClusters; i++) { - infoMap.put("remote_" + i, randomResolveClusterInfo()); + String key = "remote_" + i; + infoMap.put(key, randomResolveClusterInfo(getResolveClusterInfoFromResponse(key, existingResponse))); } return infoMap; } - private ResolveClusterInfo randomResolveClusterInfo() { + static ResolveClusterInfo randomResolveClusterInfo() { int val = randomIntBetween(1, 3); return switch (val) { case 1 -> new ResolveClusterInfo(false, randomBoolean()); @@ -52,6 +69,6 @@ protected Writeable.Reader instanceReader() { @Override protected ResolveClusterActionResponse mutateInstance(ResolveClusterActionResponse response) { - return new ResolveClusterActionResponse(randomResolveClusterInfoMap()); + return new ResolveClusterActionResponse(randomResolveClusterInfoMap(response)); } } From 85cd2043176ee2cfb8078ad685be2e1156b0c71d Mon Sep 17 00:00:00 2001 From: Craig Taverner Date: Tue, 5 Mar 2024 18:04:09 +0100 Subject: [PATCH 05/27] Fix automatic generation of spatial function types files (#105766) * Fix automatic generation of spatial function types files The automatic mapping of spatial function names from class names was not working for spatial types, so the automatic generation of these files did not happen, and in fact existing files were deleted. In addition, the generation of aggregation functions types does not yet exist at all, so the st_centroid.asciidoc file was always deleted. Until such support exists, this files contents will be moved back into the function definition file. The railroad diagrams for syntax are now also created, however, not all functions in the documentation actually use these, and certainly none of the `TO_*` type-casting functions do, so we'll not include links to them from the docs, and leave that to the docs team to decide. Personally, while these diagrams are pretty, they contain no additional informational content, and in fact give a cluttered impression to the documentation visual appeal. * Refined to use an annotation which is more generic --- .../functions/signature/to_cartesianpoint.svg | 1 + .../functions/signature/to_cartesianshape.svg | 1 + .../esql/functions/signature/to_geopoint.svg | 1 + .../esql/functions/signature/to_geoshape.svg | 1 + .../esql/functions/st_centroid.asciidoc | 7 ++++- .../esql/functions/types/st_centroid.asciidoc | 6 ---- .../function/AbstractFunctionTestCase.java | 13 ++++++-- .../expression/function/FunctionName.java | 30 +++++++++++++++++++ .../scalar/convert/ToCartesianPointTests.java | 2 ++ .../scalar/convert/ToCartesianShapeTests.java | 2 ++ .../scalar/convert/ToGeoPointTests.java | 2 ++ .../scalar/convert/ToGeoShapeTests.java | 2 ++ 12 files changed, 58 insertions(+), 10 deletions(-) create mode 100644 docs/reference/esql/functions/signature/to_cartesianpoint.svg create mode 100644 docs/reference/esql/functions/signature/to_cartesianshape.svg create mode 100644 docs/reference/esql/functions/signature/to_geopoint.svg create mode 100644 docs/reference/esql/functions/signature/to_geoshape.svg delete mode 100644 docs/reference/esql/functions/types/st_centroid.asciidoc create mode 100644 x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/FunctionName.java diff --git a/docs/reference/esql/functions/signature/to_cartesianpoint.svg b/docs/reference/esql/functions/signature/to_cartesianpoint.svg new file mode 100644 index 0000000000000..44484e8321e2f --- /dev/null +++ b/docs/reference/esql/functions/signature/to_cartesianpoint.svg @@ -0,0 +1 @@ +TO_CARTESIANPOINT(v) \ No newline at end of file diff --git a/docs/reference/esql/functions/signature/to_cartesianshape.svg b/docs/reference/esql/functions/signature/to_cartesianshape.svg new file mode 100644 index 0000000000000..c16ce9a6c15bc --- /dev/null +++ b/docs/reference/esql/functions/signature/to_cartesianshape.svg @@ -0,0 +1 @@ +TO_CARTESIANSHAPE(v) \ No newline at end of file diff --git a/docs/reference/esql/functions/signature/to_geopoint.svg b/docs/reference/esql/functions/signature/to_geopoint.svg new file mode 100644 index 0000000000000..444817aa388b9 --- /dev/null +++ b/docs/reference/esql/functions/signature/to_geopoint.svg @@ -0,0 +1 @@ +TO_GEOPOINT(v) \ No newline at end of file diff --git a/docs/reference/esql/functions/signature/to_geoshape.svg b/docs/reference/esql/functions/signature/to_geoshape.svg new file mode 100644 index 0000000000000..91b02332ad806 --- /dev/null +++ b/docs/reference/esql/functions/signature/to_geoshape.svg @@ -0,0 +1 @@ +TO_GEOSHAPE(v) \ No newline at end of file diff --git a/docs/reference/esql/functions/st_centroid.asciidoc b/docs/reference/esql/functions/st_centroid.asciidoc index abed1e71eab8f..cee0c85d5cb45 100644 --- a/docs/reference/esql/functions/st_centroid.asciidoc +++ b/docs/reference/esql/functions/st_centroid.asciidoc @@ -15,4 +15,9 @@ include::{esql-specs}/spatial.csv-spec[tag=st_centroid-airports-result] Supported types: -include::types/st_centroid.asciidoc[] +[%header.monospaced.styled,format=dsv,separator=|] +|=== +v | result +geo_point | geo_point +cartesian_point | cartesian_point +|=== diff --git a/docs/reference/esql/functions/types/st_centroid.asciidoc b/docs/reference/esql/functions/types/st_centroid.asciidoc deleted file mode 100644 index cbafb9d0fa6dc..0000000000000 --- a/docs/reference/esql/functions/types/st_centroid.asciidoc +++ /dev/null @@ -1,6 +0,0 @@ -[%header.monospaced.styled,format=dsv,separator=|] -|=== -v | result -geo_point | geo_point -cartesian_point | cartesian_point -|=== diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/AbstractFunctionTestCase.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/AbstractFunctionTestCase.java index dded86fdd8aee..612861b2889a4 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/AbstractFunctionTestCase.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/AbstractFunctionTestCase.java @@ -1088,7 +1088,7 @@ public static void renderTypesTable() throws IOException { renderTypesTable(EsqlFunctionRegistry.description(definition).argNames()); return; } - LogManager.getLogger(getTestClass()).info("Skipping rendering types because the function isn't registered"); + LogManager.getLogger(getTestClass()).info("Skipping rendering types because the function '" + name + "' isn't registered"); } private static void renderTypesTable(List argNames) throws IOException { @@ -1116,12 +1116,18 @@ private static void renderTypesTable(List argNames) throws IOException { [%header.monospaced.styled,format=dsv,separator=|] |=== """ + header + "\n" + table.stream().collect(Collectors.joining("\n")) + "\n|===\n"; - LogManager.getLogger(getTestClass()).info("Writing function types:\n{}", rendered); + LogManager.getLogger(getTestClass()).info("Writing function types for [{}]:\n{}", functionName(), rendered); writeToTempDir("types", rendered, "asciidoc"); } private static String functionName() { - return StringUtils.camelCaseToUnderscore(getTestClass().getSimpleName().replace("Tests", "")).toLowerCase(Locale.ROOT); + Class testClass = getTestClass(); + if (testClass.isAnnotationPresent(FunctionName.class)) { + FunctionName functionNameAnnotation = testClass.getAnnotation(FunctionName.class); + return functionNameAnnotation.value(); + } else { + return StringUtils.camelCaseToUnderscore(testClass.getSimpleName().replace("Tests", "")).toLowerCase(Locale.ROOT); + } } private static FunctionDefinition definition(String name) { @@ -1178,6 +1184,7 @@ private static void writeToTempDir(String subdir, String str, String extension) Files.createDirectories(dir); Path file = dir.resolve(functionName() + "." + extension); Files.writeString(file, str); + LogManager.getLogger(getTestClass()).info("Wrote function types for [{}] to file: {}", functionName(), file); } private final List breakers = Collections.synchronizedList(new ArrayList<>()); diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/FunctionName.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/FunctionName.java new file mode 100644 index 0000000000000..b4a5d3bdc2b92 --- /dev/null +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/FunctionName.java @@ -0,0 +1,30 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.esql.expression.function; + +import java.lang.annotation.ElementType; +import java.lang.annotation.Retention; +import java.lang.annotation.RetentionPolicy; +import java.lang.annotation.Target; + +/** + * Tests that extend AbstractFunctionTestCase can use this annotation to specify the name of the function + * to use when generating documentation files while running tests. + * If this is not used, the name will be deduced from the test class name, by removing the "Test" suffix, and converting + * the class name to snake case. This annotation can be used to override that behavior, for cases where the deduced name + * is not correct. For example, in Elasticsearch the class name for `GeoPoint` capitalizes the `P` in `Point`, but the + * function name is `to_geopoint`, not `to_geo_point`. In some cases, even when compatible class names are used, + * like `StX` for the function `st_x`, the annotation is needed because the name deduction does not allow only a single + * character after the underscore. + */ +@Retention(RetentionPolicy.RUNTIME) +@Target(ElementType.TYPE) +public @interface FunctionName { + /** The function name to use in generating documentation files while running tests */ + String value(); +} diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/convert/ToCartesianPointTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/convert/ToCartesianPointTests.java index 88910320c962e..4eadf88992582 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/convert/ToCartesianPointTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/convert/ToCartesianPointTests.java @@ -13,6 +13,7 @@ import org.apache.lucene.util.BytesRef; import org.elasticsearch.geo.ShapeTestUtils; import org.elasticsearch.xpack.esql.expression.function.AbstractFunctionTestCase; +import org.elasticsearch.xpack.esql.expression.function.FunctionName; import org.elasticsearch.xpack.esql.expression.function.TestCaseSupplier; import org.elasticsearch.xpack.esql.type.EsqlDataTypes; import org.elasticsearch.xpack.ql.expression.Expression; @@ -26,6 +27,7 @@ import static org.elasticsearch.xpack.ql.util.SpatialCoordinateTypes.CARTESIAN; +@FunctionName("to_cartesianpoint") public class ToCartesianPointTests extends AbstractFunctionTestCase { public ToCartesianPointTests(@Name("TestCase") Supplier testCaseSupplier) { this.testCase = testCaseSupplier.get(); diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/convert/ToCartesianShapeTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/convert/ToCartesianShapeTests.java index 117968de5148f..ad92b6578d71b 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/convert/ToCartesianShapeTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/convert/ToCartesianShapeTests.java @@ -13,6 +13,7 @@ import org.apache.lucene.util.BytesRef; import org.elasticsearch.geo.GeometryTestUtils; import org.elasticsearch.xpack.esql.expression.function.AbstractFunctionTestCase; +import org.elasticsearch.xpack.esql.expression.function.FunctionName; import org.elasticsearch.xpack.esql.expression.function.TestCaseSupplier; import org.elasticsearch.xpack.esql.type.EsqlDataTypes; import org.elasticsearch.xpack.ql.expression.Expression; @@ -26,6 +27,7 @@ import static org.elasticsearch.xpack.ql.util.SpatialCoordinateTypes.CARTESIAN; +@FunctionName("to_cartesianshape") public class ToCartesianShapeTests extends AbstractFunctionTestCase { public ToCartesianShapeTests(@Name("TestCase") Supplier testCaseSupplier) { this.testCase = testCaseSupplier.get(); diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/convert/ToGeoPointTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/convert/ToGeoPointTests.java index 4a5534e1d5d1a..342325a63d96e 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/convert/ToGeoPointTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/convert/ToGeoPointTests.java @@ -13,6 +13,7 @@ import org.apache.lucene.util.BytesRef; import org.elasticsearch.geo.GeometryTestUtils; import org.elasticsearch.xpack.esql.expression.function.AbstractFunctionTestCase; +import org.elasticsearch.xpack.esql.expression.function.FunctionName; import org.elasticsearch.xpack.esql.expression.function.TestCaseSupplier; import org.elasticsearch.xpack.esql.type.EsqlDataTypes; import org.elasticsearch.xpack.ql.expression.Expression; @@ -26,6 +27,7 @@ import static org.elasticsearch.xpack.ql.util.SpatialCoordinateTypes.GEO; +@FunctionName("to_geopoint") public class ToGeoPointTests extends AbstractFunctionTestCase { public ToGeoPointTests(@Name("TestCase") Supplier testCaseSupplier) { this.testCase = testCaseSupplier.get(); diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/convert/ToGeoShapeTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/convert/ToGeoShapeTests.java index 15db74d71d21f..290d0a08db725 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/convert/ToGeoShapeTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/convert/ToGeoShapeTests.java @@ -13,6 +13,7 @@ import org.apache.lucene.util.BytesRef; import org.elasticsearch.geo.GeometryTestUtils; import org.elasticsearch.xpack.esql.expression.function.AbstractFunctionTestCase; +import org.elasticsearch.xpack.esql.expression.function.FunctionName; import org.elasticsearch.xpack.esql.expression.function.TestCaseSupplier; import org.elasticsearch.xpack.esql.type.EsqlDataTypes; import org.elasticsearch.xpack.ql.expression.Expression; @@ -26,6 +27,7 @@ import static org.elasticsearch.xpack.ql.util.SpatialCoordinateTypes.GEO; +@FunctionName("to_geoshape") public class ToGeoShapeTests extends AbstractFunctionTestCase { public ToGeoShapeTests(@Name("TestCase") Supplier testCaseSupplier) { this.testCase = testCaseSupplier.get(); From aba546d023cbb5ea88082e1fad4d837b43827880 Mon Sep 17 00:00:00 2001 From: Stef Nestor <26751266+stefnestor@users.noreply.github.com> Date: Tue, 5 Mar 2024 10:10:58 -0700 Subject: [PATCH 06/27] (API+) CAT Nodes alias for shard header to match CAT Allocation (#105847) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 👋 howdy, team! Will you kindly consider adding `shards` as an alias to `shard_stats.total_count` column for CAT Nodes to match its naming from CAT Allocation? (The tests returned clean without changes which I hadn't expected, so please let me know if I missed something.) **Example**: To avoid running the default CAT Nodes & CAT Allocation separately, you can run CAT Nodes ``` GET _cat/nodes?v&s=master,name&h=name,id,master,node.role,cpu,heap.percent,disk.*,sstc,uptime ``` Where `sstc` is `shards` from CAT Allocation. This is a 👶 API (+ its doc) change to make the output more intuitive. --- docs/changelog/105847.yaml | 5 +++++ docs/reference/cat/nodes.asciidoc | 2 +- .../org/elasticsearch/rest/action/cat/RestNodesAction.java | 2 +- 3 files changed, 7 insertions(+), 2 deletions(-) create mode 100644 docs/changelog/105847.yaml diff --git a/docs/changelog/105847.yaml b/docs/changelog/105847.yaml new file mode 100644 index 0000000000000..a731395bc9a81 --- /dev/null +++ b/docs/changelog/105847.yaml @@ -0,0 +1,5 @@ +pr: 105847 +summary: (API+) CAT Nodes alias for shard header to match CAT Allocation +area: Stats +type: enhancement +issues: [] diff --git a/docs/reference/cat/nodes.asciidoc b/docs/reference/cat/nodes.asciidoc index b670ee26a20a9..da1ed532e41fa 100644 --- a/docs/reference/cat/nodes.asciidoc +++ b/docs/reference/cat/nodes.asciidoc @@ -318,7 +318,7 @@ Time spent in suggest, such as `0`. `suggest.total`, `suto`, `suggestTotal`:: Number of suggest operations, such as `0`. -`shard_stats.total_count`, `sstc`, `shardStatsTotalCount`:: +`shard_stats.total_count`, `sstc`, `shards`, `shardStatsTotalCount`:: Number of shards assigned. `mappings.total_count`, `mtc`, `mappingsTotalCount`:: diff --git a/server/src/main/java/org/elasticsearch/rest/action/cat/RestNodesAction.java b/server/src/main/java/org/elasticsearch/rest/action/cat/RestNodesAction.java index 39045a99aa4a2..9b70776551ba6 100644 --- a/server/src/main/java/org/elasticsearch/rest/action/cat/RestNodesAction.java +++ b/server/src/main/java/org/elasticsearch/rest/action/cat/RestNodesAction.java @@ -316,7 +316,7 @@ protected Table getTableWithHeader(final RestRequest request) { table.addCell( "shard_stats.total_count", - "alias:sstc,shardStatsTotalCount;default:false;text-align:right;desc:number of shards assigned" + "alias:sstc,shards,shardStatsTotalCount;default:false;text-align:right;desc:number of shards assigned" ); table.addCell("mappings.total_count", "alias:mtc,mappingsTotalCount;default:false;text-align:right;desc:number of mappings"); From 2fbdc33dcf2cb1efbbd893557b4c64003a4087f7 Mon Sep 17 00:00:00 2001 From: Lee Hinman Date: Tue, 5 Mar 2024 18:20:00 +0100 Subject: [PATCH 07/27] Wait forever for IndexTemplateRegistry asset installation (#105985) Previously we would wait one minute for templates, ILM policies, and pipelines to be installed. This commit changes the timeout to use `TimeValue.MAX_VALUE` so that they should continue to wait until either the asset is install, or the master fails over. --- docs/changelog/105985.yaml | 5 +++++ .../xpack/core/template/IndexTemplateRegistry.java | 12 ++++++------ 2 files changed, 11 insertions(+), 6 deletions(-) create mode 100644 docs/changelog/105985.yaml diff --git a/docs/changelog/105985.yaml b/docs/changelog/105985.yaml new file mode 100644 index 0000000000000..2f2a8c1394070 --- /dev/null +++ b/docs/changelog/105985.yaml @@ -0,0 +1,5 @@ +pr: 105985 +summary: Wait forever for `IndexTemplateRegistry` asset installation +area: Indices APIs +type: enhancement +issues: [] diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/template/IndexTemplateRegistry.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/template/IndexTemplateRegistry.java index db5746f5c1b47..e189116b0179c 100644 --- a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/template/IndexTemplateRegistry.java +++ b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/template/IndexTemplateRegistry.java @@ -465,7 +465,7 @@ private void putLegacyTemplate(final IndexTemplateConfig config, final AtomicBoo final String templateName = config.getTemplateName(); PutIndexTemplateRequest request = new PutIndexTemplateRequest(templateName).source(config.loadBytes(), XContentType.JSON); - request.masterNodeTimeout(TimeValue.timeValueMinutes(1)); + request.masterNodeTimeout(TimeValue.MAX_VALUE); executeAsyncWithOrigin( client.threadPool().getThreadContext(), getOrigin(), @@ -498,7 +498,7 @@ private void putComponentTemplate(final String templateName, final ComponentTemp final Executor executor = threadPool.generic(); executor.execute(() -> { PutComponentTemplateAction.Request request = new PutComponentTemplateAction.Request(templateName).componentTemplate(template); - request.masterNodeTimeout(TimeValue.timeValueMinutes(1)); + request.masterNodeTimeout(TimeValue.MAX_VALUE); executeAsyncWithOrigin( client.threadPool().getThreadContext(), getOrigin(), @@ -538,7 +538,7 @@ private void putComposableTemplate( executor.execute(() -> { TransportPutComposableIndexTemplateAction.Request request = new TransportPutComposableIndexTemplateAction.Request(templateName) .indexTemplate(indexTemplate); - request.masterNodeTimeout(TimeValue.timeValueMinutes(1)); + request.masterNodeTimeout(TimeValue.MAX_VALUE); executeAsyncWithOrigin( client.threadPool().getThreadContext(), getOrigin(), @@ -615,7 +615,7 @@ private void putPolicy(final LifecyclePolicy policy, final AtomicBoolean creatio final Executor executor = threadPool.generic(); executor.execute(() -> { PutLifecycleRequest request = new PutLifecycleRequest(policy); - request.masterNodeTimeout(TimeValue.timeValueMinutes(1)); + request.masterNodeTimeout(TimeValue.MAX_VALUE); executeAsyncWithOrigin( client.threadPool().getThreadContext(), getOrigin(), @@ -727,7 +727,7 @@ private void putIngestPipeline(final IngestPipelineConfig pipelineConfig, final pipelineConfig.loadConfig(), pipelineConfig.getXContentType() ); - request.masterNodeTimeout(TimeValue.timeValueMinutes(1)); + request.masterNodeTimeout(TimeValue.MAX_VALUE); executeAsyncWithOrigin( client.threadPool().getThreadContext(), @@ -815,7 +815,7 @@ public void onFailure(Exception e) { ); RolloverRequest request = new RolloverRequest(rolloverTarget, null); request.lazy(true); - request.masterNodeTimeout(TimeValue.timeValueMinutes(1)); + request.masterNodeTimeout(TimeValue.MAX_VALUE); executeAsyncWithOrigin( client.threadPool().getThreadContext(), getOrigin(), From 71ffd35b6bed3fa87b5c645907f45e824defcfc2 Mon Sep 17 00:00:00 2001 From: Nik Everett Date: Tue, 5 Mar 2024 12:33:17 -0500 Subject: [PATCH 08/27] ESQL: fix single valued query tests (#105986) In some cases the tests for our lucene query that makes sure a field is single-valued was asserting incorrect things about the stats that come from the query. That was failing the test from time to time. This fixes the assertion in those cases. Closes #105918 --- .../querydsl/query/SingleValueQueryTests.java | 97 ++++++++++++------- 1 file changed, 62 insertions(+), 35 deletions(-) diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/querydsl/query/SingleValueQueryTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/querydsl/query/SingleValueQueryTests.java index 55e8ba164ba70..6465e73417ae2 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/querydsl/query/SingleValueQueryTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/querydsl/query/SingleValueQueryTests.java @@ -50,7 +50,7 @@ interface Setup { List> build(RandomIndexWriter iw) throws IOException; - void assertStats(SingleValueQuery.Builder builder, boolean subHasTwoPhase); + void assertStats(SingleValueQuery.Builder builder, YesNoSometimes subHasTwoPhase); } @ParametersFactory @@ -74,7 +74,7 @@ public SingleValueQueryTests(Setup setup) { } public void testMatchAll() throws IOException { - testCase(new SingleValueQuery(new MatchAll(Source.EMPTY), "foo").asBuilder(), false, false, this::runCase); + testCase(new SingleValueQuery(new MatchAll(Source.EMPTY), "foo").asBuilder(), YesNoSometimes.NO, YesNoSometimes.NO, this::runCase); } @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/105952") @@ -82,8 +82,8 @@ public void testMatchSome() throws IOException { int max = between(1, 100); testCase( new SingleValueQuery.Builder(new RangeQueryBuilder("i").lt(max), "foo", new SingleValueQuery.Stats(), Source.EMPTY), - false, - false, + YesNoSometimes.SOMETIMES, + YesNoSometimes.NO, (fieldValues, count) -> runCase(fieldValues, count, null, max, false) ); } @@ -96,8 +96,8 @@ public void testSubPhrase() throws IOException { new SingleValueQuery.Stats(), Source.EMPTY ), - false, - true, + YesNoSometimes.NO, + YesNoSometimes.YES, this::runCase ); } @@ -105,8 +105,8 @@ public void testSubPhrase() throws IOException { public void testMatchNone() throws IOException { testCase( new SingleValueQuery.Builder(new MatchNoneQueryBuilder(), "foo", new SingleValueQuery.Stats(), Source.EMPTY), - true, - false, + YesNoSometimes.YES, + YesNoSometimes.NO, (fieldValues, count) -> assertThat(count, equalTo(0)) ); } @@ -114,8 +114,8 @@ public void testMatchNone() throws IOException { public void testRewritesToMatchNone() throws IOException { testCase( new SingleValueQuery.Builder(new TermQueryBuilder("missing", 0), "foo", new SingleValueQuery.Stats(), Source.EMPTY), - true, - false, + YesNoSometimes.YES, + YesNoSometimes.NO, (fieldValues, count) -> assertThat(count, equalTo(0)) ); } @@ -123,8 +123,8 @@ public void testRewritesToMatchNone() throws IOException { public void testNotMatchAll() throws IOException { testCase( new SingleValueQuery(new MatchAll(Source.EMPTY), "foo").negate(Source.EMPTY).asBuilder(), - true, - false, + YesNoSometimes.YES, + YesNoSometimes.NO, (fieldValues, count) -> assertThat(count, equalTo(0)) ); } @@ -132,19 +132,18 @@ public void testNotMatchAll() throws IOException { public void testNotMatchNone() throws IOException { testCase( new SingleValueQuery(new MatchAll(Source.EMPTY).negate(Source.EMPTY), "foo").negate(Source.EMPTY).asBuilder(), - false, - false, + YesNoSometimes.NO, + YesNoSometimes.NO, this::runCase ); } - @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/105918") public void testNotMatchSome() throws IOException { int max = between(1, 100); testCase( new SingleValueQuery(new RangeQuery(Source.EMPTY, "i", null, false, max, false, null), "foo").negate(Source.EMPTY).asBuilder(), - false, - true, + YesNoSometimes.SOMETIMES, + YesNoSometimes.SOMETIMES, (fieldValues, count) -> runCase(fieldValues, count, max, 100, true) ); } @@ -193,8 +192,18 @@ private void runCase(List> fieldValues, int count) { runCase(fieldValues, count, null, null, false); } - private void testCase(SingleValueQuery.Builder builder, boolean rewritesToMatchNone, boolean subHasTwoPhase, TestCase testCase) - throws IOException { + enum YesNoSometimes { + YES, + NO, + SOMETIMES; + } + + private void testCase( + SingleValueQuery.Builder builder, + YesNoSometimes rewritesToMatchNone, + YesNoSometimes subHasTwoPhase, + TestCase testCase + ) throws IOException { MapperService mapper = createMapperService(mapping(setup::mapping)); try (Directory d = newDirectory(); RandomIndexWriter iw = new RandomIndexWriter(random(), d)) { List> fieldValues = setup.build(iw); @@ -203,7 +212,7 @@ private void testCase(SingleValueQuery.Builder builder, boolean rewritesToMatchN QueryBuilder rewritten = builder.rewrite(ctx); Query query = rewritten.toQuery(ctx); testCase.run(fieldValues, ctx.searcher().count(query)); - if (rewritesToMatchNone) { + if (rewritesToMatchNone == YesNoSometimes.YES) { assertThat(rewritten, instanceOf(MatchNoneQueryBuilder.class)); assertThat(builder.stats().missingField(), equalTo(0)); assertThat(builder.stats().rewrittenToMatchNone(), equalTo(1)); @@ -219,7 +228,9 @@ private void testCase(SingleValueQuery.Builder builder, boolean rewritesToMatchN assertThat(builder.stats().rewrittenToMatchNone(), equalTo(0)); setup.assertStats(builder, subHasTwoPhase); } - assertThat(builder.stats().noNextScorer(), equalTo(0)); + if (rewritesToMatchNone != YesNoSometimes.SOMETIMES) { + assertThat(builder.stats().noNextScorer(), equalTo(0)); + } } } } @@ -302,7 +313,7 @@ private List docFor(int i, Iterable values) { } @Override - public void assertStats(SingleValueQuery.Builder builder, boolean subHasTwoPhase) { + public void assertStats(SingleValueQuery.Builder builder, YesNoSometimes subHasTwoPhase) { assertThat(builder.stats().missingField(), equalTo(0)); switch (fieldType) { case "long", "integer", "short", "byte", "double", "float" -> { @@ -314,12 +325,20 @@ public void assertStats(SingleValueQuery.Builder builder, boolean subHasTwoPhase if (multivaluedField || empty) { assertThat(builder.stats().numericSingle(), greaterThanOrEqualTo(0)); - if (subHasTwoPhase) { - assertThat(builder.stats().numericMultiNoApprox(), equalTo(0)); - assertThat(builder.stats().numericMultiApprox(), greaterThan(0)); - } else { - assertThat(builder.stats().numericMultiNoApprox(), greaterThan(0)); - assertThat(builder.stats().numericMultiApprox(), equalTo(0)); + switch (subHasTwoPhase) { + case YES -> { + assertThat(builder.stats().numericMultiNoApprox(), equalTo(0)); + assertThat(builder.stats().numericMultiApprox(), greaterThan(0)); + } + case NO -> { + assertThat(builder.stats().numericMultiNoApprox(), greaterThan(0)); + assertThat(builder.stats().numericMultiApprox(), equalTo(0)); + } + case SOMETIMES -> { + assertThat(builder.stats().numericMultiNoApprox() + builder.stats().numericMultiApprox(), greaterThan(0)); + assertThat(builder.stats().numericMultiNoApprox(), greaterThanOrEqualTo(0)); + assertThat(builder.stats().numericMultiApprox(), greaterThanOrEqualTo(0)); + } } } else { assertThat(builder.stats().numericSingle(), greaterThan(0)); @@ -335,12 +354,20 @@ public void assertStats(SingleValueQuery.Builder builder, boolean subHasTwoPhase assertThat(builder.stats().bytesNoApprox(), equalTo(0)); if (multivaluedField || empty) { assertThat(builder.stats().ordinalsSingle(), greaterThanOrEqualTo(0)); - if (subHasTwoPhase) { - assertThat(builder.stats().ordinalsMultiNoApprox(), equalTo(0)); - assertThat(builder.stats().ordinalsMultiApprox(), greaterThan(0)); - } else { - assertThat(builder.stats().ordinalsMultiNoApprox(), greaterThan(0)); - assertThat(builder.stats().ordinalsMultiApprox(), equalTo(0)); + switch (subHasTwoPhase) { + case YES -> { + assertThat(builder.stats().ordinalsMultiNoApprox(), equalTo(0)); + assertThat(builder.stats().ordinalsMultiApprox(), greaterThan(0)); + } + case NO -> { + assertThat(builder.stats().ordinalsMultiNoApprox(), greaterThan(0)); + assertThat(builder.stats().ordinalsMultiApprox(), equalTo(0)); + } + case SOMETIMES -> { + assertThat(builder.stats().ordinalsMultiNoApprox() + builder.stats().ordinalsMultiApprox(), greaterThan(0)); + assertThat(builder.stats().ordinalsMultiNoApprox(), greaterThanOrEqualTo(0)); + assertThat(builder.stats().ordinalsMultiApprox(), greaterThanOrEqualTo(0)); + } } } else { assertThat(builder.stats().ordinalsSingle(), greaterThan(0)); @@ -373,7 +400,7 @@ public List> build(RandomIndexWriter iw) throws IOException { } @Override - public void assertStats(SingleValueQuery.Builder builder, boolean subHasTwoPhase) { + public void assertStats(SingleValueQuery.Builder builder, YesNoSometimes subHasTwoPhase) { assertThat(builder.stats().missingField(), equalTo(1)); assertThat(builder.stats().numericSingle(), equalTo(0)); assertThat(builder.stats().numericMultiNoApprox(), equalTo(0)); From 680774e7f276ae403e62b9af9c610483ccb25344 Mon Sep 17 00:00:00 2001 From: Jonathan Buttner <56361221+jonathan-buttner@users.noreply.github.com> Date: Tue, 5 Mar 2024 12:42:08 -0500 Subject: [PATCH 09/27] [ML] Refactor inference service retry and queuing logic (#105526) * Refactoring http sender factory to be a nested class * Moving retrier and adding in request manager logic * Working tests * additional renaming and tests * Fixing merge issues * Cleaning up code * Removing interfaces --------- Co-authored-by: Elastic Machine --- .../xpack/inference/InferencePlugin.java | 13 +- .../action/cohere/CohereActionCreator.java | 2 +- .../action/cohere/CohereEmbeddingsAction.java | 46 +-- .../action/huggingface/HuggingFaceAction.java | 40 +-- .../action/openai/OpenAiEmbeddingsAction.java | 31 +- .../external/http/HttpClientManager.java | 1 - .../external/http/RequestExecutor.java | 11 +- .../{Retrier.java => RequestSender.java} | 15 +- .../external/http/retry/RetrySettings.java | 112 +++++- .../http/retry/RetryingHttpSender.java | 111 ++++-- ...ereEmbeddingsExecutableRequestCreator.java | 55 +++ .../sender/ExecutableInferenceRequest.java | 44 +++ .../http/sender/ExecutableRequestCreator.java | 29 ++ .../http/sender/HttpRequestSender.java | 192 ++++++++++ .../http/sender/HttpRequestSenderFactory.java | 175 --------- .../HuggingFaceExecutableRequestCreator.java | 64 ++++ .../http/sender/InferenceRequest.java | 45 +++ .../external/http/sender/NoopTask.java | 36 +- ...nAiEmbeddingsExecutableRequestCreator.java | 67 ++++ .../external/http/sender/RejectableTask.java | 12 + .../http/sender/RequestExecutorService.java | 80 ++--- .../external/http/sender/RequestTask.java | 163 +++------ .../external/http/sender/Sender.java | 13 +- .../http/sender/SingleRequestManager.java | 47 +++ .../external/openai/OpenAiClient.java | 47 --- .../inference/services/SenderService.java | 4 +- .../services/cohere/CohereService.java | 4 +- .../huggingface/HuggingFaceBaseService.java | 4 +- .../huggingface/HuggingFaceService.java | 4 +- .../elser/HuggingFaceElserService.java | 4 +- .../services/openai/OpenAiService.java | 4 +- .../elasticsearch/xpack/inference/Utils.java | 4 +- .../cohere/CohereActionCreatorTests.java | 4 +- .../cohere/CohereEmbeddingsActionTests.java | 23 +- .../HuggingFaceActionCreatorTests.java | 55 ++- .../huggingface/HuggingFaceActionTests.java | 9 +- .../openai/OpenAiActionCreatorTests.java | 178 +++++++++- .../openai/OpenAiEmbeddingsActionTests.java | 24 +- .../http/retry/RetrySettingsTests.java | 4 +- .../http/retry/RetryingHttpSenderTests.java | 335 ++++++++---------- .../sender/ExecutableRequestCreatorTests.java | 49 +++ ...Tests.java => HttpRequestSenderTests.java} | 134 +++++-- ...beddingsExecutableRequestCreatorTests.java | 40 +++ .../sender/RequestExecutorServiceTests.java | 148 ++++---- .../http/sender/RequestTaskTests.java | 208 ++--------- .../sender/SingleRequestManagerTests.java | 27 ++ .../external/openai/OpenAiClientTests.java | 297 ---------------- .../openai/OpenAiEmbeddingsRequestTests.java | 2 +- .../services/SenderServiceTests.java | 8 +- .../services/ServiceComponentsTests.java | 4 + .../services/cohere/CohereServiceTests.java | 21 +- .../HuggingFaceBaseServiceTests.java | 7 +- .../huggingface/HuggingFaceServiceTests.java | 12 +- .../services/openai/OpenAiServiceTests.java | 21 +- .../OpenAiEmbeddingsModelTests.java | 18 + 55 files changed, 1685 insertions(+), 1422 deletions(-) rename x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/retry/{Retrier.java => RequestSender.java} (53%) create mode 100644 x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/CohereEmbeddingsExecutableRequestCreator.java create mode 100644 x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/ExecutableInferenceRequest.java create mode 100644 x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/ExecutableRequestCreator.java create mode 100644 x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/HttpRequestSender.java delete mode 100644 x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/HttpRequestSenderFactory.java create mode 100644 x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/HuggingFaceExecutableRequestCreator.java create mode 100644 x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/InferenceRequest.java create mode 100644 x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/OpenAiEmbeddingsExecutableRequestCreator.java create mode 100644 x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/RejectableTask.java create mode 100644 x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/SingleRequestManager.java delete mode 100644 x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/openai/OpenAiClient.java create mode 100644 x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/http/sender/ExecutableRequestCreatorTests.java rename x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/http/sender/{HttpRequestSenderFactoryTests.java => HttpRequestSenderTests.java} (53%) create mode 100644 x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/http/sender/OpenAiEmbeddingsExecutableRequestCreatorTests.java create mode 100644 x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/http/sender/SingleRequestManagerTests.java delete mode 100644 x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/openai/OpenAiClientTests.java diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferencePlugin.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferencePlugin.java index 1c5e5d4e9ef94..c598a58d014f9 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferencePlugin.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/InferencePlugin.java @@ -50,7 +50,7 @@ import org.elasticsearch.xpack.inference.external.http.HttpClientManager; import org.elasticsearch.xpack.inference.external.http.HttpSettings; import org.elasticsearch.xpack.inference.external.http.retry.RetrySettings; -import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSenderFactory; +import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSender; import org.elasticsearch.xpack.inference.external.http.sender.RequestExecutorServiceSettings; import org.elasticsearch.xpack.inference.logging.ThrottlerManager; import org.elasticsearch.xpack.inference.registry.ModelRegistryImpl; @@ -95,7 +95,7 @@ public class InferencePlugin extends Plugin implements ActionPlugin, ExtensibleP public static final String NAME = "inference"; public static final String UTILITY_THREAD_POOL_NAME = "inference_utility"; private final Settings settings; - private final SetOnce httpFactory = new SetOnce<>(); + private final SetOnce httpFactory = new SetOnce<>(); private final SetOnce serviceComponents = new SetOnce<>(); private final SetOnce inferenceServiceRegistry = new SetOnce<>(); @@ -144,11 +144,10 @@ public Collection createComponents(PluginServices services) { var truncator = new Truncator(settings, services.clusterService()); serviceComponents.set(new ServiceComponents(services.threadPool(), throttlerManager, settings, truncator)); - var httpRequestSenderFactory = new HttpRequestSenderFactory( - services.threadPool(), + var httpRequestSenderFactory = new HttpRequestSender.Factory( + serviceComponents.get(), HttpClientManager.create(settings, services.threadPool(), services.clusterService(), throttlerManager), - services.clusterService(), - settings + services.clusterService() ); httpFactory.set(httpRequestSenderFactory); @@ -241,7 +240,7 @@ public List> getSettings() { return Stream.of( HttpSettings.getSettings(), HttpClientManager.getSettings(), - HttpRequestSenderFactory.HttpRequestSender.getSettings(), + HttpRequestSender.getSettings(), ThrottlerManager.getSettings(), RetrySettings.getSettingsDefinitions(), Truncator.getSettings(), diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/cohere/CohereActionCreator.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/cohere/CohereActionCreator.java index 0fb5ca9283fae..91db5e691cb61 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/cohere/CohereActionCreator.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/cohere/CohereActionCreator.java @@ -32,6 +32,6 @@ public CohereActionCreator(Sender sender, ServiceComponents serviceComponents) { public ExecutableAction create(CohereEmbeddingsModel model, Map taskSettings, InputType inputType) { var overriddenModel = CohereEmbeddingsModel.of(model, taskSettings, inputType); - return new CohereEmbeddingsAction(sender, overriddenModel, serviceComponents); + return new CohereEmbeddingsAction(sender, overriddenModel); } } diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/cohere/CohereEmbeddingsAction.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/cohere/CohereEmbeddingsAction.java index ae66496abbb1f..1f50f0ae6bc57 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/cohere/CohereEmbeddingsAction.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/cohere/CohereEmbeddingsAction.java @@ -7,21 +7,12 @@ package org.elasticsearch.xpack.inference.external.action.cohere; -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; import org.elasticsearch.ElasticsearchException; import org.elasticsearch.action.ActionListener; import org.elasticsearch.inference.InferenceServiceResults; import org.elasticsearch.xpack.inference.external.action.ExecutableAction; -import org.elasticsearch.xpack.inference.external.cohere.CohereAccount; -import org.elasticsearch.xpack.inference.external.cohere.CohereResponseHandler; -import org.elasticsearch.xpack.inference.external.http.retry.ResponseHandler; -import org.elasticsearch.xpack.inference.external.http.retry.RetrySettings; -import org.elasticsearch.xpack.inference.external.http.retry.RetryingHttpSender; +import org.elasticsearch.xpack.inference.external.http.sender.CohereEmbeddingsExecutableRequestCreator; import org.elasticsearch.xpack.inference.external.http.sender.Sender; -import org.elasticsearch.xpack.inference.external.request.cohere.CohereEmbeddingsRequest; -import org.elasticsearch.xpack.inference.external.response.cohere.CohereEmbeddingsResponseEntity; -import org.elasticsearch.xpack.inference.services.ServiceComponents; import org.elasticsearch.xpack.inference.services.cohere.embeddings.CohereEmbeddingsModel; import java.util.List; @@ -32,51 +23,32 @@ import static org.elasticsearch.xpack.inference.external.action.ActionUtils.wrapFailuresInElasticsearchException; public class CohereEmbeddingsAction implements ExecutableAction { - private static final Logger logger = LogManager.getLogger(CohereEmbeddingsAction.class); - private static final ResponseHandler HANDLER = createEmbeddingsHandler(); - - private final CohereAccount account; - private final CohereEmbeddingsModel model; private final String failedToSendRequestErrorMessage; - private final RetryingHttpSender sender; + private final Sender sender; + private final CohereEmbeddingsExecutableRequestCreator requestCreator; - public CohereEmbeddingsAction(Sender sender, CohereEmbeddingsModel model, ServiceComponents serviceComponents) { - this.model = Objects.requireNonNull(model); - this.account = new CohereAccount( - this.model.getServiceSettings().getCommonSettings().getUri(), - this.model.getSecretSettings().apiKey() - ); + public CohereEmbeddingsAction(Sender sender, CohereEmbeddingsModel model) { + Objects.requireNonNull(model); + this.sender = Objects.requireNonNull(sender); this.failedToSendRequestErrorMessage = constructFailedToSendRequestMessage( - this.model.getServiceSettings().getCommonSettings().getUri(), + model.getServiceSettings().getCommonSettings().getUri(), "Cohere embeddings" ); - this.sender = new RetryingHttpSender( - Objects.requireNonNull(sender), - serviceComponents.throttlerManager(), - logger, - new RetrySettings(serviceComponents.settings()), - serviceComponents.threadPool() - ); + requestCreator = new CohereEmbeddingsExecutableRequestCreator(model); } @Override public void execute(List input, ActionListener listener) { try { - CohereEmbeddingsRequest request = new CohereEmbeddingsRequest(account, input, model); ActionListener wrappedListener = wrapFailuresInElasticsearchException( failedToSendRequestErrorMessage, listener ); - - sender.send(request, HANDLER, wrappedListener); + sender.send(requestCreator, input, wrappedListener); } catch (ElasticsearchException e) { listener.onFailure(e); } catch (Exception e) { listener.onFailure(createInternalServerError(e, failedToSendRequestErrorMessage)); } } - - private static ResponseHandler createEmbeddingsHandler() { - return new CohereResponseHandler("cohere text embedding", CohereEmbeddingsResponseEntity::fromResponse); - } } diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/huggingface/HuggingFaceAction.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/huggingface/HuggingFaceAction.java index 67c5fda5f83a0..928d396c991f8 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/huggingface/HuggingFaceAction.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/huggingface/HuggingFaceAction.java @@ -7,19 +7,13 @@ package org.elasticsearch.xpack.inference.external.action.huggingface; -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; import org.elasticsearch.ElasticsearchException; import org.elasticsearch.action.ActionListener; import org.elasticsearch.inference.InferenceServiceResults; -import org.elasticsearch.xpack.inference.common.Truncator; import org.elasticsearch.xpack.inference.external.action.ExecutableAction; import org.elasticsearch.xpack.inference.external.http.retry.ResponseHandler; -import org.elasticsearch.xpack.inference.external.http.retry.RetrySettings; -import org.elasticsearch.xpack.inference.external.http.retry.RetryingHttpSender; +import org.elasticsearch.xpack.inference.external.http.sender.HuggingFaceExecutableRequestCreator; import org.elasticsearch.xpack.inference.external.http.sender.Sender; -import org.elasticsearch.xpack.inference.external.huggingface.HuggingFaceAccount; -import org.elasticsearch.xpack.inference.external.request.huggingface.HuggingFaceInferenceRequest; import org.elasticsearch.xpack.inference.services.ServiceComponents; import org.elasticsearch.xpack.inference.services.huggingface.HuggingFaceModel; @@ -27,19 +21,13 @@ import java.util.Objects; import static org.elasticsearch.core.Strings.format; -import static org.elasticsearch.xpack.inference.common.Truncator.truncate; import static org.elasticsearch.xpack.inference.external.action.ActionUtils.createInternalServerError; import static org.elasticsearch.xpack.inference.external.action.ActionUtils.wrapFailuresInElasticsearchException; public class HuggingFaceAction implements ExecutableAction { - private static final Logger logger = LogManager.getLogger(HuggingFaceAction.class); - - private final HuggingFaceAccount account; private final String errorMessage; - private final RetryingHttpSender sender; - private final ResponseHandler responseHandler; - private final Truncator truncator; - private final HuggingFaceModel model; + private final Sender sender; + private final HuggingFaceExecutableRequestCreator requestCreator; public HuggingFaceAction( Sender sender, @@ -50,34 +38,20 @@ public HuggingFaceAction( ) { Objects.requireNonNull(serviceComponents); Objects.requireNonNull(requestType); - - this.model = Objects.requireNonNull(model); - this.responseHandler = Objects.requireNonNull(responseHandler); - this.sender = new RetryingHttpSender( - Objects.requireNonNull(sender), - serviceComponents.throttlerManager(), - logger, - new RetrySettings(serviceComponents.settings()), - serviceComponents.threadPool() - ); - this.account = new HuggingFaceAccount(model.getUri(), model.getApiKey()); - this.errorMessage = format( + this.sender = Objects.requireNonNull(sender); + requestCreator = new HuggingFaceExecutableRequestCreator(model, responseHandler, serviceComponents.truncator()); + errorMessage = format( "Failed to send Hugging Face %s request from inference entity id [%s]", requestType, model.getInferenceEntityId() ); - this.truncator = Objects.requireNonNull(serviceComponents.truncator()); } @Override public void execute(List input, ActionListener listener) { try { - var truncatedInput = truncate(input, model.getTokenLimit()); - - HuggingFaceInferenceRequest request = new HuggingFaceInferenceRequest(truncator, account, truncatedInput, model); ActionListener wrappedListener = wrapFailuresInElasticsearchException(errorMessage, listener); - - sender.send(request, responseHandler, wrappedListener); + sender.send(requestCreator, input, wrappedListener); } catch (ElasticsearchException e) { listener.onFailure(e); } catch (Exception e) { diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/openai/OpenAiEmbeddingsAction.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/openai/OpenAiEmbeddingsAction.java index 2e804dfeb6a4f..d5f083ac8aa90 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/openai/OpenAiEmbeddingsAction.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/openai/OpenAiEmbeddingsAction.java @@ -10,52 +10,39 @@ import org.elasticsearch.ElasticsearchException; import org.elasticsearch.action.ActionListener; import org.elasticsearch.inference.InferenceServiceResults; -import org.elasticsearch.xpack.inference.common.Truncator; import org.elasticsearch.xpack.inference.external.action.ExecutableAction; +import org.elasticsearch.xpack.inference.external.http.sender.OpenAiEmbeddingsExecutableRequestCreator; import org.elasticsearch.xpack.inference.external.http.sender.Sender; -import org.elasticsearch.xpack.inference.external.openai.OpenAiAccount; -import org.elasticsearch.xpack.inference.external.openai.OpenAiClient; -import org.elasticsearch.xpack.inference.external.request.openai.OpenAiEmbeddingsRequest; import org.elasticsearch.xpack.inference.services.ServiceComponents; import org.elasticsearch.xpack.inference.services.openai.embeddings.OpenAiEmbeddingsModel; import java.util.List; import java.util.Objects; -import static org.elasticsearch.xpack.inference.common.Truncator.truncate; import static org.elasticsearch.xpack.inference.external.action.ActionUtils.constructFailedToSendRequestMessage; import static org.elasticsearch.xpack.inference.external.action.ActionUtils.createInternalServerError; import static org.elasticsearch.xpack.inference.external.action.ActionUtils.wrapFailuresInElasticsearchException; public class OpenAiEmbeddingsAction implements ExecutableAction { - private final OpenAiAccount account; - private final OpenAiClient client; - private final OpenAiEmbeddingsModel model; private final String errorMessage; - private final Truncator truncator; + private final OpenAiEmbeddingsExecutableRequestCreator requestCreator; + private final Sender sender; public OpenAiEmbeddingsAction(Sender sender, OpenAiEmbeddingsModel model, ServiceComponents serviceComponents) { - this.model = Objects.requireNonNull(model); - this.account = new OpenAiAccount( - this.model.getServiceSettings().uri(), - this.model.getServiceSettings().organizationId(), - this.model.getSecretSettings().apiKey() - ); - this.client = new OpenAiClient(Objects.requireNonNull(sender), Objects.requireNonNull(serviceComponents)); - this.errorMessage = constructFailedToSendRequestMessage(this.model.getServiceSettings().uri(), "OpenAI embeddings"); - this.truncator = Objects.requireNonNull(serviceComponents.truncator()); + Objects.requireNonNull(serviceComponents); + Objects.requireNonNull(model); + this.sender = Objects.requireNonNull(sender); + requestCreator = new OpenAiEmbeddingsExecutableRequestCreator(model, serviceComponents.truncator()); + errorMessage = constructFailedToSendRequestMessage(model.getServiceSettings().uri(), "OpenAI embeddings"); } @Override public void execute(List input, ActionListener listener) { try { - var truncatedInput = truncate(input, model.getServiceSettings().maxInputTokens()); - - OpenAiEmbeddingsRequest request = new OpenAiEmbeddingsRequest(truncator, account, truncatedInput, model); ActionListener wrappedListener = wrapFailuresInElasticsearchException(errorMessage, listener); - client.send(request, wrappedListener); + sender.send(requestCreator, input, wrappedListener); } catch (ElasticsearchException e) { listener.onFailure(e); } catch (Exception e) { diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/HttpClientManager.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/HttpClientManager.java index 7cc4a3cb24502..ab3a8a8c0e043 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/HttpClientManager.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/HttpClientManager.java @@ -37,7 +37,6 @@ public class HttpClientManager implements Closeable { */ public static final Setting MAX_CONNECTIONS = Setting.intSetting( "xpack.inference.http.max_connections", - // TODO pick a reasonable values here 20, // default 1, // min Setting.Property.NodeScope, diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/RequestExecutor.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/RequestExecutor.java index 5c8fa62ba88f9..77b4d49d62b9f 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/RequestExecutor.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/RequestExecutor.java @@ -10,8 +10,10 @@ import org.elasticsearch.action.ActionListener; import org.elasticsearch.core.Nullable; import org.elasticsearch.core.TimeValue; -import org.elasticsearch.xpack.inference.external.request.HttpRequest; +import org.elasticsearch.inference.InferenceServiceResults; +import org.elasticsearch.xpack.inference.external.http.sender.ExecutableRequestCreator; +import java.util.List; import java.util.concurrent.TimeUnit; public interface RequestExecutor { @@ -25,5 +27,10 @@ public interface RequestExecutor { boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedException; - void execute(HttpRequest request, @Nullable TimeValue timeout, ActionListener listener); + void execute( + ExecutableRequestCreator requestCreator, + List input, + @Nullable TimeValue timeout, + ActionListener listener + ); } diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/retry/Retrier.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/retry/RequestSender.java similarity index 53% rename from x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/retry/Retrier.java rename to x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/retry/RequestSender.java index 2e2ba03345a3b..8244e5ad29e95 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/retry/Retrier.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/retry/RequestSender.java @@ -7,10 +7,21 @@ package org.elasticsearch.xpack.inference.external.http.retry; +import org.apache.http.client.protocol.HttpClientContext; +import org.apache.logging.log4j.Logger; import org.elasticsearch.action.ActionListener; import org.elasticsearch.inference.InferenceServiceResults; import org.elasticsearch.xpack.inference.external.request.Request; -public interface Retrier { - void send(Request request, ResponseHandler responseHandler, ActionListener listener); +import java.util.function.Supplier; + +public interface RequestSender { + void send( + Logger logger, + Request request, + HttpClientContext context, + Supplier hasRequestTimedOutFunction, + ResponseHandler responseHandler, + ActionListener listener + ); } diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/retry/RetrySettings.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/retry/RetrySettings.java index 040903a35ab08..35e50e557cc83 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/retry/RetrySettings.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/retry/RetrySettings.java @@ -7,6 +7,7 @@ package org.elasticsearch.xpack.inference.external.http.retry; +import org.elasticsearch.cluster.service.ClusterService; import org.elasticsearch.common.settings.Setting; import org.elasticsearch.common.settings.Settings; import org.elasticsearch.core.TimeValue; @@ -15,43 +16,128 @@ public class RetrySettings { - public static final Setting RETRY_INITIAL_DELAY_SETTING = Setting.timeSetting( + static final Setting RETRY_INITIAL_DELAY_SETTING = Setting.timeSetting( "xpack.inference.http.retry.initial_delay", TimeValue.timeValueSeconds(1), Setting.Property.NodeScope, Setting.Property.Dynamic ); - public static final Setting RETRY_MAX_DELAY_BOUND_SETTING = Setting.timeSetting( + static final Setting RETRY_MAX_DELAY_BOUND_SETTING = Setting.timeSetting( "xpack.inference.http.retry.max_delay_bound", TimeValue.timeValueSeconds(5), Setting.Property.NodeScope, Setting.Property.Dynamic ); - public static final Setting RETRY_TIMEOUT_SETTING = Setting.timeSetting( + static final Setting RETRY_TIMEOUT_SETTING = Setting.timeSetting( "xpack.inference.http.retry.timeout", TimeValue.timeValueSeconds(30), Setting.Property.NodeScope, Setting.Property.Dynamic ); - private final InternalSettings internalSettings; + static final Setting RETRY_DEBUG_FREQUENCY_MODE_SETTING = Setting.enumSetting( + DebugFrequencyMode.class, + "xpack.inference.http.retry.debug_frequency_mode", + DebugFrequencyMode.OFF, + Setting.Property.NodeScope, + Setting.Property.Dynamic + ); + + static final Setting RETRY_DEBUG_FREQUENCY_AMOUNT_SETTING = Setting.timeSetting( + "xpack.inference.http.retry.debug_frequency_amount", + TimeValue.timeValueMinutes(5), + Setting.Property.NodeScope, + Setting.Property.Dynamic + ); + + private volatile TimeValue initialDelay; + private volatile TimeValue maxDelayBound; + private volatile TimeValue timeout; + private volatile DebugFrequencyMode debugMode; + private volatile TimeValue debugFrequency; + + public RetrySettings(Settings settings, ClusterService clusterService) { + initialDelay = RETRY_INITIAL_DELAY_SETTING.get(settings); + maxDelayBound = RETRY_MAX_DELAY_BOUND_SETTING.get(settings); + timeout = RETRY_TIMEOUT_SETTING.get(settings); + debugMode = RETRY_DEBUG_FREQUENCY_MODE_SETTING.get(settings); + debugFrequency = RETRY_DEBUG_FREQUENCY_AMOUNT_SETTING.get(settings); + + addSettingsUpdateConsumers(clusterService); + } + + private void addSettingsUpdateConsumers(ClusterService clusterService) { + clusterService.getClusterSettings().addSettingsUpdateConsumer(RETRY_INITIAL_DELAY_SETTING, this::setInitialDelay); + clusterService.getClusterSettings().addSettingsUpdateConsumer(RETRY_MAX_DELAY_BOUND_SETTING, this::setMaxDelayBound); + clusterService.getClusterSettings().addSettingsUpdateConsumer(RETRY_TIMEOUT_SETTING, this::setTimeout); + clusterService.getClusterSettings().addSettingsUpdateConsumer(RETRY_DEBUG_FREQUENCY_MODE_SETTING, this::setDebugMode); + clusterService.getClusterSettings().addSettingsUpdateConsumer(RETRY_DEBUG_FREQUENCY_AMOUNT_SETTING, this::setDebugFrequencyAmount); + } + + private void setInitialDelay(TimeValue initialDelay) { + this.initialDelay = initialDelay; + } - public RetrySettings(Settings settings) { - var initialDelay = RETRY_INITIAL_DELAY_SETTING.get(settings); - var maxDelayBound = RETRY_MAX_DELAY_BOUND_SETTING.get(settings); - var timeoutValue = RETRY_TIMEOUT_SETTING.get(settings); - this.internalSettings = new InternalSettings(initialDelay, maxDelayBound, timeoutValue); + private void setMaxDelayBound(TimeValue maxDelayBound) { + this.maxDelayBound = maxDelayBound; } - public record InternalSettings(TimeValue initialDelay, TimeValue maxDelayBound, TimeValue timeoutValue) {} + private void setTimeout(TimeValue timeout) { + this.timeout = timeout; + } + + private void setDebugMode(DebugFrequencyMode debugMode) { + this.debugMode = debugMode; + } - public InternalSettings getSettings() { - return internalSettings; + private void setDebugFrequencyAmount(TimeValue debugFrequency) { + this.debugFrequency = debugFrequency; } public static List> getSettingsDefinitions() { - return List.of(RETRY_INITIAL_DELAY_SETTING, RETRY_MAX_DELAY_BOUND_SETTING, RETRY_TIMEOUT_SETTING); + return List.of( + RETRY_INITIAL_DELAY_SETTING, + RETRY_MAX_DELAY_BOUND_SETTING, + RETRY_TIMEOUT_SETTING, + RETRY_DEBUG_FREQUENCY_MODE_SETTING, + RETRY_DEBUG_FREQUENCY_AMOUNT_SETTING + ); + } + + TimeValue getInitialDelay() { + return initialDelay; + } + + TimeValue getMaxDelayBound() { + return maxDelayBound; + } + + TimeValue getTimeout() { + return timeout; + } + + DebugFrequencyMode getDebugMode() { + return debugMode; + } + + TimeValue getDebugFrequency() { + return debugFrequency; + } + + enum DebugFrequencyMode { + /** + * Indicates that the debug messages should be logged every time + */ + ON, + /** + * Indicates that the debug messages should never be logged + */ + OFF, + /** + * Indicates that the debug messages should be logged on an interval + */ + INTERVAL } } diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/retry/RetryingHttpSender.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/retry/RetryingHttpSender.java index d8476c7c583d5..ffe10ffe3b6ae 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/retry/RetryingHttpSender.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/retry/RetryingHttpSender.java @@ -7,7 +7,9 @@ package org.elasticsearch.xpack.inference.external.http.retry; +import org.apache.http.client.protocol.HttpClientContext; import org.apache.logging.log4j.Logger; +import org.elasticsearch.ElasticsearchException; import org.elasticsearch.ElasticsearchStatusException; import org.elasticsearch.ExceptionsHelper; import org.elasticsearch.action.ActionListener; @@ -15,8 +17,8 @@ import org.elasticsearch.inference.InferenceServiceResults; import org.elasticsearch.rest.RestStatus; import org.elasticsearch.threadpool.ThreadPool; +import org.elasticsearch.xpack.inference.external.http.HttpClient; import org.elasticsearch.xpack.inference.external.http.HttpResult; -import org.elasticsearch.xpack.inference.external.http.sender.Sender; import org.elasticsearch.xpack.inference.external.request.Request; import org.elasticsearch.xpack.inference.logging.ThrottlerManager; @@ -24,40 +26,37 @@ import java.net.UnknownHostException; import java.util.Objects; import java.util.concurrent.Executor; +import java.util.function.Supplier; import static org.elasticsearch.core.Strings.format; import static org.elasticsearch.xpack.inference.InferencePlugin.UTILITY_THREAD_POOL_NAME; -public class RetryingHttpSender implements Retrier { - private final Sender sender; +public class RetryingHttpSender implements RequestSender { + private final HttpClient httpClient; private final ThrottlerManager throttlerManager; - private final Logger logger; private final RetrySettings retrySettings; private final ThreadPool threadPool; private final Executor executor; public RetryingHttpSender( - Sender sender, + HttpClient httpClient, ThrottlerManager throttlerManager, - Logger logger, RetrySettings retrySettings, ThreadPool threadPool ) { - this(sender, throttlerManager, logger, retrySettings, threadPool, threadPool.executor(UTILITY_THREAD_POOL_NAME)); + this(httpClient, throttlerManager, retrySettings, threadPool, threadPool.executor(UTILITY_THREAD_POOL_NAME)); } // For testing only RetryingHttpSender( - Sender sender, + HttpClient httpClient, ThrottlerManager throttlerManager, - Logger logger, RetrySettings retrySettings, ThreadPool threadPool, Executor executor ) { - this.sender = Objects.requireNonNull(sender); + this.httpClient = Objects.requireNonNull(httpClient); this.throttlerManager = Objects.requireNonNull(throttlerManager); - this.logger = Objects.requireNonNull(logger); this.retrySettings = Objects.requireNonNull(retrySettings); this.threadPool = Objects.requireNonNull(threadPool); this.executor = Objects.requireNonNull(executor); @@ -66,23 +65,41 @@ public RetryingHttpSender( private class InternalRetrier extends RetryableAction { private Request request; private final ResponseHandler responseHandler; - - InternalRetrier(Request request, ResponseHandler responseHandler, ActionListener listener) { + private final Logger logger; + private final HttpClientContext context; + private final Supplier hasRequestCompletedFunction; + + InternalRetrier( + Logger logger, + Request request, + HttpClientContext context, + Supplier hasRequestCompletedFunction, + ResponseHandler responseHandler, + ActionListener listener + ) { super( - logger, + Objects.requireNonNull(logger), threadPool, - retrySettings.getSettings().initialDelay(), - retrySettings.getSettings().maxDelayBound(), - retrySettings.getSettings().timeoutValue(), + retrySettings.getInitialDelay(), + retrySettings.getMaxDelayBound(), + retrySettings.getTimeout(), listener, executor ); - this.request = request; - this.responseHandler = responseHandler; + this.logger = logger; + this.request = Objects.requireNonNull(request); + this.context = Objects.requireNonNull(context); + this.responseHandler = Objects.requireNonNull(responseHandler); + this.hasRequestCompletedFunction = Objects.requireNonNull(hasRequestCompletedFunction); } @Override public void tryAction(ActionListener listener) { + // A timeout likely occurred so let's stop attempting to execute the request + if (hasRequestCompletedFunction.get()) { + return; + } + ActionListener responseListener = ActionListener.wrap(result -> { try { responseHandler.validateResponse(throttlerManager, logger, request, result); @@ -90,25 +107,21 @@ public void tryAction(ActionListener listener) { listener.onResponse(inferenceResults); } catch (Exception e) { - logException(request, result, responseHandler.getRequestType(), e); + logException(logger, request, result, responseHandler.getRequestType(), e); listener.onFailure(e); } }, e -> { - logException(request, responseHandler.getRequestType(), e); + logException(logger, request, responseHandler.getRequestType(), e); listener.onFailure(transformIfRetryable(e)); }); - sender.send(request.createHttpRequest(), responseListener); - } + try { + httpClient.send(request.createHttpRequest(), context, responseListener); + } catch (Exception e) { + logException(logger, request, responseHandler.getRequestType(), e); - @Override - public boolean shouldRetry(Exception e) { - if (e instanceof Retryable retry) { - request = retry.rebuildRequest(request); - return retry.shouldRetry(); + listener.onFailure(wrapWithElasticsearchException(e, request.getInferenceEntityId())); } - - return false; } /** @@ -135,15 +148,45 @@ private Exception transformIfRetryable(Exception e) { return exceptionToReturn; } + + private Exception wrapWithElasticsearchException(Exception e, String inferenceEntityId) { + var transformedException = transformIfRetryable(e); + + if (transformedException instanceof ElasticsearchException) { + return transformedException; + } + + return new ElasticsearchException( + format("Http client failed to send request from inference entity id [%s]", inferenceEntityId), + transformedException + ); + } + + @Override + public boolean shouldRetry(Exception e) { + if (e instanceof Retryable retry) { + request = retry.rebuildRequest(request); + return retry.shouldRetry(); + } + + return false; + } } @Override - public void send(Request request, ResponseHandler responseHandler, ActionListener listener) { - InternalRetrier retrier = new InternalRetrier(request, responseHandler, listener); + public void send( + Logger logger, + Request request, + HttpClientContext context, + Supplier hasRequestTimedOutFunction, + ResponseHandler responseHandler, + ActionListener listener + ) { + InternalRetrier retrier = new InternalRetrier(logger, request, context, hasRequestTimedOutFunction, responseHandler, listener); retrier.run(); } - private void logException(Request request, String requestType, Exception exception) { + private void logException(Logger logger, Request request, String requestType, Exception exception) { var causeException = ExceptionsHelper.unwrapCause(exception); throttlerManager.warn( @@ -153,7 +196,7 @@ private void logException(Request request, String requestType, Exception excepti ); } - private void logException(Request request, HttpResult result, String requestType, Exception exception) { + private void logException(Logger logger, Request request, HttpResult result, String requestType, Exception exception) { var causeException = ExceptionsHelper.unwrapCause(exception); throttlerManager.warn( diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/CohereEmbeddingsExecutableRequestCreator.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/CohereEmbeddingsExecutableRequestCreator.java new file mode 100644 index 0000000000000..b0fdc800a64da --- /dev/null +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/CohereEmbeddingsExecutableRequestCreator.java @@ -0,0 +1,55 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.inference.external.http.sender; + +import org.apache.http.client.protocol.HttpClientContext; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.elasticsearch.action.ActionListener; +import org.elasticsearch.inference.InferenceServiceResults; +import org.elasticsearch.xpack.inference.external.cohere.CohereAccount; +import org.elasticsearch.xpack.inference.external.cohere.CohereResponseHandler; +import org.elasticsearch.xpack.inference.external.http.retry.RequestSender; +import org.elasticsearch.xpack.inference.external.http.retry.ResponseHandler; +import org.elasticsearch.xpack.inference.external.request.cohere.CohereEmbeddingsRequest; +import org.elasticsearch.xpack.inference.external.response.cohere.CohereEmbeddingsResponseEntity; +import org.elasticsearch.xpack.inference.services.cohere.embeddings.CohereEmbeddingsModel; + +import java.util.List; +import java.util.Objects; +import java.util.function.Supplier; + +public class CohereEmbeddingsExecutableRequestCreator implements ExecutableRequestCreator { + private static final Logger logger = LogManager.getLogger(CohereEmbeddingsExecutableRequestCreator.class); + private static final ResponseHandler HANDLER = createEmbeddingsHandler(); + + private static ResponseHandler createEmbeddingsHandler() { + return new CohereResponseHandler("cohere text embedding", CohereEmbeddingsResponseEntity::fromResponse); + } + + private final CohereAccount account; + private final CohereEmbeddingsModel model; + + public CohereEmbeddingsExecutableRequestCreator(CohereEmbeddingsModel model) { + this.model = Objects.requireNonNull(model); + account = new CohereAccount(this.model.getServiceSettings().getCommonSettings().getUri(), this.model.getSecretSettings().apiKey()); + } + + @Override + public Runnable create( + List input, + RequestSender requestSender, + Supplier hasRequestCompletedFunction, + HttpClientContext context, + ActionListener listener + ) { + CohereEmbeddingsRequest request = new CohereEmbeddingsRequest(account, input, model); + + return new ExecutableInferenceRequest(requestSender, logger, request, context, HANDLER, hasRequestCompletedFunction, listener); + } +} diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/ExecutableInferenceRequest.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/ExecutableInferenceRequest.java new file mode 100644 index 0000000000000..53f30773cbfe3 --- /dev/null +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/ExecutableInferenceRequest.java @@ -0,0 +1,44 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.inference.external.http.sender; + +import org.apache.http.client.protocol.HttpClientContext; +import org.apache.logging.log4j.Logger; +import org.elasticsearch.ElasticsearchException; +import org.elasticsearch.action.ActionListener; +import org.elasticsearch.common.Strings; +import org.elasticsearch.inference.InferenceServiceResults; +import org.elasticsearch.xpack.inference.external.http.retry.RequestSender; +import org.elasticsearch.xpack.inference.external.http.retry.ResponseHandler; +import org.elasticsearch.xpack.inference.external.request.Request; + +import java.util.function.Supplier; + +record ExecutableInferenceRequest( + RequestSender requestSender, + Logger logger, + Request request, + HttpClientContext context, + ResponseHandler responseHandler, + Supplier hasFinished, + ActionListener listener +) implements Runnable { + + @Override + public void run() { + var inferenceEntityId = request.createHttpRequest().inferenceEntityId(); + + try { + requestSender.send(logger, request, context, hasFinished, responseHandler, listener); + } catch (Exception e) { + var errorMessage = Strings.format("Failed to send request from inference entity id [%s]", inferenceEntityId); + logger.warn(errorMessage, e); + listener.onFailure(new ElasticsearchException(errorMessage, e)); + } + } +} diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/ExecutableRequestCreator.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/ExecutableRequestCreator.java new file mode 100644 index 0000000000000..96455ca4b1cb1 --- /dev/null +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/ExecutableRequestCreator.java @@ -0,0 +1,29 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.inference.external.http.sender; + +import org.apache.http.client.protocol.HttpClientContext; +import org.elasticsearch.action.ActionListener; +import org.elasticsearch.inference.InferenceServiceResults; +import org.elasticsearch.xpack.inference.external.http.retry.RequestSender; + +import java.util.List; +import java.util.function.Supplier; + +/** + * A contract for constructing a {@link Runnable} to handle sending an inference request to a 3rd party service. + */ +public interface ExecutableRequestCreator { + Runnable create( + List input, + RequestSender requestSender, + Supplier hasRequestCompletedFunction, + HttpClientContext context, + ActionListener listener + ); +} diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/HttpRequestSender.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/HttpRequestSender.java new file mode 100644 index 0000000000000..0131bf2989f6f --- /dev/null +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/HttpRequestSender.java @@ -0,0 +1,192 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.inference.external.http.sender; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.elasticsearch.action.ActionListener; +import org.elasticsearch.cluster.service.ClusterService; +import org.elasticsearch.common.settings.Setting; +import org.elasticsearch.common.settings.Settings; +import org.elasticsearch.core.Nullable; +import org.elasticsearch.core.TimeValue; +import org.elasticsearch.inference.InferenceServiceResults; +import org.elasticsearch.threadpool.ThreadPool; +import org.elasticsearch.xpack.inference.external.http.HttpClientManager; +import org.elasticsearch.xpack.inference.external.http.retry.RetrySettings; +import org.elasticsearch.xpack.inference.external.http.retry.RetryingHttpSender; +import org.elasticsearch.xpack.inference.services.ServiceComponents; + +import java.io.IOException; +import java.util.List; +import java.util.Objects; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; + +import static org.elasticsearch.core.Strings.format; +import static org.elasticsearch.xpack.inference.InferencePlugin.UTILITY_THREAD_POOL_NAME; + +/** + * A class for providing a more friendly interface for sending an inference request to a 3rd party service. + */ +public class HttpRequestSender implements Sender { + + /** + * A helper class for constructing a {@link HttpRequestSender}. + */ + public static class Factory { + private final ServiceComponents serviceComponents; + private final HttpClientManager httpClientManager; + private final ClusterService clusterService; + private final SingleRequestManager requestManager; + + public Factory(ServiceComponents serviceComponents, HttpClientManager httpClientManager, ClusterService clusterService) { + this.serviceComponents = Objects.requireNonNull(serviceComponents); + this.httpClientManager = Objects.requireNonNull(httpClientManager); + this.clusterService = Objects.requireNonNull(clusterService); + + var requestSender = new RetryingHttpSender( + this.httpClientManager.getHttpClient(), + serviceComponents.throttlerManager(), + new RetrySettings(serviceComponents.settings(), clusterService), + serviceComponents.threadPool() + ); + requestManager = new SingleRequestManager(requestSender); + } + + public Sender createSender(String serviceName) { + return new HttpRequestSender( + serviceName, + serviceComponents.threadPool(), + httpClientManager, + clusterService, + serviceComponents.settings(), + requestManager + ); + } + } + + private static final Logger logger = LogManager.getLogger(HttpRequestSender.class); + private static final TimeValue START_COMPLETED_WAIT_TIME = TimeValue.timeValueSeconds(5); + + /** + * The maximum time a request can take. The timer starts once a request is enqueued and continues until a response is + * received from the 3rd party service. This encompasses the time the request might just sit in the queue waiting to be sent + * if another request is already waiting for a connection lease from the connection pool. + */ + public static final Setting MAX_REQUEST_TIMEOUT = Setting.timeSetting( + "xpack.inference.http.max_request_timeout", + TimeValue.timeValueSeconds(30), + Setting.Property.NodeScope, + Setting.Property.Dynamic + ); + + private final ThreadPool threadPool; + private final HttpClientManager manager; + private final RequestExecutorService service; + private final AtomicBoolean started = new AtomicBoolean(false); + private volatile TimeValue maxRequestTimeout; + private final CountDownLatch startCompleted = new CountDownLatch(2); + + private HttpRequestSender( + String serviceName, + ThreadPool threadPool, + HttpClientManager httpClientManager, + ClusterService clusterService, + Settings settings, + SingleRequestManager requestManager + ) { + this.threadPool = Objects.requireNonNull(threadPool); + this.manager = Objects.requireNonNull(httpClientManager); + service = new RequestExecutorService( + serviceName, + threadPool, + startCompleted, + new RequestExecutorServiceSettings(settings, clusterService), + requestManager + ); + + this.maxRequestTimeout = MAX_REQUEST_TIMEOUT.get(settings); + addSettingsUpdateConsumers(clusterService); + } + + private void addSettingsUpdateConsumers(ClusterService clusterService) { + clusterService.getClusterSettings().addSettingsUpdateConsumer(MAX_REQUEST_TIMEOUT, this::setMaxRequestTimeout); + } + + // Default for testing + void setMaxRequestTimeout(TimeValue maxRequestTimeout) { + logger.debug(() -> format("Max request timeout updated to [%s] for service [%s]", maxRequestTimeout, service)); + this.maxRequestTimeout = maxRequestTimeout; + } + + /** + * Start various internal services. This is required before sending requests. + */ + public void start() { + if (started.compareAndSet(false, true)) { + // The manager must be started before the executor service. That way we guarantee that the http client + // is ready prior to the service attempting to use the http client to send a request + manager.start(); + threadPool.executor(UTILITY_THREAD_POOL_NAME).execute(service::start); + startCompleted.countDown(); + } + } + + @Override + public void close() throws IOException { + manager.close(); + service.shutdown(); + } + + /** + * Send a request at some point in the future. The timeout used is retrieved from the settings. + * @param requestCreator a factory for creating a request to be sent to a 3rd party service + * @param input the list of string input to send in the request + * @param timeout the maximum time the request should wait for a response before timing out. If null, the timeout is ignored. + * The queuing logic may still throw a timeout if it fails to send the request because it couldn't get a leased + * @param listener a listener to handle the response + */ + public void send( + ExecutableRequestCreator requestCreator, + List input, + @Nullable TimeValue timeout, + ActionListener listener + ) { + assert started.get() : "call start() before sending a request"; + waitForStartToComplete(); + service.execute(requestCreator, input, timeout, listener); + } + + private void waitForStartToComplete() { + try { + if (startCompleted.await(START_COMPLETED_WAIT_TIME.getSeconds(), TimeUnit.SECONDS) == false) { + throw new IllegalStateException("Http sender startup did not complete in time"); + } + } catch (InterruptedException e) { + throw new IllegalStateException("Http sender interrupted while waiting for startup to complete"); + } + } + + /** + * Send a request at some point in the future. The timeout used is retrieved from the settings. + * @param requestCreator a factory for creating a request to be sent to a 3rd party service + * @param input the list of string input to send in the request + * @param listener a listener to handle the response + */ + public void send(ExecutableRequestCreator requestCreator, List input, ActionListener listener) { + assert started.get() : "call start() before sending a request"; + waitForStartToComplete(); + service.execute(requestCreator, input, maxRequestTimeout, listener); + } + + public static List> getSettings() { + return List.of(MAX_REQUEST_TIMEOUT); + } +} diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/HttpRequestSenderFactory.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/HttpRequestSenderFactory.java deleted file mode 100644 index c773f57933415..0000000000000 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/HttpRequestSenderFactory.java +++ /dev/null @@ -1,175 +0,0 @@ -/* - * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one - * or more contributor license agreements. Licensed under the Elastic License - * 2.0; you may not use this file except in compliance with the Elastic License - * 2.0. - */ - -package org.elasticsearch.xpack.inference.external.http.sender; - -import org.apache.http.client.methods.HttpUriRequest; -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; -import org.elasticsearch.action.ActionListener; -import org.elasticsearch.cluster.service.ClusterService; -import org.elasticsearch.common.settings.Setting; -import org.elasticsearch.common.settings.Settings; -import org.elasticsearch.core.Nullable; -import org.elasticsearch.core.TimeValue; -import org.elasticsearch.threadpool.ThreadPool; -import org.elasticsearch.xpack.inference.external.http.HttpClientManager; -import org.elasticsearch.xpack.inference.external.http.HttpResult; -import org.elasticsearch.xpack.inference.external.request.HttpRequest; - -import java.io.IOException; -import java.util.List; -import java.util.Objects; -import java.util.concurrent.CountDownLatch; -import java.util.concurrent.TimeUnit; -import java.util.concurrent.atomic.AtomicBoolean; - -import static org.elasticsearch.core.Strings.format; -import static org.elasticsearch.xpack.inference.InferencePlugin.UTILITY_THREAD_POOL_NAME; - -/** - * A helper class for constructing a {@link HttpRequestSender}. - */ -public class HttpRequestSenderFactory { - private final ThreadPool threadPool; - private final HttpClientManager httpClientManager; - private final ClusterService clusterService; - private final Settings settings; - - public HttpRequestSenderFactory( - ThreadPool threadPool, - HttpClientManager httpClientManager, - ClusterService clusterService, - Settings settings - ) { - this.threadPool = Objects.requireNonNull(threadPool); - this.httpClientManager = Objects.requireNonNull(httpClientManager); - this.clusterService = Objects.requireNonNull(clusterService); - this.settings = Objects.requireNonNull(settings); - } - - public Sender createSender(String serviceName) { - return new HttpRequestSender(serviceName, threadPool, httpClientManager, clusterService, settings); - } - - /** - * A class for providing a more friendly interface for sending an {@link HttpUriRequest}. This leverages the queuing logic for sending - * a request. - */ - public static final class HttpRequestSender implements Sender { - private static final Logger logger = LogManager.getLogger(HttpRequestSender.class); - private static final TimeValue START_COMPLETED_WAIT_TIME = TimeValue.timeValueSeconds(5); - - /** - * The maximum time a request can take. The timer starts once a request is enqueued and continues until a response is - * received from the 3rd party service. This encompasses the time the request might just sit in the queue waiting to be sent - * if another request is already waiting for a connection lease from the connection pool. - */ - public static final Setting MAX_REQUEST_TIMEOUT = Setting.timeSetting( - "xpack.inference.http.max_request_timeout", - TimeValue.timeValueSeconds(30), - Setting.Property.NodeScope, - Setting.Property.Dynamic - ); - - private final ThreadPool threadPool; - private final HttpClientManager manager; - private final RequestExecutorService service; - private final AtomicBoolean started = new AtomicBoolean(false); - private volatile TimeValue maxRequestTimeout; - private final CountDownLatch startCompleted = new CountDownLatch(2); - - private HttpRequestSender( - String serviceName, - ThreadPool threadPool, - HttpClientManager httpClientManager, - ClusterService clusterService, - Settings settings - ) { - this.threadPool = Objects.requireNonNull(threadPool); - this.manager = Objects.requireNonNull(httpClientManager); - service = new RequestExecutorService( - serviceName, - manager.getHttpClient(), - threadPool, - startCompleted, - new RequestExecutorServiceSettings(settings, clusterService) - ); - - this.maxRequestTimeout = MAX_REQUEST_TIMEOUT.get(settings); - addSettingsUpdateConsumers(clusterService); - } - - private void addSettingsUpdateConsumers(ClusterService clusterService) { - clusterService.getClusterSettings().addSettingsUpdateConsumer(MAX_REQUEST_TIMEOUT, this::setMaxRequestTimeout); - } - - // Default for testing - void setMaxRequestTimeout(TimeValue maxRequestTimeout) { - logger.debug(() -> format("Max request timeout updated to [%s] for service [%s]", maxRequestTimeout, service)); - this.maxRequestTimeout = maxRequestTimeout; - } - - /** - * Start various internal services. This is required before sending requests. - */ - public void start() { - if (started.compareAndSet(false, true)) { - // The manager must be started before the executor service. That way we guarantee that the http client - // is ready prior to the service attempting to use the http client to send a request - manager.start(); - threadPool.executor(UTILITY_THREAD_POOL_NAME).execute(service::start); - startCompleted.countDown(); - } - } - - @Override - public void close() throws IOException { - manager.close(); - service.shutdown(); - } - - /** - * Send a request at some point in the future with a timeout specified. - * @param request the http request to send - * @param timeout the maximum time the request should wait for a response before timing out. If null, the timeout is ignored. - * The queuing logic may still throw a timeout if it fails to send the request because it couldn't get a leased - * connection from the connection pool - * @param listener a listener to handle the response - */ - public void send(HttpRequest request, @Nullable TimeValue timeout, ActionListener listener) { - assert started.get() : "call start() before sending a request"; - waitForStartToComplete(); - service.execute(request, timeout, listener); - } - - private void waitForStartToComplete() { - try { - if (startCompleted.await(START_COMPLETED_WAIT_TIME.getSeconds(), TimeUnit.SECONDS) == false) { - throw new IllegalStateException("Http sender startup did not complete in time"); - } - } catch (InterruptedException e) { - throw new IllegalStateException("Http sender interrupted while waiting for startup to complete"); - } - } - - /** - * Send a request at some point in the future. The timeout used is retrieved from the settings. - * @param request the http request to send - * @param listener a listener to handle the response - */ - public void send(HttpRequest request, ActionListener listener) { - assert started.get() : "call start() before sending a request"; - waitForStartToComplete(); - service.execute(request, maxRequestTimeout, listener); - } - - public static List> getSettings() { - return List.of(MAX_REQUEST_TIMEOUT); - } - } -} diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/HuggingFaceExecutableRequestCreator.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/HuggingFaceExecutableRequestCreator.java new file mode 100644 index 0000000000000..62558fe6071ac --- /dev/null +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/HuggingFaceExecutableRequestCreator.java @@ -0,0 +1,64 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.inference.external.http.sender; + +import org.apache.http.client.protocol.HttpClientContext; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.elasticsearch.action.ActionListener; +import org.elasticsearch.inference.InferenceServiceResults; +import org.elasticsearch.xpack.inference.common.Truncator; +import org.elasticsearch.xpack.inference.external.http.retry.RequestSender; +import org.elasticsearch.xpack.inference.external.http.retry.ResponseHandler; +import org.elasticsearch.xpack.inference.external.huggingface.HuggingFaceAccount; +import org.elasticsearch.xpack.inference.external.request.huggingface.HuggingFaceInferenceRequest; +import org.elasticsearch.xpack.inference.services.huggingface.HuggingFaceModel; + +import java.util.List; +import java.util.Objects; +import java.util.function.Supplier; + +import static org.elasticsearch.xpack.inference.common.Truncator.truncate; + +public class HuggingFaceExecutableRequestCreator implements ExecutableRequestCreator { + private static final Logger logger = LogManager.getLogger(HuggingFaceExecutableRequestCreator.class); + + private final HuggingFaceModel model; + private final HuggingFaceAccount account; + private final ResponseHandler responseHandler; + private final Truncator truncator; + + public HuggingFaceExecutableRequestCreator(HuggingFaceModel model, ResponseHandler responseHandler, Truncator truncator) { + this.model = Objects.requireNonNull(model); + account = new HuggingFaceAccount(model.getUri(), model.getApiKey()); + this.responseHandler = Objects.requireNonNull(responseHandler); + this.truncator = Objects.requireNonNull(truncator); + } + + @Override + public Runnable create( + List input, + RequestSender requestSender, + Supplier hasRequestCompletedFunction, + HttpClientContext context, + ActionListener listener + ) { + var truncatedInput = truncate(input, model.getTokenLimit()); + var request = new HuggingFaceInferenceRequest(truncator, account, truncatedInput, model); + + return new ExecutableInferenceRequest( + requestSender, + logger, + request, + context, + responseHandler, + hasRequestCompletedFunction, + listener + ); + } +} diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/InferenceRequest.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/InferenceRequest.java new file mode 100644 index 0000000000000..ed77e4b207a94 --- /dev/null +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/InferenceRequest.java @@ -0,0 +1,45 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.inference.external.http.sender; + +import org.elasticsearch.action.ActionListener; +import org.elasticsearch.inference.InferenceServiceResults; + +import java.util.List; +import java.util.function.Supplier; + +/** + * A contract for defining a request sent to a 3rd party service. + */ +public interface InferenceRequest { + + /** + * Returns the creator that handles building an executable request based on the input provided. + */ + ExecutableRequestCreator getRequestCreator(); + + /** + * Returns the text input associated with this request. + */ + List getInput(); + + /** + * Returns the listener to notify of the results. + */ + ActionListener getListener(); + + /** + * Returns whether the request has completed. Returns true if from a failure, success, or a timeout. + */ + boolean hasCompleted(); + + /** + * Returns a {@link Supplier} to determine if the request has completed. + */ + Supplier getRequestCompletedFunction(); +} diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/NoopTask.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/NoopTask.java index c5e533eb7d8fe..6cdcd38d224a9 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/NoopTask.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/NoopTask.java @@ -7,13 +7,41 @@ package org.elasticsearch.xpack.inference.external.http.sender; -import org.elasticsearch.common.util.concurrent.AbstractRunnable; +import org.elasticsearch.action.ActionListener; +import org.elasticsearch.inference.InferenceServiceResults; -class NoopTask extends AbstractRunnable { +import java.util.List; +import java.util.function.Supplier; + +class NoopTask implements RejectableTask { + + @Override + public ExecutableRequestCreator getRequestCreator() { + return null; + } + + @Override + public List getInput() { + return null; + } @Override - public void onFailure(Exception e) {} + public ActionListener getListener() { + return null; + } @Override - protected void doRun() throws Exception {} + public boolean hasCompleted() { + return true; + } + + @Override + public Supplier getRequestCompletedFunction() { + return () -> true; + } + + @Override + public void onRejection(Exception e) { + + } } diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/OpenAiEmbeddingsExecutableRequestCreator.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/OpenAiEmbeddingsExecutableRequestCreator.java new file mode 100644 index 0000000000000..708e67944441c --- /dev/null +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/OpenAiEmbeddingsExecutableRequestCreator.java @@ -0,0 +1,67 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.inference.external.http.sender; + +import org.apache.http.client.protocol.HttpClientContext; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.elasticsearch.action.ActionListener; +import org.elasticsearch.inference.InferenceServiceResults; +import org.elasticsearch.xpack.inference.common.Truncator; +import org.elasticsearch.xpack.inference.external.http.retry.RequestSender; +import org.elasticsearch.xpack.inference.external.http.retry.ResponseHandler; +import org.elasticsearch.xpack.inference.external.openai.OpenAiAccount; +import org.elasticsearch.xpack.inference.external.openai.OpenAiResponseHandler; +import org.elasticsearch.xpack.inference.external.request.openai.OpenAiEmbeddingsRequest; +import org.elasticsearch.xpack.inference.external.response.openai.OpenAiEmbeddingsResponseEntity; +import org.elasticsearch.xpack.inference.services.openai.embeddings.OpenAiEmbeddingsModel; + +import java.util.List; +import java.util.Objects; +import java.util.function.Supplier; + +import static org.elasticsearch.xpack.inference.common.Truncator.truncate; + +public class OpenAiEmbeddingsExecutableRequestCreator implements ExecutableRequestCreator { + + private static final Logger logger = LogManager.getLogger(OpenAiEmbeddingsExecutableRequestCreator.class); + + private static final ResponseHandler HANDLER = createEmbeddingsHandler(); + + private static ResponseHandler createEmbeddingsHandler() { + return new OpenAiResponseHandler("openai text embedding", OpenAiEmbeddingsResponseEntity::fromResponse); + } + + private final Truncator truncator; + private final OpenAiEmbeddingsModel model; + private final OpenAiAccount account; + + public OpenAiEmbeddingsExecutableRequestCreator(OpenAiEmbeddingsModel model, Truncator truncator) { + this.model = Objects.requireNonNull(model); + this.account = new OpenAiAccount( + this.model.getServiceSettings().uri(), + this.model.getServiceSettings().organizationId(), + this.model.getSecretSettings().apiKey() + ); + this.truncator = Objects.requireNonNull(truncator); + } + + @Override + public Runnable create( + List input, + RequestSender requestSender, + Supplier hasRequestCompletedFunction, + HttpClientContext context, + ActionListener listener + ) { + var truncatedInput = truncate(input, model.getServiceSettings().maxInputTokens()); + OpenAiEmbeddingsRequest request = new OpenAiEmbeddingsRequest(truncator, account, truncatedInput, model); + + return new ExecutableInferenceRequest(requestSender, logger, request, context, HANDLER, hasRequestCompletedFunction, listener); + } +} diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/RejectableTask.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/RejectableTask.java new file mode 100644 index 0000000000000..3da64d5491a60 --- /dev/null +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/RejectableTask.java @@ -0,0 +1,12 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.inference.external.http.sender; + +interface RejectableTask extends InferenceRequest { + void onRejection(Exception e); +} diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/RequestExecutorService.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/RequestExecutorService.java index f844787455290..00c28e8afc069 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/RequestExecutorService.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/RequestExecutorService.java @@ -11,17 +11,14 @@ import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; import org.elasticsearch.action.ActionListener; -import org.elasticsearch.common.util.concurrent.AbstractRunnable; import org.elasticsearch.common.util.concurrent.EsRejectedExecutionException; import org.elasticsearch.core.Nullable; import org.elasticsearch.core.Strings; import org.elasticsearch.core.TimeValue; +import org.elasticsearch.inference.InferenceServiceResults; import org.elasticsearch.threadpool.ThreadPool; import org.elasticsearch.xpack.inference.common.AdjustableCapacityBlockingQueue; -import org.elasticsearch.xpack.inference.external.http.HttpClient; -import org.elasticsearch.xpack.inference.external.http.HttpResult; import org.elasticsearch.xpack.inference.external.http.RequestExecutor; -import org.elasticsearch.xpack.inference.external.request.HttpRequest; import java.util.ArrayList; import java.util.List; @@ -36,8 +33,7 @@ import static org.elasticsearch.core.Strings.format; /** - * An {@link java.util.concurrent.ExecutorService} for queuing and executing {@link RequestTask} containing - * {@link org.apache.http.client.methods.HttpUriRequest}. This class is useful because the + * A service for queuing and executing {@link RequestTask}. This class is useful because the * {@link org.apache.http.impl.nio.conn.PoolingNHttpClientConnectionManager} will block when leasing a connection if no * connections are available. To avoid blocking the inference transport threads, this executor will queue up the * requests until connections are available. @@ -48,12 +44,11 @@ * {@link org.apache.http.client.config.RequestConfig.Builder#setConnectionRequestTimeout} for more info. */ class RequestExecutorService implements RequestExecutor { - - private static final AdjustableCapacityBlockingQueue.QueueCreator QUEUE_CREATOR = + private static final AdjustableCapacityBlockingQueue.QueueCreator QUEUE_CREATOR = new AdjustableCapacityBlockingQueue.QueueCreator<>() { @Override - public BlockingQueue create(int capacity) { - BlockingQueue queue; + public BlockingQueue create(int capacity) { + BlockingQueue queue; if (capacity <= 0) { queue = create(); } else { @@ -64,41 +59,30 @@ public BlockingQueue create(int capacity) { } @Override - public BlockingQueue create() { + public BlockingQueue create() { return new LinkedBlockingQueue<>(); } }; private static final Logger logger = LogManager.getLogger(RequestExecutorService.class); private final String serviceName; - private final AdjustableCapacityBlockingQueue queue; + private final AdjustableCapacityBlockingQueue queue; private final AtomicBoolean running = new AtomicBoolean(true); private final CountDownLatch terminationLatch = new CountDownLatch(1); private final HttpClientContext httpContext; - private final HttpClient httpClient; private final ThreadPool threadPool; private final CountDownLatch startupLatch; private final BlockingQueue controlQueue = new LinkedBlockingQueue<>(); + private final SingleRequestManager requestManager; RequestExecutorService( String serviceName, - HttpClient httpClient, ThreadPool threadPool, @Nullable CountDownLatch startupLatch, - RequestExecutorServiceSettings settings + RequestExecutorServiceSettings settings, + SingleRequestManager requestManager ) { - this(serviceName, httpClient, threadPool, QUEUE_CREATOR, startupLatch, settings); - } - - private static BlockingQueue buildQueue(int capacity) { - BlockingQueue queue; - if (capacity <= 0) { - queue = new LinkedBlockingQueue<>(); - } else { - queue = new LinkedBlockingQueue<>(capacity); - } - - return queue; + this(serviceName, threadPool, QUEUE_CREATOR, startupLatch, settings, requestManager); } /** @@ -106,18 +90,18 @@ private static BlockingQueue buildQueue(int capacity) { */ RequestExecutorService( String serviceName, - HttpClient httpClient, ThreadPool threadPool, - AdjustableCapacityBlockingQueue.QueueCreator createQueue, + AdjustableCapacityBlockingQueue.QueueCreator createQueue, @Nullable CountDownLatch startupLatch, - RequestExecutorServiceSettings settings + RequestExecutorServiceSettings settings, + SingleRequestManager requestManager ) { this.serviceName = Objects.requireNonNull(serviceName); - this.httpClient = Objects.requireNonNull(httpClient); this.threadPool = Objects.requireNonNull(threadPool); this.httpContext = HttpClientContext.create(); this.queue = new AdjustableCapacityBlockingQueue<>(createQueue, settings.getQueueCapacity()); this.startupLatch = startupLatch; + this.requestManager = Objects.requireNonNull(requestManager); Objects.requireNonNull(settings); settings.registerQueueCapacityCallback(this::onCapacityChange); @@ -179,7 +163,7 @@ private void signalStartInitiated() { */ private void handleTasks() throws InterruptedException { try { - AbstractRunnable task = queue.take(); + RejectableTask task = queue.take(); var command = controlQueue.poll(); if (command != null) { @@ -200,9 +184,9 @@ private void handleTasks() throws InterruptedException { } } - private void executeTask(AbstractRunnable task) { + private void executeTask(RejectableTask task) { try { - task.run(); + requestManager.execute(task, httpContext); } catch (Exception e) { logger.warn(format("Http executor service [%s] failed to execute request [%s]", serviceName, task), e); } @@ -212,7 +196,7 @@ private synchronized void notifyRequestsOfShutdown() { assert isShutdown() : "Requests should only be notified if the executor is shutting down"; try { - List notExecuted = new ArrayList<>(); + List notExecuted = new ArrayList<>(); queue.drainTo(notExecuted); rejectTasks(notExecuted, this::rejectTaskBecauseOfShutdown); @@ -221,7 +205,7 @@ private synchronized void notifyRequestsOfShutdown() { } } - private void rejectTaskBecauseOfShutdown(AbstractRunnable task) { + private void rejectTaskBecauseOfShutdown(RejectableTask task) { try { task.onRejection( new EsRejectedExecutionException( @@ -236,7 +220,7 @@ private void rejectTaskBecauseOfShutdown(AbstractRunnable task) { } } - private void rejectTasks(List tasks, Consumer rejectionFunction) { + private void rejectTasks(List tasks, Consumer rejectionFunction) { for (var task : tasks) { rejectionFunction.accept(task); } @@ -270,16 +254,22 @@ public boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedE } /** - * Send the request at some point in the future. + * Execute the request at some point in the future. * - * @param request the http request to send - * @param timeout the maximum time to wait for this request to complete (failing or succeeding). Once the time elapses, the - * listener::onFailure is called with a {@link org.elasticsearch.ElasticsearchTimeoutException}. - * If null, then the request will wait forever - * @param listener an {@link ActionListener} for the response or failure + * @param requestCreator the http request to send + * @param input the text to perform inference on + * @param timeout the maximum time to wait for this request to complete (failing or succeeding). Once the time elapses, the + * listener::onFailure is called with a {@link org.elasticsearch.ElasticsearchTimeoutException}. + * If null, then the request will wait forever + * @param listener an {@link ActionListener} for the response or failure */ - public void execute(HttpRequest request, @Nullable TimeValue timeout, ActionListener listener) { - RequestTask task = new RequestTask(request, httpClient, httpContext, timeout, threadPool, listener); + public void execute( + ExecutableRequestCreator requestCreator, + List input, + @Nullable TimeValue timeout, + ActionListener listener + ) { + var task = new RequestTask(requestCreator, input, timeout, threadPool, listener); if (isShutdown()) { EsRejectedExecutionException rejected = new EsRejectedExecutionException( diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/RequestTask.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/RequestTask.java index cc65d16af652c..970366f7765dd 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/RequestTask.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/RequestTask.java @@ -7,157 +7,96 @@ package org.elasticsearch.xpack.inference.external.http.sender; -import org.apache.http.client.protocol.HttpClientContext; -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; -import org.elasticsearch.ElasticsearchException; import org.elasticsearch.ElasticsearchTimeoutException; import org.elasticsearch.action.ActionListener; -import org.elasticsearch.common.util.concurrent.AbstractRunnable; +import org.elasticsearch.action.support.ListenerTimeouts; +import org.elasticsearch.common.Strings; import org.elasticsearch.core.Nullable; import org.elasticsearch.core.TimeValue; -import org.elasticsearch.threadpool.Scheduler; +import org.elasticsearch.inference.InferenceServiceResults; import org.elasticsearch.threadpool.ThreadPool; -import org.elasticsearch.xpack.inference.external.http.HttpClient; -import org.elasticsearch.xpack.inference.external.http.HttpResult; -import org.elasticsearch.xpack.inference.external.request.HttpRequest; +import java.util.List; import java.util.Objects; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.function.Supplier; -import static org.elasticsearch.core.Strings.format; import static org.elasticsearch.xpack.inference.InferencePlugin.UTILITY_THREAD_POOL_NAME; -class RequestTask extends AbstractRunnable { - private static final Logger logger = LogManager.getLogger(RequestTask.class); - private static final Scheduler.Cancellable NOOP_TIMEOUT_HANDLER = createDefaultHandler(); +class RequestTask implements RejectableTask { - private final HttpRequest request; - private final ActionListener listener; - private final Scheduler.Cancellable timeoutHandler; - private final AtomicBoolean notified = new AtomicBoolean(); - private final TimeValue timeout; - private final Runnable command; + private final AtomicBoolean finished = new AtomicBoolean(); + private final ExecutableRequestCreator requestCreator; + private final List input; + private final ActionListener listener; RequestTask( - HttpRequest request, - HttpClient httpClient, - HttpClientContext context, + ExecutableRequestCreator requestCreator, + List input, @Nullable TimeValue timeout, ThreadPool threadPool, - ActionListener listener + ActionListener listener ) { - this.request = Objects.requireNonNull(request); - this.listener = Objects.requireNonNull(listener); - this.timeout = timeout; - this.timeoutHandler = startTimer(threadPool, timeout); - this.command = threadPool.getThreadContext() - .preserveContext( - new Command( - Objects.requireNonNull(httpClient), - this.request, - Objects.requireNonNull(context), - ActionListener.wrap(this::onSuccess, this::onFailure) - ) - ); + this.requestCreator = Objects.requireNonNull(requestCreator); + this.input = Objects.requireNonNull(input); + this.listener = getListener(Objects.requireNonNull(listener), timeout, Objects.requireNonNull(threadPool)); } - private Scheduler.Cancellable startTimer(ThreadPool threadPool, TimeValue timeout) { - Objects.requireNonNull(threadPool); + private ActionListener getListener( + ActionListener origListener, + @Nullable TimeValue timeout, + ThreadPool threadPool + ) { + ActionListener notificationListener = ActionListener.wrap(result -> { + finished.set(true); + origListener.onResponse(result); + }, e -> { + finished.set(true); + origListener.onFailure(e); + }); if (timeout == null) { - return NOOP_TIMEOUT_HANDLER; + return notificationListener; } - return threadPool.schedule(this::onTimeout, timeout, threadPool.executor(UTILITY_THREAD_POOL_NAME)); - } - - private void onTimeout() { - assert timeout != null : "timeout must be defined to use a timeout handler"; - logger.debug( - () -> format( - "Request from inference entity id [%s] timed out after [%s] while waiting to be executed", - request.inferenceEntityId(), - timeout + return ListenerTimeouts.wrapWithTimeout( + threadPool, + timeout, + threadPool.executor(UTILITY_THREAD_POOL_NAME), + notificationListener, + (ignored) -> notificationListener.onFailure( + new ElasticsearchTimeoutException(Strings.format("Request timed out waiting to be sent after [%s]", timeout)) ) ); - notifyOfResult( - () -> listener.onFailure( - new ElasticsearchTimeoutException(format("Request timed out waiting to be executed after [%s]", timeout)) - ) - ); - } - - private void notifyOfResult(Runnable runnable) { - if (notified.compareAndSet(false, true)) { - runnable.run(); - } } @Override - public void onFailure(Exception e) { - timeoutHandler.cancel(); - notifyOfResult(() -> listener.onFailure(e)); + public boolean hasCompleted() { + return finished.get(); } @Override - protected void doRun() { - try { - command.run(); - } catch (Exception e) { - String message = format("Failed while executing request from inference entity id [%s]", request.inferenceEntityId()); - logger.warn(message, e); - onFailure(new ElasticsearchException(message, e)); - } + public Supplier getRequestCompletedFunction() { + return this::hasCompleted; } - private void onSuccess(HttpResult result) { - timeoutHandler.cancel(); - notifyOfResult(() -> listener.onResponse(result)); + @Override + public List getInput() { + return input; } @Override - public String toString() { - return request.inferenceEntityId(); + public ActionListener getListener() { + return listener; } - private static Scheduler.Cancellable createDefaultHandler() { - return new Scheduler.Cancellable() { - @Override - public boolean cancel() { - return true; - } - - @Override - public boolean isCancelled() { - return true; - } - }; + @Override + public void onRejection(Exception e) { + listener.onFailure(e); } - private record Command( - HttpClient httpClient, - HttpRequest requestToSend, - HttpClientContext context, - ActionListener resultListener - ) implements Runnable { - - @Override - public void run() { - try { - httpClient.send(requestToSend, context, resultListener); - } catch (Exception e) { - logger.warn( - format("Failed to send request from inference entity id [%s] via the http client", requestToSend.inferenceEntityId()), - e - ); - resultListener.onFailure( - new ElasticsearchException( - format("Failed to send request from inference entity id [%s]", requestToSend.inferenceEntityId()), - e - ) - ); - } - } + @Override + public ExecutableRequestCreator getRequestCreator() { + return requestCreator; } } diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/Sender.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/Sender.java index f1a0e112219fd..0272f4b0e351c 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/Sender.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/Sender.java @@ -10,15 +10,20 @@ import org.elasticsearch.action.ActionListener; import org.elasticsearch.core.Nullable; import org.elasticsearch.core.TimeValue; -import org.elasticsearch.xpack.inference.external.http.HttpResult; -import org.elasticsearch.xpack.inference.external.request.HttpRequest; +import org.elasticsearch.inference.InferenceServiceResults; import java.io.Closeable; +import java.util.List; public interface Sender extends Closeable { void start(); - void send(HttpRequest request, ActionListener listener); + void send( + ExecutableRequestCreator requestCreator, + List input, + @Nullable TimeValue timeout, + ActionListener listener + ); - void send(HttpRequest request, @Nullable TimeValue timeout, ActionListener listener); + void send(ExecutableRequestCreator requestCreator, List input, ActionListener listener); } diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/SingleRequestManager.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/SingleRequestManager.java new file mode 100644 index 0000000000000..ecd12814d0877 --- /dev/null +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/SingleRequestManager.java @@ -0,0 +1,47 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.inference.external.http.sender; + +import org.apache.http.client.protocol.HttpClientContext; +import org.elasticsearch.xpack.inference.external.http.retry.RetryingHttpSender; + +import java.util.Objects; + +/** + * Handles executing a single inference request at a time. + */ +public class SingleRequestManager { + + protected RetryingHttpSender requestSender; + + public SingleRequestManager(RetryingHttpSender requestSender) { + this.requestSender = Objects.requireNonNull(requestSender); + } + + public void execute(InferenceRequest inferenceRequest, HttpClientContext context) { + if (isNoopRequest(inferenceRequest) || inferenceRequest.hasCompleted()) { + return; + } + + inferenceRequest.getRequestCreator() + .create( + inferenceRequest.getInput(), + requestSender, + inferenceRequest.getRequestCompletedFunction(), + context, + inferenceRequest.getListener() + ) + .run(); + } + + private static boolean isNoopRequest(InferenceRequest inferenceRequest) { + return inferenceRequest.getRequestCreator() == null + || inferenceRequest.getInput() == null + || inferenceRequest.getListener() == null; + } +} diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/openai/OpenAiClient.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/openai/OpenAiClient.java deleted file mode 100644 index cb82616587091..0000000000000 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/openai/OpenAiClient.java +++ /dev/null @@ -1,47 +0,0 @@ -/* - * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one - * or more contributor license agreements. Licensed under the Elastic License - * 2.0; you may not use this file except in compliance with the Elastic License - * 2.0. - */ - -package org.elasticsearch.xpack.inference.external.openai; - -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; -import org.elasticsearch.action.ActionListener; -import org.elasticsearch.inference.InferenceServiceResults; -import org.elasticsearch.xpack.inference.external.http.retry.ResponseHandler; -import org.elasticsearch.xpack.inference.external.http.retry.RetrySettings; -import org.elasticsearch.xpack.inference.external.http.retry.RetryingHttpSender; -import org.elasticsearch.xpack.inference.external.http.sender.Sender; -import org.elasticsearch.xpack.inference.external.request.openai.OpenAiEmbeddingsRequest; -import org.elasticsearch.xpack.inference.external.response.openai.OpenAiEmbeddingsResponseEntity; -import org.elasticsearch.xpack.inference.services.ServiceComponents; - -import java.io.IOException; - -public class OpenAiClient { - private static final Logger logger = LogManager.getLogger(OpenAiClient.class); - private static final ResponseHandler EMBEDDINGS_HANDLER = createEmbeddingsHandler(); - - private final RetryingHttpSender sender; - - public OpenAiClient(Sender sender, ServiceComponents serviceComponents) { - this.sender = new RetryingHttpSender( - sender, - serviceComponents.throttlerManager(), - logger, - new RetrySettings(serviceComponents.settings()), - serviceComponents.threadPool() - ); - } - - public void send(OpenAiEmbeddingsRequest request, ActionListener listener) throws IOException { - sender.send(request, EMBEDDINGS_HANDLER, listener); - } - - private static ResponseHandler createEmbeddingsHandler() { - return new OpenAiResponseHandler("openai text embedding", OpenAiEmbeddingsResponseEntity::fromResponse); - } -} diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/SenderService.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/SenderService.java index 96378b109ae2d..98b004cd1aa7f 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/SenderService.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/SenderService.java @@ -15,7 +15,7 @@ import org.elasticsearch.inference.InferenceServiceResults; import org.elasticsearch.inference.InputType; import org.elasticsearch.inference.Model; -import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSenderFactory; +import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSender; import org.elasticsearch.xpack.inference.external.http.sender.Sender; import java.io.IOException; @@ -27,7 +27,7 @@ public abstract class SenderService implements InferenceService { private final Sender sender; private final ServiceComponents serviceComponents; - public SenderService(HttpRequestSenderFactory factory, ServiceComponents serviceComponents) { + public SenderService(HttpRequestSender.Factory factory, ServiceComponents serviceComponents) { Objects.requireNonNull(factory); sender = factory.createSender(name()); this.serviceComponents = Objects.requireNonNull(serviceComponents); diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/cohere/CohereService.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/cohere/CohereService.java index 9502acdaf93e5..172a71bd45434 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/cohere/CohereService.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/cohere/CohereService.java @@ -23,7 +23,7 @@ import org.elasticsearch.inference.TaskType; import org.elasticsearch.rest.RestStatus; import org.elasticsearch.xpack.inference.external.action.cohere.CohereActionCreator; -import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSenderFactory; +import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSender; import org.elasticsearch.xpack.inference.services.SenderService; import org.elasticsearch.xpack.inference.services.ServiceComponents; import org.elasticsearch.xpack.inference.services.ServiceUtils; @@ -42,7 +42,7 @@ public class CohereService extends SenderService { public static final String NAME = "cohere"; - public CohereService(HttpRequestSenderFactory factory, ServiceComponents serviceComponents) { + public CohereService(HttpRequestSender.Factory factory, ServiceComponents serviceComponents) { super(factory, serviceComponents); } diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceBaseService.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceBaseService.java index 9cd8c285b406e..20994f3cc8e1e 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceBaseService.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceBaseService.java @@ -17,7 +17,7 @@ import org.elasticsearch.inference.ModelSecrets; import org.elasticsearch.inference.TaskType; import org.elasticsearch.xpack.inference.external.action.huggingface.HuggingFaceActionCreator; -import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSenderFactory; +import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSender; import org.elasticsearch.xpack.inference.services.SenderService; import org.elasticsearch.xpack.inference.services.ServiceComponents; @@ -32,7 +32,7 @@ public abstract class HuggingFaceBaseService extends SenderService { - public HuggingFaceBaseService(HttpRequestSenderFactory factory, ServiceComponents serviceComponents) { + public HuggingFaceBaseService(HttpRequestSender.Factory factory, ServiceComponents serviceComponents) { super(factory, serviceComponents); } diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceService.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceService.java index 2d2f4667478d5..838d3dc857fbc 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceService.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceService.java @@ -15,7 +15,7 @@ import org.elasticsearch.inference.Model; import org.elasticsearch.inference.TaskType; import org.elasticsearch.rest.RestStatus; -import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSenderFactory; +import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSender; import org.elasticsearch.xpack.inference.services.ServiceComponents; import org.elasticsearch.xpack.inference.services.ServiceUtils; import org.elasticsearch.xpack.inference.services.huggingface.elser.HuggingFaceElserModel; @@ -26,7 +26,7 @@ public class HuggingFaceService extends HuggingFaceBaseService { public static final String NAME = "hugging_face"; - public HuggingFaceService(HttpRequestSenderFactory factory, ServiceComponents serviceComponents) { + public HuggingFaceService(HttpRequestSender.Factory factory, ServiceComponents serviceComponents) { super(factory, serviceComponents); } diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/huggingface/elser/HuggingFaceElserService.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/huggingface/elser/HuggingFaceElserService.java index 3cc2ca5ed60a5..2587b2737e164 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/huggingface/elser/HuggingFaceElserService.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/huggingface/elser/HuggingFaceElserService.java @@ -13,7 +13,7 @@ import org.elasticsearch.core.Nullable; import org.elasticsearch.inference.TaskType; import org.elasticsearch.rest.RestStatus; -import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSenderFactory; +import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSender; import org.elasticsearch.xpack.inference.services.ServiceComponents; import org.elasticsearch.xpack.inference.services.huggingface.HuggingFaceBaseService; import org.elasticsearch.xpack.inference.services.huggingface.HuggingFaceModel; @@ -23,7 +23,7 @@ public class HuggingFaceElserService extends HuggingFaceBaseService { public static final String NAME = "hugging_face_elser"; - public HuggingFaceElserService(HttpRequestSenderFactory factory, ServiceComponents serviceComponents) { + public HuggingFaceElserService(HttpRequestSender.Factory factory, ServiceComponents serviceComponents) { super(factory, serviceComponents); } diff --git a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/openai/OpenAiService.java b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/openai/OpenAiService.java index 5062bba8e7eac..234328de67efe 100644 --- a/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/openai/OpenAiService.java +++ b/x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/openai/OpenAiService.java @@ -24,7 +24,7 @@ import org.elasticsearch.inference.TaskType; import org.elasticsearch.rest.RestStatus; import org.elasticsearch.xpack.inference.external.action.openai.OpenAiActionCreator; -import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSenderFactory; +import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSender; import org.elasticsearch.xpack.inference.services.SenderService; import org.elasticsearch.xpack.inference.services.ServiceComponents; import org.elasticsearch.xpack.inference.services.ServiceUtils; @@ -45,7 +45,7 @@ public class OpenAiService extends SenderService { public static final String NAME = "openai"; - public OpenAiService(HttpRequestSenderFactory factory, ServiceComponents serviceComponents) { + public OpenAiService(HttpRequestSender.Factory factory, ServiceComponents serviceComponents) { super(factory, serviceComponents); } diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/Utils.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/Utils.java index 5b7ffb3c8153e..96650bcca565e 100644 --- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/Utils.java +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/Utils.java @@ -16,7 +16,7 @@ import org.elasticsearch.xpack.inference.external.http.HttpClientManager; import org.elasticsearch.xpack.inference.external.http.HttpSettings; import org.elasticsearch.xpack.inference.external.http.retry.RetrySettings; -import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSenderFactory; +import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSender; import org.elasticsearch.xpack.inference.external.http.sender.RequestExecutorServiceSettings; import org.elasticsearch.xpack.inference.logging.ThrottlerManager; @@ -39,7 +39,7 @@ public static ClusterService mockClusterService(Settings settings) { var registeredSettings = Stream.of( HttpSettings.getSettings(), HttpClientManager.getSettings(), - HttpRequestSenderFactory.HttpRequestSender.getSettings(), + HttpRequestSender.getSettings(), ThrottlerManager.getSettings(), RetrySettings.getSettingsDefinitions(), Truncator.getSettings(), diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/action/cohere/CohereActionCreatorTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/action/cohere/CohereActionCreatorTests.java index e7cfc784db117..66ef9910a2649 100644 --- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/action/cohere/CohereActionCreatorTests.java +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/action/cohere/CohereActionCreatorTests.java @@ -19,7 +19,7 @@ import org.elasticsearch.threadpool.ThreadPool; import org.elasticsearch.xcontent.XContentType; import org.elasticsearch.xpack.inference.external.http.HttpClientManager; -import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSenderFactory; +import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSenderTests; import org.elasticsearch.xpack.inference.logging.ThrottlerManager; import org.elasticsearch.xpack.inference.services.cohere.CohereTruncation; import org.elasticsearch.xpack.inference.services.cohere.embeddings.CohereEmbeddingType; @@ -67,7 +67,7 @@ public void shutdown() throws IOException { } public void testCreate_CohereEmbeddingsModel() throws IOException { - var senderFactory = new HttpRequestSenderFactory(threadPool, clientManager, mockClusterServiceEmpty(), Settings.EMPTY); + var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager); try (var sender = senderFactory.createSender("test_service")) { sender.start(); diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/action/cohere/CohereEmbeddingsActionTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/action/cohere/CohereEmbeddingsActionTests.java index 7fd33f7bba58f..b504744bfe5f3 100644 --- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/action/cohere/CohereEmbeddingsActionTests.java +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/action/cohere/CohereEmbeddingsActionTests.java @@ -23,7 +23,7 @@ import org.elasticsearch.xcontent.XContentType; import org.elasticsearch.xpack.inference.external.http.HttpClientManager; import org.elasticsearch.xpack.inference.external.http.HttpResult; -import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSenderFactory; +import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSenderTests; import org.elasticsearch.xpack.inference.external.http.sender.Sender; import org.elasticsearch.xpack.inference.external.request.cohere.CohereUtils; import org.elasticsearch.xpack.inference.logging.ThrottlerManager; @@ -47,7 +47,6 @@ import static org.elasticsearch.xpack.inference.external.http.Utils.entityAsMap; import static org.elasticsearch.xpack.inference.external.http.Utils.getUrl; import static org.elasticsearch.xpack.inference.results.TextEmbeddingResultsTests.buildExpectation; -import static org.elasticsearch.xpack.inference.services.ServiceComponentsTests.createWithEmptySettings; import static org.hamcrest.Matchers.equalTo; import static org.hamcrest.Matchers.hasSize; import static org.hamcrest.Matchers.is; @@ -77,9 +76,9 @@ public void shutdown() throws IOException { } public void testExecute_ReturnsSuccessfulResponse() throws IOException { - var senderFactory = new HttpRequestSenderFactory(threadPool, clientManager, mockClusterServiceEmpty(), Settings.EMPTY); + var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager); - try (var sender = senderFactory.createSender("test_service")) { + try (var sender = HttpRequestSenderTests.createSenderWithSingleRequestManager(senderFactory, "test_service")) { sender.start(); String responseJson = """ @@ -158,9 +157,9 @@ public void testExecute_ReturnsSuccessfulResponse() throws IOException { } public void testExecute_ReturnsSuccessfulResponse_ForInt8ResponseType() throws IOException { - var senderFactory = new HttpRequestSenderFactory(threadPool, clientManager, mockClusterServiceEmpty(), Settings.EMPTY); + var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager); - try (var sender = senderFactory.createSender("test_service")) { + try (var sender = HttpRequestSenderTests.createSenderWithSingleRequestManager(senderFactory, "test_service")) { sender.start(); String responseJson = """ @@ -253,7 +252,7 @@ public void testExecute_ThrowsURISyntaxException_ForInvalidUrl() throws IOExcept public void testExecute_ThrowsElasticsearchException() { var sender = mock(Sender.class); - doThrow(new ElasticsearchException("failed")).when(sender).send(any(), any()); + doThrow(new ElasticsearchException("failed")).when(sender).send(any(), any(), any()); var action = createAction(getUrl(webServer), "secret", CohereEmbeddingsTaskSettings.EMPTY_SETTINGS, null, null, sender); @@ -274,7 +273,7 @@ public void testExecute_ThrowsElasticsearchException_WhenSenderOnFailureIsCalled listener.onFailure(new IllegalStateException("failed")); return Void.TYPE; - }).when(sender).send(any(), any()); + }).when(sender).send(any(), any(), any()); var action = createAction(getUrl(webServer), "secret", CohereEmbeddingsTaskSettings.EMPTY_SETTINGS, null, null, sender); @@ -298,7 +297,7 @@ public void testExecute_ThrowsElasticsearchException_WhenSenderOnFailureIsCalled listener.onFailure(new IllegalStateException("failed")); return Void.TYPE; - }).when(sender).send(any(), any()); + }).when(sender).send(any(), any(), any()); var action = createAction(null, "secret", CohereEmbeddingsTaskSettings.EMPTY_SETTINGS, null, null, sender); @@ -312,7 +311,7 @@ public void testExecute_ThrowsElasticsearchException_WhenSenderOnFailureIsCalled public void testExecute_ThrowsException() { var sender = mock(Sender.class); - doThrow(new IllegalArgumentException("failed")).when(sender).send(any(), any()); + doThrow(new IllegalArgumentException("failed")).when(sender).send(any(), any(), any()); var action = createAction(getUrl(webServer), "secret", CohereEmbeddingsTaskSettings.EMPTY_SETTINGS, null, null, sender); @@ -329,7 +328,7 @@ public void testExecute_ThrowsException() { public void testExecute_ThrowsExceptionWithNullUrl() { var sender = mock(Sender.class); - doThrow(new IllegalArgumentException("failed")).when(sender).send(any(), any()); + doThrow(new IllegalArgumentException("failed")).when(sender).send(any(), any(), any()); var action = createAction(null, "secret", CohereEmbeddingsTaskSettings.EMPTY_SETTINGS, null, null, sender); @@ -351,7 +350,7 @@ private CohereEmbeddingsAction createAction( ) { var model = CohereEmbeddingsModelTests.createModel(url, apiKey, taskSettings, 1024, 1024, modelName, embeddingType); - return new CohereEmbeddingsAction(sender, model, createWithEmptySettings(threadPool)); + return new CohereEmbeddingsAction(sender, model); } } diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/action/huggingface/HuggingFaceActionCreatorTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/action/huggingface/HuggingFaceActionCreatorTests.java index 95b69f1231e9d..6334c669d0c1f 100644 --- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/action/huggingface/HuggingFaceActionCreatorTests.java +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/action/huggingface/HuggingFaceActionCreatorTests.java @@ -20,7 +20,7 @@ import org.elasticsearch.xcontent.XContentType; import org.elasticsearch.xpack.inference.common.TruncatorTests; import org.elasticsearch.xpack.inference.external.http.HttpClientManager; -import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSenderFactory; +import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSenderTests; import org.elasticsearch.xpack.inference.logging.ThrottlerManager; import org.elasticsearch.xpack.inference.results.SparseEmbeddingResultsTests; import org.elasticsearch.xpack.inference.results.TextEmbeddingResultsTests; @@ -31,7 +31,6 @@ import org.junit.Before; import java.io.IOException; -import java.net.URISyntaxException; import java.util.List; import java.util.Map; import java.util.concurrent.TimeUnit; @@ -72,7 +71,7 @@ public void shutdown() throws IOException { @SuppressWarnings("unchecked") public void testExecute_ReturnsSuccessfulResponse_ForElserAction() throws IOException { - var senderFactory = new HttpRequestSenderFactory(threadPool, clientManager, mockClusterServiceEmpty(), Settings.EMPTY); + var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager); try (var sender = senderFactory.createSender("test_service")) { sender.start(); @@ -121,8 +120,14 @@ public void testExecute_ReturnsSuccessfulResponse_ForElserAction() throws IOExce } @SuppressWarnings("unchecked") - public void testSend_FailsFromInvalidResponseFormat_ForElserAction() throws IOException, URISyntaxException { - var senderFactory = new HttpRequestSenderFactory(threadPool, clientManager, mockClusterServiceEmpty(), Settings.EMPTY); + public void testSend_FailsFromInvalidResponseFormat_ForElserAction() throws IOException { + // timeout as zero for no retries + var settings = buildSettingsWithRetryFields( + TimeValue.timeValueMillis(1), + TimeValue.timeValueMinutes(1), + TimeValue.timeValueSeconds(0) + ); + var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager, settings); try (var sender = senderFactory.createSender("test_service")) { sender.start(); @@ -147,17 +152,7 @@ public void testSend_FailsFromInvalidResponseFormat_ForElserAction() throws IOEx var model = HuggingFaceElserModelTests.createModel(getUrl(webServer), "secret"); var actionCreator = new HuggingFaceActionCreator( sender, - new ServiceComponents( - threadPool, - mockThrottlerManager(), - // timeout as zero for no retries - buildSettingsWithRetryFields( - TimeValue.timeValueMillis(1), - TimeValue.timeValueMinutes(1), - TimeValue.timeValueSeconds(0) - ), - TruncatorTests.createTruncator() - ) + new ServiceComponents(threadPool, mockThrottlerManager(), settings, TruncatorTests.createTruncator()) ); var action = actionCreator.create(model); @@ -188,7 +183,7 @@ public void testSend_FailsFromInvalidResponseFormat_ForElserAction() throws IOEx @SuppressWarnings("unchecked") public void testExecute_ReturnsSuccessfulResponse_ForEmbeddingsAction() throws IOException { - var senderFactory = new HttpRequestSenderFactory(threadPool, clientManager, mockClusterServiceEmpty(), Settings.EMPTY); + var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager); try (var sender = senderFactory.createSender("test_service")) { sender.start(); @@ -233,8 +228,14 @@ public void testExecute_ReturnsSuccessfulResponse_ForEmbeddingsAction() throws I } @SuppressWarnings("unchecked") - public void testSend_FailsFromInvalidResponseFormat_ForEmbeddingsAction() throws IOException, URISyntaxException { - var senderFactory = new HttpRequestSenderFactory(threadPool, clientManager, mockClusterServiceEmpty(), Settings.EMPTY); + public void testSend_FailsFromInvalidResponseFormat_ForEmbeddingsAction() throws IOException { + // timeout as zero for no retries + var settings = buildSettingsWithRetryFields( + TimeValue.timeValueMillis(1), + TimeValue.timeValueMinutes(1), + TimeValue.timeValueSeconds(0) + ); + var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager, settings); try (var sender = senderFactory.createSender("test_service")) { sender.start(); @@ -257,17 +258,7 @@ public void testSend_FailsFromInvalidResponseFormat_ForEmbeddingsAction() throws var model = HuggingFaceEmbeddingsModelTests.createModel(getUrl(webServer), "secret"); var actionCreator = new HuggingFaceActionCreator( sender, - new ServiceComponents( - threadPool, - mockThrottlerManager(), - // timeout as zero for no retries - buildSettingsWithRetryFields( - TimeValue.timeValueMillis(1), - TimeValue.timeValueMinutes(1), - TimeValue.timeValueSeconds(0) - ), - TruncatorTests.createTruncator() - ) + new ServiceComponents(threadPool, mockThrottlerManager(), settings, TruncatorTests.createTruncator()) ); var action = actionCreator.create(model); @@ -297,7 +288,7 @@ public void testSend_FailsFromInvalidResponseFormat_ForEmbeddingsAction() throws } public void testExecute_ReturnsSuccessfulResponse_AfterTruncating() throws IOException { - var senderFactory = new HttpRequestSenderFactory(threadPool, clientManager, mockClusterServiceEmpty(), Settings.EMPTY); + var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager); try (var sender = senderFactory.createSender("test_service")) { sender.start(); @@ -362,7 +353,7 @@ public void testExecute_ReturnsSuccessfulResponse_AfterTruncating() throws IOExc } public void testExecute_TruncatesInputBeforeSending() throws IOException { - var senderFactory = new HttpRequestSenderFactory(threadPool, clientManager, mockClusterServiceEmpty(), Settings.EMPTY); + var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager); try (var sender = senderFactory.createSender("test_service")) { sender.start(); diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/action/huggingface/HuggingFaceActionTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/action/huggingface/HuggingFaceActionTests.java index 25b05327a21b7..7b332e8c6634d 100644 --- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/action/huggingface/HuggingFaceActionTests.java +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/action/huggingface/HuggingFaceActionTests.java @@ -16,7 +16,6 @@ import org.elasticsearch.test.ESTestCase; import org.elasticsearch.threadpool.ThreadPool; import org.elasticsearch.xpack.inference.common.TruncatorTests; -import org.elasticsearch.xpack.inference.external.http.HttpResult; import org.elasticsearch.xpack.inference.external.http.retry.AlwaysRetryingResponseHandler; import org.elasticsearch.xpack.inference.external.http.sender.Sender; import org.elasticsearch.xpack.inference.logging.ThrottlerManager; @@ -54,7 +53,7 @@ public void shutdown() throws IOException { public void testExecute_ThrowsElasticsearchException_WhenSenderThrows() { var sender = mock(Sender.class); - doThrow(new ElasticsearchException("failed")).when(sender).send(any(), any()); + doThrow(new ElasticsearchException("failed")).when(sender).send(any(), any(), any()); var action = createAction(URL, sender); @@ -71,11 +70,11 @@ public void testExecute_ThrowsElasticsearchException_WhenSenderOnFailureIsCalled doAnswer(invocation -> { @SuppressWarnings("unchecked") - ActionListener listener = (ActionListener) invocation.getArguments()[1]; + ActionListener listener = (ActionListener) invocation.getArguments()[1]; listener.onFailure(new IllegalStateException("failed")); return Void.TYPE; - }).when(sender).send(any(), any()); + }).when(sender).send(any(), any(), any()); var action = createAction(URL, sender, "inferenceEntityId"); @@ -92,7 +91,7 @@ public void testExecute_ThrowsElasticsearchException_WhenSenderOnFailureIsCalled public void testExecute_ThrowsException() { var sender = mock(Sender.class); - doThrow(new IllegalArgumentException("failed")).when(sender).send(any(), any()); + doThrow(new IllegalArgumentException("failed")).when(sender).send(any(), any(), any()); var action = createAction(URL, sender, "inferenceEntityId"); diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/action/openai/OpenAiActionCreatorTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/action/openai/OpenAiActionCreatorTests.java index cf1a569548143..a844061fa48e1 100644 --- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/action/openai/OpenAiActionCreatorTests.java +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/action/openai/OpenAiActionCreatorTests.java @@ -8,6 +8,7 @@ package org.elasticsearch.xpack.inference.external.action.openai; import org.apache.http.HttpHeaders; +import org.elasticsearch.ElasticsearchStatusException; import org.elasticsearch.action.support.PlainActionFuture; import org.elasticsearch.common.Strings; import org.elasticsearch.common.settings.Settings; @@ -19,7 +20,7 @@ import org.elasticsearch.threadpool.ThreadPool; import org.elasticsearch.xcontent.XContentType; import org.elasticsearch.xpack.inference.external.http.HttpClientManager; -import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSenderFactory; +import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSenderTests; import org.elasticsearch.xpack.inference.logging.ThrottlerManager; import org.junit.After; import org.junit.Before; @@ -28,10 +29,12 @@ import java.util.List; import java.util.concurrent.TimeUnit; +import static org.elasticsearch.core.Strings.format; import static org.elasticsearch.xpack.inference.Utils.inferenceUtilityPool; import static org.elasticsearch.xpack.inference.Utils.mockClusterServiceEmpty; import static org.elasticsearch.xpack.inference.external.http.Utils.entityAsMap; import static org.elasticsearch.xpack.inference.external.http.Utils.getUrl; +import static org.elasticsearch.xpack.inference.external.http.retry.RetrySettingsTests.buildSettingsWithRetryFields; import static org.elasticsearch.xpack.inference.external.request.openai.OpenAiUtils.ORGANIZATION_HEADER; import static org.elasticsearch.xpack.inference.results.TextEmbeddingResultsTests.buildExpectation; import static org.elasticsearch.xpack.inference.services.ServiceComponentsTests.createWithEmptySettings; @@ -63,7 +66,7 @@ public void shutdown() throws IOException { } public void testCreate_OpenAiEmbeddingsModel() throws IOException { - var senderFactory = new HttpRequestSenderFactory(threadPool, clientManager, mockClusterServiceEmpty(), Settings.EMPTY); + var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager); try (var sender = senderFactory.createSender("test_service")) { sender.start(); @@ -115,8 +118,173 @@ public void testCreate_OpenAiEmbeddingsModel() throws IOException { } } + public void testCreate_OpenAiEmbeddingsModel_WithoutUser() throws IOException { + var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager); + + try (var sender = senderFactory.createSender("test_service")) { + sender.start(); + + String responseJson = """ + { + "object": "list", + "data": [ + { + "object": "embedding", + "index": 0, + "embedding": [ + 0.0123, + -0.0123 + ] + } + ], + "model": "text-embedding-ada-002-v2", + "usage": { + "prompt_tokens": 8, + "total_tokens": 8 + } + } + """; + webServer.enqueue(new MockResponse().setResponseCode(200).setBody(responseJson)); + + var model = createModel(getUrl(webServer), "org", "secret", "model", null); + var actionCreator = new OpenAiActionCreator(sender, createWithEmptySettings(threadPool)); + var overriddenTaskSettings = getRequestTaskSettingsMap(null); + var action = actionCreator.create(model, overriddenTaskSettings); + + PlainActionFuture listener = new PlainActionFuture<>(); + action.execute(List.of("abc"), listener); + + var result = listener.actionGet(TIMEOUT); + + assertThat(result.asMap(), is(buildExpectation(List.of(List.of(0.0123F, -0.0123F))))); + assertThat(webServer.requests(), hasSize(1)); + assertNull(webServer.requests().get(0).getUri().getQuery()); + assertThat(webServer.requests().get(0).getHeader(HttpHeaders.CONTENT_TYPE), equalTo(XContentType.JSON.mediaType())); + assertThat(webServer.requests().get(0).getHeader(HttpHeaders.AUTHORIZATION), equalTo("Bearer secret")); + assertThat(webServer.requests().get(0).getHeader(ORGANIZATION_HEADER), equalTo("org")); + + var requestMap = entityAsMap(webServer.requests().get(0).getBody()); + assertThat(requestMap.size(), is(2)); + assertThat(requestMap.get("input"), is(List.of("abc"))); + assertThat(requestMap.get("model"), is("model")); + } + } + + public void testCreate_OpenAiEmbeddingsModel_WithoutOrganization() throws IOException { + var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager); + + try (var sender = senderFactory.createSender("test_service")) { + sender.start(); + + String responseJson = """ + { + "object": "list", + "data": [ + { + "object": "embedding", + "index": 0, + "embedding": [ + 0.0123, + -0.0123 + ] + } + ], + "model": "text-embedding-ada-002-v2", + "usage": { + "prompt_tokens": 8, + "total_tokens": 8 + } + } + """; + webServer.enqueue(new MockResponse().setResponseCode(200).setBody(responseJson)); + + var model = createModel(getUrl(webServer), null, "secret", "model", null); + var actionCreator = new OpenAiActionCreator(sender, createWithEmptySettings(threadPool)); + var overriddenTaskSettings = getRequestTaskSettingsMap("overridden_user"); + var action = actionCreator.create(model, overriddenTaskSettings); + + PlainActionFuture listener = new PlainActionFuture<>(); + action.execute(List.of("abc"), listener); + + var result = listener.actionGet(TIMEOUT); + + assertThat(result.asMap(), is(buildExpectation(List.of(List.of(0.0123F, -0.0123F))))); + assertThat(webServer.requests(), hasSize(1)); + assertNull(webServer.requests().get(0).getUri().getQuery()); + assertThat(webServer.requests().get(0).getHeader(HttpHeaders.CONTENT_TYPE), equalTo(XContentType.JSON.mediaType())); + assertThat(webServer.requests().get(0).getHeader(HttpHeaders.AUTHORIZATION), equalTo("Bearer secret")); + assertNull(webServer.requests().get(0).getHeader(ORGANIZATION_HEADER)); + + var requestMap = entityAsMap(webServer.requests().get(0).getBody()); + assertThat(requestMap.size(), is(3)); + assertThat(requestMap.get("input"), is(List.of("abc"))); + assertThat(requestMap.get("model"), is("model")); + assertThat(requestMap.get("user"), is("overridden_user")); + } + } + + public void testCreate_OpenAiEmbeddingsModel_FailsFromInvalidResponseFormat() throws IOException { + // timeout as zero for no retries + var settings = buildSettingsWithRetryFields( + TimeValue.timeValueMillis(1), + TimeValue.timeValueMinutes(1), + TimeValue.timeValueSeconds(0) + ); + var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager, settings); + + try (var sender = senderFactory.createSender("test_service")) { + sender.start(); + + String responseJson = """ + { + "object": "list", + "data_does_not_exist": [ + { + "object": "embedding", + "index": 0, + "embedding": [ + 0.0123, + -0.0123 + ] + } + ], + "model": "text-embedding-ada-002-v2", + "usage": { + "prompt_tokens": 8, + "total_tokens": 8 + } + } + """; + webServer.enqueue(new MockResponse().setResponseCode(200).setBody(responseJson)); + + var model = createModel(getUrl(webServer), null, "secret", "model", null); + var actionCreator = new OpenAiActionCreator(sender, createWithEmptySettings(threadPool)); + var overriddenTaskSettings = getRequestTaskSettingsMap("overridden_user"); + var action = actionCreator.create(model, overriddenTaskSettings); + + PlainActionFuture listener = new PlainActionFuture<>(); + action.execute(List.of("abc"), listener); + + var thrownException = expectThrows(ElasticsearchStatusException.class, () -> listener.actionGet(TIMEOUT)); + assertThat(thrownException.getMessage(), is(format("Failed to send OpenAI embeddings request to [%s]", getUrl(webServer)))); + assertThat(thrownException.getCause().getMessage(), is("Failed to find required field [data] in OpenAI embeddings response")); + + assertThat(webServer.requests(), hasSize(1)); + assertNull(webServer.requests().get(0).getUri().getQuery()); + assertThat(webServer.requests().get(0).getHeader(HttpHeaders.CONTENT_TYPE), equalTo(XContentType.JSON.mediaType())); + assertThat(webServer.requests().get(0).getHeader(HttpHeaders.AUTHORIZATION), equalTo("Bearer secret")); + assertNull(webServer.requests().get(0).getHeader(ORGANIZATION_HEADER)); + + var requestMap = entityAsMap(webServer.requests().get(0).getBody()); + assertThat(requestMap.size(), is(3)); + assertThat(requestMap.get("input"), is(List.of("abc"))); + assertThat(requestMap.get("model"), is("model")); + assertThat(requestMap.get("user"), is("overridden_user")); + } + } + public void testExecute_ReturnsSuccessfulResponse_AfterTruncating_From413StatusCode() throws IOException { - var senderFactory = new HttpRequestSenderFactory(threadPool, clientManager, mockClusterServiceEmpty(), Settings.EMPTY); + var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager); try (var sender = senderFactory.createSender("test_service")) { sender.start(); @@ -199,7 +367,7 @@ public void testExecute_ReturnsSuccessfulResponse_AfterTruncating_From413StatusC } public void testExecute_ReturnsSuccessfulResponse_AfterTruncating_From400StatusCode() throws IOException { - var senderFactory = new HttpRequestSenderFactory(threadPool, clientManager, mockClusterServiceEmpty(), Settings.EMPTY); + var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager); try (var sender = senderFactory.createSender("test_service")) { sender.start(); @@ -282,7 +450,7 @@ public void testExecute_ReturnsSuccessfulResponse_AfterTruncating_From400StatusC } public void testExecute_TruncatesInputBeforeSending() throws IOException { - var senderFactory = new HttpRequestSenderFactory(threadPool, clientManager, mockClusterServiceEmpty(), Settings.EMPTY); + var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager); try (var sender = senderFactory.createSender("test_service")) { sender.start(); diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/action/openai/OpenAiEmbeddingsActionTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/action/openai/OpenAiEmbeddingsActionTests.java index 6bc8e2d61d579..c803121e6da79 100644 --- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/action/openai/OpenAiEmbeddingsActionTests.java +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/action/openai/OpenAiEmbeddingsActionTests.java @@ -21,10 +21,10 @@ import org.elasticsearch.threadpool.ThreadPool; import org.elasticsearch.xcontent.XContentType; import org.elasticsearch.xpack.inference.external.http.HttpClientManager; -import org.elasticsearch.xpack.inference.external.http.HttpResult; -import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSenderFactory; +import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSender; import org.elasticsearch.xpack.inference.external.http.sender.Sender; import org.elasticsearch.xpack.inference.logging.ThrottlerManager; +import org.elasticsearch.xpack.inference.services.ServiceComponentsTests; import org.junit.After; import org.junit.Before; @@ -70,7 +70,11 @@ public void shutdown() throws IOException { } public void testExecute_ReturnsSuccessfulResponse() throws IOException { - var senderFactory = new HttpRequestSenderFactory(threadPool, clientManager, mockClusterServiceEmpty(), Settings.EMPTY); + var senderFactory = new HttpRequestSender.Factory( + ServiceComponentsTests.createWithEmptySettings(threadPool), + clientManager, + mockClusterServiceEmpty() + ); try (var sender = senderFactory.createSender("test_service")) { sender.start(); @@ -131,7 +135,7 @@ public void testExecute_ThrowsURISyntaxException_ForInvalidUrl() throws IOExcept public void testExecute_ThrowsElasticsearchException() { var sender = mock(Sender.class); - doThrow(new ElasticsearchException("failed")).when(sender).send(any(), any()); + doThrow(new ElasticsearchException("failed")).when(sender).send(any(), any(), any()); var action = createAction(getUrl(webServer), "org", "secret", "model", "user", sender); @@ -148,11 +152,11 @@ public void testExecute_ThrowsElasticsearchException_WhenSenderOnFailureIsCalled doAnswer(invocation -> { @SuppressWarnings("unchecked") - ActionListener listener = (ActionListener) invocation.getArguments()[1]; + ActionListener listener = (ActionListener) invocation.getArguments()[1]; listener.onFailure(new IllegalStateException("failed")); return Void.TYPE; - }).when(sender).send(any(), any()); + }).when(sender).send(any(), any(), any()); var action = createAction(getUrl(webServer), "org", "secret", "model", "user", sender); @@ -169,11 +173,11 @@ public void testExecute_ThrowsElasticsearchException_WhenSenderOnFailureIsCalled doAnswer(invocation -> { @SuppressWarnings("unchecked") - ActionListener listener = (ActionListener) invocation.getArguments()[1]; + ActionListener listener = (ActionListener) invocation.getArguments()[1]; listener.onFailure(new IllegalStateException("failed")); return Void.TYPE; - }).when(sender).send(any(), any()); + }).when(sender).send(any(), any(), any()); var action = createAction(null, "org", "secret", "model", "user", sender); @@ -187,7 +191,7 @@ public void testExecute_ThrowsElasticsearchException_WhenSenderOnFailureIsCalled public void testExecute_ThrowsException() { var sender = mock(Sender.class); - doThrow(new IllegalArgumentException("failed")).when(sender).send(any(), any()); + doThrow(new IllegalArgumentException("failed")).when(sender).send(any(), any(), any()); var action = createAction(getUrl(webServer), "org", "secret", "model", "user", sender); @@ -201,7 +205,7 @@ public void testExecute_ThrowsException() { public void testExecute_ThrowsExceptionWithNullUrl() { var sender = mock(Sender.class); - doThrow(new IllegalArgumentException("failed")).when(sender).send(any(), any()); + doThrow(new IllegalArgumentException("failed")).when(sender).send(any(), any(), any()); var action = createAction(null, "org", "secret", "model", "user", sender); diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/http/retry/RetrySettingsTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/http/retry/RetrySettingsTests.java index 940205a663337..2c63e085a9937 100644 --- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/http/retry/RetrySettingsTests.java +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/http/retry/RetrySettingsTests.java @@ -11,6 +11,8 @@ import org.elasticsearch.core.TimeValue; import org.elasticsearch.test.ESTestCase; +import static org.elasticsearch.xpack.inference.Utils.mockClusterServiceEmpty; + public class RetrySettingsTests extends ESTestCase { /** @@ -24,7 +26,7 @@ public static RetrySettings createDefaultRetrySettings() { public static RetrySettings createRetrySettings(TimeValue initialDelay, TimeValue maxDelayBound, TimeValue timeout) { var settings = buildSettingsWithRetryFields(initialDelay, maxDelayBound, timeout); - return new RetrySettings(settings); + return new RetrySettings(settings, mockClusterServiceEmpty()); } public static Settings buildSettingsWithRetryFields(TimeValue initialDelay, TimeValue maxDelayBound, TimeValue timeout) { diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/http/retry/RetryingHttpSenderTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/http/retry/RetryingHttpSenderTests.java index 8d60c2f5bfa48..30bd40bdcc111 100644 --- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/http/retry/RetryingHttpSenderTests.java +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/http/retry/RetryingHttpSenderTests.java @@ -10,7 +10,9 @@ import org.apache.http.ConnectionClosedException; import org.apache.http.HttpResponse; import org.apache.http.StatusLine; +import org.apache.http.client.protocol.HttpClientContext; import org.apache.logging.log4j.Logger; +import org.elasticsearch.ElasticsearchException; import org.elasticsearch.ElasticsearchStatusException; import org.elasticsearch.action.ActionListener; import org.elasticsearch.action.support.PlainActionFuture; @@ -19,14 +21,15 @@ import org.elasticsearch.core.TimeValue; import org.elasticsearch.inference.InferenceServiceResults; import org.elasticsearch.test.ESTestCase; +import org.elasticsearch.xpack.inference.external.http.HttpClient; import org.elasticsearch.xpack.inference.external.http.HttpResult; -import org.elasticsearch.xpack.inference.external.http.sender.Sender; import org.elasticsearch.xpack.inference.external.request.HttpRequestTests; import org.elasticsearch.xpack.inference.external.request.Request; import org.elasticsearch.xpack.inference.logging.ThrottlerManager; import org.junit.Before; import org.mockito.stubbing.Answer; +import java.io.IOException; import java.net.UnknownHostException; import static org.elasticsearch.xpack.inference.external.http.retry.RetrySettingsTests.createDefaultRetrySettings; @@ -50,17 +53,17 @@ public void init() throws Exception { taskQueue = new DeterministicTaskQueue(); } - public void testSend_CallsSenderAgain_AfterValidateResponseThrowsAnException() { - var sender = mock(Sender.class); + public void testSend_CallsSenderAgain_AfterValidateResponseThrowsAnException() throws IOException { + var httpClient = mock(HttpClient.class); var httpResponse = mockHttpResponse(); doAnswer(invocation -> { @SuppressWarnings("unchecked") - ActionListener listener = (ActionListener) invocation.getArguments()[1]; + ActionListener listener = (ActionListener) invocation.getArguments()[2]; listener.onResponse(new HttpResult(httpResponse, new byte[0])); return Void.TYPE; - }).when(sender).send(any(), any()); + }).when(httpClient).send(any(), any(), any()); var inferenceResults = mock(InferenceServiceResults.class); Answer answer = (invocation) -> inferenceResults; @@ -71,72 +74,58 @@ public void testSend_CallsSenderAgain_AfterValidateResponseThrowsAnException() { // bounded wild card list, thenAnswer must be used instead. when(handler.parseResult(any(), any())).thenAnswer(answer); - var retrier = new RetryingHttpSender( - sender, - mock(ThrottlerManager.class), - mock(Logger.class), - createDefaultRetrySettings(), - taskQueue.getThreadPool(), - EsExecutors.DIRECT_EXECUTOR_SERVICE - ); + var retrier = createRetrier(httpClient); var listener = new PlainActionFuture(); - executeTasks(() -> retrier.send(mockRequest(), handler, listener), 1); + executeTasks(() -> retrier.send(mock(Logger.class), mockRequest(), HttpClientContext.create(), () -> false, handler, listener), 1); assertThat(listener.actionGet(TIMEOUT), is(inferenceResults)); - verify(sender, times(2)).send(any(), any()); - verifyNoMoreInteractions(sender); + verify(httpClient, times(2)).send(any(), any(), any()); + verifyNoMoreInteractions(httpClient); } - public void testSend_CallsSenderAgain_WhenAFailureStatusCodeIsReturned() { + public void testSend_CallsSenderAgain_WhenAFailureStatusCodeIsReturned() throws IOException { var statusLine = mock(StatusLine.class); when(statusLine.getStatusCode()).thenReturn(300).thenReturn(200); var httpResponse = mock(HttpResponse.class); when(httpResponse.getStatusLine()).thenReturn(statusLine); - var sender = mock(Sender.class); + var httpClient = mock(HttpClient.class); doAnswer(invocation -> { @SuppressWarnings("unchecked") - ActionListener listener = (ActionListener) invocation.getArguments()[1]; + ActionListener listener = (ActionListener) invocation.getArguments()[2]; listener.onResponse(new HttpResult(httpResponse, new byte[] { 'a' })); return Void.TYPE; - }).when(sender).send(any(), any()); + }).when(httpClient).send(any(), any(), any()); var inferenceResults = mock(InferenceServiceResults.class); var handler = new AlwaysRetryingResponseHandler("test", result -> inferenceResults); - var retrier = new RetryingHttpSender( - sender, - mock(ThrottlerManager.class), - mock(Logger.class), - createDefaultRetrySettings(), - taskQueue.getThreadPool(), - EsExecutors.DIRECT_EXECUTOR_SERVICE - ); + var retrier = createRetrier(httpClient); var listener = new PlainActionFuture(); - executeTasks(() -> retrier.send(mockRequest(), handler, listener), 1); + executeTasks(() -> retrier.send(mock(Logger.class), mockRequest(), HttpClientContext.create(), () -> false, handler, listener), 1); assertThat(listener.actionGet(TIMEOUT), is(inferenceResults)); - verify(sender, times(2)).send(any(), any()); - verifyNoMoreInteractions(sender); + verify(httpClient, times(2)).send(any(), any(), any()); + verifyNoMoreInteractions(httpClient); } - public void testSend_CallsSenderAgain_WhenParsingFailsOnce() { - var sender = mock(Sender.class); + public void testSend_CallsSenderAgain_WhenParsingFailsOnce() throws IOException { + var httpClient = mock(HttpClient.class); var httpResponse = mockHttpResponse(); doAnswer(invocation -> { @SuppressWarnings("unchecked") - ActionListener listener = (ActionListener) invocation.getArguments()[1]; + ActionListener listener = (ActionListener) invocation.getArguments()[2]; listener.onResponse(new HttpResult(httpResponse, new byte[] { 'a' })); return Void.TYPE; - }).when(sender).send(any(), any()); + }).when(httpClient).send(any(), any(), any()); var inferenceResults = mock(InferenceServiceResults.class); Answer answer = (invocation) -> inferenceResults; @@ -144,34 +133,27 @@ public void testSend_CallsSenderAgain_WhenParsingFailsOnce() { var handler = mock(ResponseHandler.class); when(handler.parseResult(any(), any())).thenThrow(new RetryException(true, "failed")).thenAnswer(answer); - var retrier = new RetryingHttpSender( - sender, - mock(ThrottlerManager.class), - mock(Logger.class), - createDefaultRetrySettings(), - taskQueue.getThreadPool(), - EsExecutors.DIRECT_EXECUTOR_SERVICE - ); + var retrier = createRetrier(httpClient); var listener = new PlainActionFuture(); - executeTasks(() -> retrier.send(mockRequest(), handler, listener), 1); + executeTasks(() -> retrier.send(mock(Logger.class), mockRequest(), HttpClientContext.create(), () -> false, handler, listener), 1); assertThat(listener.actionGet(TIMEOUT), is(inferenceResults)); - verify(sender, times(2)).send(any(), any()); - verifyNoMoreInteractions(sender); + verify(httpClient, times(2)).send(any(), any(), any()); + verifyNoMoreInteractions(httpClient); } - public void testSend_DoesNotCallSenderAgain_WhenParsingFailsWithNonRetryableException() { - var sender = mock(Sender.class); + public void testSend_DoesNotCallSenderAgain_WhenParsingFailsWithNonRetryableException() throws IOException { + var httpClient = mock(HttpClient.class); var httpResponse = mockHttpResponse(); doAnswer(invocation -> { @SuppressWarnings("unchecked") - ActionListener listener = (ActionListener) invocation.getArguments()[1]; + ActionListener listener = (ActionListener) invocation.getArguments()[2]; listener.onResponse(new HttpResult(httpResponse, new byte[] { 'a' })); return Void.TYPE; - }).when(sender).send(any(), any()); + }).when(httpClient).send(any(), any(), any()); var inferenceResults = mock(InferenceServiceResults.class); Answer answer = (invocation) -> inferenceResults; @@ -179,41 +161,34 @@ public void testSend_DoesNotCallSenderAgain_WhenParsingFailsWithNonRetryableExce var handler = mock(ResponseHandler.class); when(handler.parseResult(any(), any())).thenThrow(new IllegalStateException("failed")).thenAnswer(answer); - var retrier = new RetryingHttpSender( - sender, - mock(ThrottlerManager.class), - mock(Logger.class), - createDefaultRetrySettings(), - taskQueue.getThreadPool(), - EsExecutors.DIRECT_EXECUTOR_SERVICE - ); + var retrier = createRetrier(httpClient); var listener = new PlainActionFuture(); - executeTasks(() -> retrier.send(mockRequest(), handler, listener), 0); + executeTasks(() -> retrier.send(mock(Logger.class), mockRequest(), HttpClientContext.create(), () -> false, handler, listener), 0); var thrownException = expectThrows(IllegalStateException.class, () -> listener.actionGet(TIMEOUT)); assertThat(thrownException.getMessage(), is("failed")); - verify(sender, times(1)).send(any(), any()); - verifyNoMoreInteractions(sender); + verify(httpClient, times(1)).send(any(), any(), any()); + verifyNoMoreInteractions(httpClient); } - public void testSend_CallsSenderAgain_WhenHttpResultListenerCallsOnFailureOnce() { - var sender = mock(Sender.class); + public void testSend_CallsSenderAgain_WhenHttpResultListenerCallsOnFailureOnce() throws IOException { + var httpClient = mock(HttpClient.class); doAnswer(invocation -> { @SuppressWarnings("unchecked") - ActionListener listener = (ActionListener) invocation.getArguments()[1]; + ActionListener listener = (ActionListener) invocation.getArguments()[2]; listener.onFailure(new RetryException(true, "failed")); return Void.TYPE; }).doAnswer(invocation -> { @SuppressWarnings("unchecked") - ActionListener listener = (ActionListener) invocation.getArguments()[1]; + ActionListener listener = (ActionListener) invocation.getArguments()[2]; listener.onResponse(new HttpResult(mock(HttpResponse.class), new byte[] { 'a' })); return Void.TYPE; - }).when(sender).send(any(), any()); + }).when(httpClient).send(any(), any(), any()); var inferenceResults = mock(InferenceServiceResults.class); Answer answer = (invocation) -> inferenceResults; @@ -221,39 +196,32 @@ public void testSend_CallsSenderAgain_WhenHttpResultListenerCallsOnFailureOnce() var handler = mock(ResponseHandler.class); when(handler.parseResult(any(), any())).thenAnswer(answer); - var retrier = new RetryingHttpSender( - sender, - mock(ThrottlerManager.class), - mock(Logger.class), - createDefaultRetrySettings(), - taskQueue.getThreadPool(), - EsExecutors.DIRECT_EXECUTOR_SERVICE - ); + var retrier = createRetrier(httpClient); var listener = new PlainActionFuture(); - executeTasks(() -> retrier.send(mockRequest(), handler, listener), 1); + executeTasks(() -> retrier.send(mock(Logger.class), mockRequest(), HttpClientContext.create(), () -> false, handler, listener), 1); assertThat(listener.actionGet(TIMEOUT), is(inferenceResults)); - verify(sender, times(2)).send(any(), any()); - verifyNoMoreInteractions(sender); + verify(httpClient, times(2)).send(any(), any(), any()); + verifyNoMoreInteractions(httpClient); } - public void testSend_CallsSenderAgain_WhenHttpResultListenerCallsOnFailureOnce_WithContentTooLargeException() { - var sender = mock(Sender.class); + public void testSend_CallsSenderAgain_WhenHttpResultListenerCallsOnFailureOnce_WithContentTooLargeException() throws IOException { + var httpClient = mock(HttpClient.class); doAnswer(invocation -> { @SuppressWarnings("unchecked") - ActionListener listener = (ActionListener) invocation.getArguments()[1]; + ActionListener listener = (ActionListener) invocation.getArguments()[2]; listener.onFailure(new ContentTooLargeException(new IllegalStateException("failed"))); return Void.TYPE; }).doAnswer(invocation -> { @SuppressWarnings("unchecked") - ActionListener listener = (ActionListener) invocation.getArguments()[1]; + ActionListener listener = (ActionListener) invocation.getArguments()[2]; listener.onResponse(new HttpResult(mock(HttpResponse.class), new byte[] { 'a' })); return Void.TYPE; - }).when(sender).send(any(), any()); + }).when(httpClient).send(any(), any(), any()); var inferenceResults = mock(InferenceServiceResults.class); Answer answer = (invocation) -> inferenceResults; @@ -261,39 +229,32 @@ public void testSend_CallsSenderAgain_WhenHttpResultListenerCallsOnFailureOnce_W var handler = mock(ResponseHandler.class); when(handler.parseResult(any(), any())).thenAnswer(answer); - var retrier = new RetryingHttpSender( - sender, - mock(ThrottlerManager.class), - mock(Logger.class), - createDefaultRetrySettings(), - taskQueue.getThreadPool(), - EsExecutors.DIRECT_EXECUTOR_SERVICE - ); + var retrier = createRetrier(httpClient); var listener = new PlainActionFuture(); - executeTasks(() -> retrier.send(mockRequest(), handler, listener), 1); + executeTasks(() -> retrier.send(mock(Logger.class), mockRequest(), HttpClientContext.create(), () -> false, handler, listener), 1); assertThat(listener.actionGet(TIMEOUT), is(inferenceResults)); - verify(sender, times(2)).send(any(), any()); - verifyNoMoreInteractions(sender); + verify(httpClient, times(2)).send(any(), any(), any()); + verifyNoMoreInteractions(httpClient); } - public void testSend_CallsSenderAgain_WhenHttpResultListenerCallsOnFailureOnceWithConnectionClosedException() { - var sender = mock(Sender.class); + public void testSend_CallsSenderAgain_WhenHttpResultListenerCallsOnFailureOnceWithConnectionClosedException() throws IOException { + var httpClient = mock(HttpClient.class); doAnswer(invocation -> { @SuppressWarnings("unchecked") - ActionListener listener = (ActionListener) invocation.getArguments()[1]; + ActionListener listener = (ActionListener) invocation.getArguments()[2]; listener.onFailure(new ConnectionClosedException("failed")); return Void.TYPE; }).doAnswer(invocation -> { @SuppressWarnings("unchecked") - ActionListener listener = (ActionListener) invocation.getArguments()[1]; + ActionListener listener = (ActionListener) invocation.getArguments()[2]; listener.onResponse(new HttpResult(mock(HttpResponse.class), new byte[] { 'a' })); return Void.TYPE; - }).when(sender).send(any(), any()); + }).when(httpClient).send(any(), any(), any()); var inferenceResults = mock(InferenceServiceResults.class); Answer answer = (invocation) -> inferenceResults; @@ -301,33 +262,26 @@ public void testSend_CallsSenderAgain_WhenHttpResultListenerCallsOnFailureOnceWi var handler = mock(ResponseHandler.class); when(handler.parseResult(any(), any())).thenAnswer(answer); - var retrier = new RetryingHttpSender( - sender, - mock(ThrottlerManager.class), - mock(Logger.class), - createDefaultRetrySettings(), - taskQueue.getThreadPool(), - EsExecutors.DIRECT_EXECUTOR_SERVICE - ); + var retrier = createRetrier(httpClient); var listener = new PlainActionFuture(); - executeTasks(() -> retrier.send(mockRequest(), handler, listener), 1); + executeTasks(() -> retrier.send(mock(Logger.class), mockRequest(), HttpClientContext.create(), () -> false, handler, listener), 1); assertThat(listener.actionGet(TIMEOUT), is(inferenceResults)); - verify(sender, times(2)).send(any(), any()); - verifyNoMoreInteractions(sender); + verify(httpClient, times(2)).send(any(), any(), any()); + verifyNoMoreInteractions(httpClient); } - public void testSend_ReturnsFailure_WhenHttpResultListenerCallsOnFailureOnceWithUnknownHostException() { - var sender = mock(Sender.class); + public void testSend_ReturnsFailure_WhenHttpResultListenerCallsOnFailureOnceWithUnknownHostException() throws IOException { + var httpClient = mock(HttpClient.class); doAnswer(invocation -> { @SuppressWarnings("unchecked") - ActionListener listener = (ActionListener) invocation.getArguments()[1]; + ActionListener listener = (ActionListener) invocation.getArguments()[2]; listener.onFailure(new UnknownHostException("failed")); return Void.TYPE; - }).when(sender).send(any(), any()); + }).when(httpClient).send(any(), any(), any()); var inferenceResults = mock(InferenceServiceResults.class); Answer answer = (invocation) -> inferenceResults; @@ -335,37 +289,55 @@ public void testSend_ReturnsFailure_WhenHttpResultListenerCallsOnFailureOnceWith var handler = mock(ResponseHandler.class); when(handler.parseResult(any(), any())).thenAnswer(answer); - var retrier = new RetryingHttpSender( - sender, - mock(ThrottlerManager.class), - mock(Logger.class), - createDefaultRetrySettings(), - taskQueue.getThreadPool(), - EsExecutors.DIRECT_EXECUTOR_SERVICE - ); + var retrier = createRetrier(httpClient); var listener = new PlainActionFuture(); - executeTasks(() -> retrier.send(mockRequest(), handler, listener), 0); + executeTasks(() -> retrier.send(mock(Logger.class), mockRequest(), HttpClientContext.create(), () -> false, handler, listener), 0); var thrownException = expectThrows(ElasticsearchStatusException.class, () -> listener.actionGet(TIMEOUT)); assertThat(thrownException.getMessage(), is("Invalid host [null], please check that the URL is correct.")); - verify(sender, times(1)).send(any(), any()); - verifyNoMoreInteractions(sender); + verify(httpClient, times(1)).send(any(), any(), any()); + verifyNoMoreInteractions(httpClient); + } + + public void testSend_ReturnsElasticsearchExceptionFailure_WhenTheHttpClientThrowsAnIllegalStateException() throws IOException { + var httpClient = mock(HttpClient.class); + + doAnswer(invocation -> { throw new IllegalStateException("failed"); }).when(httpClient).send(any(), any(), any()); + + var inferenceResults = mock(InferenceServiceResults.class); + Answer answer = (invocation) -> inferenceResults; + + var handler = mock(ResponseHandler.class); + when(handler.parseResult(any(), any())).thenAnswer(answer); + + var retrier = createRetrier(httpClient); + + var listener = new PlainActionFuture(); + executeTasks( + () -> retrier.send(mock(Logger.class), mockRequest("id"), HttpClientContext.create(), () -> false, handler, listener), + 0 + ); + + var thrownException = expectThrows(ElasticsearchException.class, () -> listener.actionGet(TIMEOUT)); + assertThat(thrownException.getMessage(), is("Http client failed to send request from inference entity id [id]")); + verify(httpClient, times(1)).send(any(), any(), any()); + verifyNoMoreInteractions(httpClient); } - public void testSend_ReturnsFailure_WhenValidateResponseThrowsAnException_AfterOneRetry() { + public void testSend_ReturnsFailure_WhenValidateResponseThrowsAnException_AfterOneRetry() throws IOException { var httpResponse = mock(HttpResponse.class); when(httpResponse.getStatusLine()).thenReturn(mock(StatusLine.class)); - var sender = mock(Sender.class); + var sender = mock(HttpClient.class); doAnswer(invocation -> { @SuppressWarnings("unchecked") - ActionListener listener = (ActionListener) invocation.getArguments()[1]; + ActionListener listener = (ActionListener) invocation.getArguments()[2]; listener.onResponse(new HttpResult(httpResponse, new byte[0])); return Void.TYPE; - }).when(sender).send(any(), any()); + }).when(sender).send(any(), any(), any()); var inferenceResults = mock(InferenceServiceResults.class); Answer answer = (invocation) -> inferenceResults; @@ -376,40 +348,33 @@ public void testSend_ReturnsFailure_WhenValidateResponseThrowsAnException_AfterO .validateResponse(any(), any(), any(), any()); when(handler.parseResult(any(), any())).thenAnswer(answer); - var retrier = new RetryingHttpSender( - sender, - mock(ThrottlerManager.class), - mock(Logger.class), - createDefaultRetrySettings(), - taskQueue.getThreadPool(), - EsExecutors.DIRECT_EXECUTOR_SERVICE - ); + var retrier = createRetrier(sender); var listener = new PlainActionFuture(); - executeTasks(() -> retrier.send(mockRequest(), handler, listener), 1); + executeTasks(() -> retrier.send(mock(Logger.class), mockRequest(), HttpClientContext.create(), () -> false, handler, listener), 1); var thrownException = expectThrows(IllegalStateException.class, () -> listener.actionGet(TIMEOUT)); assertThat(thrownException.getMessage(), is("failed again")); assertThat(thrownException.getSuppressed().length, is(1)); assertThat(thrownException.getSuppressed()[0].getMessage(), is("failed")); - verify(sender, times(2)).send(any(), any()); + verify(sender, times(2)).send(any(), any(), any()); verifyNoMoreInteractions(sender); } - public void testSend_ReturnsFailure_WhenValidateResponseThrowsAnElasticsearchException_AfterOneRetry() { + public void testSend_ReturnsFailure_WhenValidateResponseThrowsAnElasticsearchException_AfterOneRetry() throws IOException { var httpResponse = mock(HttpResponse.class); when(httpResponse.getStatusLine()).thenReturn(mock(StatusLine.class)); - var sender = mock(Sender.class); + var httpClient = mock(HttpClient.class); doAnswer(invocation -> { @SuppressWarnings("unchecked") - ActionListener listener = (ActionListener) invocation.getArguments()[1]; + ActionListener listener = (ActionListener) invocation.getArguments()[2]; listener.onResponse(new HttpResult(httpResponse, new byte[0])); return Void.TYPE; - }).when(sender).send(any(), any()); + }).when(httpClient).send(any(), any(), any()); var inferenceResults = mock(InferenceServiceResults.class); Answer answer = (invocation) -> inferenceResults; @@ -420,101 +385,74 @@ public void testSend_ReturnsFailure_WhenValidateResponseThrowsAnElasticsearchExc .validateResponse(any(), any(), any(), any()); when(handler.parseResult(any(), any())).thenAnswer(answer); - var retrier = new RetryingHttpSender( - sender, - mock(ThrottlerManager.class), - mock(Logger.class), - createDefaultRetrySettings(), - taskQueue.getThreadPool(), - EsExecutors.DIRECT_EXECUTOR_SERVICE - ); + var retrier = createRetrier(httpClient); var listener = new PlainActionFuture(); - executeTasks(() -> retrier.send(mockRequest(), handler, listener), 1); + executeTasks(() -> retrier.send(mock(Logger.class), mockRequest(), HttpClientContext.create(), () -> false, handler, listener), 1); var thrownException = expectThrows(RetryException.class, () -> listener.actionGet(TIMEOUT)); assertThat(thrownException.getMessage(), is("failed again")); assertThat(thrownException.getSuppressed().length, is(1)); assertThat(thrownException.getSuppressed()[0].getMessage(), is("failed")); - verify(sender, times(2)).send(any(), any()); - verifyNoMoreInteractions(sender); + verify(httpClient, times(2)).send(any(), any(), any()); + verifyNoMoreInteractions(httpClient); } - public void testSend_ReturnsFailure_WhenHttpResultsListenerCallsOnFailure_AfterOneRetry() { - var httpResponse = mock(HttpResponse.class); - when(httpResponse.getStatusLine()).thenReturn(mock(StatusLine.class)); - - var sender = mock(Sender.class); + public void testSend_ReturnsFailure_WhenHttpResultsListenerCallsOnFailure_AfterOneRetry() throws IOException { + var httpClient = mock(HttpClient.class); doAnswer(invocation -> { @SuppressWarnings("unchecked") - ActionListener listener = (ActionListener) invocation.getArguments()[1]; + ActionListener listener = (ActionListener) invocation.getArguments()[2]; listener.onFailure(new RetryException(true, "failed")); return Void.TYPE; }).doAnswer(invocation -> { @SuppressWarnings("unchecked") - ActionListener listener = (ActionListener) invocation.getArguments()[1]; + ActionListener listener = (ActionListener) invocation.getArguments()[2]; listener.onFailure(new RetryException(false, "failed again")); return Void.TYPE; - }).when(sender).send(any(), any()); + }).when(httpClient).send(any(), any(), any()); var handler = mock(ResponseHandler.class); - var retrier = new RetryingHttpSender( - sender, - mock(ThrottlerManager.class), - mock(Logger.class), - createDefaultRetrySettings(), - taskQueue.getThreadPool(), - EsExecutors.DIRECT_EXECUTOR_SERVICE - ); + var retrier = createRetrier(httpClient); var listener = new PlainActionFuture(); - executeTasks(() -> retrier.send(mockRequest(), handler, listener), 1); + executeTasks(() -> retrier.send(mock(Logger.class), mockRequest(), HttpClientContext.create(), () -> false, handler, listener), 1); var thrownException = expectThrows(RetryException.class, () -> listener.actionGet(TIMEOUT)); assertThat(thrownException.getMessage(), is("failed again")); assertThat(thrownException.getSuppressed().length, is(1)); assertThat(thrownException.getSuppressed()[0].getMessage(), is("failed")); - verify(sender, times(2)).send(any(), any()); - verifyNoMoreInteractions(sender); + verify(httpClient, times(2)).send(any(), any(), any()); + verifyNoMoreInteractions(httpClient); } - public void testSend_ReturnsFailure_WhenHttpResultsListenerCallsOnFailure_WithNonRetryableException() { - var httpResponse = mock(HttpResponse.class); - when(httpResponse.getStatusLine()).thenReturn(mock(StatusLine.class)); - - var sender = mock(Sender.class); + public void testSend_ReturnsFailure_WhenHttpResultsListenerCallsOnFailure_WithNonRetryableException() throws IOException { + var httpClient = mock(HttpClient.class); doAnswer(invocation -> { @SuppressWarnings("unchecked") - ActionListener listener = (ActionListener) invocation.getArguments()[1]; + ActionListener listener = (ActionListener) invocation.getArguments()[2]; listener.onFailure(new IllegalStateException("failed")); return Void.TYPE; - }).when(sender).send(any(), any()); + }).when(httpClient).send(any(), any(), any()); var handler = mock(ResponseHandler.class); - var retrier = new RetryingHttpSender( - sender, - mock(ThrottlerManager.class), - mock(Logger.class), - createDefaultRetrySettings(), - taskQueue.getThreadPool(), - EsExecutors.DIRECT_EXECUTOR_SERVICE - ); + var retrier = createRetrier(httpClient); var listener = new PlainActionFuture(); - executeTasks(() -> retrier.send(mockRequest(), handler, listener), 0); + executeTasks(() -> retrier.send(mock(Logger.class), mockRequest(), HttpClientContext.create(), () -> false, handler, listener), 0); var thrownException = expectThrows(IllegalStateException.class, () -> listener.actionGet(TIMEOUT)); assertThat(thrownException.getMessage(), is("failed")); assertThat(thrownException.getSuppressed().length, is(0)); - verify(sender, times(1)).send(any(), any()); - verifyNoMoreInteractions(sender); + verify(httpClient, times(1)).send(any(), any(), any()); + verifyNoMoreInteractions(httpClient); } private static HttpResponse mockHttpResponse() { @@ -540,10 +478,25 @@ private void executeTasks(Runnable runnable, int retries) { } private static Request mockRequest() { + return mockRequest("inferenceEntityId"); + } + + private static Request mockRequest(String inferenceEntityId) { var request = mock(Request.class); when(request.truncate()).thenReturn(request); - when(request.createHttpRequest()).thenReturn(HttpRequestTests.createMock("inferenceEntityId")); + when(request.createHttpRequest()).thenReturn(HttpRequestTests.createMock(inferenceEntityId)); + when(request.getInferenceEntityId()).thenReturn(inferenceEntityId); return request; } + + private RetryingHttpSender createRetrier(HttpClient httpClient) { + return new RetryingHttpSender( + httpClient, + mock(ThrottlerManager.class), + createDefaultRetrySettings(), + taskQueue.getThreadPool(), + EsExecutors.DIRECT_EXECUTOR_SERVICE + ); + } } diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/http/sender/ExecutableRequestCreatorTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/http/sender/ExecutableRequestCreatorTests.java new file mode 100644 index 0000000000000..24f930a818134 --- /dev/null +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/http/sender/ExecutableRequestCreatorTests.java @@ -0,0 +1,49 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.inference.external.http.sender; + +import org.apache.http.client.protocol.HttpClientContext; +import org.apache.logging.log4j.Logger; +import org.elasticsearch.action.support.PlainActionFuture; +import org.elasticsearch.xpack.inference.external.http.retry.RequestSender; +import org.elasticsearch.xpack.inference.external.http.retry.ResponseHandler; +import org.elasticsearch.xpack.inference.external.request.RequestTests; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyList; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +public class ExecutableRequestCreatorTests { + public static ExecutableRequestCreator createMock() { + var mockCreator = mock(ExecutableRequestCreator.class); + when(mockCreator.create(anyList(), any(), any(), any(), any())).thenReturn(() -> {}); + + return mockCreator; + } + + public static ExecutableRequestCreator createMock(RequestSender requestSender) { + return createMock(requestSender, "id"); + } + + public static ExecutableRequestCreator createMock(RequestSender requestSender, String modelId) { + var mockCreator = mock(ExecutableRequestCreator.class); + when(mockCreator.create(anyList(), any(), any(), any(), any())).thenReturn(() -> { + requestSender.send( + mock(Logger.class), + RequestTests.mockRequest(modelId), + HttpClientContext.create(), + () -> false, + mock(ResponseHandler.class), + new PlainActionFuture<>() + ); + }); + + return mockCreator; + } +} diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/http/sender/HttpRequestSenderFactoryTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/http/sender/HttpRequestSenderTests.java similarity index 53% rename from x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/http/sender/HttpRequestSenderFactoryTests.java rename to x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/http/sender/HttpRequestSenderTests.java index 6b085f8dd80a7..79b17f8dff29d 100644 --- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/http/sender/HttpRequestSenderFactoryTests.java +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/http/sender/HttpRequestSenderTests.java @@ -13,6 +13,7 @@ import org.elasticsearch.common.settings.Settings; import org.elasticsearch.common.util.concurrent.ThreadContext; import org.elasticsearch.core.TimeValue; +import org.elasticsearch.inference.InferenceServiceResults; import org.elasticsearch.test.ESTestCase; import org.elasticsearch.test.http.MockResponse; import org.elasticsearch.test.http.MockWebServer; @@ -21,14 +22,13 @@ import org.elasticsearch.xcontent.XContentType; import org.elasticsearch.xpack.inference.external.http.HttpClient; import org.elasticsearch.xpack.inference.external.http.HttpClientManager; -import org.elasticsearch.xpack.inference.external.http.HttpResult; -import org.elasticsearch.xpack.inference.external.request.HttpRequestTests; import org.elasticsearch.xpack.inference.logging.ThrottlerManager; +import org.elasticsearch.xpack.inference.services.ServiceComponentsTests; import org.junit.After; import org.junit.Before; import java.io.IOException; -import java.nio.charset.StandardCharsets; +import java.util.List; import java.util.concurrent.ExecutorService; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicReference; @@ -36,7 +36,10 @@ import static org.elasticsearch.core.Strings.format; import static org.elasticsearch.xpack.inference.Utils.inferenceUtilityPool; import static org.elasticsearch.xpack.inference.Utils.mockClusterServiceEmpty; -import static org.elasticsearch.xpack.inference.external.http.HttpClientTests.createHttpPost; +import static org.elasticsearch.xpack.inference.external.http.Utils.entityAsMap; +import static org.elasticsearch.xpack.inference.external.http.Utils.getUrl; +import static org.elasticsearch.xpack.inference.external.request.openai.OpenAiUtils.ORGANIZATION_HEADER; +import static org.elasticsearch.xpack.inference.results.TextEmbeddingResultsTests.buildExpectation; import static org.hamcrest.Matchers.equalTo; import static org.hamcrest.Matchers.hasSize; import static org.hamcrest.Matchers.instanceOf; @@ -47,7 +50,7 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; -public class HttpRequestSenderFactoryTests extends ESTestCase { +public class HttpRequestSenderTests extends ESTestCase { private static final TimeValue TIMEOUT = new TimeValue(30, TimeUnit.SECONDS); private final MockWebServer webServer = new MockWebServer(); private ThreadPool threadPool; @@ -79,36 +82,63 @@ public void testCreateSender_SendsRequestAndReceivesResponse() throws Exception try (var sender = senderFactory.createSender("test_service")) { sender.start(); - int responseCode = randomIntBetween(200, 203); - String body = randomAlphaOfLengthBetween(2, 8096); - webServer.enqueue(new MockResponse().setResponseCode(responseCode).setBody(body)); - - String paramKey = randomAlphaOfLength(3); - String paramValue = randomAlphaOfLength(3); - var httpPost = createHttpPost(webServer.getPort(), paramKey, paramValue); - - PlainActionFuture listener = new PlainActionFuture<>(); - sender.send(httpPost, null, listener); + String responseJson = """ + { + "object": "list", + "data": [ + { + "object": "embedding", + "index": 0, + "embedding": [ + 0.0123, + -0.0123 + ] + } + ], + "model": "text-embedding-ada-002-v2", + "usage": { + "prompt_tokens": 8, + "total_tokens": 8 + } + } + """; + webServer.enqueue(new MockResponse().setResponseCode(200).setBody(responseJson)); + + PlainActionFuture listener = new PlainActionFuture<>(); + sender.send( + OpenAiEmbeddingsExecutableRequestCreatorTests.makeCreator(getUrl(webServer), null, "key", "model", null), + List.of("abc"), + listener + ); var result = listener.actionGet(TIMEOUT); + assertThat(result.asMap(), is(buildExpectation(List.of(List.of(0.0123F, -0.0123F))))); - assertThat(result.response().getStatusLine().getStatusCode(), equalTo(responseCode)); - assertThat(new String(result.body(), StandardCharsets.UTF_8), is(body)); assertThat(webServer.requests(), hasSize(1)); - assertThat(webServer.requests().get(0).getUri().getPath(), equalTo(httpPost.httpRequestBase().getURI().getPath())); - assertThat(webServer.requests().get(0).getUri().getQuery(), equalTo(paramKey + "=" + paramValue)); + assertNull(webServer.requests().get(0).getUri().getQuery()); assertThat(webServer.requests().get(0).getHeader(HttpHeaders.CONTENT_TYPE), equalTo(XContentType.JSON.mediaType())); + assertThat(webServer.requests().get(0).getHeader(HttpHeaders.AUTHORIZATION), equalTo("Bearer key")); + assertNull(webServer.requests().get(0).getHeader(ORGANIZATION_HEADER)); + + var requestMap = entityAsMap(webServer.requests().get(0).getBody()); + assertThat(requestMap.size(), is(2)); + assertThat(requestMap.get("input"), is(List.of("abc"))); + assertThat(requestMap.get("model"), is("model")); } } public void testHttpRequestSender_Throws_WhenCallingSendBeforeStart() throws Exception { - var senderFactory = new HttpRequestSenderFactory(threadPool, clientManager, mockClusterServiceEmpty(), Settings.EMPTY); + var senderFactory = new HttpRequestSender.Factory( + ServiceComponentsTests.createWithEmptySettings(threadPool), + clientManager, + mockClusterServiceEmpty() + ); try (var sender = senderFactory.createSender("test_service")) { - PlainActionFuture listener = new PlainActionFuture<>(); + PlainActionFuture listener = new PlainActionFuture<>(); var thrownException = expectThrows( AssertionError.class, - () -> sender.send(HttpRequestTests.createMock("inferenceEntityId"), listener) + () -> sender.send(ExecutableRequestCreatorTests.createMock(), List.of(), listener) ); assertThat(thrownException.getMessage(), is("call start() before sending a request")); } @@ -118,23 +148,27 @@ public void testHttpRequestSender_Throws_WhenATimeoutOccurs() throws Exception { var mockManager = mock(HttpClientManager.class); when(mockManager.getHttpClient()).thenReturn(mock(HttpClient.class)); - var senderFactory = new HttpRequestSenderFactory(threadPool, mockManager, mockClusterServiceEmpty(), Settings.EMPTY); + var senderFactory = new HttpRequestSender.Factory( + ServiceComponentsTests.createWithEmptySettings(threadPool), + mockManager, + mockClusterServiceEmpty() + ); try (var sender = senderFactory.createSender("test_service")) { - assertThat(sender, instanceOf(HttpRequestSenderFactory.HttpRequestSender.class)); + assertThat(sender, instanceOf(HttpRequestSender.class)); // hack to get around the sender interface so we can set the timeout directly - var httpSender = (HttpRequestSenderFactory.HttpRequestSender) sender; + var httpSender = (HttpRequestSender) sender; httpSender.setMaxRequestTimeout(TimeValue.timeValueNanos(1)); sender.start(); - PlainActionFuture listener = new PlainActionFuture<>(); - sender.send(HttpRequestTests.createMock("inferenceEntityId"), TimeValue.timeValueNanos(1), listener); + PlainActionFuture listener = new PlainActionFuture<>(); + sender.send(ExecutableRequestCreatorTests.createMock(), List.of(), TimeValue.timeValueNanos(1), listener); var thrownException = expectThrows(ElasticsearchTimeoutException.class, () -> listener.actionGet(TIMEOUT)); assertThat( thrownException.getMessage(), - is(format("Request timed out waiting to be executed after [%s]", TimeValue.timeValueNanos(1))) + is(format("Request timed out waiting to be sent after [%s]", TimeValue.timeValueNanos(1))) ); } } @@ -143,24 +177,28 @@ public void testHttpRequestSenderWithTimeout_Throws_WhenATimeoutOccurs() throws var mockManager = mock(HttpClientManager.class); when(mockManager.getHttpClient()).thenReturn(mock(HttpClient.class)); - var senderFactory = new HttpRequestSenderFactory(threadPool, mockManager, mockClusterServiceEmpty(), Settings.EMPTY); + var senderFactory = new HttpRequestSender.Factory( + ServiceComponentsTests.createWithEmptySettings(threadPool), + mockManager, + mockClusterServiceEmpty() + ); try (var sender = senderFactory.createSender("test_service")) { sender.start(); - PlainActionFuture listener = new PlainActionFuture<>(); - sender.send(HttpRequestTests.createMock("id"), TimeValue.timeValueNanos(1), listener); + PlainActionFuture listener = new PlainActionFuture<>(); + sender.send(ExecutableRequestCreatorTests.createMock(), List.of(), TimeValue.timeValueNanos(1), listener); var thrownException = expectThrows(ElasticsearchTimeoutException.class, () -> listener.actionGet(TIMEOUT)); assertThat( thrownException.getMessage(), - is(format("Request timed out waiting to be executed after [%s]", TimeValue.timeValueNanos(1))) + is(format("Request timed out waiting to be sent after [%s]", TimeValue.timeValueNanos(1))) ); } } - private static HttpRequestSenderFactory createSenderFactory(HttpClientManager clientManager, AtomicReference threadRef) { + private static HttpRequestSender.Factory createSenderFactory(HttpClientManager clientManager, AtomicReference threadRef) { var mockExecutorService = mock(ExecutorService.class); doAnswer(invocation -> { Runnable runnable = (Runnable) invocation.getArguments()[0]; @@ -175,6 +213,34 @@ private static HttpRequestSenderFactory createSenderFactory(HttpClientManager cl when(mockThreadPool.getThreadContext()).thenReturn(new ThreadContext(Settings.EMPTY)); when(mockThreadPool.schedule(any(Runnable.class), any(), any())).thenReturn(mock(Scheduler.ScheduledCancellable.class)); - return new HttpRequestSenderFactory(mockThreadPool, clientManager, mockClusterServiceEmpty(), Settings.EMPTY); + return new HttpRequestSender.Factory( + ServiceComponentsTests.createWithEmptySettings(mockThreadPool), + clientManager, + mockClusterServiceEmpty() + ); + } + + public static HttpRequestSender.Factory createSenderFactory(ThreadPool threadPool, HttpClientManager httpClientManager) { + return new HttpRequestSender.Factory( + ServiceComponentsTests.createWithEmptySettings(threadPool), + httpClientManager, + mockClusterServiceEmpty() + ); + } + + public static HttpRequestSender.Factory createSenderFactory( + ThreadPool threadPool, + HttpClientManager httpClientManager, + Settings settings + ) { + return new HttpRequestSender.Factory( + ServiceComponentsTests.createWithSettings(threadPool, settings), + httpClientManager, + mockClusterServiceEmpty() + ); + } + + public static Sender createSenderWithSingleRequestManager(HttpRequestSender.Factory factory, String serviceName) { + return factory.createSender(serviceName); } } diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/http/sender/OpenAiEmbeddingsExecutableRequestCreatorTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/http/sender/OpenAiEmbeddingsExecutableRequestCreatorTests.java new file mode 100644 index 0000000000000..53537a3ff77c2 --- /dev/null +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/http/sender/OpenAiEmbeddingsExecutableRequestCreatorTests.java @@ -0,0 +1,40 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.inference.external.http.sender; + +import org.elasticsearch.core.Nullable; +import org.elasticsearch.xpack.inference.common.TruncatorTests; + +import static org.elasticsearch.xpack.inference.services.openai.embeddings.OpenAiEmbeddingsModelTests.createModel; + +public class OpenAiEmbeddingsExecutableRequestCreatorTests { + public static OpenAiEmbeddingsExecutableRequestCreator makeCreator( + String url, + @Nullable String org, + String apiKey, + String modelName, + @Nullable String user + ) { + var model = createModel(url, org, apiKey, modelName, user); + + return new OpenAiEmbeddingsExecutableRequestCreator(model, TruncatorTests.createTruncator()); + } + + public static OpenAiEmbeddingsExecutableRequestCreator makeCreator( + String url, + @Nullable String org, + String apiKey, + String modelName, + @Nullable String user, + String inferenceEntityId + ) { + var model = createModel(url, org, apiKey, modelName, user, inferenceEntityId); + + return new OpenAiEmbeddingsExecutableRequestCreator(model, TruncatorTests.createTruncator()); + } +} diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/http/sender/RequestExecutorServiceTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/http/sender/RequestExecutorServiceTests.java index ef8731746e187..ebad28095294b 100644 --- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/http/sender/RequestExecutorServiceTests.java +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/http/sender/RequestExecutorServiceTests.java @@ -11,20 +11,19 @@ import org.elasticsearch.ElasticsearchTimeoutException; import org.elasticsearch.action.support.PlainActionFuture; import org.elasticsearch.common.Strings; -import org.elasticsearch.common.util.concurrent.AbstractRunnable; import org.elasticsearch.common.util.concurrent.EsRejectedExecutionException; import org.elasticsearch.core.Nullable; import org.elasticsearch.core.TimeValue; +import org.elasticsearch.inference.InferenceServiceResults; import org.elasticsearch.test.ESTestCase; import org.elasticsearch.threadpool.ThreadPool; -import org.elasticsearch.xpack.inference.external.http.HttpClient; -import org.elasticsearch.xpack.inference.external.http.HttpResult; -import org.elasticsearch.xpack.inference.external.request.HttpRequestTests; +import org.elasticsearch.xpack.inference.external.http.retry.RetryingHttpSender; import org.junit.After; import org.junit.Before; import org.mockito.ArgumentCaptor; import java.io.IOException; +import java.util.List; import java.util.concurrent.BlockingQueue; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; @@ -36,7 +35,6 @@ import static org.elasticsearch.core.Strings.format; import static org.elasticsearch.xpack.inference.Utils.inferenceUtilityPool; import static org.elasticsearch.xpack.inference.common.AdjustableCapacityBlockingQueueTests.mockQueueCreator; -import static org.elasticsearch.xpack.inference.external.http.HttpClientTests.createHttpPost; import static org.elasticsearch.xpack.inference.external.http.sender.RequestExecutorServiceSettingsTests.createRequestExecutorServiceSettings; import static org.elasticsearch.xpack.inference.external.http.sender.RequestExecutorServiceSettingsTests.createRequestExecutorServiceSettingsEmpty; import static org.hamcrest.Matchers.instanceOf; @@ -70,7 +68,7 @@ public void testQueueSize_IsEmpty() { public void testQueueSize_IsOne() { var service = createRequestExecutorServiceWithMocks(); - service.execute(HttpRequestTests.createMock("inferenceEntityId"), null, new PlainActionFuture<>()); + service.execute(ExecutableRequestCreatorTests.createMock(), List.of(), null, new PlainActionFuture<>()); assertThat(service.queueSize(), is(1)); } @@ -83,7 +81,7 @@ public void testIsTerminated_IsFalse() { public void testIsTerminated_IsTrue() throws InterruptedException { var latch = new CountDownLatch(1); - var service = createRequestExecutorService(null, latch); + var service = createRequestExecutorService(latch, mock(RetryingHttpSender.class)); service.shutdown(); service.start(); @@ -96,19 +94,24 @@ public void testIsTerminated_AfterStopFromSeparateThread() throws Exception { var waitToShutdown = new CountDownLatch(1); var waitToReturnFromSend = new CountDownLatch(1); - var mockHttpClient = mock(HttpClient.class); + var requestSender = mock(RetryingHttpSender.class); doAnswer(invocation -> { waitToShutdown.countDown(); waitToReturnFromSend.await(TIMEOUT.getSeconds(), TimeUnit.SECONDS); return Void.TYPE; - }).when(mockHttpClient).send(any(), any(), any()); + }).when(requestSender).send(any(), any(), any(), any(), any(), any()); - var service = createRequestExecutorService(mockHttpClient, null); + var service = createRequestExecutorService(null, requestSender); Future executorTermination = submitShutdownRequest(waitToShutdown, waitToReturnFromSend, service); - PlainActionFuture listener = new PlainActionFuture<>(); - service.execute(HttpRequestTests.createMock("inferenceEntityId"), null, listener); + PlainActionFuture listener = new PlainActionFuture<>(); + service.execute( + OpenAiEmbeddingsExecutableRequestCreatorTests.makeCreator("url", null, "key", "id", null), + List.of(), + null, + listener + ); service.start(); @@ -127,8 +130,8 @@ public void testSend_AfterShutdown_Throws() { service.shutdown(); - var listener = new PlainActionFuture(); - service.execute(HttpRequestTests.createMock("inferenceEntityId"), null, listener); + var listener = new PlainActionFuture(); + service.execute(ExecutableRequestCreatorTests.createMock(), List.of(), null, listener); var thrownException = expectThrows(EsRejectedExecutionException.class, () -> listener.actionGet(TIMEOUT)); @@ -142,15 +145,15 @@ public void testSend_AfterShutdown_Throws() { public void testSend_Throws_WhenQueueIsFull() { var service = new RequestExecutorService( "test_service", - mock(HttpClient.class), threadPool, null, - RequestExecutorServiceSettingsTests.createRequestExecutorServiceSettings(1) + RequestExecutorServiceSettingsTests.createRequestExecutorServiceSettings(1), + new SingleRequestManager(mock(RetryingHttpSender.class)) ); - service.execute(HttpRequestTests.createMock("inferenceEntityId"), null, new PlainActionFuture<>()); - var listener = new PlainActionFuture(); - service.execute(HttpRequestTests.createMock("inferenceEntityId"), null, listener); + service.execute(ExecutableRequestCreatorTests.createMock(), List.of(), null, new PlainActionFuture<>()); + var listener = new PlainActionFuture(); + service.execute(ExecutableRequestCreatorTests.createMock(), List.of(), null, listener); var thrownException = expectThrows(EsRejectedExecutionException.class, () -> listener.actionGet(TIMEOUT)); @@ -161,27 +164,28 @@ public void testSend_Throws_WhenQueueIsFull() { assertFalse(thrownException.isExecutorShutdown()); } - public void testTaskThrowsError_CallsOnFailure() throws Exception { - var httpClient = mock(HttpClient.class); + public void testTaskThrowsError_CallsOnFailure() { + var requestSender = mock(RetryingHttpSender.class); - var service = createRequestExecutorService(httpClient, null); + var service = createRequestExecutorService(null, requestSender); doAnswer(invocation -> { service.shutdown(); throw new IllegalArgumentException("failed"); - }).when(httpClient).send(any(), any(), any()); + }).when(requestSender).send(any(), any(), any(), any(), any(), any()); - PlainActionFuture listener = new PlainActionFuture<>(); + PlainActionFuture listener = new PlainActionFuture<>(); - var request = createHttpPost(0, "a", "b"); - service.execute(request, null, listener); + service.execute( + OpenAiEmbeddingsExecutableRequestCreatorTests.makeCreator("url", null, "key", "id", null), + List.of(), + null, + listener + ); service.start(); var thrownException = expectThrows(ElasticsearchException.class, () -> listener.actionGet(TIMEOUT)); - assertThat( - thrownException.getMessage(), - is(format("Failed to send request from inference entity id [%s]", request.inferenceEntityId())) - ); + assertThat(thrownException.getMessage(), is(format("Failed to send request from inference entity id [%s]", "id"))); assertThat(thrownException.getCause(), instanceOf(IllegalArgumentException.class)); assertTrue(service.isTerminated()); } @@ -200,22 +204,23 @@ public void testShutdown_AllowsMultipleCalls() { public void testSend_CallsOnFailure_WhenRequestTimesOut() { var service = createRequestExecutorServiceWithMocks(); - var listener = new PlainActionFuture(); - service.execute(HttpRequestTests.createMock("inferenceEntityId"), TimeValue.timeValueNanos(1), listener); + var listener = new PlainActionFuture(); + service.execute(ExecutableRequestCreatorTests.createMock(), List.of(), TimeValue.timeValueNanos(1), listener); var thrownException = expectThrows(ElasticsearchTimeoutException.class, () -> listener.actionGet(TIMEOUT)); assertThat( thrownException.getMessage(), - is(format("Request timed out waiting to be executed after [%s]", TimeValue.timeValueNanos(1))) + is(format("Request timed out waiting to be sent after [%s]", TimeValue.timeValueNanos(1))) ); } public void testSend_NotifiesTasksOfShutdown() { var service = createRequestExecutorServiceWithMocks(); - var listener = new PlainActionFuture(); - service.execute(HttpRequestTests.createMock("inferenceEntityId"), null, listener); + var listener = new PlainActionFuture(); + service.execute(ExecutableRequestCreatorTests.createMock(), List.of(), null, listener); + service.shutdown(); service.start(); @@ -231,15 +236,15 @@ public void testSend_NotifiesTasksOfShutdown() { public void testQueueTake_DoesNotCauseServiceToTerminate_WhenItThrows() throws InterruptedException { @SuppressWarnings("unchecked") - BlockingQueue queue = mock(LinkedBlockingQueue.class); + BlockingQueue queue = mock(LinkedBlockingQueue.class); var service = new RequestExecutorService( getTestName(), - mock(HttpClient.class), threadPool, mockQueueCreator(queue), null, - createRequestExecutorServiceSettingsEmpty() + createRequestExecutorServiceSettingsEmpty(), + new SingleRequestManager(mock(RetryingHttpSender.class)) ); when(queue.take()).thenThrow(new ElasticsearchException("failed")).thenAnswer(invocation -> { @@ -254,16 +259,16 @@ public void testQueueTake_DoesNotCauseServiceToTerminate_WhenItThrows() throws I public void testQueueTake_ThrowingInterruptedException_TerminatesService() throws Exception { @SuppressWarnings("unchecked") - BlockingQueue queue = mock(LinkedBlockingQueue.class); + BlockingQueue queue = mock(LinkedBlockingQueue.class); when(queue.take()).thenThrow(new InterruptedException("failed")); var service = new RequestExecutorService( getTestName(), - mock(HttpClient.class), threadPool, mockQueueCreator(queue), null, - createRequestExecutorServiceSettingsEmpty() + createRequestExecutorServiceSettingsEmpty(), + new SingleRequestManager(mock(RetryingHttpSender.class)) ); Future executorTermination = threadPool.generic().submit(() -> { @@ -281,17 +286,17 @@ public void testQueueTake_ThrowingInterruptedException_TerminatesService() throw } public void testQueueTake_RejectsTask_WhenServiceShutsDown() throws Exception { - var mockTask = mock(AbstractRunnable.class); + var mockTask = mock(RejectableTask.class); @SuppressWarnings("unchecked") - BlockingQueue queue = mock(LinkedBlockingQueue.class); + BlockingQueue queue = mock(LinkedBlockingQueue.class); var service = new RequestExecutorService( "test_service", - mock(HttpClient.class), threadPool, mockQueueCreator(queue), null, - createRequestExecutorServiceSettingsEmpty() + createRequestExecutorServiceSettingsEmpty(), + new SingleRequestManager(mock(RetryingHttpSender.class)) ); doAnswer(invocation -> { @@ -316,17 +321,17 @@ public void testQueueTake_RejectsTask_WhenServiceShutsDown() throws Exception { assertTrue(rejectionException.isExecutorShutdown()); } - public void testChangingCapacity_SetsCapacityToTwo() throws ExecutionException, InterruptedException, TimeoutException, IOException { - var httpClient = mock(HttpClient.class); + public void testChangingCapacity_SetsCapacityToTwo() throws ExecutionException, InterruptedException, TimeoutException { + var requestSender = mock(RetryingHttpSender.class); var settings = createRequestExecutorServiceSettings(1); - var service = new RequestExecutorService("test_service", httpClient, threadPool, null, settings); + var service = new RequestExecutorService("test_service", threadPool, null, settings, new SingleRequestManager(requestSender)); - service.execute(HttpRequestTests.createMock("inferenceEntityId"), null, new PlainActionFuture<>()); + service.execute(ExecutableRequestCreatorTests.createMock(requestSender), List.of(), null, new PlainActionFuture<>()); assertThat(service.queueSize(), is(1)); - PlainActionFuture listener = new PlainActionFuture<>(); - service.execute(HttpRequestTests.createMock("inferenceEntityId"), null, listener); + PlainActionFuture listener = new PlainActionFuture<>(); + service.execute(ExecutableRequestCreatorTests.createMock(requestSender), List.of(), null, listener); var thrownException = expectThrows(EsRejectedExecutionException.class, () -> listener.actionGet(TIMEOUT)); assertThat( @@ -343,7 +348,7 @@ public void testChangingCapacity_SetsCapacityToTwo() throws ExecutionException, waitToShutdown.countDown(); waitToReturnFromSend.await(TIMEOUT.getSeconds(), TimeUnit.SECONDS); return Void.TYPE; - }).when(httpClient).send(any(), any(), any()); + }).when(requestSender).send(any(), any(), any(), any(), any(), any()); Future executorTermination = submitShutdownRequest(waitToShutdown, waitToReturnFromSend, service); @@ -354,18 +359,18 @@ public void testChangingCapacity_SetsCapacityToTwo() throws ExecutionException, assertThat(service.remainingQueueCapacity(), is(2)); } - public void testChangingCapacity_DoesNotRejectsOverflowTasks_BecauseOfQueueFull() throws IOException, ExecutionException, - InterruptedException, TimeoutException { - var httpClient = mock(HttpClient.class); + public void testChangingCapacity_DoesNotRejectsOverflowTasks_BecauseOfQueueFull() throws ExecutionException, InterruptedException, + TimeoutException { + var requestSender = mock(RetryingHttpSender.class); var settings = createRequestExecutorServiceSettings(3); - var service = new RequestExecutorService("test_service", httpClient, threadPool, null, settings); + var service = new RequestExecutorService("test_service", threadPool, null, settings, new SingleRequestManager(requestSender)); - service.execute(HttpRequestTests.createMock("inferenceEntityId"), null, new PlainActionFuture<>()); - service.execute(HttpRequestTests.createMock("inferenceEntityId"), null, new PlainActionFuture<>()); + service.execute(ExecutableRequestCreatorTests.createMock(requestSender), List.of(), null, new PlainActionFuture<>()); + service.execute(ExecutableRequestCreatorTests.createMock(requestSender), List.of(), null, new PlainActionFuture<>()); - PlainActionFuture listener = new PlainActionFuture<>(); - service.execute(HttpRequestTests.createMock("inferenceEntityId"), null, listener); + PlainActionFuture listener = new PlainActionFuture<>(); + service.execute(ExecutableRequestCreatorTests.createMock(requestSender), List.of(), null, listener); assertThat(service.queueSize(), is(3)); settings.setQueueCapacity(1); @@ -377,7 +382,7 @@ public void testChangingCapacity_DoesNotRejectsOverflowTasks_BecauseOfQueueFull( waitToShutdown.countDown(); waitToReturnFromSend.await(TIMEOUT.getSeconds(), TimeUnit.SECONDS); return Void.TYPE; - }).when(httpClient).send(any(), any(), any()); + }).when(requestSender).send(any(), any(), any(), any(), any(), any()); Future executorTermination = submitShutdownRequest(waitToShutdown, waitToReturnFromSend, service); @@ -401,16 +406,16 @@ public void testChangingCapacity_DoesNotRejectsOverflowTasks_BecauseOfQueueFull( public void testChangingCapacity_ToZero_SetsQueueCapacityToUnbounded() throws IOException, ExecutionException, InterruptedException, TimeoutException { - var httpClient = mock(HttpClient.class); + var requestSender = mock(RetryingHttpSender.class); var settings = createRequestExecutorServiceSettings(1); - var service = new RequestExecutorService("test_service", httpClient, threadPool, null, settings); + var service = new RequestExecutorService("test_service", threadPool, null, settings, new SingleRequestManager(requestSender)); - service.execute(HttpRequestTests.createMock("inferenceEntityId"), null, new PlainActionFuture<>()); + service.execute(ExecutableRequestCreatorTests.createMock(requestSender), List.of(), null, new PlainActionFuture<>()); assertThat(service.queueSize(), is(1)); - PlainActionFuture listener = new PlainActionFuture<>(); - service.execute(HttpRequestTests.createMock("inferenceEntityId"), null, listener); + PlainActionFuture listener = new PlainActionFuture<>(); + service.execute(ExecutableRequestCreatorTests.createMock(requestSender), List.of(), null, listener); var thrownException = expectThrows(EsRejectedExecutionException.class, () -> listener.actionGet(TIMEOUT)); assertThat( @@ -427,7 +432,7 @@ public void testChangingCapacity_ToZero_SetsQueueCapacityToUnbounded() throws IO waitToShutdown.countDown(); waitToReturnFromSend.await(TIMEOUT.getSeconds(), TimeUnit.SECONDS); return Void.TYPE; - }).when(httpClient).send(any(), any(), any()); + }).when(requestSender).send(any(), any(), any(), any(), any(), any()); Future executorTermination = submitShutdownRequest(waitToShutdown, waitToReturnFromSend, service); @@ -458,17 +463,16 @@ private Future submitShutdownRequest( } private RequestExecutorService createRequestExecutorServiceWithMocks() { - return createRequestExecutorService(null, null); + return createRequestExecutorService(null, mock(RetryingHttpSender.class)); } - private RequestExecutorService createRequestExecutorService(@Nullable HttpClient httpClient, @Nullable CountDownLatch startupLatch) { - var httpClientToUse = httpClient == null ? mock(HttpClient.class) : httpClient; + private RequestExecutorService createRequestExecutorService(@Nullable CountDownLatch startupLatch, RetryingHttpSender requestSender) { return new RequestExecutorService( "test_service", - httpClientToUse, threadPool, startupLatch, - createRequestExecutorServiceSettingsEmpty() + createRequestExecutorServiceSettingsEmpty(), + new SingleRequestManager(requestSender) ); } } diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/http/sender/RequestTaskTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/http/sender/RequestTaskTests.java index eaf1a0ac267cf..5c35d8ce49b60 100644 --- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/http/sender/RequestTaskTests.java +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/http/sender/RequestTaskTests.java @@ -7,30 +7,19 @@ package org.elasticsearch.xpack.inference.external.http.sender; -import org.apache.http.HttpHeaders; -import org.apache.http.HttpResponse; -import org.apache.http.client.protocol.HttpClientContext; -import org.elasticsearch.ElasticsearchException; import org.elasticsearch.ElasticsearchTimeoutException; import org.elasticsearch.action.ActionListener; import org.elasticsearch.action.support.PlainActionFuture; import org.elasticsearch.core.TimeValue; +import org.elasticsearch.inference.InferenceServiceResults; import org.elasticsearch.test.ESTestCase; -import org.elasticsearch.test.http.MockResponse; -import org.elasticsearch.test.http.MockWebServer; import org.elasticsearch.threadpool.Scheduler; import org.elasticsearch.threadpool.ThreadPool; -import org.elasticsearch.xcontent.XContentType; -import org.elasticsearch.xpack.inference.external.http.HttpClient; -import org.elasticsearch.xpack.inference.external.http.HttpResult; -import org.elasticsearch.xpack.inference.external.request.HttpRequestTests; -import org.elasticsearch.xpack.inference.logging.ThrottlerManager; import org.junit.After; import org.junit.Before; import org.mockito.ArgumentCaptor; -import java.io.IOException; -import java.nio.charset.StandardCharsets; +import java.util.List; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; import java.util.concurrent.TimeUnit; @@ -38,16 +27,9 @@ import static org.elasticsearch.core.Strings.format; import static org.elasticsearch.xpack.inference.Utils.inferenceUtilityPool; -import static org.elasticsearch.xpack.inference.external.http.HttpClientTests.createConnectionManager; -import static org.elasticsearch.xpack.inference.external.http.HttpClientTests.createHttpPost; -import static org.elasticsearch.xpack.inference.external.http.HttpClientTests.emptyHttpSettings; -import static org.hamcrest.Matchers.equalTo; -import static org.hamcrest.Matchers.hasSize; -import static org.hamcrest.Matchers.instanceOf; import static org.hamcrest.Matchers.is; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.doAnswer; -import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; @@ -57,134 +39,65 @@ public class RequestTaskTests extends ESTestCase { private static final TimeValue TIMEOUT = new TimeValue(30, TimeUnit.SECONDS); - private final MockWebServer webServer = new MockWebServer(); private ThreadPool threadPool; @Before public void init() throws Exception { - webServer.start(); threadPool = createThreadPool(inferenceUtilityPool()); } @After public void shutdown() { terminate(threadPool); - webServer.close(); } - public void testDoRun_SendsRequestAndReceivesResponse() throws Exception { - int responseCode = randomIntBetween(200, 203); - String body = randomAlphaOfLengthBetween(2, 8096); - webServer.enqueue(new MockResponse().setResponseCode(responseCode).setBody(body)); - - String paramKey = randomAlphaOfLength(3); - String paramValue = randomAlphaOfLength(3); - var httpPost = createHttpPost(webServer.getPort(), paramKey, paramValue); - - try (var httpClient = HttpClient.create(emptyHttpSettings(), threadPool, createConnectionManager(), mock(ThrottlerManager.class))) { - httpClient.start(); - - PlainActionFuture listener = new PlainActionFuture<>(); - var requestTask = new RequestTask(httpPost, httpClient, HttpClientContext.create(), null, threadPool, listener); - requestTask.doRun(); - var result = listener.actionGet(TIMEOUT); - - assertThat(result.response().getStatusLine().getStatusCode(), equalTo(responseCode)); - assertThat(new String(result.body(), StandardCharsets.UTF_8), is(body)); - assertThat(webServer.requests(), hasSize(1)); - assertThat(webServer.requests().get(0).getUri().getPath(), equalTo(httpPost.httpRequestBase().getURI().getPath())); - assertThat(webServer.requests().get(0).getUri().getQuery(), equalTo(paramKey + "=" + paramValue)); - assertThat(webServer.requests().get(0).getHeader(HttpHeaders.CONTENT_TYPE), equalTo(XContentType.JSON.mediaType())); - } - } - - public void testDoRun_SendThrowsIOException() throws Exception { - var httpClient = mock(HttpClient.class); - doThrow(new IOException("exception")).when(httpClient).send(any(), any(), any()); - - String paramKey = randomAlphaOfLength(3); - String paramValue = randomAlphaOfLength(3); - var httpPost = createHttpPost(webServer.getPort(), paramKey, paramValue); - - PlainActionFuture listener = new PlainActionFuture<>(); - var requestTask = new RequestTask(httpPost, httpClient, HttpClientContext.create(), null, threadPool, listener); - requestTask.doRun(); - - var thrownException = expectThrows(ElasticsearchException.class, () -> listener.actionGet(TIMEOUT)); - assertThat( - thrownException.getMessage(), - is(format("Failed to send request from inference entity id [%s]", httpPost.inferenceEntityId())) - ); - } - - public void testRequest_DoesNotCallOnFailureForTimeout_AfterSendThrowsIllegalArgumentException() throws Exception { + public void testExecuting_DoesNotCallOnFailureForTimeout_AfterIllegalArgumentException() { AtomicReference onTimeout = new AtomicReference<>(); var mockThreadPool = mockThreadPoolForTimeout(onTimeout); - var httpClient = mock(HttpClient.class); - doThrow(new IllegalArgumentException("failed")).when(httpClient).send(any(), any(), any()); - - var httpPost = createHttpPost(webServer.getPort(), "a", "b"); - @SuppressWarnings("unchecked") - ActionListener listener = mock(ActionListener.class); + ActionListener listener = mock(ActionListener.class); var requestTask = new RequestTask( - httpPost, - httpClient, - HttpClientContext.create(), + OpenAiEmbeddingsExecutableRequestCreatorTests.makeCreator("url", null, "key", "model", null, "id"), + List.of("abc"), TimeValue.timeValueMillis(1), mockThreadPool, listener ); - requestTask.doRun(); - - ArgumentCaptor argument = ArgumentCaptor.forClass(Exception.class); - verify(listener, times(1)).onFailure(argument.capture()); - assertThat( - argument.getValue().getMessage(), - is(format("Failed to send request from inference entity id [%s]", httpPost.inferenceEntityId())) - ); - assertThat(argument.getValue(), instanceOf(ElasticsearchException.class)); - assertThat(argument.getValue().getCause(), instanceOf(IllegalArgumentException.class)); + requestTask.getListener().onFailure(new IllegalArgumentException("failed")); + verify(listener, times(1)).onFailure(any()); + assertTrue(requestTask.hasCompleted()); + assertTrue(requestTask.getRequestCompletedFunction().get()); onTimeout.get().run(); verifyNoMoreInteractions(listener); } public void testRequest_ReturnsTimeoutException() { - var httpClient = mock(HttpClient.class); - PlainActionFuture listener = new PlainActionFuture<>(); + PlainActionFuture listener = new PlainActionFuture<>(); var requestTask = new RequestTask( - HttpRequestTests.createMock("inferenceEntityId"), - httpClient, - HttpClientContext.create(), + OpenAiEmbeddingsExecutableRequestCreatorTests.makeCreator("url", null, "key", "model", null, "id"), + List.of("abc"), TimeValue.timeValueMillis(1), threadPool, listener ); - requestTask.doRun(); var thrownException = expectThrows(ElasticsearchTimeoutException.class, () -> listener.actionGet(TIMEOUT)); assertThat( thrownException.getMessage(), - is(format("Request timed out waiting to be executed after [%s]", TimeValue.timeValueMillis(1))) + is(format("Request timed out waiting to be sent after [%s]", TimeValue.timeValueMillis(1))) ); + assertTrue(requestTask.hasCompleted()); + assertTrue(requestTask.getRequestCompletedFunction().get()); } public void testRequest_DoesNotCallOnFailureTwiceWhenTimingOut() throws Exception { - var httpClient = mock(HttpClient.class); - doAnswer(invocation -> { - @SuppressWarnings("unchecked") - ActionListener listener = (ActionListener) invocation.getArguments()[2]; - listener.onFailure(new ElasticsearchException("failed")); - return Void.TYPE; - }).when(httpClient).send(any(), any(), any()); - @SuppressWarnings("unchecked") - ActionListener listener = mock(ActionListener.class); + ActionListener listener = mock(ActionListener.class); var calledOnFailureLatch = new CountDownLatch(1); doAnswer(invocation -> { calledOnFailureLatch.countDown(); @@ -192,9 +105,8 @@ public void testRequest_DoesNotCallOnFailureTwiceWhenTimingOut() throws Exceptio }).when(listener).onFailure(any()); var requestTask = new RequestTask( - HttpRequestTests.createMock("inferenceEntityId"), - httpClient, - HttpClientContext.create(), + OpenAiEmbeddingsExecutableRequestCreatorTests.makeCreator("url", null, "key", "model", null, "id"), + List.of("abc"), TimeValue.timeValueMillis(1), threadPool, listener @@ -206,25 +118,18 @@ public void testRequest_DoesNotCallOnFailureTwiceWhenTimingOut() throws Exceptio verify(listener, times(1)).onFailure(argument.capture()); assertThat( argument.getValue().getMessage(), - is(format("Request timed out waiting to be executed after [%s]", TimeValue.timeValueMillis(1))) + is(format("Request timed out waiting to be sent after [%s]", TimeValue.timeValueMillis(1))) ); + assertTrue(requestTask.hasCompleted()); + assertTrue(requestTask.getRequestCompletedFunction().get()); - requestTask.doRun(); + requestTask.getListener().onFailure(new IllegalArgumentException("failed")); verifyNoMoreInteractions(listener); } public void testRequest_DoesNotCallOnResponseAfterTimingOut() throws Exception { - var httpClient = mock(HttpClient.class); - doAnswer(invocation -> { - @SuppressWarnings("unchecked") - ActionListener listener = (ActionListener) invocation.getArguments()[2]; - var result = new HttpResult(mock(HttpResponse.class), new byte[0]); - listener.onResponse(result); - return Void.TYPE; - }).when(httpClient).send(any(), any(), any()); - @SuppressWarnings("unchecked") - ActionListener listener = mock(ActionListener.class); + ActionListener listener = mock(ActionListener.class); var calledOnFailureLatch = new CountDownLatch(1); doAnswer(invocation -> { calledOnFailureLatch.countDown(); @@ -232,9 +137,8 @@ public void testRequest_DoesNotCallOnResponseAfterTimingOut() throws Exception { }).when(listener).onFailure(any()); var requestTask = new RequestTask( - HttpRequestTests.createMock("inferenceEntityId"), - httpClient, - HttpClientContext.create(), + OpenAiEmbeddingsExecutableRequestCreatorTests.makeCreator("url", null, "key", "model", null, "id"), + List.of("abc"), TimeValue.timeValueMillis(1), threadPool, listener @@ -246,44 +150,12 @@ public void testRequest_DoesNotCallOnResponseAfterTimingOut() throws Exception { verify(listener, times(1)).onFailure(argument.capture()); assertThat( argument.getValue().getMessage(), - is(format("Request timed out waiting to be executed after [%s]", TimeValue.timeValueMillis(1))) - ); - - requestTask.doRun(); - verifyNoMoreInteractions(listener); - } - - public void testRequest_DoesNotCallOnFailureForTimeout_AfterAlreadyCallingOnFailure() throws Exception { - AtomicReference onTimeout = new AtomicReference<>(); - var mockThreadPool = mockThreadPoolForTimeout(onTimeout); - - var httpClient = mock(HttpClient.class); - doAnswer(invocation -> { - @SuppressWarnings("unchecked") - ActionListener listener = (ActionListener) invocation.getArguments()[2]; - listener.onFailure(new ElasticsearchException("failed")); - return Void.TYPE; - }).when(httpClient).send(any(), any(), any()); - - @SuppressWarnings("unchecked") - ActionListener listener = mock(ActionListener.class); - - var requestTask = new RequestTask( - HttpRequestTests.createMock("inferenceEntityId"), - httpClient, - HttpClientContext.create(), - TimeValue.timeValueMillis(1), - mockThreadPool, - listener + is(format("Request timed out waiting to be sent after [%s]", TimeValue.timeValueMillis(1))) ); + assertTrue(requestTask.hasCompleted()); + assertTrue(requestTask.getRequestCompletedFunction().get()); - requestTask.doRun(); - - ArgumentCaptor argument = ArgumentCaptor.forClass(Exception.class); - verify(listener, times(1)).onFailure(argument.capture()); - assertThat(argument.getValue().getMessage(), is("failed")); - - onTimeout.get().run(); + requestTask.getListener().onResponse(mock(InferenceServiceResults.class)); verifyNoMoreInteractions(listener); } @@ -291,29 +163,21 @@ public void testRequest_DoesNotCallOnFailureForTimeout_AfterAlreadyCallingOnResp AtomicReference onTimeout = new AtomicReference<>(); var mockThreadPool = mockThreadPoolForTimeout(onTimeout); - var httpClient = mock(HttpClient.class); - doAnswer(invocation -> { - @SuppressWarnings("unchecked") - ActionListener listener = (ActionListener) invocation.getArguments()[2]; - listener.onResponse(new HttpResult(mock(HttpResponse.class), new byte[0])); - return Void.TYPE; - }).when(httpClient).send(any(), any(), any()); - @SuppressWarnings("unchecked") - ActionListener listener = mock(ActionListener.class); + ActionListener listener = mock(ActionListener.class); var requestTask = new RequestTask( - HttpRequestTests.createMock("inferenceEntityId"), - httpClient, - HttpClientContext.create(), + OpenAiEmbeddingsExecutableRequestCreatorTests.makeCreator("url", null, "key", "model", null, "id"), + List.of("abc"), TimeValue.timeValueMillis(1), mockThreadPool, listener ); - requestTask.doRun(); - + requestTask.getListener().onResponse(mock(InferenceServiceResults.class)); verify(listener, times(1)).onResponse(any()); + assertTrue(requestTask.hasCompleted()); + assertTrue(requestTask.getRequestCompletedFunction().get()); onTimeout.get().run(); verifyNoMoreInteractions(listener); diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/http/sender/SingleRequestManagerTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/http/sender/SingleRequestManagerTests.java new file mode 100644 index 0000000000000..ab8bf244a4d2c --- /dev/null +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/http/sender/SingleRequestManagerTests.java @@ -0,0 +1,27 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.inference.external.http.sender; + +import org.apache.http.client.protocol.HttpClientContext; +import org.elasticsearch.test.ESTestCase; +import org.elasticsearch.xpack.inference.external.http.retry.RetryingHttpSender; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verifyNoInteractions; +import static org.mockito.Mockito.when; + +public class SingleRequestManagerTests extends ESTestCase { + public void testExecute_DoesNotCallRequestCreatorCreate_WhenInputIsNull() { + var requestCreator = mock(ExecutableRequestCreator.class); + var request = mock(InferenceRequest.class); + when(request.getRequestCreator()).thenReturn(requestCreator); + + new SingleRequestManager(mock(RetryingHttpSender.class)).execute(mock(InferenceRequest.class), HttpClientContext.create()); + verifyNoInteractions(requestCreator); + } +} diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/openai/OpenAiClientTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/openai/OpenAiClientTests.java deleted file mode 100644 index bb9612f01d8ff..0000000000000 --- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/openai/OpenAiClientTests.java +++ /dev/null @@ -1,297 +0,0 @@ -/* - * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one - * or more contributor license agreements. Licensed under the Elastic License - * 2.0; you may not use this file except in compliance with the Elastic License - * 2.0. - */ - -package org.elasticsearch.xpack.inference.external.openai; - -import org.apache.http.HttpHeaders; -import org.elasticsearch.ElasticsearchException; -import org.elasticsearch.action.support.PlainActionFuture; -import org.elasticsearch.common.settings.Settings; -import org.elasticsearch.core.TimeValue; -import org.elasticsearch.inference.InferenceServiceResults; -import org.elasticsearch.test.ESTestCase; -import org.elasticsearch.test.http.MockResponse; -import org.elasticsearch.test.http.MockWebServer; -import org.elasticsearch.threadpool.ThreadPool; -import org.elasticsearch.xcontent.XContentType; -import org.elasticsearch.xpack.inference.common.TruncatorTests; -import org.elasticsearch.xpack.inference.external.http.HttpClientManager; -import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSenderFactory; -import org.elasticsearch.xpack.inference.external.http.sender.Sender; -import org.elasticsearch.xpack.inference.services.ServiceComponents; -import org.junit.After; -import org.junit.Before; - -import java.io.IOException; -import java.net.URISyntaxException; -import java.util.List; -import java.util.concurrent.TimeUnit; - -import static org.elasticsearch.core.Strings.format; -import static org.elasticsearch.xpack.inference.Utils.inferenceUtilityPool; -import static org.elasticsearch.xpack.inference.Utils.mockClusterServiceEmpty; -import static org.elasticsearch.xpack.inference.external.http.Utils.entityAsMap; -import static org.elasticsearch.xpack.inference.external.http.Utils.getUrl; -import static org.elasticsearch.xpack.inference.external.http.retry.RetrySettingsTests.buildSettingsWithRetryFields; -import static org.elasticsearch.xpack.inference.external.request.openai.OpenAiEmbeddingsRequestTests.createRequest; -import static org.elasticsearch.xpack.inference.external.request.openai.OpenAiUtils.ORGANIZATION_HEADER; -import static org.elasticsearch.xpack.inference.logging.ThrottlerManagerTests.mockThrottlerManager; -import static org.elasticsearch.xpack.inference.results.TextEmbeddingResultsTests.buildExpectation; -import static org.elasticsearch.xpack.inference.services.ServiceComponentsTests.createWithEmptySettings; -import static org.hamcrest.Matchers.equalTo; -import static org.hamcrest.Matchers.hasSize; -import static org.hamcrest.Matchers.is; -import static org.mockito.ArgumentMatchers.any; -import static org.mockito.Mockito.doThrow; -import static org.mockito.Mockito.mock; - -public class OpenAiClientTests extends ESTestCase { - private static final TimeValue TIMEOUT = new TimeValue(30, TimeUnit.SECONDS); - private final MockWebServer webServer = new MockWebServer(); - private ThreadPool threadPool; - private HttpClientManager clientManager; - - @Before - public void init() throws Exception { - webServer.start(); - threadPool = createThreadPool(inferenceUtilityPool()); - clientManager = HttpClientManager.create(Settings.EMPTY, threadPool, mockClusterServiceEmpty(), mockThrottlerManager()); - } - - @After - public void shutdown() throws IOException { - clientManager.close(); - terminate(threadPool); - webServer.close(); - } - - public void testSend_SuccessfulResponse() throws IOException, URISyntaxException { - var senderFactory = new HttpRequestSenderFactory(threadPool, clientManager, mockClusterServiceEmpty(), Settings.EMPTY); - - try (var sender = senderFactory.createSender("test_service")) { - sender.start(); - - String responseJson = """ - { - "object": "list", - "data": [ - { - "object": "embedding", - "index": 0, - "embedding": [ - 0.0123, - -0.0123 - ] - } - ], - "model": "text-embedding-ada-002-v2", - "usage": { - "prompt_tokens": 8, - "total_tokens": 8 - } - } - """; - webServer.enqueue(new MockResponse().setResponseCode(200).setBody(responseJson)); - - OpenAiClient openAiClient = new OpenAiClient(sender, createWithEmptySettings(threadPool)); - - PlainActionFuture listener = new PlainActionFuture<>(); - openAiClient.send(createRequest(getUrl(webServer), "org", "secret", "abc", "model", "user"), listener); - - var result = listener.actionGet(TIMEOUT); - - assertThat(result.asMap(), is(buildExpectation(List.of(List.of(0.0123F, -0.0123F))))); - - assertThat(webServer.requests(), hasSize(1)); - assertNull(webServer.requests().get(0).getUri().getQuery()); - assertThat(webServer.requests().get(0).getHeader(HttpHeaders.CONTENT_TYPE), equalTo(XContentType.JSON.mediaType())); - assertThat(webServer.requests().get(0).getHeader(HttpHeaders.AUTHORIZATION), equalTo("Bearer secret")); - assertThat(webServer.requests().get(0).getHeader(ORGANIZATION_HEADER), equalTo("org")); - - var requestMap = entityAsMap(webServer.requests().get(0).getBody()); - assertThat(requestMap.size(), is(3)); - assertThat(requestMap.get("input"), is(List.of("abc"))); - assertThat(requestMap.get("model"), is("model")); - assertThat(requestMap.get("user"), is("user")); - } - } - - public void testSend_SuccessfulResponse_WithoutUser() throws IOException, URISyntaxException { - var senderFactory = new HttpRequestSenderFactory(threadPool, clientManager, mockClusterServiceEmpty(), Settings.EMPTY); - - try (var sender = senderFactory.createSender("test_service")) { - sender.start(); - - String responseJson = """ - { - "object": "list", - "data": [ - { - "object": "embedding", - "index": 0, - "embedding": [ - 0.0123, - -0.0123 - ] - } - ], - "model": "text-embedding-ada-002-v2", - "usage": { - "prompt_tokens": 8, - "total_tokens": 8 - } - } - """; - webServer.enqueue(new MockResponse().setResponseCode(200).setBody(responseJson)); - - OpenAiClient openAiClient = new OpenAiClient(sender, createWithEmptySettings(threadPool)); - - PlainActionFuture listener = new PlainActionFuture<>(); - openAiClient.send(createRequest(getUrl(webServer), "org", "secret", "abc", "model", null), listener); - - var result = listener.actionGet(TIMEOUT); - - assertThat(result.asMap(), is(buildExpectation(List.of(List.of(0.0123F, -0.0123F))))); - - assertThat(webServer.requests(), hasSize(1)); - assertNull(webServer.requests().get(0).getUri().getQuery()); - assertThat(webServer.requests().get(0).getHeader(HttpHeaders.CONTENT_TYPE), equalTo(XContentType.JSON.mediaType())); - assertThat(webServer.requests().get(0).getHeader(HttpHeaders.AUTHORIZATION), equalTo("Bearer secret")); - assertThat(webServer.requests().get(0).getHeader(ORGANIZATION_HEADER), equalTo("org")); - - var requestMap = entityAsMap(webServer.requests().get(0).getBody()); - assertThat(requestMap.size(), is(2)); - assertThat(requestMap.get("input"), is(List.of("abc"))); - assertThat(requestMap.get("model"), is("model")); - } - } - - public void testSend_SuccessfulResponse_WithoutOrganization() throws IOException, URISyntaxException { - var senderFactory = new HttpRequestSenderFactory(threadPool, clientManager, mockClusterServiceEmpty(), Settings.EMPTY); - - try (var sender = senderFactory.createSender("test_service")) { - sender.start(); - - String responseJson = """ - { - "object": "list", - "data": [ - { - "object": "embedding", - "index": 0, - "embedding": [ - 0.0123, - -0.0123 - ] - } - ], - "model": "text-embedding-ada-002-v2", - "usage": { - "prompt_tokens": 8, - "total_tokens": 8 - } - } - """; - webServer.enqueue(new MockResponse().setResponseCode(200).setBody(responseJson)); - - OpenAiClient openAiClient = new OpenAiClient(sender, createWithEmptySettings(threadPool)); - - PlainActionFuture listener = new PlainActionFuture<>(); - openAiClient.send(createRequest(getUrl(webServer), null, "secret", "abc", "model", null), listener); - - var result = listener.actionGet(TIMEOUT); - - assertThat(result.asMap(), is(buildExpectation(List.of(List.of(0.0123F, -0.0123F))))); - - assertThat(webServer.requests(), hasSize(1)); - assertNull(webServer.requests().get(0).getUri().getQuery()); - assertThat(webServer.requests().get(0).getHeader(HttpHeaders.CONTENT_TYPE), equalTo(XContentType.JSON.mediaType())); - assertThat(webServer.requests().get(0).getHeader(HttpHeaders.AUTHORIZATION), equalTo("Bearer secret")); - assertNull(webServer.requests().get(0).getHeader(ORGANIZATION_HEADER)); - - var requestMap = entityAsMap(webServer.requests().get(0).getBody()); - assertThat(requestMap.size(), is(2)); - assertThat(requestMap.get("input"), is(List.of("abc"))); - assertThat(requestMap.get("model"), is("model")); - } - } - - public void testSend_FailsFromInvalidResponseFormat() throws IOException, URISyntaxException { - var senderFactory = new HttpRequestSenderFactory(threadPool, clientManager, mockClusterServiceEmpty(), Settings.EMPTY); - - try (var sender = senderFactory.createSender("test_service")) { - sender.start(); - - String responseJson = """ - { - "object": "list", - "data_does_not_exist": [ - { - "object": "embedding", - "index": 0, - "embedding": [ - 0.0123, - -0.0123 - ] - } - ], - "model": "text-embedding-ada-002-v2", - "usage": { - "prompt_tokens": 8, - "total_tokens": 8 - } - } - """; - webServer.enqueue(new MockResponse().setResponseCode(200).setBody(responseJson)); - - OpenAiClient openAiClient = new OpenAiClient( - sender, - new ServiceComponents( - threadPool, - mockThrottlerManager(), - // timeout as zero for no retries - buildSettingsWithRetryFields( - TimeValue.timeValueMillis(1), - TimeValue.timeValueMinutes(1), - TimeValue.timeValueSeconds(0) - ), - TruncatorTests.createTruncator() - ) - ); - - PlainActionFuture listener = new PlainActionFuture<>(); - openAiClient.send(createRequest(getUrl(webServer), "org", "secret", "abc", "model", "user"), listener); - - var thrownException = expectThrows(IllegalStateException.class, () -> listener.actionGet(TIMEOUT)); - assertThat(thrownException.getMessage(), is(format("Failed to find required field [data] in OpenAI embeddings response"))); - - assertThat(webServer.requests(), hasSize(1)); - assertNull(webServer.requests().get(0).getUri().getQuery()); - assertThat(webServer.requests().get(0).getHeader(HttpHeaders.CONTENT_TYPE), equalTo(XContentType.JSON.mediaType())); - assertThat(webServer.requests().get(0).getHeader(HttpHeaders.AUTHORIZATION), equalTo("Bearer secret")); - assertThat(webServer.requests().get(0).getHeader(ORGANIZATION_HEADER), equalTo("org")); - - var requestMap = entityAsMap(webServer.requests().get(0).getBody()); - assertThat(requestMap.size(), is(3)); - assertThat(requestMap.get("input"), is(List.of("abc"))); - assertThat(requestMap.get("model"), is("model")); - assertThat(requestMap.get("user"), is("user")); - } - } - - public void testSend_ThrowsException() throws URISyntaxException, IOException { - var sender = mock(Sender.class); - doThrow(new ElasticsearchException("failed")).when(sender).send(any(), any()); - - OpenAiClient openAiClient = new OpenAiClient(sender, createWithEmptySettings(threadPool)); - PlainActionFuture listener = new PlainActionFuture<>(); - openAiClient.send(createRequest(getUrl(webServer), "org", "secret", "abc", "model", "user"), listener); - - var thrownException = expectThrows(ElasticsearchException.class, () -> listener.actionGet(TIMEOUT)); - assertThat(thrownException.getMessage(), is("failed")); - } -} diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/request/openai/OpenAiEmbeddingsRequestTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/request/openai/OpenAiEmbeddingsRequestTests.java index 4c4c40e9c1056..ebff1c5e096e8 100644 --- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/request/openai/OpenAiEmbeddingsRequestTests.java +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/request/openai/OpenAiEmbeddingsRequestTests.java @@ -116,7 +116,7 @@ public static OpenAiEmbeddingsRequest createRequest( String model, @Nullable String user ) { - var embeddingsModel = OpenAiEmbeddingsModelTests.createModel(url, org, apiKey, model, user, null); + var embeddingsModel = OpenAiEmbeddingsModelTests.createModel(url, org, apiKey, model, user, (Integer) null); var account = new OpenAiAccount(embeddingsModel.getServiceSettings().uri(), org, embeddingsModel.getSecretSettings().apiKey()); return new OpenAiEmbeddingsRequest( diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/SenderServiceTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/SenderServiceTests.java index 82d53cfb09037..5c438644a18c5 100644 --- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/SenderServiceTests.java +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/SenderServiceTests.java @@ -19,7 +19,7 @@ import org.elasticsearch.inference.TaskType; import org.elasticsearch.test.ESTestCase; import org.elasticsearch.threadpool.ThreadPool; -import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSenderFactory; +import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSender; import org.elasticsearch.xpack.inference.external.http.sender.Sender; import org.junit.After; import org.junit.Before; @@ -57,7 +57,7 @@ public void shutdown() throws IOException { public void testStart_InitializesTheSender() throws IOException { var sender = mock(Sender.class); - var factory = mock(HttpRequestSenderFactory.class); + var factory = mock(HttpRequestSender.Factory.class); when(factory.createSender(anyString())).thenReturn(sender); try (var service = new TestSenderService(factory, createWithEmptySettings(threadPool))) { @@ -77,7 +77,7 @@ public void testStart_InitializesTheSender() throws IOException { public void testStart_CallingStartTwiceKeepsSameSenderReference() throws IOException { var sender = mock(Sender.class); - var factory = mock(HttpRequestSenderFactory.class); + var factory = mock(HttpRequestSender.Factory.class); when(factory.createSender(anyString())).thenReturn(sender); try (var service = new TestSenderService(factory, createWithEmptySettings(threadPool))) { @@ -98,7 +98,7 @@ public void testStart_CallingStartTwiceKeepsSameSenderReference() throws IOExcep } private static final class TestSenderService extends SenderService { - TestSenderService(HttpRequestSenderFactory factory, ServiceComponents serviceComponents) { + TestSenderService(HttpRequestSender.Factory factory, ServiceComponents serviceComponents) { super(factory, serviceComponents); } diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/ServiceComponentsTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/ServiceComponentsTests.java index 77713fbfc30a5..fd568bf7f15da 100644 --- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/ServiceComponentsTests.java +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/ServiceComponentsTests.java @@ -18,4 +18,8 @@ public class ServiceComponentsTests extends ESTestCase { public static ServiceComponents createWithEmptySettings(ThreadPool threadPool) { return new ServiceComponents(threadPool, mockThrottlerManager(), Settings.EMPTY, TruncatorTests.createTruncator()); } + + public static ServiceComponents createWithSettings(ThreadPool threadPool, Settings settings) { + return new ServiceComponents(threadPool, mockThrottlerManager(), settings, TruncatorTests.createTruncator()); + } } diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/cohere/CohereServiceTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/cohere/CohereServiceTests.java index 9c2722e68efd6..356da0ece08af 100644 --- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/cohere/CohereServiceTests.java +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/cohere/CohereServiceTests.java @@ -28,7 +28,8 @@ import org.elasticsearch.threadpool.ThreadPool; import org.elasticsearch.xcontent.XContentType; import org.elasticsearch.xpack.inference.external.http.HttpClientManager; -import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSenderFactory; +import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSender; +import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSenderTests; import org.elasticsearch.xpack.inference.external.http.sender.Sender; import org.elasticsearch.xpack.inference.logging.ThrottlerManager; import org.elasticsearch.xpack.inference.services.cohere.embeddings.CohereEmbeddingType; @@ -93,7 +94,6 @@ public void shutdown() throws IOException { public void testParseRequestConfig_CreatesACohereEmbeddingsModel() throws IOException { try (var service = createCohereService()) { - ActionListener modelListener = ActionListener.wrap(model -> { MatcherAssert.assertThat(model, instanceOf(CohereEmbeddingsModel.class)); @@ -125,7 +125,6 @@ public void testParseRequestConfig_CreatesACohereEmbeddingsModel() throws IOExce public void testParseRequestConfig_ThrowsUnsupportedModelType() throws IOException { try (var service = createCohereService()) { - var failureListener = getModelListenerForException( ElasticsearchStatusException.class, "The [cohere] service does not support task type [sparse_embedding]" @@ -577,7 +576,7 @@ public void testParsePersistedConfig_NotThrowWhenAnExtraKeyExistsInTaskSettings( public void testInfer_ThrowsErrorWhenModelIsNotCohereModel() throws IOException { var sender = mock(Sender.class); - var factory = mock(HttpRequestSenderFactory.class); + var factory = mock(HttpRequestSender.Factory.class); when(factory.createSender(anyString())).thenReturn(sender); var mockModel = getInvalidModel("model_id", "service_name"); @@ -602,7 +601,7 @@ public void testInfer_ThrowsErrorWhenModelIsNotCohereModel() throws IOException } public void testInfer_SendsRequest() throws IOException { - var senderFactory = new HttpRequestSenderFactory(threadPool, clientManager, mockClusterServiceEmpty(), Settings.EMPTY); + var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager); try (var service = new CohereService(senderFactory, createWithEmptySettings(threadPool))) { @@ -662,7 +661,7 @@ public void testInfer_SendsRequest() throws IOException { } public void testCheckModelConfig_UpdatesDimensions() throws IOException { - var senderFactory = new HttpRequestSenderFactory(threadPool, clientManager, mockClusterServiceEmpty(), Settings.EMPTY); + var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager); try (var service = new CohereService(senderFactory, createWithEmptySettings(threadPool))) { @@ -725,7 +724,7 @@ public void testCheckModelConfig_UpdatesDimensions() throws IOException { } public void testInfer_UnauthorisedResponse() throws IOException { - var senderFactory = new HttpRequestSenderFactory(threadPool, clientManager, mockClusterServiceEmpty(), Settings.EMPTY); + var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager); try (var service = new CohereService(senderFactory, createWithEmptySettings(threadPool))) { @@ -756,7 +755,7 @@ public void testInfer_UnauthorisedResponse() throws IOException { } public void testInfer_SetsInputTypeToIngest_FromInferParameter_WhenTaskSettingsAreEmpty() throws IOException { - var senderFactory = new HttpRequestSenderFactory(threadPool, clientManager, mockClusterServiceEmpty(), Settings.EMPTY); + var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager); try (var service = new CohereService(senderFactory, createWithEmptySettings(threadPool))) { @@ -817,7 +816,7 @@ public void testInfer_SetsInputTypeToIngest_FromInferParameter_WhenTaskSettingsA public void testInfer_SetsInputTypeToIngestFromInferParameter_WhenModelSettingIsNull_AndRequestTaskSettingsIsSearch() throws IOException { - var senderFactory = new HttpRequestSenderFactory(threadPool, clientManager, mockClusterServiceEmpty(), Settings.EMPTY); + var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager); try (var service = new CohereService(senderFactory, createWithEmptySettings(threadPool))) { @@ -883,7 +882,7 @@ public void testInfer_SetsInputTypeToIngestFromInferParameter_WhenModelSettingIs } public void testInfer_DoesNotSetInputType_WhenNotPresentInTaskSettings_AndUnspecifiedIsPassedInRequest() throws IOException { - var senderFactory = new HttpRequestSenderFactory(threadPool, clientManager, mockClusterServiceEmpty(), Settings.EMPTY); + var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager); try (var service = new CohereService(senderFactory, createWithEmptySettings(threadPool))) { @@ -957,7 +956,7 @@ private Map getRequestConfigMap( } private CohereService createCohereService() { - return new CohereService(mock(HttpRequestSenderFactory.class), createWithEmptySettings(threadPool)); + return new CohereService(mock(HttpRequestSender.Factory.class), createWithEmptySettings(threadPool)); } private PeristedConfig getPersistedConfigMap( diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceBaseServiceTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceBaseServiceTests.java index 345aa1a80e5bd..cd896cb18440a 100644 --- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceBaseServiceTests.java +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceBaseServiceTests.java @@ -16,7 +16,7 @@ import org.elasticsearch.inference.TaskType; import org.elasticsearch.test.ESTestCase; import org.elasticsearch.threadpool.ThreadPool; -import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSenderFactory; +import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSender; import org.elasticsearch.xpack.inference.external.http.sender.Sender; import org.elasticsearch.xpack.inference.services.ServiceComponents; import org.junit.After; @@ -57,7 +57,7 @@ public void shutdown() throws IOException { public void testInfer_ThrowsErrorWhenModelIsNotHuggingFaceModel() throws IOException { var sender = mock(Sender.class); - var factory = mock(HttpRequestSenderFactory.class); + var factory = mock(HttpRequestSender.Factory.class); when(factory.createSender(anyString())).thenReturn(sender); var mockModel = getInvalidModel("model_id", "service_name"); @@ -82,8 +82,7 @@ public void testInfer_ThrowsErrorWhenModelIsNotHuggingFaceModel() throws IOExcep } private static final class TestService extends HuggingFaceBaseService { - - TestService(HttpRequestSenderFactory factory, ServiceComponents serviceComponents) { + TestService(HttpRequestSender.Factory factory, ServiceComponents serviceComponents) { super(factory, serviceComponents); } diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceServiceTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceServiceTests.java index 23d6bd17e48d1..c4c49065cd79c 100644 --- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceServiceTests.java +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/huggingface/HuggingFaceServiceTests.java @@ -28,7 +28,8 @@ import org.elasticsearch.threadpool.ThreadPool; import org.elasticsearch.xcontent.XContentType; import org.elasticsearch.xpack.inference.external.http.HttpClientManager; -import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSenderFactory; +import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSender; +import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSenderTests; import org.elasticsearch.xpack.inference.logging.ThrottlerManager; import org.elasticsearch.xpack.inference.results.SparseEmbeddingResultsTests; import org.elasticsearch.xpack.inference.services.huggingface.elser.HuggingFaceElserModel; @@ -83,7 +84,6 @@ public void shutdown() throws IOException { public void testParseRequestConfig_CreatesAnEmbeddingsModel() throws IOException { try (var service = createHuggingFaceService()) { - ActionListener modelVerificationActionListener = ActionListener.wrap((model) -> { assertThat(model, instanceOf(HuggingFaceEmbeddingsModel.class)); @@ -408,7 +408,7 @@ public void testParsePersistedConfig_DoesNotThrowWhenAnExtraKeyExistsInTaskSetti } public void testInfer_SendsEmbeddingsRequest() throws IOException { - var senderFactory = new HttpRequestSenderFactory(threadPool, clientManager, mockClusterServiceEmpty(), Settings.EMPTY); + var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager); try (var service = new HuggingFaceService(senderFactory, createWithEmptySettings(threadPool))) { @@ -446,7 +446,7 @@ public void testInfer_SendsEmbeddingsRequest() throws IOException { } public void testInfer_SendsElserRequest() throws IOException { - var senderFactory = new HttpRequestSenderFactory(threadPool, clientManager, mockClusterServiceEmpty(), Settings.EMPTY); + var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager); try (var service = new HuggingFaceService(senderFactory, createWithEmptySettings(threadPool))) { @@ -488,7 +488,7 @@ public void testInfer_SendsElserRequest() throws IOException { } public void testCheckModelConfig_IncludesMaxTokens() throws IOException { - var senderFactory = new HttpRequestSenderFactory(threadPool, clientManager, mockClusterServiceEmpty(), Settings.EMPTY); + var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager); try (var service = new HuggingFaceService(senderFactory, createWithEmptySettings(threadPool))) { @@ -513,7 +513,7 @@ public void testCheckModelConfig_IncludesMaxTokens() throws IOException { } private HuggingFaceService createHuggingFaceService() { - return new HuggingFaceService(mock(HttpRequestSenderFactory.class), createWithEmptySettings(threadPool)); + return new HuggingFaceService(mock(HttpRequestSender.Factory.class), createWithEmptySettings(threadPool)); } private Map getRequestConfigMap(Map serviceSettings, Map secretSettings) { diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/openai/OpenAiServiceTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/openai/OpenAiServiceTests.java index e97040ed7d795..d819b2b243872 100644 --- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/openai/OpenAiServiceTests.java +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/openai/OpenAiServiceTests.java @@ -29,7 +29,8 @@ import org.elasticsearch.threadpool.ThreadPool; import org.elasticsearch.xcontent.XContentType; import org.elasticsearch.xpack.inference.external.http.HttpClientManager; -import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSenderFactory; +import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSender; +import org.elasticsearch.xpack.inference.external.http.sender.HttpRequestSenderTests; import org.elasticsearch.xpack.inference.external.http.sender.Sender; import org.elasticsearch.xpack.inference.logging.ThrottlerManager; import org.elasticsearch.xpack.inference.services.ServiceFields; @@ -604,7 +605,7 @@ public void testParsePersistedConfig_NotThrowWhenAnExtraKeyExistsInTaskSettings( public void testInfer_ThrowsErrorWhenModelIsNotOpenAiModel() throws IOException { var sender = mock(Sender.class); - var factory = mock(HttpRequestSenderFactory.class); + var factory = mock(HttpRequestSender.Factory.class); when(factory.createSender(anyString())).thenReturn(sender); var mockModel = getInvalidModel("model_id", "service_name"); @@ -629,7 +630,7 @@ public void testInfer_ThrowsErrorWhenModelIsNotOpenAiModel() throws IOException } public void testInfer_SendsRequest() throws IOException { - var senderFactory = new HttpRequestSenderFactory(threadPool, clientManager, mockClusterServiceEmpty(), Settings.EMPTY); + var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager); try (var service = new OpenAiService(senderFactory, createWithEmptySettings(threadPool))) { @@ -677,7 +678,7 @@ public void testInfer_SendsRequest() throws IOException { } public void testCheckModelConfig_IncludesMaxTokens() throws IOException { - var senderFactory = new HttpRequestSenderFactory(threadPool, clientManager, mockClusterServiceEmpty(), Settings.EMPTY); + var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager); try (var service = new OpenAiService(senderFactory, createWithEmptySettings(threadPool))) { @@ -718,7 +719,7 @@ public void testCheckModelConfig_IncludesMaxTokens() throws IOException { } public void testCheckModelConfig_ThrowsIfEmbeddingSizeDoesNotMatchValueSetByUser() throws IOException { - var senderFactory = new HttpRequestSenderFactory(threadPool, clientManager, mockClusterServiceEmpty(), Settings.EMPTY); + var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager); try (var service = new OpenAiService(senderFactory, createWithEmptySettings(threadPool))) { @@ -769,7 +770,7 @@ public void testCheckModelConfig_ThrowsIfEmbeddingSizeDoesNotMatchValueSetByUser public void testCheckModelConfig_ReturnsModelWithDimensionsSetTo2_AndDocProductSet_IfDimensionsSetByUser_ButSetToNull() throws IOException { - var senderFactory = new HttpRequestSenderFactory(threadPool, clientManager, mockClusterServiceEmpty(), Settings.EMPTY); + var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager); try (var service = new OpenAiService(senderFactory, createWithEmptySettings(threadPool))) { @@ -827,7 +828,7 @@ public void testCheckModelConfig_ReturnsModelWithDimensionsSetTo2_AndDocProductS public void testCheckModelConfig_ReturnsModelWithSameDimensions_AndDocProductSet_IfDimensionsSetByUser_AndTheyMatchReturnedSize() throws IOException { - var senderFactory = new HttpRequestSenderFactory(threadPool, clientManager, mockClusterServiceEmpty(), Settings.EMPTY); + var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager); try (var service = new OpenAiService(senderFactory, createWithEmptySettings(threadPool))) { @@ -886,7 +887,7 @@ public void testCheckModelConfig_ReturnsModelWithSameDimensions_AndDocProductSet } public void testCheckModelConfig_ReturnsNewModelReference_AndDoesNotSendDimensionsField_WhenNotSetByUser() throws IOException { - var senderFactory = new HttpRequestSenderFactory(threadPool, clientManager, mockClusterServiceEmpty(), Settings.EMPTY); + var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager); try (var service = new OpenAiService(senderFactory, createWithEmptySettings(threadPool))) { @@ -952,7 +953,7 @@ public void testCheckModelConfig_ReturnsNewModelReference_AndDoesNotSendDimensio } public void testInfer_UnauthorisedResponse() throws IOException { - var senderFactory = new HttpRequestSenderFactory(threadPool, clientManager, mockClusterServiceEmpty(), Settings.EMPTY); + var senderFactory = HttpRequestSenderTests.createSenderFactory(threadPool, clientManager); try (var service = new OpenAiService(senderFactory, createWithEmptySettings(threadPool))) { @@ -1007,7 +1008,7 @@ public void testMoveModelFromTaskToServiceSettings_AlreadyMoved() { } private OpenAiService createOpenAiService() { - return new OpenAiService(mock(HttpRequestSenderFactory.class), createWithEmptySettings(threadPool)); + return new OpenAiService(mock(HttpRequestSender.Factory.class), createWithEmptySettings(threadPool)); } private Map getRequestConfigMap( diff --git a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/openai/embeddings/OpenAiEmbeddingsModelTests.java b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/openai/embeddings/OpenAiEmbeddingsModelTests.java index 01b60fdb896d0..db5febef1dab2 100644 --- a/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/openai/embeddings/OpenAiEmbeddingsModelTests.java +++ b/x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/openai/embeddings/OpenAiEmbeddingsModelTests.java @@ -47,6 +47,24 @@ public void testOverrideWith_NullMap() { assertThat(overriddenModel, sameInstance(model)); } + public static OpenAiEmbeddingsModel createModel( + String url, + @Nullable String org, + String apiKey, + String modelName, + @Nullable String user, + String inferenceEntityId + ) { + return new OpenAiEmbeddingsModel( + inferenceEntityId, + TaskType.TEXT_EMBEDDING, + "service", + new OpenAiEmbeddingsServiceSettings(modelName, url, org, SimilarityMeasure.DOT_PRODUCT, 1536, null, false), + new OpenAiEmbeddingsTaskSettings(user), + new DefaultSecretSettings(new SecureString(apiKey.toCharArray())) + ); + } + public static OpenAiEmbeddingsModel createModel( String url, @Nullable String org, From d72665a207a9ae3074ed391ab0407fd397b7dfd9 Mon Sep 17 00:00:00 2001 From: John Verwolf Date: Tue, 5 Mar 2024 10:47:47 -0800 Subject: [PATCH 10/27] Bugfix: Disable eager loading BitSetFilterCache on Stateless Indexing Nodes (#105791) The BitSetFilterCache is used for search traffic, which is not served on stateless indexing nodes. Thus, we can disable it and save memory. --- docs/changelog/105791.yaml | 5 ++ .../index/cache/bitset/BitsetFilterCache.java | 18 +++++- .../cache/bitset/BitSetFilterCacheTests.java | 57 +++++++++++++++++++ 3 files changed, 79 insertions(+), 1 deletion(-) create mode 100644 docs/changelog/105791.yaml diff --git a/docs/changelog/105791.yaml b/docs/changelog/105791.yaml new file mode 100644 index 0000000000000..f18b5e6b8fdd7 --- /dev/null +++ b/docs/changelog/105791.yaml @@ -0,0 +1,5 @@ +pr: 105791 +summary: "Bugfix: Disable eager loading `BitSetFilterCache` on Indexing Nodes" +area: Search +type: bug +issues: [] diff --git a/server/src/main/java/org/elasticsearch/index/cache/bitset/BitsetFilterCache.java b/server/src/main/java/org/elasticsearch/index/cache/bitset/BitsetFilterCache.java index f1f03eff88d08..f8bc40a395472 100644 --- a/server/src/main/java/org/elasticsearch/index/cache/bitset/BitsetFilterCache.java +++ b/server/src/main/java/org/elasticsearch/index/cache/bitset/BitsetFilterCache.java @@ -24,6 +24,8 @@ import org.apache.lucene.util.BitDocIdSet; import org.apache.lucene.util.BitSet; import org.elasticsearch.ExceptionsHelper; +import org.elasticsearch.cluster.node.DiscoveryNode; +import org.elasticsearch.cluster.node.DiscoveryNodeRole; import org.elasticsearch.common.cache.Cache; import org.elasticsearch.common.cache.CacheBuilder; import org.elasticsearch.common.cache.RemovalListener; @@ -55,6 +57,8 @@ import java.util.concurrent.ExecutionException; import java.util.concurrent.Executor; +import static org.elasticsearch.index.IndexSettings.INDEX_FAST_REFRESH_SETTING; + /** * This is a cache for {@link BitDocIdSet} based filters and is unbounded by size or time. *

@@ -92,10 +96,22 @@ public BitsetFilterCache(IndexSettings indexSettings, Listener listener) { throw new IllegalArgumentException("listener must not be null"); } this.index = indexSettings.getIndex(); - this.loadRandomAccessFiltersEagerly = indexSettings.getValue(INDEX_LOAD_RANDOM_ACCESS_FILTERS_EAGERLY_SETTING); + this.loadRandomAccessFiltersEagerly = shouldLoadRandomAccessFiltersEagerly(indexSettings); this.listener = listener; } + static boolean shouldLoadRandomAccessFiltersEagerly(IndexSettings settings) { + boolean loadFiltersEagerlySetting = settings.getValue(INDEX_LOAD_RANDOM_ACCESS_FILTERS_EAGERLY_SETTING); + boolean isStateless = DiscoveryNode.isStateless(settings.getNodeSettings()); + if (isStateless) { + return DiscoveryNode.hasRole(settings.getNodeSettings(), DiscoveryNodeRole.INDEX_ROLE) + && loadFiltersEagerlySetting + && INDEX_FAST_REFRESH_SETTING.get(settings.getSettings()); + } else { + return loadFiltersEagerlySetting; + } + } + public static BitSet bitsetFromQuery(Query query, LeafReaderContext context) throws IOException { final IndexReaderContext topLevelContext = ReaderUtil.getTopLevelContext(context); final IndexSearcher searcher = new IndexSearcher(topLevelContext); diff --git a/server/src/test/java/org/elasticsearch/index/cache/bitset/BitSetFilterCacheTests.java b/server/src/test/java/org/elasticsearch/index/cache/bitset/BitSetFilterCacheTests.java index 1c164e898426d..6d72649e90764 100644 --- a/server/src/test/java/org/elasticsearch/index/cache/bitset/BitSetFilterCacheTests.java +++ b/server/src/test/java/org/elasticsearch/index/cache/bitset/BitSetFilterCacheTests.java @@ -28,19 +28,27 @@ import org.apache.lucene.store.Directory; import org.apache.lucene.util.Accountable; import org.apache.lucene.util.BitSet; +import org.elasticsearch.cluster.metadata.IndexMetadata; +import org.elasticsearch.cluster.node.DiscoveryNodeRole; import org.elasticsearch.common.lucene.index.ElasticsearchDirectoryReader; import org.elasticsearch.common.settings.Settings; import org.elasticsearch.core.IOUtils; +import org.elasticsearch.index.Index; import org.elasticsearch.index.IndexSettings; import org.elasticsearch.index.shard.ShardId; import org.elasticsearch.lucene.util.MatchAllBitSet; +import org.elasticsearch.node.NodeRoleSettings; import org.elasticsearch.test.ESTestCase; import org.elasticsearch.test.IndexSettingsModule; import java.io.IOException; +import java.util.List; import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicLong; +import static org.elasticsearch.cluster.node.DiscoveryNode.STATELESS_ENABLED_SETTING_NAME; +import static org.elasticsearch.index.IndexSettings.INDEX_FAST_REFRESH_SETTING; +import static org.elasticsearch.index.cache.bitset.BitsetFilterCache.INDEX_LOAD_RANDOM_ACCESS_FILTERS_EAGERLY_SETTING; import static org.hamcrest.Matchers.equalTo; import static org.hamcrest.Matchers.greaterThan; import static org.hamcrest.Matchers.hasSize; @@ -259,4 +267,53 @@ public void onRemoval(ShardId shardId, Accountable accountable) { } } + public void testShouldLoadRandomAccessFiltersEagerly() { + var values = List.of(true, false); + for (var hasIndexRole : values) { + for (var indexFastRefresh : values) { + for (var loadFiltersEagerly : values) { + for (var isStateless : values) { + if (isStateless) { + assertEquals( + loadFiltersEagerly && indexFastRefresh && hasIndexRole, + BitsetFilterCache.shouldLoadRandomAccessFiltersEagerly( + bitsetFilterCacheSettings(isStateless, hasIndexRole, loadFiltersEagerly, indexFastRefresh) + ) + ); + } else { + assertEquals( + loadFiltersEagerly, + BitsetFilterCache.shouldLoadRandomAccessFiltersEagerly( + bitsetFilterCacheSettings(isStateless, hasIndexRole, loadFiltersEagerly, indexFastRefresh) + ) + ); + } + } + } + } + } + } + + private IndexSettings bitsetFilterCacheSettings( + boolean isStateless, + boolean hasIndexRole, + boolean loadFiltersEagerly, + boolean indexFastRefresh + ) { + var indexSettingsBuilder = Settings.builder().put(INDEX_LOAD_RANDOM_ACCESS_FILTERS_EAGERLY_SETTING.getKey(), loadFiltersEagerly); + if (isStateless) indexSettingsBuilder.put(INDEX_FAST_REFRESH_SETTING.getKey(), indexFastRefresh); + + var nodeSettingsBuilder = Settings.builder() + .putList( + NodeRoleSettings.NODE_ROLES_SETTING.getKey(), + hasIndexRole ? DiscoveryNodeRole.INDEX_ROLE.roleName() : DiscoveryNodeRole.SEARCH_ROLE.roleName() + ) + .put(STATELESS_ENABLED_SETTING_NAME, isStateless); + + return IndexSettingsModule.newIndexSettings( + new Index("index", IndexMetadata.INDEX_UUID_NA_VALUE), + indexSettingsBuilder.build(), + nodeSettingsBuilder.build() + ); + } } From e8039b9ecb2451752ac5377c44a6a0c662087a9f Mon Sep 17 00:00:00 2001 From: Nik Everett Date: Tue, 5 Mar 2024 16:56:28 -0500 Subject: [PATCH 11/27] ESQL: Reenable svq tests (#105996) We fixed the test failure in #105986 but this snuck in. Closes #105952 --- .../xpack/esql/querydsl/query/SingleValueQueryTests.java | 1 - 1 file changed, 1 deletion(-) diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/querydsl/query/SingleValueQueryTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/querydsl/query/SingleValueQueryTests.java index 6465e73417ae2..1324b3977786a 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/querydsl/query/SingleValueQueryTests.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/querydsl/query/SingleValueQueryTests.java @@ -77,7 +77,6 @@ public void testMatchAll() throws IOException { testCase(new SingleValueQuery(new MatchAll(Source.EMPTY), "foo").asBuilder(), YesNoSometimes.NO, YesNoSometimes.NO, this::runCase); } - @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/105952") public void testMatchSome() throws IOException { int max = between(1, 100); testCase( From c0d35f1e77dd7031abf039ad7d8f3a686e1fe954 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Przemys=C5=82aw=20Witek?= Date: Wed, 6 Mar 2024 09:33:11 +0100 Subject: [PATCH 12/27] [Transform] Cleanup and simplify reading transform settings (#105554) --- .../transform/transforms/SettingsConfig.java | 36 +++-- .../transform/transforms/TransformConfig.java | 2 +- .../TransformEffectiveSettings.java | 86 +++++++++++ .../transforms/SettingsConfigTests.java | 32 +++-- .../TransformEffectiveSettingsTests.java | 135 ++++++++++++++++++ .../integration/TransformProgressIT.java | 4 +- .../TransportPreviewTransformAction.java | 3 +- .../action/TransportStartTransformAction.java | 7 +- .../TimeBasedCheckpointProvider.java | 3 +- .../transform/persistence/TransformIndex.java | 5 +- .../transforms/ClientTransformIndexer.java | 10 +- .../transforms/TransformFailureHandler.java | 38 ++--- .../transforms/TransformIndexer.java | 5 +- .../transform/transforms/pivot/Pivot.java | 10 +- .../transforms/pivot/SchemaUtil.java | 9 +- .../action/TransformConfigLinterTests.java | 4 +- .../ClientTransformIndexerTests.java | 3 +- .../AggregationSchemaAndResultTests.java | 4 +- .../transforms/pivot/PivotTests.java | 16 +-- 19 files changed, 305 insertions(+), 107 deletions(-) create mode 100644 x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/transform/transforms/TransformEffectiveSettings.java create mode 100644 x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/transform/transforms/TransformEffectiveSettingsTests.java diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/transform/transforms/SettingsConfig.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/transform/transforms/SettingsConfig.java index 9b0fa3876819b..1557f2843b6af 100644 --- a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/transform/transforms/SettingsConfig.java +++ b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/transform/transforms/SettingsConfig.java @@ -31,6 +31,9 @@ import static org.elasticsearch.xcontent.ConstructingObjectParser.optionalConstructorArg; public class SettingsConfig implements Writeable, ToXContentObject { + + public static final SettingsConfig EMPTY = new SettingsConfig(null, null, null, null, null, null, null, (Integer) null); + public static final ConstructingObjectParser STRICT_PARSER = createParser(false); public static final ConstructingObjectParser LENIENT_PARSER = createParser(true); @@ -110,10 +113,6 @@ private static ConstructingObjectParser createParser(boole private final Integer numFailureRetries; private final Integer unattended; - public SettingsConfig() { - this(null, null, (Integer) null, (Integer) null, (Integer) null, (Integer) null, (Integer) null, (Integer) null); - } - public SettingsConfig( Integer maxPageSearchSize, Float docsPerSecond, @@ -136,7 +135,7 @@ public SettingsConfig( ); } - SettingsConfig( + private SettingsConfig( Integer maxPageSearchSize, Float docsPerSecond, Integer datesAsEpochMillis, @@ -188,51 +187,51 @@ public Float getDocsPerSecond() { return docsPerSecond; } - public Boolean getDatesAsEpochMillis() { + Boolean getDatesAsEpochMillis() { return datesAsEpochMillis != null ? datesAsEpochMillis > 0 : null; } - public Integer getDatesAsEpochMillisForUpdate() { + Integer getDatesAsEpochMillisForUpdate() { return datesAsEpochMillis; } - public Boolean getAlignCheckpoints() { + Boolean getAlignCheckpoints() { return alignCheckpoints != null ? (alignCheckpoints > 0) || (alignCheckpoints == DEFAULT_ALIGN_CHECKPOINTS) : null; } - public Integer getAlignCheckpointsForUpdate() { + Integer getAlignCheckpointsForUpdate() { return alignCheckpoints; } - public Boolean getUsePit() { + Boolean getUsePit() { return usePit != null ? (usePit > 0) || (usePit == DEFAULT_USE_PIT) : null; } - public Integer getUsePitForUpdate() { + Integer getUsePitForUpdate() { return usePit; } - public Boolean getDeduceMappings() { + Boolean getDeduceMappings() { return deduceMappings != null ? (deduceMappings > 0) || (deduceMappings == DEFAULT_DEDUCE_MAPPINGS) : null; } - public Integer getDeduceMappingsForUpdate() { + Integer getDeduceMappingsForUpdate() { return deduceMappings; } - public Integer getNumFailureRetries() { + Integer getNumFailureRetries() { return numFailureRetries != null ? (numFailureRetries == DEFAULT_NUM_FAILURE_RETRIES ? null : numFailureRetries) : null; } - public Integer getNumFailureRetriesForUpdate() { + Integer getNumFailureRetriesForUpdate() { return numFailureRetries; } - public Boolean getUnattended() { + Boolean getUnattended() { return unattended != null ? (unattended == DEFAULT_UNATTENDED) ? null : (unattended > 0) : null; } - public Integer getUnattendedForUpdate() { + Integer getUnattendedForUpdate() { return unattended; } @@ -495,7 +494,7 @@ public Builder setNumFailureRetries(Integer numFailureRetries) { * An explicit `null` resets to default. * * @param unattended true if this is a unattended transform. - * @return the {@link Builder} with usePit set. + * @return the {@link Builder} with unattended set. */ public Builder setUnattended(Boolean unattended) { this.unattended = unattended == null ? DEFAULT_UNATTENDED : unattended ? 1 : 0; @@ -545,7 +544,6 @@ public Builder update(SettingsConfig update) { if (update.getUnattendedForUpdate() != null) { this.unattended = update.getUnattendedForUpdate().equals(DEFAULT_UNATTENDED) ? null : update.getUnattendedForUpdate(); } - return this; } diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/transform/transforms/TransformConfig.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/transform/transforms/TransformConfig.java index d89eb9b397180..fb782bdae0068 100644 --- a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/transform/transforms/TransformConfig.java +++ b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/transform/transforms/TransformConfig.java @@ -234,7 +234,7 @@ public TransformConfig( this.pivotConfig = pivotConfig; this.latestConfig = latestConfig; this.description = description; - this.settings = settings == null ? new SettingsConfig() : settings; + this.settings = settings == null ? SettingsConfig.EMPTY : settings; this.metadata = metadata; this.retentionPolicyConfig = retentionPolicyConfig; if (this.description != null && this.description.length() > MAX_DESCRIPTION_LENGTH) { diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/transform/transforms/TransformEffectiveSettings.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/transform/transforms/TransformEffectiveSettings.java new file mode 100644 index 0000000000000..3d4b8ccc64d89 --- /dev/null +++ b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/transform/transforms/TransformEffectiveSettings.java @@ -0,0 +1,86 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.core.transform.transforms; + +import org.elasticsearch.xpack.core.transform.TransformConfigVersion; + +public final class TransformEffectiveSettings { + + private TransformEffectiveSettings() {} + + /** + * Determines if the transform should write dates as epoch millis based on settings and version. + * + * @param settings transform's settings + * @return whether or not the transform is unattended + */ + public static boolean writeDatesAsEpochMillis(SettingsConfig settings, TransformConfigVersion version) { + // defines how dates are written, if not specified in settings + // < 7.11 as epoch millis + // >= 7.11 as string + // note: it depends on the version when the transform has been created, not the version of the code + return settings.getDatesAsEpochMillis() != null + ? settings.getDatesAsEpochMillis() + : version.before(TransformConfigVersion.V_7_11_0); + } + + /** + * Determines if aligning checkpoints is disabled for this transform based on settings. + * + * @param settings transform's settings + * @return whether or not aligning checkpoints is disabled for this transform + */ + public static boolean isAlignCheckpointsDisabled(SettingsConfig settings) { + return Boolean.FALSE.equals(settings.getAlignCheckpoints()); + } + + /** + * Determines if pit is disabled for this transform based on settings. + * + * @param settings transform's settings + * @return whether or not pit is disabled for this transform + */ + public static boolean isPitDisabled(SettingsConfig settings) { + return Boolean.FALSE.equals(settings.getUsePit()); + } + + /** + * Determines if mappings deduction is disabled for this transform based on settings. + * + * @param settings transform's settings + * @return whether or not mappings deduction is disabled for this transform + */ + public static boolean isDeduceMappingsDisabled(SettingsConfig settings) { + return Boolean.FALSE.equals(settings.getDeduceMappings()); + } + + /** + * Determines the appropriate number of retries. + *

+ * The number of retries are read from the config or if not read from the context which is based on a cluster wide default. + * If the transform runs in unattended mode, the number of retries is always indefinite. + * + * @param settings transform's settings + * @return the number of retries or -1 if retries are indefinite + */ + public static int getNumFailureRetries(SettingsConfig settings, int defaultNumFailureRetries) { + return isUnattended(settings) ? -1 + : settings.getNumFailureRetries() != null ? settings.getNumFailureRetries() + : defaultNumFailureRetries; + } + + /** + * Determines if the transform is unattended based on settings. + * + * @param settings transform's settings + * @return whether or not the transform is unattended + */ + public static boolean isUnattended(SettingsConfig settings) { + return Boolean.TRUE.equals(settings.getUnattended()); + } +} diff --git a/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/transform/transforms/SettingsConfigTests.java b/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/transform/transforms/SettingsConfigTests.java index 62b9e2e48a907..6bedd60d582dd 100644 --- a/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/transform/transforms/SettingsConfigTests.java +++ b/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/transform/transforms/SettingsConfigTests.java @@ -33,32 +33,30 @@ public class SettingsConfigTests extends AbstractSerializingTransformTestCase instanceReader() { } public void testExplicitNullParsing() throws IOException { - // explicit null assertThat(fromString("{\"max_page_search_size\" : null}").getMaxPageSearchSize(), equalTo(-1)); // not set @@ -119,6 +116,11 @@ public void testExplicitNullParsing() throws IOException { assertThat(fromString("{\"num_failure_retries\" : null}").getNumFailureRetriesForUpdate(), equalTo(-2)); assertNull(fromString("{}").getNumFailureRetries()); assertNull(fromString("{}").getNumFailureRetriesForUpdate()); + + assertNull(fromString("{\"unattended\" : null}").getUnattended()); + assertThat(fromString("{\"unattended\" : null}").getUnattendedForUpdate(), equalTo(-1)); + assertNull(fromString("{}").getUnattended()); + assertNull(fromString("{}").getUnattendedForUpdate()); } public void testUpdateMaxPageSearchSizeUsingBuilder() throws IOException { diff --git a/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/transform/transforms/TransformEffectiveSettingsTests.java b/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/transform/transforms/TransformEffectiveSettingsTests.java new file mode 100644 index 0000000000000..98726d8dbf272 --- /dev/null +++ b/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/transform/transforms/TransformEffectiveSettingsTests.java @@ -0,0 +1,135 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.core.transform.transforms; + +import org.elasticsearch.test.ESTestCase; +import org.elasticsearch.xpack.core.transform.TransformConfigVersion; + +public class TransformEffectiveSettingsTests extends ESTestCase { + + public void testWriteDatesAsEpochMillis() { + SettingsConfig settingsConfig = new SettingsConfig.Builder().build(); + assertFalse(TransformEffectiveSettings.writeDatesAsEpochMillis(settingsConfig, TransformConfigVersion.V_7_11_0)); + assertTrue(TransformEffectiveSettings.writeDatesAsEpochMillis(settingsConfig, TransformConfigVersion.V_7_10_1)); + + settingsConfig = new SettingsConfig.Builder().setDatesAsEpochMillis(null).build(); + assertFalse(TransformEffectiveSettings.writeDatesAsEpochMillis(settingsConfig, TransformConfigVersion.V_7_11_0)); + // Note that the result is not the same as if we just left "setDatesAsEpochMillis" unset in the builder! + assertFalse(TransformEffectiveSettings.writeDatesAsEpochMillis(settingsConfig, TransformConfigVersion.V_7_10_1)); + + settingsConfig = new SettingsConfig.Builder().setDatesAsEpochMillis(false).build(); + assertFalse(TransformEffectiveSettings.writeDatesAsEpochMillis(settingsConfig, TransformConfigVersion.V_7_11_0)); + assertFalse(TransformEffectiveSettings.writeDatesAsEpochMillis(settingsConfig, TransformConfigVersion.V_7_10_1)); + + settingsConfig = new SettingsConfig.Builder().setDatesAsEpochMillis(true).build(); + assertTrue(TransformEffectiveSettings.writeDatesAsEpochMillis(settingsConfig, TransformConfigVersion.V_7_11_0)); + assertTrue(TransformEffectiveSettings.writeDatesAsEpochMillis(settingsConfig, TransformConfigVersion.V_7_10_1)); + } + + public void testIsAlignCheckpointsDisabled() { + SettingsConfig settingsConfig = new SettingsConfig.Builder().build(); + assertFalse(TransformEffectiveSettings.isAlignCheckpointsDisabled(settingsConfig)); + + settingsConfig = new SettingsConfig.Builder().setAlignCheckpoints(null).build(); + assertFalse(TransformEffectiveSettings.isAlignCheckpointsDisabled(settingsConfig)); + + settingsConfig = new SettingsConfig.Builder().setAlignCheckpoints(false).build(); + assertTrue(TransformEffectiveSettings.isAlignCheckpointsDisabled(settingsConfig)); + + settingsConfig = new SettingsConfig.Builder().setAlignCheckpoints(true).build(); + assertFalse(TransformEffectiveSettings.isAlignCheckpointsDisabled(settingsConfig)); + } + + public void testIsPitDisabled() { + SettingsConfig settingsConfig = new SettingsConfig.Builder().build(); + assertFalse(TransformEffectiveSettings.isPitDisabled(settingsConfig)); + + settingsConfig = new SettingsConfig.Builder().setUsePit(null).build(); + assertFalse(TransformEffectiveSettings.isPitDisabled(settingsConfig)); + + settingsConfig = new SettingsConfig.Builder().setUsePit(false).build(); + assertTrue(TransformEffectiveSettings.isPitDisabled(settingsConfig)); + + settingsConfig = new SettingsConfig.Builder().setUsePit(true).build(); + assertFalse(TransformEffectiveSettings.isPitDisabled(settingsConfig)); + } + + public void testIsDeduceMappingsDisabled() { + SettingsConfig settingsConfig = new SettingsConfig.Builder().build(); + assertFalse(TransformEffectiveSettings.isDeduceMappingsDisabled(settingsConfig)); + + settingsConfig = new SettingsConfig.Builder().setDeduceMappings(null).build(); + assertFalse(TransformEffectiveSettings.isDeduceMappingsDisabled(settingsConfig)); + + settingsConfig = new SettingsConfig.Builder().setDeduceMappings(false).build(); + assertTrue(TransformEffectiveSettings.isDeduceMappingsDisabled(settingsConfig)); + + settingsConfig = new SettingsConfig.Builder().setDeduceMappings(true).build(); + assertFalse(TransformEffectiveSettings.isDeduceMappingsDisabled(settingsConfig)); + } + + public void testGetNumFailureRetries() { + SettingsConfig settingsConfig = new SettingsConfig.Builder().build(); + assertEquals(10, TransformEffectiveSettings.getNumFailureRetries(settingsConfig, 10)); + + settingsConfig = new SettingsConfig.Builder().setNumFailureRetries(null).build(); + assertEquals(10, TransformEffectiveSettings.getNumFailureRetries(settingsConfig, 10)); + + settingsConfig = new SettingsConfig.Builder().setNumFailureRetries(-1).build(); + assertEquals(-1, TransformEffectiveSettings.getNumFailureRetries(settingsConfig, 10)); + + settingsConfig = new SettingsConfig.Builder().setNumFailureRetries(0).build(); + assertEquals(0, TransformEffectiveSettings.getNumFailureRetries(settingsConfig, 10)); + + settingsConfig = new SettingsConfig.Builder().setNumFailureRetries(1).build(); + assertEquals(1, TransformEffectiveSettings.getNumFailureRetries(settingsConfig, 10)); + + settingsConfig = new SettingsConfig.Builder().setNumFailureRetries(10).build(); + assertEquals(10, TransformEffectiveSettings.getNumFailureRetries(settingsConfig, 10)); + + settingsConfig = new SettingsConfig.Builder().setNumFailureRetries(100).build(); + assertEquals(100, TransformEffectiveSettings.getNumFailureRetries(settingsConfig, 10)); + } + + public void testGetNumFailureRetries_Unattended() { + SettingsConfig settingsConfig = new SettingsConfig.Builder().setUnattended(true).build(); + assertEquals(-1, TransformEffectiveSettings.getNumFailureRetries(settingsConfig, 10)); + + settingsConfig = new SettingsConfig.Builder().setUnattended(true).setNumFailureRetries(null).build(); + assertEquals(-1, TransformEffectiveSettings.getNumFailureRetries(settingsConfig, 10)); + + settingsConfig = new SettingsConfig.Builder().setUnattended(true).setNumFailureRetries(-1).build(); + assertEquals(-1, TransformEffectiveSettings.getNumFailureRetries(settingsConfig, 10)); + + settingsConfig = new SettingsConfig.Builder().setUnattended(true).setNumFailureRetries(0).build(); + assertEquals(-1, TransformEffectiveSettings.getNumFailureRetries(settingsConfig, 10)); + + settingsConfig = new SettingsConfig.Builder().setUnattended(true).setNumFailureRetries(1).build(); + assertEquals(-1, TransformEffectiveSettings.getNumFailureRetries(settingsConfig, 10)); + + settingsConfig = new SettingsConfig.Builder().setUnattended(true).setNumFailureRetries(10).build(); + assertEquals(-1, TransformEffectiveSettings.getNumFailureRetries(settingsConfig, 10)); + + settingsConfig = new SettingsConfig.Builder().setUnattended(true).setNumFailureRetries(100).build(); + assertEquals(-1, TransformEffectiveSettings.getNumFailureRetries(settingsConfig, 10)); + } + + public void testIsUnattended() { + SettingsConfig settingsConfig = new SettingsConfig.Builder().build(); + assertFalse(TransformEffectiveSettings.isUnattended(settingsConfig)); + + settingsConfig = new SettingsConfig.Builder().setUnattended(null).build(); + assertFalse(TransformEffectiveSettings.isUnattended(settingsConfig)); + + settingsConfig = new SettingsConfig.Builder().setUnattended(false).build(); + assertFalse(TransformEffectiveSettings.isUnattended(settingsConfig)); + + settingsConfig = new SettingsConfig.Builder().setUnattended(true).build(); + assertTrue(TransformEffectiveSettings.isUnattended(settingsConfig)); + } +} diff --git a/x-pack/plugin/transform/src/internalClusterTest/java/org/elasticsearch/xpack/transform/integration/TransformProgressIT.java b/x-pack/plugin/transform/src/internalClusterTest/java/org/elasticsearch/xpack/transform/integration/TransformProgressIT.java index c62ff49ae6865..dbe09663abc20 100644 --- a/x-pack/plugin/transform/src/internalClusterTest/java/org/elasticsearch/xpack/transform/integration/TransformProgressIT.java +++ b/x-pack/plugin/transform/src/internalClusterTest/java/org/elasticsearch/xpack/transform/integration/TransformProgressIT.java @@ -160,7 +160,7 @@ public void assertGetProgress(int userWithMissingBuckets) throws Exception { null ); - Pivot pivot = new Pivot(pivotConfig, new SettingsConfig(), TransformConfigVersion.CURRENT, Collections.emptySet()); + Pivot pivot = new Pivot(pivotConfig, SettingsConfig.EMPTY, TransformConfigVersion.CURRENT, Collections.emptySet()); TransformProgress progress = getProgress(pivot, getProgressQuery(pivot, config.getSource().getIndex(), null)); @@ -188,7 +188,7 @@ public void assertGetProgress(int userWithMissingBuckets) throws Exception { Collections.singletonMap("every_50", new HistogramGroupSource("missing_field", null, missingBucket, 50.0)) ); pivotConfig = new PivotConfig(histgramGroupConfig, aggregationConfig, null); - pivot = new Pivot(pivotConfig, new SettingsConfig(), TransformConfigVersion.CURRENT, Collections.emptySet()); + pivot = new Pivot(pivotConfig, SettingsConfig.EMPTY, TransformConfigVersion.CURRENT, Collections.emptySet()); progress = getProgress( pivot, diff --git a/x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/action/TransportPreviewTransformAction.java b/x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/action/TransportPreviewTransformAction.java index 79644fac07579..f14ac9a534f28 100644 --- a/x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/action/TransportPreviewTransformAction.java +++ b/x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/action/TransportPreviewTransformAction.java @@ -52,6 +52,7 @@ import org.elasticsearch.xpack.core.transform.transforms.SyncConfig; import org.elasticsearch.xpack.core.transform.transforms.TransformConfig; import org.elasticsearch.xpack.core.transform.transforms.TransformDestIndexSettings; +import org.elasticsearch.xpack.core.transform.transforms.TransformEffectiveSettings; import org.elasticsearch.xpack.transform.TransformExtensionHolder; import org.elasticsearch.xpack.transform.persistence.TransformIndex; import org.elasticsearch.xpack.transform.transforms.Function; @@ -289,7 +290,7 @@ private void getPreview( }, listener::onFailure); ActionListener> deduceMappingsListener = ActionListener.wrap(deducedMappings -> { - if (Boolean.FALSE.equals(settingsConfig.getDeduceMappings())) { + if (TransformEffectiveSettings.isDeduceMappingsDisabled(settingsConfig)) { mappings.set(emptyMap()); } else { mappings.set(deducedMappings); diff --git a/x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/action/TransportStartTransformAction.java b/x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/action/TransportStartTransformAction.java index 825d0b8d12119..01359f351f07a 100644 --- a/x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/action/TransportStartTransformAction.java +++ b/x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/action/TransportStartTransformAction.java @@ -39,6 +39,7 @@ import org.elasticsearch.xpack.core.transform.action.ValidateTransformAction; import org.elasticsearch.xpack.core.transform.transforms.AuthorizationState; import org.elasticsearch.xpack.core.transform.transforms.TransformConfig; +import org.elasticsearch.xpack.core.transform.transforms.TransformEffectiveSettings; import org.elasticsearch.xpack.core.transform.transforms.TransformState; import org.elasticsearch.xpack.core.transform.transforms.TransformTaskParams; import org.elasticsearch.xpack.core.transform.transforms.TransformTaskState; @@ -187,7 +188,7 @@ protected void masterOperation( // <3> If the destination index exists, start the task, otherwise deduce our mappings for the destination index and create it ActionListener validationListener = ActionListener.wrap(validationResponse -> { - if (Boolean.TRUE.equals(transformConfigHolder.get().getSettings().getUnattended())) { + if (TransformEffectiveSettings.isUnattended(transformConfigHolder.get().getSettings())) { logger.debug( () -> format("[%s] Skip dest index creation as this is an unattended transform", transformConfigHolder.get().getId()) ); @@ -205,7 +206,7 @@ protected void masterOperation( createOrGetIndexListener ); }, e -> { - if (Boolean.TRUE.equals(transformConfigHolder.get().getSettings().getUnattended())) { + if (TransformEffectiveSettings.isUnattended(transformConfigHolder.get().getSettings())) { logger.debug( () -> format("[%s] Skip dest index creation as this is an unattended transform", transformConfigHolder.get().getId()) ); @@ -268,7 +269,7 @@ protected void masterOperation( ActionListener getTransformListener = ActionListener.wrap(config -> { transformConfigHolder.set(config); - if (Boolean.TRUE.equals(config.getSettings().getUnattended())) { + if (TransformEffectiveSettings.isUnattended(config.getSettings())) { // We do not fail the _start request of the unattended transform due to permission issues, // we just let it run fetchAuthStateListener.onResponse(null); diff --git a/x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/checkpoint/TimeBasedCheckpointProvider.java b/x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/checkpoint/TimeBasedCheckpointProvider.java index ec4cc2dcbcbf4..f49d5fc96f3ab 100644 --- a/x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/checkpoint/TimeBasedCheckpointProvider.java +++ b/x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/checkpoint/TimeBasedCheckpointProvider.java @@ -22,6 +22,7 @@ import org.elasticsearch.xpack.core.transform.transforms.TimeSyncConfig; import org.elasticsearch.xpack.core.transform.transforms.TransformCheckpoint; import org.elasticsearch.xpack.core.transform.transforms.TransformConfig; +import org.elasticsearch.xpack.core.transform.transforms.TransformEffectiveSettings; import org.elasticsearch.xpack.core.transform.transforms.pivot.DateHistogramGroupSource; import org.elasticsearch.xpack.core.transform.transforms.pivot.SingleGroupSource; import org.elasticsearch.xpack.transform.notifications.TransformAuditor; @@ -109,7 +110,7 @@ public void createNextCheckpoint(final TransformCheckpoint lastCheckpoint, final * @return function aligning the given timestamp with date histogram interval */ private static Function createAlignTimestampFunction(TransformConfig transformConfig) { - if (Boolean.FALSE.equals(transformConfig.getSettings().getAlignCheckpoints())) { + if (TransformEffectiveSettings.isAlignCheckpointsDisabled(transformConfig.getSettings())) { return identity(); } // In case of transforms created before aligning timestamp optimization was introduced we assume the default was "false". diff --git a/x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/persistence/TransformIndex.java b/x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/persistence/TransformIndex.java index fe3d4ede898bc..e3d9fa3aff671 100644 --- a/x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/persistence/TransformIndex.java +++ b/x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/persistence/TransformIndex.java @@ -33,6 +33,7 @@ import org.elasticsearch.xpack.core.transform.transforms.DestAlias; import org.elasticsearch.xpack.core.transform.transforms.TransformConfig; import org.elasticsearch.xpack.core.transform.transforms.TransformDestIndexSettings; +import org.elasticsearch.xpack.core.transform.transforms.TransformEffectiveSettings; import org.elasticsearch.xpack.transform.notifications.TransformAuditor; import java.time.Clock; @@ -128,7 +129,7 @@ public static void createDestinationIndex( // <2> Set up destination index aliases, regardless whether the destination index was created by the transform or by the user ActionListener createDestinationIndexListener = ActionListener.wrap(createdDestinationIndex -> { if (createdDestinationIndex) { - String message = Boolean.FALSE.equals(config.getSettings().getDeduceMappings()) + String message = TransformEffectiveSettings.isDeduceMappingsDisabled(config.getSettings()) ? "Created destination index [" + destinationIndex + "]." : "Created destination index [" + destinationIndex + "] with deduced mappings."; auditor.info(config.getId(), message); @@ -139,7 +140,7 @@ public static void createDestinationIndex( if (dest.length == 0) { TransformDestIndexSettings generatedDestIndexSettings = createTransformDestIndexSettings( destIndexSettings, - Boolean.FALSE.equals(config.getSettings().getDeduceMappings()) ? emptyMap() : destIndexMappings, + TransformEffectiveSettings.isDeduceMappingsDisabled(config.getSettings()) ? emptyMap() : destIndexMappings, config.getId(), Clock.systemUTC() ); diff --git a/x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/transforms/ClientTransformIndexer.java b/x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/transforms/ClientTransformIndexer.java index 1634f417924c0..c68c73fd71d9e 100644 --- a/x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/transforms/ClientTransformIndexer.java +++ b/x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/transforms/ClientTransformIndexer.java @@ -50,6 +50,7 @@ import org.elasticsearch.xpack.core.transform.transforms.SettingsConfig; import org.elasticsearch.xpack.core.transform.transforms.TransformCheckpoint; import org.elasticsearch.xpack.core.transform.transforms.TransformConfig; +import org.elasticsearch.xpack.core.transform.transforms.TransformEffectiveSettings; import org.elasticsearch.xpack.core.transform.transforms.TransformIndexerPosition; import org.elasticsearch.xpack.core.transform.transforms.TransformIndexerStats; import org.elasticsearch.xpack.core.transform.transforms.TransformProgress; @@ -131,17 +132,12 @@ class ClientTransformIndexer extends TransformIndexer { // TODO: move into context constructor context.setShouldStopAtCheckpoint(shouldStopAtCheckpoint); - if (transformConfig.getSettings().getUsePit() != null) { - disablePit = transformConfig.getSettings().getUsePit() == false; - } + disablePit = TransformEffectiveSettings.isPitDisabled(transformConfig.getSettings()); } @Override public void applyNewSettings(SettingsConfig newSettings) { - if (newSettings.getUsePit() != null) { - disablePit = newSettings.getUsePit() == false; - } - + disablePit = TransformEffectiveSettings.isPitDisabled(newSettings); super.applyNewSettings(newSettings); } diff --git a/x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/transforms/TransformFailureHandler.java b/x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/transforms/TransformFailureHandler.java index c7e0eda5ca5e6..337d3c5820c07 100644 --- a/x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/transforms/TransformFailureHandler.java +++ b/x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/transforms/TransformFailureHandler.java @@ -17,12 +17,11 @@ import org.elasticsearch.script.ScriptException; import org.elasticsearch.xpack.core.transform.TransformMessages; import org.elasticsearch.xpack.core.transform.transforms.SettingsConfig; +import org.elasticsearch.xpack.core.transform.transforms.TransformEffectiveSettings; import org.elasticsearch.xpack.core.transform.utils.ExceptionsHelper; import org.elasticsearch.xpack.transform.notifications.TransformAuditor; import org.elasticsearch.xpack.transform.utils.ExceptionRootCauseFinder; -import java.util.Optional; - import static org.elasticsearch.core.Strings.format; import static org.elasticsearch.xpack.core.common.notifications.Level.INFO; import static org.elasticsearch.xpack.core.common.notifications.Level.WARNING; @@ -59,32 +58,28 @@ void handleIndexerFailure(Exception exception, SettingsConfig settingsConfig) { // more detailed reporting in the handlers and below logger.atDebug().withThrowable(exception).log("[{}] transform encountered an exception", transformId); Throwable unwrappedException = ExceptionsHelper.findSearchExceptionRootCause(exception); - boolean unattended = Boolean.TRUE.equals(settingsConfig.getUnattended()); + boolean unattended = TransformEffectiveSettings.isUnattended(settingsConfig); + int numFailureRetries = TransformEffectiveSettings.getNumFailureRetries(settingsConfig, context.getNumFailureRetries()); if (unwrappedException instanceof CircuitBreakingException e) { handleCircuitBreakingException(e, unattended); } else if (unwrappedException instanceof ScriptException e) { handleScriptException(e, unattended); } else if (unwrappedException instanceof BulkIndexingException e) { - handleBulkIndexingException(e, unattended, getNumFailureRetries(settingsConfig)); + handleBulkIndexingException(e, unattended, numFailureRetries); } else if (unwrappedException instanceof ClusterBlockException e) { // gh#89802 always retry for a cluster block exception, because a cluster block should be temporary. - retry(e, e.getDetailedMessage(), unattended, getNumFailureRetries(settingsConfig)); + retry(e, e.getDetailedMessage(), unattended, numFailureRetries); } else if (unwrappedException instanceof SearchPhaseExecutionException e) { // The reason of a SearchPhaseExecutionException unfortunately contains a full stack trace. // Instead of displaying that to the user, get the cause's message instead. - retry(e, e.getCause() != null ? e.getCause().getMessage() : null, unattended, getNumFailureRetries(settingsConfig)); + retry(e, e.getCause() != null ? e.getCause().getMessage() : null, unattended, numFailureRetries); } else if (unwrappedException instanceof ElasticsearchException e) { - handleElasticsearchException(e, unattended, getNumFailureRetries(settingsConfig)); + handleElasticsearchException(e, unattended, numFailureRetries); } else if (unwrappedException instanceof IllegalArgumentException e) { handleIllegalArgumentException(e, unattended); } else { - retry( - unwrappedException, - ExceptionRootCauseFinder.getDetailedMessage(unwrappedException), - unattended, - getNumFailureRetries(settingsConfig) - ); + retry(unwrappedException, ExceptionRootCauseFinder.getDetailedMessage(unwrappedException), unattended, numFailureRetries); } } @@ -98,7 +93,7 @@ void handleIndexerFailure(Exception exception, SettingsConfig settingsConfig) { boolean handleStatePersistenceFailure(Exception e, SettingsConfig settingsConfig) { // we use the same setting for retries, however a separate counter, because the failure // counter for search/index gets reset after a successful bulk index request - int numFailureRetries = getNumFailureRetries(settingsConfig); + int numFailureRetries = TransformEffectiveSettings.getNumFailureRetries(settingsConfig, context.getNumFailureRetries()); int failureCount = context.incrementAndGetStatePersistenceFailureCount(e); @@ -273,19 +268,4 @@ private void fail(Throwable exception, String failureMessage) { // note: logging and audit is done as part of context.markAsFailed context.markAsFailed(exception, failureMessage); } - - /** - * Get the number of retries. - *

- * The number of retries are read from the config or if not read from the context which is based on a cluster wide - * default. If the transform runs in unattended mode, the number of retries is always indefinite. - * - * @param settingsConfig the setting config - * @return the number of retries or -1 if retries are indefinite - */ - private int getNumFailureRetries(SettingsConfig settingsConfig) { - return Boolean.TRUE.equals(settingsConfig.getUnattended()) - ? -1 - : Optional.ofNullable(settingsConfig.getNumFailureRetries()).orElse(context.getNumFailureRetries()); - } } diff --git a/x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/transforms/TransformIndexer.java b/x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/transforms/TransformIndexer.java index ff52f5e267655..38bd231e3e76a 100644 --- a/x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/transforms/TransformIndexer.java +++ b/x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/transforms/TransformIndexer.java @@ -36,6 +36,7 @@ import org.elasticsearch.xpack.core.transform.transforms.SettingsConfig; import org.elasticsearch.xpack.core.transform.transforms.TransformCheckpoint; import org.elasticsearch.xpack.core.transform.transforms.TransformConfig; +import org.elasticsearch.xpack.core.transform.transforms.TransformEffectiveSettings; import org.elasticsearch.xpack.core.transform.transforms.TransformIndexerPosition; import org.elasticsearch.xpack.core.transform.transforms.TransformIndexerStats; import org.elasticsearch.xpack.core.transform.transforms.TransformProgress; @@ -334,7 +335,7 @@ protected void onStart(long now, ActionListener listener) { }, listener::onFailure); var shouldMaybeCreateDestIndexForUnattended = context.getCheckpoint() == 0 - && Boolean.TRUE.equals(transformConfig.getSettings().getUnattended()); + && TransformEffectiveSettings.isUnattended(transformConfig.getSettings()); ActionListener> fieldMappingsListener = ActionListener.wrap(destIndexMappings -> { if (destIndexMappings.isEmpty() == false) { @@ -413,7 +414,7 @@ protected void onStart(long now, ActionListener listener) { hasSourceChanged = true; listener.onFailure(failure); })); - } else if (context.getCheckpoint() == 0 && Boolean.TRUE.equals(transformConfig.getSettings().getUnattended())) { + } else if (context.getCheckpoint() == 0 && TransformEffectiveSettings.isUnattended(transformConfig.getSettings())) { // this transform runs in unattended mode and has never run, to go on validate(changedSourceListener); } else { diff --git a/x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/transforms/pivot/Pivot.java b/x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/transforms/pivot/Pivot.java index 0d4dbcb6c2094..8c134b92c02af 100644 --- a/x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/transforms/pivot/Pivot.java +++ b/x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/transforms/pivot/Pivot.java @@ -27,6 +27,7 @@ import org.elasticsearch.xpack.core.transform.TransformMessages; import org.elasticsearch.xpack.core.transform.transforms.SettingsConfig; import org.elasticsearch.xpack.core.transform.transforms.SourceConfig; +import org.elasticsearch.xpack.core.transform.transforms.TransformEffectiveSettings; import org.elasticsearch.xpack.core.transform.transforms.TransformIndexerStats; import org.elasticsearch.xpack.core.transform.transforms.TransformProgress; import org.elasticsearch.xpack.core.transform.transforms.pivot.PivotConfig; @@ -132,14 +133,7 @@ protected Stream> extractResults( TransformIndexerStats transformIndexerStats, TransformProgress transformProgress ) { - // defines how dates are written, if not specified in settings - // < 7.11 as epoch millis - // >= 7.11 as string - // note: it depends on the version when the transform has been created, not the version of the code - boolean datesAsEpoch = settings.getDatesAsEpochMillis() != null - ? settings.getDatesAsEpochMillis() - : version.before(TransformConfigVersion.V_7_11_0); - + boolean datesAsEpoch = TransformEffectiveSettings.writeDatesAsEpochMillis(settings, version); return AggregationResultUtils.extractCompositeAggregationResults( agg, config.getGroupConfig(), diff --git a/x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/transforms/pivot/SchemaUtil.java b/x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/transforms/pivot/SchemaUtil.java index 48b156ce39fc2..d5e0351a8822e 100644 --- a/x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/transforms/pivot/SchemaUtil.java +++ b/x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/transforms/pivot/SchemaUtil.java @@ -24,6 +24,7 @@ import org.elasticsearch.xpack.core.ClientHelper; import org.elasticsearch.xpack.core.transform.transforms.SettingsConfig; import org.elasticsearch.xpack.core.transform.transforms.SourceConfig; +import org.elasticsearch.xpack.core.transform.transforms.TransformEffectiveSettings; import org.elasticsearch.xpack.core.transform.transforms.pivot.PivotConfig; import java.math.BigDecimal; @@ -167,7 +168,7 @@ public static void deduceMappings( sourceMappings -> listener.onResponse( resolveMappings( transformId, - Boolean.FALSE.equals(settingsConfig.getDeduceMappings()) == false, + TransformEffectiveSettings.isDeduceMappingsDisabled(settingsConfig), aggregationSourceFieldNames, aggregationTypes, fieldNamesForGrouping, @@ -207,7 +208,7 @@ public static void getDestinationFieldMappings( private static Map resolveMappings( String transformId, - boolean deduceMappings, + boolean deduceMappingsDisabled, Map aggregationSourceFieldNames, Map aggregationTypes, Map fieldNamesForGrouping, @@ -244,7 +245,7 @@ private static Map resolveMappings( targetMapping.put(targetFieldName, destinationMapping); } else { logger.log( - deduceMappings ? Level.WARN : Level.INFO, + deduceMappingsDisabled ? Level.INFO : Level.WARN, "[{}] Failed to deduce mapping for [{}], fall back to dynamic mapping. " + "Create the destination index with complete mappings first to avoid deducing the mappings", transformId, @@ -260,7 +261,7 @@ private static Map resolveMappings( targetMapping.put(targetFieldName, destinationMapping); } else { logger.log( - deduceMappings ? Level.WARN : Level.INFO, + deduceMappingsDisabled ? Level.INFO : Level.WARN, "[{}] Failed to deduce mapping for [{}], fall back to keyword. " + "Create the destination index with complete mappings first to avoid deducing the mappings", transformId, diff --git a/x-pack/plugin/transform/src/test/java/org/elasticsearch/xpack/transform/action/TransformConfigLinterTests.java b/x-pack/plugin/transform/src/test/java/org/elasticsearch/xpack/transform/action/TransformConfigLinterTests.java index 3006717bd843b..288ec8fc7a3d7 100644 --- a/x-pack/plugin/transform/src/test/java/org/elasticsearch/xpack/transform/action/TransformConfigLinterTests.java +++ b/x-pack/plugin/transform/src/test/java/org/elasticsearch/xpack/transform/action/TransformConfigLinterTests.java @@ -42,7 +42,7 @@ public void testGetWarnings_Pivot_WithScriptBasedRuntimeFields() { AggregationConfigTests.randomAggregationConfig(), null ); - Function function = new Pivot(pivotConfig, new SettingsConfig(), TransformConfigVersion.CURRENT, Collections.emptySet()); + Function function = new Pivot(pivotConfig, SettingsConfig.EMPTY, TransformConfigVersion.CURRENT, Collections.emptySet()); SourceConfig sourceConfig = SourceConfigTests.randomSourceConfig(); assertThat(TransformConfigLinter.getWarnings(function, sourceConfig, null), is(empty())); @@ -117,7 +117,7 @@ public void testGetWarnings_Pivot_CouldNotFindAnyOptimization() { AggregationConfigTests.randomAggregationConfig(), null ); - Function function = new Pivot(pivotConfig, new SettingsConfig(), TransformConfigVersion.CURRENT, Collections.emptySet()); + Function function = new Pivot(pivotConfig, SettingsConfig.EMPTY, TransformConfigVersion.CURRENT, Collections.emptySet()); SourceConfig sourceConfig = SourceConfigTests.randomSourceConfig(); SyncConfig syncConfig = TimeSyncConfigTests.randomTimeSyncConfig(); assertThat( diff --git a/x-pack/plugin/transform/src/test/java/org/elasticsearch/xpack/transform/transforms/ClientTransformIndexerTests.java b/x-pack/plugin/transform/src/test/java/org/elasticsearch/xpack/transform/transforms/ClientTransformIndexerTests.java index 43a8f35cfeafe..017fe3d289b0c 100644 --- a/x-pack/plugin/transform/src/test/java/org/elasticsearch/xpack/transform/transforms/ClientTransformIndexerTests.java +++ b/x-pack/plugin/transform/src/test/java/org/elasticsearch/xpack/transform/transforms/ClientTransformIndexerTests.java @@ -47,6 +47,7 @@ import org.elasticsearch.xpack.core.transform.transforms.TransformCheckpoint; import org.elasticsearch.xpack.core.transform.transforms.TransformConfig; import org.elasticsearch.xpack.core.transform.transforms.TransformConfigTests; +import org.elasticsearch.xpack.core.transform.transforms.TransformEffectiveSettings; import org.elasticsearch.xpack.core.transform.transforms.TransformIndexerPosition; import org.elasticsearch.xpack.core.transform.transforms.TransformIndexerStats; import org.elasticsearch.xpack.core.transform.transforms.TransformProgress; @@ -309,7 +310,7 @@ public void testDisablePit() throws InterruptedException { } TransformConfig config = configBuilder.build(); - boolean pitEnabled = config.getSettings().getUsePit() == null || config.getSettings().getUsePit(); + boolean pitEnabled = TransformEffectiveSettings.isPitDisabled(config.getSettings()) == false; try (var threadPool = createThreadPool()) { final var client = new PitMockClient(threadPool, true); diff --git a/x-pack/plugin/transform/src/test/java/org/elasticsearch/xpack/transform/transforms/pivot/AggregationSchemaAndResultTests.java b/x-pack/plugin/transform/src/test/java/org/elasticsearch/xpack/transform/transforms/pivot/AggregationSchemaAndResultTests.java index 5943a9007fb7c..1eb86b813f260 100644 --- a/x-pack/plugin/transform/src/test/java/org/elasticsearch/xpack/transform/transforms/pivot/AggregationSchemaAndResultTests.java +++ b/x-pack/plugin/transform/src/test/java/org/elasticsearch/xpack/transform/transforms/pivot/AggregationSchemaAndResultTests.java @@ -153,7 +153,7 @@ public void testBasic() throws InterruptedException { client, emptyMap(), "my-transform", - new SettingsConfig(), + SettingsConfig.EMPTY, pivotConfig, new SourceConfig(new String[] { "source-index" }), listener @@ -233,7 +233,7 @@ public void testNested() throws InterruptedException { client, emptyMap(), "my-transform", - new SettingsConfig(), + SettingsConfig.EMPTY, pivotConfig, new SourceConfig(new String[] { "source-index" }), listener diff --git a/x-pack/plugin/transform/src/test/java/org/elasticsearch/xpack/transform/transforms/pivot/PivotTests.java b/x-pack/plugin/transform/src/test/java/org/elasticsearch/xpack/transform/transforms/pivot/PivotTests.java index 5d58ac9904482..0a030d26016f7 100644 --- a/x-pack/plugin/transform/src/test/java/org/elasticsearch/xpack/transform/transforms/pivot/PivotTests.java +++ b/x-pack/plugin/transform/src/test/java/org/elasticsearch/xpack/transform/transforms/pivot/PivotTests.java @@ -125,14 +125,14 @@ protected NamedXContentRegistry xContentRegistry() { public void testValidateExistingIndex() throws Exception { SourceConfig source = new SourceConfig("existing_source_index"); - Function pivot = new Pivot(getValidPivotConfig(), new SettingsConfig(), TransformConfigVersion.CURRENT, Collections.emptySet()); + Function pivot = new Pivot(getValidPivotConfig(), SettingsConfig.EMPTY, TransformConfigVersion.CURRENT, Collections.emptySet()); assertValidTransform(client, source, pivot); } public void testValidateNonExistingIndex() throws Exception { SourceConfig source = new SourceConfig("non_existing_source_index"); - Function pivot = new Pivot(getValidPivotConfig(), new SettingsConfig(), TransformConfigVersion.CURRENT, Collections.emptySet()); + Function pivot = new Pivot(getValidPivotConfig(), SettingsConfig.EMPTY, TransformConfigVersion.CURRENT, Collections.emptySet()); assertInvalidTransform(client, source, pivot); } @@ -142,7 +142,7 @@ public void testInitialPageSize() throws Exception { Function pivot = new Pivot( new PivotConfig(GroupConfigTests.randomGroupConfig(), getValidAggregationConfig(), expectedPageSize), - new SettingsConfig(), + SettingsConfig.EMPTY, TransformConfigVersion.CURRENT, Collections.emptySet() ); @@ -150,7 +150,7 @@ public void testInitialPageSize() throws Exception { pivot = new Pivot( new PivotConfig(GroupConfigTests.randomGroupConfig(), getValidAggregationConfig(), null), - new SettingsConfig(), + SettingsConfig.EMPTY, TransformConfigVersion.CURRENT, Collections.emptySet() ); @@ -164,7 +164,7 @@ public void testSearchFailure() throws Exception { // search has failures although they might just be temporary SourceConfig source = new SourceConfig("existing_source_index_with_failing_shards"); - Function pivot = new Pivot(getValidPivotConfig(), new SettingsConfig(), TransformConfigVersion.CURRENT, Collections.emptySet()); + Function pivot = new Pivot(getValidPivotConfig(), SettingsConfig.EMPTY, TransformConfigVersion.CURRENT, Collections.emptySet()); assertInvalidTransform(client, source, pivot); } @@ -177,7 +177,7 @@ public void testValidateAllSupportedAggregations() throws Exception { Function pivot = new Pivot( getValidPivotConfig(aggregationConfig), - new SettingsConfig(), + SettingsConfig.EMPTY, TransformConfigVersion.CURRENT, Collections.emptySet() ); @@ -191,7 +191,7 @@ public void testValidateAllUnsupportedAggregations() throws Exception { Function pivot = new Pivot( getValidPivotConfig(aggregationConfig), - new SettingsConfig(), + SettingsConfig.EMPTY, TransformConfigVersion.CURRENT, Collections.emptySet() ); @@ -233,7 +233,7 @@ public void testGetPerformanceCriticalFields() throws IOException { assertThat(groupConfig.validate(null), is(nullValue())); PivotConfig pivotConfig = new PivotConfig(groupConfig, AggregationConfigTests.randomAggregationConfig(), null); - Function pivot = new Pivot(pivotConfig, new SettingsConfig(), TransformConfigVersion.CURRENT, Collections.emptySet()); + Function pivot = new Pivot(pivotConfig, SettingsConfig.EMPTY, TransformConfigVersion.CURRENT, Collections.emptySet()); assertThat(pivot.getPerformanceCriticalFields(), contains("field-A", "field-B", "field-C")); } From 099a5a9b923bd67b584606c7554978fed6e79fa0 Mon Sep 17 00:00:00 2001 From: Jedr Blaszyk Date: Wed, 6 Mar 2024 09:38:42 +0100 Subject: [PATCH 13/27] [Connector API] Fix default ordering in SyncJob list endpoint (#105945) --- docs/changelog/105945.yaml | 5 ++++ .../entsearch/470_connector_sync_job_list.yml | 19 ++++++++------ .../syncjob/ConnectorSyncJobIndexService.java | 4 +-- .../ConnectorSyncJobIndexServiceTests.java | 26 ++++++++++--------- 4 files changed, 32 insertions(+), 22 deletions(-) create mode 100644 docs/changelog/105945.yaml diff --git a/docs/changelog/105945.yaml b/docs/changelog/105945.yaml new file mode 100644 index 0000000000000..ec76faf6ef76f --- /dev/null +++ b/docs/changelog/105945.yaml @@ -0,0 +1,5 @@ +pr: 105945 +summary: "[Connector API] Fix default ordering in `SyncJob` list endpoint" +area: Application +type: bug +issues: [] diff --git a/x-pack/plugin/ent-search/qa/rest/src/yamlRestTest/resources/rest-api-spec/test/entsearch/470_connector_sync_job_list.yml b/x-pack/plugin/ent-search/qa/rest/src/yamlRestTest/resources/rest-api-spec/test/entsearch/470_connector_sync_job_list.yml index 8d23850f49840..82d9a18bb51e9 100644 --- a/x-pack/plugin/ent-search/qa/rest/src/yamlRestTest/resources/rest-api-spec/test/entsearch/470_connector_sync_job_list.yml +++ b/x-pack/plugin/ent-search/qa/rest/src/yamlRestTest/resources/rest-api-spec/test/entsearch/470_connector_sync_job_list.yml @@ -50,10 +50,10 @@ setup: - match: { count: 3 } - # Ascending order by creation_date for results - - match: { results.0.id: $sync-job-one-id } + # Descending order by creation_date for results + - match: { results.0.id: $sync-job-three-id } - match: { results.1.id: $sync-job-two-id } - - match: { results.2.id: $sync-job-three-id } + - match: { results.2.id: $sync-job-one-id } --- "List Connector Sync Jobs - with from": @@ -84,9 +84,9 @@ setup: - match: { count: 3 } - # Ascending order by creation_date for results + # Descending order by creation_date for results - match: { results.0.id: $sync-job-two-id } - - match: { results.1.id: $sync-job-three-id } + - match: { results.1.id: $sync-job-one-id } --- "List Connector Sync Jobs - with size": @@ -117,7 +117,8 @@ setup: - match: { count: 3 } - - match: { results.0.id: $sync-job-one-id } + # Descending order by creation_date for results + - match: { results.0.id: $sync-job-three-id } --- "List Connector Sync Jobs - Get pending jobs": @@ -216,9 +217,11 @@ setup: connector_sync_job.list: connector_id: connector-one job_type: full,incremental + + # Descending order by creation_date for results - match: { count: 2 } - - match: { results.0.id: $sync-job-one-id } - - match: { results.1.id: $sync-job-two-id } + - match: { results.0.id: $sync-job-two-id } + - match: { results.1.id: $sync-job-one-id } --- "List Connector Sync Jobs - with invalid job type": diff --git a/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/application/connector/syncjob/ConnectorSyncJobIndexService.java b/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/application/connector/syncjob/ConnectorSyncJobIndexService.java index 3ac598fd58ee8..d1d345840874f 100644 --- a/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/application/connector/syncjob/ConnectorSyncJobIndexService.java +++ b/x-pack/plugin/ent-search/src/main/java/org/elasticsearch/xpack/application/connector/syncjob/ConnectorSyncJobIndexService.java @@ -293,7 +293,7 @@ public void cancelConnectorSyncJob(String connectorSyncJobId, ActionListener Date: Wed, 6 Mar 2024 09:25:52 +0000 Subject: [PATCH 14/27] Make sure we test the listener is called (#105914) --- .../org/elasticsearch/xpack/core/ilm/DeleteStepTests.java | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ilm/DeleteStepTests.java b/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ilm/DeleteStepTests.java index 5851ebe2fb3c9..7445e82da3ecf 100644 --- a/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ilm/DeleteStepTests.java +++ b/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ilm/DeleteStepTests.java @@ -20,6 +20,7 @@ import org.mockito.Mockito; import java.util.List; +import java.util.concurrent.atomic.AtomicBoolean; import static org.elasticsearch.test.ActionListenerUtils.anyActionListener; import static org.hamcrest.Matchers.is; @@ -158,14 +159,17 @@ public void testPerformActionCallsFailureListenerIfIndexIsTheDataStreamWriteInde .metadata(Metadata.builder().put(index1, false).put(sourceIndexMetadata, false).put(dataStream).build()) .build(); + AtomicBoolean listenerCalled = new AtomicBoolean(false); createRandomInstance().performDuringNoSnapshot(sourceIndexMetadata, clusterState, new ActionListener<>() { @Override public void onResponse(Void complete) { + listenerCalled.set(true); fail("unexpected listener callback"); } @Override public void onFailure(Exception e) { + listenerCalled.set(true); assertThat( e.getMessage(), is( @@ -180,5 +184,7 @@ public void onFailure(Exception e) { ); } }); + + assertThat(listenerCalled.get(), is(true)); } } From c70956ac167908c111df33ccb845d5601f99cebc Mon Sep 17 00:00:00 2001 From: Jan Kuipers <148754765+jan-elastic@users.noreply.github.com> Date: Wed, 6 Mar 2024 11:03:42 +0100 Subject: [PATCH 15/27] [text structure] Find field and message structure endpoints (#105660) * Extract AbstractFindStructureRequest * Extract FindStructureResponse * Extract RestFindStructureRequestParser * FindFieldStructure endpoint * FindMessageStructure endpoint * Improve FindTextStructureResponseTests * REST API spec + YAML REST tests * Lint fixes * Remove POST find_field_structure * Update docs/changelog/105660.yaml * Update changelog * Fix text_structure.find_field_structure.json * Fix find_field_structure yaml rest test * Fix FindTextStructureResponseTests * Fix YAML tests with security * Remove unreachable code * DelimitedTextStructureFinder::createFromMessages * NdJsonTextStructureFinderFactory::createFromMessages * XmlTextStructureFinderFactory::createFromMessages * LogTextStructureFinderFactory::createFromMessages * Lint fixes * Add createFromMessages to TextStructureFinderFactory interface * Wire createFromMessages in the endpoints * Uppercase UTF-8 * REST test for semi-structured messages * Restrict query params to applicable endpoints * typo * Polish thread scheduling * Propagate parent task in search request * No header row for find message/field structure * Expose findTextStructure more consistently * Move text structure query params to shared doc * Rename "find structure API" -> "find text structure API" * Find message structure API docs * Find field structure docs * Maybe fix docs error? * bugfix * Fix docs? * Fix find-field-structure test from docs * Improve docs * Add param documents_to_sample to docs * improve docs --- docs/changelog/105660.yaml | 5 + .../apis/find-field-structure.asciidoc | 316 ++++++++++++++ .../apis/find-message-structure.asciidoc | 292 +++++++++++++ .../apis/find-structure-shared.asciidoc | 215 ++++++++++ .../apis/find-structure.asciidoc | 201 +-------- .../text-structure/apis/index.asciidoc | 4 + .../text_structure.find_field_structure.json | 90 ++++ ...text_structure.find_message_structure.json | 80 ++++ .../action/AbstractFindStructureRequest.java | 377 +++++++++++++++++ .../action/FindFieldStructureAction.java | 98 +++++ .../action/FindMessageStructureAction.java | 97 +++++ .../action/FindStructureAction.java | 389 +----------------- .../action/FindStructureResponse.java | 61 +++ .../FindTextStructureActionResponseTests.java | 29 -- .../FindTextStructureResponseTests.java | 33 ++ .../xpack/security/operator/Constants.java | 2 + .../text_structure/find_field_structure.yml | 63 +++ .../text_structure/find_message_structure.yml | 56 +++ .../text-structure-with-security/build.gradle | 2 +- .../qa/text-structure-with-security/roles.yml | 12 + .../textstructure/TextStructurePlugin.java | 15 +- .../rest/RestFindFieldStructureAction.java | 51 +++ .../rest/RestFindMessageStructureAction.java | 55 +++ .../rest/RestFindStructureAction.java | 38 +- .../RestFindStructureArgumentsParser.java | 73 ++++ .../DelimitedTextStructureFinder.java | 54 +-- .../DelimitedTextStructureFinderFactory.java | 40 +- .../LogTextStructureFinder.java | 141 ++++++- .../LogTextStructureFinderFactory.java | 13 + .../NdJsonTextStructureFinderFactory.java | 23 ++ .../TextStructureFinderFactory.java | 9 + .../TextStructureFinderManager.java | 91 +++- .../TextStructureOverrides.java | 3 +- .../XmlTextStructureFinderFactory.java | 64 ++- .../TransportFindFieldStructureAction.java | 94 +++++ .../TransportFindMessageStructureAction.java | 56 +++ .../TransportFindStructureAction.java | 38 +- ...imitedTextStructureFinderFactoryTests.java | 18 + .../DelimitedTextStructureFinderTests.java | 24 ++ .../LogTextStructureFinderTests.java | 16 + ...NdJsonTextStructureFinderFactoryTests.java | 18 + .../XmlTextStructureFinderFactoryTests.java | 18 + 42 files changed, 2640 insertions(+), 734 deletions(-) create mode 100644 docs/changelog/105660.yaml create mode 100644 docs/reference/text-structure/apis/find-field-structure.asciidoc create mode 100644 docs/reference/text-structure/apis/find-message-structure.asciidoc create mode 100644 docs/reference/text-structure/apis/find-structure-shared.asciidoc create mode 100644 rest-api-spec/src/main/resources/rest-api-spec/api/text_structure.find_field_structure.json create mode 100644 rest-api-spec/src/main/resources/rest-api-spec/api/text_structure.find_message_structure.json create mode 100644 x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/textstructure/action/AbstractFindStructureRequest.java create mode 100644 x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/textstructure/action/FindFieldStructureAction.java create mode 100644 x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/textstructure/action/FindMessageStructureAction.java create mode 100644 x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/textstructure/action/FindStructureResponse.java delete mode 100644 x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/textstructure/action/FindTextStructureActionResponseTests.java create mode 100644 x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/textstructure/action/FindTextStructureResponseTests.java create mode 100644 x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/text_structure/find_field_structure.yml create mode 100644 x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/text_structure/find_message_structure.yml create mode 100644 x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/rest/RestFindFieldStructureAction.java create mode 100644 x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/rest/RestFindMessageStructureAction.java create mode 100644 x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/rest/RestFindStructureArgumentsParser.java create mode 100644 x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/transport/TransportFindFieldStructureAction.java create mode 100644 x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/transport/TransportFindMessageStructureAction.java diff --git a/docs/changelog/105660.yaml b/docs/changelog/105660.yaml new file mode 100644 index 0000000000000..1b30a25417906 --- /dev/null +++ b/docs/changelog/105660.yaml @@ -0,0 +1,5 @@ +pr: 105660 +summary: "Text structure endpoints to determine the structure of a list of messages and of an indexed field" +area: Machine Learning +type: feature +issues: [] diff --git a/docs/reference/text-structure/apis/find-field-structure.asciidoc b/docs/reference/text-structure/apis/find-field-structure.asciidoc new file mode 100644 index 0000000000000..6788ddf7f42be --- /dev/null +++ b/docs/reference/text-structure/apis/find-field-structure.asciidoc @@ -0,0 +1,316 @@ +[role="xpack"] +[[find-field-structure]] += Find field structure API + +Finds the structure of a field in an Elasticsearch index. + +[discrete] +[[find-field-structure-request]] +== {api-request-title} + +`GET _text_structure/find_field_structure` + +[discrete] +[[find-field-structure-prereqs]] +== {api-prereq-title} + +* If the {es} {security-features} are enabled, you must have `monitor_text_structure` or +`monitor` cluster privileges to use this API. See +<>. + +[discrete] +[[find-field-structure-desc]] +== {api-description-title} + +This API provides a starting point for extracting further information from log messages +already ingested into {es}. For example, if you have ingested data into a very simple +index that has just `@timestamp` and `message` fields, you can use this API to +see what common structure exists in the `message` field. + +The response from the API contains: + +* Sample messages. +* Statistics that reveal the most common values for all fields detected within +the text and basic numeric statistics for numeric fields. +* Information about the structure of the text, which is useful when you write +ingest configurations to index it or similarly formatted text. +* Appropriate mappings for an {es} index, which you could use to ingest the text. + +All this information can be calculated by the structure finder with no guidance. +However, you can optionally override some of the decisions about the text +structure by specifying one or more query parameters. + +Details of the output can be seen in the <>. + +If the structure finder produces unexpected results, +specify the `explain` query parameter and an `explanation` will appear in +the response. It helps determine why the returned structure was +chosen. + +[discrete] +[[find-field-structure-query-parms]] +== {api-query-parms-title} + +`index`:: +(Required, string) The name of the index containing the field. + +`field`:: +(Required, string) The name of the field that's analyzed. + +include::{es-repo-dir}/text-structure/apis/find-structure-shared.asciidoc[tag=param-column-names] +include::{es-repo-dir}/text-structure/apis/find-structure-shared.asciidoc[tag=param-delimiter] + +`documents_to_sample`:: +(Optional, unsigned integer) The number of documents to include in the structural +analysis. The minimum is 2; the default is 1000. + +include::{es-repo-dir}/text-structure/apis/find-structure-shared.asciidoc[tag=param-explain] +include::{es-repo-dir}/text-structure/apis/find-structure-shared.asciidoc[tag=param-format] +include::{es-repo-dir}/text-structure/apis/find-structure-shared.asciidoc[tag=param-grok-pattern] +include::{es-repo-dir}/text-structure/apis/find-structure-shared.asciidoc[tag=param-ecs-compatibility] +include::{es-repo-dir}/text-structure/apis/find-structure-shared.asciidoc[tag=param-quote] +include::{es-repo-dir}/text-structure/apis/find-structure-shared.asciidoc[tag=param-should-trim-fields] +include::{es-repo-dir}/text-structure/apis/find-structure-shared.asciidoc[tag=param-timeout] +include::{es-repo-dir}/text-structure/apis/find-structure-shared.asciidoc[tag=param-timestamp-field] +include::{es-repo-dir}/text-structure/apis/find-structure-shared.asciidoc[tag=param-timestamp-format] + +[discrete] +[[find-field-structure-examples]] +== {api-examples-title} + +[discrete] +[[find-field-structure-example]] +=== Analyzing Elasticsearch log files + +Suppose you have a list of {es} log messages in an index. +You can analyze them with the `find_field_structure` endpoint as follows: + +[source,console] +---- +POST _bulk?refresh=true +{"index":{"_index":"test-logs"}} +{"message":"[2024-03-05T10:52:36,256][INFO ][o.a.l.u.VectorUtilPanamaProvider] [laptop] Java vector incubator API enabled; uses preferredBitSize=128"} +{"index":{"_index":"test-logs"}} +{"message":"[2024-03-05T10:52:41,038][INFO ][o.e.p.PluginsService ] [laptop] loaded module [repository-url]"} +{"index":{"_index":"test-logs"}} +{"message":"[2024-03-05T10:52:41,042][INFO ][o.e.p.PluginsService ] [laptop] loaded module [rest-root]"} +{"index":{"_index":"test-logs"}} +{"message":"[2024-03-05T10:52:41,043][INFO ][o.e.p.PluginsService ] [laptop] loaded module [x-pack-core]"} +{"index":{"_index":"test-logs"}} +{"message":"[2024-03-05T10:52:41,043][INFO ][o.e.p.PluginsService ] [laptop] loaded module [x-pack-redact]"} +{"index":{"_index":"test-logs"}} +{"message":"[2024-03-05T10:52:41,043][INFO ][o.e.p.PluginsService ] [laptop] loaded module [ingest-user-agent]"} +{"index":{"_index":"test-logs"}} +{"message":"[2024-03-05T10:52:41,044][INFO ][o.e.p.PluginsService ] [laptop] loaded module [x-pack-monitoring]"} +{"index":{"_index":"test-logs"}} +{"message":"[2024-03-05T10:52:41,044][INFO ][o.e.p.PluginsService ] [laptop] loaded module [repository-s3]"} +{"index":{"_index":"test-logs"}} +{"message":"[2024-03-05T10:52:41,044][INFO ][o.e.p.PluginsService ] [laptop] loaded module [x-pack-analytics]"} +{"index":{"_index":"test-logs"}} +{"message":"[2024-03-05T10:52:41,044][INFO ][o.e.p.PluginsService ] [laptop] loaded module [x-pack-ent-search]"} +{"index":{"_index":"test-logs"}} +{"message":"[2024-03-05T10:52:41,044][INFO ][o.e.p.PluginsService ] [laptop] loaded module [x-pack-autoscaling]"} +{"index":{"_index":"test-logs"}} +{"message":"[2024-03-05T10:52:41,044][INFO ][o.e.p.PluginsService ] [laptop] loaded module [lang-painless]]"} +{"index":{"_index":"test-logs"}} +{"message":"[2024-03-05T10:52:41,059][INFO ][o.e.p.PluginsService ] [laptop] loaded module [lang-expression]"} +{"index":{"_index":"test-logs"}} +{"message":"[2024-03-05T10:52:41,059][INFO ][o.e.p.PluginsService ] [laptop] loaded module [x-pack-eql]"} +{"index":{"_index":"test-logs"}} +{"message":"[2024-03-05T10:52:43,291][INFO ][o.e.e.NodeEnvironment ] [laptop] heap size [16gb], compressed ordinary object pointers [true]"} +{"index":{"_index":"test-logs"}} +{"message":"[2024-03-05T10:52:46,098][INFO ][o.e.x.s.Security ] [laptop] Security is enabled"} +{"index":{"_index":"test-logs"}} +{"message":"[2024-03-05T10:52:47,227][INFO ][o.e.x.p.ProfilingPlugin ] [laptop] Profiling is enabled"} +{"index":{"_index":"test-logs"}} +{"message":"[2024-03-05T10:52:47,259][INFO ][o.e.x.p.ProfilingPlugin ] [laptop] profiling index templates will not be installed or reinstalled"} +{"index":{"_index":"test-logs"}} +{"message":"[2024-03-05T10:52:47,755][INFO ][o.e.i.r.RecoverySettings ] [laptop] using rate limit [40mb] with [default=40mb, read=0b, write=0b, max=0b]"} +{"index":{"_index":"test-logs"}} +{"message":"[2024-03-05T10:52:47,787][INFO ][o.e.d.DiscoveryModule ] [laptop] using discovery type [multi-node] and seed hosts providers [settings]"} +{"index":{"_index":"test-logs"}} +{"message":"[2024-03-05T10:52:49,188][INFO ][o.e.n.Node ] [laptop] initialized"} +{"index":{"_index":"test-logs"}} +{"message":"[2024-03-05T10:52:49,199][INFO ][o.e.n.Node ] [laptop] starting ..."} + +GET _text_structure/find_field_structure?index=test-logs&field=message +---- +// TEST + +If the request does not encounter errors, you receive the following result: + +[source,console-result] +---- +{ + "num_lines_analyzed" : 22, + "num_messages_analyzed" : 22, + "sample_start" : "[2024-03-05T10:52:36,256][INFO ][o.a.l.u.VectorUtilPanamaProvider] [laptop] Java vector incubator API enabled; uses preferredBitSize=128\n[2024-03-05T10:52:41,038][INFO ][o.e.p.PluginsService ] [laptop] loaded module [repository-url]\n", <3> + "charset" : "UTF-8", + "format" : "semi_structured_text", + "multiline_start_pattern" : "^\\[\\b\\d{4}-\\d{2}-\\d{2}[T ]\\d{2}:\\d{2}", + "grok_pattern" : "\\[%{TIMESTAMP_ISO8601:timestamp}\\]\\[%{LOGLEVEL:loglevel} \\]\\[.*", + "ecs_compatibility" : "disabled", + "timestamp_field" : "timestamp", + "joda_timestamp_formats" : [ + "ISO8601" + ], + "java_timestamp_formats" : [ + "ISO8601" + ], + "need_client_timezone" : true, + "mappings" : { + "properties" : { + "@timestamp" : { + "type" : "date" + }, + "loglevel" : { + "type" : "keyword" + }, + "message" : { + "type" : "text" + } + } + }, + "ingest_pipeline" : { + "description" : "Ingest pipeline created by text structure finder", + "processors" : [ + { + "grok" : { + "field" : "message", + "patterns" : [ + "\\[%{TIMESTAMP_ISO8601:timestamp}\\]\\[%{LOGLEVEL:loglevel} \\]\\[.*" + ], + "ecs_compatibility" : "disabled" + } + }, + { + "date" : { + "field" : "timestamp", + "timezone" : "{{ event.timezone }}", + "formats" : [ + "ISO8601" + ] + } + }, + { + "remove" : { + "field" : "timestamp" + } + } + ] + }, + "field_stats" : { + "loglevel" : { + "count" : 22, + "cardinality" : 1, + "top_hits" : [ + { + "value" : "INFO", + "count" : 22 + } + ] + }, + "message" : { + "count" : 22, + "cardinality" : 22, + "top_hits" : [ + { + "value" : "[2024-03-05T10:52:36,256][INFO ][o.a.l.u.VectorUtilPanamaProvider] [laptop] Java vector incubator API enabled; uses preferredBitSize=128", + "count" : 1 + }, + { + "value" : "[2024-03-05T10:52:41,038][INFO ][o.e.p.PluginsService ] [laptop] loaded module [repository-url]", + "count" : 1 + }, + { + "value" : "[2024-03-05T10:52:41,042][INFO ][o.e.p.PluginsService ] [laptop] loaded module [rest-root]", + "count" : 1 + }, + { + "value" : "[2024-03-05T10:52:41,043][INFO ][o.e.p.PluginsService ] [laptop] loaded module [ingest-user-agent]", + "count" : 1 + }, + { + "value" : "[2024-03-05T10:52:41,043][INFO ][o.e.p.PluginsService ] [laptop] loaded module [x-pack-core]", + "count" : 1 + }, + { + "value" : "[2024-03-05T10:52:41,043][INFO ][o.e.p.PluginsService ] [laptop] loaded module [x-pack-redact]", + "count" : 1 + }, + { + "value" : "[2024-03-05T10:52:41,044][INFO ][o.e.p.PluginsService ] [laptop] loaded module [lang-painless]]", + "count" : 1 + }, + { + "value" : "[2024-03-05T10:52:41,044][INFO ][o.e.p.PluginsService ] [laptop] loaded module [repository-s3]", + "count" : 1 + }, + { + "value" : "[2024-03-05T10:52:41,044][INFO ][o.e.p.PluginsService ] [laptop] loaded module [x-pack-analytics]", + "count" : 1 + }, + { + "value" : "[2024-03-05T10:52:41,044][INFO ][o.e.p.PluginsService ] [laptop] loaded module [x-pack-autoscaling]", + "count" : 1 + } + ] + }, + "timestamp" : { + "count" : 22, + "cardinality" : 14, + "earliest" : "2024-03-05T10:52:36,256", + "latest" : "2024-03-05T10:52:49,199", + "top_hits" : [ + { + "value" : "2024-03-05T10:52:41,044", + "count" : 6 + }, + { + "value" : "2024-03-05T10:52:41,043", + "count" : 3 + }, + { + "value" : "2024-03-05T10:52:41,059", + "count" : 2 + }, + { + "value" : "2024-03-05T10:52:36,256", + "count" : 1 + }, + { + "value" : "2024-03-05T10:52:41,038", + "count" : 1 + }, + { + "value" : "2024-03-05T10:52:41,042", + "count" : 1 + }, + { + "value" : "2024-03-05T10:52:43,291", + "count" : 1 + }, + { + "value" : "2024-03-05T10:52:46,098", + "count" : 1 + }, + { + "value" : "2024-03-05T10:52:47,227", + "count" : 1 + }, + { + "value" : "2024-03-05T10:52:47,259", + "count" : 1 + } + ] + } + } +} +---- +// TESTRESPONSE[s/"sample_start" : ".*",/"sample_start" : "$body.sample_start",/] +// The substitution is because the text is pre-processed by the test harness, +// so the fields may get reordered in the JSON the endpoint sees + +For a detailed description of the response format, or for additional examples +on ingesting delimited text (such as CSV) or newline-delimited JSON, refer to the +<>. diff --git a/docs/reference/text-structure/apis/find-message-structure.asciidoc b/docs/reference/text-structure/apis/find-message-structure.asciidoc new file mode 100644 index 0000000000000..085f65b852126 --- /dev/null +++ b/docs/reference/text-structure/apis/find-message-structure.asciidoc @@ -0,0 +1,292 @@ +[role="xpack"] +[[find-message-structure]] += Find messages structure API + +Finds the structure of a list of text messages. + +[discrete] +[[find-message-structure-request]] +== {api-request-title} + +`GET _text_structure/find_message_structure` + +`POST _text_structure/find_message_structure` + +[discrete] +[[find-message-structure-prereqs]] +== {api-prereq-title} + +* If the {es} {security-features} are enabled, you must have `monitor_text_structure` or +`monitor` cluster privileges to use this API. See +<>. + +[discrete] +[[find-message-structure-desc]] +== {api-description-title} + +This API provides a starting point for ingesting data into {es} in a format that +is suitable for subsequent use with other {stack} functionality. Use this +API in preference to `find_structure` when your input text has already been +split up into separate messages by some other process. + +The response from the API contains: + +* Sample messages. +* Statistics that reveal the most common values for all fields detected within +the text and basic numeric statistics for numeric fields. +* Information about the structure of the text, which is useful when you write +ingest configurations to index it or similarly formatted text. +* Appropriate mappings for an {es} index, which you could use to ingest the text. + +All this information can be calculated by the structure finder with no guidance. +However, you can optionally override some of the decisions about the text +structure by specifying one or more query parameters. + +Details of the output can be seen in the <>. + +If the structure finder produces unexpected results, +specify the `explain` query parameter and an `explanation` will appear in +the response. It helps determine why the returned structure was +chosen. + +[discrete] +[[find-message-structure-query-parms]] +== {api-query-parms-title} + +include::{es-repo-dir}/text-structure/apis/find-structure-shared.asciidoc[tag=param-column-names] +include::{es-repo-dir}/text-structure/apis/find-structure-shared.asciidoc[tag=param-delimiter] +include::{es-repo-dir}/text-structure/apis/find-structure-shared.asciidoc[tag=param-explain] +include::{es-repo-dir}/text-structure/apis/find-structure-shared.asciidoc[tag=param-format] +include::{es-repo-dir}/text-structure/apis/find-structure-shared.asciidoc[tag=param-grok-pattern] +include::{es-repo-dir}/text-structure/apis/find-structure-shared.asciidoc[tag=param-ecs-compatibility] +include::{es-repo-dir}/text-structure/apis/find-structure-shared.asciidoc[tag=param-quote] +include::{es-repo-dir}/text-structure/apis/find-structure-shared.asciidoc[tag=param-should-trim-fields] +include::{es-repo-dir}/text-structure/apis/find-structure-shared.asciidoc[tag=param-timeout] +include::{es-repo-dir}/text-structure/apis/find-structure-shared.asciidoc[tag=param-timestamp-field] +include::{es-repo-dir}/text-structure/apis/find-structure-shared.asciidoc[tag=param-timestamp-format] + +[discrete] +[[find-message-structure-request-body]] +== {api-request-body-title} + +`messages`:: +(Required, array of strings) +The list of messages you want to analyze. + +[discrete] +[[find-message-structure-examples]] +== {api-examples-title} + +[discrete] +[[find-message-structure-example]] +=== Analyzing Elasticsearch log files + +Suppose you have a list of {es} logs messages. +You can send it to the `find_message_structure` endpoint as follows: + +[source,console] +---- +POST _text_structure/find_message_structure +{ + "messages": [ + "[2024-03-05T10:52:36,256][INFO ][o.a.l.u.VectorUtilPanamaProvider] [laptop] Java vector incubator API enabled; uses preferredBitSize=128", + "[2024-03-05T10:52:41,038][INFO ][o.e.p.PluginsService ] [laptop] loaded module [repository-url]", + "[2024-03-05T10:52:41,042][INFO ][o.e.p.PluginsService ] [laptop] loaded module [rest-root]", + "[2024-03-05T10:52:41,043][INFO ][o.e.p.PluginsService ] [laptop] loaded module [x-pack-core]", + "[2024-03-05T10:52:41,043][INFO ][o.e.p.PluginsService ] [laptop] loaded module [x-pack-redact]", + "[2024-03-05T10:52:41,043][INFO ][o.e.p.PluginsService ] [laptop] loaded module [ingest-user-agent]", + "[2024-03-05T10:52:41,044][INFO ][o.e.p.PluginsService ] [laptop] loaded module [x-pack-monitoring]", + "[2024-03-05T10:52:41,044][INFO ][o.e.p.PluginsService ] [laptop] loaded module [repository-s3]", + "[2024-03-05T10:52:41,044][INFO ][o.e.p.PluginsService ] [laptop] loaded module [x-pack-analytics]", + "[2024-03-05T10:52:41,044][INFO ][o.e.p.PluginsService ] [laptop] loaded module [x-pack-ent-search]", + "[2024-03-05T10:52:41,044][INFO ][o.e.p.PluginsService ] [laptop] loaded module [x-pack-autoscaling]", + "[2024-03-05T10:52:41,044][INFO ][o.e.p.PluginsService ] [laptop] loaded module [lang-painless]]", + "[2024-03-05T10:52:41,059][INFO ][o.e.p.PluginsService ] [laptop] loaded module [lang-expression]", + "[2024-03-05T10:52:41,059][INFO ][o.e.p.PluginsService ] [laptop] loaded module [x-pack-eql]", + "[2024-03-05T10:52:43,291][INFO ][o.e.e.NodeEnvironment ] [laptop] heap size [16gb], compressed ordinary object pointers [true]", + "[2024-03-05T10:52:46,098][INFO ][o.e.x.s.Security ] [laptop] Security is enabled", + "[2024-03-05T10:52:47,227][INFO ][o.e.x.p.ProfilingPlugin ] [laptop] Profiling is enabled", + "[2024-03-05T10:52:47,259][INFO ][o.e.x.p.ProfilingPlugin ] [laptop] profiling index templates will not be installed or reinstalled", + "[2024-03-05T10:52:47,755][INFO ][o.e.i.r.RecoverySettings ] [laptop] using rate limit [40mb] with [default=40mb, read=0b, write=0b, max=0b]", + "[2024-03-05T10:52:47,787][INFO ][o.e.d.DiscoveryModule ] [laptop] using discovery type [multi-node] and seed hosts providers [settings]", + "[2024-03-05T10:52:49,188][INFO ][o.e.n.Node ] [laptop] initialized", + "[2024-03-05T10:52:49,199][INFO ][o.e.n.Node ] [laptop] starting ..." + ] +} +---- +// TEST + +If the request does not encounter errors, you receive the following result: + +[source,console-result] +---- +{ + "num_lines_analyzed" : 22, + "num_messages_analyzed" : 22, + "sample_start" : "[2024-03-05T10:52:36,256][INFO ][o.a.l.u.VectorUtilPanamaProvider] [laptop] Java vector incubator API enabled; uses preferredBitSize=128\n[2024-03-05T10:52:41,038][INFO ][o.e.p.PluginsService ] [laptop] loaded module [repository-url]\n", <3> + "charset" : "UTF-8", + "format" : "semi_structured_text", + "multiline_start_pattern" : "^\\[\\b\\d{4}-\\d{2}-\\d{2}[T ]\\d{2}:\\d{2}", + "grok_pattern" : "\\[%{TIMESTAMP_ISO8601:timestamp}\\]\\[%{LOGLEVEL:loglevel} \\]\\[.*", + "ecs_compatibility" : "disabled", + "timestamp_field" : "timestamp", + "joda_timestamp_formats" : [ + "ISO8601" + ], + "java_timestamp_formats" : [ + "ISO8601" + ], + "need_client_timezone" : true, + "mappings" : { + "properties" : { + "@timestamp" : { + "type" : "date" + }, + "loglevel" : { + "type" : "keyword" + }, + "message" : { + "type" : "text" + } + } + }, + "ingest_pipeline" : { + "description" : "Ingest pipeline created by text structure finder", + "processors" : [ + { + "grok" : { + "field" : "message", + "patterns" : [ + "\\[%{TIMESTAMP_ISO8601:timestamp}\\]\\[%{LOGLEVEL:loglevel} \\]\\[.*" + ], + "ecs_compatibility" : "disabled" + } + }, + { + "date" : { + "field" : "timestamp", + "timezone" : "{{ event.timezone }}", + "formats" : [ + "ISO8601" + ] + } + }, + { + "remove" : { + "field" : "timestamp" + } + } + ] + }, + "field_stats" : { + "loglevel" : { + "count" : 22, + "cardinality" : 1, + "top_hits" : [ + { + "value" : "INFO", + "count" : 22 + } + ] + }, + "message" : { + "count" : 22, + "cardinality" : 22, + "top_hits" : [ + { + "value" : "[2024-03-05T10:52:36,256][INFO ][o.a.l.u.VectorUtilPanamaProvider] [laptop] Java vector incubator API enabled; uses preferredBitSize=128", + "count" : 1 + }, + { + "value" : "[2024-03-05T10:52:41,038][INFO ][o.e.p.PluginsService ] [laptop] loaded module [repository-url]", + "count" : 1 + }, + { + "value" : "[2024-03-05T10:52:41,042][INFO ][o.e.p.PluginsService ] [laptop] loaded module [rest-root]", + "count" : 1 + }, + { + "value" : "[2024-03-05T10:52:41,043][INFO ][o.e.p.PluginsService ] [laptop] loaded module [ingest-user-agent]", + "count" : 1 + }, + { + "value" : "[2024-03-05T10:52:41,043][INFO ][o.e.p.PluginsService ] [laptop] loaded module [x-pack-core]", + "count" : 1 + }, + { + "value" : "[2024-03-05T10:52:41,043][INFO ][o.e.p.PluginsService ] [laptop] loaded module [x-pack-redact]", + "count" : 1 + }, + { + "value" : "[2024-03-05T10:52:41,044][INFO ][o.e.p.PluginsService ] [laptop] loaded module [lang-painless]]", + "count" : 1 + }, + { + "value" : "[2024-03-05T10:52:41,044][INFO ][o.e.p.PluginsService ] [laptop] loaded module [repository-s3]", + "count" : 1 + }, + { + "value" : "[2024-03-05T10:52:41,044][INFO ][o.e.p.PluginsService ] [laptop] loaded module [x-pack-analytics]", + "count" : 1 + }, + { + "value" : "[2024-03-05T10:52:41,044][INFO ][o.e.p.PluginsService ] [laptop] loaded module [x-pack-autoscaling]", + "count" : 1 + } + ] + }, + "timestamp" : { + "count" : 22, + "cardinality" : 14, + "earliest" : "2024-03-05T10:52:36,256", + "latest" : "2024-03-05T10:52:49,199", + "top_hits" : [ + { + "value" : "2024-03-05T10:52:41,044", + "count" : 6 + }, + { + "value" : "2024-03-05T10:52:41,043", + "count" : 3 + }, + { + "value" : "2024-03-05T10:52:41,059", + "count" : 2 + }, + { + "value" : "2024-03-05T10:52:36,256", + "count" : 1 + }, + { + "value" : "2024-03-05T10:52:41,038", + "count" : 1 + }, + { + "value" : "2024-03-05T10:52:41,042", + "count" : 1 + }, + { + "value" : "2024-03-05T10:52:43,291", + "count" : 1 + }, + { + "value" : "2024-03-05T10:52:46,098", + "count" : 1 + }, + { + "value" : "2024-03-05T10:52:47,227", + "count" : 1 + }, + { + "value" : "2024-03-05T10:52:47,259", + "count" : 1 + } + ] + } + } +} +---- +// TESTRESPONSE + +For a detailed description of the response format, or for additional examples +on ingesting delimited text (such as CSV) or newline-delimited JSON, refer to the +<>. diff --git a/docs/reference/text-structure/apis/find-structure-shared.asciidoc b/docs/reference/text-structure/apis/find-structure-shared.asciidoc new file mode 100644 index 0000000000000..67a85dd072a9a --- /dev/null +++ b/docs/reference/text-structure/apis/find-structure-shared.asciidoc @@ -0,0 +1,215 @@ +tag::param-charset[] +`charset`:: +(Optional, string) The text's character set. It must be a character set that is +supported by the JVM that {es} uses. For example, `UTF-8`, `UTF-16LE`, +`windows-1252`, or `EUC-JP`. If this parameter is not specified, the structure +finder chooses an appropriate character set. +end::param-charset[] + +tag::param-column-names[] +`column_names`:: +(Optional, string) If you have set `format` to `delimited`, you can specify the +column names in a comma-separated list. If this parameter is not specified, the +structure finder uses the column names from the header row of the text. If the +text does not have a header row, columns are named "column1", "column2", +"column3", etc. +end::param-column-names[] + +tag::param-delimiter[] +`delimiter`:: +(Optional, string) If you have set `format` to `delimited`, you can specify the +character used to delimit the values in each row. Only a single character is +supported; the delimiter cannot have multiple characters. By default, the API +considers the following possibilities: comma, tab, semi-colon, and pipe (`|`). +In this default scenario, all rows must have the same number of fields for the +delimited format to be detected. If you specify a delimiter, up to 10% of the +rows can have a different number of columns than the first row. +end::param-delimiter[] + +tag::param-explain[] +`explain`:: +(Optional, Boolean) If `true`, the response includes a +field named `explanation`, which is an array of strings that indicate how the +structure finder produced its result. The default value is `false`. +end::param-explain[] + +tag::param-format[] +`format`:: +(Optional, string) The high level structure of the text. Valid values are +`ndjson`, `xml`, `delimited`, and `semi_structured_text`. By default, the API +chooses the format. In this default scenario, all rows must have the same number +of fields for a delimited format to be detected. If the `format` is set to +`delimited` and the `delimiter` is not set, however, the API tolerates up to 5% +of rows that have a different number of columns than the first row. +end::param-format[] + +tag::param-grok-pattern[] +`grok_pattern`:: +(Optional, string) If you have set `format` to `semi_structured_text`, you can +specify a Grok pattern that is used to extract fields from every message in the +text. The name of the timestamp field in the Grok pattern must match what is +specified in the `timestamp_field` parameter. If that parameter is not +specified, the name of the timestamp field in the Grok pattern must match +"timestamp". If `grok_pattern` is not specified, the structure finder creates a +Grok pattern. +end::param-grok-pattern[] + +tag::param-ecs-compatibility[] +`ecs_compatibility`:: +(Optional, string) The mode of compatibility with ECS compliant Grok patterns. +Use this parameter to specify whether to use ECS Grok patterns instead of +legacy ones when the structure finder creates a Grok pattern. Valid values +are `disabled` and `v1`. The default value is `disabled`. This setting primarily +has an impact when a whole message Grok pattern such as `%{CATALINALOG}` +matches the input. If the structure finder identifies a common structure but +has no idea of meaning then generic field names such as `path`, `ipaddress`, +`field1` and `field2` are used in the `grok_pattern` output, with the intention +that a user who knows the meanings rename these fields before using it. +end::param-ecs-compatibility[] + +tag::param-has-header-row[] +`has_header_row`:: +(Optional, Boolean) If you have set `format` to `delimited`, you can use this +parameter to indicate whether the column names are in the first row of the text. +If this parameter is not specified, the structure finder guesses based on the +similarity of the first row of the text to other rows. +end::param-has-header-row[] + +tag::param-line-merge-size-limit[] +`line_merge_size_limit`:: +(Optional, unsigned integer) The maximum number of characters in a message when +lines are merged to form messages while analyzing semi-structured text. The +default is `10000`. If you have extremely long messages you may need to increase +this, but be aware that this may lead to very long processing times if the way +to group lines into messages is misdetected. +end::param-line-merge-size-limit[] + +tag::param-lines-to-sample[] +`lines_to_sample`:: +(Optional, unsigned integer) The number of lines to include in the structural +analysis, starting from the beginning of the text. The minimum is 2; the default +is `1000`. If the value of this parameter is greater than the number of lines in +the text, the analysis proceeds (as long as there are at least two lines in the +text) for all of the lines. ++ +-- +NOTE: The number of lines and the variation of the lines affects the speed of +the analysis. For example, if you upload text where the first 1000 lines +are all variations on the same message, the analysis will find more commonality +than would be seen with a bigger sample. If possible, however, it is more +efficient to upload sample text with more variety in the first 1000 lines than +to request analysis of 100000 lines to achieve some variety. + +-- +end::param-lines-to-sample[] + +tag::param-quote[] +`quote`:: +(Optional, string) If you have set `format` to `delimited`, you can specify the +character used to quote the values in each row if they contain newlines or the +delimiter character. Only a single character is supported. If this parameter is +not specified, the default value is a double quote (`"`). If your delimited text +format does not use quoting, a workaround is to set this argument to a character +that does not appear anywhere in the sample. +end::param-quote[] + +tag::param-should-trim-fields[] +`should_trim_fields`:: +(Optional, Boolean) If you have set `format` to `delimited`, you can specify +whether values between delimiters should have whitespace trimmed from them. If +this parameter is not specified and the delimiter is pipe (`|`), the default +value is `true`. Otherwise, the default value is `false`. +end::param-should-trim-fields[] + +tag::param-timeout[] +`timeout`:: +(Optional, <>) Sets the maximum amount of time that the +structure analysis may take. If the analysis is still running when the timeout +expires then it will be stopped. The default value is 25 seconds. +end::param-timeout[] + +tag::param-timestamp-field[] +`timestamp_field`:: +(Optional, string) The name of the field that contains the primary timestamp of +each record in the text. In particular, if the text were ingested into an index, +this is the field that would be used to populate the `@timestamp` field. ++ +-- +If the `format` is `semi_structured_text`, this field must match the name of the +appropriate extraction in the `grok_pattern`. Therefore, for semi-structured +text, it is best not to specify this parameter unless `grok_pattern` is +also specified. + +For structured text, if you specify this parameter, the field must exist +within the text. + +If this parameter is not specified, the structure finder makes a decision about +which field (if any) is the primary timestamp field. For structured text, +it is not compulsory to have a timestamp in the text. +-- +end::param-timestamp-field[] + +tag::param-timestamp-format[] +`timestamp_format`:: +(Optional, string) The Java time format of the timestamp field in the text. ++ +-- +Only a subset of Java time format letter groups are supported: + +* `a` +* `d` +* `dd` +* `EEE` +* `EEEE` +* `H` +* `HH` +* `h` +* `M` +* `MM` +* `MMM` +* `MMMM` +* `mm` +* `ss` +* `XX` +* `XXX` +* `yy` +* `yyyy` +* `zzz` + +Additionally `S` letter groups (fractional seconds) of length one to nine are +supported providing they occur after `ss` and separated from the `ss` by a `.`, +`,` or `:`. Spacing and punctuation is also permitted with the exception of `?`, +newline and carriage return, together with literal text enclosed in single +quotes. For example, `MM/dd HH.mm.ss,SSSSSS 'in' yyyy` is a valid override +format. + +One valuable use case for this parameter is when the format is semi-structured +text, there are multiple timestamp formats in the text, and you know which +format corresponds to the primary timestamp, but you do not want to specify the +full `grok_pattern`. Another is when the timestamp format is one that the +structure finder does not consider by default. + +If this parameter is not specified, the structure finder chooses the best +format from a built-in set. + +If the special value `null` is specified the structure finder will not look +for a primary timestamp in the text. When the format is semi-structured text +this will result in the structure finder treating the text as single-line +messages. + +The following table provides the appropriate `timeformat` values for some example timestamps: + +|=== +| Timeformat | Presentation + +| yyyy-MM-dd HH:mm:ssZ | 2019-04-20 13:15:22+0000 +| EEE, d MMM yyyy HH:mm:ss Z | Sat, 20 Apr 2019 13:15:22 +0000 +| dd.MM.yy HH:mm:ss.SSS | 20.04.19 13:15:22.285 +|=== + +Refer to +https://docs.oracle.com/javase/8/docs/api/java/time/format/DateTimeFormatter.html[the Java date/time format documentation] +for more information about date and time format syntax. + +-- +end::param-timestamp-format[] diff --git a/docs/reference/text-structure/apis/find-structure.asciidoc b/docs/reference/text-structure/apis/find-structure.asciidoc index a65f87290b0a8..b49b0f3526689 100644 --- a/docs/reference/text-structure/apis/find-structure.asciidoc +++ b/docs/reference/text-structure/apis/find-structure.asciidoc @@ -1,6 +1,6 @@ [role="xpack"] [[find-structure]] -= Find structure API += Find text structure API Finds the structure of text. The text must contain data that is suitable to be ingested into the @@ -55,190 +55,21 @@ chosen. [[find-structure-query-parms]] == {api-query-parms-title} -`charset`:: -(Optional, string) The text's character set. It must be a character set that is -supported by the JVM that {es} uses. For example, `UTF-8`, `UTF-16LE`, -`windows-1252`, or `EUC-JP`. If this parameter is not specified, the structure -finder chooses an appropriate character set. - -`column_names`:: -(Optional, string) If you have set `format` to `delimited`, you can specify the -column names in a comma-separated list. If this parameter is not specified, the -structure finder uses the column names from the header row of the text. If the -text does not have a header role, columns are named "column1", "column2", -"column3", etc. - -`delimiter`:: -(Optional, string) If you have set `format` to `delimited`, you can specify the -character used to delimit the values in each row. Only a single character is -supported; the delimiter cannot have multiple characters. By default, the API -considers the following possibilities: comma, tab, semi-colon, and pipe (`|`). -In this default scenario, all rows must have the same number of fields for the -delimited format to be detected. If you specify a delimiter, up to 10% of the -rows can have a different number of columns than the first row. - -`explain`:: -(Optional, Boolean) If this parameter is set to `true`, the response includes a -field named `explanation`, which is an array of strings that indicate how the -structure finder produced its result. The default value is `false`. - -`format`:: -(Optional, string) The high level structure of the text. Valid values are -`ndjson`, `xml`, `delimited`, and `semi_structured_text`. By default, the API -chooses the format. In this default scenario, all rows must have the same number -of fields for a delimited format to be detected. If the `format` is set to -`delimited` and the `delimiter` is not set, however, the API tolerates up to 5% -of rows that have a different number of columns than the first row. - -`grok_pattern`:: -(Optional, string) If you have set `format` to `semi_structured_text`, you can -specify a Grok pattern that is used to extract fields from every message in the -text. The name of the timestamp field in the Grok pattern must match what is -specified in the `timestamp_field` parameter. If that parameter is not -specified, the name of the timestamp field in the Grok pattern must match -"timestamp". If `grok_pattern` is not specified, the structure finder creates a -Grok pattern. - -`ecs_compatibility`:: -(Optional, string) The mode of compatibility with ECS compliant Grok patterns. -Use this parameter to specify whether to use ECS Grok patterns instead of -legacy ones when the structure finder creates a Grok pattern. Valid values -are `disabled` and `v1`. The default value is `disabled`. This setting primarily -has an impact when a whole message Grok pattern such as `%{CATALINALOG}` -matches the input. If the structure finder identifies a common structure but -has no idea of meaning then generic field names such as `path`, `ipaddress`, -`field1` and `field2` are used in the `grok_pattern` output, with the intention -that a user who knows the meanings rename these fields before using it. -`has_header_row`:: -(Optional, Boolean) If you have set `format` to `delimited`, you can use this -parameter to indicate whether the column names are in the first row of the text. -If this parameter is not specified, the structure finder guesses based on the -similarity of the first row of the text to other rows. - -`line_merge_size_limit`:: -(Optional, unsigned integer) The maximum number of characters in a message when -lines are merged to form messages while analyzing semi-structured text. The -default is `10000`. If you have extremely long messages you may need to increase -this, but be aware that this may lead to very long processing times if the way -to group lines into messages is misdetected. - -`lines_to_sample`:: -(Optional, unsigned integer) The number of lines to include in the structural -analysis, starting from the beginning of the text. The minimum is 2; the default -is `1000`. If the value of this parameter is greater than the number of lines in -the text, the analysis proceeds (as long as there are at least two lines in the -text) for all of the lines. -+ --- -NOTE: The number of lines and the variation of the lines affects the speed of -the analysis. For example, if you upload text where the first 1000 lines -are all variations on the same message, the analysis will find more commonality -than would be seen with a bigger sample. If possible, however, it is more -efficient to upload sample text with more variety in the first 1000 lines than -to request analysis of 100000 lines to achieve some variety. - --- - -`quote`:: -(Optional, string) If you have set `format` to `delimited`, you can specify the -character used to quote the values in each row if they contain newlines or the -delimiter character. Only a single character is supported. If this parameter is -not specified, the default value is a double quote (`"`). If your delimited text -format does not use quoting, a workaround is to set this argument to a character -that does not appear anywhere in the sample. - -`should_trim_fields`:: -(Optional, Boolean) If you have set `format` to `delimited`, you can specify -whether values between delimiters should have whitespace trimmed from them. If -this parameter is not specified and the delimiter is pipe (`|`), the default -value is `true`. Otherwise, the default value is `false`. - -`timeout`:: -(Optional, <>) Sets the maximum amount of time that the -structure analysis make take. If the analysis is still running when the timeout -expires then it will be aborted. The default value is 25 seconds. - -`timestamp_field`:: -(Optional, string) The name of the field that contains the primary timestamp of -each record in the text. In particular, if the text were ingested into an index, -this is the field that would be used to populate the `@timestamp` field. -+ --- -If the `format` is `semi_structured_text`, this field must match the name of the -appropriate extraction in the `grok_pattern`. Therefore, for semi-structured -text, it is best not to specify this parameter unless `grok_pattern` is -also specified. - -For structured text, if you specify this parameter, the field must exist -within the text. - -If this parameter is not specified, the structure finder makes a decision about -which field (if any) is the primary timestamp field. For structured text, -it is not compulsory to have a timestamp in the text. --- - -`timestamp_format`:: -(Optional, string) The Java time format of the timestamp field in the text. -+ --- -Only a subset of Java time format letter groups are supported: - -* `a` -* `d` -* `dd` -* `EEE` -* `EEEE` -* `H` -* `HH` -* `h` -* `M` -* `MM` -* `MMM` -* `MMMM` -* `mm` -* `ss` -* `XX` -* `XXX` -* `yy` -* `yyyy` -* `zzz` - -Additionally `S` letter groups (fractional seconds) of length one to nine are -supported providing they occur after `ss` and separated from the `ss` by a `.`, -`,` or `:`. Spacing and punctuation is also permitted with the exception of `?`, -newline and carriage return, together with literal text enclosed in single -quotes. For example, `MM/dd HH.mm.ss,SSSSSS 'in' yyyy` is a valid override -format. - -One valuable use case for this parameter is when the format is semi-structured -text, there are multiple timestamp formats in the text, and you know which -format corresponds to the primary timestamp, but you do not want to specify the -full `grok_pattern`. Another is when the timestamp format is one that the -structure finder does not consider by default. - -If this parameter is not specified, the structure finder chooses the best -format from a built-in set. - -If the special value `null` is specified the structure finder will not look -for a primary timestamp in the text. When the format is semi-structured text -this will result in the structure finder treating the text as single-line -messages. - -The following table provides the appropriate `timeformat` values for some example timestamps: - -|=== -| Timeformat | Presentation - -| yyyy-MM-dd HH:mm:ssZ | 2019-04-20 13:15:22+0000 -| EEE, d MMM yyyy HH:mm:ss Z | Sat, 20 Apr 2019 13:15:22 +0000 -| dd.MM.yy HH:mm:ss.SSS | 20.04.19 13:15:22.285 -|=== - -See -https://docs.oracle.com/javase/8/docs/api/java/time/format/DateTimeFormatter.html[the Java date/time format documentation] -for more information about date and time format syntax. - --- +include::{es-repo-dir}/text-structure/apis/find-structure-shared.asciidoc[tag=param-charset] +include::{es-repo-dir}/text-structure/apis/find-structure-shared.asciidoc[tag=param-column-names] +include::{es-repo-dir}/text-structure/apis/find-structure-shared.asciidoc[tag=param-delimiter] +include::{es-repo-dir}/text-structure/apis/find-structure-shared.asciidoc[tag=param-explain] +include::{es-repo-dir}/text-structure/apis/find-structure-shared.asciidoc[tag=param-format] +include::{es-repo-dir}/text-structure/apis/find-structure-shared.asciidoc[tag=param-grok-pattern] +include::{es-repo-dir}/text-structure/apis/find-structure-shared.asciidoc[tag=param-ecs-compatibility] +include::{es-repo-dir}/text-structure/apis/find-structure-shared.asciidoc[tag=param-has-header-row] +include::{es-repo-dir}/text-structure/apis/find-structure-shared.asciidoc[tag=param-line-merge-size-limit] +include::{es-repo-dir}/text-structure/apis/find-structure-shared.asciidoc[tag=param-lines-to-sample] +include::{es-repo-dir}/text-structure/apis/find-structure-shared.asciidoc[tag=param-quote] +include::{es-repo-dir}/text-structure/apis/find-structure-shared.asciidoc[tag=param-should-trim-fields] +include::{es-repo-dir}/text-structure/apis/find-structure-shared.asciidoc[tag=param-timeout] +include::{es-repo-dir}/text-structure/apis/find-structure-shared.asciidoc[tag=param-timestamp-field] +include::{es-repo-dir}/text-structure/apis/find-structure-shared.asciidoc[tag=param-timestamp-format] [discrete] [[find-structure-request-body]] diff --git a/docs/reference/text-structure/apis/index.asciidoc b/docs/reference/text-structure/apis/index.asciidoc index 8628badba7e78..9f4af120690f7 100644 --- a/docs/reference/text-structure/apis/index.asciidoc +++ b/docs/reference/text-structure/apis/index.asciidoc @@ -4,8 +4,12 @@ You can use the following APIs to find text structures: +* <> +* <> * <> * <> +include::find-field-structure.asciidoc[leveloffset=+2] +include::find-message-structure.asciidoc[leveloffset=+2] include::find-structure.asciidoc[leveloffset=+2] include::test-grok-pattern.asciidoc[leveloffset=+2] diff --git a/rest-api-spec/src/main/resources/rest-api-spec/api/text_structure.find_field_structure.json b/rest-api-spec/src/main/resources/rest-api-spec/api/text_structure.find_field_structure.json new file mode 100644 index 0000000000000..f82e2ca2d190f --- /dev/null +++ b/rest-api-spec/src/main/resources/rest-api-spec/api/text_structure.find_field_structure.json @@ -0,0 +1,90 @@ +{ + "text_structure.find_field_structure":{ + "documentation":{ + "url":"https://www.elastic.co/guide/en/elasticsearch/reference/current/find-field-structure.html", + "description":"Finds the structure of a text field in an index." + }, + "stability":"stable", + "visibility":"public", + "headers":{ + "accept": ["application/json"] + }, + "url":{ + "paths":[ + { + "path":"/_text_structure/find_field_structure", + "methods":["GET"] + } + ] + }, + "params":{ + "index":{ + "type":"string", + "description":"The index containing the analyzed field", + "required":true + }, + "field":{ + "type":"string", + "description":"The field that should be analyzed", + "required":true + }, + "documents_to_sample":{ + "type":"int", + "description":"How many documents should be included in the analysis", + "default":1000 + }, + "timeout":{ + "type":"time", + "description":"Timeout after which the analysis will be aborted", + "default":"25s" + }, + "format":{ + "type":"enum", + "options":[ + "ndjson", + "xml", + "delimited", + "semi_structured_text" + ], + "description":"Optional parameter to specify the high level file format" + }, + "column_names":{ + "type":"list", + "description":"Optional parameter containing a comma separated list of the column names for a delimited file" + }, + "delimiter":{ + "type":"string", + "description":"Optional parameter to specify the delimiter character for a delimited file - must be a single character" + }, + "quote":{ + "type":"string", + "description":"Optional parameter to specify the quote character for a delimited file - must be a single character" + }, + "should_trim_fields":{ + "type":"boolean", + "description":"Optional parameter to specify whether the values between delimiters in a delimited file should have whitespace trimmed from them" + }, + "grok_pattern":{ + "type":"string", + "description":"Optional parameter to specify the Grok pattern that should be used to extract fields from messages in a semi-structured text file" + }, + "ecs_compatibility":{ + "type":"string", + "description":"Optional parameter to specify the compatibility mode with ECS Grok patterns - may be either 'v1' or 'disabled'" + }, + "timestamp_field":{ + "type":"string", + "description":"Optional parameter to specify the timestamp field in the file" + }, + "timestamp_format":{ + "type":"string", + "description":"Optional parameter to specify the timestamp format in the file - may be either a Joda or Java time format" + }, + "explain":{ + "type":"boolean", + "description":"Whether to include a commentary on how the structure was derived", + "default":false + } + } + } +} diff --git a/rest-api-spec/src/main/resources/rest-api-spec/api/text_structure.find_message_structure.json b/rest-api-spec/src/main/resources/rest-api-spec/api/text_structure.find_message_structure.json new file mode 100644 index 0000000000000..d839e4b048f7d --- /dev/null +++ b/rest-api-spec/src/main/resources/rest-api-spec/api/text_structure.find_message_structure.json @@ -0,0 +1,80 @@ +{ + "text_structure.find_message_structure":{ + "documentation":{ + "url":"https://www.elastic.co/guide/en/elasticsearch/reference/current/find-message-structure.html", + "description":"Finds the structure of a list of messages. The messages must contain data that is suitable to be ingested into Elasticsearch." + }, + "stability":"stable", + "visibility":"public", + "headers":{ + "accept": [ "application/json"], + "content_type": ["application/json"] + }, + "url":{ + "paths":[ + { + "path":"/_text_structure/find_message_structure", + "methods":["GET", "POST"] + } + ] + }, + "params":{ + "timeout":{ + "type":"time", + "description":"Timeout after which the analysis will be aborted", + "default":"25s" + }, + "format":{ + "type":"enum", + "options":[ + "ndjson", + "xml", + "delimited", + "semi_structured_text" + ], + "description":"Optional parameter to specify the high level file format" + }, + "column_names":{ + "type":"list", + "description":"Optional parameter containing a comma separated list of the column names for a delimited file" + }, + "delimiter":{ + "type":"string", + "description":"Optional parameter to specify the delimiter character for a delimited file - must be a single character" + }, + "quote":{ + "type":"string", + "description":"Optional parameter to specify the quote character for a delimited file - must be a single character" + }, + "should_trim_fields":{ + "type":"boolean", + "description":"Optional parameter to specify whether the values between delimiters in a delimited file should have whitespace trimmed from them" + }, + "grok_pattern":{ + "type":"string", + "description":"Optional parameter to specify the Grok pattern that should be used to extract fields from messages in a semi-structured text file" + }, + "ecs_compatibility":{ + "type":"string", + "description":"Optional parameter to specify the compatibility mode with ECS Grok patterns - may be either 'v1' or 'disabled'" + }, + "timestamp_field":{ + "type":"string", + "description":"Optional parameter to specify the timestamp field in the file" + }, + "timestamp_format":{ + "type":"string", + "description":"Optional parameter to specify the timestamp format in the file - may be either a Joda or Java time format" + }, + "explain":{ + "type":"boolean", + "description":"Whether to include a commentary on how the structure was derived", + "default":false + } + }, + "body":{ + "description":"JSON object with one field [messages], containing an array of messages to be analyzed", + "required":true + } + } +} diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/textstructure/action/AbstractFindStructureRequest.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/textstructure/action/AbstractFindStructureRequest.java new file mode 100644 index 0000000000000..e06ffd3b95a05 --- /dev/null +++ b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/textstructure/action/AbstractFindStructureRequest.java @@ -0,0 +1,377 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.core.textstructure.action; + +import org.elasticsearch.TransportVersions; +import org.elasticsearch.action.ActionRequest; +import org.elasticsearch.action.ActionRequestValidationException; +import org.elasticsearch.common.io.stream.StreamInput; +import org.elasticsearch.common.io.stream.StreamOutput; +import org.elasticsearch.core.TimeValue; +import org.elasticsearch.grok.GrokBuiltinPatterns; +import org.elasticsearch.xcontent.ParseField; +import org.elasticsearch.xpack.core.textstructure.structurefinder.TextStructure; + +import java.io.IOException; +import java.util.Arrays; +import java.util.List; +import java.util.Locale; +import java.util.Objects; + +import static org.elasticsearch.action.ValidateActions.addValidationError; + +public abstract class AbstractFindStructureRequest extends ActionRequest { + + public static final int MIN_SAMPLE_LINE_COUNT = 2; + + public static final ParseField LINES_TO_SAMPLE = new ParseField("lines_to_sample"); + public static final ParseField DOCUMENTS_TO_SAMPLE = new ParseField("documents_to_sample"); + public static final ParseField LINE_MERGE_SIZE_LIMIT = new ParseField("line_merge_size_limit"); + public static final ParseField TIMEOUT = new ParseField("timeout"); + public static final ParseField CHARSET = TextStructure.CHARSET; + public static final ParseField FORMAT = TextStructure.FORMAT; + public static final ParseField COLUMN_NAMES = TextStructure.COLUMN_NAMES; + public static final ParseField HAS_HEADER_ROW = TextStructure.HAS_HEADER_ROW; + public static final ParseField DELIMITER = TextStructure.DELIMITER; + public static final ParseField QUOTE = TextStructure.QUOTE; + public static final ParseField SHOULD_TRIM_FIELDS = TextStructure.SHOULD_TRIM_FIELDS; + public static final ParseField GROK_PATTERN = TextStructure.GROK_PATTERN; + // This one is plural in FileStructure, but singular in FileStructureOverrides + public static final ParseField TIMESTAMP_FORMAT = new ParseField("timestamp_format"); + public static final ParseField TIMESTAMP_FIELD = TextStructure.TIMESTAMP_FIELD; + + public static final ParseField ECS_COMPATIBILITY = TextStructure.ECS_COMPATIBILITY; + + private static final String ARG_INCOMPATIBLE_WITH_FORMAT_TEMPLATE = "[%s] may only be specified if [" + + FORMAT.getPreferredName() + + "] is [%s]"; + + private Integer linesToSample; + private Integer lineMergeSizeLimit; + private TimeValue timeout; + private String charset; + private TextStructure.Format format; + private List columnNames; + private Boolean hasHeaderRow; + private Character delimiter; + private Character quote; + private Boolean shouldTrimFields; + private String grokPattern; + private String ecsCompatibility; + private String timestampFormat; + private String timestampField; + + AbstractFindStructureRequest() {} + + AbstractFindStructureRequest(StreamInput in) throws IOException { + super(in); + linesToSample = in.readOptionalVInt(); + lineMergeSizeLimit = in.readOptionalVInt(); + timeout = in.readOptionalTimeValue(); + charset = in.readOptionalString(); + format = in.readBoolean() ? in.readEnum(TextStructure.Format.class) : null; + columnNames = in.readBoolean() ? in.readStringCollectionAsList() : null; + hasHeaderRow = in.readOptionalBoolean(); + delimiter = in.readBoolean() ? (char) in.readVInt() : null; + quote = in.readBoolean() ? (char) in.readVInt() : null; + shouldTrimFields = in.readOptionalBoolean(); + grokPattern = in.readOptionalString(); + if (in.getTransportVersion().onOrAfter(TransportVersions.V_8_5_0)) { + ecsCompatibility = in.readOptionalString(); + } else { + ecsCompatibility = null; + } + timestampFormat = in.readOptionalString(); + timestampField = in.readOptionalString(); + } + + public Integer getLinesToSample() { + return linesToSample; + } + + public void setLinesToSample(Integer linesToSample) { + this.linesToSample = linesToSample; + } + + public Integer getLineMergeSizeLimit() { + return lineMergeSizeLimit; + } + + public void setLineMergeSizeLimit(Integer lineMergeSizeLimit) { + this.lineMergeSizeLimit = lineMergeSizeLimit; + } + + public TimeValue getTimeout() { + return timeout; + } + + public void setTimeout(TimeValue timeout) { + this.timeout = timeout; + } + + public String getCharset() { + return charset; + } + + public void setCharset(String charset) { + this.charset = (charset == null || charset.isEmpty()) ? null : charset; + } + + public TextStructure.Format getFormat() { + return format; + } + + public void setFormat(TextStructure.Format format) { + this.format = format; + } + + public void setFormat(String format) { + this.format = (format == null || format.isEmpty()) ? null : TextStructure.Format.fromString(format); + } + + public List getColumnNames() { + return columnNames; + } + + public void setColumnNames(List columnNames) { + this.columnNames = (columnNames == null || columnNames.isEmpty()) ? null : columnNames; + } + + public void setColumnNames(String[] columnNames) { + this.columnNames = (columnNames == null || columnNames.length == 0) ? null : Arrays.asList(columnNames); + } + + public Boolean getHasHeaderRow() { + return hasHeaderRow; + } + + public void setHasHeaderRow(Boolean hasHeaderRow) { + this.hasHeaderRow = hasHeaderRow; + } + + public Character getDelimiter() { + return delimiter; + } + + public void setDelimiter(Character delimiter) { + this.delimiter = delimiter; + } + + public void setDelimiter(String delimiter) { + if (delimiter == null || delimiter.isEmpty()) { + this.delimiter = null; + } else if (delimiter.length() == 1) { + this.delimiter = delimiter.charAt(0); + } else { + throw new IllegalArgumentException(DELIMITER.getPreferredName() + " must be a single character"); + } + } + + public Character getQuote() { + return quote; + } + + public void setQuote(Character quote) { + this.quote = quote; + } + + public void setQuote(String quote) { + if (quote == null || quote.isEmpty()) { + this.quote = null; + } else if (quote.length() == 1) { + this.quote = quote.charAt(0); + } else { + throw new IllegalArgumentException(QUOTE.getPreferredName() + " must be a single character"); + } + } + + public Boolean getShouldTrimFields() { + return shouldTrimFields; + } + + public void setShouldTrimFields(Boolean shouldTrimFields) { + this.shouldTrimFields = shouldTrimFields; + } + + public String getGrokPattern() { + return grokPattern; + } + + public void setGrokPattern(String grokPattern) { + this.grokPattern = (grokPattern == null || grokPattern.isEmpty()) ? null : grokPattern; + } + + public String getEcsCompatibility() { + return ecsCompatibility; + } + + public void setEcsCompatibility(String ecsCompatibility) { + this.ecsCompatibility = (ecsCompatibility == null || ecsCompatibility.isEmpty()) ? null : ecsCompatibility; + } + + public String getTimestampFormat() { + return timestampFormat; + } + + public void setTimestampFormat(String timestampFormat) { + this.timestampFormat = (timestampFormat == null || timestampFormat.isEmpty()) ? null : timestampFormat; + } + + public String getTimestampField() { + return timestampField; + } + + public void setTimestampField(String timestampField) { + this.timestampField = (timestampField == null || timestampField.isEmpty()) ? null : timestampField; + } + + private static ActionRequestValidationException addIncompatibleArgError( + ParseField arg, + TextStructure.Format format, + ActionRequestValidationException validationException + ) { + return addValidationError( + String.format(Locale.ROOT, ARG_INCOMPATIBLE_WITH_FORMAT_TEMPLATE, arg.getPreferredName(), format), + validationException + ); + } + + @Override + public ActionRequestValidationException validate() { + ActionRequestValidationException validationException = null; + if (linesToSample != null && linesToSample < MIN_SAMPLE_LINE_COUNT) { + validationException = addValidationError( + "[" + LINES_TO_SAMPLE.getPreferredName() + "] must be at least [" + MIN_SAMPLE_LINE_COUNT + "] if specified", + validationException + ); + } + if (lineMergeSizeLimit != null && lineMergeSizeLimit <= 0) { + validationException = addValidationError( + "[" + LINE_MERGE_SIZE_LIMIT.getPreferredName() + "] must be positive if specified", + validationException + ); + } + if (format != TextStructure.Format.DELIMITED) { + if (columnNames != null) { + validationException = addIncompatibleArgError(COLUMN_NAMES, TextStructure.Format.DELIMITED, validationException); + } + if (hasHeaderRow != null) { + validationException = addIncompatibleArgError(HAS_HEADER_ROW, TextStructure.Format.DELIMITED, validationException); + } + if (delimiter != null) { + validationException = addIncompatibleArgError(DELIMITER, TextStructure.Format.DELIMITED, validationException); + } + if (quote != null) { + validationException = addIncompatibleArgError(QUOTE, TextStructure.Format.DELIMITED, validationException); + } + if (shouldTrimFields != null) { + validationException = addIncompatibleArgError(SHOULD_TRIM_FIELDS, TextStructure.Format.DELIMITED, validationException); + } + } + if (format != TextStructure.Format.SEMI_STRUCTURED_TEXT) { + if (grokPattern != null) { + validationException = addIncompatibleArgError(GROK_PATTERN, TextStructure.Format.SEMI_STRUCTURED_TEXT, validationException); + } + } + + if (ecsCompatibility != null && GrokBuiltinPatterns.isValidEcsCompatibilityMode(ecsCompatibility) == false) { + validationException = addValidationError( + "[" + + ECS_COMPATIBILITY.getPreferredName() + + "] must be one of [" + + String.join(", ", GrokBuiltinPatterns.ECS_COMPATIBILITY_MODES) + + "] if specified", + validationException + ); + } + + return validationException; + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + super.writeTo(out); + out.writeOptionalVInt(linesToSample); + out.writeOptionalVInt(lineMergeSizeLimit); + out.writeOptionalTimeValue(timeout); + out.writeOptionalString(charset); + if (format == null) { + out.writeBoolean(false); + } else { + out.writeBoolean(true); + out.writeEnum(format); + } + if (columnNames == null) { + out.writeBoolean(false); + } else { + out.writeBoolean(true); + out.writeStringCollection(columnNames); + } + out.writeOptionalBoolean(hasHeaderRow); + if (delimiter == null) { + out.writeBoolean(false); + } else { + out.writeBoolean(true); + out.writeVInt(delimiter); + } + if (quote == null) { + out.writeBoolean(false); + } else { + out.writeBoolean(true); + out.writeVInt(quote); + } + out.writeOptionalBoolean(shouldTrimFields); + out.writeOptionalString(grokPattern); + if (out.getTransportVersion().onOrAfter(TransportVersions.V_8_5_0)) { + out.writeOptionalString(ecsCompatibility); + } + out.writeOptionalString(timestampFormat); + out.writeOptionalString(timestampField); + } + + @Override + public int hashCode() { + return Objects.hash( + linesToSample, + lineMergeSizeLimit, + timeout, + charset, + format, + columnNames, + hasHeaderRow, + delimiter, + grokPattern, + ecsCompatibility, + timestampFormat, + timestampField + ); + } + + @Override + public boolean equals(Object other) { + if (this == other) { + return true; + } + if (other == null || getClass() != other.getClass()) { + return false; + } + AbstractFindStructureRequest that = (AbstractFindStructureRequest) other; + return Objects.equals(this.linesToSample, that.linesToSample) + && Objects.equals(this.lineMergeSizeLimit, that.lineMergeSizeLimit) + && Objects.equals(this.timeout, that.timeout) + && Objects.equals(this.charset, that.charset) + && Objects.equals(this.format, that.format) + && Objects.equals(this.columnNames, that.columnNames) + && Objects.equals(this.hasHeaderRow, that.hasHeaderRow) + && Objects.equals(this.delimiter, that.delimiter) + && Objects.equals(this.grokPattern, that.grokPattern) + && Objects.equals(this.ecsCompatibility, that.ecsCompatibility) + && Objects.equals(this.timestampFormat, that.timestampFormat) + && Objects.equals(this.timestampField, that.timestampField); + } +} diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/textstructure/action/FindFieldStructureAction.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/textstructure/action/FindFieldStructureAction.java new file mode 100644 index 0000000000000..2e6f3af312e2b --- /dev/null +++ b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/textstructure/action/FindFieldStructureAction.java @@ -0,0 +1,98 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ +package org.elasticsearch.xpack.core.textstructure.action; + +import org.elasticsearch.action.ActionRequestValidationException; +import org.elasticsearch.action.ActionType; +import org.elasticsearch.common.Strings; +import org.elasticsearch.common.io.stream.StreamInput; +import org.elasticsearch.common.io.stream.StreamOutput; +import org.elasticsearch.xcontent.ParseField; + +import java.io.IOException; +import java.util.Objects; + +import static org.elasticsearch.action.ValidateActions.addValidationError; + +public class FindFieldStructureAction extends ActionType { + + public static final FindFieldStructureAction INSTANCE = new FindFieldStructureAction(); + public static final String NAME = "cluster:monitor/text_structure/find_field_structure"; + + private FindFieldStructureAction() { + super(NAME); + } + + public static class Request extends AbstractFindStructureRequest { + + public static final ParseField INDEX = new ParseField("index"); + public static final ParseField FIELD = new ParseField("field"); + + private String index; + private String field; + + public Request() {} + + public Request(StreamInput in) throws IOException { + super(in); + index = in.readString(); + field = in.readString(); + } + + public String getIndex() { + return index; + } + + public void setIndex(String index) { + this.index = index; + } + + public String getField() { + return field; + } + + public void setField(String field) { + this.field = field; + } + + @Override + public ActionRequestValidationException validate() { + ActionRequestValidationException validationException = super.validate(); + if (Strings.isNullOrEmpty(index)) { + validationException = addValidationError("index must be specified", validationException); + } + if (Strings.isNullOrEmpty(field)) { + validationException = addValidationError("field must be specified", validationException); + } + return validationException; + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + super.writeTo(out); + out.writeString(index); + out.writeString(field); + } + + @Override + public int hashCode() { + return Objects.hash(super.hashCode(), field, index); + } + + @Override + public boolean equals(Object other) { + if (this == other) { + return true; + } + if (other == null || getClass() != other.getClass()) { + return false; + } + Request that = (Request) other; + return super.equals(other) && Objects.equals(this.index, that.index) && Objects.equals(this.field, that.field); + } + } +} diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/textstructure/action/FindMessageStructureAction.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/textstructure/action/FindMessageStructureAction.java new file mode 100644 index 0000000000000..49035b36ff42c --- /dev/null +++ b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/textstructure/action/FindMessageStructureAction.java @@ -0,0 +1,97 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ +package org.elasticsearch.xpack.core.textstructure.action; + +import org.elasticsearch.action.ActionRequestValidationException; +import org.elasticsearch.action.ActionType; +import org.elasticsearch.common.io.stream.StreamInput; +import org.elasticsearch.common.io.stream.StreamOutput; +import org.elasticsearch.xcontent.ObjectParser; +import org.elasticsearch.xcontent.ParseField; +import org.elasticsearch.xcontent.XContentParser; + +import java.io.IOException; +import java.util.List; +import java.util.Objects; + +import static org.elasticsearch.action.ValidateActions.addValidationError; + +public class FindMessageStructureAction extends ActionType { + + public static final FindMessageStructureAction INSTANCE = new FindMessageStructureAction(); + public static final String NAME = "cluster:monitor/text_structure/find_message_structure"; + + private FindMessageStructureAction() { + super(NAME); + } + + public static class Request extends AbstractFindStructureRequest { + + public static final ParseField MESSAGES = new ParseField("messages"); + + private List messages; + + private static final ObjectParser PARSER = createParser(); + + private static ObjectParser createParser() { + ObjectParser parser = new ObjectParser<>("text_structure/find_message_structure", false, Request::new); + parser.declareStringArray(Request::setMessages, MESSAGES); + return parser; + } + + public Request() {} + + public Request(StreamInput in) throws IOException { + super(in); + messages = in.readStringCollectionAsList(); + } + + public static Request parseRequest(XContentParser parser) throws IOException { + return PARSER.parse(parser, null); + } + + public List getMessages() { + return messages; + } + + public void setMessages(List messages) { + this.messages = messages; + } + + @Override + public ActionRequestValidationException validate() { + ActionRequestValidationException validationException = super.validate(); + if (messages == null || messages.isEmpty()) { + validationException = addValidationError("messages must be specified", validationException); + } + return validationException; + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + super.writeTo(out); + out.writeStringCollection(messages); + } + + @Override + public int hashCode() { + return Objects.hash(super.hashCode(), messages); + } + + @Override + public boolean equals(Object other) { + if (this == other) { + return true; + } + if (other == null || getClass() != other.getClass()) { + return false; + } + Request that = (Request) other; + return super.equals(other) && Objects.equals(this.messages, that.messages); + } + } +} diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/textstructure/action/FindStructureAction.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/textstructure/action/FindStructureAction.java index 98bdff8cbced7..15aa3be46a675 100644 --- a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/textstructure/action/FindStructureAction.java +++ b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/textstructure/action/FindStructureAction.java @@ -6,290 +6,37 @@ */ package org.elasticsearch.xpack.core.textstructure.action; -import org.elasticsearch.TransportVersions; -import org.elasticsearch.action.ActionRequest; import org.elasticsearch.action.ActionRequestValidationException; -import org.elasticsearch.action.ActionResponse; import org.elasticsearch.action.ActionType; import org.elasticsearch.common.bytes.BytesReference; import org.elasticsearch.common.io.stream.StreamInput; import org.elasticsearch.common.io.stream.StreamOutput; -import org.elasticsearch.common.io.stream.Writeable; -import org.elasticsearch.core.TimeValue; -import org.elasticsearch.grok.GrokBuiltinPatterns; -import org.elasticsearch.xcontent.ParseField; -import org.elasticsearch.xcontent.ToXContentObject; -import org.elasticsearch.xcontent.XContentBuilder; -import org.elasticsearch.xpack.core.textstructure.structurefinder.TextStructure; import java.io.IOException; -import java.util.Arrays; -import java.util.List; -import java.util.Locale; import java.util.Objects; import static org.elasticsearch.action.ValidateActions.addValidationError; -public class FindStructureAction extends ActionType { +public class FindStructureAction extends ActionType { public static final FindStructureAction INSTANCE = new FindStructureAction(); public static final String NAME = "cluster:monitor/text_structure/findstructure"; - public static final int MIN_SAMPLE_LINE_COUNT = 2; - private FindStructureAction() { super(NAME); } - public static class Response extends ActionResponse implements ToXContentObject, Writeable { - - private final TextStructure textStructure; - - public Response(TextStructure textStructure) { - this.textStructure = textStructure; - } - - Response(StreamInput in) throws IOException { - super(in); - textStructure = new TextStructure(in); - } - - @Override - public void writeTo(StreamOutput out) throws IOException { - textStructure.writeTo(out); - } - - @Override - public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { - textStructure.toXContent(builder, params); - return builder; - } - - @Override - public int hashCode() { - return Objects.hash(textStructure); - } - - @Override - public boolean equals(Object other) { - - if (this == other) { - return true; - } - - if (other == null || getClass() != other.getClass()) { - return false; - } - - FindStructureAction.Response that = (FindStructureAction.Response) other; - return Objects.equals(textStructure, that.textStructure); - } - } - - public static class Request extends ActionRequest { - - public static final ParseField LINES_TO_SAMPLE = new ParseField("lines_to_sample"); - public static final ParseField LINE_MERGE_SIZE_LIMIT = new ParseField("line_merge_size_limit"); - public static final ParseField TIMEOUT = new ParseField("timeout"); - public static final ParseField CHARSET = TextStructure.CHARSET; - public static final ParseField FORMAT = TextStructure.FORMAT; - public static final ParseField COLUMN_NAMES = TextStructure.COLUMN_NAMES; - public static final ParseField HAS_HEADER_ROW = TextStructure.HAS_HEADER_ROW; - public static final ParseField DELIMITER = TextStructure.DELIMITER; - public static final ParseField QUOTE = TextStructure.QUOTE; - public static final ParseField SHOULD_TRIM_FIELDS = TextStructure.SHOULD_TRIM_FIELDS; - public static final ParseField GROK_PATTERN = TextStructure.GROK_PATTERN; - // This one is plural in FileStructure, but singular in FileStructureOverrides - public static final ParseField TIMESTAMP_FORMAT = new ParseField("timestamp_format"); - public static final ParseField TIMESTAMP_FIELD = TextStructure.TIMESTAMP_FIELD; + public static class Request extends AbstractFindStructureRequest { - public static final ParseField ECS_COMPATIBILITY = TextStructure.ECS_COMPATIBILITY; - - private static final String ARG_INCOMPATIBLE_WITH_FORMAT_TEMPLATE = "[%s] may only be specified if [" - + FORMAT.getPreferredName() - + "] is [%s]"; - - private Integer linesToSample; - private Integer lineMergeSizeLimit; - private TimeValue timeout; - private String charset; - private TextStructure.Format format; - private List columnNames; - private Boolean hasHeaderRow; - private Character delimiter; - private Character quote; - private Boolean shouldTrimFields; - private String grokPattern; - private String ecsCompatibility; - private String timestampFormat; - private String timestampField; private BytesReference sample; public Request() {} public Request(StreamInput in) throws IOException { super(in); - linesToSample = in.readOptionalVInt(); - lineMergeSizeLimit = in.readOptionalVInt(); - timeout = in.readOptionalTimeValue(); - charset = in.readOptionalString(); - format = in.readBoolean() ? in.readEnum(TextStructure.Format.class) : null; - columnNames = in.readBoolean() ? in.readStringCollectionAsList() : null; - hasHeaderRow = in.readOptionalBoolean(); - delimiter = in.readBoolean() ? (char) in.readVInt() : null; - quote = in.readBoolean() ? (char) in.readVInt() : null; - shouldTrimFields = in.readOptionalBoolean(); - grokPattern = in.readOptionalString(); - if (in.getTransportVersion().onOrAfter(TransportVersions.V_8_5_0)) { - ecsCompatibility = in.readOptionalString(); - } else { - ecsCompatibility = null; - } - timestampFormat = in.readOptionalString(); - timestampField = in.readOptionalString(); sample = in.readBytesReference(); } - public Integer getLinesToSample() { - return linesToSample; - } - - public void setLinesToSample(Integer linesToSample) { - this.linesToSample = linesToSample; - } - - public Integer getLineMergeSizeLimit() { - return lineMergeSizeLimit; - } - - public void setLineMergeSizeLimit(Integer lineMergeSizeLimit) { - this.lineMergeSizeLimit = lineMergeSizeLimit; - } - - public TimeValue getTimeout() { - return timeout; - } - - public void setTimeout(TimeValue timeout) { - this.timeout = timeout; - } - - public String getCharset() { - return charset; - } - - public void setCharset(String charset) { - this.charset = (charset == null || charset.isEmpty()) ? null : charset; - } - - public TextStructure.Format getFormat() { - return format; - } - - public void setFormat(TextStructure.Format format) { - this.format = format; - } - - public void setFormat(String format) { - this.format = (format == null || format.isEmpty()) ? null : TextStructure.Format.fromString(format); - } - - public List getColumnNames() { - return columnNames; - } - - public void setColumnNames(List columnNames) { - this.columnNames = (columnNames == null || columnNames.isEmpty()) ? null : columnNames; - } - - public void setColumnNames(String[] columnNames) { - this.columnNames = (columnNames == null || columnNames.length == 0) ? null : Arrays.asList(columnNames); - } - - public Boolean getHasHeaderRow() { - return hasHeaderRow; - } - - public void setHasHeaderRow(Boolean hasHeaderRow) { - this.hasHeaderRow = hasHeaderRow; - } - - public Character getDelimiter() { - return delimiter; - } - - public void setDelimiter(Character delimiter) { - this.delimiter = delimiter; - } - - public void setDelimiter(String delimiter) { - if (delimiter == null || delimiter.isEmpty()) { - this.delimiter = null; - } else if (delimiter.length() == 1) { - this.delimiter = delimiter.charAt(0); - } else { - throw new IllegalArgumentException(DELIMITER.getPreferredName() + " must be a single character"); - } - } - - public Character getQuote() { - return quote; - } - - public void setQuote(Character quote) { - this.quote = quote; - } - - public void setQuote(String quote) { - if (quote == null || quote.isEmpty()) { - this.quote = null; - } else if (quote.length() == 1) { - this.quote = quote.charAt(0); - } else { - throw new IllegalArgumentException(QUOTE.getPreferredName() + " must be a single character"); - } - } - - public Boolean getShouldTrimFields() { - return shouldTrimFields; - } - - public void setShouldTrimFields(Boolean shouldTrimFields) { - this.shouldTrimFields = shouldTrimFields; - } - - public String getGrokPattern() { - return grokPattern; - } - - public void setGrokPattern(String grokPattern) { - this.grokPattern = (grokPattern == null || grokPattern.isEmpty()) ? null : grokPattern; - } - - public String getEcsCompatibility() { - return ecsCompatibility; - } - - public void setEcsCompatibility(String ecsCompatibility) { - this.ecsCompatibility = (ecsCompatibility == null || ecsCompatibility.isEmpty()) ? null : ecsCompatibility; - } - - public String getTimestampFormat() { - return timestampFormat; - } - - public void setTimestampFormat(String timestampFormat) { - this.timestampFormat = (timestampFormat == null || timestampFormat.isEmpty()) ? null : timestampFormat; - } - - public String getTimestampField() { - return timestampField; - } - - public void setTimestampField(String timestampField) { - this.timestampField = (timestampField == null || timestampField.isEmpty()) ? null : timestampField; - } - public BytesReference getSample() { return sample; } @@ -298,70 +45,9 @@ public void setSample(BytesReference sample) { this.sample = sample; } - private static ActionRequestValidationException addIncompatibleArgError( - ParseField arg, - TextStructure.Format format, - ActionRequestValidationException validationException - ) { - return addValidationError( - String.format(Locale.ROOT, ARG_INCOMPATIBLE_WITH_FORMAT_TEMPLATE, arg.getPreferredName(), format), - validationException - ); - } - @Override public ActionRequestValidationException validate() { - ActionRequestValidationException validationException = null; - if (linesToSample != null && linesToSample < MIN_SAMPLE_LINE_COUNT) { - validationException = addValidationError( - "[" + LINES_TO_SAMPLE.getPreferredName() + "] must be at least [" + MIN_SAMPLE_LINE_COUNT + "] if specified", - validationException - ); - } - if (lineMergeSizeLimit != null && lineMergeSizeLimit <= 0) { - validationException = addValidationError( - "[" + LINE_MERGE_SIZE_LIMIT.getPreferredName() + "] must be positive if specified", - validationException - ); - } - if (format != TextStructure.Format.DELIMITED) { - if (columnNames != null) { - validationException = addIncompatibleArgError(COLUMN_NAMES, TextStructure.Format.DELIMITED, validationException); - } - if (hasHeaderRow != null) { - validationException = addIncompatibleArgError(HAS_HEADER_ROW, TextStructure.Format.DELIMITED, validationException); - } - if (delimiter != null) { - validationException = addIncompatibleArgError(DELIMITER, TextStructure.Format.DELIMITED, validationException); - } - if (quote != null) { - validationException = addIncompatibleArgError(QUOTE, TextStructure.Format.DELIMITED, validationException); - } - if (shouldTrimFields != null) { - validationException = addIncompatibleArgError(SHOULD_TRIM_FIELDS, TextStructure.Format.DELIMITED, validationException); - } - } - if (format != TextStructure.Format.SEMI_STRUCTURED_TEXT) { - if (grokPattern != null) { - validationException = addIncompatibleArgError( - GROK_PATTERN, - TextStructure.Format.SEMI_STRUCTURED_TEXT, - validationException - ); - } - } - - if (ecsCompatibility != null && GrokBuiltinPatterns.isValidEcsCompatibilityMode(ecsCompatibility) == false) { - validationException = addValidationError( - "[" - + ECS_COMPATIBILITY.getPreferredName() - + "] must be one of [" - + String.join(", ", GrokBuiltinPatterns.ECS_COMPATIBILITY_MODES) - + "] if specified", - validationException - ); - } - + ActionRequestValidationException validationException = super.validate(); if (sample == null || sample.length() == 0) { validationException = addValidationError("sample must be specified", validationException); } @@ -371,89 +57,24 @@ public ActionRequestValidationException validate() { @Override public void writeTo(StreamOutput out) throws IOException { super.writeTo(out); - out.writeOptionalVInt(linesToSample); - out.writeOptionalVInt(lineMergeSizeLimit); - out.writeOptionalTimeValue(timeout); - out.writeOptionalString(charset); - if (format == null) { - out.writeBoolean(false); - } else { - out.writeBoolean(true); - out.writeEnum(format); - } - if (columnNames == null) { - out.writeBoolean(false); - } else { - out.writeBoolean(true); - out.writeStringCollection(columnNames); - } - out.writeOptionalBoolean(hasHeaderRow); - if (delimiter == null) { - out.writeBoolean(false); - } else { - out.writeBoolean(true); - out.writeVInt(delimiter); - } - if (quote == null) { - out.writeBoolean(false); - } else { - out.writeBoolean(true); - out.writeVInt(quote); - } - out.writeOptionalBoolean(shouldTrimFields); - out.writeOptionalString(grokPattern); - if (out.getTransportVersion().onOrAfter(TransportVersions.V_8_5_0)) { - out.writeOptionalString(ecsCompatibility); - } - out.writeOptionalString(timestampFormat); - out.writeOptionalString(timestampField); out.writeBytesReference(sample); } @Override public int hashCode() { - return Objects.hash( - linesToSample, - lineMergeSizeLimit, - timeout, - charset, - format, - columnNames, - hasHeaderRow, - delimiter, - grokPattern, - ecsCompatibility, - timestampFormat, - timestampField, - sample - ); + return Objects.hash(super.hashCode(), sample); } @Override public boolean equals(Object other) { - if (this == other) { return true; } - if (other == null || getClass() != other.getClass()) { return false; } - Request that = (Request) other; - return Objects.equals(this.linesToSample, that.linesToSample) - && Objects.equals(this.lineMergeSizeLimit, that.lineMergeSizeLimit) - && Objects.equals(this.timeout, that.timeout) - && Objects.equals(this.charset, that.charset) - && Objects.equals(this.format, that.format) - && Objects.equals(this.columnNames, that.columnNames) - && Objects.equals(this.hasHeaderRow, that.hasHeaderRow) - && Objects.equals(this.delimiter, that.delimiter) - && Objects.equals(this.grokPattern, that.grokPattern) - && Objects.equals(this.ecsCompatibility, that.ecsCompatibility) - && Objects.equals(this.timestampFormat, that.timestampFormat) - && Objects.equals(this.timestampField, that.timestampField) - && Objects.equals(this.sample, that.sample); + return super.equals(other) && Objects.equals(this.sample, that.sample); } } } diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/textstructure/action/FindStructureResponse.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/textstructure/action/FindStructureResponse.java new file mode 100644 index 0000000000000..5848c2cbd0a1d --- /dev/null +++ b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/textstructure/action/FindStructureResponse.java @@ -0,0 +1,61 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.core.textstructure.action; + +import org.elasticsearch.action.ActionResponse; +import org.elasticsearch.common.io.stream.StreamInput; +import org.elasticsearch.common.io.stream.StreamOutput; +import org.elasticsearch.common.io.stream.Writeable; +import org.elasticsearch.xcontent.ToXContentObject; +import org.elasticsearch.xcontent.XContentBuilder; +import org.elasticsearch.xpack.core.textstructure.structurefinder.TextStructure; + +import java.io.IOException; +import java.util.Objects; + +public class FindStructureResponse extends ActionResponse implements ToXContentObject, Writeable { + + private final TextStructure textStructure; + + public FindStructureResponse(TextStructure textStructure) { + this.textStructure = textStructure; + } + + FindStructureResponse(StreamInput in) throws IOException { + super(in); + textStructure = new TextStructure(in); + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + textStructure.writeTo(out); + } + + @Override + public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { + textStructure.toXContent(builder, params); + return builder; + } + + @Override + public int hashCode() { + return Objects.hash(textStructure); + } + + @Override + public boolean equals(Object other) { + if (this == other) { + return true; + } + if (other == null || getClass() != other.getClass()) { + return false; + } + FindStructureResponse that = (FindStructureResponse) other; + return Objects.equals(textStructure, that.textStructure); + } +} diff --git a/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/textstructure/action/FindTextStructureActionResponseTests.java b/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/textstructure/action/FindTextStructureActionResponseTests.java deleted file mode 100644 index 31dbfc7dccff3..0000000000000 --- a/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/textstructure/action/FindTextStructureActionResponseTests.java +++ /dev/null @@ -1,29 +0,0 @@ -/* - * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one - * or more contributor license agreements. Licensed under the Elastic License - * 2.0; you may not use this file except in compliance with the Elastic License - * 2.0. - */ -package org.elasticsearch.xpack.core.textstructure.action; - -import org.elasticsearch.common.io.stream.Writeable; -import org.elasticsearch.test.AbstractWireSerializingTestCase; -import org.elasticsearch.xpack.core.textstructure.structurefinder.TextStructureTests; - -public class FindTextStructureActionResponseTests extends AbstractWireSerializingTestCase { - - @Override - protected FindStructureAction.Response createTestInstance() { - return new FindStructureAction.Response(TextStructureTests.createTestFileStructure()); - } - - @Override - protected FindStructureAction.Response mutateInstance(FindStructureAction.Response instance) { - return null;// TODO implement https://github.com/elastic/elasticsearch/issues/25929 - } - - @Override - protected Writeable.Reader instanceReader() { - return FindStructureAction.Response::new; - } -} diff --git a/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/textstructure/action/FindTextStructureResponseTests.java b/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/textstructure/action/FindTextStructureResponseTests.java new file mode 100644 index 0000000000000..887d75e3751c5 --- /dev/null +++ b/x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/textstructure/action/FindTextStructureResponseTests.java @@ -0,0 +1,33 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ +package org.elasticsearch.xpack.core.textstructure.action; + +import org.elasticsearch.common.io.stream.Writeable; +import org.elasticsearch.test.AbstractWireSerializingTestCase; +import org.elasticsearch.xpack.core.textstructure.structurefinder.TextStructureTests; + +public class FindTextStructureResponseTests extends AbstractWireSerializingTestCase { + + @Override + protected FindStructureResponse createTestInstance() { + return new FindStructureResponse(TextStructureTests.createTestFileStructure()); + } + + @Override + protected FindStructureResponse mutateInstance(FindStructureResponse response) { + FindStructureResponse newResponse; + do { + newResponse = createTestInstance(); + } while (response.equals(newResponse)); + return newResponse; + } + + @Override + protected Writeable.Reader instanceReader() { + return FindStructureResponse::new; + } +} diff --git a/x-pack/plugin/security/qa/operator-privileges-tests/src/javaRestTest/java/org/elasticsearch/xpack/security/operator/Constants.java b/x-pack/plugin/security/qa/operator-privileges-tests/src/javaRestTest/java/org/elasticsearch/xpack/security/operator/Constants.java index 2d743f562df8e..2250411fa7882 100644 --- a/x-pack/plugin/security/qa/operator-privileges-tests/src/javaRestTest/java/org/elasticsearch/xpack/security/operator/Constants.java +++ b/x-pack/plugin/security/qa/operator-privileges-tests/src/javaRestTest/java/org/elasticsearch/xpack/security/operator/Constants.java @@ -348,6 +348,8 @@ public class Constants { "cluster:monitor/task", "cluster:monitor/task/get", "cluster:monitor/tasks/lists", + "cluster:monitor/text_structure/find_field_structure", + "cluster:monitor/text_structure/find_message_structure", "cluster:monitor/text_structure/findstructure", "cluster:monitor/text_structure/test_grok_pattern", "cluster:monitor/transform/get", diff --git a/x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/text_structure/find_field_structure.yml b/x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/text_structure/find_field_structure.yml new file mode 100644 index 0000000000000..c2e9dbea1600a --- /dev/null +++ b/x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/text_structure/find_field_structure.yml @@ -0,0 +1,63 @@ +setup: + - do: + indices.create: + index: airlines + body: + mappings: + properties: + message: + type: text + - do: + bulk: + refresh: true + body: + - index: + _index: airlines + - message: "{\"airline\": \"AAL\", \"responsetime\": 132.2046, \"sourcetype\": \"text-structure-test\", \"time\": 1403481600}" + - index: + _index: airlines + - message: "{\"airline\": \"JZA\", \"responsetime\": 990.4628, \"sourcetype\": \"text-structure-test\", \"time\": 1403481700}" + - index: + _index: airlines + - message: "{\"airline\": \"AAL\", \"responsetime\": 134.2046, \"sourcetype\": \"text-structure-test\", \"time\": 1403481800}" +--- +"Field structure finder with JSON messages": + - do: + text_structure.find_field_structure: + index: airlines + field: message + documents_to_sample: 3 + timeout: 10s + - match: { num_lines_analyzed: 3 } + - match: { num_messages_analyzed: 3 } + - match: { charset: "UTF-8" } + - match: { has_byte_order_marker: null } + - match: { format: ndjson } + - match: { timestamp_field: time } + - match: { joda_timestamp_formats.0: UNIX } + - match: { java_timestamp_formats.0: UNIX } + - match: { need_client_timezone: false } + - match: { mappings.properties.airline.type: keyword } + - match: { mappings.properties.responsetime.type: double } + - match: { mappings.properties.sourcetype.type: keyword } + - match: { mappings.properties.time.type: date } + - match: { mappings.properties.time.format: epoch_second } + - match: { ingest_pipeline.description: "Ingest pipeline created by text structure finder" } + - match: { ingest_pipeline.processors.0.date.field: time } + - match: { ingest_pipeline.processors.0.date.formats.0: UNIX } + - match: { field_stats.airline.count: 3 } + - match: { field_stats.airline.cardinality: 2 } + - match: { field_stats.responsetime.count: 3 } + - match: { field_stats.responsetime.cardinality: 3 } + - match: { field_stats.responsetime.min_value: 132.2046 } + - match: { field_stats.responsetime.max_value: 990.4628 } + # Not asserting on field_stats.responsetime.mean as it's a recurring decimal + # so its representation in the response could cause spurious failures + - match: { field_stats.responsetime.median_value: 134.2046 } + - match: { field_stats.sourcetype.count: 3 } + - match: { field_stats.sourcetype.cardinality: 1 } + - match: { field_stats.time.count: 3 } + - match: { field_stats.time.cardinality: 3 } + - match: { field_stats.time.earliest: "1403481600" } + - match: { field_stats.time.latest: "1403481800" } + - is_false: explanation diff --git a/x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/text_structure/find_message_structure.yml b/x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/text_structure/find_message_structure.yml new file mode 100644 index 0000000000000..b1000510f2972 --- /dev/null +++ b/x-pack/plugin/src/yamlRestTest/resources/rest-api-spec/test/text_structure/find_message_structure.yml @@ -0,0 +1,56 @@ +"Messages structure finder with JSON messages": + - do: + text_structure.find_message_structure: + timeout: 10s + body: + messages: + - "{\"airline\": \"AAL\", \"responsetime\": 132.2046, \"sourcetype\": \"text-structure-test\", \"time\": 1403481600}" + - "{\"airline\": \"JZA\", \"responsetime\": 990.4628, \"sourcetype\": \"text-structure-test\", \"time\": 1403481700}" + - "{\"airline\": \"AAL\", \"responsetime\": 134.2046, \"sourcetype\": \"text-structure-test\", \"time\": 1403481800}" + - match: { num_lines_analyzed: 3 } + - match: { num_messages_analyzed: 3 } + - match: { charset: "UTF-8" } + - match: { has_byte_order_marker: null } + - match: { format: ndjson } + - match: { timestamp_field: time } + - match: { joda_timestamp_formats.0: UNIX } + - match: { java_timestamp_formats.0: UNIX } + - match: { need_client_timezone: false } + - match: { mappings.properties.airline.type: keyword } + - match: { mappings.properties.responsetime.type: double } + - match: { mappings.properties.sourcetype.type: keyword } + - match: { mappings.properties.time.type: date } + - match: { mappings.properties.time.format: epoch_second } + - match: { ingest_pipeline.description: "Ingest pipeline created by text structure finder" } + - match: { ingest_pipeline.processors.0.date.field: time } + - match: { ingest_pipeline.processors.0.date.formats.0: UNIX } + - match: { field_stats.airline.count: 3 } + - match: { field_stats.airline.cardinality: 2 } + - match: { field_stats.responsetime.count: 3 } + - match: { field_stats.responsetime.cardinality: 3 } + - match: { field_stats.responsetime.min_value: 132.2046 } + - match: { field_stats.responsetime.max_value: 990.4628 } + # Not asserting on field_stats.responsetime.mean as it's a recurring decimal + # so its representation in the response could cause spurious failures + - match: { field_stats.responsetime.median_value: 134.2046 } + - match: { field_stats.sourcetype.count: 3 } + - match: { field_stats.sourcetype.cardinality: 1 } + - match: { field_stats.time.count: 3 } + - match: { field_stats.time.cardinality: 3 } + - match: { field_stats.time.earliest: "1403481600" } + - match: { field_stats.time.latest: "1403481800" } + - is_false: explanation +--- +"Messages structure finder with log messages": + - do: + text_structure.find_message_structure: + timeout: 10s + body: + messages: + - "2019-05-16 16:56:14 line 1 abcdefghijklmnopqrstuvwxyz" + - "2019-05-16 16:56:14 line 2 abcdefghijklmnopqrstuvwxyz\ncontinuation...\ncontinuation...\n" + - "2019-05-16 16:56:14 line 3 abcdefghijklmnopqrstuvwxyz" + - match: { num_lines_analyzed: 3 } + - match: { num_messages_analyzed: 3 } + - match: { format: semi_structured_text } + - match: { grok_pattern: "%{TIMESTAMP_ISO8601:timestamp} .*? %{INT:field} .*" } diff --git a/x-pack/plugin/text-structure/qa/text-structure-with-security/build.gradle b/x-pack/plugin/text-structure/qa/text-structure-with-security/build.gradle index 5fc76885aa7eb..1e592615da1f2 100644 --- a/x-pack/plugin/text-structure/qa/text-structure-with-security/build.gradle +++ b/x-pack/plugin/text-structure/qa/text-structure-with-security/build.gradle @@ -9,7 +9,7 @@ dependencies { restResources { restApi { // needed for template installation, etc. - include '_common', 'indices', 'text_structure' + include '_common', 'bulk', 'indices', 'text_structure' } restTests { includeXpack 'text_structure' diff --git a/x-pack/plugin/text-structure/qa/text-structure-with-security/roles.yml b/x-pack/plugin/text-structure/qa/text-structure-with-security/roles.yml index 7eff54728320a..7095acb3c60a1 100644 --- a/x-pack/plugin/text-structure/qa/text-structure-with-security/roles.yml +++ b/x-pack/plugin/text-structure/qa/text-structure-with-security/roles.yml @@ -6,3 +6,15 @@ minimal: # This is always required because the REST client uses it to find the version of # Elasticsearch it's talking to - cluster:monitor/main + indices: + # Give all users involved in these tests access to the indices where the data to + # be analyzed is stored. + - names: [ 'airlines' ] + privileges: + - create_index + - indices:admin/refresh + - read + - write + - view_index_metadata + - indices:data/write/bulk + - indices:data/write/index diff --git a/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/TextStructurePlugin.java b/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/TextStructurePlugin.java index 2a2fe1ea5a55a..07e49989b9f09 100644 --- a/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/TextStructurePlugin.java +++ b/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/TextStructurePlugin.java @@ -21,10 +21,16 @@ import org.elasticsearch.plugins.Plugin; import org.elasticsearch.rest.RestController; import org.elasticsearch.rest.RestHandler; +import org.elasticsearch.xpack.core.textstructure.action.FindFieldStructureAction; +import org.elasticsearch.xpack.core.textstructure.action.FindMessageStructureAction; import org.elasticsearch.xpack.core.textstructure.action.FindStructureAction; import org.elasticsearch.xpack.core.textstructure.action.TestGrokPatternAction; +import org.elasticsearch.xpack.textstructure.rest.RestFindFieldStructureAction; +import org.elasticsearch.xpack.textstructure.rest.RestFindMessageStructureAction; import org.elasticsearch.xpack.textstructure.rest.RestFindStructureAction; import org.elasticsearch.xpack.textstructure.rest.RestTestGrokPatternAction; +import org.elasticsearch.xpack.textstructure.transport.TransportFindFieldStructureAction; +import org.elasticsearch.xpack.textstructure.transport.TransportFindMessageStructureAction; import org.elasticsearch.xpack.textstructure.transport.TransportFindStructureAction; import org.elasticsearch.xpack.textstructure.transport.TransportTestGrokPatternAction; @@ -53,12 +59,19 @@ public List getRestHandlers( Supplier nodesInCluster, Predicate clusterSupportsFeature ) { - return Arrays.asList(new RestFindStructureAction(), new RestTestGrokPatternAction()); + return Arrays.asList( + new RestFindFieldStructureAction(), + new RestFindMessageStructureAction(), + new RestFindStructureAction(), + new RestTestGrokPatternAction() + ); } @Override public List> getActions() { return Arrays.asList( + new ActionHandler<>(FindFieldStructureAction.INSTANCE, TransportFindFieldStructureAction.class), + new ActionHandler<>(FindMessageStructureAction.INSTANCE, TransportFindMessageStructureAction.class), new ActionHandler<>(FindStructureAction.INSTANCE, TransportFindStructureAction.class), new ActionHandler<>(TestGrokPatternAction.INSTANCE, TransportTestGrokPatternAction.class) ); diff --git a/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/rest/RestFindFieldStructureAction.java b/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/rest/RestFindFieldStructureAction.java new file mode 100644 index 0000000000000..0f81a4fc9726b --- /dev/null +++ b/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/rest/RestFindFieldStructureAction.java @@ -0,0 +1,51 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ +package org.elasticsearch.xpack.textstructure.rest; + +import org.elasticsearch.client.internal.node.NodeClient; +import org.elasticsearch.rest.BaseRestHandler; +import org.elasticsearch.rest.RestRequest; +import org.elasticsearch.rest.Scope; +import org.elasticsearch.rest.ServerlessScope; +import org.elasticsearch.rest.action.RestToXContentListener; +import org.elasticsearch.xpack.core.textstructure.action.FindFieldStructureAction; +import org.elasticsearch.xpack.core.textstructure.structurefinder.TextStructure; + +import java.util.Collections; +import java.util.List; +import java.util.Set; + +import static org.elasticsearch.rest.RestRequest.Method.GET; +import static org.elasticsearch.xpack.textstructure.TextStructurePlugin.BASE_PATH; + +@ServerlessScope(Scope.INTERNAL) +public class RestFindFieldStructureAction extends BaseRestHandler { + + @Override + public List routes() { + return List.of(new Route(GET, BASE_PATH + "find_field_structure")); + } + + @Override + public String getName() { + return "text_structure_find_field_structure_action"; + } + + @Override + protected RestChannelConsumer prepareRequest(RestRequest restRequest, NodeClient client) { + FindFieldStructureAction.Request request = new FindFieldStructureAction.Request(); + RestFindStructureArgumentsParser.parse(restRequest, request); + request.setIndex(restRequest.param(FindFieldStructureAction.Request.INDEX.getPreferredName())); + request.setField(restRequest.param(FindFieldStructureAction.Request.FIELD.getPreferredName())); + return channel -> client.execute(FindFieldStructureAction.INSTANCE, request, new RestToXContentListener<>(channel)); + } + + @Override + protected Set responseParams() { + return Collections.singleton(TextStructure.EXPLAIN); + } +} diff --git a/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/rest/RestFindMessageStructureAction.java b/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/rest/RestFindMessageStructureAction.java new file mode 100644 index 0000000000000..cc607dbdcd646 --- /dev/null +++ b/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/rest/RestFindMessageStructureAction.java @@ -0,0 +1,55 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ +package org.elasticsearch.xpack.textstructure.rest; + +import org.elasticsearch.client.internal.node.NodeClient; +import org.elasticsearch.rest.BaseRestHandler; +import org.elasticsearch.rest.RestRequest; +import org.elasticsearch.rest.Scope; +import org.elasticsearch.rest.ServerlessScope; +import org.elasticsearch.rest.action.RestToXContentListener; +import org.elasticsearch.xcontent.XContentParser; +import org.elasticsearch.xpack.core.textstructure.action.FindMessageStructureAction; +import org.elasticsearch.xpack.core.textstructure.structurefinder.TextStructure; + +import java.io.IOException; +import java.util.Collections; +import java.util.List; +import java.util.Set; + +import static org.elasticsearch.rest.RestRequest.Method.GET; +import static org.elasticsearch.rest.RestRequest.Method.POST; +import static org.elasticsearch.xpack.textstructure.TextStructurePlugin.BASE_PATH; + +@ServerlessScope(Scope.INTERNAL) +public class RestFindMessageStructureAction extends BaseRestHandler { + + @Override + public List routes() { + return List.of(new Route(GET, BASE_PATH + "find_message_structure"), new Route(POST, BASE_PATH + "find_message_structure")); + } + + @Override + public String getName() { + return "text_structure_find_message_structure_action"; + } + + @Override + protected RestChannelConsumer prepareRequest(RestRequest restRequest, NodeClient client) throws IOException { + FindMessageStructureAction.Request request; + try (XContentParser parser = restRequest.contentOrSourceParamParser()) { + request = FindMessageStructureAction.Request.parseRequest(parser); + } + RestFindStructureArgumentsParser.parse(restRequest, request); + return channel -> client.execute(FindMessageStructureAction.INSTANCE, request, new RestToXContentListener<>(channel)); + } + + @Override + protected Set responseParams() { + return Collections.singleton(TextStructure.EXPLAIN); + } +} diff --git a/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/rest/RestFindStructureAction.java b/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/rest/RestFindStructureAction.java index 94aee3c2a5f49..65325f2268ed2 100644 --- a/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/rest/RestFindStructureAction.java +++ b/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/rest/RestFindStructureAction.java @@ -9,7 +9,6 @@ import org.elasticsearch.ElasticsearchParseException; import org.elasticsearch.client.internal.node.NodeClient; import org.elasticsearch.core.RestApiVersion; -import org.elasticsearch.core.TimeValue; import org.elasticsearch.rest.BaseRestHandler; import org.elasticsearch.rest.RestRequest; import org.elasticsearch.rest.Scope; @@ -17,12 +16,10 @@ import org.elasticsearch.rest.action.RestToXContentListener; import org.elasticsearch.xpack.core.textstructure.action.FindStructureAction; import org.elasticsearch.xpack.core.textstructure.structurefinder.TextStructure; -import org.elasticsearch.xpack.textstructure.structurefinder.TextStructureFinderManager; import java.util.Collections; import java.util.List; import java.util.Set; -import java.util.concurrent.TimeUnit; import static org.elasticsearch.rest.RestRequest.Method.POST; import static org.elasticsearch.xpack.textstructure.TextStructurePlugin.BASE_PATH; @@ -30,8 +27,6 @@ @ServerlessScope(Scope.INTERNAL) public class RestFindStructureAction extends BaseRestHandler { - private static final TimeValue DEFAULT_TIMEOUT = new TimeValue(25, TimeUnit.SECONDS); - @Override public List routes() { return List.of( @@ -46,38 +41,9 @@ public String getName() { @Override protected RestChannelConsumer prepareRequest(RestRequest restRequest, NodeClient client) { - FindStructureAction.Request request = new FindStructureAction.Request(); - request.setLinesToSample( - restRequest.paramAsInt( - FindStructureAction.Request.LINES_TO_SAMPLE.getPreferredName(), - TextStructureFinderManager.DEFAULT_IDEAL_SAMPLE_LINE_COUNT - ) - ); - request.setLineMergeSizeLimit( - restRequest.paramAsInt( - FindStructureAction.Request.LINE_MERGE_SIZE_LIMIT.getPreferredName(), - TextStructureFinderManager.DEFAULT_LINE_MERGE_SIZE_LIMIT - ) - ); - request.setTimeout( - TimeValue.parseTimeValue( - restRequest.param(FindStructureAction.Request.TIMEOUT.getPreferredName()), - DEFAULT_TIMEOUT, - FindStructureAction.Request.TIMEOUT.getPreferredName() - ) - ); - request.setCharset(restRequest.param(FindStructureAction.Request.CHARSET.getPreferredName())); - request.setFormat(restRequest.param(FindStructureAction.Request.FORMAT.getPreferredName())); - request.setColumnNames(restRequest.paramAsStringArray(FindStructureAction.Request.COLUMN_NAMES.getPreferredName(), null)); - request.setHasHeaderRow(restRequest.paramAsBoolean(FindStructureAction.Request.HAS_HEADER_ROW.getPreferredName(), null)); - request.setDelimiter(restRequest.param(FindStructureAction.Request.DELIMITER.getPreferredName())); - request.setQuote(restRequest.param(FindStructureAction.Request.QUOTE.getPreferredName())); - request.setShouldTrimFields(restRequest.paramAsBoolean(FindStructureAction.Request.SHOULD_TRIM_FIELDS.getPreferredName(), null)); - request.setGrokPattern(restRequest.param(FindStructureAction.Request.GROK_PATTERN.getPreferredName())); - request.setEcsCompatibility(restRequest.param(FindStructureAction.Request.ECS_COMPATIBILITY.getPreferredName())); - request.setTimestampFormat(restRequest.param(FindStructureAction.Request.TIMESTAMP_FORMAT.getPreferredName())); - request.setTimestampField(restRequest.param(FindStructureAction.Request.TIMESTAMP_FIELD.getPreferredName())); + RestFindStructureArgumentsParser.parse(restRequest, request); + if (restRequest.hasContent()) { request.setSample(restRequest.content()); } else { diff --git a/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/rest/RestFindStructureArgumentsParser.java b/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/rest/RestFindStructureArgumentsParser.java new file mode 100644 index 0000000000000..bd6fe553fc447 --- /dev/null +++ b/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/rest/RestFindStructureArgumentsParser.java @@ -0,0 +1,73 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.textstructure.rest; + +import org.elasticsearch.core.TimeValue; +import org.elasticsearch.rest.RestRequest; +import org.elasticsearch.xpack.core.textstructure.action.AbstractFindStructureRequest; +import org.elasticsearch.xpack.core.textstructure.action.FindFieldStructureAction; +import org.elasticsearch.xpack.core.textstructure.action.FindMessageStructureAction; +import org.elasticsearch.xpack.core.textstructure.action.FindStructureAction; +import org.elasticsearch.xpack.core.textstructure.structurefinder.TextStructure; +import org.elasticsearch.xpack.textstructure.structurefinder.TextStructureFinderManager; + +import java.util.concurrent.TimeUnit; + +public class RestFindStructureArgumentsParser { + + private static final TimeValue DEFAULT_TIMEOUT = new TimeValue(25, TimeUnit.SECONDS); + + static void parse(RestRequest restRequest, AbstractFindStructureRequest request) { + if (request instanceof FindStructureAction.Request) { + request.setLinesToSample( + restRequest.paramAsInt( + FindStructureAction.Request.LINES_TO_SAMPLE.getPreferredName(), + TextStructureFinderManager.DEFAULT_IDEAL_SAMPLE_LINE_COUNT + ) + ); + request.setLineMergeSizeLimit( + restRequest.paramAsInt( + FindStructureAction.Request.LINE_MERGE_SIZE_LIMIT.getPreferredName(), + TextStructureFinderManager.DEFAULT_LINE_MERGE_SIZE_LIMIT + ) + ); + request.setCharset(restRequest.param(FindStructureAction.Request.CHARSET.getPreferredName())); + request.setHasHeaderRow(restRequest.paramAsBoolean(FindStructureAction.Request.HAS_HEADER_ROW.getPreferredName(), null)); + } else if (request instanceof FindFieldStructureAction.Request) { + request.setLinesToSample( + restRequest.paramAsInt( + FindStructureAction.Request.DOCUMENTS_TO_SAMPLE.getPreferredName(), + TextStructureFinderManager.DEFAULT_IDEAL_SAMPLE_LINE_COUNT + ) + ); + } + + request.setTimeout( + TimeValue.parseTimeValue( + restRequest.param(FindStructureAction.Request.TIMEOUT.getPreferredName()), + DEFAULT_TIMEOUT, + FindStructureAction.Request.TIMEOUT.getPreferredName() + ) + ); + request.setFormat(restRequest.param(FindStructureAction.Request.FORMAT.getPreferredName())); + request.setColumnNames(restRequest.paramAsStringArray(FindStructureAction.Request.COLUMN_NAMES.getPreferredName(), null)); + request.setDelimiter(restRequest.param(FindStructureAction.Request.DELIMITER.getPreferredName())); + request.setQuote(restRequest.param(FindStructureAction.Request.QUOTE.getPreferredName())); + request.setShouldTrimFields(restRequest.paramAsBoolean(FindStructureAction.Request.SHOULD_TRIM_FIELDS.getPreferredName(), null)); + request.setGrokPattern(restRequest.param(FindStructureAction.Request.GROK_PATTERN.getPreferredName())); + request.setEcsCompatibility(restRequest.param(FindStructureAction.Request.ECS_COMPATIBILITY.getPreferredName())); + request.setTimestampFormat(restRequest.param(FindStructureAction.Request.TIMESTAMP_FORMAT.getPreferredName())); + request.setTimestampField(restRequest.param(FindStructureAction.Request.TIMESTAMP_FIELD.getPreferredName())); + + if (request instanceof FindMessageStructureAction.Request || request instanceof FindFieldStructureAction.Request) { + if (TextStructure.Format.DELIMITED.equals(request.getFormat())) { + request.setHasHeaderRow(false); + } + } + } +} diff --git a/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/structurefinder/DelimitedTextStructureFinder.java b/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/structurefinder/DelimitedTextStructureFinder.java index 6d7faaadae433..7fc6db9cb5c6f 100644 --- a/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/structurefinder/DelimitedTextStructureFinder.java +++ b/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/structurefinder/DelimitedTextStructureFinder.java @@ -44,7 +44,7 @@ public class DelimitedTextStructureFinder implements TextStructureFinder { private final List sampleMessages; private final TextStructure structure; - static DelimitedTextStructureFinder makeDelimitedTextStructureFinder( + static DelimitedTextStructureFinder createFromSample( List explanation, String sample, String charsetName, @@ -590,6 +590,36 @@ static boolean lineHasUnescapedQuote(String line, CsvPreference csvPreference) { return false; } + static boolean canCreateFromMessages( + List explanation, + List messages, + int minFieldsPerRow, + CsvPreference csvPreference, + String formatName, + double allowedFractionOfBadLines + ) { + for (String message : messages) { + try (CsvListReader csvReader = new CsvListReader(new StringReader(message), csvPreference)) { + if (csvReader.read() == null) { + explanation.add(format("Not %s because message with no lines: [%s]", formatName, message)); + return false; + } + if (csvReader.read() != null) { + explanation.add(format("Not %s because message with multiple lines: [%s]", formatName, message)); + return false; + } + } catch (IOException e) { + explanation.add(format("Not %s because there was a parsing exception: [%s]", formatName, e.getMessage())); + return false; + } + } + + // Every line contains a single valid delimited message, so + // we can safely concatenate and run the logic for a sample. + String sample = String.join("\n", messages); + return canCreateFromSample(explanation, sample, minFieldsPerRow, csvPreference, formatName, allowedFractionOfBadLines); + } + static boolean canCreateFromSample( List explanation, String sample, @@ -598,7 +628,6 @@ static boolean canCreateFromSample( String formatName, double allowedFractionOfBadLines ) { - // Logstash's CSV parser won't tolerate fields where just part of the // value is quoted, whereas SuperCSV will, hence this extra check String[] sampleLines = sample.split("\n"); @@ -619,7 +648,6 @@ static boolean canCreateFromSample( try (CsvListReader csvReader = new CsvListReader(new StringReader(sample), csvPreference)) { int fieldsInFirstRow = -1; - int fieldsInLastRow = -1; List illFormattedRows = new ArrayList<>(); int numberOfRows = 0; @@ -643,7 +671,6 @@ static boolean canCreateFromSample( ); return false; } - fieldsInLastRow = fieldsInFirstRow; continue; } @@ -676,26 +703,7 @@ static boolean canCreateFromSample( ); return false; } - continue; } - - fieldsInLastRow = fieldsInThisRow; - } - - if (fieldsInLastRow > fieldsInFirstRow) { - explanation.add( - "Not " - + formatName - + " because last row has more fields than first row: [" - + fieldsInFirstRow - + "] and [" - + fieldsInLastRow - + "]" - ); - return false; - } - if (fieldsInLastRow < fieldsInFirstRow) { - --numberOfRows; } } catch (SuperCsvException e) { // Tolerate an incomplete last row diff --git a/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/structurefinder/DelimitedTextStructureFinderFactory.java b/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/structurefinder/DelimitedTextStructureFinderFactory.java index f809665199fea..5f09fdb437fe4 100644 --- a/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/structurefinder/DelimitedTextStructureFinderFactory.java +++ b/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/structurefinder/DelimitedTextStructureFinderFactory.java @@ -67,6 +67,22 @@ public boolean canCreateFromSample(List explanation, String sample, doub ); } + public boolean canCreateFromMessages(List explanation, List messages, double allowedFractionOfBadLines) { + String formatName = switch ((char) csvPreference.getDelimiterChar()) { + case ',' -> "CSV"; + case '\t' -> "TSV"; + default -> Character.getName(csvPreference.getDelimiterChar()).toLowerCase(Locale.ROOT) + " delimited values"; + }; + return DelimitedTextStructureFinder.canCreateFromMessages( + explanation, + messages, + minFieldsPerRow, + csvPreference, + formatName, + allowedFractionOfBadLines + ); + } + @Override public TextStructureFinder createFromSample( List explanation, @@ -78,7 +94,7 @@ public TextStructureFinder createFromSample( TimeoutChecker timeoutChecker ) throws IOException { CsvPreference adjustedCsvPreference = new CsvPreference.Builder(csvPreference).maxLinesPerRow(lineMergeSizeLimit).build(); - return DelimitedTextStructureFinder.makeDelimitedTextStructureFinder( + return DelimitedTextStructureFinder.createFromSample( explanation, sample, charsetName, @@ -89,4 +105,26 @@ public TextStructureFinder createFromSample( timeoutChecker ); } + + public TextStructureFinder createFromMessages( + List explanation, + List messages, + TextStructureOverrides overrides, + TimeoutChecker timeoutChecker + ) throws IOException { + // DelimitedTextStructureFinderFactory::canCreateFromMessages already + // checked that every line contains a single valid delimited message, + // so we can safely concatenate and run the logic for a sample. + String sample = String.join("\n", messages); + return DelimitedTextStructureFinder.createFromSample( + explanation, + sample, + "UTF-8", + null, + csvPreference, + trimFields, + overrides, + timeoutChecker + ); + } } diff --git a/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/structurefinder/LogTextStructureFinder.java b/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/structurefinder/LogTextStructureFinder.java index 4e01d32645008..c9ca6002b6c03 100644 --- a/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/structurefinder/LogTextStructureFinder.java +++ b/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/structurefinder/LogTextStructureFinder.java @@ -36,7 +36,6 @@ private static LogTextStructureFinder makeSingleLineLogTextStructureFinder( String[] sampleLines, String charsetName, Boolean hasByteOrderMarker, - int lineMergeSizeLimit, TextStructureOverrides overrides, TimeoutChecker timeoutChecker ) { @@ -108,12 +107,9 @@ private static LogTextStructureFinder makeSingleLineLogTextStructureFinder( return new LogTextStructureFinder(sampleMessages, structure); } - private static LogTextStructureFinder makeMultiLineLogTextStructureFinder( + private static TimestampFormatFinder getTimestampFormatFinder( List explanation, String[] sampleLines, - String charsetName, - Boolean hasByteOrderMarker, - int lineMergeSizeLimit, TextStructureOverrides overrides, TimeoutChecker timeoutChecker ) { @@ -145,15 +141,20 @@ private static LogTextStructureFinder makeMultiLineLogTextStructureFinder( + timestampFormatFinder.getJavaTimestampFormats() ); + return timestampFormatFinder; + } + + private static Tuple, Integer> getSampleMessages( + String multiLineRegex, + String[] sampleLines, + int lineMergeSizeLimit, + TimeoutChecker timeoutChecker + ) { List sampleMessages = new ArrayList<>(); - StringBuilder preamble = new StringBuilder(); int linesConsumed = 0; StringBuilder message = null; int linesInMessage = 0; - String multiLineRegex = createMultiLineMessageStartRegex( - timestampFormatFinder.getPrefaces(), - timestampFormatFinder.getSimplePattern().pattern() - ); + Pattern multiLinePattern = Pattern.compile(multiLineRegex); for (String sampleLine : sampleLines) { if (multiLinePattern.matcher(sampleLine).find()) { @@ -195,9 +196,6 @@ private static LogTextStructureFinder makeMultiLineLogTextStructureFinder( } } timeoutChecker.check("multi-line message determination"); - if (sampleMessages.size() < 2) { - preamble.append(sampleLine).append('\n'); - } } // Don't add the last message, as it might be partial and mess up subsequent pattern finding @@ -209,8 +207,24 @@ private static LogTextStructureFinder makeMultiLineLogTextStructureFinder( ); } - // null to allow GC before Grok pattern search - sampleLines = null; + return new Tuple<>(sampleMessages, linesConsumed); + } + + private static LogTextStructureFinder makeMultiLineLogTextStructureFinder( + List explanation, + List sampleMessages, + String charsetName, + Boolean hasByteOrderMarker, + TextStructureOverrides overrides, + int linesConsumed, + TimestampFormatFinder timestampFormatFinder, + String multiLineRegex, + TimeoutChecker timeoutChecker + ) { + StringBuilder preamble = new StringBuilder(); + for (int i = 0; i < sampleMessages.size() && i < 2; i++) { + preamble.append(sampleMessages.get(i)).append('\n'); + } TextStructure.Builder structureBuilder = new TextStructure.Builder(TextStructure.Format.SEMI_STRUCTURED_TEXT).setCharset( charsetName @@ -300,6 +314,80 @@ private static LogTextStructureFinder makeMultiLineLogTextStructureFinder( return new LogTextStructureFinder(sampleMessages, structure); } + private static LogTextStructureFinder makeMultiLineLogTextStructureFinder( + List explanation, + String[] sampleLines, + String charsetName, + Boolean hasByteOrderMarker, + int lineMergeSizeLimit, + TextStructureOverrides overrides, + TimeoutChecker timeoutChecker + ) { + TimestampFormatFinder timestampFormatFinder = getTimestampFormatFinder(explanation, sampleLines, overrides, timeoutChecker); + + String multiLineRegex = createMultiLineMessageStartRegex( + timestampFormatFinder.getPrefaces(), + timestampFormatFinder.getSimplePattern().pattern() + ); + + Tuple, Integer> sampleMessagesAndLinesConsumed = getSampleMessages( + multiLineRegex, + sampleLines, + lineMergeSizeLimit, + timeoutChecker + ); + List sampleMessages = sampleMessagesAndLinesConsumed.v1(); + int linesConsumed = sampleMessagesAndLinesConsumed.v2(); + + // null to allow GC before Grok pattern search + sampleLines = null; + + return makeMultiLineLogTextStructureFinder( + explanation, + sampleMessages, + charsetName, + hasByteOrderMarker, + overrides, + linesConsumed, + timestampFormatFinder, + multiLineRegex, + timeoutChecker + ); + } + + private static LogTextStructureFinder makeMultiLineLogTextStructureFinder( + List explanation, + List messages, + String charsetName, + Boolean hasByteOrderMarker, + TextStructureOverrides overrides, + TimeoutChecker timeoutChecker + ) { + TimestampFormatFinder timestampFormatFinder = getTimestampFormatFinder( + explanation, + messages.toArray(new String[0]), + overrides, + timeoutChecker + ); + + String multiLineRegex = createMultiLineMessageStartRegex( + timestampFormatFinder.getPrefaces(), + timestampFormatFinder.getSimplePattern().pattern() + ); + + return makeMultiLineLogTextStructureFinder( + explanation, + messages, + charsetName, + hasByteOrderMarker, + overrides, + messages.size(), + timestampFormatFinder, + multiLineRegex, + timeoutChecker + ); + } + static LogTextStructureFinder makeLogTextStructureFinder( List explanation, String sample, @@ -316,7 +404,6 @@ static LogTextStructureFinder makeLogTextStructureFinder( sampleLines, charsetName, hasByteOrderMarker, - lineMergeSizeLimit, overrides, timeoutChecker ); @@ -333,6 +420,28 @@ static LogTextStructureFinder makeLogTextStructureFinder( } } + static LogTextStructureFinder makeLogTextStructureFinder( + List explanation, + List messages, + String charsetName, + Boolean hasByteOrderMarker, + TextStructureOverrides overrides, + TimeoutChecker timeoutChecker + ) { + if (TextStructureUtils.NULL_TIMESTAMP_FORMAT.equals(overrides.getTimestampFormat())) { + return makeSingleLineLogTextStructureFinder( + explanation, + messages.toArray(new String[0]), + charsetName, + hasByteOrderMarker, + overrides, + timeoutChecker + ); + } else { + return makeMultiLineLogTextStructureFinder(explanation, messages, charsetName, hasByteOrderMarker, overrides, timeoutChecker); + } + } + private LogTextStructureFinder(List sampleMessages, TextStructure structure) { this.sampleMessages = Collections.unmodifiableList(sampleMessages); this.structure = structure; diff --git a/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/structurefinder/LogTextStructureFinderFactory.java b/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/structurefinder/LogTextStructureFinderFactory.java index d3978946ce908..24532e9fdaae4 100644 --- a/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/structurefinder/LogTextStructureFinderFactory.java +++ b/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/structurefinder/LogTextStructureFinderFactory.java @@ -40,6 +40,10 @@ public boolean canCreateFromSample(List explanation, String sample, doub return true; } + public boolean canCreateFromMessages(List explanation, List messages, double allowedFractionOfBadLines) { + return true; + } + @Override public TextStructureFinder createFromSample( List explanation, @@ -60,4 +64,13 @@ public TextStructureFinder createFromSample( timeoutChecker ); } + + public TextStructureFinder createFromMessages( + List explanation, + List messages, + TextStructureOverrides overrides, + TimeoutChecker timeoutChecker + ) { + return LogTextStructureFinder.makeLogTextStructureFinder(explanation, messages, "UTF-8", null, overrides, timeoutChecker); + } } diff --git a/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/structurefinder/NdJsonTextStructureFinderFactory.java b/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/structurefinder/NdJsonTextStructureFinderFactory.java index 5afba653dde6c..c98010d12e2fb 100644 --- a/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/structurefinder/NdJsonTextStructureFinderFactory.java +++ b/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/structurefinder/NdJsonTextStructureFinderFactory.java @@ -72,6 +72,16 @@ public boolean canCreateFromSample(List explanation, String sample, doub return true; } + public boolean canCreateFromMessages(List explanation, List messages, double allowedFractionOfBadLines) { + for (String message : messages) { + if (message.contains("\n")) { + explanation.add("Not NDJSON because message contains multiple lines: [" + message + "]"); + return false; + } + } + return canCreateFromSample(explanation, String.join("\n", messages), allowedFractionOfBadLines); + } + @Override public TextStructureFinder createFromSample( List explanation, @@ -92,6 +102,19 @@ public TextStructureFinder createFromSample( ); } + public TextStructureFinder createFromMessages( + List explanation, + List messages, + TextStructureOverrides overrides, + TimeoutChecker timeoutChecker + ) throws IOException { + // NdJsonTextStructureFinderFactory::canCreateFromMessages already + // checked that every line contains a single valid JSON message, + // so we can safely concatenate and run the logic for a sample. + String sample = String.join("\n", messages); + return NdJsonTextStructureFinder.makeNdJsonTextStructureFinder(explanation, sample, "UTF-8", null, overrides, timeoutChecker); + } + private static class ContextPrintingStringReader extends StringReader { private final String str; diff --git a/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/structurefinder/TextStructureFinderFactory.java b/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/structurefinder/TextStructureFinderFactory.java index 63970dd2c58d9..1e8317400d09d 100644 --- a/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/structurefinder/TextStructureFinderFactory.java +++ b/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/structurefinder/TextStructureFinderFactory.java @@ -33,6 +33,8 @@ public interface TextStructureFinderFactory { */ boolean canCreateFromSample(List explanation, String sample, double allowedFractionOfBadLines); + boolean canCreateFromMessages(List explanation, List messages, double allowedFractionOfBadMessages); + /** * Create an object representing the structure of some text. * @param explanation List of reasons for making decisions. May contain items when passed and new reasons @@ -56,4 +58,11 @@ TextStructureFinder createFromSample( TextStructureOverrides overrides, TimeoutChecker timeoutChecker ) throws Exception; + + TextStructureFinder createFromMessages( + List explanation, + List messages, + TextStructureOverrides overrides, + TimeoutChecker timeoutChecker + ) throws Exception; } diff --git a/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/structurefinder/TextStructureFinderManager.java b/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/structurefinder/TextStructureFinderManager.java index c0a100fbb280d..899f6c9108060 100644 --- a/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/structurefinder/TextStructureFinderManager.java +++ b/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/structurefinder/TextStructureFinderManager.java @@ -13,7 +13,7 @@ import org.elasticsearch.ElasticsearchTimeoutException; import org.elasticsearch.core.TimeValue; import org.elasticsearch.core.Tuple; -import org.elasticsearch.xpack.core.textstructure.action.FindStructureAction; +import org.elasticsearch.xpack.core.textstructure.action.AbstractFindStructureRequest; import org.elasticsearch.xpack.core.textstructure.structurefinder.TextStructure; import java.io.BufferedInputStream; @@ -310,7 +310,7 @@ public TextStructureFinder findTextStructure(Integer idealSampleLineCount, Integ * Given a stream of text data, determine its structure. * @param idealSampleLineCount Ideally, how many lines from the stream will be read to determine the structure? * If the stream has fewer lines then an attempt will still be made, providing at - * least {@link FindStructureAction#MIN_SAMPLE_LINE_COUNT} lines can be read. If + * least {@link AbstractFindStructureRequest#MIN_SAMPLE_LINE_COUNT} lines can be read. If * null the value of {@link #DEFAULT_IDEAL_SAMPLE_LINE_COUNT} will be used. * @param lineMergeSizeLimit Maximum number of characters permitted when lines are merged to create messages. * If null the value of {@link #DEFAULT_LINE_MERGE_SIZE_LIMIT} will be used. @@ -383,11 +383,11 @@ public TextStructureFinder findTextStructure( sampleReader = charsetMatch.getReader(); } - assert idealSampleLineCount >= FindStructureAction.MIN_SAMPLE_LINE_COUNT; + assert idealSampleLineCount >= AbstractFindStructureRequest.MIN_SAMPLE_LINE_COUNT; Tuple sampleInfo = sampleText( sampleReader, charsetName, - FindStructureAction.MIN_SAMPLE_LINE_COUNT, + AbstractFindStructureRequest.MIN_SAMPLE_LINE_COUNT, idealSampleLineCount, timeoutChecker ); @@ -413,6 +413,23 @@ public TextStructureFinder findTextStructure( } } + public TextStructureFinder findTextStructure(List messages, TextStructureOverrides overrides, TimeValue timeout) + throws Exception { + List explanation = new ArrayList<>(); + try (TimeoutChecker timeoutChecker = new TimeoutChecker("structure analysis", timeout, scheduler)) { + return makeBestStructureFinder(explanation, messages, overrides, timeoutChecker); + } catch (Exception e) { + // Add a dummy exception containing the explanation so far - this can be invaluable for troubleshooting as incorrect + // decisions made early on in the structure analysis can result in seemingly crazy decisions or timeouts later on + if (explanation.isEmpty() == false) { + e.addSuppressed( + new ElasticsearchException(explanation.stream().collect(Collectors.joining("]\n[", "Explanation so far:\n[", "]\n"))) + ); + } + throw e; + } + } + CharsetMatch findCharset(List explanation, InputStream inputStream, TimeoutChecker timeoutChecker) throws Exception { // We need an input stream that supports mark and reset, so wrap the argument @@ -551,24 +568,12 @@ CharsetMatch findCharset(List explanation, InputStream inputStream, Time ); } - TextStructureFinder makeBestStructureFinder( - List explanation, - String sample, - String charsetName, - Boolean hasByteOrderMarker, - int lineMergeSizeLimit, - TextStructureOverrides overrides, - TimeoutChecker timeoutChecker - ) throws Exception { - + List getFactories(TextStructureOverrides overrides) { Character delimiter = overrides.getDelimiter(); Character quote = overrides.getQuote(); Boolean shouldTrimFields = overrides.getShouldTrimFields(); List factories; - double allowedFractionOfBadLines = 0.0; if (delimiter != null) { - allowedFractionOfBadLines = DelimitedTextStructureFinderFactory.DELIMITER_OVERRIDDEN_ALLOWED_FRACTION_OF_BAD_LINES; - // If a precise delimiter is specified, we only need one structure finder // factory, and we'll tolerate as little as one column in the input factories = Collections.singletonList( @@ -581,8 +586,6 @@ TextStructureFinder makeBestStructureFinder( ); } else if (quote != null || shouldTrimFields != null || TextStructure.Format.DELIMITED.equals(overrides.getFormat())) { - allowedFractionOfBadLines = DelimitedTextStructureFinderFactory.FORMAT_OVERRIDDEN_ALLOWED_FRACTION_OF_BAD_LINES; - // The delimiter is not specified, but some other aspect of delimited text is, // so clone our default delimited factories altering the overridden values factories = ORDERED_STRUCTURE_FACTORIES.stream() @@ -599,6 +602,34 @@ TextStructureFinder makeBestStructureFinder( } + return factories; + } + + private double getAllowedFractionOfBadLines(TextStructureOverrides overrides) { + Character delimiter = overrides.getDelimiter(); + Character quote = overrides.getQuote(); + Boolean shouldTrimFields = overrides.getShouldTrimFields(); + if (delimiter != null) { + return DelimitedTextStructureFinderFactory.DELIMITER_OVERRIDDEN_ALLOWED_FRACTION_OF_BAD_LINES; + } else if (quote != null || shouldTrimFields != null || TextStructure.Format.DELIMITED.equals(overrides.getFormat())) { + return DelimitedTextStructureFinderFactory.FORMAT_OVERRIDDEN_ALLOWED_FRACTION_OF_BAD_LINES; + } else { + return 0.0; + } + } + + TextStructureFinder makeBestStructureFinder( + List explanation, + String sample, + String charsetName, + Boolean hasByteOrderMarker, + int lineMergeSizeLimit, + TextStructureOverrides overrides, + TimeoutChecker timeoutChecker + ) throws Exception { + List factories = getFactories(overrides); + double allowedFractionOfBadLines = getAllowedFractionOfBadLines(overrides); + for (TextStructureFinderFactory factory : factories) { timeoutChecker.check("high level format detection"); if (factory.canCreateFromSample(explanation, sample, allowedFractionOfBadLines)) { @@ -620,6 +651,28 @@ TextStructureFinder makeBestStructureFinder( ); } + private TextStructureFinder makeBestStructureFinder( + List explanation, + List messages, + TextStructureOverrides overrides, + TimeoutChecker timeoutChecker + ) throws Exception { + List factories = getFactories(overrides); + double allowedFractionOfBadLines = getAllowedFractionOfBadLines(overrides); + + for (TextStructureFinderFactory factory : factories) { + timeoutChecker.check("high level format detection"); + if (factory.canCreateFromMessages(explanation, messages, allowedFractionOfBadLines)) { + return factory.createFromMessages(explanation, messages, overrides, timeoutChecker); + } + } + + throw new IllegalArgumentException( + "Input did not match " + + ((overrides.getFormat() == null) ? "any known formats" : "the specified format [" + overrides.getFormat() + "]") + ); + } + private Tuple sampleText(Reader reader, String charsetName, int minLines, int maxLines, TimeoutChecker timeoutChecker) throws IOException { diff --git a/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/structurefinder/TextStructureOverrides.java b/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/structurefinder/TextStructureOverrides.java index 5ba4e464508f1..303cb2a59ea16 100644 --- a/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/structurefinder/TextStructureOverrides.java +++ b/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/structurefinder/TextStructureOverrides.java @@ -6,6 +6,7 @@ */ package org.elasticsearch.xpack.textstructure.structurefinder; +import org.elasticsearch.xpack.core.textstructure.action.AbstractFindStructureRequest; import org.elasticsearch.xpack.core.textstructure.action.FindStructureAction; import org.elasticsearch.xpack.core.textstructure.structurefinder.TextStructure; @@ -37,7 +38,7 @@ public class TextStructureOverrides { private final String ecsCompatibility; - public TextStructureOverrides(FindStructureAction.Request request) { + public TextStructureOverrides(AbstractFindStructureRequest request) { this( request.getCharset(), diff --git a/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/structurefinder/XmlTextStructureFinderFactory.java b/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/structurefinder/XmlTextStructureFinderFactory.java index 10f65564c3dde..2f56c73616866 100644 --- a/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/structurefinder/XmlTextStructureFinderFactory.java +++ b/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/structurefinder/XmlTextStructureFinderFactory.java @@ -46,7 +46,42 @@ public boolean canFindFormat(TextStructure.Format format) { */ @Override public boolean canCreateFromSample(List explanation, String sample, double allowedFractionOfBadLines) { + int completeDocCount = parseXml(explanation, sample); + if (completeDocCount == -1) { + return false; + } + if (completeDocCount == 0) { + explanation.add("Not XML because sample didn't contain a complete document"); + return false; + } + explanation.add("Deciding sample is XML"); + return true; + } + + public boolean canCreateFromMessages(List explanation, List messages, double allowedFractionOfBadLines) { + for (String message : messages) { + int completeDocCount = parseXml(explanation, message); + if (completeDocCount == -1) { + return false; + } + if (completeDocCount == 0) { + explanation.add("Not XML because a message didn't contain a complete document"); + return false; + } + if (completeDocCount > 1) { + explanation.add("Not XML because a message contains a multiple documents"); + return false; + } + } + explanation.add("Deciding sample is XML"); + return true; + } + /** + * Tries to parse the sample as XML. + * @return -1 if invalid, otherwise the number of complete docs + */ + private int parseXml(List explanation, String sample) { int completeDocCount = 0; String commonRootElementName = null; String remainder = sample.trim(); @@ -80,14 +115,14 @@ public boolean canCreateFromSample(List explanation, String sample, doub + rootElementName + "]" ); - return false; + return -1; } } break; case XMLStreamReader.END_ELEMENT: if (--nestingLevel < 0) { explanation.add("Not XML because an end element occurs before a start element"); - return false; + return -1; } break; } @@ -111,7 +146,7 @@ public boolean canCreateFromSample(List explanation, String sample, doub + remainder + "]" ); - return false; + return -1; } } endPos += location.getColumnNumber() - 1; @@ -125,17 +160,11 @@ public boolean canCreateFromSample(List explanation, String sample, doub } } catch (IOException | XMLStreamException e) { explanation.add("Not XML because there was a parsing exception: [" + e.getMessage().replaceAll("\\s?\r?\n\\s?", " ") + "]"); - return false; + return -1; } } - if (completeDocCount == 0) { - explanation.add("Not XML because sample didn't contain a complete document"); - return false; - } - - explanation.add("Deciding sample is XML"); - return true; + return completeDocCount; } @Override @@ -157,4 +186,17 @@ public TextStructureFinder createFromSample( timeoutChecker ); } + + public TextStructureFinder createFromMessages( + List explanation, + List messages, + TextStructureOverrides overrides, + TimeoutChecker timeoutChecker + ) throws IOException, ParserConfigurationException, SAXException { + // XmlTextStructureFinderFactory::canCreateFromMessages already + // checked that every message contains a single valid XML document, + // so we can safely concatenate and run the logic for a sample. + String sample = String.join("\n", messages); + return XmlTextStructureFinder.makeXmlTextStructureFinder(explanation, sample, "UTF-8", null, overrides, timeoutChecker); + } } diff --git a/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/transport/TransportFindFieldStructureAction.java b/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/transport/TransportFindFieldStructureAction.java new file mode 100644 index 0000000000000..43a990f6f565b --- /dev/null +++ b/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/transport/TransportFindFieldStructureAction.java @@ -0,0 +1,94 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ +package org.elasticsearch.xpack.textstructure.transport; + +import org.elasticsearch.action.ActionListener; +import org.elasticsearch.action.search.SearchResponse; +import org.elasticsearch.action.support.ActionFilters; +import org.elasticsearch.action.support.HandledTransportAction; +import org.elasticsearch.client.internal.Client; +import org.elasticsearch.client.internal.ParentTaskAssigningClient; +import org.elasticsearch.common.inject.Inject; +import org.elasticsearch.index.query.QueryBuilders; +import org.elasticsearch.tasks.Task; +import org.elasticsearch.tasks.TaskId; +import org.elasticsearch.threadpool.ThreadPool; +import org.elasticsearch.transport.TransportService; +import org.elasticsearch.xpack.core.ml.utils.MapHelper; +import org.elasticsearch.xpack.core.textstructure.action.AbstractFindStructureRequest; +import org.elasticsearch.xpack.core.textstructure.action.FindFieldStructureAction; +import org.elasticsearch.xpack.core.textstructure.action.FindStructureResponse; +import org.elasticsearch.xpack.textstructure.structurefinder.TextStructureFinder; +import org.elasticsearch.xpack.textstructure.structurefinder.TextStructureFinderManager; +import org.elasticsearch.xpack.textstructure.structurefinder.TextStructureOverrides; + +import java.util.Arrays; +import java.util.List; +import java.util.Objects; +import java.util.stream.Collectors; + +public class TransportFindFieldStructureAction extends HandledTransportAction { + + private final Client client; + private final TransportService transportService; + private final ThreadPool threadPool; + + @Inject + public TransportFindFieldStructureAction( + TransportService transportService, + ActionFilters actionFilters, + Client client, + ThreadPool threadPool + ) { + super(FindFieldStructureAction.NAME, transportService, actionFilters, FindFieldStructureAction.Request::new, threadPool.generic()); + this.client = client; + this.transportService = transportService; + this.threadPool = threadPool; + } + + @Override + protected void doExecute(Task task, FindFieldStructureAction.Request request, ActionListener listener) { + TaskId taskId = new TaskId(transportService.getLocalNode().getId(), task.getId()); + new ParentTaskAssigningClient(client, taskId).prepareSearch(request.getIndex()) + .setSize(request.getLinesToSample()) + .setFetchSource(true) + .setQuery(QueryBuilders.existsQuery(request.getField())) + .setFetchSource(new String[] { request.getField() }, null) + .execute(ActionListener.wrap(searchResponse -> { + long hitCount = searchResponse.getHits().getHits().length; + if (hitCount < AbstractFindStructureRequest.MIN_SAMPLE_LINE_COUNT) { + listener.onFailure( + new IllegalArgumentException("Input contained too few lines [" + hitCount + "] to obtain a meaningful sample") + ); + return; + } + List messages = getMessages(searchResponse, request.getField()); + try { + listener.onResponse(buildTextStructureResponse(messages, request)); + } catch (Exception e) { + listener.onFailure(e); + } + }, listener::onFailure)); + } + + private List getMessages(SearchResponse searchResponse, String field) { + return Arrays.stream(searchResponse.getHits().getHits()) + .map(hit -> MapHelper.dig(field, Objects.requireNonNull(hit.getSourceAsMap())).toString()) + .collect(Collectors.toList()); + } + + private FindStructureResponse buildTextStructureResponse(List messages, FindFieldStructureAction.Request request) + throws Exception { + TextStructureFinderManager structureFinderManager = new TextStructureFinderManager(threadPool.scheduler()); + TextStructureFinder textStructureFinder = structureFinderManager.findTextStructure( + messages, + new TextStructureOverrides(request), + request.getTimeout() + ); + return new FindStructureResponse(textStructureFinder.getStructure()); + } +} diff --git a/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/transport/TransportFindMessageStructureAction.java b/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/transport/TransportFindMessageStructureAction.java new file mode 100644 index 0000000000000..79c21b3cea306 --- /dev/null +++ b/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/transport/TransportFindMessageStructureAction.java @@ -0,0 +1,56 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ +package org.elasticsearch.xpack.textstructure.transport; + +import org.elasticsearch.action.ActionListener; +import org.elasticsearch.action.support.ActionFilters; +import org.elasticsearch.action.support.HandledTransportAction; +import org.elasticsearch.common.inject.Inject; +import org.elasticsearch.tasks.Task; +import org.elasticsearch.threadpool.ThreadPool; +import org.elasticsearch.transport.TransportService; +import org.elasticsearch.xpack.core.textstructure.action.FindMessageStructureAction; +import org.elasticsearch.xpack.core.textstructure.action.FindStructureResponse; +import org.elasticsearch.xpack.textstructure.structurefinder.TextStructureFinder; +import org.elasticsearch.xpack.textstructure.structurefinder.TextStructureFinderManager; +import org.elasticsearch.xpack.textstructure.structurefinder.TextStructureOverrides; + +public class TransportFindMessageStructureAction extends HandledTransportAction { + + private final ThreadPool threadPool; + + @Inject + public TransportFindMessageStructureAction(TransportService transportService, ActionFilters actionFilters, ThreadPool threadPool) { + super( + FindMessageStructureAction.NAME, + transportService, + actionFilters, + FindMessageStructureAction.Request::new, + threadPool.generic() + ); + this.threadPool = threadPool; + } + + @Override + protected void doExecute(Task task, FindMessageStructureAction.Request request, ActionListener listener) { + try { + listener.onResponse(buildTextStructureResponse(request)); + } catch (Exception e) { + listener.onFailure(e); + } + } + + private FindStructureResponse buildTextStructureResponse(FindMessageStructureAction.Request request) throws Exception { + TextStructureFinderManager structureFinderManager = new TextStructureFinderManager(threadPool.scheduler()); + TextStructureFinder textStructureFinder = structureFinderManager.findTextStructure( + request.getMessages(), + new TextStructureOverrides(request), + request.getTimeout() + ); + return new FindStructureResponse(textStructureFinder.getStructure()); + } +} diff --git a/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/transport/TransportFindStructureAction.java b/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/transport/TransportFindStructureAction.java index 8bf0f1cd4395f..4257a36bc150a 100644 --- a/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/transport/TransportFindStructureAction.java +++ b/x-pack/plugin/text-structure/src/main/java/org/elasticsearch/xpack/textstructure/transport/TransportFindStructureAction.java @@ -10,53 +10,38 @@ import org.elasticsearch.action.support.ActionFilters; import org.elasticsearch.action.support.HandledTransportAction; import org.elasticsearch.common.inject.Inject; -import org.elasticsearch.common.util.concurrent.EsExecutors; import org.elasticsearch.tasks.Task; import org.elasticsearch.threadpool.ThreadPool; import org.elasticsearch.transport.TransportService; import org.elasticsearch.xpack.core.textstructure.action.FindStructureAction; +import org.elasticsearch.xpack.core.textstructure.action.FindStructureResponse; import org.elasticsearch.xpack.textstructure.structurefinder.TextStructureFinder; import org.elasticsearch.xpack.textstructure.structurefinder.TextStructureFinderManager; import org.elasticsearch.xpack.textstructure.structurefinder.TextStructureOverrides; import java.io.InputStream; -import static org.elasticsearch.threadpool.ThreadPool.Names.GENERIC; - -public class TransportFindStructureAction extends HandledTransportAction { +public class TransportFindStructureAction extends HandledTransportAction { private final ThreadPool threadPool; @Inject public TransportFindStructureAction(TransportService transportService, ActionFilters actionFilters, ThreadPool threadPool) { - super( - FindStructureAction.NAME, - transportService, - actionFilters, - FindStructureAction.Request::new, - EsExecutors.DIRECT_EXECUTOR_SERVICE - ); + super(FindStructureAction.NAME, transportService, actionFilters, FindStructureAction.Request::new, threadPool.generic()); this.threadPool = threadPool; } @Override - protected void doExecute(Task task, FindStructureAction.Request request, ActionListener listener) { - - // As determining the text structure might take a while, we run - // in a different thread to avoid blocking the network thread. - threadPool.executor(GENERIC).execute(() -> { - try { - listener.onResponse(buildTextStructureResponse(request)); - } catch (Exception e) { - listener.onFailure(e); - } - }); + protected void doExecute(Task task, FindStructureAction.Request request, ActionListener listener) { + try { + listener.onResponse(buildTextStructureResponse(request)); + } catch (Exception e) { + listener.onFailure(e); + } } - private FindStructureAction.Response buildTextStructureResponse(FindStructureAction.Request request) throws Exception { - + private FindStructureResponse buildTextStructureResponse(FindStructureAction.Request request) throws Exception { TextStructureFinderManager structureFinderManager = new TextStructureFinderManager(threadPool.scheduler()); - try (InputStream sampleStream = request.getSample().streamInput()) { TextStructureFinder textStructureFinder = structureFinderManager.findTextStructure( request.getLinesToSample(), @@ -65,8 +50,7 @@ private FindStructureAction.Response buildTextStructureResponse(FindStructureAct new TextStructureOverrides(request), request.getTimeout() ); - - return new FindStructureAction.Response(textStructureFinder.getStructure()); + return new FindStructureResponse(textStructureFinder.getStructure()); } } } diff --git a/x-pack/plugin/text-structure/src/test/java/org/elasticsearch/xpack/textstructure/structurefinder/DelimitedTextStructureFinderFactoryTests.java b/x-pack/plugin/text-structure/src/test/java/org/elasticsearch/xpack/textstructure/structurefinder/DelimitedTextStructureFinderFactoryTests.java index cd8c451ee0547..e28de72202460 100644 --- a/x-pack/plugin/text-structure/src/test/java/org/elasticsearch/xpack/textstructure/structurefinder/DelimitedTextStructureFinderFactoryTests.java +++ b/x-pack/plugin/text-structure/src/test/java/org/elasticsearch/xpack/textstructure/structurefinder/DelimitedTextStructureFinderFactoryTests.java @@ -6,6 +6,9 @@ */ package org.elasticsearch.xpack.textstructure.structurefinder; +import java.util.Arrays; +import java.util.List; + public class DelimitedTextStructureFinderFactoryTests extends TextStructureTestCase { private final TextStructureFinderFactory csvFactory = new DelimitedTextStructureFinderFactory(',', '"', 2, false); @@ -40,6 +43,21 @@ public void testCanCreateCsvFromSampleGivenText() { assertFalse(csvFactory.canCreateFromSample(explanation, TEXT_SAMPLE, 0.0)); } + public void testCanCreateCsvFromMessagesCsv() { + List messages = Arrays.asList(CSV_SAMPLE.split("\n")); + assertTrue(csvFactory.canCreateFromMessages(explanation, messages, 0.0)); + } + + public void testCanCreateCsvFromMessagesCsv_multipleCsvRowsPerMessage() { + List messages = List.of(CSV_SAMPLE, CSV_SAMPLE, CSV_SAMPLE); + assertFalse(csvFactory.canCreateFromMessages(explanation, messages, 0.0)); + } + + public void testCanCreateCsvFromMessagesCsv_emptyMessages() { + List messages = List.of("", "", ""); + assertFalse(csvFactory.canCreateFromMessages(explanation, messages, 0.0)); + } + // TSV - no need to check NDJSON, XML or CSV because they come earlier in the order we check formats public void testCanCreateTsvFromSampleGivenTsv() { diff --git a/x-pack/plugin/text-structure/src/test/java/org/elasticsearch/xpack/textstructure/structurefinder/DelimitedTextStructureFinderTests.java b/x-pack/plugin/text-structure/src/test/java/org/elasticsearch/xpack/textstructure/structurefinder/DelimitedTextStructureFinderTests.java index 478994178c5bc..62e06af809711 100644 --- a/x-pack/plugin/text-structure/src/test/java/org/elasticsearch/xpack/textstructure/structurefinder/DelimitedTextStructureFinderTests.java +++ b/x-pack/plugin/text-structure/src/test/java/org/elasticsearch/xpack/textstructure/structurefinder/DelimitedTextStructureFinderTests.java @@ -790,6 +790,30 @@ public void testCreateConfigsGivenDotInFieldName() throws Exception { assertEquals(Collections.singleton("properties"), structure.getMappings().keySet()); } + public void testCreateFromMessages() throws Exception { + List messages = List.of("a,b,c", "d,e,f", "g,h,i"); + assertTrue(csvFactory.canCreateFromMessages(explanation, messages, 0.0)); + TextStructureFinder structureFinder = csvFactory.createFromMessages( + explanation, + messages, + TextStructureOverrides.EMPTY_OVERRIDES, + NOOP_TIMEOUT_CHECKER + ); + TextStructure structure = structureFinder.getStructure(); + assertEquals(TextStructure.Format.DELIMITED, structure.getFormat()); + assertEquals(3, structure.getNumMessagesAnalyzed()); + } + + public void testCreateFromMessages_multipleRowPerMessage() { + List messages = List.of("a,b,c\nd,e,f", "g,h,i"); + assertFalse(csvFactory.canCreateFromMessages(explanation, messages, 0.0)); + } + + public void testCreateFromMessages_emptyMessage() { + List messages = List.of("a,b,c", "", "d,e,f"); + assertFalse(csvFactory.canCreateFromMessages(explanation, messages, 0.0)); + } + public void testFindHeaderFromSampleGivenHeaderInSample() throws IOException { String withHeader = """ time,airline,responsetime,sourcetype diff --git a/x-pack/plugin/text-structure/src/test/java/org/elasticsearch/xpack/textstructure/structurefinder/LogTextStructureFinderTests.java b/x-pack/plugin/text-structure/src/test/java/org/elasticsearch/xpack/textstructure/structurefinder/LogTextStructureFinderTests.java index 4ee651408af56..484fde023be6b 100644 --- a/x-pack/plugin/text-structure/src/test/java/org/elasticsearch/xpack/textstructure/structurefinder/LogTextStructureFinderTests.java +++ b/x-pack/plugin/text-structure/src/test/java/org/elasticsearch/xpack/textstructure/structurefinder/LogTextStructureFinderTests.java @@ -11,6 +11,7 @@ import org.elasticsearch.xpack.core.textstructure.structurefinder.TextStructure; import java.util.Collections; +import java.util.List; import java.util.Map; import java.util.Set; import java.util.stream.Collectors; @@ -106,6 +107,21 @@ public void testCreateConfigsGivenElasticsearchLog() throws Exception { assertTrue(keys.contains("@timestamp")); } + public void testCreateFromMessages() throws Exception { + List messages = List.of(TEXT_SAMPLE.split("\n")); + assertTrue(factory.canCreateFromMessages(explanation, messages, 0.0)); + + TextStructureFinder structureFinder = factory.createFromMessages( + explanation, + messages, + TextStructureOverrides.EMPTY_OVERRIDES, + NOOP_TIMEOUT_CHECKER + ); + + TextStructure structure = structureFinder.getStructure(); + assertEquals("\\[%{TIMESTAMP_ISO8601:timestamp}\\]\\[%{LOGLEVEL:loglevel} \\]\\[.*", structure.getGrokPattern()); + } + public void testCreateConfigsGivenElasticsearchLogWithNoTimestamps() throws Exception { assertTrue(factory.canCreateFromSample(explanation, TEXT_WITH_NO_TIMESTAMPS_SAMPLE, 0.0)); diff --git a/x-pack/plugin/text-structure/src/test/java/org/elasticsearch/xpack/textstructure/structurefinder/NdJsonTextStructureFinderFactoryTests.java b/x-pack/plugin/text-structure/src/test/java/org/elasticsearch/xpack/textstructure/structurefinder/NdJsonTextStructureFinderFactoryTests.java index 85baf238630bb..dac202df8e811 100644 --- a/x-pack/plugin/text-structure/src/test/java/org/elasticsearch/xpack/textstructure/structurefinder/NdJsonTextStructureFinderFactoryTests.java +++ b/x-pack/plugin/text-structure/src/test/java/org/elasticsearch/xpack/textstructure/structurefinder/NdJsonTextStructureFinderFactoryTests.java @@ -6,6 +6,9 @@ */ package org.elasticsearch.xpack.textstructure.structurefinder; +import java.util.Arrays; +import java.util.List; + public class NdJsonTextStructureFinderFactoryTests extends TextStructureTestCase { private final TextStructureFinderFactory factory = new NdJsonTextStructureFinderFactory(); @@ -15,6 +18,21 @@ public void testCanCreateFromSampleGivenNdJson() { assertTrue(factory.canCreateFromSample(explanation, NDJSON_SAMPLE, 0.0)); } + public void testCanCreateFromMessages() { + List messages = Arrays.asList(NDJSON_SAMPLE.split("\n")); + assertTrue(factory.canCreateFromMessages(explanation, messages, 0.0)); + } + + public void testCanCreateFromMessages_multipleJsonLinesPerMessage() { + List messages = List.of(NDJSON_SAMPLE, NDJSON_SAMPLE, NDJSON_SAMPLE); + assertFalse(factory.canCreateFromMessages(explanation, messages, 0.0)); + } + + public void testCanCreateFromMessages_emptyMessages() { + List messages = List.of("", "", ""); + assertFalse(factory.canCreateFromMessages(explanation, messages, 0.0)); + } + public void testCanCreateFromSampleGivenXml() { assertFalse(factory.canCreateFromSample(explanation, XML_SAMPLE, 0.0)); diff --git a/x-pack/plugin/text-structure/src/test/java/org/elasticsearch/xpack/textstructure/structurefinder/XmlTextStructureFinderFactoryTests.java b/x-pack/plugin/text-structure/src/test/java/org/elasticsearch/xpack/textstructure/structurefinder/XmlTextStructureFinderFactoryTests.java index ea92420a1ea5a..7340c0c3dff00 100644 --- a/x-pack/plugin/text-structure/src/test/java/org/elasticsearch/xpack/textstructure/structurefinder/XmlTextStructureFinderFactoryTests.java +++ b/x-pack/plugin/text-structure/src/test/java/org/elasticsearch/xpack/textstructure/structurefinder/XmlTextStructureFinderFactoryTests.java @@ -6,6 +6,9 @@ */ package org.elasticsearch.xpack.textstructure.structurefinder; +import java.util.Arrays; +import java.util.List; + public class XmlTextStructureFinderFactoryTests extends TextStructureTestCase { private final TextStructureFinderFactory factory = new XmlTextStructureFinderFactory(); @@ -17,6 +20,21 @@ public void testCanCreateFromSampleGivenXml() { assertTrue(factory.canCreateFromSample(explanation, XML_SAMPLE, 0.0)); } + public void testCanCreateFromMessages() { + List messages = Arrays.asList(XML_SAMPLE.split("\n\n")); + assertTrue(factory.canCreateFromMessages(explanation, messages, 0.0)); + } + + public void testCanCreateFromMessages_multipleXmlDocsPerMessage() { + List messages = List.of(XML_SAMPLE, XML_SAMPLE, XML_SAMPLE); + assertFalse(factory.canCreateFromMessages(explanation, messages, 0.0)); + } + + public void testCanCreateFromMessages_emptyMessages() { + List messages = List.of("", "", ""); + assertFalse(factory.canCreateFromMessages(explanation, messages, 0.0)); + } + public void testCanCreateFromSampleGivenCsv() { assertFalse(factory.canCreateFromSample(explanation, CSV_SAMPLE, 0.0)); From fa735a9b772842d643b7c8eb4248c550db8262c3 Mon Sep 17 00:00:00 2001 From: David Roberts Date: Wed, 6 Mar 2024 10:08:28 +0000 Subject: [PATCH 16/27] [ML] Fix `categorize_text` aggregation nested under empty buckets (#105987) Previously the `categorize_text` aggregation could throw an exception if nested as a sub-aggregation of another aggregation that produced empty buckets at the end of its results. This change avoids this possibility. Fixes #105836 --- docs/changelog/105987.yaml | 6 ++ .../CategorizeTextAggregator.java | 3 +- .../CategorizeTextAggregatorTests.java | 84 +++++++++++++++++++ 3 files changed, 92 insertions(+), 1 deletion(-) create mode 100644 docs/changelog/105987.yaml diff --git a/docs/changelog/105987.yaml b/docs/changelog/105987.yaml new file mode 100644 index 0000000000000..d09a6907c72bf --- /dev/null +++ b/docs/changelog/105987.yaml @@ -0,0 +1,6 @@ +pr: 105987 +summary: Fix `categorize_text` aggregation nested under empty buckets +area: Machine Learning +type: bug +issues: + - 105836 diff --git a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/aggs/categorization/CategorizeTextAggregator.java b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/aggs/categorization/CategorizeTextAggregator.java index 520d554379cfc..cedaced0f57ee 100644 --- a/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/aggs/categorization/CategorizeTextAggregator.java +++ b/x-pack/plugin/ml/src/main/java/org/elasticsearch/xpack/ml/aggs/categorization/CategorizeTextAggregator.java @@ -113,7 +113,8 @@ protected void doClose() { public InternalAggregation[] buildAggregations(long[] ordsToCollect) throws IOException { Bucket[][] topBucketsPerOrd = new Bucket[ordsToCollect.length][]; for (int ordIdx = 0; ordIdx < ordsToCollect.length; ordIdx++) { - final TokenListCategorizer categorizer = categorizers.get(ordsToCollect[ordIdx]); + final long ord = ordsToCollect[ordIdx]; + final TokenListCategorizer categorizer = (ord < categorizers.size()) ? categorizers.get(ord) : null; if (categorizer == null) { topBucketsPerOrd[ordIdx] = new Bucket[0]; continue; diff --git a/x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/aggs/categorization/CategorizeTextAggregatorTests.java b/x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/aggs/categorization/CategorizeTextAggregatorTests.java index cb5b98af29d57..29f298894477a 100644 --- a/x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/aggs/categorization/CategorizeTextAggregatorTests.java +++ b/x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/aggs/categorization/CategorizeTextAggregatorTests.java @@ -322,6 +322,90 @@ public void testCategorizationWithSubAggsManyDocs() throws Exception { ); } + public void testCategorizationAsSubAggWithExtendedBounds() throws Exception { + // Test with more buckets than we have data for (via extended bounds in the histogram config). + // This will confirm that we don't try to read beyond the end of arrays of categorizers. + int numHistoBuckets = 50; + HistogramAggregationBuilder aggBuilder = new HistogramAggregationBuilder("histo").field(NUMERIC_FIELD_NAME) + .interval(1) + .extendedBounds(0, numHistoBuckets - 1) + .subAggregation(new CategorizeTextAggregationBuilder("my_agg", TEXT_FIELD_NAME)); + testCase(CategorizeTextAggregatorTests::writeTestDocs, (InternalHistogram histo) -> { + assertThat(histo.getBuckets(), hasSize(numHistoBuckets)); + assertThat(histo.getBuckets().get(0).getDocCount(), equalTo(2L)); + assertThat(histo.getBuckets().get(0).getKeyAsString(), equalTo("0.0")); + InternalCategorizationAggregation categorizationAggregation = histo.getBuckets().get(0).getAggregations().get("my_agg"); + assertThat(categorizationAggregation.getBuckets().get(0).getDocCount(), equalTo(1L)); + assertThat( + categorizationAggregation.getBuckets().get(0).getKeyAsString(), + equalTo("Failed to shutdown error org.aaaa.bbbb.Cccc line caused by foo exception") + ); + assertThat(categorizationAggregation.getBuckets().get(0).getSerializableCategory().maxMatchingStringLen(), equalTo(84)); + assertThat( + categorizationAggregation.getBuckets().get(0).getSerializableCategory().getRegex(), + equalTo(".*?Failed.+?to.+?shutdown.+?error.+?org\\.aaaa\\.bbbb\\.Cccc.+?line.+?caused.+?by.+?foo.+?exception.*?") + ); + assertThat(categorizationAggregation.getBuckets().get(1).getDocCount(), equalTo(1L)); + assertThat(categorizationAggregation.getBuckets().get(1).getKeyAsString(), equalTo("Node started")); + assertThat(categorizationAggregation.getBuckets().get(1).getSerializableCategory().maxMatchingStringLen(), equalTo(15)); + assertThat(categorizationAggregation.getBuckets().get(1).getSerializableCategory().getRegex(), equalTo(".*?Node.+?started.*?")); + assertThat(histo.getBuckets().get(1).getDocCount(), equalTo(1L)); + assertThat(histo.getBuckets().get(1).getKeyAsString(), equalTo("1.0")); + categorizationAggregation = histo.getBuckets().get(1).getAggregations().get("my_agg"); + assertThat(categorizationAggregation.getBuckets().get(0).getDocCount(), equalTo(1L)); + assertThat(categorizationAggregation.getBuckets().get(0).getKeyAsString(), equalTo("Node started")); + assertThat(categorizationAggregation.getBuckets().get(0).getSerializableCategory().maxMatchingStringLen(), equalTo(15)); + assertThat(categorizationAggregation.getBuckets().get(0).getSerializableCategory().getRegex(), equalTo(".*?Node.+?started.*?")); + assertThat(histo.getBuckets().get(2).getDocCount(), equalTo(1L)); + assertThat(histo.getBuckets().get(2).getKeyAsString(), equalTo("2.0")); + categorizationAggregation = histo.getBuckets().get(2).getAggregations().get("my_agg"); + assertThat(categorizationAggregation.getBuckets().get(0).getDocCount(), equalTo(1L)); + assertThat(categorizationAggregation.getBuckets().get(0).getKeyAsString(), equalTo("Node started")); + assertThat(categorizationAggregation.getBuckets().get(0).getSerializableCategory().maxMatchingStringLen(), equalTo(15)); + assertThat(categorizationAggregation.getBuckets().get(0).getSerializableCategory().getRegex(), equalTo(".*?Node.+?started.*?")); + assertThat(histo.getBuckets().get(3).getDocCount(), equalTo(1L)); + assertThat(histo.getBuckets().get(3).getKeyAsString(), equalTo("3.0")); + categorizationAggregation = histo.getBuckets().get(3).getAggregations().get("my_agg"); + assertThat(categorizationAggregation.getBuckets().get(0).getDocCount(), equalTo(1L)); + assertThat(categorizationAggregation.getBuckets().get(0).getKeyAsString(), equalTo("Node started")); + assertThat(categorizationAggregation.getBuckets().get(0).getSerializableCategory().maxMatchingStringLen(), equalTo(15)); + assertThat(categorizationAggregation.getBuckets().get(0).getSerializableCategory().getRegex(), equalTo(".*?Node.+?started.*?")); + assertThat(histo.getBuckets().get(4).getDocCount(), equalTo(2L)); + assertThat(histo.getBuckets().get(4).getKeyAsString(), equalTo("4.0")); + categorizationAggregation = histo.getBuckets().get(4).getAggregations().get("my_agg"); + assertThat(categorizationAggregation.getBuckets().get(0).getDocCount(), equalTo(1L)); + assertThat( + categorizationAggregation.getBuckets().get(0).getKeyAsString(), + equalTo("Failed to shutdown error org.aaaa.bbbb.Cccc line caused by foo exception") + ); + assertThat(categorizationAggregation.getBuckets().get(0).getSerializableCategory().maxMatchingStringLen(), equalTo(84)); + assertThat( + categorizationAggregation.getBuckets().get(0).getSerializableCategory().getRegex(), + equalTo(".*?Failed.+?to.+?shutdown.+?error.+?org\\.aaaa\\.bbbb\\.Cccc.+?line.+?caused.+?by.+?foo.+?exception.*?") + ); + assertThat(categorizationAggregation.getBuckets().get(1).getDocCount(), equalTo(1L)); + assertThat(categorizationAggregation.getBuckets().get(1).getKeyAsString(), equalTo("Node started")); + assertThat(categorizationAggregation.getBuckets().get(1).getSerializableCategory().maxMatchingStringLen(), equalTo(15)); + assertThat(categorizationAggregation.getBuckets().get(1).getSerializableCategory().getRegex(), equalTo(".*?Node.+?started.*?")); + assertThat(histo.getBuckets().get(5).getDocCount(), equalTo(1L)); + assertThat(histo.getBuckets().get(5).getKeyAsString(), equalTo("5.0")); + categorizationAggregation = histo.getBuckets().get(5).getAggregations().get("my_agg"); + assertThat(categorizationAggregation.getBuckets().get(0).getDocCount(), equalTo(1L)); + assertThat(categorizationAggregation.getBuckets().get(0).getKeyAsString(), equalTo("Node started")); + assertThat(categorizationAggregation.getBuckets().get(0).getSerializableCategory().maxMatchingStringLen(), equalTo(15)); + assertThat(categorizationAggregation.getBuckets().get(0).getSerializableCategory().getRegex(), equalTo(".*?Node.+?started.*?")); + for (int bucket = 6; bucket < numHistoBuckets; ++bucket) { + assertThat(histo.getBuckets().get(bucket).getDocCount(), equalTo(0L)); + } + }, + new AggTestConfig( + aggBuilder, + new TextFieldMapper.TextFieldType(TEXT_FIELD_NAME, randomBoolean()), + longField(NUMERIC_FIELD_NAME) + ) + ); + } + private static void writeTestDocs(RandomIndexWriter w) throws IOException { w.addDocument( Arrays.asList( From 882b92ab6092a8ea6b53a64e90cb12e812241483 Mon Sep 17 00:00:00 2001 From: Andrei Dan Date: Wed, 6 Mar 2024 10:12:08 +0000 Subject: [PATCH 17/27] Add service for computing the optimal number of shards for data streams (#105498) This adds the `DataStreamAutoShardingService` that will compute the optimal number of shards for a data stream and return a recommendation as to when to apply it (a time interval we call cool down which is 0 when the auto sharding recommendation can be applied immediately). This also introduces a `DataStreamAutoShardingEvent` object that will be stored in the data stream metadata to indicate the last auto sharding event that was applied to a data stream and its cluster state representation looks like so: ``` "auto_sharding": { "trigger_index_name": ".ds-logs-nginx-2024.02.12-000002", "target_number_of_shards": 3, "event_timestamp": 1707739707954 } ``` The auto sharding service is not used in this PR, so the auto sharding event will not be stored in the data stream metadata, but the required infrastructure to configure it is in place. --- .../datastreams/DataStreamIT.java | 3 +- .../DataStreamIndexSettingsProviderTests.java | 3 +- .../UpdateTimeSeriesRangeServiceTests.java | 3 +- .../action/GetDataStreamsResponseTests.java | 6 +- .../DataStreamLifecycleServiceTests.java | 3 +- server/src/main/java/module-info.java | 1 + .../org/elasticsearch/TransportVersions.java | 1 + .../datastreams/GetDataStreamAction.java | 17 + .../autosharding/AutoShardingResult.java | 59 ++ .../autosharding/AutoShardingType.java | 21 + .../DataStreamAutoShardingService.java | 415 ++++++++++ .../cluster/metadata/DataStream.java | 106 ++- .../metadata/DataStreamAutoShardingEvent.java | 84 ++ .../MetadataCreateDataStreamService.java | 3 +- .../metadata/MetadataDataStreamsService.java | 6 +- .../snapshots/RestoreService.java | 3 +- .../DataStreamAutoShardingServiceTests.java | 771 ++++++++++++++++++ .../DataStreamAutoShardingEventTests.java | 62 ++ .../cluster/metadata/DataStreamTests.java | 124 ++- .../MetadataDataStreamsServiceTests.java | 3 +- .../metadata/DataStreamTestHelper.java | 11 +- .../ccr/action/TransportPutFollowAction.java | 6 +- ...StreamLifecycleUsageTransportActionIT.java | 3 +- .../LicensedWriteLoadForecaster.java | 27 +- .../LicensedWriteLoadForecasterTests.java | 59 -- 25 files changed, 1681 insertions(+), 119 deletions(-) create mode 100644 server/src/main/java/org/elasticsearch/action/datastreams/autosharding/AutoShardingResult.java create mode 100644 server/src/main/java/org/elasticsearch/action/datastreams/autosharding/AutoShardingType.java create mode 100644 server/src/main/java/org/elasticsearch/action/datastreams/autosharding/DataStreamAutoShardingService.java create mode 100644 server/src/main/java/org/elasticsearch/cluster/metadata/DataStreamAutoShardingEvent.java create mode 100644 server/src/test/java/org/elasticsearch/action/datastreams/autosharding/DataStreamAutoShardingServiceTests.java create mode 100644 server/src/test/java/org/elasticsearch/cluster/metadata/DataStreamAutoShardingEventTests.java diff --git a/modules/data-streams/src/internalClusterTest/java/org/elasticsearch/datastreams/DataStreamIT.java b/modules/data-streams/src/internalClusterTest/java/org/elasticsearch/datastreams/DataStreamIT.java index 34f1701a595de..6c06511ccfbd1 100644 --- a/modules/data-streams/src/internalClusterTest/java/org/elasticsearch/datastreams/DataStreamIT.java +++ b/modules/data-streams/src/internalClusterTest/java/org/elasticsearch/datastreams/DataStreamIT.java @@ -1790,7 +1790,8 @@ public ClusterState execute(ClusterState currentState) throws Exception { original.getIndexMode(), original.getLifecycle(), original.isFailureStore(), - original.getFailureIndices() + original.getFailureIndices(), + null ); brokenDataStreamHolder.set(broken); return ClusterState.builder(currentState) diff --git a/modules/data-streams/src/test/java/org/elasticsearch/datastreams/DataStreamIndexSettingsProviderTests.java b/modules/data-streams/src/test/java/org/elasticsearch/datastreams/DataStreamIndexSettingsProviderTests.java index c65854903f7a9..01ad1bb09b20f 100644 --- a/modules/data-streams/src/test/java/org/elasticsearch/datastreams/DataStreamIndexSettingsProviderTests.java +++ b/modules/data-streams/src/test/java/org/elasticsearch/datastreams/DataStreamIndexSettingsProviderTests.java @@ -314,7 +314,8 @@ public void testGetAdditionalIndexSettingsDataStreamAlreadyCreatedTimeSettingsMi IndexMode.TIME_SERIES, ds.getLifecycle(), ds.isFailureStore(), - ds.getFailureIndices() + ds.getFailureIndices(), + null ) ); Metadata metadata = mb.build(); diff --git a/modules/data-streams/src/test/java/org/elasticsearch/datastreams/UpdateTimeSeriesRangeServiceTests.java b/modules/data-streams/src/test/java/org/elasticsearch/datastreams/UpdateTimeSeriesRangeServiceTests.java index dbb48ea3ddc26..abd5132edde16 100644 --- a/modules/data-streams/src/test/java/org/elasticsearch/datastreams/UpdateTimeSeriesRangeServiceTests.java +++ b/modules/data-streams/src/test/java/org/elasticsearch/datastreams/UpdateTimeSeriesRangeServiceTests.java @@ -153,7 +153,8 @@ public void testUpdateTimeSeriesTemporalRange_NoUpdateBecauseReplicated() { d.getIndexMode(), d.getLifecycle(), d.isFailureStore(), - d.getFailureIndices() + d.getFailureIndices(), + null ) ) .build(); diff --git a/modules/data-streams/src/test/java/org/elasticsearch/datastreams/action/GetDataStreamsResponseTests.java b/modules/data-streams/src/test/java/org/elasticsearch/datastreams/action/GetDataStreamsResponseTests.java index 13054379dd666..e200ff7cba2e1 100644 --- a/modules/data-streams/src/test/java/org/elasticsearch/datastreams/action/GetDataStreamsResponseTests.java +++ b/modules/data-streams/src/test/java/org/elasticsearch/datastreams/action/GetDataStreamsResponseTests.java @@ -89,7 +89,8 @@ public void testResponseIlmAndDataStreamLifecycleRepresentation() throws Excepti IndexMode.STANDARD, new DataStreamLifecycle(), true, - failureStores + failureStores, + null ); String ilmPolicyName = "rollover-30days"; @@ -198,7 +199,8 @@ public void testResponseIlmAndDataStreamLifecycleRepresentation() throws Excepti IndexMode.STANDARD, new DataStreamLifecycle(null, null, false), true, - failureStores + failureStores, + null ); String ilmPolicyName = "rollover-30days"; diff --git a/modules/data-streams/src/test/java/org/elasticsearch/datastreams/lifecycle/DataStreamLifecycleServiceTests.java b/modules/data-streams/src/test/java/org/elasticsearch/datastreams/lifecycle/DataStreamLifecycleServiceTests.java index 15f526d0a06d6..d0456d669663d 100644 --- a/modules/data-streams/src/test/java/org/elasticsearch/datastreams/lifecycle/DataStreamLifecycleServiceTests.java +++ b/modules/data-streams/src/test/java/org/elasticsearch/datastreams/lifecycle/DataStreamLifecycleServiceTests.java @@ -295,7 +295,8 @@ public void testRetentionNotExecutedForTSIndicesWithinTimeBounds() { dataStream.getIndexMode(), DataStreamLifecycle.newBuilder().dataRetention(0L).build(), dataStream.isFailureStore(), - dataStream.getFailureIndices() + dataStream.getFailureIndices(), + null ) ); clusterState = ClusterState.builder(clusterState).metadata(builder).build(); diff --git a/server/src/main/java/module-info.java b/server/src/main/java/module-info.java index 36a940af63c61..9c142d18034c0 100644 --- a/server/src/main/java/module-info.java +++ b/server/src/main/java/module-info.java @@ -381,6 +381,7 @@ opens org.elasticsearch.common.logging to org.apache.logging.log4j.core; exports org.elasticsearch.action.datastreams.lifecycle; + exports org.elasticsearch.action.datastreams.autosharding; exports org.elasticsearch.action.downsample; exports org.elasticsearch.plugins.internal to diff --git a/server/src/main/java/org/elasticsearch/TransportVersions.java b/server/src/main/java/org/elasticsearch/TransportVersions.java index 055fcb6d9cf7b..d484da5ba506c 100644 --- a/server/src/main/java/org/elasticsearch/TransportVersions.java +++ b/server/src/main/java/org/elasticsearch/TransportVersions.java @@ -135,6 +135,7 @@ static TransportVersion def(int id) { public static final TransportVersion ML_MODEL_IN_SERVICE_SETTINGS = def(8_595_00_0); public static final TransportVersion RANDOM_AGG_SHARD_SEED = def(8_596_00_0); public static final TransportVersion ESQL_TIMINGS = def(8_597_00_0); + public static final TransportVersion DATA_STREAM_AUTO_SHARDING_EVENT = def(8_598_00_0); /* * STOP! READ THIS FIRST! No, really, diff --git a/server/src/main/java/org/elasticsearch/action/datastreams/GetDataStreamAction.java b/server/src/main/java/org/elasticsearch/action/datastreams/GetDataStreamAction.java index 079c8f8b01ceb..8c469f7dffc4d 100644 --- a/server/src/main/java/org/elasticsearch/action/datastreams/GetDataStreamAction.java +++ b/server/src/main/java/org/elasticsearch/action/datastreams/GetDataStreamAction.java @@ -18,6 +18,7 @@ import org.elasticsearch.cluster.SimpleDiffable; import org.elasticsearch.cluster.health.ClusterHealthStatus; import org.elasticsearch.cluster.metadata.DataStream; +import org.elasticsearch.cluster.metadata.DataStreamAutoShardingEvent; import org.elasticsearch.common.io.stream.StreamInput; import org.elasticsearch.common.io.stream.StreamOutput; import org.elasticsearch.common.io.stream.Writeable; @@ -37,6 +38,7 @@ import java.util.Objects; import static org.elasticsearch.TransportVersions.V_8_11_X; +import static org.elasticsearch.cluster.metadata.DataStream.AUTO_SHARDING_FIELD; public class GetDataStreamAction extends ActionType { @@ -179,6 +181,10 @@ public static class DataStreamInfo implements SimpleDiffable, To public static final ParseField TEMPORAL_RANGES = new ParseField("temporal_ranges"); public static final ParseField TEMPORAL_RANGE_START = new ParseField("start"); public static final ParseField TEMPORAL_RANGE_END = new ParseField("end"); + public static final ParseField TIME_SINCE_LAST_AUTO_SHARD_EVENT = new ParseField("time_since_last_auto_shard_event"); + public static final ParseField TIME_SINCE_LAST_AUTO_SHARD_EVENT_MILLIS = new ParseField( + "time_since_last_auto_shard_event_millis" + ); private final DataStream dataStream; private final ClusterHealthStatus dataStreamStatus; @@ -348,6 +354,17 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params, @Nulla if (DataStream.isFailureStoreEnabled()) { builder.field(DataStream.FAILURE_STORE_FIELD.getPreferredName(), dataStream.isFailureStore()); } + if (dataStream.getAutoShardingEvent() != null) { + DataStreamAutoShardingEvent autoShardingEvent = dataStream.getAutoShardingEvent(); + builder.startObject(AUTO_SHARDING_FIELD.getPreferredName()); + autoShardingEvent.toXContent(builder, params); + builder.humanReadableField( + TIME_SINCE_LAST_AUTO_SHARD_EVENT_MILLIS.getPreferredName(), + TIME_SINCE_LAST_AUTO_SHARD_EVENT.getPreferredName(), + autoShardingEvent.getTimeSinceLastAutoShardingEvent(System::currentTimeMillis) + ); + builder.endObject(); + } if (timeSeries != null) { builder.startObject(TIME_SERIES.getPreferredName()); builder.startArray(TEMPORAL_RANGES.getPreferredName()); diff --git a/server/src/main/java/org/elasticsearch/action/datastreams/autosharding/AutoShardingResult.java b/server/src/main/java/org/elasticsearch/action/datastreams/autosharding/AutoShardingResult.java new file mode 100644 index 0000000000000..7bbd3291caf3a --- /dev/null +++ b/server/src/main/java/org/elasticsearch/action/datastreams/autosharding/AutoShardingResult.java @@ -0,0 +1,59 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0 and the Server Side Public License, v 1; you may not use this file except + * in compliance with, at your election, the Elastic License 2.0 or the Server + * Side Public License, v 1. + */ + +package org.elasticsearch.action.datastreams.autosharding; + +import org.elasticsearch.core.Nullable; +import org.elasticsearch.core.TimeValue; + +import java.util.Arrays; + +import static org.elasticsearch.action.datastreams.autosharding.AutoShardingType.COOLDOWN_PREVENTED_DECREASE; +import static org.elasticsearch.action.datastreams.autosharding.AutoShardingType.COOLDOWN_PREVENTED_INCREASE; + +/** + * Represents an auto sharding recommendation. It includes the current and target number of shards together with a remaining cooldown + * period that needs to lapse before the current recommendation should be applied. + *

+ * If auto sharding is not applicable for a data stream (e.g. due to + * {@link DataStreamAutoShardingService#DATA_STREAMS_AUTO_SHARDING_EXCLUDES_SETTING}) the target number of shards will be -1 and cool down + * remaining {@link TimeValue#MAX_VALUE}. + */ +public record AutoShardingResult( + AutoShardingType type, + int currentNumberOfShards, + int targetNumberOfShards, + TimeValue coolDownRemaining, + @Nullable Double writeLoad +) { + + static final String COOLDOWN_PREVENTING_TYPES = Arrays.toString( + new AutoShardingType[] { COOLDOWN_PREVENTED_DECREASE, COOLDOWN_PREVENTED_INCREASE } + ); + + public AutoShardingResult { + if (type.equals(AutoShardingType.INCREASE_SHARDS) || type.equals(AutoShardingType.DECREASE_SHARDS)) { + if (coolDownRemaining.equals(TimeValue.ZERO) == false) { + throw new IllegalArgumentException( + "The increase/decrease shards events must have a cooldown period of zero. Use one of [" + + COOLDOWN_PREVENTING_TYPES + + "] types indead" + ); + } + } + } + + public static final AutoShardingResult NOT_APPLICABLE_RESULT = new AutoShardingResult( + AutoShardingType.NOT_APPLICABLE, + -1, + -1, + TimeValue.MAX_VALUE, + null + ); + +} diff --git a/server/src/main/java/org/elasticsearch/action/datastreams/autosharding/AutoShardingType.java b/server/src/main/java/org/elasticsearch/action/datastreams/autosharding/AutoShardingType.java new file mode 100644 index 0000000000000..50d3027abbc88 --- /dev/null +++ b/server/src/main/java/org/elasticsearch/action/datastreams/autosharding/AutoShardingType.java @@ -0,0 +1,21 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0 and the Server Side Public License, v 1; you may not use this file except + * in compliance with, at your election, the Elastic License 2.0 or the Server + * Side Public License, v 1. + */ + +package org.elasticsearch.action.datastreams.autosharding; + +/** + * Represents the type of recommendation the auto sharding service provided. + */ +public enum AutoShardingType { + INCREASE_SHARDS, + DECREASE_SHARDS, + COOLDOWN_PREVENTED_INCREASE, + COOLDOWN_PREVENTED_DECREASE, + NO_CHANGE_REQUIRED, + NOT_APPLICABLE +} diff --git a/server/src/main/java/org/elasticsearch/action/datastreams/autosharding/DataStreamAutoShardingService.java b/server/src/main/java/org/elasticsearch/action/datastreams/autosharding/DataStreamAutoShardingService.java new file mode 100644 index 0000000000000..e830f538d222f --- /dev/null +++ b/server/src/main/java/org/elasticsearch/action/datastreams/autosharding/DataStreamAutoShardingService.java @@ -0,0 +1,415 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0 and the Server Side Public License, v 1; you may not use this file except + * in compliance with, at your election, the Elastic License 2.0 or the Server + * Side Public License, v 1. + */ + +package org.elasticsearch.action.datastreams.autosharding; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.elasticsearch.cluster.ClusterState; +import org.elasticsearch.cluster.metadata.DataStream; +import org.elasticsearch.cluster.metadata.IndexMetadata; +import org.elasticsearch.cluster.metadata.IndexMetadataStats; +import org.elasticsearch.cluster.metadata.IndexWriteLoad; +import org.elasticsearch.cluster.metadata.Metadata; +import org.elasticsearch.cluster.service.ClusterService; +import org.elasticsearch.common.regex.Regex; +import org.elasticsearch.common.settings.Setting; +import org.elasticsearch.common.settings.Settings; +import org.elasticsearch.core.Nullable; +import org.elasticsearch.core.TimeValue; +import org.elasticsearch.features.FeatureService; +import org.elasticsearch.features.NodeFeature; +import org.elasticsearch.index.Index; + +import java.util.List; +import java.util.Objects; +import java.util.OptionalDouble; +import java.util.function.Function; +import java.util.function.LongSupplier; + +import static org.elasticsearch.action.datastreams.autosharding.AutoShardingResult.NOT_APPLICABLE_RESULT; + +/** + * Calculates the optimal number of shards the data stream write index should have based on the indexing load. + */ +public class DataStreamAutoShardingService { + + private static final Logger logger = LogManager.getLogger(DataStreamAutoShardingService.class); + public static final String DATA_STREAMS_AUTO_SHARDING_ENABLED = "data_streams.auto_sharding.enabled"; + + public static final NodeFeature DATA_STREAM_AUTO_SHARDING_FEATURE = new NodeFeature("data_stream.auto_sharding"); + + public static final Setting> DATA_STREAMS_AUTO_SHARDING_EXCLUDES_SETTING = Setting.listSetting( + "data_streams.auto_sharding.excludes", + List.of("*"), + Function.identity(), + Setting.Property.Dynamic, + Setting.Property.NodeScope + ); + + /** + * Represents the minimum amount of time between two scaling events if the next event will increase the number of shards. + * We've chosen a value of 4.5minutes by default, just lower than the data stream lifecycle poll interval so we can increase shards with + * every DSL run, but we don't want it to be lower/0 as data stream lifecycle might run more often than the poll interval in case of + * a master failover. + */ + public static final Setting DATA_STREAMS_AUTO_SHARDING_INCREASE_SHARDS_COOLDOWN = Setting.timeSetting( + "data_streams.auto_sharding.increase_shards.cooldown", + TimeValue.timeValueSeconds(270), + TimeValue.timeValueSeconds(0), + Setting.Property.Dynamic, + Setting.Property.NodeScope + ); + + /** + * Represents the minimum amount of time between two scaling events if the next event will reduce the number of shards. + */ + public static final Setting DATA_STREAMS_AUTO_SHARDING_DECREASE_SHARDS_COOLDOWN = Setting.timeSetting( + "data_streams.auto_sharding.decrease_shards.cooldown", + TimeValue.timeValueDays(3), + TimeValue.timeValueSeconds(0), + Setting.Property.Dynamic, + Setting.Property.NodeScope + ); + + /** + * Represents the minimum number of write threads we expect a node to have in the environments where auto sharding will be enabled. + */ + public static final Setting CLUSTER_AUTO_SHARDING_MIN_WRITE_THREADS = Setting.intSetting( + "cluster.auto_sharding.min_write_threads", + 2, + 1, + Setting.Property.Dynamic, + Setting.Property.NodeScope + ); + + /** + * Represents the maximum number of write threads we expect a node to have in the environments where auto sharding will be enabled. + */ + public static final Setting CLUSTER_AUTO_SHARDING_MAX_WRITE_THREADS = Setting.intSetting( + "cluster.auto_sharding.max_write_threads", + 32, + 1, + Setting.Property.Dynamic, + Setting.Property.NodeScope + ); + private final ClusterService clusterService; + private final boolean isAutoShardingEnabled; + private final FeatureService featureService; + private final LongSupplier nowSupplier; + private volatile TimeValue increaseShardsCooldown; + private volatile TimeValue reduceShardsCooldown; + private volatile int minWriteThreads; + private volatile int maxWriteThreads; + private volatile List dataStreamExcludePatterns; + + public DataStreamAutoShardingService( + Settings settings, + ClusterService clusterService, + FeatureService featureService, + LongSupplier nowSupplier + ) { + this.clusterService = clusterService; + this.isAutoShardingEnabled = settings.getAsBoolean(DATA_STREAMS_AUTO_SHARDING_ENABLED, false); + this.increaseShardsCooldown = DATA_STREAMS_AUTO_SHARDING_INCREASE_SHARDS_COOLDOWN.get(settings); + this.reduceShardsCooldown = DATA_STREAMS_AUTO_SHARDING_DECREASE_SHARDS_COOLDOWN.get(settings); + this.minWriteThreads = CLUSTER_AUTO_SHARDING_MIN_WRITE_THREADS.get(settings); + this.maxWriteThreads = CLUSTER_AUTO_SHARDING_MAX_WRITE_THREADS.get(settings); + this.dataStreamExcludePatterns = DATA_STREAMS_AUTO_SHARDING_EXCLUDES_SETTING.get(settings); + this.featureService = featureService; + this.nowSupplier = nowSupplier; + } + + public void init() { + clusterService.getClusterSettings() + .addSettingsUpdateConsumer(DATA_STREAMS_AUTO_SHARDING_INCREASE_SHARDS_COOLDOWN, this::updateIncreaseShardsCooldown); + clusterService.getClusterSettings() + .addSettingsUpdateConsumer(DATA_STREAMS_AUTO_SHARDING_DECREASE_SHARDS_COOLDOWN, this::updateReduceShardsCooldown); + clusterService.getClusterSettings().addSettingsUpdateConsumer(CLUSTER_AUTO_SHARDING_MIN_WRITE_THREADS, this::updateMinWriteThreads); + clusterService.getClusterSettings().addSettingsUpdateConsumer(CLUSTER_AUTO_SHARDING_MAX_WRITE_THREADS, this::updateMaxWriteThreads); + clusterService.getClusterSettings() + .addSettingsUpdateConsumer(DATA_STREAMS_AUTO_SHARDING_EXCLUDES_SETTING, this::updateDataStreamExcludePatterns); + } + + /** + * Computes the optimal number of shards for the provided data stream according to the write index's indexing load (to check if we must + * increase the number of shards, whilst the heuristics for decreasing the number of shards _might_ use the provided write indexing + * load). + * The result type will indicate the recommendation of the auto sharding service : + * - not applicable if the data stream is excluded from auto sharding as configured by + * {@link #DATA_STREAMS_AUTO_SHARDING_EXCLUDES_SETTING} or if the auto sharding functionality is disabled according to + * {@link #DATA_STREAMS_AUTO_SHARDING_EXCLUDES_SETTING}, or if the cluster doesn't have the feature available + * - increase number of shards if the optimal number of shards it deems necessary for the provided data stream is GT the current number + * of shards + * - decrease the number of shards if the optimal number of shards it deems necessary for the provided data stream is LT the current + * number of shards + * + * If the recommendation is to INCREASE/DECREASE shards the reported cooldown period will be TimeValue.ZERO. + * If the auto sharding service thinks the number of shards must be changed but it can't recommend a change due to the cooldown + * period not lapsing, the result will be of type {@link AutoShardingType#COOLDOWN_PREVENTED_INCREASE} or + * {@link AutoShardingType#COOLDOWN_PREVENTED_INCREASE} with the remaining cooldown configured and the number of shards that should + * be configured for the data stream once the remaining cooldown lapses as the target number of shards. + * + * The NOT_APPLICABLE type result will report a cooldown period of TimeValue.MAX_VALUE. + * + * The NO_CHANGE_REQUIRED type will potentially report the remaining cooldown always report a cool down period of TimeValue.ZERO (as + * there'll be no new auto sharding event) + */ + public AutoShardingResult calculate(ClusterState state, DataStream dataStream, @Nullable Double writeIndexLoad) { + Metadata metadata = state.metadata(); + if (isAutoShardingEnabled == false) { + logger.debug("Data stream auto sharding service is not enabled."); + return NOT_APPLICABLE_RESULT; + } + + if (featureService.clusterHasFeature(state, DataStreamAutoShardingService.DATA_STREAM_AUTO_SHARDING_FEATURE) == false) { + logger.debug( + "Data stream auto sharding service cannot compute the optimal number of shards for data stream [{}] because the cluster " + + "doesn't have the auto sharding feature", + dataStream.getName() + ); + return NOT_APPLICABLE_RESULT; + } + + if (dataStreamExcludePatterns.stream().anyMatch(pattern -> Regex.simpleMatch(pattern, dataStream.getName()))) { + logger.debug( + "Data stream [{}] is excluded from auto sharding via the [{}] setting", + dataStream.getName(), + DATA_STREAMS_AUTO_SHARDING_EXCLUDES_SETTING.getKey() + ); + return NOT_APPLICABLE_RESULT; + } + + if (writeIndexLoad == null) { + logger.debug( + "Data stream auto sharding service cannot compute the optimal number of shards for data stream [{}] as the write index " + + "load is not available", + dataStream.getName() + ); + return NOT_APPLICABLE_RESULT; + } + return innerCalculate(metadata, dataStream, writeIndexLoad, nowSupplier); + } + + private AutoShardingResult innerCalculate(Metadata metadata, DataStream dataStream, double writeIndexLoad, LongSupplier nowSupplier) { + // increasing the number of shards is calculated solely based on the index load of the write index + IndexMetadata writeIndex = metadata.index(dataStream.getWriteIndex()); + assert writeIndex != null : "the data stream write index must exist in the provided cluster metadata"; + AutoShardingResult increaseShardsResult = getIncreaseShardsResult(dataStream, writeIndexLoad, nowSupplier, writeIndex); + return Objects.requireNonNullElseGet( + increaseShardsResult, + () -> getDecreaseShardsResult( + metadata, + dataStream, + writeIndexLoad, + nowSupplier, + writeIndex, + getRemainingDecreaseShardsCooldown(metadata, dataStream) + ) + ); + + } + + @Nullable + private AutoShardingResult getIncreaseShardsResult( + DataStream dataStream, + double writeIndexLoad, + LongSupplier nowSupplier, + IndexMetadata writeIndex + ) { + // increasing the number of shards is calculated solely based on the index load of the write index + long optimalShardCount = computeOptimalNumberOfShards(minWriteThreads, maxWriteThreads, writeIndexLoad); + if (optimalShardCount > writeIndex.getNumberOfShards()) { + TimeValue timeSinceLastAutoShardingEvent = dataStream.getAutoShardingEvent() != null + ? dataStream.getAutoShardingEvent().getTimeSinceLastAutoShardingEvent(nowSupplier) + : TimeValue.MAX_VALUE; + + TimeValue coolDownRemaining = TimeValue.timeValueMillis( + Math.max(0L, increaseShardsCooldown.millis() - timeSinceLastAutoShardingEvent.millis()) + ); + logger.debug( + "data stream autosharding service recommends increasing the number of shards from [{}] to [{}] after [{}] cooldown for " + + "data stream [{}]", + writeIndex.getNumberOfShards(), + optimalShardCount, + coolDownRemaining, + dataStream.getName() + ); + return new AutoShardingResult( + coolDownRemaining.equals(TimeValue.ZERO) ? AutoShardingType.INCREASE_SHARDS : AutoShardingType.COOLDOWN_PREVENTED_INCREASE, + writeIndex.getNumberOfShards(), + Math.toIntExact(optimalShardCount), + coolDownRemaining, + writeIndexLoad + ); + } + return null; + } + + /** + * Calculates the amount of time remaining before we can consider reducing the number of shards. + * This reference for the remaining time math is either the time since the last auto sharding event (if available) or otherwise the + * oldest index in the data stream. + */ + private TimeValue getRemainingDecreaseShardsCooldown(Metadata metadata, DataStream dataStream) { + Index oldestBackingIndex = dataStream.getIndices().get(0); + IndexMetadata oldestIndexMeta = metadata.getIndexSafe(oldestBackingIndex); + + return dataStream.getAutoShardingEvent() == null + // without a pre-existing auto sharding event we wait until the oldest index has been created longer than the decrease_shards + // cool down period "ago" so we don't immediately reduce the number of shards after a data stream is created + ? TimeValue.timeValueMillis( + Math.max(0L, oldestIndexMeta.getCreationDate() + reduceShardsCooldown.millis() - nowSupplier.getAsLong()) + ) + : TimeValue.timeValueMillis( + Math.max( + 0L, + reduceShardsCooldown.millis() - dataStream.getAutoShardingEvent() + .getTimeSinceLastAutoShardingEvent(nowSupplier) + .millis() + ) + ); + } + + private AutoShardingResult getDecreaseShardsResult( + Metadata metadata, + DataStream dataStream, + double writeIndexLoad, + LongSupplier nowSupplier, + IndexMetadata writeIndex, + TimeValue remainingReduceShardsCooldown + ) { + double maxIndexLoadWithinCoolingPeriod = getMaxIndexLoadWithinCoolingPeriod( + metadata, + dataStream, + writeIndexLoad, + reduceShardsCooldown, + nowSupplier + ); + + logger.trace( + "calculating the optimal number of shards for a potential decrease in number of shards for data stream [{}] with the" + + " max indexing load [{}] over the decrease shards cool down period", + dataStream.getName(), + maxIndexLoadWithinCoolingPeriod + ); + long optimalShardCount = computeOptimalNumberOfShards(minWriteThreads, maxWriteThreads, maxIndexLoadWithinCoolingPeriod); + if (optimalShardCount < writeIndex.getNumberOfShards()) { + logger.debug( + "data stream autosharding service recommends decreasing the number of shards from [{}] to [{}] after [{}] cooldown for " + + "data stream [{}]", + writeIndex.getNumberOfShards(), + optimalShardCount, + remainingReduceShardsCooldown, + dataStream.getName() + ); + + // we should reduce the number of shards + return new AutoShardingResult( + remainingReduceShardsCooldown.equals(TimeValue.ZERO) + ? AutoShardingType.DECREASE_SHARDS + : AutoShardingType.COOLDOWN_PREVENTED_DECREASE, + writeIndex.getNumberOfShards(), + Math.toIntExact(optimalShardCount), + remainingReduceShardsCooldown, + maxIndexLoadWithinCoolingPeriod + ); + } + + logger.trace( + "data stream autosharding service recommends maintaining the number of shards [{}] for data stream [{}]", + writeIndex.getNumberOfShards(), + dataStream.getName() + ); + return new AutoShardingResult( + AutoShardingType.NO_CHANGE_REQUIRED, + writeIndex.getNumberOfShards(), + writeIndex.getNumberOfShards(), + TimeValue.ZERO, + maxIndexLoadWithinCoolingPeriod + ); + } + + // Visible for testing + static long computeOptimalNumberOfShards(int minNumberWriteThreads, int maxNumberWriteThreads, double indexingLoad) { + return Math.max( + Math.min(roundUp(indexingLoad / (minNumberWriteThreads / 2.0)), 3), + roundUp(indexingLoad / (maxNumberWriteThreads / 2.0)) + ); + } + + private static long roundUp(double value) { + return (long) Math.ceil(value); + } + + // Visible for testing + /** + * Calculates the maximum write index load observed for the provided data stream across all the backing indices that were created + * during the provide {@param coolingPeriod} (note: to cover the entire cooling period, the backing index created before the cooling + * period is also considered). + */ + static double getMaxIndexLoadWithinCoolingPeriod( + Metadata metadata, + DataStream dataStream, + double writeIndexLoad, + TimeValue coolingPeriod, + LongSupplier nowSupplier + ) { + // for reducing the number of shards we look at more than just the write index + List writeLoadsWithinCoolingPeriod = DataStream.getIndicesWithinMaxAgeRange( + dataStream, + metadata::getIndexSafe, + coolingPeriod, + nowSupplier + ) + .stream() + .filter(index -> index.equals(dataStream.getWriteIndex()) == false) + .map(metadata::index) + .filter(Objects::nonNull) + .map(IndexMetadata::getStats) + .filter(Objects::nonNull) + .map(IndexMetadataStats::writeLoad) + .filter(Objects::nonNull) + .toList(); + + // assume the current write index load is the highest observed and look back to find the actual maximum + double maxIndexLoadWithinCoolingPeriod = writeIndexLoad; + for (IndexWriteLoad writeLoad : writeLoadsWithinCoolingPeriod) { + double totalIndexLoad = 0; + for (int shardId = 0; shardId < writeLoad.numberOfShards(); shardId++) { + final OptionalDouble writeLoadForShard = writeLoad.getWriteLoadForShard(shardId); + totalIndexLoad += writeLoadForShard.orElse(0); + } + + if (totalIndexLoad > maxIndexLoadWithinCoolingPeriod) { + maxIndexLoadWithinCoolingPeriod = totalIndexLoad; + } + } + return maxIndexLoadWithinCoolingPeriod; + } + + void updateIncreaseShardsCooldown(TimeValue scaleUpCooldown) { + this.increaseShardsCooldown = scaleUpCooldown; + } + + void updateReduceShardsCooldown(TimeValue scaleDownCooldown) { + this.reduceShardsCooldown = scaleDownCooldown; + } + + void updateMinWriteThreads(int minNumberWriteThreads) { + this.minWriteThreads = minNumberWriteThreads; + } + + void updateMaxWriteThreads(int maxNumberWriteThreads) { + this.maxWriteThreads = maxNumberWriteThreads; + } + + private void updateDataStreamExcludePatterns(List newExcludePatterns) { + this.dataStreamExcludePatterns = newExcludePatterns; + } +} diff --git a/server/src/main/java/org/elasticsearch/cluster/metadata/DataStream.java b/server/src/main/java/org/elasticsearch/cluster/metadata/DataStream.java index 1bcfdba1d16f4..66cef1ea49af0 100644 --- a/server/src/main/java/org/elasticsearch/cluster/metadata/DataStream.java +++ b/server/src/main/java/org/elasticsearch/cluster/metadata/DataStream.java @@ -70,6 +70,7 @@ public final class DataStream implements SimpleDiffable, ToXContentO public static final FeatureFlag FAILURE_STORE_FEATURE_FLAG = new FeatureFlag("failure_store"); public static final TransportVersion ADDED_FAILURE_STORE_TRANSPORT_VERSION = TransportVersions.V_8_12_0; + public static final TransportVersion ADDED_AUTO_SHARDING_EVENT_VERSION = TransportVersions.DATA_STREAM_AUTO_SHARDING_EVENT; public static boolean isFailureStoreEnabled() { return FAILURE_STORE_FEATURE_FLAG.isEnabled(); @@ -113,6 +114,8 @@ public static boolean isFailureStoreEnabled() { private final boolean rolloverOnWrite; private final boolean failureStore; private final List failureIndices; + @Nullable + private final DataStreamAutoShardingEvent autoShardingEvent; public DataStream( String name, @@ -126,7 +129,8 @@ public DataStream( IndexMode indexMode, DataStreamLifecycle lifecycle, boolean failureStore, - List failureIndices + List failureIndices, + @Nullable DataStreamAutoShardingEvent autoShardingEvent ) { this( name, @@ -142,7 +146,8 @@ public DataStream( lifecycle, failureStore, failureIndices, - false + false, + autoShardingEvent ); } @@ -159,7 +164,8 @@ public DataStream( DataStreamLifecycle lifecycle, boolean failureStore, List failureIndices, - boolean rolloverOnWrite + boolean rolloverOnWrite, + @Nullable DataStreamAutoShardingEvent autoShardingEvent ) { this( name, @@ -175,7 +181,8 @@ public DataStream( lifecycle, failureStore, failureIndices, - rolloverOnWrite + rolloverOnWrite, + autoShardingEvent ); } @@ -194,7 +201,8 @@ public DataStream( DataStreamLifecycle lifecycle, boolean failureStore, List failureIndices, - boolean rolloverOnWrite + boolean rolloverOnWrite, + @Nullable DataStreamAutoShardingEvent autoShardingEvent ) { this.name = name; this.indices = List.copyOf(indices); @@ -213,6 +221,7 @@ public DataStream( this.failureIndices = failureIndices; assert assertConsistent(this.indices); this.rolloverOnWrite = rolloverOnWrite; + this.autoShardingEvent = autoShardingEvent; } // mainly available for testing @@ -227,7 +236,7 @@ public DataStream( boolean allowCustomRouting, IndexMode indexMode ) { - this(name, indices, generation, metadata, hidden, replicated, system, allowCustomRouting, indexMode, null, false, List.of()); + this(name, indices, generation, metadata, hidden, replicated, system, allowCustomRouting, indexMode, null, false, List.of(), null); } private static boolean assertConsistent(List indices) { @@ -412,6 +421,13 @@ public DataStreamLifecycle getLifecycle() { return lifecycle; } + /** + * Returns the latest auto sharding event that happened for this data stream + */ + public DataStreamAutoShardingEvent getAutoShardingEvent() { + return autoShardingEvent; + } + /** * Performs a rollover on a {@code DataStream} instance and returns a new instance containing * the updated list of backing indices and incremented generation. @@ -456,7 +472,8 @@ public DataStream unsafeRollover(Index writeIndex, long generation, boolean time indexMode, lifecycle, failureStore, - failureIndices + failureIndices, + autoShardingEvent ); } @@ -534,7 +551,8 @@ public DataStream removeBackingIndex(Index index) { indexMode, lifecycle, failureStore, - failureIndices + failureIndices, + autoShardingEvent ); } @@ -579,7 +597,8 @@ public DataStream replaceBackingIndex(Index existingBackingIndex, Index newBacki indexMode, lifecycle, failureStore, - failureIndices + failureIndices, + autoShardingEvent ); } @@ -639,7 +658,8 @@ public DataStream addBackingIndex(Metadata clusterMetadata, Index index) { indexMode, lifecycle, failureStore, - failureIndices + failureIndices, + autoShardingEvent ); } @@ -658,7 +678,8 @@ public DataStream promoteDataStream() { lifecycle, failureStore, failureIndices, - rolloverOnWrite + rolloverOnWrite, + autoShardingEvent ); } @@ -694,7 +715,8 @@ public DataStream snapshot(Collection indicesInSnapshot) { indexMode, lifecycle, failureStore, - failureIndices + failureIndices, + autoShardingEvent ); } @@ -909,7 +931,10 @@ public DataStream(StreamInput in) throws IOException { in.getTransportVersion().onOrAfter(TransportVersions.V_8_9_X) ? in.readOptionalWriteable(DataStreamLifecycle::new) : null, in.getTransportVersion().onOrAfter(DataStream.ADDED_FAILURE_STORE_TRANSPORT_VERSION) ? in.readBoolean() : false, in.getTransportVersion().onOrAfter(DataStream.ADDED_FAILURE_STORE_TRANSPORT_VERSION) ? readIndices(in) : List.of(), - in.getTransportVersion().onOrAfter(TransportVersions.LAZY_ROLLOVER_ADDED) ? in.readBoolean() : false + in.getTransportVersion().onOrAfter(TransportVersions.LAZY_ROLLOVER_ADDED) ? in.readBoolean() : false, + in.getTransportVersion().onOrAfter(DataStream.ADDED_AUTO_SHARDING_EVENT_VERSION) + ? in.readOptionalWriteable(DataStreamAutoShardingEvent::new) + : null ); } @@ -953,6 +978,9 @@ public void writeTo(StreamOutput out) throws IOException { if (out.getTransportVersion().onOrAfter(TransportVersions.LAZY_ROLLOVER_ADDED)) { out.writeBoolean(rolloverOnWrite); } + if (out.getTransportVersion().onOrAfter(DataStream.ADDED_AUTO_SHARDING_EVENT_VERSION)) { + out.writeOptionalWriteable(autoShardingEvent); + } } public static final ParseField NAME_FIELD = new ParseField("name"); @@ -969,13 +997,14 @@ public void writeTo(StreamOutput out) throws IOException { public static final ParseField FAILURE_STORE_FIELD = new ParseField("failure_store"); public static final ParseField FAILURE_INDICES_FIELD = new ParseField("failure_indices"); public static final ParseField ROLLOVER_ON_WRITE_FIELD = new ParseField("rollover_on_write"); + public static final ParseField AUTO_SHARDING_FIELD = new ParseField("auto_sharding"); @SuppressWarnings("unchecked") private static final ConstructingObjectParser PARSER = new ConstructingObjectParser<>("data_stream", args -> { // Fields behind a feature flag need to be parsed last otherwise the parser will fail when the feature flag is disabled. // Until the feature flag is removed we keep them separately to be mindful of this. - boolean failureStoreEnabled = DataStream.isFailureStoreEnabled() && args[11] != null && (boolean) args[11]; - List failureStoreIndices = DataStream.isFailureStoreEnabled() && args[12] != null ? (List) args[12] : List.of(); + boolean failureStoreEnabled = DataStream.isFailureStoreEnabled() && args[12] != null && (boolean) args[12]; + List failureStoreIndices = DataStream.isFailureStoreEnabled() && args[13] != null ? (List) args[13] : List.of(); return new DataStream( (String) args[0], (List) args[1], @@ -989,7 +1018,8 @@ public void writeTo(StreamOutput out) throws IOException { (DataStreamLifecycle) args[9], failureStoreEnabled, failureStoreIndices, - args[10] != null && (boolean) args[10] + args[10] != null && (boolean) args[10], + (DataStreamAutoShardingEvent) args[11] ); }); @@ -1013,6 +1043,11 @@ public void writeTo(StreamOutput out) throws IOException { PARSER.declareString(ConstructingObjectParser.optionalConstructorArg(), INDEX_MODE); PARSER.declareObject(ConstructingObjectParser.optionalConstructorArg(), (p, c) -> DataStreamLifecycle.fromXContent(p), LIFECYCLE); PARSER.declareBoolean(ConstructingObjectParser.optionalConstructorArg(), ROLLOVER_ON_WRITE_FIELD); + PARSER.declareObject( + ConstructingObjectParser.optionalConstructorArg(), + (p, c) -> DataStreamAutoShardingEvent.fromXContent(p), + AUTO_SHARDING_FIELD + ); // The fields behind the feature flag should always be last. if (DataStream.isFailureStoreEnabled()) { PARSER.declareBoolean(ConstructingObjectParser.optionalConstructorArg(), FAILURE_STORE_FIELD); @@ -1067,6 +1102,11 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params, @Nulla lifecycle.toXContent(builder, params, rolloverConfiguration); } builder.field(ROLLOVER_ON_WRITE_FIELD.getPreferredName(), rolloverOnWrite); + if (autoShardingEvent != null) { + builder.startObject(AUTO_SHARDING_FIELD.getPreferredName()); + autoShardingEvent.toXContent(builder, params); + builder.endObject(); + } builder.endObject(); return builder; } @@ -1088,7 +1128,8 @@ public boolean equals(Object o) { && Objects.equals(lifecycle, that.lifecycle) && failureStore == that.failureStore && failureIndices.equals(that.failureIndices) - && rolloverOnWrite == that.rolloverOnWrite; + && rolloverOnWrite == that.rolloverOnWrite + && Objects.equals(autoShardingEvent, that.autoShardingEvent); } @Override @@ -1106,7 +1147,8 @@ public int hashCode() { lifecycle, failureStore, failureIndices, - rolloverOnWrite + rolloverOnWrite, + autoShardingEvent ); } @@ -1169,6 +1211,34 @@ public DataStream getParentDataStream() { "strict_date_optional_time_nanos||strict_date_optional_time||epoch_millis" ); + /** + * Returns the indices created within the {@param maxIndexAge} interval. Note that this strives to cover + * the entire {@param maxIndexAge} interval so one backing index created before the specified age will also + * be return. + */ + public static List getIndicesWithinMaxAgeRange( + DataStream dataStream, + Function indexProvider, + TimeValue maxIndexAge, + LongSupplier nowSupplier + ) { + final List dataStreamIndices = dataStream.getIndices(); + final long currentTimeMillis = nowSupplier.getAsLong(); + // Consider at least 1 index (including the write index) for cases where rollovers happen less often than maxIndexAge + int firstIndexWithinAgeRange = Math.max(dataStreamIndices.size() - 2, 0); + for (int i = 0; i < dataStreamIndices.size(); i++) { + Index index = dataStreamIndices.get(i); + final IndexMetadata indexMetadata = indexProvider.apply(index); + final long indexAge = currentTimeMillis - indexMetadata.getCreationDate(); + if (indexAge < maxIndexAge.getMillis()) { + // We need to consider the previous index too in order to cover the entire max-index-age range. + firstIndexWithinAgeRange = i == 0 ? 0 : i - 1; + break; + } + } + return dataStreamIndices.subList(firstIndexWithinAgeRange, dataStreamIndices.size()); + } + private static Instant getTimeStampFromRaw(Object rawTimestamp) { try { if (rawTimestamp instanceof Long lTimestamp) { diff --git a/server/src/main/java/org/elasticsearch/cluster/metadata/DataStreamAutoShardingEvent.java b/server/src/main/java/org/elasticsearch/cluster/metadata/DataStreamAutoShardingEvent.java new file mode 100644 index 0000000000000..ff143681827ca --- /dev/null +++ b/server/src/main/java/org/elasticsearch/cluster/metadata/DataStreamAutoShardingEvent.java @@ -0,0 +1,84 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0 and the Server Side Public License, v 1; you may not use this file except + * in compliance with, at your election, the Elastic License 2.0 or the Server + * Side Public License, v 1. + */ + +package org.elasticsearch.cluster.metadata; + +import org.elasticsearch.cluster.Diff; +import org.elasticsearch.cluster.SimpleDiffable; +import org.elasticsearch.common.io.stream.StreamInput; +import org.elasticsearch.common.io.stream.StreamOutput; +import org.elasticsearch.core.TimeValue; +import org.elasticsearch.xcontent.ConstructingObjectParser; +import org.elasticsearch.xcontent.ParseField; +import org.elasticsearch.xcontent.ToXContentFragment; +import org.elasticsearch.xcontent.XContentBuilder; +import org.elasticsearch.xcontent.XContentParser; + +import java.io.IOException; +import java.util.function.LongSupplier; + +/** + * Represents the last auto sharding event that occured for a data stream. + */ +public record DataStreamAutoShardingEvent(String triggerIndexName, int targetNumberOfShards, long timestamp) + implements + SimpleDiffable, + ToXContentFragment { + + public static final ParseField TRIGGER_INDEX_NAME = new ParseField("trigger_index_name"); + public static final ParseField TARGET_NUMBER_OF_SHARDS = new ParseField("target_number_of_shards"); + public static final ParseField EVENT_TIME = new ParseField("event_time"); + public static final ParseField EVENT_TIME_MILLIS = new ParseField("event_time_millis"); + + public static final ConstructingObjectParser PARSER = new ConstructingObjectParser<>( + "auto_sharding", + false, + (args, unused) -> new DataStreamAutoShardingEvent((String) args[0], (int) args[1], (long) args[2]) + ); + + static { + PARSER.declareString(ConstructingObjectParser.optionalConstructorArg(), TRIGGER_INDEX_NAME); + PARSER.declareInt(ConstructingObjectParser.optionalConstructorArg(), TARGET_NUMBER_OF_SHARDS); + PARSER.declareLong(ConstructingObjectParser.optionalConstructorArg(), EVENT_TIME_MILLIS); + } + + public static DataStreamAutoShardingEvent fromXContent(XContentParser parser) throws IOException { + return PARSER.parse(parser, null); + } + + static Diff readDiffFrom(StreamInput in) throws IOException { + return SimpleDiffable.readDiffFrom(DataStreamAutoShardingEvent::new, in); + } + + DataStreamAutoShardingEvent(StreamInput in) throws IOException { + this(in.readString(), in.readVInt(), in.readVLong()); + } + + @Override + public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { + builder.field(TRIGGER_INDEX_NAME.getPreferredName(), triggerIndexName); + builder.field(TARGET_NUMBER_OF_SHARDS.getPreferredName(), targetNumberOfShards); + builder.humanReadableField( + EVENT_TIME_MILLIS.getPreferredName(), + EVENT_TIME.getPreferredName(), + TimeValue.timeValueMillis(timestamp) + ); + return builder; + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + out.writeString(triggerIndexName); + out.writeVInt(targetNumberOfShards); + out.writeVLong(timestamp); + } + + public TimeValue getTimeSinceLastAutoShardingEvent(LongSupplier now) { + return TimeValue.timeValueMillis(Math.max(0L, now.getAsLong() - timestamp)); + } +} diff --git a/server/src/main/java/org/elasticsearch/cluster/metadata/MetadataCreateDataStreamService.java b/server/src/main/java/org/elasticsearch/cluster/metadata/MetadataCreateDataStreamService.java index d500a8b8e6876..20b28edef5ca2 100644 --- a/server/src/main/java/org/elasticsearch/cluster/metadata/MetadataCreateDataStreamService.java +++ b/server/src/main/java/org/elasticsearch/cluster/metadata/MetadataCreateDataStreamService.java @@ -314,7 +314,8 @@ static ClusterState createDataStream( indexMode, lifecycle == null && isDslOnlyMode ? DataStreamLifecycle.DEFAULT : lifecycle, template.getDataStreamTemplate().hasFailureStore(), - failureIndices + failureIndices, + null ); Metadata.Builder builder = Metadata.builder(currentState.metadata()).put(newDataStream); diff --git a/server/src/main/java/org/elasticsearch/cluster/metadata/MetadataDataStreamsService.java b/server/src/main/java/org/elasticsearch/cluster/metadata/MetadataDataStreamsService.java index 6b81aa230f0d9..4006bc8d1a94a 100644 --- a/server/src/main/java/org/elasticsearch/cluster/metadata/MetadataDataStreamsService.java +++ b/server/src/main/java/org/elasticsearch/cluster/metadata/MetadataDataStreamsService.java @@ -212,7 +212,8 @@ static ClusterState updateDataLifecycle( dataStream.getIndexMode(), lifecycle, dataStream.isFailureStore(), - dataStream.getFailureIndices() + dataStream.getFailureIndices(), + dataStream.getAutoShardingEvent() ) ); } @@ -249,7 +250,8 @@ public static ClusterState setRolloverOnWrite(ClusterState currentState, String dataStream.getLifecycle(), dataStream.isFailureStore(), dataStream.getFailureIndices(), - rolloverOnWrite + rolloverOnWrite, + dataStream.getAutoShardingEvent() ) ); return ClusterState.builder(currentState).metadata(builder.build()).build(); diff --git a/server/src/main/java/org/elasticsearch/snapshots/RestoreService.java b/server/src/main/java/org/elasticsearch/snapshots/RestoreService.java index 9ac76e653b640..4b6e3f30fe6fa 100644 --- a/server/src/main/java/org/elasticsearch/snapshots/RestoreService.java +++ b/server/src/main/java/org/elasticsearch/snapshots/RestoreService.java @@ -716,7 +716,8 @@ static DataStream updateDataStream(DataStream dataStream, Metadata.Builder metad dataStream.getIndexMode(), dataStream.getLifecycle(), dataStream.isFailureStore(), - dataStream.getFailureIndices() + dataStream.getFailureIndices(), + dataStream.getAutoShardingEvent() ); } diff --git a/server/src/test/java/org/elasticsearch/action/datastreams/autosharding/DataStreamAutoShardingServiceTests.java b/server/src/test/java/org/elasticsearch/action/datastreams/autosharding/DataStreamAutoShardingServiceTests.java new file mode 100644 index 0000000000000..674b3e855e912 --- /dev/null +++ b/server/src/test/java/org/elasticsearch/action/datastreams/autosharding/DataStreamAutoShardingServiceTests.java @@ -0,0 +1,771 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0 and the Server Side Public License, v 1; you may not use this file except + * in compliance with, at your election, the Elastic License 2.0 or the Server + * Side Public License, v 1. + */ + +package org.elasticsearch.action.datastreams.autosharding; + +import org.elasticsearch.action.admin.indices.rollover.MaxAgeCondition; +import org.elasticsearch.action.admin.indices.rollover.RolloverInfo; +import org.elasticsearch.cluster.ClusterName; +import org.elasticsearch.cluster.ClusterState; +import org.elasticsearch.cluster.metadata.DataStream; +import org.elasticsearch.cluster.metadata.DataStreamAutoShardingEvent; +import org.elasticsearch.cluster.metadata.IndexMetadata; +import org.elasticsearch.cluster.metadata.IndexMetadataStats; +import org.elasticsearch.cluster.metadata.IndexWriteLoad; +import org.elasticsearch.cluster.metadata.Metadata; +import org.elasticsearch.cluster.service.ClusterService; +import org.elasticsearch.common.settings.ClusterSettings; +import org.elasticsearch.common.settings.Setting; +import org.elasticsearch.common.settings.Settings; +import org.elasticsearch.core.Nullable; +import org.elasticsearch.core.TimeValue; +import org.elasticsearch.features.FeatureService; +import org.elasticsearch.features.FeatureSpecification; +import org.elasticsearch.features.NodeFeature; +import org.elasticsearch.index.Index; +import org.elasticsearch.index.IndexMode; +import org.elasticsearch.index.IndexVersion; +import org.elasticsearch.test.ESTestCase; +import org.elasticsearch.threadpool.TestThreadPool; +import org.elasticsearch.threadpool.ThreadPool; +import org.junit.After; +import org.junit.Before; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.function.Function; + +import static org.elasticsearch.action.datastreams.autosharding.AutoShardingResult.NOT_APPLICABLE_RESULT; +import static org.elasticsearch.action.datastreams.autosharding.AutoShardingType.COOLDOWN_PREVENTED_DECREASE; +import static org.elasticsearch.action.datastreams.autosharding.AutoShardingType.COOLDOWN_PREVENTED_INCREASE; +import static org.elasticsearch.action.datastreams.autosharding.AutoShardingType.DECREASE_SHARDS; +import static org.elasticsearch.action.datastreams.autosharding.AutoShardingType.INCREASE_SHARDS; +import static org.elasticsearch.action.datastreams.autosharding.AutoShardingType.NO_CHANGE_REQUIRED; +import static org.elasticsearch.test.ClusterServiceUtils.createClusterService; +import static org.hamcrest.Matchers.is; + +public class DataStreamAutoShardingServiceTests extends ESTestCase { + + private ClusterService clusterService; + private ThreadPool threadPool; + private DataStreamAutoShardingService service; + private long now; + String dataStreamName; + + @Before + public void setupService() { + threadPool = new TestThreadPool(getTestName()); + Set> builtInClusterSettings = new HashSet<>(ClusterSettings.BUILT_IN_CLUSTER_SETTINGS); + builtInClusterSettings.add(DataStreamAutoShardingService.CLUSTER_AUTO_SHARDING_MIN_WRITE_THREADS); + builtInClusterSettings.add(DataStreamAutoShardingService.CLUSTER_AUTO_SHARDING_MAX_WRITE_THREADS); + builtInClusterSettings.add(DataStreamAutoShardingService.DATA_STREAMS_AUTO_SHARDING_INCREASE_SHARDS_COOLDOWN); + builtInClusterSettings.add(DataStreamAutoShardingService.DATA_STREAMS_AUTO_SHARDING_DECREASE_SHARDS_COOLDOWN); + builtInClusterSettings.add( + Setting.boolSetting( + DataStreamAutoShardingService.DATA_STREAMS_AUTO_SHARDING_ENABLED, + false, + Setting.Property.Dynamic, + Setting.Property.NodeScope + ) + ); + ClusterSettings clusterSettings = new ClusterSettings(Settings.EMPTY, builtInClusterSettings); + clusterService = createClusterService(threadPool, clusterSettings); + now = System.currentTimeMillis(); + service = new DataStreamAutoShardingService( + Settings.builder() + .put(DataStreamAutoShardingService.DATA_STREAMS_AUTO_SHARDING_ENABLED, true) + .putList(DataStreamAutoShardingService.DATA_STREAMS_AUTO_SHARDING_EXCLUDES_SETTING.getKey(), List.of()) + .build(), + clusterService, + new FeatureService(List.of(new FeatureSpecification() { + @Override + public Set getFeatures() { + return Set.of(DataStreamAutoShardingService.DATA_STREAM_AUTO_SHARDING_FEATURE); + } + })), + () -> now + ); + dataStreamName = randomAlphaOfLengthBetween(10, 100); + logger.info("-> data stream name is [{}]", dataStreamName); + } + + @After + public void cleanup() { + clusterService.close(); + threadPool.shutdownNow(); + } + + public void testCalculateValidations() { + Metadata.Builder builder = Metadata.builder(); + DataStream dataStream = createDataStream( + builder, + dataStreamName, + 1, + now, + List.of(now - 3000, now - 2000, now - 1000), + getWriteLoad(1, 2.0), + null + ); + builder.put(dataStream); + ClusterState state = ClusterState.builder(ClusterName.DEFAULT) + .nodeFeatures( + Map.of( + "n1", + Set.of(DataStreamAutoShardingService.DATA_STREAM_AUTO_SHARDING_FEATURE.id()), + "n2", + Set.of(DataStreamAutoShardingService.DATA_STREAM_AUTO_SHARDING_FEATURE.id()) + ) + ) + .metadata(builder) + .build(); + + { + // autosharding disabled + DataStreamAutoShardingService disabledAutoshardingService = new DataStreamAutoShardingService( + Settings.EMPTY, + clusterService, + new FeatureService(List.of(new FeatureSpecification() { + @Override + public Set getFeatures() { + return Set.of(DataStreamAutoShardingService.DATA_STREAM_AUTO_SHARDING_FEATURE); + } + })), + System::currentTimeMillis + ); + + AutoShardingResult autoShardingResult = disabledAutoshardingService.calculate(state, dataStream, 2.0); + assertThat(autoShardingResult, is(NOT_APPLICABLE_RESULT)); + } + + { + // cluster doesn't have feature + ClusterState stateNoFeature = ClusterState.builder(ClusterName.DEFAULT).metadata(Metadata.builder()).build(); + + DataStreamAutoShardingService noFeatureService = new DataStreamAutoShardingService( + Settings.builder() + .put(DataStreamAutoShardingService.DATA_STREAMS_AUTO_SHARDING_ENABLED, true) + .putList(DataStreamAutoShardingService.DATA_STREAMS_AUTO_SHARDING_EXCLUDES_SETTING.getKey(), List.of()) + .build(), + clusterService, + new FeatureService(List.of()), + () -> now + ); + + AutoShardingResult autoShardingResult = noFeatureService.calculate(stateNoFeature, dataStream, 2.0); + assertThat(autoShardingResult, is(NOT_APPLICABLE_RESULT)); + } + + { + // patterns are configured to exclude the current data stream + DataStreamAutoShardingService noFeatureService = new DataStreamAutoShardingService( + Settings.builder() + .put(DataStreamAutoShardingService.DATA_STREAMS_AUTO_SHARDING_ENABLED, true) + .putList( + DataStreamAutoShardingService.DATA_STREAMS_AUTO_SHARDING_EXCLUDES_SETTING.getKey(), + List.of("foo", dataStreamName + "*") + ) + .build(), + clusterService, + new FeatureService(List.of()), + () -> now + ); + + AutoShardingResult autoShardingResult = noFeatureService.calculate(state, dataStream, 2.0); + assertThat(autoShardingResult, is(NOT_APPLICABLE_RESULT)); + } + + { + // null write load passed + AutoShardingResult autoShardingResult = service.calculate(state, dataStream, null); + assertThat(autoShardingResult, is(NOT_APPLICABLE_RESULT)); + } + } + + public void testCalculateIncreaseShardingRecommendations() { + // the input is a data stream with 5 backing indices with 1 shard each + // all 4 backing indices have a write load of 2.0 + // we'll recreate it across the test and add an auto sharding event as we iterate + { + Metadata.Builder builder = Metadata.builder(); + Function dataStreamSupplier = (autoShardingEvent) -> createDataStream( + builder, + dataStreamName, + 1, + now, + List.of(now - 10_000, now - 7000, now - 5000, now - 2000, now - 1000), + getWriteLoad(1, 2.0), + autoShardingEvent + ); + + DataStream dataStream = dataStreamSupplier.apply(null); + builder.put(dataStream); + ClusterState state = ClusterState.builder(ClusterName.DEFAULT) + .nodeFeatures( + Map.of( + "n1", + Set.of(DataStreamAutoShardingService.DATA_STREAM_AUTO_SHARDING_FEATURE.id()), + "n2", + Set.of(DataStreamAutoShardingService.DATA_STREAM_AUTO_SHARDING_FEATURE.id()) + ) + ) + .metadata(builder) + .build(); + + AutoShardingResult autoShardingResult = service.calculate(state, dataStream, 2.5); + assertThat(autoShardingResult.type(), is(INCREASE_SHARDS)); + // no pre-existing scaling event so the cool down must be zero + assertThat(autoShardingResult.coolDownRemaining(), is(TimeValue.ZERO)); + assertThat(autoShardingResult.targetNumberOfShards(), is(3)); + } + + { + // let's add a pre-existing sharding event so that we'll return some cool down period that's preventing an INCREASE_SHARDS + // event so the result type we're expecting is COOLDOWN_PREVENTED_INCREASE + Metadata.Builder builder = Metadata.builder(); + Function dataStreamSupplier = (autoShardingEvent) -> createDataStream( + builder, + dataStreamName, + 1, + now, + List.of(now - 10_000, now - 7000, now - 5000, now - 2000, now - 1000), + getWriteLoad(1, 2.0), + autoShardingEvent + ); + + // generation 4 triggered an auto sharding event to 2 shards + DataStream dataStream = dataStreamSupplier.apply( + new DataStreamAutoShardingEvent(DataStream.getDefaultBackingIndexName(dataStreamName, 4), 2, now - 1005) + ); + builder.put(dataStream); + ClusterState state = ClusterState.builder(ClusterName.DEFAULT) + .nodeFeatures( + Map.of( + "n1", + Set.of(DataStreamAutoShardingService.DATA_STREAM_AUTO_SHARDING_FEATURE.id()), + "n2", + Set.of(DataStreamAutoShardingService.DATA_STREAM_AUTO_SHARDING_FEATURE.id()) + ) + ) + .metadata(builder) + .build(); + + AutoShardingResult autoShardingResult = service.calculate(state, dataStream, 2.5); + assertThat(autoShardingResult.type(), is(COOLDOWN_PREVENTED_INCREASE)); + // no pre-existing scaling event so the cool down must be zero + assertThat(autoShardingResult.targetNumberOfShards(), is(3)); + // it's been 1005 millis since the last auto sharding event and the cool down is 270secoinds (270_000 millis) + assertThat(autoShardingResult.coolDownRemaining(), is(TimeValue.timeValueMillis(268995))); + } + + { + // let's test a subsequent increase in the number of shards after a previos auto sharding event + Metadata.Builder builder = Metadata.builder(); + Function dataStreamSupplier = (autoShardingEvent) -> createDataStream( + builder, + dataStreamName, + 1, + now, + List.of(now - 10_000_000, now - 7_000_000, now - 2_000_000, now - 1_000_000, now - 1000), + getWriteLoad(1, 2.0), + autoShardingEvent + ); + + // generation 3 triggered an increase in shards event to 2 shards + DataStream dataStream = dataStreamSupplier.apply( + new DataStreamAutoShardingEvent(DataStream.getDefaultBackingIndexName(dataStreamName, 4), 2, now - 2_000_100) + ); + builder.put(dataStream); + ClusterState state = ClusterState.builder(ClusterName.DEFAULT) + .nodeFeatures( + Map.of( + "n1", + Set.of(DataStreamAutoShardingService.DATA_STREAM_AUTO_SHARDING_FEATURE.id()), + "n2", + Set.of(DataStreamAutoShardingService.DATA_STREAM_AUTO_SHARDING_FEATURE.id()) + ) + ) + .metadata(builder) + .build(); + + AutoShardingResult autoShardingResult = service.calculate(state, dataStream, 2.5); + assertThat(autoShardingResult.type(), is(INCREASE_SHARDS)); + // no pre-existing scaling event so the cool down must be zero + assertThat(autoShardingResult.targetNumberOfShards(), is(3)); + assertThat(autoShardingResult.coolDownRemaining(), is(TimeValue.ZERO)); + } + } + + public void testCalculateDecreaseShardingRecommendations() { + // the input is a data stream with 5 backing indices with 3 shards each + { + // testing a decrease shards events prevented by the cool down period not lapsing due to the oldest generation index being + // "too new" (i.e. the cool down period hasn't lapsed since the oldest generation index) + Metadata.Builder builder = Metadata.builder(); + Function dataStreamSupplier = (autoShardingEvent) -> createDataStream( + builder, + dataStreamName, + 3, + now, + List.of(now - 10_000, now - 7000, now - 5000, now - 2000, now - 1000), + getWriteLoad(3, 0.25), + autoShardingEvent + ); + + DataStream dataStream = dataStreamSupplier.apply(null); + builder.put(dataStream); + ClusterState state = ClusterState.builder(ClusterName.DEFAULT) + .nodeFeatures( + Map.of( + "n1", + Set.of(DataStreamAutoShardingService.DATA_STREAM_AUTO_SHARDING_FEATURE.id()), + "n2", + Set.of(DataStreamAutoShardingService.DATA_STREAM_AUTO_SHARDING_FEATURE.id()) + ) + ) + .metadata(builder) + .build(); + + AutoShardingResult autoShardingResult = service.calculate(state, dataStream, 1.0); + // the cooldown period for the decrease shards event hasn't lapsed since the data stream was created + assertThat(autoShardingResult.type(), is(COOLDOWN_PREVENTED_DECREASE)); + assertThat(autoShardingResult.coolDownRemaining(), is(TimeValue.timeValueMillis(TimeValue.timeValueDays(3).millis() - 10_000))); + // based on the write load of 0.75 we should be reducing the number of shards to 1 + assertThat(autoShardingResult.targetNumberOfShards(), is(1)); + } + + { + Metadata.Builder builder = Metadata.builder(); + Function dataStreamSupplier = (autoShardingEvent) -> createDataStream( + builder, + dataStreamName, + 3, + now, + List.of( + now - TimeValue.timeValueDays(21).getMillis(), + now - TimeValue.timeValueDays(15).getMillis(), + now - TimeValue.timeValueDays(4).getMillis(), + now - TimeValue.timeValueDays(2).getMillis(), + now - 1000 + ), + getWriteLoad(3, 0.333), + autoShardingEvent + ); + + DataStream dataStream = dataStreamSupplier.apply(null); + builder.put(dataStream); + ClusterState state = ClusterState.builder(ClusterName.DEFAULT) + .nodeFeatures( + Map.of( + "n1", + Set.of(DataStreamAutoShardingService.DATA_STREAM_AUTO_SHARDING_FEATURE.id()), + "n2", + Set.of(DataStreamAutoShardingService.DATA_STREAM_AUTO_SHARDING_FEATURE.id()) + ) + ) + .metadata(builder) + .build(); + + AutoShardingResult autoShardingResult = service.calculate(state, dataStream, 1.0); + assertThat(autoShardingResult.type(), is(DECREASE_SHARDS)); + assertThat(autoShardingResult.targetNumberOfShards(), is(1)); + // no pre-existing auto sharding event however we have old enough backing indices (older than the cooldown period) so we can + // make a decision to reduce the number of shards + assertThat(autoShardingResult.coolDownRemaining(), is(TimeValue.ZERO)); + } + + { + // let's test a decrease in number of shards after a previous decrease event + Metadata.Builder builder = Metadata.builder(); + Function dataStreamSupplier = (autoShardingEvent) -> createDataStream( + builder, + dataStreamName, + 3, + now, + List.of( + now - TimeValue.timeValueDays(21).getMillis(), + now - TimeValue.timeValueDays(15).getMillis(), // triggers auto sharding event + now - TimeValue.timeValueDays(4).getMillis(), + now - TimeValue.timeValueDays(2).getMillis(), + now - 1000 + ), + getWriteLoad(3, 0.333), + autoShardingEvent + ); + + // generation 2 triggered a decrease in shards event to 2 shards + DataStream dataStream = dataStreamSupplier.apply( + new DataStreamAutoShardingEvent( + DataStream.getDefaultBackingIndexName(dataStreamName, 2), + 2, + now - TimeValue.timeValueDays(4).getMillis() + ) + ); + builder.put(dataStream); + ClusterState state = ClusterState.builder(ClusterName.DEFAULT) + .nodeFeatures( + Map.of( + "n1", + Set.of(DataStreamAutoShardingService.DATA_STREAM_AUTO_SHARDING_FEATURE.id()), + "n2", + Set.of(DataStreamAutoShardingService.DATA_STREAM_AUTO_SHARDING_FEATURE.id()) + ) + ) + .metadata(builder) + .build(); + + AutoShardingResult autoShardingResult = service.calculate(state, dataStream, 1.0); + assertThat(autoShardingResult.type(), is(DECREASE_SHARDS)); + assertThat(autoShardingResult.targetNumberOfShards(), is(1)); + assertThat(autoShardingResult.coolDownRemaining(), is(TimeValue.ZERO)); + } + + { + // let's test a decrease in number of shards that's prevented by the cool down period due to a previous sharding event + // the expected result type here is COOLDOWN_PREVENTED_DECREASE + Metadata.Builder builder = Metadata.builder(); + Function dataStreamSupplier = (autoShardingEvent) -> createDataStream( + builder, + dataStreamName, + 3, + now, + List.of( + now - TimeValue.timeValueDays(21).getMillis(), + now - TimeValue.timeValueDays(2).getMillis(), // triggers auto sharding event + now - TimeValue.timeValueDays(1).getMillis(), + now - 1000 + ), + getWriteLoad(3, 0.25), + autoShardingEvent + ); + + // generation 2 triggered a decrease in shards event to 2 shards + DataStream dataStream = dataStreamSupplier.apply( + new DataStreamAutoShardingEvent( + DataStream.getDefaultBackingIndexName(dataStreamName, 2), + 2, + now - TimeValue.timeValueDays(2).getMillis() + ) + ); + builder.put(dataStream); + ClusterState state = ClusterState.builder(ClusterName.DEFAULT) + .nodeFeatures( + Map.of( + "n1", + Set.of(DataStreamAutoShardingService.DATA_STREAM_AUTO_SHARDING_FEATURE.id()), + "n2", + Set.of(DataStreamAutoShardingService.DATA_STREAM_AUTO_SHARDING_FEATURE.id()) + ) + ) + .metadata(builder) + .build(); + + AutoShardingResult autoShardingResult = service.calculate(state, dataStream, 1.0); + assertThat(autoShardingResult.type(), is(COOLDOWN_PREVENTED_DECREASE)); + assertThat(autoShardingResult.targetNumberOfShards(), is(1)); + assertThat(autoShardingResult.coolDownRemaining(), is(TimeValue.timeValueDays(1))); + } + + { + // no change required + Metadata.Builder builder = Metadata.builder(); + Function dataStreamSupplier = (autoShardingEvent) -> createDataStream( + builder, + dataStreamName, + 3, + now, + List.of( + now - TimeValue.timeValueDays(21).getMillis(), + now - TimeValue.timeValueDays(15).getMillis(), + now - TimeValue.timeValueDays(4).getMillis(), + now - TimeValue.timeValueDays(2).getMillis(), + now - 1000 + ), + getWriteLoad(3, 1.333), + autoShardingEvent + ); + + // generation 2 triggered a decrease in shards event to 2 shards + DataStream dataStream = dataStreamSupplier.apply(null); + builder.put(dataStream); + ClusterState state = ClusterState.builder(ClusterName.DEFAULT) + .nodeFeatures( + Map.of( + "n1", + Set.of(DataStreamAutoShardingService.DATA_STREAM_AUTO_SHARDING_FEATURE.id()), + "n2", + Set.of(DataStreamAutoShardingService.DATA_STREAM_AUTO_SHARDING_FEATURE.id()) + ) + ) + .metadata(builder) + .build(); + + AutoShardingResult autoShardingResult = service.calculate(state, dataStream, 4.0); + assertThat(autoShardingResult.type(), is(NO_CHANGE_REQUIRED)); + assertThat(autoShardingResult.targetNumberOfShards(), is(3)); + assertThat(autoShardingResult.coolDownRemaining(), is(TimeValue.ZERO)); + } + } + + public void testComputeOptimalNumberOfShards() { + int minWriteThreads = 2; + int maxWriteThreads = 32; + { + // the small values will be very common so let's randomise to make sure we never go below 1L + double indexingLoad = randomDoubleBetween(0.0001, 1.0, true); + logger.info("-> indexingLoad {}", indexingLoad); + assertThat(DataStreamAutoShardingService.computeOptimalNumberOfShards(minWriteThreads, maxWriteThreads, indexingLoad), is(1L)); + } + + { + double indexingLoad = 2.0; + logger.info("-> indexingLoad {}", indexingLoad); + assertThat(DataStreamAutoShardingService.computeOptimalNumberOfShards(minWriteThreads, maxWriteThreads, indexingLoad), is(2L)); + } + + { + // there's a broad range of popular values (a write index starting to be very busy, using between 3 and all of the 32 write + // threads, so let's randomise this too to make sure we stay at 3 recommended shards) + double indexingLoad = randomDoubleBetween(3.0002, 32.0, true); + logger.info("-> indexingLoad {}", indexingLoad); + + assertThat(DataStreamAutoShardingService.computeOptimalNumberOfShards(minWriteThreads, maxWriteThreads, indexingLoad), is(3L)); + } + + { + double indexingLoad = 49.0; + logger.info("-> indexingLoad {}", indexingLoad); + assertThat(DataStreamAutoShardingService.computeOptimalNumberOfShards(minWriteThreads, maxWriteThreads, indexingLoad), is(4L)); + } + + { + double indexingLoad = 70.0; + logger.info("-> indexingLoad {}", indexingLoad); + assertThat(DataStreamAutoShardingService.computeOptimalNumberOfShards(minWriteThreads, maxWriteThreads, indexingLoad), is(5L)); + } + + { + double indexingLoad = 100.0; + logger.info("-> indexingLoad {}", indexingLoad); + assertThat(DataStreamAutoShardingService.computeOptimalNumberOfShards(minWriteThreads, maxWriteThreads, indexingLoad), is(7L)); + } + + { + double indexingLoad = 180.0; + logger.info("-> indexingLoad {}", indexingLoad); + assertThat(DataStreamAutoShardingService.computeOptimalNumberOfShards(minWriteThreads, maxWriteThreads, indexingLoad), is(12L)); + } + } + + public void testGetMaxIndexLoadWithinCoolingPeriod() { + final TimeValue coolingPeriod = TimeValue.timeValueDays(3); + + final Metadata.Builder metadataBuilder = Metadata.builder(); + final int numberOfBackingIndicesOutsideCoolingPeriod = randomIntBetween(3, 10); + final int numberOfBackingIndicesWithinCoolingPeriod = randomIntBetween(3, 10); + final List backingIndices = new ArrayList<>(); + final String dataStreamName = "logs"; + long now = System.currentTimeMillis(); + + // to cover the entire cooling period we'll also include the backing index right before the index age calculation + // this flag makes that index have a very low or very high write load + boolean lastIndexBeforeCoolingPeriodHasLowWriteLoad = randomBoolean(); + for (int i = 0; i < numberOfBackingIndicesOutsideCoolingPeriod; i++) { + long creationDate = now - (coolingPeriod.millis() * 2); + IndexMetadata indexMetadata = createIndexMetadata( + DataStream.getDefaultBackingIndexName(dataStreamName, backingIndices.size(), creationDate), + 1, + getWriteLoad(1, 999.0), + creationDate + ); + + if (lastIndexBeforeCoolingPeriodHasLowWriteLoad) { + indexMetadata = createIndexMetadata( + DataStream.getDefaultBackingIndexName(dataStreamName, backingIndices.size(), creationDate), + 1, + getWriteLoad(1, 1.0), + creationDate + ); + } + backingIndices.add(indexMetadata.getIndex()); + metadataBuilder.put(indexMetadata, false); + } + + for (int i = 0; i < numberOfBackingIndicesWithinCoolingPeriod; i++) { + final long createdAt = now - (coolingPeriod.getMillis() / 2); + IndexMetadata indexMetadata; + if (i == numberOfBackingIndicesWithinCoolingPeriod - 1) { + indexMetadata = createIndexMetadata( + DataStream.getDefaultBackingIndexName(dataStreamName, backingIndices.size(), createdAt), + 3, + getWriteLoad(3, 5.0), // max write index within cooling period + createdAt + ); + } else { + indexMetadata = createIndexMetadata( + DataStream.getDefaultBackingIndexName(dataStreamName, backingIndices.size(), createdAt), + 3, + getWriteLoad(3, 3.0), // each backing index has a write load of 9.0 + createdAt + ); + } + backingIndices.add(indexMetadata.getIndex()); + metadataBuilder.put(indexMetadata, false); + } + + final String writeIndexName = DataStream.getDefaultBackingIndexName(dataStreamName, backingIndices.size()); + final IndexMetadata writeIndexMetadata = createIndexMetadata(writeIndexName, 3, getWriteLoad(3, 1.0), System.currentTimeMillis()); + backingIndices.add(writeIndexMetadata.getIndex()); + metadataBuilder.put(writeIndexMetadata, false); + + final DataStream dataStream = new DataStream( + dataStreamName, + backingIndices, + backingIndices.size(), + Collections.emptyMap(), + false, + false, + false, + false, + IndexMode.STANDARD + ); + + metadataBuilder.put(dataStream); + + double maxIndexLoadWithinCoolingPeriod = DataStreamAutoShardingService.getMaxIndexLoadWithinCoolingPeriod( + metadataBuilder.build(), + dataStream, + 3.0, + coolingPeriod, + () -> now + ); + // to cover the entire cooldown period, the last index before the cooling period is taken into account + assertThat(maxIndexLoadWithinCoolingPeriod, is(lastIndexBeforeCoolingPeriodHasLowWriteLoad ? 15.0 : 999.0)); + } + + public void testAutoShardingResultValidation() { + { + // throws exception when constructed using types that shouldn't report cooldowns + expectThrows( + IllegalArgumentException.class, + () -> new AutoShardingResult(INCREASE_SHARDS, 1, 3, TimeValue.timeValueSeconds(3), 3.0) + ); + + expectThrows( + IllegalArgumentException.class, + () -> new AutoShardingResult(DECREASE_SHARDS, 3, 1, TimeValue.timeValueSeconds(3), 1.0) + ); + + } + + { + // we can successfully create results with cooldown period for the designated types + AutoShardingResult cooldownPreventedIncrease = new AutoShardingResult( + COOLDOWN_PREVENTED_INCREASE, + 1, + 3, + TimeValue.timeValueSeconds(3), + 3.0 + ); + assertThat(cooldownPreventedIncrease.coolDownRemaining(), is(TimeValue.timeValueSeconds(3))); + + AutoShardingResult cooldownPreventedDecrease = new AutoShardingResult( + COOLDOWN_PREVENTED_DECREASE, + 3, + 1, + TimeValue.timeValueSeconds(7), + 1.0 + ); + assertThat(cooldownPreventedDecrease.coolDownRemaining(), is(TimeValue.timeValueSeconds(7))); + } + } + + private DataStream createDataStream( + Metadata.Builder builder, + String dataStreamName, + int numberOfShards, + Long now, + List indicesCreationDate, + IndexWriteLoad backingIndicesWriteLoad, + @Nullable DataStreamAutoShardingEvent autoShardingEvent + ) { + final List backingIndices = new ArrayList<>(); + int backingIndicesCount = indicesCreationDate.size(); + for (int k = 0; k < indicesCreationDate.size(); k++) { + long createdAt = indicesCreationDate.get(k); + IndexMetadata.Builder indexMetaBuilder; + if (k < backingIndicesCount - 1) { + indexMetaBuilder = IndexMetadata.builder( + createIndexMetadata( + DataStream.getDefaultBackingIndexName(dataStreamName, k + 1), + numberOfShards, + backingIndicesWriteLoad, + createdAt + ) + ); + // add rollover info only for non-write indices + MaxAgeCondition rolloverCondition = new MaxAgeCondition(TimeValue.timeValueMillis(now - 2000L)); + indexMetaBuilder.putRolloverInfo(new RolloverInfo(dataStreamName, List.of(rolloverCondition), now - 2000L)); + } else { + // write index + indexMetaBuilder = IndexMetadata.builder( + createIndexMetadata(DataStream.getDefaultBackingIndexName(dataStreamName, k + 1), numberOfShards, null, createdAt) + ); + } + IndexMetadata indexMetadata = indexMetaBuilder.build(); + builder.put(indexMetadata, false); + backingIndices.add(indexMetadata.getIndex()); + } + return new DataStream( + dataStreamName, + backingIndices, + backingIndicesCount, + null, + false, + false, + false, + false, + null, + null, + false, + List.of(), + autoShardingEvent + ); + } + + private IndexMetadata createIndexMetadata( + String indexName, + int numberOfShards, + @Nullable IndexWriteLoad indexWriteLoad, + long createdAt + ) { + return IndexMetadata.builder(indexName) + .settings( + Settings.builder() + .put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, numberOfShards) + .put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 0) + .put(IndexMetadata.SETTING_VERSION_CREATED, IndexVersion.current()) + .build() + ) + .stats(indexWriteLoad == null ? null : new IndexMetadataStats(indexWriteLoad, 1, 1)) + .creationDate(createdAt) + .build(); + } + + private IndexWriteLoad getWriteLoad(int numberOfShards, double shardWriteLoad) { + IndexWriteLoad.Builder builder = IndexWriteLoad.builder(numberOfShards); + for (int shardId = 0; shardId < numberOfShards; shardId++) { + builder.withShardWriteLoad(shardId, shardWriteLoad, randomLongBetween(1, 10)); + } + return builder.build(); + } + +} diff --git a/server/src/test/java/org/elasticsearch/cluster/metadata/DataStreamAutoShardingEventTests.java b/server/src/test/java/org/elasticsearch/cluster/metadata/DataStreamAutoShardingEventTests.java new file mode 100644 index 0000000000000..925c204fa5b27 --- /dev/null +++ b/server/src/test/java/org/elasticsearch/cluster/metadata/DataStreamAutoShardingEventTests.java @@ -0,0 +1,62 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0 and the Server Side Public License, v 1; you may not use this file except + * in compliance with, at your election, the Elastic License 2.0 or the Server + * Side Public License, v 1. + */ + +package org.elasticsearch.cluster.metadata; + +import org.elasticsearch.cluster.Diff; +import org.elasticsearch.common.io.stream.Writeable; +import org.elasticsearch.test.ESTestCase; +import org.elasticsearch.test.SimpleDiffableSerializationTestCase; +import org.elasticsearch.xcontent.XContentParser; + +import java.io.IOException; + +public class DataStreamAutoShardingEventTests extends SimpleDiffableSerializationTestCase { + + @Override + protected DataStreamAutoShardingEvent doParseInstance(XContentParser parser) throws IOException { + return DataStreamAutoShardingEvent.fromXContent(parser); + } + + @Override + protected Writeable.Reader instanceReader() { + return DataStreamAutoShardingEvent::new; + } + + @Override + protected DataStreamAutoShardingEvent createTestInstance() { + return DataStreamAutoShardingEventTests.randomInstance(); + } + + @Override + protected DataStreamAutoShardingEvent mutateInstance(DataStreamAutoShardingEvent instance) { + String triggerIndex = instance.triggerIndexName(); + long timestamp = instance.timestamp(); + int targetNumberOfShards = instance.targetNumberOfShards(); + switch (randomInt(2)) { + case 0 -> triggerIndex = randomValueOtherThan(triggerIndex, () -> randomAlphaOfLengthBetween(10, 50)); + case 1 -> timestamp = randomValueOtherThan(timestamp, ESTestCase::randomNonNegativeLong); + case 2 -> targetNumberOfShards = randomValueOtherThan(targetNumberOfShards, ESTestCase::randomNonNegativeInt); + } + return new DataStreamAutoShardingEvent(triggerIndex, targetNumberOfShards, timestamp); + } + + static DataStreamAutoShardingEvent randomInstance() { + return new DataStreamAutoShardingEvent(randomAlphaOfLengthBetween(10, 40), randomNonNegativeInt(), randomNonNegativeLong()); + } + + @Override + protected DataStreamAutoShardingEvent makeTestChanges(DataStreamAutoShardingEvent testInstance) { + return mutateInstance(testInstance); + } + + @Override + protected Writeable.Reader> diffReader() { + return DataStreamAutoShardingEvent::readDiffFrom; + } +} diff --git a/server/src/test/java/org/elasticsearch/cluster/metadata/DataStreamTests.java b/server/src/test/java/org/elasticsearch/cluster/metadata/DataStreamTests.java index 9f7d6b49b0844..7e8e9805b54e7 100644 --- a/server/src/test/java/org/elasticsearch/cluster/metadata/DataStreamTests.java +++ b/server/src/test/java/org/elasticsearch/cluster/metadata/DataStreamTests.java @@ -39,6 +39,7 @@ import java.time.temporal.ChronoUnit; import java.util.ArrayList; import java.util.Arrays; +import java.util.Collections; import java.util.HashMap; import java.util.HashSet; import java.util.List; @@ -93,7 +94,8 @@ protected DataStream mutateInstance(DataStream instance) { var lifecycle = instance.getLifecycle(); var failureStore = instance.isFailureStore(); var failureIndices = instance.getFailureIndices(); - switch (between(0, 10)) { + var autoShardingEvent = instance.getAutoShardingEvent(); + switch (between(0, 11)) { case 0 -> name = randomAlphaOfLength(10); case 1 -> indices = randomValueOtherThan(List.of(), DataStreamTestHelper::randomIndexInstances); case 2 -> generation = instance.getGeneration() + randomIntBetween(1, 10); @@ -130,6 +132,15 @@ protected DataStream mutateInstance(DataStream instance) { failureStore = true; } } + case 11 -> { + autoShardingEvent = randomBoolean() && autoShardingEvent != null + ? null + : new DataStreamAutoShardingEvent( + indices.get(indices.size() - 1).getName(), + randomIntBetween(1, 10), + randomMillisUpToYear9999() + ); + } } return new DataStream( @@ -144,7 +155,8 @@ protected DataStream mutateInstance(DataStream instance) { indexMode, lifecycle, failureStore, - failureIndices + failureIndices, + autoShardingEvent ); } @@ -201,7 +213,8 @@ public void testRolloverUpgradeToTsdbDataStream() { indexMode, ds.getLifecycle(), ds.isFailureStore(), - ds.getFailureIndices() + ds.getFailureIndices(), + ds.getAutoShardingEvent() ); var newCoordinates = ds.nextWriteIndexAndGeneration(Metadata.EMPTY_METADATA); @@ -228,7 +241,8 @@ public void testRolloverDowngradeToRegularDataStream() { IndexMode.TIME_SERIES, ds.getLifecycle(), ds.isFailureStore(), - ds.getFailureIndices() + ds.getFailureIndices(), + ds.getAutoShardingEvent() ); var newCoordinates = ds.nextWriteIndexAndGeneration(Metadata.EMPTY_METADATA); @@ -590,7 +604,8 @@ public void testSnapshot() { preSnapshotDataStream.getIndexMode(), preSnapshotDataStream.getLifecycle(), preSnapshotDataStream.isFailureStore(), - preSnapshotDataStream.getFailureIndices() + preSnapshotDataStream.getFailureIndices(), + preSnapshotDataStream.getAutoShardingEvent() ); var reconciledDataStream = postSnapshotDataStream.snapshot( @@ -634,7 +649,8 @@ public void testSnapshotWithAllBackingIndicesRemoved() { preSnapshotDataStream.getIndexMode(), preSnapshotDataStream.getLifecycle(), preSnapshotDataStream.isFailureStore(), - preSnapshotDataStream.getFailureIndices() + preSnapshotDataStream.getFailureIndices(), + preSnapshotDataStream.getAutoShardingEvent() ); assertNull(postSnapshotDataStream.snapshot(preSnapshotDataStream.getIndices().stream().map(Index::getName).toList())); @@ -1654,7 +1670,8 @@ public void testXContentSerializationWithRollover() throws IOException { lifecycle, failureStore, failureIndices, - false + false, + null ); try (XContentBuilder builder = XContentBuilder.builder(XContentType.JSON.xContent())) { @@ -1671,6 +1688,99 @@ public void testXContentSerializationWithRollover() throws IOException { } } + public void testGetIndicesWithinMaxAgeRange() { + final TimeValue maxIndexAge = TimeValue.timeValueDays(7); + + final Metadata.Builder metadataBuilder = Metadata.builder(); + final int numberOfBackingIndicesOlderThanMinAge = randomIntBetween(0, 10); + final int numberOfBackingIndicesWithinMinAnge = randomIntBetween(0, 10); + final int numberOfShards = 1; + final List backingIndices = new ArrayList<>(); + final String dataStreamName = "logs-es"; + final List backingIndicesOlderThanMinAge = new ArrayList<>(); + for (int i = 0; i < numberOfBackingIndicesOlderThanMinAge; i++) { + long creationDate = System.currentTimeMillis() - maxIndexAge.millis() * 2; + final IndexMetadata indexMetadata = createIndexMetadata( + DataStream.getDefaultBackingIndexName(dataStreamName, backingIndices.size(), creationDate), + randomIndexWriteLoad(numberOfShards), + creationDate + ); + backingIndices.add(indexMetadata.getIndex()); + backingIndicesOlderThanMinAge.add(indexMetadata.getIndex()); + metadataBuilder.put(indexMetadata, false); + } + + final List backingIndicesWithinMinAge = new ArrayList<>(); + for (int i = 0; i < numberOfBackingIndicesWithinMinAnge; i++) { + final long createdAt = System.currentTimeMillis() - (maxIndexAge.getMillis() / 2); + final IndexMetadata indexMetadata = createIndexMetadata( + DataStream.getDefaultBackingIndexName(dataStreamName, backingIndices.size(), createdAt), + randomIndexWriteLoad(numberOfShards), + createdAt + ); + backingIndices.add(indexMetadata.getIndex()); + backingIndicesWithinMinAge.add(indexMetadata.getIndex()); + metadataBuilder.put(indexMetadata, false); + } + + final String writeIndexName = DataStream.getDefaultBackingIndexName(dataStreamName, backingIndices.size()); + final IndexMetadata writeIndexMetadata = createIndexMetadata(writeIndexName, null, System.currentTimeMillis()); + backingIndices.add(writeIndexMetadata.getIndex()); + metadataBuilder.put(writeIndexMetadata, false); + + final DataStream dataStream = new DataStream( + dataStreamName, + backingIndices, + backingIndices.size(), + Collections.emptyMap(), + false, + false, + false, + false, + randomBoolean() ? IndexMode.STANDARD : IndexMode.TIME_SERIES + ); + + metadataBuilder.put(dataStream); + + final List indicesWithinMaxAgeRange = DataStream.getIndicesWithinMaxAgeRange( + dataStream, + metadataBuilder::getSafe, + maxIndexAge, + System::currentTimeMillis + ); + + final List expectedIndicesWithinMaxAgeRange = new ArrayList<>(); + if (numberOfBackingIndicesOlderThanMinAge > 0) { + expectedIndicesWithinMaxAgeRange.add(backingIndicesOlderThanMinAge.get(backingIndicesOlderThanMinAge.size() - 1)); + } + expectedIndicesWithinMaxAgeRange.addAll(backingIndicesWithinMinAge); + expectedIndicesWithinMaxAgeRange.add(writeIndexMetadata.getIndex()); + + assertThat(indicesWithinMaxAgeRange, is(equalTo(expectedIndicesWithinMaxAgeRange))); + } + + private IndexWriteLoad randomIndexWriteLoad(int numberOfShards) { + IndexWriteLoad.Builder builder = IndexWriteLoad.builder(numberOfShards); + for (int shardId = 0; shardId < numberOfShards; shardId++) { + builder.withShardWriteLoad(shardId, randomDoubleBetween(0, 64, true), randomLongBetween(1, 10)); + } + return builder.build(); + } + + private IndexMetadata createIndexMetadata(String indexName, IndexWriteLoad indexWriteLoad, long createdAt) { + return IndexMetadata.builder(indexName) + .settings( + Settings.builder() + .put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1) + .put(IndexMetadata.SETTING_NUMBER_OF_REPLICAS, 0) + .put(IndexMetadata.SETTING_VERSION_CREATED, IndexVersion.current()) + .build() + ) + .stats(indexWriteLoad == null ? null : new IndexMetadataStats(indexWriteLoad, 1, 1)) + .creationDate(createdAt) + .build(); + } + private record DataStreamMetadata(Long creationTimeInMillis, Long rolloverTimeInMillis, Long originationTimeInMillis) { public static DataStreamMetadata dataStreamMetadata(Long creationTimeInMillis, Long rolloverTimeInMillis) { return new DataStreamMetadata(creationTimeInMillis, rolloverTimeInMillis, null); diff --git a/server/src/test/java/org/elasticsearch/cluster/metadata/MetadataDataStreamsServiceTests.java b/server/src/test/java/org/elasticsearch/cluster/metadata/MetadataDataStreamsServiceTests.java index ba3b1a7387110..71306d7fe0aef 100644 --- a/server/src/test/java/org/elasticsearch/cluster/metadata/MetadataDataStreamsServiceTests.java +++ b/server/src/test/java/org/elasticsearch/cluster/metadata/MetadataDataStreamsServiceTests.java @@ -356,7 +356,8 @@ public void testRemoveBrokenBackingIndexReference() { original.getIndexMode(), original.getLifecycle(), original.isFailureStore(), - original.getFailureIndices() + original.getFailureIndices(), + original.getAutoShardingEvent() ); var brokenState = ClusterState.builder(state).metadata(Metadata.builder(state.getMetadata()).put(broken).build()).build(); diff --git a/test/framework/src/main/java/org/elasticsearch/cluster/metadata/DataStreamTestHelper.java b/test/framework/src/main/java/org/elasticsearch/cluster/metadata/DataStreamTestHelper.java index 3a47e0885f2d2..8fc02bb8e808c 100644 --- a/test/framework/src/main/java/org/elasticsearch/cluster/metadata/DataStreamTestHelper.java +++ b/test/framework/src/main/java/org/elasticsearch/cluster/metadata/DataStreamTestHelper.java @@ -74,6 +74,7 @@ import static org.elasticsearch.test.ESTestCase.randomAlphaOfLength; import static org.elasticsearch.test.ESTestCase.randomBoolean; import static org.elasticsearch.test.ESTestCase.randomFrom; +import static org.elasticsearch.test.ESTestCase.randomIntBetween; import static org.elasticsearch.test.ESTestCase.randomMap; import static org.elasticsearch.test.ESTestCase.randomMillisUpToYear9999; import static org.mockito.ArgumentMatchers.any; @@ -136,7 +137,8 @@ public static DataStream newInstance( null, lifecycle, failureStores.size() > 0, - failureStores + failureStores, + null ); } @@ -307,7 +309,14 @@ public static DataStream randomInstance(String dataStreamName, LongSupplier time randomBoolean() ? DataStreamLifecycle.newBuilder().dataRetention(randomMillisUpToYear9999()).build() : null, failureStore, failureIndices, + randomBoolean(), randomBoolean() + ? new DataStreamAutoShardingEvent( + indices.get(indices.size() - 1).getName(), + randomIntBetween(1, 10), + randomMillisUpToYear9999() + ) + : null ); } diff --git a/x-pack/plugin/ccr/src/main/java/org/elasticsearch/xpack/ccr/action/TransportPutFollowAction.java b/x-pack/plugin/ccr/src/main/java/org/elasticsearch/xpack/ccr/action/TransportPutFollowAction.java index b06ff73e29960..c3dd30bd2f242 100644 --- a/x-pack/plugin/ccr/src/main/java/org/elasticsearch/xpack/ccr/action/TransportPutFollowAction.java +++ b/x-pack/plugin/ccr/src/main/java/org/elasticsearch/xpack/ccr/action/TransportPutFollowAction.java @@ -334,7 +334,8 @@ static DataStream updateLocalDataStream( remoteDataStream.getIndexMode(), remoteDataStream.getLifecycle(), remoteDataStream.isFailureStore(), - remoteDataStream.getFailureIndices() + remoteDataStream.getFailureIndices(), + remoteDataStream.getAutoShardingEvent() ); } else { if (localDataStream.isReplicated() == false) { @@ -387,7 +388,8 @@ static DataStream updateLocalDataStream( localDataStream.getIndexMode(), localDataStream.getLifecycle(), localDataStream.isFailureStore(), - localDataStream.getFailureIndices() + localDataStream.getFailureIndices(), + localDataStream.getAutoShardingEvent() ); } } diff --git a/x-pack/plugin/core/src/internalClusterTest/java/org/elasticsearch/xpack/core/action/DataStreamLifecycleUsageTransportActionIT.java b/x-pack/plugin/core/src/internalClusterTest/java/org/elasticsearch/xpack/core/action/DataStreamLifecycleUsageTransportActionIT.java index c102470628a00..bc97623c76970 100644 --- a/x-pack/plugin/core/src/internalClusterTest/java/org/elasticsearch/xpack/core/action/DataStreamLifecycleUsageTransportActionIT.java +++ b/x-pack/plugin/core/src/internalClusterTest/java/org/elasticsearch/xpack/core/action/DataStreamLifecycleUsageTransportActionIT.java @@ -134,7 +134,8 @@ public void testAction() throws Exception { IndexMode.STANDARD, lifecycle, false, - List.of() + List.of(), + null ); dataStreamMap.put(dataStream.getName(), dataStream); } diff --git a/x-pack/plugin/write-load-forecaster/src/main/java/org/elasticsearch/xpack/writeloadforecaster/LicensedWriteLoadForecaster.java b/x-pack/plugin/write-load-forecaster/src/main/java/org/elasticsearch/xpack/writeloadforecaster/LicensedWriteLoadForecaster.java index c1126df228cfe..d4a85ce859b2b 100644 --- a/x-pack/plugin/write-load-forecaster/src/main/java/org/elasticsearch/xpack/writeloadforecaster/LicensedWriteLoadForecaster.java +++ b/x-pack/plugin/write-load-forecaster/src/main/java/org/elasticsearch/xpack/writeloadforecaster/LicensedWriteLoadForecaster.java @@ -76,7 +76,13 @@ public Metadata.Builder withWriteLoadForecastForWriteIndex(String dataStreamName clearPreviousForecast(dataStream, metadata); - final List indicesWriteLoadWithinMaxAgeRange = getIndicesWithinMaxAgeRange(dataStream, metadata).stream() + final List indicesWriteLoadWithinMaxAgeRange = DataStream.getIndicesWithinMaxAgeRange( + dataStream, + metadata::getSafe, + maxIndexAge, + threadPool::absoluteTimeInMillis + ) + .stream() .filter(index -> index.equals(dataStream.getWriteIndex()) == false) .map(metadata::getSafe) .map(IndexMetadata::getStats) @@ -134,25 +140,6 @@ static OptionalDouble forecastIndexWriteLoad(List indicesWriteLo return totalShardUptime == 0 ? OptionalDouble.empty() : OptionalDouble.of(totalWeightedWriteLoad / totalShardUptime); } - // Visible for testing - List getIndicesWithinMaxAgeRange(DataStream dataStream, Metadata.Builder metadata) { - final List dataStreamIndices = dataStream.getIndices(); - final long currentTimeMillis = threadPool.absoluteTimeInMillis(); - // Consider at least 1 index (including the write index) for cases where rollovers happen less often than maxIndexAge - int firstIndexWithinAgeRange = Math.max(dataStreamIndices.size() - 2, 0); - for (int i = 0; i < dataStreamIndices.size(); i++) { - Index index = dataStreamIndices.get(i); - final IndexMetadata indexMetadata = metadata.getSafe(index); - final long indexAge = currentTimeMillis - indexMetadata.getCreationDate(); - if (indexAge < maxIndexAge.getMillis()) { - // We need to consider the previous index too in order to cover the entire max-index-age range. - firstIndexWithinAgeRange = i == 0 ? 0 : i - 1; - break; - } - } - return dataStreamIndices.subList(firstIndexWithinAgeRange, dataStreamIndices.size()); - } - @Override @SuppressForbidden(reason = "This is the only place where IndexMetadata#getForecastedWriteLoad is allowed to be used") public OptionalDouble getForecastedWriteLoad(IndexMetadata indexMetadata) { diff --git a/x-pack/plugin/write-load-forecaster/src/test/java/org/elasticsearch/xpack/writeloadforecaster/LicensedWriteLoadForecasterTests.java b/x-pack/plugin/write-load-forecaster/src/test/java/org/elasticsearch/xpack/writeloadforecaster/LicensedWriteLoadForecasterTests.java index 38e754c802983..c7efb27509ef7 100644 --- a/x-pack/plugin/write-load-forecaster/src/test/java/org/elasticsearch/xpack/writeloadforecaster/LicensedWriteLoadForecasterTests.java +++ b/x-pack/plugin/write-load-forecaster/src/test/java/org/elasticsearch/xpack/writeloadforecaster/LicensedWriteLoadForecasterTests.java @@ -287,65 +287,6 @@ public void testWriteLoadForecast() { } } - public void testGetIndicesWithinMaxAgeRange() { - final TimeValue maxIndexAge = TimeValue.timeValueDays(7); - final LicensedWriteLoadForecaster writeLoadForecaster = new LicensedWriteLoadForecaster(() -> true, threadPool, maxIndexAge); - - final Metadata.Builder metadataBuilder = Metadata.builder(); - final int numberOfBackingIndicesOlderThanMinAge = randomIntBetween(0, 10); - final int numberOfBackingIndicesWithinMinAnge = randomIntBetween(0, 10); - final int numberOfShards = 1; - final List backingIndices = new ArrayList<>(); - final String dataStreamName = "logs-es"; - final List backingIndicesOlderThanMinAge = new ArrayList<>(); - for (int i = 0; i < numberOfBackingIndicesOlderThanMinAge; i++) { - long creationDate = System.currentTimeMillis() - maxIndexAge.millis() * 2; - final IndexMetadata indexMetadata = createIndexMetadata( - DataStream.getDefaultBackingIndexName(dataStreamName, backingIndices.size(), creationDate), - numberOfShards, - randomIndexWriteLoad(numberOfShards), - creationDate - ); - backingIndices.add(indexMetadata.getIndex()); - backingIndicesOlderThanMinAge.add(indexMetadata.getIndex()); - metadataBuilder.put(indexMetadata, false); - } - - final List backingIndicesWithinMinAge = new ArrayList<>(); - for (int i = 0; i < numberOfBackingIndicesWithinMinAnge; i++) { - final long createdAt = System.currentTimeMillis() - (maxIndexAge.getMillis() / 2); - final IndexMetadata indexMetadata = createIndexMetadata( - DataStream.getDefaultBackingIndexName(dataStreamName, backingIndices.size(), createdAt), - numberOfShards, - randomIndexWriteLoad(numberOfShards), - createdAt - ); - backingIndices.add(indexMetadata.getIndex()); - backingIndicesWithinMinAge.add(indexMetadata.getIndex()); - metadataBuilder.put(indexMetadata, false); - } - - final String writeIndexName = DataStream.getDefaultBackingIndexName(dataStreamName, backingIndices.size()); - final IndexMetadata writeIndexMetadata = createIndexMetadata(writeIndexName, numberOfShards, null, System.currentTimeMillis()); - backingIndices.add(writeIndexMetadata.getIndex()); - metadataBuilder.put(writeIndexMetadata, false); - - final DataStream dataStream = createDataStream(dataStreamName, backingIndices); - - metadataBuilder.put(dataStream); - - final List indicesWithinMaxAgeRange = writeLoadForecaster.getIndicesWithinMaxAgeRange(dataStream, metadataBuilder); - - final List expectedIndicesWithinMaxAgeRange = new ArrayList<>(); - if (numberOfBackingIndicesOlderThanMinAge > 0) { - expectedIndicesWithinMaxAgeRange.add(backingIndicesOlderThanMinAge.get(backingIndicesOlderThanMinAge.size() - 1)); - } - expectedIndicesWithinMaxAgeRange.addAll(backingIndicesWithinMinAge); - expectedIndicesWithinMaxAgeRange.add(writeIndexMetadata.getIndex()); - - assertThat(indicesWithinMaxAgeRange, is(equalTo(expectedIndicesWithinMaxAgeRange))); - } - private IndexWriteLoad randomIndexWriteLoad(int numberOfShards) { IndexWriteLoad.Builder builder = IndexWriteLoad.builder(numberOfShards); for (int shardId = 0; shardId < numberOfShards; shardId++) { From 8d93a934f6924b3359d08a7cf81db6f732a57ce6 Mon Sep 17 00:00:00 2001 From: Craig Taverner Date: Wed, 6 Mar 2024 11:19:54 +0100 Subject: [PATCH 18/27] Add two new OGC functions ST_X and ST_Y (#105768) * Add two new OGC functions ST_X and ST_Y Recently Nik did work that involved extracting the X and Y coordinates from geo_point data using `to_string(field)` followed by a DISSECT command to re-parse the string to get the X and Y coordinates. This is much more efficiently achieved using existing known OGC functions `ST_X` and `ST_Y`. * Update docs/changelog/105768.yaml * Fixed invalid changelog yaml * Fixed mixed cluster tests * Fixed tests and added docs * Removed false impression that these functions were different for geo/cartesian With the use of WKB as the core type in the compute engine, many spatial functions are actually the same between these two types, so we should not give the impression they are different. * Code review comments and reduced object creation. * Revert temporary StringUtils hack, and fix bug in x/y extraction from WKB * Revert object creation reduction * Fixed mistakes in documentation --- docs/changelog/105768.yaml | 5 + .../esql/esql-functions-operators.asciidoc | 4 + .../esql/functions/signature/st_x.svg | 1 + .../esql/functions/signature/st_y.svg | 1 + .../esql/functions/spatial-functions.asciidoc | 16 +++ docs/reference/esql/functions/st_x.asciidoc | 33 +++++ docs/reference/esql/functions/st_y.asciidoc | 33 +++++ .../esql/functions/types/st_x.asciidoc | 6 + .../esql/functions/types/st_y.asciidoc | 6 + .../src/main/resources/show.csv-spec | 10 +- .../src/main/resources/spatial.csv-spec | 44 ++++++ .../scalar/spatial/StXFromWKBEvaluator.java | 127 ++++++++++++++++++ .../scalar/spatial/StYFromWKBEvaluator.java | 127 ++++++++++++++++++ .../function/EsqlFunctionRegistry.java | 4 + .../function/scalar/spatial/StX.java | 73 ++++++++++ .../function/scalar/spatial/StY.java | 73 ++++++++++ .../xpack/esql/io/stream/PlanNamedTypes.java | 6 + .../function/AbstractFunctionTestCase.java | 3 +- .../function/scalar/spatial/StXTests.java | 50 +++++++ .../function/scalar/spatial/StYTests.java | 50 +++++++ .../xpack/ql/util/SpatialCoordinateTypes.java | 16 ++- 21 files changed, 683 insertions(+), 5 deletions(-) create mode 100644 docs/changelog/105768.yaml create mode 100644 docs/reference/esql/functions/signature/st_x.svg create mode 100644 docs/reference/esql/functions/signature/st_y.svg create mode 100644 docs/reference/esql/functions/spatial-functions.asciidoc create mode 100644 docs/reference/esql/functions/st_x.asciidoc create mode 100644 docs/reference/esql/functions/st_y.asciidoc create mode 100644 docs/reference/esql/functions/types/st_x.asciidoc create mode 100644 docs/reference/esql/functions/types/st_y.asciidoc create mode 100644 x-pack/plugin/esql/src/main/java/generated/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StXFromWKBEvaluator.java create mode 100644 x-pack/plugin/esql/src/main/java/generated/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StYFromWKBEvaluator.java create mode 100644 x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StX.java create mode 100644 x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StY.java create mode 100644 x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StXTests.java create mode 100644 x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StYTests.java diff --git a/docs/changelog/105768.yaml b/docs/changelog/105768.yaml new file mode 100644 index 0000000000000..49d7f1f15c453 --- /dev/null +++ b/docs/changelog/105768.yaml @@ -0,0 +1,5 @@ +pr: 105768 +summary: Add two new OGC functions ST_X and ST_Y +area: "ES|QL" +type: enhancement +issues: [] diff --git a/docs/reference/esql/esql-functions-operators.asciidoc b/docs/reference/esql/esql-functions-operators.asciidoc index 375bb4ee9dd00..a1ad512fbe512 100644 --- a/docs/reference/esql/esql-functions-operators.asciidoc +++ b/docs/reference/esql/esql-functions-operators.asciidoc @@ -21,6 +21,9 @@ include::functions/string-functions.asciidoc[tag=string_list] <>:: include::functions/date-time-functions.asciidoc[tag=date_list] +<>:: +include::functions/spatial-functions.asciidoc[tag=spatial_list] + <>:: include::functions/type-conversion-functions.asciidoc[tag=type_list] @@ -37,6 +40,7 @@ include::functions/aggregation-functions.asciidoc[] include::functions/math-functions.asciidoc[] include::functions/string-functions.asciidoc[] include::functions/date-time-functions.asciidoc[] +include::functions/spatial-functions.asciidoc[] include::functions/type-conversion-functions.asciidoc[] include::functions/conditional-functions-and-expressions.asciidoc[] include::functions/mv-functions.asciidoc[] diff --git a/docs/reference/esql/functions/signature/st_x.svg b/docs/reference/esql/functions/signature/st_x.svg new file mode 100644 index 0000000000000..d6fac8a96505a --- /dev/null +++ b/docs/reference/esql/functions/signature/st_x.svg @@ -0,0 +1 @@ +ST_X(point) \ No newline at end of file diff --git a/docs/reference/esql/functions/signature/st_y.svg b/docs/reference/esql/functions/signature/st_y.svg new file mode 100644 index 0000000000000..c6dc23724d59c --- /dev/null +++ b/docs/reference/esql/functions/signature/st_y.svg @@ -0,0 +1 @@ +ST_Y(point) \ No newline at end of file diff --git a/docs/reference/esql/functions/spatial-functions.asciidoc b/docs/reference/esql/functions/spatial-functions.asciidoc new file mode 100644 index 0000000000000..d99fe36191a31 --- /dev/null +++ b/docs/reference/esql/functions/spatial-functions.asciidoc @@ -0,0 +1,16 @@ +[[esql-spatial-functions]] +==== {esql} spatial functions + +++++ +Spatial functions +++++ + +{esql} supports these spatial functions: + +// tag::spatial_list[] +* <> +* <> +// end::spatial_list[] + +include::st_x.asciidoc[] +include::st_y.asciidoc[] diff --git a/docs/reference/esql/functions/st_x.asciidoc b/docs/reference/esql/functions/st_x.asciidoc new file mode 100644 index 0000000000000..0f40a66417f9f --- /dev/null +++ b/docs/reference/esql/functions/st_x.asciidoc @@ -0,0 +1,33 @@ +[discrete] +[[esql-st_x]] +=== `ST_X` + +*Syntax* + +[.text-center] +image::esql/functions/signature/st_x.svg[Embedded,opts=inline] + +*Parameters* + +`point`:: +Expression of type `geo_point` or `cartesian_point`. If `null`, the function returns `null`. + +*Description* + +Extracts the `x` coordinate from the supplied point. +If the points is of type `geo_point` this is equivalent to extracting the `longitude` value. + +*Supported types* + +include::types/st_x.asciidoc[] + +*Example* + +[source.merge.styled,esql] +---- +include::{esql-specs}/spatial.csv-spec[tag=st_x_y] +---- +[%header.monospaced.styled,format=dsv,separator=|] +|=== +include::{esql-specs}/spatial.csv-spec[tag=st_x_y-result] +|=== diff --git a/docs/reference/esql/functions/st_y.asciidoc b/docs/reference/esql/functions/st_y.asciidoc new file mode 100644 index 0000000000000..e876852228d83 --- /dev/null +++ b/docs/reference/esql/functions/st_y.asciidoc @@ -0,0 +1,33 @@ +[discrete] +[[esql-st_y]] +=== `ST_Y` + +*Syntax* + +[.text-center] +image::esql/functions/signature/st_y.svg[Embedded,opts=inline] + +*Parameters* + +`point`:: +Expression of type `geo_point` or `cartesian_point`. If `null`, the function returns `null`. + +*Description* + +Extracts the `y` coordinate from the supplied point. +If the points is of type `geo_point` this is equivalent to extracting the `latitude` value. + +*Supported types* + +include::types/st_y.asciidoc[] + +*Example* + +[source.merge.styled,esql] +---- +include::{esql-specs}/spatial.csv-spec[tag=st_x_y] +---- +[%header.monospaced.styled,format=dsv,separator=|] +|=== +include::{esql-specs}/spatial.csv-spec[tag=st_x_y-result] +|=== diff --git a/docs/reference/esql/functions/types/st_x.asciidoc b/docs/reference/esql/functions/types/st_x.asciidoc new file mode 100644 index 0000000000000..94ed4b296f1d4 --- /dev/null +++ b/docs/reference/esql/functions/types/st_x.asciidoc @@ -0,0 +1,6 @@ +[%header.monospaced.styled,format=dsv,separator=|] +|=== +point | result +cartesian_point | double +geo_point | double +|=== diff --git a/docs/reference/esql/functions/types/st_y.asciidoc b/docs/reference/esql/functions/types/st_y.asciidoc new file mode 100644 index 0000000000000..94ed4b296f1d4 --- /dev/null +++ b/docs/reference/esql/functions/types/st_y.asciidoc @@ -0,0 +1,6 @@ +[%header.monospaced.styled,format=dsv,separator=|] +|=== +point | result +cartesian_point | double +geo_point | double +|=== diff --git a/x-pack/plugin/esql/qa/testFixtures/src/main/resources/show.csv-spec b/x-pack/plugin/esql/qa/testFixtures/src/main/resources/show.csv-spec index 6887a1bbe9069..3f2d87c6d7a08 100644 --- a/x-pack/plugin/esql/qa/testFixtures/src/main/resources/show.csv-spec +++ b/x-pack/plugin/esql/qa/testFixtures/src/main/resources/show.csv-spec @@ -68,6 +68,8 @@ sinh |"double sinh(n:double|integer|long|unsigned_long)"|n split |"keyword split(str:keyword|text, delim:keyword|text)" |[str, delim] |["keyword|text", "keyword|text"] |["", ""] |keyword | "Split a single valued string into multiple strings." | [false, false] | false | false sqrt |"double sqrt(n:double|integer|long|unsigned_long)" |n |"double|integer|long|unsigned_long" | "" |double | "Returns the square root of a number." | false | false | false st_centroid |"geo_point|cartesian_point st_centroid(field:geo_point|cartesian_point)" |field |"geo_point|cartesian_point" | "" |"geo_point|cartesian_point" | "The centroid of a spatial field." | false | false | true +st_x |"double st_x(point:geo_point|cartesian_point)" |point |"geo_point|cartesian_point" | "" |double | "Extracts the x-coordinate from a point geometry." | false | false | false +st_y |"double st_y(point:geo_point|cartesian_point)" |point |"geo_point|cartesian_point" | "" |double | "Extracts the y-coordinate from a point geometry." | false | false | false starts_with |"boolean starts_with(str:keyword|text, prefix:keyword|text)" |[str, prefix] |["keyword|text", "keyword|text"] |["", ""] |boolean | "Returns a boolean that indicates whether a keyword string starts with another string" | [false, false] | false | false substring |"keyword substring(str:keyword|text, start:integer, ?length:integer)" |[str, start, length] |["keyword|text", "integer", "integer"] |["", "", ""] |keyword | "Returns a substring of a string, specified by a start position and an optional length" | [false, false, true]| false | false sum |"long sum(field:double|integer|long)" |field |"double|integer|long" | "" |long | "The sum of a numeric field." | false | false | true @@ -103,7 +105,7 @@ trim |"keyword|text trim(str:keyword|text)" ; -showFunctionsSynopsis#[skip:-8.12.99] +showFunctionsSynopsis#[skip:-8.13.99] show functions | keep synopsis; synopsis:keyword @@ -165,6 +167,8 @@ double pi() "keyword split(str:keyword|text, delim:keyword|text)" "double sqrt(n:double|integer|long|unsigned_long)" "geo_point|cartesian_point st_centroid(field:geo_point|cartesian_point)" +"double st_x(point:geo_point|cartesian_point)" +"double st_y(point:geo_point|cartesian_point)" "boolean starts_with(str:keyword|text, prefix:keyword|text)" "keyword substring(str:keyword|text, start:integer, ?length:integer)" "long sum(field:double|integer|long)" @@ -216,9 +220,9 @@ sinh | "double sinh(n:double|integer|long|unsigned_long)" // see https://github.com/elastic/elasticsearch/issues/102120 -countFunctions#[skip:-8.12.99] +countFunctions#[skip:-8.13.99] show functions | stats a = count(*), b = count(*), c = count(*) | mv_expand c; a:long | b:long | c:long -90 | 90 | 90 +92 | 92 | 92 ; diff --git a/x-pack/plugin/esql/qa/testFixtures/src/main/resources/spatial.csv-spec b/x-pack/plugin/esql/qa/testFixtures/src/main/resources/spatial.csv-spec index 5c4aae740910b..1eb4d82b5fcc2 100644 --- a/x-pack/plugin/esql/qa/testFixtures/src/main/resources/spatial.csv-spec +++ b/x-pack/plugin/esql/qa/testFixtures/src/main/resources/spatial.csv-spec @@ -69,6 +69,30 @@ c:geo_point POINT(39.58327988510707 20.619513023697994) ; +centroidFromString4#[skip:-8.13.99, reason:st_x and st_y added in 8.14] +ROW wkt = ["POINT(42.97109629958868 14.7552534006536)", "POINT(75.80929149873555 22.72774917539209)", "POINT(-0.030548143003023033 24.37553649504829)"] +| MV_EXPAND wkt +| EVAL pt = TO_GEOPOINT(wkt) +| STATS c = ST_CENTROID(pt) +| EVAL x = ST_X(c), y = ST_Y(c); + +c:geo_point | x:double | y:double +POINT(39.58327988510707 20.619513023697994) | 39.58327988510707 | 20.619513023697994 +; + +stXFromString#[skip:-8.13.99, reason:st_x and st_y added in 8.14] +// tag::st_x_y[] +ROW point = TO_GEOPOINT("POINT(42.97109629958868 14.7552534006536)") +| EVAL x = ST_X(point), y = ST_Y(point) +// end::st_x_y[] +; + +// tag::st_x_y-result[] +point:geo_point | x:double | y:double +POINT(42.97109629958868 14.7552534006536) | 42.97109629958868 | 14.7552534006536 +// end::st_x_y-result[] +; + simpleLoad#[skip:-8.12.99, reason:spatial type geo_point improved precision in 8.13] FROM airports | WHERE scalerank == 9 | SORT abbrev | WHERE length(name) > 12; @@ -87,6 +111,17 @@ WIIT | Bandar Lampung | POINT(105.2667 -5.45) | Indonesia ZAH | Zāhedān | POINT(60.8628 29.4964) | Iran | POINT(60.900708564915 29.4752941956573) | Zahedan Int'l | 9 | mid ; +stXFromAirportsSupportsNull#[skip:-8.13.99, reason:st_x and st_y added in 8.14] +FROM airports +| EVAL x = FLOOR(ABS(ST_X(city_location))/200), y = FLOOR(ABS(ST_Y(city_location))/100) +| STATS c = count(*) BY x, y +; + +c:long | x:double | y:double +872 | 0.0 | 0.0 +19 | null | null +; + centroidFromAirports#[skip:-8.12.99, reason:st_centroid added in 8.13] // tag::st_centroid-airports[] FROM airports @@ -399,6 +434,15 @@ c:cartesian_point POINT(3949.163965353159 1078.2645465797348) ; +stXFromCartesianString#[skip:-8.13.99, reason:st_x and st_y added in 8.14] +ROW point = TO_CARTESIANPOINT("POINT(4297.10986328125 -1475.530029296875)") +| EVAL x = ST_X(point), y = ST_Y(point) +; + +point:cartesian_point | x:double | y:double +POINT(4297.10986328125 -1475.530029296875) | 4297.10986328125 | -1475.530029296875 +; + simpleCartesianLoad#[skip:-8.12.99, reason:spatial type cartesian_point improved precision in 8.13] FROM airports_web | WHERE scalerank == 9 | SORT abbrev | WHERE length(name) > 12; diff --git a/x-pack/plugin/esql/src/main/java/generated/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StXFromWKBEvaluator.java b/x-pack/plugin/esql/src/main/java/generated/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StXFromWKBEvaluator.java new file mode 100644 index 0000000000000..937eedc1d8fe0 --- /dev/null +++ b/x-pack/plugin/esql/src/main/java/generated/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StXFromWKBEvaluator.java @@ -0,0 +1,127 @@ +// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one +// or more contributor license agreements. Licensed under the Elastic License +// 2.0; you may not use this file except in compliance with the Elastic License +// 2.0. +package org.elasticsearch.xpack.esql.expression.function.scalar.spatial; + +import java.lang.IllegalArgumentException; +import java.lang.Override; +import java.lang.String; +import org.apache.lucene.util.BytesRef; +import org.elasticsearch.compute.data.Block; +import org.elasticsearch.compute.data.BytesRefBlock; +import org.elasticsearch.compute.data.BytesRefVector; +import org.elasticsearch.compute.data.DoubleBlock; +import org.elasticsearch.compute.data.Vector; +import org.elasticsearch.compute.operator.DriverContext; +import org.elasticsearch.compute.operator.EvalOperator; +import org.elasticsearch.xpack.esql.expression.function.scalar.convert.AbstractConvertFunction; +import org.elasticsearch.xpack.ql.tree.Source; + +/** + * {@link EvalOperator.ExpressionEvaluator} implementation for {@link StX}. + * This class is generated. Do not edit it. + */ +public final class StXFromWKBEvaluator extends AbstractConvertFunction.AbstractEvaluator { + public StXFromWKBEvaluator(EvalOperator.ExpressionEvaluator field, Source source, + DriverContext driverContext) { + super(driverContext, field, source); + } + + @Override + public String name() { + return "StXFromWKB"; + } + + @Override + public Block evalVector(Vector v) { + BytesRefVector vector = (BytesRefVector) v; + int positionCount = v.getPositionCount(); + BytesRef scratchPad = new BytesRef(); + if (vector.isConstant()) { + try { + return driverContext.blockFactory().newConstantDoubleBlockWith(evalValue(vector, 0, scratchPad), positionCount); + } catch (IllegalArgumentException e) { + registerException(e); + return driverContext.blockFactory().newConstantNullBlock(positionCount); + } + } + try (DoubleBlock.Builder builder = driverContext.blockFactory().newDoubleBlockBuilder(positionCount)) { + for (int p = 0; p < positionCount; p++) { + try { + builder.appendDouble(evalValue(vector, p, scratchPad)); + } catch (IllegalArgumentException e) { + registerException(e); + builder.appendNull(); + } + } + return builder.build(); + } + } + + private static double evalValue(BytesRefVector container, int index, BytesRef scratchPad) { + BytesRef value = container.getBytesRef(index, scratchPad); + return StX.fromWellKnownBinary(value); + } + + @Override + public Block evalBlock(Block b) { + BytesRefBlock block = (BytesRefBlock) b; + int positionCount = block.getPositionCount(); + try (DoubleBlock.Builder builder = driverContext.blockFactory().newDoubleBlockBuilder(positionCount)) { + BytesRef scratchPad = new BytesRef(); + for (int p = 0; p < positionCount; p++) { + int valueCount = block.getValueCount(p); + int start = block.getFirstValueIndex(p); + int end = start + valueCount; + boolean positionOpened = false; + boolean valuesAppended = false; + for (int i = start; i < end; i++) { + try { + double value = evalValue(block, i, scratchPad); + if (positionOpened == false && valueCount > 1) { + builder.beginPositionEntry(); + positionOpened = true; + } + builder.appendDouble(value); + valuesAppended = true; + } catch (IllegalArgumentException e) { + registerException(e); + } + } + if (valuesAppended == false) { + builder.appendNull(); + } else if (positionOpened) { + builder.endPositionEntry(); + } + } + return builder.build(); + } + } + + private static double evalValue(BytesRefBlock container, int index, BytesRef scratchPad) { + BytesRef value = container.getBytesRef(index, scratchPad); + return StX.fromWellKnownBinary(value); + } + + public static class Factory implements EvalOperator.ExpressionEvaluator.Factory { + private final Source source; + + private final EvalOperator.ExpressionEvaluator.Factory field; + + public Factory(EvalOperator.ExpressionEvaluator.Factory field, Source source) { + this.field = field; + this.source = source; + } + + @Override + public StXFromWKBEvaluator get(DriverContext context) { + return new StXFromWKBEvaluator(field.get(context), source, context); + } + + @Override + public String toString() { + return "StXFromWKBEvaluator[field=" + field + "]"; + } + } +} diff --git a/x-pack/plugin/esql/src/main/java/generated/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StYFromWKBEvaluator.java b/x-pack/plugin/esql/src/main/java/generated/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StYFromWKBEvaluator.java new file mode 100644 index 0000000000000..33405f6db5998 --- /dev/null +++ b/x-pack/plugin/esql/src/main/java/generated/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StYFromWKBEvaluator.java @@ -0,0 +1,127 @@ +// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one +// or more contributor license agreements. Licensed under the Elastic License +// 2.0; you may not use this file except in compliance with the Elastic License +// 2.0. +package org.elasticsearch.xpack.esql.expression.function.scalar.spatial; + +import java.lang.IllegalArgumentException; +import java.lang.Override; +import java.lang.String; +import org.apache.lucene.util.BytesRef; +import org.elasticsearch.compute.data.Block; +import org.elasticsearch.compute.data.BytesRefBlock; +import org.elasticsearch.compute.data.BytesRefVector; +import org.elasticsearch.compute.data.DoubleBlock; +import org.elasticsearch.compute.data.Vector; +import org.elasticsearch.compute.operator.DriverContext; +import org.elasticsearch.compute.operator.EvalOperator; +import org.elasticsearch.xpack.esql.expression.function.scalar.convert.AbstractConvertFunction; +import org.elasticsearch.xpack.ql.tree.Source; + +/** + * {@link EvalOperator.ExpressionEvaluator} implementation for {@link StY}. + * This class is generated. Do not edit it. + */ +public final class StYFromWKBEvaluator extends AbstractConvertFunction.AbstractEvaluator { + public StYFromWKBEvaluator(EvalOperator.ExpressionEvaluator field, Source source, + DriverContext driverContext) { + super(driverContext, field, source); + } + + @Override + public String name() { + return "StYFromWKB"; + } + + @Override + public Block evalVector(Vector v) { + BytesRefVector vector = (BytesRefVector) v; + int positionCount = v.getPositionCount(); + BytesRef scratchPad = new BytesRef(); + if (vector.isConstant()) { + try { + return driverContext.blockFactory().newConstantDoubleBlockWith(evalValue(vector, 0, scratchPad), positionCount); + } catch (IllegalArgumentException e) { + registerException(e); + return driverContext.blockFactory().newConstantNullBlock(positionCount); + } + } + try (DoubleBlock.Builder builder = driverContext.blockFactory().newDoubleBlockBuilder(positionCount)) { + for (int p = 0; p < positionCount; p++) { + try { + builder.appendDouble(evalValue(vector, p, scratchPad)); + } catch (IllegalArgumentException e) { + registerException(e); + builder.appendNull(); + } + } + return builder.build(); + } + } + + private static double evalValue(BytesRefVector container, int index, BytesRef scratchPad) { + BytesRef value = container.getBytesRef(index, scratchPad); + return StY.fromWellKnownBinary(value); + } + + @Override + public Block evalBlock(Block b) { + BytesRefBlock block = (BytesRefBlock) b; + int positionCount = block.getPositionCount(); + try (DoubleBlock.Builder builder = driverContext.blockFactory().newDoubleBlockBuilder(positionCount)) { + BytesRef scratchPad = new BytesRef(); + for (int p = 0; p < positionCount; p++) { + int valueCount = block.getValueCount(p); + int start = block.getFirstValueIndex(p); + int end = start + valueCount; + boolean positionOpened = false; + boolean valuesAppended = false; + for (int i = start; i < end; i++) { + try { + double value = evalValue(block, i, scratchPad); + if (positionOpened == false && valueCount > 1) { + builder.beginPositionEntry(); + positionOpened = true; + } + builder.appendDouble(value); + valuesAppended = true; + } catch (IllegalArgumentException e) { + registerException(e); + } + } + if (valuesAppended == false) { + builder.appendNull(); + } else if (positionOpened) { + builder.endPositionEntry(); + } + } + return builder.build(); + } + } + + private static double evalValue(BytesRefBlock container, int index, BytesRef scratchPad) { + BytesRef value = container.getBytesRef(index, scratchPad); + return StY.fromWellKnownBinary(value); + } + + public static class Factory implements EvalOperator.ExpressionEvaluator.Factory { + private final Source source; + + private final EvalOperator.ExpressionEvaluator.Factory field; + + public Factory(EvalOperator.ExpressionEvaluator.Factory field, Source source) { + this.field = field; + this.source = source; + } + + @Override + public StYFromWKBEvaluator get(DriverContext context) { + return new StYFromWKBEvaluator(field.get(context), source, context); + } + + @Override + public String toString() { + return "StYFromWKBEvaluator[field=" + field + "]"; + } + } +} diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/EsqlFunctionRegistry.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/EsqlFunctionRegistry.java index 395a9ed16dc67..ede3633c1b3e8 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/EsqlFunctionRegistry.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/EsqlFunctionRegistry.java @@ -75,6 +75,8 @@ import org.elasticsearch.xpack.esql.expression.function.scalar.multivalue.MvMin; import org.elasticsearch.xpack.esql.expression.function.scalar.multivalue.MvSum; import org.elasticsearch.xpack.esql.expression.function.scalar.nulls.Coalesce; +import org.elasticsearch.xpack.esql.expression.function.scalar.spatial.StX; +import org.elasticsearch.xpack.esql.expression.function.scalar.spatial.StY; import org.elasticsearch.xpack.esql.expression.function.scalar.string.Concat; import org.elasticsearch.xpack.esql.expression.function.scalar.string.EndsWith; import org.elasticsearch.xpack.esql.expression.function.scalar.string.LTrim; @@ -174,6 +176,8 @@ private FunctionDefinition[][] functions() { def(Now.class, Now::new, "now") }, // spatial new FunctionDefinition[] { def(SpatialCentroid.class, SpatialCentroid::new, "st_centroid") }, + new FunctionDefinition[] { def(StX.class, StX::new, "st_x") }, + new FunctionDefinition[] { def(StY.class, StY::new, "st_y") }, // conditional new FunctionDefinition[] { def(Case.class, Case::new, "case") }, // null diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StX.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StX.java new file mode 100644 index 0000000000000..f86be9290fed1 --- /dev/null +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StX.java @@ -0,0 +1,73 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.esql.expression.function.scalar.spatial; + +import org.apache.lucene.util.BytesRef; +import org.elasticsearch.compute.ann.ConvertEvaluator; +import org.elasticsearch.compute.operator.EvalOperator; +import org.elasticsearch.xpack.esql.expression.function.FunctionInfo; +import org.elasticsearch.xpack.esql.expression.function.Param; +import org.elasticsearch.xpack.esql.expression.function.scalar.UnaryScalarFunction; +import org.elasticsearch.xpack.ql.expression.Expression; +import org.elasticsearch.xpack.ql.expression.TypeResolutions; +import org.elasticsearch.xpack.ql.tree.NodeInfo; +import org.elasticsearch.xpack.ql.tree.Source; +import org.elasticsearch.xpack.ql.type.DataType; + +import java.util.List; +import java.util.function.Function; + +import static org.elasticsearch.xpack.esql.expression.EsqlTypeResolutions.isSpatialPoint; +import static org.elasticsearch.xpack.ql.type.DataTypes.DOUBLE; +import static org.elasticsearch.xpack.ql.util.SpatialCoordinateTypes.UNSPECIFIED; + +/** + * Extracts the x-coordinate from a point geometry. + * For cartesian geometries, the x-coordinate is the first coordinate. + * For geographic geometries, the x-coordinate is the longitude. + * The function `st_x` is defined in the OGC Simple Feature Access standard. + * Alternatively it is well described in PostGIS documentation at PostGIS:ST_X. + */ +public class StX extends UnaryScalarFunction { + @FunctionInfo(returnType = "double", description = "Extracts the x-coordinate from a point geometry.") + public StX(Source source, @Param(name = "point", type = { "geo_point", "cartesian_point" }) Expression field) { + super(source, field); + } + + @Override + protected Expression.TypeResolution resolveType() { + return isSpatialPoint(field(), sourceText(), TypeResolutions.ParamOrdinal.DEFAULT); + } + + @Override + public EvalOperator.ExpressionEvaluator.Factory toEvaluator( + Function toEvaluator + ) { + return new StXFromWKBEvaluator.Factory(toEvaluator.apply(field()), source()); + } + + @Override + public DataType dataType() { + return DOUBLE; + } + + @Override + public Expression replaceChildren(List newChildren) { + return new StX(source(), newChildren.get(0)); + } + + @Override + protected NodeInfo info() { + return NodeInfo.create(this, StX::new, field()); + } + + @ConvertEvaluator(extraName = "FromWKB", warnExceptions = { IllegalArgumentException.class }) + static double fromWellKnownBinary(BytesRef in) { + return UNSPECIFIED.wkbAsPoint(in).getX(); + } +} diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StY.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StY.java new file mode 100644 index 0000000000000..759c23c73374a --- /dev/null +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StY.java @@ -0,0 +1,73 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.esql.expression.function.scalar.spatial; + +import org.apache.lucene.util.BytesRef; +import org.elasticsearch.compute.ann.ConvertEvaluator; +import org.elasticsearch.compute.operator.EvalOperator; +import org.elasticsearch.xpack.esql.expression.function.FunctionInfo; +import org.elasticsearch.xpack.esql.expression.function.Param; +import org.elasticsearch.xpack.esql.expression.function.scalar.UnaryScalarFunction; +import org.elasticsearch.xpack.ql.expression.Expression; +import org.elasticsearch.xpack.ql.expression.TypeResolutions; +import org.elasticsearch.xpack.ql.tree.NodeInfo; +import org.elasticsearch.xpack.ql.tree.Source; +import org.elasticsearch.xpack.ql.type.DataType; + +import java.util.List; +import java.util.function.Function; + +import static org.elasticsearch.xpack.esql.expression.EsqlTypeResolutions.isSpatialPoint; +import static org.elasticsearch.xpack.ql.type.DataTypes.DOUBLE; +import static org.elasticsearch.xpack.ql.util.SpatialCoordinateTypes.UNSPECIFIED; + +/** + * Extracts the y-coordinate from a point geometry. + * For cartesian geometries, the y-coordinate is the second coordinate. + * For geographic geometries, the y-coordinate is the latitude. + * The function `st_y` is defined in the OGC Simple Feature Access standard. + * Alternatively it is well described in PostGIS documentation at PostGIS:ST_Y. + */ +public class StY extends UnaryScalarFunction { + @FunctionInfo(returnType = "double", description = "Extracts the y-coordinate from a point geometry.") + public StY(Source source, @Param(name = "point", type = { "geo_point", "cartesian_point" }) Expression field) { + super(source, field); + } + + @Override + protected TypeResolution resolveType() { + return isSpatialPoint(field(), sourceText(), TypeResolutions.ParamOrdinal.DEFAULT); + } + + @Override + public EvalOperator.ExpressionEvaluator.Factory toEvaluator( + Function toEvaluator + ) { + return new StYFromWKBEvaluator.Factory(toEvaluator.apply(field()), source()); + } + + @Override + public DataType dataType() { + return DOUBLE; + } + + @Override + public Expression replaceChildren(List newChildren) { + return new StY(source(), newChildren.get(0)); + } + + @Override + protected NodeInfo info() { + return NodeInfo.create(this, StY::new, field()); + } + + @ConvertEvaluator(extraName = "FromWKB", warnExceptions = { IllegalArgumentException.class }) + static double fromWellKnownBinary(BytesRef in) { + return UNSPECIFIED.wkbAsPoint(in).getY(); + } +} diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/io/stream/PlanNamedTypes.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/io/stream/PlanNamedTypes.java index 95892ac42e587..3ca5f2f5868ba 100644 --- a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/io/stream/PlanNamedTypes.java +++ b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/io/stream/PlanNamedTypes.java @@ -99,6 +99,8 @@ import org.elasticsearch.xpack.esql.expression.function.scalar.multivalue.MvMin; import org.elasticsearch.xpack.esql.expression.function.scalar.multivalue.MvSum; import org.elasticsearch.xpack.esql.expression.function.scalar.nulls.Coalesce; +import org.elasticsearch.xpack.esql.expression.function.scalar.spatial.StX; +import org.elasticsearch.xpack.esql.expression.function.scalar.spatial.StY; import org.elasticsearch.xpack.esql.expression.function.scalar.string.Concat; import org.elasticsearch.xpack.esql.expression.function.scalar.string.EndsWith; import org.elasticsearch.xpack.esql.expression.function.scalar.string.LTrim; @@ -340,6 +342,8 @@ public static List namedTypeEntries() { of(ESQL_UNARY_SCLR_CLS, Sin.class, PlanNamedTypes::writeESQLUnaryScalar, PlanNamedTypes::readESQLUnaryScalar), of(ESQL_UNARY_SCLR_CLS, Sinh.class, PlanNamedTypes::writeESQLUnaryScalar, PlanNamedTypes::readESQLUnaryScalar), of(ESQL_UNARY_SCLR_CLS, Sqrt.class, PlanNamedTypes::writeESQLUnaryScalar, PlanNamedTypes::readESQLUnaryScalar), + of(ESQL_UNARY_SCLR_CLS, StX.class, PlanNamedTypes::writeESQLUnaryScalar, PlanNamedTypes::readESQLUnaryScalar), + of(ESQL_UNARY_SCLR_CLS, StY.class, PlanNamedTypes::writeESQLUnaryScalar, PlanNamedTypes::readESQLUnaryScalar), of(ESQL_UNARY_SCLR_CLS, Tan.class, PlanNamedTypes::writeESQLUnaryScalar, PlanNamedTypes::readESQLUnaryScalar), of(ESQL_UNARY_SCLR_CLS, Tanh.class, PlanNamedTypes::writeESQLUnaryScalar, PlanNamedTypes::readESQLUnaryScalar), of(ESQL_UNARY_SCLR_CLS, ToBoolean.class, PlanNamedTypes::writeESQLUnaryScalar, PlanNamedTypes::readESQLUnaryScalar), @@ -1248,6 +1252,8 @@ static void writeBinaryLogic(PlanStreamOutput out, BinaryLogic binaryLogic) thro entry(name(Sin.class), Sin::new), entry(name(Sinh.class), Sinh::new), entry(name(Sqrt.class), Sqrt::new), + entry(name(StX.class), StX::new), + entry(name(StY.class), StY::new), entry(name(Tan.class), Tan::new), entry(name(Tanh.class), Tanh::new), entry(name(ToBoolean.class), ToBoolean::new), diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/AbstractFunctionTestCase.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/AbstractFunctionTestCase.java index 612861b2889a4..4d44d3111c094 100644 --- a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/AbstractFunctionTestCase.java +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/AbstractFunctionTestCase.java @@ -967,7 +967,8 @@ protected static String typeErrorMessage(boolean includeOrdinal, List testCaseSupplier) { + this.testCase = testCaseSupplier.get(); + } + + @ParametersFactory + public static Iterable parameters() { + String expectedEvaluator = "StXFromWKBEvaluator[field=Attribute[channel=0]]"; + final List suppliers = new ArrayList<>(); + TestCaseSupplier.forUnaryGeoPoint(suppliers, expectedEvaluator, DOUBLE, StXTests::valueOf, List.of()); + TestCaseSupplier.forUnaryCartesianPoint(suppliers, expectedEvaluator, DOUBLE, StXTests::valueOf, List.of()); + return parameterSuppliersFromTypedData(errorsForCasesWithoutExamples(anyNullIsNull(true, suppliers))); + } + + private static double valueOf(BytesRef wkb) { + return UNSPECIFIED.wkbAsPoint(wkb).getX(); + } + + @Override + protected Expression build(Source source, List args) { + return new StX(source, args.get(0)); + } +} diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StYTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StYTests.java new file mode 100644 index 0000000000000..9416b7ba8cad4 --- /dev/null +++ b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/scalar/spatial/StYTests.java @@ -0,0 +1,50 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.esql.expression.function.scalar.spatial; + +import com.carrotsearch.randomizedtesting.annotations.Name; +import com.carrotsearch.randomizedtesting.annotations.ParametersFactory; + +import org.apache.lucene.util.BytesRef; +import org.elasticsearch.xpack.esql.expression.function.AbstractFunctionTestCase; +import org.elasticsearch.xpack.esql.expression.function.FunctionName; +import org.elasticsearch.xpack.esql.expression.function.TestCaseSupplier; +import org.elasticsearch.xpack.ql.expression.Expression; +import org.elasticsearch.xpack.ql.tree.Source; + +import java.util.ArrayList; +import java.util.List; +import java.util.function.Supplier; + +import static org.elasticsearch.xpack.ql.type.DataTypes.DOUBLE; +import static org.elasticsearch.xpack.ql.util.SpatialCoordinateTypes.UNSPECIFIED; + +@FunctionName("st_y") +public class StYTests extends AbstractFunctionTestCase { + public StYTests(@Name("TestCase") Supplier testCaseSupplier) { + this.testCase = testCaseSupplier.get(); + } + + @ParametersFactory + public static Iterable parameters() { + String expectedEvaluator = "StYFromWKBEvaluator[field=Attribute[channel=0]]"; + final List suppliers = new ArrayList<>(); + TestCaseSupplier.forUnaryGeoPoint(suppliers, expectedEvaluator, DOUBLE, StYTests::valueOf, List.of()); + TestCaseSupplier.forUnaryCartesianPoint(suppliers, expectedEvaluator, DOUBLE, StYTests::valueOf, List.of()); + return parameterSuppliersFromTypedData(errorsForCasesWithoutExamples(anyNullIsNull(true, suppliers))); + } + + private static double valueOf(BytesRef wkb) { + return UNSPECIFIED.wkbAsPoint(wkb).getY(); + } + + @Override + protected Expression build(Source source, List args) { + return new StY(source, args.get(0)); + } +} diff --git a/x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/util/SpatialCoordinateTypes.java b/x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/util/SpatialCoordinateTypes.java index 6508a67f7e785..32bd76cf84e19 100644 --- a/x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/util/SpatialCoordinateTypes.java +++ b/x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/util/SpatialCoordinateTypes.java @@ -56,6 +56,15 @@ public long pointAsLong(double x, double y) { final long yi = XYEncodingUtils.encode((float) y); return (yi & 0xFFFFFFFFL) | xi << 32; } + }, + UNSPECIFIED { + public Point longAsPoint(long encoded) { + throw new UnsupportedOperationException("Cannot convert long to point without specifying coordinate type"); + } + + public long pointAsLong(double x, double y) { + throw new UnsupportedOperationException("Cannot convert point to long without specifying coordinate type"); + } }; public abstract Point longAsPoint(long encoded); @@ -63,9 +72,14 @@ public long pointAsLong(double x, double y) { public abstract long pointAsLong(double x, double y); public long wkbAsLong(BytesRef wkb) { + Point point = wkbAsPoint(wkb); + return pointAsLong(point.getX(), point.getY()); + } + + public Point wkbAsPoint(BytesRef wkb) { Geometry geometry = WellKnownBinary.fromWKB(GeometryValidator.NOOP, false, wkb.bytes, wkb.offset, wkb.length); if (geometry instanceof Point point) { - return pointAsLong(point.getX(), point.getY()); + return point; } else { throw new IllegalArgumentException("Unsupported geometry: " + geometry.type()); } From eeecdbf87b08ec03682dabe7def7272055119252 Mon Sep 17 00:00:00 2001 From: Ievgen Degtiarenko Date: Wed, 6 Mar 2024 11:28:23 +0100 Subject: [PATCH 19/27] Additional trace logging for desired balance computer (#105910) --- .../allocation/decider/DiskThresholdDeciderIT.java | 3 ++- .../allocator/DesiredBalanceComputer.java | 14 ++++++++++++-- .../snapshots/SnapshotShardSizeInfo.java | 5 +++++ 3 files changed, 19 insertions(+), 3 deletions(-) diff --git a/server/src/internalClusterTest/java/org/elasticsearch/cluster/routing/allocation/decider/DiskThresholdDeciderIT.java b/server/src/internalClusterTest/java/org/elasticsearch/cluster/routing/allocation/decider/DiskThresholdDeciderIT.java index 7b9f89b60ed94..56eacb1bc41b5 100644 --- a/server/src/internalClusterTest/java/org/elasticsearch/cluster/routing/allocation/decider/DiskThresholdDeciderIT.java +++ b/server/src/internalClusterTest/java/org/elasticsearch/cluster/routing/allocation/decider/DiskThresholdDeciderIT.java @@ -162,7 +162,8 @@ public void testRestoreSnapshotAllocationDoesNotExceedWatermark() throws Excepti } @TestIssueLogging( - value = "org.elasticsearch.cluster.routing.allocation.allocator.DesiredBalanceReconciler:DEBUG," + value = "org.elasticsearch.cluster.routing.allocation.allocator.DesiredBalanceComputer:TRACE," + + "org.elasticsearch.cluster.routing.allocation.allocator.DesiredBalanceReconciler:DEBUG," + "org.elasticsearch.cluster.routing.allocation.allocator.DesiredBalanceShardsAllocator:TRACE", issueUrl = "https://github.com/elastic/elasticsearch/issues/105331" ) diff --git a/server/src/main/java/org/elasticsearch/cluster/routing/allocation/allocator/DesiredBalanceComputer.java b/server/src/main/java/org/elasticsearch/cluster/routing/allocation/allocator/DesiredBalanceComputer.java index effd5ec110c44..3a26bbcc7b280 100644 --- a/server/src/main/java/org/elasticsearch/cluster/routing/allocation/allocator/DesiredBalanceComputer.java +++ b/server/src/main/java/org/elasticsearch/cluster/routing/allocation/allocator/DesiredBalanceComputer.java @@ -77,7 +77,18 @@ public DesiredBalance compute( Predicate isFresh ) { - logger.debug("Recomputing desired balance for [{}]", desiredBalanceInput.index()); + if (logger.isTraceEnabled()) { + logger.trace( + "Recomputing desired balance for [{}]: {}, {}, {}, {}", + desiredBalanceInput.index(), + previousDesiredBalance, + desiredBalanceInput.routingAllocation().routingNodes().toString(), + desiredBalanceInput.routingAllocation().clusterInfo().toString(), + desiredBalanceInput.routingAllocation().snapshotShardSizeInfo().toString() + ); + } else { + logger.debug("Recomputing desired balance for [{}]", desiredBalanceInput.index()); + } final var routingAllocation = desiredBalanceInput.routingAllocation().mutableCloneForSimulation(); final var routingNodes = routingAllocation.routingNodes(); @@ -283,7 +294,6 @@ public DesiredBalance compute( hasChanges = true; clusterInfoSimulator.simulateShardStarted(shardRouting); routingNodes.startShard(logger, shardRouting, changes, 0L); - logger.trace("starting shard {}", shardRouting); } } } diff --git a/server/src/main/java/org/elasticsearch/snapshots/SnapshotShardSizeInfo.java b/server/src/main/java/org/elasticsearch/snapshots/SnapshotShardSizeInfo.java index 29ae2d1c5da4b..3bd5431c7be63 100644 --- a/server/src/main/java/org/elasticsearch/snapshots/SnapshotShardSizeInfo.java +++ b/server/src/main/java/org/elasticsearch/snapshots/SnapshotShardSizeInfo.java @@ -46,4 +46,9 @@ public long getShardSize(ShardRouting shardRouting, long fallback) { } return shardSize; } + + @Override + public String toString() { + return "SnapshotShardSizeInfo{snapshotShardSizes=" + snapshotShardSizes + '}'; + } } From 6fbf9892aa6964151c7491ba8e6b5a00336bba84 Mon Sep 17 00:00:00 2001 From: Mary Gouseti Date: Wed, 6 Mar 2024 14:04:18 +0200 Subject: [PATCH 20/27] Add the parameter `failure_store` to multi-target syntax APIs (#105386) In this PR we introduce a new query parameter behind the failure store feature flag. The query param, `faliure_store` allows the multi-syntax supporting APIs to choose the failure store indices as well. If an API should not support failure store, the `allowFailureStore` flag should be `false`. --- .../anomaly-detection/apis/put-job.asciidoc | 1 + .../datastreams/FailureStoreQueryParamIT.java | 216 +++++++++ .../elasticsearch/ElasticsearchException.java | 7 + .../org/elasticsearch/TransportVersions.java | 1 + .../admin/indices/get/GetIndexRequest.java | 12 +- .../mapping/put/PutMappingRequest.java | 19 +- .../action/support/IndicesOptions.java | 455 +++++++++++++++--- .../master/info/ClusterInfoRequest.java | 5 + .../cluster/metadata/DataStream.java | 27 ++ .../metadata/IndexNameExpressionResolver.java | 136 +++++- .../cluster/metadata/Metadata.java | 8 + .../FailureIndexNotSupportedException.java | 37 ++ .../ExceptionSerializationTests.java | 2 + .../indices/get/GetIndexRequestTests.java | 15 + .../action/support/IndicesOptionsTests.java | 52 +- .../cluster/metadata/DataStreamTests.java | 158 ++++++ .../IndexNameExpressionResolverTests.java | 198 +++++++- .../metadata/DataStreamTestHelper.java | 4 + .../datafeed/DatafeedNodeSelectorTests.java | 13 +- .../xpack/ql/index/IndexResolver.java | 16 +- 20 files changed, 1261 insertions(+), 121 deletions(-) create mode 100644 modules/data-streams/src/javaRestTest/java/org/elasticsearch/datastreams/FailureStoreQueryParamIT.java create mode 100644 server/src/main/java/org/elasticsearch/indices/FailureIndexNotSupportedException.java diff --git a/docs/reference/ml/anomaly-detection/apis/put-job.asciidoc b/docs/reference/ml/anomaly-detection/apis/put-job.asciidoc index 97120ff1873ae..1ab5de76a94b0 100644 --- a/docs/reference/ml/anomaly-detection/apis/put-job.asciidoc +++ b/docs/reference/ml/anomaly-detection/apis/put-job.asciidoc @@ -537,3 +537,4 @@ The API returns the following results: // TESTRESPONSE[s/"job_version" : "8.4.0"/"job_version" : $body.job_version/] // TESTRESPONSE[s/1656087283340/$body.$_path/] // TESTRESPONSE[s/"superuser"/"_es_test_root"/] +// TESTRESPONSE[s/"ignore_throttled" : true/"ignore_throttled" : true,"failure_store":"false"/] diff --git a/modules/data-streams/src/javaRestTest/java/org/elasticsearch/datastreams/FailureStoreQueryParamIT.java b/modules/data-streams/src/javaRestTest/java/org/elasticsearch/datastreams/FailureStoreQueryParamIT.java new file mode 100644 index 0000000000000..a6b235e8d566f --- /dev/null +++ b/modules/data-streams/src/javaRestTest/java/org/elasticsearch/datastreams/FailureStoreQueryParamIT.java @@ -0,0 +1,216 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0 and the Server Side Public License, v 1; you may not use this file except + * in compliance with, at your election, the Elastic License 2.0 or the Server + * Side Public License, v 1. + */ + +package org.elasticsearch.datastreams; + +import org.elasticsearch.client.Request; +import org.elasticsearch.client.Response; +import org.elasticsearch.client.ResponseException; +import org.junit.Before; + +import java.io.IOException; +import java.util.List; +import java.util.Map; + +import static org.hamcrest.Matchers.equalTo; +import static org.hamcrest.Matchers.is; + +/** + * This should be a yaml test, but in order to write one we would need to expose the new parameter in the rest-api-spec. + * We do not want to do that until the feature flag is removed. For this reason, we temporarily, test the affected APIs here. + * Please convert this to a yaml test when the feature flag is removed. + */ +public class FailureStoreQueryParamIT extends DisabledSecurityDataStreamTestCase { + + private static final String DATA_STREAM_NAME = "failure-data-stream"; + private String backingIndex; + private String failureStoreIndex; + + @SuppressWarnings("unchecked") + @Before + public void setup() throws IOException { + Request putComposableIndexTemplateRequest = new Request("POST", "/_index_template/ds-template"); + putComposableIndexTemplateRequest.setJsonEntity(""" + { + "index_patterns": ["failure-data-stream"], + "template": { + "settings": { + "number_of_replicas": 0 + } + }, + "data_stream": { + "failure_store": true + } + } + """); + assertOK(client().performRequest(putComposableIndexTemplateRequest)); + + assertOK(client().performRequest(new Request("PUT", "/_data_stream/" + DATA_STREAM_NAME))); + ensureGreen(DATA_STREAM_NAME); + + final Response dataStreamResponse = client().performRequest(new Request("GET", "/_data_stream/" + DATA_STREAM_NAME)); + List dataStreams = (List) entityAsMap(dataStreamResponse).get("data_streams"); + assertThat(dataStreams.size(), is(1)); + Map dataStream = (Map) dataStreams.get(0); + assertThat(dataStream.get("name"), equalTo(DATA_STREAM_NAME)); + List backingIndices = getBackingIndices(dataStream); + assertThat(backingIndices.size(), is(1)); + List failureStore = getFailureStore(dataStream); + assertThat(failureStore.size(), is(1)); + backingIndex = backingIndices.get(0); + failureStoreIndex = failureStore.get(0); + } + + public void testGetIndexApi() throws IOException { + { + final Response indicesResponse = client().performRequest(new Request("GET", "/" + DATA_STREAM_NAME)); + Map indices = entityAsMap(indicesResponse); + assertThat(indices.size(), is(2)); + assertThat(indices.containsKey(backingIndex), is(true)); + assertThat(indices.containsKey(failureStoreIndex), is(true)); + } + { + final Response indicesResponse = client().performRequest(new Request("GET", "/" + DATA_STREAM_NAME + "?failure_store=false")); + Map indices = entityAsMap(indicesResponse); + assertThat(indices.size(), is(1)); + assertThat(indices.containsKey(backingIndex), is(true)); + } + { + final Response indicesResponse = client().performRequest(new Request("GET", "/" + DATA_STREAM_NAME + "?failure_store=only")); + Map indices = entityAsMap(indicesResponse); + assertThat(indices.size(), is(1)); + assertThat(indices.containsKey(failureStoreIndex), is(true)); + } + } + + @SuppressWarnings("unchecked") + public void testGetIndexStatsApi() throws IOException { + { + final Response statsResponse = client().performRequest(new Request("GET", "/" + DATA_STREAM_NAME + "/_stats")); + Map indices = (Map) entityAsMap(statsResponse).get("indices"); + assertThat(indices.size(), is(1)); + assertThat(indices.containsKey(backingIndex), is(true)); + } + { + final Response statsResponse = client().performRequest( + new Request("GET", "/" + DATA_STREAM_NAME + "/_stats?failure_store=true") + ); + Map indices = (Map) entityAsMap(statsResponse).get("indices"); + assertThat(indices.size(), is(2)); + assertThat(indices.containsKey(backingIndex), is(true)); + assertThat(indices.containsKey(failureStoreIndex), is(true)); + } + { + final Response statsResponse = client().performRequest( + new Request("GET", "/" + DATA_STREAM_NAME + "/_stats?failure_store=only") + ); + Map indices = (Map) entityAsMap(statsResponse).get("indices"); + assertThat(indices.size(), is(1)); + assertThat(indices.containsKey(failureStoreIndex), is(true)); + } + } + + public void testGetIndexSettingsApi() throws IOException { + { + final Response indicesResponse = client().performRequest(new Request("GET", "/" + DATA_STREAM_NAME + "/_settings")); + Map indices = entityAsMap(indicesResponse); + assertThat(indices.size(), is(1)); + assertThat(indices.containsKey(backingIndex), is(true)); + } + { + final Response indicesResponse = client().performRequest( + new Request("GET", "/" + DATA_STREAM_NAME + "/_settings?failure_store=true") + ); + Map indices = entityAsMap(indicesResponse); + assertThat(indices.size(), is(2)); + assertThat(indices.containsKey(backingIndex), is(true)); + assertThat(indices.containsKey(failureStoreIndex), is(true)); + } + { + final Response indicesResponse = client().performRequest( + new Request("GET", "/" + DATA_STREAM_NAME + "/_settings?failure_store=only") + ); + Map indices = entityAsMap(indicesResponse); + assertThat(indices.size(), is(1)); + assertThat(indices.containsKey(failureStoreIndex), is(true)); + } + } + + public void testGetIndexMappingApi() throws IOException { + { + final Response indicesResponse = client().performRequest(new Request("GET", "/" + DATA_STREAM_NAME + "/_mapping")); + Map indices = entityAsMap(indicesResponse); + assertThat(indices.size(), is(1)); + assertThat(indices.containsKey(backingIndex), is(true)); + } + { + final Response indicesResponse = client().performRequest( + new Request("GET", "/" + DATA_STREAM_NAME + "/_mapping?failure_store=true") + ); + Map indices = entityAsMap(indicesResponse); + assertThat(indices.size(), is(2)); + assertThat(indices.containsKey(backingIndex), is(true)); + assertThat(indices.containsKey(failureStoreIndex), is(true)); + } + { + final Response indicesResponse = client().performRequest( + new Request("GET", "/" + DATA_STREAM_NAME + "/_mapping?failure_store=only") + ); + Map indices = entityAsMap(indicesResponse); + assertThat(indices.size(), is(1)); + assertThat(indices.containsKey(failureStoreIndex), is(true)); + } + } + + @SuppressWarnings("unchecked") + public void testPutIndexMappingApi() throws IOException { + { + final Request mappingRequest = new Request("PUT", "/" + DATA_STREAM_NAME + "/_mapping"); + mappingRequest.setJsonEntity(""" + { + "properties": { + "email": { + "type": "keyword" + } + } + } + """); + assertAcknowledged(client().performRequest(mappingRequest)); + } + { + final Request mappingRequest = new Request("PUT", "/" + DATA_STREAM_NAME + "/_mapping?failure_store=true"); + mappingRequest.setJsonEntity(""" + { + "properties": { + "email": { + "type": "keyword" + } + } + } + """); + ResponseException responseException = expectThrows(ResponseException.class, () -> client().performRequest(mappingRequest)); + Map response = entityAsMap(responseException.getResponse()); + assertThat(((Map) response.get("error")).get("reason"), is("failure index not supported")); + } + } + + private List getBackingIndices(Map response) { + return getIndices(response, "indices"); + } + + private List getFailureStore(Map response) { + return getIndices(response, "failure_indices"); + + } + + @SuppressWarnings("unchecked") + private List getIndices(Map response, String fieldName) { + List> indices = (List>) response.get(fieldName); + return indices.stream().map(index -> index.get("index_name")).toList(); + } +} diff --git a/server/src/main/java/org/elasticsearch/ElasticsearchException.java b/server/src/main/java/org/elasticsearch/ElasticsearchException.java index 33566203bb99a..83e5375546b63 100644 --- a/server/src/main/java/org/elasticsearch/ElasticsearchException.java +++ b/server/src/main/java/org/elasticsearch/ElasticsearchException.java @@ -29,6 +29,7 @@ import org.elasticsearch.index.mapper.DocumentParsingException; import org.elasticsearch.index.shard.ShardId; import org.elasticsearch.indices.AutoscalingMissedIndicesUpdateException; +import org.elasticsearch.indices.FailureIndexNotSupportedException; import org.elasticsearch.indices.recovery.RecoveryCommitTooNewException; import org.elasticsearch.ingest.GraphStructureException; import org.elasticsearch.rest.ApiNotAvailableException; @@ -1910,6 +1911,12 @@ private enum ElasticsearchExceptionHandle { GraphStructureException::new, 177, TransportVersions.INGEST_GRAPH_STRUCTURE_EXCEPTION + ), + FAILURE_INDEX_NOT_SUPPORTED_EXCEPTION( + FailureIndexNotSupportedException.class, + FailureIndexNotSupportedException::new, + 178, + TransportVersions.ADD_FAILURE_STORE_INDICES_OPTIONS ); final Class exceptionClass; diff --git a/server/src/main/java/org/elasticsearch/TransportVersions.java b/server/src/main/java/org/elasticsearch/TransportVersions.java index d484da5ba506c..ec3971a48a649 100644 --- a/server/src/main/java/org/elasticsearch/TransportVersions.java +++ b/server/src/main/java/org/elasticsearch/TransportVersions.java @@ -136,6 +136,7 @@ static TransportVersion def(int id) { public static final TransportVersion RANDOM_AGG_SHARD_SEED = def(8_596_00_0); public static final TransportVersion ESQL_TIMINGS = def(8_597_00_0); public static final TransportVersion DATA_STREAM_AUTO_SHARDING_EVENT = def(8_598_00_0); + public static final TransportVersion ADD_FAILURE_STORE_INDICES_OPTIONS = def(8_599_00_0); /* * STOP! READ THIS FIRST! No, really, diff --git a/server/src/main/java/org/elasticsearch/action/admin/indices/get/GetIndexRequest.java b/server/src/main/java/org/elasticsearch/action/admin/indices/get/GetIndexRequest.java index 1c2598d70998a..a550350c20f6b 100644 --- a/server/src/main/java/org/elasticsearch/action/admin/indices/get/GetIndexRequest.java +++ b/server/src/main/java/org/elasticsearch/action/admin/indices/get/GetIndexRequest.java @@ -9,7 +9,9 @@ package org.elasticsearch.action.admin.indices.get; import org.elasticsearch.action.ActionRequestValidationException; +import org.elasticsearch.action.support.IndicesOptions; import org.elasticsearch.action.support.master.info.ClusterInfoRequest; +import org.elasticsearch.cluster.metadata.DataStream; import org.elasticsearch.common.io.stream.StreamInput; import org.elasticsearch.common.io.stream.StreamOutput; import org.elasticsearch.common.util.ArrayUtils; @@ -92,7 +94,15 @@ public static Feature[] fromRequest(RestRequest request) { private transient boolean includeDefaults = false; public GetIndexRequest() { - + super( + DataStream.isFailureStoreEnabled() + ? IndicesOptions.builder(IndicesOptions.strictExpandOpen()) + .failureStoreOptions( + IndicesOptions.FailureStoreOptions.builder().includeRegularIndices(true).includeFailureIndices(true) + ) + .build() + : IndicesOptions.strictExpandOpen() + ); } public GetIndexRequest(StreamInput in) throws IOException { diff --git a/server/src/main/java/org/elasticsearch/action/admin/indices/mapping/put/PutMappingRequest.java b/server/src/main/java/org/elasticsearch/action/admin/indices/mapping/put/PutMappingRequest.java index 45532d8024f87..a2787e1a55fd7 100644 --- a/server/src/main/java/org/elasticsearch/action/admin/indices/mapping/put/PutMappingRequest.java +++ b/server/src/main/java/org/elasticsearch/action/admin/indices/mapping/put/PutMappingRequest.java @@ -66,7 +66,24 @@ public class PutMappingRequest extends AcknowledgedRequest im private String[] indices; - private IndicesOptions indicesOptions = IndicesOptions.fromOptions(false, false, true, true); + private IndicesOptions indicesOptions = IndicesOptions.builder() + .concreteTargetOptions(IndicesOptions.ConcreteTargetOptions.ERROR_WHEN_UNAVAILABLE_TARGETS) + .wildcardOptions( + IndicesOptions.WildcardOptions.builder() + .matchOpen(true) + .matchClosed(true) + .includeHidden(false) + .allowEmptyExpressions(false) + .resolveAliases(true) + ) + .gatekeeperOptions( + IndicesOptions.GatekeeperOptions.builder() + .allowClosedIndices(true) + .allowAliasToMultipleIndices(true) + .ignoreThrottled(false) + .allowFailureIndices(false) + ) + .build(); private String source; private String origin = ""; diff --git a/server/src/main/java/org/elasticsearch/action/support/IndicesOptions.java b/server/src/main/java/org/elasticsearch/action/support/IndicesOptions.java index 6e94ea11c652d..3b03b1cf0a4f6 100644 --- a/server/src/main/java/org/elasticsearch/action/support/IndicesOptions.java +++ b/server/src/main/java/org/elasticsearch/action/support/IndicesOptions.java @@ -8,8 +8,11 @@ package org.elasticsearch.action.support; import org.elasticsearch.ElasticsearchParseException; +import org.elasticsearch.TransportVersions; +import org.elasticsearch.cluster.metadata.DataStream; import org.elasticsearch.common.io.stream.StreamInput; import org.elasticsearch.common.io.stream.StreamOutput; +import org.elasticsearch.common.io.stream.Writeable; import org.elasticsearch.common.logging.DeprecationCategory; import org.elasticsearch.common.logging.DeprecationLogger; import org.elasticsearch.core.Nullable; @@ -40,12 +43,27 @@ * target does not exist. * @param wildcardOptions, applies only to wildcard expressions and defines how the wildcards will be expanded and if it will * be acceptable to have expressions that results to no indices. - * @param generalOptions, applies to all the resolved indices and defines if throttled will be included and if certain type of - * aliases or indices are allowed, or they will throw an error. + * @param gatekeeperOptions, applies to all the resolved indices and defines if throttled will be included and if certain type of + * aliases or indices are allowed, or they will throw an error. It acts as a gatekeeper when an action + * does not support certain options. + * @param failureStoreOptions, applies to all indices already matched and controls the type of indices that will be returned. Currently, + * there are two types, data stream failure indices (only certain data streams have them) and data stream + * backing indices or stand-alone indices. */ -public record IndicesOptions(ConcreteTargetOptions concreteTargetOptions, WildcardOptions wildcardOptions, GeneralOptions generalOptions) - implements - ToXContentFragment { +public record IndicesOptions( + ConcreteTargetOptions concreteTargetOptions, + WildcardOptions wildcardOptions, + GatekeeperOptions gatekeeperOptions, + FailureStoreOptions failureStoreOptions +) implements ToXContentFragment { + + public IndicesOptions( + ConcreteTargetOptions concreteTargetOptions, + WildcardOptions wildcardOptions, + GatekeeperOptions gatekeeperOptions + ) { + this(concreteTargetOptions, wildcardOptions, gatekeeperOptions, FailureStoreOptions.DEFAULT); + } public static IndicesOptions.Builder builder() { return new Builder(); @@ -286,20 +304,28 @@ public static Builder builder(WildcardOptions wildcardOptions) { } /** - * These options apply on all indices that have been selected by the other Options. It can either filter the response or - * define what type of indices or aliases are not allowed which will result in an error response. + * The "gatekeeper" options apply on all indices that have been selected by the other Options. It contains two type of flags: + * - The "allow*" flags, which purpose is to enable actions to define certain conditions that need to apply on the concrete indices + * they accept. For example, single-index actions will set allowAliasToMultipleIndices to false, while search will not accept a + * closed index etc. These options are not configurable by the end-user. + * - The ignoreThrottled flag, which is a depricared flag that will filter out frozen indices. * @param allowAliasToMultipleIndices, allow aliases to multiple indices, true by default. * @param allowClosedIndices, allow closed indices, true by default. - * @param ignoreThrottled, filters out throttled (aka frozen indices), defaults to true. + * @param allowFailureIndices, allow failure indices in the response, true by default + * @param ignoreThrottled, filters out throttled (aka frozen indices), defaults to true. This is deprecated and the only one + * that only filters and never throws an error. */ - public record GeneralOptions(boolean allowAliasToMultipleIndices, boolean allowClosedIndices, @Deprecated boolean ignoreThrottled) - implements - ToXContentFragment { + public record GatekeeperOptions( + boolean allowAliasToMultipleIndices, + boolean allowClosedIndices, + boolean allowFailureIndices, + @Deprecated boolean ignoreThrottled + ) implements ToXContentFragment { public static final String IGNORE_THROTTLED = "ignore_throttled"; - public static final GeneralOptions DEFAULT = new GeneralOptions(true, true, false); + public static final GatekeeperOptions DEFAULT = new GatekeeperOptions(true, true, true, false); - public static GeneralOptions parseParameter(Object ignoreThrottled, GeneralOptions defaultOptions) { + public static GatekeeperOptions parseParameter(Object ignoreThrottled, GatekeeperOptions defaultOptions) { if (ignoreThrottled == null && defaultOptions != null) { return defaultOptions; } @@ -316,15 +342,17 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws public static class Builder { private boolean allowAliasToMultipleIndices; private boolean allowClosedIndices; + private boolean allowFailureIndices; private boolean ignoreThrottled; public Builder() { this(DEFAULT); } - Builder(GeneralOptions options) { + Builder(GatekeeperOptions options) { allowAliasToMultipleIndices = options.allowAliasToMultipleIndices; allowClosedIndices = options.allowClosedIndices; + allowFailureIndices = options.allowFailureIndices; ignoreThrottled = options.ignoreThrottled; } @@ -346,6 +374,15 @@ public Builder allowClosedIndices(boolean allowClosedIndices) { return this; } + /** + * Failure indices are accepted when true, otherwise the resolution will throw an error. + * Defaults to true. + */ + public Builder allowFailureIndices(boolean allowFailureIndices) { + this.allowFailureIndices = allowFailureIndices; + return this; + } + /** * Throttled indices will not be included in the result. Defaults to false. */ @@ -354,8 +391,8 @@ public Builder ignoreThrottled(boolean ignoreThrottled) { return this; } - public GeneralOptions build() { - return new GeneralOptions(allowAliasToMultipleIndices, allowClosedIndices, ignoreThrottled); + public GatekeeperOptions build() { + return new GatekeeperOptions(allowAliasToMultipleIndices, allowClosedIndices, allowFailureIndices, ignoreThrottled); } } @@ -363,8 +400,102 @@ public static Builder builder() { return new Builder(); } - public static Builder builder(GeneralOptions generalOptions) { - return new Builder(generalOptions); + public static Builder builder(GatekeeperOptions gatekeeperOptions) { + return new Builder(gatekeeperOptions); + } + } + + /** + * Applies to all indices already matched and controls the type of indices that will be returned. There are two types, data stream + * failure indices (only certain data streams have them) and data stream backing indices or stand-alone indices. + * @param includeRegularIndices, when true regular or data stream backing indices will be retrieved. + * @param includeFailureIndices, when true data stream failure indices will be included. + */ + public record FailureStoreOptions(boolean includeRegularIndices, boolean includeFailureIndices) + implements + Writeable, + ToXContentFragment { + + public static final String FAILURE_STORE = "failure_store"; + public static final String INCLUDE_ALL = "true"; + public static final String INCLUDE_ONLY_REGULAR_INDICES = "false"; + public static final String INCLUDE_ONLY_FAILURE_INDICES = "only"; + + public static final FailureStoreOptions DEFAULT = new FailureStoreOptions(true, false); + + public static FailureStoreOptions read(StreamInput in) throws IOException { + return new FailureStoreOptions(in.readBoolean(), in.readBoolean()); + } + + public static FailureStoreOptions parseParameters(Object failureStoreValue, FailureStoreOptions defaultOptions) { + if (failureStoreValue == null) { + return defaultOptions; + } + FailureStoreOptions.Builder builder = defaultOptions == null + ? new FailureStoreOptions.Builder() + : new FailureStoreOptions.Builder(defaultOptions); + return switch (failureStoreValue.toString()) { + case INCLUDE_ALL -> builder.includeRegularIndices(true).includeFailureIndices(true).build(); + case INCLUDE_ONLY_REGULAR_INDICES -> builder.includeRegularIndices(true).includeFailureIndices(false).build(); + case INCLUDE_ONLY_FAILURE_INDICES -> builder.includeRegularIndices(false).includeFailureIndices(true).build(); + default -> throw new IllegalArgumentException("No valid " + FAILURE_STORE + " value [" + failureStoreValue + "]"); + }; + } + + @Override + public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { + return builder.field(FAILURE_STORE, displayValue()); + } + + public String displayValue() { + if (includeRegularIndices && includeFailureIndices) { + return INCLUDE_ALL; + } else if (includeRegularIndices) { + return INCLUDE_ONLY_REGULAR_INDICES; + } + return INCLUDE_ONLY_FAILURE_INDICES; + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + out.writeBoolean(includeRegularIndices); + out.writeBoolean(includeFailureIndices); + } + + public static class Builder { + private boolean includeRegularIndices; + private boolean includeFailureIndices; + + public Builder() { + this(DEFAULT); + } + + Builder(FailureStoreOptions options) { + includeRegularIndices = options.includeRegularIndices; + includeFailureIndices = options.includeFailureIndices; + } + + public Builder includeRegularIndices(boolean includeRegularIndices) { + this.includeRegularIndices = includeRegularIndices; + return this; + } + + public Builder includeFailureIndices(boolean includeFailureIndices) { + this.includeFailureIndices = includeFailureIndices; + return this; + } + + public FailureStoreOptions build() { + return new FailureStoreOptions(includeRegularIndices, includeFailureIndices); + } + } + + public static Builder builder() { + return new Builder(); + } + + public static Builder builder(FailureStoreOptions failureStoreOptions) { + return new Builder(failureStoreOptions); } } @@ -400,9 +531,10 @@ private enum Option { EXCLUDE_ALIASES, ALLOW_EMPTY_WILDCARD_EXPRESSIONS, ERROR_WHEN_ALIASES_TO_MULTIPLE_INDICES, - ERROR_WHEN_CLOSED_INDICES, - IGNORE_THROTTLED + IGNORE_THROTTLED, + + ALLOW_FAILURE_INDICES // Added in 8.14 } private static final DeprecationLogger DEPRECATION_LOGGER = DeprecationLogger.getLogger(IndicesOptions.class); @@ -415,7 +547,8 @@ private enum Option { public static final IndicesOptions DEFAULT = new IndicesOptions( ConcreteTargetOptions.ERROR_WHEN_UNAVAILABLE_TARGETS, WildcardOptions.DEFAULT, - GeneralOptions.DEFAULT + GatekeeperOptions.DEFAULT, + FailureStoreOptions.DEFAULT ); public static final IndicesOptions STRICT_EXPAND_OPEN = IndicesOptions.builder() @@ -428,7 +561,14 @@ private enum Option { .allowEmptyExpressions(true) .resolveAliases(true) ) - .generalOptions(GeneralOptions.builder().allowAliasToMultipleIndices(true).allowClosedIndices(true).ignoreThrottled(false)) + .gatekeeperOptions( + GatekeeperOptions.builder() + .allowAliasToMultipleIndices(true) + .allowClosedIndices(true) + .allowFailureIndices(true) + .ignoreThrottled(false) + ) + .failureStoreOptions(FailureStoreOptions.builder().includeRegularIndices(true).includeFailureIndices(false)) .build(); public static final IndicesOptions LENIENT_EXPAND_OPEN = IndicesOptions.builder() .concreteTargetOptions(ConcreteTargetOptions.ALLOW_UNAVAILABLE_TARGETS) @@ -440,7 +580,14 @@ private enum Option { .allowEmptyExpressions(true) .resolveAliases(true) ) - .generalOptions(GeneralOptions.builder().allowAliasToMultipleIndices(true).allowClosedIndices(true).ignoreThrottled(false)) + .gatekeeperOptions( + GatekeeperOptions.builder() + .allowAliasToMultipleIndices(true) + .allowClosedIndices(true) + .allowFailureIndices(true) + .ignoreThrottled(false) + ) + .failureStoreOptions(FailureStoreOptions.builder().includeRegularIndices(true).includeFailureIndices(false)) .build(); public static final IndicesOptions LENIENT_EXPAND_OPEN_HIDDEN = IndicesOptions.builder() .concreteTargetOptions(ConcreteTargetOptions.ALLOW_UNAVAILABLE_TARGETS) @@ -452,7 +599,14 @@ private enum Option { .allowEmptyExpressions(true) .resolveAliases(true) ) - .generalOptions(GeneralOptions.builder().allowAliasToMultipleIndices(true).allowClosedIndices(true).ignoreThrottled(false)) + .gatekeeperOptions( + GatekeeperOptions.builder() + .allowAliasToMultipleIndices(true) + .allowClosedIndices(true) + .allowFailureIndices(true) + .ignoreThrottled(false) + ) + .failureStoreOptions(FailureStoreOptions.builder().includeRegularIndices(true).includeFailureIndices(false)) .build(); public static final IndicesOptions LENIENT_EXPAND_OPEN_CLOSED = IndicesOptions.builder() .concreteTargetOptions(ConcreteTargetOptions.ALLOW_UNAVAILABLE_TARGETS) @@ -464,14 +618,28 @@ private enum Option { .allowEmptyExpressions(true) .resolveAliases(true) ) - .generalOptions(GeneralOptions.builder().allowAliasToMultipleIndices(true).allowClosedIndices(true).ignoreThrottled(false)) + .gatekeeperOptions( + GatekeeperOptions.builder() + .allowAliasToMultipleIndices(true) + .allowClosedIndices(true) + .allowFailureIndices(true) + .ignoreThrottled(false) + ) + .failureStoreOptions(FailureStoreOptions.builder().includeRegularIndices(true).includeFailureIndices(false)) .build(); public static final IndicesOptions LENIENT_EXPAND_OPEN_CLOSED_HIDDEN = IndicesOptions.builder() .concreteTargetOptions(ConcreteTargetOptions.ALLOW_UNAVAILABLE_TARGETS) .wildcardOptions( WildcardOptions.builder().matchOpen(true).matchClosed(true).includeHidden(true).allowEmptyExpressions(true).resolveAliases(true) ) - .generalOptions(GeneralOptions.builder().allowAliasToMultipleIndices(true).allowClosedIndices(true).ignoreThrottled(false)) + .gatekeeperOptions( + GatekeeperOptions.builder() + .allowAliasToMultipleIndices(true) + .allowClosedIndices(true) + .allowFailureIndices(true) + .ignoreThrottled(false) + ) + .failureStoreOptions(FailureStoreOptions.builder().includeRegularIndices(true).includeFailureIndices(false)) .build(); public static final IndicesOptions STRICT_EXPAND_OPEN_CLOSED = IndicesOptions.builder() .concreteTargetOptions(ConcreteTargetOptions.ERROR_WHEN_UNAVAILABLE_TARGETS) @@ -483,14 +651,28 @@ private enum Option { .allowEmptyExpressions(true) .resolveAliases(true) ) - .generalOptions(GeneralOptions.builder().allowAliasToMultipleIndices(true).allowClosedIndices(true).ignoreThrottled(false)) + .gatekeeperOptions( + GatekeeperOptions.builder() + .allowAliasToMultipleIndices(true) + .allowClosedIndices(true) + .allowFailureIndices(true) + .ignoreThrottled(false) + ) + .failureStoreOptions(FailureStoreOptions.builder().includeRegularIndices(true).includeFailureIndices(false)) .build(); public static final IndicesOptions STRICT_EXPAND_OPEN_CLOSED_HIDDEN = IndicesOptions.builder() .concreteTargetOptions(ConcreteTargetOptions.ERROR_WHEN_UNAVAILABLE_TARGETS) .wildcardOptions( WildcardOptions.builder().matchOpen(true).matchClosed(true).includeHidden(true).allowEmptyExpressions(true).resolveAliases(true) ) - .generalOptions(GeneralOptions.builder().allowAliasToMultipleIndices(true).allowClosedIndices(true).ignoreThrottled(false)) + .gatekeeperOptions( + GatekeeperOptions.builder() + .allowAliasToMultipleIndices(true) + .allowClosedIndices(true) + .allowFailureIndices(true) + .ignoreThrottled(false) + ) + .failureStoreOptions(FailureStoreOptions.builder().includeRegularIndices(true).includeFailureIndices(false)) .build(); public static final IndicesOptions STRICT_EXPAND_OPEN_FORBID_CLOSED = IndicesOptions.builder() .concreteTargetOptions(ConcreteTargetOptions.ERROR_WHEN_UNAVAILABLE_TARGETS) @@ -502,7 +684,14 @@ private enum Option { .allowEmptyExpressions(true) .resolveAliases(true) ) - .generalOptions(GeneralOptions.builder().allowClosedIndices(false).allowAliasToMultipleIndices(true).ignoreThrottled(false)) + .gatekeeperOptions( + GatekeeperOptions.builder() + .allowClosedIndices(false) + .allowAliasToMultipleIndices(true) + .allowFailureIndices(true) + .ignoreThrottled(false) + ) + .failureStoreOptions(FailureStoreOptions.builder().includeRegularIndices(true).includeFailureIndices(false)) .build(); public static final IndicesOptions STRICT_EXPAND_OPEN_HIDDEN_FORBID_CLOSED = IndicesOptions.builder() .concreteTargetOptions(ConcreteTargetOptions.ERROR_WHEN_UNAVAILABLE_TARGETS) @@ -514,7 +703,14 @@ private enum Option { .allowEmptyExpressions(true) .resolveAliases(true) ) - .generalOptions(GeneralOptions.builder().allowClosedIndices(false).allowAliasToMultipleIndices(true).ignoreThrottled(false)) + .gatekeeperOptions( + GatekeeperOptions.builder() + .allowClosedIndices(false) + .allowAliasToMultipleIndices(true) + .allowFailureIndices(true) + .ignoreThrottled(false) + ) + .failureStoreOptions(FailureStoreOptions.builder().includeRegularIndices(true).includeFailureIndices(false)) .build(); public static final IndicesOptions STRICT_EXPAND_OPEN_FORBID_CLOSED_IGNORE_THROTTLED = IndicesOptions.builder() .concreteTargetOptions(ConcreteTargetOptions.ERROR_WHEN_UNAVAILABLE_TARGETS) @@ -526,7 +722,14 @@ private enum Option { .allowEmptyExpressions(true) .resolveAliases(true) ) - .generalOptions(GeneralOptions.builder().ignoreThrottled(true).allowClosedIndices(false).allowAliasToMultipleIndices(true)) + .gatekeeperOptions( + GatekeeperOptions.builder() + .ignoreThrottled(true) + .allowClosedIndices(false) + .allowFailureIndices(true) + .allowAliasToMultipleIndices(true) + ) + .failureStoreOptions(FailureStoreOptions.builder().includeRegularIndices(true).includeFailureIndices(false)) .build(); public static final IndicesOptions STRICT_SINGLE_INDEX_NO_EXPAND_FORBID_CLOSED = IndicesOptions.builder() .concreteTargetOptions(ConcreteTargetOptions.ERROR_WHEN_UNAVAILABLE_TARGETS) @@ -538,7 +741,14 @@ private enum Option { .allowEmptyExpressions(true) .resolveAliases(true) ) - .generalOptions(GeneralOptions.builder().allowAliasToMultipleIndices(false).allowClosedIndices(false).ignoreThrottled(false)) + .gatekeeperOptions( + GatekeeperOptions.builder() + .allowAliasToMultipleIndices(false) + .allowClosedIndices(false) + .allowFailureIndices(true) + .ignoreThrottled(false) + ) + .failureStoreOptions(FailureStoreOptions.builder().includeRegularIndices(true).includeFailureIndices(false)) .build(); public static final IndicesOptions STRICT_NO_EXPAND_FORBID_CLOSED = IndicesOptions.builder() .concreteTargetOptions(ConcreteTargetOptions.ERROR_WHEN_UNAVAILABLE_TARGETS) @@ -550,7 +760,14 @@ private enum Option { .allowEmptyExpressions(true) .resolveAliases(true) ) - .generalOptions(GeneralOptions.builder().allowClosedIndices(false).allowAliasToMultipleIndices(true).ignoreThrottled(false)) + .gatekeeperOptions( + GatekeeperOptions.builder() + .allowClosedIndices(false) + .allowAliasToMultipleIndices(true) + .allowFailureIndices(true) + .ignoreThrottled(false) + ) + .failureStoreOptions(FailureStoreOptions.builder().includeRegularIndices(true).includeFailureIndices(false)) .build(); /** @@ -604,14 +821,21 @@ public boolean expandWildcardsHidden() { * @return Whether execution on closed indices is allowed. */ public boolean forbidClosedIndices() { - return generalOptions.allowClosedIndices() == false; + return gatekeeperOptions.allowClosedIndices() == false; + } + + /** + * @return Whether execution on closed indices is allowed. + */ + public boolean allowFailureIndices() { + return gatekeeperOptions.allowFailureIndices(); } /** * @return whether aliases pointing to multiple indices are allowed */ public boolean allowAliasesToMultipleIndices() { - return generalOptions().allowAliasToMultipleIndices(); + return gatekeeperOptions().allowAliasToMultipleIndices(); } /** @@ -625,7 +849,21 @@ public boolean ignoreAliases() { * @return whether indices that are marked as throttled should be ignored */ public boolean ignoreThrottled() { - return generalOptions().ignoreThrottled(); + return gatekeeperOptions().ignoreThrottled(); + } + + /** + * @return whether regular indices (stand-alone or backing indices) will be included in the response + */ + public boolean includeRegularIndices() { + return failureStoreOptions().includeRegularIndices(); + } + + /** + * @return whether failure indices (only supported by certain data streams) will be included in the response + */ + public boolean includeFailureIndices() { + return failureStoreOptions().includeFailureIndices(); } public void writeIndicesOptions(StreamOutput out) throws IOException { @@ -648,6 +886,11 @@ public void writeIndicesOptions(StreamOutput out) throws IOException { if (ignoreUnavailable()) { backwardsCompatibleOptions.add(Option.ALLOW_UNAVAILABLE_CONCRETE_TARGETS); } + if (out.getTransportVersion().onOrAfter(TransportVersions.ADD_FAILURE_STORE_INDICES_OPTIONS)) { + if (allowFailureIndices()) { + backwardsCompatibleOptions.add(Option.ALLOW_FAILURE_INDICES); + } + } out.writeEnumSet(backwardsCompatibleOptions); EnumSet states = EnumSet.noneOf(WildcardStates.class); @@ -661,6 +904,9 @@ public void writeIndicesOptions(StreamOutput out) throws IOException { states.add(WildcardStates.HIDDEN); } out.writeEnumSet(states); + if (out.getTransportVersion().onOrAfter(TransportVersions.ADD_FAILURE_STORE_INDICES_OPTIONS)) { + failureStoreOptions.writeTo(out); + } } public static IndicesOptions readIndicesOptions(StreamInput in) throws IOException { @@ -670,24 +916,34 @@ public static IndicesOptions readIndicesOptions(StreamInput in) throws IOExcepti options.contains(Option.ALLOW_EMPTY_WILDCARD_EXPRESSIONS), options.contains(Option.EXCLUDE_ALIASES) ); - GeneralOptions generalOptions = GeneralOptions.builder() + boolean allowFailureIndices = true; + if (in.getTransportVersion().onOrAfter(TransportVersions.ADD_FAILURE_STORE_INDICES_OPTIONS)) { + allowFailureIndices = options.contains(Option.ALLOW_FAILURE_INDICES); + } + GatekeeperOptions gatekeeperOptions = GatekeeperOptions.builder() .allowClosedIndices(options.contains(Option.ERROR_WHEN_CLOSED_INDICES) == false) .allowAliasToMultipleIndices(options.contains(Option.ERROR_WHEN_ALIASES_TO_MULTIPLE_INDICES) == false) + .allowFailureIndices(allowFailureIndices) .ignoreThrottled(options.contains(Option.IGNORE_THROTTLED)) .build(); + FailureStoreOptions failureStoreOptions = in.getTransportVersion().onOrAfter(TransportVersions.ADD_FAILURE_STORE_INDICES_OPTIONS) + ? FailureStoreOptions.read(in) + : FailureStoreOptions.DEFAULT; return new IndicesOptions( options.contains(Option.ALLOW_UNAVAILABLE_CONCRETE_TARGETS) ? ConcreteTargetOptions.ALLOW_UNAVAILABLE_TARGETS : ConcreteTargetOptions.ERROR_WHEN_UNAVAILABLE_TARGETS, wildcardOptions, - generalOptions + gatekeeperOptions, + failureStoreOptions ); } public static class Builder { private ConcreteTargetOptions concreteTargetOptions; private WildcardOptions wildcardOptions; - private GeneralOptions generalOptions; + private GatekeeperOptions gatekeeperOptions; + private FailureStoreOptions failureStoreOptions; Builder() { this(DEFAULT); @@ -696,7 +952,8 @@ public static class Builder { Builder(IndicesOptions indicesOptions) { concreteTargetOptions = indicesOptions.concreteTargetOptions; wildcardOptions = indicesOptions.wildcardOptions; - generalOptions = indicesOptions.generalOptions; + gatekeeperOptions = indicesOptions.gatekeeperOptions; + failureStoreOptions = indicesOptions.failureStoreOptions; } public Builder concreteTargetOptions(ConcreteTargetOptions concreteTargetOptions) { @@ -714,18 +971,28 @@ public Builder wildcardOptions(WildcardOptions.Builder wildcardOptions) { return this; } - public Builder generalOptions(GeneralOptions generalOptions) { - this.generalOptions = generalOptions; + public Builder gatekeeperOptions(GatekeeperOptions gatekeeperOptions) { + this.gatekeeperOptions = gatekeeperOptions; + return this; + } + + public Builder gatekeeperOptions(GatekeeperOptions.Builder generalOptions) { + this.gatekeeperOptions = generalOptions.build(); return this; } - public Builder generalOptions(GeneralOptions.Builder generalOptions) { - this.generalOptions = generalOptions.build(); + public Builder failureStoreOptions(FailureStoreOptions failureStoreOptions) { + this.failureStoreOptions = failureStoreOptions; + return this; + } + + public Builder failureStoreOptions(FailureStoreOptions.Builder failureStoreOptions) { + this.failureStoreOptions = failureStoreOptions.build(); return this; } public IndicesOptions build() { - return new IndicesOptions(concreteTargetOptions, wildcardOptions, generalOptions); + return new IndicesOptions(concreteTargetOptions, wildcardOptions, gatekeeperOptions, failureStoreOptions); } } @@ -819,7 +1086,7 @@ public static IndicesOptions fromOptions( .resolveAliases(ignoreAliases == false) .allowEmptyExpressions(allowNoIndices) .build(); - final GeneralOptions generalOptions = GeneralOptions.builder() + final GatekeeperOptions gatekeeperOptions = GatekeeperOptions.builder() .allowAliasToMultipleIndices(allowAliasesToMultipleIndices) .allowClosedIndices(forbidClosedIndices == false) .ignoreThrottled(ignoreThrottled) @@ -827,12 +1094,13 @@ public static IndicesOptions fromOptions( return new IndicesOptions( ignoreUnavailable ? ConcreteTargetOptions.ALLOW_UNAVAILABLE_TARGETS : ConcreteTargetOptions.ERROR_WHEN_UNAVAILABLE_TARGETS, wildcards, - generalOptions + gatekeeperOptions, + FailureStoreOptions.DEFAULT ); } public static IndicesOptions fromRequest(RestRequest request, IndicesOptions defaultSettings) { - if (request.hasParam(GeneralOptions.IGNORE_THROTTLED)) { + if (request.hasParam(GatekeeperOptions.IGNORE_THROTTLED)) { DEPRECATION_LOGGER.warn(DeprecationCategory.API, "ignore_throttled_param", IGNORE_THROTTLED_DEPRECATION_MESSAGE); } @@ -840,19 +1108,36 @@ public static IndicesOptions fromRequest(RestRequest request, IndicesOptions def request.param(WildcardOptions.EXPAND_WILDCARDS), request.param(ConcreteTargetOptions.IGNORE_UNAVAILABLE), request.param(WildcardOptions.ALLOW_NO_INDICES), - request.param(GeneralOptions.IGNORE_THROTTLED), + request.param(GatekeeperOptions.IGNORE_THROTTLED), + DataStream.isFailureStoreEnabled() + ? request.param(FailureStoreOptions.FAILURE_STORE) + : FailureStoreOptions.INCLUDE_ONLY_REGULAR_INDICES, defaultSettings ); } public static IndicesOptions fromMap(Map map, IndicesOptions defaultSettings) { + if (DataStream.isFailureStoreEnabled()) { + return fromParameters( + map.containsKey(WildcardOptions.EXPAND_WILDCARDS) ? map.get(WildcardOptions.EXPAND_WILDCARDS) : map.get("expandWildcards"), + map.containsKey(ConcreteTargetOptions.IGNORE_UNAVAILABLE) + ? map.get(ConcreteTargetOptions.IGNORE_UNAVAILABLE) + : map.get("ignoreUnavailable"), + map.containsKey(WildcardOptions.ALLOW_NO_INDICES) ? map.get(WildcardOptions.ALLOW_NO_INDICES) : map.get("allowNoIndices"), + map.containsKey(GatekeeperOptions.IGNORE_THROTTLED) + ? map.get(GatekeeperOptions.IGNORE_THROTTLED) + : map.get("ignoreThrottled"), + map.containsKey(FailureStoreOptions.FAILURE_STORE) ? map.get(FailureStoreOptions.FAILURE_STORE) : map.get("failureStore"), + defaultSettings + ); + } return fromParameters( map.containsKey(WildcardOptions.EXPAND_WILDCARDS) ? map.get(WildcardOptions.EXPAND_WILDCARDS) : map.get("expandWildcards"), map.containsKey(ConcreteTargetOptions.IGNORE_UNAVAILABLE) ? map.get(ConcreteTargetOptions.IGNORE_UNAVAILABLE) : map.get("ignoreUnavailable"), map.containsKey(WildcardOptions.ALLOW_NO_INDICES) ? map.get(WildcardOptions.ALLOW_NO_INDICES) : map.get("allowNoIndices"), - map.containsKey(GeneralOptions.IGNORE_THROTTLED) ? map.get(GeneralOptions.IGNORE_THROTTLED) : map.get("ignoreThrottled"), + map.containsKey(GatekeeperOptions.IGNORE_THROTTLED) ? map.get(GatekeeperOptions.IGNORE_THROTTLED) : map.get("ignoreThrottled"), defaultSettings ); } @@ -866,10 +1151,22 @@ public static boolean isIndicesOptions(String name) { || "expandWildcards".equals(name) || ConcreteTargetOptions.IGNORE_UNAVAILABLE.equals(name) || "ignoreUnavailable".equals(name) - || GeneralOptions.IGNORE_THROTTLED.equals(name) + || GatekeeperOptions.IGNORE_THROTTLED.equals(name) || "ignoreThrottled".equals(name) || WildcardOptions.ALLOW_NO_INDICES.equals(name) - || "allowNoIndices".equals(name); + || "allowNoIndices".equals(name) + || (DataStream.isFailureStoreEnabled() && FailureStoreOptions.FAILURE_STORE.equals(name)) + || (DataStream.isFailureStoreEnabled() && "failureStore".equals(name)); + } + + public static IndicesOptions fromParameters( + Object wildcardsString, + Object ignoreUnavailableString, + Object allowNoIndicesString, + Object ignoreThrottled, + IndicesOptions defaultSettings + ) { + return fromParameters(wildcardsString, ignoreUnavailableString, allowNoIndicesString, ignoreThrottled, null, defaultSettings); } public static IndicesOptions fromParameters( @@ -877,20 +1174,29 @@ public static IndicesOptions fromParameters( Object ignoreUnavailableString, Object allowNoIndicesString, Object ignoreThrottled, + Object failureStoreString, IndicesOptions defaultSettings ) { - if (wildcardsString == null && ignoreUnavailableString == null && allowNoIndicesString == null && ignoreThrottled == null) { + if (wildcardsString == null + && ignoreUnavailableString == null + && allowNoIndicesString == null + && ignoreThrottled == null + && failureStoreString == null) { return defaultSettings; } WildcardOptions wildcards = WildcardOptions.parseParameters(wildcardsString, allowNoIndicesString, defaultSettings.wildcardOptions); - GeneralOptions generalOptions = GeneralOptions.parseParameter(ignoreThrottled, defaultSettings.generalOptions); + GatekeeperOptions gatekeeperOptions = GatekeeperOptions.parseParameter(ignoreThrottled, defaultSettings.gatekeeperOptions); + FailureStoreOptions failureStoreOptions = DataStream.isFailureStoreEnabled() + ? FailureStoreOptions.parseParameters(failureStoreString, defaultSettings.failureStoreOptions) + : FailureStoreOptions.DEFAULT; // note that allowAliasesToMultipleIndices is not exposed, always true (only for internal use) return IndicesOptions.builder() .concreteTargetOptions(ConcreteTargetOptions.fromParameter(ignoreUnavailableString, defaultSettings.concreteTargetOptions)) .wildcardOptions(wildcards) - .generalOptions(generalOptions) + .gatekeeperOptions(gatekeeperOptions) + .failureStoreOptions(failureStoreOptions) .build(); } @@ -898,14 +1204,18 @@ public static IndicesOptions fromParameters( public XContentBuilder toXContent(XContentBuilder builder, ToXContent.Params params) throws IOException { concreteTargetOptions.toXContent(builder, params); wildcardOptions.toXContent(builder, params); - generalOptions.toXContent(builder, params); + gatekeeperOptions.toXContent(builder, params); + if (DataStream.isFailureStoreEnabled()) { + failureStoreOptions.toXContent(builder, params); + } return builder; } private static final ParseField EXPAND_WILDCARDS_FIELD = new ParseField(WildcardOptions.EXPAND_WILDCARDS); private static final ParseField IGNORE_UNAVAILABLE_FIELD = new ParseField(ConcreteTargetOptions.IGNORE_UNAVAILABLE); - private static final ParseField IGNORE_THROTTLED_FIELD = new ParseField(GeneralOptions.IGNORE_THROTTLED).withAllDeprecated(); + private static final ParseField IGNORE_THROTTLED_FIELD = new ParseField(GatekeeperOptions.IGNORE_THROTTLED).withAllDeprecated(); private static final ParseField ALLOW_NO_INDICES_FIELD = new ParseField(WildcardOptions.ALLOW_NO_INDICES); + private static final ParseField FAILURE_STORE_FIELD = new ParseField(FailureStoreOptions.FAILURE_STORE); public static IndicesOptions fromXContent(XContentParser parser) throws IOException { return fromXContent(parser, null); @@ -914,8 +1224,9 @@ public static IndicesOptions fromXContent(XContentParser parser) throws IOExcept public static IndicesOptions fromXContent(XContentParser parser, @Nullable IndicesOptions defaults) throws IOException { boolean parsedWildcardStates = false; WildcardOptions.Builder wildcards = defaults == null ? null : WildcardOptions.builder(defaults.wildcardOptions()); - GeneralOptions.Builder generalOptions = GeneralOptions.builder() - .ignoreThrottled(defaults != null && defaults.generalOptions().ignoreThrottled()); + GatekeeperOptions.Builder generalOptions = GatekeeperOptions.builder() + .ignoreThrottled(defaults != null && defaults.gatekeeperOptions().ignoreThrottled()); + FailureStoreOptions failureStoreOptions = defaults == null ? FailureStoreOptions.DEFAULT : defaults.failureStoreOptions(); Boolean allowNoIndices = defaults == null ? null : defaults.allowNoIndices(); Boolean ignoreUnavailable = defaults == null ? null : defaults.ignoreUnavailable(); Token token = parser.currentToken() == Token.START_OBJECT ? parser.currentToken() : parser.nextToken(); @@ -965,13 +1276,16 @@ public static IndicesOptions fromXContent(XContentParser parser, @Nullable Indic allowNoIndices = parser.booleanValue(); } else if (IGNORE_THROTTLED_FIELD.match(currentFieldName, parser.getDeprecationHandler())) { generalOptions.ignoreThrottled(parser.booleanValue()); - } else { - throw new ElasticsearchParseException( - "could not read indices options. unexpected index option [" + currentFieldName + "]" - ); - } + } else if (DataStream.isFailureStoreEnabled() + && FAILURE_STORE_FIELD.match(currentFieldName, parser.getDeprecationHandler())) { + failureStoreOptions = FailureStoreOptions.parseParameters(parser.text(), failureStoreOptions); + } else { + throw new ElasticsearchParseException( + "could not read indices options. Unexpected index option [" + currentFieldName + "]" + ); + } } else { - throw new ElasticsearchParseException("could not read indices options. unexpected object field [" + currentFieldName + "]"); + throw new ElasticsearchParseException("could not read indices options. Unexpected object field [" + currentFieldName + "]"); } } @@ -994,7 +1308,8 @@ public static IndicesOptions fromXContent(XContentParser parser, @Nullable Indic return IndicesOptions.builder() .concreteTargetOptions(new ConcreteTargetOptions(ignoreUnavailable)) .wildcardOptions(wildcards) - .generalOptions(generalOptions) + .gatekeeperOptions(generalOptions) + .failureStoreOptions(failureStoreOptions) .build(); } @@ -1108,6 +1423,14 @@ public String toString() { + ignoreAliases() + ", ignore_throttled=" + ignoreThrottled() + + (DataStream.isFailureStoreEnabled() + ? ", include_regular_indices=" + + includeRegularIndices() + + ", include_failure_indices=" + + includeFailureIndices() + + ", allow_failure_indices=" + + allowFailureIndices() + : "") + ']'; } } diff --git a/server/src/main/java/org/elasticsearch/action/support/master/info/ClusterInfoRequest.java b/server/src/main/java/org/elasticsearch/action/support/master/info/ClusterInfoRequest.java index 22f0da70137af..00384852d1472 100644 --- a/server/src/main/java/org/elasticsearch/action/support/master/info/ClusterInfoRequest.java +++ b/server/src/main/java/org/elasticsearch/action/support/master/info/ClusterInfoRequest.java @@ -28,6 +28,11 @@ public abstract class ClusterInfoRequest failureIndices; + private volatile Set failureStoreLookup; @Nullable private final DataStreamAutoShardingEvent autoShardingEvent; @@ -282,6 +283,32 @@ public Index getWriteIndex() { return indices.get(indices.size() - 1); } + /** + * @return the write failure index if the failure store is enabled and there is already at least one failure, null otherwise + */ + @Nullable + public Index getFailureStoreWriteIndex() { + return isFailureStore() == false || failureIndices.isEmpty() ? null : failureIndices.get(failureIndices.size() - 1); + } + + /** + * Returns true if the index name provided belongs to a failure store index. + * This method builds a local Set with all the failure store index names and then checks if it contains the name. + * This will perform better if there are multiple indices of this data stream checked. + */ + public boolean isFailureStoreIndex(String indexName) { + if (failureStoreLookup == null) { + // There is a chance this will be calculated twice, but it's a relatively cheap action, + // so it's not worth synchronising + if (failureIndices == null || failureIndices.isEmpty()) { + failureStoreLookup = Set.of(); + } else { + failureStoreLookup = failureIndices.stream().map(Index::getName).collect(Collectors.toSet()); + } + } + return failureStoreLookup.contains(indexName); + } + public boolean rolloverOnWrite() { return rolloverOnWrite; } diff --git a/server/src/main/java/org/elasticsearch/cluster/metadata/IndexNameExpressionResolver.java b/server/src/main/java/org/elasticsearch/cluster/metadata/IndexNameExpressionResolver.java index 4c3318d8d2f6a..e8e8ca767cc34 100644 --- a/server/src/main/java/org/elasticsearch/cluster/metadata/IndexNameExpressionResolver.java +++ b/server/src/main/java/org/elasticsearch/cluster/metadata/IndexNameExpressionResolver.java @@ -31,6 +31,7 @@ import org.elasticsearch.index.IndexNotFoundException; import org.elasticsearch.index.IndexVersion; import org.elasticsearch.index.IndexVersions; +import org.elasticsearch.indices.FailureIndexNotSupportedException; import org.elasticsearch.indices.IndexClosedException; import org.elasticsearch.indices.InvalidIndexNameException; import org.elasticsearch.indices.SystemIndices; @@ -354,16 +355,19 @@ Index[] concreteIndices(Context context, String... indexExpressions) { + " indices without one being designated as a write index" ); } - if (addIndex(writeIndex, null, context)) { - concreteIndicesResult.add(writeIndex); + if (indexAbstraction.isDataStreamRelated()) { + DataStream dataStream = indicesLookup.get(indexAbstraction.getWriteIndex().getName()).getParentDataStream(); + resolveWriteIndexForDataStreams(context, dataStream, concreteIndicesResult); + } else { + if (addIndex(writeIndex, null, context)) { + concreteIndicesResult.add(writeIndex); + } } } else if (indexAbstraction.getType() == Type.DATA_STREAM && context.isResolveToWriteIndex()) { - Index writeIndex = indexAbstraction.getWriteIndex(); - if (addIndex(writeIndex, null, context)) { - concreteIndicesResult.add(writeIndex); - } + resolveWriteIndexForDataStreams(context, (DataStream) indexAbstraction, concreteIndicesResult); } else { - if (indexAbstraction.getIndices().size() > 1 && context.getOptions().allowAliasesToMultipleIndices() == false) { + if (resolvesToMoreThanOneIndex(indexAbstraction, context) + && context.getOptions().allowAliasesToMultipleIndices() == false) { String[] indexNames = new String[indexAbstraction.getIndices().size()]; int i = 0; for (Index indexName : indexAbstraction.getIndices()) { @@ -379,11 +383,27 @@ Index[] concreteIndices(Context context, String... indexExpressions) { ); } - for (Index index : indexAbstraction.getIndices()) { - if (shouldTrackConcreteIndex(context, context.getOptions(), index)) { - concreteIndicesResult.add(index); + if (indexAbstraction.getType() == Type.DATA_STREAM) { + resolveIndicesForDataStream(context, (DataStream) indexAbstraction, concreteIndicesResult); + } else if (indexAbstraction.getType() == Type.ALIAS + && indexAbstraction.isDataStreamRelated() + && DataStream.isFailureStoreEnabled() + && context.getOptions().includeFailureIndices()) { + // Collect the data streams involved + Set aliasDataStreams = new HashSet<>(); + for (Index index : indexAbstraction.getIndices()) { + aliasDataStreams.add(indicesLookup.get(index.getName()).getParentDataStream()); + } + for (DataStream dataStream : aliasDataStreams) { + resolveIndicesForDataStream(context, dataStream, concreteIndicesResult); + } + } else { + for (Index index : indexAbstraction.getIndices()) { + if (shouldTrackConcreteIndex(context, context.getOptions(), index)) { + concreteIndicesResult.add(index); + } + } } - } } } @@ -394,6 +414,67 @@ Index[] concreteIndices(Context context, String... indexExpressions) { return concreteIndicesResult.toArray(Index.EMPTY_ARRAY); } + private static void resolveIndicesForDataStream(Context context, DataStream dataStream, Set concreteIndicesResult) { + if (shouldIncludeRegularIndices(context.getOptions())) { + for (Index index : dataStream.getIndices()) { + if (shouldTrackConcreteIndex(context, context.getOptions(), index)) { + concreteIndicesResult.add(index); + } + } + } + if (shouldIncludeFailureIndices(context.getOptions(), dataStream)) { + // We short-circuit here, if failure indices are not allowed and they can be skipped + if (context.getOptions().allowFailureIndices() || context.getOptions().ignoreUnavailable() == false) { + for (Index index : dataStream.getFailureIndices()) { + if (shouldTrackConcreteIndex(context, context.getOptions(), index)) { + concreteIndicesResult.add(index); + } + } + } + } + } + + private static void resolveWriteIndexForDataStreams(Context context, DataStream dataStream, Set concreteIndicesResult) { + if (shouldIncludeRegularIndices(context.getOptions())) { + Index writeIndex = dataStream.getWriteIndex(); + if (addIndex(writeIndex, null, context)) { + concreteIndicesResult.add(writeIndex); + } + } + if (shouldIncludeFailureIndices(context.getOptions(), dataStream)) { + Index failureStoreWriteIndex = dataStream.getFailureStoreWriteIndex(); + if (failureStoreWriteIndex != null && addIndex(failureStoreWriteIndex, null, context)) { + if (context.options.allowFailureIndices() == false) { + throw new FailureIndexNotSupportedException(failureStoreWriteIndex); + } + concreteIndicesResult.add(failureStoreWriteIndex); + } + } + } + + private static boolean shouldIncludeRegularIndices(IndicesOptions indicesOptions) { + return DataStream.isFailureStoreEnabled() == false || indicesOptions.includeRegularIndices(); + } + + private static boolean shouldIncludeFailureIndices(IndicesOptions indicesOptions, DataStream dataStream) { + return DataStream.isFailureStoreEnabled() && indicesOptions.includeFailureIndices() && dataStream.isFailureStore(); + } + + private static boolean resolvesToMoreThanOneIndex(IndexAbstraction indexAbstraction, Context context) { + if (indexAbstraction.getType() == Type.DATA_STREAM) { + DataStream dataStream = (DataStream) indexAbstraction; + int count = 0; + if (shouldIncludeRegularIndices(context.getOptions())) { + count += dataStream.getIndices().size(); + } + if (shouldIncludeFailureIndices(context.getOptions(), dataStream)) { + count += dataStream.getFailureIndices().size(); + } + return count > 1; + } + return indexAbstraction.getIndices().size() > 1; + } + private void checkSystemIndexAccess(Context context, Set concreteIndices) { final Predicate systemIndexAccessPredicate = context.getSystemIndexAccessPredicate(); if (systemIndexAccessPredicate == Predicates.always()) { @@ -485,6 +566,21 @@ private static boolean shouldTrackConcreteIndex(Context context, IndicesOptions // Exclude this one as it's a net-new system index, and we explicitly don't want those. return false; } + if (DataStream.isFailureStoreEnabled()) { + IndexAbstraction indexAbstraction = context.getState().metadata().getIndicesLookup().get(index.getName()); + if (context.options.allowFailureIndices() == false) { + DataStream parentDataStream = indexAbstraction.getParentDataStream(); + if (parentDataStream != null && parentDataStream.isFailureStore()) { + if (parentDataStream.isFailureStoreIndex(index.getName())) { + if (options.ignoreUnavailable()) { + return false; + } else { + throw new FailureIndexNotSupportedException(index); + } + } + } + } + } final IndexMetadata imd = context.state.metadata().index(index); if (imd.getState() == IndexMetadata.State.CLOSE) { if (options.forbidClosedIndices() && options.ignoreUnavailable() == false) { @@ -1309,7 +1405,7 @@ private static Map filterIndicesLookupForSuffixWildcar /** * Return the {@code Stream} of open and/or closed index names for the given {@param resources}. - * Datastreams and aliases are interpreted to refer to multiple indices, + * Data streams and aliases are interpreted to refer to multiple indices, * then all index resources are filtered by their open/closed status. */ private static Stream expandToOpenClosed(Context context, Stream resources) { @@ -1320,7 +1416,18 @@ private static Stream expandToOpenClosed(Context context, Stream indicesStateStream = indexAbstraction.getIndices().stream().map(context.state.metadata()::index); + Stream indicesStateStream = Stream.of(); + if (shouldIncludeRegularIndices(context.getOptions())) { + indicesStateStream = indexAbstraction.getIndices().stream().map(context.state.metadata()::index); + } + if (indexAbstraction.getType() == Type.DATA_STREAM + && shouldIncludeFailureIndices(context.getOptions(), (DataStream) indexAbstraction)) { + DataStream dataStream = (DataStream) indexAbstraction; + indicesStateStream = Stream.concat( + indicesStateStream, + dataStream.getFailureIndices().stream().map(context.state.metadata()::index) + ); + } if (excludeState != null) { indicesStateStream = indicesStateStream.filter(indexMeta -> indexMeta.getState() != excludeState); } @@ -1362,6 +1469,9 @@ private static List resolveEmptyOrTrivialWildcardWithAllowedSystemIndice } private static String[] resolveEmptyOrTrivialWildcardToAllIndices(IndicesOptions options, Metadata metadata) { + if (shouldIncludeRegularIndices(options) == false) { + return Strings.EMPTY_ARRAY; + } if (options.expandWildcardsOpen() && options.expandWildcardsClosed() && options.expandWildcardsHidden()) { return metadata.getConcreteAllIndices(); } else if (options.expandWildcardsOpen() && options.expandWildcardsClosed()) { diff --git a/server/src/main/java/org/elasticsearch/cluster/metadata/Metadata.java b/server/src/main/java/org/elasticsearch/cluster/metadata/Metadata.java index 4d76ead90e12a..b450251ff7e3f 100644 --- a/server/src/main/java/org/elasticsearch/cluster/metadata/Metadata.java +++ b/server/src/main/java/org/elasticsearch/cluster/metadata/Metadata.java @@ -2597,6 +2597,9 @@ private static void collectIndices( private static boolean assertContainsIndexIfDataStream(DataStream parent, IndexMetadata indexMetadata) { assert parent == null || parent.getIndices().stream().anyMatch(index -> indexMetadata.getIndex().getName().equals(index.getName())) + || (DataStream.isFailureStoreEnabled() + && parent.isFailureStore() + && parent.getFailureIndices().stream().anyMatch(index -> indexMetadata.getIndex().getName().equals(index.getName()))) : "Expected data stream [" + parent.getName() + "] to contain index " + indexMetadata.getIndex(); return true; } @@ -2618,6 +2621,11 @@ private static void collectDataStreams( for (Index i : dataStream.getIndices()) { indexToDataStreamLookup.put(i.getName(), dataStream); } + if (DataStream.isFailureStoreEnabled() && dataStream.isFailureStore()) { + for (Index i : dataStream.getFailureIndices()) { + indexToDataStreamLookup.put(i.getName(), dataStream); + } + } } } diff --git a/server/src/main/java/org/elasticsearch/indices/FailureIndexNotSupportedException.java b/server/src/main/java/org/elasticsearch/indices/FailureIndexNotSupportedException.java new file mode 100644 index 0000000000000..90fdd364b7035 --- /dev/null +++ b/server/src/main/java/org/elasticsearch/indices/FailureIndexNotSupportedException.java @@ -0,0 +1,37 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0 and the Server Side Public License, v 1; you may not use this file except + * in compliance with, at your election, the Elastic License 2.0 or the Server + * Side Public License, v 1. + */ + +package org.elasticsearch.indices; + +import org.elasticsearch.ElasticsearchException; +import org.elasticsearch.common.io.stream.StreamInput; +import org.elasticsearch.index.Index; +import org.elasticsearch.rest.RestStatus; + +import java.io.IOException; + +/** + * Exception indicating that one or more requested indices are failure indices. + */ +public final class FailureIndexNotSupportedException extends ElasticsearchException { + + public FailureIndexNotSupportedException(Index index) { + super("failure index not supported"); + setIndex(index); + } + + public FailureIndexNotSupportedException(StreamInput in) throws IOException { + super(in); + } + + @Override + public RestStatus status() { + return RestStatus.BAD_REQUEST; + } + +} diff --git a/server/src/test/java/org/elasticsearch/ExceptionSerializationTests.java b/server/src/test/java/org/elasticsearch/ExceptionSerializationTests.java index 1265a4e7f96db..b8091b50b5dd8 100644 --- a/server/src/test/java/org/elasticsearch/ExceptionSerializationTests.java +++ b/server/src/test/java/org/elasticsearch/ExceptionSerializationTests.java @@ -60,6 +60,7 @@ import org.elasticsearch.index.shard.ShardId; import org.elasticsearch.index.shard.ShardNotInPrimaryModeException; import org.elasticsearch.indices.AutoscalingMissedIndicesUpdateException; +import org.elasticsearch.indices.FailureIndexNotSupportedException; import org.elasticsearch.indices.IndexTemplateMissingException; import org.elasticsearch.indices.InvalidIndexTemplateException; import org.elasticsearch.indices.recovery.PeerRecoveryNotFound; @@ -827,6 +828,7 @@ public void testIds() { ids.put(175, AutoscalingMissedIndicesUpdateException.class); ids.put(176, SearchTimeoutException.class); ids.put(177, GraphStructureException.class); + ids.put(178, FailureIndexNotSupportedException.class); Map, Integer> reverse = new HashMap<>(); for (Map.Entry> entry : ids.entrySet()) { diff --git a/server/src/test/java/org/elasticsearch/action/admin/indices/get/GetIndexRequestTests.java b/server/src/test/java/org/elasticsearch/action/admin/indices/get/GetIndexRequestTests.java index 21c5d0bee47e9..d77be7c45e416 100644 --- a/server/src/test/java/org/elasticsearch/action/admin/indices/get/GetIndexRequestTests.java +++ b/server/src/test/java/org/elasticsearch/action/admin/indices/get/GetIndexRequestTests.java @@ -8,6 +8,7 @@ package org.elasticsearch.action.admin.indices.get; +import org.elasticsearch.action.support.IndicesOptions; import org.elasticsearch.core.Strings; import org.elasticsearch.rest.RestRequest; import org.elasticsearch.rest.RestRequestTests; @@ -71,4 +72,18 @@ public void testInvalidFeatures() { IllegalArgumentException e = expectThrows(IllegalArgumentException.class, () -> GetIndexRequest.Feature.fromRequest(request)); assertThat(e.getMessage(), containsString(Strings.format("Invalid features specified [%s]", String.join(",", invalidFeatures)))); } + + public void testIndicesOptions() { + GetIndexRequest getIndexRequest = new GetIndexRequest(); + assertThat( + getIndexRequest.indicesOptions().concreteTargetOptions(), + equalTo(IndicesOptions.strictExpandOpen().concreteTargetOptions()) + ); + assertThat(getIndexRequest.indicesOptions().wildcardOptions(), equalTo(IndicesOptions.strictExpandOpen().wildcardOptions())); + assertThat(getIndexRequest.indicesOptions().gatekeeperOptions(), equalTo(IndicesOptions.strictExpandOpen().gatekeeperOptions())); + assertThat( + getIndexRequest.indicesOptions().failureStoreOptions(), + equalTo(IndicesOptions.FailureStoreOptions.builder().includeRegularIndices(true).includeFailureIndices(true).build()) + ); + } } diff --git a/server/src/test/java/org/elasticsearch/action/support/IndicesOptionsTests.java b/server/src/test/java/org/elasticsearch/action/support/IndicesOptionsTests.java index deec53de59326..297ebbae6c85a 100644 --- a/server/src/test/java/org/elasticsearch/action/support/IndicesOptionsTests.java +++ b/server/src/test/java/org/elasticsearch/action/support/IndicesOptionsTests.java @@ -9,7 +9,8 @@ package org.elasticsearch.action.support; import org.elasticsearch.action.support.IndicesOptions.ConcreteTargetOptions; -import org.elasticsearch.action.support.IndicesOptions.GeneralOptions; +import org.elasticsearch.action.support.IndicesOptions.FailureStoreOptions; +import org.elasticsearch.action.support.IndicesOptions.GatekeeperOptions; import org.elasticsearch.action.support.IndicesOptions.WildcardOptions; import org.elasticsearch.common.bytes.BytesReference; import org.elasticsearch.common.io.stream.BytesStreamOutput; @@ -40,17 +41,25 @@ public class IndicesOptionsTests extends ESTestCase { public void testSerialization() throws Exception { int iterations = randomIntBetween(5, 20); for (int i = 0; i < iterations; i++) { - IndicesOptions indicesOptions = IndicesOptions.fromOptions( - randomBoolean(), - randomBoolean(), - randomBoolean(), - randomBoolean(), - randomBoolean(), - randomBoolean(), - randomBoolean(), - randomBoolean(), - randomBoolean() - ); + IndicesOptions indicesOptions = IndicesOptions.builder() + .wildcardOptions( + WildcardOptions.builder() + .matchOpen(randomBoolean()) + .matchClosed(randomBoolean()) + .includeHidden(randomBoolean()) + .allowEmptyExpressions(randomBoolean()) + .resolveAliases(randomBoolean()) + ) + .gatekeeperOptions( + GatekeeperOptions.builder() + .ignoreThrottled(randomBoolean()) + .allowAliasToMultipleIndices(randomBoolean()) + .allowClosedIndices(randomBoolean()) + ) + .failureStoreOptions( + FailureStoreOptions.builder().includeRegularIndices(randomBoolean()).includeFailureIndices(randomBoolean()) + ) + .build(); BytesStreamOutput output = new BytesStreamOutput(); indicesOptions.writeIndicesOptions(output); @@ -58,16 +67,7 @@ public void testSerialization() throws Exception { StreamInput streamInput = output.bytes().streamInput(); IndicesOptions indicesOptions2 = IndicesOptions.readIndicesOptions(streamInput); - assertThat(indicesOptions2.ignoreUnavailable(), equalTo(indicesOptions.ignoreUnavailable())); - assertThat(indicesOptions2.allowNoIndices(), equalTo(indicesOptions.allowNoIndices())); - assertThat(indicesOptions2.expandWildcardsOpen(), equalTo(indicesOptions.expandWildcardsOpen())); - assertThat(indicesOptions2.expandWildcardsClosed(), equalTo(indicesOptions.expandWildcardsClosed())); - assertThat(indicesOptions2.expandWildcardsHidden(), equalTo(indicesOptions.expandWildcardsHidden())); - - assertThat(indicesOptions2.forbidClosedIndices(), equalTo(indicesOptions.forbidClosedIndices())); - assertThat(indicesOptions2.allowAliasesToMultipleIndices(), equalTo(indicesOptions.allowAliasesToMultipleIndices())); - - assertEquals(indicesOptions2.ignoreAliases(), indicesOptions.ignoreAliases()); + assertThat(indicesOptions2, equalTo(indicesOptions)); } } @@ -343,9 +343,10 @@ public void testToXContent() throws IOException { randomBoolean(), randomBoolean() ); - GeneralOptions generalOptions = new GeneralOptions(randomBoolean(), randomBoolean(), randomBoolean()); + GatekeeperOptions gatekeeperOptions = new GatekeeperOptions(randomBoolean(), randomBoolean(), randomBoolean(), randomBoolean()); + FailureStoreOptions failureStoreOptions = new IndicesOptions.FailureStoreOptions(randomBoolean(), randomBoolean()); - IndicesOptions indicesOptions = new IndicesOptions(concreteTargetOptions, wildcardOptions, generalOptions); + IndicesOptions indicesOptions = new IndicesOptions(concreteTargetOptions, wildcardOptions, gatekeeperOptions, failureStoreOptions); XContentType type = randomFrom(XContentType.values()); BytesReference xContentBytes = toXContentBytes(indicesOptions, type); @@ -359,7 +360,8 @@ public void testToXContent() throws IOException { assertThat(((List) map.get("expand_wildcards")).contains("hidden"), equalTo(wildcardOptions.includeHidden())); assertThat(map.get("ignore_unavailable"), equalTo(concreteTargetOptions.allowUnavailableTargets())); assertThat(map.get("allow_no_indices"), equalTo(wildcardOptions.allowEmptyExpressions())); - assertThat(map.get("ignore_throttled"), equalTo(generalOptions.ignoreThrottled())); + assertThat(map.get("ignore_throttled"), equalTo(gatekeeperOptions.ignoreThrottled())); + assertThat(map.get("failure_store"), equalTo(failureStoreOptions.displayValue())); } public void testFromXContent() throws IOException { diff --git a/server/src/test/java/org/elasticsearch/cluster/metadata/DataStreamTests.java b/server/src/test/java/org/elasticsearch/cluster/metadata/DataStreamTests.java index 7e8e9805b54e7..8b6a0fcb55c5b 100644 --- a/server/src/test/java/org/elasticsearch/cluster/metadata/DataStreamTests.java +++ b/server/src/test/java/org/elasticsearch/cluster/metadata/DataStreamTests.java @@ -1781,6 +1781,164 @@ private IndexMetadata createIndexMetadata(String indexName, IndexWriteLoad index .build(); } + public void testWriteFailureIndex() { + boolean hidden = randomBoolean(); + boolean system = hidden && randomBoolean(); + DataStream noFailureStoreDataStream = new DataStream( + randomAlphaOfLength(10), + randomIndexInstances(), + randomNonNegativeInt(), + null, + hidden, + randomBoolean(), + system, + System::currentTimeMillis, + randomBoolean(), + randomBoolean() ? IndexMode.STANDARD : IndexMode.TIME_SERIES, + DataStreamLifecycleTests.randomLifecycle(), + false, + null, + randomBoolean(), + null + ); + assertThat(noFailureStoreDataStream.getFailureStoreWriteIndex(), nullValue()); + + DataStream failureStoreDataStreamWithEmptyFailureIndices = new DataStream( + randomAlphaOfLength(10), + randomIndexInstances(), + randomNonNegativeInt(), + null, + hidden, + randomBoolean(), + system, + System::currentTimeMillis, + randomBoolean(), + randomBoolean() ? IndexMode.STANDARD : IndexMode.TIME_SERIES, + DataStreamLifecycleTests.randomLifecycle(), + true, + List.of(), + randomBoolean(), + null + ); + assertThat(failureStoreDataStreamWithEmptyFailureIndices.getFailureStoreWriteIndex(), nullValue()); + + List failureIndices = randomIndexInstances(); + String dataStreamName = randomAlphaOfLength(10); + Index writeFailureIndex = new Index( + getDefaultBackingIndexName(dataStreamName, randomNonNegativeInt()), + UUIDs.randomBase64UUID(LuceneTestCase.random()) + ); + failureIndices.add(writeFailureIndex); + DataStream failureStoreDataStream = new DataStream( + dataStreamName, + randomIndexInstances(), + randomNonNegativeInt(), + null, + hidden, + randomBoolean(), + system, + System::currentTimeMillis, + randomBoolean(), + randomBoolean() ? IndexMode.STANDARD : IndexMode.TIME_SERIES, + DataStreamLifecycleTests.randomLifecycle(), + true, + failureIndices, + randomBoolean(), + null + ); + assertThat(failureStoreDataStream.getFailureStoreWriteIndex(), is(writeFailureIndex)); + } + + public void testIsFailureIndex() { + boolean hidden = randomBoolean(); + boolean system = hidden && randomBoolean(); + List backingIndices = randomIndexInstances(); + DataStream noFailureStoreDataStream = new DataStream( + randomAlphaOfLength(10), + backingIndices, + randomNonNegativeInt(), + null, + hidden, + randomBoolean(), + system, + System::currentTimeMillis, + randomBoolean(), + randomBoolean() ? IndexMode.STANDARD : IndexMode.TIME_SERIES, + DataStreamLifecycleTests.randomLifecycle(), + false, + null, + randomBoolean(), + null + ); + assertThat( + noFailureStoreDataStream.isFailureStoreIndex(backingIndices.get(randomIntBetween(0, backingIndices.size() - 1)).getName()), + is(false) + ); + + backingIndices = randomIndexInstances(); + DataStream failureStoreDataStreamWithEmptyFailureIndices = new DataStream( + randomAlphaOfLength(10), + backingIndices, + randomNonNegativeInt(), + null, + hidden, + randomBoolean(), + system, + System::currentTimeMillis, + randomBoolean(), + randomBoolean() ? IndexMode.STANDARD : IndexMode.TIME_SERIES, + DataStreamLifecycleTests.randomLifecycle(), + true, + List.of(), + randomBoolean(), + null + ); + assertThat( + failureStoreDataStreamWithEmptyFailureIndices.isFailureStoreIndex( + backingIndices.get(randomIntBetween(0, backingIndices.size() - 1)).getName() + ), + is(false) + ); + + backingIndices = randomIndexInstances(); + List failureIndices = randomIndexInstances(); + String dataStreamName = randomAlphaOfLength(10); + Index writeFailureIndex = new Index( + getDefaultBackingIndexName(dataStreamName, randomNonNegativeInt()), + UUIDs.randomBase64UUID(LuceneTestCase.random()) + ); + failureIndices.add(writeFailureIndex); + DataStream failureStoreDataStream = new DataStream( + dataStreamName, + backingIndices, + randomNonNegativeInt(), + null, + hidden, + randomBoolean(), + system, + System::currentTimeMillis, + randomBoolean(), + randomBoolean() ? IndexMode.STANDARD : IndexMode.TIME_SERIES, + DataStreamLifecycleTests.randomLifecycle(), + true, + failureIndices, + randomBoolean(), + null + ); + assertThat(failureStoreDataStream.isFailureStoreIndex(writeFailureIndex.getName()), is(true)); + assertThat( + failureStoreDataStream.isFailureStoreIndex(failureIndices.get(randomIntBetween(0, failureIndices.size() - 1)).getName()), + is(true) + ); + assertThat( + failureStoreDataStreamWithEmptyFailureIndices.isFailureStoreIndex( + backingIndices.get(randomIntBetween(0, backingIndices.size() - 1)).getName() + ), + is(false) + ); + assertThat(failureStoreDataStreamWithEmptyFailureIndices.isFailureStoreIndex(randomAlphaOfLength(10)), is(false)); + } + private record DataStreamMetadata(Long creationTimeInMillis, Long rolloverTimeInMillis, Long originationTimeInMillis) { public static DataStreamMetadata dataStreamMetadata(Long creationTimeInMillis, Long rolloverTimeInMillis) { return new DataStreamMetadata(creationTimeInMillis, rolloverTimeInMillis, null); diff --git a/server/src/test/java/org/elasticsearch/cluster/metadata/IndexNameExpressionResolverTests.java b/server/src/test/java/org/elasticsearch/cluster/metadata/IndexNameExpressionResolverTests.java index c043734d15cdf..a1eeceba8a390 100644 --- a/server/src/test/java/org/elasticsearch/cluster/metadata/IndexNameExpressionResolverTests.java +++ b/server/src/test/java/org/elasticsearch/cluster/metadata/IndexNameExpressionResolverTests.java @@ -29,6 +29,7 @@ import org.elasticsearch.index.IndexNotFoundException; import org.elasticsearch.index.IndexSettings; import org.elasticsearch.index.IndexVersion; +import org.elasticsearch.indices.FailureIndexNotSupportedException; import org.elasticsearch.indices.IndexClosedException; import org.elasticsearch.indices.InvalidIndexNameException; import org.elasticsearch.indices.SystemIndexDescriptor; @@ -53,6 +54,7 @@ import static org.elasticsearch.cluster.metadata.DataStreamTestHelper.backingIndexEqualTo; import static org.elasticsearch.cluster.metadata.DataStreamTestHelper.createBackingIndex; +import static org.elasticsearch.cluster.metadata.DataStreamTestHelper.createFailureStore; import static org.elasticsearch.cluster.metadata.DataStreamTestHelper.newInstance; import static org.elasticsearch.cluster.metadata.IndexMetadata.INDEX_HIDDEN_SETTING; import static org.elasticsearch.common.util.set.Sets.newHashSet; @@ -2294,7 +2296,7 @@ public void testIgnoreThrottled() { new IndicesOptions( IndicesOptions.ConcreteTargetOptions.ERROR_WHEN_UNAVAILABLE_TARGETS, IndicesOptions.WildcardOptions.DEFAULT, - IndicesOptions.GeneralOptions.builder().ignoreThrottled(true).build() + IndicesOptions.GatekeeperOptions.builder().ignoreThrottled(true).build() ), "ind*", "test-index" @@ -2697,6 +2699,200 @@ public void testDataStreams() { } } + public void testDataStreamsWithFailureStore() { + final String dataStreamName = "my-data-stream"; + IndexMetadata index1 = createBackingIndex(dataStreamName, 1, epochMillis).build(); + IndexMetadata index2 = createBackingIndex(dataStreamName, 2, epochMillis).build(); + IndexMetadata failureIndex1 = createFailureStore(dataStreamName, 1, epochMillis).build(); + IndexMetadata failureIndex2 = createFailureStore(dataStreamName, 2, epochMillis).build(); + IndexMetadata otherIndex = indexBuilder("my-other-index", Settings.EMPTY).state(State.OPEN).build(); + + Metadata.Builder mdBuilder = Metadata.builder() + .put(index1, false) + .put(index2, false) + .put(failureIndex1, false) + .put(failureIndex2, false) + .put(otherIndex, false) + .put( + newInstance( + dataStreamName, + List.of(index1.getIndex(), index2.getIndex()), + List.of(failureIndex1.getIndex(), failureIndex2.getIndex()) + ) + ); + ClusterState state = ClusterState.builder(new ClusterName("_name")).metadata(mdBuilder).build(); + + // Test default with an exact data stream name + { + IndicesOptions indicesOptions = IndicesOptions.STRICT_EXPAND_OPEN; + Index[] result = indexNameExpressionResolver.concreteIndices(state, indicesOptions, true, "my-data-stream"); + assertThat(result.length, equalTo(2)); + assertThat(result[0].getName(), equalTo(DataStream.getDefaultBackingIndexName(dataStreamName, 1, epochMillis))); + assertThat(result[1].getName(), equalTo(DataStream.getDefaultBackingIndexName(dataStreamName, 2, epochMillis))); + } + + // Test include failure store with an exact data stream name + { + IndicesOptions indicesOptions = IndicesOptions.builder(IndicesOptions.STRICT_EXPAND_OPEN) + .failureStoreOptions(IndicesOptions.FailureStoreOptions.builder().includeRegularIndices(true).includeFailureIndices(true)) + .build(); + Index[] result = indexNameExpressionResolver.concreteIndices(state, indicesOptions, true, "my-data-stream"); + assertThat(result.length, equalTo(4)); + assertThat(result[0].getName(), equalTo(DataStream.getDefaultBackingIndexName(dataStreamName, 1, epochMillis))); + assertThat(result[1].getName(), equalTo(DataStream.getDefaultBackingIndexName(dataStreamName, 2, epochMillis))); + assertThat(result[2].getName(), equalTo(DataStream.getDefaultFailureStoreName(dataStreamName, 1, epochMillis))); + assertThat(result[3].getName(), equalTo(DataStream.getDefaultFailureStoreName(dataStreamName, 2, epochMillis))); + } + + // Test include failure store while we do not allow failure indices and ignore unavailable + // We expect that they will be skipped + { + IndicesOptions indicesOptions = IndicesOptions.builder(IndicesOptions.STRICT_EXPAND_OPEN) + .failureStoreOptions(IndicesOptions.FailureStoreOptions.builder().includeRegularIndices(true).includeFailureIndices(true)) + .gatekeeperOptions(IndicesOptions.GatekeeperOptions.builder().allowFailureIndices(false).build()) + .concreteTargetOptions(IndicesOptions.ConcreteTargetOptions.ALLOW_UNAVAILABLE_TARGETS) + .build(); + Index[] result = indexNameExpressionResolver.concreteIndices(state, indicesOptions, true, "my-data-stream"); + assertThat(result.length, equalTo(2)); + assertThat(result[0].getName(), equalTo(DataStream.getDefaultBackingIndexName(dataStreamName, 1, epochMillis))); + assertThat(result[1].getName(), equalTo(DataStream.getDefaultBackingIndexName(dataStreamName, 2, epochMillis))); + } + + // Test include failure store while we do not allow failure indices + // We expect an error + { + IndicesOptions indicesOptions = IndicesOptions.builder(IndicesOptions.STRICT_EXPAND_OPEN) + .failureStoreOptions(IndicesOptions.FailureStoreOptions.builder().includeRegularIndices(true).includeFailureIndices(true)) + .gatekeeperOptions(IndicesOptions.GatekeeperOptions.builder().allowFailureIndices(false).build()) + .build(); + FailureIndexNotSupportedException failureIndexNotSupportedException = expectThrows( + FailureIndexNotSupportedException.class, + () -> indexNameExpressionResolver.concreteIndices(state, indicesOptions, true, "my-data-stream") + ); + assertThat( + failureIndexNotSupportedException.getIndex().getName(), + equalTo(DataStream.getDefaultFailureStoreName(dataStreamName, 1, epochMillis)) + ); + } + + // Test only failure store with an exact data stream name + { + IndicesOptions indicesOptions = IndicesOptions.builder(IndicesOptions.STRICT_EXPAND_OPEN) + .failureStoreOptions(IndicesOptions.FailureStoreOptions.builder().includeRegularIndices(false).includeFailureIndices(true)) + .build(); + Index[] result = indexNameExpressionResolver.concreteIndices(state, indicesOptions, true, "my-data-stream"); + assertThat(result.length, equalTo(2)); + assertThat(result[0].getName(), equalTo(DataStream.getDefaultFailureStoreName(dataStreamName, 1, epochMillis))); + assertThat(result[1].getName(), equalTo(DataStream.getDefaultFailureStoreName(dataStreamName, 2, epochMillis))); + } + + // Test default without any expressions + { + IndicesOptions indicesOptions = IndicesOptions.STRICT_EXPAND_OPEN; + Index[] result = indexNameExpressionResolver.concreteIndices(state, indicesOptions, true); + assertThat(result.length, equalTo(3)); + List indexNames = Arrays.stream(result).map(Index::getName).toList(); + assertThat( + indexNames, + containsInAnyOrder( + DataStream.getDefaultBackingIndexName(dataStreamName, 2, epochMillis), + DataStream.getDefaultBackingIndexName(dataStreamName, 1, epochMillis), + otherIndex.getIndex().getName() + ) + ); + } + + // Test include failure store without any expressions + { + IndicesOptions indicesOptions = IndicesOptions.builder(IndicesOptions.STRICT_EXPAND_OPEN) + .failureStoreOptions(IndicesOptions.FailureStoreOptions.builder().includeRegularIndices(true).includeFailureIndices(true)) + .build(); + Index[] result = indexNameExpressionResolver.concreteIndices(state, indicesOptions, true); + assertThat(result.length, equalTo(5)); + List indexNames = Arrays.stream(result).map(Index::getName).toList(); + assertThat( + indexNames, + containsInAnyOrder( + DataStream.getDefaultBackingIndexName(dataStreamName, 2, epochMillis), + DataStream.getDefaultBackingIndexName(dataStreamName, 1, epochMillis), + DataStream.getDefaultFailureStoreName(dataStreamName, 2, epochMillis), + DataStream.getDefaultFailureStoreName(dataStreamName, 1, epochMillis), + otherIndex.getIndex().getName() + ) + ); + } + + // Test only failure store without any expressions + { + IndicesOptions indicesOptions = IndicesOptions.builder(IndicesOptions.STRICT_EXPAND_OPEN) + .failureStoreOptions(IndicesOptions.FailureStoreOptions.builder().includeRegularIndices(false).includeFailureIndices(true)) + .build(); + Index[] result = indexNameExpressionResolver.concreteIndices(state, indicesOptions, true); + assertThat(result.length, equalTo(2)); + List indexNames = Arrays.stream(result).map(Index::getName).toList(); + assertThat( + indexNames, + containsInAnyOrder( + DataStream.getDefaultFailureStoreName(dataStreamName, 2, epochMillis), + DataStream.getDefaultFailureStoreName(dataStreamName, 1, epochMillis) + ) + ); + } + + // Test default with wildcard expression + { + IndicesOptions indicesOptions = IndicesOptions.STRICT_EXPAND_OPEN; + Index[] result = indexNameExpressionResolver.concreteIndices(state, indicesOptions, true, "my-*"); + assertThat(result.length, equalTo(3)); + List indexNames = Arrays.stream(result).map(Index::getName).toList(); + assertThat( + indexNames, + containsInAnyOrder( + DataStream.getDefaultBackingIndexName(dataStreamName, 2, epochMillis), + DataStream.getDefaultBackingIndexName(dataStreamName, 1, epochMillis), + otherIndex.getIndex().getName() + ) + ); + } + + // Test include failure store with wildcard expression + { + IndicesOptions indicesOptions = IndicesOptions.builder(IndicesOptions.STRICT_EXPAND_OPEN) + .failureStoreOptions(IndicesOptions.FailureStoreOptions.builder().includeRegularIndices(true).includeFailureIndices(true)) + .build(); + Index[] result = indexNameExpressionResolver.concreteIndices(state, indicesOptions, true, "my-*"); + assertThat(result.length, equalTo(5)); + List indexNames = Arrays.stream(result).map(Index::getName).toList(); + assertThat( + indexNames, + containsInAnyOrder( + DataStream.getDefaultBackingIndexName(dataStreamName, 2, epochMillis), + DataStream.getDefaultBackingIndexName(dataStreamName, 1, epochMillis), + DataStream.getDefaultFailureStoreName(dataStreamName, 2, epochMillis), + DataStream.getDefaultFailureStoreName(dataStreamName, 1, epochMillis), + otherIndex.getIndex().getName() + ) + ); + } + + // Test only failure store with wildcard expression + { + IndicesOptions indicesOptions = IndicesOptions.builder(IndicesOptions.STRICT_EXPAND_OPEN) + .failureStoreOptions(IndicesOptions.FailureStoreOptions.builder().includeRegularIndices(false).includeFailureIndices(true)) + .build(); + Index[] result = indexNameExpressionResolver.concreteIndices(state, indicesOptions, true, "my-*"); + assertThat(result.length, equalTo(2)); + List indexNames = Arrays.stream(result).map(Index::getName).toList(); + assertThat( + indexNames, + containsInAnyOrder( + DataStream.getDefaultFailureStoreName(dataStreamName, 2, epochMillis), + DataStream.getDefaultFailureStoreName(dataStreamName, 1, epochMillis) + ) + ); + } + } + public void testDataStreamAliases() { String dataStream1 = "my-data-stream-1"; IndexMetadata index1 = createBackingIndex(dataStream1, 1, epochMillis).build(); diff --git a/test/framework/src/main/java/org/elasticsearch/cluster/metadata/DataStreamTestHelper.java b/test/framework/src/main/java/org/elasticsearch/cluster/metadata/DataStreamTestHelper.java index 8fc02bb8e808c..1cc5006fe0018 100644 --- a/test/framework/src/main/java/org/elasticsearch/cluster/metadata/DataStreamTestHelper.java +++ b/test/framework/src/main/java/org/elasticsearch/cluster/metadata/DataStreamTestHelper.java @@ -91,6 +91,10 @@ public static DataStream newInstance(String name, List indices) { return newInstance(name, indices, indices.size(), null); } + public static DataStream newInstance(String name, List indices, List failureIndices) { + return newInstance(name, indices, indices.size(), null, false, null, failureIndices); + } + public static DataStream newInstance(String name, List indices, long generation, Map metadata) { return newInstance(name, indices, generation, metadata, false); } diff --git a/x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/datafeed/DatafeedNodeSelectorTests.java b/x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/datafeed/DatafeedNodeSelectorTests.java index 69d8663478b36..5ddba7519eef2 100644 --- a/x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/datafeed/DatafeedNodeSelectorTests.java +++ b/x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/datafeed/DatafeedNodeSelectorTests.java @@ -338,8 +338,8 @@ public void testIndexDoesntExist() { "cannot start datafeed [datafeed_id] because it failed resolving indices given [not_foo] and " + "indices_options [IndicesOptions[ignore_unavailable=false, allow_no_indices=true, expand_wildcards_open=true, " + "expand_wildcards_closed=false, expand_wildcards_hidden=false, allow_aliases_to_multiple_indices=true, " - + "forbid_closed_indices=true, ignore_aliases=false, ignore_throttled=true]] " - + "with exception [no such index [not_foo]]" + + "forbid_closed_indices=true, ignore_aliases=false, ignore_throttled=true, include_regular_indices=true, " + + "include_failure_indices=false, allow_failure_indices=true]] with exception [no such index [not_foo]]" ) ); @@ -361,8 +361,9 @@ public void testIndexDoesntExist() { + "[cannot start datafeed [datafeed_id] because it failed resolving " + "indices given [not_foo] and indices_options [IndicesOptions[ignore_unavailable=false, allow_no_indices=true, " + "expand_wildcards_open=true, expand_wildcards_closed=false, expand_wildcards_hidden=false, " - + "allow_aliases_to_multiple_indices=true, forbid_closed_indices=true, ignore_aliases=false, ignore_throttled=true" - + "]] with exception [no such index [not_foo]]]" + + "allow_aliases_to_multiple_indices=true, forbid_closed_indices=true, ignore_aliases=false, ignore_throttled=true, " + + "include_regular_indices=true, include_failure_indices=false, allow_failure_indices=true]] " + + "with exception [no such index [not_foo]]]" ) ); } @@ -527,8 +528,8 @@ public void testSelectNode_GivenJobOpeningAndIndexDoesNotExist() { + "[cannot start datafeed [datafeed_id] because it failed resolving indices given [not_foo] and " + "indices_options [IndicesOptions[ignore_unavailable=false, allow_no_indices=true, expand_wildcards_open=true, " + "expand_wildcards_closed=false, expand_wildcards_hidden=false, allow_aliases_to_multiple_indices=true, " - + "forbid_closed_indices=true, ignore_aliases=false, ignore_throttled=true]] " - + "with exception [no such index [not_foo]]]" + + "forbid_closed_indices=true, ignore_aliases=false, ignore_throttled=true, include_regular_indices=true, " + + "include_failure_indices=false, allow_failure_indices=true]] with exception [no such index [not_foo]]]" ) ); } diff --git a/x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/index/IndexResolver.java b/x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/index/IndexResolver.java index 9009baf188cac..318f3888ac9b3 100644 --- a/x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/index/IndexResolver.java +++ b/x-pack/plugin/ql/src/main/java/org/elasticsearch/xpack/ql/index/IndexResolver.java @@ -124,8 +124,8 @@ public String toString() { .allowEmptyExpressions(true) .resolveAliases(false) ) - .generalOptions( - IndicesOptions.GeneralOptions.builder().ignoreThrottled(true).allowClosedIndices(true).allowAliasToMultipleIndices(true) + .gatekeeperOptions( + IndicesOptions.GatekeeperOptions.builder().ignoreThrottled(true).allowClosedIndices(true).allowAliasToMultipleIndices(true) ) .build(); private static final IndicesOptions FROZEN_INDICES_OPTIONS = IndicesOptions.builder() @@ -138,8 +138,8 @@ public String toString() { .allowEmptyExpressions(true) .resolveAliases(false) ) - .generalOptions( - IndicesOptions.GeneralOptions.builder().ignoreThrottled(false).allowClosedIndices(true).allowAliasToMultipleIndices(true) + .gatekeeperOptions( + IndicesOptions.GatekeeperOptions.builder().ignoreThrottled(false).allowClosedIndices(true).allowAliasToMultipleIndices(true) ) .build(); @@ -153,8 +153,8 @@ public String toString() { .allowEmptyExpressions(true) .resolveAliases(true) ) - .generalOptions( - IndicesOptions.GeneralOptions.builder().ignoreThrottled(true).allowClosedIndices(true).allowAliasToMultipleIndices(true) + .gatekeeperOptions( + IndicesOptions.GatekeeperOptions.builder().ignoreThrottled(true).allowClosedIndices(true).allowAliasToMultipleIndices(true) ) .build(); public static final IndicesOptions FIELD_CAPS_FROZEN_INDICES_OPTIONS = IndicesOptions.builder() @@ -167,8 +167,8 @@ public String toString() { .allowEmptyExpressions(true) .resolveAliases(true) ) - .generalOptions( - IndicesOptions.GeneralOptions.builder().ignoreThrottled(false).allowClosedIndices(true).allowAliasToMultipleIndices(true) + .gatekeeperOptions( + IndicesOptions.GatekeeperOptions.builder().ignoreThrottled(false).allowClosedIndices(true).allowAliasToMultipleIndices(true) ) .build(); From d9d291b8a3b215358cfc365e2d9067fd6f00dd54 Mon Sep 17 00:00:00 2001 From: Yang Wang Date: Wed, 6 Mar 2024 23:34:18 +1100 Subject: [PATCH 21/27] Minor improvement for SparseFileTracker toString (#106004) Include both length and complete in toString for easier introspection. Also add a comment to the choice of using a string constant (`file`) for the description. --- .../org/elasticsearch/blobcache/common/SparseFileTracker.java | 2 +- .../elasticsearch/blobcache/shared/SharedBlobCacheService.java | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/x-pack/plugin/blob-cache/src/main/java/org/elasticsearch/blobcache/common/SparseFileTracker.java b/x-pack/plugin/blob-cache/src/main/java/org/elasticsearch/blobcache/common/SparseFileTracker.java index 5eb146102cd76..e9be9577063cf 100644 --- a/x-pack/plugin/blob-cache/src/main/java/org/elasticsearch/blobcache/common/SparseFileTracker.java +++ b/x-pack/plugin/blob-cache/src/main/java/org/elasticsearch/blobcache/common/SparseFileTracker.java @@ -525,7 +525,7 @@ private boolean invariant() { @Override public String toString() { - return "SparseFileTracker[" + description + ']'; + return "SparseFileTracker{description=" + description + ", length=" + length + ", complete=" + complete + '}'; } /** diff --git a/x-pack/plugin/blob-cache/src/main/java/org/elasticsearch/blobcache/shared/SharedBlobCacheService.java b/x-pack/plugin/blob-cache/src/main/java/org/elasticsearch/blobcache/shared/SharedBlobCacheService.java index 1f6f075a2b2af..d4c7c04c5b26e 100644 --- a/x-pack/plugin/blob-cache/src/main/java/org/elasticsearch/blobcache/shared/SharedBlobCacheService.java +++ b/x-pack/plugin/blob-cache/src/main/java/org/elasticsearch/blobcache/shared/SharedBlobCacheService.java @@ -715,6 +715,7 @@ class CacheFileRegion extends EvictableRefCounted { CacheFileRegion(RegionKey regionKey, int regionSize) { this.regionKey = regionKey; assert regionSize > 0; + // NOTE we use a constant string for description to avoid consume extra heap space tracker = new SparseFileTracker("file", regionSize); } From 38168407ef6200ab3f1df55553a004f7a254ec56 Mon Sep 17 00:00:00 2001 From: Joe Gallo Date: Wed, 6 Mar 2024 07:45:13 -0500 Subject: [PATCH 22/27] Docs typo fix (#105835) (#106002) Co-authored-by: MikhailBerezhanov <35196259+MikhailBerezhanov@users.noreply.github.com> --- docs/reference/data-streams/lifecycle/index.asciidoc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/reference/data-streams/lifecycle/index.asciidoc b/docs/reference/data-streams/lifecycle/index.asciidoc index ef5558817885e..bf861df7c80d4 100644 --- a/docs/reference/data-streams/lifecycle/index.asciidoc +++ b/docs/reference/data-streams/lifecycle/index.asciidoc @@ -36,8 +36,8 @@ each data stream and performs the following steps: automatically tail merges the index. Data stream lifecycle executes a merge operation that only targets the long tail of small segments instead of the whole shard. As the segments are organised into tiers of exponential sizes, merging the long tail of small segments is only a -fraction of the cost of force mergeing to a single segment. The small segments would usually -hold the most recent data so tail mergeing will focus the merging resources on the higher-value +fraction of the cost of force merging to a single segment. The small segments would usually +hold the most recent data so tail merging will focus the merging resources on the higher-value data that is most likely to keep being queried. 4. If <> is configured it will execute all the configured downsampling rounds. From 61a50339a67e807687554074373125a2be99b25d Mon Sep 17 00:00:00 2001 From: Craig Taverner Date: Wed, 6 Mar 2024 14:05:47 +0100 Subject: [PATCH 23/27] For cartesian values we are even more lenient with extremely large values (#106014) --- .../lucene/spatial/CentroidCalculatorTests.java | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/server/src/test/java/org/elasticsearch/lucene/spatial/CentroidCalculatorTests.java b/server/src/test/java/org/elasticsearch/lucene/spatial/CentroidCalculatorTests.java index d15ea1ac2e469..7a5cb5de49bdc 100644 --- a/server/src/test/java/org/elasticsearch/lucene/spatial/CentroidCalculatorTests.java +++ b/server/src/test/java/org/elasticsearch/lucene/spatial/CentroidCalculatorTests.java @@ -423,13 +423,12 @@ private CentroidMatcher(double x, double y, double weight, double weightFactor) } private Matcher matchDouble(double value) { - if (value > 1e20 || value < 1e20) { - // Very large values have floating point errors, so instead of an absolute value, we use a relative one - return closeTo(value, Math.abs(value / 1e10)); - } else { - // Most data (notably geo data) has values within bounds, and an absolute delta makes more sense. - return closeTo(value, DELTA); - } + // Very large values have floating point errors, so instead of an absolute value, we use a relative one + // Most data (notably geo data) has values within bounds, and an absolute delta makes more sense. + double delta = (value > 1e28 || value < -1e28) ? Math.abs(value / 1e6) + : (value > 1e20 || value < -1e20) ? Math.abs(value / 1e10) + : DELTA; + return closeTo(value, delta); } @Override From c21b23c2641fbc65648d19719cd61ffe9556dd11 Mon Sep 17 00:00:00 2001 From: Armin Braun Date: Wed, 6 Mar 2024 14:29:56 +0100 Subject: [PATCH 24/27] Reduce overhead of (source-less) SearchHit (#105659) Reduce overhead of `SearchHit` a little. No need for any real ref-counting if there's neither source nor nested hits. Same goes for `SearchHits` which don't have to be ref-counted if their contents aren't. Also, don't create pointless unmodifiable maps wrapping the empty singleton for highlight fields and use the singleton for the empty search sort values. --- .../org/elasticsearch/search/SearchHit.java | 19 +++++++++++++------ .../org/elasticsearch/search/SearchHits.java | 11 +++++++++-- .../search/SearchSortValues.java | 18 +++++++++++++----- .../search/SearchSortValuesTests.java | 2 +- 4 files changed, 36 insertions(+), 14 deletions(-) diff --git a/server/src/main/java/org/elasticsearch/search/SearchHit.java b/server/src/main/java/org/elasticsearch/search/SearchHit.java index fe11aa8af39f4..60ced289929a0 100644 --- a/server/src/main/java/org/elasticsearch/search/SearchHit.java +++ b/server/src/main/java/org/elasticsearch/search/SearchHit.java @@ -204,7 +204,7 @@ public SearchHit( this.innerHits = innerHits; this.documentFields = documentFields; this.metaFields = metaFields; - this.refCounted = refCounted == null ? LeakTracker.wrap(new SimpleRefCounted()) : ALWAYS_REFERENCED; + this.refCounted = refCounted == null ? LeakTracker.wrap(new SimpleRefCounted()) : refCounted; } public static SearchHit readFrom(StreamInput in, boolean pooled) throws IOException { @@ -233,8 +233,10 @@ public static SearchHit readFrom(StreamInput in, boolean pooled) throws IOExcept } final Map documentFields = in.readMap(DocumentField::new); final Map metaFields = in.readMap(DocumentField::new); - final Map highlightFields = in.readMapValues(HighlightField::new, HighlightField::name); - final SearchSortValues sortValues = new SearchSortValues(in); + Map highlightFields = in.readMapValues(HighlightField::new, HighlightField::name); + highlightFields = highlightFields.isEmpty() ? null : unmodifiableMap(highlightFields); + + final SearchSortValues sortValues = SearchSortValues.readFrom(in); final Map matchedQueries; if (in.getTransportVersion().onOrAfter(TransportVersions.V_8_8_0)) { @@ -257,12 +259,17 @@ public static SearchHit readFrom(StreamInput in, boolean pooled) throws IOExcept index = shardTarget.getIndex(); clusterAlias = shardTarget.getClusterAlias(); } + + boolean isPooled = pooled && source != null; final Map innerHits; int size = in.readVInt(); if (size > 0) { innerHits = Maps.newMapWithExpectedSize(size); for (int i = 0; i < size; i++) { - innerHits.put(in.readString(), SearchHits.readFrom(in, pooled)); + var key = in.readString(); + var nestedHits = SearchHits.readFrom(in, pooled); + innerHits.put(key, nestedHits); + isPooled = isPooled || nestedHits.isPooled(); } } else { innerHits = null; @@ -277,7 +284,7 @@ public static SearchHit readFrom(StreamInput in, boolean pooled) throws IOExcept seqNo, primaryTerm, source, - unmodifiableMap(highlightFields), + highlightFields, sortValues, matchedQueries, explanation, @@ -288,7 +295,7 @@ public static SearchHit readFrom(StreamInput in, boolean pooled) throws IOExcept innerHits, documentFields, metaFields, - pooled ? null : ALWAYS_REFERENCED + isPooled ? null : ALWAYS_REFERENCED ); } diff --git a/server/src/main/java/org/elasticsearch/search/SearchHits.java b/server/src/main/java/org/elasticsearch/search/SearchHits.java index ce8ccf4b7f0e6..d559fc60fa72d 100644 --- a/server/src/main/java/org/elasticsearch/search/SearchHits.java +++ b/server/src/main/java/org/elasticsearch/search/SearchHits.java @@ -132,24 +132,31 @@ public static SearchHits readFrom(StreamInput in, boolean pooled) throws IOExcep final float maxScore = in.readFloat(); int size = in.readVInt(); final SearchHit[] hits; + boolean isPooled = false; if (size == 0) { hits = EMPTY; } else { hits = new SearchHit[size]; for (int i = 0; i < hits.length; i++) { - hits[i] = SearchHit.readFrom(in, pooled); + var hit = SearchHit.readFrom(in, pooled); + hits[i] = hit; + isPooled = isPooled || hit.isPooled(); } } var sortFields = in.readOptionalArray(Lucene::readSortField, SortField[]::new); var collapseField = in.readOptionalString(); var collapseValues = in.readOptionalArray(Lucene::readSortValue, Object[]::new); - if (pooled) { + if (isPooled) { return new SearchHits(hits, totalHits, maxScore, sortFields, collapseField, collapseValues); } else { return unpooled(hits, totalHits, maxScore, sortFields, collapseField, collapseValues); } } + public boolean isPooled() { + return refCounted != ALWAYS_REFERENCED; + } + @Override public void writeTo(StreamOutput out) throws IOException { assert hasReferences(); diff --git a/server/src/main/java/org/elasticsearch/search/SearchSortValues.java b/server/src/main/java/org/elasticsearch/search/SearchSortValues.java index b82e6632ca9ec..38bc705bdf5ae 100644 --- a/server/src/main/java/org/elasticsearch/search/SearchSortValues.java +++ b/server/src/main/java/org/elasticsearch/search/SearchSortValues.java @@ -32,8 +32,7 @@ public class SearchSortValues implements ToXContentFragment, Writeable { private final Object[] rawSortValues; SearchSortValues(Object[] sortValues) { - this.formattedSortValues = Objects.requireNonNull(sortValues, "sort values must not be empty"); - this.rawSortValues = EMPTY_ARRAY; + this(Objects.requireNonNull(sortValues, "sort values must not be empty"), EMPTY_ARRAY); } public SearchSortValues(Object[] rawSortValues, DocValueFormat[] sortValueFormats) { @@ -52,9 +51,18 @@ public SearchSortValues(Object[] rawSortValues, DocValueFormat[] sortValueFormat } } - SearchSortValues(StreamInput in) throws IOException { - this.formattedSortValues = in.readArray(Lucene::readSortValue, Object[]::new); - this.rawSortValues = in.readArray(Lucene::readSortValue, Object[]::new); + public static SearchSortValues readFrom(StreamInput in) throws IOException { + Object[] formattedSortValues = in.readArray(Lucene::readSortValue, Object[]::new); + Object[] rawSortValues = in.readArray(Lucene::readSortValue, Object[]::new); + if (formattedSortValues.length == 0 && rawSortValues.length == 0) { + return EMPTY; + } + return new SearchSortValues(formattedSortValues, rawSortValues); + } + + private SearchSortValues(Object[] formattedSortValues, Object[] rawSortValues) { + this.formattedSortValues = formattedSortValues; + this.rawSortValues = rawSortValues; } @Override diff --git a/server/src/test/java/org/elasticsearch/search/SearchSortValuesTests.java b/server/src/test/java/org/elasticsearch/search/SearchSortValuesTests.java index ac9ae1da0fddd..e21ae8af04b77 100644 --- a/server/src/test/java/org/elasticsearch/search/SearchSortValuesTests.java +++ b/server/src/test/java/org/elasticsearch/search/SearchSortValuesTests.java @@ -90,7 +90,7 @@ protected SearchSortValues createTestInstance() { @Override protected Writeable.Reader instanceReader() { - return SearchSortValues::new; + return SearchSortValues::readFrom; } @Override From f468024c25ed7f373e209322b2f2b9ced25a100e Mon Sep 17 00:00:00 2001 From: Armin Braun Date: Wed, 6 Mar 2024 14:30:24 +0100 Subject: [PATCH 25/27] Fix duplication around transport code in IndexBasedTransformConfigManager (#105907) Drying up the logic in this class further and making use of `delegateAndWrap` in an effort to narrow down possible sources for the remaining very rare leaks we observe around `SearchResponse`. --- .../IndexBasedTransformConfigManager.java | 475 ++++++++---------- 1 file changed, 209 insertions(+), 266 deletions(-) diff --git a/x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/persistence/IndexBasedTransformConfigManager.java b/x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/persistence/IndexBasedTransformConfigManager.java index 1d44ed5a1f8ef..40eb2e2ad294a 100644 --- a/x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/persistence/IndexBasedTransformConfigManager.java +++ b/x-pack/plugin/transform/src/main/java/org/elasticsearch/xpack/transform/persistence/IndexBasedTransformConfigManager.java @@ -14,6 +14,9 @@ import org.elasticsearch.ResourceAlreadyExistsException; import org.elasticsearch.ResourceNotFoundException; import org.elasticsearch.action.ActionListener; +import org.elasticsearch.action.ActionRequest; +import org.elasticsearch.action.ActionResponse; +import org.elasticsearch.action.ActionType; import org.elasticsearch.action.DocWriteRequest; import org.elasticsearch.action.admin.indices.delete.DeleteIndexRequest; import org.elasticsearch.action.admin.indices.delete.TransportDeleteIndexAction; @@ -60,6 +63,7 @@ import org.elasticsearch.xcontent.XContentFactory; import org.elasticsearch.xcontent.XContentParser; import org.elasticsearch.xcontent.XContentType; +import org.elasticsearch.xpack.core.ClientHelper; import org.elasticsearch.xpack.core.action.util.ExpandedIdsMatcher; import org.elasticsearch.xpack.core.action.util.PageParams; import org.elasticsearch.xpack.core.transform.TransformField; @@ -76,10 +80,10 @@ import java.util.HashSet; import java.util.List; import java.util.Set; +import java.util.function.BiConsumer; import static org.elasticsearch.index.mapper.MapperService.SINGLE_MAPPING_NAME; import static org.elasticsearch.xpack.core.ClientHelper.TRANSFORM_ORIGIN; -import static org.elasticsearch.xpack.core.ClientHelper.executeAsyncWithOrigin; /** * Place of all interactions with the internal transforms index. For configuration and mappings see @link{TransformInternalIndex} @@ -135,9 +139,7 @@ public void putTransformCheckpoint(TransformCheckpoint checkpoint, ActionListene .id(TransformCheckpoint.documentId(checkpoint.getTransformId(), checkpoint.getCheckpoint())) .source(source); - executeAsyncWithOrigin(client, TRANSFORM_ORIGIN, TransportIndexAction.TYPE, indexRequest, ActionListener.wrap(r -> { - listener.onResponse(true); - }, listener::onFailure)); + executeAsyncWithOrigin(TransportIndexAction.TYPE, indexRequest, listener.delegateFailureAndWrap((l, r) -> l.onResponse(true))); } catch (IOException e) { // not expected to happen but for the sake of completeness listener.onFailure(e); @@ -180,22 +182,16 @@ public void deleteOldTransformConfigurations(String transformId, ActionListener< ) ); - executeAsyncWithOrigin( - client, - TRANSFORM_ORIGIN, - DeleteByQueryAction.INSTANCE, - deleteByQueryRequest, - ActionListener.wrap(response -> { - if ((response.getBulkFailures().isEmpty() && response.getSearchFailures().isEmpty()) == false) { - Tuple statusAndReason = getStatusAndReason(response); - listener.onFailure( - new ElasticsearchStatusException(statusAndReason.v2().getMessage(), statusAndReason.v1(), statusAndReason.v2()) - ); - return; - } - listener.onResponse(true); - }, listener::onFailure) - ); + executeAsyncWithOrigin(DeleteByQueryAction.INSTANCE, deleteByQueryRequest, listener.delegateFailureAndWrap((l, response) -> { + if ((response.getBulkFailures().isEmpty() && response.getSearchFailures().isEmpty()) == false) { + Tuple statusAndReason = getStatusAndReason(response); + l.onFailure( + new ElasticsearchStatusException(statusAndReason.v2().getMessage(), statusAndReason.v1(), statusAndReason.v2()) + ); + return; + } + l.onResponse(true); + })); } @Override @@ -212,22 +208,7 @@ public void deleteOldTransformStoredDocuments(String transformId, ActionListener .filter(QueryBuilders.termQuery("_id", TransformStoredDoc.documentId(transformId))) ) ); - executeAsyncWithOrigin( - client, - TRANSFORM_ORIGIN, - DeleteByQueryAction.INSTANCE, - deleteByQueryRequest, - ActionListener.wrap(response -> { - if ((response.getBulkFailures().isEmpty() && response.getSearchFailures().isEmpty()) == false) { - Tuple statusAndReason = getStatusAndReason(response); - listener.onFailure( - new ElasticsearchStatusException(statusAndReason.v2().getMessage(), statusAndReason.v1(), statusAndReason.v2()) - ); - return; - } - listener.onResponse(response.getDeleted()); - }, listener::onFailure) - ); + deleteByQuery(listener, deleteByQueryRequest); } @Override @@ -247,22 +228,20 @@ public void deleteOldCheckpoints(String transformId, long deleteCheckpointsBelow ) ); logger.debug("Deleting old checkpoints using {}", deleteByQueryRequest.getSearchRequest()); - executeAsyncWithOrigin( - client, - TRANSFORM_ORIGIN, - DeleteByQueryAction.INSTANCE, - deleteByQueryRequest, - ActionListener.wrap(response -> { - if ((response.getBulkFailures().isEmpty() && response.getSearchFailures().isEmpty()) == false) { - Tuple statusAndReason = getStatusAndReason(response); - listener.onFailure( - new ElasticsearchStatusException(statusAndReason.v2().getMessage(), statusAndReason.v1(), statusAndReason.v2()) - ); - return; - } - listener.onResponse(response.getDeleted()); - }, listener::onFailure) - ); + deleteByQuery(listener, deleteByQueryRequest); + } + + private void deleteByQuery(ActionListener listener, DeleteByQueryRequest deleteByQueryRequest) { + executeAsyncWithOrigin(DeleteByQueryAction.INSTANCE, deleteByQueryRequest, listener.delegateFailureAndWrap((l, response) -> { + if ((response.getBulkFailures().isEmpty() && response.getSearchFailures().isEmpty()) == false) { + Tuple statusAndReason = getStatusAndReason(response); + l.onFailure( + new ElasticsearchStatusException(statusAndReason.v2().getMessage(), statusAndReason.v1(), statusAndReason.v2()) + ); + return; + } + l.onResponse(response.getDeleted()); + })); } @Override @@ -304,13 +283,13 @@ public void deleteOldIndices(ActionListener listener) { IndicesOptions.LENIENT_EXPAND_OPEN ); - executeAsyncWithOrigin(client, TRANSFORM_ORIGIN, TransportDeleteIndexAction.TYPE, deleteRequest, ActionListener.wrap(response -> { + executeAsyncWithOrigin(TransportDeleteIndexAction.TYPE, deleteRequest, listener.delegateFailureAndWrap((l, response) -> { if (response.isAcknowledged() == false) { - listener.onFailure(new ElasticsearchStatusException("Failed to delete internal indices", RestStatus.INTERNAL_SERVER_ERROR)); + l.onFailure(new ElasticsearchStatusException("Failed to delete internal indices", RestStatus.INTERNAL_SERVER_ERROR)); return; } - listener.onResponse(true); - }, listener::onFailure)); + l.onResponse(true); + })); } private void putTransformConfiguration( @@ -331,9 +310,7 @@ private void putTransformConfiguration( if (seqNoPrimaryTermAndIndex != null) { indexRequest.setIfSeqNo(seqNoPrimaryTermAndIndex.getSeqNo()).setIfPrimaryTerm(seqNoPrimaryTermAndIndex.getPrimaryTerm()); } - executeAsyncWithOrigin(client, TRANSFORM_ORIGIN, TransportIndexAction.TYPE, indexRequest, ActionListener.wrap(r -> { - listener.onResponse(true); - }, e -> { + executeAsyncWithOrigin(TransportIndexAction.TYPE, indexRequest, ActionListener.wrap(r -> listener.onResponse(true), e -> { if (e instanceof VersionConflictEngineException) { if (DocWriteRequest.OpType.CREATE.equals(opType)) { // we want to create the transform but it already exists listener.onFailure( @@ -378,22 +355,16 @@ public void getTransformCheckpoint(String transformId, long checkpoint, ActionLi .setAllowPartialSearchResults(false) .request(); - executeAsyncWithOrigin( - client, - TRANSFORM_ORIGIN, - TransportSearchAction.TYPE, - searchRequest, - ActionListener.wrap(searchResponse -> { - if (searchResponse.getHits().getHits().length == 0) { - // do not fail if checkpoint does not exist but return an empty checkpoint - logger.trace("found no checkpoint for transform [" + transformId + "], returning empty checkpoint"); - resultListener.onResponse(TransformCheckpoint.EMPTY); - return; - } - BytesReference source = searchResponse.getHits().getHits()[0].getSourceRef(); - parseCheckpointsLenientlyFromSource(source, transformId, resultListener); - }, resultListener::onFailure) - ); + executeAsyncWithOrigin(TransportSearchAction.TYPE, searchRequest, resultListener.delegateFailureAndWrap((l, searchResponse) -> { + if (searchResponse.getHits().getHits().length == 0) { + // do not fail if checkpoint does not exist but return an empty checkpoint + logger.trace("found no checkpoint for transform [{}], returning empty checkpoint", transformId); + l.onResponse(TransformCheckpoint.EMPTY); + return; + } + BytesReference source = searchResponse.getHits().getHits()[0].getSourceRef(); + parseCheckpointsLenientlyFromSource(source, transformId, l); + })); } @Override @@ -416,14 +387,12 @@ public void getTransformCheckpointForUpdate( .request(); executeAsyncWithOrigin( - client, - TRANSFORM_ORIGIN, TransportSearchAction.TYPE, searchRequest, - ActionListener.wrap(searchResponse -> { + checkpointAndVersionListener.delegateFailureAndWrap((l, searchResponse) -> { if (searchResponse.getHits().getHits().length == 0) { // do not fail, this _must_ be handled by the caller - checkpointAndVersionListener.onResponse(null); + l.onResponse(null); return; } SearchHit hit = searchResponse.getHits().getHits()[0]; @@ -431,17 +400,16 @@ public void getTransformCheckpointForUpdate( parseCheckpointsLenientlyFromSource( source, transformId, - ActionListener.wrap( - parsedCheckpoint -> checkpointAndVersionListener.onResponse( + l.delegateFailureAndWrap( + (ll, parsedCheckpoint) -> ll.onResponse( Tuple.tuple( parsedCheckpoint, new SeqNoPrimaryTermAndIndex(hit.getSeqNo(), hit.getPrimaryTerm(), hit.getIndex()) ) - ), - checkpointAndVersionListener::onFailure + ) ) ); - }, checkpointAndVersionListener::onFailure) + }) ); } @@ -459,22 +427,16 @@ public void getTransformConfiguration(String transformId, ActionListenerwrap(searchResponse -> { - if (searchResponse.getHits().getHits().length == 0) { - resultListener.onFailure( - new ResourceNotFoundException(TransformMessages.getMessage(TransformMessages.REST_UNKNOWN_TRANSFORM, transformId)) - ); - return; - } - BytesReference source = searchResponse.getHits().getHits()[0].getSourceRef(); - parseTransformLenientlyFromSource(source, transformId, resultListener); - }, resultListener::onFailure) - ); + executeAsyncWithOrigin(TransportSearchAction.TYPE, searchRequest, resultListener.delegateFailureAndWrap((l, searchResponse) -> { + if (searchResponse.getHits().getHits().length == 0) { + l.onFailure( + new ResourceNotFoundException(TransformMessages.getMessage(TransformMessages.REST_UNKNOWN_TRANSFORM, transformId)) + ); + return; + } + BytesReference source = searchResponse.getHits().getHits()[0].getSourceRef(); + parseTransformLenientlyFromSource(source, transformId, l); + })); } @Override @@ -495,26 +457,29 @@ public void getTransformConfigurationForUpdate( .seqNoAndPrimaryTerm(true) .request(); - executeAsyncWithOrigin(client, TRANSFORM_ORIGIN, TransportSearchAction.TYPE, searchRequest, ActionListener.wrap(searchResponse -> { - if (searchResponse.getHits().getHits().length == 0) { - configAndVersionListener.onFailure( - new ResourceNotFoundException(TransformMessages.getMessage(TransformMessages.REST_UNKNOWN_TRANSFORM, transformId)) + executeAsyncWithOrigin( + TransportSearchAction.TYPE, + searchRequest, + configAndVersionListener.delegateFailureAndWrap((l, searchResponse) -> { + if (searchResponse.getHits().getHits().length == 0) { + l.onFailure( + new ResourceNotFoundException(TransformMessages.getMessage(TransformMessages.REST_UNKNOWN_TRANSFORM, transformId)) + ); + return; + } + SearchHit hit = searchResponse.getHits().getHits()[0]; + BytesReference source = hit.getSourceRef(); + parseTransformLenientlyFromSource( + source, + transformId, + l.delegateFailureAndWrap( + (ll, config) -> ll.onResponse( + Tuple.tuple(config, new SeqNoPrimaryTermAndIndex(hit.getSeqNo(), hit.getPrimaryTerm(), hit.getIndex())) + ) + ) ); - return; - } - SearchHit hit = searchResponse.getHits().getHits()[0]; - BytesReference source = hit.getSourceRef(); - parseTransformLenientlyFromSource( - source, - transformId, - ActionListener.wrap( - config -> configAndVersionListener.onResponse( - Tuple.tuple(config, new SeqNoPrimaryTermAndIndex(hit.getSeqNo(), hit.getPrimaryTerm(), hit.getIndex())) - ), - configAndVersionListener::onFailure - ) - ); - }, configAndVersionListener::onFailure)); + }) + ); } @Override @@ -543,48 +508,40 @@ public void expandTransformIds( final ExpandedIdsMatcher requiredMatches = new ExpandedIdsMatcher(idTokens, allowNoMatch); - executeAsyncWithOrigin( - client.threadPool().getThreadContext(), - TRANSFORM_ORIGIN, - request, - ActionListener.wrap(searchResponse -> { - long totalHits = searchResponse.getHits().getTotalHits().value; - // important: preserve order - Set ids = Sets.newLinkedHashSetWithExpectedSize(searchResponse.getHits().getHits().length); - Set configs = Sets.newLinkedHashSetWithExpectedSize(searchResponse.getHits().getHits().length); - for (SearchHit hit : searchResponse.getHits().getHits()) { - try (XContentParser parser = createParser(hit)) { - TransformConfig config = TransformConfig.fromXContent(parser, null, true); - if (ids.add(config.getId())) { - configs.add(config); - } - } catch (IOException e) { - foundConfigsListener.onFailure(new ElasticsearchParseException("failed to parse search hit for ids", e)); - return; + executeAsyncWithOrigin(request, foundConfigsListener.delegateFailureAndWrap((l, searchResponse) -> { + long totalHits = searchResponse.getHits().getTotalHits().value; + // important: preserve order + Set ids = Sets.newLinkedHashSetWithExpectedSize(searchResponse.getHits().getHits().length); + Set configs = Sets.newLinkedHashSetWithExpectedSize(searchResponse.getHits().getHits().length); + for (SearchHit hit : searchResponse.getHits().getHits()) { + try (XContentParser parser = createParser(hit)) { + TransformConfig config = TransformConfig.fromXContent(parser, null, true); + if (ids.add(config.getId())) { + configs.add(config); } - } - requiredMatches.filterMatchedIds(ids); - if (requiredMatches.hasUnmatchedIds()) { - // some required Ids were not found - foundConfigsListener.onFailure( - new ResourceNotFoundException( - TransformMessages.getMessage(TransformMessages.REST_UNKNOWN_TRANSFORM, requiredMatches.unmatchedIdsString()) - ) - ); + } catch (IOException e) { + l.onFailure(new ElasticsearchParseException("failed to parse search hit for ids", e)); return; } - // if only exact ids have been given, take the count from docs to avoid potential duplicates - // in versioned indexes (like transform) - if (requiredMatches.isOnlyExact()) { - foundConfigsListener.onResponse( - new Tuple<>((long) ids.size(), Tuple.tuple(new ArrayList<>(ids), new ArrayList<>(configs))) - ); - } else { - foundConfigsListener.onResponse(new Tuple<>(totalHits, Tuple.tuple(new ArrayList<>(ids), new ArrayList<>(configs)))); - } - }, foundConfigsListener::onFailure), - client::search - ); + } + requiredMatches.filterMatchedIds(ids); + if (requiredMatches.hasUnmatchedIds()) { + // some required Ids were not found + l.onFailure( + new ResourceNotFoundException( + TransformMessages.getMessage(TransformMessages.REST_UNKNOWN_TRANSFORM, requiredMatches.unmatchedIdsString()) + ) + ); + return; + } + // if only exact ids have been given, take the count from docs to avoid potential duplicates + // in versioned indexes (like transform) + if (requiredMatches.isOnlyExact()) { + l.onResponse(new Tuple<>((long) ids.size(), Tuple.tuple(new ArrayList<>(ids), new ArrayList<>(configs)))); + } else { + l.onResponse(new Tuple<>(totalHits, Tuple.tuple(new ArrayList<>(ids), new ArrayList<>(configs)))); + } + }), client::search); } private XContentParser createParser(BytesReference source) throws IOException { @@ -601,12 +558,7 @@ private XContentParser createParser(SearchHit hit) throws IOException { @Override public void getAllTransformIds(TimeValue timeout, ActionListener> listener) { - expandAllTransformIds( - false, - MAX_RESULTS_WINDOW, - timeout, - ActionListener.wrap(r -> listener.onResponse(r.v2()), listener::onFailure) - ); + expandAllTransformIds(false, MAX_RESULTS_WINDOW, timeout, listener.delegateFailureAndWrap((l, r) -> l.onResponse(r.v2()))); } @Override @@ -616,7 +568,7 @@ public void getAllOutdatedTransformIds(TimeValue timeout, ActionListener listener) { - ActionListener deleteListener = ActionListener.wrap(dbqResponse -> { listener.onResponse(true); }, e -> { + ActionListener deleteListener = ActionListener.wrap(dbqResponse -> listener.onResponse(true), e -> { if (e.getClass() == IndexNotFoundException.class) { listener.onFailure( new ResourceNotFoundException(TransformMessages.getMessage(TransformMessages.REST_UNKNOWN_TRANSFORM, transformId)) @@ -636,7 +588,7 @@ public void resetTransform(String transformId, ActionListener listener) .query(QueryBuilders.termQuery(TransformField.ID.getPreferredName(), transformId)) .trackTotalHitsUpTo(1) ); - executeAsyncWithOrigin(client, TRANSFORM_ORIGIN, TransportSearchAction.TYPE, searchRequest, ActionListener.wrap(searchResponse -> { + executeAsyncWithOrigin(TransportSearchAction.TYPE, searchRequest, deleteListener.delegateFailureAndWrap((l, searchResponse) -> { if (searchResponse.getHits().getTotalHits().value == 0) { listener.onFailure( new ResourceNotFoundException(TransformMessages.getMessage(TransformMessages.REST_UNKNOWN_TRANSFORM, transformId)) @@ -655,8 +607,8 @@ public void resetTransform(String transformId, ActionListener listener) TransformInternalIndexConstants.INDEX_NAME_PATTERN, TransformInternalIndexConstants.INDEX_NAME_PATTERN_DEPRECATED ).setQuery(dbqQuery).setRefresh(true); - executeAsyncWithOrigin(client, TRANSFORM_ORIGIN, DeleteByQueryAction.INSTANCE, dbqRequest, deleteListener); - }, deleteListener::onFailure)); + executeAsyncWithOrigin(DeleteByQueryAction.INSTANCE, dbqRequest, l); + })); } @Override @@ -668,7 +620,7 @@ public void deleteTransform(String transformId, ActionListener listener request.setQuery(query); request.setRefresh(true); - executeAsyncWithOrigin(client, TRANSFORM_ORIGIN, DeleteByQueryAction.INSTANCE, request, ActionListener.wrap(deleteResponse -> { + executeAsyncWithOrigin(DeleteByQueryAction.INSTANCE, request, ActionListener.wrap(deleteResponse -> { if (deleteResponse.getDeleted() == 0) { listener.onFailure( new ResourceNotFoundException(TransformMessages.getMessage(TransformMessages.REST_UNKNOWN_TRANSFORM, transformId)) @@ -714,8 +666,6 @@ public void putOrUpdateTransformStoredDoc( } executeAsyncWithOrigin( - client, - TRANSFORM_ORIGIN, TransportIndexAction.TYPE, indexRequest, ActionListener.wrap( @@ -758,38 +708,30 @@ public void getTransformStoredDoc( .seqNoAndPrimaryTerm(true) .request(); - executeAsyncWithOrigin( - client, - TRANSFORM_ORIGIN, - TransportSearchAction.TYPE, - searchRequest, - ActionListener.wrap(searchResponse -> { - if (searchResponse.getHits().getHits().length == 0) { - if (allowNoMatch) { - resultListener.onResponse(null); - } else { - resultListener.onFailure( - new ResourceNotFoundException( - TransformMessages.getMessage(TransformMessages.UNKNOWN_TRANSFORM_STATS, transformId) - ) - ); - } - return; - } - SearchHit searchHit = searchResponse.getHits().getHits()[0]; - try (XContentParser parser = createParser(searchHit)) { - resultListener.onResponse( - Tuple.tuple(TransformStoredDoc.fromXContent(parser), SeqNoPrimaryTermAndIndex.fromSearchHit(searchHit)) - ); - } catch (Exception e) { - logger.error( - TransformMessages.getMessage(TransformMessages.FAILED_TO_PARSE_TRANSFORM_STATISTICS_CONFIGURATION, transformId), - e + executeAsyncWithOrigin(TransportSearchAction.TYPE, searchRequest, resultListener.delegateFailureAndWrap((l, searchResponse) -> { + if (searchResponse.getHits().getHits().length == 0) { + if (allowNoMatch) { + l.onResponse(null); + } else { + l.onFailure( + new ResourceNotFoundException(TransformMessages.getMessage(TransformMessages.UNKNOWN_TRANSFORM_STATS, transformId)) ); - resultListener.onFailure(e); } - }, resultListener::onFailure) - ); + return; + } + SearchHit searchHit = searchResponse.getHits().getHits()[0]; + try (XContentParser parser = createParser(searchHit)) { + resultListener.onResponse( + Tuple.tuple(TransformStoredDoc.fromXContent(parser), SeqNoPrimaryTermAndIndex.fromSearchHit(searchHit)) + ); + } catch (Exception e) { + logger.error( + TransformMessages.getMessage(TransformMessages.FAILED_TO_PARSE_TRANSFORM_STATISTICS_CONFIGURATION, transformId), + e + ); + resultListener.onFailure(e); + } + })); } @Override @@ -816,43 +758,50 @@ public void getTransformStoredDocs( .setTimeout(timeout) .request(); - executeAsyncWithOrigin( - client.threadPool().getThreadContext(), - TRANSFORM_ORIGIN, - searchRequest, - ActionListener.wrap(searchResponse -> { - List stats = new ArrayList<>(); - String previousId = null; - for (SearchHit hit : searchResponse.getHits().getHits()) { - // skip old versions - if (hit.getId().equals(previousId) == false) { - previousId = hit.getId(); - try (XContentParser parser = createParser(hit)) { - stats.add(TransformStoredDoc.fromXContent(parser)); - } catch (IOException e) { - listener.onFailure(new ElasticsearchParseException("failed to parse transform stats from search hit", e)); - return; - } + executeAsyncWithOrigin(searchRequest, listener.delegateFailureAndWrap((l, searchResponse) -> { + List stats = new ArrayList<>(); + String previousId = null; + for (SearchHit hit : searchResponse.getHits().getHits()) { + // skip old versions + if (hit.getId().equals(previousId) == false) { + previousId = hit.getId(); + try (XContentParser parser = createParser(hit)) { + stats.add(TransformStoredDoc.fromXContent(parser)); + } catch (IOException e) { + l.onFailure(new ElasticsearchParseException("failed to parse transform stats from search hit", e)); + return; } } - - listener.onResponse(stats); - }, listener::onFailure), - client::search - ); + } + l.onResponse(stats); + }), client::search); } @Override public void refresh(ActionListener listener) { executeAsyncWithOrigin( - client.threadPool().getThreadContext(), - TRANSFORM_ORIGIN, new RefreshRequest(TransformInternalIndexConstants.LATEST_INDEX_NAME), - ActionListener.wrap(r -> listener.onResponse(true), listener::onFailure), + listener.delegateFailureAndWrap((l, r) -> l.onResponse(true)), client.admin().indices()::refresh ); } + private void executeAsyncWithOrigin( + Request request, + ActionListener listener, + BiConsumer> consumer + ) { + ClientHelper.executeAsyncWithOrigin(client.threadPool().getThreadContext(), TRANSFORM_ORIGIN, request, listener, consumer); + } + + private void executeAsyncWithOrigin( + ActionType action, + Request request, + ActionListener listener + ) { + ClientHelper.executeAsyncWithOrigin(client, TRANSFORM_ORIGIN, action, request, listener); + } + private void parseTransformLenientlyFromSource( BytesReference source, String transformId, @@ -950,51 +899,45 @@ private void recursiveExpandAllTransformIds( ) .request(); - executeAsyncWithOrigin( - client.threadPool().getThreadContext(), - TRANSFORM_ORIGIN, - request, - ActionListener.wrap(searchResponse -> { - long totalHits = total; - String idOfLastHit = lastId; - - for (SearchHit hit : searchResponse.getHits().getHits()) { - String id = hit.field(TransformField.ID.getPreferredName()).getValue(); - - // paranoia - if (Strings.isNullOrEmpty(id)) { - continue; - } + executeAsyncWithOrigin(request, listener.delegateFailureAndWrap((l, searchResponse) -> { + long totalHits = total; + String idOfLastHit = lastId; - // only count hits if looking for outdated transforms - if (filterForOutdated && hit.getIndex().equals(TransformInternalIndexConstants.LATEST_INDEX_VERSIONED_NAME)) { - ++totalHits; - } else if (id.equals(idOfLastHit) == false && collectedIds.add(id)) { - ++totalHits; - } - idOfLastHit = id; + for (SearchHit hit : searchResponse.getHits().getHits()) { + String id = hit.field(TransformField.ID.getPreferredName()).getValue(); + + // paranoia + if (Strings.isNullOrEmpty(id)) { + continue; } - if (searchResponse.getHits().getHits().length == page.getSize()) { - PageParams nextPage = new PageParams(page.getFrom() + page.getSize(), maxResultWindow); - - recursiveExpandAllTransformIds( - collectedIds, - totalHits, - filterForOutdated, - maxResultWindow, - idOfLastHit, - nextPage, - timeout, - listener - ); - return; + // only count hits if looking for outdated transforms + if (filterForOutdated && hit.getIndex().equals(TransformInternalIndexConstants.LATEST_INDEX_VERSIONED_NAME)) { + ++totalHits; + } else if (id.equals(idOfLastHit) == false && collectedIds.add(id)) { + ++totalHits; } + idOfLastHit = id; + } - listener.onResponse(new Tuple<>(totalHits, collectedIds)); - }, listener::onFailure), - client::search - ); + if (searchResponse.getHits().getHits().length == page.getSize()) { + PageParams nextPage = new PageParams(page.getFrom() + page.getSize(), maxResultWindow); + + recursiveExpandAllTransformIds( + collectedIds, + totalHits, + filterForOutdated, + maxResultWindow, + idOfLastHit, + nextPage, + timeout, + l + ); + return; + } + + l.onResponse(new Tuple<>(totalHits, collectedIds)); + }), client::search); } private static Tuple getStatusAndReason(final BulkByScrollResponse response) { From 263a017ca9d7f9c3cc25f6bd5f6f4514b9129523 Mon Sep 17 00:00:00 2001 From: Benjamin Trent Date: Wed, 6 Mar 2024 08:46:13 -0500 Subject: [PATCH 26/27] Fix bug when nested knn pre-filter might match nested docs (#105994) When using a pre-filter with nested kNN vectors, its treated like a top-level filter. Meaning, it is applied over parent document fields. However, there are times when a query filter is applied that may or may not match internal nested or non-nested docs. We failed to handle this case correctly and users would receive an error. closes: https://github.com/elastic/elasticsearch/issues/105901 --- docs/changelog/105994.yaml | 5 ++ .../search.vectors/100_knn_nested_search.yml | 84 ++++++++++++++++++ .../130_knn_query_nested_search.yml | 87 +++++++++++++++++++ .../search/vectors/KnnVectorQueryBuilder.java | 23 +++-- 4 files changed, 193 insertions(+), 6 deletions(-) create mode 100644 docs/changelog/105994.yaml diff --git a/docs/changelog/105994.yaml b/docs/changelog/105994.yaml new file mode 100644 index 0000000000000..ef9889d0a47af --- /dev/null +++ b/docs/changelog/105994.yaml @@ -0,0 +1,5 @@ +pr: 105994 +summary: Fix bug when nested knn pre-filter might match nested docs +area: Vector Search +type: bug +issues: [] diff --git a/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/100_knn_nested_search.yml b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/100_knn_nested_search.yml index c69e22d274c8e..6c6c75990b0f5 100644 --- a/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/100_knn_nested_search.yml +++ b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/100_knn_nested_search.yml @@ -323,3 +323,87 @@ setup: - match: {hits.total.value: 3} - is_true : profile +--- +"nested kNN search with filter that might match nested docs": + - skip: + version: ' - 8.13.99' + reason: 'bugfix for matching non-nested docs in 8.14' + + - do: + indices.create: + index: nested_text + body: + mappings: + properties: + range: + type: long + other_nested_thing: + type: nested + properties: + text: + type: text + paragraphs: + type: nested + properties: + other_nested_thing: + type: nested + properties: + text: + type: text + vector: + type: dense_vector + dims: 2 + index: true + similarity: cosine + vector: + type: dense_vector + dims: 2 + index: true + similarity: cosine + - do: + index: + index: nested_text + id: "1" + body: + publish_date: "1" + paragraphs: + - vector: [1, 1] + text: "some text" + - vector: [1, 2] + text: "some text" + other_nested_thing: + - text: "some text" + vector: [1, 2] + - do: + index: + index: nested_text + id: "2" + body: + paragraphs: + - vector: [2, 1] + text: "some text" + - vector: [2, 2] + text: "some text" + other_nested_thing: + - text: "some text" + vector: [ 1, 2 ] + - do: + indices.refresh: {} + + - do: + search: + index: nested_text + body: + knn: + field: paragraphs.vector + query_vector: [1, 2] + num_candidates: 10 + k: 10 + filter: + bool: + must_not: + exists: + field: publish_date + + - match: {hits.total.value: 1} + - match: {hits.hits.0._id: "2"} diff --git a/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/130_knn_query_nested_search.yml b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/130_knn_query_nested_search.yml index 5d07c0c8b5f9d..53cc7eb064270 100644 --- a/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/130_knn_query_nested_search.yml +++ b/rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/search.vectors/130_knn_query_nested_search.yml @@ -319,3 +319,90 @@ setup: # Rabbit only has one passage vector - match: {hits.hits.4.fields.name.0: "rabbit.jpg"} - length: { hits.hits.4.inner_hits.nested.hits.hits: 1 } +--- +"nested kNN query search with filter that might match nested docs": + - skip: + version: ' - 8.13.99' + reason: 'bugfix for matching non-nested docs in 8.14' + + - do: + indices.create: + index: nested_text + body: + mappings: + properties: + range: + type: long + other_nested_thing: + type: nested + properties: + text: + type: text + paragraphs: + type: nested + properties: + other_nested_thing: + type: nested + properties: + text: + type: text + vector: + type: dense_vector + dims: 2 + index: true + similarity: cosine + vector: + type: dense_vector + dims: 2 + index: true + similarity: cosine + - do: + index: + index: nested_text + id: "1" + body: + publish_date: "1" + paragraphs: + - vector: [1, 1] + text: "some text" + - vector: [1, 2] + text: "some text" + other_nested_thing: + - text: "some text" + vector: [1, 2] + - do: + index: + index: nested_text + id: "2" + body: + paragraphs: + - vector: [2, 1] + text: "some text" + - vector: [2, 2] + text: "some text" + other_nested_thing: + - text: "some text" + vector: [ 1, 2 ] + - do: + indices.refresh: {} + + - do: + search: + index: nested_text + body: + query: + nested: + path: paragraphs + query: + knn: + field: paragraphs.vector + query_vector: [1, 2] + num_candidates: 10 + filter: + bool: + must_not: + exists: + field: publish_date + + - match: {hits.total.value: 1} + - match: {hits.hits.0._id: "2"} diff --git a/server/src/main/java/org/elasticsearch/search/vectors/KnnVectorQueryBuilder.java b/server/src/main/java/org/elasticsearch/search/vectors/KnnVectorQueryBuilder.java index 96a16013ab874..7e65cd19638ce 100644 --- a/server/src/main/java/org/elasticsearch/search/vectors/KnnVectorQueryBuilder.java +++ b/server/src/main/java/org/elasticsearch/search/vectors/KnnVectorQueryBuilder.java @@ -28,6 +28,7 @@ import org.elasticsearch.index.query.QueryBuilder; import org.elasticsearch.index.query.QueryRewriteContext; import org.elasticsearch.index.query.SearchExecutionContext; +import org.elasticsearch.index.search.NestedHelper; import org.elasticsearch.xcontent.ConstructingObjectParser; import org.elasticsearch.xcontent.ObjectParser; import org.elasticsearch.xcontent.ParseField; @@ -274,7 +275,6 @@ protected Query doToQuery(SearchExecutionContext context) throws IOException { ); } - final BitSetProducer parentFilter; BooleanQuery.Builder builder = new BooleanQuery.Builder(); for (QueryBuilder query : this.filterQueries) { builder.add(query.toQuery(context), BooleanClause.Occur.FILTER); @@ -289,6 +289,8 @@ protected Query doToQuery(SearchExecutionContext context) throws IOException { String parentPath = context.nestedLookup().getNestedParent(fieldName); if (parentPath != null) { + final BitSetProducer parentBitSet; + final Query parentFilter; NestedObjectMapper originalObjectMapper = context.nestedScope().getObjectMapper(); if (originalObjectMapper != null) { try { @@ -296,19 +298,28 @@ protected Query doToQuery(SearchExecutionContext context) throws IOException { context.nestedScope().previousLevel(); NestedObjectMapper objectMapper = context.nestedScope().getObjectMapper(); parentFilter = objectMapper == null - ? context.bitsetFilter(Queries.newNonNestedFilter(context.indexVersionCreated())) - : context.bitsetFilter(objectMapper.nestedTypeFilter()); + ? Queries.newNonNestedFilter(context.indexVersionCreated()) + : objectMapper.nestedTypeFilter(); } finally { context.nestedScope().nextLevel(originalObjectMapper); } } else { // we are NOT in a nested context, coming from the top level knn search - parentFilter = context.bitsetFilter(Queries.newNonNestedFilter(context.indexVersionCreated())); + parentFilter = Queries.newNonNestedFilter(context.indexVersionCreated()); } + parentBitSet = context.bitsetFilter(parentFilter); if (filterQuery != null) { - filterQuery = new ToChildBlockJoinQuery(filterQuery, parentFilter); + NestedHelper nestedHelper = new NestedHelper(context.nestedLookup(), context::isFieldMapped); + // We treat the provided filter as a filter over PARENT documents, so if it might match nested documents + // we need to adjust it. + if (nestedHelper.mightMatchNestedDocs(filterQuery)) { + // Ensure that the query only returns parent documents matching `filterQuery` + filterQuery = Queries.filtered(filterQuery, parentFilter); + } + // Now join the filterQuery & parentFilter to provide the matching blocks of children + filterQuery = new ToChildBlockJoinQuery(filterQuery, parentBitSet); } - return vectorFieldType.createKnnQuery(queryVector, adjustedNumCands, filterQuery, vectorSimilarity, parentFilter); + return vectorFieldType.createKnnQuery(queryVector, adjustedNumCands, filterQuery, vectorSimilarity, parentBitSet); } return vectorFieldType.createKnnQuery(queryVector, adjustedNumCands, filterQuery, vectorSimilarity, null); } From 78df37cbf934380aba8e39cb6d62398593bb25bc Mon Sep 17 00:00:00 2001 From: Nikolaj Volgushev Date: Wed, 6 Mar 2024 14:50:48 +0100 Subject: [PATCH 27/27] API key APIs customization (#105671) This PR makes request translation and building customizable for our API Key Management APIs, create, update, and bulk update. Relates: ES-7827 --- .../apikey/BulkUpdateApiKeyRequest.java | 1 + .../BulkUpdateApiKeyRequestTranslator.java | 63 ++++++++++++++++++ .../apikey/CreateApiKeyRequestBuilder.java | 52 ++++++++------- .../CreateApiKeyRequestBuilderFactory.java | 22 +++++++ .../action/apikey/UpdateApiKeyRequest.java | 1 + .../apikey/UpdateApiKeyRequestTranslator.java | 66 +++++++++++++++++++ .../xpack/security/Security.java | 25 +++++-- .../apikey/RestBulkUpdateApiKeyAction.java | 43 +++--------- .../action/apikey/RestCreateApiKeyAction.java | 18 ++--- .../action/apikey/RestUpdateApiKeyAction.java | 47 +++---------- .../apikey/RestCreateApiKeyActionTests.java | 7 +- .../apikey/RestUpdateApiKeyActionTests.java | 3 +- 12 files changed, 242 insertions(+), 106 deletions(-) create mode 100644 x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/security/action/apikey/BulkUpdateApiKeyRequestTranslator.java create mode 100644 x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/security/action/apikey/CreateApiKeyRequestBuilderFactory.java create mode 100644 x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/security/action/apikey/UpdateApiKeyRequestTranslator.java diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/security/action/apikey/BulkUpdateApiKeyRequest.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/security/action/apikey/BulkUpdateApiKeyRequest.java index f915781c6211a..534c874438e3f 100644 --- a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/security/action/apikey/BulkUpdateApiKeyRequest.java +++ b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/security/action/apikey/BulkUpdateApiKeyRequest.java @@ -49,4 +49,5 @@ public BulkUpdateApiKeyRequest(StreamInput in) throws IOException { public ApiKey.Type getType() { return ApiKey.Type.REST; } + } diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/security/action/apikey/BulkUpdateApiKeyRequestTranslator.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/security/action/apikey/BulkUpdateApiKeyRequestTranslator.java new file mode 100644 index 0000000000000..57a5848970b2e --- /dev/null +++ b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/security/action/apikey/BulkUpdateApiKeyRequestTranslator.java @@ -0,0 +1,63 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.core.security.action.apikey; + +import org.elasticsearch.common.CheckedBiFunction; +import org.elasticsearch.core.TimeValue; +import org.elasticsearch.rest.RestRequest; +import org.elasticsearch.xcontent.ConstructingObjectParser; +import org.elasticsearch.xcontent.ParseField; +import org.elasticsearch.xcontent.XContentParser; +import org.elasticsearch.xpack.core.security.authz.RoleDescriptor; + +import java.io.IOException; +import java.util.List; +import java.util.Map; + +import static org.elasticsearch.xcontent.ConstructingObjectParser.constructorArg; +import static org.elasticsearch.xcontent.ConstructingObjectParser.optionalConstructorArg; + +public interface BulkUpdateApiKeyRequestTranslator { + BulkUpdateApiKeyRequest translate(RestRequest request) throws IOException; + + class Default implements BulkUpdateApiKeyRequestTranslator { + private static final ConstructingObjectParser PARSER = createParser( + (n, p) -> RoleDescriptor.parse(n, p, false) + ); + + @SuppressWarnings("unchecked") + protected static ConstructingObjectParser createParser( + CheckedBiFunction roleDescriptorParser + ) { + final ConstructingObjectParser parser = new ConstructingObjectParser<>( + "bulk_update_api_key_request", + a -> new BulkUpdateApiKeyRequest( + (List) a[0], + (List) a[1], + (Map) a[2], + TimeValue.parseTimeValue((String) a[3], null, "expiration") + ) + ); + parser.declareStringArray(constructorArg(), new ParseField("ids")); + parser.declareNamedObjects(optionalConstructorArg(), (p, c, n) -> { + p.nextToken(); + return roleDescriptorParser.apply(n, p); + }, new ParseField("role_descriptors")); + parser.declareObject(optionalConstructorArg(), (p, c) -> p.map(), new ParseField("metadata")); + parser.declareString(optionalConstructorArg(), new ParseField("expiration")); + return parser; + } + + @Override + public BulkUpdateApiKeyRequest translate(RestRequest request) throws IOException { + try (XContentParser parser = request.contentParser()) { + return PARSER.parse(parser, null); + } + } + } +} diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/security/action/apikey/CreateApiKeyRequestBuilder.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/security/action/apikey/CreateApiKeyRequestBuilder.java index 2747dc47058f8..5c156ab4e6166 100644 --- a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/security/action/apikey/CreateApiKeyRequestBuilder.java +++ b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/security/action/apikey/CreateApiKeyRequestBuilder.java @@ -9,6 +9,7 @@ import org.elasticsearch.action.ActionRequestBuilder; import org.elasticsearch.action.support.WriteRequest; import org.elasticsearch.client.internal.ElasticsearchClient; +import org.elasticsearch.common.CheckedBiFunction; import org.elasticsearch.common.bytes.BytesReference; import org.elasticsearch.common.xcontent.LoggingDeprecationHandler; import org.elasticsearch.common.xcontent.XContentHelper; @@ -29,30 +30,34 @@ /** * Request builder for populating a {@link CreateApiKeyRequest} */ -public final class CreateApiKeyRequestBuilder extends ActionRequestBuilder { +public class CreateApiKeyRequestBuilder extends ActionRequestBuilder { + private static final ConstructingObjectParser PARSER = createParser( + (n, p) -> RoleDescriptor.parse(n, p, false) + ); @SuppressWarnings("unchecked") - static final ConstructingObjectParser PARSER = new ConstructingObjectParser<>( - "api_key_request", - false, - (args, v) -> { - return new CreateApiKeyRequest( + protected static ConstructingObjectParser createParser( + CheckedBiFunction roleDescriptorParser + ) { + ConstructingObjectParser parser = new ConstructingObjectParser<>( + "api_key_request", + false, + (args, v) -> new CreateApiKeyRequest( (String) args[0], (List) args[1], TimeValue.parseTimeValue((String) args[2], null, "expiration"), (Map) args[3] - ); - } - ); + ) + ); - static { - PARSER.declareString(constructorArg(), new ParseField("name")); - PARSER.declareNamedObjects(optionalConstructorArg(), (p, c, n) -> { + parser.declareString(constructorArg(), new ParseField("name")); + parser.declareNamedObjects(optionalConstructorArg(), (p, c, n) -> { p.nextToken(); - return RoleDescriptor.parse(n, p, false); + return roleDescriptorParser.apply(n, p); }, new ParseField("role_descriptors")); - PARSER.declareString(optionalConstructorArg(), new ParseField("expiration")); - PARSER.declareObject(optionalConstructorArg(), (p, c) -> p.map(), new ParseField("metadata")); + parser.declareString(optionalConstructorArg(), new ParseField("expiration")); + parser.declareObject(optionalConstructorArg(), (p, c) -> p.map(), new ParseField("metadata")); + return parser; } public CreateApiKeyRequestBuilder(ElasticsearchClient client) { @@ -85,6 +90,15 @@ public CreateApiKeyRequestBuilder setMetadata(Map metadata) { } public CreateApiKeyRequestBuilder source(BytesReference source, XContentType xContentType) throws IOException { + CreateApiKeyRequest createApiKeyRequest = parse(source, xContentType); + setName(createApiKeyRequest.getName()); + setRoleDescriptors(createApiKeyRequest.getRoleDescriptors()); + setExpiration(createApiKeyRequest.getExpiration()); + setMetadata(createApiKeyRequest.getMetadata()); + return this; + } + + protected CreateApiKeyRequest parse(BytesReference source, XContentType xContentType) throws IOException { try ( XContentParser parser = XContentHelper.createParserNotCompressed( LoggingDeprecationHandler.XCONTENT_PARSER_CONFIG, @@ -92,14 +106,8 @@ public CreateApiKeyRequestBuilder source(BytesReference source, XContentType xCo xContentType ) ) { - CreateApiKeyRequest createApiKeyRequest = parse(parser); - setName(createApiKeyRequest.getName()); - setRoleDescriptors(createApiKeyRequest.getRoleDescriptors()); - setExpiration(createApiKeyRequest.getExpiration()); - setMetadata(createApiKeyRequest.getMetadata()); - + return parse(parser); } - return this; } public static CreateApiKeyRequest parse(XContentParser parser) throws IOException { diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/security/action/apikey/CreateApiKeyRequestBuilderFactory.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/security/action/apikey/CreateApiKeyRequestBuilderFactory.java new file mode 100644 index 0000000000000..ff5592e339634 --- /dev/null +++ b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/security/action/apikey/CreateApiKeyRequestBuilderFactory.java @@ -0,0 +1,22 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.core.security.action.apikey; + +import org.elasticsearch.client.internal.Client; + +public interface CreateApiKeyRequestBuilderFactory { + CreateApiKeyRequestBuilder create(Client client, boolean restrictRequest); + + class Default implements CreateApiKeyRequestBuilderFactory { + @Override + public CreateApiKeyRequestBuilder create(Client client, boolean restrictRequest) { + assert false == restrictRequest; + return new CreateApiKeyRequestBuilder(client); + } + } +} diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/security/action/apikey/UpdateApiKeyRequest.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/security/action/apikey/UpdateApiKeyRequest.java index c5c8bcc4fc87a..9b1e9194d59fd 100644 --- a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/security/action/apikey/UpdateApiKeyRequest.java +++ b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/security/action/apikey/UpdateApiKeyRequest.java @@ -38,4 +38,5 @@ public UpdateApiKeyRequest(StreamInput in) throws IOException { public ApiKey.Type getType() { return ApiKey.Type.REST; } + } diff --git a/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/security/action/apikey/UpdateApiKeyRequestTranslator.java b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/security/action/apikey/UpdateApiKeyRequestTranslator.java new file mode 100644 index 0000000000000..f70732dd50990 --- /dev/null +++ b/x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/security/action/apikey/UpdateApiKeyRequestTranslator.java @@ -0,0 +1,66 @@ +/* + * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one + * or more contributor license agreements. Licensed under the Elastic License + * 2.0; you may not use this file except in compliance with the Elastic License + * 2.0. + */ + +package org.elasticsearch.xpack.core.security.action.apikey; + +import org.elasticsearch.common.CheckedBiFunction; +import org.elasticsearch.core.TimeValue; +import org.elasticsearch.rest.RestRequest; +import org.elasticsearch.xcontent.ConstructingObjectParser; +import org.elasticsearch.xcontent.ParseField; +import org.elasticsearch.xcontent.XContentParser; +import org.elasticsearch.xpack.core.security.authz.RoleDescriptor; + +import java.io.IOException; +import java.util.List; +import java.util.Map; + +import static org.elasticsearch.xcontent.ConstructingObjectParser.optionalConstructorArg; + +public interface UpdateApiKeyRequestTranslator { + UpdateApiKeyRequest translate(RestRequest request) throws IOException; + + class Default implements UpdateApiKeyRequestTranslator { + private static final ConstructingObjectParser PARSER = createParser((n, p) -> RoleDescriptor.parse(n, p, false)); + + @SuppressWarnings("unchecked") + protected static ConstructingObjectParser createParser( + CheckedBiFunction roleDescriptorParser + ) { + final ConstructingObjectParser parser = new ConstructingObjectParser<>( + "update_api_key_request_payload", + a -> new Payload( + (List) a[0], + (Map) a[1], + TimeValue.parseTimeValue((String) a[2], null, "expiration") + ) + ); + parser.declareNamedObjects(optionalConstructorArg(), (p, c, n) -> { + p.nextToken(); + return roleDescriptorParser.apply(n, p); + }, new ParseField("role_descriptors")); + parser.declareObject(optionalConstructorArg(), (p, c) -> p.map(), new ParseField("metadata")); + parser.declareString(optionalConstructorArg(), new ParseField("expiration")); + return parser; + } + + @Override + public UpdateApiKeyRequest translate(RestRequest request) throws IOException { + // Note that we use `ids` here even though we only support a single ID. This is because the route where this translator is used + // shares a path prefix with `RestClearApiKeyCacheAction` and our current REST implementation requires that path params have the + // same wildcard if their paths share a prefix + final String apiKeyId = request.param("ids"); + if (false == request.hasContent()) { + return UpdateApiKeyRequest.usingApiKeyId(apiKeyId); + } + final Payload payload = PARSER.parse(request.contentParser(), null); + return new UpdateApiKeyRequest(apiKeyId, payload.roleDescriptors, payload.metadata, payload.expiration); + } + + protected record Payload(List roleDescriptors, Map metadata, TimeValue expiration) {} + } +} diff --git a/x-pack/plugin/security/src/main/java/org/elasticsearch/xpack/security/Security.java b/x-pack/plugin/security/src/main/java/org/elasticsearch/xpack/security/Security.java index 3beff69849a58..219f645a92bbe 100644 --- a/x-pack/plugin/security/src/main/java/org/elasticsearch/xpack/security/Security.java +++ b/x-pack/plugin/security/src/main/java/org/elasticsearch/xpack/security/Security.java @@ -119,13 +119,16 @@ import org.elasticsearch.xpack.core.security.action.ClearSecurityCacheAction; import org.elasticsearch.xpack.core.security.action.DelegatePkiAuthenticationAction; import org.elasticsearch.xpack.core.security.action.apikey.BulkUpdateApiKeyAction; +import org.elasticsearch.xpack.core.security.action.apikey.BulkUpdateApiKeyRequestTranslator; import org.elasticsearch.xpack.core.security.action.apikey.CreateApiKeyAction; +import org.elasticsearch.xpack.core.security.action.apikey.CreateApiKeyRequestBuilderFactory; import org.elasticsearch.xpack.core.security.action.apikey.CreateCrossClusterApiKeyAction; import org.elasticsearch.xpack.core.security.action.apikey.GetApiKeyAction; import org.elasticsearch.xpack.core.security.action.apikey.GrantApiKeyAction; import org.elasticsearch.xpack.core.security.action.apikey.InvalidateApiKeyAction; import org.elasticsearch.xpack.core.security.action.apikey.QueryApiKeyAction; import org.elasticsearch.xpack.core.security.action.apikey.UpdateApiKeyAction; +import org.elasticsearch.xpack.core.security.action.apikey.UpdateApiKeyRequestTranslator; import org.elasticsearch.xpack.core.security.action.apikey.UpdateCrossClusterApiKeyAction; import org.elasticsearch.xpack.core.security.action.enrollment.KibanaEnrollmentAction; import org.elasticsearch.xpack.core.security.action.enrollment.NodeEnrollmentAction; @@ -561,6 +564,9 @@ public class Security extends Plugin private final SetOnce scriptServiceReference = new SetOnce<>(); private final SetOnce operatorOnlyRegistry = new SetOnce<>(); private final SetOnce putRoleRequestBuilderFactory = new SetOnce<>(); + private final SetOnce createApiKeyRequestBuilderFactory = new SetOnce<>(); + private final SetOnce updateApiKeyRequestTranslator = new SetOnce<>(); + private final SetOnce bulkUpdateApiKeyRequestTranslator = new SetOnce<>(); private final SetOnce getBuiltinPrivilegesResponseTranslator = new SetOnce<>(); private final SetOnce fileRolesStore = new SetOnce<>(); private final SetOnce operatorPrivilegesService = new SetOnce<>(); @@ -821,10 +827,18 @@ Collection createComponents( if (putRoleRequestBuilderFactory.get() == null) { putRoleRequestBuilderFactory.set(new PutRoleRequestBuilderFactory.Default()); } - + if (createApiKeyRequestBuilderFactory.get() == null) { + createApiKeyRequestBuilderFactory.set(new CreateApiKeyRequestBuilderFactory.Default()); + } if (getBuiltinPrivilegesResponseTranslator.get() == null) { getBuiltinPrivilegesResponseTranslator.set(new GetBuiltinPrivilegesResponseTranslator.Default()); } + if (updateApiKeyRequestTranslator.get() == null) { + updateApiKeyRequestTranslator.set(new UpdateApiKeyRequestTranslator.Default()); + } + if (bulkUpdateApiKeyRequestTranslator.get() == null) { + bulkUpdateApiKeyRequestTranslator.set(new BulkUpdateApiKeyRequestTranslator.Default()); + } final Map, ActionListener>>> customRoleProviders = new LinkedHashMap<>(); for (SecurityExtension extension : securityExtensions) { @@ -1456,10 +1470,10 @@ public List getRestHandlers( new RestGetPrivilegesAction(settings, getLicenseState()), new RestPutPrivilegesAction(settings, getLicenseState()), new RestDeletePrivilegesAction(settings, getLicenseState()), - new RestCreateApiKeyAction(settings, getLicenseState()), + new RestCreateApiKeyAction(settings, getLicenseState(), createApiKeyRequestBuilderFactory.get()), new RestCreateCrossClusterApiKeyAction(settings, getLicenseState()), - new RestUpdateApiKeyAction(settings, getLicenseState()), - new RestBulkUpdateApiKeyAction(settings, getLicenseState()), + new RestUpdateApiKeyAction(settings, getLicenseState(), updateApiKeyRequestTranslator.get()), + new RestBulkUpdateApiKeyAction(settings, getLicenseState(), bulkUpdateApiKeyRequestTranslator.get()), new RestUpdateCrossClusterApiKeyAction(settings, getLicenseState()), new RestGrantApiKeyAction(settings, getLicenseState()), new RestInvalidateApiKeyAction(settings, getLicenseState()), @@ -2039,6 +2053,9 @@ public void loadExtensions(ExtensionLoader loader) { loadSingletonExtensionAndSetOnce(loader, operatorOnlyRegistry, OperatorOnlyRegistry.class); loadSingletonExtensionAndSetOnce(loader, putRoleRequestBuilderFactory, PutRoleRequestBuilderFactory.class); loadSingletonExtensionAndSetOnce(loader, getBuiltinPrivilegesResponseTranslator, GetBuiltinPrivilegesResponseTranslator.class); + loadSingletonExtensionAndSetOnce(loader, updateApiKeyRequestTranslator, UpdateApiKeyRequestTranslator.class); + loadSingletonExtensionAndSetOnce(loader, bulkUpdateApiKeyRequestTranslator, BulkUpdateApiKeyRequestTranslator.class); + loadSingletonExtensionAndSetOnce(loader, createApiKeyRequestBuilderFactory, CreateApiKeyRequestBuilderFactory.class); } private void loadSingletonExtensionAndSetOnce(ExtensionLoader loader, SetOnce setOnce, Class clazz) { diff --git a/x-pack/plugin/security/src/main/java/org/elasticsearch/xpack/security/rest/action/apikey/RestBulkUpdateApiKeyAction.java b/x-pack/plugin/security/src/main/java/org/elasticsearch/xpack/security/rest/action/apikey/RestBulkUpdateApiKeyAction.java index 584ad08704ddd..97ee7cc50a7d5 100644 --- a/x-pack/plugin/security/src/main/java/org/elasticsearch/xpack/security/rest/action/apikey/RestBulkUpdateApiKeyAction.java +++ b/x-pack/plugin/security/src/main/java/org/elasticsearch/xpack/security/rest/action/apikey/RestBulkUpdateApiKeyAction.java @@ -9,53 +9,32 @@ import org.elasticsearch.client.internal.node.NodeClient; import org.elasticsearch.common.settings.Settings; -import org.elasticsearch.core.TimeValue; import org.elasticsearch.license.XPackLicenseState; import org.elasticsearch.rest.RestRequest; import org.elasticsearch.rest.Scope; import org.elasticsearch.rest.ServerlessScope; import org.elasticsearch.rest.action.RestToXContentListener; -import org.elasticsearch.xcontent.ConstructingObjectParser; -import org.elasticsearch.xcontent.ParseField; -import org.elasticsearch.xcontent.XContentParser; import org.elasticsearch.xpack.core.security.action.apikey.BulkUpdateApiKeyAction; import org.elasticsearch.xpack.core.security.action.apikey.BulkUpdateApiKeyRequest; -import org.elasticsearch.xpack.core.security.authz.RoleDescriptor; +import org.elasticsearch.xpack.core.security.action.apikey.BulkUpdateApiKeyRequestTranslator; import java.io.IOException; import java.util.List; -import java.util.Map; import static org.elasticsearch.rest.RestRequest.Method.POST; -import static org.elasticsearch.xcontent.ConstructingObjectParser.constructorArg; -import static org.elasticsearch.xcontent.ConstructingObjectParser.optionalConstructorArg; @ServerlessScope(Scope.PUBLIC) public final class RestBulkUpdateApiKeyAction extends ApiKeyBaseRestHandler { - @SuppressWarnings("unchecked") - static final ConstructingObjectParser PARSER = new ConstructingObjectParser<>( - "bulk_update_api_key_request", - a -> new BulkUpdateApiKeyRequest( - (List) a[0], - (List) a[1], - (Map) a[2], - TimeValue.parseTimeValue((String) a[3], null, "expiration") - ) - ); + private final BulkUpdateApiKeyRequestTranslator requestTranslator; - static { - PARSER.declareStringArray(constructorArg(), new ParseField("ids")); - PARSER.declareNamedObjects(optionalConstructorArg(), (p, c, n) -> { - p.nextToken(); - return RoleDescriptor.parse(n, p, false); - }, new ParseField("role_descriptors")); - PARSER.declareObject(optionalConstructorArg(), (p, c) -> p.map(), new ParseField("metadata")); - PARSER.declareString(optionalConstructorArg(), new ParseField("expiration")); - } - - public RestBulkUpdateApiKeyAction(final Settings settings, final XPackLicenseState licenseState) { + public RestBulkUpdateApiKeyAction( + final Settings settings, + final XPackLicenseState licenseState, + final BulkUpdateApiKeyRequestTranslator requestTranslator + ) { super(settings, licenseState); + this.requestTranslator = requestTranslator; } @Override @@ -70,9 +49,7 @@ public String getName() { @Override protected RestChannelConsumer innerPrepareRequest(final RestRequest request, final NodeClient client) throws IOException { - try (XContentParser parser = request.contentParser()) { - final BulkUpdateApiKeyRequest parsed = PARSER.parse(parser, null); - return channel -> client.execute(BulkUpdateApiKeyAction.INSTANCE, parsed, new RestToXContentListener<>(channel)); - } + final BulkUpdateApiKeyRequest bulkUpdateApiKeyRequest = requestTranslator.translate(request); + return channel -> client.execute(BulkUpdateApiKeyAction.INSTANCE, bulkUpdateApiKeyRequest, new RestToXContentListener<>(channel)); } } diff --git a/x-pack/plugin/security/src/main/java/org/elasticsearch/xpack/security/rest/action/apikey/RestCreateApiKeyAction.java b/x-pack/plugin/security/src/main/java/org/elasticsearch/xpack/security/rest/action/apikey/RestCreateApiKeyAction.java index 2cb5a15f1e0f2..217afdb3cfea2 100644 --- a/x-pack/plugin/security/src/main/java/org/elasticsearch/xpack/security/rest/action/apikey/RestCreateApiKeyAction.java +++ b/x-pack/plugin/security/src/main/java/org/elasticsearch/xpack/security/rest/action/apikey/RestCreateApiKeyAction.java @@ -16,6 +16,7 @@ import org.elasticsearch.rest.ServerlessScope; import org.elasticsearch.rest.action.RestToXContentListener; import org.elasticsearch.xpack.core.security.action.apikey.CreateApiKeyRequestBuilder; +import org.elasticsearch.xpack.core.security.action.apikey.CreateApiKeyRequestBuilderFactory; import org.elasticsearch.xpack.security.authc.ApiKeyService; import java.io.IOException; @@ -30,13 +31,16 @@ @ServerlessScope(Scope.PUBLIC) public final class RestCreateApiKeyAction extends ApiKeyBaseRestHandler { + private final CreateApiKeyRequestBuilderFactory builderFactory; + /** - * @param settings the node's settings - * @param licenseState the license state that will be used to determine if - * security is licensed + * @param settings the node's settings + * @param licenseState the license state that will be used to determine if + * security is licensed */ - public RestCreateApiKeyAction(Settings settings, XPackLicenseState licenseState) { + public RestCreateApiKeyAction(Settings settings, XPackLicenseState licenseState, CreateApiKeyRequestBuilderFactory builderFactory) { super(settings, licenseState); + this.builderFactory = builderFactory; } @Override @@ -51,10 +55,8 @@ public String getName() { @Override protected RestChannelConsumer innerPrepareRequest(final RestRequest request, final NodeClient client) throws IOException { - CreateApiKeyRequestBuilder builder = new CreateApiKeyRequestBuilder(client).source( - request.requiredContent(), - request.getXContentType() - ); + CreateApiKeyRequestBuilder builder = builderFactory.create(client, request.hasParam(RestRequest.PATH_RESTRICTED)) + .source(request.requiredContent(), request.getXContentType()); String refresh = request.param("refresh"); if (refresh != null) { builder.setRefreshPolicy(WriteRequest.RefreshPolicy.parse(request.param("refresh"))); diff --git a/x-pack/plugin/security/src/main/java/org/elasticsearch/xpack/security/rest/action/apikey/RestUpdateApiKeyAction.java b/x-pack/plugin/security/src/main/java/org/elasticsearch/xpack/security/rest/action/apikey/RestUpdateApiKeyAction.java index d64e7f4007387..0fe0f3df0715f 100644 --- a/x-pack/plugin/security/src/main/java/org/elasticsearch/xpack/security/rest/action/apikey/RestUpdateApiKeyAction.java +++ b/x-pack/plugin/security/src/main/java/org/elasticsearch/xpack/security/rest/action/apikey/RestUpdateApiKeyAction.java @@ -9,49 +9,31 @@ import org.elasticsearch.client.internal.node.NodeClient; import org.elasticsearch.common.settings.Settings; -import org.elasticsearch.core.TimeValue; import org.elasticsearch.license.XPackLicenseState; import org.elasticsearch.rest.RestRequest; import org.elasticsearch.rest.Scope; import org.elasticsearch.rest.ServerlessScope; import org.elasticsearch.rest.action.RestToXContentListener; -import org.elasticsearch.xcontent.ConstructingObjectParser; -import org.elasticsearch.xcontent.ParseField; import org.elasticsearch.xpack.core.security.action.apikey.UpdateApiKeyAction; import org.elasticsearch.xpack.core.security.action.apikey.UpdateApiKeyRequest; -import org.elasticsearch.xpack.core.security.authz.RoleDescriptor; +import org.elasticsearch.xpack.core.security.action.apikey.UpdateApiKeyRequestTranslator; import java.io.IOException; import java.util.List; -import java.util.Map; import static org.elasticsearch.rest.RestRequest.Method.PUT; -import static org.elasticsearch.xcontent.ConstructingObjectParser.optionalConstructorArg; @ServerlessScope(Scope.PUBLIC) public final class RestUpdateApiKeyAction extends ApiKeyBaseRestHandler { + private final UpdateApiKeyRequestTranslator requestTranslator; - @SuppressWarnings("unchecked") - static final ConstructingObjectParser PARSER = new ConstructingObjectParser<>( - "update_api_key_request_payload", - a -> new Payload( - (List) a[0], - (Map) a[1], - TimeValue.parseTimeValue((String) a[2], null, "expiration") - ) - ); - - static { - PARSER.declareNamedObjects(optionalConstructorArg(), (p, c, n) -> { - p.nextToken(); - return RoleDescriptor.parse(n, p, false); - }, new ParseField("role_descriptors")); - PARSER.declareObject(optionalConstructorArg(), (p, c) -> p.map(), new ParseField("metadata")); - PARSER.declareString(optionalConstructorArg(), new ParseField("expiration")); - } - - public RestUpdateApiKeyAction(final Settings settings, final XPackLicenseState licenseState) { + public RestUpdateApiKeyAction( + final Settings settings, + final XPackLicenseState licenseState, + final UpdateApiKeyRequestTranslator requestTranslator + ) { super(settings, licenseState); + this.requestTranslator = requestTranslator; } @Override @@ -66,17 +48,8 @@ public String getName() { @Override protected RestChannelConsumer innerPrepareRequest(final RestRequest request, final NodeClient client) throws IOException { - // Note that we use `ids` here even though we only support a single id. This is because this route shares a path prefix with - // `RestClearApiKeyCacheAction` and our current REST implementation requires that path params have the same wildcard if their paths - // share a prefix - final var apiKeyId = request.param("ids"); - final var payload = request.hasContent() == false ? new Payload(null, null, null) : PARSER.parse(request.contentParser(), null); - return channel -> client.execute( - UpdateApiKeyAction.INSTANCE, - new UpdateApiKeyRequest(apiKeyId, payload.roleDescriptors, payload.metadata, payload.expiration), - new RestToXContentListener<>(channel) - ); + final UpdateApiKeyRequest updateApiKeyRequest = requestTranslator.translate(request); + return channel -> client.execute(UpdateApiKeyAction.INSTANCE, updateApiKeyRequest, new RestToXContentListener<>(channel)); } - record Payload(List roleDescriptors, Map metadata, TimeValue expiration) {} } diff --git a/x-pack/plugin/security/src/test/java/org/elasticsearch/xpack/security/rest/action/apikey/RestCreateApiKeyActionTests.java b/x-pack/plugin/security/src/test/java/org/elasticsearch/xpack/security/rest/action/apikey/RestCreateApiKeyActionTests.java index 0ab9533e62d4c..d487eab9f7887 100644 --- a/x-pack/plugin/security/src/test/java/org/elasticsearch/xpack/security/rest/action/apikey/RestCreateApiKeyActionTests.java +++ b/x-pack/plugin/security/src/test/java/org/elasticsearch/xpack/security/rest/action/apikey/RestCreateApiKeyActionTests.java @@ -31,6 +31,7 @@ import org.elasticsearch.xcontent.XContentType; import org.elasticsearch.xpack.core.security.action.apikey.ApiKey; import org.elasticsearch.xpack.core.security.action.apikey.CreateApiKeyRequest; +import org.elasticsearch.xpack.core.security.action.apikey.CreateApiKeyRequestBuilderFactory; import org.elasticsearch.xpack.core.security.action.apikey.CreateApiKeyResponse; import java.time.Duration; @@ -105,7 +106,11 @@ public void doE } } }; - final RestCreateApiKeyAction restCreateApiKeyAction = new RestCreateApiKeyAction(Settings.EMPTY, mockLicenseState); + final RestCreateApiKeyAction restCreateApiKeyAction = new RestCreateApiKeyAction( + Settings.EMPTY, + mockLicenseState, + new CreateApiKeyRequestBuilderFactory.Default() + ); restCreateApiKeyAction.handleRequest(restRequest, restChannel, client); final RestResponse restResponse = responseSetOnce.get(); diff --git a/x-pack/plugin/security/src/test/java/org/elasticsearch/xpack/security/rest/action/apikey/RestUpdateApiKeyActionTests.java b/x-pack/plugin/security/src/test/java/org/elasticsearch/xpack/security/rest/action/apikey/RestUpdateApiKeyActionTests.java index eb0b7bea1a5fb..c349ad57a486c 100644 --- a/x-pack/plugin/security/src/test/java/org/elasticsearch/xpack/security/rest/action/apikey/RestUpdateApiKeyActionTests.java +++ b/x-pack/plugin/security/src/test/java/org/elasticsearch/xpack/security/rest/action/apikey/RestUpdateApiKeyActionTests.java @@ -16,6 +16,7 @@ import org.elasticsearch.xcontent.XContentType; import org.elasticsearch.xpack.core.XPackSettings; import org.elasticsearch.xpack.core.security.action.apikey.UpdateApiKeyRequest; +import org.elasticsearch.xpack.core.security.action.apikey.UpdateApiKeyRequestTranslator; import org.elasticsearch.xpack.core.security.action.apikey.UpdateApiKeyResponse; import org.junit.Before; @@ -34,7 +35,7 @@ public void init() { final Settings settings = Settings.builder().put(XPackSettings.SECURITY_ENABLED.getKey(), true).build(); final XPackLicenseState licenseState = mock(XPackLicenseState.class); requestHolder = new AtomicReference<>(); - restAction = new RestUpdateApiKeyAction(settings, licenseState); + restAction = new RestUpdateApiKeyAction(settings, licenseState, new UpdateApiKeyRequestTranslator.Default()); controller().registerHandler(restAction); verifyingClient.setExecuteVerifier(((actionType, actionRequest) -> { assertThat(actionRequest, instanceOf(UpdateApiKeyRequest.class));