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/// The default JSON key the per-message attribute map is surfaced under.
22pub const DEFAULT_ATTRIBUTES_KEY: &str = "__attributes";
23
24/// Configuration for [`PubsubSource`](crate::PubsubSource).
25#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
26pub struct PubsubSourceConfig {
27    /// Subscription id (short name, not the fully-qualified path). The client
28    /// forms `projects/<project>/subscriptions/<id>` from the connection's
29    /// project.
30    pub subscription: String,
31
32    /// Project / endpoint / emulator / credentials (flattened).
33    #[serde(flatten)]
34    pub connection: PubsubConnection,
35
36    /// Payload decoding for the message `data` field.
37    #[serde(default)]
38    pub value_format: ValueFormat,
39
40    /// JSON key the message attribute map is surfaced under. Default
41    /// `__attributes`.
42    #[serde(default = "default_attributes_key")]
43    pub attributes_key: String,
44
45    /// Messages requested per `pull` RPC (1–1000). Default 100.
46    #[serde(default = "default_max_messages_per_pull")]
47    pub max_messages_per_pull: usize,
48
49    /// Stop after this many seconds without a new message. At least one of
50    /// `idle_termination_secs` and `max_messages` must be set (mirrors the
51    /// Kafka / Kinesis sources) — a batch run must terminate.
52    #[serde(default, skip_serializing_if = "Option::is_none")]
53    pub idle_termination_secs: Option<u64>,
54    /// Stop after this many messages in total.
55    #[serde(default, skip_serializing_if = "Option::is_none")]
56    pub max_messages: Option<usize>,
57
58    /// Records per emitted [`StreamPage`](faucet_core::StreamPage). `0` is the
59    /// "no batching" sentinel: one page per drain. Default 1000. A page's
60    /// messages are acked once the pipeline has durably written the page.
61    #[serde(default = "default_batch_size")]
62    pub batch_size: usize,
63}
64
65fn default_attributes_key() -> String {
66    DEFAULT_ATTRIBUTES_KEY.to_string()
67}
68fn default_max_messages_per_pull() -> usize {
69    100
70}
71fn default_batch_size() -> usize {
72    faucet_core::DEFAULT_BATCH_SIZE
73}
74
75impl PubsubSourceConfig {
76    /// Minimal config with defaults for everything but the subscription id.
77    pub fn new(subscription: impl Into<String>) -> Self {
78        Self {
79            subscription: subscription.into(),
80            connection: PubsubConnection::default(),
81            value_format: ValueFormat::default(),
82            attributes_key: default_attributes_key(),
83            max_messages_per_pull: default_max_messages_per_pull(),
84            idle_termination_secs: None,
85            max_messages: None,
86            batch_size: default_batch_size(),
87        }
88    }
89
90    /// Fail-fast validation, called from `PubsubSource::new`.
91    pub fn validate(&self) -> Result<(), faucet_core::FaucetError> {
92        use faucet_core::FaucetError;
93        if self.subscription.trim().is_empty() {
94            return Err(FaucetError::Config(
95                "pubsub source: subscription must not be empty".into(),
96            ));
97        }
98        if self.attributes_key.trim().is_empty() {
99            return Err(FaucetError::Config(
100                "pubsub source: attributes_key must not be empty".into(),
101            ));
102        }
103        if self.max_messages_per_pull == 0 || self.max_messages_per_pull > 1000 {
104            return Err(FaucetError::Config(format!(
105                "pubsub source: max_messages_per_pull must be 1..=1000 (got {})",
106                self.max_messages_per_pull
107            )));
108        }
109        faucet_core::validate_batch_size(self.batch_size)?;
110        if self.idle_termination_secs.is_none() && self.max_messages.is_none() {
111            return Err(FaucetError::Config(
112                "pubsub source: set at least one of idle_termination_secs / max_messages so a \
113                 run can terminate (mirrors the kafka / kinesis sources)"
114                    .into(),
115            ));
116        }
117        if self.idle_termination_secs == Some(0) {
118            return Err(FaucetError::Config(
119                "pubsub source: idle_termination_secs must be at least 1".into(),
120            ));
121        }
122        if self.max_messages == Some(0) {
123            return Err(FaucetError::Config(
124                "pubsub source: max_messages must be at least 1".into(),
125            ));
126        }
127        Ok(())
128    }
129}
130
131#[cfg(test)]
132mod tests {
133    use super::*;
134    use faucet_common_pubsub::PubsubCredentials;
135
136    fn valid() -> PubsubSourceConfig {
137        let mut c = PubsubSourceConfig::new("orders-sub");
138        c.max_messages = Some(100);
139        c
140    }
141
142    #[test]
143    fn defaults_are_sensible() {
144        let c = PubsubSourceConfig::new("orders-sub");
145        assert_eq!(c.value_format, ValueFormat::Json);
146        assert_eq!(c.attributes_key, "__attributes");
147        assert_eq!(c.max_messages_per_pull, 100);
148        assert_eq!(c.batch_size, faucet_core::DEFAULT_BATCH_SIZE);
149        assert!(c.idle_termination_secs.is_none() && c.max_messages.is_none());
150    }
151
152    #[test]
153    fn validation_bounds() {
154        valid().validate().unwrap();
155
156        let mut c = valid();
157        c.subscription = "  ".into();
158        assert!(c.validate().is_err());
159
160        let mut c = valid();
161        c.attributes_key = String::new();
162        assert!(c.validate().is_err());
163
164        let mut c = valid();
165        c.max_messages_per_pull = 0;
166        assert!(c.validate().is_err());
167        c.max_messages_per_pull = 1001;
168        assert!(c.validate().is_err());
169
170        let mut c = valid();
171        c.batch_size = faucet_core::MAX_BATCH_SIZE + 1;
172        assert!(c.validate().is_err());
173
174        let mut c = valid();
175        c.idle_termination_secs = None;
176        c.max_messages = None;
177        let err = c.validate().unwrap_err();
178        assert!(err.to_string().contains("idle_termination_secs"), "{err}");
179
180        let mut c = valid();
181        c.idle_termination_secs = Some(0);
182        assert!(c.validate().is_err());
183        let mut c = valid();
184        c.max_messages = Some(0);
185        assert!(c.validate().is_err());
186    }
187
188    #[test]
189    fn full_config_parses_from_yaml() {
190        let yaml = r#"
191subscription: orders-sub
192project_id: my-proj
193emulator_host: "localhost:8085"
194credentials: { type: anonymous }
195value_format: string
196attributes_key: attrs
197max_messages_per_pull: 250
198idle_termination_secs: 5
199max_messages: 500
200batch_size: 100
201"#;
202        let c: PubsubSourceConfig = serde_yaml::from_str(yaml).unwrap();
203        c.validate().unwrap();
204        assert_eq!(c.value_format, ValueFormat::String);
205        assert_eq!(c.attributes_key, "attrs");
206        assert_eq!(c.connection.project_id.as_deref(), Some("my-proj"));
207        assert_eq!(c.connection.credentials, PubsubCredentials::Anonymous);
208        assert_eq!(c.max_messages_per_pull, 250);
209    }
210}