#3322 Async server output reading

Former-commit-id: 54122f0e

plugins/org.jkiss.dbeaver.core/src/org/jkiss/dbeaver/runtime/sql/SQLQueryJob.java

@@ -147,6 +147,10 @@ public class SQLQueryJob extends DataSourceJob
         return lastGoodQuery;
     }
 
+    public DBCStatement getCurrentStatement() {
+        return curStatement;
+    }
+
     private boolean hasLimits()
     {
         return rsOffset >= 0 && rsMaxRows > 0;

plugins/org.jkiss.dbeaver.core/src/org/jkiss/dbeaver/ui/editors/sql/SQLEditor.java

@@ -251,14 +251,12 @@ public class SQLEditor extends SQLEditorBase implements
             EditorUtils.setInputDataSource(input, container, true);
         }
 
-        checkConnected(false, status -> {
-            DBeaverUI.asyncExec(() -> {
-                if (!status.isOK()) {
-                    DBUserInterface.getInstance().showError("Can't connect to database", "Error connecting to datasource", status);
-                }
-                setFocus();
-            });
-        });
+        checkConnected(false, status -> DBeaverUI.asyncExec(() -> {
+            if (!status.isOK()) {
+                DBUserInterface.getInstance().showError("Can't connect to database", "Error connecting to datasource", status);
+            }
+            setFocus();
+        }));
         setPartName(getEditorName());
 
         fireDataSourceChange();
@@ -440,17 +438,17 @@ public class SQLEditor extends SQLEditorBase implements
 
     @Nullable
     @Override
-    public Object getAdapter(Class required)
+    public <T> T getAdapter(Class<T> required)
     {
         if (required == IFindReplaceTarget.class) {
-            return findReplaceTarget;
+            return required.cast(findReplaceTarget);
         }
         ResultSetViewer resultsView = getActiveResultSetViewer();
         if (resultsView != null) {
             if (required == ResultSetViewer.class) {
-                return resultsView;
+                return required.cast(resultsView);
             }
-            Object adapter = resultsView.getAdapter(required);
+            T adapter = resultsView.getAdapter(required);
             if (adapter != null) {
                 return adapter;
             }
@@ -509,7 +507,6 @@ public class SQLEditor extends SQLEditorBase implements
     /**
      * Sets focus in current editor.
      * This function is called on drag-n-drop and some other operations
-     * @return
      */
     @Override
     public boolean validateEditorInputState() {
@@ -1022,8 +1019,7 @@ public class SQLEditor extends SQLEditorBase implements
                 DBPDataSource dataSource = getDataSource();
                 if (dataSource instanceof SQLDataSource) {
                     List<SQLScriptElement> xQueries = new ArrayList<>(elements.size());
-                    for (int i = 0; i < elements.size(); i++) {
-                        SQLScriptElement element = elements.get(i);
+                    for (SQLScriptElement element : elements) {
                         if (element instanceof SQLQuery) {
                             SQLQuery query = transformer.transformQuery((SQLDataSource) dataSource, getSyntaxManager(), (SQLQuery) element);
                             if (!CommonUtils.isEmpty(query.getParameters())) {
@@ -1074,12 +1070,8 @@ public class SQLEditor extends SQLEditorBase implements
                             status);
                         return;
                     }
-                    updateExecutionContext(new Runnable() {
-                        @Override
-                        public void run() {
-                            DBeaverUI.syncExec(() -> processQueries(queries, newTab, export, false, queryListener));
-                        }
-                    });
+                    updateExecutionContext(() -> DBeaverUI.syncExec(() ->
+                        processQueries(queries, newTab, export, false, queryListener)));
                 };
                 if (!checkSession(connectListener)) {
                     return;
@@ -1484,7 +1476,7 @@ public class SQLEditor extends SQLEditorBase implements
 
     private void showScriptPositionRuler(boolean show)
     {
-        IColumnSupport columnSupport = (IColumnSupport) getAdapter(IColumnSupport.class);
+        IColumnSupport columnSupport = getAdapter(IColumnSupport.class);
         if (columnSupport != null) {
             RulerColumnDescriptor positionColumn = RulerColumnRegistry.getDefault().getColumnDescriptor(ScriptPositionColumn.ID);
             columnSupport.setColumnVisible(positionColumn, show);
@@ -1614,10 +1606,10 @@ public class SQLEditor extends SQLEditorBase implements
 
     public class QueryProcessor implements SQLResultsConsumer {
 
-        private SQLQueryJob curJob;
+        private volatile SQLQueryJob curJob;
         private AtomicInteger curJobRunning = new AtomicInteger(0);
         private final List<QueryResultsContainer> resultContainers = new ArrayList<>();
-        private DBDDataReceiver curDataReceiver = null;
+        private volatile DBDDataReceiver curDataReceiver = null;
 
         QueryProcessor() {
             // Create first (default) results provider
@@ -2084,13 +2076,8 @@ public class SQLEditor extends SQLEditorBase implements
         }
 
         @Override
-        public boolean isPersisted()
-        {
-            if (dataContainer != null) {
-                return dataContainer.isPersisted();
-            } else {
-                return true;
-            }
+        public boolean isPersisted() {
+            return dataContainer == null || dataContainer.isPersisted();
         }
 
         @NotNull
@@ -2194,9 +2181,8 @@ public class SQLEditor extends SQLEditorBase implements
             try {
                 boolean isInExecute = getTotalQueryRunning() > 0;
                 if (!isInExecute) {
-                    DBeaverUI.asyncExec(() -> {
-                        setTitleImage(DBeaverIcons.getImage(UIIcon.SQL_SCRIPT_EXECUTE));
-                    });
+                    DBeaverUI.asyncExec(() ->
+                        setTitleImage(DBeaverIcons.getImage(UIIcon.SQL_SCRIPT_EXECUTE)));
                 }
                 queryProcessor.curJobRunning.incrementAndGet();
                 synchronized (runningQueries) {
@@ -2228,9 +2214,8 @@ public class SQLEditor extends SQLEditorBase implements
                 }
                 queryProcessor.curJobRunning.decrementAndGet();
                 if (getTotalQueryRunning() <= 0) {
-                    DBeaverUI.asyncExec(() -> {
-                        setTitleImage(editorImage);
-                    });
+                    DBeaverUI.asyncExec(() ->
+                        setTitleImage(editorImage));
                 }
 
                 if (isDisposed()) {
@@ -2248,7 +2233,7 @@ public class SQLEditor extends SQLEditorBase implements
         }
 
         private void processQueryResult(DBCSession session, SQLQueryResult result) {
-            dumpServerOutput(result);
+            dumpQueryServerOutput(result);
             if (!scriptMode) {
                 runPostExecuteActions(result);
             }
@@ -2393,18 +2378,18 @@ public class SQLEditor extends SQLEditorBase implements
         }
     }
 
-    private void dumpServerOutput(@Nullable SQLQueryResult result) {
+    private void dumpQueryServerOutput(@Nullable SQLQueryResult result) {
         final DBCExecutionContext executionContext = getExecutionContext();
         if (executionContext != null) {
             final DBPDataSource dataSource = executionContext.getDataSource();
             // Dump server output
             DBCServerOutputReader outputReader = DBUtils.getAdapter(DBCServerOutputReader.class, dataSource);
             if (outputReader == null && result != null) {
-                outputReader = new DefaultServerOutputReader(result);
+                outputReader = new DefaultServerOutputReader();
             }
             if (outputReader != null && outputReader.isServerOutputEnabled()) {
                 synchronized (serverOutputs) {
-                    serverOutputs.add(new ServerOutputInfo(outputReader, executionContext));
+                    serverOutputs.add(new ServerOutputInfo(outputReader, executionContext, result));
                 }
             }
         }
@@ -2479,10 +2464,12 @@ public class SQLEditor extends SQLEditorBase implements
     private static class ServerOutputInfo {
         private final DBCServerOutputReader outputReader;
         private final DBCExecutionContext executionContext;
+        private final SQLQueryResult result;
 
-        ServerOutputInfo(DBCServerOutputReader outputReader, DBCExecutionContext executionContext) {
+        ServerOutputInfo(DBCServerOutputReader outputReader, DBCExecutionContext executionContext, SQLQueryResult result) {
             this.outputReader = outputReader;
             this.executionContext = executionContext;
+            this.result = result;
         }
     }
 
@@ -2513,31 +2500,58 @@ public class SQLEditor extends SQLEditorBase implements
                 serverOutputs.clear();
             }
 
+            final StringWriter dump = new StringWriter();
             for (ServerOutputInfo info : outputs) {
-                final StringWriter dump = new StringWriter();
                 try {
-                    info.outputReader.readServerOutput(monitor, info.executionContext, new PrintWriter(dump, true));
-                    final String dumpString = dump.toString();
-                    if (!dumpString.isEmpty()) {
-                        DBeaverUI.asyncExec(() -> {
-                            if (outputViewer.isDisposed()) {
-                                return;
-                            }
-                            try {
-                                IOUtils.copyText(new StringReader(dumpString), outputViewer.getOutputWriter());
-                            } catch (IOException e) {
-                                log.error(e);
-                            }
-                            if (outputViewer.isHasNewOutput()) {
-                                outputViewer.scrollToEnd();
-                                updateOutputViewerIcon(true);
-                            }
-                        });
-                    }
+                    info.outputReader.readServerOutput(monitor, info.executionContext, info.result, null, new PrintWriter(dump, true));
                 } catch (Exception e) {
                     log.error(e);
                 }
             }
+
+            {
+                // Check running queries for async output
+                DBCServerOutputReader outputReader = null;
+                final DBCExecutionContext executionContext = getExecutionContext();
+                if (executionContext != null) {
+                    final DBPDataSource dataSource = executionContext.getDataSource();
+                    // Dump server output
+                    outputReader = DBUtils.getAdapter(DBCServerOutputReader.class, dataSource);
+                }
+                if (outputReader != null && outputReader.isAsyncOutputReadSupported()) {
+                    for (QueryProcessor qp : queryProcessors) {
+                        SQLQueryJob queryJob = qp.curJob;
+                        if (queryJob != null) {
+                            DBCStatement statement = queryJob.getCurrentStatement();
+                            if (statement != null) {
+                                try {
+                                    outputReader.readServerOutput(monitor, executionContext, null, statement, new PrintWriter(dump, true));
+                                } catch (DBCException e) {
+                                    log.error(e);
+                                }
+                            }
+                        }
+                    }
+                }
+            }
+
+            final String dumpString = dump.toString();
+            if (!dumpString.isEmpty()) {
+                DBeaverUI.asyncExec(() -> {
+                    if (outputViewer.isDisposed()) {
+                        return;
+                    }
+                    try {
+                        IOUtils.copyText(new StringReader(dumpString), outputViewer.getOutputWriter());
+                    } catch (IOException e) {
+                        log.error(e);
+                    }
+                    if (outputViewer.isHasNewOutput()) {
+                        outputViewer.scrollToEnd();
+                        updateOutputViewerIcon(true);
+                    }
+                });
+            }
         }
     }
 }

plugins/org.jkiss.dbeaver.ext.oracle/src/org/jkiss/dbeaver/ext/oracle/model/OracleDataSource.java

@@ -40,6 +40,7 @@ import org.jkiss.dbeaver.model.impl.jdbc.cache.JDBCStructCache;
 import org.jkiss.dbeaver.model.meta.Association;
 import org.jkiss.dbeaver.model.runtime.DBRProgressMonitor;
 import org.jkiss.dbeaver.model.sql.SQLConstants;
+import org.jkiss.dbeaver.model.sql.SQLQueryResult;
 import org.jkiss.dbeaver.model.sql.SQLState;
 import org.jkiss.dbeaver.model.struct.*;
 import org.jkiss.dbeaver.runtime.ui.DBUserInterface;
@@ -709,6 +710,10 @@ public class OracleDataSource extends JDBCDataSource
         }
 
         @Override
+        public boolean isAsyncOutputReadSupported() {
+            return false;
+        }
+
         public void enableServerOutput(DBRProgressMonitor monitor, DBCExecutionContext context, boolean enable) throws DBCException {
             String sql = enable ?
                 "BEGIN DBMS_OUTPUT.ENABLE(" + OracleConstants.MAXIMUM_DBMS_OUTPUT_SIZE + "); END;" :
@@ -721,7 +726,7 @@ public class OracleDataSource extends JDBCDataSource
         }
 
         @Override
-        public void readServerOutput(DBRProgressMonitor monitor, DBCExecutionContext context, PrintWriter output) throws DBCException {
+        public void readServerOutput(@NotNull DBRProgressMonitor monitor, @NotNull DBCExecutionContext context, @Nullable SQLQueryResult queryResult, @Nullable DBCStatement statement, @NotNull PrintWriter output) throws DBCException {
             try (JDBCSession session = (JDBCSession) context.openSession(monitor, DBCExecutionPurpose.UTIL, "Read DBMS output")) {
                 try (CallableStatement getLineProc = session.getOriginal().prepareCall("{CALL DBMS_OUTPUT.GET_LINE(?, ?)}")) {
                     getLineProc.registerOutParameter(1, java.sql.Types.VARCHAR);

plugins/org.jkiss.dbeaver.ext.postgresql/src/org/jkiss/dbeaver/ext/postgresql/model/PostgreDataSource.java

@@ -34,6 +34,7 @@ import org.jkiss.dbeaver.model.exec.jdbc.*;
 import org.jkiss.dbeaver.model.exec.plan.DBCPlan;
 import org.jkiss.dbeaver.model.exec.plan.DBCPlanStyle;
 import org.jkiss.dbeaver.model.exec.plan.DBCQueryPlanner;
+import org.jkiss.dbeaver.model.impl.AsyncServerOutputReader;
 import org.jkiss.dbeaver.model.impl.jdbc.JDBCDataSource;
 import org.jkiss.dbeaver.model.impl.jdbc.JDBCExecutionContext;
 import org.jkiss.dbeaver.model.impl.jdbc.JDBCUtils;
@@ -412,12 +413,9 @@ public class PostgreDataSource extends JDBCDataSource implements DBSObjectSelect
     {
         if (adapter == DBSStructureAssistant.class) {
             return adapter.cast(new PostgreStructureAssistant(this));
+        } else if (adapter == DBCServerOutputReader.class) {
+            return adapter.cast(new AsyncServerOutputReader());
         }
-/*
-        else if (adapter == DBAServerSessionManager.class) {
-            return new PostgreSessionManager(this);
-        }
-*/
         return super.getAdapter(adapter);
     }
 

plugins/org.jkiss.dbeaver.model/src/org/jkiss/dbeaver/model/exec/DBCServerOutputReader.java

@@ -16,9 +16,11 @@
  */
 package org.jkiss.dbeaver.model.exec;
 
-import org.jkiss.dbeaver.model.DBPCloseableObject;
+import org.jkiss.code.NotNull;
+import org.jkiss.code.Nullable;
 import org.jkiss.dbeaver.model.DBPObject;
 import org.jkiss.dbeaver.model.runtime.DBRProgressMonitor;
+import org.jkiss.dbeaver.model.sql.SQLQueryResult;
 
 import java.io.PrintWriter;
 
@@ -29,8 +31,21 @@ public interface DBCServerOutputReader extends DBPObject
 {
     boolean isServerOutputEnabled();
 
-    void enableServerOutput(DBRProgressMonitor monitor, DBCExecutionContext context, boolean enable) throws DBCException;
+    /**
+     * If async output reading is supported then SQL job will read output during statement execution.
+     */
+    boolean isAsyncOutputReadSupported();
 
-    void readServerOutput(DBRProgressMonitor monitor, DBCExecutionContext context, PrintWriter output)
+    /**
+     * Reads server output messages.
+     * Only @queryResult or @statement can be specified. Non-null statement means async output reading.
+     * Output for statement can be requested only if @isAsyncOutputReadSupported returns true.
+     */
+    void readServerOutput(
+        @NotNull DBRProgressMonitor monitor,
+        @NotNull DBCExecutionContext context,
+        @Nullable SQLQueryResult queryResult,
+        @Nullable DBCStatement statement,
+        @NotNull PrintWriter output)
         throws DBCException;
 }

plugins/org.jkiss.dbeaver.model/src/org/jkiss/dbeaver/model/impl/AsyncServerOutputReader.java

+/*
+ * DBeaver - Universal Database Manager
+ * Copyright (C) 2010-2018 Serge Rider (serge@jkiss.org)
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.jkiss.dbeaver.model.impl;
+
+import org.jkiss.dbeaver.model.exec.DBCException;
+import org.jkiss.dbeaver.model.exec.DBCExecutionContext;
+import org.jkiss.dbeaver.model.exec.DBCStatement;
+import org.jkiss.dbeaver.model.runtime.DBRProgressMonitor;
+import org.jkiss.dbeaver.model.sql.SQLQueryResult;
+
+import java.io.PrintWriter;
+import java.util.Arrays;
+
+public class AsyncServerOutputReader extends DefaultServerOutputReader {
+        @Override
+        public boolean isAsyncOutputReadSupported() {
+            return true;
+        }
+
+        @Override
+        public void readServerOutput(DBRProgressMonitor monitor, DBCExecutionContext context, SQLQueryResult queryResult, DBCStatement statement, PrintWriter output) throws DBCException {
+            if (statement == null) {
+                super.readServerOutput(monitor, context, queryResult, null, output);
+            } else {
+                Throwable[] statementWarnings = statement.getStatementWarnings();
+                if (statementWarnings != null && statementWarnings.length > 0) {
+                    dumpWarnings(output, Arrays.asList(statementWarnings));
+                }
+            }
+        }
+    }
+

plugins/org.jkiss.dbeaver.model/src/org/jkiss/dbeaver/model/impl/DefaultServerOutputReader.java

 /*
  * DBeaver - Universal Database Manager
- * Copyright (C) 2010-2017 Serge Rider (serge@jkiss.org)
+ * Copyright (C) 2010-2018 Serge Rider (serge@jkiss.org)
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -16,9 +16,12 @@
  */
 package org.jkiss.dbeaver.model.impl;
 
+import org.jkiss.code.NotNull;
+import org.jkiss.code.Nullable;
 import org.jkiss.dbeaver.model.exec.DBCException;
 import org.jkiss.dbeaver.model.exec.DBCExecutionContext;
 import org.jkiss.dbeaver.model.exec.DBCServerOutputReader;
+import org.jkiss.dbeaver.model.exec.DBCStatement;
 import org.jkiss.dbeaver.model.runtime.DBRProgressMonitor;
 import org.jkiss.dbeaver.model.sql.SQLQueryResult;
 import org.jkiss.utils.CommonUtils;
@@ -33,25 +36,22 @@ import java.util.List;
  */
 public class DefaultServerOutputReader implements DBCServerOutputReader
 {
-    private final SQLQueryResult queryResult;
-
-    public DefaultServerOutputReader(SQLQueryResult queryResult) {
-        this.queryResult = queryResult;
-    }
-
     @Override
     public boolean isServerOutputEnabled() {
         return true;
     }
 
     @Override
-    public void enableServerOutput(DBRProgressMonitor monitor, DBCExecutionContext context, boolean enable) throws DBCException {
-        // do nothing
+    public boolean isAsyncOutputReadSupported() {
+        return false;
     }
 
     @Override
-    public void readServerOutput(DBRProgressMonitor monitor, DBCExecutionContext context, PrintWriter output) throws DBCException {
-        List<Throwable> warnings = queryResult.getWarnings();
+    public void readServerOutput(@NotNull DBRProgressMonitor monitor, @NotNull DBCExecutionContext context, @Nullable SQLQueryResult queryResult, @Nullable DBCStatement statement, @NotNull PrintWriter output) throws DBCException {
+        dumpWarnings(output, queryResult.getWarnings());
+    }
+
+    protected void dumpWarnings(@NotNull PrintWriter output, List<Throwable> warnings) {
         if (warnings != null && warnings.size() > 0) {
             for (Throwable warning : warnings) {
                 if (warning instanceof SQLException) {