ethanlin01x commented on code in PR #3996: URL: https://github.com/apache/iggy/pull/3996#discussion_r3944270086
########## core/connectors/sinks/dynamodb_sink/src/lib.rs: ########## @@ -0,0 +1,1348 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you 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 async_trait::async_trait; +use aws_config::BehaviorVersion; +use aws_sdk_dynamodb::Client; +use aws_sdk_dynamodb::config::{Credentials, Region}; +use aws_sdk_dynamodb::error::{ProvideErrorMetadata, SdkError}; +use aws_sdk_dynamodb::primitives::Blob; +use aws_sdk_dynamodb::types::{ + AttributeValue, KeySchemaElement, KeyType, PutRequest, WriteRequest, +}; +use humantime::Duration as HumanDuration; +use iggy_connector_sdk::retry::{exponential_backoff, jitter}; +use iggy_connector_sdk::{ + ConsumedMessage, Error, MessagesMetadata, Payload, Sink, TopicMetadata, sink_connector, +}; +use secrecy::{ExposeSecret, SecretString}; +use serde::Deserialize; +use simd_json::{OwnedValue, StaticNode}; +use std::collections::HashMap; +use std::fmt::Write; +use std::str::FromStr; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::time::Duration; +use tracing::{debug, error, info, warn}; + +sink_connector!(DynamoDbSink); + +/// `BatchWriteItem` rejects any request carrying more than 25 write requests. +const MAX_BATCH_WRITE_ITEMS: usize = 25; +/// DynamoDB rejects items larger than 400 KB. +const MAX_ITEM_SIZE: usize = 400 * 1024; +const DEFAULT_PARTITION_KEY_FIELD: &str = "iggy_id"; +const DEFAULT_MAX_RETRIES: u32 = 3; +const DEFAULT_RETRY_DELAY: &str = "500ms"; +const DEFAULT_MAX_RETRY_DELAY: &str = "5s"; +const PAYLOAD_FIELD: &str = "payload"; +const CREDENTIALS_PROVIDER_NAME: &str = "iggy-dynamodb-sink"; + +#[derive(Debug)] +pub struct DynamoDbSink { + pub id: u32, + client: Option<Client>, + config: DynamoDbSinkConfig, + partition_key_field: String, + sort_key_field: Option<String>, + batch_size: usize, + include_metadata: bool, + include_checksum: bool, + include_origin_timestamp: bool, + max_item_size: usize, + max_retries: u32, + retry_delay: Duration, + max_retry_delay: Duration, + verbose: bool, + items_written: AtomicU64, + items_skipped: AtomicU64, + items_deduplicated: AtomicU64, + write_errors: AtomicU64, +} + +/// Only `Deserialize` - nothing serializes a plugin config back out, and the +/// missing impl is what keeps the credentials unserializable. +#[derive(Debug, Clone, Deserialize)] +pub struct DynamoDbSinkConfig { + pub table: String, + pub region: Option<String>, + pub endpoint: Option<String>, + pub access_key_id: Option<SecretString>, + pub secret_access_key: Option<SecretString>, + pub session_token: Option<SecretString>, + pub partition_key_field: Option<String>, + pub sort_key_field: Option<String>, + pub batch_size: Option<u32>, + pub include_metadata: Option<bool>, + pub include_checksum: Option<bool>, + pub include_origin_timestamp: Option<bool>, + pub max_item_size: Option<usize>, + pub max_retries: Option<u32>, + pub retry_delay: Option<String>, + pub max_retry_delay: Option<String>, + pub verbose_logging: Option<bool>, +} + +impl DynamoDbSink { + pub fn new(id: u32, config: DynamoDbSinkConfig) -> Self { + let partition_key_field = config + .partition_key_field + .clone() + .unwrap_or_else(|| DEFAULT_PARTITION_KEY_FIELD.to_owned()); + let sort_key_field = config.sort_key_field.clone(); + let batch_size = config + .batch_size + .unwrap_or(MAX_BATCH_WRITE_ITEMS as u32) + .clamp(1, MAX_BATCH_WRITE_ITEMS as u32) as usize; + let include_metadata = config.include_metadata.unwrap_or(true); + let include_checksum = config.include_checksum.unwrap_or(true); + let include_origin_timestamp = config.include_origin_timestamp.unwrap_or(true); + let max_item_size = config + .max_item_size + .unwrap_or(MAX_ITEM_SIZE) + .min(MAX_ITEM_SIZE); + let max_retries = config.max_retries.unwrap_or(DEFAULT_MAX_RETRIES); + let retry_delay = parse_duration(config.retry_delay.as_deref(), DEFAULT_RETRY_DELAY); + let mut max_retry_delay = + parse_duration(config.max_retry_delay.as_deref(), DEFAULT_MAX_RETRY_DELAY); + if max_retry_delay < retry_delay { + warn!( + "DynamoDB sink ID: {id} has max_retry_delay below retry_delay, raising it to the retry delay" + ); + max_retry_delay = retry_delay; + } + let verbose = config.verbose_logging.unwrap_or(false); + + DynamoDbSink { + id, + client: None, + config, + partition_key_field, + sort_key_field, + batch_size, + include_metadata, + include_checksum, + include_origin_timestamp, + max_item_size, + max_retries, + retry_delay, + max_retry_delay, + verbose, + items_written: AtomicU64::new(0), + items_skipped: AtomicU64::new(0), + items_deduplicated: AtomicU64::new(0), + write_errors: AtomicU64::new(0), + } + } +} + +#[async_trait] +impl Sink for DynamoDbSink { + async fn open(&mut self) -> Result<(), Error> { + info!( + "Opening DynamoDB sink connector with ID: {}, table: {}", + self.id, self.config.table + ); + let client = self.build_client().await?; + let description = client + .describe_table() + .table_name(&self.config.table) + .send() + .await + .map_err(|error| { + Error::InitError(format!( + "DynamoDB table '{}' is not reachable, error: {}", + self.config.table, + describe_sdk_error(&error) + )) + })?; + self.validate_key_schema( + description + .table + .and_then(|table| table.key_schema) + .unwrap_or_default(), + )?; + + self.client = Some(client); + info!( + "Opened DynamoDB sink connector with ID: {}, table: {}", + self.id, self.config.table + ); + Ok(()) + } + + async fn consume( + &self, + topic_metadata: &TopicMetadata, + messages_metadata: MessagesMetadata, + messages: Vec<ConsumedMessage>, + ) -> Result<(), Error> { + self.write_messages(topic_metadata, &messages_metadata, messages) + .await + } + + async fn close(&mut self) -> Result<(), Error> { + info!("Closing DynamoDB sink connector with ID: {}", self.id); + self.client.take(); + info!( + "Closed DynamoDB sink connector with ID: {}, written: {}, skipped: {}, deduplicated: {}, errors: {}", + self.id, + self.items_written.load(Ordering::Relaxed), + self.items_skipped.load(Ordering::Relaxed), + self.items_deduplicated.load(Ordering::Relaxed), + self.write_errors.load(Ordering::Relaxed) + ); + Ok(()) + } +} + +impl DynamoDbSink { + async fn build_client(&self) -> Result<Client, Error> { + if self.config.access_key_id.is_some() != self.config.secret_access_key.is_some() { + return Err(Error::InvalidConfigValue( + "Partially configured credentials. You must provide both access_key_id \ + and secret_access_key, or omit both." + .to_owned(), + )); + } + + let mut loader = aws_config::defaults(BehaviorVersion::latest()); + if let Some(region) = &self.config.region { + loader = loader.region(Region::new(region.clone())); + } + if let Some(endpoint) = &self.config.endpoint { + info!("Using custom DynamoDB endpoint: {endpoint}"); + loader = loader.endpoint_url(endpoint); + } Review Comment: Fixed in 84caa3a2f -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
