beanqueue 0.1.0__py3-none-any.whl

beanqueue-0.1.0.dist-info/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Launch Platform
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

beanqueue-0.1.0.dist-info/METADATA

@@ -0,0 +1,23 @@
+Metadata-Version: 2.1
+Name: beanqueue
+Version: 0.1.0
+Summary: BeanQueue or BQ for short, PostgreSQL SKIP LOCK based worker queue library
+License: MIT
+Author: Fang-Pen Lin
+Author-email: fangpen@launchplatform.com
+Requires-Python: >=3.11,<4.0
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Requires-Dist: click (>=8.1.7,<9.0.0)
+Requires-Dist: dependency-injector (>=4.41.0,<5.0.0)
+Requires-Dist: pg-activity (>=3.5.1,<4.0.0)
+Requires-Dist: pydantic-settings (>=2.2.1,<3.0.0)
+Requires-Dist: sqlalchemy (>=2.0.30,<3.0.0)
+Requires-Dist: venusian (>=3.1.0,<4.0.0)
+Description-Content-Type: text/markdown
+
+# bq
+BeanQueue or BQ for short, PostgreSQL SKIP LOCK based worker queue library
+

beanqueue-0.1.0.dist-info/RECORD

@@ -0,0 +1,23 @@
+bq/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+bq/cmds/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+bq/cmds/create_tables.py,sha256=lI9GmovZT7-_ws4jq8xV-0pHStiyPZoE6Ygb0Xn1q2c,609
+bq/cmds/process.py,sha256=9CE2PCzdl3WAdx1whlcqmDwuQujeuh6rvrLjEAcXlAs,5879
+bq/cmds/submit.py,sha256=WhzitLawbR_hEjLqCuvr2yhKJcwGNplarntpAhjGBb4,1148
+bq/config.py,sha256=LUEwzodrt6H32bC2Fn8_J8oQQQ9dSFZnIg2Uv_EMLus,1508
+bq/container.py,sha256=NQLpc2zEap8LW2XkqRQjGtzyZslPL5iDnEZKBnH67eo,1558
+bq/db/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+bq/db/base.py,sha256=0LXS0WlLr1KAHGZng46SmUwJam8m2AxVDBqK1Yzwx8w,88
+bq/db/session.py,sha256=ife8ocHxCT7EbDgyEG38TFhP4l4wQv_Bo16bxt9XCos,157
+bq/models/__init__.py,sha256=bibCa4EgF9MtkmSSqa15C7Jc7tN9P6LS78YOBkvX72s,114
+bq/models/helpers.py,sha256=vf8IvREk2Lfc_LsdxWfUgVc99Fx1K2RsWVj9Xi4Dhbg,184
+bq/models/task.py,sha256=hVkYgt-TPFdKmbmL9E_AUbWiZlk0nkq7ZGWfkyoZICk,3604
+bq/models/worker.py,sha256=EuC8k6g9l_xddjXxhMqXmHt6uSO86VZFs9hJH1DQUi8,1940
+bq/processors/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+bq/processors/registry.py,sha256=sc9rfkWAJFFh5p8K-Vv1z2GCzgqoTLX3yNMJ3j4F15k,4554
+bq/services/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+bq/services/dispatch.py,sha256=uSssmqGKCxsr3ZzbglXs01XZNi3nv3zUfO2d88u3Wv0,3463
+bq/services/worker.py,sha256=-NDP-6dA40OYdw00sbS3jLqU2PSY62oL_1gPVK5ijMQ,2554
+beanqueue-0.1.0.dist-info/LICENSE,sha256=sNYpr-_bmGA6hBtD7qtmZmVgQH6HC6KXlx7tOOOjFq8,1072
+beanqueue-0.1.0.dist-info/METADATA,sha256=YRbok6yaosRGobhhsnQYWTiRouSeytUIBOSfV82jseY,834
+beanqueue-0.1.0.dist-info/WHEEL,sha256=FMvqSimYX_P7y0a7UY-_Mc83r5zkBZsCYPm7Lr0Bsq4,88
+beanqueue-0.1.0.dist-info/RECORD,,

beanqueue-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: poetry-core 1.8.1
+Root-Is-Purelib: true
+Tag: py3-none-any

bq/cmds/create_tables.py

@@ -0,0 +1,25 @@
+import logging
+
+import click
+from dependency_injector.wiring import inject
+from dependency_injector.wiring import Provide
+from sqlalchemy.engine import Engine
+
+from .. import models  # noqa
+from ..container import Container
+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__":
+    container = Container()
+    container.wire(modules=[__name__])
+    main()

bq/cmds/process.py

@@ -0,0 +1,178 @@
+import functools
+import importlib
+import logging
+import platform
+import sys
+import threading
+import time
+import typing
+import uuid
+
+import click
+from dependency_injector.wiring import inject
+from dependency_injector.wiring import Provide
+from sqlalchemy import func
+from sqlalchemy.orm import Session as DBSession
+
+from .. import models
+from ..config import Config
+from ..container import Container
+from ..processors.registry import collect
+from ..services.dispatch import DispatchService
+from ..services.worker import WorkerService
+
+
+def update_workers(
+    make_session: typing.Callable[[], DBSession],
+    worker_id: uuid.UUID,
+    heartbeat_period: int,
+    heartbeat_timeout: int,
+):
+    db: DBSession = make_session()
+    worker_service = WorkerService(session=db)
+    dispatch_service = DispatchService(session=db)
+    current_worker = db.get(models.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,
+    )
+    while True:
+        dead_workers = worker_service.fetch_dead_workers(timeout=heartbeat_timeout)
+        task_count = worker_service.reschedule_dead_tasks(
+            dead_workers.with_entities(models.Worker.id)
+        )
+        found_dead_worker = False
+        for dead_worker in dead_workers:
+            found_dead_worker = True
+            logger.info(
+                "Found dead worker %s (name=%s), reschedule %s dead tasks in channels %s",
+                dead_worker.id,
+                dead_worker.name,
+                task_count,
+                dead_worker.channels,
+            )
+            dispatch_service.notify(dead_worker.channels)
+        if found_dead_worker:
+            db.commit()
+
+        time.sleep(heartbeat_period)
+        current_worker.last_heartbeat = func.now()
+        db.add(current_worker)
+        db.commit()
+
+
+@inject
+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],
+):
+    logger = logging.getLogger(__name__)
+
+    if not channels:
+        channels = ["default"]
+
+    if not config.PROCESSOR_PACKAGES:
+        logger.error("No PROCESSOR_PACKAGES provided")
+        sys.exit(-1)
+
+    logger.info("Scanning packages %s", config.PROCESSOR_PACKAGES)
+    pkgs = list(map(importlib.import_module, config.PROCESSOR_PACKAGES))
+    registry = collect(pkgs)
+    for channel, module_processors in registry.processors.items():
+        logger.info("Collected processors with channel %r", channel)
+        for module, func_processors in module_processors.items():
+            for processor in func_processors.values():
+                logger.info(
+                    "  Processor module %r, processor %r", module, processor.name
+                )
+
+    worker = models.Worker(name=platform.node(), channels=channels)
+    db.add(worker)
+    dispatch_service.listen(channels)
+    db.commit()
+
+    logger.info("Created worker %s, name=%s", worker.id, worker.name)
+    logger.info("Processing tasks in channels = %s ...", channels)
+
+    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",
+    )
+    worker_update_thread.daemon = True
+    worker_update_thread.start()
+
+    worker_id = worker.id
+
+    try:
+        while True:
+            while True:
+                tasks = dispatch_service.dispatch(
+                    channels,
+                    worker_id=worker_id,
+                    limit=config.BATCH_SIZE,
+                ).all()
+                for task in tasks:
+                    logger.info(
+                        "Processing task %s, channel=%s, module=%s, func=%s",
+                        task.id,
+                        task.channel,
+                        task.module,
+                        task.func_name,
+                    )
+                    # TODO: support processor pool and other approaches to dispatch the workload
+                    registry.process(task)
+                if not tasks:
+                    # we should try to keep dispatching until we cannot find tasks
+                    break
+                else:
+                    db.commit()
+            # we will not see notifications in a transaction, need to close the transaction first before entering
+            # polling
+            db.close()
+            try:
+                for notification in dispatch_service.poll(timeout=config.POLL_TIMEOUT):
+                    logger.debug("Receive notification %s", notification)
+            except TimeoutError:
+                logger.debug("Poll timeout, try again")
+                continue
+    except (SystemExit, KeyboardInterrupt):
+        db.rollback()
+        logger.info("Shutting down ...")
+        worker_update_thread.join(5)
+
+    worker.state = models.WorkerState.SHUTDOWN
+    db.add(worker)
+    task_count = worker_service.reschedule_dead_tasks([worker.id])
+    logger.info("Reschedule %s tasks", task_count)
+    dispatch_service.notify(channels)
+    db.commit()
+
+    logger.info("Shutdown gracefully")
+
+
+@click.command()
+@click.argument("channels", nargs=-1)
+def main(
+    channels: tuple[str, ...],
+):
+    process_tasks(channels)
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.INFO)
+    container = Container()
+    container.wire(modules=[__name__])
+    main()

bq/cmds/submit.py

@@ -0,0 +1,48 @@
+import json
+import logging
+
+import click
+from dependency_injector.wiring import inject
+from dependency_injector.wiring import Provide
+
+from .. import models
+from ..container import Container
+from ..db.session import Session
+
+
+@click.command()
+@click.argument("channel", nargs=1)
+@click.argument("module", nargs=1)
+@click.argument("func", nargs=1)
+@click.option(
+    "-k", "--kwargs", type=str, help="Keyword arguments as JSON", default=None
+)
+@inject
+def main(
+    channel: str,
+    module: str,
+    func: str,
+    kwargs: str | None,
+    db: Session = Provide[Container.session],
+):
+    logging.basicConfig(level=logging.INFO)
+    logger = logging.getLogger(__name__)
+
+    logger.info(
+        "Submit task with channel=%s, module=%s, func=%s", channel, module, func
+    )
+    kwargs_value = {}
+    if kwargs:
+        kwargs_value = json.loads(kwargs)
+    task = models.Task(
+        channel=channel, module=module, func_name=func, kwargs=kwargs_value
+    )
+    db.add(task)
+    db.commit()
+    logger.info("Done, submit task %s", task.id)
+
+
+if __name__ == "__main__":
+    container = Container()
+    container.wire(modules=[__name__])
+    main()

bq/config.py

@@ -0,0 +1,47 @@
+import typing
+
+from pydantic import field_validator
+from pydantic import PostgresDsn
+from pydantic import ValidationInfo
+from pydantic_settings import BaseSettings
+from pydantic_settings import SettingsConfigDict
+
+
+class Config(BaseSettings):
+    # Packages to scan for processor functions
+    PROCESSOR_PACKAGES: list[str] = []
+
+    # Size of tasks batch to fetch each time from the database
+    BATCH_SIZE: int = 1
+
+    # How long we should poll before timeout in seconds
+    POLL_TIMEOUT: int = 60
+
+    # Interval of worker heartbeat update cycle in seconds
+    WORKER_HEARTBEAT_PERIOD: int = 30
+
+    # Timeout of worker heartbeat in seconds
+    WORKER_HEARTBEAT_TIMEOUT: int = 100
+
+    POSTGRES_SERVER: str = "localhost"
+    POSTGRES_USER: str = "bq"
+    POSTGRES_PASSWORD: str = ""
+    POSTGRES_DB: str = "bq"
+    # The URL of postgresql database to connect
+    DATABASE_URL: typing.Optional[PostgresDsn] = None
+
+    @field_validator("DATABASE_URL", mode="before")
+    def assemble_db_connection(
+        cls, v: typing.Optional[str], info: ValidationInfo
+    ) -> typing.Any:
+        if isinstance(v, str):
+            return v
+        return PostgresDsn.build(
+            scheme="postgresql",
+            username=info.data.get("POSTGRES_USER"),
+            password=info.data.get("POSTGRES_PASSWORD"),
+            host=info.data.get("POSTGRES_SERVER"),
+            path=f"{info.data.get('POSTGRES_DB') or ''}",
+        )
+
+    model_config = SettingsConfigDict(case_sensitive=True, env_prefix="BQ_")

bq/container.py

@@ -0,0 +1,54 @@
+import functools
+import typing
+
+from dependency_injector import containers
+from dependency_injector import providers
+from sqlalchemy import create_engine
+from sqlalchemy import Engine
+from sqlalchemy.orm import Session as DBSession
+from sqlalchemy.pool import SingletonThreadPool
+
+from .config import Config
+from .db.session import SessionMaker
+from .services.dispatch import DispatchService
+from .services.worker import WorkerService
+
+
+def make_db_engine(config: Config) -> Engine:
+    return create_engine(str(config.DATABASE_URL), poolclass=SingletonThreadPool)
+
+
+def make_session_factory(engine: Engine) -> typing.Callable:
+    return functools.partial(SessionMaker, bind=engine)
+
+
+def make_session(factory: typing.Callable) -> DBSession:
+    return factory()
+
+
+def make_dispatch_service(session: DBSession) -> DispatchService:
+    return DispatchService(session)
+
+
+def make_worker_service(session: DBSession) -> WorkerService:
+    return WorkerService(session)
+
+
+class Container(containers.DeclarativeContainer):
+    config = providers.Singleton(Config)
+
+    db_engine: Engine = providers.Singleton(make_db_engine, config=config)
+
+    session_factory: typing.Callable = providers.Singleton(
+        make_session_factory, engine=db_engine
+    )
+
+    session: DBSession = providers.Singleton(make_session, factory=session_factory)
+
+    dispatch_service: DispatchService = providers.Singleton(
+        make_dispatch_service, session=session
+    )
+
+    worker_service: WorkerService = providers.Singleton(
+        make_worker_service, session=session
+    )

bq/db/base.py

@@ -0,0 +1,5 @@
+from sqlalchemy.orm import DeclarativeBase
+
+
+class Base(DeclarativeBase):
+    pass

bq/db/session.py

@@ -0,0 +1,5 @@
+from sqlalchemy.orm import scoped_session
+from sqlalchemy.orm import sessionmaker
+
+SessionMaker = sessionmaker()
+Session = scoped_session(SessionMaker)

bq/models/__init__.py

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

bq/models/helpers.py

@@ -0,0 +1,5 @@
+import typing
+
+
+def make_repr_attrs(items: typing.Sequence[typing.Tuple[str, typing.Any]]) -> str:
+    return " ".join(map(lambda item: "=".join([item[0], str(item[1])]), items))

bq/models/task.py

@@ -0,0 +1,120 @@
+import enum
+
+from sqlalchemy import Column
+from sqlalchemy import Connection
+from sqlalchemy import DateTime
+from sqlalchemy import Enum
+from sqlalchemy import event
+from sqlalchemy import ForeignKey
+from sqlalchemy import func
+from sqlalchemy import inspect
+from sqlalchemy import String
+from sqlalchemy.dialects.postgresql import JSONB
+from sqlalchemy.dialects.postgresql import UUID
+from sqlalchemy.orm import Mapper
+from sqlalchemy.orm import relationship
+
+from ..db.base import Base
+from .helpers import make_repr_attrs
+
+
+class TaskState(enum.Enum):
+    # task just created, not scheduled yet
+    PENDING = "PENDING"
+    # a worker is processing the task right now
+    PROCESSING = "PROCESSING"
+    # the task is done
+    DONE = "DONE"
+    # the task is failed
+    FAILED = "FAILED"
+
+
+class Task(Base):
+    id = 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(
+        Enum(TaskState),
+        nullable=False,
+        default=TaskState.PENDING,
+        server_default=TaskState.PENDING.value,
+        index=True,
+    )
+    # channel for workers and job creator to listen/notify
+    channel = Column(String, nullable=False, index=True)
+    # module of the processor function
+    module = Column(String, nullable=False)
+    # func name of the processor func
+    func_name = Column(String, nullable=False)
+    # keyword arguments
+    kwargs = Column(JSONB, nullable=True)
+    # Result of the task
+    result = Column(JSONB, nullable=True)
+    # Error message
+    error_message = Column(String, nullable=True)
+    # created datetime of the task
+    created_at = Column(
+        DateTime(timezone=True), nullable=False, server_default=func.now()
+    )
+
+    worker = relationship("Worker", back_populates="tasks", uselist=False)
+
+    __tablename__ = "bq_tasks"
+
+    def __repr__(self) -> str:
+        items = [
+            ("id", self.id),
+            ("state", self.state),
+            ("channel", self.channel),
+        ]
+        return f"<{self.__class__.__name__} {make_repr_attrs(items)}>"
+
+
+def notify_if_needed(connection: Connection, task: Task):
+    session = inspect(task).session
+    transaction = session.get_transaction()
+    if transaction is not None:
+        key = "_notified_channels"
+        if hasattr(transaction, key):
+            notified_channels = getattr(transaction, key)
+        else:
+            notified_channels = set()
+            setattr(transaction, key, notified_channels)
+
+        if task.channel in notified_channels:
+            # already notified, skip
+            return
+        notified_channels.add(task.channel)
+
+    quoted_channel = connection.dialect.identifier_preparer.quote_identifier(
+        task.channel
+    )
+    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:
+        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:
+        return
+    notify_if_needed(connection, target)

bq/models/worker.py

@@ -0,0 +1,69 @@
+import enum
+
+from sqlalchemy import Column
+from sqlalchemy import DateTime
+from sqlalchemy import Enum
+from sqlalchemy import func
+from sqlalchemy import String
+from sqlalchemy.dialects.postgresql import ARRAY
+from sqlalchemy.dialects.postgresql import UUID
+from sqlalchemy.orm import relationship
+
+from ..db.base import Base
+from .helpers import make_repr_attrs
+
+
+class WorkerState(enum.Enum):
+    # the worker is running
+    RUNNING = "RUNNING"
+    # the worker shuts down normally
+    SHUTDOWN = "SHUTDOWN"
+    # The worker has no heartbeat for a while
+    NO_HEARTBEAT = "NO_HEARTBEAT"
+
+
+class Worker(Base):
+    id = Column(
+        UUID(as_uuid=True), primary_key=True, server_default=func.gen_random_uuid()
+    )
+    # current state of the worker
+    state = Column(
+        Enum(WorkerState),
+        nullable=False,
+        default=WorkerState.RUNNING,
+        server_default=WorkerState.RUNNING.value,
+        index=True,
+    )
+    # name of the worker
+    name = Column(String, nullable=False)
+    # the channels we are processing
+    channels = Column(ARRAY(String), nullable=False)
+    # last heartbeat of this worker
+    last_heartbeat = Column(
+        DateTime(timezone=True),
+        nullable=False,
+        server_default=func.now(),
+        index=True,
+    )
+    # created datetime of the worker
+    created_at = Column(
+        DateTime(timezone=True), nullable=False, server_default=func.now()
+    )
+
+    tasks = relationship(
+        "Task",
+        back_populates="worker",
+        cascade="all,delete",
+        order_by="Task.created_at",
+    )
+
+    __tablename__ = "bq_workers"
+
+    def __repr__(self) -> str:
+        items = [
+            ("id", self.id),
+            ("name", self.name),
+            ("channels", self.channels),
+            ("state", self.state),
+        ]
+        return f"<{self.__class__.__name__} {make_repr_attrs(items)}>"

bq/processors/registry.py

@@ -0,0 +1,136 @@
+import collections
+import dataclasses
+import inspect
+import logging
+import typing
+
+import venusian
+from sqlalchemy.orm import object_session
+
+from bq import models
+
+BQ_PROCESSOR_CATEGORY = "bq_processor"
+
+
+@dataclasses.dataclass(frozen=True)
+class Processor:
+    channel: str
+    module: str
+    name: str
+    func: typing.Callable
+    # should we auto complete the task or not
+    auto_complete: bool = True
+    # should we auto rollback the transaction when encounter unhandled exception
+    auto_rollback_on_exc: bool = True
+
+
+class ProcessorHelper:
+    def __init__(self, processor: Processor, task_cls: typing.Type = models.Task):
+        self._processor = processor
+        self._task_cls = task_cls
+
+    def __call__(self, *args, **kwargs):
+        return self._processor.func(*args, **kwargs)
+
+    def run(self, **kwargs) -> models.Task:
+        return self._task_cls(
+            channel=self._processor.channel,
+            module=self._processor.module,
+            func_name=self._processor.name,
+            kwargs=kwargs,
+        )
+
+
+def process_task(task: models.Task, processor: Processor):
+    logger = logging.getLogger(__name__)
+    db = object_session(task)
+    func_signature = inspect.signature(processor.func)
+    base_kwargs = {}
+    if "task" in func_signature.parameters:
+        base_kwargs["task"] = task
+    if "db" in func_signature.parameters:
+        base_kwargs["db"] = db
+    with db.begin_nested() as 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:
+                savepoint.rollback()
+            # TODO: add error event
+            task.state = models.TaskState.FAILED
+            task.error_message = str(exc)
+            db.add(task)
+            return
+    if processor.auto_complete:
+        logger.info("Task %s auto complete", task.id)
+        task.state = models.TaskState.DONE
+        task.result = result
+        db.add(task)
+    return result
+
+
+class Registry:
+    def __init__(self):
+        self.logger = logging.getLogger(__name__)
+        self.processors = collections.defaultdict(lambda: collections.defaultdict(dict))
+
+    def add(self, processor: Processor):
+        self.processors[processor.channel][processor.module][processor.name] = processor
+
+    def process(self, task: models.Task) -> typing.Any:
+        modules = self.processors.get(task.channel, {})
+        functions = modules.get(task.module, {})
+        processor = functions.get(task.func_name)
+        db = object_session(task)
+        if processor is None:
+            self.logger.error(
+                "Cannot find processor for task %s with module=%s, func=%s",
+                task.id,
+                task.module,
+                task.func_name,
+            )
+            # TODO: add error event
+            task.state = models.TaskState.FAILED
+            task.error_message = f"Cannot find processor for task with module={task.module}, func={task.func_name}"
+            db.add(task)
+            return
+        return process_task(task, processor)
+
+
+def processor(
+    channel: str,
+    auto_complete: bool = True,
+    auto_rollback_on_exc: bool = True,
+    task_cls: typing.Type = models.Task,
+) -> typing.Callable:
+    def decorator(wrapped: typing.Callable):
+        processor = Processor(
+            module=wrapped.__module__,
+            name=wrapped.__name__,
+            channel=channel,
+            func=wrapped,
+            auto_complete=auto_complete,
+            auto_rollback_on_exc=auto_rollback_on_exc,
+        )
+        helper_obj = ProcessorHelper(processor, task_cls=task_cls)
+
+        def callback(scanner: venusian.Scanner, name: str, ob: typing.Callable):
+            if processor.name != name:
+                raise ValueError("Name is not the same")
+            scanner.registry.add(processor)
+
+        venusian.attach(helper_obj, callback, category=BQ_PROCESSOR_CATEGORY)
+        return helper_obj
+
+    return decorator
+
+
+def collect(packages: list[typing.Any], registry: Registry | None = None) -> Registry:
+    if registry is None:
+        registry = Registry()
+    scanner = venusian.Scanner(registry=registry)
+    for package in packages:
+        scanner.scan(package, categories=(BQ_PROCESSOR_CATEGORY,))
+    return registry

bq/services/dispatch.py

@@ -0,0 +1,96 @@
+import dataclasses
+import select
+import typing
+import uuid
+
+from sqlalchemy.orm import Query
+
+from .. import models
+from ..db.session import Session
+
+
+@dataclasses.dataclass(frozen=True)
+class Notification:
+    pid: int
+    channel: str
+    payload: typing.Optional[str] = None
+
+
+class DispatchService:
+    def __init__(self, session: Session):
+        self.session = session
+
+    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)
+            .limit(limit)
+            .with_for_update(skip_locked=True)
+        )
+
+    def make_update_query(self, task_query: typing.Any, worker_id: uuid.UUID):
+        return (
+            models.Task.__table__.update()
+            .where(models.Task.id.in_(task_query))
+            .values(
+                state=models.TaskState.PROCESSING,
+                worker_id=worker_id,
+            )
+            .returning(models.Task.id)
+        )
+
+    def dispatch(
+        self, channels: typing.Sequence[str], worker_id: uuid.UUID, limit: int = 1
+    ) -> Query:
+        task_query = self.make_task_query(channels, limit=limit)
+        task_subquery = task_query.scalar_subquery()
+        task_ids = [
+            item[0]
+            for item in self.session.execute(
+                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
+        #       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))
+
+    def listen(self, channels: typing.Sequence[str]):
+        conn = self.session.connection()
+        for channel in channels:
+            quoted_channel = conn.dialect.identifier_preparer.quote_identifier(channel)
+            conn.exec_driver_sql(f"LISTEN {quoted_channel}")
+
+    def poll(self, timeout: int = 5) -> typing.Generator[Notification, None, None]:
+        conn = self.session.connection()
+        driver_conn = conn.connection.driver_connection
+
+        def pop_notifies():
+            while driver_conn.notifies:
+                notify = driver_conn.notifies.pop(0)
+                yield Notification(
+                    pid=notify.pid,
+                    channel=notify.channel,
+                    payload=notify.payload,
+                )
+
+        # poll first to see if there's anything already
+        driver_conn.poll()
+        if driver_conn.notifies:
+            yield from pop_notifies()
+        else:
+            # okay, nothing, let's select and wait for new stuff
+            if select.select([driver_conn], [], [], timeout) == ([], [], []):
+                # nope, nothing, times out
+                raise TimeoutError("Timeout waiting for new notifications")
+            else:
+                # yep, we got something
+                driver_conn.poll()
+                yield from pop_notifies()
+
+    def notify(self, channels: typing.Sequence[str]):
+        conn = self.session.connection()
+        for channel in channels:
+            quoted_channel = conn.dialect.identifier_preparer.quote_identifier(channel)
+            conn.exec_driver_sql(f"NOTIFY {quoted_channel}")

bq/services/worker.py

@@ -0,0 +1,69 @@
+import datetime
+import typing
+
+from sqlalchemy import func
+from sqlalchemy.orm import Query
+from sqlalchemy.orm import Session
+
+from .. import models
+
+
+class WorkerService:
+    def __init__(self, session: Session):
+        self.session = session
+
+    def update_heartbeat(self, worker: models.Worker):
+        worker.last_heartbeat = func.now()
+        self.session.add(worker)
+
+    def make_dead_worker_query(self, timeout: int, limit: int = 5) -> Query:
+        return (
+            self.session.query(models.Worker.id)
+            .filter(
+                models.Worker.last_heartbeat
+                < (func.now() - datetime.timedelta(seconds=timeout))
+            )
+            .filter(models.Worker.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))
+            .values(
+                state=models.WorkerState.NO_HEARTBEAT,
+            )
+            .returning(models.Worker.id)
+        )
+
+    def fetch_dead_workers(self, timeout: int, limit: int = 5) -> Query:
+        dead_worker_query = self.make_dead_worker_query(timeout=timeout, limit=limit)
+        dead_worker_subquery = dead_worker_query.scalar_subquery()
+        worker_ids = [
+            item[0]
+            for item in self.session.execute(
+                self.make_update_dead_worker_query(dead_worker_subquery)
+            )
+        ]
+        # 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)
+        )
+
+    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)
+            .values(
+                state=models.TaskState.PENDING,
+            )
+        )
+
+    def reschedule_dead_tasks(self, worker_query: typing.Any) -> int:
+        update_dead_task_query = self.make_update_tasks_query(worker_query=worker_query)
+        res = self.session.execute(update_dead_task_query)
+        return res.rowcount