1use 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#[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 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 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 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#[derive(Debug, Clone, Default, PartialEq, Eq)]
249pub struct PublishOptions {
250 pub key: Option<Bytes>,
252 pub headers: Vec<(String, Option<Bytes>)>,
254 pub partition: Option<i32>,
256 pub idempotency_key: Option<String>,
258}
259
260#[derive(Debug, Clone, Copy, PartialEq, Eq)]
262pub struct RecordResult {
263 pub partition: i32,
265 pub offset: i64,
267 pub deduplicated: bool,
269}
270
271#[derive(Debug, Clone, PartialEq, Eq)]
273pub struct CloudEvent {
274 pub id: String,
276 pub source: String,
278 pub type_: String,
280 pub specversion: String,
282 pub datacontenttype: Option<String>,
284 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#[derive(Debug, Clone, PartialEq)]
311pub struct Filter {
312 pub path: String,
314 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#[derive(Debug, Clone, PartialEq, Eq)]
339pub struct Inbound {
340 pub topic: String,
342 pub partition: i32,
344 pub offset: i64,
346 pub value: Bytes,
348 pub headers: Vec<(String, Option<Bytes>)>,
350}
351
352#[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}