singlestoredb 1.16.1__py3-none-any.whl

singlestoredb/fusion/registry.py

@@ -0,0 +1,249 @@
+#!/usr/bin/env python3
+import re
+import sys
+from typing import Any
+from typing import Dict
+from typing import List
+from typing import Optional
+from typing import Tuple
+from typing import Type
+from typing import Union
+
+from . import result
+from .. import connection
+from ..config import get_option
+from .handler import SQLHandler
+
+_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), x))
+    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 get_option('fusion.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 get_option('fusion.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)
+
+
+class ShowFusionCommandsHandler(SQLHandler):
+    """
+    SHOW FUSION COMMANDS [ like ];
+
+    # LIKE pattern
+    like = LIKE '<pattern>'
+
+    Description
+    -----------
+    Displays a list of all the Fusion commands.
+
+    Arguments
+    ---------
+    * `<pattern>``: A pattern similar to SQL LIKE clause. Uses ``%`` as
+      the wildcard character.
+
+    Remarks
+    -------
+    * Use the ``LIKE`` clause to specify a pattern and return only the
+      commands that match the specified pattern.
+
+    Example
+    -------
+    The following command returns all the Fusion commands that start
+    with 'SHOW'::
+
+        SHOW FUSION COMMANDS LIKE 'SHOW%'
+
+    See Also
+    --------
+    * ``SHOW FUSION HELP``
+
+    """
+
+    def run(self, params: Dict[str, Any]) -> Optional[result.FusionSQLResult]:
+        res = result.FusionSQLResult()
+        res.add_field('Command', result.STRING)
+
+        data: List[Tuple[Any, ...]] = []
+        for _, v in sorted(_handlers.items()):
+            data.append((v.syntax.lstrip(),))
+
+        res.set_rows(data)
+
+        if params['like']:
+            res = res.like(Command=params['like'])
+
+        return res
+
+
+ShowFusionCommandsHandler.register()
+
+
+class ShowFusionGrammarHandler(SQLHandler):
+    """
+    SHOW FUSION GRAMMAR for_query;
+
+    # Query to show grammar for
+    for_query = FOR '<query>'
+
+    Description
+    -----------
+    Show the full grammar of a Fusion SQL command for a given query.
+
+    Arguments
+    ---------
+    * ``<command>``: A Fusion command.
+
+    Example
+    -------
+    The following command displays the grammar for the
+    ``CREATE WORKSPACE`` Fusion command::
+
+        SHOW FUSION GRAMMAR FOR 'CREATE WORKSPACE';
+
+    """
+
+    def run(self, params: Dict[str, Any]) -> Optional[result.FusionSQLResult]:
+        res = result.FusionSQLResult()
+        res.add_field('Grammar', result.STRING)
+        handler = get_handler(params['for_query'])
+        data: List[Tuple[Any, ...]] = []
+        if handler is not None:
+            data.append((handler._grammar,))
+        res.set_rows(data)
+        return res
+
+
+ShowFusionGrammarHandler.register()
+
+
+class ShowFusionHelpHandler(SQLHandler):
+    """
+    SHOW FUSION HELP for_command;
+
+    # Command to show help for
+    for_command = FOR '<command>'
+
+    Description
+    -----------
+    Displays the documentation for a Fusion command.
+
+    Arguments
+    ---------
+    * ``<command>``: A Fusion command.
+
+    Example
+    -------
+    The following command displays the documentation for
+    the ``CREATE WORKSPACE`` Fusion command.
+
+        SHOW FUSION HELP FOR 'CREATE WORKSPACE';
+
+    """
+
+    def run(self, params: Dict[str, Any]) -> Optional[result.FusionSQLResult]:
+        handler = get_handler(params['for_command'])
+        if handler is not None:
+            try:
+                from IPython.display import display
+                from IPython.display import Markdown
+                display(Markdown(handler.help))
+            except Exception:
+                print(handler.help)
+        else:
+            print(
+                f'No handler found for command \'{params["for_command"]}\'',
+                file=sys.stderr,
+            )
+        return None
+
+
+ShowFusionHelpHandler.register()

singlestoredb/fusion/result.py

@@ -0,0 +1,399 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import re
+from collections.abc import Iterable
+from typing import Any
+from typing import List
+from typing import Optional
+from typing import Tuple
+from typing import Union
+
+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 FusionSQLColumn(object):
+    """Column accessor for a FusionSQLResult."""
+
+    def __init__(self, result: FusionSQLResult, index: int) -> None:
+        self._result = result
+        self._index = index
+
+    def __getitem__(self, index: Any) -> Any:
+        return self._result.rows[index][self._index]
+
+    def __iter__(self) -> Iterable[Any]:
+        def gen() -> Iterable[Any]:
+            for row in iter(self._result):
+                yield row[self._index]
+        return gen()
+
+
+class FieldIndexDict(dict):  # type: ignore
+    """Case-insensitive dictionary for column name lookups."""
+
+    def __getitem__(self, key: str) -> int:
+        return super().__getitem__(key.lower())
+
+    def __setitem__(self, key: str, value: int) -> None:
+        super().__setitem__(key.lower(), value)
+
+    def __contains__(self, key: object) -> bool:
+        if not isinstance(key, str):
+            return False
+        return super().__contains__(str(key).lower())
+
+    def copy(self) -> FieldIndexDict:
+        out = type(self)()
+        for k, v in self.items():
+            out[k.lower()] = v
+        return out
+
+
+class FusionSQLResult(object):
+    """Result for Fusion SQL commands."""
+
+    def __init__(self) -> None:
+        self.connection: Any = None
+        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._field_indexes: FieldIndexDict = FieldIndexDict()
+        self._row_idx: int = -1
+
+    def copy(self) -> FusionSQLResult:
+        """Copy the result."""
+        out = type(self)()
+        for k, v in vars(self).items():
+            if isinstance(v, list):
+                setattr(out, k, list(v))
+            elif isinstance(v, dict):
+                setattr(out, k, v.copy())
+            else:
+                setattr(out, k, v)
+        return out
+
+    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 format_results(self, connection: connection.Connection) -> None:
+        """
+        Format the results using the connection converters and options.
+
+        Parameters
+        ----------
+        connection : Connection
+            The connection containing the converters and options
+
+        """
+        self.converters = []
+
+        for item in self.description:
+            self.converters.append((
+                item.charset,
+                connection.decoders.get(item.type_code),
+            ))
+
+        # Convert values
+        for i, row in enumerate(self.rows):
+            new_row = []
+            for (_, converter), value in zip(self.converters, row):
+                new_row.append(converter(value) if converter is not None else value)
+            self.rows[i] = tuple(new_row)
+
+        self.rows[:] = format_results(
+            connection._results_type, self.description, self.rows,
+        )
+
+    def __iter__(self) -> Iterable[Tuple[Any, ...]]:
+        return iter(self.rows)
+
+    def __len__(self) -> int:
+        return len(self.rows)
+
+    def __getitem__(self, key: Any) -> Tuple[Any, ...]:
+        if isinstance(key, str):
+            return self.__getattr__(key)
+        return self.rows[key]
+
+    def __getattr__(self, name: str) -> Any:
+        return FusionSQLColumn(self, self._field_indexes[name])
+
+    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))
+        self._field_indexes[name] = len(self.fields) - 1
+        self.converters.append((encoding, None))
+
+    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.
+
+        """
+        self.rows = list(data)
+        self.affected_rows = 0
+
+    def like(self, **kwargs: str) -> FusionSQLResult:
+        """
+        Return a new result containing only rows that match all `kwargs` like patterns.
+
+        Parameters
+        ----------
+        **kwargs : str
+            Each parameter name corresponds to a column name in the result. The value
+            of the parameters is a LIKE pattern to match.
+
+        Returns
+        -------
+        FusionSQLResult
+
+        """
+        likers = []
+        for k, v in kwargs.items():
+            if k not in self._field_indexes:
+                raise KeyError(f'field name does not exist in results: {k}')
+            if not v:
+                continue
+            regex = re.compile(
+                '^{}$'.format(
+                    re.sub(r'\\%', r'.*', re.sub(r'([^\w])', r'\\\1', v)),
+                ), flags=re.I | re.M,
+            )
+            likers.append((self._field_indexes[k], regex))
+
+        filtered_rows = []
+        for row in self.rows:
+            found = True
+            for i, liker in likers:
+                if row[i] is None or not liker.match(row[i]):
+                    found = False
+                    break
+            if found:
+                filtered_rows.append(row)
+
+        out = self.copy()
+        out.rows[:] = filtered_rows
+        return out
+
+    like_all = like
+
+    def like_any(self, **kwargs: str) -> FusionSQLResult:
+        """
+        Return a new result containing only rows that match any `kwargs` like patterns.
+
+        Parameters
+        ----------
+        **kwargs : str
+            Each parameter name corresponds to a column name in the result. The value
+            of the parameters is a LIKE pattern to match.
+
+        Returns
+        -------
+        FusionSQLResult
+
+        """
+        likers = []
+        for k, v in kwargs.items():
+            if k not in self._field_indexes:
+                raise KeyError(f'field name does not exist in results: {k}')
+            if not v:
+                continue
+            regex = re.compile(
+                '^{}$'.format(
+                    re.sub(r'\\%', r'.*', re.sub(r'([^\w])', r'\\\1', v)),
+                ), flags=re.I | re.M,
+            )
+            likers.append((self._field_indexes[k], regex))
+
+        filtered_rows = []
+        for row in self.rows:
+            found = False
+            for i, liker in likers:
+                if liker.match(row[i]):
+                    found = True
+                    break
+            if found:
+                filtered_rows.append(row)
+
+        out = self.copy()
+        out.rows[:] = filtered_rows
+        return out
+
+    def filter(self, **kwargs: str) -> FusionSQLResult:
+        """
+        Return a new result containing only rows that match all `kwargs` values.
+
+        Parameters
+        ----------
+        **kwargs : str
+            Each parameter name corresponds to a column name in the result. The value
+            of the parameters is the value to match.
+
+        Returns
+        -------
+        FusionSQLResult
+
+        """
+        if not kwargs:
+            return self.copy()
+
+        values = []
+        for k, v in kwargs.items():
+            if k not in self._field_indexes:
+                raise KeyError(f'field name does not exist in results: {k}')
+            values.append((self._field_indexes[k], v))
+
+        filtered_rows = []
+        for row in self.rows:
+            found = True
+            for i, val in values:
+                if row[0] != val:
+                    found = False
+                    break
+            if found:
+                filtered_rows.append(row)
+
+        out = self.copy()
+        out.rows[:] = filtered_rows
+        return out
+
+    def limit(self, n_rows: int) -> FusionSQLResult:
+        """
+        Return a new result containing only `n_rows` rows.
+
+        Parameters
+        ----------
+        n_rows : int
+            The number of rows to limit the result to
+
+        Returns
+        -------
+        FusionSQLResult
+
+        """
+        out = self.copy()
+        if n_rows:
+            out.rows[:] = out.rows[:n_rows]
+        return out
+
+    def sort_by(
+        self,
+        by: Union[str, List[str]],
+        ascending: Union[bool, List[bool]] = True,
+    ) -> FusionSQLResult:
+        """
+        Return a new result with rows sorted in specified order.
+
+        Parameters
+        ----------
+        by : str or List[str]
+            Name or names of columns to sort by
+        ascending : bool or List[bool], optional
+            Should the sort order be ascending? If not all sort columns
+            use the same ordering, a list of booleans can be supplied to
+            indicate the order for each column.
+
+        Returns
+        -------
+        FusionSQLResult
+
+        """
+        if not by:
+            return self.copy()
+
+        if isinstance(by, str):
+            by = [by]
+        by = list(reversed(by))
+
+        if isinstance(ascending, bool):
+            ascending = [ascending]
+        ascending = list(reversed(ascending))
+
+        out = self.copy()
+        for i, byvar in enumerate(by):
+            out.rows.sort(
+                key=lambda x: (
+                    0 if x[self._field_indexes[byvar]] is None else 1,
+                    x[self._field_indexes[byvar]],
+                ),
+                reverse=not ascending[i],
+            )
+        return out
+
+    order_by = sort_by

singlestoredb/http/__init__.py

@@ -0,0 +1,27 @@
+from ..exceptions import DatabaseError
+from ..exceptions import DataError
+from ..exceptions import Error
+from ..exceptions import IntegrityError
+from ..exceptions import InterfaceError
+from ..exceptions import InternalError
+from ..exceptions import ManagementError
+from ..exceptions import NotSupportedError
+from ..exceptions import OperationalError
+from ..exceptions import ProgrammingError
+from ..exceptions import Warning
+from ..types import BINARY
+from ..types import Binary
+from ..types import Date
+from ..types import DateFromTicks
+from ..types import DATETIME
+from ..types import NUMBER
+from ..types import ROWID
+from ..types import STRING
+from ..types import Time
+from ..types import TimeFromTicks
+from ..types import Timestamp
+from ..types import TimestampFromTicks
+from .connection import apilevel
+from .connection import connect
+from .connection import paramstyle
+from .connection import threadsafety