Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
89 changes: 86 additions & 3 deletions python/fluxqueue/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,9 +32,8 @@ def task(
"""
Mark a function as a FluxQueue task.

This returns a decorator. When you apply it to a function, calling that
function will enqueue a task in Redis instead of running the function
immediately. The actual work is done later by the worker.
When you apply to a function that function is being marked as a task function.
Running it will enqueue the task and then the worker will execute it.

Parameters
----------
Expand All @@ -48,6 +47,28 @@ def task(
`max_retries`:
Maximum number of retries the worker will attempt for this task
before treating it as dead.

Example
---

```py
@fluxqueue.task()
async def send_email_task(name: str, username: str, email: str):
async with get_email_client() as client:
await send_email(
email_client=client,
to_email=email,
subject="Welcome to FluxQueue",
config=email_config,
)
```

Enqueueing the task
---

```py
await send_email_task(name, username, email)
```
"""

@overload
Expand Down Expand Up @@ -78,6 +99,68 @@ def task_with_context(
queue: str = "default",
max_retries: int = 3,
):
"""
Mark a function as a FluxQueue task with context.

This decorator works like the `task` decorator but adds support for the `Context` class.
The function must accept a context as its first argument, with `Context` (or a subclass of `Context`) as the type hint.
When decorated, the context argument is automatically injected by the worker and is no longer
part of the function's public signature - users calling the function do not need to provide it.

Parameters
----------
`name`:
Optional explicit task name. If not set, a name is derived from the
function name.

`queue`:
Name of the queue to push tasks to. Defaults to `"default"`.

`max_retries`:
Maximum number of retries the worker will attempt for this task
before treating it as dead.

Example
---

```py
class DbContext(Context):
def __init__(self) -> None:
super().__init__()

def _get_local_session(self) -> async_sessionmaker[AsyncSession]:
if not self.thread_storage.get("session"):
engine = create_async_engine(SQLALCHEMY_DATABASE_URL)
self.thread_storage["session"] = async_sessionmaker(
bind=engine, expire_on_commit=False
)

return self.thread_storage["session"]

@asynccontextmanager
async def session_context(self):
local_session = self._get_local_session()
async with local_session() as session:
try:
yield session
await session.commit()
except Exception:
await session.rollback()
raise

@fluxqueue.task_with_context()
async def create_user_task(ctx: DbContext, email: str, username: str):
async with ctx.session_context() as db_session:
user = User(
email=email,
username=username
)
db_session.add(user)

await create_user_task
```
"""

@overload
def decorator(func: Callable[Concatenate[C, P], None]) -> Callable[P, None]: ...

Expand Down
17 changes: 17 additions & 0 deletions python/fluxqueue/context.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,23 @@


class Context:
"""
Base execution context for FluxQueue tasks.

Provides a dual-layer storage system designed for high-performance
distributed task execution:

- Worker Layer (`thread_storage`): Persists across the lifetime of an
individual worker thread. Use this for heavy resource pooling
(e.g., DB engines, HTTP clients) to avoid re-initialization overhead.
- Task Layer (`metadata`): Isolated to a single task execution via
ContextVars. Provides read-only access to the current task's
unique identity and execution state.

This class can be used directly for basic metadata access or subclassed
to provide domain-specific resources.
"""

__fluxqueue_context__: str | None = None

def __init__(self) -> None:
Expand Down