Add KeySuffix and CompatibilityMode to CosmosDbPartitionedStorage (#1468)

* add compatibility mode and backward compat

* remove semaphore and use promises, fix tests

* Update libraries/botbuilder-azure/src/cosmosDbPartitionedStorage.ts

Co-Authored-By: Michael Richardson <40401643+mdrichardson@users.noreply.github.com>

Author: 2 people

libraries/botbuilder-azure/src/cosmosDbKeyEscape.ts

@@ -8,8 +8,10 @@
 var crypto = require('crypto');
 
 export namespace CosmosDbKeyEscape {
-    // Per the CosmosDB Docs, there is a max key length of 255. 
-    // https://docs.microsoft.com/en-us/azure/cosmos-db/faq#table
+    // Older libraries had a max key length of 255.
+    // The limit is now 1023. In this library, 255 remains the default for backwards compat.
+    // To override this behavior, and use the longer limit, set cosmosDbPartitionedStorageOptions.compatibilityMode to false.
+    // https://docs.microsoft.com/en-us/azure/cosmos-db/concepts-limits#per-item-limits
     const maxKeyLength = 255;
     const illegalKeys: string[] = ['\\', '?', '/', '#', '\t', '\n', '\r', '*'];
     const illegalKeyCharacterReplacementMap: Map<string, string> =
@@ -27,8 +29,11 @@ export namespace CosmosDbKeyEscape {
      * The following characters are restricted and cannot be used in the Id property: '/', '\', '?', '#'
      * More information at https://docs.microsoft.com/en-us/dotnet/api/microsoft.azure.documents.resource.id?view=azure-dotnet#remarks
       * @param key The provided key to be escaped.
+      * @param keySuffix The string to add a the end of all RowKeys.
+      * @param compatibilityMode True if keys should be truncated in order to support previous CosmosDb
+      * max key length of 255.  This behavior can be overridden by setting cosmosDbPartitionedStorageOptions.compatibilityMode to false.
       */
-    export function escapeKey(key: string): string {
+    export function escapeKey(key: string, keySuffix?: string, compatibilityMode?: boolean): string {
         if (!key) {
             throw new Error('The \'key\' parameter is required.');
         }
@@ -39,7 +44,7 @@ export namespace CosmosDbKeyEscape {
 
         // If there are no illegal characters return immediately and avoid any further processing/allocations
         if (firstIllegalCharIndex === -1) {
-            return truncateKey(key);
+            return truncateKey(`${key}${keySuffix || ''}`, compatibilityMode);
         }
 
         let sanitizedKey = keySplitted.reduce(
@@ -48,10 +53,14 @@ export namespace CosmosDbKeyEscape {
             ''
         );
 
-        return truncateKey(sanitizedKey);
+        return truncateKey(`${sanitizedKey}${keySuffix || ''}`, compatibilityMode);
     }
 
-    function truncateKey(key: string): string {
+    function truncateKey(key: string, truncateKeysForCompatibility?: boolean): string {
+        if (truncateKeysForCompatibility === false) {
+            return key;
+        }
+
         if (key.length > maxKeyLength) {
             const hash = crypto.createHash('sha256');
             hash.update(key);

libraries/botbuilder-azure/src/cosmosDbPartitionedStorage.ts

@@ -9,9 +9,9 @@
 import { Storage, StoreItems } from 'botbuilder';
 import { Container, CosmosClient, CosmosClientOptions } from '@azure/cosmos';
 import { CosmosDbKeyEscape } from './cosmosDbKeyEscape';
-import * as semaphore from 'semaphore';
+import { DoOnce } from './doOnce';
 
-const _semaphore: semaphore.Semaphore = semaphore(1);
+const _doOnce: DoOnce<Container> = new DoOnce<Container>();
 
 /**
  * Cosmos DB Partitioned Storage Options.
@@ -41,6 +41,26 @@ export interface CosmosDbPartitionedStorageOptions {
      * The throughput set when creating the Container. Defaults to 400.
      */
     containerThroughput?: number;
+    /**
+     * The suffix to be added to every key. See cosmosDbKeyEscape.escapeKey
+     * 
+     * Note: compatibilityMode must be set to 'false' to use a KeySuffix.
+     * When KeySuffix is used, keys will NOT be truncated but an exception will
+     * be thrown if the key length is longer than allowed by CosmosDb.
+     * 
+     * The keySuffix must contain only valid ComosDb key characters.
+     * (e.g. not: '\\', '?', '/', '#', '*')
+     */
+    keySuffix?: string;
+    /**
+    * Early version of CosmosDb had a max key length of 255.  Keys longer than
+    * this were truncated in cosmosDbKeyEscape.escapeKey.  This remains the default
+    * behavior of cosmosDbPartitionedStorage, but can be overridden by setting
+    * compatibilityMode to false. 
+    * 
+    * compatibilityMode cannot be true if keySuffix is used.
+    */
+    compatibilityMode?: boolean;
 }
 
 /**
@@ -93,6 +113,7 @@ export class CosmosDbPartitionedStorage implements Storage {
     private container: Container;
     private readonly cosmosDbStorageOptions: CosmosDbPartitionedStorageOptions;
     private client: CosmosClient;
+    private compatabilityModePartitionKey: boolean = false;
 
     /**
      * Initializes a new instance of the <see cref="CosmosDbPartitionedStorage"/> class.
@@ -106,6 +127,22 @@ export class CosmosDbPartitionedStorage implements Storage {
         if (!cosmosDbStorageOptions.authKey) { throw new ReferenceError('authKey for CosmosDB is required.'); }
         if (!cosmosDbStorageOptions.databaseId) { throw new ReferenceError('databaseId is for CosmosDB required.'); }
         if (!cosmosDbStorageOptions.containerId) { throw new ReferenceError('containerId for CosmosDB is required.'); }
+        // In order to support collections previously restricted to max key length of 255, we default
+        // compatabilityMode to 'true'.  No compatibilityMode is opt-in only.
+        if (typeof cosmosDbStorageOptions.compatibilityMode === "undefined") {
+            cosmosDbStorageOptions.compatibilityMode = true;
+        }
+        if (cosmosDbStorageOptions.keySuffix) {
+            if (cosmosDbStorageOptions.compatibilityMode){
+                throw new ReferenceError('compatibilityMode cannot be true while using a keySuffix.');
+            }
+            // In order to reduce key complexity, we do not allow invalid characters in a KeySuffix
+            // If the keySuffix has invalid characters, the escaped key will not match
+            const suffixEscaped = CosmosDbKeyEscape.escapeKey(cosmosDbStorageOptions.keySuffix);
+            if (cosmosDbStorageOptions.keySuffix !== suffixEscaped) {
+                throw new ReferenceError(`Cannot use invalid Row Key characters: ${cosmosDbStorageOptions.keySuffix} in keySuffix`);
+            }
+        }
 
         this.cosmosDbStorageOptions = cosmosDbStorageOptions;
     }
@@ -120,9 +157,9 @@ export class CosmosDbPartitionedStorage implements Storage {
 
         await Promise.all(keys.map(async (k: string): Promise<void> => {
             try {
-                const escapedKey = CosmosDbKeyEscape.escapeKey(k);
+                const escapedKey = CosmosDbKeyEscape.escapeKey(k, this.cosmosDbStorageOptions.keySuffix, this.cosmosDbStorageOptions.compatibilityMode);
 
-                const readItemResponse = await this.container.item(escapedKey, escapedKey).read<DocumentStoreItem>();
+                const readItemResponse = await this.container.item(escapedKey, this.getPartitionKey(escapedKey)).read<DocumentStoreItem>();
                 const documentStoreItem = readItemResponse.resource;
                 if (documentStoreItem) {
                     storeItems[documentStoreItem.realId] = documentStoreItem.document;
@@ -159,7 +196,7 @@ export class CosmosDbPartitionedStorage implements Storage {
             // The ETag information is updated as an _etag attribute in the document metadata.
             delete changesCopy.eTag;
             const documentChange = new DocumentStoreItem({
-                id: CosmosDbKeyEscape.escapeKey(k),
+                id: CosmosDbKeyEscape.escapeKey(k, this.cosmosDbStorageOptions.keySuffix, this.cosmosDbStorageOptions.compatibilityMode),
                 realId: k,
                 document: changesCopy
             });
@@ -180,10 +217,10 @@ export class CosmosDbPartitionedStorage implements Storage {
         await this.initialize();
 
         await Promise.all(keys.map(async (k: string): Promise<void> => {
-            const escapedKey = CosmosDbKeyEscape.escapeKey(k);
+            const escapedKey = CosmosDbKeyEscape.escapeKey(k, this.cosmosDbStorageOptions.keySuffix, this.cosmosDbStorageOptions.compatibilityMode);
             try {
                 await this.container
-                    .item(escapedKey, escapedKey)
+                    .item(escapedKey, this.getPartitionKey(escapedKey))
                     .delete();
             } catch (err) {
                 // If trying to delete a document that doesn't exist, do nothing. Otherwise, throw
@@ -200,35 +237,63 @@ export class CosmosDbPartitionedStorage implements Storage {
      */
     public async initialize(): Promise<void> {
         if (!this.container) {
-
             if (!this.client) {
                 this.client = new CosmosClient({
                     endpoint: this.cosmosDbStorageOptions.cosmosDbEndpoint,
                     key: this.cosmosDbStorageOptions.authKey,
                     ...this.cosmosDbStorageOptions.cosmosClientOptions,
                 });
             }
+            this.container = await _doOnce.waitFor(async ()=> await this.getOrCreateContainer());
+        }
+    }
+    
+    private async getOrCreateContainer(): Promise<Container> {
+        let createIfNotExists = !this.cosmosDbStorageOptions.compatibilityMode;
+        let container;
+        if(this.cosmosDbStorageOptions.compatibilityMode) {
+            try {
+                container = await this.client
+                                    .database(this.cosmosDbStorageOptions.databaseId)
+                                    .container(this.cosmosDbStorageOptions.containerId);
+                const partitionKeyPath = await container.readPartitionKeyDefinition();
+                const paths = partitionKeyPath.resource.paths;
+                if(paths) {
+                    // Containers created with CosmosDbStorage had no partition key set, so the default was '/_partitionKey'.
+                    if (paths.indexOf('/_partitionKey') !== -1) {
+                        this.compatabilityModePartitionKey = true;
+                    }
+                    else if (paths.indexOf(DocumentStoreItem.partitionKeyPath) === -1) {
+                        // We are not supporting custom Partition Key Paths.
+                        new Error(`Custom Partition Key Paths are not supported. ${this.cosmosDbStorageOptions.containerId} has a custom Partition Key Path of ${paths[0]}.`);
+                    }
+                } else {
+                    this.compatabilityModePartitionKey = true;
+                }
+                return container;
+            } catch (err) {
+                createIfNotExists = true;
+            }
+        }
 
-            if (!this.container) {
-                this.container = await new Promise((resolve: Function): void => {
-                    _semaphore.take(async (): Promise<void> => {
-                        const result = await this.client
-                            .database(this.cosmosDbStorageOptions.databaseId)
-                            .containers.createIfNotExists({
-                                id: this.cosmosDbStorageOptions.containerId,
-                                partitionKey: {
-                                    paths: [DocumentStoreItem.partitionKeyPath]
-                                },
-                                throughput: this.cosmosDbStorageOptions.containerThroughput
-                            });
-                        _semaphore.leave();
-                        resolve(result.container);
-                    });
+        if (createIfNotExists) {
+            const result = await this.client
+                .database(this.cosmosDbStorageOptions.databaseId)
+                .containers.createIfNotExists({
+                    id: this.cosmosDbStorageOptions.containerId,
+                    partitionKey: {
+                        paths: [DocumentStoreItem.partitionKeyPath]
+                    },
+                    throughput: this.cosmosDbStorageOptions.containerThroughput
                 });
-            }
+            return result.container;
         }
     }
 
+    private getPartitionKey(key) {
+        return this.compatabilityModePartitionKey ? undefined : key;
+    }
+
     /**
      * The Cosmos JS SDK doesn't return very descriptive errors and not all errors contain a body. 
      * This provides more detailed errors and err['message'] prevents ReferenceError
@@ -240,4 +305,4 @@ export class CosmosDbPartitionedStorage implements Storage {
         err['message'] = `[${ prependedMessage }] ${ err['message'] }`;
         throw err;
     }
-}
+}

libraries/botbuilder-azure/src/doOnce.ts

@@ -0,0 +1,19 @@
+/**
+ * @module botbuilder-azure
+ */
+/**
+ * Copyright (c) Microsoft Corporation. All rights reserved.
+ * Licensed under the MIT License.
+ */
+
+export class DoOnce<T> {
+    private task: Promise<T>;
+
+    public waitFor(fn: () => Promise<T>): Promise<T> {
+        if (!this.task) {
+            this.task = fn();
+        }
+
+        return this.task;
+    }
+}