Skip to content

Commit

Permalink
Exposes curent action name through action context
Browse files Browse the repository at this point in the history
This allows you to get the current action name along with other
application-level metadata (sequence ID, partition_key, etc...)

See #501 for details
  • Loading branch information
elijahbenizzy committed Jan 27, 2025
1 parent 55ecfb8 commit 575ba94
Show file tree
Hide file tree
Showing 5 changed files with 12 additions and 4 deletions.
5 changes: 3 additions & 2 deletions burr/core/application.py
Original file line number Diff line number Diff line change
Expand Up @@ -543,8 +543,6 @@ class ApplicationContext(AbstractContextManager, ApplicationIdentifiers):
"""Application context. This is anything your node might need to know about the application.
Often used for recursive tracking.
Note this is also a context manager (allowing you to pass context to sub-applications).
To access this object in a running application, you can use the `__context` variable in the
action signature:
Expand All @@ -556,6 +554,7 @@ class ApplicationContext(AbstractContextManager, ApplicationIdentifiers):
def my_action(state: State, __context: ApplicationContext) -> State:
app_id = __context.app_id
partition_key = __context.partition_key
current_action_name = __context.action_name # Access the current action name
...
"""
Expand All @@ -564,6 +563,7 @@ def my_action(state: State, __context: ApplicationContext) -> State:
parallel_executor_factory: Callable[[], Executor]
state_initializer: Optional[BaseStateLoader]
state_persister: Optional[BaseStateSaver]
action_name: Optional[str] # Store just the action name

@staticmethod
def get() -> Optional["ApplicationContext"]:
Expand Down Expand Up @@ -865,6 +865,7 @@ def _context_factory(self, action: Action, sequence_id: int) -> ApplicationConte
parallel_executor_factory=self._parallel_executor_factory,
state_initializer=self._state_initializer,
state_persister=self._state_persister,
action_name=action.name if action else None, # Pass just the action name
)

def _step(
Expand Down
4 changes: 3 additions & 1 deletion docs/concepts/actions.rst
Original file line number Diff line number Diff line change
Expand Up @@ -110,7 +110,7 @@ Will require the inputs to be passed in at runtime. See below for how to do that
This means that the application does not *need* the inputs to be set.

Note: to access ``app_id`` and ``partition_key`` in your running application, you can have the :py:class:`ApplicationContext <burr.core.application.ApplicationContext>`
Note: to access application-level metadata such as ``app_id``, ``partition_key``, ``sequence_id``, and ``action_name`` in your running application, you can have the :py:class:`ApplicationContext <burr.core.application.ApplicationContext>`
injected into your Burr Actions. This is done by adding ``__context`` to the action signature:

.. code-block:: python
Expand All @@ -121,6 +121,8 @@ injected into your Burr Actions. This is done by adding ``__context`` to the act
def my_action(state: State, __context: ApplicationContext) -> State:
app_id = __context.app_id
partition_key = __context.partition_key
action_name = __context.action_name
sequence_id = __context.sequence_id
...
Expand Down
4 changes: 3 additions & 1 deletion docs/concepts/state-persistence.rst
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,7 @@ Note that ``partition_key`` can be `None` if this is not relevant. A UUID is alw

You set these values using the :py:meth:`with_identifiers() <burr.core.application.ApplicationBuilder.with_identifiers>` method.

Note: to access ``app_id`` and ``partition_key`` in your running application, you can have the :py:class:`ApplicationContext <burr.core.application.ApplicationContext>`
Note: to access application-level metadata such as ``app_id``, ``partition_key``, ``sequence_id``, and ``action_name`` in your running application, you can have the :py:class:`ApplicationContext <burr.core.application.ApplicationContext>`
injected into your Burr Actions. This is done by adding ``__context`` to the action signature:

.. code-block:: python
Expand All @@ -54,6 +54,8 @@ injected into your Burr Actions. This is done by adding ``__context`` to the act
def my_action(state: State, __context: ApplicationContext) -> State:
app_id = __context.app_id
partition_key = __context.partition_key
action_name = __context.action_name
sequence_id = __context.sequence_id
...
Expand Down
2 changes: 2 additions & 0 deletions tests/core/test_application.py
Original file line number Diff line number Diff line change
Expand Up @@ -1353,6 +1353,7 @@ def test_action(state: State, __context: ApplicationContext) -> State:
assert __context.sequence_id == 0
assert __context.partition_key == PARTITION_KEY
assert __context.app_id == APP_ID
assert __context.action_name == "test_action"
return state

app = (
Expand All @@ -1379,6 +1380,7 @@ def test_action(state: State, __context: ApplicationContext) -> State:
assert __context.sequence_id == 0
assert __context.partition_key == PARTITION_KEY
assert __context.app_id == APP_ID
assert __context.action_name == "test_action"
return state

app = (
Expand Down
1 change: 1 addition & 0 deletions tests/core/test_parallelism.py
Original file line number Diff line number Diff line change
Expand Up @@ -1201,6 +1201,7 @@ def reads(self) -> list[str]:
state_persister=persister,
state_initializer=persister,
parallel_executor_factory=lambda: concurrent.futures.ThreadPoolExecutor(),
action_name=action.name,
),
inputs={},
)
Expand Down

0 comments on commit 575ba94

Please sign in to comment.