scalar-labs · feeblefakie · Oct 23, 2023 · Oct 18, 2023 · Oct 19, 2023 · Oct 19, 2023
diff --git a/.../integration-test/java/com/scalar/db/storage/cassandra/CassandraAdminIntegrationTest.java b/.../integration-test/java/com/scalar/db/storage/cassandra/CassandraAdminIntegrationTest.java
@@ -1,6 +1,7 @@
 package com.scalar.db.storage.cassandra;
 
 import com.scalar.db.api.DistributedStorageAdminIntegrationTestBase;
+import com.scalar.db.util.AdminTestUtils;
 import java.util.Collections;
 import java.util.Map;
 import java.util.Properties;
@@ -15,4 +16,9 @@ protected Properties getProperties(String testName) {
   protected Map<String, String> getCreationOptions() {
     return Collections.singletonMap(CassandraAdmin.REPLICATION_FACTOR, "1");
   }
+
+  @Override
+  protected AdminTestUtils getAdminTestUtils(String testName) {
+    return new CassandraAdminTestUtils(getProperties(testName));
+  }
 }
diff --git a/...ava/com/scalar/db/storage/cassandra/ConsensusCommitAdminIntegrationTestWithCassandra.java b/...ava/com/scalar/db/storage/cassandra/ConsensusCommitAdminIntegrationTestWithCassandra.java
@@ -1,6 +1,7 @@
 package com.scalar.db.storage.cassandra;
 
 import com.scalar.db.transaction.consensuscommit.ConsensusCommitAdminIntegrationTestBase;
+import com.scalar.db.util.AdminTestUtils;
 import java.util.Collections;
 import java.util.Map;
 import java.util.Properties;
@@ -16,4 +17,8 @@ protected Properties getProps(String testName) {
   protected Map<String, String> getCreationOptions() {
     return Collections.singletonMap(CassandraAdmin.REPLICATION_FACTOR, "1");
   }
+
+  protected AdminTestUtils getAdminTestUtils(String testName) {
+    return new CassandraAdminTestUtils(getProperties(testName));
+  }
 }
diff --git a/...test/java/com/scalar/db/storage/cosmos/ConsensusCommitAdminIntegrationTestWithCosmos.java b/...test/java/com/scalar/db/storage/cosmos/ConsensusCommitAdminIntegrationTestWithCosmos.java
@@ -1,6 +1,7 @@
 package com.scalar.db.storage.cosmos;
 
 import com.scalar.db.transaction.consensuscommit.ConsensusCommitAdminIntegrationTestBase;
+import com.scalar.db.util.AdminTestUtils;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Properties;
@@ -37,4 +38,8 @@ private String getNamespace(String namespace) {
   protected Map<String, String> getCreationOptions() {
     return CosmosEnv.getCreationOptions();
   }
+
+  protected AdminTestUtils getAdminTestUtils(String testName) {
+    return new CosmosAdminTestUtils(getProperties(testName));
+  }
 }
diff --git a/core/src/integration-test/java/com/scalar/db/storage/cosmos/CosmosAdminIntegrationTest.java b/core/src/integration-test/java/com/scalar/db/storage/cosmos/CosmosAdminIntegrationTest.java
@@ -1,6 +1,7 @@
 package com.scalar.db.storage.cosmos;
 
 import com.scalar.db.api.DistributedStorageAdminIntegrationTestBase;
+import com.scalar.db.util.AdminTestUtils;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Properties;
@@ -36,4 +37,8 @@ private String getNamespace(String namespace) {
   protected Map<String, String> getCreationOptions() {
     return CosmosEnv.getCreationOptions();
   }
+
+  protected AdminTestUtils getAdminTestUtils(String testName) {
+    return new CosmosAdminTestUtils(getProperties(testName));
+  }
 }
diff --git a/...test/java/com/scalar/db/storage/dynamo/ConsensusCommitAdminIntegrationTestWithDynamo.java b/...test/java/com/scalar/db/storage/dynamo/ConsensusCommitAdminIntegrationTestWithDynamo.java
@@ -1,6 +1,7 @@
 package com.scalar.db.storage.dynamo;
 
 import com.scalar.db.transaction.consensuscommit.ConsensusCommitAdminIntegrationTestBase;
+import com.scalar.db.util.AdminTestUtils;
 import java.util.Map;
 import java.util.Properties;
 
@@ -21,4 +22,8 @@ protected Map<String, String> getCreationOptions() {
   protected boolean isIndexOnBooleanColumnSupported() {
     return false;
   }
+
+  protected AdminTestUtils getAdminTestUtils(String testName) {
+    return new DynamoAdminTestUtils(getProperties(testName));
+  }
 }
diff --git a/core/src/integration-test/java/com/scalar/db/storage/dynamo/DynamoAdminIntegrationTest.java b/core/src/integration-test/java/com/scalar/db/storage/dynamo/DynamoAdminIntegrationTest.java
@@ -1,6 +1,7 @@
 package com.scalar.db.storage.dynamo;
 
 import com.scalar.db.api.DistributedStorageAdminIntegrationTestBase;
+import com.scalar.db.util.AdminTestUtils;
 import java.util.Map;
 import java.util.Properties;
 
@@ -20,4 +21,8 @@ protected Map<String, String> getCreationOptions() {
   protected boolean isIndexOnBooleanColumnSupported() {
     return false;
   }
+
+  protected AdminTestUtils getAdminTestUtils(String testName) {
+    return new DynamoAdminTestUtils(getProperties(testName));
+  }
 }
diff --git a/.../java/com/scalar/db/storage/jdbc/ConsensusCommitAdminIntegrationTestWithJdbcDatabase.java b/.../java/com/scalar/db/storage/jdbc/ConsensusCommitAdminIntegrationTestWithJdbcDatabase.java
@@ -1,6 +1,7 @@
 package com.scalar.db.storage.jdbc;
 
 import com.scalar.db.transaction.consensuscommit.ConsensusCommitAdminIntegrationTestBase;
+import com.scalar.db.util.AdminTestUtils;
 import java.util.Properties;
 
 public class ConsensusCommitAdminIntegrationTestWithJdbcDatabase
@@ -10,4 +11,8 @@ public class ConsensusCommitAdminIntegrationTestWithJdbcDatabase
   protected Properties getProps(String testName) {
     return JdbcEnv.getProperties(testName);
   }
+
+  protected AdminTestUtils getAdminTestUtils(String testName) {
+    return new JdbcAdminTestUtils(getProperties(testName));
+  }
 }
diff --git a/core/src/integration-test/java/com/scalar/db/storage/jdbc/JdbcAdminIntegrationTest.java b/core/src/integration-test/java/com/scalar/db/storage/jdbc/JdbcAdminIntegrationTest.java
@@ -1,6 +1,7 @@
 package com.scalar.db.storage.jdbc;
 
 import com.scalar.db.api.DistributedStorageAdminIntegrationTestBase;
+import com.scalar.db.util.AdminTestUtils;
 import java.util.Properties;
 
 public class JdbcAdminIntegrationTest extends DistributedStorageAdminIntegrationTestBase {
@@ -9,4 +10,8 @@ public class JdbcAdminIntegrationTest extends DistributedStorageAdminIntegration
   protected Properties getProperties(String testName) {
     return JdbcEnv.getProperties(testName);
   }
+
+  protected AdminTestUtils getAdminTestUtils(String testName) {
+    return new JdbcAdminTestUtils(getProperties(testName));
+  }
 }
diff --git a/...gration-test/java/com/scalar/db/transaction/jdbc/JdbcTransactionAdminIntegrationTest.java b/...gration-test/java/com/scalar/db/transaction/jdbc/JdbcTransactionAdminIntegrationTest.java
@@ -4,6 +4,7 @@
 import com.scalar.db.config.DatabaseConfig;
 import com.scalar.db.exception.storage.ExecutionException;
 import com.scalar.db.storage.jdbc.JdbcEnv;
+import com.scalar.db.util.AdminTestUtils;
 import java.util.Properties;
 import org.junit.jupiter.api.Disabled;
 import org.junit.jupiter.api.Test;
@@ -71,4 +72,15 @@ public void dropCoordinatorTables_IfExist_CoordinatorTablesDoNotExist_ShouldNotT
       throws ExecutionException {
     super.dropCoordinatorTables_IfExist_CoordinatorTablesDoNotExist_ShouldNotThrowAnyException();
   }
+
+  @Test
+  @Disabled("JDBC Transaction Admin does not support upgrade()")
+  @Override
+  public void
+      upgrade_WhenMetadataTableExistsButNotNamespacesTable_ShouldCreateNamespacesTableAndImportExistingNamespaces() {}
+
+  @Override
+  protected AdminTestUtils getAdminTestUtils(String testName) {
+    throw new UnsupportedOperationException();
+  }
 }
diff --git a/core/src/main/java/com/scalar/db/api/Admin.java b/core/src/main/java/com/scalar/db/api/Admin.java
@@ -431,4 +431,14 @@ void addNewColumnToTable(String namespace, String table, String columnName, Data
    * @throws ExecutionException if the operation fails
    */
   Set<String> getNamespaceNames() throws ExecutionException;
+
+  /**
+   * Upgrades the Scalar DB environment to support the latest version of the Scalar DB API.
+   * Typically, you will be requested, as indicated on the release notes, to run this method after
+   * updating the Scalar DB version of your application environment.
+   *
+   * @param options options to upgrade
+   * @throws ExecutionException if the operation fails
+   */
+  void upgrade(Map<String, String> options) throws ExecutionException;
 }
diff --git a/core/src/main/java/com/scalar/db/common/CheckedDistributedStorageAdmin.java b/core/src/main/java/com/scalar/db/common/CheckedDistributedStorageAdmin.java
@@ -327,6 +327,15 @@ public void addRawColumnToTable(
     }
   }
 
+  @Override
+  public void upgrade(Map<String, String> options) throws ExecutionException {
+    try {
+      admin.upgrade(options);
+    } catch (ExecutionException e) {
+      throw new ExecutionException("Upgrading the ScalarDB environment failed", e);
+    }
+  }
+
   @Override
   public void close() {
     admin.close();

diff --git a/core/src/main/java/com/scalar/db/service/AdminService.java b/core/src/main/java/com/scalar/db/service/AdminService.java
@@ -122,6 +122,11 @@ public void repairNamespace(String namespace, Map<String, String> options)
     admin.repairNamespace(namespace, options);
   }
 
+  @Override
+  public void upgrade(Map<String, String> options) throws ExecutionException {
+    admin.upgrade(options);
+  }
+
   @Override
   public void close() {
     admin.close();

diff --git a/core/src/main/java/com/scalar/db/storage/cassandra/CassandraAdmin.java b/core/src/main/java/com/scalar/db/storage/cassandra/CassandraAdmin.java
@@ -406,6 +406,26 @@ public Set<String> getNamespaceNames() throws ExecutionException {
     }
   }
 
+  @Override
+  public void upgrade(Map<String, String> options) throws ExecutionException {
+    try {
+      createKeyspace(metadataKeyspace, options, true);
+      createNamespacesTableIfNotExists();
+      // Retrieve user keyspace and filter out system ones. A downside is that this may include
+      // keyspace not created by Scalar DB.
+      Set<String> userKeyspaces =
+          clusterManager.getSession().getCluster().getMetadata().getKeyspaces().stream()
+              .map(KeyspaceMetadata::getName)
+              .filter(name -> !name.startsWith("system") && !name.equals(metadataKeyspace))
+              .collect(Collectors.toSet());
+      for (String userKeyspace : userKeyspaces) {
+        upsertIntoNamespacesTable(userKeyspace);
+      }
+    } catch (RuntimeException e) {
+      throw new ExecutionException("Upgrading the ScalarDB environment failed", e);
+    }
+  }
+
   private void createNamespacesTableIfNotExists() {
     String createTableQuery =
         SchemaBuilder.createTable(

diff --git a/core/src/main/java/com/scalar/db/storage/cosmos/CosmosAdmin.java b/core/src/main/java/com/scalar/db/storage/cosmos/CosmosAdmin.java
@@ -647,6 +647,33 @@ public Set<String> getNamespaceNames() throws ExecutionException {
     }
   }
 
+  @Override
+  public void upgrade(Map<String, String> options) throws ExecutionException {
+    try {
+      if (!tableMetadataContainerExists()) {
+        return;
+      }
+      createMetadataDatabaseAndNamespaceContainerIfNotExists();
+
+      // Upsert namespace of existing tables in the "namespaces" container
+      getTableMetadataContainer()
+          .queryItems(
+              "SELECT container.id FROM container",
+              new CosmosQueryRequestOptions(),
+              CosmosTableMetadata.class)
+          .stream()
+          .map(
+              tableMetadata ->
+                  tableMetadata.getId().substring(0, tableMetadata.getId().indexOf('.')))
+          .distinct()
+          .forEach(
+              namespaceName ->
+                  getNamespacesContainer().upsertItem(new CosmosNamespace(namespaceName)));
+    } catch (RuntimeException e) {
+      throw new ExecutionException("Upgrading the ScalarDB environemnt failed", e);
+    }
+  }
+
   private void createMetadataDatabaseAndNamespaceContainerIfNotExists() {
     ThroughputProperties manualThroughput =
         ThroughputProperties.createManualThroughput(Integer.parseInt(DEFAULT_REQUEST_UNIT));

diff --git a/core/src/main/java/com/scalar/db/storage/dynamo/DynamoAdmin.java b/core/src/main/java/com/scalar/db/storage/dynamo/DynamoAdmin.java
@@ -1363,6 +1363,50 @@ public Set<String> getNamespaceNames() throws ExecutionException {
     }
   }
 
+  @Override
+  public void upgrade(Map<String, String> options) throws ExecutionException {
+    if (!metadataTableExists()) {
+      return;
+    }
+    boolean noBackup = Boolean.parseBoolean(options.getOrDefault(NO_BACKUP, DEFAULT_NO_BACKUP));
+    createNamespacesTableIfNotExists(noBackup);
+    try {
+      for (Namespace namespace : getNamespacesOfExistingTables()) {
+        upsertIntoNamespacesTable(namespace);
+      }
+    } catch (ExecutionException e) {
+      throw new ExecutionException("Upgrading the ScalarDB environment failed", e);
+    }
+  }
+
+  private Set<Namespace> getNamespacesOfExistingTables() throws ExecutionException {
+    Set<Namespace> namespaceNames = new HashSet<>();
+    Map<String, AttributeValue> lastEvaluatedKey = null;
+    do {
+      ScanResponse scanResponse;
+      try {
+        scanResponse =
+            client.scan(
+                ScanRequest.builder()
+                    .tableName(ScalarDbUtils.getFullTableName(metadataNamespace, METADATA_TABLE))
+                    .exclusiveStartKey(lastEvaluatedKey)
+                    .build());
+      } catch (RuntimeException e) {
+        throw new ExecutionException(
+            "failed to retrieve the namespaces names of existing tables", e);
+      }
+      lastEvaluatedKey = scanResponse.lastEvaluatedKey();
+
+      for (Map<String, AttributeValue> tableMetadata : scanResponse.items()) {
+        String fullTableName = tableMetadata.get(METADATA_ATTR_TABLE).s();
+        String namespaceName = fullTableName.substring(0, fullTableName.indexOf('.'));
+        namespaceNames.add(Namespace.of(namespaceName));
+      }
+    } while (!lastEvaluatedKey.isEmpty());
+
+    return namespaceNames;
+  }
+
   private void createNamespacesTableIfNotExists(boolean noBackup) throws ExecutionException {
     try {
       if (!namespacesTableExists()) {

diff --git a/core/src/main/java/com/scalar/db/storage/dynamo/Namespace.java b/core/src/main/java/com/scalar/db/storage/dynamo/Namespace.java
@@ -1,5 +1,7 @@
 package com.scalar.db.storage.dynamo;
 
+import java.util.Objects;
+
 public class Namespace {
   private final String prefix;
   private final String name;
@@ -35,4 +37,21 @@ public String nonPrefixed() {
   public String toString() {
     return prefixed();
   }
+
+  @Override
+  public boolean equals(Object o) {
+    if (this == o) {
+      return true;
+    }
+    if (!(o instanceof Namespace)) {
+      return false;
+    }
+    Namespace namespace = (Namespace) o;
+    return Objects.equals(prefix, namespace.prefix) && Objects.equals(name, namespace.name);
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hash(prefix, name);
+  }
 }