ch_query.py

import lzma

import requests


DEFAULT_CREDS = {
  "url": "http://localhost:8123/",
  "database": None,
  "user": None,
  "password": None,
}


def ch_query(
    query: str,
    params: Optional[Dict[str, Any]] = None,
    creds: Optional[Dict[str, Any]] = None
) -> pd.DataFrame:
    """Send SELECT query to clickhouse and return result as pandas.DataFrame.

    Example:

    >>> ch_query("SELECT * FROM table WHERE int_column = {id:UInt8} and string_column = {phrase:String}",
    ...          params={"id": 1, "phrase": "Hello, world!"})
    """
    if creds is None:
        creds = DEFAULT_CREDS
    request_params = {
        "query": query,
        "enable_http_compression": 1,
    }
    if creds.get("database") is not None:
        request_params["database"] = creds["database"]
    if params is not None:
        for k, v in params.items():
            request_params[f"param_{k}"] = v
    headers = {
        'Accept-Encoding': 'xz',
        'X-ClickHouse-Format': 'Parquet',
    }
    if creds.get("user") is not None:
        headers['X-ClickHouse-User'] = creds['user']
    if creds.get("password") is not None:
        headers['X-ClickHouse-Key'] = creds['password']
    response = requests.post(creds["url"], params=request_params, headers=headers)
    data = response.content
    if response.headers.get("Content-Encoding") == 'xz':
        data = lzma.decompress(data)
    if response.status_code == 200:
        return pd.read_parquet(BytesIO(data))
    else:
        response.raise_for_status()
        raise RuntimeError("got %d response: %s" % (response.status_code, data))