1use 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
11pub(crate) const NO_BATCHING_SEARCH_SIZE: usize = 10_000;
15
16pub struct ElasticsearchSource {
18 config: ElasticsearchSourceConfig,
19 client: Client,
20 auth_provider: Option<SharedAuthProvider>,
24}
25
26impl ElasticsearchSource {
27 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 pub fn with_auth_provider(mut self, provider: SharedAuthProvider) -> Self {
45 self.auth_provider = Some(provider);
46 self
47 }
48
49 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 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 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 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 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 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 let auth = self.resolve_auth().await?;
160
161 let mut all_records = Vec::new();
162
163 let page_size = if self.config.batch_size == 0 {
167 NO_BATCHING_SEARCH_SIZE
168 } else {
169 self.config.batch_size
170 };
171
172 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 while let Some(ref sid) = scroll_id {
198 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 if page_records.is_empty() {
228 break;
229 }
230
231 scroll_id = Self::extract_scroll_id(&body);
233 all_records.append(&mut page_records);
234 }
235
236 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 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 let auth = self.resolve_auth().await?;
285
286 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 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 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 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 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 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 async fn discover(&self) -> Result<Vec<faucet_core::DatasetDescriptor>, FaucetError> {
421 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
459fn 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
473fn 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
492fn 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
520fn 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
545struct 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 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 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
618fn 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 #[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 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 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}