tusharagrahari commented on code in PR #3568: URL: https://github.com/apache/iggy/pull/3568#discussion_r3482894567
########## core/connectors/sources/mysql_source/src/lib.rs: ########## @@ -0,0 +1,1358 @@ +// 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::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"; + +#[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, +} + +#[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>, +} + +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") | Some("jsonb") | Some("jsonb_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 { + last_poll_time: DateTime<Utc>, + 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(), + )); + } + self.connect().await?; + + info!( + "MySQL source connector with ID: {} opened successfully", + self.id + ); + Ok(()) + } + + async fn poll(&self) -> Result<ProducedMessages, Error> { + let poll_interval = self.poll_interval; + tokio::time::sleep(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, + }; + + let persisted_state = 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 { + last_poll_time: Utc::now(), + tracking_offsets: HashMap::new(), + processed_rows: 0, + })), + verbose, + retry_delay, + poll_interval, + } + } + + 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> { + match format { + PayloadFormat::Bytea => { + let bytes: Option<Vec<u8>> = row + .try_get(column_index) + .map_err(|_| Error::InvalidRecord)?; + Ok(bytes.unwrap_or_default()) + } + PayloadFormat::Text => { + let text: Option<String> = row + .try_get(column_index) + .map_err(|_| Error::InvalidRecord)?; + Ok(text.unwrap_or_default().into_bytes()) + } + PayloadFormat::JsonDirect => { + let json_value: Option<serde_json::Value> = row + .try_get(column_index) + .map_err(|_| Error::InvalidRecord)?; + simd_json::to_vec(&json_value.unwrap_or(serde_json::Value::Null)) + .map_err(|_| Error::InvalidRecord) + } + 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 now = Utc::now(); + + query + .replace("$table", table) + .replace("$offset", &offset_value) + .replace("$limit", &batch_size.to_string()) + .replace("$now", &now.to_rfc3339()) + .replace("$now_unix", &now.timestamp().to_string()) + } + + 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.clone() + } else { + format!( + "'{}'", + id.replace('\\', "\\\\") + .replace('\'', "''") + .replace('\0', "") + ) + } + }) + .collect::<Vec<_>>() + .join(", "); + + if self.config.delete_after_read.unwrap_or(false) { + let delete_query = + format!("DELETE FROM {quoted_table} WHERE {quoted_pk} IN ({ids_list})"); + + if self.verbose { + info!("Deleting {} processed rows from '{table}'", ids.len()); + } else { + debug!("Deleting {} processed rows from '{table}'", ids.len()); + } + + sqlx::query(sqlx::AssertSqlSafe(delete_query)) + .execute(pool) + .await + .map_err(|e| { + error!("Failed to delete processed rows: {e}"); + Error::InvalidRecord + })?; + } else if let Some(processed_col) = &self.config.processed_column { + let quoted_processed = quote_identifier(processed_col)?; + let update_query = format!( + "UPDATE {quoted_table} SET {quoted_processed} = TRUE WHERE {quoted_pk} IN ({ids_list})" + ); + + if self.verbose { + info!("Marking {} rows as processed in '{table}'", ids.len()); + } else { + debug!("Marking {} rows as processed in '{table}'", ids.len()); + } + + sqlx::query(sqlx::AssertSqlSafe(update_query)) + .execute(pool) + .await + .map_err(|e| { + error!("Failed to mark rows as processed: {e}"); + Error::InvalidRecord + })?; + } + + Ok(()) + } + + async fn poll_tables(&self) -> Result<Vec<ProducedMessage>, Error> { + let pool = self.get_pool()?; + let mut messages = Vec::new(); + + 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), + }; + + // Collect state updates to apply after processing + let mut state_updates: Vec<(String, String)> = Vec::new(); + let mut total_processed: u64 = 0; + + for table in &self.config.tables { + let table_config = RowProcessingConfig { + table, + ..row_config + }; + + // Get last offset with minimal lock time + let last_offset = { + let state = self.state.lock().await; + state.tracking_offsets.get(table).cloned() + }; + + let query = if let Some(custom_query) = &self.config.custom_query { + self.validate_custom_query(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 max_offset: Option<String> = None; + let mut processed_ids: Vec<String> = Vec::new(); + + let mut count_per_table = 0; + for row in rows { + let processed = self.process_row(&row, &table_config)?; + + if let Some(pk) = processed.row_pk { + processed_ids.push(pk); + } + if let Some(offset) = processed.max_offset { + max_offset = Some(offset); + } + + messages.push(processed.message); + count_per_table += 1; + total_processed += 1; + } + + // Database I/O without holding the lock + if !processed_ids.is_empty() { + self.mark_or_delete_processed_rows(pool, table, pk_column, &processed_ids) + .await?; + } Review Comment: valid concern, I hit this too. The delete-before-send window exists because the SDK has no post-send ack primitive, not a conscious design choice. I modeled it on postgres_source as the reference, but that's context not a justification. Proper fix needs a post-send hook on the Source trait — a SDK-level change. Happy to open a separate issue for it if useful. -- 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]
