modified: frontend/src/pages/DatabaseSchema.tsx

frontend/src/pages/DatabaseSchema.tsx

@@ -55,22 +55,39 @@ function TableNode({ data }: { data: TableInfo & Record<string, unknown> }) {
             className="px-3 py-1.5 text-xs flex items-center gap-2 hover:bg-gray-50 cursor-default relative"
             title={col.comment || undefined}
           >
+            {/* Handles on both sides for dynamic edge routing */}
             {pkColumns.has(col.name) && (
-              <Handle
-                type="target"
-                position={Position.Left}
-                id={`target-${col.name}`}
-                className="!bg-yellow-500 !w-2 !h-2 !border-0 !-left-1"
-              />
+              <>
+                <Handle
+                  type="target"
+                  position={Position.Left}
+                  id={`target-left-${col.name}`}
+                  className="!bg-yellow-500 !w-2 !h-2 !border-0 !-left-1"
+                />
+                <Handle
+                  type="target"
+                  position={Position.Right}
+                  id={`target-right-${col.name}`}
+                  className="!bg-yellow-500 !w-2 !h-2 !border-0 !-right-1"
+                />
+              </>
             )}
 
             {fkColumns.has(col.name) && (
-              <Handle
-                type="source"
-                position={Position.Right}
-                id={`source-${col.name}`}
-                className="!bg-blue-500 !w-2 !h-2 !border-0 !-right-1"
-              />
+              <>
+                <Handle
+                  type="source"
+                  position={Position.Left}
+                  id={`source-left-${col.name}`}
+                  className="!bg-blue-500 !w-2 !h-2 !border-0 !-left-1"
+                />
+                <Handle
+                  type="source"
+                  position={Position.Right}
+                  id={`source-right-${col.name}`}
+                  className="!bg-blue-500 !w-2 !h-2 !border-0 !-right-1"
+                />
+              </>
             )}
 
             <div className="flex items-center gap-1 flex-shrink-0">
@@ -132,7 +149,67 @@ function calculateTableWeights(schema: SchemaData): Map<string, number> {
   return weights;
 }
 
-function getLayoutedElements(schema: SchemaData): { nodes: Node[]; edges: Edge[] } {
+// Build edges with optimal handle selection based on node positions
+function buildEdgesWithPositions(
+  schema: SchemaData,
+  nodePositions: Map<string, { x: number; y: number; width: number }>
+): Edge[] {
+  const edges: Edge[] = [];
+
+  schema.tables.forEach((table) => {
+    const sourceId = `${table.schema}.${table.name}`;
+    const sourcePos = nodePositions.get(sourceId);
+    if (!sourcePos) return;
+
+    table.foreign_keys.forEach((fk) => {
+      let targetTable = schema.tables.find(t =>
+        t.name === fk.references_table && t.schema === table.schema
+      );
+      if (!targetTable) {
+        targetTable = schema.tables.find(t => t.name === fk.references_table);
+      }
+
+      if (targetTable) {
+        const targetId = `${targetTable.schema}.${targetTable.name}`;
+        const targetPos = nodePositions.get(targetId);
+        if (!targetPos) return;
+
+        // Calculate center positions
+        const sourceCenterX = sourcePos.x + sourcePos.width / 2;
+        const targetCenterX = targetPos.x + targetPos.width / 2;
+
+        // Choose handles based on relative positions
+        const sourceOnRight = sourceCenterX < targetCenterX;
+        const sourceHandle = sourceOnRight
+          ? `source-right-${fk.column}`
+          : `source-left-${fk.column}`;
+        const targetHandle = sourceOnRight
+          ? `target-left-${fk.references_column}`
+          : `target-right-${fk.references_column}`;
+
+        edges.push({
+          id: `${fk.constraint_name}`,
+          source: sourceId,
+          target: targetId,
+          sourceHandle,
+          targetHandle,
+          type: 'smoothstep',
+          style: { stroke: '#3b82f6', strokeWidth: 1.5 },
+          markerEnd: {
+            type: MarkerType.ArrowClosed,
+            color: '#3b82f6',
+            width: 15,
+            height: 15,
+          },
+        });
+      }
+    });
+  });
+
+  return edges;
+}
+
+function getLayoutedElements(schema: SchemaData): { nodes: Node[]; edges: Edge[]; nodePositions: Map<string, { x: number; y: number; width: number }> } {
   const g = new Dagre.graphlib.Graph().setDefaultEdgeLabel(() => ({}));
 
   const weights = calculateTableWeights(schema);
@@ -144,19 +221,20 @@ function getLayoutedElements(schema: SchemaData): { nodes: Node[]; edges: Edge[]
   });
 
   g.setGraph({
-    rankdir: 'TB',  // Top-to-bottom for wider spread
-    nodesep: 40,    // Horizontal spacing between nodes
-    ranksep: 80,    // Vertical spacing between ranks
+    rankdir: 'TB',
+    nodesep: 40,
+    ranksep: 80,
     marginx: 30,
     marginy: 30,
-    ranker: 'network-simplex',  // Better distribution across plane
+    ranker: 'network-simplex',
   });
 
+  const nodeWidth = 280;
+
   sortedTables.forEach((table) => {
     const nodeId = `${table.schema}.${table.name}`;
-    const width = 280;
     const height = getNodeHeight(table);
-    g.setNode(nodeId, { width, height });
+    g.setNode(nodeId, { width: nodeWidth, height });
   });
 
   schema.tables.forEach((table) => {
@@ -172,57 +250,27 @@ function getLayoutedElements(schema: SchemaData): { nodes: Node[]; edges: Edge[]
 
   Dagre.layout(g);
 
+  const nodePositions = new Map<string, { x: number; y: number; width: number }>();
+
   const nodes: Node[] = schema.tables.map((table) => {
     const nodeId = `${table.schema}.${table.name}`;
     const nodeWithPosition = g.node(nodeId);
+    const x = nodeWithPosition.x - nodeWithPosition.width / 2;
+    const y = nodeWithPosition.y - nodeWithPosition.height / 2;
+
+    nodePositions.set(nodeId, { x, y, width: nodeWidth });
 
     return {
       id: nodeId,
       type: 'table',
-      position: {
-        x: nodeWithPosition.x - nodeWithPosition.width / 2,
-        y: nodeWithPosition.y - nodeWithPosition.height / 2,
-      },
+      position: { x, y },
       data: { ...table } as TableInfo & Record<string, unknown>,
     };
   });
 
-  const edges: Edge[] = [];
+  const edges = buildEdgesWithPositions(schema, nodePositions);
 
-  schema.tables.forEach((table) => {
-    const sourceId = `${table.schema}.${table.name}`;
-
-    table.foreign_keys.forEach((fk) => {
-      let targetTable = schema.tables.find(t =>
-        t.name === fk.references_table && t.schema === table.schema
-      );
-      if (!targetTable) {
-        targetTable = schema.tables.find(t => t.name === fk.references_table);
-      }
-
-      if (targetTable) {
-        const targetId = `${targetTable.schema}.${targetTable.name}`;
-
-        edges.push({
-          id: `${fk.constraint_name}`,
-          source: sourceId,
-          target: targetId,
-          sourceHandle: `source-${fk.column}`,
-          targetHandle: `target-${fk.references_column}`,
-          type: 'smoothstep',
-          style: { stroke: '#3b82f6', strokeWidth: 1.5 },
-          markerEnd: {
-            type: MarkerType.ArrowClosed,
-            color: '#3b82f6',
-            width: 15,
-            height: 15,
-          },
-        });
-      }
-    });
-  });
-
-  return { nodes, edges };
+  return { nodes, edges, nodePositions };
 }
 
 export default function DatabaseSchema() {
@@ -264,8 +312,10 @@ export default function DatabaseSchema() {
     const results: SearchResult[] = [];
 
     schema.tables.forEach(table => {
+      const tableNameLower = table.name.toLowerCase();
+
       // Search by table name
-      if (table.name.toLowerCase().includes(query)) {
+      if (tableNameLower.includes(query)) {
         results.push({
           type: 'table',
           tableName: table.name,
@@ -313,6 +363,35 @@ export default function DatabaseSchema() {
       });
     });
 
+    // Calculate relevance score (lower = better match)
+    const getRelevanceScore = (result: SearchResult): number => {
+      const searchTarget = result.columnName?.toLowerCase() || result.tableName.toLowerCase();
+
+      // Exact match
+      if (searchTarget === query) return 0;
+      // Starts with query
+      if (searchTarget.startsWith(query)) return 1;
+      // Query position in string (earlier = better)
+      const position = searchTarget.indexOf(query);
+      return 2 + position;
+    };
+
+    // Sort: by type (table -> column -> comment), then by relevance
+    const typeOrder: Record<SearchResult['type'], number> = {
+      'table': 0,
+      'column': 1,
+      'comment': 2,
+    };
+
+    results.sort((a, b) => {
+      // First by type
+      const typeDiff = typeOrder[a.type] - typeOrder[b.type];
+      if (typeDiff !== 0) return typeDiff;
+
+      // Then by relevance within same type
+      return getRelevanceScore(a) - getRelevanceScore(b);
+    });
+
     return results.slice(0, 50);
   }, [schema, searchQuery]);
 
@@ -351,11 +430,55 @@ export default function DatabaseSchema() {
     };
   }, [schema, selectedTables, getRelatedTables]);
 
-  const { nodes: layoutedNodes, edges: layoutedEdges } = useMemo(() => {
-    if (!filteredSchema || filteredSchema.tables.length === 0) return { nodes: [], edges: [] };
+  const [nodes, setNodes, onNodesChange] = useNodesState<Node>([]);
+  const [edges, setEdges, onEdgesChange] = useEdgesState<Edge>([]);
+  const [nodePositions, setNodePositions] = useState<Map<string, { x: number; y: number; width: number }>>(new Map());
+
+  const { nodes: layoutedNodes, edges: layoutedEdges, nodePositions: initialPositions } = useMemo(() => {
+    if (!filteredSchema || filteredSchema.tables.length === 0) {
+      return { nodes: [], edges: [], nodePositions: new Map() };
+    }
     return getLayoutedElements(filteredSchema);
   }, [filteredSchema]);
 
+  // Update positions when layout changes
+  useEffect(() => {
+    setNodePositions(initialPositions);
+  }, [initialPositions]);
+
+  // Reset and rebuild when filtered schema changes
+  useEffect(() => {
+    // Clear first to avoid stale edges
+    setNodes([]);
+    setEdges([]);
+
+    // Then set new layout after a tick
+    if (layoutedNodes.length > 0) {
+      requestAnimationFrame(() => {
+        setNodes(layoutedNodes);
+        // Edges will be set by nodePositions effect
+      });
+    }
+  }, [layoutedNodes, layoutedEdges, setNodes, setEdges]);
+
+  // Recalculate edges when positions change
+  useEffect(() => {
+    if (!filteredSchema || nodePositions.size === 0) return;
+    const newEdges = buildEdgesWithPositions(filteredSchema, nodePositions);
+    setEdges(newEdges);
+  }, [nodePositions, filteredSchema, setEdges]);
+
+  // Recalculate edges when nodes are dragged
+  const handleNodeDragStop = useCallback((_event: React.MouseEvent, node: Node) => {
+    if (!filteredSchema) return;
+
+    setNodePositions(prev => {
+      const updated = new Map(prev);
+      updated.set(node.id, { x: node.position.x, y: node.position.y, width: 280 });
+      return updated;
+    });
+  }, [filteredSchema]);
+
   const addTable = (tableName: string) => {
     setSelectedTables(prev => new Set([...prev, tableName]));
     setSearchQuery('');
@@ -373,24 +496,6 @@ export default function DatabaseSchema() {
     setSelectedTables(new Set());
   };
 
-  const [nodes, setNodes, onNodesChange] = useNodesState<Node>([]);
-  const [edges, setEdges, onEdgesChange] = useEdgesState<Edge>([]);
-
-  // Reset and rebuild when filtered schema changes
-  useEffect(() => {
-    // Clear first to avoid stale edges
-    setNodes([]);
-    setEdges([]);
-
-    // Then set new layout after a tick
-    if (layoutedNodes.length > 0) {
-      requestAnimationFrame(() => {
-        setNodes(layoutedNodes);
-        setEdges(layoutedEdges);
-      });
-    }
-  }, [layoutedNodes, layoutedEdges, setNodes, setEdges]);
-
   const getSearchIcon = (type: SearchResult['type']) => {
     switch (type) {
       case 'table':
@@ -548,6 +653,7 @@ export default function DatabaseSchema() {
             edges={edges}
             onNodesChange={onNodesChange}
             onEdgesChange={onEdgesChange}
+            onNodeDragStop={handleNodeDragStop}
             nodeTypes={nodeTypes}
             fitView
             fitViewOptions={{ padding: 0.2 }}