faucet_source_elasticsearch/
config.rs

1//! Elasticsearch source configuration.
2
3use faucet_core::{AuthSpec, DEFAULT_BATCH_SIZE, FaucetError, validate_batch_size};
4use schemars::JsonSchema;
5use serde::{Deserialize, Serialize};
6use serde_json::{Value, json};
7
8pub use faucet_common_elasticsearch::ElasticsearchAuth;
9
10/// Configuration for the Elasticsearch search source.
11#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
12pub struct ElasticsearchSourceConfig {
13    /// Base URL of the Elasticsearch cluster (e.g. `"http://localhost:9200"`).
14    pub base_url: String,
15    /// Index name to search.
16    pub index: String,
17    /// Elasticsearch query DSL. Defaults to `{"match_all": {}}`.
18    pub query: Value,
19    /// Scroll context timeout (e.g. `"1m"`). Defaults to `"1m"`.
20    pub scroll_timeout: String,
21    /// Authentication: either inline (`{ type, config }`) or a `{ ref: <name> }`
22    /// pointer to a shared provider in the CLI's top-level `auth:` catalog.
23    pub auth: AuthSpec<ElasticsearchAuth>,
24    /// Maximum number of scroll pages to fetch. `None` means no limit.
25    pub max_pages: Option<usize>,
26    /// Records per emitted [`StreamPage`](faucet_core::StreamPage), which is
27    /// also the `size` parameter passed to the Elasticsearch scroll API
28    /// (`POST /{index}/_search?scroll={timeout}&size={batch_size}`). Each
29    /// scroll response becomes exactly one `StreamPage`. Defaults to
30    /// [`DEFAULT_BATCH_SIZE`].
31    ///
32    /// `batch_size = 0` is the "no batching" sentinel: the source issues a
33    /// single non-scroll `_search` request with `size = 10_000` (the default
34    /// `index.max_result_window`) and emits one `StreamPage`. Use it for
35    /// small indices or for sinks (e.g. SQL `COPY`, BigQuery load jobs) that
36    /// prefer one large request to many small ones. Indices that have raised
37    /// their `max_result_window` will still cap at 10_000 — raise this knob
38    /// or switch back to scroll if you need more.
39    #[serde(default = "default_batch_size")]
40    pub batch_size: usize,
41}
42
43fn default_batch_size() -> usize {
44    DEFAULT_BATCH_SIZE
45}
46
47impl ElasticsearchSourceConfig {
48    /// Create a new config with the required fields and sensible defaults.
49    pub fn new(base_url: impl Into<String>, index: impl Into<String>) -> Self {
50        Self {
51            base_url: base_url.into().trim_end_matches('/').to_string(),
52            index: index.into(),
53            query: json!({"match_all": {}}),
54            scroll_timeout: "1m".to_string(),
55            auth: AuthSpec::Inline(ElasticsearchAuth::None),
56            max_pages: None,
57            batch_size: DEFAULT_BATCH_SIZE,
58        }
59    }
60
61    /// Set the Elasticsearch query DSL.
62    pub fn query(mut self, q: Value) -> Self {
63        self.query = q;
64        self
65    }
66
67    /// Set the scroll context timeout (e.g. `"5m"`).
68    pub fn scroll_timeout(mut self, t: impl Into<String>) -> Self {
69        self.scroll_timeout = t.into();
70        self
71    }
72
73    /// Set the authentication method.
74    pub fn auth(mut self, a: ElasticsearchAuth) -> Self {
75        self.auth = AuthSpec::Inline(a);
76        self
77    }
78
79    /// Set the maximum number of scroll pages to fetch.
80    pub fn max_pages(mut self, n: usize) -> Self {
81        self.max_pages = Some(n);
82        self
83    }
84
85    /// Set the per-page document count for both the scroll API's `size`
86    /// parameter and the emitted [`StreamPage`](faucet_core::StreamPage)
87    /// size.
88    ///
89    /// Pass `0` to opt out of scroll entirely — the source will issue a
90    /// single `_search` with `size = 10_000` and emit one page.
91    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
92        self.batch_size = batch_size;
93        self
94    }
95
96    /// Validate the config at load time so a bad config fails fast with a typed
97    /// [`FaucetError::Config`] instead of surfacing deep in a run: rejects an
98    /// out-of-range `batch_size` (`> MAX_BATCH_SIZE`) and an empty `base_url` or
99    /// `index`.
100    pub fn validate(&self) -> Result<(), FaucetError> {
101        if self.base_url.trim().is_empty() {
102            return Err(FaucetError::Config(
103                "Elasticsearch source requires a non-empty `base_url`".into(),
104            ));
105        }
106        if self.index.trim().is_empty() {
107            return Err(FaucetError::Config(
108                "Elasticsearch source requires a non-empty `index`".into(),
109            ));
110        }
111        validate_batch_size(self.batch_size)?;
112        Ok(())
113    }
114}
115
116#[cfg(test)]
117mod tests {
118    use super::*;
119
120    #[test]
121    fn default_config() {
122        let config = ElasticsearchSourceConfig::new("http://localhost:9200", "my_index");
123        assert_eq!(config.base_url, "http://localhost:9200");
124        assert_eq!(config.index, "my_index");
125        assert_eq!(config.query, json!({"match_all": {}}));
126        assert_eq!(config.scroll_timeout, "1m");
127        assert!(config.max_pages.is_none());
128    }
129
130    #[test]
131    fn builder_methods() {
132        let config = ElasticsearchSourceConfig::new("http://es:9200/", "idx")
133            .query(json!({"term": {"status": "active"}}))
134            .scroll_timeout("5m")
135            .max_pages(10)
136            .auth(ElasticsearchAuth::Bearer {
137                token: "tok".into(),
138            });
139        assert_eq!(config.base_url, "http://es:9200");
140        assert_eq!(config.scroll_timeout, "5m");
141        assert_eq!(config.max_pages, Some(10));
142    }
143
144    #[test]
145    fn batch_size_defaults_to_default_batch_size() {
146        let config = ElasticsearchSourceConfig::new("http://localhost:9200", "idx");
147        assert_eq!(config.batch_size, faucet_core::DEFAULT_BATCH_SIZE);
148    }
149
150    #[test]
151    fn with_batch_size_overrides_default() {
152        let config =
153            ElasticsearchSourceConfig::new("http://localhost:9200", "idx").with_batch_size(500);
154        assert_eq!(config.batch_size, 500);
155    }
156
157    #[test]
158    fn batch_size_zero_is_accepted_as_no_batching_sentinel() {
159        let config =
160            ElasticsearchSourceConfig::new("http://localhost:9200", "idx").with_batch_size(0);
161        assert_eq!(config.batch_size, 0);
162        assert!(faucet_core::validate_batch_size(config.batch_size).is_ok());
163    }
164
165    #[test]
166    fn batch_size_above_max_is_rejected_by_validate_batch_size() {
167        let config = ElasticsearchSourceConfig::new("http://localhost:9200", "idx")
168            .with_batch_size(faucet_core::MAX_BATCH_SIZE + 1);
169        assert!(faucet_core::validate_batch_size(config.batch_size).is_err());
170    }
171
172    #[test]
173    fn batch_size_deserializes_from_json() {
174        let json = r#"{
175            "base_url": "http://localhost:9200",
176            "index": "idx",
177            "query": {"match_all": {}},
178            "scroll_timeout": "1m",
179            "auth": {"type": "none"},
180            "batch_size": 250
181        }"#;
182        let config: ElasticsearchSourceConfig = serde_json::from_str(json).unwrap();
183        assert_eq!(config.batch_size, 250);
184        assert!(matches!(config.auth, faucet_core::AuthSpec::Inline(_)));
185    }
186
187    #[test]
188    fn batch_size_defaults_when_missing_from_json() {
189        let json = r#"{
190            "base_url": "http://localhost:9200",
191            "index": "idx",
192            "query": {"match_all": {}},
193            "scroll_timeout": "1m",
194            "auth": {"type": "none"}
195        }"#;
196        let config: ElasticsearchSourceConfig = serde_json::from_str(json).unwrap();
197        assert_eq!(config.batch_size, faucet_core::DEFAULT_BATCH_SIZE);
198    }
199
200    #[test]
201    fn validate_accepts_valid_config() {
202        assert!(
203            ElasticsearchSourceConfig::new("http://localhost:9200", "idx")
204                .validate()
205                .is_ok()
206        );
207    }
208
209    #[test]
210    fn validate_rejects_oversized_batch_size() {
211        let config = ElasticsearchSourceConfig::new("http://localhost:9200", "idx")
212            .with_batch_size(faucet_core::MAX_BATCH_SIZE + 1);
213        assert!(matches!(config.validate(), Err(FaucetError::Config(_))));
214    }
215
216    #[test]
217    fn validate_rejects_empty_base_url() {
218        assert!(matches!(
219            ElasticsearchSourceConfig::new("  ", "idx").validate(),
220            Err(FaucetError::Config(_))
221        ));
222    }
223
224    #[test]
225    fn validate_rejects_empty_index() {
226        assert!(matches!(
227            ElasticsearchSourceConfig::new("http://localhost:9200", "").validate(),
228            Err(FaucetError::Config(_))
229        ));
230    }
231}
faucet_source_elasticsearch/config.rs

faucet_source_elasticsearch/
config.rs
