singlestoredb 1.10.0__py3-none-any.whl → 1.12.0__py3-none-any.whl

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

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import datetime
+import glob
 import io
 import os
 import re
@@ -23,9 +24,9 @@ from .utils import PathLike
 from .utils import to_datetime
 from .utils import vars_to_str
 
-
 PERSONAL_SPACE = 'personal'
 SHARED_SPACE = 'shared'
+MODELS_SPACE = 'models'
 
 
 class FilesObject(object):
@@ -35,8 +36,8 @@ class FilesObject(object):
     It can belong to either a workspace stage or personal/shared space.
 
     This object is not instantiated directly. It is used in the results
-    of various operations in ``WorkspaceGroup.stage``, ``FilesManager.personal_space``
-    and ``FilesManager.shared_space`` methods.
+    of various operations in ``WorkspaceGroup.stage``, ``FilesManager.personal_space``,
+    ``FilesManager.shared_space`` and ``FilesManager.models_space`` methods.
 
     """
 
@@ -513,6 +514,11 @@ class FilesManager(Manager):
         """Return the shared file space."""
         return FileSpace(SHARED_SPACE, self)
 
+    @property
+    def models_space(self) -> FileSpace:
+        """Return the models file space."""
+        return FileSpace(MODELS_SPACE, self)
+
 
 def manage_files(
     access_token: Optional[str] = None,
@@ -551,7 +557,8 @@ class FileSpace(FileLocation):
     FileSpace manager.
 
     This object is not instantiated directly.
-    It is returned by ``FilesManager.personal_space`` or ``FilesManager.shared_space``.
+    It is returned by ``FilesManager.personal_space``, ``FilesManager.shared_space``
+    or ``FileManger.models_space``.
 
     """
 
@@ -687,10 +694,36 @@ class FileSpace(FileLocation):
             ignore all '*.pyc' files in the directory tree
 
         """
-        raise ManagementError(
-            msg='Operation not supported: directories are currently not allowed '
-                'in Files API',
-        )
+        if not os.path.isdir(local_path):
+            raise NotADirectoryError(f'local path is not a directory: {local_path}')
+
+        if not path:
+            path = local_path
+
+        ignore_files = set()
+        if ignore:
+            if isinstance(ignore, list):
+                for item in ignore:
+                    ignore_files.update(glob.glob(str(item), recursive=recursive))
+            else:
+                ignore_files.update(glob.glob(str(ignore), recursive=recursive))
+
+        for dir_path, _, files in os.walk(str(local_path)):
+            for fname in files:
+                if ignore_files and fname in ignore_files:
+                    continue
+
+                local_file_path = os.path.join(dir_path, fname)
+                remote_path = os.path.join(
+                    path,
+                    local_file_path.lstrip(str(local_path)),
+                )
+                self.upload_file(
+                    local_path=local_file_path,
+                    path=remote_path,
+                    overwrite=overwrite,
+                )
+        return self.info(path)
 
     def _upload(
         self,
@@ -875,15 +908,30 @@ class FileSpace(FileLocation):
                 return False
             raise
 
-    def _list_root_dir(self) -> List[str]:
+    def _listdir(self, path: PathLike, *, recursive: bool = False) -> List[str]:
         """
-        Return the names of files in the root directory.
+        Return the names of files in a directory.
+
         Parameters
         ----------
+        path : Path or str
+            Path to the folder
+        recursive : bool, optional
+            Should folders be listed recursively?
+
         """
         res = self._manager._get(
-            f'files/fs/{self._location}',
+            f'files/fs/{self._location}/{path}',
         ).json()
+
+        if recursive:
+            out = []
+            for item in res['content'] or []:
+                out.append(item['path'])
+                if item['type'] == 'directory':
+                    out.extend(self._listdir(item['path'], recursive=recursive))
+            return out
+
         return [x['path'] for x in res['content'] or []]
 
     def listdir(
@@ -905,13 +953,17 @@ class FileSpace(FileLocation):
         List[str]
 
         """
-        if path == '' or path == '/':
-            return self._list_root_dir()
+        path = re.sub(r'^(\./|/)+', r'', str(path))
+        path = re.sub(r'/+$', r'', path) + '/'
 
-        raise ManagementError(
-            msg='Operation not supported: directories are currently not allowed '
-                'in Files API',
-        )
+        if not self.is_dir(path):
+            raise NotADirectoryError(f'path is not a directory: {path}')
+
+        out = self._listdir(path, recursive=recursive)
+        if path != '/':
+            path_n = len(path.split('/')) - 1
+            out = ['/'.join(x.split('/')[path_n:]) for x in out]
+        return out
 
     def download_file(
         self,
@@ -973,17 +1025,28 @@ class FileSpace(FileLocation):
         Parameters
         ----------
         path : Path or str
-            Path to the file
+            Directory path
         local_path : Path or str
             Path to local directory target location
         overwrite : bool, optional
             Should an existing directory / files be overwritten if they exist?
 
         """
-        raise ManagementError(
-            msg='Operation not supported: directories are currently not allowed '
-                'in Files API',
-        )
+
+        if local_path is not None and not overwrite and os.path.exists(local_path):
+            raise OSError('target path already exists; use overwrite=True to replace')
+
+        if not self.is_dir(path):
+            raise NotADirectoryError(f'path is not a directory: {path}')
+
+        files = self.listdir(path, recursive=True)
+        for f in files:
+            remote_path = os.path.join(path, f)
+            if self.is_dir(remote_path):
+                continue
+            target = os.path.normpath(os.path.join(local_path, f))
+            os.makedirs(os.path.dirname(target), exist_ok=True)
+            self.download_file(remote_path, target, overwrite=overwrite)
 
     def remove(self, path: PathLike) -> None:
         """
@@ -1010,10 +1073,10 @@ class FileSpace(FileLocation):
             Path to the file location
 
         """
-        raise ManagementError(
-            msg='Operation not supported: directories are currently not allowed '
-                'in Files API',
-        )
+        if not self.is_dir(path):
+            raise NotADirectoryError('path is not a directory')
+
+        self._manager._delete(f'files/fs/{self._location}/{path}')
 
     def rmdir(self, path: PathLike) -> None:
         """

singlestoredb/mysql/connection.py

@@ -625,6 +625,9 @@ class Connection(BaseConnection):
 
         from .. import __version__ as VERSION_STRING
 
+        if 'SINGLESTOREDB_WORKLOAD_TYPE' in os.environ:
+            VERSION_STRING += '+' + os.environ['SINGLESTOREDB_WORKLOAD_TYPE']
+
         self._connect_attrs = {
             '_os': str(sys.platform),
             '_pid': str(os.getpid()),
@@ -729,6 +732,7 @@ class Connection(BaseConnection):
             return
         if self._closed:
             raise err.Error('Already closed')
+        events.unsubscribe(self._handle_event)
         self._closed = True
         if self._sock is None:
             return
@@ -985,11 +989,18 @@ class Connection(BaseConnection):
 
     def set_character_set(self, charset, collation=None):
         """
-        Set charaset (and collation) on the server.
+        Set session charaset (and collation) on the server.
 
-        Send "SET NAMES charset [COLLATE collation]" query.
+        Send "SET [COLLATION|CHARACTER_SET]_SERVER = [collation|charset]" query.
         Update Connection.encoding based on charset.
 
+        If charset/collation are being set to utf8mb4, the corresponding global
+        variables (COLLATION_SERVER and CHARACTER_SET_SERVER) must be also set
+        to utf8mb4. This is true by default for SingleStore 8.7+. For previuous
+        versions or non-default setting user must manully run the query
+        `SET global collation_connection = utf8mb4_general_ci`
+        replacing utf8mb4_general_ci with {collation}.
+
         Parameters
         ----------
         charset : str
@@ -1002,9 +1013,9 @@ class Connection(BaseConnection):
         encoding = charset_by_name(charset).encoding
 
         if collation:
-            query = f'SET NAMES {charset} COLLATE {collation}'
+            query = f'SET COLLATION_SERVER={collation}'
         else:
-            query = f'SET NAMES {charset}'
+            query = f'SET CHARACTER_SET_SERVER={charset}'
         self._execute_command(COMMAND.COM_QUERY, query)
         self._read_packet()
         self.charset = charset
@@ -1108,19 +1119,6 @@ class Connection(BaseConnection):
             self._get_server_information()
             self._request_authentication()
 
-            # Send "SET NAMES" query on init for:
-            # - Ensure charaset (and collation) is set to the server.
-            #   - collation_id in handshake packet may be ignored.
-            # - If collation is not specified, we don't know what is server's
-            #   default collation for the charset. For example, default collation
-            #   of utf8mb4 is:
-            #   - MySQL 5.7, MariaDB 10.x: utf8mb4_general_ci
-            #   - MySQL 8.0: utf8mb4_0900_ai_ci
-            #
-            # Reference:
-            # - https://github.com/PyMySQL/PyMySQL/issues/1092
-            # - https://github.com/wagtail/wagtail/issues/9477
-            # - https://zenn.dev/methane/articles/2023-mysql-collation (Japanese)
             self.set_character_set(self.charset, self.collation)
 
             if self.sql_mode is not None:
@@ -1846,7 +1844,7 @@ class MySQLResult:
 
     def _read_row_from_packet(self, packet):
         row = []
-        for encoding, converter in self.converters:
+        for i, (encoding, converter) in enumerate(self.converters):
             try:
                 data = packet.read_length_coded_string()
             except IndexError:
@@ -1855,7 +1853,15 @@ class MySQLResult:
                 break
             if data is not None:
                 if encoding is not None:
-                    data = data.decode(encoding, errors=self.encoding_errors)
+                    try:
+                        data = data.decode(encoding, errors=self.encoding_errors)
+                    except UnicodeDecodeError:
+                        raise UnicodeDecodeError(
+                            'failed to decode string value in column '
+                            f"'{self.fields[i].name}' using encoding '{encoding}'; " +
+                            "use the 'encoding_errors' option on the connection " +
+                            'to specify how to handle this error',
+                        )
                 if DEBUG:
                     print('DEBUG: DATA = ', data)
                 if converter is not None: