ryerraguntla commented on code in PR #4730:
URL: https://github.com/apache/iggy/pull/4730#discussion_r4501525464


##########
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 => {
+                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 check_tiebreaker_mapping(&self, client: &OpenSearch) -> 
Result<(), Error> {
+        let field = self.tiebreaker_field();
+        if field == DEFAULT_TIEBREAKER_FIELD {
+            return Ok(());
+        }
+        let index = &self.config.index;
+        let body = retry_async(
+            self.open_retry_policy,
+            &format!("{CONNECTOR_NAME} ID {} tiebreaker mapping check", 
self.id),
+            RequestFailure::is_transient,
+            || async {
+                let response = client
+                    .indices()
+                    .get_field_mapping(
+                        
opensearch::indices::IndicesGetFieldMappingParts::IndexFields(
+                            &[index],
+                            &[field],
+                        ),
+                    )
+                    .send()
+                    .await
+                    .map_err(RequestFailure::transport)?;
+                RequestFailure::read_text(response).await
+            },
+        )
+        .await
+        .map_err(|failure| {
+            Error::InitError(format!(
+                "Failed to read the mapping of '{field}': {}",
+                failure.into_error()
+            ))
+        })?;
+        let mapping: Value = serde_json::from_str(&body).map_err(|error| {
+            Error::InitError(format!("Failed to parse the mapping of 
'{field}': {error}"))
+        })?;
+        ensure_keyword_with_doc_values(&mapping, field)
+    }
+
+    async fn send_search_with_retry(
+        &self,
+        client: &OpenSearch,
+        search_body: &Value,
+    ) -> Result<String, 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::read_text(response).await
+            },
+        )
+        .await
+        .map_err(|failure| failure.into_error().into_search_error())
+    }
+
+    async fn search_documents(&self, client: &OpenSearch) -> 
Result<SearchOutcome, Error> {
+        let (last_sort, watermark_hits, scan_cursor) = {
+            let state = self.state.lock().await;
+            (
+                state.last_sort.clone(),
+                state.watermark_hits.clone(),
+                state.scan_cursor.clone(),
+            )
+        };
+        let watermark = last_sort.as_deref().and_then(<[Value]>::first);
+        let published: HashSet<&HitKey> = watermark_hits.iter().collect();
+        let mut cursor = scan_cursor.or_else(|| watermark.map(restart_cursor));
+
+        let page = loop {
+            let hits = self.fetch_hits(client, cursor.take()).await?;
+            let page = self.read_page(hits, watermark, &published)?;
+            if !page.fresh.is_empty() || !page.full {
+                break page;
+            }
+            // A full page of hits published before: step past them within 
this scan.
+            cursor = page.sorts.last().cloned();
+        };
+
+        let moved = match page.fresh.last() {
+            Some(&last) => {
+                let sort = page.sorts[last].clone();
+                let mut hits = if watermark == sort.first() {
+                    watermark_hits.clone()
+                } else {
+                    Vec::new()
+                };
+                hits.extend(
+                    page.fresh
+                        .iter()
+                        .filter(|&&position| page.sorts[position].first() == 
sort.first())
+                        .map(|&position| page.keys[position].clone()),
+                );
+                if hits.len() > MAX_WATERMARK_HITS {
+                    return Err(Error::SchemaMismatch(format!(
+                        "more than {MAX_WATERMARK_HITS} documents share one 
value of the sort \
+                         field; use a timestamp_field with a finer resolution"
+                    )));
+                }
+                Some((sort, hits))
+            }
+            None => None,
+        };
+        let (last_sort, watermark_hits) = moved.unzip();
+
+        Ok(SearchOutcome {
+            scan_cursor: if page.full {
+                page.sorts.last().cloned()
+            } else {
+                None
+            },
+            messages: page.messages,
+            last_sort,
+            watermark_hits: watermark_hits.unwrap_or_default(),
+            last_poll_timestamp: page.last_poll_timestamp,
+            batch_bytes: page.batch_bytes,
+        })
+    }
+
+    async fn fetch_hits(
+        &self,
+        client: &OpenSearch,
+        cursor: Option<Vec<Value>>,
+    ) -> Result<Vec<Value>, Error> {
+        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) = cursor {
+            search_body["search_after"] = Value::Array(cursor);
+        }
+
+        let body = self.send_search_with_retry(client, &search_body).await?;
+        let mut response_body: Value = serde_json::from_str(&body).map_err(|e| 
{
+            Error::PermanentHttpError(format!("Failed to parse search 
response: {e}"))
+        })?;
+        ensure_complete_response(&response_body)?;
+
+        Ok(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())
+    }
+
+    fn read_page(
+        &self,
+        hits: Vec<Value>,
+        watermark: Option<&Value>,
+        published: &HashSet<&HitKey>,
+    ) -> Result<Page, Error> {
+        let timestamp_field = self.timestamp_field()?;
+        let mut page = Page {
+            full: hits.len() >= self.batch_size(),
+            messages: Vec::with_capacity(hits.len()),
+            sorts: Vec::with_capacity(hits.len()),
+            keys: Vec::with_capacity(hits.len()),
+            fresh: Vec::new(),
+            last_poll_timestamp: None,
+            batch_bytes: 0,
+        };
+
+        for mut hit in hits {
+            let Some(sort) = hit
+                .get_mut("sort")
+                .and_then(Value::as_array_mut)
+                .filter(|a| !a.is_empty())
+                .map(std::mem::take)
+            else {
+                return Err(Error::SchemaMismatch(format!(
+                    "OpenSearch hit {:?} has no sort tuple; the cursor cannot 
move past it \
+                     without risking skipped documents",
+                    hit.get("_id").and_then(Value::as_str)
+                )));
+            };
+
+            if sort.first().and_then(Value::as_i64) == 
Some(MISSING_LAST_SORT_VALUE) {
+                return Err(Error::SchemaMismatch(format!(
+                    "OpenSearch hit {:?} sorted on the missing-value sentinel; 
a cursor there \
+                     would skip every document indexed later",
+                    hit.get("_id").and_then(Value::as_str)
+                )));
+            }
+
+            let key = hit_key(&hit, &sort, &self.config.index, 
self.tiebreaker_field())?;
+            let already_published = watermark == sort.first() && 
published.contains(&key);
+            if !already_published {
+                page.fresh.push(page.sorts.len());
+                match hit.get_mut("_source").map(Value::take) {
+                    Some(source) => {
+                        self.push_message(&mut page, &hit, &key, &source, 
timestamp_field)?;
+                    }
+                    None => warn!(
+                        connector_id = self.id,
+                        hit_id = key.id,
+                        hit_index = key.index,
+                        "Skipping OpenSearch hit without _source; it has no 
body to publish"
+                    ),
+                }
+            }
+            page.sorts.push(sort);
+            page.keys.push(key);
+        }
+        Ok(page)
+    }
+
+    fn push_message(
+        &self,
+        page: &mut Page,
+        hit: &Value,
+        key: &HitKey,
+        source: &Value,
+        timestamp_field: &str,
+    ) -> Result<(), Error> {
+        let document_timestamp =
+            field_value(source, 
timestamp_field).and_then(parse_document_timestamp);
+        if let Some(timestamp_utc) = document_timestamp {
+            page.last_poll_timestamp = Some(timestamp_utc);
+        }
+
+        let payload = serde_json::to_vec(source)
+            .map_err(|e| Error::Serialization(format!("Failed to serialize 
document: {e}")))?;
+        page.batch_bytes += payload.len() as u64;
+
+        // `index` may be a pattern, so the concrete index of each hit keeps 
equal `_id`s
+        // in different indices from sharing a message id.
+        page.messages.push(ProducedMessage {
+            id: Some(message_id(&key.index, &key.id)),
+            headers: document_headers(
+                self.id,
+                hit.get("_index").and_then(Value::as_str),
+                hit.get("_id").and_then(Value::as_str),
+            ),
+            checksum: None,
+            timestamp: None,
+            origin_timestamp: None,
+            payload,
+        });
+        Ok(())
+    }
+
+    async fn finalize_poll(
+        &self,
+        outcome: SearchOutcome,
+        processing_time_ms: f64,
+    ) -> Result<(Vec<ProducedMessage>, Option<ConnectorState>), Error> {
+        let mut state = self.state.lock().await.clone();
+        state.total_documents_published += outcome.messages.len();
+        state.poll_count += 1;
+        let cursor_moved = outcome.last_sort.is_some();
+        if outcome.last_sort.is_some() {
+            state.last_sort = outcome.last_sort;
+            state.watermark_hits = outcome.watermark_hits;
+        }
+        state.scan_cursor = outcome.scan_cursor;
+        if let Some(timestamp) = outcome.last_poll_timestamp {
+            state.last_poll_timestamp = Some(timestamp);
+        }
+        state.processing_stats.total_bytes_processed += outcome.batch_bytes;
+
+        if outcome.messages.is_empty() {
+            state.processing_stats.empty_polls_count += 1;
+        } else {
+            state.processing_stats.successful_polls_count += 1;
+            state.processing_stats.last_successful_poll = Some(Utc::now());
+        }
+
+        let total_polls = state.processing_stats.successful_polls_count
+            + state.processing_stats.empty_polls_count;
+        state.processing_stats.avg_batch_processing_time_ms += 
(processing_time_ms
+            - state.processing_stats.avg_batch_processing_time_ms)
+            / total_polls as f64;
+
+        let produced_count = outcome.messages.len();
+        let total_documents_published = state.total_documents_published;
+        let messages = outcome.messages;
+        // A moved cursor waits for the batch result so a Nack can undo it, 
including a page
+        // whose hits were all skipped for lacking `_source`. A poll that saw 
no hits has no
+        // cursor to protect, so its counters commit directly.
+        let persisted_state = if cursor_moved {
+            let persisted = self.serialize_state(&state).ok_or_else(|| {
+                Error::Serialization("failed to serialize source 
state".to_string())
+            })?;
+            *self.pending_state.lock().await = Some(state);
+            Some(persisted)
+        } else {
+            *self.state.lock().await = state;
+            None
+        };
+
+        if self.verbose {
+            info!(
+                "OpenSearch source connector ID: {} produced {produced_count} 
messages. \
+                 Total published: {total_documents_published}",
+                self.id
+            );
+        } else {
+            debug!(
+                "OpenSearch source connector ID: {} produced {produced_count} 
messages. \
+                 Total published: {total_documents_published}",
+                self.id
+            );
+        }
+
+        Ok((messages, persisted_state))
+    }
+
+    #[cfg(test)]
+    fn client_initialized(&self) -> bool {
+        self.client.is_some()
+    }
+
+    #[cfg(test)]
+    async fn test_metrics(&self) -> (usize, usize, usize, usize) {
+        let state = self.state.lock().await;
+        (
+            state.total_documents_published,
+            state.poll_count,
+            state.error_count,
+            state.processing_stats.empty_polls_count,
+        )
+    }
+
+    #[cfg(test)]
+    async fn test_last_sort(&self) -> Option<Vec<Value>> {
+        self.state.lock().await.last_sort.clone()
+    }
+
+    #[cfg(test)]
+    async fn test_last_poll_timestamp(&self) -> Option<DateTime<Utc>> {
+        self.state.lock().await.last_poll_timestamp
+    }
+
+    async fn handle_poll_error(&self, error: Error) -> 
Result<ProducedMessages, Error> {
+        // Only cluster-side trouble opens the breaker. A permanent error (bad 
query, missing
+        // sort tuple) is a configuration or mapping problem that skipping 
polls cannot fix,
+        // and an open breaker would hide it behind a quiet empty batch.
+        if matches!(error, Error::Connection(_) | Error::HttpRequestFailed(_)) 
{
+            self.circuit_breaker.record_failure().await;
+        }
+        let mut state = self.state.lock().await;
+        // error_count and last_error accumulate in State and are captured by 
the next
+        // finalize_poll call (success path). Both runtime ConnectorState and 
file state
+        // preserve these values across restarts.
+        state.error_count += 1;
+        state.last_error = Some(truncated(error.to_string()));
+        drop(state);
+        error!(
+            "{CONNECTOR_NAME} connector ID: {} poll failed: {error}",
+            self.id
+        );
+        Err(error)
+    }
+}
+
+// `last_error` is persisted with every batch, and an OpenSearch error body 
can be large.
+fn truncated(message: String) -> String {
+    match message.char_indices().nth(MAX_LAST_ERROR_CHARS) {
+        Some((end, _)) => format!("{}...", &message[..end]),
+        None => message,
+    }
+}
+
+// Sorting on any other field type would fail or load the field into heap, and 
a field that is
+// missing from one matched index gives its documents no tie-break at all.
+fn ensure_keyword_with_doc_values(mappings: &Value, field: &str) -> Result<(), 
Error> {
+    let Some(indices) = mappings.as_object() else {
+        return Ok(());
+    };
+    for (index, entry) in indices {
+        let Some(mapping) = entry
+            .get("mappings")
+            .and_then(|mappings| mappings.get(field))
+            .and_then(|mapped| mapped.get("mapping"))
+            .and_then(Value::as_object)
+            .and_then(|by_leaf_name| by_leaf_name.values().next())
+        else {
+            return Err(Error::InvalidConfigValue(format!(
+                "tiebreaker_field '{field}' is not mapped in index '{index}'"
+            )));
+        };
+        if mapping.get("type").and_then(Value::as_str) != Some("keyword") {
+            return Err(Error::InvalidConfigValue(format!(
+                "tiebreaker_field '{field}' must be of type keyword in index 
'{index}'"
+            )));
+        }
+        if mapping.get("doc_values").and_then(Value::as_bool) == Some(false) {
+            return Err(Error::InvalidConfigValue(format!(
+                "tiebreaker_field '{field}' must have doc_values enabled in 
index '{index}'"
+            )));
+        }
+    }
+    Ok(())
+}
+
+/// Identifies a hit by the value it is ordered on, so documents that share an 
`_id` (custom
+/// routing does not make `_id` unique across shards) stay distinct once 
`tiebreaker_field`
+/// names a unique field.
+fn hit_key(
+    hit: &Value,
+    sort: &[Value],
+    default_index: &str,
+    tiebreaker_field: &str,
+) -> Result<HitKey, Error> {
+    let value = if tiebreaker_field == DEFAULT_TIEBREAKER_FIELD {
+        hit.get("_id")
+    } else {
+        sort.get(TIEBREAKER_SORT_POSITION)
+    };
+    let id = value.and_then(Value::as_str).ok_or_else(|| {
+        Error::SchemaMismatch(format!(
+            "OpenSearch hit has no {tiebreaker_field} value; it cannot be told 
apart from \
+             published documents"
+        ))
+    })?;
+    let index = hit
+        .get("_index")
+        .and_then(Value::as_str)
+        .unwrap_or(default_index);
+    Ok(HitKey {
+        index: index.to_owned(),
+        id: id.to_owned(),
+    })
+}
+
+// Lowest possible tuple at the watermark, so the scan sees every hit with 
that sort value.
+fn restart_cursor(watermark: &Value) -> Vec<Value> {
+    std::iter::once(watermark.clone())
+        .chain(std::iter::repeat_n(json!(""), SORT_KEY_COUNT - 1))
+        .collect()
+}
+
+// A partial page can omit hits sorted before the last returned one. The 
cursor would then move
+// past them for good, so such a response must fail the poll before any cursor 
update.
+fn ensure_complete_response(body: &Value) -> Result<(), Error> {
+    if body.get("timed_out").and_then(Value::as_bool) == Some(true) {
+        return Err(Error::HttpRequestFailed(
+            "OpenSearch search timed out; results may be partial".to_string(),
+        ));
+    }
+    let failed_shards = body
+        .get("_shards")
+        .and_then(|shards| shards.get("failed"))
+        .and_then(Value::as_u64)
+        .unwrap_or(0);
+    if failed_shards > 0 {
+        return Err(Error::HttpRequestFailed(format!(
+            "OpenSearch search failed on {failed_shards} shards; results may 
be partial"
+        )));
+    }
+    Ok(())
+}
+
+fn restore_state(id: u32, state: Option<ConnectorState>) -> (State, 
Option<String>) {

Review Comment:
   Fixed



-- 
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]

Reply via email to