[GOBBLIN-2159] Adding support for partition level copy in Iceberg distcp

gobblin-data-management/src/main/java/org/apache/gobblin/data/management/copy/iceberg/BaseIcebergCatalog.java

@@ -21,10 +21,12 @@
 
 import org.apache.hadoop.conf.Configuration;
 import org.apache.iceberg.CatalogUtil;
+import org.apache.iceberg.Table;
 import org.apache.iceberg.TableOperations;
 import org.apache.iceberg.catalog.Catalog;
 import org.apache.iceberg.catalog.TableIdentifier;
 import org.apache.gobblin.dataset.DatasetConstants;
+import org.apache.iceberg.exceptions.NoSuchTableException;
 
 /**
  * Base implementation of {@link IcebergCatalog} to access {@link IcebergTable} and the
@@ -41,9 +43,15 @@ protected BaseIcebergCatalog(String catalogName, Class<? extends Catalog> compan
   }
 
   @Override
-  public IcebergTable openTable(String dbName, String tableName) {
+  public IcebergTable openTable(String dbName, String tableName) throws IcebergTable.TableNotFoundException {
     TableIdentifier tableId = TableIdentifier.of(dbName, tableName);
-    return new IcebergTable(tableId, calcDatasetDescriptorName(tableId), getDatasetDescriptorPlatform(), createTableOperations(tableId), this.getCatalogUri());
+    try {
+      return new IcebergTable(tableId, calcDatasetDescriptorName(tableId), getDatasetDescriptorPlatform(),
+          createTableOperations(tableId), this.getCatalogUri(), loadTableInstance(tableId));
+    } catch (NoSuchTableException ex) {
+      // defend against `org.apache.iceberg.catalog.Catalog::loadTable` throwing inside some `@Override` of `loadTableInstance`
+      throw new IcebergTable.TableNotFoundException(tableId);
+    }
   }
 
   protected Catalog createCompanionCatalog(Map<String, String> properties, Configuration configuration) {
@@ -67,4 +75,6 @@ protected String getDatasetDescriptorPlatform() {
   }
 
   protected abstract TableOperations createTableOperations(TableIdentifier tableId);
+
+  protected abstract Table loadTableInstance(TableIdentifier tableId);
 }

gobblin-data-management/src/main/java/org/apache/gobblin/data/management/copy/iceberg/IcebergCatalog.java

@@ -29,10 +29,10 @@
 public interface IcebergCatalog {
 
   /** @return table identified by `dbName` and `tableName` */
-  IcebergTable openTable(String dbName, String tableName);
+  IcebergTable openTable(String dbName, String tableName) throws IcebergTable.TableNotFoundException;
 
   /** @return table identified by `tableId` */
-  default IcebergTable openTable(TableIdentifier tableId) {
+  default IcebergTable openTable(TableIdentifier tableId) throws IcebergTable.TableNotFoundException {
     // CHALLENGE: clearly better to implement in the reverse direction - `openTable(String, String)` in terms of `openTable(TableIdentifier)` -
     // but challenging to do at this point, with multiple derived classes already "in the wild" that implement `openTable(String, String)`
     return openTable(tableId.namespace().toString(), tableId.name());

gobblin-data-management/src/main/java/org/apache/gobblin/data/management/copy/iceberg/IcebergDatasetFinder.java

@@ -85,7 +85,7 @@ public String getConfigPrefix() {
   }
 
   protected final FileSystem sourceFs;
-  private final Properties properties;
+  protected final Properties properties;
 
   /**
    * Finds all {@link IcebergDataset}s in the file system using the Iceberg Catalog.
@@ -153,7 +153,7 @@ protected IcebergDataset createIcebergDataset(IcebergCatalog sourceIcebergCatalo
     IcebergTable destIcebergTable = destinationIcebergCatalog.openTable(destDbName, destTableName);
     // TODO: Rethink strategy to enforce dest iceberg table
     Preconditions.checkArgument(destinationIcebergCatalog.tableAlreadyExists(destIcebergTable), String.format("Missing Destination Iceberg Table: {%s}.{%s}", destDbName, destTableName));
-    return new IcebergDataset(srcIcebergTable, destIcebergTable, properties, fs, getConfigShouldCopyMetadataPath(properties));
+    return createSpecificDataset(srcIcebergTable, destIcebergTable, properties, fs, getConfigShouldCopyMetadataPath(properties));
   }
 
   protected static IcebergCatalog createIcebergCatalog(Properties properties, CatalogLocation location) throws IOException {
@@ -165,6 +165,11 @@ protected static IcebergCatalog createIcebergCatalog(Properties properties, Cata
     return IcebergCatalogFactory.create(icebergCatalogClassName, catalogProperties, configuration);
   }
 
+  protected IcebergDataset createSpecificDataset(IcebergTable srcIcebergTable, IcebergTable destIcebergTable, Properties properties, FileSystem fs, boolean shouldIncludeMetadataPath)
+      throws IOException {
+    return new IcebergDataset(srcIcebergTable, destIcebergTable, properties, fs, shouldIncludeMetadataPath);
+  }
+
   protected static boolean getConfigShouldCopyMetadataPath(Properties properties) {
     return Boolean.valueOf(properties.getProperty(ICEBERG_DATASET_SHOULD_COPY_METADATA_PATH, DEFAULT_ICEBERG_DATASET_SHOULD_COPY_METADATA_PATH));
   }

gobblin-data-management/src/main/java/org/apache/gobblin/data/management/copy/iceberg/IcebergHiveCatalog.java

@@ -21,6 +21,7 @@
 
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.hive.conf.HiveConf;
+import org.apache.iceberg.Table;
 import org.apache.iceberg.TableOperations;
 import org.apache.iceberg.catalog.TableIdentifier;
 import org.apache.iceberg.hive.HiveCatalog;
@@ -61,4 +62,9 @@ protected TableOperations createTableOperations(TableIdentifier tableId) {
   public boolean tableAlreadyExists(IcebergTable icebergTable) {
     return hc.tableExists(icebergTable.getTableId());
   }
+
+  @Override
+  protected Table loadTableInstance(TableIdentifier tableId) {
+    return hc.loadTable(tableId);
+  }
 }

gobblin-data-management/src/main/java/org/apache/gobblin/data/management/copy/iceberg/IcebergOverwritePartitionsStep.java

@@ -0,0 +1,164 @@
+/*
+ * 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.gobblin.data.management.copy.iceberg;
+
+import java.io.IOException;
+import java.time.Duration;
+import java.util.List;
+import java.util.Optional;
+import java.util.Properties;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.ExecutionException;
+
+import org.apache.iceberg.DataFile;
+import org.apache.iceberg.catalog.TableIdentifier;
+import org.apache.iceberg.util.SerializationUtil;
+
+import com.github.rholder.retry.Attempt;
+import com.github.rholder.retry.RetryException;
+import com.github.rholder.retry.RetryListener;
+import com.github.rholder.retry.Retryer;
+import com.google.common.collect.ImmutableMap;
+import com.typesafe.config.Config;
+import com.typesafe.config.ConfigFactory;
+
+import lombok.extern.slf4j.Slf4j;
+
+import org.apache.gobblin.commit.CommitStep;
+import org.apache.gobblin.util.retry.RetryerFactory;
+
+import static org.apache.gobblin.util.retry.RetryerFactory.RETRY_INTERVAL_MS;
+import static org.apache.gobblin.util.retry.RetryerFactory.RETRY_TIMES;
+import static org.apache.gobblin.util.retry.RetryerFactory.RETRY_TYPE;
+import static org.apache.gobblin.util.retry.RetryerFactory.RetryType;
+
+/**
+ * Commit step for overwriting partitions in an Iceberg table.
+ * <p>
+ * This class implements the {@link CommitStep} interface and provides functionality to overwrite
+ * partitions in the destination Iceberg table using serialized data files.
+ * </p>
+ */
+@Slf4j
+public class IcebergOverwritePartitionsStep implements CommitStep {
+  private final String destTableIdStr;
+  private final Properties properties;
+  private final byte[] serializedDataFiles;
+  private final String partitionColName;
+  private final String partitionValue;
+  public static final String OVERWRITE_PARTITIONS_RETRYER_CONFIG_PREFIX = IcebergDatasetFinder.ICEBERG_DATASET_PREFIX +
+      ".catalog.overwrite.partitions.retries";
+  private static final Config RETRYER_FALLBACK_CONFIG = ConfigFactory.parseMap(ImmutableMap.of(
+      RETRY_INTERVAL_MS, TimeUnit.SECONDS.toMillis(3L),
+      RETRY_TIMES, 3,
+      RETRY_TYPE, RetryType.FIXED_ATTEMPT.name()));
+
+  /**
+   * Constructs an {@code IcebergReplacePartitionsStep} with the specified parameters.
+   *
+   * @param destTableIdStr the identifier of the destination table as a string
+   * @param serializedDataFiles [from List<DataFiles>] the serialized data files to be used for replacing partitions
+   * @param properties the properties containing configuration
+   */
+  public IcebergOverwritePartitionsStep(String destTableIdStr, String partitionColName, String partitionValue, byte[] serializedDataFiles, Properties properties) {
+    this.destTableIdStr = destTableIdStr;
+    this.partitionColName = partitionColName;
+    this.partitionValue = partitionValue;
+    this.serializedDataFiles = serializedDataFiles;
+    this.properties = properties;
+  }
+
+  @Override
+  public boolean isCompleted() {
+    return false;
+  }
+
+  /**
+   * Executes the partition replacement in the destination Iceberg table.
+   * Also, have retry mechanism as done in {@link IcebergRegisterStep#execute()}
+   *
+   * @throws IOException if an I/O error occurs during execution
+   */
+  @Override
+  public void execute() throws IOException {
+    // Unlike IcebergRegisterStep::execute, which validates dest table metadata has not changed between copy entity
+    // generation and the post-copy commit, do no such validation here, so dest table writes may continue throughout
+    // our copying. any new data written in the meanwhile to THE SAME partitions we are about to overwrite will be
+    // clobbered and replaced by the copy entities from our execution.
+    IcebergTable destTable = createDestinationCatalog().openTable(TableIdentifier.parse(this.destTableIdStr));
+    List<DataFile> dataFiles = SerializationUtil.deserializeFromBytes(this.serializedDataFiles);
+    try {
+      log.info("~{}~ Starting partition overwrite - partition: {}; value: {}; numDataFiles: {}; path[0]: {}",
+          this.destTableIdStr,
+          this.partitionColName,
+          this.partitionValue,
+          dataFiles.size(),
+          dataFiles.get(0).path()
+      );
+      Retryer<Void> overwritePartitionsRetryer = createOverwritePartitionsRetryer();
+      overwritePartitionsRetryer.call(() -> {
+        destTable.overwritePartition(dataFiles, this.partitionColName, this.partitionValue);
+        return null;
+      });
+      log.info("~{}~ Successful partition overwrite - partition: {}; value: {}",
+          this.destTableIdStr,
+          this.partitionColName,
+          this.partitionValue
+      );
+    } catch (ExecutionException executionException) {
+      String msg = String.format("~%s~ Failed to overwrite partitions", this.destTableIdStr);
+      log.error(msg, executionException);
+      throw new RuntimeException(msg, executionException.getCause());
+    } catch (RetryException retryException) {
+      String interruptedNote = Thread.currentThread().isInterrupted() ? "... then interrupted" : "";
+      String msg = String.format("~%s~ Failure attempting to overwrite partition [num failures: %d] %s",
+          this.destTableIdStr,
+          retryException.getNumberOfFailedAttempts(),
+          interruptedNote);
+      Throwable informativeException = retryException.getLastFailedAttempt().hasException()
+          ? retryException.getLastFailedAttempt().getExceptionCause()
+          : retryException;
+      log.error(msg, informativeException);
+      throw new RuntimeException(msg, informativeException);
+    }
+  }
+
+  protected IcebergCatalog createDestinationCatalog() throws IOException {
+    return IcebergDatasetFinder.createIcebergCatalog(this.properties, IcebergDatasetFinder.CatalogLocation.DESTINATION);
+  }
+
+  private Retryer<Void> createOverwritePartitionsRetryer() {
+    Config config = ConfigFactory.parseProperties(this.properties);
+    Config retryerOverridesConfig = config.hasPath(IcebergOverwritePartitionsStep.OVERWRITE_PARTITIONS_RETRYER_CONFIG_PREFIX)
+        ? config.getConfig(IcebergOverwritePartitionsStep.OVERWRITE_PARTITIONS_RETRYER_CONFIG_PREFIX)
+        : ConfigFactory.empty();
+
+    return RetryerFactory.newInstance(retryerOverridesConfig.withFallback(RETRYER_FALLBACK_CONFIG), Optional.of(new RetryListener() {
+      @Override
+      public <V> void onRetry(Attempt<V> attempt) {
+        if (attempt.hasException()) {
+          String msg = String.format("~%s~ Exception while overwriting partitions [attempt: %d; elapsed: %s]",
+              destTableIdStr,
+              attempt.getAttemptNumber(),
+              Duration.ofMillis(attempt.getDelaySinceFirstAttempt()).toString());
+          log.warn(msg, attempt.getExceptionCause());
+        }
+      }
+    }));
+  }
+}