beaker-notebook 2.0.0__py3-none-any.whl

beaker_kernel/__init__.py

@@ -0,0 +1,37 @@
+"""
+Compatibility shim: redirects any `beaker_kernel[.X]` import to `beaker_notebook[.X]`.
+"""
+import importlib
+import sys
+from importlib.abc import MetaPathFinder, Loader
+from importlib.machinery import ModuleSpec
+
+_OLD = "beaker_kernel"
+_NEW = "beaker_notebook"
+
+
+class _RenameLoader(Loader):
+    def create_module(self, spec):
+        target = _NEW + spec.name[len(_OLD):]
+        module = importlib.import_module(target)
+        sys.modules[spec.name] = module
+        return module
+
+    def exec_module(self, module):
+        pass
+
+
+class _RenameFinder(MetaPathFinder):
+    def find_spec(self, fullname, path=None, target=None):
+        if fullname == _OLD or fullname.startswith(_OLD + "."):
+            return ModuleSpec(fullname, _RenameLoader(), is_package=True)
+        return None
+
+
+# Install once.
+if not any(isinstance(f, _RenameFinder) for f in sys.meta_path):
+    sys.meta_path.insert(0, _RenameFinder())
+
+# Make `beaker_kernel` itself be `beaker_notebook`.
+import beaker_notebook as _bn
+sys.modules[__name__] = _bn

beaker_notebook/__init__.py

@@ -0,0 +1,172 @@
+from beaker_notebook.util.refactor import redirect_imports, ImportRedirectMap
+
+redirect_map: ImportRedirectMap = {
+    'beaker_kernel.service.admin_utils': {
+        'build_edges_map': 'beaker_kernel.lib.admin',
+        'build_proc_info': 'beaker_kernel.lib.admin',
+        'fetch_kernel_info': 'beaker_kernel.lib.admin',
+        'fetch_system_stats': 'beaker_kernel.lib.admin',
+    },
+    'beaker_kernel.service.api': {},
+    'beaker_kernel.service.api.handlers': {
+        'add_handler_prefix': 'beaker_kernel.app.api.handlers',
+        'find_api_handlers': 'beaker_kernel.app.api.handlers',
+        'register_api_handlers': 'beaker_kernel.app.api.handlers',
+    },
+    'beaker_kernel.service.api.integrations': {
+        'BeakerAPIMixin': 'beaker_kernel.app.api.integrations',
+        'IntegrationHandler': 'beaker_kernel.app.api.integrations',
+        'IntegrationResourceHandler': 'beaker_kernel.app.api.integrations',
+        'get_context': 'beaker_kernel.app.api.integrations',
+    },
+    'beaker_kernel.service.api.notebook': {
+        'NotebookHandler': 'beaker_kernel.app.api.notebook'
+    },
+    'beaker_kernel.service.auth': {
+        'BeakerUser': 'beaker_kernel.services.auth',
+        'current_request': 'beaker_kernel.services.auth',
+        'current_user': 'beaker_kernel.services.auth',
+        'BeakerAuthorizer': 'beaker_kernel.services.auth',
+        'BeakerIdentityProvider': 'beaker_kernel.services.auth',
+        'BeakerPermission': 'beaker_kernel.services.auth',
+        'BeakerRole': 'beaker_kernel.services.auth',
+        'RoleBasedUser': 'beaker_kernel.services.auth',
+    },
+    'beaker_kernel.service.auth.dummy': {
+        'DummyAuthorizer': 'beaker_kernel.services.auth.dummy',
+        'DummyIdentityProvider': 'beaker_kernel.services.auth.dummy',
+    },
+    'beaker_kernel.service.auth.notebook': {
+        'NotebookAuthorizer': 'beaker_kernel.services.auth.notebook',
+        'NotebookIdentityProvider': 'beaker_kernel.services.auth.notebook',
+    },
+    'beaker_kernel.service.base': {
+        'BaseBeakerApp': 'beaker_kernel.app.base',
+        'BeakerKernelManager': 'beaker_kernel.services.kernel.manager',
+        'BeakerKernelMappingManager': 'beaker_kernel.services.kernel.mappingmanager',
+        'BeakerKernelSpecManager': 'beaker_kernel.services.kernel.spec',
+        'BeakerSessionManager': 'beaker_kernel.services.session',
+    },
+    'beaker_kernel.service.dev': {
+        'BeakerWatchDog': 'beaker_kernel.app.dev_app',
+        'DevBeakerJupyterApp': 'beaker_kernel.app.dev_app',
+        'create_observer': 'beaker_kernel.app.dev_app',
+    },
+    'beaker_kernel.service.handlers': {
+        'ConfigController': 'beaker_kernel.app.handlers',
+        'ConfigHandler': 'beaker_kernel.app.handlers',
+        'ContextHandler': 'beaker_kernel.app.handlers',
+        'ExportAsHandler': 'beaker_kernel.app.handlers',
+        'PageHandler': 'beaker_kernel.app.handlers',
+        'StatsHandler': 'beaker_kernel.app.handlers',
+        'register_handlers': 'beaker_kernel.app.handlers',
+        'request_log_handler': 'beaker_kernel.app.handlers',
+        'sanitize_env': 'beaker_kernel.app.handlers',
+    },
+    'beaker_kernel.service.notebook': {
+        'BeakerNotebookApp': 'beaker_kernel.app.notebook_app'
+    },
+    'beaker_kernel.service.server': {
+        'BeakerServerApp': 'beaker_kernel.app.server_app'
+    },
+    'beaker_kernel.service.storage': {},
+    'beaker_kernel.service.storage.base': {
+        'BaseBeakerContentsManager': 'beaker_kernel.services.storage.base',
+        'BeakerLocalContentsHandler': 'beaker_kernel.services.storage.base',
+        'BeakerLocalContentsManager': 'beaker_kernel.services.storage.base',
+        'BeakerStorageManager': 'beaker_kernel.services.storage.base',
+        'with_hidden_files': 'beaker_kernel.services.storage.base',
+    },
+    'beaker_kernel.service.storage.notebook': {
+        'BaseNotebookManager': 'beaker_kernel.services.storage.notebook',
+        'BrowserLocalDataNotebookManager': 'beaker_kernel.services.storage.notebook',
+        'FileNotebookManager': 'beaker_kernel.services.storage.notebook',
+        'NotebookInfo': 'beaker_kernel.services.storage.notebook',
+        'with_hidden_files': 'beaker_kernel.services.storage.notebook',
+    },
+    'beaker_notebook.service.admin_utils': {
+        'build_edges_map': 'beaker_notebook.lib.admin',
+        'build_proc_info': 'beaker_notebook.lib.admin',
+        'fetch_kernel_info': 'beaker_notebook.lib.admin',
+        'fetch_system_stats': 'beaker_notebook.lib.admin',
+    },
+    'beaker_notebook.service.api': {},
+    'beaker_notebook.service.api.handlers': {
+        'add_handler_prefix': 'beaker_notebook.app.api.handlers',
+        'find_api_handlers': 'beaker_notebook.app.api.handlers',
+        'register_api_handlers': 'beaker_notebook.app.api.handlers',
+    },
+    'beaker_notebook.service.api.integrations': {
+        'BeakerAPIMixin': 'beaker_notebook.app.api.integrations',
+        'IntegrationHandler': 'beaker_notebook.app.api.integrations',
+        'IntegrationResourceHandler': 'beaker_notebook.app.api.integrations',
+        'get_context': 'beaker_notebook.app.api.integrations',
+    },
+    'beaker_notebook.service.api.notebook': {
+        'NotebookHandler': 'beaker_notebook.app.api.notebook'
+    },
+    'beaker_notebook.service.auth': {
+        'BeakerUser': 'beaker_notebook.services.auth',
+        'current_request': 'beaker_notebook.services.auth',
+        'current_user': 'beaker_notebook.services.auth',
+        'BeakerAuthorizer': 'beaker_notebook.services.auth',
+        'BeakerIdentityProvider': 'beaker_notebook.services.auth',
+        'BeakerPermission': 'beaker_notebook.services.auth',
+        'BeakerRole': 'beaker_notebook.services.auth',
+        'RoleBasedUser': 'beaker_notebook.services.auth',
+    },
+    'beaker_notebook.service.auth.dummy': {
+        'DummyAuthorizer': 'beaker_notebook.services.auth.dummy',
+        'DummyIdentityProvider': 'beaker_notebook.services.auth.dummy',
+    },
+    'beaker_notebook.service.auth.notebook': {
+        'NotebookAuthorizer': 'beaker_notebook.services.auth.notebook',
+        'NotebookIdentityProvider': 'beaker_notebook.services.auth.notebook',
+    },
+    'beaker_notebook.service.base': {
+        'BaseBeakerApp': 'beaker_notebook.app.base',
+        'BeakerKernelManager': 'beaker_notebook.services.kernel.manager',
+        'BeakerKernelMappingManager': 'beaker_notebook.services.kernel.mappingmanager',
+        'BeakerKernelSpecManager': 'beaker_notebook.services.kernel.spec',
+        'BeakerSessionManager': 'beaker_notebook.services.session',
+    },
+    'beaker_notebook.service.dev': {
+        'BeakerWatchDog': 'beaker_notebook.app.dev_app',
+        'DevBeakerJupyterApp': 'beaker_notebook.app.dev_app',
+        'create_observer': 'beaker_notebook.app.dev_app',
+    },
+    'beaker_notebook.service.handlers': {
+        'ConfigController': 'beaker_notebook.app.handlers',
+        'ConfigHandler': 'beaker_notebook.app.handlers',
+        'ContextHandler': 'beaker_notebook.app.handlers',
+        'ExportAsHandler': 'beaker_notebook.app.handlers',
+        'PageHandler': 'beaker_notebook.app.handlers',
+        'StatsHandler': 'beaker_notebook.app.handlers',
+        'register_handlers': 'beaker_notebook.app.handlers',
+        'request_log_handler': 'beaker_notebook.app.handlers',
+        'sanitize_env': 'beaker_notebook.app.handlers',
+    },
+    'beaker_notebook.service.notebook': {
+        'BeakerNotebookApp': 'beaker_notebook.app.notebook_app'
+    },
+    'beaker_notebook.service.server': {
+        'BeakerServerApp': 'beaker_notebook.app.server_app'
+    },
+    'beaker_notebook.service.storage': {},
+    'beaker_notebook.service.storage.base': {
+        'BaseBeakerContentsManager': 'beaker_notebook.services.storage.base',
+        'BeakerLocalContentsHandler': 'beaker_notebook.services.storage.base',
+        'BeakerLocalContentsManager': 'beaker_notebook.services.storage.base',
+        'BeakerStorageManager': 'beaker_notebook.services.storage.base',
+        'with_hidden_files': 'beaker_notebook.services.storage.base',
+    },
+    'beaker_notebook.service.storage.notebook': {
+        'BaseNotebookManager': 'beaker_notebook.services.storage.notebook',
+        'BrowserLocalDataNotebookManager': 'beaker_notebook.services.storage.notebook',
+        'FileNotebookManager': 'beaker_notebook.services.storage.notebook',
+        'NotebookInfo': 'beaker_notebook.services.storage.notebook',
+        'with_hidden_files': 'beaker_notebook.services.storage.notebook',
+    },
+}
+
+redirect_imports(redirect_map)

beaker_notebook/app/__init__.py

@@ -0,0 +1,13 @@
+# Application interface layer
+
+def __getattr__(name):
+    if name == "BaseBeakerApp":
+        from .base import BaseBeakerApp
+        return BaseBeakerApp
+    elif name == "BeakerServerApp":
+        from .server_app import BeakerServerApp
+        return BeakerServerApp
+    elif name == "BeakerNotebookApp":
+        from .notebook_app import BeakerNotebookApp
+        return BeakerNotebookApp
+    raise AttributeError(f"module '{__name__}' has no attribute '{name}'")

beaker_notebook/app/api/handlers.py

@@ -0,0 +1,84 @@
+import importlib
+import logging
+import os
+import sys
+import urllib.parse
+from typing import get_origin, get_args
+from dataclasses import is_dataclass, asdict
+from typing import TYPE_CHECKING, get_origin, get_args, GenericAlias, Union, Generic, Generator, Optional, Any
+
+
+from beaker_notebook.lib.autodiscovery import autodiscover
+from beaker_notebook.lib.app import BeakerApp
+from beaker_notebook.lib.context import BeakerContext
+from beaker_notebook.lib.subkernel import BeakerSubkernel
+from beaker_notebook.lib.agent_tasks import summarize
+from beaker_notebook.lib.config import config, locate_config, Config, Table, Choice, recursiveOptionalUpdate, reset_config
+from beaker_notebook.lib import admin
+
+if TYPE_CHECKING:
+    from beaker_notebook.app.base import BaseBeakerApp
+
+logger = logging.getLogger(__name__)
+
+PREFIX = '/beaker/'
+
+def find_api_handlers(base=None) -> Generator[tuple[str, Any, str], None, None]:
+    """Discover and yield API handlers from registered Beaker extensions.
+
+    This function uses the Beaker extension autodiscovery mechanism to find
+    all extensions that may provide API handlers. It then iterates through
+    each extension's handlers and yields them one by one.
+
+    Yields
+    ------
+    tuple[str, Any, str]
+        A tuple containing the URL pattern, handler class, and optional name
+        for each discovered API handler.
+    """
+    if base is None:
+        package = __package__
+        base_dir = os.path.dirname(str(sys.modules[package].__file__))
+    else:
+        package = None
+        try:
+            mod = importlib.import_module(base)
+            base_dir = os.path.dirname(str(mod.__file__))
+            package = mod.__package__
+        except (ImportError, TypeError):
+            logger.error(f"Could not import module {base} for API handler discovery")
+            return
+
+    if base_dir is None or package is None:
+        logger.error(f"Invalid base parameter: {type(base)}: {base}")
+        return
+
+    for f in os.listdir(base_dir):
+        if f.endswith('.py') and f != '__init__.py' and f != 'handlers.py':
+            s = f'{package}.{f[:-3]}'
+            try:
+                mod = importlib.import_module(s)
+                if "handlers" in dir(mod):
+                    logger.debug(f"Found handlers in {s}")
+                    for handlers in getattr(mod, "handlers"):
+                        yield handlers
+            except ImportError as e:
+                logger.warning(f"Failed to import API handler module '{s}': {e}")
+                logger.warning(f"Skipping API handler module '{s}' due to import error")
+                continue
+
+
+def add_handler_prefix(prefix: str, handler_tuple: tuple[str, Any, str]):
+    path, rest = handler_tuple[0], handler_tuple[1:]
+    if not prefix.endswith('/'):
+        prefix = prefix + '/'
+    if path.startswith('/'):
+        path = path.lstrip('/')
+    return (urllib.parse.urljoin(prefix, path), *rest)
+
+
+def register_api_handlers(app: "BaseBeakerApp"):
+    app.handlers.extend([
+        add_handler_prefix(PREFIX, handler)
+        for handler in find_api_handlers()
+    ])

beaker_notebook/app/api/integrations.py

@@ -0,0 +1,299 @@
+import json
+import logging
+import typing
+import asyncio
+from inspect import isawaitable, iscoroutinefunction
+from queue import Empty
+
+from jupyter_server.base.handlers import JupyterHandler
+from jupyter_server.services.kernels.kernelmanager import AsyncMappingKernelManager
+from jupyter_server.services.sessions.sessionmanager import SessionManager
+
+from beaker_notebook.lib.autodiscovery import autodiscover
+
+import tornado
+
+logger = logging.getLogger(__name__)
+
+from abc import ABC, abstractmethod
+
+from functools import lru_cache
+
+context_cache = None
+
+@lru_cache
+def get_context(context_slug: str):
+    global context_cache
+    if context_cache is None:
+        context_cache = autodiscover("contexts")
+    return context_cache.get(context_slug)
+
+
+# Resource Handlers
+
+# class ResourceHandler(ABC):
+#     KEY: str
+
+#     @abstractmethod
+#     async def get(self, integration_id, resource_id=None):
+#         pass
+
+#     @abstractmethod
+#     async def new(self, integration_id, payload):
+#         pass
+
+#     @abstractmethod
+#     async def replace(self, integration_id, resource_id, payload):
+#         pass
+
+#     @abstractmethod
+#     async def delete(self, integration_id, resource_id):
+#         pass
+
+
+# class FileResourceHandler(ResourceHandler):
+#     KEY="file"
+
+#     async def get(self, integration_id, resource_id=None):
+#         pass
+
+#     async def new(self, integration_id, payload):
+#         return await super().new(integration_id, payload)
+
+#     async def replace(self, integration_id, resource_id, payload):
+#         return await super().replace(integration_id, resource_id, payload)
+
+#     async def delete(self, integration_id, resource_id):
+#         return await super().delete(integration_id, resource_id)
+
+
+# class ExampleResourceHandler(ResourceHandler):
+#     KEY="example"
+
+#     async def get(self,  integration_id, resource_id=None):
+#         pass
+
+#     async def new(self, integration_id, payload):
+#         return await super().new(context_slug, integration_id, payload)
+
+#     async def replace(self,  integration_id, resource_id, payload):
+#         return await super().replace(context_slug, integration_id, resource_id, payload)
+
+#     async def delete(self, integration_id, resource_id):
+#         return await super().delete(context_slug, integration_id, resource_id)
+
+
+# Resources: list[ResourceHandler] = [
+#     FileResourceHandler,
+#     ExampleResourceHandler,
+# ]
+
+# ResourceMap = {
+#     cls.KEY: cls for cls in Resources
+# }
+
+
+class BeakerAPIMixin:
+    session_manager: SessionManager
+    kernel_manager: AsyncMappingKernelManager
+
+    def stringify_serialization(self, obj):
+        return json.loads(json.dumps(obj, default=str))
+
+    async def call_in_context(
+        self,
+        session_id: str | None,
+        target: str,
+        function: str,
+        args: typing.Optional[list] = None,
+        kwargs: typing.Optional[dict] = None,
+    ):
+        if args is None:
+            args = []
+        if kwargs is None:
+            kwargs = {}
+
+        kernel = None
+        retries = 0
+        max_retries = 5
+        interval = 3
+        while True:
+            try:
+                # TODO: session manager ready future
+                # TODO: ensure it's defined
+                session = await self.session_manager.get_session(name=session_id)
+                break
+            except Exception as err:
+                logger.warning(f"Failed to get session, retrying: {retries}/{max_retries} {err}")
+                await asyncio.sleep(interval)
+                retries += 1
+                if retries >= max_retries:
+                    logger.error(f"Failed after {max_retries} retries. Giving up")
+                    return None
+
+        if session is not None:
+            kernel = self.kernel_manager.get_kernel(session["kernel"]["id"])
+        if not kernel:
+            return None
+
+        client = kernel.client()
+        try:
+            content = {
+                "target": target,
+                "function": function,
+                "args": args,
+                "kwargs": kwargs
+            }
+            client.session.send(
+                stream=client.shell_channel.socket,
+                msg_or_type="call_in_context_request",
+                content=content,
+                track=True,
+                metadata=None,
+            )
+            if iscoroutinefunction(client.get_shell_msg):
+                result = await client.get_shell_msg(timeout=30) # timeout is in seconds, not milliseconds
+            else:
+                result = client.get_shell_msg(timeout=30)
+        except Empty as err:
+            raise TimeoutError(err)
+        finally:
+            client.stop_channels()
+        return result["content"]["return"]
+
+# Integration Handler
+
+class IntegrationHandler(BeakerAPIMixin, JupyterHandler):
+    """
+    Handles fetching and adding integrations.
+    """
+    def set_default_headers(self):
+        self.set_header("Content-Type", "application/json")
+
+    async def head(self, session_id=None, integration_id=None):
+        integrations = await self.call_in_context(
+            session_id=session_id,
+            target="context",
+            function="list_integrations"
+        ) or {}
+        self.write({"integrations": list(integrations.keys())})
+
+    async def get(self, session_id=None, integration_id=None):
+        if integration_id is None:
+            integrations = await self.call_in_context(
+                session_id=session_id,
+                target="context",
+                function="list_integrations"
+            ) or {}
+            self.write({"integrations": self.stringify_serialization(integrations)})
+        else:
+            integration = await self.call_in_context(
+                session_id=session_id,
+                target=f"integration:{integration_id}",
+                function="get_integration",
+                kwargs={"integration_id": integration_id}
+            )
+            self.write(self.stringify_serialization(integration))
+
+    async def post(self, session_id=None, integration_id=None):
+        body = tornado.escape.json_decode(self.request.body)
+        try:
+            if integration_id is not None:
+                body["integration_id"] = integration_id
+                result = await self.call_in_context(
+                    session_id=session_id,
+                    target=f"integration:{integration_id}",
+                    function="update_integration",
+                    kwargs=body
+                ) or {}
+            else:
+                provider_id = body.pop("provider")
+                result = await self.call_in_context(
+                    session_id=session_id,
+                    target=f"provider:{provider_id}",
+                    function="add_integration",
+                    kwargs=body
+                ) or {}
+            self.write(self.stringify_serialization(result))
+        except TimeoutError as e:
+            raise tornado.web.HTTPError(status_code=504, log_message="Kernel took too long to repond to request")
+        except Exception as e:
+            raise tornado.web.HTTPError(status_code=500, log_message=str(e))
+
+
+class IntegrationResourceHandler(BeakerAPIMixin, JupyterHandler):
+    """
+    Handles fetching and adding resources belonging to an integration.
+    """
+    def set_default_headers(self):
+        self.set_header("Content-Type", "application/json")
+
+    async def head(self, session_id=None, integration_id=None, resource_type=None, resource_id=None):
+        pass
+
+    async def get(self, session_id=None, integration_id=None, resource_type=None, resource_id=None):
+        if resource_id is None:
+            base_kwargs = {"integration_id": integration_id}
+            result = await self.call_in_context(
+                session_id=session_id,
+                target=f"integration:{integration_id}",
+                function="list_resources",
+                kwargs=base_kwargs | (
+                    {"resource_type": resource_type}
+                    if resource_type is not None and resource_type != "all"
+                    else {}
+                )
+            )
+            self.write({"resources": self.stringify_serialization(result)})
+        else:
+            try:
+                result = await self.call_in_context(
+                    session_id=session_id,
+                    target=f"integration:{integration_id}",
+                    function="get_resource",
+                    kwargs={"integration_id": integration_id, "resource_id": resource_id}
+                )
+                self.write(self.stringify_serialization(result))
+            except TimeoutError as e:
+                raise tornado.web.HTTPError(status_code=504, log_message="Kernel took too long to repond to request")
+            except Exception as e:
+                raise tornado.web.HTTPError(status_code=500, log_message=str(e))
+
+    async def post(self, session_id=None, integration_id=None, resource_type=None, resource_id=None):
+        function = "add_resource" if resource_id is None else "update_resource"
+        kwargs = {"integration_id": integration_id}
+        if function == "update_resource":
+            kwargs["resource_id"] = resource_id
+        kwargs.update(tornado.escape.json_decode(self.request.body))
+        try:
+            result = await self.call_in_context(
+                session_id=session_id,
+                target=f"integration:{integration_id}",
+                function=function,
+                kwargs=kwargs
+            )
+            self.write(self.stringify_serialization(result))
+        except TimeoutError as e:
+            raise tornado.web.HTTPError(status_code=504, log_message="Kernel took too long to repond to request")
+        except Exception as e:
+            raise tornado.web.HTTPError(status_code=500, log_message=str(e))
+
+
+    async def delete(self, session_id=None, integration_id=None, resource_type=None, resource_id=None):
+        try:
+            await self.call_in_context(
+                session_id=session_id,
+                target=f"integration:{integration_id}",
+                function="remove_resource",
+                kwargs={"resource_id": resource_id, "integration_id": integration_id}
+            )
+            self.write({})
+        except TimeoutError as e:
+            raise tornado.web.HTTPError(status_code=504, log_message="Kernel took too long to repond to request")
+        except Exception as e:
+            raise tornado.web.HTTPError(status_code=500, log_message=str(e))
+
+handlers = [
+    (r'integrations/(?P<session_id>[\w\d-]+)/?(?P<integration_id>[\w\d-]+)?', IntegrationHandler),
+    (r'integrations/(?P<session_id>[\w\d-]+)/(?P<integration_id>[\w\d-]+)/(?P<resource_type>\w+)/?(?P<resource_id>[\w\d-]+)?', IntegrationResourceHandler),
+]