risingwave_connector/sink/snowflake_redshift/
mod.rs

1// Copyright 2025 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::fmt::Write;
16use std::sync::Arc;
17use std::time::{SystemTime, UNIX_EPOCH};
18
19use bytes::BytesMut;
20use opendal::Operator;
21use risingwave_common::array::{ArrayImpl, DataChunk, Op, PrimitiveArray, StreamChunk, Utf8Array};
22use risingwave_common::catalog::{ColumnDesc, ColumnId, Schema};
23use risingwave_common::row::Row;
24use risingwave_common::types::DataType;
25use serde_json::{Map, Value};
26use thiserror_ext::AsReport;
27use uuid::Uuid;
28
29use crate::sink::encoder::{
30    JsonEncoder, JsonbHandlingMode, RowEncoder, TimeHandlingMode, TimestampHandlingMode,
31    TimestamptzHandlingMode,
32};
33use crate::sink::file_sink::opendal_sink::FileSink;
34use crate::sink::file_sink::s3::{S3Common, S3Sink};
35use crate::sink::remote::CoordinatedRemoteSinkWriter;
36use crate::sink::writer::SinkWriter;
37use crate::sink::{Result, SinkError, SinkParam, SinkWriterMetrics, SinkWriterParam};
38
39pub mod redshift;
40pub mod snowflake;
41
42pub const __ROW_ID: &str = "__row_id";
43pub const __OP: &str = "__op";
44
45pub struct AugmentedRow {
46    row_encoder: JsonEncoder,
47    current_epoch: u64,
48    current_row_count: usize,
49    is_append_only: bool,
50}
51
52impl AugmentedRow {
53    pub fn new(current_epoch: u64, is_append_only: bool, schema: Schema) -> Self {
54        let row_encoder = JsonEncoder::new(
55            schema,
56            None,
57            crate::sink::encoder::DateHandlingMode::String,
58            TimestampHandlingMode::String,
59            TimestamptzHandlingMode::UtcString,
60            TimeHandlingMode::String,
61            JsonbHandlingMode::String,
62        );
63        Self {
64            row_encoder,
65            current_epoch,
66            current_row_count: 0,
67            is_append_only,
68        }
69    }
70
71    pub fn reset_epoch(&mut self, current_epoch: u64) {
72        if self.is_append_only || current_epoch == self.current_epoch {
73            return;
74        }
75        self.current_epoch = current_epoch;
76        self.current_row_count = 0;
77    }
78
79    pub fn augmented_row(&mut self, row: impl Row, op: Op) -> Result<Map<String, Value>> {
80        let mut row = self.row_encoder.encode(row)?;
81        if self.is_append_only {
82            return Ok(row);
83        }
84        self.current_row_count += 1;
85        let row_id = format!("{}_{}", self.current_epoch, self.current_row_count);
86        row.insert(__ROW_ID.to_owned(), Value::String(row_id));
87        row.insert(
88            __OP.to_owned(),
89            Value::Number(serde_json::Number::from(op.to_i16())),
90        );
91        Ok(row)
92    }
93}
94
95pub struct AugmentedChunk {
96    current_epoch: u64,
97    current_row_count: usize,
98    is_append_only: bool,
99}
100
101impl AugmentedChunk {
102    pub fn new(current_epoch: u64, is_append_only: bool) -> Self {
103        Self {
104            current_epoch,
105            current_row_count: 0,
106            is_append_only,
107        }
108    }
109
110    pub fn reset_epoch(&mut self, current_epoch: u64) {
111        if self.is_append_only || current_epoch == self.current_epoch {
112            return;
113        }
114        self.current_epoch = current_epoch;
115        self.current_row_count = 0;
116    }
117
118    pub fn augmented_chunk(&mut self, chunk: StreamChunk) -> Result<StreamChunk> {
119        if self.is_append_only {
120            return Ok(chunk);
121        }
122        let (data_chunk, ops) = chunk.into_parts();
123        let chunk_row_count = data_chunk.capacity();
124        let (columns, visibility) = data_chunk.into_parts();
125
126        let op_column = ops.iter().map(|op| op.to_i16() as i32).collect::<Vec<_>>();
127        let row_column_strings: Vec<String> = (0..chunk_row_count)
128            .map(|i| format!("{}_{}", self.current_epoch, self.current_row_count + i))
129            .collect();
130
131        let row_column_refs: Vec<&str> = row_column_strings.iter().map(|s| s.as_str()).collect();
132        self.current_row_count += chunk_row_count;
133
134        let mut arrays: Vec<Arc<ArrayImpl>> = columns;
135        arrays.push(Arc::new(ArrayImpl::Utf8(Utf8Array::from_iter(
136            row_column_refs,
137        ))));
138        arrays.push(Arc::new(ArrayImpl::Int32(
139            PrimitiveArray::<i32>::from_iter(op_column),
140        )));
141
142        let chunk = DataChunk::new(arrays, visibility);
143        let ops = vec![Op::Insert; chunk_row_count];
144        let chunk = StreamChunk::from_parts(ops, chunk);
145        Ok(chunk)
146    }
147}
148
149pub struct SnowflakeRedshiftSinkS3Writer {
150    s3_config: S3Common,
151    s3_operator: Operator,
152    augmented_row: AugmentedRow,
153    opendal_writer_path: Option<(opendal::Writer, String)>,
154    target_table_name: String,
155}
156
157impl SnowflakeRedshiftSinkS3Writer {
158    pub fn new(
159        s3_config: S3Common,
160        schema: Schema,
161        is_append_only: bool,
162        target_table_name: String,
163    ) -> Result<Self> {
164        let s3_operator = FileSink::<S3Sink>::new_s3_sink(&s3_config)?;
165        Ok(Self {
166            s3_config,
167            s3_operator,
168            opendal_writer_path: None,
169            augmented_row: AugmentedRow::new(0, is_append_only, schema),
170            target_table_name,
171        })
172    }
173
174    pub fn begin_epoch(&mut self, epoch: u64) -> Result<()> {
175        self.augmented_row.reset_epoch(epoch);
176        Ok(())
177    }
178
179    pub async fn write_batch(&mut self, chunk: StreamChunk) -> Result<()> {
180        if self.opendal_writer_path.is_none() {
181            let opendal_writer_path = build_opendal_writer_path(
182                &self.s3_config,
183                &self.s3_operator,
184                None,
185                &self.target_table_name,
186            )
187            .await?;
188            self.opendal_writer_path = Some(opendal_writer_path);
189        }
190        let mut chunk_buf = BytesMut::new();
191        for (op, row) in chunk.rows() {
192            let encoded_row = self.augmented_row.augmented_row(row, op)?;
193            writeln!(chunk_buf, "{}", Value::Object(encoded_row)).unwrap(); // write to a `BytesMut` should never fail
194        }
195        self.opendal_writer_path
196            .as_mut()
197            .ok_or_else(|| SinkError::File("Sink writer is not created.".to_owned()))?
198            .0
199            .write(chunk_buf.freeze())
200            .await?;
201        Ok(())
202    }
203
204    pub async fn barrier(&mut self, is_checkpoint: bool) -> Result<Option<String>> {
205        if is_checkpoint && let Some((mut writer, path)) = self.opendal_writer_path.take() {
206            writer
207                .close()
208                .await
209                .map_err(|e| SinkError::File(e.to_report_string()))?;
210            Ok(Some(path))
211        } else {
212            Ok(None)
213        }
214    }
215}
216
217pub async fn build_opendal_writer_path(
218    s3_config: &S3Common,
219    operator: &Operator,
220    dir: Option<&str>,
221    target_table_name: &str,
222) -> Result<(opendal::Writer, String)> {
223    let mut base_path = s3_config.path.clone().unwrap_or("".to_owned());
224    if !base_path.ends_with('/') {
225        base_path.push('/');
226    }
227    base_path.push_str(&format!("{}/", target_table_name));
228    if let Some(dir) = dir {
229        base_path.push_str(&format!("{}/", dir));
230    }
231    let create_time = SystemTime::now()
232        .duration_since(UNIX_EPOCH)
233        .expect("Time went backwards");
234    let object_name = format!(
235        "{}{}_{}.{}",
236        base_path,
237        Uuid::new_v4(),
238        create_time.as_millis(),
239        "json",
240    );
241    let all_path = format!("s3://{}/{}", s3_config.bucket_name, object_name);
242    Ok((
243        operator.writer_with(&object_name).concurrent(8).await?,
244        all_path,
245    ))
246}
247
248/// Generic JDBC writer for both Redshift and Snowflake sinks
249pub struct SnowflakeRedshiftSinkJdbcWriter {
250    augmented_row: AugmentedChunk,
251    jdbc_sink_writer: CoordinatedRemoteSinkWriter,
252}
253
254impl SnowflakeRedshiftSinkJdbcWriter {
255    pub async fn new(
256        is_append_only: bool,
257        writer_param: SinkWriterParam,
258        mut param: SinkParam,
259        full_table_name: String,
260    ) -> Result<Self> {
261        let metrics = SinkWriterMetrics::new(&writer_param);
262        let column_descs = &mut param.columns;
263
264        // Build full table name based on connector type
265        if !is_append_only {
266            // Add CDC-specific columns for upsert mode
267            let max_column_id = column_descs
268                .iter()
269                .map(|column| column.column_id.get_id())
270                .max()
271                .unwrap_or(0);
272
273            (*column_descs).push(ColumnDesc::named(
274                __ROW_ID,
275                ColumnId::new(max_column_id + 1),
276                DataType::Varchar,
277            ));
278            (*column_descs).push(ColumnDesc::named(
279                __OP,
280                ColumnId::new(max_column_id + 2),
281                DataType::Int32,
282            ));
283        };
284
285        if let Some(schema_name) = param.properties.remove("schema") {
286            param
287                .properties
288                .insert("schema.name".to_owned(), schema_name);
289        }
290        if let Some(database_name) = param.properties.remove("database") {
291            param
292                .properties
293                .insert("database.name".to_owned(), database_name);
294        }
295        param
296            .properties
297            .insert("table.name".to_owned(), full_table_name.clone());
298        param
299            .properties
300            .insert("type".to_owned(), "append-only".to_owned());
301
302        let jdbc_sink_writer =
303            CoordinatedRemoteSinkWriter::new(param.clone(), metrics.clone()).await?;
304
305        Ok(Self {
306            augmented_row: AugmentedChunk::new(0, is_append_only),
307            jdbc_sink_writer,
308        })
309    }
310
311    pub async fn begin_epoch(&mut self, epoch: u64) -> Result<()> {
312        self.augmented_row.reset_epoch(epoch);
313        self.jdbc_sink_writer.begin_epoch(epoch).await?;
314        Ok(())
315    }
316
317    pub async fn write_batch(&mut self, chunk: StreamChunk) -> Result<()> {
318        let chunk = self.augmented_row.augmented_chunk(chunk)?;
319        self.jdbc_sink_writer.write_batch(chunk).await?;
320        Ok(())
321    }
322
323    pub async fn barrier(&mut self, is_checkpoint: bool) -> Result<()> {
324        self.jdbc_sink_writer.barrier(is_checkpoint).await?;
325        Ok(())
326    }
327
328    pub async fn abort(&mut self) -> Result<()> {
329        // TODO: abort should clean up all the data written in this epoch
330        self.jdbc_sink_writer.abort().await?;
331        Ok(())
332    }
333}
risingwave_connector/sink/snowflake_redshift/mod.rs

risingwave_connector/sink/snowflake_redshift/
mod.rs
