18585923528

Committed 17 Oct 2025 07:36AM UTC coverage: 83.21% (-0.2%) from 83.365%

Build # 18585923528

Build Type

Pull #1068

github

Committed by

web-flow

Commit Message

Merge 5051fe076 into 6cc42274c

Pull Request Pull Request #1068: feat: update data connectors when resuming a session

Run Details

47 of 118 new or added lines in 4 files covered. (39.83%)

11 existing lines in 5 files now uncovered.

21787 of 26183 relevant lines covered (83.21%)

1.49 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

63.5

/components/renku_data_services/notebooks/models.py

"""Basic models for amalthea sessions."""

from dataclasses import dataclass, field
from enum import StrEnum
from pathlib import Path
from typing import Any, cast

from kubernetes.client import V1ObjectMeta, V1Secret
from pydantic import AliasGenerator, BaseModel, Field, Json
from ulid import ULID

from renku_data_services.data_connectors.models import DataConnectorSecret
from renku_data_services.errors import errors
from renku_data_services.errors.errors import ProgrammingError
from renku_data_services.notebooks.api.schemas.cloud_storage import RCloneStorageRequestOverride
from renku_data_services.notebooks.crs import (
    AmaltheaSessionV1Alpha1,
    DataSource,
    ExtraContainer,
    ExtraVolume,
    ExtraVolumeMount,
    InitContainer,
    SecretRef,
)


@dataclass
class SessionEnvVar:
    """Environment variables for an amalthea session."""

    name: str
    value: str


@dataclass
class SessionUserSecrets:
    """User secret mounted in an amalthea session."""

    mount_path: Path
    user_secret_ids: list[str]


class _AmaltheaSessionAnnotations(BaseModel):
    class Config:
        extra = "allow"
        alias_generator = AliasGenerator(
            alias=lambda field_name: f"renku.io/{field_name}",
        )

    session_launcher_id: str | None = None
    project_id: str | None = None
    user_secrets_mount_path: str | None = None
    user_secrets_ids: Json[list[str]] = Field(default_factory=list)
    env_variable_names: Json[list[str]] = Field(default_factory=list)


class _MetadataValidation(BaseModel):
    class Config:
        extra = "allow"

    name: str
    annotations: _AmaltheaSessionAnnotations
    labels: dict[str, str] = Field(default_factory=dict)
    namespace: str | None = None


class AmaltheaSessionManifest:
    """The manifest for an amalthea session."""

    def __init__(self, manifest: AmaltheaSessionV1Alpha1) -> None:
        self._manifest = manifest
        self._metadata = _MetadataValidation.model_validate(self._manifest.metadata)

    def __repr__(self) -> str:
        return f"{self.__class__}(name={self._metadata.name})"

    @property
    def env_vars(self) -> dict[str, SessionEnvVar]:
        """Extract the environment variables from a manifest."""
        output: dict[str, SessionEnvVar] = {}
        assert self._manifest.spec
        for env in self._manifest.spec.session.env or []:
            if env.value is None:
                continue
            output[env.name] = SessionEnvVar(env.name, env.value)
        return output

    @property
    def requested_env_vars(self) -> dict[str, SessionEnvVar]:
        """The environment variables requested."""
        requested_names = self._metadata.annotations.env_variable_names
        return {ikey: ival for ikey, ival in self.env_vars.items() if ikey in requested_names}


@dataclass
class ExtraSecret:
    """Specification for a K8s secret and its coresponding volumes and mounts."""

    secret: V1Secret = field(repr=False)
    volume: ExtraVolume | None = None
    volume_mount: ExtraVolumeMount | None = None
    adopt: bool = True

    def __post_init__(self) -> None:
        if not self.secret.metadata:
            raise errors.ValidationError(message="The secret in Extra secret is missing its metadata.")
        if isinstance(self.secret.metadata, V1ObjectMeta):
            secret_name = cast(str | None, self.secret.metadata.name)
        else:
            secret_name = cast(str | None, self.secret.metadata.get("name"))
        if not isinstance(secret_name, str):
            raise errors.ValidationError(message="The secret name in Extra secret is not a string.")
        if len(secret_name) == 0:
            raise errors.ValidationError(message="The secret name in Extra secret is empty.")
        self.__secret_name = secret_name

    def key_ref(self, key: str | None = None) -> SecretRef:
        """Get an amalthea secret key reference."""
        meta = self.secret.metadata
        if not meta:
            raise ProgrammingError(message="Cannot get reference to a secret that does not have metadata.")
        secret_name = meta.name
        if not secret_name:
            raise ProgrammingError(message="Cannot get reference to a secret that does not have a name.")
        data = self.secret.data or {}
        string_data = self.secret.string_data or {}
        if key is not None and key not in data and key not in string_data:
            raise KeyError(f"Cannot find the key {key} in the secret with name {secret_name}")
        return SecretRef(key=key, name=secret_name, adopt=self.adopt)

    def ref(self) -> SecretRef:
        """Get an amalthea reference to the whole secret."""
        meta = self.secret.metadata
        if not meta:
            raise ProgrammingError(message="Cannot get reference to a secret that does not have metadata.")
        secret_name = meta.name
        if not secret_name:
            raise ProgrammingError(message="Cannot get reference to a secret that does not have a name.")
        return SecretRef(name=secret_name, adopt=self.adopt)

    @property
    def name(self) -> str:
        """Return the name of the secret."""
        return self.__secret_name


@dataclass(frozen=True, kw_only=True)
class SessionExtraResources:
    """Represents extra resources to add to an amalthea session."""

    annotations: dict[str, str] = field(default_factory=dict)
    containers: list[ExtraContainer] = field(default_factory=list)
    data_connector_secrets: dict[str, list[DataConnectorSecret]] = field(default_factory=dict)
    data_sources: list[DataSource] = field(default_factory=list)
    init_containers: list[InitContainer] = field(default_factory=list)
    secrets: list[ExtraSecret] = field(default_factory=list)
    volume_mounts: list[ExtraVolumeMount] = field(default_factory=list)
    volumes: list[ExtraVolume] = field(default_factory=list)

    def concat(self, added_extras: "SessionExtraResources | None") -> "SessionExtraResources":
        """Concatenates these session extras with more session extras."""
        if added_extras is None:
            return self
        annotations: dict[str, str] = dict()
        annotations.update(self.annotations)
        annotations.update(added_extras.annotations)
        data_connector_secrets: dict[str, list[DataConnectorSecret]] = dict()
        data_connector_secrets.update(self.data_connector_secrets)
        data_connector_secrets.update(added_extras.data_connector_secrets)
        return SessionExtraResources(
            annotations=annotations,
            containers=self.containers + added_extras.containers,
            data_connector_secrets=data_connector_secrets,
            data_sources=self.data_sources + added_extras.data_sources,
            init_containers=self.init_containers + added_extras.init_containers,
            secrets=self.secrets + added_extras.secrets,
            volume_mounts=self.volume_mounts + added_extras.volume_mounts,
            volumes=self.volumes + added_extras.volumes,
        )


@dataclass(eq=True, kw_only=True)
class SessionDataConnectorOverride(RCloneStorageRequestOverride):
    """Model for a data connector override."""

    skip: bool
    data_connector_id: ULID
    configuration: dict[str, Any] | None
    source_path: str | None
    target_path: str | None
    readonly: bool | None


@dataclass(frozen=True, eq=True, kw_only=True)
class SessionLaunchRequest:
    """Model for requesting a session launch."""

    launcher_id: ULID
    disk_storage: int | None
    resource_class_id: int | None
    data_connectors_overrides: list[SessionDataConnectorOverride] | None
    env_variable_overrides: list[SessionEnvVar] | None


class SessionState(StrEnum):
    """Session state."""

    running = "running"
    hibernated = "hibernated"


@dataclass(frozen=True, eq=True, kw_only=True)
class SessionPatchRequest:
    """Model for patching a session."""

    resource_class_id: int | None
    state: SessionState | None

1	"""Basic models for amalthea sessions."""
2
3	from dataclasses import dataclass, field	2✔
4	from enum import StrEnum	2✔
5	from pathlib import Path	2✔
6	from typing import Any, cast	2✔
7
8	from kubernetes.client import V1ObjectMeta, V1Secret	2✔
9	from pydantic import AliasGenerator, BaseModel, Field, Json	2✔
10	from ulid import ULID	2✔
11
12	from renku_data_services.data_connectors.models import DataConnectorSecret	2✔
13	from renku_data_services.errors import errors	2✔
14	from renku_data_services.errors.errors import ProgrammingError	2✔
15	from renku_data_services.notebooks.api.schemas.cloud_storage import RCloneStorageRequestOverride	2✔
16	from renku_data_services.notebooks.crs import (	2✔
17	AmaltheaSessionV1Alpha1,
18	DataSource,
19	ExtraContainer,
20	ExtraVolume,
21	ExtraVolumeMount,
22	InitContainer,
23	SecretRef,
24	)
25
26
27	@dataclass	2✔
28	class SessionEnvVar:	2✔
29	"""Environment variables for an amalthea session."""
30
31	name: str	2✔
32	value: str	2✔
33
34
35	@dataclass	2✔
36	class SessionUserSecrets:	2✔
37	"""User secret mounted in an amalthea session."""
38
39	mount_path: Path	2✔
40	user_secret_ids: list[str]	2✔
41
42
43	class _AmaltheaSessionAnnotations(BaseModel):	2✔
44	class Config:	2✔
45	extra = "allow"	2✔
46	alias_generator = AliasGenerator(	2✔
47	alias=lambda field_name: f"renku.io/{field_name}",
48	)
49
50	session_launcher_id: str \| None = None	2✔
51	project_id: str \| None = None	2✔
52	user_secrets_mount_path: str \| None = None	2✔
53	user_secrets_ids: Json[list[str]] = Field(default_factory=list)	2✔
54	env_variable_names: Json[list[str]] = Field(default_factory=list)	2✔
55
56
57	class _MetadataValidation(BaseModel):	2✔
58	class Config:	2✔
59	extra = "allow"	2✔
60
61	name: str	2✔
62	annotations: _AmaltheaSessionAnnotations	2✔
63	labels: dict[str, str] = Field(default_factory=dict)	2✔
64	namespace: str \| None = None	2✔
65
66
67	class AmaltheaSessionManifest:	2✔
68	"""The manifest for an amalthea session."""
69
70	def __init__(self, manifest: AmaltheaSessionV1Alpha1) -> None:	2✔
71	self._manifest = manifest	×
72	self._metadata = _MetadataValidation.model_validate(self._manifest.metadata)	×
73
74	def __repr__(self) -> str:	2✔
75	return f"{self.__class__}(name={self._metadata.name})"	×
76
77	@property	2✔
78	def env_vars(self) -> dict[str, SessionEnvVar]:	2✔
79	"""Extract the environment variables from a manifest."""
80	output: dict[str, SessionEnvVar] = {}	×
81	assert self._manifest.spec	×
82	for env in self._manifest.spec.session.env or []:	×
83	if env.value is None:	×
84	continue	×
85	output[env.name] = SessionEnvVar(env.name, env.value)	×
86	return output	×
87
88	@property	2✔
89	def requested_env_vars(self) -> dict[str, SessionEnvVar]:	2✔
90	"""The environment variables requested."""
91	requested_names = self._metadata.annotations.env_variable_names	×
92	return {ikey: ival for ikey, ival in self.env_vars.items() if ikey in requested_names}	×
93
94
95	@dataclass	2✔
96	class ExtraSecret:	2✔
97	"""Specification for a K8s secret and its coresponding volumes and mounts."""
98
99	secret: V1Secret = field(repr=False)	2✔
100	volume: ExtraVolume \| None = None	2✔
101	volume_mount: ExtraVolumeMount \| None = None	2✔
102	adopt: bool = True	2✔
103
104	def __post_init__(self) -> None:	2✔
105	if not self.secret.metadata:	×
106	raise errors.ValidationError(message="The secret in Extra secret is missing its metadata.")	×
107	if isinstance(self.secret.metadata, V1ObjectMeta):	×
108	secret_name = cast(str \| None, self.secret.metadata.name)	×
109	else:
110	secret_name = cast(str \| None, self.secret.metadata.get("name"))	×
111	if not isinstance(secret_name, str):	×
112	raise errors.ValidationError(message="The secret name in Extra secret is not a string.")	×
113	if len(secret_name) == 0:	×
114	raise errors.ValidationError(message="The secret name in Extra secret is empty.")	×
115	self.__secret_name = secret_name	×
116
117	def key_ref(self, key: str \| None = None) -> SecretRef:	2✔
118	"""Get an amalthea secret key reference."""
119	meta = self.secret.metadata	×
120	if not meta:	×
121	raise ProgrammingError(message="Cannot get reference to a secret that does not have metadata.")	×
122	secret_name = meta.name	×
123	if not secret_name:	×
124	raise ProgrammingError(message="Cannot get reference to a secret that does not have a name.")	×
125	data = self.secret.data or {}	×
126	string_data = self.secret.string_data or {}	×
127	if key is not None and key not in data and key not in string_data:	×
128	raise KeyError(f"Cannot find the key {key} in the secret with name {secret_name}")	×
129	return SecretRef(key=key, name=secret_name, adopt=self.adopt)	×
130
131	def ref(self) -> SecretRef:	2✔
132	"""Get an amalthea reference to the whole secret."""
133	meta = self.secret.metadata	×
134	if not meta:	×
135	raise ProgrammingError(message="Cannot get reference to a secret that does not have metadata.")	×
136	secret_name = meta.name	×
137	if not secret_name:	×
138	raise ProgrammingError(message="Cannot get reference to a secret that does not have a name.")	×
139	return SecretRef(name=secret_name, adopt=self.adopt)	×
140
141	@property	2✔
142	def name(self) -> str:	2✔
143	"""Return the name of the secret."""
144	return self.__secret_name	×
145
146
147	@dataclass(frozen=True, kw_only=True)	2✔
148	class SessionExtraResources:	2✔
149	"""Represents extra resources to add to an amalthea session."""
150
151	annotations: dict[str, str] = field(default_factory=dict)	2✔
152	containers: list[ExtraContainer] = field(default_factory=list)	2✔
153	data_connector_secrets: dict[str, list[DataConnectorSecret]] = field(default_factory=dict)	2✔
154	data_sources: list[DataSource] = field(default_factory=list)	2✔
155	init_containers: list[InitContainer] = field(default_factory=list)	2✔
156	secrets: list[ExtraSecret] = field(default_factory=list)	2✔
157	volume_mounts: list[ExtraVolumeMount] = field(default_factory=list)	2✔
158	volumes: list[ExtraVolume] = field(default_factory=list)	2✔
159
160	def concat(self, added_extras: "SessionExtraResources \| None") -> "SessionExtraResources":	2✔
161	"""Concatenates these session extras with more session extras."""
162	if added_extras is None:	×
163	return self	×
NEW 164	annotations: dict[str, str] = dict()	×
NEW 165	annotations.update(self.annotations)	×
NEW 166	annotations.update(added_extras.annotations)	×
167	data_connector_secrets: dict[str, list[DataConnectorSecret]] = dict()	×
168	data_connector_secrets.update(self.data_connector_secrets)	×
169	data_connector_secrets.update(added_extras.data_connector_secrets)	×
170	return SessionExtraResources(	×
171	annotations=annotations,
172	containers=self.containers + added_extras.containers,
173	data_connector_secrets=data_connector_secrets,
174	data_sources=self.data_sources + added_extras.data_sources,
175	init_containers=self.init_containers + added_extras.init_containers,
176	secrets=self.secrets + added_extras.secrets,
177	volume_mounts=self.volume_mounts + added_extras.volume_mounts,
178	volumes=self.volumes + added_extras.volumes,
179	)
180
181
182	@dataclass(eq=True, kw_only=True)	2✔
183	class SessionDataConnectorOverride(RCloneStorageRequestOverride):	2✔
184	"""Model for a data connector override."""
185
186	skip: bool	2✔
187	data_connector_id: ULID	2✔
188	configuration: dict[str, Any] \| None	2✔
189	source_path: str \| None	2✔
190	target_path: str \| None	2✔
191	readonly: bool \| None	2✔
192
193
194	@dataclass(frozen=True, eq=True, kw_only=True)	2✔
195	class SessionLaunchRequest:	2✔
196	"""Model for requesting a session launch."""
197
198	launcher_id: ULID	2✔
199	disk_storage: int \| None	2✔
200	resource_class_id: int \| None	2✔
201	data_connectors_overrides: list[SessionDataConnectorOverride] \| None	2✔
202	env_variable_overrides: list[SessionEnvVar] \| None	2✔
203
204
205	class SessionState(StrEnum):	2✔
206	"""Session state."""
207
208	running = "running"	2✔
209	hibernated = "hibernated"	2✔
210
211
212	@dataclass(frozen=True, eq=True, kw_only=True)	2✔
213	class SessionPatchRequest:	2✔
214	"""Model for patching a session."""
215
216	resource_class_id: int \| None	2✔
217	state: SessionState \| None	2✔

SwissDataScienceCenter / renku-data-services / 18585923528

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous