singlestoredb 1.16.1__py3-none-any.whl

singlestoredb/management/manager.py

@@ -0,0 +1,373 @@
+#!/usr/bin/env python
+"""SingleStoreDB Base Manager."""
+import os
+import sys
+import time
+from typing import Any
+from typing import Dict
+from typing import List
+from typing import Optional
+from typing import Union
+from urllib.parse import urljoin
+
+import jwt
+import requests
+
+from .. import config
+from ..exceptions import ManagementError
+from ..exceptions import OperationalError
+from .utils import get_token
+
+
+def set_organization(kwargs: Dict[str, Any]) -> None:
+    """Set the organization ID in the dictionary."""
+    if kwargs.get('params', {}).get('organizationID', None):
+        return
+
+    org = os.environ.get('SINGLESTOREDB_ORGANIZATION')
+    if org:
+        if 'params' not in kwargs:
+            kwargs['params'] = {}
+        kwargs['params']['organizationID'] = org
+
+
+def is_jwt(token: str) -> bool:
+    """Is the given token a JWT?"""
+    try:
+        jwt.decode(token, options={'verify_signature': False})
+        return True
+    except jwt.DecodeError:
+        return False
+
+
+class Manager(object):
+    """SingleStoreDB manager base class."""
+
+    #: Management API version if none is specified.
+    default_version = config.get_option('management.version') or 'v1'
+
+    #: Base URL if none is specified.
+    default_base_url = config.get_option('management.base_url') \
+        or 'https://api.singlestore.com'
+
+    #: Object type
+    obj_type = ''
+
+    def __init__(
+        self, access_token: Optional[str] = None, version: Optional[str] = None,
+        base_url: Optional[str] = None, *, organization_id: Optional[str] = None,
+    ):
+        from .. import __version__ as client_version
+        new_access_token = (
+            access_token or get_token()
+        )
+        if not new_access_token:
+            raise ManagementError(msg='No management token was configured.')
+
+        self._is_jwt = not access_token and new_access_token and is_jwt(new_access_token)
+        self._sess = requests.Session()
+        self._sess.headers.update({
+            'Authorization': f'Bearer {new_access_token}',
+            'Content-Type': 'application/json',
+            'Accept': 'application/json',
+            'User-Agent': f'SingleStoreDB-Python/{client_version}',
+        })
+
+        self._base_url = urljoin(
+            base_url
+            or config.get_option('management.base_url')
+            or type(self).default_base_url,
+            version or type(self).default_version,
+        ) + '/'
+
+        self._params: Dict[str, str] = {}
+        if organization_id:
+            self._params['organizationID'] = organization_id
+
+    def _check(
+        self, res: requests.Response, url: str, params: Dict[str, Any],
+    ) -> requests.Response:
+        """
+        Check the HTTP response status code and raise an exception as needed.
+
+        Parameters
+        ----------
+        res : requests.Response
+            HTTP response to check
+
+        Returns
+        -------
+        requests.Response
+
+        """
+        if config.get_option('debug.queries'):
+            print(os.path.join(self._base_url, url), params, file=sys.stderr)
+        if res.status_code >= 400:
+            txt = res.text.strip()
+            msg = f'{txt}: /{url}'
+            if params:
+                new_params = params.copy()
+                if 'json' in new_params:
+                    for k, v in new_params['json'].items():
+                        if 'password' in k.lower() and v:
+                            new_params['json'][k] = '*' * len(v)
+                msg += ': {}'.format(str(new_params))
+            raise ManagementError(errno=res.status_code, msg=msg, response=txt)
+        return res
+
+    def _doit(
+        self,
+        method: str,
+        path: str,
+        *args: Any,
+        **kwargs: Any,
+    ) -> requests.Response:
+        """Perform HTTP request."""
+        # Refresh the JWT as needed
+        if self._is_jwt:
+            self._sess.headers.update({'Authorization': f'Bearer {get_token()}'})
+        return getattr(self._sess, method.lower())(
+            urljoin(self._base_url, path), *args, **kwargs,
+        )
+
+    def _get(self, path: str, *args: Any, **kwargs: Any) -> requests.Response:
+        """
+        Invoke a GET request.
+
+        Parameters
+        ----------
+        path : str
+            Path of the resource
+        *args : positional arguments, optional
+            Arguments to add to the GET request
+        **kwargs : keyword arguments, optional
+            Keyword arguments to add to the GET request
+
+        Returns
+        -------
+        requests.Response
+
+        """
+        if self._params:
+            params = dict(self._params)
+            params.update(kwargs.get('params', {}))
+            kwargs['params'] = params
+        set_organization(kwargs)
+        return self._check(self._doit('get', path, *args, **kwargs), path, kwargs)
+
+    def _post(self, path: str, *args: Any, **kwargs: Any) -> requests.Response:
+        """
+        Invoke a POST request.
+
+        Parameters
+        ----------
+        path : str
+            Path of the resource
+        *args : positional arguments, optional
+            Arguments to add to the POST request
+        **kwargs : keyword arguments, optional
+            Keyword arguments to add to the POST request
+
+        Returns
+        -------
+        requests.Response
+
+        """
+        if self._params:
+            params = dict(self._params)
+            params.update(kwargs.get('params', {}))
+            kwargs['params'] = params
+        set_organization(kwargs)
+        return self._check(self._doit('post', path, *args, **kwargs), path, kwargs)
+
+    def _put(self, path: str, *args: Any, **kwargs: Any) -> requests.Response:
+        """
+        Invoke a PUT request.
+
+        Parameters
+        ----------
+        path : str
+            Path of the resource
+        *args : positional arguments, optional
+            Arguments to add to the POST request
+        **kwargs : keyword arguments, optional
+            Keyword arguments to add to the POST request
+
+        Returns
+        -------
+        requests.Response
+
+        """
+        if self._params:
+            params = dict(self._params)
+            params.update(kwargs.get('params', {}))
+            kwargs['params'] = params
+        set_organization(kwargs)
+        return self._check(self._doit('put', path, *args, **kwargs), path, kwargs)
+
+    def _delete(self, path: str, *args: Any, **kwargs: Any) -> requests.Response:
+        """
+        Invoke a DELETE request.
+
+        Parameters
+        ----------
+        path : str
+            Path of the resource
+        *args : positional arguments, optional
+            Arguments to add to the DELETE request
+        **kwargs : keyword arguments, optional
+            Keyword arguments to add to the DELETE request
+
+        Returns
+        -------
+        requests.Response
+
+        """
+        if self._params:
+            params = dict(self._params)
+            params.update(kwargs.get('params', {}))
+            kwargs['params'] = params
+        set_organization(kwargs)
+        return self._check(self._doit('delete', path, *args, **kwargs), path, kwargs)
+
+    def _patch(self, path: str, *args: Any, **kwargs: Any) -> requests.Response:
+        """
+        Invoke a PATCH request.
+
+        Parameters
+        ----------
+        path : str
+            Path of the resource
+        *args : positional arguments, optional
+            Arguments to add to the PATCH request
+        **kwargs : keyword arguments, optional
+            Keyword arguments to add to the PATCH request
+
+        Returns
+        -------
+        requests.Response
+
+        """
+        if self._params:
+            params = dict(self._params)
+            params.update(kwargs.get('params', {}))
+            kwargs['params'] = params
+        set_organization(kwargs)
+        return self._check(self._doit('patch', path, *args, **kwargs), path, kwargs)
+
+    def _wait_on_state(
+        self,
+        out: Any,
+        state: Union[str, List[str]],
+        interval: int = 20,
+        timeout: int = 600,
+    ) -> Any:
+        """
+        Wait on server state before continuing.
+
+        Parameters
+        ----------
+        out : Any
+            Current object
+        state : str or List[str]
+            State(s) to wait for
+        interval : int, optional
+            Interval between each server poll
+        timeout : int, optional
+            Maximum time to wait before raising an exception
+
+        Raises
+        ------
+        ManagementError
+            If timeout is reached
+
+        Returns
+        -------
+        Same object type as `out`
+
+        """
+        states = [
+            x.lower().strip()
+            for x in (isinstance(state, str) and [state] or state)
+        ]
+
+        if getattr(out, 'state', None) is None:
+            raise ManagementError(
+                msg='{} object does not have a `state` attribute'.format(
+                    type(out).__name__,
+                ),
+            )
+
+        while True:
+            if getattr(out, 'state').lower() in states:
+                break
+            if timeout <= 0:
+                raise ManagementError(
+                    msg=f'Exceeded waiting time for {self.obj_type} to become '
+                        '{}.'.format(', '.join(states)),
+                )
+            time.sleep(interval)
+            timeout -= interval
+            out = getattr(self, f'get_{self.obj_type}')(out.id)
+
+        return out
+
+    def _wait_on_endpoint(
+        self,
+        out: Any,
+        interval: int = 10,
+        timeout: int = 300,
+    ) -> Any:
+        """
+        Wait for the endpoint to be ready by attempting to connect.
+
+        Parameters
+        ----------
+        out : Any
+            Workspace object with a connect method
+        interval : int, optional
+            Interval between each connection attempt (default: 10 seconds)
+        timeout : int, optional
+            Maximum time to wait before raising an exception (default: 300 seconds)
+
+        Raises
+        ------
+        ManagementError
+            If timeout is reached or endpoint is not available
+
+        Returns
+        -------
+        Same object type as `out`
+
+        """
+        # Only wait if workload type is set which means we are in the
+        # notebook environment. Outside of the environment, the endpoint
+        # may not be reachable directly.
+        if not os.environ.get('SINGLESTOREDB_WORKLOAD_TYPE', ''):
+            return out
+
+        if not hasattr(out, 'connect') or not out.connect:
+            raise ManagementError(
+                msg=f'{type(out).__name__} object does not have a valid endpoint',
+            )
+
+        while True:
+            try:
+                # Try to establish a connection to the endpoint using context manager
+                with out.connect(connect_timeout=5):
+                    pass
+            except Exception as exc:
+                # If we get an 'access denied' error, that means that the server is
+                # up and we just aren't authenticating.
+                if isinstance(exc, OperationalError) and exc.errno == 1045:
+                    break
+                # If connection fails, check timeout and retry
+                if timeout <= 0:
+                    raise ManagementError(
+                        msg=f'Exceeded waiting time for {self.obj_type} endpoint '
+                            'to become ready',
+                    )
+                time.sleep(interval)
+                timeout -= interval
+
+        return out

singlestoredb/management/organization.py

@@ -0,0 +1,226 @@
+#!/usr/bin/env python
+"""SingleStoreDB Cloud Organization."""
+import datetime
+from typing import Dict
+from typing import List
+from typing import Optional
+from typing import Union
+
+from ..exceptions import ManagementError
+from .inference_api import InferenceAPIManager
+from .job import JobsManager
+from .manager import Manager
+from .utils import vars_to_str
+
+
+def listify(x: Union[str, List[str]]) -> List[str]:
+    if isinstance(x, list):
+        return x
+    return [x]
+
+
+def stringify(x: Union[str, List[str]]) -> str:
+    if isinstance(x, list):
+        return x[0]
+    return x
+
+
+class Secret(object):
+    """
+    SingleStoreDB secrets definition.
+
+    This object is not directly instantiated. It is used in results
+    of API calls on the :class:`Organization`. See :meth:`Organization.get_secret`.
+    """
+
+    def __init__(
+        self,
+        id: str,
+        name: str,
+        created_by: str,
+        created_at: Union[str, datetime.datetime],
+        last_updated_by: str,
+        last_updated_at: Union[str, datetime.datetime],
+        value: Optional[str] = None,
+        deleted_by: Optional[str] = None,
+        deleted_at: Optional[Union[str, datetime.datetime]] = None,
+    ):
+        # UUID of the secret
+        self.id = id
+
+        # Name of the secret
+        self.name = name
+
+        # Value of the secret
+        self.value = value
+
+        # User who created the secret
+        self.created_by = created_by
+
+        # Time when the secret was created
+        self.created_at = created_at
+
+        # UUID of the user who last updated the secret
+        self.last_updated_by = last_updated_by
+
+        # Time when the secret was last updated
+        self.last_updated_at = last_updated_at
+
+        # UUID of the user who deleted the secret
+        self.deleted_by = deleted_by
+
+        # Time when the secret was deleted
+        self.deleted_at = deleted_at
+
+    @classmethod
+    def from_dict(cls, obj: Dict[str, str]) -> 'Secret':
+        """
+        Construct a Secret from a dictionary of values.
+
+        Parameters
+        ----------
+        obj : dict
+            Dictionary of values
+
+        Returns
+        -------
+        :class:`Secret`
+
+        """
+        out = cls(
+            id=obj['secretID'],
+            name=obj['name'],
+            created_by=obj['createdBy'],
+            created_at=obj['createdAt'],
+            last_updated_by=obj['lastUpdatedBy'],
+            last_updated_at=obj['lastUpdatedAt'],
+            value=obj.get('value'),
+            deleted_by=obj.get('deletedBy'),
+            deleted_at=obj.get('deletedAt'),
+        )
+
+        return out
+
+    def __str__(self) -> str:
+        """Return string representation."""
+        return vars_to_str(self)
+
+    def __repr__(self) -> str:
+        """Return string representation."""
+        return str(self)
+
+
+class Organization(object):
+    """
+    Organization in SingleStoreDB Cloud portal.
+
+    This object is not directly instantiated. It is used in results
+    of ``WorkspaceManager`` API calls.
+
+    See Also
+    --------
+    :attr:`WorkspaceManager.organization`
+
+    """
+
+    id: str
+    name: str
+    firewall_ranges: List[str]
+
+    def __init__(self, id: str, name: str, firewall_ranges: List[str]):
+        """Use :attr:`WorkspaceManager.organization` instead."""
+        #: Unique ID of the organization
+        self.id = id
+
+        #: Name of the organization
+        self.name = name
+
+        #: Firewall ranges of the organization
+        self.firewall_ranges = list(firewall_ranges)
+
+        self._manager: Optional[Manager] = None
+
+    def __str__(self) -> str:
+        """Return string representation."""
+        return vars_to_str(self)
+
+    def __repr__(self) -> str:
+        """Return string representation."""
+        return str(self)
+
+    def get_secret(self, name: str) -> Secret:
+        if self._manager is None:
+            raise ManagementError(msg='Organization not initialized')
+
+        res = self._manager._get('secrets', params=dict(name=name))
+
+        secrets = [Secret.from_dict(item) for item in res.json()['secrets']]
+
+        if len(secrets) == 0:
+            raise ManagementError(msg=f'Secret {name} not found')
+
+        if len(secrets) > 1:
+            raise ManagementError(msg=f'Multiple secrets found for {name}')
+
+        return secrets[0]
+
+    @classmethod
+    def from_dict(
+        cls,
+        obj: Dict[str, Union[str, List[str]]],
+        manager: Manager,
+    ) -> 'Organization':
+        """
+        Convert dictionary to an ``Organization`` object.
+
+        Parameters
+        ----------
+        obj : dict
+            Key-value pairs to retrieve organization information from
+        manager : WorkspaceManager, optional
+            The WorkspaceManager the Organization belongs to
+
+        Returns
+        -------
+        :class:`Organization`
+
+        """
+        out = cls(
+            id=stringify(obj['orgID']),
+            name=stringify(obj.get('name', '<unknown>')),
+            firewall_ranges=listify(obj.get('firewallRanges', [])),
+        )
+        out._manager = manager
+        return out
+
+    @property
+    def jobs(self) -> JobsManager:
+        """
+        Retrieve a SingleStoreDB scheduled job manager.
+
+        Parameters
+        ----------
+        manager : WorkspaceManager, optional
+            The WorkspaceManager the JobsManager belongs to
+
+        Returns
+        -------
+        :class:`JobsManager`
+        """
+        return JobsManager(self._manager)
+
+    @property
+    def inference_apis(self) -> InferenceAPIManager:
+        """
+        Retrieve a SingleStoreDB inference api manager.
+
+        Parameters
+        ----------
+        manager : WorkspaceManager, optional
+            The WorkspaceManager the InferenceAPIManager belongs to
+
+        Returns
+        -------
+        :class:`InferenceAPIManager`
+        """
+        return InferenceAPIManager(self._manager)