Proper close

airbytehq · Oct 10, 2023 · 2671973 · 2671973
1 parent 3be9554
commit 2671973
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 3 deletions.
diff --git a/...query/src/main/java/io/airbyte/integrations/destination/bigquery/BigQueryDestination.java b/...query/src/main/java/io/airbyte/integrations/destination/bigquery/BigQueryDestination.java
@@ -391,9 +391,7 @@ private SerializedAirbyteMessageConsumer getStandardRecordConsumer(final BigQuer
         },
         (hasFailed) -> {
           try {
-            writeConfigs.forEach((streamId, uploader) -> {
-              uploader.closeAfterPush();
-            });
+            writeConfigs.forEach((streamId, uploader) -> uploader.closeWithoutState(hasFailed));
             Thread.sleep(30 * 1000); // 30 seconds
             typerDeduper.typeAndDedupe();
             typerDeduper.commitFinalTables();

diff --git a/.../java/io/airbyte/integrations/destination/bigquery/uploader/AbstractBigQueryUploader.java b/.../java/io/airbyte/integrations/destination/bigquery/uploader/AbstractBigQueryUploader.java
@@ -107,6 +107,20 @@ public void close(final boolean hasFailed, final Consumer<AirbyteMessage> output
     }
   }
 
+  public void closeWithoutState(final boolean hasFailed) {
+    try {
+      recordFormatter.printAndCleanFieldFails();
+
+      this.writer.close(hasFailed);
+
+      this.postProcessAction(hasFailed);
+    } catch (final Exception e) {
+      LOGGER.error(String.format("Failed to close %s writer, \n details: %s", this, e.getMessage()));
+      printHeapMemoryConsumption();
+      throw new RuntimeException(e);
+    }
+  }
+
   public void closeAfterPush() {
     try {
       this.writer.close(false);