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
« prev ^ index » next coverage.py v7.8.0, created at 2025-05-06 10:20 +0200
1from __future__ import annotations
3from contextlib import asynccontextmanager, contextmanager, suppress
4from pathlib import Path
5from typing import Any, ClassVar, Literal, Optional, Self
7import grpc
8import yaml
9from anyio.from_thread import start_blocking_portal
10from pydantic import BaseModel, ConfigDict, Field, RootModel
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
20class ExporterConfigV1Alpha1DriverInstanceProxy(BaseModel):
21 ref: str
24class ExporterConfigV1Alpha1DriverInstanceComposite(BaseModel):
25 children: dict[str, ExporterConfigV1Alpha1DriverInstance] = Field(default_factory=dict)
28class ExporterConfigV1Alpha1DriverInstanceBase(BaseModel):
29 type: str
30 config: dict[str, Any] = Field(default_factory=dict)
31 children: dict[str, ExporterConfigV1Alpha1DriverInstance] = Field(default_factory=dict)
34class ExporterConfigV1Alpha1DriverInstance(RootModel):
35 root: (
36 ExporterConfigV1Alpha1DriverInstanceBase
37 | ExporterConfigV1Alpha1DriverInstanceComposite
38 | ExporterConfigV1Alpha1DriverInstanceProxy
39 )
41 def instantiate(self) -> Driver:
42 match self.root:
43 case ExporterConfigV1Alpha1DriverInstanceBase():
44 driver_class = import_class(self.root.type, [], True)
46 children = {name: child.instantiate() for name, child in self.root.children.items()}
48 return driver_class(children=children, **self.root.config)
50 case ExporterConfigV1Alpha1DriverInstanceComposite():
51 from jumpstarter_driver_composite.driver import Composite
53 children = {name: child.instantiate() for name, child in self.root.children.items()}
55 return Composite(children=children)
57 case ExporterConfigV1Alpha1DriverInstanceProxy():
58 from jumpstarter_driver_composite.driver import Proxy
60 return Proxy(ref=self.root.ref)
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))
67 @classmethod
68 def from_str(cls, config: str) -> ExporterConfigV1Alpha1DriverInstance:
69 return cls.model_validate(yaml.safe_load(config))
72class ExporterConfigV1Alpha1(BaseModel):
73 BASE_PATH: ClassVar[Path] = Path("/etc/jumpstarter/exporters")
75 alias: str = Field(default="default")
77 apiVersion: Literal["jumpstarter.dev/v1alpha1"] = Field(default="jumpstarter.dev/v1alpha1")
78 kind: Literal["ExporterConfig"] = Field(default="ExporterConfig")
79 metadata: ObjectMeta
81 endpoint: str
82 tls: TLSConfigV1Alpha1 = Field(default_factory=TLSConfigV1Alpha1)
83 token: str
84 grpcOptions: dict[str, str | int] | None = Field(default_factory=dict)
86 export: dict[str, ExporterConfigV1Alpha1DriverInstance] = Field(default_factory=dict)
88 path: Path | None = Field(default=None)
90 @classmethod
91 def _get_path(cls, alias: str):
92 return (cls.BASE_PATH / alias).with_suffix(".yaml")
94 @classmethod
95 def exists(cls, alias: str):
96 return cls._get_path(alias).exists()
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
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
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
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)
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
135 @classmethod
136 def delete(cls, alias: str) -> Path:
137 path = cls._get_path(alias)
138 path.unlink(missing_ok=True)
139 return path
141 @asynccontextmanager
142 async def serve_unix_async(self):
143 # dynamic import to avoid circular imports
144 from jumpstarter.exporter import Session
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
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
158 async def serve(self):
159 # dynamic import to avoid circular imports
160 from jumpstarter.exporter import Exporter
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)
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()
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")
183 def dump_json(self):
184 return self.model_dump_json(indent=4, by_alias=True)
186 def dump_yaml(self):
187 return yaml.safe_dump(self.model_dump(mode="json", by_alias=True), indent=2)
189 model_config = ConfigDict(arbitrary_types_allowed=True, populate_by_name=True)