connect/tests/integration_tests/test_native.py

import decimal
import uuid
from collections.abc import Callable
from datetime import date, datetime
from ipaddress import IPv4Address, IPv6Address

from clickhouse_connect.datatypes.format import clear_default_format, set_default_formats, set_read_format
from clickhouse_connect.driver import Client


def test_low_card(param_client: Client, call, table_context: Callable):
    with table_context("native_test", ["key LowCardinality(Int32)", "value_1 LowCardinality(String)"]):
        call(param_client.insert, "native_test", [[55, "TV1"], [-578328, "TV38882"], [57372, "Kabc/defXX"]])
        result = call(param_client.query, "SELECT * FROM native_test WHERE value_1 LIKE '%abc/def%'")
        assert len(result.result_set) == 1


def test_low_card_uuid(param_client: Client, call, table_context: Callable):
    with table_context("low_card_uuid", ["dt Date", "low_card_uuid LowCardinality(UUID)"]):
        data = ([date(2023, 1, 1), "80397B00E0B248AFAF34AE11A5546A3B"], [date(2024, 1, 1), "70397B00-E0B2-48AF-AF34-AE11A5546A3B"])
        call(param_client.insert, "low_card_uuid", data)
        result = call(param_client.query, "SELECT * FROM low_card_uuid order by dt").result_set
        assert len(result) == 2
        assert str(result[0][1]) == "80397b00-e0b2-48af-af34-ae11a5546a3b"
        assert str(result[1][1]) == "70397b00-e0b2-48af-af34-ae11a5546a3b"


def test_bare_datetime64(param_client: Client, call, table_context: Callable):
    with table_context("bare_datetime64_test", ["key UInt32", "dt64 DateTime64"]):
        call(
            param_client.insert,
            "bare_datetime64_test",
            [[1, datetime(2023, 3, 25, 10, 5, 44, 772402)], [2, datetime.now()], [3, datetime(1965, 10, 15, 12, 0, 0)]],
        )
        result = call(param_client.query, "SELECT * FROM bare_datetime64_test ORDER BY key").result_rows
        assert result[0][0] == 1
        assert result[0][1] == datetime(2023, 3, 25, 10, 5, 44, 772000)
        assert result[2][1] == datetime(1965, 10, 15, 12, 0, 0)


def test_nulls(param_client: Client, call, table_context: Callable):
    with table_context("nullable_test", ["key UInt32", "null_str Nullable(String)", "null_int Nullable(Int64)"]):
        call(
            param_client.insert,
            "nullable_test",
            [[1, None, None], [2, "nonnull", -57382882345666], [3, None, 5882374747732834], [4, "nonnull2", None]],
        )
        result = call(param_client.query, "SELECT * FROM nullable_test ORDER BY key", use_none=False).result_rows
        assert result[2] == (3, "", 5882374747732834)
        assert result[3] == (4, "nonnull2", 0)
        result = call(param_client.query, "SELECT * FROM nullable_test ORDER BY key").result_rows
        assert result[1] == (2, "nonnull", -57382882345666)
        assert result[2] == (3, None, 5882374747732834)
        assert result[3] == (4, "nonnull2", None)


def test_read_formats(param_client: Client, call, test_table_engine: str):
    call(param_client.command, "DROP TABLE IF EXISTS read_format_test")
    call(
        param_client.command,
        "CREATE TABLE read_format_test (key Int32, uuid UUID, fs FixedString(10), ipv4 IPv4,"
        + "ip_array Array(IPv6), tup Tuple(u1 UInt64, ip2 IPv4))"
        + f"Engine {test_table_engine} ORDER BY key",
    )
    uuid1 = uuid.UUID("23E45688e89B-12D3-3273-426614174000")
    uuid2 = uuid.UUID("77AA3278-3728-12d3-5372-000377723832")
    row1 = (1, uuid1, "530055777k", "10.251.30.50", ["2600::", "2001:4860:4860::8844"], (7372, "10.20.30.203"))
    row2 = (2, uuid2, "short str", "10.44.75.20", ["74:382::3332", "8700:5200::5782:3992"], (7320, "252.18.4.50"))
    call(param_client.insert, "read_format_test", [row1, row2])

    result = call(param_client.query, "SELECT * FROM read_format_test;;;").result_set
    assert result[0][1] == uuid1
    assert result[1][3] == IPv4Address("10.44.75.20")
    assert result[0][2] == b"\x35\x33\x30\x30\x35\x35\x37\x37\x37\x6b"
    assert result[0][5]["u1"] == 7372
    assert result[0][5]["ip2"] == IPv4Address("10.20.30.203")

    set_default_formats("uuid", "string", "ip*", "string", "FixedString", "string")
    result = call(param_client.query, "SELECT * FROM read_format_test").result_set
    assert result[0][1] == "23e45688-e89b-12d3-3273-426614174000"
    assert result[1][3] == "10.44.75.20"
    assert result[0][2] == "530055777k"
    assert result[0][4][1] == "2001:4860:4860::8844"

    clear_default_format("ip*")
    result = call(param_client.query, "SELECT * FROM read_format_test").result_set
    assert result[0][1] == "23e45688-e89b-12d3-3273-426614174000"
    assert result[1][3] == IPv4Address("10.44.75.20")
    assert result[0][4][1] == IPv6Address("2001:4860:4860::8844")
    assert result[0][2] == "530055777k"

    # Test query formats
    result = call(param_client.query, "SELECT * FROM read_format_test", query_formats={"IP*": "string", "tup": "json"}).result_set
    assert result[1][3] == "10.44.75.20"
    assert result[0][5] == b'{"u1":7372,"ip2":"10.20.30.203"}'

    # Ensure that the query format clears
    result = call(param_client.query, "SELECT * FROM read_format_test").result_set
    assert result[1][3] == IPv4Address("10.44.75.20")
    assert result[0][5]["ip2"] == IPv4Address("10.20.30.203")

    # Test column formats
    result = call(param_client.query, "SELECT * FROM read_format_test", column_formats={"ipv4": "string", "tup": "tuple"}).result_set
    assert result[1][3] == "10.44.75.20"
    assert result[0][5][1] == IPv4Address("10.20.30.203")

    # Ensure that the column format clears
    result = call(param_client.query, "SELECT * FROM read_format_test").result_set
    assert result[1][3] == IPv4Address("10.44.75.20")
    assert result[0][5]["ip2"] == IPv4Address("10.20.30.203")

    # Test sub column formats
    set_read_format("tuple", "tuple")
    result = call(param_client.query, "SELECT * FROM read_format_test", column_formats={"tup": {"ip*": "string"}}).result_set
    assert result[0][5][1] == "10.20.30.203"

    set_read_format("tuple", "native")
    result = call(param_client.query, "SELECT * FROM read_format_test", column_formats={"tup": {"ip*": "string"}}).result_set
    assert result[0][5]["ip2"] == "10.20.30.203"


def test_tuple_inserts(param_client: Client, call, table_context: Callable):
    with table_context(
        "insert_tuple_test",
        ["key Int32", "named Tuple(fl Float64, `ns space` Nullable(String))", "unnamed Tuple(Float64, Nullable(String))"],
    ):
        data = [[1, (3.55, "str1"), (555, None)], [2, (-43.2, None), (0, "str2")]]
        call(param_client.insert, "insert_tuple_test", data, settings={"insert_deduplication_token": 5772})

        data = [[1, {"fl": 3.55, "ns space": "str1"}, (555, None)], [2, {"fl": -43.2}, (0, "str2")]]
        call(param_client.insert, "insert_tuple_test", data, settings={"insert_deduplication_token": 5773})
        query_result = call(param_client.query, "SELECT * FROM insert_tuple_test ORDER BY key").result_rows
        assert len(query_result) == 4
        assert query_result[0] == query_result[1]
        assert query_result[2] == query_result[3]


def test_agg_function(param_client: Client, call, table_context: Callable):
    with table_context(
        "agg_func_test",
        ["key Int32", "str SimpleAggregateFunction(any, String)", "lc_str SimpleAggregateFunction(any, LowCardinality(String))"],
        engine="AggregatingMergeTree",
    ):
        call(param_client.insert, "agg_func_test", [(1, "str", "lc_str")])
        row = call(param_client.query, "SELECT str, lc_str FROM agg_func_test").first_row
        assert row[0] == "str"
        assert row[1] == "lc_str"


def test_decimal_rounding(param_client: Client, call, table_context: Callable):
    test_vals = [732.4, 75.57, 75.49, 40.16]
    with table_context("test_decimal", ["key Int32, value Decimal(10, 2)"]):
        call(param_client.insert, "test_decimal", [[ix, x] for ix, x in enumerate(test_vals)])
        values = call(param_client.query, "SELECT value FROM test_decimal").result_columns[0]
    with decimal.localcontext() as dec_ctx:
        dec_ctx.prec = 10
        assert [decimal.Decimal(str(x)) for x in test_vals] == values


def test_empty_maps(param_client: Client, call):
    result = call(param_client.query, "select Cast(([],[]), 'Map(String, Map(String, String))')")
    assert result.first_row[0] == {}


def test_fixed_str_padding(param_client: Client, call, table_context: Callable):
    table = "test_fixed_str_padding"
    with table_context(table, "key Int32, value FixedString(3)"):
        call(param_client.insert, table, [[1, "abc"]])
        call(param_client.insert, table, [[2, "a"]])
        call(param_client.insert, table, [[3, ""]])
        result = call(param_client.query, f"select * from {table} ORDER BY key")
        assert result.result_columns[1] == [b"abc", b"a\x00\x00", b"\x00\x00\x00"]


def test_nonstandard_column_names(param_client: Client, call, table_context: Callable):
    table = "пример_кириллица"
    with table_context(table, "колонка String") as t:
        call(param_client.insert, t.table, (("привет",),))
        result = call(param_client.query, f"SELECT * FROM {t.table}").result_set
        assert result[0][0] == "привет"