singlestoredb 1.1.0__cp38-abi3-win32.whl → 1.2.0__cp38-abi3-win32.whl

singlestoredb/management/workspace.py

@@ -34,9 +34,53 @@ from .utils import ttl_property
 from .utils import vars_to_str
 
 
+def get_organization() -> Organization:
+    """Get the organization."""
+    return manage_workspaces().organization
+
+
 def get_secret(name: str) -> str:
     """Get a secret from the organization."""
-    return manage_workspaces().organization.get_secret(name).value
+    return get_organization().get_secret(name).value
+
+
+def get_workspace_group(
+    workspace_group: Optional[Union[WorkspaceGroup, str]] = None,
+) -> WorkspaceGroup:
+    """Get the stage for the workspace group."""
+    if isinstance(workspace_group, WorkspaceGroup):
+        return workspace_group
+    elif workspace_group:
+        return manage_workspaces().workspace_groups[workspace_group]
+    elif 'SINGLESTOREDB_WORKSPACE_GROUP' in os.environ:
+        return manage_workspaces().workspace_groups[
+            os.environ['SINGLESTOREDB_WORKSPACE_GROUP']
+        ]
+    raise RuntimeError('no workspace group specified')
+
+
+def get_stage(
+    workspace_group: Optional[Union[WorkspaceGroup, str]] = None,
+) -> Stage:
+    """Get the stage for the workspace group."""
+    return get_workspace_group(workspace_group).stage
+
+
+def get_workspace(
+    workspace_group: Optional[Union[WorkspaceGroup, str]] = None,
+    workspace: Optional[Union[Workspace, str]] = None,
+) -> Workspace:
+    """Get the workspaces for a workspace_group."""
+    if isinstance(workspace, Workspace):
+        return workspace
+    wg = get_workspace_group(workspace_group)
+    if workspace:
+        return wg.workspaces[workspace]
+    elif 'SINGLESTOREDB_WORKSPACE' in os.environ:
+        return wg.workspaces[
+            os.environ['SINGLESTOREDB_WORKSPACE']
+        ]
+    raise RuntimeError('no workspace group specified')
 
 
 class StageObject(object):
@@ -926,6 +970,21 @@ class Workspace(object):
 
     """
 
+    name: str
+    id: str
+    group_id: str
+    size: str
+    state: str
+    created_at: Optional[datetime.datetime]
+    terminated_at: Optional[datetime.datetime]
+    endpoint: Optional[str]
+    auto_suspend: Optional[Dict[str, Any]]
+    cache_config: Optional[int]
+    deployment_type: Optional[str]
+    resume_attachments: Optional[List[Dict[str, Any]]]
+    scaling_progress: Optional[int]
+    last_resumed_at: Optional[datetime.datetime]
+
     def __init__(
         self,
         name: str,
@@ -939,9 +998,9 @@ class Workspace(object):
         auto_suspend: Optional[Dict[str, Any]] = None,
         cache_config: Optional[int] = None,
         deployment_type: Optional[str] = None,
-        resume_attachments: Optional[Dict[str, Any]] = None,
+        resume_attachments: Optional[List[Dict[str, Any]]] = None,
         scaling_progress: Optional[int] = None,
-        last_resumed_at: Optional[str] = None,
+        last_resumed_at: Optional[Union[str, datetime.datetime]] = None,
     ):
         #: Name of the workspace
         self.name = name
@@ -981,7 +1040,11 @@ class Workspace(object):
         self.deployment_type = deployment_type
 
         #: Database attachments
-        self.resume_attachments = camel_to_snake_dict(resume_attachments)
+        self.resume_attachments = [
+            camel_to_snake_dict(x)  # type: ignore
+            for x in resume_attachments or []
+            if x is not None
+        ]
 
         #: Current progress percentage for scaling the workspace
         self.scaling_progress = scaling_progress
@@ -1248,8 +1311,18 @@ class WorkspaceGroup(object):
 
     """
 
+    name: str
+    id: str
+    created_at: Optional[datetime.datetime]
+    region: Optional[Region]
+    firewall_ranges: List[str]
+    terminated_at: Optional[datetime.datetime]
+    allow_all_traffic: bool
+
     def __init__(
-        self, name: str, id: str,
+        self,
+        name: str,
+        id: str,
         created_at: Union[str, datetime.datetime],
         region: Optional[Region],
         firewall_ranges: List[str],
@@ -1275,7 +1348,7 @@ class WorkspaceGroup(object):
         self.terminated_at = to_datetime(terminated_at)
 
         #: Should all traffic be allowed?
-        self.allow_all_traffic = allow_all_traffic
+        self.allow_all_traffic = allow_all_traffic or False
 
         self._manager: Optional[WorkspaceManager] = None
 

singlestoredb/mysql/connection.py

@@ -259,6 +259,8 @@ class Connection(BaseConnection):
         uploading data?
     track_env : bool, optional
         Should the connection track the SINGLESTOREDB_URL environment variable?
+    enable_extended_data_types : bool, optional
+        Should extended data types (BSON, vector) be enabled?
 
     See `Connection <https://www.python.org/dev/peps/pep-0249/#connection-objects>`_
     in the specification.
@@ -328,6 +330,7 @@ class Connection(BaseConnection):
         inf_as_null=None,
         encoding_errors='strict',
         track_env=False,
+        enable_extended_data_types=True,
     ):
         BaseConnection.__init__(**dict(locals()))
 
@@ -609,6 +612,7 @@ class Connection(BaseConnection):
 
         self._in_sync = False
         self._track_env = bool(track_env) or self.host == 'singlestore.com'
+        self._enable_extended_data_types = enable_extended_data_types
 
         if defer_connect or self._track_env:
             self._sock = None
@@ -1070,6 +1074,14 @@ class Connection(BaseConnection):
                 c.execute('SET sql_mode=%s', (self.sql_mode,))
                 c.close()
 
+            if self._enable_extended_data_types:
+                c = self.cursor()
+                try:
+                    c.execute('SET @@SESSION.enable_extended_types_metadata=on')
+                except self.OperationalError:
+                    pass
+                c.close()
+
             if self.init_command is not None:
                 c = self.cursor()
                 c.execute(self.init_command)

singlestoredb/mysql/constants/EXTENDED_TYPE.py

@@ -0,0 +1,3 @@
+NONE = 0
+BSON = 1
+VECTOR = 2

singlestoredb/mysql/constants/FIELD_TYPE.py

@@ -30,3 +30,19 @@ GEOMETRY = 255
 CHAR = TINY
 INTERVAL = ENUM
 BOOL = TINY
+
+# SingleStoreDB-specific.
+# Only enabled when enable_extended_types_metadata=1 in the server.
+BSON = 1001
+FLOAT32_VECTOR_JSON = 2001
+FLOAT64_VECTOR_JSON = 2002
+INT8_VECTOR_JSON = 2003
+INT16_VECTOR_JSON = 2004
+INT32_VECTOR_JSON = 2005
+INT64_VECTOR_JSON = 2006
+FLOAT32_VECTOR = 3001
+FLOAT64_VECTOR = 3002
+INT8_VECTOR = 3003
+INT16_VECTOR = 3004
+INT32_VECTOR = 3005
+INT64_VECTOR = 3006

singlestoredb/mysql/constants/VECTOR_TYPE.py

@@ -0,0 +1,6 @@
+FLOAT32 = 1
+FLOAT64 = 2
+INT8 = 3
+INT16 = 4
+INT32 = 5
+INT64 = 6

singlestoredb/mysql/cursors.py

@@ -417,10 +417,11 @@ class Cursor(BaseCursor):
         if self.rowcount == 18446744073709551615:
             self.rowcount = -1
         self._description = result.description
-        self._format_schema = get_schema(
-            self.connection._results_type,
-            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
 

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']

singlestoredb/tests/test.sql

@@ -606,4 +606,53 @@ INSERT INTO `badutf8` VALUES ('🥷🧙👻.eth');
 INSERT INTO `badutf8` VALUES ('🥒rick.eth');
 
 
+CREATE TABLE IF NOT EXISTS `f32_vectors` (
+    id INT(11),
+    a VECTOR(3)
+);
+INSERT INTO f32_vectors VALUES(1, '[0.267261237,0.534522474,0.801783681]');
+INSERT INTO f32_vectors VALUES(2, '[0.371390671,0.557085991,0.742781341]');
+INSERT INTO f32_vectors VALUES(3, '[-0.424264073,-0.565685451,0.707106829]');
+
+CREATE TABLE IF NOT EXISTS `f64_vectors` (
+    id INT(11),
+    a VECTOR(3, F64)
+);
+INSERT INTO f64_vectors VALUES(1, '[0.267261237,0.534522474,0.801783681]');
+INSERT INTO f64_vectors VALUES(2, '[0.371390671,0.557085991,0.742781341]');
+INSERT INTO f64_vectors VALUES(3, '[-0.424264073,-0.565685451,0.707106829]');
+
+CREATE TABLE `i8_vectors` (
+    id INT(11),
+    a VECTOR(3, I8)
+);
+INSERT INTO i8_vectors VALUES(1, '[1, 2, 3]');
+INSERT INTO i8_vectors VALUES(2, '[4, 5, 6]');
+INSERT INTO i8_vectors VALUES(3, '[-1, -4, 8]');
+
+CREATE TABLE `i16_vectors` (
+    id INT(11),
+    a VECTOR(3, I16)
+);
+INSERT INTO i16_vectors VALUES(1, '[1, 2, 3]');
+INSERT INTO i16_vectors VALUES(2, '[4, 5, 6]');
+INSERT INTO i16_vectors VALUES(3, '[-1, -4, 8]');
+
+CREATE TABLE `i32_vectors` (
+    id INT(11),
+    a VECTOR(3, I32)
+);
+INSERT INTO i32_vectors VALUES(1, '[1, 2, 3]');
+INSERT INTO i32_vectors VALUES(2, '[4, 5, 6]');
+INSERT INTO i32_vectors VALUES(3, '[-1, -4, 8]');
+
+CREATE TABLE `i64_vectors` (
+    id INT(11),
+    a VECTOR(3, I64)
+);
+INSERT INTO i64_vectors VALUES(1, '[1, 2, 3]');
+INSERT INTO i64_vectors VALUES(2, '[4, 5, 6]');
+INSERT INTO i64_vectors VALUES(3, '[-1, -4, 8]');
+
+
 COMMIT;