Coverage for /Users/ajo/work/jumpstarter/jumpstarter/packages/jumpstarter/jumpstarter/client/grpc.py: 52%
120 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 __future__ import annotations
3from dataclasses import dataclass, field
4from datetime import datetime, timedelta
6import yaml
7from google.protobuf import duration_pb2, field_mask_pb2, json_format
8from grpc.aio import Channel
9from jumpstarter_protocol import client_pb2, client_pb2_grpc, kubernetes_pb2
10from pydantic import BaseModel, ConfigDict, Field, field_serializer
12from jumpstarter.common.grpc import translate_grpc_exceptions
15def parse_identifier(identifier: str, kind: str) -> (str, str):
16 segments = identifier.split("/")
17 if len(segments) != 4:
18 raise ValueError("incorrect number of segments in identifier, expecting 4, got {}".format(len(segments)))
19 if segments[0] != "namespaces":
20 raise ValueError("incorrect first segment in identifier, expecting namespaces, got {}".format(segments[0]))
21 if segments[2] != kind:
22 raise ValueError("incorrect third segment in identifier, expecting {}, got {}".format(kind, segments[2]))
23 return segments[1], segments[3]
26def parse_client_identifier(identifier: str) -> (str, str):
27 return parse_identifier(identifier, "clients")
30def parse_exporter_identifier(identifier: str) -> (str, str):
31 return parse_identifier(identifier, "exporters")
34def parse_lease_identifier(identifier: str) -> (str, str):
35 return parse_identifier(identifier, "leases")
38class Exporter(BaseModel):
39 namespace: str
40 name: str
41 labels: dict[str, str]
43 @classmethod
44 def from_protobuf(cls, data: client_pb2.Exporter) -> Exporter:
45 namespace, name = parse_exporter_identifier(data.name)
46 return cls(namespace=namespace, name=name, labels=data.labels)
49class Lease(BaseModel):
50 namespace: str
51 name: str
52 selector: str
53 duration: timedelta
54 client: str
55 exporter: str
56 conditions: list[kubernetes_pb2.Condition]
57 effective_begin_time: datetime | None = None
59 model_config = ConfigDict(
60 arbitrary_types_allowed=True,
61 ser_json_timedelta="float",
62 )
64 @field_serializer("conditions")
65 def serialize_conditions(self, conditions: list[kubernetes_pb2.Condition], _info):
66 return [json_format.MessageToDict(condition) for condition in conditions]
68 @classmethod
69 def from_protobuf(cls, data: client_pb2.Lease) -> Lease:
70 namespace, name = parse_lease_identifier(data.name)
72 _, client = parse_client_identifier(data.client)
73 if data.exporter != "":
74 _, exporter = parse_exporter_identifier(data.exporter)
75 else:
76 exporter = ""
78 effective_begin_time = None
79 if data.effective_begin_time:
80 effective_begin_time = data.effective_begin_time.ToDatetime(
81 tzinfo=datetime.now().astimezone().tzinfo,
82 )
84 return cls(
85 namespace=namespace,
86 name=name,
87 selector=data.selector,
88 duration=data.duration.ToTimedelta(),
89 client=client,
90 exporter=exporter,
91 effective_begin_time=effective_begin_time,
92 conditions=data.conditions,
93 )
95 def dump_json(self):
96 return self.model_dump_json(indent=4, by_alias=True)
98 def dump_yaml(self):
99 return yaml.safe_dump(self.model_dump(mode="json", by_alias=True), indent=2)
102class ExporterList(BaseModel):
103 exporters: list[Exporter]
104 next_page_token: str | None = Field(exclude=True)
106 @classmethod
107 def from_protobuf(cls, data: client_pb2.ListExportersResponse) -> ExporterList:
108 return cls(
109 exporters=list(map(Exporter.from_protobuf, data.exporters)),
110 next_page_token=data.next_page_token,
111 )
113 def dump_json(self):
114 return self.model_dump_json(indent=4, by_alias=True)
116 def dump_yaml(self):
117 return yaml.safe_dump(self.model_dump(mode="json", by_alias=True), indent=2)
120class LeaseList(BaseModel):
121 leases: list[Lease]
122 next_page_token: str | None = Field(exclude=True)
124 @classmethod
125 def from_protobuf(cls, data: client_pb2.ListLeasesResponse) -> LeaseList:
126 return cls(
127 leases=list(map(Lease.from_protobuf, data.leases)),
128 next_page_token=data.next_page_token,
129 )
131 def dump_json(self):
132 return self.model_dump_json(indent=4, by_alias=True)
134 def dump_yaml(self):
135 return yaml.safe_dump(self.model_dump(mode="json", by_alias=True), indent=2)
138@dataclass(kw_only=True, slots=True)
139class ClientService:
140 channel: Channel
141 namespace: str
142 stub: client_pb2_grpc.ClientServiceStub = field(init=False)
144 def __post_init__(self):
145 self.stub = client_pb2_grpc.ClientServiceStub(channel=self.channel)
147 async def GetExporter(self, *, name: str):
148 with translate_grpc_exceptions():
149 exporter = await self.stub.GetExporter(
150 client_pb2.GetExporterRequest(
151 name="namespaces/{}/exporters/{}".format(self.namespace, name),
152 )
153 )
154 return Exporter.from_protobuf(exporter)
156 async def ListExporters(
157 self,
158 *,
159 page_size: int | None = None,
160 page_token: str | None = None,
161 filter: str | None = None,
162 ):
163 with translate_grpc_exceptions():
164 exporters = await self.stub.ListExporters(
165 client_pb2.ListExportersRequest(
166 parent="namespaces/{}".format(self.namespace),
167 page_size=page_size,
168 page_token=page_token,
169 filter=filter,
170 )
171 )
172 return ExporterList.from_protobuf(exporters)
174 async def GetLease(self, *, name: str):
175 with translate_grpc_exceptions():
176 lease = await self.stub.GetLease(
177 client_pb2.GetLeaseRequest(
178 name="namespaces/{}/leases/{}".format(self.namespace, name),
179 )
180 )
181 return Lease.from_protobuf(lease)
183 async def ListLeases(
184 self,
185 *,
186 page_size: int | None = None,
187 page_token: str | None = None,
188 filter: str | None = None,
189 ):
190 with translate_grpc_exceptions():
191 leases = await self.stub.ListLeases(
192 client_pb2.ListLeasesRequest(
193 parent="namespaces/{}".format(self.namespace),
194 page_size=page_size,
195 page_token=page_token,
196 filter=filter,
197 )
198 )
199 return LeaseList.from_protobuf(leases)
201 async def CreateLease(
202 self,
203 *,
204 selector: str,
205 duration: timedelta,
206 ):
207 duration_pb = duration_pb2.Duration()
208 duration_pb.FromTimedelta(duration)
210 with translate_grpc_exceptions():
211 lease = await self.stub.CreateLease(
212 client_pb2.CreateLeaseRequest(
213 parent="namespaces/{}".format(self.namespace),
214 lease=client_pb2.Lease(
215 duration=duration_pb,
216 selector=selector,
217 ),
218 )
219 )
220 return Lease.from_protobuf(lease)
222 async def UpdateLease(
223 self,
224 *,
225 name: str,
226 duration: timedelta,
227 ):
228 duration_pb = duration_pb2.Duration()
229 duration_pb.FromTimedelta(duration)
231 update_mask = field_mask_pb2.FieldMask()
232 update_mask.FromJsonString("duration")
234 with translate_grpc_exceptions():
235 lease = await self.stub.UpdateLease(
236 client_pb2.UpdateLeaseRequest(
237 lease=client_pb2.Lease(
238 name="namespaces/{}/leases/{}".format(self.namespace, name),
239 duration=duration_pb,
240 ),
241 update_mask=update_mask,
242 )
243 )
244 return Lease.from_protobuf(lease)
246 async def DeleteLease(self, *, name: str):
247 with translate_grpc_exceptions():
248 await self.stub.DeleteLease(
249 client_pb2.DeleteLeaseRequest(
250 name="namespaces/{}/leases/{}".format(self.namespace, name),
251 )
252 )