[Backport 2.14] Fixing model group parsing and restoring context

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))

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))

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";

src/main/java/org/opensearch/flowframework/util/EncryptorUtils.java

@@ -250,7 +250,6 @@
 
             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 @@
                         .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 @@
                     }));
 
                 } 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 @@
             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 {

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;