From d4b5296aa0fada5a8fb554a3bd360547829f49c1 Mon Sep 17 00:00:00 2001
From: sergeyteleshev <iamsergeyteleshev@gmail.com>
Date: Wed, 10 Apr 2024 17:58:28 +0200
Subject: [PATCH 01/36] CB-4270 adds cancel option for load total count of rows
 in table

---
 .../DatabaseDataModel/DatabaseDataSource.ts   |  1 +
 .../DatabaseDataModel/IDatabaseDataSource.ts  |  1 +
 .../src/ResultSetDataSource.ts                | 60 +++++++++++++++----
 3 files changed, 50 insertions(+), 12 deletions(-)

diff --git a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
index 07629988be..9a54f18f9c 100644
--- a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
@@ -374,6 +374,7 @@ export abstract class DatabaseDataSource<TOptions, TResult extends IDatabaseData
 
   abstract dispose(): Promise<void>;
   abstract loadTotalCount(resultIndex: number): Promise<void>;
+  abstract cancelLoadTotalCount(resultIndex: number): Promise<void>;
 
   async requestDataAction(): Promise<TResult[] | null> {
     this.prevOptions = toJS(this.options);
diff --git a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/IDatabaseDataSource.ts b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/IDatabaseDataSource.ts
index bf6745f95f..07f70db247 100644
--- a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/IDatabaseDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/IDatabaseDataSource.ts
@@ -83,6 +83,7 @@ export interface IDatabaseDataSource<TOptions, TResult extends IDatabaseDataResu
   setExecutionContext: (context: IConnectionExecutionContext | null) => this;
   setTotalCount: (resultIndex: number, count: number) => this;
   loadTotalCount: (resultIndex: number) => Promise<void>;
+  cancelLoadTotalCount: (resultIndex: number) => Promise<void>;
 
   retry: () => Promise<void>;
   /** Allows to perform an asynchronous action on the data source, this action will wait previous action to finish and save or load requests.
diff --git a/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts b/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
index be89670a59..9c5859d842 100644
--- a/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
@@ -5,19 +5,42 @@
  * Licensed under the Apache License, Version 2.0.
  * you may not use this file except in compliance with the License.
  */
+import { makeObservable, observable } from 'mobx';
+
 import type { IServiceInjector } from '@cloudbeaver/core-di';
-import type { AsyncTaskInfoService, GraphQLService } from '@cloudbeaver/core-sdk';
+import type { AsyncTask, AsyncTaskInfoService, GraphQLService } from '@cloudbeaver/core-sdk';
 
 import { DatabaseDataSource } from './DatabaseDataModel/DatabaseDataSource';
 import type { IDatabaseResultSet } from './DatabaseDataModel/IDatabaseResultSet';
 
+interface ResultSetDataSourceState {
+  totalCountTasks: Map<number, AsyncTask>; // Map<resultIndex, AsyncTask>
+}
+
 export abstract class ResultSetDataSource<TOptions> extends DatabaseDataSource<TOptions, IDatabaseResultSet> {
+  readonly totalCountTasks: Map<number, AsyncTask>;
+
   constructor(
     readonly serviceInjector: IServiceInjector,
     protected graphQLService: GraphQLService,
     protected asyncTaskInfoService: AsyncTaskInfoService,
   ) {
     super(serviceInjector);
+
+    this.totalCountTasks = new Map();
+
+    makeObservable<ResultSetDataSourceState>(this, {
+      totalCountTasks: observable.shallow,
+    });
+  }
+
+  async cancelLoadTotalCount(resultIndex: number) {
+    const task = this.totalCountTasks.get(resultIndex);
+
+    if (task && !task.cancelled) {
+      await this.asyncTaskInfoService.cancel(task.id);
+      this.totalCountTasks.delete(resultIndex);
+    }
   }
 
   async loadTotalCount(resultIndex: number) {
@@ -45,16 +68,29 @@ export abstract class ResultSetDataSource<TOptions> extends DatabaseDataSource<T
       return taskInfo;
     });
 
-    const count = await executionContext.run(
-      async () => {
-        const info = await this.asyncTaskInfoService.run(task);
-        const { count } = await this.graphQLService.sdk.getSqlRowDataCountResult({ taskId: info.id });
-        return count;
-      },
-      () => this.asyncTaskInfoService.cancel(task.id),
-      () => this.asyncTaskInfoService.remove(task.id),
-    );
-
-    this.setTotalCount(resultIndex, count);
+    this.totalCountTasks.set(resultIndex, task);
+
+    try {
+      const count = await executionContext.run(
+        async () => {
+          const info = await this.asyncTaskInfoService.run(task);
+
+          if (!task.cancelled && !info.running) {
+            const { count } = await this.graphQLService.sdk.getSqlRowDataCountResult({ taskId: info.id });
+            return count;
+          }
+
+          return;
+        },
+        () => this.asyncTaskInfoService.cancel(task.id),
+        () => this.asyncTaskInfoService.remove(task.id),
+      );
+
+      if (count !== undefined) {
+        this.setTotalCount(resultIndex, count);
+      }
+    } finally {
+      this.totalCountTasks.delete(resultIndex);
+    }
   }
 }

From a6da7c5d3f3fc690443f65939438edd28c97277a Mon Sep 17 00:00:00 2001
From: sergeyteleshev <iamsergeyteleshev@gmail.com>
Date: Thu, 11 Apr 2024 08:32:05 +0200
Subject: [PATCH 02/36] CB-4270 adds ui for load total count of rows in table

---
 .../src/ToolsPanel/ToolsAction.tsx            | 19 ++++++++++++++++---
 .../public/icons/data_cancel.svg              | 10 ++++++++++
 .../TableFooter/TableFooterRowCount.tsx       | 19 +++++++++++++++++++
 .../plugin-data-viewer/src/locales/en.ts      |  8 ++++++++
 .../plugin-data-viewer/src/locales/it.ts      |  8 ++++++++
 .../plugin-data-viewer/src/locales/ru.ts      |  8 ++++++++
 .../plugin-data-viewer/src/locales/zh.ts      |  8 ++++++++
 7 files changed, 77 insertions(+), 3 deletions(-)
 create mode 100644 webapp/packages/plugin-data-viewer/public/icons/data_cancel.svg

diff --git a/webapp/packages/core-blocks/src/ToolsPanel/ToolsAction.tsx b/webapp/packages/core-blocks/src/ToolsPanel/ToolsAction.tsx
index 325e2a659b..a771005f9e 100644
--- a/webapp/packages/core-blocks/src/ToolsPanel/ToolsAction.tsx
+++ b/webapp/packages/core-blocks/src/ToolsPanel/ToolsAction.tsx
@@ -19,10 +19,23 @@ interface Props extends Omit<ButtonHTMLAttributes<any>, 'onClick'> {
   viewBox?: string;
   loading?: boolean;
   className?: string;
+  disableLoading?: boolean;
+  disableDisabled?: boolean;
   onClick?: (event: React.MouseEvent<HTMLButtonElement>) => Promise<any> | any;
 }
 
-export const ToolsAction: React.FC<Props> = function ToolsAction({ icon, viewBox, disabled, loading, children, className, onClick, ...rest }) {
+export const ToolsAction: React.FC<Props> = function ToolsAction({
+  icon,
+  viewBox,
+  disabled,
+  loading,
+  children,
+  className,
+  onClick,
+  disableLoading = false,
+  disableDisabled = false,
+  ...rest
+}) {
   const styles = useS(style);
   const [loadingState, setLoadingState] = useState(false);
 
@@ -35,8 +48,8 @@ export const ToolsAction: React.FC<Props> = function ToolsAction({ icon, viewBox
     }
   }
 
-  loading = useStateDelay(loading || loadingState, 300);
-  disabled = disabled || loadingState;
+  loading = useStateDelay(!disableLoading && (loading || loadingState), 300);
+  disabled = !disableDisabled && (disabled || loadingState);
 
   return (
     <button type="button" disabled={disabled} onClick={handleClick} {...rest} className={s(styles, { button: true }, className)}>
diff --git a/webapp/packages/plugin-data-viewer/public/icons/data_cancel.svg b/webapp/packages/plugin-data-viewer/public/icons/data_cancel.svg
new file mode 100644
index 0000000000..9a20bf546c
--- /dev/null
+++ b/webapp/packages/plugin-data-viewer/public/icons/data_cancel.svg
@@ -0,0 +1,10 @@
+<svg width="16" height="16" viewBox="0 0 16 16" fill="none" xmlns="http://www.w3.org/2000/svg">
+<g clip-path="url(#clip0)">
+<path fill-rule="evenodd" clip-rule="evenodd" d="M4.3988 3.13864C4.05709 2.79693 3.50307 2.79693 3.16136 3.13864C2.81965 3.48035 2.81965 4.03437 3.16136 4.37608L6.78528 8L3.16136 11.6239C2.81965 11.9656 2.81965 12.5197 3.16136 12.8614C3.50307 13.2031 4.05709 13.2031 4.3988 12.8614L8.02272 9.23744L11.6466 12.8614C11.9883 13.2031 12.5424 13.2031 12.8841 12.8614C13.2258 12.5197 13.2258 11.9656 12.8841 11.6239L9.26015 8L12.8841 4.37608C13.2258 4.03437 13.2258 3.48035 12.8841 3.13864C12.5424 2.79693 11.9883 2.79693 11.6466 3.13864L8.02272 6.76256L4.3988 3.13864Z" fill="#E49388"/>
+</g>
+<defs>
+<clipPath id="clip0">
+<rect width="16" height="16" fill="white"/>
+</clipPath>
+</defs>
+</svg>
diff --git a/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TableFooterRowCount.tsx b/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TableFooterRowCount.tsx
index 59181a8fc8..3f8994950d 100644
--- a/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TableFooterRowCount.tsx
+++ b/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TableFooterRowCount.tsx
@@ -45,10 +45,29 @@ export const TableFooterRowCount: React.FC<Props> = observer(function TableFoote
     }
   }
 
+  async function cancelTotalCount() {
+    try {
+      setLoading(false);
+      await model.source.cancelLoadTotalCount(resultIndex);
+    } catch (e: any) {
+      notificationService.logError({ title: 'data_viewer_total_count_cancel_failed', message: String(e) });
+    }
+  }
+
   const currentCount = result.loadedFully ? result.count : `${result.count}+`;
   const count = result.totalCount ?? currentCount;
   const disabled = getComputed(() => model.isLoading() || model.isDisabled(resultIndex));
 
+  if (loading) {
+    return styled(tableFooterMenuStyles)(
+      <div className={classes.wrapper} title={translate('ui_processing_cancel')}>
+        <ToolsAction icon="/icons/data_cancel.svg" viewBox="0 0 32 32" disableLoading disableDisabled onClick={cancelTotalCount}>
+          {translate('ui_processing_cancel')}
+        </ToolsAction>
+      </div>,
+    );
+  }
+
   return styled(tableFooterMenuStyles)(
     <div className={classes.wrapper} title={translate('data_viewer_total_count_tooltip')}>
       <ToolsAction disabled={disabled} loading={loading} icon="/icons/data_row_count.svg" viewBox="0 0 32 32" onClick={loadTotalCount}>
diff --git a/webapp/packages/plugin-data-viewer/src/locales/en.ts b/webapp/packages/plugin-data-viewer/src/locales/en.ts
index 8252fcf1f7..66dfdee9d6 100644
--- a/webapp/packages/plugin-data-viewer/src/locales/en.ts
+++ b/webapp/packages/plugin-data-viewer/src/locales/en.ts
@@ -1,3 +1,10 @@
+/*
+ * CloudBeaver - Cloud Database Manager
+ * Copyright (C) 2020-2024 DBeaver Corp and others
+ *
+ * Licensed under the Apache License, Version 2.0.
+ * you may not use this file except in compliance with the License.
+ */
 export default [
   ['plugin_data_viewer_data_viewer_settings_group', 'Data Viewer'],
   ['table_header_sql_expression', 'Enter a SQL expression to filter results, e.g. column_name=10'],
@@ -46,6 +53,7 @@ export default [
   ['data_viewer_refresh_result_set', 'Refresh result set'],
   ['data_viewer_total_count_tooltip', 'Get total count'],
   ['data_viewer_total_count_failed', 'Failed to get total count'],
+  ['data_viewer_total_count_cancel_failed', 'Failed to cancel getting of total count'],
   ['data_viewer_model_not_loaded', 'Table model is not loaded'],
   ['settings_data_editor', 'Data Editor'],
   ['settings_data_editor_disable_edit_name', 'Disable Edit'],
diff --git a/webapp/packages/plugin-data-viewer/src/locales/it.ts b/webapp/packages/plugin-data-viewer/src/locales/it.ts
index 8a952d8490..770ca50392 100644
--- a/webapp/packages/plugin-data-viewer/src/locales/it.ts
+++ b/webapp/packages/plugin-data-viewer/src/locales/it.ts
@@ -1,3 +1,10 @@
+/*
+ * CloudBeaver - Cloud Database Manager
+ * Copyright (C) 2020-2024 DBeaver Corp and others
+ *
+ * Licensed under the Apache License, Version 2.0.
+ * you may not use this file except in compliance with the License.
+ */
 export default [
   ['plugin_data_viewer_data_viewer_settings_group', 'Data Viewer'],
   ['table_header_sql_expression', 'Digita una espresione SQL per filtrare i risultati'],
@@ -39,6 +46,7 @@ export default [
   ['data_viewer_refresh_result_set', 'Refresh result set'],
   ['data_viewer_total_count_tooltip', 'Get total count'],
   ['data_viewer_total_count_failed', 'Failed to get total count'],
+  ['data_viewer_total_count_cancel_failed', 'Failed to cancel getting of total count'],
   ['data_viewer_model_not_loaded', 'Table model is not loaded'],
   ['settings_data_editor', 'Data Editor'],
   ['settings_data_editor_disable_edit_name', 'Disable Edit'],
diff --git a/webapp/packages/plugin-data-viewer/src/locales/ru.ts b/webapp/packages/plugin-data-viewer/src/locales/ru.ts
index 8792885dc2..616b5d9f84 100644
--- a/webapp/packages/plugin-data-viewer/src/locales/ru.ts
+++ b/webapp/packages/plugin-data-viewer/src/locales/ru.ts
@@ -1,3 +1,10 @@
+/*
+ * CloudBeaver - Cloud Database Manager
+ * Copyright (C) 2020-2024 DBeaver Corp and others
+ *
+ * Licensed under the Apache License, Version 2.0.
+ * you may not use this file except in compliance with the License.
+ */
 export default [
   ['plugin_data_viewer_data_viewer_settings_group', 'Просмотр данных'],
   ['table_header_sql_expression', 'Введите SQL выражение чтобы отфильтровать результаты, например, column_name=10'],
@@ -40,6 +47,7 @@ export default [
   ['data_viewer_refresh_result_set', 'Обновить резалт сет'],
   ['data_viewer_total_count_tooltip', 'Получить количество записей'],
   ['data_viewer_total_count_failed', 'Не удалось получить количество записей'],
+  ['data_viewer_total_count_cancel_failed', 'Не удалось отменить получение количества записей'],
   ['data_viewer_model_not_loaded', 'Не удалось загрузить модель таблицы'],
   ['settings_data_editor', 'Редактор данных'],
   ['settings_data_editor_disable_edit_name', 'Отключить редактирование'],
diff --git a/webapp/packages/plugin-data-viewer/src/locales/zh.ts b/webapp/packages/plugin-data-viewer/src/locales/zh.ts
index 7c8babb48c..f8a0895f5a 100644
--- a/webapp/packages/plugin-data-viewer/src/locales/zh.ts
+++ b/webapp/packages/plugin-data-viewer/src/locales/zh.ts
@@ -1,3 +1,10 @@
+/*
+ * CloudBeaver - Cloud Database Manager
+ * Copyright (C) 2020-2024 DBeaver Corp and others
+ *
+ * Licensed under the Apache License, Version 2.0.
+ * you may not use this file except in compliance with the License.
+ */
 export default [
   ['plugin_data_viewer_data_viewer_settings_group', 'Data Viewer'],
   ['table_header_sql_expression', '输入SQL表达式以过滤结果'],
@@ -46,6 +53,7 @@ export default [
   ['data_viewer_refresh_result_set', 'Refresh result set'],
   ['data_viewer_total_count_failed', 'Failed to get total count'],
   ['data_viewer_total_count_tooltip', 'Get total count'],
+  ['data_viewer_total_count_cancel_failed', 'Failed to cancel getting of total count'],
   ['data_viewer_model_not_loaded', 'Table model is not loaded'],
   ['settings_data_editor', 'Data Editor'],
   ['settings_data_editor_disable_edit_name', 'Disable Edit'],

From de7c6e49c57ac24709bd11881535c44fecb2c157 Mon Sep 17 00:00:00 2001
From: sergeyteleshev <iamsergeyteleshev@gmail.com>
Date: Thu, 11 Apr 2024 12:22:48 +0200
Subject: [PATCH 03/36] CB-4270 fix: task cancels in a true way

---
 .../DatabaseDataModel/IDatabaseDataSource.ts  |  2 +
 .../src/ResultSetDataSource.ts                | 60 +++++++++----------
 .../TableFooter/TableFooterRowCount.tsx       | 13 +++-
 .../plugin-data-viewer/src/locales/en.ts      |  3 +-
 .../plugin-data-viewer/src/locales/it.ts      |  3 +-
 .../plugin-data-viewer/src/locales/ru.ts      |  3 +-
 .../plugin-data-viewer/src/locales/zh.ts      |  3 +-
 7 files changed, 52 insertions(+), 35 deletions(-)

diff --git a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/IDatabaseDataSource.ts b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/IDatabaseDataSource.ts
index 07f70db247..63f35c7aed 100644
--- a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/IDatabaseDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/IDatabaseDataSource.ts
@@ -7,6 +7,7 @@
  */
 import type { IConnectionExecutionContext } from '@cloudbeaver/core-connections';
 import type { IServiceInjector } from '@cloudbeaver/core-di';
+import type { ITask } from '@cloudbeaver/core-executor';
 import type { ResultDataFormat } from '@cloudbeaver/core-sdk';
 
 import type { IDatabaseDataAction, IDatabaseDataActionClass, IDatabaseDataActionInterface } from './IDatabaseDataAction';
@@ -47,6 +48,7 @@ export interface IDatabaseDataSource<TOptions, TResult extends IDatabaseDataResu
   readonly cancelled: boolean;
   readonly serviceInjector: IServiceInjector;
   readonly outdated: boolean;
+  readonly tasks?: Map<string, ITask<number>>;
 
   isLoadable: () => boolean;
   isReadonly: (resultIndex: number) => boolean;
diff --git a/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts b/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
index 9c5859d842..f814ee47c1 100644
--- a/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
@@ -8,17 +8,18 @@
 import { makeObservable, observable } from 'mobx';
 
 import type { IServiceInjector } from '@cloudbeaver/core-di';
-import type { AsyncTask, AsyncTaskInfoService, GraphQLService } from '@cloudbeaver/core-sdk';
+import type { ITask } from '@cloudbeaver/core-executor';
+import type { AsyncTaskInfoService, GraphQLService } from '@cloudbeaver/core-sdk';
 
 import { DatabaseDataSource } from './DatabaseDataModel/DatabaseDataSource';
 import type { IDatabaseResultSet } from './DatabaseDataModel/IDatabaseResultSet';
 
 interface ResultSetDataSourceState {
-  totalCountTasks: Map<number, AsyncTask>; // Map<resultIndex, AsyncTask>
+  tasks: Map<string, ITask<number>>;
 }
 
 export abstract class ResultSetDataSource<TOptions> extends DatabaseDataSource<TOptions, IDatabaseResultSet> {
-  readonly totalCountTasks: Map<number, AsyncTask>;
+  readonly tasks: Map<string, ITask<number>>;
 
   constructor(
     readonly serviceInjector: IServiceInjector,
@@ -27,20 +28,23 @@ export abstract class ResultSetDataSource<TOptions> extends DatabaseDataSource<T
   ) {
     super(serviceInjector);
 
-    this.totalCountTasks = new Map();
+    this.tasks = new Map();
 
     makeObservable<ResultSetDataSourceState>(this, {
-      totalCountTasks: observable.shallow,
+      tasks: observable.shallow,
     });
   }
 
   async cancelLoadTotalCount(resultIndex: number) {
-    const task = this.totalCountTasks.get(resultIndex);
+    const result = this.getResult(resultIndex);
 
-    if (task && !task.cancelled) {
-      await this.asyncTaskInfoService.cancel(task.id);
-      this.totalCountTasks.delete(resultIndex);
+    if (!result?.id) {
+      throw new Error('Result id must be provided');
     }
+
+    const task = this.tasks.get(result.id);
+    task?.cancel();
+    this.tasks.delete(result.id);
   }
 
   async loadTotalCount(resultIndex: number) {
@@ -57,7 +61,7 @@ export abstract class ResultSetDataSource<TOptions> extends DatabaseDataSource<T
       throw new Error('Result id must be provided');
     }
 
-    const task = this.asyncTaskInfoService.create(async () => {
+    const asyncTask = this.asyncTaskInfoService.create(async () => {
       const { taskInfo } = await this.graphQLService.sdk.asyncSqlRowDataCount({
         resultsId: result.id!,
         connectionId: executionContextInfo.connectionId,
@@ -68,29 +72,25 @@ export abstract class ResultSetDataSource<TOptions> extends DatabaseDataSource<T
       return taskInfo;
     });
 
-    this.totalCountTasks.set(resultIndex, task);
+    const task = executionContext.run(
+      async () => {
+        const info = await this.asyncTaskInfoService.run(asyncTask);
+        const { count } = await this.graphQLService.sdk.getSqlRowDataCountResult({ taskId: info.id });
+
+        return count;
+      },
+      () => this.asyncTaskInfoService.cancel(asyncTask.id),
+      () => this.asyncTaskInfoService.remove(asyncTask.id),
+    );
+
+    this.tasks.set(result.id, task);
 
     try {
-      const count = await executionContext.run(
-        async () => {
-          const info = await this.asyncTaskInfoService.run(task);
-
-          if (!task.cancelled && !info.running) {
-            const { count } = await this.graphQLService.sdk.getSqlRowDataCountResult({ taskId: info.id });
-            return count;
-          }
-
-          return;
-        },
-        () => this.asyncTaskInfoService.cancel(task.id),
-        () => this.asyncTaskInfoService.remove(task.id),
-      );
-
-      if (count !== undefined) {
-        this.setTotalCount(resultIndex, count);
-      }
+      const count = await task;
+
+      this.setTotalCount(resultIndex, count);
     } finally {
-      this.totalCountTasks.delete(resultIndex);
+      this.tasks.delete(result.id);
     }
   }
 }
diff --git a/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TableFooterRowCount.tsx b/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TableFooterRowCount.tsx
index 3f8994950d..0e56fde831 100644
--- a/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TableFooterRowCount.tsx
+++ b/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TableFooterRowCount.tsx
@@ -50,7 +50,18 @@ export const TableFooterRowCount: React.FC<Props> = observer(function TableFoote
       setLoading(false);
       await model.source.cancelLoadTotalCount(resultIndex);
     } catch (e: any) {
-      notificationService.logError({ title: 'data_viewer_total_count_cancel_failed', message: String(e) });
+      const result = model.getResult(resultIndex);
+      const cancelled = Boolean(result?.id && model.source.tasks?.get(result?.id)?.cancelled);
+
+      if (cancelled) {
+        notificationService.logError({
+          title: 'data_viewer_total_count_cancel_failed_title',
+          message: 'data_viewer_total_count_cancel_failed_desc',
+        });
+        return;
+      }
+
+      notificationService.logException(e, 'data_viewer_total_count_cancel_failed_title', typeof e === 'string' ? e : undefined);
     }
   }
 
diff --git a/webapp/packages/plugin-data-viewer/src/locales/en.ts b/webapp/packages/plugin-data-viewer/src/locales/en.ts
index 66dfdee9d6..327f6567c0 100644
--- a/webapp/packages/plugin-data-viewer/src/locales/en.ts
+++ b/webapp/packages/plugin-data-viewer/src/locales/en.ts
@@ -53,7 +53,8 @@ export default [
   ['data_viewer_refresh_result_set', 'Refresh result set'],
   ['data_viewer_total_count_tooltip', 'Get total count'],
   ['data_viewer_total_count_failed', 'Failed to get total count'],
-  ['data_viewer_total_count_cancel_failed', 'Failed to cancel getting of total count'],
+  ['data_viewer_total_count_cancel_failed_title', 'Failed to cancel getting of total count'],
+  ['data_viewer_total_count_cancel_failed_desc', 'Counting already finished'],
   ['data_viewer_model_not_loaded', 'Table model is not loaded'],
   ['settings_data_editor', 'Data Editor'],
   ['settings_data_editor_disable_edit_name', 'Disable Edit'],
diff --git a/webapp/packages/plugin-data-viewer/src/locales/it.ts b/webapp/packages/plugin-data-viewer/src/locales/it.ts
index 770ca50392..45ebfe9a0b 100644
--- a/webapp/packages/plugin-data-viewer/src/locales/it.ts
+++ b/webapp/packages/plugin-data-viewer/src/locales/it.ts
@@ -46,7 +46,8 @@ export default [
   ['data_viewer_refresh_result_set', 'Refresh result set'],
   ['data_viewer_total_count_tooltip', 'Get total count'],
   ['data_viewer_total_count_failed', 'Failed to get total count'],
-  ['data_viewer_total_count_cancel_failed', 'Failed to cancel getting of total count'],
+  ['data_viewer_total_count_cancel_failed_title', 'Failed to cancel getting of total count'],
+  ['data_viewer_total_count_cancel_failed_desc', 'Counting already finished'],
   ['data_viewer_model_not_loaded', 'Table model is not loaded'],
   ['settings_data_editor', 'Data Editor'],
   ['settings_data_editor_disable_edit_name', 'Disable Edit'],
diff --git a/webapp/packages/plugin-data-viewer/src/locales/ru.ts b/webapp/packages/plugin-data-viewer/src/locales/ru.ts
index 616b5d9f84..b13ea2580b 100644
--- a/webapp/packages/plugin-data-viewer/src/locales/ru.ts
+++ b/webapp/packages/plugin-data-viewer/src/locales/ru.ts
@@ -47,7 +47,8 @@ export default [
   ['data_viewer_refresh_result_set', 'Обновить резалт сет'],
   ['data_viewer_total_count_tooltip', 'Получить количество записей'],
   ['data_viewer_total_count_failed', 'Не удалось получить количество записей'],
-  ['data_viewer_total_count_cancel_failed', 'Не удалось отменить получение количества записей'],
+  ['data_viewer_total_count_cancel_failed_title', 'Не удалось отменить получение количества записей'],
+  ['data_viewer_total_count_cancel_failed_desc', 'Подсчет уже завершен'],
   ['data_viewer_model_not_loaded', 'Не удалось загрузить модель таблицы'],
   ['settings_data_editor', 'Редактор данных'],
   ['settings_data_editor_disable_edit_name', 'Отключить редактирование'],
diff --git a/webapp/packages/plugin-data-viewer/src/locales/zh.ts b/webapp/packages/plugin-data-viewer/src/locales/zh.ts
index f8a0895f5a..908b4b3986 100644
--- a/webapp/packages/plugin-data-viewer/src/locales/zh.ts
+++ b/webapp/packages/plugin-data-viewer/src/locales/zh.ts
@@ -53,7 +53,8 @@ export default [
   ['data_viewer_refresh_result_set', 'Refresh result set'],
   ['data_viewer_total_count_failed', 'Failed to get total count'],
   ['data_viewer_total_count_tooltip', 'Get total count'],
-  ['data_viewer_total_count_cancel_failed', 'Failed to cancel getting of total count'],
+  ['data_viewer_total_count_cancel_failed_title', 'Failed to cancel getting of total count'],
+  ['data_viewer_total_count_cancel_failed_desc', 'Counting already finished'],
   ['data_viewer_model_not_loaded', 'Table model is not loaded'],
   ['settings_data_editor', 'Data Editor'],
   ['settings_data_editor_disable_edit_name', 'Disable Edit'],

From bfcfaa0315ac6e702d9d5191ccc6fb9c281132f1 Mon Sep 17 00:00:00 2001
From: sergeyteleshev <iamsergeyteleshev@gmail.com>
Date: Thu, 11 Apr 2024 12:25:50 +0200
Subject: [PATCH 04/36] CB-4270 cleanup

---
 .../src/DatabaseDataModel/DatabaseDataSource.ts                | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
index 9a54f18f9c..c0fc4f56a0 100644
--- a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
@@ -9,7 +9,7 @@ import { action, makeObservable, observable, toJS } from 'mobx';
 
 import type { IConnectionExecutionContext } from '@cloudbeaver/core-connections';
 import type { IServiceInjector } from '@cloudbeaver/core-di';
-import { Task } from '@cloudbeaver/core-executor';
+import { ITask, Task } from '@cloudbeaver/core-executor';
 import { ResultDataFormat } from '@cloudbeaver/core-sdk';
 
 import { DatabaseDataActions } from './DatabaseDataActions';
@@ -33,6 +33,7 @@ export abstract class DatabaseDataSource<TOptions, TResult extends IDatabaseData
   error: Error | null;
   executionContext: IConnectionExecutionContext | null;
   outdated: boolean;
+  abstract tasks?: Map<string, ITask<number>>;
 
   get canCancel(): boolean {
     if (this.activeTask instanceof Task) {

From 7c424c60cde084e191805bc94353c3cdd9effa11 Mon Sep 17 00:00:00 2001
From: sergeyteleshev <iamsergeyteleshev@gmail.com>
Date: Thu, 11 Apr 2024 13:58:51 +0200
Subject: [PATCH 05/36] CB-4270 cancel load total count on close table tab

---
 .../packages/plugin-data-viewer/src/DataViewerTableService.ts  | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/webapp/packages/plugin-data-viewer/src/DataViewerTableService.ts b/webapp/packages/plugin-data-viewer/src/DataViewerTableService.ts
index 0d38871517..ba74ba1e97 100644
--- a/webapp/packages/plugin-data-viewer/src/DataViewerTableService.ts
+++ b/webapp/packages/plugin-data-viewer/src/DataViewerTableService.ts
@@ -45,6 +45,9 @@ export class DataViewerTableService {
 
     if (model) {
       this.tableViewerStorageService.remove(tableId);
+      model.getResults().forEach((_, index) => {
+        model.source.cancelLoadTotalCount(index);
+      });
       await model.dispose();
     }
   }

From 110c1d81e188c4bc51b23a65c6d70f91f7e80ce4 Mon Sep 17 00:00:00 2001
From: sergeyteleshev <iamsergeyteleshev@gmail.com>
Date: Thu, 11 Apr 2024 13:59:22 +0200
Subject: [PATCH 06/36] CB-4270 cancel load total count on close sql editor tab

---
 .../src/SqlEditorTabService.ts                                 | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/webapp/packages/plugin-sql-editor-navigation-tab/src/SqlEditorTabService.ts b/webapp/packages/plugin-sql-editor-navigation-tab/src/SqlEditorTabService.ts
index a4d7857fd1..76836c0cff 100644
--- a/webapp/packages/plugin-sql-editor-navigation-tab/src/SqlEditorTabService.ts
+++ b/webapp/packages/plugin-sql-editor-navigation-tab/src/SqlEditorTabService.ts
@@ -493,17 +493,18 @@ export class SqlEditorTabService extends Bootstrap {
 
   private async handleTabCloseSilent(editorTab: ITab<ISqlEditorTabState>) {
     const dataSource = this.sqlDataSourceService.get(editorTab.handlerState.editorId);
+    this.sqlResultTabsService.removeResultTabs(editorTab.handlerState);
 
     if (dataSource?.executionContext) {
       await this.sqlEditorService.destroyContext(dataSource.executionContext);
     }
 
     await this.sqlDataSourceService.destroySilent(editorTab.handlerState.editorId);
-    this.sqlResultTabsService.removeResultTabs(editorTab.handlerState);
   }
 
   private async handleTabClose(editorTab: ITab<ISqlEditorTabState>) {
     const dataSource = this.sqlDataSourceService.get(editorTab.handlerState.editorId);
+    this.sqlResultTabsService.removeResultTabs(editorTab.handlerState);
 
     if (dataSource?.executionContext) {
       await this.sqlEditorService.destroyContext(dataSource.executionContext);

From d58b80cd208639bbc84f098765f4cf57a17055f9 Mon Sep 17 00:00:00 2001
From: sergeyteleshev <iamsergeyteleshev@gmail.com>
Date: Thu, 11 Apr 2024 13:59:59 +0200
Subject: [PATCH 07/36] CB-4270 fix: cancel close only task if it is not
 already cancelled

---
 .../packages/plugin-data-viewer/src/ResultSetDataSource.ts  | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts b/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
index f814ee47c1..d3957bcfe3 100644
--- a/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
@@ -43,7 +43,11 @@ export abstract class ResultSetDataSource<TOptions> extends DatabaseDataSource<T
     }
 
     const task = this.tasks.get(result.id);
-    task?.cancel();
+
+    if (!task?.cancelled) {
+      task?.cancel();
+    }
+
     this.tasks.delete(result.id);
   }
 

From 9eea88c1370601535cf0a30ae971029278c4e153 Mon Sep 17 00:00:00 2001
From: sergeyteleshev <iamsergeyteleshev@gmail.com>
Date: Thu, 11 Apr 2024 14:56:36 +0200
Subject: [PATCH 08/36] CB-4270 chore: tasks -> cancelLoadTotalCountTasks

---
 .../src/DatabaseDataModel/DatabaseDataSource.ts  |  5 ++++-
 .../src/DatabaseDataModel/IDatabaseDataSource.ts |  2 +-
 .../src/ResultSetDataSource.ts                   | 16 ++++++++--------
 .../TableFooter/TableFooterRowCount.tsx          |  2 +-
 4 files changed, 14 insertions(+), 11 deletions(-)

diff --git a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
index c0fc4f56a0..6460924616 100644
--- a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
@@ -33,7 +33,7 @@ export abstract class DatabaseDataSource<TOptions, TResult extends IDatabaseData
   error: Error | null;
   executionContext: IConnectionExecutionContext | null;
   outdated: boolean;
-  abstract tasks?: Map<string, ITask<number>>;
+  abstract cancelLoadTotalCountTasks?: Map<string, ITask<number>>;
 
   get canCancel(): boolean {
     if (this.activeTask instanceof Task) {
@@ -194,6 +194,9 @@ export abstract class DatabaseDataSource<TOptions, TResult extends IDatabaseData
   }
 
   setResults(results: TResult[]): this {
+    results.forEach((_, index) => {
+      this.cancelLoadTotalCount(index);
+    });
     results = observable(results);
     this.actions.updateResults(results);
     this.results = results;
diff --git a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/IDatabaseDataSource.ts b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/IDatabaseDataSource.ts
index 63f35c7aed..cb958d7463 100644
--- a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/IDatabaseDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/IDatabaseDataSource.ts
@@ -48,7 +48,7 @@ export interface IDatabaseDataSource<TOptions, TResult extends IDatabaseDataResu
   readonly cancelled: boolean;
   readonly serviceInjector: IServiceInjector;
   readonly outdated: boolean;
-  readonly tasks?: Map<string, ITask<number>>;
+  readonly cancelLoadTotalCountTasks?: Map<string, ITask<number>>;
 
   isLoadable: () => boolean;
   isReadonly: (resultIndex: number) => boolean;
diff --git a/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts b/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
index d3957bcfe3..4751ee538c 100644
--- a/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
@@ -15,11 +15,11 @@ import { DatabaseDataSource } from './DatabaseDataModel/DatabaseDataSource';
 import type { IDatabaseResultSet } from './DatabaseDataModel/IDatabaseResultSet';
 
 interface ResultSetDataSourceState {
-  tasks: Map<string, ITask<number>>;
+  cancelLoadTotalCountTasks: Map<string, ITask<number>>;
 }
 
 export abstract class ResultSetDataSource<TOptions> extends DatabaseDataSource<TOptions, IDatabaseResultSet> {
-  readonly tasks: Map<string, ITask<number>>;
+  readonly cancelLoadTotalCountTasks: Map<string, ITask<number>>;
 
   constructor(
     readonly serviceInjector: IServiceInjector,
@@ -28,10 +28,10 @@ export abstract class ResultSetDataSource<TOptions> extends DatabaseDataSource<T
   ) {
     super(serviceInjector);
 
-    this.tasks = new Map();
+    this.cancelLoadTotalCountTasks = new Map();
 
     makeObservable<ResultSetDataSourceState>(this, {
-      tasks: observable.shallow,
+      cancelLoadTotalCountTasks: observable.shallow,
     });
   }
 
@@ -42,13 +42,13 @@ export abstract class ResultSetDataSource<TOptions> extends DatabaseDataSource<T
       throw new Error('Result id must be provided');
     }
 
-    const task = this.tasks.get(result.id);
+    const task = this.cancelLoadTotalCountTasks.get(result.id);
 
     if (!task?.cancelled) {
       task?.cancel();
     }
 
-    this.tasks.delete(result.id);
+    this.cancelLoadTotalCountTasks.delete(result.id);
   }
 
   async loadTotalCount(resultIndex: number) {
@@ -87,14 +87,14 @@ export abstract class ResultSetDataSource<TOptions> extends DatabaseDataSource<T
       () => this.asyncTaskInfoService.remove(asyncTask.id),
     );
 
-    this.tasks.set(result.id, task);
+    this.cancelLoadTotalCountTasks.set(result.id, task);
 
     try {
       const count = await task;
 
       this.setTotalCount(resultIndex, count);
     } finally {
-      this.tasks.delete(result.id);
+      this.cancelLoadTotalCountTasks.delete(result.id);
     }
   }
 }
diff --git a/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TableFooterRowCount.tsx b/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TableFooterRowCount.tsx
index 0e56fde831..6a65144c3c 100644
--- a/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TableFooterRowCount.tsx
+++ b/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TableFooterRowCount.tsx
@@ -51,7 +51,7 @@ export const TableFooterRowCount: React.FC<Props> = observer(function TableFoote
       await model.source.cancelLoadTotalCount(resultIndex);
     } catch (e: any) {
       const result = model.getResult(resultIndex);
-      const cancelled = Boolean(result?.id && model.source.tasks?.get(result?.id)?.cancelled);
+      const cancelled = Boolean(result?.id && model.source.cancelLoadTotalCountTasks?.get(result?.id)?.cancelled);
 
       if (cancelled) {
         notificationService.logError({

From 9ddb968b41e8e75eb72293217b0eddd3b0cce7b0 Mon Sep 17 00:00:00 2001
From: sergeyteleshev <iamsergeyteleshev@gmail.com>
Date: Thu, 11 Apr 2024 15:29:36 +0200
Subject: [PATCH 09/36] CB-4270 stops counting rows when tab closes

---
 .../src/DatabaseDataModel/DatabaseDataSource.ts             | 3 ---
 .../src/SqlResultTabs/SqlQueryResultService.ts              | 6 ++++++
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
index 6460924616..498e9e506c 100644
--- a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
@@ -194,9 +194,6 @@ export abstract class DatabaseDataSource<TOptions, TResult extends IDatabaseData
   }
 
   setResults(results: TResult[]): this {
-    results.forEach((_, index) => {
-      this.cancelLoadTotalCount(index);
-    });
     results = observable(results);
     this.actions.updateResults(results);
     this.results = results;
diff --git a/webapp/packages/plugin-sql-editor/src/SqlResultTabs/SqlQueryResultService.ts b/webapp/packages/plugin-sql-editor/src/SqlResultTabs/SqlQueryResultService.ts
index 5d0d8d31ac..20531984c2 100644
--- a/webapp/packages/plugin-sql-editor/src/SqlResultTabs/SqlQueryResultService.ts
+++ b/webapp/packages/plugin-sql-editor/src/SqlResultTabs/SqlQueryResultService.ts
@@ -153,6 +153,12 @@ export class SqlQueryResultService {
     const resultTab = state.resultTabs.find(resultTab => resultTab.tabId === tabId);
     const group = state.resultGroups.find(group => group.groupId === resultTab?.groupId);
 
+    const model = group && this.tableViewerStorageService.get(group.modelId);
+
+    if (model && resultTab) {
+      model.source.cancelLoadTotalCount(resultTab.indexInResultSet);
+    }
+
     if (resultTab && group) {
       state.resultTabs.splice(state.resultTabs.indexOf(resultTab), 1);
 

From bcc43ce02bd7a42a1f9e29df8cae83a9334bb27b Mon Sep 17 00:00:00 2001
From: sergeyteleshev <iamsergeyteleshev@gmail.com>
Date: Thu, 11 Apr 2024 17:05:03 +0200
Subject: [PATCH 10/36] CB-4270 reverts incorrect calls of cancelLoadTotalCount

---
 .../packages/plugin-data-viewer/src/DataViewerTableService.ts | 4 +---
 .../src/SqlEditorTabService.ts                                | 3 +--
 2 files changed, 2 insertions(+), 5 deletions(-)

diff --git a/webapp/packages/plugin-data-viewer/src/DataViewerTableService.ts b/webapp/packages/plugin-data-viewer/src/DataViewerTableService.ts
index ba74ba1e97..ccb4e7c235 100644
--- a/webapp/packages/plugin-data-viewer/src/DataViewerTableService.ts
+++ b/webapp/packages/plugin-data-viewer/src/DataViewerTableService.ts
@@ -45,9 +45,7 @@ export class DataViewerTableService {
 
     if (model) {
       this.tableViewerStorageService.remove(tableId);
-      model.getResults().forEach((_, index) => {
-        model.source.cancelLoadTotalCount(index);
-      });
+      model.source.cancelLoadTotalCount(0);
       await model.dispose();
     }
   }
diff --git a/webapp/packages/plugin-sql-editor-navigation-tab/src/SqlEditorTabService.ts b/webapp/packages/plugin-sql-editor-navigation-tab/src/SqlEditorTabService.ts
index 76836c0cff..a4d7857fd1 100644
--- a/webapp/packages/plugin-sql-editor-navigation-tab/src/SqlEditorTabService.ts
+++ b/webapp/packages/plugin-sql-editor-navigation-tab/src/SqlEditorTabService.ts
@@ -493,18 +493,17 @@ export class SqlEditorTabService extends Bootstrap {
 
   private async handleTabCloseSilent(editorTab: ITab<ISqlEditorTabState>) {
     const dataSource = this.sqlDataSourceService.get(editorTab.handlerState.editorId);
-    this.sqlResultTabsService.removeResultTabs(editorTab.handlerState);
 
     if (dataSource?.executionContext) {
       await this.sqlEditorService.destroyContext(dataSource.executionContext);
     }
 
     await this.sqlDataSourceService.destroySilent(editorTab.handlerState.editorId);
+    this.sqlResultTabsService.removeResultTabs(editorTab.handlerState);
   }
 
   private async handleTabClose(editorTab: ITab<ISqlEditorTabState>) {
     const dataSource = this.sqlDataSourceService.get(editorTab.handlerState.editorId);
-    this.sqlResultTabsService.removeResultTabs(editorTab.handlerState);
 
     if (dataSource?.executionContext) {
       await this.sqlEditorService.destroyContext(dataSource.executionContext);

From 064a73739a419c7ce01b5f2f862cd7c8317d597f Mon Sep 17 00:00:00 2001
From: sergeyteleshev <iamsergeyteleshev@gmail.com>
Date: Fri, 12 Apr 2024 09:07:59 +0200
Subject: [PATCH 11/36] CB-4270 reverts ToolsActions unneeded props

---
 .../src/ToolsPanel/ToolsAction.tsx            | 19 +-----
 .../TableFooter/TableFooterRowCount.tsx       | 66 +++++++++++--------
 .../plugin-data-viewer/src/locales/en.ts      |  1 -
 .../plugin-data-viewer/src/locales/it.ts      |  1 -
 .../plugin-data-viewer/src/locales/ru.ts      |  1 -
 .../plugin-data-viewer/src/locales/zh.ts      |  1 -
 6 files changed, 43 insertions(+), 46 deletions(-)

diff --git a/webapp/packages/core-blocks/src/ToolsPanel/ToolsAction.tsx b/webapp/packages/core-blocks/src/ToolsPanel/ToolsAction.tsx
index a771005f9e..325e2a659b 100644
--- a/webapp/packages/core-blocks/src/ToolsPanel/ToolsAction.tsx
+++ b/webapp/packages/core-blocks/src/ToolsPanel/ToolsAction.tsx
@@ -19,23 +19,10 @@ interface Props extends Omit<ButtonHTMLAttributes<any>, 'onClick'> {
   viewBox?: string;
   loading?: boolean;
   className?: string;
-  disableLoading?: boolean;
-  disableDisabled?: boolean;
   onClick?: (event: React.MouseEvent<HTMLButtonElement>) => Promise<any> | any;
 }
 
-export const ToolsAction: React.FC<Props> = function ToolsAction({
-  icon,
-  viewBox,
-  disabled,
-  loading,
-  children,
-  className,
-  onClick,
-  disableLoading = false,
-  disableDisabled = false,
-  ...rest
-}) {
+export const ToolsAction: React.FC<Props> = function ToolsAction({ icon, viewBox, disabled, loading, children, className, onClick, ...rest }) {
   const styles = useS(style);
   const [loadingState, setLoadingState] = useState(false);
 
@@ -48,8 +35,8 @@ export const ToolsAction: React.FC<Props> = function ToolsAction({
     }
   }
 
-  loading = useStateDelay(!disableLoading && (loading || loadingState), 300);
-  disabled = !disableDisabled && (disabled || loadingState);
+  loading = useStateDelay(loading || loadingState, 300);
+  disabled = disabled || loadingState;
 
   return (
     <button type="button" disabled={disabled} onClick={handleClick} {...rest} className={s(styles, { button: true }, className)}>
diff --git a/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TableFooterRowCount.tsx b/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TableFooterRowCount.tsx
index 6a65144c3c..41bbcc34d4 100644
--- a/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TableFooterRowCount.tsx
+++ b/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TableFooterRowCount.tsx
@@ -24,16 +24,9 @@ interface Props {
 }
 
 export const TableFooterRowCount: React.FC<Props> = observer(function TableFooterRowCount({ resultIndex, model }) {
-  const translate = useTranslate();
   const notificationService = useService(NotificationService);
   const [loading, setLoading] = useState(false);
 
-  const result = model.getResult(resultIndex);
-
-  if (!result) {
-    return null;
-  }
-
   async function loadTotalCount() {
     try {
       setLoading(true);
@@ -53,35 +46,56 @@ export const TableFooterRowCount: React.FC<Props> = observer(function TableFoote
       const result = model.getResult(resultIndex);
       const cancelled = Boolean(result?.id && model.source.cancelLoadTotalCountTasks?.get(result?.id)?.cancelled);
 
-      if (cancelled) {
-        notificationService.logError({
-          title: 'data_viewer_total_count_cancel_failed_title',
-          message: 'data_viewer_total_count_cancel_failed_desc',
-        });
-        return;
+      if (!cancelled) {
+        notificationService.logException(e, 'data_viewer_total_count_cancel_failed_title', typeof e === 'string' ? e : undefined);
       }
-
-      notificationService.logException(e, 'data_viewer_total_count_cancel_failed_title', typeof e === 'string' ? e : undefined);
     }
   }
 
-  const currentCount = result.loadedFully ? result.count : `${result.count}+`;
-  const count = result.totalCount ?? currentCount;
+  if (loading) {
+    return <CancelTotalCountAction onClick={cancelTotalCount} />;
+  }
+
+  return <TotalCountAction loading={loading} resultIndex={resultIndex} model={model} onClick={loadTotalCount} />;
+});
+
+const CancelTotalCountAction = observer(function CancelTotalCountAction({ onClick }: { onClick: VoidFunction }) {
+  const translate = useTranslate();
+
+  return styled(tableFooterMenuStyles)(
+    <div className={classes.wrapper} title={translate('ui_processing_cancel')}>
+      <ToolsAction icon="/icons/data_cancel.svg" viewBox="0 0 32 32" onClick={onClick}>
+        {translate('ui_processing_cancel')}
+      </ToolsAction>
+    </div>,
+  );
+});
+
+const TotalCountAction = observer(function TotalCountAction({
+  onClick,
+  loading,
+  resultIndex,
+  model,
+}: {
+  onClick: VoidFunction;
+  loading: boolean;
+  resultIndex: number;
+  model: IDatabaseDataModel<any, IDatabaseResultSet>;
+}) {
+  const result = model.getResult(resultIndex);
+  const translate = useTranslate();
   const disabled = getComputed(() => model.isLoading() || model.isDisabled(resultIndex));
 
-  if (loading) {
-    return styled(tableFooterMenuStyles)(
-      <div className={classes.wrapper} title={translate('ui_processing_cancel')}>
-        <ToolsAction icon="/icons/data_cancel.svg" viewBox="0 0 32 32" disableLoading disableDisabled onClick={cancelTotalCount}>
-          {translate('ui_processing_cancel')}
-        </ToolsAction>
-      </div>,
-    );
+  if (!result) {
+    return null;
   }
 
+  const currentCount = result.loadedFully ? result.count : `${result.count}+`;
+  const count = result.totalCount ?? currentCount;
+
   return styled(tableFooterMenuStyles)(
     <div className={classes.wrapper} title={translate('data_viewer_total_count_tooltip')}>
-      <ToolsAction disabled={disabled} loading={loading} icon="/icons/data_row_count.svg" viewBox="0 0 32 32" onClick={loadTotalCount}>
+      <ToolsAction disabled={disabled} loading={loading} icon="/icons/data_row_count.svg" viewBox="0 0 32 32" onClick={onClick}>
         <span>{count}</span>
       </ToolsAction>
     </div>,
diff --git a/webapp/packages/plugin-data-viewer/src/locales/en.ts b/webapp/packages/plugin-data-viewer/src/locales/en.ts
index 327f6567c0..256c8dadf6 100644
--- a/webapp/packages/plugin-data-viewer/src/locales/en.ts
+++ b/webapp/packages/plugin-data-viewer/src/locales/en.ts
@@ -54,7 +54,6 @@ export default [
   ['data_viewer_total_count_tooltip', 'Get total count'],
   ['data_viewer_total_count_failed', 'Failed to get total count'],
   ['data_viewer_total_count_cancel_failed_title', 'Failed to cancel getting of total count'],
-  ['data_viewer_total_count_cancel_failed_desc', 'Counting already finished'],
   ['data_viewer_model_not_loaded', 'Table model is not loaded'],
   ['settings_data_editor', 'Data Editor'],
   ['settings_data_editor_disable_edit_name', 'Disable Edit'],
diff --git a/webapp/packages/plugin-data-viewer/src/locales/it.ts b/webapp/packages/plugin-data-viewer/src/locales/it.ts
index 45ebfe9a0b..0d2d5487fe 100644
--- a/webapp/packages/plugin-data-viewer/src/locales/it.ts
+++ b/webapp/packages/plugin-data-viewer/src/locales/it.ts
@@ -47,7 +47,6 @@ export default [
   ['data_viewer_total_count_tooltip', 'Get total count'],
   ['data_viewer_total_count_failed', 'Failed to get total count'],
   ['data_viewer_total_count_cancel_failed_title', 'Failed to cancel getting of total count'],
-  ['data_viewer_total_count_cancel_failed_desc', 'Counting already finished'],
   ['data_viewer_model_not_loaded', 'Table model is not loaded'],
   ['settings_data_editor', 'Data Editor'],
   ['settings_data_editor_disable_edit_name', 'Disable Edit'],
diff --git a/webapp/packages/plugin-data-viewer/src/locales/ru.ts b/webapp/packages/plugin-data-viewer/src/locales/ru.ts
index b13ea2580b..1428ae7971 100644
--- a/webapp/packages/plugin-data-viewer/src/locales/ru.ts
+++ b/webapp/packages/plugin-data-viewer/src/locales/ru.ts
@@ -48,7 +48,6 @@ export default [
   ['data_viewer_total_count_tooltip', 'Получить количество записей'],
   ['data_viewer_total_count_failed', 'Не удалось получить количество записей'],
   ['data_viewer_total_count_cancel_failed_title', 'Не удалось отменить получение количества записей'],
-  ['data_viewer_total_count_cancel_failed_desc', 'Подсчет уже завершен'],
   ['data_viewer_model_not_loaded', 'Не удалось загрузить модель таблицы'],
   ['settings_data_editor', 'Редактор данных'],
   ['settings_data_editor_disable_edit_name', 'Отключить редактирование'],
diff --git a/webapp/packages/plugin-data-viewer/src/locales/zh.ts b/webapp/packages/plugin-data-viewer/src/locales/zh.ts
index 908b4b3986..4f6a411e46 100644
--- a/webapp/packages/plugin-data-viewer/src/locales/zh.ts
+++ b/webapp/packages/plugin-data-viewer/src/locales/zh.ts
@@ -54,7 +54,6 @@ export default [
   ['data_viewer_total_count_failed', 'Failed to get total count'],
   ['data_viewer_total_count_tooltip', 'Get total count'],
   ['data_viewer_total_count_cancel_failed_title', 'Failed to cancel getting of total count'],
-  ['data_viewer_total_count_cancel_failed_desc', 'Counting already finished'],
   ['data_viewer_model_not_loaded', 'Table model is not loaded'],
   ['settings_data_editor', 'Data Editor'],
   ['settings_data_editor_disable_edit_name', 'Disable Edit'],

From 2e5f4f8391c3316b567798f084af0c158d1215a5 Mon Sep 17 00:00:00 2001
From: sergeyteleshev <iamsergeyteleshev@gmail.com>
Date: Fri, 12 Apr 2024 10:15:59 +0200
Subject: [PATCH 12/36] CB-4270 pr fixes

---
 .../src/DataViewerTableService.ts             |  1 -
 .../DatabaseDataModel/DatabaseDataSource.ts   |  7 ++--
 .../DatabaseDataModel/IDatabaseDataSource.ts  |  6 ++--
 .../src/ResultSetDataSource.ts                | 32 ++++++++-----------
 .../TableFooter/TableFooterRowCount.tsx       |  7 ++--
 .../SqlResultTabs/SqlQueryResultService.ts    |  6 ----
 6 files changed, 24 insertions(+), 35 deletions(-)

diff --git a/webapp/packages/plugin-data-viewer/src/DataViewerTableService.ts b/webapp/packages/plugin-data-viewer/src/DataViewerTableService.ts
index ccb4e7c235..0d38871517 100644
--- a/webapp/packages/plugin-data-viewer/src/DataViewerTableService.ts
+++ b/webapp/packages/plugin-data-viewer/src/DataViewerTableService.ts
@@ -45,7 +45,6 @@ export class DataViewerTableService {
 
     if (model) {
       this.tableViewerStorageService.remove(tableId);
-      model.source.cancelLoadTotalCount(0);
       await model.dispose();
     }
   }
diff --git a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
index 498e9e506c..22ea760d6a 100644
--- a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
@@ -33,7 +33,7 @@ export abstract class DatabaseDataSource<TOptions, TResult extends IDatabaseData
   error: Error | null;
   executionContext: IConnectionExecutionContext | null;
   outdated: boolean;
-  abstract cancelLoadTotalCountTasks?: Map<string, ITask<number>>;
+  cancelLoadTotalCountTask: ITask<number> | null;
 
   get canCancel(): boolean {
     if (this.activeTask instanceof Task) {
@@ -60,6 +60,7 @@ export abstract class DatabaseDataSource<TOptions, TResult extends IDatabaseData
 
   constructor(serviceInjector: IServiceInjector) {
     this.serviceInjector = serviceInjector;
+    this.cancelLoadTotalCountTask = null;
     this.actions = new DatabaseDataActions(this);
     this.access = DatabaseDataAccessMode.Default;
     this.results = [];
@@ -373,9 +374,9 @@ export abstract class DatabaseDataSource<TOptions, TResult extends IDatabaseData
   abstract request(prevResults: TResult[]): TResult[] | Promise<TResult[]>;
   abstract save(prevResults: TResult[]): Promise<TResult[]> | TResult[];
 
-  abstract dispose(): Promise<void>;
+  abstract dispose(): void;
   abstract loadTotalCount(resultIndex: number): Promise<void>;
-  abstract cancelLoadTotalCount(resultIndex: number): Promise<void>;
+  abstract cancelLoadTotalCount(): ITask<number> | null;
 
   async requestDataAction(): Promise<TResult[] | null> {
     this.prevOptions = toJS(this.options);
diff --git a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/IDatabaseDataSource.ts b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/IDatabaseDataSource.ts
index cb958d7463..0e4a09b05d 100644
--- a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/IDatabaseDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/IDatabaseDataSource.ts
@@ -48,7 +48,7 @@ export interface IDatabaseDataSource<TOptions, TResult extends IDatabaseDataResu
   readonly cancelled: boolean;
   readonly serviceInjector: IServiceInjector;
   readonly outdated: boolean;
-  readonly cancelLoadTotalCountTasks?: Map<string, ITask<number>>;
+  readonly cancelLoadTotalCountTask: ITask<number> | null;
 
   isLoadable: () => boolean;
   isReadonly: (resultIndex: number) => boolean;
@@ -85,7 +85,7 @@ export interface IDatabaseDataSource<TOptions, TResult extends IDatabaseDataResu
   setExecutionContext: (context: IConnectionExecutionContext | null) => this;
   setTotalCount: (resultIndex: number, count: number) => this;
   loadTotalCount: (resultIndex: number) => Promise<void>;
-  cancelLoadTotalCount: (resultIndex: number) => Promise<void>;
+  cancelLoadTotalCount: () => ITask<number> | null;
 
   retry: () => Promise<void>;
   /** Allows to perform an asynchronous action on the data source, this action will wait previous action to finish and save or load requests.
@@ -97,5 +97,5 @@ export interface IDatabaseDataSource<TOptions, TResult extends IDatabaseDataResu
   cancel: () => Promise<void> | void;
   clearError: () => this;
   resetData: () => this;
-  dispose: () => Promise<void>;
+  dispose: () => Promise<void> | void;
 }
diff --git a/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts b/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
index 4751ee538c..4ceba97056 100644
--- a/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
@@ -15,11 +15,11 @@ import { DatabaseDataSource } from './DatabaseDataModel/DatabaseDataSource';
 import type { IDatabaseResultSet } from './DatabaseDataModel/IDatabaseResultSet';
 
 interface ResultSetDataSourceState {
-  cancelLoadTotalCountTasks: Map<string, ITask<number>>;
+  cancelLoadTotalCountTask: ITask<number> | null;
 }
 
 export abstract class ResultSetDataSource<TOptions> extends DatabaseDataSource<TOptions, IDatabaseResultSet> {
-  readonly cancelLoadTotalCountTasks: Map<string, ITask<number>>;
+  cancelLoadTotalCountTask: ITask<number> | null;
 
   constructor(
     readonly serviceInjector: IServiceInjector,
@@ -28,27 +28,19 @@ export abstract class ResultSetDataSource<TOptions> extends DatabaseDataSource<T
   ) {
     super(serviceInjector);
 
-    this.cancelLoadTotalCountTasks = new Map();
+    this.cancelLoadTotalCountTask = null;
 
     makeObservable<ResultSetDataSourceState>(this, {
-      cancelLoadTotalCountTasks: observable.shallow,
+      cancelLoadTotalCountTask: observable.ref,
     });
   }
 
-  async cancelLoadTotalCount(resultIndex: number) {
-    const result = this.getResult(resultIndex);
-
-    if (!result?.id) {
-      throw new Error('Result id must be provided');
+  cancelLoadTotalCount() {
+    if (!this.cancelLoadTotalCountTask?.cancelled && this.results.length === 0) {
+      this.cancelLoadTotalCountTask?.cancel();
     }
 
-    const task = this.cancelLoadTotalCountTasks.get(result.id);
-
-    if (!task?.cancelled) {
-      task?.cancel();
-    }
-
-    this.cancelLoadTotalCountTasks.delete(result.id);
+    return this.cancelLoadTotalCountTask;
   }
 
   async loadTotalCount(resultIndex: number) {
@@ -87,14 +79,18 @@ export abstract class ResultSetDataSource<TOptions> extends DatabaseDataSource<T
       () => this.asyncTaskInfoService.remove(asyncTask.id),
     );
 
-    this.cancelLoadTotalCountTasks.set(result.id, task);
+    this.cancelLoadTotalCountTask = task;
 
     try {
       const count = await task;
 
       this.setTotalCount(resultIndex, count);
     } finally {
-      this.cancelLoadTotalCountTasks.delete(result.id);
+      this.cancelLoadTotalCountTask = null;
     }
   }
+
+  dispose(): void {
+    this.cancelLoadTotalCount();
+  }
 }
diff --git a/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TableFooterRowCount.tsx b/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TableFooterRowCount.tsx
index 41bbcc34d4..c6bd62ef9b 100644
--- a/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TableFooterRowCount.tsx
+++ b/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TableFooterRowCount.tsx
@@ -38,13 +38,12 @@ export const TableFooterRowCount: React.FC<Props> = observer(function TableFoote
     }
   }
 
-  async function cancelTotalCount() {
+  function cancelTotalCount() {
     try {
       setLoading(false);
-      await model.source.cancelLoadTotalCount(resultIndex);
+      model.source.cancelLoadTotalCount();
     } catch (e: any) {
-      const result = model.getResult(resultIndex);
-      const cancelled = Boolean(result?.id && model.source.cancelLoadTotalCountTasks?.get(result?.id)?.cancelled);
+      const cancelled = model.source.cancelLoadTotalCountTask?.cancelled;
 
       if (!cancelled) {
         notificationService.logException(e, 'data_viewer_total_count_cancel_failed_title', typeof e === 'string' ? e : undefined);
diff --git a/webapp/packages/plugin-sql-editor/src/SqlResultTabs/SqlQueryResultService.ts b/webapp/packages/plugin-sql-editor/src/SqlResultTabs/SqlQueryResultService.ts
index 20531984c2..5d0d8d31ac 100644
--- a/webapp/packages/plugin-sql-editor/src/SqlResultTabs/SqlQueryResultService.ts
+++ b/webapp/packages/plugin-sql-editor/src/SqlResultTabs/SqlQueryResultService.ts
@@ -153,12 +153,6 @@ export class SqlQueryResultService {
     const resultTab = state.resultTabs.find(resultTab => resultTab.tabId === tabId);
     const group = state.resultGroups.find(group => group.groupId === resultTab?.groupId);
 
-    const model = group && this.tableViewerStorageService.get(group.modelId);
-
-    if (model && resultTab) {
-      model.source.cancelLoadTotalCount(resultTab.indexInResultSet);
-    }
-
     if (resultTab && group) {
       state.resultTabs.splice(state.resultTabs.indexOf(resultTab), 1);
 

From 96f4ea0bfc2380870a6b6c4620cb14b07e511749 Mon Sep 17 00:00:00 2001
From: sergeyteleshev <iamsergeyteleshev@gmail.com>
Date: Fri, 12 Apr 2024 10:52:41 +0200
Subject: [PATCH 13/36] =?UTF-8?q?=D0=A1B-4270=20cancel=20load=20total=20co?=
 =?UTF-8?q?unt=20for=20refresh=20action?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../src/DatabaseDataModel/DatabaseDataModel.ts    |  1 +
 .../src/DatabaseDataModel/DatabaseDataSource.ts   |  1 +
 .../plugin-data-viewer/src/ResultSetDataSource.ts | 15 ---------------
 3 files changed, 2 insertions(+), 15 deletions(-)

diff --git a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataModel.ts b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataModel.ts
index d800d469ed..4cebeeda00 100644
--- a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataModel.ts
+++ b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataModel.ts
@@ -129,6 +129,7 @@ export class DatabaseDataModel<TOptions, TResult extends IDatabaseDataResult = I
   }
 
   async refresh(concurrent?: boolean): Promise<void> {
+    this.source.cancelLoadTotalCount();
     if (concurrent) {
       await this.source.refreshData();
       return;
diff --git a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
index 22ea760d6a..8e6fd87b6f 100644
--- a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
@@ -97,6 +97,7 @@ export abstract class DatabaseDataSource<TOptions, TResult extends IDatabaseData
       prevOptions: observable,
       options: observable,
       requestInfo: observable,
+      cancelLoadTotalCountTask: observable.ref,
       error: observable.ref,
       executionContext: observable,
       disabled: observable,
diff --git a/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts b/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
index 4ceba97056..4980bc52ad 100644
--- a/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
@@ -5,34 +5,19 @@
  * Licensed under the Apache License, Version 2.0.
  * you may not use this file except in compliance with the License.
  */
-import { makeObservable, observable } from 'mobx';
-
 import type { IServiceInjector } from '@cloudbeaver/core-di';
-import type { ITask } from '@cloudbeaver/core-executor';
 import type { AsyncTaskInfoService, GraphQLService } from '@cloudbeaver/core-sdk';
 
 import { DatabaseDataSource } from './DatabaseDataModel/DatabaseDataSource';
 import type { IDatabaseResultSet } from './DatabaseDataModel/IDatabaseResultSet';
 
-interface ResultSetDataSourceState {
-  cancelLoadTotalCountTask: ITask<number> | null;
-}
-
 export abstract class ResultSetDataSource<TOptions> extends DatabaseDataSource<TOptions, IDatabaseResultSet> {
-  cancelLoadTotalCountTask: ITask<number> | null;
-
   constructor(
     readonly serviceInjector: IServiceInjector,
     protected graphQLService: GraphQLService,
     protected asyncTaskInfoService: AsyncTaskInfoService,
   ) {
     super(serviceInjector);
-
-    this.cancelLoadTotalCountTask = null;
-
-    makeObservable<ResultSetDataSourceState>(this, {
-      cancelLoadTotalCountTask: observable.ref,
-    });
   }
 
   cancelLoadTotalCount() {

From 8417093cf2a7eb46c8e00c2cc45cc46ba67090c4 Mon Sep 17 00:00:00 2001
From: sergeyteleshev <iamsergeyteleshev@gmail.com>
Date: Fri, 12 Apr 2024 11:08:22 +0200
Subject: [PATCH 14/36] CB-4270 fix: place cancelLoadTotalCount to
 DatabaseDataSource

---
 .../src/DatabaseDataModel/DatabaseDataSource.ts    | 14 ++++++++++++--
 .../plugin-data-viewer/src/ResultSetDataSource.ts  | 12 ------------
 2 files changed, 12 insertions(+), 14 deletions(-)

diff --git a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
index 8e6fd87b6f..2e639825f2 100644
--- a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
@@ -372,12 +372,22 @@ export abstract class DatabaseDataSource<TOptions, TResult extends IDatabaseData
     return this;
   }
 
+  cancelLoadTotalCount(): ITask<number> | null {
+    if (!this.cancelLoadTotalCountTask?.cancelled) {
+      this.cancelLoadTotalCountTask?.cancel();
+    }
+
+    return this.cancelLoadTotalCountTask;
+  }
+
+  dispose(): Promise<void> | void {
+    this.cancelLoadTotalCount();
+  }
+
   abstract request(prevResults: TResult[]): TResult[] | Promise<TResult[]>;
   abstract save(prevResults: TResult[]): Promise<TResult[]> | TResult[];
 
-  abstract dispose(): void;
   abstract loadTotalCount(resultIndex: number): Promise<void>;
-  abstract cancelLoadTotalCount(): ITask<number> | null;
 
   async requestDataAction(): Promise<TResult[] | null> {
     this.prevOptions = toJS(this.options);
diff --git a/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts b/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
index 4980bc52ad..962af46447 100644
--- a/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
@@ -20,14 +20,6 @@ export abstract class ResultSetDataSource<TOptions> extends DatabaseDataSource<T
     super(serviceInjector);
   }
 
-  cancelLoadTotalCount() {
-    if (!this.cancelLoadTotalCountTask?.cancelled && this.results.length === 0) {
-      this.cancelLoadTotalCountTask?.cancel();
-    }
-
-    return this.cancelLoadTotalCountTask;
-  }
-
   async loadTotalCount(resultIndex: number) {
     const executionContext = this.executionContext;
     const executionContextInfo = this.executionContext?.context;
@@ -74,8 +66,4 @@ export abstract class ResultSetDataSource<TOptions> extends DatabaseDataSource<T
       this.cancelLoadTotalCountTask = null;
     }
   }
-
-  dispose(): void {
-    this.cancelLoadTotalCount();
-  }
 }

From 13f3eae8fb7ccd3588d43b0cf706b622cbbcc634 Mon Sep 17 00:00:00 2001
From: sergeyteleshev <iamsergeyteleshev@gmail.com>
Date: Fri, 12 Apr 2024 11:09:14 +0200
Subject: [PATCH 15/36] CB-4270 fix: cancel load rows in refresh data in data
 source

---
 .../src/DatabaseDataModel/DatabaseDataModel.ts                  | 1 -
 .../src/DatabaseDataModel/DatabaseDataSource.ts                 | 2 ++
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataModel.ts b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataModel.ts
index 4cebeeda00..d800d469ed 100644
--- a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataModel.ts
+++ b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataModel.ts
@@ -129,7 +129,6 @@ export class DatabaseDataModel<TOptions, TResult extends IDatabaseDataResult = I
   }
 
   async refresh(concurrent?: boolean): Promise<void> {
-    this.source.cancelLoadTotalCount();
     if (concurrent) {
       await this.source.refreshData();
       return;
diff --git a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
index 2e639825f2..4784c179c9 100644
--- a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
@@ -328,6 +328,8 @@ export abstract class DatabaseDataSource<TOptions, TResult extends IDatabaseData
   }
 
   async refreshData(): Promise<void> {
+    this.cancelLoadTotalCount();
+
     if (this.prevOptions) {
       this.options = toJS(this.prevOptions);
     }

From 62d282f157018aebd1818ae32038598b086197a4 Mon Sep 17 00:00:00 2001
From: sergeyteleshev <iamsergeyteleshev@gmail.com>
Date: Fri, 12 Apr 2024 11:29:18 +0200
Subject: [PATCH 16/36] CB-4270 cancels load on close result set tab

---
 .../plugin-sql-editor/src/SqlResultTabs/SqlQueryResultService.ts | 1 +
 1 file changed, 1 insertion(+)

diff --git a/webapp/packages/plugin-sql-editor/src/SqlResultTabs/SqlQueryResultService.ts b/webapp/packages/plugin-sql-editor/src/SqlResultTabs/SqlQueryResultService.ts
index 5d0d8d31ac..74c59f9310 100644
--- a/webapp/packages/plugin-sql-editor/src/SqlResultTabs/SqlQueryResultService.ts
+++ b/webapp/packages/plugin-sql-editor/src/SqlResultTabs/SqlQueryResultService.ts
@@ -163,6 +163,7 @@ export class SqlQueryResultService {
 
         // TODO: we need to dispose table model, but don't close execution context, so now we only
         const model = this.tableViewerStorageService.get(group.modelId);
+        model?.source.cancelLoadTotalCount();
         // model?.dispose();
 
         if (model?.isLoading()) {

From 10623cead066ccbba92f5df00f70d587ca98b742 Mon Sep 17 00:00:00 2001
From: sergeyteleshev <iamsergeyteleshev@gmail.com>
Date: Fri, 12 Apr 2024 12:03:14 +0200
Subject: [PATCH 17/36] CB-4270 refactor: move cancelLoadTotalCount to
 resultsetdatasource

---
 .../src/DatabaseDataModel/DatabaseDataSource.ts    | 14 ++------------
 .../plugin-data-viewer/src/ResultSetDataSource.ts  | 13 +++++++++++++
 2 files changed, 15 insertions(+), 12 deletions(-)

diff --git a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
index 4784c179c9..1a3ed3ac41 100644
--- a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
@@ -374,21 +374,11 @@ export abstract class DatabaseDataSource<TOptions, TResult extends IDatabaseData
     return this;
   }
 
-  cancelLoadTotalCount(): ITask<number> | null {
-    if (!this.cancelLoadTotalCountTask?.cancelled) {
-      this.cancelLoadTotalCountTask?.cancel();
-    }
-
-    return this.cancelLoadTotalCountTask;
-  }
-
-  dispose(): Promise<void> | void {
-    this.cancelLoadTotalCount();
-  }
-
   abstract request(prevResults: TResult[]): TResult[] | Promise<TResult[]>;
   abstract save(prevResults: TResult[]): Promise<TResult[]> | TResult[];
 
+  abstract cancelLoadTotalCount(): ITask<number> | null;
+  abstract dispose(): Promise<void> | void;
   abstract loadTotalCount(resultIndex: number): Promise<void>;
 
   async requestDataAction(): Promise<TResult[] | null> {
diff --git a/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts b/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
index 962af46447..eb24e16365 100644
--- a/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
@@ -6,6 +6,7 @@
  * you may not use this file except in compliance with the License.
  */
 import type { IServiceInjector } from '@cloudbeaver/core-di';
+import type { ITask } from '@cloudbeaver/core-executor';
 import type { AsyncTaskInfoService, GraphQLService } from '@cloudbeaver/core-sdk';
 
 import { DatabaseDataSource } from './DatabaseDataModel/DatabaseDataSource';
@@ -20,6 +21,18 @@ export abstract class ResultSetDataSource<TOptions> extends DatabaseDataSource<T
     super(serviceInjector);
   }
 
+  dispose(): void | Promise<void> {
+    this.cancelLoadTotalCount();
+  }
+
+  cancelLoadTotalCount(): ITask<number> | null {
+    if (!this.cancelLoadTotalCountTask?.cancelled) {
+      this.cancelLoadTotalCountTask?.cancel();
+    }
+
+    return this.cancelLoadTotalCountTask;
+  }
+
   async loadTotalCount(resultIndex: number) {
     const executionContext = this.executionContext;
     const executionContextInfo = this.executionContext?.context;

From 1fbb20c1f5cbbbf180db845a5ea630477f98d401 Mon Sep 17 00:00:00 2001
From: sergeyteleshev <iamsergeyteleshev@gmail.com>
Date: Fri, 12 Apr 2024 15:33:01 +0200
Subject: [PATCH 18/36] CB-4270 cancel load total count execute only for
 resultsetdatasource

---
 .../src/DatabaseDataModel/DatabaseDataSource.ts             | 2 --
 .../packages/plugin-data-viewer/src/ResultSetDataSource.ts  | 6 ++++++
 2 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
index 1a3ed3ac41..02f7aa469c 100644
--- a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
@@ -328,8 +328,6 @@ export abstract class DatabaseDataSource<TOptions, TResult extends IDatabaseData
   }
 
   async refreshData(): Promise<void> {
-    this.cancelLoadTotalCount();
-
     if (this.prevOptions) {
       this.options = toJS(this.prevOptions);
     }
diff --git a/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts b/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
index eb24e16365..277d728a25 100644
--- a/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
@@ -25,6 +25,12 @@ export abstract class ResultSetDataSource<TOptions> extends DatabaseDataSource<T
     this.cancelLoadTotalCount();
   }
 
+  async refreshData(): Promise<void> {
+    this.cancelLoadTotalCount();
+
+    await super.refreshData();
+  }
+
   cancelLoadTotalCount(): ITask<number> | null {
     if (!this.cancelLoadTotalCountTask?.cancelled) {
       this.cancelLoadTotalCountTask?.cancel();

From c96508eb81399de2dfab8136edba710a3848cf15 Mon Sep 17 00:00:00 2001
From: sergeyteleshev <iamsergeyteleshev@gmail.com>
Date: Fri, 12 Apr 2024 15:58:00 +0200
Subject: [PATCH 19/36] CB-4270 adds constraint for TableFooterRowCount

---
 webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts | 2 +-
 .../src/TableViewer/TableFooter/TableFooter.tsx               | 4 +++-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts b/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
index 277d728a25..e5175208c0 100644
--- a/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
@@ -21,7 +21,7 @@ export abstract class ResultSetDataSource<TOptions> extends DatabaseDataSource<T
     super(serviceInjector);
   }
 
-  dispose(): void | Promise<void> {
+  dispose(): void {
     this.cancelLoadTotalCount();
   }
 
diff --git a/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TableFooter.tsx b/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TableFooter.tsx
index fc61db52d2..146bbff93a 100644
--- a/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TableFooter.tsx
+++ b/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TableFooter.tsx
@@ -13,6 +13,7 @@ import { Form, getComputed, ToolsPanel } from '@cloudbeaver/core-blocks';
 import type { IDataContext } from '@cloudbeaver/core-data-context';
 import { useService } from '@cloudbeaver/core-di';
 
+import { ResultSetConstraintAction } from '../../DatabaseDataModel/Actions/ResultSet/ResultSetConstraintAction';
 import type { IDatabaseDataModel } from '../../DatabaseDataModel/IDatabaseDataModel';
 import { DataViewerSettingsService } from '../../DataViewerSettingsService';
 import { AutoRefreshButton } from './AutoRefresh/AutoRefreshButton';
@@ -94,6 +95,7 @@ export const TableFooter = observer<Props>(function TableFooter({ resultIndex, m
   }, [model.countGain]);
 
   const disabled = getComputed(() => model.isLoading() || model.isDisabled(resultIndex));
+  const constraint = model.getResult(resultIndex) ? model.source.getAction(resultIndex, ResultSetConstraintAction) : null;
 
   return styled(tableFooterStyles)(
     <ToolsPanel type="secondary">
@@ -116,7 +118,7 @@ export const TableFooter = observer<Props>(function TableFooter({ resultIndex, m
           />
         </Form>
       </count>
-      <TableFooterRowCount model={model} resultIndex={resultIndex} />
+      {constraint?.supported && <TableFooterRowCount model={model} resultIndex={resultIndex} />}
       <TableFooterMenu model={model} resultIndex={resultIndex} simple={simple} context={context} />
       {model.source.requestInfo.requestMessage.length > 0 && (
         <time>

From 12c3d99991b2c006f88ff5f55a5efcb1913a783b Mon Sep 17 00:00:00 2001
From: sergeyteleshev <iamsergeyteleshev@gmail.com>
Date: Fri, 12 Apr 2024 16:21:57 +0200
Subject: [PATCH 20/36] CB-4270 fix: do not load new value if cancelled rows
 loading

---
 .../plugin-data-viewer/src/ResultSetDataSource.ts        | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts b/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
index e5175208c0..e22d9b4f61 100644
--- a/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
@@ -67,6 +67,11 @@ export abstract class ResultSetDataSource<TOptions> extends DatabaseDataSource<T
     const task = executionContext.run(
       async () => {
         const info = await this.asyncTaskInfoService.run(asyncTask);
+
+        if (this.cancelLoadTotalCountTask?.cancelled) {
+          return this.count;
+        }
+
         const { count } = await this.graphQLService.sdk.getSqlRowDataCountResult({ taskId: info.id });
 
         return count;
@@ -80,7 +85,9 @@ export abstract class ResultSetDataSource<TOptions> extends DatabaseDataSource<T
     try {
       const count = await task;
 
-      this.setTotalCount(resultIndex, count);
+      if (count !== this.count) {
+        this.setTotalCount(resultIndex, count);
+      }
     } finally {
       this.cancelLoadTotalCountTask = null;
     }

From f8bffa253ad4c83c41108c0f582fd7db8770b7bf Mon Sep 17 00:00:00 2001
From: sergeyteleshev <iamsergeyteleshev@gmail.com>
Date: Fri, 12 Apr 2024 16:22:35 +0200
Subject: [PATCH 21/36] CB-4270 cleanup

---
 .../src/DatabaseDataModel/DatabaseDataSource.ts               | 2 +-
 webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts | 4 +---
 2 files changed, 2 insertions(+), 4 deletions(-)

diff --git a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
index 02f7aa469c..dc162eaa52 100644
--- a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
@@ -261,7 +261,7 @@ export abstract class DatabaseDataSource<TOptions, TResult extends IDatabaseData
   setTotalCount(resultIndex: number, count: number): this {
     const result = this.getResult(resultIndex);
 
-    if (result) {
+    if (result && result.totalCount !== count) {
       result.totalCount = count;
     }
     return this;
diff --git a/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts b/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
index e22d9b4f61..ec16d60a0a 100644
--- a/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
@@ -85,9 +85,7 @@ export abstract class ResultSetDataSource<TOptions> extends DatabaseDataSource<T
     try {
       const count = await task;
 
-      if (count !== this.count) {
-        this.setTotalCount(resultIndex, count);
-      }
+      this.setTotalCount(resultIndex, count);
     } finally {
       this.cancelLoadTotalCountTask = null;
     }

From 24f2e1225acfd2d8d599fa3c7eae8cc1aace8dfc Mon Sep 17 00:00:00 2001
From: sergeyteleshev <iamsergeyteleshev@gmail.com>
Date: Fri, 12 Apr 2024 16:23:11 +0200
Subject: [PATCH 22/36] Revert "CB-4270 cleanup"

This reverts commit f8bffa253ad4c83c41108c0f582fd7db8770b7bf.
---
 .../src/DatabaseDataModel/DatabaseDataSource.ts               | 2 +-
 webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts | 4 +++-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
index dc162eaa52..02f7aa469c 100644
--- a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
@@ -261,7 +261,7 @@ export abstract class DatabaseDataSource<TOptions, TResult extends IDatabaseData
   setTotalCount(resultIndex: number, count: number): this {
     const result = this.getResult(resultIndex);
 
-    if (result && result.totalCount !== count) {
+    if (result) {
       result.totalCount = count;
     }
     return this;
diff --git a/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts b/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
index ec16d60a0a..e22d9b4f61 100644
--- a/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
@@ -85,7 +85,9 @@ export abstract class ResultSetDataSource<TOptions> extends DatabaseDataSource<T
     try {
       const count = await task;
 
-      this.setTotalCount(resultIndex, count);
+      if (count !== this.count) {
+        this.setTotalCount(resultIndex, count);
+      }
     } finally {
       this.cancelLoadTotalCountTask = null;
     }

From 95fa2edfe72e2e49044a6aefe3918e4f8f65185a Mon Sep 17 00:00:00 2001
From: sergeyteleshev <iamsergeyteleshev@gmail.com>
Date: Fri, 12 Apr 2024 17:17:44 +0200
Subject: [PATCH 23/36] CB-4270 pr fixes

---
 .../TableFooter/CancelTotalCountAction.tsx    | 26 +++++++++
 .../TableFooter/TableFooterRowCount.tsx       | 55 +------------------
 .../TableFooter/TotalCountAction.tsx          | 47 ++++++++++++++++
 .../plugin-data-viewer/src/locales/en.ts      |  1 -
 .../plugin-data-viewer/src/locales/it.ts      |  1 -
 .../plugin-data-viewer/src/locales/ru.ts      |  1 -
 .../plugin-data-viewer/src/locales/zh.ts      |  1 -
 7 files changed, 76 insertions(+), 56 deletions(-)
 create mode 100644 webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/CancelTotalCountAction.tsx
 create mode 100644 webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TotalCountAction.tsx

diff --git a/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/CancelTotalCountAction.tsx b/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/CancelTotalCountAction.tsx
new file mode 100644
index 0000000000..ab4280cd74
--- /dev/null
+++ b/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/CancelTotalCountAction.tsx
@@ -0,0 +1,26 @@
+/*
+ * CloudBeaver - Cloud Database Manager
+ * Copyright (C) 2020-2024 DBeaver Corp and others
+ *
+ * Licensed under the Apache License, Version 2.0.
+ * you may not use this file except in compliance with the License.
+ */
+import { observer } from 'mobx-react-lite';
+import styled from 'reshadow';
+
+import { ToolsAction, useTranslate } from '@cloudbeaver/core-blocks';
+
+import { tableFooterMenuStyles } from './TableFooterMenu/TableFooterMenuItem';
+import classes from './TableFooterRowCount.m.css';
+
+export const CancelTotalCountAction = observer(function CancelTotalCountAction({ onClick }: { onClick: VoidFunction }) {
+  const translate = useTranslate();
+
+  return styled(tableFooterMenuStyles)(
+    <div className={classes.wrapper} title={translate('ui_processing_cancel')}>
+      <ToolsAction icon="/icons/data_cancel.svg" viewBox="0 0 32 32" onClick={onClick}>
+        {translate('ui_processing_cancel')}
+      </ToolsAction>
+    </div>,
+  );
+});
diff --git a/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TableFooterRowCount.tsx b/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TableFooterRowCount.tsx
index c6bd62ef9b..d7b6116920 100644
--- a/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TableFooterRowCount.tsx
+++ b/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TableFooterRowCount.tsx
@@ -7,16 +7,14 @@
  */
 import { observer } from 'mobx-react-lite';
 import { useState } from 'react';
-import styled from 'reshadow';
 
-import { getComputed, ToolsAction, useTranslate } from '@cloudbeaver/core-blocks';
 import { useService } from '@cloudbeaver/core-di';
 import { NotificationService } from '@cloudbeaver/core-events';
 
 import type { IDatabaseDataModel } from '../../DatabaseDataModel/IDatabaseDataModel';
 import type { IDatabaseResultSet } from '../../DatabaseDataModel/IDatabaseResultSet';
-import { tableFooterMenuStyles } from './TableFooterMenu/TableFooterMenuItem';
-import classes from './TableFooterRowCount.m.css';
+import { CancelTotalCountAction } from './CancelTotalCountAction';
+import { TotalCountAction } from './TotalCountAction';
 
 interface Props {
   resultIndex: number;
@@ -43,11 +41,7 @@ export const TableFooterRowCount: React.FC<Props> = observer(function TableFoote
       setLoading(false);
       model.source.cancelLoadTotalCount();
     } catch (e: any) {
-      const cancelled = model.source.cancelLoadTotalCountTask?.cancelled;
-
-      if (!cancelled) {
-        notificationService.logException(e, 'data_viewer_total_count_cancel_failed_title', typeof e === 'string' ? e : undefined);
-      }
+      notificationService.logException(e);
     }
   }
 
@@ -57,46 +51,3 @@ export const TableFooterRowCount: React.FC<Props> = observer(function TableFoote
 
   return <TotalCountAction loading={loading} resultIndex={resultIndex} model={model} onClick={loadTotalCount} />;
 });
-
-const CancelTotalCountAction = observer(function CancelTotalCountAction({ onClick }: { onClick: VoidFunction }) {
-  const translate = useTranslate();
-
-  return styled(tableFooterMenuStyles)(
-    <div className={classes.wrapper} title={translate('ui_processing_cancel')}>
-      <ToolsAction icon="/icons/data_cancel.svg" viewBox="0 0 32 32" onClick={onClick}>
-        {translate('ui_processing_cancel')}
-      </ToolsAction>
-    </div>,
-  );
-});
-
-const TotalCountAction = observer(function TotalCountAction({
-  onClick,
-  loading,
-  resultIndex,
-  model,
-}: {
-  onClick: VoidFunction;
-  loading: boolean;
-  resultIndex: number;
-  model: IDatabaseDataModel<any, IDatabaseResultSet>;
-}) {
-  const result = model.getResult(resultIndex);
-  const translate = useTranslate();
-  const disabled = getComputed(() => model.isLoading() || model.isDisabled(resultIndex));
-
-  if (!result) {
-    return null;
-  }
-
-  const currentCount = result.loadedFully ? result.count : `${result.count}+`;
-  const count = result.totalCount ?? currentCount;
-
-  return styled(tableFooterMenuStyles)(
-    <div className={classes.wrapper} title={translate('data_viewer_total_count_tooltip')}>
-      <ToolsAction disabled={disabled} loading={loading} icon="/icons/data_row_count.svg" viewBox="0 0 32 32" onClick={onClick}>
-        <span>{count}</span>
-      </ToolsAction>
-    </div>,
-  );
-});
diff --git a/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TotalCountAction.tsx b/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TotalCountAction.tsx
new file mode 100644
index 0000000000..754d0c7e27
--- /dev/null
+++ b/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TotalCountAction.tsx
@@ -0,0 +1,47 @@
+/*
+ * CloudBeaver - Cloud Database Manager
+ * Copyright (C) 2020-2024 DBeaver Corp and others
+ *
+ * Licensed under the Apache License, Version 2.0.
+ * you may not use this file except in compliance with the License.
+ */
+import { observer } from 'mobx-react-lite';
+import styled from 'reshadow';
+
+import { getComputed, ToolsAction, useTranslate } from '@cloudbeaver/core-blocks';
+
+import type { IDatabaseDataModel } from '../../DatabaseDataModel/IDatabaseDataModel';
+import type { IDatabaseResultSet } from '../../DatabaseDataModel/IDatabaseResultSet';
+import { tableFooterMenuStyles } from './TableFooterMenu/TableFooterMenuItem';
+import classes from './TableFooterRowCount.m.css';
+
+export const TotalCountAction = observer(function TotalCountAction({
+  onClick,
+  loading,
+  resultIndex,
+  model,
+}: {
+  onClick: VoidFunction;
+  loading: boolean;
+  resultIndex: number;
+  model: IDatabaseDataModel<any, IDatabaseResultSet>;
+}) {
+  const result = model.getResult(resultIndex);
+  const translate = useTranslate();
+  const disabled = getComputed(() => model.isLoading() || model.isDisabled(resultIndex));
+
+  if (!result) {
+    return null;
+  }
+
+  const currentCount = result.loadedFully ? result.count : `${result.count}+`;
+  const count = result.totalCount ?? currentCount;
+
+  return styled(tableFooterMenuStyles)(
+    <div className={classes.wrapper} title={translate('data_viewer_total_count_tooltip')}>
+      <ToolsAction disabled={disabled} loading={loading} icon="/icons/data_row_count.svg" viewBox="0 0 32 32" onClick={onClick}>
+        <span>{count}</span>
+      </ToolsAction>
+    </div>,
+  );
+});
diff --git a/webapp/packages/plugin-data-viewer/src/locales/en.ts b/webapp/packages/plugin-data-viewer/src/locales/en.ts
index 256c8dadf6..0e525a78f9 100644
--- a/webapp/packages/plugin-data-viewer/src/locales/en.ts
+++ b/webapp/packages/plugin-data-viewer/src/locales/en.ts
@@ -53,7 +53,6 @@ export default [
   ['data_viewer_refresh_result_set', 'Refresh result set'],
   ['data_viewer_total_count_tooltip', 'Get total count'],
   ['data_viewer_total_count_failed', 'Failed to get total count'],
-  ['data_viewer_total_count_cancel_failed_title', 'Failed to cancel getting of total count'],
   ['data_viewer_model_not_loaded', 'Table model is not loaded'],
   ['settings_data_editor', 'Data Editor'],
   ['settings_data_editor_disable_edit_name', 'Disable Edit'],
diff --git a/webapp/packages/plugin-data-viewer/src/locales/it.ts b/webapp/packages/plugin-data-viewer/src/locales/it.ts
index 0d2d5487fe..7c1a13b91c 100644
--- a/webapp/packages/plugin-data-viewer/src/locales/it.ts
+++ b/webapp/packages/plugin-data-viewer/src/locales/it.ts
@@ -46,7 +46,6 @@ export default [
   ['data_viewer_refresh_result_set', 'Refresh result set'],
   ['data_viewer_total_count_tooltip', 'Get total count'],
   ['data_viewer_total_count_failed', 'Failed to get total count'],
-  ['data_viewer_total_count_cancel_failed_title', 'Failed to cancel getting of total count'],
   ['data_viewer_model_not_loaded', 'Table model is not loaded'],
   ['settings_data_editor', 'Data Editor'],
   ['settings_data_editor_disable_edit_name', 'Disable Edit'],
diff --git a/webapp/packages/plugin-data-viewer/src/locales/ru.ts b/webapp/packages/plugin-data-viewer/src/locales/ru.ts
index 1428ae7971..cd1e69c78a 100644
--- a/webapp/packages/plugin-data-viewer/src/locales/ru.ts
+++ b/webapp/packages/plugin-data-viewer/src/locales/ru.ts
@@ -47,7 +47,6 @@ export default [
   ['data_viewer_refresh_result_set', 'Обновить резалт сет'],
   ['data_viewer_total_count_tooltip', 'Получить количество записей'],
   ['data_viewer_total_count_failed', 'Не удалось получить количество записей'],
-  ['data_viewer_total_count_cancel_failed_title', 'Не удалось отменить получение количества записей'],
   ['data_viewer_model_not_loaded', 'Не удалось загрузить модель таблицы'],
   ['settings_data_editor', 'Редактор данных'],
   ['settings_data_editor_disable_edit_name', 'Отключить редактирование'],
diff --git a/webapp/packages/plugin-data-viewer/src/locales/zh.ts b/webapp/packages/plugin-data-viewer/src/locales/zh.ts
index 4f6a411e46..368f2f6e8f 100644
--- a/webapp/packages/plugin-data-viewer/src/locales/zh.ts
+++ b/webapp/packages/plugin-data-viewer/src/locales/zh.ts
@@ -53,7 +53,6 @@ export default [
   ['data_viewer_refresh_result_set', 'Refresh result set'],
   ['data_viewer_total_count_failed', 'Failed to get total count'],
   ['data_viewer_total_count_tooltip', 'Get total count'],
-  ['data_viewer_total_count_cancel_failed_title', 'Failed to cancel getting of total count'],
   ['data_viewer_model_not_loaded', 'Table model is not loaded'],
   ['settings_data_editor', 'Data Editor'],
   ['settings_data_editor_disable_edit_name', 'Disable Edit'],

From dee02956889bebe4ac67e9c5d1f844145e5dcd28 Mon Sep 17 00:00:00 2001
From: sergeyteleshev <iamsergeyteleshev@gmail.com>
Date: Mon, 15 Apr 2024 12:14:09 +0200
Subject: [PATCH 24/36] CB-4270 pr fixes

---
 .../DatabaseDataModel/DatabaseDataModel.ts    |  5 +--
 .../DatabaseDataModel/DatabaseDataSource.ts   | 12 +++----
 .../DatabaseDataModel/IDatabaseDataSource.ts  |  8 ++---
 .../src/ResultSetDataSource.ts                | 32 ++++++-------------
 .../TableFooter/CancelTotalCountAction.tsx    |  6 +++-
 .../TableFooter/TableFooterRowCount.tsx       |  4 +--
 .../TableFooter/TotalCountAction.tsx          | 11 +++----
 .../SqlResultTabs/SqlQueryResultService.ts    |  1 -
 8 files changed, 34 insertions(+), 45 deletions(-)

diff --git a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataModel.ts b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataModel.ts
index d800d469ed..00d2f36a14 100644
--- a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataModel.ts
+++ b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataModel.ts
@@ -154,8 +154,9 @@ export class DatabaseDataModel<TOptions, TResult extends IDatabaseDataResult = I
     }
   }
 
-  cancel(): Promise<void> | void {
-    return this.source.cancel();
+  async cancel(): Promise<void> {
+    await this.source.cancelLoadTotalCount();
+    return await this.source.cancel();
   }
 
   resetData(): void {
diff --git a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
index 02f7aa469c..27dad8c500 100644
--- a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
@@ -33,7 +33,7 @@ export abstract class DatabaseDataSource<TOptions, TResult extends IDatabaseData
   error: Error | null;
   executionContext: IConnectionExecutionContext | null;
   outdated: boolean;
-  cancelLoadTotalCountTask: ITask<number> | null;
+  totalCountRequestTask: ITask<number> | null;
 
   get canCancel(): boolean {
     if (this.activeTask instanceof Task) {
@@ -60,7 +60,7 @@ export abstract class DatabaseDataSource<TOptions, TResult extends IDatabaseData
 
   constructor(serviceInjector: IServiceInjector) {
     this.serviceInjector = serviceInjector;
-    this.cancelLoadTotalCountTask = null;
+    this.totalCountRequestTask = null;
     this.actions = new DatabaseDataActions(this);
     this.access = DatabaseDataAccessMode.Default;
     this.results = [];
@@ -97,7 +97,7 @@ export abstract class DatabaseDataSource<TOptions, TResult extends IDatabaseData
       prevOptions: observable,
       options: observable,
       requestInfo: observable,
-      cancelLoadTotalCountTask: observable.ref,
+      totalCountRequestTask: observable.ref,
       error: observable.ref,
       executionContext: observable,
       disabled: observable,
@@ -375,9 +375,9 @@ export abstract class DatabaseDataSource<TOptions, TResult extends IDatabaseData
   abstract request(prevResults: TResult[]): TResult[] | Promise<TResult[]>;
   abstract save(prevResults: TResult[]): Promise<TResult[]> | TResult[];
 
-  abstract cancelLoadTotalCount(): ITask<number> | null;
-  abstract dispose(): Promise<void> | void;
-  abstract loadTotalCount(resultIndex: number): Promise<void>;
+  abstract dispose(): Promise<void>;
+  abstract loadTotalCount(resultIndex: number): Promise<ITask<number> | null>;
+  abstract cancelLoadTotalCount(): Promise<ITask<number> | null>;
 
   async requestDataAction(): Promise<TResult[] | null> {
     this.prevOptions = toJS(this.options);
diff --git a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/IDatabaseDataSource.ts b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/IDatabaseDataSource.ts
index 0e4a09b05d..69e37eef3f 100644
--- a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/IDatabaseDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/IDatabaseDataSource.ts
@@ -48,7 +48,7 @@ export interface IDatabaseDataSource<TOptions, TResult extends IDatabaseDataResu
   readonly cancelled: boolean;
   readonly serviceInjector: IServiceInjector;
   readonly outdated: boolean;
-  readonly cancelLoadTotalCountTask: ITask<number> | null;
+  readonly totalCountRequestTask: ITask<number> | null;
 
   isLoadable: () => boolean;
   isReadonly: (resultIndex: number) => boolean;
@@ -84,8 +84,8 @@ export interface IDatabaseDataSource<TOptions, TResult extends IDatabaseDataResu
   setSupportedDataFormats: (dataFormats: ResultDataFormat[]) => this;
   setExecutionContext: (context: IConnectionExecutionContext | null) => this;
   setTotalCount: (resultIndex: number, count: number) => this;
-  loadTotalCount: (resultIndex: number) => Promise<void>;
-  cancelLoadTotalCount: () => ITask<number> | null;
+  loadTotalCount: (resultIndex: number) => Promise<ITask<number> | null>;
+  cancelLoadTotalCount: () => Promise<ITask<number> | null>;
 
   retry: () => Promise<void>;
   /** Allows to perform an asynchronous action on the data source, this action will wait previous action to finish and save or load requests.
@@ -97,5 +97,5 @@ export interface IDatabaseDataSource<TOptions, TResult extends IDatabaseDataResu
   cancel: () => Promise<void> | void;
   clearError: () => this;
   resetData: () => this;
-  dispose: () => Promise<void> | void;
+  dispose: () => Promise<void>;
 }
diff --git a/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts b/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
index e22d9b4f61..3f8cdb5440 100644
--- a/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
@@ -21,22 +21,14 @@ export abstract class ResultSetDataSource<TOptions> extends DatabaseDataSource<T
     super(serviceInjector);
   }
 
-  dispose(): void {
-    this.cancelLoadTotalCount();
+  async dispose(): Promise<void> {
+    await this.cancelLoadTotalCount();
   }
 
-  async refreshData(): Promise<void> {
-    this.cancelLoadTotalCount();
+  async cancelLoadTotalCount(): Promise<ITask<number> | null> {
+    await this.totalCountRequestTask?.cancel();
 
-    await super.refreshData();
-  }
-
-  cancelLoadTotalCount(): ITask<number> | null {
-    if (!this.cancelLoadTotalCountTask?.cancelled) {
-      this.cancelLoadTotalCountTask?.cancel();
-    }
-
-    return this.cancelLoadTotalCountTask;
+    return this.totalCountRequestTask;
   }
 
   async loadTotalCount(resultIndex: number) {
@@ -68,10 +60,6 @@ export abstract class ResultSetDataSource<TOptions> extends DatabaseDataSource<T
       async () => {
         const info = await this.asyncTaskInfoService.run(asyncTask);
 
-        if (this.cancelLoadTotalCountTask?.cancelled) {
-          return this.count;
-        }
-
         const { count } = await this.graphQLService.sdk.getSqlRowDataCountResult({ taskId: info.id });
 
         return count;
@@ -80,16 +68,16 @@ export abstract class ResultSetDataSource<TOptions> extends DatabaseDataSource<T
       () => this.asyncTaskInfoService.remove(asyncTask.id),
     );
 
-    this.cancelLoadTotalCountTask = task;
+    this.totalCountRequestTask = task;
 
     try {
       const count = await task;
 
-      if (count !== this.count) {
-        this.setTotalCount(resultIndex, count);
-      }
+      this.setTotalCount(resultIndex, count);
     } finally {
-      this.cancelLoadTotalCountTask = null;
+      this.totalCountRequestTask = null;
     }
+
+    return this.totalCountRequestTask;
   }
 }
diff --git a/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/CancelTotalCountAction.tsx b/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/CancelTotalCountAction.tsx
index ab4280cd74..c4b447f839 100644
--- a/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/CancelTotalCountAction.tsx
+++ b/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/CancelTotalCountAction.tsx
@@ -13,7 +13,11 @@ import { ToolsAction, useTranslate } from '@cloudbeaver/core-blocks';
 import { tableFooterMenuStyles } from './TableFooterMenu/TableFooterMenuItem';
 import classes from './TableFooterRowCount.m.css';
 
-export const CancelTotalCountAction = observer(function CancelTotalCountAction({ onClick }: { onClick: VoidFunction }) {
+interface Props {
+  onClick: VoidFunction;
+}
+
+export const CancelTotalCountAction = observer(function CancelTotalCountAction({ onClick }: Props) {
   const translate = useTranslate();
 
   return styled(tableFooterMenuStyles)(
diff --git a/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TableFooterRowCount.tsx b/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TableFooterRowCount.tsx
index d7b6116920..fee254d8c1 100644
--- a/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TableFooterRowCount.tsx
+++ b/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TableFooterRowCount.tsx
@@ -36,10 +36,10 @@ export const TableFooterRowCount: React.FC<Props> = observer(function TableFoote
     }
   }
 
-  function cancelTotalCount() {
+  async function cancelTotalCount() {
     try {
       setLoading(false);
-      model.source.cancelLoadTotalCount();
+      await model.source.cancelLoadTotalCount();
     } catch (e: any) {
       notificationService.logException(e);
     }
diff --git a/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TotalCountAction.tsx b/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TotalCountAction.tsx
index 754d0c7e27..7e1587a032 100644
--- a/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TotalCountAction.tsx
+++ b/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TotalCountAction.tsx
@@ -15,17 +15,14 @@ import type { IDatabaseResultSet } from '../../DatabaseDataModel/IDatabaseResult
 import { tableFooterMenuStyles } from './TableFooterMenu/TableFooterMenuItem';
 import classes from './TableFooterRowCount.m.css';
 
-export const TotalCountAction = observer(function TotalCountAction({
-  onClick,
-  loading,
-  resultIndex,
-  model,
-}: {
+interface Props {
   onClick: VoidFunction;
   loading: boolean;
   resultIndex: number;
   model: IDatabaseDataModel<any, IDatabaseResultSet>;
-}) {
+}
+
+export const TotalCountAction = observer(function TotalCountAction({ onClick, loading, resultIndex, model }: Props) {
   const result = model.getResult(resultIndex);
   const translate = useTranslate();
   const disabled = getComputed(() => model.isLoading() || model.isDisabled(resultIndex));
diff --git a/webapp/packages/plugin-sql-editor/src/SqlResultTabs/SqlQueryResultService.ts b/webapp/packages/plugin-sql-editor/src/SqlResultTabs/SqlQueryResultService.ts
index 74c59f9310..5d0d8d31ac 100644
--- a/webapp/packages/plugin-sql-editor/src/SqlResultTabs/SqlQueryResultService.ts
+++ b/webapp/packages/plugin-sql-editor/src/SqlResultTabs/SqlQueryResultService.ts
@@ -163,7 +163,6 @@ export class SqlQueryResultService {
 
         // TODO: we need to dispose table model, but don't close execution context, so now we only
         const model = this.tableViewerStorageService.get(group.modelId);
-        model?.source.cancelLoadTotalCount();
         // model?.dispose();
 
         if (model?.isLoading()) {

From 708d658ca5c33aabce79375b95774ba85f22f202 Mon Sep 17 00:00:00 2001
From: sergeyteleshev <iamsergeyteleshev@gmail.com>
Date: Mon, 15 Apr 2024 14:13:27 +0200
Subject: [PATCH 25/36] CB-4270 deletes cross icon icon

---
 .../public/icons/data_cancel.svg              | 10 ---------
 .../TableFooter/CancelTotalCountAction.m.css  | 22 +++++++++++++++++++
 .../TableFooter/CancelTotalCountAction.tsx    | 13 +++++------
 3 files changed, 28 insertions(+), 17 deletions(-)
 delete mode 100644 webapp/packages/plugin-data-viewer/public/icons/data_cancel.svg
 create mode 100644 webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/CancelTotalCountAction.m.css

diff --git a/webapp/packages/plugin-data-viewer/public/icons/data_cancel.svg b/webapp/packages/plugin-data-viewer/public/icons/data_cancel.svg
deleted file mode 100644
index 9a20bf546c..0000000000
--- a/webapp/packages/plugin-data-viewer/public/icons/data_cancel.svg
+++ /dev/null
@@ -1,10 +0,0 @@
-<svg width="16" height="16" viewBox="0 0 16 16" fill="none" xmlns="http://www.w3.org/2000/svg">
-<g clip-path="url(#clip0)">
-<path fill-rule="evenodd" clip-rule="evenodd" d="M4.3988 3.13864C4.05709 2.79693 3.50307 2.79693 3.16136 3.13864C2.81965 3.48035 2.81965 4.03437 3.16136 4.37608L6.78528 8L3.16136 11.6239C2.81965 11.9656 2.81965 12.5197 3.16136 12.8614C3.50307 13.2031 4.05709 13.2031 4.3988 12.8614L8.02272 9.23744L11.6466 12.8614C11.9883 13.2031 12.5424 13.2031 12.8841 12.8614C13.2258 12.5197 13.2258 11.9656 12.8841 11.6239L9.26015 8L12.8841 4.37608C13.2258 4.03437 13.2258 3.48035 12.8841 3.13864C12.5424 2.79693 11.9883 2.79693 11.6466 3.13864L8.02272 6.76256L4.3988 3.13864Z" fill="#E49388"/>
-</g>
-<defs>
-<clipPath id="clip0">
-<rect width="16" height="16" fill="white"/>
-</clipPath>
-</defs>
-</svg>
diff --git a/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/CancelTotalCountAction.m.css b/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/CancelTotalCountAction.m.css
new file mode 100644
index 0000000000..207f55175e
--- /dev/null
+++ b/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/CancelTotalCountAction.m.css
@@ -0,0 +1,22 @@
+/*
+ * CloudBeaver - Cloud Database Manager
+ * Copyright (C) 2020-2024 DBeaver Corp and others
+ *
+ * Licensed under the Apache License, Version 2.0.
+ * you may not use this file except in compliance with the License.
+ */
+.action {
+  composes: theme-ripple from global;
+  cursor: pointer;
+  user-select: none;
+}
+
+.icon {
+  composes: theme-text-error from global;
+}
+
+.cancelText {
+  text-transform: uppercase;
+  font-weight: 700;
+  font-size: 12px;
+}
diff --git a/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/CancelTotalCountAction.tsx b/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/CancelTotalCountAction.tsx
index c4b447f839..d5a9d49189 100644
--- a/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/CancelTotalCountAction.tsx
+++ b/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/CancelTotalCountAction.tsx
@@ -8,10 +8,10 @@
 import { observer } from 'mobx-react-lite';
 import styled from 'reshadow';
 
-import { ToolsAction, useTranslate } from '@cloudbeaver/core-blocks';
+import { Container, IconButton, useTranslate } from '@cloudbeaver/core-blocks';
 
+import styles from './CancelTotalCountAction.m.css';
 import { tableFooterMenuStyles } from './TableFooterMenu/TableFooterMenuItem';
-import classes from './TableFooterRowCount.m.css';
 
 interface Props {
   onClick: VoidFunction;
@@ -21,10 +21,9 @@ export const CancelTotalCountAction = observer(function CancelTotalCountAction({
   const translate = useTranslate();
 
   return styled(tableFooterMenuStyles)(
-    <div className={classes.wrapper} title={translate('ui_processing_cancel')}>
-      <ToolsAction icon="/icons/data_cancel.svg" viewBox="0 0 32 32" onClick={onClick}>
-        {translate('ui_processing_cancel')}
-      </ToolsAction>
-    </div>,
+    <Container className={styles.action} title={translate('ui_processing_cancel')} noWrap center zeroBasis keepSize onClick={onClick}>
+      <IconButton className={styles.icon} name="cross" title={translate('ui_processing_cancel')} viewBox="0 0 32 32" />
+      <span className={styles.cancelText}>{translate('ui_processing_cancel')}</span>
+    </Container>,
   );
 });

From 4a3445eb3518e5f3e4c371056754ea41f7689b05 Mon Sep 17 00:00:00 2001
From: sergeyteleshev <iamsergeyteleshev@gmail.com>
Date: Mon, 15 Apr 2024 16:11:47 +0200
Subject: [PATCH 26/36] CB-4270 DataSources has correct tree of calls for
 cancel and dispose functions

---
 .../plugin-data-viewer/src/ContainerDataSource.ts   |  7 ++-----
 .../src/DatabaseDataModel/DatabaseDataModel.ts      |  3 +--
 .../src/DatabaseDataModel/DatabaseDataSource.ts     |  5 ++++-
 .../plugin-data-viewer/src/ResultSetDataSource.ts   |  6 +++++-
 .../plugin-sql-editor/src/QueryDataSource.ts        | 13 +++++++------
 .../src/SqlResultTabs/SqlQueryResultService.ts      |  4 +---
 6 files changed, 20 insertions(+), 18 deletions(-)

diff --git a/webapp/packages/plugin-data-viewer/src/ContainerDataSource.ts b/webapp/packages/plugin-data-viewer/src/ContainerDataSource.ts
index 431471c8cf..b70b9eb6b3 100644
--- a/webapp/packages/plugin-data-viewer/src/ContainerDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/ContainerDataSource.ts
@@ -68,9 +68,7 @@ export class ContainerDataSource extends ResultSetDataSource<IDataContainerOptio
   }
 
   async cancel(): Promise<void> {
-    if (this.currentTask) {
-      await this.currentTask.cancel();
-    }
+    await Promise.all([this.currentTask?.cancel(), super.cancel()]);
   }
 
   async request(prevResults: IDatabaseResultSet[]): Promise<IDatabaseResultSet[]> {
@@ -221,8 +219,7 @@ export class ContainerDataSource extends ResultSetDataSource<IDataContainerOptio
   }
 
   async dispose(): Promise<void> {
-    await this.closeResults(this.results);
-    await this.executionContext?.destroy();
+    await Promise.all([this.closeResults(this.results), this.cancel(), super.dispose()]);
   }
 
   private async closeResults(results: IDatabaseResultSet[]) {
diff --git a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataModel.ts b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataModel.ts
index 00d2f36a14..32192e3554 100644
--- a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataModel.ts
+++ b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataModel.ts
@@ -155,8 +155,7 @@ export class DatabaseDataModel<TOptions, TResult extends IDatabaseDataResult = I
   }
 
   async cancel(): Promise<void> {
-    await this.source.cancelLoadTotalCount();
-    return await this.source.cancel();
+    await this.source.cancel();
   }
 
   resetData(): void {
diff --git a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
index 27dad8c500..22c12065d3 100644
--- a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
@@ -372,10 +372,13 @@ export abstract class DatabaseDataSource<TOptions, TResult extends IDatabaseData
     return this;
   }
 
+  async dispose(): Promise<void> {
+    await Promise.all([this.cancel(), this.executionContext?.destroy()]);
+  }
+
   abstract request(prevResults: TResult[]): TResult[] | Promise<TResult[]>;
   abstract save(prevResults: TResult[]): Promise<TResult[]> | TResult[];
 
-  abstract dispose(): Promise<void>;
   abstract loadTotalCount(resultIndex: number): Promise<ITask<number> | null>;
   abstract cancelLoadTotalCount(): Promise<ITask<number> | null>;
 
diff --git a/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts b/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
index 3f8cdb5440..cbab6e4a04 100644
--- a/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
@@ -21,8 +21,12 @@ export abstract class ResultSetDataSource<TOptions> extends DatabaseDataSource<T
     super(serviceInjector);
   }
 
+  async cancel(): Promise<void> {
+    await Promise.all([this.cancelLoadTotalCount(), super.cancel()]);
+  }
+
   async dispose(): Promise<void> {
-    await this.cancelLoadTotalCount();
+    await Promise.all([this.cancel(), super.dispose()]);
   }
 
   async cancelLoadTotalCount(): Promise<ITask<number> | null> {
diff --git a/webapp/packages/plugin-sql-editor/src/QueryDataSource.ts b/webapp/packages/plugin-sql-editor/src/QueryDataSource.ts
index 0b5ddcc1cd..a76072e7a3 100644
--- a/webapp/packages/plugin-sql-editor/src/QueryDataSource.ts
+++ b/webapp/packages/plugin-sql-editor/src/QueryDataSource.ts
@@ -49,7 +49,11 @@ export class QueryDataSource<TOptions extends IDataQueryOptions = IDataQueryOpti
     return this.currentTask?.cancelled || false;
   }
 
-  constructor(readonly serviceInjector: IServiceInjector, graphQLService: GraphQLService, asyncTaskInfoService: AsyncTaskInfoService) {
+  constructor(
+    readonly serviceInjector: IServiceInjector,
+    graphQLService: GraphQLService,
+    asyncTaskInfoService: AsyncTaskInfoService,
+  ) {
     super(serviceInjector, graphQLService, asyncTaskInfoService);
 
     this.currentTask = null;
@@ -80,9 +84,7 @@ export class QueryDataSource<TOptions extends IDataQueryOptions = IDataQueryOpti
   }
 
   async cancel(): Promise<void> {
-    if (this.currentTask) {
-      await this.currentTask.cancel();
-    }
+    await Promise.all([this.currentTask?.cancel(), super.cancel()]);
   }
 
   async save(prevResults: IDatabaseResultSet[]): Promise<IDatabaseResultSet[]> {
@@ -293,7 +295,6 @@ export class QueryDataSource<TOptions extends IDataQueryOptions = IDataQueryOpti
   }
 
   async dispose(): Promise<void> {
-    await this.closeResults(this.results);
-    await this.cancel();
+    await Promise.all([this.closeResults(this.results), this.cancel(), super.dispose()]);
   }
 }
diff --git a/webapp/packages/plugin-sql-editor/src/SqlResultTabs/SqlQueryResultService.ts b/webapp/packages/plugin-sql-editor/src/SqlResultTabs/SqlQueryResultService.ts
index 5d0d8d31ac..9094462f96 100644
--- a/webapp/packages/plugin-sql-editor/src/SqlResultTabs/SqlQueryResultService.ts
+++ b/webapp/packages/plugin-sql-editor/src/SqlResultTabs/SqlQueryResultService.ts
@@ -165,9 +165,7 @@ export class SqlQueryResultService {
         const model = this.tableViewerStorageService.get(group.modelId);
         // model?.dispose();
 
-        if (model?.isLoading()) {
-          model.cancel();
-        }
+        model?.cancel();
 
         this.tableViewerStorageService.remove(group.modelId);
       }

From 815cd26c6a26d63a055fafed3c4887e8d60dd4e6 Mon Sep 17 00:00:00 2001
From: sergeyteleshev <iamsergeyteleshev@gmail.com>
Date: Mon, 15 Apr 2024 16:26:03 +0200
Subject: [PATCH 27/36] CB-4270 cleanup

---
 .../TableViewer/TableFooter/CancelTotalCountAction.m.css    | 1 +
 .../src/TableViewer/TableFooter/CancelTotalCountAction.tsx  | 6 ++----
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/CancelTotalCountAction.m.css b/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/CancelTotalCountAction.m.css
index 207f55175e..fe86b1288c 100644
--- a/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/CancelTotalCountAction.m.css
+++ b/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/CancelTotalCountAction.m.css
@@ -7,6 +7,7 @@
  */
 .action {
   composes: theme-ripple from global;
+  padding: 0 8px;
   cursor: pointer;
   user-select: none;
 }
diff --git a/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/CancelTotalCountAction.tsx b/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/CancelTotalCountAction.tsx
index d5a9d49189..eea1e570f5 100644
--- a/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/CancelTotalCountAction.tsx
+++ b/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/CancelTotalCountAction.tsx
@@ -6,12 +6,10 @@
  * you may not use this file except in compliance with the License.
  */
 import { observer } from 'mobx-react-lite';
-import styled from 'reshadow';
 
 import { Container, IconButton, useTranslate } from '@cloudbeaver/core-blocks';
 
 import styles from './CancelTotalCountAction.m.css';
-import { tableFooterMenuStyles } from './TableFooterMenu/TableFooterMenuItem';
 
 interface Props {
   onClick: VoidFunction;
@@ -20,10 +18,10 @@ interface Props {
 export const CancelTotalCountAction = observer(function CancelTotalCountAction({ onClick }: Props) {
   const translate = useTranslate();
 
-  return styled(tableFooterMenuStyles)(
+  return (
     <Container className={styles.action} title={translate('ui_processing_cancel')} noWrap center zeroBasis keepSize onClick={onClick}>
       <IconButton className={styles.icon} name="cross" title={translate('ui_processing_cancel')} viewBox="0 0 32 32" />
       <span className={styles.cancelText}>{translate('ui_processing_cancel')}</span>
-    </Container>,
+    </Container>
   );
 });

From 8b185d5aba2e4ac4cfe2f31e222822f4e929310d Mon Sep 17 00:00:00 2001
From: Aleksei Potsetsuev <wrouds@gmail.com>
Date: Mon, 15 Apr 2024 22:27:34 +0800
Subject: [PATCH 28/36] chore: add experimental cancellable promise
 implementation

---
 .../src/TaskScheduler/CancelError.ts          |  15 +
 .../src/TaskScheduler/ITaskNew.ts             |  43 ++
 .../src/TaskScheduler/TaskNew.ts              | 386 ++++++++++++++++++
 3 files changed, 444 insertions(+)
 create mode 100644 webapp/packages/core-executor/src/TaskScheduler/CancelError.ts
 create mode 100644 webapp/packages/core-executor/src/TaskScheduler/ITaskNew.ts
 create mode 100644 webapp/packages/core-executor/src/TaskScheduler/TaskNew.ts

diff --git a/webapp/packages/core-executor/src/TaskScheduler/CancelError.ts b/webapp/packages/core-executor/src/TaskScheduler/CancelError.ts
new file mode 100644
index 0000000000..797be2a5a9
--- /dev/null
+++ b/webapp/packages/core-executor/src/TaskScheduler/CancelError.ts
@@ -0,0 +1,15 @@
+/*
+ * CloudBeaver - Cloud Database Manager
+ * Copyright (C) 2020-2024 DBeaver Corp and others
+ *
+ * Licensed under the Apache License, Version 2.0.
+ * you may not use this file except in compliance with the License.
+ */
+
+export class CancelError extends Error {
+  constructor(message?: string, options?: ErrorOptions) {
+    super(message, options);
+
+    this.name = 'Cancel Error';
+  }
+}
diff --git a/webapp/packages/core-executor/src/TaskScheduler/ITaskNew.ts b/webapp/packages/core-executor/src/TaskScheduler/ITaskNew.ts
new file mode 100644
index 0000000000..dc0327e77d
--- /dev/null
+++ b/webapp/packages/core-executor/src/TaskScheduler/ITaskNew.ts
@@ -0,0 +1,43 @@
+/*
+ * CloudBeaver - Cloud Database Manager
+ * Copyright (C) 2020-2024 DBeaver Corp and others
+ *
+ * Licensed under the Apache License, Version 2.0.
+ * you may not use this file except in compliance with the License.
+ */
+
+export interface ITaskConstructor {
+  readonly prototype: ITask<any>;
+  new <T>(executor: TaskExecutor<T>): ITask<T>;
+
+  reject<T = never>(reason?: any): ITask<T>;
+  cancel<T = never>(reason?: any): ITask<T>;
+  resolve(): ITask<void>;
+  resolve<T>(value: T): ITask<Awaited<T>>;
+  resolve<T>(value: T | PromiseLike<T>): ITask<Awaited<T>>;
+  fromPromise<T>(promise: Promise<T>): ITask<T>;
+  all<T extends readonly unknown[] | []>(values: T): ITask<{ -readonly [P in keyof T]: Awaited<T[P]> }>;
+  allSettled<T extends readonly unknown[] | []>(
+    values: T,
+  ): ITask<{ -readonly [P in keyof T]: { status: 'fulfilled'; value: Awaited<T[P]> } | { status: 'rejected' | 'cancelled'; reason: any } }>;
+}
+
+declare const ITask: ITaskConstructor;
+
+export interface ITask<TValue> extends Promise<TValue>, PromiseLike<TValue> {
+  readonly then: <TResult1 = TValue, TResult2 = never, TResult3 = never>(
+    onFulfilled?: ((value: TValue) => TResult1 | PromiseLike<TResult1>) | null,
+    onRejected?: ((reason: any, cancelled: boolean) => TResult2 | PromiseLike<TResult2>) | null,
+  ) => ITask<TResult1 | TResult2 | TResult3>;
+  readonly catch: <TResult = never>(
+    onRejected?: ((reason: any, cancelled: boolean) => TResult | PromiseLike<TResult>) | null,
+  ) => ITask<TValue | TResult>;
+  readonly finally: (onfinally?: (() => void) | null) => ITask<TValue>;
+  readonly cancel: (reason?: any) => ITask<void>;
+}
+
+export type TaskResolveFn<T> = (reason: T | PromiseLike<T>) => void;
+export type TaskRejectFn = (reason?: any, cancelled?: boolean) => void;
+export type TaskCancelFn = (reason?: any) => ITask<void> | void;
+
+export type TaskExecutor<T> = (resolve: TaskResolveFn<T>, reject: TaskRejectFn, registerCancel: (cancel: TaskCancelFn) => void) => void;
diff --git a/webapp/packages/core-executor/src/TaskScheduler/TaskNew.ts b/webapp/packages/core-executor/src/TaskScheduler/TaskNew.ts
new file mode 100644
index 0000000000..4a53de791e
--- /dev/null
+++ b/webapp/packages/core-executor/src/TaskScheduler/TaskNew.ts
@@ -0,0 +1,386 @@
+/*
+ * CloudBeaver - Cloud Database Manager
+ * Copyright (C) 2020-2024 DBeaver Corp and others
+ *
+ * Licensed under the Apache License, Version 2.0.
+ * you may not use this file except in compliance with the License.
+ */
+import { CancelError } from './CancelError';
+import type { ITask, TaskCancelFn, TaskExecutor, TaskRejectFn, TaskResolveFn } from './ITaskNew';
+
+export const NOOP: VoidFunction = Object.freeze(() => {});
+
+// experimental cancellable promise implementation
+export class Task<T> implements ITask<T> {
+  get [Symbol.toStringTag](): string {
+    return 'Task';
+  }
+
+  protected status: 'pending' | 'fulfilled' | 'rejected';
+  private value: any;
+  private fulfillCallbacks: Array<(value: T | PromiseLike<T>) => void>;
+  private rejectCallbacks: Array<(reason?: any, cancelled?: boolean) => void>;
+  private cancelling: boolean;
+  private cancelCallback: TaskCancelFn;
+  private parentTask: Task<any> | null;
+
+  constructor(executor: TaskExecutor<T>) {
+    this.status = 'pending';
+    this.internalCancel = this.internalCancel.bind(this);
+    this.internalResolve = this.internalResolve.bind(this);
+    this.internalReject = this.internalReject.bind(this);
+    this.registerCancel = this.registerCancel.bind(this);
+    this.cancel = this.cancel.bind(this);
+
+    this.fulfillCallbacks = [];
+    this.rejectCallbacks = [];
+    this.cancelCallback = this.internalCancel;
+    this.cancelling = false;
+    this.parentTask = null;
+
+    try {
+      executor(this.internalResolve, this.internalReject, this.registerCancel);
+    } catch (err) {
+      this.internalReject(err);
+    }
+  }
+
+  then<TResult1 = T, TResult2 = never>(
+    onFulfilled?: ((value: T) => TResult1 | PromiseLike<TResult1>) | null,
+    onRejected?: ((reason: any, cancelled: boolean) => TResult2 | PromiseLike<TResult2>) | null,
+  ): Task<TResult1 | TResult2> {
+    const task = new Task<TResult1 | TResult2>(NOOP);
+
+    task.parentTask = this;
+
+    this.subscribeResolve(value => {
+      queueMicrotask(() => {
+        if (typeof onFulfilled !== 'function') {
+          task.internalResolve(value as TResult1);
+          return;
+        }
+
+        console.warn('generate microtask');
+        this.executeValue(
+          value,
+          value => {
+            this.resolvePromise(task, onFulfilled(value as any));
+            // this.promiseResolution(task, onFulfilled(value as any), resolve, reject, registerCancel);
+          },
+          task.internalReject,
+        );
+      });
+    });
+
+    this.subscribeReject((reason, cancelled) => {
+      queueMicrotask(() => {
+        if (typeof onRejected !== 'function') {
+          task.internalReject(reason, cancelled);
+          return;
+        }
+
+        console.warn('generate microtask');
+        this.executeValue(
+          reason,
+          (reason: any) => {
+            this.resolvePromise(task, onRejected(reason, cancelled ?? false) as TResult2);
+            // this.promiseResolution(task, onRejected(reason, cancelled ?? false) as TResult2, resolve, reject, registerCancel);
+          },
+          task.internalReject,
+        );
+      });
+    });
+
+    return task;
+  }
+
+  catch<TResult = never>(onRejected?: ((reason: any, cancelled: boolean) => TResult | PromiseLike<TResult>) | null): Task<T | TResult> {
+    return this.then(undefined, onRejected);
+  }
+
+  finally(onFinally?: (() => void) | null): Task<T> {
+    return this.then(
+      value => {
+        onFinally?.();
+        return value;
+      },
+      reason => {
+        onFinally?.();
+        throw reason;
+      },
+    );
+  }
+
+  cancel(reason?: any): Task<void> {
+    if (this.cancelling) {
+      throw new Error('Task is already cancelling');
+    }
+    this.cancelling = true;
+
+    if (this.parentTask && this.parentTask.status === 'pending') {
+      return this.parentTask.cancel(reason).finally(() => {
+        this.cancelling = false;
+      });
+    }
+
+    return Task.resolve(this.cancelCallback(reason)).finally(() => {
+      this.cancelling = false;
+    });
+  }
+
+  protected internalCancel(reason?: any): void {
+    this.internalReject(reason, true);
+  }
+
+  protected internalResolve(value: T | PromiseLike<T>) {
+    if (this.status === 'pending') {
+      this.status = 'fulfilled';
+      this.value = value;
+
+      this.executeTasks(this.fulfillCallbacks, value, undefined);
+    }
+  }
+
+  protected internalReject(reason?: any, cancelled?: boolean) {
+    if (this.status === 'pending') {
+      cancelled ||= reason instanceof CancelError;
+      if (cancelled && !(reason instanceof CancelError)) {
+        reason = new CancelError(reason, { cause: reason });
+      }
+      this.status = 'rejected';
+      this.value = reason;
+
+      this.executeTasks(this.rejectCallbacks, reason, cancelled);
+    }
+  }
+
+  executeTasks<T>(tasks: ((value: T, cancelled?: boolean) => void)[], value: T, cancelled: any) {
+    this.fulfillCallbacks = [];
+    this.rejectCallbacks = [];
+    this.cancelCallback = this.internalCancel;
+
+    for (const onRejectedCallback of tasks) {
+      onRejectedCallback(value, cancelled);
+    }
+  }
+
+  private registerCancel(cancel: TaskCancelFn): void {
+    let copy = this.cancelCallback;
+    if (copy === this.internalCancel) {
+      copy = NOOP;
+    }
+
+    this.cancelCallback = reason => Task.resolve(copy(reason)).then(() => cancel(reason));
+  }
+
+  private subscribeResolve(resolve: TaskResolveFn<T>): void {
+    switch (this.status) {
+      case 'pending':
+        {
+          this.fulfillCallbacks.push(resolve);
+        }
+        break;
+      case 'fulfilled':
+        resolve(this.value);
+        break;
+    }
+  }
+
+  private subscribeReject(reject: TaskRejectFn): void {
+    switch (this.status) {
+      case 'pending':
+        {
+          this.rejectCallbacks.push(reject);
+        }
+        break;
+      case 'rejected':
+        reject(this.value, this.value instanceof CancelError);
+        break;
+    }
+  }
+
+  static reject<T = never>(reason?: any): Task<T> {
+    return new Task<T>((_, reject) => {
+      reject(reason);
+    });
+  }
+
+  static resolve(): Task<void>;
+  static resolve<T>(value: T): Task<Awaited<T>>;
+  static resolve<T>(value: T | PromiseLike<T>): Task<Awaited<T>>;
+  static resolve<T = void>(value?: T | PromiseLike<T>): Task<Awaited<T>> {
+    if (isPromiseLike(value)) {
+      return value as Task<Awaited<T>>;
+    } else {
+      return new Task(resolve => {
+        resolve(value as any);
+      });
+    }
+  }
+
+  static cancel<T = never>(reason?: any): Task<T> {
+    return new Task<T>((_, reject) => {
+      reject(reason, true);
+    });
+  }
+
+  static fromPromise<T>(promise: PromiseLike<T>): Task<T> {
+    return new Task<T>((resolve, reject) => {
+      promise.then(resolve, reject);
+    });
+  }
+
+  static all<T extends readonly unknown[] | []>(values: T): Task<{ -readonly [P in keyof T]: Awaited<T[P]> }> {
+    return new Task((resolve, reject) => {
+      let counter = 0;
+      const result = [] as { -readonly [P in keyof T]: Awaited<T[P]> };
+
+      if (values.length === 0) {
+        resolve([] as any);
+        return;
+      }
+
+      for (let i = 0; i < values.length; i++) {
+        Task.resolve(values[i]).then(
+          res => {
+            result[i] = res;
+            counter += 1;
+            if (counter === values.length) {
+              resolve(result);
+            }
+          },
+          reason => {
+            reject(reason);
+          },
+        );
+      }
+    });
+  }
+
+  static allSettled<T extends readonly unknown[] | []>(
+    values: T,
+  ): Task<{ -readonly [P in keyof T]: { status: 'fulfilled'; value: Awaited<T[P]> } | { status: 'rejected' | 'cancelled'; reason: any } }> {
+    return new Task(resolve => {
+      let counter = 0;
+      const result = [] as {
+        -readonly [P in keyof T]: { status: 'fulfilled'; value: Awaited<T[P]> } | { status: 'rejected' | 'cancelled'; reason: any };
+      };
+
+      if (values.length === 0) {
+        resolve([] as any);
+        return;
+      }
+
+      for (let i = 0; i < values.length; i++) {
+        Task.resolve(values[i]).then(
+          res => {
+            result[i] = { status: 'fulfilled', value: res };
+            counter += 1;
+            if (counter === values.length) {
+              resolve(result);
+            }
+          },
+          (reason, cancelled) => {
+            result[i] = { status: cancelled ? 'cancelled' : 'rejected', reason };
+            counter += 1;
+            if (counter === values.length) {
+              resolve(result);
+            }
+          },
+        );
+      }
+    });
+  }
+
+  private resolvePromise<T>(promise: Task<any>, value: T | PromiseLike<T>): void {
+    this.executeValue(
+      value,
+      value => {
+        if (parent === value) {
+          throw new TypeError('Cannot resolve task with itself');
+        }
+
+        if (value instanceof Task || value instanceof Promise) {
+          if (value instanceof Task) {
+            promise.registerCancel(reason => value.cancel(reason));
+          }
+          value.then(promise.internalResolve, promise.internalReject);
+          return;
+        }
+
+        if (value && (typeof value === 'object' || typeof value === 'function')) {
+          const then = (value as PromiseLike<any>).then;
+
+          if (typeof then === 'function') {
+            (then as (...args: any[]) => any).call(
+              value,
+              (r: any) => {
+                promise.resolvePromise(value as any, r);
+              },
+              promise.internalReject,
+            );
+            return;
+          }
+        }
+
+        promise.internalResolve(value);
+      },
+      promise.internalReject,
+    );
+  }
+
+  private promiseResolution<T>(
+    parent: Task<any>,
+    value: T | PromiseLike<T>,
+    resolve: (value: T) => void,
+    reject: TaskRejectFn,
+    registerCancel: (cancel: TaskCancelFn) => void,
+  ): void {
+    this.executeValue(
+      value,
+      value => {
+        if (parent === value) {
+          throw new TypeError('Cannot resolve task with itself');
+        }
+
+        if (value instanceof Task || value instanceof Promise) {
+          if (value instanceof Task) {
+            registerCancel(reason => value.cancel(reason));
+          }
+          value.then(resolve, reject);
+          return;
+        }
+
+        if (value && (typeof value === 'object' || typeof value === 'function')) {
+          const then = (value as PromiseLike<any>).then;
+
+          if (typeof then === 'function') {
+            (then as (...args: any[]) => any).call(
+              value,
+              (r: any) => {
+                this.promiseResolution(value as any, r, resolve, reject, registerCancel);
+              },
+              reject,
+            );
+            return;
+          }
+        }
+
+        resolve(value as any);
+      },
+      reject,
+    );
+  }
+
+  private executeValue<T>(value: T, resolve: (value: T) => void, reject: TaskRejectFn) {
+    try {
+      resolve(value);
+    } catch (err) {
+      reject(err, err instanceof CancelError);
+    }
+  }
+}
+
+function isPromiseLike<T>(value: any): value is PromiseLike<T> {
+  return value && typeof value === 'object' && 'then' in value;
+}

From 46b89da2f32751631fd831bdf653e1833683cf62 Mon Sep 17 00:00:00 2001
From: sergeyteleshev <iamsergeyteleshev@gmail.com>
Date: Mon, 15 Apr 2024 20:20:30 +0200
Subject: [PATCH 29/36] CB-4270 pr fixes: remove unneeded disposes + close
 results in close result tab flow

---
 .../packages/plugin-data-viewer/src/ContainerDataSource.ts  | 6 +-----
 .../src/DatabaseDataModel/DatabaseDataSource.ts             | 3 ++-
 .../src/DatabaseDataModel/IDatabaseDataSource.ts            | 2 ++
 .../packages/plugin-data-viewer/src/ResultSetDataSource.ts  | 4 ----
 webapp/packages/plugin-sql-editor/src/QueryDataSource.ts    | 6 +-----
 .../src/SqlResultTabs/SqlQueryResultService.ts              | 1 +
 6 files changed, 7 insertions(+), 15 deletions(-)

diff --git a/webapp/packages/plugin-data-viewer/src/ContainerDataSource.ts b/webapp/packages/plugin-data-viewer/src/ContainerDataSource.ts
index b70b9eb6b3..75a25c8cd8 100644
--- a/webapp/packages/plugin-data-viewer/src/ContainerDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/ContainerDataSource.ts
@@ -218,11 +218,7 @@ export class ContainerDataSource extends ResultSetDataSource<IDataContainerOptio
     return prevResults;
   }
 
-  async dispose(): Promise<void> {
-    await Promise.all([this.closeResults(this.results), this.cancel(), super.dispose()]);
-  }
-
-  private async closeResults(results: IDatabaseResultSet[]) {
+  async closeResults(results: IDatabaseResultSet[]) {
     await this.connectionExecutionContextService.load();
 
     if (!this.executionContext?.context) {
diff --git a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
index 22c12065d3..75953dbd32 100644
--- a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
@@ -373,7 +373,8 @@ export abstract class DatabaseDataSource<TOptions, TResult extends IDatabaseData
   }
 
   async dispose(): Promise<void> {
-    await Promise.all([this.cancel(), this.executionContext?.destroy()]);
+    await this.cancel();
+    await this.executionContext?.destroy();
   }
 
   abstract request(prevResults: TResult[]): TResult[] | Promise<TResult[]>;
diff --git a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/IDatabaseDataSource.ts b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/IDatabaseDataSource.ts
index 69e37eef3f..dafafb1966 100644
--- a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/IDatabaseDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/IDatabaseDataSource.ts
@@ -13,6 +13,7 @@ import type { ResultDataFormat } from '@cloudbeaver/core-sdk';
 import type { IDatabaseDataAction, IDatabaseDataActionClass, IDatabaseDataActionInterface } from './IDatabaseDataAction';
 import type { IDatabaseDataActions } from './IDatabaseDataActions';
 import type { IDatabaseDataResult } from './IDatabaseDataResult';
+import type { IDatabaseResultSet } from './IDatabaseResultSet';
 
 export interface IRequestInfo {
   readonly originalQuery: string;
@@ -98,4 +99,5 @@ export interface IDatabaseDataSource<TOptions, TResult extends IDatabaseDataResu
   clearError: () => this;
   resetData: () => this;
   dispose: () => Promise<void>;
+  closeResults?: (results: IDatabaseResultSet[]) => Promise<void>;
 }
diff --git a/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts b/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
index cbab6e4a04..46f7e653ca 100644
--- a/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
@@ -25,10 +25,6 @@ export abstract class ResultSetDataSource<TOptions> extends DatabaseDataSource<T
     await Promise.all([this.cancelLoadTotalCount(), super.cancel()]);
   }
 
-  async dispose(): Promise<void> {
-    await Promise.all([this.cancel(), super.dispose()]);
-  }
-
   async cancelLoadTotalCount(): Promise<ITask<number> | null> {
     await this.totalCountRequestTask?.cancel();
 
diff --git a/webapp/packages/plugin-sql-editor/src/QueryDataSource.ts b/webapp/packages/plugin-sql-editor/src/QueryDataSource.ts
index a76072e7a3..cb50959bfb 100644
--- a/webapp/packages/plugin-sql-editor/src/QueryDataSource.ts
+++ b/webapp/packages/plugin-sql-editor/src/QueryDataSource.ts
@@ -254,7 +254,7 @@ export class QueryDataSource<TOptions extends IDataQueryOptions = IDataQueryOpti
     }
   }
 
-  private async closeResults(results: IDatabaseResultSet[]) {
+  async closeResults(results: IDatabaseResultSet[]) {
     if (!this.executionContext?.context) {
       return;
     }
@@ -293,8 +293,4 @@ export class QueryDataSource<TOptions extends IDataQueryOptions = IDataQueryOpti
       data: result.resultSet,
     }));
   }
-
-  async dispose(): Promise<void> {
-    await Promise.all([this.closeResults(this.results), this.cancel(), super.dispose()]);
-  }
 }
diff --git a/webapp/packages/plugin-sql-editor/src/SqlResultTabs/SqlQueryResultService.ts b/webapp/packages/plugin-sql-editor/src/SqlResultTabs/SqlQueryResultService.ts
index 9094462f96..1c2eee7a68 100644
--- a/webapp/packages/plugin-sql-editor/src/SqlResultTabs/SqlQueryResultService.ts
+++ b/webapp/packages/plugin-sql-editor/src/SqlResultTabs/SqlQueryResultService.ts
@@ -166,6 +166,7 @@ export class SqlQueryResultService {
         // model?.dispose();
 
         model?.cancel();
+        model?.source.closeResults?.(model.getResults());
 
         this.tableViewerStorageService.remove(group.modelId);
       }

From c665a13a8bf8c1a373e25f6e3924c94a32ef9ecb Mon Sep 17 00:00:00 2001
From: sergeyteleshev <iamsergeyteleshev@gmail.com>
Date: Tue, 16 Apr 2024 08:49:00 +0200
Subject: [PATCH 30/36] CB-4270 removes closeResults abstraction leak in
 ResultSetDataSource extended classes

---
 .../src/useGroupingDataModel.ts               |  5 ++--
 .../src/ContainerDataSource.ts                | 28 ++-----------------
 .../DatabaseDataModel/IDatabaseDataSource.ts  |  2 --
 .../src/ResultSetDataSource.ts                | 26 +++++++++++++++++
 .../plugin-sql-editor/src/QueryDataSource.ts  | 27 ++----------------
 .../SqlResultTabs/SqlQueryResultService.ts    |  8 ++++--
 .../src/SqlResultTabs/SqlQueryService.ts      | 14 ++++++++--
 7 files changed, 51 insertions(+), 59 deletions(-)

diff --git a/webapp/packages/plugin-data-viewer-result-set-grouping/src/useGroupingDataModel.ts b/webapp/packages/plugin-data-viewer-result-set-grouping/src/useGroupingDataModel.ts
index d008c625d1..034dd3efef 100644
--- a/webapp/packages/plugin-data-viewer-result-set-grouping/src/useGroupingDataModel.ts
+++ b/webapp/packages/plugin-data-viewer-result-set-grouping/src/useGroupingDataModel.ts
@@ -9,7 +9,7 @@ import { reaction } from 'mobx';
 import { useEffect } from 'react';
 
 import { useObjectRef, useResource } from '@cloudbeaver/core-blocks';
-import { ConnectionInfoResource, createConnectionParam } from '@cloudbeaver/core-connections';
+import { ConnectionExecutionContextService, ConnectionInfoResource, createConnectionParam } from '@cloudbeaver/core-connections';
 import { App, useService } from '@cloudbeaver/core-di';
 import { AsyncTaskInfoService, GraphQLService } from '@cloudbeaver/core-sdk';
 import { isObjectsEqual } from '@cloudbeaver/core-utils';
@@ -39,6 +39,7 @@ export function useGroupingDataModel(
   const graphQLService = useService(GraphQLService);
   const asyncTaskInfoService = useService(AsyncTaskInfoService);
   const dataViewerSettingsService = useService(DataViewerSettingsService);
+  const connectionExecutionContextService = useService(ConnectionExecutionContextService);
 
   const executionContext = sourceModel.source.executionContext;
   const contextInfo = executionContext?.context;
@@ -49,7 +50,7 @@ export function useGroupingDataModel(
 
   const model = useObjectRef(
     () => {
-      const source = new GroupingDataSource(app.getServiceInjector(), graphQLService, asyncTaskInfoService);
+      const source = new GroupingDataSource(app.getServiceInjector(), graphQLService, asyncTaskInfoService, connectionExecutionContextService);
 
       const model = tableViewerStorageService.add(new DatabaseDataModel(source));
 
diff --git a/webapp/packages/plugin-data-viewer/src/ContainerDataSource.ts b/webapp/packages/plugin-data-viewer/src/ContainerDataSource.ts
index 75a25c8cd8..52256d6195 100644
--- a/webapp/packages/plugin-data-viewer/src/ContainerDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/ContainerDataSource.ts
@@ -46,9 +46,9 @@ export class ContainerDataSource extends ResultSetDataSource<IDataContainerOptio
     serviceInjector: IServiceInjector,
     graphQLService: GraphQLService,
     asyncTaskInfoService: AsyncTaskInfoService,
-    private readonly connectionExecutionContextService: ConnectionExecutionContextService,
+    connectionExecutionContextService: ConnectionExecutionContextService,
   ) {
-    super(serviceInjector, graphQLService, asyncTaskInfoService);
+    super(serviceInjector, graphQLService, asyncTaskInfoService, connectionExecutionContextService);
 
     this.currentTask = null;
     this.executionContext = null;
@@ -218,30 +218,6 @@ export class ContainerDataSource extends ResultSetDataSource<IDataContainerOptio
     return prevResults;
   }
 
-  async closeResults(results: IDatabaseResultSet[]) {
-    await this.connectionExecutionContextService.load();
-
-    if (!this.executionContext?.context) {
-      return;
-    }
-
-    for (const result of results) {
-      if (result.id === null) {
-        continue;
-      }
-      try {
-        await this.graphQLService.sdk.closeResult({
-          projectId: result.projectId,
-          connectionId: result.connectionId,
-          contextId: result.contextId,
-          resultId: result.id,
-        });
-      } catch (exception: any) {
-        console.log(`Error closing result (${result.id}):`, exception);
-      }
-    }
-  }
-
   private transformResults(executionContextInfo: IConnectionExecutionContextInfo, results: SqlQueryResults[], limit: number): IDatabaseResultSet[] {
     return results.map<IDatabaseResultSet>((result, index) => ({
       id: result.resultSet?.id || '0',
diff --git a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/IDatabaseDataSource.ts b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/IDatabaseDataSource.ts
index dafafb1966..69e37eef3f 100644
--- a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/IDatabaseDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/IDatabaseDataSource.ts
@@ -13,7 +13,6 @@ import type { ResultDataFormat } from '@cloudbeaver/core-sdk';
 import type { IDatabaseDataAction, IDatabaseDataActionClass, IDatabaseDataActionInterface } from './IDatabaseDataAction';
 import type { IDatabaseDataActions } from './IDatabaseDataActions';
 import type { IDatabaseDataResult } from './IDatabaseDataResult';
-import type { IDatabaseResultSet } from './IDatabaseResultSet';
 
 export interface IRequestInfo {
   readonly originalQuery: string;
@@ -99,5 +98,4 @@ export interface IDatabaseDataSource<TOptions, TResult extends IDatabaseDataResu
   clearError: () => this;
   resetData: () => this;
   dispose: () => Promise<void>;
-  closeResults?: (results: IDatabaseResultSet[]) => Promise<void>;
 }
diff --git a/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts b/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
index 46f7e653ca..749b1051c3 100644
--- a/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
@@ -5,6 +5,7 @@
  * Licensed under the Apache License, Version 2.0.
  * you may not use this file except in compliance with the License.
  */
+import type { ConnectionExecutionContextService } from '@cloudbeaver/core-connections';
 import type { IServiceInjector } from '@cloudbeaver/core-di';
 import type { ITask } from '@cloudbeaver/core-executor';
 import type { AsyncTaskInfoService, GraphQLService } from '@cloudbeaver/core-sdk';
@@ -17,6 +18,7 @@ export abstract class ResultSetDataSource<TOptions> extends DatabaseDataSource<T
     readonly serviceInjector: IServiceInjector,
     protected graphQLService: GraphQLService,
     protected asyncTaskInfoService: AsyncTaskInfoService,
+    protected connectionExecutionContextService: ConnectionExecutionContextService,
   ) {
     super(serviceInjector);
   }
@@ -80,4 +82,28 @@ export abstract class ResultSetDataSource<TOptions> extends DatabaseDataSource<T
 
     return this.totalCountRequestTask;
   }
+
+  async closeResults(results: IDatabaseResultSet[]): Promise<void> {
+    await this.connectionExecutionContextService.load();
+
+    if (!this.executionContext?.context) {
+      return;
+    }
+
+    for (const result of results) {
+      if (result.id === null) {
+        continue;
+      }
+      try {
+        await this.graphQLService.sdk.closeResult({
+          projectId: result.projectId,
+          connectionId: result.connectionId,
+          contextId: result.contextId,
+          resultId: result.id,
+        });
+      } catch (exception: any) {
+        console.log(`Error closing result (${result.id}):`, exception);
+      }
+    }
+  }
 }
diff --git a/webapp/packages/plugin-sql-editor/src/QueryDataSource.ts b/webapp/packages/plugin-sql-editor/src/QueryDataSource.ts
index cb50959bfb..e27b1ece9f 100644
--- a/webapp/packages/plugin-sql-editor/src/QueryDataSource.ts
+++ b/webapp/packages/plugin-sql-editor/src/QueryDataSource.ts
@@ -7,7 +7,7 @@
  */
 import { makeObservable, observable } from 'mobx';
 
-import type { IConnectionExecutionContextInfo } from '@cloudbeaver/core-connections';
+import type { ConnectionExecutionContextService, IConnectionExecutionContextInfo } from '@cloudbeaver/core-connections';
 import type { IServiceInjector } from '@cloudbeaver/core-di';
 import type { ITask } from '@cloudbeaver/core-executor';
 import {
@@ -53,8 +53,9 @@ export class QueryDataSource<TOptions extends IDataQueryOptions = IDataQueryOpti
     readonly serviceInjector: IServiceInjector,
     graphQLService: GraphQLService,
     asyncTaskInfoService: AsyncTaskInfoService,
+    connectionExecutionContextService: ConnectionExecutionContextService,
   ) {
-    super(serviceInjector, graphQLService, asyncTaskInfoService);
+    super(serviceInjector, graphQLService, asyncTaskInfoService, connectionExecutionContextService);
 
     this.currentTask = null;
     this.requestInfo = {
@@ -254,28 +255,6 @@ export class QueryDataSource<TOptions extends IDataQueryOptions = IDataQueryOpti
     }
   }
 
-  async closeResults(results: IDatabaseResultSet[]) {
-    if (!this.executionContext?.context) {
-      return;
-    }
-
-    for (const result of results) {
-      if (result.id === null) {
-        continue;
-      }
-      try {
-        await this.graphQLService.sdk.closeResult({
-          projectId: result.projectId,
-          connectionId: result.connectionId,
-          contextId: result.contextId,
-          resultId: result.id,
-        });
-      } catch (exception: any) {
-        console.log(`Error closing result (${result.id}):`, exception);
-      }
-    }
-  }
-
   private transformResults(executionContextInfo: IConnectionExecutionContextInfo, results: SqlQueryResults[], limit: number): IDatabaseResultSet[] {
     return results.map<IDatabaseResultSet>((result, index) => ({
       id: result.resultSet?.id || null,
diff --git a/webapp/packages/plugin-sql-editor/src/SqlResultTabs/SqlQueryResultService.ts b/webapp/packages/plugin-sql-editor/src/SqlResultTabs/SqlQueryResultService.ts
index 1c2eee7a68..d0992815da 100644
--- a/webapp/packages/plugin-sql-editor/src/SqlResultTabs/SqlQueryResultService.ts
+++ b/webapp/packages/plugin-sql-editor/src/SqlResultTabs/SqlQueryResultService.ts
@@ -7,7 +7,7 @@
  */
 import { injectable } from '@cloudbeaver/core-di';
 import { uuid } from '@cloudbeaver/core-utils';
-import { IDatabaseDataModel, IDatabaseResultSet, TableViewerStorageService } from '@cloudbeaver/plugin-data-viewer';
+import { IDatabaseDataModel, IDatabaseResultSet, ResultSetDataSource, TableViewerStorageService } from '@cloudbeaver/plugin-data-viewer';
 
 import type { IResultGroup, IResultTab, ISqlEditorTabState, IStatisticsTab } from '../ISqlEditorTabState';
 import type { IDataQueryOptions } from '../QueryDataSource';
@@ -165,8 +165,10 @@ export class SqlQueryResultService {
         const model = this.tableViewerStorageService.get(group.modelId);
         // model?.dispose();
 
-        model?.cancel();
-        model?.source.closeResults?.(model.getResults());
+        if (model?.source instanceof ResultSetDataSource) {
+          model?.source.closeResults?.(model.getResults());
+          model?.cancel();
+        }
 
         this.tableViewerStorageService.remove(group.modelId);
       }
diff --git a/webapp/packages/plugin-sql-editor/src/SqlResultTabs/SqlQueryService.ts b/webapp/packages/plugin-sql-editor/src/SqlResultTabs/SqlQueryService.ts
index 26269a08c7..86c172d35c 100644
--- a/webapp/packages/plugin-sql-editor/src/SqlResultTabs/SqlQueryService.ts
+++ b/webapp/packages/plugin-sql-editor/src/SqlResultTabs/SqlQueryService.ts
@@ -130,7 +130,12 @@ export class SqlQueryService {
     let tabGroup = this.sqlQueryResultService.getSelectedGroup(editorState);
 
     if (inNewTab || !tabGroup) {
-      source = new QueryDataSource(this.app.getServiceInjector(), this.graphQLService, this.asyncTaskInfoService);
+      source = new QueryDataSource(
+        this.app.getServiceInjector(),
+        this.graphQLService,
+        this.asyncTaskInfoService,
+        this.connectionExecutionContextService,
+      );
       model = this.tableViewerStorageService.add(new DatabaseDataModel(source));
       this.dataViewerDataChangeConfirmationService.trackTableDataUpdate(model.id);
       tabGroup = this.sqlQueryResultService.createGroup(editorState, model.id, query);
@@ -216,7 +221,12 @@ export class SqlQueryService {
       options?.onQueryExecutionStart?.(query, i);
 
       if (!model || !source) {
-        source = new QueryDataSource(this.app.getServiceInjector(), this.graphQLService, this.asyncTaskInfoService);
+        source = new QueryDataSource(
+          this.app.getServiceInjector(),
+          this.graphQLService,
+          this.asyncTaskInfoService,
+          this.connectionExecutionContextService,
+        );
         model = this.tableViewerStorageService.add(new DatabaseDataModel(source));
         this.dataViewerDataChangeConfirmationService.trackTableDataUpdate(model.id);
       }

From a37d3327215ad0c84246b45ee2fd6a6bf85bc193 Mon Sep 17 00:00:00 2001
From: sergeyteleshev <iamsergeyteleshev@gmail.com>
Date: Tue, 16 Apr 2024 09:32:50 +0200
Subject: [PATCH 31/36] CB-4270 removes unused context in data sources

---
 .../src/useGroupingDataModel.ts                    |  5 ++---
 .../plugin-data-viewer/src/ContainerDataSource.ts  |  4 ++--
 .../plugin-data-viewer/src/ResultSetDataSource.ts  |  4 ----
 .../plugin-sql-editor/src/QueryDataSource.ts       |  5 ++---
 .../src/SqlResultTabs/SqlQueryResultService.ts     |  4 ++--
 .../src/SqlResultTabs/SqlQueryService.ts           | 14 ++------------
 6 files changed, 10 insertions(+), 26 deletions(-)

diff --git a/webapp/packages/plugin-data-viewer-result-set-grouping/src/useGroupingDataModel.ts b/webapp/packages/plugin-data-viewer-result-set-grouping/src/useGroupingDataModel.ts
index 034dd3efef..d008c625d1 100644
--- a/webapp/packages/plugin-data-viewer-result-set-grouping/src/useGroupingDataModel.ts
+++ b/webapp/packages/plugin-data-viewer-result-set-grouping/src/useGroupingDataModel.ts
@@ -9,7 +9,7 @@ import { reaction } from 'mobx';
 import { useEffect } from 'react';
 
 import { useObjectRef, useResource } from '@cloudbeaver/core-blocks';
-import { ConnectionExecutionContextService, ConnectionInfoResource, createConnectionParam } from '@cloudbeaver/core-connections';
+import { ConnectionInfoResource, createConnectionParam } from '@cloudbeaver/core-connections';
 import { App, useService } from '@cloudbeaver/core-di';
 import { AsyncTaskInfoService, GraphQLService } from '@cloudbeaver/core-sdk';
 import { isObjectsEqual } from '@cloudbeaver/core-utils';
@@ -39,7 +39,6 @@ export function useGroupingDataModel(
   const graphQLService = useService(GraphQLService);
   const asyncTaskInfoService = useService(AsyncTaskInfoService);
   const dataViewerSettingsService = useService(DataViewerSettingsService);
-  const connectionExecutionContextService = useService(ConnectionExecutionContextService);
 
   const executionContext = sourceModel.source.executionContext;
   const contextInfo = executionContext?.context;
@@ -50,7 +49,7 @@ export function useGroupingDataModel(
 
   const model = useObjectRef(
     () => {
-      const source = new GroupingDataSource(app.getServiceInjector(), graphQLService, asyncTaskInfoService, connectionExecutionContextService);
+      const source = new GroupingDataSource(app.getServiceInjector(), graphQLService, asyncTaskInfoService);
 
       const model = tableViewerStorageService.add(new DatabaseDataModel(source));
 
diff --git a/webapp/packages/plugin-data-viewer/src/ContainerDataSource.ts b/webapp/packages/plugin-data-viewer/src/ContainerDataSource.ts
index 52256d6195..6f44808b58 100644
--- a/webapp/packages/plugin-data-viewer/src/ContainerDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/ContainerDataSource.ts
@@ -46,9 +46,9 @@ export class ContainerDataSource extends ResultSetDataSource<IDataContainerOptio
     serviceInjector: IServiceInjector,
     graphQLService: GraphQLService,
     asyncTaskInfoService: AsyncTaskInfoService,
-    connectionExecutionContextService: ConnectionExecutionContextService,
+    protected connectionExecutionContextService: ConnectionExecutionContextService,
   ) {
-    super(serviceInjector, graphQLService, asyncTaskInfoService, connectionExecutionContextService);
+    super(serviceInjector, graphQLService, asyncTaskInfoService);
 
     this.currentTask = null;
     this.executionContext = null;
diff --git a/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts b/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
index 749b1051c3..78efd06a15 100644
--- a/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
@@ -5,7 +5,6 @@
  * Licensed under the Apache License, Version 2.0.
  * you may not use this file except in compliance with the License.
  */
-import type { ConnectionExecutionContextService } from '@cloudbeaver/core-connections';
 import type { IServiceInjector } from '@cloudbeaver/core-di';
 import type { ITask } from '@cloudbeaver/core-executor';
 import type { AsyncTaskInfoService, GraphQLService } from '@cloudbeaver/core-sdk';
@@ -18,7 +17,6 @@ export abstract class ResultSetDataSource<TOptions> extends DatabaseDataSource<T
     readonly serviceInjector: IServiceInjector,
     protected graphQLService: GraphQLService,
     protected asyncTaskInfoService: AsyncTaskInfoService,
-    protected connectionExecutionContextService: ConnectionExecutionContextService,
   ) {
     super(serviceInjector);
   }
@@ -84,8 +82,6 @@ export abstract class ResultSetDataSource<TOptions> extends DatabaseDataSource<T
   }
 
   async closeResults(results: IDatabaseResultSet[]): Promise<void> {
-    await this.connectionExecutionContextService.load();
-
     if (!this.executionContext?.context) {
       return;
     }
diff --git a/webapp/packages/plugin-sql-editor/src/QueryDataSource.ts b/webapp/packages/plugin-sql-editor/src/QueryDataSource.ts
index e27b1ece9f..2ebc42f6e3 100644
--- a/webapp/packages/plugin-sql-editor/src/QueryDataSource.ts
+++ b/webapp/packages/plugin-sql-editor/src/QueryDataSource.ts
@@ -7,7 +7,7 @@
  */
 import { makeObservable, observable } from 'mobx';
 
-import type { ConnectionExecutionContextService, IConnectionExecutionContextInfo } from '@cloudbeaver/core-connections';
+import type { IConnectionExecutionContextInfo } from '@cloudbeaver/core-connections';
 import type { IServiceInjector } from '@cloudbeaver/core-di';
 import type { ITask } from '@cloudbeaver/core-executor';
 import {
@@ -53,9 +53,8 @@ export class QueryDataSource<TOptions extends IDataQueryOptions = IDataQueryOpti
     readonly serviceInjector: IServiceInjector,
     graphQLService: GraphQLService,
     asyncTaskInfoService: AsyncTaskInfoService,
-    connectionExecutionContextService: ConnectionExecutionContextService,
   ) {
-    super(serviceInjector, graphQLService, asyncTaskInfoService, connectionExecutionContextService);
+    super(serviceInjector, graphQLService, asyncTaskInfoService);
 
     this.currentTask = null;
     this.requestInfo = {
diff --git a/webapp/packages/plugin-sql-editor/src/SqlResultTabs/SqlQueryResultService.ts b/webapp/packages/plugin-sql-editor/src/SqlResultTabs/SqlQueryResultService.ts
index d0992815da..a83db3c0f1 100644
--- a/webapp/packages/plugin-sql-editor/src/SqlResultTabs/SqlQueryResultService.ts
+++ b/webapp/packages/plugin-sql-editor/src/SqlResultTabs/SqlQueryResultService.ts
@@ -166,8 +166,8 @@ export class SqlQueryResultService {
         // model?.dispose();
 
         if (model?.source instanceof ResultSetDataSource) {
-          model?.source.closeResults?.(model.getResults());
-          model?.cancel();
+          model.source.closeResults(model.getResults());
+          model.cancel();
         }
 
         this.tableViewerStorageService.remove(group.modelId);
diff --git a/webapp/packages/plugin-sql-editor/src/SqlResultTabs/SqlQueryService.ts b/webapp/packages/plugin-sql-editor/src/SqlResultTabs/SqlQueryService.ts
index 86c172d35c..26269a08c7 100644
--- a/webapp/packages/plugin-sql-editor/src/SqlResultTabs/SqlQueryService.ts
+++ b/webapp/packages/plugin-sql-editor/src/SqlResultTabs/SqlQueryService.ts
@@ -130,12 +130,7 @@ export class SqlQueryService {
     let tabGroup = this.sqlQueryResultService.getSelectedGroup(editorState);
 
     if (inNewTab || !tabGroup) {
-      source = new QueryDataSource(
-        this.app.getServiceInjector(),
-        this.graphQLService,
-        this.asyncTaskInfoService,
-        this.connectionExecutionContextService,
-      );
+      source = new QueryDataSource(this.app.getServiceInjector(), this.graphQLService, this.asyncTaskInfoService);
       model = this.tableViewerStorageService.add(new DatabaseDataModel(source));
       this.dataViewerDataChangeConfirmationService.trackTableDataUpdate(model.id);
       tabGroup = this.sqlQueryResultService.createGroup(editorState, model.id, query);
@@ -221,12 +216,7 @@ export class SqlQueryService {
       options?.onQueryExecutionStart?.(query, i);
 
       if (!model || !source) {
-        source = new QueryDataSource(
-          this.app.getServiceInjector(),
-          this.graphQLService,
-          this.asyncTaskInfoService,
-          this.connectionExecutionContextService,
-        );
+        source = new QueryDataSource(this.app.getServiceInjector(), this.graphQLService, this.asyncTaskInfoService);
         model = this.tableViewerStorageService.add(new DatabaseDataModel(source));
         this.dataViewerDataChangeConfirmationService.trackTableDataUpdate(model.id);
       }

From 6f4482dc7164e500c97b685c597ab01abe808215 Mon Sep 17 00:00:00 2001
From: sergeyteleshev <iamsergeyteleshev@gmail.com>
Date: Thu, 18 Apr 2024 11:57:18 +0200
Subject: [PATCH 32/36] CB-4270 adds loader for CancelTotalCountAction on
 cancel task action

---
 .../TableFooter/CancelTotalCountAction.m.css  |  4 +++
 .../TableFooter/CancelTotalCountAction.tsx    | 29 +++++++++++++++----
 .../TableFooter/TableFooterRowCount.tsx       |  5 ++--
 3 files changed, 31 insertions(+), 7 deletions(-)

diff --git a/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/CancelTotalCountAction.m.css b/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/CancelTotalCountAction.m.css
index fe86b1288c..7e3b80a67d 100644
--- a/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/CancelTotalCountAction.m.css
+++ b/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/CancelTotalCountAction.m.css
@@ -21,3 +21,7 @@
   font-weight: 700;
   font-size: 12px;
 }
+
+.loader {
+  margin-right: 8px !important;
+}
diff --git a/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/CancelTotalCountAction.tsx b/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/CancelTotalCountAction.tsx
index eea1e570f5..a2165f5c6c 100644
--- a/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/CancelTotalCountAction.tsx
+++ b/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/CancelTotalCountAction.tsx
@@ -7,21 +7,40 @@
  */
 import { observer } from 'mobx-react-lite';
 
-import { Container, IconButton, useTranslate } from '@cloudbeaver/core-blocks';
+import { Container, IconButton, Loader, s, useS, useTranslate } from '@cloudbeaver/core-blocks';
 
 import styles from './CancelTotalCountAction.m.css';
 
 interface Props {
   onClick: VoidFunction;
+  loading: boolean;
 }
 
-export const CancelTotalCountAction = observer(function CancelTotalCountAction({ onClick }: Props) {
+export const CancelTotalCountAction = observer(function CancelTotalCountAction({ onClick, loading }: Props) {
   const translate = useTranslate();
+  const style = useS(styles);
+
+  function handleClick() {
+    if (loading) {
+      return;
+    }
+
+    onClick();
+  }
 
   return (
-    <Container className={styles.action} title={translate('ui_processing_cancel')} noWrap center zeroBasis keepSize onClick={onClick}>
-      <IconButton className={styles.icon} name="cross" title={translate('ui_processing_cancel')} viewBox="0 0 32 32" />
-      <span className={styles.cancelText}>{translate('ui_processing_cancel')}</span>
+    <Container
+      className={s(style, { action: true })}
+      title={loading ? translate('ui_processing_canceling') : translate('ui_processing_cancel')}
+      noWrap
+      center
+      zeroBasis
+      keepSize
+      onClick={handleClick}
+    >
+      {loading && <Loader className={s(style, { loader: true })} small />}
+      {!loading && <IconButton disabled={loading} className={s(style, { icon: true })} name="cross" viewBox="0 0 32 32" />}
+      <span className={s(style, { cancelText: true })}>{translate('ui_processing_cancel')}</span>
     </Container>
   );
 });
diff --git a/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TableFooterRowCount.tsx b/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TableFooterRowCount.tsx
index fee254d8c1..afeaca70d9 100644
--- a/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TableFooterRowCount.tsx
+++ b/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TableFooterRowCount.tsx
@@ -10,6 +10,7 @@ import { useState } from 'react';
 
 import { useService } from '@cloudbeaver/core-di';
 import { NotificationService } from '@cloudbeaver/core-events';
+import { isNotNullDefined } from '@cloudbeaver/core-utils';
 
 import type { IDatabaseDataModel } from '../../DatabaseDataModel/IDatabaseDataModel';
 import type { IDatabaseResultSet } from '../../DatabaseDataModel/IDatabaseResultSet';
@@ -45,8 +46,8 @@ export const TableFooterRowCount: React.FC<Props> = observer(function TableFoote
     }
   }
 
-  if (loading) {
-    return <CancelTotalCountAction onClick={cancelTotalCount} />;
+  if (isNotNullDefined(model.source.totalCountRequestTask)) {
+    return <CancelTotalCountAction loading={model.source.totalCountRequestTask.cancelled} onClick={cancelTotalCount} />;
   }
 
   return <TotalCountAction loading={loading} resultIndex={resultIndex} model={model} onClick={loadTotalCount} />;

From d894010d791935f5b9762319319da66854973be3 Mon Sep 17 00:00:00 2001
From: Ainur <ainur.iagudin@dbeaver.com>
Date: Thu, 18 Apr 2024 13:29:28 +0200
Subject: [PATCH 33/36] CB-4270 cancel row count job fix

---
 .../src/io/cloudbeaver/service/sql/impl/WebServiceSQL.java      | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/server/bundles/io.cloudbeaver.server/src/io/cloudbeaver/service/sql/impl/WebServiceSQL.java b/server/bundles/io.cloudbeaver.server/src/io/cloudbeaver/service/sql/impl/WebServiceSQL.java
index d6760fb16c..6658935c90 100644
--- a/server/bundles/io.cloudbeaver.server/src/io/cloudbeaver/service/sql/impl/WebServiceSQL.java
+++ b/server/bundles/io.cloudbeaver.server/src/io/cloudbeaver/service/sql/impl/WebServiceSQL.java
@@ -557,7 +557,7 @@ public void run(DBRProgressMonitor monitor) throws InvocationTargetException, In
                 try {
                     monitor.beginTask("Get row data count", 1);
                     WebSQLResultsInfo results = contextInfo.getResults(resultsId);
-                    long rowCount = DBUtils.readRowCount(webSession.getProgressMonitor(), contextInfo.getProcessor().getExecutionContext(), results.getDataContainer(), null, this);
+                    long rowCount = DBUtils.readRowCount(monitor, contextInfo.getProcessor().getExecutionContext(), results.getDataContainer(), null, this);
                     this.result = "Row data count completed";
                     this.extendedResults = rowCount;
                 } catch (Throwable e) {

From f7b932800e3b6e6d92abc51f650482f4d126f6d5 Mon Sep 17 00:00:00 2001
From: sergeyteleshev <iamsergeyteleshev@gmail.com>
Date: Thu, 18 Apr 2024 20:31:14 +0200
Subject: [PATCH 34/36] CB-4270 cancel row counter with notification

---
 .../DatabaseDataModel/DatabaseDataSource.ts    |  2 +-
 .../DatabaseDataModel/IDatabaseDataSource.ts   |  2 +-
 .../src/ResultSetDataSource.ts                 | 11 +++--------
 .../TableFooter/TableFooterRowCount.tsx        | 18 +++++++++++++-----
 .../plugin-data-viewer/src/locales/en.ts       |  2 ++
 .../plugin-data-viewer/src/locales/it.ts       |  2 ++
 .../plugin-data-viewer/src/locales/ru.ts       |  2 ++
 .../plugin-data-viewer/src/locales/zh.ts       |  2 ++
 8 files changed, 26 insertions(+), 15 deletions(-)

diff --git a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
index 75953dbd32..f5278b30c8 100644
--- a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/DatabaseDataSource.ts
@@ -380,7 +380,7 @@ export abstract class DatabaseDataSource<TOptions, TResult extends IDatabaseData
   abstract request(prevResults: TResult[]): TResult[] | Promise<TResult[]>;
   abstract save(prevResults: TResult[]): Promise<TResult[]> | TResult[];
 
-  abstract loadTotalCount(resultIndex: number): Promise<ITask<number> | null>;
+  abstract loadTotalCount(resultIndex: number): Promise<ITask<number>>;
   abstract cancelLoadTotalCount(): Promise<ITask<number> | null>;
 
   async requestDataAction(): Promise<TResult[] | null> {
diff --git a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/IDatabaseDataSource.ts b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/IDatabaseDataSource.ts
index 69e37eef3f..7a7f9baa4f 100644
--- a/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/IDatabaseDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/DatabaseDataModel/IDatabaseDataSource.ts
@@ -84,7 +84,7 @@ export interface IDatabaseDataSource<TOptions, TResult extends IDatabaseDataResu
   setSupportedDataFormats: (dataFormats: ResultDataFormat[]) => this;
   setExecutionContext: (context: IConnectionExecutionContext | null) => this;
   setTotalCount: (resultIndex: number, count: number) => this;
-  loadTotalCount: (resultIndex: number) => Promise<ITask<number> | null>;
+  loadTotalCount: (resultIndex: number) => Promise<ITask<number>>;
   cancelLoadTotalCount: () => Promise<ITask<number> | null>;
 
   retry: () => Promise<void>;
diff --git a/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts b/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
index 78efd06a15..52b35ebe43 100644
--- a/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
@@ -31,7 +31,7 @@ export abstract class ResultSetDataSource<TOptions> extends DatabaseDataSource<T
     return this.totalCountRequestTask;
   }
 
-  async loadTotalCount(resultIndex: number) {
+  async loadTotalCount(resultIndex: number): Promise<ITask<number>> {
     const executionContext = this.executionContext;
     const executionContextInfo = this.executionContext?.context;
 
@@ -70,13 +70,8 @@ export abstract class ResultSetDataSource<TOptions> extends DatabaseDataSource<T
 
     this.totalCountRequestTask = task;
 
-    try {
-      const count = await task;
-
-      this.setTotalCount(resultIndex, count);
-    } finally {
-      this.totalCountRequestTask = null;
-    }
+    const count = await task;
+    this.setTotalCount(resultIndex, count);
 
     return this.totalCountRequestTask;
   }
diff --git a/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TableFooterRowCount.tsx b/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TableFooterRowCount.tsx
index afeaca70d9..f9abdac440 100644
--- a/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TableFooterRowCount.tsx
+++ b/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TableFooterRowCount.tsx
@@ -31,7 +31,14 @@ export const TableFooterRowCount: React.FC<Props> = observer(function TableFoote
       setLoading(true);
       await model.source.loadTotalCount(resultIndex);
     } catch (exception: any) {
-      notificationService.logException(exception, 'data_viewer_total_count_failed');
+      if (model.source.totalCountRequestTask?.cancelled) {
+        notificationService.logInfo({
+          title: 'data_viewer_total_count_canceled_title',
+          message: 'data_viewer_total_count_canceled_message',
+        });
+      } else {
+        notificationService.logException(exception, 'data_viewer_total_count_failed');
+      }
     } finally {
       setLoading(false);
     }
@@ -39,15 +46,16 @@ export const TableFooterRowCount: React.FC<Props> = observer(function TableFoote
 
   async function cancelTotalCount() {
     try {
-      setLoading(false);
       await model.source.cancelLoadTotalCount();
     } catch (e: any) {
-      notificationService.logException(e);
+      if (!model.source.totalCountRequestTask?.cancelled) {
+        notificationService.logException(e);
+      }
     }
   }
 
-  if (isNotNullDefined(model.source.totalCountRequestTask)) {
-    return <CancelTotalCountAction loading={model.source.totalCountRequestTask.cancelled} onClick={cancelTotalCount} />;
+  if (loading) {
+    return <CancelTotalCountAction loading={Boolean(model.source.totalCountRequestTask?.cancelled)} onClick={cancelTotalCount} />;
   }
 
   return <TotalCountAction loading={loading} resultIndex={resultIndex} model={model} onClick={loadTotalCount} />;
diff --git a/webapp/packages/plugin-data-viewer/src/locales/en.ts b/webapp/packages/plugin-data-viewer/src/locales/en.ts
index 0e525a78f9..c1eba19dd4 100644
--- a/webapp/packages/plugin-data-viewer/src/locales/en.ts
+++ b/webapp/packages/plugin-data-viewer/src/locales/en.ts
@@ -52,6 +52,8 @@ export default [
   ['data_viewer_script_preview_error_title', "Can't get the script"],
   ['data_viewer_refresh_result_set', 'Refresh result set'],
   ['data_viewer_total_count_tooltip', 'Get total count'],
+  ['data_viewer_total_count_canceled_title', 'Total count canceled'],
+  ['data_viewer_total_count_canceled_message', 'Statement was cancelled due to user request'],
   ['data_viewer_total_count_failed', 'Failed to get total count'],
   ['data_viewer_model_not_loaded', 'Table model is not loaded'],
   ['settings_data_editor', 'Data Editor'],
diff --git a/webapp/packages/plugin-data-viewer/src/locales/it.ts b/webapp/packages/plugin-data-viewer/src/locales/it.ts
index 7c1a13b91c..a0b3a9f7ad 100644
--- a/webapp/packages/plugin-data-viewer/src/locales/it.ts
+++ b/webapp/packages/plugin-data-viewer/src/locales/it.ts
@@ -45,6 +45,8 @@ export default [
   ['data_viewer_presentation_value_content_paste_error', 'Cannot load full text'],
   ['data_viewer_refresh_result_set', 'Refresh result set'],
   ['data_viewer_total_count_tooltip', 'Get total count'],
+  ['data_viewer_total_count_canceled_title', 'Total count canceled'],
+  ['data_viewer_total_count_canceled_message', 'Statement was cancelled due to user request'],
   ['data_viewer_total_count_failed', 'Failed to get total count'],
   ['data_viewer_model_not_loaded', 'Table model is not loaded'],
   ['settings_data_editor', 'Data Editor'],
diff --git a/webapp/packages/plugin-data-viewer/src/locales/ru.ts b/webapp/packages/plugin-data-viewer/src/locales/ru.ts
index cd1e69c78a..76629ec648 100644
--- a/webapp/packages/plugin-data-viewer/src/locales/ru.ts
+++ b/webapp/packages/plugin-data-viewer/src/locales/ru.ts
@@ -46,6 +46,8 @@ export default [
   ['data_viewer_script_preview_error_title', 'Не удалось получить скрипт'],
   ['data_viewer_refresh_result_set', 'Обновить резалт сет'],
   ['data_viewer_total_count_tooltip', 'Получить количество записей'],
+  ['data_viewer_total_count_canceled_title', 'Получение количества записей отменено'],
+  ['data_viewer_total_count_canceled_message', 'Запрос был отменен пользователем'],
   ['data_viewer_total_count_failed', 'Не удалось получить количество записей'],
   ['data_viewer_model_not_loaded', 'Не удалось загрузить модель таблицы'],
   ['settings_data_editor', 'Редактор данных'],
diff --git a/webapp/packages/plugin-data-viewer/src/locales/zh.ts b/webapp/packages/plugin-data-viewer/src/locales/zh.ts
index 368f2f6e8f..cb02a92fc8 100644
--- a/webapp/packages/plugin-data-viewer/src/locales/zh.ts
+++ b/webapp/packages/plugin-data-viewer/src/locales/zh.ts
@@ -53,6 +53,8 @@ export default [
   ['data_viewer_refresh_result_set', 'Refresh result set'],
   ['data_viewer_total_count_failed', 'Failed to get total count'],
   ['data_viewer_total_count_tooltip', 'Get total count'],
+  ['data_viewer_total_count_canceled_title', 'Total count canceled'],
+  ['data_viewer_total_count_canceled_message', 'Statement was cancelled due to user request'],
   ['data_viewer_model_not_loaded', 'Table model is not loaded'],
   ['settings_data_editor', 'Data Editor'],
   ['settings_data_editor_disable_edit_name', 'Disable Edit'],

From ea2a1cff696c186f461f5e220957d2fbb9cc225e Mon Sep 17 00:00:00 2001
From: sergeyteleshev <iamsergeyteleshev@gmail.com>
Date: Fri, 19 Apr 2024 13:07:36 +0200
Subject: [PATCH 35/36] CB-4270 pr fixes

---
 .../TableViewer/TableFooter/CancelTotalCountAction.m.css  | 8 ++++----
 .../TableViewer/TableFooter/CancelTotalCountAction.tsx    | 2 +-
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/CancelTotalCountAction.m.css b/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/CancelTotalCountAction.m.css
index 7e3b80a67d..8531a5867f 100644
--- a/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/CancelTotalCountAction.m.css
+++ b/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/CancelTotalCountAction.m.css
@@ -12,6 +12,10 @@
   user-select: none;
 }
 
+.action .loader {
+  margin-right: 8px;
+}
+
 .icon {
   composes: theme-text-error from global;
 }
@@ -21,7 +25,3 @@
   font-weight: 700;
   font-size: 12px;
 }
-
-.loader {
-  margin-right: 8px !important;
-}
diff --git a/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/CancelTotalCountAction.tsx b/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/CancelTotalCountAction.tsx
index a2165f5c6c..d1b0d19dc6 100644
--- a/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/CancelTotalCountAction.tsx
+++ b/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/CancelTotalCountAction.tsx
@@ -31,7 +31,7 @@ export const CancelTotalCountAction = observer(function CancelTotalCountAction({
   return (
     <Container
       className={s(style, { action: true })}
-      title={loading ? translate('ui_processing_canceling') : translate('ui_processing_cancel')}
+      title={translate(loading ? 'ui_processing_canceling' : 'ui_processing_cancel')}
       noWrap
       center
       zeroBasis

From 8e7f48b8f8e7b636dfd51a32f458bdaa4729c4d5 Mon Sep 17 00:00:00 2001
From: sergeyteleshev <iamsergeyteleshev@gmail.com>
Date: Fri, 19 Apr 2024 16:36:11 +0200
Subject: [PATCH 36/36] CB-4270 pr fixes

---
 webapp/packages/plugin-data-viewer/src/ContainerDataSource.ts  | 3 ++-
 webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts  | 3 ++-
 .../src/TableViewer/TableFooter/CancelTotalCountAction.tsx     | 2 +-
 .../src/TableViewer/TableFooter/TotalCountAction.tsx           | 2 +-
 webapp/packages/plugin-sql-editor/src/QueryDataSource.ts       | 3 ++-
 5 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/webapp/packages/plugin-data-viewer/src/ContainerDataSource.ts b/webapp/packages/plugin-data-viewer/src/ContainerDataSource.ts
index 6f44808b58..1e6faa5d96 100644
--- a/webapp/packages/plugin-data-viewer/src/ContainerDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/ContainerDataSource.ts
@@ -68,7 +68,8 @@ export class ContainerDataSource extends ResultSetDataSource<IDataContainerOptio
   }
 
   async cancel(): Promise<void> {
-    await Promise.all([this.currentTask?.cancel(), super.cancel()]);
+    await super.cancel();
+    await this.currentTask?.cancel();
   }
 
   async request(prevResults: IDatabaseResultSet[]): Promise<IDatabaseResultSet[]> {
diff --git a/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts b/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
index 52b35ebe43..04aa079a52 100644
--- a/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
+++ b/webapp/packages/plugin-data-viewer/src/ResultSetDataSource.ts
@@ -22,7 +22,8 @@ export abstract class ResultSetDataSource<TOptions> extends DatabaseDataSource<T
   }
 
   async cancel(): Promise<void> {
-    await Promise.all([this.cancelLoadTotalCount(), super.cancel()]);
+    await super.cancel();
+    await this.cancelLoadTotalCount();
   }
 
   async cancelLoadTotalCount(): Promise<ITask<number> | null> {
diff --git a/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/CancelTotalCountAction.tsx b/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/CancelTotalCountAction.tsx
index d1b0d19dc6..28502a16d0 100644
--- a/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/CancelTotalCountAction.tsx
+++ b/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/CancelTotalCountAction.tsx
@@ -16,7 +16,7 @@ interface Props {
   loading: boolean;
 }
 
-export const CancelTotalCountAction = observer(function CancelTotalCountAction({ onClick, loading }: Props) {
+export const CancelTotalCountAction = observer<Props>(function CancelTotalCountAction({ onClick, loading }) {
   const translate = useTranslate();
   const style = useS(styles);
 
diff --git a/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TotalCountAction.tsx b/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TotalCountAction.tsx
index 7e1587a032..7a4259728d 100644
--- a/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TotalCountAction.tsx
+++ b/webapp/packages/plugin-data-viewer/src/TableViewer/TableFooter/TotalCountAction.tsx
@@ -22,7 +22,7 @@ interface Props {
   model: IDatabaseDataModel<any, IDatabaseResultSet>;
 }
 
-export const TotalCountAction = observer(function TotalCountAction({ onClick, loading, resultIndex, model }: Props) {
+export const TotalCountAction = observer<Props>(function TotalCountAction({ onClick, loading, resultIndex, model }) {
   const result = model.getResult(resultIndex);
   const translate = useTranslate();
   const disabled = getComputed(() => model.isLoading() || model.isDisabled(resultIndex));
diff --git a/webapp/packages/plugin-sql-editor/src/QueryDataSource.ts b/webapp/packages/plugin-sql-editor/src/QueryDataSource.ts
index 2ebc42f6e3..5cb9acace8 100644
--- a/webapp/packages/plugin-sql-editor/src/QueryDataSource.ts
+++ b/webapp/packages/plugin-sql-editor/src/QueryDataSource.ts
@@ -84,7 +84,8 @@ export class QueryDataSource<TOptions extends IDataQueryOptions = IDataQueryOpti
   }
 
   async cancel(): Promise<void> {
-    await Promise.all([this.currentTask?.cancel(), super.cancel()]);
+    await super.cancel();
+    await this.currentTask?.cancel();
   }
 
   async save(prevResults: IDatabaseResultSet[]): Promise<IDatabaseResultSet[]> {