beanqueue 0.1.2__tar.gz → 0.1.3__tar.gz

{beanqueue-0.1.2 → beanqueue-0.1.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: beanqueue
-Version: 0.1.2
+Version: 0.1.3
 Summary: BeanQueue or BQ for short, PostgreSQL SKIP LOCK based worker queue library
 License: MIT
 Author: Fang-Pen Lin
@@ -19,7 +19,7 @@ Requires-Dist: venusian (>=3.1.0,<4.0.0)
 Description-Content-Type: text/markdown
 
 # BeanQueue  [![CircleCI](https://dl.circleci.com/status-badge/img/gh/LaunchPlatform/bq/tree/master.svg?style=svg)](https://dl.circleci.com/status-badge/redirect/gh/LaunchPlatform/beanhub-extract/tree/master)
-BeanQueue, a lightweight worker queue framework based on [SQLAlchemy](https://www.sqlalchemy.org/), [PostgreSQL SKIP LOCKED queries](https://www.2ndquadrant.com/en/blog/what-is-select-skip-locked-for-in-postgresql-9-5/) and [NOTIFY](https://www.postgresql.org/docs/current/sql-notify.html) / [LISTEN](https://www.postgresql.org/docs/current/sql-listen.html) statements.
+BeanQueue, a lightweight worker queue framework based on [SQLAlchemy](https://www.sqlalchemy.org/), PostgreSQL [SKIP LOCKED queries](https://www.2ndquadrant.com/en/blog/what-is-select-skip-locked-for-in-postgresql-9-5/) and [NOTIFY](https://www.postgresql.org/docs/current/sql-notify.html) / [LISTEN](https://www.postgresql.org/docs/current/sql-listen.html) statements.
 
 **Notice**: Still in its early stage, we built this for [BeanHub](https://beanhub.io)'s internal usage. May change rapidly. Use at your own risk for now.
 
@@ -29,7 +29,7 @@ BeanQueue, a lightweight worker queue framework based on [SQLAlchemy](https://ww
 - **Easy-to-deploy**: Only rely on PostgreSQL
 - **Easy-to-use**: Provide command line tools for processing tasks, also helpers for generating tasks models
 - **Auto-notify**: Notify will automatically be generated and send for inserted or update tasks
-- **Worker heartbeat and auto-reschedule**: Each worker keeps updating heartbeat, if one is dead, the others will reschedule the tasks
+- **Worker heartbeat and auto-reschedule**: Each worker keeps updating heartbeat, if one is found dead, the others will reschedule the tasks
 - **Customizable**: Use it as an library and build your own worker queue
 - **Native DB operations**: Commit your tasks with other db entries altogether without worrying about data inconsistent issue
 
@@ -46,14 +46,13 @@ You can define a task processor like this
 ```python
 from sqlalchemy.orm import Session
 
-from bq.processors.registry import processor
-from bq import models
-from .. import my_models
+import bq
+from .. import models
 from .. import image_utils
 
-@processor(channel="images")
-def resize_image(db: Session, task: models.Task, width: int, height: int):
-    image = db.query(my_models.Image).filter(my_models.Image.task == task).one()
+@bq.processor(channel="images")
+def resize_image(db: Session, task: bq.Task, width: int, height: int):
+    image = db.query(models.Image).filter(models.Image.task == task).one()
     image_utils.resize(image, size=(width, height))
     db.add(image)
     # by default the `processor` decorator has `auto_complete` flag turns on,
@@ -63,21 +62,21 @@ def resize_image(db: Session, task: models.Task, width: int, height: int):
 The `db` and `task` keyword arguments are optional.
 If you don't need to access the task object, you can simply define the function without these two parameters.
 
-To submit a task, you can either use `bq.models.Task` model object to construct the task object, insert into the
+To submit a task, you can either use `bq.Task` model object to construct the task object, insert into the
 database session and commit.
 
 ```python
-from bq import models
+import bq
 from .db import Session
-from .. import my_models
+from .. import models
 
 db = Session()
-task = models.Task(
+task = bq.Task(
     channel="files",
     module="my_pkgs.files.processors",
     name="upload_to_s3_for_backup",
 )
-file = my_models.File(
+file = models.File(
     task=task,
     blob_name="...",
 )
@@ -112,6 +111,7 @@ To run the worker, you can do this:
 BQ_PROCESSOR_PACKAGES='["my_pkgs.processors"]' python -m bq.cmds.process images
 ```
 
+The `BQ_PROCESSOR_PACKAGES` is a JSON list contains the Python packages where you define your processors (the functions you decorated with `bq.processors.registry.processor`).
 To submit a task for testing purpose, you can do
 
 ```bash
@@ -136,24 +136,104 @@ If you want to configure BeanQueue programmatically for the command lines, you c
 For example:
 
 ```python
-import bq.cmds.process
-from bq.container import Container
-from bq.config import Config
-
-container = Container()
-container.wire(modules=[bq.cmds.process])
-with container.config.override(
-    Config(
-        PROCESSOR_PACKAGES=["my_pkgs.processors"],
-        DATABASE_URL="postgresql://...",
-        BATCH_SIZE=10,
-    )
-):
-    bq.cmds.process.process_tasks(channels=("images",))
+import bq
+from bq.cmds.process import process_tasks
+from .my_config import config
+
+container = bq.Container()
+container.wire(packages=[bq])
+config = bq.Config(
+    PROCESSOR_PACKAGES=["my_pkgs.processors"],
+    DATABASE_URL=str(config.DATABASE_URL),
+    BATCH_SIZE=10,
+)
+with container.config.override(config):
+    process_tasks(channels=("images",))
 ```
 
 Many other behaviors of this framework can also be modified by overriding the container defined at [bq/container.py](bq/container.py).
 
+### Define your own tables
+
+BeanQueue is designed to be as customizable as much as possible.
+Of course, you can define your own SQLAlchemy model instead of using the ones we provided. 
+
+To make defining your own `Task` model or `Worker` model much easier, you can use our mixin classes:
+
+- `bq.TaskModelMixin`: provides task model columns
+- `bq.TaskModelRefWorkerMixin`: provides foreign key column and relationship to `bq.Worker`
+- `bq.WorkerModelMixin`: provides worker model columns
+- `bq.WorkerRefMixin`: provides relationship to `bq.Task`
+
+Here's an example for defining your own Task model:
+
+```python
+import uuid
+
+import bq
+from sqlalchemy import ForeignKey
+from sqlalchemy.dialects.postgresql import UUID
+from sqlalchemy.orm import Mapped
+from sqlalchemy.orm import mapped_column
+from sqlalchemy.orm import relationship
+
+from .base_class import Base
+
+
+class Task(bq.TaskModelMixin, Base):
+    __tablename__ = "task"
+    worker_id: Mapped[uuid.UUID] = mapped_column(
+        UUID(as_uuid=True),
+        ForeignKey("worker.id", onupdate="CASCADE"),
+        nullable=True,
+        index=True,
+    )
+
+    worker: Mapped["Worker"] = relationship(
+        "Worker", back_populates="tasks", uselist=False
+    )
+```
+
+To make task insert and update with state changing to `PENDING` send out NOTIFY "channel" statement automatically, you can also use `bq.models.task.listen_events` helper to register our SQLAlchemy event handlers automatically like this
+
+```python
+from bq.models.task import listen_events
+listen_events(Task)
+```
+
+You just see how easy it is to define your Task model. Now, here's an example for defining your own Worker model:
+
+```python
+import bq
+from sqlalchemy.orm import Mapped
+from sqlalchemy.orm import relationship
+
+from .base_class import Base
+
+
+class Worker(bq.WorkerModelMixin, Base):
+    __tablename__ = "worker"
+
+    tasks: Mapped[list["Task"]] = relationship(
+        "Task",
+        back_populates="worker",
+        cascade="all,delete",
+        order_by="Task.created_at",
+    )
+```
+
+With the model class ready, you only need to change the `TASK_MODEL` and `WORKER_MODEL` of `Config` to the full Python module name plus the class name like this.
+
+```python
+import bq
+config = bq.Config(
+    TASK_MODEL="my_pkgs.models.Task",
+    WORKER_MODEL="my_pkgs.models.Worker",
+    # ... other configs
+)
+# Override container...
+```
+
 ## Why?
 
 There are countless worker queue projects. Why make yet another one?
@@ -230,6 +310,7 @@ A modern accounting book service based on the most popular open source version c
 
 - [solid_queue](https://github.com/rails/solid_queue)
 - [postgres-tq](https://github.com/flix-tech/postgres-tq)
+- [pq](https://github.com/malthe/pq/)
 - [PgQueuer](https://github.com/janbjorge/PgQueuer)
 - [hatchet](https://github.com/hatchet-dev/hatchet)
 

{beanqueue-0.1.2 → beanqueue-0.1.3}/README.md

@@ -1,5 +1,5 @@
 # BeanQueue  [![CircleCI](https://dl.circleci.com/status-badge/img/gh/LaunchPlatform/bq/tree/master.svg?style=svg)](https://dl.circleci.com/status-badge/redirect/gh/LaunchPlatform/beanhub-extract/tree/master)
-BeanQueue, a lightweight worker queue framework based on [SQLAlchemy](https://www.sqlalchemy.org/), [PostgreSQL SKIP LOCKED queries](https://www.2ndquadrant.com/en/blog/what-is-select-skip-locked-for-in-postgresql-9-5/) and [NOTIFY](https://www.postgresql.org/docs/current/sql-notify.html) / [LISTEN](https://www.postgresql.org/docs/current/sql-listen.html) statements.
+BeanQueue, a lightweight worker queue framework based on [SQLAlchemy](https://www.sqlalchemy.org/), PostgreSQL [SKIP LOCKED queries](https://www.2ndquadrant.com/en/blog/what-is-select-skip-locked-for-in-postgresql-9-5/) and [NOTIFY](https://www.postgresql.org/docs/current/sql-notify.html) / [LISTEN](https://www.postgresql.org/docs/current/sql-listen.html) statements.
 
 **Notice**: Still in its early stage, we built this for [BeanHub](https://beanhub.io)'s internal usage. May change rapidly. Use at your own risk for now.
 
@@ -9,7 +9,7 @@ BeanQueue, a lightweight worker queue framework based on [SQLAlchemy](https://ww
 - **Easy-to-deploy**: Only rely on PostgreSQL
 - **Easy-to-use**: Provide command line tools for processing tasks, also helpers for generating tasks models
 - **Auto-notify**: Notify will automatically be generated and send for inserted or update tasks
-- **Worker heartbeat and auto-reschedule**: Each worker keeps updating heartbeat, if one is dead, the others will reschedule the tasks
+- **Worker heartbeat and auto-reschedule**: Each worker keeps updating heartbeat, if one is found dead, the others will reschedule the tasks
 - **Customizable**: Use it as an library and build your own worker queue
 - **Native DB operations**: Commit your tasks with other db entries altogether without worrying about data inconsistent issue
 
@@ -26,14 +26,13 @@ You can define a task processor like this
 ```python
 from sqlalchemy.orm import Session
 
-from bq.processors.registry import processor
-from bq import models
-from .. import my_models
+import bq
+from .. import models
 from .. import image_utils
 
-@processor(channel="images")
-def resize_image(db: Session, task: models.Task, width: int, height: int):
-    image = db.query(my_models.Image).filter(my_models.Image.task == task).one()
+@bq.processor(channel="images")
+def resize_image(db: Session, task: bq.Task, width: int, height: int):
+    image = db.query(models.Image).filter(models.Image.task == task).one()
     image_utils.resize(image, size=(width, height))
     db.add(image)
     # by default the `processor` decorator has `auto_complete` flag turns on,
@@ -43,21 +42,21 @@ def resize_image(db: Session, task: models.Task, width: int, height: int):
 The `db` and `task` keyword arguments are optional.
 If you don't need to access the task object, you can simply define the function without these two parameters.
 
-To submit a task, you can either use `bq.models.Task` model object to construct the task object, insert into the
+To submit a task, you can either use `bq.Task` model object to construct the task object, insert into the
 database session and commit.
 
 ```python
-from bq import models
+import bq
 from .db import Session
-from .. import my_models
+from .. import models
 
 db = Session()
-task = models.Task(
+task = bq.Task(
     channel="files",
     module="my_pkgs.files.processors",
     name="upload_to_s3_for_backup",
 )
-file = my_models.File(
+file = models.File(
     task=task,
     blob_name="...",
 )
@@ -92,6 +91,7 @@ To run the worker, you can do this:
 BQ_PROCESSOR_PACKAGES='["my_pkgs.processors"]' python -m bq.cmds.process images
 ```
 
+The `BQ_PROCESSOR_PACKAGES` is a JSON list contains the Python packages where you define your processors (the functions you decorated with `bq.processors.registry.processor`).
 To submit a task for testing purpose, you can do
 
 ```bash
@@ -116,24 +116,104 @@ If you want to configure BeanQueue programmatically for the command lines, you c
 For example:
 
 ```python
-import bq.cmds.process
-from bq.container import Container
-from bq.config import Config
-
-container = Container()
-container.wire(modules=[bq.cmds.process])
-with container.config.override(
-    Config(
-        PROCESSOR_PACKAGES=["my_pkgs.processors"],
-        DATABASE_URL="postgresql://...",
-        BATCH_SIZE=10,
-    )
-):
-    bq.cmds.process.process_tasks(channels=("images",))
+import bq
+from bq.cmds.process import process_tasks
+from .my_config import config
+
+container = bq.Container()
+container.wire(packages=[bq])
+config = bq.Config(
+    PROCESSOR_PACKAGES=["my_pkgs.processors"],
+    DATABASE_URL=str(config.DATABASE_URL),
+    BATCH_SIZE=10,
+)
+with container.config.override(config):
+    process_tasks(channels=("images",))
 ```
 
 Many other behaviors of this framework can also be modified by overriding the container defined at [bq/container.py](bq/container.py).
 
+### Define your own tables
+
+BeanQueue is designed to be as customizable as much as possible.
+Of course, you can define your own SQLAlchemy model instead of using the ones we provided. 
+
+To make defining your own `Task` model or `Worker` model much easier, you can use our mixin classes:
+
+- `bq.TaskModelMixin`: provides task model columns
+- `bq.TaskModelRefWorkerMixin`: provides foreign key column and relationship to `bq.Worker`
+- `bq.WorkerModelMixin`: provides worker model columns
+- `bq.WorkerRefMixin`: provides relationship to `bq.Task`
+
+Here's an example for defining your own Task model:
+
+```python
+import uuid
+
+import bq
+from sqlalchemy import ForeignKey
+from sqlalchemy.dialects.postgresql import UUID
+from sqlalchemy.orm import Mapped
+from sqlalchemy.orm import mapped_column
+from sqlalchemy.orm import relationship
+
+from .base_class import Base
+
+
+class Task(bq.TaskModelMixin, Base):
+    __tablename__ = "task"
+    worker_id: Mapped[uuid.UUID] = mapped_column(
+        UUID(as_uuid=True),
+        ForeignKey("worker.id", onupdate="CASCADE"),
+        nullable=True,
+        index=True,
+    )
+
+    worker: Mapped["Worker"] = relationship(
+        "Worker", back_populates="tasks", uselist=False
+    )
+```
+
+To make task insert and update with state changing to `PENDING` send out NOTIFY "channel" statement automatically, you can also use `bq.models.task.listen_events` helper to register our SQLAlchemy event handlers automatically like this
+
+```python
+from bq.models.task import listen_events
+listen_events(Task)
+```
+
+You just see how easy it is to define your Task model. Now, here's an example for defining your own Worker model:
+
+```python
+import bq
+from sqlalchemy.orm import Mapped
+from sqlalchemy.orm import relationship
+
+from .base_class import Base
+
+
+class Worker(bq.WorkerModelMixin, Base):
+    __tablename__ = "worker"
+
+    tasks: Mapped[list["Task"]] = relationship(
+        "Task",
+        back_populates="worker",
+        cascade="all,delete",
+        order_by="Task.created_at",
+    )
+```
+
+With the model class ready, you only need to change the `TASK_MODEL` and `WORKER_MODEL` of `Config` to the full Python module name plus the class name like this.
+
+```python
+import bq
+config = bq.Config(
+    TASK_MODEL="my_pkgs.models.Task",
+    WORKER_MODEL="my_pkgs.models.Worker",
+    # ... other configs
+)
+# Override container...
+```
+
 ## Why?
 
 There are countless worker queue projects. Why make yet another one?
@@ -210,5 +290,6 @@ A modern accounting book service based on the most popular open source version c
 
 - [solid_queue](https://github.com/rails/solid_queue)
 - [postgres-tq](https://github.com/flix-tech/postgres-tq)
+- [pq](https://github.com/malthe/pq/)
 - [PgQueuer](https://github.com/janbjorge/PgQueuer)
 - [hatchet](https://github.com/hatchet-dev/hatchet)

beanqueue-0.1.3/bq/__init__.py

@@ -0,0 +1,11 @@
+from .config import Config  # noqa
+from .container import Container  # noqa
+from .models import Task  # noqa
+from .models import TaskModelMixin
+from .models import TaskModelRefWorkerMixin
+from .models import TaskState  # noqa
+from .models import Worker  # noqa
+from .models import WorkerModelMixin  # noqa
+from .models import WorkerRefMixin  # noqa
+from .models import WorkerState  # noqa
+from .processors.registry import processor  # noqa

{beanqueue-0.1.2 → beanqueue-0.1.3}/bq/cmds/process.py

@@ -6,7 +6,6 @@ import sys
 import threading
 import time
 import typing
-import uuid
 
 import click
 from dependency_injector.wiring import inject
@@ -14,6 +13,7 @@ from dependency_injector.wiring import Provide
 from sqlalchemy import func
 from sqlalchemy.orm import Session as DBSession
 
+from .. import constants
 from .. import models
 from ..config import Config
 from ..container import Container
@@ -22,27 +22,32 @@ from ..services.dispatch import DispatchService
 from ..services.worker import WorkerService
 
 
+@inject
 def update_workers(
-    make_session: typing.Callable[[], DBSession],
-    worker_id: uuid.UUID,
-    heartbeat_period: int,
-    heartbeat_timeout: int,
+    worker_id: typing.Any,
+    config: Config = Provide[Container.config],
+    session_factory: typing.Callable = Provide[Container.session_factory],
+    make_dispatch_service: typing.Callable = Provide[Container.make_dispatch_service],
+    make_worker_service: typing.Callable = Provide[Container.make_worker_service],
 ):
-    db: DBSession = make_session()
-    worker_service = WorkerService(session=db)
-    dispatch_service = DispatchService(session=db)
-    current_worker = db.get(models.Worker, worker_id)
+    db: DBSession = session_factory()
+    worker_service: WorkerService = make_worker_service(session=db)
+    dispatch_service: DispatchService = make_dispatch_service(session=db)
+    current_worker = worker_service.get_worker(worker_id)
     logger = logging.getLogger(__name__)
     logger.info(
         "Updating worker %s with heartbeat_period=%s, heartbeat_timeout=%s",
         current_worker.id,
-        heartbeat_period,
-        heartbeat_timeout,
+        config.WORKER_HEARTBEAT_PERIOD,
+        config.WORKER_HEARTBEAT_TIMEOUT,
     )
     while True:
-        dead_workers = worker_service.fetch_dead_workers(timeout=heartbeat_timeout)
+        dead_workers = worker_service.fetch_dead_workers(
+            timeout=config.WORKER_HEARTBEAT_TIMEOUT
+        )
         task_count = worker_service.reschedule_dead_tasks(
-            dead_workers.with_entities(models.Worker.id)
+            # TODO: a better way to abstract this?
+            dead_workers.with_entities(current_worker.__class__.id)
         )
         found_dead_worker = False
         for dead_worker in dead_workers:
@@ -58,7 +63,16 @@ def update_workers(
         if found_dead_worker:
             db.commit()
 
-        time.sleep(heartbeat_period)
+        if current_worker.state != models.WorkerState.RUNNING:
+            # This probably means we are somehow very slow to update the heartbeat in time, or the timeout window
+            # is set too short. It could also be the administrator update the worker state to something else than
+            # RUNNING. Regardless the reason, let's stop processing.
+            logger.warning(
+                "Current worker %s state is %s instead of running, quit processing"
+            )
+            sys.exit(0)
+
+        time.sleep(config.WORKER_HEARTBEAT_PERIOD)
         current_worker.last_heartbeat = func.now()
         db.add(current_worker)
         db.commit()
@@ -68,7 +82,6 @@ def update_workers(
 def process_tasks(
     channels: tuple[str, ...],
     config: Config = Provide[Container.config],
-    session_factory: typing.Callable = Provide[Container.session_factory],
     db: DBSession = Provide[Container.session],
     dispatch_service: DispatchService = Provide[Container.dispatch_service],
     worker_service: WorkerService = Provide[Container.worker_service],
@@ -76,7 +89,7 @@ def process_tasks(
     logger = logging.getLogger(__name__)
 
     if not channels:
-        channels = ["default"]
+        channels = [constants.DEFAULT_CHANNEL]
 
     if not config.PROCESSOR_PACKAGES:
         logger.error("No PROCESSOR_PACKAGES provided")
@@ -93,7 +106,7 @@ def process_tasks(
                     "  Processor module %r, processor %r", module, processor.name
                 )
 
-    worker = models.Worker(name=platform.node(), channels=channels)
+    worker = worker_service.make_worker(name=platform.node(), channels=channels)
     db.add(worker)
     dispatch_service.listen(channels)
     db.commit()
@@ -104,10 +117,7 @@ def process_tasks(
     worker_update_thread = threading.Thread(
         target=functools.partial(
             update_workers,
-            make_session=session_factory,
             worker_id=worker.id,
-            heartbeat_period=config.WORKER_HEARTBEAT_PERIOD,
-            heartbeat_timeout=config.WORKER_HEARTBEAT_TIMEOUT,
         ),
         name="update_workers",
     )

{beanqueue-0.1.2 → beanqueue-0.1.3}/bq/config.py

@@ -3,6 +3,7 @@ import typing
 from pydantic import field_validator
 from pydantic import PostgresDsn
 from pydantic import ValidationInfo
+from pydantic_core import MultiHostUrl
 from pydantic_settings import BaseSettings
 from pydantic_settings import SettingsConfigDict
 
@@ -23,6 +24,12 @@ class Config(BaseSettings):
     # Timeout of worker heartbeat in seconds
     WORKER_HEARTBEAT_TIMEOUT: int = 100
 
+    # which task model to use
+    TASK_MODEL: str = "bq.Task"
+
+    # which worker model to use
+    WORKER_MODEL: str = "bq.Worker"
+
     POSTGRES_SERVER: str = "localhost"
     POSTGRES_USER: str = "bq"
     POSTGRES_PASSWORD: str = ""
@@ -36,6 +43,8 @@ class Config(BaseSettings):
     ) -> typing.Any:
         if isinstance(v, str):
             return v
+        if isinstance(v, MultiHostUrl):
+            return v
         return PostgresDsn.build(
             scheme="postgresql",
             username=info.data.get("POSTGRES_USER"),

beanqueue-0.1.3/bq/constants.py

@@ -0,0 +1,4 @@
+# the name of default channel to use if not provided
+DEFAULT_CHANNEL = "default"
+# category value for venusian to scan functions decorated with `processor`
+BQ_PROCESSOR_CATEGORY = "bq_processor"

{beanqueue-0.1.2 → beanqueue-0.1.3}/bq/container.py

@@ -1,4 +1,5 @@
 import functools
+import importlib
 import typing
 
 from dependency_injector import containers
@@ -14,6 +15,12 @@ from .services.dispatch import DispatchService
 from .services.worker import WorkerService
 
 
+def get_model_class(name: str) -> typing.Type:
+    module_name, model_name = name.rsplit(".", 1)
+    module = importlib.import_module(module_name)
+    return getattr(module, model_name)
+
+
 def make_db_engine(config: Config) -> Engine:
     return create_engine(str(config.DATABASE_URL), poolclass=SingletonThreadPool)
 
@@ -26,12 +33,16 @@ def make_session(factory: typing.Callable) -> DBSession:
     return factory()
 
 
-def make_dispatch_service(session: DBSession) -> DispatchService:
-    return DispatchService(session)
+def make_dispatch_service(config: Config, session: DBSession) -> DispatchService:
+    return DispatchService(session, task_model=get_model_class(config.TASK_MODEL))
 
 
-def make_worker_service(session: DBSession) -> WorkerService:
-    return WorkerService(session)
+def make_worker_service(config: Config, session: DBSession) -> WorkerService:
+    return WorkerService(
+        session,
+        task_model=get_model_class(config.TASK_MODEL),
+        worker_model=get_model_class(config.WORKER_MODEL),
+    )
 
 
 class Container(containers.DeclarativeContainer):
@@ -46,9 +57,21 @@ class Container(containers.DeclarativeContainer):
     session: DBSession = providers.Singleton(make_session, factory=session_factory)
 
     dispatch_service: DispatchService = providers.Singleton(
-        make_dispatch_service, session=session
+        make_dispatch_service,
+        config=config,
+        session=session,
     )
 
     worker_service: WorkerService = providers.Singleton(
-        make_worker_service, session=session
+        make_worker_service, config=config, session=session
+    )
+
+    make_dispatch_service = providers.Singleton(
+        lambda config: functools.partial(make_dispatch_service, config=config),
+        config=config,
+    )
+
+    make_worker_service = providers.Singleton(
+        lambda config: functools.partial(make_worker_service, config=config),
+        config=config,
     )

beanqueue-0.1.3/bq/models/__init__.py

@@ -0,0 +1,8 @@
+from .task import Task
+from .task import TaskModelMixin
+from .task import TaskModelRefWorkerMixin
+from .task import TaskState
+from .worker import Worker
+from .worker import WorkerModelMixin
+from .worker import WorkerRefMixin
+from .worker import WorkerState

{beanqueue-0.1.2 → beanqueue-0.1.3}/bq/models/task.py

@@ -1,6 +1,8 @@
+import datetime
 import enum
+import typing
+import uuid
 
-from sqlalchemy import Column
 from sqlalchemy import Connection
 from sqlalchemy import DateTime
 from sqlalchemy import Enum
@@ -11,6 +13,9 @@ from sqlalchemy import inspect
 from sqlalchemy import String
 from sqlalchemy.dialects.postgresql import JSONB
 from sqlalchemy.dialects.postgresql import UUID
+from sqlalchemy.orm import declared_attr
+from sqlalchemy.orm import Mapped
+from sqlalchemy.orm import mapped_column
 from sqlalchemy.orm import Mapper
 from sqlalchemy.orm import relationship
 
@@ -29,18 +34,12 @@ class TaskState(enum.Enum):
     FAILED = "FAILED"
 
 
-class Task(Base):
-    id = Column(
+class TaskModelMixin:
+    id: Mapped[uuid.UUID] = mapped_column(
         UUID(as_uuid=True), primary_key=True, server_default=func.gen_random_uuid()
     )
-    # foreign key id of assigned worker
-    worker_id = Column(
-        UUID(as_uuid=True),
-        ForeignKey("bq_workers.id", name="fk_workers_id"),
-        nullable=True,
-    )
     # current state of the task
-    state = Column(
+    state: Mapped[TaskState] = mapped_column(
         Enum(TaskState),
         nullable=False,
         default=TaskState.PENDING,
@@ -48,24 +47,37 @@ class Task(Base):
         index=True,
     )
     # channel for workers and job creator to listen/notify
-    channel = Column(String, nullable=False, index=True)
+    channel: Mapped[str] = mapped_column(String, nullable=False, index=True)
     # module of the processor function
-    module = Column(String, nullable=False)
+    module: Mapped[str] = mapped_column(String, nullable=False)
     # func name of the processor func
-    func_name = Column(String, nullable=False)
+    func_name: Mapped[str] = mapped_column(String, nullable=False)
     # keyword arguments
-    kwargs = Column(JSONB, nullable=True)
+    kwargs: Mapped[typing.Optional[typing.Any]] = mapped_column(JSONB, nullable=True)
     # Result of the task
-    result = Column(JSONB, nullable=True)
+    result: Mapped[typing.Optional[typing.Any]] = mapped_column(JSONB, nullable=True)
     # Error message
-    error_message = Column(String, nullable=True)
+    error_message: Mapped[typing.Optional[str]] = mapped_column(String, nullable=True)
     # created datetime of the task
-    created_at = Column(
+    created_at: Mapped[datetime.datetime] = mapped_column(
         DateTime(timezone=True), nullable=False, server_default=func.now()
     )
 
-    worker = relationship("Worker", back_populates="tasks", uselist=False)
 
+class TaskModelRefWorkerMixin:
+    # foreign key id of assigned worker
+    worker_id: Mapped[uuid.UUID] = mapped_column(
+        UUID(as_uuid=True),
+        ForeignKey("bq_workers.id", name="fk_workers_id"),
+        nullable=True,
+    )
+
+    @declared_attr
+    def worker(cls) -> Mapped["Worker"]:
+        return relationship("Worker", back_populates="tasks", uselist=False)
+
+
+class Task(TaskModelMixin, TaskModelRefWorkerMixin, Base):
     __tablename__ = "bq_tasks"
 
     def __repr__(self) -> str:
@@ -99,22 +111,24 @@ def notify_if_needed(connection: Connection, task: Task):
     connection.exec_driver_sql(f"NOTIFY {quoted_channel}")
 
 
-@event.listens_for(Task, "after_insert")
 def task_insert_notify(mapper: Mapper, connection: Connection, target: Task):
-    from .. import models
-
-    if target.state != models.TaskState.PENDING:
+    if target.state != TaskState.PENDING:
         return
     notify_if_needed(connection, target)
 
 
-@event.listens_for(Task, "after_update")
 def task_update_notify(mapper: Mapper, connection: Connection, target: Task):
-    from .. import models
-
     history = inspect(target).attrs.state.history
     if not history.has_changes():
         return
-    if target.state != models.TaskState.PENDING:
+    if target.state != TaskState.PENDING:
         return
     notify_if_needed(connection, target)
+
+
+def listen_events(model_cls: typing.Type):
+    event.listens_for(model_cls, "after_insert")(task_insert_notify)
+    event.listens_for(model_cls, "after_update")(task_update_notify)
+
+
+listen_events(Task)

{beanqueue-0.1.2 → beanqueue-0.1.3}/bq/models/worker.py

@@ -1,4 +1,6 @@
+import datetime
 import enum
+import uuid
 
 from sqlalchemy import Column
 from sqlalchemy import DateTime
@@ -7,6 +9,10 @@ from sqlalchemy import func
 from sqlalchemy import String
 from sqlalchemy.dialects.postgresql import ARRAY
 from sqlalchemy.dialects.postgresql import UUID
+from sqlalchemy.orm import declared_attr
+from sqlalchemy.orm import Mapped
+from sqlalchemy.orm import mapped_column
+from sqlalchemy.orm import Mapper
 from sqlalchemy.orm import relationship
 
 from ..db.base import Base
@@ -22,12 +28,12 @@ class WorkerState(enum.Enum):
     NO_HEARTBEAT = "NO_HEARTBEAT"
 
 
-class Worker(Base):
-    id = Column(
+class WorkerModelMixin:
+    id: Mapped[uuid.UUID] = mapped_column(
         UUID(as_uuid=True), primary_key=True, server_default=func.gen_random_uuid()
     )
     # current state of the worker
-    state = Column(
+    state: Mapped[WorkerState] = mapped_column(
         Enum(WorkerState),
         nullable=False,
         default=WorkerState.RUNNING,
@@ -35,28 +41,34 @@ class Worker(Base):
         index=True,
     )
     # name of the worker
-    name = Column(String, nullable=False)
+    name: Mapped[str] = mapped_column(String, nullable=False)
     # the channels we are processing
-    channels = Column(ARRAY(String), nullable=False)
+    channels: Mapped[list[str]] = mapped_column(ARRAY(String), nullable=False)
     # last heartbeat of this worker
-    last_heartbeat = Column(
+    last_heartbeat: Mapped[datetime.datetime] = mapped_column(
         DateTime(timezone=True),
         nullable=False,
         server_default=func.now(),
         index=True,
     )
     # created datetime of the worker
-    created_at = Column(
+    created_at: Mapped[datetime.datetime] = mapped_column(
         DateTime(timezone=True), nullable=False, server_default=func.now()
     )
 
-    tasks = relationship(
-        "Task",
-        back_populates="worker",
-        cascade="all,delete",
-        order_by="Task.created_at",
-    )
 
+class WorkerRefMixin:
+    @declared_attr
+    def tasks(cls) -> Mapped[list["Task"]]:
+        return relationship(
+            "Task",
+            back_populates="worker",
+            cascade="all,delete",
+            order_by="Task.created_at",
+        )
+
+
+class Worker(WorkerModelMixin, WorkerRefMixin, Base):
     __tablename__ = "bq_workers"
 
     def __repr__(self) -> str:

{beanqueue-0.1.2 → beanqueue-0.1.3}/bq/processors/registry.py

@@ -7,9 +7,8 @@ import typing
 import venusian
 from sqlalchemy.orm import object_session
 
-from bq import models
-
-BQ_PROCESSOR_CATEGORY = "bq_processor"
+from .. import constants
+from .. import models
 
 
 @dataclasses.dataclass(frozen=True)
@@ -51,9 +50,10 @@ def process_task(task: models.Task, processor: Processor):
     if "db" in func_signature.parameters:
         base_kwargs["db"] = db
     with db.begin_nested() as savepoint:
+        if "savepoint" in func_signature.parameters:
+            base_kwargs["savepoint"] = savepoint
         try:
             result = processor.func(**base_kwargs, **task.kwargs)
-            savepoint.commit()
         except Exception as exc:
             logger.error("Unhandled exception for task %s", task.id, exc_info=True)
             if processor.auto_rollback_on_exc:
@@ -100,7 +100,7 @@ class Registry:
 
 
 def processor(
-    channel: str,
+    channel: str = constants.DEFAULT_CHANNEL,
     auto_complete: bool = True,
     auto_rollback_on_exc: bool = True,
     task_cls: typing.Type = models.Task,
@@ -121,7 +121,7 @@ def processor(
                 raise ValueError("Name is not the same")
             scanner.registry.add(processor)
 
-        venusian.attach(helper_obj, callback, category=BQ_PROCESSOR_CATEGORY)
+        venusian.attach(helper_obj, callback, category=constants.BQ_PROCESSOR_CATEGORY)
         return helper_obj
 
     return decorator
@@ -132,5 +132,5 @@ def collect(packages: list[typing.Any], registry: Registry | None = None) -> Reg
         registry = Registry()
     scanner = venusian.Scanner(registry=registry)
     for package in packages:
-        scanner.scan(package, categories=(BQ_PROCESSOR_CATEGORY,))
+        scanner.scan(package, categories=(constants.BQ_PROCESSOR_CATEGORY,))
     return registry

{beanqueue-0.1.2 → beanqueue-0.1.3}/bq/services/dispatch.py

@@ -17,28 +17,29 @@ class Notification:
 
 
 class DispatchService:
-    def __init__(self, session: Session):
+    def __init__(self, session: Session, task_model: typing.Type = models.Task):
         self.session = session
+        self.task_model: typing.Type[models.Task] = task_model
 
     def make_task_query(self, channels: typing.Sequence[str], limit: int = 1) -> Query:
         return (
-            self.session.query(models.Task.id)
-            .filter(models.Task.channel.in_(channels))
-            .filter(models.Task.state == models.TaskState.PENDING)
-            .order_by(models.Task.created_at)
+            self.session.query(self.task_model.id)
+            .filter(self.task_model.channel.in_(channels))
+            .filter(self.task_model.state == models.TaskState.PENDING)
+            .order_by(self.task_model.created_at)
             .limit(limit)
             .with_for_update(skip_locked=True)
         )
 
-    def make_update_query(self, task_query: typing.Any, worker_id: uuid.UUID):
+    def make_update_query(self, task_query: typing.Any, worker_id: typing.Any):
         return (
-            models.Task.__table__.update()
-            .where(models.Task.id.in_(task_query))
+            self.task_model.__table__.update()
+            .where(self.task_model.id.in_(task_query))
             .values(
                 state=models.TaskState.PROCESSING,
                 worker_id=worker_id,
             )
-            .returning(models.Task.id)
+            .returning(self.task_model.id)
         )
 
     def dispatch(
@@ -52,9 +53,11 @@ class DispatchService:
                 self.make_update_query(task_subquery, worker_id=worker_id)
             )
         ]
-        # TODO: ideally returning with (models.Task) should return the whole model, but SQLAlchemy is returning
+        # TODO: ideally returning with (self.task_model) should return the whole model, but SQLAlchemy is returning
         #       it columns in rows. We can save a round trip if we can find out how to solve this
-        return self.session.query(models.Task).filter(models.Task.id.in_(task_ids))
+        return self.session.query(self.task_model).filter(
+            self.task_model.id.in_(task_ids)
+        )
 
     def listen(self, channels: typing.Sequence[str]):
         conn = self.session.connection()

{beanqueue-0.1.2 → beanqueue-0.1.3}/bq/services/worker.py

@@ -9,8 +9,21 @@ from .. import models
 
 
 class WorkerService:
-    def __init__(self, session: Session):
+    def __init__(
+        self,
+        session: Session,
+        task_model: typing.Type = models.Task,
+        worker_model: typing.Type = models.Worker,
+    ):
         self.session = session
+        self.task_model: typing.Type[models.Task] = task_model
+        self.worker_model: typing.Type[models.Worker] = worker_model
+
+    def get_worker(self, id: typing.Any) -> typing.Any:
+        return self.session.get(self.worker_model, id)
+
+    def make_worker(self, name: str, channels: tuple[str, ...]):
+        return self.worker_model(name=name, channels=channels)
 
     def update_heartbeat(self, worker: models.Worker):
         worker.last_heartbeat = func.now()
@@ -18,24 +31,24 @@ class WorkerService:
 
     def make_dead_worker_query(self, timeout: int, limit: int = 5) -> Query:
         return (
-            self.session.query(models.Worker.id)
+            self.session.query(self.worker_model.id)
             .filter(
-                models.Worker.last_heartbeat
+                self.worker_model.last_heartbeat
                 < (func.now() - datetime.timedelta(seconds=timeout))
             )
-            .filter(models.Worker.state == models.WorkerState.RUNNING)
+            .filter(self.worker_model.state == models.WorkerState.RUNNING)
             .limit(limit)
             .with_for_update(skip_locked=True)
         )
 
     def make_update_dead_worker_query(self, worker_query: typing.Any):
         return (
-            models.Worker.__table__.update()
-            .where(models.Worker.id.in_(worker_query))
+            self.worker_model.__table__.update()
+            .where(self.worker_model.id.in_(worker_query))
             .values(
                 state=models.WorkerState.NO_HEARTBEAT,
             )
-            .returning(models.Worker.id)
+            .returning(self.worker_model.id)
         )
 
     def fetch_dead_workers(self, timeout: int, limit: int = 5) -> Query:
@@ -49,17 +62,18 @@ class WorkerService:
         ]
         # TODO: ideally returning with (models.Task) should return the whole model, but SQLAlchemy is returning
         #       it columns in rows. We can save a round trip if we can find out how to solve this
-        return self.session.query(models.Worker).filter(
-            models.Worker.id.in_(worker_ids)
+        return self.session.query(self.worker_model).filter(
+            self.worker_model.id.in_(worker_ids)
         )
 
     def make_update_tasks_query(self, worker_query: typing.Any):
         return (
-            models.Task.__table__.update()
-            .where(models.Task.worker_id.in_(worker_query))
-            .where(models.Task.state == models.TaskState.PROCESSING)
+            self.task_model.__table__.update()
+            .where(self.task_model.worker_id.in_(worker_query))
+            .where(self.task_model.state == models.TaskState.PROCESSING)
             .values(
                 state=models.TaskState.PENDING,
+                worker_id=None,
             )
         )
 

{beanqueue-0.1.2 → beanqueue-0.1.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "beanqueue"
-version = "0.1.2"
+version = "0.1.3"
 description = "BeanQueue or BQ for short, PostgreSQL SKIP LOCK based worker queue library"
 authors = ["Fang-Pen Lin <fangpen@launchplatform.com>"]
 license = "MIT"

beanqueue-0.1.2/bq/models/__init__.py

@@ -1,4 +0,0 @@
-from .task import Task
-from .task import TaskState
-from .worker import Worker
-from .worker import WorkerState

{beanqueue-0.1.2 → beanqueue-0.1.3}/bq/cmds/create_tables.py

@@ -13,13 +13,13 @@ from ..db.base import Base
 @click.command()
 @inject
 def main(engine: Engine = Provide[Container.db_engine]):
-    logging.basicConfig(level=logging.INFO)
     logger = logging.getLogger(__name__)
     Base.metadata.create_all(bind=engine)
     logger.info("Done, tables created")
 
 
 if __name__ == "__main__":
+    logging.basicConfig(level=logging.INFO)
     container = Container()
     container.wire(modules=[__name__])
     main()

{beanqueue-0.1.2 → beanqueue-0.1.3}/bq/cmds/submit.py

@@ -25,7 +25,6 @@ def main(
     kwargs: str | None,
     db: Session = Provide[Container.session],
 ):
-    logging.basicConfig(level=logging.INFO)
     logger = logging.getLogger(__name__)
 
     logger.info(
@@ -43,6 +42,7 @@ def main(
 
 
 if __name__ == "__main__":
+    logging.basicConfig(level=logging.INFO)
     container = Container()
     container.wire(modules=[__name__])
     main()