singlestoredb 1.16.1__py3-none-any.whl

singlestoredb/apps/_process.py

@@ -0,0 +1,32 @@
+import os
+import signal
+import typing
+if typing.TYPE_CHECKING:
+    from psutil import Process
+
+
+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/apps/_python_udfs.py

@@ -0,0 +1,100 @@
+import asyncio
+import os
+import typing
+
+from ..functions.ext.asgi import Application
+from ._config import AppConfig
+from ._connection_info import UdfConnectionInfo
+from ._process import kill_process_by_port
+
+if typing.TYPE_CHECKING:
+    from ._uvicorn_util import AwaitableUvicornServer
+
+# Keep track of currently running server
+_running_server: 'typing.Optional[AwaitableUvicornServer]' = None
+
+# Maximum number of UDFs allowed
+MAX_UDFS_LIMIT = 10
+
+
+async def run_udf_app(
+    log_level: str = 'error',
+    kill_existing_app_server: bool = True,
+) -> UdfConnectionInfo:
+    global _running_server
+    from ._uvicorn_util import AwaitableUvicornServer
+
+    try:
+        import uvicorn
+    except ImportError:
+        raise ImportError('package uvicorn is required to run python udfs')
+
+    app_config = AppConfig.from_env()
+
+    if kill_existing_app_server:
+        # Shutdown the server gracefully if it was started by us.
+        # Since the uvicorn server doesn't start a new subprocess
+        # killing the process would result in kernel dying.
+        if _running_server is not None:
+            await _running_server.shutdown()
+            _running_server = None
+
+        # Kill if any other process is occupying the port
+        kill_process_by_port(app_config.listen_port)
+
+    base_url = generate_base_url(app_config)
+
+    udf_suffix = ''
+    if app_config.running_interactively:
+        udf_suffix = '_test'
+    app = Application(
+        url=base_url,
+        app_mode='managed',
+        name_suffix=udf_suffix,
+        log_level=log_level,
+    )
+
+    if not app.endpoints:
+        raise ValueError('You must define at least one function.')
+    if len(app.endpoints) > MAX_UDFS_LIMIT:
+        raise ValueError(
+            f'You can only define a maximum of {MAX_UDFS_LIMIT} functions.',
+        )
+
+    config = uvicorn.Config(
+        app,
+        host='0.0.0.0',
+        port=app_config.listen_port,
+        log_config=app.get_uvicorn_log_config(),
+    )
+
+    # Register the functions only if the app is running interactively.
+    if app_config.running_interactively:
+        app.register_functions(replace=True)
+
+    _running_server = AwaitableUvicornServer(config)
+    asyncio.create_task(_running_server.serve())
+    await _running_server.wait_for_startup()
+
+    print(f'Python UDF registered at {base_url}')
+
+    return UdfConnectionInfo(base_url, app.get_function_info())
+
+
+def generate_base_url(app_config: AppConfig) -> str:
+    if not app_config.is_gateway_enabled:
+        raise RuntimeError('Python UDFs are not available if Nova Gateway is not enabled')
+
+    if not app_config.running_interactively:
+        return app_config.base_url
+
+    # generate python udf endpoint for interactive notebooks
+    gateway_url = os.environ.get('SINGLESTOREDB_NOVA_GATEWAY_ENDPOINT')
+    if app_config.is_local_dev:
+        gateway_url = os.environ.get('SINGLESTOREDB_NOVA_GATEWAY_DEV_ENDPOINT')
+        if gateway_url is None:
+            raise RuntimeError(
+                'Missing SINGLESTOREDB_NOVA_GATEWAY_DEV_ENDPOINT environment variable.',
+            )
+
+    return f'{gateway_url}/pythonudfs/{app_config.notebook_server_id}/interactive/'

singlestoredb/apps/_stdout_supress.py

@@ -0,0 +1,30 @@
+import io
+import sys
+from typing import Optional
+
+
+class StdoutSuppressor:
+    """
+    Supresses the stdout for code executed within the context.
+    This should not be used for asynchronous or threaded executions.
+
+    ```py
+    with StdoutSupressor():
+        print("This won't be printed")
+    ```
+
+    """
+
+    def __enter__(self) -> None:
+        self.stdout = sys.stdout
+        self.buffer = io.StringIO()
+        sys.stdout = self.buffer
+
+    def __exit__(
+        self,
+        exc_type: Optional[object],
+        exc_value: Optional[Exception],
+        exc_traceback: Optional[str],
+    ) -> None:
+        del self.buffer
+        sys.stdout = self.stdout

singlestoredb/apps/_uvicorn_util.py

@@ -0,0 +1,36 @@
+import asyncio
+import socket
+from typing import List
+from typing import Optional
+try:
+    import uvicorn
+except ImportError:
+    raise ImportError('package uvicorn is required')
+
+
+class AwaitableUvicornServer(uvicorn.Server):
+    """
+    Adds `wait_for_startup` method.
+    The function (asynchronously) 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: Optional[List[socket.socket]] = 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
+
+    async def shutdown(self, sockets: Optional[list[socket.socket]] = None) -> None:
+        if self.started:
+            await super().shutdown(sockets)

singlestoredb/auth.py

@@ -0,0 +1,245 @@
+#!/usr/bin/env python
+import datetime
+from typing import Any
+from typing import List
+from typing import Optional
+from typing import Union
+
+import jwt
+
+
+# Credential types
+PASSWORD = 'password'
+JWT = 'jwt'
+BROWSER_SSO = 'browser_sso'
+
+# Single Sign-On URL
+SSO_URL = 'https://portal.singlestore.com/engine-sso'
+
+
+class JSONWebToken(object):
+    """Container for JWT information."""
+
+    def __init__(
+        self, token: str, expires: datetime.datetime,
+        email: str, username: str, url: str = SSO_URL,
+        clusters: Optional[Union[str, List[str]]] = None,
+        databases: Optional[Union[str, List[str]]] = None,
+        timeout: int = 60,
+    ):
+        self.token = token
+        self.expires = expires
+        self.email = email
+        self.username = username
+        self.model_version_number = 1
+
+        # Attributes needed for refreshing tokens
+        self.url = url
+        self.clusters = clusters
+        self.databases = databases
+        self.timeout = timeout
+
+    @classmethod
+    def from_token(cls, token: bytes, verify_signature: bool = False) -> 'JSONWebToken':
+        """Validate the contents of the JWT."""
+        info = jwt.decode(token, options={'verify_signature': verify_signature})
+
+        if not info.get('sub', None) and not info.get('username', None):
+            raise ValueError("Missing 'sub' and 'username' in claims")
+        if not info.get('email', None):
+            raise ValueError("Missing 'email' in claims")
+        if not info.get('exp', None):
+            raise ValueError("Missing 'exp' in claims")
+        try:
+            expires = datetime.datetime.fromtimestamp(info['exp'], datetime.timezone.utc)
+        except Exception as exc:
+            raise ValueError("Invalid 'exp' in claims: {}".format(str(exc)))
+
+        username = info.get('username', info.get('sub', None))
+        email = info['email']
+
+        return cls(token.decode('utf-8'), expires=expires, email=email, username=username)
+
+    def __str__(self) -> str:
+        return self.token
+
+    def __repr__(self) -> str:
+        return repr(self.token)
+
+    @property
+    def is_expired(self) -> bool:
+        """Determine if the token has expired."""
+        return self.expires >= datetime.datetime.now()
+
+    def refresh(self, force: bool = False) -> bool:
+        """
+        Refresh the token as needed.
+
+        Parameters
+        ----------
+        force : bool, optional
+           Should a new token be generated even if the existing
+           one has not expired yet?
+
+        Returns
+        -------
+        bool : Indicating whether the token was refreshed or not
+
+        """
+        if force or self.is_expired:
+            out = get_jwt(
+                self.email, url=self.url, clusters=self.clusters,
+                databases=self.databases, timeout=self.timeout,
+            )
+            self.token = out.token
+            self.expires = out.expires
+            return True
+        return False
+
+
+def _listify(s: Optional[Union[str, List[str]]]) -> Optional[str]:
+    """Return a list of strings in a comma-separated string."""
+    if s is None:
+        return None
+    if not isinstance(s, str):
+        return ','.join(s)
+    return s
+
+
+def get_jwt(
+    email: str, url: str = SSO_URL,
+    clusters: Optional[Union[str, List[str]]] = None,
+    databases: Optional[Union[str, List[str]]] = None,
+    timeout: int = 60, browser: Optional[Union[str, List[str]]] = None,
+) -> JSONWebToken:
+    """
+    Retrieve a JWT token from the SingleStoreDB single-sign-on URL.
+
+    Parameters
+    ----------
+    email : str
+        EMail of the database user
+    url : str, optional
+        The URL of the single-sign-on token generator
+    clusters : str or list[str], optional
+        The name of the cluster being connected to
+    databases : str or list[str], optional
+        The name of the database being connected to
+    timeout : int, optional
+        Number of seconds to wait before timing out the authentication request
+    browser : str or list[str], optional
+        Browser to use instead of the default. This value can be any of the
+        names specified in Python's `webbrowser` module. This includes
+        'google-chrome', 'chrome', 'chromium', 'chromium-browser', 'firefox',
+        etc. Note that at the time of this writing, Safari was not
+        compatible. If a list of names is specified, each one tried until
+        a working browser is located.
+
+    Returns
+    -------
+    JSONWebToken
+
+    """
+    import platform
+    import webbrowser
+    import time
+    import threading
+    import urllib
+    from http.server import BaseHTTPRequestHandler, HTTPServer
+
+    from .config import get_option
+
+    token = []
+    error = []
+
+    class AuthServer(BaseHTTPRequestHandler):
+
+        def log_message(self, format: str, *args: Any) -> None:
+            return
+
+        def do_POST(self) -> None:
+            content_len = int(self.headers.get('Content-Length', 0))
+            post_body = self.rfile.read(content_len)
+
+            try:
+                out = JSONWebToken.from_token(post_body)
+            except Exception as exc:
+                self.send_response(400, exc.args[0])
+                self.send_header('Content-Type', 'text/plain')
+                self.end_headers()
+                error.append(exc)
+                return
+
+            token.append(out)
+
+            self.send_response(204)
+            self.send_header('Access-Control-Allow-Origin', '*')
+            self.send_header('Content-Type', 'text/plain')
+            self.end_headers()
+
+    server = None
+
+    try:
+        server = HTTPServer(('127.0.0.1', 0), AuthServer)
+        threading.Thread(target=server.serve_forever).start()
+
+        host = server.server_address[0]
+        if isinstance(host, bytes):
+            host = host.decode('utf-8')
+
+        query = urllib.parse.urlencode({
+            k: v for k, v in dict(
+                email=email,
+                returnTo=f'http://{host}:{server.server_address[1]}',
+                db=_listify(databases),
+                cluster=_listify(clusters),
+            ).items() if v is not None
+        })
+
+        if browser is None:
+            browser = get_option('sso_browser')
+
+        # On Mac, always specify a list of browsers to check because Safari
+        # is not compatible.
+        if browser is None and platform.platform().lower().startswith('mac'):
+            browser = [
+                'chrome', 'google-chrome', 'chromium',
+                'chromium-browser', 'firefox',
+            ]
+
+        if browser and isinstance(browser, str):
+            browser = [browser]
+
+        if browser:
+            exc: Optional[Exception] = None
+            for item in browser:
+                try:
+                    webbrowser.get(item).open(f'{url}?{query}')
+                    break
+                except webbrowser.Error as wexc:
+                    exc = wexc
+                    pass
+            if exc is not None:
+                raise RuntimeError(
+                    'Could not find compatible web browser for accessing JWT',
+                )
+        else:
+            webbrowser.open(f'{url}?{query}')
+
+        for i in range(timeout * 2):
+            if error:
+                raise error[0]
+            if token:
+                out = token[0]
+                out.url = url
+                out.clusters = clusters
+                out.databases = databases
+                out.timeout = timeout
+                return out
+            time.sleep(0.5)
+
+    finally:
+        if server is not None:
+            server.shutdown()
+
+    raise RuntimeError('Timeout waiting for token')