Coverage for /Users/ajo/work/jumpstarter/jumpstarter/packages/jumpstarter/jumpstarter/config/exporter.py: 56%

124 statements  

« prev     ^ index     » next       coverage.py v7.8.0, created at 2025-05-06 10:20 +0200

1from __future__ import annotations 

2 

3from contextlib import asynccontextmanager, contextmanager, suppress 

4from pathlib import Path 

5from typing import Any, ClassVar, Literal, Optional, Self 

6 

7import grpc 

8import yaml 

9from anyio.from_thread import start_blocking_portal 

10from pydantic import BaseModel, ConfigDict, Field, RootModel 

11 

12from .common import ObjectMeta 

13from .grpc import call_credentials 

14from .tls import TLSConfigV1Alpha1 

15from jumpstarter.common.grpc import aio_secure_channel, ssl_channel_credentials 

16from jumpstarter.common.importlib import import_class 

17from jumpstarter.driver import Driver 

18 

19 

20class ExporterConfigV1Alpha1DriverInstanceProxy(BaseModel): 

21 ref: str 

22 

23 

24class ExporterConfigV1Alpha1DriverInstanceComposite(BaseModel): 

25 children: dict[str, ExporterConfigV1Alpha1DriverInstance] = Field(default_factory=dict) 

26 

27 

28class ExporterConfigV1Alpha1DriverInstanceBase(BaseModel): 

29 type: str 

30 config: dict[str, Any] = Field(default_factory=dict) 

31 children: dict[str, ExporterConfigV1Alpha1DriverInstance] = Field(default_factory=dict) 

32 

33 

34class ExporterConfigV1Alpha1DriverInstance(RootModel): 

35 root: ( 

36 ExporterConfigV1Alpha1DriverInstanceBase 

37 | ExporterConfigV1Alpha1DriverInstanceComposite 

38 | ExporterConfigV1Alpha1DriverInstanceProxy 

39 ) 

40 

41 def instantiate(self) -> Driver: 

42 match self.root: 

43 case ExporterConfigV1Alpha1DriverInstanceBase(): 

44 driver_class = import_class(self.root.type, [], True) 

45 

46 children = {name: child.instantiate() for name, child in self.root.children.items()} 

47 

48 return driver_class(children=children, **self.root.config) 

49 

50 case ExporterConfigV1Alpha1DriverInstanceComposite(): 

51 from jumpstarter_driver_composite.driver import Composite 

52 

53 children = {name: child.instantiate() for name, child in self.root.children.items()} 

54 

55 return Composite(children=children) 

56 

57 case ExporterConfigV1Alpha1DriverInstanceProxy(): 

58 from jumpstarter_driver_composite.driver import Proxy 

59 

60 return Proxy(ref=self.root.ref) 

61 

62 @classmethod 

63 def from_path(cls, path: str) -> ExporterConfigV1Alpha1DriverInstance: 

64 with open(path) as f: 

65 return cls.model_validate(yaml.safe_load(f)) 

66 

67 @classmethod 

68 def from_str(cls, config: str) -> ExporterConfigV1Alpha1DriverInstance: 

69 return cls.model_validate(yaml.safe_load(config)) 

70 

71 

72class ExporterConfigV1Alpha1(BaseModel): 

73 BASE_PATH: ClassVar[Path] = Path("/etc/jumpstarter/exporters") 

74 

75 alias: str = Field(default="default") 

76 

77 apiVersion: Literal["jumpstarter.dev/v1alpha1"] = Field(default="jumpstarter.dev/v1alpha1") 

78 kind: Literal["ExporterConfig"] = Field(default="ExporterConfig") 

79 metadata: ObjectMeta 

80 

81 endpoint: str 

82 tls: TLSConfigV1Alpha1 = Field(default_factory=TLSConfigV1Alpha1) 

83 token: str 

84 grpcOptions: dict[str, str | int] | None = Field(default_factory=dict) 

85 

86 export: dict[str, ExporterConfigV1Alpha1DriverInstance] = Field(default_factory=dict) 

87 

88 path: Path | None = Field(default=None) 

89 

90 @classmethod 

91 def _get_path(cls, alias: str): 

92 return (cls.BASE_PATH / alias).with_suffix(".yaml") 

93 

94 @classmethod 

95 def exists(cls, alias: str): 

96 return cls._get_path(alias).exists() 

97 

98 @classmethod 

99 def load_path(cls, path: Path): 

100 with path.open() as f: 

101 config = cls.model_validate(yaml.safe_load(f)) 

102 config.path = path 

103 return config 

104 

105 @classmethod 

106 def load(cls, alias: str) -> Self: 

107 config = cls.load_path(cls._get_path(alias)) 

108 config.alias = alias 

109 return config 

110 

111 @classmethod 

112 def list(cls) -> list[Self]: 

113 exporters = [] 

114 with suppress(FileNotFoundError): 

115 for entry in cls.BASE_PATH.iterdir(): 

116 exporters.append(cls.load(entry.stem)) 

117 return exporters 

118 

119 @classmethod 

120 def dump_yaml(self, config: Self) -> str: 

121 return yaml.safe_dump(config.model_dump(mode="json", exclude={"alias", "path"}), sort_keys=False) 

122 

123 @classmethod 

124 def save(cls, config: Self, path: Optional[str] = None) -> Path: 

125 # Set the config path before saving 

126 if path is None: 

127 config.path = cls._get_path(config.alias) 

128 config.path.parent.mkdir(parents=True, exist_ok=True) 

129 else: 

130 config.path = Path(path) 

131 with config.path.open(mode="w") as f: 

132 yaml.safe_dump(config.model_dump(mode="json", exclude={"alias", "path"}), f, sort_keys=False) 

133 return config.path 

134 

135 @classmethod 

136 def delete(cls, alias: str) -> Path: 

137 path = cls._get_path(alias) 

138 path.unlink(missing_ok=True) 

139 return path 

140 

141 @asynccontextmanager 

142 async def serve_unix_async(self): 

143 # dynamic import to avoid circular imports 

144 from jumpstarter.exporter import Session 

145 

146 with Session( 

147 root_device=ExporterConfigV1Alpha1DriverInstance(children=self.export).instantiate(), 

148 ) as session: 

149 async with session.serve_unix_async() as path: 

150 yield path 

151 

152 @contextmanager 

153 def serve_unix(self): 

154 with start_blocking_portal() as portal: 

155 with portal.wrap_async_context_manager(self.serve_unix_async()) as path: 

156 yield path 

157 

158 async def serve(self): 

159 # dynamic import to avoid circular imports 

160 from jumpstarter.exporter import Exporter 

161 

162 async def channel_factory(): 

163 credentials = grpc.composite_channel_credentials( 

164 await ssl_channel_credentials(self.endpoint, self.tls), 

165 call_credentials("Exporter", self.metadata, self.token), 

166 ) 

167 return aio_secure_channel(self.endpoint, credentials, self.grpcOptions) 

168 

169 async with Exporter( 

170 channel_factory=channel_factory, 

171 device_factory=ExporterConfigV1Alpha1DriverInstance(children=self.export).instantiate, 

172 tls=self.tls, 

173 grpc_options=self.grpcOptions, 

174 ) as exporter: 

175 await exporter.serve() 

176 

177 

178class ExporterConfigListV1Alpha1(BaseModel): 

179 api_version: Literal["jumpstarter.dev/v1alpha1"] = Field(alias="apiVersion", default="jumpstarter.dev/v1alpha1") 

180 items: list[ExporterConfigV1Alpha1] 

181 kind: Literal["ExporterConfigList"] = Field(default="ExporterConfigList") 

182 

183 def dump_json(self): 

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

185 

186 def dump_yaml(self): 

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

188 

189 model_config = ConfigDict(arbitrary_types_allowed=True, populate_by_name=True)