Coverage for /Users/ajo/work/jumpstarter/jumpstarter/packages/jumpstarter/jumpstarter/client/client.py: 94%
33 statements
« prev ^ index » next coverage.py v7.8.0, created at 2025-05-06 10:21 +0200
« prev ^ index » next coverage.py v7.8.0, created at 2025-05-06 10:21 +0200
1from collections import OrderedDict, defaultdict
2from contextlib import ExitStack, asynccontextmanager
3from graphlib import TopologicalSorter
4from uuid import UUID
6import grpc
7from anyio.from_thread import BlockingPortal
8from google.protobuf import empty_pb2
9from jumpstarter_protocol import jumpstarter_pb2_grpc
11from jumpstarter.client import DriverClient
12from jumpstarter.common.importlib import import_class
15@asynccontextmanager
16async def client_from_path(path: str, portal: BlockingPortal, stack: ExitStack, allow: list[str], unsafe: bool):
17 async with grpc.aio.secure_channel(
18 f"unix://{path}", grpc.local_channel_credentials(grpc.LocalConnectionType.UDS)
19 ) as channel:
20 yield await client_from_channel(channel, portal, stack, allow, unsafe)
23async def client_from_channel(
24 channel: grpc.aio.Channel,
25 portal: BlockingPortal,
26 stack: ExitStack,
27 allow: list[str],
28 unsafe: bool,
29) -> DriverClient:
30 topo = defaultdict(list)
31 last_seen = {}
32 reports = {}
33 clients = OrderedDict()
35 response = await jumpstarter_pb2_grpc.ExporterServiceStub(channel).GetReport(empty_pb2.Empty())
37 for index, report in enumerate(response.reports):
38 topo[index] = []
40 last_seen[report.uuid] = index
42 if report.parent_uuid != "":
43 parent_index = last_seen[report.parent_uuid]
44 topo[parent_index].append(index)
46 reports[index] = report
48 for index in TopologicalSorter(topo).static_order():
49 report = reports[index]
51 client_class = import_class(report.labels["jumpstarter.dev/client"], allow, unsafe)
52 client = client_class(
53 uuid=UUID(report.uuid),
54 labels=report.labels,
55 channel=channel,
56 portal=portal,
57 stack=stack.enter_context(ExitStack()),
58 children={reports[k].labels["jumpstarter.dev/name"]: clients[k] for k in topo[index]},
59 )
61 clients[index] = client
63 return clients.popitem(last=True)[1]