Fix query plan default viewport in preview UI

core/trino-web-ui/src/main/resources/webapp-preview/src/components/QueryLivePlan.tsx

@@ -13,12 +13,12 @@
  */
 import { useParams } from 'react-router-dom'
 import { useEffect, useRef, useState } from 'react'
-import { Alert, Box, CircularProgress, Grid } from '@mui/material'
-import { ReactFlow, type Edge, type Node, useNodesState, useEdgesState } from '@xyflow/react'
+import { Alert, Box, CircularProgress, Grid, Typography } from '@mui/material'
+import { ReactFlow, type Edge, type Node, useNodesState, useEdgesState, type Viewport } from '@xyflow/react'
 import '@xyflow/react/dist/style.css'
 import { queryStatusApi, QueryStatusInfo } from '../api/webapp/api.ts'
 import { QueryProgressBar } from './QueryProgressBar'
-import { nodeTypes, getLayoutedPlanFlowElements } from './flow/layout'
+import { nodeTypes, getLayoutedPlanFlowElements, getViewportFocusedOnNode } from './flow/layout'
 import { HelpMessage } from './flow/HelpMessage'
 import { getPlanFlowElements } from './flow/flowUtils'
 import { IQueryStatus, LayoutDirectionType } from './flow/types'
@@ -36,6 +36,7 @@ export const QueryLivePlan = () => {
     const [nodes, setNodes, onNodesChange] = useNodesState<Node>([])
     const [edges, setEdges, onEdgesChange] = useEdgesState<Edge>([])
     const [layoutDirection, setLayoutDirection] = useState<LayoutDirectionType>('BT')
+    const [viewport, setViewport] = useState<Viewport>()
 
     const [loading, setLoading] = useState<boolean>(true)
     const [error, setError] = useState<string | null>(null)
@@ -94,6 +95,16 @@ export const QueryLivePlan = () => {
         }
     }
 
+    const focusViewportToFirstStage = () => {
+        const viewportTarget = getViewportFocusedOnNode(nodes, {
+            targetNodeId: 'stage-0',
+            containerWidth: containerRef.current?.clientWidth ?? 0,
+        })
+        if (viewportTarget) {
+            setViewport(viewportTarget)
+        }
+    }
+
     return (
         <>
             {loading && <CircularProgress />}
@@ -115,21 +126,30 @@ export const QueryLivePlan = () => {
                                                 ref={containerRef}
                                                 sx={{ width: '100%', height: '80vh', border: '1px solid #ccc' }}
                                             >
-                                                <ReactFlow
-                                                    nodes={nodes}
-                                                    edges={edges}
-                                                    onNodesChange={onNodesChange}
-                                                    onEdgesChange={onEdgesChange}
-                                                    nodeTypes={nodeTypes}
-                                                    minZoom={0.1}
-                                                    proOptions={{ hideAttribution: true }}
-                                                    defaultViewport={{ x: 200, y: 20, zoom: 0.8 }}
-                                                >
-                                                    <HelpMessage
-                                                        layoutDirection={layoutDirection}
-                                                        onLayoutDirectionChange={setLayoutDirection}
-                                                    />
-                                                </ReactFlow>
+                                                {nodes.length > 0 ? (
+                                                    <ReactFlow
+                                                        nodes={nodes}
+                                                        edges={edges}
+                                                        onNodesChange={onNodesChange}
+                                                        onEdgesChange={onEdgesChange}
+                                                        nodeTypes={nodeTypes}
+                                                        minZoom={0.1}
+                                                        proOptions={{ hideAttribution: true }}
+                                                        viewport={viewport}
+                                                        onViewportChange={setViewport}
+                                                        fitView
+                                                    >
+                                                        <HelpMessage
+                                                            layoutDirection={layoutDirection}
+                                                            onLayoutDirectionChange={setLayoutDirection}
+                                                            onOriginClick={focusViewportToFirstStage}
+                                                        />
+                                                    </ReactFlow>
+                                                ) : (
+                                                    <Typography sx={{ p: 1 }} fontSize="small">
+                                                        Rendering...
+                                                    </Typography>
+                                                )}
                                             </Box>
                                         </Box>
                                     </Grid>

core/trino-web-ui/src/main/resources/webapp-preview/src/components/QueryStagePerformance.tsx

@@ -24,13 +24,13 @@ import {
     Select,
     SelectChangeEvent,
 } from '@mui/material'
-import { type Edge, type Node, ReactFlow, useEdgesState, useNodesState } from '@xyflow/react'
+import { type Edge, type Node, ReactFlow, useEdgesState, useNodesState, type Viewport } from '@xyflow/react'
 import { queryStatusApi, QueryStatusInfo, QueryStage } from '../api/webapp/api.ts'
 import { ApiResponse } from '../api/base.ts'
 import { Texts } from '../constant.ts'
 import { QueryProgressBar } from './QueryProgressBar.tsx'
 import { HelpMessage } from './flow/HelpMessage'
-import { nodeTypes, getLayoutedStagePerformanceElements } from './flow/layout'
+import { nodeTypes, getLayoutedStagePerformanceElements, getViewportFocusedOnNode } from './flow/layout'
 import { LayoutDirectionType } from './flow/types'
 import { getStagePerformanceFlowElements } from './flow/flowUtils.ts'
 
@@ -53,6 +53,7 @@ export const QueryStagePerformance = () => {
     const [nodes, setNodes, onNodesChange] = useNodesState<Node>([])
     const [edges, setEdges, onEdgesChange] = useEdgesState<Edge>([])
     const [layoutDirection, setLayoutDirection] = useState<LayoutDirectionType>('BT')
+    const [viewport, setViewport] = useState<Viewport>()
 
     const [loading, setLoading] = useState<boolean>(true)
     const [error, setError] = useState<string | null>(null)
@@ -135,6 +136,16 @@ export const QueryStagePerformance = () => {
         },
     }
 
+    const focusViewportToFirstPipeline = () => {
+        const viewportTarget = getViewportFocusedOnNode(nodes, {
+            targetNodeId: 'pipeline-0',
+            containerWidth: containerRef.current?.clientWidth ?? 0,
+        })
+        if (viewportTarget) {
+            setViewport(viewportTarget)
+        }
+    }
+
     const handleStageIdChange = (event: SelectChangeEvent) => {
         setStagePlanId(event.target.value as string)
     }
@@ -169,11 +180,14 @@ export const QueryStagePerformance = () => {
                                                         nodeTypes={nodeTypes}
                                                         minZoom={0.1}
                                                         proOptions={{ hideAttribution: true }}
-                                                        defaultViewport={{ x: 200, y: 20, zoom: 0.8 }}
+                                                        viewport={viewport}
+                                                        onViewportChange={setViewport}
+                                                        fitView
                                                     >
                                                         <HelpMessage
                                                             layoutDirection={layoutDirection}
                                                             onLayoutDirectionChange={setLayoutDirection}
+                                                            onOriginClick={focusViewportToFirstPipeline}
                                                             additionalContent={
                                                                 <Box sx={{ mt: 2 }}>
                                                                     <FormControl size="small" sx={{ minWidth: 200 }}>

core/trino-web-ui/src/main/resources/webapp-preview/src/components/flow/HelpMessage.tsx

@@ -11,23 +11,33 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-import { Box, Typography, ToggleButtonGroup, ToggleButton } from '@mui/material'
+import { Box, Button, Typography, ToggleButtonGroup, ToggleButton } from '@mui/material'
 import React from 'react'
 import { LayoutDirectionType } from './types'
 
 interface IHelpMessageProps {
     layoutDirection: LayoutDirectionType
     onLayoutDirectionChange: (layoutDirection: LayoutDirectionType) => void
+    onOriginClick: () => void
     additionalContent?: React.ReactNode
 }
 
-export const HelpMessage = ({ layoutDirection, onLayoutDirectionChange, additionalContent }: IHelpMessageProps) => {
+export const HelpMessage = ({
+    layoutDirection,
+    onLayoutDirectionChange,
+    onOriginClick,
+    additionalContent,
+}: IHelpMessageProps) => {
     const handleLayoutChange = (_event: React.MouseEvent<HTMLElement>, newDirection: LayoutDirectionType | null) => {
         if (newDirection !== null) {
             onLayoutDirectionChange(newDirection)
         }
     }
 
+    const handleOriginClick = () => {
+        onOriginClick()
+    }
+
     return (
         <Box
             sx={{
@@ -47,9 +57,19 @@ export const HelpMessage = ({ layoutDirection, onLayoutDirectionChange, addition
                 alignItems: 'center',
             }}
         >
-            <Typography variant="body2" color="text.secondary" sx={{ mb: 0 }}>
-                Scroll to zoom in/out
-            </Typography>
+            <Box sx={{ display: 'flex', alignItems: 'center', gap: 1, mb: 1 }}>
+                <Typography variant="body2" color="text.secondary" sx={{ mb: 0 }}>
+                    Scroll to zoom in/out
+                </Typography>
+                <Button
+                    onClick={handleOriginClick}
+                    size="small"
+                    variant="text"
+                    sx={{ minWidth: 'auto', px: 1.5, py: 0.25 }}
+                >
+                    Origin
+                </Button>
+            </Box>
 
             <ToggleButtonGroup
                 value={layoutDirection}
@@ -65,7 +85,6 @@ export const HelpMessage = ({ layoutDirection, onLayoutDirectionChange, addition
                     <Typography variant="caption">Horizontal</Typography>
                 </ToggleButton>
             </ToggleButtonGroup>
-
             {additionalContent}
         </Box>
     )

core/trino-web-ui/src/main/resources/webapp-preview/src/components/flow/layout.ts

@@ -12,7 +12,7 @@
  * limitations under the License.
  */
 import Dagre from '@dagrejs/dagre'
-import { type Edge, type Node } from '@xyflow/react'
+import { type Edge, type Node, type Viewport } from '@xyflow/react'
 import { RemoteExchangeNode } from './RemoteExchangeNode.tsx'
 import { PlanFragmentNode } from './PlanFragmentNode.tsx'
 import { OperatorNode } from './OperatorNode.tsx'
@@ -184,3 +184,24 @@ export const getLayoutedStagePerformanceElements = (nodes: Node[], edges: Edge[]
         edges,
     }
 }
+
+export const getViewportFocusedOnNode = (
+    nodes: Node[],
+    options?: { targetNodeId?: string; zoom?: number; padding?: number; containerWidth?: number }
+): Viewport | undefined => {
+    if (nodes.length === 0) {
+        return undefined
+    }
+    const { targetNodeId, zoom = 0.8, padding = 20, containerWidth = 0 } = options ?? {}
+    const targetNode = targetNodeId ? nodes.find((node) => node.id === targetNodeId) : undefined
+    const focusNode = targetNode ?? nodes[0]
+    const focusNodeWidth = focusNode.measured?.width ?? focusNode.width ?? 0
+    return {
+        // If width is known then center horizontally in container otherwise align to left
+        x: focusNodeWidth
+            ? containerWidth / 2 - (focusNode.position.x + focusNodeWidth / 2) * zoom
+            : -focusNode.position.x * zoom + padding,
+        y: -focusNode.position.y * zoom + padding,
+        zoom,
+    }
+}