tusharagrahari commented on code in PR #3568: URL: https://github.com/apache/iggy/pull/3568#discussion_r4297516046
########## core/connectors/sources/mysql_source/src/lib.rs: ########## @@ -0,0 +1,4071 @@ +// 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::SourceBatchResult, 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::cmp::Ordering as CmpOrdering; +use std::collections::{HashMap, HashSet}; +use std::str::FromStr; +use std::sync::atomic::{AtomicBool, AtomicU32, AtomicU64, 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); + +/// A disabled table is skipped before anything is logged, so without a reminder the +/// only evidence a connector went dark is the single line written when it happened. +/// At the default 10s poll interval this repeats it every ten minutes. +const DISABLED_TABLES_LOG_INTERVAL: u64 = 60; + +/// How many consecutive polls a table may end in `CursorRegression` before the repeat +/// stops reading as transient. Nothing is published in the meantime, so past this it +/// is a stall the operator has to see. +const CURSOR_REGRESSION_ERROR_THRESHOLD: u32 = 3; + +/// Upper bound on applying a batch's staged mark/delete after its Ack. The SDK waits +/// `BATCH_RESULT_TIMEOUT` (30s) for the result, so this has to finish well inside it. +const ACK_BATCH_TIMEOUT: Duration = Duration::from_secs(10); + +/// MySQL rejects a prepared statement with more placeholders than this (error 1390), and +/// the mark/delete binds one per row of the batch. +const MAX_CLEANUP_BATCH_SIZE: u32 = 65_535; + +/// How many mark/delete statements in a row may fail before the connector stops. The +/// operations stay in the persisted state and run again on restart. +const MAX_CONSECUTIVE_CLEANUP_FAILURES: u32 = 3; + +#[derive(Debug)] +pub struct MySqlSource { + pub id: u32, + pool: Option<Pool<MySql>>, + config: MySqlSourceConfig, + state: Mutex<State>, + pending_batch: Mutex<Option<PendingBatch>>, + consecutive_cleanup_failures: AtomicU32, + verified_cleanup_keys: Mutex<HashSet<String>>, + verbose: bool, + retry_delay: Duration, + poll_interval: Duration, + last_batch_full: AtomicBool, + poisoned_tables: Mutex<HashSet<String>>, + tracking_kinds: Mutex<HashMap<String, OffsetKind>>, + polls: AtomicU64, + cursor_regressions: Mutex<HashMap<String, u32>>, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +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, + tracking_value: Option<String>, + row_pk: Option<PkValue>, +} + +/// Why a table's batch could not be produced. +enum FetchError { + Ordering(String), + /// The batch opened below the cursor it was told to resume from, so the query's + /// filter and the cursor disagree about what "past the offset" means. Retried + /// rather than latched, because unlike a within-batch decrease it depends on the + /// rows as much as on the query: a `custom_query` whose WHERE admits rows the + /// cursor does not gate (an OR branch is the usual one) only fails while such a + /// row is there, so it can clear without a config change. + CursorRegression(String), + MissingTrackingColumn(String), + UnusableTrackingValue(String), + StaleOffsetKind(String), + Other(Error), +} + +impl From<Error> for FetchError { + fn from(error: Error) -> Self { + FetchError::Other(error) + } +} + +/// How MySQL ordered a tracking column, which offsets alone cannot reveal: they are +/// carried as strings, so an `INT` 10 and a `VARCHAR` "10" are indistinguishable by +/// content while the server orders them differently. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum OffsetKind { + Numeric, + Lexical, + Unknown, +} + +/// Maps an `information_schema.columns.data_type` to the order MySQL sorts it in. +/// `decimal` belongs with the numbers even though `extract_column_value` renders it +/// as a string to preserve precision. +fn offset_kind_for_data_type(data_type: &str) -> OffsetKind { + const NUMERIC: &[&str] = &[ + "tinyint", + "smallint", + "mediumint", + "int", + "integer", + "bigint", + "decimal", + "numeric", + "float", + "double", + "real", + "bit", + "year", + ]; + if NUMERIC + .iter() + .any(|numeric| data_type.eq_ignore_ascii_case(numeric)) + { + OffsetKind::Numeric + } else { + OffsetKind::Lexical + } +} + +/// Whether MySQL orders the column by its collation. The other lexically compared +/// types are the temporal ones, which render fixed-width and so compare the same way +/// as text no matter what the collation is. +fn is_collation_ordered(data_type: &str) -> bool { + const TEXT: &[&str] = &[ + "char", + "varchar", + "tinytext", + "text", + "mediumtext", + "longtext", + ]; + TEXT.iter().any(|text| data_type.eq_ignore_ascii_case(text)) +} + +fn offset_kind_for_type_name(type_name: &str) -> OffsetKind { + match type_name { + "TINYINT" | "SMALLINT" | "MEDIUMINT" | "INT" | "BIGINT" | "TINYINT UNSIGNED" + | "SMALLINT UNSIGNED" | "MEDIUMINT UNSIGNED" | "INT UNSIGNED" | "BIGINT UNSIGNED" + | "FLOAT" | "DOUBLE" | "DECIMAL" | "BIT" | "YEAR" => OffsetKind::Numeric, + "DATE" | "TIME" | "DATETIME" | "TIMESTAMP" | "CHAR" | "VARCHAR" | "TINYTEXT" | "TEXT" + | "MEDIUMTEXT" | "LONGTEXT" => OffsetKind::Lexical, + _ => OffsetKind::Unknown, + } +} + +fn unusable_tracking_type(data_type: &str) -> Option<&'static str> { + const BINARY: &[&str] = &[ + "binary", + "varbinary", + "tinyblob", + "blob", + "mediumblob", + "longblob", + ]; + + if data_type.eq_ignore_ascii_case("json") { + return Some("it yields no ordered scalar cursor"); + } + // Judged on `data_type` rather than on the `tinyint(1)` display width that spells + // BOOLEAN, because that width is not reliably reported: MySQL 8.0.19 and later drop + // it for unsigned columns while MariaDB keeps it. + if data_type.eq_ignore_ascii_case("tinyint") { + return Some("it has too few distinct values to be unique per row"); + } + if BINARY + .iter() + .any(|binary| data_type.eq_ignore_ascii_case(binary)) + { + return Some("binary values are compared as base64 text, which does not keep byte order"); + } + if data_type.eq_ignore_ascii_case("enum") || data_type.eq_ignore_ascii_case("set") { + return Some("MySQL orders these by declaration index, not by label"); + } + None +} + +/// What a `custom_query` may begin with. `TABLE` is MySQL 8.0.19's table statement, +/// and `(` opens a parenthesized query such as `(SELECT ...) UNION (SELECT ...)`. +const READ_STATEMENT_STARTS: &[&str] = &["SELECT", "WITH", "TABLE", "("]; + +/// The token that decides which statement a query is, skipping leading comments in +/// all three MySQL spellings. Returns a single character when the statement does not +/// open with a word, so a parenthesized query is still recognisable. +fn leading_statement_token(query: &str) -> Option<String> { + let mut rest = query.trim_start(); + loop { + rest = if let Some(after) = rest.strip_prefix("/*") { + after.split_once("*/")?.1.trim_start() + } else if rest.starts_with("--") || rest.starts_with('#') { + rest.split_once('\n')?.1.trim_start() + } else { + break; + }; + } + + let word: String = rest + .chars() + .take_while(|c| c.is_ascii_alphanumeric() || *c == '_') + .collect(); + if word.is_empty() { + return rest.chars().next().map(String::from); + } + Some(word.to_uppercase()) +} + +/// The order the tracking column sorts in according to the result set that just +/// came back, or `None` when the query did not project the column under that name. +fn observed_tracking_kind(row: &MySqlRow, tracking_column: &str) -> Option<OffsetKind> { + row.columns() + .iter() + .find(|column| column.name().eq_ignore_ascii_case(tracking_column)) + .map(|column| offset_kind_for_type_name(column.type_info().name())) + .filter(|kind| *kind != OffsetKind::Unknown) +} + +/// `Clone` so the passing case, which is every batch, walks the names without collecting. +fn ensure_tracking_column_projected<'a>( + projected: impl IntoIterator<Item = &'a str> + Clone, + tracking_column: &str, +) -> Result<(), FetchError> { + if projected + .clone() + .into_iter() + .any(|column| column.eq_ignore_ascii_case(tracking_column)) + { + return Ok(()); + } + Err(FetchError::MissingTrackingColumn(format!( + "the result set has no column '{tracking_column}', it returned [{}]", + projected.into_iter().collect::<Vec<_>>().join(", ") + ))) +} + +fn unusable_tracking_value(tracking_column: &str, pk: &str) -> FetchError { + FetchError::UnusableTrackingValue(format!( + "row (pk {pk}) has no usable value for tracking_column '{tracking_column}': it is NULL \ + or not an ordered scalar (tinyint(1), JSON)" + )) +} + +/// Verifies rows arrive non-decreasing in the tracking column, which is what makes +/// "cursor = last row's value" correct. The built polling query orders ascending by +/// construction, so this only ever fires on a `custom_query`, where the operator +/// owns the `ORDER BY`. +struct OrderingGuard<'a> { + tracking_column: &'a str, + kind: OffsetKind, + last: Option<String>, + /// True while `last` still holds the cursor carried in from the previous poll + /// rather than a value from this batch. Distinguishes the two violations, and + /// keeps an empty batch from re-persisting the offset it started with. + at_cursor: bool, +} + +enum OrderingViolation { + WithinBatch(String), + BelowCursor(String), +} + +impl<'a> OrderingGuard<'a> { + fn new(tracking_column: &'a str, kind: OffsetKind, cursor: Option<String>) -> Self { + OrderingGuard { + tracking_column, + kind, + at_cursor: cursor.is_some(), + last: cursor, + } + } + + fn accept(&mut self, value: String) -> Result<(), OrderingViolation> { + if let Some(previous) = &self.last + && offset_decreased(previous, &value, self.kind) + { + let column = self.tracking_column; + return Err(if self.at_cursor { + OrderingViolation::BelowCursor(format!( + "row '{value}' is below the stored offset '{previous}' for tracking_column \ + '{column}'. Filter the query with `{column} > $offset`" + )) + } else { + OrderingViolation::WithinBatch(format!( + "rows arrived out of order by tracking_column '{column}': '{previous}' \ + preceded '{value}'. Order the query ascending by '{column}'" + )) + }); + } + self.last = Some(value); + self.at_cursor = false; + Ok(()) + } + + /// `None` when no row was accepted, so an empty batch leaves the stored offset + /// untouched rather than rewriting the cursor it was seeded with. + fn into_cursor(self) -> Option<String> { + if self.at_cursor { None } else { self.last } + } +} + +struct TableBatch { + table: String, + messages: Vec<ProducedMessage>, + processed_ids: Vec<PkValue>, + 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") => PayloadFormat::JsonDirect, + _ => PayloadFormat::Json, + } + } +} + +/// A primary-key value in a form the mark/delete `WHERE pk IN (...)` can bind +/// faithfully. `Text` covers ints/strings/dates/decimals (MySQL implicitly +/// converts the bound string to the column type). `Bytes` carries raw binary +/// (BINARY/VARBINARY/BLOB) so a UUID stored as bytes matches the row instead of +/// being compared as its base64 text, which never matches. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +enum PkValue { + Text(String), + Bytes(Vec<u8>), +} + +impl PkValue { + fn as_key(&self) -> String { + match self { + PkValue::Text(text) => text.clone(), + PkValue::Bytes(bytes) => base64::engine::general_purpose::STANDARD.encode(bytes), + } + } +} + +#[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, Clone, Serialize, Deserialize)] +struct State { + tracking_offsets: HashMap<String, String>, + processed_rows: u64, + /// Mark/delete work for a batch the runtime has already sent. It is persisted with + /// the batch's offsets so a crash between the send and the cleanup reruns it. + #[serde(default)] + pending_operations: Vec<PendingOperation>, +} + +/// The rows of one table to mark or delete once their batch is acknowledged. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +struct PendingOperation { + table: String, + ids: Vec<PkValue>, + target: CleanupTarget, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +struct CleanupTarget { + key_column: String, + action: CleanupAction, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +enum CleanupAction { + Delete, + MarkProcessed { column: String }, +} + +/// The candidate state of the batch most recently handed to the runtime. Live state +/// only takes it over on Ack, so a Nack leaves offsets and rows where they were. +#[derive(Debug, Clone)] +struct PendingBatch { + state: State, + /// The runtime has delivered the batch but its cleanup has not finished yet. + acknowledged: bool, + /// Cleanup finished and the next poll must persist `state` without the operations, + /// so a restart does not run them again. + retiring: bool, +} + +struct PolledBatch { + messages: Vec<ProducedMessage>, + pending: Option<PendingBatch>, +} + +enum CleanupKeyCheck { + Unique, + NotUnique, + TableMissing, +} + +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 self.config.tables.is_empty() { + return Err(Error::InitError( + "tables must not be empty, list the table even when custom_query hardcodes it" + .to_string(), + )); + } + + if self.config.batch_size == Some(0) { + return Err(Error::InitError( + "batch_size must be greater than 0; omit it to use the default of 1000".to_string(), + )); + } + + if let Some(batch_size) = self.config.batch_size + && batch_size > MAX_CLEANUP_BATCH_SIZE + && self.cleanup_target().is_some() + { + return Err(Error::InitError(format!( + "batch_size {batch_size} exceeds {MAX_CLEANUP_BATCH_SIZE} with delete_after_read or \ + processed_column set, MySQL allows at most {MAX_CLEANUP_BATCH_SIZE} placeholders \ + in the statement that marks or deletes a batch" + ))); + } + + if self.config.batch_size == Some(1) + && self + .config + .custom_query + .as_deref() + .is_some_and(|query| query.contains("$offset")) + { + return Err(Error::InitError( + "batch_size must be greater than 1 when custom_query filters on $offset" + .to_string(), + )); + } + + if self + .config + .tracking_column + .as_deref() + .is_none_or(str::is_empty) + { + let guidance = if self.config.custom_query.is_some() { + "name the column custom_query orders by and filters with $offset" + } else { + "use a unique, monotonically increasing column" + }; + return Err(Error::InitError(format!( + "tracking_column must be set, {guidance}" + ))); + } + + 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.validate_pending_cleanup()?; + self.connect().await?; + *self.tracking_kinds.get_mut() = self.resolve_tracking_columns().await?; + self.validate_cleanup_keys().await?; + + info!( + "MySQL source connector with ID: {} opened successfully", + self.id + ); + Ok(()) + } + + async fn poll(&self) -> Result<ProducedMessages, Error> { + let schema = match self.payload_format() { + PayloadFormat::Bytea => Schema::Raw, + PayloadFormat::Text => Schema::Text, + PayloadFormat::JsonDirect | PayloadFormat::Json => Schema::Json, + }; + + if let Some(state) = self.retiring_state().await { + return Ok(ProducedMessages { + schema, + messages: Vec::new(), + state: Some(self.serialize_pending_state(&state)?), + }); + } + + // Skip the pacing delay while draining a backlog, but never while a cleanup + // retry is pending, or a failing mark/delete would spin without a pause. + let cleanup_pending = self.pending_batch.lock().await.is_some(); + if cleanup_pending || !self.last_batch_full.load(Ordering::Relaxed) { + tokio::time::sleep(self.poll_interval).await; + } + + { + let pending = self.pending_batch.lock().await; + if let Some(pending) = pending.as_ref() { + if !pending.acknowledged { + error!( + "{CONNECTOR_NAME} connector ID: {} was polled while a batch was still in flight", + self.id + ); + return Err(Error::InvalidState); + } + // An empty batch whose Ack retries the cleanup the last one left behind. + return Ok(ProducedMessages { + schema, + messages: Vec::new(), + state: None, + }); + } + } + + let polled = self.poll_tables().await?; + + let processed_rows = match &polled.pending { + Some(pending) => pending.state.processed_rows, + None => self.state.lock().await.processed_rows, + }; + if self.verbose { + info!( + "MySQL source connector ID: {} produced {} messages. Total processed: {}", + self.id, + polled.messages.len(), + processed_rows + ); + } else { + debug!( + "MySQL source connector ID: {} produced {} messages. Total processed: {}", + self.id, + polled.messages.len(), + processed_rows + ); + } + + // Idle cycles stage nothing, so return None to let the runtime skip the fsync. + let persisted_state = polled + .pending + .as_ref() + .map(|pending| self.serialize_pending_state(&pending.state)) + .transpose()?; + *self.pending_batch.lock().await = polled.pending; + + Ok(ProducedMessages { + schema, + messages: polled.messages, + state: persisted_state, + }) + } + + async fn on_batch_result(&self, result: SourceBatchResult) -> Result<(), Error> { + if result == SourceBatchResult::Nack { + // Dropping the candidate leaves the live offsets where they were, so the + // next poll reads the same rows again. Delivered batches keep their cleanup. + let mut pending = self.pending_batch.lock().await; + if pending + .as_ref() + .is_some_and(|pending| !pending.acknowledged && !pending.retiring) + { + pending.take(); + } + return Ok(()); + } + + let Some(PendingBatch { mut state, .. }) = self.pending_batch.lock().await.clone() else { + return Ok(()); + }; + + let operations = std::mem::take(&mut state.pending_operations); + let had_operations = !operations.is_empty(); + let deadline = tokio::time::Instant::now() + ACK_BATCH_TIMEOUT; + let failed_operations = self.apply_pending_operations(operations, deadline).await?; + + if failed_operations.is_empty() { + *self.state.lock().await = state.clone(); + let mut pending = self.pending_batch.lock().await; + if had_operations { + *pending = Some(PendingBatch { + state, + acknowledged: false, + retiring: true, + }); + } else { + pending.take(); + } + } else { + state.pending_operations = failed_operations; + *self.pending_batch.lock().await = Some(PendingBatch { + state, + acknowledged: true, + retiring: false, + }); + } + + Ok(()) + } + + 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)); + let state = restored_state.unwrap_or(State { + tracking_offsets: HashMap::new(), + processed_rows: 0, + pending_operations: Vec::new(), + }); + // The runtime only persists a batch's operations after sending it, so restored + // operations belong to a delivered batch and only the cleanup is left to do. + let pending_batch = (!state.pending_operations.is_empty()).then(|| PendingBatch { + state: state.clone(), + acknowledged: true, + retiring: false, + }); + MySqlSource { + id, + pool: None, + config, + state: Mutex::new(state), + pending_batch: Mutex::new(pending_batch), + consecutive_cleanup_failures: AtomicU32::new(0), + verified_cleanup_keys: Mutex::new(HashSet::new()), + verbose, + retry_delay, + poll_interval, + last_batch_full: AtomicBool::new(false), + poisoned_tables: Mutex::new(HashSet::new()), + tracking_kinds: Mutex::new(HashMap::new()), + polls: AtomicU64::new(0), + cursor_regressions: Mutex::new(HashMap::new()), + } + } + + 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) + // `format_offset_value` escapes backslashes, which under NO_BACKSLASH_ESCAPES + // would reach MySQL as literal characters and filter on the wrong cursor. + .after_connect(|connection, _| { + Box::pin(async move { + sqlx::query( + "SET SESSION sql_mode = REPLACE(@@SESSION.sql_mode, 'NO_BACKSLASH_ESCAPES', '')", + ) + .execute(&mut *connection) + .await?; + Ok(()) + }) + }) + .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 tracking_column(&self) -> &str { + self.config.tracking_column.as_deref().unwrap_or("id") + } + + /// Fail fast if a tracking column can't yield a stable, ordered scalar + /// cursor. `value_as_string` drops NULL, tinyint(1) (decoded as BOOLEAN), + /// and JSON to `None`, which silently stalls offset advancement and re-reads + /// rows every cycle. + async fn resolve_tracking_columns(&self) -> Result<HashMap<String, OffsetKind>, Error> { + // Whether the connector owns the query, and can therefore hold the schema to + // what that query will do with it. + let strict = self.config.custom_query.is_none(); + let pool = self.get_pool()?; + let tracking_column = self.tracking_column(); + let mut kinds = HashMap::with_capacity(self.config.tables.len()); + + for table in &self.config.tables { + let (schema, table_name) = table + .split_once('.') + .map(|(schema, name)| (schema.to_string(), name.to_string())) + .unwrap_or_else(|| (String::new(), table.clone())); + + let row = sqlx::query( + "SELECT is_nullable, data_type, column_type \ + FROM information_schema.columns \ + WHERE table_schema = IF(? = '', DATABASE(), ?) \ + AND table_name = ? AND column_name = ?", + ) + .bind(&schema) + .bind(&schema) + .bind(&table_name) + .bind(tracking_column) + .fetch_optional(pool) + .await + .map_err(|e| { + Error::InitError(format!( + "failed to probe schema for tracking_column '{tracking_column}' on table '{table}': {e}" + )) + })?; + + let Some(row) = row else { + // A missing columns row means either the table doesn't exist yet + // or the column is absent. Only the latter is a misconfiguration: + // a table created after the connector starts is legitimate, so + // defer validation rather than permanently disabling the source. + let table_exists = sqlx::query( + "SELECT 1 FROM information_schema.tables \ + WHERE table_schema = IF(? = '', DATABASE(), ?) \ + AND table_name = ?", + ) + .bind(&schema) + .bind(&schema) + .bind(&table_name) + .fetch_optional(pool) + .await + .map_err(|e| { + Error::InitError(format!("failed to probe existence of table '{table}': {e}")) + })? + .is_some(); + + if !table_exists { + warn!( + "table '{table}' does not exist yet; deferring tracking_column \ + '{tracking_column}' validation until the table is created" + ); + kinds.insert(table.clone(), OffsetKind::Unknown); + continue; + } + + if strict { + return Err(Error::InitError(format!( + "tracking_column '{tracking_column}' not found on table '{table}'" + ))); + } + + warn!( + "tracking_column '{tracking_column}' is not a column of table '{table}'; \ + assuming custom_query projects it. Its comparison order is resolved from \ + the first result set instead." + ); + kinds.insert(table.clone(), OffsetKind::Unknown); + continue; + }; + + // Read by position: MySQL returns information_schema column names + // uppercased (IS_NULLABLE, ...), so try_get by lowercase name misses. + let is_nullable: String = row.try_get(0).map_err(|e| { + Error::InitError(format!( + "failed to read is_nullable for table '{table}': {e}" + )) + })?; + let data_type: String = row.try_get(1).map_err(|e| { + Error::InitError(format!("failed to read data_type for table '{table}': {e}")) + })?; + let column_type: String = row.try_get(2).map_err(|e| { + Error::InitError(format!( + "failed to read column_type for table '{table}': {e}" + )) + })?; + + // The column's type survives a plain projection, so it is worth rejecting + // whatever the query looks like. A custom_query aliasing an expression onto + // a name the table also uses is the one false positive, hence the hint. + if let Some(reason) = unusable_tracking_type(&data_type) { + let aliasing_hint = if self.config.custom_query.is_some() { + " If custom_query aliases a different expression to this name, rename the \ + alias so it does not collide with the column." + } else { + "" + }; + return Err(Error::InitError(format!( + "tracking_column '{tracking_column}' on table '{table}' has type \ + '{column_type}': {reason}. Use an integer, timestamp, or string \ + column.{aliasing_hint}" + ))); + } + + if is_nullable.eq_ignore_ascii_case("YES") { + if strict { + return Err(Error::InitError(format!( + "tracking_column '{tracking_column}' on table '{table}' is nullable; \ + a NULL value stalls offset tracking. Declare it NOT NULL." + ))); + } + warn!( + "tracking_column '{tracking_column}' on table '{table}' is nullable; a row \ + whose value is NULL cannot produce a cursor and will fail that table's batch. \ + Have custom_query exclude those rows, or declare the column NOT NULL." + ); + } + + if is_collation_ordered(&data_type) { + warn!( + "tracking_column '{tracking_column}' on table '{table}' has type \ + '{column_type}': values must be unique and increase in its collation order \ + (zero-pad numbers, prefer a _bin collation)" + ); + } + + kinds.insert(table.clone(), offset_kind_for_data_type(&data_type)); + } + Ok(kinds) + } + + 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 serialize_pending_state(&self, state: &State) -> Result<ConnectorState, Error> { + self.serialize_state(state).ok_or_else(|| { + Error::Serialization(format!("failed to serialize {CONNECTOR_NAME} state")) + }) + } + + async fn retiring_state(&self) -> Option<State> { + self.pending_batch + .lock() + .await + .as_ref() + .filter(|pending| pending.retiring) + .map(|pending| pending.state.clone()) + } + + fn primary_key_column(&self) -> &str { + self.config + .primary_key_column + .as_deref() + .unwrap_or(self.tracking_column()) + } + + fn cleanup_target(&self) -> Option<CleanupTarget> { + let action = if self.config.delete_after_read.unwrap_or(false) { + CleanupAction::Delete + } else if let Some(column) = &self.config.processed_column { + CleanupAction::MarkProcessed { + column: column.clone(), + } + } else { + return None; + }; + Some(CleanupTarget { + key_column: self.primary_key_column().to_string(), + action, + }) + } + + /// Restored operations were staged under the config of the run that wrote them. + /// Running them under a different key or action would touch rows nobody read. + fn validate_pending_cleanup(&mut self) -> Result<(), Error> { + let configured = self.cleanup_target(); + for operation in &self.state.get_mut().pending_operations { + if configured.as_ref() != Some(&operation.target) { + return Err(Error::InitError(format!( + "the cleanup configuration changed while cleanup of {} delivered row(s) in \ + table '{}' is pending; restore primary_key_column '{}' and action {:?}", + operation.ids.len(), + operation.table, + operation.target.key_column, + operation.target.action + ))); + } + } + Ok(()) + } + + /// `WHERE key IN (...)` on a key that repeats marks or deletes every row sharing a + /// value, including ones that were never read. Tables that don't exist yet are + /// checked by the first poll that stages cleanup for them. + async fn validate_cleanup_keys(&mut self) -> Result<(), Error> { + let Some(target) = self.cleanup_target() else { + return Ok(()); + }; + let pool = self.get_pool()?; + let mut verified = HashSet::with_capacity(self.config.tables.len()); + for table in &self.config.tables { + let check = cleanup_key_check(pool, table, &target.key_column) + .await + .map_err(|e| { + Error::InitError(format!( + "failed to check primary_key_column '{}' on table '{table}': {e}", + target.key_column + )) + })?; + match check { + CleanupKeyCheck::Unique => { + verified.insert(table.clone()); + } + CleanupKeyCheck::NotUnique => { + return Err(Error::InitError(not_unique_key_message( + table, + &target.key_column, + ))); + } + CleanupKeyCheck::TableMissing => {} + } + } + *self.verified_cleanup_keys.get_mut() = verified; + Ok(()) + } + + async fn cleanup_key_is_unique( + &self, + pool: &Pool<MySql>, + table: &str, + key_column: &str, + ) -> Result<bool, Error> { + if self.verified_cleanup_keys.lock().await.contains(table) { + return Ok(true); + } + let check = cleanup_key_check(pool, table, key_column) + .await + .map_err(|e| { + Error::Connection(format!( + "failed to check primary_key_column '{key_column}' on table '{table}': {e}" + )) + })?; + match check { + CleanupKeyCheck::Unique => { + self.verified_cleanup_keys + .lock() + .await + .insert(table.to_string()); + Ok(true) + } + CleanupKeyCheck::NotUnique => Ok(false), + CleanupKeyCheck::TableMissing => Err(Error::Connection(format!( + "table '{table}' disappeared before its primary_key_column could be checked" + ))), + } + } + + 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 + } + } + + /// Best current knowledge of how a table's tracking values compare. + async fn tracking_kind(&self, table: &str) -> OffsetKind { + self.tracking_kinds + .lock() + .await + .get(table) + .copied() + .unwrap_or(OffsetKind::Unknown) + } + + /// Whether the stored cursor is what limits the rows a poll can see. The built + /// polling query always filters on it; a `custom_query` does so only if it + /// interpolates `$offset`. + fn cursor_gates_query(&self) -> bool { + match &self.config.custom_query { + Some(query) => query.contains("$offset"), + None => true, + } + } + + /// The offset a poll actually resumes from: the stored cursor, or the configured + /// starting point before anything has been stored. + fn effective_offset(&self, last_offset: &Option<String>) -> Option<String> { + last_offset + .clone() + .or_else(|| self.config.initial_offset.clone()) + } + + fn substitute_query_params( + &self, + query: &str, + table: &str, + last_offset: &Option<String>, + batch_size: u32, + kind: OffsetKind, + ) -> Result<String, Error> { + let offset_value = self.effective_offset(last_offset).unwrap_or_default(); + let offset = format_offset_value(&offset_value, kind); + let quoted_table = quote_qualified_identifier(table)?; + let limit = batch_size.to_string(); + let now = Utc::now(); + let now_unix = now.timestamp().to_string(); + let now = now.to_rfc3339(); + let replacements = [ + ("$table", quoted_table.as_str()), + ("$offset", offset.as_str()), + ("$limit", limit.as_str()), + ("$now_unix", now_unix.as_str()), + ("$now", now.as_str()), + ]; + + Ok(substitute_query_tokens(query, &replacements)) + } + + /// The query is substituted as text and re-executed once per configured table on + /// every poll for the life of the process, so a shape that cannot terminate does + /// not fail, it repeats. What can be judged from the text alone is judged here, + /// at startup, rather than left to the first poll. + fn validate_custom_query(&self, query: &str) -> Result<(), Error> { + let tracking_column = self.tracking_column(); + + match leading_statement_token(query) { + None => { + return Err(Error::InitError( + "custom_query has no statement: it is empty, or nothing but comments" + .to_string(), + )); + } + Some(token) if !READ_STATEMENT_STARTS.contains(&token.as_str()) => { + return Err(Error::InitError(format!( + "custom_query must be a read starting with SELECT, WITH or TABLE, found \ + '{token}'" + ))); + } + Some(_) => {} + } + + // Collapsed so `ORDER\nBY` and `ORDER BY` read the same as `ORDER BY`; the + // check hard-fails, so a formatting difference must not decide it. + let normalized = query + .split_whitespace() + .collect::<Vec<_>>() + .join(" ") + .to_uppercase(); + + if query.contains("$offset") { + if !normalized.contains("ORDER BY") { + return Err(Error::InitError(format!( + "custom_query filters on $offset but has no ORDER BY, add \ + `ORDER BY {tracking_column}`" + ))); + } + } else if !self.config.delete_after_read.unwrap_or(false) + && self.config.processed_column.is_none() + { + return Err(Error::InitError(format!( + "custom_query must filter with $offset (`WHERE {tracking_column} > $offset`), or \ + set delete_after_read or processed_column" + ))); + } + + if query.contains("$table") && self.config.tables.is_empty() { + return Err(Error::InitError( + "custom_query uses $table but no tables are configured".to_string(), + )); + } + // Without $table the same SQL is executed once per table, so identical rows + // would be emitted repeatedly, each carrying a different table's metadata, + // message ID, and tracking offset. Table identity is not recoverable from a + // static query, so reject the combination instead of guessing one. + if !query.contains("$table") && self.config.tables.len() > 1 { + return Err(Error::InitError(format!( + "custom_query has no $table placeholder but {} tables are configured ({}), use \ + $table in the query or configure a single table", + self.config.tables.len(), + self.config.tables.join(", ") + ))); + } + Ok(()) + } + + fn build_polling_query( + &self, + table: &str, + tracking_column: &str, + last_offset: &Option<String>, + batch_size: u32, + kind: OffsetKind, + ) -> 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) = self.effective_offset(last_offset) { + conditions.push(format!( + "{quoted_tracking} > {}", + format_offset_value(&offset, kind) + )); + } + + 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 apply_pending_operations( + &self, + operations: Vec<PendingOperation>, + deadline: tokio::time::Instant, + ) -> Result<Vec<PendingOperation>, Error> { + let mut failed_operations = Vec::new(); + let mut operations = operations.into_iter(); + while let Some(operation) = operations.next() { + if tokio::time::Instant::now() >= deadline { + failed_operations.push(operation); + failed_operations.extend(operations); + break; + } + let result = tokio::time::timeout_at(deadline, async { + let pool = self.get_pool()?; + self.mark_or_delete_processed_rows(pool, &operation).await + }) + .await; + match result { + Ok(Ok(())) => { + self.consecutive_cleanup_failures + .store(0, Ordering::Relaxed); + } + Ok(Err(error)) => { + self.record_cleanup_failure(&operation, &error)?; + failed_operations.push(operation); + } + Err(_) => { + self.record_cleanup_failure( + &operation, + "cleanup exceeded the 10s budget for one Ack", + )?; + failed_operations.push(operation); + } + } + } + Ok(failed_operations) + } + + fn record_cleanup_failure( + &self, + operation: &PendingOperation, + error: impl std::fmt::Display, + ) -> Result<(), Error> { + let consecutive_failures = self + .consecutive_cleanup_failures + .fetch_add(1, Ordering::Relaxed) + .saturating_add(1); + error!( + "Failed to mark or delete {} delivered row(s) in table '{}' for {CONNECTOR_NAME} \ + connector ID: {}, retrying on the next poll ({consecutive_failures} consecutive \ + failures): {error}", + operation.ids.len(), + operation.table, + self.id + ); + if consecutive_failures >= MAX_CONSECUTIVE_CLEANUP_FAILURES { + return Err(Error::Connection(format!( + "stopping {CONNECTOR_NAME} connector ID {} after {consecutive_failures} \ + consecutive mark/delete failures; the cleanup stays in the saved state and \ + runs again on restart", + self.id + ))); + } + Ok(()) + } + + async fn mark_or_delete_processed_rows( + &self, + pool: &Pool<MySql>, + operation: &PendingOperation, + ) -> Result<(), Error> { + let PendingOperation { table, ids, target } = operation; + if ids.is_empty() { + return Ok(()); + } + + let quoted_table = quote_qualified_identifier(table)?; + let quoted_pk = quote_identifier(&target.key_column)?; + let placeholders = vec!["?"; ids.len()].join(", "); + + let query = match &target.action { + CleanupAction::Delete => { + 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 ({placeholders})") Review Comment: fixed. Instead of saving a tracking value for every row, each pending cleanup now saves the batch's highest tracking value, and the cleanup only touches rows at or below it: DELETE FROM t WHERE pk IN (...) AND tracking_column <= <batch max> A replacement row always gets a newer tracking value, so the cleanup leaves it alone. This also avoids exact-equality matching on DATETIME and DECIMAL values, and the state stays small. Since this needs the tracking column to exist in the table, a computed or aliased tracking column is now rejected when delete_after_read or processed_column is set. The README covers it. Added integration tests that replay a pending cleanup after a key is reused, for both the built-in query and a custom_query. -- 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]
