Merge pull request #365 from jumpstarter-dev/duration

mangelajo · web-flow · commit 36015dd049f3 · 2025-03-20T17:31:23.000+01:00
Allow specifying duration in jmp client shell command
diff --git a/packages/jumpstarter-cli-client/jumpstarter_cli_client/client_shell.py b/packages/jumpstarter-cli-client/jumpstarter_cli_client/client_shell.py
@@ -1,23 +1,25 @@
 import sys
+from datetime import timedelta
 
 import asyncclick as click
 from jumpstarter_cli_common.exceptions import handle_exceptions
 
-from .common import opt_config, opt_selector_simple
+from .common import opt_config, opt_duration_partial, opt_selector_simple
 from jumpstarter.common.utils import launch_shell
 
 
 @click.command("shell", short_help="Spawns a shell connecting to a leased remote exporter")
 @click.option("-n", "--lease", "lease_name", type=str)
 @opt_config
 @opt_selector_simple
+@opt_duration_partial(default=timedelta(minutes=30), show_default="00:30:00")
 @handle_exceptions
-def client_shell(config, selector: str, lease_name):
+def client_shell(config, selector: str, duration: timedelta, lease_name):
     """Spawns a shell connecting to a leased remote exporter"""
 
     exit_code = 0
 
-    with config.lease(selector=selector, lease_name=lease_name) as lease:
+    with config.lease(selector=selector, lease_name=lease_name, duration=duration) as lease:
         with lease.serve_unix() as path:
             with lease.monitor():
                 exit_code = launch_shell(path, "remote", config.drivers.allow, config.drivers.unsafe)
diff --git a/packages/jumpstarter-cli-client/jumpstarter_cli_client/common.py b/packages/jumpstarter-cli-client/jumpstarter_cli_client/common.py
@@ -1,4 +1,5 @@
 from datetime import timedelta
+from functools import partial
 
 import asyncclick as click
 from pydantic import TypeAdapter
@@ -57,3 +58,20 @@ def convert(self, value, param, ctx):
 CLIENT = ClientParamType()
 
 opt_config = click.option("--client", "config", type=CLIENT, default=False, help="Name of client config")
+opt_duration_partial = partial(
+    click.option,
+    "--duration",
+    "duration",
+    type=DURATION,
+    help="""
+Accepted duration formats:
+
+\b
+PnYnMnDTnHnMnS - ISO 8601 duration format
+HH:MM:SS - time in hours, minutes, seconds
+D days, HH:MM:SS - time prefixed by X days
+D d, HH:MM:SS - time prefixed by X d
+
+See https://docs.rs/speedate/latest/speedate/ for details
+""",
+)
diff --git a/packages/jumpstarter-cli-client/jumpstarter_cli_client/create.py b/packages/jumpstarter-cli-client/jumpstarter_cli_client/create.py
@@ -9,7 +9,7 @@
 )
 from jumpstarter_cli_common.exceptions import handle_exceptions
 
-from .common import DURATION, opt_config, opt_selector_simple
+from .common import opt_config, opt_duration_partial, opt_selector_simple
 
 
 @click.group()
@@ -22,7 +22,7 @@ def create():
 @create.command(name="lease")
 @opt_config
 @opt_selector_simple
-@click.option("--duration", "duration", type=DURATION, required=True)
+@opt_duration_partial(required=True)
 @opt_output_all
 @handle_exceptions
 async def create_lease(config, selector: str, duration: timedelta, output: OutputType):
diff --git a/packages/jumpstarter/jumpstarter/client/grpc.py b/packages/jumpstarter/jumpstarter/client/grpc.py
@@ -9,6 +9,8 @@
 from jumpstarter_protocol import client_pb2, client_pb2_grpc, kubernetes_pb2
 from pydantic import BaseModel, ConfigDict, Field, field_serializer
 
+from jumpstarter.common.grpc import translate_grpc_exceptions
+
 
 def parse_identifier(identifier: str, kind: str) -> (str, str):
     segments = identifier.split("/")
@@ -143,11 +145,12 @@ def __post_init__(self):
         self.stub = client_pb2_grpc.ClientServiceStub(channel=self.channel)
 
     async def GetExporter(self, *, name: str):
-        exporter = await self.stub.GetExporter(
-            client_pb2.GetExporterRequest(
-                name="namespaces/{}/exporters/{}".format(self.namespace, name),
+        with translate_grpc_exceptions():
+            exporter = await self.stub.GetExporter(
+                client_pb2.GetExporterRequest(
+                    name="namespaces/{}/exporters/{}".format(self.namespace, name),
+                )
             )
-        )
         return Exporter.from_protobuf(exporter)
 
     async def ListExporters(
@@ -157,22 +160,24 @@ async def ListExporters(
         page_token: str | None = None,
         filter: str | None = None,
     ):
-        exporters = await self.stub.ListExporters(
-            client_pb2.ListExportersRequest(
-                parent="namespaces/{}".format(self.namespace),
-                page_size=page_size,
-                page_token=page_token,
-                filter=filter,
+        with translate_grpc_exceptions():
+            exporters = await self.stub.ListExporters(
+                client_pb2.ListExportersRequest(
+                    parent="namespaces/{}".format(self.namespace),
+                    page_size=page_size,
+                    page_token=page_token,
+                    filter=filter,
+                )
             )
-        )
         return ExporterList.from_protobuf(exporters)
 
     async def GetLease(self, *, name: str):
-        lease = await self.stub.GetLease(
-            client_pb2.GetLeaseRequest(
-                name="namespaces/{}/leases/{}".format(self.namespace, name),
+        with translate_grpc_exceptions():
+            lease = await self.stub.GetLease(
+                client_pb2.GetLeaseRequest(
+                    name="namespaces/{}/leases/{}".format(self.namespace, name),
+                )
             )
-        )
         return Lease.from_protobuf(lease)
 
     async def ListLeases(
@@ -182,14 +187,15 @@ async def ListLeases(
         page_token: str | None = None,
         filter: str | None = None,
     ):
-        leases = await self.stub.ListLeases(
-            client_pb2.ListLeasesRequest(
-                parent="namespaces/{}".format(self.namespace),
-                page_size=page_size,
-                page_token=page_token,
-                filter=filter,
+        with translate_grpc_exceptions():
+            leases = await self.stub.ListLeases(
+                client_pb2.ListLeasesRequest(
+                    parent="namespaces/{}".format(self.namespace),
+                    page_size=page_size,
+                    page_token=page_token,
+                    filter=filter,
+                )
             )
-        )
         return LeaseList.from_protobuf(leases)
 
     async def CreateLease(
@@ -201,15 +207,16 @@ async def CreateLease(
         duration_pb = duration_pb2.Duration()
         duration_pb.FromTimedelta(duration)
 
-        lease = await self.stub.CreateLease(
-            client_pb2.CreateLeaseRequest(
-                parent="namespaces/{}".format(self.namespace),
-                lease=client_pb2.Lease(
-                    duration=duration_pb,
-                    selector=selector,
-                ),
+        with translate_grpc_exceptions():
+            lease = await self.stub.CreateLease(
+                client_pb2.CreateLeaseRequest(
+                    parent="namespaces/{}".format(self.namespace),
+                    lease=client_pb2.Lease(
+                        duration=duration_pb,
+                        selector=selector,
+                    ),
+                )
             )
-        )
         return Lease.from_protobuf(lease)
 
     async def UpdateLease(
@@ -224,20 +231,22 @@ async def UpdateLease(
         update_mask = field_mask_pb2.FieldMask()
         update_mask.FromJsonString("duration")
 
-        lease = await self.stub.UpdateLease(
-            client_pb2.UpdateLeaseRequest(
-                lease=client_pb2.Lease(
-                    name="namespaces/{}/leases/{}".format(self.namespace, name),
-                    duration=duration_pb,
-                ),
-                update_mask=update_mask,
+        with translate_grpc_exceptions():
+            lease = await self.stub.UpdateLease(
+                client_pb2.UpdateLeaseRequest(
+                    lease=client_pb2.Lease(
+                        name="namespaces/{}/leases/{}".format(self.namespace, name),
+                        duration=duration_pb,
+                    ),
+                    update_mask=update_mask,
+                )
             )
-        )
         return Lease.from_protobuf(lease)
 
     async def DeleteLease(self, *, name: str):
-        await self.stub.DeleteLease(
-            client_pb2.DeleteLeaseRequest(
-                name="namespaces/{}/leases/{}".format(self.namespace, name),
+        with translate_grpc_exceptions():
+            await self.stub.DeleteLease(
+                client_pb2.DeleteLeaseRequest(
+                    name="namespaces/{}/leases/{}".format(self.namespace, name),
+                )
             )
-        )
diff --git a/packages/jumpstarter/jumpstarter/client/lease.py b/packages/jumpstarter/jumpstarter/client/lease.py
@@ -30,7 +30,7 @@
 @dataclass(kw_only=True)
 class Lease(AbstractContextManager, AbstractAsyncContextManager):
     channel: Channel
-    timeout: int = 1800
+    duration: timedelta
     selector: str
     portal: BlockingPortal
     namespace: str
@@ -51,17 +51,15 @@ def __post_init__(self):
         self.manager = self.portal.wrap_async_context_manager(self)
 
     async def _create(self):
-        duration = timedelta(seconds=self.timeout)
-
-        logger.debug("Creating lease request for selector %s for duration %s", self.selector, duration)
+        logger.debug("Creating lease request for selector %s for duration %s", self.selector, self.duration)
         with translate_grpc_exceptions():
             self.name = (
                 await self.svc.CreateLease(
                     selector=self.selector,
-                    duration=timedelta(seconds=self.timeout),
+                    duration=self.duration,
                 )
             ).name
-        logger.info("Created lease request for selector %s for duration %s", self.selector, duration)
+        logger.info("Created lease request for selector %s for duration %s", self.selector, self.duration)
 
     async def get(self):
         with translate_grpc_exceptions():
diff --git a/packages/jumpstarter/jumpstarter/config/client.py b/packages/jumpstarter/jumpstarter/config/client.py
@@ -15,7 +15,7 @@
 from .tls import TLSConfigV1Alpha1
 from jumpstarter.client.grpc import ClientService
 from jumpstarter.common.exceptions import FileNotFoundError
-from jumpstarter.common.grpc import aio_secure_channel, ssl_channel_credentials, translate_grpc_exceptions
+from jumpstarter.common.grpc import aio_secure_channel, ssl_channel_credentials
 
 
 def _allow_from_env():
@@ -61,9 +61,14 @@ async def channel(self):
         return aio_secure_channel(self.endpoint, credentials, self.grpcOptions)
 
     @contextmanager
-    def lease(self, selector: str | None = None, lease_name: str | None = None):
+    def lease(
+        self,
+        selector: str | None = None,
+        lease_name: str | None = None,
+        duration: timedelta = timedelta(minutes=30),
+    ):
         with start_blocking_portal() as portal:
-            with portal.wrap_async_context_manager(self.lease_async(selector, lease_name, portal)) as lease:
+            with portal.wrap_async_context_manager(self.lease_async(selector, lease_name, duration, portal)) as lease:
                 yield lease
 
     def get_exporter(self, name: str):
@@ -79,10 +84,6 @@ def list_exporters(
         with start_blocking_portal() as portal:
             return portal.call(self.list_exporters_async, page_size, page_token, filter)
 
-    def request_lease(self, selector: str):
-        with start_blocking_portal() as portal:
-            return portal.call(self.request_lease_async, selector, portal)
-
     def list_leases(self, filter: str):
         with start_blocking_portal() as portal:
             return portal.call(self.list_leases_async, filter)
@@ -106,14 +107,9 @@ def update_lease(self, name, duration: timedelta):
         with start_blocking_portal() as portal:
             return portal.call(self.update_lease_async, name, duration)
 
-    def release_lease(self, name):
-        with start_blocking_portal() as portal:
-            portal.call(self.release_lease_async, name)
-
     async def get_exporter_async(self, name: str):
         svc = ClientService(channel=await self.channel(), namespace=self.metadata.namespace)
-        with translate_grpc_exceptions():
-            return await svc.GetExporter(name=name)
+        return await svc.GetExporter(name=name)
 
     async def list_exporters_async(
         self,
@@ -122,70 +118,39 @@ async def list_exporters_async(
         filter: str | None = None,
     ):
         svc = ClientService(channel=await self.channel(), namespace=self.metadata.namespace)
-        with translate_grpc_exceptions():
-            return await svc.ListExporters(page_size=page_size, page_token=page_token, filter=filter)
+        return await svc.ListExporters(page_size=page_size, page_token=page_token, filter=filter)
 
     async def create_lease_async(
         self,
         selector: str,
         duration: timedelta,
     ):
         svc = ClientService(channel=await self.channel(), namespace=self.metadata.namespace)
-        with translate_grpc_exceptions():
-            return await svc.CreateLease(
-                selector=selector,
-                duration=duration,
-            )
+        return await svc.CreateLease(
+            selector=selector,
+            duration=duration,
+        )
 
     async def delete_lease_async(self, name: str):
         svc = ClientService(channel=await self.channel(), namespace=self.metadata.namespace)
-        with translate_grpc_exceptions():
-            await svc.DeleteLease(
-                name=name,
-            )
-
-    async def request_lease_async(
-        self,
-        selector: str,
-        portal: BlockingPortal,
-    ):
-        # dynamically import to avoid circular imports
-        from jumpstarter.client import Lease
-
-        lease = Lease(
-            channel=await self.channel(),
-            namespace=self.metadata.namespace,
-            name=None,
-            selector=selector,
-            portal=portal,
-            allow=self.drivers.allow,
-            unsafe=self.drivers.unsafe,
-            tls_config=self.tls,
-            grpc_options=self.grpcOptions,
+        await svc.DeleteLease(
+            name=name,
         )
-        with translate_grpc_exceptions():
-            return await lease.request_async()
 
     async def list_leases_async(self, filter: str):
         svc = ClientService(channel=await self.channel(), namespace=self.metadata.namespace)
-        with translate_grpc_exceptions():
-            return await svc.ListLeases(filter=filter)
+        return await svc.ListLeases(filter=filter)
 
     async def update_lease_async(self, name, duration: timedelta):
         svc = ClientService(channel=await self.channel(), namespace=self.metadata.namespace)
-        with translate_grpc_exceptions():
-            return await svc.UpdateLease(name=name, duration=duration)
-
-    async def release_lease_async(self, name):
-        svc = ClientService(channel=await self.channel(), namespace=self.metadata.namespace)
-        with translate_grpc_exceptions():
-            await svc.DeleteLease(name=name)
+        return await svc.UpdateLease(name=name, duration=duration)
 
     @asynccontextmanager
     async def lease_async(
         self,
         selector: str,
         lease_name: str | None,
+        duration: timedelta,
         portal: BlockingPortal,
     ):
         from jumpstarter.client import Lease
@@ -200,6 +165,7 @@ async def lease_async(
             namespace=self.metadata.namespace,
             name=lease_name,
             selector=selector,
+            duration=duration,
             portal=portal,
             allow=self.drivers.allow,
             unsafe=self.drivers.unsafe,
diff --git a/pyproject.toml b/pyproject.toml
@@ -71,6 +71,7 @@ locale = "en-us"
 
 [tool.typos.default.extend-words]
 ser = "ser"
+Pn = "Pn"
 
 [tool.coverage.run]
 omit = ["conftest.py", "test_*.py", "*_test.py", "*_pb2.py", "*_pb2_grpc.py"]
