|
| 1 | +import { |
| 2 | + isAggregationExplain, |
| 3 | + isShardedAggregationExplain, |
| 4 | + getStageCursorKey, |
| 5 | +} from 'mongodb-explain-compat'; |
| 6 | + |
| 7 | +import type { Stage } from './index'; |
| 8 | + |
| 9 | +export type ExecutionStats = { |
| 10 | + executionSuccess: boolean; |
| 11 | + nReturned: number; |
| 12 | + executionTimeMillis: number; |
| 13 | + totalKeysExamined: number; |
| 14 | + totalDocsExamined: number; |
| 15 | + executionStages: Stage; |
| 16 | + allPlansExecution: unknown[]; |
| 17 | +}; |
| 18 | + |
| 19 | +export const getExecutionStats = (explain: Stage): ExecutionStats => { |
| 20 | + const executionStats = isAggregationExplain(explain) |
| 21 | + ? _getAggregationStats(explain) |
| 22 | + : _getFindStats(explain); |
| 23 | + return executionStats; |
| 24 | +}; |
| 25 | +const _getAggregationStats = (explain: Stage): ExecutionStats => { |
| 26 | + return isShardedAggregationExplain(explain) |
| 27 | + ? _getShardedAggregationStats(explain) |
| 28 | + : _getUnshardedAggregationStats(explain); |
| 29 | +}; |
| 30 | +const _getUnshardedAggregationStats = (explain: Stage): ExecutionStats => { |
| 31 | + const firstStage = explain.stages[0]; |
| 32 | + const cursorKey = getStageCursorKey(firstStage); |
| 33 | + if (!cursorKey) { |
| 34 | + throw new Error('Can not find a cursor stage.'); |
| 35 | + } |
| 36 | + |
| 37 | + const lastStage = explain.stages[explain.stages.length - 1]; |
| 38 | + |
| 39 | + const stats = _getFindStats(firstStage[cursorKey]); |
| 40 | + stats.nReturned = lastStage.nReturned; |
| 41 | + stats.executionTimeMillis = sumArrayProp( |
| 42 | + explain.stages, |
| 43 | + 'executionTimeMillisEstimate' |
| 44 | + ); |
| 45 | + return stats; |
| 46 | +}; |
| 47 | +const _getShardedAggregationStats = (explain: Stage): ExecutionStats => { |
| 48 | + const shardStats = []; |
| 49 | + for (const shardName in explain.shards) { |
| 50 | + const stats = explain.shards[shardName].stages |
| 51 | + ? _getUnshardedAggregationStats(explain.shards[shardName]) |
| 52 | + : _getFindStats(explain.shards[shardName]); |
| 53 | + |
| 54 | + shardStats.push({ |
| 55 | + shardName, |
| 56 | + ...stats, |
| 57 | + }); |
| 58 | + } |
| 59 | + |
| 60 | + const nReturned = sumArrayProp(shardStats, 'nReturned'); |
| 61 | + const executionTimeMillis = sumArrayProp(shardStats, 'executionTimeMillis'); |
| 62 | + const totalKeysExamined = sumArrayProp(shardStats, 'totalKeysExamined'); |
| 63 | + const totalDocsExamined = sumArrayProp(shardStats, 'totalDocsExamined'); |
| 64 | + const response = { |
| 65 | + nReturned, |
| 66 | + executionTimeMillis, |
| 67 | + totalKeysExamined, |
| 68 | + totalDocsExamined, |
| 69 | + allPlansExecution: [], |
| 70 | + executionSuccess: true, |
| 71 | + executionStages: { |
| 72 | + stage: shardStats.length === 1 ? 'SINGLE_SHARD' : 'SHARD_MERGE', |
| 73 | + nReturned, |
| 74 | + executionTimeMillis, |
| 75 | + totalKeysExamined, |
| 76 | + totalDocsExamined, |
| 77 | + shards: shardStats, |
| 78 | + }, |
| 79 | + } as unknown as ExecutionStats; |
| 80 | + |
| 81 | + return response; |
| 82 | +}; |
| 83 | +const _getFindStats = (explain: Stage): ExecutionStats => { |
| 84 | + return explain.executionStats; |
| 85 | +}; |
| 86 | + |
| 87 | +function sumArrayProp<T>(arr: T[], prop: keyof T): number { |
| 88 | + return arr.reduce((acc, x) => acc + Number(x[prop] ?? 0), 0); |
| 89 | +} |
0 commit comments