singlestoredb 1.14.2__cp38-abi3-win_amd64.whl → 1.15.0__cp38-abi3-win_amd64.whl

singlestoredb/functions/ext/timer.py

@@ -0,0 +1,98 @@
+import json
+import time
+from typing import Any
+from typing import Dict
+from typing import Optional
+
+from . import utils
+
+logger = utils.get_logger('singlestoredb.functions.ext.metrics')
+
+
+class RoundedFloatEncoder(json.JSONEncoder):
+
+    def encode(self, obj: Any) -> str:
+        if isinstance(obj, dict):
+            return '{' + ', '.join(
+                f'"{k}": {self._format_value(v)}'
+                for k, v in obj.items()
+            ) + '}'
+        return super().encode(obj)
+
+    def _format_value(self, value: Any) -> str:
+        if isinstance(value, float):
+            return f'{value:.2f}'
+        return json.dumps(value)
+
+
+class Timer:
+    """
+    Timer context manager that supports nested timing using a stack.
+
+    Example
+    -------
+    timer = Timer()
+
+    with timer('total'):
+        with timer('receive_data'):
+            time.sleep(0.1)
+        with timer('parse_input'):
+            time.sleep(0.2)
+        with timer('call_function'):
+            with timer('inner_operation'):
+                time.sleep(0.05)
+            time.sleep(0.3)
+
+    print(timer.metrics)
+    # {'receive_data': 0.1, 'parse_input': 0.2, 'inner_operation': 0.05,
+    #  'call_function': 0.35, 'total': 0.65}
+
+    """
+
+    def __init__(self, **kwargs: Any) -> None:
+        self.metadata: Dict[str, Any] = kwargs
+        self.metrics: Dict[str, float] = dict()
+        self.entries: Dict[str, float] = dict()
+        self._current_key: Optional[str] = None
+        self.start_time = time.perf_counter()
+
+    def __call__(self, key: str) -> 'Timer':
+        self._current_key = key
+        return self
+
+    def __enter__(self) -> 'Timer':
+        if self._current_key is None:
+            raise ValueError(
+                "No key specified. Use timer('key_name') as context manager.",
+            )
+        self.entries[self._current_key] = time.perf_counter()
+        return self
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        key = self._current_key
+        if key and key in self.entries:
+            start = self.entries.pop(key)
+            elapsed = time.perf_counter() - start
+            self.metrics[key] = elapsed
+        self._current_key = None
+
+    async def __aenter__(self) -> 'Timer':
+        return self.__enter__()
+
+    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        self.__exit__(exc_type, exc_val, exc_tb)
+
+    def reset(self) -> None:
+        self.metrics.clear()
+        self.entries.clear()
+        self._current_key = None
+
+    def finish(self) -> None:
+        """Finish the current timing context and store the elapsed time."""
+        self.metrics['total'] = time.perf_counter() - self.start_time
+        self.log_metrics()
+
+    def log_metrics(self) -> None:
+        if self.metadata.get('function'):
+            result = dict(type='function_metrics', **self.metadata, **self.metrics)
+            logger.info(json.dumps(result, cls=RoundedFloatEncoder))

singlestoredb/functions/typing/numpy.py

@@ -0,0 +1,20 @@
+import numpy as np
+import numpy.typing as npt
+
+NDArray = npt.NDArray
+
+StringArray = StrArray = npt.NDArray[np.str_]
+BytesArray = npt.NDArray[np.bytes_]
+Float32Array = FloatArray = npt.NDArray[np.float32]
+Float64Array = DoubleArray = npt.NDArray[np.float64]
+IntArray = npt.NDArray[np.int_]
+Int8Array = npt.NDArray[np.int8]
+Int16Array = npt.NDArray[np.int16]
+Int32Array = npt.NDArray[np.int32]
+Int64Array = npt.NDArray[np.int64]
+UInt8Array = npt.NDArray[np.uint8]
+UInt16Array = npt.NDArray[np.uint16]
+UInt32Array = npt.NDArray[np.uint32]
+UInt64Array = npt.NDArray[np.uint64]
+DateTimeArray = npt.NDArray[np.datetime64]
+TimeDeltaArray = npt.NDArray[np.timedelta64]

singlestoredb/functions/typing/pandas.py

@@ -0,0 +1,2 @@
+from pandas import DataFrame  # noqa: F401
+from pandas import Series  # noqa: F401

singlestoredb/functions/typing/polars.py

@@ -0,0 +1,2 @@
+from polars import DataFrame  # noqa: F401
+from polars import Series  # noqa: F401

singlestoredb/functions/typing/pyarrow.py

@@ -0,0 +1,2 @@
+from pyarrow import Array  # noqa: F401
+from pyarrow import Table  # noqa: F401

singlestoredb/magics/run_personal.py

@@ -1,6 +1,8 @@
 import os
 import tempfile
+from pathlib import Path
 from typing import Any
+from warnings import warn
 
 from IPython.core.interactiveshell import InteractiveShell
 from IPython.core.magic import line_magic
@@ -8,6 +10,8 @@ from IPython.core.magic import Magics
 from IPython.core.magic import magics_class
 from IPython.core.magic import needs_local_scope
 from IPython.core.magic import no_var_expand
+from IPython.utils.contexts import preserve_keys
+from IPython.utils.syspathcontext import prepended_to_syspath
 from jinja2 import Template
 
 
@@ -53,4 +57,81 @@ class RunPersonalMagic(Magics):
             # Execute the SQL command
             self.shell.run_line_magic('sql', sql_command)
             # Run the downloaded file
-            self.shell.run_line_magic('run', f'"{temp_file_path}"')
+            with preserve_keys(self.shell.user_ns, '__file__'):
+                self.shell.user_ns['__file__'] = temp_file_path
+                self.safe_execfile_ipy(temp_file_path, raise_exceptions=True)
+
+    def safe_execfile_ipy(
+        self,
+        fname: str,
+        shell_futures: bool = False,
+        raise_exceptions: bool = False,
+    ) -> None:
+        """Like safe_execfile, but for .ipy or .ipynb files with IPython syntax.
+
+        Parameters
+        ----------
+        fname : str
+            The name of the file to execute.  The filename must have a
+            .ipy or .ipynb extension.
+        shell_futures : bool (False)
+            If True, the code will share future statements with the interactive
+            shell. It will both be affected by previous __future__ imports, and
+            any __future__ imports in the code will affect the shell. If False,
+            __future__ imports are not shared in either direction.
+        raise_exceptions : bool (False)
+            If True raise exceptions everywhere.  Meant for testing.
+        """
+        fpath = Path(fname).expanduser().resolve()
+
+        # Make sure we can open the file
+        try:
+            with fpath.open('rb'):
+                pass
+        except Exception:
+            warn('Could not open file <%s> for safe execution.' % fpath)
+            return
+
+        # Find things also in current directory.  This is needed to mimic the
+        # behavior of running a script from the system command line, where
+        # Python inserts the script's directory into sys.path
+        dname = str(fpath.parent)
+
+        def get_cells() -> Any:
+            """generator for sequence of code blocks to run"""
+            if fpath.suffix == '.ipynb':
+                from nbformat import read
+                nb = read(fpath, as_version=4)
+                if not nb.cells:
+                    return
+                for cell in nb.cells:
+                    if cell.cell_type == 'code':
+                        if not cell.source.strip():
+                            continue
+                        if getattr(cell, 'metadata', {}).get('language', '') == 'sql':
+                            output_redirect = getattr(
+                                cell, 'metadata', {},
+                            ).get('output_variable', '') or ''
+                            if output_redirect:
+                                output_redirect = f' {output_redirect} <<'
+                            yield f'%%sql{output_redirect}\n{cell.source}'
+                        else:
+                            yield cell.source
+            else:
+                yield fpath.read_text(encoding='utf-8')
+
+        with prepended_to_syspath(dname):
+            try:
+                for cell in get_cells():
+                    result = self.shell.run_cell(
+                        cell, silent=True, shell_futures=shell_futures,
+                    )
+                    if raise_exceptions:
+                        result.raise_error()
+                    elif not result.success:
+                        break
+            except Exception:
+                if raise_exceptions:
+                    raise
+                self.shell.showtraceback()
+                warn('Unknown failure executing file: <%s>' % fpath)

singlestoredb/magics/run_shared.py

@@ -1,6 +1,8 @@
 import os
 import tempfile
+from pathlib import Path
 from typing import Any
+from warnings import warn
 
 from IPython.core.interactiveshell import InteractiveShell
 from IPython.core.magic import line_magic
@@ -8,6 +10,8 @@ from IPython.core.magic import Magics
 from IPython.core.magic import magics_class
 from IPython.core.magic import needs_local_scope
 from IPython.core.magic import no_var_expand
+from IPython.utils.contexts import preserve_keys
+from IPython.utils.syspathcontext import prepended_to_syspath
 from jinja2 import Template
 
 
@@ -50,4 +54,81 @@ class RunSharedMagic(Magics):
             # Execute the SQL command
             self.shell.run_line_magic('sql', sql_command)
             # Run the downloaded file
-            self.shell.run_line_magic('run', f'"{temp_file_path}"')
+            with preserve_keys(self.shell.user_ns, '__file__'):
+                self.shell.user_ns['__file__'] = temp_file_path
+                self.safe_execfile_ipy(temp_file_path, raise_exceptions=True)
+
+    def safe_execfile_ipy(
+        self,
+        fname: str,
+        shell_futures: bool = False,
+        raise_exceptions: bool = False,
+    ) -> None:
+        """Like safe_execfile, but for .ipy or .ipynb files with IPython syntax.
+
+        Parameters
+        ----------
+        fname : str
+            The name of the file to execute.  The filename must have a
+            .ipy or .ipynb extension.
+        shell_futures : bool (False)
+            If True, the code will share future statements with the interactive
+            shell. It will both be affected by previous __future__ imports, and
+            any __future__ imports in the code will affect the shell. If False,
+            __future__ imports are not shared in either direction.
+        raise_exceptions : bool (False)
+            If True raise exceptions everywhere.  Meant for testing.
+        """
+        fpath = Path(fname).expanduser().resolve()
+
+        # Make sure we can open the file
+        try:
+            with fpath.open('rb'):
+                pass
+        except Exception:
+            warn('Could not open file <%s> for safe execution.' % fpath)
+            return
+
+        # Find things also in current directory.  This is needed to mimic the
+        # behavior of running a script from the system command line, where
+        # Python inserts the script's directory into sys.path
+        dname = str(fpath.parent)
+
+        def get_cells() -> Any:
+            """generator for sequence of code blocks to run"""
+            if fpath.suffix == '.ipynb':
+                from nbformat import read
+                nb = read(fpath, as_version=4)
+                if not nb.cells:
+                    return
+                for cell in nb.cells:
+                    if cell.cell_type == 'code':
+                        if not cell.source.strip():
+                            continue
+                        if getattr(cell, 'metadata', {}).get('language', '') == 'sql':
+                            output_redirect = getattr(
+                                cell, 'metadata', {},
+                            ).get('output_variable', '') or ''
+                            if output_redirect:
+                                output_redirect = f' {output_redirect} <<'
+                            yield f'%%sql{output_redirect}\n{cell.source}'
+                        else:
+                            yield cell.source
+            else:
+                yield fpath.read_text(encoding='utf-8')
+
+        with prepended_to_syspath(dname):
+            try:
+                for cell in get_cells():
+                    result = self.shell.run_cell(
+                        cell, silent=True, shell_futures=shell_futures,
+                    )
+                    if raise_exceptions:
+                        result.raise_error()
+                    elif not result.success:
+                        break
+            except Exception:
+                if raise_exceptions:
+                    raise
+                self.shell.showtraceback()
+                warn('Unknown failure executing file: <%s>' % fpath)

singlestoredb/management/__init__.py

@@ -2,6 +2,7 @@
 from .cluster import manage_cluster
 from .files import manage_files
 from .manager import get_token
+from .region import manage_regions
 from .workspace import get_organization
 from .workspace import get_secret
 from .workspace import get_stage

singlestoredb/management/region.py

@@ -4,6 +4,7 @@ from typing import Dict
 from typing import Optional
 
 from .manager import Manager
+from .utils import NamedList
 from .utils import vars_to_str
 
 
@@ -65,3 +66,94 @@ class Region(object):
         )
         out._manager = manager
         return out
+
+
+class RegionManager(Manager):
+    """
+    SingleStoreDB region manager.
+
+    This class should be instantiated using :func:`singlestoredb.manage_regions`.
+
+    Parameters
+    ----------
+    access_token : str, optional
+        The API key or other access token for the workspace management API
+    version : str, optional
+        Version of the API to use
+    base_url : str, optional
+        Base URL of the workspace management API
+
+    See Also
+    --------
+    :func:`singlestoredb.manage_regions`
+    """
+
+    #: Object type
+    obj_type = 'region'
+
+    def list_regions(self) -> NamedList[Region]:
+        """
+        List all available regions.
+
+        Returns
+        -------
+        NamedList[Region]
+            List of available regions
+
+        Raises
+        ------
+        ManagementError
+            If there is an error getting the regions
+        """
+        res = self._get('regions')
+        return NamedList(
+            [Region.from_dict(item, self) for item in res.json()],
+        )
+
+    def list_shared_tier_regions(self) -> NamedList[Region]:
+        """
+        List regions that support shared tier workspaces.
+
+        Returns
+        -------
+        NamedList[Region]
+            List of regions that support shared tier workspaces
+
+        Raises
+        ------
+        ManagementError
+            If there is an error getting the regions
+        """
+        res = self._get('regions/sharedtier')
+        return NamedList(
+            [Region.from_dict(item, self) for item in res.json()],
+        )
+
+
+def manage_regions(
+    access_token: Optional[str] = None,
+    version: Optional[str] = None,
+    base_url: Optional[str] = None,
+) -> RegionManager:
+    """
+    Retrieve a SingleStoreDB region manager.
+
+    Parameters
+    ----------
+    access_token : str, optional
+        The API key or other access token for the workspace management API
+    version : str, optional
+        Version of the API to use
+    base_url : str, optional
+        Base URL of the workspace management API
+
+    Returns
+    -------
+    :class:`RegionManager`
+
+    """
+    return RegionManager(
+        access_token=access_token,
+        version=version,
+        base_url=base_url,
+    )

singlestoredb/management/workspace.py

@@ -1301,17 +1301,24 @@ class StarterWorkspace(object):
     See Also
     --------
     :meth:`WorkspaceManager.get_starter_workspace`
+    :meth:`WorkspaceManager.create_starter_workspace`
+    :meth:`WorkspaceManager.terminate_starter_workspace`
+    :meth:`WorkspaceManager.create_starter_workspace_user`
     :attr:`WorkspaceManager.starter_workspaces`
 
     """
 
     name: str
     id: str
+    database_name: str
+    endpoint: Optional[str]
 
     def __init__(
         self,
         name: str,
         id: str,
+        database_name: str,
+        endpoint: Optional[str] = None,
     ):
         #: Name of the starter workspace
         self.name = name
@@ -1319,6 +1326,13 @@ class StarterWorkspace(object):
         #: Unique ID of the starter workspace
         self.id = id
 
+        #: Name of the database associated with the starter workspace
+        self.database_name = database_name
+
+        #: Endpoint to connect to the starter workspace. The endpoint is in the form
+        #: of ``hostname:port``
+        self.endpoint = endpoint
+
         self._manager: Optional[WorkspaceManager] = None
 
     def __str__(self) -> str:
@@ -1351,10 +1365,63 @@ class StarterWorkspace(object):
         out = cls(
             name=obj['name'],
             id=obj['virtualWorkspaceID'],
+            database_name=obj['databaseName'],
+            endpoint=obj.get('endpoint'),
         )
         out._manager = manager
         return out
 
+    def connect(self, **kwargs: Any) -> connection.Connection:
+        """
+        Create a connection to the database server for this starter workspace.
+
+        Parameters
+        ----------
+        **kwargs : keyword-arguments, optional
+            Parameters to the SingleStoreDB `connect` function except host
+            and port which are supplied by the starter workspace object
+
+        Returns
+        -------
+        :class:`Connection`
+
+        """
+        if not self.endpoint:
+            raise ManagementError(
+                msg='An endpoint has not been set in this '
+                    'starter workspace configuration',
+            )
+        # Parse endpoint as host:port
+        if ':' in self.endpoint:
+            host, port = self.endpoint.split(':', 1)
+            kwargs['host'] = host
+            kwargs['port'] = int(port)
+        else:
+            kwargs['host'] = self.endpoint
+        return connection.connect(**kwargs)
+
+    def terminate(self) -> None:
+        """Terminate the starter workspace."""
+        if self._manager is None:
+            raise ManagementError(
+                msg='No workspace manager is associated with this object.',
+            )
+        self._manager._delete(f'sharedtier/virtualWorkspaces/{self.id}')
+
+    def refresh(self) -> StarterWorkspace:
+        """Update the object to the current state."""
+        if self._manager is None:
+            raise ManagementError(
+                msg='No workspace manager is associated with this object.',
+            )
+        new_obj = self._manager.get_starter_workspace(self.id)
+        for name, value in vars(new_obj).items():
+            if isinstance(value, Mapping):
+                setattr(self, name, snake_to_camel_dict(value))
+            else:
+                setattr(self, name, value)
+        return self
+
     @property
     def organization(self) -> Organization:
         if self._manager is None:
@@ -1375,7 +1442,7 @@ class StarterWorkspace(object):
     stages = stage
 
     @property
-    def starter_workspaces(self) -> NamedList[StarterWorkspace]:
+    def starter_workspaces(self) -> NamedList['StarterWorkspace']:
         """Return a list of available starter workspaces."""
         if self._manager is None:
             raise ManagementError(
@@ -1386,6 +1453,64 @@ class StarterWorkspace(object):
             [StarterWorkspace.from_dict(item, self._manager) for item in res.json()],
         )
 
+    def create_user(
+        self,
+        user_name: str,
+        password: Optional[str] = None,
+    ) -> Dict[str, str]:
+        """
+        Create a new user for this starter workspace.
+
+        Parameters
+        ----------
+        user_name : str
+            The starter workspace user name to connect the new user to the database
+        password : str, optional
+            Password for the new user. If not provided, a password will be
+            auto-generated by the system.
+
+        Returns
+        -------
+        Dict[str, str]
+            Dictionary containing 'userID' and 'password' of the created user
+
+        Raises
+        ------
+        ManagementError
+            If no workspace manager is associated with this object.
+        """
+        if self._manager is None:
+            raise ManagementError(
+                msg='No workspace manager is associated with this object.',
+            )
+
+        payload = {
+            'userName': user_name,
+        }
+        if password is not None:
+            payload['password'] = password
+
+        res = self._manager._post(
+            f'sharedtier/virtualWorkspaces/{self.id}/users',
+            json=payload,
+        )
+
+        response_data = res.json()
+        user_id = response_data.get('userID')
+        if not user_id:
+            raise ManagementError(msg='No userID returned from API')
+
+        # Return the password provided by user or generated by API
+        returned_password = password if password is not None \
+            else response_data.get('password')
+        if not returned_password:
+            raise ManagementError(msg='No password available from API response')
+
+        return {
+            'user_id': user_id,
+            'password': returned_password,
+        }
+
 
 class Billing(object):
     """Billing information."""
@@ -1522,6 +1647,14 @@ class WorkspaceManager(Manager):
         res = self._get('regions')
         return NamedList([Region.from_dict(item, self) for item in res.json()])
 
+    @ttl_property(datetime.timedelta(hours=1))
+    def shared_tier_regions(self) -> NamedList[Region]:
+        """Return a list of regions that support shared tier workspaces."""
+        res = self._get('regions/sharedtier')
+        return NamedList(
+            [Region.from_dict(item, self) for item in res.json()],
+        )
+
     def create_workspace_group(
         self,
         name: str,
@@ -1717,6 +1850,52 @@ class WorkspaceManager(Manager):
         res = self._get(f'sharedtier/virtualWorkspaces/{id}')
         return StarterWorkspace.from_dict(res.json(), manager=self)
 
+    def create_starter_workspace(
+        self,
+        name: str,
+        database_name: str,
+        workspace_group: dict[str, str],
+    ) -> 'StarterWorkspace':
+        """
+        Create a new starter (shared tier) workspace.
+
+        Parameters
+        ----------
+        name : str
+            Name of the starter workspace
+        database_name : str
+            Name of the database for the starter workspace
+        workspace_group : dict[str, str]
+            Workspace group input (dict with keys: 'cell_id')
+
+        Returns
+        -------
+        :class:`StarterWorkspace`
+        """
+        if not workspace_group or not isinstance(workspace_group, dict):
+            raise ValueError(
+                'workspace_group must be a dict with keys: '
+                "'cell_id'",
+            )
+        if set(workspace_group.keys()) != {'cell_id'}:
+            raise ValueError("workspace_group must contain only 'cell_id'")
+
+        payload = {
+            'name': name,
+            'databaseName': database_name,
+            'workspaceGroup': {
+                'cellID': workspace_group['cell_id'],
+            },
+        }
+
+        res = self._post('sharedtier/virtualWorkspaces', json=payload)
+        virtual_workspace_id = res.json().get('virtualWorkspaceID')
+        if not virtual_workspace_id:
+            raise ManagementError(msg='No virtualWorkspaceID returned from API')
+
+        res = self._get(f'sharedtier/virtualWorkspaces/{virtual_workspace_id}')
+        return StarterWorkspace.from_dict(res.json(), self)
+
 
 def manage_workspaces(
     access_token: Optional[str] = None,