dbeaver/pro#3911 AI and SQL control commands

dbeaver · Dec 29, 2024 · af3895a · af3895a
1 parent e420a77
commit af3895a
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(webSession.getProgressMonitor(), 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(