Skip to main content

faucet_source_duckdb/
config.rs

1//! DuckDB source configuration.
2
3use faucet_core::{DEFAULT_BATCH_SIZE, FaucetError, validate_batch_size};
4use schemars::JsonSchema;
5use serde::{Deserialize, Serialize};
6
7/// Configuration for the DuckDB query source.
8#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
9pub struct DuckdbSourceConfig {
10    /// Path to the DuckDB database file, or `:memory:` for an in-memory
11    /// database. A `duckdb://` / `duckdb:` scheme prefix is accepted and
12    /// stripped.
13    pub database: String,
14    /// SQL query to execute.
15    pub query: String,
16    /// Open the database read-only. Defaults to `false` (read-write).
17    ///
18    /// Set `true` to attach to a database file another process holds open —
19    /// DuckDB permits multiple read-only connections to one file but only a
20    /// single read-write connection.
21    #[serde(default)]
22    pub read_only: bool,
23    /// Records per emitted [`StreamPage`](faucet_core::StreamPage). Rows are
24    /// drained from the DuckDB result and yielded whenever the buffer reaches
25    /// this size. Defaults to [`DEFAULT_BATCH_SIZE`].
26    ///
27    /// `batch_size = 0` is the "no batching" sentinel: the entire result set is
28    /// emitted in a single page. Useful for small lookup tables.
29    #[serde(default = "default_batch_size")]
30    pub batch_size: usize,
31}
32
33fn default_batch_size() -> usize {
34    DEFAULT_BATCH_SIZE
35}
36
37impl DuckdbSourceConfig {
38    /// Create a new config with the required database path and query.
39    pub fn new(database: impl Into<String>, query: impl Into<String>) -> Self {
40        Self {
41            database: database.into(),
42            query: query.into(),
43            read_only: false,
44            batch_size: DEFAULT_BATCH_SIZE,
45        }
46    }
47
48    /// Set the per-page row count for
49    /// [`Source::stream_pages`](faucet_core::Source::stream_pages).
50    ///
51    /// Pass `0` to opt out of batching — the entire result set is emitted in a
52    /// single [`StreamPage`](faucet_core::StreamPage).
53    pub fn with_batch_size(mut self, batch_size: usize) -> Self {
54        self.batch_size = batch_size;
55        self
56    }
57
58    /// Open the database read-only.
59    pub fn read_only(mut self, read_only: bool) -> Self {
60        self.read_only = read_only;
61        self
62    }
63
64    /// The filesystem path with any `duckdb://` / `duckdb:` scheme stripped.
65    pub(crate) fn resolved_path(&self) -> &str {
66        self.database
67            .trim_start_matches("duckdb://")
68            .trim_start_matches("duckdb:")
69    }
70
71    /// Validate the config at load time so a bad config fails fast with a typed
72    /// [`FaucetError::Config`] instead of surfacing deep in a run: rejects an
73    /// out-of-range `batch_size` (`> MAX_BATCH_SIZE`) and an empty `database` or
74    /// `query`.
75    pub fn validate(&self) -> Result<(), FaucetError> {
76        if self.database.trim().is_empty() {
77            return Err(FaucetError::Config(
78                "DuckDB source requires a non-empty `database` (a file path or `:memory:`)".into(),
79            ));
80        }
81        if self.query.trim().is_empty() {
82            return Err(FaucetError::Config(
83                "DuckDB source requires a non-empty `query`".into(),
84            ));
85        }
86        validate_batch_size(self.batch_size)?;
87        Ok(())
88    }
89}
90
91#[cfg(test)]
92mod tests {
93    use super::*;
94
95    #[test]
96    fn default_config() {
97        let config = DuckdbSourceConfig::new("data.duckdb", "SELECT * FROM events");
98        assert_eq!(config.database, "data.duckdb");
99        assert_eq!(config.query, "SELECT * FROM events");
100        assert_eq!(config.batch_size, DEFAULT_BATCH_SIZE);
101        assert!(!config.read_only);
102    }
103
104    #[test]
105    fn resolved_path_strips_scheme() {
106        assert_eq!(
107            DuckdbSourceConfig::new("duckdb:///tmp/a.duckdb", "SELECT 1").resolved_path(),
108            "/tmp/a.duckdb"
109        );
110        assert_eq!(
111            DuckdbSourceConfig::new("duckdb::memory:", "SELECT 1").resolved_path(),
112            ":memory:"
113        );
114        assert_eq!(
115            DuckdbSourceConfig::new(":memory:", "SELECT 1").resolved_path(),
116            ":memory:"
117        );
118    }
119
120    #[test]
121    fn with_batch_size_overrides_default() {
122        let config = DuckdbSourceConfig::new(":memory:", "SELECT 1").with_batch_size(500);
123        assert_eq!(config.batch_size, 500);
124    }
125
126    #[test]
127    fn batch_size_zero_is_accepted_as_no_batching_sentinel() {
128        let config = DuckdbSourceConfig::new(":memory:", "SELECT 1").with_batch_size(0);
129        assert_eq!(config.batch_size, 0);
130        assert!(faucet_core::validate_batch_size(config.batch_size).is_ok());
131    }
132
133    #[test]
134    fn batch_size_deserializes_from_json() {
135        let json = r#"{
136            "database": ":memory:",
137            "query": "SELECT 1",
138            "batch_size": 250
139        }"#;
140        let config: DuckdbSourceConfig = serde_json::from_str(json).unwrap();
141        assert_eq!(config.batch_size, 250);
142        assert!(!config.read_only);
143    }
144
145    #[test]
146    fn validate_accepts_valid_config() {
147        assert!(
148            DuckdbSourceConfig::new(":memory:", "SELECT 1")
149                .validate()
150                .is_ok()
151        );
152    }
153
154    #[test]
155    fn validate_rejects_oversized_batch_size() {
156        let config = DuckdbSourceConfig::new(":memory:", "SELECT 1")
157            .with_batch_size(faucet_core::MAX_BATCH_SIZE + 1);
158        assert!(matches!(config.validate(), Err(FaucetError::Config(_))));
159    }
160
161    #[test]
162    fn validate_rejects_empty_database() {
163        assert!(matches!(
164            DuckdbSourceConfig::new("  ", "SELECT 1").validate(),
165            Err(FaucetError::Config(_))
166        ));
167    }
168
169    #[test]
170    fn validate_rejects_empty_query() {
171        assert!(matches!(
172            DuckdbSourceConfig::new(":memory:", "").validate(),
173            Err(FaucetError::Config(_))
174        ));
175    }
176}