singlestoredb 1.14.2__cp38-abi3-win_amd64.whl → 1.15.1__cp38-abi3-win_amd64.whl

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,9 @@ from ..signature import get_signature
 from ..signature import signature_to_sql
 from ..typing import Masked
 from ..typing import Table
+from .timer import Timer
+from singlestoredb.docstring.parser import parse
+from singlestoredb.functions.dtypes import escape_name
 
 try:
     import cloudpickle
@@ -95,6 +104,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 +269,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 +317,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 +363,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 +378,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 +428,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 +440,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 +480,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 +493,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):
@@ -445,6 +540,8 @@ def make_func(
         Name of the function to create
     func : Callable
         The function to call as the endpoint
+    database : str, optional
+        The database to use for the function definition
 
     Returns
     -------
@@ -458,6 +555,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 +578,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 +605,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.get('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.
@@ -534,6 +678,8 @@ class Application(object):
     link_credentials : Dict[str, Any], optional
         The CREDENTIALS section of a LINK definition. This dictionary gets
         converted to JSON for the CREATE LINK call.
+    function_database : str, optional
+        The database to use for external function definitions.
 
     """
 
@@ -676,6 +822,7 @@ class Application(object):
     invoke_path = ('invoke',)
     show_create_function_path = ('show', 'create_function')
     show_function_info_path = ('show', 'function_info')
+    status = ('status',)
 
     def __init__(
         self,
@@ -698,6 +845,7 @@ class Application(object):
         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'),
+        function_database: Optional[str] = None,
     ) -> None:
         if link_name and (link_config or link_credentials):
             raise ValueError(
@@ -804,6 +952,7 @@ class Application(object):
         self.link_credentials = link_credentials
         self.endpoints = endpoints
         self.external_functions = external_functions
+        self.function_database = function_database
 
     async def __call__(
         self,
@@ -824,6 +973,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']
@@ -837,6 +1001,10 @@ class Application(object):
         accepts = headers.get(b'accepts', content_type)
         func_name = headers.get(b's2-ef-name', b'')
         func_endpoint = self.endpoints.get(func_name)
+        ignore_cancel = headers.get(b's2-ef-ignore-cancel', b'false') == b'true'
+
+        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] = {}
@@ -845,35 +1013,121 @@ 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.get('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(
+                    asyncio.sleep(int(1e9))
+                    if ignore_cancel else 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')
@@ -887,6 +1141,7 @@ class Application(object):
                             endpoint_info['signature'],
                             url=self.url or reflected_url,
                             data_format=self.data_format,
+                            database=self.function_database or None,
                         ),
                     )
             body = '\n'.join(syntax).encode('utf-8')
@@ -899,15 +1154,26 @@ class Application(object):
             body = json.dumps(dict(functions=functions)).encode('utf-8')
             await send(self.text_response_dict)
 
+        # Return status
+        elif method == 'GET' and path == self.status:
+            body = json.dumps(dict(status='ok')).encode('utf-8')
+            await send(self.text_response_dict)
+
         # Path not found
         else:
             body = b''
             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,
@@ -935,20 +1201,27 @@ class Application(object):
     def _locate_app_functions(self, cur: Any) -> Tuple[Set[str], Set[str]]:
         """Locate all current functions and links belonging to this app."""
         funcs, links = set(), set()
-        cur.execute('SHOW FUNCTIONS')
+        if self.function_database:
+            database_prefix = escape_name(self.function_database) + '.'
+            cur.execute(f'SHOW FUNCTIONS IN {escape_name(self.function_database)}')
+        else:
+            database_prefix = ''
+            cur.execute('SHOW FUNCTIONS')
+
         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}`')
+            cur.execute(f'SHOW CREATE FUNCTION {database_prefix}{escape_name(name)}')
             for fname, _, code, *_ in list(cur):
                 m = re.search(r" (?:\w+) (?:SERVICE|MANAGED) '([^']+)'", code)
                 if m and m.group(1) == self.url:
-                    funcs.add(fname)
+                    funcs.add(f'{database_prefix}{escape_name(fname)}')
                     if link and re.match(r'^py_ext_func_link_\S{14}$', link):
                         links.add(link)
+
         return funcs, links
 
     def get_function_info(
@@ -964,20 +1237,73 @@ class Application(object):
         functions = {}
         no_default = object()
 
-        for key, (_, info) in self.endpoints.items():
+        # 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, (func, info) in self.endpoints.items():
+            # Get info from docstring
+            doc_summary = ''
+            doc_long_description = ''
+            doc_params = {}
+            doc_returns = None
+            doc_examples = []
+            if func.__doc__:
+                try:
+                    docs = parse(func.__doc__)
+                    doc_params = {p.arg_name: p for p in docs.params}
+                    doc_returns = docs.returns
+                    if not docs.short_description and docs.long_description:
+                        doc_summary = docs.long_description or ''
+                    else:
+                        doc_summary = docs.short_description or ''
+                        doc_long_description = docs.long_description or ''
+                    for ex in docs.examples:
+                        ex_dict: Dict[str, Any] = {
+                            'description': None,
+                            'code': None,
+                            'output': None,
+                        }
+                        if ex.description:
+                            ex_dict['description'] = ex.description
+                        if ex.snippet:
+                            code, output = [], []
+                            for line in ex.snippet.split('\n'):
+                                line = line.rstrip()
+                                if re.match(r'^(\w+>|>>>|\.\.\.)', line):
+                                    code.append(line)
+                                else:
+                                    output.append(line)
+                            ex_dict['code'] = '\n'.join(code) or None
+                            ex_dict['output'] = '\n'.join(output) or None
+                        if ex.post_snippet:
+                            ex_dict['postscript'] = ex.post_snippet
+                        doc_examples.append(ex_dict)
+
+                except Exception as e:
+                    logger.warning(
+                        f'Could not parse docstring for function {key}: {e}',
+                    )
+
             if not func_name or key == func_name:
                 sig = info['signature']
                 args = []
 
                 # Function arguments
-                for a in sig.get('args', []):
+                for i, a in enumerate(sig.get('args', [])):
+                    name = a['name']
                     dtype = a['dtype']
                     nullable = '?' in dtype
                     args.append(
                         dict(
-                            name=a['name'],
+                            name=name,
                             dtype=dtype.replace('?', ''),
                             nullable=nullable,
+                            description=(doc_params[name].description or '')
+                            if name in doc_params else '',
                         ),
                     )
                     if a.get('default', no_default) is not no_default:
@@ -994,6 +1320,8 @@ class Application(object):
                         dict(
                             dtype=dtype.replace('?', ''),
                             nullable=nullable,
+                            description=doc_returns.description
+                            if doc_returns else '',
                         ),
                     )
                     if a.get('name', None):
@@ -1001,8 +1329,15 @@ 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,
+                    summary=doc_summary,
+                    long_description=doc_long_description,
+                    examples=doc_examples,
                 )
 
         return functions
@@ -1043,6 +1378,7 @@ class Application(object):
                     app_mode=self.app_mode,
                     replace=replace,
                     link=link or None,
+                    database=self.function_database or None,
                 ),
             )
 
@@ -1072,7 +1408,7 @@ class Application(object):
                 if replace:
                     funcs, links = self._locate_app_functions(cur)
                     for fname in funcs:
-                        cur.execute(f'DROP FUNCTION IF EXISTS `{fname}`')
+                        cur.execute(f'DROP FUNCTION IF EXISTS {fname}')
                     for link in links:
                         cur.execute(f'DROP LINK {link}')
                 for func in self.get_create_functions(replace=replace):
@@ -1098,7 +1434,7 @@ class Application(object):
             with conn.cursor() as cur:
                 funcs, links = self._locate_app_functions(cur)
                 for fname in funcs:
-                    cur.execute(f'DROP FUNCTION IF EXISTS `{fname}`')
+                    cur.execute(f'DROP FUNCTION IF EXISTS {fname}')
                 for link in links:
                     cur.execute(f'DROP LINK {link}')
 
@@ -1155,6 +1491,7 @@ class Application(object):
                 b'accepts': accepts[data_format.lower()],
                 b's2-ef-name': name.encode('utf-8'),
                 b's2-ef-version': data_version.encode('utf-8'),
+                b's2-ef-ignore-cancel': b'true',
             },
         )
 
@@ -1419,6 +1756,14 @@ def main(argv: Optional[List[str]] = None) -> None:
             ),
             help='Suffix to add to function names',
         )
+        parser.add_argument(
+            '--function-database', metavar='function_database',
+            default=defaults.get(
+                'function_database',
+                get_option('external_function.function_database'),
+            ),
+            help='Database to use for the function definition',
+        )
         parser.add_argument(
             'functions', metavar='module.or.func.path', nargs='*',
             help='functions or modules to export in UDF server',
@@ -1518,6 +1863,7 @@ def main(argv: Optional[List[str]] = None) -> None:
         app_mode='remote',
         name_prefix=args.name_prefix,
         name_suffix=args.name_suffix,
+        function_database=args.function_database or None,
     )
 
     funcs = app.get_create_functions(replace=args.replace_existing)