Skip to main content

faucet_source_pubsub/
config.rs

1//! Configuration for the Pub/Sub source. No I/O here.
2
3use faucet_common_pubsub::PubsubConnection;
4use schemars::JsonSchema;
5use serde::{Deserialize, Serialize};
6
7/// How each message's `data` payload is decoded into the emitted `data` field.
8#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
9#[serde(rename_all = "snake_case")]
10pub enum ValueFormat {
11    /// Parse the payload as JSON. A message that is not valid JSON fails the
12    /// stream with a typed error naming its `message_id`.
13    #[default]
14    Json,
15    /// Decode the payload as UTF-8 (invalid UTF-8 fails the message).
16    String,
17    /// Base64-encode the raw payload bytes into a JSON string.
18    Bytes,
19}
20
21/// What to do with a message whose payload does not decode under
22/// `value_format`.
23#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
24#[serde(rename_all = "snake_case")]
25pub enum OnDecodeError {
26    /// Fail the run. The message stays unacked, so Pub/Sub redelivers it and
27    /// the next run fails at it again until it is removed.
28    #[default]
29    Fail,
30    /// Drop the message: it is acked with its page and no record is emitted.
31    Skip,
32    /// Emit the record with the payload base64-encoded under `data` and the
33    /// decode error under `decode_error`.
34    Raw,
35}
36
37/// The default JSON key the per-message attribute map is surfaced under.
38pub const DEFAULT_ATTRIBUTES_KEY: &str = "__attributes";
39
40/// Configuration for [`PubsubSource`](crate::PubsubSource).
41#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
42pub struct PubsubSourceConfig {
43    /// Subscription id (short name, not the fully-qualified path). The client
44    /// forms `projects/<project>/subscriptions/<id>` from the connection's
45    /// project.
46    pub subscription: String,
47
48    /// Project / endpoint / emulator / credentials (flattened).
49    #[serde(flatten)]
50    pub connection: PubsubConnection,
51
52    /// Payload decoding for the message `data` field.
53    #[serde(default)]
54    pub value_format: ValueFormat,
55
56    /// JSON key the message attribute map is surfaced under. Default
57    /// `__attributes`. Must differ from the record's own fields (`data`,
58    /// `message_id`, `ordering_key`, `publish_time_millis`, `decode_error`).
59    #[serde(default = "default_attributes_key")]
60    pub attributes_key: String,
61
62    /// Messages requested per `pull` RPC (1–1000). Default 100.
63    #[serde(default = "default_max_messages_per_pull")]
64    pub max_messages_per_pull: usize,
65
66    /// Stop after this many seconds without a new message. At least one of
67    /// `idle_termination_secs` and `max_messages` must be set (mirrors the
68    /// Kafka / Kinesis sources) — a batch run must terminate.
69    #[serde(default, skip_serializing_if = "Option::is_none")]
70    pub idle_termination_secs: Option<u64>,
71    /// Stop after this many messages in total.
72    #[serde(default, skip_serializing_if = "Option::is_none")]
73    pub max_messages: Option<usize>,
74
75    /// Records per emitted [`StreamPage`](faucet_core::StreamPage). `0` is the
76    /// "no batching" sentinel: one page per drain. Default 1000. A page's
77    /// messages are acked once the pipeline has durably written the page.
78    #[serde(default = "default_batch_size")]
79    pub batch_size: usize,
80
81    /// What to do with a message whose payload does not decode. Default
82    /// `fail`.
83    #[serde(default)]
84    pub on_decode_error: OnDecodeError,
85
86    /// Ack deadline, in seconds, the source keeps renewing on every message it
87    /// has pulled but not yet acked (while a page is assembled and while the
88    /// sink writes it), so a slow page is not redelivered into the same run.
89    /// Renewed every 5 s (or every third of the window, if shorter). `0`
90    /// disables renewal. Range 10–600. Default 60.
91    #[serde(default = "default_ack_deadline_extension_secs")]
92    pub ack_deadline_extension_secs: u32,
93}
94
95fn default_ack_deadline_extension_secs() -> u32 {
96    60
97}
98
99/// Fields every record already carries, which `attributes_key` may not name.
100const RESERVED_RECORD_KEYS: &[&str] = &[
101    "data",
102    "message_id",
103    "ordering_key",
104    "publish_time_millis",
105    "decode_error",
106];
107
108fn default_attributes_key() -> String {
109    DEFAULT_ATTRIBUTES_KEY.to_string()
110}
111fn default_max_messages_per_pull() -> usize {
112    100
113}
114fn default_batch_size() -> usize {
115    faucet_core::DEFAULT_BATCH_SIZE
116}
117
118impl PubsubSourceConfig {
119    /// Minimal config with defaults for everything but the subscription id.
120    pub fn new(subscription: impl Into<String>) -> Self {
121        Self {
122            subscription: subscription.into(),
123            connection: PubsubConnection::default(),
124            value_format: ValueFormat::default(),
125            attributes_key: default_attributes_key(),
126            max_messages_per_pull: default_max_messages_per_pull(),
127            idle_termination_secs: None,
128            max_messages: None,
129            batch_size: default_batch_size(),
130            on_decode_error: OnDecodeError::default(),
131            ack_deadline_extension_secs: default_ack_deadline_extension_secs(),
132        }
133    }
134
135    /// Fail-fast validation, called from `PubsubSource::new`.
136    pub fn validate(&self) -> Result<(), faucet_core::FaucetError> {
137        use faucet_core::FaucetError;
138        if self.subscription.trim().is_empty() {
139            return Err(FaucetError::Config(
140                "pubsub source: subscription must not be empty".into(),
141            ));
142        }
143        if self.attributes_key.trim().is_empty() {
144            return Err(FaucetError::Config(
145                "pubsub source: attributes_key must not be empty".into(),
146            ));
147        }
148        if RESERVED_RECORD_KEYS.contains(&self.attributes_key.as_str()) {
149            return Err(FaucetError::Config(format!(
150                "pubsub source: attributes_key {:?} would replace the record's own `{}` \
151                 field; choose another name",
152                self.attributes_key, self.attributes_key
153            )));
154        }
155        if self.max_messages_per_pull == 0 || self.max_messages_per_pull > 1000 {
156            return Err(FaucetError::Config(format!(
157                "pubsub source: max_messages_per_pull must be 1..=1000 (got {})",
158                self.max_messages_per_pull
159            )));
160        }
161        faucet_core::validate_batch_size(self.batch_size)?;
162        if self.idle_termination_secs.is_none() && self.max_messages.is_none() {
163            return Err(FaucetError::Config(
164                "pubsub source: set at least one of idle_termination_secs / max_messages so a \
165                 run can terminate (mirrors the kafka / kinesis sources)"
166                    .into(),
167            ));
168        }
169        if self.idle_termination_secs == Some(0) {
170            return Err(FaucetError::Config(
171                "pubsub source: idle_termination_secs must be at least 1".into(),
172            ));
173        }
174        if self.max_messages == Some(0) {
175            return Err(FaucetError::Config(
176                "pubsub source: max_messages must be at least 1".into(),
177            ));
178        }
179        if self.ack_deadline_extension_secs != 0
180            && !(10..=600).contains(&self.ack_deadline_extension_secs)
181        {
182            return Err(FaucetError::Config(format!(
183                "pubsub source: ack_deadline_extension_secs must be 10..=600 or 0 to disable \
184                 renewal (got {})",
185                self.ack_deadline_extension_secs
186            )));
187        }
188        Ok(())
189    }
190}
191
192#[cfg(test)]
193mod tests {
194    use super::*;
195    use faucet_common_pubsub::PubsubCredentials;
196
197    fn valid() -> PubsubSourceConfig {
198        let mut c = PubsubSourceConfig::new("orders-sub");
199        c.max_messages = Some(100);
200        c
201    }
202
203    #[test]
204    fn defaults_are_sensible() {
205        let c = PubsubSourceConfig::new("orders-sub");
206        assert_eq!(c.value_format, ValueFormat::Json);
207        assert_eq!(c.attributes_key, "__attributes");
208        assert_eq!(c.max_messages_per_pull, 100);
209        assert_eq!(c.batch_size, faucet_core::DEFAULT_BATCH_SIZE);
210        assert!(c.idle_termination_secs.is_none() && c.max_messages.is_none());
211        assert_eq!(c.on_decode_error, OnDecodeError::Fail);
212        assert_eq!(c.ack_deadline_extension_secs, 60);
213    }
214
215    #[test]
216    fn ack_deadline_extension_bounds() {
217        let mut c = valid();
218        c.ack_deadline_extension_secs = 0;
219        c.validate().unwrap();
220        c.ack_deadline_extension_secs = 9;
221        assert!(c.validate().is_err());
222        c.ack_deadline_extension_secs = 601;
223        assert!(c.validate().is_err());
224        c.ack_deadline_extension_secs = 600;
225        c.validate().unwrap();
226    }
227
228    #[test]
229    fn validation_bounds() {
230        valid().validate().unwrap();
231
232        let mut c = valid();
233        c.subscription = "  ".into();
234        assert!(c.validate().is_err());
235
236        let mut c = valid();
237        c.attributes_key = String::new();
238        assert!(c.validate().is_err());
239
240        for key in [
241            "data",
242            "message_id",
243            "ordering_key",
244            "publish_time_millis",
245            "decode_error",
246        ] {
247            let mut c = valid();
248            c.attributes_key = key.into();
249            let err = c.validate().unwrap_err().to_string();
250            assert!(err.contains("would replace"), "{key}: {err}");
251        }
252
253        let mut c = valid();
254        c.max_messages_per_pull = 0;
255        assert!(c.validate().is_err());
256        c.max_messages_per_pull = 1001;
257        assert!(c.validate().is_err());
258
259        let mut c = valid();
260        c.batch_size = faucet_core::MAX_BATCH_SIZE + 1;
261        assert!(c.validate().is_err());
262
263        let mut c = valid();
264        c.idle_termination_secs = None;
265        c.max_messages = None;
266        let err = c.validate().unwrap_err();
267        assert!(err.to_string().contains("idle_termination_secs"), "{err}");
268
269        let mut c = valid();
270        c.idle_termination_secs = Some(0);
271        assert!(c.validate().is_err());
272        let mut c = valid();
273        c.max_messages = Some(0);
274        assert!(c.validate().is_err());
275    }
276
277    #[test]
278    fn full_config_parses_from_yaml() {
279        let yaml = r#"
280subscription: orders-sub
281project_id: my-proj
282emulator_host: "localhost:8085"
283credentials: { type: anonymous }
284value_format: string
285attributes_key: attrs
286max_messages_per_pull: 250
287idle_termination_secs: 5
288max_messages: 500
289batch_size: 100
290"#;
291        let c: PubsubSourceConfig = yaml_serde::from_str(yaml).unwrap();
292        c.validate().unwrap();
293        assert_eq!(c.value_format, ValueFormat::String);
294        assert_eq!(c.attributes_key, "attrs");
295        assert_eq!(c.connection.project_id.as_deref(), Some("my-proj"));
296        assert_eq!(c.connection.credentials, PubsubCredentials::Anonymous);
297        assert_eq!(c.max_messages_per_pull, 250);
298    }
299}