diff --git a/l10n/bundle.l10n.json b/l10n/bundle.l10n.json index 9ad22905..ec52f624 100644 --- a/l10n/bundle.l10n.json +++ b/l10n/bundle.l10n.json @@ -94,7 +94,6 @@ "{0} completed successfully": "{0} completed successfully", "{0} connections": "{0} connections", "{0} created": "{0} created", - "{0} documents returned": "{0} documents returned", "{0} failed: {1}": "{0} failed: {1}", "{0} file(s) were ignored because they do not match the \"*.json\" pattern.": "{0} file(s) were ignored because they do not match the \"*.json\" pattern.", "{0} inserted": "{0} inserted", @@ -185,6 +184,7 @@ "Authenticate to Connect with Your DocumentDB Cluster": "Authenticate to Connect with Your DocumentDB Cluster", "Authenticate using a username and password": "Authenticate using a username and password", "Authenticate using Microsoft Entra ID (Azure AD)": "Authenticate using Microsoft Entra ID (Azure AD)", + "Authenticating…": "Authenticating…", "Authentication configuration is missing for \"{cluster}\".": "Authentication configuration is missing for \"{cluster}\".", "Authentication data (primary connection string) is missing for \"{cluster}\".": "Authentication data (primary connection string) is missing for \"{cluster}\".", "Authentication data (properties.connectionString) is missing for \"{cluster}\".": "Authentication data (properties.connectionString) is missing for \"{cluster}\".", @@ -367,6 +367,7 @@ "DocumentDB for VS Code is not signed in to Azure": "DocumentDB for VS Code is not signed in to Azure", "DocumentDB Local": "DocumentDB Local", "DocumentDB Performance Tips": "DocumentDB Performance Tips", + "DocumentDB Scratchpad": "DocumentDB Scratchpad", "DocumentDB Scratchpad connected to {0}": "DocumentDB Scratchpad connected to {0}", "DocumentDB Scratchpad connected to {0}/{1}": "DocumentDB Scratchpad connected to {0}/{1}", "Documents": "Documents", @@ -600,6 +601,7 @@ "Info from the webview: ": "Info from the webview: ", "Information was confusing": "Information was confusing", "Initializing task...": "Initializing task...", + "Initializing…": "Initializing…", "Inserted {0} document(s).": "Inserted {0} document(s).", "Install Azure Account Extension...": "Install Azure Account Extension...", "Internal error: connectionString must be defined.": "Internal error: connectionString must be defined.", @@ -710,6 +712,7 @@ "No collection has been marked for copy. Please use \"Copy Collection...\" first to select a source collection.": "No collection has been marked for copy. Please use \"Copy Collection...\" first to select a source collection.", "No collection selected.": "No collection selected.", "No Connectivity": "No Connectivity", + "No credentials found for cluster {0}": "No credentials found for cluster {0}", "No credentials found for id {clusterId}": "No credentials found for id {clusterId}", "No credentials found for the selected cluster.": "No credentials found for the selected cluster.", "No database connected": "No database connected", @@ -743,6 +746,7 @@ "Open the VS Code Marketplace to learn more about \"{0}\"": "Open the VS Code Marketplace to learn more about \"{0}\"", "Opening DocumentDB connection…": "Opening DocumentDB connection…", "Operation cancelled.": "Operation cancelled.", + "Operation timed out after {0} seconds": "Operation timed out after {0} seconds", "Optimization Opportunities": "Optimization Opportunities", "Optimize Index Strategy": "Optimize Index Strategy", "Optimizing the index on {0} can improve query performance by better matching the query pattern.": "Optimizing the index on {0} can improve query performance by better matching the query pattern.", @@ -823,10 +827,13 @@ "Report an issue": "Report an issue", "Resource group \"{0}\" already exists in subscription \"{1}\".": "Resource group \"{0}\" already exists in subscription \"{1}\".", "Result: {0}": "Result: {0}", + "Result: Array ({0} elements)": "Result: Array ({0} elements)", + "Result: Cursor ({0} documents)": "Result: Cursor ({0} documents)", "Results found": "Results found", "Retry": "Retry", "Reusing active connection for \"{cluster}\".": "Reusing active connection for \"{cluster}\".", "Revisit connection details and try again.": "Revisit connection details and try again.", + "Right-click a database or collection in the DocumentDB panel and select \"Connect Scratchpad to this database\".": "Right-click a database or collection in the DocumentDB panel and select \"Connect Scratchpad to this database\".", "Role assignment \"{0}\" created for the {2} resource \"{1}\".": "Role assignment \"{0}\" created for the {2} resource \"{1}\".", "Role Assignment {0} created for {1}": "Role Assignment {0} created for {1}", "Role Assignment {0} failed for {1}": "Role Assignment {0} failed for {1}", @@ -834,7 +841,7 @@ "Run All": "Run All", "Run the entire file ({0}+Shift+Enter)": "Run the entire file ({0}+Shift+Enter)", "Run this block ({0}+Enter)": "Run this block ({0}+Enter)", - "Running scratchpad query…": "Running scratchpad query…", + "Running query…": "Running query…", "Running…": "Running…", "Save": "Save", "Save credentials for future connections.": "Save credentials for future connections.", @@ -1009,7 +1016,6 @@ "This will also delete {0}.": "This will also delete {0}.", "This will prevent the query planner from using this index.": "This will prevent the query planner from using this index.", "To connect to Azure resources, you need to sign in to Azure accounts.": "To connect to Azure resources, you need to sign in to Azure accounts.", - "To connect, right-click a database or collection in the DocumentDB panel and select \"Connect Scratchpad to this database\".": "To connect, right-click a database or collection in the DocumentDB panel and select \"Connect Scratchpad to this database\".", "TODO: Share the steps needed to reliably reproduce the problem. Please include actual and expected results.": "TODO: Share the steps needed to reliably reproduce the problem. Please include actual and expected results.", "Total documents to import: {0}": "Total documents to import: {0}", "Total time taken to execute the query on the server": "Total time taken to execute the query on the server", @@ -1079,6 +1085,7 @@ "What's New": "What's New", "Where to save the exported documents?": "Where to save the exported documents?", "with Popover": "with Popover", + "Worker is not running": "Worker is not running", "Working...": "Working...", "Working…": "Working…", "Would you like to open the Collection View?": "Would you like to open the Collection View?", diff --git a/package.json b/package.json index e491869b..fcbda2d5 100644 --- a/package.json +++ b/package.json @@ -551,6 +551,12 @@ "category": "DocumentDB", "command": "vscode-documentdb.command.clearSchemaCache", "title": "Clear Schema Cache" + }, + { + "//": "[Diagnostics] Show Schema Store Stats", + "category": "DocumentDB", + "command": "vscode-documentdb.command.showSchemaStoreStats", + "title": "Show Schema Store Stats" } ], "submenus": [ diff --git a/src/commands/schemaStore/showSchemaStoreStats.ts b/src/commands/schemaStore/showSchemaStoreStats.ts new file mode 100644 index 00000000..f038f34d --- /dev/null +++ b/src/commands/schemaStore/showSchemaStoreStats.ts @@ -0,0 +1,38 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import { type IActionContext } from '@microsoft/vscode-azext-utils'; +import { SchemaStore } from '../../documentdb/SchemaStore'; +import { ext } from '../../extensionVariables'; + +/** + * Command handler: Show SchemaStore statistics in the output channel. + * Displays collection count, document count, field count, and per-collection breakdown. + */ +export function showSchemaStoreStats(_context: IActionContext): void { + const store = SchemaStore.getInstance(); + const stats = store.getStats(); + + ext.outputChannel.appendLog( + `[SchemaStore] Stats: ${String(stats.collectionCount)} collections, ` + + `${String(stats.totalDocuments)} documents analyzed, ` + + `${String(stats.totalFields)} fields discovered`, + ); + + if (stats.collections.length > 0) { + for (const c of stats.collections) { + // Key format is "clusterId::db::collection" — show only db/collection + const parts = c.key.split('::'); + const displayKey = parts.length >= 3 ? `${parts[1]}/${parts[2]}` : c.key; + ext.outputChannel.appendLog( + `[SchemaStore] ${displayKey}: ${String(c.documentCount)} docs, ${String(c.fieldCount)} fields`, + ); + } + } else { + ext.outputChannel.appendLog('[SchemaStore] (empty — no schemas cached)'); + } + + ext.outputChannel.show(); +} diff --git a/src/commands/scratchpad/connectDatabase.ts b/src/commands/scratchpad/connectDatabase.ts index 263955c3..c2fb03fa 100644 --- a/src/commands/scratchpad/connectDatabase.ts +++ b/src/commands/scratchpad/connectDatabase.ts @@ -27,11 +27,12 @@ export async function connectDatabase(_context: IActionContext, node?: DatabaseI l10n.t('DocumentDB Scratchpad connected to {0}/{1}', node.cluster.name, node.databaseInfo.name), ); } else { - // No tree context — show instructions (per plan §2.3) - void vscode.window.showInformationMessage( - l10n.t( - 'To connect, right-click a database or collection in the DocumentDB panel and select "Connect Scratchpad to this database".', + // No tree context — show instructions as modal dialog + void vscode.window.showInformationMessage(l10n.t('No database connected'), { + modal: true, + detail: l10n.t( + 'Right-click a database or collection in the DocumentDB panel and select "Connect Scratchpad to this database".', ), - ); + }); } } diff --git a/src/commands/scratchpad/executeScratchpadCode.ts b/src/commands/scratchpad/executeScratchpadCode.ts index 2a305f21..98bb02a2 100644 --- a/src/commands/scratchpad/executeScratchpadCode.ts +++ b/src/commands/scratchpad/executeScratchpadCode.ts @@ -3,24 +3,51 @@ * Licensed under the MIT License. See License.txt in the project root for license information. *--------------------------------------------------------------------------------------------*/ -import { openReadOnlyContent } from '@microsoft/vscode-azext-utils'; +import { + UserCancelledError, + callWithTelemetryAndErrorHandling, + openReadOnlyContent, +} from '@microsoft/vscode-azext-utils'; import * as l10n from '@vscode/l10n'; import { type Document, type WithId } from 'mongodb'; import * as vscode from 'vscode'; +import { CredentialCache } from '../../documentdb/CredentialCache'; import { SchemaStore } from '../../documentdb/SchemaStore'; import { ScratchpadEvaluator } from '../../documentdb/scratchpad/ScratchpadEvaluator'; import { ScratchpadService } from '../../documentdb/scratchpad/ScratchpadService'; import { formatError, formatResult } from '../../documentdb/scratchpad/resultFormatter'; import { type ExecutionResult, type ScratchpadConnection } from '../../documentdb/scratchpad/types'; +import { getHostsFromConnectionString } from '../../documentdb/utils/connectionStringHelpers'; +import { addDomainInfoToProperties } from '../../documentdb/utils/getClusterMetadata'; /** Shared evaluator instance — lazily created, reused across runs. */ let evaluator: ScratchpadEvaluator | undefined; +/** + * Dispose the shared evaluator instance (kills the worker thread). + * Called during extension deactivation. + */ +export function disposeEvaluator(): void { + evaluator?.dispose(); + evaluator = undefined; +} + +/** + * Gracefully shut down the evaluator's worker thread (closes MongoClient). + * Called when the scratchpad connection is cleared or when all scratchpad editors close. + * The worker will be re-spawned lazily on the next Run. + */ +export function shutdownEvaluator(): void { + void evaluator?.shutdown(); +} + /** * Executes scratchpad code and displays the result in a read-only side panel. * Used by both `runAll` and `runSelected` commands. */ -export async function executeScratchpadCode(code: string): Promise { +export type ScratchpadRunMode = 'runAll' | 'runSelected'; + +export async function executeScratchpadCode(code: string, runMode: ScratchpadRunMode): Promise { const service = ScratchpadService.getInstance(); const connection = service.getConnection(); if (!connection) { @@ -38,54 +65,148 @@ export async function executeScratchpadCode(code: string): Promise { service.setExecuting(true); - await vscode.window.withProgress( - { - location: vscode.ProgressLocation.Notification, - title: l10n.t('Running scratchpad query…'), - cancellable: false, - }, - async () => { - const startTime = Date.now(); - try { - const result = await evaluator!.evaluate(connection, code); - const formattedOutput = formatResult(result, code, connection); - - // Feed document results to SchemaStore for cross-tab schema sharing - feedResultToSchemaStore(result, connection); - - const resultLabel = l10n.t('{0}/{1} — Results', connection.clusterDisplayName, connection.databaseName); - - await openReadOnlyContent( - { label: resultLabel, fullId: `scratchpad-results-${Date.now()}` }, - formattedOutput, - '.jsonc', - { viewColumn: vscode.ViewColumn.Beside, preserveFocus: true }, - ); - } catch (error: unknown) { - const durationMs = Date.now() - startTime; - const errorMessage = error instanceof Error ? error.message : String(error); - const formattedOutput = formatError(error, code, durationMs, connection); - - const errorLabel = l10n.t('{0}/{1} — Error', connection.clusterDisplayName, connection.databaseName); - - await openReadOnlyContent( - { label: errorLabel, fullId: `scratchpad-error-${Date.now()}` }, - formattedOutput, - '.jsonc', - { viewColumn: vscode.ViewColumn.Beside, preserveFocus: true }, - ); - - void vscode.window.showErrorMessage(l10n.t('Scratchpad execution failed: {0}', errorMessage)); - } finally { - service.setExecuting(false); - } - }, - ); + // callWithTelemetryAndErrorHandling automatically tracks: + // - duration (measured from callback start to end) + // - result: 'Succeeded' | 'Failed' | 'Canceled' + // - error / errorMessage (from thrown errors) + // We add custom properties for scratchpad-specific analytics. + await callWithTelemetryAndErrorHandling('scratchpad.execute', async (context) => { + context.errorHandling.suppressDisplay = true; // we show our own error UI + context.errorHandling.rethrow = false; + + // ── Pre-execution telemetry (known before eval) ────────────── + context.telemetry.properties.sessionId = evaluator!.sessionId ?? 'none'; + context.telemetry.properties.sessionEvalCount = String(evaluator!.sessionEvalCount); + context.telemetry.properties.authMethod = evaluator!.sessionAuthMethod ?? 'unknown'; + context.telemetry.properties.runMode = runMode; + context.telemetry.measurements.codeLineCount = code.split('\n').length; + + // Domain info — privacy-safe hashed host data for platform analytics + const domainProps: Record = {}; + collectDomainTelemetry(connection, domainProps); + Object.assign(context.telemetry.properties, domainProps); + + await vscode.window.withProgress( + { + location: vscode.ProgressLocation.Notification, + title: l10n.t('DocumentDB Scratchpad'), + cancellable: true, + }, + async (progress, token) => { + let cancelled = false; + + token.onCancellationRequested(() => { + cancelled = true; + evaluator?.killWorker(); + }); + + const startTime = Date.now(); + try { + const result = await evaluator!.evaluate(connection, code, (message) => { + progress.report({ message }); + }); + + // ── Post-execution telemetry (known after success) ──── + context.telemetry.properties.resultType = result.type ?? 'null'; + context.telemetry.measurements.initDurationMs = evaluator!.lastInitDurationMs; + // sessionId/sessionEvalCount may have changed after evaluate (if worker was spawned) + context.telemetry.properties.sessionId = evaluator!.sessionId ?? 'none'; + context.telemetry.properties.sessionEvalCount = String(evaluator!.sessionEvalCount); + context.telemetry.properties.authMethod = evaluator!.sessionAuthMethod ?? 'unknown'; + + const formattedOutput = formatResult(result, code, connection); + feedResultToSchemaStore(result, connection); + + const resultLabel = l10n.t( + '{0}/{1} — Results', + connection.clusterDisplayName, + connection.databaseName, + ); + + await openReadOnlyContent( + { label: resultLabel, fullId: `scratchpad-results-${Date.now()}` }, + formattedOutput, + '.jsonc', + { viewColumn: vscode.ViewColumn.Beside, preserveFocus: true }, + ); + + // result: 'Succeeded' is set automatically by the framework + } catch (error: unknown) { + // Update session telemetry even on failure (worker may have spawned before failing) + context.telemetry.properties.sessionId = evaluator!.sessionId ?? 'none'; + context.telemetry.properties.sessionEvalCount = String(evaluator!.sessionEvalCount); + context.telemetry.properties.authMethod = evaluator!.sessionAuthMethod ?? 'unknown'; + context.telemetry.measurements.initDurationMs = evaluator!.lastInitDurationMs; + + if (cancelled) { + // Throw UserCancelledError so framework marks result as 'Canceled' + throw new UserCancelledError('scratchpad.execute'); + } + + // Show our own error UI before re-throwing + const errorMessage = error instanceof Error ? error.message : String(error); + const durationMs = Date.now() - startTime; + const formattedOutput = formatError(error, code, durationMs, connection); + + const errorLabel = l10n.t( + '{0}/{1} — Error', + connection.clusterDisplayName, + connection.databaseName, + ); + + await openReadOnlyContent( + { label: errorLabel, fullId: `scratchpad-error-${Date.now()}` }, + formattedOutput, + '.jsonc', + { viewColumn: vscode.ViewColumn.Beside, preserveFocus: true }, + ); + + void vscode.window.showErrorMessage(l10n.t('Scratchpad execution failed: {0}', errorMessage)); + + // Re-throw so framework automatically captures result: 'Failed', + // error, and errorMessage in telemetry + throw error; + } finally { + service.setExecuting(false); + } + }, + ); + }); } +// ─── Domain telemetry ──────────────────────────────────────────────────────── + +/** + * Collects domain info from the scratchpad connection's cached credentials. + * Reuses the same hashing logic as the connection metadata telemetry. + */ +function collectDomainTelemetry( + connection: ScratchpadConnection, + properties: Record, +): void { + try { + const credentials = CredentialCache.getCredentials(connection.clusterId); + if (!credentials?.connectionString) { + return; + } + const hosts = getHostsFromConnectionString(credentials.connectionString); + addDomainInfoToProperties(hosts, properties); + } catch { + // Domain info is best-effort — don't fail telemetry if parsing fails + } +} + +/** + * Maximum number of documents to feed to SchemaStore per execution. + * If the result set is larger, a random sample of this size is used. + */ +const SCHEMA_DOC_CAP = 100; + /** * Very conservative schema feeding: only feed 'Cursor' and 'Document' result types. * Extracts documents from the printable result and adds them to SchemaStore. + * + * Caps at {@link SCHEMA_DOC_CAP} documents (randomly sampled if more). */ function feedResultToSchemaStore(result: ExecutionResult, connection: ScratchpadConnection): void { // Only feed known document-producing result types @@ -103,17 +224,53 @@ function feedResultToSchemaStore(result: ExecutionResult, connection: Scratchpad return; } - // Normalize to array - const items: unknown[] = Array.isArray(printable) ? printable : [printable]; + // CursorIterationResult from @mongosh wraps documents in { cursorHasMore, documents }. + // Only unwrap when the full wrapper shape is present to avoid false positives + // on user documents that happen to have a `documents` field. + let items: unknown[]; + if ( + typeof printable === 'object' && + !Array.isArray(printable) && + 'cursorHasMore' in printable && + typeof (printable as Record).cursorHasMore === 'boolean' && + 'documents' in printable && + Array.isArray((printable as { documents: unknown }).documents) + ) { + items = (printable as { documents: unknown[] }).documents; + } else if (Array.isArray(printable)) { + items = printable; + } else { + items = [printable]; + } // Filter to actual document objects with _id (not primitives, not nested arrays, // not projection results with _id: 0 which have artificial shapes) - const docs = items.filter( + let docs = items.filter( (d): d is WithId => d !== null && d !== undefined && typeof d === 'object' && !Array.isArray(d) && '_id' in d, ); - if (docs.length > 0) { - SchemaStore.getInstance().addDocuments(connection.clusterId, ns.db, ns.collection, docs); + if (docs.length === 0) { + return; + } + + // Cap at SCHEMA_DOC_CAP documents — randomly sample if more + if (docs.length > SCHEMA_DOC_CAP) { + docs = randomSample(docs, SCHEMA_DOC_CAP); + } + + SchemaStore.getInstance().addDocuments(connection.clusterId, ns.db, ns.collection, docs); +} + +/** + * Partial Fisher–Yates random sample of `count` items from `array`. + * Only performs `count` swaps instead of shuffling the entire array. + */ +function randomSample(array: T[], count: number): T[] { + const copy = [...array]; + for (let i = 0; i < count; i++) { + const j = i + Math.floor(Math.random() * (copy.length - i)); + [copy[i], copy[j]] = [copy[j], copy[i]]; } + return copy.slice(0, count); } diff --git a/src/commands/scratchpad/newScratchpad.ts b/src/commands/scratchpad/newScratchpad.ts index 7121a0c3..d3fcdd8a 100644 --- a/src/commands/scratchpad/newScratchpad.ts +++ b/src/commands/scratchpad/newScratchpad.ts @@ -39,6 +39,7 @@ export async function newScratchpad(_context: IActionContext, node?: DatabaseIte headerComment, '// Use Ctrl+Enter (Cmd+Enter) to run the current block', '// Use Ctrl+Shift+Enter (Cmd+Shift+Enter) to run the entire file', + '// Note: when running multiple statements, only the last result is displayed', '', `db.getCollection('${collectionName}').find({ })`, '', diff --git a/src/commands/scratchpad/runAll.ts b/src/commands/scratchpad/runAll.ts index 65016eaf..d9781730 100644 --- a/src/commands/scratchpad/runAll.ts +++ b/src/commands/scratchpad/runAll.ts @@ -32,5 +32,5 @@ export async function runAll(_context: IActionContext): Promise { return; } - await executeScratchpadCode(code); + await executeScratchpadCode(code, 'runAll'); } diff --git a/src/commands/scratchpad/runSelected.ts b/src/commands/scratchpad/runSelected.ts index 7f62b7d2..0cf30c06 100644 --- a/src/commands/scratchpad/runSelected.ts +++ b/src/commands/scratchpad/runSelected.ts @@ -68,5 +68,5 @@ export async function runSelected(_context: IActionContext, startLine?: number, return; } - await executeScratchpadCode(codeToRun); + await executeScratchpadCode(codeToRun, 'runSelected'); } diff --git a/src/documentdb/ClustersExtension.ts b/src/documentdb/ClustersExtension.ts index 64cdc6ba..f6554267 100644 --- a/src/documentdb/ClustersExtension.ts +++ b/src/documentdb/ClustersExtension.ts @@ -56,7 +56,9 @@ import { removeDiscoveryRegistry } from '../commands/removeDiscoveryRegistry/rem import { retryAuthentication } from '../commands/retryAuthentication/retryAuthentication'; import { revealView } from '../commands/revealView/revealView'; import { clearSchemaCache } from '../commands/schemaStore/clearSchemaCache'; +import { showSchemaStoreStats } from '../commands/schemaStore/showSchemaStoreStats'; import { connectDatabase } from '../commands/scratchpad/connectDatabase'; +import { disposeEvaluator, shutdownEvaluator } from '../commands/scratchpad/executeScratchpadCode'; import { newScratchpad } from '../commands/scratchpad/newScratchpad'; import { runAll } from '../commands/scratchpad/runAll'; import { runSelected } from '../commands/scratchpad/runSelected'; @@ -204,6 +206,59 @@ export class ClustersExtension implements vscode.Disposable { const scratchpadService = ScratchpadService.getInstance(); ext.context.subscriptions.push(scratchpadService); + // Register evaluator disposal for clean worker shutdown on deactivation + ext.context.subscriptions.push({ dispose: disposeEvaluator }); + + // Shut down the scratchpad worker when connection is cleared + ext.context.subscriptions.push( + scratchpadService.onDidChangeState(() => { + if (!scratchpadService.isConnected()) { + ext.outputChannel.debug('[Scratchpad] Connection cleared — shutting down worker'); + shutdownEvaluator(); + } + }), + ); + + // Shut down the scratchpad worker when the last .documentdb editor closes + ext.context.subscriptions.push( + vscode.window.tabGroups.onDidChangeTabs((event) => { + // Only react when tabs are closed + if (event.closed.length === 0) { + return; + } + + // Check if any closed tab was a scratchpad + const closedScratchpad = event.closed.some((tab) => { + const input = tab.input; + return ( + input instanceof vscode.TabInputText && + (input.uri.path.endsWith('.documentdb') || input.uri.path.endsWith('.documentdb.js')) + ); + }); + + if (!closedScratchpad) { + return; + } + + // Check if any scratchpad tabs remain open + const hasOpenScratchpad = vscode.window.tabGroups.all.some((group) => + group.tabs.some((tab) => { + const input = tab.input; + return ( + input instanceof vscode.TabInputText && + (input.uri.path.endsWith('.documentdb') || + input.uri.path.endsWith('.documentdb.js')) + ); + }), + ); + + if (!hasOpenScratchpad) { + ext.outputChannel.debug('[Scratchpad] All editors closed — shutting down worker'); + shutdownEvaluator(); + } + }), + ); + // Register CodeLens provider for scratchpad files const codeLensProvider = new ScratchpadCodeLensProvider(); ext.context.subscriptions.push(codeLensProvider); @@ -233,6 +288,11 @@ export class ClustersExtension implements vscode.Disposable { registerCommand('vscode-documentdb.command.clearSchemaCache', withCommandCorrelation(clearSchemaCache)); + registerCommand( + 'vscode-documentdb.command.showSchemaStoreStats', + withCommandCorrelation(showSchemaStoreStats), + ); + //// General Commands: registerCommandWithTreeNodeUnwrapping( diff --git a/src/documentdb/scratchpad/ScratchpadEvaluator.ts b/src/documentdb/scratchpad/ScratchpadEvaluator.ts index 012def87..fccabcb9 100644 --- a/src/documentdb/scratchpad/ScratchpadEvaluator.ts +++ b/src/documentdb/scratchpad/ScratchpadEvaluator.ts @@ -4,105 +4,510 @@ *--------------------------------------------------------------------------------------------*/ import * as l10n from '@vscode/l10n'; -import { EventEmitter } from 'events'; -import * as vm from 'vm'; +import { randomUUID } from 'crypto'; +import * as path from 'path'; import * as vscode from 'vscode'; +import { Worker } from 'worker_threads'; import { ext } from '../../extensionVariables'; -import { ClustersClient } from '../ClustersClient'; +import { CredentialCache } from '../CredentialCache'; import { type ExecutionResult, type ScratchpadConnection } from './types'; +import { + type MainToWorkerMessage, + type SerializableExecutionResult, + type SerializableMongoClientOptions, + type WorkerToMainMessage, +} from './workerTypes'; + +/** Worker lifecycle states */ +type WorkerState = 'idle' | 'spawning' | 'ready' | 'executing'; /** - * Evaluates scratchpad code in-process using the `@mongosh` pipeline, - * reusing the existing authenticated `MongoClient` from `ClustersClient`. + * Evaluates scratchpad code in a persistent worker thread. * - * A fresh `ShellInstanceState` + `ShellEvaluator` is created per execution - * to avoid variable leakage between runs. Only the `MongoClient` is reused. + * The worker owns its own database client (authenticated via credentials from + * `CredentialCache`) and stays alive between runs. This provides: + * - Infinite loop safety (main thread can kill the worker) + * - Client isolation from the Collection View + * - Zero re-auth overhead after the first run * - * Dependencies (`@mongosh/*`) are lazy-imported to avoid loading ~2-4 MB - * of Babel + async-rewriter at extension activation time. + * The public API is unchanged from the in-process evaluator: + * `evaluate(connection, code) → Promise` */ -export class ScratchpadEvaluator { +export class ScratchpadEvaluator implements vscode.Disposable { + private _worker: Worker | undefined; + private _workerState: WorkerState = 'idle'; + /** Which cluster the live worker is connected to (to detect cluster switches) */ + private _workerClusterId: string | undefined; + + /** + * Telemetry session ID — generated on worker spawn, stable across evals within the + * same worker lifecycle. Resets when the worker is terminated/respawned. + * Used to correlate multiple scratchpad runs within a single "session". + */ + private _sessionId: string | undefined; + + /** Number of eval calls completed during this worker session (for usage tracking). */ + private _sessionEvalCount: number = 0; + + /** Auth mechanism used for the current worker session (for telemetry). */ + private _sessionAuthMethod: string | undefined; + + /** Pending request correlation map: requestId → { resolve, reject } */ + private _pendingRequests = new Map< + string, + { + resolve: (value: unknown) => void; + reject: (error: Error) => void; + } + >(); + + /** Telemetry accessors — read by the command layer for telemetry properties. */ + get sessionId(): string | undefined { + return this._sessionId; + } + get sessionEvalCount(): number { + return this._sessionEvalCount; + } + get sessionAuthMethod(): string | undefined { + return this._sessionAuthMethod; + } + + /** Duration of the last worker init (spawn + auth), in ms. 0 if worker was already alive. */ + private _lastInitDurationMs: number = 0; + get lastInitDurationMs(): number { + return this._lastInitDurationMs; + } + /** * Evaluate user code against the connected database. * * @param connection - Active scratchpad connection (clusterId + databaseName). * @param code - JavaScript code string to evaluate. + * @param onProgress - Optional callback for phased progress reporting. * @returns Formatted execution result with type, printable value, and timing. */ - async evaluate(connection: ScratchpadConnection, code: string): Promise { - // Lazy-import @mongosh packages to avoid loading at activation time - const { NodeDriverServiceProvider } = await import('@mongosh/service-provider-node-driver'); - const { ShellInstanceState } = await import('@mongosh/shell-api'); - const { ShellEvaluator } = await import('@mongosh/shell-evaluator'); - - // Reuse the existing authenticated MongoClient - const client = await ClustersClient.getClient(connection.clusterId); - const mongoClient = client.getMongoClient(); - - // Create fresh shell context per execution (no variable leakage) - const bus = new EventEmitter(); - const serviceProvider = new NodeDriverServiceProvider(mongoClient, bus, { - productDocsLink: 'https://github.com/microsoft/vscode-documentdb', - productName: 'DocumentDB for VS Code Scratchpad', - }); - const instanceState = new ShellInstanceState(serviceProvider, bus); - const evaluator = new ShellEvaluator(instanceState); + async evaluate( + connection: ScratchpadConnection, + code: string, + onProgress?: (message: string) => void, + ): Promise { + // Intercept scratchpad-specific commands before they reach the worker + const trimmed = code.trim(); + const helpResult = this.tryHandleHelp(trimmed); + if (helpResult) { + return { ...helpResult, durationMs: 0 }; + } - // Set up eval context with shell globals (db, ObjectId, ISODate, etc.) - const context = {}; - instanceState.setCtx(context); + // Ensure worker is alive and connected to the right cluster + const needsSpawn = + !this._worker || this._workerState === 'idle' || this._workerClusterId !== connection.clusterId; + if (needsSpawn) { + onProgress?.(l10n.t('Initializing…')); + } - // Pre-select the target database - await evaluator.customEval( - customEvalFn, - `use(${JSON.stringify(connection.databaseName)})`, - context, - 'scratchpad', - ); + const initStartTime = Date.now(); + await this.ensureWorker(connection, onProgress); + this._lastInitDurationMs = needsSpawn ? Date.now() - initStartTime : 0; - // Execute with timeout - const timeoutMs = (vscode.workspace.getConfiguration().get(ext.settingsKeys.shellTimeout) ?? 30) * 1000; + // Send eval message and await result + onProgress?.(l10n.t('Running query…')); + const timeoutSec = vscode.workspace.getConfiguration().get(ext.settingsKeys.shellTimeout) ?? 30; + const timeoutMs = timeoutSec * 1000; - const startTime = Date.now(); + const result = await this.sendEval(connection, code, timeoutMs); + return result; + } - // Intercept scratchpad-specific commands before they reach @mongosh - const trimmed = code.trim(); - const helpResult = this.tryHandleHelp(trimmed); - if (helpResult) { - return { ...helpResult, durationMs: Date.now() - startTime }; + /** + * Gracefully shut down the worker: close the database client, then terminate thread. + * Returns after the worker has confirmed shutdown or after a timeout. + */ + async shutdown(): Promise { + if (!this._worker || this._workerState === 'idle') { + return; + } + + try { + await this.sendRequest({ type: 'shutdown', requestId: '' }, 5000); + } catch { + // Shutdown timed out or failed — force-kill + } + + this.terminateWorker(); + } + + /** + * Force-terminate the worker thread immediately. + * Used for infinite loop recovery (timeout) and cancellation. + */ + killWorker(): void { + this.terminateWorker(); + } + + dispose(): void { + this.terminateWorker(); + } + + // ─── Private: Worker lifecycle ─────────────────────────────────────────── + + /** + * Ensure a worker is alive and connected to the correct cluster. + * Spawns a new worker if needed (lazy), or kills and respawns if the + * cluster has changed. + */ + private async ensureWorker( + connection: ScratchpadConnection, + onProgress?: (message: string) => void, + ): Promise { + // If worker is alive but connected to a different cluster, shut it down + if (this._worker && this._workerClusterId !== connection.clusterId) { + this.terminateWorker(); } - const evalPromise = evaluator.customEval(customEvalFn, code, context, 'scratchpad'); - const timeoutPromise = new Promise((_resolve, reject) => { - setTimeout(() => reject(new Error(l10n.t('Execution timed out'))), timeoutMs); + // If no worker exists, spawn one + if (!this._worker || this._workerState === 'idle') { + await this.spawnWorker(connection, onProgress); + } + } + + /** + * Spawn a new worker thread and send the init message. + */ + private async spawnWorker(connection: ScratchpadConnection, onProgress?: (message: string) => void): Promise { + this._workerState = 'spawning'; + + // Resolve worker script path (same directory as the main bundle in dist/) + const workerPath = path.join(__dirname, 'scratchpadWorker.js'); + this._worker = new Worker(workerPath); + this._workerClusterId = connection.clusterId; + + // Listen for messages from the worker + this._worker.on('message', (msg: WorkerToMainMessage) => { + this.handleWorkerMessage(msg); + }); + + // Listen for worker exit (crash or termination) + this._worker.on('exit', (exitCode: number) => { + ext.outputChannel.debug(`[Scratchpad Worker] Worker exited with code ${String(exitCode)}`); + this.handleWorkerExit(); + }); + + this._worker.on('error', (error: Error) => { + ext.outputChannel.error(`[Scratchpad Worker] ${error.message}`); }); - const result = await Promise.race([evalPromise, timeoutPromise]); - const durationMs = Date.now() - startTime; + // Build init message from cached credentials and send to worker. + // If init fails (bad credentials, unreachable host, etc.), tear down + // the worker so the next evaluate() call can respawn cleanly. + try { + const initMsg = this.buildInitMessage(connection); - // customEval already runs toShellResult() internally via resultHandler, - // so `result` is already a ShellResult { type, printable, rawValue, source? }. - const shellResult = result as { - type: string | null; - printable: unknown; - source?: { namespace?: { db: string; collection: string } }; + // Start a new telemetry session for this worker lifecycle + this._sessionId = randomUUID(); + this._sessionEvalCount = 0; + this._sessionAuthMethod = initMsg.authMechanism; + + // Send init and wait for acknowledgment + onProgress?.(l10n.t('Authenticating…')); + await this.sendRequest(initMsg, 30000); + this._workerState = 'ready'; + } catch (error) { + this.terminateWorker(); + throw error; + } + } + + /** + * Build the init message from CredentialCache data. + */ + private buildInitMessage(connection: ScratchpadConnection): MainToWorkerMessage & { type: 'init' } { + const credentials = CredentialCache.getCredentials(connection.clusterId); + if (!credentials) { + throw new Error(l10n.t('No credentials found for cluster {0}', connection.clusterId)); + } + + const authMechanism = credentials.authMechanism ?? 'NativeAuth'; + + // Build connection string + let connectionString: string; + if (authMechanism === 'NativeAuth') { + connectionString = CredentialCache.getConnectionStringWithPassword(connection.clusterId); + } else { + // Entra ID: use connection string without embedded credentials + connectionString = credentials.connectionString; + } + + // Build serializable MongoClientOptions + const clientOptions: SerializableMongoClientOptions = { + serverSelectionTimeoutMS: credentials.emulatorConfiguration?.isEmulator ? 4000 : undefined, + tlsAllowInvalidCertificates: + credentials.emulatorConfiguration?.isEmulator && + credentials.emulatorConfiguration?.disableEmulatorSecurity + ? true + : undefined, }; return { - type: shellResult.type, - printable: shellResult.printable, - durationMs, - source: shellResult.source?.namespace - ? { - namespace: { - db: shellResult.source.namespace.db, - collection: shellResult.source.namespace.collection, - }, - } - : undefined, + type: 'init', + requestId: '', + connectionString, + clientOptions, + databaseName: connection.databaseName, + authMechanism: authMechanism as 'NativeAuth' | 'MicrosoftEntraID', + tenantId: credentials.entraIdConfig?.tenantId, + // TODO(F11): Read from documentDB.mongoShell.batchSize setting and wire in worker + displayBatchSize: 50, }; } + /** + * Send an eval message to the worker and await the result. + */ + private async sendEval( + connection: ScratchpadConnection, + code: string, + timeoutMs: number, + ): Promise { + this._workerState = 'executing'; + this._sessionEvalCount++; + + const evalMsg: MainToWorkerMessage = { + type: 'eval', + requestId: '', + code, + databaseName: connection.databaseName, + }; + + try { + const result = await this.sendRequest<{ result: SerializableExecutionResult }>(evalMsg, timeoutMs); + + // Deserialize the result — printable is a canonical EJSON string from the worker. + // Canonical EJSON preserves all BSON types (ObjectId, Date, Decimal128, Int32, + // Long, Double, etc.) so that SchemaAnalyzer correctly identifies field types. + const serResult = result.result; + let printable: unknown; + try { + const { EJSON } = await import('bson'); + printable = EJSON.parse(serResult.printable, { relaxed: false }); + } catch { + // Fallback to JSON.parse if EJSON fails, then raw string + try { + printable = JSON.parse(serResult.printable) as unknown; + } catch { + printable = serResult.printable; + } + } + + return { + type: serResult.type, + printable, + durationMs: serResult.durationMs, + source: serResult.source, + }; + } finally { + if (this._workerState === 'executing') { + this._workerState = 'ready'; + } + } + } + + // ─── Private: IPC request/response ─────────────────────────────────────── + + /** + * Send a message to the worker and return a promise that resolves + * when the corresponding response arrives. + */ + private sendRequest(msg: MainToWorkerMessage, timeoutMs: number): Promise { + if (!this._worker) { + return Promise.reject(new Error(l10n.t('Worker is not running'))); + } + + const requestId = randomUUID(); + const msgWithId = { ...msg, requestId }; + + return new Promise((resolve, reject) => { + this._pendingRequests.set(requestId, { + resolve: resolve as (value: unknown) => void, + reject, + }); + + // Timeout — kills the worker for safety (infinite loop protection) + const timer = setTimeout(() => { + const pending = this._pendingRequests.get(requestId); + if (pending) { + this._pendingRequests.delete(requestId); + this.killWorker(); + pending.reject( + new Error( + l10n.t('Operation timed out after {0} seconds', String(Math.round(timeoutMs / 1000))), + ), + ); + } + }, timeoutMs); + + // Store timer reference on the pending entry so we can clear it + const entry = this._pendingRequests.get(requestId)!; + const originalResolve = entry.resolve; + const originalReject = entry.reject; + entry.resolve = (value: unknown) => { + clearTimeout(timer); + originalResolve(value); + }; + entry.reject = (error: Error) => { + clearTimeout(timer); + originalReject(error); + }; + + this._worker!.postMessage(msgWithId); + }); + } + + /** + * Handle an incoming message from the worker. + */ + private handleWorkerMessage(msg: WorkerToMainMessage): void { + switch (msg.type) { + case 'initResult': { + const pending = this._pendingRequests.get(msg.requestId); + if (pending) { + this._pendingRequests.delete(msg.requestId); + if (msg.success) { + pending.resolve(undefined); + } else { + pending.reject(new Error(msg.error ?? 'Worker init failed')); + } + } + break; + } + + case 'evalResult': { + const pending = this._pendingRequests.get(msg.requestId); + if (pending) { + this._pendingRequests.delete(msg.requestId); + pending.resolve(msg); + } + break; + } + + case 'evalError': { + const pending = this._pendingRequests.get(msg.requestId); + if (pending) { + this._pendingRequests.delete(msg.requestId); + const error = new Error(msg.error); + if (msg.stack) { + error.stack = msg.stack; + } + pending.reject(error); + } + break; + } + + case 'shutdownComplete': { + const pending = this._pendingRequests.get(msg.requestId); + if (pending) { + this._pendingRequests.delete(msg.requestId); + pending.resolve(undefined); + } + break; + } + + case 'tokenRequest': { + // Entra ID: worker needs an OIDC token — delegate to main thread VS Code API + void this.handleTokenRequest(msg); + break; + } + + case 'log': { + const prefix = '[Scratchpad Worker]'; + switch (msg.level) { + case 'error': + ext.outputChannel.error(`${prefix} ${msg.message}`); + break; + case 'warn': + ext.outputChannel.warn(`${prefix} ${msg.message}`); + break; + case 'debug': + ext.outputChannel.debug(`${prefix} ${msg.message}`); + break; + default: + ext.outputChannel.trace(`${prefix} ${msg.message}`); + break; + } + break; + } + } + } + + /** + * Handle a token request from the worker (Entra ID OIDC). + * Calls VS Code's auth API on the main thread and sends the token back. + */ + private async handleTokenRequest(msg: Extract): Promise { + try { + const { getSessionFromVSCode } = await import( + // eslint-disable-next-line import/no-internal-modules + '@microsoft/vscode-azext-azureauth/out/src/getSessionFromVSCode' + ); + const session = await getSessionFromVSCode(msg.scopes as string[], msg.tenantId, { createIfNone: true }); + + if (!session) { + throw new Error('Failed to obtain Entra ID session'); + } + + const response: MainToWorkerMessage = { + type: 'tokenResponse', + requestId: msg.requestId, + accessToken: session.accessToken, + }; + this._worker?.postMessage(response); + } catch (error: unknown) { + const errorMessage = error instanceof Error ? error.message : String(error); + const response: MainToWorkerMessage = { + type: 'tokenError', + requestId: msg.requestId, + error: errorMessage, + }; + this._worker?.postMessage(response); + } + } + + // ─── Private: Worker cleanup ───────────────────────────────────────────── + + private terminateWorker(): void { + if (this._worker) { + void this._worker.terminate(); + this._worker = undefined; + } + this._workerState = 'idle'; + this._workerClusterId = undefined; + this._sessionId = undefined; + this._sessionEvalCount = 0; + this._sessionAuthMethod = undefined; + + // Reject all pending requests + for (const [, entry] of this._pendingRequests) { + entry.reject(new Error('Worker terminated')); + } + this._pendingRequests.clear(); + } + + private handleWorkerExit(): void { + this._worker = undefined; + this._workerState = 'idle'; + this._workerClusterId = undefined; + this._sessionId = undefined; + this._sessionEvalCount = 0; + this._sessionAuthMethod = undefined; + + // Reject any still-pending requests + for (const [, entry] of this._pendingRequests) { + entry.reject(new Error('Worker exited unexpectedly')); + } + this._pendingRequests.clear(); + } + + // ─── Help command interception ─────────────────────────────────────────── + /** * Handle `help` command with scratchpad-specific output. * Returns undefined if the input is not a help command. @@ -172,22 +577,10 @@ export class ScratchpadEvaluator { 'Tips:', ' • Separate code blocks with blank lines', ' • Variables persist within a block but not between separate runs', + ' • When running multiple statements, only the last result is shown', ' • Use .toArray() to get all results (default: first 20 documents)', ].join('\n'); return { type: 'Help', printable: helpText }; } } - -/** - * The eval function passed to `ShellEvaluator.customEval()`. - * Called by @mongosh with (rewrittenCode, context, filename). - * - * Uses `vm.runInContext` so that all properties on the context object - * (including getters like `db`) are accessible as globals in the code. - */ -// eslint-disable-next-line @typescript-eslint/require-await -async function customEvalFn(code: string, context: object): Promise { - const vmContext = vm.isContext(context) ? context : vm.createContext(context); - return vm.runInContext(code, vmContext) as unknown; -} diff --git a/src/documentdb/scratchpad/resultFormatter.ts b/src/documentdb/scratchpad/resultFormatter.ts index 7d270e98..2b21e368 100644 --- a/src/documentdb/scratchpad/resultFormatter.ts +++ b/src/documentdb/scratchpad/resultFormatter.ts @@ -30,21 +30,25 @@ export function formatResult(result: ExecutionResult, code: string, connection: } // Result metadata - if (result.type) { - const printable = result.printable; - if (Array.isArray(printable)) { - lines.push(`// ${l10n.t('{0} documents returned', printable.length)}`); - } else { - lines.push(`// ${l10n.t('Result: {0}', result.type)}`); - } + // Result metadata — state what we know from @mongosh, don't guess + const unwrapped = unwrapCursorResult(result.printable); + if (result.type === 'Cursor' && Array.isArray(unwrapped)) { + // Cursor with a known batch: "Result: Cursor (20 documents)" + lines.push(`// ${l10n.t('Result: Cursor ({0} documents)', unwrapped.length)}`); + } else if (result.type) { + // Typed result: "Result: Document", "Result: string", etc. + lines.push(`// ${l10n.t('Result: {0}', result.type)}`); + } else if (Array.isArray(unwrapped)) { + // Untyped array (e.g. .toArray()) — show type and count + lines.push(`// ${l10n.t('Result: Array ({0} elements)', unwrapped.length)}`); } lines.push(`// ${l10n.t('Executed in {0}ms', result.durationMs)}`); lines.push('// ─────────────────────────'); lines.push(''); - // Result value - lines.push(formatPrintable(result.printable)); + // Result value — unwrap cursor wrapper for clean output + lines.push(formatPrintable(unwrapCursorResult(result.printable))); return lines.join('\n'); } @@ -81,6 +85,30 @@ export function formatError( return lines.join('\n'); } +/** + * Unwrap CursorIterationResult from @mongosh. + * + * @mongosh's `asPrintable()` on CursorIterationResult produces + * `{ cursorHasMore: boolean, documents: unknown[] }` instead of a plain array. + * Only unwraps when the full wrapper shape is present to avoid + * false positives on user documents that happen to have a `documents` field. + */ +function unwrapCursorResult(printable: unknown): unknown { + if ( + printable !== null && + printable !== undefined && + typeof printable === 'object' && + !Array.isArray(printable) && + 'cursorHasMore' in printable && + typeof (printable as Record).cursorHasMore === 'boolean' && + 'documents' in printable && + Array.isArray((printable as { documents: unknown }).documents) + ) { + return (printable as { documents: unknown[] }).documents; + } + return printable; +} + function formatPrintable(printable: unknown): string { if (printable === undefined) { return 'undefined'; diff --git a/src/documentdb/scratchpad/scratchpadWorker.ts b/src/documentdb/scratchpad/scratchpadWorker.ts new file mode 100644 index 00000000..b33dc08e --- /dev/null +++ b/src/documentdb/scratchpad/scratchpadWorker.ts @@ -0,0 +1,307 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +/** + * Worker thread entry point for scratchpad code evaluation. + * + * This file runs in a Node.js worker_thread, isolated from the extension host. + * It owns its own database client instance (authenticated via credentials passed from + * the main thread at init time) and evaluates user code through the @mongosh pipeline. + * + * Communication with the main thread is via postMessage (structured clone). + * See workerTypes.ts for the message protocol. + */ + +import { randomUUID } from 'crypto'; +import { type MongoClientOptions, type MongoClient as MongoClientType } from 'mongodb'; +import { parentPort } from 'worker_threads'; +import { type MainToWorkerMessage, type WorkerToMainMessage } from './workerTypes'; + +if (!parentPort) { + throw new Error('scratchpadWorker.ts must be run as a worker_thread'); +} + +// ─── Worker state ──────────────────────────────────────────────────────────── + +let mongoClient: MongoClientType | undefined; +let currentDatabaseName: string | undefined; + +/** + * Cache for pending Entra ID token requests from the OIDC callback. + * The OIDC_CALLBACK in the worker sends a tokenRequest to the main thread + * and awaits the response via this map. + */ +const pendingTokenRequests = new Map void; reject: (err: Error) => void }>(); + +// ─── Logging helper ────────────────────────────────────────────────────────── + +function log(level: 'trace' | 'debug' | 'info' | 'warn' | 'error', message: string): void { + const msg: WorkerToMainMessage = { type: 'log', level, message }; + parentPort!.postMessage(msg); +} + +// ─── Message handler ───────────────────────────────────────────────────────── + +parentPort.on('message', (msg: MainToWorkerMessage) => { + switch (msg.type) { + case 'init': + void handleInit(msg).catch((err: unknown) => { + const errorMessage = err instanceof Error ? err.message : String(err); + const response: WorkerToMainMessage = { + type: 'initResult', + requestId: msg.requestId, + success: false, + error: errorMessage, + }; + parentPort!.postMessage(response); + }); + break; + + case 'eval': + void handleEval(msg).catch((err: unknown) => { + const errorMessage = err instanceof Error ? err.message : String(err); + const stack = err instanceof Error ? err.stack : undefined; + const response: WorkerToMainMessage = { + type: 'evalError', + requestId: msg.requestId, + error: errorMessage, + stack, + }; + parentPort!.postMessage(response); + }); + break; + + case 'shutdown': + void handleShutdown(msg); + break; + + case 'tokenResponse': + handleTokenResponse(msg); + break; + + case 'tokenError': + handleTokenError(msg); + break; + } +}); + +// ─── Init handler ──────────────────────────────────────────────────────────── + +async function handleInit(msg: Extract): Promise { + log('debug', `Initializing worker (auth: ${msg.authMechanism}, db: ${msg.databaseName})`); + + // Lazy-import MongoDB driver + const { MongoClient } = await import('mongodb'); + + // Build client options from the serializable subset + const options: MongoClientOptions = { + ...msg.clientOptions, + }; + + // For Entra ID, configure OIDC callback that requests tokens via IPC + if (msg.authMechanism === 'MicrosoftEntraID') { + options.authMechanism = 'MONGODB-OIDC'; + options.tls = true; + options.authMechanismProperties = { + ALLOWED_HOSTS: ['*.azure.com'], + OIDC_CALLBACK: async (): Promise<{ accessToken: string; expiresInSeconds: number }> => { + const requestId = randomUUID(); + const tokenPromise = new Promise((resolve, reject) => { + pendingTokenRequests.set(requestId, { resolve, reject }); + }); + const tokenRequest: WorkerToMainMessage = { + type: 'tokenRequest', + requestId, + scopes: ['https://ossrdbms-aad.database.windows.net/.default'], + tenantId: msg.tenantId, + }; + parentPort!.postMessage(tokenRequest); + const accessToken = await tokenPromise; + return { accessToken, expiresInSeconds: 0 }; + }, + }; + } + + // Create and connect the database client + mongoClient = new MongoClient(msg.connectionString, options); + await mongoClient.connect(); + currentDatabaseName = msg.databaseName; + + log('debug', 'Worker initialized — client connected'); + + const response: WorkerToMainMessage = { + type: 'initResult', + requestId: msg.requestId, + success: true, + }; + parentPort!.postMessage(response); +} + +// ─── Eval handler ──────────────────────────────────────────────────────────── + +async function handleEval(msg: Extract): Promise { + if (!mongoClient) { + throw new Error('Worker not initialized — call init first'); + } + + const lineCount = msg.code.split('\n').length; + log( + 'trace', + `Evaluating code (${String(lineCount)} lines, ${String(msg.code.length)} chars, db: ${msg.databaseName})`, + ); + + // Lazy-import @mongosh packages + const { EventEmitter } = await import('events'); + const vm = await import('vm'); + const { NodeDriverServiceProvider } = await import('@mongosh/service-provider-node-driver'); + const { ShellInstanceState } = await import('@mongosh/shell-api'); + const { ShellEvaluator } = await import('@mongosh/shell-evaluator'); + + const startTime = Date.now(); + + // Create fresh shell context per execution (no variable leakage between runs) + const bus = new EventEmitter(); + const serviceProvider = new NodeDriverServiceProvider(mongoClient, bus, { + productDocsLink: 'https://github.com/microsoft/vscode-documentdb', + productName: 'DocumentDB for VS Code Scratchpad', + }); + const instanceState = new ShellInstanceState(serviceProvider, bus); + const evaluator = new ShellEvaluator(instanceState); + + // Set up eval context with shell globals (db, ObjectId, ISODate, etc.) + const context = {}; + instanceState.setCtx(context); + + // The eval function using vm.runInContext for @mongosh + // eslint-disable-next-line @typescript-eslint/require-await + const customEvalFn = async (code: string, ctx: object): Promise => { + const vmContext = vm.isContext(ctx) ? ctx : vm.createContext(ctx); + return vm.runInContext(code, vmContext) as unknown; + }; + + // Switch database if different from current + if (msg.databaseName !== currentDatabaseName) { + await evaluator.customEval(customEvalFn, `use(${JSON.stringify(msg.databaseName)})`, context, 'scratchpad'); + currentDatabaseName = msg.databaseName; + } else { + // Pre-select the target database (fresh context each time) + await evaluator.customEval(customEvalFn, `use(${JSON.stringify(msg.databaseName)})`, context, 'scratchpad'); + } + + // Evaluate user code + const result = await evaluator.customEval(customEvalFn, msg.code, context, 'scratchpad'); + const durationMs = Date.now() - startTime; + + // result is a ShellResult { type, printable, rawValue, source? } + const shellResult = result as { + type: string | null; + printable: unknown; + source?: { namespace?: { db: string; collection: string } }; + }; + + // Normalize the printable value for IPC transfer. + // @mongosh's ShellEvaluator wraps cursor results as { cursorHasMore, documents } + // when running in a worker context. Extract the documents array so that the + // main thread receives a clean array (matching the in-process behavior where + // printable was a CursorIterationResult array). + let printableValue: unknown = shellResult.printable; + if ( + shellResult.type === 'Cursor' && + typeof shellResult.printable === 'object' && + shellResult.printable !== null && + 'documents' in shellResult.printable && + Array.isArray((shellResult.printable as { documents?: unknown }).documents) + ) { + printableValue = (shellResult.printable as { documents: unknown[] }).documents; + } else if (Array.isArray(shellResult.printable)) { + // Array subclass (CursorIterationResult) — normalize to plain Array + printableValue = Array.from(shellResult.printable as unknown[]); + } + + let printableStr: string; + try { + const { EJSON } = await import('bson'); + printableStr = EJSON.stringify(printableValue, { relaxed: false }); + } catch { + // Fallback: try JSON, then plain string + try { + printableStr = JSON.stringify(printableValue); + } catch { + printableStr = String(printableValue); + } + } + + log('trace', `Evaluation complete (${durationMs}ms, type: ${shellResult.type ?? 'null'})`); + + const response: WorkerToMainMessage = { + type: 'evalResult', + requestId: msg.requestId, + result: { + type: shellResult.type, + printable: printableStr, + durationMs, + source: shellResult.source?.namespace + ? { + namespace: { + db: shellResult.source.namespace.db, + collection: shellResult.source.namespace.collection, + }, + } + : undefined, + }, + }; + parentPort!.postMessage(response); +} + +// ─── Shutdown handler ──────────────────────────────────────────────────────── + +async function handleShutdown(msg: Extract): Promise { + log('debug', 'Shutting down worker — closing client'); + + try { + if (mongoClient) { + await mongoClient.close(); + mongoClient = undefined; + } + } catch (err: unknown) { + const errorMessage = err instanceof Error ? err.message : String(err); + log('warn', `Error closing client during shutdown: ${errorMessage}`); + } + + const response: WorkerToMainMessage = { + type: 'shutdownComplete', + requestId: msg.requestId, + }; + parentPort!.postMessage(response); +} + +// ─── Token response/error handlers (Entra ID) ─────────────────────────────── + +function handleTokenResponse(msg: Extract): void { + const pending = pendingTokenRequests.get(msg.requestId); + if (pending) { + pending.resolve(msg.accessToken); + pendingTokenRequests.delete(msg.requestId); + } +} + +function handleTokenError(msg: Extract): void { + const pending = pendingTokenRequests.get(msg.requestId); + if (pending) { + pending.reject(new Error(msg.error)); + pendingTokenRequests.delete(msg.requestId); + } +} + +// ─── Uncaught exception handler ────────────────────────────────────────────── + +process.on('uncaughtException', (error: Error) => { + log('error', `Uncaught exception in worker: ${error.message}\n${error.stack ?? ''}`); +}); + +process.on('unhandledRejection', (reason: unknown) => { + const message = reason instanceof Error ? reason.message : String(reason); + log('error', `Unhandled rejection in worker: ${message}`); +}); diff --git a/src/documentdb/scratchpad/workerTypes.ts b/src/documentdb/scratchpad/workerTypes.ts new file mode 100644 index 00000000..af13b0e5 --- /dev/null +++ b/src/documentdb/scratchpad/workerTypes.ts @@ -0,0 +1,125 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +/** + * IPC message types for the scratchpad worker thread. + * + * This file is shared between the main thread (ScratchpadEvaluator) and + * the worker thread (scratchpadWorker). It must have zero runtime dependencies — + * only TypeScript types and string literal unions. + * + * Communication uses Node.js `worker_threads` `postMessage()` with the + * structured clone algorithm. Functions cannot be sent — this is why + * Entra ID OIDC tokens must be requested via IPC (tokenRequest/tokenResponse). + */ + +// ─── Serializable subset of MongoClientOptions ────────────────────────────── + +/** + * Only the MongoClientOptions fields that can survive structured clone. + * Function-valued options (like OIDC_CALLBACK) are stripped before sending + * and reconstructed on the worker side. + */ +export interface SerializableMongoClientOptions { + readonly serverSelectionTimeoutMS?: number; + readonly tlsAllowInvalidCertificates?: boolean; + readonly appName?: string; + readonly tls?: boolean; +} + +// ─── Serializable execution result ────────────────────────────────────────── + +/** + * The subset of ExecutionResult that can be sent via postMessage. + * BSON types are serialized to EJSON strings by the worker before sending. + */ +export interface SerializableExecutionResult { + readonly type: string | null; + /** EJSON-serialized printable value */ + readonly printable: string; + readonly durationMs: number; + readonly source?: { + readonly namespace?: { + readonly db: string; + readonly collection: string; + }; + }; +} + +// ─── Main → Worker messages ───────────────────────────────────────────────── + +export type MainToWorkerMessage = + | { + readonly type: 'init'; + readonly requestId: string; + /** Connection string (with embedded credentials for SCRAM, without for Entra ID) */ + readonly connectionString: string; + readonly clientOptions: SerializableMongoClientOptions; + readonly databaseName: string; + readonly authMechanism: 'NativeAuth' | 'MicrosoftEntraID'; + /** Tenant ID for Entra ID clusters */ + readonly tenantId?: string; + // TODO(F11): Wire displayBatchSize end-to-end — currently sent but not read by the worker. + // See future-work.md §F11 for the plan to honor documentDB.mongoShell.batchSize. + readonly displayBatchSize: number; + } + | { + readonly type: 'eval'; + readonly requestId: string; + /** JavaScript code to evaluate */ + readonly code: string; + /** Target database name (may differ from init if user switched databases) */ + readonly databaseName: string; + } + | { + readonly type: 'shutdown'; + readonly requestId: string; + } + | { + readonly type: 'tokenResponse'; + readonly requestId: string; + readonly accessToken: string; + } + | { + readonly type: 'tokenError'; + readonly requestId: string; + readonly error: string; + }; + +// ─── Worker → Main messages ───────────────────────────────────────────────── + +export type WorkerToMainMessage = + | { + readonly type: 'initResult'; + readonly requestId: string; + readonly success: boolean; + readonly error?: string; + } + | { + readonly type: 'evalResult'; + readonly requestId: string; + readonly result: SerializableExecutionResult; + } + | { + readonly type: 'evalError'; + readonly requestId: string; + readonly error: string; + readonly stack?: string; + } + | { + readonly type: 'shutdownComplete'; + readonly requestId: string; + } + | { + readonly type: 'tokenRequest'; + readonly requestId: string; + readonly scopes: readonly string[]; + readonly tenantId?: string; + } + | { + readonly type: 'log'; + readonly level: 'trace' | 'debug' | 'info' | 'warn' | 'error'; + readonly message: string; + }; diff --git a/src/documentdb/utils/getClusterMetadata.ts b/src/documentdb/utils/getClusterMetadata.ts index 36430136..1e433f7a 100644 --- a/src/documentdb/utils/getClusterMetadata.ts +++ b/src/documentdb/utils/getClusterMetadata.ts @@ -108,6 +108,17 @@ async function fetchHostInfo(adminDb: Admin, result: ClusterMetadata): Promise): void { for (const [index, host] of hosts.entries()) { const telemetrySuffix = index > 0 ? `_h${index}` : ''; try { diff --git a/webpack.config.ext.js b/webpack.config.ext.js index 82a70551..de547515 100644 --- a/webpack.config.ext.js +++ b/webpack.config.ext.js @@ -22,6 +22,7 @@ module.exports = (env, { mode }) => { node: { __filename: false, __dirname: false }, entry: { main: './main.ts', + scratchpadWorker: './src/documentdb/scratchpad/scratchpadWorker.ts', }, output: { path: path.resolve(__dirname, 'dist'),