forked from Limework/skript-db
367 lines
15 KiB
Java
367 lines
15 KiB
Java
package com.btk5h.skriptdb.skript;
|
|
|
|
import ch.njol.skript.Skript;
|
|
import ch.njol.skript.effects.Delay;
|
|
import ch.njol.skript.lang.*;
|
|
import ch.njol.skript.variables.Variables;
|
|
import ch.njol.util.Kleenean;
|
|
import ch.njol.util.Pair;
|
|
import com.btk5h.skriptdb.SkriptDB;
|
|
import com.btk5h.skriptdb.SkriptUtil;
|
|
import com.zaxxer.hikari.HikariDataSource;
|
|
import org.bukkit.Bukkit;
|
|
import org.bukkit.event.Event;
|
|
|
|
import javax.sql.DataSource;
|
|
import javax.sql.rowset.CachedRowSet;
|
|
import javax.sql.rowset.serial.SerialBlob;
|
|
import javax.sql.rowset.serial.SerialException;
|
|
import java.io.IOException;
|
|
import java.sql.Connection;
|
|
import java.sql.PreparedStatement;
|
|
import java.sql.ResultSetMetaData;
|
|
import java.sql.SQLException;
|
|
import java.sql.Statement;
|
|
import java.util.*;
|
|
import java.util.concurrent.CompletableFuture;
|
|
import java.util.concurrent.CompletionException;
|
|
import java.util.concurrent.ExecutorService;
|
|
import java.util.concurrent.Executors;
|
|
import java.util.regex.Pattern;
|
|
|
|
/**
|
|
* Executes a statement on a database and optionally stores the result in a variable. Expressions
|
|
* embedded in the query will be escaped to avoid SQL injection.
|
|
* <p>
|
|
* If a single variable, such as `{test}`, is passed, the variable will be set to the number of
|
|
* affected rows.
|
|
* <p>
|
|
* If a list variable, such as `{test::*}`, is passed, the query result will be mapped to the list
|
|
* variable in the form `{test::<column name>::<row number>}`
|
|
*
|
|
* @name Execute Statement
|
|
* @pattern [synchronously] execute %string% (in|on) %datasource% [and store [[the] (output|result)[s]] (to|in)
|
|
* [the] [var[iable]] %-objects%]
|
|
* @example execute "select * from table" in {sql} and store the result in {output::*}
|
|
* @example execute "select * from %{table variable}%" in {sql} and store the result in {output::*}
|
|
* @since 0.1.0
|
|
*/
|
|
public class EffExecuteStatement extends Effect {
|
|
private static final ExecutorService threadPool = Executors.newFixedThreadPool(SkriptDB.getInstance().getConfig().getInt("thread-pool-size", 10));
|
|
private static final Pattern ARGUMENT_PLACEHOLDER = Pattern.compile("(?<!\\\\)\\?");
|
|
static String lastError;
|
|
|
|
static {
|
|
Skript.registerEffect(EffExecuteStatement.class,
|
|
"[quickly:quickly] execute %string% (in|on) %datasource% " +
|
|
"[with arg[ument][s] %-objects%] [and store [[the] [keys:generated keys] (output|result)[s]] (to|in) [the] [var[iable]] %-objects%]");
|
|
}
|
|
|
|
private Expression<String> query;
|
|
private Expression<HikariDataSource> dataSource;
|
|
private Expression<Object> queryArguments;
|
|
private VariableString resultVariableName;
|
|
private boolean isLocal;
|
|
private boolean isList;
|
|
private boolean quickly;
|
|
private boolean generatedKeys;
|
|
private boolean isSync = false;
|
|
|
|
@Override
|
|
protected void execute(Event e) {
|
|
DataSource ds = dataSource.getSingle(e);
|
|
//if data source isn't set
|
|
if (ds == null) {
|
|
return;
|
|
}
|
|
Pair<String, List<Object>> parsedQuery = parseQuery(e);
|
|
String baseVariable = resultVariableName != null ? resultVariableName.toString(e).toLowerCase(Locale.ENGLISH) : null;
|
|
|
|
Object locals = Variables.removeLocals(e);
|
|
|
|
//execute SQL statement
|
|
if (Bukkit.isPrimaryThread()) {
|
|
CompletableFuture.supplyAsync(() -> executeStatement(ds, baseVariable, parsedQuery), threadPool)
|
|
.whenComplete((resources, err) -> {
|
|
//handle last error syntax data
|
|
resetLastSQLError();
|
|
if (err instanceof CompletionException && err.getCause() instanceof SkriptDBQueryException) {
|
|
setLastSQLError(err.getCause().getMessage());
|
|
}
|
|
//if local variables are present
|
|
//bring back local variables
|
|
//populate SQL data into variables
|
|
if (!quickly) {
|
|
Bukkit.getScheduler().runTask(SkriptDB.getInstance(),
|
|
() -> postExecution(e, locals, resources));
|
|
} else {
|
|
postExecution(e, locals, resources);
|
|
}
|
|
});
|
|
// sync executed SQL query, same as above, just sync
|
|
} else {
|
|
isSync = true;
|
|
Map<String, Object> resources = null;
|
|
resetLastSQLError();
|
|
try {
|
|
resources = executeStatement(ds, baseVariable, parsedQuery);
|
|
} catch (SkriptDBQueryException err) {
|
|
//handle last error syntax data
|
|
setLastSQLError(err.getMessage());
|
|
}
|
|
//if local variables are present
|
|
//bring back local variables
|
|
//populate SQL data into variables
|
|
postExecution(e, locals, resources);
|
|
}
|
|
}
|
|
|
|
private void postExecution(Event e, Object locals, Map<String, Object> resources) {
|
|
if (locals != null && getNext() != null) {
|
|
Variables.setLocalVariables(e, locals);
|
|
}
|
|
if (resources != null) {
|
|
resources.forEach((name, value) -> setVariable(e, name, value));
|
|
}
|
|
TriggerItem.walk(getNext(), e);
|
|
//the line below is required to prevent memory leaks
|
|
Variables.removeLocals(e);
|
|
}
|
|
|
|
@Override
|
|
protected TriggerItem walk(Event e) {
|
|
debug(e, true);
|
|
if (!quickly || !isSync) {
|
|
Delay.addDelayedEvent(e);
|
|
}
|
|
execute(e);
|
|
return null;
|
|
}
|
|
|
|
private Pair<String, List<Object>> parseQuery(Event e) {
|
|
if (queryArguments != null) {
|
|
Object[] args = queryArguments.getArray(e);
|
|
String queryString = query.getSingle(e);
|
|
int queryArgCount = (int) ARGUMENT_PLACEHOLDER.matcher(queryString).results().count();
|
|
if (queryArgCount != args.length) {
|
|
Skript.warning(String.format("Your query has %d question marks, but you provided %d arguments. (%s) [%s]",
|
|
queryArgCount,
|
|
args.length,
|
|
queryArguments.toString(e, true),
|
|
Optional.ofNullable(getTrigger())
|
|
.map(Trigger::getDebugLabel)
|
|
.orElse("unknown")));
|
|
args = Arrays.copyOf(args, queryArgCount);
|
|
}
|
|
return new Pair<>(query.getSingle(e), Arrays.asList(args));
|
|
} else if (query instanceof VariableString && !((VariableString) query).isSimple()) {
|
|
return parseVariableQuery(e, (VariableString) query);
|
|
}
|
|
return new Pair<>(query.getSingle(e), null);
|
|
}
|
|
|
|
private Pair<String, List<Object>> parseVariableQuery(Event e, VariableString varQuery) {
|
|
StringBuilder sb = new StringBuilder();
|
|
List<Object> parameters = new LinkedList<>();
|
|
Object[] objects = SkriptUtil.getTemplateString(varQuery);
|
|
|
|
for (int i = 0; i < objects.length; i++) {
|
|
if (objects[i] instanceof String) {
|
|
sb.append(objects[i]);
|
|
} else {
|
|
Expression<?> expr = objects[i] instanceof Expression ? (Expression<?>) objects[i] : SkriptUtil.getExpressionFromInfo(objects[i]);
|
|
boolean standaloneString = isStandaloneString(objects, i);
|
|
Object expressionValue = expr.getSingle(e);
|
|
|
|
Pair<String, Object> toAppend = parseExpressionQuery(expr, expressionValue, standaloneString);
|
|
sb.append(toAppend.getFirst());
|
|
if (toAppend.getSecond() != null) {
|
|
parameters.add(toAppend.getSecond());
|
|
}
|
|
}
|
|
}
|
|
return new Pair<>(sb.toString(), parameters);
|
|
}
|
|
|
|
private Pair<String, Object> parseExpressionQuery(Expression<?> expr, Object expressionValue, boolean standaloneString) {
|
|
if (expr instanceof ExprUnsafe) {
|
|
if (standaloneString && expressionValue instanceof String) {
|
|
Skript.warning(
|
|
String.format("Unsafe may have been used unnecessarily. Try replacing 'unsafe %1$s' with %1$s",
|
|
((ExprUnsafe) expr).getRawExpression()));
|
|
}
|
|
return new Pair<>((String) expressionValue, null);
|
|
} else {
|
|
if (standaloneString) {
|
|
Skript.warning("Do not surround expressions with quotes!");
|
|
}
|
|
return new Pair<>("?", expressionValue);
|
|
}
|
|
}
|
|
|
|
private Map<String, Object> executeStatement(DataSource ds, String baseVariable, Pair<String, List<Object>> query) throws SkriptDBQueryException {
|
|
if (ds == null) {
|
|
throw new SkriptDBQueryException("Data source is not set");
|
|
}
|
|
try (Connection conn = ds.getConnection()) {
|
|
try (PreparedStatement stmt = createStatement(conn, query)) {
|
|
boolean hasResultSet = stmt.execute();
|
|
|
|
if (baseVariable != null) {
|
|
return processBaseVariable(baseVariable, stmt, hasResultSet);
|
|
}
|
|
return Map.of();
|
|
}
|
|
} catch (SQLException ex) {
|
|
throw new SkriptDBQueryException(ex.getMessage());
|
|
}
|
|
}
|
|
|
|
private Map<String, Object> processBaseVariable(String baseVariable, PreparedStatement stmt, boolean hasResultSet) throws SQLException {
|
|
if (isList) {
|
|
baseVariable = baseVariable.substring(0, baseVariable.length() - 1);
|
|
}
|
|
|
|
if (hasResultSet) {
|
|
CachedRowSet crs = SkriptDB.getRowSetFactory().createCachedRowSet();
|
|
crs.populate(generatedKeys ? stmt.getGeneratedKeys() : stmt.getResultSet());
|
|
|
|
if (isList) {
|
|
return fetchQueryResultSet(crs, baseVariable);
|
|
} else {
|
|
crs.last();
|
|
return Map.of(baseVariable, crs.getRow());
|
|
}
|
|
} else if (!isList) {
|
|
//if no results are returned and the specified variable isn't a list variable, put the affected rows count in the variable
|
|
return Map.of(baseVariable, stmt.getUpdateCount());
|
|
}
|
|
return Map.of();
|
|
}
|
|
|
|
private Map<String, Object> fetchQueryResultSet(CachedRowSet crs, String baseVariable) throws SQLException {
|
|
Map<String, Object> variableList = new HashMap<>();
|
|
ResultSetMetaData meta = crs.getMetaData();
|
|
int columnCount = meta.getColumnCount();
|
|
|
|
for (int i = 1; i <= columnCount; i++) {
|
|
String label = meta.getColumnLabel(i);
|
|
variableList.put(baseVariable + label, label);
|
|
}
|
|
|
|
int rowNumber = 1;
|
|
while (crs.next()) {
|
|
for (int i = 1; i <= columnCount; i++) {
|
|
variableList.put(baseVariable + meta.getColumnLabel(i).toLowerCase(Locale.ENGLISH)
|
|
+ Variable.SEPARATOR + rowNumber, crs.getObject(i));
|
|
}
|
|
rowNumber++;
|
|
}
|
|
return variableList;
|
|
}
|
|
|
|
private PreparedStatement createStatement(Connection conn, Pair<String, List<Object>> query) throws SQLException {
|
|
PreparedStatement stmt = generatedKeys ?
|
|
conn.prepareStatement(query.getFirst(), Statement.RETURN_GENERATED_KEYS)
|
|
: conn.prepareStatement(query.getFirst(), Statement.NO_GENERATED_KEYS);
|
|
if (query.getSecond() != null) {
|
|
Iterator<Object> iter = query.getSecond().iterator();
|
|
for (int i = 1; iter.hasNext(); i++) {
|
|
stmt.setObject(i, iter.next());
|
|
}
|
|
}
|
|
|
|
return stmt;
|
|
}
|
|
|
|
private boolean isStandaloneString(Object[] objects, int index) {
|
|
String before = getString(objects, index - 1);
|
|
String after = getString(objects, index + 1);
|
|
return before != null && before.endsWith("'") && after != null && after.endsWith("'");
|
|
}
|
|
|
|
private String getString(Object[] objects, int index) {
|
|
if (index >= 0 && index < objects.length && objects[index] instanceof String) {
|
|
return (String) objects[index];
|
|
}
|
|
return null;
|
|
}
|
|
|
|
private void setVariable(Event e, String name, Object obj) {
|
|
|
|
//fix mediumblob and similar column types, so they return a String correctly
|
|
if (obj != null) {
|
|
if (obj instanceof byte[]) {
|
|
obj = new String((byte[]) obj);
|
|
|
|
//in some servers instead of being byte array, it appears as SerialBlob (depends on mc version, 1.12.2 is bvte array, 1.16.5 SerialBlob)
|
|
} else if (obj instanceof SerialBlob) {
|
|
try {
|
|
obj = new String(((SerialBlob) obj).getBinaryStream().readAllBytes());
|
|
} catch (IOException | SerialException ex) {
|
|
ex.printStackTrace();
|
|
}
|
|
}
|
|
}
|
|
Variables.setVariable(name.toLowerCase(Locale.ENGLISH), obj, e, isLocal);
|
|
}
|
|
|
|
private static void resetLastSQLError() {
|
|
lastError = null;
|
|
}
|
|
|
|
private static void setLastSQLError(String error) {
|
|
lastError = error;
|
|
}
|
|
|
|
@Override
|
|
public String toString(Event e, boolean debug) {
|
|
return "execute " + query.toString(e, debug) + " in " + dataSource.toString(e, debug);
|
|
}
|
|
|
|
@SuppressWarnings("unchecked")
|
|
@Override
|
|
public boolean init(Expression<?>[] exprs, int matchedPattern, Kleenean isDelayed, SkriptParser.ParseResult parseResult) {
|
|
Expression<String> statementExpr = (Expression<String>) exprs[0];
|
|
if (statementExpr instanceof VariableString || statementExpr instanceof ExprUnsafe) {
|
|
query = statementExpr;
|
|
} else {
|
|
Skript.error("Database statements must be string literals. If you must use an expression, " +
|
|
"you may use \"%unsafe (your expression)%\", but keep in mind, you may be vulnerable " +
|
|
"to SQL injection attacks!");
|
|
return false;
|
|
}
|
|
dataSource = (Expression<HikariDataSource>) exprs[1];
|
|
if (exprs[2] != null) {
|
|
if (query instanceof VariableString && !((VariableString) query).isSimple()) {
|
|
Skript.warning("Your query string contains expresions, but you've also provided query arguments. Consider using `unsafe` keyword before your query.");
|
|
}
|
|
queryArguments = (Expression<Object>) exprs[2];
|
|
}
|
|
;
|
|
Expression<?> resultHolder = exprs[3];
|
|
quickly = parseResult.hasTag("quickly");
|
|
if (resultHolder instanceof Variable) {
|
|
Variable<?> varExpr = (Variable<?>) resultHolder;
|
|
resultVariableName = varExpr.getName();
|
|
isLocal = varExpr.isLocal();
|
|
isList = varExpr.isList();
|
|
generatedKeys = parseResult.hasTag("keys");
|
|
} else if (resultHolder != null) {
|
|
Skript.error(resultHolder + " is not a variable");
|
|
return false;
|
|
}
|
|
return true;
|
|
}
|
|
|
|
public static class SkriptDBQueryException extends RuntimeException {
|
|
|
|
private static final long serialVersionUID = -1869895286406538884L;
|
|
|
|
public SkriptDBQueryException(String message) {
|
|
super(message);
|
|
}
|
|
|
|
}
|
|
|
|
} |