#![cfg_attr(
not(feature = "testing"),
allow(clippy::infallible_destructuring_match)
)]
use std::future::{Future, ready};
use std::sync::Arc;
use aws_sdk_kinesis::Client;
use aws_sdk_kinesis::primitives::Blob;
#[cfg(feature = "asyncapi")]
use ruststream::asyncapi::{Binding, Bindings};
use ruststream::runtime::{MapPublisher, PublishBuilder, PublishSink};
use ruststream::{HeaderMap, OutgoingFor, PairError, PublishPolicy, Publisher, Take};
#[cfg(feature = "asyncapi")]
use serde::Serialize;
#[cfg(feature = "asyncapi")]
use crate::stream::BINDING_KEY;
use crate::broker::{ConnectedKinesisBroker, CoreCell, Link};
use crate::error::{KinesisError, sdk_err};
use crate::message::{PARTITION_KEY_HEADER, encode_envelope};
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct KinesisPublishOptions {
pub partition_key: Option<String>,
}
#[derive(Clone)]
pub struct KinesisPublisher {
cell: CoreCell,
partition_key: Option<Arc<str>>,
}
impl std::fmt::Debug for KinesisPublisher {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("KinesisPublisher").finish_non_exhaustive()
}
}
impl KinesisPublisher {
pub(crate) fn new(cell: CoreCell) -> Self {
Self {
cell,
partition_key: None,
}
}
pub(crate) fn keyed(cell: CoreCell, partition_key: Option<Arc<str>>) -> Self {
Self {
cell,
partition_key,
}
}
fn link(&self) -> Result<&Link, KinesisError> {
self.cell.get().ok_or(KinesisError::NotConnected)
}
}
pub(crate) fn resolve_partition_key(
options: Option<&KinesisPublishOptions>,
headers: &HeaderMap,
from_policy: Option<&str>,
) -> String {
if let Some(key) = options.and_then(|options| options.partition_key.as_deref()) {
return key.to_owned();
}
if let Some(named) = headers.get(PARTITION_KEY_HEADER) {
return String::from_utf8_lossy(named).into_owned();
}
from_policy.map_or_else(spread_key, str::to_owned)
}
fn spread_key() -> String {
use std::sync::atomic::{AtomicU64, Ordering};
static SEQ: AtomicU64 = AtomicU64::new(0);
format!(
"rs-{}-{}",
std::process::id(),
SEQ.fetch_add(1, Ordering::Relaxed)
)
}
impl Publisher for KinesisPublisher {
type Payload = Take;
type Error = KinesisError;
type Options = KinesisPublishOptions;
async fn publish(
&self,
msg: OutgoingFor<'_, Take>,
options: Option<&Self::Options>,
) -> Result<(), Self::Error> {
let core = match self.link()? {
Link::Aws(core) => core,
#[cfg(feature = "testing")]
Link::InProcess(bus) => {
return bus.put(msg, options, self.partition_key.as_deref());
}
};
core.ensure_open()?;
let partition_key =
resolve_partition_key(options, msg.headers(), self.partition_key.as_deref());
let (stream, payload, headers) = msg.into_parts();
let data = encode_envelope(&headers, payload);
core.client
.with(Client::put_record)
.stream_name(stream)
.partition_key(partition_key)
.data(Blob::new(data))
.send()
.await
.map(|_| ())
.map_err(|e| KinesisError::Publish {
stream: stream.to_owned(),
source: sdk_err(&e),
})
}
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
#[must_use]
pub struct KinesisPublish {
partition_key: Option<String>,
}
impl KinesisPublish {
pub fn partition_key(mut self, key: impl Into<String>) -> Self {
self.partition_key = Some(key.into());
self
}
fn key(&self) -> Option<Arc<str>> {
self.partition_key.as_deref().map(Arc::from)
}
#[cfg(feature = "asyncapi")]
fn message_binding(&self) -> Bindings {
let Some(partition_key) = self.partition_key.as_deref() else {
return Bindings::new();
};
let body = KinesisMessageBinding { partition_key };
Binding::extension(BINDING_KEY, &body)
.map_or_else(|_| Bindings::new(), |binding| Bindings::new().with(binding))
}
}
#[cfg(feature = "asyncapi")]
fn channel_binding(stream: &str) -> Bindings {
let body = KinesisPublishChannelBinding { stream };
Binding::extension(BINDING_KEY, &body)
.map_or_else(|_| Bindings::new(), |binding| Bindings::new().with(binding))
}
#[cfg(feature = "asyncapi")]
#[derive(Debug, Serialize)]
struct KinesisMessageBinding<'a> {
#[serde(rename = "partitionKey")]
partition_key: &'a str,
}
#[cfg(feature = "asyncapi")]
#[derive(Debug, Serialize)]
struct KinesisPublishChannelBinding<'a> {
stream: &'a str,
}
impl crate::sealed::PolicyKey for KinesisPublish {
fn with_partition_key(self, key: String) -> Self {
self.partition_key(key)
}
}
impl PublishPolicy<ConnectedKinesisBroker> for KinesisPublish {
type Live = KinesisPublisher;
fn pair(
self,
connected: &ConnectedKinesisBroker,
) -> impl Future<Output = Result<Self::Live, PairError>> {
ready(Ok(connected.publisher_keyed(self.key())))
}
#[cfg(feature = "asyncapi")]
fn channel_bindings(&self, channel: &str) -> Bindings {
channel_binding(channel)
}
#[cfg(feature = "asyncapi")]
fn message_bindings(&self, _channel: &str) -> Bindings {
self.message_binding()
}
}
pub trait KinesisPublishSettings: Sized {
#[allow(clippy::return_self_not_must_use)]
fn partition_key(self, key: impl Into<String>) -> Self;
}
impl<T> KinesisPublishSettings for T
where
T: MapPublisher<Policy: crate::sealed::PolicyKey>,
{
fn partition_key(self, key: impl Into<String>) -> Self {
let key = key.into();
self.map_publisher(|policy| crate::sealed::PolicyKey::with_partition_key(policy, key))
}
}
pub trait KinesisPublishSteps {
#[must_use]
fn partition_key(self, key: impl Into<String>) -> Self;
}
impl<Sink, Body, Enc, Hdrs, Dest> KinesisPublishSteps
for PublishBuilder<Sink, Body, Enc, Hdrs, Dest>
where
Sink: PublishSink<Options = KinesisPublishOptions>,
{
fn partition_key(mut self, key: impl Into<String>) -> Self {
self.options_mut()
.get_or_insert_with(KinesisPublishOptions::default)
.partition_key = Some(key.into());
self
}
}
#[cfg(all(test, feature = "testing"))]
mod tests {
use ruststream::runtime::PublishExt;
use ruststream::testing::{InProcess, TestableBroker};
use ruststream::{HeaderMap, Outgoing, OutgoingMessage, Serialized};
use super::*;
use crate::KinesisBroker;
#[derive(Outgoing, Serialized)]
struct Payload(Vec<u8>);
fn payload() -> Payload {
Payload(b"payload".to_vec())
}
async fn connected() -> ConnectedKinesisBroker {
KinesisBroker::new()
.connect_in_process()
.await
.expect("the broker connects in process")
}
#[tokio::test]
async fn the_step_carries_the_key_through_the_publish_builder() {
let broker = connected().await;
broker
.publisher()
.message(&payload())
.to("jobs")
.partition_key("tenant-acme")
.publish()
.await
.expect("the publish succeeds");
let published = broker.published("jobs");
assert_eq!(published.len(), 1);
assert_eq!(
published[0].headers().get_str(PARTITION_KEY_HEADER),
Some("tenant-acme")
);
}
#[tokio::test]
async fn the_step_wins_over_a_key_the_call_site_wrote_by_hand() {
let broker = connected().await;
let mut headers = HeaderMap::new();
headers.insert(PARTITION_KEY_HEADER, "written-by-hand");
broker
.publisher()
.message(&payload())
.to("jobs")
.with_headers(headers)
.partition_key("named-by-the-step")
.publish()
.await
.expect("the publish succeeds");
let published = broker.published("jobs");
assert_eq!(
published[0].headers().get_str(PARTITION_KEY_HEADER),
Some("named-by-the-step")
);
}
#[tokio::test]
async fn the_steps_key_survives_a_call_that_names_other_headers() {
let broker = connected().await;
let mut headers = HeaderMap::new();
headers.insert("x-tenant", "acme");
broker
.publisher()
.message(&payload())
.to("jobs")
.with_headers(headers)
.partition_key("named-by-the-step")
.publish()
.await
.expect("the publish succeeds");
let published = broker.published("jobs");
assert_eq!(
published[0].headers().get_str(PARTITION_KEY_HEADER),
Some("named-by-the-step")
);
assert_eq!(published[0].headers().get_str("x-tenant"), Some("acme"));
}
#[tokio::test]
async fn a_publish_without_a_step_keeps_the_header_route() {
let broker = connected().await;
let mut headers = HeaderMap::new();
headers.insert(PARTITION_KEY_HEADER, "by-hand");
broker
.publisher()
.publish(
OutgoingMessage::new("jobs", b"payload").with_headers(headers),
None,
)
.await
.expect("the publish succeeds");
let published = broker.published("jobs");
assert_eq!(
published[0].headers().get_str(PARTITION_KEY_HEADER),
Some("by-hand")
);
}
#[tokio::test]
async fn a_record_no_one_named_a_key_for_still_gets_one() {
let broker = connected().await;
broker
.publisher()
.message(&payload())
.to("jobs")
.publish()
.await
.expect("the publish succeeds");
let published = broker.published("jobs");
let key = published[0]
.headers()
.get_str(PARTITION_KEY_HEADER)
.expect("the record carries a partition key");
assert!(
key.starts_with("rs-"),
"expected the spreading fallback key, got {key:?}",
);
}
#[tokio::test]
async fn the_mount_sites_key_applies_until_a_publish_names_its_own() {
let broker = connected().await;
let publisher = KinesisPublish::default()
.partition_key("named-on-the-mount")
.pair(&broker)
.await
.expect("the policy pairs with the connected broker");
publisher
.message(&payload())
.to("jobs")
.publish()
.await
.expect("the publish succeeds");
publisher
.message(&payload())
.to("jobs")
.partition_key("named-on-the-call")
.publish()
.await
.expect("the publish succeeds");
let records = broker.published("jobs");
assert_eq!(
records[0].headers().get_str(PARTITION_KEY_HEADER),
Some("named-on-the-mount")
);
assert_eq!(
records[1].headers().get_str(PARTITION_KEY_HEADER),
Some("named-on-the-call")
);
}
}