Skip to content

Commit 04bc47a

Browse files
yoyamayoyama
authored
Introduce disable direct download feature in CLI. (#1275)
Introduce disable direct download feature in CLI.
1 parent 508abad commit 04bc47a

5 files changed

Lines changed: 300 additions & 8 deletions

File tree

‎digdag-cli/src/main/java/io/digdag/cli/client/ClientCommand.java‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ public abstract class ClientCommand
3131
private static final Logger logger = LoggerFactory.getLogger(ClientCommand.class);
3232

3333
private static final String DEFAULT_ENDPOINT = "http://127.0.0.1:65432";
34+
private static final String DEFAULT_DISABLE_DIRECT_DOWNLOAD = "false";
3435

3536
@Inject Injector injector;
3637

@@ -170,10 +171,12 @@ private boolean isBatchModeForVersionCheck()
170171
}
171172

172173
@VisibleForTesting
173-
static DigdagClient buildClient(String endpoint, Map<String, String> env, Properties props, boolean disableCertValidation, Map<String, String> httpHeaders, Iterable<DigdagClientConfigurator> clientConfigurators)
174+
static DigdagClient buildClient(String endpoint, Map<String, String> env, Properties props, boolean disableCertValidation,
175+
Map<String, String> httpHeaders, Iterable<DigdagClientConfigurator> clientConfigurators)
174176
throws SystemExitException
175177
{
176178
String[] fragments = endpoint.split(":", 2);
179+
boolean disableDirectDownload = Boolean.parseBoolean(props.getProperty("client.http.disable_direct_download", DEFAULT_DISABLE_DIRECT_DOWNLOAD));
177180

178181
boolean useSsl = false;
179182
if (fragments.length == 2 && fragments[1].startsWith("//")) {
@@ -219,6 +222,7 @@ static DigdagClient buildClient(String endpoint, Map<String, String> env, Proper
219222
.port(port)
220223
.ssl(useSsl)
221224
.disableCertValidation(disableCertValidation)
225+
.disableDirectDownload(disableDirectDownload)
222226
.headers(headers);
223227

224228
Optional<ProxyConfig> proxyConfig = Proxies.proxyConfigFromEnv(scheme, env);

‎digdag-client/src/main/java/io/digdag/client/DigdagClient.java‎

Lines changed: 32 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -104,6 +104,7 @@ public static class Builder
104104
private final Map<String, String> baseHeaders = new HashMap<>();
105105
private Function<Map<String, String>, Map<String, String>> headerBuilder = null;
106106
private boolean disableCertValidation;
107+
private boolean disableDirectDownload = false;
107108

108109
public Builder host(String host)
109110
{
@@ -165,6 +166,12 @@ public Builder disableCertValidation(boolean value)
165166
return this;
166167
}
167168

169+
public Builder disableDirectDownload(boolean value)
170+
{
171+
this.disableDirectDownload = value;
172+
return this;
173+
}
174+
168175
public DigdagClient build()
169176
{
170177
return new DigdagClient(this);
@@ -197,6 +204,7 @@ public static Builder builder()
197204

198205
private final Client client;
199206
private final ConfigFactory cf;
207+
private final boolean disableDirectDownload;
200208

201209
private DigdagClient(Builder builder)
202210
{
@@ -261,6 +269,9 @@ private DigdagClient(Builder builder)
261269
this.client = clientBuilder.build();
262270

263271
this.cf = new ConfigFactory(mapper);
272+
273+
this.disableDirectDownload = builder.disableDirectDownload;
274+
264275
}
265276

266277
@Override
@@ -473,7 +484,7 @@ public InputStream getProjectArchive(Id projId, String revision)
473484
WebTarget webTarget = target("/api/projects/{id}/archive")
474485
.resolveTemplate("id", projId)
475486
.queryParam("revision", revision);
476-
487+
webTarget = addDisableDirectDownloadParam(webTarget);
477488
return withFollowingRedirect(webTarget,
478489
(wt, lastResponse) -> {
479490
Invocation.Builder request = wt.request();
@@ -628,18 +639,22 @@ public RestTaskCollection getTasks(Id attemptId)
628639

629640
public RestLogFileHandleCollection getLogFileHandlesOfAttempt(Id attemptId)
630641
{
631-
return doGet(RestLogFileHandleCollection.class,
632-
target("/api/logs/{id}/files")
633-
.resolveTemplate("id", attemptId));
642+
WebTarget webTarget = target("/api/logs/{id}/files")
643+
.resolveTemplate("id", attemptId);
644+
webTarget = addDisableDirectDownloadParam(webTarget);
645+
646+
return doGet(RestLogFileHandleCollection.class, webTarget);
634647
}
635648

636649
public RestLogFileHandleCollection getLogFileHandlesOfTask(Id attemptId, String taskName)
637650
{
638651
try {
639-
return doGet(RestLogFileHandleCollection.class,
640-
target("/api/logs/{id}/files")
652+
WebTarget webTarget = target("/api/logs/{id}/files")
641653
.resolveTemplate("id", attemptId)
642-
.queryParam("task", URLEncoder.encode(taskName, "UTF-8")));
654+
.queryParam("task", URLEncoder.encode(taskName, "UTF-8"));
655+
webTarget = addDisableDirectDownloadParam(webTarget);
656+
657+
return doGet(RestLogFileHandleCollection.class, webTarget);
643658
} catch (UnsupportedEncodingException ex) {
644659
throw Throwables.propagate(ex);
645660
}
@@ -671,6 +686,16 @@ public InputStream getLogFile(Id attemptId, String fileName)
671686
.readEntity(InputStream.class);
672687
}
673688

689+
private WebTarget addDisableDirectDownloadParam(WebTarget target)
690+
{
691+
// direct_download is default true
692+
// So only set direct_download=false when disableDirectDownload == true
693+
if (disableDirectDownload) {
694+
target = target.queryParam("direct_download", false);
695+
}
696+
return target;
697+
}
698+
674699
private Response invokeWithRetry(Invocation request)
675700
{
676701
Retryer<Response> retryer = RetryerBuilder.<Response>newBuilder()

‎digdag-docs/src/command_reference.rst‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -445,10 +445,13 @@ Client-mode common options:
445445

446446
Example: ``-c digdag-server/client.properties``
447447

448+
449+
448450
You can include following parameters in ~/.config/digdag/config file:
449451

450452
* client.http.endpoint = http://HOST:PORT or https://HOST:PORT
451453
* client.http.headers.KEY = VALUE (set custom HTTP header)
454+
* client.http.disable_direct_download=true (disable direct download in `log` and `download`. effect to server v0.10.0(not yet released) or later.)
452455

453456

454457
start

‎digdag-tests/src/test/java/acceptance/CliDownloadIT.java‎

Lines changed: 140 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,140 @@
1+
package acceptance;
2+
3+
import io.netty.handler.codec.http.FullHttpRequest;
4+
import org.junit.After;
5+
import org.junit.Before;
6+
import org.junit.Rule;
7+
import org.junit.Test;
8+
import org.junit.rules.TemporaryFolder;
9+
import org.littleshoot.proxy.HttpProxyServer;
10+
import utils.CommandStatus;
11+
import utils.TemporaryDigdagServer;
12+
import utils.TestUtils;
13+
14+
import java.nio.file.Files;
15+
import java.nio.file.Path;
16+
import java.util.ArrayList;
17+
import java.util.Arrays;
18+
import java.util.Collections;
19+
import java.util.HashMap;
20+
import java.util.List;
21+
import java.util.Map;
22+
23+
import static org.hamcrest.MatcherAssert.assertThat;
24+
import static org.hamcrest.Matchers.is;
25+
import static utils.TestUtils.copyResource;
26+
import static utils.TestUtils.main;
27+
28+
public class CliDownloadIT
29+
{
30+
@Rule
31+
public TemporaryFolder folder = new TemporaryFolder();
32+
33+
public TemporaryDigdagServer server;
34+
private HttpProxyServer proxyServer;
35+
36+
private Map<String, String> env;
37+
private Path config;
38+
private Path projectDir;
39+
40+
private List<FullHttpRequest> requests = Collections.synchronizedList(new ArrayList<>());
41+
42+
@Before
43+
public void setUp()
44+
throws Exception
45+
{
46+
env = new HashMap<>();
47+
proxyServer = TestUtils.startRequestTrackingProxy(requests);
48+
server = TemporaryDigdagServer.builder()
49+
.withRandomSecretEncryptionKey()
50+
.build();
51+
server.start();
52+
projectDir = folder.getRoot().toPath().resolve("foobar");
53+
String proxyUrl = "http://" + proxyServer.getListenAddress().getHostString() + ":" + proxyServer.getListenAddress().getPort();
54+
env.put("http_proxy", proxyUrl);
55+
56+
config = folder.newFile().toPath();
57+
}
58+
59+
@Test
60+
public void disableDirectDownload()
61+
throws Exception
62+
{
63+
// Create new project
64+
CommandStatus initStatus = main("init",
65+
"-c", config.toString(),
66+
projectDir.toString());
67+
assertThat(initStatus.code(), is(0));
68+
69+
copyResource("acceptance/basic.dig", projectDir.resolve("basic.dig"));
70+
71+
{
72+
CommandStatus pushStatus = main(
73+
"push",
74+
"--project", projectDir.toString(),
75+
"test_proj",
76+
"-c", config.toString(),
77+
"-e", server.endpoint()
78+
);
79+
assertThat(pushStatus.errUtf8(), pushStatus.code(), is(0));
80+
}
81+
82+
// No configuration "client.http.disable_direct_download=true"
83+
{
84+
requests.clear();
85+
CommandStatus status = main(env,
86+
"download",
87+
"test_proj",
88+
"-c", config.toString(),
89+
"-e", server.endpoint(),
90+
"-o", folder.getRoot().toPath().resolve("test1").toString()
91+
);
92+
assertThat(status.errUtf8(), status.code(), is(0));
93+
boolean match = false;
94+
for (FullHttpRequest req : requests) {
95+
if (req.uri().matches(".*/api/projects/.*/archive.*")) {
96+
match = true;
97+
assertThat("direct_download= must not be set.",
98+
!req.uri().matches(".*direct_download=.*"));
99+
}
100+
}
101+
assertThat("No record", match);
102+
}
103+
104+
// Set configuration "client.http.disable_direct_download=true"
105+
{
106+
requests.clear();
107+
Files.write(config, Arrays.asList("client.http.disable_direct_download=true"));
108+
CommandStatus status = main(env,
109+
"download",
110+
"test_proj",
111+
"-c", config.toString(),
112+
"-e", server.endpoint(),
113+
"-o", folder.getRoot().toPath().resolve("test3").toString()
114+
);
115+
assertThat(status.errUtf8(), status.code(), is(0));
116+
boolean match = false;
117+
for (FullHttpRequest req : requests) {
118+
if (req.uri().matches(".*/api/projects/.*/archive.*")) {
119+
match = true;
120+
assertThat("direct_download= must be set.",req.uri().matches(".*direct_download=.*"));
121+
}
122+
}
123+
assertThat("No record", match);
124+
}
125+
}
126+
127+
@After
128+
public void tearDown()
129+
throws Exception
130+
{
131+
if (proxyServer != null) {
132+
proxyServer.stop();
133+
proxyServer = null;
134+
}
135+
if (server != null) {
136+
server.close();
137+
server = null;
138+
}
139+
}
140+
}

‎digdag-tests/src/test/java/acceptance/CliLogIT.java‎

Lines changed: 120 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,120 @@
1+
package acceptance;
2+
3+
import io.digdag.client.DigdagClient;
4+
import io.netty.handler.codec.http.FullHttpRequest;
5+
import org.junit.After;
6+
import org.junit.Before;
7+
import org.junit.Rule;
8+
import org.junit.Test;
9+
import org.junit.rules.TemporaryFolder;
10+
import org.littleshoot.proxy.HttpProxyServer;
11+
import utils.CommandStatus;
12+
import utils.TemporaryDigdagServer;
13+
import utils.TestUtils;
14+
15+
import java.nio.file.Files;
16+
import java.nio.file.Path;
17+
import java.util.ArrayList;
18+
import java.util.Arrays;
19+
import java.util.Collections;
20+
import java.util.HashMap;
21+
import java.util.List;
22+
import java.util.Map;
23+
24+
import static org.hamcrest.MatcherAssert.assertThat;
25+
import static org.hamcrest.Matchers.is;
26+
import static utils.TestUtils.copyResource;
27+
import static utils.TestUtils.main;
28+
29+
public class CliLogIT
30+
{
31+
@Rule
32+
public TemporaryFolder folder = new TemporaryFolder();
33+
34+
public TemporaryDigdagServer server;
35+
private HttpProxyServer proxyServer;
36+
37+
private Map<String, String> env;
38+
private Path config;
39+
private Path projectDir;
40+
41+
private List<FullHttpRequest> requests = Collections.synchronizedList(new ArrayList<>());
42+
43+
@Before
44+
public void setUp()
45+
throws Exception
46+
{
47+
env = new HashMap<>();
48+
proxyServer = TestUtils.startRequestTrackingProxy(requests);
49+
server = TemporaryDigdagServer.builder()
50+
.withRandomSecretEncryptionKey()
51+
.build();
52+
server.start();
53+
projectDir = folder.getRoot().toPath().resolve("foobar");
54+
String proxyUrl = "http://" + proxyServer.getListenAddress().getHostString() + ":" + proxyServer.getListenAddress().getPort();
55+
env.put("http_proxy", proxyUrl);
56+
57+
config = folder.newFile().toPath();
58+
59+
}
60+
61+
@Test
62+
public void disableDirectDownload()
63+
throws Exception
64+
{
65+
// No configuration "client.http.disable_direct_download=true"
66+
{
67+
requests.clear();
68+
CommandStatus status = main(env,
69+
"log",
70+
"1",
71+
"-c", config.toString(),
72+
"-e", server.endpoint()
73+
);
74+
boolean match = false;
75+
for (FullHttpRequest req :requests) {
76+
if (req.uri().matches(".*/api/logs/1/.*")) {
77+
assertThat("direct_download= must not be set.",
78+
!req.uri().matches(".*direct_download=.*"));
79+
match = true;
80+
}
81+
}
82+
assertThat("No record", match);
83+
}
84+
85+
// Set configuration "client.http.disable_direct_download=true"
86+
{
87+
requests.clear();
88+
Files.write(config, Arrays.asList("client.http.disable_direct_download=true"));
89+
CommandStatus status = main(env,
90+
"log",
91+
"3",
92+
"-c", config.toString(),
93+
"-e", server.endpoint()
94+
);
95+
boolean match = false;
96+
for (FullHttpRequest req :requests) {
97+
if (req.uri().matches(".*/api/logs/3/.*")) {
98+
assertThat("direct_download=false must be set with client.http.disable_direct_download=true",
99+
req.uri().matches(".*direct_download=false.*"));
100+
match = true;
101+
}
102+
}
103+
assertThat("No record", match);
104+
}
105+
}
106+
107+
@After
108+
public void tearDown()
109+
throws Exception
110+
{
111+
if (proxyServer != null) {
112+
proxyServer.stop();
113+
proxyServer = null;
114+
}
115+
if (server != null) {
116+
server.close();
117+
server = null;
118+
}
119+
}
120+
}

0 commit comments

Comments
 (0)