tusharagrahari commented on code in PR #3568: URL: https://github.com/apache/iggy/pull/3568#discussion_r3531965664
########## core/connectors/sources/mysql_source/src/lib.rs: ########## @@ -0,0 +1,1418 @@ +// 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>, +} + +/// 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") | 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(", "); + + 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()); + + 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()); + } + 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. A table's messages are emitted + // only after its rows are marked or deleted, so a failure here can never + // leave a table marked-but-unpublished. `mark_or_delete_processed_rows` + // retries transient errors; a permanent failure isolates that one table + // (its rows stay in MySQL and are retried next cycle) while the others + // still flow. + let mut messages = Vec::new(); + let mut state_updates: Vec<(String, String)> = Vec::new(); + let mut total_processed: u64 = 0; + + for mut batch in batches { + 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); + } + state.last_poll_time = Utc::now(); + } + + 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.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 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)?; Review Comment: using per-row skip can cause dataloss as we are updating the offset to the last polled row, for now i have emphasised on making the error message more detailed so the operator can pin pointedly know the exact reason and fix it, documented the same in the readme. As for the DLQ can take that as a fresh issue after going through the discussions around 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]
