1 files changed, 15 insertions, 2 deletions
diff --git a/src/main/java/io/trygvis/queue/TaskDao.java b/src/main/java/io/trygvis/queue/TaskDao.java
index 5d77a41..365b44b 100644
--- a/src/main/java/io/trygvis/queue/TaskDao.java
+++ b/src/main/java/io/trygvis/queue/TaskDao.java
@@ -84,14 +84,27 @@ public class TaskDao {
     }
 
     public int update(Task task) throws SQLException {
-        try (PreparedStatement stmt = c.prepareStatement("UPDATE task SET state=?, scheduled=?, last_run=?, run_count=?, completed=? WHERE id=?")) {
+        return update(task, null);
+    }
+
+    public int update(Task task, TaskState state) throws SQLException {
+        String sql = "UPDATE task SET state=?, scheduled=?, last_run=?, run_count=?, completed=? WHERE id=?";
+
+        if (state != null) {
+            sql += " AND state=?";
+        }
+
+        try (PreparedStatement stmt = c.prepareStatement(sql)) {
             int i = 1;
             stmt.setString(i++, task.state.name());
             stmt.setTimestamp(i++, new Timestamp(task.scheduled.getTime()));
             setTimestamp(stmt, i++, task.lastRun);
             stmt.setInt(i++, task.runCount);
             setTimestamp(stmt, i++, task.completed);
-            stmt.setLong(i, task.id());
+            stmt.setLong(i++, task.id());
+            if (state != null) {
+                stmt.setString(i, state.name());
+            }
             return stmt.executeUpdate();
         }
     }