Add async support

Author: kesmit13

singlestoredb/functions/decorator.py

@@ -1,3 +1,4 @@
+import asyncio
 import functools
 import inspect
 from typing import Any
@@ -19,6 +20,7 @@
 ]
 
 ReturnType = ParameterType
+UDFType = Callable[..., Any]
 
 
 def is_valid_type(obj: Any) -> bool:
@@ -101,7 +103,7 @@ def _func(
     args: Optional[ParameterType] = None,
     returns: Optional[ReturnType] = None,
     timeout: Optional[int] = None,
-) -> Callable[..., Any]:
+) -> UDFType:
     """Generic wrapper for UDF and TVF decorators."""
 
     _singlestoredb_attrs = {  # type: ignore
@@ -117,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(
@@ -143,7 +155,7 @@ def udf(
     args: Optional[ParameterType] = None,
     returns: Optional[ReturnType] = None,
     timeout: Optional[int] = None,
-) -> Callable[..., Any]:
+) -> UDFType:
     """
     Define a user-defined function (UDF).
 

singlestoredb/functions/ext/asgi.py

@@ -285,6 +285,8 @@ def build_udf_endpoint(
     """
     if returns_data_format in ['scalar', 'list']:
 
+        is_async = asyncio.iscoroutinefunction(func)
+
         async def do_func(
             cancel_event: threading.Event,
             row_ids: Sequence[int],
@@ -297,7 +299,10 @@ async def do_func(
                     raise asyncio.CancelledError(
                         'Function call was cancelled',
                     )
-                out.append(func(*row))
+                if is_async:
+                    out.append(await func(*row))
+                else:
+                    out.append(func(*row))
             return row_ids, list(zip(out))
 
         return do_func
@@ -327,6 +332,7 @@ 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,
@@ -341,9 +347,15 @@ async def do_func(
 
         # 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)])
+            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:
-            out = func()
+            if is_async:
+                out = await func()
+            else:
+                out = func()
 
         # Single masked value
         if isinstance(out, Masked):
@@ -381,6 +393,8 @@ def build_tvf_endpoint(
     """
     if returns_data_format in ['scalar', 'list']:
 
+        is_async = asyncio.iscoroutinefunction(func)
+
         async def do_func(
             cancel_event: threading.Event,
             row_ids: Sequence[int],
@@ -390,11 +404,15 @@ async def do_func(
             out_ids: List[int] = []
             out = []
             # Call function on each row of data
-            for i, res in zip(row_ids, func_map(func, rows)):
+            for i, row in zip(row_ids, rows):
                 if cancel_event.is_set():
                     raise asyncio.CancelledError(
                         'Function call was cancelled',
                     )
+                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
@@ -440,13 +458,23 @@ async def do_func(
         #        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)]),
-            )
+            if is_async:
+                res = get_dataframe_columns(
+                    await func(*[x if m else x[0] for x, m in zip(cols, masks)]),
+                )
+            else:
+                res = get_dataframe_columns(
+                    func(*[x if m else x[0] for x, m in zip(cols, masks)]),
+                )
         else:
-            res = get_dataframe_columns(func())
+            if is_async:
+                res = get_dataframe_columns(await func())
+            else:
+                res = get_dataframe_columns(func())
 
         # Generate row IDs
         if isinstance(res[0], Masked):
@@ -508,6 +536,9 @@ def make_func(
     # 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']:
@@ -927,18 +958,28 @@ async def __call__(
 
                 cancel_event = threading.Event()
 
-                func_task = asyncio.create_task(
-                    to_thread(
-                        lambda: asyncio.run(
-                            func(
-                                cancel_event,
-                                *input_handler['load'](  # type: ignore
-                                    func_info['colspec'], b''.join(data),
+                if func_info['is_async']:
+                    func_task = asyncio.create_task(
+                        func(
+                            cancel_event,
+                            *input_handler['load'](  # type: ignore
+                                func_info['colspec'], b''.join(data),
+                            ),
+                        ),
+                    )
+                else:
+                    func_task = asyncio.create_task(
+                        to_thread(
+                            lambda: asyncio.run(
+                                func(
+                                    cancel_event,
+                                    *input_handler['load'](  # type: ignore
+                                        func_info['colspec'], b''.join(data),
+                                    ),
                                 ),
                             ),
                         ),
-                    ),
-                )
+                    )
                 disconnect_task = asyncio.create_task(
                     cancel_on_disconnect(receive),
                 )
@@ -970,6 +1011,7 @@ async def __call__(
                     elif task is func_task:
                         result.extend(task.result())
 
+                print(result)
                 body = output_handler['dump'](
                     [x[1] for x in func_info['returns']], *result,  # type: ignore
                 )