Context Management

Overview

The Context Management module is the backbone for multi-tenancy, distributed tracing, and enriched logging throughout the Athomic Layer. It provides a safe and reliable way to carry request-scoped state through your application without passing arguments down the call stack.

It is built on top of Python's native contextvars, making it fully compatible with asyncio and ensuring that context is never accidentally shared between concurrent requests or tasks.

Core Use Cases

Multi-Tenancy: Automatically isolates data by carrying the tenant_id from the initial request to the database layer.
Distributed Tracing: Propagates trace_id and span_id across function calls and even to background tasks.
Enriched Logging: Allows log records to be automatically enriched with contextual data like request_id and user_id.
Consistent Keying: Provides a ContextKeyGenerator for creating standardized keys for caches, locks, and rate limiters.

Core Components

`ContextVarManager`

This is the central registry that manages the lifecycle of all context variables. It handles their creation, default values, and tracks which variables should be propagated to background tasks.

`context_vars` Module

This module defines all the context variables available in the application (e.g., tenant_id, request_id, user_id, locale) and provides simple getter and setter functions for each one. This is the primary API you will interact with.

# Example of using the getters and setters
from nala.athomic.context import context_vars

# Set a value (typically in a middleware)
token = context_vars.set_tenant_id("tenant-123")

# Get a value anywhere in the call stack
current_tenant = context_vars.get_tenant_id() # Returns "tenant-123"

# Reset the value (restores the previous state)
context_vars.get_tenant_id().reset(token)

`ExecutionContext`

This is a simple data class that captures a snapshot of all current context variables at a specific moment in time. Its main purpose is to facilitate context propagation.

Context Propagation

One of the biggest challenges in a microservices architecture is maintaining context across different processes, especially when dispatching background tasks. This module provides a simple and effective solution.

capture_context(): Before a background task is sent to a broker, this function is called. It reads all context variables marked for propagation and returns them in a serializable dictionary.
restore_context(): On the worker side, this function is used as a context manager (with restore_context(...)). It takes the dictionary captured in step 1 and temporarily applies those values to the contextvars for the duration of the task execution.

This ensures that a background task for sending an email, for example, will have the same trace_id and tenant_id as the API request that triggered it.

Example

# In your API endpoint
from nala.athomic.context import capture_context
from my_app.tasks import send_welcome_email_task

async def create_user(user_data: dict):
    # ... create user ...

    # Capture the current context before enqueuing the task
    context_to_propagate = capture_context()

    # Pass the context in the task's keyword arguments
    await send_welcome_email_task.delay(
        user_id=user.id,
        _nala_context=context_to_propagate
    )

# In your worker/task definition
from nala.athomic.context import restore_context
from nala.athomic.integration.tasks import task

@task()
async def send_welcome_email_task(user_id: int, _nala_context: dict):
    # The `run_task_with_context` decorator usually handles this automatically,
    # but this is how it works internally.
    with restore_context(_nala_context):
        # Now, inside this block, get_tenant_id() and get_trace_id()
        # will return the values from the original API request.
        await send_email_to(user_id)

Note: The @run_task_with_context decorator automates this process for you.

Contextual Key Generation

To ensure consistency in caching, locking, and rate limiting, the ContextKeyGenerator automatically creates keys that are namespaced and aware of the current context.

A key is typically generated with the following structure: static_prefix:context_prefix:namespace:logical_key

static_prefix: A global prefix (e.g., "nala").
context_prefix: Automatically includes the tenant_id and/or user_id if enabled in the configuration.
namespace: A logical grouping for the key (e.g., "cache", "ratelimit").
logical_key: The specific identifier for the resource.

from nala.athomic.context import ContextKeyGenerator

# Assume tenant_id='tenant-123' is set in the context
key_gen = ContextKeyGenerator(namespace="cache")
cache_key = key_gen.generate("user-profile", "user-456")

# Result: "nala:tenant-123:cache:user-profile:user-456"

API Reference

`nala.athomic.context.context_vars`

`clear_all_context()`

Clears all context variables to their default values.

`clear_correlation_id()`

Clears the correlation_id context variable.

`clear_feature_flags()`

Resets the feature_flags context variable to an empty dict.

`clear_locale()`

Resets the locale context variable to its default ('pt-BR').

`clear_request_id()`

Clears the request_id context variable.

`clear_role()`

Clears the role context variable.

`clear_session_id()`

Clears the session_id context variable.

`clear_source_ip()`

Clears the source_ip context variable.

`clear_span_id()`

Clears the span_id context variable.

`clear_tenant_id()`

Clears the tenant_id context variable.

`clear_timeout_cancelled()`

Resets the timeout_cancelled flag to False.

`clear_timeout_deadline()`

Clears the timeout_deadline context variable.

`clear_trace_id()`

Clears the trace_id context variable.

`clear_user_id()`

Clears the user_id context variable.

`get_correlation_id()`

Retrieves the correlation ID for linking related events across services.

`get_current_context_dic()`

Returns a dictionary with the state of all context variables.

`get_feature_flags()`

Retrieves the dictionary of active feature flags.

`get_locale()`

Retrieves the locale/language code for the current request.

`get_request_id()`

Retrieves the unique identifier for the current request context.

`get_role()`

Retrieves the role of the authenticated user for the current context.

`get_session_id()`

Retrieves the session identifier for the current context.

`get_source_ip()`

Retrieves the source IP address of the client making the request.

`get_span_id()`

Retrieves the ID for the current span within the distributed trace.

`get_tenant_id()`

Retrieves the tenant ID for the current multi-tenancy context.

`get_timeout_deadline()`

Retrieves the absolute Unix timestamp of when the current operation must complete.

`get_trace_id()`

Retrieves the distributed trace ID for the current context.

`get_user_id()`

Retrieves the authenticated user ID for the current context.

`get_vector_clock()`

Retrieves the VectorClock for the current causality context.

`is_timeout_cancelled()`

Checks if the current operation has been marked as cancelled due to timeout.

`set_correlation_id(value)`

Sets the correlation ID for the current context, returning a token to reset the change.

`set_feature_flags(value)`

Sets the feature flags dictionary for the current context, returning a token to reset the change.

`set_locale(value)`

Sets the locale for the current context, returning a token to reset the change.

`set_request_id(value)`

Sets the request ID for the current context, returning a token to reset the change.

`set_role(value)`

Sets the user role for the current context, returning a token to reset the change.

`set_session_id(value)`

Sets the session ID for the current context, returning a token to reset the change.

`set_source_ip(value)`

Sets the source IP for the current context, returning a token to reset the change.

`set_span_id(value)`

Sets the span ID for the current context, returning a token to reset the change.

`set_tenant_id(value)`

Sets the tenant ID for the current context, returning a token to reset the change.

`set_timeout_cancelled(value)`

Sets the timeout cancelled flag for the current context, returning a token to reset the change.

`set_timeout_deadline(value)`

Sets the timeout deadline for the current context, returning a token to reset the change.

`set_trace_id(value)`

Sets the trace ID for the current context, returning a token to reset the change.

`set_user_id(value)`

Sets the user ID for the current context, returning a token to reset the change.

`set_vector_clock(value)`

Sets the VectorClock for the current context, returning a token to reset the change.

`nala.athomic.context.propagation`

`capture_context()`

Captures the current values of all context variables marked for propagation.

This function iterates through the centrally registered context variables in the context_var_manager, collects the values of those flagged with propagate=True, and returns them as a dictionary. This dictionary is suitable for serialization and passing to background tasks or events.

Returns:

Type	Description
`Dict[str, Any]`	A dictionary containing the current context values to be propagated.

`restore_context(context_dict)`

A context manager to temporarily set context variables from a dictionary.

This is used on the worker side (e.g., in a background task or event handler) to re-establish the context that existed when the job was enqueued. It safely resets all variables to their previous state upon exiting the with block.

Parameters:

Name	Type	Description	Default
`context_dict`	`Optional[Dict[str, Any]]`	The dictionary of context values captured before the task was enqueued.	required

`nala.athomic.context.generator.ContextKeyGenerator`

A centralized service for creating consistent and context-sensitive keys.

This class standardizes key generation for features like caching, locking, and rate limiting. It automatically prepends keys with a static prefix and contextual information (e.g., tenant ID) based on the application's configuration, ensuring key consistency and multi-tenancy support.

Attributes:

Name	Type	Description
`settings`	`ContextSettings`	The context configuration settings.
`static_prefix`	`str`	A global prefix for all generated keys (e.g., 'nala').
`use_tenant`	`bool`	If True, the tenant ID is included in the key.
`use_user`	`bool`	If True, the user ID is included in the key.
`separator`	`str`	The character used to separate key parts.
`namespace`	`str`	The specific namespace for this key generator instance.

`init(*, namespace=None, settings=None)`

Initializes the ContextKeyGenerator.

Parameters:

Name	Type	Description	Default
`namespace`	`Optional[str]`	A specific namespace for the keys generated by this instance. Overrides the default namespace from settings if provided.	`None`
`settings`	`Optional[ContextSettings]`	The context configuration settings. If not provided, global settings are used.	`None`

`generate(*key_parts)`

Generates a final key from multiple parts with context and namespace.

The final key structure is: static_prefix:context_prefix:namespace:part1:part2:...

Parameters:

Name	Type	Description	Default
`*key_parts`	`Any`	A variable number of logical parts to be appended to the key.	`()`

Returns:

Name	Type	Description
`str`	`str`	The fully constructed, context-aware key.

`generate_for_function(func, args, kwargs)`

Generates a unique and deterministic key for a function call.

This method is ideal for function caching decorators. It creates a stable key by combining the function's qualified name with a SHA256 hash of its serialized arguments, ensuring that the same call always produces the same key.

Parameters:

Name	Type	Description	Default
`func`	`Callable`	The function being called.	required
`args`	`Tuple`	The positional arguments passed to the function.	required
`kwargs`	`Dict`	The keyword arguments passed to the function.	required

Returns:

Name	Type	Description
`str`	`str`	A unique key representing the specific function call.

`nala.athomic.context.execution.ExecutionContext`

A data class that provides a snapshot of the current execution context.

This class consolidates all context variables (e.g., tenant_id, request_id, trace_id) from the context_vars module into a single, cohesive object. Its primary purpose is to offer a clean way to access context information and to capture the context at a specific point in time for propagation to background tasks or asynchronous events.

Attributes:

Name	Type	Description
`tenant_id`	`Optional[str]`	The identifier for the current tenant.
`request_id`	`Optional[str]`	The unique ID for the current request.
`trace_id`	`Optional[str]`	The ID for the distributed trace.
`span_id`	`Optional[str]`	The ID for the current span within a trace.
`user_id`	`Optional[str]`	The identifier for the authenticated user.
`role`	`Optional[str]`	The role of the authenticated user.
`locale`	`Optional[str]`	The locale/language for the current request.
`source_ip`	`Optional[str]`	The IP address of the original client.
`session_id`	`Optional[str]`	The session identifier.
`correlation_id`	`Optional[str]`	An ID to correlate logs and events across services.
`feature_flags`	`Optional[Dict[str, bool]]`	A dictionary of active feature flags.

`init(tenant_id=None, request_id=None, trace_id=None, span_id=None, user_id=None, role=None, locale=None, source_ip=None, session_id=None, correlation_id=None, feature_flags=None)`

Initializes the ExecutionContext.

If a value for a parameter is not explicitly provided, the constructor will fetch the current value from the global context_vars. This allows for both creating a snapshot of the current context (ExecutionContext()) and creating a custom context for specific purposes (like testing or background jobs).

Parameters:

Name	Type	Description	Default
`tenant_id`	`Optional[str]`	The unique identifier for the tenant.	`None`
`request_id`	`Optional[str]`	The unique ID for the request.	`None`
`trace_id`	`Optional[str]`	The distributed trace ID.	`None`
`span_id`	`Optional[str]`	The current span ID.	`None`
`user_id`	`Optional[str]`	The identifier for the user.	`None`
`role`	`Optional[str]`	The user's role.	`None`
`locale`	`Optional[str]`	The request's locale.	`None`
`source_ip`	`Optional[str]`	The client's IP address.	`None`
`session_id`	`Optional[str]`	The session ID.	`None`
`correlation_id`	`Optional[str]`	The correlation ID.	`None`
`feature_flags`	`Optional[Dict[str, bool]]`	A dictionary of feature flags.	`None`

`repr()`

Returns a string representation of the object.

Returns:

Name	Type	Description
`str`	`str`	A developer-friendly string representation of the context.

`to_dict()`

Serializes the context object to a dictionary.

Returns:

Type	Description
`Dict[str, Any]`	Dict[str, Any]: A dictionary representation of the execution context,
`Dict[str, Any]`	suitable for logging or propagating to other services.

Context Management

Overview

Core Use Cases

Core Components

ContextVarManager

context_vars Module

ExecutionContext

Context Propagation

Example

Contextual Key Generation

API Reference

nala.athomic.context.context_vars

clear_all_context()

clear_correlation_id()

clear_feature_flags()

clear_locale()

clear_request_id()

clear_role()

clear_session_id()

clear_source_ip()

clear_span_id()

clear_tenant_id()

clear_timeout_cancelled()

clear_timeout_deadline()

clear_trace_id()

clear_user_id()

get_correlation_id()

get_current_context_dic()

get_feature_flags()

get_locale()

get_request_id()

get_role()

get_session_id()

get_source_ip()

get_span_id()

get_tenant_id()

get_timeout_deadline()

get_trace_id()

get_user_id()

get_vector_clock()

is_timeout_cancelled()

set_correlation_id(value)

set_feature_flags(value)

set_locale(value)

set_request_id(value)

set_role(value)

set_session_id(value)

set_source_ip(value)

set_span_id(value)

set_tenant_id(value)

set_timeout_cancelled(value)

set_timeout_deadline(value)

set_trace_id(value)

set_user_id(value)

set_vector_clock(value)

nala.athomic.context.propagation

capture_context()

restore_context(context_dict)

nala.athomic.context.generator.ContextKeyGenerator

__init__(*, namespace=None, settings=None)

generate(*key_parts)

generate_for_function(func, args, kwargs)

nala.athomic.context.execution.ExecutionContext

__init__(tenant_id=None, request_id=None, trace_id=None, span_id=None, user_id=None, role=None, locale=None, source_ip=None, session_id=None, correlation_id=None, feature_flags=None)

__repr__()

to_dict()

`ContextVarManager`

`context_vars` Module

`ExecutionContext`

`nala.athomic.context.context_vars`

`clear_all_context()`

`clear_correlation_id()`

`clear_feature_flags()`

`clear_locale()`

`clear_request_id()`

`clear_role()`

`clear_session_id()`

`clear_source_ip()`

`clear_span_id()`

`clear_tenant_id()`

`clear_timeout_cancelled()`

`clear_timeout_deadline()`

`clear_trace_id()`

`clear_user_id()`

`get_correlation_id()`

`get_current_context_dic()`

`get_feature_flags()`

`get_locale()`

`get_request_id()`

`get_role()`

`get_session_id()`

`get_source_ip()`

`get_span_id()`

`get_tenant_id()`

`get_timeout_deadline()`

`get_trace_id()`

`get_user_id()`

`get_vector_clock()`

`is_timeout_cancelled()`

`set_correlation_id(value)`

`set_feature_flags(value)`

`set_locale(value)`

`set_request_id(value)`

`set_role(value)`

`set_session_id(value)`

`set_source_ip(value)`

`set_span_id(value)`

`set_tenant_id(value)`

`set_timeout_cancelled(value)`

`set_timeout_deadline(value)`

`set_trace_id(value)`

`set_user_id(value)`

`set_vector_clock(value)`

`nala.athomic.context.propagation`

`capture_context()`

`restore_context(context_dict)`

`nala.athomic.context.generator.ContextKeyGenerator`

`init(*, namespace=None, settings=None)`

`generate(*key_parts)`

`generate_for_function(func, args, kwargs)`

`nala.athomic.context.execution.ExecutionContext`

`init(tenant_id=None, request_id=None, trace_id=None, span_id=None, user_id=None, role=None, locale=None, source_ip=None, session_id=None, correlation_id=None, feature_flags=None)`

`repr()`

`to_dict()`