hubcio commented on code in PR #4730: URL: https://github.com/apache/iggy/pull/4730#discussion_r4469206831
########## core/connectors/sources/opensearch_source/src/lib.rs: ########## @@ -0,0 +1,1043 @@ +// 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 iggy_common::{DateTime, HeaderKey, HeaderValue, Utc}; +use iggy_connector_sdk::retry::{CircuitBreaker, RetryPolicy, parse_duration, retry_async}; +use iggy_connector_sdk::{ + ConnectorState, Error, ProducedMessage, ProducedMessages, Schema, Source, + source::SourceBatchResult, source_connector, +}; +use opensearch::{ + OpenSearch, SearchParts, + auth::Credentials, + http::{StatusCode, Url, response::Response, transport::TransportBuilder}, +}; +use secrecy::{ExposeSecret, SecretString}; +use serde::{Deserialize, Serialize}; +use serde_json::{Value, json}; +use std::collections::BTreeMap; +use std::sync::Arc; +use std::time::Duration; +use tokio::{sync::Mutex, time::sleep}; +use tracing::{debug, error, info, warn}; + +mod retry; +use crate::retry::{ + DEFAULT_CB_COOL_DOWN, DEFAULT_CB_THRESHOLD, DEFAULT_MAX_OPEN_RETRIES, DEFAULT_MAX_RETRIES, + DEFAULT_OPEN_RETRY_MAX_DELAY, DEFAULT_RETRY_DELAY, DEFAULT_RETRY_MAX_DELAY, RequestFailure, +}; + +source_connector!(OpenSearchSource); + +const CONNECTOR_NAME: &str = "OpenSearch source"; +const DEFAULT_POLLING_INTERVAL: &str = "10s"; +const DEFAULT_BATCH_SIZE: usize = 100; +// OpenSearch rejects `size` above the default `index.max_result_window` with HTTP 400. +const MAX_BATCH_SIZE: usize = 10_000; +const DOCUMENT_ID_HEADER: &str = "opensearch.document_id"; +const INDEX_HEADER: &str = "opensearch.index"; +const MAX_HEADER_VALUE_LENGTH: usize = 255; +const MAX_LAST_ERROR_CHARS: usize = 512; +// What OpenSearch reports for a document missing the sort field when it sorts last. The sort +// asks for missing values first, so seeing this means that request option was not honored. +const MISSING_LAST_SORT_VALUE: i64 = i64::MAX; + +static DOCUMENT_ID_HEADER_KEY: std::sync::LazyLock<HeaderKey> = std::sync::LazyLock::new(|| { + HeaderKey::try_from(DOCUMENT_ID_HEADER).expect("document id header key is a valid literal") +}); +static INDEX_HEADER_KEY: std::sync::LazyLock<HeaderKey> = std::sync::LazyLock::new(|| { + HeaderKey::try_from(INDEX_HEADER).expect("index header key is a valid literal") +}); + +#[derive(Debug, Clone, Serialize, Deserialize, Default)] +struct State { + #[serde(default)] + last_poll_timestamp: Option<DateTime<Utc>>, + #[serde(default)] + total_documents_published: usize, + #[serde(default)] + poll_count: usize, + /// OpenSearch `search_after` tuple from the last hit in the previous batch. + #[serde(default)] + search_after: Option<Vec<Value>>, + #[serde(default)] + error_count: usize, + #[serde(default)] + last_error: Option<String>, + #[serde(default)] + processing_stats: ProcessingStats, +} + +#[derive(Debug, Clone, Serialize, Deserialize, Default)] +struct ProcessingStats { + #[serde(default)] + total_bytes_processed: u64, + /// Running cumulative average over the connector's lifetime, persisted and accumulated + /// across restarts. Reflects long-term throughput baseline, not session-only average. + #[serde(default)] + avg_batch_processing_time_ms: f64, + #[serde(default)] + last_successful_poll: Option<DateTime<Utc>>, + #[serde(default)] + empty_polls_count: usize, + #[serde(default)] + successful_polls_count: usize, +} + +// `Deserialize` only: nothing re-serializes a plugin config, and omitting `Serialize` +// keeps `password` from ever being written out as plaintext. +#[derive(Debug, Clone, Deserialize)] +pub struct OpenSearchSourceConfig { + pub url: String, + pub index: String, + pub username: Option<String>, + pub password: Option<SecretString>, + pub query: Option<Value>, + pub polling_interval: Option<String>, + pub batch_size: Option<usize>, + pub timestamp_field: Option<String>, + #[serde(default)] + pub verbose_logging: bool, + /// Total poll-phase attempt count (not retry count); 1 = no retries, 3 = 2 retries. + pub max_retries: Option<u32>, + pub retry_delay: Option<String>, + pub retry_max_delay: Option<String>, + /// Total open-phase attempt count (not retry count); 1 = no retries, 5 = 4 retries. + pub max_open_retries: Option<u32>, + pub open_retry_max_delay: Option<String>, + pub circuit_breaker_threshold: Option<u32>, + pub circuit_breaker_cool_down: Option<String>, +} + +#[derive(Debug)] +pub struct OpenSearchSource { + id: u32, + config: OpenSearchSourceConfig, + client: Option<OpenSearch>, + polling_interval: Duration, + /// Pre-built query body (query + size + sort). Set once in `open()` after config validation. + /// `None` before `open()` and after a failed `open()`. + search_body_base: Option<Value>, + verbose: bool, + poll_retry_policy: RetryPolicy, + open_retry_policy: RetryPolicy, + circuit_breaker: Arc<CircuitBreaker>, + state: Mutex<State>, + /// Candidate state of the batch in flight; committed on ACK, dropped on NACK. + pending_state: Mutex<Option<State>>, + /// `Some(cause)` when runtime state restore was rejected; `None` means restore succeeded. + state_restore_error: Option<String>, +} + +struct SearchOutcome { + messages: Vec<ProducedMessage>, + search_after: Option<Vec<Value>>, + last_poll_timestamp: Option<DateTime<Utc>>, + batch_bytes: u64, +} + +impl OpenSearchSource { + pub fn new(id: u32, config: OpenSearchSourceConfig, state: Option<ConnectorState>) -> Self { + let polling_interval = + parse_duration(config.polling_interval.as_deref(), DEFAULT_POLLING_INTERVAL); + let verbose = config.verbose_logging; + let (restored_state, state_restore_error) = restore_state(id, state); + + let cb_threshold = config + .circuit_breaker_threshold + .unwrap_or(DEFAULT_CB_THRESHOLD); + let cb_cool_down = parse_duration( + config.circuit_breaker_cool_down.as_deref(), + DEFAULT_CB_COOL_DOWN, + ); + let circuit_breaker = Arc::new(CircuitBreaker::new(cb_threshold, cb_cool_down)); + let retry_delay = parse_duration(config.retry_delay.as_deref(), DEFAULT_RETRY_DELAY); + let poll_retry_policy = RetryPolicy { + max_attempts: config.max_retries.unwrap_or(DEFAULT_MAX_RETRIES), + base_delay: retry_delay, + max_delay: parse_duration(config.retry_max_delay.as_deref(), DEFAULT_RETRY_MAX_DELAY), + }; + // The open phase has no base delay of its own: it shares `retry_delay` and only + // `open_retry_max_delay` caps the growth independently. + let open_retry_policy = RetryPolicy { + max_attempts: config.max_open_retries.unwrap_or(DEFAULT_MAX_OPEN_RETRIES), + base_delay: retry_delay, + max_delay: parse_duration( + config.open_retry_max_delay.as_deref(), + DEFAULT_OPEN_RETRY_MAX_DELAY, + ), + }; + + OpenSearchSource { + id, + config, + poll_retry_policy, + open_retry_policy, + circuit_breaker, + client: None, + polling_interval, + search_body_base: None, + verbose, + state: Mutex::new(restored_state), + pending_state: Mutex::new(None), + state_restore_error, + } + } + + fn serialize_state(&self, state: &State) -> Option<ConnectorState> { + ConnectorState::serialize(state, CONNECTOR_NAME, self.id) + } + + fn batch_size(&self) -> usize { + self.config.batch_size.unwrap_or(DEFAULT_BATCH_SIZE) + } + + fn timestamp_field(&self) -> Result<&str, Error> { + self.config + .timestamp_field + .as_deref() + .filter(|s| !s.is_empty()) + .ok_or_else(|| Error::InvalidConfigValue("timestamp_field is required".to_string())) + } + + async fn create_client(&self) -> Result<OpenSearch, Error> { + let url = Url::parse(&self.config.url).map_err(|error| { + Error::InvalidConfigValue(format!("Invalid OpenSearch URL: {error}")) + })?; + + let conn_pool = opensearch::http::transport::SingleNodeConnectionPool::new(url); + let mut transport_builder = TransportBuilder::new(conn_pool); + + if let (Some(username), Some(password)) = (&self.config.username, &self.config.password) { + let credentials = + Credentials::Basic(username.clone(), password.expose_secret().to_string()); + transport_builder = transport_builder.auth(credentials); + } + + let transport = transport_builder + .build() + .map_err(|error| Error::InitError(format!("Failed to build transport: {error}")))?; + + Ok(OpenSearch::new(transport)) + } + + async fn check_index_exists_with_retry(&self, client: &OpenSearch) -> Result<(), Error> { + let index = &self.config.index; + retry_async( + self.open_retry_policy, + &format!("{CONNECTOR_NAME} ID {} index exists check", self.id), + RequestFailure::is_transient, + || async { + let response = client + .indices() + .exists(opensearch::indices::IndicesExistsParts::Index(&[index])) + .send() + .await + .map_err(RequestFailure::transport)?; + RequestFailure::check(response).await.map(drop) + }, + ) + .await + .map_err(|failure| match failure.into_error() { + RequestFailure::Transport(message) => { + Error::InitError(format!("Failed to check index existence: {message}")) + } + RequestFailure::Status { status, .. } if status == StatusCode::NOT_FOUND => { + Error::InitError(format!( + "Index '{index}' does not exist or is not accessible" + )) + } + failure @ RequestFailure::Status { .. } => { + Error::InitError(format!("Index '{index}' check failed: {failure}")) + } + }) + } + + async fn send_search_with_retry( + &self, + client: &OpenSearch, + search_body: &Value, + ) -> Result<Response, Error> { + retry_async( + self.poll_retry_policy, + &format!("{CONNECTOR_NAME} ID {} search", self.id), + RequestFailure::is_transient, + || async { + let response = client + .search(SearchParts::Index(&[&self.config.index])) + .allow_partial_search_results(false) + .body(search_body) + .send() + .await + .map_err(RequestFailure::transport)?; + RequestFailure::check(response).await + }, + ) + .await + .map_err(|failure| failure.into_error().into_search_error()) + } + + async fn search_documents(&self, client: &OpenSearch) -> Result<SearchOutcome, Error> { + let state = self.state.lock().await; + let search_after = state.search_after.clone(); + drop(state); + + let timestamp_field = self.timestamp_field()?; + + let mut search_body = self + .search_body_base + .as_ref() + .ok_or_else(|| { + Error::Connection("connector not initialized; call open() first".to_string()) + })? + .clone(); + + if let Some(cursor) = search_after { + search_body["search_after"] = json!(cursor); + } + + let response = self.send_search_with_retry(client, &search_body).await?; + + let mut response_body: Value = response.json().await.map_err(|e| { + Error::PermanentHttpError(format!("Failed to parse search response: {e}")) + })?; + ensure_complete_response(&response_body)?; + + let hits: Vec<Value> = response_body + .get_mut("hits") + .and_then(|h| h.get_mut("hits")) + .and_then(|arr| arr.as_array_mut()) + .map(std::mem::take) + .unwrap_or_default(); + + let mut messages = Vec::with_capacity(hits.len()); + let mut batch_bytes = 0u64; + let mut last_sort = None; + let mut last_poll_timestamp = None; + + for hit in &hits { Review Comment: warning: the hit loop retains every parsed document alongside the complete serialized batch, increasing peak memory with document size. retain the final sort tuple separately and consume each hit by value after validation. ########## core/connectors/sources/opensearch_source/src/lib.rs: ########## @@ -0,0 +1,1043 @@ +// 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 iggy_common::{DateTime, HeaderKey, HeaderValue, Utc}; +use iggy_connector_sdk::retry::{CircuitBreaker, RetryPolicy, parse_duration, retry_async}; +use iggy_connector_sdk::{ + ConnectorState, Error, ProducedMessage, ProducedMessages, Schema, Source, + source::SourceBatchResult, source_connector, +}; +use opensearch::{ + OpenSearch, SearchParts, + auth::Credentials, + http::{StatusCode, Url, response::Response, transport::TransportBuilder}, +}; +use secrecy::{ExposeSecret, SecretString}; +use serde::{Deserialize, Serialize}; +use serde_json::{Value, json}; +use std::collections::BTreeMap; +use std::sync::Arc; +use std::time::Duration; +use tokio::{sync::Mutex, time::sleep}; +use tracing::{debug, error, info, warn}; + +mod retry; +use crate::retry::{ + DEFAULT_CB_COOL_DOWN, DEFAULT_CB_THRESHOLD, DEFAULT_MAX_OPEN_RETRIES, DEFAULT_MAX_RETRIES, + DEFAULT_OPEN_RETRY_MAX_DELAY, DEFAULT_RETRY_DELAY, DEFAULT_RETRY_MAX_DELAY, RequestFailure, +}; + +source_connector!(OpenSearchSource); + +const CONNECTOR_NAME: &str = "OpenSearch source"; +const DEFAULT_POLLING_INTERVAL: &str = "10s"; +const DEFAULT_BATCH_SIZE: usize = 100; +// OpenSearch rejects `size` above the default `index.max_result_window` with HTTP 400. +const MAX_BATCH_SIZE: usize = 10_000; +const DOCUMENT_ID_HEADER: &str = "opensearch.document_id"; +const INDEX_HEADER: &str = "opensearch.index"; +const MAX_HEADER_VALUE_LENGTH: usize = 255; +const MAX_LAST_ERROR_CHARS: usize = 512; +// What OpenSearch reports for a document missing the sort field when it sorts last. The sort +// asks for missing values first, so seeing this means that request option was not honored. +const MISSING_LAST_SORT_VALUE: i64 = i64::MAX; + +static DOCUMENT_ID_HEADER_KEY: std::sync::LazyLock<HeaderKey> = std::sync::LazyLock::new(|| { + HeaderKey::try_from(DOCUMENT_ID_HEADER).expect("document id header key is a valid literal") +}); +static INDEX_HEADER_KEY: std::sync::LazyLock<HeaderKey> = std::sync::LazyLock::new(|| { + HeaderKey::try_from(INDEX_HEADER).expect("index header key is a valid literal") +}); + +#[derive(Debug, Clone, Serialize, Deserialize, Default)] +struct State { + #[serde(default)] + last_poll_timestamp: Option<DateTime<Utc>>, + #[serde(default)] + total_documents_published: usize, + #[serde(default)] + poll_count: usize, + /// OpenSearch `search_after` tuple from the last hit in the previous batch. + #[serde(default)] + search_after: Option<Vec<Value>>, + #[serde(default)] + error_count: usize, + #[serde(default)] + last_error: Option<String>, + #[serde(default)] + processing_stats: ProcessingStats, +} + +#[derive(Debug, Clone, Serialize, Deserialize, Default)] +struct ProcessingStats { + #[serde(default)] + total_bytes_processed: u64, + /// Running cumulative average over the connector's lifetime, persisted and accumulated + /// across restarts. Reflects long-term throughput baseline, not session-only average. + #[serde(default)] + avg_batch_processing_time_ms: f64, + #[serde(default)] + last_successful_poll: Option<DateTime<Utc>>, + #[serde(default)] + empty_polls_count: usize, + #[serde(default)] + successful_polls_count: usize, +} + +// `Deserialize` only: nothing re-serializes a plugin config, and omitting `Serialize` +// keeps `password` from ever being written out as plaintext. +#[derive(Debug, Clone, Deserialize)] +pub struct OpenSearchSourceConfig { + pub url: String, + pub index: String, + pub username: Option<String>, + pub password: Option<SecretString>, + pub query: Option<Value>, + pub polling_interval: Option<String>, + pub batch_size: Option<usize>, + pub timestamp_field: Option<String>, + #[serde(default)] + pub verbose_logging: bool, + /// Total poll-phase attempt count (not retry count); 1 = no retries, 3 = 2 retries. + pub max_retries: Option<u32>, + pub retry_delay: Option<String>, + pub retry_max_delay: Option<String>, + /// Total open-phase attempt count (not retry count); 1 = no retries, 5 = 4 retries. + pub max_open_retries: Option<u32>, + pub open_retry_max_delay: Option<String>, + pub circuit_breaker_threshold: Option<u32>, + pub circuit_breaker_cool_down: Option<String>, +} + +#[derive(Debug)] +pub struct OpenSearchSource { + id: u32, + config: OpenSearchSourceConfig, + client: Option<OpenSearch>, + polling_interval: Duration, + /// Pre-built query body (query + size + sort). Set once in `open()` after config validation. + /// `None` before `open()` and after a failed `open()`. + search_body_base: Option<Value>, + verbose: bool, + poll_retry_policy: RetryPolicy, + open_retry_policy: RetryPolicy, + circuit_breaker: Arc<CircuitBreaker>, + state: Mutex<State>, + /// Candidate state of the batch in flight; committed on ACK, dropped on NACK. + pending_state: Mutex<Option<State>>, + /// `Some(cause)` when runtime state restore was rejected; `None` means restore succeeded. + state_restore_error: Option<String>, +} + +struct SearchOutcome { + messages: Vec<ProducedMessage>, + search_after: Option<Vec<Value>>, + last_poll_timestamp: Option<DateTime<Utc>>, + batch_bytes: u64, +} + +impl OpenSearchSource { + pub fn new(id: u32, config: OpenSearchSourceConfig, state: Option<ConnectorState>) -> Self { + let polling_interval = + parse_duration(config.polling_interval.as_deref(), DEFAULT_POLLING_INTERVAL); + let verbose = config.verbose_logging; + let (restored_state, state_restore_error) = restore_state(id, state); + + let cb_threshold = config + .circuit_breaker_threshold + .unwrap_or(DEFAULT_CB_THRESHOLD); + let cb_cool_down = parse_duration( + config.circuit_breaker_cool_down.as_deref(), + DEFAULT_CB_COOL_DOWN, + ); + let circuit_breaker = Arc::new(CircuitBreaker::new(cb_threshold, cb_cool_down)); + let retry_delay = parse_duration(config.retry_delay.as_deref(), DEFAULT_RETRY_DELAY); + let poll_retry_policy = RetryPolicy { + max_attempts: config.max_retries.unwrap_or(DEFAULT_MAX_RETRIES), + base_delay: retry_delay, + max_delay: parse_duration(config.retry_max_delay.as_deref(), DEFAULT_RETRY_MAX_DELAY), + }; + // The open phase has no base delay of its own: it shares `retry_delay` and only + // `open_retry_max_delay` caps the growth independently. + let open_retry_policy = RetryPolicy { + max_attempts: config.max_open_retries.unwrap_or(DEFAULT_MAX_OPEN_RETRIES), + base_delay: retry_delay, + max_delay: parse_duration( + config.open_retry_max_delay.as_deref(), + DEFAULT_OPEN_RETRY_MAX_DELAY, + ), + }; + + OpenSearchSource { + id, + config, + poll_retry_policy, + open_retry_policy, + circuit_breaker, + client: None, + polling_interval, + search_body_base: None, + verbose, + state: Mutex::new(restored_state), + pending_state: Mutex::new(None), + state_restore_error, + } + } + + fn serialize_state(&self, state: &State) -> Option<ConnectorState> { + ConnectorState::serialize(state, CONNECTOR_NAME, self.id) + } + + fn batch_size(&self) -> usize { + self.config.batch_size.unwrap_or(DEFAULT_BATCH_SIZE) + } + + fn timestamp_field(&self) -> Result<&str, Error> { + self.config + .timestamp_field + .as_deref() + .filter(|s| !s.is_empty()) + .ok_or_else(|| Error::InvalidConfigValue("timestamp_field is required".to_string())) + } + + async fn create_client(&self) -> Result<OpenSearch, Error> { + let url = Url::parse(&self.config.url).map_err(|error| { + Error::InvalidConfigValue(format!("Invalid OpenSearch URL: {error}")) + })?; + + let conn_pool = opensearch::http::transport::SingleNodeConnectionPool::new(url); + let mut transport_builder = TransportBuilder::new(conn_pool); + + if let (Some(username), Some(password)) = (&self.config.username, &self.config.password) { + let credentials = + Credentials::Basic(username.clone(), password.expose_secret().to_string()); + transport_builder = transport_builder.auth(credentials); + } + + let transport = transport_builder + .build() + .map_err(|error| Error::InitError(format!("Failed to build transport: {error}")))?; + + Ok(OpenSearch::new(transport)) + } + + async fn check_index_exists_with_retry(&self, client: &OpenSearch) -> Result<(), Error> { + let index = &self.config.index; + retry_async( + self.open_retry_policy, + &format!("{CONNECTOR_NAME} ID {} index exists check", self.id), + RequestFailure::is_transient, + || async { + let response = client + .indices() + .exists(opensearch::indices::IndicesExistsParts::Index(&[index])) + .send() + .await + .map_err(RequestFailure::transport)?; + RequestFailure::check(response).await.map(drop) + }, + ) + .await + .map_err(|failure| match failure.into_error() { + RequestFailure::Transport(message) => { + Error::InitError(format!("Failed to check index existence: {message}")) + } + RequestFailure::Status { status, .. } if status == StatusCode::NOT_FOUND => { + Error::InitError(format!( + "Index '{index}' does not exist or is not accessible" + )) + } + failure @ RequestFailure::Status { .. } => { + Error::InitError(format!("Index '{index}' check failed: {failure}")) + } + }) + } + + async fn send_search_with_retry( + &self, + client: &OpenSearch, + search_body: &Value, + ) -> Result<Response, Error> { + retry_async( + self.poll_retry_policy, + &format!("{CONNECTOR_NAME} ID {} search", self.id), + RequestFailure::is_transient, + || async { + let response = client + .search(SearchParts::Index(&[&self.config.index])) + .allow_partial_search_results(false) + .body(search_body) + .send() + .await + .map_err(RequestFailure::transport)?; + RequestFailure::check(response).await + }, + ) + .await + .map_err(|failure| failure.into_error().into_search_error()) + } + + async fn search_documents(&self, client: &OpenSearch) -> Result<SearchOutcome, Error> { + let state = self.state.lock().await; + let search_after = state.search_after.clone(); + drop(state); + + let timestamp_field = self.timestamp_field()?; + + let mut search_body = self + .search_body_base + .as_ref() + .ok_or_else(|| { + Error::Connection("connector not initialized; call open() first".to_string()) + })? + .clone(); + + if let Some(cursor) = search_after { Review Comment: critical: the saved cursor permanently skips documents that become visible later with the same timestamp and a smaller `_id`. revisit the final timestamp between scans, deduplicate documents, and keep page cursors within each scan. ########## core/connectors/sources/opensearch_source/src/lib.rs: ########## @@ -0,0 +1,1043 @@ +// 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 iggy_common::{DateTime, HeaderKey, HeaderValue, Utc}; +use iggy_connector_sdk::retry::{CircuitBreaker, RetryPolicy, parse_duration, retry_async}; +use iggy_connector_sdk::{ + ConnectorState, Error, ProducedMessage, ProducedMessages, Schema, Source, + source::SourceBatchResult, source_connector, +}; +use opensearch::{ + OpenSearch, SearchParts, + auth::Credentials, + http::{StatusCode, Url, response::Response, transport::TransportBuilder}, +}; +use secrecy::{ExposeSecret, SecretString}; +use serde::{Deserialize, Serialize}; +use serde_json::{Value, json}; +use std::collections::BTreeMap; +use std::sync::Arc; +use std::time::Duration; +use tokio::{sync::Mutex, time::sleep}; +use tracing::{debug, error, info, warn}; + +mod retry; +use crate::retry::{ + DEFAULT_CB_COOL_DOWN, DEFAULT_CB_THRESHOLD, DEFAULT_MAX_OPEN_RETRIES, DEFAULT_MAX_RETRIES, + DEFAULT_OPEN_RETRY_MAX_DELAY, DEFAULT_RETRY_DELAY, DEFAULT_RETRY_MAX_DELAY, RequestFailure, +}; + +source_connector!(OpenSearchSource); + +const CONNECTOR_NAME: &str = "OpenSearch source"; +const DEFAULT_POLLING_INTERVAL: &str = "10s"; +const DEFAULT_BATCH_SIZE: usize = 100; +// OpenSearch rejects `size` above the default `index.max_result_window` with HTTP 400. +const MAX_BATCH_SIZE: usize = 10_000; +const DOCUMENT_ID_HEADER: &str = "opensearch.document_id"; +const INDEX_HEADER: &str = "opensearch.index"; +const MAX_HEADER_VALUE_LENGTH: usize = 255; +const MAX_LAST_ERROR_CHARS: usize = 512; +// What OpenSearch reports for a document missing the sort field when it sorts last. The sort +// asks for missing values first, so seeing this means that request option was not honored. +const MISSING_LAST_SORT_VALUE: i64 = i64::MAX; + +static DOCUMENT_ID_HEADER_KEY: std::sync::LazyLock<HeaderKey> = std::sync::LazyLock::new(|| { + HeaderKey::try_from(DOCUMENT_ID_HEADER).expect("document id header key is a valid literal") +}); +static INDEX_HEADER_KEY: std::sync::LazyLock<HeaderKey> = std::sync::LazyLock::new(|| { + HeaderKey::try_from(INDEX_HEADER).expect("index header key is a valid literal") +}); + +#[derive(Debug, Clone, Serialize, Deserialize, Default)] +struct State { + #[serde(default)] + last_poll_timestamp: Option<DateTime<Utc>>, + #[serde(default)] + total_documents_published: usize, + #[serde(default)] + poll_count: usize, + /// OpenSearch `search_after` tuple from the last hit in the previous batch. + #[serde(default)] + search_after: Option<Vec<Value>>, + #[serde(default)] + error_count: usize, + #[serde(default)] + last_error: Option<String>, + #[serde(default)] + processing_stats: ProcessingStats, +} + +#[derive(Debug, Clone, Serialize, Deserialize, Default)] +struct ProcessingStats { + #[serde(default)] + total_bytes_processed: u64, + /// Running cumulative average over the connector's lifetime, persisted and accumulated + /// across restarts. Reflects long-term throughput baseline, not session-only average. + #[serde(default)] + avg_batch_processing_time_ms: f64, + #[serde(default)] + last_successful_poll: Option<DateTime<Utc>>, + #[serde(default)] + empty_polls_count: usize, + #[serde(default)] + successful_polls_count: usize, +} + +// `Deserialize` only: nothing re-serializes a plugin config, and omitting `Serialize` +// keeps `password` from ever being written out as plaintext. +#[derive(Debug, Clone, Deserialize)] +pub struct OpenSearchSourceConfig { + pub url: String, + pub index: String, + pub username: Option<String>, + pub password: Option<SecretString>, + pub query: Option<Value>, + pub polling_interval: Option<String>, + pub batch_size: Option<usize>, + pub timestamp_field: Option<String>, + #[serde(default)] + pub verbose_logging: bool, + /// Total poll-phase attempt count (not retry count); 1 = no retries, 3 = 2 retries. + pub max_retries: Option<u32>, + pub retry_delay: Option<String>, + pub retry_max_delay: Option<String>, + /// Total open-phase attempt count (not retry count); 1 = no retries, 5 = 4 retries. + pub max_open_retries: Option<u32>, + pub open_retry_max_delay: Option<String>, + pub circuit_breaker_threshold: Option<u32>, + pub circuit_breaker_cool_down: Option<String>, +} + +#[derive(Debug)] +pub struct OpenSearchSource { + id: u32, + config: OpenSearchSourceConfig, + client: Option<OpenSearch>, + polling_interval: Duration, + /// Pre-built query body (query + size + sort). Set once in `open()` after config validation. + /// `None` before `open()` and after a failed `open()`. + search_body_base: Option<Value>, + verbose: bool, + poll_retry_policy: RetryPolicy, + open_retry_policy: RetryPolicy, + circuit_breaker: Arc<CircuitBreaker>, Review Comment: simplification: the source wraps its exclusively owned circuit breaker in `Arc`, but no path clones or shares it independently. store `CircuitBreaker` directly. -- 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]
