singlestoredb 1.11.0__cp38-abi3-win32.whl → 1.12.1__cp38-abi3-win32.whl

singlestoredb/functions/signature.py

@@ -1,4 +1,5 @@
 #!/usr/bin/env python3
+import dataclasses
 import datetime
 import inspect
 import numbers
@@ -22,6 +23,12 @@ try:
 except ImportError:
     has_numpy = False
 
+try:
+    import pydantic
+    has_pydantic = True
+except ImportError:
+    has_pydantic = False
+
 from . import dtypes as dt
 from ..mysql.converters import escape_item  # type: ignore
 
@@ -243,6 +250,9 @@ def classify_dtype(dtype: Any) -> str:
     if isinstance(dtype, list):
         return '|'.join(classify_dtype(x) for x in dtype)
 
+    if isinstance(dtype, str):
+        return sql_to_dtype(dtype)
+
     # Specific types
     if dtype is None or dtype is type(None):  # noqa: E721
         return 'null'
@@ -253,6 +263,21 @@ def classify_dtype(dtype: Any) -> str:
     if dtype is bool:
         return 'bool'
 
+    if dataclasses.is_dataclass(dtype):
+        fields = dataclasses.fields(dtype)
+        item_dtypes = ','.join(
+            f'{classify_dtype(simplify_dtype(x.type))}' for x in fields
+        )
+        return f'tuple[{item_dtypes}]'
+
+    if has_pydantic and inspect.isclass(dtype) and issubclass(dtype, pydantic.BaseModel):
+        fields = dtype.model_fields.values()
+        item_dtypes = ','.join(
+            f'{classify_dtype(simplify_dtype(x.annotation))}'  # type: ignore
+            for x in fields
+        )
+        return f'tuple[{item_dtypes}]'
+
     if not inspect.isclass(dtype):
         # Check for compound types
         origin = typing.get_origin(dtype)
@@ -261,7 +286,7 @@ def classify_dtype(dtype: Any) -> str:
             if origin is Tuple:
                 args = typing.get_args(dtype)
                 item_dtypes = ','.join(classify_dtype(x) for x in args)
-                return f'tuple:{item_dtypes}'
+                return f'tuple[{item_dtypes}]'
 
             # Array types
             elif issubclass(origin, array_types):
@@ -312,7 +337,10 @@ def classify_dtype(dtype: Any) -> str:
     if is_int:
         return int_type_map.get(name, 'int64')
 
-    raise TypeError(f'unsupported type annotation: {dtype}')
+    raise TypeError(
+        f'unsupported type annotation: {dtype}; '
+        'use `args`/`returns` on the @udf/@tvf decotator to specify the data type',
+    )
 
 
 def collapse_dtypes(dtypes: Union[str, List[str]]) -> str:
@@ -428,6 +456,7 @@ def get_signature(func: Callable[..., Any], name: Optional[str] = None) -> Dict[
     args: List[Dict[str, Any]] = []
     attrs = getattr(func, '_singlestoredb_attrs', {})
     name = attrs.get('name', name if name else func.__name__)
+    function_type = attrs.get('function_type', 'udf')
     out: Dict[str, Any] = dict(name=name, args=args)
 
     arg_names = [x for x in signature.parameters]
@@ -448,6 +477,7 @@ def get_signature(func: Callable[..., Any], name: Optional[str] = None) -> Dict[
 
     args_overrides = attrs.get('args', None)
     returns_overrides = attrs.get('returns', None)
+    output_fields = attrs.get('output_fields', None)
 
     spec_diff = set(arg_names).difference(set(annotations.keys()))
 
@@ -488,7 +518,7 @@ def get_signature(func: Callable[..., Any], name: Optional[str] = None) -> Dict[
             arg_type = collapse_dtypes([
                 classify_dtype(x) for x in simplify_dtype(annotations[arg])
             ])
-            sql = dtype_to_sql(arg_type)
+            sql = dtype_to_sql(arg_type, function_type=function_type)
         args.append(dict(name=arg, dtype=arg_type, sql=sql, default=defaults[i]))
 
     if returns_overrides is None \
@@ -498,13 +528,56 @@ def get_signature(func: Callable[..., Any], name: Optional[str] = None) -> Dict[
     if isinstance(returns_overrides, str):
         sql = returns_overrides
         out_type = sql_to_dtype(sql)
+    elif isinstance(returns_overrides, list):
+        if not output_fields:
+            output_fields = [
+                string.ascii_letters[i] for i in range(len(returns_overrides))
+            ]
+        out_type = 'tuple[' + collapse_dtypes([
+            classify_dtype(x)
+            for x in simplify_dtype(returns_overrides)
+        ]).replace('|', ',') + ']'
+        sql = dtype_to_sql(
+            out_type, function_type=function_type, field_names=output_fields,
+        )
+    elif dataclasses.is_dataclass(returns_overrides):
+        out_type = collapse_dtypes([
+            classify_dtype(x)
+            for x in simplify_dtype([x.type for x in returns_overrides.fields])
+        ])
+        sql = dtype_to_sql(
+            out_type,
+            function_type=function_type,
+            field_names=[x.name for x in returns_overrides.fields],
+        )
+    elif has_pydantic and inspect.isclass(returns_overrides) \
+            and issubclass(returns_overrides, pydantic.BaseModel):
+        out_type = collapse_dtypes([
+            classify_dtype(x)
+            for x in simplify_dtype([x for x in returns_overrides.model_fields.values()])
+        ])
+        sql = dtype_to_sql(
+            out_type,
+            function_type=function_type,
+            field_names=[x for x in returns_overrides.model_fields.keys()],
+        )
     elif returns_overrides is not None and not isinstance(returns_overrides, str):
         raise TypeError(f'unrecognized type for return value: {returns_overrides}')
     else:
+        if not output_fields:
+            if dataclasses.is_dataclass(signature.return_annotation):
+                output_fields = [
+                    x.name for x in dataclasses.fields(signature.return_annotation)
+                ]
+            elif has_pydantic and inspect.isclass(signature.return_annotation) \
+                    and issubclass(signature.return_annotation, pydantic.BaseModel):
+                output_fields = list(signature.return_annotation.model_fields.keys())
         out_type = collapse_dtypes([
             classify_dtype(x) for x in simplify_dtype(signature.return_annotation)
         ])
-        sql = dtype_to_sql(out_type)
+        sql = dtype_to_sql(
+            out_type, function_type=function_type, field_names=output_fields,
+        )
     out['returns'] = dict(dtype=out_type, sql=sql, default=None)
 
     copied_keys = ['database', 'environment', 'packages', 'resources', 'replace']
@@ -559,7 +632,12 @@ def sql_to_dtype(sql: str) -> str:
     return dtype
 
 
-def dtype_to_sql(dtype: str, default: Any = None) -> str:
+def dtype_to_sql(
+    dtype: str,
+    default: Any = None,
+    field_names: Optional[List[str]] = None,
+    function_type: str = 'udf',
+) -> str:
     """
     Convert a collapsed dtype string to a SQL type.
 
@@ -569,6 +647,8 @@ def dtype_to_sql(dtype: str, default: Any = None) -> str:
         Simplified data type string
     default : Any, optional
         Default value
+    field_names : List[str], optional
+        Field names for tuple types
 
     Returns
     -------
@@ -592,7 +672,7 @@ def dtype_to_sql(dtype: str, default: Any = None) -> str:
     if dtype.startswith('array['):
         _, dtypes = dtype.split('[', 1)
         dtypes = dtypes[:-1]
-        item_dtype = dtype_to_sql(dtypes)
+        item_dtype = dtype_to_sql(dtypes, function_type=function_type)
         return f'ARRAY({item_dtype}){nullable}{default_clause}'
 
     if dtype.startswith('tuple['):
@@ -600,11 +680,22 @@ def dtype_to_sql(dtype: str, default: Any = None) -> str:
         dtypes = dtypes[:-1]
         item_dtypes = []
         for i, item in enumerate(dtypes.split(',')):
-            name = string.ascii_letters[i]
+            if field_names:
+                name = field_names[i]
+            else:
+                name = string.ascii_letters[i]
             if '=' in item:
                 name, item = item.split('=', 1)
-            item_dtypes.append(name + ' ' + dtype_to_sql(item))
-        return f'RECORD({", ".join(item_dtypes)}){nullable}{default_clause}'
+            item_dtypes.append(
+                f'`{name}` ' + dtype_to_sql(item, function_type=function_type),
+            )
+        if function_type == 'udf':
+            return f'RECORD({", ".join(item_dtypes)}){nullable}{default_clause}'
+        else:
+            return re.sub(
+                r' NOT NULL\s*$', r'',
+                f'TABLE({", ".join(item_dtypes)}){nullable}{default_clause}',
+            )
 
     return f'{sql_type_map[dtype]}{nullable}{default_clause}'
 

singlestoredb/fusion/handlers/export.py

@@ -43,8 +43,6 @@ class CreateClusterIdentity(SQLHandler):
 
     Remarks
     -------
-    * ``FROM <table>`` specifies the SingleStore table to export. The same name will
-      be used for the exported table.
     * ``CATALOG`` specifies the details of the catalog to connect to.
     * ``LINK`` specifies the details of the data storage to connect to.
 
@@ -69,6 +67,8 @@ class CreateClusterIdentity(SQLHandler):
 
     """
 
+    _enabled = False
+
     def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]:
         # Catalog
         catalog_config = json.loads(params['catalog'].get('catalog_config', '{}') or '{}')
@@ -110,11 +110,34 @@ class CreateExport(SQLHandler):
         from_table
         catalog
         storage
+        [ partition_by ]
+        [ order_by ]
+        [ properties ]
     ;
 
     # From table
     from_table = FROM <table>
 
+    # Transforms
+    _col_transform = { VOID | IDENTITY | YEAR | MONTH | DAY | HOUR } ( _transform_col )
+    _transform_col = <column>
+    _arg_transform = { BUCKET | TRUNCATE } ( _transform_col <comma> _transform_arg )
+    _transform_arg = <integer>
+    transform = { _col_transform | _arg_transform }
+
+    # Partitions
+    partition_by = PARTITION BY partition_key,...
+    partition_key = transform
+
+    # Sort order
+    order_by = ORDER BY sort_key,...
+    sort_key = transform [ direction ] [ null_order ]
+    direction = { ASC | DESC | ASCENDING | DESCENDING }
+    null_order = { NULLS_FIRST | NULLS_LAST }
+
+    # Properties
+    properties = PROPERTIES '<json>'
+
     # Catolog
     catalog = CATALOG [ _catalog_config ] [ _catalog_creds ]
     _catalog_config = CONFIG '<catalog-config>'
@@ -163,6 +186,8 @@ class CreateExport(SQLHandler):
 
     """  # noqa
 
+    _enabled = False
+
     def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]:
         # From table
         if isinstance(params['from_table'], str):
@@ -189,6 +214,32 @@ class CreateExport(SQLHandler):
         if wsg._manager is None:
             raise TypeError('no workspace manager is associated with workspace group')
 
+        partition_by = []
+        if params['partition_by']:
+            for key in params['partition_by']:
+                transform = key['partition_key']['transform']['col_transform']
+                part = {}
+                part['transform'] = transform[0].lower()
+                part['name'] = transform[-1]['transform_col']
+                partition_by.append(part)
+
+        order_by = []
+        if params['order_by'] and params['order_by']['by']:
+            for key in params['order_by']['by']:
+                transform = key['transform']['col_transform']
+                order = {}
+                order['transform'] = transform[0].lower()
+                order['name'] = transform[-1]['transform_col']
+                order['direction'] = 'ascending'
+                order['null_order'] = 'nulls_first'
+                if key.get('direction'):
+                    if 'desc' in key['direction'].lower():
+                        order['direction'] = 'descending'
+                if key.get('null_order'):
+                    if 'last' in key['null_order'].lower():
+                        order['null_order'] = 'nulls_last'
+                order_by.append(order)
+
         out = ExportService(
             wsg,
             from_database,
@@ -196,6 +247,9 @@ class CreateExport(SQLHandler):
             dict(**catalog_config, **catalog_creds),
             dict(**storage_config, **storage_creds),
             columns=None,
+            partition_by=partition_by or None,
+            order_by=order_by or None,
+            properties=json.loads(params['properties']) if params['properties'] else None,
         ).start()
 
         res = FusionSQLResult()
@@ -217,6 +271,8 @@ class ShowExport(SQLHandler):
 
     """
 
+    _enabled = False
+
     def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]:
         wsg = get_workspace_group({})
         out = ExportStatus(params['export_id'], wsg)

singlestoredb/fusion/handlers/files.py

@@ -209,7 +209,7 @@ class UploadPersonalFileHandler(UploadFileHandler):
         FROM local_path [ overwrite ];
 
     # Path to file
-    path = '<path>'
+    path = '<filename>'
 
     # Path to local file
     local_path = '<local-path>'
@@ -223,7 +223,7 @@ class UploadPersonalFileHandler(UploadFileHandler):
 
     Arguments
     ---------
-    * ``<path>``: The path in the personal/shared space where the file is uploaded.
+    * ``<filename>``: The filename in the personal/shared space where the file is uploaded.
     * ``<local-path>``: The path to the file to upload in the local
       directory.
 
@@ -237,7 +237,7 @@ class UploadPersonalFileHandler(UploadFileHandler):
     The following command uploads a file to a personal/shared space and overwrite any
     existing files at the specified path::
 
-        UPLOAD PERSONAL FILE TO '/data/stats.csv'
+        UPLOAD PERSONAL FILE TO 'stats.csv'
             FROM '/tmp/user/stats.csv' OVERWRITE;
 
     See Also
@@ -259,7 +259,7 @@ class UploadSharedFileHandler(UploadFileHandler):
         FROM local_path [ overwrite ];
 
     # Path to file
-    path = '<path>'
+    path = '<filename>'
 
     # Path to local file
     local_path = '<local-path>'
@@ -273,7 +273,7 @@ class UploadSharedFileHandler(UploadFileHandler):
 
     Arguments
     ---------
-    * ``<path>``: The path in the personal/shared space where the file is uploaded.
+    * ``<filename>``: The filename in the personal/shared space where the file is uploaded.
     * ``<local-path>``: The path to the file to upload in the local
       directory.
 
@@ -287,7 +287,7 @@ class UploadSharedFileHandler(UploadFileHandler):
     The following command uploads a file to a personal/shared space and overwrite any
     existing files at the specified path::
 
-        UPLOAD SHARED FILE TO '/data/stats.csv'
+        UPLOAD SHARED FILE TO 'stats.csv'
             FROM '/tmp/user/stats.csv' OVERWRITE;
 
     See Also

singlestoredb/fusion/handlers/utils.py

@@ -7,13 +7,11 @@ from typing import Optional
 from typing import Union
 
 from ...exceptions import ManagementError
+from ...management import files as mgmt_files
 from ...management import manage_workspaces
 from ...management.files import FilesManager
 from ...management.files import FileSpace
 from ...management.files import manage_files
-from ...management.files import MODELS_SPACE
-from ...management.files import PERSONAL_SPACE
-from ...management.files import SHARED_SPACE
 from ...management.workspace import StarterWorkspace
 from ...management.workspace import Workspace
 from ...management.workspace import WorkspaceGroup
@@ -190,6 +188,10 @@ def get_deployment(
         * params['group']['deployment_id']
         * params['in_deployment']['deployment_name']
         * params['in_deployment']['deployment_id']
+        * params['in']['in_group']['deployment_name']
+        * params['in']['in_group']['deployment_id']
+        * params['in']['in_deployment']['deployment_name']
+        * params['in']['in_deployment']['deployment_id']
 
     Or, from the SINGLESTOREDB_WORKSPACE_GROUP
     or SINGLESTOREDB_CLUSTER environment variables.
@@ -199,7 +201,9 @@ def get_deployment(
 
     deployment_name = params.get('deployment_name') or \
         (params.get('in_deployment') or {}).get('deployment_name') or \
-        (params.get('group') or {}).get('deployment_name')
+        (params.get('group') or {}).get('deployment_name') or \
+        ((params.get('in') or {}).get('in_group') or {}).get('deployment_name') or \
+        ((params.get('in') or {}).get('in_deployment') or {}).get('deployment_name')
     if deployment_name:
         workspace_groups = [
             x for x in manager.workspace_groups
@@ -239,7 +243,9 @@ def get_deployment(
 
     deployment_id = params.get('deployment_id') or \
         (params.get('in_deployment') or {}).get('deployment_id') or \
-        (params.get('group') or {}).get('deployment_id')
+        (params.get('group') or {}).get('deployment_id') or \
+        ((params.get('in') or {}).get('in_group') or {}).get('deployment_id') or \
+        ((params.get('in') or {}).get('in_deployment') or {}).get('deployment_id')
     if deployment_id:
         try:
             return manager.get_workspace_group(deployment_id)
@@ -298,11 +304,11 @@ def get_file_space(params: Dict[str, Any]) -> FileSpace:
     if file_location:
         file_location_lower_case = file_location.lower()
 
-        if file_location_lower_case == PERSONAL_SPACE:
+        if file_location_lower_case == mgmt_files.PERSONAL_SPACE:
             return manager.personal_space
-        elif file_location_lower_case == SHARED_SPACE:
+        elif file_location_lower_case == mgmt_files.SHARED_SPACE:
             return manager.shared_space
-        elif file_location_lower_case == MODELS_SPACE:
+        elif file_location_lower_case == mgmt_files.MODELS_SPACE:
             return manager.models_space
         else:
             raise ValueError(f'invalid file location: {file_location}')

singlestoredb/fusion/result.py

@@ -241,7 +241,7 @@ class FusionSQLResult(object):
         for row in self.rows:
             found = True
             for i, liker in likers:
-                if not liker.match(row[i]):
+                if row[i] is None or not liker.match(row[i]):
                     found = False
                     break
             if found:

singlestoredb/http/connection.py

@@ -972,6 +972,10 @@ class Connection(connection.Connection):
 
     def __init__(self, **kwargs: Any):
         from .. import __version__ as client_version
+
+        if 'SINGLESTOREDB_WORKLOAD_TYPE' in os.environ:
+            client_version += '+' + os.environ['SINGLESTOREDB_WORKLOAD_TYPE']
+
         connection.Connection.__init__(self, **kwargs)
 
         host = kwargs.get('host', get_option('host'))

singlestoredb/management/export.py

@@ -24,6 +24,9 @@ class ExportService(object):
     catalog_info: Dict[str, Any]
     storage_info: Dict[str, Any]
     columns: Optional[List[str]]
+    partition_by: Optional[List[Dict[str, str]]]
+    order_by: Optional[List[Dict[str, Dict[str, str]]]]
+    properties: Optional[Dict[str, Any]]
 
     def __init__(
         self,
@@ -32,7 +35,10 @@ class ExportService(object):
         table: str,
         catalog_info: Union[str, Dict[str, Any]],
         storage_info: Union[str, Dict[str, Any]],
-        columns: Optional[List[str]],
+        columns: Optional[List[str]] = None,
+        partition_by: Optional[List[Dict[str, str]]] = None,
+        order_by: Optional[List[Dict[str, Dict[str, str]]]] = None,
+        properties: Optional[Dict[str, Any]] = None,
     ):
         #: Workspace group
         self.workspace_group = workspace_group
@@ -58,6 +64,10 @@ class ExportService(object):
         else:
             self.storage_info = copy.copy(storage_info)
 
+        self.partition_by = partition_by or None
+        self.order_by = order_by or None
+        self.properties = properties or None
+
         self._manager: Optional[WorkspaceManager] = workspace_group._manager
 
     def __str__(self) -> str:
@@ -93,14 +103,27 @@ class ExportService(object):
                 msg='No workspace manager is associated with this object.',
             )
 
+        partition_spec = None
+        if self.partition_by:
+            partition_spec = dict(partitions=self.partition_by)
+
+        sort_order_spec = None
+        if self.order_by:
+            sort_order_spec = dict(keys=self.order_by)
+
         out = self._manager._post(
             f'workspaceGroups/{self.workspace_group.id}/egress/startTableEgress',
-            json=dict(
-                databaseName=self.database,
-                tableName=self.table,
-                storageInfo=self.storage_info,
-                catalogInfo=self.catalog_info,
-            ),
+            json={
+                k: v for k, v in dict(
+                    databaseName=self.database,
+                    tableName=self.table,
+                    storageInfo=self.storage_info,
+                    catalogInfo=self.catalog_info,
+                    partitionSpec=partition_spec,
+                    sortOrderSpec=sort_order_spec,
+                    properties=self.properties,
+                ).items() if v is not None
+            },
         )
 
         return ExportStatus(out.json()['egressID'], self.workspace_group)

singlestoredb/management/files.py

@@ -10,11 +10,9 @@ import re
 from abc import ABC
 from abc import abstractmethod
 from typing import Any
-from typing import BinaryIO
 from typing import Dict
 from typing import List
 from typing import Optional
-from typing import TextIO
 from typing import Union
 
 from .. import config
@@ -362,7 +360,7 @@ class FileLocation(ABC):
     @abstractmethod
     def upload_file(
         self,
-        local_path: Union[PathLike, TextIO, BinaryIO],
+        local_path: Union[PathLike, io.IOBase],
         path: PathLike,
         *,
         overwrite: bool = False,
@@ -385,7 +383,7 @@ class FileLocation(ABC):
     @abstractmethod
     def _upload(
         self,
-        content: Union[str, bytes, TextIO, BinaryIO],
+        content: Union[str, bytes, io.IOBase],
         path: PathLike,
         *,
         overwrite: bool = False,
@@ -628,7 +626,7 @@ class FileSpace(FileLocation):
 
     def upload_file(
         self,
-        local_path: Union[PathLike, TextIO, BinaryIO],
+        local_path: Union[PathLike, io.IOBase],
         path: PathLike,
         *,
         overwrite: bool = False,
@@ -646,7 +644,7 @@ class FileSpace(FileLocation):
             Should the ``path`` be overwritten if it exists already?
 
         """
-        if isinstance(local_path, (TextIO, BinaryIO)):
+        if isinstance(local_path, io.IOBase):
             pass
         elif not os.path.isfile(local_path):
             raise IsADirectoryError(f'local path is not a file: {local_path}')
@@ -657,8 +655,9 @@ class FileSpace(FileLocation):
 
             self.remove(path)
 
-        if isinstance(local_path, (TextIO, BinaryIO)):
+        if isinstance(local_path, io.IOBase):
             return self._upload(local_path, path, overwrite=overwrite)
+
         return self._upload(open(local_path, 'rb'), path, overwrite=overwrite)
 
     def upload_folder(
@@ -727,7 +726,7 @@ class FileSpace(FileLocation):
 
     def _upload(
         self,
-        content: Union[str, bytes, TextIO, BinaryIO],
+        content: Union[str, bytes, io.IOBase],
         path: PathLike,
         *,
         overwrite: bool = False,

singlestoredb/management/workspace.py

@@ -10,11 +10,9 @@ import re
 import time
 from collections.abc import Mapping
 from typing import Any
-from typing import BinaryIO
 from typing import Dict
 from typing import List
 from typing import Optional
-from typing import TextIO
 from typing import Union
 
 from .. import config
@@ -165,7 +163,7 @@ class Stage(FileLocation):
 
     def upload_file(
         self,
-        local_path: Union[PathLike, TextIO, BinaryIO],
+        local_path: Union[PathLike, io.IOBase],
         stage_path: PathLike,
         *,
         overwrite: bool = False,
@@ -183,7 +181,7 @@ class Stage(FileLocation):
             Should the ``stage_path`` be overwritten if it exists already?
 
         """
-        if isinstance(local_path, (TextIO, BinaryIO)):
+        if isinstance(local_path, io.IOBase):
             pass
         elif not os.path.isfile(local_path):
             raise IsADirectoryError(f'local path is not a file: {local_path}')
@@ -194,8 +192,9 @@ class Stage(FileLocation):
 
             self.remove(stage_path)
 
-        if isinstance(local_path, (TextIO, BinaryIO)):
+        if isinstance(local_path, io.IOBase):
             return self._upload(local_path, stage_path, overwrite=overwrite)
+
         return self._upload(open(local_path, 'rb'), stage_path, overwrite=overwrite)
 
     def upload_folder(
@@ -258,7 +257,7 @@ class Stage(FileLocation):
 
     def _upload(
         self,
-        content: Union[str, bytes, TextIO, BinaryIO],
+        content: Union[str, bytes, io.IOBase],
         stage_path: PathLike,
         *,
         overwrite: bool = False,