[carnot] Add ClickHouse source and export-sink planner IR

src/carnot/planner/compiler_state/compiler_state.h

@@ -119,7 +119,8 @@ class CompilerState : public NotCopyable {
                 int64_t max_output_rows_per_table, std::string_view result_address,
                 std::string_view result_ssl_targetname, const RedactionOptions& redaction_options,
                 std::unique_ptr<planpb::OTelEndpointConfig> endpoint_config,
-                std::unique_ptr<PluginConfig> plugin_config, DebugInfo debug_info)
+                std::unique_ptr<PluginConfig> plugin_config, DebugInfo debug_info,
+                std::unique_ptr<planpb::ClickHouseConfig> clickhouse_config = nullptr)
       : relation_map_(std::move(relation_map)),
         table_names_to_sensitive_columns_(table_names_to_sensitive_columns),
         registry_info_(registry_info),
@@ -130,7 +131,8 @@ class CompilerState : public NotCopyable {
         redaction_options_(redaction_options),
         endpoint_config_(std::move(endpoint_config)),
         plugin_config_(std::move(plugin_config)),
-        debug_info_(std::move(debug_info)) {}
+        debug_info_(std::move(debug_info)),
+        clickhouse_config_(std::move(clickhouse_config)) {}
 
   CompilerState() = delete;
 
@@ -175,6 +177,7 @@ class CompilerState : public NotCopyable {
   planpb::OTelEndpointConfig* endpoint_config() { return endpoint_config_.get(); }
   PluginConfig* plugin_config() { return plugin_config_.get(); }
   const DebugInfo& debug_info() { return debug_info_; }
+  planpb::ClickHouseConfig* clickhouse_config() { return clickhouse_config_.get(); }
 
  private:
   std::unique_ptr<RelationMap> relation_map_;
@@ -191,6 +194,7 @@ class CompilerState : public NotCopyable {
   std::unique_ptr<planpb::OTelEndpointConfig> endpoint_config_ = nullptr;
   std::unique_ptr<PluginConfig> plugin_config_ = nullptr;
   DebugInfo debug_info_;
+  std::unique_ptr<planpb::ClickHouseConfig> clickhouse_config_ = nullptr;
 };
 
 }  // namespace planner

src/carnot/planner/ir/BUILD.bazel

@@ -47,6 +47,7 @@ pl_cc_library(
         "//src/carnot/planpb:plan_pl_cc_proto",
         "//src/shared/metadata:cc_library",
         "//src/shared/metadatapb:metadata_pl_cc_proto",
+        "@com_github_clickhouse_clickhouse_cpp//:clickhouse_cpp",
         "@com_github_vinzenz_libpypa//:libpypa",
     ],
 )
@@ -67,6 +68,22 @@ pl_cc_test(
     ],
 )
 
+pl_cc_test(
+    name = "clickhouse_export_sink_ir_test",
+    srcs = ["clickhouse_export_sink_ir_test.cc"],
+    deps = [
+        "//src/carnot/planner/compiler:test_utils",
+    ],
+)
+
+pl_cc_test(
+    name = "clickhouse_source_ir_test",
+    srcs = ["clickhouse_source_ir_test.cc"],
+    deps = [
+        "//src/carnot/planner/compiler:test_utils",
+    ],
+)
+
 pl_cc_test(
     name = "pattern_match_test",
     srcs = ["pattern_match_test.cc"],

src/carnot/planner/ir/all_ir_nodes.h

@@ -20,6 +20,8 @@
 
 #include "src/carnot/planner/ir/blocking_agg_ir.h"
 #include "src/carnot/planner/ir/bool_ir.h"
+#include "src/carnot/planner/ir/clickhouse_source_ir.h"
+#include "src/carnot/planner/ir/clickhouse_export_sink_ir.h"
 #include "src/carnot/planner/ir/column_ir.h"
 #include "src/carnot/planner/ir/data_ir.h"
 #include "src/carnot/planner/ir/drop_ir.h"

src/carnot/planner/ir/clickhouse_export_sink_ir.cc

@@ -0,0 +1,135 @@
+/*
+ * 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/planner/ir/clickhouse_export_sink_ir.h"
+#include <regex>
+#include "src/carnot/planner/ir/ir.h"
+#include "src/carnot/planpb/plan.pb.h"
+
+namespace px {
+namespace carnot {
+namespace planner {
+
+StatusOr<std::vector<absl::flat_hash_set<std::string>>>
+ClickHouseExportSinkIR::RequiredInputColumns() const {
+  return std::vector<absl::flat_hash_set<std::string>>{required_column_names_};
+}
+
+Status ClickHouseExportSinkIR::Init(OperatorIR* parent, const std::string& table_name,
+                                    const std::string& clickhouse_dsn) {
+  table_name_ = table_name;
+
+  // Parse the ClickHouse DSN and initialize the config
+  PX_ASSIGN_OR_RETURN(auto config, ParseClickHouseDSN(clickhouse_dsn));
+  clickhouse_config_ = std::make_unique<planpb::ClickHouseConfig>(config);
+
+  return AddParent(parent);
+}
+
+StatusOr<planpb::ClickHouseConfig> ClickHouseExportSinkIR::ParseClickHouseDSN(
+    const std::string& dsn) {
+  // Expected format: [clickhouse://]username:password@host:port/database
+  // The clickhouse:// prefix is optional
+  std::regex dsn_regex(R"((?:clickhouse://)?([^:]+):([^@]+)@([^:]+):(\d+)/(.+))");
+  std::smatch matches;
+
+  if (!std::regex_match(dsn, matches, dsn_regex)) {
+    return error::InvalidArgument(
+        "Invalid ClickHouse DSN format. Expected: "
+        "[clickhouse://]username:password@host:port/database");
+  }
+
+  planpb::ClickHouseConfig config;
+
+  // Extract the components
+  config.set_username(matches[1].str());
+  config.set_password(matches[2].str());
+  config.set_host(matches[3].str());
+  config.set_port(std::stoi(matches[4].str()));
+  config.set_database(matches[5].str());
+
+  // hostname will be set by the runtime
+  config.set_hostname("");
+
+  return config;
+}
+
+Status ClickHouseExportSinkIR::ToProto(planpb::Operator* op) const {
+  op->set_op_type(planpb::CLICKHOUSE_EXPORT_SINK_OPERATOR);
+  auto clickhouse_op = op->mutable_clickhouse_sink_op();
+
+  // ClickHouse config must be set before calling ToProto
+  if (clickhouse_config_ == nullptr) {
+    return error::InvalidArgument("ClickHouse config not set");
+  }
+
+  // Set the ClickHouse configuration
+  *clickhouse_op->mutable_clickhouse_config() = *clickhouse_config_;
+  clickhouse_op->set_table_name(table_name_);
+
+  // Map all input columns to ClickHouse columns
+  DCHECK(is_type_resolved());
+  int64_t idx = 0;
+  for (const auto& [col_name, col_type] : *resolved_table_type()) {
+    DCHECK(col_type->IsValueType());
+    auto value_type = std::static_pointer_cast<ValueType>(col_type);
+
+    auto column_mapping = clickhouse_op->add_column_mappings();
+    column_mapping->set_input_column_index(idx);
+    column_mapping->set_clickhouse_column_name(col_name);
+    column_mapping->set_column_type(value_type->data_type());
+    idx++;
+  }
+
+  return Status::OK();
+}
+
+Status ClickHouseExportSinkIR::CopyFromNodeImpl(const IRNode* node,
+                                                absl::flat_hash_map<const IRNode*, IRNode*>*) {
+  const ClickHouseExportSinkIR* source = static_cast<const ClickHouseExportSinkIR*>(node);
+  table_name_ = source->table_name_;
+  required_column_names_ = source->required_column_names_;
+  if (source->clickhouse_config_ != nullptr) {
+    clickhouse_config_ = std::make_unique<planpb::ClickHouseConfig>(*source->clickhouse_config_);
+  }
+  return Status::OK();
+}
+
+Status ClickHouseExportSinkIR::ResolveType(CompilerState* compiler_state) {
+  DCHECK_EQ(1U, parent_types().size());
+
+  auto parent_table_type = std::static_pointer_cast<TableType>(parent_types()[0]);
+
+  // Store ClickHouse config from compiler state only if not already set by Init()
+  if (clickhouse_config_ == nullptr && compiler_state->clickhouse_config() != nullptr) {
+    clickhouse_config_ =
+        std::make_unique<planpb::ClickHouseConfig>(*compiler_state->clickhouse_config());
+  }
+
+  // Populate required column names
+  for (const auto& col_name : parent_table_type->ColumnNames()) {
+    required_column_names_.insert(col_name);
+  }
+
+  // Export sink passes through the input schema
+  return SetResolvedType(parent_table_type);
+}
+
+}  // namespace planner
+}  // namespace carnot
+}  // namespace px

src/carnot/planner/ir/clickhouse_export_sink_ir.h

@@ -0,0 +1,74 @@
+/*
+ * 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 <memory>
+#include <string>
+#include <vector>
+
+#include <absl/container/flat_hash_set.h>
+#include "src/carnot/planner/compiler_state/compiler_state.h"
+#include "src/carnot/planner/ir/column_ir.h"
+#include "src/carnot/planner/ir/operator_ir.h"
+#include "src/carnot/planpb/plan.pb.h"
+#include "src/common/base/base.h"
+
+namespace px {
+namespace carnot {
+namespace planner {
+
+/**
+ * @brief The IR representation for the ClickHouseExportSink operator.
+ * Represents a configuration to export a DataFrame to a ClickHouse database.
+ */
+class ClickHouseExportSinkIR : public OperatorIR {
+ public:
+  explicit ClickHouseExportSinkIR(int64_t id) : OperatorIR(id, IRNodeType::kClickHouseExportSink) {}
+
+  Status Init(OperatorIR* parent, const std::string& table_name, const std::string& clickhouse_dsn);
+
+  StatusOr<planpb::ClickHouseConfig> ParseClickHouseDSN(const std::string& dsn);
+
+  Status ToProto(planpb::Operator* op) const override;
+
+  Status CopyFromNodeImpl(const IRNode* node,
+                          absl::flat_hash_map<const IRNode*, IRNode*>*) override;
+
+  Status ResolveType(CompilerState* compiler_state);
+  inline bool IsBlocking() const override { return true; }
+
+  StatusOr<std::vector<absl::flat_hash_set<std::string>>> RequiredInputColumns() const override;
+
+  const std::string& table_name() const { return table_name_; }
+
+ protected:
+  StatusOr<absl::flat_hash_set<std::string>> PruneOutputColumnsToImpl(
+      const absl::flat_hash_set<std::string>& /*kept_columns*/) override {
+    return error::Unimplemented("Unexpected call to ClickHouseExportSinkIR::PruneOutputColumnsTo.");
+  }
+
+ private:
+  std::string table_name_;
+  absl::flat_hash_set<std::string> required_column_names_;
+  std::unique_ptr<planpb::ClickHouseConfig> clickhouse_config_;
+};
+
+}  // namespace planner
+}  // namespace carnot
+}  // namespace px