From c821b4ac2d24af0297614dba4a76988b0db6e317 Mon Sep 17 00:00:00 2001
From: Daniel Widdis <widdis@gmail.com>
Date: Tue, 1 Oct 2024 19:31:47 -0700
Subject: [PATCH] Reorder methods for logical relationships

Signed-off-by: Daniel Widdis <widdis@gmail.com>
---
 .../indices/FlowFrameworkIndicesHandler.java  | 24 +++++++++----------
 1 file changed, 12 insertions(+), 12 deletions(-)

diff --git a/src/main/java/org/opensearch/flowframework/indices/FlowFrameworkIndicesHandler.java b/src/main/java/org/opensearch/flowframework/indices/FlowFrameworkIndicesHandler.java
index c9fcd19ad..004d9e819 100644
--- a/src/main/java/org/opensearch/flowframework/indices/FlowFrameworkIndicesHandler.java
+++ b/src/main/java/org/opensearch/flowframework/indices/FlowFrameworkIndicesHandler.java
@@ -578,14 +578,14 @@ public <T> void canDeleteWorkflowStateDoc(
     }
 
     /**
-     * Updates a complete document in the workflow state index
+     * Updates a partial document in the workflow state index
      * @param documentId the document ID
-     * @param updatedDocument a complete document to update the global state index with
+     * @param updatedFields the fields to update the global state index with
      * @param listener action listener
      */
     public void updateFlowFrameworkSystemIndexDoc(
         String documentId,
-        ToXContentObject updatedDocument,
+        Map<String, Object> updatedFields,
         ActionListener<UpdateResponse> listener
     ) {
         if (!doesIndexExist(WORKFLOW_STATE_INDEX)) {
@@ -595,11 +595,11 @@ public void updateFlowFrameworkSystemIndexDoc(
         } else {
             try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
                 UpdateRequest updateRequest = new UpdateRequest(WORKFLOW_STATE_INDEX, documentId);
-                XContentBuilder builder = XContentFactory.jsonBuilder();
-                updatedDocument.toXContent(builder, null);
-                updateRequest.doc(builder);
+                Map<String, Object> updatedContent = new HashMap<>(updatedFields);
+                updateRequest.doc(updatedContent);
                 updateRequest.setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE);
                 updateRequest.retryOnConflict(5);
+                // TODO: decide what condition can be considered as an update conflict and add retry strategy
                 client.update(updateRequest, ActionListener.runBefore(listener, context::restore));
             } catch (Exception e) {
                 String errorMessage = "Failed to update " + WORKFLOW_STATE_INDEX + " entry : " + documentId;
@@ -610,14 +610,14 @@ public void updateFlowFrameworkSystemIndexDoc(
     }
 
     /**
-     * Updates a partial document in the workflow state index
+     * Updates a complete document in the workflow state index
      * @param documentId the document ID
-     * @param updatedFields the fields to update the global state index with
+     * @param updatedDocument a complete document to update the global state index with
      * @param listener action listener
      */
     public void updateFlowFrameworkSystemIndexDoc(
         String documentId,
-        Map<String, Object> updatedFields,
+        ToXContentObject updatedDocument,
         ActionListener<UpdateResponse> listener
     ) {
         if (!doesIndexExist(WORKFLOW_STATE_INDEX)) {
@@ -627,11 +627,11 @@ public void updateFlowFrameworkSystemIndexDoc(
         } else {
             try (ThreadContext.StoredContext context = client.threadPool().getThreadContext().stashContext()) {
                 UpdateRequest updateRequest = new UpdateRequest(WORKFLOW_STATE_INDEX, documentId);
-                Map<String, Object> updatedContent = new HashMap<>(updatedFields);
-                updateRequest.doc(updatedContent);
+                XContentBuilder builder = XContentFactory.jsonBuilder();
+                updatedDocument.toXContent(builder, null);
+                updateRequest.doc(builder);
                 updateRequest.setRefreshPolicy(WriteRequest.RefreshPolicy.IMMEDIATE);
                 updateRequest.retryOnConflict(5);
-                // TODO: decide what condition can be considered as an update conflict and add retry strategy
                 client.update(updateRequest, ActionListener.runBefore(listener, context::restore));
             } catch (Exception e) {
                 String errorMessage = "Failed to update " + WORKFLOW_STATE_INDEX + " entry : " + documentId;