Fixing schema designer and table designer profiles before init (#20753)

* Adding telemetry and prepare conn profile to table designer and schema designer.

* remove dup confirmTokenValidity call.

* Fixing ordering

* Fixing tests

* Add error handling and user feedback for Schema Designer initialization and connection preparation

- Introduced error messages for loading Schema Designer and retry options.
- Enhanced error handling in SchemaDesignerStateProvider and TableDesignerWebviewController.
- Updated locConstants for new error messages.

* Fixed

* Fixing rpc registration.

* Fixing lint issues.

* fixed name

* Adding more error handling and fixng message.

* Added loc

* Fixing error telemetry

* Made dialogs conditional.

---------

Co-authored-by: Aasim Khan <[email protected]>

Author: aasimkhan30

extensions/mssql/l10n/bundle.l10n.json

@@ -159,7 +159,7 @@
     "Generate Script": "Generate Script",
     "Publish": "Publish",
     "Preview Database Updates": "Preview Database Updates",
-    "Error loading designer": "Error loading designer",
+    "Error loading Table Designer": "Error loading Table Designer",
     "Severity": "Severity",
     "Script as Create": "Script as Create",
     "I have read the summary and understand the potential risks.": "I have read the summary and understand the potential risks.",
@@ -612,6 +612,7 @@
         "comment": ["{0} is the max length"]
     },
     "Loading Schema Designer": "Loading Schema Designer",
+    "Error loading Schema Designer": "Error loading Schema Designer",
     "Generating report, this might take a while...": "Generating report, this might take a while...",
     "{0} warnings/{0} is the number of warnings": {
         "message": "{0} warnings",

extensions/mssql/src/reactviews/common/errorDialog.tsx

@@ -0,0 +1,72 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import {
+    Button,
+    Dialog,
+    DialogActions,
+    DialogBody,
+    DialogContent,
+    DialogSurface,
+    DialogTitle,
+    makeStyles,
+    shorthands,
+} from "@fluentui/react-components";
+import { ErrorCircleRegular } from "@fluentui/react-icons";
+import React from "react";
+
+const useStyles = makeStyles({
+    dialogSurface: {
+        minWidth: "320px",
+        maxWidth: "420px",
+    },
+    title: {
+        display: "flex",
+        alignItems: "center",
+        columnGap: "8px",
+    },
+    icon: {
+        fontSize: "32px",
+    },
+    content: {
+        ...shorthands.marginBlock("16px", "0"),
+    },
+});
+
+export interface ErrorDialogProps {
+    open: boolean;
+    title: string;
+    message: string;
+    retryLabel: string;
+    onRetry: () => void;
+}
+
+export const ErrorDialog: React.FC<ErrorDialogProps> = ({
+    open,
+    title,
+    message,
+    retryLabel,
+    onRetry,
+}) => {
+    const classes = useStyles();
+    return (
+        <Dialog open={open} modalType="modal" inertTrapFocus>
+            <DialogSurface className={classes.dialogSurface}>
+                <DialogBody>
+                    <DialogTitle className={classes.title}>
+                        <ErrorCircleRegular className={classes.icon} />
+                        {title}
+                    </DialogTitle>
+                    <DialogContent className={classes.content}>{message}</DialogContent>
+                    <DialogActions>
+                        <Button appearance="primary" onClick={onRetry}>
+                            {retryLabel}
+                        </Button>
+                    </DialogActions>
+                </DialogBody>
+            </DialogSurface>
+        </Dialog>
+    );
+};

extensions/mssql/src/reactviews/common/locConstants.ts

@@ -139,7 +139,7 @@ export class LocConstants {
             generateScript: l10n.t("Generate Script"),
             publish: l10n.t("Publish"),
             previewDatabaseUpdates: l10n.t("Preview Database Updates"),
-            errorLoadingDesigner: l10n.t("Error loading designer"),
+            errorLoadingDesigner: l10n.t("Error loading Table Designer"),
             severity: l10n.t("Severity"),
             description: l10n.t("Description"),
             scriptAsCreate: l10n.t("Script as Create"),
@@ -876,6 +876,8 @@ export class LocConstants {
                     comment: ["{0} is the max length"],
                 }),
             loadingSchemaDesigner: l10n.t("Loading Schema Designer"),
+            errorLoadingSchemaDesigner: l10n.t("Error loading Schema Designer"),
+            retry: l10n.t("Retry"),
             generatingReport: l10n.t("Generating report, this might take a while..."),
             nWarnings: (warningCount: number) =>
                 l10n.t({

extensions/mssql/src/reactviews/common/utils.ts

@@ -32,6 +32,25 @@ export function formatString(str: string, ...args: any[]): string {
     return result;
 }
 
+/**
+ * Get the error message from an unknown error object
+ * @param error The error object
+ * @returns The error message
+ */
+export function getErrorMessage(error: unknown): string {
+    if (error instanceof Error) {
+        return error.message;
+    }
+    if (typeof error === "string") {
+        return error;
+    }
+    try {
+        return JSON.stringify(error);
+    } catch {
+        return "Unknown error";
+    }
+}
+
 /**
  * Get the css string representation of a ColorThemeKind
  * @param themeKind The ColorThemeKind to convert

extensions/mssql/src/reactviews/pages/SchemaDesigner/graph/SchemaDiagramFlow.tsx

@@ -82,12 +82,18 @@ export const SchemaDesignerFlow = () => {
 
     useEffect(() => {
         const intialize = async () => {
-            const { nodes, edges } = await context.initializeSchemaDesigner();
-            setSchemaNodes(nodes);
-            setRelationshipEdges(edges);
+            try {
+                const { nodes, edges } = await context.initializeSchemaDesigner();
+                setSchemaNodes(nodes);
+                setRelationshipEdges(edges);
+            } catch (error) {
+                context.log?.(`Failed to initialize schema designer: ${String(error)}`);
+                setSchemaNodes([]);
+                setRelationshipEdges([]);
+            }
         };
         void intialize();
-    }, []);
+    }, [context.initializationRequestId]);
 
     /**
      * Displays an error toast notification

extensions/mssql/src/reactviews/pages/SchemaDesigner/schemaDesignerPage.tsx

@@ -14,6 +14,7 @@ import { SchemaDesignerFindTableWidget } from "./schemaDesignerFindTables";
 import { makeStyles, Spinner } from "@fluentui/react-components";
 import { locConstants } from "../../common/locConstants";
 import { Panel, PanelGroup, PanelResizeHandle } from "react-resizable-panels";
+import { ErrorDialog } from "../../common/errorDialog";
 
 const useStyles = makeStyles({
     resizeHandle: {
@@ -43,7 +44,16 @@ export const SchemaDesignerPage = () => {
                     <PanelResizeHandle className={classes.resizeHandle} />
                     <SchemaDesignerDefinitionsPanel />
                 </PanelGroup>
-                {!context.isInitialized && <LoadingOverlay />}
+                {!context.isInitialized && !context.initializationError && <LoadingOverlay />}
+                {context?.initializationError && (
+                    <ErrorDialog
+                        open={!!context?.initializationError}
+                        title={locConstants.schemaDesigner.errorLoadingSchemaDesigner}
+                        message={context?.initializationError ?? ""}
+                        retryLabel={locConstants.schemaDesigner.retry}
+                        onRetry={context?.triggerInitialization}
+                    />
+                )}
             </MainLayout>
         </>
     );

extensions/mssql/src/reactviews/pages/SchemaDesigner/schemaDesignerStateProvider.tsx

@@ -6,7 +6,7 @@
 import { createContext, useEffect, useState } from "react";
 import { SchemaDesigner } from "../../../sharedInterfaces/schemaDesigner";
 import { useVscodeWebview } from "../../common/vscodeWebviewProvider";
-import { getCoreRPCs } from "../../common/utils";
+import { getCoreRPCs, getErrorMessage } from "../../common/utils";
 import { WebviewRpc } from "../../common/rpc";
 
 import { Edge, MarkerType, Node, ReactFlowJsonObject, useReactFlow } from "@xyflow/react";
@@ -48,6 +48,9 @@ export interface SchemaDesignerContextProps
     resetUndoRedoState: () => void;
     resetView: () => void;
     isInitialized: boolean;
+    initializationError?: string;
+    initializationRequestId: number;
+    triggerInitialization: () => void;
     renderOnlyVisibleTables: boolean;
     setRenderOnlyVisibleTables: (value: boolean) => void;
     isExporting: boolean;
@@ -79,6 +82,8 @@ const SchemaDesignerStateProvider: React.FC<SchemaDesignerProviderProps> = ({ ch
     const [schemaNames, setSchemaNames] = useState<string[]>([]);
     const reactFlow = useReactFlow();
     const [isInitialized, setIsInitialized] = useState(false);
+    const [initializationError, setInitializationError] = useState<string | undefined>(undefined);
+    const [initializationRequestId, setInitializationRequestId] = useState(0);
     const [findTableText, setFindTableText] = useState<string>("");
     const [renderOnlyVisibleTables, setRenderOnlyVisibleTables] = useState<boolean>(true);
     const [isExporting, setIsExporting] = useState<boolean>(false);
@@ -132,29 +137,44 @@ const SchemaDesignerStateProvider: React.FC<SchemaDesignerProviderProps> = ({ ch
     }, []);
 
     const initializeSchemaDesigner = async () => {
-        const model = await extensionRpc.sendRequest(
-            SchemaDesigner.InitializeSchemaDesignerRequest.type,
-        );
+        try {
+            setIsInitialized(false);
+            setInitializationError(undefined);
+            const model = await extensionRpc.sendRequest(
+                SchemaDesigner.InitializeSchemaDesignerRequest.type,
+            );
 
-        const { nodes, edges } = flowUtils.generateSchemaDesignerFlowComponents(model.schema);
+            const { nodes, edges } = flowUtils.generateSchemaDesignerFlowComponents(model.schema);
 
-        setDatatypes(model.dataTypes);
-        setSchemaNames(model.schemaNames);
-        setIsInitialized(true);
+            setDatatypes(model.dataTypes);
+            setSchemaNames(model.schemaNames);
+            setIsInitialized(true);
 
-        setTimeout(() => {
-            stateStack.setInitialState(
-                reactFlow.toObject() as ReactFlowJsonObject<
-                    Node<SchemaDesigner.Table>,
-                    Edge<SchemaDesigner.ForeignKey>
-                >,
-            );
-        });
+            setTimeout(() => {
+                stateStack.setInitialState(
+                    reactFlow.toObject() as ReactFlowJsonObject<
+                        Node<SchemaDesigner.Table>,
+                        Edge<SchemaDesigner.ForeignKey>
+                    >,
+                );
+            });
 
-        return {
-            nodes,
-            edges,
-        };
+            return {
+                nodes,
+                edges,
+            };
+        } catch (error) {
+            const errorMessage = getErrorMessage(error);
+            setInitializationError(errorMessage);
+            setIsInitialized(false);
+            throw error;
+        }
+    };
+
+    const triggerInitialization = () => {
+        setInitializationError(undefined);
+        setIsInitialized(false);
+        setInitializationRequestId((id) => id + 1);
     };
 
     // Get the script from the server
@@ -460,6 +480,9 @@ const SchemaDesignerStateProvider: React.FC<SchemaDesignerProviderProps> = ({ ch
                 setFindTableText,
                 getDefinition,
                 initializeSchemaDesigner,
+                initializationError,
+                initializationRequestId,
+                triggerInitialization,
                 saveAsFile,
                 getReport,
                 openInEditor,

extensions/mssql/src/reactviews/pages/TableDesigner/tableDesignerPage.tsx

@@ -3,16 +3,16 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { Button, Spinner, makeStyles, shorthands } from "@fluentui/react-components";
+import { Spinner, makeStyles, shorthands } from "@fluentui/react-components";
 import { useContext, useEffect, useRef } from "react";
 import * as designer from "../../../sharedInterfaces/tableDesigner";
 import { TableDesignerContext } from "./tableDesignerStateProvider";
-import { ErrorCircleRegular } from "@fluentui/react-icons";
 import { DesignerPageRibbon } from "./designerPageRibbon";
 import { DesignerMainPane } from "./designerMainPane";
 import { DesignerPropertiesPane } from "./designerPropertiesPane";
 import { DesignerResultPane } from "./designerResultPane";
 import { locConstants } from "../../common/locConstants";
+import { ErrorDialog } from "../../common/errorDialog";
 import {
     ImperativePanelHandle,
     Panel,
@@ -37,13 +37,6 @@ const useStyles = makeStyles({
         width: "100%",
         flexDirection: "column",
     },
-    errorIcon: {
-        fontSize: "100px",
-        opacity: 0.5,
-    },
-    retryButton: {
-        marginTop: "10px",
-    },
     resultPaneHandle: {
         position: "absolute",
         top: "0",
@@ -140,6 +133,8 @@ export const TableDesigner = () => {
         }
     }, [context?.state.propertiesPaneData, context?.propertiesPaneResizeInfo.isMaximized]);
 
+    const isErrorState = tableDesignerState.apiState?.initializeState === designer.LoadState.Error;
+
     return (
         <div className={classes.root}>
             {tableDesignerState.apiState?.initializeState === designer.LoadState.Loading && (
@@ -150,14 +145,14 @@ export const TableDesigner = () => {
                     />
                 </div>
             )}
-            {tableDesignerState.apiState?.initializeState === designer.LoadState.Error && (
-                <div className={classes.pageContext}>
-                    <ErrorCircleRegular className={classes.errorIcon} />
-                    <div>{locConstants.tableDesigner.errorLoadingDesigner}</div>
-                    <Button className={classes.retryButton}>
-                        {locConstants.tableDesigner.retry}
-                    </Button>
-                </div>
+            {isErrorState && (
+                <ErrorDialog
+                    open={isErrorState}
+                    title={locConstants.tableDesigner.errorLoadingDesigner}
+                    message={tableDesignerState?.initializationError ?? ""}
+                    retryLabel={locConstants.tableDesigner.retry}
+                    onRetry={() => context.initializeTableDesigner()}
+                />
             )}
             {tableDesignerState.apiState?.initializeState === designer.LoadState.Loaded && (
                 <div className={classes.mainContent}>

extensions/mssql/src/schemaDesigner/schemaDesignerWebviewController.ts

@@ -87,6 +87,8 @@ export class SchemaDesignerWebviewController extends ReactWebviewPanelController
                 TelemetryActions.Initialize,
                 undefined,
                 undefined,
+                undefined,
+                true, // include callstack in telemetry
             );
             try {
                 let sessionResponse: SchemaDesigner.CreateSessionResponse;

extensions/mssql/src/schemaDesigner/schemaDesignerWebviewManager.ts

@@ -12,6 +12,7 @@ import MainController from "../controllers/mainController";
 import * as LocConstants from "../constants/locConstants";
 import { TelemetryViews, TelemetryActions } from "../sharedInterfaces/telemetry";
 import { sendActionEvent } from "../telemetry/telemetry";
+import { IConnectionProfile } from "../models/interfaces";
 
 export class SchemaDesignerWebviewManager {
     private static instance: SchemaDesignerWebviewManager;
@@ -57,13 +58,15 @@ export class SchemaDesignerWebviewManager {
         let connectionString: string | undefined;
         let azureAccountToken: string | undefined;
         if (treeNode) {
-            const connectionInfo = treeNode.connectionProfile;
+            let connectionInfo = treeNode.connectionProfile;
+            connectionInfo = (await mainController.connectionManager.prepareConnectionInfo(
+                connectionInfo,
+            )) as IConnectionProfile;
             connectionInfo.database = databaseName;
 
             const connectionDetails =
                 await mainController.connectionManager.createConnectionDetails(connectionInfo);
 
-            await mainController.connectionManager.confirmEntraTokenValidity(connectionInfo);
             treeNode.updateConnectionProfile(connectionInfo);
 
             connectionString = await mainController.connectionManager.getConnectionString(