singlestoredb 0.9.1__cp36-abi3-win32.whl → 0.9.2__cp36-abi3-win32.whl

singlestoredb/fusion/handlers/workspace.py

@@ -0,0 +1,361 @@
+#!/usr/bin/env python3
+import datetime
+import json
+import os
+from typing import Any
+from typing import Dict
+from typing import Optional
+
+from .. import result
+from ...management import manage_workspaces
+from ...management.workspace import WorkspaceGroup
+from ..handler import SQLHandler
+
+manager = manage_workspaces(os.environ.get('SINGLESTOREDB_MANAGEMENT_TOKEN', 'DEAD'))
+
+
+def dt_isoformat(dt: Optional[datetime.datetime]) -> Optional[str]:
+    """Convert datetime to string."""
+    if dt is None:
+        return None
+    return dt.isoformat()
+
+
+def get_workspace_group(params: Dict[str, Any]) -> WorkspaceGroup:
+    """Find a workspace group matching group_id or group_name."""
+    if 'group_name' in params:
+        workspace_groups = [
+            x for x in manager.workspace_groups
+            if x.name == params['group_name']
+        ]
+        if not workspace_groups:
+            raise KeyError(
+                'no workspace group found with name "{}"'.format(params['group_name']),
+            )
+        if len(workspace_groups) > 1:
+            ids = ', '.join(x.id for x in workspace_groups)
+            raise ValueError(
+                f'more than one workspace group with given name was found: {ids}',
+            )
+        return workspace_groups[0]
+    return manager.get_workspace_group(params['group_id'])
+
+
+class ShowRegionsHandler(SQLHandler):
+    """
+    SHOW REGIONS [ like ];
+
+    # Region name pattern
+    like = LIKE '<pattern>'
+
+    """
+
+    def run(self, params: Dict[str, Any]) -> Optional[result.FusionSQLResult]:
+        res = self.create_result()
+        res.add_field('Name', result.STRING)
+        res.add_field('ID', result.STRING)
+        res.add_field('Provider', result.STRING)
+
+        is_like = self.create_like_func(params.get('like', None))
+
+        res.set_rows([(x.name, x.id, x.provider)
+                      for x in manager.regions if is_like(x.name)])
+        return res
+
+
+ShowRegionsHandler.register()
+
+
+class ShowWorkspaceGroupsHandler(SQLHandler):
+    """
+    SHOW WORKSPACE GROUPS [ like ] [ extended ];
+
+    # Workspace group name pattern
+    like = LIKE '<pattern>'
+
+    # Add additional data columns
+    extended = EXTENDED
+
+    """
+
+    def run(self, params: Dict[str, Any]) -> Optional[result.FusionSQLResult]:
+        res = self.create_result()
+        res.add_field('Name', result.STRING)
+        res.add_field('ID', result.STRING)
+        res.add_field('Region Name', result.STRING)
+        res.add_field('Firewall Ranges', result.JSON)
+
+        if params.get('extended'):
+            res.add_field('Created At', result.DATETIME)
+            res.add_field('Terminated At', result.DATETIME)
+
+            def fields(x: Any) -> Any:
+                return (
+                    x.name, x.id, x.region.name,
+                    json.dumps(x.firewall_ranges),
+                    dt_isoformat(x.created_at),
+                    dt_isoformat(x.terminated_at),
+                )
+        else:
+            def fields(x: Any) -> Any:
+                return (x.name, x.id, x.region.name, x.firewall_ranges)
+
+        is_like = self.create_like_func(params.get('like', None))
+
+        res.set_rows([fields(x) for x in manager.workspace_groups if is_like(x.name)])
+
+        return res
+
+
+ShowWorkspaceGroupsHandler.register()
+
+
+class ShowWorkspacesHandler(SQLHandler):
+    """
+    SHOW WORKSPACES IN GROUP { group_id | group_name } [ like ] [ extended ];
+
+    # ID of group
+    group_id = ID '<group-id>'
+
+    # Name of group
+    group_name = '<group-name>'
+
+    # Workspace group name pattern
+    like = LIKE '<pattern>'
+
+    # Add additional data columns
+    extended = EXTENDED
+
+    """
+
+    def run(self, params: Dict[str, Any]) -> Optional[result.FusionSQLResult]:
+        res = self.create_result()
+        res.add_field('Name', result.STRING)
+        res.add_field('ID', result.STRING)
+        res.add_field('Size', result.STRING)
+        res.add_field('State', result.STRING)
+
+        workspace_group = get_workspace_group(params)
+
+        if params.get('extended'):
+            res.add_field('Endpoint', result.STRING)
+            res.add_field('Created At', result.DATETIME)
+            res.add_field('Terminated At', result.DATETIME)
+
+            def fields(x: Any) -> Any:
+                return (
+                    x.name, x.id, x.size, x.state,
+                    x.endpoint, dt_isoformat(x.created_at),
+                    dt_isoformat(x.terminated_at),
+                )
+        else:
+            def fields(x: Any) -> Any:
+                return (x.name, x.id, x.size, x.state)
+
+        is_like = self.create_like_func(params.get('like', None))
+
+        res.set_rows([fields(x) for x in workspace_group.workspaces if is_like(x.name)])
+
+        return res
+
+
+ShowWorkspacesHandler.register()
+
+
+class CreateWorkspaceGroupHandler(SQLHandler):
+    """
+    CREATE WORKSPACE GROUP [ if_not_exists ] group_name
+        IN REGION { region_id | region_name }
+        [ with_password ]
+        [ expires_at ]
+        [ with_firewall_ranges ]
+    ;
+
+    # Only create workspace group if it doesn't exist already
+    if_not_exists = IF NOT EXISTS
+
+    # Name of the workspace group
+    group_name = '<group-name>'
+
+    # ID of region to create workspace group in
+    region_id = ID '<region-id>'
+
+    # Name of region to create workspace group in
+    region_name = '<region-name>'
+
+    # Admin password
+    with_password = WITH PASSWORD '<password>'
+
+    # Datetime or interval for expiration date/time of workspace group
+    expires_at = EXPIRES AT '<iso-datetime-or-interval>'
+
+    # Incoming IP ranges
+    with_firewall_ranges = WITH FIREWALL RANGES '<ip-range>',...
+
+    """
+
+    def run(self, params: Dict[str, Any]) -> Optional[result.FusionSQLResult]:
+        # Only create if one doesn't exist
+        if params.get('if_not_exists'):
+            try:
+                get_workspace_group(params)
+                return None
+            except (ValueError, KeyError):
+                pass
+
+        # Get region ID
+        if 'region_name' in params:
+            regs = [x for x in manager.regions if x.name == params['region_name']]
+            if not regs:
+                raise ValueError(f'no region found with name "{params["region_name"]}"')
+            if len(regs) > 1:
+                raise ValueError(
+                    f'multiple regions found with the name "{params["region_name"]}"',
+                )
+            region_id = regs[0].id
+        else:
+            region_id = params['region_id']
+
+        manager.create_workspace_group(
+            params['group_name'],
+            region=region_id,
+            admin_password=params.get('with_password'),
+            expires_at=params.get('expires_at'),
+            firewall_ranges=params.get('with_firewall_ranges', []),
+        )
+
+        return self.create_result()
+
+
+CreateWorkspaceGroupHandler.register()
+
+
+class CreateWorkspaceHandler(SQLHandler):
+    """
+    CREATE WORKSPACE [ if_not_exists ] workspace_name
+        IN GROUP { group_id | group_name }
+        WITH SIZE size [ wait_on_active ];
+
+    # Only run command if workspace doesn't already exist
+    if_not_exists = IF NOT EXISTS
+
+    # Name of the workspace
+    workspace_name = '<workspace-name>'
+
+    # ID of the group to create workspace in
+    group_id = ID '<group-id>'
+
+    # Name of the group to create workspace in
+    group_name = '<group-name>'
+
+    # Runtime size
+    size = '<size>'
+
+    # Wait for workspace to be active before continuing
+    wait_on_active = WAIT ON ACTIVE
+
+    """
+
+    def run(self, params: Dict[str, Any]) -> Optional[result.FusionSQLResult]:
+        workspace_group = get_workspace_group(params)
+
+        # Only create if one doesn't exist
+        if params.get('if_not_exists'):
+            try:
+                workspace_group.workspaces[params['workspace_name']]
+                return None
+            except KeyError:
+                pass
+
+        workspace_group.create_workspace(
+            params['workspace_name'], size=params['size'],
+            wait_on_active=params.get('wait_on_active', False),
+        )
+
+        return None
+
+
+CreateWorkspaceHandler.register()
+
+
+class DropWorkspaceGroupHandler(SQLHandler):
+    """
+    DROP WORKSPACE GROUP [ if_exists ] { group_id | group_name }
+        [ wait_on_terminated ];
+
+    # Only run command if the workspace group exists
+    if_exists = IF EXISTS
+
+    # ID of the workspace group to delete
+    group_id = ID '<group-id>'
+
+    # Name of the workspace group to delete
+    group_name = '<group-name>'
+
+    # Wait for termination to complete before continuing
+    wait_on_terminated = WAIT ON TERMINATED
+
+    """
+
+    def run(self, params: Dict[str, Any]) -> Optional[result.FusionSQLResult]:
+        try:
+            name_or_id = params.get('group_name', params.get('group_id'))
+            wg = manager.workspace_groups[name_or_id]
+            wg.terminate(wait_on_terminated=params.get('wait_on_terminated', False))
+
+        except KeyError:
+            if not params.get('if_exists'):
+                raise ValueError(f"could not find workspace group '{name_or_id}'")
+
+        return None
+
+
+DropWorkspaceGroupHandler.register()
+
+
+class DropWorkspaceHandler(SQLHandler):
+    """
+    DROP WORKSPACE [ if_exists ] { workspace_id | workspace_name }
+        IN GROUP { group_id | group_name } [ wait_on_terminated ];
+
+    # Only drop workspace if it exists
+    if_exists = IF EXISTS
+
+    # ID of workspace
+    workspace_id = ID '<workspace-id>'
+
+    # Name of workspace
+    workspace_name = '<workspace-name>'
+
+    # ID of workspace group
+    group_id = ID '<group-id>'
+
+    # Name of workspace group
+    group_name = '<group-name>'
+
+    # Wait for workspace to be terminated before continuing
+    wait_on_terminated = WAIT ON TERMINATED
+
+    """
+
+    def run(self, params: Dict[str, Any]) -> Optional[result.FusionSQLResult]:
+        try:
+            workspace_name_or_id = params.get(
+                'workspace_name', params.get('workspace_id'),
+            )
+            group_name_or_id = params.get('group_name', params.get('group_id'))
+            wg = manager.workspace_groups[group_name_or_id]
+            ws = wg.workspaces[workspace_name_or_id]
+            ws.terminate(wait_on_terminated=params.get('wait_on_terminated', False))
+
+        except KeyError:
+            if not params.get('if_exists'):
+                raise ValueError(
+                    f"could not find workspace '{workspace_name_or_id}' "
+                    f"in group '{group_name_or_id}'",
+                )
+
+        return None
+
+
+DropWorkspaceHandler.register()

singlestoredb/fusion/registry.py

@@ -0,0 +1,112 @@
+#!/usr/bin/env python3
+import os
+import re
+from typing import Any
+from typing import Dict
+from typing import Optional
+from typing import Type
+from typing import Union
+
+from . import result
+from .. import connection
+from .handler import SQLHandler
+
+_enabled = ('1', 'yes', 'on', 'enabled', 'true')
+_handlers: Dict[str, Type[SQLHandler]] = {}
+_handlers_re: Optional[Any] = None
+
+
+def register_handler(handler: Type[SQLHandler], overwrite: bool = False) -> None:
+    """
+    Register a new SQL handler.
+
+    Parameters
+    ----------
+    handler : SQLHandler subclass
+        The handler class to register
+    overwrite : bool, optional
+        Should an existing handler be overwritten if it uses the same command key?
+
+    """
+    global _handlers
+    global _handlers_re
+
+    # Build key for handler
+    key = ' '.join(x.upper() for x in handler.command_key)
+
+    # Check for existing handler with same key
+    if not overwrite and key in _handlers:
+        raise ValueError(f'command already exists, use overwrite=True to override: {key}')
+
+    # Add handler to registry
+    _handlers[key] = handler
+
+    # Build regex to detect fusion query
+    keys = sorted(_handlers.keys(), key=lambda x: -len(x[0]))
+    keys_str = '|'.join(x.replace(' ', '\\s+') for x in keys)
+    _handlers_re = re.compile(f'^\\s*({keys_str})(?:\\s+|;|$)', flags=re.I)
+
+
+def get_handler(sql: Union[str, bytes]) -> Optional[Type[SQLHandler]]:
+    """
+    Return a fusion handler for the given query.
+
+    Parameters
+    ----------
+    sql : str or bytes
+        The SQL query
+
+    Returns
+    -------
+    SQLHandler - if a matching one exists
+    None - if no matching handler could be found
+
+    """
+    if not os.environ.get('SINGLESTOREDB_ENABLE_FUSION', '').lower() in _enabled:
+        return None
+
+    if isinstance(sql, (bytes, bytearray)):
+        sql = sql.decode('utf-8')
+
+    if _handlers_re is None:
+        return None
+
+    m = _handlers_re.match(sql)
+    if m:
+        return _handlers[re.sub(r'\s+', r' ', m.group(1).strip().upper())]
+
+    return None
+
+
+def execute(
+    connection: connection.Connection,
+    sql: str,
+    handler: Optional[Type[SQLHandler]] = None,
+) -> result.FusionSQLResult:
+    """
+    Execute a SQL query in the management interface.
+
+    Parameters
+    ----------
+    connection : Connection
+        The SingleStoreDB connection object
+    sql : str
+        The SQL query
+    handler : SQLHandler, optional
+        The handler to use for the commands. If not supplied, one will be
+        looked up in the registry.
+
+    Returns
+    -------
+    FusionSQLResult
+
+    """
+    if not os.environ.get('SINGLESTOREDB_ENABLE_FUSION', '').lower() in _enabled:
+        raise RuntimeError('management API queries have not been enabled')
+
+    if handler is None:
+        handler = get_handler(sql)
+        if handler is None:
+            raise RuntimeError(f'could not find handler for query: {sql}')
+
+    return handler(connection).execute(sql)

singlestoredb/fusion/result.py

@@ -0,0 +1,120 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+from typing import Any
+from typing import List
+from typing import Optional
+from typing import Tuple
+
+from .. import connection
+from ..mysql.constants.FIELD_TYPE import BLOB  # noqa: F401
+from ..mysql.constants.FIELD_TYPE import BOOL  # noqa: F401
+from ..mysql.constants.FIELD_TYPE import DATE  # noqa: F401
+from ..mysql.constants.FIELD_TYPE import DATETIME  # noqa: F401
+from ..mysql.constants.FIELD_TYPE import DOUBLE  # noqa: F401
+from ..mysql.constants.FIELD_TYPE import JSON  # noqa: F401
+from ..mysql.constants.FIELD_TYPE import LONGLONG as INTEGER  # noqa: F401
+from ..mysql.constants.FIELD_TYPE import STRING  # noqa: F401
+from ..utils.results import Description
+from ..utils.results import format_results
+
+
+class FusionField(object):
+    """Field for PyMySQL compatibility."""
+
+    def __init__(self, name: str, flags: int, charset: int) -> None:
+        self.name = name
+        self.flags = flags
+        self.charsetnr = charset
+
+
+class FusionSQLResult:
+    """Result for Fusion SQL commands."""
+
+    def __init__(self, connection: connection.Connection, unbuffered: bool = False):
+        self.connection: Any = connection
+        self.affected_rows: Optional[int] = None
+        self.insert_id: int = 0
+        self.server_status: Optional[int] = None
+        self.warning_count: int = 0
+        self.message: Optional[str] = None
+        self.description: List[Description] = []
+        self.rows: Any = []
+        self.has_next: bool = False
+        self.unbuffered_active: bool = False
+        self.converters: List[Any] = []
+        self.fields: List[FusionField] = []
+        self._row_idx: int = -1
+
+    def _read_rowdata_packet_unbuffered(self, size: int = 1) -> Optional[List[Any]]:
+        if not self.rows:
+            return None
+
+        out = []
+
+        try:
+            for i in range(1, size + 1):
+                out.append(self.rows[self._row_idx + i])
+        except IndexError:
+            self._row_idx = -1
+            self.rows = []
+            return None
+        else:
+            self._row_idx += size
+
+        return out
+
+    def _finish_unbuffered_query(self) -> None:
+        self._row_idx = -1
+        self.rows = []
+        self.affected_rows = None
+
+    def add_field(self, name: str, dtype: int) -> None:
+        """
+        Add a new field / column to the data set.
+
+        Parameters
+        ----------
+        name : str
+            The name of the field / column
+        dtype : int
+            The MySQL field type: BLOB, BOOL, DATE, DATETIME,
+            DOUBLE, JSON, INTEGER, or STRING
+
+        """
+        charset = 0
+        if dtype in (JSON, STRING):
+            encoding = 'utf-8'
+        elif dtype == BLOB:
+            charset = 63
+            encoding = None
+        else:
+            encoding = 'ascii'
+        self.description.append(
+            Description(name, dtype, None, None, 0, 0, True, 0, charset),
+        )
+        self.fields.append(FusionField(name, 0, charset))
+        converter = self.connection.decoders.get(dtype)
+        self.converters.append((encoding, converter))
+
+    def set_rows(self, data: List[Tuple[Any, ...]]) -> None:
+        """
+        Set the rows of the result.
+
+        Parameters
+        ----------
+        data : List[Tuple[Any, ...]]
+            The data should be a list of tuples where each element of the
+            tuple corresponds to a field added to the result with
+            the :meth:`add_field` method.
+
+        """
+        # Convert values
+        for i, row in enumerate(list(data)):
+            new_row = []
+            for (_, converter), value in zip(self.converters, row):
+                new_row.append(converter(value) if converter is not None else value)
+            data[i] = tuple(new_row)
+
+        self.rows = format_results(self.connection._results_type, self.description, data)
+        self.affected_rows = 0

singlestoredb/http/connection.py

@@ -41,6 +41,7 @@ except ImportError:
     has_shapely = False
 
 from .. import connection
+from .. import fusion
 from .. import types
 from ..config import get_option
 from ..converters import converters
@@ -56,6 +57,7 @@ from ..exceptions import ProgrammingError
 from ..exceptions import Warning  # noqa: F401
 from ..utils.convert_rows import convert_rows
 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 Result
@@ -431,16 +433,79 @@ class Cursor(connection.Cursor):
             else:
                 query = query % args
 
+    def _execute_fusion_query(
+        self,
+        oper: Union[str, bytes],
+        params: Optional[Union[Sequence[Any], Dict[str, Any]]] = None,
+        handler: Any = None,
+    ) -> int:
+        oper = mogrify(oper, params)
+
+        if isinstance(oper, bytes):
+            oper = oper.decode('utf-8')
+
+        log_query(oper, None)
+
+        results_type = self._results_type
+        self._results_type = 'tuples'
+        try:
+            mgmt_res = fusion.execute(
+                self._connection,  # type: ignore
+                oper,
+                handler=handler,
+            )
+        finally:
+            self._results_type = results_type
+
+        self._descriptions.append(list(mgmt_res.description))
+        self._results.append(list(mgmt_res.rows))
+        self.rowcount = len(self._results[-1])
+
+        pymy_res = PyMyResult()
+        for field in mgmt_res.fields:
+            pymy_res.append(
+                PyMyField(
+                    field.name,
+                    field.flags,
+                    field.charsetnr,
+                ),
+            )
+
+        self._pymy_results.append(pymy_res)
+
+        if self._results and self._results[0]:
+            self._row_idx = 0
+            self._result_idx = 0
+
+        return self.rowcount
+
     def _execute(
         self, oper: str,
         params: Optional[Union[Sequence[Any], Dict[str, Any]]] = None,
         is_callproc: bool = False,
     ) -> int:
+        self._descriptions = []
+        self._results = []
+        self._pymy_results = []
+        self._row_idx = -1
+        self._result_idx = -1
+        self.rowcount = 0
+        self._expect_results = False
+
         if self._connection is None:
             raise ProgrammingError(errno=2048, msg='Connection is closed.')
 
+        sql_type = 'exec'
+        if re.match(r'^\s*(select|show|call|echo|describe|with)\s+', oper, flags=re.I):
+            self._expect_results = True
+            sql_type = 'query'
+
         self._validate_param_subs(oper, params)
 
+        handler = fusion.get_handler(oper)
+        if handler is not None:
+            return self._execute_fusion_query(oper, params, handler=handler)
+
         oper, params = self._connection._convert_params(oper, params)
 
         log_query(oper, params)
@@ -455,12 +520,6 @@ class Cursor(connection.Cursor):
         if self._connection._database:
             data['database'] = self._connection._database
 
-        self._expect_results = False
-        sql_type = 'exec'
-        if re.match(r'^\s*(select|show|call|echo|describe|with)\s+', oper, flags=re.I):
-            self._expect_results = True
-            sql_type = 'query'
-
         if sql_type == 'query':
             res = self._post('query/tuples', json=data)
         else:
@@ -479,12 +538,6 @@ class Cursor(connection.Cursor):
 
         out = json.loads(res.text)
 
-        self._descriptions = []
-        self._results = []
-        self._row_idx = -1
-        self._result_idx = -1
-        self.rowcount = 0
-
         if sql_type == 'query':
             # description: (name, type_code, display_size, internal_size,
             #               precision, scale, null_ok, column_flags, charset)