From 95abac64f6c4dd2a7f0a8e804a836fccb1f45fc3 Mon Sep 17 00:00:00 2001
From: hideki narimiya <hideki.narimiya@zozo.com>
Date: Thu, 17 Mar 2022 16:54:13 +0900
Subject: [PATCH] Fixed to be able to overwrite the parameter in the retry
 command

---
 .../main/java/io/digdag/cli/client/Retry.java | 14 ++++++++++++-
 .../src/test/java/acceptance/RetryIT.java     | 21 +++++++++++++++++--
 2 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/digdag-cli/src/main/java/io/digdag/cli/client/Retry.java b/digdag-cli/src/main/java/io/digdag/cli/client/Retry.java
index 7ebd0a05f1..acc808dfa4 100644
--- a/digdag-cli/src/main/java/io/digdag/cli/client/Retry.java
+++ b/digdag-cli/src/main/java/io/digdag/cli/client/Retry.java
@@ -10,6 +10,9 @@
 import io.digdag.client.api.RestSessionAttempt;
 import io.digdag.client.api.RestSessionAttemptRequest;
 import io.digdag.client.api.RestWorkflowDefinition;
+import io.digdag.client.config.Config;
+import io.digdag.client.config.ConfigFactory;
+import io.digdag.core.config.ConfigLoaderManager;
 
 import java.util.ArrayList;
 import java.util.HashMap;
@@ -17,6 +20,7 @@
 import java.util.Map;
 import java.util.UUID;
 
+import static io.digdag.cli.Arguments.loadParams;
 import static io.digdag.cli.SystemExitException.systemExit;
 import static java.util.Locale.ENGLISH;
 
@@ -91,6 +95,8 @@ public SystemExitException usage(String error)
         err.println("        --all                        retry all tasks");
         err.println("        --resume                     retry only non-successful tasks");
         err.println("        --resume-from <+name>        retry from a specific task");
+        err.println("    -p, --param KEY=VALUE            add a session parameter (use multiple times to set many parameters)");
+        err.println("    -P, --params-file PATH.yml       read session parameters from a YAML file");
         err.println("");
         return systemExit(error);
     }
@@ -128,11 +134,17 @@ else if (latestRevision) {
             retryAttemptName = UUID.randomUUID().toString();
         }
 
+        final ConfigFactory cf = injector.getInstance(ConfigFactory.class);
+        final ConfigLoaderManager loader = injector.getInstance(ConfigLoaderManager.class);
+
+        params = ParameterValidator.toMap(paramsList);
+        Config overrideParams = loadParams(cf, loader, loadSystemProperties(), paramsFile, params);
+
         RestSessionAttemptRequest request = RestSessionAttemptRequest.builder()
             .workflowId(workflowId)
             .sessionTime(attempt.getSessionTime().toInstant())
             .retryAttemptName(Optional.of(retryAttemptName))
-            .params(attempt.getParams())
+            .params(attempt.getParams().merge(overrideParams))
             .build();
 
         if (resumeFrom != null) {
diff --git a/digdag-tests/src/test/java/acceptance/RetryIT.java b/digdag-tests/src/test/java/acceptance/RetryIT.java
index 3c293d8cac..3894e9c25a 100644
--- a/digdag-tests/src/test/java/acceptance/RetryIT.java
+++ b/digdag-tests/src/test/java/acceptance/RetryIT.java
@@ -17,11 +17,9 @@
 
 import static java.nio.charset.StandardCharsets.UTF_8;
 import static java.util.Arrays.asList;
-import static utils.TestUtils.copyResource;
 import static utils.TestUtils.getAttemptId;
 import static utils.TestUtils.main;
 import static org.hamcrest.Matchers.is;
-import static org.hamcrest.Matchers.not;
 import static org.hamcrest.MatcherAssert.assertThat;
 
 public class RetryIT
@@ -160,6 +158,25 @@ public void testRetry()
         assertOutputExists("3-1", false);  // skipped
         assertOutputExists("3-2a", true);
         assertOutputExists("3-2b", true);
+
+        // Retry with parameters
+        Id retry5;
+        {
+            CommandStatus retryStatus = main("retry",
+                    "-c", config.toString(),
+                    "-e", server.endpoint(),
+                    "--latest-revision",
+                    "--all",
+                    "-p", "key=value",
+                    String.valueOf(originalAttemptId));
+            assertThat(retryStatus.errUtf8(), retryStatus.code(), is(0));
+            retry5 = getAttemptId(retryStatus);
+        }
+
+        // Wait for the attempt to success
+        RestSessionAttempt retry5Attempt = joinAttempt(client, retry5);
+        assertThat(retry5Attempt.getSuccess(), is(true));
+        assertThat(retry5Attempt.getParams().get("key", String.class), is("value"));
     }
 
     private void pushRevision(String resourceName, String workflowName)