From 42313add2e8d61220a83423881a0c33774284a7d Mon Sep 17 00:00:00 2001 From: saf Date: Wed, 6 Mar 2024 11:19:43 +0800 Subject: [PATCH] fix: check kernel id before delete and upgrade ws --- enterprise_gateway/services/kernels/handlers.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/enterprise_gateway/services/kernels/handlers.py b/enterprise_gateway/services/kernels/handlers.py index d5d2c590..59f85aad 100644 --- a/enterprise_gateway/services/kernels/handlers.py +++ b/enterprise_gateway/services/kernels/handlers.py @@ -146,6 +146,22 @@ def get(self, kernel_id: str): model = km.kernel_model(kernel_id) self.finish(json.dumps(model, default=date_default)) + @web.authenticated + async def delete(self, kernel_id): + """Remove a kernel.""" + self.kernel_manager.check_kernel_id(kernel_id=kernel_id) + await super().delete(kernel_id=kernel_id) + + +class ZMQChannelsHandler(jupyter_server_handlers.ZMQChannelsHandler): + """Extends the kernel websocket handler.""" + + async def get(self, kernel_id): + """Handle a get request for a kernel.""" + # Synchronize Kernel and check if it exists. + self.kernel_manager.check_kernel_id(kernel_id=kernel_id) + await super().get(kernel_id=kernel_id) + default_handlers: list[tuple] = [] for path, cls in jupyter_server_handlers.default_handlers: