Skip to main content

krabka_app_sdk/
messaging.rs

1//! Messaging module for publish, `CloudEvents` publish, and subscribe.
2
3use std::{
4    pin::Pin,
5    task::{Context, Poll},
6};
7
8use bytes::Bytes;
9use futures_util::Stream;
10use serde_json::Value;
11use tokio::sync::mpsc;
12
13use crate::{
14    client::{KrabkaClient, MockMessage},
15    error::KrabkaError,
16    pb,
17};
18
19const SUBSCRIPTION_CHANNEL_CAPACITY: usize = 32;
20
21/// Messaging module client.
22#[derive(Debug, Clone)]
23pub struct MessagingClient {
24    client: KrabkaClient,
25}
26
27impl MessagingClient {
28    pub(crate) fn new(client: KrabkaClient) -> Self {
29        Self { client }
30    }
31
32    /// Publish one raw record.
33    ///
34    /// # Errors
35    ///
36    /// Returns an error for invalid input, transport failures, or a gateway
37    /// record error.
38    pub async fn publish(
39        &self,
40        topic: impl Into<String>,
41        value: impl Into<Bytes>,
42        options: PublishOptions,
43    ) -> Result<RecordResult, KrabkaError> {
44        let topic = topic.into();
45        let value = value.into();
46        if self.client.is_mock() {
47            return self.publish_mock(topic, value, options);
48        }
49        if self.client.is_unreachable() {
50            return Err(KrabkaError::Transport("endpoint unreachable".into()));
51        }
52        if topic.is_empty() {
53            return Err(KrabkaError::InvalidArgument("topic is required".into()));
54        }
55        let request = pb::SendRequest {
56            records: vec![pb::Record {
57                topic,
58                key: options.key.map(Vec::from),
59                body: Some(pb::record::Body::Raw(value.to_vec())),
60                headers: headers_to_pb(options.headers),
61                partition: options.partition,
62                timestamp_ms: None,
63                idempotency_key: options.idempotency_key,
64                schema: None,
65            }],
66            acks: pb::Acks::All as i32,
67        };
68        let response: pb::SendResponse = self
69            .client
70            .inner
71            .connect
72            .unary("/krabka.gateway.v1.Gateway/Send", &request)
73            .await?;
74        response
75            .results
76            .into_iter()
77            .next()
78            .map(record_result_from_pb)
79            .transpose()?
80            .ok_or_else(|| KrabkaError::ServerError("gateway returned no record result".into()))
81    }
82
83    /// Publish a `CloudEvent` using the contract's binary-mode header mapping.
84    ///
85    /// # Errors
86    ///
87    /// Returns an error for invalid input, transport failures, or a gateway
88    /// record error.
89    pub async fn publish_event(
90        &self,
91        topic: impl Into<String>,
92        event: CloudEvent,
93    ) -> Result<RecordResult, KrabkaError> {
94        if event.id.is_empty() {
95            return Err(KrabkaError::InvalidArgument(
96                "CloudEvent id is required".into(),
97            ));
98        }
99        let headers = event.headers();
100        self.publish(
101            topic,
102            event.data,
103            PublishOptions {
104                headers,
105                ..PublishOptions::default()
106            },
107        )
108        .await
109    }
110
111    /// Subscribe to topics.
112    ///
113    /// # Errors
114    ///
115    /// Returns an error for invalid topic or group input, an unreachable
116    /// endpoint, or unavailable mock state.
117    pub fn subscribe(
118        &self,
119        topics: &[String],
120        group: impl Into<String>,
121        filter: Option<&Filter>,
122    ) -> Result<MessageStream, KrabkaError> {
123        let group = group.into();
124        let gateway_filter = filter.map(Filter::to_gateway_filter).transpose()?;
125        if !self.client.is_mock() {
126            return self.subscribe_live(topics, group, gateway_filter);
127        }
128        let messages = self
129            .client
130            .inner
131            .mock
132            .lock()
133            .map_err(|_| KrabkaError::ServerError("mock state lock poisoned".into()))?
134            .messages
135            .clone();
136        let (tx, rx) = mpsc::channel(messages.len().max(1));
137        for message in messages {
138            if !topics.iter().any(|topic| topic == &message.topic) {
139                continue;
140            }
141            if !filter_matches(filter, &message.value) {
142                continue;
143            }
144            let inbound = Inbound {
145                topic: message.topic,
146                partition: message.partition,
147                offset: message.offset,
148                value: message.value,
149                headers: message.headers.into_iter().collect(),
150            };
151            if tx.try_send(Ok(inbound)).is_err() {
152                break;
153            }
154        }
155        Ok(MessageStream { rx })
156    }
157
158    fn publish_mock(
159        &self,
160        topic: String,
161        value: Bytes,
162        options: PublishOptions,
163    ) -> Result<RecordResult, KrabkaError> {
164        if topic.is_empty() {
165            return Err(KrabkaError::InvalidArgument("topic is required".into()));
166        }
167        if topic == "__missing_topic" {
168            return Err(KrabkaError::NotFound("topic not found".into()));
169        }
170        let mut state = self
171            .client
172            .inner
173            .mock
174            .lock()
175            .expect("mock state lock not poisoned");
176        let offset = state
177            .messages
178            .iter()
179            .filter(|message| message.topic == topic)
180            .count()
181            .try_into()
182            .expect("message count fits i64");
183        state.messages.push(MockMessage {
184            topic,
185            partition: 0,
186            offset,
187            value,
188            headers: options.headers,
189        });
190        Ok(RecordResult {
191            partition: 0,
192            offset,
193            deduplicated: false,
194        })
195    }
196
197    fn subscribe_live(
198        &self,
199        topics: &[String],
200        group: String,
201        filter: Option<String>,
202    ) -> Result<MessageStream, KrabkaError> {
203        if self.client.is_unreachable() {
204            return Err(KrabkaError::Transport("endpoint unreachable".into()));
205        }
206        if topics.is_empty() {
207            return Err(KrabkaError::InvalidArgument(
208                "at least one topic is required".into(),
209            ));
210        }
211        if group.is_empty() {
212            return Err(KrabkaError::InvalidArgument("group is required".into()));
213        }
214
215        let request = pb::SubscribeFrame {
216            frame: Some(pb::subscribe_frame::Frame::Start(pb::SubscribeStart {
217                group_id: group,
218                topics: topics.to_vec(),
219                auto_commit: true,
220                filter: filter.unwrap_or_default(),
221            })),
222        };
223        let connect = self.client.inner.connect.clone();
224        let (tx, rx) = mpsc::channel(SUBSCRIPTION_CHANNEL_CAPACITY);
225        tokio::spawn(async move {
226            let mut inbound = match connect
227                .streaming::<_, pb::Inbound>("/krabka.gateway.v1.Gateway/Subscribe", &request)
228                .await
229            {
230                Ok(inbound) => inbound,
231                Err(error) => {
232                    let _ = tx.send(Err(KrabkaError::from(error))).await;
233                    return;
234                }
235            };
236            while let Some(message) = inbound.recv().await {
237                let mapped = message.map(inbound_from_pb).map_err(KrabkaError::from);
238                if tx.send(mapped).await.is_err() {
239                    return;
240                }
241            }
242        });
243        Ok(MessageStream { rx })
244    }
245}
246
247/// Publish options for raw records.
248#[derive(Debug, Clone, Default, PartialEq, Eq)]
249pub struct PublishOptions {
250    /// Optional record key.
251    pub key: Option<Bytes>,
252    /// Headers in publish order. Header keys may repeat and values may be null.
253    pub headers: Vec<(String, Option<Bytes>)>,
254    /// Optional partition.
255    pub partition: Option<i32>,
256    /// Optional idempotency key.
257    pub idempotency_key: Option<String>,
258}
259
260/// Publish result.
261#[derive(Debug, Clone, Copy, PartialEq, Eq)]
262pub struct RecordResult {
263    /// Partition index.
264    pub partition: i32,
265    /// Record offset.
266    pub offset: i64,
267    /// Whether the gateway deduplicated the record.
268    pub deduplicated: bool,
269}
270
271/// `CloudEvent` binary-mode payload.
272#[derive(Debug, Clone, PartialEq, Eq)]
273pub struct CloudEvent {
274    /// Event id.
275    pub id: String,
276    /// Event source.
277    pub source: String,
278    /// Event type.
279    pub type_: String,
280    /// Specversion, usually `1.0`.
281    pub specversion: String,
282    /// Optional data content type.
283    pub datacontenttype: Option<String>,
284    /// Raw event data.
285    pub data: Bytes,
286}
287
288impl CloudEvent {
289    fn headers(&self) -> Vec<(String, Option<Bytes>)> {
290        let mut headers = vec![
291            ("ce_id".into(), Some(Bytes::from(self.id.clone()))),
292            ("ce_source".into(), Some(Bytes::from(self.source.clone()))),
293            ("ce_type".into(), Some(Bytes::from(self.type_.clone()))),
294            (
295                "ce_specversion".into(),
296                Some(Bytes::from(self.specversion.clone())),
297            ),
298        ];
299        if let Some(datacontenttype) = &self.datacontenttype {
300            headers.push((
301                "content-type".into(),
302                Some(Bytes::from(datacontenttype.clone())),
303            ));
304        }
305        headers
306    }
307}
308
309/// Equality filter for structured JSON values.
310#[derive(Debug, Clone, PartialEq)]
311pub struct Filter {
312    /// JSON field path. v1 supports `$.field`.
313    pub path: String,
314    /// Expected scalar value.
315    pub value: Value,
316}
317
318impl Filter {
319    fn to_gateway_filter(&self) -> Result<String, KrabkaError> {
320        let Some(field) = self.path.strip_prefix("$.") else {
321            return Err(KrabkaError::InvalidArgument(
322                "filter path must start with $.".into(),
323            ));
324        };
325        if field.is_empty() || !field.split('.').all(valid_filter_segment) {
326            return Err(KrabkaError::InvalidArgument(
327                "filter path must contain identifier segments".into(),
328            ));
329        }
330        Ok(format!(
331            "{field} = {}",
332            gateway_filter_literal(&self.value)?
333        ))
334    }
335}
336
337/// Delivered message.
338#[derive(Debug, Clone, PartialEq, Eq)]
339pub struct Inbound {
340    /// Topic name.
341    pub topic: String,
342    /// Partition index.
343    pub partition: i32,
344    /// Record offset.
345    pub offset: i64,
346    /// Record value.
347    pub value: Bytes,
348    /// Headers in deterministic order.
349    pub headers: Vec<(String, Option<Bytes>)>,
350}
351
352/// Subscription stream.
353#[derive(Debug)]
354pub struct MessageStream {
355    rx: mpsc::Receiver<Result<Inbound, KrabkaError>>,
356}
357
358impl Stream for MessageStream {
359    type Item = Result<Inbound, KrabkaError>;
360
361    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
362        self.rx.poll_recv(cx)
363    }
364}
365
366fn headers_to_pb(headers: Vec<(String, Option<Bytes>)>) -> Vec<pb::Header> {
367    headers
368        .into_iter()
369        .map(|(key, value)| pb::Header {
370            key,
371            value: value.map(Vec::from),
372        })
373        .collect()
374}
375
376fn inbound_from_pb(inbound: pb::Inbound) -> Inbound {
377    Inbound {
378        topic: inbound.topic,
379        partition: inbound.partition,
380        offset: inbound.offset,
381        value: Bytes::from(inbound.value),
382        headers: inbound
383            .headers
384            .into_iter()
385            .map(|header| (header.key, header.value.map(Bytes::from)))
386            .collect(),
387    }
388}
389
390fn gateway_filter_literal(value: &Value) -> Result<String, KrabkaError> {
391    match value {
392        Value::String(value) => Ok(format!("'{}'", value.replace('\'', "''"))),
393        Value::Bool(value) => Ok(value.to_string()),
394        Value::Number(value) => Ok(value.to_string()),
395        Value::Null => Ok("null".into()),
396        Value::Array(_) | Value::Object(_) => Err(KrabkaError::InvalidArgument(
397            "filter value must be a string, number, boolean, or null".into(),
398        )),
399    }
400}
401
402fn valid_filter_segment(segment: &str) -> bool {
403    let mut chars = segment.chars();
404    let Some(first) = chars.next() else {
405        return false;
406    };
407    (first == '_' || first.is_ascii_alphabetic())
408        && chars.all(|character| character == '_' || character.is_ascii_alphanumeric())
409}
410
411fn record_result_from_pb(result: pb::RecordResult) -> Result<RecordResult, KrabkaError> {
412    if let Some(error) = result.error {
413        return Err(KrabkaError::from_connect_code(
414            match error.code {
415                3 => "invalid_argument",
416                5 => "not_found",
417                7 => "unauthenticated",
418                14 => "unavailable",
419                _ => "internal",
420            },
421            error.message,
422        ));
423    }
424    Ok(RecordResult {
425        partition: result.partition,
426        offset: result.offset,
427        deduplicated: result.deduplicated,
428    })
429}
430
431fn filter_matches(filter: Option<&Filter>, value: &Bytes) -> bool {
432    let Some(filter) = filter else {
433        return true;
434    };
435    let Ok(json) = serde_json::from_slice::<Value>(value) else {
436        return false;
437    };
438    let Some(field) = filter.path.strip_prefix("$.") else {
439        return false;
440    };
441    field
442        .split('.')
443        .try_fold(&json, |value, segment| value.get(segment))
444        == Some(&filter.value)
445}
446
447#[cfg(test)]
448mod tests {
449    use futures_util::StreamExt as _;
450
451    use super::*;
452
453    #[test]
454    fn cloudevent_binary_mode_headers_are_contract_shaped() {
455        let event = CloudEvent {
456            id: "1".into(),
457            source: "/svc".into(),
458            type_: "order.created".into(),
459            specversion: "1.0".into(),
460            datacontenttype: Some("application/json".into()),
461            data: Bytes::from_static(br#"{"n":7}"#),
462        };
463
464        let headers = event.headers();
465
466        assert!(
467            headers
468                .iter()
469                .any(|(k, v)| k == "ce_id" && v.as_deref() == Some(&b"1"[..]))
470        );
471        assert!(
472            headers
473                .iter()
474                .any(|(k, v)| k == "ce_source" && v.as_deref() == Some(&b"/svc"[..]))
475        );
476        assert!(
477            headers
478                .iter()
479                .any(|(k, v)| k == "ce_type" && v.as_deref() == Some(&b"order.created"[..]))
480        );
481        assert!(
482            headers
483                .iter()
484                .any(|(k, v)| k == "ce_specversion" && v.as_deref() == Some(&b"1.0"[..]))
485        );
486        assert!(
487            headers
488                .iter()
489                .any(|(k, v)| k == "content-type"
490                    && v.as_deref() == Some(&b"application/json"[..]))
491        );
492        assert!(!headers.iter().any(|(k, _)| k == "ce_datacontenttype"));
493    }
494
495    #[test]
496    fn headers_to_pb_preserves_order_duplicates_and_nulls() {
497        let headers = headers_to_pb(vec![
498            ("x".into(), Some(Bytes::from_static(b"first"))),
499            ("x".into(), None),
500            ("y".into(), Some(Bytes::from_static(b"last"))),
501        ]);
502
503        assert_eq!(
504            headers,
505            vec![
506                pb::Header {
507                    key: "x".into(),
508                    value: Some(b"first".to_vec()),
509                },
510                pb::Header {
511                    key: "x".into(),
512                    value: None,
513                },
514                pb::Header {
515                    key: "y".into(),
516                    value: Some(b"last".to_vec()),
517                },
518            ]
519        );
520    }
521
522    #[tokio::test]
523    async fn mock_subscribe_filters_matching_records() {
524        let client = KrabkaClient::new("mock://gateway").unwrap();
525        let messaging = client.messaging();
526        messaging
527            .publish(
528                "orders",
529                Bytes::from_static(br#"{"kind":"skip"}"#),
530                PublishOptions::default(),
531            )
532            .await
533            .unwrap();
534        messaging
535            .publish(
536                "orders",
537                Bytes::from_static(br#"{"kind":"keep"}"#),
538                PublishOptions::default(),
539            )
540            .await
541            .unwrap();
542
543        let mut stream = messaging
544            .subscribe(
545                &["orders".into()],
546                "g",
547                Some(&Filter {
548                    path: "$.kind".into(),
549                    value: Value::String("keep".into()),
550                }),
551            )
552            .unwrap();
553        let msg = stream.next().await.unwrap().unwrap();
554
555        assert_eq!(msg.offset, 1);
556        assert_eq!(msg.value.as_ref(), br#"{"kind":"keep"}"#);
557        assert!(stream.next().await.is_none());
558    }
559
560    #[tokio::test]
561    async fn subscription_inbox_applies_backpressure_at_capacity() {
562        let (tx, rx) = mpsc::channel(SUBSCRIPTION_CHANNEL_CAPACITY);
563        let message = Inbound {
564            topic: "orders".into(),
565            partition: 0,
566            offset: 0,
567            value: Bytes::new(),
568            headers: Vec::new(),
569        };
570        for _ in 0..SUBSCRIPTION_CHANNEL_CAPACITY {
571            tx.try_send(Ok(message.clone())).expect("inbox has room");
572        }
573        assert2::assert!(matches!(
574            tx.try_send(Ok(message.clone())),
575            Err(mpsc::error::TrySendError::Full(_))
576        ));
577
578        let mut stream = MessageStream { rx };
579        assert2::assert!(stream.next().await.is_some());
580        tx.try_send(Ok(message))
581            .expect("draining one item restores one slot");
582    }
583
584    #[test]
585    fn gateway_filter_uses_sql_standard_string_literals() {
586        let filter = Filter {
587            path: "$.path".into(),
588            value: Value::String("C:\\tmp\\O'Brien".into()),
589        };
590
591        assert_eq!(
592            filter.to_gateway_filter().unwrap(),
593            "path = 'C:\\tmp\\O''Brien'"
594        );
595    }
596
597    #[tokio::test]
598    async fn publish_event_rejects_empty_id() {
599        let client = KrabkaClient::new("mock://gateway").unwrap();
600        let error = client
601            .messaging()
602            .publish_event(
603                "events",
604                CloudEvent {
605                    id: String::new(),
606                    source: "/svc".into(),
607                    type_: "created".into(),
608                    specversion: "1.0".into(),
609                    datacontenttype: None,
610                    data: Bytes::new(),
611                },
612            )
613            .await
614            .unwrap_err();
615
616        assert_eq!(
617            error,
618            KrabkaError::InvalidArgument("CloudEvent id is required".into())
619        );
620    }
621
622    #[test]
623    fn filter_rejects_malformed_paths_and_nonscalar_values() {
624        for filter in [
625            Filter {
626                path: "kind".into(),
627                value: Value::Bool(true),
628            },
629            Filter {
630                path: "$.kind;drop".into(),
631                value: Value::Bool(true),
632            },
633            Filter {
634                path: "$.kind".into(),
635                value: Value::Array(Vec::new()),
636            },
637        ] {
638            assert!(filter.to_gateway_filter().is_err());
639        }
640    }
641}