Add query stages view to Preview Web UI

Author: koszti

core/trino-web-ui/src/main/resources/webapp-preview/src/api/webapp/api.ts

@@ -206,6 +206,79 @@ export interface Session {
     catalogProperties: { [key: string]: string | number | boolean }
 }
 
+export interface QueryStagePlan {
+    id: string
+    jsonRepresentation: string
+    root: {
+        id: string
+    }
+}
+
+export interface QueryStageStats {
+    completedDrivers: number
+    fullyBlocked: boolean
+    totalCpuTime: string
+    userMemoryReservation: string
+    queuedDrivers: number
+    runningDrivers: number
+    rawInputDataSize: string
+    rawInputPositions: number
+    bufferedDataSize: string
+    totalBlockedTime: string
+    outputDataSize: string
+    outputPositions: number
+    totalScheduledTime: string
+    failedScheduledTime: string
+    failedCpuTime: string
+    cumulativeUserMemory: number
+    totalBufferedBytes: number
+    failedCumulativeUserMemory: number
+    peakUserMemoryReservation: string
+}
+
+export interface QueryTask {
+    lastHeartbeat: string
+    needsPlan: boolean
+    estimatedMemory: string
+    outputBuffers: {
+        totalBufferedBytes: number
+    }
+    stats: {
+        createTime: string
+        elapsedTime: string
+        fullyBlocked: boolean
+        blockedDrivers: number
+        completedDrivers: number
+        queuedDrivers: number
+        peakUserMemoryReservation: string
+        runningDrivers: number
+        rawInputDataSize: string
+        rawInputPositions: number
+        totalCpuTime: string
+        userMemoryReservation: string
+    }
+    taskStatus: {
+        nodeId: string
+        taskId: string
+        self: string
+        state: string
+    }
+}
+
+export interface QueryStage {
+    coordinatorOnly: boolean
+    plan: QueryStagePlan
+    stageId: string
+    state: string
+    stageStats: QueryStageStats
+    tasks: QueryTask[]
+}
+
+export interface QueryStages {
+    outputStageId: string
+    stages: QueryStage[]
+}
+
 export interface QueryStatusInfo extends QueryInfoBase {
     session: Session
     query: string
@@ -214,6 +287,11 @@ export interface QueryStatusInfo extends QueryInfoBase {
     retryPolicy: string
     pruned: boolean
     finalQueryInfo: boolean
+    stages: QueryStages
+}
+
+export interface WorkerTaskInfo {
+    dummy: string
 }
 
 export async function statsApi(): Promise<ApiResponse<Stats>> {
@@ -228,6 +306,10 @@ export async function workerStatusApi(nodeId: string): Promise<ApiResponse<Worke
     return await api.get<WorkerStatusInfo>(`/ui/api/worker/${nodeId}/status`)
 }
 
+export async function workerTaskApi(nodeId: string, taskId: string): Promise<ApiResponse<WorkerTaskInfo>> {
+    return await api.get<WorkerTaskInfo>(`/ui/api/worker/${nodeId}/task/${taskId}`)
+}
+
 export async function queryApi(): Promise<ApiResponse<QueryInfo[]>> {
     return await api.get<QueryInfo[]>('/ui/api/query')
 }

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

@@ -29,7 +29,7 @@ import {
 } from '@mui/material'
 import { SparkLineChart } from '@mui/x-charts/SparkLineChart'
 import { Texts } from '../constant.ts'
-import { StackInfo, queryStatusApi, QueryStatusInfo, Session } from '../api/webapp/api.ts'
+import { StackInfo, queryStatusApi, QueryStatusInfo, QueryStage, QueryStages, Session } from '../api/webapp/api.ts'
 import { ApiResponse } from '../api/base.ts'
 import {
     addToHistory,
@@ -40,10 +40,12 @@ import {
     parseDataSize,
     parseDuration,
 } from '../utils/utils.ts'
+import { QueryStageCard } from './QueryStageCard'
 import { CodeBlock } from './CodeBlock.tsx'
 
 interface IQueryStatus {
     info: QueryStatusInfo | null
+    lastSnapshotStages: QueryStages | null
 
     lastScheduledTime: number
     lastCpuTime: number
@@ -66,6 +68,7 @@ export const QueryOverview = () => {
     const { queryId } = useParams()
     const initialQueryStatus: IQueryStatus = {
         info: null,
+        lastSnapshotStages: null,
 
         lastScheduledTime: 0,
         lastCpuTime: 0,
@@ -117,6 +120,7 @@ export const QueryOverview = () => {
                     const newQueryStatusInfo: QueryStatusInfo = apiResponse.data
                     setQueryStatus((prevQueryStatus) => {
                         let lastRefresh = prevQueryStatus.lastRefresh
+
                         const lastScheduledTime = prevQueryStatus.lastScheduledTime
                         const lastCpuTime = prevQueryStatus.lastCpuTime
                         const lastPhysicalInputReadTime = prevQueryStatus.lastPhysicalInputReadTime
@@ -205,6 +209,7 @@ export const QueryOverview = () => {
 
                         return {
                             info: newQueryStatusInfo,
+                            lastSnapshotStages: newQueryStatusInfo.stages,
 
                             lastScheduledTime: parseDuration(newQueryStats.totalScheduledTime) || 0,
                             lastCpuTime: parseDuration(newQueryStats.totalCpuTime) || 0,
@@ -416,6 +421,38 @@ export const QueryOverview = () => {
         }
     }
 
+    const renderStages = (taskRetriesEnabled: boolean) => {
+        const stages = queryStatus.lastSnapshotStages?.stages
+
+        return (
+            <Grid size={{ xs: 12 }}>
+                <Box sx={{ pt: 2 }}>
+                    <Typography variant="h6">Stages</Typography>
+                    <Divider />
+
+                    {stages ? (
+                        stages
+                            .slice()
+                            .sort((a, b) => a.stageId.localeCompare(b.stageId))
+                            .map((stage: QueryStage) => (
+                                <QueryStageCard
+                                    key={stage.stageId}
+                                    stage={stage}
+                                    taskRetriesEnabled={taskRetriesEnabled}
+                                />
+                            ))
+                    ) : (
+                        <>
+                            <Box sx={{ width: '100%', mt: 1 }}>
+                                <Alert severity="info">No stage information available.</Alert>
+                            </Box>
+                        </>
+                    )}
+                </Box>
+            </Grid>
+        )
+    }
+
     const taskRetriesEnabled = queryStatus.info?.retryPolicy == 'TASK'
     return (
         <>
@@ -839,6 +876,7 @@ export const QueryOverview = () => {
                             </Box>
                         </Grid>
                         {renderPreparedQuery()}
+                        {renderStages(taskRetriesEnabled)}
                     </Grid>
                 </Grid>
             )}