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

singlestoredb/management/__init__.py

@@ -3,4 +3,5 @@ from .cluster import manage_cluster
 from .manager import get_organization
 from .manager import get_token
 from .workspace import get_secret
+from .workspace import get_stage
 from .workspace import manage_workspaces

singlestoredb/management/organization.py

@@ -121,6 +121,10 @@ class Organization(object):
 
     """
 
+    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

singlestoredb/management/utils.py

@@ -294,7 +294,7 @@ def snake_to_camel_dict(
         return None
     out = {}
     for k, v in s.items():
-        if isinstance(s, Mapping):
+        if isinstance(v, Mapping):
             out[str(snake_to_camel(k))] = snake_to_camel_dict(v, cap_first=cap_first)
         else:
             out[str(snake_to_camel(k))] = v
@@ -307,7 +307,7 @@ def camel_to_snake_dict(s: Optional[Mapping[str, Any]]) -> Optional[Dict[str, An
         return None
     out = {}
     for k, v in s.items():
-        if isinstance(s, Mapping):
+        if isinstance(v, Mapping):
             out[str(camel_to_snake(k))] = camel_to_snake_dict(v)
         else:
             out[str(camel_to_snake(k))] = v

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

@@ -5,12 +5,15 @@
 # https://dev.mysql.com/doc/refman/5.5/en/error-handling.html
 import errno
 import functools
+import io
 import os
+import queue
 import socket
 import struct
 import sys
 import traceback
 import warnings
+from typing import Iterable
 
 try:
     import _singlestoredb_accel
@@ -259,6 +262,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 +333,7 @@ class Connection(BaseConnection):
         inf_as_null=None,
         encoding_errors='strict',
         track_env=False,
+        enable_extended_data_types=True,
     ):
         BaseConnection.__init__(**dict(locals()))
 
@@ -350,6 +356,7 @@ class Connection(BaseConnection):
             )
 
         self._local_infile = bool(local_infile)
+        self._local_infile_stream = None
         if self._local_infile:
             client_flag |= CLIENT.LOCAL_FILES
         if multi_statements:
@@ -609,6 +616,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
@@ -839,7 +847,7 @@ class Connection(BaseConnection):
         return self.cursorclass(self)
 
     # The following methods are INTERNAL USE ONLY (called from Cursor)
-    def query(self, sql, unbuffered=False):
+    def query(self, sql, unbuffered=False, infile_stream=None):
         """
         Run a query on the server.
 
@@ -855,8 +863,10 @@ class Connection(BaseConnection):
         else:
             if isinstance(sql, str):
                 sql = sql.encode(self.encoding, 'surrogateescape')
+            self._local_infile_stream = infile_stream
             self._execute_command(COMMAND.COM_QUERY, sql)
             self._affected_rows = self._read_query_result(unbuffered=unbuffered)
+            self._local_infile_stream = None
         return self._affected_rows
 
     def next_result(self, unbuffered=False):
@@ -1070,6 +1080,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)
@@ -1859,24 +1877,82 @@ class LoadLocalFile:
     def send_data(self):
         """Send data packets from the local file to the server"""
         if not self.connection._sock:
-            raise err.InterfaceError(0, '')
+            raise err.InterfaceError(0, 'Connection is closed')
+
         conn = self.connection
+        infile = conn._local_infile_stream
+
+        # 16KB is efficient enough
+        packet_size = min(conn.max_allowed_packet, 16 * 1024)
 
         try:
-            with open(self.filename, 'rb') as open_file:
-                packet_size = min(
-                    conn.max_allowed_packet, 16 * 1024,
-                )  # 16KB is efficient enough
-                while True:
-                    chunk = open_file.read(packet_size)
-                    if not chunk:
-                        break
-                    conn.write_packet(chunk)
-        except OSError:
-            raise err.OperationalError(
-                ER.FILE_NOT_FOUND,
-                f"Can't find file '{self.filename}'",
-            )
+
+            if self.filename in [':stream:', b':stream:']:
+
+                if infile is None:
+                    raise err.OperationalError(
+                        ER.FILE_NOT_FOUND,
+                        ':stream: specified for LOCAL INFILE, but no stream was supplied',
+                    )
+
+                # Binary IO
+                elif isinstance(infile, io.RawIOBase):
+                    while True:
+                        chunk = infile.read(packet_size)
+                        if not chunk:
+                            break
+                        conn.write_packet(chunk)
+
+                # Text IO
+                elif isinstance(infile, io.TextIOBase):
+                    while True:
+                        chunk = infile.read(packet_size)
+                        if not chunk:
+                            break
+                        conn.write_packet(chunk.encode('utf8'))
+
+                # Iterable of bytes or str
+                elif isinstance(infile, Iterable):
+                    for chunk in infile:
+                        if not chunk:
+                            continue
+                        if isinstance(chunk, str):
+                            conn.write_packet(chunk.encode('utf8'))
+                        else:
+                            conn.write_packet(chunk)
+
+                # Queue (empty value ends the iteration)
+                elif isinstance(infile, queue.Queue):
+                    while True:
+                        chunk = infile.get()
+                        if not chunk:
+                            break
+                        if isinstance(chunk, str):
+                            conn.write_packet(chunk.encode('utf8'))
+                        else:
+                            conn.write_packet(chunk)
+
+                else:
+                    raise err.OperationalError(
+                        ER.FILE_NOT_FOUND,
+                        ':stream: specified for LOCAL INFILE, ' +
+                        f'but stream type is unrecognized: {infile}',
+                    )
+
+            else:
+                try:
+                    with open(self.filename, 'rb') as open_file:
+                        while True:
+                            chunk = open_file.read(packet_size)
+                            if not chunk:
+                                break
+                            conn.write_packet(chunk)
+                except OSError:
+                    raise err.OperationalError(
+                        ER.FILE_NOT_FOUND,
+                        f"Can't find file '{self.filename!s}'",
+                    )
+
         finally:
             if not conn._closed:
                 # send the empty packet to signify we are done sending data

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

@@ -178,7 +178,7 @@ class Cursor(BaseCursor):
 
         return query
 
-    def execute(self, query, args=None):
+    def execute(self, query, args=None, infile_stream=None):
         """
         Execute a query.
 
@@ -192,6 +192,8 @@ class Cursor(BaseCursor):
             Query to execute.
         args : Sequence[Any] or Dict[str, Any] or Any, optional
             Parameters used with query. (optional)
+        infile_stream : io.BytesIO or Iterator[bytes], optional
+            Data stream for ``LOCAL INFILE`` statements
 
         Returns
         -------
@@ -205,7 +207,7 @@ class Cursor(BaseCursor):
 
         query = self.mogrify(query, args)
 
-        result = self._query(query)
+        result = self._query(query, infile_stream=infile_stream)
         self._executed = query
         return result
 
@@ -387,10 +389,10 @@ class Cursor(BaseCursor):
             raise IndexError('out of range')
         self._rownumber = r
 
-    def _query(self, q):
+    def _query(self, q, infile_stream=None):
         conn = self._get_db()
         self._clear_result()
-        conn.query(q)
+        conn.query(q, infile_stream=infile_stream)
         self._do_get_result()
         return self.rowcount
 
@@ -417,10 +419,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
 
@@ -679,10 +682,10 @@ class SSCursor(Cursor):
 
     __del__ = close
 
-    def _query(self, q):
+    def _query(self, q, infile_stream=None):
         conn = self._get_db()
         self._clear_result()
-        conn.query(q, unbuffered=True)
+        conn.query(q, unbuffered=True, infile_stream=infile_stream)
         self._do_get_result()
         return self.rowcount
 

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