feat: add enableSnappy flag to import snapshot pipeline and select th… (#3586)

* feat: add enableSnappy flag to import snapshot pipeline and select the right container image based on beam version if flag is set.

* 🦉 Updates from OwlBot post-processor

See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md

* Use DataflowServiceOptions for runner_v2 and update snapshot name in test.

* Resolve merge conflict and update E2E test.

* address comments

* fix zip data and generate script.

* 🦉 Updates from OwlBot post-processor

See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md

* upload correct zip file.

* Fix clirr and address comment.

* Filter by job name and add back pagination in snappy snapshot test.

Co-authored-by: Tracy Cui <tracycui@google.com>
Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>

Author: 3 people

bigtable-dataflow-parent/bigtable-beam-import/clirr-ignored-differences.xml

@@ -0,0 +1,9 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!-- see http://www.mojohaus.org/clirr-maven-plugin/examples/ignored-differences.html -->
+<differences>
+    <difference>
+        <differenceType>7012</differenceType>
+        <className>com/google/cloud/bigtable/beam/hbasesnapshots/ImportJobFromHbaseSnapshot$ImportOptions</className>
+        <method>*</method>
+    </difference>
+</differences>

bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/ImportJobFromHbaseSnapshot.java

@@ -22,16 +22,21 @@
 import com.google.cloud.bigtable.beam.sequencefiles.ImportJob;
 import com.google.cloud.bigtable.beam.sequencefiles.Utils;
 import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.MoreObjects;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
+import org.apache.beam.runners.dataflow.options.DataflowPipelineOptions;
 import org.apache.beam.sdk.Pipeline;
 import org.apache.beam.sdk.PipelineResult;
 import org.apache.beam.sdk.io.hadoop.format.HadoopFormatIO;
+import org.apache.beam.sdk.options.Default;
 import org.apache.beam.sdk.options.Description;
 import org.apache.beam.sdk.options.PipelineOptionsFactory;
 import org.apache.beam.sdk.transforms.Create;
 import org.apache.beam.sdk.transforms.ParDo;
 import org.apache.beam.sdk.transforms.Wait;
+import org.apache.beam.sdk.util.ReleaseInfo;
 import org.apache.beam.sdk.values.KV;
 import org.apache.beam.sdk.values.PCollection;
 import org.apache.commons.logging.Log;
@@ -66,6 +71,8 @@
 @InternalExtensionOnly
 public class ImportJobFromHbaseSnapshot {
   private static final Log LOG = LogFactory.getLog(ImportJobFromHbaseSnapshot.class);
+  private static final String CONTAINER_IMAGE_PATH_PREFIX =
+      "gcr.io/cloud-bigtable-ecosystem/unified-harness:";
 
   public interface ImportOptions extends ImportJob.ImportOptions {
     @Description("The HBase root dir where HBase snapshot files resides.")
@@ -79,6 +86,13 @@ public interface ImportOptions extends ImportJob.ImportOptions {
 
     @SuppressWarnings("unused")
     void setSnapshotName(String snapshotName);
+
+    @Description("Is importing Snappy compressed snapshot.")
+    @Default.Boolean(false)
+    Boolean getEnableSnappy();
+
+    @SuppressWarnings("unused")
+    void setEnableSnappy(Boolean enableSnappy);
   }
 
   public static void main(String[] args) throws Exception {
@@ -100,6 +114,20 @@ public static void main(String[] args) throws Exception {
   @VisibleForTesting
   static Pipeline buildPipeline(ImportOptions opts) throws Exception {
 
+    if (opts.getEnableSnappy()) {
+      DataflowPipelineOptions dataFlowOpts = opts.as(DataflowPipelineOptions.class);
+      dataFlowOpts.setSdkContainerImage(
+          CONTAINER_IMAGE_PATH_PREFIX + ReleaseInfo.getReleaseInfo().getVersion());
+
+      List<String> expOpts =
+          MoreObjects.firstNonNull(dataFlowOpts.getExperiments(), new ArrayList());
+      if (!expOpts.contains("use_runner_v2")) {
+        expOpts = new ArrayList<>(expOpts);
+        expOpts.add("use_runner_v2");
+      }
+      dataFlowOpts.setExperiments(expOpts);
+    }
+
     Pipeline pipeline = Pipeline.create(Utils.tweakOptions(opts));
     HBaseSnapshotInputConfigBuilder configurationBuilder =
         new HBaseSnapshotInputConfigBuilder()

bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/EndToEndIT.java

@@ -84,6 +84,7 @@ public class EndToEndIT {
   private static final Logger LOG = LoggerFactory.getLogger(HBaseResultToMutationFn.class);
   private static final String SNAPSHOT_FIXTURE_NAME = "EndToEndIT-snapshot.zip";
   private static final String TEST_SNAPSHOT_NAME = "test-snapshot";
+  private static final String TEST_SNAPPY_SNAPSHOT_NAME = "test-snappy-snapshot";
   private static final String CF = "cf";
 
   private TestProperties properties;
@@ -340,7 +341,7 @@ public void testHBaseSnapshotImport() throws Exception {
 
     // Validate the counters.
     Map<String, Long> counters = getCountMap(result);
-    Assert.assertEquals(counters.get("ranges_matched"), (Long) 101L);
+    Assert.assertEquals(counters.get("ranges_matched"), (Long) 100L);
     Assert.assertEquals(counters.get("ranges_not_matched"), (Long) 0L);
   }
 
@@ -403,7 +404,58 @@ public void testHBaseSnapshotImportWithCorruptions() throws Exception {
 
     // Assert that the output collection is the right one.
     Map<String, Long> counters = getCountMap(result);
-    Assert.assertEquals(counters.get("ranges_matched"), (Long) 97L);
+    Assert.assertEquals(counters.get("ranges_matched"), (Long) 96L);
     Assert.assertEquals(counters.get("ranges_not_matched"), (Long) 4L);
   }
+
+  @Test
+  public void testSnappyCompressedHBaseSnapshotImport() throws Exception {
+    // Start import
+    ImportOptions importOpts = createImportOptions();
+    importOpts.setEnableSnappy(true);
+    importOpts.setSnapshotName(TEST_SNAPPY_SNAPSHOT_NAME);
+
+    // run pipeline
+    State state = ImportJobFromHbaseSnapshot.buildPipeline(importOpts).run().waitUntilFinish();
+    Assert.assertEquals(State.DONE, state);
+
+    // check that the .restore dir used for temp files has been removed
+    // The restore directory is stored relative to the snapshot directory and contains the job name
+    String bucket = GcsPath.fromUri(hbaseSnapshotDir).getBucket();
+    String restorePathPrefix =
+        CleanupHBaseSnapshotRestoreFilesFn.getListPrefix(
+            HBaseSnapshotInputConfigBuilder.RESTORE_DIR);
+
+    List<StorageObject> allObjects = new ArrayList<>();
+    String nextToken;
+    do {
+      Objects objects = gcsUtil.listObjects(bucket, restorePathPrefix, null);
+      List<StorageObject> items = objects.getItems();
+      if (items != null) {
+        allObjects.addAll(items);
+      }
+      nextToken = objects.getNextPageToken();
+    } while (nextToken != null);
+
+    List<StorageObject> myObjects =
+        allObjects.stream()
+            .filter(o -> o.getName().contains(importOpts.getJobName()))
+            .collect(Collectors.toList());
+    Assert.assertTrue("Restore directory wasn't deleted", myObjects.isEmpty());
+
+    // Verify the import using the sync job
+    SyncTableOptions syncOpts = createSyncTableOptions();
+
+    PipelineResult result = SyncTableJob.buildPipeline(syncOpts).run();
+    state = result.waitUntilFinish();
+    Assert.assertEquals(State.DONE, state);
+
+    // Read the output files and validate that there are no mismatches.
+    Assert.assertEquals(0, readMismatchesFromOutputFiles().size());
+
+    // Validate the counters.
+    Map<String, Long> counters = getCountMap(result);
+    Assert.assertEquals(counters.get("ranges_matched"), (Long) 100L);
+    Assert.assertEquals(counters.get("ranges_not_matched"), (Long) 0L);
+  }
 }