Skip to content

Commit

Permalink
return updated rowcount when completing a task
Browse files Browse the repository at this point in the history
  • Loading branch information
Sourour Benzarti committed Jul 1, 2024
1 parent 5106ed4 commit 988cb76
Showing 1 changed file with 13 additions and 2 deletions.
15 changes: 13 additions & 2 deletions postgrestq/task_queue.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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.
Expand Down

0 comments on commit 988cb76

Please sign in to comment.