mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-21 13:44:15 +08:00
6344ca1dd4
Since we are on 3.9, we can use this nice str builtin which is more readable and more efficient. Pull Request resolved: https://github.com/pytorch/pytorch/pull/146997 Approved by: https://github.com/XuehaiPan, https://github.com/cyyever, https://github.com/jansel
41 lines
1.2 KiB
Python
41 lines
1.2 KiB
Python
import json
|
|
import os
|
|
from functools import lru_cache
|
|
from typing import Any
|
|
|
|
import clickhouse_connect # type: ignore[import]
|
|
|
|
|
|
@lru_cache(maxsize=1)
|
|
def get_clickhouse_client() -> Any:
|
|
endpoint = os.environ["CLICKHOUSE_ENDPOINT"]
|
|
# I cannot figure out why these values aren't being handled automatically
|
|
# when it is fine in the lambda
|
|
endpoint = endpoint.removeprefix("https://")
|
|
if endpoint.endswith(":8443"):
|
|
endpoint = endpoint[: -len(":8443")]
|
|
return clickhouse_connect.get_client(
|
|
host=endpoint,
|
|
user=os.environ["CLICKHOUSE_USERNAME"],
|
|
password=os.environ["CLICKHOUSE_PASSWORD"],
|
|
secure=True,
|
|
interface="https",
|
|
port=8443,
|
|
)
|
|
|
|
|
|
def query_clickhouse(query: str, params: dict[str, Any]) -> list[dict[str, Any]]:
|
|
"""
|
|
Queries ClickHouse. Returns datetime in YYYY-MM-DD HH:MM:SS format.
|
|
"""
|
|
|
|
def convert_to_json_list(res: bytes) -> list[dict[str, Any]]:
|
|
rows = []
|
|
for row in res.decode().split("\n"):
|
|
if row:
|
|
rows.append(json.loads(row))
|
|
return rows
|
|
|
|
res = get_clickhouse_client().raw_query(query, params, fmt="JSONEachRow")
|
|
return convert_to_json_list(res)
|