singlestoredb 1.16.1__py3-none-any.whl

singlestoredb/notebook/_portal.py

@@ -0,0 +1,352 @@
+#!/usr/bin/env python
+import json
+import os
+import re
+import time
+import urllib.parse
+from typing import Any
+from typing import Callable
+from typing import Dict
+from typing import List
+from typing import Optional
+from typing import Tuple
+from typing import Union
+
+from . import _objects as obj
+from ..management import workspace as mgr
+from ..utils import events
+
+try:
+    from IPython import display
+    has_ipython = True
+except ImportError:
+    has_ipython = False
+
+
+class Portal(object):
+    """SingleStore Portal information."""
+
+    def __init__(self) -> None:
+        self._connection_info: Dict[str, Any] = {}
+        self._authentication_info: Dict[str, Any] = {}
+        self._theme_info: Dict[str, Any] = {}
+        events.subscribe(self._request)
+
+    def __str__(self) -> str:
+        attrs = []
+        for name in [
+            'organization_id', 'workspace_group_id', 'workspace_id',
+            'host', 'port', 'user', 'password', 'default_database',
+        ]:
+            if name == 'password':
+                if self.password is not None:
+                    attrs.append("password='***'")
+                else:
+                    attrs.append('password=None')
+            else:
+                attrs.append(f'{name}={getattr(self, name)!r}')
+        return f'{type(self).__name__}({", ".join(attrs)})'
+
+    def __repr__(self) -> str:
+        return str(self)
+
+    def _call_javascript(
+        self,
+        func: str,
+        args: Optional[List[Any]] = None,
+        wait_on_condition: Optional[Callable[[], bool]] = None,
+        timeout_message: str = 'timed out waiting on condition',
+        wait_interval: float = 2.0,
+        timeout: float = 60.0,
+    ) -> None:
+        if not has_ipython or not func:
+            return
+
+        if not re.match(r'^[A-Z_][\w\._]*$', func, flags=re.I):
+            raise ValueError(f'function name is not valid: {func}')
+
+        args = args if args else []
+
+        code = f'''
+            if (window.singlestore && window.singlestore.portal) {{
+                window.singlestore.portal.{func}.apply(
+                    window,
+                    JSON.parse({repr(json.dumps(args))})
+                )
+            }}
+        '''
+
+        display.display(display.Javascript(code))
+
+        if wait_on_condition is not None:
+            elapsed = 0.0
+            while True:
+                if wait_on_condition():
+                    break
+                if elapsed > timeout:
+                    raise RuntimeError(timeout_message)
+                time.sleep(wait_interval)
+                elapsed += wait_interval
+
+    def _request(self, msg: Dict[str, Any]) -> None:
+        """Handle request on the control stream."""
+        func = getattr(self, '_handle_' + msg.get('name', 'unknown').split('.')[-1])
+        if func is not None:
+            func(msg.get('data', {}))
+
+    def _handle_connection_updated(self, data: Dict[str, Any]) -> None:
+        """Handle connection_updated event."""
+        self._connection_info = dict(data)
+
+    def _handle_authentication_updated(self, data: Dict[str, Any]) -> None:
+        """Handle authentication_updated event."""
+        self._authentication_info = dict(data)
+
+    def _handle_theme_updated(self, data: Dict[str, Any]) -> None:
+        """Handle theme_updated event."""
+        self._theme_info = dict(data)
+
+    def _handle_unknown(self, data: Dict[str, Any]) -> None:
+        """Handle unknown events."""
+        pass
+
+    @property
+    def organization_id(self) -> Optional[str]:
+        """Organization ID."""
+        try:
+            return self._connection_info['organization']
+        except KeyError:
+            return os.environ.get('SINGLESTOREDB_ORGANIZATION')
+
+    @property
+    def organization(self) -> obj.Organization:
+        """Organization."""
+        return obj.organization
+
+    @property
+    def stage(self) -> obj.Stage:
+        """Stage."""
+        return obj.stage
+
+    @property
+    def secrets(self) -> obj.Secrets:
+        """Secrets."""
+        return obj.secrets
+
+    @property
+    def workspace_group_id(self) -> Optional[str]:
+        """Workspace Group ID."""
+        try:
+            return self._connection_info['workspace_group']
+        except KeyError:
+            return os.environ.get('SINGLESTOREDB_WORKSPACE_GROUP')
+
+    @property
+    def workspace_group(self) -> obj.WorkspaceGroup:
+        """Workspace group."""
+        return obj.workspace_group
+
+    @workspace_group.setter
+    def workspace_group(self) -> None:
+        """Set workspace group."""
+        raise AttributeError(
+            'workspace group can not be set explictly; ' +
+            'you can only set a workspace',
+        )
+
+    @property
+    def workspace_id(self) -> Optional[str]:
+        """Workspace ID."""
+        try:
+            return self._connection_info['workspace']
+        except KeyError:
+            return os.environ.get('SINGLESTOREDB_WORKSPACE')
+
+    @property
+    def workspace(self) -> obj.Workspace:
+        """Workspace."""
+        return obj.workspace
+
+    @workspace.setter
+    def workspace(self, workspace_spec: Union[str, Tuple[str, str]]) -> None:
+        """Set workspace."""
+        if isinstance(workspace_spec, tuple):
+            # 2-element tuple: (workspace_group_id, workspace_name_or_id)
+            workspace_group_id, name_or_id = workspace_spec
+            uuid_pattern = (
+                r'[a-z0-9]{8}-[a-z0-9]{4}-[a-z0-9]{4}-[a-z0-9]{4}-[a-z0-9]{12}'
+            )
+            if re.match(uuid_pattern, name_or_id, flags=re.I):
+                w = mgr.get_workspace(name_or_id)
+            else:
+                w = mgr.get_workspace_group(workspace_group_id).workspaces[
+                    name_or_id
+                ]
+        else:
+            # String: workspace_name_or_id (existing behavior)
+            name_or_id = workspace_spec
+            uuid_pattern = (
+                r'[a-z0-9]{8}-[a-z0-9]{4}-[a-z0-9]{4}-[a-z0-9]{4}-[a-z0-9]{12}'
+            )
+            if re.match(uuid_pattern, name_or_id, flags=re.I):
+                w = mgr.get_workspace(name_or_id)
+            else:
+                w = mgr.get_workspace_group(
+                    self.workspace_group_id,
+                ).workspaces[name_or_id]
+
+        if w.state and w.state.lower() not in ['active', 'resumed']:
+            raise RuntimeError('workspace is not active')
+
+        id = w.id
+
+        self._call_javascript(
+            'changeDeployment', [id],
+            wait_on_condition=lambda: self.workspace_id == id,  # type: ignore
+            timeout_message='timeout waiting for workspace update',
+        )
+
+    deployment = workspace
+
+    @property
+    def connection(self) -> Tuple[obj.Workspace, Optional[str]]:
+        """Workspace and default database name."""
+        return self.workspace, self.default_database
+
+    @connection.setter
+    def connection(
+        self,
+        connection_spec: Union[Tuple[str, str], Tuple[str, str, str]],
+    ) -> None:
+        """Set workspace and default database name."""
+        if len(connection_spec) == 3:
+            # 3-element tuple: (workspace_group_id, workspace_name_or_id,
+            # default_database)
+            workspace_group_id, name_or_id, default_database = connection_spec
+            uuid_pattern = (
+                r'[a-z0-9]{8}-[a-z0-9]{4}-[a-z0-9]{4}-[a-z0-9]{4}-[a-z0-9]{12}'
+            )
+            if re.match(uuid_pattern, name_or_id, flags=re.I):
+                w = mgr.get_workspace(name_or_id)
+            else:
+                w = mgr.get_workspace_group(workspace_group_id).workspaces[
+                    name_or_id
+                ]
+        else:
+            # 2-element tuple: (workspace_name_or_id, default_database)
+            # existing behavior
+            name_or_id, default_database = connection_spec
+            uuid_pattern = (
+                r'[a-z0-9]{8}-[a-z0-9]{4}-[a-z0-9]{4}-[a-z0-9]{4}-[a-z0-9]{12}'
+            )
+            if re.match(uuid_pattern, name_or_id, flags=re.I):
+                w = mgr.get_workspace(name_or_id)
+            else:
+                w = mgr.get_workspace_group(
+                    self.workspace_group_id,
+                ).workspaces[name_or_id]
+
+        if w.state and w.state.lower() not in ['active', 'resumed']:
+            raise RuntimeError('workspace is not active')
+
+        id = w.id
+
+        self._call_javascript(
+            'changeConnection', [id, default_database],
+            wait_on_condition=lambda: self.workspace_id == id and
+            self.default_database == default_database,  # type: ignore
+            timeout_message='timeout waiting for workspace update',
+        )
+
+    @property
+    def cluster_id(self) -> Optional[str]:
+        """Cluster ID."""
+        try:
+            return self._connection_info['cluster']
+        except KeyError:
+            return os.environ.get('SINGLESTOREDB_CLUSTER')
+
+    def _parse_url(self) -> Dict[str, Any]:
+        url = urllib.parse.urlparse(
+            os.environ.get('SINGLESTOREDB_URL', ''),
+        )
+        return dict(
+            host=url.hostname or None,
+            port=url.port or None,
+            user=url.username or None,
+            password=url.password or None,
+            default_database=url.path.split('/')[-1] or None,
+        )
+
+    @property
+    def connection_url(self) -> Optional[str]:
+        """Connection URL."""
+        try:
+            return self._connection_info['connection_url']
+        except KeyError:
+            return os.environ.get('SINGLESTOREDB_URL')
+
+    @property
+    def connection_url_kai(self) -> Optional[str]:
+        """Kai connectionURL."""
+        try:
+            return self._connection_info.get('connection_url_kai')
+        except KeyError:
+            return os.environ.get('SINGLESTOREDB_URL_KAI')
+
+    @property
+    def host(self) -> Optional[str]:
+        """Hostname."""
+        try:
+            return self._connection_info['host']
+        except KeyError:
+            return self._parse_url()['host']
+
+    @property
+    def port(self) -> Optional[int]:
+        """Database server port."""
+        try:
+            return self._connection_info['port']
+        except KeyError:
+            return self._parse_url()['port']
+
+    @property
+    def user(self) -> Optional[str]:
+        """Username."""
+        try:
+            return self._authentication_info['user']
+        except KeyError:
+            return self._parse_url()['user']
+
+    @property
+    def password(self) -> Optional[str]:
+        """Password."""
+        try:
+            return self._authentication_info['password']
+        except KeyError:
+            return self._parse_url()['password']
+
+    @property
+    def default_database(self) -> Optional[str]:
+        """Default database."""
+        try:
+            return self._connection_info['default_database']
+        except KeyError:
+            return self._parse_url()['default_database']
+
+    @default_database.setter
+    def default_database(self, name: str) -> None:
+        """Set default database."""
+        self._call_javascript(
+            'changeDefaultDatabase', [name],
+            wait_on_condition=lambda: self.default_database == name,  # type: ignore
+            timeout_message='timeout waiting for database update',
+        )
+
+    @property
+    def version(self) -> Optional[str]:
+        """Version."""
+        return self._connection_info.get('version')
+
+
+portal = Portal()

singlestoredb/pytest.py

@@ -0,0 +1,352 @@
+#!/usr/bin/env python
+"""Pytest plugin"""
+import logging
+import os
+import socket
+import subprocess
+import time
+import uuid
+from collections.abc import Iterator
+from enum import Enum
+from typing import Optional
+
+import pytest
+
+from . import connect
+from .connection import Connection
+from .connection import Cursor
+
+
+logger = logging.getLogger(__name__)
+
+
+# How many times to attempt to connect to the container
+STARTUP_CONNECT_ATTEMPTS = 10
+# How long to wait between connection attempts
+STARTUP_CONNECT_TIMEOUT_SECONDS = 2
+# How many times to check if all connections are closed
+TEARDOWN_WAIT_ATTEMPTS = 20
+# How long to wait between checking connections
+TEARDOWN_WAIT_SECONDS = 2
+
+
+def _find_free_port() -> int:
+    """Find a free port by binding to port 0 and getting the assigned port."""
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.bind(('', 0))
+        s.listen(1)
+        return s.getsockname()[1]
+
+
+class ExecutionMode(Enum):
+    SEQUENTIAL = 1
+    LEADER = 2
+    FOLLOWER = 3
+
+
+@pytest.fixture(scope='session')
+def execution_mode() -> ExecutionMode:
+    """Determine the pytest mode for this process"""
+
+    worker = os.environ.get('PYTEST_XDIST_WORKER')
+    worker_count = os.environ.get('PYTEST_XDIST_WORKER_COUNT')
+
+    # If we're not in pytest-xdist, the mode is Sequential
+    if worker is None or worker_count is None:
+        logger.debug('XDIST environment vars not found')
+        return ExecutionMode.SEQUENTIAL
+
+    logger.debug(f'PYTEST_XDIST_WORKER == {worker}')
+    logger.debug(f'PYTEST_XDIST_WORKER_COUNT == {worker_count}')
+
+    # If we're the only worker, than the mode is Sequential
+    if worker_count == '1':
+        return ExecutionMode.SEQUENTIAL
+    else:
+        # The first worker (named "gw0") is the leader
+        # if there are multiple workers
+        if worker == 'gw0':
+            return ExecutionMode.LEADER
+        else:
+            return ExecutionMode.FOLLOWER
+
+
+@pytest.fixture(scope='session')
+def node_name() -> Iterator[str]:
+    """Determine the name of this worker node"""
+
+    worker = os.environ.get('PYTEST_XDIST_WORKER')
+
+    if worker is None:
+        logger.debug('XDIST environment vars not found')
+        yield 'master'
+    else:
+        logger.debug(f'PYTEST_XDIST_WORKER == {worker}')
+        yield worker
+
+
+class _TestContainerManager():
+    """Manages the setup and teardown of a SingleStoreDB Dev Container"""
+
+    def __init__(self) -> None:
+        # Generate unique container name using UUID and worker ID
+        worker = os.environ.get('PYTEST_XDIST_WORKER', 'master')
+        unique_id = uuid.uuid4().hex[:8]
+        self.container_name = f'singlestoredb-test-{worker}-{unique_id}'
+
+        self.dev_image_name = 'ghcr.io/singlestore-labs/singlestoredb-dev'
+
+        assert 'SINGLESTORE_LICENSE' in os.environ, 'SINGLESTORE_LICENSE not set'
+
+        self.root_password = 'Q8r4D7yXR8oqn'
+        self.environment_vars = {
+            'SINGLESTORE_LICENSE': None,
+            'ROOT_PASSWORD': f"\"{self.root_password}\"",
+            'SINGLESTORE_SET_GLOBAL_DEFAULT_PARTITIONS_PER_LEAF': '1',
+        }
+
+        # Use dynamic port allocation to avoid conflicts
+        self.mysql_port = _find_free_port()
+        self.http_port = _find_free_port()
+        self.studio_port = _find_free_port()
+        self.ports = [
+            (self.mysql_port, '3306'),    # External port -> Internal port
+            (self.http_port, '8080'),
+            (self.studio_port, '9000'),
+        ]
+
+        self.url = f'root:{self.root_password}@127.0.0.1:{self.mysql_port}'
+
+    def _container_exists(self) -> bool:
+        """Check if a container with this name already exists."""
+        try:
+            result = subprocess.run(
+                [
+                    'docker', 'ps', '-a', '--filter',
+                    f'name={self.container_name}',
+                    '--format', '{{.Names}}',
+                ],
+                capture_output=True,
+                text=True,
+                check=True,
+            )
+            return self.container_name in result.stdout
+        except subprocess.CalledProcessError:
+            return False
+
+    def _cleanup_existing_container(self) -> None:
+        """Stop and remove any existing container with the same name."""
+        if not self._container_exists():
+            return
+
+        logger.info(f'Found existing container {self.container_name}, cleaning up')
+        try:
+            # Try to stop the container (ignore if it's already stopped)
+            subprocess.run(
+                ['docker', 'stop', self.container_name],
+                capture_output=True,
+                check=False,
+            )
+            # Remove the container
+            subprocess.run(
+                ['docker', 'rm', self.container_name],
+                capture_output=True,
+                check=True,
+            )
+            logger.debug(f'Cleaned up existing container {self.container_name}')
+        except subprocess.CalledProcessError as e:
+            logger.warning(f'Failed to cleanup existing container: {e}')
+            # Continue anyway - the unique name should prevent most conflicts
+
+    def start(self) -> None:
+        # Clean up any existing container with the same name
+        self._cleanup_existing_container()
+
+        command = ' '.join(self._start_command())
+
+        logger.info(
+            f'Starting container {self.container_name} on ports {self.mysql_port}, '
+            f'{self.http_port}, {self.studio_port}',
+        )
+        try:
+            license = os.environ['SINGLESTORE_LICENSE']
+            env = {
+                'SINGLESTORE_LICENSE': license,
+            }
+            subprocess.check_call(command, shell=True, env=env)
+        except Exception as e:
+            logger.exception(e)
+            raise RuntimeError(
+                f'Failed to start container {self.container_name}. '
+                f'Command: {command}',
+            ) from e
+        logger.debug('Container started')
+
+    def _start_command(self) -> Iterator[str]:
+        yield 'docker run -d --name'
+        yield self.container_name
+        for key, value in self.environment_vars.items():
+            yield '-e'
+            if value is None:
+                yield key
+            else:
+                yield f'{key}={value}'
+
+        for external_port, internal_port in self.ports:
+            yield '-p'
+            yield f'{external_port}:{internal_port}'
+
+        yield self.dev_image_name
+
+    def print_logs(self) -> None:
+        logs_command = ['docker', 'logs', self.container_name]
+        logger.info('Getting logs')
+        logger.info(subprocess.check_output(logs_command))
+
+    def connect(self) -> Connection:
+        # Run all but one attempts trying again if they fail
+        for i in range(STARTUP_CONNECT_ATTEMPTS - 1):
+            try:
+                return connect(self.url)
+            except Exception:
+                logger.debug(f'Database not available yet (attempt #{i}).')
+                time.sleep(STARTUP_CONNECT_TIMEOUT_SECONDS)
+        else:
+            # Try one last time and report error if it fails
+            try:
+                return connect(self.url)
+            except Exception as e:
+                logger.error('Timed out while waiting to connect to database.')
+                logger.exception(e)
+                self.print_logs()
+                raise RuntimeError('Failed to connect to database') from e
+
+    def wait_till_connections_closed(self) -> None:
+        heart_beat = connect(self.url)
+        for i in range(TEARDOWN_WAIT_ATTEMPTS):
+            connections = self.get_open_connections(heart_beat)
+            if connections is None:
+                raise RuntimeError('Could not determine the number of open connections.')
+            logger.debug(
+                f'Waiting for other connections (n={connections-1}) '
+                f'to close (attempt #{i})',
+            )
+            time.sleep(TEARDOWN_WAIT_SECONDS)
+        else:
+            logger.warning('Timed out while waiting for other connections to close')
+            self.print_logs()
+
+    def get_open_connections(self, conn: Connection) -> Optional[int]:
+        for row in conn.show.status(extended=True):
+            name = row['Name']
+            value = row['Value']
+            logger.info(f'{name} = {value}')
+            if name == 'Threads_connected':
+                return int(value)
+
+        return None
+
+    def stop(self) -> None:
+        logger.info('Cleaning up SingleStore DB dev container')
+        logger.debug('Stopping container')
+        try:
+            subprocess.check_call(f'docker stop {self.container_name}', shell=True)
+        except Exception as e:
+            logger.exception(e)
+            raise RuntimeError('Failed to stop container.') from e
+
+        logger.debug('Removing container')
+        try:
+            subprocess.check_call(f'docker rm {self.container_name}', shell=True)
+        except Exception as e:
+            logger.exception(e)
+            raise RuntimeError('Failed to stop container.') from e
+
+
+@pytest.fixture(scope='session')
+def singlestoredb_test_container(
+    execution_mode: ExecutionMode,
+) -> Iterator[_TestContainerManager]:
+    """Sets up and tears down the test container"""
+
+    if not isinstance(execution_mode, ExecutionMode):
+        raise TypeError(f"Invalid execution mode '{execution_mode}'")
+
+    container_manager = _TestContainerManager()
+
+    # In sequential operation do all the steps
+    if execution_mode == ExecutionMode.SEQUENTIAL:
+        logger.debug('Not distributed')
+        container_manager.start()
+        yield container_manager
+        container_manager.stop()
+
+    # In distributed execution as leader,
+    # do the steps but wait for other workers before stopping
+    elif execution_mode == ExecutionMode.LEADER:
+        logger.debug('Distributed leader')
+        container_manager.start()
+        yield container_manager
+        container_manager.wait_till_connections_closed()
+        container_manager.stop()
+
+    # In distributed exeuction as a non-leader,
+    # don't worry about the container lifecycle
+    elif execution_mode == ExecutionMode.FOLLOWER:
+        logger.debug('Distributed follower')
+        yield container_manager
+
+
+@pytest.fixture(scope='session')
+def singlestoredb_connection(
+    singlestoredb_test_container: _TestContainerManager,
+) -> Iterator[Connection]:
+    """Creates and closes the connection"""
+
+    connection = singlestoredb_test_container.connect()
+    logger.debug('Connected to database.')
+
+    yield connection
+
+    logger.debug('Closing connection')
+    connection.close()
+
+
+class _NameAllocator():
+    """Generates unique names for each database"""
+
+    def __init__(self, id: str) -> None:
+        self.id = id
+        self.names = 0
+
+    def get_name(self) -> str:
+        name = f'x_db_{self.id}_{self.names}'
+        self.names += 1
+        return name
+
+
+@pytest.fixture(scope='session')
+def name_allocator(node_name: str) -> Iterator[_NameAllocator]:
+    """Makes a worker-local name allocator using the node name"""
+
+    yield _NameAllocator(node_name)
+
+
+@pytest.fixture
+def singlestoredb_tempdb(
+    singlestoredb_connection: Connection, name_allocator: _NameAllocator,
+) -> Iterator[Cursor]:
+    """Provides a connection to a unique temporary test database"""
+
+    assert singlestoredb_connection.is_connected(), 'Database is no longer connected'
+    db = name_allocator.get_name()
+
+    with singlestoredb_connection.cursor() as cursor:
+        logger.debug(f"Creating temporary DB \"{db}\"")
+        cursor.execute(f'CREATE DATABASE {db}')
+        cursor.execute(f'USE {db}')
+
+        yield cursor
+
+        logger.debug(f"Dropping temporary DB \"{db}\"")
+        cursor.execute(f'DROP DATABASE {db}')