singlestoredb 1.16.1__py3-none-any.whl

singlestoredb/docstring/util.py

@@ -0,0 +1,141 @@
+"""Utility functions for working with docstrings."""
+import typing as T
+from collections import ChainMap
+from inspect import Signature
+from itertools import chain
+
+from .common import DocstringMeta
+from .common import DocstringParam
+from .common import DocstringReturns  # noqa: F401
+from .common import DocstringStyle
+from .common import RenderingStyle
+from .parser import compose
+from .parser import parse
+
+_Func = T.Callable[..., T.Any]
+
+
+def combine_docstrings(
+    *others: _Func,
+    exclude: T.Iterable[T.Type[DocstringMeta]] = (),
+    style: DocstringStyle = DocstringStyle.AUTO,
+    rendering_style: RenderingStyle = RenderingStyle.COMPACT,
+) -> _Func:
+    """A function decorator that parses the docstrings from `others`,
+    programmatically combines them with the parsed docstring of the decorated
+    function, and replaces the docstring of the decorated function with the
+    composed result. Only parameters that are part of the decorated functions
+    signature are included in the combined docstring. When multiple sources for
+    a parameter or docstring metadata exists then the decorator will first
+    default to the wrapped function's value (when available) and otherwise use
+    the rightmost definition from ``others``.
+
+    The following example illustrates its usage:
+
+    >>> def fun1(a, b, c, d):
+    ...    '''short_description: fun1
+    ...
+    ...    :param a: fun1
+    ...    :param b: fun1
+    ...    :return: fun1
+    ...    '''
+    >>> def fun2(b, c, d, e):
+    ...    '''short_description: fun2
+    ...
+    ...    long_description: fun2
+    ...
+    ...    :param b: fun2
+    ...    :param c: fun2
+    ...    :param e: fun2
+    ...    '''
+    >>> @combine_docstrings(fun1, fun2)
+    >>> def decorated(a, b, c, d, e, f):
+    ...     '''
+    ...     :param e: decorated
+    ...     :param f: decorated
+    ...     '''
+    >>> print(decorated.__doc__)
+    short_description: fun2
+    <BLANKLINE>
+    long_description: fun2
+    <BLANKLINE>
+    :param a: fun1
+    :param b: fun1
+    :param c: fun2
+    :param e: fun2
+    :param f: decorated
+    :returns: fun1
+    >>> @combine_docstrings(fun1, fun2, exclude=[DocstringReturns])
+    >>> def decorated(a, b, c, d, e, f): pass
+    >>> print(decorated.__doc__)
+    short_description: fun2
+    <BLANKLINE>
+    long_description: fun2
+    <BLANKLINE>
+    :param a: fun1
+    :param b: fun1
+    :param c: fun2
+    :param e: fun2
+
+    :param others: callables from which to parse docstrings.
+    :param exclude: an iterable of ``DocstringMeta`` subclasses to exclude when
+        combining docstrings.
+    :param style: style composed docstring. The default will infer the style
+        from the decorated function.
+    :param rendering_style: The rendering style used to compose a docstring.
+    :return: the decorated function with a modified docstring.
+    """
+
+    def wrapper(func: _Func) -> _Func:
+        sig = Signature.from_callable(func)
+
+        comb_doc = parse(func.__doc__ or '')
+        docs = [parse(other.__doc__ or '') for other in others] + [comb_doc]
+        params = dict(
+            ChainMap(
+                *(
+                    {param.arg_name: param for param in doc.params}
+                    for doc in docs
+                ),
+            ),
+        )
+
+        for doc in reversed(docs):
+            if not doc.short_description:
+                continue
+            comb_doc.short_description = doc.short_description
+            comb_doc.blank_after_short_description = (
+                doc.blank_after_short_description
+            )
+            break
+
+        for doc in reversed(docs):
+            if not doc.long_description:
+                continue
+            comb_doc.long_description = doc.long_description
+            comb_doc.blank_after_long_description = (
+                doc.blank_after_long_description
+            )
+            break
+
+        combined: T.Dict[T.Type[DocstringMeta], T.List[DocstringMeta]] = {}
+        for doc in docs:
+            metas: T.Dict[T.Type[DocstringMeta], T.List[DocstringMeta]] = {}
+            for meta in doc.meta:
+                meta_type = type(meta)
+                if meta_type in exclude:
+                    continue
+                metas.setdefault(meta_type, []).append(meta)
+            for meta_type, meta_list in metas.items():
+                combined[meta_type] = meta_list
+
+        combined[DocstringParam] = [
+            params[name] for name in sig.parameters if name in params
+        ]
+        comb_doc.meta = list(chain(*combined.values()))
+        func.__doc__ = compose(
+            comb_doc, style=style, rendering_style=rendering_style,
+        )
+        return func
+
+    return wrapper

singlestoredb/exceptions.py

@@ -0,0 +1,120 @@
+#!/usr/bin/env python
+"""Database exeception classes."""
+from typing import Optional
+
+
+class MySQLError(Exception):
+    """All MySQL-related exceptions."""
+
+
+class Error(MySQLError):
+    """
+    Generic database exception.
+
+    Parameters
+    ----------
+    errno : int
+        Database error code
+    msg : str
+        Database error message
+    sqlstate : str, optional
+        SQL engine state code
+
+    """
+
+    def __init__(
+        self, errno: Optional[int] = None, msg: Optional[str] = None,
+        sqlstate: Optional[int] = None,
+    ):
+        self.errno = errno
+        self.errmsg = msg
+        self.sqlstate = sqlstate
+        super(Exception, self).__init__(errno, msg, sqlstate)
+
+    def __str__(self) -> str:
+        """Return string representation."""
+        prefix = []
+        if self.errno is not None:
+            prefix.append(f'{self.errno}')
+        if self.sqlstate is not None:
+            prefix.append(f'({self.sqlstate})')
+        if prefix and self.errmsg:
+            return ' '.join(prefix) + ': ' + self.errmsg
+        elif prefix:
+            return ' '.join(prefix)
+        elif self.errmsg:
+            return f'{self.errmsg}'
+        return 'Unknown error'
+
+    def __repr__(self) -> str:
+        """Return string representation."""
+        return str(self)
+
+    @property
+    def msg(self) -> Optional[str]:
+        """Return error message."""
+        return self.errmsg
+
+
+class Warning(Warning, MySQLError):  # type: ignore
+    """Exception for important warnings like data truncations, etc."""
+
+
+class InterfaceError(Error):
+    """Exception for errors in the database interface rather than the database."""
+
+
+class DatabaseError(Error):
+    """General exception for errors in the database."""
+
+
+class InternalError(DatabaseError):
+    """Exception for internal database errors such as out of sync transactions."""
+
+
+class OperationalError(DatabaseError):
+    """Exception for operational errors such as unexpected disconnections."""
+
+
+class ProgrammingError(DatabaseError):
+    """Exception for programming errors."""
+
+
+class IntegrityError(DatabaseError):
+    """Exception for relational integrity errors."""
+
+
+class DataError(DatabaseError):
+    """Exception for problems with processed data like division by zero."""
+
+
+class NotSupportedError(DatabaseError):
+    """Exception for using unsupported features of the database."""
+
+
+class ManagementError(Error):
+    """Exception for errors in the management API."""
+
+    def __init__(
+        self, errno: Optional[int] = None, msg: Optional[str] = None,
+        response: Optional[str] = None,
+    ):
+        self.errno = errno
+        self.errmsg = msg
+        self.response = response
+        super(Exception, self).__init__(errno, msg)
+
+    def __str__(self) -> str:
+        """Return string representation."""
+        prefix = []
+        if self.errno is not None:
+            prefix.append(f'{self.errno}')
+        if self.response is not None:
+            prefix.append(f'({self.response})')
+        if prefix and self.errmsg:
+            return ' '.join(prefix) + ': ' + self.errmsg
+        elif prefix:
+            return ' '.join(prefix)
+        elif self.errmsg:
+            return f'{self.errmsg}'
+        return 'Unknown error'

singlestoredb/functions/__init__.py

@@ -0,0 +1,16 @@
+from .decorator import udf  # noqa: F401
+from .typing import Masked  # noqa: F401
+from .typing import Table  # noqa: F401
+from .utils import pack_vector  # noqa: F401
+from .utils import pack_vectors  # noqa: F401
+from .utils import unpack_vector  # noqa: F401
+from .utils import unpack_vectors  # noqa: F401
+from .utils import VectorTypes
+
+
+F32 = VectorTypes.F32
+F64 = VectorTypes.F64
+I8 = VectorTypes.I8
+I16 = VectorTypes.I16
+I32 = VectorTypes.I32
+I64 = VectorTypes.I64

singlestoredb/functions/decorator.py

@@ -0,0 +1,201 @@
+import asyncio
+import functools
+import inspect
+from typing import Any
+from typing import Callable
+from typing import List
+from typing import Optional
+from typing import Type
+from typing import Union
+
+from . import utils
+from .dtypes import SQLString
+
+
+ParameterType = Union[
+    str,
+    Callable[..., SQLString],
+    List[Union[str, Callable[..., SQLString]]],
+    Type[Any],
+]
+
+ReturnType = ParameterType
+UDFType = Callable[..., Any]
+
+
+def is_valid_type(obj: Any) -> bool:
+    """Check if the object is a valid type for a schema definition."""
+    if not inspect.isclass(obj):
+        return False
+
+    if utils.is_typeddict(obj):
+        return True
+
+    if utils.is_namedtuple(obj):
+        return True
+
+    if utils.is_dataclass(obj):
+        return True
+
+    # We don't want to import pydantic here, so we check if
+    # the class is a subclass
+    if utils.is_pydantic(obj):
+        return True
+
+    return False
+
+
+def is_sqlstr_callable(obj: Any) -> bool:
+    """Check if the object is a valid callable for a parameter type."""
+    if not callable(obj):
+        return False
+
+    returns = utils.get_annotations(obj).get('return', None)
+
+    if inspect.isclass(returns) and issubclass(returns, SQLString):
+        return True
+
+    return False
+
+
+def expand_types(args: Any) -> Optional[List[Any]]:
+    """Expand the types for the function arguments / return values."""
+    if args is None:
+        return None
+
+    # SQL string
+    if isinstance(args, str):
+        return [args]
+
+    # List of SQL strings or callables
+    elif isinstance(args, list):
+        new_args: List[Any] = []
+        for arg in args:
+            if isinstance(arg, str):
+                new_args.append(arg)
+            elif is_sqlstr_callable(arg):
+                new_args.append(arg())
+            elif type(arg) is type:
+                new_args.append(arg)
+            elif is_valid_type(arg):
+                new_args.append(arg)
+            else:
+                raise TypeError(f'unrecognized type for parameter: {arg}')
+        return new_args
+
+    # Callable that returns a SQL string
+    elif is_sqlstr_callable(args):
+        return [args()]
+
+    # General way of accepting pydantic.BaseModel, NamedTuple, TypedDict
+    elif is_valid_type(args):
+        return [args]
+
+    elif type(args) is type:
+        return [args]
+
+    raise TypeError(f'unrecognized type for parameter: {args}')
+
+
+def _func(
+    func: Optional[Callable[..., Any]] = None,
+    *,
+    name: Optional[str] = None,
+    args: Optional[ParameterType] = None,
+    returns: Optional[ReturnType] = None,
+    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: UDFType) -> UDFType:
+
+            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)
+
+            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
+
+    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)
+
+    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(
+    func: Optional[Callable[..., Any]] = None,
+    *,
+    name: Optional[str] = None,
+    args: Optional[ParameterType] = None,
+    returns: Optional[ReturnType] = None,
+    timeout: Optional[int] = None,
+) -> UDFType:
+    """
+    Define a user-defined function (UDF).
+
+    Parameters
+    ----------
+    func : callable, optional
+        The UDF to apply parameters to
+    name : str, optional
+        The name to use for the UDF in the database
+    args : str | Type | Callable | List[str | Callable], optional
+        Specifies the data types of the function arguments. Typically,
+        the function data types are derived from the function parameter
+        annotations. These annotations can be overridden. If the function
+        takes a single type for all parameters, `args` can be set to a
+        SQL string describing all parameters. If the function takes more
+        than one parameter and all of the parameters are being manually
+        defined, a list of SQL strings may be used (one for each parameter).
+        A dictionary of SQL strings may be used to specify a parameter type
+        for a subset of parameters; the keys are the names of the
+        function parameters. Callables may also be used for datatypes. This
+        is primarily for using the functions in the ``dtypes`` module that
+        are associated with SQL types with all default options (e.g., ``dt.FLOAT``).
+    returns : str | Type | Callable | List[str | Callable] | Table, optional
+        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
+    -------
+    Callable
+
+    """
+    return _func(
+        func=func,
+        name=name,
+        args=args,
+        returns=returns,
+        timeout=timeout,
+    )