BigQuery: Add retry detection for idempotent commit handling

bigquery/src/main/java/org/apache/iceberg/gcp/bigquery/BigQueryMetastoreClient.java

@@ -27,6 +27,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import org.apache.iceberg.gcp.bigquery.util.RetryDetector;
 
 /**
  * A client of Google BigQuery Metastore functions over the BigQuery service. Uses the Google
@@ -94,6 +95,14 @@ interface BigQueryMetastoreClient {
    */
   Table create(Table table);
 
+  /**
+   * Creates and returns a new table.
+   *
+   * @param table body of the table to create
+   * @param retryDetector detector to determine if a retry is needed
+   */
+  Table create(Table table, RetryDetector retryDetector);
+
   /**
    * Returns a table.
    *
@@ -109,6 +118,15 @@ interface BigQueryMetastoreClient {
    */
   Table update(TableReference tableReference, Table table);
 
+  /**
+   * Updates the catalog table options of an Iceberg table and returns the updated table.
+   *
+   * @param tableReference full table reference
+   * @param table to patch
+   * @param retryDetector detector to determine if a retry is needed
+   */
+  Table update(TableReference tableReference, Table table, RetryDetector retryDetector);
+
   /**
    * Deletes a table.
    *

bigquery/src/main/java/org/apache/iceberg/gcp/bigquery/BigQueryMetastoreClientImpl.java

@@ -70,6 +70,7 @@
 import org.apache.iceberg.exceptions.RuntimeIOException;
 import org.apache.iceberg.exceptions.ServiceFailureException;
 import org.apache.iceberg.exceptions.ServiceUnavailableException;
+import org.apache.iceberg.gcp.bigquery.util.RetryDetector;
 import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
 import org.apache.iceberg.relocated.com.google.common.collect.Lists;
 import org.apache.iceberg.relocated.com.google.common.collect.Maps;
@@ -325,14 +326,18 @@ public List<Datasets> list(String projectId) {
 
   @Override
   public Table create(Table table) {
+    return create(table, new RetryDetector());
+  }
+
+  @Override
+  public Table create(Table table, RetryDetector retryDetector) {
     // Ensure it is an Iceberg table supported by the BigQuery metastore catalog.
     validateTable(table);
-    // TODO: Ensure table creation is idempotent when handling retries.
     Table response = null;
     try {
       response =
           BigQueryRetryHelper.runWithRetries(
-              () -> internalCreate(table),
+              retryDetector.wrap(() -> internalCreate(table)),
               bigqueryOptions.getRetrySettings(),
               BIGQUERY_EXCEPTION_HANDLER,
               bigqueryOptions.getClock(),
@@ -387,6 +392,11 @@ public Table load(TableReference tableReference) {
 
   @Override
   public Table update(TableReference tableReference, Table table) {
+    return update(tableReference, table, new RetryDetector());
+  }
+
+  @Override
+  public Table update(TableReference tableReference, Table table, RetryDetector retryDetector) {
     // Ensure it is an Iceberg table supported by the BQ metastore catalog.
     validateTable(table);
 
@@ -405,7 +415,8 @@ public Table update(TableReference tableReference, Table table) {
     try {
       response =
           BigQueryRetryHelper.runWithRetries(
-              () -> internalUpdate(tableReference, updatedTable, table.getEtag()),
+              retryDetector.wrap(
+                  () -> internalUpdate(tableReference, updatedTable, table.getEtag())),
               bigqueryOptions.getRetrySettings(),
               BIGQUERY_EXCEPTION_HANDLER,
               bigqueryOptions.getClock(),

bigquery/src/main/java/org/apache/iceberg/gcp/bigquery/BigQueryTableOperations.java

@@ -22,7 +22,6 @@
 import com.google.api.services.bigquery.model.Table;
 import com.google.api.services.bigquery.model.TableReference;
 import java.util.Map;
-import org.apache.iceberg.BaseMetastoreOperations;
 import org.apache.iceberg.BaseMetastoreTableOperations;
 import org.apache.iceberg.SnapshotSummary;
 import org.apache.iceberg.TableMetadata;
@@ -31,8 +30,11 @@
 import org.apache.iceberg.exceptions.CommitFailedException;
 import org.apache.iceberg.exceptions.CommitStateUnknownException;
 import org.apache.iceberg.exceptions.NoSuchTableException;
+import org.apache.iceberg.exceptions.RuntimeIOException;
 import org.apache.iceberg.exceptions.ValidationException;
+import org.apache.iceberg.gcp.bigquery.util.RetryDetector;
 import org.apache.iceberg.io.FileIO;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
 import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
 import org.apache.iceberg.relocated.com.google.common.collect.Maps;
 import org.slf4j.Logger;
@@ -82,48 +84,45 @@ public void doRefresh() {
   // atomically
   @Override
   public void doCommit(TableMetadata base, TableMetadata metadata) {
-    String newMetadataLocation =
-        base == null && metadata.metadataFileLocation() != null
-            ? metadata.metadataFileLocation()
-            : writeNewMetadata(metadata, currentVersion() + 1);
-    BaseMetastoreOperations.CommitStatus commitStatus =
-        BaseMetastoreOperations.CommitStatus.FAILURE;
+    CommitStatus commitStatus = CommitStatus.FAILURE;
+    RetryDetector retryDetector = new RetryDetector();
+
+    String newMetadataLocation = null;
     try {
-      if (base == null) {
-        createTable(newMetadataLocation, metadata);
+      boolean newTable = base == null;
+      newMetadataLocation = writeNewMetadataIfRequired(newTable, metadata);
+
+      if (newTable) {
+        createTable(newMetadataLocation, metadata, retryDetector);
       } else {
-        updateTable(newMetadataLocation, metadata);
+        updateTable(newMetadataLocation, metadata, retryDetector);
       }
-      commitStatus = BaseMetastoreOperations.CommitStatus.SUCCESS;
-    } catch (CommitFailedException | CommitStateUnknownException e) {
+      commitStatus = CommitStatus.SUCCESS;
+    } catch (CommitFailedException | AlreadyExistsException e) {
       throw e;
     } catch (Throwable e) {
       LOG.error("Exception thrown on commit: ", e);
-      if (e instanceof AlreadyExistsException) {
-        throw e;
+      boolean isRuntimeIOException = e instanceof RuntimeIOException;
+
+      // If retries occurred, an earlier attempt may have succeeded. If we got a
+      // RuntimeIOException, we have no way of knowing if the request reached the server.
+      // In either case, check whether the commit actually succeeded.
+      if (isRuntimeIOException || retryDetector.retried()) {
+        LOG.warn(
+            "Received unexpected failure when committing to {}, validating if commit ended up succeeding.",
+            tableName(),
+            e);
+        commitStatus = checkCommitStatus(newMetadataLocation, metadata);
       }
-      commitStatus =
-          BaseMetastoreOperations.CommitStatus.valueOf(
-              checkCommitStatus(newMetadataLocation, metadata).name());
-      if (commitStatus == BaseMetastoreOperations.CommitStatus.FAILURE) {
+
+      if (commitStatus == CommitStatus.FAILURE) {
         throw new CommitFailedException(e, "Failed to commit");
       }
-      if (commitStatus == BaseMetastoreOperations.CommitStatus.UNKNOWN) {
+      if (commitStatus == CommitStatus.UNKNOWN) {
         throw new CommitStateUnknownException(e);
       }
     } finally {
-      try {
-        if (commitStatus == BaseMetastoreOperations.CommitStatus.FAILURE) {
-          LOG.warn("Failed to commit updates to table {}", tableName());
-          io().deleteFile(newMetadataLocation);
-        }
-      } catch (RuntimeException e) {
-        LOG.error(
-            "Failed to cleanup metadata file at {} for table {}",
-            newMetadataLocation,
-            tableName(),
-            e);
-      }
+      cleanupMetadata(commitStatus, newMetadataLocation);
     }
   }
 
@@ -137,13 +136,14 @@ public FileIO io() {
     return fileIO;
   }
 
-  private void createTable(String newMetadataLocation, TableMetadata metadata) {
+  private void createTable(
+      String newMetadataLocation, TableMetadata metadata, RetryDetector retryDetector) {
     LOG.debug("Creating a new Iceberg table: {}", tableName());
     Table tableBuilder = makeNewTable(metadata, newMetadataLocation);
     tableBuilder.setTableReference(tableReference);
     addConnectionIfProvided(tableBuilder, metadata.properties());
 
-    client.create(tableBuilder);
+    client.create(tableBuilder, retryDetector);
   }
 
   private void addConnectionIfProvided(Table tableBuilder, Map<String, String> metadataProperties) {
@@ -155,7 +155,8 @@ private void addConnectionIfProvided(Table tableBuilder, Map<String, String> met
   }
 
   /** Update table properties with concurrent update detection using etag. */
-  private void updateTable(String newMetadataLocation, TableMetadata metadata) {
+  private void updateTable(
+      String newMetadataLocation, TableMetadata metadata, RetryDetector retryDetector) {
     Preconditions.checkState(
         metastoreTable != null,
         "Table %s must be loaded during refresh before commit",
@@ -171,7 +172,7 @@ private void updateTable(String newMetadataLocation, TableMetadata metadata) {
     addConnectionIfProvided(metastoreTable, metadata.properties());
 
     options.setParameters(buildTableParameters(newMetadataLocation, metadata));
-    client.update(tableReference, metastoreTable);
+    client.update(tableReference, metastoreTable, retryDetector);
     this.metastoreTable = null;
   }
 
@@ -241,4 +242,19 @@ private String loadMetadataLocationOrThrow(ExternalCatalogTableOptions tableOpti
 
     return tableOptions.getParameters().get(METADATA_LOCATION_PROP);
   }
+
+  @VisibleForTesting
+  void cleanupMetadata(CommitStatus commitStatus, String metadataLocation) {
+    try {
+      if (commitStatus == CommitStatus.FAILURE
+          && metadataLocation != null
+          && !metadataLocation.isEmpty()) {
+        // if anything went wrong, clean up the uncommitted metadata file
+        io().deleteFile(metadataLocation);
+      }
+    } catch (RuntimeException e) {
+      LOG.error(
+          "Failed to cleanup metadata file at {} for table {}", metadataLocation, tableName(), e);
+    }
+  }
 }

bigquery/src/main/java/org/apache/iceberg/gcp/bigquery/util/RetryDetector.java

@@ -0,0 +1,42 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.gcp.bigquery.util;
+
+import java.util.concurrent.Callable;
+import java.util.concurrent.atomic.AtomicInteger;
+
+/** Detects retries by counting callable invocations. */
+public class RetryDetector {
+  private final AtomicInteger attempts = new AtomicInteger(0);
+
+  public <T> Callable<T> wrap(Callable<T> delegate) {
+    return () -> {
+      attempts.getAndIncrement();
+      return delegate.call();
+    };
+  }
+
+  public boolean retried() {
+    return attempts.get() > 1;
+  }
+
+  public int attempts() {
+    return attempts.get();
+  }
+}

bigquery/src/test/java/org/apache/iceberg/gcp/bigquery/FakeBigQueryMetastoreClient.java

@@ -36,6 +36,7 @@
 import org.apache.iceberg.exceptions.NamespaceNotEmptyException;
 import org.apache.iceberg.exceptions.NoSuchNamespaceException;
 import org.apache.iceberg.exceptions.NoSuchTableException;
+import org.apache.iceberg.gcp.bigquery.util.RetryDetector;
 import org.apache.iceberg.relocated.com.google.common.collect.Maps;
 
 public class FakeBigQueryMetastoreClient implements BigQueryMetastoreClient {
@@ -149,6 +150,11 @@ public List<DatasetList.Datasets> list(String projectId) {
         .collect(Collectors.toList());
   }
 
+  @Override
+  public Table create(Table table, RetryDetector retryDetector) {
+    return create(table);
+  }
+
   @Override
   public Table create(Table table) {
     if (tables.containsKey(table.getTableReference())) {
@@ -170,6 +176,11 @@ public Table load(TableReference tableReference) {
     return table;
   }
 
+  @Override
+  public Table update(TableReference tableReference, Table table, RetryDetector retryDetector) {
+    return update(tableReference, table);
+  }
+
   @Override
   public Table update(TableReference tableReference, Table table) {
     Table existingTable = tables.get(tableReference);