table store batch to reject batches with different partitionkeys

ChangeLog.md

@@ -10,6 +10,10 @@ Blob:
 - Fix HTTP header parsing of `SubmitBatch()`. If a HTTP header has HTTP header delimiter (`:`) in its value, `SubmitBatch()` returns "400 One of the request inputs is not valid". For example, if `user-agent` header is `azsdk-cpp-storage-blobs/12.10.0-beta.1 (Darwin 23.1.0 arm64 Darwin Kernel Version 23.1.0: Mon Oct  9 21:28:12 PDT 2023; root:xnu-10002.41.9~6/RELEASE_ARM64_T8103)`, all `SubmitBatch()` requests are failed.
 - Fixed issue of blob copying succeed without 'r' permission in source blob's SAS token credential.
 
+Table:
+
+- Rejects batch opertations with differing partition keys (issue #1215)
+
 ## 2023.12 Version 3.29.0
 
 General:

src/table/batch/TableBatchOrchestrator.ts

@@ -94,7 +94,6 @@ export default class TableBatchOrchestrator {
       const batchId = uuidv4();
 
       this.checkForPartitionKey();
-
       // initialize transaction rollback capability
       await this.initTransaction(batchId);
 
@@ -722,9 +721,34 @@ export default class TableBatchOrchestrator {
     batchContextClone: any
   ) {
     if (partitionKey === undefined) {
+
       throw StorageErrorFactory.getInvalidInput(batchContextClone);
+
     }
+    this.checkForDifferingPartitionKeys(partitionKey, batchContextClone);
     this.checkForDuplicateRowKey(partitionKey, rowKey, batchContextClone);
+
+  }
+
+
+  /**
+   * Checks that the partition key is the same for all requests in the batch
+   * @param {string} partitionKey
+   * @param {*} batchContextClone
+   * @memberof TableBatchOrchestrator
+   */
+  private checkForDifferingPartitionKeys(
+    partitionKey: string,
+    batchContextClone: any
+  ) {
+    if (this.partitionKeyMap.size > 0) {
+      console.log("checking for differing partition keys");
+      if (this.partitionKeyMap.values().next().value !== partitionKey) {
+        throw StorageErrorFactory.getBatchPartitionKeyMismatch(
+          batchContextClone
+        );
+      }
+    }
   }
 
   /**

src/table/errors/StorageErrorFactory.ts

@@ -12,7 +12,7 @@ const defaultID: string = "DefaultID";
 // see: https://learn.microsoft.com/en-us/rest/api/storageservices/table-service-error-codes
 export default class StorageErrorFactory {
 
-  static getBatchDuplicateRowKey(context: Context, rowKey: string) : StorageError{
+  static getBatchDuplicateRowKey(context: Context, rowKey: string): StorageError {
     return new StorageError(
       400,
       "InvalidDuplicateRow",
@@ -23,6 +23,17 @@ export default class StorageErrorFactory {
     );
   }
 
+  static getBatchPartitionKeyMismatch(context: Context): StorageError {
+    return new StorageError(
+      400,
+      "InvalidPartitionKey",
+      `All entities in the batch must have the same partition key. `,
+      context.contextID || defaultID,
+      undefined,
+      context
+    );
+  }
+
   public static getInvalidHeaderValue(
     context: Context,
     additionalMessages?: { [key: string]: string }

tests/table/apis/table.batch.errorhandling.test.ts

@@ -604,4 +604,46 @@ describe("table Entity APIs test", () => {
 
     await tableClientrollback.deleteTable();
   });
+
+  //Skips this test because it requires modifying the Azure Data Tables client so that it does not perform a check for partition key consistency before submitting the batch.
+  it.skip("10. Batch API should fail when attempting to insert elements with different partitionkeys in same batch, @loki", async () => {
+    const partitionKey = createUniquePartitionKey("1");
+    const partitionKey2 = createUniquePartitionKey("2");
+    console.log(`PartitionKey: ${partitionKey} partitionKey2: ${partitionKey2}`);
+    const tableNameBatchError: string = getUniqueName("datatables");
+    const testEntities: TableTestEntity[] = [
+      entityFactory.createBasicEntityForTest(partitionKey),
+      entityFactory.createBasicEntityForTest(partitionKey2)
+    ];
+
+    const tableClientrollback = createAzureDataTablesClient(
+      testLocalAzuriteInstance,
+      tableNameBatchError
+    );
+
+    await tableClientrollback.createTable();
+
+    const transaction = new TableTransaction();
+    for (const testEntity of testEntities) {
+      transaction.createEntity(testEntity);
+    }
+    try {
+      await tableClientrollback.submitTransaction(
+        transaction.actions
+      );
+      assert.fail("Did not get expected 400 (InvalidInput) error.");
+    } catch (err: any) {
+      if (err.code === "ERR_ASSERTION") {
+        throw err;
+      }
+      const restErr = err as RestError;
+      assert.strictEqual(
+        restErr.statusCode,
+        400,
+        "Did not get expected 400 (InvalidInput) error."
+      );
+    }
+
+    await tableClientrollback.deleteTable();
+  });
 });