From 582d22af692e35c183572b5a78c433d7e45abc35 Mon Sep 17 00:00:00 2001 From: Jonathan Shook Date: Tue, 6 Feb 2024 23:56:04 -0600 Subject: [PATCH] part1 of aws os --- adapter-aws-opensearch/pom.xml | 6 ++ .../adapter/opensearch/OpenSearchAdapter.java | 13 +++++ .../adapter/opensearch/OpenSearchSpace.java | 58 ++++++++++++++++--- .../dispensers/CreateIndexOpDispenser.java | 29 ++++++---- .../adapter/opensearch/ops/CreateIndexOp.java | 4 +- .../main/resources/activities/osvectors.yaml | 1 + .../uniform/actions/StandardAction.java | 12 ++-- .../api/templating/TriFunction.java.new | 23 ++++++++ 8 files changed, 122 insertions(+), 24 deletions(-) create mode 100644 virtdata-api/src/main/java/io/nosqlbench/adapters/api/templating/TriFunction.java.new diff --git a/adapter-aws-opensearch/pom.xml b/adapter-aws-opensearch/pom.xml index fa297eff8..99f12d657 100644 --- a/adapter-aws-opensearch/pom.xml +++ b/adapter-aws-opensearch/pom.xml @@ -85,11 +85,17 @@ aws-crt-client 2.20.109 + software.amazon.awssdk regions 2.23.18 + + software.amazon.awssdk + auth + 2.23.19 + diff --git a/adapter-aws-opensearch/src/main/java/io/nosqlbench/adapter/opensearch/OpenSearchAdapter.java b/adapter-aws-opensearch/src/main/java/io/nosqlbench/adapter/opensearch/OpenSearchAdapter.java index 2309563da..bbdfb545f 100644 --- a/adapter-aws-opensearch/src/main/java/io/nosqlbench/adapter/opensearch/OpenSearchAdapter.java +++ b/adapter-aws-opensearch/src/main/java/io/nosqlbench/adapter/opensearch/OpenSearchAdapter.java @@ -21,16 +21,29 @@ import io.nosqlbench.adapters.api.activityimpl.uniform.BaseDriverAdapter; import io.nosqlbench.adapters.api.activityimpl.uniform.DriverSpaceCache; import io.nosqlbench.adapters.api.activityimpl.uniform.flowtypes.Op; import io.nosqlbench.nb.api.components.core.NBComponent; +import io.nosqlbench.nb.api.config.standard.NBConfigModel; import io.nosqlbench.nb.api.config.standard.NBConfiguration; import io.nosqlbench.nb.api.labels.NBLabels; +import java.util.function.Function; + public class OpenSearchAdapter extends BaseDriverAdapter { public OpenSearchAdapter(NBComponent parentComponent, NBLabels labels) { super(parentComponent, labels); } + @Override + public Function getSpaceInitializer(NBConfiguration cfg) { + return (String spaceName) -> new OpenSearchSpace(cfg); + } + @Override public OpMapper getOpMapper() { return new OpenSearchOpMapper(this); } + + @Override + public NBConfigModel getConfigModel() { + return super.getConfigModel().add(OpenSearchSpace.getConfigModel()); + } } diff --git a/adapter-aws-opensearch/src/main/java/io/nosqlbench/adapter/opensearch/OpenSearchSpace.java b/adapter-aws-opensearch/src/main/java/io/nosqlbench/adapter/opensearch/OpenSearchSpace.java index 1d830d3d2..5ce87549e 100644 --- a/adapter-aws-opensearch/src/main/java/io/nosqlbench/adapter/opensearch/OpenSearchSpace.java +++ b/adapter-aws-opensearch/src/main/java/io/nosqlbench/adapter/opensearch/OpenSearchSpace.java @@ -17,18 +17,27 @@ package io.nosqlbench.adapter.opensearch; +import com.amazonaws.auth.BasicAWSCredentials; +import com.amazonaws.auth.profile.internal.BasicProfile; +import com.amazonaws.auth.profile.internal.ProfileKeyConstants; +import com.amazonaws.auth.profile.internal.ProfileStaticCredentialsProvider; import io.nosqlbench.nb.api.config.standard.ConfigModel; import io.nosqlbench.nb.api.config.standard.NBConfigModel; import io.nosqlbench.nb.api.config.standard.NBConfiguration; import io.nosqlbench.nb.api.config.standard.Param; import org.opensearch.client.opensearch.OpenSearchClient; +import org.opensearch.client.opensearch.core.InfoResponse; import org.opensearch.client.transport.aws.AwsSdk2Transport; import org.opensearch.client.transport.aws.AwsSdk2TransportOptions; +import software.amazon.awssdk.auth.credentials.*; import software.amazon.awssdk.http.async.SdkAsyncHttpClient; import software.amazon.awssdk.http.crt.AwsCrtAsyncHttpClient; import software.amazon.awssdk.regions.Region; -public class OpenSearchSpace { +import java.io.IOException; +import java.util.Map; + +public class OpenSearchSpace implements AutoCloseable { private final NBConfiguration cfg; protected OpenSearchClient client; @@ -38,7 +47,7 @@ public class OpenSearchSpace { } public synchronized OpenSearchClient getClient() { - if (client!=null) { + if (client == null) { client = createClient(); } return client; @@ -49,12 +58,39 @@ public class OpenSearchSpace { Region selectedRegion = Region.of(region); String host = cfg.get("host"); - // https://ecj69wlzmday6hvr2586.eu-central-1.aoss.amazonaws.com - SdkAsyncHttpClient client1 = AwsCrtAsyncHttpClient.builder().build(); - AwsSdk2TransportOptions transportOptions = AwsSdk2TransportOptions.builder().build(); - AwsSdk2Transport awsSdk2Transport = new AwsSdk2Transport(client1,host,selectedRegion,transportOptions); + ProfileCredentialsProvider creds = ProfileCredentialsProvider.builder() + .profileName("686157956141_ENG-TESTENG_AWS-ENG-PERF") + .build(); + + + SdkAsyncHttpClient httpClient = + AwsCrtAsyncHttpClient.builder() + .build(); + + AwsSdk2TransportOptions transportOptions = + AwsSdk2TransportOptions.builder() + .setCredentials(creds) + .build(); + + AwsSdk2Transport awsSdk2Transport = + new AwsSdk2Transport( + httpClient, + host, + "es", + selectedRegion, + transportOptions + ); + OpenSearchClient client = new OpenSearchClient(awsSdk2Transport); + + try { + InfoResponse info = client.info(); + System.out.println(info.version().distribution() + ": " + info.version().number()); + } catch (IOException e) { + throw new RuntimeException(e); + } + return client; // // Create a new domain, update its configuration, and delete it. @@ -69,10 +105,17 @@ public class OpenSearchSpace { public static NBConfigModel getConfigModel() { return ConfigModel.of(OpenSearchSpace.class) .add(Param.required("region", String.class).setDescription("The region to connect to")) - .add(Param.required("host",String.class).setDescription("The Open Search API endpoint")) + .add(Param.required("host", String.class).setDescription("The Open Search API endpoint host")) .asReadOnly(); } + @Override + public void close() throws Exception { + if (client != null) { + client.shutdown(); + } + } + // /** // * Waits for the domain to finish processing changes. New domains typically take 15-30 minutes @@ -109,4 +152,5 @@ public class OpenSearchSpace { // System.out.println("Domain description: "+describeResponse.toString()); // } + } diff --git a/adapter-aws-opensearch/src/main/java/io/nosqlbench/adapter/opensearch/dispensers/CreateIndexOpDispenser.java b/adapter-aws-opensearch/src/main/java/io/nosqlbench/adapter/opensearch/dispensers/CreateIndexOpDispenser.java index 26c06582d..d337bab32 100644 --- a/adapter-aws-opensearch/src/main/java/io/nosqlbench/adapter/opensearch/dispensers/CreateIndexOpDispenser.java +++ b/adapter-aws-opensearch/src/main/java/io/nosqlbench/adapter/opensearch/dispensers/CreateIndexOpDispenser.java @@ -19,6 +19,7 @@ package io.nosqlbench.adapter.opensearch.dispensers; import io.nosqlbench.adapter.opensearch.OpenSearchAdapter; import io.nosqlbench.adapter.opensearch.ops.CreateIndexOp; import io.nosqlbench.adapters.api.templating.ParsedOp; +import org.opensearch.client.json.JsonData; import org.opensearch.client.opensearch.OpenSearchClient; import org.opensearch.client.opensearch._types.mapping.*; import org.opensearch.client.opensearch.indices.CreateIndexRequest; @@ -58,22 +59,30 @@ public class CreateIndexOpDispenser extends BaseOpenSearchOpDispenser { @Override public LongFunction createOpFunc(LongFunction clientF, ParsedOp op) { CreateIndexRequest.Builder eb = new CreateIndexRequest.Builder(); - LongFunction bfunc = l -> new CreateIndexRequest.Builder(); + LongFunction bfunc = l -> new CreateIndexRequest.Builder().index("testindex1"); bfunc = op.enhanceFunc(bfunc, "mappings", Map.class, this::resolveTypeMapping); LongFunction finalBfunc = bfunc; return (long l) -> new CreateIndexOp(clientF.apply(l), finalBfunc.apply(l).build()); } - private CreateIndexRequest.Builder resolveTypeMapping(CreateIndexRequest.Builder eb, Map mappings) { - TypeMapping.Builder builder = new TypeMapping.Builder() - .properties(Map.of( - "p1", new Property.Builder().knnVector(new KnnVectorProperty.Builder() - .dimension(23) - .method(b -> b.spaceType("sdf")) - .build() - ).build() - )) + // https://opensearch.org/docs/latest/search-plugins/knn/knn-index/ + private CreateIndexRequest.Builder resolveTypeMapping(CreateIndexRequest.Builder eb, Map mappings) { + TypeMapping.Builder builder = new TypeMapping.Builder().properties( + Map.of( + "p1", + new Property.Builder().knnVector(new KnnVectorProperty.Builder() + .dimension(23) + .method( + new KnnVectorMethod.Builder() + .name("hnsw") + .engine("faiss") + .spaceType("l2") + .parameters(Map.of("ef_construction", JsonData.of(256),"m",JsonData.of(8))) + .build() + ).build() + ).build() + )) .indexField(new IndexField.Builder() .enabled(true).build()) .fieldNames(new FieldNamesField.Builder() diff --git a/adapter-aws-opensearch/src/main/java/io/nosqlbench/adapter/opensearch/ops/CreateIndexOp.java b/adapter-aws-opensearch/src/main/java/io/nosqlbench/adapter/opensearch/ops/CreateIndexOp.java index a30125594..f20cce7c1 100644 --- a/adapter-aws-opensearch/src/main/java/io/nosqlbench/adapter/opensearch/ops/CreateIndexOp.java +++ b/adapter-aws-opensearch/src/main/java/io/nosqlbench/adapter/opensearch/ops/CreateIndexOp.java @@ -18,6 +18,7 @@ package io.nosqlbench.adapter.opensearch.ops; import org.opensearch.client.opensearch.OpenSearchClient; import org.opensearch.client.opensearch.indices.CreateIndexRequest; +import org.opensearch.client.opensearch.indices.CreateIndexResponse; import java.io.IOException; @@ -32,7 +33,8 @@ public class CreateIndexOp extends BaseOpenSearchOp { @Override public Object apply(long value) { try { - return client.indices().create(rq); + CreateIndexResponse response = client.indices().create(rq); + return response; } catch (IOException e) { throw new RuntimeException(e); } diff --git a/adapter-aws-opensearch/src/main/resources/activities/osvectors.yaml b/adapter-aws-opensearch/src/main/resources/activities/osvectors.yaml index 9a60b6975..9add87b00 100644 --- a/adapter-aws-opensearch/src/main/resources/activities/osvectors.yaml +++ b/adapter-aws-opensearch/src/main/resources/activities/osvectors.yaml @@ -13,3 +13,4 @@ blocks: index: create_index: TEMPLATE(indexname,vectors_index) mappings: + m1: v1 diff --git a/engine-core/src/main/java/io/nosqlbench/engine/api/activityimpl/uniform/actions/StandardAction.java b/engine-core/src/main/java/io/nosqlbench/engine/api/activityimpl/uniform/actions/StandardAction.java index f5a09523f..4e4c15504 100644 --- a/engine-core/src/main/java/io/nosqlbench/engine/api/activityimpl/uniform/actions/StandardAction.java +++ b/engine-core/src/main/java/io/nosqlbench/engine/api/activityimpl/uniform/actions/StandardAction.java @@ -95,12 +95,12 @@ public class StandardAction, R extends Op> impl dispenser.onStart(cycle); try (Timer.Context ct = executeTimer.time()) { - if (op instanceof RunnableOp) { - ((RunnableOp) op).run(); - } else if (op instanceof CycleOp) { - result = ((CycleOp) op).apply(cycle); - } else if (op instanceof ChainingOp) { - result = ((ChainingOp) op).apply(result); + if (op instanceof RunnableOp runnableOp) { + runnableOp.run(); + } else if (op instanceof CycleOp cycleOp) { + result = cycleOp.apply(cycle); + } else if (op instanceof ChainingOp chainingOp) { + result = chainingOp.apply(result); } else { throw new RuntimeException("The op implementation did not implement any active logic. Implement " + "one of [RunnableOp, CycleOp, or ChainingOp]"); diff --git a/virtdata-api/src/main/java/io/nosqlbench/adapters/api/templating/TriFunction.java.new b/virtdata-api/src/main/java/io/nosqlbench/adapters/api/templating/TriFunction.java.new new file mode 100644 index 000000000..afaf8a47c --- /dev/null +++ b/virtdata-api/src/main/java/io/nosqlbench/adapters/api/templating/TriFunction.java.new @@ -0,0 +1,23 @@ +package io.nosqlbench.adapters.api.templating; + +/* + * Copyright (c) 2022 nosqlbench + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + + +public interface TriFunction { + public A apply(A a,B b,C c); +}