fix: allow session launcher parameters to be reset (#434)

Allows the API to accept None as input for args, command and the session
launcher resource class ID so that they can be reset to their defaults
in patch endpoints.

Author: olevski

components/renku_data_services/base_models/core.py

@@ -220,4 +220,4 @@ async def authenticate(self, access_token: str, request: Request) -> AnyAPIUser:
 """
 
 RESET: ResetType = ResetType(object())
-"""The single instance of the ResetType, can be compared to similar to None, i.e. `if value is RESET`"""
+"""The single instance of the ResetType, can be compared to similar to None, i.e. `if value is RESET`"""

components/renku_data_services/session/converters.py

@@ -0,0 +1,83 @@
+"""Code used to convert from/to apispec and models."""
+
+from pathlib import PurePosixPath
+
+from renku_data_services.base_models.core import RESET, ResetType
+from renku_data_services.session import apispec, models
+
+
+def environment_update_from_patch(data: apispec.EnvironmentPatch) -> models.EnvironmentUpdate:
+    """Create an update object from an apispec or any other pydantic model."""
+    data_dict = data.model_dump(exclude_unset=True, mode="json")
+    working_directory: PurePosixPath | None = None
+    if data.working_directory is not None:
+        working_directory = PurePosixPath(data.working_directory)
+    mount_directory: PurePosixPath | None = None
+    if data.mount_directory is not None:
+        mount_directory = PurePosixPath(data.mount_directory)
+    # NOTE: If the args or command are present in the data_dict and they are None they were passed in by the user.
+    # The None specifically passed by the user indicates that the value should be removed from the DB.
+    args = RESET if "args" in data_dict and data_dict["args"] is None else data.args
+    command = RESET if "command" in data_dict and data_dict["command"] is None else data.command
+    return models.EnvironmentUpdate(
+        name=data.name,
+        description=data.description,
+        container_image=data.container_image,
+        default_url=data.default_url,
+        port=data.port,
+        working_directory=working_directory,
+        mount_directory=mount_directory,
+        uid=data.uid,
+        gid=data.gid,
+        args=args,
+        command=command,
+    )
+
+
+def launcher_update_from_patch(
+    data: apispec.SessionLauncherPatch,
+    current_launcher: models.SessionLauncher | None = None,
+) -> models.SessionLauncherUpdate:
+    """Create an update object from an apispec or any other pydantic model."""
+    data_dict = data.model_dump(exclude_unset=True, mode="json")
+    environment: str | models.EnvironmentUpdate | models.UnsavedEnvironment | None = None
+    if (
+        isinstance(data.environment, apispec.EnvironmentPatchInLauncher)
+        and current_launcher is not None
+        and current_launcher.environment.environment_kind == models.EnvironmentKind.GLOBAL
+        and data.environment.environment_kind == apispec.EnvironmentKind.CUSTOM
+    ):
+        # This means that the global environment is being swapped for a custom one,
+        # so we have to create a brand new environment, but we have to validate here.
+        validated_env = apispec.EnvironmentPostInLauncher.model_validate(data_dict["environment"])
+        environment = models.UnsavedEnvironment(
+            name=validated_env.name,
+            description=validated_env.description,
+            container_image=validated_env.container_image,
+            default_url=validated_env.default_url,
+            port=validated_env.port,
+            working_directory=PurePosixPath(validated_env.working_directory),
+            mount_directory=PurePosixPath(validated_env.mount_directory),
+            uid=validated_env.uid,
+            gid=validated_env.gid,
+            environment_kind=models.EnvironmentKind(validated_env.environment_kind.value),
+            args=validated_env.args,
+            command=validated_env.command,
+        )
+    elif isinstance(data.environment, apispec.EnvironmentPatchInLauncher):
+        environment = environment_update_from_patch(data.environment)
+    elif isinstance(data.environment, apispec.EnvironmentIdOnlyPatch):
+        environment = data.environment.id
+    resource_class_id: int | None | ResetType = None
+    if "resource_class_id" in data_dict and data_dict["resource_class_id"] is None:
+        # NOTE: This means that the resource class set in the DB should be removed so that the
+        # default resource class currently set in the CRC will be used.
+        resource_class_id = RESET
+    else:
+        resource_class_id = data_dict.get("resource_class_id")
+    return models.SessionLauncherUpdate(
+        name=data_dict.get("name"),
+        description=data_dict.get("description"),
+        environment=environment,
+        resource_class_id=resource_class_id,
+    )

components/renku_data_services/session/db.py

@@ -147,7 +147,7 @@ async def update_environment(
             res = await session.scalars(
                 select(schemas.EnvironmentORM)
                 .where(schemas.EnvironmentORM.id == str(environment_id))
-                .where(schemas.EnvironmentORM.environment_kind == models.EnvironmentKind.GLOBAL.value)
+                .where(schemas.EnvironmentORM.environment_kind == models.EnvironmentKind.GLOBAL)
             )
             environment = res.one_or_none()
             if environment is None:

components/renku_data_services/session/models.py

@@ -73,6 +73,23 @@ class Environment(UnsavedEnvironment):
     gid: int
 
 
+@dataclass(kw_only=True, frozen=True, eq=True)
+class EnvironmentUpdate:
+    """Model for the update of some or all parts of an environment."""
+
+    name: str | None = None
+    description: str | None = None
+    container_image: str | None = None
+    default_url: str | None = None
+    port: int | None = None
+    working_directory: PurePosixPath | None = None
+    mount_directory: PurePosixPath | None = None
+    uid: int | None = None
+    gid: int | None = None
+    args: list[str] | None | ResetType = None
+    command: list[str] | None | ResetType = None
+
+
 @dataclass(frozen=True, eq=True, kw_only=True)
 class EnvironmentPatch:
     """Model for changes requested on a session environment."""

test/bases/renku_data_services/data_api/test_sessions.py

@@ -3,6 +3,7 @@
 import os
 import shutil
 from asyncio import AbstractEventLoop
+from collections.abc import Iterator
 from typing import Any
 
 import pytest
@@ -18,8 +19,8 @@
 os.environ["KUBECONFIG"] = ".k3d-config.yaml"
 
 
-@pytest.fixture(scope="module", autouse=True)
-def cluster() -> K3DCluster:
+@pytest.fixture(scope="module")
+def cluster() -> Iterator[K3DCluster]:
     if shutil.which("k3d") is None:
         pytest.skip("Requires k3d for cluster creation")
 
@@ -172,10 +173,14 @@ async def test_patch_session_environment(
     env = await create_session_environment("Environment 1")
     environment_id = env["id"]
 
+    command = ["python", "test.py"]
+    args = ["arg1", "arg2"]
     payload = {
         "name": "New name",
         "description": "New description.",
         "container_image": "new_image:new_tag",
+        "command": command,
+        "args": args,
     }
 
     _, res = await sanic_client.patch(f"/api/data/environments/{environment_id}", headers=admin_headers, json=payload)
@@ -185,6 +190,14 @@ async def test_patch_session_environment(
     assert res.json.get("name") == "New name"
     assert res.json.get("description") == "New description."
     assert res.json.get("container_image") == "new_image:new_tag"
+    assert res.json.get("args") == args
+    assert res.json.get("command") == command
+
+    # Test that patching with None will reset the command and args
+    payload = {"args": None, "command": None}
+    _, res = await sanic_client.patch(f"/api/data/environments/{environment_id}", headers=admin_headers, json=payload)
+    assert res.json.get("args") is None
+    assert res.json.get("command") is None
 
 
 @pytest.mark.asyncio
@@ -515,13 +528,25 @@ async def test_patch_session_launcher_environment(
 
     # Should be able to patch some fields of the custom environment
     patch_payload = {
-        "environment": {"container_image": "nginx:latest"},
+        "environment": {"container_image": "nginx:latest", "args": ["a", "b", "c"]},
     }
     _, res = await sanic_client.patch(
         f"/api/data/session_launchers/{launcher_id}", headers=user_headers, json=patch_payload
     )
     assert res.status_code == 200, res.text
     assert res.json["environment"]["container_image"] == "nginx:latest"
+    assert res.json["environment"]["args"] == ["a", "b", "c"]
+
+    # Should be able to reset args by patching in None, pathcing a null field should do nothing
+    patch_payload = {
+        "environment": {"args": None, "command": None},
+    }
+    _, res = await sanic_client.patch(
+        f"/api/data/session_launchers/{launcher_id}", headers=user_headers, json=patch_payload
+    )
+    assert res.status_code == 200, res.text
+    assert res.json["environment"].get("args") is None
+    assert res.json["environment"].get("command") is None
 
 
 @pytest.fixture
@@ -540,6 +565,7 @@ async def test_starting_session_anonymous(
     admin_headers,
     launch_session,
     anonymous_user_headers,
+    cluster,
 ) -> None:
     _, res = await sanic_client.post(
         "/api/data/resource_pools",