Fixed to be able to overwrite the parameter in the retry command (#1725)

treasure-data · Jul 8, 2022 · 46d7507 · 46d7507
1 parent 8e3a53d
commit 46d7507
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 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
@@ -10,13 +10,17 @@
 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;
 import java.util.List;
 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
@@ -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)