ryerraguntla commented on code in PR #3568: URL: https://github.com/apache/iggy/pull/3568#discussion_r3603866423
########## 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)?; + Review Comment: 748-791,427 — hot path: per-row serde_json Map + SELECT *. **Fix**: column projection + direct serialize (v1 acceptable, matches postgres shape). -- 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]
