def test_compressed_external_table(): columns = [ clickhouse_grpc_pb2.NameAndType(name="UserID", type="UInt64"), clickhouse_grpc_pb2.NameAndType(name="UserName", type="String"), ] d1 = lz4.frame.compress(b"1\tAlex\n2\tBen\n3\tCarl\n") d2 = gzip.compress(b"4,Daniel\n5,Ethan\n") ext1 = clickhouse_grpc_pb2.ExternalTable( name="ext1", columns=columns, data=d1, format="TabSeparated", compression_type="lz4", ) ext2 = clickhouse_grpc_pb2.ExternalTable(name="ext2", columns=columns, data=d2, format="CSV", compression_type="gzip") stub = clickhouse_grpc_pb2_grpc.ClickHouseStub(main_channel) query_info = clickhouse_grpc_pb2.QueryInfo( query= "SELECT * FROM (SELECT * FROM ext1 UNION ALL SELECT * FROM ext2) ORDER BY UserID", external_tables=[ext1, ext2], ) result = stub.ExecuteQuery(query_info) assert (result.output == b"1\tAlex\n" b"2\tBen\n" b"3\tCarl\n" b"4\tDaniel\n" b"5\tEthan\n")
def query(query_text, channel): query_info = clickhouse_grpc_pb2.QueryInfo(query=query_text) stub = clickhouse_grpc_pb2_grpc.ClickHouseStub(channel) result = stub.ExecuteQuery(query_info) if result and result.HasField("exception"): raise Exception(result.exception.display_text) return result.output.decode(DEFAULT_ENCODING)
def query_common(query_text, settings={}, input_data=[], input_data_delimiter='', output_format='TabSeparated', external_tables=[], user_name='', password='', query_id='123', session_id='', stream_output=False, channel=None): if type(input_data) == str: input_data = [input_data] if not channel: channel = main_channel stub = clickhouse_grpc_pb2_grpc.ClickHouseStub(channel) def query_info(): input_data_part = input_data.pop(0) if input_data else '' return clickhouse_grpc_pb2.QueryInfo(query=query_text, settings=settings, input_data=input_data_part, input_data_delimiter=input_data_delimiter, output_format=output_format, external_tables=external_tables, user_name=user_name, password=password, query_id=query_id, session_id=session_id, next_query_info=bool(input_data)) def send_query_info(): yield query_info() while input_data: input_data_part = input_data.pop(0) yield clickhouse_grpc_pb2.QueryInfo(input_data=input_data_part, next_query_info=bool(input_data)) stream_input = len(input_data) > 1 if stream_input and stream_output: return list(stub.ExecuteQueryWithStreamIO(send_query_info())) elif stream_input: return [stub.ExecuteQueryWithStreamInput(send_query_info())] elif stream_output: return list(stub.ExecuteQueryWithStreamOutput(query_info())) else: return [stub.ExecuteQuery(query_info())]
def test_compressed_external_table(): columns = [ clickhouse_grpc_pb2.NameAndType(name='UserID', type='UInt64'), clickhouse_grpc_pb2.NameAndType(name='UserName', type='String') ] d1 = lz4.frame.compress(b'1\tAlex\n2\tBen\n3\tCarl\n') d2 = gzip.compress(b'4,Daniel\n5,Ethan\n') ext1 = clickhouse_grpc_pb2.ExternalTable(name='ext1', columns=columns, data=d1, format='TabSeparated', compression_type="lz4") ext2 = clickhouse_grpc_pb2.ExternalTable(name='ext2', columns=columns, data=d2, format='CSV', compression_type="gzip") stub = clickhouse_grpc_pb2_grpc.ClickHouseStub(main_channel) query_info = clickhouse_grpc_pb2.QueryInfo( query= "SELECT * FROM (SELECT * FROM ext1 UNION ALL SELECT * FROM ext2) ORDER BY UserID", external_tables=[ext1, ext2]) result = stub.ExecuteQuery(query_info) assert result.output == b"1\tAlex\n"\ b"2\tBen\n"\ b"3\tCarl\n"\ b"4\tDaniel\n"\ b"5\tEthan\n"
def test_result_compression(): query_info = clickhouse_grpc_pb2.QueryInfo(query="SELECT 0 FROM numbers(1000000)", result_compression=clickhouse_grpc_pb2.Compression(algorithm=clickhouse_grpc_pb2.CompressionAlgorithm.GZIP, level=clickhouse_grpc_pb2.CompressionLevel.COMPRESSION_HIGH)) stub = clickhouse_grpc_pb2_grpc.ClickHouseStub(main_channel) result = stub.ExecuteQuery(query_info) assert result.output == (b'0\n')*1000000
def test_external_table_streaming(): columns = [ clickhouse_grpc_pb2.NameAndType(name='UserID', type='UInt64'), clickhouse_grpc_pb2.NameAndType(name='UserName', type='String') ] def send_query_info(): yield clickhouse_grpc_pb2.QueryInfo( query="SELECT * FROM exts ORDER BY UserID", external_tables=[ clickhouse_grpc_pb2.ExternalTable( name='exts', columns=columns, data=b'1\tAlex\n2\tBen\n3\tCarl\n') ], next_query_info=True) yield clickhouse_grpc_pb2.QueryInfo(external_tables=[ clickhouse_grpc_pb2.ExternalTable(name='exts', data=b'4\tDaniel\n5\tEthan\n') ]) stub = clickhouse_grpc_pb2_grpc.ClickHouseStub(main_channel) result = stub.ExecuteQueryWithStreamInput(send_query_info()) assert result.output == b'1\tAlex\n'\ b'2\tBen\n'\ b'3\tCarl\n'\ b'4\tDaniel\n'\ b'5\tEthan\n'
def query_common( query_text, settings={}, input_data=[], input_data_delimiter="", output_format="TabSeparated", send_output_columns=False, external_tables=[], user_name="", password="", query_id="123", session_id="", stream_output=False, channel=None, ): if type(input_data) is not list: input_data = [input_data] if type(input_data_delimiter) is str: input_data_delimiter = input_data_delimiter.encode(DEFAULT_ENCODING) if not channel: channel = main_channel stub = clickhouse_grpc_pb2_grpc.ClickHouseStub(channel) def query_info(): input_data_part = input_data.pop(0) if input_data else b"" if type(input_data_part) is str: input_data_part = input_data_part.encode(DEFAULT_ENCODING) return clickhouse_grpc_pb2.QueryInfo( query=query_text, settings=settings, input_data=input_data_part, input_data_delimiter=input_data_delimiter, output_format=output_format, send_output_columns=send_output_columns, external_tables=external_tables, user_name=user_name, password=password, query_id=query_id, session_id=session_id, next_query_info=bool(input_data), ) def send_query_info(): yield query_info() while input_data: input_data_part = input_data.pop(0) if type(input_data_part) is str: input_data_part = input_data_part.encode(DEFAULT_ENCODING) yield clickhouse_grpc_pb2.QueryInfo( input_data=input_data_part, next_query_info=bool(input_data)) stream_input = len(input_data) > 1 if stream_input and stream_output: return list(stub.ExecuteQueryWithStreamIO(send_query_info())) elif stream_input: return [stub.ExecuteQueryWithStreamInput(send_query_info())] elif stream_output: return list(stub.ExecuteQueryWithStreamOutput(query_info())) else: return [stub.ExecuteQuery(query_info())]
def test_external_table_streaming(): columns = [ clickhouse_grpc_pb2.NameAndType(name="UserID", type="UInt64"), clickhouse_grpc_pb2.NameAndType(name="UserName", type="String"), ] def send_query_info(): yield clickhouse_grpc_pb2.QueryInfo( query="SELECT * FROM exts ORDER BY UserID", external_tables=[ clickhouse_grpc_pb2.ExternalTable( name="exts", columns=columns, data=b"1\tAlex\n2\tBen\n3\tCarl\n") ], next_query_info=True, ) yield clickhouse_grpc_pb2.QueryInfo(external_tables=[ clickhouse_grpc_pb2.ExternalTable(name="exts", data=b"4\tDaniel\n5\tEthan\n") ]) stub = clickhouse_grpc_pb2_grpc.ClickHouseStub(main_channel) result = stub.ExecuteQueryWithStreamInput(send_query_info()) assert (result.output == b"1\tAlex\n" b"2\tBen\n" b"3\tCarl\n" b"4\tDaniel\n" b"5\tEthan\n")
def test_compressed_output_gzip(): query_info = clickhouse_grpc_pb2.QueryInfo( query="SELECT 0 FROM numbers(1000)", compression_type="gzip", compression_level=6) stub = clickhouse_grpc_pb2_grpc.ClickHouseStub(main_channel) result = stub.ExecuteQuery(query_info) assert gzip.decompress(result.output) == (b'0\n') * 1000
def test_transport_compression(): query_info = clickhouse_grpc_pb2.QueryInfo( query="SELECT 0 FROM numbers(1000000)", transport_compression_type='gzip', transport_compression_level=3) stub = clickhouse_grpc_pb2_grpc.ClickHouseStub(main_channel) result = stub.ExecuteQuery(query_info) assert result.output == (b'0\n') * 1000000
def test_compressed_output_streaming(): query_info = clickhouse_grpc_pb2.QueryInfo(query="SELECT 0 FROM numbers(100000)", output_compression_type="lz4") stub = clickhouse_grpc_pb2_grpc.ClickHouseStub(main_channel) d_context = lz4.frame.create_decompression_context() data = b'' for result in stub.ExecuteQueryWithStreamOutput(query_info): d1, _, _ = lz4.frame.decompress_chunk(d_context, result.output) data += d1 assert data == (b'0\n')*100000
def test_compressed_totals_and_extremes(): query("CREATE TABLE t (x UInt8, y UInt8) ENGINE = Memory") query("INSERT INTO t VALUES (1, 2), (2, 4), (3, 2), (3, 3), (3, 4)") stub = clickhouse_grpc_pb2_grpc.ClickHouseStub(main_channel) query_info = clickhouse_grpc_pb2.QueryInfo(query="SELECT sum(x), y FROM t GROUP BY y WITH TOTALS", output_compression_type="lz4") result = stub.ExecuteQuery(query_info) assert lz4.frame.decompress(result.totals) == b'12\t0\n' query_info = clickhouse_grpc_pb2.QueryInfo(query="SELECT x, y FROM t", settings={"extremes": "1"}, output_compression_type="lz4") result = stub.ExecuteQuery(query_info) assert lz4.frame.decompress(result.extremes) == b'1\t2\n3\t4\n'
def test_cancel_while_processing_input(): query("CREATE TABLE t (a UInt8) ENGINE = Memory") def send_query_info(): yield clickhouse_grpc_pb2.QueryInfo(query="INSERT INTO t FORMAT TabSeparated", input_data="1\n2\n3\n", next_query_info=True) yield clickhouse_grpc_pb2.QueryInfo(input_data="4\n5\n6\n", next_query_info=True) yield clickhouse_grpc_pb2.QueryInfo(cancel=True) stub = clickhouse_grpc_pb2_grpc.ClickHouseStub(main_channel) result = stub.ExecuteQueryWithStreamInput(send_query_info()) assert result.cancelled == True assert result.progress.written_rows == 6 assert query("SELECT a FROM t ORDER BY a") == "1\n2\n3\n4\n5\n6\n"
def test_cancel_while_generating_output(): def send_query_info(): yield clickhouse_grpc_pb2.QueryInfo(query="SELECT number, sleep(0.2) FROM numbers(10) SETTINGS max_block_size=2") time.sleep(0.5) yield clickhouse_grpc_pb2.QueryInfo(cancel=True) stub = clickhouse_grpc_pb2_grpc.ClickHouseStub(main_channel) results = list(stub.ExecuteQueryWithStreamIO(send_query_info())) assert len(results) >= 1 assert results[-1].cancelled == True output = '' for result in results: output += result.output assert output == '0\t0\n1\t0\n2\t0\n3\t0\n'
def test_opentelemetry_context_propagation(): trace_id = "80c190b5-9dc1-4eae-82b9-6c261438c817" parent_span_id = 123 trace_state = "some custom state" trace_id_hex = trace_id.replace("-", "") parent_span_id_hex = f'{parent_span_id:0>16X}' metadata = [("traceparent", f"00-{trace_id_hex}-{parent_span_id_hex}-01"), ("tracestate", trace_state)] stub = clickhouse_grpc_pb2_grpc.ClickHouseStub(main_channel) query_info = clickhouse_grpc_pb2.QueryInfo(query="SELECT 1") result = stub.ExecuteQuery(query_info, metadata=metadata) assert result.output == b"1\n" node.query("SYSTEM FLUSH LOGS") assert node.query(f"SELECT attribute['db.statement'], attribute['clickhouse.tracestate'] FROM system.opentelemetry_span_log " f"WHERE trace_id='{trace_id}' AND parent_span_id={parent_span_id}") == "SELECT 1\tsome custom state\n"
def test_compressed_insert_query_streaming(): query("CREATE TABLE t (a UInt8) ENGINE = Memory") data = lz4.frame.compress(b'(1),(2),(3),(5),(4),(6),(7),(8),(9)') sz1 = len(data) // 3 sz2 = len(data) // 3 d1 = data[:sz1] d2 = data[sz1:sz1+sz2] d3 = data[sz1+sz2:] def send_query_info(): yield clickhouse_grpc_pb2.QueryInfo(query="INSERT INTO t VALUES", input_data=d1, input_compression_type="lz4", next_query_info=True) yield clickhouse_grpc_pb2.QueryInfo(input_data=d2, next_query_info=True) yield clickhouse_grpc_pb2.QueryInfo(input_data=d3) stub = clickhouse_grpc_pb2_grpc.ClickHouseStub(main_channel) stub.ExecuteQueryWithStreamInput(send_query_info()) assert query("SELECT a FROM t ORDER BY a") == "1\n2\n3\n4\n5\n6\n7\n8\n9\n"
def test_compressed_output(): query_info = clickhouse_grpc_pb2.QueryInfo( query="SELECT 0 FROM numbers(1000)", output_compression_type="lz4") stub = clickhouse_grpc_pb2_grpc.ClickHouseStub(main_channel) result = stub.ExecuteQuery(query_info) assert lz4.frame.decompress(result.output) == (b"0\n") * 1000