singlestoredb 1.7.2__cp38-abi3-win_amd64.whl → 1.9.0__cp38-abi3-win_amd64.whl

singlestoredb/fusion/handlers/utils.py

@@ -4,9 +4,16 @@ import os
 from typing import Any
 from typing import Dict
 from typing import Optional
+from typing import Union
 
 from ...exceptions import ManagementError
 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 PERSONAL_SPACE
+from ...management.files import SHARED_SPACE
+from ...management.workspace import StarterWorkspace
 from ...management.workspace import Workspace
 from ...management.workspace import WorkspaceGroup
 from ...management.workspace import WorkspaceManager
@@ -17,6 +24,11 @@ def get_workspace_manager() -> WorkspaceManager:
     return manage_workspaces()
 
 
+def get_files_manager() -> FilesManager:
+    """Return a new files manager."""
+    return manage_files()
+
+
 def dt_isoformat(dt: Optional[datetime.datetime]) -> Optional[str]:
     """Convert datetime to string."""
     if dt is None:
@@ -160,3 +172,139 @@ def get_workspace(params: Dict[str, Any]) -> Workspace:
         raise ValueError('clusters and shared workspaces are not currently supported')
 
     raise KeyError('no workspace was specified')
+
+
+def get_deployment(
+        params: Dict[str, Any],
+) -> Union[WorkspaceGroup, StarterWorkspace]:
+    """
+    Find a starter workspace matching deployment_id or deployment_name.
+
+    This function will get a starter workspace or ID from the
+    following parameters:
+
+        * params['deployment_name']
+        * params['deployment_id']
+        * params['group']['deployment_name']
+        * params['group']['deployment_id']
+        * params['in_deployment']['deployment_name']
+        * params['in_deployment']['deployment_id']
+
+    Or, from the SINGLESTOREDB_WORKSPACE_GROUP
+    or SINGLESTOREDB_CLUSTER environment variables.
+
+    """
+    manager = get_workspace_manager()
+
+    deployment_name = params.get('deployment_name') or \
+        (params.get('in_deployment') or {}).get('deployment_name') or \
+        (params.get('group') or {}).get('deployment_name')
+    if deployment_name:
+        workspace_groups = [
+            x for x in manager.workspace_groups
+            if x.name == deployment_name
+        ]
+
+        starter_workspaces = []
+        if not workspace_groups:
+            filtered_starter_workspaces = [
+                x for x in manager.starter_workspaces
+                if x.name == deployment_name
+            ]
+
+            if not filtered_starter_workspaces:
+                raise KeyError(
+                    f'no deployment found with name: {deployment_name}',
+                )
+
+            starter_workspaces = filtered_starter_workspaces
+
+        if len(workspace_groups) > 1:
+            ids = ', '.join(x.id for x in workspace_groups)
+            raise ValueError(
+                f'more than one workspace group with given name was found: {ids}',
+            )
+
+        if len(starter_workspaces) > 1:
+            ids = ', '.join(x.id for x in starter_workspaces)
+            raise ValueError(
+                f'more than one starter workspace with given name was found: {ids}',
+            )
+
+        if workspace_groups:
+            return workspace_groups[0]
+        else:
+            return starter_workspaces[0]
+
+    deployment_id = params.get('deployment_id') or \
+        (params.get('in_deployment') or {}).get('deployment_id') or \
+        (params.get('group') or {}).get('deployment_id')
+    if deployment_id:
+        try:
+            return manager.get_workspace_group(deployment_id)
+        except ManagementError as exc:
+            if exc.errno == 404:
+                try:
+                    return manager.get_starter_workspace(deployment_id)
+                except ManagementError as exc:
+                    if exc.errno == 404:
+                        raise KeyError(f'no deployment found with ID: {deployment_id}')
+                    raise
+            else:
+                raise
+
+    if os.environ.get('SINGLESTOREDB_WORKSPACE_GROUP'):
+        try:
+            return manager.get_workspace_group(
+                os.environ['SINGLESTOREDB_WORKSPACE_GROUP'],
+            )
+        except ManagementError as exc:
+            if exc.errno == 404:
+                raise KeyError(
+                    'no workspace found with ID: '
+                    f'{os.environ["SINGLESTOREDB_WORKSPACE_GROUP"]}',
+                )
+            raise
+
+    if os.environ.get('SINGLESTOREDB_CLUSTER'):
+        try:
+            return manager.get_starter_workspace(
+                os.environ['SINGLESTOREDB_CLUSTER'],
+            )
+        except ManagementError as exc:
+            if exc.errno == 404:
+                raise KeyError(
+                    'no starter workspace found with ID: '
+                    f'{os.environ["SINGLESTOREDB_CLUSTER"]}',
+                )
+            raise
+
+    raise KeyError('no deployment was specified')
+
+
+def get_file_space(params: Dict[str, Any]) -> FileSpace:
+    """
+    Retrieve the specified file space.
+
+    This function will get a file space from the
+    following parameters:
+
+        * params['file_location']
+    """
+    manager = get_files_manager()
+
+    file_location = params.get('file_location')
+    if file_location:
+        file_location_lower_case = file_location.lower()
+        if (
+            file_location_lower_case != PERSONAL_SPACE and
+            file_location_lower_case != SHARED_SPACE
+        ):
+            raise ValueError(f'invalid file location: {file_location}')
+
+        if file_location_lower_case == PERSONAL_SPACE:
+            return manager.personal_space
+        elif file_location_lower_case == SHARED_SPACE:
+            return manager.shared_space
+
+    raise KeyError('no file space was specified')

singlestoredb/management/__init__.py

@@ -1,5 +1,6 @@
 #!/usr/bin/env python
 from .cluster import manage_cluster
+from .files import manage_files
 from .manager import get_token
 from .workspace import get_organization
 from .workspace import get_secret

singlestoredb/management/export.py

@@ -0,0 +1,146 @@
+#!/usr/bin/env python
+"""SingleStoreDB export service."""
+from __future__ import annotations
+
+import copy
+import json
+from typing import Any
+from typing import Dict
+from typing import List
+from typing import Optional
+from typing import Union
+
+from .. import ManagementError
+from .utils import vars_to_str
+from .workspace import WorkspaceGroup
+from .workspace import WorkspaceManager
+
+
+class ExportService(object):
+    """Export service."""
+
+    database: str
+    table: str
+    catalog_info: Dict[str, Any]
+    storage_info: Dict[str, Any]
+    columns: Optional[List[str]]
+
+    def __init__(
+        self,
+        workspace_group: WorkspaceGroup,
+        database: str,
+        table: str,
+        catalog_info: Union[str, Dict[str, Any]],
+        storage_info: Union[str, Dict[str, Any]],
+        columns: Optional[List[str]],
+    ):
+        #: Workspace group
+        self.workspace_group = workspace_group
+
+        #: Name of SingleStoreDB database
+        self.database = database
+
+        #: Name of SingleStoreDB table
+        self.table = table
+
+        #: List of columns to export
+        self.columns = columns
+
+        #: Catalog
+        if isinstance(catalog_info, str):
+            self.catalog_info = json.loads(catalog_info)
+        else:
+            self.catalog_info = copy.copy(catalog_info)
+
+        #: Storage
+        if isinstance(storage_info, str):
+            self.storage_info = json.loads(storage_info)
+        else:
+            self.storage_info = copy.copy(storage_info)
+
+        self._manager: Optional[WorkspaceManager] = workspace_group._manager
+
+    def __str__(self) -> str:
+        """Return string representation."""
+        return vars_to_str(self)
+
+    def __repr__(self) -> str:
+        """Return string representation."""
+        return str(self)
+
+    def create_cluster_identity(self) -> Dict[str, Any]:
+        """Create a cluster identity."""
+        if self._manager is None:
+            raise ManagementError(
+                msg='No workspace manager is associated with this object.',
+            )
+
+        out = self._manager._post(
+            f'workspaceGroups/{self.workspace_group.id}/'
+            'egress/createEgressClusterIdentity',
+            json=dict(
+                catalogInfo=self.catalog_info,
+                storageInfo=self.storage_info,
+            ),
+        )
+
+        return out.json()
+
+    def start(self, tags: Optional[List[str]] = None) -> 'ExportStatus':
+        """Start the export process."""
+        if self._manager is None:
+            raise ManagementError(
+                msg='No workspace manager is associated with this object.',
+            )
+
+        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,
+            ),
+        )
+
+        return ExportStatus(out.json()['egressID'], self.workspace_group)
+
+
+class ExportStatus(object):
+
+    export_id: str
+
+    def __init__(self, export_id: str, workspace_group: WorkspaceGroup):
+        self.export_id = export_id
+        self.workspace_group = workspace_group
+        self._manager: Optional[WorkspaceManager] = workspace_group._manager
+
+    def _info(self) -> Dict[str, Any]:
+        """Return export status."""
+        if self._manager is None:
+            raise ManagementError(
+                msg='No workspace manager is associated with this object.',
+            )
+
+        out = self._manager._get(
+            f'workspaceGroups/{self.workspace_group.id}/egress/tableEgressStatus',
+            json=dict(egressID=self.export_id),
+        )
+
+        return out.json()
+
+    @property
+    def status(self) -> str:
+        """Return export status."""
+        return self._info().get('status', 'Unknown')
+
+    @property
+    def message(self) -> str:
+        """Return export status message."""
+        return self._info().get('statusMsg', '')
+
+    def __str__(self) -> str:
+        return self.status
+
+    def __repr__(self) -> str:
+        return self.status