Support encrypted columns in Schema Loader (#1975)

scalar-labs · Jun 25, 2024 · 2c8f025 · 2c8f025
1 parent 93a5d19
commit 2c8f025
Show file tree

Hide file tree

Showing 3 changed files with 74 additions and 12 deletions.
diff --git a/core/src/main/java/com/scalar/db/api/TableMetadata.java b/core/src/main/java/com/scalar/db/api/TableMetadata.java
@@ -1,5 +1,6 @@
 package com.scalar.db.api;
 
+import com.google.common.base.MoreObjects;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.scalar.db.api.Scan.Ordering.Order;
@@ -162,7 +163,8 @@ public boolean equals(Object o) {
         && Objects.equals(partitionKeyNames, metadata.partitionKeyNames)
         && Objects.equals(clusteringKeyNames, metadata.clusteringKeyNames)
         && Objects.equals(clusteringOrders, metadata.clusteringOrders)
-        && Objects.equals(secondaryIndexNames, metadata.secondaryIndexNames);
+        && Objects.equals(secondaryIndexNames, metadata.secondaryIndexNames)
+        && Objects.equals(encryptedColumnNames, metadata.encryptedColumnNames);
   }
 
   @Override
@@ -173,7 +175,21 @@ public int hashCode() {
         partitionKeyNames,
         clusteringKeyNames,
         clusteringOrders,
-        secondaryIndexNames);
+        secondaryIndexNames,
+        encryptedColumnNames);
+  }
+
+  @Override
+  public String toString() {
+    return MoreObjects.toStringHelper(this)
+        .add("columnNames", columnNames)
+        .add("columnDataTypes", columnDataTypes)
+        .add("partitionKeyNames", partitionKeyNames)
+        .add("clusteringKeyNames", clusteringKeyNames)
+        .add("clusteringOrders", clusteringOrders)
+        .add("secondaryIndexNames", secondaryIndexNames)
+        .add("encryptedColumnNames", encryptedColumnNames)
+        .toString();
   }
 
   /** A builder class that creates a TableMetadata instance. */

diff --git a/schema-loader/src/main/java/com/scalar/db/schemaloader/TableSchema.java b/schema-loader/src/main/java/com/scalar/db/schemaloader/TableSchema.java
@@ -96,7 +96,7 @@ protected TableMetadata buildTableMetadata(String tableFullName, JsonObject tabl
       for (JsonElement clusteringKeyRaw : clusteringKeys) {
         String clusteringKey;
         String order;
-        String[] clusteringKeyFull = clusteringKeyRaw.getAsString().split(" ", -1);
+        String[] clusteringKeyFull = clusteringKeyRaw.getAsString().split("\\s+", -1);
         if (clusteringKeyFull.length < 2) {
           clusteringKey = clusteringKeyFull[0];
           tableBuilder.addClusteringKey(clusteringKey);
@@ -128,13 +128,21 @@ protected TableMetadata buildTableMetadata(String tableFullName, JsonObject tabl
     traveledKeys.add(COLUMNS);
     for (Entry<String, JsonElement> column : columns.entrySet()) {
       String columnName = column.getKey();
-      DataType columnDataType = DATA_MAP_TYPE.get(column.getValue().getAsString().toUpperCase());
-      if (columnDataType == null) {
+
+      String[] columnDataTypeAndEncrypted = column.getValue().getAsString().split("\\s+", -1);
+      String columnDataType = columnDataTypeAndEncrypted[0];
+      boolean encrypted =
+          columnDataTypeAndEncrypted.length > 1
+              && columnDataTypeAndEncrypted[1].equalsIgnoreCase("ENCRYPTED");
+
+      DataType dataType = DATA_MAP_TYPE.get(columnDataType.toUpperCase());
+      if (dataType == null) {
         throw new IllegalArgumentException(
             CoreError.SCHEMA_LOADER_PARSE_ERROR_INVALID_COLUMN_TYPE.buildMessage(
                 tableFullName, columnName, column.getValue().getAsString()));
       }
-      tableBuilder.addColumn(columnName, columnDataType);
+
+      tableBuilder.addColumn(columnName, dataType, encrypted);
     }
 
     // Add secondary indexes

diff --git a/schema-loader/src/test/java/com/scalar/db/schemaloader/TableSchemaTest.java b/schema-loader/src/test/java/com/scalar/db/schemaloader/TableSchemaTest.java
@@ -140,12 +140,12 @@ public void buildOptions_OptionsMapFromTableDefinitionAndOptionsGiven_ShouldRetu
   }
 
   @Test
-  public void buildTableMetadata_ProperFormatTableDefinitionGiven_ShouldReturnProperTableMetadata()
-      throws SchemaLoaderException {
+  public void
+      buildTableMetadata_ProperFormatTableDefinitionGiven_ShouldReturnProperTableMetadata() {
     String tableDefinitionJson =
         "{\"transaction\": false,"
             + "\"partition-key\": [\"c1\"],"
-            + "\"clustering-key\": [\"c3\",\"c4 ASC\",\"c6 DESC\"],"
+            + "\"clustering-key\": [\"c3\",\"c4 ASC\",\"c6  DESC\"],"
             + "\"columns\": {"
             + "  \"c1\": \"INT\","
             + "  \"c2\": \"TEXT\","
@@ -182,6 +182,44 @@ public void buildTableMetadata_ProperFormatTableDefinitionGiven_ShouldReturnProp
     Assertions.assertThat(tableMetadata).isEqualTo(expectedTableMetadata);
   }
 
+  @Test
+  public void
+      buildTableMetadata_ProperFormatTableDefinitionWithEncryptedColumnsGiven_ShouldReturnProperTableMetadata() {
+    String tableDefinitionJson =
+        "{\"transaction\": false,"
+            + "\"partition-key\": [\"c1\"],"
+            + "\"clustering-key\": [\"c3\",\"c4 ASC\",\"c6  DESC\"],"
+            + "\"columns\": {"
+            + "  \"c1\": \"INT\","
+            + "  \"c2\": \"TEXT ENCRYPTED\","
+            + "  \"c3\": \"BLOB\","
+            + "  \"c4\": \"INT\","
+            + "  \"c5\": \"BOOLEAN  ENCRYPTED\","
+            + "  \"c6\": \"INT\""
+            + "}}";
+    JsonObject tableDefinition = JsonParser.parseString(tableDefinitionJson).getAsJsonObject();
+
+    TableMetadata.Builder tableBuilder = TableMetadata.newBuilder();
+    tableBuilder.addPartitionKey("c1");
+    tableBuilder.addClusteringKey("c3");
+    tableBuilder.addClusteringKey("c4");
+    tableBuilder.addClusteringKey("c6", Order.DESC);
+    tableBuilder.addColumn("c1", DataType.INT);
+    tableBuilder.addColumn("c2", DataType.TEXT, true);
+    tableBuilder.addColumn("c3", DataType.BLOB);
+    tableBuilder.addColumn("c4", DataType.INT);
+    tableBuilder.addColumn("c5", DataType.BOOLEAN, true);
+    tableBuilder.addColumn("c6", DataType.INT);
+    TableMetadata expectedTableMetadata = tableBuilder.build();
+
+    // Act
+    TableMetadata tableMetadata =
+        new TableSchema("ns.tb", tableDefinition, Collections.emptyMap()).getTableMetadata();
+
+    // Assert
+    Assertions.assertThat(tableMetadata).isEqualTo(expectedTableMetadata);
+  }
+
   @Test
   public void Table_WrongFormatTableFullNameGiven_ShouldThrowIllegalArgumentException() {
     // Arrange
@@ -194,11 +232,11 @@ public void Table_WrongFormatTableFullNameGiven_ShouldThrowIllegalArgumentExcept
   }
 
   @Test
-  public void constructor_TableDefinitionWithoutTransactionGiven_ShouldConstructProperTableSchema()
-      throws SchemaLoaderException {
+  public void
+      constructor_TableDefinitionWithoutTransactionGiven_ShouldConstructProperTableSchema() {
     String tableDefinitionJson =
         "{\"partition-key\": [\"c1\"],"
-            + "\"clustering-key\": [\"c3\",\"c4 ASC\",\"c6 DESC\"],"
+            + "\"clustering-key\": [\"c3\",\"c4 ASC\",\"c6  DESC\"],"
             + "\"columns\": {"
             + "  \"c1\": \"INT\","
             + "  \"c2\": \"TEXT\","