ryerraguntla commented on code in PR #3568: URL: https://github.com/apache/iggy/pull/3568#discussion_r3603820903
########## core/connectors/sources/mysql_source/src/lib.rs: ########## @@ -0,0 +1,1563 @@ +// 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 base64::Engine; +use chrono::{NaiveDate, NaiveDateTime, NaiveTime}; +use humantime::Duration as HumanDuration; +use iggy_common::{DateTime, Utc}; +use iggy_connector_sdk::{ + ConnectorState, Error, ProducedMessage, ProducedMessages, Schema, Source, source_connector, +}; +use secrecy::{ExposeSecret, SecretString}; +use serde::{Deserialize, Serialize}; +use sqlx::mysql::{MySqlDatabaseError, MySqlRow}; +use sqlx::{Column, MySql, Pool, Row, TypeInfo, mysql::MySqlPoolOptions}; +use std::collections::HashMap; +use std::str::FromStr; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::time::Duration; +use tokio::sync::Mutex; +use tracing::{debug, error, info, warn}; +use uuid::Uuid; + +source_connector!(MySqlSource); + +const DEFAULT_MAX_RETRIES: u32 = 3; +const DEFAULT_RETRY_DELAY: &str = "1s"; + +/// Fixed namespace so a given (table, key) always hashes to the same message id +/// across restarts. This is what makes a replayed row idempotent for downstream +/// dedup, so the value must never change once connectors are in the field. +const MESSAGE_ID_NAMESPACE: Uuid = Uuid::from_u128(0x8f3b1e6a4c9d4f2a9b7c0d1e2f3a4b5c); + +#[derive(Debug)] +pub struct MySqlSource { + pub id: u32, + pool: Option<Pool<MySql>>, + config: MySqlSourceConfig, + state: Mutex<State>, + verbose: bool, + retry_delay: Duration, + poll_interval: Duration, + last_batch_full: AtomicBool, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct MySqlSourceConfig { + #[serde(serialize_with = "iggy_common::serde_secret::serialize_secret")] + pub connection_string: SecretString, + pub tables: Vec<String>, + pub poll_interval: Option<String>, + pub batch_size: Option<u32>, + pub tracking_column: Option<String>, + pub initial_offset: Option<String>, + pub max_connections: Option<u32>, + pub custom_query: Option<String>, + pub snake_case_columns: Option<bool>, + pub include_metadata: Option<bool>, + pub delete_after_read: Option<bool>, + pub processed_column: Option<String>, + pub primary_key_column: Option<String>, + pub payload_column: Option<String>, + pub payload_format: Option<String>, + pub verbose_logging: Option<bool>, + pub max_retries: Option<u32>, + pub retry_delay: Option<String>, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum PayloadFormat { + #[default] + Json, + Bytea, + Text, + JsonDirect, +} + +struct ProcessedRow { + message: ProducedMessage, + max_offset: Option<String>, + row_pk: Option<String>, +} + +/// One table's fully processed but not-yet-committed work. Built in the +/// side-effect-free first phase of `poll_tables`, then marked/deleted and +/// published in the second phase so a table's messages are emitted only once +/// its rows are marked. +struct TableBatch { + table: String, + messages: Vec<ProducedMessage>, + processed_ids: Vec<String>, + max_offset: Option<String>, +} + +impl PayloadFormat { + fn from_config(s: Option<&str>) -> Self { + match s.map(|s| s.to_lowercase()).as_deref() { + Some("bytea") | Some("raw") => PayloadFormat::Bytea, + Some("text") => PayloadFormat::Text, + Some("json_direct") => PayloadFormat::JsonDirect, + _ => PayloadFormat::Json, + } + } +} + +#[derive(Debug, Serialize, Deserialize)] +pub struct DatabaseRecord { + pub table_name: String, + pub operation_type: String, + pub timestamp: DateTime<Utc>, + pub data: serde_json::Value, + pub old_data: Option<serde_json::Value>, +} + +#[derive(Clone, Copy)] +struct RowProcessingConfig<'a> { + table: &'a str, + tracking_column: &'a str, + pk_column: &'a str, + payload_format: PayloadFormat, + payload_col: &'a str, + snake_case_columns: bool, + include_metadata: bool, +} + +#[derive(Debug, Serialize, Deserialize)] +struct State { + tracking_offsets: HashMap<String, String>, + processed_rows: u64, +} + +const CONNECTOR_NAME: &str = "MySQL source"; + +#[async_trait] +impl Source for MySqlSource { + async fn open(&mut self) -> Result<(), Error> { + info!( + "Opening MySQL source connector with ID: {}, Tables: {:?}", + self.id, self.config.tables + ); + + if let Some(ref col) = self.config.payload_column + && !col.is_empty() + && PayloadFormat::from_config(self.config.payload_format.as_deref()) + == PayloadFormat::Json + { + return Err(Error::InitError( + "payload_format must be 'bytea', 'text', or 'json_direct' when payload_column is set" + .to_string(), + )); + } + + // custom_query never changes after startup, so validate it once here + if let Some(ref query) = self.config.custom_query { + self.validate_custom_query(query)?; + + if query.contains("$offset") && self.config.initial_offset.is_none() { + let state = self.state.lock().await; + let missing: Vec<&str> = self + .config + .tables + .iter() + .filter(|t| !state.tracking_offsets.contains_key(*t)) + .map(String::as_str) + .collect(); + if !missing.is_empty() { + return Err(Error::InitError(format!( + "custom_query uses $offset but initial_offset is not set and no stored offset exists for table(s): {}", + missing.join(", ") + ))); + } + } + } + + if self.config.delete_after_read.unwrap_or(false) && self.config.processed_column.is_some() + { + warn!( + "both delete_after_read and processed_column are set; delete_after_read takes precedence, so rows are deleted and processed_column only acts as a poll filter" + ); + } + + self.connect().await?; + + info!( + "MySQL source connector with ID: {} opened successfully", + self.id + ); + Ok(()) + } + + async fn poll(&self) -> Result<ProducedMessages, Error> { + // Skip the pacing delay while draining a backlog + if !self.last_batch_full.load(Ordering::Relaxed) { + tokio::time::sleep(self.poll_interval).await; + } + + let messages = self.poll_tables().await?; + + let state = self.state.lock().await; + if self.verbose { + info!( + "MySQL source connector ID: {} produced {} messages. Total processed: {}", + self.id, + messages.len(), + state.processed_rows + ); + } else { + debug!( + "MySQL source connector ID: {} produced {} messages. Total processed: {}", + self.id, + messages.len(), + state.processed_rows + ); + } + + let schema = match self.payload_format() { + PayloadFormat::Bytea => Schema::Raw, + PayloadFormat::Text => Schema::Text, + PayloadFormat::JsonDirect | PayloadFormat::Json => Schema::Json, + }; + + // Idle cycles produce no messages and leave offsets/processed_rows + // untouched, so return None to let the runtime skip the state fsync. + let persisted_state = if messages.is_empty() { + None + } else { + self.serialize_state(&state) + }; + + Ok(ProducedMessages { + schema, + messages, + state: persisted_state, + }) + } + + async fn close(&mut self) -> Result<(), Error> { + if let Some(pool) = self.pool.take() { + pool.close().await; + info!("MySQL connection pool closed for connector ID: {}", self.id); + } + + let state = self.state.lock().await; + info!( + "MySQL source connector ID: {} closed. Total rows processed: {}", + self.id, state.processed_rows + ); + Ok(()) + } +} + +impl MySqlSource { + pub fn new(id: u32, config: MySqlSourceConfig, state: Option<ConnectorState>) -> Self { + let verbose = config.verbose_logging.unwrap_or(false); + let restored_state = state + .and_then(|s| s.deserialize::<State>(CONNECTOR_NAME, id)) + .inspect(|s| { + info!( + "Restored state for {CONNECTOR_NAME} connector with ID: {id}. \ + Tracking offsets: {:?}, processed rows: {}", + s.tracking_offsets, s.processed_rows + ); + }); + + let delay_str = config.retry_delay.as_deref().unwrap_or(DEFAULT_RETRY_DELAY); + let retry_delay = HumanDuration::from_str(delay_str) + .map(|duration| duration.into()) + .unwrap_or_else(|_| Duration::from_secs(1)); + let interval_str = config.poll_interval.as_deref().unwrap_or("10s"); + let poll_interval = HumanDuration::from_str(interval_str) + .map(|duration| duration.into()) + .unwrap_or_else(|_| Duration::from_secs(10)); + MySqlSource { + id, + pool: None, + config, + state: Mutex::new(restored_state.unwrap_or(State { + tracking_offsets: HashMap::new(), + processed_rows: 0, + })), + verbose, + retry_delay, + poll_interval, + last_batch_full: AtomicBool::new(false), + } + } + + async fn connect(&mut self) -> Result<(), Error> { + let max_connections = self.config.max_connections.unwrap_or(10); + let redacted = redact_connection_string(self.config.connection_string.expose_secret()); + + info!("Connecting to MySQL with max {max_connections} connections: {redacted}"); + + let pool = MySqlPoolOptions::new() + .max_connections(max_connections) + .connect(self.config.connection_string.expose_secret()) + .await + .map_err(|e| Error::InitError(format!("Failed to connect to MySQL: {e}")))?; + + sqlx::query("SELECT 1") + .execute(&pool) + .await + .map_err(|e| Error::InitError(format!("Database connectivity test failed: {e}")))?; + + self.pool = Some(pool); + info!("Connected to MySQL database with {max_connections} max connections"); + Ok(()) + } + + fn payload_format(&self) -> PayloadFormat { + if let Some(ref payload_col) = self.config.payload_column + && !payload_col.is_empty() + { + return PayloadFormat::from_config(self.config.payload_format.as_deref()); + } + PayloadFormat::Json + } + + fn serialize_state(&self, state: &State) -> Option<ConnectorState> { + ConnectorState::serialize(state, CONNECTOR_NAME, self.id) + } + + fn get_pool(&self) -> Result<&Pool<MySql>, Error> { + self.pool + .as_ref() + .ok_or_else(|| Error::InitError("Database not connected".to_string())) + } + + fn extract_payload_column( + &self, + row: &MySqlRow, + column_index: usize, + format: PayloadFormat, + ) -> Result<Vec<u8>, Error> { + let column_name = row.columns()[column_index].name(); + match format { + PayloadFormat::Bytea => { + let bytes: Option<Vec<u8>> = row.try_get(column_index).map_err(|e| { + Error::InvalidRecordValue(format!( + "payload column '{column_name}' as bytea: {e}" + )) + })?; + Ok(bytes.unwrap_or_default()) + } + PayloadFormat::Text => { + let text: Option<String> = row.try_get(column_index).map_err(|e| { + Error::InvalidRecordValue(format!( + "payload column '{column_name}' as text (invalid UTF-8?): {e}" + )) + })?; + Ok(text.unwrap_or_default().into_bytes()) + } + PayloadFormat::JsonDirect => { + let json_value: Option<serde_json::Value> = + row.try_get(column_index).map_err(|e| { + Error::InvalidRecordValue(format!( + "payload column '{column_name}' as json_direct (invalid JSON?): {e}" + )) + })?; + simd_json::to_vec(&json_value.unwrap_or(serde_json::Value::Null)).map_err(|e| { + Error::InvalidRecordValue(format!( + "payload column '{column_name}': failed to serialize JSON: {e}" + )) + }) + } + PayloadFormat::Json => Err(Error::InvalidConfig), // unreachable! if payload_column is there then payload_format can never be json + } + } + + fn substitute_query_params( + &self, + query: &str, + table: &str, + last_offset: &Option<String>, + batch_size: u32, + ) -> String { + let offset_value = last_offset + .clone() + .or_else(|| self.config.initial_offset.clone()) + .unwrap_or_default(); + let offset = format_offset_value(&offset_value); + let now = Utc::now(); + + query + .replace("$table", table) + .replace("$offset", &offset) + .replace("$limit", &batch_size.to_string()) + .replace("$now_unix", &now.timestamp().to_string()) + .replace("$now", &now.to_rfc3339()) + } + + fn validate_custom_query(&self, query: &str) -> Result<(), Error> { + let query_upper = query.to_uppercase(); + if !query_upper.contains("SELECT") { + warn!("Custom query should contain SELECT statement"); + } + if query.contains("$table") && self.config.tables.is_empty() { + return Err(Error::InvalidConfig); + } + Ok(()) + } + + fn build_polling_query( + &self, + table: &str, + tracking_column: &str, + last_offset: &Option<String>, + batch_size: u32, + ) -> Result<String, Error> { + let quoted_table = quote_qualified_identifier(table)?; + let quoted_tracking = quote_identifier(tracking_column)?; + + let base_query = format!("SELECT * FROM {quoted_table}"); + + let mut conditions = Vec::new(); + + if let Some(offset) = last_offset { + conditions.push(format!( + "{quoted_tracking} > {}", + format_offset_value(offset) + )); + } else if let Some(initial) = &self.config.initial_offset { + conditions.push(format!( + "{quoted_tracking} > {}", + format_offset_value(initial) + )); + } + + if let Some(processed_col) = &self.config.processed_column { + let quoted_processed = quote_identifier(processed_col)?; + conditions.push(format!("{quoted_processed} = FALSE")); + } + + let where_clause = if conditions.is_empty() { + String::new() + } else { + format!(" WHERE {}", conditions.join(" AND ")) + }; + + let order_clause = format!(" ORDER BY {quoted_tracking} ASC"); + let limit_clause = format!(" LIMIT {batch_size}"); + + Ok(format!( + "{base_query}{where_clause}{order_clause}{limit_clause}" + )) + } + + fn get_max_retries(&self) -> u32 { + self.config.max_retries.unwrap_or(DEFAULT_MAX_RETRIES) + } + + async fn mark_or_delete_processed_rows( + &self, + pool: &Pool<MySql>, + table: &str, + pk_column: &str, + ids: &[String], + ) -> Result<(), Error> { + if ids.is_empty() { + return Ok(()); + } + + let quoted_table = quote_qualified_identifier(table)?; + let quoted_pk = quote_identifier(pk_column)?; + + let ids_list = ids + .iter() + .map(|id| { + if id.parse::<i64>().is_ok() || id.parse::<u64>().is_ok() { + id.clone() + } else { + format!( + "'{}'", + id.replace('\\', "\\\\") + .replace('\'', "''") + .replace('\0', "") + ) + } + }) + .collect::<Vec<_>>() + .join(", "); + + let query = if self.config.delete_after_read.unwrap_or(false) { + if self.verbose { + info!("Deleting {} processed rows from '{table}'", ids.len()); + } else { + debug!("Deleting {} processed rows from '{table}'", ids.len()); + } + format!("DELETE FROM {quoted_table} WHERE {quoted_pk} IN ({ids_list})") + } else if let Some(processed_col) = &self.config.processed_column { + let quoted_processed = quote_identifier(processed_col)?; + if self.verbose { + info!("Marking {} rows as processed in '{table}'", ids.len()); + } else { + debug!("Marking {} rows as processed in '{table}'", ids.len()); + } + format!( + "UPDATE {quoted_table} SET {quoted_processed} = TRUE WHERE {quoted_pk} IN ({ids_list})" + ) + } else { + // Offset-tracking only: nothing to mark or delete. + return Ok(()); + }; + + with_retry( + || sqlx::query(sqlx::AssertSqlSafe(query.as_str())).execute(pool), + self.get_max_retries(), + self.retry_delay.as_millis() as u64, + ) + .await + .map(|_| ()) + } + + async fn poll_tables(&self) -> Result<Vec<ProducedMessage>, Error> { + let pool = self.get_pool()?; + + let batch_size = self.config.batch_size.unwrap_or(1000); + let tracking_column = self.config.tracking_column.as_deref().unwrap_or("id"); + let pk_column = self + .config + .primary_key_column + .as_deref() + .unwrap_or(tracking_column); + + let row_config = RowProcessingConfig { + table: "", + tracking_column, + pk_column, + payload_format: self.payload_format(), + payload_col: self.config.payload_column.as_deref().unwrap_or(""), + snake_case_columns: self.config.snake_case_columns.unwrap_or(false), + include_metadata: self.config.include_metadata.unwrap_or(true), + }; + + // Phase 1: fetch and process every table into its own buffer without any + // side effect. A `process_row` failure is deterministic (a row that cannot + // be decoded fails identically every cycle), so a failing table is logged + // with its resume offset and skipped rather than aborting the whole poll: + // nothing has been marked yet, its rows stay in MySQL, and the remaining + // tables still make progress. The operator sees exactly which table and + // offset to fix. + let mut batches: Vec<TableBatch> = Vec::with_capacity(self.config.tables.len()); + let mut any_batch_full = false; + + for table in &self.config.tables { + // Get last offset with minimal lock time + let last_offset = { + let state = self.state.lock().await; + state.tracking_offsets.get(table).cloned() + }; + + match self + .fetch_table_batch( + table, + &row_config, + tracking_column, + batch_size, + &last_offset, + ) + .await + { + Ok(batch) => { + if self.verbose { + info!("Fetched {} rows from table '{table}'", batch.messages.len()); + } else { + debug!("Fetched {} rows from table '{table}'", batch.messages.len()); + } + // A full batch means the LIMIT was hit, so more rows are waiting. + any_batch_full |= batch.messages.len() as u32 >= batch_size; + batches.push(batch); + } + Err(error) => { + error!( + "Failed to process table '{table}' at offset {}, skipping this cycle: {error}", + last_offset.as_deref().unwrap_or("<start>") + ); + } + } + } + + // Phase 2: commit each table independently. If the mark/delete call fails + // for a table we just skip its messages this cycle instead of queuing them + // - `mark_or_delete_processed_rows` already retries transient errors, so a + // failure here is permanent and the table's rows stay in MySQL to be + // retried next cycle, without blocking the others. + // + // Doesn't cover the case where mark/delete succeeds but the runtime never + // gets to publish the batch (crash, send failure, etc) - those rows are + // gone from MySQL with nothing delivered. That's just how + // delete_after_read/processed_column work, not something this loop fixes. + let mut messages = Vec::new(); + let mut state_updates: Vec<(String, String)> = Vec::new(); + let mut total_processed: u64 = 0; + + let mark_or_delete = self.config.delete_after_read.unwrap_or(false) + || self.config.processed_column.is_some(); + + for mut batch in batches { + if mark_or_delete && !batch.messages.is_empty() && batch.processed_ids.is_empty() { + warn!( + "Table '{}': mark/delete is configured but no primary keys were extracted \ + from {} row(s); rows publish without being marked or deleted \ + (check that '{pk_column}' is projected and scalar)", + batch.table, + batch.messages.len() + ); + } + + if !batch.processed_ids.is_empty() + && let Err(error) = self + .mark_or_delete_processed_rows( + pool, + &batch.table, + pk_column, + &batch.processed_ids, + ) + .await + { + error!( + "Failed to mark or delete processed rows for table '{}', skipping this cycle: {error}", + batch.table + ); + continue; + } + + total_processed += batch.messages.len() as u64; + messages.append(&mut batch.messages); + if let Some(offset) = batch.max_offset { + state_updates.push((batch.table, offset)); + } + } + + // Apply all state updates with a single lock acquisition + { + let mut state = self.state.lock().await; + state.processed_rows += total_processed; + for (table, offset) in state_updates { + state.tracking_offsets.insert(table, offset); + } + } + + self.last_batch_full + .store(any_batch_full, Ordering::Relaxed); + Ok(messages) + } + + async fn fetch_table_batch( + &self, + table: &str, + row_config: &RowProcessingConfig<'_>, + tracking_column: &str, + batch_size: u32, + last_offset: &Option<String>, + ) -> Result<TableBatch, Error> { + let pool = self.get_pool()?; + let table_config = RowProcessingConfig { + table, + ..*row_config + }; + + let query = if let Some(custom_query) = &self.config.custom_query { + self.substitute_query_params(custom_query, table, last_offset, batch_size) + } else { + self.build_polling_query(table, tracking_column, last_offset, batch_size)? + }; + + // Database I/O without holding the lock + let rows = with_retry( + || sqlx::query(sqlx::AssertSqlSafe(query.as_str())).fetch_all(pool), + self.get_max_retries(), + self.retry_delay.as_millis() as u64, + ) + .await?; + + let mut batch = TableBatch { + table: table.to_string(), + messages: Vec::with_capacity(rows.len()), + processed_ids: Vec::new(), + max_offset: None, + }; + for row in rows { + let processed = self.process_row(&row, &table_config).map_err(|e| { + error!( + "Failed to decode row in table '{table}' (pk {}): {e}", + pk_for_log(&row, table_config.pk_column) + ); + e + })?; + + if let Some(pk) = processed.row_pk { + batch.processed_ids.push(pk); + } + if let Some(offset) = processed.max_offset { + batch.max_offset = Some(offset); + } + + batch.messages.push(processed.message); + } + + Ok(batch) + } + + fn process_row( + &self, + row: &MySqlRow, + config: &RowProcessingConfig, + ) -> Result<ProcessedRow, Error> { + let mut row_pk: Option<String> = None; + let mut max_offset: Option<String> = None; + let mut extracted_payload: Option<Vec<u8>> = None; + + // Payload column set: only extract it plus tracking/pk columns. + // Avoids extract_column_value on every other column since the data map + // built below would be discarded anyway. + if !config.payload_col.is_empty() { + for (i, column) in row.columns().iter().enumerate() { + let name = column.name(); + if name == config.payload_col { + extracted_payload = + Some(self.extract_payload_column(row, i, config.payload_format)?); + } + if name == config.tracking_column || name == config.pk_column { + let as_string = value_as_string(&extract_column_value(row, i)?); Review Comment: value_as_string only String|Number; BOOL/NULL/JSON tracking/PK ⇒ no offset/PK for mark. **Fix**: canonical scalar encoding or open-time schema probe. -- 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]
