From eba325f0023e2819c3b37a00a9a7329002b2a22b Mon Sep 17 00:00:00 2001
From: Wroud <wrouds@gmail.com>
Date: Sun, 3 Oct 2021 14:41:16 +0300
Subject: [PATCH] fix(plugin-data-viewer): CB-1250 actions update

---
 .../src/ContainerDataSource.ts                | 37 +++++++++-------
 .../DatabaseDataModel/DatabaseDataAction.ts   |  6 +++
 .../DatabaseDataModel/DatabaseDataActions.ts  |  8 ++--
 .../DatabaseDataModel/DatabaseDataSource.ts   |  1 +
 .../DatabaseDataModel/IDatabaseDataResult.ts  |  3 ++
 .../plugin-sql-editor/src/QueryDataSource.ts  | 42 ++++++++++++-------
 6 files changed, 65 insertions(+), 32 deletions(-)

diff --git a/webapp/packages/plugin-data-viewer/src/ContainerDataSource.ts b/webapp/packages/plugin-data-viewer/src/ContainerDataSource.ts
index 00c46544bb..a0fc9b3edc 100644
--- a/webapp/packages/plugin-data-viewer/src/ContainerDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/ContainerDataSource.ts
@@ -67,7 +67,11 @@ export class ContainerDataSource extends DatabaseDataSource<IDataContainerOption
 
     let firstResultId: string | undefined;
 
-    if (prevResults.length === 1) {
+    if (
+      prevResults.length === 1
+      && prevResults[0].contextId === executionContext.context!.id
+      && prevResults[0].connectionId === executionContext.context?.connectionId
+    ) {
       firstResultId = prevResults[0].id;
     }
 
@@ -107,9 +111,9 @@ export class ContainerDataSource extends DatabaseDataSource<IDataContainerOption
         return prevResults;
       }
 
-      await this.closeResults(executionContext.context!, prevResults);
+      await this.closeResults(prevResults);
 
-      return this.transformResults(response.results, limit);
+      return this.transformResults(executionContext.context!, response.results, limit);
     } catch (exception) {
       this.error = exception;
       throw exception;
@@ -123,9 +127,10 @@ export class ContainerDataSource extends DatabaseDataSource<IDataContainerOption
 
     try {
       for (const result of prevResults) {
+        const executionContextInfo = executionContext.context!;
         const updateVariables: UpdateResultsDataBatchMutationVariables = {
-          connectionId: executionContext.context!.connectionId,
-          contextId: executionContext.context!.id,
+          connectionId: executionContextInfo.connectionId,
+          contextId: executionContextInfo.id,
           resultsId: result.id,
         };
         let editor: ResultSetEditAction | DocumentEditAction | undefined;
@@ -141,7 +146,7 @@ export class ContainerDataSource extends DatabaseDataSource<IDataContainerOption
         const response = await this.graphQLService.sdk.updateResultsDataBatch(updateVariables);
 
         if (editor) {
-          const responseResult = this.transformResults(response.result.results, 0)
+          const responseResult = this.transformResults(executionContextInfo, response.result.results, 0)
             .find(newResult => newResult.id === result.id);
 
           if (responseResult) {
@@ -166,19 +171,16 @@ export class ContainerDataSource extends DatabaseDataSource<IDataContainerOption
   }
 
   async dispose(): Promise<void> {
-    if (this.executionContext?.context) {
-      await this.closeResults(this.executionContext.context, this.results);
-    }
-
+    await this.closeResults(this.results);
     await this.executionContext?.destroy();
   }
 
-  private async closeResults(context: IConnectionExecutionContextInfo, results: IDatabaseResultSet[]) {
+  private async closeResults(results: IDatabaseResultSet[]) {
     for (const result of results) {
       try {
         await this.graphQLService.sdk.closeResult({
-          connectionId: context.connectionId,
-          contextId: context.id,
+          connectionId: result.connectionId,
+          contextId: result.contextId,
           resultId: result.id,
         });
       } catch (exception) {
@@ -187,9 +189,16 @@ export class ContainerDataSource extends DatabaseDataSource<IDataContainerOption
     }
   }
 
-  private transformResults(results: SqlQueryResults[], limit: number): IDatabaseResultSet[] {
+  private transformResults(
+    executionContextInfo: IConnectionExecutionContextInfo,
+    results: SqlQueryResults[],
+    limit: number
+  ): IDatabaseResultSet[] {
     return results.map<IDatabaseResultSet>(result => ({
       id: result.resultSet?.id || '0',
+      uniqueResultId: `${executionContextInfo.connectionId}_${executionContextInfo.id}_${result.resultSet?.id || '0'}`,
+      connectionId: executionContextInfo.connectionId,
+      contextId: executionContextInfo.id,
       dataFormat: result.dataFormat!,
       updateRowCount: result.updateRowCount || 0,
       loadedFully: (result.resultSet?.rows?.length || 0) < limit,
diff --git a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataAction.ts b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataAction.ts
index 3e1d512d76..9c8960b09a 100644
--- a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataAction.ts
+++ b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataAction.ts
@@ -6,6 +6,8 @@
  * you may not use this file except in compliance with the License.
  */
 
+import { makeObservable, observable } from 'mobx';
+
 import type { IDatabaseDataAction, IDatabaseDataActionClass, IDatabaseDataActionInterface } from './IDatabaseDataAction';
 import type { IDatabaseDataResult } from './IDatabaseDataResult';
 import type { IDatabaseDataSource } from './IDatabaseDataSource';
@@ -23,6 +25,10 @@ implements IDatabaseDataAction<TOptions, TResult> {
   constructor(source: IDatabaseDataSource<TOptions, TResult>, result: TResult) {
     this.result = result;
     this.source = source;
+
+    makeObservable(this, {
+      result: observable.ref,
+    });
   }
 
   updateResult(result: TResult): void {
diff --git a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataActions.ts b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataActions.ts
index ee70047040..a4816b4c7c 100644
--- a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataActions.ts
+++ b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataActions.ts
@@ -52,7 +52,7 @@ implements IDatabaseDataActions<TOptions, TResult> {
       throw new Error('DataFormat unsupported');
     }
 
-    const actions = this.getOrCreateActionsList(result.id);
+    const actions = this.getOrCreateActionsList(result.uniqueResultId);
 
     return runInAction(() => {
       let action = actions.find(action => action instanceof Action);
@@ -70,7 +70,7 @@ implements IDatabaseDataActions<TOptions, TResult> {
         }
 
         action = new Action(this.source, result, ...depends);
-        this.addActionToList(result.id, actions, action);
+        this.addActionToList(result.uniqueResultId, actions, action);
       }
       return action as T;
     });
@@ -80,7 +80,7 @@ implements IDatabaseDataActions<TOptions, TResult> {
     result: TResult,
     Action: IDatabaseDataActionInterface<TOptions, TResult, T>
   ): T | undefined {
-    const actions = this.getActionsList(result.id);
+    const actions = this.getActionsList(result.uniqueResultId);
     const action = actions?.find(action => action instanceof Action);
 
     return action as T | undefined;
@@ -90,7 +90,7 @@ implements IDatabaseDataActions<TOptions, TResult> {
     const actionsMap = Array.from(this.actions.entries());
 
     for (const [key, actions] of actionsMap) {
-      const result = results.find(result => result.id === key);
+      const result = results.find(result => result.uniqueResultId === key);
 
       if (!result) {
         for (const action of actions) {
diff --git a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
index ad35412c59..c8605c9ba4 100644
--- a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
@@ -76,6 +76,7 @@ implements IDatabaseDataSource<TOptions, TResult> {
       activeRequest: observable.ref,
       activeSave: observable.ref,
       activeTask: observable.ref,
+      setResults: action,
       setSupportedDataFormats: action,
     });
   }
diff --git a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/IDatabaseDataResult.ts b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/IDatabaseDataResult.ts
index 6c9f249f76..1c2798b436 100644
--- a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/IDatabaseDataResult.ts
+++ b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/IDatabaseDataResult.ts
@@ -10,6 +10,9 @@ import type { ResultDataFormat } from '@cloudbeaver/core-sdk';
 
 export interface IDatabaseDataResult {
   id: string;
+  uniqueResultId: string;
+  connectionId: string;
+  contextId: string;
   dataFormat: ResultDataFormat;
   loadedFully: boolean;
   updateRowCount: number;
diff --git a/webapp/packages/plugin-sql-editor/src/QueryDataSource.ts b/webapp/packages/plugin-sql-editor/src/QueryDataSource.ts
index b29c2d3335..915d117139 100644
--- a/webapp/packages/plugin-sql-editor/src/QueryDataSource.ts
+++ b/webapp/packages/plugin-sql-editor/src/QueryDataSource.ts
@@ -62,9 +62,10 @@ export class QueryDataSource extends DatabaseDataSource<IDataQueryOptions, IData
 
     try {
       for (const result of prevResults) {
+        const executionContextInfo = this.executionContext.context;
         const updateVariables: UpdateResultsDataBatchMutationVariables = {
           connectionId: this.options.connectionId,
-          contextId: this.executionContext.context.id,
+          contextId: executionContextInfo.id,
           resultsId: result.id,
         };
         let editor: ResultSetEditAction | DocumentEditAction | undefined;
@@ -80,7 +81,7 @@ export class QueryDataSource extends DatabaseDataSource<IDataQueryOptions, IData
         const response = await this.graphQLService.sdk.updateResultsDataBatch(updateVariables);
 
         if (editor) {
-          const responseResult = this.transformResults(response.result.results, 0)
+          const responseResult = this.transformResults(executionContextInfo, response.result.results, 0)
             .find(newResult => newResult.id === result.id);
 
           if (responseResult) {
@@ -108,7 +109,11 @@ export class QueryDataSource extends DatabaseDataSource<IDataQueryOptions, IData
     return this;
   }
 
-  getResults(response: SqlExecuteInfo, limit: number): IDatabaseResultSet[] | null {
+  getResults(
+    executionContextInfo: IConnectionExecutionContextInfo,
+    response: SqlExecuteInfo,
+    limit: number
+  ): IDatabaseResultSet[] | null {
     this.requestInfo = {
       requestDuration: response.duration || 0,
       requestMessage: response.statusMessage || '',
@@ -120,7 +125,7 @@ export class QueryDataSource extends DatabaseDataSource<IDataQueryOptions, IData
       return null;
     }
 
-    return this.transformResults(response.results, limit);
+    return this.transformResults(executionContextInfo, response.results, limit);
   }
 
   async request(
@@ -139,7 +144,11 @@ export class QueryDataSource extends DatabaseDataSource<IDataQueryOptions, IData
 
     let firstResultId: string | undefined;
 
-    if (prevResults.length === 1) {
+    if (
+      prevResults.length === 1
+      && prevResults[0].contextId === executionContext.context!.id
+      && prevResults[0].connectionId === executionContext.context?.connectionId
+    ) {
       firstResultId = prevResults[0].id;
     }
 
@@ -163,14 +172,14 @@ export class QueryDataSource extends DatabaseDataSource<IDataQueryOptions, IData
     try {
       const response = await this.currentTask;
 
-      const results = this.getResults(response, limit);
+      const results = this.getResults(executionContextInfo, response, limit);
       this.clearError();
 
       if (!results) {
         return prevResults;
       }
 
-      this.closeResults(executionContextInfo, prevResults);
+      this.closeResults(prevResults);
 
       return results;
     } catch (exception) {
@@ -179,12 +188,12 @@ export class QueryDataSource extends DatabaseDataSource<IDataQueryOptions, IData
     }
   }
 
-  private async closeResults(context: IConnectionExecutionContextInfo, results: IDatabaseResultSet[]) {
+  private async closeResults(results: IDatabaseResultSet[]) {
     for (const result of results) {
       try {
         await this.graphQLService.sdk.closeResult({
-          connectionId: context.connectionId,
-          contextId: context.id,
+          connectionId: result.connectionId,
+          contextId: result.contextId,
           resultId: result.id,
         });
       } catch (exception) {
@@ -193,9 +202,16 @@ export class QueryDataSource extends DatabaseDataSource<IDataQueryOptions, IData
     }
   }
 
-  private transformResults(results: SqlQueryResults[], limit: number): IDatabaseResultSet[] {
+  private transformResults(
+    executionContextInfo: IConnectionExecutionContextInfo,
+    results: SqlQueryResults[],
+    limit: number
+  ): IDatabaseResultSet[] {
     return results.map<IDatabaseResultSet>(result => ({
       id: result.resultSet?.id || '0',
+      uniqueResultId: `${executionContextInfo.connectionId}_${executionContextInfo.id}_${result.resultSet?.id || '0'}`,
+      connectionId: executionContextInfo.connectionId,
+      contextId: executionContextInfo.id,
       dataFormat: result.dataFormat!,
       updateRowCount: result.updateRowCount || 0,
       loadedFully: (result.resultSet?.rows?.length || 0) < limit,
@@ -206,9 +222,7 @@ export class QueryDataSource extends DatabaseDataSource<IDataQueryOptions, IData
   }
 
   async dispose(): Promise<void> {
-    if (this.executionContext?.context) {
-      await this.closeResults(this.executionContext.context, this.results);
-    }
+    await this.closeResults(this.results);
     await this.cancel();
   }
 }