singlestoredb 1.0.3__cp38-abi3-win32.whl → 1.1.0__cp38-abi3-win32.whl

singlestoredb/management/workspace.py

@@ -7,8 +7,8 @@ import glob
 import io
 import os
 import re
-import socket
 import time
+from collections.abc import Mapping
 from typing import Any
 from typing import BinaryIO
 from typing import Dict
@@ -23,10 +23,12 @@ from .billing_usage import BillingUsageItem
 from .manager import Manager
 from .organization import Organization
 from .region import Region
+from .utils import camel_to_snake_dict
 from .utils import from_datetime
 from .utils import NamedList
 from .utils import PathLike
 from .utils import snake_to_camel
+from .utils import snake_to_camel_dict
 from .utils import to_datetime
 from .utils import ttl_property
 from .utils import vars_to_str
@@ -934,6 +936,12 @@ class Workspace(object):
         created_at: Union[str, datetime.datetime],
         terminated_at: Optional[Union[str, datetime.datetime]] = None,
         endpoint: Optional[str] = None,
+        auto_suspend: Optional[Dict[str, Any]] = None,
+        cache_config: Optional[int] = None,
+        deployment_type: Optional[str] = None,
+        resume_attachments: Optional[Dict[str, Any]] = None,
+        scaling_progress: Optional[int] = None,
+        last_resumed_at: Optional[str] = None,
     ):
         #: Name of the workspace
         self.name = name
@@ -963,6 +971,24 @@ class Workspace(object):
         #: Hostname (or IP address) of the workspace database server
         self.endpoint = endpoint
 
+        #: Current auto-suspend settings
+        self.auto_suspend = camel_to_snake_dict(auto_suspend)
+
+        #: Multiplier for the persistent cache
+        self.cache_config = cache_config
+
+        #: Deployment type of the workspace
+        self.deployment_type = deployment_type
+
+        #: Database attachments
+        self.resume_attachments = camel_to_snake_dict(resume_attachments)
+
+        #: Current progress percentage for scaling the workspace
+        self.scaling_progress = scaling_progress
+
+        #: Timestamp when workspace was last resumed
+        self.last_resumed_at = to_datetime(last_resumed_at)
+
         self._manager: Optional[WorkspaceManager] = None
 
     def __str__(self) -> str:
@@ -991,15 +1017,64 @@ class Workspace(object):
 
         """
         out = cls(
-            name=obj['name'], workspace_id=obj['workspaceID'],
+            name=obj['name'],
+            workspace_id=obj['workspaceID'],
             workspace_group=obj['workspaceGroupID'],
-            size=obj.get('size', 'Unknown'), state=obj['state'],
-            created_at=obj['createdAt'], terminated_at=obj.get('terminatedAt'),
+            size=obj.get('size', 'Unknown'),
+            state=obj['state'],
+            created_at=obj['createdAt'],
+            terminated_at=obj.get('terminatedAt'),
             endpoint=obj.get('endpoint'),
+            auto_suspend=obj.get('autoSuspend'),
+            cache_config=obj.get('cacheConfig'),
+            deployment_type=obj.get('deploymentType'),
+            last_resumed_at=obj.get('lastResumedAt'),
+            resume_attachments=obj.get('resumeAttachments'),
+            scaling_progress=obj.get('scalingProgress'),
         )
         out._manager = manager
         return out
 
+    def update(
+        self,
+        auto_suspend: Optional[Dict[str, Any]] = None,
+        cache_config: Optional[int] = None,
+        deployment_type: Optional[str] = None,
+        size: Optional[str] = None,
+    ) -> None:
+        """
+        Update the workspace definition.
+
+        Parameters
+        ----------
+        auto_suspend : Dict[str, Any], optional
+            Auto-suspend mode for the workspace: IDLE, SCHEDULED, DISABLED
+        cache_config : int, optional
+            Specifies the multiplier for the persistent cache associated
+            with the workspace. If specified, it enables the cache configuration
+            multiplier. It can have one of the following values: 1, 2, or 4.
+        deployment_type : str, optional
+            The deployment type that will be applied to all the workspaces
+            within the group
+        size : str, optional
+            Size of the workspace (in workspace size notation), such as "S-1".
+
+        """
+        if self._manager is None:
+            raise ManagementError(
+                msg='No workspace manager is associated with this object.',
+            )
+        data = {
+            k: v for k, v in dict(
+                autoSuspend=snake_to_camel_dict(auto_suspend),
+                cacheConfig=cache_config,
+                deploymentType=deployment_type,
+                size=size,
+            ).items() if v is not None
+        }
+        self._manager._patch(f'workspaces/{self.id}', json=data)
+        self.refresh()
+
     def refresh(self) -> Workspace:
         """Update the object to the current state."""
         if self._manager is None:
@@ -1008,7 +1083,10 @@ class Workspace(object):
             )
         new_obj = self._manager.get_workspace(self.id)
         for name, value in vars(new_obj).items():
-            setattr(self, name, value)
+            if isinstance(value, Mapping):
+                setattr(self, name, snake_to_camel_dict(value))
+            else:
+                setattr(self, name, value)
         return self
 
     def terminate(
@@ -1111,6 +1189,7 @@ class Workspace(object):
 
     def resume(
         self,
+        disable_auto_suspend: bool = False,
         wait_on_resumed: bool = False,
         wait_interval: int = 20,
         wait_timeout: int = 600,
@@ -1120,6 +1199,8 @@ class Workspace(object):
 
         Parameters
         ----------
+        disable_auto_suspend : bool, optional
+            Should auto-suspend be disabled?
         wait_on_resumed : bool, optional
             Wait for the workspace to go into 'Resumed' or 'Active' mode before returning
         wait_interval : int, optional
@@ -1137,7 +1218,10 @@ class Workspace(object):
             raise ManagementError(
                 msg='No workspace manager is associated with this object.',
             )
-        self._manager._post(f'workspaces/{self.id}/resume')
+        self._manager._post(
+            f'workspaces/{self.id}/resume',
+            json=dict(disableAutoSuspend=disable_auto_suspend),
+        )
         if wait_on_resumed:
             self._manager._wait_on_state(
                 self._manager.get_workspace(self.id),
@@ -1170,6 +1254,7 @@ class WorkspaceGroup(object):
         region: Optional[Region],
         firewall_ranges: List[str],
         terminated_at: Optional[Union[str, datetime.datetime]],
+        allow_all_traffic: Optional[bool],
     ):
         #: Name of the workspace group
         self.name = name
@@ -1189,6 +1274,9 @@ class WorkspaceGroup(object):
         #: Timestamp of when the workspace group was terminated
         self.terminated_at = to_datetime(terminated_at)
 
+        #: Should all traffic be allowed?
+        self.allow_all_traffic = allow_all_traffic
+
         self._manager: Optional[WorkspaceManager] = None
 
     def __str__(self) -> str:
@@ -1229,6 +1317,7 @@ class WorkspaceGroup(object):
             region=region,
             firewall_ranges=obj.get('firewallRanges', []),
             terminated_at=obj.get('terminatedAt'),
+            allow_all_traffic=obj.get('allowAllTraffic'),
         )
         out._manager = manager
         return out
@@ -1260,13 +1349,20 @@ class WorkspaceGroup(object):
             )
         new_obj = self._manager.get_workspace_group(self.id)
         for name, value in vars(new_obj).items():
-            setattr(self, name, value)
+            if isinstance(value, Mapping):
+                setattr(self, name, camel_to_snake_dict(value))
+            else:
+                setattr(self, name, value)
         return self
 
     def update(
-        self, name: Optional[str] = None,
-        admin_password: Optional[str] = None,
+        self,
+        name: Optional[str] = None,
         firewall_ranges: Optional[List[str]] = None,
+        admin_password: Optional[str] = None,
+        expires_at: Optional[str] = None,
+        allow_all_traffic: Optional[bool] = None,
+        update_window: Optional[Dict[str, int]] = None,
     ) -> None:
         """
         Update the workspace group definition.
@@ -1274,11 +1370,24 @@ class WorkspaceGroup(object):
         Parameters
         ----------
         name : str, optional
-            Workspace group name
-        admim_password : str, optional
-            Admin password for the workspace group
-        firewall_ranges : Sequence[str], optional
-            List of allowed incoming IP addresses
+            Name of the workspace group
+        firewall_ranges : list[str], optional
+            List of allowed CIDR ranges. An empty list indicates that all
+            inbound requests are allowed.
+        admin_password : str, optional
+            Admin password for the workspace group. If no password is supplied,
+            a password will be generated and retured in the response.
+        expires_at : str, optional
+            The timestamp of when the workspace group will expire.
+            If the expiration time is not specified,
+            the workspace group will have no expiration time.
+            At expiration, the workspace group is terminated and all the data is lost.
+            Expiration time can be specified as a timestamp or duration.
+            Example: "2021-01-02T15:04:05Z07:00", "2021-01-02", "3h30m"
+        allow_all_traffic : bool, optional
+            Allow all traffic to the workspace group
+        update_window : Dict[str, int], optional
+            Specify the day and hour of an update window: dict(day=0-6, hour=0-23)
 
         """
         if self._manager is None:
@@ -1287,8 +1396,12 @@ class WorkspaceGroup(object):
             )
         data = {
             k: v for k, v in dict(
-                name=name, adminPassword=admin_password,
+                name=name,
                 firewallRanges=firewall_ranges,
+                adminPassword=admin_password,
+                expiresAt=expires_at,
+                allowAllTraffic=allow_all_traffic,
+                updateWindow=snake_to_camel_dict(update_window),
             ).items() if v is not None
         }
         self._manager._patch(f'workspaceGroups/{self.id}', json=data)
@@ -1338,9 +1451,15 @@ class WorkspaceGroup(object):
                 wait_timeout -= wait_interval
 
     def create_workspace(
-        self, name: str, size: Optional[str] = None,
-        wait_on_active: bool = False, wait_interval: int = 10,
-        wait_timeout: int = 600, add_endpoint_to_firewall_ranges: bool = True,
+        self,
+        name: str,
+        size: Optional[str] = None,
+        auto_suspend: Optional[Dict[str, Any]] = None,
+        cache_config: Optional[int] = None,
+        enable_kai: Optional[bool] = None,
+        wait_on_active: bool = False,
+        wait_interval: int = 10,
+        wait_timeout: int = 600,
     ) -> Workspace:
         """
         Create a new workspace.
@@ -1351,6 +1470,15 @@ class WorkspaceGroup(object):
             Name of the workspace
         size : str, optional
             Workspace size in workspace size notation (S-00, S-1, etc.)
+        auto_suspend : Dict[str, Any], optional
+            Auto suspend settings for the workspace. If this field is not
+            provided, no settings will be enabled.
+        cache_config : int, optional
+            Specifies the multiplier for the persistent cache associated
+            with the workspace. If specified, it enables the cache configuration
+            multiplier. It can have one of the following values: 1, 2, or 4.
+        enable_kai : bool, optional
+            Whether to create a SingleStore Kai-enabled workspace
         wait_on_active : bool, optional
             Wait for the workspace to be active before returning
         wait_timeout : int, optional
@@ -1358,9 +1486,6 @@ class WorkspaceGroup(object):
             if wait=True
         wait_interval : int, optional
             Number of seconds between each polling interval
-        add_endpoint_to_firewall_ranges : bool, optional
-            Should the workspace endpoint be added to the workspace group
-            firewall ranges?
 
         Returns
         -------
@@ -1373,14 +1498,17 @@ class WorkspaceGroup(object):
             )
 
         out = self._manager.create_workspace(
-            name=name, workspace_group=self, size=size, wait_on_active=wait_on_active,
-            wait_interval=wait_interval, wait_timeout=wait_timeout,
+            name=name,
+            workspace_group=self,
+            size=size,
+            auto_suspend=snake_to_camel_dict(auto_suspend),
+            cache_config=cache_config,
+            enable_kai=enable_kai,
+            wait_on_active=wait_on_active,
+            wait_interval=wait_interval,
+            wait_timeout=wait_timeout,
         )
 
-        if add_endpoint_to_firewall_ranges and out.endpoint is not None:
-            ip_address = '{}/32'.format(socket.gethostbyname(out.endpoint))
-            self.update(firewall_ranges=self.firewall_ranges+[ip_address])
-
         return out
 
     @property
@@ -1525,9 +1653,15 @@ class WorkspaceManager(Manager):
         return NamedList([Region.from_dict(item, self) for item in res.json()])
 
     def create_workspace_group(
-        self, name: str, region: Union[str, Region],
-        firewall_ranges: List[str], admin_password: Optional[str] = None,
+        self,
+        name: str,
+        region: Union[str, Region],
+        firewall_ranges: List[str],
+        admin_password: Optional[str] = None,
+        backup_bucket_kms_key_id: Optional[str] = None,
+        data_bucket_kms_key_id: Optional[str] = None,
         expires_at: Optional[str] = None,
+        smart_dr: Optional[bool] = None,
         allow_all_traffic: Optional[bool] = None,
         update_window: Optional[Dict[str, int]] = None,
     ) -> WorkspaceGroup:
@@ -1546,6 +1680,17 @@ class WorkspaceManager(Manager):
         admin_password : str, optional
             Admin password for the workspace group. If no password is supplied,
             a password will be generated and retured in the response.
+        backup_bucket_kms_key_id : str, optional
+            Specifies the KMS key ID associated with the backup bucket.
+            If specified, enables Customer-Managed Encryption Keys (CMEK)
+            encryption for the backup bucket of the workspace group.
+            This feature is only supported in workspace groups deployed in AWS.
+        data_bucket_kms_key_id : str, optional
+            Specifies the KMS key ID associated with the data bucket.
+            If specified, enables Customer-Managed Encryption Keys (CMEK)
+            encryption for the data bucket and Amazon Elastic Block Store
+            (EBS) volumes of the workspace group. This feature is only supported
+            in workspace groups deployed in AWS.
         expires_at : str, optional
             The timestamp of when the workspace group will expire.
             If the expiration time is not specified,
@@ -1553,6 +1698,10 @@ class WorkspaceManager(Manager):
             At expiration, the workspace group is terminated and all the data is lost.
             Expiration time can be specified as a timestamp or duration.
             Example: "2021-01-02T15:04:05Z07:00", "2021-01-02", "3h30m"
+        smart_dr : bool, optional
+            Enables Smart Disaster Recovery (SmartDR) for the workspace group.
+            SmartDR is a disaster recovery solution that ensures seamless and
+            continuous replication of data from the primary region to a secondary region
         allow_all_traffic : bool, optional
             Allow all traffic to the workspace group
         update_window : Dict[str, int], optional
@@ -1569,18 +1718,28 @@ class WorkspaceManager(Manager):
             'workspaceGroups', json=dict(
                 name=name, regionID=region,
                 adminPassword=admin_password,
+                backupBucketKMSKeyID=backup_bucket_kms_key_id,
+                dataBucketKMSKeyID=data_bucket_kms_key_id,
                 firewallRanges=firewall_ranges or [],
                 expiresAt=expires_at,
+                smartDR=smart_dr,
                 allowAllTraffic=allow_all_traffic,
-                updateWindow=update_window,
+                updateWindow=snake_to_camel_dict(update_window),
             ),
         )
         return self.get_workspace_group(res.json()['workspaceGroupID'])
 
     def create_workspace(
-        self, name: str, workspace_group: Union[str, WorkspaceGroup],
-        size: Optional[str] = None, wait_on_active: bool = False,
-        wait_interval: int = 10, wait_timeout: int = 600,
+        self,
+        name: str,
+        workspace_group: Union[str, WorkspaceGroup],
+        size: Optional[str] = None,
+        auto_suspend: Optional[Dict[str, Any]] = None,
+        cache_config: Optional[int] = None,
+        enable_kai: Optional[bool] = None,
+        wait_on_active: bool = False,
+        wait_interval: int = 10,
+        wait_timeout: int = 600,
     ) -> Workspace:
         """
         Create a new workspace.
@@ -1593,6 +1752,15 @@ class WorkspaceManager(Manager):
             The workspace ID of the workspace
         size : str, optional
             Workspace size in workspace size notation (S-00, S-1, etc.)
+        auto_suspend : Dict[str, Any], optional
+            Auto suspend settings for the workspace. If this field is not
+            provided, no settings will be enabled.
+        cache_config : int, optional
+            Specifies the multiplier for the persistent cache associated
+            with the workspace. If specified, it enables the cache configuration
+            multiplier. It can have one of the following values: 1, 2, or 4.
+        enable_kai : bool, optional
+            Whether to create a SingleStore Kai-enabled workspace
         wait_on_active : bool, optional
             Wait for the workspace to be active before returning
         wait_timeout : int, optional
@@ -1610,14 +1778,20 @@ class WorkspaceManager(Manager):
             workspace_group = workspace_group.id
         res = self._post(
             'workspaces', json=dict(
-                name=name, workspaceGroupID=workspace_group,
+                name=name,
+                workspaceGroupID=workspace_group,
                 size=size,
+                autoSuspend=snake_to_camel_dict(auto_suspend),
+                cacheConfig=cache_config,
+                enableKai=enable_kai,
             ),
         )
         out = self.get_workspace(res.json()['workspaceID'])
         if wait_on_active:
             out = self._wait_on_state(
-                out, 'Active', interval=wait_interval,
+                out,
+                'Active',
+                interval=wait_interval,
                 timeout=wait_timeout,
             )
         return out

singlestoredb/mysql/connection.py

@@ -29,12 +29,28 @@ from .cursors import (
     DictCursorSV,
     NamedtupleCursor,
     NamedtupleCursorSV,
+    ArrowCursor,
+    ArrowCursorSV,
+    NumpyCursor,
+    NumpyCursorSV,
+    PandasCursor,
+    PandasCursorSV,
+    PolarsCursor,
+    PolarsCursorSV,
     SSCursor,
     SSCursorSV,
     SSDictCursor,
     SSDictCursorSV,
     SSNamedtupleCursor,
     SSNamedtupleCursorSV,
+    SSArrowCursor,
+    SSArrowCursorSV,
+    SSNumpyCursor,
+    SSNumpyCursorSV,
+    SSPandasCursor,
+    SSPandasCursorSV,
+    SSPolarsCursor,
+    SSPolarsCursorSV,
 )
 from .optionfile import Parser
 from .protocol import (
@@ -443,6 +459,14 @@ class Connection(BaseConnection):
                 self.cursorclass = DictCursor
             elif 'namedtuple' in self.results_type:
                 self.cursorclass = NamedtupleCursor
+            elif 'numpy' in self.results_type:
+                self.cursorclass = NumpyCursor
+            elif 'arrow' in self.results_type:
+                self.cursorclass = ArrowCursor
+            elif 'pandas' in self.results_type:
+                self.cursorclass = PandasCursor
+            elif 'polars' in self.results_type:
+                self.cursorclass = PolarsCursor
             else:
                 self.cursorclass = Cursor
         else:
@@ -450,6 +474,14 @@ class Connection(BaseConnection):
                 self.cursorclass = SSDictCursor
             elif 'namedtuple' in self.results_type:
                 self.cursorclass = SSNamedtupleCursor
+            elif 'numpy' in self.results_type:
+                self.cursorclass = SSNumpyCursor
+            elif 'arrow' in self.results_type:
+                self.cursorclass = SSArrowCursor
+            elif 'pandas' in self.results_type:
+                self.cursorclass = SSPandasCursor
+            elif 'polars' in self.results_type:
+                self.cursorclass = SSPolarsCursor
             else:
                 self.cursorclass = SSCursor
 
@@ -487,6 +519,30 @@ class Connection(BaseConnection):
             elif self.cursorclass is SSNamedtupleCursor:
                 self.cursorclass = SSNamedtupleCursorSV
                 self.results_type = 'namedtuples'
+            elif self.cursorclass is NumpyCursor:
+                self.cursorclass = NumpyCursorSV
+                self.results_type = 'numpy'
+            elif self.cursorclass is SSNumpyCursor:
+                self.cursorclass = SSNumpyCursorSV
+                self.results_type = 'numpy'
+            elif self.cursorclass is ArrowCursor:
+                self.cursorclass = ArrowCursorSV
+                self.results_type = 'arrow'
+            elif self.cursorclass is SSArrowCursor:
+                self.cursorclass = SSArrowCursorSV
+                self.results_type = 'arrow'
+            elif self.cursorclass is PandasCursor:
+                self.cursorclass = PandasCursorSV
+                self.results_type = 'pandas'
+            elif self.cursorclass is SSPandasCursor:
+                self.cursorclass = SSPandasCursorSV
+                self.results_type = 'pandas'
+            elif self.cursorclass is PolarsCursor:
+                self.cursorclass = PolarsCursorSV
+                self.results_type = 'polars'
+            elif self.cursorclass is SSPolarsCursor:
+                self.cursorclass = SSPolarsCursorSV
+                self.results_type = 'polars'
 
         self._result = None
         self._affected_rows = 0
@@ -498,9 +554,22 @@ class Connection(BaseConnection):
         if conv is None:
             conv = converters.conversions
 
+        conv = conv.copy()
+
         self.parse_json = parse_json
         self.invalid_values = (invalid_values or {}).copy()
 
+        # Disable JSON parsing for Arrow
+        if self.results_type in ['arrow']:
+            conv[245] = None
+            self.parse_json = False
+
+        # Disable date/time parsing for polars; let polars do the parsing
+        elif self.results_type in ['polars']:
+            conv[7] = None
+            conv[10] = None
+            conv[12] = None
+
         # Need for MySQLdb compatibility.
         self.encoders = {k: v for (k, v) in conv.items() if type(k) is not int}
         self.decoders = {k: v for (k, v) in conv.items() if type(k) is int}