[backend:fix] Reject engine-unsupported UDF parameter and return types MCDB-100032 #closes

ARCHITECTURE.md

@@ -744,9 +744,6 @@ The `signature.py` module maps Python types to SQL types:
 | `str` | `TEXT` |
 | `bytes` | `BLOB` |
 | `bool` | `TINYINT` |
-| `datetime.datetime` | `DATETIME` |
-| `datetime.date` | `DATE` |
-| `decimal.Decimal` | `DECIMAL` |
 | `List[T]` | `ARRAY` |
 | `Optional[T]` | Nullable variant |
 

singlestoredb/functions/decorator.py

@@ -9,6 +9,7 @@
 
 from . import utils
 from .dtypes import SQLString
+from .signature import get_signature
 
 
 ParameterType = Union[
@@ -115,37 +116,27 @@ def _func(
         ).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: UDFType) -> UDFType:
+    def decorate(func: UDFType) -> UDFType:
+        if inspect.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
+            wrapped: UDFType = functools.wraps(func)(async_wrapper)
 
-            if inspect.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)
+        else:
+            def wrapper(*args: Any, **kwargs: Any) -> UDFType:
+                return func(*args, **kwargs)  # type: ignore
+            wrapper._singlestoredb_attrs = _singlestoredb_attrs  # type: ignore
+            wrapped = 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)
+        get_signature(wrapped)
 
+        return wrapped
+
+    if func is None:
         return decorate
 
-    if inspect.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)
-
-    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(func)
 
 
 def udf(

singlestoredb/functions/ext/asgi.py

@@ -71,6 +71,7 @@
 from ...config import get_option
 from ...management.stage import get_stage
 from ...mysql.constants import FIELD_TYPE as ft
+from ..signature import get_resolved_signature
 from ..signature import get_signature
 from ..signature import signature_to_sql
 from ..typing import Masked
@@ -509,7 +510,7 @@ def get_masked_params(func: Callable[..., Any]) -> List[bool]:
         Boolean list of masked parameters
 
     """
-    params = inspect.signature(func).parameters
+    params = get_resolved_signature(func).parameters
     return [typing.get_origin(x.annotation) is Masked for x in params.values()]
 
 

singlestoredb/functions/ext/plugin/registry.py

@@ -26,6 +26,7 @@
 from singlestoredb.functions.ext import rowdat_1 as _rowdat_1
 from singlestoredb.functions.ext.rowdat_1 import dump as _dump_rowdat_1
 from singlestoredb.functions.ext.rowdat_1 import load as _load_rowdat_1
+from singlestoredb.functions.signature import get_resolved_signature
 from singlestoredb.functions.signature import get_signature
 from singlestoredb.functions.typing import Masked
 from singlestoredb.mysql.constants import FIELD_TYPE as ft
@@ -496,7 +497,7 @@ def _register_function(
 
 def _get_masked_params(func: Callable[..., Any]) -> List[bool]:
     """Determine which parameters expect (data, mask) tuples vs just data."""
-    params = inspect.signature(func).parameters
+    params = get_resolved_signature(func).parameters
     return [typing.get_origin(x.annotation) is Masked for x in params.values()]
 
 

singlestoredb/functions/signature.py

@@ -142,12 +142,6 @@ class NoDefaultType:
     'str': 'TEXT',
     'bytes': 'BLOB',
     'null': 'NULL',
-    'datetime': 'DATETIME',
-    'datetime6': 'DATETIME(6)',
-    'date': 'DATE',
-    'time': 'TIME',
-    'time6': 'TIME(6)',
-    'decimal': 'DECIMAL',
 }
 
 sql_to_type_map = {
@@ -169,14 +163,6 @@ class NoDefaultType:
     'FLOAT': 'float32',
     'DOUBLE': 'float64',
     'REAL': 'float64',
-    'DATE': 'date',
-    'TIME': 'time',
-    'TIME(6)': 'time6',
-    'DATETIME': 'datetime',
-    'DATETIME(6)': 'datetime',
-    'TIMESTAMP': 'datetime',
-    'TIMESTAMP(6)': 'datetime',
-    'YEAR': 'uint64',
     'CHAR': 'str',
     'VARCHAR': 'str',
     'TEXT': 'str',
@@ -189,9 +175,36 @@ class NoDefaultType:
     'TINYBLOB': 'bytes',
     'MEDIUMBLOB': 'bytes',
     'LONGBLOB': 'bytes',
-    'DECIMAL': 'decimal',
-    'NEWDECIMAL': 'decimal',
-    'NUMERIC': 'decimal',
+}
+
+unsupported_python_types = (
+    datetime.datetime,
+    datetime.date,
+    datetime.time,
+    datetime.timedelta,
+    decimal.Decimal,
+)
+
+unsupported_dtypes = {
+    'datetime',
+    'datetime6',
+    'date',
+    'time',
+    'time6',
+    'decimal',
+}
+
+unsupported_sql_types = {
+    'DATE',
+    'TIME',
+    'DATETIME',
+    'TIMESTAMP',
+    'YEAR',
+    'DECIMAL',
+    'NEWDECIMAL',
+    'NUMERIC',
+    'DEC',
+    'FIXED',
 }
 
 
@@ -424,22 +437,11 @@ def normalize_dtype(dtype: Any) -> str:
     if issubclass(dtype, (bytes, bytearray)):
         return 'bytes'
 
-    # Decimal
-    if issubclass(dtype, decimal.Decimal):
-        return 'decimal'
-
-    # Date / Times
-    if issubclass(dtype, datetime.datetime):
-        return 'datetime'
-    if issubclass(dtype, datetime.date):
-        return 'date'
-    if issubclass(dtype, datetime.time):
+    if issubclass(dtype, unsupported_python_types):
         raise TypeError(
-            'datetime.time is not supported for UDF annotations; '
-            'use datetime.timedelta instead',
+            f'{dtype.__module__}.{dtype.__qualname__} is not supported '
+            'for UDF annotations',
         )
-    if issubclass(dtype, datetime.timedelta):
-        return 'time'
 
     # Last resort, guess it by the name...
     name = dtype.__name__.lower()
@@ -572,10 +574,11 @@ def get_dataclass_schema(obj: Any) -> List[ParamSpec]:
         A list of parameter specifications for the dataclass fields
 
     """
+    hints = utils.get_annotations(obj)
     return [
         ParamSpec(
             name=f.name,
-            dtype=f.type,
+            dtype=hints[f.name],
             default=NO_DEFAULT if f.default is dataclasses.MISSING else f.default,
         )
         for f in dataclasses.fields(obj)
@@ -1178,6 +1181,47 @@ def vector_check(obj: Any) -> Tuple[Any, str]:
     return obj, 'scalar'
 
 
+def _is_dtype_str(annotation: str) -> bool:
+    """Is the string a dtype string (e.g. 'int64?', 'array[str]') not Python code?"""
+    dtype = annotation[:-1] if annotation.endswith('?') else annotation
+    dtype = dtype.replace('|null', '')
+    return (
+        dtype in sql_type_map
+        or dtype.startswith(('array[', 'tuple['))
+    )
+
+
+def _resolve_annotation(annotation: Any, func: Callable[..., Any]) -> Any:
+    """
+    Evaluate a string annotation in the function's module namespace.
+
+    String annotations come from ``from __future__ import annotations`` or
+    quoted forward references. Dtype strings are left as-is.
+
+    """
+    if not isinstance(annotation, str) or _is_dtype_str(annotation):
+        return annotation
+    globalns = getattr(inspect.unwrap(func), '__globals__', {})
+    try:
+        return eval(annotation, globalns)
+    except Exception as exc:
+        raise TypeError(
+            f'could not resolve type annotation {annotation!r}: {exc}',
+        ) from exc
+
+
+def get_resolved_signature(func: Callable[..., Any]) -> inspect.Signature:
+    """Return the signature of ``func`` with string annotations evaluated."""
+    signature = inspect.signature(func)
+    return signature.replace(
+        parameters=[
+            p.replace(annotation=_resolve_annotation(p.annotation, func))
+            for p in signature.parameters.values()
+        ],
+        return_annotation=_resolve_annotation(signature.return_annotation, func),
+    )
+
+
 def get_masks(func: Callable[..., Any]) -> Tuple[List[bool], List[bool]]:
     """
     Get the list of masked parameters and return values for the function.
@@ -1195,8 +1239,9 @@ def get_masks(func: Callable[..., Any]) -> Tuple[List[bool], List[bool]]:
 
 
     """
-    params = inspect.signature(func).parameters
-    returns = inspect.signature(func).return_annotation
+    signature = get_resolved_signature(func)
+    params = signature.parameters
+    returns = signature.return_annotation
 
     ret_masks = []
     if typing.get_origin(returns) is Masked:
@@ -1235,7 +1280,7 @@ def get_signature(
     Dict[str, Any]
 
     '''
-    signature = inspect.signature(func)
+    signature = get_resolved_signature(func)
     args: List[Dict[str, Any]] = []
     returns: List[Dict[str, Any]] = []
 
@@ -1251,8 +1296,6 @@ def get_signature(
         elif p.kind == inspect.Parameter.VAR_KEYWORD:
             raise TypeError('variable keyword arguments are not supported')
 
-    # TODO: Use typing.get_type_hints() for parameters / return values?
-
     # Generate the parameter type and the corresponding SQL code for that parameter
     args_schema: List[ParamSpec] = []
     args_data_formats = []
@@ -1429,6 +1472,9 @@ def sql_to_dtype(sql: str) -> str:
     elif ' UNSIGNED' in sql:
         sql_type += ' UNSIGNED'
 
+    if sql_type in unsupported_sql_types:
+        raise TypeError(f'{sql_type} is not supported for UDF annotations')
+
     try:
         dtype = sql_to_type_map[sql_type]
     except KeyError:
@@ -1515,7 +1561,15 @@ def dtype_to_sql(
                 f'TABLE({", ".join(item_dtypes)}){nullable}{default_clause}',
             )
 
-    return f'{sql_type_map[dtype]}{nullable}{default_clause}'
+    if dtype in unsupported_dtypes:
+        raise TypeError(f'{dtype} is not supported for UDF annotations')
+
+    try:
+        sql_type = sql_type_map[dtype]
+    except KeyError:
+        raise TypeError(f'unrecognized data type: {dtype}')
+
+    return f'{sql_type}{nullable}{default_clause}'
 
 
 def signature_to_sql(