Table metadata service split to storage and transaction

Author: inv-jishnu

data-loader/cli/src/main/java/com/scalar/db/dataloader/cli/command/dataexport/ExportCommand.java

@@ -27,6 +27,8 @@
 import com.scalar.db.dataloader.core.exception.KeyParsingException;
 import com.scalar.db.dataloader.core.tablemetadata.TableMetadataException;
 import com.scalar.db.dataloader.core.tablemetadata.TableMetadataService;
+import com.scalar.db.dataloader.core.tablemetadata.TableMetadataStorageService;
+import com.scalar.db.dataloader.core.tablemetadata.TableMetadataTransactionService;
 import com.scalar.db.dataloader.core.util.KeyUtils;
 import com.scalar.db.io.Key;
 import com.scalar.db.service.StorageFactory;
@@ -134,10 +136,10 @@ private TableMetadataService createTableMetadataService(
       ScalarDbMode scalarDbMode, String scalarDbPropertiesFilePath) throws IOException {
     if (scalarDbMode.equals(ScalarDbMode.TRANSACTION)) {
       TransactionFactory transactionFactory = TransactionFactory.create(scalarDbPropertiesFilePath);
-      return new TableMetadataService(transactionFactory.getTransactionAdmin());
+      return new TableMetadataTransactionService(transactionFactory.getTransactionAdmin());
     }
     StorageFactory storageFactory = StorageFactory.create(scalarDbPropertiesFilePath);
-    return new TableMetadataService(storageFactory.getStorageAdmin());
+    return new TableMetadataStorageService(storageFactory.getStorageAdmin());
   }
 
   /**

data-loader/cli/src/main/java/com/scalar/db/dataloader/cli/command/dataimport/ImportCommand.java

@@ -24,6 +24,7 @@
 import com.scalar.db.dataloader.core.dataimport.processor.ImportProcessorFactory;
 import com.scalar.db.dataloader.core.tablemetadata.TableMetadataException;
 import com.scalar.db.dataloader.core.tablemetadata.TableMetadataService;
+import com.scalar.db.dataloader.core.tablemetadata.TableMetadataStorageService;
 import com.scalar.db.dataloader.core.util.TableMetadataUtil;
 import com.scalar.db.service.StorageFactory;
 import com.scalar.db.service.TransactionFactory;
@@ -109,7 +110,7 @@ private Map<String, TableMetadata> createTableMetadataMap(
     File configFile = new File(configFilePath);
     StorageFactory storageFactory = StorageFactory.create(configFile);
     try (DistributedStorageAdmin storageAdmin = storageFactory.getStorageAdmin()) {
-      TableMetadataService tableMetadataService = new TableMetadataService(storageAdmin);
+      TableMetadataService tableMetadataService = new TableMetadataStorageService(storageAdmin);
       Map<String, TableMetadata> tableMetadataMap = new HashMap<>();
       if (controlFile != null) {
         for (ControlFileTable table : controlFile.getTables()) {

data-loader/core/src/main/java/com/scalar/db/dataloader/core/tablemetadata/TableMetadataService.java

@@ -1,35 +1,19 @@
 package com.scalar.db.dataloader.core.tablemetadata;
 
-import com.scalar.db.api.DistributedStorageAdmin;
-import com.scalar.db.api.DistributedTransactionAdmin;
 import com.scalar.db.api.TableMetadata;
 import com.scalar.db.common.error.CoreError;
 import com.scalar.db.dataloader.core.util.TableMetadataUtil;
-import com.scalar.db.exception.storage.ExecutionException;
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.Map;
-import lombok.RequiredArgsConstructor;
 
 /**
- * Service for retrieving {@link TableMetadata} from ScalarDB. Provides methods to fetch metadata
- * for individual tables or a collection of tables.
+ * Abstract base class for retrieving {@link TableMetadata} from ScalarDB. Provides shared logic for
+ * fetching metadata for a collection of tables. Subclasses must implement the specific logic for
+ * fetching individual table metadata.
  */
-@RequiredArgsConstructor
-public class TableMetadataService {
+public abstract class TableMetadataService {
 
-  private final DistributedStorageAdmin storageAdmin;
-  private final DistributedTransactionAdmin transactionAdmin;
-
-  public TableMetadataService(DistributedStorageAdmin storageAdmin) {
-    this.transactionAdmin = null;
-    this.storageAdmin = storageAdmin;
-  }
-
-  public TableMetadataService(DistributedTransactionAdmin transactionAdmin) {
-    this.transactionAdmin = transactionAdmin;
-    this.storageAdmin = null;
-  }
   /**
    * Retrieves the {@link TableMetadata} for a specific namespace and table name.
    *
@@ -41,22 +25,12 @@ public TableMetadataService(DistributedTransactionAdmin transactionAdmin) {
    */
   public TableMetadata getTableMetadata(String namespace, String tableName)
       throws TableMetadataException {
-    try {
-      TableMetadata tableMetadata = null;
-      if (storageAdmin != null) {
-        tableMetadata = storageAdmin.getTableMetadata(namespace, tableName);
-      } else if (transactionAdmin != null) {
-        tableMetadata = transactionAdmin.getTableMetadata(namespace, tableName);
-      }
-      if (tableMetadata == null) {
-        throw new TableMetadataException(
-            CoreError.DATA_LOADER_MISSING_NAMESPACE_OR_TABLE.buildMessage(namespace, tableName));
-      }
-      return tableMetadata;
-    } catch (ExecutionException e) {
+    TableMetadata metadata = getTableMetadataInternal(namespace, tableName);
+    if (metadata == null) {
       throw new TableMetadataException(
-          CoreError.DATA_LOADER_TABLE_METADATA_RETRIEVAL_FAILED.buildMessage(e.getMessage()), e);
+          CoreError.DATA_LOADER_MISSING_NAMESPACE_OR_TABLE.buildMessage(namespace, tableName));
     }
+    return metadata;
   }
 
   /**
@@ -75,15 +49,25 @@ public TableMetadata getTableMetadata(String namespace, String tableName)
   public Map<String, TableMetadata> getTableMetadata(Collection<TableMetadataRequest> requests)
       throws TableMetadataException {
     Map<String, TableMetadata> metadataMap = new HashMap<>();
-
     for (TableMetadataRequest request : requests) {
       String namespace = request.getNamespace();
       String tableName = request.getTable();
-      TableMetadata tableMetadata = getTableMetadata(namespace, tableName);
+      TableMetadata metadata = getTableMetadata(namespace, tableName);
       String key = TableMetadataUtil.getTableLookupKey(namespace, tableName);
-      metadataMap.put(key, tableMetadata);
+      metadataMap.put(key, metadata);
     }
-
     return metadataMap;
   }
+
+  /**
+   * Abstract method for retrieving table metadata for a specific namespace and table. Subclasses
+   * must implement this to define how the metadata is fetched.
+   *
+   * @param namespace The namespace of the table.
+   * @param tableName The table name.
+   * @return The {@link TableMetadata} object, or null if not found.
+   * @throws TableMetadataException if an error occurs during metadata retrieval.
+   */
+  protected abstract TableMetadata getTableMetadataInternal(String namespace, String tableName)
+      throws TableMetadataException;
 }

data-loader/core/src/main/java/com/scalar/db/dataloader/core/tablemetadata/TableMetadataStorageService.java

@@ -0,0 +1,37 @@
+package com.scalar.db.dataloader.core.tablemetadata;
+
+import com.scalar.db.api.DistributedStorageAdmin;
+import com.scalar.db.api.TableMetadata;
+import com.scalar.db.common.error.CoreError;
+import com.scalar.db.exception.storage.ExecutionException;
+import lombok.RequiredArgsConstructor;
+
+/**
+ * Implementation of {@link TableMetadataService} that retrieves table metadata using {@link
+ * DistributedStorageAdmin}.
+ */
+@RequiredArgsConstructor
+public class TableMetadataStorageService extends TableMetadataService {
+
+  private final DistributedStorageAdmin storageAdmin;
+
+  /**
+   * Retrieves the {@link TableMetadata} for a given namespace and table using the {@link
+   * DistributedStorageAdmin}.
+   *
+   * @param namespace The namespace of the table.
+   * @param tableName The name of the table.
+   * @return The {@link TableMetadata} for the specified table, or null if not found.
+   * @throws TableMetadataException If an error occurs while fetching metadata.
+   */
+  @Override
+  protected TableMetadata getTableMetadataInternal(String namespace, String tableName)
+      throws TableMetadataException {
+    try {
+      return storageAdmin.getTableMetadata(namespace, tableName);
+    } catch (ExecutionException e) {
+      throw new TableMetadataException(
+          CoreError.DATA_LOADER_TABLE_METADATA_RETRIEVAL_FAILED.buildMessage(e.getMessage()), e);
+    }
+  }
+}

data-loader/core/src/main/java/com/scalar/db/dataloader/core/tablemetadata/TableMetadataTransactionService.java

@@ -0,0 +1,37 @@
+package com.scalar.db.dataloader.core.tablemetadata;
+
+import com.scalar.db.api.DistributedTransactionAdmin;
+import com.scalar.db.api.TableMetadata;
+import com.scalar.db.common.error.CoreError;
+import com.scalar.db.exception.storage.ExecutionException;
+import lombok.RequiredArgsConstructor;
+
+/**
+ * Implementation of {@link TableMetadataService} that retrieves table metadata using {@link
+ * DistributedTransactionAdmin}.
+ */
+@RequiredArgsConstructor
+public class TableMetadataTransactionService extends TableMetadataService {
+
+  private final DistributedTransactionAdmin transactionAdmin;
+
+  /**
+   * Retrieves the {@link TableMetadata} for a given namespace and table using the {@link
+   * DistributedTransactionAdmin}.
+   *
+   * @param namespace The namespace of the table.
+   * @param tableName The name of the table.
+   * @return The {@link TableMetadata} for the specified table, or null if not found.
+   * @throws TableMetadataException If an error occurs while fetching metadata.
+   */
+  @Override
+  protected TableMetadata getTableMetadataInternal(String namespace, String tableName)
+      throws TableMetadataException {
+    try {
+      return transactionAdmin.getTableMetadata(namespace, tableName);
+    } catch (ExecutionException e) {
+      throw new TableMetadataException(
+          CoreError.DATA_LOADER_TABLE_METADATA_RETRIEVAL_FAILED.buildMessage(e.getMessage()), e);
+    }
+  }
+}

data-loader/core/src/test/java/com/scalar/db/dataloader/core/tablemetadata/TableMetadataServiceTest.java

@@ -0,0 +1,58 @@
+package com.scalar.db.dataloader.core.tablemetadata;
+
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+import com.scalar.db.api.DistributedStorageAdmin;
+import com.scalar.db.api.TableMetadata;
+import com.scalar.db.common.error.CoreError;
+import com.scalar.db.dataloader.core.UnitTestUtils;
+import com.scalar.db.exception.storage.ExecutionException;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.Map;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.mockito.Mockito;
+
+class TableMetadataStorageServiceTest {
+
+  DistributedStorageAdmin storageAdmin;
+  TableMetadataService tableMetadataService;
+
+  @BeforeEach
+  void setup() throws ExecutionException {
+    storageAdmin = Mockito.mock(DistributedStorageAdmin.class);
+    Mockito.when(storageAdmin.getTableMetadata("namespace", "table"))
+        .thenReturn(UnitTestUtils.createTestTableMetadata());
+
+    tableMetadataService = new TableMetadataStorageService(storageAdmin);
+  }
+
+  @Test
+  void getTableMetadata_withValidNamespaceAndTable_shouldReturnTableMetadataMap()
+      throws TableMetadataException {
+    Map<String, TableMetadata> expected = new HashMap<>();
+    expected.put("namespace.table", UnitTestUtils.createTestTableMetadata());
+
+    TableMetadataRequest tableMetadataRequest = new TableMetadataRequest("namespace", "table");
+    Map<String, TableMetadata> output =
+        tableMetadataService.getTableMetadata(Collections.singleton(tableMetadataRequest));
+
+    Assertions.assertEquals(expected.get("namespace.table"), output.get("namespace.table"));
+  }
+
+  @Test
+  void getTableMetadata_withInvalidNamespaceAndTable_shouldThrowException()
+      throws ExecutionException {
+    Mockito.when(storageAdmin.getTableMetadata("namespace2", "table2")).thenReturn(null);
+
+    TableMetadataRequest tableMetadataRequest = new TableMetadataRequest("namespace2", "table2");
+    assertThatThrownBy(
+            () ->
+                tableMetadataService.getTableMetadata(Collections.singleton(tableMetadataRequest)))
+        .isInstanceOf(TableMetadataException.class)
+        .hasMessage(
+            CoreError.DATA_LOADER_MISSING_NAMESPACE_OR_TABLE.buildMessage("namespace2", "table2"));
+  }
+}