fix: ensure that only one worker processing task

Author: danfimov

README.md

@@ -1,6 +1,6 @@
 [![PyPI - Python Version](https://img.shields.io/pypi/pyversions/taskiq-postgres?style=for-the-badge&logo=python)](https://pypi.org/project/taskiq-postgres/)
 [![PyPI](https://img.shields.io/pypi/v/taskiq-postgres?style=for-the-badge&logo=pypi)](https://pypi.org/project/taskiq-postgres/)
-[![Checks](https://img.shields.io/github/check-runs/danfimov/taskiq-postgres/main?style=for-the-badge&logo=pytest)](https://github.com/danfimov/taskiq-postgres)
+[![Checks](https://img.shields.io/github/check-runs/danfimov/taskiq-postgres/main?nameFilter=Tests%20(3.12)&style=for-the-badge)](https://github.com/danfimov/taskiq-postgres)
 
 <div align="center">
 <a href="https://github.com/danfimov/taskiq-postgres/"><img src="https://raw.githubusercontent.com/danfimov/taskiq-postgres/main/assets/logo.png" width=400></a>

docs/contributing.md

@@ -0,0 +1,3 @@
+---
+title: Contributing
+---

docs/tutorial/broker.md

@@ -0,0 +1,69 @@
+---
+title: Broker
+---
+
+To use broker with PostgreSQL you need to import broker and result backend from this library and provide a address for  connection. For example, lets create a file `broker.py` with the following content:
+
+=== "asyncpg"
+
+    ```python
+    import asyncio
+    from taskiq_pg.asyncpg import AsyncpgResultBackend, AsyncpgBroker
+
+
+    dsn = "postgres://postgres:postgres@localhost:5432/postgres"
+    broker = AsyncpgBroker(dsn).with_result_backend(AsyncpgResultBackend(dsn))
+
+
+    @broker.task
+    async def best_task_ever() -> None:
+        """Solve all problems in the world."""
+        await asyncio.sleep(5.5)
+        print("All problems are solved!")
+
+
+    async def main():
+        await broker.startup()
+        task = await best_task_ever.kiq()
+        print(await task.wait_result())
+        await broker.shutdown()
+
+
+    if __name__ == "__main__":
+        asyncio.run(main())
+    ```
+
+=== "psqlpy"
+
+    ```python
+    import asyncio
+    from taskiq_pg.psqlpy import PSQLPyResultBackend, PSQLPyBroker
+
+
+    dsn = "postgres://postgres:postgres@localhost:5432/postgres"
+    broker = PSQLPyBroker(dsn).with_result_backend(PSQLPyResultBackend(dsn))
+
+
+    @broker.task
+    async def best_task_ever() -> None:
+        """Solve all problems in the world."""
+        await asyncio.sleep(5.5)
+        print("All problems are solved!")
+
+
+    async def main():
+        await broker.startup()
+        task = await best_task_ever.kiq()
+        print(await task.wait_result())
+        await broker.shutdown()
+
+
+    if __name__ == "__main__":
+        asyncio.run(main())
+    ```
+
+Then you can run this file with:
+
+```bash
+python broker.py
+```

docs/tutorial/quickstart.md

@@ -0,0 +1,3 @@
+---
+title: Getting Started
+---

docs/tutorial/schedule_source.md

@@ -0,0 +1,3 @@
+---
+title: Schedule Source
+---

examples/example_with_broker.py

@@ -0,0 +1,35 @@
+"""
+How to run:
+
+    1) Run worker in one terminal:
+        uv run taskiq worker examples.example_with_broker:broker
+
+    2) Run this script in another terminal:
+        uv run python -m examples.example_with_broker
+"""
+
+import asyncio
+
+from taskiq_pg.asyncpg import AsyncpgBroker, AsyncpgResultBackend
+
+
+dsn = "postgres://taskiq_postgres:look_in_vault@localhost:5432/taskiq_postgres"
+broker = AsyncpgBroker(dsn).with_result_backend(AsyncpgResultBackend(dsn))
+
+
+@broker.task("solve_all_problems")
+async def best_task_ever() -> None:
+    """Solve all problems in the world."""
+    await asyncio.sleep(2)
+    print("All problems are solved!")
+
+
+async def main():
+    await broker.startup()
+    task = await best_task_ever.kiq()
+    print(await task.wait_result())
+    await broker.shutdown()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

mkdocs.yml

@@ -8,6 +8,10 @@ edit_uri: edit/main/docs/
 nav:
   - Overview:
       - index.md
+  - Tutorial:
+      - tutorial/quickstart.md
+      - tutorial/broker.md
+      - tutorial/schedule_source.md
   - API:
       - reference.md
   - Contributing:

pyproject.toml

@@ -143,10 +143,19 @@ ignore = [
 
     "S101",  # assert usage
     "S311",  # pseudo-random generators are not suitable for cryptographic purposes
+    "S608",
+
+    "RUF",
 ]
 "tests/test_linting.py" = [
     "S603",  # subprocess usage
 ]
+"examples/*" = [
+    "T201",
+    "D",
+    "ANN",
+    "INP001",
+]
 
 [tool.ruff.lint.isort]
 known-local-folder = ["taskiq_pg"]

src/taskiq_pg/asyncpg/broker.py

@@ -10,10 +10,10 @@
 
 from taskiq_pg._internal.broker import BasePostgresBroker
 from taskiq_pg.asyncpg.queries import (
+    CLAIM_MESSAGE_QUERY,
     CREATE_MESSAGE_TABLE_QUERY,
     DELETE_MESSAGE_QUERY,
     INSERT_MESSAGE_QUERY,
-    SELECT_MESSAGE_QUERY,
 )
 
 
@@ -142,20 +142,14 @@ async def listen(self) -> AsyncGenerator[AckableMessage, None]:
             try:
                 payload = await self._queue.get()
                 message_id = int(payload)
-                message_row = await self.read_conn.fetchrow(
-                    SELECT_MESSAGE_QUERY.format(self.table_name),
-                    message_id,
-                )
-                if message_row is None:
-                    logger.warning(
-                        "Message with id %s not found in database.",
+                async with self.write_pool.acquire() as conn:
+                    claimed = await conn.fetchrow(
+                        CLAIM_MESSAGE_QUERY.format(self.table_name),
                         message_id,
                     )
+                if claimed is None:
                     continue
-                if message_row.get("message") is None:
-                    msg = "Message row does not have 'message' column"
-                    raise ValueError(msg)
-                message_str = message_row["message"]
+                message_str = claimed["message"]
                 if not isinstance(message_str, str):
                     msg = "message is not a string"
                     raise TypeError(msg)

src/taskiq_pg/asyncpg/queries.py

@@ -37,6 +37,7 @@
     task_name VARCHAR NOT NULL,
     message TEXT NOT NULL,
     labels JSONB NOT NULL,
+    status TEXT NOT NULL DEFAULT 'pending',
     created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()
 );
 """
@@ -47,6 +48,6 @@
 RETURNING id
 """
 
-SELECT_MESSAGE_QUERY = "SELECT * FROM {} WHERE id = $1"
+CLAIM_MESSAGE_QUERY = "UPDATE {} SET status = 'processing' WHERE id = $1 AND status = 'pending' RETURNING id, message"
 
 DELETE_MESSAGE_QUERY = "DELETE FROM {} WHERE id = $1"