[Backport 2.x] Fixing model group parsing and restoring context (#696)

Fixing model group parsing and restoring context (#695) * adding context.restore when missed and fixing model group parsing * adding change log, fix casting * changed to contains key --------- (cherry picked from commit 28ea3dd) Signed-off-by: Amit Galitzky <[email protected]> Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
opensearch-project · May 1, 2024 · c356b6e · c356b6e
1 parent 99c5ca1
commit c356b6e
Show file tree

Hide file tree

Showing 5 changed files with 12 additions and 2 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -25,6 +25,7 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.1.0/)
 - Hide user and credential field from search response ([#680](https://github.com/opensearch-project/flow-framework/pull/680))
 - Throw the correct error message in status API for WorkflowSteps ([#676](https://github.com/opensearch-project/flow-framework/pull/676))
 - Delete workflow state when template is deleted and no resources exist ([#689](https://github.com/opensearch-project/flow-framework/pull/689))
+- Fixing model group parsing and restoring context ([#695] (https://github.com/opensearch-project/flow-framework/pull/695))
 
 ### Infrastructure
 - Switch macos runner to macos-13 from macos-latest since macos-latest is now arm64 ([#686](https://github.com/opensearch-project/flow-framework/pull/686))

diff --git a/release-notes/opensearch-flow-framework.release-notes-2.14.0.0.md b/release-notes/opensearch-flow-framework.release-notes-2.14.0.0.md
@@ -13,6 +13,8 @@ Compatible with OpenSearch 2.14.0
 - Hide user and credential field from search response ([#680](https://github.com/opensearch-project/flow-framework/pull/680))
 - Throw the correct error message in status API for WorkflowSteps ([#676](https://github.com/opensearch-project/flow-framework/pull/676))
 - Delete workflow state when template is deleted and no resources exist ([#689](https://github.com/opensearch-project/flow-framework/pull/689))
+- Fixing model group parsing and restoring context ([#695] (https://github.com/opensearch-project/flow-framework/pull/695))
+
 
 ### Infrastructure
 - Switch macOS runner to macos-13 from macos-latest since macos-latest is now arm64 ([#686](https://github.com/opensearch-project/flow-framework/pull/686))

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,10 @@ 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 = null;
+            if (inputs.containsKey(MODEL_ACCESS_MODE)) {
+                modelAccessMode = AccessMode.from((inputs.get(MODEL_ACCESS_MODE)).toString().toLowerCase(Locale.ROOT));
+            }
             Boolean isAddAllBackendRoles = inputs.containsKey(ADD_ALL_BACKEND_ROLES)
                 ? Booleans.parseBoolean(inputs.get(ADD_ALL_BACKEND_ROLES).toString())
                 : null;