Added Hierarchical Partition Key Support

sdk/cosmosdb/cosmos/src/ClientContext.ts

@@ -13,7 +13,7 @@ import { Constants, HTTPMethod, OperationType, ResourceType } from "./common/con
 import { getIdFromLink, getPathFromLink, parseLink } from "./common/helper";
 import { StatusCodes, SubStatusCodes } from "./common/statusCodes";
 import { Agent, CosmosClientOptions } from "./CosmosClientOptions";
-import { ConnectionPolicy, ConsistencyLevel, DatabaseAccount, PartitionKey } from "./documents";
+import { ConnectionPolicy, ConsistencyLevel, DatabaseAccount, PartitionKey, mapPartitionToInternal } from "./documents";
 import { GlobalEndpointManager } from "./globalEndpointManager";
 import { PluginConfig, PluginOn, executePlugins } from "./plugins/Plugin";
 import { FetchFunctionCallback, SqlQuerySpec } from "./queryExecutionContext";
@@ -757,7 +757,7 @@ export class ClientContext {
       options: requestContext.options,
       partitionKeyRangeId: requestContext.partitionKeyRangeId,
       useMultipleWriteLocations: this.connectionPolicy.useMultipleWriteLocations,
-      partitionKey: requestContext.partitionKey,
+      partitionKey: requestContext.partitionKey !== undefined ? mapPartitionToInternal(requestContext.partitionKey) : undefined,
     });
   }
 

sdk/cosmosdb/cosmos/src/client/Item/Item.ts

@@ -9,7 +9,7 @@ import {
   ResourceType,
   StatusCodes,
 } from "../../common";
-import { PartitionKey } from "../../documents";
+import { PartitionKey, PartitionKeyInternal, mapPartitionToInternal } from "../../documents";
 import { extractPartitionKey, undefinedPartitionKey } from "../../extractPartitionKey";
 import { RequestOptions, Response } from "../../request";
 import { PatchRequestBody } from "../../utils/patch";
@@ -24,7 +24,7 @@ import { ItemResponse } from "./ItemResponse";
  * @see {@link Items} for operations on all items; see `container.items`.
  */
 export class Item {
-  private partitionKey: PartitionKey;
+  private partitionKey: PartitionKeyInternal;
   /**
    * Returns a reference URL to the resource. Used for linking in Permissions.
    */
@@ -41,10 +41,10 @@ export class Item {
   constructor(
     public readonly container: Container,
     public readonly id: string,
-    partitionKey: PartitionKey,
+    partitionKey: PartitionKey | undefined,
     private readonly clientContext: ClientContext
   ) {
-    this.partitionKey = partitionKey;
+    this.partitionKey = partitionKey === undefined ? undefined : mapPartitionToInternal(partitionKey);
   }
 
   /**

sdk/cosmosdb/cosmos/src/client/Item/Items.ts

@@ -17,16 +17,16 @@ import { ItemResponse } from "./ItemResponse";
 import {
   Batch,
   isKeyInRange,
-  Operation,
-  getPartitionKeyToHash,
-  decorateOperation,
+  prepareOperations,
   OperationResponse,
   OperationInput,
   BulkOptions,
   decorateBatchOperation,
+  hasResource,
 } from "../../utils/batch";
-import { hashV1PartitionKey } from "../../utils/hashing/v1";
-import { hashV2PartitionKey } from "../../utils/hashing/v2";
+import { stripUndefined } from "../../utils/typeChecks";
+import { hashPartitionKey } from "../../utils/hashing/hash";
+import { PartitionKeyDefinition } from "../../documents";
 
 /**
  * @hidden
@@ -408,7 +408,8 @@ export class Items {
     const { resources: partitionKeyRanges } = await this.container
       .readPartitionKeyRanges()
       .fetchAll();
-    const { resource: definition } = await this.container.getPartitionKeyDefinition();
+    let { resource } = await this.container.readPartitionKeyDefinition();
+    const partitionDefinition = stripUndefined(resource, "PartitionKeyDefinition.");
     const batches: Batch[] = partitionKeyRanges.map((keyRange: PartitionKeyRange) => {
       return {
         min: keyRange.minInclusive,
@@ -418,19 +419,8 @@ export class Items {
         operations: [],
       };
     });
-    operations
-      .map((operation) => decorateOperation(operation, definition, options))
-      .forEach((operation: Operation, index: number) => {
-        const partitionProp = definition.paths[0].replace("/", "");
-        const isV2 = definition.version && definition.version === 2;
-        const toHashKey = getPartitionKeyToHash(operation, partitionProp);
-        const hashed = isV2 ? hashV2PartitionKey(toHashKey) : hashV1PartitionKey(toHashKey);
-        const batchForKey = batches.find((batch: Batch) => {
-          return isKeyInRange(batch.min, batch.max, hashed);
-        });
-        batchForKey.operations.push(operation);
-        batchForKey.indexes.push(index);
-      });
+
+    this.groupOperationsBasedOnPartitionKey(operations, partitionDefinition, options, batches);
 
     const path = getPathFromLink(this.container.url, ResourceType.item);
 
@@ -460,7 +450,7 @@ export class Items {
             // partition key types as well since we don't support them, so for now we throw
             if (err.code === 410) {
               throw new Error(
-                "Partition key error. Either the partitions have split or an operation has an unsupported partitionKey type"
+                "Partition key error. Either the partitions have split or an operation has an unsupported partitionKey type" +  err.message
               );
             }
             throw new Error(`Bulk request errored with: ${err.message}`);
@@ -470,6 +460,27 @@ export class Items {
     return orderedResponses;
   }
 
+  /**
+   * Function to create batches based of partition key Ranges.
+   * 1. if {@link operation} has {@link hasResource} then {@link PartitionKey} is extracted from there.
+   * @param operations 
+   * @param partitionDefinition 
+   * @param options 
+   * @param batches 
+   */
+  private groupOperationsBasedOnPartitionKey(operations: OperationInput[], partitionDefinition: PartitionKeyDefinition, options: RequestOptions | undefined, batches: Batch[]) {
+    operations
+      .forEach((operationInput, index: number) => {
+        const { operation, partitionKey } = prepareOperations(operationInput, partitionDefinition, options);
+        const hashed = hashPartitionKey(stripUndefined(partitionKey, "PartitionKey"), partitionDefinition);
+        const batchForKey = stripUndefined(batches.find((batch: Batch) => {
+          return isKeyInRange(batch.min, batch.max, hashed);
+        }), "No suitable Batch found.");
+        batchForKey.operations.push(operation);
+        batchForKey.indexes.push(index);
+      });
+  }
+
   /**
    * Execute transactional batch operations on items.
    *

sdk/cosmosdb/cosmos/src/documents/PartitionKey.ts

@@ -1,5 +1,45 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT license.
-import { PartitionKeyDefinition } from "./PartitionKeyDefinition";
 
-export type PartitionKey = PartitionKeyDefinition | string | number | unknown;
+export type PartitionKey = PrimitivePartitionKeyValue | PrimitivePartitionKeyValue[];
+
+/**
+ * Internal Representation Of Partition Key.
+ */
+export type PartitionKeyInternal = PrimitivePartitionKeyValue[];
+
+/**
+ * A primitive Partition Key value.
+ */
+export type PrimitivePartitionKeyValue =
+  | string
+  | number
+  | boolean
+  | NullPartitionType
+  | NonePartitionKey;
+
+/**
+ * The returned object represents a partition key value that allows creating and accessing items
+ * with a null value for the partition key.
+ */
+export type NullPartitionType = null;
+export const NullPartitionKeyLiteral: NullPartitionType = null;
+
+/**
+ * The returned object represents a partition key value that allows creating and accessing items
+ * without a value for partition key
+ */
+export type NonePartitionKey = {
+  [K in any]: never;
+};
+export const NonePartitionKeyLiteral: NonePartitionKey = {};
+
+/**
+ * Maps PartitionKey to InternalPartitionKey.
+ * @param partitionKey 
+ * @returns 
+ */
+export function mapPartitionToInternal(partitionKey: PartitionKey): PartitionKeyInternal {
+  if (Array.isArray(partitionKey)) return partitionKey;
+  else return [partitionKey];
+}

sdk/cosmosdb/cosmos/src/documents/PartitionKeyDefinition.ts

@@ -1,5 +1,8 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT license.
+import { PartitionKeyDefinitionVersion } from "./PartitionKeyDefinitionVersion";
+import { PartitionKeyKind } from "./PartitionKind";
+
 export interface PartitionKeyDefinition {
   /**
    * An array of paths for which data within the collection can be partitioned. Paths must not contain a wildcard or
@@ -11,6 +14,10 @@ export interface PartitionKeyDefinition {
    * An optional field, if not specified the default value is 1. To use the large partition key set the version to 2.
    * To learn about large partition keys, see [how to create containers with large partition key](https://docs.microsoft.com/en-us/azure/cosmos-db/large-partition-keys) article.
    */
-  version?: number;
+  version?: PartitionKeyDefinitionVersion;
   systemKey?: boolean;
+  /**
+   * What kind of partition key is being defined (default: "Hash")
+   */
+  kind?: PartitionKeyKind;
 }

sdk/cosmosdb/cosmos/src/documents/PartitionKeyDefinitionVersion.ts

@@ -0,0 +1,6 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+export enum PartitionKeyDefinitionVersion {
+    V1 = 1,
+    V2 = 2
+  }

sdk/cosmosdb/cosmos/src/documents/PartitionKind.ts

@@ -0,0 +1,6 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT license.
+export enum PartitionKeyKind {
+    Hash = "Hash",
+    MultiHash = "MultiHash"
+}

sdk/cosmosdb/cosmos/src/documents/index.ts

@@ -10,6 +10,8 @@ export * from "./IndexingMode";
 export * from "./IndexingPolicy";
 export * from "./IndexKind";
 export * from "./PartitionKey";
+export * from "./PartitionKeyDefinitionVersion"
+export * from "./PartitionKind"
 export * from "./PartitionKeyDefinition";
 export * from "./PermissionMode";
 export * from "./TriggerOperation";

sdk/cosmosdb/cosmos/src/extractPartitionKey.ts

@@ -1,47 +1,67 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT license.
+import { AzureLogger, createClientLogger } from "@azure/logger";
 import { parsePath } from "./common";
-import { PartitionKey, PartitionKeyDefinition } from "./documents";
+import { NonePartitionKeyLiteral, NullPartitionKeyLiteral, PartitionKeyDefinition, PartitionKeyInternal, PrimitivePartitionKeyValue } from "./documents";
+
+const logger: AzureLogger = createClientLogger("extractPartitionKey");
 
 /**
+ * Function to extract PartitionKey based on {@link PartitionKeyDefinition}
+ * from an object.
+ * Retuns
+ * 1. {@link PartitionKeyInternal[]} - if extraction is successful.
+ * 2. {@link undefined} - if either {@link partitionKeyDefinition} is not well formed
+ * or an unsupported partitionkey type is encountered.
  * @hidden
  */
 export function extractPartitionKey(
   document: unknown,
-  partitionKeyDefinition: PartitionKeyDefinition
-): PartitionKey[] {
+  partitionKeyDefinition?: PartitionKeyDefinition
+): PartitionKeyInternal | undefined {
   if (
     partitionKeyDefinition &&
     partitionKeyDefinition.paths &&
     partitionKeyDefinition.paths.length > 0
   ) {
-    const partitionKey: PartitionKey[] = [];
+    const partitionKeys: PrimitivePartitionKeyValue[] = [];
     partitionKeyDefinition.paths.forEach((path: string) => {
-      const pathParts = parsePath(path);
+      const pathParts: string[] = parsePath(path);
       let obj = document;
       for (const part of pathParts) {
-        if (typeof obj === "object" && part in obj) {
+        if (typeof obj === "object" && obj !==null && part in obj) {
           obj = (obj as Record<string, unknown>)[part];
         } else {
           obj = undefined;
           break;
         }
       }
-      partitionKey.push(obj);
+      if(typeof obj === "string" || typeof obj === "number" || typeof obj === "boolean" ) {
+        partitionKeys.push(obj);
+      } else if(obj === NullPartitionKeyLiteral) {
+        partitionKeys.push(NullPartitionKeyLiteral);
+      } else if(obj === undefined || JSON.stringify(obj) === JSON.stringify(NonePartitionKeyLiteral)) {
+        if(partitionKeyDefinition.systemKey === true) {
+          return []
+        }
+        partitionKeys.push(NonePartitionKeyLiteral);
+      } else {
+        logger.warning('Unsupported PartitionKey found.');
+        return undefined;
+      }
     });
-    if (partitionKey.length === 1 && partitionKey[0] === undefined) {
-      return undefinedPartitionKey(partitionKeyDefinition);
-    }
-    return partitionKey;
+    return partitionKeys;
   }
+  logger.warning("Unexpected Partition Key Definition Found.");
+  return undefined;
 }
 /**
  * @hidden
  */
-export function undefinedPartitionKey(partitionKeyDefinition: PartitionKeyDefinition): unknown[] {
+export function undefinedPartitionKey(partitionKeyDefinition: PartitionKeyDefinition): PartitionKeyInternal {
   if (partitionKeyDefinition.systemKey === true) {
     return [];
   } else {
-    return [{}];
+    return partitionKeyDefinition.paths.map(() => NonePartitionKeyLiteral);
   }
 }

sdk/cosmosdb/cosmos/src/request/request.ts

@@ -3,7 +3,7 @@
 import { setAuthorizationHeader } from "../auth";
 import { Constants, HTTPMethod, jsonStringifyAndEscapeNonASCII, ResourceType } from "../common";
 import { CosmosClientOptions } from "../CosmosClientOptions";
-import { PartitionKey } from "../documents";
+import { PartitionKeyInternal } from "../documents";
 import { CosmosHeaders } from "../queryExecutionContext";
 import { FeedOptions, RequestOptions } from "./index";
 import { defaultLogger } from "../common/logger";
@@ -41,7 +41,7 @@ interface GetHeadersOptions {
   options: RequestOptions & FeedOptions;
   partitionKeyRangeId?: string;
   useMultipleWriteLocations?: boolean;
-  partitionKey?: PartitionKey;
+  partitionKey?: PartitionKeyInternal;
 }
 
 const JsonContentType = "application/json";
@@ -168,9 +168,6 @@ export async function getHeaders({
   }
 
   if (partitionKey !== undefined && !headers[Constants.HttpHeaders.PartitionKey]) {
-    if (partitionKey === null || !Array.isArray(partitionKey)) {
-      partitionKey = [partitionKey as string];
-    }
     headers[Constants.HttpHeaders.PartitionKey] = jsonStringifyAndEscapeNonASCII(partitionKey);
   }