singlestoredb 1.12.4__py3-none-any.whl → 1.13.0__py3-none-any.whl

singlestoredb/functions/ext/asgi.py

@@ -26,6 +26,7 @@ import argparse
 import asyncio
 import dataclasses
 import importlib.util
+import inspect
 import io
 import itertools
 import json
@@ -36,6 +37,7 @@ import secrets
 import sys
 import tempfile
 import textwrap
+import typing
 import urllib
 import zipfile
 import zipimport
@@ -62,6 +64,8 @@ from ...config import get_option
 from ...mysql.constants import FIELD_TYPE as ft
 from ..signature import get_signature
 from ..signature import signature_to_sql
+from ..typing import Masked
+from ..typing import Table
 
 try:
     import cloudpickle
@@ -148,20 +152,100 @@ def as_tuple(x: Any) -> Any:
     if has_pydantic and isinstance(x, BaseModel):
         return tuple(x.model_dump().values())
     if dataclasses.is_dataclass(x):
-        return dataclasses.astuple(x)
-    return x
+        return dataclasses.astuple(x)  # type: ignore
+    if isinstance(x, dict):
+        return tuple(x.values())
+    return tuple(x)
 
 
 def as_list_of_tuples(x: Any) -> Any:
     """Convert object to a list of tuples."""
+    if isinstance(x, Table):
+        x = x[0]
     if isinstance(x, (list, tuple)) and len(x) > 0:
+        if isinstance(x[0], (list, tuple)):
+            return x
         if has_pydantic and isinstance(x[0], BaseModel):
             return [tuple(y.model_dump().values()) for y in x]
         if dataclasses.is_dataclass(x[0]):
             return [dataclasses.astuple(y) for y in x]
+        if isinstance(x[0], dict):
+            return [tuple(y.values()) for y in x]
+        return [(y,) for y in x]
     return x
 
 
+def get_dataframe_columns(df: Any) -> List[Any]:
+    """Return columns of data from a dataframe/table."""
+    if isinstance(df, Table):
+        if len(df) == 1:
+            df = df[0]
+        else:
+            return list(df)
+
+    if isinstance(df, Masked):
+        return [df]
+
+    if isinstance(df, tuple):
+        return list(df)
+
+    rtype = str(type(df)).lower()
+    if 'dataframe' in rtype:
+        return [df[x] for x in df.columns]
+    elif 'table' in rtype:
+        return df.columns
+    elif 'series' in rtype:
+        return [df]
+    elif 'array' in rtype:
+        return [df]
+    elif 'tuple' in rtype:
+        return list(df)
+
+    raise TypeError(
+        'Unsupported data type for dataframe columns: '
+        f'{rtype}',
+    )
+
+
+def get_array_class(data_format: str) -> Callable[..., Any]:
+    """
+    Get the array class for the current data format.
+
+    """
+    if data_format == 'polars':
+        import polars as pl
+        array_cls = pl.Series
+    elif data_format == 'arrow':
+        import pyarrow as pa
+        array_cls = pa.array
+    elif data_format == 'pandas':
+        import pandas as pd
+        array_cls = pd.Series
+    else:
+        import numpy as np
+        array_cls = np.array
+    return array_cls
+
+
+def get_masked_params(func: Callable[..., Any]) -> List[bool]:
+    """
+    Get the list of masked parameters for the function.
+
+    Parameters
+    ----------
+    func : Callable
+        The function to call as the endpoint
+
+    Returns
+    -------
+    List[bool]
+        Boolean list of masked parameters
+
+    """
+    params = inspect.signature(func).parameters
+    return [typing.get_origin(x.annotation) is Masked for x in params.values()]
+
+
 def make_func(
     name: str,
     func: Callable[..., Any],
@@ -181,104 +265,107 @@ def make_func(
     (Callable, Dict[str, Any])
 
     """
-    attrs = getattr(func, '_singlestoredb_attrs', {})
-    data_format = attrs.get('data_format') or 'python'
-    include_masks = attrs.get('include_masks', False)
-    function_type = attrs.get('function_type', 'udf').lower()
     info: Dict[str, Any] = {}
 
+    sig = get_signature(func, func_name=name)
+
+    function_type = sig.get('function_type', 'udf')
+    args_data_format = sig.get('args_data_format', 'scalar')
+    returns_data_format = sig.get('returns_data_format', 'scalar')
+
+    masks = get_masked_params(func)
+
     if function_type == 'tvf':
-        if data_format == 'python':
+        # Scalar / list types (row-based)
+        if returns_data_format in ['scalar', 'list']:
             async def do_func(
                 row_ids: Sequence[int],
                 rows: Sequence[Sequence[Any]],
-            ) -> Tuple[
-                Sequence[int],
-                List[Tuple[Any]],
-            ]:
+            ) -> Tuple[Sequence[int], List[Tuple[Any, ...]]]:
                 '''Call function on given rows of data.'''
                 out_ids: List[int] = []
                 out = []
+                # Call function on each row of data
                 for i, res in zip(row_ids, func_map(func, rows)):
                     out.extend(as_list_of_tuples(res))
                     out_ids.extend([row_ids[i]] * (len(out)-len(out_ids)))
                 return out_ids, out
 
+        # Vector formats (column-based)
         else:
-            # Vector formats use the same function wrapper
+            array_cls = get_array_class(returns_data_format)
+
             async def do_func(  # type: ignore
                 row_ids: Sequence[int],
                 cols: Sequence[Tuple[Sequence[Any], Optional[Sequence[bool]]]],
-            ) -> Tuple[Sequence[int], List[Tuple[Any, ...]]]:
+            ) -> Tuple[
+                Sequence[int],
+                List[Tuple[Sequence[Any], Optional[Sequence[bool]]]],
+            ]:
                 '''Call function on given cols of data.'''
-                if include_masks:
-                    out = func(*cols)
-                    assert isinstance(out, tuple)
-                    return row_ids, [out]
-
-                out = []
-                res = func(*[x[0] for x in cols])
-                rtype = str(type(res)).lower()
-
-                # Map tables / dataframes to a list of columns
-                if 'dataframe' in rtype:
-                    res = [res[x] for x in res.columns]
-                elif 'table' in rtype:
-                    res = res.columns
-
-                for vec in res:
-                    # C extension only supports Python objects as strings
-                    if data_format == 'numpy' and str(vec.dtype)[:2] in ['<U', '<S']:
-                        vec = vec.astype(object)
-                    out.append((vec, None))
-
                 # NOTE: There is no way to determine which row ID belongs to
                 #        each result row, so we just have to use the same
                 #        row ID for all rows in the result.
-                if data_format == 'polars':
-                    import polars as pl
-                    array_cls = pl.Series
-                elif data_format == 'arrow':
-                    import pyarrow as pa
-                    array_cls = pa.array
-                elif data_format == 'pandas':
-                    import pandas as pd
-                    array_cls = pd.Series
+
+                def build_tuple(x: Any) -> Any:
+                    return tuple(x) if isinstance(x, Masked) else (x, None)
+
+                # Call function on each column of data
+                if cols and cols[0]:
+                    res = get_dataframe_columns(
+                        func(*[x if m else x[0] for x, m in zip(cols, masks)]),
+                    )
+                else:
+                    res = get_dataframe_columns(func())
+
+                # Generate row IDs
+                if isinstance(res[0], Masked):
+                    row_ids = array_cls([row_ids[0]] * len(res[0][0]))
                 else:
-                    import numpy as np
-                    array_cls = np.array
+                    row_ids = array_cls([row_ids[0]] * len(res[0]))
 
-                return array_cls([row_ids[0]] * len(out[0][0])), out
+                return row_ids, [build_tuple(x) for x in res]
 
     else:
-        if data_format == 'python':
+        # Scalar / list types (row-based)
+        if returns_data_format in ['scalar', 'list']:
             async def do_func(
                 row_ids: Sequence[int],
                 rows: Sequence[Sequence[Any]],
-            ) -> Tuple[
-                Sequence[int],
-                List[Tuple[Any]],
-            ]:
+            ) -> Tuple[Sequence[int], List[Tuple[Any, ...]]]:
                 '''Call function on given rows of data.'''
                 return row_ids, [as_tuple(x) for x in zip(func_map(func, rows))]
 
+        # Vector formats (column-based)
         else:
-            # Vector formats use the same function wrapper
+            array_cls = get_array_class(returns_data_format)
+
             async def do_func(  # type: ignore
                 row_ids: Sequence[int],
                 cols: Sequence[Tuple[Sequence[Any], Optional[Sequence[bool]]]],
-            ) -> Tuple[Sequence[int], List[Tuple[Any, ...]]]:
+            ) -> Tuple[
+                Sequence[int],
+                List[Tuple[Sequence[Any], Optional[Sequence[bool]]]],
+            ]:
                 '''Call function on given cols of data.'''
-                if include_masks:
-                    out = func(*cols)
-                    assert isinstance(out, tuple)
-                    return row_ids, [out]
+                row_ids = array_cls(row_ids)
+
+                def build_tuple(x: Any) -> Any:
+                    return tuple(x) if isinstance(x, Masked) else (x, None)
+
+                # Call the function with `cols` as the function parameters
+                if cols and cols[0]:
+                    out = func(*[x if m else x[0] for x, m in zip(cols, masks)])
+                else:
+                    out = func()
 
-                out = func(*[x[0] for x in cols])
+                # Single masked value
+                if isinstance(out, Masked):
+                    return row_ids, [tuple(out)]
 
                 # Multiple return values
                 if isinstance(out, tuple):
-                    return row_ids, [(x, None) for x in out]
+                    return row_ids, [build_tuple(x) for x in out]
 
                 # Single return value
                 return row_ids, [(out, None)]
@@ -286,13 +373,12 @@ def make_func(
     do_func.__name__ = name
     do_func.__doc__ = func.__doc__
 
-    sig = get_signature(func, name=name)
-
     # Store signature for generating CREATE FUNCTION calls
     info['signature'] = sig
 
     # Set data format
-    info['data_format'] = data_format
+    info['args_data_format'] = args_data_format
+    info['returns_data_format'] = returns_data_format
 
     # Set function type
     info['function_type'] = function_type
@@ -306,20 +392,13 @@ def make_func(
         colspec.append((x['name'], rowdat_1_type_map[dtype]))
     info['colspec'] = colspec
 
-    def parse_return_type(s: str) -> List[str]:
-        if s.startswith('tuple['):
-            return s[6:-1].split(',')
-        if s.startswith('array[tuple['):
-            return s[12:-2].split(',')
-        return [s]
-
     # Setup return type
     returns = []
-    for x in parse_return_type(sig['returns']['dtype']):
-        dtype = x.replace('?', '')
+    for x in sig['returns']:
+        dtype = x['dtype'].replace('?', '')
         if dtype not in rowdat_1_type_map:
             raise TypeError(f'no data type mapping for {dtype}')
-        returns.append(rowdat_1_type_map[dtype])
+        returns.append((x['name'], rowdat_1_type_map[dtype]))
     info['returns'] = returns
 
     return do_func, info
@@ -371,6 +450,13 @@ class Application(object):
         headers=[(b'content-type', b'text/plain')],
     )
 
+    # Error response start
+    error_response_dict: Dict[str, Any] = dict(
+        type='http.response.start',
+        status=401,
+        headers=[(b'content-type', b'text/plain')],
+    )
+
     # JSON response start
     json_response_dict: Dict[str, Any] = dict(
         type='http.response.start',
@@ -405,7 +491,12 @@ class Application(object):
 
     # Data format + version handlers
     handlers = {
-        (b'application/octet-stream', b'1.0', 'python'): dict(
+        (b'application/octet-stream', b'1.0', 'scalar'): dict(
+            load=rowdat_1.load,
+            dump=rowdat_1.dump,
+            response=rowdat_1_response_dict,
+        ),
+        (b'application/octet-stream', b'1.0', 'list'): dict(
             load=rowdat_1.load,
             dump=rowdat_1.dump,
             response=rowdat_1_response_dict,
@@ -430,7 +521,12 @@ class Application(object):
             dump=rowdat_1.dump_arrow,
             response=rowdat_1_response_dict,
         ),
-        (b'application/json', b'1.0', 'python'): dict(
+        (b'application/json', b'1.0', 'scalar'): dict(
+            load=jdata.load,
+            dump=jdata.dump,
+            response=json_response_dict,
+        ),
+        (b'application/json', b'1.0', 'list'): dict(
             load=jdata.load,
             dump=jdata.dump,
             response=json_response_dict,
@@ -455,7 +551,7 @@ class Application(object):
             dump=jdata.dump_arrow,
             response=json_response_dict,
         ),
-        (b'application/vnd.apache.arrow.file', b'1.0', 'python'): dict(
+        (b'application/vnd.apache.arrow.file', b'1.0', 'scalar'): dict(
             load=arrow.load,
             dump=arrow.dump,
             response=arrow_response_dict,
@@ -485,6 +581,7 @@ class Application(object):
     # Valid URL paths
     invoke_path = ('invoke',)
     show_create_function_path = ('show', 'create_function')
+    show_function_info_path = ('show', 'function_info')
 
     def __init__(
         self,
@@ -505,6 +602,8 @@ class Application(object):
         link_name: Optional[str] = get_option('external_function.link_name'),
         link_config: Optional[Dict[str, Any]] = None,
         link_credentials: Optional[Dict[str, Any]] = None,
+        name_prefix: str = get_option('external_function.name_prefix'),
+        name_suffix: str = get_option('external_function.name_suffix'),
     ) -> None:
         if link_name and (link_config or link_credentials):
             raise ValueError(
@@ -561,6 +660,7 @@ class Application(object):
                         if not hasattr(x, '_singlestoredb_attrs'):
                             continue
                         name = x._singlestoredb_attrs.get('name', x.__name__)
+                        name = f'{name_prefix}{name}{name_suffix}'
                         external_functions[x.__name__] = x
                         func, info = make_func(name, x)
                         endpoints[name.encode('utf-8')] = func, info
@@ -576,6 +676,7 @@ class Application(object):
                     # Add endpoint for each exported function
                     for name, alias in get_func_names(func_names):
                         item = getattr(pkg, name)
+                        alias = f'{name_prefix}{name}{name_suffix}'
                         external_functions[name] = item
                         func, info = make_func(alias, item)
                         endpoints[alias.encode('utf-8')] = func, info
@@ -588,6 +689,7 @@ class Application(object):
                     if not hasattr(x, '_singlestoredb_attrs'):
                         continue
                     name = x._singlestoredb_attrs.get('name', x.__name__)
+                    name = f'{name_prefix}{name}{name_suffix}'
                     external_functions[x.__name__] = x
                     func, info = make_func(name, x)
                     endpoints[name.encode('utf-8')] = func, info
@@ -595,6 +697,7 @@ class Application(object):
             else:
                 alias = funcs.__name__
                 external_functions[funcs.__name__] = funcs
+                alias = f'{name_prefix}{alias}{name_suffix}'
                 func, info = make_func(alias, funcs)
                 endpoints[alias.encode('utf-8')] = func, info
 
@@ -648,7 +751,8 @@ class Application(object):
 
         # Call the endpoint
         if method == 'POST' and func is not None and path == self.invoke_path:
-            data_format = func_info['data_format']
+            args_data_format = func_info['args_data_format']
+            returns_data_format = func_info['returns_data_format']
             data = []
             more_body = True
             while more_body:
@@ -657,17 +761,24 @@ class Application(object):
                 more_body = request.get('more_body', False)
 
             data_version = headers.get(b's2-ef-version', b'')
-            input_handler = self.handlers[(content_type, data_version, data_format)]
-            output_handler = self.handlers[(accepts, data_version, data_format)]
+            input_handler = self.handlers[(content_type, data_version, args_data_format)]
+            output_handler = self.handlers[(accepts, data_version, returns_data_format)]
 
-            out = await func(
-                *input_handler['load'](  # type: ignore
-                    func_info['colspec'], b''.join(data),
-                ),
-            )
-            body = output_handler['dump'](func_info['returns'], *out)  # type: ignore
+            try:
+                out = await func(
+                    *input_handler['load'](  # type: ignore
+                        func_info['colspec'], b''.join(data),
+                    ),
+                )
+                body = output_handler['dump'](
+                    [x[1] for x in func_info['returns']], *out,  # type: ignore
+                )
+                await send(output_handler['response'])
 
-            await send(output_handler['response'])
+            except Exception as e:
+                logging.exception('Error in function call')
+                body = f'[{type(e).__name__}] {str(e).strip()}'.encode('utf-8')
+                await send(self.error_response_dict)
 
         # Handle api reflection
         elif method == 'GET' and path == self.show_create_function_path:
@@ -688,6 +799,12 @@ class Application(object):
 
             await send(self.text_response_dict)
 
+        # Return function info
+        elif method == 'GET' and (path == self.show_function_info_path or not path):
+            functions = self.get_function_info()
+            body = json.dumps(dict(functions=functions)).encode('utf-8')
+            await send(self.text_response_dict)
+
         # Path not found
         else:
             body = b''
@@ -725,21 +842,78 @@ class Application(object):
         """Locate all current functions and links belonging to this app."""
         funcs, links = set(), set()
         cur.execute('SHOW FUNCTIONS')
-        for name, ftype, _, _, _, link in list(cur):
+        for row in list(cur):
+            name, ftype, link = row[0], row[1], row[-1]
             # Only look at external functions
             if 'external' not in ftype.lower():
                 continue
             # See if function URL matches url
             cur.execute(f'SHOW CREATE FUNCTION `{name}`')
             for fname, _, code, *_ in list(cur):
-                m = re.search(r" (?:\w+) SERVICE '([^']+)'", code)
+                m = re.search(r" (?:\w+) (?:SERVICE|MANAGED) '([^']+)'", code)
                 if m and m.group(1) == self.url:
                     funcs.add(fname)
                     if link and re.match(r'^py_ext_func_link_\S{14}$', link):
                         links.add(link)
         return funcs, links
 
-    def show_create_functions(
+    def get_function_info(
+        self,
+        func_name: Optional[str] = None,
+    ) -> Dict[str, Any]:
+        """
+        Return the functions and function signature information.
+        Returns
+        -------
+        Dict[str, Any]
+        """
+        functions = {}
+        no_default = object()
+
+        for key, (_, info) in self.endpoints.items():
+            if not func_name or key == func_name:
+                sig = info['signature']
+                args = []
+
+                # Function arguments
+                for a in sig.get('args', []):
+                    dtype = a['dtype']
+                    nullable = '?' in dtype
+                    args.append(
+                        dict(
+                            name=a['name'],
+                            dtype=dtype.replace('?', ''),
+                            nullable=nullable,
+                        ),
+                    )
+                    if a.get('default', no_default) is not no_default:
+                        args[-1]['default'] = a['default']
+
+                # Return values
+                ret = sig.get('returns', [])
+                returns = []
+
+                for a in ret:
+                    dtype = a['dtype']
+                    nullable = '?' in dtype
+                    returns.append(
+                        dict(
+                            dtype=dtype.replace('?', ''),
+                            nullable=nullable,
+                        ),
+                    )
+                    if a.get('name', None):
+                        returns[-1]['name'] = a['name']
+                    if a.get('default', no_default) is not no_default:
+                        returns[-1]['default'] = a['default']
+
+                functions[sig['name']] = dict(
+                    args=args, returns=returns, function_type=info['function_type'],
+                )
+
+        return functions
+
+    def get_create_functions(
         self,
         replace: bool = False,
     ) -> List[str]:
@@ -807,7 +981,7 @@ class Application(object):
                         cur.execute(f'DROP FUNCTION IF EXISTS `{fname}`')
                     for link in links:
                         cur.execute(f'DROP LINK {link}')
-                for func in self.show_create_functions(replace=replace):
+                for func in self.get_create_functions(replace=replace):
                     cur.execute(func)
 
     def drop_functions(
@@ -1135,6 +1309,22 @@ def main(argv: Optional[List[str]] = None) -> None:
             ),
             help='logging level',
         )
+        parser.add_argument(
+            '--name-prefix', metavar='name_prefix',
+            default=defaults.get(
+                'name_prefix',
+                get_option('external_function.name_prefix'),
+            ),
+            help='Prefix to add to function names',
+        )
+        parser.add_argument(
+            '--name-suffix', metavar='name_suffix',
+            default=defaults.get(
+                'name_suffix',
+                get_option('external_function.name_suffix'),
+            ),
+            help='Suffix to add to function names',
+        )
         parser.add_argument(
             'functions', metavar='module.or.func.path', nargs='*',
             help='functions or modules to export in UDF server',
@@ -1217,6 +1407,11 @@ def main(argv: Optional[List[str]] = None) -> None:
         or defaults.get('replace_existing') \
         or get_option('external_function.replace_existing')
 
+    # Substitute in host / port if specified
+    if args.host != defaults.get('host') or args.port != defaults.get('port'):
+        u = urllib.parse.urlparse(args.url)
+        args.url = u._replace(netloc=f'{args.host}:{args.port}').geturl()
+
     # Create application from functions / module
     app = Application(
         functions=args.functions,
@@ -1227,9 +1422,11 @@ def main(argv: Optional[List[str]] = None) -> None:
         link_config=json.loads(args.link_config) or None,
         link_credentials=json.loads(args.link_credentials) or None,
         app_mode='remote',
+        name_prefix=args.name_prefix,
+        name_suffix=args.name_suffix,
     )
 
-    funcs = app.show_create_functions(replace=args.replace_existing)
+    funcs = app.get_create_functions(replace=args.replace_existing)
     if not funcs:
         raise RuntimeError('no functions specified')
 
@@ -1249,6 +1446,7 @@ def main(argv: Optional[List[str]] = None) -> None:
                 host=args.host or None,
                 port=args.port or None,
                 log_level=args.log_level,
+                lifespan='off',
             ).items() if v is not None
         }