CB-3777 api for getting row count (#2100)

* CB-3777 api for getting row count

* CB-4176 add rows total count

* CB-4177 count rows from sql queries

* CB-4176 review fixes

* CB-3777 refactor: use active task instead of currentTask in database datasource

* CB-3777 fix: don't use private params

* CB-4176 revert currentTask

* CB-4176 add error message

* CB-4176 add tooltip

* CB-4176 change locales

---------

Co-authored-by: naumov <iamemptyhuh@gmail.com>
Co-authored-by: Aleksei Potsetsuev <wrouds@gmail.com>
Co-authored-by: Evgenia Bezborodova <139753579+EvgeniaBzzz@users.noreply.github.com>

server/bundles/io.cloudbeaver.server/schema/service.sql.graphqls

@@ -356,4 +356,12 @@ extend type Mutation {
     # Returns SQLExecutionPlan
     asyncSqlExplainExecutionPlanResult(taskId: ID!): SQLExecutionPlan !
 
+    asyncSqlRowDataCount(
+        projectId: ID,
+        connectionId: ID!,
+        contextId: ID!,
+        resultsId: ID!
+    ): AsyncTaskInfo!
+
+    asyncSqlRowDataCountResult(taskId: ID!): Int!
 }

server/bundles/io.cloudbeaver.server/src/io/cloudbeaver/service/sql/DBWServiceSQL.java

@@ -157,4 +157,11 @@ public interface DBWServiceSQL extends DBWService {
                                 @NotNull List<String> columnsList,
                                 @Nullable List<String> functions,
                                 @Nullable Boolean showDuplicatesOnly) throws DBWebException;
+
+    @WebAction
+    WebAsyncTaskInfo getRowDataCount(@NotNull WebSession webSession, @NotNull WebSQLContextInfo contextInfo, @NotNull String resultsId) throws DBWebException;
+
+    @Nullable
+    @WebAction
+    Long getRowDataCountResult(@NotNull WebSession webSession, @NotNull String taskId) throws DBWebException;
 }

server/bundles/io.cloudbeaver.server/src/io/cloudbeaver/service/sql/WebSQLProcessor.java

@@ -171,7 +171,7 @@ public class WebSQLProcessor implements WebSessionProvider {
         long startTime = System.currentTimeMillis();
         WebSQLExecuteInfo executeInfo = new WebSQLExecuteInfo();
 
-        var dataContainer = new WebSQLQueryDataContainer(connection.getDataSource(), sql);
+        var dataContainer = new WebSQLQueryDataContainer(connection.getDataSource(), syntaxManager, sql);
 
         DBCExecutionContext context = getExecutionContext(dataContainer);
 

server/bundles/io.cloudbeaver.server/src/io/cloudbeaver/service/sql/WebSQLQueryDataContainer.java

@@ -18,6 +18,7 @@ package io.cloudbeaver.service.sql;
 
 import org.jkiss.code.NotNull;
 import org.jkiss.code.Nullable;
+import org.jkiss.dbeaver.DBException;
 import org.jkiss.dbeaver.Log;
 import org.jkiss.dbeaver.model.DBPContextProvider;
 import org.jkiss.dbeaver.model.DBPDataSource;
@@ -25,11 +26,14 @@ import org.jkiss.dbeaver.model.DBUtils;
 import org.jkiss.dbeaver.model.data.DBDDataFilter;
 import org.jkiss.dbeaver.model.data.DBDDataReceiver;
 import org.jkiss.dbeaver.model.exec.*;
+import org.jkiss.dbeaver.model.impl.sql.SQLQueryTransformerCount;
 import org.jkiss.dbeaver.model.sql.SQLQuery;
 import org.jkiss.dbeaver.model.sql.SQLScriptContext;
+import org.jkiss.dbeaver.model.sql.SQLSyntaxManager;
 import org.jkiss.dbeaver.model.sql.data.SQLQueryDataContainer;
 import org.jkiss.dbeaver.model.struct.DBSDataContainer;
 import org.jkiss.dbeaver.model.struct.DBSObject;
+import org.jkiss.utils.CommonUtils;
 
 import java.io.PrintWriter;
 
@@ -41,11 +45,13 @@ public class WebSQLQueryDataContainer implements DBSDataContainer, DBPContextPro
     private static final Log log = Log.getLog(WebSQLQueryDataContainer.class);
 
     private final DBPDataSource dataSource;
+    private final SQLSyntaxManager syntaxManager;
     private final String query;
     private final SQLQueryDataContainer queryDataContainer;
 
-    public WebSQLQueryDataContainer(DBPDataSource dataSource, String query) {
+    public WebSQLQueryDataContainer(DBPDataSource dataSource, SQLSyntaxManager syntaxManager, String query) {
         this.dataSource = dataSource;
+        this.syntaxManager = syntaxManager;
         this.query = query;
 
         SQLScriptContext scriptContext = new SQLScriptContext(null,
@@ -95,7 +101,12 @@ public class WebSQLQueryDataContainer implements DBSDataContainer, DBPContextPro
 
     @Override
     public long countData(@NotNull DBCExecutionSource source, @NotNull DBCSession session, @Nullable DBDDataFilter dataFilter, long flags) throws DBCException {
-        return queryDataContainer.countData(source, session, dataFilter, flags);
+        try {
+            SQLQuery countQuery = new SQLQueryTransformerCount().transformQuery(dataSource, syntaxManager, new SQLQuery(dataSource, query));
+            return DBUtils.countDataFromQuery(source, session, countQuery);
+        } catch (DBException e) {
+            throw new DBCException("Error executing row count", e);
+        }
     }
 
     @Nullable

server/bundles/io.cloudbeaver.server/src/io/cloudbeaver/service/sql/WebServiceBindingSQL.java

@@ -184,7 +184,19 @@ public class WebServiceBindingSQL extends WebServiceBindingBase<DBWServiceSQL> i
             .dataFetcher("asyncSqlExplainExecutionPlanResult", env ->
                 getService(env).asyncSqlExplainExecutionPlanResult(
                     getWebSession(env), env.getArgument("taskId")
-                ));
+                ))
+            .dataFetcher("asyncSqlRowDataCount", env ->
+                getService(env).getRowDataCount(
+                    getWebSession(env),
+                    getSQLContext(env),
+                    env.getArgument("resultsId")
+                ))
+            .dataFetcher("asyncSqlRowDataCountResult", env ->
+                getService(env).getRowDataCountResult(
+                    getWebSession(env),
+                    env.getArgument("taskId")
+            )
+        );
     }
 
     @NotNull

server/bundles/io.cloudbeaver.server/src/io/cloudbeaver/service/sql/impl/WebServiceSQL.java

@@ -31,6 +31,7 @@ import org.jkiss.dbeaver.DBException;
 import org.jkiss.dbeaver.Log;
 import org.jkiss.dbeaver.model.DBPDataSource;
 import org.jkiss.dbeaver.model.DBPDataSourceContainer;
+import org.jkiss.dbeaver.model.DBUtils;
 import org.jkiss.dbeaver.model.data.DBDAttributeBinding;
 import org.jkiss.dbeaver.model.exec.DBCException;
 import org.jkiss.dbeaver.model.exec.DBCLogicalOperator;
@@ -516,4 +517,36 @@ public class WebServiceSQL implements DBWServiceSQL {
             throw new DBWebException("Error on generating GROUP BY query", e);
         }
     }
+
+    @Override
+    public WebAsyncTaskInfo getRowDataCount(@NotNull WebSession webSession, @NotNull WebSQLContextInfo contextInfo, @NotNull String resultsId) throws DBWebException {
+        WebAsyncTaskProcessor<String> runnable = new WebAsyncTaskProcessor<String>() {
+            @Override
+            public void run(DBRProgressMonitor monitor) throws InvocationTargetException, InterruptedException {
+                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);
+                    this.result = "Row data count completed";
+                    this.extendedResults = rowCount;
+                } catch (Throwable e) {
+                    throw new InvocationTargetException(e);
+                } finally {
+                    monitor.done();
+                }
+            }
+        };
+        return contextInfo.getProcessor().getWebSession().createAndRunAsyncTask("SQL result set count rows", runnable);
+    }
+
+    @Override
+    @Nullable
+    public Long getRowDataCountResult(@NotNull WebSession webSession, @NotNull String taskId) throws DBWebException {
+        WebAsyncTaskInfo taskStatus = webSession.asyncTaskStatus(taskId, false);
+        if (taskStatus != null) {
+            return (Long) taskStatus.getExtendedResult();
+        }
+        return null;
+    }
+
 }