Merge branch 'cancel_udf' of https://github.com/singlestore-labs/singlestoredb-python into users/kaushik/udfs-temp

Author: kkampli-singlestore

.github/workflows/code-check.yml

@@ -1,4 +1,4 @@
-name: Code checks
+name: Coverage tests
 
 on:
   push:
@@ -7,7 +7,6 @@ on:
     branches: [ main ]
   workflow_dispatch:
 
-
 jobs:
   test-coverage:
     runs-on: ubuntu-latest
@@ -28,7 +27,7 @@ jobs:
       - name: Checkout code
         uses: actions/checkout@v4
         with:
-          fetch-depth: 2
+            fetch-depth: 0
 
       - name: Set up Python
         uses: actions/setup-python@v4
@@ -61,7 +60,7 @@ jobs:
             # For pushes to main/master, compare against previous commit
             BASE_COMMIT="HEAD~1"
             echo "Push to main/master: Comparing against $BASE_COMMIT"
-          else
+          else:
             # For pushes to other branches, compare against master/main
             if git rev-parse --verify origin/main >/dev/null 2>&1; then
               BASE_COMMIT="origin/main"

.github/workflows/coverage.yml

@@ -23,6 +23,8 @@ jobs:
 
     steps:
       - uses: actions/checkout@v4
+        with:
+            fetch-depth: 0
 
       - name: Set up Python
         uses: actions/setup-python@v4

singlestoredb/config.py

@@ -438,6 +438,11 @@
     environ=['SINGLESTOREDB_EXT_FUNC_PORT'],
 )
 
+register_option(
+    'external_function.timeout', 'int', check_int, 24*60*60,
+    'Specifies the timeout in seconds for processing a batch of rows.',
+    environ=['SINGLESTOREDB_EXT_FUNC_TIMEOUT'],
+)
 
 #
 # Debugging options

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:
@@ -100,38 +102,50 @@ def _func(
     name: Optional[str] = None,
     args: Optional[ParameterType] = None,
     returns: Optional[ReturnType] = None,
-) -> Callable[..., Any]:
+    timeout: Optional[int] = None,
+) -> UDFType:
     """Generic wrapper for UDF and TVF decorators."""
 
     _singlestoredb_attrs = {  # type: ignore
         k: v for k, v in dict(
             name=name,
             args=expand_types(args),
             returns=expand_types(returns),
+            timeout=timeout,
         ).items() if v is not None
     }
 
     # No func was specified, this is an uncalled decorator that will get
     # 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(
@@ -140,7 +154,8 @@ def udf(
     name: Optional[str] = None,
     args: Optional[ParameterType] = None,
     returns: Optional[ReturnType] = None,
-) -> Callable[..., Any]:
+    timeout: Optional[int] = None,
+) -> UDFType:
     """
     Define a user-defined function (UDF).
 
@@ -167,6 +182,9 @@ def udf(
         Specifies the return data type of the function. This parameter
         works the same way as `args`. If the function is a table-valued
         function, the return type should be a `Table` object.
+    timeout : int, optional
+        The timeout in seconds for the UDF execution. If not specified,
+        the global default timeout is used.
 
     Returns
     -------
@@ -178,4 +196,5 @@ def udf(
         name=name,
         args=args,
         returns=returns,
+        timeout=timeout,
     )