@Before
public void setUp() throws Exception {
assumeThat(TD_API_KEY, not(isEmptyOrNullString()));
projectDir = folder.getRoot().toPath().toAbsolutePath().normalize();
config = folder.newFile().toPath();
Files.write(config, asList("secrets.td.apikey = " + TD_API_KEY));
outfile = projectDir.resolve("outfile");
client = TDClient.newBuilder(false).setApiKey(TD_API_KEY).build();
database = "tmp_" + UUID.randomUUID().toString().replace('-', '_');
client.createDatabase(database);
table = "test";
String insertJobId =
client.submit(
TDJobRequest.newPrestoQuery(database, "create table " + table + " as select 1"));
TestUtils.expect(Duration.ofMinutes(5), jobSuccess(client, insertJobId));
String selectCountJobId =
client.submit(TDJobRequest.newPrestoQuery(database, "select count(*) from " + table));
TestUtils.expect(Duration.ofMinutes(5), jobSuccess(client, selectCountJobId));
List<ArrayNode> result = downloadResult(selectCountJobId);
assertThat(result.get(0).get(0).asInt(), is(1));
}
@Test
public void testPartialDeleteWithEpochNumbers() throws Exception {
copyResource(
"acceptance/td/td_partial_delete/td_partial_delete_with_epoch_numbers.dig",
projectDir.resolve("workflow.dig"));
CommandStatus runStatus =
main(
"run",
"-o",
projectDir.toString(),
"--config",
config.toString(),
"--project",
projectDir.toString(),
"-p",
"outfile=" + outfile,
"-p",
"database=" + database,
"-p",
"table=" + table,
"workflow.dig");
assertThat(runStatus.errUtf8(), runStatus.code(), is(0));
assertThat(Files.exists(outfile), is(true));
// Verify that table contents are deleted
String selectCountJobId =
client.submit(TDJobRequest.newPrestoQuery(database, "select count(*) from " + table));
TestUtils.expect(Duration.ofMinutes(5), jobSuccess(client, selectCountJobId));
List<ArrayNode> result = downloadResult(selectCountJobId);
assertThat(result.get(0).get(0).asInt(), is(0));
}
private static Callable<Boolean> jobSuccess(TDClient client, String jobId) {
return () -> {
TDJobSummary status = client.jobStatus(jobId);
if (status.getStatus() == TDJob.Status.SUCCESS) {
return true;
}
if (status.getStatus().isFinished()) {
fail(status.getStatus().toString());
}
return false;
};
}
private List<ArrayNode> downloadResult(String jobId) {
return client.jobResult(
jobId,
TDResultFormat.JSON,
input -> {
try {
List<String> lines = CharStreams.readLines(new InputStreamReader(input));
ObjectReader reader = objectMapper().readerFor(ArrayNode.class);
List<ArrayNode> result = new ArrayList<>();
for (String line : lines) {
result.add(reader.readValue(line));
}
return result;
} catch (IOException e) {
throw Throwables.propagate(e);
}
});
}
@After
public void deleteDatabase() throws Exception {
if (client != null && database != null) {
client.deleteDatabase(database);
}
}
