Fixing schema designer filtering and auto layout behavior. (#19475)

Author: aasimkhan30

localization/l10n/bundle.l10n.json

@@ -360,6 +360,12 @@
   "Auto Arrange Confirmation": "Auto Arrange Confirmation",
   "Auto Arrange will automatically reposition all diagram elements based on optimal layout algorithms. Any custom positioning you've created will be lost. Do you want to proceed with auto-arranging your schema diagram?": "Auto Arrange will automatically reposition all diagram elements based on optimal layout algorithms. Any custom positioning you've created will be lost. Do you want to proceed with auto-arranging your schema diagram?",
   "Filter": "Filter",
+  "Filter ({0})/{0} is the number of selected tables": {
+    "message": "Filter ({0})",
+    "comment": [
+      "{0} is the number of selected tables"
+    ]
+  },
   "Refresh": "Refresh",
   "Publish Changes": "Publish Changes",
   "Edit Table": "Edit Table",
@@ -532,6 +538,7 @@
   "Are you sure you want to delete the selected items?": "Are you sure you want to delete the selected items?",
   "Undo": "Undo",
   "Redo": "Redo",
+  "Search tables...": "Search tables...",
   "To compare two schemas, first select a source schema and target schema, then press compare.": "To compare two schemas, first select a source schema and target schema, then press compare.",
   "No schema differences were found.": "No schema differences were found.",
   "Initializing comparison, this might take a while...": "Initializing comparison, this might take a while...",

localization/xliff/vscode-mssql.xlf

@@ -850,6 +850,10 @@
     <trans-unit id="++CODE++638e249f4a15ebb84957da130701d138a6e06d88dceeaa2e6dcd91db70cc1381">
       <source xml:lang="en">Filter</source>
     </trans-unit>
+    <trans-unit id="++CODE++a76057f27f2880792cf681fd60ca7b9043463252597c22b60eb5beca29292017">
+      <source xml:lang="en">Filter ({0})</source>
+      <note>{0} is the number of selected tables</note>
+    </trans-unit>
     <trans-unit id="++CODE++3905257a1ca17a05ee44a95330afe24483f05af738cc05090d1c2522910a5671">
       <source xml:lang="en">Filter Azure subscriptions</source>
     </trans-unit>
@@ -1643,6 +1647,9 @@
     <trans-unit id="++CODE++6801207d5994165123a10ab767fd9439130fae500f8c407877ae5070843b906d">
       <source xml:lang="en">Search settings...</source>
     </trans-unit>
+    <trans-unit id="++CODE++771f5d2c0599dcf4194e3fdddd9a5212d6392f4703d29971575bb455b5287a84">
+      <source xml:lang="en">Search tables...</source>
+    </trans-unit>
     <trans-unit id="++CODE++7f55382219f0202c1b4f56deb099e2fedcec87ee73fe2edb2105df5447323bc6">
       <source xml:lang="en">Search...</source>
     </trans-unit>

src/reactviews/common/locConstants.ts

@@ -386,7 +386,17 @@ export class LocConstants {
             autoArrangeConfirmationContent: l10n.t(
                 "Auto Arrange will automatically reposition all diagram elements based on optimal layout algorithms. Any custom positioning you've created will be lost. Do you want to proceed with auto-arranging your schema diagram?",
             ),
-            filter: l10n.t("Filter"),
+            filter: (selectedTablesCount: number) => {
+                if (selectedTablesCount === 0) {
+                    return l10n.t("Filter");
+                } else {
+                    return l10n.t({
+                        message: "Filter ({0})",
+                        args: [selectedTablesCount],
+                        comment: ["{0} is the number of selected tables"],
+                    });
+                }
+            },
             clearFilter: l10n.t("Clear All"),
             applyFilter: l10n.t("Apply"),
             refresh: l10n.t("Refresh"),
@@ -586,6 +596,7 @@ export class LocConstants {
             ),
             undo: l10n.t("Undo"),
             redo: l10n.t("Redo"),
+            searchTables: l10n.t("Search tables..."),
         };
     }
 

src/reactviews/pages/SchemaDesigner/schemaDesignerStateProvider.tsx

@@ -9,7 +9,7 @@ import { useVscodeWebview } from "../../common/vscodeWebviewProvider";
 import { getCoreRPCs } from "../../common/utils";
 import { WebviewRpc } from "../../common/rpc";
 
-import { Edge, Node, ReactFlowJsonObject, useReactFlow } from "@xyflow/react";
+import { Edge, MarkerType, Node, ReactFlowJsonObject, useReactFlow } from "@xyflow/react";
 import { flowUtils, foreignKeyUtils } from "./schemaDesignerUtils";
 import eventBus from "./schemaDesignerEvents";
 import { UndoRedoStack } from "../../common/undoRedoStack";
@@ -48,6 +48,7 @@ export interface SchemaDesignerContextProps
     }>;
     closeDesigner: () => void;
     resetUndoRedoState: () => void;
+    resetView: () => void;
     isInitialized: boolean;
 }
 
@@ -215,45 +216,30 @@ const SchemaDesignerStateProvider: React.FC<SchemaDesignerProviderProps> = ({ ch
      * Adds a new table to the flow
      */
     const addTable = async (table: SchemaDesigner.Table) => {
-        const newReactFlowNode: Node<SchemaDesigner.Table> = {
-            id: table.id,
-            type: "tableNode",
-            data: { ...table },
-            position: { x: 0, y: 0 },
-        };
+        const existingNodes = reactFlow.getNodes() as Node<SchemaDesigner.Table>[];
+        const existingEdges = reactFlow.getEdges() as Edge<SchemaDesigner.ForeignKey>[];
 
-        const schemaModel = flowUtils.extractSchemaModel(
-            reactFlow.getNodes() as Node<SchemaDesigner.Table>[],
-            reactFlow.getEdges() as Edge<SchemaDesigner.ForeignKey>[],
-        );
+        const schemaModel = flowUtils.extractSchemaModel(existingNodes, existingEdges);
+
+        schemaModel.tables.push(table);
 
-        schemaModel.tables.push(newReactFlowNode.data);
         const updatedPositions = flowUtils.generateSchemaDesignerFlowComponents(schemaModel);
 
-        const nodeWithPosition = updatedPositions.nodes.find(
-            (node) => node.id === newReactFlowNode.id,
-        );
+        const nodeWithPosition = updatedPositions.nodes.find((node) => node.id === table.id);
 
         const edgesForNewTable = updatedPositions.edges.filter(
-            (edge) => edge.source === newReactFlowNode.id || edge.target === newReactFlowNode.id,
+            (edge) => edge.source === table.id || edge.target === table.id,
         );
 
         if (nodeWithPosition) {
-            nodeWithPosition.selected = true;
-            reactFlow.addNodes(nodeWithPosition);
-            reactFlow.addEdges(edgesForNewTable);
-            // update all affected positions
-            for (const updatedNode of updatedPositions.nodes) {
-                if (updatedNode.id === nodeWithPosition.id) continue;
-                reactFlow.updateNode(updatedNode.id, {
-                    position: updatedNode.position,
-                    data: updatedNode.data,
-                });
-            }
+            existingNodes.push(nodeWithPosition);
+            existingEdges.push(...edgesForNewTable);
+
+            reactFlow.setNodes(existingNodes);
+            reactFlow.setEdges(existingEdges);
             requestAnimationFrame(async () => {
                 setCenter(nodeWithPosition.id, true);
             });
-
             eventBus.emit("getScript");
             return true;
         }
@@ -264,68 +250,71 @@ const SchemaDesignerStateProvider: React.FC<SchemaDesignerProviderProps> = ({ ch
     /**
      * Updates a table in the flow
      */
-    const updateTable = async (table: SchemaDesigner.Table) => {
-        const schemaModel = flowUtils.extractSchemaModel(
-            reactFlow.getNodes() as Node<SchemaDesigner.Table>[],
-            reactFlow.getEdges() as Edge<SchemaDesigner.ForeignKey>[],
-        );
+    const updateTable = async (updatedTable: SchemaDesigner.Table) => {
+        const existingNodes = reactFlow.getNodes() as Node<SchemaDesigner.Table>[];
+        let existingEdges = reactFlow.getEdges() as Edge<SchemaDesigner.ForeignKey>[];
+
+        const existingTableNode = existingNodes.find((node) => node.id === updatedTable.id);
 
-        const tableNode = schemaModel.tables.find((node) => node.id === table.id);
-        if (!tableNode) {
-            console.warn(`Table with id ${table.id} not found`);
+        if (!existingTableNode) {
             return false;
         }
 
-        for (const t of schemaModel.tables) {
-            if (t.id !== table.id) {
-                // Update the foreign keys for other tables
-                t.foreignKeys.forEach((fk) => {
-                    if (
-                        fk.referencedSchemaName === tableNode.schema &&
-                        fk.referencedTableName === tableNode.name
-                    ) {
-                        fk.referencedSchemaName = table.schema;
-                        fk.referencedTableName = table.name;
-                    }
-                });
+        // Updating the table name and schema in all foreign keys that reference this table
+        // This is necessary because the table name and schema might have changed
+        existingEdges.forEach((edge) => {
+            if (
+                edge?.data?.referencedSchemaName === existingTableNode?.data?.schema &&
+                edge?.data?.referencedTableName === existingTableNode?.data?.name
+            ) {
+                edge.data.referencedSchemaName = updatedTable.schema;
+                edge.data.referencedTableName = updatedTable.name;
             }
-        }
-
-        const updatedTable = {
-            ...tableNode,
-            ...table,
-        };
-
-        const updatedSchema = {
-            ...schemaModel,
-            tables: schemaModel.tables.map((node) => (node.id === table.id ? updatedTable : node)),
-        };
-
-        // Delete existing edges for this table
-        const edgesToDelete = reactFlow
-            .getEdges()
-            .filter((edge) => edge.source === table.id || edge.target === table.id);
-
-        await reactFlow.deleteElements({
-            nodes: [],
-            edges: edgesToDelete,
         });
 
-        // Regenerate flow components with updated schema
-        const newFlowComponents = flowUtils.generateSchemaDesignerFlowComponents(updatedSchema);
+        // Update the table node with the new data
+        existingTableNode.data = updatedTable;
 
-        const nodeWithPosition = newFlowComponents.nodes.find((node) => node.id === table.id);
+        // Remove the existing foreign keys from the table
+        existingEdges = existingEdges.filter((edge) => edge.source !== updatedTable.id);
 
-        if (nodeWithPosition) {
-            const edgesForUpdatedTable = newFlowComponents.edges.filter(
-                (edge) => edge.source === table.id || edge.target === table.id,
+        // Add the new foreign keys to the table
+        updatedTable.foreignKeys.forEach((foreignKey) => {
+            const referencedTable = existingNodes.find(
+                (node) =>
+                    node.data.schema === foreignKey.referencedSchemaName &&
+                    node.data.name === foreignKey.referencedTableName,
             );
+            if (!referencedTable) {
+                return;
+            }
 
-            reactFlow.updateNodeData(nodeWithPosition.id, nodeWithPosition.data);
-            reactFlow.addEdges(edgesForUpdatedTable);
-            return true;
-        }
-        return false;
+            foreignKey.columns.forEach((column, index) => {
+                const referencedColumn = foreignKey.referencedColumns[index];
+                existingEdges.push({
+                    id: foreignKey.id,
+                    source: updatedTable.id,
+                    target: referencedTable.id,
+                    sourceHandle: `right-${column}`,
+                    targetHandle: `left-${referencedColumn}`,
+                    markerEnd: {
+                        type: MarkerType.ArrowClosed,
+                    },
+                    data: {
+                        ...foreignKey,
+                        referencedColumns: [referencedColumn],
+                        columns: [column],
+                    },
+                });
+            });
+        });
+
+        reactFlow.setNodes(existingNodes);
+        reactFlow.setEdges(existingEdges);
+        requestAnimationFrame(() => {
+            setCenter(updatedTable.id, true);
+        });
+        return true;
     };
 
     const deleteTable = async (table: SchemaDesigner.Table) => {
@@ -409,6 +398,14 @@ const SchemaDesignerStateProvider: React.FC<SchemaDesignerProviderProps> = ({ ch
         eventBus.emit("updateUndoRedoState", stateStack.canUndo(), stateStack.canRedo());
     };
 
+    function resetView() {
+        setTimeout(async () => {
+            await reactFlow.fitView({
+                nodes: reactFlow.getNodes().filter((node) => node.hidden !== true),
+            });
+        }, 10);
+    }
+
     return (
         <SchemaDesignerContext.Provider
             value={{
@@ -437,6 +434,7 @@ const SchemaDesignerStateProvider: React.FC<SchemaDesignerProviderProps> = ({ ch
                 isInitialized,
                 closeDesigner,
                 resetUndoRedoState,
+                resetView,
             }}>
             {children}
         </SchemaDesignerContext.Provider>

src/reactviews/pages/SchemaDesigner/schemaDesignerUtils.ts

@@ -771,6 +771,53 @@ export const flowUtils = {
     getTableHeight: (table: SchemaDesigner.Table): number =>
         LAYOUT_CONSTANTS.BASE_NODE_HEIGHT + table.columns.length * LAYOUT_CONSTANTS.COLUMN_HEIGHT,
 
+    generatePositions: (
+        nodes: Node<SchemaDesigner.Table>[],
+        edges: Edge<SchemaDesigner.ForeignKey>[],
+    ): {
+        nodes: Node<SchemaDesigner.Table>[];
+        edges: Edge<SchemaDesigner.ForeignKey>[];
+    } => {
+        const graph = new dagre.graphlib.Graph().setDefaultEdgeLabel(() => ({}));
+        graph.setGraph(LAYOUT_CONSTANTS.LAYOUT_OPTIONS);
+
+        for (const node of nodes) {
+            if (node.hidden) continue;
+            graph.setNode(node.id, {
+                width: flowUtils.getTableWidth(),
+                height: flowUtils.getTableHeight(node.data),
+            });
+        }
+
+        for (const edge of edges) {
+            if (edge.hidden) continue;
+            const sourceNode = nodes.find((n) => n.id === edge.source);
+            const targetNode = nodes.find((n) => n.id === edge.target);
+            if (!sourceNode?.hidden && !targetNode?.hidden) {
+                graph.setEdge(edge.source, edge.target);
+            }
+        }
+
+        dagre.layout(graph);
+
+        const layoutedNodes = nodes.map((node) => {
+            if (node.hidden) return node;
+            const dagreNode = graph.node(node.id);
+            return {
+                ...node,
+                position: {
+                    x: dagreNode.x - flowUtils.getTableWidth() / 2,
+                    y: dagreNode.y - flowUtils.getTableHeight(node.data) / 2,
+                },
+            };
+        });
+
+        return {
+            nodes: layoutedNodes,
+            edges,
+        };
+    },
+
     generateSchemaDesignerFlowComponents: (
         schema: SchemaDesigner.Schema,
     ): {

src/reactviews/pages/SchemaDesigner/toolbar/autoArrangeButton.tsx

@@ -29,13 +29,11 @@ export function AutoArrangeButton() {
     const autoArrange = () => {
         eventBus.emit("pushState");
         const nodes = reactFlow.getNodes() as Node<SchemaDesigner.Table>[];
-        const schema = flowUtils.extractSchemaModel(
-            nodes,
-            reactFlow.getEdges() as Edge<SchemaDesigner.ForeignKey>[],
-        );
-        const generateComponenets = flowUtils.generateSchemaDesignerFlowComponents(schema);
+        const edges = reactFlow.getEdges() as Edge<SchemaDesigner.ForeignKey>[];
+        const generateComponenets = flowUtils.generatePositions(nodes, edges);
         reactFlow.setNodes(generateComponenets.nodes);
         reactFlow.setEdges(generateComponenets.edges);
+        context.resetView();
     };
     if (!context) {
         return undefined;