hubcio commented on code in PR #3654: URL: https://github.com/apache/iggy/pull/3654#discussion_r3850365749
########## core/integration/tests/connectors/fixtures/redshift/redshift_mock/copy.rs: ########## @@ -0,0 +1,146 @@ +// 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 std::collections::HashMap; + +use sqlparser::{ + ast::ObjectName, + keywords::Keyword, + parser::{Parser, ParserError}, + tokenizer::Token, +}; + +use crate::connectors::fixtures::redshift::redshift_mock::{ + expect_word, parse_number_literal, parse_string_literal, Review Comment: same for `create.rs`. neither file is declared in `mod.rs`, so they never compile - `expect_word`, `parse_string_literal` and `parse_number_literal` don't exist anywhere in the tree. adding the `mod` lines back fails with E0432. delete both files. ########## core/connectors/sinks/redshift_sink/README.md: ########## @@ -0,0 +1,217 @@ +# Redshift Sink Connector + +Writes Apache Iggy stream messages into Amazon Redshift via S3-staged Parquet +files and a `COPY` load. + +Each connector batch is serialized to a Parquet file and uploaded to the +configured S3 bucket/prefix, then loaded into the target Redshift table with a +`COPY` statement. This makes S3 a staging area rather than a destination in +its own right — Redshift is the system of record for the data. + +Persistent load failures are at-most-once from the runtime's perspective: +messages may already be committed in Iggy before this connector exhausts its +write attempts, so failed loads are logged but not redelivered. + +## Configuration + +```toml +type = "sink" +key = "redshift" +enabled = true +version = 0 +name = "Redshift sink" +path = "../../target/release/libiggy_connector_redshift_sink" +verbose = false + +[[streams]] +stream = "user_events" +topics = ["users", "orders"] +schema = "json" +batch_length = 100 +poll_interval = "5ms" +consumer_group = "redshift_sink" + +[plugin_config] +connection_string = "postgresql://user:pass@localhost:5439/database" +target_table = "iggy_messages" +batch_size = 100 +max_connections = 10 +include_metadata = true +include_checksum = true +include_origin_timestamp = true +payload_format = "varbyte" +aws_access_key_id = "admin" +aws_secret_access_key = "password" +aws_iam_role ="aws_iam_role=arn:aws:iam::0123456789012:role/iggyRole" Review Comment: `build_copy_sql` already prepends `aws_iam_role=`, so this example ends up as `CREDENTIALS 'aws_iam_role=aws_iam_role=arn:...'` and the COPY fails. use the bare arn like `config.toml` does. ########## core/connectors/sinks/redshift_sink/src/lib.rs: ########## @@ -0,0 +1,1493 @@ +// 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. + +mod config; + +use std::{collections::BTreeMap, str::FromStr, sync::Arc, time::Duration}; + +use arrow::{ + array::{ArrayRef, BinaryBuilder, Int64Array, RecordBatch, StringArray, StringBuilder}, + datatypes::{DataType, Field, Schema}, +}; +use async_trait::async_trait; +use humantime::Duration as HumanDuration; +use iggy_connector_sdk::{ + ConsumedMessage, Error, MessagesMetadata, Sink, TopicMetadata, sink_connector, +}; +use parquet::{ + arrow::ArrowWriter, + basic::{Compression, ZstdLevel}, + file::properties::WriterProperties, +}; +use s3::{Bucket, Region, creds::Credentials}; +use secrecy::ExposeSecret; +use sqlx::{AssertSqlSafe, Pool, Postgres, Row, postgres::PgPoolOptions}; +use tokio::sync::Mutex; +use uuid::Uuid; + +use crate::config::{PayloadFormat, RedshiftSinkConfig}; + +sink_connector!(RedshiftSink); + +const DEFAULT_MAX_RETRIES: u32 = 3; +const DEFAULT_RETRY_DELAY: &str = "1s"; +const DEFAULT_MAX_CONNECTIONS: u32 = 5; +const DEFAULT_ARCHIVE_PREFIX: &str = "archive/messages"; + +#[derive(Debug)] +pub struct RedshiftSink { + pub id: u32, + config: RedshiftSinkConfig, + pool: Option<Pool<Postgres>>, + state: Mutex<State>, + verbose: bool, + bucket: Option<Box<Bucket>>, +} + +#[async_trait] +impl Sink for RedshiftSink { + async fn open(&mut self) -> Result<(), Error> { + tracing::info!( + sink_id = self.id, + table = %self.config.target_table, "opening Redshift sink connector" + ); + + // Validating config + self.config.validate()?; + + self.connect().await?; + // Ensuring tables exist + self.ensure_tables_exist().await?; + // Checking for schema drift + self.ensure_schema_match().await?; + Ok(()) + } + + async fn consume( + &self, + topic_metadata: &TopicMetadata, + messages_metadata: MessagesMetadata, + messages: Vec<ConsumedMessage>, + ) -> Result<(), Error> { + tracing::debug!( + sink_id = self.id, + count = messages.len(), + "consuming messages" + ); + self.process_messages(topic_metadata, &messages_metadata, &messages) + .await + } + + async fn close(&mut self) -> Result<(), Error> { + tracing::info!(sink_id = self.id, "closing Redshift sink connector"); + + if let Some(pool) = self.pool.take() { + pool.close().await; + + tracing::debug!(sink_id = self.id, "database pool closed"); + } + + let state = self.state.lock().await; + + tracing::info!( + sink_id = self.id, + messages_processed = state.messages_processed, + batches_loaded = state.batches_loaded, + insertion_errors = state.insertion_errors, + "Redshift sink connector closed", + ); + + Ok(()) + } +} + +impl RedshiftSink { + pub fn new(id: u32, config: RedshiftSinkConfig) -> Self { + let verbose = config.verbose_logging.unwrap_or(false); + + Self { + id, + config, + pool: None, + state: Mutex::new(State::default()), + verbose, + bucket: None, + } + } + + async fn connect(&mut self) -> Result<(), Error> { + let max_connections = self + .config + .max_connections + .unwrap_or(DEFAULT_MAX_CONNECTIONS); + + let redacted = redact_connection_string(self.config.connection_string.expose_secret()); + + tracing::info!(max_connections, dsn = %redacted, "connecting to Redshift"); + + let pool = PgPoolOptions::new() + .max_connections(max_connections) + .connect(self.config.connection_string.expose_secret()) + .await + .map_err(|e| Error::InitError(format!("Failed to connect to Redshift: {e}")))?; + + sqlx::query("SELECT 1").execute(&pool).await.map_err(|e| { + tracing::error!("Tracing failed: {:#?}", e); + Error::InitError(format!("Warehouse connectivity test failed: {e}")) + })?; + + self.pool = Some(pool); + tracing::debug!("Redshift connection pool established"); + + let region = self.build_region()?; + + let credentials = Credentials::new( + self.config + .aws_access_key_id + .as_ref() + .map(|v| v.expose_secret()), + self.config + .aws_secret_access_key + .as_ref() + .map(|v| v.expose_secret()), + None, + None, + None, + ) + .map_err(|e| { + tracing::error!("Failed to create S3 credentials: {e}"); + Error::InvalidConfig + })?; + + let mut bucket = Bucket::new(&self.config.s3_bucket, region, credentials).map_err(|e| { + tracing::error!("Failed to create S3 bucket client: {e}"); + Error::InvalidConfig + })?; + + if self.config.s3_endpoint.is_some() { + bucket = bucket.with_path_style(); + } + + self.bucket = Some(bucket); + + tracing::info!("Redshift sink connector ready"); + + Ok(()) + } + + fn build_region(&self) -> Result<Region, Error> { + if let Some(endpoint) = &self.config.s3_endpoint { + tracing::debug!(endpoint = %endpoint, "using custom S3 endpoint"); + Ok(Region::Custom { + region: self.config.aws_region.clone(), + endpoint: endpoint.clone(), + }) + } else { + Region::from_str(&self.config.aws_region).map_err(|_| Error::InvalidConfig) + } + } + + async fn ensure_tables_exist(&self) -> Result<(), Error> { + let pool = self.get_pool()?; + + let target_table = quote_identifier(&self.config.target_table)?; + let staging_table = quote_identifier(&format!("staging_{}", self.config.target_table))?; + + let payload_type = self.payload_format().sql_type(); + + let target_query = self.build_create_table_sql(&target_table)?; + + let staging_query = self.build_create_table_sql(&staging_table)?; + + tracing::debug!("ensuring staging and target tables exist"); + + sqlx::query(AssertSqlSafe(staging_query)) + .execute(pool) + .await + .map_err(|e| { + tracing::error!(error = %e); + Error::InitError(format!("Failed to create table '{staging_table}': {e}")) + })?; + + tracing::debug!("Staging table created"); + + sqlx::query(AssertSqlSafe(target_query)) + .execute(pool) + .await + .map_err(|e| { + tracing::error!(error = %e); + Error::InitError(format!("Failed to create table '{target_table}': {e}")) + })?; + + tracing::info!( + staging_table = staging_table, + target_table = target_table, + payload_type, + "staging and target tables ready" + ); + + Ok(()) + } + + // This method ensures that the target table schema matches the expected schema. + // it also verifies there is a created_at column + async fn ensure_schema_match(&self) -> Result<(), Error> { + let include_metadata = self.config.include_metadata.unwrap_or(true); + let include_checksum = self.config.include_checksum.unwrap_or(true); + let include_origin_timestamp = self.config.include_origin_timestamp.unwrap_or(true); + let target_table = quote_identifier(&self.config.target_table)?; + let staging_table = quote_identifier(&format!("staging_{}", self.config.target_table))?; + let payload_type = self.payload_format().sql_type(); + let pool = self.get_pool()?; + + let mut expected_cols: BTreeMap<&str, &str> = BTreeMap::new(); + expected_cols.insert("id", "VARCHAR"); + if include_metadata { + expected_cols.insert("iggy_offset", "VARCHAR"); + expected_cols.insert("iggy_timestamp", "VARCHAR"); + expected_cols.insert("iggy_stream", "VARCHAR"); + expected_cols.insert("iggy_topic", "VARCHAR"); + expected_cols.insert("iggy_partition_id", "BIGINT"); + } + if include_checksum { + expected_cols.insert("iggy_checksum", "VARCHAR"); + } + if include_origin_timestamp { + expected_cols.insert("iggy_origin_timestamp", "VARCHAR"); + } + expected_cols.insert("payload", Self::normalize_type(payload_type)); + expected_cols.insert("created_at", "VARCHAR"); + + let target_cols = Self::load_columns(pool, &target_table).await?; + let staging_cols = Self::load_columns(pool, &staging_table).await?; + + let mut mismatches = Self::diff_schema(&target_table, &target_cols, &expected_cols); + mismatches.extend(Self::diff_schema( + &staging_table, + &staging_cols, + &expected_cols, + )); + + tracing::info!("Mismatches: {:?}", mismatches); + + if !mismatches.is_empty() { + return Err(Error::InitError(format!( + "Schema mismatch detected:\n{}", + mismatches.join("\n") + ))); + } + + Ok(()) + } + + fn diff_schema( + table_name: &str, + actual_cols: &BTreeMap<String, String>, + expected_cols: &BTreeMap<&str, &str>, + ) -> Vec<String> { + let mut errors = Vec::new(); + + for (col_name, expected_type) in expected_cols { + match actual_cols.get(*col_name) { + None => errors.push(format!( + "{table_name}: missing column '{col_name}' (expected {expected_type})" + )), + Some(actual_type) if !Self::type_matches(actual_type, expected_type) => errors.push(format!( + "{table_name}: column '{col_name}' type mismatch — expected {expected_type}, found {actual_type}" + )), + _ => {} + } + } + errors + } + + fn type_matches(actual: &str, expected: &str) -> bool { + Self::normalize_type(actual) == Self::normalize_type(expected) + } + + async fn load_columns( + pool: &sqlx::PgPool, + table: &str, + ) -> Result<BTreeMap<String, String>, Error> { + let query = format!( + "SELECT \"column\", type FROM pg_table_def WHERE tablename = '{}'", + table.replace('"', "") + ); + + let rows = sqlx::query(AssertSqlSafe(query)) + .fetch_all(pool) + .await + .map_err(|e| Error::InitError(format!("Failed to read schema for '{table}': {e}")))?; + + if rows.is_empty() { + return Err(Error::InitError(format!( + "Table '{table}' was not found or has no visible columns" + ))); + } + + rows.into_iter() + .map(|row| { + Ok(( + row.try_get::<String, _>("column") + .map_err(|e| Error::InitError(e.to_string()))?, + Self::normalize_type( + &row.try_get::<String, _>("type") + .map_err(|e| Error::InitError(e.to_string()))?, + ) + .to_string(), + )) + }) + .collect() + } + + fn normalize_type(t: &str) -> &'static str { + let base = t.split('(').next().unwrap_or(t).trim(); + + match base.to_ascii_uppercase().as_str() { + "INTEGER" | "INT" | "INT4" => "INTEGER", + // e.g. "bigint" + "INT8" | "BIGINT" => "BIGINT", + // e.g "character varying(40)", "character varying(20)", "character varying(256)" + "TEXT" | "VARCHAR" | "CHARACTER VARYING" => "VARCHAR", + // Having bytea because of the Postgres Test + // e.g. "binary varying(64000)" + "BYTEA" | "VARBYTE" | "VARBINARY" | "BINARY VARYING" => "VARBYTE", + // e.g. "timestamp with time zone" + "TIMESTAMPTZ" | "TIMESTAMP WITH TIME ZONE" => "TIMESTAMPTZ", + _ => "UNKNOWN", + } + } + + async fn process_messages( + &self, + topic_metadata: &TopicMetadata, + messages_metadata: &MessagesMetadata, + messages: &[ConsumedMessage], + ) -> Result<(), Error> { + let batch_size = self.config.batch_size.unwrap_or(100) as usize; + + for batch in messages.chunks(batch_size) { + match self + .insert_batch(batch, topic_metadata, messages_metadata) + .await + { + Ok(path) => { + // Messages were received and ingested to Redshift + if let Some(s3_path) = path { + // Truncate the staging table + if let Err(e) = self.staging_cleanup().await { + tracing::warn!(error = %e, "failed to cleanup staging table"); + } + + // Handle archiving + if let Err(e) = self.archive_parquet(&s3_path).await { + tracing::warn!(error = %e, "failed to archive parquet file: {}", s3_path); + } + + self.state.lock().await.batches_loaded += 1 + } else { + tracing::info!("Zero messages found for processing"); + } + } + Err(e) => { + self.state.lock().await.insertion_errors += batch.len() as u64; + tracing::error!(error = %e, batch_size = batch.len(), "failed to insert batch"); + return Err(e); + } + } + } + + let mut state = self.state.lock().await; + state.messages_processed += messages.len() as u64; + + if self.verbose { + tracing::info!( + sink_id = self.id, + total_processed = state.messages_processed, + batch_received = messages.len(), + table = %self.config.target_table, + batches_loaded = state.batches_loaded, + "processed message batch" + ); + } else { + tracing::debug!( + sink_id = self.id, + total_processed = state.messages_processed, + table = %self.config.target_table, + "processed message batch" + ); + } + + Ok(()) + } + + // This function builds a parquet from messages and metadata + // It uploads the parquet to S3 and returns the path + // It then copies the parquet to Redshift target table via + // a staging table by means of a MERGE statement + // This function treats parquet-generation, uploading to s3, + // copying to staging and merging to target as atomic + // process of focus for this sink connector + async fn insert_batch( + &self, + messages: &[ConsumedMessage], + topic_metadata: &TopicMetadata, + messages_metadata: &MessagesMetadata, + ) -> Result<Option<String>, Error> { + if messages.is_empty() { + return Ok(None); + } + + let include_metadata = self.config.include_metadata.unwrap_or(true); + let include_checksum = self.config.include_checksum.unwrap_or(true); + let include_origin_timestamp = self.config.include_origin_timestamp.unwrap_or(true); + let payload_format = self.payload_format(); + + let record_batch = create_record_batch( + topic_metadata, + messages_metadata, + messages, + include_metadata, + include_checksum, + include_origin_timestamp, + payload_format, + )?; + + let content = encode_parquet(&record_batch)?; + + tracing::debug!( + bytes = content.len(), + rows = record_batch.num_rows(), + "encoded parquet batch" + ); + + let s3_path = self.upload_parquet(&content).await?; + + let schema = record_batch.schema(); + let cols = schema + .fields + .iter() + .map(|f| f.name().as_str()) + .collect::<Vec<_>>(); + + // Copy the parquet file to Redshift staging + // Cleanup + tracing::info!("copying parquet to Redshift staging"); + if let Err(e) = self.copy_parquet(&s3_path, &cols).await { + let key = s3_path + .strip_prefix(&format!("s3://{}/", self.config.s3_bucket)) + .ok_or(Error::InvalidConfigValue("Missing Cleanup S3 path".into()))?; + + self.delete_object(key).await?; + + Err(e)? + } + + tracing::info!("Redshift staging COPY completed"); + + // Do a merge into Redshift target table + self.insert_into_target(&cols).await?; + + tracing::info!("Redshift target table merge completed"); + + tracing::info!(count = messages.len(), path = %s3_path, "batch inserted into Redshift"); + + Ok(Some(s3_path)) + } + + async fn copy_parquet(&self, s3_path: &str, cols: &[&str]) -> Result<(), Error> { + let max_retries = self.get_max_retries(); + let retry_delay = self.get_retry_delay(); + let staging_table = quote_identifier(&format!("staging_{}", self.config.target_table))?; + + let sql = self.build_copy_sql(&staging_table, s3_path, &cols.join(", "))?; + let pool = self.get_pool()?; + + tracing::debug!(table = %self.config.target_table, s3_path, "issuing Redshift COPY"); + + retry_with_backoff( + "Redshift COPY", + max_retries, + retry_delay, + is_transient_error, + || async { + sqlx::query(AssertSqlSafe(sql.as_str())) + .execute(pool) + .await + .map(|_| ()) + }, + ) + .await?; + + tracing::debug!(staging_table = staging_table, "Redshift COPY completed"); + + Ok(()) + } + + async fn insert_into_target(&self, cols: &[&str]) -> Result<(), Error> { + let max_retries = self.get_max_retries(); + let retry_delay = self.get_retry_delay(); + let target_table = quote_identifier(&self.config.target_table)?; + let staging_table = quote_identifier(&format!("staging_{}", self.config.target_table))?; + let sql = self.build_insert_sql(cols, &staging_table, &target_table); + + let pool = self.get_pool()?; + + tracing::debug!(table = %self.config.target_table, "issuing Redshift MERGE"); + + retry_with_backoff( + "Redshift INSERT", + max_retries, + retry_delay, + is_transient_error, + || async { + sqlx::query(AssertSqlSafe(sql.as_str())) + .execute(pool) + .await + .map(|_| ()) + }, + ) + .await?; + + tracing::debug!(staging_table = %staging_table, target_table = %target_table, "Redshift INSERT completed"); + + Ok(()) + } + + async fn staging_cleanup(&self) -> Result<(), Error> { + let max_retries = self.get_max_retries(); + let retry_delay = self.get_retry_delay(); + let staging_table = quote_identifier(&format!("staging_{}", self.config.target_table))?; + let sql = self.build_truncate_sql(&staging_table); + let pool = self.get_pool()?; + + tracing::debug!(table = %self.config.target_table, "issuing Redshift TRUNCATE"); + + retry_with_backoff( + "Redshift TRUNCATE", + max_retries, + retry_delay, + is_transient_error, + || async { + sqlx::query(AssertSqlSafe(sql.as_str())) + .execute(pool) + .await + .map(|_| ()) + }, + ) + .await?; + + tracing::debug!(table = %self.config.target_table, "Redshift TRUNCATE completed"); + + Ok(()) + } + + fn build_create_table_sql(&self, table_name: &str) -> Result<String, Error> { + let include_metadata = self.config.include_metadata.unwrap_or(true); + let include_checksum = self.config.include_checksum.unwrap_or(true); + let include_origin_timestamp = self.config.include_origin_timestamp.unwrap_or(true); + let payload_type = self.payload_format().sql_type(); + + let mut query = format!("CREATE TABLE IF NOT EXISTS {table_name} (id VARCHAR(40)"); + + if include_metadata { + query.push_str(", iggy_offset VARCHAR(20), iggy_timestamp VARCHAR(20), iggy_stream VARCHAR, iggy_topic VARCHAR, iggy_partition_id BIGINT"); + } + + if include_checksum { + query.push_str(", iggy_checksum VARCHAR"); + } + + if include_origin_timestamp { + query.push_str(", iggy_origin_timestamp VARCHAR(20)"); + } + + query.push_str(&format!(", payload {payload_type}")); + query.push_str(", created_at VARCHAR);"); + + Ok(query) + } + + fn build_copy_sql( + &self, + staging_table: &str, + s3_path: &str, + cols: &str, + ) -> Result<String, Error> { + // Redshift allows this from the docs + // https://docs.aws.amazon.com/redshift/latest/dg/r_COPY_command_examples.html + let iam_role = quote_identifier(&self.config.aws_iam_role)?.replace('"', ""); + + let region = quote_identifier(&self.config.aws_region)?.replace('"', ""); + + Ok(format!( + "COPY {} ({}) FROM '{}' CREDENTIALS 'aws_iam_role={}' FORMAT AS PARQUET REGION '{}';", Review Comment: `REGION` is not supported with parquet COPY. the aws notes on columnar COPY (https://docs.aws.amazon.com/redshift/latest/dg/copy-usage_notes-copy-from-columnar.html) say it fails with "REGION argument is not supported for PARQUET based COPY", and it's missing from their supported parameter list. the mock can't catch this, `execute_copy` only checks for `Parquet` and ignores the rest. drop the clause - COPY defaults to the cluster region and the bucket has to be in that region anyway. also, the thread only has DDL screenshots so far, one pasted real COPY output would settle it. -- 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]
