feat: Deleted documents are highlighted

src/cosmosdb/session/DocumentSession.ts

@@ -25,10 +25,10 @@ import { ext } from '../../extensionVariables';
 import { type Channel } from '../../panels/Communication/Channel/Channel';
 import { getErrorMessage } from '../../panels/Communication/Channel/CommonChannel';
 import { getConfirmationAsInSettings } from '../../utils/dialogs/getConfirmation';
-import { arePartitionKeysEqual, extractPartitionKey } from '../../utils/document';
+import { arePartitionKeysEqual, extractPartitionKey, getDocumentId } from '../../utils/document';
 import { getCosmosDBKeyCredential } from '../CosmosDBCredential';
 import { type NoSqlQueryConnection } from '../NoSqlQueryConnection';
-import { type CosmosDBRecord, type CosmosDBRecordIdentifier } from '../types/queryResult';
+import { type CosmosDBRecord, type CosmosDBRecordIdentifier, type QueryResultRecord } from '../types/queryResult';
 import { withClaimsChallengeHandling } from '../withClaimsChallengeHandling';
 
 /**
@@ -318,8 +318,8 @@ export class DocumentSession {
         });
     }
 
-    public async delete(documentId: CosmosDBRecordIdentifier): Promise<void> {
-        await callWithTelemetryAndErrorHandling('cosmosDB.nosql.document.session.delete', async (context) => {
+    public async delete(documentId: CosmosDBRecordIdentifier): Promise<boolean | undefined> {
+        return callWithTelemetryAndErrorHandling('cosmosDB.nosql.document.session.delete', async (context) => {
             this.setTelemetryProperties(context);
 
             if (this.isDisposed) {
@@ -340,7 +340,14 @@ export class DocumentSession {
                 return;
             }
 
-            await this.deleteInternal(documentId, context);
+            return this.deleteInternal(documentId, context);
+        });
+    }
+
+    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);
         });
     }
 
@@ -384,8 +391,8 @@ export class DocumentSession {
         return false;
     }
 
-    public async bulkDelete(documentIds: CosmosDBRecordIdentifier[]): Promise<void> {
-        await callWithTelemetryAndErrorHandling('cosmosDB.nosql.document.session.bulkDelete', async (context) => {
+    public async bulkDelete(documentIds: CosmosDBRecordIdentifier[]): Promise<DeleteStatus | void> {
+        return callWithTelemetryAndErrorHandling('cosmosDB.nosql.document.session.bulkDelete', async (context) => {
             this.setTelemetryProperties(context);
 
             if (this.isDisposed) {
@@ -437,9 +444,8 @@ export class DocumentSession {
                         const abortController = new AbortController();
                         const abortSignal = abortController.signal;
 
-                        token.onCancellationRequested(async () => {
+                        token.onCancellationRequested(() => {
                             status.aborted = true;
-                            await sendResponse();
                             abortController.abort();
                         });
 
@@ -487,7 +493,9 @@ export class DocumentSession {
                     );
                 }
 
-                return sendResponse();
+                await sendResponse();
+
+                return status;
             } catch (error) {
                 await this.errorHandling(error, context);
             }

src/cosmosdb/session/QuerySession.ts

@@ -18,9 +18,11 @@ import { type NoSqlQueryConnection } from '../NoSqlQueryConnection';
 import {
     DEFAULT_EXECUTION_TIMEOUT,
     DEFAULT_PAGE_SIZE,
+    type CosmosDBRecordIdentifier,
     type QueryMetadata,
     type QueryResultRecord,
 } from '../types/queryResult';
+import { DocumentSession } from './DocumentSession';
 import { QuerySessionResult } from './QuerySessionResult';
 
 export class QuerySession {
@@ -236,6 +238,144 @@ export class QuerySession {
         });
     }
 
+    public async getDocumentId(row: number): Promise<QueryResultRecord | undefined> {
+        return callWithTelemetryAndErrorHandling(
+            'cosmosDB.nosql.queryEditor.session.getDocumentId',
+            async (context) => {
+                this.setTelemetryProperties(context);
+
+                if (this.isDisposed) {
+                    throw new Error(l10n.t('Session is disposed'));
+                }
+
+                if (!this.iterator) {
+                    throw new Error(l10n.t('Session is not running! Please run the session first'));
+                }
+
+                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> {
+        return callWithTelemetryAndErrorHandling(
+            'cosmosDB.nosql.queryEditor.session.deleteDocument',
+            async (context) => {
+                this.setTelemetryProperties(context);
+
+                if (this.isDisposed) {
+                    throw new Error(l10n.t('Session is disposed'));
+                }
+
+                if (!this.iterator) {
+                    throw new Error(l10n.t('Session is not running! Please run the session first'));
+                }
+
+                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 async bulkDelete(rows: number[]): Promise<void> {
+        return callWithTelemetryAndErrorHandling('cosmosDB.nosql.queryEditor.session.bulkDelete', async (context) => {
+            this.setTelemetryProperties(context);
+
+            if (this.isDisposed) {
+                throw new Error(l10n.t('Session is disposed'));
+            }
+
+            if (!this.iterator) {
+                throw new Error(l10n.t('Session is not running! Please run the session first'));
+            }
+
+            const result = this.sessionResult.getResult(this.currentIteration);
+            if (!result) {
+                throw new Error('No result found for current iteration');
+            }
+
+            const session = new DocumentSession(this.connection, this.channel);
+            const documents: Array<{ documentId: CosmosDBRecordIdentifier; row: number }> = [];
+
+            for (const row of rows) {
+                const document = result.documents[row];
+                if (!document) {
+                    throw new Error(`No document found for row: ${row}`);
+                }
+
+                const documentId = await session.getDocumentId(document);
+                if (!documentId) {
+                    throw new Error('Document id not found');
+                }
+
+                documents.push({ documentId, row });
+            }
+
+            const documentsToDelete = documents.map((doc) => doc.documentId);
+            const status = await session.bulkDelete(documentsToDelete);
+
+            if (status && status.deleted.length > 0) {
+                result.deletedDocuments = [
+                    ...result.deletedDocuments,
+                    ...status.deleted.map((documentId) => {
+                        const doc = documents.find((document) => document.documentId === documentId);
+                        return doc ? doc.row : -1;
+                    }),
+                ].filter((row) => row !== -1);
+
+                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/cosmosdb/session/QuerySessionResult.ts

@@ -64,6 +64,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;
     }
@@ -85,6 +86,7 @@ export class QuerySessionResult {
                 requestCharge: result.requestCharge,
                 roundTrips: result.roundTrips,
                 hasMoreResults: result.hasMoreResults,
+                deletedDocuments: result.deletedDocuments,
 
                 query: this.query,
             };

src/cosmosdb/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

@@ -12,13 +12,8 @@ import { getThemedIconPath } from '../constants';
 import { getCosmosDBKeyCredential } from '../cosmosdb/CosmosDBCredential';
 import { getCosmosClient } from '../cosmosdb/getCosmosClient';
 import { getNoSqlQueryConnection, type NoSqlQueryConnection } from '../cosmosdb/NoSqlQueryConnection';
-import { DocumentSession } from '../cosmosdb/session/DocumentSession';
 import { QuerySession } from '../cosmosdb/session/QuerySession';
-import {
-    type CosmosDBRecordIdentifier,
-    type QueryMetadata,
-    type SerializedQueryResult,
-} from '../cosmosdb/types/queryResult';
+import { type QueryMetadata, type SerializedQueryResult } from '../cosmosdb/types/queryResult';
 import { withClaimsChallengeHandling } from '../cosmosdb/withClaimsChallengeHandling';
 import { StorageNames, StorageService, type StorageItem } from '../services/storageService';
 import { toStringUniversal } from '../utils/convertors';
@@ -163,11 +158,15 @@ 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);
             case 'deleteDocuments':
-                return this.deleteDocuments(payload.params[0] as CosmosDBRecordIdentifier[]);
+                return this.deleteDocuments(payload.params[0] as string, payload.params[1] as number[]);
             case 'provideFeedback':
                 return this.provideFeedback();
             case 'saveCSV':
@@ -565,52 +564,67 @@ export class QueryEditorTab extends BaseTab {
         void promptAfterActionEventually(ExperienceKind.NoSQL, UsageImpact.Medium, callbackId);
     }
 
-    private async openDocument(mode: string, documentId?: CosmosDBRecordIdentifier): Promise<void> {
+    private async openDocument(executionId: string, mode: string, row?: number): Promise<void> {
         const callbackId = 'cosmosDB.nosql.queryEditor.openDocument';
-        await callWithTelemetryAndErrorHandling(callbackId, () => {
+        await callWithTelemetryAndErrorHandling(callbackId, async () => {
             if (!this.connection) {
                 throw new Error(l10n.t('No connection'));
             }
 
-            if (!documentId && mode !== 'add') {
+            if (!row && mode !== 'add') {
                 throw new Error(l10n.t('Impossible to open an item without an id'));
             }
 
             if (mode !== 'edit' && mode !== 'view' && mode !== 'add') {
                 throw new Error(l10n.t('Invalid mode: {0}', 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());
         });
         void promptAfterActionEventually(ExperienceKind.NoSQL, UsageImpact.Medium, callbackId);
     }
 
-    private async deleteDocument(documentId: CosmosDBRecordIdentifier): Promise<void> {
+    private async deleteDocument(executionId: string, row: number): Promise<void> {
         const callbackId = 'cosmosDB.nosql.queryEditor.deleteDocument';
         await callWithTelemetryAndErrorHandling(callbackId, async () => {
             if (!this.connection) {
                 throw new Error(l10n.t('No connection'));
             }
 
-            if (!documentId) {
+            const session = this.sessions.get(executionId);
+            if (!session) {
+                throw new Error(`No session found for executionId: ${executionId}`);
+            }
+
+            if (!row) {
                 throw new Error(l10n.t('Impossible to delete an item without an id'));
             }
 
-            const session = new DocumentSession(this.connection, this.channel);
-            await session.delete(documentId);
+            return session.deleteDocument(row);
         });
         void promptAfterActionEventually(ExperienceKind.NoSQL, UsageImpact.Medium, callbackId);
     }
 
-    private async deleteDocuments(documentIds: CosmosDBRecordIdentifier[]): Promise<void> {
+    private async deleteDocuments(executionId: string, rows: number[]): Promise<void> {
         const callbackId = 'cosmosDB.nosql.queryEditor.deleteDocuments';
         await callWithTelemetryAndErrorHandling(callbackId, async () => {
             if (!this.connection) {
                 throw new Error(l10n.t('No connection'));
             }
 
-            const session = new DocumentSession(this.connection, this.channel);
-            await session.bulkDelete(documentIds);
+            const session = this.sessions.get(executionId);
+            if (!session) {
+                throw new Error(`No session found for executionId: ${executionId}`);
+            }
+
+            return session.bulkDelete(rows);
         });
         void promptAfterActionEventually(ExperienceKind.NoSQL, UsageImpact.Medium, callbackId);
     }