Fix

Author: ZorTik

core/src/main/java/me/zort/sqllib/SQLiteDatabaseConnectionImpl.java

@@ -1,5 +1,6 @@
 package me.zort.sqllib;
 
+import me.zort.sqllib.api.Query;
 import me.zort.sqllib.api.data.QueryResult;
 import me.zort.sqllib.api.data.QueryRowsResult;
 import me.zort.sqllib.api.data.Row;
@@ -136,4 +137,11 @@ public UpsertQuery upsert(@Nullable String table) {
         throw new IllegalStatementOperationException("Default upsert is not supported by SQLite!");
     }
 
+    @Override
+    public QueryResult exec(Query query) {
+        if (query instanceof UpsertQuery && ((UpsertQuery) query).getAssignedSaveObject() != null)
+            return save(((UpsertQuery) query).getTable(), ((UpsertQuery) query).getAssignedSaveObject());
+
+        return super.exec(query);
+    }
 }

core/src/main/java/me/zort/sqllib/internal/query/UpsertQuery.java

@@ -1,5 +1,7 @@
 package me.zort.sqllib.internal.query;
 
+import lombok.Getter;
+import lombok.Setter;
 import me.zort.sqllib.SQLDatabaseConnection;
 import me.zort.sqllib.internal.query.part.SetStatement;
 import org.jetbrains.annotations.Nullable;
@@ -8,6 +10,10 @@
 
 public class UpsertQuery extends InsertQuery {
 
+    @Setter
+    @Getter
+    private Object assignedSaveObject = null; // Can be null, only relevant in SQLite mode.
+
     public UpsertQuery(SQLDatabaseConnection connection) {
         super(connection);
     }

core/src/main/java/me/zort/sqllib/mapping/builder/NativeQueryBuilder.java

@@ -1,5 +1,7 @@
 package me.zort.sqllib.mapping.builder;
 
+import lombok.Getter;
+import me.zort.sqllib.api.Executive;
 import me.zort.sqllib.api.SQLConnection;
 import me.zort.sqllib.internal.query.QueryDetails;
 import me.zort.sqllib.internal.query.QueryNode;
@@ -20,7 +22,7 @@ public class NativeQueryBuilder implements QueryAnnotation.QueryBuilder<Annotati
     @Override
     public QueryNode<?> build(SQLConnection connection, Annotation queryAnnotation, Method method, ParameterPair[] parameters) {
         PlaceholderMapper mapper = new PlaceholderMapper(parameters);
-        return new QueryNode<QueryNode<?>>(null, new ArrayList<>(), QueryPriority.GENERAL) {
+        return new LocalQueryNodeExecutive(connection) {
             private int currPhIndex = 0;
             @Override
             public QueryDetails buildQueryDetails() {
@@ -59,4 +61,15 @@ private String nextPlaceholder() {
         };
     }
 
+    private static abstract class LocalQueryNodeExecutive extends QueryNode<QueryNode<?>> implements Executive {
+
+        @Getter
+        private final SQLConnection connection;
+
+        public LocalQueryNodeExecutive(SQLConnection connection) {
+            super(null, new ArrayList<>(), QueryPriority.GENERAL);
+            this.connection = connection;
+        }
+    }
+
 }

core/src/main/java/me/zort/sqllib/mapping/builder/SaveQueryBuilder.java

@@ -20,8 +20,11 @@ public QueryNode<?> build(SQLConnection connection, Save queryAnnotation, Method
 
         String table = Table.Util.getFromContext(method, parameters);
 
-        UpsertQuery query = ((SQLDatabaseConnectionImpl) connection).save(getSaveableObject(parameters));
+        Object saveableObject = getSaveableObject(parameters);
+
+        UpsertQuery query = ((SQLDatabaseConnectionImpl) connection).save(saveableObject);
         query.table(table);
+        query.setAssignedSaveObject(saveableObject);
 
         return query;
     }