From 9bbf1046d9516126d00acbcbd6796c11f1ce9a4e Mon Sep 17 00:00:00 2001
From: JasonYeMSFT <chuye@microsoft.com>
Date: Thu, 21 Sep 2023 10:11:40 -0700
Subject: [PATCH] Customize partition key validation rule for Gremlin DB
 (#2184)

---
 src/docdb/tree/DocDBDatabaseTreeItemBase.ts | 36 ++++++++++++---------
 src/graph/tree/GraphDatabaseTreeItem.ts     | 26 +++++++++++++++
 2 files changed, 47 insertions(+), 15 deletions(-)

diff --git a/src/docdb/tree/DocDBDatabaseTreeItemBase.ts b/src/docdb/tree/DocDBDatabaseTreeItemBase.ts
index 2297d6fc0..d8707a76c 100644
--- a/src/docdb/tree/DocDBDatabaseTreeItemBase.ts
+++ b/src/docdb/tree/DocDBDatabaseTreeItemBase.ts
@@ -77,16 +77,7 @@ export abstract class DocDBDatabaseTreeItemBase extends DocDBTreeItemBase<Contai
             id: containerName
         };
 
-        let partitionKey: string | undefined = await context.ui.showInputBox({
-            prompt: 'Enter the partition key for the collection, or leave blank for fixed size.',
-            stepName: 'partitionKeyForCollection',
-            validateInput: validatePartitionKey,
-            placeHolder: 'e.g. address/zipCode'
-        });
-
-        if (partitionKey && partitionKey.length && partitionKey[0] !== '/') {
-            partitionKey = '/' + partitionKey;
-        }
+        const partitionKey = await this.getNewPartitionKey(context);
         if (partitionKey) {
             containerDefinition.partitionKey = {
                 paths: [partitionKey]
@@ -113,13 +104,28 @@ export abstract class DocDBDatabaseTreeItemBase extends DocDBTreeItemBase<Contai
 
         return this.initChild(nonNullProp(container, 'resource'));
     }
-}
 
-function validatePartitionKey(key: string): string | undefined | null {
-    if (/[#?\\]/.test(key)) {
-        return "Cannot contain these characters: ?,#,\\, etc.";
+    protected async getNewPartitionKey(context: IActionContext): Promise<string | undefined> {
+        let partitionKey: string | undefined = await context.ui.showInputBox({
+            prompt: 'Enter the partition key for the collection, or leave blank for fixed size.',
+            stepName: 'partitionKeyForCollection',
+            validateInput: this.validatePartitionKey,
+            placeHolder: 'e.g. /address/zipCode'
+        });
+
+        if (partitionKey && partitionKey.length && partitionKey[0] !== '/') {
+            partitionKey = '/' + partitionKey;
+        }
+
+        return partitionKey;
+    }
+
+    protected validatePartitionKey(key: string): string | undefined {
+        if (/[#?\\]/.test(key)) {
+            return "Cannot contain these characters: ?,#,\\, etc.";
+        }
+        return undefined;
     }
-    return undefined;
 }
 
 function validateThroughput(isFixed: boolean, input: string): string | undefined | null {
diff --git a/src/graph/tree/GraphDatabaseTreeItem.ts b/src/graph/tree/GraphDatabaseTreeItem.ts
index 6965edaf3..1c5332700 100644
--- a/src/graph/tree/GraphDatabaseTreeItem.ts
+++ b/src/graph/tree/GraphDatabaseTreeItem.ts
@@ -4,6 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { ContainerDefinition, CosmosClient, Database, DatabaseDefinition, Resource } from '@azure/cosmos';
+import { IActionContext } from '@microsoft/vscode-azext-utils';
 import { DocDBDatabaseTreeItemBase } from '../../docdb/tree/DocDBDatabaseTreeItemBase';
 import { IGremlinEndpoint } from '../../vscode-cosmosdbgraph.api';
 import { getPossibleGremlinEndpoints } from '../gremlinEndpoints';
@@ -36,4 +37,29 @@ export class GraphDatabaseTreeItem extends DocDBDatabaseTreeItemBase {
         return client.database(this.id);
 
     }
+
+    protected override async getNewPartitionKey(context: IActionContext): Promise<string | undefined> {
+        let partitionKey: string | undefined = await context.ui.showInputBox({
+            prompt: 'Enter the partition key for the collection, or leave blank for fixed size.',
+            stepName: 'partitionKeyForCollection',
+            validateInput: this.validatePartitionKey,
+            placeHolder: 'e.g. /address'
+        });
+
+        if (partitionKey && partitionKey.length && partitionKey[0] !== '/') {
+            partitionKey = '/' + partitionKey;
+        }
+
+        return partitionKey;
+    }
+
+    protected validatePartitionKey(key: string): string | undefined {
+        if (/[#?\\]/.test(key)) {
+            return "Cannot contain these characters: ?,#,\\, etc.";
+        }
+        if (/.+\//.test(key)) {
+            return "Cannot be a nested path";
+        }
+        return undefined;
+    }
 }