[#515] improvement(core): Add UT to test restart for RocksDBBackend (…

…#515) ### What changes were proposed in this pull request? Add UT to test restart and reopen for `RocksDBBackend` ### Why are the changes needed? Fully test `RocksDBBackend` to verify if it works well. Fix: #514 ### Does this PR introduce _any_ user-facing change? N/A ### How was this patch tested? UT
apache · Oct 16, 2023 · fa7889f · fa7889f
1 parent 16d0075
commit fa7889f
Show file tree

Hide file tree

Showing 2 changed files with 87 additions and 3 deletions.
diff --git a/core/src/main/java/com/datastrato/graviton/storage/RandomIdGenerator.java b/core/src/main/java/com/datastrato/graviton/storage/RandomIdGenerator.java
@@ -8,18 +8,18 @@
 import java.util.UUID;
 
 /**
- * Random id generator. This is used to generate random id for entities. Please see {@link
+ * Random id generator. This is used to generate random ids for entities. Please see {@link
  * com.datastrato.graviton.meta.BaseMetalake#ID} for more details.
  */
 public class RandomIdGenerator implements IdGenerator {
 
-  public static final RandomIdGenerator INSTACNE = new RandomIdGenerator();
+  public static final RandomIdGenerator INSTANCE = new RandomIdGenerator();
 
   public static final long MAX_ID = 0x7fffffffffffffffL;
 
   @Override
   public long nextId() {
-    // Make sure this is a postive number.
+    // Make sure this is a positive number.
     return UUID.randomUUID().getLeastSignificantBits() & MAX_ID;
   }
 }
diff --git a/core/src/test/java/com/datastrato/graviton/storage/kv/TestKvEntityStorage.java b/core/src/test/java/com/datastrato/graviton/storage/kv/TestKvEntityStorage.java
@@ -93,6 +93,90 @@ public TableEntity createTableEntity(Namespace namespace, String name, AuditInfo
         .build();
   }
 
+  @Test
+  void testRestart() throws IOException {
+    Config config = Mockito.mock(Config.class);
+    Mockito.when(config.get(ENTITY_STORE)).thenReturn("kv");
+    Mockito.when(config.get(ENTITY_KV_STORE)).thenReturn(DEFAULT_ENTITY_KV_STORE);
+    Mockito.when(config.get(Configs.ENTITY_SERDE)).thenReturn("proto");
+    Mockito.when(config.get(ENTRY_KV_ROCKSDB_BACKEND_PATH)).thenReturn(ROCKS_DB_STORE_PATH);
+
+    Assertions.assertEquals(ROCKS_DB_STORE_PATH, config.get(ENTRY_KV_ROCKSDB_BACKEND_PATH));
+    AuditInfo auditInfo =
+        new AuditInfo.Builder().withCreator("creator").withCreateTime(Instant.now()).build();
+
+    try (EntityStore store = EntityStoreFactory.createEntityStore(config)) {
+      store.initialize(config);
+      Assertions.assertTrue(store instanceof KvEntityStore);
+      store.setSerDe(EntitySerDeFactory.createEntitySerDe(config.get(Configs.ENTITY_SERDE)));
+
+      BaseMetalake metalake = createBaseMakeLake("metalake", auditInfo);
+      CatalogEntity catalog = createCatalog(Namespace.of("metalake"), "catalog", auditInfo);
+      CatalogEntity catalogCopy = createCatalog(Namespace.of("metalake"), "catalogCopy", auditInfo);
+
+      SchemaEntity schema1 =
+          createSchemaEntity(Namespace.of("metalake", "catalog"), "schema1", auditInfo);
+      TableEntity table1 =
+          createTableEntity(Namespace.of("metalake", "catalog", "schema1"), "table1", auditInfo);
+
+      // Store all entities
+      store.put(metalake);
+      store.put(catalog);
+      store.put(catalogCopy);
+      store.put(schema1);
+      store.put(table1);
+
+      Assertions.assertDoesNotThrow(
+          () -> store.get(NameIdentifier.of("metalake"), EntityType.METALAKE, BaseMetalake.class));
+      Assertions.assertDoesNotThrow(
+          () ->
+              store.get(
+                  NameIdentifier.of("metalake", "catalog"),
+                  EntityType.CATALOG,
+                  CatalogEntity.class));
+      Assertions.assertDoesNotThrow(
+          () ->
+              store.get(
+                  NameIdentifier.of("metalake", "catalog", "schema1"),
+                  EntityType.SCHEMA,
+                  SchemaEntity.class));
+      Assertions.assertDoesNotThrow(
+          () ->
+              store.get(
+                  NameIdentifier.of("metalake", "catalog", "schema1", "table1"),
+                  EntityType.TABLE,
+                  TableEntity.class));
+    }
+
+    // It will automatically close the store we create before, then we reopen the entity store
+    try (EntityStore store = EntityStoreFactory.createEntityStore(config)) {
+      store.initialize(config);
+      Assertions.assertTrue(store instanceof KvEntityStore);
+      store.setSerDe(EntitySerDeFactory.createEntitySerDe(config.get(Configs.ENTITY_SERDE)));
+
+      Assertions.assertDoesNotThrow(
+          () -> store.get(NameIdentifier.of("metalake"), EntityType.METALAKE, BaseMetalake.class));
+      Assertions.assertDoesNotThrow(
+          () ->
+              store.get(
+                  NameIdentifier.of("metalake", "catalog"),
+                  EntityType.CATALOG,
+                  CatalogEntity.class));
+      Assertions.assertDoesNotThrow(
+          () ->
+              store.get(
+                  NameIdentifier.of("metalake", "catalog", "schema1"),
+                  EntityType.SCHEMA,
+                  SchemaEntity.class));
+      Assertions.assertDoesNotThrow(
+          () ->
+              store.get(
+                  NameIdentifier.of("metalake", "catalog", "schema1", "table1"),
+                  EntityType.TABLE,
+                  TableEntity.class));
+    }
+  }
+
   @Test
   void testEntityUpdate() throws Exception {
     Config config = Mockito.mock(Config.class);