From 99be47d4d301b48cc6fdd795d89f42d1f0fb2c84 Mon Sep 17 00:00:00 2001
From: szumielxd <43210079+szumielxd@users.noreply.github.com>
Date: Thu, 30 Jan 2025 02:41:32 +0100
Subject: [PATCH] Added additional arguments info in prepared statement warning

---
 .../skriptdb/skript/EffExecuteStatement.java  | 34 +++++++++++++------
 1 file changed, 23 insertions(+), 11 deletions(-)

diff --git a/src/main/java/com/btk5h/skriptdb/skript/EffExecuteStatement.java b/src/main/java/com/btk5h/skriptdb/skript/EffExecuteStatement.java
index 70ca721..ad6bfb5 100644
--- a/src/main/java/com/btk5h/skriptdb/skript/EffExecuteStatement.java
+++ b/src/main/java/com/btk5h/skriptdb/skript/EffExecuteStatement.java
@@ -21,8 +21,10 @@ 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;
@@ -52,7 +54,7 @@ public class EffExecuteStatement extends Effect {
     static {
         Skript.registerEffect(EffExecuteStatement.class,
                 "[quickly:quickly] execute %string% (in|on) %datasource% " +
-                        "[with arg[ument][s] %-objects%] [and store [[the] (output|result)[s]] (to|in) [the] [var[iable]] %-objects%]");
+                        "[with arg[ument][s] %-objects%] [and store [[the] [keys:generated keys] (output|result)[s]] (to|in) [the] [var[iable]] %-objects%]");
     }
 
     private Expression<String> query;
@@ -62,6 +64,7 @@ public class EffExecuteStatement extends Effect {
     private boolean isLocal;
     private boolean isList;
     private boolean quickly;
+    private boolean generatedKeys;
     private boolean isSync = false;
 
     @Override
@@ -82,8 +85,8 @@ public class EffExecuteStatement extends Effect {
 		            .whenComplete((resources, err) -> {
 		            	//handle last error syntax data
 		            	resetLastSQLError();
-		            	if (err instanceof SkriptDBQueryException) {
-	            			setLastSQLError(err.getMessage());
+		            	if (err instanceof CompletionException && err.getCause() instanceof SkriptDBQueryException) {
+	            			setLastSQLError(err.getCause().getMessage());
 	            		}
 		                //if local variables are present
 		                //bring back local variables
@@ -99,9 +102,9 @@ public class EffExecuteStatement extends Effect {
         } else {
             isSync = true;
             Map<String, Object> resources = null;
-			try {
+            resetLastSQLError();
+            try {
                 resources = executeStatement(ds, baseVariable, parsedQuery);
-                resetLastSQLError();
 			} catch (SkriptDBQueryException err) {
 	            //handle last error syntax data
 				setLastSQLError(err.getMessage());
@@ -141,10 +144,16 @@ public class EffExecuteStatement extends Effect {
             String queryString = query.getSingle(e);
             int queryArgCount = (int) ARGUMENT_PLACEHOLDER.matcher(queryString).results().count();
             if (queryArgCount != args.length) {
-                Skript.warning("Your query has %d question marks, but you provided %d arguments.");
+                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), List.of(args));
+            return new Pair<>(query.getSingle(e), Arrays.asList(args));
         } else if (query instanceof VariableString && !((VariableString) query).isSimple()) {
             return parseVariableQuery(e, (VariableString) query);
         }
@@ -209,20 +218,19 @@ public class EffExecuteStatement extends Effect {
     }
     
     private Map<String, 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());
+            crs.populate(generatedKeys ? stmt.getGeneratedKeys() : stmt.getResultSet());
 
             if (isList) {
                 return fetchQueryResultSet(crs, baseVariable);
             } else {
                 crs.last();
-                variableList.put(baseVariable, crs.getRow());
+                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
@@ -253,7 +261,9 @@ public class EffExecuteStatement extends Effect {
     }
 
     private PreparedStatement createStatement(Connection conn, Pair<String, List<Object>> query) throws SQLException {
-        PreparedStatement stmt = conn.prepareStatement(query.getFirst());
+        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++) {
@@ -328,6 +338,7 @@ public class EffExecuteStatement extends Effect {
             }
             queryArguments = (Expression<Object>) exprs[2];
         }
+        ;
         Expression<?> resultHolder = exprs[3];
         quickly = parseResult.hasTag("quickly");
         if (resultHolder instanceof Variable) {
@@ -335,6 +346,7 @@ public class EffExecuteStatement extends Effect {
             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;