risingwave_connector/sink/
dynamodb.rs

1// Copyright 2024 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::collections::{BTreeMap, BTreeSet, HashMap};
16
17use anyhow::{Context, anyhow};
18use aws_sdk_dynamodb as dynamodb;
19use aws_sdk_dynamodb::client::Client;
20use aws_smithy_types::Blob;
21use dynamodb::types::{AttributeValue, KeySchemaElement, TableStatus, WriteRequest};
22use futures::TryFutureExt;
23use risingwave_common::array::{Op, RowRef, StreamChunk};
24use risingwave_common::catalog::Schema;
25use risingwave_common::row::Row as _;
26use risingwave_common::types::{DataType, ScalarRefImpl, ToText};
27use risingwave_common::util::iter_util::ZipEqDebug;
28use serde::Deserialize;
29use serde_with::{DisplayFromStr, serde_as};
30use with_options::WithOptions;
31use write_chunk_future::{DynamoDbPayloadWriter, WriteChunkFuture};
32
33use super::writer::{
34    AsyncTruncateLogSinkerOf, AsyncTruncateSinkWriter, AsyncTruncateSinkWriterExt,
35};
36use super::{Result, Sink, SinkError, SinkParam, SinkWriterParam};
37use crate::connector_common::AwsAuthProps;
38use crate::enforce_secret::EnforceSecret;
39use crate::error::ConnectorResult;
40use crate::sink::log_store::DeliveryFutureManagerAddFuture;
41
42pub const DYNAMO_DB_SINK: &str = "dynamodb";
43
44#[serde_as]
45#[derive(Deserialize, Debug, Clone, WithOptions)]
46pub struct DynamoDbConfig {
47    #[serde(rename = "table", alias = "dynamodb.table")]
48    pub table: String,
49
50    #[serde(rename = "dynamodb.max_batch_rows", default = "default_max_batch_rows")]
51    #[serde_as(as = "DisplayFromStr")]
52    #[deprecated]
53    pub max_batch_rows: usize,
54
55    #[serde(flatten)]
56    pub aws_auth_props: AwsAuthProps,
57
58    #[serde(
59        rename = "dynamodb.max_batch_item_nums",
60        default = "default_max_batch_item_nums"
61    )]
62    #[serde_as(as = "DisplayFromStr")]
63    pub max_batch_item_nums: usize,
64
65    #[serde(
66        rename = "dynamodb.max_future_send_nums",
67        default = "default_max_future_send_nums"
68    )]
69    #[serde_as(as = "DisplayFromStr")]
70    pub max_future_send_nums: usize,
71
72    #[serde(
73        rename = "dynamodb.batch_write_retry_times",
74        default = "default_batch_write_retry_times"
75    )]
76    #[serde_as(as = "DisplayFromStr")]
77    pub batch_write_retry_times: usize,
78
79    #[serde(
80        rename = "dynamodb.batch_write_retry_backoff_ms",
81        default = "default_batch_write_retry_backoff_ms"
82    )]
83    #[serde_as(as = "DisplayFromStr")]
84    pub batch_write_retry_backoff_ms: u64,
85
86    #[serde(flatten)]
87    pub unknown_fields: std::collections::HashMap<String, String>,
88}
89
90crate::impl_sink_unknown_fields!(DynamoDbConfig);
91
92impl EnforceSecret for DynamoDbConfig {
93    fn enforce_one(prop: &str) -> crate::error::ConnectorResult<()> {
94        AwsAuthProps::enforce_one(prop)
95    }
96}
97
98fn default_max_batch_item_nums() -> usize {
99    25
100}
101
102fn default_max_future_send_nums() -> usize {
103    256
104}
105
106fn default_batch_write_retry_times() -> usize {
107    3
108}
109
110fn default_batch_write_retry_backoff_ms() -> u64 {
111    100
112}
113
114fn default_max_batch_rows() -> usize {
115    1024
116}
117
118impl DynamoDbConfig {
119    pub async fn build_client(&self) -> ConnectorResult<Client> {
120        let config = &self.aws_auth_props;
121        let aws_config = config.build_config().await?;
122
123        Ok(Client::new(&aws_config))
124    }
125
126    fn from_btreemap(values: BTreeMap<String, String>) -> Result<Self> {
127        serde_json::from_value::<DynamoDbConfig>(serde_json::to_value(values).unwrap())
128            .map_err(|e| SinkError::Config(anyhow!(e)))
129    }
130}
131
132#[derive(Clone, Debug)]
133pub struct DynamoDbSink {
134    pub config: DynamoDbConfig,
135    schema: Schema,
136    pk_indices: Vec<usize>,
137}
138
139impl EnforceSecret for DynamoDbSink {
140    fn enforce_secret<'a>(
141        prop_iter: impl Iterator<Item = &'a str>,
142    ) -> crate::error::ConnectorResult<()> {
143        for prop in prop_iter {
144            DynamoDbConfig::enforce_one(prop)?;
145        }
146        Ok(())
147    }
148}
149
150impl Sink for DynamoDbSink {
151    type LogSinker = AsyncTruncateLogSinkerOf<DynamoDbSinkWriter>;
152
153    const SINK_NAME: &'static str = DYNAMO_DB_SINK;
154
155    crate::impl_validate_sink_unknown_fields!();
156
157    async fn validate(&self) -> Result<()> {
158        risingwave_common::license::Feature::DynamoDbSink
159            .check_available()
160            .map_err(|e| anyhow::anyhow!(e))?;
161        let client = (self.config.build_client().await)
162            .context("validate DynamoDB sink error")
163            .map_err(SinkError::DynamoDb)?;
164
165        let table_name = &self.config.table;
166        let output = client
167            .describe_table()
168            .table_name(table_name)
169            .send()
170            .await
171            .map_err(|e| anyhow!(e))?;
172        let Some(table) = output.table else {
173            return Err(SinkError::DynamoDb(anyhow!(
174                "table {} not found",
175                table_name
176            )));
177        };
178        if !matches!(table.table_status(), Some(TableStatus::Active)) {
179            return Err(SinkError::DynamoDb(anyhow!(
180                "table {} is not active",
181                table_name
182            )));
183        }
184        let rw_pk_names = rw_pk_names(&self.schema, &self.pk_indices)?;
185        let dynamodb_keys = dynamodb_key_schema_names(table_name, table.key_schema())?;
186        validate_pk_matches_dynamodb_key_schema(table_name, &rw_pk_names, &dynamodb_keys)?;
187
188        Ok(())
189    }
190
191    async fn new_log_sinker(&self, _writer_param: SinkWriterParam) -> Result<Self::LogSinker> {
192        Ok(
193            DynamoDbSinkWriter::new(self.config.clone(), self.schema.clone())
194                .await?
195                .into_log_sinker(self.config.max_future_send_nums),
196        )
197    }
198}
199
200impl TryFrom<SinkParam> for DynamoDbSink {
201    type Error = SinkError;
202
203    fn try_from(param: SinkParam) -> std::result::Result<Self, Self::Error> {
204        let schema = param.schema();
205        let pk_indices = param.downstream_pk_or_empty();
206        let config = DynamoDbConfig::from_btreemap(param.properties)?;
207
208        Ok(Self {
209            config,
210            schema,
211            pk_indices,
212        })
213    }
214}
215
216#[derive(Debug)]
217struct DynamoDbRequest {
218    inner: WriteRequest,
219    key_items: Vec<String>,
220}
221
222impl DynamoDbRequest {
223    fn extract_key(&self) -> Option<&HashMap<String, AttributeValue>> {
224        match (&self.inner.put_request(), &self.inner.delete_request()) {
225            (Some(put_req), None) => Some(&put_req.item),
226            (None, Some(del_req)) => Some(&del_req.key),
227            _ => None,
228        }
229    }
230
231    fn has_same_pk(&self, other: &Self) -> bool {
232        if self.key_items.is_empty() {
233            return false;
234        }
235
236        let Some(key) = self.extract_key() else {
237            return false;
238        };
239        let Some(other_key) = other.extract_key() else {
240            return false;
241        };
242
243        self.key_items.iter().all(|key_item| {
244            matches!(
245                (key.get(key_item), other_key.get(key_item)),
246                (Some(value), Some(other_value)) if value == other_value
247            )
248        })
249    }
250}
251
252pub struct DynamoDbSinkWriter {
253    payload_writer: DynamoDbPayloadWriter,
254    formatter: DynamoDbFormatter,
255    max_future_send_nums: usize,
256}
257
258impl DynamoDbSinkWriter {
259    pub async fn new(config: DynamoDbConfig, schema: Schema) -> Result<Self> {
260        let client = config.build_client().await?;
261        let table_name = &config.table;
262        let output = client
263            .describe_table()
264            .table_name(table_name)
265            .send()
266            .await
267            .map_err(|e| anyhow!(e))?;
268        let Some(table) = output.table else {
269            return Err(SinkError::DynamoDb(anyhow!(
270                "table {} not found",
271                table_name
272            )));
273        };
274        let dynamodb_keys = dynamodb_key_schema_names(table_name, table.key_schema())?;
275
276        let payload_writer = DynamoDbPayloadWriter {
277            client,
278            table: config.table.clone(),
279            dynamodb_keys,
280            max_batch_item_nums: config.max_batch_item_nums,
281            batch_write_retry_times: config.batch_write_retry_times,
282            batch_write_retry_backoff_ms: config.batch_write_retry_backoff_ms,
283        };
284
285        Ok(Self {
286            payload_writer,
287            formatter: DynamoDbFormatter { schema },
288            max_future_send_nums: config.max_future_send_nums,
289        })
290    }
291
292    fn write_chunk_inner(&mut self, chunk: StreamChunk) -> Result<WriteChunkFuture> {
293        let mut request_items = Vec::new();
294        for (op, row) in chunk.rows() {
295            let items = self.formatter.format_row(row)?;
296            match op {
297                Op::Insert | Op::UpdateInsert => {
298                    self.payload_writer
299                        .write_one_insert(items, &mut request_items);
300                }
301                Op::Delete => {
302                    self.payload_writer
303                        .write_one_delete(items, &mut request_items);
304                }
305                Op::UpdateDelete => {}
306            }
307        }
308        Ok(self
309            .payload_writer
310            .write_chunk(request_items, self.max_future_send_nums))
311    }
312}
313
314impl AsyncTruncateSinkWriter for DynamoDbSinkWriter {
315    type DeliveryFuture = WriteChunkFuture;
316
317    async fn write_chunk<'a>(
318        &'a mut self,
319        chunk: StreamChunk,
320        _add_future: DeliveryFutureManagerAddFuture<'a, Self::DeliveryFuture>,
321    ) -> Result<()> {
322        self.write_chunk_inner(chunk)?.map_ok(|_| ()).await?;
323        Ok(())
324    }
325}
326
327struct DynamoDbFormatter {
328    schema: Schema,
329}
330
331impl DynamoDbFormatter {
332    fn format_row(&self, row: RowRef<'_>) -> Result<HashMap<String, AttributeValue>> {
333        row.iter()
334            .zip_eq_debug((self.schema.clone()).into_fields())
335            .map(|(scalar, field)| {
336                map_data(scalar, &field.data_type()).map(|attr| (field.name, attr))
337            })
338            .collect()
339    }
340}
341
342fn map_data(scalar_ref: Option<ScalarRefImpl<'_>>, data_type: &DataType) -> Result<AttributeValue> {
343    let Some(scalar_ref) = scalar_ref else {
344        return Ok(AttributeValue::Null(true));
345    };
346    let attr = match data_type {
347        DataType::Int16
348        | DataType::Int32
349        | DataType::Int64
350        | DataType::Int256
351        | DataType::Float32
352        | DataType::Float64
353        | DataType::Decimal
354        | DataType::Serial => AttributeValue::N(scalar_ref.to_text_with_type(data_type)),
355        // TODO: jsonb as dynamic type (https://github.com/risingwavelabs/risingwave/issues/11699)
356        DataType::Varchar
357        | DataType::Interval
358        | DataType::Date
359        | DataType::Time
360        | DataType::Timestamp
361        | DataType::Timestamptz
362        | DataType::Jsonb => AttributeValue::S(scalar_ref.to_text_with_type(data_type)),
363        DataType::Boolean => AttributeValue::Bool(scalar_ref.into_bool()),
364        DataType::Bytea => AttributeValue::B(Blob::new(scalar_ref.into_bytea())),
365        DataType::List(lt) => {
366            let list_attr = scalar_ref
367                .into_list()
368                .iter()
369                .map(|x| map_data(x, lt.elem()))
370                .collect::<Result<Vec<_>>>()?;
371            AttributeValue::L(list_attr)
372        }
373        DataType::Struct(st) => {
374            let mut map = HashMap::with_capacity(st.len());
375            for (sub_datum_ref, (name, data_type)) in scalar_ref
376                .into_struct()
377                .iter_fields_ref()
378                .zip_eq_debug(st.iter())
379            {
380                let attr = map_data(sub_datum_ref, data_type)?;
381                map.insert(name.to_owned(), attr);
382            }
383            AttributeValue::M(map)
384        }
385        DataType::Map(_m) => {
386            return Err(SinkError::DynamoDb(anyhow!("map is not supported yet")));
387        }
388        DataType::Vector(_) => {
389            return Err(SinkError::DynamoDb(anyhow!("vector is not supported yet")));
390        }
391    };
392    Ok(attr)
393}
394
395fn rw_pk_names(schema: &Schema, pk_indices: &[usize]) -> Result<Vec<String>> {
396    pk_indices
397        .iter()
398        .map(|pk_idx| {
399            schema
400                .fields()
401                .get(*pk_idx)
402                .map(|field| field.name.clone())
403                .ok_or_else(|| {
404                    SinkError::DynamoDb(anyhow!(
405                        "RisingWave primary key column index {} is out of range",
406                        pk_idx
407                    ))
408                })
409        })
410        .collect()
411}
412
413fn dynamodb_key_schema_names(
414    table_name: &str,
415    key_schema: &[KeySchemaElement],
416) -> Result<Vec<String>> {
417    if key_schema.is_empty() {
418        return Err(SinkError::DynamoDb(anyhow!(
419            "table {} key schema is empty",
420            table_name
421        )));
422    }
423
424    Ok(key_schema
425        .iter()
426        .map(|key_element| key_element.attribute_name().to_owned())
427        .collect())
428}
429
430fn validate_pk_matches_dynamodb_key_schema(
431    table_name: &str,
432    rw_pk_names: &[String],
433    dynamodb_keys: &[String],
434) -> Result<()> {
435    let rw_pk_set = rw_pk_names.iter().collect::<BTreeSet<_>>();
436    let dynamodb_key_set = dynamodb_keys.iter().collect::<BTreeSet<_>>();
437    if rw_pk_names.len() != dynamodb_keys.len() || rw_pk_set != dynamodb_key_set {
438        return Err(SinkError::DynamoDb(anyhow!(
439            "DynamoDB table {} primary key {:?} must match RisingWave primary key {:?}",
440            table_name,
441            dynamodb_keys,
442            rw_pk_names
443        )));
444    }
445
446    Ok(())
447}
448
449mod write_chunk_future {
450    use std::collections::HashMap;
451    use std::time::Duration;
452
453    use anyhow::anyhow;
454    use aws_sdk_dynamodb as dynamodb;
455    use aws_sdk_dynamodb::client::Client;
456    use dynamodb::types::{
457        AttributeValue, DeleteRequest, PutRequest, ReturnConsumedCapacity,
458        ReturnItemCollectionMetrics, WriteRequest,
459    };
460    use futures::{FutureExt, StreamExt, TryFuture, TryStreamExt, stream};
461    use itertools::Itertools;
462    use maplit::hashmap;
463    use tokio::time::sleep;
464    use tokio_retry::strategy::{ExponentialBackoff, jitter};
465
466    use super::{DynamoDbRequest, SinkError};
467
468    const MAX_BATCH_WRITE_RETRY_DELAY_MS: u64 = 2000;
469    const MAX_BATCH_WRITE_CONCURRENCY: usize = 256;
470
471    pub struct DynamoDbPayloadWriter {
472        pub client: Client,
473        pub table: String,
474        pub dynamodb_keys: Vec<String>,
475        pub max_batch_item_nums: usize,
476        pub batch_write_retry_times: usize,
477        pub batch_write_retry_backoff_ms: u64,
478    }
479
480    pub type WriteChunkFuture = impl TryFuture<Ok = (), Error = SinkError> + Unpin + Send + 'static;
481
482    impl DynamoDbPayloadWriter {
483        pub fn write_one_insert(
484            &mut self,
485            item: HashMap<String, AttributeValue>,
486            request_items: &mut Vec<DynamoDbRequest>,
487        ) {
488            let put_req = PutRequest::builder().set_item(Some(item)).build().unwrap();
489            let req = WriteRequest::builder().put_request(put_req).build();
490            self.write_one_req(req, request_items);
491        }
492
493        pub fn write_one_delete(
494            &mut self,
495            key: HashMap<String, AttributeValue>,
496            request_items: &mut Vec<DynamoDbRequest>,
497        ) {
498            let key = key
499                .into_iter()
500                .filter(|(k, _)| self.dynamodb_keys.contains(k))
501                .collect();
502            let del_req = DeleteRequest::builder().set_key(Some(key)).build().unwrap();
503            let req = WriteRequest::builder().delete_request(del_req).build();
504            self.write_one_req(req, request_items);
505        }
506
507        pub fn write_one_req(
508            &mut self,
509            req: WriteRequest,
510            request_items: &mut Vec<DynamoDbRequest>,
511        ) {
512            let r_req = DynamoDbRequest {
513                inner: req,
514                key_items: self.dynamodb_keys.clone(),
515            };
516            request_items.retain(|item| !item.has_same_pk(&r_req));
517            request_items.push(r_req);
518        }
519
520        #[define_opaque(WriteChunkFuture)]
521        pub fn write_chunk(
522            &mut self,
523            request_items: Vec<DynamoDbRequest>,
524            max_future_send_nums: usize,
525        ) -> WriteChunkFuture {
526            let client = self.client.clone();
527            let table = self.table.clone();
528            let max_batch_item_nums = self.max_batch_item_nums;
529            let batch_write_retry_times = self.batch_write_retry_times;
530            let batch_write_retry_backoff_ms = self.batch_write_retry_backoff_ms;
531            async move {
532                let chunks = request_items
533                    .into_iter()
534                    .map(|r| r.inner)
535                    .chunks(max_batch_item_nums)
536                    .into_iter()
537                    .map(|chunk| chunk.collect::<Vec<_>>())
538                    .collect_vec();
539                let max_future_send_nums =
540                    max_future_send_nums.clamp(1, MAX_BATCH_WRITE_CONCURRENCY);
541                stream::iter(chunks.into_iter().map(|req_items| {
542                    let client = client.clone();
543                    let table = table.clone();
544                    async move {
545                        let mut req_items = req_items;
546                        let mut retry_count = 0;
547                        let mut retry_backoff = ExponentialBackoff::from_millis(
548                            batch_write_retry_backoff_ms,
549                        )
550                        .factor(2)
551                        .max_delay(Duration::from_millis(MAX_BATCH_WRITE_RETRY_DELAY_MS))
552                        .map(jitter)
553                        .take(batch_write_retry_times);
554
555                        loop {
556                            let return_consumed_capacity = if retry_count == 0 {
557                                ReturnConsumedCapacity::None
558                            } else {
559                                ReturnConsumedCapacity::Total
560                            };
561                            let reqs = hashmap! {
562                                table.clone() => req_items.clone(),
563                            };
564                            let result = client
565                                .batch_write_item()
566                                .set_request_items(Some(reqs))
567                                .return_consumed_capacity(return_consumed_capacity)
568                                .return_item_collection_metrics(ReturnItemCollectionMetrics::None)
569                                .send()
570                                .await;
571
572                            match result {
573                                Ok(output) => {
574                                    let unprocessed_items =
575                                        output.unprocessed_items().cloned().unwrap_or_default();
576                                    if unprocessed_items.is_empty() {
577                                        if retry_count > 0 {
578                                            tracing::warn!(
579                                                retry_count,
580                                                consumed_capacity = ?output.consumed_capacity(),
581                                                "DynamoDB batch write retry succeeded"
582                                            );
583                                        }
584                                        return Ok(());
585                                    }
586
587                                    req_items = unprocessed_items.into_values().flatten().collect();
588                                    if retry_count >= batch_write_retry_times {
589                                        return Err(SinkError::DynamoDb(anyhow!(
590                                            "failed to write {} unprocessed items to DynamoDB sink after {} retries",
591                                            req_items.len(),
592                                            batch_write_retry_times,
593                                        )));
594                                    }
595                                }
596                                Err(e) => {
597                                    return Err(SinkError::DynamoDb(
598                                        anyhow!(e).context("failed to write items to DynamoDB sink"),
599                                    ));
600                                }
601                            }
602
603                            retry_count += 1;
604                            let Some(delay) = retry_backoff.next() else {
605                                return Err(SinkError::DynamoDb(anyhow!(
606                                    "failed to write {} unprocessed items to DynamoDB sink after {} retries",
607                                    req_items.len(),
608                                    batch_write_retry_times,
609                                )));
610                            };
611                            tracing::warn!(
612                                retry_count,
613                                delay_ms = delay.as_millis(),
614                                unprocessed_items_count = req_items.len(),
615                                "retrying DynamoDB batch write"
616                            );
617                            sleep(delay).await;
618                        }
619                    }
620                }))
621                .buffer_unordered(max_future_send_nums)
622                .try_collect::<Vec<_>>()
623                .await?;
624                Ok(())
625            }
626            .boxed()
627        }
628    }
629}
630
631#[cfg(test)]
632mod tests {
633    use aws_sdk_dynamodb::types::{DeleteRequest, KeyType, PutRequest};
634
635    use super::*;
636
637    fn dynamodb_put_request(
638        items: impl IntoIterator<Item = (&'static str, &'static str)>,
639    ) -> DynamoDbRequest {
640        let item = dynamodb_items(items);
641        let put_req = PutRequest::builder().set_item(Some(item)).build().unwrap();
642        DynamoDbRequest {
643            inner: WriteRequest::builder().put_request(put_req).build(),
644            key_items: vec!["pk".to_owned(), "sk".to_owned()],
645        }
646    }
647
648    fn dynamodb_delete_request(
649        items: impl IntoIterator<Item = (&'static str, &'static str)>,
650    ) -> DynamoDbRequest {
651        let key = dynamodb_items(items);
652        let delete_req = DeleteRequest::builder().set_key(Some(key)).build().unwrap();
653        DynamoDbRequest {
654            inner: WriteRequest::builder().delete_request(delete_req).build(),
655            key_items: vec!["pk".to_owned(), "sk".to_owned()],
656        }
657    }
658
659    fn dynamodb_items(
660        items: impl IntoIterator<Item = (&'static str, &'static str)>,
661    ) -> HashMap<String, AttributeValue> {
662        items
663            .into_iter()
664            .map(|(k, v)| (k.to_owned(), AttributeValue::S(v.to_owned())))
665            .collect()
666    }
667
668    #[test]
669    fn dynamodb_request_compares_pk_by_key_attribute() {
670        let req = dynamodb_put_request([("pk", "a"), ("sk", "b")]);
671        let swapped_values = dynamodb_put_request([("pk", "b"), ("sk", "a")]);
672        let same_pk = dynamodb_put_request([("pk", "a"), ("sk", "b")]);
673
674        assert!(!req.has_same_pk(&swapped_values));
675        assert!(req.has_same_pk(&same_pk));
676    }
677
678    #[test]
679    fn dynamodb_request_empty_key_items_never_match() {
680        let mut req = dynamodb_put_request([("pk", "a"), ("sk", "b")]);
681        let same_pk = dynamodb_put_request([("pk", "a"), ("sk", "b")]);
682        req.key_items.clear();
683
684        assert!(!req.has_same_pk(&same_pk));
685    }
686
687    #[test]
688    fn dynamodb_request_compares_put_and_delete_by_composite_pk() {
689        let put = dynamodb_put_request([("pk", "a"), ("sk", "b"), ("value", "1")]);
690        let same_pk_delete = dynamodb_delete_request([("pk", "a"), ("sk", "b")]);
691        let different_hash_key_delete = dynamodb_delete_request([("pk", "x"), ("sk", "b")]);
692        let different_range_key_delete = dynamodb_delete_request([("pk", "a"), ("sk", "x")]);
693
694        assert!(put.has_same_pk(&same_pk_delete));
695        assert!(same_pk_delete.has_same_pk(&put));
696        assert!(!put.has_same_pk(&different_hash_key_delete));
697        assert!(!put.has_same_pk(&different_range_key_delete));
698    }
699
700    #[test]
701    fn dynamodb_key_schema_empty_errors() {
702        let err = dynamodb_key_schema_names("test_table", &[]).unwrap_err();
703
704        assert!(
705            err.to_string()
706                .contains("table test_table key schema is empty")
707        );
708    }
709
710    #[test]
711    fn dynamodb_key_schema_must_match_rw_pk() {
712        let dynamodb_keys = ["pk".to_owned(), "sk".to_owned()];
713        let same_rw_pk = ["sk".to_owned(), "pk".to_owned()];
714        let extra_rw_pk = ["pk".to_owned(), "sk".to_owned(), "extra".to_owned()];
715        let different_rw_pk = ["pk".to_owned(), "other".to_owned()];
716
717        validate_pk_matches_dynamodb_key_schema("test_table", &same_rw_pk, &dynamodb_keys).unwrap();
718
719        assert!(
720            validate_pk_matches_dynamodb_key_schema("test_table", &extra_rw_pk, &dynamodb_keys)
721                .unwrap_err()
722                .to_string()
723                .contains("must match RisingWave primary key")
724        );
725        assert!(
726            validate_pk_matches_dynamodb_key_schema("test_table", &different_rw_pk, &dynamodb_keys)
727                .unwrap_err()
728                .to_string()
729                .contains("must match RisingWave primary key")
730        );
731    }
732
733    #[test]
734    fn dynamodb_key_schema_names_uses_explicit_schema() {
735        let key_schema = vec![
736            KeySchemaElement::builder()
737                .attribute_name("pk")
738                .key_type(KeyType::Hash)
739                .build()
740                .unwrap(),
741            KeySchemaElement::builder()
742                .attribute_name("sk")
743                .key_type(KeyType::Range)
744                .build()
745                .unwrap(),
746        ];
747
748        assert_eq!(
749            dynamodb_key_schema_names("test_table", &key_schema).unwrap(),
750            vec!["pk".to_owned(), "sk".to_owned()]
751        );
752    }
753}
risingwave_connector/sink/dynamodb.rs

risingwave_connector/sink/
dynamodb.rs
