Skip to content

Commit 9846e6e

Browse files
dhrubo-osdhrubo-os
authored
adding multi-tenancy to config api and master key related changes (#3439)
* adding multi-tenancy to config api and master key related changes Signed-off-by: Dhrubo Saha <dhrubo@amazon.com> * adding more unit tests Signed-off-by: Dhrubo Saha <dhrubo@amazon.com> --------- Signed-off-by: Dhrubo Saha <dhrubo@amazon.com>
1 parent 25e6c9a commit 9846e6e

File tree

54 files changed

+814
-271
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

54 files changed

+814
-271
lines changed

‎common/build.gradle

+2
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,8 @@ dependencies {
3939
compileOnly("com.fasterxml.jackson.core:jackson-annotations:${versions.jackson}")
4040
compileOnly("com.fasterxml.jackson.core:jackson-databind:${versions.jackson_databind}")
4141
compileOnly group: 'com.networknt' , name: 'json-schema-validator', version: '1.4.0'
42+
// Multi-tenant SDK Client
43+
compileOnly "org.opensearch:opensearch-remote-metadata-sdk:${opensearch_build}"
4244
}
4345

4446
lombok {

‎common/src/main/java/org/opensearch/ml/common/MLConfig.java

+18-3
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@
66
package org.opensearch.ml.common;
77

88
import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
9+
import static org.opensearch.ml.common.CommonValue.TENANT_ID_FIELD;
10+
import static org.opensearch.ml.common.CommonValue.VERSION_2_19_0;
911

1012
import java.io.IOException;
1113
import java.time.Instant;
@@ -57,6 +59,7 @@ public class MLConfig implements ToXContentObject, Writeable {
5759
private final Instant createTime;
5860
private Instant lastUpdateTime;
5961
private Instant lastUpdatedTime;
62+
private final String tenantId;
6063

6164
@Builder(toBuilder = true)
6265
public MLConfig(
@@ -66,7 +69,8 @@ public MLConfig(
6669
Configuration mlConfiguration,
6770
Instant createTime,
6871
Instant lastUpdateTime,
69-
Instant lastUpdatedTime
72+
Instant lastUpdatedTime,
73+
String tenantId
7074
) {
7175
this.type = type;
7276
this.configType = configType;
@@ -75,6 +79,7 @@ public MLConfig(
7579
this.createTime = createTime;
7680
this.lastUpdateTime = lastUpdateTime;
7781
this.lastUpdatedTime = lastUpdatedTime;
82+
this.tenantId = tenantId;
7883
}
7984

8085
public MLConfig(StreamInput input) throws IOException {
@@ -92,6 +97,7 @@ public MLConfig(StreamInput input) throws IOException {
9297
}
9398
lastUpdatedTime = input.readOptionalInstant();
9499
}
100+
this.tenantId = streamInputVersion.onOrAfter(VERSION_2_19_0) ? input.readOptionalString() : null;
95101
}
96102

97103
@Override
@@ -116,6 +122,9 @@ public void writeTo(StreamOutput out) throws IOException {
116122
}
117123
out.writeOptionalInstant(lastUpdatedTime);
118124
}
125+
if (streamOutputVersion.onOrAfter(VERSION_2_19_0)) {
126+
out.writeOptionalString(tenantId);
127+
}
119128
}
120129

121130
@Override
@@ -133,12 +142,14 @@ public XContentBuilder toXContent(XContentBuilder xContentBuilder, Params params
133142
if (lastUpdateTime != null || lastUpdatedTime != null) {
134143
builder.field(LAST_UPDATE_TIME_FIELD, lastUpdatedTime == null ? lastUpdateTime.toEpochMilli() : lastUpdatedTime.toEpochMilli());
135144
}
145+
if (tenantId != null) {
146+
builder.field(TENANT_ID_FIELD, tenantId);
147+
}
136148
return builder.endObject();
137149
}
138150

139151
public static MLConfig fromStream(StreamInput in) throws IOException {
140-
MLConfig mlConfig = new MLConfig(in);
141-
return mlConfig;
152+
return new MLConfig(in);
142153
}
143154

144155
public static MLConfig parse(XContentParser parser) throws IOException {
@@ -149,6 +160,7 @@ public static MLConfig parse(XContentParser parser) throws IOException {
149160
Instant createTime = null;
150161
Instant lastUpdateTime = null;
151162
Instant lastUpdatedTime = null;
163+
String tenantId = null;
152164

153165
ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.currentToken(), parser);
154166
while (parser.nextToken() != XContentParser.Token.END_OBJECT) {
@@ -177,6 +189,8 @@ public static MLConfig parse(XContentParser parser) throws IOException {
177189
case LAST_UPDATED_TIME_FIELD:
178190
lastUpdatedTime = Instant.ofEpochMilli(parser.longValue());
179191
break;
192+
case TENANT_ID_FIELD:
193+
tenantId = parser.textOrNull();
180194
default:
181195
parser.skipChildren();
182196
break;
@@ -191,6 +205,7 @@ public static MLConfig parse(XContentParser parser) throws IOException {
191205
.createTime(createTime)
192206
.lastUpdateTime(lastUpdateTime)
193207
.lastUpdatedTime(lastUpdatedTime)
208+
.tenantId(tenantId)
194209
.build();
195210
}
196211
}

‎common/src/main/java/org/opensearch/ml/common/connector/Connector.java

+4-4
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616
import java.util.List;
1717
import java.util.Map;
1818
import java.util.Optional;
19-
import java.util.function.Function;
19+
import java.util.function.BiFunction;
2020
import java.util.regex.Matcher;
2121
import java.util.regex.Pattern;
2222

@@ -79,9 +79,9 @@ public interface Connector extends ToXContentObject, Writeable {
7979

8080
<T> T createPayload(String action, Map<String, String> parameters);
8181

82-
void decrypt(String action, Function<String, String> function);
82+
void decrypt(String action, BiFunction<String, String, String> function, String tenantId);
8383

84-
void encrypt(Function<String, String> function);
84+
void encrypt(BiFunction<String, String, String> function, String tenantId);
8585

8686
Connector cloneConnector();
8787

@@ -91,7 +91,7 @@ public interface Connector extends ToXContentObject, Writeable {
9191

9292
void writeTo(StreamOutput out) throws IOException;
9393

94-
void update(MLCreateConnectorInput updateContent, Function<String, String> function);
94+
void update(MLCreateConnectorInput updateContent, BiFunction<String, String, String> function);
9595

9696
<T> void parseResponse(T orElse, List<ModelTensor> modelTensors, boolean b) throws IOException;
9797

‎common/src/main/java/org/opensearch/ml/common/connector/HttpConnector.java

+7-7
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@
2121
import java.util.List;
2222
import java.util.Map;
2323
import java.util.Optional;
24-
import java.util.function.Function;
24+
import java.util.function.BiFunction;
2525
import java.util.regex.Matcher;
2626
import java.util.regex.Pattern;
2727

@@ -300,7 +300,7 @@ public void writeTo(StreamOutput out) throws IOException {
300300
}
301301

302302
@Override
303-
public void update(MLCreateConnectorInput updateContent, Function<String, String> function) {
303+
public void update(MLCreateConnectorInput updateContent, BiFunction<String, String, String> function) {
304304
if (updateContent.getName() != null) {
305305
this.name = updateContent.getName();
306306
}
@@ -318,7 +318,7 @@ public void update(MLCreateConnectorInput updateContent, Function<String, String
318318
}
319319
if (updateContent.getCredential() != null && !updateContent.getCredential().isEmpty()) {
320320
this.credential = updateContent.getCredential();
321-
encrypt(function);
321+
encrypt(function, this.tenantId);
322322
}
323323
if (updateContent.getActions() != null) {
324324
this.actions = updateContent.getActions();
@@ -377,10 +377,10 @@ private List<String> findStringParametersWithNullDefaultValue(String input) {
377377
}
378378

379379
@Override
380-
public void decrypt(String action, Function<String, String> function) {
380+
public void decrypt(String action, BiFunction<String, String, String> function, String tenantId) {
381381
Map<String, String> decrypted = new HashMap<>();
382382
for (String key : credential.keySet()) {
383-
decrypted.put(key, function.apply(credential.get(key)));
383+
decrypted.put(key, function.apply(credential.get(key), tenantId));
384384
}
385385
this.decryptedCredential = decrypted;
386386
Optional<ConnectorAction> connectorAction = findAction(action);
@@ -400,9 +400,9 @@ public Connector cloneConnector() {
400400
}
401401

402402
@Override
403-
public void encrypt(Function<String, String> function) {
403+
public void encrypt(BiFunction<String, String, String> function, String tenantId) {
404404
for (String key : credential.keySet()) {
405-
String encrypted = function.apply(credential.get(key));
405+
String encrypted = function.apply(credential.get(key), tenantId);
406406
credential.put(key, encrypted);
407407
}
408408
}

‎common/src/main/java/org/opensearch/ml/common/transport/config/MLConfigGetRequest.java

+11-1
Original file line numberDiff line numberDiff line change
@@ -6,12 +6,14 @@
66
package org.opensearch.ml.common.transport.config;
77

88
import static org.opensearch.action.ValidateActions.addValidationError;
9+
import static org.opensearch.ml.common.CommonValue.VERSION_2_19_0;
910

1011
import java.io.ByteArrayInputStream;
1112
import java.io.ByteArrayOutputStream;
1213
import java.io.IOException;
1314
import java.io.UncheckedIOException;
1415

16+
import org.opensearch.Version;
1517
import org.opensearch.action.ActionRequest;
1618
import org.opensearch.action.ActionRequestValidationException;
1719
import org.opensearch.core.common.io.stream.InputStreamStreamInput;
@@ -26,21 +28,29 @@
2628
public class MLConfigGetRequest extends ActionRequest {
2729

2830
String configId;
31+
String tenantId;
2932

3033
@Builder
31-
public MLConfigGetRequest(String configId) {
34+
public MLConfigGetRequest(String configId, String tenantId) {
3235
this.configId = configId;
36+
this.tenantId = tenantId;
3337
}
3438

3539
public MLConfigGetRequest(StreamInput in) throws IOException {
3640
super(in);
41+
Version streamInputVersion = in.getVersion();
3742
this.configId = in.readString();
43+
this.tenantId = streamInputVersion.onOrAfter(VERSION_2_19_0) ? in.readOptionalString() : null;
3844
}
3945

4046
@Override
4147
public void writeTo(StreamOutput out) throws IOException {
4248
super.writeTo(out);
49+
Version streamOutputVersion = out.getVersion();
4350
out.writeString(this.configId);
51+
if (streamOutputVersion.onOrAfter(VERSION_2_19_0)) {
52+
out.writeOptionalString(tenantId);
53+
}
4454
}
4555

4656
@Override

‎common/src/main/java/org/opensearch/ml/common/utils/StringUtils.java

+20
Original file line numberDiff line numberDiff line change
@@ -8,9 +8,12 @@
88
import java.nio.ByteBuffer;
99
import java.nio.charset.StandardCharsets;
1010
import java.security.AccessController;
11+
import java.security.MessageDigest;
12+
import java.security.NoSuchAlgorithmException;
1113
import java.security.PrivilegedActionException;
1214
import java.security.PrivilegedExceptionAction;
1315
import java.util.ArrayList;
16+
import java.util.Base64;
1417
import java.util.HashMap;
1518
import java.util.HashSet;
1619
import java.util.List;
@@ -477,4 +480,21 @@ public static void validateSchema(String schemaString, String instanceString) {
477480
throw new OpenSearchParseException("Schema validation failed: " + e.getMessage(), e);
478481
}
479482
}
483+
484+
public static String hashString(String input) {
485+
try {
486+
// Create a MessageDigest instance for SHA-256
487+
MessageDigest digest = MessageDigest.getInstance("SHA-256");
488+
489+
// Perform the hashing and get the byte array
490+
byte[] hashBytes = digest.digest(input.getBytes());
491+
492+
// Convert the byte array to a Base64 encoded string
493+
return Base64.getUrlEncoder().encodeToString(hashBytes);
494+
495+
} catch (NoSuchAlgorithmException e) {
496+
throw new RuntimeException("Error: Unable to compute hash", e);
497+
}
498+
}
499+
480500
}

0 commit comments

Comments
 (0)