1use core::num::NonZero;
16use std::collections::BTreeMap;
17use std::fmt::Write;
18use std::time::Duration;
19
20use anyhow::anyhow;
21use bytes::BytesMut;
22use phf::{Set, phf_set};
23use risingwave_common::array::StreamChunk;
24use risingwave_common::catalog::Schema;
25use risingwave_common::types::DataType;
26use risingwave_pb::connector_service::sink_metadata::SerializedMetadata;
27use risingwave_pb::connector_service::{SinkMetadata, sink_metadata};
28use risingwave_pb::stream_plan::PbSinkSchemaChange;
29use serde::Deserialize;
30use serde_json::json;
31use serde_with::{DisplayFromStr, serde_as};
32use thiserror_ext::AsReport;
33use tokio::sync::mpsc::{UnboundedSender, unbounded_channel};
34use tokio::time::{MissedTickBehavior, interval};
35use tonic::async_trait;
36use tracing::warn;
37use with_options::WithOptions;
38
39use crate::connector_common::IcebergSinkCompactionUpdate;
40use crate::enforce_secret::EnforceSecret;
41use crate::sink::catalog::SinkId;
42use crate::sink::coordinate::CoordinatedLogSinker;
43use crate::sink::file_sink::opendal_sink::FileSink;
44use crate::sink::file_sink::s3::{S3Common, S3Sink};
45use crate::sink::jdbc_jni_client::{self, JdbcJniClient};
46use crate::sink::snowflake_redshift::{
47 __OP, __ROW_ID, SnowflakeRedshiftSinkJdbcWriter, SnowflakeRedshiftSinkS3Writer,
48 build_opendal_writer_path,
49};
50use crate::sink::writer::SinkWriter;
51use crate::sink::{
52 Result, SinglePhaseCommitCoordinator, Sink, SinkCommitCoordinator, SinkError, SinkParam,
53};
54
55pub const REDSHIFT_SINK: &str = "redshift";
56
57pub fn build_full_table_name(schema_name: Option<&str>, table_name: &str) -> String {
58 if let Some(schema_name) = schema_name {
59 format!(r#""{}"."{}""#, schema_name, table_name)
60 } else {
61 format!(r#""{}""#, table_name)
62 }
63}
64
65fn build_alter_add_column_sql(
66 schema_name: Option<&str>,
67 table_name: &str,
68 columns: &Vec<(String, String)>,
69) -> String {
70 let full_table_name = build_full_table_name(schema_name, table_name);
71 jdbc_jni_client::build_alter_add_column_sql(&full_table_name, columns, false)
73}
74
75#[serde_as]
76#[derive(Debug, Clone, Deserialize, WithOptions)]
77pub struct RedShiftConfig {
78 #[serde(rename = "jdbc.url")]
79 pub jdbc_url: String,
80
81 #[serde(rename = "user")]
82 pub username: Option<String>,
83
84 #[serde(rename = "password")]
85 pub password: Option<String>,
86
87 #[serde(rename = "schema")]
88 pub schema: Option<String>,
89
90 #[serde(rename = "intermediate.schema.name")]
91 pub intermediate_schema: Option<String>,
92
93 #[serde(rename = "table.name")]
94 pub table: String,
95
96 #[serde(rename = "intermediate.table.name")]
97 pub cdc_table: Option<String>,
98
99 #[serde(default)]
100 #[serde(rename = "create_table_if_not_exists")]
101 #[serde_as(as = "DisplayFromStr")]
102 pub create_table_if_not_exists: bool,
103
104 #[serde(default = "default_target_interval_schedule")]
105 #[serde(rename = "write.target.interval.seconds")]
106 #[serde_as(as = "DisplayFromStr")]
107 pub writer_target_interval_seconds: u64,
108
109 #[serde(default = "default_intermediate_interval_schedule")]
110 #[serde(rename = "write.intermediate.interval.seconds")]
111 #[serde_as(as = "DisplayFromStr")]
112 pub write_intermediate_interval_seconds: u64,
113
114 #[serde(default = "default_batch_insert_rows")]
115 #[serde(rename = "batch.insert.rows")]
116 #[serde_as(as = "DisplayFromStr")]
117 pub batch_insert_rows: u32,
118
119 #[serde(default = "default_with_s3")]
120 #[serde(rename = "with_s3")]
121 #[serde_as(as = "DisplayFromStr")]
122 pub with_s3: bool,
123
124 #[serde(flatten)]
125 pub s3_inner: Option<S3Common>,
126
127 #[serde(flatten)]
128 pub unknown_fields: std::collections::HashMap<String, String>,
129}
130
131crate::impl_sink_unknown_fields!(RedShiftConfig);
132
133fn default_target_interval_schedule() -> u64 {
134 3600 }
136
137fn default_intermediate_interval_schedule() -> u64 {
138 1800 }
140
141fn default_batch_insert_rows() -> u32 {
142 4096 }
144
145fn default_with_s3() -> bool {
146 true
147}
148
149impl RedShiftConfig {
150 pub fn from_btreemap(properties: BTreeMap<String, String>) -> Result<Self> {
151 serde_json::from_value::<RedShiftConfig>(serde_json::to_value(properties).unwrap())
152 .map_err(|e| SinkError::Config(anyhow!(e)))
153 }
154
155 pub fn build_client(&self) -> Result<JdbcJniClient> {
156 let mut jdbc_url = self.jdbc_url.clone();
157 if let Some(username) = &self.username {
158 jdbc_url = format!("{}?user={}", jdbc_url, username);
159 }
160 if let Some(password) = &self.password {
161 jdbc_url = format!("{}&password={}", jdbc_url, password);
162 }
163 JdbcJniClient::new(jdbc_url)
164 }
165}
166
167#[derive(Debug)]
168pub struct RedshiftSink {
169 config: RedShiftConfig,
170 param: SinkParam,
171 is_append_only: bool,
172 schema: Schema,
173 pk_indices: Vec<usize>,
174}
175impl EnforceSecret for RedshiftSink {
176 const ENFORCE_SECRET_PROPERTIES: Set<&'static str> = phf_set! {
177 "user",
178 "password",
179 "jdbc.url"
180 };
181}
182
183impl TryFrom<SinkParam> for RedshiftSink {
184 type Error = SinkError;
185
186 fn try_from(param: SinkParam) -> std::result::Result<Self, Self::Error> {
187 let config = RedShiftConfig::from_btreemap(param.properties.clone())?;
188 let is_append_only = param.sink_type.is_append_only();
189 let schema = param.schema();
190 let pk_indices = param.downstream_pk_or_empty();
191 Ok(Self {
192 config,
193 param,
194 is_append_only,
195 schema,
196 pk_indices,
197 })
198 }
199}
200
201impl Sink for RedshiftSink {
202 type LogSinker = CoordinatedLogSinker<RedShiftSinkWriter>;
203
204 const SINK_NAME: &'static str = REDSHIFT_SINK;
205
206 crate::impl_validate_sink_unknown_fields!();
207
208 async fn validate(&self) -> Result<()> {
209 if self.config.create_table_if_not_exists {
210 let client = self.config.build_client()?;
211 let schema = self.param.schema();
212 let build_table_sql = build_create_table_sql(
213 self.config.schema.as_deref(),
214 &self.config.table,
215 &schema,
216 false,
217 )?;
218 client.execute_sql_sync(vec![build_table_sql]).await?;
219 if !self.is_append_only {
220 let cdc_table = self.config.cdc_table.as_ref().ok_or_else(|| {
221 SinkError::Config(anyhow!(
222 "intermediate.table.name is required for append-only sink"
223 ))
224 })?;
225 let cdc_schema_for_create = self
226 .config
227 .intermediate_schema
228 .as_deref()
229 .or(self.config.schema.as_deref());
230 let build_cdc_table_sql =
231 build_create_table_sql(cdc_schema_for_create, cdc_table, &schema, true)?;
232 client.execute_sql_sync(vec![build_cdc_table_sql]).await?;
233 }
234 }
235 Ok(())
236 }
237
238 fn support_schema_change() -> bool {
239 true
240 }
241
242 async fn new_log_sinker(
243 &self,
244 writer_param: crate::sink::SinkWriterParam,
245 ) -> Result<Self::LogSinker> {
246 let writer = RedShiftSinkWriter::new(
247 self.config.clone(),
248 self.is_append_only,
249 writer_param.clone(),
250 self.param.clone(),
251 )
252 .await?;
253 CoordinatedLogSinker::new(
254 &writer_param,
255 self.param.clone(),
256 writer,
257 NonZero::new(1).unwrap(),
258 )
259 .await
260 }
261
262 fn is_coordinated_sink(&self) -> bool {
263 true
264 }
265
266 async fn new_coordinator(
267 &self,
268 _iceberg_compact_stat_sender: Option<UnboundedSender<IcebergSinkCompactionUpdate>>,
269 ) -> Result<SinkCommitCoordinator> {
270 let pk_column_names: Vec<_> = self
271 .schema
272 .fields
273 .iter()
274 .enumerate()
275 .filter(|(index, _)| self.pk_indices.contains(index))
276 .map(|(_, field)| field.name.clone())
277 .collect();
278 if pk_column_names.is_empty() && !self.is_append_only {
279 return Err(SinkError::Config(anyhow!(
280 "Primary key columns not found. Please set the `primary_key` column in the sink properties, or ensure that the sink contains the primary key columns from the upstream."
281 )));
282 }
283 let all_column_names = self
284 .schema
285 .fields
286 .iter()
287 .map(|field| field.name.clone())
288 .collect();
289 let coordinator = RedshiftSinkCommitter::new(
290 self.config.clone(),
291 self.is_append_only,
292 &pk_column_names,
293 &all_column_names,
294 self.param.sink_id,
295 )?;
296 Ok(SinkCommitCoordinator::SinglePhase(Box::new(coordinator)))
297 }
298}
299
300pub enum RedShiftSinkWriter {
301 S3(SnowflakeRedshiftSinkS3Writer),
302 Jdbc(SnowflakeRedshiftSinkJdbcWriter),
303}
304
305impl RedShiftSinkWriter {
306 pub async fn new(
307 config: RedShiftConfig,
308 is_append_only: bool,
309 writer_param: super::SinkWriterParam,
310 mut param: SinkParam,
311 ) -> Result<Self> {
312 let schema = param.schema();
313 if config.with_s3 {
314 let s3_writer = SnowflakeRedshiftSinkS3Writer::new(
315 config.s3_inner.ok_or_else(|| {
316 SinkError::Config(anyhow!("S3 configuration is required for S3 sink"))
317 })?,
318 schema,
319 is_append_only,
320 config.table,
321 )?;
322 Ok(Self::S3(s3_writer))
323 } else {
324 let (writer_schema, writer_table) = if is_append_only {
325 (config.schema.clone(), config.table.clone())
326 } else {
327 (
328 config.intermediate_schema.clone().or(config.schema.clone()),
329 config.cdc_table.clone().ok_or_else(|| {
330 SinkError::Config(anyhow!(
331 "intermediate.table.name is required for non-append-only sink"
332 ))
333 })?,
334 )
335 };
336 param.properties.remove("schema");
337 param.properties.remove("schema.name");
338 if let Some(writer_schema) = writer_schema {
339 param.properties.insert("schema".to_owned(), writer_schema);
340 }
341 let jdbc_writer = SnowflakeRedshiftSinkJdbcWriter::new(
342 is_append_only,
343 writer_param,
344 param,
345 writer_table,
346 )
347 .await?;
348 Ok(Self::Jdbc(jdbc_writer))
349 }
350 }
351}
352
353#[async_trait]
354impl SinkWriter for RedShiftSinkWriter {
355 type CommitMetadata = Option<SinkMetadata>;
356
357 async fn begin_epoch(&mut self, epoch: u64) -> Result<()> {
358 match self {
359 Self::S3(writer) => writer.begin_epoch(epoch),
360 Self::Jdbc(writer) => writer.begin_epoch(epoch).await,
361 }
362 }
363
364 async fn write_batch(&mut self, chunk: StreamChunk) -> Result<()> {
365 match self {
366 Self::S3(writer) => writer.write_batch(chunk).await,
367 Self::Jdbc(writer) => writer.write_batch(chunk).await,
368 }
369 }
370
371 async fn barrier(&mut self, is_checkpoint: bool) -> Result<Option<SinkMetadata>> {
372 let metadata = match self {
373 Self::S3(writer) => {
374 if let Some(path) = writer.barrier(is_checkpoint).await? {
375 path.into_bytes()
376 } else {
377 vec![]
378 }
379 }
380 Self::Jdbc(writer) => {
381 writer.barrier(is_checkpoint).await?;
382 vec![]
383 }
384 };
385 Ok(Some(SinkMetadata {
386 metadata: Some(sink_metadata::Metadata::Serialized(SerializedMetadata {
387 metadata,
388 })),
389 }))
390 }
391
392 async fn abort(&mut self) -> Result<()> {
393 if let Self::Jdbc(writer) = self {
394 writer.abort().await
395 } else {
396 Ok(())
397 }
398 }
399}
400
401pub struct RedshiftSinkCommitter {
402 config: RedShiftConfig,
403 client: JdbcJniClient,
404 sink_id: SinkId,
405 pk_column_names: Vec<String>,
406 all_column_names: Vec<String>,
407 writer_target_interval_seconds: u64,
408 write_intermediate_interval_seconds: u64,
409 is_append_only: bool,
410 periodic_task_handle: Option<tokio::task::JoinHandle<()>>,
411 shutdown_sender: Option<tokio::sync::mpsc::UnboundedSender<()>>,
412}
413
414impl RedshiftSinkCommitter {
415 pub fn new(
416 config: RedShiftConfig,
417 is_append_only: bool,
418 pk_column_names: &Vec<String>,
419 all_column_names: &Vec<String>,
420 sink_id: SinkId,
421 ) -> Result<Self> {
422 let client = config.build_client()?;
423 let writer_target_interval_seconds = config.writer_target_interval_seconds;
424 let write_intermediate_interval_seconds = config.write_intermediate_interval_seconds;
425
426 let (periodic_task_handle, shutdown_sender) = match (is_append_only, config.with_s3) {
427 (true, true) | (false, _) => {
428 let task_client = config.build_client()?;
429 let config = config.clone();
430 let (shutdown_sender, shutdown_receiver) = unbounded_channel();
431 let target_schema_name = config.schema.as_deref();
432 let effective_cdc_schema =
433 config.intermediate_schema.as_deref().or(target_schema_name);
434 let merge_into_sql = if !is_append_only {
435 Some(build_create_merge_into_task_sql(
436 effective_cdc_schema,
437 target_schema_name,
438 config.cdc_table.as_ref().ok_or_else(|| {
439 SinkError::Config(anyhow!(
440 "intermediate.table.name is required for non-append-only sink"
441 ))
442 })?,
443 &config.table,
444 pk_column_names,
445 all_column_names,
446 ))
447 } else {
448 None
449 };
450 let periodic_task_handle = tokio::spawn(async move {
451 Self::run_periodic_query_task(
452 task_client,
453 merge_into_sql,
454 config.with_s3,
455 writer_target_interval_seconds,
456 write_intermediate_interval_seconds,
457 sink_id,
458 config,
459 is_append_only,
460 shutdown_receiver,
461 )
462 .await;
463 });
464 (Some(periodic_task_handle), Some(shutdown_sender))
465 }
466 _ => (None, None),
467 };
468
469 Ok(Self {
470 client,
471 config,
472 sink_id,
473 pk_column_names: pk_column_names.clone(),
474 all_column_names: all_column_names.clone(),
475 is_append_only,
476 writer_target_interval_seconds,
477 write_intermediate_interval_seconds,
478 periodic_task_handle,
479 shutdown_sender,
480 })
481 }
482
483 async fn flush_manifest_to_redshift(
484 client: &JdbcJniClient,
485 config: &RedShiftConfig,
486 s3_inner: &S3Common,
487 is_append_only: bool,
488 ) -> Result<()> {
489 let s3_operator = FileSink::<S3Sink>::new_s3_sink(s3_inner)?;
490 let mut manifest_path = s3_inner.path.clone().unwrap_or("".to_owned());
491 if !manifest_path.ends_with('/') {
492 manifest_path.push('/');
493 }
494 manifest_path.push_str(&format!("{}/", config.table));
495 manifest_path.push_str("manifest/");
496 let manifests = s3_operator
497 .list(&manifest_path)
498 .await?
499 .into_iter()
500 .map(|e| e.path().to_owned())
501 .collect::<Vec<_>>();
502 for manifest in &manifests {
503 Self::copy_into_from_s3_to_redshift(client, config, s3_inner, is_append_only, manifest)
504 .await?;
505 }
506 s3_operator.delete_iter(manifests).await?;
507 Ok(())
508 }
509
510 async fn write_manifest_to_s3(
511 s3_inner: &S3Common,
512 paths: Vec<String>,
513 table: &str,
514 ) -> Result<String> {
515 let manifest_entries: Vec<_> = paths
516 .into_iter()
517 .map(|path| json!({ "url": path, "mandatory": true }))
518 .collect();
519 let s3_operator = FileSink::<S3Sink>::new_s3_sink(s3_inner)?;
520 let (mut writer, manifest_path) =
521 build_opendal_writer_path(s3_inner, &s3_operator, Some("manifest"), table).await?;
522 let manifest_json = json!({ "entries": manifest_entries });
523 let mut chunk_buf = BytesMut::new();
524 writeln!(chunk_buf, "{}", manifest_json).unwrap();
525 writer.write(chunk_buf.freeze()).await?;
526 writer.close().await.map_err(|e| {
527 SinkError::Redshift(anyhow!(
528 "Failed to close manifest writer: {}",
529 e.to_report_string()
530 ))
531 })?;
532 Ok(manifest_path)
533 }
534
535 pub async fn copy_into_from_s3_to_redshift(
536 client: &JdbcJniClient,
537 config: &RedShiftConfig,
538 s3_inner: &S3Common,
539 is_append_only: bool,
540 manifest: &str,
541 ) -> Result<()> {
542 let all_path = format!("s3://{}/{}", s3_inner.bucket_name, manifest);
543
544 let (table, schema_name) = if is_append_only {
545 (&config.table, config.schema.as_deref())
546 } else {
547 (
548 config.cdc_table.as_ref().ok_or_else(|| {
549 SinkError::Config(anyhow!(
550 "intermediate.table.name is required for non-append-only sink"
551 ))
552 })?,
553 config
554 .intermediate_schema
555 .as_deref()
556 .or(config.schema.as_deref()),
557 )
558 };
559 let copy_into_sql = build_copy_into_sql(
560 schema_name,
561 table,
562 &all_path,
563 &s3_inner.access,
564 &s3_inner.secret,
565 &s3_inner.assume_role,
566 )?;
567 client.execute_sql_sync(vec![copy_into_sql]).await?;
568 Ok(())
569 }
570
571 async fn run_periodic_query_task(
572 client: JdbcJniClient,
573 merge_into_sql: Option<Vec<String>>,
574 need_copy_into: bool,
575 writer_target_interval_seconds: u64,
576 write_intermediate_interval_seconds: u64,
577 sink_id: SinkId,
578 config: RedShiftConfig,
579 is_append_only: bool,
580 mut shutdown_receiver: tokio::sync::mpsc::UnboundedReceiver<()>,
581 ) {
582 let mut copy_timer = interval(Duration::from_secs(write_intermediate_interval_seconds));
583 copy_timer.set_missed_tick_behavior(MissedTickBehavior::Skip);
584 let mut merge_timer = interval(Duration::from_secs(writer_target_interval_seconds));
585 merge_timer.set_missed_tick_behavior(MissedTickBehavior::Skip);
586
587 loop {
588 tokio::select! {
589 _ = shutdown_receiver.recv() => break,
590 _ = merge_timer.tick(), if merge_into_sql.is_some() => {
591 if let Some(sql) = &merge_into_sql && let Err(e) = client.execute_sql_sync(sql.clone()).await {
592 tracing::warn!("Failed to execute periodic query for table {}: {}", config.table, e.as_report());
593 }
594 },
595 _ = copy_timer.tick(), if need_copy_into => {
596 if let Err(e) = async {
597 let s3_inner = config.s3_inner.as_ref().ok_or_else(|| {
598 SinkError::Config(anyhow!("S3 configuration is required for redshift s3 sink"))
599 })?;
600 Self::flush_manifest_to_redshift(&client, &config,s3_inner, is_append_only).await?;
601 Ok::<(),SinkError>(())
602 }.await {
603 tracing::error!("Failed to execute copy into task for sink id {}: {}", sink_id, e.as_report());
604 }
605 }
606 }
607 }
608 tracing::info!("Periodic query task stopped for sink id {}", sink_id);
609 }
610}
611
612impl Drop for RedshiftSinkCommitter {
613 fn drop(&mut self) {
614 if let Some(shutdown_sender) = &self.shutdown_sender
616 && let Err(e) = shutdown_sender.send(())
617 {
618 tracing::warn!(
619 "Failed to send shutdown signal to periodic task: {}",
620 e.as_report()
621 );
622 }
623 tracing::info!("RedshiftSinkCommitter dropped, periodic task stopped");
624 }
625}
626
627#[async_trait]
628impl SinglePhaseCommitCoordinator for RedshiftSinkCommitter {
629 async fn init(&mut self) -> Result<()> {
630 if self.config.with_s3 {
631 Self::flush_manifest_to_redshift(
632 &self.client,
633 &self.config,
634 self.config.s3_inner.as_ref().ok_or_else(|| {
635 SinkError::Config(anyhow!("S3 configuration is required for redshift s3 sink"))
636 })?,
637 self.is_append_only,
638 )
639 .await?;
640 }
641 Ok(())
642 }
643
644 async fn commit_data(&mut self, epoch: u64, metadata: Vec<SinkMetadata>) -> Result<()> {
645 if let Some(handle) = &self.periodic_task_handle {
646 let is_finished = handle.is_finished();
647 if is_finished {
648 let handle = self.periodic_task_handle.take().unwrap();
649 handle.await.map_err(|e| {
650 SinkError::Redshift(anyhow!(
651 "Periodic task for sink id {} panicked: {}",
652 self.sink_id,
653 e.to_report_string()
654 ))
655 })?;
656 }
657 };
658 let paths = metadata
659 .into_iter()
660 .filter(|m| {
661 if let Some(sink_metadata::Metadata::Serialized(SerializedMetadata { metadata })) =
662 &m.metadata
663 {
664 !metadata.is_empty()
665 } else {
666 false
667 }
668 })
669 .map(|metadata| {
670 let path = if let Some(sink_metadata::Metadata::Serialized(SerializedMetadata {
671 metadata,
672 })) = metadata.metadata
673 {
674 String::from_utf8(metadata).map_err(|e| SinkError::Config(anyhow!(e)))
675 } else {
676 Err(SinkError::Config(anyhow!("Invalid metadata format")))
677 }?;
678 Ok(path)
679 })
680 .collect::<Result<Vec<_>>>()?;
681
682 if !paths.is_empty() {
684 let s3_inner = self.config.s3_inner.as_ref().ok_or_else(|| {
685 SinkError::Config(anyhow!("S3 configuration is required for S3 sink"))
686 })?;
687 {
688 Self::write_manifest_to_s3(s3_inner, paths, &self.config.table).await?;
689 }
690 tracing::info!(
691 "Manifest file written to S3 for sink id {} at epoch {}",
692 self.sink_id,
693 epoch
694 );
695 }
696 Ok(())
697 }
698
699 async fn commit_schema_change(
700 &mut self,
701 _epoch: u64,
702 schema_change: PbSinkSchemaChange,
703 ) -> Result<()> {
704 use risingwave_pb::stream_plan::sink_schema_change::PbOp as SinkSchemaChangeOp;
705 let schema_change_op = schema_change
706 .op
707 .ok_or_else(|| SinkError::Coordinator(anyhow!("Invalid schema change operation")))?;
708 let SinkSchemaChangeOp::AddColumns(add_columns) = schema_change_op else {
709 return Err(SinkError::Coordinator(anyhow!(
710 "Only AddColumns schema change is supported for Redshift sink"
711 )));
712 };
713 if let Some(shutdown_sender) = &self.shutdown_sender {
714 shutdown_sender
716 .send(())
717 .map_err(|e| SinkError::Config(anyhow!(e)))?;
718 }
719 let sql = build_alter_add_column_sql(
720 self.config.schema.as_deref(),
721 &self.config.table,
722 &add_columns
723 .fields
724 .iter()
725 .map(|f| {
726 let dt = DataType::from(f.data_type.as_ref().unwrap());
727 Ok((f.name.clone(), convert_redshift_data_type(&dt)?))
728 })
729 .collect::<Result<Vec<_>>>()?,
730 );
731 let check_column_exists = |e: anyhow::Error| {
732 let err_str = e.to_report_string();
733 if regex::Regex::new(".+ of relation .+ already exists")
734 .unwrap()
735 .find(&err_str)
736 .is_none()
737 {
738 return Err(e);
739 }
740 warn!("redshift sink columns already exists. skipped");
741 Ok(())
742 };
743 self.client
744 .execute_sql_sync(vec![sql.clone()])
745 .await
746 .or_else(check_column_exists)?;
747 let merge_into_sql = if !self.is_append_only {
748 let cdc_table_name = self.config.cdc_table.as_ref().ok_or_else(|| {
749 SinkError::Config(anyhow!(
750 "intermediate.table.name is required for non-append-only sink"
751 ))
752 })?;
753 let sql = build_alter_add_column_sql(
754 self.config
755 .intermediate_schema
756 .as_deref()
757 .or(self.config.schema.as_deref()),
758 cdc_table_name,
759 &add_columns
760 .fields
761 .iter()
762 .map(|f| {
763 let dt = DataType::from(f.data_type.as_ref().unwrap());
764 Ok((f.name.clone(), convert_redshift_data_type(&dt)?))
765 })
766 .collect::<Result<Vec<_>>>()?,
767 );
768 self.client
769 .execute_sql_sync(vec![sql.clone()])
770 .await
771 .or_else(check_column_exists)?;
772 self.all_column_names
773 .extend(add_columns.fields.iter().map(|f| f.name.clone()));
774 let target_schema_name = self.config.schema.as_deref();
775 let effective_cdc_schema = self
776 .config
777 .intermediate_schema
778 .as_deref()
779 .or(target_schema_name);
780 let merge_into_sql = build_create_merge_into_task_sql(
781 effective_cdc_schema,
782 target_schema_name,
783 self.config.cdc_table.as_ref().ok_or_else(|| {
784 SinkError::Config(anyhow!(
785 "intermediate.table.name is required for non-append-only sink"
786 ))
787 })?,
788 &self.config.table,
789 &self.pk_column_names,
790 &self.all_column_names,
791 );
792 Some(merge_into_sql)
793 } else {
794 None
795 };
796
797 if let Some(shutdown_sender) = self.shutdown_sender.take() {
798 let _ = shutdown_sender.send(());
799 }
800 if let Some(periodic_task_handle) = self.periodic_task_handle.take() {
801 let _ = periodic_task_handle.await;
802 }
803
804 let (shutdown_sender, shutdown_receiver) = unbounded_channel();
805 let client = self.client.clone();
806
807 let writer_target_interval_seconds = self.writer_target_interval_seconds;
808 let write_intermediate_interval_seconds = self.write_intermediate_interval_seconds;
809 let config = self.config.clone();
810 let sink_id = self.sink_id;
811 let is_append_only = self.is_append_only;
812 let periodic_task_handle = tokio::spawn(async move {
813 Self::run_periodic_query_task(
814 client,
815 merge_into_sql,
816 config.with_s3,
817 writer_target_interval_seconds,
818 write_intermediate_interval_seconds,
819 sink_id,
820 config,
821 is_append_only,
822 shutdown_receiver,
823 )
824 .await;
825 });
826 self.shutdown_sender = Some(shutdown_sender);
827 self.periodic_task_handle = Some(periodic_task_handle);
828
829 Ok(())
830 }
831}
832
833pub fn build_create_table_sql(
834 schema_name: Option<&str>,
835 table_name: &str,
836 schema: &Schema,
837 need_op_and_row_id: bool,
838) -> Result<String> {
839 let mut columns: Vec<String> = schema
840 .fields
841 .iter()
842 .map(|field| {
843 let data_type = convert_redshift_data_type(&field.data_type)?;
844 Ok(format!("{} {}", field.name, data_type))
845 })
846 .collect::<Result<Vec<String>>>()?;
847 if need_op_and_row_id {
848 columns.push(format!("{} VARCHAR(MAX)", __ROW_ID));
849 columns.push(format!("{} INT", __OP));
850 }
851 let columns_str = columns.join(", ");
852 let full_table_name = build_full_table_name(schema_name, table_name);
853 Ok(format!(
854 "CREATE TABLE IF NOT EXISTS {} ({})",
855 full_table_name, columns_str
856 ))
857}
858
859fn convert_redshift_data_type(data_type: &DataType) -> Result<String> {
860 let data_type = match data_type {
861 DataType::Int16 => "SMALLINT".to_owned(),
862 DataType::Int32 => "INTEGER".to_owned(),
863 DataType::Int64 => "BIGINT".to_owned(),
864 DataType::Float32 => "REAL".to_owned(),
865 DataType::Float64 => "FLOAT".to_owned(),
866 DataType::Boolean => "BOOLEAN".to_owned(),
867 DataType::Varchar => "VARCHAR(MAX)".to_owned(),
868 DataType::Date => "DATE".to_owned(),
869 DataType::Timestamp => "TIMESTAMP".to_owned(),
870 DataType::Timestamptz => "TIMESTAMPTZ".to_owned(),
871 DataType::Jsonb => "VARCHAR(MAX)".to_owned(),
872 DataType::Decimal => "DECIMAL(38, 10)".to_owned(),
876 DataType::Time => "TIME".to_owned(),
877 _ => {
878 return Err(SinkError::Config(anyhow!(
879 "Dont support auto create table for datatype: {}",
880 data_type
881 )));
882 }
883 };
884 Ok(data_type)
885}
886
887fn build_create_merge_into_task_sql(
888 cdc_schema_name: Option<&str>,
889 target_schema_name: Option<&str>,
890 cdc_table_name: &str,
891 target_table_name: &str,
892 pk_column_names: &Vec<String>,
893 all_column_names: &Vec<String>,
894) -> Vec<String> {
895 let cdc_table_name = build_full_table_name(cdc_schema_name, cdc_table_name);
896 let target_table_name = build_full_table_name(target_schema_name, target_table_name);
897 let pk_names_str = pk_column_names.join(", ");
898 let pk_names_eq_str = pk_column_names
899 .iter()
900 .map(|name| format!("{target_table_name}.{name} = source.{name}", name = name))
901 .collect::<Vec<String>>()
902 .join(" AND ");
903 let all_column_names_set_str = all_column_names
904 .iter()
905 .map(|name| format!("{name} = source.{name}", name = name))
906 .collect::<Vec<String>>()
907 .join(", ");
908 let all_column_names_str = all_column_names.join(", ");
909 let all_column_names_insert_str = all_column_names
910 .iter()
911 .map(|name| format!("source.{name}", name = name))
912 .collect::<Vec<String>>()
913 .join(", ");
914
915 vec![
916 format!(
917 r#"
918 CREATE TEMP TABLE max_id_table AS
919 SELECT COALESCE(MAX({redshift_sink_row_id}), '0') AS max_row_id
920 FROM {cdc_table_name};
921 "#,
922 redshift_sink_row_id = __ROW_ID,
923 cdc_table_name = cdc_table_name,
924 ),
925 format!(
926 r#"
927 DELETE FROM {target_table_name}
928 USING (
929 SELECT *
930 FROM (
931 SELECT *, ROW_NUMBER() OVER (
932 PARTITION BY {pk_names_str}
933 ORDER BY {redshift_sink_row_id} DESC
934 ) AS dedupe_id
935 FROM {cdc_table_name}, max_id_table
936 WHERE {cdc_table_name}.{redshift_sink_row_id} <= max_id_table.max_row_id
937 ) AS subquery
938 WHERE dedupe_id = 1 AND {redshift_sink_op} IN (2, 4)
939 ) AS source
940 WHERE {pk_names_eq_str};
941 "#,
942 target_table_name = target_table_name,
943 pk_names_str = pk_names_str,
944 redshift_sink_row_id = __ROW_ID,
945 cdc_table_name = cdc_table_name,
946 redshift_sink_op = __OP,
947 pk_names_eq_str = pk_names_eq_str,
948 ),
949 format!(
950 r#"
951 MERGE INTO {target_table_name}
952 USING (
953 SELECT *
954 FROM (
955 SELECT *, ROW_NUMBER() OVER (
956 PARTITION BY {pk_names_str}
957 ORDER BY {redshift_sink_row_id} DESC
958 ) AS dedupe_id
959 FROM {cdc_table_name}, max_id_table
960 WHERE {cdc_table_name}.{redshift_sink_row_id} <= max_id_table.max_row_id
961 ) AS subquery
962 WHERE dedupe_id = 1 AND {redshift_sink_op} IN (1, 3)
963 ) AS source
964 ON {pk_names_eq_str}
965 WHEN MATCHED THEN
966 UPDATE SET {all_column_names_set_str}
967 WHEN NOT MATCHED THEN
968 INSERT ({all_column_names_str}) VALUES ({all_column_names_insert_str});
969 "#,
970 target_table_name = target_table_name,
971 pk_names_str = pk_names_str,
972 redshift_sink_row_id = __ROW_ID,
973 cdc_table_name = cdc_table_name,
974 redshift_sink_op = __OP,
975 pk_names_eq_str = pk_names_eq_str,
976 all_column_names_set_str = all_column_names_set_str,
977 all_column_names_str = all_column_names_str,
978 all_column_names_insert_str = all_column_names_insert_str,
979 ),
980 format!(
981 r#"
982 DELETE FROM {cdc_table_name}
983 USING max_id_table
984 WHERE {cdc_table_name}.{redshift_sink_row_id} <= max_id_table.max_row_id;
985 "#,
986 cdc_table_name = cdc_table_name,
987 redshift_sink_row_id = __ROW_ID,
988 ),
989 "DROP TABLE IF EXISTS max_id_table;".to_owned(),
990 ]
991}
992
993fn build_copy_into_sql(
994 schema_name: Option<&str>,
995 table_name: &str,
996 manifest_dir: &str,
997 access_key: &Option<String>,
998 secret_key: &Option<String>,
999 assume_role: &Option<String>,
1000) -> Result<String> {
1001 let table_name = build_full_table_name(schema_name, table_name);
1002 let credentials = if let Some(assume_role) = assume_role {
1003 &format!("aws_iam_role={}", assume_role)
1004 } else if let (Some(access_key), Some(secret_key)) = (access_key, secret_key) {
1005 &format!(
1006 "aws_access_key_id={};aws_secret_access_key={}",
1007 access_key, secret_key
1008 )
1009 } else {
1010 return Err(SinkError::Config(anyhow!(
1011 "Either assume_role or access_key and secret_key must be provided for Redshift COPY command"
1012 )));
1013 };
1014 Ok(format!(
1015 r#"
1016 COPY {table_name}
1017 FROM '{manifest_dir}'
1018 CREDENTIALS '{credentials}'
1019 FORMAT AS JSON 'auto'
1020 DATEFORMAT 'auto'
1021 TIMEFORMAT 'auto'
1022 MANIFEST;
1023 "#,
1024 table_name = table_name,
1025 manifest_dir = manifest_dir,
1026 credentials = credentials
1027 ))
1028}
1029
1030#[cfg(test)]
1031mod tests {
1032 use super::*;
1033
1034 #[test]
1035 fn test_convert_redshift_decimal_data_type() {
1036 assert_eq!(
1037 convert_redshift_data_type(&DataType::Decimal).unwrap(),
1038 "DECIMAL(38, 10)"
1039 );
1040 }
1041}