singlestoredb 1.16.1__py3-none-any.whl

singlestoredb/management/cluster.py

@@ -0,0 +1,462 @@
+#!/usr/bin/env python
+"""SingleStoreDB Cluster Management."""
+import datetime
+import warnings
+from typing import Any
+from typing import Dict
+from typing import List
+from typing import Optional
+from typing import Union
+
+from .. import config
+from .. import connection
+from ..exceptions import ManagementError
+from .manager import Manager
+from .region import Region
+from .utils import NamedList
+from .utils import to_datetime
+from .utils import vars_to_str
+
+
+class Cluster(object):
+    """
+    SingleStoreDB cluster definition.
+
+    This object is not instantiated directly. It is used in the results
+    of API calls on the :class:`ClusterManager`. Clusters are created using
+    :meth:`ClusterManager.create_cluster`, or existing clusters are accessed by either
+    :attr:`ClusterManager.clusters` or by calling :meth:`ClusterManager.get_cluster`.
+
+    See Also
+    --------
+    :meth:`ClusterManager.create_cluster`
+    :meth:`ClusterManager.get_cluster`
+    :attr:`ClusterManager.clusters`
+
+    """
+
+    def __init__(
+        self, name: str, id: str, region: Region, size: str,
+        units: float, state: str, version: str,
+        created_at: Union[str, datetime.datetime],
+        expires_at: Optional[Union[str, datetime.datetime]] = None,
+        firewall_ranges: Optional[List[str]] = None,
+        terminated_at: Optional[Union[str, datetime.datetime]] = None,
+        endpoint: Optional[str] = None,
+    ):
+        """Use :attr:`ClusterManager.clusters` or :meth:`ClusterManager.get_cluster`."""
+        #: Name of the cluster
+        self.name = name.strip()
+
+        #: Unique ID of the cluster
+        self.id = id
+
+        #: Region of the cluster (see :class:`Region`)
+        self.region = region
+
+        #: Size of the cluster in cluster size notation (S-00, S-1, etc.)
+        self.size = size
+
+        #: Size of the cluster in units such as 0.25, 1.0, etc.
+        self.units = units
+
+        #: State of the cluster: PendingCreation, Transitioning, Active,
+        #: Terminated, Suspended, Resuming, Failed
+        self.state = state.strip()
+
+        #: Version of the SingleStoreDB server
+        self.version = version.strip()
+
+        #: Timestamp of when the cluster was created
+        self.created_at = to_datetime(created_at)
+
+        #: Timestamp of when the cluster expires
+        self.expires_at = to_datetime(expires_at)
+
+        #: List of allowed incoming IP addresses / ranges
+        self.firewall_ranges = firewall_ranges
+
+        #: Timestamp of when the cluster was terminated
+        self.terminated_at = to_datetime(terminated_at)
+
+        #: Hostname (or IP address) of the cluster database server
+        self.endpoint = endpoint
+
+        self._manager: Optional[ClusterManager] = None
+
+    def __str__(self) -> str:
+        """Return string representation."""
+        return vars_to_str(self)
+
+    def __repr__(self) -> str:
+        """Return string representation."""
+        return str(self)
+
+    @classmethod
+    def from_dict(cls, obj: Dict[str, Any], manager: 'ClusterManager') -> 'Cluster':
+        """
+        Construct a Cluster from a dictionary of values.
+
+        Parameters
+        ----------
+        obj : dict
+            Dictionary of values
+        manager : ClusterManager, optional
+            The ClusterManager the Cluster belongs to
+
+        Returns
+        -------
+        :class:`Cluster`
+
+        """
+        out = cls(
+            name=obj['name'], id=obj['clusterID'],
+            region=Region.from_dict(obj['region'], manager),
+            size=obj.get('size', 'Unknown'), units=obj.get('units', float('nan')),
+            state=obj['state'], version=obj['version'],
+            created_at=obj['createdAt'], expires_at=obj.get('expiresAt'),
+            firewall_ranges=obj.get('firewallRanges'),
+            terminated_at=obj.get('terminatedAt'),
+            endpoint=obj.get('endpoint'),
+        )
+        out._manager = manager
+        return out
+
+    def refresh(self) -> 'Cluster':
+        """Update the object to the current state."""
+        if self._manager is None:
+            raise ManagementError(
+                msg='No cluster manager is associated with this object.',
+            )
+        new_obj = self._manager.get_cluster(self.id)
+        for name, value in vars(new_obj).items():
+            setattr(self, name, value)
+        return self
+
+    def update(
+        self, name: Optional[str] = None,
+        admin_password: Optional[str] = None,
+        expires_at: Optional[str] = None,
+        size: Optional[str] = None, firewall_ranges: Optional[List[str]] = None,
+    ) -> None:
+        """
+        Update the cluster definition.
+
+        Parameters
+        ----------
+        name : str, optional
+            Cluster name
+        admim_password : str, optional
+            Admin password for the cluster
+        expires_at : str, optional
+            Timestamp when the cluster expires
+        size : str, optional
+            Cluster size in cluster size notation (S-00, S-1, etc.)
+        firewall_ranges : Sequence[str], optional
+            List of allowed incoming IP addresses
+
+        """
+        if self._manager is None:
+            raise ManagementError(
+                msg='No cluster manager is associated with this object.',
+            )
+        data = {
+            k: v for k, v in dict(
+                name=name, adminPassword=admin_password,
+                expiresAt=expires_at, size=size,
+                firewallRanges=firewall_ranges,
+            ).items() if v is not None
+        }
+        self._manager._patch(f'clusters/{self.id}', json=data)
+        self.refresh()
+
+    def suspend(
+        self,
+        wait_on_suspended: bool = False,
+        wait_interval: int = 20,
+        wait_timeout: int = 600,
+    ) -> None:
+        """
+        Suspend the cluster.
+
+        Parameters
+        ----------
+        wait_on_suspended : bool, optional
+            Wait for the cluster to go into 'Suspended' mode before returning
+        wait_interval : int, optional
+            Number of seconds between each server check
+        wait_timeout : int, optional
+            Total number of seconds to check server before giving up
+
+        Raises
+        ------
+        ManagementError
+            If timeout is reached
+
+        """
+        if self._manager is None:
+            raise ManagementError(
+                msg='No cluster manager is associated with this object.',
+            )
+        self._manager._post(
+            f'clusters/{self.id}/suspend',
+            headers={'Content-Type': 'application/x-www-form-urlencoded'},
+        )
+        if wait_on_suspended:
+            self._manager._wait_on_state(
+                self._manager.get_cluster(self.id),
+                'Suspended', interval=wait_interval, timeout=wait_timeout,
+            )
+            self.refresh()
+
+    def resume(
+        self,
+        wait_on_resumed: bool = False,
+        wait_interval: int = 20,
+        wait_timeout: int = 600,
+    ) -> None:
+        """
+        Resume the cluster.
+
+        Parameters
+        ----------
+        wait_on_resumed : bool, optional
+            Wait for the cluster to go into 'Resumed' or 'Active' mode before returning
+        wait_interval : int, optional
+            Number of seconds between each server check
+        wait_timeout : int, optional
+            Total number of seconds to check server before giving up
+
+        Raises
+        ------
+        ManagementError
+            If timeout is reached
+
+        """
+        if self._manager is None:
+            raise ManagementError(
+                msg='No cluster manager is associated with this object.',
+            )
+        self._manager._post(
+            f'clusters/{self.id}/resume',
+            headers={'Content-Type': 'application/x-www-form-urlencoded'},
+        )
+        if wait_on_resumed:
+            self._manager._wait_on_state(
+                self._manager.get_cluster(self.id),
+                ['Resumed', 'Active'], interval=wait_interval, timeout=wait_timeout,
+            )
+            self.refresh()
+
+    def terminate(
+        self,
+        wait_on_terminated: bool = False,
+        wait_interval: int = 10,
+        wait_timeout: int = 600,
+    ) -> None:
+        """
+        Terminate the cluster.
+
+        Parameters
+        ----------
+        wait_on_terminated : bool, optional
+            Wait for the cluster to go into 'Terminated' mode before returning
+        wait_interval : int, optional
+            Number of seconds between each server check
+        wait_timeout : int, optional
+            Total number of seconds to check server before giving up
+
+        Raises
+        ------
+        ManagementError
+            If timeout is reached
+
+        """
+        if self._manager is None:
+            raise ManagementError(
+                msg='No cluster manager is associated with this object.',
+            )
+        self._manager._delete(f'clusters/{self.id}')
+        if wait_on_terminated:
+            self._manager._wait_on_state(
+                self._manager.get_cluster(self.id),
+                'Terminated', interval=wait_interval, timeout=wait_timeout,
+            )
+            self.refresh()
+
+    def connect(self, **kwargs: Any) -> connection.Connection:
+        """
+        Create a connection to the database server for this cluster.
+
+        Parameters
+        ----------
+        **kwargs : keyword-arguments, optional
+            Parameters to the SingleStoreDB `connect` function except host
+            and port which are supplied by the cluster object
+
+        Returns
+        -------
+        :class:`Connection`
+
+        """
+        if not self.endpoint:
+            raise ManagementError(
+                msg='An endpoint has not been set in '
+                'this cluster configuration',
+            )
+        kwargs['host'] = self.endpoint
+        return connection.connect(**kwargs)
+
+
+class ClusterManager(Manager):
+    """
+    SingleStoreDB cluster manager.
+
+    This class should be instantiated using :func:`singlestoredb.manage_cluster`.
+
+    Parameters
+    ----------
+    access_token : str, optional
+        The API key or other access token for the cluster management API
+    version : str, optional
+        Version of the API to use
+    base_url : str, optional
+        Base URL of the cluster management API
+
+    See Also
+    --------
+    :func:`singlestoredb.manage_cluster`
+
+    """
+
+    #: Cluster management API version if none is specified.
+    default_version = 'v0beta'
+
+    #: Base URL if none is specified.
+    default_base_url = config.get_option('management.base_url') \
+        or 'https://api.singlestore.com'
+
+    #: Object type
+    obj_type = 'cluster'
+
+    @property
+    def clusters(self) -> NamedList[Cluster]:
+        """Return a list of available clusters."""
+        res = self._get('clusters')
+        return NamedList([Cluster.from_dict(item, self) for item in res.json()])
+
+    @property
+    def regions(self) -> NamedList[Region]:
+        """Return a list of available regions."""
+        res = self._get('regions')
+        return NamedList([Region.from_dict(item, self) for item in res.json()])
+
+    def create_cluster(
+        self, name: str, region: Union[str, Region], admin_password: str,
+        firewall_ranges: List[str], expires_at: Optional[str] = None,
+        size: Optional[str] = None, plan: Optional[str] = None,
+        wait_on_active: bool = False, wait_timeout: int = 600,
+        wait_interval: int = 20,
+    ) -> Cluster:
+        """
+        Create a new cluster.
+
+        Parameters
+        ----------
+        name : str
+            Name of the cluster
+        region : str or Region
+            The region ID of the cluster
+        admin_password : str
+            Admin password for the cluster
+        firewall_ranges : Sequence[str], optional
+            List of allowed incoming IP addresses
+        expires_at : str, optional
+            Timestamp of when the cluster expires
+        size : str, optional
+            Cluster size in cluster size notation (S-00, S-1, etc.)
+        plan : str, optional
+            Internal use only
+        wait_on_active : bool, optional
+            Wait for the cluster to be active before returning
+        wait_timeout : int, optional
+            Maximum number of seconds to wait before raising an exception
+            if wait=True
+        wait_interval : int, optional
+            Number of seconds between each polling interval
+
+        Returns
+        -------
+        :class:`Cluster`
+
+        """
+        if isinstance(region, Region) and region.id:
+            region = region.id
+        res = self._post(
+            'clusters', json=dict(
+                name=name, regionID=region, adminPassword=admin_password,
+                expiresAt=expires_at, size=size, firewallRanges=firewall_ranges,
+                plan=plan,
+            ),
+        )
+        out = self.get_cluster(res.json()['clusterID'])
+        if wait_on_active:
+            out = self._wait_on_state(
+                out, 'Active', interval=wait_interval,
+                timeout=wait_timeout,
+            )
+        return out
+
+    def get_cluster(self, id: str) -> Cluster:
+        """
+        Retrieve a cluster definition.
+
+        Parameters
+        ----------
+        id : str
+            ID of the cluster
+
+        Returns
+        -------
+        :class:`Cluster`
+
+        """
+        res = self._get(f'clusters/{id}')
+        return Cluster.from_dict(res.json(), manager=self)
+
+
+def manage_cluster(
+    access_token: Optional[str] = None,
+    version: Optional[str] = None,
+    base_url: Optional[str] = None,
+    *,
+    organization_id: Optional[str] = None,
+) -> ClusterManager:
+    """
+    Retrieve a SingleStoreDB cluster manager.
+
+    Parameters
+    ----------
+    access_token : str, optional
+        The API key or other access token for the cluster management API
+    version : str, optional
+        Version of the API to use
+    base_url : str, optional
+        Base URL of the cluster management API
+    organization_id: str, optional
+        ID of organization, if using a JWT for authentication
+
+    Returns
+    -------
+    :class:`ClusterManager`
+
+    """
+    warnings.warn(
+        'The cluster management API is deprecated; '
+        'use manage_workspaces instead.',
+        category=DeprecationWarning,
+    )
+    return ClusterManager(
+        access_token=access_token, base_url=base_url,
+        version=version, organization_id=organization_id,
+    )