singlestoredb 0.4.0__py3-none-any.whl → 1.0.4__py3-none-any.whl

singlestoredb/converters.py

@@ -1,8 +1,8 @@
 #!/usr/bin/env python
 """Data value conversion utilities."""
 import datetime
+import re
 from base64 import b64decode
-from datetime import timezone
 from decimal import Decimal
 from json import loads as json_loads
 from typing import Any
@@ -13,160 +13,246 @@ from typing import Optional
 from typing import Set
 from typing import Union
 
+try:
+    import shapely.wkt
+    has_shapely = True
+except ImportError:
+    has_shapely = False
+
+try:
+    import pygeos
+    has_pygeos = True
+except ImportError:
+    has_pygeos = False
+
+
+# Cache fromisoformat methods if they exist
+_dt_datetime_fromisoformat = None
+if hasattr(datetime.datetime, 'fromisoformat'):
+    _dt_datetime_fromisoformat = datetime.datetime.fromisoformat  # type: ignore
+_dt_time_fromisoformat = None
+if hasattr(datetime.time, 'fromisoformat'):
+    _dt_time_fromisoformat = datetime.time.fromisoformat  # type: ignore
+_dt_date_fromisoformat = None
+if hasattr(datetime.date, 'fromisoformat'):
+    _dt_date_fromisoformat = datetime.date.fromisoformat  # type: ignore
+
+
+def _convert_second_fraction(s: str) -> int:
+    if not s:
+        return 0
+    # Pad zeros to ensure the fraction length in microseconds
+    s = s.ljust(6, '0')
+    return int(s[:6])
+
+
+DATETIME_RE = re.compile(
+    r'(\d{1,4})-(\d{1,2})-(\d{1,2})[T ](\d{1,2}):(\d{1,2}):(\d{1,2})(?:.(\d{1,6}))?',
+)
+
+ZERO_DATETIMES = set([
+    '0000-00-00 00:00:00',
+    '0000-00-00 00:00:00.000',
+    '0000-00-00 00:00:00.000000',
+    '0000-00-00T00:00:00',
+    '0000-00-00T00:00:00.000',
+    '0000-00-00T00:00:00.000000',
+])
+ZERO_DATES = set([
+    '0000-00-00',
+])
+
+
+def datetime_fromisoformat(
+    obj: Union[str, bytes, bytearray],
+) -> Union[datetime.datetime, str, None]:
+    """Returns a DATETIME or TIMESTAMP column value as a datetime object:
+
+      >>> datetime_fromisoformat('2007-02-25 23:06:20')
+      datetime.datetime(2007, 2, 25, 23, 6, 20)
+      >>> datetime_fromisoformat('2007-02-25T23:06:20')
+      datetime.datetime(2007, 2, 25, 23, 6, 20)
+
+    Illegal values are returned as str or None:
+
+      >>> datetime_fromisoformat('2007-02-31T23:06:20')
+      '2007-02-31T23:06:20'
+      >>> datetime_fromisoformat('0000-00-00 00:00:00')
+      None
 
-def _parse_isoformat_date(dtstr: str) -> List[Any]:
-    # It is assumed that this function will only be called with a
-    # string of length exactly 10, and (though this is not used) ASCII-only
-    year = int(dtstr[0:4])
-    if dtstr[4] != '-':
-        raise ValueError('Invalid date separator: %s' % dtstr[4])
-
-    month = int(dtstr[5:7])
-
-    if dtstr[7] != '-':
-        raise ValueError('Invalid date separator')
-
-    day = int(dtstr[8:10])
-
-    return [year, month, day]
-
-
-def _parse_hh_mm_ss_ff(tstr: str) -> List[Any]:
-    # Parses things of the form HH[:MM[:SS[.fff[fff]]]]
-    len_str = len(tstr)
+    """
+    if isinstance(obj, (bytes, bytearray)):
+        obj = obj.decode('ascii')
 
-    time_comps = [0, 0, 0, 0]
-    pos = 0
-    for comp in range(0, 3):
-        if (len_str - pos) < 2:
-            raise ValueError('Incomplete time component')
+    if obj in ZERO_DATETIMES:
+        return None
 
-        time_comps[comp] = int(tstr[pos:pos+2])
+    # Use datetime methods if possible
+    if _dt_datetime_fromisoformat is not None:
+        try:
+            if ' ' in obj or 'T' in obj:
+                return _dt_datetime_fromisoformat(obj)
+            if _dt_date_fromisoformat is not None:
+                date = _dt_date_fromisoformat(obj)
+                return datetime.datetime(date.year, date.month, date.day)
+        except ValueError:
+            return obj
 
-        pos += 2
-        next_char = tstr[pos:pos+1]
+    m = DATETIME_RE.match(obj)
+    if not m:
+        mdate = date_fromisoformat(obj)
+        if type(mdate) is str:
+            return mdate
+        return datetime.datetime(mdate.year, mdate.month, mdate.day)  # type: ignore
 
-        if not next_char or comp >= 2:
-            break
+    try:
+        groups = list(m.groups())
+        groups[-1] = _convert_second_fraction(groups[-1])
+        return datetime.datetime(*[int(x) for x in groups])  # type: ignore
+    except ValueError:
+        mdate = date_fromisoformat(obj)
+        if type(mdate) is str:
+            return mdate
+        return datetime.datetime(mdate.year, mdate.month, mdate.day)  # type: ignore
 
-        if next_char != ':':
-            raise ValueError('Invalid time separator: %c' % next_char)
 
-        pos += 1
+TIMEDELTA_RE = re.compile(r'(-)?(\d{1,3}):(\d{1,2}):(\d{1,2})(?:.(\d{1,6}))?')
 
-    if pos < len_str:
-        if tstr[pos] != '.':
-            raise ValueError('Invalid microsecond component')
-        else:
-            pos += 1
 
-            len_remainder = len_str - pos
-            if len_remainder not in (3, 6):
-                raise ValueError('Invalid microsecond component')
+def timedelta_fromisoformat(
+    obj: Union[str, bytes, bytearray],
+) -> Union[datetime.timedelta, str, None]:
+    """Returns a TIME column as a timedelta object:
 
-            time_comps[3] = int(tstr[pos:])
-            if len_remainder == 3:
-                time_comps[3] *= 1000
+      >>> timedelta_fromisoformat('25:06:17')
+      datetime.timedelta(days=1, seconds=3977)
+      >>> timedelta_fromisoformat('-25:06:17')
+      datetime.timedelta(days=-2, seconds=82423)
 
-    return time_comps
+    Illegal values are returned as string:
 
+      >>> timedelta_fromisoformat('random crap')
+      'random crap'
 
-def _parse_isoformat_time(tstr: str) -> List[Any]:
-    # Format supported is HH[:MM[:SS[.fff[fff]]]][+HH:MM[:SS[.ffffff]]]
-    len_str = len(tstr)
-    if len_str < 2:
-        raise ValueError('Isoformat time too short')
+    Note that MySQL always returns TIME columns as (+|-)HH:MM:SS, but
+    can accept values as (+|-)DD HH:MM:SS. The latter format will not
+    be parsed correctly by this function.
+    """
+    if isinstance(obj, (bytes, bytearray)):
+        obj = obj.decode('ascii')
 
-    # This is equivalent to re.search('[+-]', tstr), but faster
-    tz_pos = (tstr.find('-') + 1 or tstr.find('+') + 1)
-    timestr = tstr[:tz_pos-1] if tz_pos > 0 else tstr
+    m = TIMEDELTA_RE.match(obj)
+    if not m:
+        return obj
 
-    time_comps = _parse_hh_mm_ss_ff(timestr)
+    try:
+        groups = list(m.groups())
+        groups[-1] = _convert_second_fraction(groups[-1])
+        negate = -1 if groups[0] else 1
+        hours, minutes, seconds, microseconds = groups[1:]
+
+        tdelta = (
+            datetime.timedelta(
+                hours=int(hours),
+                minutes=int(minutes),
+                seconds=int(seconds),
+                microseconds=int(microseconds),
+            )
+            * negate
+        )
+        return tdelta
+    except ValueError:
+        return obj
 
-    tzi = None
-    if tz_pos > 0:
-        tzstr = tstr[tz_pos:]
 
-        # Valid time zone strings are:
-        # HH:MM               len: 5
-        # HH:MM:SS            len: 8
-        # HH:MM:SS.ffffff     len: 15
+TIME_RE = re.compile(r'(\d{1,2}):(\d{1,2}):(\d{1,2})(?:.(\d{1,6}))?')
 
-        if len(tzstr) not in (5, 8, 15):
-            raise ValueError('Malformed time zone string')
 
-        tz_comps = _parse_hh_mm_ss_ff(tzstr)
-        if all(x == 0 for x in tz_comps):
-            tzi = timezone.utc
-        else:
-            tzsign = -1 if tstr[tz_pos - 1] == '-' else 1
+def time_fromisoformat(
+    obj: Union[str, bytes, bytearray],
+) -> Union[datetime.time, str, None]:
+    """Returns a TIME column as a time object:
 
-            td = datetime.timedelta(
-                hours=tz_comps[0], minutes=tz_comps[1],
-                seconds=tz_comps[2], microseconds=tz_comps[3],
-            )
+      >>> time_fromisoformat('15:06:17')
+      datetime.time(15, 6, 17)
 
-            tzi = timezone(tzsign * td)
+    Illegal values are returned as str:
 
-    time_comps.append(tzi)
+      >>> time_fromisoformat('-25:06:17')
+      '-25:06:17'
+      >>> time_fromisoformat('random crap')
+      'random crap'
 
-    return time_comps
+    Note that MySQL always returns TIME columns as (+|-)HH:MM:SS, but
+    can accept values as (+|-)DD HH:MM:SS. The latter format will not
+    be parsed correctly by this function.
 
+    Also note that MySQL's TIME column corresponds more closely to
+    Python's timedelta and not time. However if you want TIME columns
+    to be treated as time-of-day and not a time offset, then you can
+    use set this function as the converter for FIELD_TYPE.TIME.
+    """
+    if isinstance(obj, (bytes, bytearray)):
+        obj = obj.decode('ascii')
 
-def datetime_fromisoformat(date_string: str) -> datetime.datetime:
-    """Construct a datetime from the output of datetime.isoformat()."""
-    if not isinstance(date_string, str):
-        raise TypeError('fromisoformat: argument must be str')
+    # Use datetime methods if possible
+    if _dt_time_fromisoformat is not None:
+        try:
+            return _dt_time_fromisoformat(obj)
+        except ValueError:
+            return obj
 
-    # Split this at the separator
-    dstr = date_string[0:10]
-    tstr = date_string[11:]
+    m = TIME_RE.match(obj)
+    if not m:
+        return obj
 
     try:
-        date_components = _parse_isoformat_date(dstr)
+        groups = list(m.groups())
+        groups[-1] = _convert_second_fraction(groups[-1])
+        hours, minutes, seconds, microseconds = groups
+        return datetime.time(
+            hour=int(hours),
+            minute=int(minutes),
+            second=int(seconds),
+            microsecond=int(microseconds),
+        )
     except ValueError:
-        raise ValueError(f'Invalid isoformat string: {date_string!r}')
+        return obj
 
-    if tstr:
-        try:
-            time_components = _parse_isoformat_time(tstr)
-        except ValueError:
-            raise ValueError(f'Invalid isoformat string: {date_string!r}')
-    else:
-        time_components = [0, 0, 0, 0, None]
 
-    return datetime.datetime(*(date_components + time_components))
+def date_fromisoformat(
+    obj: Union[str, bytes, bytearray],
+) -> Union[datetime.date, str, None]:
+    """Returns a DATE column as a date object:
 
+      >>> date_fromisoformat('2007-02-26')
+      datetime.date(2007, 2, 26)
 
-def date_fromisoformat(date_string: str) -> datetime.date:
-    """Construct a date from the output of date.isoformat()."""
-    if not isinstance(date_string, str):
-        raise TypeError('fromisoformat: argument must be str')
+    Illegal values are returned as str or None:
 
-    try:
-        assert len(date_string) == 10
-        return datetime.date(*_parse_isoformat_date(date_string))
-    except Exception:
-        raise ValueError(f'Invalid isoformat string: {date_string!r}')
+      >>> date_fromisoformat('2007-02-31')
+      '2007-02-31'
+      >>> date_fromisoformat('0000-00-00')
+      None
 
+    """
+    if isinstance(obj, (bytes, bytearray)):
+        obj = obj.decode('ascii')
 
-def time_fromisoformat(time_string: str) -> datetime.time:
-    """Construct a time from the output of isoformat()."""
-    if not isinstance(time_string, str):
-        raise TypeError('fromisoformat: argument must be str')
-
-    try:
-        return datetime.time(*_parse_isoformat_time(time_string))
-    except Exception:
-        raise ValueError(f'Invalid isoformat string: {time_string!r}')
+    if obj in ZERO_DATES:
+        return None
 
+    # Use datetime methods if possible
+    if _dt_date_fromisoformat is not None:
+        try:
+            return _dt_date_fromisoformat(obj)
+        except ValueError:
+            return obj
 
-# datetime_fromisoformat = datetime.datetime.fromisoformat
-# time_fromisoformat = datetime.time.fromisoformat
-# date_fromisoformat = datetime.date.fromisoformat
-datetime_min = datetime.datetime.min
-date_min = datetime.date.min
-datetime_combine = datetime.datetime.combine
+    try:
+        return datetime.date(*[int(x) for x in obj.split('-', 2)])
+    except ValueError:
+        return obj
 
 
 def identity(x: Any) -> Optional[Any]:
@@ -264,7 +350,7 @@ def decimal_or_none(x: Any) -> Optional[Decimal]:
     return Decimal(x)
 
 
-def date_or_none(x: Optional[str]) -> Optional[datetime.date]:
+def date_or_none(x: Optional[str]) -> Optional[Union[datetime.date, str]]:
     """
     Convert value to a date.
 
@@ -283,13 +369,10 @@ def date_or_none(x: Optional[str]) -> Optional[datetime.date]:
     """
     if x is None:
         return None
-    try:
-        return date_fromisoformat(x)
-    except ValueError:
-        return None
+    return date_fromisoformat(x)
 
 
-def time_or_none(x: Optional[str]) -> Optional[datetime.timedelta]:
+def timedelta_or_none(x: Optional[str]) -> Optional[Union[datetime.timedelta, str]]:
     """
     Convert value to a timedelta.
 
@@ -308,13 +391,32 @@ def time_or_none(x: Optional[str]) -> Optional[datetime.timedelta]:
     """
     if x is None:
         return None
-    try:
-        return datetime_combine(date_min, time_fromisoformat(x)) - datetime_min
-    except ValueError:
+    return timedelta_fromisoformat(x)
+
+
+def time_or_none(x: Optional[str]) -> Optional[Union[datetime.time, str]]:
+    """
+    Convert value to a time.
+
+    Parameters
+    ----------
+    x : Any
+        Arbitrary value
+
+    Returns
+    -------
+    datetime.time
+        If value can be cast to a time
+    None
+        If input value is None
+
+    """
+    if x is None:
         return None
+    return time_fromisoformat(x)
 
 
-def datetime_or_none(x: Optional[str]) -> Optional[datetime.datetime]:
+def datetime_or_none(x: Optional[str]) -> Optional[Union[datetime.datetime, str]]:
     """
     Convert value to a datetime.
 
@@ -333,10 +435,7 @@ def datetime_or_none(x: Optional[str]) -> Optional[datetime.datetime]:
     """
     if x is None:
         return None
-    try:
-        return datetime_fromisoformat(x)
-    except ValueError:
-        return None
+    return datetime_fromisoformat(x)
 
 
 def none(x: Any) -> None:
@@ -413,14 +512,18 @@ def geometry_or_none(x: Optional[str]) -> Optional[Any]:
 
     Returns
     -------
-    ???
+    shapely object or pygeos object or str
         If value is valid geometry value
     None
-        If input value is None
+        If input value is None or empty
 
     """
-    if x is None:
+    if x is None or not x:
         return None
+    if has_shapely:
+        return shapely.wkt.loads(x)
+    if has_pygeos:
+        return pygeos.io.from_wkt(x)
     return x
 
 
@@ -437,7 +540,7 @@ converters: Dict[int, Callable[..., Any]] = {
     8: int_or_none,
     9: int_or_none,
     10: date_or_none,
-    11: time_or_none,
+    11: timedelta_or_none,
     12: datetime_or_none,
     13: int_or_none,
     14: date_or_none,

singlestoredb/exceptions.py

@@ -3,7 +3,11 @@
 from typing import Optional
 
 
-class Error(Exception):
+class MySQLError(Exception):
+    """All MySQL-related exceptions."""
+
+
+class Error(MySQLError):
     """
     Generic database exception.
 
@@ -52,7 +56,7 @@ class Error(Exception):
         return self.errmsg
 
 
-class Warning(Exception):
+class Warning(Warning, MySQLError):  # type: ignore
     """Exception for important warnings like data truncations, etc."""
 
 
@@ -90,3 +94,27 @@ class NotSupportedError(DatabaseError):
 
 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 @@
+from .decorator import udf  # noqa: F401

singlestoredb/functions/decorator.py

@@ -0,0 +1,142 @@
+import functools
+from typing import Any
+from typing import Callable
+from typing import Dict
+from typing import List
+from typing import Optional
+from typing import Union
+
+from .dtypes import DataType
+
+
+def listify(x: Any) -> List[Any]:
+    """Make sure sure value is a list."""
+    if x is None:
+        return []
+    if isinstance(x, (list, tuple, set)):
+        return list(x)
+    return [x]
+
+
+def udf(
+    func: Optional[Callable[..., Any]] = None,
+    *,
+    name: Optional[str] = None,
+    args: Optional[Union[DataType, List[DataType], Dict[str, DataType]]] = None,
+    returns: Optional[str] = None,
+    data_format: Optional[str] = None,
+    include_masks: bool = False,
+) -> Callable[..., Any]:
+    """
+    Apply attributes to a 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 | Callable | List[str | Callable] | Dict[str, 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, optional
+        Specifies the return data type of the function. If not specified,
+        the type annotation from the function is used.
+    data_format : str, optional
+        The data format of each parameter: python, pandas, arrow, polars
+    include_masks : bool, optional
+        Should boolean masks be included with each input parameter to indicate
+        which elements are NULL? This is only used when a input parameters are
+        configured to a vector type (numpy, pandas, polars, arrow).
+
+    Returns
+    -------
+    Callable
+
+    """
+    if args is None:
+        pass
+    elif isinstance(args, (list, tuple)):
+        args = list(args)
+        for i, item in enumerate(args):
+            if callable(item):
+                args[i] = item()
+        for item in args:
+            if not isinstance(item, str):
+                raise TypeError(f'unrecognized type for parameter: {item}')
+    elif isinstance(args, dict):
+        args = dict(args)
+        for k, v in list(args.items()):
+            if callable(v):
+                args[k] = v()
+        for item in args.values():
+            if not isinstance(item, str):
+                raise TypeError(f'unrecognized type for parameter: {item}')
+    elif callable(args):
+        args = args()
+    elif isinstance(args, str):
+        args = args
+    else:
+        raise TypeError(f'unrecognized data type for args: {args}')
+
+    if returns is None:
+        pass
+    elif callable(returns):
+        returns = returns()
+    elif isinstance(returns, str):
+        returns = returns
+    else:
+        raise TypeError(f'unrecognized return type: {returns}')
+
+    if returns is not None and not isinstance(returns, str):
+        raise TypeError(f'unrecognized return type: {returns}')
+
+    if include_masks and data_format == 'python':
+        raise RuntimeError(
+            'include_masks is only valid when using '
+            'vectors for input parameters',
+        )
+
+    _singlestoredb_attrs = {  # type: ignore
+        k: v for k, v in dict(
+            name=name,
+            args=args,
+            returns=returns,
+            data_format=data_format,
+            include_masks=include_masks,
+        ).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 wrapper(*args: Any, **kwargs: Any) -> Callable[..., Any]:
+                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
+
+    return functools.wraps(func)(wrapper)
+
+
+udf.pandas = functools.partial(udf, data_format='pandas')  # type: ignore
+udf.polars = functools.partial(udf, data_format='polars')  # type: ignore
+udf.arrow = functools.partial(udf, data_format='arrow')  # type: ignore
+udf.numpy = functools.partial(udf, data_format='numpy')  # type: ignore