Add medium int range checks; add organization field to WorkspaceManager

singlestore-labs · Jan 30, 2024 · ecd2435 · ecd2435
1 parent 73a7596
commit ecd2435
Show file tree

Hide file tree

Showing 4 changed files with 46 additions and 13 deletions.
diff --git a/docs/src/api.rst b/docs/src/api.rst
@@ -171,6 +171,7 @@ create new ones.
    :toctree: generated/
 
    WorkspaceManager
+   WorkspaceManager.organization
    WorkspaceManager.workspace_groups
    WorkspaceManager.regions
    WorkspaceManager.create_workspace_group

diff --git a/singlestoredb/functions/ext/rowdat_1.py b/singlestoredb/functions/ext/rowdat_1.py
@@ -82,6 +82,7 @@
     ft.FLOAT: 4,
     ft.DOUBLE: 8,
 }
+medium_int_types = set([ft.INT24, -ft.INT24])
 int_types = set([
     ft.TINY, -ft.TINY, ft.SHORT, -ft.SHORT, ft.INT24, -ft.INT24,
     ft.LONG, -ft.LONG, ft.LONGLONG, -ft.LONGLONG,
@@ -369,6 +370,16 @@ def _dump(
                     out.write(struct.pack(numeric_formats[rtype], default))
                 else:
                     if rtype in int_types:
+                        if rtype == ft.INT24:
+                            if int(value) > 8388607 or int(value) < -8388608:
+                                raise ValueError(
+                                    'value is outside range of MEDIUMINT',
+                                )
+                        elif rtype == -ft.INT24:
+                            if int(value) > 16777215 or int(value) < 0:
+                                raise ValueError(
+                                    'value is outside range of UNSIGNED MEDIUMINT',
+                                )
                         out.write(struct.pack(numeric_formats[rtype], int(value)))
                     else:
                         out.write(struct.pack(numeric_formats[rtype], float(value)))
@@ -437,6 +448,16 @@ def _dump_vectors(
                         out.write(struct.pack(numeric_formats[rtype], default))
                     else:
                         if rtype in int_types:
+                            if rtype == ft.INT24:
+                                if int(value) > 8388607 or int(value) < -8388608:
+                                    raise ValueError(
+                                        'value is outside range of MEDIUMINT',
+                                    )
+                            elif rtype == -ft.INT24:
+                                if int(value) > 16777215 or int(value) < 0:
+                                    raise ValueError(
+                                        'value is outside range of UNSIGNED MEDIUMINT',
+                                    )
                             out.write(struct.pack(numeric_formats[rtype], int(value)))
                         else:
                             out.write(struct.pack(numeric_formats[rtype], float(value)))
@@ -697,18 +718,16 @@ def _dump_arrow_accel(
 
 
 if not has_accel:
-    _load_accel = _load
-    _dump_accel = _dump
-    load = _load
-    dump = _dump
-    load_pandas = _load_pandas
-    dump_pandas = _dump_pandas
-    load_numpy = _load_numpy
-    dump_numpy = _dump_numpy
-    load_arrow = _load_arrow
-    dump_arrow = _dump_arrow
-    load_polars = _load_polars
-    dump_polars = _dump_polars
+    load = _load_accel = _load
+    dump = _dump_accel = _dump
+    load_pandas = _load_pandas_accel = _load_pandas  # noqa: F811
+    dump_pandas = _dump_pandas_accel = _dump_pandas  # noqa: F811
+    load_numpy = _load_numpy_accel = _load_numpy  # noqa: F811
+    dump_numpy = _dump_numpy_accel = _dump_numpy  # noqa: F811
+    load_arrow = _load_arrow_accel = _load_arrow  # noqa: F811
+    dump_arrow = _dump_arrow_accel = _dump_arrow  # noqa: F811
+    load_polars = _load_polars_accel = _load_polars  # noqa: F811
+    dump_polars = _dump_polars_accel = _dump_polars  # noqa: F811
 
 else:
     _load_accel = _singlestoredb_accel.load_rowdat_1

diff --git a/singlestoredb/management/workspace.py b/singlestoredb/management/workspace.py
@@ -1222,6 +1222,14 @@ def from_dict(
         out._manager = manager
         return out
 
+    @property
+    def organization(self) -> Organization:
+        if self._manager is None:
+            raise ManagementError(
+                msg='No workspace manager is associated with this object.',
+            )
+        return self._manager.organization
+
     @property
     def stage(self) -> Stage:
         """Stage manager."""
@@ -1479,6 +1487,11 @@ def organizations(self) -> Organizations:
         """Return the organizations."""
         return Organizations(self)
 
+    @property
+    def organization(self) -> Organization:
+        """ Return the current organization."""
+        return self.organizations.current
+
     @property
     def billing(self) -> Billing:
         """Return the current billing information."""

diff --git a/singlestoredb/tests/test_ext_func_data.py b/singlestoredb/tests/test_ext_func_data.py
@@ -240,7 +240,7 @@
     ],
     [
         100, 100, 32700, 32800, 2147483600, 2147483800, 100.0,
-        100.0, 100, 100, 2147483600, 2147483800, 'bye', b'bye',
+        100.0, 100, 100, 8388600, 16777200, 'bye', b'bye',
     ],
     [
         120, 130, 254, 254, 254, 254, 3.14159, 3.14159,