mirror of
https://github.com/hanzo-ds/connect.git
synced 2026-10-30 20:31:24 +00:00
Joe SpadolaandGitHub
55dcd7db4c
* migrate from pylint to ruff * ruff format/lint switch over * add formatting commit * update changelog
64 lines
2.8 KiB
Python
64 lines
2.8 KiB
Python
from collections.abc import Callable
|
|
from pathlib import Path
|
|
|
|
from clickhouse_connect.driver import Client
|
|
from clickhouse_connect.driver.tools import insert_file, insert_file_async
|
|
|
|
|
|
def test_csv_upload(param_client: Client, call, table_context: Callable, client_mode):
|
|
data_file = f"{Path(__file__).parent}/movies.csv.gz"
|
|
with table_context("test_csv_upload", ["movie String", "year UInt16", "rating Decimal32(3)"]):
|
|
# Use appropriate insert_file version based on client type
|
|
if client_mode == "async":
|
|
call(
|
|
insert_file_async,
|
|
param_client,
|
|
"test_csv_upload",
|
|
data_file,
|
|
settings={"input_format_allow_errors_ratio": 0.2, "input_format_allow_errors_num": 5},
|
|
)
|
|
else: # Sync client
|
|
insert_file(
|
|
param_client,
|
|
"test_csv_upload",
|
|
data_file,
|
|
settings={"input_format_allow_errors_ratio": 0.2, "input_format_allow_errors_num": 5},
|
|
)
|
|
res = call(param_client.query, "SELECT count() as count, sum(rating) as rating, max(year) as year FROM test_csv_upload").first_item
|
|
assert res["count"] == 248
|
|
assert res["year"] == 2022
|
|
|
|
|
|
def test_parquet_upload(param_client: Client, call, client_mode, table_context: Callable):
|
|
data_file = f"{Path(__file__).parent}/movies.parquet"
|
|
with table_context("test_parquet_upload", ["movie String", "year UInt16", "rating Float64"]):
|
|
if client_mode == "async":
|
|
call(
|
|
insert_file_async,
|
|
param_client,
|
|
"test_parquet_upload",
|
|
data_file,
|
|
"Parquet",
|
|
settings={"output_format_parquet_string_as_string": 1},
|
|
)
|
|
else:
|
|
insert_file(param_client, "test_parquet_upload", data_file, "Parquet", settings={"output_format_parquet_string_as_string": 1})
|
|
res = call(
|
|
param_client.query, "SELECT count() as count, sum(rating) as rating, max(year) as year FROM test_parquet_upload"
|
|
).first_item
|
|
assert res["count"] == 250
|
|
assert res["year"] == 2022
|
|
|
|
|
|
def test_json_insert(param_client: Client, call, client_mode, table_context: Callable):
|
|
data_file = f"{Path(__file__).parent}/json_test.ndjson"
|
|
with table_context("test_json_upload", ["key UInt16", "flt_val Float64", "int_val Int8"]):
|
|
if client_mode == "async":
|
|
call(insert_file_async, param_client, "test_json_upload", data_file, "JSONEachRow")
|
|
else:
|
|
insert_file(param_client, "test_json_upload", data_file, "JSONEachRow")
|
|
res = call(param_client.query, "SELECT * FROM test_json_upload ORDER BY key").result_rows
|
|
assert res[1][0] == 17
|
|
assert res[1][1] == 5.3
|
|
assert res[1][2] == 121
|