mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
d48eb58d1d
Pull Request resolved: https://github.com/pytorch/pytorch/pull/145606 Approved by: https://github.com/malfet ghstack dependencies: #144546
40 lines
1.2 KiB
Python
40 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://")
|
|
endpoint = endpoint.removesuffix(":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)
|