singlestoredb 1.14.1__cp38-abi3-win32.whl → 1.15.0__cp38-abi3-win32.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/fusion/handler.py

@@ -33,7 +33,7 @@ CORE_GRAMMAR = r'''
     close_paren = ws* ")" ws*
     open_repeats = ws* ~r"[\(\[\{]" ws*
     close_repeats = ws* ~r"[\)\]\}]" ws*
-    select = ~r"SELECT"i ws+ ~r".+" ws*
+    statement = ~r"[\s\S]*" ws*
     table = ~r"(?:([A-Za-z0-9_\-]+)|`([^\`]+)`)(?:\.(?:([A-Za-z0-9_\-]+)|`([^\`]+)`))?" ws*
     column = ~r"(?:([A-Za-z0-9_\-]+)|`([^\`]+)`)(?:\.(?:([A-Za-z0-9_\-]+)|`([^\`]+)`))?" ws*
     link_name = ~r"(?:([A-Za-z0-9_\-]+)|`([^\`]+)`)(?:\.(?:([A-Za-z0-9_\-]+)|`([^\`]+)`))?" ws*
@@ -77,6 +77,7 @@ BUILTINS = {
     '<file-type>': r'''
     file_type = { FILE | FOLDER }
     ''',
+    '<statement>': '',
 }
 
 BUILTIN_DEFAULTS = {  # type: ignore
@@ -627,6 +628,18 @@ class SQLHandler(NodeVisitor):
         cls.compile()
         registry.register_handler(cls, overwrite=overwrite)
 
+    def create_result(self) -> result.FusionSQLResult:
+        """
+        Create a new result object.
+
+        Returns
+        -------
+        FusionSQLResult
+            A new result object for this handler
+
+        """
+        return result.FusionSQLResult()
+
     def execute(self, sql: str) -> result.FusionSQLResult:
         """
         Parse the SQL and invoke the handler method.
@@ -746,9 +759,9 @@ class SQLHandler(NodeVisitor):
         _, out, *_ = visited_children
         return out
 
-    def visit_select(self, node: Node, visited_children: Iterable[Any]) -> Any:
-        out = ' '.join(flatten(visited_children))
-        return {'select': out}
+    def visit_statement(self, node: Node, visited_children: Iterable[Any]) -> Any:
+        out = ' '.join(flatten(visited_children)).strip()
+        return {'statement': out}
 
     def visit_order_by(self, node: Node, visited_children: Iterable[Any]) -> Any:
         """Handle ORDER BY."""

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/export.py

@@ -215,7 +215,7 @@ class ExportService(object):
                 msg='Export ID is not set. You must start the export first.',
             )
 
-        self._manager._post(
+        self._manager._delete(
             f'workspaceGroups/{self.workspace_group.id}/egress/dropTableEgress',
             json=dict(egressID=self.export_id),
         )

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,
+    )