Make workspace group optional if it is specified in the environment

Author: kesmit13

singlestoredb/fusion/handlers/utils.py

@@ -1,5 +1,6 @@
 #!/usr/bin/env python
 import datetime
+import os
 from typing import Any
 from typing import Dict
 from typing import Optional
@@ -26,15 +27,16 @@ 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']:
+    group_name = params.get('group_name') or params.get('in_group', {}).get('group_name')
+    if group_name:
         workspace_groups = [
             x for x in manager.workspace_groups
-            if x.name == params['group_name']
+            if x.name == group_name
         ]
 
         if not workspace_groups:
             raise KeyError(
-                'no workspace group found with name "{}"'.format(params['group_name']),
+                'no workspace group found with name "{}"'.format(group_name),
             )
 
         if len(workspace_groups) > 1:
@@ -45,4 +47,16 @@ def get_workspace_group(params: Dict[str, Any]) -> WorkspaceGroup:
 
         return workspace_groups[0]
 
-    return manager.get_workspace_group(params['group_id'])
+    group_id = params.get('group_id') or params.get('in_group', {}).get('group_id')
+    if group_id:
+        return manager.get_workspace_group(group_id)
+
+    if os.environ.get('SINGLESTOREDB_WORKSPACE_GROUP'):
+        return manager.get_workspace_group(
+            os.environ['SINGLESTOREDB_WORKSPACE_GROUP'],
+        )
+
+    if os.environ.get('SINGLESTOREDB_CLUSTER'):
+        raise ValueError('clusters and shared workspaces are not currently supported')
+
+    raise KeyError('no workspace group was specified')

singlestoredb/fusion/handlers/workspace.py

@@ -80,8 +80,10 @@ def fields(x: Any) -> Any:
 
 class ShowWorkspacesHandler(SQLHandler):
     """
-    SHOW WORKSPACES IN GROUP { group_id | group_name }
-        [ <like> ] [ <extended> ] [ <order-by> ] [ <limit> ];
+    SHOW WORKSPACES [ in_group ] [ <like> ] [ <extended> ] [ <order-by> ] [ <limit> ];
+
+    # Workspace group
+    in_group = IN GROUP { group_id | group_name }
 
     # ID of group
     group_id = ID '<group-id>'
@@ -198,10 +200,12 @@ def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]:
 
 class CreateWorkspaceHandler(SQLHandler):
     """
-    CREATE WORKSPACE [ if_not_exists ] workspace_name
-        IN GROUP { group_id | group_name }
+    CREATE WORKSPACE [ if_not_exists ] workspace_name [ in_group ]
         WITH SIZE size [ wait_on_active ];
 
+    # Create workspace in workspace group
+    in_group = IN GROUP { group_id | group_name }
+
     # Only run command if workspace doesn't already exist
     if_not_exists = IF NOT EXISTS
 
@@ -264,15 +268,15 @@ class DropWorkspaceGroupHandler(SQLHandler):
     """
 
     def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]:
-        manager = get_workspace_manager()
-
         try:
-            name_or_id = params['group_name'] or params['group_id']
-            wg = manager.workspace_groups[name_or_id]
-            wg.terminate(wait_on_terminated=params['wait_on_terminated'])
+            workspace_group = get_workspace_group(params)
+            if workspace_group.terminated_at is not None:
+                raise KeyError
+            workspace_group.terminate(wait_on_terminated=params['wait_on_terminated'])
 
         except KeyError:
             if not params['if_exists']:
+                name_or_id = params['group_id'] or params['group_name']
                 raise KeyError(f"could not find workspace group '{name_or_id}'")
 
         return None
@@ -284,7 +288,10 @@ def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]:
 class DropWorkspaceHandler(SQLHandler):
     """
     DROP WORKSPACE [ if_exists ] { workspace_id | workspace_name }
-        IN GROUP { group_id | group_name } [ wait_on_terminated ];
+        [ in_group ] [ wait_on_terminated ];
+
+    # Workspace group
+    in_group = IN GROUP { group_id | group_name }
 
     # Only drop workspace if it exists
     if_exists = IF EXISTS
@@ -307,16 +314,15 @@ class DropWorkspaceHandler(SQLHandler):
     """
 
     def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]:
-        manager = get_workspace_manager()
-
         try:
+            workspace_group = get_workspace_group(params)
             workspace_name_or_id = params['workspace_name'] or params['workspace_id']
-            group_name_or_id = params['group_name'] or params['group_id']
-            wg = manager.workspace_groups[group_name_or_id]
-            ws = wg.workspaces[workspace_name_or_id]
+            ws = workspace_group.workspaces[workspace_name_or_id]
             ws.terminate(wait_on_terminated=params['wait_on_terminated'])
 
         except KeyError:
+            group_name_or_id = params['in_group'].get('group_id', None) or \
+                               params['in_group'].get('group_name', None)
             if not params['if_exists']:
                 raise KeyError(
                     f"could not find workspace '{workspace_name_or_id}' "