use crate::config::{ElasticsearchAuth, ElasticsearchSourceConfig};
use async_trait::async_trait;
use faucet_core::util::{DEFAULT_ERROR_BODY_MAX_LEN, check_http_response};
use faucet_core::{AuthSpec, FaucetError, SharedAuthProvider, Stream, StreamPage};
use reqwest::Client;
use serde_json::{Value, json};
use std::pin::Pin;
pub(crate) const NO_BATCHING_SEARCH_SIZE: usize = 10_000;
pub struct ElasticsearchSource {
config: ElasticsearchSourceConfig,
client: Client,
auth_provider: Option<SharedAuthProvider>,
}
impl ElasticsearchSource {
pub fn new(config: ElasticsearchSourceConfig) -> Result<Self, FaucetError> {
faucet_core::validate_batch_size(config.batch_size)?;
Ok(Self {
config,
client: Client::new(),
auth_provider: None,
})
}
pub fn with_auth_provider(mut self, provider: SharedAuthProvider) -> Self {
self.auth_provider = Some(provider);
self
}
async fn resolve_auth(&self) -> Result<ElasticsearchAuth, FaucetError> {
if let Some(p) = &self.auth_provider {
return faucet_common_elasticsearch::credential_to_auth(p.credential().await?);
}
match &self.config.auth {
AuthSpec::Inline(a) => Ok(a.clone()),
AuthSpec::Reference(r) => Err(FaucetError::Auth(format!(
"auth references provider '{}' but no provider was supplied",
r.name
))),
}
}
fn apply_auth_value(
req: reqwest::RequestBuilder,
auth: &ElasticsearchAuth,
) -> reqwest::RequestBuilder {
match auth {
ElasticsearchAuth::None => req,
ElasticsearchAuth::Basic { username, password } => {
req.basic_auth(username, Some(password))
}
ElasticsearchAuth::Bearer { token } => req.bearer_auth(token),
ElasticsearchAuth::ApiKey { key } => {
req.header("Authorization", format!("ApiKey {key}"))
}
}
}
fn extract_hits(body: &Value) -> Vec<Value> {
body.get("hits")
.and_then(|h| h.get("hits"))
.and_then(|h| h.as_array())
.map(|hits| {
hits.iter()
.filter_map(|hit| hit.get("_source").cloned())
.collect()
})
.unwrap_or_default()
}
fn extract_scroll_id(body: &Value) -> Option<String> {
body.get("_scroll_id")
.and_then(|v| v.as_str())
.map(|s| s.to_string())
}
async fn clear_scroll(&self, scroll_id: &str) {
let url = format!("{}/_search/scroll", self.config.base_url);
let req = self
.client
.delete(&url)
.json(&json!({"scroll_id": scroll_id}));
let auth = match self.resolve_auth().await {
Ok(a) => a,
Err(e) => {
tracing::warn!(error = %e, "failed to resolve auth for scroll cleanup");
return;
}
};
let req = Self::apply_auth_value(req, &auth);
if let Err(e) = req.send().await {
tracing::warn!(error = %e, "failed to clear Elasticsearch scroll context");
}
}
fn resolve_index_and_query(
&self,
context: &std::collections::HashMap<String, Value>,
) -> Result<(String, Value), FaucetError> {
let index = if context.is_empty() {
self.config.index.clone()
} else {
faucet_core::util::substitute_context(&self.config.index, context)
};
let query = if context.is_empty() {
self.config.query.clone()
} else {
let s = serde_json::to_string(&self.config.query)
.map_err(|e| FaucetError::Config(format!("failed to serialize query: {e}")))?;
let s = faucet_core::util::substitute_context_json(&s, context);
serde_json::from_str(&s).map_err(|e| {
FaucetError::Config(format!("failed to parse substituted query: {e}"))
})?
};
Ok((index, query))
}
}
#[async_trait]
impl faucet_core::Source for ElasticsearchSource {
async fn fetch_with_context(
&self,
context: &std::collections::HashMap<String, serde_json::Value>,
) -> Result<Vec<Value>, FaucetError> {
let (index, query) = self.resolve_index_and_query(context)?;
let auth = self.resolve_auth().await?;
let mut all_records = Vec::new();
let page_size = if self.config.batch_size == 0 {
NO_BATCHING_SEARCH_SIZE
} else {
self.config.batch_size
};
let url = format!(
"{}/{}/_search?scroll={}&size={}",
self.config.base_url, index, self.config.scroll_timeout, page_size
);
let req = self.client.post(&url).json(&json!({"query": query}));
let req = Self::apply_auth_value(req, &auth);
let resp = req.send().await?;
let resp = check_http_response(resp, DEFAULT_ERROR_BODY_MAX_LEN).await?;
let body: Value = resp.json().await?;
let mut records = Self::extract_hits(&body);
let mut scroll_id = Self::extract_scroll_id(&body);
let mut pages_fetched: usize = 1;
tracing::debug!(
records = records.len(),
page = pages_fetched,
"Elasticsearch initial search"
);
all_records.append(&mut records);
while let Some(ref sid) = scroll_id {
if let Some(max) = self.config.max_pages
&& pages_fetched >= max
{
tracing::debug!(max_pages = max, "max_pages reached, stopping scroll");
break;
}
let scroll_url = format!("{}/_search/scroll", self.config.base_url);
let req = self.client.post(&scroll_url).json(&json!({
"scroll": self.config.scroll_timeout,
"scroll_id": sid,
}));
let req = Self::apply_auth_value(req, &auth);
let resp = req.send().await?;
let resp = check_http_response(resp, DEFAULT_ERROR_BODY_MAX_LEN).await?;
let body: Value = resp.json().await?;
let mut page_records = Self::extract_hits(&body);
pages_fetched += 1;
tracing::debug!(
records = page_records.len(),
page = pages_fetched,
"Elasticsearch scroll page"
);
if page_records.is_empty() {
break;
}
scroll_id = Self::extract_scroll_id(&body);
all_records.append(&mut page_records);
}
if let Some(ref sid) = scroll_id {
self.clear_scroll(sid).await;
}
tracing::debug!(
total_records = all_records.len(),
pages = pages_fetched,
"Elasticsearch fetch complete"
);
Ok(all_records)
}
fn stream_pages<'a>(
&'a self,
context: &'a std::collections::HashMap<String, Value>,
_batch_size: usize,
) -> Pin<Box<dyn Stream<Item = Result<StreamPage, FaucetError>> + Send + 'a>> {
let batch_size = self.config.batch_size;
Box::pin(async_stream::try_stream! {
let (index, query) = self.resolve_index_and_query(context)?;
let auth = self.resolve_auth().await?;
if batch_size == 0 {
let url = format!(
"{}/{}/_search?size={}",
self.config.base_url, index, NO_BATCHING_SEARCH_SIZE
);
let req = self.client.post(&url).json(&json!({"query": query}));
let req = Self::apply_auth_value(req, &auth);
let resp = req.send().await?;
let resp = check_http_response(resp, DEFAULT_ERROR_BODY_MAX_LEN).await?;
let body: Value = resp.json().await?;
let records = Self::extract_hits(&body);
tracing::info!(
docs = records.len(),
batch_size = 0,
"Elasticsearch source stream complete (no-batching path)",
);
yield StreamPage { records, bookmark: None };
return;
}
let mut guard = ScrollGuard::new(
self.config.base_url.clone(),
self.client.clone(),
auth.clone(),
);
let url = format!(
"{}/{}/_search?scroll={}&size={}",
self.config.base_url, index, self.config.scroll_timeout, batch_size
);
let req = self.client.post(&url).json(&json!({"query": query}));
let req = Self::apply_auth_value(req, &auth);
let resp = req.send().await?;
let resp = check_http_response(resp, DEFAULT_ERROR_BODY_MAX_LEN).await?;
let body: Value = resp.json().await?;
let records = Self::extract_hits(&body);
guard.update(Self::extract_scroll_id(&body));
let mut pages_emitted: usize = 0;
let mut total = records.len();
pages_emitted += 1;
let is_final = records.is_empty()
|| guard.scroll_id().is_none()
|| matches!(self.config.max_pages, Some(max) if pages_emitted >= max);
yield StreamPage { records, bookmark: None };
if is_final {
guard.disarm_if_done();
tracing::info!(
docs = total,
pages = pages_emitted,
batch_size,
"Elasticsearch source stream complete",
);
return;
}
while let Some(sid) = guard.scroll_id().map(|s| s.to_string()) {
let scroll_url = format!("{}/_search/scroll", self.config.base_url);
let req = self.client.post(&scroll_url).json(&json!({
"scroll": self.config.scroll_timeout,
"scroll_id": sid,
}));
let req = Self::apply_auth_value(req, &auth);
let resp = req.send().await?;
let resp = check_http_response(resp, DEFAULT_ERROR_BODY_MAX_LEN).await?;
let body: Value = resp.json().await?;
let records = Self::extract_hits(&body);
guard.update(Self::extract_scroll_id(&body));
pages_emitted += 1;
total += records.len();
let is_empty = records.is_empty();
let hit_cap = matches!(self.config.max_pages, Some(max) if pages_emitted >= max);
if is_empty {
break;
}
yield StreamPage { records, bookmark: None };
if hit_cap {
tracing::debug!(
max_pages = self.config.max_pages.unwrap_or(0),
"max_pages reached, stopping scroll"
);
break;
}
}
tracing::info!(
docs = total,
pages = pages_emitted,
batch_size,
"Elasticsearch source stream complete",
);
guard.disarm_if_done();
})
}
fn config_schema(&self) -> serde_json::Value {
serde_json::to_value(faucet_core::schema_for!(ElasticsearchSourceConfig))
.expect("schema serialization")
}
fn dataset_uri(&self) -> String {
format!(
"{}/{}",
faucet_core::redact_uri_credentials(&self.config.base_url),
self.config.index
)
}
fn supports_discover(&self) -> bool {
true
}
async fn discover(&self) -> Result<Vec<faucet_core::DatasetDescriptor>, FaucetError> {
let auth = self.resolve_auth().await?;
let url = format!(
"{}/_cat/indices?format=json&h=index,docs.count",
self.config.base_url
);
let req = Self::apply_auth_value(self.client.get(&url), &auth);
let resp = req.send().await.map_err(|e| {
FaucetError::Source(format!("elasticsearch: catalog discovery failed: {e}"))
})?;
let resp = check_http_response(resp, DEFAULT_ERROR_BODY_MAX_LEN).await?;
let cat: Value = resp.json().await.map_err(|e| {
FaucetError::Source(format!("elasticsearch: catalog discovery failed: {e}"))
})?;
let entries = parse_cat_indices(&cat);
let mut datasets = Vec::with_capacity(entries.len());
for (index, doc_count) in entries {
let url = format!("{}/{}/_mapping", self.config.base_url, index);
let req = Self::apply_auth_value(self.client.get(&url), &auth);
let resp = req.send().await.map_err(|e| {
FaucetError::Source(format!(
"elasticsearch: catalog discovery failed (mapping for {index:?}): {e}"
))
})?;
let resp = check_http_response(resp, DEFAULT_ERROR_BODY_MAX_LEN).await?;
let body: Value = resp.json().await.map_err(|e| {
FaucetError::Source(format!(
"elasticsearch: catalog discovery failed (mapping for {index:?}): {e}"
))
})?;
datasets.push(descriptor_for_index(&index, doc_count, &body));
}
Ok(datasets)
}
}
fn es_type_to_json_type(es_type: &str) -> &'static str {
match es_type {
"long" | "integer" | "short" | "byte" => "integer",
"double" | "float" | "half_float" | "scaled_float" => "number",
"boolean" => "boolean",
"object" | "nested" => "object",
_ => "string",
}
}
fn mapping_to_schema(mappings: &Value) -> Value {
let mut properties = serde_json::Map::new();
if let Some(fields) = mappings.get("properties").and_then(Value::as_object) {
for (name, spec) in fields {
let ty = match spec.get("type").and_then(Value::as_str) {
Some(t) => es_type_to_json_type(t),
None => "object",
};
properties.insert(name.clone(), json!({ "type": ty }));
}
}
json!({ "type": "object", "properties": Value::Object(properties) })
}
fn parse_cat_indices(cat: &Value) -> Vec<(String, Option<u64>)> {
let mut entries: Vec<(String, Option<u64>)> = cat
.as_array()
.map(|rows| {
rows.iter()
.filter_map(|row| {
let index = row.get("index")?.as_str()?;
if index.starts_with('.') {
return None;
}
let doc_count = row.get("docs.count").and_then(|v| {
v.as_str()
.and_then(|s| s.parse::<u64>().ok())
.or_else(|| v.as_u64())
});
Some((index.to_string(), doc_count))
})
.collect()
})
.unwrap_or_default();
entries.sort_by(|a, b| a.0.cmp(&b.0));
entries
}
fn descriptor_for_index(
index: &str,
doc_count: Option<u64>,
mapping_body: &Value,
) -> faucet_core::DatasetDescriptor {
let empty = json!({});
let mappings = mapping_body
.get(index)
.or_else(|| mapping_body.as_object().and_then(|o| o.values().next()))
.and_then(|entry| entry.get("mappings"))
.unwrap_or(&empty);
let mut descriptor =
faucet_core::DatasetDescriptor::new(index, "index", json!({ "index": index }))
.with_schema(mapping_to_schema(mappings));
if let Some(rows) = doc_count {
descriptor = descriptor.with_estimated_rows(rows);
}
descriptor
}
struct ScrollGuard {
base_url: String,
client: Client,
auth: ElasticsearchAuth,
scroll_id: Option<String>,
}
impl ScrollGuard {
fn new(base_url: String, client: Client, auth: ElasticsearchAuth) -> Self {
Self {
base_url,
client,
auth,
scroll_id: None,
}
}
fn scroll_id(&self) -> Option<&str> {
self.scroll_id.as_deref()
}
fn update(&mut self, new_id: Option<String>) {
if let Some(id) = new_id {
self.scroll_id = Some(id);
}
}
fn disarm_if_done(&mut self) {
if let Some(sid) = self.scroll_id.take() {
let base_url = self.base_url.clone();
let auth = self.auth.clone();
let client = self.client.clone();
tokio::spawn(async move {
let url = format!("{base_url}/_search/scroll");
let req = client.delete(&url).json(&json!({"scroll_id": sid}));
let req = apply_auth_to(req, &auth);
if let Err(e) = req.send().await {
tracing::warn!(error = %e, "failed to clear Elasticsearch scroll context");
}
});
}
}
}
impl Drop for ScrollGuard {
fn drop(&mut self) {
if let Some(sid) = self.scroll_id.take() {
let base_url = self.base_url.clone();
let auth = self.auth.clone();
let client = self.client.clone();
tokio::spawn(async move {
let url = format!("{base_url}/_search/scroll");
let req = client.delete(&url).json(&json!({"scroll_id": sid}));
let req = apply_auth_to(req, &auth);
if let Err(e) = req.send().await {
tracing::warn!(
error = %e,
"failed to clear Elasticsearch scroll context (drop path)",
);
}
});
}
}
}
fn apply_auth_to(
req: reqwest::RequestBuilder,
auth: &ElasticsearchAuth,
) -> reqwest::RequestBuilder {
match auth {
ElasticsearchAuth::None => req,
ElasticsearchAuth::Basic { username, password } => req.basic_auth(username, Some(password)),
ElasticsearchAuth::Bearer { token } => req.bearer_auth(token),
ElasticsearchAuth::ApiKey { key } => req.header("Authorization", format!("ApiKey {key}")),
}
}
#[cfg(test)]
mod tests {
use super::*;
use faucet_core::Source;
#[test]
fn new_rejects_out_of_range_batch_size() {
let mut config = ElasticsearchSourceConfig::new("http://localhost:9200", "idx");
config.batch_size = faucet_core::MAX_BATCH_SIZE + 1;
match ElasticsearchSource::new(config) {
Err(FaucetError::Config(m)) => assert!(m.contains("batch_size"), "got: {m}"),
_ => panic!("expected a batch_size Config error"),
}
}
#[test]
fn dataset_uri_returns_base_url_slash_index() {
let config = ElasticsearchSourceConfig::new("http://localhost:9200", "my_index");
let source = ElasticsearchSource::new(config).unwrap();
assert_eq!(source.dataset_uri(), "http://localhost:9200/my_index");
}
#[test]
fn dataset_uri_strips_credentials() {
let config =
ElasticsearchSourceConfig::new("http://user:secret@es.example.com:9200", "logs");
let source = ElasticsearchSource::new(config).unwrap();
assert_eq!(source.dataset_uri(), "http://es.example.com:9200/logs");
}
#[test]
fn es_types_map_to_json_types() {
for (es, want) in [
("long", "integer"),
("integer", "integer"),
("short", "integer"),
("byte", "integer"),
("double", "number"),
("float", "number"),
("half_float", "number"),
("scaled_float", "number"),
("boolean", "boolean"),
("object", "object"),
("nested", "object"),
("text", "string"),
("keyword", "string"),
("date", "string"),
("ip", "string"),
("geo_point", "string"),
] {
assert_eq!(es_type_to_json_type(es), want, "for ES type {es:?}");
}
}
#[test]
fn mapping_to_schema_covers_scalar_object_and_nested_fields() {
let mappings = json!({
"properties": {
"id": {"type": "long"},
"total": {"type": "scaled_float", "scaling_factor": 100},
"note": {"type": "text", "fields": {"keyword": {"type": "keyword"}}},
"active": {"type": "boolean"},
"customer": {"properties": {"name": {"type": "text"}}},
"meta": {"type": "nested", "properties": {"k": {"type": "keyword"}}},
}
});
let schema = mapping_to_schema(&mappings);
assert_eq!(schema["type"], "object");
let props = &schema["properties"];
assert_eq!(props["id"]["type"], "integer");
assert_eq!(props["total"]["type"], "number");
assert_eq!(props["note"]["type"], "string");
assert_eq!(props["active"]["type"], "boolean");
assert_eq!(
props["customer"]["type"], "object",
"type-less field with nested properties is an object"
);
assert_eq!(props["meta"]["type"], "object");
}
#[test]
fn mapping_to_schema_empty_mappings_yield_empty_properties() {
assert_eq!(
mapping_to_schema(&json!({})),
json!({"type": "object", "properties": {}})
);
assert_eq!(
mapping_to_schema(&Value::Null),
json!({"type": "object", "properties": {}})
);
}
#[test]
fn parse_cat_indices_skips_system_and_parses_counts() {
let cat = json!([
{"index": "orders", "docs.count": "1200"},
{"index": ".kibana_1", "docs.count": "3"},
{"index": "logs", "docs.count": "n/a"},
{"index": "metrics"},
{"index": "numeric", "docs.count": 7},
]);
let entries = parse_cat_indices(&cat);
assert_eq!(
entries,
vec![
("logs".to_string(), None),
("metrics".to_string(), None),
("numeric".to_string(), Some(7)),
("orders".to_string(), Some(1200)),
],
"system index skipped, unparsable/missing counts → None, sorted"
);
}
#[test]
fn parse_cat_indices_non_array_is_empty() {
assert!(parse_cat_indices(&json!({"error": "nope"})).is_empty());
assert!(parse_cat_indices(&Value::Null).is_empty());
}
#[test]
fn descriptor_for_index_builds_full_descriptor() {
let body = json!({
"orders": {"mappings": {"properties": {"id": {"type": "long"}}}}
});
let d = descriptor_for_index("orders", Some(1200), &body);
assert_eq!(d.name, "orders");
assert_eq!(d.kind, "index");
assert_eq!(d.config_patch, json!({"index": "orders"}));
assert_eq!(d.estimated_rows, Some(1200));
let schema = d.schema.as_ref().expect("schema");
assert_eq!(schema["properties"]["id"]["type"], "integer");
}
#[test]
fn descriptor_for_index_falls_back_to_first_mapping_entry() {
let body = json!({
"orders-000001": {"mappings": {"properties": {"id": {"type": "long"}}}}
});
let d = descriptor_for_index("orders", None, &body);
assert_eq!(d.estimated_rows, None);
let schema = d.schema.as_ref().expect("schema");
assert_eq!(schema["properties"]["id"]["type"], "integer");
}
#[test]
fn descriptor_for_index_missing_mappings_yields_empty_schema() {
let d = descriptor_for_index("orders", Some(0), &json!({}));
assert_eq!(
d.schema,
Some(json!({"type": "object", "properties": {}})),
"no mappings → empty object schema, never a panic"
);
}
#[test]
fn source_advertises_discover() {
let config = ElasticsearchSourceConfig::new("http://localhost:9200", "idx");
let source = ElasticsearchSource::new(config).unwrap();
assert!(source.supports_discover());
}
}