interactive

Author: kkampli-singlestore

singlestoredb/functions/ext/asgi.py

@@ -66,6 +66,8 @@
 from ..signature import signature_to_sql
 from ..typing import Masked
 from ..typing import Table
+from ...config import get_option
+
 
 try:
     import cloudpickle
@@ -273,24 +275,12 @@ def build_udf_endpoint(
     """
     if returns_data_format in ['scalar', 'list']:
 
-        is_async = asyncio.iscoroutinefunction(func)
-
         async def do_func(
             row_ids: Sequence[int],
             rows: Sequence[Sequence[Any]],
         ) -> Tuple[Sequence[int], List[Tuple[Any, ...]]]:
             '''Call function on given rows of data.'''
-            out = []
-            for row in rows:
-                if cancel_event.is_set():
-                    raise asyncio.CancelledError(
-                        'Function call was cancelled',
-                    )
-                if is_async:
-                    out.append(await func(*row))
-                else:
-                    out.append(func(*row))
-            return row_ids, list(zip(out))
+            return row_ids, [as_tuple(x) for x in zip(func_map(func, rows))]
 
         return do_func
 
@@ -319,7 +309,6 @@ 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(
         row_ids: Sequence[int],
@@ -332,17 +321,18 @@ async def do_func(
         row_ids = array_cls(row_ids)
 
         # Call the function with `cols` as the function parameters
+        is_async = inspect.iscoroutinefunction(func) or inspect.iscoroutinefunction(getattr(func, "__wrapped__", None))
         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)])
+                out = await asyncio.to_thread(func, *[x if m else x[0] for x, m in zip(cols, masks)])
         else:
             if is_async:
                 out = await func()
             else:
-                out = func()
-
+                out = await asyncio.to_thread(func())
+                
         # Single masked value
         if isinstance(out, Masked):
             return row_ids, [tuple(out)]
@@ -379,8 +369,6 @@ def build_tvf_endpoint(
     """
     if returns_data_format in ['scalar', 'list']:
 
-        is_async = asyncio.iscoroutinefunction(func)
-
         async def do_func(
             row_ids: Sequence[int],
             rows: Sequence[Sequence[Any]],
@@ -389,15 +377,7 @@ async def do_func(
             out_ids: List[int] = []
             out = []
             # Call function on each row of data
-            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)
+            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)))
             return out_ids, out
@@ -442,23 +422,13 @@ 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]:
-            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)]),
-                )
+            res = get_dataframe_columns(
+                func(*[x if m else x[0] for x, m in zip(cols, masks)]),
+            )
         else:
-            if is_async:
-                res = get_dataframe_columns(await func())
-            else:
-                res = get_dataframe_columns(func())
+            res = get_dataframe_columns(func())
 
         # Generate row IDs
         if isinstance(res[0], Masked):
@@ -516,9 +486,6 @@ def make_func(
     # Set function type
     info['function_type'] = function_type
 
-    # Set async flag
-    info['is_async'] = asyncio.iscoroutinefunction(func)
-
     # Setup argument types for rowdat_1 parser
     colspec = []
     for x in sig['args']:
@@ -901,64 +868,11 @@ async def __call__(
             output_handler = self.handlers[(accepts, data_version, returns_data_format)]
 
             try:
-                result = []
-
-                cancel_event = threading.Event()
-
-                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),
-                )
-                timeout_task = asyncio.create_task(
-                    cancel_on_timeout(func_info['timeout']),
-                )
-
-                all_tasks = [func_task, disconnect_task, timeout_task]
-
-                done, pending = await asyncio.wait(
-                    all_tasks, return_when=asyncio.FIRST_COMPLETED,
+                out = await func(
+                    *input_handler['load'](  # type: ignore
+                        func_info['colspec'], b''.join(data),
+                    ),
                 )
-
-                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())
-
-                print(result)
                 body = output_handler['dump'](
                     [x[1] for x in func_info['returns']], *out,  # type: ignore
                 )
@@ -1065,6 +979,11 @@ def get_function_info(
         for (_, info), sql in zip(self.endpoints.values(), create_sqls):
             sig = info['signature']
             sql_map[sig['name']] = sql
+        
+        connection_info = {}
+        workspace_group_id = os.environ.get('SINGLESTOREDB_WORKSPACE_GROUP')
+        connection_info['workspace_group_id'] = workspace_group_id
+
 
         for key, (_, info) in self.endpoints.items():
             if not func_name or key == func_name:
@@ -1111,7 +1030,10 @@ def get_function_info(
                     sql_statement=sql,
                 )
 
-        return functions
+        return {
+            'functions': functions,
+            'connection_info': connection_info
+        }
 
     def get_create_functions(
         self,