Backport to branch(3.13) : Add validation for primary key columns in Cosmos DB adapter (#2324)

Co-authored-by: Toshihiro <[email protected]>

Author: feeblefakie

core/src/main/java/com/scalar/db/common/error/CoreError.java

@@ -647,6 +647,13 @@ public enum CoreError implements ScalarDbError {
           + "If you want to modify a condition, please use clearConditions() to remove all existing conditions first",
       "",
       ""),
+  COSMOS_PRIMARY_KEY_CONTAINS_ILLEGAL_CHARACTER(
+      Category.USER_ERROR,
+      "0145",
+      "The value of the column %s in the primary key contains an illegal character. "
+          + "Primary-key columns must not contain any of the following characters in Cosmos DB: ':', '/', '\\', '#', '?'. Value: %s",
+      "",
+      ""),
 
   //
   // Errors for the concurrency error category

core/src/main/java/com/scalar/db/storage/cosmos/ConcatenationVisitor.java

@@ -14,7 +14,8 @@
 import javax.annotation.concurrent.NotThreadSafe;
 
 /**
- * A visitor class to make a concatenated key string for the partition key
+ * A visitor class to make a concatenated key string for the partition key. This uses a colon as a
+ * key separator, so the text column value should not contain colons.
  *
  * @author Yuji Ito
  */
@@ -27,7 +28,6 @@ public ConcatenationVisitor() {
   }
 
   public String build() {
-    // TODO What if the string or blob value includes `:`?
     return String.join(":", values);
   }
 

core/src/main/java/com/scalar/db/storage/cosmos/CosmosOperationChecker.java

@@ -2,36 +2,129 @@
 
 import com.scalar.db.api.ConditionalExpression;
 import com.scalar.db.api.Delete;
+import com.scalar.db.api.Get;
 import com.scalar.db.api.Mutation;
+import com.scalar.db.api.Operation;
 import com.scalar.db.api.Put;
+import com.scalar.db.api.Scan;
 import com.scalar.db.api.TableMetadata;
 import com.scalar.db.common.TableMetadataManager;
 import com.scalar.db.common.checker.OperationChecker;
 import com.scalar.db.common.error.CoreError;
 import com.scalar.db.config.DatabaseConfig;
 import com.scalar.db.exception.storage.ExecutionException;
+import com.scalar.db.io.BigIntColumn;
+import com.scalar.db.io.BlobColumn;
+import com.scalar.db.io.BooleanColumn;
+import com.scalar.db.io.ColumnVisitor;
 import com.scalar.db.io.DataType;
+import com.scalar.db.io.DoubleColumn;
+import com.scalar.db.io.FloatColumn;
+import com.scalar.db.io.IntColumn;
+import com.scalar.db.io.TextColumn;
 
 public class CosmosOperationChecker extends OperationChecker {
+
+  private static final char[] ILLEGAL_CHARACTERS_IN_PRIMARY_KEY = {
+    // Colons are not allowed in primary-key columns due to the `ConcatenationVisitor` limitation.
+    ':',
+
+    // The following characters are not allowed in primary-key columns because they are restricted
+    // and cannot be used in the `Id` property of a Cosmos DB document. For more information, see:
+    // https://learn.microsoft.com/en-us/dotnet/api/microsoft.azure.cosmos.databaseproperties.id?view=azure-dotnet#remarks
+    '/',
+    '\\',
+    '#',
+    '?'
+  };
+
+  private static final ColumnVisitor PRIMARY_KEY_COLUMN_CHECKER =
+      new ColumnVisitor() {
+        @Override
+        public void visit(BooleanColumn column) {}
+
+        @Override
+        public void visit(IntColumn column) {}
+
+        @Override
+        public void visit(BigIntColumn column) {}
+
+        @Override
+        public void visit(FloatColumn column) {}
+
+        @Override
+        public void visit(DoubleColumn column) {}
+
+        @Override
+        public void visit(TextColumn column) {
+          String value = column.getTextValue();
+          assert value != null;
+
+          for (char illegalCharacter : ILLEGAL_CHARACTERS_IN_PRIMARY_KEY) {
+            if (value.indexOf(illegalCharacter) != -1) {
+              throw new IllegalArgumentException(
+                  CoreError.COSMOS_PRIMARY_KEY_CONTAINS_ILLEGAL_CHARACTER.buildMessage(
+                      column.getName(), value));
+            }
+          }
+        }
+
+        @Override
+        public void visit(BlobColumn column) {}
+      };
+
   public CosmosOperationChecker(
       DatabaseConfig databaseConfig, TableMetadataManager metadataManager) {
     super(databaseConfig, metadataManager);
   }
 
+  @Override
+  public void check(Get get) throws ExecutionException {
+    super.check(get);
+    checkPrimaryKey(get);
+  }
+
+  @Override
+  public void check(Scan scan) throws ExecutionException {
+    super.check(scan);
+    checkPrimaryKey(scan);
+    scan.getStartClusteringKey()
+        .ifPresent(
+            c -> c.getColumns().forEach(column -> column.accept(PRIMARY_KEY_COLUMN_CHECKER)));
+    scan.getEndClusteringKey()
+        .ifPresent(
+            c -> c.getColumns().forEach(column -> column.accept(PRIMARY_KEY_COLUMN_CHECKER)));
+  }
+
   @Override
   public void check(Put put) throws ExecutionException {
     super.check(put);
+    checkPrimaryKey(put);
+
     TableMetadata metadata = getTableMetadata(put);
     checkCondition(put, metadata);
   }
 
   @Override
   public void check(Delete delete) throws ExecutionException {
     super.check(delete);
+    checkPrimaryKey(delete);
+
     TableMetadata metadata = getTableMetadata(delete);
     checkCondition(delete, metadata);
   }
 
+  private void checkPrimaryKey(Operation operation) {
+    operation
+        .getPartitionKey()
+        .getColumns()
+        .forEach(column -> column.accept(PRIMARY_KEY_COLUMN_CHECKER));
+    operation
+        .getClusteringKey()
+        .ifPresent(
+            c -> c.getColumns().forEach(column -> column.accept(PRIMARY_KEY_COLUMN_CHECKER)));
+  }
+
   private void checkCondition(Mutation mutation, TableMetadata metadata) {
     if (!mutation.getCondition().isPresent()) {
       return;

core/src/main/java/com/scalar/db/transaction/consensuscommit/CoordinatorGroupCommitter.java

@@ -38,7 +38,7 @@ public static boolean isEnabled(ConsensusCommitConfig config) {
   public static class CoordinatorGroupCommitKeyManipulator
       implements KeyManipulator<String, String, String, String, String> {
     private static final int PRIMARY_KEY_SIZE = 24;
-    private static final char DELIMITER = ':';
+    private static final char DELIMITER = '$';
     private static final int MAX_FULL_KEY_SIZE = 64;
     private static final int MAX_CHILD_KEY_SIZE =
         MAX_FULL_KEY_SIZE - PRIMARY_KEY_SIZE - 1 /* delimiter */;