mirror of
https://github.com/hanzo-ds/connect.git
synced 2026-10-30 20:31:24 +00:00
zeekay
75dd1132bd
The wire protocol was the last holdout and it was not cosmetic: the Hanzo Datastore server reads X-Datastore-User, and answers a request bearing X-ClickHouse-User with 401. Verified against the live server, both directions. So a client still sending the old headers is not merely off-brand, it cannot authenticate over HTTP. Everything else follows: types, dialects, error modules, docstrings and doc links now name Datastore, and doc URLs point at docs.hanzo.ai/datastore rather than the upstream vendor's site. Zero occurrences of the old name remain. Still cut from the pinned upstream version, so behaviour is unchanged.
728 lines
42 KiB
Python
728 lines
42 KiB
Python
import asyncio
|
|
import io
|
|
import os
|
|
from concurrent.futures.thread import ThreadPoolExecutor
|
|
from datetime import tzinfo
|
|
from typing import Optional, Union, Dict, Any, Sequence, Iterable, Generator, BinaryIO
|
|
|
|
from datastore_connect.driver.client import Client
|
|
from datastore_connect.driver.common import StreamContext
|
|
from datastore_connect.driver.httpclient import HttpClient
|
|
from datastore_connect.driver.external import ExternalData
|
|
from datastore_connect.driver.query import QueryContext, QueryResult
|
|
from datastore_connect.driver.summary import QuerySummary
|
|
from datastore_connect.datatypes.base import DatastoreType
|
|
from datastore_connect.driver.insert import InsertContext
|
|
|
|
|
|
# pylint: disable=too-many-public-methods,too-many-instance-attributes,too-many-arguments,too-many-positional-arguments,too-many-locals
|
|
class AsyncClient:
|
|
"""
|
|
AsyncClient is a wrapper around the Datastore Client object that allows for async calls to the Datastore server.
|
|
Internally, each of the methods that uses IO is wrapped in a call to EventLoop.run_in_executor.
|
|
"""
|
|
|
|
def __init__(self, *, client: Client, executor_threads: int = 0):
|
|
if isinstance(client, HttpClient):
|
|
client.headers['User-Agent'] = client.headers['User-Agent'].replace('mode:sync;', 'mode:async;')
|
|
self.client = client
|
|
if executor_threads == 0:
|
|
executor_threads = min(32, (os.cpu_count() or 1) + 4) # Mimic the default behavior
|
|
self.executor = ThreadPoolExecutor(max_workers=executor_threads)
|
|
|
|
def set_client_setting(self, key, value):
|
|
"""
|
|
Set a datastore setting for the client after initialization. If a setting is not recognized by Datastore,
|
|
or the setting is identified as "read_only", this call will either throw a Programming exception or attempt
|
|
to send the setting anyway based on the common setting 'invalid_setting_action'.
|
|
:param key: Datastore setting name
|
|
:param value: Datastore setting value
|
|
"""
|
|
self.client.set_client_setting(key=key, value=value)
|
|
|
|
def get_client_setting(self, key) -> Optional[str]:
|
|
"""
|
|
:param key: The setting key
|
|
:return: The string value of the setting, if it exists, or None
|
|
"""
|
|
return self.client.get_client_setting(key=key)
|
|
|
|
def set_access_token(self, access_token: str):
|
|
"""
|
|
Set the Datastore access token for the client
|
|
:param access_token: Access token string
|
|
"""
|
|
self.client.set_access_token(access_token)
|
|
|
|
def min_version(self, version_str: str) -> bool:
|
|
"""
|
|
Determine whether the connected server is at least the submitted version
|
|
For Altinity Stable versions like 22.8.15.25.altinitystable
|
|
the last condition in the first list comprehension expression is added
|
|
:param version_str: A version string consisting of up to 4 integers delimited by dots
|
|
:return: True if version_str is greater than the server_version, False if less than
|
|
"""
|
|
return self.client.min_version(version_str)
|
|
|
|
async def close(self):
|
|
"""
|
|
Subclass implementation to close the connection to the server/deallocate the client
|
|
"""
|
|
self.client.close()
|
|
await asyncio.to_thread(self.executor.shutdown, True)
|
|
|
|
async def query(self,
|
|
query: Optional[str] = None,
|
|
parameters: Optional[Union[Sequence, Dict[str, Any]]] = None,
|
|
settings: Optional[Dict[str, Any]] = None,
|
|
query_formats: Optional[Dict[str, str]] = None,
|
|
column_formats: Optional[Dict[str, Union[str, Dict[str, str]]]] = None,
|
|
encoding: Optional[str] = None,
|
|
use_none: Optional[bool] = None,
|
|
column_oriented: Optional[bool] = None,
|
|
use_numpy: Optional[bool] = None,
|
|
max_str_len: Optional[int] = None,
|
|
context: QueryContext = None,
|
|
query_tz: Optional[Union[str, tzinfo]] = None,
|
|
column_tzs: Optional[Dict[str, Union[str, tzinfo]]] = None,
|
|
external_data: Optional[ExternalData] = None,
|
|
transport_settings: Optional[Dict[str, str]] = None) -> QueryResult:
|
|
"""
|
|
Main query method for SELECT, DESCRIBE and other SQL statements that return a result matrix.
|
|
For parameters, see the create_query_context method.
|
|
:return: QueryResult -- data and metadata from response
|
|
"""
|
|
|
|
def _query():
|
|
return self.client.query(query=query, parameters=parameters, settings=settings, query_formats=query_formats,
|
|
column_formats=column_formats, encoding=encoding, use_none=use_none,
|
|
column_oriented=column_oriented, use_numpy=use_numpy, max_str_len=max_str_len,
|
|
context=context, query_tz=query_tz, column_tzs=column_tzs,
|
|
external_data=external_data, transport_settings=transport_settings)
|
|
|
|
loop = asyncio.get_running_loop()
|
|
result = await loop.run_in_executor(self.executor, _query)
|
|
return result
|
|
|
|
async def query_column_block_stream(self,
|
|
query: Optional[str] = None,
|
|
parameters: Optional[Union[Sequence, Dict[str, Any]]] = None,
|
|
settings: Optional[Dict[str, Any]] = None,
|
|
query_formats: Optional[Dict[str, str]] = None,
|
|
column_formats: Optional[Dict[str, Union[str, Dict[str, str]]]] = None,
|
|
encoding: Optional[str] = None,
|
|
use_none: Optional[bool] = None,
|
|
context: QueryContext = None,
|
|
query_tz: Optional[Union[str, tzinfo]] = None,
|
|
column_tzs: Optional[Dict[str, Union[str, tzinfo]]] = None,
|
|
external_data: Optional[ExternalData] = None,
|
|
transport_settings: Optional[Dict[str, str]] = None,
|
|
) -> StreamContext:
|
|
"""
|
|
Variation of main query method that returns a stream of column oriented blocks.
|
|
For parameters, see the create_query_context method.
|
|
:return: StreamContext -- Iterable stream context that returns column oriented blocks
|
|
"""
|
|
|
|
def _query_column_block_stream():
|
|
return self.client.query_column_block_stream(query=query, parameters=parameters, settings=settings,
|
|
query_formats=query_formats, column_formats=column_formats,
|
|
encoding=encoding, use_none=use_none, context=context,
|
|
query_tz=query_tz, column_tzs=column_tzs,
|
|
external_data=external_data, transport_settings=transport_settings)
|
|
|
|
loop = asyncio.get_running_loop()
|
|
result = await loop.run_in_executor(self.executor, _query_column_block_stream)
|
|
return result
|
|
|
|
async def query_row_block_stream(self,
|
|
query: Optional[str] = None,
|
|
parameters: Optional[Union[Sequence, Dict[str, Any]]] = None,
|
|
settings: Optional[Dict[str, Any]] = None,
|
|
query_formats: Optional[Dict[str, str]] = None,
|
|
column_formats: Optional[Dict[str, Union[str, Dict[str, str]]]] = None,
|
|
encoding: Optional[str] = None,
|
|
use_none: Optional[bool] = None,
|
|
context: QueryContext = None,
|
|
query_tz: Optional[Union[str, tzinfo]] = None,
|
|
column_tzs: Optional[Dict[str, Union[str, tzinfo]]] = None,
|
|
external_data: Optional[ExternalData] = None,
|
|
transport_settings: Optional[Dict[str, str]] = None) -> StreamContext:
|
|
"""
|
|
Variation of main query method that returns a stream of row oriented blocks.
|
|
For parameters, see the create_query_context method.
|
|
:return: StreamContext -- Iterable stream context that returns blocks of rows
|
|
"""
|
|
|
|
def _query_row_block_stream():
|
|
return self.client.query_row_block_stream(query=query, parameters=parameters, settings=settings,
|
|
query_formats=query_formats, column_formats=column_formats,
|
|
encoding=encoding, use_none=use_none, context=context,
|
|
query_tz=query_tz, column_tzs=column_tzs,
|
|
external_data=external_data, transport_settings=transport_settings)
|
|
|
|
loop = asyncio.get_running_loop()
|
|
result = await loop.run_in_executor(self.executor, _query_row_block_stream)
|
|
return result
|
|
|
|
async def query_rows_stream(self,
|
|
query: Optional[str] = None,
|
|
parameters: Optional[Union[Sequence, Dict[str, Any]]] = None,
|
|
settings: Optional[Dict[str, Any]] = None,
|
|
query_formats: Optional[Dict[str, str]] = None,
|
|
column_formats: Optional[Dict[str, Union[str, Dict[str, str]]]] = None,
|
|
encoding: Optional[str] = None,
|
|
use_none: Optional[bool] = None,
|
|
context: QueryContext = None,
|
|
query_tz: Optional[Union[str, tzinfo]] = None,
|
|
column_tzs: Optional[Dict[str, Union[str, tzinfo]]] = None,
|
|
external_data: Optional[ExternalData] = None,
|
|
transport_settings: Optional[Dict[str, str]] = None) -> StreamContext:
|
|
"""
|
|
Variation of main query method that returns a stream of row oriented blocks.
|
|
For parameters, see the create_query_context method.
|
|
:return: StreamContext -- Iterable stream context that returns blocks of rows
|
|
"""
|
|
|
|
def _query_rows_stream():
|
|
return self.client.query_rows_stream(query=query, parameters=parameters, settings=settings,
|
|
query_formats=query_formats, column_formats=column_formats,
|
|
encoding=encoding, use_none=use_none, context=context,
|
|
query_tz=query_tz, column_tzs=column_tzs,
|
|
external_data=external_data, transport_settings=transport_settings)
|
|
|
|
loop = asyncio.get_running_loop()
|
|
result = await loop.run_in_executor(self.executor, _query_rows_stream)
|
|
return result
|
|
|
|
async def raw_query(self,
|
|
query: str,
|
|
parameters: Optional[Union[Sequence, Dict[str, Any]]] = None,
|
|
settings: Optional[Dict[str, Any]] = None,
|
|
fmt: str = None,
|
|
use_database: bool = True,
|
|
external_data: Optional[ExternalData] = None,
|
|
transport_settings: Optional[Dict[str, str]] = None) -> bytes:
|
|
"""
|
|
Query method that simply returns the raw Datastore format bytes.
|
|
:param query: Query statement/format string
|
|
:param parameters: Optional dictionary used to format the query
|
|
:param settings: Optional dictionary of Datastore settings (key/string values)
|
|
:param fmt: Datastore output format
|
|
:param use_database Send the database parameter to Datastore so the command will be executed in the client
|
|
database context
|
|
:param external_data External data to send with the query
|
|
:param transport_settings: Optional dictionary of transport level settings (HTTP headers, etc.)
|
|
:return: bytes representing raw Datastore return value based on format
|
|
"""
|
|
|
|
def _raw_query():
|
|
return self.client.raw_query(query=query, parameters=parameters, settings=settings, fmt=fmt,
|
|
use_database=use_database, external_data=external_data,
|
|
transport_settings=transport_settings)
|
|
|
|
loop = asyncio.get_running_loop()
|
|
result = await loop.run_in_executor(self.executor, _raw_query)
|
|
return result
|
|
|
|
async def raw_stream(self, query: str,
|
|
parameters: Optional[Union[Sequence, Dict[str, Any]]] = None,
|
|
settings: Optional[Dict[str, Any]] = None,
|
|
fmt: str = None,
|
|
use_database: bool = True,
|
|
external_data: Optional[ExternalData] = None,
|
|
transport_settings: Optional[Dict[str, str]] = None) -> io.IOBase:
|
|
"""
|
|
Query method that returns the result as an io.IOBase iterator.
|
|
:param query: Query statement/format string
|
|
:param parameters: Optional dictionary used to format the query
|
|
:param settings: Optional dictionary of Datastore settings (key/string values)
|
|
:param fmt: Datastore output format
|
|
:param use_database Send the database parameter to Datastore so the command will be executed in the client
|
|
database context
|
|
:param external_data External data to send with the query
|
|
:param transport_settings: Optional dictionary of transport level settings (HTTP headers, etc.)
|
|
:return: io.IOBase stream/iterator for the result
|
|
"""
|
|
|
|
def _raw_stream():
|
|
return self.client.raw_stream(query=query, parameters=parameters, settings=settings, fmt=fmt,
|
|
use_database=use_database, external_data=external_data, transport_settings=transport_settings)
|
|
|
|
loop = asyncio.get_running_loop()
|
|
result = await loop.run_in_executor(self.executor, _raw_stream)
|
|
return result
|
|
|
|
async def query_np(self,
|
|
query: Optional[str] = None,
|
|
parameters: Optional[Union[Sequence, Dict[str, Any]]] = None,
|
|
settings: Optional[Dict[str, Any]] = None,
|
|
query_formats: Optional[Dict[str, str]] = None,
|
|
column_formats: Optional[Dict[str, str]] = None,
|
|
encoding: Optional[str] = None,
|
|
use_none: Optional[bool] = None,
|
|
max_str_len: Optional[int] = None,
|
|
context: QueryContext = None,
|
|
external_data: Optional[ExternalData] = None,
|
|
transport_settings: Optional[Dict[str, str]] = None):
|
|
"""
|
|
Query method that returns the results as a numpy array.
|
|
For parameter values, see the create_query_context method.
|
|
:return: Numpy array representing the result set
|
|
"""
|
|
|
|
def _query_np():
|
|
return self.client.query_np(query=query, parameters=parameters, settings=settings,
|
|
query_formats=query_formats, column_formats=column_formats, encoding=encoding,
|
|
use_none=use_none, max_str_len=max_str_len, context=context,
|
|
external_data=external_data, transport_settings=transport_settings)
|
|
|
|
loop = asyncio.get_running_loop()
|
|
result = await loop.run_in_executor(self.executor, _query_np)
|
|
return result
|
|
|
|
async def query_np_stream(self,
|
|
query: Optional[str] = None,
|
|
parameters: Optional[Union[Sequence, Dict[str, Any]]] = None,
|
|
settings: Optional[Dict[str, Any]] = None,
|
|
query_formats: Optional[Dict[str, str]] = None,
|
|
column_formats: Optional[Dict[str, str]] = None,
|
|
encoding: Optional[str] = None,
|
|
use_none: Optional[bool] = None,
|
|
max_str_len: Optional[int] = None,
|
|
context: QueryContext = None,
|
|
external_data: Optional[ExternalData] = None,
|
|
transport_settings: Optional[Dict[str, str]] = None) -> StreamContext:
|
|
"""
|
|
Query method that returns the results as a stream of numpy arrays.
|
|
For parameter values, see the create_query_context method.
|
|
:return: Generator that yield a numpy array per block representing the result set
|
|
"""
|
|
|
|
def _query_np_stream():
|
|
return self.client.query_np_stream(query=query, parameters=parameters, settings=settings,
|
|
query_formats=query_formats, column_formats=column_formats,
|
|
encoding=encoding, use_none=use_none, max_str_len=max_str_len,
|
|
context=context, external_data=external_data, transport_settings=transport_settings)
|
|
|
|
loop = asyncio.get_running_loop()
|
|
result = await loop.run_in_executor(self.executor, _query_np_stream)
|
|
return result
|
|
|
|
async def query_df(self,
|
|
query: Optional[str] = None,
|
|
parameters: Optional[Union[Sequence, Dict[str, Any]]] = None,
|
|
settings: Optional[Dict[str, Any]] = None,
|
|
query_formats: Optional[Dict[str, str]] = None,
|
|
column_formats: Optional[Dict[str, str]] = None,
|
|
encoding: Optional[str] = None,
|
|
use_none: Optional[bool] = None,
|
|
max_str_len: Optional[int] = None,
|
|
use_na_values: Optional[bool] = None,
|
|
query_tz: Optional[str] = None,
|
|
column_tzs: Optional[Dict[str, Union[str, tzinfo]]] = None,
|
|
context: QueryContext = None,
|
|
external_data: Optional[ExternalData] = None,
|
|
use_extended_dtypes: Optional[bool] = None,
|
|
transport_settings: Optional[Dict[str, str]] = None):
|
|
"""
|
|
Query method that results the results as a pandas dataframe.
|
|
For parameter values, see the create_query_context method.
|
|
:return: Pandas dataframe representing the result set
|
|
"""
|
|
|
|
def _query_df():
|
|
return self.client.query_df(query=query, parameters=parameters, settings=settings,
|
|
query_formats=query_formats, column_formats=column_formats, encoding=encoding,
|
|
use_none=use_none, max_str_len=max_str_len, use_na_values=use_na_values,
|
|
query_tz=query_tz, column_tzs=column_tzs, context=context,
|
|
external_data=external_data, use_extended_dtypes=use_extended_dtypes,
|
|
transport_settings=transport_settings)
|
|
|
|
loop = asyncio.get_running_loop()
|
|
result = await loop.run_in_executor(self.executor, _query_df)
|
|
return result
|
|
|
|
async def query_df_stream(self,
|
|
query: Optional[str] = None,
|
|
parameters: Optional[Union[Sequence, Dict[str, Any]]] = None,
|
|
settings: Optional[Dict[str, Any]] = None,
|
|
query_formats: Optional[Dict[str, str]] = None,
|
|
column_formats: Optional[Dict[str, str]] = None,
|
|
encoding: Optional[str] = None,
|
|
use_none: Optional[bool] = None,
|
|
max_str_len: Optional[int] = None,
|
|
use_na_values: Optional[bool] = None,
|
|
query_tz: Optional[str] = None,
|
|
column_tzs: Optional[Dict[str, Union[str, tzinfo]]] = None,
|
|
context: QueryContext = None,
|
|
external_data: Optional[ExternalData] = None,
|
|
use_extended_dtypes: Optional[bool] = None,
|
|
transport_settings: Optional[Dict[str, str]] = None) -> StreamContext:
|
|
"""
|
|
Query method that returns the results as a StreamContext.
|
|
For parameter values, see the create_query_context method.
|
|
:return: Generator that yields a Pandas dataframe per block representing the result set
|
|
"""
|
|
|
|
def _query_df_stream():
|
|
return self.client.query_df_stream(query=query, parameters=parameters, settings=settings,
|
|
query_formats=query_formats, column_formats=column_formats,
|
|
encoding=encoding,
|
|
use_none=use_none, max_str_len=max_str_len, use_na_values=use_na_values,
|
|
query_tz=query_tz, column_tzs=column_tzs, context=context,
|
|
external_data=external_data, use_extended_dtypes=use_extended_dtypes,
|
|
transport_settings=transport_settings)
|
|
|
|
loop = asyncio.get_running_loop()
|
|
result = await loop.run_in_executor(self.executor, _query_df_stream)
|
|
return result
|
|
|
|
def create_query_context(self,
|
|
query: Optional[Union[str, bytes]] = None,
|
|
parameters: Optional[Union[Sequence, Dict[str, Any]]] = None,
|
|
settings: Optional[Dict[str, Any]] = None,
|
|
query_formats: Optional[Dict[str, str]] = None,
|
|
column_formats: Optional[Dict[str, Union[str, Dict[str, str]]]] = None,
|
|
encoding: Optional[str] = None,
|
|
use_none: Optional[bool] = None,
|
|
column_oriented: Optional[bool] = None,
|
|
use_numpy: Optional[bool] = False,
|
|
max_str_len: Optional[int] = 0,
|
|
context: Optional[QueryContext] = None,
|
|
query_tz: Optional[Union[str, tzinfo]] = None,
|
|
column_tzs: Optional[Dict[str, Union[str, tzinfo]]] = None,
|
|
use_na_values: Optional[bool] = None,
|
|
streaming: bool = False,
|
|
as_pandas: bool = False,
|
|
external_data: Optional[ExternalData] = None,
|
|
use_extended_dtypes: Optional[bool] = None,
|
|
transport_settings: Optional[Dict[str, str]] = None) -> QueryContext:
|
|
"""
|
|
Creates or updates a reusable QueryContext object
|
|
:param query: Query statement/format string
|
|
:param parameters: Optional dictionary used to format the query
|
|
:param settings: Optional dictionary of Datastore settings (key/string values)
|
|
:param query_formats: See QueryContext __init__ docstring
|
|
:param column_formats: See QueryContext __init__ docstring
|
|
:param encoding: See QueryContext __init__ docstring
|
|
:param use_none: Use None for Datastore NULL instead of default values. Note that using None in Numpy
|
|
arrays will force the numpy array dtype to 'object', which is often inefficient. This effect also
|
|
will impact the performance of Pandas dataframes.
|
|
:param column_oriented: Deprecated. Controls orientation of the QueryResult result_set property
|
|
:param use_numpy: Return QueryResult columns as one-dimensional numpy arrays
|
|
:param max_str_len: Limit returned Datastore String values to this length, which allows a Numpy
|
|
structured array even with Datastore variable length String columns. If 0, Numpy arrays for
|
|
String columns will always be object arrays
|
|
:param context: An existing QueryContext to be updated with any provided parameter values
|
|
:param query_tz Either a string or a pytz tzinfo object. (Strings will be converted to tzinfo objects).
|
|
Values for any DateTime or DateTime64 column in the query will be converted to Python datetime.datetime
|
|
objects with the selected timezone
|
|
:param column_tzs A dictionary of column names to tzinfo objects (or strings that will be converted to
|
|
tzinfo objects). The timezone will be applied to datetime objects returned in the query
|
|
:param use_na_values: Deprecated alias for use_advanced_dtypes
|
|
:param as_pandas Return the result columns as pandas.Series objects
|
|
:param streaming Marker used to correctly configure streaming queries
|
|
:param external_data Datastore "external data" to send with query
|
|
:param use_extended_dtypes: Only relevant to Pandas Dataframe queries. Use Pandas "missing types", such as
|
|
pandas.NA and pandas.NaT for Datastore NULL values, as well as extended Pandas dtypes such as IntegerArray
|
|
and StringArray. Defaulted to True for query_df methods
|
|
:param transport_settings: Optional dictionary of transport level settings (HTTP headers, etc.)
|
|
:return: Reusable QueryContext
|
|
"""
|
|
|
|
return self.client.create_query_context(query=query, parameters=parameters, settings=settings,
|
|
query_formats=query_formats, column_formats=column_formats,
|
|
encoding=encoding, use_none=use_none,
|
|
column_oriented=column_oriented,
|
|
use_numpy=use_numpy, max_str_len=max_str_len, context=context,
|
|
query_tz=query_tz, column_tzs=column_tzs,
|
|
use_na_values=use_na_values,
|
|
streaming=streaming, as_pandas=as_pandas,
|
|
external_data=external_data,
|
|
use_extended_dtypes=use_extended_dtypes,
|
|
transport_settings=transport_settings)
|
|
|
|
async def query_arrow(self,
|
|
query: str,
|
|
parameters: Optional[Union[Sequence, Dict[str, Any]]] = None,
|
|
settings: Optional[Dict[str, Any]] = None,
|
|
use_strings: Optional[bool] = None,
|
|
external_data: Optional[ExternalData] = None,
|
|
transport_settings: Optional[Dict[str, str]] = None):
|
|
"""
|
|
Query method using the Datastore Arrow format to return a PyArrow table
|
|
:param query: Query statement/format string
|
|
:param parameters: Optional dictionary used to format the query
|
|
:param settings: Optional dictionary of Datastore settings (key/string values)
|
|
:param use_strings: Convert Datastore String type to Arrow string type (instead of binary)
|
|
:param external_data Datastore "external data" to send with query
|
|
:param transport_settings: Optional dictionary of transport level settings (HTTP headers, etc.)
|
|
:return: PyArrow.Table
|
|
"""
|
|
|
|
def _query_arrow():
|
|
return self.client.query_arrow(query=query, parameters=parameters, settings=settings,
|
|
use_strings=use_strings, external_data=external_data,
|
|
transport_settings=transport_settings)
|
|
|
|
loop = asyncio.get_running_loop()
|
|
result = await loop.run_in_executor(self.executor, _query_arrow)
|
|
return result
|
|
|
|
async def query_arrow_stream(self,
|
|
query: str,
|
|
parameters: Optional[Union[Sequence, Dict[str, Any]]] = None,
|
|
settings: Optional[Dict[str, Any]] = None,
|
|
use_strings: Optional[bool] = None,
|
|
external_data: Optional[ExternalData] = None,
|
|
transport_settings: Optional[Dict[str, str]] = None) -> StreamContext:
|
|
"""
|
|
Query method that returns the results as a stream of Arrow tables
|
|
:param query: Query statement/format string
|
|
:param parameters: Optional dictionary used to format the query
|
|
:param settings: Optional dictionary of Datastore settings (key/string values)
|
|
:param use_strings: Convert Datastore String type to Arrow string type (instead of binary)
|
|
:param external_data Datastore "external data" to send with query
|
|
:param transport_settings: Optional dictionary of transport level settings (HTTP headers, etc.)
|
|
:return: Generator that yields a PyArrow.Table for per block representing the result set
|
|
"""
|
|
|
|
def _query_arrow_stream():
|
|
return self.client.query_arrow_stream(query=query, parameters=parameters, settings=settings,
|
|
use_strings=use_strings, external_data=external_data,
|
|
transport_settings=transport_settings)
|
|
|
|
loop = asyncio.get_running_loop()
|
|
result = await loop.run_in_executor(self.executor, _query_arrow_stream)
|
|
return result
|
|
|
|
async def command(self,
|
|
cmd: str,
|
|
parameters: Optional[Union[Sequence, Dict[str, Any]]] = None,
|
|
data: Union[str, bytes] = None,
|
|
settings: Dict[str, Any] = None,
|
|
use_database: bool = True,
|
|
external_data: Optional[ExternalData] = None,
|
|
transport_settings: Optional[Dict[str, str]] = None) -> Union[str, int, Sequence[str], QuerySummary]:
|
|
"""
|
|
Client method that returns a single value instead of a result set
|
|
:param cmd: Datastore query/command as a python format string
|
|
:param parameters: Optional dictionary of key/values pairs to be formatted
|
|
:param data: Optional 'data' for the command (for INSERT INTO in particular)
|
|
:param settings: Optional dictionary of Datastore settings (key/string values)
|
|
:param use_database: Send the database parameter to Datastore so the command will be executed in the client
|
|
database context. Otherwise, no database will be specified with the command. This is useful for determining
|
|
the default user database
|
|
:param external_data Datastore "external data" to send with command/query
|
|
:param transport_settings: Optional dictionary of transport level settings (HTTP headers, etc.)
|
|
:return: Decoded response from Datastore as either a string, int, or sequence of strings, or QuerySummary
|
|
if no data returned
|
|
"""
|
|
|
|
def _command():
|
|
return self.client.command(cmd=cmd, parameters=parameters, data=data, settings=settings,
|
|
use_database=use_database, external_data=external_data,
|
|
transport_settings=transport_settings)
|
|
|
|
loop = asyncio.get_running_loop()
|
|
result = await loop.run_in_executor(self.executor, _command)
|
|
return result
|
|
|
|
async def ping(self) -> bool:
|
|
"""
|
|
Validate the connection, does not throw an Exception (see debug logs)
|
|
:return: Datastore server is up and reachable
|
|
"""
|
|
|
|
def _ping():
|
|
return self.client.ping()
|
|
|
|
loop = asyncio.get_running_loop()
|
|
result = await loop.run_in_executor(self.executor, _ping)
|
|
return result
|
|
|
|
async def insert(self,
|
|
table: Optional[str] = None,
|
|
data: Sequence[Sequence[Any]] = None,
|
|
column_names: Union[str, Iterable[str]] = '*',
|
|
database: Optional[str] = None,
|
|
column_types: Sequence[DatastoreType] = None,
|
|
column_type_names: Sequence[str] = None,
|
|
column_oriented: bool = False,
|
|
settings: Optional[Dict[str, Any]] = None,
|
|
context: InsertContext = None,
|
|
transport_settings: Optional[Dict[str, str]] = None) -> QuerySummary:
|
|
"""
|
|
Method to insert multiple rows/data matrix of native Python objects. If context is specified arguments
|
|
other than data are ignored
|
|
:param table: Target table
|
|
:param data: Sequence of sequences of Python data
|
|
:param column_names: Ordered list of column names or '*' if column types should be retrieved from the
|
|
Datastore table definition
|
|
:param database: Target database -- will use client default database if not specified.
|
|
:param column_types: Datastore column types. If set then column data does not need to be retrieved from
|
|
the server
|
|
:param column_type_names: Datastore column type names. If set then column data does not need to be
|
|
retrieved from the server
|
|
:param column_oriented: If true the data is already "pivoted" in column form
|
|
:param settings: Optional dictionary of Datastore settings (key/string values)
|
|
:param context: Optional reusable insert context to allow repeated inserts into the same table with
|
|
different data batches
|
|
:param transport_settings: Optional dictionary of transport level settings (HTTP headers, etc.)
|
|
:return: QuerySummary with summary information, throws exception if insert fails
|
|
"""
|
|
|
|
def _insert():
|
|
return self.client.insert(table=table, data=data, column_names=column_names, database=database,
|
|
column_types=column_types, column_type_names=column_type_names,
|
|
column_oriented=column_oriented, settings=settings, context=context,
|
|
transport_settings=transport_settings)
|
|
|
|
loop = asyncio.get_running_loop()
|
|
result = await loop.run_in_executor(self.executor, _insert)
|
|
return result
|
|
|
|
async def insert_df(self, table: str = None,
|
|
df=None,
|
|
database: Optional[str] = None,
|
|
settings: Optional[Dict] = None,
|
|
column_names: Optional[Sequence[str]] = None,
|
|
column_types: Sequence[DatastoreType] = None,
|
|
column_type_names: Sequence[str] = None,
|
|
context: InsertContext = None,
|
|
transport_settings: Optional[Dict[str, str]] = None) -> QuerySummary:
|
|
"""
|
|
Insert a pandas DataFrame into Datastore. If context is specified arguments other than df are ignored
|
|
:param table: Datastore table
|
|
:param df: two-dimensional pandas dataframe
|
|
:param database: Optional Datastore database
|
|
:param settings: Optional dictionary of Datastore settings (key/string values)
|
|
:param column_names: An optional list of Datastore column names. If not set, the DataFrame column names
|
|
will be used
|
|
:param column_types: Datastore column types. If set then column data does not need to be retrieved from
|
|
the server
|
|
:param column_type_names: Datastore column type names. If set then column data does not need to be
|
|
retrieved from the server
|
|
:param context: Optional reusable insert context to allow repeated inserts into the same table with
|
|
different data batches
|
|
:param transport_settings: Optional dictionary of transport level settings (HTTP headers, etc.)
|
|
:return: QuerySummary with summary information, throws exception if insert fails
|
|
"""
|
|
|
|
def _insert_df():
|
|
return self.client.insert_df(table=table, df=df, database=database, settings=settings,
|
|
column_names=column_names,
|
|
column_types=column_types, column_type_names=column_type_names,
|
|
context=context, transport_settings=transport_settings)
|
|
|
|
loop = asyncio.get_running_loop()
|
|
result = await loop.run_in_executor(self.executor, _insert_df)
|
|
return result
|
|
|
|
async def insert_arrow(self, table: str,
|
|
arrow_table, database: str = None,
|
|
settings: Optional[Dict] = None,
|
|
transport_settings: Optional[Dict[str, str]] = None) -> QuerySummary:
|
|
"""
|
|
Insert a PyArrow table DataFrame into Datastore using raw Arrow format
|
|
:param table: Datastore table
|
|
:param arrow_table: PyArrow Table object
|
|
:param database: Optional Datastore database
|
|
:param settings: Optional dictionary of Datastore settings (key/string values)
|
|
:param transport_settings: Optional dictionary of transport level settings (HTTP headers, etc.)
|
|
:return: QuerySummary with summary information, throws exception if insert fails
|
|
"""
|
|
|
|
def _insert_arrow():
|
|
return self.client.insert_arrow(table=table, arrow_table=arrow_table, database=database,
|
|
settings=settings, transport_settings=transport_settings)
|
|
|
|
loop = asyncio.get_running_loop()
|
|
result = await loop.run_in_executor(self.executor, _insert_arrow)
|
|
return result
|
|
|
|
async def create_insert_context(self,
|
|
table: str,
|
|
column_names: Optional[Union[str, Sequence[str]]] = None,
|
|
database: Optional[str] = None,
|
|
column_types: Sequence[DatastoreType] = None,
|
|
column_type_names: Sequence[str] = None,
|
|
column_oriented: bool = False,
|
|
settings: Optional[Dict[str, Any]] = None,
|
|
data: Optional[Sequence[Sequence[Any]]] = None,
|
|
transport_settings: Optional[Dict[str, str]] = None) -> InsertContext:
|
|
"""
|
|
Builds a reusable insert context to hold state for a duration of an insert
|
|
:param table: Target table
|
|
:param database: Target database. If not set, uses the client default database
|
|
:param column_names: Optional ordered list of column names. If not set, all columns ('*') will be assumed
|
|
in the order specified by the table definition
|
|
:param database: Target database -- will use client default database if not specified
|
|
:param column_types: Datastore column types. Optional Sequence of DatastoreType objects. If neither column
|
|
types nor column type names are set, actual column types will be retrieved from the server.
|
|
:param column_type_names: Datastore column type names. Specified column types by name string
|
|
:param column_oriented: If true the data is already "pivoted" in column form
|
|
:param settings: Optional dictionary of Datastore settings (key/string values)
|
|
:param data: Initial dataset for insert
|
|
:param transport_settings: Optional dictionary of transport level settings (HTTP headers, etc.)
|
|
:return: Reusable insert context
|
|
"""
|
|
|
|
def _create_insert_context():
|
|
return self.client.create_insert_context(table=table, column_names=column_names, database=database,
|
|
column_types=column_types, column_type_names=column_type_names,
|
|
column_oriented=column_oriented, settings=settings, data=data,
|
|
transport_settings=transport_settings)
|
|
|
|
loop = asyncio.get_running_loop()
|
|
result = await loop.run_in_executor(self.executor, _create_insert_context)
|
|
return result
|
|
|
|
async def data_insert(self, context: InsertContext) -> QuerySummary:
|
|
"""
|
|
Subclass implementation of the data insert
|
|
:context: InsertContext parameter object
|
|
:return: No return, throws an exception if the insert fails
|
|
"""
|
|
|
|
def _data_insert():
|
|
return self.client.data_insert(context=context)
|
|
|
|
loop = asyncio.get_running_loop()
|
|
result = await loop.run_in_executor(self.executor, _data_insert)
|
|
return result
|
|
|
|
async def raw_insert(self, table: str,
|
|
column_names: Optional[Sequence[str]] = None,
|
|
insert_block: Union[str, bytes, Generator[bytes, None, None], BinaryIO] = None,
|
|
settings: Optional[Dict] = None,
|
|
fmt: Optional[str] = None,
|
|
compression: Optional[str] = None,
|
|
transport_settings: Optional[Dict[str, str]] = None) -> QuerySummary:
|
|
"""
|
|
Insert data already formatted in a bytes object
|
|
:param table: Table name (whether qualified with the database name or not)
|
|
:param column_names: Sequence of column names
|
|
:param insert_block: Binary or string data already in a recognized Datastore format
|
|
:param settings: Optional dictionary of Datastore settings (key/string values)
|
|
:param compression: Recognized Datastore `Accept-Encoding` header compression value
|
|
:param transport_settings: Optional dictionary of transport level settings (HTTP headers, etc.)
|
|
:param fmt: Valid datastore format
|
|
"""
|
|
|
|
def _raw_insert():
|
|
return self.client.raw_insert(table=table, column_names=column_names, insert_block=insert_block,
|
|
settings=settings, fmt=fmt, compression=compression,
|
|
transport_settings=transport_settings)
|
|
|
|
loop = asyncio.get_running_loop()
|
|
result = await loop.run_in_executor(self.executor, _raw_insert)
|
|
return result
|
|
|
|
async def __aenter__(self) -> "AsyncClient":
|
|
return self
|
|
|
|
async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
|
|
await self.close()
|