Coverage for /Users/ajo/work/jumpstarter/jumpstarter/packages/jumpstarter/jumpstarter/common/streams.py: 75%
28 statements
« prev ^ index » next coverage.py v7.8.0, created at 2025-05-06 10:20 +0200
« prev ^ index » next coverage.py v7.8.0, created at 2025-05-06 10:20 +0200
1from contextlib import asynccontextmanager
2from typing import Annotated, Literal, Union
3from uuid import UUID
5import grpc
6from jumpstarter_protocol.jumpstarter.v1 import router_pb2_grpc
7from pydantic import BaseModel, Field, Json
9from jumpstarter.common.grpc import aio_secure_channel, ssl_channel_credentials
10from jumpstarter.streams.common import forward_stream
11from jumpstarter.streams.router import RouterStream
14class ResourceStreamRequest(BaseModel):
15 kind: Literal["resource"] = "resource"
16 uuid: UUID
19class DriverStreamRequest(BaseModel):
20 kind: Literal["driver"] = "driver"
21 uuid: UUID
22 method: str
25StreamRequest = Annotated[
26 Union[ResourceStreamRequest, DriverStreamRequest],
27 Field(discriminator="kind"),
28]
31class StreamRequestMetadata(BaseModel):
32 request: Json[StreamRequest]
35@asynccontextmanager
36async def connect_router_stream(endpoint, token, stream, tls_config, grpc_options):
37 credentials = grpc.composite_channel_credentials(
38 await ssl_channel_credentials(endpoint, tls_config),
39 grpc.access_token_call_credentials(token),
40 )
42 async with aio_secure_channel(endpoint, credentials, grpc_options) as channel:
43 router = router_pb2_grpc.RouterServiceStub(channel)
44 context = router.Stream(metadata=())
45 async with RouterStream(context=context) as s:
46 async with forward_stream(s, stream):
47 yield