Set all tablets availability to ONDEMAND on import

DomGarguilo · DomGarguilo · commit 59785cd5122c · 2024-02-16T16:25:30.000-05:00
diff --git a/server/manager/src/main/java/org/apache/accumulo/manager/tableOps/tableImport/PopulateMetadataTable.java b/server/manager/src/main/java/org/apache/accumulo/manager/tableOps/tableImport/PopulateMetadataTable.java
@@ -56,6 +56,7 @@
 import org.apache.accumulo.manager.tableOps.ManagerRepo;
 import org.apache.accumulo.server.fs.VolumeManager;
 import org.apache.accumulo.server.util.MetadataTableUtil;
+import org.apache.hadoop.fs.FSDataInputStream;
 import org.apache.hadoop.fs.Path;
 import org.apache.hadoop.io.Text;
 import org.slf4j.Logger;
@@ -103,7 +104,8 @@ public Repo<Manager> call(FateId fateId, Manager manager) throws Exception {
     try (
         BatchWriter mbw =
             manager.getContext().createBatchWriter(AccumuloTable.METADATA.tableName());
-        ZipInputStream zis = new ZipInputStream(fs.open(path))) {
+        FSDataInputStream fsDataInputStream = fs.open(path);
+        ZipInputStream zis = new ZipInputStream(fsDataInputStream)) {
 
       Map<String,String> fileNameMappings = new HashMap<>();
       for (ImportedTableInfo.DirectoryMapping dm : tableInfo.directories) {
@@ -115,8 +117,11 @@ public Repo<Manager> call(FateId fateId, Manager manager) throws Exception {
 
       ZipEntry zipEntry;
       while ((zipEntry = zis.getNextEntry()) != null) {
-        if (zipEntry.getName().equals(Constants.EXPORT_METADATA_FILE)) {
-          DataInputStream in = new DataInputStream(new BufferedInputStream(zis));
+        if (!zipEntry.getName().equals(Constants.EXPORT_METADATA_FILE)) {
+          continue;
+        }
+        try (BufferedInputStream bufferedInputStream = new BufferedInputStream(zis);
+            DataInputStream in = new DataInputStream(bufferedInputStream)) {
 
           Key key = new Key();
           Value val = new Value();
@@ -125,8 +130,6 @@ public Repo<Manager> call(FateId fateId, Manager manager) throws Exception {
           Text currentRow = null;
           int dirCount = 0;
 
-          boolean sawTabletAvailability = false;
-
           while (true) {
             key.readFields(in);
             val.readFields(in);
@@ -166,11 +169,9 @@ public Repo<Manager> call(FateId fateId, Manager manager) throws Exception {
             if (m == null || !currentRow.equals(metadataRow)) {
 
               if (m != null) {
-                if (!sawTabletAvailability) {
-                  // add a default tablet availability
-                  TabletColumnFamily.AVAILABILITY_COLUMN.put(m,
-                      TabletAvailabilityUtil.toValue(TabletAvailability.ONDEMAND));
-                }
+                // add a default tablet availability
+                TabletColumnFamily.AVAILABILITY_COLUMN.put(m,
+                    TabletAvailabilityUtil.toValue(TabletAvailability.ONDEMAND));
                 mbw.addMutation(m);
               }
 
@@ -183,21 +184,15 @@ public Repo<Manager> call(FateId fateId, Manager manager) throws Exception {
               m = new Mutation(metadataRow);
               ServerColumnFamily.DIRECTORY_COLUMN.put(m, new Value(tabletDir));
               currentRow = metadataRow;
-              sawTabletAvailability = false;
             }
 
-            if (TabletColumnFamily.AVAILABILITY_COLUMN.hasColumns(key)) {
-              sawTabletAvailability = true;
-            }
             m.put(key.getColumnFamily(), cq, val);
 
             if (endRow == null && TabletColumnFamily.PREV_ROW_COLUMN.hasColumns(key)) {
 
-              if (!sawTabletAvailability) {
-                // add a default tablet availability
-                TabletColumnFamily.AVAILABILITY_COLUMN.put(m,
-                    TabletAvailabilityUtil.toValue(TabletAvailability.ONDEMAND));
-              }
+              // add a default tablet availability
+              TabletColumnFamily.AVAILABILITY_COLUMN.put(m,
+                  TabletAvailabilityUtil.toValue(TabletAvailability.ONDEMAND));
 
               mbw.addMutation(m);
               break; // it is the last column in the last row
diff --git a/test/src/main/java/org/apache/accumulo/test/ImportExportIT.java b/test/src/main/java/org/apache/accumulo/test/ImportExportIT.java
@@ -20,6 +20,8 @@
 
 import static org.apache.accumulo.core.Constants.IMPORT_MAPPINGS_FILE;
 import static org.apache.accumulo.core.util.LazySingletons.RANDOM;
+import static org.apache.accumulo.test.TableOperationsIT.setExpectedTabletAvailability;
+import static org.apache.accumulo.test.TableOperationsIT.verifyTabletAvailabilites;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
@@ -33,20 +35,24 @@
 import java.io.InputStreamReader;
 import java.nio.file.Paths;
 import java.time.Duration;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map.Entry;
 import java.util.Set;
+import java.util.SortedSet;
 
 import org.apache.accumulo.cluster.AccumuloCluster;
 import org.apache.accumulo.core.Constants;
 import org.apache.accumulo.core.client.Accumulo;
 import org.apache.accumulo.core.client.AccumuloClient;
 import org.apache.accumulo.core.client.BatchWriter;
 import org.apache.accumulo.core.client.Scanner;
+import org.apache.accumulo.core.client.admin.AvailabilityForTablet;
 import org.apache.accumulo.core.client.admin.CompactionConfig;
 import org.apache.accumulo.core.client.admin.ImportConfiguration;
+import org.apache.accumulo.core.client.admin.TabletAvailability;
 import org.apache.accumulo.core.data.Key;
 import org.apache.accumulo.core.data.Mutation;
 import org.apache.accumulo.core.data.Range;
@@ -75,6 +81,8 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import com.google.common.collect.Sets;
+
 /**
  * ImportTable didn't correctly place absolute paths in metadata. This resulted in the imported
  * table only being usable when the actual HDFS directory for Accumulo was the same as
@@ -367,6 +375,144 @@ public void testExportImportOffline(boolean fenced) throws Exception {
     }
   }
 
+  /**
+   * Ensure all tablets in an imported table are ONDEMAND.
+   *
+   * Create a table with multiple tablets, each with a different tablet availability. Export the
+   * table. Import the table and make sure that all tablets on the imported table have the ONDEMAND
+   * tablet availability.
+   *
+   * This test case stitches together code from TableOperationsIT to create the table, set up the
+   * tablets and verify. The code to export then import the table is from
+   * ImportExportIT.testExportImportOffline()
+   */
+  @Test
+  public void testImportedTableIsOnDemand() throws Exception {
+
+    try (AccumuloClient client = Accumulo.newClient().from(getClientProps()).build()) {
+      String[] tableNames = getUniqueNames(2);
+      String srcTable = tableNames[0], destTable = tableNames[1];
+
+      client.tableOperations().create(srcTable);
+      String srcTableId = client.tableOperations().tableIdMap().get(srcTable);
+
+      // add split 'h' and 'q'. Leave first as ONDEMAND, set second to UNHOSTED, and third to HOSTED
+      SortedSet<Text> splits = Sets.newTreeSet(Arrays.asList(new Text("h"), new Text("q")));
+      client.tableOperations().addSplits(srcTable, splits);
+      Range range = new Range(new Text("h"), false, new Text("q"), true);
+      client.tableOperations().setTabletAvailability(srcTable, range, TabletAvailability.UNHOSTED);
+      range = new Range(new Text("q"), false, null, true);
+      client.tableOperations().setTabletAvailability(srcTable, range, TabletAvailability.HOSTED);
+
+      // verify
+      List<AvailabilityForTablet> expectedTabletAvailability = new ArrayList<>();
+      setExpectedTabletAvailability(expectedTabletAvailability, srcTableId, "h", null,
+          TabletAvailability.ONDEMAND);
+      setExpectedTabletAvailability(expectedTabletAvailability, srcTableId, "q", "h",
+          TabletAvailability.UNHOSTED);
+      setExpectedTabletAvailability(expectedTabletAvailability, srcTableId, null, "q",
+          TabletAvailability.HOSTED);
+      verifyTabletAvailabilites(client, srcTable, new Range(), expectedTabletAvailability);
+
+      // Add a split within each of the existing tablets. Adding 'd', 'm', and 'v'
+      splits = Sets.newTreeSet(Arrays.asList(new Text("d"), new Text("m"), new Text("v")));
+      client.tableOperations().addSplits(srcTable, splits);
+
+      // verify results
+      expectedTabletAvailability.clear();
+      setExpectedTabletAvailability(expectedTabletAvailability, srcTableId, "d", null,
+          TabletAvailability.ONDEMAND);
+      setExpectedTabletAvailability(expectedTabletAvailability, srcTableId, "h", "d",
+          TabletAvailability.ONDEMAND);
+      setExpectedTabletAvailability(expectedTabletAvailability, srcTableId, "m", "h",
+          TabletAvailability.UNHOSTED);
+      setExpectedTabletAvailability(expectedTabletAvailability, srcTableId, "q", "m",
+          TabletAvailability.UNHOSTED);
+      setExpectedTabletAvailability(expectedTabletAvailability, srcTableId, "v", "q",
+          TabletAvailability.HOSTED);
+      setExpectedTabletAvailability(expectedTabletAvailability, srcTableId, null, "v",
+          TabletAvailability.HOSTED);
+      verifyTabletAvailabilites(client, srcTable, new Range(), expectedTabletAvailability);
+
+      // Make a directory we can use to throw the export and import directories
+      // Must exist on the filesystem the cluster is running.
+      FileSystem fs = cluster.getFileSystem();
+      log.info("Using FileSystem: " + fs);
+      Path baseDir = new Path(cluster.getTemporaryPath(), getClass().getName());
+      fs.deleteOnExit(baseDir);
+      if (fs.exists(baseDir)) {
+        log.info("{} exists on filesystem, deleting", baseDir);
+        assertTrue(fs.delete(baseDir, true), "Failed to deleted " + baseDir);
+      }
+      log.info("Creating {}", baseDir);
+      assertTrue(fs.mkdirs(baseDir), "Failed to create " + baseDir);
+      Path exportDir = new Path(baseDir, "export");
+      fs.deleteOnExit(exportDir);
+      Path importDirA = new Path(baseDir, "import-a");
+      Path importDirB = new Path(baseDir, "import-b");
+      fs.deleteOnExit(importDirA);
+      fs.deleteOnExit(importDirB);
+      for (Path p : new Path[] {exportDir, importDirA, importDirB}) {
+        assertTrue(fs.mkdirs(p), "Failed to create " + p);
+      }
+
+      Set<String> importDirs = Set.of(importDirA.toString(), importDirB.toString());
+
+      Path[] importDirAry = new Path[] {importDirA, importDirB};
+
+      log.info("Exporting table to {}", exportDir);
+      log.info("Importing table from {}", importDirs);
+
+      // test fast fail offline check
+      assertThrows(IllegalStateException.class,
+          () -> client.tableOperations().exportTable(srcTable, exportDir.toString()));
+
+      // Offline the table
+      client.tableOperations().offline(srcTable, true);
+      // Then export it
+      client.tableOperations().exportTable(srcTable, exportDir.toString());
+
+      // Make sure the distcp.txt file that exporttable creates is available
+      Path distcp = new Path(exportDir, "distcp.txt");
+      fs.deleteOnExit(distcp);
+      assertTrue(fs.exists(distcp), "Distcp file doesn't exist");
+      FSDataInputStream is = fs.open(distcp);
+      BufferedReader reader = new BufferedReader(new InputStreamReader(is));
+
+      // Copy each file that was exported to one of the imports directory
+      String line;
+
+      while ((line = reader.readLine()) != null) {
+        Path p = new Path(line.substring(5));
+        assertTrue(fs.exists(p), "File doesn't exist: " + p);
+        Path importDir = importDirAry[RANDOM.get().nextInt(importDirAry.length)];
+        Path dest = new Path(importDir, p.getName());
+        assertFalse(fs.exists(dest), "Did not expect " + dest + " to exist");
+        FileUtil.copy(fs, p, fs, dest, false, fs.getConf());
+      }
+
+      reader.close();
+
+      log.info("Import dir A: {}", Arrays.toString(fs.listStatus(importDirA)));
+      log.info("Import dir B: {}", Arrays.toString(fs.listStatus(importDirB)));
+
+      // Import the exported data into a new table
+      client.tableOperations().importTable(destTable, importDirs, ImportConfiguration.empty());
+
+      // Get the table ID for the table that the importtable command created
+      final String destTableId = client.tableOperations().tableIdMap().get(destTable);
+      assertNotNull(destTableId);
+
+      // Get all `file` colfams from the metadata table for the new table
+      log.info("Imported into table with ID: {}", destTableId);
+
+      client.tableOperations().getTabletInformation(destTable, new Range())
+          .forEach(tabletInformation -> assertEquals(TabletAvailability.ONDEMAND,
+              tabletInformation.getTabletAvailability(),
+              "Expected all tablets in imported table to be ONDEMAND"));
+    }
+  }
+
   private boolean verifyMappingsFile(String destTableId) throws IOException {
     AccumuloCluster cluster = getCluster();
     assertTrue(cluster instanceof MiniAccumuloClusterImpl);
diff --git a/test/src/main/java/org/apache/accumulo/test/TableOperationsIT.java b/test/src/main/java/org/apache/accumulo/test/TableOperationsIT.java
@@ -83,7 +83,7 @@
 
 public class TableOperationsIT extends AccumuloClusterHarness {
 
-  private AccumuloClient accumuloClient;
+  private static AccumuloClient accumuloClient;
   private static final int MAX_TABLE_NAME_LEN = 1024;
 
   @Override
@@ -727,9 +727,14 @@ private void verifyTablesWithSplits(String tableName, Map<String,String> idMap,
         expectedTabletAvailability);
   }
 
-  private void verifyTabletAvailabilites(String tableName, Range range,
+  public static void verifyTabletAvailabilites(String tableName, Range range,
       List<AvailabilityForTablet> expectedAvailability) throws TableNotFoundException {
-    List<TabletInformation> tabletInfo = accumuloClient.tableOperations()
+    verifyTabletAvailabilites(accumuloClient, tableName, range, expectedAvailability);
+  }
+
+  public static void verifyTabletAvailabilites(AccumuloClient client, String tableName, Range range,
+      List<AvailabilityForTablet> expectedAvailability) throws TableNotFoundException {
+    List<TabletInformation> tabletInfo = client.tableOperations()
         .getTabletInformation(tableName, range).collect(Collectors.toList());
     assertEquals(expectedAvailability.size(), tabletInfo.size());
     for (var i = 0; i < expectedAvailability.size(); i++) {
@@ -739,7 +744,7 @@ private void verifyTabletAvailabilites(String tableName, Range range,
     }
   }
 
-  private void setExpectedTabletAvailability(List<AvailabilityForTablet> expected, String id,
+  public static void setExpectedTabletAvailability(List<AvailabilityForTablet> expected, String id,
       String endRow, String prevEndRow, TabletAvailability availability) {
     KeyExtent ke = new KeyExtent(TableId.of(id), endRow == null ? null : new Text(endRow),
         prevEndRow == null ? null : new Text(prevEndRow));
