Added option to provide PreparedStatement-like query argument binding

src/main/java/com/btk5h/skriptdb/skript/EffExecuteStatement.java

@@ -25,6 +25,7 @@ import java.util.*;
 import java.util.concurrent.CompletableFuture;
 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
@@ -45,17 +46,20 @@ import java.util.concurrent.Executors;
  */
 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,
                 "execute %string% (in|on) %datasource% " +
-                        "[and store [[the] (output|result)[s]] (to|in) [the] [var[iable]] %-objects%]", "quickly execute %string% (in|on) %datasource% " +
+                        "[with arg[ument][s] %-objects%] [and store [[the] (output|result)[s]] (to|in) [the] [var[iable]] %-objects%]",
-                        "[and store [[the] (output|result)[s]] (to|in) [the] [var[iable]] %-objects%]");
+                "quickly execute %string% (in|on) %datasource% " +
+                        "[with arg[ument][s] %-objects%] [and store [[the] (output|result)[s]] (to|in) [the] [var[iable]] %-objects%]");
     }
 
     private Expression<String> query;
     private Expression<HikariDataSource> dataSource;
+    private Expression<Object> queryArguments;
     private VariableString var;
     private boolean isLocal;
     private boolean isList;
@@ -157,142 +161,144 @@ public class EffExecuteStatement extends Effect {
     }
 
     private Pair<String, List<Object>> parseQuery(Event e) {
-        if (!(query instanceof VariableString)) {
+        if (queryArguments != null) {
-            return new Pair<>(query.getSingle(e), null);
+            Object[] args = queryArguments.getArray(e);
-        }
+            String queryString = query.getSingle(e);
-        VariableString q = (VariableString) query;
+            int queryArgCount = (int) ARGUMENT_PLACEHOLDER.matcher(queryString).results().count();
-        if (q.isSimple()) {
+            if (queryArgCount != args.length) {
-            return new Pair<>(q.toString(e), null);
+                Skript.warning("Your query has %d question marks, but you provided %d arguments.");
+                args = Arrays.copyOf(args, queryArgCount);
+            }
+            return new Pair<>(query.getSingle(e), List.of(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 ArrayList<>();
+        List<Object> parameters = new LinkedList<>();
-        Object[] objects = SkriptUtil.getTemplateString(q);
+        Object[] objects = SkriptUtil.getTemplateString(varQuery);
         
         for (int i = 0; i < objects.length; i++) {
-            Object o = objects[i];
+            if (objects[i] instanceof String) {
-            if (o instanceof String) {
+                sb.append(objects[i]);
-                sb.append(o);
             } else {
-                Expression<?> expr;
+                Expression<?> expr = objects[i] instanceof Expression ? (Expression<?>) objects[i] : SkriptUtil.getExpressionFromInfo(objects[i]);
-                if (o instanceof Expression)
+                boolean standaloneString = isStandaloneString(objects, i);
-                    expr = (Expression<?>) o;
-                else
-                    expr = SkriptUtil.getExpressionFromInfo(o);
-
-                String before = getString(objects, i - 1);
-                String after = getString(objects, i + 1);
-                boolean standaloneString = false;
-
-                if (before != null && after != null) {
-                    if (before.endsWith("'") && after.endsWith("'")) {
-                        standaloneString = true;
-                    }
-                }
-
                 Object expressionValue = expr.getSingle(e);
 
-                if (expr instanceof ExprUnsafe) {
+                Pair<String, Object> toAppend = parseExpressionQuery(expr, expressionValue, standaloneString);
-                    sb.append(expressionValue);
+                sb.append(toAppend.getFirst());
-
+                if (toAppend.getSecond() != null) {
-                    if (standaloneString && expressionValue instanceof String) {
+                    parameters.add(toAppend.getSecond());
-                        String rawExpression = ((ExprUnsafe) expr).getRawExpression();
-                        Skript.warning(
-                                String.format("Unsafe may have been used unnecessarily. Try replacing 'unsafe %1$s' with %1$s",
-                                        rawExpression));
-                    }
-                } else {
-                    parameters.add(expressionValue);
-                    sb.append('?');
-
-                    if (standaloneString) {
-                        Skript.warning("Do not surround expressions with quotes!");
-                    }
                 }
             }
         }
         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 Object executeStatement(DataSource ds, String baseVariable, Pair<String, List<Object>> query) {
         if (ds == null) {
             return "Data source is not set";
         }
-        Map<String, Object> variableList = new HashMap<>();
+        try (Connection conn = ds.getConnection()) {
-        try (Connection conn = ds.getConnection();
+            try (PreparedStatement stmt = createStatement(conn, query)) {
-             PreparedStatement stmt = createStatement(conn, query)) {
+                boolean hasResultSet = stmt.execute();
 
-            boolean hasResultSet = stmt.execute();
+                if (baseVariable != null) {
-
+                    return processBaseVariable(baseVariable, stmt, hasResultSet);
-            if (baseVariable != null) {
-                if (isList) {
-                    baseVariable = baseVariable.substring(0, baseVariable.length() - 1);
-                }
-
-                if (hasResultSet) {
-                    CachedRowSet crs = SkriptDB.getRowSetFactory().createCachedRowSet();
-                    crs.populate(stmt.getResultSet());
-
-                    if (isList) {
-                        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;
-                        try {
-                            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++;
-                            }
-                        } catch (SQLException ex) {
-                            return ex.getMessage();
-                        }
-                    } else {
-                        crs.last();
-                        variableList.put(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
-                    variableList.put(baseVariable, stmt.getUpdateCount());
                 }
+                return Map.of();
             }
         } catch (SQLException ex) {
             return ex.getMessage();
         }
+    }
+    
+    private Object processBaseVariable(String baseVariable, PreparedStatement stmt, boolean hasResultSet) throws SQLException {
+        Map<String, Object> variableList = new HashMap<>();
+        if (isList) {
+            baseVariable = baseVariable.substring(0, baseVariable.length() - 1);
+        }
+
+        if (hasResultSet) {
+            CachedRowSet crs = SkriptDB.getRowSetFactory().createCachedRowSet();
+            crs.populate(stmt.getResultSet());
+
+            if (isList) {
+                return fetchQueryResultSet(crs, baseVariable);
+            } else {
+                crs.last();
+                variableList.put(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 = conn.prepareStatement(query.getFirst());
-        List<Object> parameters = query.getSecond();
+        if (query.getSecond() != null) {
-
+            Iterator<Object> iter = query.getSecond().iterator();
-        if (parameters != null) {
+            for (int i = 1; iter.hasNext(); i++) {
-            for (int i = 0; i < parameters.size(); i++) {
+                stmt.setObject(i, iter.next());
-                stmt.setObject(i + 1, parameters.get(i));
             }
         }
 
         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) {
+        if (index >= 0 && index < objects.length && objects[index] instanceof String) {
-            return null;
+            return (String) objects[index];
         }
-
-        Object object = objects[index];
-
-        if (object instanceof String) {
-            return (String) object;
-        }
-
         return null;
     }
 
@@ -334,15 +340,21 @@ public class EffExecuteStatement extends Effect {
             return false;
         }
         dataSource = (Expression<HikariDataSource>) exprs[1];
-        Expression<?> expr = exprs[2];
+        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 = matchedPattern == 1;
-        if (expr instanceof Variable) {
+        if (resultHolder instanceof Variable) {
-            Variable<?> varExpr = (Variable<?>) expr;
+            Variable<?> varExpr = (Variable<?>) resultHolder;
             var = varExpr.getName();
             isLocal = varExpr.isLocal();
             isList = varExpr.isList();
-        } else if (expr != null) {
+        } else if (resultHolder != null) {
-            Skript.error(expr + " is not a variable");
+            Skript.error(resultHolder + " is not a variable");
             return false;
         }
         return true;