diff --git a/src/main/java/org/embulk/output/bigquery_java/BigqueryJavaOutputPlugin.java b/src/main/java/org/embulk/output/bigquery_java/BigqueryJavaOutputPlugin.java index c7c1510..888e0b9 100644 --- a/src/main/java/org/embulk/output/bigquery_java/BigqueryJavaOutputPlugin.java +++ b/src/main/java/org/embulk/output/bigquery_java/BigqueryJavaOutputPlugin.java @@ -50,114 +50,112 @@ public ConfigDiff transaction( BigqueryTaskBuilder.build(task); BigqueryClient client = new BigqueryClient(task, schema); autoCreate(task, client); - client.storeCachedSrcFieldsIfNeed(); - - control.run(task.dump()); - this.writers.values().forEach(BigqueryFileWriter::close); - logger.info("embulk-output-bigquery: finish to create intermediate files"); try { - paths = BigqueryUtil.getIntermediateFiles(task); - } catch (Exception e) { - logger.info(e.getMessage()); - throw new RuntimeException(e); - } - if (paths.isEmpty()) { - logger.info("embulk-output-bigquery: Nothing for transfer"); - client.createTableIfNotExist(task.getTable()); - - switch (task.getMode()) { - case "merge": - case "append": - case "replace": - case "delete_in_advance": - if (task.getTempTable().isPresent()) { - client.deleteTable(task.getTempTable().get()); - } - break; - } - - return CONFIG_MAPPER_FACTORY.newConfigDiff(); - } - - logger.debug( - "embulk-output-bigquery: LOAD IN PARALLEL {}", - paths.stream().map(Path::toString).collect(Collectors.joining("\n"))); - - // transfer data to BQ from files - ExecutorService executor = Executors.newFixedThreadPool(paths.size()); - List> jobStatisticFutures = new ArrayList<>(); - List statistics = new ArrayList<>(); + client.storeCachedSrcFieldsIfNeed(); - for (Path path : paths) { - Future jobStatisticsFuture = - executor.submit(new BigqueryJobRunner(task, schema, path)); - jobStatisticFutures.add(jobStatisticsFuture); - } + control.run(task.dump()); + this.writers.values().forEach(BigqueryFileWriter::close); + logger.info("embulk-output-bigquery: finish to create intermediate files"); - for (Future jobStatisticFuture : jobStatisticFutures) { try { - statistics.add((JobStatistics.LoadStatistics) jobStatisticFuture.get()); + paths = BigqueryUtil.getIntermediateFiles(task); } catch (Exception e) { + logger.info(e.getMessage()); throw new RuntimeException(e); } - } - BigqueryTransactionReport report = - getTransactionReport(task, client, statistics, this.writers.values()); - if (task.getAbortOnError().get() && !task.getIsSkipJobResultCheck()) { - if (report.getNumInputRows().compareTo(report.getNumOutputRows()) != 0) { - String msg = - String.format( - "ABORT: `num_input_rows (%d)` and `num_output_rows (%d)` does not match", - report.getNumInputRows(), report.getNumOutputRows()); - throw new RuntimeException(msg); + if (paths.isEmpty()) { + logger.info("embulk-output-bigquery: Nothing for transfer"); + client.createTableIfNotExist(task.getTable()); + return CONFIG_MAPPER_FACTORY.newConfigDiff(); } - } - if (task.getMode().equals("append") && task.getBeforeLoad().isPresent()) { - logger.info("embulk-output-bigquery: before_load will be executed"); - logger.info("embulk-output-bigquery: {}", task.getBeforeLoad().get()); - client.executeQuery(task.getBeforeLoad().get()); - } + logger.debug( + "embulk-output-bigquery: LOAD IN PARALLEL {}", + paths.stream().map(Path::toString).collect(Collectors.joining("\n"))); - if (task.getMode().equals("replace_backup")) { - if (task.getOldTable().isPresent()) { - client.copy( - task.getTable(), - task.getOldTable().get(), - task.getOldDataset().orElse(client.destinationDataset), - JobInfo.WriteDisposition.WRITE_TRUNCATE); + // transfer data to BQ from files + ExecutorService executor = Executors.newFixedThreadPool(paths.size()); + List> jobStatisticFutures = new ArrayList<>(); + List statistics = new ArrayList<>(); + + for (Path path : paths) { + Future jobStatisticsFuture = + executor.submit(new BigqueryJobRunner(task, schema, path)); + jobStatisticFutures.add(jobStatisticsFuture); } - } - if (task.getTempTable().isPresent()) { - if (task.getMode().equals("merge")) { - client.merge( - task.getTempTable().get(), - task.getTable(), - task.getMergeKeys().orElse(Collections.emptyList()), - task.getMergeRule().orElse(Collections.emptyList())); - } else if (task.getMode().equals("append")) { - client.copy( - task.getTempTable().get(), task.getTable(), JobInfo.WriteDisposition.WRITE_APPEND); - } else { - client.copy( - task.getTempTable().get(), task.getTable(), JobInfo.WriteDisposition.WRITE_TRUNCATE); + + for (Future jobStatisticFuture : jobStatisticFutures) { + try { + statistics.add((JobStatistics.LoadStatistics) jobStatisticFuture.get()); + } catch (Exception e) { + throw new RuntimeException(e); + } + } + BigqueryTransactionReport report = + getTransactionReport(task, client, statistics, this.writers.values()); + if (task.getAbortOnError().get() && !task.getIsSkipJobResultCheck()) { + if (report.getNumInputRows().compareTo(report.getNumOutputRows()) != 0) { + String msg = + String.format( + "ABORT: `num_input_rows (%d)` and `num_output_rows (%d)` does not match", + report.getNumInputRows(), report.getNumOutputRows()); + throw new RuntimeException(msg); + } } - client.deleteTable(task.getTempTable().get()); - } - client.updateTableIfNeed(); + if (task.getMode().equals("append") && task.getBeforeLoad().isPresent()) { + logger.info("embulk-output-bigquery: before_load will be executed"); + logger.info("embulk-output-bigquery: {}", task.getBeforeLoad().get()); + client.executeQuery(task.getBeforeLoad().get()); + } - if (task.getDeleteFromLocalWhenJobEnd()) { - paths.forEach(p -> p.toFile().delete()); - } else { - paths.forEach( - p -> { - File intermediateFile = new File(p.toString()); - if (intermediateFile.exists()) { - logger.info("embulk-output-bigquery: keep {}", p.toString()); - } - }); + if (task.getMode().equals("replace_backup")) { + if (task.getOldTable().isPresent()) { + client.copy( + task.getTable(), + task.getOldTable().get(), + task.getOldDataset().orElse(client.destinationDataset), + JobInfo.WriteDisposition.WRITE_TRUNCATE); + } + } + if (task.getTempTable().isPresent()) { + if (task.getMode().equals("merge")) { + client.merge( + task.getTempTable().get(), + task.getTable(), + task.getMergeKeys().orElse(Collections.emptyList()), + task.getMergeRule().orElse(Collections.emptyList())); + } else if (task.getMode().equals("append")) { + client.copy( + task.getTempTable().get(), task.getTable(), JobInfo.WriteDisposition.WRITE_APPEND); + } else { + client.copy( + task.getTempTable().get(), task.getTable(), JobInfo.WriteDisposition.WRITE_TRUNCATE); + } + } + + client.updateTableIfNeed(); + } finally { + try { + if (task.getTempTable().isPresent()) { + client.deleteTable(task.getTempTable().get()); + } + } finally { + if (paths != null) { + if (task.getDeleteFromLocalWhenJobEnd()) { + paths.forEach(p -> p.toFile().delete()); + } else { + paths.forEach( + p -> { + File intermediateFile = new File(p.toString()); + if (intermediateFile.exists()) { + logger.info("embulk-output-bigquery: keep {}", p.toString()); + } + }); + } + } + } } return CONFIG_MAPPER_FACTORY.newConfigDiff(); diff --git a/src/test/java/org/embulk/output/bigquery_java/TestBigqueryJavaOutputPluginWithMockServer.java b/src/test/java/org/embulk/output/bigquery_java/TestBigqueryJavaOutputPluginWithMockServer.java index 8fffeb4..2a81071 100644 --- a/src/test/java/org/embulk/output/bigquery_java/TestBigqueryJavaOutputPluginWithMockServer.java +++ b/src/test/java/org/embulk/output/bigquery_java/TestBigqueryJavaOutputPluginWithMockServer.java @@ -34,6 +34,7 @@ import static org.junit.Assert.assertTrue; import java.util.Arrays; +import java.util.Collections; import java.util.List; import java.util.function.Function; import okhttp3.mockwebserver.MockResponse; @@ -140,8 +141,8 @@ public void testRunAppendDirectMode() throws Exception { @Test public void testRunAppendMode() throws Exception { // GET dataset, POST temp table, load into temp table (POST job + GET status), GET temp table - // row count (getTransactionReport), copy temp to final table (POST job + GET status), DELETE - // temp table, GET table (updateTableIfNeed). + // row count (getTransactionReport), copy temp to final table (POST job + GET status), GET + // table (updateTableIfNeed), DELETE temp table. List requests = runWithMockServer( c -> c.set("mode", "append"), @@ -152,8 +153,8 @@ public void testRunAppendMode() throws Exception { tableResponseWithNumRows(1), createCopyJobResponse("testjob"), waitForCopyJobResponse("testjob"), - deleteResponse(), - tableResponse()); + tableResponse(), + deleteResponse()); assertEquals(9, requests.size()); @@ -178,9 +179,9 @@ public void testRunAppendMode() throws Exception { assertGetJobStatus(requests.get(6), "testjob"); - assertDeleteTable(requests.get(7), tempTableId); + assertGetTable(requests.get(7), "table"); // updateTableIfNeed() - assertGetTable(requests.get(8), "table"); + assertDeleteTable(requests.get(8), tempTableId); } @Test @@ -195,8 +196,8 @@ public void testRunReplaceMode() throws Exception { tableResponseWithNumRows(1), createCopyJobResponse("testjob"), waitForCopyJobResponse("testjob"), - deleteResponse(), - tableResponse()); + tableResponse(), + deleteResponse()); assertEquals(9, requests.size()); @@ -221,17 +222,17 @@ public void testRunReplaceMode() throws Exception { assertGetJobStatus(requests.get(6), "testjob"); - assertDeleteTable(requests.get(7), tempTableId); + assertGetTable(requests.get(7), "table"); // updateTableIfNeed() - assertGetTable(requests.get(8), "table"); + assertDeleteTable(requests.get(8), tempTableId); } @Test public void testRunReplaceModeRestoresRetainedDescriptionAndPolicyTags() throws Exception { // With retain_column_descriptions/retain_column_policy_tags on, isNeedUpdateTable() is true: // storeCachedSrcFieldsIfNeed() GETs the (pre-existing) destination table right after autoCreate - // creates the temp table, and updateTableIfNeed() PATCHes the destination afterward to restore - // the cached description/policy tag onto its post-replace schema. + // creates the temp table, and updateTableIfNeed() PATCHes the destination afterward (before the + // temp table delete) to restore the cached description/policy tag onto its post-replace schema. List requests = runWithMockServer( c -> @@ -246,9 +247,9 @@ public void testRunReplaceModeRestoresRetainedDescriptionAndPolicyTags() throws tableResponseWithNumRows(1), createCopyJobResponse("testjob"), waitForCopyJobResponse("testjob"), - deleteResponse(), tableResponseWithNullableC0(), - tableResponse()); + tableResponse(), + deleteResponse()); assertEquals(11, requests.size()); @@ -275,22 +276,22 @@ public void testRunReplaceModeRestoresRetainedDescriptionAndPolicyTags() throws assertGetJobStatus(requests.get(7), "testjob"); - assertDeleteTable(requests.get(8), tempTableId); - - assertGetTable(requests.get(9), "table"); // updateTableIfNeed() + assertGetTable(requests.get(8), "table"); // updateTableIfNeed() - RecordedRequest patchRequest = requests.get(10); + RecordedRequest patchRequest = requests.get(9); assertPatchTable(patchRequest, "table"); assertFieldDescriptionAndPolicyTag( firstSchemaField(requestBodyJson(patchRequest)), "old-description", "old-policy-tag"); + + assertDeleteTable(requests.get(10), tempTableId); } @Test public void testRunReplaceModeGivesUpAfterMaxLoadRetries() throws Exception { // With retries=1, load()'s own RetryExecutor allows only 2 attempts (the initial attempt plus // one retry) before giving up; two consecutive internalError job failures exhaust that budget - // and the whole plugin run fails, so autoCreate's temp table create is the only other request. - // TODO: the temp table should always be deleted, even when the run fails here. + // and the whole plugin run fails. The temp table delete is guaranteed via finally, so it still + // runs even though the load itself failed. List requests = runWithMockServerExpectingFailure( c -> c.set("mode", "replace").set("retries", 1), @@ -301,13 +302,15 @@ public void testRunReplaceModeGivesUpAfterMaxLoadRetries() throws Exception { jobResponse("load-job-1", "load", LOAD_CONFIG_BODY, "RUNNING", null), jobResponse("load-job-1", "load", LOAD_CONFIG_BODY, "DONE", "internalError"), jobResponse("load-job-2", "load", LOAD_CONFIG_BODY, "RUNNING", null), - jobResponse("load-job-2", "load", LOAD_CONFIG_BODY, "DONE", "internalError")); + jobResponse("load-job-2", "load", LOAD_CONFIG_BODY, "DONE", "internalError"), + deleteResponse()); - assertEquals(6, requests.size()); + assertEquals(7, requests.size()); assertGetDataset(requests.get(0)); assertPostTables(requests.get(1)); + String tempTableId = tableIdOf(requestBodyJson(requests.get(1)), "tableReference"); assertPostJobs(requests.get(2)); @@ -316,16 +319,15 @@ public void testRunReplaceModeGivesUpAfterMaxLoadRetries() throws Exception { assertPostJobs(requests.get(4)); assertGetJobStatus(requests.get(5), "load-job-2"); + + assertDeleteTable(requests.get(6), tempTableId); } @Test - public void testRunReplaceModeSkipsSchemaUpdateWhenTempTableDeleteFails() throws Exception { - // Ruby always updates the schema first and deletes the temp table afterward, so the temp - // table is always cleaned up. Java deletes the temp table first and only then calls - // updateTableIfNeed(), so when the delete fails, the run aborts before the schema update is - // ever attempted, even though retain_column_descriptions is on here. - // TODO: update the schema before deleting the temp table, like ruby does, so a delete - // failure doesn't also block the schema update. + public void testRunReplaceModeStillUpdatesSchemaWhenTempTableDeleteFails() throws Exception { + // Like ruby, the schema update now runs before the temp table delete, and the delete is + // guaranteed via finally, so a delete failure no longer blocks the schema update from being + // attempted (it still surfaces the delete failure as the run's exception, though). List requests = runWithMockServerExpectingFailure( c -> c.set("mode", "replace").set("retain_column_descriptions", true), @@ -339,9 +341,10 @@ public void testRunReplaceModeSkipsSchemaUpdateWhenTempTableDeleteFails() throws tableResponseWithNumRows(1), createCopyJobResponse("testjob"), waitForCopyJobResponse("testjob"), - errorResponse(400, "boom", "invalid")); + tableResponse(), // updateTableIfNeed(), no schema so it returns before patching + errorResponse(400, "boom", "invalid")); // delete temp table (finally) - assertEquals(9, requests.size()); + assertEquals(10, requests.size()); assertGetDataset(requests.get(0)); @@ -360,7 +363,69 @@ public void testRunReplaceModeSkipsSchemaUpdateWhenTempTableDeleteFails() throws assertGetJobStatus(requests.get(7), "testjob"); - assertDeleteTable(requests.get(8), tempTableId); + assertGetTable(requests.get(8), "table"); // updateTableIfNeed() + + assertDeleteTable(requests.get(9), tempTableId); + } + + @Test + public void testRunReplaceModeStillDeletesTempTableWhenStoreCachedSrcFieldsIfNeedFails() + throws Exception { + // storeCachedSrcFieldsIfNeed() runs first inside the (now widened) try block, right after + // autoCreate() creates the temp table. Even when it fails outright, the temp table delete in + // finally still runs, guaranteeing cleanup. + List requests = + runWithMockServerExpectingFailure( + c -> c.set("mode", "replace").set("retain_column_descriptions", true), + RuntimeException.class, + "(?s).*boom.*", + datasetResponse(), + tableResponse(), + errorResponse(400, "boom", "invalid"), // storeCachedSrcFieldsIfNeed() + deleteResponse()); + + assertEquals(4, requests.size()); + + assertGetDataset(requests.get(0)); + + assertPostTables(requests.get(1)); + String tempTableId = tableIdOf(requestBodyJson(requests.get(1)), "tableReference"); + + assertGetTable(requests.get(2), "table"); // storeCachedSrcFieldsIfNeed(), fails with "boom" + + assertDeleteTable(requests.get(3), tempTableId); + } + + @Test + public void testRunReplaceModeDeletesTempTableWhenPathsAreEmpty() throws Exception { + // With zero input records, BigqueryPageOutput#add() is never called, so no writer is ever + // registered and no intermediate file gets created: paths.isEmpty() is true, and transaction() + // returns right after creating the destination table, before ever reaching the load/copy + // logic. The temp table (already created by autoCreate()) must still be deleted via the + // guaranteed finally block. + List requests = + BigqueryMockWebServerTestUtil.runWithMockServer( + embulk, + testFolder, + loadTestHostConfig(embulk, c -> c.set("mode", "replace")), + Collections.singletonList("c0:string"), + datasetResponse(), + tableResponse(), + tableResponse(), + deleteResponse()); + + assertEquals(4, requests.size()); + + assertGetDataset(requests.get(0)); + + assertPostTables(requests.get(1)); + String tempTableId = tableIdOf(requestBodyJson(requests.get(1)), "tableReference"); + assertMatches(tempTableId, "LOAD_TEMP_.*_table"); + + assertPostTables(requests.get(2)); + assertEquals("table", tableIdOf(requestBodyJson(requests.get(2)), "tableReference")); + + assertDeleteTable(requests.get(3), tempTableId); } @Test @@ -379,8 +444,8 @@ public void testRunDeleteInAdvanceMode() throws Exception { tableResponseWithNumRows(1), createCopyJobResponse("testjob"), waitForCopyJobResponse("testjob"), - deleteResponse(), - tableResponse()); + tableResponse(), + deleteResponse()); assertEquals(10, requests.size()); @@ -407,9 +472,9 @@ public void testRunDeleteInAdvanceMode() throws Exception { assertGetJobStatus(requests.get(7), "testjob"); - assertDeleteTable(requests.get(8), tempTableId); + assertGetTable(requests.get(8), "table"); // updateTableIfNeed() - assertGetTable(requests.get(9), "table"); + assertDeleteTable(requests.get(9), tempTableId); } @Test @@ -430,8 +495,8 @@ public void testRunMergeMode() throws Exception { tableResponseWithNumRows(1), createQueryJobResponse("testjob"), waitForQueryJobResponse("testjob"), - deleteResponse(), - tableResponse()); + tableResponse(), + deleteResponse()); assertEquals(10, requests.size()); @@ -459,8 +524,8 @@ public void testRunMergeMode() throws Exception { assertGetJobStatus(requests.get(7), "testjob"); - assertDeleteTable(requests.get(8), tempTableId); + assertGetTable(requests.get(8), "table"); // updateTableIfNeed() - assertGetTable(requests.get(9), "table"); + assertDeleteTable(requests.get(9), tempTableId); } }