1 files changed, 9 insertions, 12 deletions
diff --git a/src/main/java/io/trygvis/queue/TaskDao.java b/src/main/java/io/trygvis/queue/TaskDao.java
index 8b58585..9adec8f 100644
--- a/src/main/java/io/trygvis/queue/TaskDao.java
+++ b/src/main/java/io/trygvis/queue/TaskDao.java
@@ -57,10 +57,11 @@ public class TaskDao {
         }
     }
 
-    public List<Task> findByNameAndCompletedIsNull(String name) throws SQLException {
-        try (PreparedStatement stmt = c.prepareStatement("SELECT " + fields + " FROM task WHERE queue=? AND completed IS NULL")) {
+    public List<Task> findByQueueAndState(String queue, TaskState state) throws SQLException {
+        try (PreparedStatement stmt = c.prepareStatement("SELECT " + fields + " FROM task WHERE queue=? AND state=?")) {
             int i = 1;
-            stmt.setString(i, name);
+            stmt.setString(i++, queue);
+            stmt.setString(i, state.name());
             ResultSet rs = stmt.executeQuery();
             List<Task> list = new ArrayList<>();
             while (rs.next()) {
@@ -70,7 +71,7 @@ public class TaskDao {
         }
     }
 
-    public void update(Task task) throws SQLException {
+    public int update(Task task) throws SQLException {
         try (PreparedStatement stmt = c.prepareStatement("UPDATE task SET state=?, scheduled=?, last_run=?, run_count=?, completed=? WHERE id=?")) {
             int i = 1;
             stmt.setString(i++, task.state.name());
@@ -79,19 +80,15 @@ public class TaskDao {
             stmt.setInt(i++, task.runCount);
             setTimestamp(stmt, i++, task.completed);
             stmt.setLong(i, task.id());
-            stmt.executeUpdate();
+            return stmt.executeUpdate();
         }
     }
 
-    public void setState(List<Task> tasks, TaskState state) throws SQLException {
-        Long[] ids = new Long[tasks.size()];
-        for (int i = 0, tasksSize = tasks.size(); i < tasksSize; i++) {
-            ids[i] = tasks.get(i).id();
-        }
-        try (PreparedStatement stmt = c.prepareStatement("UPDATE task SET state=? WHERE id = ANY (?)")) {
+    public void setState(Task task, TaskState state) throws SQLException {
+        try (PreparedStatement stmt = c.prepareStatement("UPDATE task SET state=? WHERE id = ?")) {
             int i = 1;
             stmt.setString(i++, state.name());
-            stmt.setObject(i, c.createArrayOf("bigint", ids));
+            stmt.setLong(i, task.id());
             stmt.executeUpdate();
         }
     }