feat(sink): implement snowflake sink

src/connector/src/sink/mod.rs

@@ -32,6 +32,8 @@ pub mod nats;
 pub mod pulsar;
 pub mod redis;
 pub mod remote;
+pub mod snowflake;
+pub mod snowflake_connector;
 pub mod starrocks;
 pub mod test_sink;
 pub mod trivial;
@@ -88,6 +90,7 @@ macro_rules! for_all_sinks {
  { HttpJava, $crate::sink::remote::HttpJavaSink },
  { Doris, $crate::sink::doris::DorisSink },
  { Starrocks, $crate::sink::starrocks::StarrocksSink },
+ { Snowflake, $crate::sink::snowflake::SnowflakeSink },
  { DeltaLake, $crate::sink::deltalake::DeltaLakeSink },
  { BigQuery, $crate::sink::big_query::BigQuerySink },
  { Test, $crate::sink::test_sink::TestSink },
@@ -515,6 +518,8 @@ pub enum SinkError {
  ),
  #[error("Starrocks error: {0}")]
  Starrocks(String),
+ #[error("Snowflake error: {0}")]
+ Snowflake(String),
  #[error("Pulsar error: {0}")]
  Pulsar(
  #[source]

src/connector/src/sink/snowflake.rs

@@ -0,0 +1,155 @@
+// Copyright 2024 RisingWave Labs
+//
+// 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.
+
+use std::collections::HashMap;
+use std::sync::Arc;
+
+use anyhow::anyhow;
+use async_trait::async_trait;
+use risingwave_common::array::StreamChunk;
+use risingwave_common::buffer::Bitmap;
+use risingwave_common::catalog::Schema;
+use serde::Deserialize;
+use serde_derive::Serialize;
+use serde_json::Value;
+use serde_with::serde_as;
+use with_options::WithOptions;
+
+use super::encoder::JsonEncoder;
+use super::writer::LogSinkerOf;
+use super::{SinkError, SinkParam};
+use crate::sink::{DummySinkCommitCoordinator, Result, Sink, SinkWriter, SinkWriterParam};
+
+pub const SNOWFLAKE_SINK: &str = "snowflake";
+
+// TODO: add comments
+#[derive(Deserialize, Debug, Clone, WithOptions)]
+pub struct SnowflakeCommon {
+ #[serde(rename = "snowflake.database")]
+ pub database: String,
+
+ #[serde(rename = "snowflake.database.schema")]
+ pub schema: String,
+
+ #[serde(rename = "snowflake.database.schema.pipe")]
+ pub pipe: String,
+
+ #[serde(rename = "snowflake.account")]
+ pub account: String,
+
+ #[serde(rename = "snowflake.private.key")]
+ pub private_key: String,
+
+ #[serde(rename = "snowflake.private.key.passphrase")]
+ pub private_key_passphrase: Option<String>,
+
+ #[serde(rename = "snowflake.role")]
+ pub role: String,
+
+ #[serde(rename = "snowflake.jwt_token")]
+ pub jwt_token: String,
+
+ #[serde(rename = "snowflake.s3")]
+ pub s3: String,
+}
+
+#[serde_as]
+#[derive(Clone, Debug, Deserialize, WithOptions)]
+pub struct SnowflakeConfig {
+ #[serde(flatten)]
+ pub common: SnowflakeCommon,
+}
+
+impl SnowflakeConfig {
+ pub fn from_hashmap(properties: HashMap<String, String>) -> Result<Self> {
+ let config =
+ serde_json::from_value::<SnowflakeConfig>(serde_json::to_value(properties).unwrap())
+ .map_err(|e| SinkError::Config(anyhow!(e)))?;
+ Ok(config)
+ }
+}
+
+#[derive(Debug)]
+pub struct SnowflakeSink {
+ pub config: SnowflakeConfig,
+ schema: Schema,
+ pk_indices: Vec<usize>,
+ is_append_only: bool,
+}
+
+impl Sink for SnowflakeSink {
+ type Coordinator = DummySinkCommitCoordinator;
+ type LogSinker = LogSinkerOf<SnowflakeSinkWriter>;
+
+ const SINK_NAME: &'static str = SNOWFLAKE_SINK;
+
+ async fn new_log_sinker(&self, writer_param: SinkWriterParam) -> Result<Self::LogSinker> {
+ todo!()
+ }
+
+ async fn validate(&self) -> Result<()> {
+ todo!()
+ }
+}
+
+pub struct SnowflakeSinkWriter {
+ pub config: SnowflakeSink,
+ schema: Schema,
+ pk_indices: Vec<usize>,
+ is_append_only: bool,
+ client: Option<SnowflakeClient>,
+ row_encoder: JsonEncoder,
+}
+
+impl TryFrom<SinkParam> for SnowflakeSink {
+ type Error = SinkError;
+
+ fn try_from(param: SinkParam) -> std::result::Result<Self, Self::Error> {
+ let schema = param.schema();
+ let config = SnowflakeConfig::from_hashmap(param.properties)?;
+ Ok(SnowflakeSink {
+ config,
+ schema,
+ pk_indices: param.downstream_pk,
+ is_append_only: param.sink_type.is_append_only(),
+ })
+ }
+}
+
+#[async_trait]
+impl SinkWriter for SnowflakeSinkWriter {
+ async fn begin_epoch(&mut self, epoch: u64) -> Result<()> {
+ todo!()
+ }
+
+ async fn write_batch(&mut self, chunk: StreamChunk) -> Result<()> {
+ todo!()
+ }
+
+ async fn barrier(&mut self, is_checkpoint: bool) -> Result<Self::CommitMetadata> {
+ todo!()
+ }
+
+ async fn abort(&mut self) -> Result<()> {
+ Ok(())
+ }
+
+ async fn update_vnode_bitmap(&mut self, _vnode_bitmap: Arc<Bitmap>) -> Result<()> {
+ Ok(())
+ }
+}
+
+pub struct SnowflakeClient {}
+
+impl SnowflakeClient {}

src/connector/src/sink/snowflake_connector.rs

@@ -0,0 +1,51 @@
+use http::request::Builder;
+use hyper::body::{Body, Sender};
+use hyper::client::HttpConnector;
+use hyper::{body, Client, Request, StatusCode};
+use hyper_tls::HttpsConnector;
+
+use std::collections::HashMap;
+
+use super::{Result, SinkError};
+
+const SNOWFLAKE_HOST_ADDR: &str = "snowflakecomputing.com";
+const SNOWFLAKE_REQUEST_ID: &str = "RW_SNOWFLAKE_SINK";
+
+#[derive(Debug)]
+pub struct SnowflakeInserterBuilder {
+ url: String,
+ header: HashMap<String, String>,
+}
+
+impl SnowflakeInserterBuilder {
+ pub fn new(account: String, db: String, schema: String, pipe: String, header: HashMap<String, String>) -> Self {
+ // TODO: ensure if we need user to *explicitly* provide the request id
+ let url = format!("https://{}.{}/v1/data/pipes/{}.{}.{}/insertFiles?request_id={}",
+ account,
+ SNOWFLAKE_HOST_ADDR,
+ db,
+ schema.
+ pipe,
+ SNOWFLAKE_REQUEST_ID);
+
+ Self {
+ url,
+ header,
+ }
+ }
+
+ fn build_request_and_client() -> (Builder, Client<HttpsConnector<HttpConnector>>) {
+
+ }
+
+ pub async fn build(&self) -> Result<SnowflakeInserter> {
+
+ }
+}
+
+#[derive(Debug)]
+pub struct SnowflakeInserter {
+ sender: Option<Sender>,
+ join_handle: Option<JoinHandle<Result<Vec<u8>>,
+ buffer: BytesMut,
+}