Skip to content

Commit

Permalink
Update SQL plugin for core refactor (opensearch-project#1571)
Browse files Browse the repository at this point in the history
Signed-off-by: MaxKsyunz <[email protected]>
Signed-off-by: Yury-Fridlyand <[email protected]>
  • Loading branch information
Yury-Fridlyand committed Jun 27, 2023
1 parent 6532a10 commit 8558e90
Show file tree
Hide file tree
Showing 7 changed files with 44 additions and 61 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -12,8 +12,6 @@
import java.util.Objects;
import org.opensearch.cluster.metadata.MappingMetadata;
import org.opensearch.cluster.metadata.Metadata;
import org.opensearch.common.collect.ImmutableOpenMap;
import org.opensearch.sql.legacy.domain.Field;

/**
* Index mappings in the cluster.
Expand Down Expand Up @@ -51,7 +49,7 @@ public IndexMappings(Metadata metaData) {
indexMetaData -> new FieldMappings(indexMetaData.mapping()));
}

public IndexMappings(ImmutableOpenMap<String, MappingMetadata> mappings) {
public IndexMappings(Map<String, MappingMetadata> mappings) {
this.indexMappings = buildMappings(mappings, FieldMappings::new);
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -6,12 +6,10 @@

package org.opensearch.sql.legacy.esdomain.mapping;

import com.carrotsearch.hppc.cursors.ObjectObjectCursor;
import com.google.common.collect.ImmutableMap;
import java.util.Collection;
import java.util.Map;
import java.util.function.Function;
import org.opensearch.common.collect.ImmutableOpenMap;
import java.util.stream.Collectors;

/**
* Mappings interface to provide default implementation (minimal set of Map methods) for subclass in hierarchy.
Expand Down Expand Up @@ -47,13 +45,10 @@ default boolean isEmpty() {
Map<String, T> data();

/**
* Convert OpenSearch ImmutableOpenMap<String, X> to JDK Map<String, Y> by applying function: Y func(X)
* Build a map from an existing map by applying provided function to each value.
*/
default <X, Y> Map<String, Y> buildMappings(ImmutableOpenMap<String, X> mappings, Function<X, Y> func) {
ImmutableMap.Builder<String, Y> builder = ImmutableMap.builder();
for (ObjectObjectCursor<String, X> mapping : mappings) {
builder.put(mapping.key, func.apply(mapping.value));
}
return builder.build();
default <X, Y> Map<String, Y> buildMappings(Map<String, X> mappings, Function<X, Y> func) {
return mappings.entrySet().stream().collect(
Collectors.toUnmodifiableMap(Map.Entry::getKey, func.compose(Map.Entry::getValue)));
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,6 @@

package org.opensearch.sql.legacy.executor.format;

import com.carrotsearch.hppc.cursors.ObjectObjectCursor;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
Expand All @@ -15,7 +14,6 @@
import org.opensearch.action.admin.indices.get.GetIndexResponse;
import org.opensearch.client.Client;
import org.opensearch.cluster.metadata.MappingMetadata;
import org.opensearch.common.collect.ImmutableOpenMap;
import org.opensearch.sql.legacy.domain.IndexStatement;
import org.opensearch.sql.legacy.executor.format.DataRows.Row;
import org.opensearch.sql.legacy.executor.format.Schema.Column;
Expand Down Expand Up @@ -79,14 +77,14 @@ private List<Column> loadColumns() {
private List<Row> loadRows() {
List<Row> rows = new ArrayList<>();
GetIndexResponse indexResponse = (GetIndexResponse) queryResult;
ImmutableOpenMap<String, MappingMetadata> indexMappings = indexResponse.getMappings();
Map<String, MappingMetadata> indexMappings = indexResponse.getMappings();

// Iterate through indices in indexMappings
for (ObjectObjectCursor<String, MappingMetadata> indexCursor : indexMappings) {
String index = indexCursor.key;
for (Entry<String, MappingMetadata> indexCursor : indexMappings.entrySet()) {
String index = indexCursor.getKey();

if (matchesPatternIfRegex(index, statement.getIndexPattern())) {
rows.addAll(loadIndexData(index, indexCursor.value.getSourceAsMap()));
rows.addAll(loadIndexData(index, indexCursor.getValue().getSourceAsMap()));
}
}
return rows;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
import java.io.IOException;
import java.sql.SQLFeatureNotSupportedException;
import java.util.List;
import java.util.Map;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import org.mockito.stubbing.Answer;
Expand Down Expand Up @@ -227,9 +228,9 @@ public static ClusterService mockClusterService(String mappings) {
when(mockService.state()).thenReturn(mockState);
when(mockState.metadata()).thenReturn(mockMetaData);
try {
ImmutableOpenMap.Builder<String, MappingMetadata> builder = ImmutableOpenMap.builder();
builder.put(TestsConstants.TEST_INDEX_BANK, IndexMetadata.fromXContent(createParser(mappings)).mapping());
when(mockMetaData.findMappings(any(), any())).thenReturn(builder.build());
when(mockMetaData.findMappings(any(), any())).thenReturn(
Map.of(TestsConstants.TEST_INDEX_BANK, IndexMetadata.fromXContent(
createParser(mappings)).mapping()));
}
catch (IOException e) {
throw new IllegalStateException(e);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -14,15 +14,14 @@
import static org.opensearch.sql.legacy.util.CheckScriptContents.mockIndexNameExpressionResolver;
import static org.opensearch.sql.legacy.util.CheckScriptContents.mockPluginSettings;

import com.google.common.collect.ImmutableMap;
import java.io.IOException;
import java.util.Map;
import java.util.stream.Collectors;
import org.opensearch.cluster.ClusterState;
import org.opensearch.cluster.metadata.IndexMetadata;
import org.opensearch.cluster.metadata.MappingMetadata;
import org.opensearch.cluster.metadata.Metadata;
import org.opensearch.cluster.service.ClusterService;
import org.opensearch.common.collect.ImmutableOpenMap;
import org.opensearch.sql.legacy.esdomain.LocalClusterState;

/**
Expand Down Expand Up @@ -139,57 +138,54 @@ public class MultipleIndexClusterUtils {
"}";

public static void mockMultipleIndexEnv() {
mockLocalClusterState(new ImmutableMap.Builder<String, ImmutableOpenMap<String, MappingMetadata>>()
.put(INDEX_ACCOUNT_1, buildIndexMapping(INDEX_ACCOUNT_1, INDEX_ACCOUNT_1_MAPPING))
.put(INDEX_ACCOUNT_2, buildIndexMapping(INDEX_ACCOUNT_2, INDEX_ACCOUNT_2_MAPPING))
.put(INDEX_ACCOUNT_ALL, buildIndexMapping(new ImmutableMap.Builder<String, String>()
.put(INDEX_ACCOUNT_1, INDEX_ACCOUNT_1_MAPPING)
.put(INDEX_ACCOUNT_2, INDEX_ACCOUNT_2_MAPPING)
.build()))
.build());
mockLocalClusterState(
Map.of(INDEX_ACCOUNT_1, buildIndexMapping(INDEX_ACCOUNT_1, INDEX_ACCOUNT_1_MAPPING),
INDEX_ACCOUNT_2, buildIndexMapping(INDEX_ACCOUNT_2, INDEX_ACCOUNT_2_MAPPING),
INDEX_ACCOUNT_ALL, buildIndexMapping(Map.of(INDEX_ACCOUNT_1, INDEX_ACCOUNT_1_MAPPING,
INDEX_ACCOUNT_2, INDEX_ACCOUNT_2_MAPPING))));
}

public static void mockLocalClusterState(Map<String, ImmutableOpenMap<String,MappingMetadata>> indexMapping) {
public static void mockLocalClusterState(Map<String, Map<String,MappingMetadata>> indexMapping) {
LocalClusterState.state().setClusterService(mockClusterService(indexMapping));
LocalClusterState.state().setResolver(mockIndexNameExpressionResolver());
LocalClusterState.state().setPluginSettings(mockPluginSettings());
}


public static ClusterService mockClusterService(Map<String, ImmutableOpenMap<String,MappingMetadata>> indexMapping) {
public static ClusterService mockClusterService(Map<String, Map<String,MappingMetadata>>
indexMapping) {
ClusterService mockService = mock(ClusterService.class);
ClusterState mockState = mock(ClusterState.class);
Metadata mockMetaData = mock(Metadata.class);

when(mockService.state()).thenReturn(mockState);
when(mockState.metadata()).thenReturn(mockMetaData);
try {
for (Map.Entry<String, ImmutableOpenMap<String, MappingMetadata>> entry : indexMapping.entrySet()) {
when(mockMetaData.findMappings(eq(new String[]{entry.getKey()}), any())).thenReturn(entry.getValue());
for (var entry : indexMapping.entrySet()) {
when(mockMetaData.findMappings(eq(new String[]{entry.getKey()}), any()))
.thenReturn(entry.getValue());
}
} catch (IOException e) {
throw new IllegalStateException(e);
}
return mockService;
}

private static ImmutableOpenMap<String, MappingMetadata> buildIndexMapping(Map<String, String> indexMapping) {
try {
ImmutableOpenMap.Builder<String, MappingMetadata> builder = ImmutableOpenMap.builder();
for (Map.Entry<String, String> entry : indexMapping.entrySet()) {
builder.put(entry.getKey(), IndexMetadata.fromXContent(createParser(entry.getValue())).mapping());
private static Map<String, MappingMetadata> buildIndexMapping(Map<String, String> indexMapping) {
return indexMapping.entrySet().stream().collect(Collectors.toUnmodifiableMap(
Map.Entry::getKey, e -> {
try {
return IndexMetadata.fromXContent(createParser(e.getValue())).mapping();
} catch (IOException ex) {
throw new IllegalStateException(ex);
}
return builder.build();
} catch (IOException e) {
throw new IllegalStateException(e);
}
}));

}

private static ImmutableOpenMap<String, MappingMetadata> buildIndexMapping(String index, String mapping) {
private static Map<String, MappingMetadata> buildIndexMapping(String index, String mapping) {
try {
ImmutableOpenMap.Builder<String, MappingMetadata> builder = ImmutableOpenMap.builder();
builder.put(index, IndexMetadata.fromXContent(createParser(mapping)).mapping());
return builder.build();
return Map.of(index, IndexMetadata.fromXContent(createParser(mapping)).mapping());
} catch (IOException e) {
throw new IllegalStateException(e);
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,6 @@
import com.carrotsearch.hppc.cursors.ObjectObjectCursor;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Streams;
import java.util.Arrays;
import java.util.Collection;
import java.util.List;
Expand Down Expand Up @@ -89,9 +88,9 @@ public Map<String, IndexMapping> getIndexMappings(String... indexExpression) {
.prepareGetMappings(indexExpression)
.setLocal(true)
.get();
return Streams.stream(mappingsResponse.mappings().iterator())
.collect(Collectors.toMap(cursor -> cursor.key,
cursor -> new IndexMapping(cursor.value)));
return mappingsResponse.mappings().entrySet().stream().collect(Collectors.toUnmodifiableMap(
Map.Entry::getKey,
cursor -> new IndexMapping(cursor.getValue())));
} catch (IndexNotFoundException e) {
// Re-throw directly to be treated as client error finally
throw e;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -437,16 +437,12 @@ public void mockNodeClientIndicesMappings(String indexName, String mappings) {
.setLocal(anyBoolean())
.get()).thenReturn(mockResponse);
try {
ImmutableOpenMap<String, MappingMetadata> metadata;
Map<String, MappingMetadata> metadata;
if (mappings.isEmpty()) {
when(emptyMapping.getSourceAsMap()).thenReturn(ImmutableMap.of());
metadata =
new ImmutableOpenMap.Builder<String, MappingMetadata>()
.fPut(indexName, emptyMapping)
.build();
when(emptyMapping.getSourceAsMap()).thenReturn(Map.of());
metadata = Map.of(indexName, emptyMapping);
} else {
metadata = new ImmutableOpenMap.Builder<String, MappingMetadata>().fPut(indexName,
IndexMetadata.fromXContent(createParser(mappings)).mapping()).build();
metadata = Map.of(indexName, IndexMetadata.fromXContent(createParser(mappings)).mapping());
}
when(mockResponse.mappings()).thenReturn(metadata);
} catch (IOException e) {
Expand Down

0 comments on commit 8558e90

Please sign in to comment.