feat: Ability to import HBase Snapshot data into Cloud Bigtable using Dataflow

bigtable-dataflow-parent/bigtable-beam-import/README.md

@@ -42,4 +42,4 @@ java -jar bigtable-beam-import-1.14.1-shaded.jar import \
     --maxNumWorkers=[3x number of nodes] \
     --zone=[zone of your cluster]
 ```
-[//]: # ({x-version-update-end})
+[//]: # ({x-version-update-end})

bigtable-dataflow-parent/bigtable-beam-import/pom.xml

@@ -25,7 +25,7 @@ limitations under the License.
   <artifactId>bigtable-beam-import</artifactId>
 
   <properties>
-    <mainClass>com.google.cloud.bigtable.beam.sequencefiles.Main</mainClass>
+    <mainClass>com.google.cloud.bigtable.beam.Main</mainClass>
   </properties>
 
   <!-- Adding this to resolve version conflict within beam sdk-->
@@ -49,6 +49,12 @@ limitations under the License.
       <groupId>${project.groupId}</groupId>
       <artifactId>bigtable-hbase-beam</artifactId>
       <version>${project.version}</version>
+      <exclusions>
+        <exclusion>
+          <groupId>org.apache.hbase</groupId>
+          <artifactId>hbase-shaded-client</artifactId>
+        </exclusion>
+      </exclusions>
     </dependency>
     <dependency>
       <groupId>com.google.cloud.bigtable</groupId>
@@ -67,6 +73,10 @@ limitations under the License.
           <groupId>io.opencensus</groupId>
           <artifactId>*</artifactId>
         </exclusion>
+        <exclusion>
+          <groupId>org.apache.hbase</groupId>
+          <artifactId>hbase-shaded-client</artifactId>
+        </exclusion>
       </exclusions>
     </dependency>
 
@@ -87,10 +97,18 @@ limitations under the License.
       <artifactId>beam-sdks-java-io-hadoop-common</artifactId>
       <version>${beam.version}</version>
     </dependency>
+    <dependency>
+      <groupId>org.apache.beam</groupId>
+      <artifactId>beam-sdks-java-io-hadoop-format</artifactId>
+      <version>${beam.version}</version>
+    </dependency>
 
+    <!-- For HBase 2.x, this should be hbase-mapreduce
+    https://hbase.apache.org/2.1/book.html#export
+    -->
     <dependency>
       <groupId>org.apache.hbase</groupId>
-      <artifactId>hbase-shaded-client</artifactId>
+      <artifactId>hbase-shaded-server</artifactId>
       <version>${hbase.version}</version>
     </dependency>
 
@@ -104,7 +122,7 @@ limitations under the License.
     <dependency>
       <groupId>com.google.guava</groupId>
       <artifactId>guava</artifactId>
-      <version>${guava.version}</version>
+      <version>${gcs-guava.version}</version>
     </dependency>
 
     <!-- TODO: check if commons-codec was transitively updated to 1.13 and okhttp was updated to 2.7.5 when upgrading-->
@@ -149,6 +167,21 @@ limitations under the License.
       <artifactId>slf4j-api</artifactId>
       <version>${slf4j.version}</version>
     </dependency>
+    <!-- https://mvnrepository.com/artifact/com.google.cloud.bigdataoss/gcs-connector -->
+    <dependency>
+      <groupId>com.google.cloud.bigdataoss</groupId>
+      <artifactId>gcs-connector</artifactId>
+      <version>hadoop2-2.1.4</version>
+      <classifier>shaded</classifier>
+    </dependency>
+
+    <!-- https://mvnrepository.com/artifact/com.google.apis/google-api-services-storage -->
+    <dependency>
+      <groupId>com.google.apis</groupId>
+      <artifactId>google-api-services-storage</artifactId>
+      <version>v1-rev171-1.25.0</version>
+    </dependency>
+
 
     <!-- Test -->
     <dependency>
@@ -181,6 +214,12 @@ limitations under the License.
       <version>${junit.version}</version>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>org.apache.hbase</groupId>
+      <artifactId>hbase-shaded-testing-util</artifactId>
+      <version>${hbase.version}</version>
+      <scope>test</scope>
+    </dependency>
   </dependencies>
 
   <build>
@@ -265,6 +304,16 @@ limitations under the License.
           <transformers>
             <transformer implementation="org.apache.maven.plugins.shade.resource.ServicesResourceTransformer" />
           </transformers>
+	  <filters>
+	      <filter>
+		  <artifact>*:*</artifact>
+		  <excludes>
+		      <exclude>META-INF/*.SF</exclude>
+		      <exclude>META-INF/*.DSA</exclude>
+		      <exclude>META-INF/*.RSA</exclude>
+		  </excludes>
+	      </filter>
+	  </filters>
         </configuration>
       </plugin>
 
@@ -376,5 +425,33 @@ limitations under the License.
         </plugins>
       </build>
     </profile>
+
+    <profile>
+      <id>hbasesnapshotsIntegrationTest</id>
+      <build>
+        <plugins>
+          <plugin>
+            <groupId>org.apache.maven.plugins</groupId>
+            <artifactId>maven-failsafe-plugin</artifactId>
+            <executions>
+              <execution>
+                <id>hbasesnapshots-integration-test</id>
+                <goals>
+                  <goal>integration-test</goal>
+                </goals>
+                <phase>integration-test</phase>
+                <configuration>
+                  <forkCount>1</forkCount>
+                  <includes>
+                    <include>**/hbasesnapshots/*IT.java</include>
+                  </includes>
+                  <useSystemClassLoader>false</useSystemClassLoader>
+                </configuration>
+              </execution>
+            </executions>
+          </plugin>
+        </plugins>
+      </build>
+    </profile>
   </profiles>
 </project>

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

@@ -0,0 +1,79 @@
+/*
+ * Copyright 2021 Google LLC
+ *
+ * Licensed 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 com.google.cloud.bigtable.beam;
+
+import com.google.bigtable.repackaged.com.google.api.core.InternalApi;
+import com.google.bigtable.repackaged.com.google.api.core.InternalExtensionOnly;
+import com.google.cloud.bigtable.beam.hbasesnapshots.ImportJobFromHbaseSnapshot;
+import com.google.cloud.bigtable.beam.sequencefiles.CreateTableHelper;
+import com.google.cloud.bigtable.beam.sequencefiles.ExportJob;
+import com.google.cloud.bigtable.beam.sequencefiles.ImportJob;
+import java.io.File;
+import java.net.URISyntaxException;
+import java.util.Arrays;
+
+/** Entry point for create-table/import/export job submission. */
+@InternalExtensionOnly
+final class Main {
+  /** For internal use only - public for technical reasons. */
+  @InternalApi("For internal usage only")
+  public Main() {}
+
+  public static void main(String[] args) throws Exception {
+    if (args.length < 1) {
+      usage();
+      System.exit(1);
+    }
+
+    String[] subArgs = Arrays.copyOfRange(args, 1, args.length);
+
+    switch (args[0]) {
+      case "export":
+        ExportJob.main(subArgs);
+        break;
+      case "import":
+        ImportJob.main(subArgs);
+        break;
+      case "importsnapshot":
+        ImportJobFromHbaseSnapshot.main(subArgs);
+        break;
+      case "create-table":
+        CreateTableHelper.main(subArgs);
+        break;
+      default:
+        usage();
+        System.exit(1);
+    }
+  }
+
+  private static void usage() {
+    String jarName;
+
+    try {
+      jarName =
+          new File(Main.class.getProtectionDomain().getCodeSource().getLocation().toURI().getPath())
+              .getName();
+    } catch (URISyntaxException e) {
+      jarName = "<jar>";
+    }
+
+    System.out.printf(
+        "java -jar %s <action> <action_params>\n"
+            + "Where <action> can be 'export', 'import' , 'importsnapshot' or 'create-table'. To get further help, run: \n"
+            + "java -jar %s <action> --help\n",
+        jarName, jarName);
+  }
+}

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

@@ -0,0 +1,84 @@
+/*
+ * Copyright 2021 Google LLC
+ *
+ * Licensed 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 com.google.cloud.bigtable.beam.hbasesnapshots;
+
+import com.google.api.services.storage.model.Objects;
+import com.google.common.base.Preconditions;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import org.apache.beam.sdk.extensions.gcp.options.GcpOptions;
+import org.apache.beam.sdk.extensions.gcp.util.GcsUtil;
+import org.apache.beam.sdk.extensions.gcp.util.gcsfs.GcsPath;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.values.KV;
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
+
+/**
+ * A {@link DoFn} that could be used for cleaning up temp files generated during HBase snapshot
+ * scans in Google Cloud Storage(GCS) bucket via GCS connector.
+ */
+class CleanupHBaseSnapshotRestoreFilesFn extends DoFn<KV<String, String>, Boolean> {
+  private static final Log LOG = LogFactory.getLog(CleanupHBaseSnapshotRestoreFilesFn.class);
+
+  @ProcessElement
+  public void processElement(ProcessContext context) throws IOException {
+    KV<String, String> elem = context.element();
+
+    String hbaseSnapshotDir = elem.getKey();
+    String restorePath = elem.getValue();
+    String prefix = getListPrefix(restorePath);
+    String bucketName = getWorkingBucketName(hbaseSnapshotDir);
+    Preconditions.checkState(
+        !prefix.isEmpty() && !hbaseSnapshotDir.contains(String.format("%s/%s", bucketName, prefix)),
+        "restore folder should not be empty or a subfolder of hbaseSnapshotSourceDir");
+    GcpOptions gcpOptions = context.getPipelineOptions().as(GcpOptions.class);
+    GcsUtil gcsUtil = new GcsUtil.GcsUtilFactory().create(gcpOptions);
+
+    String pageToken = null;
+    List<String> results = new ArrayList<>();
+    do {
+      Objects objects = gcsUtil.listObjects(bucketName, prefix, pageToken);
+      if (objects.getItems() == null) {
+        break;
+      }
+
+      objects.getItems().stream()
+          .map(storageObject -> GcsPath.fromObject(storageObject).toString())
+          .forEach(results::add);
+      pageToken = objects.getNextPageToken();
+    } while (pageToken != null);
+    gcsUtil.remove(results);
+    context.output(true);
+  }
+
+  public static String getWorkingBucketName(String hbaseSnapshotDir) {
+    Preconditions.checkArgument(
+        hbaseSnapshotDir.startsWith(GcsPath.SCHEME),
+        "snapshot folder must be hosted in a GCS bucket ");
+
+    return GcsPath.fromUri(hbaseSnapshotDir).getBucket();
+  }
+  // getListPrefix convert absolute restorePath in a Hadoop filesystem
+  // to a match prefix in a GCS bucket
+  public static String getListPrefix(String restorePath) {
+    Preconditions.checkArgument(
+        restorePath.startsWith("/"),
+        "restore folder must be an absolute path in current filesystem");
+    return restorePath.substring(1);
+  }
+}