lsabi commented on code in PR #3851: URL: https://github.com/apache/iggy/pull/3851#discussion_r3789014420
########## core/connectors/sources/mssql_source/src/lib.rs: ########## @@ -0,0 +1,2103 @@ +// 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 regex::Regex; +use async_trait::async_trait; +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 tiberius::{Client, Config, AuthMethod, Query, Row}; +use std::collections::HashMap; +use std::str::FromStr; +use std::time::Duration; +use tokio::sync::Mutex; +use tokio::net::TcpStream; +use tokio_util::compat::{Compat, TokioAsyncWriteCompatExt}; +use tracing::{debug, error, info, warn}; +use uuid::Uuid; + + +source_connector!(MSSQLSource); + +const DEFAULT_MAX_RETRIES: u32 = 3; +const DEFAULT_RETRY_DELAY: &str = "1s"; + + +// Define a new type for the client with the generic type incorporated +pub type DBClient = Client<Compat<TcpStream>>; + +pub(crate) type Lsn = [u8; 10]; + +/// Convert LSN parameter into a binary string +pub(crate) fn lsn_to_hex(lsn: &Lsn) -> String { + + let mut hex_string = String::with_capacity(lsn.len() * 2); + // Append the leading 0s and push a hex upper case representation + hex_string.push_str("0x"); + for byte in lsn.into_iter() { + hex_string.push_str(&format!("{:02X}", byte)); + } + hex_string +} + + +#[derive(Debug)] +pub struct MSSQLSource { + pub id: u32, + client: Mutex<Option<DBClient>>, + config: MSSQLSourceConfig, + state: Mutex<State>, + verbose: bool, + retry_delay: Duration, + poll_interval: Duration, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct MSSQLSourceConfig { + // Example jdbc:sqlserver://localhost;encrypt=true;user=MyUserName;password=<password>; + // See https://learn.microsoft.com/en-us/sql/connect/jdbc/building-the-connection-url?view=sql-server-ver15 + #[serde(serialize_with = "iggy_common::serde_secret::serialize_secret")] + pub connection_string: SecretString, + pub database: String, + pub mode: String, + 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 enable_cdc: Option<bool>, + pub custom_query: Option<String>, + pub snake_case_columns: Option<bool>, + pub include_metadata: Option<bool>, + pub capture_operations: Option<Vec<String>>, + 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>, + pub cdc_schema: Option<String>, + pub cdc_role: Option<String>, + pub capture_table_columns: HashMap<String, String> +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] +pub enum PayloadFormat { + #[default] + Json, + Bytea, + Text, + JsonDirect, +} + +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)] +struct State { + last_poll_lsn: Option<Lsn>, + start_lsn: Option<Lsn>, + end_lsn: Option<Lsn>, + last_poll_time: DateTime<Utc>, + tracking_offsets: HashMap<String, String>, + processed_rows: u64, +} + +#[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, +} + +struct ProcessedRow { + message: ProducedMessage, + max_offset: Option<String>, + row_pk: Option<String>, +} + +const CONNECTOR_NAME: &str = "MSSQL source"; + +impl MSSQLSource { + pub fn new(id: u32, config: MSSQLSourceConfig, 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)); + MSSQLSource { + id, + client: Mutex::new(None), + config, + state: Mutex::new(restored_state.unwrap_or(State { + last_poll_lsn: None, + start_lsn: None, + end_lsn: None, + last_poll_time: Utc::now(), + tracking_offsets: HashMap::new(), + processed_rows: 0 + })), + verbose, + retry_delay, + poll_interval + } + } + + fn serialize_state(&self, state: &State) -> Option<ConnectorState> { + ConnectorState::serialize(state, CONNECTOR_NAME, self.id) + } +} + +#[async_trait] +impl Source for MSSQLSource { + async fn open(&mut self) -> Result<(), Error> { + info!( + "Opening MSSQL source connector with ID: {}. Mode: {}, Tables: {:?}", + self.id, self.config.mode, self.config.tables + ); + + self.connect().await?; + + match self.config.mode.as_str() { + "cdc" => { + self.setup_cdc().await?; + info!( + "MSSQL CDC mode enabled for connector ID: {}", + self.id + ); + } + "polling" => { + info!( + "MSSQL polling mode not enabled for connector ID: {}", + self.id + ); + info!("Poll interval: {:?}", self.poll_interval); + } + _ => { + return Err(Error::InitError(format!( + "Invalid mode '{}'. Supported modes: 'polling', 'cdc'", + self.config.mode + ))); + } + } + + info!( + "MSSQL 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 = match self.config.mode.as_str() { + // "polling" => self.poll_tables().await?, + "cdc" => self.poll_cdc().await?, + _ => { + error!("Invalid mode: {}", self.config.mode); + return Err(Error::InvalidConfig); + } + }; + + let state = self.state.lock().await; + if self.verbose { + info!( + "MSSQL source connector ID: {} produced {} messages. Total processed: {}", + self.id, + messages.len(), + state.processed_rows + ); + } else { + debug!( + "MSSQL 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> { + let mut lock = self.client.lock().await; + + let mut client = lock.take(); + + let _ = match client { + Some(x) => x.close().await, + None => Ok(()) + }; + info!( + "MSSQL connection closed for connector ID: {}", + self.id + ); + + let state = self.state.lock().await; + info!( + "MSSQL source connector ID: {} closed. Total rows processed: {}", + self.id, state.processed_rows + ); + Ok(()) + } +} + +impl MSSQLSource { + async fn connect(&mut self) -> Result<(), Error> { + + let redacted = redact_connection_string(self.config.connection_string.expose_secret()); + + info!("Connecting to MSSQL with 1 connection: {redacted}"); + let config = Config::from_jdbc_string(self.config.connection_string.expose_secret()) + .map_err(|e| Error::InitError(format!("Invalid connection string configuration: {e}")))?; + + let tcp = TcpStream::connect(config.get_addr()) + .await + .map_err(|e| Error::InitError( + format!("Cannot connect to TCP stream: {e}")))?; + + if tcp.set_nodelay(true).is_ok() { } else { + warn!("Cannot set no delay on the TCP socket!"); + } + + // To be able to use Tokio's tcp, we're using the `compat_write` from + // the `TokioAsyncWriteCompatExt` to get a stream compatible with the + // traits from the `futures` crate. + let mut client = Client::connect( + config, + tcp.compat_write() + ).await + .map_err(|e| Error::InitError(format!("Failed to connect to MSSQL server: {e}")))?; + + client.query("SELECT 1", &[]) + .await + .map_err(|e| Error::InitError(format!("Database connectivity test failed: {e}")))?; + + { + let mut lock = self.client.lock().await; + *lock = Some(client); + } + + info!("Connected to MSSQL database"); + Ok(()) + } + + async fn setup_cdc(&mut self) -> Result<(), Error> { + + let database = self.config.database.clone(); + let cdc_tables = if self.config.tables.is_empty() { + let res = self.query(Query::new("SELECT name FROM sys.tables WHERE is_tracked_by_cdc = 1;")) + .await?; + res + .into_iter() + .map(|r| r.get::<&str, _>("name").unwrap().to_string()) + .collect() + } else { + self.config.tables.clone() + }.join(","); + + let mut query = Query::new("SELECT name, is_tracked_by_cdc FROM sys.tables WHERE is_tracked_by_cdc = 0 AND name IN (@P1);"); + query.bind(&cdc_tables); + let non_capture_tables = self.query(query).await?; + + if !non_capture_tables.is_empty() { + warn!("Not all desired tables are monitored by CDC!"); + if !self.config.enable_cdc.unwrap_or(false) { + error!("Update config to enable_cdc=true in order to monitor them!"); + return Ok(()); + } + warn!("Adding tables to the CDC monitor"); + let schema = self.config.cdc_schema.clone().unwrap_or("dbo".to_string()); + for row in non_capture_tables { + let table = row.get("name").unwrap_or(""); + let cdc_role = match &self.config.cdc_role { + Some(x) => format!("N'{x}'"), + None => "NULL".to_string() + }; + let query = format!(r#" + EXEC sys.sp_cdc_enable_table + @source_schema = N'{schema}', + @source_name = N'{table}', + @role_name = {cdc_role}"#); + + self.query(Query::new(query)).await?; + } + } + + // For each monitored table, get the capture columns + let table_columns = self.query(Query::new( + format!(r#" + SELECT OBJECT_NAME(ct.object_id, DB_ID('{database}')) as table_name, c.object_id, ct.capture_instance, ct.start_lsn, c.columns FROM [{database}].[cdc].[change_tables] as ct JOIN (SELECT object_id, STRING_AGG(column_name, ',') as columns FROM [{database}].[cdc].[captured_columns] GROUP BY object_id) as c ON c.object_id = ct.object_id;"# Review Comment: Since MSSQL has the source table and the change table with different names (by default `_ct` get's appended), either the user inserts the change table name or, there's need to make an internal conversion (while keeping a mapping of the names). Which one is the preferred one? -- 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]
