1717import typing
1818from contextlib import asynccontextmanager
1919
20- from .client_types import RestateClient , RestateClientSendHandle , HttpError
20+ from .client_types import RestateClient , RestateClientSendHandle , RestateScopedClient , HttpError
2121
2222from .context import HandlerType
2323from .serde import BytesSerde , JsonSerde , Serde
@@ -36,6 +36,9 @@ def __init__(self, client: httpx.AsyncClient, headers: typing.Optional[dict] = N
3636 self .headers = headers or {}
3737 self .client = client
3838
39+ def scope (self , scope : str ) -> RestateScopedClient :
40+ return ScopedClient (self , scope )
41+
3942 async def do_call (
4043 self ,
4144 tpe : HandlerType [I , O ],
@@ -46,6 +49,8 @@ async def do_call(
4649 idempotency_key : str | None = None ,
4750 headers : typing .Dict [str , str ] | None = None ,
4851 force_json_output : bool = False ,
52+ scope : str | None = None ,
53+ limit_key : str | None = None ,
4954 ) -> O :
5055 """Make an RPC call to the given handler"""
5156 target_handler = handler_from_callable (tpe )
@@ -77,6 +82,8 @@ async def do_call(
7782 send = send ,
7883 idempotency_key = idempotency_key ,
7984 headers = headers ,
85+ scope = scope ,
86+ limit_key = limit_key ,
8087 )
8188
8289 async def do_raw_call (
@@ -91,6 +98,8 @@ async def do_raw_call(
9198 send : bool = False ,
9299 idempotency_key : str | None = None ,
93100 headers : typing .Dict [str , str ] | None = None ,
101+ scope : str | None = None ,
102+ limit_key : str | None = None ,
94103 ) -> O :
95104 """Make an RPC call to the given handler"""
96105 parameter = input_serde .serialize (input_param )
@@ -112,6 +121,8 @@ async def do_raw_call(
112121 key = key ,
113122 delay = ms ,
114123 idempotency_key = idempotency_key ,
124+ scope = scope ,
125+ limit_key = limit_key ,
115126 )
116127 return output_serde .deserialize (res ) # type: ignore
117128
@@ -126,21 +137,37 @@ async def post(
126137 key : str | None = None ,
127138 delay : int | None = None ,
128139 idempotency_key : str | None = None ,
140+ scope : str | None = None ,
141+ limit_key : str | None = None ,
129142 ) -> bytes :
130143 """
131144 Send a POST request to the Restate service.
132145 """
133- endpoint = service
134- if key :
135- endpoint += f"/{ key } "
136- endpoint += f"/{ handler } "
137- if send :
138- endpoint += "/send"
139- if delay is not None :
146+ if scope is not None :
147+ # Scoped invocations use the dedicated ingress path:
148+ # restate/scope/{scope}/call/{service}[/{key}]/{handler}
149+ # restate/scope/{scope}/send/{service}[/{key}]/{handler}
150+ verb = "send" if send else "call"
151+ endpoint = f"restate/scope/{ scope } /{ verb } /{ service } "
152+ if key :
153+ endpoint += f"/{ key } "
154+ endpoint += f"/{ handler } "
155+ if send and delay is not None :
140156 endpoint = endpoint + f"?delay={ delay } "
157+ else :
158+ endpoint = service
159+ if key :
160+ endpoint += f"/{ key } "
161+ endpoint += f"/{ handler } "
162+ if send :
163+ endpoint += "/send"
164+ if delay is not None :
165+ endpoint = endpoint + f"?delay={ delay } "
141166 dict_headers = dict (headers ) if headers is not None else {}
142167 if idempotency_key is not None :
143168 dict_headers ["Idempotency-Key" ] = idempotency_key
169+ if limit_key is not None :
170+ dict_headers ["x-restate-limit-key" ] = limit_key
144171 res = await self .client .post (endpoint , headers = dict_headers , content = content )
145172 if res .status_code >= 400 :
146173 raise HttpError (res .status_code , res .reason_phrase , res .text )
@@ -250,6 +277,8 @@ async def generic_call(
250277 key : str | None = None ,
251278 idempotency_key : str | None = None ,
252279 headers : typing .Dict [str , str ] | None = None ,
280+ scope : str | None = None ,
281+ limit_key : str | None = None ,
253282 ) -> bytes :
254283 serde = BytesSerde ()
255284 call_handle = await self .do_raw_call (
@@ -261,6 +290,8 @@ async def generic_call(
261290 key = key ,
262291 idempotency_key = idempotency_key ,
263292 headers = headers ,
293+ scope = scope ,
294+ limit_key = limit_key ,
264295 )
265296 return call_handle
266297
@@ -273,6 +304,8 @@ async def generic_send(
273304 send_delay : timedelta | None = None ,
274305 idempotency_key : str | None = None ,
275306 headers : typing .Dict [str , str ] | None = None ,
307+ scope : str | None = None ,
308+ limit_key : str | None = None ,
276309 ) -> RestateClientSendHandle :
277310 serde = BytesSerde ()
278311 output_serde : Serde [dict ] = JsonSerde ()
@@ -288,11 +321,110 @@ async def generic_send(
288321 send = True ,
289322 idempotency_key = idempotency_key ,
290323 headers = headers ,
324+ scope = scope ,
325+ limit_key = limit_key ,
291326 )
292327
293328 return RestateClientSendHandle (send_handle_json .get ("invocationId" , "" ), 200 ) # TODO: verify
294329
295330
331+ class ScopedClient (RestateScopedClient ):
332+ """
333+ A scoped client returned by ``client.scope(scope_key)``.
334+
335+ Re-dispatches to the underlying :class:`Client` with the captured scope and a
336+ per-call ``limit_key``.
337+ """
338+
339+ def __init__ (self , client : Client , scope_key : str ):
340+ self .client = client
341+ self .scope_key = scope_key
342+
343+ async def service_call (
344+ self ,
345+ tpe : HandlerType [I , O ],
346+ arg : I ,
347+ limit_key : str | None = None ,
348+ idempotency_key : str | None = None ,
349+ headers : typing .Dict [str , str ] | None = None ,
350+ ) -> O :
351+ return await self .client .do_call (
352+ tpe ,
353+ arg ,
354+ idempotency_key = idempotency_key ,
355+ headers = headers ,
356+ scope = self .scope_key ,
357+ limit_key = limit_key ,
358+ )
359+
360+ async def service_send (
361+ self ,
362+ tpe : HandlerType [I , O ],
363+ arg : I ,
364+ send_delay : typing .Optional [timedelta ] = None ,
365+ limit_key : str | None = None ,
366+ idempotency_key : str | None = None ,
367+ headers : typing .Dict [str , str ] | None = None ,
368+ ) -> RestateClientSendHandle :
369+ send_handle = await self .client .do_call (
370+ tpe ,
371+ parameter = arg ,
372+ send = True ,
373+ send_delay = send_delay ,
374+ idempotency_key = idempotency_key ,
375+ headers = headers ,
376+ force_json_output = True ,
377+ scope = self .scope_key ,
378+ limit_key = limit_key ,
379+ )
380+ send = typing .cast (typing .Dict [str , str ], send_handle )
381+ return RestateClientSendHandle (send .get ("invocationId" , "" ), 200 )
382+
383+ async def workflow_call (
384+ self ,
385+ tpe : HandlerType [I , O ],
386+ key : str ,
387+ arg : I ,
388+ limit_key : str | None = None ,
389+ idempotency_key : str | None = None ,
390+ headers : typing .Dict [str , str ] | None = None ,
391+ ) -> O :
392+ return await self .client .do_call (
393+ tpe ,
394+ arg ,
395+ key ,
396+ idempotency_key = idempotency_key ,
397+ headers = headers ,
398+ scope = self .scope_key ,
399+ limit_key = limit_key ,
400+ )
401+
402+ async def workflow_send (
403+ self ,
404+ tpe : HandlerType [I , O ],
405+ key : str ,
406+ arg : I ,
407+ send_delay : typing .Optional [timedelta ] = None ,
408+ limit_key : str | None = None ,
409+ idempotency_key : str | None = None ,
410+ headers : typing .Dict [str , str ] | None = None ,
411+ ) -> RestateClientSendHandle :
412+ send_handle = await self .client .do_call (
413+ tpe ,
414+ parameter = arg ,
415+ key = key ,
416+ send = True ,
417+ send_delay = send_delay ,
418+ idempotency_key = idempotency_key ,
419+ headers = headers ,
420+ force_json_output = True ,
421+ scope = self .scope_key ,
422+ limit_key = limit_key ,
423+ )
424+ send = typing .cast (typing .Dict [str , str ], send_handle )
425+ return RestateClientSendHandle (send .get ("invocationId" , "" ), 200 )
426+
427+
296428@asynccontextmanager
297429async def create_client (
298430 ingress : str , headers : typing .Optional [dict ] = None
0 commit comments