More Fusion fixes and testing

Author: kesmit13

docs/src/post-process-html

@@ -55,6 +55,9 @@ for file in sorted(txt_files):
         # Change ShowAccessor to Connection.show
         txt = re.sub(r'>ShowAccessor\.', r'>Connection.show.', txt)
 
+        # Change workspace.Stages to workspace.stages
+        txt = re.sub(r'>workspace\.Stages\.', r'>workspace.stages.', txt)
+
         # Fix class / method links
         txt = re.sub(
             r'(<a\s+[^>]+>)?(\s*<code[^>]*>\s*<span\s+class="pre">\s*)([\w\.]+)(\s*</span>\s*</code>)',

singlestoredb/fusion/handler.py

@@ -20,14 +20,15 @@
 from ..connection import Connection
 
 CORE_GRAMMAR = r'''
-    ws = ~r"(\s*(/\*.*\*/)*\s*)*"
+    ws = ~r"(\s+|(\s*/\*.*\*/\s*)+)"
     qs = ~r"\"([^\"]*)\"|'([^\']*)'|`([^\`]*)`|([A-Za-z0-9_\-\.]+)"
     number = ~r"[-+]?(\d*\.)?\d+(e[-+]?\d+)?"i
     integer = ~r"-?\d+"
-    comma = ws "," ws
-    open_paren = ws "(" ws
-    close_paren = ws ")" ws
-    select = ~r"SELECT"i ws ~r".+" ws
+    comma = ws* "," ws*
+    eq = ws* "=" ws*
+    open_paren = ws* "(" ws*
+    close_paren = ws* ")" ws*
+    select = ~r"SELECT"i ws+ ~r".+" ws*
 '''
 
 BUILTINS = {
@@ -56,7 +57,7 @@
 
 def get_keywords(grammar: str) -> Tuple[str, ...]:
     """Return all all-caps words from the beginning of the line."""
-    m = re.match(r'^\s*((?:[A-Z0-9_]+|=)(\s+|$|;))+', grammar)
+    m = re.match(r'^\s*((?:[A-Z0-9_]+)(\s+|$|;))+', grammar)
     if not m:
         return tuple()
     return tuple(re.split(r'\s+', m.group(0).replace(';', '').strip()))
@@ -86,7 +87,7 @@ def process_alternates(m: Any) -> str:
 def process_repeats(m: Any) -> str:
     """Add repeated patterns."""
     sql = m.group(1).strip()
-    return f'open_paren? {sql} ws ( comma {sql} ws )* close_paren?'
+    return f'open_paren? {sql} ws* ( comma {sql} ws* )* close_paren?'
 
 
 def lower_and_regex(m: Any) -> str:
@@ -277,7 +278,7 @@ def process_grammar(grammar: str) -> Tuple[Grammar, Tuple[str, ...], Dict[str, A
         sql = re.sub(r"'[^']+'", r'qs', sql)
 
         # Convert special characters to literal tokens
-        sql = re.sub(r'([=]) ', r" '\1' ", sql)
+        sql = re.sub(r'([=]) ', r' eq ', sql)
 
         # Convert [...] groups to (...)*
         sql = re.sub(r'\[([^\]]+)\]', process_optional, sql)
@@ -310,6 +311,15 @@ def process_grammar(grammar: str) -> Tuple[Grammar, Tuple[str, ...], Dict[str, A
 
         # Make sure every operation ends with ws
         sql = re.sub(r'\s+ws\s+ws$', r' ws', sql + ' ws')
+        sql = re.sub(r'(\s+ws)*\s+ws\*$', r' ws*', sql)
+        sql = re.sub(r'\s+ws$', r' ws*', sql)
+        sql = re.sub(r'\s+ws\s+\(', r' ws* (', sql)
+        sql = re.sub(r'\)\s+ws\s+', r') ws* ', sql)
+        sql = re.sub(r'\s+ws\s+', r' ws+ ', sql)
+        sql = re.sub(r'\?\s+ws\+', r'? ws*', sql)
+
+        # Remove extra ws around eq
+        sql = re.sub(r'ws\+\s*eq\b', r'eq', sql)
 
         out.append(f'{op} = {sql}')
 
@@ -324,7 +334,7 @@ def process_grammar(grammar: str) -> Tuple[Grammar, Tuple[str, ...], Dict[str, A
         rules[k] = v
 
     cmds = ' / '.join(x for x in rules if x.endswith('_cmd'))
-    cmds = f'init = ws ( {cmds} ) ws ";"? ws\n'
+    cmds = f'init = ws* ( {cmds} ) ws* ";"? ws*\n'
 
     return Grammar(cmds + CORE_GRAMMAR + '\n'.join(out)), command_key, rule_info, help_txt
 
@@ -500,6 +510,10 @@ def visit_ws(self, node: Node, visited_children: Iterable[Any]) -> Any:
         """Whitespace and comments."""
         return
 
+    def visit_eq(self, node: Node, visited_children: Iterable[Any]) -> Any:
+        """Equals sign."""
+        return
+
     def visit_comma(self, node: Node, visited_children: Iterable[Any]) -> Any:
         """Single comma."""
         return

singlestoredb/fusion/handlers/stages.py

@@ -0,0 +1,192 @@
+#!/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 { group_id | group_name } [ at_path ]
+        [ <like> ] [ <order-by> ] [ <limit> ] [ recursive ] [ extended ];
+
+    group_id = ID '<group-id>'
+    group_name = '<group-name>'
+    at_path = AT '<path>'
+    recursive = RECURSIVE
+    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.stages.listdir(
+                params['at_path'] or '/',
+                recursive=params['recursive'],
+            ):
+                info = wg.stages.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.stages.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 local_path TO STAGE stage_path
+        IN GROUP { group_id | group_name } [ overwrite ];
+
+    local_path = '<local-path>'
+    stage_path = '<stage-path>'
+    group_id = ID '<group-id>'
+    group_name = '<group-name>'
+    overwrite = OVERWRITE
+
+    """
+
+    def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]:
+        wg = get_workspace_group(params)
+        wg.stages.upload_file(
+            params['local_path'], params['stage_path'],
+            overwrite=params['overwrite'],
+        )
+        return None
+
+
+UploadStageFileHandler.register()
+
+
+class DownloadStageFileHandler(SQLHandler):
+    """
+    DOWNLOAD STAGE FILE stage_path IN GROUP { group_id | group_name }
+        [ local_path ] [ overwrite ] [ encoding ];
+
+    stage_path = '<stage-path>'
+    group_id = ID '<group-id>'
+    group_name = '<group-name>'
+    local_path = TO '<local-path>'
+    overwrite = OVERWRITE
+    encoding = ENCODING '<encoding>'
+
+    """
+
+    def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]:
+        wg = get_workspace_group(params)
+
+        out = wg.stages.download(
+            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()
+
+
+class DropStageFileHandler(SQLHandler):
+    """
+    DROP STAGE FILE stage_path IN GROUP { group_id | group_name };
+
+    stage_path = '<stage-path>'
+    group_id = ID '<group-id>'
+    group_name = '<group-name>'
+
+    """
+
+    def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]:
+        wg = get_workspace_group(params)
+        wg.stages.remove(params['stage_path'])
+        return None
+
+
+DropStageFileHandler.register()
+
+
+class DropStageFolderHandler(SQLHandler):
+    """
+    DROP STAGE FOLDER stage_path IN GROUP { group_id | group_name } [ recursive ];
+
+    stage_path = '<stage-path>'
+    group_id = ID '<group-id>'
+    group_name = '<group-name>'
+    recursive = RECURSIVE
+
+    """
+
+    def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]:
+        wg = get_workspace_group(params)
+        if params['recursive']:
+            wg.stages.rmdir(params['stage_path'])
+        else:
+            wg.stages.removedirs(params['stage_path'])
+        return None
+
+
+DropStageFolderHandler.register()
+
+
+class CreateStageFolderHandler(SQLHandler):
+    """
+    CREATE STAGE FOLDER stage_path IN GROUP { group_id | group_name } [ overwrite ];
+
+    stage_path = '<stage-path>'
+    group_id = ID '<group-id>'
+    group_name = '<group-name>'
+    overwrite = OVERWRITE
+
+    """
+
+    def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]:
+        wg = get_workspace_group(params)
+        wg.stages.mkdir(params['stage_path'], overwrite=params['overwrite'])
+        return None
+
+
+CreateStageFolderHandler.register()

singlestoredb/fusion/handlers/utils.py

@@ -0,0 +1,76 @@
+#!/usr/bin/env python
+import datetime
+import os
+from typing import Any
+from typing import Dict
+from typing import Optional
+from urllib.parse import urlparse
+
+import jwt
+
+from ...config import get_option
+from ...management import manage_workspaces
+from ...management.workspace import WorkspaceGroup
+from ...management.workspace import WorkspaceManager
+
+
+def get_management_token() -> Optional[str]:
+    """Return the token for the Management API."""
+    # See if an API key is configured
+    tok = get_option('management.token')
+    if tok:
+        return tok
+
+    # See if the connection URL contains a JWT
+    url = os.environ.get('SINGLESTOREDB_URL')
+    if not url:
+        return None
+
+    urlp = urlparse(url, scheme='singlestoredb', allow_fragments=True)
+    if urlp.password:
+        try:
+            jwt.decode(urlp.password, options={'verify_signature': False})
+            return urlp.password
+        except jwt.DecodeError:
+            pass
+
+    # Didn't find a key anywhere
+    return None
+
+
+def get_workspace_manager() -> WorkspaceManager:
+    """Return a new workspace manager."""
+    return manage_workspaces(get_management_token())
+
+
+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."""
+    manager = get_workspace_manager()
+
+    if params['group_name']:
+        workspace_groups = [
+            x for x in manager.workspace_groups
+            if x.name == params['group_name']
+        ]
+
+        if not workspace_groups:
+            raise KeyError(
+                'no workspace group found with name "{}"'.format(params['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]
+
+    return manager.get_workspace_group(params['group_id'])