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

1from __future__ import annotations 

2 

3from dataclasses import dataclass, field 

4from datetime import datetime, timedelta 

5 

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 

11 

12from jumpstarter.common.grpc import translate_grpc_exceptions 

13 

14 

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] 

24 

25 

26def parse_client_identifier(identifier: str) -> (str, str): 

27 return parse_identifier(identifier, "clients") 

28 

29 

30def parse_exporter_identifier(identifier: str) -> (str, str): 

31 return parse_identifier(identifier, "exporters") 

32 

33 

34def parse_lease_identifier(identifier: str) -> (str, str): 

35 return parse_identifier(identifier, "leases") 

36 

37 

38class Exporter(BaseModel): 

39 namespace: str 

40 name: str 

41 labels: dict[str, str] 

42 

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) 

47 

48 

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 

58 

59 model_config = ConfigDict( 

60 arbitrary_types_allowed=True, 

61 ser_json_timedelta="float", 

62 ) 

63 

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] 

67 

68 @classmethod 

69 def from_protobuf(cls, data: client_pb2.Lease) -> Lease: 

70 namespace, name = parse_lease_identifier(data.name) 

71 

72 _, client = parse_client_identifier(data.client) 

73 if data.exporter != "": 

74 _, exporter = parse_exporter_identifier(data.exporter) 

75 else: 

76 exporter = "" 

77 

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 ) 

83 

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 ) 

94 

95 def dump_json(self): 

96 return self.model_dump_json(indent=4, by_alias=True) 

97 

98 def dump_yaml(self): 

99 return yaml.safe_dump(self.model_dump(mode="json", by_alias=True), indent=2) 

100 

101 

102class ExporterList(BaseModel): 

103 exporters: list[Exporter] 

104 next_page_token: str | None = Field(exclude=True) 

105 

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 ) 

112 

113 def dump_json(self): 

114 return self.model_dump_json(indent=4, by_alias=True) 

115 

116 def dump_yaml(self): 

117 return yaml.safe_dump(self.model_dump(mode="json", by_alias=True), indent=2) 

118 

119 

120class LeaseList(BaseModel): 

121 leases: list[Lease] 

122 next_page_token: str | None = Field(exclude=True) 

123 

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 ) 

130 

131 def dump_json(self): 

132 return self.model_dump_json(indent=4, by_alias=True) 

133 

134 def dump_yaml(self): 

135 return yaml.safe_dump(self.model_dump(mode="json", by_alias=True), indent=2) 

136 

137 

138@dataclass(kw_only=True, slots=True) 

139class ClientService: 

140 channel: Channel 

141 namespace: str 

142 stub: client_pb2_grpc.ClientServiceStub = field(init=False) 

143 

144 def __post_init__(self): 

145 self.stub = client_pb2_grpc.ClientServiceStub(channel=self.channel) 

146 

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) 

155 

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) 

173 

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) 

182 

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) 

200 

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) 

209 

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) 

221 

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) 

230 

231 update_mask = field_mask_pb2.FieldMask() 

232 update_mask.FromJsonString("duration") 

233 

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) 

245 

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 )