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

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

Author: feeblefakie

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,35 +2,130 @@
 
 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.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(
+                  String.format(
+                      "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",
+                      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;