singlestoredb 1.15.0__py3-none-any.whl → 1.15.1__py3-none-any.whl

singlestoredb/docstring/util.py

@@ -0,0 +1,141 @@
+"""Utility functions for working with docstrings."""
+import typing as T
+from collections import ChainMap
+from inspect import Signature
+from itertools import chain
+
+from .common import DocstringMeta
+from .common import DocstringParam
+from .common import DocstringReturns  # noqa: F401
+from .common import DocstringStyle
+from .common import RenderingStyle
+from .parser import compose
+from .parser import parse
+
+_Func = T.Callable[..., T.Any]
+
+
+def combine_docstrings(
+    *others: _Func,
+    exclude: T.Iterable[T.Type[DocstringMeta]] = (),
+    style: DocstringStyle = DocstringStyle.AUTO,
+    rendering_style: RenderingStyle = RenderingStyle.COMPACT,
+) -> _Func:
+    """A function decorator that parses the docstrings from `others`,
+    programmatically combines them with the parsed docstring of the decorated
+    function, and replaces the docstring of the decorated function with the
+    composed result. Only parameters that are part of the decorated functions
+    signature are included in the combined docstring. When multiple sources for
+    a parameter or docstring metadata exists then the decorator will first
+    default to the wrapped function's value (when available) and otherwise use
+    the rightmost definition from ``others``.
+
+    The following example illustrates its usage:
+
+    >>> def fun1(a, b, c, d):
+    ...    '''short_description: fun1
+    ...
+    ...    :param a: fun1
+    ...    :param b: fun1
+    ...    :return: fun1
+    ...    '''
+    >>> def fun2(b, c, d, e):
+    ...    '''short_description: fun2
+    ...
+    ...    long_description: fun2
+    ...
+    ...    :param b: fun2
+    ...    :param c: fun2
+    ...    :param e: fun2
+    ...    '''
+    >>> @combine_docstrings(fun1, fun2)
+    >>> def decorated(a, b, c, d, e, f):
+    ...     '''
+    ...     :param e: decorated
+    ...     :param f: decorated
+    ...     '''
+    >>> print(decorated.__doc__)
+    short_description: fun2
+    <BLANKLINE>
+    long_description: fun2
+    <BLANKLINE>
+    :param a: fun1
+    :param b: fun1
+    :param c: fun2
+    :param e: fun2
+    :param f: decorated
+    :returns: fun1
+    >>> @combine_docstrings(fun1, fun2, exclude=[DocstringReturns])
+    >>> def decorated(a, b, c, d, e, f): pass
+    >>> print(decorated.__doc__)
+    short_description: fun2
+    <BLANKLINE>
+    long_description: fun2
+    <BLANKLINE>
+    :param a: fun1
+    :param b: fun1
+    :param c: fun2
+    :param e: fun2
+
+    :param others: callables from which to parse docstrings.
+    :param exclude: an iterable of ``DocstringMeta`` subclasses to exclude when
+        combining docstrings.
+    :param style: style composed docstring. The default will infer the style
+        from the decorated function.
+    :param rendering_style: The rendering style used to compose a docstring.
+    :return: the decorated function with a modified docstring.
+    """
+
+    def wrapper(func: _Func) -> _Func:
+        sig = Signature.from_callable(func)
+
+        comb_doc = parse(func.__doc__ or '')
+        docs = [parse(other.__doc__ or '') for other in others] + [comb_doc]
+        params = dict(
+            ChainMap(
+                *(
+                    {param.arg_name: param for param in doc.params}
+                    for doc in docs
+                ),
+            ),
+        )
+
+        for doc in reversed(docs):
+            if not doc.short_description:
+                continue
+            comb_doc.short_description = doc.short_description
+            comb_doc.blank_after_short_description = (
+                doc.blank_after_short_description
+            )
+            break
+
+        for doc in reversed(docs):
+            if not doc.long_description:
+                continue
+            comb_doc.long_description = doc.long_description
+            comb_doc.blank_after_long_description = (
+                doc.blank_after_long_description
+            )
+            break
+
+        combined: T.Dict[T.Type[DocstringMeta], T.List[DocstringMeta]] = {}
+        for doc in docs:
+            metas: T.Dict[T.Type[DocstringMeta], T.List[DocstringMeta]] = {}
+            for meta in doc.meta:
+                meta_type = type(meta)
+                if meta_type in exclude:
+                    continue
+                metas.setdefault(meta_type, []).append(meta)
+            for meta_type, meta_list in metas.items():
+                combined[meta_type] = meta_list
+
+        combined[DocstringParam] = [
+            params[name] for name in sig.parameters if name in params
+        ]
+        comb_doc.meta = list(chain(*combined.values()))
+        func.__doc__ = compose(
+            comb_doc, style=style, rendering_style=rendering_style,
+        )
+        return func
+
+    return wrapper

singlestoredb/functions/decorator.py

@@ -45,23 +45,20 @@ def is_valid_type(obj: Any) -> bool:
     return False
 
 
-def is_valid_callable(obj: Any) -> bool:
+def is_sqlstr_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):
+    if inspect.isclass(returns) and issubclass(returns, SQLString):
         return True
 
-    raise TypeError(
-        f'callable {obj} must return a str, '
-        f'but got {returns}',
-    )
+    return False
 
 
-def expand_types(args: Any) -> Optional[Union[List[str], Type[Any]]]:
+def expand_types(args: Any) -> Optional[List[Any]]:
     """Expand the types for the function arguments / return values."""
     if args is None:
         return None
@@ -70,28 +67,32 @@ def expand_types(args: Any) -> Optional[Union[List[str], Type[Any]]]:
     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 = []
+        new_args: List[Any] = []
         for arg in args:
             if isinstance(arg, str):
                 new_args.append(arg)
-            elif callable(arg):
+            elif is_sqlstr_callable(arg):
                 new_args.append(arg())
+            elif type(arg) is type:
+                new_args.append(arg)
+            elif is_valid_type(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]
+    elif is_sqlstr_callable(args):
+        return [args()]
+
+    # General way of accepting pydantic.BaseModel, NamedTuple, TypedDict
+    elif is_valid_type(args):
+        return [args]
+
+    elif type(args) is type:
+        return [args]
 
     raise TypeError(f'unrecognized type for parameter: {args}')
 

singlestoredb/functions/ext/asgi.py

@@ -73,6 +73,8 @@ 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
@@ -538,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
     -------
@@ -615,7 +619,7 @@ async def cancel_on_disconnect(
     """Cancel request if client disconnects."""
     while True:
         message = await receive()
-        if message['type'] == 'http.disconnect':
+        if message.get('type', '') == 'http.disconnect':
             raise asyncio.CancelledError(
                 'Function call was cancelled by client',
             )
@@ -674,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.
 
     """
 
@@ -816,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,
@@ -838,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(
@@ -944,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,
@@ -992,6 +1001,7 @@ 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']
@@ -1021,7 +1031,7 @@ class Application(object):
             with timer('receive_data'):
                 while more_body:
                     request = await receive()
-                    if request['type'] == 'http.disconnect':
+                    if request.get('type', '') == 'http.disconnect':
                         raise RuntimeError('client disconnected')
                     data.append(request['body'])
                     more_body = request.get('more_body', False)
@@ -1051,7 +1061,8 @@ class Application(object):
                     ),
                 )
                 disconnect_task = asyncio.create_task(
-                    cancel_on_disconnect(receive),
+                    asyncio.sleep(int(1e9))
+                    if ignore_cancel else cancel_on_disconnect(receive),
                 )
                 timeout_task = asyncio.create_task(
                     cancel_on_timeout(func_info['timeout']),
@@ -1130,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')
@@ -1142,6 +1154,11 @@ 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''
@@ -1184,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(
@@ -1220,20 +1244,66 @@ class Application(object):
             sig = info['signature']
             sql_map[sig['name']] = sql
 
-        for key, (_, info) in self.endpoints.items():
+        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:
@@ -1250,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):
@@ -1263,6 +1335,9 @@ class Application(object):
                     returns=returns,
                     function_type=info['function_type'],
                     sql_statement=sql,
+                    summary=doc_summary,
+                    long_description=doc_long_description,
+                    examples=doc_examples,
                 )
 
         return functions
@@ -1303,6 +1378,7 @@ class Application(object):
                     app_mode=self.app_mode,
                     replace=replace,
                     link=link or None,
+                    database=self.function_database or None,
                 ),
             )
 
@@ -1332,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):
@@ -1358,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}')
 
@@ -1415,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',
             },
         )
 
@@ -1679,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',
@@ -1778,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)