ClickHouse/tests/integration/test_table_function_redis/test.py

import datetime

import redis
import pytest
import sys
import struct

from helpers.client import QueryRuntimeException
from helpers.cluster import ClickHouseCluster
from helpers.test_tools import TSV

cluster = ClickHouseCluster(__file__)

node = cluster.add_instance("node", with_redis=True)


@pytest.fixture(scope="module")
def started_cluster():
    try:
        cluster.start()
        yield cluster
    finally:
        cluster.shutdown()


def get_redis_connection(db_id=0):
    client = redis.Redis(
        host='localhost', port=cluster.redis_port, password="clickhouse", db=db_id
    )
    return client


def get_address_for_ch():
    return cluster.redis_host + ':6379'


# see SerializationString.serializeBinary
def serialize_binary_for_string(x):
    var_uint_max = (1 << 63) - 1
    buf = bytearray()
    # write length
    length = len(x)
    # length = (length << 1) ^ (length >> 63)
    if length > var_uint_max:
        raise ValueError("Value too large for varint encoding")
    for i in range(9):
        byte = length & 0x7F
        if length > 0x7F:
            byte |= 0x80
        buf += (bytes([byte]))
        length >>= 7
        if not length:
            break
    # write data
    buf += x.encode('utf-8')
    return bytes(buf)


# see SerializationNumber.serializeBinary
def serialize_binary_for_uint32(x):
    buf = bytearray()
    packed_num = struct.pack('I', x)
    buf += packed_num
    if sys.byteorder != 'little':
        buf.reverse()
    return bytes(buf)


def test_simple_select(started_cluster):
    client = get_redis_connection()
    address = get_address_for_ch()

    # clean all
    client.flushall()

    data = {}
    for i in range(100):
        packed = serialize_binary_for_string(str(i))
        data[packed] = packed

    client.mset(data)
    client.close()

    response = TSV.toMat(node.query(
        f"""
        SELECT 
            key, value 
        FROM 
            redis('{address}', 'key', 'key String, value String', 0, 'clickhouse', 10) 
        WHERE 
            key='0' 
        FORMAT TSV
        """))

    assert (len(response) == 1)
    assert (response[0] == ['0', '0'])

    response = TSV.toMat(node.query(
        f"""
        SELECT 
            * 
        FROM 
            redis('{address}', 'key', 'key String, value String', 0, 'clickhouse', 10) 
        ORDER BY 
            key 
        FORMAT TSV
        """))

    assert (len(response) == 100)
    assert (response[0] == ['0', '0'])


def test_create_table(started_cluster):
    client = get_redis_connection()
    address = get_address_for_ch()

    # clean all
    client.flushall()
    client.close()

    node.query(
        f"""
        SELECT 
            *
        FROM 
            redis('{address}', 'k', 'k String, v UInt32', 0, 'clickhouse', 10) 
        """)

    # illegal data type
    with pytest.raises(QueryRuntimeException):
        node.query(
            f"""
            SELECT 
                *
            FROM 
                redis('{address}', 'k', 'k not_exist_type, v String', 0, 'clickhouse', 10) 
            """)

    # illegal key
    with pytest.raises(QueryRuntimeException):
        node.query(
            f"""
            SELECT 
                *
            FROM 
                redis('{address}', 'not_exist_key', 'k not_exist_type, v String', 0, 'clickhouse', 10) 
            """)


def test_data_type(started_cluster):
    client = get_redis_connection()
    address = get_address_for_ch()

    # string
    client.flushall()
    value = serialize_binary_for_string('0')
    client.set(value, value)

    response = TSV.toMat(node.query(
        f"""
        SELECT
            *
        FROM
            redis('{address}', 'k', 'k String, v String', 0, 'clickhouse', 10)
        WHERE
            k='0'
        FORMAT TSV
        """))

    assert (len(response) == 1)
    assert (response[0] == ['0', '0'])

    # number
    client.flushall()
    value = serialize_binary_for_uint32(0)
    client.set(value, value)

    response = TSV.toMat(node.query(
        f"""
        SELECT
            *
        FROM
            redis('{address}', 'k', 'k UInt32, v UInt32', 0, 'clickhouse', 10)
        WHERE
            k=0
        FORMAT TSV
        """))

    assert (len(response) == 1)
    assert (response[0] == ['0', '0'])

    # datetime
    client.flushall()
    # clickhouse store datatime as uint32 in internal
    dt = datetime.datetime(2023, 6, 1, 0, 0, 0)
    seconds_since_epoch = dt.timestamp()
    value = serialize_binary_for_uint32(int(seconds_since_epoch))
    client.set(value, value)

    response = TSV.toMat(node.query(
        f"""
        SELECT
            *
        FROM
            redis('{address}', 'k', 'k DateTime, v DateTime', 0, 'clickhouse', 10)
        WHERE
            k='2023-06-01 00:00:00'
        FORMAT TSV
        """))

    assert (len(response) == 1)
    assert (response[0] == ['2023-06-01 00:00:00', '2023-06-01 00:00:00'])