snapshots = null;
+ if (snapshotsProperty != null) {
+ snapshots =
+ (sourceDir != null && SnapshotUtils.isRegex(snapshotsProperty))
+ ? SnapshotUtils.getSnapshotsFromSnapshotPath(
+ sourceDir, gcsOptions.getGcsUtil(), snapshotsProperty)
+ : SnapshotUtils.getSnapshotsFromString(snapshotsProperty);
+ }
+
+ ImportConfig importConfig = new ImportConfig();
+ importConfig.setSourcepath(sourceDir);
+ if (snapshots != null) {
+ importConfig.setSnapshotsFromMap(snapshots);
+ }
+ importConfig.validate();
+ SnapshotUtils.setRestorePath(System.getProperty(RESTORE_PATH_PROPERTY), importConfig);
+
+ return importConfig;
+ }
+
+ @VisibleForTesting
+ static ImportConfig buildImportConfigFromConfigFile(String configFilePath) throws Exception {
+ Gson gson = new GsonBuilder().create();
+ ImportConfig importConfig =
+ gson.fromJson(SnapshotUtils.readFileContents(configFilePath), ImportConfig.class);
+ Preconditions.checkNotNull(importConfig, "ImportConfig parsed from file cannot be null.");
+ importConfig.validate();
+ SnapshotUtils.setRestorePath(importConfig.getRestorepath(), importConfig);
+ return importConfig;
+ }
+
+ @VisibleForTesting
+ /**
+ * Creates a copy of Snasphsot from the source path into restore path.
+ *
+ * @param snapshotConfig - Snapshot Configuration
+ * @throws IOException
+ */
+ static void restoreSnapshot(SnapshotConfig snapshotConfig) throws IOException {
+ Path sourcePath = snapshotConfig.getSourcePath();
+ Path restorePath = snapshotConfig.getRestorePath();
+ Configuration configuration = snapshotConfig.getConfiguration();
+ LOG.info(
+ String.format("RestoreSnapshot - sourcePath:%s restorePath: %s", sourcePath, restorePath));
+ FileSystem fileSystem = sourcePath.getFileSystem(configuration);
+ if (fileSystem.exists(restorePath)) {
+ LOG.info(
+ String.format(
+ "Restore path %s already exists, deleting it for idempotency", restorePath));
+ fileSystem.delete(restorePath, true);
+ }
+ RestoreSnapshotHelper.copySnapshotForScanner(
+ configuration, fileSystem, sourcePath, restorePath, snapshotConfig.getSnapshotName());
+ }
+}
diff --git a/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/ImportJobFromHbaseSnapshot.java b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/ImportJobFromHbaseSnapshot.java
index 064736a04b..40747862b7 100644
--- a/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/ImportJobFromHbaseSnapshot.java
+++ b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/ImportJobFromHbaseSnapshot.java
@@ -15,21 +15,38 @@
*/
package com.google.cloud.bigtable.beam.hbasesnapshots;
-import com.google.bigtable.repackaged.com.google.api.core.InternalExtensionOnly;
+import com.google.api.core.InternalExtensionOnly;
import com.google.cloud.bigtable.beam.CloudBigtableIO;
+import com.google.cloud.bigtable.beam.CloudBigtableTableConfiguration;
import com.google.cloud.bigtable.beam.TemplateUtils;
+import com.google.cloud.bigtable.beam.hbasesnapshots.conf.HBaseSnapshotInputConfigBuilder;
+import com.google.cloud.bigtable.beam.hbasesnapshots.conf.ImportConfig;
+import com.google.cloud.bigtable.beam.hbasesnapshots.conf.SnapshotConfig;
+import com.google.cloud.bigtable.beam.hbasesnapshots.dofn.CleanupHBaseSnapshotRestoreFilesFn;
+import com.google.cloud.bigtable.beam.hbasesnapshots.dofn.CleanupRestoredSnapshotsFn;
+import com.google.cloud.bigtable.beam.hbasesnapshots.dofn.RestoreSnapshotFn;
+import com.google.cloud.bigtable.beam.hbasesnapshots.transforms.ListRegions;
+import com.google.cloud.bigtable.beam.hbasesnapshots.transforms.ReadRegions;
import com.google.cloud.bigtable.beam.sequencefiles.HBaseResultToMutationFn;
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.gson.Gson;
+import com.google.gson.GsonBuilder;
+import java.io.IOException;
import java.util.Arrays;
import java.util.List;
+import java.util.Map;
+import org.apache.beam.runners.dataflow.options.DataflowPipelineDebugOptions;
import org.apache.beam.sdk.Pipeline;
import org.apache.beam.sdk.PipelineResult;
+import org.apache.beam.sdk.extensions.gcp.options.GcsOptions;
+import org.apache.beam.sdk.io.FileSystems;
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.options.ValueProvider;
import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.transforms.ParDo;
import org.apache.beam.sdk.transforms.Wait;
@@ -37,6 +54,7 @@
import org.apache.beam.sdk.values.PCollection;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
+import org.apache.hadoop.hbase.client.Mutation;
import org.apache.hadoop.hbase.client.Result;
import org.apache.hadoop.hbase.io.ImmutableBytesWritable;
@@ -63,11 +81,32 @@
* Note that in the case of job failures, the temp files generated in the .restore-$JOB_NAME
* directory under the snapshot export bucket will not get deleted. Hence one need to either launch
* a replacement job with the same jobName to re-run the job or manually delete this directory.
+ * Additionally, it is highly recommended to set a GCS Lifecycle TTL (e.g., 7 days) on the bucket
+ * used for the restore path to automatically clean up any orphaned files.
+ *
+ * Running Parallel Sharded Jobs: To import a snapshot in parallel shards (using {@code
+ * --numShards} and {@code --shardIndex}), you must first run a single restore step to copy the
+ * files to the restore path (e.g., using this job with {@code --performOnlyRestoreStep=true} or
+ * using {@link HBaseSnapshotRestoreTool}). Once the restore is complete, launch the parallel
+ * sharded import jobs concurrently, making sure to set {@code --skipRestoreStep=true} and {@code
+ * --deleteRestoredSnapshots=false} on all shards. This prevents concurrent shards from deleting the
+ * restore path while other shards are still reading.
*/
@InternalExtensionOnly
public class ImportJobFromHbaseSnapshot {
private static final Log LOG = LogFactory.getLog(ImportJobFromHbaseSnapshot.class);
+ @VisibleForTesting
+ static final String MISSING_SNAPSHOT_SOURCEPATH =
+ "Source Path containing hbase snapshots must be specified.";
+
+ @VisibleForTesting
+ static final String MISSING_SNAPSHOT_NAMES =
+ "Snapshots must be specified. Allowed values are '*' (indicating all snapshots under source"
+ + " path) or 'prefix*' (snapshots matching certain prefix) or"
+ + " 'snapshotname1:tablename1,snapshotname2:tablename2' (comma seperated list of"
+ + " snapshots)";
+
public interface ImportOptions extends ImportJob.ImportOptions {
@Description("The HBase root dir where HBase snapshot files resides.")
String getHbaseSnapshotSourceDir();
@@ -87,24 +126,298 @@ public interface ImportOptions extends ImportJob.ImportOptions {
@SuppressWarnings("unused")
void setEnableSnappy(Boolean enableSnappy);
+
+ @Description("Path to config file containing snapshot source path/snapshot names.")
+ String getImportConfigFilePath();
+
+ void setImportConfigFilePath(String value);
+
+ @Description(
+ "Snapshots to be imported. Can be '*', 'prefix*' or 'snap1,snap2' or"
+ + " 'snap1:table1,snap2:table2'.")
+ String getSnapshots();
+
+ void setSnapshots(String value);
+
+ @Description("Specifies whether to use dynamic splitting while reading hbase region.")
+ @Default.Boolean(true)
+ boolean getUseDynamicSplitting();
+
+ void setUseDynamicSplitting(boolean value);
+
+ @Description("Specifies the threshold for number of cells per mutation written.")
+ @Default.Integer(100_000 - 1)
+ int getMaxMutationsPerRequestThreshold();
+
+ void setMaxMutationsPerRequestThreshold(int value);
+
+ @Description(
+ "Specifies whether to filter large rows that exceed FilterLargeRowsThresholdBytes should be"
+ + " logged and dropped.")
+ @Default.Boolean(false)
+ boolean getFilterLargeRows();
+
+ void setFilterLargeRows(boolean value);
+
+ @Description(
+ "Specifies the size in bytes of a row that should be logged and dropped before loading to"
+ + " Bigtable.")
+ @Default.Long(256 * 1024 * 1024)
+ long getFilterLargeRowsThresholdBytes();
+
+ void setFilterLargeRowsThresholdBytes(long value);
+
+ @Description(
+ "Specifies whether to filter large cells that exceed FilterLargeCellsThresholdBytes should"
+ + " be logged and dropped.")
+ @Default.Boolean(true)
+ boolean getFilterLargeCells();
+
+ void setFilterLargeCells(boolean value);
+
+ @Description(
+ "Specifies the size in bytes of a cell that should be logged and dropped before loading to"
+ + " Bigtable.")
+ @Default.Integer(100 * 1024 * 1024)
+ int getFilterLargeCellsThresholdBytes();
+
+ void setFilterLargeCellsThresholdBytes(int value);
+
+ @Description(
+ "Specifies whether to filter large row keys that exceed FilterLargeRowKeysThresholdBytes"
+ + " should be logged and dropped.")
+ @Default.Boolean(false)
+ boolean getFilterLargeRowKeys();
+
+ void setFilterLargeRowKeys(boolean value);
+
+ @Description(
+ "Drops wide rows exceeding MaxMutationsPerRequestThreshold to prevent atomicity loss from"
+ + " splitting (losing data), otherwise splits them to preserve data.")
+ @Default.Boolean(false)
+ boolean getFilterWideRows();
+
+ void setFilterWideRows(boolean value);
+
+ @Description(
+ "Specifies the size in bytes of a row key that should be logged and dropped before loading"
+ + " to Bigtable.")
+ @Default.Integer(4 * 1024)
+ int getFilterLargeRowKeysThresholdBytes();
+
+ void setFilterLargeRowKeysThresholdBytes(int value);
+
+ @Description(
+ "Specifies the number of shards to use when loading the snapshot. "
+ + "If set, shardIndex must also be set.")
+ Integer getNumShards();
+
+ void setNumShards(Integer value);
+
+ @Description("Specifies the shard index from [0, numShards) that this load represents.")
+ Integer getShardIndex();
+
+ void setShardIndex(Integer value);
+
+ @Description("Specifies the path to the restored Snapshot files.")
+ String getRestorePath();
+
+ void setRestorePath(String value);
+
+ @Description(
+ "Specifies whether the snapshots restored should be deleted. Note: Cleanup is best-effort"
+ + " and will not fail the job if deletion fails after retries. When running parallel"
+ + " sharded jobs concurrently, this must be set to false to prevent a shard from"
+ + " deleting the files while other shards are still reading.")
+ @Default.Boolean(false)
+ Boolean getDeleteRestoredSnapshots();
+
+ void setDeleteRestoredSnapshots(Boolean value);
+
+ @Description(
+ "Specifies whether the restore step should be skipped. When running parallel sharded jobs"
+ + " concurrently, this must be set to true (after running a single separate restore"
+ + " step beforehand) to prevent concurrent restore steps from deleting and corrupting"
+ + " active files.")
+ @Default.Boolean(false)
+ Boolean getSkipRestoreStep();
+
+ void setSkipRestoreStep(Boolean value);
+
+ @Description("Specifies whether to perform only restore step.")
+ @Default.Boolean(false)
+ Boolean getPerformOnlyRestoreStep();
+
+ void setPerformOnlyRestoreStep(Boolean value);
}
public static void main(String[] args) throws Exception {
PipelineOptionsFactory.register(ImportOptions.class);
- ImportOptions opts =
+ ImportOptions options =
PipelineOptionsFactory.fromArgs(args).withValidation().as(ImportOptions.class);
+ // To determine the Google Cloud Storage file scheme (gs://)
+ FileSystems.setDefaultPipelineOptions(options);
+
LOG.info("Building Pipeline");
- Pipeline pipeline = buildPipeline(opts);
+ Pipeline pipeline = null;
+ ImportConfig importConfig = null;
+ // Maintain Backward compatibility until deprecation
+ if (options.getSnapshotName() != null && !options.getSnapshotName().isEmpty()) {
+ pipeline = buildPipeline(options);
+ } else {
+ importConfig =
+ options.getImportConfigFilePath() != null
+ ? buildImportConfigFromConfigFile(options.getImportConfigFilePath())
+ : buildImportConfigFromPipelineOptions(options, options.as(GcsOptions.class));
+
+ LOG.info(
+ String.format(
+ "SourcePath:%s, RestorePath:%s",
+ importConfig.getSourcepath(), importConfig.getRestorepath()));
+ pipeline = buildPipelineWithMultipleSnapshots(options, importConfig);
+ }
+
LOG.info("Running Pipeline");
PipelineResult result = pipeline.run();
-
- if (opts.getWait()) {
+ if (options.getWait()) {
Utils.waitForPipelineToFinish(result);
}
}
+ @VisibleForTesting
+ static ImportConfig buildImportConfigFromConfigFile(String configFilePath) throws Exception {
+ Gson gson = new GsonBuilder().create();
+ ImportConfig importConfig =
+ gson.fromJson(SnapshotUtils.readFileContents(configFilePath), ImportConfig.class);
+ importConfig.validate();
+ SnapshotUtils.setRestorePath(importConfig.getRestorepath(), importConfig);
+ return importConfig;
+ }
+
+ @VisibleForTesting
+ static ImportConfig buildImportConfigFromPipelineOptions(
+ ImportOptions options, GcsOptions gcsOptions) throws IOException {
+ String sourceDir = options.getHbaseSnapshotSourceDir();
+ String snapshotsProperty = options.getSnapshots();
+ Map snapshots = null;
+ if (snapshotsProperty != null) {
+ snapshots =
+ (sourceDir != null && SnapshotUtils.isRegex(snapshotsProperty))
+ ? SnapshotUtils.getSnapshotsFromSnapshotPath(
+ sourceDir, gcsOptions.getGcsUtil(), snapshotsProperty)
+ : SnapshotUtils.getSnapshotsFromString(snapshotsProperty);
+ }
+
+ ImportConfig importConfig = new ImportConfig();
+ importConfig.setSourcepath(sourceDir);
+ if (snapshots != null) {
+ importConfig.setSnapshotsFromMap(snapshots);
+ }
+ importConfig.validate();
+ SnapshotUtils.setRestorePath(options.getRestorePath(), importConfig);
+ return importConfig;
+ }
+
+ /**
+ * Builds the pipeline that supports loading multiple snapshots to BigTable.
+ *
+ * @param options - Pipeline options
+ * @param importConfig - Configuration representing snapshot source path, list of snapshots etc
+ * @return
+ * @throws Exception
+ */
+ static Pipeline buildPipelineWithMultipleSnapshots(
+ ImportOptions options, ImportConfig importConfig) throws Exception {
+ Map configurations =
+ SnapshotUtils.getConfiguration(
+ options.getRunner().getSimpleName(),
+ options.getProject(),
+ importConfig.getSourcepath(),
+ importConfig.getHbaseConfiguration());
+
+ List snapshotConfigs =
+ SnapshotUtils.buildSnapshotConfigs(
+ importConfig.getSnapshots(),
+ configurations,
+ options.getProject(),
+ importConfig.getSourcepath(),
+ importConfig.getRestorepath());
+ DataflowPipelineDebugOptions debugOptions = options.as(DataflowPipelineDebugOptions.class);
+ // Disable GC thrashing detection to prevent Dataflow from prematurely killing
+ // workers during memory-intensive HBase snapshot scans.
+ debugOptions.setGCThrashingPercentagePerPeriod(100.00);
+
+ Pipeline pipeline = Pipeline.create(debugOptions);
+
+ PCollection restoredSnapshots =
+ pipeline.apply("Read Snapshot Configs", Create.of(snapshotConfigs));
+ if (!options.getSkipRestoreStep()) {
+ restoredSnapshots =
+ restoredSnapshots.apply("Restore Snapshots", ParDo.of(new RestoreSnapshotFn()));
+ }
+ if (options.getPerformOnlyRestoreStep()) {
+ return pipeline;
+ }
+ // Read records from hbase region files and write to Bigtable
+ PCollection>> hbaseRecords =
+ restoredSnapshots
+ .apply("List Regions", new ListRegions())
+ .apply(
+ "Read Regions",
+ new ReadRegions(
+ options.getUseDynamicSplitting(),
+ options.getMaxMutationsPerRequestThreshold(),
+ options.getFilterLargeRows(),
+ options.getFilterLargeRowsThresholdBytes(),
+ options.getFilterLargeCells(),
+ options.getFilterLargeCellsThresholdBytes(),
+ options.getFilterLargeRowKeys(),
+ options.getFilterLargeRowKeysThresholdBytes(),
+ options.getFilterWideRows(),
+ options.getNumShards(),
+ options.getShardIndex()));
+
+ options.setBigtableTableId(ValueProvider.StaticValueProvider.of("NA"));
+ CloudBigtableTableConfiguration bigtableConfiguration =
+ TemplateUtils.buildImportConfig(options, "HBaseSnapshotImportJob");
+ if (importConfig.getBigtableConfiguration() != null) {
+ CloudBigtableTableConfiguration.Builder builder = bigtableConfiguration.toBuilder();
+ for (Map.Entry entry : importConfig.getBigtableConfiguration().entrySet()) {
+ builder = builder.withConfiguration(entry.getKey(), entry.getValue());
+ }
+ bigtableConfiguration = builder.build();
+ }
+
+ hbaseRecords.apply(
+ "Write to BigTable", CloudBigtableIO.writeToMultipleTables(bigtableConfiguration));
+
+ // Clean up all the temporary restored snapshot HLinks after reading all the data
+ if (options.getDeleteRestoredSnapshots()) {
+ restoredSnapshots
+ .apply(Wait.on(hbaseRecords))
+ .apply(
+ "Clean restored files",
+ ParDo.of(
+ new CleanupRestoredSnapshotsFn(
+ importConfig.getBackoffInitialIntervalInMillis(),
+ importConfig.getBackoffMaxIntervalInMillis(),
+ importConfig.getBackoffMaxretries())));
+ }
+
+ return pipeline;
+ }
+
+ /**
+ * Builds the pipeline that supports loading single snapshot to BigTable. Maintained for backward
+ * compatiablity and will be deprecated merging the functionality to
+ * buildPipelineWithMultipleSnapshots method.
+ *
+ * @param opts - Pipeline options
+ * @return
+ * @throws Exception
+ */
@VisibleForTesting
static Pipeline buildPipeline(ImportOptions opts) throws Exception {
Pipeline pipeline = Pipeline.create(Utils.tweakOptions(opts));
diff --git a/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/SnapshotUtils.java b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/SnapshotUtils.java
new file mode 100644
index 0000000000..e2f4e2ecf9
--- /dev/null
+++ b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/SnapshotUtils.java
@@ -0,0 +1,436 @@
+/*
+ * Copyright 2026 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.core.InternalApi;
+import com.google.api.services.storage.model.Objects;
+import com.google.api.services.storage.model.StorageObject;
+import com.google.cloud.bigtable.beam.hbasesnapshots.conf.ImportConfig;
+import com.google.cloud.bigtable.beam.hbasesnapshots.conf.SnapshotConfig;
+import com.google.common.base.Joiner;
+import com.google.common.io.CharStreams;
+import java.io.File;
+import java.io.IOException;
+import java.io.Reader;
+import java.net.URI;
+import java.net.URISyntaxException;
+import java.nio.channels.Channels;
+import java.nio.charset.StandardCharsets;
+import java.time.Instant;
+import java.time.ZoneId;
+import java.time.format.DateTimeFormatter;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.UUID;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+import java.util.stream.Collectors;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.extensions.gcp.util.GcsUtil;
+import org.apache.beam.sdk.extensions.gcp.util.gcsfs.GcsPath;
+import org.apache.beam.sdk.io.FileSystems;
+import org.apache.beam.sdk.util.BackOff;
+import org.apache.beam.sdk.util.FluentBackoff;
+import org.apache.beam.sdk.util.Sleeper;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.hbase.HBaseConfiguration;
+import org.apache.hadoop.hbase.HConstants;
+import org.joda.time.Duration;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Contains various helper methods to handle different tasks associated with importing of hbase
+ * snapshots
+ */
+@InternalApi("For internal usage only")
+public class SnapshotUtils {
+ private static final Logger LOG = LoggerFactory.getLogger(SnapshotUtils.class);
+ private static final String DIRECTRUNNER = "DirectRunner";
+ private static final String SNAPSHOT_MANIFEST_DIRECTORY = ".hbase-snapshot";
+ private static final String GCS_SCHEME = "gs";
+ private static final Sleeper sleeper = Sleeper.DEFAULT;
+
+ private SnapshotUtils() {}
+
+ private static String getParentDirectory(String hbaseSnapshotSourceDirectory) {
+ URI hbaseSnapshotSourceUri;
+ try {
+ hbaseSnapshotSourceUri = new URI(hbaseSnapshotSourceDirectory);
+ } catch (URISyntaxException e) {
+ throw new RuntimeException(
+ String.format(
+ "Invalid file path format for snapshot source directory: %s. Valid paths should have"
+ + " file scheme (gs://, file://)",
+ hbaseSnapshotSourceDirectory));
+ }
+
+ if (hbaseSnapshotSourceUri.getScheme() != null
+ && hbaseSnapshotSourceUri.getScheme().equals(GCS_SCHEME)) // i.e Cloud Storage file system
+ {
+ return GcsPath.fromUri(hbaseSnapshotSourceUri).getParent().toString();
+ }
+
+ return new File(hbaseSnapshotSourceDirectory).getParent();
+ }
+
+ static String removeSuffixSlashIfExists(String directory) {
+ return directory.endsWith("/") ? directory.substring(0, directory.length() - 1) : directory;
+ }
+
+ static String appendCurrentTimestamp(String directory) {
+ DateTimeFormatter formatter =
+ DateTimeFormatter.ofPattern("yyyyMMddHHmm").withZone(ZoneId.of("UTC"));
+ String uuid = UUID.randomUUID().toString();
+ return String.join(
+ "/", removeSuffixSlashIfExists(directory), formatter.format(Instant.now()) + "-" + uuid);
+ }
+
+ static String getNamedDirectory(String sourceDirectory, String subFoldername) {
+ String parentDirectory = getParentDirectory(sourceDirectory);
+ if (parentDirectory == null) {
+ throw new IllegalArgumentException(
+ "Source directory has no parent directory: " + sourceDirectory);
+ }
+ parentDirectory = removeSuffixSlashIfExists(parentDirectory);
+ return appendCurrentTimestamp(String.join("/", parentDirectory, subFoldername));
+ }
+
+ /** Builds the configuration combining default and user provided values. */
+ static Map getConfiguration(
+ String runner,
+ String project,
+ String sourcedir,
+ @Nullable Map hbaseConfiguration) {
+ Map configurations = new HashMap<>();
+
+ configurations.put(HConstants.HBASE_DIR, sourcedir);
+ configurations.put(
+ "fs.AbstractFileSystem.gs.impl", "com.google.cloud.hadoop.fs.gcs.GoogleHadoopFS");
+ configurations.put("fs.gs.project.id", project);
+ configurations.put("google.cloud.auth.service.account.enable", "true");
+
+ if (runner == null || runner.equals(DIRECTRUNNER)) {
+ // https://github.com/GoogleCloudDataproc/hadoop-connectors/blob/master/gcs/CONFIGURATION.md#authentication
+ configurations.put("fs.gs.auth.type", "APPLICATION_DEFAULT");
+ }
+
+ // Update the default configurations with user supplied configuration values
+ if (hbaseConfiguration != null) {
+ configurations.putAll(hbaseConfiguration);
+ }
+ return configurations;
+ }
+
+ public static Configuration getHBaseConfiguration(Map configurations) {
+ if (configurations == null) {
+ configurations = new java.util.HashMap<>();
+ }
+ return createHbaseConfiguration(configurations);
+ }
+
+ private static Configuration createHbaseConfiguration(Map configurations) {
+ LOG.info("Create HBase Configuration instance");
+ Configuration hbaseConfiguration = HBaseConfiguration.create();
+ for (Map.Entry entry : configurations.entrySet())
+ hbaseConfiguration.set(entry.getKey(), entry.getValue());
+ return hbaseConfiguration;
+ }
+
+ /**
+ * Build Snapshot Configurations.
+ *
+ * @param snapshotdetails - Snapshot details representing hbase snapshot name and corresponding
+ * bigtable table name.
+ * @param configurations - BigTable Configurations
+ * @param projectId - Google Cloud Project Id
+ * @param sourcePath - Source path containing snapshot files
+ * @param restorePath - Path snapshot files gets stored during job runs.
+ * @return
+ */
+ static List buildSnapshotConfigs(
+ List snapshotdetails,
+ Map configurations,
+ String projectId,
+ String sourcePath,
+ String restorePath) {
+
+ return snapshotdetails.stream()
+ .map(
+ snapshotInfo ->
+ SnapshotConfig.builder()
+ .setProjectId(projectId)
+ .setSourceLocation(sourcePath)
+ .setRestoreLocation(restorePath + "/" + snapshotInfo.getSnapshotName())
+ .setSnapshotName(snapshotInfo.getSnapshotName())
+ .setTableName(snapshotInfo.getbigtableTableName())
+ .setConfigurationDetails(configurations)
+ .build())
+ .collect(Collectors.toList());
+ }
+
+ public static BackOff createBackOff(
+ long backoffInitialIntervalInMillis, long backoffMaxIntervalInMillis, int maxRetries) {
+ return FluentBackoff.DEFAULT
+ .withInitialBackoff(Duration.millis(backoffInitialIntervalInMillis))
+ .withMaxRetries(maxRetries)
+ .withMaxBackoff(Duration.millis(backoffMaxIntervalInMillis))
+ .backoff();
+ }
+
+ /**
+ * Creates restore path based on the input configuration
+ *
+ * @param importConfig - Job Configuration
+ */
+ public static void setRestorePath(ImportConfig importConfig) {
+ importConfig.setRestorepath(
+ formatRestorePath(importConfig.getRestorepath(), importConfig.getSourcepath()));
+ }
+
+ /**
+ * Creates restore path based on the input configuration
+ *
+ * @param restorePath - Restore path of the job.
+ * @param importConfig - Import config where we will set the restorePath property.
+ */
+ public static void setRestorePath(String restorePath, ImportConfig importConfig) {
+ if (restorePath != null) {
+ importConfig.setRestorepath(restorePath);
+ return;
+ }
+ importConfig.setRestorepath(
+ formatRestorePath(importConfig.getRestorepath(), importConfig.getSourcepath()));
+ }
+
+ /**
+ * Parses the provided input to generate snapshot names and corresponding bigtable names. For
+ * single snapshot names the following are valid formats: If both Snapshotname and bigtablename
+ * are same then only snapshotname can be provided If bigtablename is different then should be
+ * provided in the format snapshotname:bigtablename Multiple snapshots can be provided in
+ * snapshot1:table1,snapshot2:table2 format or snapshot1,snapshot2 format
+ *
+ * @param snapshotNames - Snapshot names and corresponding bigtable table names.
+ */
+ public static Map getSnapshotsFromString(String snapshotNames) {
+ Map snapshots = new HashMap<>();
+ for (String snapshotInfo : snapshotNames.split(",")) {
+ snapshotInfo = snapshotInfo.trim();
+ if (snapshotInfo.isEmpty()) {
+ continue;
+ }
+ String[] snapshotWithTableName = snapshotInfo.split(":");
+ if (snapshotWithTableName.length == 2) {
+ String snap = snapshotWithTableName[0].trim();
+ String table = snapshotWithTableName[1].trim();
+ if (snap.isEmpty() || table.isEmpty()) {
+ throw new IllegalArgumentException(
+ "Snapshot name and table name cannot be empty: " + snapshotInfo);
+ }
+ snapshots.put(snap, table);
+ } else if (snapshotWithTableName.length == 1) {
+ String snap = snapshotWithTableName[0].trim();
+ if (snap.isEmpty()) {
+ throw new IllegalArgumentException("Snapshot name cannot be empty: " + snapshotInfo);
+ }
+ snapshots.put(snap, snap);
+ } else {
+ throw new IllegalArgumentException(
+ "Invalid specification format for snapshots. Expected format is"
+ + " snapshot1:table1,snapshot2:table2");
+ }
+ }
+ return snapshots;
+ }
+
+ public static String formatRestorePath(String providedPath, String hbaseSnapshotsPath) {
+ return providedPath == null
+ ? SnapshotUtils.getNamedDirectory(hbaseSnapshotsPath, "restore")
+ : SnapshotUtils.appendCurrentTimestamp(providedPath);
+ }
+
+ /**
+ * Read list of Snapshot names from Snapshot Source Path
+ *
+ * @param importSnapshotpath - Path representing the snapshot source directory
+ * @param gcsUtil - GCS Instance
+ * @param prefix - Specific prefix to be matched or '*' for all files.
+ * @return
+ * @throws IOException
+ */
+ public static Map getSnapshotsFromSnapshotPath(
+ String importSnapshotpath, GcsUtil gcsUtil, String prefix) throws IOException {
+
+ importSnapshotpath =
+ Joiner.on("/")
+ .join(removeSuffixSlashIfExists(importSnapshotpath), SNAPSHOT_MANIFEST_DIRECTORY);
+
+ // Determine the filesystem scheme. The import path can be a GCS path (gs://) or a local
+ // filesystem path (file:// or raw absolute/relative path) to support running the offline
+ // HBaseSnapshotRestoreTool and running local tests.
+ URI uri;
+ try {
+ uri = new URI(importSnapshotpath);
+ } catch (URISyntaxException e) {
+ uri = null;
+ }
+ boolean isGcs = uri != null && GCS_SCHEME.equalsIgnoreCase(uri.getScheme());
+ Map snapshots = new HashMap<>();
+
+ if (isGcs) {
+ // Build GCS path from given string e.g:
+ // gs://sym-bucket/snapshots/20220309230526/.hbase-snapshot
+ GcsPath gcsPath = GcsPath.fromUri(importSnapshotpath);
+ String gcsPrefix = gcsPath.getObject();
+ Pattern prefixPattern = compilePrefixPattern(prefix);
+
+ // Optimize prefix matching by passing simple globs or literals to GCS API directly,
+ // reducing the number of objects fetched and processed in memory.
+ if (prefix.equals("*")) {
+ // List all
+ } else if (prefix.endsWith("*")
+ && !prefix.substring(0, prefix.length() - 1).contains("*")
+ && !prefix.contains("+")
+ && !prefix.contains("?")) {
+ // Simple glob like "prefix*"
+ String literal = prefix.substring(0, prefix.length() - 1);
+ gcsPrefix = gcsPrefix + "/" + literal;
+ } else if (!isRegex(prefix)) {
+ // Literal match
+ gcsPrefix = gcsPrefix + "/" + prefix;
+ }
+
+ List allObjects = new ArrayList<>();
+ String pageToken = null;
+ do {
+ Objects objects = gcsUtil.listObjects(gcsPath.getBucket(), gcsPrefix, pageToken);
+ if (objects == null || objects.getItems() == null) {
+ break;
+ }
+ allObjects.addAll(objects.getItems());
+ pageToken = objects.getNextPageToken();
+ } while (pageToken != null);
+
+ if (allObjects.isEmpty()) {
+ LOG.warn("Snapshot path {} does not contain any snapshots", importSnapshotpath);
+ return snapshots;
+ }
+
+ // Build a pattern for object portion e.g if path is
+ // gs://sym-bucket/snapshots/20220309230526/.hbase-snapshot
+ // the object portion would be snapshots/60G/20220309230526/.hbase-snapshot
+ Pattern pathPattern =
+ Pattern.compile(String.format("%s/(.+?/)", Pattern.quote(gcsPath.getObject())));
+ Matcher pathMatcher = null;
+ String snapshotName = null;
+ for (StorageObject object : allObjects) {
+ if (object == null || object.getName() == null) {
+ continue;
+ }
+ pathMatcher = pathPattern.matcher(object.getName());
+ if (pathMatcher.find()) {
+ // Group 1 represents the snapshot directory name along with suffix slash (e.g:
+ // snapshot1/)
+ snapshotName = pathMatcher.group(1).replace("/", "");
+ if (prefix.equals("*") || prefixPattern.matcher(snapshotName).find()) {
+ snapshots.put(snapshotName, snapshotName);
+ }
+ }
+ }
+ } else {
+ // Use local file system to find snapshots
+ File manifestDir =
+ (uri != null && uri.getScheme() != null && uri.getScheme().equalsIgnoreCase("file"))
+ ? new File(uri.getPath())
+ : new File(importSnapshotpath);
+
+ if (!manifestDir.exists() || !manifestDir.isDirectory()) {
+ LOG.warn("Snapshot path {} does not exist or is not a directory", importSnapshotpath);
+ return snapshots;
+ }
+ File[] files = manifestDir.listFiles();
+ if (files == null || files.length == 0) {
+ LOG.warn("Snapshot path {} does not contain any snapshots", importSnapshotpath);
+ return snapshots;
+ }
+
+ Pattern prefixPattern = compilePrefixPattern(prefix);
+ for (File file : files) {
+ if (file.isDirectory()) {
+ String snapshotName = file.getName();
+ if (prefix.equals("*") || prefixPattern.matcher(snapshotName).find()) {
+ snapshots.put(snapshotName, snapshotName);
+ }
+ }
+ }
+ }
+
+ return snapshots;
+ }
+
+ private static Pattern compilePrefixPattern(String prefix) {
+ if (prefix.equals("*")) {
+ return null;
+ } else if (prefix.endsWith("*")
+ && !prefix.substring(0, prefix.length() - 1).contains("*")
+ && !prefix.contains("+")
+ && !prefix.contains("?")) {
+ // Simple glob like "prefix*"
+ String literal = prefix.substring(0, prefix.length() - 1);
+ return Pattern.compile("^" + Pattern.quote(literal) + ".*");
+ } else if (!isRegex(prefix)) {
+ // Literal match
+ return Pattern.compile("^" + Pattern.quote(prefix) + "$");
+ } else {
+ // Complex regex
+ return Pattern.compile(prefix);
+ }
+ }
+
+ /**
+ * Reads the contents of file
+ *
+ * @param filePath - Path of the file.
+ * @return
+ * @throws IOException
+ */
+ public static String readFileContents(String filePath) throws IOException {
+ try (Reader reader =
+ Channels.newReader(
+ FileSystems.open(FileSystems.matchSingleFileSpec(filePath).resourceId()),
+ StandardCharsets.UTF_8.name())) {
+ return CharStreams.toString(reader);
+ }
+ }
+
+ /**
+ * Check if the given value contains any character in given meta characters list
+ *
+ * @param data - text value
+ * @return
+ */
+ public static boolean isRegex(String data) {
+ String[] metaChars = {"*", "+", "?"};
+ return Arrays.stream(metaChars).anyMatch(data::contains);
+ }
+
+ public static Sleeper getSleeper() {
+ return sleeper;
+ }
+}
diff --git a/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/coders/RegionConfigCoder.java b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/coders/RegionConfigCoder.java
new file mode 100644
index 0000000000..17886d4ee1
--- /dev/null
+++ b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/coders/RegionConfigCoder.java
@@ -0,0 +1,100 @@
+/*
+ * Copyright 2026 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.coders;
+
+import com.google.api.core.InternalApi;
+import com.google.cloud.bigtable.beam.hbasesnapshots.conf.RegionConfig;
+import com.google.cloud.bigtable.beam.hbasesnapshots.conf.SnapshotConfig;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.OutputStream;
+import java.util.Collections;
+import java.util.List;
+import org.apache.beam.sdk.coders.ByteArrayCoder;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.SerializableCoder;
+import org.apache.beam.sdk.coders.VarLongCoder;
+import org.apache.hadoop.hbase.client.RegionInfo;
+import org.apache.hadoop.hbase.client.TableDescriptor;
+import org.apache.hadoop.hbase.shaded.protobuf.ProtobufUtil;
+import org.apache.hadoop.hbase.shaded.protobuf.generated.HBaseProtos;
+import org.apache.hadoop.hbase.shaded.protobuf.generated.HBaseProtos.TableSchema;
+
+/** Implementation of {@link Coder} for encoding and decoding of {@link RegionConfig} */
+@InternalApi("For internal usage only")
+public class RegionConfigCoder extends Coder {
+ private static final VarLongCoder longCoder = VarLongCoder.of();
+
+ private static final Coder snapshotConfigCoder =
+ SerializableCoder.of(SnapshotConfig.class);
+ private static final Coder byteArrayCoder = ByteArrayCoder.of();
+
+ @Override
+ public void encode(RegionConfig value, OutputStream outStream) throws IOException {
+ // 1. Encode SnapshotConfig using standard SerializableCoder
+ snapshotConfigCoder.encode(value.getSnapshotConfig(), outStream);
+
+ // 2. Encode RegionInfo by converting to HBase protobuf and then to byte array
+ HBaseProtos.RegionInfo regionInfo = ProtobufUtil.toRegionInfo(value.getRegionInfo());
+ byteArrayCoder.encode(regionInfo.toByteArray(), outStream);
+
+ // 3. Encode TableDescriptor by converting to HBase protobuf and then to byte array
+ HBaseProtos.TableSchema tableSchema = ProtobufUtil.toTableSchema(value.getTableDescriptor());
+ byteArrayCoder.encode(tableSchema.toByteArray(), outStream);
+
+ // 4. Encode region size using variable-length long coder
+ longCoder.encode(value.getRegionSize(), outStream);
+ }
+
+ @Override
+ public RegionConfig decode(InputStream inStream) throws IOException {
+ // 1. Decode SnapshotConfig
+ SnapshotConfig snapshotConfig = snapshotConfigCoder.decode(inStream);
+
+ // 2. Decode RegionInfo from bytes via HBase protobuf
+ byte[] regionInfoBytes = byteArrayCoder.decode(inStream);
+ RegionInfo regionInfo =
+ ProtobufUtil.toRegionInfo(HBaseProtos.RegionInfo.parseFrom(regionInfoBytes));
+
+ // 3. Decode TableDescriptor from bytes via HBase protobuf
+ byte[] tableSchemaBytes = byteArrayCoder.decode(inStream);
+ TableDescriptor tableDescriptor =
+ ProtobufUtil.toTableDescriptor(TableSchema.parseFrom(tableSchemaBytes));
+
+ // 4. Decode region size
+ Long regionSize = longCoder.decode(inStream);
+
+ return RegionConfig.builder()
+ .setSnapshotConfig(snapshotConfig)
+ .setRegionInfo(regionInfo)
+ .setTableDescriptor(tableDescriptor)
+ .setRegionSize(regionSize)
+ .build();
+ }
+
+ @Override
+ public List> getCoderArguments() {
+ return Collections.emptyList();
+ }
+
+ @Override
+ public void verifyDeterministic() throws Coder.NonDeterministicException {
+ throw new Coder.NonDeterministicException(
+ this,
+ "RegionConfigCoder is non-deterministic because it encodes SnapshotConfig using"
+ + " SerializableCoder.");
+ }
+}
diff --git a/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/coders/package-info.java b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/coders/package-info.java
new file mode 100644
index 0000000000..c24c1f5d50
--- /dev/null
+++ b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/coders/package-info.java
@@ -0,0 +1,21 @@
+/*
+ * Copyright 2026 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.
+ */
+
+/**
+ * Contains coders to handle serialization and deserialization for different classes used in the
+ * pipeline.
+ */
+package com.google.cloud.bigtable.beam.hbasesnapshots.coders;
diff --git a/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/HBaseSnapshotInputConfigBuilder.java b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/conf/HBaseSnapshotInputConfigBuilder.java
similarity index 96%
rename from bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/HBaseSnapshotInputConfigBuilder.java
rename to bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/conf/HBaseSnapshotInputConfigBuilder.java
index 62b7a81ad5..a243f93edc 100644
--- a/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/HBaseSnapshotInputConfigBuilder.java
+++ b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/conf/HBaseSnapshotInputConfigBuilder.java
@@ -13,8 +13,9 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package com.google.cloud.bigtable.beam.hbasesnapshots;
+package com.google.cloud.bigtable.beam.hbasesnapshots.conf;
+import com.google.api.core.InternalApi;
import com.google.common.base.Preconditions;
import org.apache.beam.sdk.extensions.gcp.util.gcsfs.GcsPath;
import org.apache.commons.logging.Log;
@@ -37,7 +38,8 @@
* hosted in Google Cloud Storage(GCS) bucket via GCS connector. It uses {@link
* TableSnapshotInputFormat} for reading HBase snapshots.
*/
-class HBaseSnapshotInputConfigBuilder {
+@InternalApi("For internal usage only")
+public class HBaseSnapshotInputConfigBuilder {
private static final Log LOG = LogFactory.getLog(HBaseSnapshotInputConfigBuilder.class);
// Batch size used for HBase snapshot scans
diff --git a/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/conf/ImportConfig.java b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/conf/ImportConfig.java
new file mode 100644
index 0000000000..9b52579795
--- /dev/null
+++ b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/conf/ImportConfig.java
@@ -0,0 +1,319 @@
+/*
+ * Copyright 2026 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.conf;
+
+import com.google.api.core.InternalApi;
+import com.google.common.base.Preconditions;
+import com.google.gson.TypeAdapter;
+import com.google.gson.annotations.JsonAdapter;
+import com.google.gson.annotations.SerializedName;
+import com.google.gson.stream.JsonReader;
+import com.google.gson.stream.JsonWriter;
+import java.io.IOException;
+import java.io.Serializable;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+/**
+ * Class representing the job configuration loading the different input values and combinations of
+ * snapshot names (such as all snapshots or matching prefix or explicit names) provided.
+ */
+@InternalApi("For internal usage only")
+public final class ImportConfig implements Serializable {
+ private final long DEFAULT_BACKOFF_INITIAL_INTERVAL_MILLIS = 5000; // 5 seconds
+ private final long DEFAULT_BACKOFF_MAX_INTERVAL_MILLIS = 3 * 60 * 1000; // 180 seconds
+ private final int DEFAULT_BACKOFF_MAX_RETRIES = 3;
+
+ @JsonAdapter(SnapshotInfoJsonAdapter.class)
+ @SerializedName("snapshots")
+ public List snapshotInfos;
+
+ private String sourcepath;
+ private String restorepath;
+ private String runstatuspath;
+ private long backoffInitialIntervalInMillis =
+ DEFAULT_BACKOFF_INITIAL_INTERVAL_MILLIS; // Defaults to 5 seconds
+ private long backoffMaxIntervalInMillis = DEFAULT_BACKOFF_MAX_INTERVAL_MILLIS; // 60 seconds
+ private int backoffMaxretries = DEFAULT_BACKOFF_MAX_RETRIES;
+ private Map hbaseConfiguration;
+ private Map bigtableConfiguration;
+
+ public void setSnapshotsFromMap(Map snapshots) {
+ this.snapshotInfos =
+ snapshots.entrySet().stream()
+ .map(entry -> new SnapshotInfo(entry.getKey(), entry.getValue()))
+ .collect(Collectors.toList());
+ }
+
+ public String getSourcepath() {
+ return this.sourcepath;
+ }
+
+ public void setSourcepath(String sourcepath) {
+ this.sourcepath = sourcepath;
+ }
+
+ public String getRestorepath() {
+ return restorepath;
+ }
+
+ public void setRestorepath(String restorepath) {
+ this.restorepath = restorepath;
+ }
+
+ public String getRunstatuspath() {
+ return runstatuspath;
+ }
+
+ public void setRunstatuspath(String runstatuspath) {
+ this.runstatuspath = runstatuspath;
+ }
+
+ public long getBackoffInitialIntervalInMillis() {
+ return backoffInitialIntervalInMillis;
+ }
+
+ public void setBackoffInitialIntervalInMillis(long backoffInitialIntervalInMillis) {
+ Preconditions.checkArgument(
+ backoffInitialIntervalInMillis > 0,
+ "backoffInitialIntervalInMillis must be greater than 0: %s",
+ backoffInitialIntervalInMillis);
+ Preconditions.checkArgument(
+ backoffInitialIntervalInMillis <= this.backoffMaxIntervalInMillis,
+ "backoffInitialIntervalInMillis (%s) must be less than or equal to"
+ + " backoffMaxIntervalInMillis (%s)",
+ backoffInitialIntervalInMillis,
+ this.backoffMaxIntervalInMillis);
+ this.backoffInitialIntervalInMillis = backoffInitialIntervalInMillis;
+ }
+
+ public long getBackoffMaxIntervalInMillis() {
+ return this.backoffMaxIntervalInMillis;
+ }
+
+ public void setBackoffMaxIntervalInMillis(long backoffMaxIntervalInMillis) {
+ Preconditions.checkArgument(
+ backoffMaxIntervalInMillis > 0,
+ "backoffMaxIntervalInMillis must be greater than 0: %s",
+ backoffMaxIntervalInMillis);
+ Preconditions.checkArgument(
+ this.backoffInitialIntervalInMillis <= backoffMaxIntervalInMillis,
+ "backoffInitialIntervalInMillis (%s) must be less than or equal to"
+ + " backoffMaxIntervalInMillis (%s)",
+ this.backoffInitialIntervalInMillis,
+ backoffMaxIntervalInMillis);
+ this.backoffMaxIntervalInMillis = backoffMaxIntervalInMillis;
+ }
+
+ public int getBackoffMaxretries() {
+ return this.backoffMaxretries;
+ }
+
+ public void setBackoffMaxretries(int backoffMaxretries) {
+ Preconditions.checkArgument(
+ backoffMaxretries >= 0, "backoffMaxretries must be non-negative: %s", backoffMaxretries);
+ this.backoffMaxretries = backoffMaxretries;
+ }
+
+ public void validate() {
+ Preconditions.checkNotNull(
+ sourcepath, "Source Path containing hbase snapshots must be specified.");
+ Preconditions.checkArgument(
+ !sourcepath.trim().isEmpty(), "Source Path containing hbase snapshots must be specified.");
+ Preconditions.checkArgument(
+ restorepath == null || !restorepath.trim().isEmpty(),
+ "restorepath cannot be empty if specified");
+ Preconditions.checkArgument(
+ runstatuspath == null || !runstatuspath.trim().isEmpty(),
+ "runstatuspath cannot be empty if specified");
+
+ java.net.URI sourceUri;
+ try {
+ sourceUri = new java.net.URI(sourcepath);
+ } catch (java.net.URISyntaxException e) {
+ throw new IllegalArgumentException("sourcepath is not a valid URI: " + sourcepath, e);
+ }
+
+ if (restorepath != null) {
+ try {
+ new java.net.URI(restorepath);
+ } catch (java.net.URISyntaxException e) {
+ throw new IllegalArgumentException("restorepath is not a valid URI: " + restorepath, e);
+ }
+ }
+
+ if (runstatuspath != null) {
+ try {
+ new java.net.URI(runstatuspath);
+ } catch (java.net.URISyntaxException e) {
+ throw new IllegalArgumentException("runstatuspath is not a valid URI: " + runstatuspath, e);
+ }
+ }
+
+ if (restorepath == null) {
+ if (sourceUri.getScheme() != null && sourceUri.getScheme().equalsIgnoreCase("gs")) {
+ String path = sourceUri.getPath();
+ Preconditions.checkArgument(
+ path != null && !path.isEmpty() && !path.equals("/"),
+ "sourcepath must have a parent directory to auto-generate restorepath: %s",
+ sourcepath);
+ } else {
+ Preconditions.checkArgument(
+ new java.io.File(sourcepath).getParent() != null,
+ "sourcepath must have a parent directory to auto-generate restorepath: %s",
+ sourcepath);
+ }
+ }
+
+ Preconditions.checkNotNull(
+ snapshotInfos,
+ "Snapshots must be specified. Allowed values are '*' (indicating all snapshots under source"
+ + " path) or 'prefix*' (snapshots matching certain prefix) or"
+ + " 'snapshotname1:tablename1,snapshotname2:tablename2' (comma seperated list of"
+ + " snapshots)");
+ Preconditions.checkArgument(
+ !snapshotInfos.isEmpty(),
+ "Snapshots must be specified. Allowed values are '*' (indicating all snapshots under source"
+ + " path) or 'prefix*' (snapshots matching certain prefix) or"
+ + " 'snapshotname1:tablename1,snapshotname2:tablename2' (comma seperated list of"
+ + " snapshots)");
+ java.util.Set snapshotNames = new java.util.HashSet<>();
+ for (SnapshotInfo snapshotInfo : snapshotInfos) {
+ Preconditions.checkArgument(
+ snapshotInfo.getSnapshotName() != null
+ && !snapshotInfo.getSnapshotName().trim().isEmpty(),
+ "snapshotName inside snapshots cannot be null or empty");
+ Preconditions.checkArgument(
+ snapshotNames.add(snapshotInfo.getSnapshotName()),
+ "Duplicate snapshot name detected: %s",
+ snapshotInfo.getSnapshotName());
+ Preconditions.checkArgument(
+ snapshotInfo.getbigtableTableName() != null
+ && !snapshotInfo.getbigtableTableName().trim().isEmpty(),
+ "bigtableTableName inside snapshots cannot be null or empty");
+ }
+ Preconditions.checkArgument(
+ backoffInitialIntervalInMillis > 0,
+ "backoffInitialIntervalInMillis must be greater than 0: %s",
+ backoffInitialIntervalInMillis);
+ Preconditions.checkArgument(
+ backoffMaxIntervalInMillis > 0,
+ "backoffMaxIntervalInMillis must be greater than 0: %s",
+ backoffMaxIntervalInMillis);
+ Preconditions.checkArgument(
+ backoffInitialIntervalInMillis <= backoffMaxIntervalInMillis,
+ "backoffInitialIntervalInMillis (%s) must be less than or equal to"
+ + " backoffMaxIntervalInMillis (%s)",
+ backoffInitialIntervalInMillis,
+ backoffMaxIntervalInMillis);
+ Preconditions.checkArgument(
+ backoffMaxretries >= 0, "backoffMaxretries must be non-negative: %s", backoffMaxretries);
+
+ if (hbaseConfiguration != null) {
+ for (Map.Entry entry : hbaseConfiguration.entrySet()) {
+ Preconditions.checkArgument(
+ entry.getKey() != null && !entry.getKey().trim().isEmpty(),
+ "hbaseConfiguration keys cannot be null or empty");
+ Preconditions.checkArgument(
+ entry.getValue() != null, "hbaseConfiguration values cannot be null");
+ }
+ }
+ if (bigtableConfiguration != null) {
+ for (Map.Entry entry : bigtableConfiguration.entrySet()) {
+ Preconditions.checkArgument(
+ entry.getKey() != null && !entry.getKey().trim().isEmpty(),
+ "bigtableConfiguration keys cannot be null or empty");
+ Preconditions.checkArgument(
+ entry.getValue() != null, "bigtableConfiguration values cannot be null");
+ }
+ }
+ }
+
+ public List getSnapshots() {
+ return this.snapshotInfos;
+ }
+
+ public void setSnapshots(List snapshots) {
+ this.snapshotInfos = snapshots;
+ }
+
+ public Map getHbaseConfiguration() {
+ return this.hbaseConfiguration;
+ }
+
+ public void setHbaseConfiguration(Map hbaseConfiguration) {
+ this.hbaseConfiguration = hbaseConfiguration;
+ }
+
+ public Map getBigtableConfiguration() {
+ return bigtableConfiguration;
+ }
+
+ public void setBigtableConfiguration(Map bigtableConfiguration) {
+ this.bigtableConfiguration = bigtableConfiguration;
+ }
+
+ public static class SnapshotInfo implements Serializable {
+ private final String snapshotName;
+ private final String bigtableTableName;
+
+ public SnapshotInfo(String snapshotName, String tableName) {
+ this.snapshotName = snapshotName;
+ this.bigtableTableName = tableName;
+ }
+
+ public String getSnapshotName() {
+ return snapshotName;
+ }
+
+ public String getbigtableTableName() {
+ return bigtableTableName;
+ }
+ }
+
+ static class SnapshotInfoJsonAdapter extends TypeAdapter> {
+
+ @Override
+ public void write(JsonWriter jsonWriter, List snapshotInfos) throws IOException {
+ jsonWriter.beginObject();
+ if (snapshotInfos != null) {
+ snapshotInfos.forEach(
+ snapshotInfo -> {
+ try {
+ jsonWriter.name(snapshotInfo.getSnapshotName());
+ jsonWriter.value(snapshotInfo.getbigtableTableName());
+ } catch (IOException e) {
+ throw new RuntimeException(e);
+ }
+ });
+ }
+ jsonWriter.endObject();
+ }
+
+ @Override
+ public List read(JsonReader jsonReader) throws IOException {
+ List snapshotInfoList = new ArrayList<>();
+ jsonReader.beginObject();
+ while (jsonReader.hasNext()) {
+ snapshotInfoList.add(new SnapshotInfo(jsonReader.nextName(), jsonReader.nextString()));
+ }
+ jsonReader.endObject();
+ return snapshotInfoList;
+ }
+ }
+}
diff --git a/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/conf/RegionConfig.java b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/conf/RegionConfig.java
new file mode 100644
index 0000000000..8117772a4c
--- /dev/null
+++ b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/conf/RegionConfig.java
@@ -0,0 +1,68 @@
+/*
+ * Copyright 2026 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.conf;
+
+import com.google.api.core.InternalApi;
+import com.google.auto.value.AutoValue;
+import com.google.cloud.bigtable.beam.hbasesnapshots.coders.RegionConfigCoder;
+import javax.annotation.Nullable;
+import org.apache.beam.sdk.coders.DefaultCoder;
+import org.apache.hadoop.hbase.client.RegionInfo;
+import org.apache.hadoop.hbase.client.TableDescriptor;
+
+/**
+ * A {@link AutoValue} class representing the region configuration enclosing {@link SnapshotConfig},
+ * hbase region info and hbase table descriptor.
+ */
+@DefaultCoder(RegionConfigCoder.class)
+@AutoValue
+@InternalApi("For internal usage only")
+public abstract class RegionConfig {
+ public static Builder builder() {
+ return new AutoValue_RegionConfig.Builder();
+ }
+
+ /**
+ * Returns an optional label for this configuration, typically used for testing or debugging. This
+ * is not the official HBase region name.
+ */
+ @Nullable
+ public abstract String getName();
+
+ public abstract SnapshotConfig getSnapshotConfig();
+
+ public abstract RegionInfo getRegionInfo();
+
+ public abstract TableDescriptor getTableDescriptor();
+
+ public abstract Long getRegionSize();
+
+ @AutoValue.Builder
+ public abstract static class Builder {
+
+ public abstract Builder setName(String value);
+
+ public abstract Builder setSnapshotConfig(SnapshotConfig value);
+
+ public abstract Builder setRegionInfo(RegionInfo value);
+
+ public abstract Builder setTableDescriptor(TableDescriptor value);
+
+ public abstract Builder setRegionSize(Long value);
+
+ public abstract RegionConfig build();
+ }
+}
diff --git a/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/conf/SnapshotConfig.java b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/conf/SnapshotConfig.java
new file mode 100644
index 0000000000..ce2bc0b1da
--- /dev/null
+++ b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/conf/SnapshotConfig.java
@@ -0,0 +1,100 @@
+/*
+ * Copyright 2026 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.conf;
+
+import com.google.api.core.InternalApi;
+import com.google.auto.value.AutoValue;
+import com.google.cloud.bigtable.beam.hbasesnapshots.SnapshotUtils;
+import java.io.Serializable;
+import java.util.Map;
+import org.apache.beam.sdk.schemas.AutoValueSchema;
+import org.apache.beam.sdk.schemas.annotations.DefaultSchema;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.Path;
+
+/** A {@link AutoValue} class representing the configuration associated with each snapshot. */
+@InternalApi("For internal usage only")
+@DefaultSchema(AutoValueSchema.class)
+@AutoValue
+public abstract class SnapshotConfig implements Serializable {
+
+ /** Creates a new builder for {@link SnapshotConfig}. */
+ public static Builder builder() {
+ return new AutoValue_SnapshotConfig.Builder();
+ }
+
+ /** Returns the Google Cloud project ID. */
+ public abstract String getProjectId();
+
+ /** Returns the GCS source location of the HBase snapshot. */
+ public abstract String getSourceLocation();
+
+ /** Returns the GCS source path as a Hadoop {@link Path}. */
+ public Path getSourcePath() {
+ return new Path(getSourceLocation());
+ }
+
+ /** Returns the GCS restore path as a Hadoop {@link Path}. */
+ public Path getRestorePath() {
+ return new Path(getRestoreLocation());
+ }
+
+ /** Returns the name of the HBase snapshot. */
+ public abstract String getSnapshotName();
+
+ /** Returns the name of the Bigtable table to import into. */
+ public abstract String getTableName();
+
+ /** Returns the GCS location where the snapshot is restored temporarily. */
+ public abstract String getRestoreLocation();
+
+ /** Returns the additional configuration details as a map. */
+ public abstract Map getConfigurationDetails();
+
+ /** Returns the Hadoop {@link Configuration} derived from configuration details. */
+ public Configuration getConfiguration() {
+ return SnapshotUtils.getHBaseConfiguration(getConfigurationDetails());
+ }
+
+ /** Returns a builder initialized with the current values of this instance. */
+ public abstract Builder toBuilder();
+
+ /** Builder for {@link SnapshotConfig}. */
+ @AutoValue.Builder
+ public abstract static class Builder {
+
+ /** Sets the Google Cloud project ID. */
+ public abstract Builder setProjectId(String projectId);
+
+ /** Sets the GCS source location of the HBase snapshot. */
+ public abstract Builder setSourceLocation(String value);
+
+ /** Sets the name of the HBase snapshot. */
+ public abstract Builder setSnapshotName(String value);
+
+ /** Sets the name of the Bigtable table to import into. */
+ public abstract Builder setTableName(String value);
+
+ /** Sets the GCS location where the snapshot is restored temporarily. */
+ public abstract Builder setRestoreLocation(String value);
+
+ /** Sets the additional configuration details. */
+ public abstract Builder setConfigurationDetails(Map configuration);
+
+ /** Builds the {@link SnapshotConfig} instance. */
+ public abstract SnapshotConfig build();
+ }
+}
diff --git a/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/conf/SnapshotKey.java b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/conf/SnapshotKey.java
new file mode 100644
index 0000000000..64b7059d17
--- /dev/null
+++ b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/conf/SnapshotKey.java
@@ -0,0 +1,41 @@
+/*
+ * Copyright 2026 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.conf;
+
+import com.google.api.core.InternalApi;
+import com.google.auto.value.AutoValue;
+import java.io.Serializable;
+import org.apache.beam.sdk.schemas.AutoValueSchema;
+import org.apache.beam.sdk.schemas.annotations.DefaultSchema;
+import org.apache.beam.sdk.schemas.annotations.SchemaCreate;
+
+/**
+ * A lightweight key representing a snapshot and its target table, used to avoid serializing full
+ * configuration details for every row.
+ */
+@InternalApi("For internal usage only")
+@DefaultSchema(AutoValueSchema.class)
+@AutoValue
+public abstract class SnapshotKey implements Serializable {
+ @SchemaCreate
+ public static SnapshotKey create(String snapshotName, String tableName) {
+ return new AutoValue_SnapshotKey(snapshotName, tableName);
+ }
+
+ public abstract String getSnapshotName();
+
+ public abstract String getTableName();
+}
diff --git a/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/conf/package-info.java b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/conf/package-info.java
new file mode 100644
index 0000000000..5a7bb988eb
--- /dev/null
+++ b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/conf/package-info.java
@@ -0,0 +1,17 @@
+/*
+ * Copyright 2026 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 contains configuration classes used in the pipeline. */
+package com.google.cloud.bigtable.beam.hbasesnapshots.conf;
diff --git a/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/CleanupHBaseSnapshotRestoreFilesFn.java b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/dofn/CleanupHBaseSnapshotRestoreFilesFn.java
similarity index 93%
rename from bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/CleanupHBaseSnapshotRestoreFilesFn.java
rename to bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/dofn/CleanupHBaseSnapshotRestoreFilesFn.java
index e0bdca69d5..ea65781bdf 100644
--- a/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/CleanupHBaseSnapshotRestoreFilesFn.java
+++ b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/dofn/CleanupHBaseSnapshotRestoreFilesFn.java
@@ -13,8 +13,9 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package com.google.cloud.bigtable.beam.hbasesnapshots;
+package com.google.cloud.bigtable.beam.hbasesnapshots.dofn;
+import com.google.api.core.InternalApi;
import com.google.api.services.storage.model.Objects;
import com.google.common.base.Preconditions;
import java.io.IOException;
@@ -32,9 +33,27 @@
* 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, Boolean> {
+@InternalApi("For internal usage only")
+public class CleanupHBaseSnapshotRestoreFilesFn extends DoFn, Boolean> {
private static final Log LOG = LogFactory.getLog(CleanupHBaseSnapshotRestoreFilesFn.class);
+ 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);
+ }
+
@ProcessElement
public void processElement(ProcessContext context) throws IOException {
KV elem = context.element();
@@ -65,20 +84,4 @@ public void processElement(ProcessContext context) throws IOException {
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);
- }
}
diff --git a/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/dofn/CleanupRestoredSnapshotsFn.java b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/dofn/CleanupRestoredSnapshotsFn.java
new file mode 100644
index 0000000000..35332114e2
--- /dev/null
+++ b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/dofn/CleanupRestoredSnapshotsFn.java
@@ -0,0 +1,109 @@
+/*
+ * Copyright 2026 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.dofn;
+
+import com.google.api.core.InternalApi;
+import com.google.cloud.bigtable.beam.hbasesnapshots.conf.SnapshotConfig;
+import com.google.common.annotations.VisibleForTesting;
+import java.io.IOException;
+import org.apache.beam.sdk.metrics.Counter;
+import org.apache.beam.sdk.metrics.Metrics;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.util.BackOff;
+import org.apache.beam.sdk.util.FluentBackoff;
+import org.apache.beam.sdk.util.Sleeper;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.FileSystem;
+import org.apache.hadoop.fs.Path;
+import org.joda.time.Duration;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** A {@link DoFn} for cleaning up files from restore path generated during job run. */
+@InternalApi("For internal usage only")
+public class CleanupRestoredSnapshotsFn extends DoFn {
+
+ private static final Logger LOG = LoggerFactory.getLogger(CleanupRestoredSnapshotsFn.class);
+ private final Counter failedCleanups =
+ Metrics.counter(CleanupRestoredSnapshotsFn.class, "failed_cleanups");
+
+ private final long initialBackoffMillis;
+ private final long maxBackoffMillis;
+ private final int maxRetries;
+
+ public CleanupRestoredSnapshotsFn(
+ long initialBackoffMillis, long maxBackoffMillis, int maxRetries) {
+ this.initialBackoffMillis = initialBackoffMillis;
+ this.maxBackoffMillis = maxBackoffMillis;
+ this.maxRetries = maxRetries;
+ }
+
+ @ProcessElement
+ public void processElement(
+ @Element SnapshotConfig snapshotConfig, OutputReceiver outputReceiver)
+ throws IOException {
+
+ Path restorePath = snapshotConfig.getRestorePath();
+ Configuration configuration = snapshotConfig.getConfiguration();
+ FileSystem fileSystem = restorePath.getFileSystem(configuration);
+
+ FluentBackoff backoff =
+ FluentBackoff.DEFAULT
+ .withInitialBackoff(Duration.millis(initialBackoffMillis))
+ .withMaxBackoff(Duration.millis(maxBackoffMillis))
+ .withMaxRetries(maxRetries);
+
+ Sleeper sleeper = getSleeper();
+ BackOff executionBackoff = backoff.backoff();
+
+ while (true) {
+ try {
+ cleanupSnapshot(snapshotConfig, fileSystem, restorePath);
+ return; // Success
+ } catch (IOException ex) {
+ long nextSleep = executionBackoff.nextBackOffMillis();
+
+ if (nextSleep == BackOff.STOP) {
+ LOG.error(
+ "Failed to cleanup snapshot after retries. Manual cleanup required for path: {}",
+ restorePath,
+ ex);
+ failedCleanups.inc();
+ return; // Give up but don't fail the job
+ }
+
+ LOG.warn("Cleanup failed, retrying in {} ms. Error: {}", nextSleep, ex.getMessage());
+ try {
+ sleeper.sleep(nextSleep);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new IOException("Interrupted during retry sleep", e);
+ }
+ }
+ }
+ }
+
+ @VisibleForTesting
+ Sleeper getSleeper() {
+ return Sleeper.DEFAULT;
+ }
+
+ @VisibleForTesting
+ void cleanupSnapshot(SnapshotConfig snapshotConfig, FileSystem fileSystem, Path restorePath)
+ throws IOException {
+ fileSystem.delete(restorePath, true);
+ }
+}
diff --git a/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/dofn/CreateBigtableMutationsFn.java b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/dofn/CreateBigtableMutationsFn.java
new file mode 100644
index 0000000000..9e9cf36e82
--- /dev/null
+++ b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/dofn/CreateBigtableMutationsFn.java
@@ -0,0 +1,232 @@
+/*
+ * Copyright 2026 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.dofn;
+
+import com.google.api.core.InternalApi;
+import com.google.cloud.bigtable.beam.hbasesnapshots.conf.SnapshotKey;
+import com.google.common.base.Preconditions;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import org.apache.beam.sdk.metrics.Counter;
+import org.apache.beam.sdk.metrics.Metrics;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.values.KV;
+import org.apache.hadoop.hbase.Cell;
+import org.apache.hadoop.hbase.CellUtil;
+import org.apache.hadoop.hbase.client.Mutation;
+import org.apache.hadoop.hbase.client.Put;
+import org.apache.hadoop.hbase.client.Result;
+import org.apache.hadoop.hbase.util.Bytes;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** A {@link DoFn} class for converting Hbase {@link Result} to list of Hbase {@link Mutation}s. */
+@InternalApi("For internal usage only")
+public class CreateBigtableMutationsFn
+ extends DoFn, KV>> {
+ private static final Logger LOG = LoggerFactory.getLogger(CreateBigtableMutationsFn.class);
+
+ private final Counter droppedRows =
+ Metrics.counter(CreateBigtableMutationsFn.class, "droppedRows");
+ private final Counter droppedCells =
+ Metrics.counter(CreateBigtableMutationsFn.class, "droppedCells");
+ private final Counter droppedRowKeys =
+ Metrics.counter(CreateBigtableMutationsFn.class, "droppedRowKeys");
+ private final int maxMutationsPerRequestThreshold;
+
+ private final boolean filterLargeRows;
+ private final long filterLargeRowThresholdBytes;
+
+ private final boolean filterLargeCells;
+ private final int filterLargeCellThresholdBytes;
+
+ private final boolean filterLargeRowKeys;
+ private final int filterLargeRowKeysThresholdBytes;
+ private final boolean filterWideRows;
+
+ public CreateBigtableMutationsFn(
+ int maxMutationsPerRequestThreshold,
+ boolean filterLargeRows,
+ long filterLargeRowThresholdBytes,
+ boolean filterLargeCells,
+ int filterLargeCellThresholdBytes,
+ boolean filterLargeRowKeys,
+ int filterLargeRowKeysThresholdBytes,
+ boolean filterWideRows) {
+
+ Preconditions.checkArgument(
+ maxMutationsPerRequestThreshold > 0,
+ "maxMutationsPerRequestThreshold must be greater than 0");
+ this.maxMutationsPerRequestThreshold = maxMutationsPerRequestThreshold;
+
+ this.filterLargeRows = filterLargeRows;
+ this.filterLargeRowThresholdBytes = filterLargeRowThresholdBytes;
+
+ this.filterLargeCells = filterLargeCells;
+ this.filterLargeCellThresholdBytes = filterLargeCellThresholdBytes;
+
+ this.filterLargeRowKeys = filterLargeRowKeys;
+ this.filterLargeRowKeysThresholdBytes = filterLargeRowKeysThresholdBytes;
+ this.filterWideRows = filterWideRows;
+ }
+
+ @ProcessElement
+ public void processElement(
+ @Element KV element,
+ OutputReceiver>> outputReceiver)
+ throws IOException {
+ if (element == null
+ || element.getKey() == null
+ || element.getValue() == null
+ || element.getValue().isEmpty()) {
+ return;
+ }
+
+ // Extract metadata for routing and logging.
+ String targetTable = element.getKey().getTableName();
+ String snapshotName = element.getKey().getSnapshotName();
+ byte[] rowKey = element.getValue().getRow();
+
+ // Apply filters and chunk cells into Mutations.
+ // Returns null if the entire row should be dropped.
+ List mutations =
+ convertAndValidateThresholds(rowKey, element.getValue().listCells(), snapshotName);
+
+ // Output the mutations mapped to the target table name.
+ if (mutations != null && !mutations.isEmpty()) {
+ outputReceiver.output(KV.of(targetTable, mutations));
+ }
+ }
+
+ /**
+ * Converts a list of HBase cells for a specific row into a list of Bigtable mutations. It also
+ * applies various filtering rules based on row key size, cell size, total row size, and wide row
+ * thresholds.
+ *
+ * The process follows these steps: 1. Validate row key size. 2. Iterate over cells and filter
+ * large cells if enabled. 3. Accumulate total row size and filter large rows if enabled. 4. Chunk
+ * cells into multiple Put requests if they exceed the mutation count threshold.
+ *
+ *
Note on limits: When specific filters (e.g., large row/cell/row key filtering) are disabled
+ * or not met, this method does not fail-fast locally. The decision to reject requests that
+ * violate Bigtable constraints (such as payload size or mutation limits) is deferred to the
+ * underlying Cloud Bigtable client and server to avoid duplicating validation logic.
+ *
+ * @param rowKey the row key of the row being processed
+ * @param cells the list of cells for this row
+ * @param snapshotName the name of the snapshot for logging purposes
+ * @return a list of Mutations (Puts), or null if the entire row should be dropped
+ * @throws IOException if an error occurs
+ */
+ private List convertAndValidateThresholds(
+ byte[] rowKey, List| cells, String snapshotName) throws IOException {
+
+ // 1. Check row key size. We do this first to fail-fast and avoid allocations,
+ // since the row key size is constant and independent of cell processing.
+ if (filterLargeRowKeys && rowKey.length > filterLargeRowKeysThresholdBytes) {
+ LOG.warn(
+ "For snapshot "
+ + snapshotName
+ + ": Dropping row, row key length, "
+ + rowKey.length
+ + ", exceeds filter length threshold, "
+ + filterLargeRowKeysThresholdBytes
+ + ", row key: "
+ + Bytes.toStringBinary(rowKey));
+ droppedRowKeys.inc();
+ return null; // Signal skip
+ }
+
+ List mutations = new ArrayList<>();
+ Put put = null;
+ int chunkCellCount = 0; // Tracks cells in the current Put
+ int totalCellCount = 0; // Tracks total cells in the row
+ long totalByteSize = 0L; // Tracks estimated serialized size of cells in the current row
+ boolean loggedLargeCellForRow = false;
+
+ // Iterate over all cells in the row to build mutations. We apply cell-level and
+ // row-level filters during iteration, and split the row into multiple Put requests
+ // if the number of cells exceeds maxMutationsPerRequestThreshold.
+ for (Cell cell : cells) {
+
+ // 2. Handle large cells first. We do this before accumulating row size so that
+ // cells we intend to drop anyway don't unfairly cause the entire row to exceed
+ // the row size threshold.
+ if (filterLargeCells && cell.getValueLength() > filterLargeCellThresholdBytes) {
+ if (!loggedLargeCellForRow) {
+ LOG.warn(
+ "For snapshot "
+ + snapshotName
+ + ": Dropping large cells for row "
+ + Bytes.toStringBinary(rowKey)
+ + ". At least one cell exceeds threshold "
+ + filterLargeCellThresholdBytes);
+ loggedLargeCellForRow = true;
+ }
+ droppedCells.inc();
+ continue; // Skip this cell, do NOT add to totalByteSize or mutations
+ }
+
+ // 3. Accumulate size and check row size.
+ // Drops the entire row if it exceeds the row size threshold.
+ totalByteSize += CellUtil.estimatedSerializedSizeOf(cell);
+ if (filterLargeRows && totalByteSize > filterLargeRowThresholdBytes) {
+ LOG.warn(
+ "For snapshot "
+ + snapshotName
+ + ": Dropping row, row length, "
+ + totalByteSize
+ + ", exceeds filter length threshold, "
+ + filterLargeRowThresholdBytes
+ + ", row key: "
+ + Bytes.toStringBinary(rowKey));
+ droppedRows.inc();
+ return null; // Signal skip for the entire row
+ }
+
+ // 4. Chunk cells into multiple Put requests to avoid exceeding Bigtable's limit
+ // of 100,000 mutations per request. If filterWideRows is enabled and the total
+ // cells in the row exceed the threshold, we drop the entire row to avoid loss
+ // of atomicity.
+ if (chunkCellCount == maxMutationsPerRequestThreshold || chunkCellCount == 0) {
+ if (totalCellCount > 0 && filterWideRows) {
+ LOG.warn(
+ "For snapshot "
+ + snapshotName
+ + ": Dropping wide row, cell count exceeds threshold "
+ + maxMutationsPerRequestThreshold
+ + ", row key: "
+ + Bytes.toStringBinary(rowKey));
+ droppedRows.inc();
+ return null; // Signal skip
+ }
+ chunkCellCount = 0;
+ put = new Put(rowKey);
+ mutations.add(put);
+ }
+ put.add(cell);
+ chunkCellCount++;
+ totalCellCount++;
+ }
+
+ if (mutations.isEmpty()) {
+ droppedRows.inc();
+ return null;
+ }
+ return mutations;
+ }
+}
diff --git a/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/dofn/HBaseRegionScanner.java b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/dofn/HBaseRegionScanner.java
new file mode 100644
index 0000000000..63ba782a02
--- /dev/null
+++ b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/dofn/HBaseRegionScanner.java
@@ -0,0 +1,158 @@
+/*
+ * Copyright 2026 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.dofn;
+
+import com.google.api.core.InternalApi;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.FileSystem;
+import org.apache.hadoop.fs.Path;
+import org.apache.hadoop.hbase.Cell;
+import org.apache.hadoop.hbase.client.IsolationLevel;
+import org.apache.hadoop.hbase.client.RegionInfo;
+import org.apache.hadoop.hbase.client.Result;
+import org.apache.hadoop.hbase.client.Scan;
+import org.apache.hadoop.hbase.client.TableDescriptor;
+import org.apache.hadoop.hbase.client.TableDescriptorBuilder;
+import org.apache.hadoop.hbase.regionserver.HRegion;
+import org.apache.hadoop.hbase.regionserver.RegionScanner;
+import org.apache.hadoop.hbase.regionserver.RegionServerServices;
+import org.apache.hadoop.hbase.util.CommonFSUtils;
+import org.apache.hadoop.hbase.wal.WAL;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A workalike for {@link org.apache.hadoop.hbase.client.ClientSideRegionScanner}.
+ *
+ * It serves the same purpose, but skips block and mobFile cache initialization. Those caches
+ * dont appear to useful for the import job and leak threads on shutdown
+ */
+@InternalApi("For internal usage only")
+public class HBaseRegionScanner implements AutoCloseable {
+ private static final Logger LOG = LoggerFactory.getLogger(HBaseRegionScanner.class);
+
+ private static final String DEFAULT_HFILE_BLOCK_CACHE_POLICY = "IndexOnlyLRU";
+ private static final long DEFAULT_HFILE_BLOCK_CACHE_SIZE = 33554432L; // 32 MB
+ private static final int DEFAULT_COMPACTION_THRESHOLD = 10000;
+ private static final long DEFAULT_CACHE_FLUSH_INTERVAL = 0;
+ private static final int DEFAULT_CLIENT_RETRIES = 3;
+
+ private HRegion region;
+ private RegionScanner scanner;
+ private final List cells;
+ private boolean regionOperationStarted = false;
+ private boolean hasMore = true;
+
+ public HBaseRegionScanner(
+ Configuration originalConf,
+ FileSystem fs,
+ Path rootDir,
+ TableDescriptor htd,
+ RegionInfo hri,
+ Scan scan)
+ throws IOException {
+ Configuration conf = new Configuration(originalConf);
+ conf.set("hfile.block.cache.policy", DEFAULT_HFILE_BLOCK_CACHE_POLICY);
+ // Set a small default block cache size (32MB). Since we are performing a strictly
+ // sequential scan of the snapshot, we read each block once and don't need a large
+ // cache for data blocks. 32MB is sufficient to hold HFile index blocks while avoiding
+ // OutOfMemory errors on memory-constrained Dataflow workers.
+ conf.setIfUnset(
+ "hfile.onheap.block.cache.fixed.size", String.valueOf(DEFAULT_HFILE_BLOCK_CACHE_SIZE));
+ conf.unset("hbase.bucketcache.ioengine");
+ // Setting a huge compaction threshold (10000) effectively disables compactions.
+ // Since snapshots are read-only, compactions are useless, and background threads
+ // can fail to shut down cleanly during DoFn teardown, causing thread leaks.
+ conf.setInt("hbase.hstore.compactionThreshold", DEFAULT_COMPACTION_THRESHOLD);
+ // Set flush interval to 0 to disable periodic flushes.
+ conf.setLong("hbase.regionserver.optionalcacheflushinterval", DEFAULT_CACHE_FLUSH_INTERVAL);
+ conf.setInt("hbase.client.retries.number", DEFAULT_CLIENT_RETRIES);
+
+ scan.setIsolationLevel(IsolationLevel.READ_UNCOMMITTED);
+ htd = TableDescriptorBuilder.newBuilder(htd).setReadOnly(true).build();
+ this.region =
+ HRegion.newHRegion(
+ CommonFSUtils.getTableDir(rootDir, htd.getTableName()),
+ (WAL) null,
+ fs,
+ conf,
+ hri,
+ htd,
+ (RegionServerServices) null);
+ this.region.setRestoredRegion(true);
+
+ // Wrap in try-catch to ensure close() is called on failure, avoiding leaks.
+ try {
+ this.region.initialize();
+ this.scanner = this.region.getScanner(scan);
+ this.cells = new ArrayList<>();
+
+ this.region.startRegionOperation();
+ this.regionOperationStarted = true;
+ } catch (Throwable t) {
+ LOG.error("Failed to initialize HBaseRegionScanner", t);
+ close();
+ if (t instanceof IOException) {
+ throw (IOException) t;
+ }
+ throw new IOException("Failed to initialize HBaseRegionScanner", t);
+ }
+ }
+
+ @Override
+ public void close() {
+ if (this.scanner != null) {
+ try {
+ this.scanner.close();
+ this.scanner = null;
+ } catch (IOException var3) {
+ LOG.warn("Exception while closing scanner", var3);
+ }
+ }
+
+ if (this.region != null) {
+ try {
+ if (this.regionOperationStarted) {
+ try {
+ this.region.closeRegionOperation();
+ } catch (IOException e) {
+ LOG.warn("Exception while closing region operation", e);
+ }
+ }
+ this.region.close(true);
+ this.region = null;
+ } catch (IOException var2) {
+ LOG.warn("Exception while closing region", var2);
+ }
+ }
+ }
+
+ public Result next() throws IOException {
+ while (this.hasMore) {
+ this.cells.clear();
+ this.hasMore = this.scanner.nextRaw(this.cells);
+
+ if (!this.cells.isEmpty()) {
+ return Result.create(this.cells);
+ }
+ }
+
+ return null;
+ }
+}
diff --git a/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/dofn/ReadSnapshotRegionFn.java b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/dofn/ReadSnapshotRegionFn.java
new file mode 100644
index 0000000000..83a212e776
--- /dev/null
+++ b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/dofn/ReadSnapshotRegionFn.java
@@ -0,0 +1,219 @@
+/*
+ * Copyright 2026 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.dofn;
+
+import com.google.api.core.InternalApi;
+import com.google.cloud.bigtable.beam.hbasesnapshots.SnapshotUtils;
+import com.google.cloud.bigtable.beam.hbasesnapshots.conf.RegionConfig;
+import com.google.cloud.bigtable.beam.hbasesnapshots.conf.SnapshotConfig;
+import com.google.cloud.bigtable.beam.hbasesnapshots.conf.SnapshotKey;
+import com.google.cloud.bigtable.beam.hbasesnapshots.transforms.HbaseRegionSplitTracker;
+import java.util.HashMap;
+import java.util.Map;
+import org.apache.beam.sdk.io.range.ByteKey;
+import org.apache.beam.sdk.io.range.ByteKeyRange;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.display.DisplayData;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.values.KV;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.FileSystem;
+import org.apache.hadoop.fs.Path;
+import org.apache.hadoop.hbase.client.IsolationLevel;
+import org.apache.hadoop.hbase.client.Result;
+import org.apache.hadoop.hbase.client.Scan;
+import org.apache.hadoop.hbase.util.RegionSplitter;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** A Splittable {@link DoFn} for reading the records from each region. */
+@InternalApi("For internal usage only")
+public class ReadSnapshotRegionFn extends DoFn> {
+ private static final Logger LOG = LoggerFactory.getLogger(ReadSnapshotRegionFn.class);
+
+ private static final long BYTES_PER_SPLIT = 512 * 1024 * 1024; // 512 MB
+ private static final long BYTES_PER_GB = 1024 * 1024 * 1024;
+
+ private final boolean useDynamicSplitting;
+ private transient Map, Configuration> configCache;
+
+ public ReadSnapshotRegionFn(boolean useDynamicSplitting) {
+ this.useDynamicSplitting = useDynamicSplitting;
+ }
+
+ // Beam reuses DoFn instances for multiple elements. Initializing the cache in @Setup
+ // ensures that we only create the cache once per DoFn instance lifecycle (per worker thread),
+ // avoiding heavy XML parsing overhead for Configuration while also avoiding static state
+ // and ensuring thread safety since Beam isolates DoFn instances.
+ @Setup
+ public void setup() {
+ configCache = new HashMap<>();
+ }
+
+ @ProcessElement
+ public void processElement(
+ @Element RegionConfig regionConfig,
+ OutputReceiver> outputReceiver,
+ RestrictionTracker tracker)
+ throws Exception {
+
+ Map configDetails = regionConfig.getSnapshotConfig().getConfigurationDetails();
+ Configuration configuration =
+ configCache.computeIfAbsent(configDetails, SnapshotUtils::getHBaseConfiguration);
+
+ SnapshotConfig sc = regionConfig.getSnapshotConfig();
+ SnapshotKey snapshotKey = SnapshotKey.create(sc.getSnapshotName(), sc.getTableName());
+
+ // Use try-with-resources to ensure scanner is closed and resources are released.
+ try (HBaseRegionScanner scanner =
+ newScanner(regionConfig, tracker.currentRestriction(), configuration)) {
+ for (Result result = scanner.next(); result != null; result = scanner.next()) {
+ if (tracker.tryClaim(ByteKey.copyFrom(result.getRow()))) {
+ outputReceiver.output(KV.of(snapshotKey, result));
+ } else {
+ // According to Splittable DoFn contract, when tryClaim returns false,
+ // we must terminate processing immediately and return.
+ return;
+ }
+ }
+ // Signal completion of the range. ByteKeyRangeTracker uses EMPTY to mark the end of the
+ // range.
+ // See:
+ // https://github.com/apache/beam/blob/2c4d2c6de4dca6b5954c529c2d40c031a8b74f60/sdks/java/core/src/main/java/org/apache/beam/sdk/transforms/splittabledofn/ByteKeyRangeTracker.java#L37
+ tracker.tryClaim(ByteKey.EMPTY);
+ }
+ }
+
+ HBaseRegionScanner newScanner(
+ RegionConfig regionConfig, ByteKeyRange byteKeyRange, Configuration configuration)
+ throws Exception {
+ // Create an HBase Scan bounded by the restriction's start and end keys.
+ // Use READ_UNCOMMITTED for performance since snapshots are read-only.
+ // Disable block caching as we are doing a full scan and reuse is unlikely.
+ Scan scan =
+ new Scan()
+ .withStartRow(byteKeyRange.getStartKey().getBytes())
+ .withStopRow(byteKeyRange.getEndKey().getBytes())
+ .setIsolationLevel(IsolationLevel.READ_UNCOMMITTED)
+ .setCacheBlocks(false);
+
+ SnapshotConfig snapshotConfig = regionConfig.getSnapshotConfig();
+
+ Path sourcePath = snapshotConfig.getSourcePath();
+ Path restorePath = snapshotConfig.getRestorePath();
+ FileSystem fileSystem = sourcePath.getFileSystem(configuration);
+
+ return new HBaseRegionScanner(
+ configuration,
+ fileSystem,
+ restorePath,
+ regionConfig.getTableDescriptor(),
+ regionConfig.getRegionInfo(),
+ scan);
+ }
+
+ @GetInitialRestriction
+ public ByteKeyRange getInitialRange(@Element RegionConfig regionConfig) {
+ // The initial restriction is the full key range of the HBase region.
+ return ByteKeyRange.of(
+ ByteKey.copyFrom(regionConfig.getRegionInfo().getStartKey()),
+ ByteKey.copyFrom(regionConfig.getRegionInfo().getEndKey()));
+ }
+
+ @GetSize
+ public double getSize(@Element RegionConfig regionConfig) {
+ // Return the default split size as the work estimate for this restriction.
+ return BYTES_PER_SPLIT;
+ }
+
+ @NewTracker
+ public HbaseRegionSplitTracker newTracker(
+ @Element RegionConfig regionConfig, @Restriction ByteKeyRange range) {
+ return new HbaseRegionSplitTracker(
+ regionConfig.getSnapshotConfig().getSnapshotName(),
+ regionConfig.getRegionInfo().getEncodedName(),
+ range,
+ useDynamicSplitting);
+ }
+
+ @SplitRestriction
+ public void splitRestriction(
+ @Element RegionConfig regionConfig,
+ @Restriction ByteKeyRange range,
+ OutputReceiver outputReceiver) {
+ byte[] originalEndKey = regionConfig.getRegionInfo().getEndKey();
+ if (originalEndKey == null || originalEndKey.length == 0) {
+ LOG.info(
+ "Skipping splitting for boundary region: {}",
+ regionConfig.getRegionInfo().getEncodedName());
+ outputReceiver.output(range);
+ return;
+ }
+
+ try {
+ int numSplits = getSplits(regionConfig.getRegionSize());
+ LOG.info(
+ "Splitting Initial Restriction for SnapshotName: {} - regionname:{} - regionsize(GB):{}"
+ + " - Splits: {}",
+ regionConfig.getSnapshotConfig().getSnapshotName(),
+ regionConfig.getRegionInfo().getEncodedName(),
+ (double) regionConfig.getRegionSize() / BYTES_PER_GB,
+ numSplits);
+ if (numSplits > 1) {
+ RegionSplitter.UniformSplit uniformSplit = new RegionSplitter.UniformSplit();
+ byte[] startKey = range.getStartKey().getBytes();
+ byte[] endKey = range.getEndKey().getBytes();
+
+ // Handle empty start key if it's the absolute first region
+ if (startKey.length == 0) {
+ startKey = new byte[endKey.length];
+ }
+
+ byte[][] splits = uniformSplit.split(startKey, endKey, numSplits, true);
+
+ // Preserve the absolute start boundary if necessary
+ if (range.getStartKey().isEmpty()) {
+ splits[0] = new byte[0];
+ }
+
+ for (int i = 0; i < splits.length - 1; i++) {
+ outputReceiver.output(
+ ByteKeyRange.of(ByteKey.copyFrom(splits[i]), ByteKey.copyFrom(splits[i + 1])));
+ }
+ } else {
+ outputReceiver.output(range);
+ }
+ } catch (Exception ex) {
+ LOG.warn(
+ "Unable to split range for region:{} in Snapshot:{}",
+ regionConfig.getRegionInfo().getEncodedName(),
+ regionConfig.getSnapshotConfig().getSnapshotName(),
+ ex);
+ outputReceiver.output(range);
+ }
+ }
+
+ private int getSplits(long sizeInBytes) {
+ return (int) Math.ceil((double) sizeInBytes / BYTES_PER_SPLIT);
+ }
+
+ @Override
+ public void populateDisplayData(DisplayData.Builder builder) {
+ builder.add(
+ DisplayData.item("DynamicSplitting", useDynamicSplitting ? "Enabled" : "Disabled")
+ .withLabel("Dynamic Splitting"));
+ }
+}
diff --git a/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/dofn/RestoreSnapshotFn.java b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/dofn/RestoreSnapshotFn.java
new file mode 100644
index 0000000000..8eff6a7196
--- /dev/null
+++ b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/dofn/RestoreSnapshotFn.java
@@ -0,0 +1,65 @@
+/*
+ * Copyright 2026 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.dofn;
+
+import com.google.api.core.InternalApi;
+import com.google.cloud.bigtable.beam.hbasesnapshots.conf.SnapshotConfig;
+import java.io.IOException;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.FileSystem;
+import org.apache.hadoop.fs.Path;
+import org.apache.hadoop.hbase.snapshot.RestoreSnapshotHelper;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link DoFn} for preprocessing the Snapshot files to restore HLinks & References prior to
+ * reading the snapshot.
+ */
+@InternalApi("For internal usage only")
+public class RestoreSnapshotFn extends DoFn {
+ private static final Logger LOG = LoggerFactory.getLogger(RestoreSnapshotFn.class);
+
+ @ProcessElement
+ public void processElement(
+ @Element SnapshotConfig snapshotConfig, OutputReceiver outputReceiver)
+ throws IOException {
+ restoreSnapshot(snapshotConfig);
+ outputReceiver.output(snapshotConfig);
+ }
+
+ /**
+ * Creates a copy of Snasphsot from the source path into restore path.
+ *
+ * @param snapshotConfig - Snapshot Configuration
+ * @throws IOException
+ */
+ void restoreSnapshot(SnapshotConfig snapshotConfig) throws IOException {
+ Path sourcePath = snapshotConfig.getSourcePath();
+ Path restorePath = snapshotConfig.getRestorePath();
+ Configuration configuration = snapshotConfig.getConfiguration();
+ LOG.info("RestoreSnapshot - sourcePath:{} restorePath: {}", sourcePath, restorePath);
+ FileSystem fileSystem = sourcePath.getFileSystem(configuration);
+ // Delete restore path if it exists to ensure idempotency on bundle retries.
+ if (fileSystem.exists(restorePath)) {
+ LOG.info("Restore path {} already exists, deleting it for idempotency", restorePath);
+ fileSystem.delete(restorePath, true);
+ }
+ RestoreSnapshotHelper.copySnapshotForScanner(
+ configuration, fileSystem, sourcePath, restorePath, snapshotConfig.getSnapshotName());
+ }
+}
diff --git a/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/dofn/package-info.java b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/dofn/package-info.java
new file mode 100644
index 0000000000..d92ac7b6aa
--- /dev/null
+++ b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/dofn/package-info.java
@@ -0,0 +1,20 @@
+/*
+ * Copyright 2026 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.
+ */
+/**
+ * Contains all the {@link org.apache.beam.sdk.transforms.DoFn} implementations used in the
+ * pipeline.
+ */
+package com.google.cloud.bigtable.beam.hbasesnapshots.dofn;
diff --git a/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/transforms/HbaseRegionSplitTracker.java b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/transforms/HbaseRegionSplitTracker.java
new file mode 100644
index 0000000000..33432bc98e
--- /dev/null
+++ b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/transforms/HbaseRegionSplitTracker.java
@@ -0,0 +1,89 @@
+/*
+ * Copyright 2026 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.transforms;
+
+import com.google.api.core.InternalApi;
+import org.apache.beam.sdk.io.range.ByteKey;
+import org.apache.beam.sdk.io.range.ByteKeyRange;
+import org.apache.beam.sdk.transforms.splittabledofn.ByteKeyRangeTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.SplitResult;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link RestrictionTracker} wrapping the {@link ByteKeyRangeTracker} for controlled execution of
+ * dynamic splitting.
+ */
+@InternalApi("For internal usage only")
+public class HbaseRegionSplitTracker extends RestrictionTracker
+ implements RestrictionTracker.HasProgress {
+
+ private static final Logger LOG = LoggerFactory.getLogger(HbaseRegionSplitTracker.class);
+
+ private final String snapshotName;
+
+ private final String regionName;
+ private final ByteKeyRangeTracker byteKeyRangeTracker;
+
+ private final boolean enableDynamicSplitting;
+
+ public HbaseRegionSplitTracker(
+ String snapshotName, String regionName, ByteKeyRange range, boolean enableDynamicSplitting) {
+ this.snapshotName = snapshotName;
+ this.regionName = regionName;
+ this.byteKeyRangeTracker = ByteKeyRangeTracker.of(range);
+ this.enableDynamicSplitting = enableDynamicSplitting;
+ }
+
+ @Override
+ public ByteKeyRange currentRestriction() {
+ return this.byteKeyRangeTracker.currentRestriction();
+ }
+
+ @Override
+ public SplitResult trySplit(double fractionOfRemainder) {
+ LOG.info(
+ "Splitting restriction for region:{} in snapshot:{}", this.regionName, this.snapshotName);
+
+ return enableDynamicSplitting ? this.byteKeyRangeTracker.trySplit(fractionOfRemainder) : null;
+ }
+
+ @Override
+ public boolean tryClaim(ByteKey key) {
+ return this.byteKeyRangeTracker.tryClaim(key);
+ }
+
+ @Override
+ public void checkDone() throws IllegalStateException {
+ this.byteKeyRangeTracker.checkDone();
+ }
+
+ @Override
+ public RestrictionTracker.IsBounded isBounded() {
+ return this.byteKeyRangeTracker.isBounded();
+ }
+
+ @Override
+ public String toString() {
+ return this.byteKeyRangeTracker.toString();
+ }
+
+ @Override
+ public Progress getProgress() {
+ return this.byteKeyRangeTracker.getProgress();
+ }
+}
diff --git a/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/transforms/ListRegions.java b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/transforms/ListRegions.java
new file mode 100644
index 0000000000..3e6415ddf7
--- /dev/null
+++ b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/transforms/ListRegions.java
@@ -0,0 +1,129 @@
+/*
+ * Copyright 2026 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.transforms;
+
+import com.google.api.core.InternalApi;
+import com.google.cloud.bigtable.beam.hbasesnapshots.SnapshotUtils;
+import com.google.cloud.bigtable.beam.hbasesnapshots.coders.RegionConfigCoder;
+import com.google.cloud.bigtable.beam.hbasesnapshots.conf.RegionConfig;
+import com.google.cloud.bigtable.beam.hbasesnapshots.conf.SnapshotConfig;
+import com.google.common.annotations.VisibleForTesting;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.transforms.Reshuffle;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.FileSystem;
+import org.apache.hadoop.fs.Path;
+import org.apache.hadoop.hbase.client.RegionInfo;
+import org.apache.hadoop.hbase.client.TableDescriptor;
+import org.apache.hadoop.hbase.mapreduce.TableSnapshotInputFormatImpl;
+import org.apache.hadoop.hbase.shaded.protobuf.generated.SnapshotProtos;
+import org.apache.hadoop.hbase.snapshot.SnapshotManifest;
+import org.apache.hbase.thirdparty.com.google.protobuf.ByteString;
+
+/**
+ * A {@link PTransform} for listing the regions from snapshot manifest and builds the {@link
+ * RegionConfig} instances
+ */
+@InternalApi("For internal usage only")
+public class ListRegions
+ extends PTransform, PCollection> {
+
+ @VisibleForTesting
+ static class ListRegionsFn extends DoFn {
+
+ // Beam reuses DoFn instances for multiple elements. Initializing the cache in @Setup
+ // ensures that we only create the cache once per DoFn instance lifecycle (per worker thread),
+ // avoiding heavy XML parsing overhead for Configuration while also avoiding static state
+ // and ensuring thread safety since Beam isolates DoFn instances.
+ @Setup
+ public void setup() {
+ configCache = new HashMap<>();
+ }
+
+ private transient Map, Configuration> configCache;
+
+ private Map computeRegionSize(SnapshotManifest snapshotManifest) {
+ return snapshotManifest.getRegionManifests().stream()
+ .collect(
+ Collectors.toMap(
+ region -> region.getRegionInfo().getStartKey(),
+ region ->
+ region.getFamilyFilesList().stream()
+ .flatMap(family -> family.getStoreFilesList().stream())
+ .mapToLong(SnapshotProtos.SnapshotRegionManifest.StoreFile::getFileSize)
+ .sum(),
+ (size1, size2) -> size1 + size2));
+ }
+
+ /**
+ * Reads snapshot file manifest and lists all the regions including the size.
+ *
+ * @param snapshotConfig - Snapshot Configuration containing source path.
+ * @param outputReceiver
+ * @throws Exception
+ */
+ @ProcessElement
+ public void processElement(
+ @Element SnapshotConfig snapshotConfig, OutputReceiver outputReceiver)
+ throws Exception {
+
+ Map configDetails = snapshotConfig.getConfigurationDetails();
+ Configuration configuration =
+ configCache.computeIfAbsent(configDetails, SnapshotUtils::getHBaseConfiguration);
+ Path sourcePath = snapshotConfig.getSourcePath();
+ FileSystem fileSystem = sourcePath.getFileSystem(configuration);
+ SnapshotManifest snapshotManifest =
+ TableSnapshotInputFormatImpl.getSnapshotManifest(
+ configuration, snapshotConfig.getSnapshotName(), sourcePath, fileSystem);
+
+ Map regionsSize = computeRegionSize(snapshotManifest);
+ TableDescriptor tableDescriptor = snapshotManifest.getTableDescriptor();
+
+ // Extract region information from the snapshot manifest.
+ List regionInfos =
+ TableSnapshotInputFormatImpl.getRegionInfosFromManifest(snapshotManifest);
+
+ // Emit a RegionConfig for each region in the snapshot.
+ regionInfos.stream()
+ .map(
+ regionInfo ->
+ RegionConfig.builder()
+ .setSnapshotConfig(snapshotConfig)
+ .setTableDescriptor(tableDescriptor)
+ .setRegionInfo(regionInfo)
+ .setRegionSize(
+ regionsSize.getOrDefault(
+ ByteString.copyFrom(regionInfo.getStartKey()), 0L))
+ .build())
+ .forEach(outputReceiver::output);
+ }
+ }
+
+ @Override
+ public PCollection expand(PCollection snapshotconfigs) {
+ return snapshotconfigs
+ .apply("List Regions", ParDo.of(new ListRegionsFn()))
+ .setCoder(new RegionConfigCoder())
+ .apply(Reshuffle.viaRandomKey());
+ }
+}
diff --git a/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/transforms/ReadRegions.java b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/transforms/ReadRegions.java
new file mode 100644
index 0000000000..572cd4331d
--- /dev/null
+++ b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/transforms/ReadRegions.java
@@ -0,0 +1,163 @@
+/*
+ * Copyright 2026 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.transforms;
+
+import com.google.api.core.InternalApi;
+import com.google.cloud.bigtable.beam.hbasesnapshots.conf.RegionConfig;
+import com.google.cloud.bigtable.beam.hbasesnapshots.conf.SnapshotKey;
+import com.google.cloud.bigtable.beam.hbasesnapshots.dofn.CreateBigtableMutationsFn;
+import com.google.cloud.bigtable.beam.hbasesnapshots.dofn.ReadSnapshotRegionFn;
+import com.google.common.base.Preconditions;
+import java.math.BigInteger;
+import org.apache.beam.sdk.Pipeline;
+import org.apache.beam.sdk.coders.CannotProvideCoderException;
+import org.apache.beam.sdk.coders.Coder;
+import org.apache.beam.sdk.coders.KvCoder;
+import org.apache.beam.sdk.schemas.NoSuchSchemaException;
+import org.apache.beam.sdk.schemas.SchemaCoder;
+import org.apache.beam.sdk.transforms.Filter;
+import org.apache.beam.sdk.transforms.PTransform;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.values.KV;
+import org.apache.beam.sdk.values.PCollection;
+import org.apache.beam.sdk.values.TypeDescriptor;
+import org.apache.hadoop.hbase.client.Mutation;
+import org.apache.hadoop.hbase.client.Result;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * A {@link PTransform} for reading the records from each region and creates Hbase {@link Mutation}
+ * instances. Each region will be split into configured size (512 MB) and pipeline option {@link
+ * ImportSnapshots.ImportSnapshotsOptions#getUseDynamicSplitting() useDynamicSplitting} can be used
+ * to control whether each split needs to be subdivided further or not.
+ */
+@InternalApi("For internal usage only")
+public class ReadRegions
+ extends PTransform, PCollection>>> {
+
+ private static final Logger LOG = LoggerFactory.getLogger(ReadRegions.class);
+
+ private final boolean useDynamicSplitting;
+
+ private final int maxMutationsPerRequestThreshold;
+
+ private final boolean filterLargeRows;
+ private final long filterLargeRowThresholdBytes;
+
+ private final boolean filterLargeCells;
+ private final int filterLargeCellThresholdBytes;
+
+ private final boolean filterLargeRowKeys;
+ private final int filterLargeRowKeysThresholdBytes;
+ private final boolean filterWideRows;
+
+ private final Integer numShards;
+ private final Integer shardIndex;
+
+ public ReadRegions(
+ boolean useDynamicSplitting,
+ int maxMutationsPerRequestThreshold,
+ boolean filterLargeRows,
+ long filterLargeRowThresholdBytes,
+ boolean filterLargeCells,
+ int filterLargeCellThresholdBytes,
+ boolean filterLargeRowKeys,
+ int filterLargeRowKeysThresholdBytes,
+ boolean filterWideRows,
+ Integer numShards,
+ Integer shardIndex) {
+ this.useDynamicSplitting = useDynamicSplitting;
+ this.maxMutationsPerRequestThreshold = maxMutationsPerRequestThreshold;
+
+ this.filterLargeRows = filterLargeRows;
+ this.filterLargeRowThresholdBytes = filterLargeRowThresholdBytes;
+
+ this.filterLargeCells = filterLargeCells;
+ this.filterLargeCellThresholdBytes = filterLargeCellThresholdBytes;
+
+ this.filterLargeRowKeys = filterLargeRowKeys;
+ this.filterLargeRowKeysThresholdBytes = filterLargeRowKeysThresholdBytes;
+ this.filterWideRows = filterWideRows;
+
+ if (numShards != null) {
+ Preconditions.checkArgument(
+ shardIndex != null, "if numShards is set, shardIndex must also be");
+ Preconditions.checkArgument(
+ shardIndex >= 0 && shardIndex < numShards, "shardIndex must be between [0, numShards)");
+ }
+
+ this.numShards = numShards;
+ this.shardIndex = shardIndex;
+ }
+
+ static boolean isRegionSelected(RegionConfig rc, int numShards, int shardIndex) {
+ byte[] regionName = rc.getRegionInfo().getEncodedNameAsBytes();
+ // Use signum=1 to force the byte array to be interpreted as a positive magnitude,
+ // avoiding negative numbers and potential sharding skews.
+ long remainder = new BigInteger(1, regionName).mod(BigInteger.valueOf(numShards)).longValue();
+ return remainder == shardIndex;
+ }
+
+ @Override
+ public PCollection>> expand(
+ PCollection regionConfig) {
+ // 1. Resolve coders for the output KV type.
+ Pipeline pipeline = regionConfig.getPipeline();
+ SchemaCoder snapshotKeySchemaCoder;
+ Coder hbaseResultCoder;
+ try {
+ snapshotKeySchemaCoder = pipeline.getSchemaRegistry().getSchemaCoder(SnapshotKey.class);
+ hbaseResultCoder = pipeline.getCoderRegistry().getCoder(TypeDescriptor.of(Result.class));
+ } catch (CannotProvideCoderException | NoSuchSchemaException e) {
+ throw new RuntimeException(e);
+ }
+
+ // 2. Apply sharding if enabled. Encoded name is an MD5 hash, so it provides
+ // good distribution across shards.
+ PCollection maybeShardedRegions = regionConfig;
+ if (numShards != null) {
+ maybeShardedRegions =
+ regionConfig.apply(
+ "Select regions for shard",
+ Filter.by(
+ rc -> {
+ boolean shouldTake = isRegionSelected(rc, numShards, shardIndex);
+ LOG.info(
+ "Region {} was {} due to sharding",
+ rc.getRegionInfo().getRegionNameAsString(),
+ shouldTake ? "taken" : "skipped");
+ return shouldTake;
+ }));
+ }
+
+ return maybeShardedRegions
+ .apply("Read snapshot region", ParDo.of(new ReadSnapshotRegionFn(this.useDynamicSplitting)))
+ .setCoder(KvCoder.of(snapshotKeySchemaCoder, hbaseResultCoder))
+ .apply(
+ "Create Mutation",
+ ParDo.of(
+ new CreateBigtableMutationsFn(
+ this.maxMutationsPerRequestThreshold,
+ this.filterLargeRows,
+ this.filterLargeRowThresholdBytes,
+ this.filterLargeCells,
+ this.filterLargeCellThresholdBytes,
+ this.filterLargeRowKeys,
+ this.filterLargeRowKeysThresholdBytes,
+ this.filterWideRows)));
+ }
+}
diff --git a/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/transforms/package-info.java b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/transforms/package-info.java
new file mode 100644
index 0000000000..17f27d4046
--- /dev/null
+++ b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/hbasesnapshots/transforms/package-info.java
@@ -0,0 +1,17 @@
+/*
+ * Copyright 2026 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 contains all the {@link org.apache.beam.sdk.transforms.PTransform} implementations. */
+package com.google.cloud.bigtable.beam.hbasesnapshots.transforms;
diff --git a/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/sequencefiles/ImportJob.java b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/sequencefiles/ImportJob.java
index ec56b7dab7..197176544e 100644
--- a/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/sequencefiles/ImportJob.java
+++ b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/sequencefiles/ImportJob.java
@@ -126,6 +126,18 @@ public interface ImportOptions extends GcpOptions {
@SuppressWarnings("unused")
void setMutationThrottleLatencyMs(ValueProvider throttleMs);
+ @Description("Maximum number of inflight RPCs.")
+ @Default.Integer(100)
+ ValueProvider getMaxInflightRpcs();
+
+ void setMaxInflightRpcs(ValueProvider value);
+
+ @Description("Bulk mutation close timeout in minutes.")
+ @Default.Integer(30)
+ ValueProvider getBulkMutationCloseTimeoutMinutes();
+
+ void setBulkMutationCloseTimeoutMinutes(ValueProvider value);
+
// When creating a template, this flag must be set to false.
@Description("Wait for pipeline to finish.")
@Default.Boolean(true)
diff --git a/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/validation/BufferedHadoopHashTableSource.java b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/validation/BufferedHadoopHashTableSource.java
index e1219f88be..cdff1a2ae1 100644
--- a/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/validation/BufferedHadoopHashTableSource.java
+++ b/bigtable-dataflow-parent/bigtable-beam-import/src/main/java/com/google/cloud/bigtable/beam/validation/BufferedHadoopHashTableSource.java
@@ -50,7 +50,6 @@ class BufferedHadoopHashTableSource extends BoundedSource>> CODER =
KvCoder.of(StringUtf8Coder.of(), ListCoder.of(RangeHashCoder.of()));
- ;
// Max number of RangeHashes to buffer.
private final int maxBufferSize;
diff --git a/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/TemplateUtilsTest.java b/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/TemplateUtilsTest.java
new file mode 100644
index 0000000000..ea317350f8
--- /dev/null
+++ b/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/TemplateUtilsTest.java
@@ -0,0 +1,102 @@
+/*
+ * Copyright 2026 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 static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+import com.google.cloud.bigtable.beam.sequencefiles.ImportJob.ImportOptions;
+import com.google.cloud.bigtable.hbase.BigtableOptionsFactory;
+import com.google.cloud.bigtable.hbase.wrappers.BigtableHBaseSettings;
+import java.util.Map;
+import java.util.concurrent.TimeUnit;
+import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.apache.beam.sdk.options.ValueProvider;
+import org.apache.beam.sdk.options.ValueProvider.StaticValueProvider;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Tests the {@link TemplateUtils} to ensure correct configuration building. */
+@RunWith(JUnit4.class)
+public class TemplateUtilsTest {
+
+ /**
+ * Tests that {@link TemplateUtils#buildImportConfig} sets correct default values for Bigtable
+ * configuration.
+ */
+ @Test
+ public void testBuildImportConfig_defaults() {
+ ImportOptions opts = PipelineOptionsFactory.as(ImportOptions.class);
+ opts.setBigtableProject(StaticValueProvider.of("my-project"));
+ opts.setBigtableInstanceId(StaticValueProvider.of("my-instance"));
+ opts.setBigtableTableId(StaticValueProvider.of("my-table"));
+
+ CloudBigtableTableConfiguration config = TemplateUtils.buildImportConfig(opts, "my-agent");
+
+ assertEquals("my-project", config.getProjectId());
+ assertEquals("my-instance", config.getInstanceId());
+ assertEquals("my-table", config.getTableId());
+
+ Map> configMap = config.getConfiguration();
+ assertEquals("my-agent", configMap.get(BigtableOptionsFactory.CUSTOM_USER_AGENT_KEY).get());
+ assertTrue(configMap.containsKey(BigtableOptionsFactory.MAX_INFLIGHT_RPCS_KEY));
+ assertTrue(
+ configMap.containsKey(BigtableHBaseSettings.BULK_MUTATION_CLOSE_TIMEOUT_MILLISECONDS));
+ }
+
+ /**
+ * Tests that {@link TemplateUtils#buildImportConfig} correctly overrides defaults with custom
+ * values provided in options.
+ */
+ @Test
+ public void testBuildImportConfig_customValues() {
+ ImportOptions opts = PipelineOptionsFactory.as(ImportOptions.class);
+ opts.setBigtableProject(StaticValueProvider.of("my-project"));
+ opts.setBigtableInstanceId(StaticValueProvider.of("my-instance"));
+ opts.setBigtableTableId(StaticValueProvider.of("my-table"));
+ opts.setMaxInflightRpcs(StaticValueProvider.of(200));
+ opts.setBulkMutationCloseTimeoutMinutes(StaticValueProvider.of(60));
+
+ CloudBigtableTableConfiguration config = TemplateUtils.buildImportConfig(opts, "my-agent");
+
+ Map> configMap = config.getConfiguration();
+ assertEquals("200", configMap.get(BigtableOptionsFactory.MAX_INFLIGHT_RPCS_KEY).get());
+ assertEquals(
+ String.valueOf(TimeUnit.MINUTES.toMillis(60)),
+ configMap.get(BigtableHBaseSettings.BULK_MUTATION_CLOSE_TIMEOUT_MILLISECONDS).get());
+ }
+
+ /**
+ * Tests that {@link TemplateUtils#buildImportConfig} uses default value 30 when custom bulk
+ * mutation timeout is null.
+ */
+ @Test
+ public void testBuildImportConfig_nullBulkMutationTimeout() {
+ ImportOptions opts = PipelineOptionsFactory.as(ImportOptions.class);
+ opts.setBigtableProject(StaticValueProvider.of("my-project"));
+ opts.setBigtableInstanceId(StaticValueProvider.of("my-instance"));
+ opts.setBigtableTableId(StaticValueProvider.of("my-table"));
+ opts.setBulkMutationCloseTimeoutMinutes(StaticValueProvider.of(null));
+
+ CloudBigtableTableConfiguration config = TemplateUtils.buildImportConfig(opts, "my-agent");
+
+ Map> configMap = config.getConfiguration();
+ assertEquals(
+ String.valueOf(TimeUnit.MINUTES.toMillis(30)),
+ configMap.get(BigtableHBaseSettings.BULK_MUTATION_CLOSE_TIMEOUT_MILLISECONDS).get());
+ }
+}
diff --git a/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/EndToEndIT.java b/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/EndToEndIT.java
index 1f3758dc21..acbbd1edf2 100644
--- a/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/EndToEndIT.java
+++ b/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/EndToEndIT.java
@@ -19,6 +19,8 @@
import com.google.api.services.storage.model.StorageObject;
import com.google.bigtable.repackaged.com.google.gson.Gson;
import com.google.cloud.bigtable.beam.hbasesnapshots.ImportJobFromHbaseSnapshot.ImportOptions;
+import com.google.cloud.bigtable.beam.hbasesnapshots.conf.HBaseSnapshotInputConfigBuilder;
+import com.google.cloud.bigtable.beam.hbasesnapshots.dofn.CleanupHBaseSnapshotRestoreFilesFn;
import com.google.cloud.bigtable.beam.sequencefiles.HBaseResultToMutationFn;
import com.google.cloud.bigtable.beam.test_env.EnvSetup;
import com.google.cloud.bigtable.beam.test_env.TestProperties;
@@ -50,6 +52,7 @@
import org.apache.beam.sdk.PipelineResult;
import org.apache.beam.sdk.PipelineResult.State;
import org.apache.beam.sdk.extensions.gcp.options.GcpOptions;
+import org.apache.beam.sdk.extensions.gcp.options.GcsOptions;
import org.apache.beam.sdk.extensions.gcp.util.GcsUtil;
import org.apache.beam.sdk.extensions.gcp.util.gcsfs.GcsPath;
import org.apache.beam.sdk.metrics.MetricQueryResults;
@@ -84,6 +87,7 @@
public class EndToEndIT {
private static final Logger LOG = LoggerFactory.getLogger(HBaseResultToMutationFn.class);
+ // The snapshot fixture used in these tests is approximately 78 KB zipped.
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";
@@ -315,7 +319,7 @@ public void testHBaseSnapshotImport() throws Exception {
String bucket = GcsPath.fromUri(hbaseSnapshotDir).getBucket();
String restorePathPrefix =
CleanupHBaseSnapshotRestoreFilesFn.getListPrefix(
- HBaseSnapshotInputConfigBuilder.RESTORE_DIR);
+ HBaseSnapshotInputConfigBuilder.RESTORE_DIR + importOpts.getJobName());
List allObjects = new ArrayList<>();
String nextToken;
do {
@@ -428,7 +432,7 @@ public void testSnappyCompressedHBaseSnapshotImport() throws Exception {
String bucket = GcsPath.fromUri(hbaseSnapshotDir).getBucket();
String restorePathPrefix =
CleanupHBaseSnapshotRestoreFilesFn.getListPrefix(
- HBaseSnapshotInputConfigBuilder.RESTORE_DIR);
+ HBaseSnapshotInputConfigBuilder.RESTORE_DIR + importOpts.getJobName());
List allObjects = new ArrayList<>();
String nextToken;
@@ -462,4 +466,63 @@ public void testSnappyCompressedHBaseSnapshotImport() throws Exception {
Assert.assertEquals(counters.get("ranges_matched"), (Long) 100L);
Assert.assertEquals(counters.get("ranges_not_matched"), (Long) 0L);
}
+
+ @Test
+ public void testHBaseSnapshotImportWithSharding() throws Exception {
+ int numShards = 4;
+ String sharedRestorePath = workDir + "shared-restore/";
+
+ for (int i = 0; i < numShards; i++) {
+ ImportOptions importOpts = createImportOptions();
+ importOpts.setSnapshotName(null);
+ importOpts.setSnapshots(TEST_SNAPSHOT_NAME + ":" + tableId);
+ importOpts.setNumShards(numShards);
+ importOpts.setShardIndex(i);
+ importOpts.setRestorePath(sharedRestorePath);
+ importOpts.setDeleteRestoredSnapshots(false);
+ if (i > 0) {
+ // Skip restore step for subsequent shards as they can share the restored files from the
+ // first shard.
+ importOpts.setSkipRestoreStep(true);
+ }
+
+ com.google.cloud.bigtable.beam.hbasesnapshots.conf.ImportConfig importConfig =
+ ImportJobFromHbaseSnapshot.buildImportConfigFromPipelineOptions(
+ importOpts, importOpts.as(GcsOptions.class));
+
+ PipelineResult pipelineResult =
+ ImportJobFromHbaseSnapshot.buildPipelineWithMultipleSnapshots(importOpts, importConfig)
+ .run();
+ State state = pipelineResult.waitUntilFinish();
+ if (state != State.DONE) {
+ System.err.println("Pipeline failed for shard " + i + "! State: " + state);
+ if (pipelineResult instanceof org.apache.beam.runners.dataflow.DataflowPipelineJob) {
+ System.err.println(
+ "Dataflow Job ID: "
+ + ((org.apache.beam.runners.dataflow.DataflowPipelineJob) pipelineResult)
+ .getJobId());
+ } else {
+ System.err.println("Pipeline Result: " + pipelineResult);
+ }
+ }
+ Assert.assertEquals(State.DONE, state);
+ }
+
+ SyncTableOptions syncOpts = createSyncTableOptions();
+
+ PipelineResult result = SyncTableJob.buildPipeline(syncOpts).run();
+ State state = result.waitUntilFinish();
+ Assert.assertEquals(State.DONE, state);
+
+ Assert.assertEquals(0, readMismatchesFromOutputFiles().size());
+
+ Map counters = getCountMap(result);
+ Assert.assertEquals(counters.get("ranges_matched"), (Long) 100L);
+ Assert.assertEquals(counters.get("ranges_not_matched"), (Long) 0L);
+
+ final List paths = gcsUtil.expand(GcsPath.fromUri(sharedRestorePath + "*"));
+ if (!paths.isEmpty()) {
+ gcsUtil.remove(paths.stream().map(GcsPath::toString).collect(Collectors.toList()));
+ }
+ }
}
diff --git a/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/HBaseSnapshotInputConfigBuilderTest.java b/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/HBaseSnapshotInputConfigBuilderTest.java
index fb5346f72a..039bd5b3e2 100644
--- a/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/HBaseSnapshotInputConfigBuilderTest.java
+++ b/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/HBaseSnapshotInputConfigBuilderTest.java
@@ -18,6 +18,7 @@
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.fail;
+import com.google.cloud.bigtable.beam.hbasesnapshots.conf.HBaseSnapshotInputConfigBuilder;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.hbase.mapreduce.TableSnapshotInputFormat;
import org.apache.hadoop.mapreduce.InputFormat;
diff --git a/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/HBaseSnapshotRestoreToolTest.java b/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/HBaseSnapshotRestoreToolTest.java
new file mode 100644
index 0000000000..6d48b5dee3
--- /dev/null
+++ b/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/HBaseSnapshotRestoreToolTest.java
@@ -0,0 +1,75 @@
+/*
+ * Copyright 2026 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 static org.junit.Assert.assertEquals;
+
+import com.google.cloud.bigtable.beam.hbasesnapshots.conf.ImportConfig;
+import java.util.HashMap;
+import java.util.Map;
+import org.apache.beam.sdk.extensions.gcp.options.GcsOptions;
+import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class HBaseSnapshotRestoreToolTest {
+
+ private Map originalProperties = new HashMap<>();
+
+ @Before
+ public void setUp() {
+ originalProperties.put("project", System.getProperty("project"));
+ originalProperties.put("hbaseSnapshotSourceDir", System.getProperty("hbaseSnapshotSourceDir"));
+ originalProperties.put("snapshots", System.getProperty("snapshots"));
+ originalProperties.put("restorePath", System.getProperty("restorePath"));
+ }
+
+ @After
+ public void tearDown() {
+ for (Map.Entry entry : originalProperties.entrySet()) {
+ if (entry.getValue() != null) {
+ System.setProperty(entry.getKey(), entry.getValue());
+ } else {
+ System.clearProperty(entry.getKey());
+ }
+ }
+ }
+
+ /**
+ * Tests that {@link HBaseSnapshotRestoreTool#buildImportConfigFromArgs} correctly parses system
+ * properties into an {@link ImportConfig}.
+ */
+ @Test
+ public void testBuildImportConfigFromArgs() throws Exception {
+ System.setProperty("project", "my-project");
+ System.setProperty("hbaseSnapshotSourceDir", "gs://my-bucket/snapshots");
+ System.setProperty("snapshots", "snap1:table1");
+ System.setProperty("restorePath", "gs://my-bucket/restore");
+
+ GcsOptions options = PipelineOptionsFactory.create().as(GcsOptions.class);
+ ImportConfig config = HBaseSnapshotRestoreTool.buildImportConfigFromArgs(options);
+
+ assertEquals("gs://my-bucket/snapshots", config.getSourcepath());
+ assertEquals("gs://my-bucket/restore", config.getRestorepath());
+ assertEquals(1, config.getSnapshots().size());
+ assertEquals("snap1", config.getSnapshots().get(0).getSnapshotName());
+ assertEquals("table1", config.getSnapshots().get(0).getbigtableTableName());
+ }
+}
diff --git a/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/ImportJobFromHbaseSnapshotTest.java b/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/ImportJobFromHbaseSnapshotTest.java
new file mode 100644
index 0000000000..434b92ab47
--- /dev/null
+++ b/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/ImportJobFromHbaseSnapshotTest.java
@@ -0,0 +1,200 @@
+/*
+ * Copyright 2026 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 static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.is;
+import static org.hamcrest.Matchers.notNullValue;
+
+import com.google.api.services.storage.model.Objects;
+import com.google.api.services.storage.model.StorageObject;
+import com.google.cloud.bigtable.beam.hbasesnapshots.conf.ImportConfig;
+import java.io.File;
+import java.util.Arrays;
+import java.util.List;
+import org.apache.beam.sdk.extensions.gcp.options.GcsOptions;
+import org.apache.beam.sdk.extensions.gcp.util.GcsUtil;
+import org.junit.ClassRule;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.ExpectedException;
+import org.junit.rules.TemporaryFolder;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.Mock;
+import org.mockito.Mockito;
+import org.mockito.junit.MockitoJUnit;
+import org.mockito.junit.MockitoRule;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** Test cases for the {@link ImportJobFromHbaseSnapshot} class. */
+@RunWith(JUnit4.class)
+public class ImportJobFromHbaseSnapshotTest {
+ private static final Logger LOG = LoggerFactory.getLogger(ImportJobFromHbaseSnapshotTest.class);
+
+ @ClassRule public static TemporaryFolder tempFolder = new TemporaryFolder();
+ @Rule public final ExpectedException expectedException = ExpectedException.none();
+
+ @Rule public final MockitoRule mockito = MockitoJUnit.rule();
+ @Mock GcsOptions gcsOptions;
+ @Mock GcsUtil gcsUtilMock;
+ @Mock Objects gcsObjects;
+
+ /**
+ * Tests that {@link ImportJobFromHbaseSnapshot#buildImportConfigFromPipelineOptions} throws
+ * exception when source path is missing.
+ */
+ @Test
+ public void testBuildImportConfigWithMissingSourcePathThrowsException() throws Exception {
+ ImportJobFromHbaseSnapshot.ImportOptions options =
+ SnapshotTestHelper.getPipelineOptions(
+ new String[] {
+ "--snapshots='bookmark-2099:bookmark,malwarescanstate-9087:malwarescan'"
+ });
+
+ expectedException.expect(NullPointerException.class);
+ expectedException.expectMessage(ImportJobFromHbaseSnapshot.MISSING_SNAPSHOT_SOURCEPATH);
+ ImportJobFromHbaseSnapshot.buildImportConfigFromPipelineOptions(options, gcsOptions);
+ }
+
+ /**
+ * Tests that {@link ImportJobFromHbaseSnapshot#buildImportConfigFromPipelineOptions} throws
+ * exception when snapshots are missing.
+ */
+ @Test
+ public void testBuildImportConfigWithMissingSnapshotsThrowsException() throws Exception {
+ ImportJobFromHbaseSnapshot.ImportOptions options =
+ SnapshotTestHelper.getPipelineOptions(
+ new String[] {"--hbaseSnapshotSourceDir=gs://bucket/data/"});
+
+ expectedException.expect(NullPointerException.class);
+ expectedException.expectMessage(ImportJobFromHbaseSnapshot.MISSING_SNAPSHOT_NAMES);
+ ImportJobFromHbaseSnapshot.buildImportConfigFromPipelineOptions(options, gcsOptions);
+ }
+
+ /**
+ * Tests that {@link ImportJobFromHbaseSnapshot#buildImportConfigFromPipelineOptions} correctly
+ * parses snapshots from string.
+ */
+ @Test
+ public void testBuildImportConfigFromSnapshotsString() throws Exception {
+ String sourcePath = "gs://bucket/data/";
+ ImportJobFromHbaseSnapshot.ImportOptions options =
+ SnapshotTestHelper.getPipelineOptions(
+ new String[] {
+ "--hbaseSnapshotSourceDir=" + sourcePath,
+ "--snapshots='bookmark-2099:bookmark,malwarescanstate-9087:malwarescan'"
+ });
+
+ ImportConfig importConfig =
+ ImportJobFromHbaseSnapshot.buildImportConfigFromPipelineOptions(options, gcsOptions);
+ assertThat(importConfig.getSourcepath(), is(sourcePath));
+ assertThat(importConfig.getRestorepath(), notNullValue());
+ assertThat(importConfig.getSnapshots().size(), is(2));
+ }
+
+ private void setUpGcsObjectMocks(List fakeStorageObjects) throws Exception {
+ Mockito.when(gcsObjects.getItems()).thenReturn(fakeStorageObjects);
+ Mockito.when(gcsUtilMock.listObjects(Mockito.anyString(), Mockito.anyString(), Mockito.any()))
+ .thenReturn(gcsObjects);
+ }
+
+ /**
+ * Tests that {@link ImportJobFromHbaseSnapshot#buildImportConfigFromPipelineOptions} handles
+ * wildcard for all snapshots.
+ */
+ @Test
+ public void testBuildImportConfigForAllSnapshots() throws Exception {
+ String baseObjectPath = "snapshots/20220309230526";
+ String importSnapshotpath = String.format("gs://sym-bucket/%s", baseObjectPath);
+ ImportJobFromHbaseSnapshot.ImportOptions options =
+ SnapshotTestHelper.getPipelineOptions(
+ new String[] {"--hbaseSnapshotSourceDir=" + importSnapshotpath, "--snapshots=*"});
+ Mockito.when(gcsOptions.getGcsUtil()).thenReturn(gcsUtilMock);
+
+ List snapshotList = Arrays.asList("audit-events", "dlpInfo", "ce-metrics-manifest");
+ List fakeStorageObjects =
+ SnapshotTestHelper.createFakeStorageObjects(baseObjectPath, snapshotList);
+ setUpGcsObjectMocks(fakeStorageObjects);
+
+ ImportConfig importConfig =
+ ImportJobFromHbaseSnapshot.buildImportConfigFromPipelineOptions(options, gcsOptions);
+ assertThat(importConfig.getSourcepath(), is(importSnapshotpath));
+ assertThat(importConfig.getRestorepath(), notNullValue());
+ assertThat(importConfig.getSnapshots().size(), is(snapshotList.size()));
+ }
+
+ /**
+ * Tests that {@link ImportJobFromHbaseSnapshot#buildImportConfigFromConfigFile} throws exception
+ * when source path is missing in JSON.
+ */
+ @Test
+ public void testBuildImportConfigFromJsonFileWithMissingPathThrowsException() throws Exception {
+ String config =
+ "{\n"
+ + " \"snapshots\": {\n"
+ + " \"snap_demo1\": \"snap_demo1\",\n"
+ + " \"snap_demo2\": \"snap_demo2\"\n"
+ + " }\n"
+ + "}";
+ File file = tempFolder.newFile();
+ SnapshotTestHelper.writeToFile(file.getAbsolutePath(), config);
+ ImportJobFromHbaseSnapshot.ImportOptions options =
+ SnapshotTestHelper.getPipelineOptions(
+ new String[] {"--importConfigFilePath=" + file.getAbsolutePath()});
+
+ expectedException.expect(NullPointerException.class);
+ expectedException.expectMessage(ImportJobFromHbaseSnapshot.MISSING_SNAPSHOT_SOURCEPATH);
+
+ ImportConfig importConfig =
+ ImportJobFromHbaseSnapshot.buildImportConfigFromConfigFile(
+ options.getImportConfigFilePath());
+ }
+
+ /**
+ * Tests that {@link ImportJobFromHbaseSnapshot#buildImportConfigFromConfigFile} correctly parses
+ * config from JSON file.
+ */
+ @Test
+ public void testBuildImportConfigFromJsonFile() throws Exception {
+ String importSnapshotpath = "gs://sym-datastore/snapshots/data/snap_demo";
+ String restoreSnapshotpath = "gs://sym-datastore/snapshots/data/restore";
+ String config =
+ String.format(
+ "{\n"
+ + " \"sourcepath\": \"%s\",\n"
+ + " \"restorepath\": \"%s\",\n"
+ + " \"snapshots\": {\n"
+ + " \"snap_demo1\": \"demo1\",\n"
+ + " \"snap_demo2\": \"demo2\"\n"
+ + " }\n"
+ + "}",
+ importSnapshotpath, restoreSnapshotpath);
+
+ File file = tempFolder.newFile();
+ SnapshotTestHelper.writeToFile(file.getAbsolutePath(), config);
+ ImportJobFromHbaseSnapshot.ImportOptions options =
+ SnapshotTestHelper.getPipelineOptions(
+ new String[] {"--importConfigFilePath=" + file.getAbsolutePath()});
+ ImportConfig importConfig =
+ ImportJobFromHbaseSnapshot.buildImportConfigFromConfigFile(
+ options.getImportConfigFilePath());
+ assertThat(importConfig.getSourcepath(), is(importSnapshotpath));
+ assertThat(importConfig.getRestorepath().startsWith(restoreSnapshotpath), is(true));
+ assertThat(importConfig.getSnapshots().get(0).getbigtableTableName(), is("demo1"));
+ }
+}
diff --git a/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/SnapshotTestHelper.java b/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/SnapshotTestHelper.java
new file mode 100644
index 0000000000..e17a13ec41
--- /dev/null
+++ b/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/SnapshotTestHelper.java
@@ -0,0 +1,120 @@
+/*
+ * Copyright 2026 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.StorageObject;
+import com.google.cloud.bigtable.beam.hbasesnapshots.conf.SnapshotConfig;
+import com.google.common.base.Joiner;
+import com.google.common.io.ByteStreams;
+import java.io.ByteArrayInputStream;
+import java.io.IOException;
+import java.nio.channels.Channels;
+import java.nio.channels.ReadableByteChannel;
+import java.nio.channels.WritableByteChannel;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import org.apache.beam.sdk.io.FileSystems;
+import org.apache.beam.sdk.io.fs.ResourceId;
+import org.apache.beam.sdk.options.PipelineOptionsFactory;
+import org.apache.beam.sdk.util.MimeTypes;
+
+/** Contains various helper methods to handle different tasks while executing tests. */
+public class SnapshotTestHelper {
+ private SnapshotTestHelper() {}
+
+ /**
+ * Helper to generate files for testing.
+ *
+ * @param filePath The path to the file to write.
+ * @param fileContents The content to write.
+ * @throws IOException If an error occurs while creating or writing the file.
+ */
+ static void writeToFile(String filePath, String fileContents) throws IOException {
+
+ ResourceId resourceId = FileSystems.matchNewResource(filePath, false);
+
+ // Write the file contents to the channel and close.
+ try (ReadableByteChannel readChannel =
+ Channels.newChannel(new ByteArrayInputStream(fileContents.getBytes()))) {
+ try (WritableByteChannel writeChannel = FileSystems.create(resourceId, MimeTypes.TEXT)) {
+ ByteStreams.copy(readChannel, writeChannel);
+ }
+ }
+ }
+
+ /**
+ * @param restorePath - Path to which snapshots will be restored temporarily
+ * @return SnapshotConfig - Returns the snapshot config
+ */
+ public static SnapshotConfig newSnapshotConfig(String restorePath) {
+ return newSnapshotConfig("testsourcepath", restorePath);
+ }
+
+ public static SnapshotConfig newSnapshotConfig(String sourcePath, String restorePath) {
+ return SnapshotConfig.builder()
+ .setProjectId("testproject")
+ .setSourceLocation(sourcePath)
+ .setRestoreLocation(restorePath)
+ .setSnapshotName("testsnapshot")
+ .setTableName("testtable")
+ .setConfigurationDetails(new HashMap())
+ .build();
+ }
+
+ /**
+ * Helper method providing pipeline options.
+ *
+ * @param args list of pipeline arguments.
+ */
+ static ImportJobFromHbaseSnapshot.ImportOptions getPipelineOptions(String[] args) {
+ return PipelineOptionsFactory.fromArgs(args).as(ImportJobFromHbaseSnapshot.ImportOptions.class);
+ }
+
+ /**
+ * Creates Fake Storage Objects
+ *
+ * @param basePath File System base path
+ * @param objectNames List of object names
+ * @return List of matching Storage objects
+ */
+ static List createFakeStorageObjects(String basePath, List objectNames) {
+ if (objectNames == null) return null;
+
+ List storageObjects = new ArrayList<>();
+ objectNames.forEach(
+ name -> {
+ StorageObject object = new StorageObject();
+ String path = Joiner.on("/").join(basePath, ".hbase-snapshot", name, ".snapshotinfo");
+ object.setId(path);
+ object.setName(path);
+ storageObjects.add(object);
+ });
+
+ return storageObjects;
+ }
+
+ static Map buildMapFromList(String[] values) {
+ if (values.length % 2 != 0)
+ throw new IllegalArgumentException(
+ "Input should contain even number of values to represent both"
+ + " key and value for the map.");
+ Map data = new HashMap<>();
+ for (int i = 0; i < values.length; i += 2) data.put(values[i], values[i + 1]);
+ return data;
+ }
+}
diff --git a/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/SnapshotUtilsTest.java b/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/SnapshotUtilsTest.java
new file mode 100644
index 0000000000..8ae7b27807
--- /dev/null
+++ b/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/SnapshotUtilsTest.java
@@ -0,0 +1,360 @@
+/*
+ * Copyright 2026 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 static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.*;
+
+import com.google.api.services.storage.model.Objects;
+import com.google.api.services.storage.model.StorageObject;
+import com.google.cloud.bigtable.beam.hbasesnapshots.conf.ImportConfig;
+import com.google.cloud.bigtable.beam.hbasesnapshots.conf.SnapshotConfig;
+import java.io.IOException;
+import java.time.Instant;
+import java.time.ZoneId;
+import java.time.format.DateTimeFormatter;
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+import org.apache.beam.sdk.extensions.gcp.util.GcsUtil;
+import org.apache.hadoop.conf.Configuration;
+import org.junit.ClassRule;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.Mock;
+import org.mockito.Mockito;
+import org.mockito.junit.MockitoJUnit;
+import org.mockito.junit.MockitoRule;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** Test cases for the {@link SnapshotUtils} class. */
+@RunWith(JUnit4.class)
+public class SnapshotUtilsTest {
+ private static final Logger LOG = LoggerFactory.getLogger(SnapshotUtilsTest.class);
+
+ @ClassRule public static TemporaryFolder tempFolder = new TemporaryFolder();
+ // Preferred way to instantiate mocks in JUnit4 is via the JUnit rule MockitoJUnit
+ @Rule public final MockitoRule mockito = MockitoJUnit.rule();
+ @Mock GcsUtil gcsUtilMock;
+ @Mock Objects gcsObjects;
+
+ /**
+ * Tests that {@link SnapshotUtils#removeSuffixSlashIfExists} correctly removes trailing slashes.
+ */
+ @Test
+ public void testRemoveSuffixSlashIfExists() {
+ String path = "gs://bucket/prefix";
+
+ assertThat(SnapshotUtils.removeSuffixSlashIfExists(path), is(path));
+ assertThat(SnapshotUtils.removeSuffixSlashIfExists(path + "/"), is(path));
+ }
+
+ /**
+ * Tests that {@link SnapshotUtils#appendCurrentTimestamp} appends a valid timestamp to the path.
+ */
+ @Test
+ public void testAppendCurrentTimestamp() {
+ String path = "gs://bucket/prefix";
+ DateTimeFormatter formatter =
+ DateTimeFormatter.ofPattern("yyyyMMddHHmm").withZone(ZoneId.of("UTC"));
+ long currentTime = Long.parseLong(formatter.format(Instant.now()));
+ String returnVal = SnapshotUtils.appendCurrentTimestamp(path).replace(path + "/", "");
+ long returnTime = Long.parseLong(returnVal.split("-")[0]);
+ assertThat((returnTime - currentTime), lessThan(2L));
+ }
+
+ /** Tests that {@link SnapshotUtils#getNamedDirectory} returns correct directory path. */
+ @Test
+ public void testgetNamedDirectory() {
+ String path = "gs://bucket/subdir1";
+ String subFolder = "subdir2";
+ String expectedPath = "gs://bucket/subdir2";
+ String retValue = SnapshotUtils.getNamedDirectory(path, subFolder);
+ assertThat(retValue.startsWith(expectedPath), is(true));
+ }
+
+ /**
+ * Tests that {@link SnapshotUtils#getConfiguration} returns correct config for DataflowRunner.
+ */
+ @Test
+ public void testGetConfigurationWithDataflowRunner() {
+ String projectId = "testproject";
+ Map configurations =
+ SnapshotUtils.getConfiguration("DataflowRunner", projectId, "/path/to/sourcedir", null);
+ assertThat(configurations.get("fs.gs.project.id"), is(projectId));
+ assertThat(configurations.get("fs.gs.auth.type"), nullValue());
+ }
+
+ /** Tests that {@link SnapshotUtils#getConfiguration} returns correct config for DirectRunner. */
+ @Test
+ public void testGetConfigurationWithDirectRunner() {
+ Map hbaseConfiguration =
+ SnapshotTestHelper.buildMapFromList(
+ new String[] {"fs.AbstractFileSystem.gs.impl", "org.apache.hadoop.fs.hdfs"});
+ Map configurations =
+ SnapshotUtils.getConfiguration(
+ "DirectRunner", "testproject", "/path/to/sourcedir", hbaseConfiguration);
+ assertThat(
+ configurations.get("fs.AbstractFileSystem.gs.impl"),
+ is(hbaseConfiguration.get("fs.AbstractFileSystem.gs.impl")));
+ assertThat(configurations.get("fs.gs.auth.type"), is("APPLICATION_DEFAULT"));
+ }
+
+ /**
+ * Tests that {@link SnapshotUtils#getHBaseConfiguration} converts map to Hadoop Configuration.
+ */
+ @Test
+ public void testGetHbaseConfiguration() {
+ Map configurations =
+ SnapshotTestHelper.buildMapFromList(
+ new String[] {"throttling.enable", "true", "throttling.threshold.ms", "200"});
+ Configuration hbaseConfiguration = SnapshotUtils.getHBaseConfiguration(configurations);
+ assertThat(hbaseConfiguration.getBoolean("throttling.enable", false), is(true));
+ assertThat(hbaseConfiguration.get("throttling.threshold.ms"), is("200"));
+ }
+
+ /** Tests that {@link SnapshotUtils#buildSnapshotConfigs} creates correct SnapshotConfig list. */
+ @Test
+ public void testBuildSnapshotConfigs() {
+ String projectId = "testproject";
+ String sourcePath = "/path/to/sourcedir";
+ String restorePath = "/path/to/restoredir";
+ List snapshotInfoList =
+ Arrays.asList(
+ new ImportConfig.SnapshotInfo("snapdemo", "btdemo"),
+ new ImportConfig.SnapshotInfo("bookcontent-9087", "bookcontent"));
+
+ Map conbfiguration =
+ SnapshotTestHelper.buildMapFromList(
+ new String[] {"bigtable.row.size", "100", "bigtable.auth.type", "private"});
+
+ List snapshotConfigs =
+ SnapshotUtils.buildSnapshotConfigs(
+ snapshotInfoList, new HashMap<>(), projectId, sourcePath, restorePath);
+
+ assertThat(snapshotConfigs.size(), is(2));
+ assertThat(snapshotConfigs.get(0).getProjectId(), is(projectId));
+ assertThat(snapshotConfigs.get(0).getSnapshotName(), is("snapdemo"));
+ assertThat(snapshotConfigs.get(1).getSourceLocation(), is(sourcePath));
+ assertThat(snapshotConfigs.get(1).getTableName(), is("bookcontent"));
+ }
+
+ /** Tests that {@link SnapshotUtils#getSnapshotsFromString} parses single snapshot correctly. */
+ @Test
+ public void testGetSnapshotsFromStringReturnsSameTableName() {
+ String snapshotsWithBigtableTableName = "bookmark-2099";
+ Map snapshots =
+ SnapshotUtils.getSnapshotsFromString(snapshotsWithBigtableTableName);
+ assertThat(snapshots.size(), is(equalTo(1)));
+ assertThat(snapshots.get("bookmark-2099"), is("bookmark-2099"));
+ }
+
+ /**
+ * Tests that {@link SnapshotUtils#getSnapshotsFromString} parses multiple snapshots correctly.
+ */
+ @Test
+ public void testGetSnapshotsFromStringReturnsMultipleTables() {
+ String snapshotsWithBigtableTableName = "snapshot1,snapshot2,snapshot3:mytable3,snapshot4";
+ Map snapshots =
+ SnapshotUtils.getSnapshotsFromString(snapshotsWithBigtableTableName);
+ assertThat(snapshots.size(), is(equalTo(4)));
+ assertThat(snapshots.get("snapshot1"), is("snapshot1"));
+ assertThat(snapshots.get("snapshot2"), is("snapshot2"));
+ assertThat(snapshots.get("snapshot3"), is("mytable3"));
+ assertThat(snapshots.get("snapshot4"), is("snapshot4"));
+ }
+
+ /**
+ * Tests that {@link SnapshotUtils#getSnapshotsFromString} parses snapshot and table mapping
+ * correctly.
+ */
+ @Test
+ public void testGetSnapshotsFromStringReturnsParsedValues() {
+ String snapshotsWithBigtableTableName =
+ "bookmark-2099:bookmark,malwarescanstate-9087:malwarescan";
+ Map snapshots =
+ SnapshotUtils.getSnapshotsFromString(snapshotsWithBigtableTableName);
+ assertThat(snapshots.size(), is(equalTo(2)));
+ assertThat(snapshots.get("malwarescanstate-9087"), is("malwarescan"));
+ }
+
+ /** Tests that {@link SnapshotUtils#getSnapshotsFromString} throws exception on invalid format. */
+ @Test(expected = IllegalArgumentException.class)
+ public void testGetSnapshotsFromStringThrowsException() {
+ String snapshotsWithBigtableTableName =
+ "bookmark-2099:bookmark,malwarescanstate-9087:malwarescan:snapdemo1";
+ Map snapshots =
+ SnapshotUtils.getSnapshotsFromString(snapshotsWithBigtableTableName);
+ }
+
+ private void setUpGcsObjectMocks(List fakeStorageObjects) throws IOException {
+ Mockito.when(gcsObjects.getItems()).thenReturn(fakeStorageObjects);
+ Mockito.when(gcsUtilMock.listObjects(Mockito.anyString(), Mockito.anyString(), Mockito.any()))
+ .thenReturn(gcsObjects);
+ }
+
+ private Map getMatchingSnapshotsFromSnapshotPath(
+ List snapshotList, String prefix) throws IOException {
+ String baseObjectPath = "snapshots/20220309230526";
+ String importSnapshotpath = String.format("gs://sym-bucket/%s", baseObjectPath);
+ List fakeStorageObjects =
+ SnapshotTestHelper.createFakeStorageObjects(baseObjectPath, snapshotList);
+ setUpGcsObjectMocks(fakeStorageObjects);
+ return SnapshotUtils.getSnapshotsFromSnapshotPath(importSnapshotpath, gcsUtilMock, prefix);
+ }
+
+ /**
+ * Tests that {@link SnapshotUtils#getSnapshotsFromSnapshotPath} returns all snapshots when using
+ * wildcard.
+ */
+ @Test
+ public void testgetAllSnapshotsFromSnapshotPath() throws IOException {
+ List snapshotList = Arrays.asList("audit-events", "dlpInfo", "ce-metrics-manifest");
+ Map snapshots = getMatchingSnapshotsFromSnapshotPath(snapshotList, "*");
+ assertThat(snapshots.size(), is(equalTo(3)));
+ assertThat(snapshots.keySet(), containsInAnyOrder(snapshotList.toArray(new String[0])));
+ }
+
+ /**
+ * Tests that {@link SnapshotUtils#getSnapshotsFromSnapshotPath} returns subset of snapshots
+ * matching regex.
+ */
+ @Test
+ public void testgetSubSetSnapshotsFromSnapshotPath() throws IOException {
+ List snapshotList =
+ Arrays.asList(
+ "snapshot-audit-events",
+ "snapshot-attachments",
+ "snapshot-ce-metrics-manifest",
+ "snapshot-attachments-streams");
+ Map snapshots =
+ getMatchingSnapshotsFromSnapshotPath(snapshotList, ".*attachments.*");
+ List expectedResult =
+ snapshotList.stream().filter(e -> e.contains("attachments")).collect(Collectors.toList());
+ assertThat(snapshots.size(), is(equalTo(expectedResult.size())));
+ assertThat(snapshots.keySet(), containsInAnyOrder(expectedResult.toArray(new String[0])));
+ }
+
+ /**
+ * Tests that {@link SnapshotUtils#getSnapshotsFromSnapshotPath} returns empty map when no objects
+ * found.
+ */
+ @Test
+ public void testgetSubSetSnapshotsFromSnapshotPathReturnsEmptyMap() throws IOException {
+ Map snapshots = getMatchingSnapshotsFromSnapshotPath(null, "*");
+ assertThat(snapshots.size(), is(0));
+ }
+
+ /**
+ * Tests that {@link SnapshotUtils#getSnapshotsFromSnapshotPath} handles GCS pagination correctly.
+ */
+ @Test
+ public void testGetSnapshotsFromSnapshotPath_pagination() throws IOException {
+ String baseObjectPath = "snapshots/20220309230526";
+ String importSnapshotpath = String.format("gs://sym-bucket/%s", baseObjectPath);
+
+ StorageObject obj1 =
+ new StorageObject()
+ .setId(baseObjectPath + "/.hbase-snapshot/snap1/file1")
+ .setName(baseObjectPath + "/.hbase-snapshot/snap1/file1");
+ StorageObject obj2 =
+ new StorageObject()
+ .setId(baseObjectPath + "/.hbase-snapshot/snap2/file2")
+ .setName(baseObjectPath + "/.hbase-snapshot/snap2/file2");
+
+ Objects page1 = new Objects().setItems(Arrays.asList(obj1)).setNextPageToken("token");
+ Objects page2 = new Objects().setItems(Arrays.asList(obj2)).setNextPageToken(null);
+
+ Mockito.when(
+ gcsUtilMock.listObjects(
+ Mockito.eq("sym-bucket"), Mockito.anyString(), Mockito.isNull()))
+ .thenReturn(page1);
+ Mockito.when(
+ gcsUtilMock.listObjects(
+ Mockito.eq("sym-bucket"), Mockito.anyString(), Mockito.eq("token")))
+ .thenReturn(page2);
+
+ Map snapshots =
+ SnapshotUtils.getSnapshotsFromSnapshotPath(importSnapshotpath, gcsUtilMock, "*");
+
+ assertThat(snapshots.size(), is(2));
+ assertThat(snapshots.keySet(), containsInAnyOrder("snap1", "snap2"));
+ }
+
+ /**
+ * Tests that {@link SnapshotUtils#getSnapshotsFromSnapshotPath} returns matching snapshot for
+ * literal string.
+ */
+ @Test
+ public void testGetSnapshotsFromSnapshotPath_literal() throws IOException {
+ List snapshotList = Arrays.asList("snap1", "snap2");
+ Map snapshots = getMatchingSnapshotsFromSnapshotPath(snapshotList, "snap1");
+ assertThat(snapshots.size(), is(equalTo(1)));
+ assertThat(snapshots.get("snap1"), is("snap1"));
+ }
+
+ /**
+ * Tests that {@link SnapshotUtils#getSnapshotsFromSnapshotPath} returns matching snapshots for
+ * simple glob pattern.
+ */
+ @Test
+ public void testGetSnapshotsFromSnapshotPath_simpleGlob() throws IOException {
+ List snapshotList = Arrays.asList("snap1", "snap2", "othersnap");
+ Map snapshots = getMatchingSnapshotsFromSnapshotPath(snapshotList, "snap*");
+ assertThat(snapshots.size(), is(equalTo(2)));
+ assertThat(snapshots.keySet(), containsInAnyOrder("snap1", "snap2"));
+ }
+
+ /**
+ * Tests that {@link SnapshotUtils#getSnapshotsFromSnapshotPath} correctly lists snapshots from a
+ * local directory.
+ */
+ @Test
+ public void testGetSnapshotsFromLocalPath() throws IOException {
+ java.io.File localBase = tempFolder.newFolder("local_snapshots");
+ java.io.File manifestDir = new java.io.File(localBase, ".hbase-snapshot");
+ org.junit.Assert.assertTrue(manifestDir.mkdirs());
+
+ // Create folders representing snapshots
+ org.junit.Assert.assertTrue(new java.io.File(manifestDir, "snap1").mkdirs());
+ org.junit.Assert.assertTrue(new java.io.File(manifestDir, "snap2").mkdirs());
+ org.junit.Assert.assertTrue(new java.io.File(manifestDir, "othersnap").mkdirs());
+
+ // 1. Wildcard match
+ Map snapshots =
+ SnapshotUtils.getSnapshotsFromSnapshotPath(localBase.getAbsolutePath(), null, "*");
+ assertThat(snapshots.size(), is(equalTo(3)));
+ assertThat(snapshots.keySet(), containsInAnyOrder("snap1", "snap2", "othersnap"));
+
+ // 2. Glob match
+ snapshots =
+ SnapshotUtils.getSnapshotsFromSnapshotPath(localBase.getAbsolutePath(), null, "snap*");
+ assertThat(snapshots.size(), is(equalTo(2)));
+ assertThat(snapshots.keySet(), containsInAnyOrder("snap1", "snap2"));
+
+ // 3. Literal match
+ snapshots =
+ SnapshotUtils.getSnapshotsFromSnapshotPath(localBase.getAbsolutePath(), null, "snap1");
+ assertThat(snapshots.size(), is(equalTo(1)));
+ assertThat(snapshots.keySet(), contains("snap1"));
+ }
+}
diff --git a/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/coders/RegionConfigCoderTest.java b/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/coders/RegionConfigCoderTest.java
new file mode 100644
index 0000000000..27899ef2a0
--- /dev/null
+++ b/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/coders/RegionConfigCoderTest.java
@@ -0,0 +1,105 @@
+/*
+ * Copyright 2026 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.coders;
+
+import static org.junit.Assert.assertArrayEquals;
+import static org.junit.Assert.assertEquals;
+
+import com.google.cloud.bigtable.beam.hbasesnapshots.SnapshotTestHelper;
+import com.google.cloud.bigtable.beam.hbasesnapshots.conf.RegionConfig;
+import com.google.cloud.bigtable.beam.hbasesnapshots.conf.SnapshotConfig;
+import java.io.ByteArrayInputStream;
+import java.io.ByteArrayOutputStream;
+import org.apache.hadoop.hbase.TableName;
+import org.apache.hadoop.hbase.client.RegionInfo;
+import org.apache.hadoop.hbase.client.RegionInfoBuilder;
+import org.apache.hadoop.hbase.client.TableDescriptor;
+import org.apache.hadoop.hbase.client.TableDescriptorBuilder;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Tests the {@link RegionConfigCoder} to ensure correct serialization and deserialization. */
+@RunWith(JUnit4.class)
+public class RegionConfigCoderTest {
+
+ /**
+ * Tests that {@link RegionConfigCoder} correctly encodes and decodes a {@link RegionConfig}
+ * object with empty region start and end keys.
+ */
+ @Test
+ public void testEncodeDecode_emptyKeys() throws Exception {
+ SnapshotConfig snapshotConfig = SnapshotTestHelper.newSnapshotConfig("my-restore-path");
+
+ TableName tableName = TableName.valueOf("my-table");
+ RegionInfo regionInfo = RegionInfoBuilder.newBuilder(tableName).build();
+ TableDescriptor tableDescriptor = TableDescriptorBuilder.newBuilder(tableName).build();
+
+ RegionConfig originalConfig =
+ RegionConfig.builder()
+ .setSnapshotConfig(snapshotConfig)
+ .setRegionInfo(regionInfo)
+ .setTableDescriptor(tableDescriptor)
+ .setRegionSize(12345L)
+ .build();
+
+ RegionConfigCoder coder = new RegionConfigCoder();
+
+ ByteArrayOutputStream outStream = new ByteArrayOutputStream();
+ coder.encode(originalConfig, outStream);
+
+ ByteArrayInputStream inStream = new ByteArrayInputStream(outStream.toByteArray());
+ RegionConfig decodedConfig = coder.decode(inStream);
+
+ assertEquals(originalConfig, decodedConfig);
+ }
+
+ /**
+ * Tests that {@link RegionConfigCoder} correctly encodes and decodes a {@link RegionConfig}
+ * object containing region start and end keys which are populated with raw bytes.
+ */
+ @Test
+ public void testEncodeDecode_populatedKeys() throws Exception {
+ SnapshotConfig snapshotConfig = SnapshotTestHelper.newSnapshotConfig("my-restore-path");
+
+ TableName tableName = TableName.valueOf("my-table");
+ byte[] startKey = "start-key-abc".getBytes();
+ byte[] endKey = "end-key-xyz".getBytes();
+ RegionInfo regionInfo =
+ RegionInfoBuilder.newBuilder(tableName).setStartKey(startKey).setEndKey(endKey).build();
+ TableDescriptor tableDescriptor = TableDescriptorBuilder.newBuilder(tableName).build();
+
+ RegionConfig originalConfig =
+ RegionConfig.builder()
+ .setSnapshotConfig(snapshotConfig)
+ .setRegionInfo(regionInfo)
+ .setTableDescriptor(tableDescriptor)
+ .setRegionSize(12345L)
+ .build();
+
+ RegionConfigCoder coder = new RegionConfigCoder();
+
+ ByteArrayOutputStream outStream = new ByteArrayOutputStream();
+ coder.encode(originalConfig, outStream);
+
+ ByteArrayInputStream inStream = new ByteArrayInputStream(outStream.toByteArray());
+ RegionConfig decodedConfig = coder.decode(inStream);
+
+ assertEquals(originalConfig, decodedConfig);
+ assertArrayEquals(startKey, decodedConfig.getRegionInfo().getStartKey());
+ assertArrayEquals(endKey, decodedConfig.getRegionInfo().getEndKey());
+ }
+}
diff --git a/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/conf/ImportConfigTest.java b/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/conf/ImportConfigTest.java
new file mode 100644
index 0000000000..2fe0d836fa
--- /dev/null
+++ b/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/conf/ImportConfigTest.java
@@ -0,0 +1,392 @@
+/*
+ * Copyright 2026 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.conf;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.fail;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class ImportConfigTest {
+
+ private ImportConfig createValidConfig() {
+ ImportConfig config = new ImportConfig();
+ config.setSourcepath("gs://my-bucket/hbase-export");
+ java.util.Map snapshots = new java.util.HashMap<>();
+ snapshots.put("snapshot1", "table1");
+ config.setSnapshotsFromMap(snapshots);
+ return config;
+ }
+
+ @Test
+ public void testDefaultConfigIsValid() {
+ ImportConfig config = createValidConfig();
+ config.validate();
+ assertEquals(5000L, config.getBackoffInitialIntervalInMillis());
+ assertEquals(180000L, config.getBackoffMaxIntervalInMillis());
+ assertEquals(3, config.getBackoffMaxretries());
+ }
+
+ @Test
+ public void testSetValidBackoffValues() {
+ ImportConfig config = createValidConfig();
+ config.setBackoffInitialIntervalInMillis(1000L);
+ config.setBackoffMaxIntervalInMillis(5000L);
+ config.setBackoffMaxretries(5);
+
+ assertEquals(1000L, config.getBackoffInitialIntervalInMillis());
+ assertEquals(5000L, config.getBackoffMaxIntervalInMillis());
+ assertEquals(5, config.getBackoffMaxretries());
+ }
+
+ @Test
+ public void testInitialIntervalGreaterThanMaxInterval_ThrowsException() {
+ ImportConfig config = createValidConfig();
+ try {
+ config.setBackoffInitialIntervalInMillis(200000L); // Default max is 180000L
+ fail("Should have thrown IllegalArgumentException");
+ } catch (IllegalArgumentException expected) {
+ // Expected
+ }
+
+ try {
+ config.setBackoffMaxIntervalInMillis(1000L); // Default initial is 5000L
+ fail("Should have thrown IllegalArgumentException");
+ } catch (IllegalArgumentException expected) {
+ // Expected
+ }
+ }
+
+ @Test
+ public void testNegativeAndZeroValues_ThrowsException() {
+ ImportConfig config = createValidConfig();
+ try {
+ config.setBackoffInitialIntervalInMillis(-10L);
+ fail("Should have thrown IllegalArgumentException");
+ } catch (IllegalArgumentException expected) {
+ // Expected
+ }
+
+ try {
+ config.setBackoffInitialIntervalInMillis(0L);
+ fail("Should have thrown IllegalArgumentException");
+ } catch (IllegalArgumentException expected) {
+ // Expected
+ }
+
+ try {
+ config.setBackoffMaxIntervalInMillis(-10L);
+ fail("Should have thrown IllegalArgumentException");
+ } catch (IllegalArgumentException expected) {
+ // Expected
+ }
+
+ try {
+ config.setBackoffMaxIntervalInMillis(0L);
+ fail("Should have thrown IllegalArgumentException");
+ } catch (IllegalArgumentException expected) {
+ // Expected
+ }
+
+ try {
+ config.setBackoffMaxretries(-1);
+ fail("Should have thrown IllegalArgumentException");
+ } catch (IllegalArgumentException expected) {
+ // Expected
+ }
+ }
+
+ @Test
+ public void testValidateMethodCatchesDivergentValues() {
+ ImportConfig config = createValidConfig();
+
+ config.setBackoffMaxIntervalInMillis(10000L);
+ config.setBackoffInitialIntervalInMillis(2000L);
+ config.validate(); // Should pass
+ }
+
+ @Test
+ public void testValidate_EmptySourcePath_ThrowsException() {
+ ImportConfig config = createValidConfig();
+ config.setSourcepath(" ");
+ try {
+ config.validate();
+ fail("Should have thrown IllegalArgumentException for blank sourcepath");
+ } catch (IllegalArgumentException expected) {
+ assertEquals(
+ "Source Path containing hbase snapshots must be specified.", expected.getMessage());
+ }
+ }
+
+ @Test
+ public void testValidate_EmptyRestorePath_ThrowsException() {
+ ImportConfig config = createValidConfig();
+ config.setRestorepath(" ");
+ try {
+ config.validate();
+ fail("Should have thrown IllegalArgumentException for blank restorepath");
+ } catch (IllegalArgumentException expected) {
+ assertEquals("restorepath cannot be empty if specified", expected.getMessage());
+ }
+
+ // null restorepath should be valid (will use default path during setup)
+ config.setRestorepath(null);
+ config.validate();
+ }
+
+ @Test
+ public void testValidate_EmptyRunStatusPath_ThrowsException() {
+ ImportConfig config = createValidConfig();
+ config.setRunstatuspath(" ");
+ try {
+ config.validate();
+ fail("Should have thrown IllegalArgumentException for blank runstatuspath");
+ } catch (IllegalArgumentException expected) {
+ assertEquals("runstatuspath cannot be empty if specified", expected.getMessage());
+ }
+
+ config.setRunstatuspath(null);
+ config.validate();
+ }
+
+ @Test
+ public void testValidate_EmptySnapshots_ThrowsException() {
+ ImportConfig config = createValidConfig();
+ config.setSnapshots(new java.util.ArrayList<>());
+ try {
+ config.validate();
+ fail("Should have thrown IllegalArgumentException for empty snapshots");
+ } catch (IllegalArgumentException expected) {
+ assertEquals(
+ "Snapshots must be specified. Allowed values are '*' (indicating all snapshots under"
+ + " source path) or 'prefix*' (snapshots matching certain prefix) or"
+ + " 'snapshotname1:tablename1,snapshotname2:tablename2' (comma seperated list of"
+ + " snapshots)",
+ expected.getMessage());
+ }
+ }
+
+ @Test
+ public void testValidate_EmptySnapshotNameOrTableName_ThrowsException() {
+ ImportConfig config = createValidConfig();
+
+ // Empty snapshot name
+ java.util.List invalidSnapshots = new java.util.ArrayList<>();
+ invalidSnapshots.add(new ImportConfig.SnapshotInfo("", "table1"));
+ config.setSnapshots(invalidSnapshots);
+ try {
+ config.validate();
+ fail("Should have thrown IllegalArgumentException for empty snapshot name");
+ } catch (IllegalArgumentException expected) {
+ assertEquals("snapshotName inside snapshots cannot be null or empty", expected.getMessage());
+ }
+
+ // Empty table name
+ invalidSnapshots.clear();
+ invalidSnapshots.add(new ImportConfig.SnapshotInfo("snapshot1", " "));
+ config.setSnapshots(invalidSnapshots);
+ try {
+ config.validate();
+ fail("Should have thrown IllegalArgumentException for empty table name");
+ } catch (IllegalArgumentException expected) {
+ assertEquals(
+ "bigtableTableName inside snapshots cannot be null or empty", expected.getMessage());
+ }
+ }
+
+ @Test
+ public void testValidate_InvalidConfigurationKeysOrValues_ThrowsException() {
+ ImportConfig config = createValidConfig();
+
+ // Invalid hbaseConfiguration key
+ java.util.Map hbaseConfig = new java.util.HashMap<>();
+ hbaseConfig.put("", "value");
+ config.setHbaseConfiguration(hbaseConfig);
+ try {
+ config.validate();
+ fail("Should have thrown IllegalArgumentException for empty hbaseConfiguration key");
+ } catch (IllegalArgumentException expected) {
+ assertEquals("hbaseConfiguration keys cannot be null or empty", expected.getMessage());
+ }
+
+ // Invalid hbaseConfiguration value
+ hbaseConfig.clear();
+ hbaseConfig.put("key", null);
+ config.setHbaseConfiguration(hbaseConfig);
+ try {
+ config.validate();
+ fail("Should have thrown IllegalArgumentException for null hbaseConfiguration value");
+ } catch (IllegalArgumentException expected) {
+ assertEquals("hbaseConfiguration values cannot be null", expected.getMessage());
+ }
+
+ // Reset hbaseConfig to valid
+ config.setHbaseConfiguration(null);
+
+ // Invalid bigtableConfiguration key
+ java.util.Map btConfig = new java.util.HashMap<>();
+ btConfig.put(" ", "value");
+ config.setBigtableConfiguration(btConfig);
+ try {
+ config.validate();
+ fail("Should have thrown IllegalArgumentException for empty bigtableConfiguration key");
+ } catch (IllegalArgumentException expected) {
+ assertEquals("bigtableConfiguration keys cannot be null or empty", expected.getMessage());
+ }
+
+ // Invalid bigtableConfiguration value
+ btConfig.clear();
+ btConfig.put("key", null);
+ config.setBigtableConfiguration(btConfig);
+ try {
+ config.validate();
+ fail("Should have thrown IllegalArgumentException for null bigtableConfiguration value");
+ } catch (IllegalArgumentException expected) {
+ assertEquals("bigtableConfiguration values cannot be null", expected.getMessage());
+ }
+ }
+
+ @Test
+ public void testValidate_InvalidURIs_ThrowsException() {
+ ImportConfig config = createValidConfig();
+
+ // Invalid sourcepath URI
+ config.setSourcepath("gs://bucket[with]invalidchars");
+ try {
+ config.validate();
+ fail("Should have thrown IllegalArgumentException for invalid sourcepath URI");
+ } catch (IllegalArgumentException expected) {
+ // Expected, URISyntaxException message
+ }
+
+ // Invalid restorepath URI
+ config = createValidConfig();
+ config.setRestorepath("gs://bucket[with]invalidchars");
+ try {
+ config.validate();
+ fail("Should have thrown IllegalArgumentException for invalid restorepath URI");
+ } catch (IllegalArgumentException expected) {
+ // Expected
+ }
+
+ // Invalid runstatuspath URI
+ config = createValidConfig();
+ config.setRunstatuspath("gs://bucket[with]invalidchars");
+ try {
+ config.validate();
+ fail("Should have thrown IllegalArgumentException for invalid runstatuspath URI");
+ } catch (IllegalArgumentException expected) {
+ // Expected
+ }
+ }
+
+ @Test
+ public void testValidate_AutoRestorepathLacksParent_ThrowsException() {
+ ImportConfig config = createValidConfig();
+
+ // gs:// scheme without object path (i.e. bucket root)
+ config.setSourcepath("gs://my-bucket");
+ config.setRestorepath(null);
+ try {
+ config.validate();
+ fail(
+ "Should have thrown IllegalArgumentException for gs bucket root sourcepath when"
+ + " restorepath is null");
+ } catch (IllegalArgumentException expected) {
+ assertEquals(
+ "sourcepath must have a parent directory to auto-generate restorepath: gs://my-bucket",
+ expected.getMessage());
+ }
+
+ config.setSourcepath("gs://my-bucket/");
+ try {
+ config.validate();
+ fail(
+ "Should have thrown IllegalArgumentException for gs bucket root sourcepath (with slash)"
+ + " when restorepath is null");
+ } catch (IllegalArgumentException expected) {
+ assertEquals(
+ "sourcepath must have a parent directory to auto-generate restorepath: gs://my-bucket/",
+ expected.getMessage());
+ }
+
+ // local relative path without parent
+ config.setSourcepath("sourcedir");
+ try {
+ config.validate();
+ fail(
+ "Should have thrown IllegalArgumentException for relative path without parent when"
+ + " restorepath is null");
+ } catch (IllegalArgumentException expected) {
+ assertEquals(
+ "sourcepath must have a parent directory to auto-generate restorepath: sourcedir",
+ expected.getMessage());
+ }
+
+ // local relative path with parent should succeed
+ config.setSourcepath("parent/sourcedir");
+ config.validate(); // Should pass
+ }
+
+ @Test
+ public void testValidate_DuplicateSnapshotName_ThrowsException() {
+ ImportConfig config = createValidConfig();
+ java.util.List snapshots = new java.util.ArrayList<>();
+ snapshots.add(new ImportConfig.SnapshotInfo("snap1", "table1"));
+ snapshots.add(new ImportConfig.SnapshotInfo("snap1", "table2"));
+ config.setSnapshots(snapshots);
+
+ try {
+ config.validate();
+ fail("Should have thrown IllegalArgumentException for duplicate snapshot name");
+ } catch (IllegalArgumentException expected) {
+ assertEquals("Duplicate snapshot name detected: snap1", expected.getMessage());
+ }
+ }
+
+ @Test
+ public void testValidate_NullSourcePath_ThrowsNullPointerException() {
+ ImportConfig config = createValidConfig();
+ config.setSourcepath(null);
+ try {
+ config.validate();
+ fail("Should have thrown NullPointerException for null sourcepath");
+ } catch (NullPointerException expected) {
+ assertEquals(
+ "Source Path containing hbase snapshots must be specified.", expected.getMessage());
+ }
+ }
+
+ @Test
+ public void testValidate_NullSnapshots_ThrowsNullPointerException() {
+ ImportConfig config = createValidConfig();
+ config.setSnapshots(null);
+ try {
+ config.validate();
+ fail("Should have thrown NullPointerException for null snapshots");
+ } catch (NullPointerException expected) {
+ assertEquals(
+ "Snapshots must be specified. Allowed values are '*' (indicating all snapshots under"
+ + " source path) or 'prefix*' (snapshots matching certain prefix) or"
+ + " 'snapshotname1:tablename1,snapshotname2:tablename2' (comma seperated list of"
+ + " snapshots)",
+ expected.getMessage());
+ }
+ }
+}
diff --git a/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/CleanupHBaseSnapshotRestoreFilesFnTest.java b/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/dofn/CleanupHBaseSnapshotRestoreFilesFnTest.java
similarity index 92%
rename from bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/CleanupHBaseSnapshotRestoreFilesFnTest.java
rename to bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/dofn/CleanupHBaseSnapshotRestoreFilesFnTest.java
index 0183f856f1..ed74fd66da 100644
--- a/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/CleanupHBaseSnapshotRestoreFilesFnTest.java
+++ b/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/dofn/CleanupHBaseSnapshotRestoreFilesFnTest.java
@@ -13,11 +13,12 @@
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-package com.google.cloud.bigtable.beam.hbasesnapshots;
+package com.google.cloud.bigtable.beam.hbasesnapshots.dofn;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertThrows;
+import com.google.cloud.bigtable.beam.hbasesnapshots.conf.HBaseSnapshotInputConfigBuilder;
import java.util.UUID;
import org.junit.Test;
diff --git a/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/dofn/CleanupRestoredSnapshotsFnTest.java b/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/dofn/CleanupRestoredSnapshotsFnTest.java
new file mode 100644
index 0000000000..f68d1f6c0f
--- /dev/null
+++ b/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/dofn/CleanupRestoredSnapshotsFnTest.java
@@ -0,0 +1,124 @@
+/*
+ * Copyright 2026 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.dofn;
+
+import com.google.cloud.bigtable.beam.hbasesnapshots.SnapshotTestHelper;
+import com.google.cloud.bigtable.beam.hbasesnapshots.conf.SnapshotConfig;
+import java.io.File;
+import java.io.IOException;
+import org.apache.beam.sdk.testing.TestPipeline;
+import org.apache.beam.sdk.transforms.Create;
+import org.apache.beam.sdk.transforms.ParDo;
+import org.apache.beam.sdk.util.Sleeper;
+import org.junit.Assert;
+import org.junit.ClassRule;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/** Tests the {@link CleanupRestoredSnapshotsFn} functionality. */
+@RunWith(JUnit4.class)
+public class CleanupRestoredSnapshotsFnTest {
+ private static final Logger LOG = LoggerFactory.getLogger(CleanupRestoredSnapshotsFnTest.class);
+
+ @Rule public final transient TestPipeline pipeline = TestPipeline.create();
+ @ClassRule public static TemporaryFolder tempFolder = new TemporaryFolder();
+
+ @Test
+ public void testDeleteRestoredSnapshot() throws Exception {
+ File restoreDir = tempFolder.newFolder();
+ if (restoreDir.exists()) {
+ LOG.info("Created temp folder: {}", restoreDir.getAbsolutePath());
+ SnapshotConfig snapshotConfig =
+ SnapshotTestHelper.newSnapshotConfig(restoreDir.getAbsolutePath());
+ new CleanupRestoredSnapshotsFn(5000, 180000, 3).processElement(snapshotConfig, null);
+ Assert.assertFalse(restoreDir.exists());
+ } else {
+ LOG.warn(
+ "Skipping CleanUpRestoredSnapshotsTest since temporary file was unable to be created in"
+ + " restore path: {}",
+ restoreDir.getAbsolutePath());
+ }
+ }
+
+ /**
+ * Tests CleanupRestoredSnapshots with invalid path to verify exception is handled internally
+ *
+ * @throws Exception
+ */
+ // Use a custom subclass to override the Sleeper and avoid sleeping for ~35 seconds
+ // during tests that trigger the retry loop (due to hardcoded exponential backoff).
+ private static class FastCleanupRestoredSnapshots extends CleanupRestoredSnapshotsFn {
+ public FastCleanupRestoredSnapshots() {
+ super(5000, 180000, 3);
+ }
+
+ @Override
+ Sleeper getSleeper() {
+ return new Sleeper() {
+ @Override
+ public void sleep(long millis) throws InterruptedException {
+ // Do nothing!
+ }
+ };
+ }
+ }
+
+ private static class NpeCleanupRestoredSnapshots extends CleanupRestoredSnapshotsFn {
+ public NpeCleanupRestoredSnapshots() {
+ super(5000, 180000, 3);
+ }
+
+ @Override
+ void cleanupSnapshot(
+ SnapshotConfig snapshotConfig,
+ org.apache.hadoop.fs.FileSystem fileSystem,
+ org.apache.hadoop.fs.Path restorePath)
+ throws IOException {
+ throw new NullPointerException("Simulated NPE");
+ }
+ }
+
+ @Test
+ public void testDeleteRestoredSnapshotWithInvalidPath() throws Exception {
+ pipeline
+ .apply("CreateInput", Create.of(SnapshotTestHelper.newSnapshotConfig("invalid_path")))
+ .apply("DeleteSnapshot", ParDo.of(new FastCleanupRestoredSnapshots()));
+
+ // The pipeline should run successfully without throwing an exception.
+ // The CleanupRestoredSnapshotsFn DoFn handles exceptions internally (logs them after retries)
+ // and does not fail the job.
+ pipeline.run();
+ }
+
+ @Test
+ public void testDeleteRestoredSnapshotWithNPE() throws Exception {
+ pipeline
+ .apply("CreateInput", Create.of(SnapshotTestHelper.newSnapshotConfig("any_path")))
+ .apply("DeleteSnapshot", ParDo.of(new NpeCleanupRestoredSnapshots()));
+
+ try {
+ pipeline.run();
+ Assert.fail("Expected PipelineExecutionException");
+ } catch (org.apache.beam.sdk.Pipeline.PipelineExecutionException e) {
+ Assert.assertTrue(e.getCause() instanceof NullPointerException);
+ }
+ }
+}
diff --git a/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/dofn/CreateBigtableMutationsFnTest.java b/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/dofn/CreateBigtableMutationsFnTest.java
new file mode 100644
index 0000000000..2b05a237f4
--- /dev/null
+++ b/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/dofn/CreateBigtableMutationsFnTest.java
@@ -0,0 +1,448 @@
+/*
+ * Copyright 2026 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.dofn;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+
+import com.google.cloud.bigtable.beam.hbasesnapshots.conf.SnapshotConfig;
+import com.google.cloud.bigtable.beam.hbasesnapshots.conf.SnapshotKey;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.Iterator;
+import java.util.List;
+import org.apache.beam.sdk.metrics.MetricName;
+import org.apache.beam.sdk.metrics.MetricsContainer;
+import org.apache.beam.sdk.metrics.MetricsEnvironment;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.values.KV;
+import org.apache.hadoop.hbase.Cell;
+import org.apache.hadoop.hbase.CellUtil;
+import org.apache.hadoop.hbase.KeyValue;
+import org.apache.hadoop.hbase.client.Mutation;
+import org.apache.hadoop.hbase.client.Result;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mockito;
+
+/** Tests the {@link CreateBigtableMutationsFn} functionality. */
+@RunWith(JUnit4.class)
+public class CreateBigtableMutationsFnTest {
+
+ private SnapshotConfig snapshotConfig;
+ private SnapshotKey snapshotKey;
+ private Result result;
+ private DoFn.OutputReceiver>> receiver;
+
+ @Before
+ public void setUp() {
+ snapshotConfig =
+ SnapshotConfig.builder()
+ .setProjectId("project")
+ .setSourceLocation("source")
+ .setSnapshotName("my-snapshot")
+ .setTableName("my-table")
+ .setRestoreLocation("restore")
+ .setConfigurationDetails(Collections.emptyMap())
+ .build();
+ snapshotKey = SnapshotKey.create("my-snapshot", "my-table");
+ receiver = mock(DoFn.OutputReceiver.class);
+ }
+
+ /**
+ * Tests that {@link CreateBigtableMutations#processElement} successfully converts a simple HBase
+ * result to Bigtable mutations.
+ */
+ @Test
+ public void testProcessElement() throws Exception {
+ byte[] rowKey = "row-key".getBytes();
+ Cell cell = new KeyValue(rowKey, "cf".getBytes(), "qual".getBytes(), "val".getBytes());
+
+ Result localResult = Result.create(Collections.singletonList(cell));
+
+ CreateBigtableMutationsFn fn =
+ new CreateBigtableMutationsFn(100, false, 0L, false, 0, false, 0, false);
+
+ fn.processElement(KV.of(snapshotKey, localResult), receiver);
+ ArgumentCaptor>> captor = ArgumentCaptor.forClass(KV.class);
+ verify(receiver, times(1)).output(captor.capture());
+ KV> kv = captor.getValue();
+ assertEquals("my-table", kv.getKey());
+ Iterator mutationIt = kv.getValue().iterator();
+ assertEquals(true, mutationIt.hasNext());
+ Mutation mutation = mutationIt.next();
+ assertEquals("row-key", new String(mutation.getRow()));
+ assertEquals(false, mutationIt.hasNext());
+ }
+
+ /** Tests that {@link CreateBigtableMutations#processElement} skips results with no cells. */
+ @Test
+ public void testProcessElement_emptyCells() throws Exception {
+ Result localResult = Result.create(Collections.emptyList());
+
+ CreateBigtableMutationsFn fn =
+ new CreateBigtableMutationsFn(100, false, 0L, false, 0, false, 0, false);
+
+ fn.processElement(KV.of(snapshotKey, localResult), receiver);
+ verify(receiver, times(0)).output(Mockito.any());
+ }
+
+ /**
+ * Tests that {@link CreateBigtableMutations#processElement} filters out cells that exceed the
+ * configured size threshold.
+ */
+ @Test
+ public void testProcessElement_filterLargeCell() throws Exception {
+ byte[] rowKey = "row-key".getBytes();
+ byte[] largeValue = new byte[1000];
+ Cell cell = new KeyValue(rowKey, "cf".getBytes(), "qual".getBytes(), largeValue);
+
+ result = Result.create(Collections.singletonList(cell));
+
+ CreateBigtableMutationsFn fn =
+ new CreateBigtableMutationsFn(
+ 100, false, 0L, true, 500, // filterLargeCells enabled, threshold 500
+ false, 0, false);
+
+ fn.processElement(KV.of(snapshotKey, result), receiver);
+
+ verify(receiver, times(0)).output(Mockito.any());
+ }
+
+ /**
+ * Tests that {@link CreateBigtableMutations#processElement} drops the row if the total size of
+ * its cells exceeds the configured threshold.
+ */
+ @Test
+ public void testProcessElement_filterLargeRow() throws Exception {
+ byte[] rowKey = "row-key".getBytes();
+ Cell cell = new KeyValue(rowKey, "cf".getBytes(), "qual".getBytes(), "val".getBytes());
+
+ result = Result.create(Collections.singletonList(cell));
+
+ CreateBigtableMutationsFn fn =
+ new CreateBigtableMutationsFn(
+ 100, true, 10L, false, 0, false, 0, false // filterLargeRows enabled, threshold 10 bytes
+ );
+
+ fn.processElement(KV.of(snapshotKey, result), receiver);
+
+ verify(receiver, times(0)).output(Mockito.any());
+ }
+
+ /**
+ * Tests that {@link CreateBigtableMutations#processElement} drops the row if the row key length
+ * exceeds the configured threshold.
+ */
+ @Test
+ public void testProcessElement_filterLargeRowKey() throws Exception {
+ byte[] rowKey = "row-key".getBytes(); // 7 bytes
+ Cell cell = new KeyValue(rowKey, "cf".getBytes(), "qual".getBytes(), "val".getBytes());
+
+ result = Result.create(Collections.singletonList(cell));
+
+ CreateBigtableMutationsFn fn =
+ new CreateBigtableMutationsFn(
+ 100,
+ false,
+ 0L,
+ false,
+ 0,
+ true,
+ 5,
+ false // filterLargeRowKeys enabled, threshold 5 bytes
+ );
+
+ fn.processElement(KV.of(snapshotKey, result), receiver);
+
+ verify(receiver, times(0)).output(Mockito.any());
+ }
+
+ /**
+ * Tests that {@link CreateBigtableMutations#processElement} splits mutations into multiple Put
+ * requests if the number of cells exceeds the threshold.
+ */
+ @Test
+ public void testProcessElement_splitMutations() throws Exception {
+ byte[] rowKey = "row-key".getBytes();
+ Cell cell1 = new KeyValue(rowKey, "cf".getBytes(), "qual1".getBytes(), "val1".getBytes());
+ Cell cell2 = new KeyValue(rowKey, "cf".getBytes(), "qual2".getBytes(), "val2".getBytes());
+
+ List cells = new ArrayList<>();
+ cells.add(cell1);
+ cells.add(cell2);
+
+ result = Result.create(cells);
+
+ CreateBigtableMutationsFn fn =
+ new CreateBigtableMutationsFn(
+ 1, // maxMutationsPerRequestThreshold = 1
+ false, 0L, false, 0, false, 0, false);
+
+ fn.processElement(KV.of(snapshotKey, result), receiver);
+
+ ArgumentCaptor>> captor = ArgumentCaptor.forClass(KV.class);
+
+ verify(receiver, times(1)).output(captor.capture());
+
+ KV> output = captor.getValue();
+ assertNotNull(output);
+
+ Iterator iterator = output.getValue().iterator();
+ assertEquals(true, iterator.hasNext());
+ Mutation mutation1 = iterator.next();
+ assertEquals("row-key", new String(mutation1.getRow()));
+ // Validate cells in mutation1
+ List cells1 = mutation1.getFamilyCellMap().values().iterator().next();
+ assertEquals(1, cells1.size());
+ Cell c1 = cells1.get(0);
+ assertEquals(
+ "qual1",
+ new String(c1.getQualifierArray(), c1.getQualifierOffset(), c1.getQualifierLength()));
+
+ assertEquals(true, iterator.hasNext());
+ Mutation mutation2 = iterator.next();
+ assertEquals("row-key", new String(mutation2.getRow()));
+ // Validate cells in mutation2
+ List cells2 = mutation2.getFamilyCellMap().values().iterator().next();
+ assertEquals(1, cells2.size());
+ Cell c2 = cells2.get(0);
+ assertEquals(
+ "qual2",
+ new String(c2.getQualifierArray(), c2.getQualifierOffset(), c2.getQualifierLength()));
+
+ assertEquals(false, iterator.hasNext());
+ }
+
+ /**
+ * Tests that {@link CreateBigtableMutationsFn#processElement} drops wide rows if filterWideRows
+ * is enabled and cell count exceeds threshold.
+ */
+ @Test
+ public void testProcessElement_filterWideRows_exceedsThreshold() throws Exception {
+ byte[] rowKey = "row-key".getBytes();
+ Cell cell1 = new KeyValue(rowKey, "cf".getBytes(), "qual1".getBytes(), "val1".getBytes());
+ Cell cell2 = new KeyValue(rowKey, "cf".getBytes(), "qual2".getBytes(), "val2".getBytes());
+
+ List cells = new ArrayList<>();
+ cells.add(cell1);
+ cells.add(cell2);
+
+ result = Result.create(cells);
+
+ CreateBigtableMutationsFn fn =
+ new CreateBigtableMutationsFn(
+ 1, // maxMutationsPerRequestThreshold = 1
+ false, 0L, false, 0, false, 0, true); // filterWideRows enabled
+
+ fn.processElement(KV.of(snapshotKey, result), receiver);
+
+ verify(receiver, times(0)).output(Mockito.any());
+ }
+
+ /**
+ * Tests that {@link CreateBigtableMutationsFn#processElement} does NOT drop wide rows if
+ * filterWideRows is enabled but cell count is exactly at threshold.
+ */
+ @Test
+ public void testProcessElement_filterWideRows_exactThreshold() throws Exception {
+ byte[] rowKey = "row-key".getBytes();
+ Cell cell1 = new KeyValue(rowKey, "cf".getBytes(), "qual1".getBytes(), "val1".getBytes());
+ Cell cell2 = new KeyValue(rowKey, "cf".getBytes(), "qual2".getBytes(), "val2".getBytes());
+
+ List cells = new ArrayList<>();
+ cells.add(cell1);
+ cells.add(cell2);
+
+ result = Result.create(cells);
+
+ CreateBigtableMutationsFn fn =
+ new CreateBigtableMutationsFn(
+ 2, // maxMutationsPerRequestThreshold = 2
+ false, 0L, false, 0, false, 0, true); // filterWideRows enabled
+
+ fn.processElement(KV.of(snapshotKey, result), receiver);
+
+ verify(receiver, times(1)).output(Mockito.any());
+ }
+
+ /**
+ * Tests that {@link CreateBigtableMutationsFn#processElement} does NOT drop the row if a large
+ * cell is filtered out first, and the remaining cells are within the row size threshold.
+ */
+ @Test
+ public void testProcessElement_largeCellFiltered_rowNotDropped() throws Exception {
+ byte[] rowKey = "row-key".getBytes();
+ Cell cell1 = new KeyValue(rowKey, "cf".getBytes(), "qual1".getBytes(), "val1".getBytes());
+ byte[] largeValue = new byte[1000];
+ Cell cell2 = new KeyValue(rowKey, "cf".getBytes(), "qual2".getBytes(), largeValue);
+
+ List cells = new ArrayList<>();
+ cells.add(cell1);
+ cells.add(cell2);
+
+ result = Result.create(cells);
+
+ CreateBigtableMutationsFn fn =
+ new CreateBigtableMutationsFn(
+ 100, true,
+ 200L, // filterLargeRows enabled, threshold 200 bytes (cell1 fits, cell1+cell2 doesn't)
+ true, 500, // filterLargeCells enabled, threshold 500 bytes (cell2 is filtered)
+ false, 0, false);
+
+ fn.processElement(KV.of(snapshotKey, result), receiver);
+
+ ArgumentCaptor>> captor = ArgumentCaptor.forClass(KV.class);
+ verify(receiver, times(1)).output(captor.capture());
+
+ KV> output = captor.getValue();
+ Iterator iterator = output.getValue().iterator();
+ Mutation mutation = iterator.next();
+ List outCells = mutation.getFamilyCellMap().values().iterator().next();
+ assertEquals(1, outCells.size());
+ assertEquals(
+ "qual1",
+ new String(
+ outCells.get(0).getQualifierArray(),
+ outCells.get(0).getQualifierOffset(),
+ outCells.get(0).getQualifierLength()));
+ }
+
+ /**
+ * Tests that {@link CreateBigtableMutationsFn#processElement} drops the row if total size exceeds
+ * threshold, even if it would have been split due to cell count.
+ */
+ @Test
+ public void testProcessElement_splitMutations_butRowTooLarge() throws Exception {
+ byte[] rowKey = "row-key".getBytes();
+ Cell cell1 = new KeyValue(rowKey, "cf".getBytes(), "qual1".getBytes(), "val1".getBytes());
+ Cell cell2 = new KeyValue(rowKey, "cf".getBytes(), "qual2".getBytes(), "val2".getBytes());
+
+ List cells = new ArrayList<>();
+ cells.add(cell1);
+ cells.add(cell2);
+
+ result = Result.create(cells);
+
+ long cell1Size = CellUtil.estimatedSerializedSizeOf(cell1);
+
+ CreateBigtableMutationsFn fn =
+ new CreateBigtableMutationsFn(
+ 1, // maxMutationsPerRequestThreshold = 1 (forces split)
+ true,
+ cell1Size + 10, // filterLargeRows enabled, fails on cell2
+ false,
+ 0,
+ false,
+ 0,
+ false);
+
+ fn.processElement(KV.of(snapshotKey, result), receiver);
+
+ verify(receiver, times(0)).output(Mockito.any()); // Should be dropped!
+ }
+
+ /**
+ * Tests that {@link CreateBigtableMutationsFn#processElement} filters out multiple large cells in
+ * the same row, and only logs once.
+ */
+ @Test
+ public void testProcessElement_filterMultipleLargeCells() throws Exception {
+ byte[] rowKey = "row-key".getBytes();
+ byte[] largeValue = new byte[1000];
+ Cell cell1 = new KeyValue(rowKey, "cf".getBytes(), "qual1".getBytes(), largeValue);
+ Cell cell2 = new KeyValue(rowKey, "cf".getBytes(), "qual2".getBytes(), largeValue);
+ Cell cell3 =
+ new KeyValue(rowKey, "cf".getBytes(), "qual3".getBytes(), "val".getBytes()); // Normal cell
+
+ List cells = new ArrayList<>();
+ cells.add(cell1);
+ cells.add(cell2);
+ cells.add(cell3);
+
+ result = Result.create(cells);
+
+ CreateBigtableMutationsFn fn =
+ new CreateBigtableMutationsFn(
+ 100, false, 0L, true, 500, // filterLargeCells enabled, threshold 500
+ false, 0, false);
+
+ fn.processElement(KV.of(snapshotKey, result), receiver);
+
+ ArgumentCaptor>> captor = ArgumentCaptor.forClass(KV.class);
+ verify(receiver, times(1)).output(captor.capture());
+
+ KV> output = captor.getValue();
+ Iterator iterator = output.getValue().iterator();
+ Mutation mutation = iterator.next();
+ List outCells = mutation.getFamilyCellMap().values().iterator().next();
+ assertEquals(1, outCells.size()); // Only cell3 should remain!
+ assertEquals(
+ "qual3",
+ new String(
+ outCells.get(0).getQualifierArray(),
+ outCells.get(0).getQualifierOffset(),
+ outCells.get(0).getQualifierLength()));
+ }
+
+ /**
+ * Tests that {@link CreateBigtableMutations#processElement} increments droppedRows counter when
+ * every cell in the row is filtered out due to exceeding the large cell threshold.
+ */
+ @Test
+ public void testProcessElement_filterLargeCell_incrementsDroppedRows() throws Exception {
+ MetricsContainer mockContainer = mock(MetricsContainer.class);
+ org.apache.beam.sdk.metrics.Counter mockCounter =
+ mock(org.apache.beam.sdk.metrics.Counter.class);
+ Mockito.when(mockContainer.getCounter(Mockito.any(MetricName.class))).thenReturn(mockCounter);
+ MetricsEnvironment.setCurrentContainer(mockContainer);
+
+ try {
+ byte[] rowKey = "row-key".getBytes();
+ byte[] largeValue = new byte[1000];
+ Cell cell = new KeyValue(rowKey, "cf".getBytes(), "qual".getBytes(), largeValue);
+
+ result = Result.create(Collections.singletonList(cell));
+
+ CreateBigtableMutationsFn fn =
+ new CreateBigtableMutationsFn(
+ 100, false, 0L, true, 500, // filterLargeCells enabled, threshold 500
+ false, 0, false);
+
+ fn.processElement(KV.of(snapshotKey, result), receiver);
+
+ // Verify that both droppedCells and droppedRows were requested and incremented
+ ArgumentCaptor captor = ArgumentCaptor.forClass(MetricName.class);
+ verify(mockContainer, times(2)).getCounter(captor.capture());
+
+ List requestedMetrics = captor.getAllValues();
+ assertEquals(2, requestedMetrics.size());
+
+ assertEquals("droppedCells", requestedMetrics.get(0).getName());
+ assertEquals("droppedRows", requestedMetrics.get(1).getName());
+
+ verify(mockCounter, times(2)).inc(1L);
+ } finally {
+ MetricsEnvironment.setCurrentContainer(null);
+ }
+ }
+}
diff --git a/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/dofn/HBaseRegionScannerTest.java b/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/dofn/HBaseRegionScannerTest.java
new file mode 100644
index 0000000000..d6524d7f0d
--- /dev/null
+++ b/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/dofn/HBaseRegionScannerTest.java
@@ -0,0 +1,195 @@
+/*
+ * Copyright 2026 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.dofn;
+
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.mockStatic;
+
+import java.io.IOException;
+import java.util.List;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.FileSystem;
+import org.apache.hadoop.fs.Path;
+import org.apache.hadoop.hbase.Cell;
+import org.apache.hadoop.hbase.TableName;
+import org.apache.hadoop.hbase.client.RegionInfo;
+import org.apache.hadoop.hbase.client.Result;
+import org.apache.hadoop.hbase.client.Scan;
+import org.apache.hadoop.hbase.client.TableDescriptor;
+import org.apache.hadoop.hbase.client.TableDescriptorBuilder;
+import org.apache.hadoop.hbase.regionserver.HRegion;
+import org.apache.hadoop.hbase.regionserver.RegionScanner;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.MockedStatic;
+import org.mockito.Mockito;
+
+/** Tests the {@link HBaseRegionScanner} functionality. */
+@RunWith(JUnit4.class)
+public class HBaseRegionScannerTest {
+
+ private Configuration conf;
+ private FileSystem fs;
+ private Path rootDir;
+ private TableDescriptor htd;
+ private RegionInfo hri;
+ private Scan scan;
+ private HRegion region;
+ private RegionScanner scanner;
+ private MockedStatic mockedHRegion;
+
+ @Before
+ public void setUp() throws Exception {
+ conf = new Configuration();
+ fs = mock(FileSystem.class);
+ rootDir = new Path("gs://bucket/data");
+ htd = TableDescriptorBuilder.newBuilder(TableName.valueOf("my-table")).build();
+ hri = mock(RegionInfo.class);
+ scan = new Scan();
+
+ region = mock(HRegion.class);
+ // HRegion.getScanner returns RegionScannerImpl which is package-private.
+ // We try to mock it by name, falling back to the interface if not found.
+ try {
+ scanner =
+ (RegionScanner)
+ mock(Class.forName("org.apache.hadoop.hbase.regionserver.RegionScannerImpl"));
+ } catch (ClassNotFoundException e) {
+ scanner = mock(RegionScanner.class);
+ }
+
+ Mockito.doReturn(scanner).when(region).getScanner(scan);
+
+ // Mock the static newHRegion method to avoid real file system access and
+ // region initialization which requires a lot of setup.
+ mockedHRegion = mockStatic(HRegion.class);
+ mockedHRegion
+ .when(
+ () ->
+ HRegion.newHRegion(
+ Mockito.any(Path.class),
+ Mockito.any(),
+ Mockito.any(FileSystem.class),
+ Mockito.any(Configuration.class),
+ Mockito.any(RegionInfo.class),
+ Mockito.any(TableDescriptor.class),
+ Mockito.any()))
+ .thenReturn(region);
+ }
+
+ @After
+ public void tearDown() {
+ mockedHRegion.close();
+ }
+
+ /**
+ * Tests that {@link HBaseRegionScanner#next()} successfully reads and returns records from the
+ * region.
+ */
+ @Test
+ public void testNext() throws Exception {
+ Cell cell = mock(Cell.class);
+
+ // Mock scanner.nextRaw(values) to populate the list and return false (no more rows)
+ Mockito.doAnswer(
+ invocation -> {
+ List list = invocation.getArgument(0);
+ list.add(cell);
+ return false;
+ })
+ .when(scanner)
+ .nextRaw(Mockito.anyList());
+
+ HBaseRegionScanner hbaseScanner = new HBaseRegionScanner(conf, fs, rootDir, htd, hri, scan);
+
+ Result result = hbaseScanner.next();
+ assertNotNull(result);
+ assertFalse(result.isEmpty());
+
+ Result resultNull = hbaseScanner.next();
+ assertNull(resultNull);
+ }
+
+ /**
+ * Tests that {@link HBaseRegionScanner#next()} skips empty results and continues scanning if the
+ * underlying scanner has more rows.
+ */
+ @Test
+ public void testNext_SkipsEmptyResults() throws Exception {
+ Cell cell = mock(Cell.class);
+
+ Mockito.doAnswer(
+ new org.mockito.stubbing.Answer() {
+ private int count = 0;
+
+ @Override
+ public Boolean answer(org.mockito.invocation.InvocationOnMock invocation) {
+ List| list = invocation.getArgument(0);
+ if (count == 0) {
+ count++;
+ return true; // hasMore = true, but empty list
+ } else {
+ list.add(cell);
+ return false; // hasMore = false, with data
+ }
+ }
+ })
+ .when(scanner)
+ .nextRaw(Mockito.anyList());
+
+ HBaseRegionScanner hbaseScanner = new HBaseRegionScanner(conf, fs, rootDir, htd, hri, scan);
+
+ Result result = hbaseScanner.next();
+ assertNotNull(result);
+ assertFalse(result.isEmpty());
+ }
+
+ /**
+ * Tests that {@link HBaseRegionScanner#close()} releases all resources by closing the scanner and
+ * region.
+ */
+ @Test
+ public void testClose() throws Exception {
+ HBaseRegionScanner hbaseScanner = new HBaseRegionScanner(conf, fs, rootDir, htd, hri, scan);
+
+ hbaseScanner.close();
+
+ Mockito.verify(scanner, Mockito.times(1)).close();
+ Mockito.verify(region, Mockito.times(1)).closeRegionOperation();
+ Mockito.verify(region, Mockito.times(1)).close(true);
+ }
+
+ @Test
+ public void testConstructor_InitializationFailure() throws Exception {
+ Mockito.doThrow(new IOException("Initialization failed")).when(region).initialize();
+
+ try {
+ new HBaseRegionScanner(conf, fs, rootDir, htd, hri, scan);
+ org.junit.Assert.fail("Expected IOException");
+ } catch (IOException e) {
+ org.junit.Assert.assertEquals("Initialization failed", e.getMessage());
+ }
+
+ Mockito.verify(region, Mockito.times(1)).close(true);
+ Mockito.verify(region, Mockito.never()).closeRegionOperation();
+ }
+}
diff --git a/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/dofn/ReadSnapshotRegionFnTest.java b/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/dofn/ReadSnapshotRegionFnTest.java
new file mode 100644
index 0000000000..a7595a8213
--- /dev/null
+++ b/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/dofn/ReadSnapshotRegionFnTest.java
@@ -0,0 +1,345 @@
+/*
+ * Copyright 2026 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.dofn;
+
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import com.google.cloud.bigtable.beam.hbasesnapshots.conf.RegionConfig;
+import com.google.cloud.bigtable.beam.hbasesnapshots.conf.SnapshotConfig;
+import com.google.cloud.bigtable.beam.hbasesnapshots.conf.SnapshotKey;
+import java.util.Collections;
+import org.apache.beam.sdk.io.range.ByteKey;
+import org.apache.beam.sdk.io.range.ByteKeyRange;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.values.KV;
+import org.apache.hadoop.hbase.client.Result;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Tests the {@link ReadSnapshotRegionFn} functionality. */
+@RunWith(JUnit4.class)
+public class ReadSnapshotRegionFnTest {
+
+ private RegionConfig regionConfig;
+ private SnapshotConfig snapshotConfig;
+ private org.apache.hadoop.hbase.client.RegionInfo regionInfo;
+
+ @Before
+ public void setUp() {
+ snapshotConfig =
+ SnapshotConfig.builder()
+ .setProjectId("test-project")
+ .setSourceLocation("gs://test-bucket/source")
+ .setSnapshotName("test-snapshot")
+ .setTableName("test-table")
+ .setRestoreLocation("gs://test-bucket/restore")
+ .setConfigurationDetails(Collections.emptyMap())
+ .build();
+
+ regionInfo =
+ org.apache.hadoop.hbase.client.RegionInfoBuilder.newBuilder(
+ org.apache.hadoop.hbase.TableName.valueOf("test-table"))
+ .setRegionId(12345L)
+ .build();
+
+ regionConfig =
+ RegionConfig.builder()
+ .setSnapshotConfig(snapshotConfig)
+ .setRegionInfo(regionInfo)
+ .setTableDescriptor(
+ org.apache.hadoop.hbase.client.TableDescriptorBuilder.newBuilder(
+ org.apache.hadoop.hbase.TableName.valueOf("test-table"))
+ .build())
+ .setRegionSize(100L)
+ .build();
+ }
+
+ /**
+ * Tests that {@link ReadSnapshotRegionFn#processElement} successfully reads and outputs all
+ * records from the region when no splitting occurs.
+ */
+ @Test
+ public void testProcessElement() throws Exception {
+ ByteKeyRange range =
+ ByteKeyRange.of(ByteKey.copyFrom("a".getBytes()), ByteKey.copyFrom("z".getBytes()));
+ RestrictionTracker tracker =
+ new com.google.cloud.bigtable.beam.hbasesnapshots.transforms.HbaseRegionSplitTracker(
+ "test-snapshot", "test-region", range, true);
+
+ HBaseRegionScanner scanner = mock(HBaseRegionScanner.class);
+ Result result1 = mock(Result.class);
+ Result result2 = mock(Result.class);
+ when(result1.getRow()).thenReturn("b".getBytes());
+ when(result2.getRow()).thenReturn("c".getBytes());
+
+ when(scanner.next()).thenReturn(result1, result2, null);
+
+ ReadSnapshotRegionFn fn = new TestReadSnapshotRegionFn(true, scanner);
+ fn.setup();
+
+ DoFn.OutputReceiver> receiver = mock(DoFn.OutputReceiver.class);
+
+ fn.processElement(regionConfig, receiver, tracker);
+
+ SnapshotKey key = SnapshotKey.create("test-snapshot", "test-table");
+ verify(receiver, times(1)).output(KV.of(key, result1));
+ verify(receiver, times(1)).output(KV.of(key, result2));
+ }
+
+ /**
+ * Tests that {@link ReadSnapshotRegionFn#processElement} stops processing when the restriction
+ * tracker refuses a claim (e.g. when the restriction has been split).
+ */
+ @Test
+ public void testProcessElement_StopOnClaimFailure() throws Exception {
+ ByteKeyRange range =
+ ByteKeyRange.of(ByteKey.copyFrom("a".getBytes()), ByteKey.copyFrom("z".getBytes()));
+ RestrictionTracker tracker =
+ new com.google.cloud.bigtable.beam.hbasesnapshots.transforms.HbaseRegionSplitTracker(
+ "test-snapshot", "test-region", range, true);
+
+ HBaseRegionScanner scanner = mock(HBaseRegionScanner.class);
+ Result result1 = mock(Result.class);
+ Result result2 = mock(Result.class);
+ when(result1.getRow()).thenReturn("b".getBytes());
+ when(result2.getRow()).thenReturn("c".getBytes());
+
+ when(scanner.next())
+ .thenAnswer(
+ new org.mockito.stubbing.Answer() {
+ private int count = 0;
+
+ @Override
+ public Result answer(org.mockito.invocation.InvocationOnMock invocation) {
+ count++;
+ if (count == 1) return result1;
+ if (count == 2) {
+ tracker.trySplit(0.0);
+ return result2;
+ }
+ return null;
+ }
+ });
+
+ ReadSnapshotRegionFn fn = new TestReadSnapshotRegionFn(true, scanner);
+ fn.setup();
+
+ DoFn.OutputReceiver> receiver = mock(DoFn.OutputReceiver.class);
+
+ fn.processElement(regionConfig, receiver, tracker);
+
+ SnapshotKey key = SnapshotKey.create("test-snapshot", "test-table");
+ verify(receiver, times(1)).output(KV.of(key, result1));
+ verify(receiver, times(0)).output(KV.of(key, result2));
+ }
+
+ /**
+ * Tests that {@link ReadSnapshotRegionFn#splitRestriction} does not split the restriction when
+ * the region size is smaller than the split threshold.
+ */
+ @Test
+ public void testSplitRestriction_NoSplit() throws Exception {
+ RegionConfig testRegionConfig =
+ RegionConfig.builder()
+ .setSnapshotConfig(snapshotConfig)
+ .setRegionInfo(regionInfo)
+ .setTableDescriptor(regionConfig.getTableDescriptor())
+ .setRegionSize(100L * 1024 * 1024) // 100 MB
+ .build();
+
+ ByteKeyRange range =
+ ByteKeyRange.of(ByteKey.copyFrom("a".getBytes()), ByteKey.copyFrom("z".getBytes()));
+
+ DoFn.OutputReceiver receiver = mock(DoFn.OutputReceiver.class);
+
+ ReadSnapshotRegionFn fn = new ReadSnapshotRegionFn(true);
+ fn.splitRestriction(testRegionConfig, range, receiver);
+
+ verify(receiver, times(1)).output(range);
+ }
+
+ /**
+ * Tests that {@link ReadSnapshotRegionFn#splitRestriction} splits the restriction into multiple
+ * sub-ranges when the region size exceeds the split threshold.
+ */
+ @Test
+ public void testSplitRestriction_WithSplit() throws Exception {
+ org.apache.hadoop.hbase.client.RegionInfo hriWithEndKey =
+ org.apache.hadoop.hbase.client.RegionInfoBuilder.newBuilder(
+ org.apache.hadoop.hbase.TableName.valueOf("test-table"))
+ .setRegionId(12345L)
+ .setStartKey("a".getBytes())
+ .setEndKey("z".getBytes())
+ .build();
+
+ RegionConfig testRegionConfig =
+ RegionConfig.builder()
+ .setSnapshotConfig(snapshotConfig)
+ .setRegionInfo(hriWithEndKey)
+ .setTableDescriptor(regionConfig.getTableDescriptor())
+ .setRegionSize(1500L * 1024 * 1024) // ~1.5 GB -> 3 splits
+ .build();
+
+ ByteKeyRange range =
+ ByteKeyRange.of(ByteKey.copyFrom("a".getBytes()), ByteKey.copyFrom("z".getBytes()));
+
+ DoFn.OutputReceiver receiver = mock(DoFn.OutputReceiver.class);
+
+ ReadSnapshotRegionFn fn = new ReadSnapshotRegionFn(true);
+ fn.splitRestriction(testRegionConfig, range, receiver);
+
+ verify(receiver, times(3)).output(org.mockito.ArgumentMatchers.any(ByteKeyRange.class));
+ }
+
+ /**
+ * Tests that {@link ReadSnapshotRegionFn#splitRestriction} falls back to the original restriction
+ * when an exception occurs during splitting.
+ */
+ @Test
+ public void testSplitRestriction_Exception() throws Exception {
+ RegionConfig testRegionConfig = new ThrowingRegionConfig(regionConfig);
+
+ ByteKeyRange range =
+ ByteKeyRange.of(ByteKey.copyFrom("a".getBytes()), ByteKey.copyFrom("z".getBytes()));
+
+ DoFn.OutputReceiver receiver = mock(DoFn.OutputReceiver.class);
+
+ ReadSnapshotRegionFn fn = new ReadSnapshotRegionFn(true);
+ fn.splitRestriction(testRegionConfig, range, receiver);
+
+ verify(receiver, times(1)).output(range);
+ }
+
+ static class ThrowingRegionConfig extends RegionConfig {
+ private final RegionConfig delegate;
+
+ public ThrowingRegionConfig(RegionConfig delegate) {
+ this.delegate = delegate;
+ }
+
+ @Override
+ public String getName() {
+ return delegate.getName();
+ }
+
+ @Override
+ public SnapshotConfig getSnapshotConfig() {
+ return delegate.getSnapshotConfig();
+ }
+
+ @Override
+ public org.apache.hadoop.hbase.client.RegionInfo getRegionInfo() {
+ return delegate.getRegionInfo();
+ }
+
+ @Override
+ public org.apache.hadoop.hbase.client.TableDescriptor getTableDescriptor() {
+ return delegate.getTableDescriptor();
+ }
+
+ @Override
+ public Long getRegionSize() {
+ throw new RuntimeException("test exception");
+ }
+ }
+
+ static class TestReadSnapshotRegionFn extends ReadSnapshotRegionFn {
+ private final HBaseRegionScanner mockScanner;
+
+ public TestReadSnapshotRegionFn(boolean useDynamicSplitting, HBaseRegionScanner mockScanner) {
+ super(useDynamicSplitting);
+ this.mockScanner = mockScanner;
+ }
+
+ @Override
+ HBaseRegionScanner newScanner(
+ RegionConfig regionConfig,
+ ByteKeyRange byteKeyRange,
+ org.apache.hadoop.conf.Configuration configuration) {
+ return mockScanner;
+ }
+ }
+
+ /**
+ * Tests that {@link ReadSnapshotRegionFn#splitRestriction} handles empty start keys (infinity) by
+ * padding them to match the length of the end key.
+ */
+ @Test
+ public void testSplitRestriction_EmptyStartKey() throws Exception {
+ org.apache.hadoop.hbase.client.RegionInfo hriEmptyStart =
+ org.apache.hadoop.hbase.client.RegionInfoBuilder.newBuilder(
+ org.apache.hadoop.hbase.TableName.valueOf("test-table"))
+ .setRegionId(12345L)
+ .setStartKey(new byte[0])
+ .setEndKey("z".getBytes())
+ .build();
+
+ RegionConfig testRegionConfig =
+ RegionConfig.builder()
+ .setSnapshotConfig(snapshotConfig)
+ .setRegionInfo(hriEmptyStart)
+ .setTableDescriptor(regionConfig.getTableDescriptor())
+ .setRegionSize(1500L * 1024 * 1024) // ~1.5 GB -> 3 splits
+ .build();
+
+ ByteKeyRange range = ByteKeyRange.of(ByteKey.EMPTY, ByteKey.copyFrom("z".getBytes()));
+
+ DoFn.OutputReceiver receiver = mock(DoFn.OutputReceiver.class);
+
+ ReadSnapshotRegionFn fn = new ReadSnapshotRegionFn(true);
+ fn.splitRestriction(testRegionConfig, range, receiver);
+
+ verify(receiver, times(3)).output(org.mockito.ArgumentMatchers.any(ByteKeyRange.class));
+ }
+
+ /**
+ * Tests that {@link ReadSnapshotRegionFn#splitRestriction} skips splitting for boundary regions
+ * with empty end keys.
+ */
+ @Test
+ public void testSplitRestriction_EmptyEndKey() throws Exception {
+ org.apache.hadoop.hbase.client.RegionInfo hriEmptyEnd =
+ org.apache.hadoop.hbase.client.RegionInfoBuilder.newBuilder(
+ org.apache.hadoop.hbase.TableName.valueOf("test-table"))
+ .setRegionId(12345L)
+ .setStartKey("a".getBytes())
+ .setEndKey(new byte[0])
+ .build();
+
+ RegionConfig testRegionConfig =
+ RegionConfig.builder()
+ .setSnapshotConfig(snapshotConfig)
+ .setRegionInfo(hriEmptyEnd)
+ .setTableDescriptor(regionConfig.getTableDescriptor())
+ .setRegionSize(1500L * 1024 * 1024) // ~1.5 GB -> 3 splits
+ .build();
+
+ ByteKeyRange range = ByteKeyRange.of(ByteKey.copyFrom("a".getBytes()), ByteKey.EMPTY);
+
+ DoFn.OutputReceiver receiver = mock(DoFn.OutputReceiver.class);
+
+ ReadSnapshotRegionFn fn = new ReadSnapshotRegionFn(true);
+ fn.splitRestriction(testRegionConfig, range, receiver);
+
+ verify(receiver, times(1)).output(org.mockito.ArgumentMatchers.any(ByteKeyRange.class));
+ }
+}
diff --git a/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/dofn/RestoreSnapshotFnTest.java b/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/dofn/RestoreSnapshotFnTest.java
new file mode 100644
index 0000000000..b2b0085ed1
--- /dev/null
+++ b/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/dofn/RestoreSnapshotFnTest.java
@@ -0,0 +1,58 @@
+/*
+ * Copyright 2026 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.dofn;
+
+import com.google.cloud.bigtable.beam.hbasesnapshots.SnapshotTestHelper;
+import com.google.cloud.bigtable.beam.hbasesnapshots.conf.SnapshotConfig;
+import org.apache.hadoop.hbase.snapshot.RestoreSnapshotHelper;
+import org.junit.ClassRule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.MockedStatic;
+import org.mockito.Mockito;
+
+/** Tests the {@link RestoreSnapshotFn} functionality. */
+@RunWith(JUnit4.class)
+public class RestoreSnapshotFnTest {
+
+ @ClassRule public static TemporaryFolder tempFolder = new TemporaryFolder();
+
+ @Test
+ public void testRestoreSnapshot() throws Exception {
+
+ SnapshotConfig snapshotConfig =
+ SnapshotTestHelper.newSnapshotConfig(
+ tempFolder.newFolder().getAbsolutePath(), tempFolder.newFolder().getAbsolutePath());
+
+ try (MockedStatic restoreSnapshotHelper =
+ Mockito.mockStatic(RestoreSnapshotHelper.class)) {
+ restoreSnapshotHelper
+ .when(
+ () ->
+ RestoreSnapshotHelper.copySnapshotForScanner(
+ snapshotConfig.getConfiguration(),
+ null,
+ snapshotConfig.getSourcePath(),
+ snapshotConfig.getRestorePath(),
+ snapshotConfig.getSnapshotName()))
+ .thenReturn(null);
+
+ new RestoreSnapshotFn().restoreSnapshot(snapshotConfig);
+ }
+ }
+}
diff --git a/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/transforms/HbaseRegionSplitTrackerTest.java b/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/transforms/HbaseRegionSplitTrackerTest.java
new file mode 100644
index 0000000000..db3c328507
--- /dev/null
+++ b/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/transforms/HbaseRegionSplitTrackerTest.java
@@ -0,0 +1,129 @@
+/*
+ * Copyright 2026 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.transforms;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertNull;
+import static org.junit.Assert.assertTrue;
+import static org.junit.Assert.fail;
+
+import org.apache.beam.sdk.io.range.ByteKey;
+import org.apache.beam.sdk.io.range.ByteKeyRange;
+import org.apache.beam.sdk.transforms.splittabledofn.RestrictionTracker;
+import org.apache.beam.sdk.transforms.splittabledofn.SplitResult;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/**
+ * Tests the {@link HbaseRegionSplitTracker} to ensure correct delegation and dynamic splitting
+ * behavior.
+ */
+@RunWith(JUnit4.class)
+public class HbaseRegionSplitTrackerTest {
+
+ private ByteKeyRange range;
+ private HbaseRegionSplitTracker trackerWithDynamicSplitting;
+ private HbaseRegionSplitTracker trackerWithoutDynamicSplitting;
+
+ @Before
+ public void setUp() {
+ range = ByteKeyRange.of(ByteKey.copyFrom("a".getBytes()), ByteKey.copyFrom("z".getBytes()));
+ trackerWithDynamicSplitting = new HbaseRegionSplitTracker("snapshot", "region", range, true);
+ trackerWithoutDynamicSplitting =
+ new HbaseRegionSplitTracker("snapshot", "region", range, false);
+ }
+
+ /** Tests that {@link HbaseRegionSplitTracker#currentRestriction()} returns the initial range. */
+ @Test
+ public void testCurrentRestriction() {
+ assertEquals(range, trackerWithDynamicSplitting.currentRestriction());
+ assertEquals(range, trackerWithoutDynamicSplitting.currentRestriction());
+ }
+
+ /** Tests that {@link HbaseRegionSplitTracker#tryClaim(ByteKey)} succeeds for keys in range. */
+ @Test
+ public void testTryClaim() {
+ assertTrue(trackerWithDynamicSplitting.tryClaim(ByteKey.copyFrom("b".getBytes())));
+ assertTrue(trackerWithoutDynamicSplitting.tryClaim(ByteKey.copyFrom("b".getBytes())));
+
+ // Claim outside range should fail
+ assertFalse(trackerWithDynamicSplitting.tryClaim(ByteKey.copyFrom("~".getBytes())));
+ assertFalse(trackerWithoutDynamicSplitting.tryClaim(ByteKey.copyFrom("~".getBytes())));
+ }
+
+ /**
+ * Tests that {@link HbaseRegionSplitTracker#trySplit(double)} succeeds when dynamic splitting is
+ * enabled.
+ */
+ @Test
+ public void testTrySplit_enabled() {
+ // Claim something to make remainder smaller
+ trackerWithDynamicSplitting.tryClaim(ByteKey.copyFrom("m".getBytes()));
+
+ SplitResult split = trackerWithDynamicSplitting.trySplit(0.5);
+ assertNotNull(split);
+ assertNotNull(split.getPrimary());
+ assertNotNull(split.getResidual());
+ }
+
+ /**
+ * Tests that {@link HbaseRegionSplitTracker#trySplit(double)} returns null when dynamic splitting
+ * is disabled.
+ */
+ @Test
+ public void testTrySplit_disabled() {
+ trackerWithoutDynamicSplitting.tryClaim(ByteKey.copyFrom("m".getBytes()));
+
+ SplitResult split = trackerWithoutDynamicSplitting.trySplit(0.5);
+ assertNull(split);
+ }
+
+ /**
+ * Tests that {@link HbaseRegionSplitTracker#checkDone()} throws exception if not all keys
+ * claimed.
+ */
+ @Test
+ public void testCheckDone_failsIfNotDone() {
+ // Should fail if not done
+ try {
+ trackerWithDynamicSplitting.checkDone();
+ fail("Expected IllegalStateException");
+ } catch (IllegalStateException e) {
+ // Expected
+ }
+ }
+
+ /** Tests that {@link HbaseRegionSplitTracker#checkDone()} succeeds when all keys claimed. */
+ @Test
+ public void testCheckDone_success() {
+ // Claim the start key to initialize
+ trackerWithDynamicSplitting.tryClaim(ByteKey.copyFrom("a".getBytes()));
+ // Claim the end key to mark as done
+ trackerWithDynamicSplitting.tryClaim(ByteKey.copyFrom("z".getBytes()));
+
+ trackerWithDynamicSplitting.checkDone(); // Should not throw
+ }
+
+ /** Tests that {@link HbaseRegionSplitTracker#isBounded()} returns BOUNDED. */
+ @Test
+ public void testIsBounded() {
+ assertEquals(RestrictionTracker.IsBounded.BOUNDED, trackerWithDynamicSplitting.isBounded());
+ }
+}
diff --git a/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/transforms/ListRegionsTest.java b/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/transforms/ListRegionsTest.java
new file mode 100644
index 0000000000..d2fc8f687e
--- /dev/null
+++ b/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/transforms/ListRegionsTest.java
@@ -0,0 +1,202 @@
+/*
+ * Copyright 2026 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.transforms;
+
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.mockStatic;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import com.google.cloud.bigtable.beam.hbasesnapshots.conf.RegionConfig;
+import com.google.cloud.bigtable.beam.hbasesnapshots.conf.SnapshotConfig;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import org.apache.beam.sdk.transforms.DoFn;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.FileSystem;
+import org.apache.hadoop.fs.Path;
+import org.apache.hadoop.hbase.client.RegionInfo;
+import org.apache.hadoop.hbase.client.TableDescriptor;
+import org.apache.hadoop.hbase.mapreduce.TableSnapshotInputFormatImpl;
+import org.apache.hadoop.hbase.shaded.protobuf.generated.HBaseProtos;
+import org.apache.hadoop.hbase.shaded.protobuf.generated.SnapshotProtos;
+import org.apache.hadoop.hbase.snapshot.SnapshotManifest;
+import org.apache.hbase.thirdparty.com.google.protobuf.ByteString;
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+import org.mockito.MockedStatic;
+import org.mockito.Mockito;
+
+/**
+ * Tests the {@link ListRegions} transform to ensure it correctly lists regions from a snapshot
+ * manifest.
+ */
+@RunWith(JUnit4.class)
+public class ListRegionsTest {
+
+ private SnapshotConfig snapshotConfig;
+ private Configuration conf;
+ private Path sourcePath;
+ private FileSystem fs;
+ private SnapshotManifest manifest;
+ private TableDescriptor td;
+ private DoFn.OutputReceiver receiver;
+
+ @Before
+ public void setUp() {
+ snapshotConfig = mock(SnapshotConfig.class);
+ conf = new Configuration();
+ sourcePath = new Path("file:///tmp/data");
+ fs = mock(FileSystem.class);
+ manifest = mock(SnapshotManifest.class);
+ td = mock(TableDescriptor.class);
+ receiver = mock(DoFn.OutputReceiver.class);
+
+ when(snapshotConfig.getConfiguration()).thenReturn(conf);
+ when(snapshotConfig.getSourcePath()).thenReturn(sourcePath);
+ when(snapshotConfig.getSnapshotName()).thenReturn("test-snapshot");
+ when(manifest.getTableDescriptor()).thenReturn(td);
+ }
+
+ /** Tests that {@link ListRegions.ListRegionsFn} handles snapshots with no regions. */
+ @Test
+ public void testListRegionsFn_emptySnapshot() throws Exception {
+ when(manifest.getRegionManifests()).thenReturn(Collections.emptyList());
+ List regionInfos = Collections.emptyList();
+
+ try (MockedStatic mockedStatic =
+ mockStatic(TableSnapshotInputFormatImpl.class)) {
+ mockedStatic
+ .when(
+ () ->
+ TableSnapshotInputFormatImpl.getSnapshotManifest(
+ Mockito.any(), Mockito.anyString(), Mockito.any(), Mockito.any()))
+ .thenReturn(manifest);
+ mockedStatic
+ .when(() -> TableSnapshotInputFormatImpl.getRegionInfosFromManifest(manifest))
+ .thenReturn(regionInfos);
+
+ ListRegions.ListRegionsFn fn = new ListRegions.ListRegionsFn();
+ fn.setup();
+ fn.processElement(snapshotConfig, receiver);
+
+ verify(receiver, times(0)).output(Mockito.any());
+ }
+ }
+
+ /** Tests that {@link ListRegions.ListRegionsFn} correctly lists a single region. */
+ @Test
+ public void testListRegionsFn_regularSnapshot() throws Exception {
+ RegionInfo ri = mock(RegionInfo.class);
+ when(ri.getRegionId()).thenReturn(123L);
+ when(ri.getStartKey()).thenReturn(new byte[] {1});
+
+ SnapshotProtos.SnapshotRegionManifest regionManifest =
+ mock(SnapshotProtos.SnapshotRegionManifest.class);
+ HBaseProtos.RegionInfo protoRegionInfo = mock(HBaseProtos.RegionInfo.class);
+ when(protoRegionInfo.getRegionId()).thenReturn(123L);
+ when(protoRegionInfo.getStartKey()).thenReturn(ByteString.copyFrom(new byte[] {1}));
+ when(regionManifest.getRegionInfo()).thenReturn(protoRegionInfo);
+
+ SnapshotProtos.SnapshotRegionManifest.FamilyFiles familyFiles =
+ mock(SnapshotProtos.SnapshotRegionManifest.FamilyFiles.class);
+ SnapshotProtos.SnapshotRegionManifest.StoreFile storeFile =
+ mock(SnapshotProtos.SnapshotRegionManifest.StoreFile.class);
+ when(storeFile.getFileSize()).thenReturn(1000L);
+ when(familyFiles.getStoreFilesList()).thenReturn(Collections.singletonList(storeFile));
+ when(regionManifest.getFamilyFilesList()).thenReturn(Collections.singletonList(familyFiles));
+
+ when(manifest.getRegionManifests()).thenReturn(Collections.singletonList(regionManifest));
+ List regionInfos = Collections.singletonList(ri);
+
+ try (MockedStatic mockedStatic =
+ mockStatic(TableSnapshotInputFormatImpl.class)) {
+ mockedStatic
+ .when(
+ () ->
+ TableSnapshotInputFormatImpl.getSnapshotManifest(
+ Mockito.any(), Mockito.anyString(), Mockito.any(), Mockito.any()))
+ .thenReturn(manifest);
+ mockedStatic
+ .when(() -> TableSnapshotInputFormatImpl.getRegionInfosFromManifest(manifest))
+ .thenReturn(regionInfos);
+
+ ListRegions.ListRegionsFn fn = new ListRegions.ListRegionsFn();
+ fn.setup();
+ fn.processElement(snapshotConfig, receiver);
+
+ verify(receiver, times(1)).output(Mockito.any(RegionConfig.class));
+ }
+ }
+
+ /** Tests that {@link ListRegions.ListRegionsFn} correctly lists multiple regions. */
+ @Test
+ public void testListRegionsFn_multipleRegions() throws Exception {
+ RegionInfo ri1 = mock(RegionInfo.class);
+ when(ri1.getRegionId()).thenReturn(123L);
+ when(ri1.getStartKey()).thenReturn(new byte[] {1});
+ RegionInfo ri2 = mock(RegionInfo.class);
+ when(ri2.getRegionId()).thenReturn(456L);
+ when(ri2.getStartKey()).thenReturn(new byte[] {2});
+
+ SnapshotProtos.SnapshotRegionManifest regionManifest1 =
+ mock(SnapshotProtos.SnapshotRegionManifest.class);
+ HBaseProtos.RegionInfo protoRegionInfo1 = mock(HBaseProtos.RegionInfo.class);
+ when(protoRegionInfo1.getRegionId()).thenReturn(123L);
+ when(protoRegionInfo1.getStartKey()).thenReturn(ByteString.copyFrom(new byte[] {1}));
+ when(regionManifest1.getRegionInfo()).thenReturn(protoRegionInfo1);
+
+ SnapshotProtos.SnapshotRegionManifest regionManifest2 =
+ mock(SnapshotProtos.SnapshotRegionManifest.class);
+ HBaseProtos.RegionInfo protoRegionInfo2 = mock(HBaseProtos.RegionInfo.class);
+ when(protoRegionInfo2.getRegionId()).thenReturn(456L);
+ when(protoRegionInfo2.getStartKey()).thenReturn(ByteString.copyFrom(new byte[] {2}));
+ when(regionManifest2.getRegionInfo()).thenReturn(protoRegionInfo2);
+
+ List manifests = new ArrayList<>();
+ manifests.add(regionManifest1);
+ manifests.add(regionManifest2);
+
+ when(manifest.getRegionManifests()).thenReturn(manifests);
+
+ List regionInfos = new ArrayList<>();
+ regionInfos.add(ri1);
+ regionInfos.add(ri2);
+
+ try (MockedStatic mockedStatic =
+ mockStatic(TableSnapshotInputFormatImpl.class)) {
+ mockedStatic
+ .when(
+ () ->
+ TableSnapshotInputFormatImpl.getSnapshotManifest(
+ Mockito.any(), Mockito.anyString(), Mockito.any(), Mockito.any()))
+ .thenReturn(manifest);
+ mockedStatic
+ .when(() -> TableSnapshotInputFormatImpl.getRegionInfosFromManifest(manifest))
+ .thenReturn(regionInfos);
+
+ ListRegions.ListRegionsFn fn = new ListRegions.ListRegionsFn();
+ fn.setup();
+ fn.processElement(snapshotConfig, receiver);
+
+ verify(receiver, times(2)).output(Mockito.any(RegionConfig.class));
+ }
+ }
+}
diff --git a/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/transforms/ReadRegionsTest.java b/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/transforms/ReadRegionsTest.java
new file mode 100644
index 0000000000..e73d8f7b25
--- /dev/null
+++ b/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/hbasesnapshots/transforms/ReadRegionsTest.java
@@ -0,0 +1,160 @@
+/*
+ * Copyright 2026 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.transforms;
+
+import static org.junit.Assert.assertTrue;
+
+import com.google.cloud.bigtable.beam.hbasesnapshots.conf.RegionConfig;
+import com.google.cloud.bigtable.beam.hbasesnapshots.conf.SnapshotConfig;
+import java.math.BigInteger;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import org.apache.beam.sdk.testing.PAssert;
+import org.apache.beam.sdk.testing.TestPipeline;
+import org.apache.beam.sdk.transforms.Create;
+import org.apache.beam.sdk.transforms.Filter;
+import org.apache.hadoop.hbase.TableName;
+import org.apache.hadoop.hbase.client.ColumnFamilyDescriptorBuilder;
+import org.apache.hadoop.hbase.client.RegionInfo;
+import org.apache.hadoop.hbase.client.RegionInfoBuilder;
+import org.apache.hadoop.hbase.client.TableDescriptor;
+import org.apache.hadoop.hbase.client.TableDescriptorBuilder;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+/** Tests the {@link ReadRegions} transform. */
+@RunWith(JUnit4.class)
+public class ReadRegionsTest {
+
+ @Rule public final transient TestPipeline pipeline = TestPipeline.create();
+
+ /** Tests that constructor throws exception when numShards is set but shardIndex is missing. */
+ @Test(expected = IllegalArgumentException.class)
+ public void testInvalidShardConfig_missingShardIndex() {
+ new ReadRegions(
+ true,
+ 100,
+ false,
+ 0L,
+ false,
+ 0,
+ false,
+ 0,
+ false, // filterWideRows
+ 2, // numShards
+ null // shardIndex missing
+ );
+ }
+
+ /** Tests that constructor throws exception when shardIndex is out of bounds. */
+ @Test(expected = IllegalArgumentException.class)
+ public void testInvalidShardConfig_shardIndexTooHigh() {
+ new ReadRegions(true, 100, false, 0L, false, 0, false, 0, false, 2, 2);
+ }
+
+ /** Tests the basic sharding logic math. */
+ @Test
+ public void testShardingLogic() {
+ // Verify the logic used for sharding
+ byte[] regionName = "someEncodedRegionName".getBytes();
+ int numShards = 4;
+
+ long remainder = new BigInteger(regionName).mod(BigInteger.valueOf(numShards)).longValue();
+ assertTrue(remainder >= 0 && remainder < numShards);
+ }
+
+ /**
+ * Tests that the sharding logic correctly selects or skips regions based on their encoded name,
+ * numShards, and shardIndex using TestPipeline.
+ */
+ @Test
+ public void testShardingWithPipeline() {
+ // 1. Setup configuration and table descriptor
+ SnapshotConfig snapshotConfig =
+ SnapshotConfig.builder()
+ .setProjectId("project")
+ .setSourceLocation("source")
+ .setSnapshotName("snapshot")
+ .setTableName("table")
+ .setRestoreLocation("restore")
+ .setConfigurationDetails(Collections.emptyMap())
+ .build();
+
+ TableDescriptor td =
+ TableDescriptorBuilder.newBuilder(TableName.valueOf("table"))
+ .setColumnFamily(ColumnFamilyDescriptorBuilder.newBuilder("cf".getBytes()).build())
+ .build();
+
+ // 2. Create test regions
+ RegionInfo ri1 =
+ RegionInfoBuilder.newBuilder(TableName.valueOf("table")).setRegionId(1).build();
+
+ RegionConfig rc1 =
+ RegionConfig.builder()
+ .setSnapshotConfig(snapshotConfig)
+ .setRegionInfo(ri1)
+ .setTableDescriptor(td)
+ .setRegionSize(100L)
+ .setName("region1")
+ .build();
+
+ RegionInfo ri2 =
+ RegionInfoBuilder.newBuilder(TableName.valueOf("table")).setRegionId(2).build();
+
+ RegionConfig rc2 =
+ RegionConfig.builder()
+ .setSnapshotConfig(snapshotConfig)
+ .setRegionInfo(ri2)
+ .setTableDescriptor(td)
+ .setRegionSize(100L)
+ .setName("region2")
+ .build();
+
+ // 3. Calculate expected shards for the regions to determine expected output
+ int shard1 =
+ new BigInteger(1, ri1.getEncodedNameAsBytes()).mod(BigInteger.valueOf(4)).intValue();
+ int shard2 =
+ new BigInteger(1, ri2.getEncodedNameAsBytes()).mod(BigInteger.valueOf(4)).intValue();
+
+ int targetShard = shard1;
+
+ // 4. Apply the sharding filter in a pipeline
+ org.apache.beam.sdk.values.PCollection input =
+ pipeline.apply(
+ Create.of(rc1, rc2)
+ .withCoder(
+ new com.google.cloud.bigtable.beam.hbasesnapshots.coders.RegionConfigCoder()));
+
+ org.apache.beam.sdk.values.PCollection output =
+ input.apply(Filter.by(rc -> ReadRegions.isRegionSelected(rc, 4, targetShard)));
+
+ // 5. Verify that only regions belonging to the target shard are selected
+ List expected = new ArrayList<>();
+ expected.add(rc1);
+ if (shard2 == targetShard) {
+ expected.add(rc2);
+ }
+
+ PAssert.that(output).containsInAnyOrder(expected);
+
+ // PAssert is added as a transform to the pipeline graph. We must call pipeline.run()
+ // to actually execute the pipeline and evaluate the assertions.
+ pipeline.run();
+ }
+}
diff --git a/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/sequencefiles/SequenceFileSourceTest.java b/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/sequencefiles/SequenceFileSourceTest.java
index 5d1715ccab..87f6a4c085 100644
--- a/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/sequencefiles/SequenceFileSourceTest.java
+++ b/bigtable-dataflow-parent/bigtable-beam-import/src/test/java/com/google/cloud/bigtable/beam/sequencefiles/SequenceFileSourceTest.java
@@ -16,12 +16,8 @@
package com.google.cloud.bigtable.beam.sequencefiles;
import static org.hamcrest.MatcherAssert.assertThat;
-import static org.hamcrest.Matchers.containsInAnyOrder;
-import static org.hamcrest.Matchers.hasSize;
-import static org.hamcrest.Matchers.instanceOf;
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertNotEquals;
-import static org.junit.Assert.assertTrue;
+import static org.hamcrest.Matchers.*;
+import static org.junit.Assert.*;
import com.google.common.collect.Lists;
import java.io.File;
| | | | | | | | | | | | | | |