Skip to main content

faucet_source_elasticsearch/
stream.rs

1//! Elasticsearch scroll-based search source.
2
3use crate::config::{ElasticsearchAuth, ElasticsearchSourceConfig};
4use async_trait::async_trait;
5use faucet_core::util::{DEFAULT_ERROR_BODY_MAX_LEN, check_http_response};
6use faucet_core::{AuthSpec, FaucetError, SharedAuthProvider, Stream, StreamPage};
7use reqwest::Client;
8use serde_json::{Value, json};
9use std::pin::Pin;
10
11/// `size` used by the [`Source::stream_pages`] non-scroll fallback (when
12/// `batch_size = 0`). Mirrors Elasticsearch's default `index.max_result_window`
13/// so the request stays within ES's out-of-the-box cap.
14pub(crate) const NO_BATCHING_SEARCH_SIZE: usize = 10_000;
15
16/// A source that reads documents from an Elasticsearch index using the scroll API.
17pub struct ElasticsearchSource {
18    config: ElasticsearchSourceConfig,
19    client: Client,
20    /// Optional shared auth provider. When set it takes precedence over inline
21    /// auth. Injected by the CLI (to resolve `auth: { ref }`) or directly by
22    /// library callers who want to share one token across multiple sources.
23    auth_provider: Option<SharedAuthProvider>,
24}
25
26impl ElasticsearchSource {
27    /// Create a new Elasticsearch source from the given configuration.
28    /// Construction does no I/O; it fails only on an invalid config (an
29    /// out-of-range `batch_size`).
30    pub fn new(config: ElasticsearchSourceConfig) -> Result<Self, FaucetError> {
31        faucet_core::validate_batch_size(config.batch_size)?;
32        Ok(Self {
33            config,
34            client: Client::new(),
35            auth_provider: None,
36        })
37    }
38
39    /// Attach a shared [`AuthProvider`](faucet_core::AuthProvider). When set,
40    /// the provider supplies the credential for every request (taking precedence
41    /// over inline auth), so several sources can share one token with
42    /// single-flight refresh. Used by the CLI to resolve `auth: { ref }`, and
43    /// by library callers who inject one provider into many sources.
44    pub fn with_auth_provider(mut self, provider: SharedAuthProvider) -> Self {
45        self.auth_provider = Some(provider);
46        self
47    }
48
49    /// Resolve the effective [`ElasticsearchAuth`] for the current request.
50    ///
51    /// Resolution order:
52    /// 1. If a shared provider is attached, call it and map the credential.
53    /// 2. Otherwise use the inline auth from config.
54    /// 3. If the config is a `Reference` with no provider, return an error.
55    async fn resolve_auth(&self) -> Result<ElasticsearchAuth, FaucetError> {
56        if let Some(p) = &self.auth_provider {
57            return faucet_common_elasticsearch::credential_to_auth(p.credential().await?);
58        }
59        match &self.config.auth {
60            AuthSpec::Inline(a) => Ok(a.clone()),
61            AuthSpec::Reference(r) => Err(FaucetError::Auth(format!(
62                "auth references provider '{}' but no provider was supplied",
63                r.name
64            ))),
65        }
66    }
67
68    /// Apply an [`ElasticsearchAuth`] to a request builder.
69    fn apply_auth_value(
70        req: reqwest::RequestBuilder,
71        auth: &ElasticsearchAuth,
72    ) -> reqwest::RequestBuilder {
73        match auth {
74            ElasticsearchAuth::None => req,
75            ElasticsearchAuth::Basic { username, password } => {
76                req.basic_auth(username, Some(password))
77            }
78            ElasticsearchAuth::Bearer { token } => req.bearer_auth(token),
79            ElasticsearchAuth::ApiKey { key } => {
80                req.header("Authorization", format!("ApiKey {key}"))
81            }
82        }
83    }
84
85    /// Extract `hits.hits[*]._source` from an Elasticsearch search response.
86    fn extract_hits(body: &Value) -> Vec<Value> {
87        body.get("hits")
88            .and_then(|h| h.get("hits"))
89            .and_then(|h| h.as_array())
90            .map(|hits| {
91                hits.iter()
92                    .filter_map(|hit| hit.get("_source").cloned())
93                    .collect()
94            })
95            .unwrap_or_default()
96    }
97
98    /// Extract the `_scroll_id` from an Elasticsearch response.
99    fn extract_scroll_id(body: &Value) -> Option<String> {
100        body.get("_scroll_id")
101            .and_then(|v| v.as_str())
102            .map(|s| s.to_string())
103    }
104
105    /// Clear a scroll context. Best-effort: errors are logged but not propagated.
106    async fn clear_scroll(&self, scroll_id: &str) {
107        let url = format!("{}/_search/scroll", self.config.base_url);
108        let req = self
109            .client
110            .delete(&url)
111            .json(&json!({"scroll_id": scroll_id}));
112        let auth = match self.resolve_auth().await {
113            Ok(a) => a,
114            Err(e) => {
115                tracing::warn!(error = %e, "failed to resolve auth for scroll cleanup");
116                return;
117            }
118        };
119        let req = Self::apply_auth_value(req, &auth);
120
121        if let Err(e) = req.send().await {
122            tracing::warn!(error = %e, "failed to clear Elasticsearch scroll context");
123        }
124    }
125
126    /// Resolve the index and query under the supplied parent context. Returns
127    /// `(index, query)`.
128    fn resolve_index_and_query(
129        &self,
130        context: &std::collections::HashMap<String, Value>,
131    ) -> Result<(String, Value), FaucetError> {
132        let index = if context.is_empty() {
133            self.config.index.clone()
134        } else {
135            faucet_core::util::substitute_context(&self.config.index, context)
136        };
137        let query = if context.is_empty() {
138            self.config.query.clone()
139        } else {
140            let s = serde_json::to_string(&self.config.query)
141                .map_err(|e| FaucetError::Config(format!("failed to serialize query: {e}")))?;
142            let s = faucet_core::util::substitute_context_json(&s, context);
143            serde_json::from_str(&s).map_err(|e| {
144                FaucetError::Config(format!("failed to parse substituted query: {e}"))
145            })?
146        };
147        Ok((index, query))
148    }
149}
150
151#[async_trait]
152impl faucet_core::Source for ElasticsearchSource {
153    async fn fetch_with_context(
154        &self,
155        context: &std::collections::HashMap<String, serde_json::Value>,
156    ) -> Result<Vec<Value>, FaucetError> {
157        let (index, query) = self.resolve_index_and_query(context)?;
158        // Resolve auth once; reuse the same credential across all scroll pages.
159        let auth = self.resolve_auth().await?;
160
161        let mut all_records = Vec::new();
162
163        // `batch_size = 0` is the "no batching" sentinel. Interpolating it
164        // directly as `size=0` would make Elasticsearch return zero hits, so
165        // map it to the same large page size the streaming path uses (#78/#33).
166        let page_size = if self.config.batch_size == 0 {
167            NO_BATCHING_SEARCH_SIZE
168        } else {
169            self.config.batch_size
170        };
171
172        // Initial search request with scroll.
173        let url = format!(
174            "{}/{}/_search?scroll={}&size={}",
175            self.config.base_url, index, self.config.scroll_timeout, page_size
176        );
177        let req = self.client.post(&url).json(&json!({"query": query}));
178        let req = Self::apply_auth_value(req, &auth);
179
180        let resp = req.send().await?;
181        let resp = check_http_response(resp, DEFAULT_ERROR_BODY_MAX_LEN).await?;
182        let body: Value = resp.json().await?;
183
184        let mut records = Self::extract_hits(&body);
185        let mut scroll_id = Self::extract_scroll_id(&body);
186        let mut pages_fetched: usize = 1;
187
188        tracing::debug!(
189            records = records.len(),
190            page = pages_fetched,
191            "Elasticsearch initial search"
192        );
193
194        all_records.append(&mut records);
195
196        // Scroll loop.
197        while let Some(ref sid) = scroll_id {
198            // Check max_pages limit.
199            if let Some(max) = self.config.max_pages
200                && pages_fetched >= max
201            {
202                tracing::debug!(max_pages = max, "max_pages reached, stopping scroll");
203                break;
204            }
205
206            let scroll_url = format!("{}/_search/scroll", self.config.base_url);
207            let req = self.client.post(&scroll_url).json(&json!({
208                "scroll": self.config.scroll_timeout,
209                "scroll_id": sid,
210            }));
211            let req = Self::apply_auth_value(req, &auth);
212
213            let resp = req.send().await?;
214            let resp = check_http_response(resp, DEFAULT_ERROR_BODY_MAX_LEN).await?;
215            let body: Value = resp.json().await?;
216
217            let mut page_records = Self::extract_hits(&body);
218            pages_fetched += 1;
219
220            tracing::debug!(
221                records = page_records.len(),
222                page = pages_fetched,
223                "Elasticsearch scroll page"
224            );
225
226            // Stop when no more hits are returned.
227            if page_records.is_empty() {
228                break;
229            }
230
231            // Update scroll_id for the next iteration.
232            scroll_id = Self::extract_scroll_id(&body);
233            all_records.append(&mut page_records);
234        }
235
236        // Clear the scroll context (best-effort).
237        if let Some(ref sid) = scroll_id {
238            self.clear_scroll(sid).await;
239        }
240
241        tracing::debug!(
242            total_records = all_records.len(),
243            pages = pages_fetched,
244            "Elasticsearch fetch complete"
245        );
246
247        Ok(all_records)
248    }
249
250    /// Stream documents from Elasticsearch as scroll pages, one
251    /// [`StreamPage`] per scroll response. Bounds client-side memory at
252    /// O(batch_size) regardless of the index's total document count.
253    ///
254    /// The trait-level `batch_size` argument is ignored in favour of
255    /// [`ElasticsearchSourceConfig::batch_size`] — the config is the
256    /// user-facing knob the README documents, and routing the
257    /// pipeline-supplied hint through it would silently override an explicit
258    /// config value.
259    ///
260    /// When `batch_size = 0` the source issues a single non-scroll
261    /// `_search?size=10_000` and emits exactly one page. The scroll API is
262    /// not used and no scroll context needs to be cleared.
263    ///
264    /// The Elasticsearch search source has no incremental-replication mode
265    /// today, so every emitted page carries `bookmark: None`.
266    ///
267    /// **Scroll-context cleanup is mandatory.** On every exit path — clean
268    /// drain, `max_pages` truncation, mid-stream HTTP error, or consumer
269    /// dropping the stream — the open `_scroll_id` is sent to
270    /// `DELETE _search/scroll` so the cluster does not leak server-side
271    /// state. Cleanup runs inside a guard whose `Drop` impl spawns the
272    /// delete request, so even cancellation at any `.await` point still
273    /// releases the context.
274    fn stream_pages<'a>(
275        &'a self,
276        context: &'a std::collections::HashMap<String, Value>,
277        _batch_size: usize,
278    ) -> Pin<Box<dyn Stream<Item = Result<StreamPage, FaucetError>> + Send + 'a>> {
279        let batch_size = self.config.batch_size;
280
281        Box::pin(async_stream::try_stream! {
282            let (index, query) = self.resolve_index_and_query(context)?;
283            // Resolve auth once; reuse across all scroll pages and cleanup.
284            let auth = self.resolve_auth().await?;
285
286            // batch_size == 0: single non-scroll _search with size = max_result_window default.
287            if batch_size == 0 {
288                let url = format!(
289                    "{}/{}/_search?size={}",
290                    self.config.base_url, index, NO_BATCHING_SEARCH_SIZE
291                );
292                let req = self.client.post(&url).json(&json!({"query": query}));
293                let req = Self::apply_auth_value(req, &auth);
294                let resp = req.send().await?;
295                let resp = check_http_response(resp, DEFAULT_ERROR_BODY_MAX_LEN).await?;
296                let body: Value = resp.json().await?;
297                let records = Self::extract_hits(&body);
298                tracing::info!(
299                    docs = records.len(),
300                    batch_size = 0,
301                    "Elasticsearch source stream complete (no-batching path)",
302                );
303                yield StreamPage { records, bookmark: None };
304                return;
305            }
306
307            // Scroll path. Wire up a guard so the scroll context is always
308            // cleared, even on early-return / error / drop.
309            // Pass the already-resolved auth so the guard's spawned cleanup
310            // tasks never need to call async auth resolution.
311            let mut guard = ScrollGuard::new(
312                self.config.base_url.clone(),
313                self.client.clone(),
314                auth.clone(),
315            );
316
317            let url = format!(
318                "{}/{}/_search?scroll={}&size={}",
319                self.config.base_url, index, self.config.scroll_timeout, batch_size
320            );
321            let req = self.client.post(&url).json(&json!({"query": query}));
322            let req = Self::apply_auth_value(req, &auth);
323            let resp = req.send().await?;
324            let resp = check_http_response(resp, DEFAULT_ERROR_BODY_MAX_LEN).await?;
325            let body: Value = resp.json().await?;
326
327            let records = Self::extract_hits(&body);
328            guard.update(Self::extract_scroll_id(&body));
329            let mut pages_emitted: usize = 0;
330            let mut total = records.len();
331
332            // The initial search always counts as page 1, even when it
333            // returns zero hits — emit it and move on.
334            pages_emitted += 1;
335            let is_final = records.is_empty()
336                || guard.scroll_id().is_none()
337                || matches!(self.config.max_pages, Some(max) if pages_emitted >= max);
338            yield StreamPage { records, bookmark: None };
339            if is_final {
340                guard.disarm_if_done();
341                tracing::info!(
342                    docs = total,
343                    pages = pages_emitted,
344                    batch_size,
345                    "Elasticsearch source stream complete",
346                );
347                return;
348            }
349
350            // Scroll loop.
351            while let Some(sid) = guard.scroll_id().map(|s| s.to_string()) {
352                let scroll_url = format!("{}/_search/scroll", self.config.base_url);
353                let req = self.client.post(&scroll_url).json(&json!({
354                    "scroll": self.config.scroll_timeout,
355                    "scroll_id": sid,
356                }));
357                let req = Self::apply_auth_value(req, &auth);
358                let resp = req.send().await?;
359                let resp = check_http_response(resp, DEFAULT_ERROR_BODY_MAX_LEN).await?;
360                let body: Value = resp.json().await?;
361
362                let records = Self::extract_hits(&body);
363                guard.update(Self::extract_scroll_id(&body));
364                pages_emitted += 1;
365                total += records.len();
366
367                let is_empty = records.is_empty();
368                let hit_cap = matches!(self.config.max_pages, Some(max) if pages_emitted >= max);
369
370                if is_empty {
371                    // Final empty page — ES uses an empty hits array as the
372                    // end-of-scroll sentinel. Drop it; nothing to emit.
373                    break;
374                }
375
376                yield StreamPage { records, bookmark: None };
377
378                if hit_cap {
379                    tracing::debug!(
380                        max_pages = self.config.max_pages.unwrap_or(0),
381                        "max_pages reached, stopping scroll"
382                    );
383                    break;
384                }
385            }
386
387            tracing::info!(
388                docs = total,
389                pages = pages_emitted,
390                batch_size,
391                "Elasticsearch source stream complete",
392            );
393
394            // Successful drain — let the guard clean up the scroll id (if any).
395            guard.disarm_if_done();
396        })
397    }
398
399    fn config_schema(&self) -> serde_json::Value {
400        serde_json::to_value(faucet_core::schema_for!(ElasticsearchSourceConfig))
401            .expect("schema serialization")
402    }
403
404    fn dataset_uri(&self) -> String {
405        format!(
406            "{}/{}",
407            faucet_core::redact_uri_credentials(&self.config.base_url),
408            self.config.index
409        )
410    }
411
412    fn supports_discover(&self) -> bool {
413        true
414    }
415
416    /// Enumerate the cluster's non-system indices via
417    /// `GET _cat/indices?format=json` (names + doc counts) and
418    /// `GET /<index>/_mapping` (field types). Catalog metadata only — no
419    /// document scan.
420    async fn discover(&self) -> Result<Vec<faucet_core::DatasetDescriptor>, FaucetError> {
421        // Resolve auth once; reuse across the _cat and _mapping requests.
422        let auth = self.resolve_auth().await?;
423
424        let url = format!(
425            "{}/_cat/indices?format=json&h=index,docs.count",
426            self.config.base_url
427        );
428        let req = Self::apply_auth_value(self.client.get(&url), &auth);
429        let resp = req.send().await.map_err(|e| {
430            FaucetError::Source(format!("elasticsearch: catalog discovery failed: {e}"))
431        })?;
432        let resp = check_http_response(resp, DEFAULT_ERROR_BODY_MAX_LEN).await?;
433        let cat: Value = resp.json().await.map_err(|e| {
434            FaucetError::Source(format!("elasticsearch: catalog discovery failed: {e}"))
435        })?;
436
437        let entries = parse_cat_indices(&cat);
438        let mut datasets = Vec::with_capacity(entries.len());
439        for (index, doc_count) in entries {
440            let url = format!("{}/{}/_mapping", self.config.base_url, index);
441            let req = Self::apply_auth_value(self.client.get(&url), &auth);
442            let resp = req.send().await.map_err(|e| {
443                FaucetError::Source(format!(
444                    "elasticsearch: catalog discovery failed (mapping for {index:?}): {e}"
445                ))
446            })?;
447            let resp = check_http_response(resp, DEFAULT_ERROR_BODY_MAX_LEN).await?;
448            let body: Value = resp.json().await.map_err(|e| {
449                FaucetError::Source(format!(
450                    "elasticsearch: catalog discovery failed (mapping for {index:?}): {e}"
451                ))
452            })?;
453            datasets.push(descriptor_for_index(&index, doc_count, &body));
454        }
455        Ok(datasets)
456    }
457}
458
459/// Map an Elasticsearch mapping field type to a JSON-Schema type. ES mappings
460/// carry no nullability, so types stay scalar (never nullable-wrapped).
461/// Everything not listed (text, keyword, date, ip, …) serializes as a JSON
462/// string in `_source`, so `string` is the safe over-approximation.
463fn es_type_to_json_type(es_type: &str) -> &'static str {
464    match es_type {
465        "long" | "integer" | "short" | "byte" => "integer",
466        "double" | "float" | "half_float" | "scaled_float" => "number",
467        "boolean" => "boolean",
468        "object" | "nested" => "object",
469        _ => "string",
470    }
471}
472
473/// Convert one index's `mappings` object (`{"properties": {field: spec, …}}`)
474/// into an [`infer_schema`](faucet_core::schema::infer_schema)-shaped object
475/// schema at top-level-column granularity. A field spec without a scalar
476/// `type` is an object field (Elasticsearch's default for mapping entries
477/// that carry only nested `properties`). Pure.
478fn mapping_to_schema(mappings: &Value) -> Value {
479    let mut properties = serde_json::Map::new();
480    if let Some(fields) = mappings.get("properties").and_then(Value::as_object) {
481        for (name, spec) in fields {
482            let ty = match spec.get("type").and_then(Value::as_str) {
483                Some(t) => es_type_to_json_type(t),
484                None => "object",
485            };
486            properties.insert(name.clone(), json!({ "type": ty }));
487        }
488    }
489    json!({ "type": "object", "properties": Value::Object(properties) })
490}
491
492/// Parse a `GET _cat/indices?format=json` response into `(index, doc_count)`
493/// pairs, skipping system indices (leading `.`). Doc counts arrive as strings
494/// in `_cat` JSON — an unparsable or missing count yields `None`. Sorted by
495/// index name for deterministic output. Pure.
496fn parse_cat_indices(cat: &Value) -> Vec<(String, Option<u64>)> {
497    let mut entries: Vec<(String, Option<u64>)> = cat
498        .as_array()
499        .map(|rows| {
500            rows.iter()
501                .filter_map(|row| {
502                    let index = row.get("index")?.as_str()?;
503                    if index.starts_with('.') {
504                        return None;
505                    }
506                    let doc_count = row.get("docs.count").and_then(|v| {
507                        v.as_str()
508                            .and_then(|s| s.parse::<u64>().ok())
509                            .or_else(|| v.as_u64())
510                    });
511                    Some((index.to_string(), doc_count))
512                })
513                .collect()
514        })
515        .unwrap_or_default();
516    entries.sort_by(|a, b| a.0.cmp(&b.0));
517    entries
518}
519
520/// Build the [`DatasetDescriptor`](faucet_core::DatasetDescriptor) for one
521/// index from its `_cat` row and its `GET /<index>/_mapping` response body
522/// (`{"<index>": {"mappings": {…}}}`). The body's top-level key is the
523/// concrete index name, which can differ from the requested one when the
524/// request resolved through an alias — fall back to the first entry. Pure.
525fn descriptor_for_index(
526    index: &str,
527    doc_count: Option<u64>,
528    mapping_body: &Value,
529) -> faucet_core::DatasetDescriptor {
530    let empty = json!({});
531    let mappings = mapping_body
532        .get(index)
533        .or_else(|| mapping_body.as_object().and_then(|o| o.values().next()))
534        .and_then(|entry| entry.get("mappings"))
535        .unwrap_or(&empty);
536    let mut descriptor =
537        faucet_core::DatasetDescriptor::new(index, "index", json!({ "index": index }))
538            .with_schema(mapping_to_schema(mappings));
539    if let Some(rows) = doc_count {
540        descriptor = descriptor.with_estimated_rows(rows);
541    }
542    descriptor
543}
544
545/// RAII guard that owns the active scroll id and clears it on drop.
546///
547/// Holds a pre-resolved [`ElasticsearchAuth`] (not `AuthSpec`) so the drop-path
548/// spawned cleanup tasks never need to perform async auth resolution.
549struct ScrollGuard {
550    base_url: String,
551    client: Client,
552    auth: ElasticsearchAuth,
553    scroll_id: Option<String>,
554}
555
556impl ScrollGuard {
557    fn new(base_url: String, client: Client, auth: ElasticsearchAuth) -> Self {
558        Self {
559            base_url,
560            client,
561            auth,
562            scroll_id: None,
563        }
564    }
565
566    fn scroll_id(&self) -> Option<&str> {
567        self.scroll_id.as_deref()
568    }
569
570    fn update(&mut self, new_id: Option<String>) {
571        if let Some(id) = new_id {
572            self.scroll_id = Some(id);
573        }
574    }
575
576    /// Called when the stream drained cleanly. Spawns cleanup as a detached
577    /// task and disarms the drop fallback.
578    fn disarm_if_done(&mut self) {
579        if let Some(sid) = self.scroll_id.take() {
580            let base_url = self.base_url.clone();
581            let auth = self.auth.clone();
582            let client = self.client.clone();
583            tokio::spawn(async move {
584                let url = format!("{base_url}/_search/scroll");
585                let req = client.delete(&url).json(&json!({"scroll_id": sid}));
586                let req = apply_auth_to(req, &auth);
587                if let Err(e) = req.send().await {
588                    tracing::warn!(error = %e, "failed to clear Elasticsearch scroll context");
589                }
590            });
591        }
592    }
593}
594
595impl Drop for ScrollGuard {
596    fn drop(&mut self) {
597        if let Some(sid) = self.scroll_id.take() {
598            // Error / cancellation path. Spawn so cleanup survives the
599            // stream future being dropped mid-await.
600            let base_url = self.base_url.clone();
601            let auth = self.auth.clone();
602            let client = self.client.clone();
603            tokio::spawn(async move {
604                let url = format!("{base_url}/_search/scroll");
605                let req = client.delete(&url).json(&json!({"scroll_id": sid}));
606                let req = apply_auth_to(req, &auth);
607                if let Err(e) = req.send().await {
608                    tracing::warn!(
609                        error = %e,
610                        "failed to clear Elasticsearch scroll context (drop path)",
611                    );
612                }
613            });
614        }
615    }
616}
617
618/// Apply an [`ElasticsearchAuth`] to a request builder. Standalone so
619/// spawned cleanup tasks can use it without holding a source reference.
620fn apply_auth_to(
621    req: reqwest::RequestBuilder,
622    auth: &ElasticsearchAuth,
623) -> reqwest::RequestBuilder {
624    match auth {
625        ElasticsearchAuth::None => req,
626        ElasticsearchAuth::Basic { username, password } => req.basic_auth(username, Some(password)),
627        ElasticsearchAuth::Bearer { token } => req.bearer_auth(token),
628        ElasticsearchAuth::ApiKey { key } => req.header("Authorization", format!("ApiKey {key}")),
629    }
630}
631
632#[cfg(test)]
633mod tests {
634    use super::*;
635    use faucet_core::Source;
636
637    #[test]
638    fn new_rejects_out_of_range_batch_size() {
639        let mut config = ElasticsearchSourceConfig::new("http://localhost:9200", "idx");
640        config.batch_size = faucet_core::MAX_BATCH_SIZE + 1;
641        match ElasticsearchSource::new(config) {
642            Err(FaucetError::Config(m)) => assert!(m.contains("batch_size"), "got: {m}"),
643            _ => panic!("expected a batch_size Config error"),
644        }
645    }
646
647    #[test]
648    fn dataset_uri_returns_base_url_slash_index() {
649        let config = ElasticsearchSourceConfig::new("http://localhost:9200", "my_index");
650        let source = ElasticsearchSource::new(config).unwrap();
651        assert_eq!(source.dataset_uri(), "http://localhost:9200/my_index");
652    }
653
654    #[test]
655    fn dataset_uri_strips_credentials() {
656        let config =
657            ElasticsearchSourceConfig::new("http://user:secret@es.example.com:9200", "logs");
658        let source = ElasticsearchSource::new(config).unwrap();
659        assert_eq!(source.dataset_uri(), "http://es.example.com:9200/logs");
660    }
661
662    // ── discover: pure mapping/_cat parsing (#211) ──────────────────────────
663
664    #[test]
665    fn es_types_map_to_json_types() {
666        for (es, want) in [
667            ("long", "integer"),
668            ("integer", "integer"),
669            ("short", "integer"),
670            ("byte", "integer"),
671            ("double", "number"),
672            ("float", "number"),
673            ("half_float", "number"),
674            ("scaled_float", "number"),
675            ("boolean", "boolean"),
676            ("object", "object"),
677            ("nested", "object"),
678            ("text", "string"),
679            ("keyword", "string"),
680            ("date", "string"),
681            ("ip", "string"),
682            ("geo_point", "string"),
683        ] {
684            assert_eq!(es_type_to_json_type(es), want, "for ES type {es:?}");
685        }
686    }
687
688    #[test]
689    fn mapping_to_schema_covers_scalar_object_and_nested_fields() {
690        // Real-shaped mapping: scalar types, a text field with a keyword
691        // sub-field, and an object field declared only via nested properties.
692        let mappings = json!({
693            "properties": {
694                "id": {"type": "long"},
695                "total": {"type": "scaled_float", "scaling_factor": 100},
696                "note": {"type": "text", "fields": {"keyword": {"type": "keyword"}}},
697                "active": {"type": "boolean"},
698                "customer": {"properties": {"name": {"type": "text"}}},
699                "meta": {"type": "nested", "properties": {"k": {"type": "keyword"}}},
700            }
701        });
702        let schema = mapping_to_schema(&mappings);
703        assert_eq!(schema["type"], "object");
704        let props = &schema["properties"];
705        assert_eq!(props["id"]["type"], "integer");
706        assert_eq!(props["total"]["type"], "number");
707        assert_eq!(props["note"]["type"], "string");
708        assert_eq!(props["active"]["type"], "boolean");
709        assert_eq!(
710            props["customer"]["type"], "object",
711            "type-less field with nested properties is an object"
712        );
713        assert_eq!(props["meta"]["type"], "object");
714    }
715
716    #[test]
717    fn mapping_to_schema_empty_mappings_yield_empty_properties() {
718        assert_eq!(
719            mapping_to_schema(&json!({})),
720            json!({"type": "object", "properties": {}})
721        );
722        assert_eq!(
723            mapping_to_schema(&Value::Null),
724            json!({"type": "object", "properties": {}})
725        );
726    }
727
728    #[test]
729    fn parse_cat_indices_skips_system_and_parses_counts() {
730        let cat = json!([
731            {"index": "orders", "docs.count": "1200"},
732            {"index": ".kibana_1", "docs.count": "3"},
733            {"index": "logs", "docs.count": "n/a"},
734            {"index": "metrics"},
735            {"index": "numeric", "docs.count": 7},
736        ]);
737        let entries = parse_cat_indices(&cat);
738        assert_eq!(
739            entries,
740            vec![
741                ("logs".to_string(), None),
742                ("metrics".to_string(), None),
743                ("numeric".to_string(), Some(7)),
744                ("orders".to_string(), Some(1200)),
745            ],
746            "system index skipped, unparsable/missing counts → None, sorted"
747        );
748    }
749
750    #[test]
751    fn parse_cat_indices_non_array_is_empty() {
752        assert!(parse_cat_indices(&json!({"error": "nope"})).is_empty());
753        assert!(parse_cat_indices(&Value::Null).is_empty());
754    }
755
756    #[test]
757    fn descriptor_for_index_builds_full_descriptor() {
758        let body = json!({
759            "orders": {"mappings": {"properties": {"id": {"type": "long"}}}}
760        });
761        let d = descriptor_for_index("orders", Some(1200), &body);
762        assert_eq!(d.name, "orders");
763        assert_eq!(d.kind, "index");
764        assert_eq!(d.config_patch, json!({"index": "orders"}));
765        assert_eq!(d.estimated_rows, Some(1200));
766        let schema = d.schema.as_ref().expect("schema");
767        assert_eq!(schema["properties"]["id"]["type"], "integer");
768    }
769
770    #[test]
771    fn descriptor_for_index_falls_back_to_first_mapping_entry() {
772        // The _mapping response is keyed by the *concrete* index name, which
773        // differs from the requested name when it resolved via an alias.
774        let body = json!({
775            "orders-000001": {"mappings": {"properties": {"id": {"type": "long"}}}}
776        });
777        let d = descriptor_for_index("orders", None, &body);
778        assert_eq!(d.estimated_rows, None);
779        let schema = d.schema.as_ref().expect("schema");
780        assert_eq!(schema["properties"]["id"]["type"], "integer");
781    }
782
783    #[test]
784    fn descriptor_for_index_missing_mappings_yields_empty_schema() {
785        let d = descriptor_for_index("orders", Some(0), &json!({}));
786        assert_eq!(
787            d.schema,
788            Some(json!({"type": "object", "properties": {}})),
789            "no mappings → empty object schema, never a panic"
790        );
791    }
792
793    #[test]
794    fn source_advertises_discover() {
795        let config = ElasticsearchSourceConfig::new("http://localhost:9200", "idx");
796        let source = ElasticsearchSource::new(config).unwrap();
797        assert!(source.supports_discover());
798    }
799}