Merge pull request #763 from bennyz/custom-lease-name

cli: add support for requesting custom lease name

Author: bennyz

packages/jumpstarter-cli/jumpstarter_cli/create.py

@@ -22,9 +22,17 @@ def create():
 @opt_selector
 @opt_duration_partial(required=True)
 @opt_begin_time
+@click.option(
+    "--lease-id",
+    type=str,
+    default=None,
+    help="Optional lease ID to request (if not provided, server will generate one)",
+)
 @opt_output_all
 @handle_exceptions_with_reauthentication(relogin_client)
-def create_lease(config, selector: str, duration: timedelta, begin_time: datetime | None, output: OutputType):
+def create_lease(
+    config, selector: str, duration: timedelta, begin_time: datetime | None, lease_id: str | None, output: OutputType
+):
     """
     Create a lease
 
@@ -48,8 +56,14 @@ def create_lease(config, selector: str, duration: timedelta, begin_time: datetim
         $$ exit
         $ jmp delete lease "${JMP_LEASE}"
 
+    You can also specify a unique custom lease ID:
+
+    .. code-block:: bash
+
+        $ jmp create lease -l foo=bar --duration 1d --lease-id my-custom-lease-id
+
     """
 
-    lease = config.create_lease(selector=selector, duration=duration, begin_time=begin_time)
+    lease = config.create_lease(selector=selector, duration=duration, begin_time=begin_time, lease_id=lease_id)
 
     model_print(lease, output)

packages/jumpstarter/jumpstarter/client/client.py

@@ -59,8 +59,8 @@ async def client_from_channel(
             portal=portal,
             stack=stack.enter_context(ExitStack()),
             children={reports[k].labels["jumpstarter.dev/name"]: clients[k] for k in topo[index]},
-            description=getattr(report, 'description', None) or None,
-            methods_description=getattr(report, 'methods_description', {}) or {},
+            description=getattr(report, "description", None) or None,
+            methods_description=getattr(report, "methods_description", {}) or {},
         )
 
         clients[index] = client

packages/jumpstarter/jumpstarter/client/decorators.py

@@ -34,20 +34,21 @@ def on():
     :param kwargs: Keyword arguments passed to DriverClickGroup
     :return: Decorator that creates a DriverClickGroup
     """
+
     def decorator(f: Callable) -> DriverClickGroup:
         # Use function docstring if no help= provided
-        if 'help' not in kwargs or kwargs['help'] is None:
+        if "help" not in kwargs or kwargs["help"] is None:
             if f.__doc__:
-                kwargs['help'] = f.__doc__.strip()
+                kwargs["help"] = f.__doc__.strip()
 
         # Server description overrides Click defaults
-        if getattr(client, 'description', None):
-            kwargs['help'] = client.description
+        if getattr(client, "description", None):
+            kwargs["help"] = client.description
 
         group = DriverClickGroup(client, name=f.__name__, callback=f, **kwargs)
 
         # Transfer Click parameters attached by decorators like @click.option
-        group.params = getattr(f, '__click_params__', [])
+        group.params = getattr(f, "__click_params__", [])
 
         return group
 
@@ -74,8 +75,8 @@ def ssh(args):
     :return: click.command decorator
     """
     # Server description overrides Click's defaults (help= parameter or docstring)
-    if getattr(client, 'description', None):
-        kwargs['help'] = client.description
+    if getattr(client, "description", None):
+        kwargs["help"] = client.description
 
     return click.command(**kwargs)
 
@@ -89,13 +90,14 @@ def __init__(self, client: "DriverClient", *args: Any, **kwargs: Any) -> None:
 
     def command(self, *args: Any, **kwargs: Any) -> Callable:
         """Command decorator with server methods_description override support."""
+
         def decorator(f: Callable) -> click.Command:
-            name = kwargs.get('name')
+            name = kwargs.get("name")
             if not name:
-                name = f.__name__.lower().replace('_', '-')
+                name = f.__name__.lower().replace("_", "-")
 
             if name in self.client.methods_description:
-                kwargs['help'] = self.client.methods_description[name]
+                kwargs["help"] = self.client.methods_description[name]
 
             return super(DriverClickGroup, self).command(*args, **kwargs)(f)
 

packages/jumpstarter/jumpstarter/client/grpc.py

@@ -275,11 +275,7 @@ def model_dump_json(self, **kwargs):
         if not self.include_online:
             exclude_fields.add("online")
 
-        data = {
-            "exporters": [
-                exporter.model_dump(mode="json", exclude=exclude_fields) for exporter in self.exporters
-            ]
-        }
+        data = {"exporters": [exporter.model_dump(mode="json", exclude=exclude_fields) for exporter in self.exporters]}
         return json.dumps(data, **json_kwargs)
 
     def model_dump(self, **kwargs):
@@ -289,11 +285,8 @@ def model_dump(self, **kwargs):
         if not self.include_online:
             exclude_fields.add("online")
 
-        return {
-            "exporters": [
-                exporter.model_dump(mode="json", exclude=exclude_fields) for exporter in self.exporters
-            ]
-        }
+        return {"exporters": [exporter.model_dump(mode="json", exclude=exclude_fields) for exporter in self.exporters]}
+
 
 class LeaseList(BaseModel):
     leases: list[Lease]
@@ -390,6 +383,7 @@ async def CreateLease(
         selector: str,
         duration: timedelta,
         begin_time: datetime | None = None,
+        lease_id: str | None = None,
     ):
         duration_pb = duration_pb2.Duration()
         duration_pb.FromTimedelta(duration)
@@ -409,6 +403,7 @@ async def CreateLease(
                 client_pb2.CreateLeaseRequest(
                     parent="namespaces/{}".format(self.namespace),
                     lease=lease_pb,
+                    lease_id=lease_id or "",
                 )
             )
         return Lease.from_protobuf(lease)

packages/jumpstarter/jumpstarter/client/grpc_test.py

@@ -63,12 +63,7 @@ class TestAddExporterRow:
     def create_test_exporter(self, online=True, labels=None):
         if labels is None:
             labels = {"env": "test", "type": "device"}
-        return Exporter(
-            namespace="default",
-            name="test-exporter",
-            labels=labels,
-            online=online
-        )
+        return Exporter(namespace="default", name="test-exporter", labels=labels, online=online)
 
     def test_basic_row(self):
         table = Table()
@@ -119,11 +114,16 @@ def test_row_with_all_options(self):
 
 
 class TestExporterList:
-    def create_test_lease(self, client="test-client", status="Active",
-                          effective_begin_time=datetime(2023, 1, 1, 10, 0, 0),
-                          effective_duration=timedelta(hours=1),
-                          begin_time=None, duration=timedelta(hours=1),
-                          effective_end_time=None):
+    def create_test_lease(
+        self,
+        client="test-client",
+        status="Active",
+        effective_begin_time=datetime(2023, 1, 1, 10, 0, 0),
+        effective_duration=timedelta(hours=1),
+        begin_time=None,
+        duration=timedelta(hours=1),
+        effective_end_time=None,
+    ):
         lease = Mock(spec=Lease)
         lease.client = client
         lease.get_status.return_value = status
@@ -135,12 +135,7 @@ def create_test_lease(self, client="test-client", status="Active",
         return lease
 
     def test_exporter_without_lease(self):
-        exporter = Exporter(
-            namespace="default",
-            name="test-exporter",
-            labels={"type": "device"},
-            online=True
-        )
+        exporter = Exporter(namespace="default", name="test-exporter", labels={"type": "device"}, online=True)
 
         table = Table()
         Exporter.rich_add_columns(table)
@@ -152,11 +147,7 @@ def test_exporter_without_lease(self):
     def test_exporter_with_lease_no_display(self):
         lease = self.create_test_lease()
         exporter = Exporter(
-            namespace="default",
-            name="test-exporter",
-            labels={"type": "device"},
-            online=True,
-            lease=lease
+            namespace="default", name="test-exporter", labels={"type": "device"}, online=True, lease=lease
         )
 
         table = Table()
@@ -170,11 +161,7 @@ def test_exporter_with_lease_no_display(self):
     def test_exporter_with_lease_display(self):
         lease = self.create_test_lease()
         exporter = Exporter(
-            namespace="default",
-            name="test-exporter",
-            labels={"type": "device"},
-            online=True,
-            lease=lease
+            namespace="default", name="test-exporter", labels={"type": "device"}, online=True, lease=lease
         )
 
         table = Table()
@@ -198,12 +185,7 @@ def test_exporter_with_lease_display(self):
         assert "2023-01-01 11:00:00" in output  # Expected release: begin_time (10:00:00) + duration (1h)
 
     def test_exporter_without_lease_but_show_leases(self):
-        exporter = Exporter(
-            namespace="default",
-            name="test-exporter",
-            labels={"type": "device"},
-            online=True
-        )
+        exporter = Exporter(namespace="default", name="test-exporter", labels={"type": "device"}, online=True)
 
         table = Table()
         options = WithOptions(show_leases=True)
@@ -228,19 +210,11 @@ def test_exporter_without_lease_but_show_leases(self):
     def test_exporter_online_status_display(self):
         """Test that online status icons are correctly displayed"""
         # Test online exporter
-        exporter_online = Exporter(
-            namespace="default",
-            name="online-exporter",
-            labels={"type": "device"},
-            online=True
-        )
+        exporter_online = Exporter(namespace="default", name="online-exporter", labels={"type": "device"}, online=True)
 
         # Test offline exporter
         exporter_offline = Exporter(
-            namespace="default",
-            name="offline-exporter",
-            labels={"type": "device"},
-            online=False
+            namespace="default", name="offline-exporter", labels={"type": "device"}, online=False
         )
 
         # Test with online status display enabled
@@ -262,26 +236,22 @@ def test_exporter_online_status_display(self):
         assert "online-exporter" in output
         assert "offline-exporter" in output
         assert "yes" in output  # Should show "yes" for online
-        assert "no" in output   # Should show "no" for offline
+        assert "no" in output  # Should show "no" for offline
 
     def test_exporter_all_features_display(self):
         """Test all display features together: online status + lease info"""
         lease = self.create_test_lease(client="full-test-client", status="Active")
 
         # Create exporters with different combinations of online/lease status
         exporter_online_with_lease = Exporter(
-            namespace="default",
-            name="online-with-lease",
-            labels={"env": "prod"},
-            online=True,
-            lease=lease
+            namespace="default", name="online-with-lease", labels={"env": "prod"}, online=True, lease=lease
         )
 
         exporter_offline_no_lease = Exporter(
             namespace="default",
             name="offline-no-lease",
             labels={"env": "dev"},
-            online=False
+            online=False,
             # No lease
         )
 
@@ -306,7 +276,7 @@ def test_exporter_all_features_display(self):
         assert "env=prod" in output
         assert "env=dev" in output
         assert "yes" in output  # Online indicator
-        assert "no" in output   # Offline indicator
+        assert "no" in output  # Offline indicator
         assert "full-test-client" in output  # Lease client
         assert "Active" in output  # Lease status
         assert "Available" in output  # Available status for no lease
@@ -317,14 +287,10 @@ def test_exporter_lease_info_extraction(self):
         lease = self.create_test_lease(
             client="my-client",
             status="Expired",
-            effective_end_time=datetime(2023, 1, 1, 11, 0, 0)  # Ended after 1 hour
+            effective_end_time=datetime(2023, 1, 1, 11, 0, 0),  # Ended after 1 hour
         )
         exporter = Exporter(
-            namespace="default",
-            name="test-exporter",
-            labels={"type": "device"},
-            online=True,
-            lease=lease
+            namespace="default", name="test-exporter", labels={"type": "device"}, online=True, lease=lease
         )
 
         # Manually verify the lease data that would be extracted
@@ -359,7 +325,7 @@ def test_exporter_no_lease_info_extraction(self):
             namespace="default",
             name="test-exporter",
             labels={"type": "device"},
-            online=True
+            online=True,
             # No lease attached
         )
 
@@ -380,16 +346,12 @@ def test_exporter_scheduled_lease_expected_release(self):
             client="my-client",
             status="Scheduled",
             effective_begin_time=None,  # Not started yet
-            effective_duration=None,    # Not started yet
+            effective_duration=None,  # Not started yet
             begin_time=datetime(2023, 1, 1, 10, 0, 0),
-            duration=timedelta(hours=1)
+            duration=timedelta(hours=1),
         )
         exporter = Exporter(
-            namespace="default",
-            name="test-exporter",
-            labels={"type": "device"},
-            online=True,
-            lease=lease
+            namespace="default", name="test-exporter", labels={"type": "device"}, online=True, lease=lease
         )
 
         # Test the table display with scheduled lease
@@ -412,5 +374,3 @@ def test_exporter_scheduled_lease_expected_release(self):
         assert "my-client" in output
         assert "Scheduled" in output
         assert "2023-01-01 11:00:00" in output  # begin_time (10:00) + duration (1h)
-
-

packages/jumpstarter/jumpstarter/client/lease.py

@@ -72,6 +72,7 @@ async def _create(self):
                 await self.svc.CreateLease(
                     selector=self.selector,
                     duration=self.duration,
+                    lease_id=self.name,
                 )
             ).name
         logger.info("Acquiring lease %s for selector %s for duration %s", self.name, self.selector, self.duration)
@@ -388,9 +389,7 @@ def update_status(self, message: str, force: bool = False):
             # Throttle updates to at most every 5 minutes unless forced
             now = datetime.now()
             should_log = (
-                force
-                or self._last_log_time is None
-                or (now - self._last_log_time) >= self._log_throttle_interval
+                force or self._last_log_time is None or (now - self._last_log_time) >= self._log_throttle_interval
             )
 
             if should_log: