dbeaver/pro#3911 AI and SQL control commands (#3166)

* dbeaver/pro#3911 AI and SQL control commands * dbeaver/pro#3911 AI and SQL control commands --------- Co-authored-by: kseniaguzeeva <112612526+kseniaguzeeva@users.noreply.github.com>
dbeaver · Dec 30, 2024 · 827865a · 827865a
1 parent 70b6fc2
commit 827865a
Showing 1 changed file with 9 additions and 4 deletions.
diff --git a/server/bundles/io.cloudbeaver.server/src/io/cloudbeaver/service/sql/WebSQLProcessor.java b/server/bundles/io.cloudbeaver.server/src/io/cloudbeaver/service/sql/WebSQLProcessor.java
@@ -200,10 +200,15 @@ public WebSQLExecuteInfo processQuery(
             SQLScriptElement element = SQLScriptParser.extractActiveQuery(parserContext, 0, sql.length());
 
             if (element instanceof SQLControlCommand command) {
-                dataContainer.getScriptContext().executeControlCommand(command);
-                WebSQLQueryResults stats = new WebSQLQueryResults(webSession, dataFormat);
-                executeInfo.setResults(new WebSQLQueryResults[]{stats});
-            } else if (element instanceof SQLQuery sqlQuery) {
+                SQLControlResult controlResult = dataContainer.getScriptContext().executeControlCommand(monitor, command);
+                if (controlResult.getTransformed() != null) {
+                    element = controlResult.getTransformed();
+                } else {
+                    WebSQLQueryResults stats = new WebSQLQueryResults(webSession, dataFormat);
+                    executeInfo.setResults(new WebSQLQueryResults[]{stats});
+                }
+            }
+            if (element instanceof SQLQuery sqlQuery) {
                 DBExecUtils.tryExecuteRecover(monitor, connection.getDataSource(), param -> {
                     try (DBCSession session = context.openSession(monitor, resolveQueryPurpose(dataFilter), "Execute SQL")) {
                         AbstractExecutionSource source = new AbstractExecutionSource(