CB-4090 fix: catch throttle errors (#2651)

* CB-4090 fix: catch throttle errors

* CB-4090 fix: catch throttle errors

* CB-4090 fix: detect closed contexts

* CB-4090 fix: close results properly

---------

Co-authored-by: Evgenia Bezborodova <139753579+EvgeniaBzzz@users.noreply.github.com>

webapp/packages/plugin-data-viewer/src/ContainerDataSource.ts

@@ -138,8 +138,6 @@ export class ContainerDataSource extends ResultSetDataSource<IDataContainerOptio
 
       this.clearError();
 
-      await this.closeResults(prevResults);
-
       return this.transformResults(executionContext.context!, response.results, limit);
     } catch (exception: any) {
       this.error = exception;

webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataModel.ts

@@ -162,9 +162,9 @@ export class DatabaseDataModel<TOptions, TResult extends IDatabaseDataResult = I
     this.source.resetData();
   }
 
-  async dispose(): Promise<void> {
+  async dispose(keepExecutionContext = false): Promise<void> {
     await this.onDispose.execute();
-    await this.source.dispose();
+    await this.source.dispose(keepExecutionContext);
   }
 
   async requestSaveAction(action: () => Promise<void> | void): Promise<void> {

webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts

@@ -372,9 +372,11 @@ export abstract class DatabaseDataSource<TOptions, TResult extends IDatabaseData
     return this;
   }
 
-  async dispose(): Promise<void> {
+  async dispose(keepExecutionContext = false): Promise<void> {
     await this.cancel();
-    await this.executionContext?.destroy();
+    if (!keepExecutionContext) {
+      await this.executionContext?.destroy();
+    }
   }
 
   abstract request(prevResults: TResult[]): TResult[] | Promise<TResult[]>;

webapp/packages/plugin-data-viewer/src/DatabaseDataModel/IDatabaseDataModel.ts

@@ -57,5 +57,5 @@ export interface IDatabaseDataModel<TOptions = any, TResult extends IDatabaseDat
   requestDataPortion: (offset: number, count: number) => Promise<void>;
   cancel: () => Promise<void> | void;
   resetData: () => void;
-  dispose: () => Promise<void>;
+  dispose: (keepExecutionContext?: boolean) => Promise<void>;
 }

webapp/packages/plugin-data-viewer/src/DatabaseDataModel/IDatabaseDataSource.ts

@@ -97,5 +97,5 @@ export interface IDatabaseDataSource<TOptions, TResult extends IDatabaseDataResu
   cancel: () => Promise<void> | void;
   clearError: () => this;
   resetData: () => this;
-  dispose: () => Promise<void>;
+  dispose: (keepExecutionContext?: boolean) => Promise<void>;
 }

webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts

@@ -77,13 +77,26 @@ export abstract class ResultSetDataSource<TOptions> extends DatabaseDataSource<T
     return this.totalCountRequestTask;
   }
 
-  async closeResults(results: IDatabaseResultSet[]): Promise<void> {
+  setResults(results: IDatabaseResultSet[]): this {
+    this.closeResults(this.results);
+    return super.setResults(results);
+  }
+
+  async dispose(keepExecutionContext?: boolean): Promise<void> {
+    if (keepExecutionContext) {
+      await this.closeResults(this.results);
+    }
+    return super.dispose(keepExecutionContext);
+  }
+
+  private async closeResults(results: IDatabaseResultSet[]): Promise<void> {
     if (!this.executionContext?.context) {
       return;
     }
 
     for (const result of results) {
-      if (result.id === null) {
+      // TODO: it's better to track that context is closed with subscription
+      if (result.id === null || result.contextId !== this.executionContext.context.id) {
         continue;
       }
       try {