Source code for myscaledb.db.client

import asyncio
import functools
import logging
from types import TracebackType
from typing import Any, Optional, Type, Dict, List, Iterator

from myscaledb.async_db.client import BaseClient as AsyncChClient
from myscaledb.common.exceptions import ClientError
from myscaledb.common.records import Record
import nest_asyncio
nest_asyncio.apply()

def iterate_async_to_sync(async_iterate, loop):

    async_iterate = async_iterate.__aiter__()

    async def get_next():
        try:
            obj = await async_iterate.__anext__()
            return False, obj
        except StopAsyncIteration:
            return True, None

    while True:
        try:
            done, obj = loop.run_until_complete(get_next())
        except RuntimeError:
            logging.warning("create new event loop...")
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            done, obj = loop.run_until_complete(get_next())
        if done:
            break
        yield obj


def async_to_sync():
    def handle_exception(loop, context):
        # context["message"] will always be there; but context["exception"] may not
        msg = context.get("exception", context["message"])
        logging.error("Caught exception in async method: %s", msg)
        raise ClientError(msg)

    def decorator(func):
        @functools.wraps(func)
        def wrapper(*_args, **_kwargs):
            func_wrapped = func(*_args, **_kwargs)
            if asyncio.iscoroutine(func_wrapped):
                try:
                    loop = asyncio.get_event_loop()
                except RuntimeError:
                    logging.warning("create new event loop...")
                    loop = asyncio.new_event_loop()
                    asyncio.set_event_loop(loop)
                loop.set_exception_handler(handle_exception)
                return loop.run_until_complete(func_wrapped)
            return func_wrapped

        return wrapper

    return decorator


[docs]class Client(AsyncChClient): """Client connection class. Usage: .. code-block:: python from myscaledb import Client def main(): client = Client() alive = client.is_alive() print(f"Is MyScale alive? -> {alive}") if __name__ == '__main__': main() :param aiohttp.ClientSession session: aiohttp client session. Please, use one session and one Client for all connections in your app. :param str url: MyScale server url. Need full path, like http://localhost:8123/. :param str user: User name for authorization. :param str password: Password for authorization. :param str database: Database name. :param bool compress_response: Pass True if you want MyScale to compress its responses with gzip. They will be decompressed automatically. But overall it will be slightly slower. :param **settings: Any settings from https://clickhouse.yandex/docs/en/operations/settings """ def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) # Context manager support def __enter__(self) -> 'Client': if not self.is_alive(): raise ClientError("Client is already closed.") return self def __exit__( self, exc_type: Optional[Type[BaseException]], exc_val: Optional[BaseException], exc_tb: Optional[TracebackType], ) -> None: self.close() def __del__(self): self.close()
[docs] @async_to_sync() async def close(self) -> None: """Close the session""" await self._aclose()
@async_to_sync() async def is_alive(self) -> bool: return await self._is_alive()
[docs] def execute( self, query: str, *args, json: bool = False, params: Optional[Dict[str, Any]] = None, query_id: str = None, ) -> None: """Execute query. Returns None. :param str query: MyScale query string. :param args: Arguments for insert queries. :param bool json: Execute query in JSONEachRow mode. :param Optional[Dict[str, Any]] params: Params to escape inside query string. :param str query_id: MyScale query_id. Usage: .. code-block:: python client.execute( "CREATE TABLE t (a UInt8, b Tuple(Date, Nullable(Float32))) ENGINE = Memory" ) client.execute( "INSERT INTO t VALUES", (1, (dt.date(2018, 9, 7), None)), (2, (dt.date(2018, 9, 8), 3.14)), ) client.execute( "SELECT * FROM t WHERE a={u8}", params={"u8": 12} ) :return: Nothing. """ for _ in iterate_async_to_sync(self._execute( query, *args, json=json, query_params=params, query_id=query_id ), asyncio.get_event_loop()): return None
[docs] def fetch( self, query: str, *args, json: bool = False, params: Optional[Dict[str, Any]] = None, query_id: str = None, decode: bool = True, ) -> List[Record]: """Execute query and fetch all rows from query result at once in a list. :param query: MyScale query string. :param bool json: Execute query in JSONEachRow mode. :param Optional[Dict[str, Any]] params: Params to escape inside query string. :param str query_id: MyScale query_id. :param decode: Decode to python types. If False, returns bytes for each field instead. Usage: .. code-block:: python all_rows = client.fetch("SELECT * FROM t") :return: All rows from query. """ return [ row for row in iterate_async_to_sync(self._execute( query, *args, json=json, query_params=params, query_id=query_id, decode=decode, ), asyncio.get_event_loop()) ]
[docs] def fetchrow( self, query: str, *args, json: bool = False, params: Optional[Dict[str, Any]] = None, query_id: str = None, decode: bool = True, ) -> Optional[Record]: """Execute query and fetch first row from query result or None. :param query: MyScale query string. :param bool json: Execute query in JSONEachRow mode. :param Optional[Dict[str, Any]] params: Params to escape inside query string. :param str query_id: MyScale query_id. :param decode: Decode to python types. If False, returns bytes for each field instead. Usage: .. code-block:: python row = client.fetchrow("SELECT * FROM t WHERE a=1") assert row[0] == 1 assert row["b"] == (dt.date(2018, 9, 7), None) :return: First row from query or None if there no results. """ for row in iterate_async_to_sync(self._execute( query, *args, json=json, query_params=params, query_id=query_id, decode=decode, ), asyncio.get_event_loop()): return row return None
[docs] def fetchval( self, query: str, *args, json: bool = False, params: Optional[Dict[str, Any]] = None, query_id: str = None, decode: bool = True, ) -> Any: """Execute query and fetch first value of the first row from query result or None. :param query: MyScale query string. :param bool json: Execute query in JSONEachRow mode. :param Optional[Dict[str, Any]] params: Params to escape inside query string. :param str query_id: MyScale query_id. :param decode: Decode to python types. If False, returns bytes for each field instead. Usage: .. code-block:: python val = client.fetchval("SELECT b FROM t WHERE a=2") assert val == (dt.date(2018, 9, 8), 3.14) :return: First value of the first row or None if there no results. """ for row in iterate_async_to_sync(self._execute( query, *args, json=json, query_params=params, query_id=query_id, decode=decode, ), asyncio.get_event_loop()): if row: return row[0] return None
[docs] def iterate( self, query: str, *args, json: bool = False, params: Optional[Dict[str, Any]] = None, query_id: str = None, decode: bool = True, ) -> Iterator[Record]: """Async generator by all rows from query result. :param str query: MyScale query string. :param bool json: Execute query in JSONEachRow mode. :param Optional[Dict[str, Any]] params: Params to escape inside query string. :param str query_id: MyScale query_id. :param decode: Decode to python types. If False, returns bytes for each field instead. Usage: .. code-block:: python for row in client.iterate( "SELECT number, number*2 FROM system.numbers LIMIT 10000" ): assert row[0] * 2 == row[1] for row in client.iterate( "SELECT number, number*2 FROM system.numbers LIMIT {numbers_limit}", params={"numbers_limit": 10000} ): assert row[0] * 2 == row[1] :return: Rows one by one. """ for row in iterate_async_to_sync(self._execute( query, *args, json=json, query_params=params, query_id=query_id, decode=decode, ), asyncio.get_event_loop()): yield row