Skip to content

Commit 68f0d86

Browse files
takemikamitakemikami
authored
Merge pull request #7 from takemikami/support_v0.10.0
support digdag v0.10.0
2 parents 05c3d35 + 600d46e commit 68f0d86

File tree

4 files changed

+138
-110
lines changed

4 files changed

+138
-110
lines changed

‎build.gradle

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ plugins {
1111
group = 'com.github.takemikami'
1212
version = '0.0.2'
1313

14-
def digdagVersion = '0.9.25'
14+
def digdagVersion = '0.10.0'
1515

1616
repositories {
1717
mavenCentral()

‎gradle/wrapper/gradle-wrapper.properties

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
distributionBase=GRADLE_USER_HOME
22
distributionPath=wrapper/dists
3-
distributionUrl=https\://services.gradle.org/distributions/gradle-4.8-bin.zip
3+
distributionUrl=https\://services.gradle.org/distributions/gradle-6.7-bin.zip
44
zipStoreBase=GRADLE_USER_HOME
55
zipStorePath=wrapper/dists

‎src/main/java/com/takemikami/github/digdag/plugin/shresult/ShResultOperatorFactory.java

Lines changed: 136 additions & 65 deletions
Original file line numberDiff line numberDiff line change
@@ -3,32 +3,39 @@
33
import com.fasterxml.jackson.core.type.TypeReference;
44
import com.fasterxml.jackson.databind.JsonNode;
55
import com.fasterxml.jackson.databind.ObjectMapper;
6+
import com.fasterxml.jackson.databind.node.ObjectNode;
67
import com.google.common.base.Throwables;
78
import com.google.common.collect.ImmutableList;
9+
import com.google.common.collect.Maps;
810
import com.google.inject.Inject;
911
import io.digdag.client.config.Config;
10-
import io.digdag.client.config.ConfigException;
12+
import io.digdag.client.config.ConfigElement;
1113
import io.digdag.client.config.ConfigFactory;
14+
import io.digdag.spi.CommandContext;
1215
import io.digdag.spi.CommandExecutor;
16+
import io.digdag.spi.CommandRequest;
17+
import io.digdag.spi.CommandStatus;
1318
import io.digdag.spi.Operator;
1419
import io.digdag.spi.OperatorContext;
1520
import io.digdag.spi.OperatorFactory;
16-
import io.digdag.spi.PrivilegedVariables;
1721
import io.digdag.spi.TaskExecutionException;
1822
import io.digdag.spi.TaskResult;
1923
import io.digdag.util.BaseOperator;
24+
import io.digdag.util.CommandOperators;
2025
import io.digdag.util.UserSecretTemplate;
21-
import java.io.BufferedWriter;
26+
import java.io.ByteArrayOutputStream;
2227
import java.io.IOException;
23-
import java.io.OutputStreamWriter;
28+
import java.io.PrintStream;
2429
import java.io.Writer;
30+
import java.nio.charset.StandardCharsets;
31+
import java.nio.file.Files;
32+
import java.nio.file.Path;
33+
import java.time.Duration;
2534
import java.util.ArrayList;
2635
import java.util.HashMap;
2736
import java.util.LinkedList;
2837
import java.util.List;
2938
import java.util.Map;
30-
import java.util.regex.Pattern;
31-
import org.apache.commons.io.IOUtils;
3239
import org.slf4j.Logger;
3340
import org.slf4j.LoggerFactory;
3441

@@ -37,8 +44,6 @@ public class ShResultOperatorFactory implements OperatorFactory {
3744

3845
private static Logger logger = LoggerFactory.getLogger(ShResultOperatorFactory.class);
3946

40-
private static Pattern VALID_ENV_KEY = Pattern.compile("[a-zA-Z_][a-zA-Z_0-9]*");
41-
4247
private final CommandExecutor exec;
4348

4449
@Inject
@@ -58,93 +63,164 @@ public Operator newOperator(OperatorContext operatorContext) {
5863

5964
class ShResultOperator extends BaseOperator {
6065

66+
// TODO extract as config params.
67+
final int scriptPollInterval = (int) Duration.ofSeconds(10).getSeconds();
68+
6169
public ShResultOperator(OperatorContext context) {
6270
super(context);
6371
}
6472

6573
@Override
6674
public TaskResult runTask() {
67-
Config params = request.getConfig()
75+
final Config state = request.getConfig();
76+
Config params = state.mergeDefault(request.getConfig().getNestedOrGetEmpty("sh"));
77+
78+
// save System.out
79+
PrintStream console = System.out;
80+
81+
try {
82+
// prepare capture stdout
83+
final ByteArrayOutputStream baos = new ByteArrayOutputStream();
84+
final String utf8 = StandardCharsets.UTF_8.name();
85+
PrintStream ps = new PrintStream(baos);
86+
System.setOut(ps);
87+
88+
// run script
89+
runCode(state);
90+
91+
// capture stdout
92+
String stdoutData = baos.toString(utf8);
93+
94+
String varName = params.get("destination_variable", String.class);
95+
String stdoutFormat = params.get("stdout_format", String.class);
96+
97+
ConfigFactory cf = request.getConfig().getFactory();
98+
Config storeParams = cf.create();
99+
storeParams.set(varName, createVariableObjectFromStdout(stdoutData, stdoutFormat));
100+
101+
return TaskResult.defaultBuilder(request)
102+
.storeParams(storeParams)
103+
.build();
104+
} catch (IOException | InterruptedException e) {
105+
throw Throwables.propagate(e);
106+
} finally {
107+
// restore System.out
108+
System.setOut(console);
109+
}
110+
}
111+
112+
private void runCode(final Config state)
113+
throws IOException, InterruptedException {
114+
final Config params = request.getConfig()
68115
.mergeDefault(request.getConfig().getNestedOrGetEmpty("sh"));
116+
final Path projectPath = workspace.getProjectPath();
117+
final CommandContext commandContext = buildCommandContext(projectPath);
118+
119+
final CommandStatus status;
120+
if (!state.has("commandStatus")) {
121+
// Run the code since command state doesn't exist
122+
status = runCommand(params, commandContext);
123+
} else {
124+
// Check the status of the running command
125+
final ObjectNode previousStatusJson = state.get("commandStatus", ObjectNode.class);
126+
status = exec.poll(commandContext, previousStatusJson);
127+
}
69128

70-
List<String> shell = params.getListOrEmpty("shell", String.class);
71-
if (shell.isEmpty()) {
129+
if (status.isFinished()) {
130+
final int statusCode = status.getStatusCode();
131+
if (statusCode != 0) {
132+
// Remove the polling state after fetching the result so that the result fetch can be retried
133+
// without resubmitting the code.
134+
state.remove("commandStatus");
135+
throw new RuntimeException("Command failed with code " + statusCode);
136+
}
137+
return;
138+
} else {
139+
state.set("commandStatus", status);
140+
throw TaskExecutionException.ofNextPolling(scriptPollInterval, ConfigElement.copyOf(state));
141+
}
142+
}
143+
144+
private CommandStatus runCommand(final Config params, final CommandContext commandContext)
145+
throws IOException, InterruptedException {
146+
final Path tempDir = workspace
147+
.createTempDir(String.format("digdag-sh-%d-", request.getTaskId()));
148+
final Path workingDirectory = workspace.getPath(); // absolute
149+
final Path runnerPath = tempDir.resolve("runner.sh"); // absolute
150+
151+
final List<String> shell;
152+
if (params.has("shell")) {
153+
shell = params.getListOrEmpty("shell", String.class);
154+
} else {
72155
shell = ImmutableList.of("/bin/sh");
73156
}
74-
String command = UserSecretTemplate.of(params.get("_command", String.class))
75-
.format(context.getSecrets());
76157

77-
ProcessBuilder pb = new ProcessBuilder(shell);
78-
pb.directory(workspace.getPath().toFile());
158+
final ImmutableList.Builder<String> cmdline = ImmutableList.builder();
159+
if (params.has("shell")) {
160+
cmdline.addAll(shell);
161+
} else {
162+
cmdline.addAll(shell);
163+
}
164+
cmdline.add(workingDirectory.relativize(runnerPath).toString()); // relative
165+
166+
final String shScript = UserSecretTemplate.of(params.get("_command", String.class))
167+
.format(context.getSecrets());
79168

80-
final Map<String, String> env = pb.environment();
169+
final Map<String, String> environments = Maps.newHashMap();
81170
params.getKeys()
82171
.forEach(key -> {
83-
if (isValidEnvKey(key)) {
172+
if (CommandOperators.isValidEnvKey(key)) {
84173
JsonNode value = params.get(key, JsonNode.class);
85174
String string;
86175
if (value.isTextual()) {
87176
string = value.textValue();
88177
} else {
89178
string = value.toString();
90179
}
91-
env.put(key, string);
180+
environments.put(key, string);
92181
} else {
93182
logger.trace("Ignoring invalid env var key: {}", key);
94183
}
95184
});
96185

97186
// Set up process environment according to env config. This can also refer to secrets.
98-
collectEnvironmentVariables(env, context.getPrivilegedVariables());
99-
100-
String stdoutData;
101-
int ecode;
102-
try {
103-
Process p = exec.start(workspace.getPath(), request, pb);
104-
105-
// feed command to stdin
106-
try (Writer writer = new BufferedWriter(new OutputStreamWriter(p.getOutputStream()))) {
107-
writer.write(command);
108-
}
109-
110-
ecode = p.waitFor();
111-
112-
// keep stdout
113-
stdoutData = IOUtils.toString(p.getInputStream());
114-
115-
// dump stderr
116-
String stderrData = IOUtils.toString(p.getErrorStream());
117-
logger.info(stderrData);
187+
CommandOperators.collectEnvironmentVariables(environments, context.getPrivilegedVariables());
118188

119-
} catch (IOException | InterruptedException ex) {
120-
throw Throwables.propagate(ex);
189+
// Write script content to runnerPath
190+
try (Writer writer = Files.newBufferedWriter(runnerPath)) {
191+
writer.write(shScript);
121192
}
122193

123-
if (ecode != 0) {
124-
throw new TaskExecutionException("Command failed with code " + ecode);
125-
}
126-
127-
String varName = params.get("destination_variable", String.class);
128-
String stdoutFormat = params.get("stdout_format", String.class);
129-
130-
ConfigFactory cf = request.getConfig().getFactory();
131-
Config storeParams = cf.create();
194+
final CommandRequest commandRequest = buildCommandRequest(commandContext, workingDirectory,
195+
tempDir, environments, cmdline.build());
196+
return exec.run(commandContext, commandRequest);
132197

133-
storeParams.set(varName, createVariableObjectFromStdout(stdoutData, stdoutFormat));
198+
// TaskExecutionException could not be thrown here to poll the task by non-blocking for process-base
199+
// command executor. Because they will be bounded by the _instance_ where the command was executed
200+
// first.
201+
}
134202

135-
return TaskResult.defaultBuilder(request)
136-
.storeParams(storeParams)
203+
private CommandContext buildCommandContext(final Path projectPath) {
204+
return CommandContext.builder()
205+
.localProjectPath(projectPath)
206+
.taskRequest(this.request)
137207
.build();
138208
}
139-
}
140209

141-
public static void collectEnvironmentVariables(Map<String, String> env,
142-
PrivilegedVariables variables) {
143-
for (String name : variables.getKeys()) {
144-
if (!VALID_ENV_KEY.matcher(name).matches()) {
145-
throw new ConfigException("Invalid _env key name: " + name);
146-
}
147-
env.put(name, variables.get(name));
210+
private CommandRequest buildCommandRequest(final CommandContext commandContext,
211+
final Path workingDirectory,
212+
final Path tempDir,
213+
final Map<String, String> environments,
214+
final List<String> cmdline) {
215+
final Path projectPath = commandContext.getLocalProjectPath();
216+
final Path relativeWorkingDirectory = projectPath.relativize(workingDirectory); // relative
217+
final Path ioDirectory = projectPath.relativize(tempDir); // relative
218+
return CommandRequest.builder()
219+
.workingDirectory(relativeWorkingDirectory)
220+
.environments(environments)
221+
.commandLine(cmdline)
222+
.ioDirectory(ioDirectory)
223+
.build();
148224
}
149225
}
150226

@@ -190,9 +266,4 @@ public static Object createVariableObjectFromStdout(
190266
}
191267
return null;
192268
}
193-
194-
private static boolean isValidEnvKey(String key) {
195-
return VALID_ENV_KEY.matcher(key).matches();
196-
}
197-
198269
}

‎src/test/java/com/takemikami/github/digdag/plugin/shresult/ShResultOperatorFactoryTest.java

Lines changed: 0 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -62,47 +62,4 @@ public void testCreateVariableObjectFromStdoutSpaceDelimited() throws Exception
6262
assertEquals("bar", obj.get(3));
6363
}
6464

65-
@Test
66-
public void testCollectEnvironmentVariablesSuccess() throws Exception {
67-
Map<String, String> env = new HashMap<>();
68-
HashMap<String, String> map = new HashMap<>();
69-
map.put("key", "value");
70-
PrivilegedVariables variables = new TestPrivilegedVariables(map);
71-
ShResultOperatorFactory.collectEnvironmentVariables(env, variables);
72-
assertEquals("value", env.get("key"));
73-
}
74-
75-
@Test(expected = ConfigException.class)
76-
public void testCollectEnvironmentVariablesExceptionInvalidKey() throws Exception {
77-
Map<String, String> env = new HashMap<>();
78-
HashMap<String, String> map = new HashMap<>();
79-
map.put("123", "value");
80-
PrivilegedVariables variables = new TestPrivilegedVariables(map);
81-
ShResultOperatorFactory.collectEnvironmentVariables(env, variables);
82-
}
83-
84-
class TestPrivilegedVariables implements PrivilegedVariables {
85-
86-
private Map<String, String> map;
87-
88-
public TestPrivilegedVariables(Map<String, String> map) {
89-
this.map = map;
90-
}
91-
92-
@Override
93-
public String get(String s) {
94-
return map.get(s);
95-
}
96-
97-
@Override
98-
public Optional<String> getOptional(String s) {
99-
return Optional.of(get(s));
100-
}
101-
102-
@Override
103-
public List<String> getKeys() {
104-
return new ArrayList<String>(map.keySet());
105-
}
106-
}
107-
10865
}

0 commit comments

Comments
 (0)