singlestoredb 1.14.2__cp38-abi3-macosx_10_9_universal2.whl → 1.15.0__cp38-abi3-macosx_10_9_universal2.whl

singlestoredb/__init__.py

@@ -13,7 +13,7 @@ Examples
 
 """
 
-__version__ = '1.14.2'
+__version__ = '1.15.0'
 
 from typing import Any
 
@@ -25,7 +25,7 @@ from .exceptions import (
     DataError, ManagementError,
 )
 from .management import (
-    manage_cluster, manage_workspaces, manage_files,
+    manage_cluster, manage_workspaces, manage_files, manage_regions,
 )
 from .types import (
     Date, Time, Timestamp, DateFromTicks, TimeFromTicks, TimestampFromTicks,

singlestoredb/apps/_python_udfs.py

@@ -15,7 +15,6 @@ _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:
@@ -55,8 +54,9 @@ async def run_udf_app(
     )
     _running_server = AwaitableUvicornServer(config)
 
-    # Register the functions
-    app.register_functions(replace=replace_existing)
+    # Register the functions only if the app is running interactively.
+    if app_config.running_interactively:
+        app.register_functions(replace=True)
 
     asyncio.create_task(_running_server.serve())
     await _running_server.wait_for_startup()

singlestoredb/config.py

@@ -438,6 +438,11 @@ register_option(
     environ=['SINGLESTOREDB_EXT_FUNC_PORT'],
 )
 
+register_option(
+    'external_function.timeout', 'int', check_int, 24*60*60,
+    'Specifies the timeout in seconds for processing a batch of rows.',
+    environ=['SINGLESTOREDB_EXT_FUNC_TIMEOUT'],
+)
 
 #
 # Debugging options

singlestoredb/functions/decorator.py

@@ -1,3 +1,4 @@
+import asyncio
 import functools
 import inspect
 from typing import Any
@@ -19,6 +20,7 @@ ParameterType = Union[
 ]
 
 ReturnType = ParameterType
+UDFType = Callable[..., Any]
 
 
 def is_valid_type(obj: Any) -> bool:
@@ -100,7 +102,8 @@ def _func(
     name: Optional[str] = None,
     args: Optional[ParameterType] = None,
     returns: Optional[ReturnType] = None,
-) -> Callable[..., Any]:
+    timeout: Optional[int] = None,
+) -> UDFType:
     """Generic wrapper for UDF and TVF decorators."""
 
     _singlestoredb_attrs = {  # type: ignore
@@ -108,6 +111,7 @@ def _func(
             name=name,
             args=expand_types(args),
             returns=expand_types(returns),
+            timeout=timeout,
         ).items() if v is not None
     }
 
@@ -115,23 +119,33 @@ def _func(
     # called later, so the wrapper much be created with the func passed
     # in at that time.
     if func is None:
-        def decorate(func: Callable[..., Any]) -> Callable[..., Any]:
+        def decorate(func: UDFType) -> UDFType:
 
-            def wrapper(*args: Any, **kwargs: Any) -> Callable[..., Any]:
-                return func(*args, **kwargs)  # type: ignore
+            if asyncio.iscoroutinefunction(func):
+                async def async_wrapper(*args: Any, **kwargs: Any) -> UDFType:
+                    return await func(*args, **kwargs)  # type: ignore
+                async_wrapper._singlestoredb_attrs = _singlestoredb_attrs  # type: ignore
+                return functools.wraps(func)(async_wrapper)
 
-            wrapper._singlestoredb_attrs = _singlestoredb_attrs  # type: ignore
-
-            return functools.wraps(func)(wrapper)
+            else:
+                def wrapper(*args: Any, **kwargs: Any) -> UDFType:
+                    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]:
-        return func(*args, **kwargs)  # type: ignore
-
-    wrapper._singlestoredb_attrs = _singlestoredb_attrs  # type: ignore
+    if asyncio.iscoroutinefunction(func):
+        async def async_wrapper(*args: Any, **kwargs: Any) -> UDFType:
+            return await func(*args, **kwargs)  # type: ignore
+        async_wrapper._singlestoredb_attrs = _singlestoredb_attrs  # type: ignore
+        return functools.wraps(func)(async_wrapper)
 
-    return functools.wraps(func)(wrapper)
+    else:
+        def wrapper(*args: Any, **kwargs: Any) -> UDFType:
+            return func(*args, **kwargs)  # type: ignore
+        wrapper._singlestoredb_attrs = _singlestoredb_attrs  # type: ignore
+        return functools.wraps(func)(wrapper)
 
 
 def udf(
@@ -140,7 +154,8 @@ def udf(
     name: Optional[str] = None,
     args: Optional[ParameterType] = None,
     returns: Optional[ReturnType] = None,
-) -> Callable[..., Any]:
+    timeout: Optional[int] = None,
+) -> UDFType:
     """
     Define a user-defined function (UDF).
 
@@ -167,6 +182,9 @@ def udf(
         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.
+    timeout : int, optional
+        The timeout in seconds for the UDF execution. If not specified,
+        the global default timeout is used.
 
     Returns
     -------
@@ -178,4 +196,5 @@ def udf(
         name=name,
         args=args,
         returns=returns,
+        timeout=timeout,
     )

singlestoredb/functions/ext/asgi.py

@@ -24,7 +24,10 @@ Example
 """
 import argparse
 import asyncio
+import contextvars
 import dataclasses
+import datetime
+import functools
 import importlib.util
 import inspect
 import io
@@ -37,8 +40,11 @@ import secrets
 import sys
 import tempfile
 import textwrap
+import threading
+import time
 import typing
 import urllib
+import uuid
 import zipfile
 import zipimport
 from types import ModuleType
@@ -66,6 +72,7 @@ from ..signature import get_signature
 from ..signature import signature_to_sql
 from ..typing import Masked
 from ..typing import Table
+from .timer import Timer
 
 try:
     import cloudpickle
@@ -95,6 +102,15 @@ else:
     func_map = itertools.starmap
 
 
+async def to_thread(
+    func: Any, /, *args: Any, **kwargs: Dict[str, Any],
+) -> Any:
+    loop = asyncio.get_running_loop()
+    ctx = contextvars.copy_context()
+    func_call = functools.partial(ctx.run, func, *args, **kwargs)
+    return await loop.run_in_executor(None, func_call)
+
+
 # Use negative values to indicate unsigned ints / binary data / usec time precision
 rowdat_1_type_map = {
     'bool': ft.LONGLONG,
@@ -251,6 +267,32 @@ def build_tuple(x: Any) -> Any:
     return tuple(x) if isinstance(x, Masked) else (x, None)
 
 
+def cancel_on_event(
+    cancel_event: threading.Event,
+) -> None:
+    """
+    Cancel the function call if the cancel event is set.
+
+    Parameters
+    ----------
+    cancel_event : threading.Event
+        The event to check for cancellation
+
+    Raises
+    ------
+    asyncio.CancelledError
+        If the cancel event is set
+
+    """
+    if cancel_event.is_set():
+        task = asyncio.current_task()
+        if task is not None:
+            task.cancel()
+        raise asyncio.CancelledError(
+            'Function call was cancelled by client',
+        )
+
+
 def build_udf_endpoint(
     func: Callable[..., Any],
     returns_data_format: str,
@@ -273,12 +315,24 @@ def build_udf_endpoint(
     """
     if returns_data_format in ['scalar', 'list']:
 
+        is_async = asyncio.iscoroutinefunction(func)
+
         async def do_func(
+            cancel_event: threading.Event,
+            timer: Timer,
             row_ids: Sequence[int],
             rows: Sequence[Sequence[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))]
+            out = []
+            async with timer('call_function'):
+                for row in rows:
+                    cancel_on_event(cancel_event)
+                    if is_async:
+                        out.append(await func(*row))
+                    else:
+                        out.append(func(*row))
+            return row_ids, list(zip(out))
 
         return do_func
 
@@ -307,8 +361,11 @@ def build_vector_udf_endpoint(
     """
     masks = get_masked_params(func)
     array_cls = get_array_class(returns_data_format)
+    is_async = asyncio.iscoroutinefunction(func)
 
     async def do_func(
+        cancel_event: threading.Event,
+        timer: Timer,
         row_ids: Sequence[int],
         cols: Sequence[Tuple[Sequence[Any], Optional[Sequence[bool]]]],
     ) -> Tuple[
@@ -319,10 +376,19 @@ def build_vector_udf_endpoint(
         row_ids = array_cls(row_ids)
 
         # 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()
+        async with timer('call_function'):
+            if cols and cols[0]:
+                if is_async:
+                    out = await func(*[x if m else x[0] for x, m in zip(cols, masks)])
+                else:
+                    out = func(*[x if m else x[0] for x, m in zip(cols, masks)])
+            else:
+                if is_async:
+                    out = await func()
+                else:
+                    out = func()
+
+        cancel_on_event(cancel_event)
 
         # Single masked value
         if isinstance(out, Masked):
@@ -360,7 +426,11 @@ def build_tvf_endpoint(
     """
     if returns_data_format in ['scalar', 'list']:
 
+        is_async = asyncio.iscoroutinefunction(func)
+
         async def do_func(
+            cancel_event: threading.Event,
+            timer: Timer,
             row_ids: Sequence[int],
             rows: Sequence[Sequence[Any]],
         ) -> Tuple[Sequence[int], List[Tuple[Any, ...]]]:
@@ -368,9 +438,15 @@ def build_tvf_endpoint(
             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)))
+            async with timer('call_function'):
+                for i, row in zip(row_ids, rows):
+                    cancel_on_event(cancel_event)
+                    if is_async:
+                        res = await func(*row)
+                    else:
+                        res = func(*row)
+                    out.extend(as_list_of_tuples(res))
+                    out_ids.extend([row_ids[i]] * (len(out)-len(out_ids)))
             return out_ids, out
 
         return do_func
@@ -402,6 +478,8 @@ def build_vector_tvf_endpoint(
     array_cls = get_array_class(returns_data_format)
 
     async def do_func(
+        cancel_event: threading.Event,
+        timer: Timer,
         row_ids: Sequence[int],
         cols: Sequence[Tuple[Sequence[Any], Optional[Sequence[bool]]]],
     ) -> Tuple[
@@ -413,13 +491,28 @@ def build_vector_tvf_endpoint(
         #        each result row, so we just have to use the same
         #        row ID for all rows in the result.
 
+        is_async = asyncio.iscoroutinefunction(func)
+
         # 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())
+        async with timer('call_function'):
+            if cols and cols[0]:
+                if is_async:
+                    func_res = await func(
+                        *[x if m else x[0] for x, m in zip(cols, masks)],
+                    )
+                else:
+                    func_res = func(
+                        *[x if m else x[0] for x, m in zip(cols, masks)],
+                    )
+            else:
+                if is_async:
+                    func_res = await func()
+                else:
+                    func_res = func()
+
+        res = get_dataframe_columns(func_res)
+
+        cancel_on_event(cancel_event)
 
         # Generate row IDs
         if isinstance(res[0], Masked):
@@ -458,6 +551,10 @@ def make_func(
     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')
+    timeout = (
+        func._singlestoredb_attrs.get('timeout') or  # type: ignore
+        get_option('external_function.timeout')
+    )
 
     if function_type == 'tvf':
         do_func = build_tvf_endpoint(func, returns_data_format)
@@ -477,6 +574,12 @@ def make_func(
     # Set function type
     info['function_type'] = function_type
 
+    # Set timeout
+    info['timeout'] = max(timeout, 1)
+
+    # Set async flag
+    info['is_async'] = asyncio.iscoroutinefunction(func)
+
     # Setup argument types for rowdat_1 parser
     colspec = []
     for x in sig['args']:
@@ -498,6 +601,43 @@ def make_func(
     return do_func, info
 
 
+async def cancel_on_timeout(timeout: int) -> None:
+    """Cancel request if it takes too long."""
+    await asyncio.sleep(timeout)
+    raise asyncio.CancelledError(
+        'Function call was cancelled due to timeout',
+    )
+
+
+async def cancel_on_disconnect(
+    receive: Callable[..., Awaitable[Any]],
+) -> None:
+    """Cancel request if client disconnects."""
+    while True:
+        message = await receive()
+        if message['type'] == 'http.disconnect':
+            raise asyncio.CancelledError(
+                'Function call was cancelled by client',
+            )
+
+
+async def cancel_all_tasks(tasks: Iterable[asyncio.Task[Any]]) -> None:
+    """Cancel all tasks."""
+    for task in tasks:
+        task.cancel()
+    await asyncio.gather(*tasks, return_exceptions=True)
+
+
+def start_counter() -> float:
+    """Start a timer and return the start time."""
+    return time.perf_counter()
+
+
+def end_counter(start: float) -> float:
+    """End a timer and return the elapsed time."""
+    return time.perf_counter() - start
+
+
 class Application(object):
     """
     Create an external function application.
@@ -824,6 +964,21 @@ class Application(object):
             Function to send response information
 
         '''
+        request_id = str(uuid.uuid4())
+
+        timer = Timer(
+            id=request_id,
+            timestamp=datetime.datetime.now(
+                datetime.timezone.utc,
+            ).strftime('%Y-%m-%dT%H:%M:%S.%fZ'),
+        )
+        call_timer = Timer(
+            id=request_id,
+            timestamp=datetime.datetime.now(
+                datetime.timezone.utc,
+            ).strftime('%Y-%m-%dT%H:%M:%S.%fZ'),
+        )
+
         assert scope['type'] == 'http'
 
         method = scope['method']
@@ -838,6 +993,9 @@ class Application(object):
         func_name = headers.get(b's2-ef-name', b'')
         func_endpoint = self.endpoints.get(func_name)
 
+        timer.metadata['function'] = func_name.decode('utf-8') if func_name else ''
+        call_timer.metadata['function'] = timer.metadata['function']
+
         func = None
         func_info: Dict[str, Any] = {}
         if func_endpoint is not None:
@@ -845,35 +1003,120 @@ class Application(object):
 
         # Call the endpoint
         if method == 'POST' and func is not None and path == self.invoke_path:
+
+            logger.info(
+                json.dumps({
+                    'type': 'function_call',
+                    'id': request_id,
+                    'name': func_name.decode('utf-8'),
+                    'content_type': content_type.decode('utf-8'),
+                    'accepts': accepts.decode('utf-8'),
+                }),
+            )
+
             args_data_format = func_info['args_data_format']
             returns_data_format = func_info['returns_data_format']
             data = []
             more_body = True
-            while more_body:
-                request = await receive()
-                data.append(request['body'])
-                more_body = request.get('more_body', False)
+            with timer('receive_data'):
+                while more_body:
+                    request = await receive()
+                    if request['type'] == 'http.disconnect':
+                        raise RuntimeError('client disconnected')
+                    data.append(request['body'])
+                    more_body = request.get('more_body', False)
 
             data_version = headers.get(b's2-ef-version', b'')
             input_handler = self.handlers[(content_type, data_version, args_data_format)]
             output_handler = self.handlers[(accepts, data_version, returns_data_format)]
 
             try:
-                out = await func(
-                    *input_handler['load'](  # type: ignore
+                all_tasks = []
+                result = []
+
+                cancel_event = threading.Event()
+
+                with timer('parse_input'):
+                    inputs = input_handler['load'](  # type: ignore
                         func_info['colspec'], b''.join(data),
+                    )
+
+                func_task = asyncio.create_task(
+                    func(cancel_event, call_timer, *inputs)
+                    if func_info['is_async']
+                    else to_thread(
+                        lambda: asyncio.run(
+                            func(cancel_event, call_timer, *inputs),
+                        ),
                     ),
                 )
-                body = output_handler['dump'](
-                    [x[1] for x in func_info['returns']], *out,  # type: ignore
+                disconnect_task = asyncio.create_task(
+                    cancel_on_disconnect(receive),
                 )
+                timeout_task = asyncio.create_task(
+                    cancel_on_timeout(func_info['timeout']),
+                )
+
+                all_tasks += [func_task, disconnect_task, timeout_task]
+
+                async with timer('function_wrapper'):
+                    done, pending = await asyncio.wait(
+                        all_tasks, return_when=asyncio.FIRST_COMPLETED,
+                    )
+
+                await cancel_all_tasks(pending)
+
+                for task in done:
+                    if task is disconnect_task:
+                        cancel_event.set()
+                        raise asyncio.CancelledError(
+                            'Function call was cancelled by client disconnect',
+                        )
+
+                    elif task is timeout_task:
+                        cancel_event.set()
+                        raise asyncio.TimeoutError(
+                            'Function call was cancelled due to timeout',
+                        )
+
+                    elif task is func_task:
+                        result.extend(task.result())
+
+                with timer('format_output'):
+                    body = output_handler['dump'](
+                        [x[1] for x in func_info['returns']], *result,  # type: ignore
+                    )
+
                 await send(output_handler['response'])
 
+            except asyncio.TimeoutError:
+                logging.exception(
+                    'Timeout in function call: ' + func_name.decode('utf-8'),
+                )
+                body = (
+                    '[TimeoutError] Function call timed out after ' +
+                    str(func_info['timeout']) +
+                    ' seconds'
+                ).encode('utf-8')
+                await send(self.error_response_dict)
+
+            except asyncio.CancelledError:
+                logging.exception(
+                    'Function call cancelled: ' + func_name.decode('utf-8'),
+                )
+                body = b'[CancelledError] Function call was cancelled'
+                await send(self.error_response_dict)
+
             except Exception as e:
-                logging.exception('Error in function call')
+                logging.exception(
+                    'Error in function call: ' + func_name.decode('utf-8'),
+                )
                 body = f'[{type(e).__name__}] {str(e).strip()}'.encode('utf-8')
                 await send(self.error_response_dict)
 
+            finally:
+                await cancel_all_tasks(all_tasks)
+
         # Handle api reflection
         elif method == 'GET' and path == self.show_create_function_path:
             host = headers.get(b'host', b'localhost:80')
@@ -905,9 +1148,15 @@ class Application(object):
             await send(self.path_not_found_response_dict)
 
         # Send body
-        out = self.body_response_dict.copy()
-        out['body'] = body
-        await send(out)
+        with timer('send_response'):
+            out = self.body_response_dict.copy()
+            out['body'] = body
+            await send(out)
+
+        for k, v in call_timer.metrics.items():
+            timer.metrics[k] = v
+
+        timer.finish()
 
     def _create_link(
         self,
@@ -964,6 +1213,13 @@ class Application(object):
         functions = {}
         no_default = object()
 
+        # Generate CREATE FUNCTION SQL for each function using get_create_functions
+        create_sqls = self.get_create_functions(replace=True)
+        sql_map = {}
+        for (_, info), sql in zip(self.endpoints.values(), create_sqls):
+            sig = info['signature']
+            sql_map[sig['name']] = sql
+
         for key, (_, info) in self.endpoints.items():
             if not func_name or key == func_name:
                 sig = info['signature']
@@ -1001,8 +1257,12 @@ class Application(object):
                     if a.get('default', no_default) is not no_default:
                         returns[-1]['default'] = a['default']
 
+                sql = sql_map.get(sig['name'], '')
                 functions[sig['name']] = dict(
-                    args=args, returns=returns, function_type=info['function_type'],
+                    args=args,
+                    returns=returns,
+                    function_type=info['function_type'],
+                    sql_statement=sql,
                 )
 
         return functions