singlestoredb 1.0.4__py3-none-any.whl → 1.2.0__py3-none-any.whl

singlestoredb/mysql/cursors.py

@@ -4,7 +4,9 @@ from collections import namedtuple
 
 from . import err
 from ..connection import Cursor as BaseCursor
+from ..utils import results
 from ..utils.debug import log_query
+from ..utils.results import get_schema
 
 
 #: Regular expression for :meth:`Cursor.executemany`.
@@ -45,6 +47,7 @@ class Cursor(BaseCursor):
         self._connection = connection
         self.warning_count = 0
         self._description = None
+        self._format_schema = None
         self._rownumber = 0
         self.rowcount = -1
         self.arraysize = 1
@@ -62,6 +65,10 @@ class Cursor(BaseCursor):
     def description(self):
         return self._description
 
+    @property
+    def _schema(self):
+        return self._format_schema
+
     @property
     def connection(self):
         return self._connection
@@ -394,6 +401,7 @@ class Cursor(BaseCursor):
         self.rowcount = 0
         self.warning_count = 0
         self._description = None
+        self._format_schema = None
         self.lastrowid = None
         self._rows = None
 
@@ -409,6 +417,11 @@ class Cursor(BaseCursor):
         if self.rowcount == 18446744073709551615:
             self.rowcount = -1
         self._description = result.description
+        if self._description:
+            self._format_schema = get_schema(
+                self.connection._results_type,
+                result.description,
+            )
         self.lastrowid = result.insert_id
         self._rows = result.rows
 
@@ -440,6 +453,134 @@ class CursorSV(Cursor):
     """Cursor class for C extension."""
 
 
+class ArrowCursorMixin:
+    """Fetch methods for Arrow Tables."""
+
+    def fetchone(self):
+        return results.results_to_arrow(
+            self.description, super().fetchone(), single=True, schema=self._schema,
+        )
+
+    def fetchall(self):
+        return results.results_to_arrow(
+            self.description, super().fetchall(), schema=self._schema,
+        )
+
+    def fetchall_unbuffered(self):
+        return results.results_to_arrow(
+            self.description, super().fetchall_unbuffered(), schema=self._schema,
+        )
+
+    def fetchmany(self, size=None):
+        return results.results_to_arrow(
+            self.description, super().fetchmany(size), schema=self._schema,
+        )
+
+
+class ArrowCursor(ArrowCursorMixin, Cursor):
+    """A cursor which returns results as an Arrow Table."""
+
+
+class ArrowCursorSV(ArrowCursorMixin, CursorSV):
+    """A cursor which returns results as an Arrow Table for C extension."""
+
+
+class NumpyCursorMixin:
+    """Fetch methods for numpy arrays."""
+
+    def fetchone(self):
+        return results.results_to_numpy(
+            self.description, super().fetchone(), single=True, schema=self._schema,
+        )
+
+    def fetchall(self):
+        return results.results_to_numpy(
+            self.description, super().fetchall(), schema=self._schema,
+        )
+
+    def fetchall_unbuffered(self):
+        return results.results_to_numpy(
+            self.description, super().fetchall_unbuffered(), schema=self._schema,
+        )
+
+    def fetchmany(self, size=None):
+        return results.results_to_numpy(
+            self.description, super().fetchmany(size), schema=self._schema,
+        )
+
+
+class NumpyCursor(NumpyCursorMixin, Cursor):
+    """A cursor which returns results as a numpy array."""
+
+
+class NumpyCursorSV(NumpyCursorMixin, CursorSV):
+    """A cursor which returns results as a numpy array for C extension."""
+
+
+class PandasCursorMixin:
+    """Fetch methods for pandas DataFrames."""
+
+    def fetchone(self):
+        return results.results_to_pandas(
+            self.description, super().fetchone(), single=True, schema=self._schema,
+        )
+
+    def fetchall(self):
+        return results.results_to_pandas(
+            self.description, super().fetchall(), schema=self._schema,
+        )
+
+    def fetchall_unbuffered(self):
+        return results.results_to_pandas(
+            self.description, super().fetchall_unbuffered(), schema=self._schema,
+        )
+
+    def fetchmany(self, size=None):
+        return results.results_to_pandas(
+            self.description, super().fetchmany(size), schema=self._schema,
+        )
+
+
+class PandasCursor(PandasCursorMixin, Cursor):
+    """A cursor which returns results as a pandas DataFrame."""
+
+
+class PandasCursorSV(PandasCursorMixin, CursorSV):
+    """A cursor which returns results as a pandas DataFrame for C extension."""
+
+
+class PolarsCursorMixin:
+    """Fetch methods for polars DataFrames."""
+
+    def fetchone(self):
+        return results.results_to_polars(
+            self.description, super().fetchone(), single=True, schema=self._schema,
+        )
+
+    def fetchall(self):
+        return results.results_to_polars(
+            self.description, super().fetchall(), schema=self._schema,
+        )
+
+    def fetchall_unbuffered(self):
+        return results.results_to_polars(
+            self.description, super().fetchall_unbuffered(), schema=self._schema,
+        )
+
+    def fetchmany(self, size=None):
+        return results.results_to_polars(
+            self.description, super().fetchmany(size), schema=self._schema,
+        )
+
+
+class PolarsCursor(PolarsCursorMixin, Cursor):
+    """A cursor which returns results as a polars DataFrame."""
+
+
+class PolarsCursorSV(PolarsCursorMixin, CursorSV):
+    """A cursor which returns results as a polars DataFrame for C extension."""
+
+
 class DictCursorMixin:
     # You can override this to use OrderedDict or other dict-like types.
     dict_type = dict
@@ -698,16 +839,48 @@ class SSCursorSV(SSCursor):
 
 
 class SSDictCursor(DictCursorMixin, SSCursor):
-    """An unbuffered cursor, which returns results as a dictionary"""
+    """An unbuffered cursor, which returns results as a dictionary."""
 
 
 class SSDictCursorSV(SSCursorSV):
-    """An unbuffered cursor for the C extension, which returns a dictionary"""
+    """An unbuffered cursor for the C extension, which returns a dictionary."""
 
 
 class SSNamedtupleCursor(NamedtupleCursorMixin, SSCursor):
-    """An unbuffered cursor, which returns results as a named tuple"""
+    """An unbuffered cursor, which returns results as a named tuple."""
 
 
 class SSNamedtupleCursorSV(SSCursorSV):
-    """An unbuffered cursor for the C extension, which returns results as a named tuple"""
+    """An unbuffered cursor for the C extension, which returns results as named tuple."""
+
+
+class SSArrowCursor(ArrowCursorMixin, SSCursor):
+    """An unbuffered cursor, which returns results as an Arrow Table."""
+
+
+class SSArrowCursorSV(ArrowCursorMixin, SSCursorSV):
+    """An unbuffered cursor, which returns results as an Arrow Table (accelerated)."""
+
+
+class SSNumpyCursor(NumpyCursorMixin, SSCursor):
+    """An unbuffered cursor, which returns results as a numpy array."""
+
+
+class SSNumpyCursorSV(NumpyCursorMixin, SSCursorSV):
+    """An unbuffered cursor, which returns results as a numpy array (accelerated)."""
+
+
+class SSPandasCursor(PandasCursorMixin, SSCursor):
+    """An unbuffered cursor, which returns results as a pandas DataFrame."""
+
+
+class SSPandasCursorSV(PandasCursorMixin, SSCursorSV):
+    """An unbuffered cursor, which returns results as a pandas DataFrame (accelerated)."""
+
+
+class SSPolarsCursor(PolarsCursorMixin, SSCursor):
+    """An unbuffered cursor, which returns results as a polars DataFrame."""
+
+
+class SSPolarsCursorSV(PolarsCursorMixin, SSCursorSV):
+    """An unbuffered cursor, which returns results as a polars DataFrame (accelerated)."""

singlestoredb/mysql/protocol.py

@@ -8,8 +8,10 @@ from . import err
 from ..config import get_option
 from ..utils.results import Description
 from .charset import MBLENGTH
+from .constants import EXTENDED_TYPE
 from .constants import FIELD_TYPE
 from .constants import SERVER_STATUS
+from .constants import VECTOR_TYPE
 
 
 DEBUG = get_option('debug.connection')
@@ -264,16 +266,63 @@ class FieldDescriptorPacket(MysqlPacket):
         self.org_table = self.read_length_coded_string().decode(encoding)
         self.name = self.read_length_coded_string().decode(encoding)
         self.org_name = self.read_length_coded_string().decode(encoding)
+        n_bytes = 0
         (
+            n_bytes,
             self.charsetnr,
             self.length,
             self.type_code,
             self.flags,
             self.scale,
-        ) = self.read_struct('<xHIBHBxx')
+        ) = self.read_struct('<BHIBHBxx')
+
         # 'default' is a length coded binary and is still in the buffer?
         # not used for normal result sets...
 
+        # Extended types
+        if n_bytes > 12:
+            ext_type_code = self.read_uint8()
+            if ext_type_code == EXTENDED_TYPE.NONE:
+                pass
+            elif ext_type_code == EXTENDED_TYPE.BSON:
+                self.type_code = FIELD_TYPE.BSON
+            elif ext_type_code == EXTENDED_TYPE.VECTOR:
+                (self.length, vec_type) = self.read_struct('<IB')
+                if vec_type == VECTOR_TYPE.FLOAT32:
+                    if self.charsetnr == 63:
+                        self.type_code = FIELD_TYPE.FLOAT32_VECTOR
+                    else:
+                        self.type_code = FIELD_TYPE.FLOAT32_VECTOR_JSON
+                elif vec_type == VECTOR_TYPE.FLOAT64:
+                    if self.charsetnr == 63:
+                        self.type_code = FIELD_TYPE.FLOAT64_VECTOR
+                    else:
+                        self.type_code = FIELD_TYPE.FLOAT64_VECTOR_JSON
+                elif vec_type == VECTOR_TYPE.INT8:
+                    if self.charsetnr == 63:
+                        self.type_code = FIELD_TYPE.INT8_VECTOR
+                    else:
+                        self.type_code = FIELD_TYPE.INT8_VECTOR_JSON
+                elif vec_type == VECTOR_TYPE.INT16:
+                    if self.charsetnr == 63:
+                        self.type_code = FIELD_TYPE.INT16_VECTOR
+                    else:
+                        self.type_code = FIELD_TYPE.INT16_VECTOR_JSON
+                elif vec_type == VECTOR_TYPE.INT32:
+                    if self.charsetnr == 63:
+                        self.type_code = FIELD_TYPE.INT32_VECTOR
+                    else:
+                        self.type_code = FIELD_TYPE.INT32_VECTOR_JSON
+                elif vec_type == VECTOR_TYPE.INT64:
+                    if self.charsetnr == 63:
+                        self.type_code = FIELD_TYPE.INT64_VECTOR
+                    else:
+                        self.type_code = FIELD_TYPE.INT64_VECTOR_JSON
+                else:
+                    raise TypeError(f'unrecognized vector data type: {vec_type}')
+            else:
+                raise TypeError(f'unrecognized extended data type: {ext_type_code}')
+
     def description(self):
         """Provides a 7-item tuple compatible with the Python PEP249 DB Spec."""
         return Description(

singlestoredb/notebook/__init__.py

@@ -0,0 +1,15 @@
+#!/usr/bin/env python
+import os as _os
+import warnings as _warnings
+
+from ._objects import organization  # noqa: F401
+from ._objects import secrets  # noqa: F401
+from ._objects import stage  # noqa: F401
+from ._objects import workspace  # noqa: F401
+from ._objects import workspace_group  # noqa: F401
+
+if 'SINGLESTOREDB_ORGANIZATION' not in _os.environ:
+    _warnings.warn(
+        'This package is intended for use in the SingleStoreDB notebook environment',
+        RuntimeWarning,
+    )

singlestoredb/notebook/_objects.py

@@ -0,0 +1,212 @@
+#!/usr/bin/env python
+import functools
+from typing import Any
+
+from ..management import workspace as _ws
+
+
+class Secrets(object):
+    """Wrapper for accessing secrets as object attributes."""
+
+    def __getattr__(self, name: str) -> str:
+        if name.startswith('_ipython') or name.startswith('_repr_'):
+            raise AttributeError(name)
+        return _ws.get_secret(name)
+
+    def __getitem__(self, name: str) -> str:
+        return _ws.get_secret(name)
+
+
+class Stage(object):
+
+    def __new__(cls, *args: Any, **kwargs: Any) -> Any:
+        # We are remapping the methods and attributes here so that
+        # autocomplete still works in Jupyter / IPython, but we
+        # bypass the real method / attribute calls and apply them
+        # to the currently selected stage.
+        for name in [x for x in dir(_ws.Stage) if not x.startswith('_')]:
+            if name in ['from_dict', 'refresh', 'update']:
+                continue
+            attr = getattr(_ws.Stage, name)
+
+            def make_wrapper(m: str, is_method: bool = False) -> Any:
+                if is_method:
+                    def wrap(self: Stage, *a: Any, **kw: Any) -> Any:
+                        return getattr(_ws.get_stage(), m)(*a, **kw)
+                    return functools.update_wrapper(wrap, attr)
+                else:
+                    def wrap(self: Stage, *a: Any, **kw: Any) -> Any:
+                        return getattr(_ws.get_stage(), m)
+                    return property(functools.update_wrapper(wrap, attr))
+
+            setattr(cls, name, make_wrapper(m=name, is_method=callable(attr)))
+
+        for name in [
+            x for x in _ws.Stage.__annotations__.keys()
+            if not x.startswith('_')
+        ]:
+
+            def make_wrapper(m: str, is_method: bool = False) -> Any:
+                def wrap(self: Stage) -> Any:
+                    return getattr(_ws.get_stage(), m)
+                return property(functools.update_wrapper(wrap, attr))
+
+            setattr(cls, name, make_wrapper(m=name))
+
+        cls.__doc__ = _ws.Stage.__doc__
+
+        return super().__new__(cls, *args, **kwargs)
+
+
+class WorkspaceGroup(object):
+
+    def __new__(cls, *args: Any, **kwargs: Any) -> Any:
+        # We are remapping the methods and attributes here so that
+        # autocomplete still works in Jupyter / IPython, but we
+        # bypass the real method / attribute calls and apply them
+        # to the currently selected workspace group.
+        for name in [x for x in dir(_ws.WorkspaceGroup) if not x.startswith('_')]:
+            if name in ['from_dict', 'refresh', 'update']:
+                continue
+
+            attr = getattr(_ws.WorkspaceGroup, name)
+
+            def make_wrapper(m: str, is_method: bool = False) -> Any:
+                if is_method:
+                    def wrap(self: WorkspaceGroup, *a: Any, **kw: Any) -> Any:
+                        return getattr(_ws.get_workspace_group(), m)(*a, **kw)
+                    return functools.update_wrapper(wrap, attr)
+                else:
+                    def wrap(self: WorkspaceGroup, *a: Any, **kw: Any) -> Any:
+                        return getattr(_ws.get_workspace_group(), m)
+                    return property(functools.update_wrapper(wrap, attr))
+
+            setattr(cls, name, make_wrapper(m=name, is_method=callable(attr)))
+
+        for name in [
+            x for x in _ws.WorkspaceGroup.__annotations__.keys()
+            if not x.startswith('_')
+        ]:
+
+            def make_wrapper(m: str, is_method: bool = False) -> Any:
+                def wrap(self: WorkspaceGroup) -> Any:
+                    return getattr(_ws.get_workspace_group(), m)
+                return property(functools.update_wrapper(wrap, attr))
+
+            setattr(cls, name, make_wrapper(m=name))
+
+        cls.__doc__ = _ws.WorkspaceGroup.__doc__
+
+        return super().__new__(cls, *args, **kwargs)
+
+    def __str__(self) -> str:
+        return _ws.get_workspace_group().__str__()
+
+    def __repr__(self) -> str:
+        return _ws.get_workspace_group().__repr__()
+
+
+class Workspace(object):
+
+    def __new__(cls, *args: Any, **kwargs: Any) -> Any:
+        # We are remapping the methods and attributes here so that
+        # autocomplete still works in Jupyter / IPython, but we
+        # bypass the real method / attribute calls and apply them
+        # to the currently selected workspace.
+        for name in [x for x in dir(_ws.Workspace) if not x.startswith('_')]:
+            if name in ['from_dict', 'refresh', 'update']:
+                continue
+
+            attr = getattr(_ws.Workspace, name)
+
+            def make_wrapper(m: str, is_method: bool = False) -> Any:
+                if is_method:
+                    def wrap(self: Workspace, *a: Any, **kw: Any) -> Any:
+                        return getattr(_ws.get_workspace(), m)(*a, **kw)
+                    return functools.update_wrapper(wrap, attr)
+                else:
+                    def wrap(self: Workspace, *a: Any, **kw: Any) -> Any:
+                        return getattr(_ws.get_workspace(), m)
+                    return property(functools.update_wrapper(wrap, attr))
+
+            setattr(cls, name, make_wrapper(m=name, is_method=callable(attr)))
+
+        for name in [
+            x for x in _ws.Workspace.__annotations__.keys()
+            if not x.startswith('_')
+        ]:
+
+            def make_wrapper(m: str, is_method: bool = False) -> Any:
+                def wrap(self: Workspace) -> Any:
+                    return getattr(_ws.get_workspace(), m)
+                return property(functools.update_wrapper(wrap, attr))
+
+            setattr(cls, name, make_wrapper(m=name))
+
+        cls.__doc__ = _ws.Workspace.__doc__
+
+        return super().__new__(cls, *args, **kwargs)
+
+    def __str__(self) -> str:
+        return _ws.get_workspace().__str__()
+
+    def __repr__(self) -> str:
+        return _ws.get_workspace().__repr__()
+
+
+class Organization(object):
+
+    def __new__(cls, *args: Any, **kwargs: Any) -> Any:
+        # We are remapping the methods and attributes here so that
+        # autocomplete still works in Jupyter / IPython, but we
+        # bypass the real method / attribute calls and apply them
+        # to the currently selected organization.
+        for name in [x for x in dir(_ws.Organization) if not x.startswith('_')]:
+            if name in ['from_dict', 'refresh', 'update']:
+                continue
+
+            attr = getattr(_ws.Organization, name)
+
+            def make_wrapper(m: str, is_method: bool = False) -> Any:
+                if is_method:
+                    def wrap(self: Organization, *a: Any, **kw: Any) -> Any:
+                        return getattr(_ws.get_organization(), m)(*a, **kw)
+                    return functools.update_wrapper(wrap, attr)
+                else:
+                    def wrap(self: Organization, *a: Any, **kw: Any) -> Any:
+                        return getattr(_ws.get_organization(), m)
+                    return property(functools.update_wrapper(wrap, attr))
+
+            setattr(cls, name, make_wrapper(m=name, is_method=callable(attr)))
+
+        for name in [
+            x for x in _ws.Organization.__annotations__.keys()
+            if not x.startswith('_')
+        ]:
+
+            def make_wrapper(m: str, is_method: bool = False) -> Any:
+                def wrap(self: Organization) -> Any:
+                    return getattr(_ws.get_organization(), m)
+                return property(functools.update_wrapper(wrap, attr))
+
+            setattr(cls, name, make_wrapper(m=name))
+
+        cls.__doc__ = _ws.Organization.__doc__
+
+        return super().__new__(cls, *args, **kwargs)
+
+    def __str__(self) -> str:
+        return _ws.get_organization().__str__()
+
+    def __repr__(self) -> str:
+        return _ws.get_organization().__repr__()
+
+
+secrets = Secrets()
+stage = Stage()
+organization = Organization()
+workspace_group = WorkspaceGroup()
+workspace = Workspace()
+
+
+__all__ = ['secrets', 'stage', 'workspace', 'workspace_group', 'organization']