ryerraguntla commented on code in PR #4730: URL: https://github.com/apache/iggy/pull/4730#discussion_r4501522806
########## core/connectors/sources/opensearch_source/src/lib.rs: ########## @@ -0,0 +1,1364 @@ +// 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, transport::TransportBuilder}, +}; +use secrecy::{ExposeSecret, SecretString}; +use serde::{Deserialize, Serialize}; +use serde_json::{Value, json}; +use std::collections::{BTreeMap, HashSet}; +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; +const DEFAULT_REQUEST_TIMEOUT: &str = "30s"; +const DEFAULT_TIEBREAKER_FIELD: &str = "_id"; +// Sort clause length: timestamp, tie-breaker, `_index`. +const SORT_KEY_COUNT: usize = 3; +const TIEBREAKER_SORT_POSITION: usize = 1; +// Bounds the persisted state when many documents share the newest timestamp. +const MAX_WATERMARK_HITS: usize = 10_000; +// 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, + /// Sort tuple of the newest published hit. Its first element is the watermark: the next + /// scan starts there instead of after the tuple, so a hit that became searchable late is + /// still found. + #[serde(default)] + last_sort: Option<Vec<Value>>, + #[serde(default)] + error_count: usize, + #[serde(default)] + last_error: Option<String>, + #[serde(default)] + processing_stats: ProcessingStats, + /// Hits published at the watermark, skipped when the next scan reads them again. The + /// state is MessagePack without field names, so new fields must stay at the end. + #[serde(default)] + watermark_hits: Vec<HitKey>, + /// Page cursor of the running scan. Not persisted: a restart rescans from the watermark. + #[serde(skip)] + scan_cursor: Option<Vec<Value>>, +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)] +struct HitKey { + index: String, + id: String, +} + +#[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>, + /// Unique `keyword` field with doc values that breaks timestamp ties. Defaults to `_id`, + /// which OpenSearch sorts through heap-backed fielddata. + pub tiebreaker_field: Option<String>, + /// Limit per HTTP request, from connecting to the last byte of the response body. + pub request_timeout: 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, + request_timeout: 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: 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>, + /// `None` when the scan found nothing new to publish or skip. + last_sort: Option<Vec<Value>>, + watermark_hits: Vec<HitKey>, + scan_cursor: Option<Vec<Value>>, + last_poll_timestamp: Option<DateTime<Utc>>, + batch_bytes: u64, +} + +struct Page { + messages: Vec<ProducedMessage>, + sorts: Vec<Vec<Value>>, + keys: Vec<HitKey>, + /// Positions in `sorts` and `keys` of hits that were not published before. + fresh: Vec<usize>, + last_poll_timestamp: Option<DateTime<Utc>>, + batch_bytes: u64, + full: bool, +} + +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 request_timeout = + parse_duration(config.request_timeout.as_deref(), DEFAULT_REQUEST_TIMEOUT); + 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 = 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, + request_timeout, + 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 tiebreaker_field(&self) -> &str { + self.config + .tiebreaker_field + .as_deref() + .unwrap_or(DEFAULT_TIEBREAKER_FIELD) + } + + 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); + // Without a timeout a stalled response never returns, so neither retry nor the circuit + // breaker would ever see a failure. + let mut transport_builder = TransportBuilder::new(conn_pool).timeout(self.request_timeout); + + 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 => { Review Comment: Could not reproduce, -- 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]
