add notebook functions for starting dashboard and cloud functions

singlestoredb/notebook/__init__.py

@@ -2,6 +2,8 @@
 import os as _os
 import warnings as _warnings
 
+from ._apps import run_cloud_function  # noqa: F401
+from ._apps import run_dashboard  # noqa: F401
 from ._objects import organization  # noqa: F401
 from ._objects import secrets  # noqa: F401
 from ._objects import stage  # noqa: F401

singlestoredb/notebook/_apps.py

@@ -0,0 +1,109 @@
+import asyncio
+import os
+import signal
+import typing
+import urllib.parse
+
+from ._portal import portal
+if typing.TYPE_CHECKING:
+    from plotly.graph_objs import Figure
+    from fastapi import FastAPI
+    from psutil import Process
+
+
+def run_dashboard(
+    figure: 'Figure',
+    kill_existing_server: bool = True,
+    host: str = '0.0.0.0',
+    debug: bool = False,
+) -> None:
+    try:
+        import dash
+    except ImportError:
+        raise ImportError('package dash is required to run dashboards')
+
+    port = portal.app_listen_port
+    app_url = portal.app_url
+
+    if port is None or app_url is None:
+        raise RuntimeError(
+           'Portal not fully initialized. '
+           'Is the code running outside SingleStoreDB notebook environment?',
+        )
+
+    if kill_existing_server:
+        _kill_process_by_port(port)
+
+    base_path = urllib.parse.urlparse(app_url).path
+
+    app = dash.Dash(requests_pathname_prefix=base_path)
+    app.layout = dash.html.Div([
+        dash.dcc.Graph(figure=figure),
+    ])
+
+    app.run(host=host, debug=debug, port=str(port), jupyter_mode='external')
+
+    print(f"Dash app available at {app_url}")
+
+
+async def run_cloud_function(
+    app: 'FastAPI',
+    kill_existing_server: bool = True,
+    host: str = '0.0.0.0',
+    log_level: str = 'error',
+) -> None:
+    from ._uvicorn_util import AwaitableUvicornServer
+    try:
+        import uvicorn
+    except ImportError:
+        raise ImportError('package uvicorn is required to run cloud functions')
+
+    port = portal.app_listen_port
+    app_url = portal.app_url
+
+    if port is None or app_url is None:
+        raise RuntimeError(
+           'Portal not fully initialized. '
+           ' Is the code running outside SingleStoreDB notebook environment?',
+        )
+
+    if kill_existing_server:
+        _kill_process_by_port(port)
+
+    base_path = urllib.parse.urlparse(app_url).path
+    app.root_path = base_path
+
+    config = uvicorn.Config(app, host=host, port=port, log_level=log_level)
+    server = AwaitableUvicornServer(config)
+
+    asyncio.create_task(server.serve())
+    await server.wait_for_startup()
+
+    print(f"Cloud function available at {app_url}")
+
+
+def _kill_process_by_port(port: int) -> None:
+    existing_process = _find_process_by_port(port)
+    kernel_pid = os.getpid()
+    # Make sure we are not killing current kernel
+    if existing_process is not None and kernel_pid != existing_process.pid:
+        print(f"Killing process {existing_process.pid} which is using port {port}")
+        os.kill(existing_process.pid, signal.SIGKILL)
+
+
+def _find_process_by_port(port: int) -> 'Process | None':
+    try:
+        import psutil
+    except ImportError:
+        raise ImportError('package psutil is required')
+
+    for proc in psutil.process_iter(['pid']):
+        try:
+            connections = proc.connections()
+            for conn in connections:
+                if conn.laddr.port == port:
+                    return proc
+        except psutil.AccessDenied:
+            pass
+
+    return None

singlestoredb/notebook/_portal.py

@@ -304,6 +304,19 @@ def default_database(self, name: str) -> None:
             timeout_message='timeout waiting for database update',
         )
 
+    @property
+    def app_listen_port(self) -> Optional[int]:
+        """Port to use for running SingleStore App."""
+        port = os.environ.get('SINGLESTOREDB_APP_LISTEN_PORT')
+        if port is not None:
+            return int(port)
+        return None
+
+    @property
+    def app_url(self) -> Optional[str]:
+        """SingleStore App URL."""
+        return os.environ.get('SINGLESTOREDB_APP_URL')
+
     @property
     def version(self) -> Optional[str]:
         """Version."""

singlestoredb/notebook/_uvicorn_util.py

@@ -0,0 +1,30 @@
+import asyncio
+import socket
+try:
+    import uvicorn
+except ImportError:
+    raise ImportError('package uvicorn is required')
+
+
+class AwaitableUvicornServer(uvicorn.Server):
+    """
+    Adds `wait_for_startup` method.
+    The function (asynchornously) blocks until the server
+    starts listening or throws an error.
+    """
+
+    def __init__(self, config: 'uvicorn.Config') -> None:
+        super().__init__(config)
+        self._startup_future = asyncio.get_event_loop().create_future()
+
+    async def startup(self, sockets: list[socket.socket] | None = None) -> None:
+        try:
+            result = await super().startup(sockets)
+            self._startup_future.set_result(True)
+            return result
+        except Exception as error:
+            self._startup_future.set_exception(error)
+            raise error
+
+    async def wait_for_startup(self) -> None:
+        await self._startup_future

test-requirements.txt

@@ -1,4 +1,6 @@
 coverage
+dash
+fastapi
 pandas
 parameterized
 polars