singlestoredb 0.3.3__py3-none-any.whl → 1.0.3__py3-none-any.whl

singlestoredb/pytest.py

@@ -0,0 +1,283 @@
+#!/usr/bin/env python
+"""Pytest plugin"""
+import logging
+import os
+import subprocess
+import time
+from enum import Enum
+from typing import Iterator
+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
+
+
+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:
+        self.container_name = 'singlestoredb-test-container'
+        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',
+        }
+
+        self.ports = ['3306', '8080', '9000']
+
+        self.url = f'root:{self.root_password}@127.0.0.1:3306'
+
+    def start(self) -> None:
+        command = ' '.join(self._start_command())
+
+        logger.info(f'Starting container {self.container_name}')
+        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(
+                'Failed to start container. '
+                'Is one already running?',
+            ) 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 port in self.ports:
+            yield '-p'
+            yield f'{port}:{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}')

singlestoredb/tests/ext_funcs/__init__.py

@@ -0,0 +1,385 @@
+#!/usr/bin/env python3
+# type: ignore
+from typing import Optional
+from typing import Tuple
+
+from singlestoredb.functions.decorator import udf
+from singlestoredb.functions.dtypes import BIGINT
+from singlestoredb.functions.dtypes import FLOAT
+from singlestoredb.functions.dtypes import MEDIUMINT
+from singlestoredb.functions.dtypes import SMALLINT
+from singlestoredb.functions.dtypes import TINYINT
+from singlestoredb.functions.dtypes import VARCHAR
+
+
+@udf
+def double_mult(x: float, y: float) -> float:
+    return x * y
+
+
+@udf.pandas
+def pandas_double_mult(x: float, y: float) -> float:
+    return x * y
+
+
+@udf.numpy
+def numpy_double_mult(x: float, y: float) -> float:
+    return x * y
+
+
+@udf.arrow
+def arrow_double_mult(x: float, y: float) -> float:
+    import pyarrow.compute as pc
+    return pc.multiply(x, y)
+
+
+@udf.polars
+def polars_double_mult(x: float, y: float) -> float:
+    return x * y
+
+
+@udf
+def nullable_double_mult(x: Optional[float], y: Optional[float]) -> Optional[float]:
+    if x is None or y is None:
+        return None
+    return x * y
+
+
+@udf(args=[FLOAT(nullable=False), FLOAT(nullable=False)], returns=FLOAT(nullable=False))
+def float_mult(x: float, y: float) -> float:
+    return x * y
+
+
+@udf(args=[FLOAT(nullable=True), FLOAT(nullable=True)], returns=FLOAT(nullable=True))
+def nullable_float_mult(x: Optional[float], y: Optional[float]) -> Optional[float]:
+    if x is None or y is None:
+        return None
+    return x * y
+
+
+def _int_mult(x: int, y: int) -> int:
+    if x is None or y is None:
+        return None
+    return x * y
+
+
+def _arrow_int_mult(x: int, y: int) -> int:
+    import pyarrow.compute as pc
+    return pc.multiply(x, y)
+
+
+def _int_mult_with_masks(x: Tuple[int, bool], y: Tuple[int, bool]) -> Tuple[int, bool]:
+    x_data, x_nulls = x
+    y_data, y_nulls = y
+    return (x_data * y_data, x_nulls | y_nulls)
+
+
+def _arrow_int_mult_with_masks(
+    x: Tuple[int, bool],
+    y: Tuple[int, bool],
+) -> Tuple[int, bool]:
+    import pyarrow.compute as pc
+    x_data, x_nulls = x
+    y_data, y_nulls = y
+    return (pc.multiply(x_data, y_data), pc.or_(x_nulls, y_nulls))
+
+
+int_mult = udf(_int_mult, name='int_mult')
+
+tinyint_mult = udf(
+    _int_mult,
+    name='tinyint_mult',
+    args=[TINYINT(nullable=False), TINYINT(nullable=False)],
+    returns=TINYINT(nullable=False),
+)
+
+pandas_tinyint_mult = udf.pandas(
+    _int_mult,
+    name='pandas_tinyint_mult',
+    args=[TINYINT(nullable=False), TINYINT(nullable=False)],
+    returns=TINYINT(nullable=False),
+)
+
+polars_tinyint_mult = udf.polars(
+    _int_mult,
+    name='polars_tinyint_mult',
+    args=[TINYINT(nullable=False), TINYINT(nullable=False)],
+    returns=TINYINT(nullable=False),
+)
+
+numpy_tinyint_mult = udf.numpy(
+    _int_mult,
+    name='numpy_tinyint_mult',
+    args=[TINYINT(nullable=False), TINYINT(nullable=False)],
+    returns=TINYINT(nullable=False),
+)
+
+arrow_tinyint_mult = udf.arrow(
+    _arrow_int_mult,
+    name='arrow_tinyint_mult',
+    args=[TINYINT(nullable=False), TINYINT(nullable=False)],
+    returns=TINYINT(nullable=False),
+)
+
+smallint_mult = udf(
+    _int_mult,
+    name='smallint_mult',
+    args=[SMALLINT(nullable=False), SMALLINT(nullable=False)],
+    returns=SMALLINT(nullable=False),
+)
+
+pandas_smallint_mult = udf.pandas(
+    _int_mult,
+    name='pandas_smallint_mult',
+    args=[SMALLINT(nullable=False), SMALLINT(nullable=False)],
+    returns=SMALLINT(nullable=False),
+)
+
+polars_smallint_mult = udf.polars(
+    _int_mult,
+    name='polars_smallint_mult',
+    args=[SMALLINT(nullable=False), SMALLINT(nullable=False)],
+    returns=SMALLINT(nullable=False),
+)
+
+numpy_smallint_mult = udf.numpy(
+    _int_mult,
+    name='numpy_smallint_mult',
+    args=[SMALLINT(nullable=False), SMALLINT(nullable=False)],
+    returns=SMALLINT(nullable=False),
+)
+
+arrow_smallint_mult = udf.arrow(
+    _arrow_int_mult,
+    name='arrow_smallint_mult',
+    args=[SMALLINT(nullable=False), SMALLINT(nullable=False)],
+    returns=SMALLINT(nullable=False),
+)
+
+mediumint_mult = udf(
+    _int_mult,
+    name='mediumint_mult',
+    args=[MEDIUMINT(nullable=False), MEDIUMINT(nullable=False)],
+    returns=MEDIUMINT(nullable=False),
+)
+
+pandas_mediumint_mult = udf.pandas(
+    _int_mult,
+    name='pandas_mediumint_mult',
+    args=[MEDIUMINT(nullable=False), MEDIUMINT(nullable=False)],
+    returns=MEDIUMINT(nullable=False),
+)
+
+polars_mediumint_mult = udf.polars(
+    _int_mult,
+    name='polars_mediumint_mult',
+    args=[MEDIUMINT(nullable=False), MEDIUMINT(nullable=False)],
+    returns=MEDIUMINT(nullable=False),
+)
+
+numpy_mediumint_mult = udf.numpy(
+    _int_mult,
+    name='numpy_mediumint_mult',
+    args=[MEDIUMINT(nullable=False), MEDIUMINT(nullable=False)],
+    returns=MEDIUMINT(nullable=False),
+)
+
+arrow_mediumint_mult = udf.arrow(
+    _arrow_int_mult,
+    name='arrow_mediumint_mult',
+    args=[MEDIUMINT(nullable=False), MEDIUMINT(nullable=False)],
+    returns=MEDIUMINT(nullable=False),
+)
+
+bigint_mult = udf(
+    _int_mult,
+    name='bigint_mult',
+    args=[BIGINT(nullable=False), BIGINT(nullable=False)],
+    returns=BIGINT(nullable=False),
+)
+
+pandas_bigint_mult = udf.pandas(
+    _int_mult,
+    name='pandas_bigint_mult',
+    args=[BIGINT(nullable=False), BIGINT(nullable=False)],
+    returns=BIGINT(nullable=False),
+)
+
+polars_bigint_mult = udf.polars(
+    _int_mult,
+    name='polars_bigint_mult',
+    args=[BIGINT(nullable=False), BIGINT(nullable=False)],
+    returns=BIGINT(nullable=False),
+)
+
+numpy_bigint_mult = udf.numpy(
+    _int_mult,
+    name='numpy_bigint_mult',
+    args=[BIGINT(nullable=False), BIGINT(nullable=False)],
+    returns=BIGINT(nullable=False),
+)
+
+arrow_bigint_mult = udf.arrow(
+    _arrow_int_mult,
+    name='arrow_bigint_mult',
+    args=[BIGINT(nullable=False), BIGINT(nullable=False)],
+    returns=BIGINT(nullable=False),
+)
+
+nullable_tinyint_mult = udf(
+    _int_mult,
+    name='nullable_tinyint_mult',
+    args=[TINYINT, TINYINT],
+    returns=TINYINT,
+)
+
+pandas_nullable_tinyint_mult = udf.pandas(
+    _int_mult,
+    name='pandas_nullable_tinyint_mult',
+    args=[TINYINT, TINYINT],
+    returns=TINYINT,
+)
+
+pandas_nullable_tinyint_mult_with_masks = udf.pandas(
+    _int_mult_with_masks,
+    name='pandas_nullable_tinyint_mult_with_masks',
+    args=[TINYINT, TINYINT],
+    returns=TINYINT,
+    include_masks=True,
+)
+
+polars_nullable_tinyint_mult = udf.polars(
+    _int_mult,
+    name='polars_nullable_tinyint_mult',
+    args=[TINYINT, TINYINT],
+    returns=TINYINT,
+)
+
+polars_nullable_tinyint_mult_with_masks = udf.polars(
+    _int_mult_with_masks,
+    name='polars_nullable_tinyint_mult_with_masks',
+    args=[TINYINT, TINYINT],
+    returns=TINYINT,
+    include_masks=True,
+)
+
+numpy_nullable_tinyint_mult = udf.numpy(
+    _int_mult,
+    name='numpy_nullable_tinyint_mult',
+    args=[TINYINT, TINYINT],
+    returns=TINYINT,
+)
+
+numpy_nullable_tinyint_mult_with_masks = udf.numpy(
+    _int_mult_with_masks,
+    name='numpy_nullable_tinyint_mult_with_masks',
+    args=[TINYINT, TINYINT],
+    returns=TINYINT,
+    include_masks=True,
+)
+
+arrow_nullable_tinyint_mult = udf.arrow(
+    _arrow_int_mult,
+    name='arrow_nullable_tinyint_mult',
+    args=[TINYINT, TINYINT],
+    returns=TINYINT,
+)
+
+arrow_nullable_tinyint_mult_with_masks = udf.arrow(
+    _arrow_int_mult_with_masks,
+    name='arrow_nullable_tinyint_mult_with_masks',
+    args=[TINYINT, TINYINT],
+    returns=TINYINT,
+    include_masks=True,
+)
+
+nullable_smallint_mult = udf(
+    _int_mult,
+    name='nullable_smallint_mult',
+    args=[SMALLINT, SMALLINT],
+    returns=SMALLINT,
+)
+
+nullable_mediumint_mult = udf(
+    _int_mult,
+    name='nullable_mediumint_mult',
+    args=[MEDIUMINT, MEDIUMINT],
+    returns=MEDIUMINT,
+)
+
+nullable_bigint_mult = udf(
+    _int_mult,
+    name='nullable_bigint_mult',
+    args=[BIGINT, BIGINT],
+    returns=BIGINT,
+)
+
+numpy_nullable_bigint_mult = udf.numpy(
+    _int_mult,
+    name='numpy_nullable_bigint_mult',
+    args=[BIGINT, BIGINT],
+    returns=BIGINT,
+)
+
+numpy_nullable_bigint_mult_with_masks = udf.numpy(
+    _int_mult_with_masks,
+    name='numpy_nullable_bigint_mult',
+    args=[BIGINT, BIGINT],
+    returns=BIGINT,
+    include_masks=True,
+)
+
+
+@udf
+def nullable_int_mult(x: Optional[int], y: Optional[int]) -> Optional[int]:
+    if x is None or y is None:
+        return None
+    return x * y
+
+
+@udf
+def string_mult(x: str, times: int) -> str:
+    return x * times
+
+
+@udf.pandas
+def pandas_string_mult(x: str, times: int) -> str:
+    return x * times
+
+
+@udf.numpy
+def numpy_string_mult(x: str, times: int) -> str:
+    return x * times
+
+
+# @udf.polars
+# def polars_string_mult(x: str, times: int) -> str:
+#     print(type(x), x, type(times), times)
+#     return x * times
+
+
+# @udf.arrow
+# def arrow_string_mult(x: str, times: int) -> str:
+#     print(type(x), x, type(times), times)
+#     import pyarrow.compute as pc
+#     return pc.multiply(x, times)
+#     return x * times
+
+
+@udf
+def nullable_string_mult(x: Optional[str], times: Optional[int]) -> Optional[str]:
+    if x is None or times is None:
+        return None
+    return x * times
+
+
+@udf(args=dict(x=VARCHAR(20, nullable=False)))
+def varchar_mult(x: str, times: int) -> str:
+    return x * times
+
+
+@udf(args=dict(x=VARCHAR(20, nullable=True)))
+def nullable_varchar_mult(x: Optional[str], times: Optional[int]) -> Optional[str]:
+    if x is None or times is None:
+        return None
+    return x * times