singlestoredb 1.16.1__py3-none-any.whl

singlestoredb/__init__.py

@@ -0,0 +1,75 @@
+#!/usr/bin/env python
+"""
+SingleStoreDB module.
+
+Examples
+--------
+>>> import singlestoredb as s2
+>>> conn = s2.connect('user:password@host/dbname')
+>>> cur = conn.cursor()
+>>> cur.execute('select * from customers')
+>>> for row in cur:
+...     print(row)
+
+"""
+
+__version__ = '1.16.1'
+
+from typing import Any
+
+from .config import options, get_option, set_option, describe_option
+from .connection import connect, apilevel, threadsafety, paramstyle
+from .exceptions import (
+    Warning, Error, InterfaceError, DatabaseError, OperationalError,
+    IntegrityError, InternalError, ProgrammingError, NotSupportedError,
+    DataError, ManagementError,
+)
+from .management import (
+    manage_cluster, manage_workspaces, manage_files, manage_regions,
+)
+from .types import (
+    Date, Time, Timestamp, DateFromTicks, TimeFromTicks, TimestampFromTicks,
+    Binary, STRING, BINARY, NUMBER, DATETIME, ROWID,
+)
+# These are only loaded if the singlestore-vectorstore package is available
+try:
+    from .vectorstore import (
+        vector_db, IndexInterface, IndexList, IndexModel, MatchTypedDict,
+        Metric, IndexStatsTypedDict, NamespaceStatsTypedDict, Vector,
+        VectorDictMetadataValue, VectorMetadataTypedDict, VectorTuple,
+        VectorTupleWithMetadata, DeletionProtection, AndFilter, EqFilter,
+        ExactMatchFilter, FilterTypedDict, GteFilter, GtFilter, InFilter,
+        LteFilter, LtFilter, NeFilter, NinFilter, OrFilter, SimpleFilter,
+    )
+except (ImportError, ModuleNotFoundError):
+    pass
+
+
+#
+# This function is defined here to prevent the side-effect of
+# attempting to load the SQLAlchemy dialect in the core SDK.
+#
+def create_engine(*args: Any, **kwargs: Any) -> Any:
+    """
+    Create an SQLAlchemy engine for SingleStoreDB.
+
+    Parameters
+    ----------
+    **kwargs : Any
+        The parameters taken here are the same as for
+        `sqlalchemy.create_engine`. However, this function can be
+        called without any parameters in order to inherit parameters
+        set by environment variables or parameters set in by
+        options in Python code.
+
+    See Also
+    --------
+    `sqlalchemy.create_engine`
+
+    Returns
+    -------
+    SQLAlchemy engine
+
+    """
+    from .alchemy import create_engine
+    return create_engine(*args, **kwargs)

singlestoredb/ai/__init__.py

@@ -0,0 +1,2 @@
+from .chat import SingleStoreChatFactory  # noqa: F401
+from .embeddings import SingleStoreEmbeddingsFactory  # noqa: F401

singlestoredb/ai/chat.py

@@ -0,0 +1,139 @@
+import os
+from typing import Any
+from typing import Callable
+from typing import Optional
+from typing import Union
+
+import httpx
+
+from singlestoredb import manage_workspaces
+
+try:
+    from langchain_openai import ChatOpenAI
+except ImportError:
+    raise ImportError(
+        'Could not import langchain_openai python package. '
+        'Please install it with `pip install langchain_openai`.',
+    )
+
+try:
+    from langchain_aws import ChatBedrockConverse
+except ImportError:
+    raise ImportError(
+        'Could not import langchain-aws python package. '
+        'Please install it with `pip install langchain-aws`.',
+    )
+
+import boto3
+from botocore import UNSIGNED
+from botocore.config import Config
+
+
+def SingleStoreChatFactory(
+    model_name: str,
+    api_key: Optional[str] = None,
+    streaming: bool = True,
+    http_client: Optional[httpx.Client] = None,
+    obo_token_getter: Optional[Callable[[], Optional[str]]] = None,
+    **kwargs: Any,
+) -> Union[ChatOpenAI, ChatBedrockConverse]:
+    """Return a chat model instance (ChatOpenAI or ChatBedrockConverse).
+    """
+    inference_api_manager = (
+        manage_workspaces().organizations.current.inference_apis
+    )
+    info = inference_api_manager.get(model_name=model_name)
+    token_env = os.environ.get('SINGLESTOREDB_USER_TOKEN')
+    token = api_key if api_key is not None else token_env
+
+    if info.hosting_platform == 'Amazon':
+        # Instantiate Bedrock client
+        cfg_kwargs = {
+            'signature_version': UNSIGNED,
+            'retries': {'max_attempts': 1, 'mode': 'standard'},
+        }
+        # Extract timeouts from http_client if provided
+        t = http_client.timeout if http_client is not None else None
+        connect_timeout = None
+        read_timeout = None
+        if t is not None:
+            if isinstance(t, httpx.Timeout):
+                if t.connect is not None:
+                    connect_timeout = float(t.connect)
+                if t.read is not None:
+                    read_timeout = float(t.read)
+                if connect_timeout is None and read_timeout is not None:
+                    connect_timeout = read_timeout
+                if read_timeout is None and connect_timeout is not None:
+                    read_timeout = connect_timeout
+            elif isinstance(t, (int, float)):
+                connect_timeout = float(t)
+                read_timeout = float(t)
+        if read_timeout is not None:
+            cfg_kwargs['read_timeout'] = read_timeout
+        if connect_timeout is not None:
+            cfg_kwargs['connect_timeout'] = connect_timeout
+
+        cfg = Config(**cfg_kwargs)
+        client = boto3.client(
+            'bedrock-runtime',
+            endpoint_url=info.connection_url,
+            region_name='us-east-1',
+            aws_access_key_id='placeholder',
+            aws_secret_access_key='placeholder',
+            config=cfg,
+        )
+
+        def _inject_headers(request: Any, **_ignored: Any) -> None:
+            """Inject dynamic auth/OBO headers prior to Bedrock sending."""
+            if obo_token_getter is not None:
+                obo_val = obo_token_getter()
+                if obo_val:
+                    request.headers['X-S2-OBO'] = obo_val
+            if token:
+                request.headers['Authorization'] = f'Bearer {token}'
+            request.headers.pop('X-Amz-Date', None)
+            request.headers.pop('X-Amz-Security-Token', None)
+
+        emitter = client._endpoint._event_emitter
+        emitter.register_first(
+            'before-send.bedrock-runtime.Converse',
+            _inject_headers,
+        )
+        emitter.register_first(
+            'before-send.bedrock-runtime.ConverseStream',
+            _inject_headers,
+        )
+        emitter.register_first(
+            'before-send.bedrock-runtime.InvokeModel',
+            _inject_headers,
+        )
+        emitter.register_first(
+            'before-send.bedrock-runtime.InvokeModelWithResponseStream',
+            _inject_headers,
+        )
+
+        return ChatBedrockConverse(
+            model_id=model_name,
+            endpoint_url=info.connection_url,
+            region_name='us-east-1',
+            aws_access_key_id='placeholder',
+            aws_secret_access_key='placeholder',
+            disable_streaming=not streaming,
+            client=client,
+            **kwargs,
+        )
+
+    # OpenAI / Azure OpenAI path
+    openai_kwargs = dict(
+        base_url=info.connection_url,
+        api_key=token,
+        model=model_name,
+        streaming=streaming,
+    )
+    if http_client is not None:
+        openai_kwargs['http_client'] = http_client
+    return ChatOpenAI(
+        **openai_kwargs,
+        **kwargs,
+    )

singlestoredb/ai/embeddings.py

@@ -0,0 +1,128 @@
+import os
+from typing import Any
+from typing import Callable
+from typing import Optional
+from typing import Union
+
+import httpx
+
+from singlestoredb import manage_workspaces
+
+try:
+    from langchain_openai import OpenAIEmbeddings
+except ImportError:
+    raise ImportError(
+        'Could not import langchain_openai python package. '
+        'Please install it with `pip install langchain_openai`.',
+    )
+
+try:
+    from langchain_aws import BedrockEmbeddings
+except ImportError:
+    raise ImportError(
+        'Could not import langchain-aws python package. '
+        'Please install it with `pip install langchain-aws`.',
+    )
+
+import boto3
+from botocore import UNSIGNED
+from botocore.config import Config
+
+
+def SingleStoreEmbeddingsFactory(
+    model_name: str,
+    api_key: Optional[str] = None,
+    http_client: Optional[httpx.Client] = None,
+    obo_token_getter: Optional[Callable[[], Optional[str]]] = None,
+    **kwargs: Any,
+) -> Union[OpenAIEmbeddings, BedrockEmbeddings]:
+    """Return an embeddings model instance (OpenAIEmbeddings or BedrockEmbeddings).
+    """
+    inference_api_manager = (
+        manage_workspaces().organizations.current.inference_apis
+    )
+    info = inference_api_manager.get(model_name=model_name)
+    token_env = os.environ.get('SINGLESTOREDB_USER_TOKEN')
+    token = api_key if api_key is not None else token_env
+
+    if info.hosting_platform == 'Amazon':
+        # Instantiate Bedrock client
+        cfg_kwargs = {
+            'signature_version': UNSIGNED,
+            'retries': {'max_attempts': 1, 'mode': 'standard'},
+        }
+        # Extract timeouts from http_client if provided
+        t = http_client.timeout if http_client is not None else None
+        connect_timeout = None
+        read_timeout = None
+        if t is not None:
+            if isinstance(t, httpx.Timeout):
+                if t.connect is not None:
+                    connect_timeout = float(t.connect)
+                if t.read is not None:
+                    read_timeout = float(t.read)
+                if connect_timeout is None and read_timeout is not None:
+                    connect_timeout = read_timeout
+                if read_timeout is None and connect_timeout is not None:
+                    read_timeout = connect_timeout
+            elif isinstance(t, (int, float)):
+                connect_timeout = float(t)
+                read_timeout = float(t)
+        if read_timeout is not None:
+            cfg_kwargs['read_timeout'] = read_timeout
+        if connect_timeout is not None:
+            cfg_kwargs['connect_timeout'] = connect_timeout
+
+        cfg = Config(**cfg_kwargs)
+        client = boto3.client(
+            'bedrock-runtime',
+            endpoint_url=info.connection_url,
+            region_name='us-east-1',
+            aws_access_key_id='placeholder',
+            aws_secret_access_key='placeholder',
+            config=cfg,
+        )
+
+        def _inject_headers(request: Any, **_ignored: Any) -> None:
+            """Inject dynamic auth/OBO headers prior to Bedrock sending."""
+            if obo_token_getter is not None:
+                obo_val = obo_token_getter()
+                if obo_val:
+                    request.headers['X-S2-OBO'] = obo_val
+            if token:
+                request.headers['Authorization'] = f'Bearer {token}'
+            request.headers.pop('X-Amz-Date', None)
+            request.headers.pop('X-Amz-Security-Token', None)
+
+        emitter = client._endpoint._event_emitter
+        emitter.register_first(
+            'before-send.bedrock-runtime.InvokeModel',
+            _inject_headers,
+        )
+        emitter.register_first(
+            'before-send.bedrock-runtime.InvokeModelWithResponseStream',
+            _inject_headers,
+        )
+
+        return BedrockEmbeddings(
+            model_id=model_name,
+            endpoint_url=info.connection_url,
+            region_name='us-east-1',
+            aws_access_key_id='placeholder',
+            aws_secret_access_key='placeholder',
+            client=client,
+            **kwargs,
+        )
+
+    # OpenAI / Azure OpenAI path
+    openai_kwargs = dict(
+        base_url=info.connection_url,
+        api_key=token,
+        model=model_name,
+    )
+    if http_client is not None:
+        openai_kwargs['http_client'] = http_client
+    return OpenAIEmbeddings(
+        **openai_kwargs,
+        **kwargs,
+    )

singlestoredb/alchemy/__init__.py

@@ -0,0 +1,90 @@
+#!/usr/bin/env python3
+import inspect
+from typing import Any
+from urllib.parse import quote
+
+try:
+    import sqlalchemy
+    from sqlalchemy_singlestoredb import *  # noqa: F403, F401
+    has_sqlalchemy = True
+except ImportError:
+    import warnings
+    warnings.warn(
+        'sqlalchemy_singlestoredb must be installed to use this module',
+        RuntimeWarning,
+    )
+    has_sqlalchemy = False
+
+from ..connection import build_params
+from ..connection import connect
+
+
+def create_engine(*args: Any, **kwargs: Any) -> Any:
+    """
+    Create an SQLAlchemy engine for SingleStoreDB.
+
+    Parameters
+    ----------
+    **kwargs : Any
+        The parameters taken here are the same as for
+        `sqlalchemy.create_engine`. However, this function can be
+        called without any parameters in order to inherit parameters
+        set by environment variables or parameters set in by
+        options in Python code.
+
+    See Also
+    --------
+    `sqlalchemy.create_engine`
+
+    Returns
+    -------
+    SQLAlchemy engine
+
+    """
+    if not has_sqlalchemy:
+        raise RuntimeError('sqlalchemy_singlestoredb package is not installed')
+
+    if len(args) > 1:
+        raise ValueError(
+            '`args` can only have a single element '
+            'containing the database URL',
+        )
+
+    if args:
+        kwargs['host'] = args[0]
+
+    conn_params = {}
+    sa_params = {}
+
+    conn_args = inspect.getfullargspec(connect).args
+
+    for key, value in kwargs.items():
+        if key in conn_args:
+            conn_params[key] = value
+        else:
+            sa_params[key] = value
+
+    params = build_params(**conn_params)
+    driver = params.pop('driver', None)
+    host = params.pop('host')
+    port = params.pop('port')
+    user = params.pop('user', None)
+    password = params.pop('password', None)
+    database = params.pop('database', '')
+
+    if not driver:
+        driver = 'singlestoredb+mysql'
+    elif not driver.startswith('singlestoredb'):
+        driver = f'singlestoredb+{driver}'
+
+    if user is not None and password is not None:
+        url = f'{driver}://{quote(user)}:{quote(password)}@' \
+              f'{host}:{port}/{quote(database)}'
+    elif user is not None:
+        url = f'{driver}://{quote(user)}@{host}:{port}/{quote(database)}'
+    elif password is not None:
+        url = f'{driver}://:{quote(password)}@{host}:{port}/{quote(database)}'
+    else:
+        url = f'{driver}://{host}:{port}/{quote(database)}'
+
+    return sqlalchemy.create_engine(url, connect_args=params, **sa_params)

singlestoredb/apps/__init__.py

@@ -0,0 +1,3 @@
+from ._cloud_functions import run_function_app  # noqa: F401
+from ._dashboards import run_dashboard_app  # noqa: F401
+from ._python_udfs import run_udf_app  # noqa: F401

singlestoredb/apps/_cloud_functions.py

@@ -0,0 +1,90 @@
+import asyncio
+import textwrap
+import typing
+
+from ._config import AppConfig
+from ._connection_info import ConnectionInfo
+from ._process import kill_process_by_port
+
+if typing.TYPE_CHECKING:
+    from fastapi import FastAPI
+    from ._uvicorn_util import AwaitableUvicornServer
+
+# Keep track of currently running server
+_running_server: 'typing.Optional[AwaitableUvicornServer]' = None
+
+
+async def run_function_app(
+    app: 'FastAPI',
+    log_level: str = 'error',
+    kill_existing_app_server: bool = True,
+) -> ConnectionInfo:
+    global _running_server
+    from ._uvicorn_util import AwaitableUvicornServer
+
+    try:
+        import uvicorn
+    except ImportError:
+        raise ImportError('package uvicorn is required to run cloud functions')
+    try:
+        import fastapi
+    except ImportError:
+        raise ImportError('package fastapi is required to run cloud functions')
+
+    if not isinstance(app, fastapi.FastAPI):
+        raise TypeError('app is not an instance of FastAPI')
+
+    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)
+
+    # Add `GET /` route, used for liveness check
+    @app.get('/')
+    def ping() -> str:
+        return 'Success!'
+
+    app.root_path = app_config.base_path
+
+    config = uvicorn.Config(
+        app,
+        host='0.0.0.0',
+        port=app_config.listen_port,
+        log_level=log_level,
+    )
+    _running_server = AwaitableUvicornServer(config)
+
+    asyncio.create_task(_running_server.serve())
+    await _running_server.wait_for_startup()
+
+    connection_info = ConnectionInfo(app_config.base_url, app_config.token)
+
+    if app_config.running_interactively:
+        if app_config.is_gateway_enabled:
+            print(
+                'Cloud function available at '
+                f'{app_config.base_url}docs?authToken={app_config.token}',
+            )
+        else:
+            curl_header = f'-H "Authorization: Bearer {app_config.token}"'
+            curl_example = f'curl "{app_config.base_url}" {curl_header}'
+            print(
+                textwrap.dedent(f"""
+                  Cloud function available at {app_config.base_url}
+
+                  Auth Token: {app_config.token}
+
+                  Curl example: {curl_example}
+
+              """).strip(),
+            )
+
+    return connection_info

singlestoredb/apps/_config.py

@@ -0,0 +1,72 @@
+import os
+from dataclasses import dataclass
+from typing import Optional
+
+
+@dataclass
+class AppConfig:
+    listen_port: int
+    base_url: str
+    base_path: str
+    notebook_server_id: str
+    app_token: Optional[str]
+    user_token: Optional[str]
+    running_interactively: bool
+    is_gateway_enabled: bool
+    is_local_dev: bool
+
+    @staticmethod
+    def _read_variable(name: str) -> str:
+        value = os.environ.get(name)
+        if value is None:
+            raise RuntimeError(
+                f'Missing {name} environment variable. '
+                'Is the code running outside SingleStoreDB notebook environment?',
+            )
+        return value
+
+    @classmethod
+    def from_env(cls) -> 'AppConfig':
+        port = cls._read_variable('SINGLESTOREDB_APP_LISTEN_PORT')
+        base_url = cls._read_variable('SINGLESTOREDB_APP_BASE_URL')
+        base_path = cls._read_variable('SINGLESTOREDB_APP_BASE_PATH')
+        notebook_server_id = cls._read_variable('SINGLESTOREDB_NOTEBOOK_SERVER_ID')
+        is_local_dev_env_var = cls._read_variable('SINGLESTOREDB_IS_LOCAL_DEV')
+
+        workload_type = os.environ.get('SINGLESTOREDB_WORKLOAD_TYPE')
+        running_interactively = workload_type == 'InteractiveNotebook'
+
+        is_gateway_enabled = 'SINGLESTOREDB_NOVA_GATEWAY_ENDPOINT' in os.environ
+
+        app_token = os.environ.get('SINGLESTOREDB_APP_TOKEN')
+        user_token = os.environ.get('SINGLESTOREDB_USER_TOKEN')
+
+        # Make sure the required variables are present
+        # and present useful error message if not
+        if running_interactively:
+            if is_gateway_enabled:
+                app_token = cls._read_variable('SINGLESTOREDB_APP_TOKEN')
+            else:
+                user_token = cls._read_variable('SINGLESTOREDB_USER_TOKEN')
+
+        return cls(
+            listen_port=int(port),
+            base_url=base_url,
+            base_path=base_path,
+            notebook_server_id=notebook_server_id,
+            app_token=app_token,
+            user_token=user_token,
+            running_interactively=running_interactively,
+            is_gateway_enabled=is_gateway_enabled,
+            is_local_dev=is_local_dev_env_var == 'true',
+        )
+
+    @property
+    def token(self) -> Optional[str]:
+        """
+        Returns None if running non-interactively
+        """
+        if self.is_gateway_enabled:
+            return self.app_token
+        else:
+            return self.user_token

singlestoredb/apps/_connection_info.py

@@ -0,0 +1,18 @@
+from dataclasses import dataclass
+from typing import Any
+from typing import Dict
+from typing import Optional
+
+
+@dataclass
+class ConnectionInfo:
+    url: str
+
+    # Only present in interactive mode
+    token: Optional[str]
+
+
+@dataclass
+class UdfConnectionInfo:
+    url: str
+    functions: Dict[str, Any]

singlestoredb/apps/_dashboards.py

@@ -0,0 +1,47 @@
+import typing
+
+from ._config import AppConfig
+from ._process import kill_process_by_port
+from ._stdout_supress import StdoutSuppressor
+from singlestoredb.apps._connection_info import ConnectionInfo
+
+if typing.TYPE_CHECKING:
+    from dash import Dash
+
+
+async def run_dashboard_app(
+    app: 'Dash',
+    debug: bool = False,
+    kill_existing_app_server: bool = True,
+) -> ConnectionInfo:
+    try:
+        from dash import Dash
+    except ImportError:
+        raise ImportError('package dash is required to run dashboards')
+
+    if not isinstance(app, Dash):
+        raise TypeError('app is not an instance of Dash App')
+
+    app_config = AppConfig.from_env()
+
+    if kill_existing_app_server:
+        kill_process_by_port(app_config.listen_port)
+
+    if app.config.requests_pathname_prefix is None or \
+            app.config.requests_pathname_prefix != app_config.base_path:
+        raise RuntimeError('''
+requests_pathname_prefix of the Dash App is invalid. Please set
+requests_pathname_prefix=os.environ['SINGLESTOREDB_APP_BASE_PATH']
+while initializing the Dash App and retry''')
+
+    with StdoutSuppressor():
+        app.run(
+            host='0.0.0.0',
+            debug=debug,
+            port=str(app_config.listen_port),
+            jupyter_mode='external',
+        )
+
+    if app_config.running_interactively:
+        print(f'Dash app available at {app_config.base_url}?authToken={app_config.token}')
+    return ConnectionInfo(app_config.base_url, app_config.token)