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

singlestoredb/fusion/handlers/stage.py

@@ -0,0 +1,257 @@
+#!/usr/bin/env python3
+from typing import Any
+from typing import Dict
+from typing import Optional
+
+from .. import result
+from ..handler import SQLHandler
+from ..result import FusionSQLResult
+from .utils import dt_isoformat
+from .utils import get_workspace_group
+
+
+class ShowStageFilesHandler(SQLHandler):
+    """
+    SHOW STAGE FILES [ in_group ] [ at_path ] [ <like> ] [ <order-by> ]
+                     [ <limit> ] [ recursive ] [ extended ];
+
+    # Workspace group
+    in_group = IN GROUP { group_id | group_name }
+
+    # ID of group
+    group_id = ID '<group-id>'
+
+    # Name of group
+    group_name = '<group-name>'
+
+    # Stage path to list
+    at_path = AT '<path>'
+
+    # Should the listing be recursive?
+    recursive = RECURSIVE
+
+    # Should extended attributes be shown?
+    extended = EXTENDED
+
+    """
+
+    def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]:
+        wg = get_workspace_group(params)
+
+        res = FusionSQLResult()
+        res.add_field('Name', result.STRING)
+
+        if params['extended']:
+            res.add_field('Type', result.STRING)
+            res.add_field('Size', result.INTEGER)
+            res.add_field('Writable', result.STRING)
+            res.add_field('CreatedAt', result.DATETIME)
+            res.add_field('LastModifiedAt', result.DATETIME)
+
+            files = []
+            for x in wg.stage.listdir(
+                params['at_path'] or '/',
+                recursive=params['recursive'],
+            ):
+                info = wg.stage.info(x)
+                files.append(
+                    tuple([
+                        x, info.type, info.size or 0, info.writable,
+                        dt_isoformat(info.created_at),
+                        dt_isoformat(info.last_modified_at),
+                    ]),
+                )
+            res.set_rows(files)
+
+        else:
+            res.set_rows([(x,) for x in wg.stage.listdir(
+                params['at_path'] or '/',
+                recursive=params['recursive'],
+            )])
+
+        if params['like']:
+            res = res.like(Name=params['like'])
+
+        return res.order_by(**params['order_by']).limit(params['limit'])
+
+
+ShowStageFilesHandler.register(overwrite=True)
+
+
+class UploadStageFileHandler(SQLHandler):
+    """
+    UPLOAD FILE TO STAGE stage_path [ in_group ] FROM local_path [ overwrite ];
+
+    # Path to stage file
+    stage_path = '<stage-path>'
+
+    # Workspace group
+    in_group = IN GROUP { group_id | group_name }
+
+    # ID of group
+    group_id = ID '<group-id>'
+
+    # Name of group
+    group_name = '<group-name>'
+
+    # Path to local file
+    local_path = '<local-path>'
+
+    # Should an existing file be overwritten?
+    overwrite = OVERWRITE
+
+    """
+
+    def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]:
+        wg = get_workspace_group(params)
+        wg.stage.upload_file(
+            params['local_path'], params['stage_path'],
+            overwrite=params['overwrite'],
+        )
+        return None
+
+
+UploadStageFileHandler.register(overwrite=True)
+
+
+class DownloadStageFileHandler(SQLHandler):
+    """
+    DOWNLOAD STAGE FILE stage_path [ in_group ] [ local_path ]
+                                   [ overwrite ] [ encoding ];
+
+    # Path to stage file
+    stage_path = '<stage-path>'
+
+    # Workspace group
+    in_group = IN GROUP { group_id | group_name }
+
+    # ID of group
+    group_id = ID '<group-id>'
+
+    # Name of group
+    group_name = '<group-name>'
+
+    # Path to local file
+    local_path = TO '<local-path>'
+
+    # Should an existing file be overwritten?
+    overwrite = OVERWRITE
+
+    # File encoding
+    encoding = ENCODING '<encoding>'
+
+    """
+
+    def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]:
+        wg = get_workspace_group(params)
+
+        out = wg.stage.download_file(
+            params['stage_path'],
+            local_path=params['local_path'] or None,
+            overwrite=params['overwrite'],
+            encoding=params['encoding'] or None,
+        )
+
+        if not params['local_path']:
+            res = FusionSQLResult()
+            if params['encoding']:
+                res.add_field('Data', result.STRING)
+            else:
+                res.add_field('Data', result.BLOB)
+            res.set_rows([(out,)])
+            return res
+
+        return None
+
+
+DownloadStageFileHandler.register(overwrite=True)
+
+
+class DropStageFileHandler(SQLHandler):
+    """
+    DROP STAGE FILE stage_path [ in_group ];
+
+    # Path to stage file
+    stage_path = '<stage-path>'
+
+    # Workspace group
+    in_group = IN GROUP { group_id | group_name }
+
+    # ID of group
+    group_id = ID '<group-id>'
+
+    # Name of group
+    group_name = '<group-name>'
+
+    """
+
+    def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]:
+        wg = get_workspace_group(params)
+        wg.stage.remove(params['stage_path'])
+        return None
+
+
+DropStageFileHandler.register(overwrite=True)
+
+
+class DropStageFolderHandler(SQLHandler):
+    """
+    DROP STAGE FOLDER stage_path [ in_group ] [ recursive ];
+
+    # Path to stage folder
+    stage_path = '<stage-path>'
+
+    # Workspace group
+    in_group = IN GROUP { group_id | group_name }
+
+    # ID of group
+    group_id = ID '<group-id>'
+
+    # Name of group
+    group_name = '<group-name>'
+
+    # Should folers be deleted recursively?
+    recursive = RECURSIVE
+
+    """
+
+    def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]:
+        wg = get_workspace_group(params)
+        if params['recursive']:
+            wg.stage.removedirs(params['stage_path'])
+        else:
+            wg.stage.rmdir(params['stage_path'])
+        return None
+
+
+DropStageFolderHandler.register(overwrite=True)
+
+
+class CreateStageFolderHandler(SQLHandler):
+    """
+    CREATE STAGE FOLDER stage_path [ in_group ] [ overwrite ];
+
+    # Workspace group
+    in_group = IN GROUP { group_id | group_name }
+
+    # ID of group
+    group_id = ID '<group-id>'
+
+    # Name of group
+    group_name = '<group-name>'
+
+    # Path to stage folder
+    stage_path = '<stage-path>'
+
+    # Should an existing folder be overwritten?
+    overwrite = OVERWRITE
+
+    """
+
+    def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]:
+        wg = get_workspace_group(params)
+        wg.stage.mkdir(params['stage_path'], overwrite=params['overwrite'])
+        return None
+
+
+CreateStageFolderHandler.register(overwrite=True)

singlestoredb/fusion/handlers/utils.py

@@ -0,0 +1,162 @@
+#!/usr/bin/env python
+import datetime
+import os
+from typing import Any
+from typing import Dict
+from typing import Optional
+
+from ...exceptions import ManagementError
+from ...management import manage_workspaces
+from ...management.workspace import Workspace
+from ...management.workspace import WorkspaceGroup
+from ...management.workspace import WorkspaceManager
+
+
+def get_workspace_manager() -> WorkspaceManager:
+    """Return a new workspace manager."""
+    return manage_workspaces()
+
+
+def dt_isoformat(dt: Optional[datetime.datetime]) -> Optional[str]:
+    """Convert datetime to string."""
+    if dt is None:
+        return None
+    return dt.isoformat()
+
+
+def get_workspace_group(params: Dict[str, Any]) -> WorkspaceGroup:
+    """
+    Find a workspace group matching group_id or group_name.
+
+    This function will get a workspace group or ID from the
+    following parameters:
+
+        * params['group_name']
+        * params['group_id']
+        * params['group']['group_name']
+        * params['group']['group_id']
+        * params['in_group']['group_name']
+        * params['in_group']['group_id']
+
+    Or, from the SINGLESTOREDB_WORKSPACE_GROUP environment variable.
+
+    """
+    manager = get_workspace_manager()
+
+    group_name = params.get('group_name') or \
+        (params.get('in_group') or {}).get('group_name') or \
+        (params.get('group') or {}).get('group_name')
+    if group_name:
+        workspace_groups = [
+            x for x in manager.workspace_groups
+            if x.name == group_name
+        ]
+
+        if not workspace_groups:
+            raise KeyError(
+                f'no workspace group found with name: {group_name}',
+            )
+
+        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}',
+            )
+
+        return workspace_groups[0]
+
+    group_id = params.get('group_id') or \
+        (params.get('in_group') or {}).get('group_id') or \
+        (params.get('group') or {}).get('group_id')
+    if group_id:
+        try:
+            return manager.get_workspace_group(group_id)
+        except ManagementError as exc:
+            if exc.errno == 404:
+                raise KeyError(f'no workspace group found with ID: {group_id}')
+            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'):
+        raise ValueError('clusters and shared workspaces are not currently supported')
+
+    raise KeyError('no workspace group was specified')
+
+
+def get_workspace(params: Dict[str, Any]) -> Workspace:
+    """
+    Retrieve the specified workspace.
+
+    This function will get a workspace group or ID from the
+    following parameters:
+
+        * params['workspace_name']
+        * params['workspace_id']
+        * params['workspace']['workspace_name']
+        * params['workspace']['workspace_id']
+
+    Or, from the SINGLESTOREDB_WORKSPACE environment variable.
+
+    """
+    manager = get_workspace_manager()
+    workspace_name = params.get('workspace_name') or \
+        (params.get('workspace') or {}).get('workspace_name')
+    if workspace_name:
+        wg = get_workspace_group(params)
+        workspaces = [
+            x for x in wg.workspaces
+            if x.name == workspace_name
+        ]
+
+        if not workspaces:
+            raise KeyError(
+                f'no workspace found with name: {workspace_name}',
+            )
+
+        if len(workspaces) > 1:
+            ids = ', '.join(x.id for x in workspaces)
+            raise ValueError(
+                f'more than one workspace with given name was found: {ids}',
+            )
+
+        return workspaces[0]
+
+    workspace_id = params.get('workspace_id') or \
+        (params.get('workspace') or {}).get('workspace_id')
+    if workspace_id:
+        try:
+            return manager.get_workspace(workspace_id)
+        except ManagementError as exc:
+            if exc.errno == 404:
+                raise KeyError(f'no workspace found with ID: {workspace_id}')
+            raise
+
+    if os.environ.get('SINGLESTOREDB_WORKSPACE'):
+        try:
+            return manager.get_workspace(
+                os.environ['SINGLESTOREDB_WORKSPACE'],
+            )
+        except ManagementError as exc:
+            if exc.errno == 404:
+                raise KeyError(
+                    'no workspace found with ID: '
+                    f'{os.environ["SINGLESTOREDB_WORKSPACE"]}',
+                )
+            raise
+
+    if os.environ.get('SINGLESTOREDB_CLUSTER'):
+        raise ValueError('clusters and shared workspaces are not currently supported')
+
+    raise KeyError('no workspace was specified')