singlestoredb 1.4.3__cp38-abi3-win32.whl → 1.6.0__cp38-abi3-win32.whl

singlestoredb/fusion/handlers/workspace.py

@@ -13,6 +13,73 @@ from .utils import get_workspace_group
 from .utils import get_workspace_manager
 
 
+class UseWorkspaceHandler(SQLHandler):
+    """
+    USE WORKSPACE workspace [ with_database ];
+
+    # Workspace
+    workspace = { workspace_id | workspace_name | current_workspace }
+
+    # ID of workspace
+    workspace_id = ID '<workspace-id>'
+
+    # Name of workspace
+    workspace_name = '<workspace-name>'
+
+    # Current workspace
+    current_workspace = @@CURRENT
+
+    # Name of database
+    with_database = WITH DATABASE 'database-name'
+
+    Description
+    -----------
+    Change the workspace and database in the notebook.
+
+    Arguments
+    ---------
+    * ``<workspace-id>``: The ID of the workspace to delete.
+    * ``<workspace-name>``: The name of the workspace to delete.
+
+    Remarks
+    -------
+    * If you want to specify a database in the current workspace,
+      the workspace name can be specified as ``@@CURRENT``.
+    * Specify the ``WITH DATABASE`` clause to select a default
+      database for the session.
+    * This command only works in a notebook session in the
+      Managed Service.
+
+    Example
+    -------
+    The following command sets the workspace to ``examplews`` and
+    select 'dbname' as the default database::
+
+        USE WORKSPACE 'examplews' WITH DATABASE 'dbname';
+
+    """
+    def run(self, params: Dict[str, Any]) -> Optional[FusionSQLResult]:
+        from singlestoredb.notebook import portal
+        if params['workspace'].get('current_workspace'):
+            if params.get('with_database'):
+                portal.default_database = params['with_database']
+        elif params.get('with_database'):
+            if params['workspace'].get('workspace_name'):
+                portal.connection = params['workspace']['workspace_name'], \
+                                    params['with_database']
+            else:
+                portal.connection = params['workspace']['workspace_id'], \
+                                    params['with_database']
+        elif params['workspace'].get('workspace_name'):
+            portal.workspace = params['workspace']['workspace_name']
+        else:
+            portal.workspace = params['workspace']['workspace_id']
+        return None
+
+
+UseWorkspaceHandler.register(overwrite=True)
+
+
 class ShowRegionsHandler(SQLHandler):
     """
     SHOW REGIONS [ <like> ]
@@ -416,8 +483,9 @@ class CreateWorkspaceHandler(SQLHandler):
     size = '<size>'
 
     # Auto-suspend
-    auto_suspend = AUTO SUSPEND suspend_after_seconds SECONDS suspend_type
-    suspend_after_seconds = AFTER <integer>
+    auto_suspend = AUTO SUSPEND AFTER suspend_after_value suspend_after_units suspend_type
+    suspend_after_value = <integer>
+    suspend_after_units = { SECONDS | MINUTES | HOURS | DAYS }
     suspend_type = WITH TYPE { IDLE | SCHEDULED | DISABLED }
 
     # Enable Kai
@@ -442,7 +510,7 @@ class CreateWorkspaceHandler(SQLHandler):
       in which the workspace is created.
     * ``<workspace_size>``: The size of the workspace in workspace size notation,
       for example "S-1".
-    * ``<suspend_time>``: The time (in seconds) after which the workspace is
+    * ``<suspend_time>``: The time (in given units) after which the workspace is
       suspended, according to the specified auto-suspend type.
     * ``<multiplier>``: The multiplier for the persistent cache associated with
       the workspace.
@@ -486,9 +554,17 @@ class CreateWorkspaceHandler(SQLHandler):
 
         auto_suspend = None
         if params['auto_suspend']:
+            mult = dict(
+                SECONDS=1,
+                MINUTES=60,
+                HOURS=60*60,
+                DAYS=60*60*24,
+            )
+            val = params['auto_suspend'][0]['suspend_after_value']
+            val = val * mult[params['auto_suspend'][1]['suspend_after_units'].upper()]
             auto_suspend = dict(
-                suspend_after_seconds=params['auto_suspend'][0]['suspend_after_seconds'],
-                suspend_type=params['auto_suspend'][-1]['suspend_type'].upper(),
+                suspend_after_seconds=val,
+                suspend_type=params['auto_suspend'][2]['suspend_type'].upper(),
             )
 
         workspace_group.create_workspace(

singlestoredb/http/connection.py

@@ -1243,5 +1243,7 @@ def connect(
     inf_as_null: Optional[bool] = None,
     encoding_errors: Optional[str] = None,
     track_env: Optional[bool] = None,
+    enable_extended_data_types: Optional[bool] = None,
+    vector_data_format: Optional[str] = None,
 ) -> Connection:
     return Connection(**dict(locals()))