adding context.restore when missed and fixing model group parsing

Signed-off-by: Amit Galitzky <[email protected]>
amitgalitz · May 1, 2024 · 3b94d22 · 3b94d22
1 parent ff072de
commit 3b94d22
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 2 deletions.
diff --git a/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java b/src/main/java/org/opensearch/flowframework/transport/CreateWorkflowTransportAction.java
@@ -311,6 +311,7 @@ void checkMaxWorkflows(TimeValue requestTimeOut, Integer maxWorkflow, ActionList
             try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
                 logger.info("Querying existing workflows to count the max");
                 client.search(searchRequest, ActionListener.wrap(searchResponse -> {
+                    context.restore();
                     internalListener.onResponse(searchResponse.getHits().getTotalHits().value < maxWorkflow);
                 }, exception -> {
                     String errorMessage = "Unable to fetch the workflows";

diff --git a/src/main/java/org/opensearch/flowframework/util/EncryptorUtils.java b/src/main/java/org/opensearch/flowframework/util/EncryptorUtils.java
@@ -250,7 +250,6 @@ public void initializeMasterKey(ActionListener<Boolean> listener) {
 
             GetRequest getRequest = new GetRequest(CONFIG_INDEX).id(MASTER_KEY);
             client.get(getRequest, ActionListener.wrap(getResponse -> {
-
                 if (!getResponse.isExists()) {
 
                     // Generate new key and index
@@ -260,6 +259,7 @@ public void initializeMasterKey(ActionListener<Boolean> listener) {
                         .setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE);
 
                     client.index(masterKeyIndexRequest, ActionListener.wrap(indexResponse -> {
+                        context.restore();
                         // Set generated key to master
                         logger.info("Config has been initialized successfully");
                         this.masterKey = generatedKey;
@@ -270,6 +270,7 @@ public void initializeMasterKey(ActionListener<Boolean> listener) {
                     }));
 
                 } else {
+                    context.restore();
                     // Set existing key to master
                     logger.info("Config has already been initialized");
                     this.masterKey = (String) getResponse.getSourceAsMap().get(MASTER_KEY);
@@ -300,6 +301,7 @@ void initializeMasterKeyIfAbsent() {
             try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
                 GetRequest getRequest = new GetRequest(CONFIG_INDEX).id(MASTER_KEY);
                 client.get(getRequest, ActionListener.wrap(response -> {
+                    context.restore();
                     if (response.isExists()) {
                         this.masterKey = (String) response.getSourceAsMap().get(MASTER_KEY);
                     } else {

diff --git a/src/main/java/org/opensearch/flowframework/workflow/RegisterModelGroupStep.java b/src/main/java/org/opensearch/flowframework/workflow/RegisterModelGroupStep.java
@@ -28,6 +28,7 @@
 
 import java.util.Collections;
 import java.util.List;
+import java.util.Locale;
 import java.util.Map;
 import java.util.Set;
 
@@ -142,7 +143,7 @@ public void onFailure(Exception ex) {
             String modelGroupName = (String) inputs.get(NAME_FIELD);
             String description = (String) inputs.get(DESCRIPTION_FIELD);
             List<String> backendRoles = getBackendRoles(inputs);
-            AccessMode modelAccessMode = (AccessMode) inputs.get(MODEL_ACCESS_MODE);
+            AccessMode modelAccessMode = AccessMode.from(((String) inputs.get(MODEL_ACCESS_MODE)).toLowerCase(Locale.ROOT));
             Boolean isAddAllBackendRoles = inputs.containsKey(ADD_ALL_BACKEND_ROLES)
                 ? Booleans.parseBoolean(inputs.get(ADD_ALL_BACKEND_ROLES).toString())
                 : null;