From 182cb2ef22b4869e9c8ee7e2f9e440959dc709ae Mon Sep 17 00:00:00 2001
From: Dmitrii Shilov <dshilov@microsoft.com>
Date: Tue, 23 Sep 2025 18:19:39 +0200
Subject: [PATCH] feat: Enhance restore container functionality with error
 handling and UI updates

---
 src/Contracts/DataExplorerMessagesContract.ts |  4 ++++
 src/Contracts/FabricMessageTypes.ts           |  1 +
 src/Explorer/ContextMenuButtonFactory.tsx     | 14 +++++++++++++-
 src/Platform/Fabric/FabricUtil.ts             |  6 ++++++
 src/Platform/Hosted/extractFeatures.ts        |  2 ++
 5 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/src/Contracts/DataExplorerMessagesContract.ts b/src/Contracts/DataExplorerMessagesContract.ts
index c017bffa8..bcb946449 100644
--- a/src/Contracts/DataExplorerMessagesContract.ts
+++ b/src/Contracts/DataExplorerMessagesContract.ts
@@ -22,6 +22,10 @@ export type DataExploreMessageV3 =
   | {
       type: FabricMessageTypes.OpenSettings;
       settingsId: string;
+    }
+  | {
+      type: FabricMessageTypes.RestoreContainer;
+      params: [];
     };
 export interface GetCosmosTokenMessageOptions {
   verb: "connect" | "delete" | "get" | "head" | "options" | "patch" | "post" | "put" | "trace";
diff --git a/src/Contracts/FabricMessageTypes.ts b/src/Contracts/FabricMessageTypes.ts
index 02871ca47..3b5a946ce 100644
--- a/src/Contracts/FabricMessageTypes.ts
+++ b/src/Contracts/FabricMessageTypes.ts
@@ -7,6 +7,7 @@ export enum FabricMessageTypes {
   GetAccessToken = "GetAccessToken",
   Ready = "Ready",
   OpenSettings = "OpenSettings",
+  RestoreContainer = "RestoreContainer",
 }
 
 export interface AuthorizationToken {
diff --git a/src/Explorer/ContextMenuButtonFactory.tsx b/src/Explorer/ContextMenuButtonFactory.tsx
index 3f6a795ee..57884c3bb 100644
--- a/src/Explorer/ContextMenuButtonFactory.tsx
+++ b/src/Explorer/ContextMenuButtonFactory.tsx
@@ -7,7 +7,7 @@ import {
   AddGlobalSecondaryIndexPanelProps,
 } from "Explorer/Panes/AddGlobalSecondaryIndexPanel/AddGlobalSecondaryIndexPanel";
 import { useDatabases } from "Explorer/useDatabases";
-import { isFabric, isFabricNative } from "Platform/Fabric/FabricUtil";
+import { isFabric, isFabricNative, openRestoreContainerDialog } from "Platform/Fabric/FabricUtil";
 import { Action } from "Shared/Telemetry/TelemetryConstants";
 import { traceOpen } from "Shared/Telemetry/TelemetryProcessor";
 import { ReactTabKind, useTabs } from "hooks/useTabs";
@@ -35,6 +35,7 @@ import StoredProcedure from "./Tree/StoredProcedure";
 import Trigger from "./Tree/Trigger";
 import UserDefinedFunction from "./Tree/UserDefinedFunction";
 import { useSelectedNode } from "./useSelectedNode";
+import { extractFeatures } from "../Platform/Hosted/extractFeatures";
 
 export interface CollectionContextMenuButtonParams {
   databaseId: string;
@@ -60,6 +61,17 @@ export const createDatabaseContextMenu = (container: Explorer, databaseId: strin
     },
   ];
 
+  if (isFabric() && !userContext.fabricContext?.isReadOnly) {
+    const features = extractFeatures();
+    if (features?.enableRestoreContainer) {
+      items.push({
+        iconSrc: AddCollectionIcon,
+        onClick: () => openRestoreContainerDialog(),
+        label: `Restore ${getCollectionName()}`,
+      });
+    }
+  }
+
   if (!isFabricNative() && (userContext.apiType !== "Tables" || userContext.features.enableSDKoperations)) {
     items.push({
       iconSrc: DeleteDatabaseIcon,
diff --git a/src/Platform/Fabric/FabricUtil.ts b/src/Platform/Fabric/FabricUtil.ts
index c2332e4bc..ce5f4f0ab 100644
--- a/src/Platform/Fabric/FabricUtil.ts
+++ b/src/Platform/Fabric/FabricUtil.ts
@@ -101,6 +101,12 @@ const requestAndStoreAccessToken = async (): Promise<void> => {
   });
 };
 
+export const openRestoreContainerDialog = (): void => {
+  if (configContext.platform === Platform.Fabric) {
+    sendCachedDataMessage(FabricMessageTypes.RestoreContainer, []);
+  }
+};
+
 /**
  * Check token validity and schedule a refresh if necessary
  * @param tokenTimestamp
diff --git a/src/Platform/Hosted/extractFeatures.ts b/src/Platform/Hosted/extractFeatures.ts
index 685930b43..60b8b4db1 100644
--- a/src/Platform/Hosted/extractFeatures.ts
+++ b/src/Platform/Hosted/extractFeatures.ts
@@ -40,6 +40,7 @@ export type Features = {
   readonly enablePriorityBasedExecution: boolean;
   readonly disableConnectionStringLogin: boolean;
   readonly enableCloudShell: boolean;
+  readonly enableRestoreContainer: boolean; // only for Fabric
 
   // can be set via both flight and feature flag
   autoscaleDefault: boolean;
@@ -111,6 +112,7 @@ export function extractFeatures(given = new URLSearchParams(window.location.sear
     copilotChatFixedMonacoEditorHeight: "true" === get("copilotchatfixedmonacoeditorheight"),
     enablePriorityBasedExecution: "true" === get("enableprioritybasedexecution"),
     disableConnectionStringLogin: "true" === get("disableconnectionstringlogin"),
+    enableRestoreContainer: "true" === get("enablerestorecontainer"),
     enableCloudShell: true,
   };
 }