feat: Deleted documents are highlighted

src/docdb/session/DocumentSession.ts

@@ -19,11 +19,11 @@ import vscode from 'vscode';
 import { ext } from '../../extensionVariables';
 import { type Channel } from '../../panels/Communication/Channel/Channel';
 import { getErrorMessage } from '../../panels/Communication/Channel/CommonChannel';
-import { extractPartitionKey } from '../../utils/document';
+import { extractPartitionKey, getDocumentId } from '../../utils/document';
 import { localize } from '../../utils/localize';
 import { type NoSqlQueryConnection } from '../NoSqlCodeLensProvider';
 import { getCosmosClient, type CosmosDBCredential } from '../getCosmosClient';
-import { type CosmosDbRecord, type CosmosDbRecordIdentifier } from '../types/queryResult';
+import { type CosmosDbRecord, type CosmosDbRecordIdentifier, type QueryResultRecord } from '../types/queryResult';
 
 export class DocumentSession {
     public readonly id: string;
@@ -221,43 +221,60 @@ export class DocumentSession {
         });
     }
 
-    public async delete(documentId: CosmosDbRecordIdentifier): Promise<void> {
-        await callWithTelemetryAndErrorHandling('cosmosDB.nosql.document.session.delete', async (context) => {
-            this.setTelemetryProperties(context);
-
-            if (this.isDisposed) {
-                throw new Error('Session is disposed');
-            }
+    // Returns true if document was deleted, false if not, undefined if exception occurred
+    public async delete(documentId: CosmosDbRecordIdentifier): Promise<boolean | undefined> {
+        return callWithTelemetryAndErrorHandling<boolean | undefined>(
+            'cosmosDB.nosql.document.session.delete',
+            async (context) => {
+                this.setTelemetryProperties(context);
 
-            if (documentId.id === undefined) {
-                throw new Error('Document id is required');
-            }
+                if (this.isDisposed) {
+                    throw new Error('Session is disposed');
+                }
 
-            try {
-                const result = await this.client
-                    .database(this.databaseId)
-                    .container(this.containerId)
-                    .item(documentId.id, documentId.partitionKey)
-                    .delete({
-                        abortSignal: this.abortController.signal,
-                    });
+                if (documentId.id === undefined) {
+                    throw new Error('Document id is required');
+                }
 
-                if (result?.statusCode === 204) {
-                    await this.channel.postMessage({
-                        type: 'event',
-                        name: 'documentDeleted',
-                        params: [this.id, documentId],
-                    });
-                } else {
-                    await this.channel.postMessage({
-                        type: 'event',
-                        name: 'documentError',
-                        params: [this.id, 'Document deletion failed'],
-                    });
+                try {
+                    const result = await this.client
+                        .database(this.databaseId)
+                        .container(this.containerId)
+                        .item(documentId.id, documentId.partitionKey)
+                        .delete({
+                            abortSignal: this.abortController.signal,
+                        });
+
+                    if (result?.statusCode === 204) {
+                        await this.channel.postMessage({
+                            type: 'event',
+                            name: 'documentDeleted',
+                            params: [this.id, documentId],
+                        });
+
+                        return true;
+                    } else {
+                        await this.channel.postMessage({
+                            type: 'event',
+                            name: 'documentError',
+                            params: [this.id, 'Document deletion failed'],
+                        });
+
+                        return false;
+                    }
+                } catch (error) {
+                    await this.errorHandling(error, context);
                 }
-            } catch (error) {
-                await this.errorHandling(error, context);
-            }
+
+                return undefined;
+            },
+        );
+    }
+
+    public async getDocumentId(document: QueryResultRecord): Promise<CosmosDbRecordIdentifier | undefined> {
+        return callWithTelemetryAndErrorHandling('cosmosDB.nosql.document.session.getDocumentId', async () => {
+            const partitionKey = await this.getPartitionKey();
+            return getDocumentId(document, partitionKey);
         });
     }
 

src/docdb/session/QuerySession.ts

@@ -20,6 +20,7 @@ import {
     type QueryResultRecord,
     type ResultViewMetadata,
 } from '../types/queryResult';
+import { DocumentSession } from './DocumentSession';
 import { QuerySessionResult } from './QuerySessionResult';
 
 export class QuerySession {
@@ -234,6 +235,51 @@ export class QuerySession {
         });
     }
 
+    public getDocumentId(row: number): Promise<QueryResultRecord | undefined> {
+        const result = this.sessionResult.getResult(this.currentIteration);
+        if (!result) {
+            throw new Error('No result found for current iteration');
+        }
+
+        const document = result.documents[row];
+        if (!document) {
+            throw new Error(`No document found for row: ${row}`);
+        }
+
+        const session = new DocumentSession(this.connection, this.channel);
+        return session.getDocumentId(document);
+    }
+
+    public async deleteDocument(row: number): Promise<void> {
+        const result = this.sessionResult.getResult(this.currentIteration);
+        if (!result) {
+            throw new Error('No result found for current iteration');
+        }
+
+        const document = result.documents[row];
+        if (!document) {
+            throw new Error(`No document found for row: ${row}`);
+        }
+
+        const session = new DocumentSession(this.connection, this.channel);
+        const documentId = await session.getDocumentId(document);
+
+        if (!documentId) {
+            throw new Error('Document id not found');
+        }
+
+        const isDeleted = await session.delete(documentId);
+        if (isDeleted) {
+            result.deletedDocuments.push(row);
+
+            await this.channel.postMessage({
+                type: 'event',
+                name: 'queryResults',
+                params: [this.id, this.sessionResult.getSerializedResult(this.currentIteration), this.currentIteration],
+            });
+        }
+    }
+
     public dispose(): void {
         this.isDisposed = true;
         this.abortController?.abort();

src/docdb/session/QuerySessionResult.ts

@@ -63,6 +63,7 @@ export class QuerySessionResult {
             requestCharge: response.requestCharge,
             roundTrips: 1, // TODO: Is it required field? Query Pages Until Content Present
             hasMoreResults: response.hasMoreResults,
+            deletedDocuments: [],
         });
         this.hasMoreResults = response.hasMoreResults;
     }
@@ -84,6 +85,7 @@ export class QuerySessionResult {
                 requestCharge: result.requestCharge,
                 roundTrips: result.roundTrips,
                 hasMoreResults: result.hasMoreResults,
+                deletedDocuments: result.deletedDocuments,
 
                 query: this.query,
             };

src/docdb/types/queryResult.ts

@@ -51,6 +51,7 @@ export type QueryResult = {
     requestCharge: number;
     roundTrips: number;
     hasMoreResults: boolean;
+    deletedDocuments: number[];
 };
 
 export type SerializedQueryMetrics = {
@@ -81,6 +82,7 @@ export type SerializedQueryResult = {
     requestCharge: number;
     roundTrips: number;
     hasMoreResults: boolean;
+    deletedDocuments: number[];
 
     query: string; // The query that was executed
 };

src/panels/QueryEditorTab.ts

@@ -10,9 +10,8 @@ import * as vscode from 'vscode';
 import { getNoSqlQueryConnection } from '../docdb/commands/connectNoSqlContainer';
 import { getCosmosClientByConnection } from '../docdb/getCosmosClient';
 import { type NoSqlQueryConnection } from '../docdb/NoSqlCodeLensProvider';
-import { DocumentSession } from '../docdb/session/DocumentSession';
 import { QuerySession } from '../docdb/session/QuerySession';
-import { type CosmosDbRecordIdentifier, type ResultViewMetadata } from '../docdb/types/queryResult';
+import { type ResultViewMetadata } from '../docdb/types/queryResult';
 import * as vscodeUtil from '../utils/vscodeUtils';
 import { BaseTab, type CommandPayload } from './BaseTab';
 import { DocumentTab } from './DocumentTab';
@@ -126,9 +125,13 @@ export class QueryEditorTab extends BaseTab {
             case 'firstPage':
                 return this.firstPage(payload.params[0] as string);
             case 'openDocument':
-                return this.openDocument(payload.params[0] as string, payload.params[1] as CosmosDbRecordIdentifier);
+                return this.openDocument(
+                    payload.params[0] as string,
+                    payload.params[1] as string,
+                    payload.params[2] as number,
+                );
             case 'deleteDocument':
-                return this.deleteDocument(payload.params[0] as CosmosDbRecordIdentifier);
+                return this.deleteDocument(payload.params[0] as string, payload.params[1] as number);
         }
 
         return super.getCommand(payload);
@@ -337,36 +340,43 @@ export class QueryEditorTab extends BaseTab {
         });
     }
 
-    private async openDocument(mode: string, documentId?: CosmosDbRecordIdentifier): Promise<void> {
-        await callWithTelemetryAndErrorHandling('cosmosDB.nosql.queryEditor.openDocument', () => {
+    private async openDocument(executionId: string, mode: string, row?: number): Promise<void> {
+        await callWithTelemetryAndErrorHandling('cosmosDB.nosql.queryEditor.openDocument', async () => {
             if (!this.connection) {
                 throw new Error('No connection');
             }
 
-            if (!documentId && mode !== 'add') {
+            if (!row && mode !== 'add') {
                 throw new Error('Impossible to open a document without an id');
             }
 
             if (mode !== 'edit' && mode !== 'view' && mode !== 'add') {
                 throw new Error(`Invalid mode: ${mode}`);
             }
 
+            const session = this.sessions.get(executionId);
+            if (!session) {
+                throw new Error(`No session found for executionId: ${executionId}`);
+            }
+
+            const documentId = row ? await session.getDocumentId(row) : undefined;
+
             DocumentTab.render(this.connection, mode, documentId, this.getNextViewColumn());
         });
     }
 
-    private async deleteDocument(documentId: CosmosDbRecordIdentifier): Promise<void> {
+    private async deleteDocument(executionId: string, row: number): Promise<void> {
         await callWithTelemetryAndErrorHandling('cosmosDB.nosql.queryEditor.deleteDocument', async () => {
             if (!this.connection) {
                 throw new Error('No connection');
             }
 
-            if (!documentId) {
-                throw new Error('Impossible to open a document without an id');
+            const session = this.sessions.get(executionId);
+            if (!session) {
+                throw new Error(`No session found for executionId: ${executionId}`);
             }
 
-            const session = new DocumentSession(this.connection, this.channel);
-            await session.delete(documentId);
+            await session.deleteDocument(row);
         });
     }
 

src/utils/convertors.ts

@@ -20,6 +20,7 @@ export type TableRecord = Record<string, string> & { __id: string };
 export type TableData = {
     headers: string[];
     dataset: TableRecord[];
+    deletedRows: number[];
 };
 
 /**
@@ -45,20 +46,13 @@ export const queryResultToJSON = (queryResult: SerializedQueryResult | null, sel
         return '';
     }
 
-    if (selection) {
-        const selectedDocs = queryResult.documents
-            .map((doc, index) => {
-                if (!selection.includes(index)) {
-                    return null;
-                }
-                return doc;
-            })
-            .filter((doc) => doc !== null);
-
-        return JSON.stringify(selectedDocs, null, 4);
-    }
+    const notDeletedRows = queryResult.documents
+        .map((_, index) => index)
+        .filter((index) => !queryResult.deletedDocuments.includes(index));
+    const selectedRows = selection ? notDeletedRows.filter((index) => selection.includes(index)) : notDeletedRows;
+    const selectedDocs = queryResult.documents.filter((_, index) => selectedRows.includes(index));
 
-    return JSON.stringify(queryResult.documents, null, 4);
+    return JSON.stringify(selectedDocs, null, 4);
 };
 
 export const queryResultToTree = (
@@ -322,7 +316,7 @@ export const queryResultToTable = (
     reorderColumns?: boolean,
     showServiceColumns?: boolean,
 ): TableData => {
-    let result: TableData = { headers: [], dataset: [] };
+    let result: TableData = { headers: [], dataset: [], deletedRows: [] };
 
     if (!queryResult) {
         return result;
@@ -337,11 +331,13 @@ export const queryResultToTable = (
         result = {
             headers: getTableHeadersWithRecordIdentifyColumns(queryResult.documents, partitionKey),
             dataset: getTableDatasetWithRecordIdentifyColumns(queryResult.documents, partitionKey),
+            deletedRows: queryResult.deletedDocuments,
         };
     } else {
         result = {
             headers: getTableHeaders(queryResult.documents),
             dataset: getTableDataset(queryResult.documents),
+            deletedRows: queryResult.deletedDocuments,
         };
     }
 
@@ -517,9 +513,12 @@ export const queryResultToCsv = (
     const tableView = queryResultToTable(queryResult, partitionKey);
     const headers = tableView.headers.join(',');
 
-    if (selection) {
-        tableView.dataset = tableView.dataset.filter((_, index) => selection.includes(index));
-    }
+    const notDeletedRows = tableView.dataset
+        .map((_, index) => index)
+        .filter((index) => !tableView.deletedRows.includes(index));
+    const selectedRows = selection ? notDeletedRows.filter((index) => selection.includes(index)) : notDeletedRows;
+
+    tableView.dataset = tableView.dataset.filter((_, index) => selectedRows.includes(index));
 
     const rows = tableView.dataset
         .map((row) => {

src/utils/document.ts

@@ -22,7 +22,7 @@ export const extractPartitionKey = (document: ItemDefinition, partitionKey: Part
                 // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
                 interim = interim[prop];
             } else {
-                return null; // It is not correct to return null, in other cases it should exception
+                return null; // It is not correct to return null, in other cases it should be exception
             }
         }
         if (

src/webviews/QueryEditor/ResultPanel/ResultPanelToolbar.tsx

@@ -97,7 +97,11 @@ const ToolbarGroupSave = ({ selectedTab }: ResultToolbarProps) => {
         const filename = `${state.dbName}_${state.collectionName}_${state.currentQueryResult?.activityId ?? 'query'}`;
         if (selectedTab === 'result__tab') {
             await dispatcher.saveToFile(
-                queryResultToCsv(state.currentQueryResult, state.partitionKey),
+                queryResultToCsv(
+                    state.currentQueryResult,
+                    state.partitionKey,
+                    hasSelection ? state.selectedRows : undefined,
+                ),
                 `${filename}_result`,
                 'csv',
             );
@@ -111,7 +115,11 @@ const ToolbarGroupSave = ({ selectedTab }: ResultToolbarProps) => {
     async function onSaveAsJSON() {
         const filename = `${state.dbName}_${state.collectionName}_${state.currentQueryResult?.activityId ?? 'query'}`;
         if (selectedTab === 'result__tab') {
-            await dispatcher.saveToFile(queryResultToJSON(state.currentQueryResult), `${filename}_result`, 'json');
+            await dispatcher.saveToFile(
+                queryResultToJSON(state.currentQueryResult, hasSelection ? state.selectedRows : undefined),
+                `${filename}_result`,
+                'json',
+            );
         }
 
         if (selectedTab === 'stats__tab') {