diff --git a/example_rules/models/base.sml b/example_rules/models/base.sml index 1b21ced..312ce12 100644 --- a/example_rules/models/base.sml +++ b/example_rules/models/base.sml @@ -12,3 +12,5 @@ EventType: Entity[str] = EntityJson( ActionName=GetActionName() +ActionId=GetActionId() + diff --git a/osprey_worker/src/osprey/engine/executor/execution_context.py b/osprey_worker/src/osprey/engine/executor/execution_context.py index 9454bf6..88c00fb 100644 --- a/osprey_worker/src/osprey/engine/executor/execution_context.py +++ b/osprey_worker/src/osprey/engine/executor/execution_context.py @@ -172,6 +172,10 @@ class ExecutionContext: """Returns the action name that the execution context is currently being invoked upon.""" return self._action.action_name + def get_action_id(self) -> int: + """Returns the action id of the event that the execution context is currently being invoked upon.""" + return self._action.action_id + def get_action_time(self) -> datetime: """Returns the time of the action that the execution context is currently being invoked upon.""" return self._action.timestamp diff --git a/osprey_worker/src/osprey/engine/stdlib/udfs/get_action_id.py b/osprey_worker/src/osprey/engine/stdlib/udfs/get_action_id.py new file mode 100644 index 0000000..aae2a3e --- /dev/null +++ b/osprey_worker/src/osprey/engine/stdlib/udfs/get_action_id.py @@ -0,0 +1,15 @@ +from ._prelude import ArgumentsBase, ExecutionContext, UDFBase +from .categories import UdfCategories + + +class Arguments(ArgumentsBase): + pass + + +class GetActionId(UDFBase[Arguments, int]): + """Returns the Action ID of the event being processed.""" + + category = UdfCategories.ENGINE + + def execute(self, execution_context: ExecutionContext, arguments: Arguments) -> int: + return execution_context.get_action_id() diff --git a/osprey_worker/src/osprey/engine/stdlib/udfs/tests/test_get_action_id.py b/osprey_worker/src/osprey/engine/stdlib/udfs/tests/test_get_action_id.py new file mode 100644 index 0000000..e060e48 --- /dev/null +++ b/osprey_worker/src/osprey/engine/stdlib/udfs/tests/test_get_action_id.py @@ -0,0 +1,16 @@ +import pytest +from osprey.engine.ast_validator.validators.validate_call_kwargs import ValidateCallKwargs +from osprey.engine.conftest import ExecuteFunction +from osprey.engine.stdlib.udfs.get_action_id import GetActionId +from osprey.engine.udf.registry import UDFRegistry + +pytestmark = [ + pytest.mark.use_validators([ValidateCallKwargs]), + pytest.mark.use_udf_registry(UDFRegistry.with_udfs(GetActionId)), +] + + +def test_execute(execute: ExecuteFunction) -> None: + data = execute('ActionId = GetActionId()', data={}, action_id=42) + + assert data == {'ActionId': 42} diff --git a/osprey_worker/src/osprey/worker/_stdlibplugin/udf_register.py b/osprey_worker/src/osprey/worker/_stdlibplugin/udf_register.py index 1565b93..b647bd5 100644 --- a/osprey_worker/src/osprey/worker/_stdlibplugin/udf_register.py +++ b/osprey_worker/src/osprey/worker/_stdlibplugin/udf_register.py @@ -12,6 +12,7 @@ from osprey.engine.stdlib.udfs.experiments import ( InExperiment, ) from osprey.engine.stdlib.udfs.extract_cookie import ExtractCookie +from osprey.engine.stdlib.udfs.get_action_id import GetActionId from osprey.engine.stdlib.udfs.get_action_name import GetActionName from osprey.engine.stdlib.udfs.import_ import Import from osprey.engine.stdlib.udfs.ip_network import IpNetwork @@ -86,6 +87,7 @@ def register_udfs() -> Sequence[Type[UDFBase[Any, Any]]]: ExperimentsBucketAssignment, InExperiment, ExtractCookie, + GetActionId, GetActionName, HasLabel, Import,