From 82e7df67daa6dc4e4735aa420e52e0ffab1663de Mon Sep 17 00:00:00 2001 From: Shubham Kumar Date: Wed, 4 Nov 2026 12:28:40 +0530 Subject: [PATCH] Reject engine-unsupported UDF parameter and return types Reject date, time, datetime, timestamp, year, and decimal annotations and SQL type overrides at both @udf decoration time and signature extraction before emitting DDL that SingleStore rejects with ERROR 1706. Also resolve postponed string annotations in function signatures, dataclass schemas, and masked-parameter detection. --- ARCHITECTURE.md | 3 - singlestoredb/functions/decorator.py | 43 +++--- singlestoredb/functions/ext/asgi.py | 3 +- .../functions/ext/plugin/registry.py | 3 +- singlestoredb/functions/signature.py | 130 +++++++++++++----- singlestoredb/tests/test_udf.py | 93 ++++++++++--- .../tests/test_udf_future_annotations.py | 113 +++++++++++++++ singlestoredb/tests/test_udf_returns.py | 25 +--- 8 files changed, 304 insertions(+), 109 deletions(-) create mode 100644 singlestoredb/tests/test_udf_future_annotations.py diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index 499b65550..85e170214 100644 --- a/ARCHITECTURE.md +++ b/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 | diff --git a/singlestoredb/functions/decorator.py b/singlestoredb/functions/decorator.py index 5239f27f6..72dee8839 100644 --- a/singlestoredb/functions/decorator.py +++ b/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( diff --git a/singlestoredb/functions/ext/asgi.py b/singlestoredb/functions/ext/asgi.py index 94d95e31e..99b2588f1 100755 --- a/singlestoredb/functions/ext/asgi.py +++ b/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()] diff --git a/singlestoredb/functions/ext/plugin/registry.py b/singlestoredb/functions/ext/plugin/registry.py index fa5096803..f33d8ead5 100644 --- a/singlestoredb/functions/ext/plugin/registry.py +++ b/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()] diff --git a/singlestoredb/functions/signature.py b/singlestoredb/functions/signature.py index 346116b4f..737fa2bf9 100644 --- a/singlestoredb/functions/signature.py +++ b/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( diff --git a/singlestoredb/tests/test_udf.py b/singlestoredb/tests/test_udf.py index ebb0eae0b..a0e2a44d4 100755 --- a/singlestoredb/tests/test_udf.py +++ b/singlestoredb/tests/test_udf.py @@ -3,6 +3,7 @@ """SingleStoreDB UDF testing.""" import dataclasses import datetime +import decimal import re import unittest from typing import List @@ -232,23 +233,76 @@ def foo(x: set) -> None: ... def test_datetimes(self): - # Datetime - def foo(x: datetime.datetime) -> None: ... - assert to_sql(foo) == '`foo`(`x` DATETIME NOT NULL) RETURNS TINYINT NULL' + for py_type, expected_name in [ + (datetime.datetime, 'datetime.datetime'), + (datetime.date, 'datetime.date'), + (datetime.time, 'datetime.time'), + (datetime.timedelta, 'datetime.timedelta'), + (decimal.Decimal, 'decimal.Decimal'), + ]: + def foo_param(x: py_type) -> None: ... + with self.assertRaisesRegex(TypeError, expected_name): + to_sql(foo_param) + + def foo_opt_param(x: Optional[py_type]) -> None: ... + with self.assertRaisesRegex(TypeError, expected_name): + to_sql(foo_opt_param) + + def foo_ret() -> py_type: ... + with self.assertRaisesRegex(TypeError, expected_name): + to_sql(foo_ret) + + with self.assertRaisesRegex(TypeError, expected_name): + @udf + def foo_udf_param(x: py_type) -> int: ... + + with self.assertRaisesRegex(TypeError, expected_name): + @udf + def foo_udf_ret(x: int) -> py_type: ... + + with self.assertRaisesRegex(TypeError, expected_name): + @udf() + def foo_udf_call(x: Optional[py_type]) -> int: ... - # Date - def foo(x: datetime.date) -> None: ... - assert to_sql(foo) == '`foo`(`x` DATE NOT NULL) RETURNS TINYINT NULL' - - # Time - def foo(x: datetime.timedelta) -> None: ... - assert to_sql(foo) == '`foo`(`x` TIME NOT NULL) RETURNS TINYINT NULL' - - # Datetime + Date def foo(x: Union[datetime.datetime, datetime.date]) -> None: ... with self.assertRaises(TypeError): to_sql(foo) + for sql_override, expected_sql in [ + (dt.DATE, 'DATE'), + (dt.DATE(), 'DATE'), + (dt.DATE(nullable=False), 'DATE'), + (dt.TIME, 'TIME'), + (dt.TIME(), 'TIME'), + (dt.TIME(6), 'TIME'), + (dt.DATETIME, 'DATETIME'), + (dt.DATETIME(), 'DATETIME'), + (dt.DATETIME(6), 'DATETIME'), + (dt.TIMESTAMP, 'TIMESTAMP'), + (dt.TIMESTAMP(), 'TIMESTAMP'), + (dt.TIMESTAMP(6), 'TIMESTAMP'), + (dt.YEAR, 'YEAR'), + (dt.YEAR(), 'YEAR'), + (dt.DECIMAL(10, 2), 'DECIMAL'), + (dt.DEC(10, 2), 'DEC'), + (dt.FIXED(10, 2), 'FIXED'), + (dt.NUMERIC(10, 2), 'NUMERIC'), + ('DECIMAL', 'DECIMAL'), + ]: + with self.assertRaisesRegex(TypeError, expected_sql): + @udf(args=sql_override) + def foo_arg_override(x: int) -> int: ... + + with self.assertRaisesRegex(TypeError, expected_sql): + @udf(returns=sql_override) + def foo_ret_override(x: int) -> int: ... + + @udf(args=dt.TEXT, returns=dt.TEXT) + def foo_supported_override(x: datetime.date) -> datetime.date: + return x + assert to_sql(foo_supported_override) == \ + '`foo_supported_override`(`x` TEXT NULL) RETURNS TEXT NULL' + def test_numerics(self): # # Bools @@ -405,10 +459,9 @@ def foo(x: int) -> int: ... assert to_sql(foo) == '`foo`(`x` BIGINT NOT NULL) RETURNS SMALLINT NOT NULL' # Override multiple params with one type - @udf(args=dt.SMALLINT(nullable=False)) - def foo(x: int, y: float, z: np.int8) -> int: ... with self.assertRaises(ValueError): - to_sql(foo) + @udf(args=dt.SMALLINT(nullable=False)) + def foo(x: int, y: float, z: np.int8) -> int: ... # Override with list @udf(args=[dt.SMALLINT, dt.FLOAT, dt.CHAR(30)]) @@ -418,16 +471,14 @@ def foo(x: int, y: float, z: str) -> int: ... '`z` CHAR(30) NULL) RETURNS BIGINT NOT NULL' # Override with too short of a list - @udf(args=[dt.SMALLINT, dt.FLOAT]) - def foo(x: int, y: float, z: str) -> int: ... with self.assertRaises(ValueError): - to_sql(foo) + @udf(args=[dt.SMALLINT, dt.FLOAT]) + def foo(x: int, y: float, z: str) -> int: ... # Override with too long of a list - @udf(args=[dt.SMALLINT, dt.FLOAT, dt.CHAR(30), dt.TEXT]) - def foo(x: int, y: float, z: str) -> int: ... with self.assertRaises(ValueError): - to_sql(foo) + @udf(args=[dt.SMALLINT, dt.FLOAT, dt.CHAR(30), dt.TEXT]) + def foo(x: int, y: float, z: str) -> int: ... # Override with list @udf(args=[dt.SMALLINT, dt.FLOAT, dt.CHAR(30)]) diff --git a/singlestoredb/tests/test_udf_future_annotations.py b/singlestoredb/tests/test_udf_future_annotations.py new file mode 100644 index 000000000..611cbce00 --- /dev/null +++ b/singlestoredb/tests/test_udf_future_annotations.py @@ -0,0 +1,113 @@ +#!/usr/bin/env python +# type: ignore +"""UDF signatures in a module using postponed evaluation of annotations.""" +from __future__ import annotations + +import dataclasses +import datetime +import unittest +from datetime import date +from typing import List +from typing import NamedTuple +from typing import Optional + +import numpy as np +import numpy.typing as npt + +from ..functions import dtypes as dt +from ..functions import Masked +from ..functions import signature as sig +from ..functions import Table +from ..functions import udf +from ..functions.ext.asgi import get_masked_params +from ..functions.ext.plugin import registry + + +def to_sql(x): + out = sig.signature_to_sql(sig.get_signature(x)) + out = out.split('EXTERNAL FUNCTION ')[1] + return out.split('AS REMOTE')[0].strip() + + +class Row(NamedTuple): + a: int + b: Optional[str] + + +@dataclasses.dataclass +class DataclassRow: + a: int + b: Optional[str] = None + + +class TestFutureAnnotations(unittest.TestCase): + + def test_scalars(self): + @udf + def foo(x: Optional[int], y: float = 1.5) -> Optional[str]: ... + assert to_sql(foo) == \ + '`foo`(`x` BIGINT NULL, `y` DOUBLE NOT NULL DEFAULT 1.5e0) ' \ + 'RETURNS TEXT NULL' + + def test_masked_vectors(self): + @udf + def foo( + x: npt.NDArray[np.float32], + ) -> Masked[npt.NDArray[np.float32]]: ... + assert to_sql(foo) == '`foo`(`x` FLOAT NOT NULL) RETURNS FLOAT NULL' + + def test_masked_params(self): + @udf + def foo( + x: Masked[npt.NDArray[np.float32]], + y: npt.NDArray[np.float32], + ) -> npt.NDArray[np.float32]: ... + assert to_sql(foo) == \ + '`foo`(`x` FLOAT NULL, `y` FLOAT NOT NULL) RETURNS FLOAT NOT NULL' + assert get_masked_params(foo) == [True, False] + assert registry._get_masked_params(foo) == [True, False] + + def test_table(self): + @udf + def foo(x: int) -> Table[List[Row]]: ... + assert to_sql(foo) == \ + '`foo`(`x` BIGINT NOT NULL) ' \ + 'RETURNS TABLE(`a` BIGINT NOT NULL, `b` TEXT NULL)' + + def test_dataclass_table(self): + @udf + def foo(x: int) -> Table[List[DataclassRow]]: ... + assert to_sql(foo) == \ + '`foo`(`x` BIGINT NOT NULL) ' \ + 'RETURNS TABLE(`a` BIGINT NOT NULL, `b` TEXT NULL)' + + def test_dtype_strings(self): + @udf + def foo(x: 'str') -> 'array[int64]': ... # noqa: F821 + assert to_sql(foo) == \ + '`foo`(`x` TEXT NOT NULL) RETURNS ARRAY(BIGINT NOT NULL) NOT NULL' + + def test_unsupported_types(self): + with self.assertRaisesRegex(TypeError, 'datetime.date'): + @udf + def foo(x: datetime.date) -> int: ... + + with self.assertRaisesRegex(TypeError, 'datetime.date'): + @udf + def foo_short(x: date) -> int: ... + + @udf(args=dt.TEXT, returns=dt.TEXT) + def bar(x: datetime.date) -> datetime.date: ... + assert to_sql(bar) == '`bar`(`x` TEXT NULL) RETURNS TEXT NULL' + + @udf(args=dt.TEXT, returns=dt.TEXT) + def bar_short(x: date) -> date: ... + assert to_sql(bar_short) == '`bar_short`(`x` TEXT NULL) RETURNS TEXT NULL' + + def test_unresolvable(self): + class Local(NamedTuple): + a: int + + with self.assertRaisesRegex(TypeError, 'could not resolve.*Local'): + @udf + def foo(x: int) -> Table[List[Local]]: ... diff --git a/singlestoredb/tests/test_udf_returns.py b/singlestoredb/tests/test_udf_returns.py index f1daa58de..b44572191 100644 --- a/singlestoredb/tests/test_udf_returns.py +++ b/singlestoredb/tests/test_udf_returns.py @@ -138,17 +138,10 @@ def foo_h_(x: str) -> Table[List[str]]: '`foo_h_`(`x` TEXT NOT NULL) RETURNS TABLE(`a` TEXT NOT NULL)' # Plain TVF using multiple lists -- Illegal! - @udf - def foo_h(x: str) -> Table[List[int], List[str]]: - return Table(list(range(3)), [x] * 3) - - foo_h_out = foo_h('cat') - - assert type(foo_h_out) is Table - assert foo_h_out == Table([0, 1, 2], ['cat', 'cat', 'cat']) - with self.assertRaises(TypeError): - to_sql(foo_h) + @udf + def foo_h(x: str) -> Table[List[int], List[str]]: + return Table(list(range(3)), [x] * 3) # Plain TVF using lists of NamedTuples @udf @@ -361,16 +354,10 @@ def foo_s(x: str) -> Table[pa.Table]: ) # Vectorized TVF using lists -- Illegal! - @udf - def foo_t(x: List[str]) -> Table[List[int], List[str]]: - return Table(list(range(len(x))), x) - - foo_t_out = foo_t(['cat', 'dog', 'monkey']) - - assert type(foo_t_out) is Table - assert foo_t_out == Table([0, 1, 2], ['cat', 'dog', 'monkey']) with self.assertRaises(TypeError): - to_sql(foo_t) + @udf + def foo_t(x: List[str]) -> Table[List[int], List[str]]: + return Table(list(range(len(x))), x) # Vectorized TVF using pandas Series @udf(args=Parameters, returns=TVFTuple)