From 988cb767798b7755b4a1c1130d32d14a45b10bae Mon Sep 17 00:00:00 2001
From: Sourour Benzarti <sourour.benzarti@flixbus.com>
Date: Mon, 1 Jul 2024 18:36:18 +0200
Subject: [PATCH] return updated rowcount when completing a task

---
 postgrestq/task_queue.py | 15 +++++++++++++--
 1 file changed, 13 insertions(+), 2 deletions(-)

diff --git a/postgrestq/task_queue.py b/postgrestq/task_queue.py
index f9c8def..be7922a 100644
--- a/postgrestq/task_queue.py
+++ b/postgrestq/task_queue.py
@@ -394,7 +394,7 @@ def get_many(self, amount: int) -> Sequence[
             conn.commit()
             return ret
 
-    def complete(self, task_id: Optional[UUID]) -> None:
+    def complete(self, task_id: Optional[UUID]) -> int:
         """Mark a task as completed.
 
         Marks a task as completed by setting completed_at column by
@@ -409,20 +409,31 @@ def complete(self, task_id: Optional[UUID]) -> None:
         task_id : UUID | None
             the task ID
 
+        Returns
+        -------
+        the number of updated rows: int
+
         """
         logger.info(f"Marking task {task_id} as completed")
         conn = self.conn
+        count = 0
         with conn.cursor() as cur:
             cur.execute(
                 sql.SQL(
                     """
                 UPDATE {}
                 SET completed_at = current_timestamp
-                WHERE id = %s"""
+                WHERE id = %s
+                AND completed_at is NULL"""
                 ).format(sql.Identifier(self._table_name)),
                 (task_id,),
             )
+            count = cur.rowcount
+            if count == 0:
+                logger.info(f"Task {task_id} was already completed")
+
             conn.commit()
+        return count
 
     def is_empty(self) -> bool:
         """Check if the task queue is empty.