Merge branch 'udf-apis' into users/snarayanan/udfs

Author: srinathnarayanan

singlestoredb/config.py

@@ -407,6 +407,18 @@
     environ=['SINGLESTOREDB_EXT_FUNC_LOG_LEVEL'],
 )
 
+register_option(
+    'external_function.name_prefix', 'string', check_str, '',
+    'Prefix to add to external function names.',
+    environ=['SINGLESTOREDB_EXT_FUNC_NAME_PREFIX'],
+)
+
+register_option(
+    'external_function.name_suffix', 'string', check_str, '',
+    'Suffix to add to external function names.',
+    environ=['SINGLESTOREDB_EXT_FUNC_NAME_SUFFIX'],
+)
+
 register_option(
     'external_function.connection', 'string', check_str,
     os.environ.get('SINGLESTOREDB_URL') or None,

singlestoredb/functions/ext/asgi.py

@@ -468,6 +468,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,
@@ -488,6 +489,8 @@ def __init__(
         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(
@@ -544,6 +547,7 @@ def __init__(
                         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
@@ -559,6 +563,7 @@ def __init__(
                     # 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
@@ -571,12 +576,14 @@ def __init__(
                     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
 
             else:
                 alias = funcs.__name__
+                alias = f'{name_prefix}{alias}{name_suffix}'
                 external_functions[funcs.__name__] = funcs
                 func, info = make_func(alias, funcs)
                 endpoints[alias.encode('utf-8')] = func, info
@@ -671,6 +678,12 @@ async def __call__(
 
             await send(self.text_response_dict)
 
+        # Return function info
+        elif method == 'GET' and path == self.show_function_info_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''
@@ -715,14 +728,77 @@ def _locate_app_functions(self, cur: Any) -> Tuple[Set[str], Set[str]]:
             # 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]
+
+        """
+        returns: Dict[str, Any] = {}
+        functions = {}
+
+        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,
+                        ),
+                    )
+
+                # Record / table return types
+                if sig['returns']['dtype'].startswith('tuple['):
+                    fields = []
+                    dtypes = sig['returns']['dtype'][6:-1].split(',')
+                    field_names = sig['returns']['field_names']
+                    for i, dtype in enumerate(dtypes):
+                        nullable = '?' in dtype
+                        dtype = dtype.replace('?', '')
+                        fields.append(
+                            dict(
+                                name=field_names[i],
+                                dtype=dtype,
+                                nullable=nullable,
+                            ),
+                        )
+                    returns = dict(
+                        dtype='table' if info['function_type'] == 'tvf' else 'struct',
+                        fields=fields,
+                    )
+
+                # Atomic return types
+                else:
+                    returns = dict(
+                        dtype=sig['returns'].get('dtype').replace('?', ''),
+                        nullable='?' in sig['returns'].get('dtype', ''),
+                    )
+
+                functions[sig['name']] = dict(args=args, returns=returns)
+
+        return functions
+
+    def get_create_functions(
         self,
         replace: bool = False,
     ) -> List[str]:
@@ -790,7 +866,7 @@ def register_functions(
                         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(
@@ -1118,6 +1194,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',
@@ -1210,9 +1302,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')
 

singlestoredb/functions/ext/mmap.py

@@ -338,7 +338,7 @@ def main(argv: Optional[List[str]] = None) -> None:
         app_mode='collocated',
     )
 
-    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')
 

singlestoredb/functions/signature.py

@@ -487,13 +487,15 @@ def get_signature(func: Callable[..., Any], name: Optional[str] = None) -> Dict[
             'missing annotations for {} in {}'
             .format(', '.join(spec_diff), name),
         )
+
     elif isinstance(args_overrides, dict):
         for s in spec_diff:
             if s not in args_overrides:
                 raise TypeError(
                     'missing annotations for {} in {}'
                     .format(', '.join(spec_diff), name),
                 )
+
     elif isinstance(args_overrides, list):
         if len(arg_names) != len(args_overrides):
             raise TypeError(
@@ -502,6 +504,7 @@ def get_signature(func: Callable[..., Any], name: Optional[str] = None) -> Dict[
             )
 
     for i, arg in enumerate(arg_names):
+
         if isinstance(args_overrides, list):
             sql = args_overrides[i]
             arg_type = sql_to_dtype(sql)
@@ -528,6 +531,7 @@ def get_signature(func: Callable[..., Any], name: Optional[str] = None) -> Dict[
     if isinstance(returns_overrides, str):
         sql = returns_overrides
         out_type = sql_to_dtype(sql)
+
     elif isinstance(returns_overrides, list):
         if not output_fields:
             output_fields = [
@@ -540,29 +544,35 @@ def get_signature(func: Callable[..., Any], name: Optional[str] = None) -> Dict[
         sql = dtype_to_sql(
             out_type, function_type=function_type, field_names=output_fields,
         )
+
     elif dataclasses.is_dataclass(returns_overrides):
         out_type = collapse_dtypes([
             classify_dtype(x)
             for x in simplify_dtype([x.type for x in returns_overrides.fields])
         ])
+        output_fields = [x.name for x in returns_overrides.fields]
         sql = dtype_to_sql(
             out_type,
             function_type=function_type,
-            field_names=[x.name for x in returns_overrides.fields],
+            field_names=output_fields,
         )
+
     elif has_pydantic and inspect.isclass(returns_overrides) \
             and issubclass(returns_overrides, pydantic.BaseModel):
         out_type = collapse_dtypes([
             classify_dtype(x)
             for x in simplify_dtype([x for x in returns_overrides.model_fields.values()])
         ])
+        output_fields = [x for x in returns_overrides.model_fields.keys()]
         sql = dtype_to_sql(
             out_type,
             function_type=function_type,
-            field_names=[x for x in returns_overrides.model_fields.keys()],
+            field_names=output_fields,
         )
+
     elif returns_overrides is not None and not isinstance(returns_overrides, str):
         raise TypeError(f'unrecognized type for return value: {returns_overrides}')
+
     else:
         if not output_fields:
             if dataclasses.is_dataclass(signature.return_annotation):
@@ -572,13 +582,26 @@ def get_signature(func: Callable[..., Any], name: Optional[str] = None) -> Dict[
             elif has_pydantic and inspect.isclass(signature.return_annotation) \
                     and issubclass(signature.return_annotation, pydantic.BaseModel):
                 output_fields = list(signature.return_annotation.model_fields.keys())
+
         out_type = collapse_dtypes([
             classify_dtype(x) for x in simplify_dtype(signature.return_annotation)
         ])
+
+        if not output_fields:
+            output_fields = [
+                string.ascii_letters[i] for i in range(out_type.count(',')+1)
+            ]
+
         sql = dtype_to_sql(
             out_type, function_type=function_type, field_names=output_fields,
         )
-    out['returns'] = dict(dtype=out_type, sql=sql, default=None)
+
+    out['returns'] = dict(
+                         dtype=out_type,
+                         sql=sql,
+                         default=None,
+                         field_names=output_fields,
+    )
 
     copied_keys = ['database', 'environment', 'packages', 'resources', 'replace']
     for key in copied_keys: