Fix

Author: brfrn169

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

@@ -78,9 +78,15 @@ public void check(Get get, TransactionContext context) throws ExecutionException
     // Additional checks for SERIALIZABLE isolation level
     if (context.isolation == Isolation.SERIALIZABLE) {
       // Don't allow index gets
-      if (ScalarDbUtils.isSecondaryIndexSpecified(get, metadata.getTableMetadata())) {
-        throw new IllegalArgumentException(
-            CoreError.CONSENSUS_COMMIT_INDEX_GET_NOT_ALLOWED_IN_SERIALIZABLE.buildMessage());
+      TableMetadata tableMetadata = metadata.getTableMetadata();
+      if (ScalarDbUtils.isSecondaryIndexSpecified(get, tableMetadata)) {
+        // If the index column is part of the primary key, it's allowed
+        String indexKeyColumnName = get.getPartitionKey().getColumns().get(0).getName();
+        if (!tableMetadata.getPartitionKeyNames().contains(indexKeyColumnName)
+            && !tableMetadata.getClusteringKeyNames().contains(indexKeyColumnName)) {
+          throw new IllegalArgumentException(
+              CoreError.CONSENSUS_COMMIT_INDEX_GET_NOT_ALLOWED_IN_SERIALIZABLE.buildMessage());
+        }
       }
     }
   }
@@ -137,9 +143,15 @@ public void check(Scan scan, TransactionContext context) throws ExecutionExcepti
     // Additional checks for SERIALIZABLE isolation level
     if (context.isolation == Isolation.SERIALIZABLE) {
       // Don't allow index scans
-      if (ScalarDbUtils.isSecondaryIndexSpecified(scan, metadata.getTableMetadata())) {
-        throw new IllegalArgumentException(
-            CoreError.CONSENSUS_COMMIT_INDEX_SCAN_NOT_ALLOWED_IN_SERIALIZABLE.buildMessage());
+      TableMetadata tableMetadata = metadata.getTableMetadata();
+      if (ScalarDbUtils.isSecondaryIndexSpecified(scan, tableMetadata)) {
+        // If the index column is part of the primary key, it's allowed
+        String indexKeyColumnName = scan.getPartitionKey().getColumns().get(0).getName();
+        if (!tableMetadata.getPartitionKeyNames().contains(indexKeyColumnName)
+            && !tableMetadata.getClusteringKeyNames().contains(indexKeyColumnName)) {
+          throw new IllegalArgumentException(
+              CoreError.CONSENSUS_COMMIT_INDEX_SCAN_NOT_ALLOWED_IN_SERIALIZABLE.buildMessage());
+        }
       }
 
       // If the scan is a cross-partition scan (ScanAll), don't allow conditions on indexed columns

core/src/test/java/com/scalar/db/transaction/consensuscommit/ConsensusCommitOperationCheckerTest.java

@@ -284,6 +284,48 @@ public void checkForGet_WithSecondaryIndexInSerializable_ShouldThrowIllegalArgum
         .isInstanceOf(IllegalArgumentException.class);
   }
 
+  @Test
+  public void checkForGet_WithIndexKeyUsingPartitionKeyInSerializable_ShouldNotThrowException() {
+    // Arrange
+    TableMetadata metadata =
+        TableMetadata.newBuilder()
+            .addColumn("pk", DataType.INT)
+            .addColumn("col", DataType.INT)
+            .addPartitionKey("pk")
+            .addSecondaryIndex("pk")
+            .build();
+    when(tableMetadata.getTableMetadata()).thenReturn(metadata);
+
+    Get get = Get.newBuilder().namespace("ns").table("tbl").indexKey(Key.ofInt("pk", 100)).build();
+    TransactionContext context =
+        new TransactionContext("txId", null, Isolation.SERIALIZABLE, false, false);
+
+    // Act Assert
+    assertThatCode(() -> checker.check(get, context)).doesNotThrowAnyException();
+  }
+
+  @Test
+  public void checkForGet_WithIndexKeyUsingClusteringKeyInSerializable_ShouldNotThrowException() {
+    // Arrange
+    TableMetadata metadata =
+        TableMetadata.newBuilder()
+            .addColumn("pk", DataType.INT)
+            .addColumn("ck", DataType.INT)
+            .addColumn("col", DataType.INT)
+            .addPartitionKey("pk")
+            .addClusteringKey("ck")
+            .addSecondaryIndex("ck")
+            .build();
+    when(tableMetadata.getTableMetadata()).thenReturn(metadata);
+
+    Get get = Get.newBuilder().namespace("ns").table("tbl").indexKey(Key.ofInt("ck", 100)).build();
+    TransactionContext context =
+        new TransactionContext("txId", null, Isolation.SERIALIZABLE, false, false);
+
+    // Act Assert
+    assertThatCode(() -> checker.check(get, context)).doesNotThrowAnyException();
+  }
+
   @Test
   public void checkForGet_ValidGet_ShouldNotThrowException() {
     // Arrange
@@ -402,6 +444,50 @@ public void checkForScan_WithSecondaryIndexInSerializable_ShouldThrowIllegalArgu
         .isInstanceOf(IllegalArgumentException.class);
   }
 
+  @Test
+  public void checkForScan_WithIndexKeyUsingPartitionKeyInSerializable_ShouldNotThrowException() {
+    // Arrange
+    TableMetadata metadata =
+        TableMetadata.newBuilder()
+            .addColumn("pk", DataType.INT)
+            .addColumn("col", DataType.INT)
+            .addPartitionKey("pk")
+            .addSecondaryIndex("pk")
+            .build();
+    when(tableMetadata.getTableMetadata()).thenReturn(metadata);
+
+    Scan scan =
+        Scan.newBuilder().namespace("ns").table("tbl").indexKey(Key.ofInt("pk", 100)).build();
+    TransactionContext context =
+        new TransactionContext("txId", null, Isolation.SERIALIZABLE, false, false);
+
+    // Act Assert
+    assertThatCode(() -> checker.check(scan, context)).doesNotThrowAnyException();
+  }
+
+  @Test
+  public void checkForScan_WithIndexKeyUsingClusteringKeyInSerializable_ShouldNotThrowException() {
+    // Arrange
+    TableMetadata metadata =
+        TableMetadata.newBuilder()
+            .addColumn("pk", DataType.INT)
+            .addColumn("ck", DataType.INT)
+            .addColumn("col", DataType.INT)
+            .addPartitionKey("pk")
+            .addClusteringKey("ck")
+            .addSecondaryIndex("ck")
+            .build();
+    when(tableMetadata.getTableMetadata()).thenReturn(metadata);
+
+    Scan scan =
+        Scan.newBuilder().namespace("ns").table("tbl").indexKey(Key.ofInt("ck", 100)).build();
+    TransactionContext context =
+        new TransactionContext("txId", null, Isolation.SERIALIZABLE, false, false);
+
+    // Act Assert
+    assertThatCode(() -> checker.check(scan, context)).doesNotThrowAnyException();
+  }
+
   @Test
   public void
       checkForScan_ScanAllWithConditionOnIndexedColumnInSerializable_ShouldThrowIllegalArgumentException() {