From 8021dc31506ddb3c2efd8dc104ae5e5aec981464 Mon Sep 17 00:00:00 2001 From: Dom Del Nano Date: Wed, 23 Sep 2026 22:17:40 -0700 Subject: [PATCH] [carnot] Add ClickHouse source and export-sink execution nodes Signed-off-by: Dom Del Nano --- src/carnot/carnot.cc | 2 + src/carnot/exec/BUILD.bazel | 50 ++ .../exec/clickhouse_export_sink_node.cc | 181 +++++ src/carnot/exec/clickhouse_export_sink_node.h | 55 ++ .../exec/clickhouse_export_sink_node_test.cc | 432 +++++++++++ src/carnot/exec/clickhouse_source_node.cc | 715 ++++++++++++++++++ src/carnot/exec/clickhouse_source_node.h | 108 +++ .../exec/clickhouse_source_node_test.cc | 415 ++++++++++ src/carnot/exec/exec_graph.cc | 10 + src/carnot/exec/exec_state.h | 2 + src/common/ch/BUILD.bazel | 28 + src/common/ch/utils.cc | 68 ++ src/common/ch/utils.h | 50 ++ 13 files changed, 2116 insertions(+) create mode 100644 src/carnot/exec/clickhouse_export_sink_node.cc create mode 100644 src/carnot/exec/clickhouse_export_sink_node.h create mode 100644 src/carnot/exec/clickhouse_export_sink_node_test.cc create mode 100644 src/carnot/exec/clickhouse_source_node.cc create mode 100644 src/carnot/exec/clickhouse_source_node.h create mode 100644 src/carnot/exec/clickhouse_source_node_test.cc create mode 100644 src/common/ch/BUILD.bazel create mode 100644 src/common/ch/utils.cc create mode 100644 src/common/ch/utils.h diff --git a/src/carnot/carnot.cc b/src/carnot/carnot.cc index a466bb5194d..1a1f72ff74a 100644 --- a/src/carnot/carnot.cc +++ b/src/carnot/carnot.cc @@ -181,6 +181,8 @@ Status CarnotImpl::RegisterUDFsInPlanFragment(exec::ExecState* exec_state, plan: .OnUDTFSource(no_op) .OnEmptySource(no_op) .OnOTelSink(no_op) + .OnClickHouseSource(no_op) + .OnClickHouseExportSink(no_op) .Walk(pf); } diff --git a/src/carnot/exec/BUILD.bazel b/src/carnot/exec/BUILD.bazel index 3f8732f705a..6e073426956 100644 --- a/src/carnot/exec/BUILD.bazel +++ b/src/carnot/exec/BUILD.bazel @@ -33,6 +33,7 @@ pl_cc_library( ], ), hdrs = [ + "clickhouse_source_node.h", "exec_node.h", "exec_state.h", ], @@ -42,10 +43,12 @@ pl_cc_library( "//src/carnot/plan:cc_library", "//src/carnot/planpb:plan_pl_cc_proto", "//src/carnot/udf:cc_library", + "//src/common/ch:cc_library", "//src/common/uuid:cc_library", "//src/shared/types:cc_library", "//src/table_store/table:cc_library", "@com_github_apache_arrow//:arrow", + "@com_github_clickhouse_clickhouse_cpp//:clickhouse_cpp", "@com_github_grpc_grpc//:grpc++", "@com_github_opentelemetry_proto//:logs_service_grpc_cc", "@com_github_opentelemetry_proto//:metrics_service_grpc_cc", @@ -300,3 +303,50 @@ pl_cc_test( "@com_github_grpc_grpc//:grpc++_test", ], ) + +pl_cc_test( + name = "clickhouse_source_node_test", + timeout = "long", + srcs = ["clickhouse_source_node_test.cc"], + data = [ + "//src/stirling/source_connectors/socket_tracer/testing/container_images:clickhouse.tar", + ], + flaky = True, + tags = [ + "exclusive", + "requires_root", + ], + deps = [ + ":cc_library", + ":exec_node_test_helpers", + ":test_utils", + "//src/carnot/planpb:plan_testutils", + "//src/common/ch:cc_library", + "//src/common/testing/test_utils:cc_library", + "@com_github_clickhouse_clickhouse_cpp//:clickhouse_cpp", + ], +) + +pl_cc_test( + name = "clickhouse_export_sink_node_test", + timeout = "long", + srcs = ["clickhouse_export_sink_node_test.cc"], + data = [ + "//src/stirling/source_connectors/socket_tracer/testing/container_images:clickhouse.tar", + ], + flaky = True, + tags = [ + "exclusive", + "requires_root", + ], + deps = [ + ":cc_library", + ":exec_node_test_helpers", + ":test_utils", + "//src/carnot/plan:cc_library", + "//src/carnot/planpb:plan_pl_cc_proto", + "//src/common/ch:cc_library", + "//src/common/testing/test_utils:cc_library", + "@com_github_clickhouse_clickhouse_cpp//:clickhouse_cpp", + ], +) diff --git a/src/carnot/exec/clickhouse_export_sink_node.cc b/src/carnot/exec/clickhouse_export_sink_node.cc new file mode 100644 index 00000000000..15d928b99d7 --- /dev/null +++ b/src/carnot/exec/clickhouse_export_sink_node.cc @@ -0,0 +1,181 @@ +/* + * Copyright 2018- The Pixie Authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "src/carnot/exec/clickhouse_export_sink_node.h" + +#include +#include +#include + +#include +#include +#include +#include "glog/logging.h" +#include "src/carnot/planpb/plan.pb.h" +#include "src/common/base/macros.h" +#include "src/common/ch/utils.h" +#include "src/shared/types/typespb/types.pb.h" +#include "src/table_store/table_store.h" + +namespace px { +namespace carnot { +namespace exec { + +using table_store::schema::RowBatch; +using table_store::schema::RowDescriptor; + +std::string ClickHouseExportSinkNode::DebugStringImpl() { + return absl::Substitute("Exec::ClickHouseExportSinkNode: $0", plan_node_->DebugString()); +} + +Status ClickHouseExportSinkNode::InitImpl(const plan::Operator& plan_node) { + CHECK(plan_node.op_type() == planpb::OperatorType::CLICKHOUSE_EXPORT_SINK_OPERATOR); + if (input_descriptors_.size() != 1) { + return error::InvalidArgument( + "ClickHouse Export operator expects a single input relation, got $0", + input_descriptors_.size()); + } + + input_descriptor_ = std::make_unique(input_descriptors_[0]); + const auto* sink_plan_node = static_cast(&plan_node); + plan_node_ = std::make_unique(*sink_plan_node); + return Status::OK(); +} + +Status ClickHouseExportSinkNode::PrepareImpl(ExecState*) { return Status::OK(); } + +Status ClickHouseExportSinkNode::OpenImpl(ExecState* /*exec_state*/) { + // Connect to ClickHouse using config from plan node + const auto& config = plan_node_->clickhouse_config(); + + clickhouse::ClientOptions options; + options.SetHost(config.host()); + options.SetPort(config.port()); + options.SetUser(config.username()); + options.SetPassword(config.password()); + options.SetDefaultDatabase(config.database()); + + PX_ASSIGN_OR_RETURN(clickhouse_client_, ch::CreateClient(options)); + + return Status::OK(); +} + +Status ClickHouseExportSinkNode::CloseImpl(ExecState* exec_state) { + if (sent_eos_) { + return Status::OK(); + } + + LOG(INFO) << absl::Substitute( + "Closing ClickHouseExportSinkNode $0 in query $1 before receiving EOS", plan_node_->id(), + exec_state->query_id().str()); + + return Status::OK(); +} + +Status ClickHouseExportSinkNode::ConsumeNextImpl(ExecState* /*exec_state*/, const RowBatch& rb, + size_t /*parent_index*/) { + // Skip insertion if the batch is empty + if (rb.num_rows() == 0) { + if (rb.eos()) { + sent_eos_ = true; + } + return Status::OK(); + } + + // Build an INSERT query with the data from the row batch + clickhouse::Block block; + + // Create columns based on the column mappings + for (const auto& mapping : plan_node_->column_mappings()) { + auto arrow_col = rb.ColumnAt(mapping.input_column_index()); + int64_t num_rows = arrow_col->length(); + + // Create ClickHouse column based on data type + switch (mapping.column_type()) { + case types::INT64: { + auto col = std::make_shared(); + for (int64_t i = 0; i < num_rows; ++i) { + col->Append(types::GetValueFromArrowArray(arrow_col.get(), i)); + } + block.AppendColumn(mapping.clickhouse_column_name(), col); + break; + } + case types::FLOAT64: { + auto col = std::make_shared(); + for (int64_t i = 0; i < num_rows; ++i) { + col->Append(types::GetValueFromArrowArray(arrow_col.get(), i)); + } + block.AppendColumn(mapping.clickhouse_column_name(), col); + break; + } + case types::STRING: { + auto col = std::make_shared(); + for (int64_t i = 0; i < num_rows; ++i) { + col->Append(types::GetValueFromArrowArray(arrow_col.get(), i)); + } + block.AppendColumn(mapping.clickhouse_column_name(), col); + break; + } + case types::TIME64NS: { + auto col = std::make_shared(9); + for (int64_t i = 0; i < num_rows; ++i) { + int64_t ns_val = types::GetValueFromArrowArray(arrow_col.get(), i); + col->Append(ns_val); + } + block.AppendColumn(mapping.clickhouse_column_name(), col); + break; + } + case types::BOOLEAN: { + auto col = std::make_shared(); + for (int64_t i = 0; i < num_rows; ++i) { + col->Append(types::GetValueFromArrowArray(arrow_col.get(), i) ? 1 : 0); + } + block.AppendColumn(mapping.clickhouse_column_name(), col); + break; + } + case types::UINT128: { + // UINT128 is exported as STRING (UUID format) + auto col = std::make_shared(); + for (int64_t i = 0; i < num_rows; ++i) { + auto val = types::GetValueFromArrowArray(arrow_col.get(), i); + std::string uuid_str = + sole::rebuild(absl::Uint128High64(val), absl::Uint128Low64(val)).str(); + col->Append(uuid_str); + } + block.AppendColumn(mapping.clickhouse_column_name(), col); + break; + } + default: + return error::InvalidArgument("Unsupported data type for ClickHouse export: $0", + types::ToString(mapping.column_type())); + } + } + + // Insert the block into ClickHouse + PX_RETURN_IF_ERROR(ch::Insert(clickhouse_client_.get(), plan_node_->table_name(), block)); + + if (rb.eos()) { + sent_eos_ = true; + } + + return Status::OK(); +} + +} // namespace exec +} // namespace carnot +} // namespace px diff --git a/src/carnot/exec/clickhouse_export_sink_node.h b/src/carnot/exec/clickhouse_export_sink_node.h new file mode 100644 index 00000000000..26478afe037 --- /dev/null +++ b/src/carnot/exec/clickhouse_export_sink_node.h @@ -0,0 +1,55 @@ +/* + * Copyright 2018- The Pixie Authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include +#include +#include +#include + +#include "src/carnot/exec/exec_node.h" +#include "src/carnot/planpb/plan.pb.h" +#include "src/common/base/base.h" +#include "src/shared/types/types.h" + +namespace px { +namespace carnot { +namespace exec { + +class ClickHouseExportSinkNode : public SinkNode { + public: + virtual ~ClickHouseExportSinkNode() = default; + + protected: + std::string DebugStringImpl() override; + Status InitImpl(const plan::Operator& plan_node) override; + Status PrepareImpl(ExecState* exec_state) override; + Status OpenImpl(ExecState* exec_state) override; + Status CloseImpl(ExecState* exec_state) override; + Status ConsumeNextImpl(ExecState* exec_state, const table_store::schema::RowBatch& rb, + size_t parent_index) override; + + private: + std::unique_ptr input_descriptor_; + std::unique_ptr clickhouse_client_; + std::unique_ptr plan_node_; +}; + +} // namespace exec +} // namespace carnot +} // namespace px diff --git a/src/carnot/exec/clickhouse_export_sink_node_test.cc b/src/carnot/exec/clickhouse_export_sink_node_test.cc new file mode 100644 index 00000000000..98faead0cbd --- /dev/null +++ b/src/carnot/exec/clickhouse_export_sink_node_test.cc @@ -0,0 +1,432 @@ +/* + * Copyright 2018- The Pixie Authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "src/carnot/exec/clickhouse_export_sink_node.h" + +#include + +#include + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#include "src/carnot/exec/test_utils.h" +#include "src/carnot/plan/operators.h" +#include "src/carnot/planpb/plan.pb.h" +#include "src/carnot/udf/registry.h" +#include "src/common/ch/utils.h" +#include "src/common/event/time_system.h" +#include "src/common/testing/test_utils/container_runner.h" +#include "src/common/testing/testing.h" +#include "src/shared/metadata/metadata_state.h" +#include "src/shared/types/types.h" + +namespace px { +namespace carnot { +namespace exec { + +using table_store::schema::RowDescriptor; + +class ClickHouseExportSinkNodeTest : public ::testing::Test { + protected: + static constexpr char kClickHouseImage[] = + "src/stirling/source_connectors/socket_tracer/testing/container_images/clickhouse.tar"; + static constexpr char kClickHouseReadyMessage[] = "Ready for connections"; + static constexpr int kClickHousePort = 9000; + + void SetUp() override { + // Set up function registry and exec state + func_registry_ = std::make_unique("test_registry"); + auto table_store = std::make_shared(); + exec_state_ = std::make_unique( + func_registry_.get(), table_store, MockResultSinkStubGenerator, MockMetricsStubGenerator, + MockTraceStubGenerator, MockLogStubGenerator, sole::uuid4(), nullptr); + + // Create a minimal agent metadata state for test execution + auto metadata_state = + std::make_shared("test_host", // hostname + 1, // asid + getpid(), // pid + sole::uuid4(), // agent_id + "", // pod_name + sole::uuid4(), // vizier_id + "test_vizier", // vizier_name + "", // vizier_namespace + time_system_.get()); // time_system + exec_state_->set_metadata_state(metadata_state); + + // Start ClickHouse container + clickhouse_server_ = + std::make_unique(px::testing::BazelRunfilePath(kClickHouseImage), + "clickhouse_export_test", kClickHouseReadyMessage); + + std::vector options = { + absl::Substitute("--publish=$0:$0", kClickHousePort), + "--env=CLICKHOUSE_PASSWORD=test_password", + "--network=host", + }; + + ASSERT_OK(clickhouse_server_->Run(std::chrono::seconds{60}, options, {}, true, + std::chrono::seconds{300})); + + // Give ClickHouse time to initialize + std::this_thread::sleep_for(std::chrono::seconds(5)); + + // Create ClickHouse client for verification + SetupClickHouseClient(); + } + + void TearDown() override { + if (client_) { + client_.reset(); + } + } + + void SetupClickHouseClient() { + clickhouse::ClientOptions client_options; + client_options.SetHost("localhost"); + client_options.SetPort(kClickHousePort); + client_options.SetUser("default"); + client_options.SetPassword("test_password"); + client_options.SetDefaultDatabase("default"); + + const int kMaxRetries = 5; + Status connect_status; + for (int i = 0; i < kMaxRetries; ++i) { + LOG(INFO) << "Attempting to connect to ClickHouse (attempt " << (i + 1) << "/" << kMaxRetries + << ")..."; + auto client_or = ch::CreateClient(client_options); + if (client_or.ok()) { + client_ = client_or.ConsumeValueOrDie(); + connect_status = ch::Execute(client_.get(), "SELECT 1"); + if (connect_status.ok()) { + return; + } + } else { + connect_status = client_or.status(); + } + LOG(WARNING) << "Failed to connect: " << connect_status.msg(); + client_.reset(); + if (i < kMaxRetries - 1) { + std::this_thread::sleep_for(std::chrono::seconds(2)); + } + } + FAIL() << "Failed to connect to ClickHouse after " << kMaxRetries + << " attempts: " << connect_status.msg(); + } + + void CreateExportTable(const std::string& table_name) { + ASSERT_OK(ch::Execute(client_.get(), absl::Substitute("DROP TABLE IF EXISTS $0", table_name))); + + ASSERT_OK(ch::Execute(client_.get(), absl::Substitute(R"( + CREATE TABLE $0 ( + time_ DateTime64(9), + hostname String, + count Int64, + latency Float64 + ) ENGINE = MergeTree() + ORDER BY time_ + )", + table_name))); + + LOG(INFO) << "Export table created successfully: " << table_name; + } + + std::vector> QueryTable(const std::string& query) { + std::vector> results; + + EXPECT_OK(ch::Select(client_.get(), query, [&](const clickhouse::Block& block) { + for (size_t row_idx = 0; row_idx < block.GetRowCount(); ++row_idx) { + std::vector row; + for (size_t col_idx = 0; col_idx < block.GetColumnCount(); ++col_idx) { + auto col = block[col_idx]; + std::string value; + + if (auto int_col = col->As()) { + value = std::to_string((*int_col)[row_idx]); + } else if (auto uint_col = col->As()) { + value = std::to_string((*uint_col)[row_idx]); + } else if (auto float_col = col->As()) { + value = std::to_string((*float_col)[row_idx]); + } else if (auto str_col = col->As()) { + value = (*str_col)[row_idx]; + } else if (auto dt_col = col->As()) { + value = std::to_string((*dt_col)[row_idx]); + } else { + value = ""; + } + + row.push_back(value); + } + results.push_back(row); + } + })); + + return results; + } + + std::unique_ptr CreatePlanNode( + const std::string& table_name) { + planpb::Operator op; + op.set_op_type(planpb::CLICKHOUSE_EXPORT_SINK_OPERATOR); + auto* ch_op = op.mutable_clickhouse_sink_op(); + + auto* config = ch_op->mutable_clickhouse_config(); + config->set_host("localhost"); + config->set_port(kClickHousePort); + config->set_username("default"); + config->set_password("test_password"); + config->set_database("default"); + + ch_op->set_table_name(table_name); + + // Add column mappings + auto* mapping0 = ch_op->add_column_mappings(); + mapping0->set_input_column_index(0); + mapping0->set_clickhouse_column_name("time_"); + mapping0->set_column_type(types::TIME64NS); + + auto* mapping1 = ch_op->add_column_mappings(); + mapping1->set_input_column_index(1); + mapping1->set_clickhouse_column_name("hostname"); + mapping1->set_column_type(types::STRING); + + auto* mapping2 = ch_op->add_column_mappings(); + mapping2->set_input_column_index(2); + mapping2->set_clickhouse_column_name("count"); + mapping2->set_column_type(types::INT64); + + auto* mapping3 = ch_op->add_column_mappings(); + mapping3->set_input_column_index(3); + mapping3->set_clickhouse_column_name("latency"); + mapping3->set_column_type(types::FLOAT64); + + auto plan_node = std::make_unique(1); + EXPECT_OK(plan_node->Init(op.clickhouse_sink_op())); + + return plan_node; + } + + std::unique_ptr clickhouse_server_; + std::unique_ptr client_; + std::unique_ptr exec_state_; + std::unique_ptr func_registry_; + std::unique_ptr time_system_ = std::make_unique(); +}; + +TEST_F(ClickHouseExportSinkNodeTest, BasicExport) { + const std::string table_name = "export_test_basic"; + CreateExportTable(table_name); + + auto plan_node = CreatePlanNode(table_name); + + RowDescriptor input_rd({types::TIME64NS, types::STRING, types::INT64, types::FLOAT64}); + + auto tester = exec::ExecNodeTester( + *plan_node, RowDescriptor({}), {input_rd}, exec_state_.get()); + + // Send a multi-row batch followed by several single-row batches. + auto multi_row_rb = + RowBatchBuilder(input_rd, 2, /*eow*/ false, /*eos*/ false) + .AddColumn({1000000000000000000LL, 2000000000000000000LL}) + .AddColumn({"host0", "host1"}) + .AddColumn({0, 100}) + .AddColumn({0.5, 1.5}) + .get(); + tester.ConsumeNext(multi_row_rb, 0, 0); + + for (int i = 2; i < 5; ++i) { + bool is_last = (i == 4); + auto rb = RowBatchBuilder(input_rd, 1, /*eow*/ is_last, /*eos*/ is_last) + .AddColumn({(i + 1) * 1000000000000000000LL}) + .AddColumn({absl::Substitute("host$0", i)}) + .AddColumn({i * 100}) + .AddColumn({i + 0.5}) + .get(); + tester.ConsumeNext(rb, 0, 0); + } + + tester.Close(); + + // Verify every row from every batch was inserted with the correct values. + auto results = QueryTable( + absl::Substitute("SELECT hostname, count, latency FROM $0 ORDER BY time_", table_name)); + + ASSERT_EQ(results.size(), 5); + for (int i = 0; i < 5; ++i) { + EXPECT_EQ(results[i][0], absl::Substitute("host$0", i)); + EXPECT_EQ(results[i][1], absl::StrCat(i * 100)); + EXPECT_THAT(results[i][2], ::testing::StartsWith(absl::StrCat(i, ".5"))); + } +} + +TEST_F(ClickHouseExportSinkNodeTest, EmptyBatch) { + const std::string table_name = "export_test_empty"; + CreateExportTable(table_name); + + auto plan_node = CreatePlanNode(table_name); + + RowDescriptor input_rd({types::TIME64NS, types::STRING, types::INT64, types::FLOAT64}); + + auto tester = exec::ExecNodeTester( + *plan_node, RowDescriptor({}), {input_rd}, exec_state_.get()); + + // Send only EOS batch + auto rb = RowBatchBuilder(input_rd, 0, /*eow*/ true, /*eos*/ true) + .AddColumn({}) + .AddColumn({}) + .AddColumn({}) + .AddColumn({}) + .get(); + + tester.ConsumeNext(rb, 0, 0); + tester.Close(); + + // Verify no data was inserted + auto results = QueryTable(absl::Substitute("SELECT COUNT(*) FROM $0", table_name)); + + ASSERT_EQ(results.size(), 1); + EXPECT_EQ(results[0][0], "0"); +} + +TEST_F(ClickHouseExportSinkNodeTest, UINT128Export) { + const std::string table_name = "export_test_uint128"; + + // Create table with String column for UUID + ASSERT_OK(ch::Execute(client_.get(), absl::Substitute("DROP TABLE IF EXISTS $0", table_name))); + + ASSERT_OK(ch::Execute(client_.get(), absl::Substitute(R"( + CREATE TABLE $0 ( + time_ DateTime64(9), + upid String, + hostname String, + value Int64 + ) ENGINE = MergeTree() + ORDER BY time_ + )", + table_name))); + + LOG(INFO) << "UINT128 export table created successfully: " << table_name; + + // Create plan node for UINT128 test + planpb::Operator op; + op.set_op_type(planpb::CLICKHOUSE_EXPORT_SINK_OPERATOR); + auto* ch_op = op.mutable_clickhouse_sink_op(); + + auto* config = ch_op->mutable_clickhouse_config(); + config->set_host("localhost"); + config->set_port(kClickHousePort); + config->set_username("default"); + config->set_password("test_password"); + config->set_database("default"); + + ch_op->set_table_name(table_name); + + // Add column mappings + auto* mapping0 = ch_op->add_column_mappings(); + mapping0->set_input_column_index(0); + mapping0->set_clickhouse_column_name("time_"); + mapping0->set_column_type(types::TIME64NS); + + auto* mapping1 = ch_op->add_column_mappings(); + mapping1->set_input_column_index(1); + mapping1->set_clickhouse_column_name("upid"); + mapping1->set_column_type(types::UINT128); + + auto* mapping2 = ch_op->add_column_mappings(); + mapping2->set_input_column_index(2); + mapping2->set_clickhouse_column_name("hostname"); + mapping2->set_column_type(types::STRING); + + auto* mapping3 = ch_op->add_column_mappings(); + mapping3->set_input_column_index(3); + mapping3->set_clickhouse_column_name("value"); + mapping3->set_column_type(types::INT64); + + auto plan_node = std::make_unique(1); + EXPECT_OK(plan_node->Init(op.clickhouse_sink_op())); + + // Define input schema + RowDescriptor input_rd({types::TIME64NS, types::UINT128, types::STRING, types::INT64}); + + // Create node tester + auto tester = exec::ExecNodeTester( + *plan_node, RowDescriptor({}), {input_rd}, exec_state_.get()); + + // Create test UUIDs + auto uuid1 = sole::uuid4(); + auto uuid2 = sole::uuid4(); + auto uuid3 = sole::uuid4(); + + absl::uint128 upid1 = absl::MakeUint128(uuid1.ab, uuid1.cd); + absl::uint128 upid2 = absl::MakeUint128(uuid2.ab, uuid2.cd); + absl::uint128 upid3 = absl::MakeUint128(uuid3.ab, uuid3.cd); + + // Create test data with UINT128 values + auto rb1 = RowBatchBuilder(input_rd, 2, /*eow*/ false, /*eos*/ false) + .AddColumn({1000000000000000000LL, 2000000000000000000LL}) + .AddColumn({upid1, upid2}) + .AddColumn({"host1", "host2"}) + .AddColumn({100, 200}) + .get(); + + auto rb2 = RowBatchBuilder(input_rd, 1, /*eow*/ true, /*eos*/ true) + .AddColumn({3000000000000000000LL}) + .AddColumn({upid3}) + .AddColumn({"host3"}) + .AddColumn({300}) + .get(); + + // Send data to sink + tester.ConsumeNext(rb1, 0, 0); + tester.ConsumeNext(rb2, 0, 0); + tester.Close(); + + // Verify data was inserted and UINT128 values were converted to UUID strings + auto results = QueryTable( + absl::Substitute("SELECT upid, hostname, value FROM $0 ORDER BY time_", table_name)); + + ASSERT_EQ(results.size(), 3); + + // Check that UINT128 values were converted to valid UUID strings + EXPECT_EQ(results[0][0], uuid1.str()); + EXPECT_EQ(results[0][1], "host1"); + EXPECT_EQ(results[0][2], "100"); + + EXPECT_EQ(results[1][0], uuid2.str()); + EXPECT_EQ(results[1][1], "host2"); + EXPECT_EQ(results[1][2], "200"); + + EXPECT_EQ(results[2][0], uuid3.str()); + EXPECT_EQ(results[2][1], "host3"); + EXPECT_EQ(results[2][2], "300"); +} + +} // namespace exec +} // namespace carnot +} // namespace px diff --git a/src/carnot/exec/clickhouse_source_node.cc b/src/carnot/exec/clickhouse_source_node.cc new file mode 100644 index 00000000000..e05f48125ee --- /dev/null +++ b/src/carnot/exec/clickhouse_source_node.cc @@ -0,0 +1,715 @@ +/* + * Copyright 2018- The Pixie Authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "src/carnot/exec/clickhouse_source_node.h" + +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +#include "src/carnot/planpb/plan.pb.h" +#include "src/common/base/base.h" +#include "src/common/ch/utils.h" +#include "src/shared/types/arrow_adapter.h" +#include "src/shared/types/types.h" + +namespace px { +namespace carnot { +namespace exec { + +namespace { + +// Returns the position of the first case-insensitive, whitespace-delimited occurrence of keyword +// in query (pointing at the preceding whitespace), or std::string::npos. This is a heuristic: it +// does not understand string literals or subqueries. +size_t FindKeyword(std::string_view query, std::string_view keyword) { + std::string normalized = absl::AsciiStrToLower(query); + std::replace_if(normalized.begin(), normalized.end(), absl::ascii_isspace, ' '); + return normalized.find(absl::StrCat(" ", keyword, " ")); +} + +// Returns the position of the earliest of keywords that occurs after `after`, or the query length +// if there are none. +size_t FindClauseEnd(std::string_view query, size_t after, + std::initializer_list keywords) { + size_t end = query.length(); + for (const auto& keyword : keywords) { + size_t pos = FindKeyword(query.substr(after), keyword); + if (pos != std::string::npos) { + end = std::min(end, after + pos); + } + } + return end; +} + +std::string EscapeStringLiteral(std::string_view value) { + return absl::StrReplaceAll(value, {{"\\", "\\\\"}, {"'", "\\'"}}); +} + +} // namespace + +std::string ClickHouseSourceNode::DebugStringImpl() { + return absl::Substitute("Exec::ClickHouseSourceNode: ", base_query_, + output_descriptor_->DebugString()); +} + +Status ClickHouseSourceNode::InitImpl(const plan::Operator& plan_node) { + CHECK(plan_node.op_type() == planpb::OperatorType::CLICKHOUSE_SOURCE_OPERATOR); + const auto* source_plan_node = static_cast(&plan_node); + + // Copy the plan node to local object + plan_node_ = std::make_unique(*source_plan_node); + + // Extract connection parameters from plan node + host_ = plan_node_->host(); + port_ = plan_node_->port(); + username_ = plan_node_->username(); + password_ = plan_node_->password(); + database_ = plan_node_->database(); + base_query_ = plan_node_->query(); + batch_size_ = plan_node_->batch_size(); + streaming_ = plan_node_->streaming(); + + // Initialize cursor state + current_offset_ = 0; + has_more_data_ = true; + current_block_index_ = 0; + + // Extract time filtering parameters from plan node + timestamp_column_ = plan_node_->timestamp_column(); + partition_column_ = plan_node_->partition_column(); + + // Convert start/end times from nanoseconds to whole seconds for ClickHouse DateTime. A row at + // second S is at S * 1e9 ns, so the start bound rounds up and the end bound rounds down to + // avoid including rows outside of [start_time, end_time]. + if (plan_node_->start_time() > 0) { + start_time_ = + std::chrono::ceil(std::chrono::nanoseconds(plan_node_->start_time())) + .count(); + } + if (plan_node_->end_time() > 0) { + end_time_ = + std::chrono::floor(std::chrono::nanoseconds(plan_node_->end_time())) + .count(); + } + + // The source node appends its own LIMIT/OFFSET for pagination, which would conflict with a + // LIMIT in the base query. + if (FindKeyword(base_query_, "limit") != std::string::npos) { + return error::InvalidArgument("ClickHouse source query must not contain a LIMIT clause: $0", + base_query_); + } + + return Status::OK(); +} + +Status ClickHouseSourceNode::PrepareImpl(ExecState*) { return Status::OK(); } + +Status ClickHouseSourceNode::OpenImpl(ExecState* exec_state) { + if (!partition_column_.empty()) { + auto metadata_state = exec_state->metadata_state(); + if (metadata_state == nullptr) { + return error::FailedPrecondition( + "Agent metadata state is required to filter on partition column $0", partition_column_); + } + hostname_ = metadata_state->hostname(); + } + + // Create ClickHouse client + clickhouse::ClientOptions options; + options.SetHost(host_); + options.SetPort(port_); + options.SetUser(username_); + options.SetPassword(password_); + options.SetDefaultDatabase(database_); + + PX_ASSIGN_OR_RETURN(client_, ch::CreateClient(options)); + + return Status::OK(); +} + +Status ClickHouseSourceNode::CloseImpl(ExecState*) { + client_.reset(); + current_batch_blocks_.clear(); + + // Reset cursor state + current_offset_ = 0; + current_block_index_ = 0; + has_more_data_ = true; + + return Status::OK(); +} + +StatusOr ClickHouseSourceNode::ClickHouseTypeToPixieType( + const clickhouse::TypeRef& ch_type) { + const auto& type_name = ch_type->GetName(); + + // Integer types - Pixie only supports INT64 + if (type_name == "UInt8" || type_name == "UInt16" || type_name == "UInt32" || + type_name == "UInt64" || type_name == "Int8" || type_name == "Int16" || + type_name == "Int32" || type_name == "Int64") { + return types::DataType::INT64; + } + + // UInt128 + if (type_name == "UInt128") { + return types::DataType::UINT128; + } + + // Floating point types - Pixie only supports FLOAT64 + if (type_name == "Float32" || type_name == "Float64") { + return types::DataType::FLOAT64; + } + + // String types + if (type_name == "String" || type_name == "FixedString") { + return types::DataType::STRING; + } + + // Date/time types + if (type_name == "DateTime" || type_name.find("DateTime64") == 0) { + return types::DataType::TIME64NS; + } + + // Boolean + if (type_name == "Bool") { + return types::DataType::BOOLEAN; + } + + return error::InvalidArgument("Unsupported ClickHouse type: $0", type_name); +} + +StatusOr> ClickHouseSourceNode::ConvertClickHouseBlockToRowBatch( + const clickhouse::Block& block, bool /*is_last_block*/) { + auto num_rows = block.GetRowCount(); + auto num_cols = block.GetColumnCount(); + + // Create output row descriptor if this is the first block + if (current_block_index_ == 0) { + std::vector col_types; + for (size_t i = 0; i < num_cols; ++i) { + PX_ASSIGN_OR_RETURN(auto pixie_type, ClickHouseTypeToPixieType(block[i]->Type())); + col_types.push_back(pixie_type); + } + // TODO(ddelnano): Consider fetching column names from the plan or from ClickHouse metadata + } + + auto row_batch = std::make_unique(*output_descriptor_, num_rows); + + // Convert each column + for (size_t col_idx = 0; col_idx < num_cols; ++col_idx) { + const auto& ch_column = block[col_idx]; + const auto& type_name = ch_column->Type()->GetName(); + + auto expected_type = output_descriptor_->type(col_idx); + + // Special case: String in ClickHouse that should be UINT128 in Pixie + if (type_name == "String" && expected_type == types::DataType::UINT128) { + auto typed_col = ch_column->As(); + auto builder = + types::MakeArrowBuilder(types::DataType::UINT128, arrow::default_memory_pool()); + PX_RETURN_IF_ERROR(builder->Reserve(num_rows)); + + for (size_t i = 0; i < num_rows; ++i) { + std::string value(typed_col->At(i)); + + // Parse "high:low" format + size_t colon_pos = value.find(':'); + if (colon_pos == std::string::npos) { + return error::InvalidArgument("Invalid UINT128 string format: $0 (expected high:low)", + value); + } + + uint64_t high = std::stoull(value.substr(0, colon_pos)); + uint64_t low = std::stoull(value.substr(colon_pos + 1)); + absl::uint128 uint128_val = absl::MakeUint128(high, low); + + PX_RETURN_IF_ERROR( + table_store::schema::CopyValue(builder.get(), uint128_val)); + } + + std::shared_ptr array; + PX_RETURN_IF_ERROR(builder->Finish(&array)); + PX_RETURN_IF_ERROR(row_batch->AddColumn(array)); + + continue; + } + + // Integer types - all map to INT64 in Pixie + if (type_name == "UInt8") { + auto typed_col = ch_column->As(); + arrow::BooleanBuilder builder; + PX_RETURN_IF_ERROR(builder.Reserve(num_rows)); + for (size_t i = 0; i < num_rows; ++i) { + builder.UnsafeAppend(typed_col->At(i) != 0); + } + std::shared_ptr array; + PX_RETURN_IF_ERROR(builder.Finish(&array)); + PX_RETURN_IF_ERROR(row_batch->AddColumn(array)); + } else if (type_name == "UInt16") { + auto typed_col = ch_column->As(); + arrow::Int64Builder builder; + PX_RETURN_IF_ERROR(builder.Reserve(num_rows)); + for (size_t i = 0; i < num_rows; ++i) { + builder.UnsafeAppend(static_cast(typed_col->At(i))); + } + std::shared_ptr array; + PX_RETURN_IF_ERROR(builder.Finish(&array)); + PX_RETURN_IF_ERROR(row_batch->AddColumn(array)); + } else if (type_name == "UInt32") { + auto typed_col = ch_column->As(); + arrow::Int64Builder builder; + PX_RETURN_IF_ERROR(builder.Reserve(num_rows)); + for (size_t i = 0; i < num_rows; ++i) { + builder.UnsafeAppend(static_cast(typed_col->At(i))); + } + std::shared_ptr array; + PX_RETURN_IF_ERROR(builder.Finish(&array)); + PX_RETURN_IF_ERROR(row_batch->AddColumn(array)); + } else if (type_name == "UInt64") { + auto typed_col = ch_column->As(); + arrow::Int64Builder builder; + PX_RETURN_IF_ERROR(builder.Reserve(num_rows)); + for (size_t i = 0; i < num_rows; ++i) { + builder.UnsafeAppend(static_cast(typed_col->At(i))); + } + std::shared_ptr array; + PX_RETURN_IF_ERROR(builder.Finish(&array)); + PX_RETURN_IF_ERROR(row_batch->AddColumn(array)); + } else if (type_name == "Int8") { + auto typed_col = ch_column->As(); + arrow::Int64Builder builder; + PX_RETURN_IF_ERROR(builder.Reserve(num_rows)); + for (size_t i = 0; i < num_rows; ++i) { + builder.UnsafeAppend(static_cast(typed_col->At(i))); + } + std::shared_ptr array; + PX_RETURN_IF_ERROR(builder.Finish(&array)); + PX_RETURN_IF_ERROR(row_batch->AddColumn(array)); + } else if (type_name == "Int16") { + auto typed_col = ch_column->As(); + arrow::Int64Builder builder; + PX_RETURN_IF_ERROR(builder.Reserve(num_rows)); + for (size_t i = 0; i < num_rows; ++i) { + builder.UnsafeAppend(static_cast(typed_col->At(i))); + } + std::shared_ptr array; + PX_RETURN_IF_ERROR(builder.Finish(&array)); + PX_RETURN_IF_ERROR(row_batch->AddColumn(array)); + } else if (type_name == "Int32") { + auto typed_col = ch_column->As(); + arrow::Int64Builder builder; + PX_RETURN_IF_ERROR(builder.Reserve(num_rows)); + for (size_t i = 0; i < num_rows; ++i) { + builder.UnsafeAppend(static_cast(typed_col->At(i))); + } + std::shared_ptr array; + PX_RETURN_IF_ERROR(builder.Finish(&array)); + PX_RETURN_IF_ERROR(row_batch->AddColumn(array)); + } else if (type_name == "Int64") { + auto typed_col = ch_column->As(); + arrow::Int64Builder builder; + PX_RETURN_IF_ERROR(builder.Reserve(num_rows)); + for (size_t i = 0; i < num_rows; ++i) { + builder.UnsafeAppend(typed_col->At(i)); + } + std::shared_ptr array; + PX_RETURN_IF_ERROR(builder.Finish(&array)); + PX_RETURN_IF_ERROR(row_batch->AddColumn(array)); + } else if (type_name == "String") { + auto typed_col = ch_column->As(); + arrow::StringBuilder builder; + PX_RETURN_IF_ERROR(builder.Reserve(num_rows)); + + for (size_t i = 0; i < num_rows; ++i) { + std::string value(typed_col->At(i)); + PX_RETURN_IF_ERROR(builder.Append(value)); + } + + std::shared_ptr array; + PX_RETURN_IF_ERROR(builder.Finish(&array)); + PX_RETURN_IF_ERROR(row_batch->AddColumn(array)); + + } else if (type_name == "Float32") { + auto typed_col = ch_column->As(); + arrow::DoubleBuilder builder; + PX_RETURN_IF_ERROR(builder.Reserve(num_rows)); + for (size_t i = 0; i < num_rows; ++i) { + builder.UnsafeAppend(static_cast(typed_col->At(i))); + } + std::shared_ptr array; + PX_RETURN_IF_ERROR(builder.Finish(&array)); + PX_RETURN_IF_ERROR(row_batch->AddColumn(array)); + } else if (type_name == "Float64") { + auto typed_col = ch_column->As(); + arrow::DoubleBuilder builder; + PX_RETURN_IF_ERROR(builder.Reserve(num_rows)); + for (size_t i = 0; i < num_rows; ++i) { + builder.UnsafeAppend(typed_col->At(i)); + } + std::shared_ptr array; + PX_RETURN_IF_ERROR(builder.Finish(&array)); + PX_RETURN_IF_ERROR(row_batch->AddColumn(array)); + } else if (type_name == "Bool") { + auto typed_col = ch_column->As(); + arrow::BooleanBuilder builder; + PX_RETURN_IF_ERROR(builder.Reserve(num_rows)); + for (size_t i = 0; i < num_rows; ++i) { + builder.UnsafeAppend(typed_col->At(i) != 0); + } + std::shared_ptr array; + PX_RETURN_IF_ERROR(builder.Finish(&array)); + PX_RETURN_IF_ERROR(row_batch->AddColumn(array)); + } else if (type_name == "DateTime") { + auto typed_col = ch_column->As(); + arrow::Time64Builder builder(arrow::time64(arrow::TimeUnit::NANO), + arrow::default_memory_pool()); + PX_RETURN_IF_ERROR(builder.Reserve(num_rows)); + + for (size_t i = 0; i < num_rows; ++i) { + // Convert DateTime (seconds since epoch) to nanoseconds + int64_t ns = static_cast(typed_col->At(i)) * 1000000000LL; + builder.UnsafeAppend(ns); + } + + std::shared_ptr array; + PX_RETURN_IF_ERROR(builder.Finish(&array)); + PX_RETURN_IF_ERROR(row_batch->AddColumn(array)); + + } else if (type_name.find("DateTime64") == 0) { + auto typed_col = ch_column->As(); + arrow::Time64Builder builder(arrow::time64(arrow::TimeUnit::NANO), + arrow::default_memory_pool()); + PX_RETURN_IF_ERROR(builder.Reserve(num_rows)); + + // DateTime64(P) stores ticks of 10^-P seconds. ClickHouse caps P at 9 (nanoseconds). + const size_t precision = typed_col->GetPrecision(); + if (precision > 9) { + return error::InvalidArgument("Unsupported DateTime64 precision: $0", precision); + } + int64_t ns_per_tick = 1; + for (size_t p = precision; p < 9; ++p) { + ns_per_tick *= 10; + } + + for (size_t i = 0; i < num_rows; ++i) { + builder.UnsafeAppend(typed_col->At(i) * ns_per_tick); + } + + std::shared_ptr array; + PX_RETURN_IF_ERROR(builder.Finish(&array)); + PX_RETURN_IF_ERROR(row_batch->AddColumn(array)); + + } else { + return error::InvalidArgument("Unsupported ClickHouse type for conversion: $0", type_name); + } + } + + // Set end-of-window and end-of-stream flags + // Don't set them here - they should be set in GenerateNextImpl + row_batch->set_eow(false); + row_batch->set_eos(false); + + return row_batch; +} + +std::string ClickHouseSourceNode::BuildQuery() { + std::string query = base_query_; + std::vector conditions; + + // Add time filtering if start/end times are specified and timestamp column is set + if (!timestamp_column_.empty()) { + if (start_time_.has_value()) { + conditions.push_back(absl::Substitute("$0 >= $1", timestamp_column_, start_time_.value())); + } + if (end_time_.has_value()) { + conditions.push_back(absl::Substitute("$0 <= $1", timestamp_column_, end_time_.value())); + } + } + + // Add partition column filtering if specified + if (!partition_column_.empty()) { + // Since queries originate from PEMs, the partition column assumes that + // each GEM (DaemonSet) will query for its data set using its hostname. + conditions.push_back( + absl::Substitute("$0 = '$1'", partition_column_, EscapeStringLiteral(hostname_))); + } + + if (!conditions.empty()) { + std::string conditions_clause = absl::StrJoin(conditions, " AND "); + size_t where_pos = FindKeyword(query, "where"); + + if (where_pos != std::string::npos) { + // Parenthesize the existing condition so that an OR within it can't bypass the added + // conditions. + size_t cond_start = where_pos + std::string_view(" where ").length(); + size_t cond_end = FindClauseEnd(query, cond_start, {"group by", "having", "order by"}); + query.insert(cond_end, absl::StrCat(") AND ", conditions_clause)); + query.insert(cond_start, "("); + } else { + size_t insert_pos = FindClauseEnd(query, 0, {"group by", "having", "order by"}); + query.insert(insert_pos, absl::StrCat(" WHERE ", conditions_clause)); + } + } + + // Add ORDER BY clause if needed + if (FindKeyword(query, "order by") == std::string::npos) { + if (!timestamp_column_.empty()) { + query += absl::Substitute(" ORDER BY $0", timestamp_column_); + } else { + // Fall back to ordering by first column for consistent pagination + query += " ORDER BY 1"; + } + } + + // Add LIMIT and OFFSET for pagination + query += absl::Substitute(" LIMIT $0 OFFSET $1", batch_size_, current_offset_); + + return query; +} + +Status ClickHouseSourceNode::ExecuteBatchQuery() { + // Clear previous batch results + current_batch_blocks_.clear(); + current_block_index_ = 0; + + if (!has_more_data_) { + return Status::OK(); + } + + std::string query = BuildQuery(); + VLOG(1) << "Executing ClickHouse query: " << query; + + size_t rows_received = 0; + PX_RETURN_IF_ERROR( + ch::Select(client_.get(), query, [this, &rows_received](const clickhouse::Block& block) { + // Only store non-empty blocks + if (block.GetRowCount() > 0) { + VLOG(1) << "Received block with " << block.GetRowCount() << " rows"; + current_batch_blocks_.push_back(block); + rows_received += block.GetRowCount(); + } + })); + + VLOG(1) << "Total rows received: " << rows_received << ", batch size: " << batch_size_; + + // Update cursor state + current_offset_ += rows_received; + if (rows_received < batch_size_) { + // We got fewer rows than requested, so no more data available + has_more_data_ = false; + } + + return Status::OK(); +} + +Status ClickHouseSourceNode::GenerateNextImpl(ExecState* exec_state) { + // If we need to fetch more data + if (current_block_index_ >= current_batch_blocks_.size()) { + current_block_index_ = 0; + current_batch_blocks_.clear(); + + if (!has_more_data_) { + // No more data available - send empty batch with eos=true + PX_ASSIGN_OR_RETURN(auto empty_batch, + RowBatch::WithZeroRows(*output_descriptor_, true, true)); + PX_RETURN_IF_ERROR(SendRowBatchToChildren(exec_state, *empty_batch)); + return Status::OK(); + } + + // Fetch next batch from ClickHouse + PX_RETURN_IF_ERROR(ExecuteBatchQuery()); + + // If still no blocks after fetching, we're done + if (current_batch_blocks_.empty()) { + PX_ASSIGN_OR_RETURN(auto empty_batch, + RowBatch::WithZeroRows(*output_descriptor_, true, true)); + PX_RETURN_IF_ERROR(SendRowBatchToChildren(exec_state, *empty_batch)); + return Status::OK(); + } + } + + // Calculate total rows in all blocks + size_t total_rows = 0; + for (const auto& block : current_batch_blocks_) { + total_rows += block.GetRowCount(); + } + + // Create a merged RowBatch + auto merged_batch = std::make_unique(*output_descriptor_, total_rows); + + // Process each column + for (size_t col_idx = 0; col_idx < output_descriptor_->size(); ++col_idx) { + // Get the data type from output descriptor + auto data_type = output_descriptor_->type(col_idx); + + // Create appropriate builder based on data type + std::shared_ptr builder; + switch (data_type) { + case types::DataType::INT64: + builder = std::make_shared(); + break; + case types::DataType::UINT128: + builder = types::MakeArrowBuilder(types::DataType::UINT128, arrow::default_memory_pool()); + break; + case types::DataType::FLOAT64: + builder = std::make_shared(); + break; + case types::DataType::STRING: + builder = std::make_shared(); + break; + case types::DataType::BOOLEAN: + builder = std::make_shared(); + break; + case types::DataType::TIME64NS: + builder = std::make_shared(arrow::time64(arrow::TimeUnit::NANO), + arrow::default_memory_pool()); + break; + default: + return error::InvalidArgument("Unsupported data type for column $0", col_idx); + } + + PX_RETURN_IF_ERROR(builder->Reserve(total_rows)); + + for (const auto& block : current_batch_blocks_) { + PX_ASSIGN_OR_RETURN(auto row_batch, ConvertClickHouseBlockToRowBatch(block, false)); + auto array = row_batch->ColumnAt(col_idx); + + switch (data_type) { + case types::DataType::INT64: { + auto typed_array = std::static_pointer_cast(array); + auto typed_builder = std::static_pointer_cast(builder); + for (int i = 0; i < typed_array->length(); i++) { + if (typed_array->IsNull(i)) { + PX_RETURN_IF_ERROR(typed_builder->AppendNull()); + } else { + typed_builder->UnsafeAppend(typed_array->Value(i)); + } + } + break; + } + case types::DataType::UINT128: { + auto typed_array = std::static_pointer_cast(array); + for (int i = 0; i < typed_array->length(); i++) { + if (typed_array->IsNull(i)) { + PX_RETURN_IF_ERROR(builder->AppendNull()); + } else { + auto val = types::GetValueFromArrowArray(array.get(), i); + PX_RETURN_IF_ERROR( + table_store::schema::CopyValue(builder.get(), val)); + } + } + break; + } + case types::DataType::TIME64NS: { + auto typed_array = std::static_pointer_cast(array); + auto typed_builder = std::static_pointer_cast(builder); + for (int i = 0; i < typed_array->length(); i++) { + if (typed_array->IsNull(i)) { + PX_RETURN_IF_ERROR(typed_builder->AppendNull()); + } else { + typed_builder->UnsafeAppend(typed_array->Value(i)); + } + } + break; + } + case types::DataType::FLOAT64: { + auto typed_array = std::static_pointer_cast(array); + auto typed_builder = std::static_pointer_cast(builder); + for (int i = 0; i < typed_array->length(); i++) { + if (typed_array->IsNull(i)) { + PX_RETURN_IF_ERROR(typed_builder->AppendNull()); + } else { + typed_builder->UnsafeAppend(typed_array->Value(i)); + } + } + break; + } + case types::DataType::STRING: { + auto typed_array = std::static_pointer_cast(array); + auto typed_builder = std::static_pointer_cast(builder); + for (int i = 0; i < typed_array->length(); i++) { + if (typed_array->IsNull(i)) { + PX_RETURN_IF_ERROR(typed_builder->AppendNull()); + } else { + PX_RETURN_IF_ERROR(typed_builder->Append(typed_array->GetString(i))); + } + } + break; + } + case types::DataType::BOOLEAN: { + auto typed_array = std::static_pointer_cast(array); + auto typed_builder = std::static_pointer_cast(builder); + for (int i = 0; i < typed_array->length(); i++) { + if (typed_array->IsNull(i)) { + PX_RETURN_IF_ERROR(typed_builder->AppendNull()); + } else { + typed_builder->UnsafeAppend(typed_array->Value(i)); + } + } + break; + } + default: + return error::InvalidArgument("Unsupported data type for column $0", col_idx); + } + } + + // Finish building and add column + std::shared_ptr merged_array; + PX_RETURN_IF_ERROR(builder->Finish(&merged_array)); + PX_RETURN_IF_ERROR(merged_batch->AddColumn(merged_array)); + } + + // Set proper end-of-window and end-of-stream flags + bool is_last_batch = !has_more_data_; + if (is_last_batch) { + merged_batch->set_eow(true); + merged_batch->set_eos(true); + } else { + merged_batch->set_eow(false); + merged_batch->set_eos(false); + } + + rows_processed_ += merged_batch->num_rows(); + bytes_processed_ += merged_batch->NumBytes(); + + PX_RETURN_IF_ERROR(SendRowBatchToChildren(exec_state, *merged_batch)); + + current_block_index_ = current_batch_blocks_.size(); + + return Status::OK(); +} + +bool ClickHouseSourceNode::NextBatchReady() { + // We're ready if we have blocks in current batch or if we can fetch more data + return (current_block_index_ < current_batch_blocks_.size()) || has_more_data_; +} + +} // namespace exec +} // namespace carnot +} // namespace px diff --git a/src/carnot/exec/clickhouse_source_node.h b/src/carnot/exec/clickhouse_source_node.h new file mode 100644 index 00000000000..d5297249c5b --- /dev/null +++ b/src/carnot/exec/clickhouse_source_node.h @@ -0,0 +1,108 @@ +/* + * Copyright 2018- The Pixie Authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include + +#include +#include +#include +#include + +#include "src/carnot/exec/exec_node.h" +#include "src/carnot/exec/exec_state.h" +#include "src/carnot/plan/operators.h" +#include "src/common/base/base.h" +#include "src/common/base/status.h" +#include "src/shared/types/types.h" +#include "src/table_store/schema/row_batch.h" + +namespace px { +namespace carnot { +namespace exec { + +using table_store::schema::RowBatch; +using table_store::schema::RowDescriptor; + +class ClickHouseSourceNode : public SourceNode { + public: + ClickHouseSourceNode() = default; + virtual ~ClickHouseSourceNode() = default; + + bool NextBatchReady() override; + + protected: + std::string DebugStringImpl() override; + Status InitImpl(const plan::Operator& plan_node) override; + Status PrepareImpl(ExecState* exec_state) override; + Status OpenImpl(ExecState* exec_state) override; + Status CloseImpl(ExecState* exec_state) override; + Status GenerateNextImpl(ExecState* exec_state) override; + + private: + // Convert ClickHouse column types to Pixie data types + StatusOr ClickHouseTypeToPixieType(const clickhouse::TypeRef& ch_type); + + // Convert ClickHouse block to Pixie RowBatch + StatusOr> ConvertClickHouseBlockToRowBatch( + const clickhouse::Block& block, bool is_last_block); + + // Execute a batch query + Status ExecuteBatchQuery(); + + // Build the query with time filtering and pagination + std::string BuildQuery(); + + // Connection information + std::string host_; + int port_; + std::string username_; + std::string password_; + std::string database_; + std::string base_query_; + + // Batch size and cursor tracking + size_t batch_size_ = 1024; + size_t current_offset_ = 0; + bool has_more_data_ = true; + + // Time filtering + std::optional start_time_; + std::optional end_time_; + std::string timestamp_column_; // Column to use for timestamp-based filtering and ordering + std::string partition_column_; // Column used for partitioning + std::string hostname_; // Agent hostname matched against partition_column_ + + // ClickHouse client + std::unique_ptr client_; + + // Current batch results + std::vector current_batch_blocks_; + size_t current_block_index_ = 0; + + // Streaming support + bool streaming_ = false; + + // Plan node + std::unique_ptr plan_node_; +}; + +} // namespace exec +} // namespace carnot +} // namespace px diff --git a/src/carnot/exec/clickhouse_source_node_test.cc b/src/carnot/exec/clickhouse_source_node_test.cc new file mode 100644 index 00000000000..8798da6eaf1 --- /dev/null +++ b/src/carnot/exec/clickhouse_source_node_test.cc @@ -0,0 +1,415 @@ +/* + * Copyright 2018- The Pixie Authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "src/carnot/exec/clickhouse_source_node.h" + +#include + +#include + +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +#include "src/carnot/exec/test_utils.h" +#include "src/carnot/planpb/plan.pb.h" +#include "src/carnot/planpb/test_proto.h" +#include "src/carnot/udf/registry.h" +#include "src/common/ch/utils.h" +#include "src/common/event/time_system.h" +#include "src/common/testing/test_utils/container_runner.h" +#include "src/common/testing/testing.h" +#include "src/shared/metadata/metadata_state.h" +#include "src/shared/types/types.h" + +namespace px { +namespace carnot { +namespace exec { + +using table_store::schema::RowDescriptor; + +class ClickHouseSourceNodeTest : public ::testing::Test { + protected: + static constexpr char kClickHouseImage[] = + "src/stirling/source_connectors/socket_tracer/testing/container_images/clickhouse.tar"; + static constexpr char kClickHouseReadyMessage[] = "Ready for connections"; + static constexpr int kClickHousePort = 9000; + + void SetUp() override { + // Set up function registry and exec state + func_registry_ = std::make_unique("test_registry"); + auto table_store = std::make_shared(); + exec_state_ = std::make_unique( + func_registry_.get(), table_store, MockResultSinkStubGenerator, MockMetricsStubGenerator, + MockTraceStubGenerator, MockLogStubGenerator, sole::uuid4(), nullptr); + + // Create a minimal agent metadata state for test execution + auto metadata_state = + std::make_shared("test_host", // hostname + 1, // asid + getpid(), // pid + sole::uuid4(), // agent_id + "", // pod_name + sole::uuid4(), // vizier_id + "test_vizier", // vizier_name + "", // vizier_namespace + time_system_.get()); // time_system + exec_state_->set_metadata_state(metadata_state); + + clickhouse_server_ = + std::make_unique(px::testing::BazelRunfilePath(kClickHouseImage), + "clickhouse_test", kClickHouseReadyMessage); + + std::vector options = { + absl::Substitute("--publish=$0:$0", kClickHousePort), + "--env=CLICKHOUSE_PASSWORD=test_password", + "--network=host", + }; + + ASSERT_OK(clickhouse_server_->Run(std::chrono::seconds{60}, options, {}, true, + std::chrono::seconds{300})); + + // Give ClickHouse time to initialize + std::this_thread::sleep_for(std::chrono::seconds(5)); + + SetupClickHouseClient(); + CreateTestTable(); + } + + void TearDown() override { + if (client_) { + client_.reset(); + } + } + + void SetupClickHouseClient() { + clickhouse::ClientOptions client_options; + client_options.SetHost("localhost"); + client_options.SetPort(kClickHousePort); + client_options.SetUser("default"); + client_options.SetPassword("test_password"); + client_options.SetDefaultDatabase("default"); + + const int kMaxRetries = 5; + Status connect_status; + for (int i = 0; i < kMaxRetries; ++i) { + LOG(INFO) << "Attempting to connect to ClickHouse (attempt " << (i + 1) << "/" << kMaxRetries + << ")..."; + auto client_or = ch::CreateClient(client_options); + if (client_or.ok()) { + client_ = client_or.ConsumeValueOrDie(); + connect_status = ch::Execute(client_.get(), "SELECT 1"); + if (connect_status.ok()) { + return; + } + } else { + connect_status = client_or.status(); + } + LOG(WARNING) << "Failed to connect: " << connect_status.msg(); + client_.reset(); + if (i < kMaxRetries - 1) { + std::this_thread::sleep_for(std::chrono::seconds(2)); + } + } + FAIL() << "Failed to connect to ClickHouse after " << kMaxRetries + << " attempts: " << connect_status.msg(); + } + + void CreateTestTable() { + ASSERT_OK(ch::Execute(client_.get(), "DROP TABLE IF EXISTS test_table")); + + ASSERT_OK(ch::Execute(client_.get(), R"( + CREATE TABLE test_table ( + id UInt64, + name String, + value Float64, + timestamp DateTime, + partition_key String + ) ENGINE = MergeTree() + PARTITION BY (timestamp, partition_key) + ORDER BY timestamp + )")); + + auto id_col = std::make_shared(); + auto name_col = std::make_shared(); + auto value_col = std::make_shared(); + auto timestamp_col = std::make_shared(); + auto partition_key_col = std::make_shared(); + + // Add test data with increasing timestamps + std::time_t base_time = std::time(nullptr) - 3600; // Start 1 hour ago + id_col->Append(1); + name_col->Append("test1"); + value_col->Append(10.5); + timestamp_col->Append(base_time); + partition_key_col->Append("partition_a"); + + id_col->Append(2); + name_col->Append("test2"); + value_col->Append(20.5); + timestamp_col->Append(base_time + 1800); // 30 minutes later + partition_key_col->Append("partition_a"); + + id_col->Append(3); + name_col->Append("test3"); + value_col->Append(30.5); + timestamp_col->Append(base_time + 3600); // 1 hour later + partition_key_col->Append("partition_b"); + + clickhouse::Block block; + block.AppendColumn("id", id_col); + block.AppendColumn("name", name_col); + block.AppendColumn("value", value_col); + block.AppendColumn("timestamp", timestamp_col); + block.AppendColumn("partition_key", partition_key_col); + + ASSERT_OK(ch::Insert(client_.get(), "test_table", block)); + + LOG(INFO) << "Test table created and populated successfully"; + } + + // Creates a source operator for query over the (id, name, value) columns with a time range + // covering all of the test data. + std::unique_ptr CreatePlanNode( + const std::string& query, + std::function modify = {}) { + planpb::Operator op; + op.set_op_type(planpb::OperatorType::CLICKHOUSE_SOURCE_OPERATOR); + auto* ch_op = op.mutable_clickhouse_source_op(); + ch_op->set_host("localhost"); + ch_op->set_port(kClickHousePort); + ch_op->set_username("default"); + ch_op->set_password("test_password"); + ch_op->set_database("default"); + ch_op->set_query(query); + ch_op->set_batch_size(1024); + ch_op->set_streaming(false); + ch_op->add_column_names("id"); + ch_op->add_column_names("name"); + ch_op->add_column_names("value"); + ch_op->add_column_types(types::DataType::INT64); + ch_op->add_column_types(types::DataType::STRING); + ch_op->add_column_types(types::DataType::FLOAT64); + ch_op->set_timestamp_column("timestamp"); + ch_op->set_start_time(1000000000000000000LL); // Year 2001 in nanoseconds + ch_op->set_end_time(9223372036854775807LL); // Max int64 + + if (modify) { + modify(ch_op); + } + return plan::ClickHouseSourceOperator::FromProto(op, 1); + } + + std::unique_ptr clickhouse_server_; + std::unique_ptr client_; + std::unique_ptr exec_state_; + std::unique_ptr func_registry_; + std::unique_ptr time_system_ = std::make_unique(); +}; + +TEST_F(ClickHouseSourceNodeTest, BasicQuery) { + // Create ClickHouse source operator proto + auto op_proto = planpb::testutils::CreateClickHouseSourceOperatorPB(); + std::unique_ptr plan_node = + plan::ClickHouseSourceOperator::FromProto(op_proto, 1); + + RowDescriptor output_rd( + {types::DataType::INT64, types::DataType::STRING, types::DataType::FLOAT64}); + + auto tester = exec::ExecNodeTester( + *plan_node, output_rd, std::vector({}), exec_state_.get()); + + // Verify state machine behavior + EXPECT_TRUE(tester.node()->HasBatchesRemaining()); + + // First batch should return 2 rows (batch_size = 2) + tester.GenerateNextResult().ExpectRowBatch( + RowBatchBuilder(output_rd, 2, /*eow*/ false, /*eos*/ false) + .AddColumn({1, 2}) + .AddColumn({"test1", "test2"}) + .AddColumn({10.5, 20.5}) + .get()); + + // Second batch should return remaining 1 row with eos + EXPECT_TRUE(tester.node()->HasBatchesRemaining()); + tester.GenerateNextResult().ExpectRowBatch( + RowBatchBuilder(output_rd, 1, /*eow*/ true, /*eos*/ true) + .AddColumn({3}) + .AddColumn({"test3"}) + .AddColumn({30.5}) + .get()); + + EXPECT_FALSE(tester.node()->HasBatchesRemaining()); + tester.Close(); + + EXPECT_EQ(3, tester.node()->RowsProcessed()); + EXPECT_GT(tester.node()->BytesProcessed(), 0); +} + +TEST_F(ClickHouseSourceNodeTest, EmptyResultSet) { + // Create a table with no data + ASSERT_OK(ch::Execute(client_.get(), "DROP TABLE IF EXISTS empty_table")); + ASSERT_OK(ch::Execute(client_.get(), R"( + CREATE TABLE empty_table ( + id UInt64, + name String, + value Float64, + timestamp DateTime, + partition_key String + ) ENGINE = MergeTree() + PARTITION BY (timestamp, partition_key) + ORDER BY timestamp + )")); + + auto plan_node = CreatePlanNode("SELECT id, name, value FROM empty_table"); + RowDescriptor output_rd( + {types::DataType::INT64, types::DataType::STRING, types::DataType::FLOAT64}); + + auto tester = exec::ExecNodeTester( + *plan_node, output_rd, std::vector({}), exec_state_.get()); + + EXPECT_TRUE(tester.node()->HasBatchesRemaining()); + + // Should return empty batch with eos=true + tester.GenerateNextResult().ExpectRowBatch( + RowBatchBuilder(output_rd, 0, /*eow*/ true, /*eos*/ true) + .AddColumn({}) + .AddColumn({}) + .AddColumn({}) + .get()); + + EXPECT_FALSE(tester.node()->HasBatchesRemaining()); + tester.Close(); + + EXPECT_EQ(0, tester.node()->RowsProcessed()); + EXPECT_EQ(0, tester.node()->BytesProcessed()); +} + +TEST_F(ClickHouseSourceNodeTest, FilteredQuery) { + auto plan_node = + CreatePlanNode("SELECT id, name, value FROM test_table WHERE value > 15.0 ORDER BY id"); + RowDescriptor output_rd( + {types::DataType::INT64, types::DataType::STRING, types::DataType::FLOAT64}); + + auto tester = exec::ExecNodeTester( + *plan_node, output_rd, std::vector({}), exec_state_.get()); + + EXPECT_TRUE(tester.node()->HasBatchesRemaining()); + + // Should return all filtered results in one batch (2 rows < batch_size) + tester.GenerateNextResult().ExpectRowBatch( + RowBatchBuilder(output_rd, 2, /*eow*/ true, /*eos*/ true) + .AddColumn({2, 3}) + .AddColumn({"test2", "test3"}) + .AddColumn({20.5, 30.5}) + .get()); + + EXPECT_FALSE(tester.node()->HasBatchesRemaining()); + tester.Close(); + + EXPECT_EQ(2, tester.node()->RowsProcessed()); + EXPECT_GT(tester.node()->BytesProcessed(), 0); +} + +TEST_F(ClickHouseSourceNodeTest, FiltersApplyToEntireOrCondition) { + // Only id 3 is within the last 15 minutes. If the time filter were appended without + // parenthesizing the existing condition, id 1 would also match via the OR. + int64_t start_time_ns = (std::time(nullptr) - 900) * 1000000000LL; + auto plan_node = CreatePlanNode( + "SELECT id, name, value FROM test_table WHERE id = 1 OR id = 3 ORDER BY id", + [&](planpb::ClickHouseSourceOperator* ch_op) { ch_op->set_start_time(start_time_ns); }); + RowDescriptor output_rd( + {types::DataType::INT64, types::DataType::STRING, types::DataType::FLOAT64}); + + auto tester = exec::ExecNodeTester( + *plan_node, output_rd, std::vector({}), exec_state_.get()); + + tester.GenerateNextResult().ExpectRowBatch( + RowBatchBuilder(output_rd, 1, /*eow*/ true, /*eos*/ true) + .AddColumn({3}) + .AddColumn({"test3"}) + .AddColumn({30.5}) + .get()); + tester.Close(); +} + +TEST_F(ClickHouseSourceNodeTest, MultiLineQuery) { + auto plan_node = CreatePlanNode(R"( + SELECT id, name, value + FROM test_table + WHERE value > 15.0 + ORDER BY id + )"); + RowDescriptor output_rd( + {types::DataType::INT64, types::DataType::STRING, types::DataType::FLOAT64}); + + auto tester = exec::ExecNodeTester( + *plan_node, output_rd, std::vector({}), exec_state_.get()); + + tester.GenerateNextResult().ExpectRowBatch( + RowBatchBuilder(output_rd, 2, /*eow*/ true, /*eos*/ true) + .AddColumn({2, 3}) + .AddColumn({"test2", "test3"}) + .AddColumn({20.5, 30.5}) + .get()); + tester.Close(); +} + +TEST_F(ClickHouseSourceNodeTest, PartitionFilterUsesAgentHostname) { + // The agent metadata state's hostname is "test_host" (see SetUp). + ASSERT_OK(ch::Execute(client_.get(), + "INSERT INTO test_table VALUES (4, 'test4', 40.5, now(), 'test_host')")); + auto plan_node = CreatePlanNode("SELECT id, name, value FROM test_table", + [](planpb::ClickHouseSourceOperator* ch_op) { + ch_op->set_partition_column("partition_key"); + }); + RowDescriptor output_rd( + {types::DataType::INT64, types::DataType::STRING, types::DataType::FLOAT64}); + + auto tester = exec::ExecNodeTester( + *plan_node, output_rd, std::vector({}), exec_state_.get()); + + tester.GenerateNextResult().ExpectRowBatch( + RowBatchBuilder(output_rd, 1, /*eow*/ true, /*eos*/ true) + .AddColumn({4}) + .AddColumn({"test4"}) + .AddColumn({40.5}) + .get()); + tester.Close(); +} + +TEST_F(ClickHouseSourceNodeTest, RejectsQueryWithLimit) { + auto plan_node = CreatePlanNode("SELECT id, name, value FROM test_table LIMIT 10"); + RowDescriptor output_rd( + {types::DataType::INT64, types::DataType::STRING, types::DataType::FLOAT64}); + + ClickHouseSourceNode node; + EXPECT_NOT_OK(node.Init(*plan_node, output_rd, {})); +} + +} // namespace exec +} // namespace carnot +} // namespace px diff --git a/src/carnot/exec/exec_graph.cc b/src/carnot/exec/exec_graph.cc index 705cf381e38..ecfbba8b49e 100644 --- a/src/carnot/exec/exec_graph.cc +++ b/src/carnot/exec/exec_graph.cc @@ -24,6 +24,8 @@ #include #include "src/carnot/exec/agg_node.h" +#include "src/carnot/exec/clickhouse_export_sink_node.h" +#include "src/carnot/exec/clickhouse_source_node.h" #include "src/carnot/exec/empty_source_node.h" #include "src/carnot/exec/equijoin_node.h" #include "src/carnot/exec/exec_node.h" @@ -108,6 +110,14 @@ Status ExecutionGraph::Init(table_store::schema::Schema* schema, plan::PlanState .OnOTelSink([&](auto& node) { return OnOperatorImpl(node, &descriptors); }) + .OnClickHouseSource([&](auto& node) { + return OnOperatorImpl(node, + &descriptors); + }) + .OnClickHouseExportSink([&](auto& node) { + return OnOperatorImpl( + node, &descriptors); + }) .Walk(pf_); } diff --git a/src/carnot/exec/exec_state.h b/src/carnot/exec/exec_state.h index 444d9298d06..2419e651dc9 100644 --- a/src/carnot/exec/exec_state.h +++ b/src/carnot/exec/exec_state.h @@ -207,6 +207,8 @@ class ExecState { } } + std::shared_ptr metadata_state() const { return metadata_state_; } + void set_metadata_state(std::shared_ptr metadata_state) { metadata_state_ = metadata_state; } diff --git a/src/common/ch/BUILD.bazel b/src/common/ch/BUILD.bazel new file mode 100644 index 00000000000..8b91a71726d --- /dev/null +++ b/src/common/ch/BUILD.bazel @@ -0,0 +1,28 @@ +# Copyright 2018- The Pixie Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# SPDX-License-Identifier: Apache-2.0 + +load("//bazel:pl_build_system.bzl", "pl_cc_library") + +package(default_visibility = ["//src:__subpackages__"]) + +pl_cc_library( + name = "cc_library", + srcs = ["utils.cc"], + hdrs = ["utils.h"], + deps = [ + "@com_github_clickhouse_clickhouse_cpp//:clickhouse_cpp", + ], +) diff --git a/src/common/ch/utils.cc b/src/common/ch/utils.cc new file mode 100644 index 00000000000..88d0f82d196 --- /dev/null +++ b/src/common/ch/utils.cc @@ -0,0 +1,68 @@ +/* + * Copyright 2018- The Pixie Authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "src/common/ch/utils.h" + +#include +#include +#include +#include + +namespace px { +namespace ch { + +StatusOr> CreateClient( + const ::clickhouse::ClientOptions& options) { + try { + return std::make_unique<::clickhouse::Client>(options); + } catch (const std::exception& e) { + return error::Internal("Failed to create ClickHouse client: $0", e.what()); + } +} + +Status Execute(::clickhouse::Client* client, const std::string& query) { + try { + client->Execute(query); + } catch (const std::exception& e) { + return error::Internal("Failed to execute ClickHouse statement: $0", e.what()); + } + return Status::OK(); +} + +Status Select(::clickhouse::Client* client, const std::string& query, + std::function block_cb) { + try { + client->Select(query, std::move(block_cb)); + } catch (const std::exception& e) { + return error::Internal("Failed to execute ClickHouse query: $0", e.what()); + } + return Status::OK(); +} + +Status Insert(::clickhouse::Client* client, const std::string& table_name, + const ::clickhouse::Block& block) { + try { + client->Insert(table_name, block); + } catch (const std::exception& e) { + return error::Internal("Failed to insert into ClickHouse table $0: $1", table_name, e.what()); + } + return Status::OK(); +} + +} // namespace ch +} // namespace px diff --git a/src/common/ch/utils.h b/src/common/ch/utils.h new file mode 100644 index 00000000000..4e67983c388 --- /dev/null +++ b/src/common/ch/utils.h @@ -0,0 +1,50 @@ +/* + * Copyright 2018- The Pixie Authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include + +#include +#include +#include + +#include "src/common/base/base.h" + +namespace px { +namespace ch { + +// Status-returning wrappers around the exception-throwing clickhouse-cpp client API. + +// Creates a ClickHouse client from the given options. +StatusOr> CreateClient( + const ::clickhouse::ClientOptions& options); + +// Executes a statement that returns no rows (e.g. DDL or a connection check). +Status Execute(::clickhouse::Client* client, const std::string& query); + +// Executes a SELECT query, invoking block_cb for each returned block. +Status Select(::clickhouse::Client* client, const std::string& query, + std::function block_cb); + +// Inserts a block into table_name. +Status Insert(::clickhouse::Client* client, const std::string& table_name, + const ::clickhouse::Block& block); + +} // namespace ch +} // namespace px