diff --git a/plugins/org.jkiss.dbeaver.model.ai/META-INF/MANIFEST.MF b/plugins/org.jkiss.dbeaver.model.ai/META-INF/MANIFEST.MF index 71ab52681ca8..a5c6d215e0de 100644 --- a/plugins/org.jkiss.dbeaver.model.ai/META-INF/MANIFEST.MF +++ b/plugins/org.jkiss.dbeaver.model.ai/META-INF/MANIFEST.MF @@ -18,6 +18,7 @@ Bundle-ActivationPolicy: lazy Bundle-ClassPath: . Automatic-Module-Name: org.jkiss.dbeaver.model.ai Require-Bundle: org.jkiss.dbeaver.model, + org.jkiss.dbeaver.model.sql, org.jkiss.dbeaver.registry, org.jkiss.bundle.gpt3, com.google.gson diff --git a/plugins/org.jkiss.dbeaver.model.ai/plugin.xml b/plugins/org.jkiss.dbeaver.model.ai/plugin.xml index 5087b3d8147a..1b869eb3fca3 100644 --- a/plugins/org.jkiss.dbeaver.model.ai/plugin.xml +++ b/plugins/org.jkiss.dbeaver.model.ai/plugin.xml @@ -20,4 +20,8 @@ + + + + diff --git a/plugins/org.jkiss.dbeaver.model.ai/src/org/jkiss/dbeaver/model/ai/AITextUtils.java b/plugins/org.jkiss.dbeaver.model.ai/src/org/jkiss/dbeaver/model/ai/AITextUtils.java index b2c365bb22d6..4ca35e7c73ad 100644 --- a/plugins/org.jkiss.dbeaver.model.ai/src/org/jkiss/dbeaver/model/ai/AITextUtils.java +++ b/plugins/org.jkiss.dbeaver.model.ai/src/org/jkiss/dbeaver/model/ai/AITextUtils.java @@ -18,16 +18,25 @@ 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.DBPDataSource; +import org.jkiss.dbeaver.model.DBUtils; import org.jkiss.dbeaver.model.ai.completion.DAICompletionMessage; +import org.jkiss.dbeaver.model.app.DBPProject; +import org.jkiss.dbeaver.model.runtime.DBRProgressMonitor; import org.jkiss.dbeaver.model.sql.SQLUtils; +import org.jkiss.dbeaver.model.struct.DBSEntity; import org.jkiss.dbeaver.runtime.DBWorkbench; import java.util.ArrayList; import java.util.List; +import java.util.Set; // All these ideally should be a part of a given AI engine public class AITextUtils { + private static final Log log = Log.getLog(AITextUtils.class); + private AITextUtils() { // prevents instantiation } @@ -106,4 +115,39 @@ public static MessageChunk[] splitIntoChunks(@NotNull String text) { return chunks.toArray(MessageChunk[]::new); } + + @NotNull + public static List loadCustomEntities( + @NotNull DBRProgressMonitor monitor, + @NotNull DBPDataSource dataSource, + @NotNull Set ids + ) { + monitor.beginTask("Load custom entities", ids.size()); + try { + return loadCheckedEntitiesById(monitor, dataSource.getContainer().getProject(), ids); + } catch (Exception e) { + log.error(e); + return List.of(); + } finally { + monitor.done(); + } + } + + @NotNull + private static List loadCheckedEntitiesById( + @NotNull DBRProgressMonitor monitor, + @NotNull DBPProject project, + @NotNull Set ids + ) throws DBException { + final List output = new ArrayList<>(); + + for (String id : ids) { + if (DBUtils.findObjectById(monitor, project, id) instanceof DBSEntity entity) { + output.add(entity); + } + monitor.worked(1); + } + + return output; + } } diff --git a/plugins/org.jkiss.dbeaver.model.ai/src/org/jkiss/dbeaver/model/ai/commands/AIOutputSeverity.java b/plugins/org.jkiss.dbeaver.model.ai/src/org/jkiss/dbeaver/model/ai/commands/AIOutputSeverity.java new file mode 100644 index 000000000000..31527c1e3574 --- /dev/null +++ b/plugins/org.jkiss.dbeaver.model.ai/src/org/jkiss/dbeaver/model/ai/commands/AIOutputSeverity.java @@ -0,0 +1,41 @@ +/* + * DBeaver - Universal Database Manager + * Copyright (C) 2010-2024 DBeaver Corp and others + * + * 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.ai.commands; + +import org.jkiss.code.NotNull; +import org.jkiss.dbeaver.model.exec.output.DBCOutputSeverity; + +enum AIOutputSeverity implements DBCOutputSeverity { + PROMPT("AI"); + + private final String name; + + AIOutputSeverity(@NotNull String name) { + this.name = name; + } + + @NotNull + @Override + public String getName() { + return name; + } + + @Override + public boolean isForced() { + return true; + } +} diff --git a/plugins/org.jkiss.dbeaver.model.ai/src/org/jkiss/dbeaver/model/ai/commands/SQLCommandAI.java b/plugins/org.jkiss.dbeaver.model.ai/src/org/jkiss/dbeaver/model/ai/commands/SQLCommandAI.java new file mode 100644 index 000000000000..955766390257 --- /dev/null +++ b/plugins/org.jkiss.dbeaver.model.ai/src/org/jkiss/dbeaver/model/ai/commands/SQLCommandAI.java @@ -0,0 +1,128 @@ +/* + * DBeaver - Universal Database Manager + * Copyright (C) 2010-2024 DBeaver Corp and others + * + * 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.ai.commands; + +import org.jkiss.code.NotNull; +import org.jkiss.dbeaver.DBException; +import org.jkiss.dbeaver.Log; +import org.jkiss.dbeaver.model.DBPDataSourceContainer; +import org.jkiss.dbeaver.model.ai.*; +import org.jkiss.dbeaver.model.ai.completion.*; +import org.jkiss.dbeaver.model.ai.format.IAIFormatter; +import org.jkiss.dbeaver.model.logical.DBSLogicalDataSource; +import org.jkiss.dbeaver.model.runtime.DBRProgressMonitor; +import org.jkiss.dbeaver.model.sql.*; +import org.jkiss.dbeaver.runtime.DBWorkbench; +import org.jkiss.utils.CommonUtils; + +import java.util.Arrays; +import java.util.List; +import java.util.stream.Collectors; + +/** + * Control command handler + */ +public class SQLCommandAI implements SQLControlCommandHandler { + + private static final Log log = Log.getLog(SQLCommandAI.class); + + @NotNull + @Override + public SQLControlResult handleCommand(@NotNull DBRProgressMonitor monitor, @NotNull SQLControlCommand command, @NotNull SQLScriptContext scriptContext) throws DBException { + if (command.getDataSource() == null) { + throw new DBException("Not connected to database"); + } + AISettings aiSettings = AISettingsRegistry.getInstance().getSettings(); + if (aiSettings.isAiDisabled()) { + throw new DBException("AI services are disabled"); + } + DAICompletionEngine engine = AIEngineRegistry.getInstance().getCompletionEngine( + aiSettings.getActiveEngine()); + + String prompt = command.getParameter(); + if (CommonUtils.isEmptyTrimmed(prompt)) { + throw new DBException("Empty AI prompt"); + } + + IAIFormatter formatter = AIFormatterRegistry.getInstance().getFormatter(AIConstants.CORE_FORMATTER); + + final DBSLogicalDataSource dataSource = new DBSLogicalDataSource( + command.getDataSourceContainer(), "AI logical wrapper", null); + + DBPDataSourceContainer dataSourceContainer = dataSource.getDataSourceContainer(); + DAICompletionSettings completionSettings = new DAICompletionSettings(dataSourceContainer); + if (!DBWorkbench.getPlatform().getApplication().isHeadlessMode() && !completionSettings.isMetaTransferConfirmed()) { + if (DBWorkbench.getPlatformUI().confirmAction("Do you confirm AI usage", + "Do you confirm AI usage for '" + dataSourceContainer.getName() + "'?" + )) { + completionSettings.setMetaTransferConfirmed(true); + completionSettings.saveSettings(); + } else { + throw new DBException("AI services restricted for '" + dataSourceContainer.getName() + "'"); + } + } + DAICompletionScope scope = completionSettings.getScope(); + DAICompletionContext.Builder contextBuilder = new DAICompletionContext.Builder() + .setScope(scope) + .setDataSource(dataSource) + .setExecutionContext(scriptContext.getExecutionContext()); + if (scope == DAICompletionScope.CUSTOM) { + contextBuilder.setCustomEntities( + AITextUtils.loadCustomEntities( + monitor, + command.getDataSource(), + Arrays.stream(completionSettings.getCustomObjectIds()).collect(Collectors.toSet())) + ); + } + final DAICompletionContext aiContext = contextBuilder.build(); + + DAICompletionSession aiSession = new DAICompletionSession(); + aiSession.add(new DAICompletionMessage(DAICompletionMessage.Role.USER, prompt)); + + List responses = engine.performSessionCompletion( + monitor, + aiContext, + aiSession, + formatter, + true); + + DAICompletionResponse response = responses.get(0); + MessageChunk[] messageChunks = AITextUtils.splitIntoChunks( + CommonUtils.notEmpty(response.getResultCompletion())); + + String finalSQL = null; + StringBuilder messages = new StringBuilder(); + for (MessageChunk chunk : messageChunks) { + if (chunk instanceof MessageChunk.Code code) { + finalSQL = code.text(); + } else if (chunk instanceof MessageChunk.Text text) { + messages.append(text.text()); + } + } + if (finalSQL == null) { + if (!messages.isEmpty()) { + throw new DBException(messages.toString()); + } + throw new DBException("Empty AI completion for '" + prompt + "'"); + } + + scriptContext.getOutputWriter().println(AIOutputSeverity.PROMPT, prompt + " ==> " + finalSQL + "\n"); + return SQLControlResult.transform( + new SQLQuery(command.getDataSource(), finalSQL)); + } + +} diff --git a/plugins/org.jkiss.dbeaver.model.sql/src/org/jkiss/dbeaver/model/sql/SQLControlCommandHandler.java b/plugins/org.jkiss.dbeaver.model.sql/src/org/jkiss/dbeaver/model/sql/SQLControlCommandHandler.java index 6b2f5b05652a..fd4bdf198060 100644 --- a/plugins/org.jkiss.dbeaver.model.sql/src/org/jkiss/dbeaver/model/sql/SQLControlCommandHandler.java +++ b/plugins/org.jkiss.dbeaver.model.sql/src/org/jkiss/dbeaver/model/sql/SQLControlCommandHandler.java @@ -16,7 +16,9 @@ */ package org.jkiss.dbeaver.model.sql; +import org.jkiss.code.NotNull; import org.jkiss.dbeaver.DBException; +import org.jkiss.dbeaver.model.runtime.DBRProgressMonitor; /** * Control command handler @@ -24,12 +26,16 @@ public interface SQLControlCommandHandler { /** - * + * @param monitor * @param command command * @param scriptContext script context * @return false if command failed and execution has to be stopped */ - boolean handleCommand(SQLControlCommand command, SQLScriptContext scriptContext) + @NotNull + SQLControlResult handleCommand( + @NotNull DBRProgressMonitor monitor, + @NotNull SQLControlCommand command, + @NotNull SQLScriptContext scriptContext) throws DBException; } diff --git a/plugins/org.jkiss.dbeaver.model.sql/src/org/jkiss/dbeaver/model/sql/SQLControlResult.java b/plugins/org.jkiss.dbeaver.model.sql/src/org/jkiss/dbeaver/model/sql/SQLControlResult.java new file mode 100644 index 000000000000..0df57fd7e7a8 --- /dev/null +++ b/plugins/org.jkiss.dbeaver.model.sql/src/org/jkiss/dbeaver/model/sql/SQLControlResult.java @@ -0,0 +1,48 @@ +/* + * DBeaver - Universal Database Manager + * Copyright (C) 2010-2024 DBeaver Corp and others + * + * 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.sql; + +/** + * Control command result. + * + * It may finish with no extra information or with parameters: + * - message: will be shown in UI + * - error: execution error will be shown in UI + */ +public class SQLControlResult { + + public static SQLControlResult success() { + return new SQLControlResult(); + } + + public static SQLControlResult transform(SQLScriptElement element) { + return new SQLControlResult(element); + } + + private SQLScriptElement transformed; + + private SQLControlResult() { + } + + private SQLControlResult(SQLScriptElement transformed) { + this.transformed = transformed; + } + + public SQLScriptElement getTransformed() { + return transformed; + } +} diff --git a/plugins/org.jkiss.dbeaver.model.sql/src/org/jkiss/dbeaver/model/sql/SQLScriptContext.java b/plugins/org.jkiss.dbeaver.model.sql/src/org/jkiss/dbeaver/model/sql/SQLScriptContext.java index 05c92fc5c1c3..16dd409a0bd6 100644 --- a/plugins/org.jkiss.dbeaver.model.sql/src/org/jkiss/dbeaver/model/sql/SQLScriptContext.java +++ b/plugins/org.jkiss.dbeaver.model.sql/src/org/jkiss/dbeaver/model/sql/SQLScriptContext.java @@ -32,6 +32,7 @@ import org.jkiss.dbeaver.model.exec.output.DBCOutputWriter; import org.jkiss.dbeaver.model.impl.OutputWriterAdapter; import org.jkiss.dbeaver.model.impl.sql.AbstractSQLDialect; +import org.jkiss.dbeaver.model.runtime.DBRProgressMonitor; import org.jkiss.dbeaver.model.sql.registry.SQLCommandHandlerDescriptor; import org.jkiss.dbeaver.model.sql.registry.SQLCommandsRegistry; import org.jkiss.dbeaver.model.sql.registry.SQLQueryParameterRegistry; @@ -267,15 +268,16 @@ public void setIgnoreParameters(boolean ignoreParameters) { this.ignoreParameters = ignoreParameters; } - public boolean executeControlCommand(SQLControlCommand command) throws DBException { + @NotNull + public SQLControlResult executeControlCommand(DBRProgressMonitor monitor, SQLControlCommand command) throws DBException { if (command.isEmptyCommand()) { - return true; + return SQLControlResult.success(); } SQLCommandHandlerDescriptor commandHandler = SQLCommandsRegistry.getInstance().getCommandHandler(command.getCommandId()); if (commandHandler == null) { throw new DBException("Command '" + command.getCommand() + "' not supported"); } - return commandHandler.createHandler().handleCommand(command, this); + return commandHandler.createHandler().handleCommand(monitor, command, this); } public void copyFrom(SQLScriptContext context) { diff --git a/plugins/org.jkiss.dbeaver.model.sql/src/org/jkiss/dbeaver/model/sql/commands/SQLCommandEcho.java b/plugins/org.jkiss.dbeaver.model.sql/src/org/jkiss/dbeaver/model/sql/commands/SQLCommandEcho.java index a9f2f4dbad82..655546798673 100644 --- a/plugins/org.jkiss.dbeaver.model.sql/src/org/jkiss/dbeaver/model/sql/commands/SQLCommandEcho.java +++ b/plugins/org.jkiss.dbeaver.model.sql/src/org/jkiss/dbeaver/model/sql/commands/SQLCommandEcho.java @@ -16,9 +16,12 @@ */ package org.jkiss.dbeaver.model.sql.commands; +import org.jkiss.code.NotNull; import org.jkiss.dbeaver.DBException; +import org.jkiss.dbeaver.model.runtime.DBRProgressMonitor; import org.jkiss.dbeaver.model.sql.SQLControlCommand; import org.jkiss.dbeaver.model.sql.SQLControlCommandHandler; +import org.jkiss.dbeaver.model.sql.SQLControlResult; import org.jkiss.dbeaver.model.sql.SQLScriptContext; import org.jkiss.dbeaver.model.sql.eval.ScriptVariablesResolver; import org.jkiss.dbeaver.utils.GeneralUtils; @@ -28,15 +31,16 @@ */ public class SQLCommandEcho implements SQLControlCommandHandler { + @NotNull @Override - public boolean handleCommand(SQLControlCommand command, SQLScriptContext scriptContext) throws DBException { + public SQLControlResult handleCommand(@NotNull DBRProgressMonitor monitor, @NotNull SQLControlCommand command, @NotNull SQLScriptContext scriptContext) throws DBException { String parameter = command.getParameter(); if (parameter != null) { parameter = GeneralUtils.replaceVariables(parameter, new ScriptVariablesResolver(scriptContext), true); } scriptContext.getOutputWriter().println(null, parameter); - return true; + return SQLControlResult.success(); } } diff --git a/plugins/org.jkiss.dbeaver.model.sql/src/org/jkiss/dbeaver/model/sql/commands/SQLCommandExport.java b/plugins/org.jkiss.dbeaver.model.sql/src/org/jkiss/dbeaver/model/sql/commands/SQLCommandExport.java index b5f513bd8f15..ce5ca44a1903 100644 --- a/plugins/org.jkiss.dbeaver.model.sql/src/org/jkiss/dbeaver/model/sql/commands/SQLCommandExport.java +++ b/plugins/org.jkiss.dbeaver.model.sql/src/org/jkiss/dbeaver/model/sql/commands/SQLCommandExport.java @@ -17,12 +17,11 @@ package org.jkiss.dbeaver.model.sql.commands; import com.google.gson.Gson; +import org.jkiss.code.NotNull; import org.jkiss.dbeaver.DBException; import org.jkiss.dbeaver.model.data.json.JSONUtils; -import org.jkiss.dbeaver.model.sql.SQLControlCommand; -import org.jkiss.dbeaver.model.sql.SQLControlCommandHandler; -import org.jkiss.dbeaver.model.sql.SQLPragmaHandler; -import org.jkiss.dbeaver.model.sql.SQLScriptContext; +import org.jkiss.dbeaver.model.runtime.DBRProgressMonitor; +import org.jkiss.dbeaver.model.sql.*; import java.io.StringReader; import java.util.Map; @@ -32,8 +31,9 @@ */ public class SQLCommandExport implements SQLControlCommandHandler { + @NotNull @Override - public boolean handleCommand(SQLControlCommand command, SQLScriptContext scriptContext) throws DBException { + public SQLControlResult handleCommand(@NotNull DBRProgressMonitor monitor, @NotNull SQLControlCommand command, @NotNull SQLScriptContext scriptContext) throws DBException { final Map params; try { @@ -44,6 +44,6 @@ public boolean handleCommand(SQLControlCommand command, SQLScriptContext scriptC scriptContext.setPragma(SQLPragmaHandler.PRAGMA_EXPORT, params); - return true; + return SQLControlResult.success(); } } diff --git a/plugins/org.jkiss.dbeaver.model.sql/src/org/jkiss/dbeaver/model/sql/commands/SQLCommandSet.java b/plugins/org.jkiss.dbeaver.model.sql/src/org/jkiss/dbeaver/model/sql/commands/SQLCommandSet.java index 63f92eff45bb..4911c16d8890 100644 --- a/plugins/org.jkiss.dbeaver.model.sql/src/org/jkiss/dbeaver/model/sql/commands/SQLCommandSet.java +++ b/plugins/org.jkiss.dbeaver.model.sql/src/org/jkiss/dbeaver/model/sql/commands/SQLCommandSet.java @@ -19,10 +19,8 @@ import org.jkiss.code.NotNull; import org.jkiss.dbeaver.DBException; import org.jkiss.dbeaver.model.exec.DBCException; -import org.jkiss.dbeaver.model.sql.SQLControlCommand; -import org.jkiss.dbeaver.model.sql.SQLControlCommandHandler; -import org.jkiss.dbeaver.model.sql.SQLDialect; -import org.jkiss.dbeaver.model.sql.SQLScriptContext; +import org.jkiss.dbeaver.model.runtime.DBRProgressMonitor; +import org.jkiss.dbeaver.model.sql.*; import org.jkiss.dbeaver.model.sql.parser.rules.ScriptParameterRule; import org.jkiss.dbeaver.utils.GeneralUtils; import org.jkiss.utils.CommonUtils; @@ -32,8 +30,9 @@ */ public class SQLCommandSet implements SQLControlCommandHandler { + @NotNull @Override - public boolean handleCommand(SQLControlCommand command, SQLScriptContext scriptContext) throws DBException { + public SQLControlResult handleCommand(@NotNull DBRProgressMonitor monitor, @NotNull SQLControlCommand command, @NotNull SQLScriptContext scriptContext) throws DBException { SQLDialect sqlDialect = scriptContext.getExecutionContext().getDataSource().getSQLDialect(); String parameter = command.getParameter().stripLeading(); int varNameEnd = ScriptParameterRule.tryConsumeParameterName(sqlDialect, parameter, 0); @@ -46,7 +45,7 @@ public boolean handleCommand(SQLControlCommand command, SQLScriptContext scriptC throw new DBCException("Bad set syntax. Expected syntax:\n@set varName = value or expression"); } String shouldBeEmpty = parameter.substring(varNameEnd, divPos).trim(); - if (shouldBeEmpty.length() > 0) { + if (!shouldBeEmpty.isEmpty()) { throw new DBCException( "Unexpected characters " + shouldBeEmpty + " after the variable name " + varName + ". " + "Expected syntax:\n@set varName = value or expression" @@ -56,7 +55,7 @@ public boolean handleCommand(SQLControlCommand command, SQLScriptContext scriptC varValue = GeneralUtils.replaceVariables(varValue, name -> CommonUtils.toString(scriptContext.getVariable(name)), true); scriptContext.setVariable(varName, varValue); - return true; + return SQLControlResult.success(); } /* diff --git a/plugins/org.jkiss.dbeaver.model.sql/src/org/jkiss/dbeaver/model/sql/commands/SQLCommandUnset.java b/plugins/org.jkiss.dbeaver.model.sql/src/org/jkiss/dbeaver/model/sql/commands/SQLCommandUnset.java index 91b8536ca57c..0b7542e291f2 100644 --- a/plugins/org.jkiss.dbeaver.model.sql/src/org/jkiss/dbeaver/model/sql/commands/SQLCommandUnset.java +++ b/plugins/org.jkiss.dbeaver.model.sql/src/org/jkiss/dbeaver/model/sql/commands/SQLCommandUnset.java @@ -16,12 +16,11 @@ */ package org.jkiss.dbeaver.model.sql.commands; +import org.jkiss.code.NotNull; import org.jkiss.dbeaver.DBException; import org.jkiss.dbeaver.model.exec.DBCException; -import org.jkiss.dbeaver.model.sql.SQLControlCommand; -import org.jkiss.dbeaver.model.sql.SQLControlCommandHandler; -import org.jkiss.dbeaver.model.sql.SQLDialect; -import org.jkiss.dbeaver.model.sql.SQLScriptContext; +import org.jkiss.dbeaver.model.runtime.DBRProgressMonitor; +import org.jkiss.dbeaver.model.sql.*; import org.jkiss.dbeaver.model.sql.parser.rules.ScriptParameterRule; /** @@ -29,8 +28,9 @@ */ public class SQLCommandUnset implements SQLControlCommandHandler { + @NotNull @Override - public boolean handleCommand(SQLControlCommand command, SQLScriptContext scriptContext) throws DBException { + public SQLControlResult handleCommand(@NotNull DBRProgressMonitor monitor, @NotNull SQLControlCommand command, @NotNull SQLScriptContext scriptContext) throws DBException { SQLDialect sqlDialect = scriptContext.getExecutionContext().getDataSource().getSQLDialect(); String parameter = command.getParameter().trim(); @@ -42,7 +42,7 @@ public boolean handleCommand(SQLControlCommand command, SQLScriptContext scriptC String varName = SQLCommandSet.prepareVarName(sqlDialect, parameter.substring(0, varNameEnd)); scriptContext.removeVariable(varName); - return true; + return SQLControlResult.success(); } } diff --git a/plugins/org.jkiss.dbeaver.model.sql/src/org/jkiss/dbeaver/model/sql/exec/SQLScriptProcessor.java b/plugins/org.jkiss.dbeaver.model.sql/src/org/jkiss/dbeaver/model/sql/exec/SQLScriptProcessor.java index 4f22204faf53..15004c7c2f56 100644 --- a/plugins/org.jkiss.dbeaver.model.sql/src/org/jkiss/dbeaver/model/sql/exec/SQLScriptProcessor.java +++ b/plugins/org.jkiss.dbeaver.model.sql/src/org/jkiss/dbeaver/model/sql/exec/SQLScriptProcessor.java @@ -170,10 +170,15 @@ public void runScript(DBRProgressMonitor monitor) throws DBCException { } private boolean executeSingleQuery(@NotNull DBCSession session, @NotNull SQLScriptElement element) { - if (element instanceof SQLControlCommand) { + if (element instanceof SQLControlCommand controlCommand) { log.debug(STAT_LOG_PREFIX + "Execute command\n" + element.getText()); try { - return scriptContext.executeControlCommand((SQLControlCommand) element); + SQLControlResult controlResult = scriptContext.executeControlCommand(session.getProgressMonitor(), controlCommand); + if (controlResult.getTransformed() != null) { + element = controlResult.getTransformed(); + } else { + return true; + } } catch (Throwable e) { if (!(e instanceof DBException)) { log.error("Unexpected error while processing SQL command", e); @@ -182,7 +187,10 @@ private boolean executeSingleQuery(@NotNull DBCSession session, @NotNull SQLScri return false; } } - SQLQuery sqlQuery = (SQLQuery) element; + if (!(element instanceof SQLQuery sqlQuery)) { + log.error("Unsupported SQL element type: " + element); + return false; + } scriptContext.fillQueryParameters(sqlQuery, () -> dataReceiver, true); lastError = null; diff --git a/plugins/org.jkiss.dbeaver.model/src/org/jkiss/dbeaver/model/exec/output/DBCOutputSeverity.java b/plugins/org.jkiss.dbeaver.model/src/org/jkiss/dbeaver/model/exec/output/DBCOutputSeverity.java index 5891a10a2cbe..568d7eec9b99 100644 --- a/plugins/org.jkiss.dbeaver.model/src/org/jkiss/dbeaver/model/exec/output/DBCOutputSeverity.java +++ b/plugins/org.jkiss.dbeaver.model/src/org/jkiss/dbeaver/model/exec/output/DBCOutputSeverity.java @@ -19,4 +19,7 @@ import org.jkiss.dbeaver.model.DBPNamedObject; public interface DBCOutputSeverity extends DBPNamedObject { + default boolean isForced() { + return false; + } } diff --git a/plugins/org.jkiss.dbeaver.model/src/org/jkiss/dbeaver/model/sql/SQLControlCommand.java b/plugins/org.jkiss.dbeaver.model/src/org/jkiss/dbeaver/model/sql/SQLControlCommand.java index aedd800e1f00..574748ef3a7c 100644 --- a/plugins/org.jkiss.dbeaver.model/src/org/jkiss/dbeaver/model/sql/SQLControlCommand.java +++ b/plugins/org.jkiss.dbeaver.model/src/org/jkiss/dbeaver/model/sql/SQLControlCommand.java @@ -19,6 +19,7 @@ import org.jkiss.code.NotNull; import org.jkiss.dbeaver.model.DBPDataSource; +import org.jkiss.dbeaver.model.DBPDataSourceContainer; /** * SQL control command @@ -59,6 +60,10 @@ public SQLControlCommand(DBPDataSource dataSource, SQLSyntaxManager syntaxManage this.commandId = commandId == null ? command : commandId; } + public DBPDataSourceContainer getDataSourceContainer() { + return dataSource == null ? null : dataSource.getContainer(); + } + public DBPDataSource getDataSource() { return dataSource; } @@ -120,4 +125,5 @@ public void reset() { public String toString() { return text; } + } diff --git a/plugins/org.jkiss.dbeaver.ui.editors.data/src/org/jkiss/dbeaver/ui/controls/resultset/panel/grouping/GroupingResultsContainer.java b/plugins/org.jkiss.dbeaver.ui.editors.data/src/org/jkiss/dbeaver/ui/controls/resultset/panel/grouping/GroupingResultsContainer.java index 50f7a733ac6d..d93b664d8e92 100644 --- a/plugins/org.jkiss.dbeaver.ui.editors.data/src/org/jkiss/dbeaver/ui/controls/resultset/panel/grouping/GroupingResultsContainer.java +++ b/plugins/org.jkiss.dbeaver.ui.editors.data/src/org/jkiss/dbeaver/ui/controls/resultset/panel/grouping/GroupingResultsContainer.java @@ -153,20 +153,26 @@ boolean removeGroupingAttribute(List attributes) { public void addGroupingFunctions(List functions) { for (String func : functions) { - func = DBUtils.getUnQuotedIdentifier(getDataContainer().getDataSource(), func); - if (!groupFunctions.contains(func)) { - groupFunctions.add(func); + DBPDataSource dataSource = getDataContainer().getDataSource(); + if (dataSource != null) { + func = DBUtils.getUnQuotedIdentifier(dataSource, func); + if (!groupFunctions.contains(func)) { + groupFunctions.add(func); + } } } } public boolean removeGroupingFunction(List attributes) { boolean changed = false; - for (String func : attributes) { - func = DBUtils.getUnQuotedIdentifier(getDataContainer().getDataSource(), func); - if (groupFunctions.contains(func)) { - groupFunctions.remove(func); - changed = true; + DBPDataSource dataSource = getDataContainer().getDataSource(); + if (dataSource != null) { + for (String func : attributes) { + func = DBUtils.getUnQuotedIdentifier(dataSource, func); + if (groupFunctions.contains(func)) { + groupFunctions.remove(func); + changed = true; + } } } return changed; diff --git a/plugins/org.jkiss.dbeaver.ui.editors.sql.ai/src/org/jkiss/dbeaver/ui/editors/sql/ai/controls/ScopeSelectorControl.java b/plugins/org.jkiss.dbeaver.ui.editors.sql.ai/src/org/jkiss/dbeaver/ui/editors/sql/ai/controls/ScopeSelectorControl.java index 5ebe9a27ec9a..e272b8947922 100644 --- a/plugins/org.jkiss.dbeaver.ui.editors.sql.ai/src/org/jkiss/dbeaver/ui/editors/sql/ai/controls/ScopeSelectorControl.java +++ b/plugins/org.jkiss.dbeaver.ui.editors.sql.ai/src/org/jkiss/dbeaver/ui/editors/sql/ai/controls/ScopeSelectorControl.java @@ -25,13 +25,12 @@ import org.eclipse.swt.widgets.*; 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.DBPDataSource; import org.jkiss.dbeaver.model.DBUtils; +import org.jkiss.dbeaver.model.ai.AITextUtils; import org.jkiss.dbeaver.model.ai.completion.DAICompletionScope; import org.jkiss.dbeaver.model.ai.completion.DAICompletionSettings; -import org.jkiss.dbeaver.model.app.DBPProject; import org.jkiss.dbeaver.model.exec.DBCExecutionContext; import org.jkiss.dbeaver.model.logical.DBSLogicalDataSource; import org.jkiss.dbeaver.model.navigator.DBNDatabaseNode; @@ -151,7 +150,7 @@ public DAICompletionScope getScope() { @NotNull public List getCustomEntities(@NotNull DBRProgressMonitor monitor) { - return loadCustomEntities(monitor, executionContext.getDataSource(), checkedObjectIds); + return AITextUtils.loadCustomEntities(monitor, executionContext.getDataSource(), checkedObjectIds); } @NotNull @@ -206,7 +205,7 @@ public static Set chooseCustomEntities( try { // Find nodes of already selected objects context.run(true, true, monitor -> { - for (DBSEntity entity : loadCustomEntities(monitor, dataSource, ids)) { + for (DBSEntity entity : AITextUtils.loadCustomEntities(monitor, dataSource, ids)) { DBNDatabaseNode node = navigator.getNodeByObject(monitor, entity, true); if (node != null) { nodes.add(node); @@ -240,41 +239,6 @@ public static Set chooseCustomEntities( .collect(Collectors.toSet()); } - @NotNull - public static List loadCustomEntities( - @NotNull DBRProgressMonitor monitor, - @NotNull DBPDataSource dataSource, - @NotNull Set ids - ) { - monitor.beginTask("Load custom entities", ids.size()); - try { - return loadCheckedEntitiesById(monitor, dataSource.getContainer().getProject(), ids); - } catch (Exception e) { - log.error(e); - return List.of(); - } finally { - monitor.done(); - } - } - - @NotNull - private static List loadCheckedEntitiesById( - @NotNull DBRProgressMonitor monitor, - @NotNull DBPProject project, - @NotNull Set ids - ) throws DBException { - final List output = new ArrayList<>(); - - for (String id : ids) { - if (DBUtils.findObjectById(monitor, project, id) instanceof DBSEntity entity) { - output.add(entity); - } - monitor.worked(1); - } - - return output; - } - public void changeScope(@NotNull DAICompletionScope scope) { if (scope == DAICompletionScope.CUSTOM) { Set ids = chooseCustomEntities( diff --git a/plugins/org.jkiss.dbeaver.ui.editors.sql/src/org/jkiss/dbeaver/ui/editors/sql/SQLEditor.java b/plugins/org.jkiss.dbeaver.ui.editors.sql/src/org/jkiss/dbeaver/ui/editors/sql/SQLEditor.java index 36165a477aac..4368015cb7e1 100644 --- a/plugins/org.jkiss.dbeaver.ui.editors.sql/src/org/jkiss/dbeaver/ui/editors/sql/SQLEditor.java +++ b/plugins/org.jkiss.dbeaver.ui.editors.sql/src/org/jkiss/dbeaver/ui/editors/sql/SQLEditor.java @@ -87,10 +87,7 @@ import org.jkiss.dbeaver.model.qm.QMTransactionState; import org.jkiss.dbeaver.model.qm.QMUtils; import org.jkiss.dbeaver.model.rm.RMConstants; -import org.jkiss.dbeaver.model.runtime.AbstractJob; -import org.jkiss.dbeaver.model.runtime.DBRProgressListener; -import org.jkiss.dbeaver.model.runtime.DBRProgressMonitor; -import org.jkiss.dbeaver.model.runtime.DBRRunnableWithProgress; +import org.jkiss.dbeaver.model.runtime.*; import org.jkiss.dbeaver.model.sql.*; import org.jkiss.dbeaver.model.sql.data.SQLQueryDataContainer; import org.jkiss.dbeaver.model.sql.transformers.SQLQueryTransformerCount; @@ -3769,8 +3766,15 @@ public void cancelJob() { } } - boolean processQueries(SQLScriptContext scriptContext, final List queries, boolean forceScript, final boolean fetchResults, boolean export, boolean closeTabOnError, SQLQueryListener queryListener) - { + boolean processQueries( + SQLScriptContext scriptContext, + final List queries, + boolean forceScript, + final boolean fetchResults, + boolean export, + boolean closeTabOnError, + SQLQueryListener queryListener + ) { if (queries.isEmpty()) { // Nothing to process return false; @@ -3805,12 +3809,15 @@ boolean processQueries(SQLScriptContext scriptContext, final List null, false); SQLQueryDataContainer dataContainer = new SQLQueryDataContainer(SQLEditor.this, query, scriptContext, log); diff --git a/plugins/org.jkiss.dbeaver.ui.editors.sql/src/org/jkiss/dbeaver/ui/editors/sql/SQLEditorOutputViewer.java b/plugins/org.jkiss.dbeaver.ui.editors.sql/src/org/jkiss/dbeaver/ui/editors/sql/SQLEditorOutputViewer.java index 9a621989845a..519bdba3e6f2 100644 --- a/plugins/org.jkiss.dbeaver.ui.editors.sql/src/org/jkiss/dbeaver/ui/editors/sql/SQLEditorOutputViewer.java +++ b/plugins/org.jkiss.dbeaver.ui.editors.sql/src/org/jkiss/dbeaver/ui/editors/sql/SQLEditorOutputViewer.java @@ -95,8 +95,12 @@ public void println(@Nullable DBCOutputSeverity severity, @Nullable String messa if (message == null) { return; } - if (severity == null || severities.contains(severity)) { - viewer.getOutputWriter().println(message); + if (severity == null || severity.isForced() || severities.contains(severity)) { + PrintWriter writer = viewer.getOutputWriter(); + if (severity != null && severity.isForced()) { + writer.print("[" + severity.getName() + "] "); + } + writer.println(message); } records.offer(new OutputRecord(severity, message)); if (records.size() > MAX_RECORDS) { @@ -145,8 +149,7 @@ private void updateControls() { final DBPDataSource dataSource = executionContext != null ? executionContext.getDataSource() : null; final DBCServerOutputReader reader = DBUtils.getAdapter(DBCServerOutputReader.class, dataSource); - if (reader instanceof DBCServerOutputReaderExt) { - final DBCServerOutputReaderExt readerExt = (DBCServerOutputReaderExt) reader; + if (reader instanceof DBCServerOutputReaderExt readerExt) { final DBCOutputSeverity[] supportedSeverities = readerExt.getSupportedSeverities(executionContext); severities.addAll(List.of(supportedSeverities)); @@ -166,7 +169,7 @@ private void filterOutput() { final String filter = filterText.getText().trim(); for (OutputRecord record : records) { - if (record.severity != null && !severities.contains(record.severity)) { + if (record.severity != null && !record.severity.isForced() && !severities.contains(record.severity)) { continue; } if (!filter.isEmpty() && !record.line.contains(filter)) { @@ -186,10 +189,10 @@ public ConfigureSeverityAction() { filterMenu.setRemoveAllWhenShown(true); filterMenu.addMenuListener(manager -> { final DBCServerOutputReader reader = DBUtils.getAdapter(DBCServerOutputReader.class, executionContext.getDataSource()); - if (!(reader instanceof DBCServerOutputReaderExt)) { + if (!(reader instanceof DBCServerOutputReaderExt readerExt)) { return; } - for (DBCOutputSeverity severity : ((DBCServerOutputReaderExt) reader).getSupportedSeverities(executionContext)) { + for (DBCOutputSeverity severity : readerExt.getSupportedSeverities(executionContext)) { manager.add(new ToggleSeverityAction(severity)); } }); diff --git a/plugins/org.jkiss.dbeaver.ui.editors.sql/src/org/jkiss/dbeaver/ui/editors/sql/commands/SQLCommandInclude.java b/plugins/org.jkiss.dbeaver.ui.editors.sql/src/org/jkiss/dbeaver/ui/editors/sql/commands/SQLCommandInclude.java index 4ccc8bb78180..0026d09b0977 100644 --- a/plugins/org.jkiss.dbeaver.ui.editors.sql/src/org/jkiss/dbeaver/ui/editors/sql/commands/SQLCommandInclude.java +++ b/plugins/org.jkiss.dbeaver.ui.editors.sql/src/org/jkiss/dbeaver/ui/editors/sql/commands/SQLCommandInclude.java @@ -18,11 +18,13 @@ import org.eclipse.ui.*; import org.eclipse.ui.ide.IDEEncoding; +import org.jkiss.code.NotNull; import org.jkiss.dbeaver.DBException; import org.jkiss.dbeaver.Log; import org.jkiss.dbeaver.model.DBUtils; import org.jkiss.dbeaver.model.exec.DBCSession; import org.jkiss.dbeaver.model.exec.DBCStatistics; +import org.jkiss.dbeaver.model.runtime.DBRProgressMonitor; import org.jkiss.dbeaver.model.sql.*; import org.jkiss.dbeaver.model.sql.eval.ScriptVariablesResolver; import org.jkiss.dbeaver.ui.UIUtils; @@ -55,8 +57,9 @@ public static String getResourceEncoding() { return CommonUtils.isEmpty(resourceEncoding) ? GeneralUtils.getDefaultFileEncoding() : resourceEncoding; } + @NotNull @Override - public boolean handleCommand(SQLControlCommand command, final SQLScriptContext scriptContext) throws DBException { + public SQLControlResult handleCommand(@NotNull DBRProgressMonitor monitor, @NotNull SQLControlCommand command, @NotNull final SQLScriptContext scriptContext) throws DBException { String fileName = command.getParameter(); if (CommonUtils.isEmpty(fileName)) { throw new DBException("Empty input file"); @@ -132,7 +135,7 @@ public boolean handleCommand(SQLControlCommand command, final SQLScriptContext s } } - return true; + return SQLControlResult.success(); } private static class IncludeScriptListener implements SQLQueryListener { diff --git a/plugins/org.jkiss.dbeaver.ui.editors.sql/src/org/jkiss/dbeaver/ui/editors/sql/execute/SQLQueryJob.java b/plugins/org.jkiss.dbeaver.ui.editors.sql/src/org/jkiss/dbeaver/ui/editors/sql/execute/SQLQueryJob.java index 8957492ea5ba..40d3fe855986 100644 --- a/plugins/org.jkiss.dbeaver.ui.editors.sql/src/org/jkiss/dbeaver/ui/editors/sql/execute/SQLQueryJob.java +++ b/plugins/org.jkiss.dbeaver.ui.editors.sql/src/org/jkiss/dbeaver/ui/editors/sql/execute/SQLQueryJob.java @@ -239,6 +239,7 @@ protected IStatus run(DBRProgressMonitor monitor) fetchResultSetNumber = resultSetNumber; boolean runNext = executeSingleQuery(session, query, true); + if (txnManager != null && txnManager.isSupportsTransactions() && !oldAutoCommit && commitType != SQLScriptCommitType.AUTOCOMMIT && query instanceof SQLQuery sqlQuery @@ -372,11 +373,13 @@ protected void handleTransactionStatements( } } - private boolean executeSingleQuery(@NotNull DBCSession session, @NotNull SQLScriptElement element, final boolean fireEvents) - { - - if (!scriptContext.getPragmas().isEmpty() && element instanceof SQLQuery) { - final SQLQueryDataContainer container = new SQLQueryDataContainer(this::getExecutionContext, (SQLQuery) element, scriptContext, log); + private boolean executeSingleQuery( + @NotNull DBCSession session, + @NotNull SQLScriptElement element, + final boolean fireEvents + ) { + if (!scriptContext.getPragmas().isEmpty() && element instanceof SQLQuery query) { + final SQLQueryDataContainer container = new SQLQueryDataContainer(this::getExecutionContext, query, scriptContext, log); for (var it = scriptContext.getPragmas().entrySet().iterator(); it.hasNext(); ) { final Map.Entry> entry = it.next(); @@ -403,9 +406,14 @@ private boolean executeSingleQuery(@NotNull DBCSession session, @NotNull SQLScri } } } - if (element instanceof SQLControlCommand) { + if (element instanceof SQLControlCommand controlCommand) { try { - return scriptContext.executeControlCommand((SQLControlCommand)element); + SQLControlResult controlResult = scriptContext.executeControlCommand(session.getProgressMonitor(), controlCommand); + if (controlResult.getTransformed() != null) { + element = controlResult.getTransformed(); + } else { + return true; + } } catch (Throwable e) { if (!(e instanceof DBException)) { log.error("Unexpected error while processing SQL command", e); @@ -413,11 +421,16 @@ private boolean executeSingleQuery(@NotNull DBCSession session, @NotNull SQLScri lastError = e; return false; } finally { - statistics.addStatementsCount(); - statistics.addMessage("Command " + ((SQLControlCommand) element).getCommand() + " processed"); + if (element instanceof SQLControlCommand finalCommand) { + statistics.addStatementsCount(); + statistics.addMessage("Command " + finalCommand.getCommand() + " processed"); + } } } - SQLQuery sqlQuery = (SQLQuery) element; + if (!(element instanceof SQLQuery sqlQuery)) { + log.error("Unsupported SQL element type: " + element); + return false; + } lastError = null; if (!skipConfirmation && getDataSourceContainer().getConnectionConfiguration().getConnectionType().isConfirmExecute()) {