diff --git a/packages/compass-aggregations/src/modules/aggregation.spec.ts b/packages/compass-aggregations/src/modules/aggregation.spec.ts index 1a548461916..03d3ba8263b 100644 --- a/packages/compass-aggregations/src/modules/aggregation.spec.ts +++ b/packages/compass-aggregations/src/modules/aggregation.spec.ts @@ -24,6 +24,7 @@ import { defaultPreferencesInstance } from 'compass-preferences-model'; import { createNoopLogger } from '@mongodb-js/compass-logging/provider'; import { createNoopTrack } from '@mongodb-js/compass-telemetry/provider'; import type { AggregationsStore } from '../stores/store'; +import { PipelineBuilder } from './pipeline-builder/pipeline-builder'; const getMockedStore = ( aggregation: AggregateState, @@ -69,6 +70,7 @@ describe('aggregation module', function () { it('runs an aggregation', async function () { const mockDocuments = [{ id: 1 }, { id: 2 }]; + const stopPreview = spy(PipelineBuilder.prototype, 'stopPreview'); const store: AggregationsStore = ( await configureStore( { pipeline: [] }, @@ -80,9 +82,13 @@ describe('aggregation module', function () { ) ).plugin.store; + stopPreview.resetHistory(); await store.dispatch(runAggregation() as any); const aggregation = store.getState().aggregation; + expect(stopPreview).to.have.been.calledOnce; + stopPreview.restore(); + expect(omit(aggregation, 'documents')).to.deep.equal({ pipeline: [], isLast: true, diff --git a/packages/compass-aggregations/src/modules/aggregation.ts b/packages/compass-aggregations/src/modules/aggregation.ts index a49a997c454..2f867a6eb61 100644 --- a/packages/compass-aggregations/src/modules/aggregation.ts +++ b/packages/compass-aggregations/src/modules/aggregation.ts @@ -29,6 +29,10 @@ import type { DataService } from '../modules/data-service'; import toNS from 'mongodb-ns'; import type { PreferencesAccess } from 'compass-preferences-model'; +// Used in the `comment` on the aggregate command to help identify +// the operation in server logs and currentOp. +export const RUN_AGGREGATION_COMMENT = 'Run aggregation'; + const WRITE_STAGE_LINK = { $merge: 'https://www.mongodb.com/docs/manual/reference/operator/aggregation/merge/', @@ -331,6 +335,8 @@ export const runAggregation = (): PipelineBuilderThunkAction> => { return; } + pipelineBuilder.stopPreview(); + void dispatch(fetchExplainForPipeline()); dispatch({ type: ActionTypes.RunAggregation, @@ -454,6 +460,7 @@ const fetchAggregationData = ( const options: AggregateOptions = { maxTimeMS: maxTimeMS ?? DEFAULT_MAX_TIME_MS, collation: collation ?? undefined, + comment: RUN_AGGREGATION_COMMENT, }; const lastStage = pipeline[pipeline.length - 1]; diff --git a/packages/compass-aggregations/src/modules/pipeline-builder/pipeline-preview-manager.ts b/packages/compass-aggregations/src/modules/pipeline-builder/pipeline-preview-manager.ts index ee8dcf25e7d..e12642935ee 100644 --- a/packages/compass-aggregations/src/modules/pipeline-builder/pipeline-preview-manager.ts +++ b/packages/compass-aggregations/src/modules/pipeline-builder/pipeline-preview-manager.ts @@ -14,6 +14,10 @@ import isEqual from 'lodash/isEqual'; import type { DataService } from '../data-service'; import type { PreferencesAccess } from 'compass-preferences-model'; +// Used in the `comment` on the aggregate command to help identify +// the operation in server logs and currentOp. +export const PREVIEW_AGGREGATION_COMMENT = 'Aggregation preview'; + export const DEFAULT_SAMPLE_SIZE = 100000; export const DEFAULT_PREVIEW_LIMIT = 10; @@ -120,7 +124,10 @@ export class PipelinePreviewManager { previewSize, totalDocumentCount, }), - options, + options: { + comment: PREVIEW_AGGREGATION_COMMENT, + ...options, + }, }); this.queue.delete(idx); return result; diff --git a/packages/compass-e2e-tests/tests/collection-aggregations-tab.test.ts b/packages/compass-e2e-tests/tests/collection-aggregations-tab.test.ts index d9c9d75f91d..9ef627b5a31 100644 --- a/packages/compass-e2e-tests/tests/collection-aggregations-tab.test.ts +++ b/packages/compass-e2e-tests/tests/collection-aggregations-tab.test.ts @@ -995,21 +995,45 @@ describe('Collection aggregations tab', function () { }); it('supports cancelling long-running aggregations', async function () { - if (isTestingWebAtlasCloud()) { + // Fall back to a $function sleep in earlier server versions. + // The driver disconnect cancel doesn't always kill CPU bound op + // on earlier versions. + const useSleepSlowQuery = serverSatisfies('<8.0.0'); + + if (useSleepSlowQuery && isTestingWebAtlasCloud()) { // No $function on the free tier, skipping this test. return this.skip(); } - const unsubscribeAllowWarnings = allowServerWarnings( - 8996503, // Allow "$function is deprecated" warning - (l: LogEntry) => { - return ( - l.id === 23799 && ['Interrupted'].includes(l.attr?.error?.codeName) - ); - } - ); + // We tag each aggregation command with identifiable `comment`s. + const RUN_AGGREGATION_COMMENT = 'Run aggregation'; + const PREVIEW_AGGREGATION_COMMENT = 'Aggregation preview'; + + let sawRunAggregationInterrupt = false; + const unsubscribeAllowWarnings = useSleepSlowQuery + ? allowServerWarnings( + 8996503, // Allow "$function is deprecated" warning + (l: LogEntry) => { + return l.id === 23799 && l.attr?.error?.codeName === 'Interrupted'; + } + ) + : allowServerWarnings((l: LogEntry) => { + const comment = l.attr?.cmd?.comment as string | undefined; + const matches = + l.id === 23799 && + l.attr?.error?.codeName === 'Interrupted' && + comment !== undefined && + [RUN_AGGREGATION_COMMENT, PREVIEW_AGGREGATION_COMMENT].includes( + comment + ); + if (matches && comment === RUN_AGGREGATION_COMMENT) { + sawRunAggregationInterrupt = true; + } + return matches; + }); try { - const slowQuery = `{ + const slowQuery = useSleepSlowQuery + ? `{ sleep: { $function: { body: function () { @@ -1019,7 +1043,42 @@ describe('Collection aggregations tab', function () { lang: "js", }, }, - }`; + }` + : // Nesting this $reduce N times will give runtime of 1000 ^ N, so with + // N = 5 it is basically infinite. + `{ + slow: { + $reduce: { + input: {$range: [0, 1000]}, + initialValue: 0, + in: { + $reduce: { + input: {$range: [0, 1000]}, + initialValue: 0, + in: { + $reduce: { + input: {$range: [0, 1000]}, + initialValue: 0, + in: { + $reduce: { + input: {$range: [0, 1000]}, + initialValue: 0, + in: { + $reduce: { + input: {$range: [0, 1000]}, + initialValue: 0, + in: {$add: ['$$value', 1]} + } + } + } + } + } + } + } + } + } + } +}`; // Set first stage to a very slow $addFields await browser.selectStageOperator(0, '$addFields'); @@ -1037,6 +1096,13 @@ describe('Collection aggregations tab', function () { // load anything and dismissed "Loading" banner) const emptyResultsBanner = browser.$(Selectors.AggregationEmptyResults); await emptyResultsBanner.waitForDisplayed(); + + if (!useSleepSlowQuery) { + // Wait until the server has logged the run aggregation's + // cancellation error before we remove the allowlist, so we don't + // leak it into other tests. + await browser.waitUntil(() => sawRunAggregationInterrupt); + } } finally { unsubscribeAllowWarnings(); }