singlestoredb 1.12.4__cp38-abi3-win_amd64.whl → 1.13.0__cp38-abi3-win_amd64.whl

singlestoredb/__init__.py

@@ -13,7 +13,7 @@ Examples
 
 """
 
-__version__ = '1.12.4'
+__version__ = '1.13.0'
 
 from typing import Any
 

singlestoredb/apps/__init__.py

@@ -1,2 +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/_config.py

@@ -8,10 +8,12 @@ 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:
@@ -28,6 +30,8 @@ class 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'
@@ -49,10 +53,12 @@ class AppConfig:
             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

singlestoredb/apps/_connection_info.py

@@ -1,4 +1,6 @@
 from dataclasses import dataclass
+from typing import Any
+from typing import Dict
 from typing import Optional
 
 
@@ -8,3 +10,9 @@ class ConnectionInfo:
 
     # Only present in interactive mode
     token: Optional[str]
+
+
+@dataclass
+class UdfConnectionInfo:
+    url: str
+    functions: Dict[str, Any]

singlestoredb/apps/_python_udfs.py

@@ -0,0 +1,85 @@
+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
+
+
+async def run_udf_app(
+    replace_existing: bool,
+    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)
+
+    config = uvicorn.Config(
+        app,
+        host='0.0.0.0',
+        port=app_config.listen_port,
+        log_level=log_level,
+    )
+    _running_server = AwaitableUvicornServer(config)
+
+    # Register the functions
+    app.register_functions(replace=replace_existing)
+
+    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/config.py

@@ -317,7 +317,7 @@ register_option(
     'external_function.app_mode', 'string',
     functools.partial(
         check_str,
-        valid_values=['remote', 'collocated'],
+        valid_values=['remote', 'collocated', 'managed'],
     ),
     'remote',
     'Specifies the mode of operation of the external function application.',
@@ -407,6 +407,18 @@ register_option(
     environ=['SINGLESTOREDB_EXT_FUNC_LOG_LEVEL'],
 )
 
+register_option(
+    'external_function.name_prefix', 'string', check_str, '',
+    'Prefix to add to external function names.',
+    environ=['SINGLESTOREDB_EXT_FUNC_NAME_PREFIX'],
+)
+
+register_option(
+    'external_function.name_suffix', 'string', check_str, '',
+    'Suffix to add to external function names.',
+    environ=['SINGLESTOREDB_EXT_FUNC_NAME_SUFFIX'],
+)
+
 register_option(
     'external_function.connection', 'string', check_str,
     os.environ.get('SINGLESTOREDB_URL') or None,
@@ -415,7 +427,7 @@ register_option(
 )
 
 register_option(
-    'external_function.host', 'string', check_str, '127.0.0.1',
+    'external_function.host', 'string', check_str, 'localhost',
     'Specifies the host to bind the server to.',
     environ=['SINGLESTOREDB_EXT_FUNC_HOST'],
 )

singlestoredb/functions/__init__.py

@@ -1,2 +1,12 @@
-from .decorator import tvf  # noqa: F401
 from .decorator import udf  # noqa: F401
+from .typing import Masked  # noqa: F401
+from .typing import Table  # noqa: F401
+from .utils import VectorTypes
+
+
+F32 = VectorTypes.F32
+F64 = VectorTypes.F64
+I8 = VectorTypes.I8
+I16 = VectorTypes.I16
+I32 = VectorTypes.I32
+I64 = VectorTypes.I64

singlestoredb/functions/decorator.py

@@ -1,183 +1,113 @@
-import dataclasses
-import datetime
 import functools
 import inspect
 from typing import Any
 from typing import Callable
-from typing import Dict
 from typing import List
 from typing import Optional
-from typing import Tuple
+from typing import Type
 from typing import Union
 
-from . import dtypes
-from .dtypes import DataType
-from .signature import simplify_dtype
-
-try:
-    import pydantic
-    has_pydantic = True
-except ImportError:
-    has_pydantic = False
-
-python_type_map: Dict[Any, Callable[..., str]] = {
-    str: dtypes.TEXT,
-    int: dtypes.BIGINT,
-    float: dtypes.DOUBLE,
-    bool: dtypes.BOOL,
-    bytes: dtypes.BINARY,
-    bytearray: dtypes.BINARY,
-    datetime.datetime: dtypes.DATETIME,
-    datetime.date: dtypes.DATE,
-    datetime.timedelta: dtypes.TIME,
-}
-
-
-def listify(x: Any) -> List[Any]:
-    """Make sure sure value is a list."""
-    if x is None:
-        return []
-    if isinstance(x, (list, tuple, set)):
-        return list(x)
-    return [x]
-
-
-def process_annotation(annotation: Any) -> Tuple[Any, bool]:
-    types = simplify_dtype(annotation)
-    if isinstance(types, list):
-        nullable = False
-        if type(None) in types:
-            nullable = True
-            types = [x for x in types if x is not type(None)]
-        if len(types) > 1:
-            raise ValueError(f'multiple types not supported: {annotation}')
-        return types[0], nullable
-    return types, True
-
-
-def process_types(params: Any) -> Any:
-    if params is None:
-        return params, []
-
-    elif isinstance(params, (list, tuple)):
-        params = list(params)
-        for i, item in enumerate(params):
-            if params[i] in python_type_map:
-                params[i] = python_type_map[params[i]]()
-            elif callable(item):
-                params[i] = item()
-        for item in params:
-            if not isinstance(item, str):
-                raise TypeError(f'unrecognized type for parameter: {item}')
-        return params, []
-
-    elif isinstance(params, dict):
-        names = []
-        params = dict(params)
-        for k, v in list(params.items()):
-            names.append(k)
-            if params[k] in python_type_map:
-                params[k] = python_type_map[params[k]]()
-            elif callable(v):
-                params[k] = v()
-        for item in params.values():
-            if not isinstance(item, str):
-                raise TypeError(f'unrecognized type for parameter: {item}')
-        return params, names
-
-    elif dataclasses.is_dataclass(params):
-        names = []
-        out = []
-        for item in dataclasses.fields(params):
-            typ, nullable = process_annotation(item.type)
-            sql_type = process_types(typ)[0]
-            if not nullable:
-                sql_type = sql_type.replace('NULL', 'NOT NULL')
-            out.append(sql_type)
-            names.append(item.name)
-        return out, names
-
-    elif has_pydantic and inspect.isclass(params) \
-            and issubclass(params, pydantic.BaseModel):
-        names = []
-        out = []
-        for name, item in params.model_fields.items():
-            typ, nullable = process_annotation(item.annotation)
-            sql_type = process_types(typ)[0]
-            if not nullable:
-                sql_type = sql_type.replace('NULL', 'NOT NULL')
-            out.append(sql_type)
-            names.append(name)
-        return out, names
-
-    elif params in python_type_map:
-        return python_type_map[params](), []
-
-    elif callable(params):
-        return params(), []
-
-    elif isinstance(params, str):
-        return params, []
-
-    raise TypeError(f'unrecognized data type for args: {params}')
+from . import utils
+from .dtypes import SQLString
+
+
+ParameterType = Union[
+    str,
+    Callable[..., SQLString],
+    List[Union[str, Callable[..., SQLString]]],
+    Type[Any],
+]
+
+ReturnType = ParameterType
+
+
+def is_valid_type(obj: Any) -> bool:
+    """Check if the object is a valid type for a schema definition."""
+    if not inspect.isclass(obj):
+        return False
+
+    if utils.is_typeddict(obj):
+        return True
+
+    if utils.is_namedtuple(obj):
+        return True
+
+    if utils.is_dataclass(obj):
+        return True
+
+    # We don't want to import pydantic here, so we check if
+    # the class is a subclass
+    if utils.is_pydantic(obj):
+        return True
+
+    return False
+
+
+def is_valid_callable(obj: Any) -> bool:
+    """Check if the object is a valid callable for a parameter type."""
+    if not callable(obj):
+        return False
+
+    returns = utils.get_annotations(obj).get('return', None)
+
+    if inspect.isclass(returns) and issubclass(returns, str):
+        return True
+
+    raise TypeError(
+        f'callable {obj} must return a str, '
+        f'but got {returns}',
+    )
+
+
+def expand_types(args: Any) -> Optional[Union[List[str], Type[Any]]]:
+    """Expand the types for the function arguments / return values."""
+    if args is None:
+        return None
+
+    # SQL string
+    if isinstance(args, str):
+        return [args]
+
+    # General way of accepting pydantic.BaseModel, NamedTuple, TypedDict
+    elif is_valid_type(args):
+        return args
+
+    # List of SQL strings or callables
+    elif isinstance(args, list):
+        new_args = []
+        for arg in args:
+            if isinstance(arg, str):
+                new_args.append(arg)
+            elif callable(arg):
+                new_args.append(arg())
+            else:
+                raise TypeError(f'unrecognized type for parameter: {arg}')
+        return new_args
+
+    # Callable that returns a SQL string
+    elif is_valid_callable(args):
+        out = args()
+        if not isinstance(out, str):
+            raise TypeError(f'unrecognized type for parameter: {args}')
+        return [out]
+
+    raise TypeError(f'unrecognized type for parameter: {args}')
 
 
 def _func(
     func: Optional[Callable[..., Any]] = None,
     *,
     name: Optional[str] = None,
-    args: Optional[
-        Union[
-            DataType,
-            List[DataType],
-            Dict[str, DataType],
-            'pydantic.BaseModel',
-            type,
-        ]
-    ] = None,
-    returns: Optional[
-        Union[
-            str,
-            List[DataType],
-            List[type],
-            'pydantic.BaseModel',
-            type,
-        ]
-    ] = None,
-    data_format: Optional[str] = None,
-    include_masks: bool = False,
-    function_type: str = 'udf',
-    output_fields: Optional[List[str]] = None,
+    args: Optional[ParameterType] = None,
+    returns: Optional[ReturnType] = None,
 ) -> Callable[..., Any]:
     """Generic wrapper for UDF and TVF decorators."""
-    args, _ = process_types(args)
-    returns, fields = process_types(returns)
-
-    if not output_fields and fields:
-        output_fields = fields
-
-    if isinstance(returns, list) \
-            and isinstance(output_fields, list) \
-            and len(output_fields) != len(returns):
-        raise ValueError(
-            'The number of output fields must match the number of return types',
-        )
-
-    if include_masks and data_format == 'python':
-        raise RuntimeError(
-            'include_masks is only valid when using '
-            'vectors for input parameters',
-        )
 
     _singlestoredb_attrs = {  # type: ignore
         k: v for k, v in dict(
             name=name,
-            args=args,
-            returns=returns,
-            data_format=data_format,
-            include_masks=include_masks,
-            function_type=function_type,
-            output_fields=output_fields or None,
+            args=expand_types(args),
+            returns=expand_types(returns),
         ).items() if v is not None
     }
 
@@ -186,10 +116,14 @@ def _func(
     # in at that time.
     if func is None:
         def decorate(func: Callable[..., Any]) -> Callable[..., Any]:
+
             def wrapper(*args: Any, **kwargs: Any) -> Callable[..., Any]:
                 return func(*args, **kwargs)  # type: ignore
+
             wrapper._singlestoredb_attrs = _singlestoredb_attrs  # type: ignore
+
             return functools.wraps(func)(wrapper)
+
         return decorate
 
     def wrapper(*args: Any, **kwargs: Any) -> Callable[..., Any]:
@@ -204,13 +138,11 @@ def udf(
     func: Optional[Callable[..., Any]] = None,
     *,
     name: Optional[str] = None,
-    args: Optional[Union[DataType, List[DataType], Dict[str, DataType]]] = None,
-    returns: Optional[Union[str, List[DataType], List[type]]] = None,
-    data_format: Optional[str] = None,
-    include_masks: bool = False,
+    args: Optional[ParameterType] = None,
+    returns: Optional[ReturnType] = None,
 ) -> Callable[..., Any]:
     """
-    Apply attributes to a UDF.
+    Define a user-defined function (UDF).
 
     Parameters
     ----------
@@ -218,71 +150,7 @@ def udf(
         The UDF to apply parameters to
     name : str, optional
         The name to use for the UDF in the database
-    args : str | Callable | List[str | Callable] | Dict[str, str | Callable], optional
-        Specifies the data types of the function arguments. Typically,
-        the function data types are derived from the function parameter
-        annotations. These annotations can be overridden. If the function
-        takes a single type for all parameters, `args` can be set to a
-        SQL string describing all parameters. If the function takes more
-        than one parameter and all of the parameters are being manually
-        defined, a list of SQL strings may be used (one for each parameter).
-        A dictionary of SQL strings may be used to specify a parameter type
-        for a subset of parameters; the keys are the names of the
-        function parameters. Callables may also be used for datatypes. This
-        is primarily for using the functions in the ``dtypes`` module that
-        are associated with SQL types with all default options (e.g., ``dt.FLOAT``).
-    returns : str, optional
-        Specifies the return data type of the function. If not specified,
-        the type annotation from the function is used.
-    data_format : str, optional
-        The data format of each parameter: python, pandas, arrow, polars
-    include_masks : bool, optional
-        Should boolean masks be included with each input parameter to indicate
-        which elements are NULL? This is only used when a input parameters are
-        configured to a vector type (numpy, pandas, polars, arrow).
-
-    Returns
-    -------
-    Callable
-
-    """
-    return _func(
-        func=func,
-        name=name,
-        args=args,
-        returns=returns,
-        data_format=data_format,
-        include_masks=include_masks,
-        function_type='udf',
-    )
-
-
-udf.pandas = functools.partial(udf, data_format='pandas')  # type: ignore
-udf.polars = functools.partial(udf, data_format='polars')  # type: ignore
-udf.arrow = functools.partial(udf, data_format='arrow')  # type: ignore
-udf.numpy = functools.partial(udf, data_format='numpy')  # type: ignore
-
-
-def tvf(
-    func: Optional[Callable[..., Any]] = None,
-    *,
-    name: Optional[str] = None,
-    args: Optional[Union[DataType, List[DataType], Dict[str, DataType]]] = None,
-    returns: Optional[Union[str, List[DataType], List[type]]] = None,
-    data_format: Optional[str] = None,
-    include_masks: bool = False,
-    output_fields: Optional[List[str]] = None,
-) -> Callable[..., Any]:
-    """
-    Apply attributes to a TVF.
-
-    Parameters
-    ----------
-    func : callable, optional
-        The TVF to apply parameters to
-    name : str, optional
-        The name to use for the TVF in the database
-    args : str | Callable | List[str | Callable] | Dict[str, str | Callable], optional
+    args : str | Type | Callable | List[str | Callable], optional
         Specifies the data types of the function arguments. Typically,
         the function data types are derived from the function parameter
         annotations. These annotations can be overridden. If the function
@@ -295,18 +163,10 @@ def tvf(
         function parameters. Callables may also be used for datatypes. This
         is primarily for using the functions in the ``dtypes`` module that
         are associated with SQL types with all default options (e.g., ``dt.FLOAT``).
-    returns : str, optional
-        Specifies the return data type of the function. If not specified,
-        the type annotation from the function is used.
-    data_format : str, optional
-        The data format of each parameter: python, pandas, arrow, polars
-    include_masks : bool, optional
-        Should boolean masks be included with each input parameter to indicate
-        which elements are NULL? This is only used when a input parameters are
-        configured to a vector type (numpy, pandas, polars, arrow).
-    output_fields : List[str], optional
-        The names of the output fields for the TVF. If not specified, the
-        names are generated.
+    returns : str | Type | Callable | List[str | Callable] | Table, optional
+        Specifies the return data type of the function. This parameter
+        works the same way as `args`. If the function is a table-valued
+        function, the return type should be a `Table` object.
 
     Returns
     -------
@@ -318,14 +178,4 @@ def tvf(
         name=name,
         args=args,
         returns=returns,
-        data_format=data_format,
-        include_masks=include_masks,
-        function_type='tvf',
-        output_fields=output_fields,
     )
-
-
-tvf.pandas = functools.partial(tvf, data_format='pandas')  # type: ignore
-tvf.polars = functools.partial(tvf, data_format='polars')  # type: ignore
-tvf.arrow = functools.partial(tvf, data_format='arrow')  # type: ignore
-tvf.numpy = functools.partial(tvf, data_format='numpy')  # type: ignore