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

singlestoredb/http/connection.py

@@ -62,6 +62,7 @@ from ..utils.debug import log_query
 from ..utils.mogrify import mogrify
 from ..utils.results import Description
 from ..utils.results import format_results
+from ..utils.results import get_schema
 from ..utils.results import Result
 
 
@@ -333,6 +334,7 @@ class Cursor(connection.Cursor):
         self._row_idx: int = -1
         self._result_idx: int = -1
         self._descriptions: List[List[Description]] = []
+        self._schemas: List[Dict[str, Any]] = []
         self.arraysize: int = get_option('results.arraysize')
         self.rowcount: int = 0
         self.lastrowid: Optional[int] = None
@@ -355,6 +357,14 @@ class Cursor(connection.Cursor):
             return self._descriptions[self._result_idx]
         return None
 
+    @property
+    def _schema(self) -> Optional[Any]:
+        if not self._schemas:
+            return None
+        if self._result_idx >= 0 and self._result_idx < len(self._schemas):
+            return self._schemas[self._result_idx]
+        return None
+
     def _post(self, path: str, *args: Any, **kwargs: Any) -> requests.Response:
         """
         Invoke a POST request on the HTTP connection.
@@ -460,6 +470,7 @@ class Cursor(connection.Cursor):
             self._results_type = results_type
 
         self._descriptions.append(list(mgmt_res.description))
+        self._schemas.append(get_schema(self._results_type, list(mgmt_res.description)))
         self._results.append(list(mgmt_res.rows))
         self.rowcount = len(self._results[-1])
 
@@ -487,6 +498,7 @@ class Cursor(connection.Cursor):
         is_callproc: bool = False,
     ) -> int:
         self._descriptions = []
+        self._schemas = []
         self._results = []
         self._pymy_results = []
         self._row_idx = -1
@@ -571,6 +583,20 @@ class Cursor(connection.Cursor):
                     if isinstance(k, int):
                         http_converters[k] = v
 
+            # Make JSON a string for Arrow
+            if 'arrow' in self._results_type:
+                def json_to_str(x: Any) -> Optional[str]:
+                    if x is None:
+                        return None
+                    return json.dumps(x)
+                http_converters[245] = json_to_str
+
+            # Don't convert date/times in polars
+            elif 'polars' in self._results_type:
+                http_converters.pop(7, None)
+                http_converters.pop(10, None)
+                http_converters.pop(12, None)
+
             results = out['results']
 
             # Convert data to Python types
@@ -616,6 +642,7 @@ class Cursor(connection.Cursor):
                         )
                         pymy_res.append(PyMyField(col['name'], flags, charset))
                     self._descriptions.append(description)
+                    self._schemas.append(get_schema(self._results_type, description))
 
                     rows = convert_rows(result.get('rows', []), convs)
 
@@ -659,6 +686,7 @@ class Cursor(connection.Cursor):
         rowcount = 0
         if args is not None and len(args) > 0:
             description = []
+            schema = {}
             # Detect dataframes
             if hasattr(args, 'itertuples'):
                 argiter = args.itertuples(index=False)  # type: ignore
@@ -668,11 +696,14 @@ class Cursor(connection.Cursor):
                 self.execute(query, params)
                 if self._descriptions:
                     description = self._descriptions[-1]
+                if self._schemas:
+                    schema = self._schemas[-1]
                 if self._rows is not None:
                     results.append(self._rows)
                 rowcount += self.rowcount
             self._results = results
             self._descriptions = [description for _ in range(len(results))]
+            self._schemas = [schema for _ in range(len(results))]
         else:
             self.execute(query)
             rowcount += self.rowcount
@@ -721,6 +752,7 @@ class Cursor(connection.Cursor):
             self._results_type,
             self.description or [],
             out, single=True,
+            schema=self._schema,
         )
 
     def fetchmany(
@@ -752,7 +784,10 @@ class Cursor(connection.Cursor):
             size = max(int(size), 1)
         out = self._rows[self._row_idx:self._row_idx+size]
         self._row_idx += len(out)
-        return format_results(self._results_type, self.description or [], out)
+        return format_results(
+            self._results_type, self.description or [],
+            out, schema=self._schema,
+        )
 
     def fetchall(self) -> Result:
         """
@@ -774,7 +809,10 @@ class Cursor(connection.Cursor):
             return tuple()
         out = list(self._rows[self._row_idx:])
         self._row_idx = len(out)
-        return format_results(self._results_type, self.description or [], out)
+        return format_results(
+            self._results_type, self.description or [],
+            out, schema=self._schema,
+        )
 
     def nextset(self) -> Optional[bool]:
         """Skip to the next available result set."""

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

@@ -29,12 +29,28 @@ from .cursors import (
     DictCursorSV,
     NamedtupleCursor,
     NamedtupleCursorSV,
+    ArrowCursor,
+    ArrowCursorSV,
+    NumpyCursor,
+    NumpyCursorSV,
+    PandasCursor,
+    PandasCursorSV,
+    PolarsCursor,
+    PolarsCursorSV,
     SSCursor,
     SSCursorSV,
     SSDictCursor,
     SSDictCursorSV,
     SSNamedtupleCursor,
     SSNamedtupleCursorSV,
+    SSArrowCursor,
+    SSArrowCursorSV,
+    SSNumpyCursor,
+    SSNumpyCursorSV,
+    SSPandasCursor,
+    SSPandasCursorSV,
+    SSPolarsCursor,
+    SSPolarsCursorSV,
 )
 from .optionfile import Parser
 from .protocol import (
@@ -243,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.
@@ -312,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()))
 
@@ -443,6 +462,14 @@ class Connection(BaseConnection):
                 self.cursorclass = DictCursor
             elif 'namedtuple' in self.results_type:
                 self.cursorclass = NamedtupleCursor
+            elif 'numpy' in self.results_type:
+                self.cursorclass = NumpyCursor
+            elif 'arrow' in self.results_type:
+                self.cursorclass = ArrowCursor
+            elif 'pandas' in self.results_type:
+                self.cursorclass = PandasCursor
+            elif 'polars' in self.results_type:
+                self.cursorclass = PolarsCursor
             else:
                 self.cursorclass = Cursor
         else:
@@ -450,6 +477,14 @@ class Connection(BaseConnection):
                 self.cursorclass = SSDictCursor
             elif 'namedtuple' in self.results_type:
                 self.cursorclass = SSNamedtupleCursor
+            elif 'numpy' in self.results_type:
+                self.cursorclass = SSNumpyCursor
+            elif 'arrow' in self.results_type:
+                self.cursorclass = SSArrowCursor
+            elif 'pandas' in self.results_type:
+                self.cursorclass = SSPandasCursor
+            elif 'polars' in self.results_type:
+                self.cursorclass = SSPolarsCursor
             else:
                 self.cursorclass = SSCursor
 
@@ -487,6 +522,30 @@ class Connection(BaseConnection):
             elif self.cursorclass is SSNamedtupleCursor:
                 self.cursorclass = SSNamedtupleCursorSV
                 self.results_type = 'namedtuples'
+            elif self.cursorclass is NumpyCursor:
+                self.cursorclass = NumpyCursorSV
+                self.results_type = 'numpy'
+            elif self.cursorclass is SSNumpyCursor:
+                self.cursorclass = SSNumpyCursorSV
+                self.results_type = 'numpy'
+            elif self.cursorclass is ArrowCursor:
+                self.cursorclass = ArrowCursorSV
+                self.results_type = 'arrow'
+            elif self.cursorclass is SSArrowCursor:
+                self.cursorclass = SSArrowCursorSV
+                self.results_type = 'arrow'
+            elif self.cursorclass is PandasCursor:
+                self.cursorclass = PandasCursorSV
+                self.results_type = 'pandas'
+            elif self.cursorclass is SSPandasCursor:
+                self.cursorclass = SSPandasCursorSV
+                self.results_type = 'pandas'
+            elif self.cursorclass is PolarsCursor:
+                self.cursorclass = PolarsCursorSV
+                self.results_type = 'polars'
+            elif self.cursorclass is SSPolarsCursor:
+                self.cursorclass = SSPolarsCursorSV
+                self.results_type = 'polars'
 
         self._result = None
         self._affected_rows = 0
@@ -498,9 +557,22 @@ class Connection(BaseConnection):
         if conv is None:
             conv = converters.conversions
 
+        conv = conv.copy()
+
         self.parse_json = parse_json
         self.invalid_values = (invalid_values or {}).copy()
 
+        # Disable JSON parsing for Arrow
+        if self.results_type in ['arrow']:
+            conv[245] = None
+            self.parse_json = False
+
+        # Disable date/time parsing for polars; let polars do the parsing
+        elif self.results_type in ['polars']:
+            conv[7] = None
+            conv[10] = None
+            conv[12] = None
+
         # Need for MySQLdb compatibility.
         self.encoders = {k: v for (k, v) in conv.items() if type(k) is not int}
         self.decoders = {k: v for (k, v) in conv.items() if type(k) is int}
@@ -540,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
@@ -1001,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