diff --git a/.github/workflows/build.yaml b/.github/workflows/build.yaml index d84e957..c1fa0ae 100644 --- a/.github/workflows/build.yaml +++ b/.github/workflows/build.yaml @@ -117,13 +117,13 @@ jobs: uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7 with: name: test-results-${{ matrix.os }}-${{ matrix.python-version }} - path: '**/junit.xml' + path: junit.xml if: matrix.os == 'ubuntu-latest' && matrix.python-version == '3.11' - name: Publish Unit Test Results uses: EnricoMi/publish-unit-test-result-action@d0a4676d0e0b938bc201470d88276b7c74c712b3 # v2.24.0 with: - files: '**/junit.xml' + files: junit.xml if: matrix.os == 'ubuntu-latest' && matrix.python-version == '3.11' - name: Upload coverage diff --git a/csp_gateway/client/client.py b/csp_gateway/client/client.py index 30f2e6b..f69ebe7 100644 --- a/csp_gateway/client/client.py +++ b/csp_gateway/client/client.py @@ -18,7 +18,8 @@ from threading import Thread from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Tuple, Union, cast -from httpx import AsyncClient as httpx_AsyncClient, Response, get as GET, post as POST +import httpx +from httpx import AsyncClient as httpx_AsyncClient, Response, get as GET, patch as PATCH, post as POST, put as PUT from jsonref import replace_refs from nest_asyncio import apply as applyAsyncioNesting from packaging import version @@ -421,6 +422,7 @@ def __init__(self, config: GatewayClientConfig = None, **kwargs) -> None: "lookup": set(), "next": set(), "send": set(), + "stage": set(), "state": set(), } ) @@ -485,6 +487,7 @@ def _initialize(self) -> None: ("lookup", "/lookup/"), ("next", "/next/"), ("send", "/send/"), + ("stage", "/stage/"), ("state", "/state/"), ): if "{channel}" in path: @@ -534,6 +537,17 @@ def _buildroutews(self, route: str) -> str: auth = f"{sep}token={self.config.api_key}" return f"{host}{self.config.api_route}/{route}{auth}" + def _stage_id_param(self, staging_ids: Optional[List[str]]) -> Dict[str, Any]: + """Convert staging_ids to the query param dict for the stage REST API. + + - None -> no 'id' param (server interprets as None) + - [] -> id="" (server interprets as empty list) + - ["a", "b"] -> id="a,b" + """ + if staging_ids is None: + return {} + return {"id": ",".join(staging_ids)} + def _handle_response(self, resp: Response, route: str) -> ResponseType: try: resp_json = cast(Dict[str, Any], resp.json()) @@ -615,6 +629,81 @@ async def _postasync( await client.post(resolved_route, params={**params, **extra_params}, json=data, timeout=timeout, **self.http_args), route=route ) + def _delete( + self, + route: str, + params: Dict[str, Any] = None, + data: Dict[str, Any] = None, + timeout: float = _DEFAULT_TIMEOUT, + ) -> ResponseType: + params = params or {} + resolved_route, extra_params = self._buildroute(route) + kwargs: Dict[str, Any] = {"params": {**params, **extra_params}, "timeout": timeout, **self.http_args} + if data is not None: + kwargs["json"] = data + return self._handle_response(httpx.request("DELETE", resolved_route, **kwargs), route=route) + + async def _deleteasync( + self, + route: str, + params: Dict[str, Any] = None, + data: Dict[str, Any] = None, + timeout: float = _DEFAULT_TIMEOUT, + ) -> ResponseType: + params = params or {} + resolved_route, extra_params = self._buildroute(route) + kwargs: Dict[str, Any] = {"params": {**params, **extra_params}, "timeout": timeout, **self.http_args} + if data is not None: + kwargs["json"] = data + async with httpx_AsyncClient() as client: + return self._handle_response(await client.request("DELETE", resolved_route, **kwargs), route=route) + + def _patch( + self, + route: str, + params: Dict[str, Any] = None, + timeout: float = _DEFAULT_TIMEOUT, + ) -> ResponseType: + params = params or {} + resolved_route, extra_params = self._buildroute(route) + return self._handle_response(PATCH(resolved_route, params={**params, **extra_params}, timeout=timeout, **self.http_args), route=route) + + async def _patchasync( + self, + route: str, + params: Dict[str, Any] = None, + timeout: float = _DEFAULT_TIMEOUT, + ) -> ResponseType: + params = params or {} + resolved_route, extra_params = self._buildroute(route) + async with httpx_AsyncClient() as client: + return self._handle_response( + await client.patch(resolved_route, params={**params, **extra_params}, timeout=timeout, **self.http_args), route=route + ) + + def _put( + self, + route: str, + params: Dict[str, Any] = None, + timeout: float = _DEFAULT_TIMEOUT, + ) -> ResponseType: + params = params or {} + resolved_route, extra_params = self._buildroute(route) + return self._handle_response(PUT(resolved_route, params={**params, **extra_params}, timeout=timeout, **self.http_args), route=route) + + async def _putasync( + self, + route: str, + params: Dict[str, Any] = None, + timeout: float = _DEFAULT_TIMEOUT, + ) -> ResponseType: + params = params or {} + resolved_route, extra_params = self._buildroute(route) + async with httpx_AsyncClient() as client: + return self._handle_response( + await client.put(resolved_route, params={**params, **extra_params}, timeout=timeout, **self.http_args), route=route + ) + def _stream( self, channels: Optional[List[Union[str, Tuple[str, str]]]] = None, @@ -798,6 +887,39 @@ def send(self, field: str = "", data: Any = None, timeout: float = _DEFAULT_TIME @abstractmethod def state(self, field: str = "", timeout: float = _DEFAULT_TIMEOUT, query: Optional[Query] = None) -> ResponseType: ... + @abstractmethod + def stage( + self, + field: str, + method: str = "list", + data: Any = None, + staging_ids: Optional[List[str]] = None, + timeout: float = _DEFAULT_TIMEOUT, + ) -> ResponseType: ... + + @abstractmethod + def stage_add(self, field: str, data: Any = None, staging_ids: Optional[List[str]] = None, timeout: float = _DEFAULT_TIMEOUT) -> ResponseType: ... + + @abstractmethod + def stage_new(self, field: str, timeout: float = _DEFAULT_TIMEOUT) -> ResponseType: ... + + @abstractmethod + def stage_remove( + self, field: str, data: Any = None, staging_ids: Optional[List[str]] = None, timeout: float = _DEFAULT_TIMEOUT + ) -> ResponseType: ... + + @abstractmethod + def stage_clear(self, field: str, staging_ids: Optional[List[str]] = None, timeout: float = _DEFAULT_TIMEOUT) -> ResponseType: ... + + @abstractmethod + def stage_release(self, field: str, staging_ids: Optional[List[str]] = None, timeout: float = _DEFAULT_TIMEOUT) -> ResponseType: ... + + @abstractmethod + def stage_list(self, field: str, staging_id: Optional[str] = None, timeout: float = _DEFAULT_TIMEOUT) -> ResponseType: ... + + @abstractmethod + def stage_lookup(self, field: str, staging_id: Optional[str] = None, timeout: float = _DEFAULT_TIMEOUT) -> ResponseType: ... + @abstractmethod def stream(self, channels: Optional[List[Union[str, Tuple[str, str]]]] = None, timeout: Optional[float] = None): ... @@ -885,6 +1007,132 @@ def state( self.config.return_type = old_return_type return res + # --- Staging methods --- + + @_raiseIfNotMounted + def stage( + self, + field: str, + method: str = "list", + data: Any = None, + staging_ids: Optional[List[str]] = None, + timeout: float = _DEFAULT_TIMEOUT, + ) -> ResponseType: + """Raw staging API — dispatches to the appropriate HTTP method. + + Args: + field: Channel name with staging enabled. + method: One of "add", "remove", "release", "list", "lookup". + data: Struct data (as dict) for add/remove operations. + staging_ids: Optional list of staging IDs to target. + timeout: Request timeout. + """ + params = self._stage_id_param(staging_ids) + route = "stage/{}".format(field) + if method == "add": + return self._post(route, params=params, data=data, timeout=timeout) + elif method == "remove": + return self._delete(route, params=params, data=data, timeout=timeout) + elif method == "release": + return self._patch(route, params=params, timeout=timeout) + elif method == "list": + if staging_ids and len(staging_ids) == 1: + params = {"id": staging_ids[0]} + return self._get(route, params=params, timeout=timeout) + elif method == "lookup": + if staging_ids and len(staging_ids) == 1: + params = {"id": staging_ids[0]} + return self._put(route, params=params, timeout=timeout) + else: + raise ValueError(f"Unknown stage method: {method}") + + def stage_add( + self, + field: str, + data: Any = None, + staging_ids: Optional[List[str]] = None, + timeout: float = _DEFAULT_TIMEOUT, + ) -> ResponseType: + """Add a struct to staging area(s). + + Args: + field: Channel name. + data: Struct data (as dict). If None, creates an empty staging area. + staging_ids: Target staging IDs. None = latest/create new, [] = all. + """ + return self.stage(field, method="add", data=data, staging_ids=staging_ids, timeout=timeout) + + def stage_new(self, field: str, timeout: float = _DEFAULT_TIMEOUT) -> ResponseType: + """Create a new empty staging area. + + Convenience for ``stage_add(field, data=None, staging_ids=None)``. + """ + return self.stage(field, method="add", data=None, staging_ids=None, timeout=timeout) + + def stage_remove( + self, + field: str, + data: Any = None, + staging_ids: Optional[List[str]] = None, + timeout: float = _DEFAULT_TIMEOUT, + ) -> ResponseType: + """Remove struct(s) from staging area(s). + + Args: + field: Channel name. + data: Struct to remove. If None, clears entire staging(s). + staging_ids: Target staging IDs. None = latest, [] = all. + """ + return self.stage(field, method="remove", data=data, staging_ids=staging_ids, timeout=timeout) + + def stage_clear(self, field: str, staging_ids: Optional[List[str]] = None, timeout: float = _DEFAULT_TIMEOUT) -> ResponseType: + """Clear all structs from staging area(s) without releasing. + + Convenience for ``stage_remove(field, data=None, staging_ids=...)``. + + Args: + field: Channel name. + staging_ids: Which stagings to clear. None = latest, [] = all. + """ + return self.stage(field, method="remove", data=None, staging_ids=staging_ids, timeout=timeout) + + def stage_release(self, field: str, staging_ids: Optional[List[str]] = None, timeout: float = _DEFAULT_TIMEOUT) -> ResponseType: + """Release staged structs into the channel. + + Args: + field: Channel name. + staging_ids: Which stagings to release. None = all. + """ + return self.stage(field, method="release", staging_ids=staging_ids, timeout=timeout) + + def stage_list(self, field: str, staging_id: Optional[str] = None, timeout: float = _DEFAULT_TIMEOUT) -> ResponseType: + """List active staging areas with their contents. + + Collects the active staging IDs, then looks up each one's contents. + Returns a dict of staging_id -> list of structs (same format as stage_lookup). + + Args: + field: Channel name. + staging_id: If provided, only include this specific staging ID. + """ + ids = [staging_id] if staging_id else None + id_result = self.stage(field, method="list", staging_ids=ids, timeout=timeout) + # id_result is a list of active staging IDs (or a dict with id -> exists) + if not id_result: + return {} + # Look up contents for all active stages + return self.stage(field, method="lookup", staging_ids=None, timeout=timeout) + + def stage_lookup(self, field: str, staging_id: Optional[str] = None, timeout: float = _DEFAULT_TIMEOUT) -> ResponseType: + """Look up contents of staging area(s). + + Args: + field: Channel name. + staging_id: If provided, look up only this specific staging area. + """ + ids = [staging_id] if staging_id else None + return self.stage(field, method="lookup", staging_ids=ids, timeout=timeout) + def stream(self, channels: Optional[List[Union[str, Tuple[str, str]]]] = None, callback: Callable = None, timeout: Optional[float] = None): """Stream data from specified channels with optional key filtering for dict baskets. @@ -1032,6 +1280,7 @@ class SyncGatewayClient(SyncGatewayClientMixin, BaseGatewayClient): - lookup - next - send + - stage (stage_add, stage_new, stage_remove, stage_clear, stage_release, stage_list, stage_lookup) - state These methods will depend on the configuration of the corresponding GatewayServer. @@ -1075,6 +1324,72 @@ async def state(self, field: str = "", timeout: float = _DEFAULT_TIMEOUT, query: params = None if query is None else {"query": query.model_dump_json()} return await self._getasync("{}/{}".format("state", field), timeout=timeout, params=params) + @_raiseIfNotMounted + async def stage( + self, + field: str, + method: str = "list", + data: Any = None, + staging_ids: Optional[List[str]] = None, + timeout: float = _DEFAULT_TIMEOUT, + ) -> ResponseType: + """Raw staging API — dispatches to the appropriate HTTP method.""" + params = self._stage_id_param(staging_ids) + route = "stage/{}".format(field) + if method == "add": + return await self._postasync(route, params=params, data=data, timeout=timeout) + elif method == "remove": + return await self._deleteasync(route, params=params, data=data, timeout=timeout) + elif method == "release": + return await self._patchasync(route, params=params, timeout=timeout) + elif method == "list": + if staging_ids and len(staging_ids) == 1: + params = {"id": staging_ids[0]} + return await self._getasync(route, params=params, timeout=timeout) + elif method == "lookup": + if staging_ids and len(staging_ids) == 1: + params = {"id": staging_ids[0]} + return await self._putasync(route, params=params, timeout=timeout) + else: + raise ValueError(f"Unknown stage method: {method}") + + async def stage_add( + self, field: str, data: Any = None, staging_ids: Optional[List[str]] = None, timeout: float = _DEFAULT_TIMEOUT + ) -> ResponseType: + """Add a struct to staging area(s).""" + return await self.stage(field, method="add", data=data, staging_ids=staging_ids, timeout=timeout) + + async def stage_new(self, field: str, timeout: float = _DEFAULT_TIMEOUT) -> ResponseType: + """Create a new empty staging area.""" + return await self.stage(field, method="add", data=None, staging_ids=None, timeout=timeout) + + async def stage_remove( + self, field: str, data: Any = None, staging_ids: Optional[List[str]] = None, timeout: float = _DEFAULT_TIMEOUT + ) -> ResponseType: + """Remove struct(s) from staging area(s).""" + return await self.stage(field, method="remove", data=data, staging_ids=staging_ids, timeout=timeout) + + async def stage_clear(self, field: str, staging_ids: Optional[List[str]] = None, timeout: float = _DEFAULT_TIMEOUT) -> ResponseType: + """Clear all structs from staging area(s) without releasing.""" + return await self.stage(field, method="remove", data=None, staging_ids=staging_ids, timeout=timeout) + + async def stage_release(self, field: str, staging_ids: Optional[List[str]] = None, timeout: float = _DEFAULT_TIMEOUT) -> ResponseType: + """Release staged structs into the channel.""" + return await self.stage(field, method="release", staging_ids=staging_ids, timeout=timeout) + + async def stage_list(self, field: str, staging_id: Optional[str] = None, timeout: float = _DEFAULT_TIMEOUT) -> ResponseType: + """List active staging areas with their contents.""" + ids = [staging_id] if staging_id else None + id_result = await self.stage(field, method="list", staging_ids=ids, timeout=timeout) + if not id_result: + return {} + return await self.stage(field, method="lookup", staging_ids=None, timeout=timeout) + + async def stage_lookup(self, field: str, staging_id: Optional[str] = None, timeout: float = _DEFAULT_TIMEOUT) -> ResponseType: + """Look up contents of staging area(s).""" + ids = [staging_id] if staging_id else None + return await self.stage(field, method="lookup", staging_ids=ids, timeout=timeout) + async def stream(self, channels: List[Union[str, Tuple[str, str]]] = None, timeout: Optional[float] = None): """Stream data from specified channels with optional key filtering for dict baskets. @@ -1149,12 +1464,14 @@ class AsyncGatewayClient(AsyncGatewayClientMixin, BaseGatewayClient): An asynchronous Gateway client. This client will expose async methods: - - controls - - last - - lookup - - next - - send - - state + + - controls + - last + - lookup + - next + - send + - stage (stage_add, stage_new, stage_remove, stage_clear, stage_release, stage_list, stage_lookup) + - state These methods will depend on the configuration of the corresponding GatewayServer. @@ -1174,5 +1491,5 @@ class AsyncGatewayClient(AsyncGatewayClientMixin, BaseGatewayClient): GatewayClient = SyncGatewayClient ClientConfig = GatewayClientConfig -GatewayClientConfiguration = GatewayClientConfig ClientConfiguration = GatewayClientConfig +GatewayClientConfiguration = GatewayClientConfig diff --git a/csp_gateway/server/demo/omnibus.py b/csp_gateway/server/demo/omnibus.py index 4354dd3..4be688e 100644 --- a/csp_gateway/server/demo/omnibus.py +++ b/csp_gateway/server/demo/omnibus.py @@ -33,6 +33,7 @@ MountPerspectiveTables, MountRestRoutes, MountWebSocketRoutes, + Stage, State, ) from csp_gateway.server.config import load_gateway @@ -115,11 +116,21 @@ class ExampleGatewayChannels(GatewayChannels): example_list: ts[List[ExampleData]] = None never_ticks: ts[ExampleData] = None - s_example: ts[State[ExampleData]] = None + # State fields can be added via annotation or the `set_state` API in the module's `connect` method + example_with_state: Annotated[ts[ExampleData], State(("id", "x"))] = None + example_with_state_multiple: Annotated[ + ts[ExampleData], + State(("id", "x")), + State(("id", "y"), alias="example_with_state_alternative"), + ] = None + # example: Set in connect with a different schema, to show flexibility of state definition basket: Dict[ExampleEnum, ts[ExampleData]] = None str_basket: Dict[str, ts[ExampleData]] = None + # Staging can be added via annotation or the `set_stage` API in the module's `connect` method + example_with_stage: Annotated[ts[ExampleData], Stage()] = None + # FIXME # basket_list: Dict[ExampleEnum, ts[[ExampleData]]] = None # NOTE: this second one is not populated with data so will 404 @@ -179,6 +190,9 @@ def connect(self, channels: ExampleGatewayChannels): # Channels set via `set_channel` channels.set_channel(ExampleGatewayChannels.example, data) channels.set_channel(ExampleGatewayChannels.example_list, data_list) + channels.set_channel(ExampleGatewayChannels.example_with_state, data) + channels.set_channel(ExampleGatewayChannels.example_with_state_multiple, data) + channels.set_channel(ExampleGatewayChannels.example_with_stage, data) # Generic channel for sending data from non-csp sources channels.add_send_channel(ExampleGatewayChannels.example) @@ -187,8 +201,14 @@ def connect(self, channels: ExampleGatewayChannels): channels.add_send_channel(ExampleGatewayChannels.basket, ExampleEnum.C) channels.add_send_channel(ExampleGatewayChannels.basket) - # Rudimentary state accumulation via `set_state` - channels.set_state(ExampleGatewayChannels.example, "id") + # State accumulation via `set_state` + channels.set_state( + ExampleGatewayChannels.example, + ("id",), + ) + + # Staging via `set_stage` — enables stage_add/remove/release/list/lookup APIs + channels.set_stage(ExampleGatewayChannels.example) # Create some data streams for dict baskets data_a = self.subscribe( diff --git a/csp_gateway/server/gateway/csp/__init__.py b/csp_gateway/server/gateway/csp/__init__.py index 07204cd..a5ebc43 100644 --- a/csp_gateway/server/gateway/csp/__init__.py +++ b/csp_gateway/server/gateway/csp/__init__.py @@ -10,4 +10,5 @@ on_request_node_dict_basket, ) from .module import Module +from .stage import * from .state import * diff --git a/csp_gateway/server/gateway/csp/channels.py b/csp_gateway/server/gateway/csp/channels.py index d01d572..31cdb14 100644 --- a/csp_gateway/server/gateway/csp/channels.py +++ b/csp_gateway/server/gateway/csp/channels.py @@ -9,6 +9,7 @@ DefaultDict, Dict, List, + NamedTuple, Optional, Set, Tuple, @@ -47,7 +48,8 @@ named_wait_for_next_node, named_wait_for_next_node_dict_basket, ) -from .state import State, build_track_state_node +from .stage import Stage, _StageManager, build_staging_node +from .state import State, _StateManager, build_track_state_node if TYPE_CHECKING: from csp_gateway.utils import Query @@ -60,6 +62,24 @@ log = getLogger(__name__) +class _StateSpec(NamedTuple): + """Describes a state collection attached to a Channels instance. + + ``source_field`` is the channel name whose edge feeds the state, or None + when the state was registered via ``set_state`` with a raw edge. + """ + + source_field: Optional[str] + keyby: Union[str, Tuple[str, ...]] + indexer: Optional[Union[str, int]] = None + + +def _normalize_keyby(keyby: Union[str, Tuple[str, ...], list]) -> Tuple[str, ...]: + if isinstance(keyby, (list, tuple)): + return tuple(keyby) + return (keyby,) + + class _SnapshotModelBaseClass(BaseModel): model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid", coerce_numbers_to_str=True) @@ -166,6 +186,8 @@ def __new__(mcs: Any, name: Any, bases: Any, namespace: Any, **kwargs: Any) -> A _add_field_attributes(cls) ts_pydantic_field_types = {} + declared_states: Dict[str, _StateSpec] = {} + declared_stages: List[str] = [] for field_name, field_type in cls.model_fields.items(): # Validate that timeseries types contain structs or list of structs outer_type = field_type.annotation @@ -173,6 +195,26 @@ def __new__(mcs: Any, name: Any, bases: Any, namespace: Any, **kwargs: Any) -> A if ts_pydantic_field_type is not None: ts_pydantic_field_types[field_name] = ts_pydantic_field_type + # Collect State(...) annotation markers from Annotated metadata. + for meta in getattr(field_type, "metadata", ()): # pydantic 2 FieldInfo.metadata + if isinstance(meta, State): + alias = meta._meta_alias or field_name + if alias in declared_states: + raise ValueError( + f"Duplicate state alias '{alias}' on {name}: already declared on field '{declared_states[alias].source_field}'" + ) + declared_states[alias] = _StateSpec( + source_field=field_name, + keyby=_normalize_keyby(meta._meta_keyby), + indexer=meta._meta_indexer, + ) + elif isinstance(meta, Stage): + if field_name not in declared_stages: + declared_stages.append(field_name) + + cls._declared_states = declared_states + cls._declared_stages = declared_stages + ts_pydantic_field_types[_CSP_ENGINE_CYCLE_TIMESTAMP_FIELD] = (Optional[datetime], None) dynamic_pydantic_model = create_model("_snapshot_model", __base__=_SnapshotModelBaseClass, **ts_pydantic_field_types) cls._snapshot_model = dynamic_pydantic_model @@ -188,11 +230,20 @@ class Channels(BaseModel, metaclass=ChannelsMetaclass): The names of the channels match the names that are used via the different APIs, i.e. REST, WebSockets, Perspective, etc. It is expected that developers interact with channels through these APIs (or via get_channel/set_channel in csp). - Channels that begin with ``s_`` are "state" channels, meaning that they represent a collection of messages, typically - the last message grouped by some key (i.e. security id). These are not meant to be interacted with directly, but rather - through the "state" part of the REST API. + State collections are declared via ``Annotated[ts[X], State(keyby=..., indexer=..., alias=...)]`` + on a channel field (auto-wired from the channel's edge), or registered at module + connect time via ``set_state(field_or_edge, keyby, indexer=None)``. State + collections are exposed through the state part of the REST API and + ``state``/``query`` helpers. """ + # Populated by ChannelsMetaclass from Annotated[ts[X], State(...)] markers. + # alias -> _StateSpec(source_field, keyby, indexer) + _declared_states: Dict[str, _StateSpec] = {} + + # Populated by ChannelsMetaclass from Annotated[ts[X], Stage()] markers. + _declared_stages: List[str] = [] + model_config = dict(arbitrary_types_allowed=True) # (for FeedbackOutputDef) _finalized: bool = PrivateAttr(default=False) @@ -206,12 +257,28 @@ class Channels(BaseModel, metaclass=ChannelsMetaclass): _override_blocks: Dict[Any, Optional[datetime]] = PrivateAttr(default_factory=dict) _feedbacks: Dict[int, FeedbackOutputDef] = PrivateAttr(default_factory=dict) + # alias -> _StateSpec (annotation-declared and set_state-registered combined) + _states: Dict[str, _StateSpec] = PrivateAttr(default_factory=dict) + # (alias, indexer) -> ConcurrentFutureAdapter trigger _state_requests: Dict[Tuple[str, Optional[Union[str, int]]], Any] = PrivateAttr(default_factory=dict) - _state_keybys: Dict[Tuple[str, Optional[Union[str, int]]], Tuple[str, ...]] = PrivateAttr(default_factory=dict) + # (alias, indexer) -> bound state Edge + _state_edges: Dict[Tuple[str, Optional[Union[str, int]]], Any] = PrivateAttr(default_factory=dict) + # (alias, indexer) -> DelayedEdge handed out by get_state for a state that has + # been declared (via Module.dynamic_state_channels) but not yet wired by set_state. + # The DelayedEdge is bound to the real state Edge once set_state runs, so modules + # may call get_state before the owning module's connect() runs. + _delayed_state_edges: Dict[Tuple[str, Optional[Union[str, int]]], DelayedEdge] = PrivateAttr(default_factory=dict) + # alias -> element type for pre-declared but not-yet-wired dynamic states. + _pending_state_element_types: Dict[str, type] = PrivateAttr(default_factory=dict) _last_requests: Dict[Tuple[str, Optional[Union[str, int]]], Any] = PrivateAttr(default_factory=dict) _next_requests: Dict[Tuple[str, Optional[Union[str, int]]], Any] = PrivateAttr(default_factory=dict) _send_channels: Dict[Tuple[str, Optional[Union[str, int]]], Any] = PrivateAttr(default_factory=dict) + # Staging: channel_name -> _StageManager instance + _stages: Dict[str, _StageManager] = PrivateAttr(default_factory=dict) + # Staging: channel_name -> GenericPushAdapter (release trigger) + _stage_triggers: Dict[str, Any] = PrivateAttr(default_factory=dict) + # inside context, the module being attached _module_being_attached: Any = PrivateAttr(None) # inside context, the requirements of the module being attached @@ -234,6 +301,20 @@ class Channels(BaseModel, metaclass=ChannelsMetaclass): # Might be public in the future. _null_ts: List[Tuple[str, Optional[str]]] = PrivateAttr(default_factory=list) + def model_post_init(self, __context: Any) -> None: + # Seed instance state registry with class-level declarations from annotations. + for alias, spec in self.__class__._declared_states.items(): + self._states[alias] = spec + + @classmethod + def state_aliases(cls) -> List[str]: + """Return the list of state aliases declared on the class via annotations.""" + return list(cls._declared_states.keys()) + + def all_state_aliases(self) -> List[str]: + """Return all known state aliases (declared + dynamically registered).""" + return list(self._states.keys()) + def dynamic_keys(self) -> Optional[Dict[str, List[Any]]]: """Define dynamic dictionary keys by field, driven by data from the channels.""" ... @@ -341,6 +422,11 @@ def _finalize(self) -> None: who_requires_id[requires].add(id(module)) if not self._finalized: + # Auto-wire state collections declared via annotations. + self._wire_declared_states() + # Auto-wire staging declared via annotations. + self._wire_declared_stages() + # first ensure everything is provided for ( field, @@ -436,6 +522,60 @@ def _finalize(self) -> None: self._finalized = True log.debug(f"Feedback count: {self._feedback_count}") + def _wire_declared_states(self) -> None: + for alias, spec in list(self._states.items()): + if spec.source_field is None: + continue # registered via set_state, already wired + if (alias, spec.indexer) in self._state_edges: + continue + + # Skip if no module provides this channel (e.g. all setters disabled) + if not self._delayed_edge_providers.get(spec.source_field): + continue + + tstype = self.get_outer_type(spec.source_field) + if is_dict_basket(tstype): + if spec.indexer is None: + raise NotImplementedError( + f"Annotation-declared state '{alias}' on dict basket '{spec.source_field}' " + f"requires an indexer (set indexer=... in State(...))" + ) + edge = self.get_channel(spec.source_field, indexer=spec.indexer) + else: + edge = self.get_channel(spec.source_field) + self._wire_state_edge(alias, edge, spec.keyby, spec.indexer) + + def _wire_declared_stages(self) -> None: + """Auto-wire staging for channels declared via Annotated[ts[X], Stage()].""" + for field_name in self.__class__._declared_stages: + if field_name in self._stages: + continue # already enabled via set_stage in connect + + if not self._delayed_edge_providers.get(field_name): + continue # no provider for this channel + + # Determine element type from the channel's ts type + tstype = self.get_outer_type(field_name) + if is_dict_basket(tstype): + element_type = get_dict_basket_value_type(tstype) + elif isTsType(tstype): + element_type = tstype.typ + else: + continue + + # Unwrap List[T] -> T + from typing import get_args as _get_args, get_origin as _get_origin + + if _get_origin(element_type) is list: + element_type = _get_args(element_type)[0] + + stage, push_adapter = build_staging_node(element_type) + self._stages[field_name] = stage + self._stage_triggers[field_name] = push_adapter + + # Wire the push adapter's output as an additional provider for this channel + self._delayed_edge_providers[field_name].append((None, push_adapter.out())) + def _bind_delayed_channel(self, field, list_of_edges_and_modules, indexer=None): tstype = self.get_outer_type(field) # Make sure a getter node exists first @@ -603,10 +743,6 @@ def get_channel( return getattr(self, field) - @classmethod - def is_state_field(cls, field): - return field.startswith("s_") - @classmethod def get_outer_type(cls, field): return cls.model_fields[field].annotation @@ -620,9 +756,6 @@ def set_channel( # add to graph self._add_field_to_graph(field, self._module_being_attached, True, indexer) - # TODO fix ugly state field stuff - is_state_field = self.is_state_field(field) - tstype = self.get_outer_type(field) if is_dict_basket(tstype): _is_dict_basket = True @@ -669,7 +802,7 @@ def set_channel( else: edge_tstypes = [edge.tstype] # type: ignore[union-attr] - if not all(edge_tstype == gateway_tstype for edge_tstype in edge_tstypes) and not is_state_field: + if not all(edge_tstype == gateway_tstype for edge_tstype in edge_tstypes): raise TypeError("Edge type incorrect for {}: should be {}, found {}".format(field, gateway_tstype, edge_tstypes[0])) module = self._module_being_attached @@ -688,61 +821,146 @@ def set_channel( self._set_last(field) self._set_next(field) - def _ensure_state_field(self, field: str) -> str: - if not field.startswith("s_"): - return "s_{}".format(field) - return field - def set_state( self, - field: str, + field_or_edge: Union[Edge, str], keyby: Union[str, Tuple[str, ...]], - indexer: Union[str, int] = None, + indexer: Optional[Union[str, int]] = None, ) -> None: - # grab state version of field - state_field = self._ensure_state_field(field) - - # Bail if already setup - if (state_field, indexer) in self._state_requests: - return - - # First ensure edge is constructed - edge = self.get_channel(field, indexer=indexer) - - # And ensure the state edge is constructed - self.get_state(state_field, indexer=indexer) + """Register a state collection on a channel. - if isinstance(edge, Edge): - # instantiate state node - if get_origin(edge.tstype.typ) is list: - edge_type_name = get_args(edge.tstype.typ)[0].__name__ - state_edge = build_track_state_node(csp.unroll(edge), keyby) - else: - edge_type_name = edge.tstype.typ.__name__ - state_edge = build_track_state_node(edge, keyby) + The first argument may be either: - state_edge.nodedef.__name__ = "State[{}]".format(edge_type_name) + - a channel field name (``str``) — the edge is resolved via + :meth:`get_channel`, and the state is exposed under that field name; + - a csp ``Edge`` previously registered via :meth:`set_channel` — the + field name is recovered by reverse lookup of the registered edge. - # register for use inside other csp nodes - self.set_channel(state_field, state_edge, indexer=indexer) + ``keyby`` and ``indexer`` describe how ticks are accumulated. If the + state is already registered, this is a no-op when ``keyby``/``indexer`` + match; otherwise a ``ValueError`` is raised. + """ + if isinstance(field_or_edge, str): + field = field_or_edge + edge = self.get_channel(field, indexer=indexer) if indexer is not None else self.get_channel(field) + elif isinstance(field_or_edge, Edge): + edge = field_or_edge + field = self._find_field_for_edge(edge) + if field is None: + raise ValueError( + "set_state could not resolve a channel name from the given edge; " + "register it via set_channel first, or pass the channel field name (str)." + ) + else: + raise TypeError("set_state expects a channel field name (str) or a csp Edge as the first argument; got {}".format(type(field_or_edge))) + + keyby = _normalize_keyby(keyby) + existing = self._states.get(field) + if existing is not None: + if existing.keyby != keyby or existing.indexer != indexer: + raise ValueError( + f"State '{field}' already registered with " + f"keyby={existing.keyby!r}, indexer={existing.indexer!r}; " + f"cannot redefine with keyby={keyby!r}, indexer={indexer!r}" + ) + if (field, indexer) in self._state_edges: + return # already wired + + # If this was pre-declared via dynamic_state_channels, it is no longer pending + # once set_state is called for it. + self._pending_state_element_types.pop(field, None) + + self._states[field] = _StateSpec(source_field=field, keyby=keyby, indexer=indexer) + self._wire_state_edge(field, edge, keyby, indexer) + + def _find_field_for_edge(self, edge: Edge) -> Optional[str]: + """Reverse-lookup a channel field name for a previously-set edge.""" + for field, providers in self._delayed_edge_providers.items(): + for _module, provided in providers: + if provided is edge: + return field + if isinstance(provided, dict): + for v in provided.values(): + if v is edge: + return field + return None - # setup ad-hoc querying - trigger = ConcurrentFutureAdapter(name="RequestState<{}>".format(edge_type_name)) + def _wire_state_edge( + self, + field: str, + edge: Edge, + keyby: Union[str, Tuple[str, ...]], + indexer: Optional[Union[str, int]], + ) -> None: + if get_origin(edge.tstype.typ) is list: + edge_type_name = get_args(edge.tstype.typ)[0].__name__ + state_edge = build_track_state_node(csp.unroll(edge), keyby) + else: + edge_type_name = edge.tstype.typ.__name__ + state_edge = build_track_state_node(edge, keyby) + + state_edge.nodedef.__name__ = "State[{}]".format(edge_type_name) + + trigger = ConcurrentFutureAdapter(name="RequestState<{}>".format(edge_type_name)) + named_on_request_node("QueryState<{}>".format(edge_type_name))(state_edge, trigger.out()) + + self._state_requests[field, indexer] = trigger + # If a DelayedEdge was previously handed out by get_state (because another module + # called get_state before this module's set_state), bind it now so the consumer's + # edge resolves to the real state node. + delayed = self._delayed_state_edges.get((field, indexer)) + if delayed is not None: + delayed.bind(state_edge) + self._state_edges[field, indexer] = delayed + else: + self._state_edges[field, indexer] = state_edge - named_on_request_node("QueryState<{}>".format(edge_type_name))(state_edge, trigger.out()) + def _declare_dynamic_state(self, field: str, element_type: type) -> None: + """Pre-register a dynamically-created state channel so that :meth:`get_state` + for ``field`` returns a :class:`DelayedEdge` before the owning module's + ``connect`` calls :meth:`set_state`. - # register the trigger - self._state_requests[state_field, indexer] = trigger - self._state_keybys[state_field, indexer] = (keyby,) if isinstance(keyby, str) else tuple(keyby) - else: - # TODO - raise NotImplementedError() + ``element_type`` is the unwrapped element type of the state (i.e. ``T`` for a + channel typed as either ``ts[T]`` or ``ts[List[T]]``). The eventual + :meth:`set_state` call provides ``keyby``/``indexer`` and binds the real state + edge into the previously-handed-out :class:`DelayedEdge`. + """ + if field in self._states or field in self._pending_state_element_types: + return + self._pending_state_element_types[field] = element_type def get_state(self, field: str, indexer: Union[str, int] = None) -> Any: - # grab state version of field - state_field = self._ensure_state_field(field) - - return self.get_channel(state_field, indexer=indexer) + """Return the underlying state Edge for ``field`` (csp-graph use).""" + if (field, indexer) in self._state_edges: + return self._state_edges[field, indexer] + if field in self._pending_state_element_types: + # Pre-declared by Module.dynamic_state_channels but the owning module has + # not yet called set_state. Hand out a DelayedEdge that will be bound once + # set_state runs (order-independent across modules). + delayed = self._delayed_state_edges.get((field, indexer)) + if delayed is None: + element_type = self._pending_state_element_types[field] + delayed = DelayedEdge(ts[_StateManager[element_type]]) + delayed.__name__ = "s_{}".format(field) + self._delayed_state_edges[field, indexer] = delayed + return delayed + if field not in self._states: + raise NoProviderException("Unknown state: {}".format(field)) + spec = self._states[field] + if spec.source_field is None: + raise NoProviderException("State '{}' (indexer={}) has not been wired yet".format(field, indexer)) + # Lazily wire annotation-declared state on first access + tstype = self.get_outer_type(spec.source_field) + if is_dict_basket(tstype): + if spec.indexer is None: + raise NotImplementedError( + f"Annotation-declared state '{field}' on dict basket '{spec.source_field}' requires an indexer (set indexer=... in State(...))" + ) + edge = self.get_channel(spec.source_field, indexer=spec.indexer) + else: + edge = self.get_channel(spec.source_field) + self._wire_state_edge(field, edge, spec.keyby, spec.indexer) + return self._state_edges[field, indexer] def _set_last(self, field: str, indexer: Union[str, int] = None) -> None: # Bail if already setup @@ -925,17 +1143,13 @@ def next(self, field: str, indexer: Union[str, int] = None, *, timeout=None) -> return result def state(self, field: str, indexer: Union[str, int] = None, *, timeout=None) -> Any: - # grab state version of field - state_field = self._ensure_state_field(field) + self._check(field, self._state_requests, "state", indexer=indexer) - self._check(state_field, self._state_requests, "state", indexer=indexer) - - # TODO checks for state tracking # TODO make sure not called from inside graph context - state_edge = self._state_requests[state_field, indexer] + trigger = self._state_requests[field, indexer] # trigger request for state into graph - future = state_edge.push_tick() + future = trigger.push_tick() # wait for result return future.result(timeout=timeout) @@ -961,6 +1175,130 @@ def send(self, field: str, value: Any, indexer: Union[str, int] = None) -> None: # basket, push into first item of tuple send_channel[0].push_tick(value) + # ------------------------------------------------------------------ + # Staging API + # ------------------------------------------------------------------ + + def set_stage(self, field: str) -> None: + """Enable staging mode for a channel. + + Once staging is enabled, the channel gains stage_add/stage_remove/ + stage_release/stage_list/stage_lookup methods. Released items are + injected into the channel as a single tick via a feedback edge. + + Must be called during module connect (before finalization). + """ + if field in self._stages: + return # already enabled + + if not self._validate_field_name(field): + raise AttributeError("{} has no channel: {}".format(self.__class__, field)) + + # Determine element type from the channel's ts type + tstype = self.get_outer_type(field) + if is_dict_basket(tstype): + element_type = get_dict_basket_value_type(tstype) + elif isTsType(tstype): + element_type = tstype.typ + else: + raise TypeError(f"Cannot enable staging on non-timeseries field: {field}") + + # Unwrap List[T] -> T + from typing import get_args as _get_args, get_origin as _get_origin + + if _get_origin(element_type) is list: + element_type = _get_args(element_type)[0] + + stage, push_adapter = build_staging_node(element_type) + self._stages[field] = stage + self._stage_triggers[field] = push_adapter + + # Wire the push adapter's output as an additional provider for this channel + module = self._module_being_attached + self._delayed_edge_providers[field].append((module, push_adapter.out())) + + def staged_channels(self) -> List[str]: + """Return list of channel names that have staging enabled.""" + return list(self._stages.keys()) + + def stage_add( + self, + field: str, + struct: Any = None, + staging_ids: Optional[List[str]] = None, + ) -> List[str]: + """Add a struct to staging area(s) for a channel. + + See STAGE.md for full semantics. + """ + if field not in self._stages: + raise NoProviderException(f"No staging enabled for channel: {field}") + return self._stages[field].stage_add(struct, staging_ids) + + def stage_remove( + self, + field: str, + struct: Any = None, + staging_ids: Optional[List[str]] = None, + ) -> List[str]: + """Remove struct(s) from staging area(s). + + See STAGE.md for full semantics. + """ + if field not in self._stages: + raise NoProviderException(f"No staging enabled for channel: {field}") + return self._stages[field].stage_remove(struct, staging_ids) + + def stage_release( + self, + field: str, + staging_ids: Optional[List[str]] = None, + ) -> Dict[str, List[Any]]: + """Release staged structs into the channel. + + Released items are pushed into the csp graph individually. + Returns dict mapping staging_id -> list of released structs. + """ + if field not in self._stages: + raise NoProviderException(f"No staging enabled for channel: {field}") + + stage = self._stages[field] + released = stage.stage_release(staging_ids) + + # Push each released item into the graph via the push adapter + push_adapter = self._stage_triggers[field] + for items in released.values(): + for item in items: + push_adapter.push_tick(item) + + return released + + def stage_list( + self, + field: str, + staging_id: Optional[str] = None, + ) -> List[str]: + """List staging IDs for a channel. + + See STAGE.md for full semantics. + """ + if field not in self._stages: + raise NoProviderException(f"No staging enabled for channel: {field}") + return self._stages[field].stage_list(staging_id) + + def stage_lookup( + self, + field: str, + staging_id: Optional[str] = None, + ) -> Dict[str, List[Any]]: + """Look up contents of staging area(s). + + Returns dict mapping staging_id -> list of structs. + """ + if field not in self._stages: + raise NoProviderException(f"No staging enabled for channel: {field}") + return self._stages[field].stage_lookup(staging_id) + def _check(self, field: str, where: Dict, kind: str, indexer: Union[str, int] = None) -> None: if (field, indexer) not in where: # TODO should only be called once the graph is started diff --git a/csp_gateway/server/gateway/csp/factory.py b/csp_gateway/server/gateway/csp/factory.py index 5ba3fa7..c5017df 100644 --- a/csp_gateway/server/gateway/csp/factory.py +++ b/csp_gateway/server/gateway/csp/factory.py @@ -1,5 +1,5 @@ from datetime import datetime -from typing import Generic, List, Optional, Type +from typing import Generic, List, Optional, Type, get_args, get_origin from ccflow import BaseModel from csp.impl.enum import Enum @@ -57,6 +57,24 @@ def build(self, channels: ChannelsType) -> ChannelsType: if self.block_set_channels_until is not None: channels._block_set_channels_until = self.block_set_channels_until + # Pre-declare dynamically-created state channels so that modules calling get_state + # during connect() are independent of the order in which set_state is called by the + # owning module. The owning module includes the name in both dynamic_channels() + # (which provides the element type) and dynamic_state_channels(). + for node in enabled_modules: + declared = node.dynamic_state_channels() if hasattr(node, "dynamic_state_channels") else None + if not declared: + continue + owned_channels = node.dynamic_channels() or {} + for name in declared: + if name not in owned_channels: + raise ValueError( + f"Module {type(node).__name__} declared '{name}' in dynamic_state_channels() but did not include it in dynamic_channels()" + ) + t = owned_channels[name] + element_type = get_args(t)[0] if get_origin(t) is list else t + channels._declare_dynamic_state(name, element_type) + # Wire in each edge Provider. # The implementation of set_channel will handle multiplexing streams for node in enabled_modules: diff --git a/csp_gateway/server/gateway/csp/module.py b/csp_gateway/server/gateway/csp/module.py index e434573..e9746aa 100644 --- a/csp_gateway/server/gateway/csp/module.py +++ b/csp_gateway/server/gateway/csp/module.py @@ -49,8 +49,13 @@ def dynamic_channels(self) -> Optional[Dict[str, Union[Type[GatewayStruct], Type ... def dynamic_state_channels(self) -> Optional[Set[str]]: - """ - The set of dynamic channels that have state. + """The subset of :meth:`dynamic_channels` for which this module will also call + :meth:`GatewayChannels.set_state` from within :meth:`connect`. + + Declaring a name here lets *other* modules call :meth:`GatewayChannels.get_state` + on the channel from their own ``connect`` regardless of the order in which modules + are connected; the returned state edge is bound to the real state node once the + owning module's ``connect`` calls ``set_state``. """ ... diff --git a/csp_gateway/server/gateway/csp/stage.py b/csp_gateway/server/gateway/csp/stage.py new file mode 100644 index 0000000..84d77fa --- /dev/null +++ b/csp_gateway/server/gateway/csp/stage.py @@ -0,0 +1,281 @@ +"""Core staging support for csp-gateway channels. + +A channel with staging enabled accumulates structs into named "staging areas" +before they are released into the main channel. This allows batch preparation +and atomic release of groups of structs. + +See STAGE.md for the full API specification. +""" + +import threading +from typing import Dict, List, Optional + +from csp.impl.genericpushadapter import GenericPushAdapter + +from csp_gateway.utils import GatewayStruct + +__all__ = ( + "Stage", + "Staging", + "build_staging_node", +) + + +class Stage: + """Annotation marker for declaring staging on a channel. + + Usage:: + + class MyChannels(GatewayChannels): + orders: Annotated[ts[OrderStruct], Stage()] = None + + This is equivalent to calling ``channels.set_stage("orders")`` in the + module's ``connect`` method. + """ + + def __init__(self) -> None: + pass + + +class Staging(GatewayStruct): + """A staging group represented as a GatewayStruct. + + Each StagingArea gets a unique id and timestamp automatically from + GatewayStruct. It holds a list of struct instances that have been staged + but not yet released. The staging_id is simply ``self.id``. + + Note: ``items`` uses ``list`` (untyped) rather than ``List[GatewayStruct]`` + because CSP's C++ struct layer enforces strict type matching on typed lists, + and not all channel structs inherit from GatewayStruct in the struct hierarchy. + """ + + items: list = [] + + def add(self, struct) -> None: + self.items.append(struct) + + def remove(self, struct) -> bool: + """Remove a struct by id. Returns True if found and removed.""" + for i, item in enumerate(self.items): + if item.id == struct.id: + self.items.pop(i) + return True + return False + + def clear(self) -> list: + """Remove all items and return them.""" + items = self.items[:] + self.items = [] + return items + + def lookup(self) -> list: + """Return a copy of the items list.""" + return self.items[:] + + +class _StageManager: + """Manages multiple staging areas for a single channel. + + Thread-safe: all mutations are guarded by a lock. + """ + + def __init__(self): + self._lock = threading.Lock() + self._areas: Dict[str, Staging] = {} + self._released: Dict[str, Staging] = {} + + @property + def staging_ids(self) -> List[str]: + with self._lock: + return list(self._areas.keys()) + + def stage_add( + self, + struct: Optional[GatewayStruct] = None, + staging_ids: Optional[List[str]] = None, + ) -> List[str]: + """Add a struct to staging area(s). + + Returns the list of staging IDs affected. + See STAGE.md for the full semantics. + """ + with self._lock: + if struct is None and staging_ids is not None and len(staging_ids) > 0: + # None, [staging_id]: error + raise ValueError("Cannot specify staging_ids without a struct to add") + + if struct is None: + # None, None or None, []: create a new empty staging + area = Staging() + self._areas[area.id] = area + return [area.id] + + if staging_ids is None: + # struct, None: if staging exists, add to latest; else create new + if self._areas: + latest_id = list(self._areas.keys())[-1] + self._areas[latest_id].add(struct) + return [latest_id] + else: + area = Staging() + area.add(struct) + self._areas[area.id] = area + return [area.id] + + if len(staging_ids) == 0: + # struct, []: add to all existing, or create new if none + if self._areas: + affected = [] + for sid, area in self._areas.items(): + # Add only if not already present + if not any(item.id == struct.id for item in area.items): + area.add(struct) + affected.append(sid) + if not affected: + # Already in all, create a new one + area = Staging() + area.add(struct) + self._areas[area.id] = area + return [area.id] + return affected + else: + area = Staging() + area.add(struct) + self._areas[area.id] = area + return [area.id] + + # struct, [staging_id, ...]: add to specified stagings + affected = [] + for sid in staging_ids: + if sid not in self._areas: + raise KeyError(f"Staging ID not found: {sid}") + self._areas[sid].add(struct) + affected.append(sid) + return affected + + def stage_remove( + self, + struct: Optional[GatewayStruct] = None, + staging_ids: Optional[List[str]] = None, + ) -> List[str]: + """Remove struct(s) from staging area(s). + + Returns the list of staging IDs affected. + See STAGE.md for the full semantics. + """ + with self._lock: + if struct is None and staging_ids is not None and len(staging_ids) == 0: + # None, []: clear all stagings + affected = list(self._areas.keys()) + self._areas.clear() + return affected + + if struct is None and staging_ids is None: + # None, None: clear latest staging + if not self._areas: + return [] + latest_id = list(self._areas.keys())[-1] + del self._areas[latest_id] + return [latest_id] + + if struct is None and staging_ids is not None and len(staging_ids) > 0: + # None, [staging_id]: clear all structs from given staging + affected = [] + for sid in staging_ids: + if sid in self._areas: + self._areas[sid].clear() + affected.append(sid) + return affected + + if struct is not None and staging_ids is None: + # struct, None: remove from latest staging containing it + for sid in reversed(list(self._areas.keys())): + if self._areas[sid].remove(struct): + return [sid] + return [] + + if struct is not None and staging_ids is not None and len(staging_ids) == 0: + # struct, []: remove from all stagings + affected = [] + for sid, area in self._areas.items(): + if area.remove(struct): + affected.append(sid) + return affected + + # struct, [staging_id]: remove from specific staging + affected = [] + for sid in staging_ids: + if sid in self._areas: + if self._areas[sid].remove(struct): + affected.append(sid) + return affected + + def stage_release( + self, + staging_ids: Optional[List[str]] = None, + ) -> Dict[str, List[GatewayStruct]]: + """Release staged structs. + + Returns a dict mapping staging_id -> list of released structs. + Released stagings are moved to the released archive. + """ + with self._lock: + if staging_ids is None: + # Release all + released = {} + for sid, area in list(self._areas.items()): + released[sid] = area.items[:] + self._released[sid] = area + self._areas.clear() + return released + + released = {} + for sid in staging_ids: + if sid in self._areas: + released[sid] = self._areas[sid].items[:] + self._released[sid] = self._areas[sid] + del self._areas[sid] + return released + + def stage_list( + self, + staging_id: Optional[str] = None, + ) -> List[str]: + """List staging IDs, or verify a specific one exists.""" + with self._lock: + if staging_id is None: + return list(self._areas.keys()) + if staging_id in self._areas: + return [staging_id] + return [] + + def stage_lookup( + self, + staging_id: Optional[str] = None, + ) -> Dict[str, List[GatewayStruct]]: + """Look up contents of staging area(s), including released stages. + + Returns dict mapping staging_id -> list of structs. + """ + with self._lock: + if staging_id is None: + result = {sid: area.lookup() for sid, area in self._areas.items()} + result.update({sid: area.lookup() for sid, area in self._released.items()}) + return result + if staging_id in self._areas: + return {staging_id: self._areas[staging_id].lookup()} + if staging_id in self._released: + return {staging_id: self._released[staging_id].lookup()} + return {} + + +def build_staging_node(element_type: type) -> tuple: + """Build a _StageManager instance and its associated push adapter for releases. + + Returns (stage, push_adapter) where: + - stage: the _StageManager instance for managing staging areas + - push_adapter: GenericPushAdapter[element_type] to push released items into the graph + """ + stage = _StageManager() + push_adapter = GenericPushAdapter(element_type, name=f"StagingRelease<{element_type.__name__}>") + return stage, push_adapter diff --git a/csp_gateway/server/gateway/csp/state.py b/csp_gateway/server/gateway/csp/state.py index 4d906fe..242aaeb 100644 --- a/csp_gateway/server/gateway/csp/state.py +++ b/csp_gateway/server/gateway/csp/state.py @@ -8,7 +8,7 @@ from enum import Enum as PyEnum from functools import lru_cache from pprint import pformat -from typing import Any, Deque, Dict, List, Tuple, Union +from typing import Any, Deque, Dict, List, Optional, Tuple, Union import csp import duckdb @@ -135,6 +135,16 @@ def disable_duckdb_state() -> None: _USE_DUCKDB_STATE = False +def _resolve_keyby_attr(record: Any, path: str) -> Any: + """Resolve a (possibly dotted) attribute path on ``record``, returning ``None`` if any segment is missing.""" + obj = record + for segment in path.split("."): + if obj is None: + return None + obj = getattr(obj, segment, None) + return obj + + class StateType(CoreEnum): UNKNOWN = 0 DEFAULT = 1 @@ -209,7 +219,7 @@ def insert(self, record: Any) -> None: for subkey in self._keyby: # extract the key from the record - subkey_to_use = getattr(record, subkey, None) + subkey_to_use = _resolve_keyby_attr(record, subkey) if subkey == self._keyby[-1]: # Put the element there if last @@ -468,7 +478,7 @@ def insert(self, record: Struct) -> None: obj_id = None for subkey in self._keyby: # extract the key from the record - subkey_to_use = getattr(record, subkey, None) + subkey_to_use = _resolve_keyby_attr(record, subkey) if subkey == self._keyby[-1]: if subkey_to_use not in place.keys(): @@ -587,23 +597,64 @@ def get_duckdb_schema_struct(cls: Struct) -> Tuple[Dict, bool]: return (new_type_info, use_duckdb) -# NOTE: NEVER access State object directly, always access through the __class_getitem__ API -class State(BaseState): - def __init__(self, keyby: Union[Tuple[str, ...], str] = ("id",)) -> None: - """Switch case between different state specializations based on the type of the records""" +# NOTE: For runtime state instances, always use the _StateManager[] API. +# State() called directly is the annotation marker form used in +# `Annotated[ts[X], State(keyby=..., indexer=..., alias=...)]`. +# State[T] is sugar that delegates to _StateManager[T] for backward compatibility. + + +class State: + """Annotation marker for declaring state on a channel. + + Usage:: + + class MyChannels(GatewayChannels): + orders: Annotated[ts[OrderStruct], State(keyby=("id", "x"))] = None + + This is equivalent to calling ``channels.set_state("orders", ("id", "x"))`` + in the module's ``connect`` method. + + For backward compatibility, ``State[T]`` delegates to ``_StateManager[T]`` + to create runtime state instances. + """ + + _meta_keyby: Union[Tuple[str, ...], str] = ("id",) + _meta_indexer: Optional[Union[str, int]] = None + _meta_alias: Optional[str] = None + + def __init__( + self, + keyby: Union[Tuple[str, ...], str] = ("id",), + indexer: Optional[Union[str, int]] = None, + alias: Optional[str] = None, + ) -> None: + self._meta_keyby = keyby + self._meta_indexer = indexer + self._meta_alias = alias + + +class _StateManager(BaseState, State): + """Runtime state container that dispatches to DefaultState or DuckDBState. + + Use via ``_StateManager[T](keyby=...)`` to create a parameterized instance. + Inherits from State so that ``isinstance(mgr, State)`` is True. + """ + + def __init__( + self, + keyby: Union[Tuple[str, ...], str] = ("id",), + ) -> None: + typ = getattr(self, "_typ", None) + if typ is None: + raise TypeError("_StateManager must be parameterized with a type: _StateManager[MyStruct](keyby=...)") global _USE_DUCKDB_STATE - try: - typ = self._typ - if _USE_DUCKDB_STATE and isinstance(typ, type) and issubclass(typ, Struct): - schema, use_duckdb = get_duckdb_schema_struct(typ) - if use_duckdb: - self._state_impl = DuckDBState(typ, keyby, schema) - self._state_type = StateType.DUCKDB - return - except AttributeError: - log.warning("Do not create object directly from State use the State[] API instead for performance reasons") - log.warning("Using DefaultStateClass") + if _USE_DUCKDB_STATE and isinstance(typ, type) and issubclass(typ, Struct): + schema, use_duckdb = get_duckdb_schema_struct(typ) + if use_duckdb: + self._state_impl = DuckDBState(typ, keyby, schema) + self._state_type = StateType.DUCKDB + return self._state_impl = DefaultState(keyby) self._state_type = StateType.DEFAULT @@ -641,18 +692,22 @@ def __repr__(self) -> str: @classmethod @lru_cache() def __class_getitem__(cls: type, typ: type) -> Any: - new_cls = type("State[{}]".format(typ.__name__), (State,), {}) + new_cls = type("_StateManager[{}]".format(typ.__name__), (_StateManager,), {}) new_cls._typ = typ return new_cls +# Attach __class_getitem__ to State for backward-compat: State[T] -> _StateManager[T] +State.__class_getitem__ = classmethod(lru_cache()(lambda cls, typ: _StateManager[typ])) # type: ignore[attr-defined] + + def build_track_state_node(edge: Any, keyby: Union[str, Tuple[str, ...]]) -> Any: @csp.node def _track_state_node( # type: ignore[no-untyped-def] ts: ts[edge.tstype.typ], keyby: object - ) -> ts[State[edge.tstype.typ]]: + ) -> ts[_StateManager[edge.tstype.typ]]: with csp.state(): - s_tracker = State[edge.tstype.typ](keyby) + s_tracker = _StateManager[edge.tstype.typ](keyby) if csp.ticked(ts): s_tracker.insert(ts) csp.output(s_tracker) diff --git a/csp_gateway/server/gateway/gateway.py b/csp_gateway/server/gateway/gateway.py index 9bab60a..ea7b4eb 100644 --- a/csp_gateway/server/gateway/gateway.py +++ b/csp_gateway/server/gateway/gateway.py @@ -6,13 +6,12 @@ from datetime import datetime, timedelta from socket import gethostname from time import sleep -from typing import Any, Callable, List, Optional, Type, Union, get_args, get_origin +from typing import Any, Callable, List, Optional, Type, Union import csp from csp import ts from pydantic import Field, PrivateAttr, create_model, model_validator -from csp_gateway.server.gateway import State from csp_gateway.server.settings import Settings from .csp import Channels, ChannelsFactory, ChannelsType, Module @@ -113,7 +112,6 @@ def _instantiate_dynamic_channel(self, modules: List[Module[GatewayChannels]], c for m in modules: module_dynamic_channels = m.dynamic_channels() if hasattr(m, "dynamic_channels") else None if module_dynamic_channels: - channels_with_state = m.dynamic_state_channels() if hasattr(m, "dynamic_state_channels") else None for n, t in module_dynamic_channels.items(): existing_type = dynamic_channels.get(n, None) if existing_type is not None: @@ -121,11 +119,6 @@ def _instantiate_dynamic_channel(self, modules: List[Module[GatewayChannels]], c raise ValueError(f"Conflicting types for dynamic channel {n}.") dynamic_channels[n] = t - if channels_with_state and n in channels_with_state: - if get_origin(t) is list: - t = get_args(t)[0] - - dynamic_channels[f"s_{n}"] = State[t] if dynamic_channels: dynamic_channel_kwargs = {n: (ts[t], None) for n, t in dynamic_channels.items()} diff --git a/csp_gateway/server/modules/logging/stdlib.py b/csp_gateway/server/modules/logging/stdlib.py index ec55b40..59d56af 100644 --- a/csp_gateway/server/modules/logging/stdlib.py +++ b/csp_gateway/server/modules/logging/stdlib.py @@ -383,7 +383,7 @@ def _convert_log_level(cls, v: Union[str, int]) -> int: def connect(self, channels: ChannelsType): logger_to_use = logging.getLogger(self.log_name) - for field in self.selection.select_from(channels, state_channels=self.log_states): + for field in self.selection.select_from(channels): data = channels.get_channel(field) # list baskets not supported yet if isinstance(data, dict): @@ -393,6 +393,11 @@ def connect(self, channels: ChannelsType): edge = channels.get_channel(field) csp.log(self.log_level, field, edge, logger=logger_to_use) + if self.log_states: + for alias in self.selection.select_from(channels, state_channels=True): + edge = channels.get_state(alias) + csp.log(self.log_level, f"s_{alias}", edge, logger=logger_to_use) + # Backwards compatibility alias StdlibLogging = Logging diff --git a/csp_gateway/server/modules/web/mount.py b/csp_gateway/server/modules/web/mount.py index a28e37a..a609407 100644 --- a/csp_gateway/server/modules/web/mount.py +++ b/csp_gateway/server/modules/web/mount.py @@ -19,6 +19,7 @@ class MountRestRoutes(GatewayModule): mount_last: ChannelSelection = Field(default_factory=ChannelSelection, description="Channels to mount for last operations. Defaults to all") mount_next: ChannelSelection = Field(default_factory=ChannelSelection, description="Channels to mount for next operations. Defaults to all") mount_send: ChannelSelection = Field(default_factory=ChannelSelection, description="Channels to mount for send operations. Defaults to all") + mount_stage: ChannelSelection = Field(default_factory=ChannelSelection, description="Channels to mount for stage operations. Defaults to all") mount_state: ChannelSelection = Field(default_factory=ChannelSelection, description="Channels to mount for state operations. Defaults to all") mount_lookup: ChannelSelection = Field(default_factory=ChannelSelection, description="Channels to mount for lookup operations. Defaults to all") @@ -39,6 +40,7 @@ def rest(self, app: GatewayWebApp) -> None: self._mount_last(app) self._mount_next(app) self._mount_send(app) + self._mount_stage(app) self._mount_state(app) self._mount_lookup(app) @@ -89,9 +91,23 @@ def _mount_send(self, app: GatewayWebApp) -> None: app.add_send_available_channels(seen_channels) + def _mount_stage(self, app: GatewayWebApp) -> None: + staged_channels = set(app.gateway.channels.staged_channels()) + if not staged_channels: + return + + seen_channels = set() + for channel in staged_channels: + if channel not in seen_channels: + seen_channels.add(channel) + app.add_stage_api(channel) + + app.add_stage_available_channels(seen_channels) + def _mount_state(self, app: GatewayWebApp) -> None: selection = ChannelSelection() if self.force_mount_all else self.mount_state - channels_set = set(selection.select_from(app.gateway.channels_model, state_channels=True)) + # Use the instance so dynamically-registered set_state aliases are included. + channels_set = set(selection.select_from(app.gateway.channels, state_channels=True)) seen_channels = set() # Bind every wire @@ -110,7 +126,7 @@ def _mount_state(self, app: GatewayWebApp) -> None: missing_channels = channels_set - seen_channels if missing_channels: - log.info(f"Requested channels missing state routes: {list(channel[2:] for channel in missing_channels)}") + log.info(f"Requested channels missing state routes: {list(missing_channels)}") app.add_state_available_channels(seen_channels) diff --git a/csp_gateway/server/shared/channel_selection.py b/csp_gateway/server/shared/channel_selection.py index 1191f62..89bba44 100644 --- a/csp_gateway/server/shared/channel_selection.py +++ b/csp_gateway/server/shared/channel_selection.py @@ -68,21 +68,24 @@ def select_from( fields = channels.fields() if self.include is None else self.include return list(dict.fromkeys([field for field in fields if field not in self.exclude])) + # State aliases live in their own namespace alongside channel fields. + if state_channels: + if isinstance(channels, type): + # called with the Channels class (declared aliases only) + aliases = list(channels._declared_states.keys()) + else: + aliases = channels.all_state_aliases() + if self.include is not None: + ordered = [a for a in self.include if a in aliases and a not in self.exclude] + else: + ordered = [a for a in aliases if a not in self.exclude] + return list(dict.fromkeys(ordered)) + for idx, field in enumerate(channels.fields()): # avoid duplicates if field in names: continue - # TODO not this `s_` business... - # Return state channels or regular channels - if state_field := field.startswith("s_"): - if not state_channels: - continue - field = field[2:] - else: - if state_channels: - continue - # Check whether static outer_type = channels.get_outer_type(field) if is_dict_basket(outer_type) or isTsType(outer_type): @@ -103,9 +106,6 @@ def select_from( if field in self.exclude: continue - if state_field: - field = f"s_{field}" - names[field] = idx return [k for k, _ in sorted(names.items(), key=lambda x: x[1])] diff --git a/csp_gateway/server/web/app.py b/csp_gateway/server/web/app.py index 770ee22..8cb7735 100644 --- a/csp_gateway/server/web/app.py +++ b/csp_gateway/server/web/app.py @@ -38,6 +38,8 @@ add_next_routes, add_send_available_channels, add_send_routes, + add_stage_available_channels, + add_stage_routes, add_state_available_channels, add_state_routes, ) @@ -132,6 +134,7 @@ def __init__( "lookup": APIRouter(), "next": APIRouter(), "send": APIRouter(), + "stage": APIRouter(), "state": APIRouter(), } @@ -427,6 +430,12 @@ def add_api(self) -> None: tags=["Requests"], dependencies=self._middlewares, ) + api_router.include_router( + self.get_router("stage"), + prefix="/stage", + tags=["Stage"], + dependencies=self._middlewares, + ) api_router.include_router( self.get_router("state"), prefix="/state", @@ -530,44 +539,60 @@ def add_send_available_channels(self, fields: Optional[Set[str]] = None) -> None add_send_available_channels(api_router=api_router, fields=fields) def add_state_api(self, field: str) -> None: - api_router = self.get_router("state") + """Mount REST routes for the given state ``field``. - # Prune s_ from start - name_without_state = field[2:] + ``field`` must be a known state name on the gateway's channels — either + declared via ``Annotated[..., State(...)]`` or registered dynamically + via ``set_state`` during a module's ``connect``. + """ + api_router = self.get_router("state") - dict_basket = self._is_dict_basket_field(field=name_without_state) + spec = self.gateway.channels._states.get(field) or self.gateway.channels_model._declared_states.get(field) + if spec is None: + raise ValueError("Unknown state '{}' on {}".format(field, self.gateway.channels_model.__name__)) - if dict_basket: - dict_basket_key_type, model = dict_basket - subroute_key = dict_basket_key_type + if spec.source_field is not None: + dict_basket = self._is_dict_basket_field(field=spec.source_field) + if dict_basket: + dict_basket_key_type, model = dict_basket + # If the annotation pinned an indexer, expose as a non-keyed route on that one key. + subroute_key = None if spec.indexer is not None else dict_basket_key_type + else: + model = self._get_field_pydantic_type(spec.source_field) + subroute_key = None else: - model = self._get_field_pydantic_type(name_without_state) + # set_state-registered: source is a raw edge with a known type. + state_edge = self.gateway.channels._state_edges.get((field, spec.indexer)) + model = None subroute_key = None - - # Look up the keyby/indexer this state was registered with (if any), - # so we can surface it in the route's OpenAPI description. - keybys = self.gateway.channels._state_keybys - keyby = () - indexer = None - for (registered_field, registered_indexer), registered_keyby in keybys.items(): - if registered_field == field: - keyby = registered_keyby - indexer = registered_indexer - break + if state_edge is not None: + inner = state_edge.tstype.typ + # Unwrap _StateManager[T] -> T for the response model + model = getattr(inner, "_typ", inner) add_state_routes( api_router=api_router, field=field, model=model, subroute_key=subroute_key, - keyby=keyby, - indexer=indexer, + keyby=tuple(spec.keyby) if spec.keyby else (), + indexer=spec.indexer, ) def add_state_available_channels(self, fields: Optional[Set[str]] = None) -> None: api_router = self.get_router("state") add_state_available_channels(api_router=api_router, fields=fields) + def add_stage_api(self, field: str) -> None: + """Mount REST routes for staging on a channel.""" + api_router = self.get_router("stage") + model = self._get_field_pydantic_type(field) + add_stage_routes(api_router=api_router, field=field, model=model) + + def add_stage_available_channels(self, fields: Optional[Set[str]] = None) -> None: + api_router = self.get_router("stage") + add_stage_available_channels(api_router=api_router, fields=fields) + def add_controls_api(self, field: str) -> None: api_router = self.get_router("controls") add_controls_routes(api_router, field=field) diff --git a/csp_gateway/server/web/routes/__init__.py b/csp_gateway/server/web/routes/__init__.py index 6f26484..8b8c791 100644 --- a/csp_gateway/server/web/routes/__init__.py +++ b/csp_gateway/server/web/routes/__init__.py @@ -4,4 +4,5 @@ from .next import * from .send import * from .shared import * +from .stage import * from .state import * diff --git a/csp_gateway/server/web/routes/stage.py b/csp_gateway/server/web/routes/stage.py new file mode 100644 index 0000000..09895b0 --- /dev/null +++ b/csp_gateway/server/web/routes/stage.py @@ -0,0 +1,252 @@ +"""REST routes for the staging API. + +Exposes endpoints under ``/api/v1/stage/`` for managing staged structs +on channels that have staging enabled. + +See STAGE.md for the full API specification. +""" + +import json +import logging +from typing import Dict, List, Optional, Set, Union + +from fastapi import APIRouter, Body, HTTPException, Query, Request +from fastapi.responses import Response +from pydantic import BaseModel + +from csp_gateway.utils import NoProviderException + +from ..utils import get_default_responses + +log = logging.getLogger(__name__) + +__all__ = ( + "add_stage_routes", + "add_stage_available_channels", +) + + +def _serialize_staging_result(result: Dict[str, List]) -> Response: + """Serialize a staging result dict (staging_id -> list of structs) to JSON Response.""" + serialized = {} + for sid, items in result.items(): + serialized[sid] = [json.loads(item.type_adapter().dump_json(item)) for item in items] + return Response( + content=json.dumps(serialized), + media_type="application/json", + ) + + +def add_stage_routes( + api_router: APIRouter, + field: str, + model: Union[BaseModel, List[BaseModel]] = None, +) -> None: + """Mount REST routes for staging on a single channel ``field``.""" + + # POST /stage/ — stage_add + async def stage_add( + request: Request, + id: Optional[str] = Query(None, description="Comma-separated staging IDs to add to"), + data: Optional[model] = Body(None), + ): + """Add a struct to staging area(s). + + - Empty body, no id: create a new empty staging + - Body with struct, no id: add to latest staging or create new + - Body with struct, id=: add to specified staging(s) + """ + try: + staging_ids = _parse_staging_ids(id) + affected = request.app.gateway.channels.stage_add(field, struct=data, staging_ids=staging_ids) + # Return the affected stagings with their contents + result = request.app.gateway.channels.stage_lookup(field) + filtered = {sid: result.get(sid, []) for sid in affected} + return _serialize_staging_result(filtered) + except NoProviderException: + raise HTTPException(status_code=404, detail=f"Staging not enabled for channel: {field}") + except (ValueError, KeyError) as e: + raise HTTPException(status_code=400, detail=str(e)) + + api_router.post( + f"/{field}", + responses=get_default_responses(), + name=f"Stage Add {field}", + description=f"Add struct(s) to staging for channel `{field}`.", + )(stage_add) + + if "_" in field: + api_router.post( + f"/{field.replace('_', '-')}", + responses=get_default_responses(), + include_in_schema=False, + )(stage_add) + + # DELETE /stage/ — stage_remove + async def stage_remove( + request: Request, + id: Optional[str] = Query(None, description="Comma-separated staging IDs to remove from"), + data: Optional[model] = Body(None), + ): + """Remove struct(s) from staging area(s). + + - Empty body, no id: clear latest staging + - Empty body, id=: clear all stagings + - Empty body, id=: clear structs from specific staging + - Body with struct, no id: remove struct from latest staging containing it + - Body with struct, id=: remove struct from all stagings + - Body with struct, id=: remove struct from specific staging + """ + try: + staging_ids = _parse_staging_ids(id) + affected = request.app.gateway.channels.stage_remove(field, struct=data, staging_ids=staging_ids) + # Return affected staging IDs mapped to their remaining contents + result = request.app.gateway.channels.stage_lookup(field) + filtered = {sid: result.get(sid, []) for sid in affected} + return _serialize_staging_result(filtered) + except NoProviderException: + raise HTTPException(status_code=404, detail=f"Staging not enabled for channel: {field}") + except (ValueError, KeyError) as e: + raise HTTPException(status_code=400, detail=str(e)) + + api_router.delete( + f"/{field}", + responses=get_default_responses(), + name=f"Stage Remove {field}", + description=f"Remove struct(s) from staging for channel `{field}`.", + )(stage_remove) + + if "_" in field: + api_router.delete( + f"/{field.replace('_', '-')}", + responses=get_default_responses(), + include_in_schema=False, + )(stage_remove) + + # PATCH /stage/ — stage_release + async def stage_release( + request: Request, + id: Optional[str] = Query(None, description="Comma-separated staging IDs to release"), + ): + """Release staged structs into the channel. + + - No id: release all stagings + - id=: release specific staging(s) + """ + try: + staging_ids = _parse_staging_ids(id) + released = request.app.gateway.channels.stage_release(field, staging_ids=staging_ids) + return _serialize_staging_result(released) + except NoProviderException: + raise HTTPException(status_code=404, detail=f"Staging not enabled for channel: {field}") + except (ValueError, KeyError) as e: + raise HTTPException(status_code=400, detail=str(e)) + + api_router.patch( + f"/{field}", + responses=get_default_responses(), + name=f"Stage Release {field}", + description=f"Release staged structs into channel `{field}`.", + )(stage_release) + + if "_" in field: + api_router.patch( + f"/{field.replace('_', '-')}", + responses=get_default_responses(), + include_in_schema=False, + )(stage_release) + + # GET /stage/ — stage_list + async def stage_list( + request: Request, + id: Optional[str] = Query(None, description="Specific staging ID to check"), + ) -> List[str]: + """List staging IDs for a channel. + + - No id: list all staging IDs + - id=: check if specific staging exists + """ + try: + return request.app.gateway.channels.stage_list(field, staging_id=id) + except NoProviderException: + raise HTTPException(status_code=404, detail=f"Staging not enabled for channel: {field}") + + api_router.get( + f"/{field}", + responses=get_default_responses(), + response_model=List[str], + name=f"Stage List {field}", + description=f"List staging IDs for channel `{field}`.", + )(stage_list) + + if "_" in field: + api_router.get( + f"/{field.replace('_', '-')}", + responses=get_default_responses(), + response_model=List[str], + include_in_schema=False, + )(stage_list) + + # PUT /stage/ — stage_lookup + async def stage_lookup( + request: Request, + id: Optional[str] = Query(None, description="Specific staging ID to look up"), + ): + """Look up contents of staging area(s). + + - No id: list all staging contents + - id=: list contents of specific staging + """ + try: + result = request.app.gateway.channels.stage_lookup(field, staging_id=id) + return _serialize_staging_result(result) + except NoProviderException: + raise HTTPException(status_code=404, detail=f"Staging not enabled for channel: {field}") + + api_router.put( + f"/{field}", + responses=get_default_responses(), + name=f"Stage Lookup {field}", + description=f"Look up staged contents for channel `{field}`.", + )(stage_lookup) + + if "_" in field: + api_router.put( + f"/{field.replace('_', '-')}", + responses=get_default_responses(), + include_in_schema=False, + )(stage_lookup) + + +def add_stage_available_channels( + api_router: APIRouter, + fields: Optional[Set[str]] = None, +) -> None: + """Mount the top-level GET /stage/ route listing all staged channels.""" + + @api_router.get( + "/", + responses=get_default_responses(), + response_model=List[str], + ) + async def get_staged_channels(request: Request) -> List[str]: + """Return the list of all channels with staging enabled.""" + all_staged = request.app.gateway.channels.staged_channels() + if fields is not None: + return sorted(name for name in all_staged if name in fields) + return sorted(all_staged) + + +def _parse_staging_ids(id_param: Optional[str]) -> Optional[List[str]]: + """Parse the comma-separated id query parameter. + + Returns: + - None if id_param is None (not provided) + - [] if id_param is empty string + - list of IDs otherwise + """ + if id_param is None: + return None + if id_param == "": + return [] + return [s.strip() for s in id_param.split(",") if s.strip()] diff --git a/csp_gateway/server/web/routes/state.py b/csp_gateway/server/web/routes/state.py index 6be4bfa..88cc20a 100644 --- a/csp_gateway/server/web/routes/state.py +++ b/csp_gateway/server/web/routes/state.py @@ -24,12 +24,20 @@ def add_state_routes( keyby: Tuple[str, ...] = (), indexer: Optional[Union[str, int]] = None, ) -> None: + """Mount REST routes for a single state ``field``. + + The state mechanism is now name-based; there is no implicit ``s_`` prefix. + Routes are exposed at ``/state/`` (and a variant with ``_`` + replaced by ``-`` for URL friendliness). + + ``keyby`` / ``indexer`` are surfaced in the route description so the + accumulation semantics are visible in the OpenAPI schema. + """ if model and get_origin(model) is list: list_model = model else: list_model = List[model] - # Get the fully qualified type name for the description fq_type_name = get_fully_qualified_type_name(model) keyby_lines = [] @@ -40,92 +48,42 @@ def add_state_routes( keyby_prefix = ("\n\n".join(keyby_lines) + "\n\n") if keyby_lines else "" if subroute_key: - # Prune s_ from start - name_without_state = field[2:] async def get_state(key: str, query: Optional[Query] = query_json(), request: Request = None) -> list_model: # type: ignore[valid-type] """ Get state value on a dictionary basket channel, where `key` is the key of the dictionary basket. If such a key does not exist or is not mounted, this endpoint will raise a `404` error. - States may be queried by certain conditions. Currently only filtering is supported. Filters can be used to evaluate - objects in state and compare them to either scalar values or other attributes on the object. Here are some simple examples - using the demo application provided with `csp-gateway`: - - ``` - # Filter only records where `record.x` == 5 - api/v1/state/example?query={"filters":[{"attr":"x","by":{"value":5,"where":"=="}}]} - - - # Filter only records where `record.x` < 10 - /api/v1/state/example?query={"filters":[{"attr":"x","by":{"value":10,"where":"<"}}]} - - # Filter only records where `record.timestamp` < "2023-03-30T14:45:26.394000" - /api/v1/state/example?query={"filters":[{"attr":"timestamp","by":{"when":"2023-03-30T14:45:26.394000","where":"<"}}]} - - # Filter only records where `record.id` < `record.y` - /api/v1/state/example?query={"filters":[{"attr":"id","by":{"attr":"y","where":"<"}}]} - - # Filter only records where `record.x` < 50 and `record.x` >= 30 - /api/v1/state/example?query={"filters":[{"attr":"x","by":{"value":50,"where":"<"}},{"attr":"x","by":{"value":30,"where":">="}}]} - ``` + States may be queried by certain conditions. See `/state/` for query examples. """ - channel = request.app.gateway.channels._ensure_state_field(field) - - if not hasattr(request.app.gateway.channels, channel): - raise HTTPException( - status_code=404, - detail="State channel not found: {}".format(channel), - ) - - # FIXME this wont work, is dicts try: - res = request.app.gateway.channels.query( - getattr(request.app.gateway.channels_model, channel), - key, - query=query, - ) + res = request.app.gateway.channels.query(field, key, query=query) except NoProviderException: raise HTTPException( status_code=404, - detail="Channel not found: {}/{}".format(field, key), + detail="State not found: {}/{}".format(field, key), ) return prepare_response(res, is_list_model=True) api_router.get( - "/{}/{{key:path}}".format(name_without_state), + "/{}/{{key:path}}".format(field), responses=get_default_responses(), response_model=list_model, # type: ignore[valid-type] - name="Get State {} by key".format(name_without_state), + name="Get State {} by key".format(field), description=keyby_prefix + cleandoc(get_state.__doc__ or ""), openapi_extra={"type_": fq_type_name} if fq_type_name else None, )(get_state) - api_router.get( - "/{}/{{key:path}}".format(name_without_state.replace("_", "-")), - responses=get_default_responses(), - response_model=list_model, # type: ignore[valid-type] - include_in_schema=False, - )(get_state) - - api_router.get( - "/{}/{{key:path}}".format(field), - responses=get_default_responses(), - response_model=list_model, # type: ignore[valid-type] - include_in_schema=False, - )(get_state) - - api_router.get( - "/{}/{{key:path}}".format(field.replace("_", "-")), - responses=get_default_responses(), - response_model=list_model, # type: ignore[valid-type] - include_in_schema=False, - )(get_state) + if "_" in field: + api_router.get( + "/{}/{{key:path}}".format(field.replace("_", "-")), + responses=get_default_responses(), + response_model=list_model, # type: ignore[valid-type] + include_in_schema=False, + )(get_state) if model: - # Prune s_ from start - name_without_state = field[2:] async def get_state(query: Optional[Query] = query_json(), request: Request = None) -> list_model: # type: ignore[misc, valid-type] """Get state value on a non-dict basket channel. This endpoint will flatten the state structure and return a list of the elements. @@ -137,71 +95,47 @@ async def get_state(query: Optional[Query] = query_json(), request: Request = No ``` # Filter only records where `record.x` == 5 - api/v1/state/example?query={"filters":[{"attr":"x","by":{"value":5,"where":"=="}}]} + api/v1/state/example_with_state?query={"filters":[{"attr":"x","by":{"value":5,"where":"=="}}]} # Filter only records where `record.x` < 10 - /api/v1/state/example?query={"filters":[{"attr":"x","by":{"value":10,"where":"<"}}]} + /api/v1/state/example_with_state?query={"filters":[{"attr":"x","by":{"value":10,"where":"<"}}]} # Filter only records where `record.timestamp` < "2023-03-30T14:45:26.394000" - /api/v1/state/example?query={"filters":[{"attr":"timestamp","by":{"when":"2023-03-30T14:45:26.394000","where":"<"}}]} + /api/v1/state/example_with_state?query={"filters":[{"attr":"timestamp","by":{"when":"2023-03-30T14:45:26.394000","where":"<"}}]} # Filter only records where `record.id` < `record.y` - /api/v1/state/example?query={"filters":[{"attr":"id","by":{"attr":"y","where":"<"}}]} + /api/v1/state/example_with_state?query={"filters":[{"attr":"id","by":{"attr":"y","where":"<"}}]} # Filter only records where `record.x` < 50 and `record.x` >= 30 - /api/v1/state/example?query={"filters":[{"attr":"x","by":{"value":50,"where":"<"}},{"attr":"x","by":{"value":30,"where":">="}}]} + /api/v1/state/example_with_state?query={"filters":[{"attr":"x","by":{"value":50,"where":"<"}},{"attr":"x","by":{"value":30,"where":">="}}]} ``` """ - channel = request.app.gateway.channels._ensure_state_field(field) - - if not hasattr(request.app.gateway.channels, channel): - raise HTTPException( - status_code=404, - detail="State channel not found: {}".format(channel), - ) - try: - res = request.app.gateway.channels.query( - getattr(request.app.gateway.channels_model, channel), - query=query, - ) + res = request.app.gateway.channels.query(field, query=query) except NoProviderException: raise HTTPException( status_code=404, - detail="Channel not found: {}".format(channel), + detail="State not found: {}".format(field), ) return prepare_response(res, is_list_model=True) api_router.get( - "/{}".format(name_without_state), + "/{}".format(field), responses=get_default_responses(), response_model=list_model, # type: ignore[valid-type] - name="Get State {}".format(name_without_state), + name="Get State {}".format(field), description=keyby_prefix + cleandoc(get_state.__doc__ or ""), openapi_extra={"type_": fq_type_name} if fq_type_name else None, )(get_state) - api_router.get( - "/{}".format(name_without_state.replace("_", "-")), - responses=get_default_responses(), - response_model=list_model, # type: ignore[valid-type] - include_in_schema=False, - )(get_state) - - api_router.get( - "/{}".format(field), - responses=get_default_responses(), - response_model=list_model, # type: ignore[valid-type] - include_in_schema=False, - )(get_state) - - api_router.get( - "/{}".format(field.replace("_", "-")), - responses=get_default_responses(), - response_model=list_model, # type: ignore[valid-type] - include_in_schema=False, - )(get_state) + if "_" in field: + api_router.get( + "/{}".format(field.replace("_", "-")), + responses=get_default_responses(), + response_model=list_model, # type: ignore[valid-type] + include_in_schema=False, + )(get_state) def add_state_available_channels( @@ -215,13 +149,10 @@ def add_state_available_channels( ) async def get_state(request: Request) -> List[str]: """ - This endpoint will return a list of string values of all available state channels under the `/state` route. + Return the list of all available state names under the `/state` route. - Note: This endpoint does not support filtering + Note: This endpoint does not support filtering. """ - # TODO: Change state channel stuff return sorted( - field[2:] - for field in ChannelSelection().select_from(request.app.gateway.channels, state_channels=True) - if fields is None or field in fields + name for name in ChannelSelection().select_from(request.app.gateway.channels, state_channels=True) if fields is None or name in fields ) diff --git a/csp_gateway/server/web/templates/channels_graph.html.j2 b/csp_gateway/server/web/templates/channels_graph.html.j2 index 6b6d9b6..ae61fdf 100644 --- a/csp_gateway/server/web/templates/channels_graph.html.j2 +++ b/csp_gateway/server/web/templates/channels_graph.html.j2 @@ -9,7 +9,7 @@ "dagre-d3-es": "https://cdn.jsdelivr.net/npm/dagre-d3-es/+esm" } } - + @@ -35,7 +35,7 @@ top: 0; display: flex; } - + div.dagred3 svg { position: absolute; left: 0; @@ -44,7 +44,7 @@ top: 0; display: flex; } - + .dagred3 .node rect, .dagred3 .node circle, .dagred3 .node ellipse, @@ -52,21 +52,21 @@ stroke: #555; fill: #fff; } - + .dagred3 .edgePath path { stroke: #555; fill: transparent; stroke-width: 1.5px; } - + .dagred3 .edgeLabel foreignObject { overflow: visible; } - + .dagred3 .edgeLabel foreignObject u { text-decoration: none; } - + .dagred3 .node text { pointer-events: none; } @@ -74,27 +74,24 @@ \ No newline at end of file diff --git a/csp_gateway/testing/harness.py b/csp_gateway/testing/harness.py index e4f3204..c0f98c5 100644 --- a/csp_gateway/testing/harness.py +++ b/csp_gateway/testing/harness.py @@ -418,10 +418,7 @@ def _run_test_events(self, event: ts[BaseGatewayTestEvent], channel_values: Dict if channel not in s_values: s_values[channel] = [] - if isinstance(channel, str) and channel.startswith("s_"): - s_values[channel].append((csp.now(), value.query())) - else: - s_values[channel].append((csp.now(), value)) + s_values[channel].append((csp.now(), value)) if csp.ticked(event): try: diff --git a/csp_gateway/testing/shared_helpful_classes.py b/csp_gateway/testing/shared_helpful_classes.py index 729da8c..060115b 100644 --- a/csp_gateway/testing/shared_helpful_classes.py +++ b/csp_gateway/testing/shared_helpful_classes.py @@ -1,5 +1,5 @@ from datetime import timedelta -from typing import Dict, List, Optional, Type +from typing import Annotated, Dict, List, Optional, Type import csp import numpy as np @@ -47,10 +47,8 @@ class MyGatewayChannels(GatewayChannels): my_static: float = 0.0 my_static_dict: Dict[str, float] = {} my_static_list: List[str] = [] - my_channel: ts[MyStruct] = None - s_my_channel: ts[State[MyStruct]] = None - my_list_channel: ts[List[MyStruct]] = None - s_my_list_channel: ts[State[MyStruct]] = None + my_channel: Annotated[ts[MyStruct], State(keyby="id")] = None + my_list_channel: Annotated[ts[List[MyStruct]], State(keyby="id")] = None my_enum_basket: Dict[MyEnum, ts[MyStruct]] = None my_str_basket: Dict[str, ts[MyStruct]] = None my_enum_basket_list: Dict[MyEnum, ts[List[MyStruct]]] = None @@ -58,8 +56,7 @@ class MyGatewayChannels(GatewayChannels): class MySmallGatewayChannels(GatewayChannels): - example: ts[int] = None - s_example: ts[State[int]] = None + example: Annotated[ts[int], State(keyby="id")] = None my_str_basket: Dict[str, ts[float]] = None @@ -88,9 +85,7 @@ def connect(self, channels: MyGatewayChannels) -> None: self.my_list_data, ) channels.add_send_channel(MyGatewayChannels.my_channel) - channels.set_state(MyGatewayChannels.my_channel, "id") channels.add_send_channel(MyGatewayChannels.my_list_channel) - channels.set_state(MyGatewayChannels.my_list_channel, "id") channels.set_channel(MyGatewayChannels.my_array_channel, csp.const(np.array([1.0, 2.0]))) @@ -180,7 +175,6 @@ def connect(self, channels: MySmallGatewayChannels): self.subscribe(csp.timer(interval=self.interval, value=True)), ) channels.add_send_channel(MySmallGatewayChannels.example) - channels.set_state(MySmallGatewayChannels.example, "id") class MyDictBasketModule(GatewayModule): diff --git a/csp_gateway/tests/server/gateway/csp/test_channels.py b/csp_gateway/tests/server/gateway/csp/test_channels.py index d312f90..c222b92 100644 --- a/csp_gateway/tests/server/gateway/csp/test_channels.py +++ b/csp_gateway/tests/server/gateway/csp/test_channels.py @@ -1,5 +1,5 @@ from datetime import date, datetime -from typing import Dict, List, Optional, get_type_hints +from typing import Annotated, Dict, List, Optional, get_type_hints import csp import numpy as np @@ -58,10 +58,8 @@ class MyGatewayChannels(GatewayChannels): my_static: float = 0.0 my_static_dict: Dict[str, float] = {} my_static_list: List[str] = [] - my_channel: ts[MyStruct] = None - s_my_channel: ts[State[MyStruct]] = None - my_list_channel: ts[List[MyStruct]] = None - s_my_list_channel: ts[State[MyStruct]] = None + my_channel: Annotated[ts[MyStruct], State(keyby="id")] = None + my_list_channel: Annotated[ts[List[MyStruct]], State(keyby="id")] = None my_enum_basket: Dict[MyEnum, ts[MyStruct]] = None my_str_basket: Dict[str, ts[MyStruct]] = None my_enum_basket_list: Dict[MyEnum, ts[List[MyStruct]]] = None diff --git a/csp_gateway/tests/server/gateway/csp/test_stage.py b/csp_gateway/tests/server/gateway/csp/test_stage.py new file mode 100644 index 0000000..e23f4e3 --- /dev/null +++ b/csp_gateway/tests/server/gateway/csp/test_stage.py @@ -0,0 +1,497 @@ +"""Tests for the staging functionality.""" + +from datetime import datetime, timedelta +from typing import Annotated, Type + +import csp +import pytest +from csp import ts + +from csp_gateway import ( + Channels, + Gateway, + GatewayChannels, + GatewayModule, + GatewayStruct, + State, +) +from csp_gateway.server.gateway.csp.stage import Stage, Staging, _StageManager +from csp_gateway.testing import GatewayTestHarness +from csp_gateway.utils import NoProviderException + +# --- Test structures --- + + +class OrderStruct(GatewayStruct): + symbol: str = "" + quantity: int = 0 + price: float = 0.0 + + +class StagedChannels(GatewayChannels): + orders: ts[OrderStruct] = None + + +class StagedWithStateChannels(GatewayChannels): + orders: Annotated[ts[OrderStruct], State(keyby="id")] = None + + +class StagingModule(GatewayModule): + """Module that enables staging on the orders channel.""" + + def connect(self, channels: StagedChannels) -> None: + channels.set_channel(StagedChannels.orders, csp.null_ts(OrderStruct)) + channels.set_stage(StagedChannels.orders) + channels.add_send_channel(StagedChannels.orders) + + def shutdown(self) -> None: + pass + + +class StagingWithStateModule(GatewayModule): + """Module that enables staging + state on orders channel.""" + + def connect(self, channels: StagedWithStateChannels) -> None: + channels.set_channel(StagedWithStateChannels.orders, csp.null_ts(OrderStruct)) + channels.set_stage(StagedWithStateChannels.orders) + channels.add_send_channel(StagedWithStateChannels.orders) + + def shutdown(self) -> None: + pass + + +# --- Unit Tests for _StageManager class --- + + +class TestStagingArea: + def test_is_gateway_struct(self): + area = Staging() + assert isinstance(area, GatewayStruct) + assert area.id # auto-generated + assert area.timestamp + + def test_add_and_lookup(self): + area = Staging() + s = OrderStruct(symbol="AAPL", quantity=100, price=150.0) + area.add(s) + assert area.lookup() == [s] + + def test_remove(self): + area = Staging() + s = OrderStruct(symbol="AAPL", quantity=100, price=150.0) + area.add(s) + assert area.remove(s) is True + assert area.lookup() == [] + + def test_remove_not_found(self): + area = Staging() + s = OrderStruct(symbol="AAPL", quantity=100, price=150.0) + assert area.remove(s) is False + + def test_clear(self): + area = Staging() + s1 = OrderStruct(symbol="AAPL", quantity=100, price=150.0) + s2 = OrderStruct(symbol="GOOG", quantity=50, price=2800.0) + area.add(s1) + area.add(s2) + cleared = area.clear() + assert cleared == [s1, s2] + assert area.lookup() == [] + + +class TestStageManager: + def test_stage_add_none_none_creates_empty(self): + stage = _StageManager() + result = stage.stage_add(None, None) + assert len(result) == 1 + assert stage.stage_list() == result + + def test_stage_add_none_empty_list_creates_empty(self): + stage = _StageManager() + result = stage.stage_add(None, []) + assert len(result) == 1 + + def test_stage_add_struct_none_creates_or_appends(self): + stage = _StageManager() + s = OrderStruct(symbol="AAPL", quantity=100, price=150.0) + # No existing staging -> creates new + result = stage.stage_add(s, None) + assert len(result) == 1 + sid = result[0] + contents = stage.stage_lookup(sid) + assert contents[sid] == [s] + + # Existing staging -> appends to latest + s2 = OrderStruct(symbol="GOOG", quantity=50, price=2800.0) + result2 = stage.stage_add(s2, None) + assert result2 == [sid] + contents = stage.stage_lookup(sid) + assert len(contents[sid]) == 2 + + def test_stage_add_struct_empty_list_adds_to_all(self): + stage = _StageManager() + stage.stage_add(None, None) # create first + stage.stage_add(None, None) # create second + ids = stage.stage_list() + assert len(ids) == 2 + + s = OrderStruct(symbol="AAPL", quantity=100, price=150.0) + result = stage.stage_add(s, []) + assert set(result) == set(ids) + + def test_stage_add_struct_specific_ids(self): + stage = _StageManager() + ids = stage.stage_add(None, None) + sid = ids[0] + s = OrderStruct(symbol="AAPL", quantity=100, price=150.0) + result = stage.stage_add(s, [sid]) + assert result == [sid] + contents = stage.stage_lookup(sid) + assert contents[sid] == [s] + + def test_stage_add_none_with_ids_errors(self): + stage = _StageManager() + ids = stage.stage_add(None, None) + with pytest.raises(ValueError): + stage.stage_add(None, ids) + + def test_stage_add_nonexistent_id_errors(self): + stage = _StageManager() + s = OrderStruct(symbol="AAPL", quantity=100, price=150.0) + with pytest.raises(KeyError): + stage.stage_add(s, ["nonexistent"]) + + def test_stage_remove_none_none_clears_latest(self): + stage = _StageManager() + s = OrderStruct(symbol="AAPL", quantity=100, price=150.0) + ids1 = stage.stage_add(s, None) + ids2 = stage.stage_add(None, None) + + result = stage.stage_remove(None, None) + assert result == ids2 + assert stage.stage_list() == ids1 + + def test_stage_remove_none_empty_clears_all(self): + stage = _StageManager() + stage.stage_add(None, None) + stage.stage_add(None, None) + result = stage.stage_remove(None, []) + assert len(result) == 2 + assert stage.stage_list() == [] + + def test_stage_remove_none_specific_clears_that_staging(self): + stage = _StageManager() + s = OrderStruct(symbol="AAPL", quantity=100, price=150.0) + ids = stage.stage_add(s, None) + stage.stage_remove(None, ids) + # Staging still exists but is empty + contents = stage.stage_lookup(ids[0]) + assert contents[ids[0]] == [] + + def test_stage_remove_struct_none_removes_from_latest(self): + stage = _StageManager() + s = OrderStruct(symbol="AAPL", quantity=100, price=150.0) + ids = stage.stage_add(s, None) + result = stage.stage_remove(s, None) + assert result == ids + contents = stage.stage_lookup(ids[0]) + assert contents[ids[0]] == [] + + def test_stage_remove_struct_empty_removes_from_all(self): + stage = _StageManager() + s = OrderStruct(symbol="AAPL", quantity=100, price=150.0) + stage.stage_add(s, None) + stage.stage_add(None, None) + # Add to second staging too + all_ids = stage.stage_list() + stage.stage_add(s, [all_ids[1]]) + + result = stage.stage_remove(s, []) + assert set(result) == set(all_ids) + + def test_stage_release_all(self): + stage = _StageManager() + s1 = OrderStruct(symbol="AAPL", quantity=100, price=150.0) + s2 = OrderStruct(symbol="GOOG", quantity=50, price=2800.0) + stage.stage_add(s1, None) + stage.stage_add(None, None) + ids = stage.stage_list() + stage.stage_add(s2, [ids[1]]) + + released = stage.stage_release(None) + assert ids[0] in released + assert ids[1] in released + assert released[ids[0]] == [s1] + assert released[ids[1]] == [s2] + assert stage.stage_list() == [] + + def test_stage_release_specific(self): + stage = _StageManager() + s = OrderStruct(symbol="AAPL", quantity=100, price=150.0) + ids = stage.stage_add(s, None) + stage.stage_add(None, None) + + released = stage.stage_release(ids) + assert ids[0] in released + assert released[ids[0]] == [s] + # Second staging still exists + remaining = stage.stage_list() + assert len(remaining) == 1 + assert ids[0] not in remaining + + def test_stage_list(self): + stage = _StageManager() + assert stage.stage_list() == [] + ids1 = stage.stage_add(None, None) + ids2 = stage.stage_add(None, None) + assert stage.stage_list() == ids1 + ids2 + + def test_stage_list_specific(self): + stage = _StageManager() + ids = stage.stage_add(None, None) + assert stage.stage_list(ids[0]) == ids + assert stage.stage_list("nonexistent") == [] + + def test_stage_lookup_all(self): + stage = _StageManager() + s = OrderStruct(symbol="AAPL", quantity=100, price=150.0) + ids = stage.stage_add(s, None) + result = stage.stage_lookup() + assert ids[0] in result + assert result[ids[0]] == [s] + + def test_stage_lookup_specific(self): + stage = _StageManager() + s = OrderStruct(symbol="AAPL", quantity=100, price=150.0) + ids = stage.stage_add(s, None) + result = stage.stage_lookup(ids[0]) + assert result == {ids[0]: [s]} + + def test_stage_lookup_nonexistent(self): + stage = _StageManager() + result = stage.stage_lookup("nonexistent") + assert result == {} + + +# --- Integration tests with Gateway --- + + +class TestStagingChannels: + def _build_channels(self) -> StagedChannels: + channels = StagedChannels() + with channels._connection_context("StagingTest"): + channels.set_stage(StagedChannels.orders) + return channels + + def test_set_stage_basic(self): + channels = self._build_channels() + assert "orders" in channels.staged_channels() + + def test_stage_add_and_lookup(self): + channels = self._build_channels() + s = OrderStruct(symbol="AAPL", quantity=100, price=150.0) + ids = channels.stage_add("orders", s) + assert len(ids) == 1 + contents = channels.stage_lookup("orders", ids[0]) + assert contents[ids[0]] == [s] + + def test_stage_release_and_clear(self): + channels = self._build_channels() + s = OrderStruct(symbol="AAPL", quantity=100, price=150.0) + channels.stage_add("orders", s) + removed = channels.stage_remove("orders", s, []) + assert len(removed) == 1 + assert channels.stage_list("orders") == removed + looked_up = channels.stage_lookup("orders", removed[0]) + assert looked_up[removed[0]] == [] + + def test_stage_not_enabled_raises(self): + channels = self._build_channels() + with pytest.raises(NoProviderException): + channels.stage_add("nonexistent", None) + + def test_stage_with_state_channels_model(self): + channels = StagedWithStateChannels() + with channels._connection_context("StagingStateTest"): + channels.set_stage(StagedWithStateChannels.orders) + assert "orders" in channels.staged_channels() + + +class TestStagingAnnotation: + """Tests for Stage() annotation marker — symmetric with State().""" + + def test_annotation_declares_staging(self): + """Stage() in Annotated auto-wires staging during finalization.""" + + class AnnotatedChannels(GatewayChannels): + orders: Annotated[ts[OrderStruct], Stage()] = None + + assert "orders" in AnnotatedChannels._declared_stages + + def test_annotation_wires_staging_via_harness(self): + """Annotation-declared staging is wired during finalization like State.""" + + class AnnotatedChannels(GatewayChannels): + orders: Annotated[ts[OrderStruct], Stage()] = None + + class AnnotatedModule(GatewayModule): + def connect(self, channels: AnnotatedChannels) -> None: + channels.set_channel(AnnotatedChannels.orders, csp.null_ts(OrderStruct)) + + class AnnotatedGateway(Gateway): + channels_model: Type[Channels] = AnnotatedChannels # type: ignore[assignment] + + import socket + + from csp_gateway import GatewaySettings, MountRestRoutes + + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("", 0)) + s.listen(1) + port = s.getsockname()[1] + + module = AnnotatedModule() + gateway = AnnotatedGateway( + modules=[module, MountRestRoutes(force_mount_all=True)], + channels=AnnotatedChannels(), + settings=GatewaySettings(PORT=port), + ) + gateway.start(rest=True, _in_test=True) + try: + assert "orders" in gateway.channels.staged_channels() + # Can use staging API + s = OrderStruct(symbol="AAPL", quantity=100, price=150.0) + ids = gateway.channels.stage_add("orders", s) + assert len(ids) == 1 + finally: + gateway.stop() + + def test_set_stage_and_annotation_dont_conflict(self): + """If a channel has both annotation and set_stage, only one stage is created.""" + + class DualChannels(GatewayChannels): + orders: Annotated[ts[OrderStruct], Stage()] = None + + class DualModule(GatewayModule): + def connect(self, channels: DualChannels) -> None: + channels.set_channel(DualChannels.orders, csp.null_ts(OrderStruct)) + channels.set_stage(DualChannels.orders) + + class DualGateway(Gateway): + channels_model: Type[Channels] = DualChannels # type: ignore[assignment] + + import socket + + from csp_gateway import GatewaySettings, MountRestRoutes + + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("", 0)) + s.listen(1) + port = s.getsockname()[1] + + module = DualModule() + gateway = DualGateway( + modules=[module, MountRestRoutes(force_mount_all=True)], + channels=DualChannels(), + settings=GatewaySettings(PORT=port), + ) + gateway.start(rest=True, _in_test=True) + try: + assert "orders" in gateway.channels.staged_channels() + assert gateway.channels.staged_channels().count("orders") == 1 + finally: + gateway.stop() + + +# --- Test with harness --- + + +class TestStagingHarness: + def test_harness_with_staging(self): + """Test that the harness works with a staging-enabled module.""" + + class HarnessChannels(GatewayChannels): + orders: ts[OrderStruct] = None + + class HarnessStagingModule(GatewayModule): + def connect(self, channels: HarnessChannels) -> None: + channels.set_stage(HarnessChannels.orders) + + def shutdown(self) -> None: + pass + + class HarnessGateway(Gateway): + channels_model: Type[Channels] = HarnessChannels # type: ignore[assignment] + + harness = GatewayTestHarness( + test_channels=["orders"], + ) + s = OrderStruct(symbol="AAPL", quantity=100, price=150.0) + harness.send("orders", s) + harness.delay(timedelta(seconds=1)) + harness.assert_equal("orders", s) + + module = HarnessStagingModule() + gateway = HarnessGateway(modules=[module, harness], channels=HarnessChannels()) + csp.run(gateway.graph, starttime=datetime(2020, 1, 1), endtime=timedelta(seconds=5)) + + def test_stage_release_pushes_to_channel(self): + """Test that stage_release via the gateway pushes items into the csp channel. + + Note: GenericPushAdapter.push_tick only works in realtime mode. + We use gateway.start(rest=True, _in_test=True) which runs csp on a thread. + """ + + class ReleaseChannels(GatewayChannels): + orders: ts[OrderStruct] = None + + class ReleaseStagingModule(GatewayModule): + def connect(self, channels: ReleaseChannels) -> None: + channels.set_stage(ReleaseChannels.orders) + channels.add_send_channel(ReleaseChannels.orders) + + def shutdown(self) -> None: + pass + + class ReleaseGateway(Gateway): + channels_model: Type[Channels] = ReleaseChannels # type: ignore[assignment] + + import socket + + from csp_gateway import GatewaySettings, MountRestRoutes + + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.bind(("", 0)) + s.listen(1) + port = s.getsockname()[1] + + module = ReleaseStagingModule() + gateway = ReleaseGateway( + modules=[module, MountRestRoutes(force_mount_all=True)], + channels=ReleaseChannels(), + settings=GatewaySettings(PORT=port), + ) + + gateway.start(rest=True, _in_test=True) + try: + # Stage items + s1 = OrderStruct(symbol="AAPL", quantity=100, price=150.0) + s2 = OrderStruct(symbol="GOOG", quantity=50, price=2800.0) + ids = gateway.channels.stage_add("orders", s1) + gateway.channels.stage_add("orders", s2, staging_ids=ids) + + # Verify staged content + contents = gateway.channels.stage_lookup("orders", ids[0]) + assert len(contents[ids[0]]) == 2 + + # Release - this calls push_tick internally + released = gateway.channels.stage_release("orders", staging_ids=ids) + assert ids[0] in released + assert len(released[ids[0]]) == 2 + assert released[ids[0]][0].symbol == "AAPL" + assert released[ids[0]][1].symbol == "GOOG" + + # After release, staging should be empty + assert gateway.channels.stage_list("orders") == [] + finally: + gateway.stop() diff --git a/csp_gateway/tests/server/gateway/csp/test_state.py b/csp_gateway/tests/server/gateway/csp/test_state.py index 9f80e7a..8f1b1d3 100644 --- a/csp_gateway/tests/server/gateway/csp/test_state.py +++ b/csp_gateway/tests/server/gateway/csp/test_state.py @@ -99,6 +99,22 @@ def test_state_keyby_unset(): assert res == [ts] +def test_state_keyby_dot_path(): + for specialization, struct_type in zip(_STATES, _STRUCTS): + s = State[struct_type](keyby="g.suba") + assert s.state_type() == specialization + ts1 = struct_type(a=0, b="hello", g=CspSubStruct(suba=1)) + ts2 = struct_type(a=1, b="world", g=CspSubStruct(suba=2)) + s.insert(ts1) + s.insert(ts2) + assert sorted(s.query(), key=lambda r: r.g.suba) == [ts1, ts2] + + # Overwriting with same nested key replaces the record. + ts3 = struct_type(a=2, b="again", g=CspSubStruct(suba=1)) + s.insert(ts3) + assert sorted(s.query(), key=lambda r: r.g.suba) == [ts3, ts2] + + def test_state_keyby_two(): for specialization, struct_type in zip(_STATES, _STRUCTS): s = State[struct_type](keyby=("a", "b")) diff --git a/csp_gateway/tests/server/gateway/test_gateway.py b/csp_gateway/tests/server/gateway/test_gateway.py index 51c2656..dd28439 100644 --- a/csp_gateway/tests/server/gateway/test_gateway.py +++ b/csp_gateway/tests/server/gateway/test_gateway.py @@ -3,7 +3,7 @@ import time from datetime import datetime, timedelta from io import StringIO -from typing import Any, Callable, Dict, List, Optional, Set, Type, Union +from typing import Annotated, Any, Callable, Dict, List, Optional, Set, Type, Union import csp import numpy as np @@ -50,10 +50,8 @@ class MyGatewayChannels(GatewayChannels): my_static_dict: Dict[str, float] = {} my_static_list: List[str] = [] my_static_dict_of_objects: Dict[str, Any] = {} - my_channel: ts[MyStruct] = None - s_my_channel: ts[State[MyStruct]] = None - my_list_channel: ts[List[MyStruct]] = None - s_my_list_channel: ts[State[MyStruct]] = None + my_channel: Annotated[ts[MyStruct], State(keyby="id")] = None + my_list_channel: Annotated[ts[List[MyStruct]], State(keyby="id")] = None my_enum_basket: Dict[MyEnum, ts[MyStruct]] = None my_str_basket: Dict[str, ts[MyStruct]] = None my_enum_basket_list: Dict[MyEnum, ts[List[MyStruct]]] = None @@ -84,9 +82,7 @@ def connect(self, channels: MyGatewayChannels) -> None: self.my_list_data, ) channels.add_send_channel(MyGatewayChannels.my_channel) - channels.set_state(MyGatewayChannels.my_channel, "id") channels.add_send_channel(MyGatewayChannels.my_list_channel) - channels.set_state(MyGatewayChannels.my_list_channel, "id") channels.set_channel(MyGatewayChannels.my_array_channel, csp.const(np.array([1.0, 2.0]))) @@ -766,14 +762,10 @@ def dynamic_state_channels(self) -> Optional[Set[str]]: return {self.scalar_channel_name, self.list_channel_name} def connect(self, channels: MyGatewayChannels) -> None: - channels.set_channel( - self.scalar_channel_name, - csp.const(MyStruct(foo=1.0)), - ) - channels.set_channel( - self.list_channel_name, - csp.const([MyStruct(foo=2.0), MyStruct(foo=3.0)]), - ) + scalar_edge = csp.const(MyStruct(foo=1.0)) + list_edge = csp.const([MyStruct(foo=2.0), MyStruct(foo=3.0)]) + channels.set_channel(self.scalar_channel_name, scalar_edge) + channels.set_channel(self.list_channel_name, list_edge) channels.set_state(self.scalar_channel_name, keyby="id") channels.set_state(self.list_channel_name, keyby="id") if self.connect_channels_assertion: @@ -795,11 +787,11 @@ def connect(self, channels: MyGatewayChannels) -> None: ) csp.add_graph_output( f"s_{self.scalar_channel_name}", - channels.get_channel(f"s_{self.scalar_channel_name}"), + channels.get_state(self.scalar_channel_name), ) csp.add_graph_output( f"s_{self.list_channel_name}", - channels.get_channel(f"s_{self.list_channel_name}"), + channels.get_state(self.list_channel_name), ) @@ -873,6 +865,34 @@ def test_conflicting_dynamic_channels(): csp.build_graph(MyGateway(modules=[setter_1, setter_2], channels=MyGatewayChannels()).graph) +@pytest.mark.parametrize("module_order", ["setter_first", "getter_first"]) +def test_dynamic_channels_state_module_order_independence(module_order): + """Cross-module get_state/set_state must work regardless of connect() order. + + The setter module declares a dynamic channel and registers state on it from + inside connect(). The getter module calls get_state() on that name from + inside its own connect(). Whichever module's connect() runs first, the + getter must end up wired to the real state edge via DelayedEdge binding. + """ + setter = MySetModuleDynamicChannels( + scalar_channel_name="my_dynamic_scalar_channel", + list_channel_name="my_dynamic_list_channel", + ) + getter = MyGetModuleDynamicChannels( + scalar_channel_name="my_dynamic_scalar_channel", + list_channel_name="my_dynamic_list_channel", + ) + modules = [setter, getter] if module_order == "setter_first" else [getter, setter] + gateway = MyGateway(modules=modules, channels=MyGatewayChannels()) + + out = csp.run(gateway.graph, starttime=datetime(2020, 1, 1), endtime=timedelta(1)) + assert len(out["my_dynamic_scalar_channel"]) == 1 + assert out["my_dynamic_scalar_channel"][0][1].foo == 1.0 + assert len(out["my_dynamic_list_channel"]) == 1 + assert len(out["s_my_dynamic_scalar_channel"]) == 1 + assert len(out["s_my_dynamic_list_channel"]) == 2 + + def test_dynamic_channels_same_static_properties(): my_object = {} channel = MyGatewayChannels(my_static_dict_of_objects={"a": my_object}) @@ -990,9 +1010,7 @@ def test_gateway_channels_fields_classmethod(): "my_static_list", "my_static_dict_of_objects", "my_channel", - "s_my_channel", "my_list_channel", - "s_my_list_channel", "my_enum_basket", "my_str_basket", "my_enum_basket_list", diff --git a/csp_gateway/tests/server/modules/io/test_json.py b/csp_gateway/tests/server/modules/io/test_json.py index 96960d4..2a35274 100644 --- a/csp_gateway/tests/server/modules/io/test_json.py +++ b/csp_gateway/tests/server/modules/io/test_json.py @@ -1,5 +1,5 @@ from datetime import datetime, timedelta -from typing import Dict, Type +from typing import Annotated, Dict, Type import csp import orjson @@ -31,9 +31,8 @@ class MyStrStruct(GatewayStruct): class MySmallGatewayChannels(GatewayChannels): - example: ts[int] = None + example: Annotated[ts[int], State(keyby="id")] = None struct_with_str: ts[MyStrStruct] = None - s_example: ts[State[int]] = None my_str_basket: Dict[str, ts[float]] = None diff --git a/csp_gateway/tests/server/modules/io/test_json_converter.py b/csp_gateway/tests/server/modules/io/test_json_converter.py index 8947ec8..93401ae 100644 --- a/csp_gateway/tests/server/modules/io/test_json_converter.py +++ b/csp_gateway/tests/server/modules/io/test_json_converter.py @@ -74,9 +74,7 @@ def connect(self, channels: MyGatewayChannels) -> None: self.my_list_data, ) channels.add_send_channel(MyGatewayChannels.my_channel) - channels.set_state(MyGatewayChannels.my_channel, "id") channels.add_send_channel(MyGatewayChannels.my_list_channel) - channels.set_state(MyGatewayChannels.my_list_channel, "id") if self.by_key: channels.set_channel(MyGatewayChannels.my_enum_basket, self.my_data, MyEnum.ONE) @@ -415,7 +413,6 @@ def test_encode_filter_channels(by_key): by_key=by_key, ) channels = MyGatewayChannels.fields() - assert MyGatewayChannels.s_my_channel in channels assert MyGatewayChannels.my_array_channel in channels json_encoder = MyJsonEncoder(channels_list=channels) gateway = MyGateway(modules=[setter, json_encoder], channels=MyGatewayChannels()) @@ -429,7 +426,6 @@ def test_encode_filter_channels(by_key): assert snapshot_model.my_enum_basket[MyEnum.ONE].foo == 1.0 assert snapshot_model.my_enum_basket[MyEnum.TWO].foo == 2.0 assert snapshot_model.my_enum_basket_list[MyEnum.ONE] == snapshot_model.my_enum_basket_list[MyEnum.TWO] - assert MyGatewayChannels.s_my_channel not in type(snapshot_model).model_fields assert getattr(snapshot_model, _CSP_ENGINE_CYCLE_TIMESTAMP_FIELD) == datetime(2020, 1, 1) diff --git a/csp_gateway/tests/server/modules/logging/test_printing.py b/csp_gateway/tests/server/modules/logging/test_printing.py index ef61862..99fa300 100644 --- a/csp_gateway/tests/server/modules/logging/test_printing.py +++ b/csp_gateway/tests/server/modules/logging/test_printing.py @@ -1,5 +1,5 @@ from datetime import datetime, timedelta -from typing import Dict, List +from typing import Annotated, Dict, List import csp import numpy as np @@ -21,10 +21,8 @@ class MyPrintGatewayChannels(GatewayChannels): my_static: float = 0.0 my_static_dict: Dict[str, float] = {} my_static_list: List[str] = [] - my_channel: ts[int] = None - s_my_channel: ts[State[int]] = None - my_list_channel: ts[[int]] = None - s_my_list_channel: ts[State[int]] = None + my_channel: Annotated[ts[int], State(keyby="id")] = None + my_list_channel: Annotated[ts[[int]], State(keyby="id")] = None my_enum_basket: Dict[MyEnum, ts[int]] = None my_str_basket: Dict[str, ts[int]] = None my_enum_basket_list: Dict[MyEnum, ts[[int]]] = None @@ -48,9 +46,7 @@ def connect(self, channels: MyPrintGatewayChannels) -> None: csp.const(self.my_list_data), ) channels.add_send_channel(MyPrintGatewayChannels.my_channel) - channels.set_state(MyPrintGatewayChannels.my_channel, "id") channels.add_send_channel(MyPrintGatewayChannels.my_list_channel) - channels.set_state(MyPrintGatewayChannels.my_list_channel, "id") channels.set_channel(MyPrintGatewayChannels.my_array_channel, csp.const(np.array([1.0, 2.0]))) diff --git a/csp_gateway/tests/server/shared/test_channel_selection.py b/csp_gateway/tests/server/shared/test_channel_selection.py index eae254e..d06d0a8 100644 --- a/csp_gateway/tests/server/shared/test_channel_selection.py +++ b/csp_gateway/tests/server/shared/test_channel_selection.py @@ -18,8 +18,8 @@ def test_selection(): assert selection.select_from(MyGatewayChannels) == channels assert selection.select_from(MyGatewayChannels, state_channels=True) == [ - "s_my_channel", - "s_my_list_channel", + "my_channel", + "my_list_channel", ] assert selection.select_from(MyGatewayChannels, static_fields=True) == static @@ -30,19 +30,18 @@ def test_selection_duplicates(): selection = ChannelSelection(include=channels + channels + static) assert selection.select_from(MyGatewayChannels) == channels - assert selection.select_from(MyGatewayChannels, state_channels=True) == ["s_my_channel"] + assert selection.select_from(MyGatewayChannels, state_channels=True) == ["my_channel"] assert selection.select_from(MyGatewayChannels, static_fields=True) == static def test_selection_all_fields(): - channels = ["my_enum_basket", "my_channel", "s_my_channel"] + channels = ["my_enum_basket", "my_channel"] static = ["my_static", "my_static_dict"] selection = ChannelSelection(include=channels + static) assert selection.select_from(MyGatewayChannels, all_fields=True) == [ "my_enum_basket", "my_channel", - "s_my_channel", "my_static", "my_static_dict", ] @@ -54,7 +53,7 @@ def test_selection_include(): selection = ChannelSelection(include=channels + static) assert selection.select_from(MyGatewayChannels) == channels - assert selection.select_from(MyGatewayChannels, state_channels=True) == ["s_my_channel"] + assert selection.select_from(MyGatewayChannels, state_channels=True) == ["my_channel"] assert selection.select_from(MyGatewayChannels, static_fields=True) == static assert selection.select_from(MyGatewayChannels, all_fields=True) == [ "my_channel", @@ -75,7 +74,7 @@ def test_selection_exclude(): static = ["my_static_dict", "my_static_list"] assert selection.select_from(MyGatewayChannels) == channels - assert selection.select_from(MyGatewayChannels, state_channels=True) == ["s_my_list_channel"] + assert selection.select_from(MyGatewayChannels, state_channels=True) == ["my_list_channel"] assert selection.select_from(MyGatewayChannels, static_fields=True) == static diff --git a/csp_gateway/tests/server/web/test_stage_routes.py b/csp_gateway/tests/server/web/test_stage_routes.py new file mode 100644 index 0000000..066c413 --- /dev/null +++ b/csp_gateway/tests/server/web/test_stage_routes.py @@ -0,0 +1,162 @@ +from typing import Type + +import csp +from csp import ts +from fastapi.testclient import TestClient + +from csp_gateway import ( + Channels, + Gateway, + GatewayChannels, + GatewayModule, + GatewaySettings, + GatewayStruct, + MountRestRoutes, +) + + +class StageOrder(GatewayStruct): + symbol: str = "" + quantity: int = 0 + price: float = 0.0 + + +class StageChannels(GatewayChannels): + orders: ts[StageOrder] = None + + +class StageGateway(Gateway): + channels_model: Type[Channels] = StageChannels # type: ignore[assignment] + + +class StageModule(GatewayModule): + def connect(self, channels: StageChannels) -> None: + channels.set_channel(StageChannels.orders, csp.null_ts(StageOrder)) + channels.set_stage(StageChannels.orders) + + def shutdown(self) -> None: + pass + + +def test_stage_routes_basic_flow(free_port): + gateway = StageGateway( + modules=[ + StageModule(), + MountRestRoutes(force_mount_all=True), + ], + channels=StageChannels(), + settings=GatewaySettings(PORT=free_port), + ) + + gateway.start(rest=True, _in_test=True) + client = TestClient(gateway.web_app.get_fastapi()) + try: + # List staged channels + response = client.get("/api/v1/stage/") + assert response.status_code == 200 + assert "orders" in response.json() + + # stage_add: create new staging with an item + payload = {"symbol": "AAPL", "quantity": 10, "price": 190.5} + response = client.post("/api/v1/stage/orders", json=payload) + assert response.status_code == 200 + add_result = response.json() + assert len(add_result) == 1 + staging_id = list(add_result.keys())[0] + assert add_result[staging_id][0]["symbol"] == "AAPL" + + # stage_list: ensure staging is present + response = client.get("/api/v1/stage/orders") + assert response.status_code == 200 + assert staging_id in response.json() + + # stage_lookup specific staging + response = client.put(f"/api/v1/stage/orders?id={staging_id}") + assert response.status_code == 200 + lookup_result = response.json() + assert staging_id in lookup_result + assert lookup_result[staging_id][0]["quantity"] == 10 + + # stage_release specific staging + response = client.patch(f"/api/v1/stage/orders?id={staging_id}") + assert response.status_code == 200 + release_result = response.json() + assert staging_id in release_result + assert release_result[staging_id][0]["symbol"] == "AAPL" + + # stage_list after release should no longer include released ID + response = client.get("/api/v1/stage/orders") + assert response.status_code == 200 + assert staging_id not in response.json() + finally: + gateway.stop() + + +def test_stage_routes_full_lifecycle(free_port): + """Test the full stage lifecycle: new, add, remove, lookup, list, release.""" + gateway = StageGateway( + modules=[ + StageModule(), + MountRestRoutes(force_mount_all=True), + ], + channels=StageChannels(), + settings=GatewaySettings(PORT=free_port), + ) + + gateway.start(rest=True, _in_test=True) + client = TestClient(gateway.web_app.get_fastapi()) + try: + # stage_new: POST with no body creates empty staging + response = client.post("/api/v1/stage/orders") + assert response.status_code == 200 + result = response.json() + assert len(result) == 1 + staging_id = list(result.keys())[0] + assert result[staging_id] == [] + + # stage_add: POST with body adds to latest staging + payload = {"symbol": "AAPL", "quantity": 10, "price": 190.5} + response = client.post(f"/api/v1/stage/orders?id={staging_id}", json=payload) + assert response.status_code == 200 + result = response.json() + assert result[staging_id][0]["symbol"] == "AAPL" + + # stage_add: add second item + payload2 = {"symbol": "MSFT", "quantity": 5, "price": 400.0} + response = client.post(f"/api/v1/stage/orders?id={staging_id}", json=payload2) + assert response.status_code == 200 + result = response.json() + assert len(result[staging_id]) == 2 + + # stage_list + response = client.get("/api/v1/stage/orders") + assert response.status_code == 200 + assert staging_id in response.json() + + # stage_lookup + response = client.put(f"/api/v1/stage/orders?id={staging_id}") + assert response.status_code == 200 + lookup = response.json() + assert len(lookup[staging_id]) == 2 + + # stage_remove: DELETE with body removes specific struct + item_to_remove = lookup[staging_id][0] + response = client.request("DELETE", f"/api/v1/stage/orders?id={staging_id}", json=item_to_remove) + assert response.status_code == 200 + result = response.json() + assert len(result[staging_id]) == 1 + assert result[staging_id][0]["symbol"] == "MSFT" + + # stage_release + response = client.patch(f"/api/v1/stage/orders?id={staging_id}") + assert response.status_code == 200 + release_result = response.json() + assert staging_id in release_result + assert release_result[staging_id][0]["symbol"] == "MSFT" + + # After release, staging gone + response = client.get("/api/v1/stage/orders") + assert response.status_code == 200 + assert staging_id not in response.json() + finally: + gateway.stop() diff --git a/csp_gateway/tests/server/web/test_webserver.py b/csp_gateway/tests/server/web/test_webserver.py index 3761707..eb31628 100644 --- a/csp_gateway/tests/server/web/test_webserver.py +++ b/csp_gateway/tests/server/web/test_webserver.py @@ -127,12 +127,12 @@ def test_openapi_includes_fully_qualified_type_name(self, rest_client: TestClien f"Expected fully qualified type name in /api/v1/lookup/example/{{id}} type_, got: {lookup_example_type}" ) - # Test /api/v1/state/example route - state_example_path = paths.get("/api/v1/state/example", {}) + # Test /api/v1/state/example_with_state route + state_example_path = paths.get("/api/v1/state/example_with_state", {}) state_example_get = state_example_path.get("get", {}) state_example_type = state_example_get.get("type_", "") assert state_example_type == "csp_gateway.server.demo.omnibus.ExampleData", ( - f"Expected fully qualified type name in /api/v1/state/example type_, got: {state_example_type}" + f"Expected fully qualified type name in /api/v1/state/example_with_state type_, got: {state_example_type}" ) # Test /api/v1/controls/heartbeat route @@ -359,7 +359,7 @@ def test_csp_state(self, rest_client: TestClient): last_data = self._wait_for_data(rest_client=rest_client) # Get state data - response = rest_client.get("/api/v1/state/example?token=test") + response = rest_client.get("/api/v1/state/example_with_state?token=test") assert response.status_code == 200 state_data = response.json() @@ -371,7 +371,7 @@ def test_csp_state_query(self, rest_client: TestClient): self._wait_for_data(rest_client=rest_client) # Get state data - response = rest_client.get('/api/v1/state/example?token=test&query={"filters":[{"attr":"x","by":{"value":1,"where":"=="}}]}') + response = rest_client.get('/api/v1/state/example_with_state?token=test&query={"filters":[{"attr":"x","by":{"value":1,"where":"=="}}]}') assert response.status_code == 200 state_data = response.json() @@ -508,6 +508,9 @@ def test_csp_toplevel_last(self, rest_client: TestClient): "controls", "example", "example_list", + "example_with_stage", + "example_with_state", + "example_with_state_multiple", "never_ticks", "str_basket", ] @@ -522,6 +525,9 @@ def test_csp_toplevel_next(self, rest_client: TestClient): "controls", "example", "example_list", + "example_with_stage", + "example_with_state", + "example_with_state_multiple", "never_ticks", "str_basket", ] @@ -530,7 +536,12 @@ def test_csp_toplevel_state(self, rest_client: TestClient): self._wait_for_data(rest_client=rest_client) response_state = rest_client.get("/api/v1/state?token=test") assert response_state.status_code == 200 - assert sorted(response_state.json()) == ["example"] + assert sorted(response_state.json()) == [ + "example", + "example_with_state", + "example_with_state_alternative", + "example_with_state_multiple", + ] def test_csp_toplevel_send(self, rest_client: TestClient): self._wait_for_data(rest_client=rest_client) @@ -557,7 +568,12 @@ def test_gateway_client(self, mock_get, mock_post, rest_client: TestClient, gate assert sorted(list(gateway_client.openapi_spec.keys())) == ["components", "info", "openapi", "paths"] assert "/api/v1/last/example" in gateway_client.openapi_spec["paths"].keys() assert "/api/v1/last/example_list" in gateway_client.openapi_spec["paths"].keys() - assert response_state == ["example"] + assert sorted(response_state) == [ + "example", + "example_with_state", + "example_with_state_alternative", + "example_with_state_multiple", + ] for route in ["example", "example_list"]: data = gateway_client.last(route) @@ -602,7 +618,7 @@ def test_gateway_client(self, mock_get, mock_post, rest_client: TestClient, gate send_data["y"], ) - return_data = gateway_client.state("example", query=Query(filters=[Filter(attr="x", by={"value": 1, "where": "=="})])) + return_data = gateway_client.state("example_with_state", query=Query(filters=[Filter(attr="x", by={"value": 1, "where": "=="})])) assert isinstance(return_data, list) return_datum = return_data[0] assert "id" in return_datum @@ -798,6 +814,9 @@ def test_perspective_tables(self, rest_client: TestClient): "controls", "example", "example_list", + "example_with_stage", + "example_with_state", + "example_with_state_multiple", "my_custom_table", # "never_ticks", # NOTE: never_ticks has no data, so table is omitted "str_basket", @@ -831,6 +850,9 @@ def test_stream(self, rest_client: TestClient): "controls", "example", "example_list", + "example_with_stage", + "example_with_state", + "example_with_state_multiple", "heartbeat", "never_ticks", "str_basket/a", diff --git a/docs/wiki/API.md b/docs/wiki/API.md index 00adbcc..a57508d 100644 --- a/docs/wiki/API.md +++ b/docs/wiki/API.md @@ -17,6 +17,7 @@ As described in [Overview#Channels](Overview#Channels), the `csp-gateway` REST A - **last** (`GET`, `/api/v1/last/`): Get the last tick of data on a channel - **next** (`GET`, `/api/v1/next/`): Wait for the next tick of data on a channel: **WARNING**: blocks, and can often be misused into race conditions - **state** (`GET`, `/api/v1/state/`): Get the accumulated state for any channel +- **stage** (`POST`/`DELETE`/`PATCH`/`GET`/`PUT`, `/api/v1/stage/`): Manage staged structs on a channel before releasing them into the graph - **send** (`POST`, `/api/v1/send/`): Send a new datum as a tick into the running csp graph - **lookup** (`POST`, `/api/v1/lookup//`): Lookup an individual GatewayStruct by its required `id` field @@ -32,10 +33,39 @@ As described in [Overview#Channels](Overview#Channels), the `csp-gateway` REST A ## State -A `GatewayModule` can call `set_state` in its `connect` method to allow for this API to be available. +State in `csp-gateway` exposes the accumulated history of a channel via the REST API as `/api/v1/state/`. State is collected by one or more instance attributes into an in-memory [DuckDB](https://duckdb.org/) instance. -For example, suppose I have the following type: +There are two ways to declare state on a channel: + +1. **Annotation** on the channel definition, using `State(keyby, alias=...)`: + + ```python + from typing import Annotated + from csp_gateway import State, GatewayChannels, ts + + class MyChannels(GatewayChannels): + # implicit alias "example_with_state" (matches the field name) + example_with_state: Annotated[ts[ExampleData], State(("id", "x"))] = None + + # multiple state views on a single channel, each addressable by alias + example_multi: Annotated[ + ts[ExampleData], + State(("id", "x")), # alias = "example_multi" + State(("id", "y"), alias="example_multi_alt"), + ] = None + ``` + +1. **Imperative** call from a `GatewayModule`'s `connect` method: + + ```python + def connect(self, channels: MyChannels): + edge = ... # produce a csp.ts edge + channels.set_channel(MyChannels.example, edge) + channels.set_state(edge, "example_from_connect", ("id", "x", "z")) + ``` + +For example, given: ```python class ExampleData(GatewayStruct): @@ -44,12 +74,7 @@ class ExampleData(GatewayStruct): z: str ``` -If my `GatewayModule` called `set_state("example", ("x",))`, state would be collected as the last tick of `ExampleData` per each unique value of `x`. If called with `set_state("example", ("x", "y"))`, it would be collected as the last tick per each unique pair `x,y`, etc. - -> [!IMPORTANT] -> -> This code and API will likely change a bit as we allow for more granular collection of records, -> and expose more DuckDB functionality. +Declaring `State(("x",))` collects the last tick of `ExampleData` per each unique value of `x`. Declaring `State(("x", "y"))` collects the last tick per each unique pair `(x, y)`, etc. ## Query @@ -59,25 +84,99 @@ Here are some examples from the autodocumentation illustrating the use of filter ```raw # Filter only records where `record.x` == 5 -api/v1/state/example?query={"filters":[{"attr":"x","by":{"value":5,"where":"=="}}]} +api/v1/state/example_with_state?query={"filters":[{"attr":"x","by":{"value":5,"where":"=="}}]} # Filter only records where `record.x` < 10 -/api/v1/state/example?query={"filters":[{"attr":"x","by":{"value":10,"where":"<"}}]} +/api/v1/state/example_with_state?query={"filters":[{"attr":"x","by":{"value":10,"where":"<"}}]} # Filter only records where `record.timestamp` < "2023-03-30T14:45:26.394000" -/api/v1/state/example?query={"filters":[{"attr":"timestamp","by":{"when":"2023-03-30T14:45:26.394000","where":"<"}}]} +/api/v1/state/example_with_state?query={"filters":[{"attr":"timestamp","by":{"when":"2023-03-30T14:45:26.394000","where":"<"}}]} # Filter only records where `record.id` < `record.y` -/api/v1/state/example?query={"filters":[{"attr":"id","by":{"attr":"y","where":"<"}}]} +/api/v1/state/example_with_state?query={"filters":[{"attr":"id","by":{"attr":"y","where":"<"}}]} # Filter only records where `record.x` < 50 and `record.x` >= 30 -/api/v1/state/example?query={"filters":[{"attr":"x","by":{"value":50,"where":"<"}},{"attr":"x","by":{"value":30,"where":">="}}]} +/api/v1/state/example_with_state?query={"filters":[{"attr":"x","by":{"value":50,"where":"<"}},{"attr":"x","by":{"value":30,"where":">="}}]} ``` > [!IMPORTANT] > > This code and API will likely change a bit as we expose more DuckDB functionality. +## Staging + +Staging in `csp-gateway` allows you to accumulate structs into named "staging areas" on a channel before atomically releasing them into the graph. This is useful for batch preparation workflows where you want to assemble a group of items and then release them all at once. + +Staging is exposed via the REST API under `/api/v1/stage/`. + +There are two ways to enable staging on a channel: + +1. **Annotation** on the channel definition, using `Stage()`: + + ```python + from typing import Annotated + from csp_gateway import Stage, GatewayChannels, ts + + class MyChannels(GatewayChannels): + orders: Annotated[ts[OrderData], Stage()] = None + ``` + +1. **Imperative** call from a `GatewayModule`'s `connect` method: + + ```python + def connect(self, channels: MyChannels): + channels.set_channel(MyChannels.orders, edge) + channels.set_stage(MyChannels.orders) + ``` + +### REST Endpoints + +| Method | Endpoint | Description | +| -------- | ------------------------- | ------------------------------------------------- | +| `GET` | `/api/v1/stage/` | List channels that have staging enabled | +| `POST` | `/api/v1/stage/` | Add a struct to staging (or create empty staging) | +| `DELETE` | `/api/v1/stage/` | Remove struct(s) from staging | +| `PATCH` | `/api/v1/stage/` | Release staged structs into the channel | +| `GET` | `/api/v1/stage/` | List active staging IDs | +| `PUT` | `/api/v1/stage/` | Look up contents of staging area(s) | + +All endpoints accept an optional `id` query parameter to target specific staging areas: + +- **No `id` param**: operate on the latest staging area (or all, depending on the method) +- **`id=`** (empty string): operate on all staging areas +- **`id=abc,def`**: operate on specific staging area(s) by ID + +### Lifecycle + +1. **Create** a new staging area: `POST /api/v1/stage/` with no body +1. **Add** structs: `POST /api/v1/stage/?id=` with a JSON body +1. **Review** contents: `PUT /api/v1/stage/?id=` +1. **Release** into the graph: `PATCH /api/v1/stage/?id=` + +Released structs are pushed into the CSP graph as individual ticks on the channel. Released staging areas remain accessible via `PUT` (lookup) for historical reference, but no longer appear in `GET` (list). + +### Examples + +```bash +# List channels with staging enabled +curl http://localhost:8000/api/v1/stage/ + +# Create an empty staging area +curl -X POST http://localhost:8000/api/v1/stage/orders +# Returns: {"": []} + +# Add a struct to the latest staging +curl -X POST http://localhost:8000/api/v1/stage/orders \ + -H "Content-Type: application/json" \ + -d '{"symbol": "AAPL", "quantity": 10, "price": 190.5}' + +# Look up contents of all staging areas (including released) +curl -X PUT http://localhost:8000/api/v1/stage/orders + +# Release a specific staging into the graph +curl -X PATCH "http://localhost:8000/api/v1/stage/orders?id=" +``` + ## Websockets In addition to a REST API, a more `csp`-like streaming API is also available via Websockets when the `MountWebSocketRoutes` module is included in a `Gateway` instance. diff --git a/docs/wiki/Client.md b/docs/wiki/Client.md index 582a431..a8c2fea 100644 --- a/docs/wiki/Client.md +++ b/docs/wiki/Client.md @@ -29,6 +29,14 @@ A client as a small number of general-purpose methods. In alphabetical order: - `lookup`: lookup a piece of data by `id` - `next`: get the next ticked data value on a channel - `send`: send some data onto a channel +- `stage`: raw staging API dispatcher (add/remove/release/list/lookup) +- `stage_add`: add a struct to staging area(s) +- `stage_new`: create a new empty staging area +- `stage_remove`: remove struct(s) from staging area(s) +- `stage_clear`: clear all structs from staging area(s) without releasing +- `stage_release`: release staged structs into the channel +- `stage_list`: list active staging areas with their contents +- `stage_lookup`: look up contents of all staging areas (including released) - `state`: get the value of a given channel's state accumulator Additionally, a client has some streaming methods available when websockets are configured: @@ -282,11 +290,11 @@ client.state().as_json() ``` ``` -['example'] +['example', 'example_with_state', 'example_with_state_alternative', 'example_with_state_multiple'] ``` ```python -client.state("example").as_pandas_df().tail() +client.state("example_with_state").as_pandas_df().tail() # We note that there are a large number of columns in the above dataframe. # This is because `mapping` is a dict with different keys for eery row. @@ -503,7 +511,7 @@ client.send( } ) -client.state("example").as_pandas_df().tail() +client.state("example_with_state").as_pandas_df().tail() ```
@@ -723,6 +731,56 @@ File /isilon/home/nk12433/bitbucket/csp-gateway/csp_gateway/client/client.py:293 ServerUnprocessableException: [{'loc': ['body'], 'msg': 'value is not a valid list', 'type': 'type_error.list'}, {'loc': ['body', 'x'], 'msg': 'value must be non-negative.', 'type': 'value_error'}] ``` +## Staging + +Channels with staging enabled allow you to accumulate structs into staging areas before releasing them into the graph. The client provides several convenience methods for this workflow. + +```python +# Create a new empty staging area +result = client.stage_new("example_with_stage") +# Returns: {"": []} +staging_id = list(result.keys())[0] +``` + +```python +# Add a struct to a specific staging area +client.stage_add("example_with_stage", data={"x": 42, "y": "hello"}, staging_ids=[staging_id]) +# Returns: {"": [{"x": 42, "y": "hello", "id": "...", ...}]} +``` + +```python +# List active staging areas with their contents +client.stage_list("example_with_stage") +# Returns: {"": []} +``` + +```python +# Look up contents of all staging areas (including previously released) +client.stage_lookup("example_with_stage") +# Returns: {"": [], "": [], ...} +``` + +```python +# Release staged structs into the channel (pushes them as ticks) +client.stage_release("example_with_stage", staging_ids=[staging_id]) +# Returns: {"": []} +``` + +```python +# Remove specific structs or clear entire staging areas +client.stage_remove("example_with_stage", data=item_dict, staging_ids=[staging_id]) +client.stage_clear("example_with_stage", staging_ids=[staging_id]) +``` + +The raw `stage()` method is also available for direct control: + +```python +client.stage("example_with_stage", method="add", data={"x": 1}, staging_ids=None) +client.stage("example_with_stage", method="list") +client.stage("example_with_stage", method="lookup") +client.stage("example_with_stage", method="release", staging_ids=[staging_id]) +``` + ## Next The running `GatewayServer` is a synchronous system, and we're interacting it via asynchronous `REST` requests. However, we can still perform actions like "wait for the next tick". This can be dangerous and lead to race conditions, but it can still be useful in certain circumstances. diff --git a/docs/wiki/Modules.md b/docs/wiki/Modules.md index 2d56cb1..3a88e97 100644 --- a/docs/wiki/Modules.md +++ b/docs/wiki/Modules.md @@ -946,6 +946,7 @@ modules: - **mount_last** (`ChannelSelection`): channels to include in last routes - **mount_next** (`ChannelSelection`): channels to include in next routes - **mount_send** (`ChannelSelection`): channels to include in send routes +- **mount_stage** (`ChannelSelection`): channels to include in stage routes (automatically includes channels with `Stage()` annotation or `set_stage` call) - **mount_state** (`ChannelSelection`): channels to include in state routes - **mount_lookup** (`ChannelSelection`): channels to include in lookup routes diff --git a/docs/wiki/Overview.md b/docs/wiki/Overview.md index 9485dc2..c680d27 100644 --- a/docs/wiki/Overview.md +++ b/docs/wiki/Overview.md @@ -165,8 +165,29 @@ Recall from above: `GatewayChannels` also have the ability to create and maintain state, which is particularly useful in the REST API. State is managed by an in-memory [DuckDB](https://duckdb.org/) instance. -- `set_state(self, field: str, keyby: Union[str, Tuple[str, ...]], indexer: Union[str, int] = None) -> None`: Collect state of an edge by a certain attribute indexer -- `get_state(self, field: str, indexer: Union[str, int] = None) -> Any`: Get the current state collected on an edge as a ticking object +State can be declared two ways: + +1. **Annotation** (on the channel definition): + + ```python + from typing import Annotated + from csp_gateway import State, GatewayChannels + + class MyChannels(GatewayChannels): + example_with_state: Annotated[ts[ExampleData], State(("id", "x"))] = None + # multiple State() annotations are allowed; each may set an `alias` + example_multi: Annotated[ + ts[ExampleData], + State(("id", "x")), + State(("id", "y"), alias="example_multi_alt"), + ] = None + ``` + +1. **Imperative**, from a module's `connect` method via `set_state`: + + - `set_state(self, edge: Edge, field: str, keyby: Union[str, Tuple[str, ...]], indexer: Union[str, int] = None) -> None`: Collect state of `edge` by a certain attribute, exposed under `field` + +- `get_state(self, field: str, indexer: Union[str, int] = None) -> Any`: Get the current state collected for a state `field` as a ticking object `GatewayChannels` can tell the REST API to allow a certain channel to be sent in via POST requests: @@ -177,7 +198,7 @@ This is the primary API for both the [REST API](API) and the [integrated Python - `last(self, field: str, indexer: Union[str, int] = None, *, timeout=None) -> None`: Get the last ticked value on a channel - `next(self, field: str, indexer: Union[str, int] = None, *, timeout=None) -> None`: Get (wait for) the next ticked value on a channel -- `state(self, field: str, indexer: Union[str, int] = None, *, timeout=None) -> Any`: Get the state collected on a channel +- `state(self, field: str, indexer: Union[str, int] = None, *, timeout=None) -> Any`: Get the state collected for a state `field` - `query(self, field: str, indexer: Union[str, int] = None, query: "Query" = None) -> Any`: Query the state on a channel (see [API](API) for more about `Query`) - `send(self, field: str, value: Any, indexer: Union[str, int] = None) -> None`: Send new data to a channel, from outside the `csp` graph