ruststream-kinesis 0.7.1

Amazon Kinesis Data Streams broker implementation for the RustStream messaging framework.
Documentation
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
//! [`KinesisPublisher`], its [`KinesisPublish`] policy, and the crate's publish steps.

// Without the `testing` feature a connection link has one variant, so a `match` on it has a
// single arm; the match stays so that the in-process arm has its place when the feature is on.
#![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};

/// The per-message publish settings of this broker: what one call varies, against what the mount
/// site fixed on the policy.
///
/// Every field is optional, so a call says only what it changes, and a field no step touched
/// keeps what the policy holds. A handler body that names a setting bounds its injected slot on
/// this type - `Out<impl Publisher<Options = KinesisPublishOptions>, Marker>` - and takes the
/// steps themselves from [`KinesisPublishSteps`].
///
/// # Examples
///
/// ```
/// # #[cfg(feature = "testing")]
/// # mod demo {
/// use std::error::Error;
///
/// use ruststream::testing::TestApp;
/// use ruststream_kinesis::prelude::*;
/// use serde::{Deserialize, Serialize};
///
/// #[derive(Deserialize, Serialize, Outgoing)]
/// pub struct Order {
///     pub id: u64,
/// }
///
/// #[derive(Serialize, Outgoing)]
/// struct Journalled {
///     order: u64,
/// }
///
/// /// The slot the journal entries leave through.
/// #[derive(OutSlot)]
/// #[publishes(Journalled)]
/// struct Journal;
///
/// /// Every entry of one tenant lands on one shard, so they stay in order.
/// #[subscriber(KinesisStream::new("orders"))]
/// async fn record(
///     order: &Order,
///     Out(journal): Out<impl Publisher<Options = KinesisPublishOptions>, Journal>,
/// ) -> HandlerOutcome {
///     let sent = journal
///         .message(&Journalled { order: order.id })
///         .to("orders.journal")
///         .partition_key("tenant-acme")
///         .publish()
///         .await;
///     if sent.is_err() {
///         return HandlerOutcome::retry();
///     }
///     HandlerOutcome::ack()
/// }
///
/// /// The app `main` runs.
/// pub fn app() -> RustStream {
///     RustStream::new(AppInfo::new("orders", "0.1.0")).with_broker(KinesisBroker::new(), |b| {
///         b.include(record).out(Journal, Publish::default()).build();
///     })
/// }
///
/// pub async fn keys_each_entry() -> Result<(), Box<dyn Error>> {
///     let tb = TestApp::start(app()).await?;
///     tb.broker::<KinesisBroker>()
///         .message(&Order { id: 7 })
///         .to("orders")
///         .publish()
///         .await?;
///
///     tb.out::<Journal>()
///         .assert_called_once()
///         .with_options(&KinesisPublishOptions {
///             partition_key: Some("tenant-acme".to_owned()),
///         });
///     Ok(())
/// }
/// # }
/// # #[cfg(feature = "testing")]
/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
/// #     tokio::runtime::Builder::new_multi_thread()
/// #         .enable_all()
/// #         .build()?
/// #         .block_on(demo::keys_each_entry())
/// # }
/// # #[cfg(not(feature = "testing"))]
/// # fn main() {}
/// ```
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct KinesisPublishOptions {
    /// The record's partition key, named by
    /// [`partition_key`](KinesisPublishSteps::partition_key) on the publish builder.
    pub partition_key: Option<String>,
}

/// Publishes records to Kinesis streams; the destination is the stream's name.
///
/// The record's partition key is the unit of shard routing and per-key ordering. A publish names
/// it with [`partition_key`](KinesisPublishSteps::partition_key) on the publish builder; a record
/// that names none takes the key the mount site named on the policy
/// ([`KinesisPublish::partition_key`]), and without that a process-unique key spreads the records
/// across the shards. User headers beyond the partition key travel in a small conditional envelope
/// (Kinesis records carry only a data blob and a partition key); plain payloads stay unenveloped,
/// so any Kinesis consumer reads them, and a record from another producer reads back as
/// headerless. The envelope is the magic `RSK2`, a big-endian `u32` length of the header block,
/// the block, and the payload; each header is a `u32`-prefixed name and a `u32`-prefixed value,
/// so a value comes back byte for byte. Records in the earlier `RSK1` text envelope still read
/// back with their headers. One publish is one `PutRecord` call.
/// Buildable before `connect` and usable until `shutdown`; afterwards every publish reports
/// [`KinesisError::NotConnected`].
#[derive(Clone)]
pub struct KinesisPublisher {
    cell: CoreCell,
    /// The partition key the mount site named on the policy this publisher was paired from,
    /// used by every record that does not name one of its own.
    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,
        }
    }

    /// The publisher a policy pairs into: the same connection, plus what the mount site named
    /// on that policy.
    pub(crate) fn keyed(cell: CoreCell, partition_key: Option<Arc<str>>) -> Self {
        Self {
            cell,
            partition_key,
        }
    }

    /// What this publisher speaks over, once `connect` (or the harness) has filled the cell.
    fn link(&self) -> Result<&Link, KinesisError> {
        self.cell.get().ok_or(KinesisError::NotConnected)
    }
}

/// The record's partition key, in the order the four places that can name one resolve.
///
/// The call's step first, then the `partition-key` header a call site wrote by hand (the portable
/// spelling, and what a message crossing from another broker carries), then the key the mount site
/// named on the policy - the only place a reply and an injected slot can be given one - and last a
/// process-unique key, because a record cannot go out without one.
///
/// One function for both transports: the in-process mode answers the key the service would, never
/// one of its own.
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 {
    /// A record's data is a `Blob`, a vector the client keeps for the request, so the publisher
    /// takes the buffer the framework wrote: a publish that carries no header hands it straight
    /// through, and one that does spends it as the body of the envelope.
    type Payload = Take;
    type Error = KinesisError;
    /// The partition key is the one setting a Kinesis record varies per message, and
    /// [`KinesisPublishSteps`] is where a call site sets it.
    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());
        // The resolved key rides the record's own field, which is this transport's wire for it:
        // a delivery reads it back into `PARTITION_KEY_HEADER`, where the cross-broker
        // `Partitioned` capability finds it. Stamping the header here as well would cost a map
        // copy per publish for a value the envelope drops on the next line.
        //
        // The stream name is the caller's string and outlives the message the record is built
        // from.
        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),
            })
    }
}

/// The publish policy for [`KinesisPublisher`]: pure declaration, constructible anywhere,
/// paired with the connected broker by the runtime after `connect`.
///
/// # Examples
///
/// ```
/// # mod demo {
/// use ruststream_kinesis::KinesisPublish;
/// use ruststream_kinesis::prelude::*;
/// use serde::{Deserialize, Serialize};
///
/// #[derive(Deserialize)]
/// struct Order {
///     id: u64,
/// }
///
/// #[derive(Serialize, Outgoing)]
/// struct Journalled {
///     order: u64,
/// }
///
/// /// The slot the journal entries leave through.
/// #[derive(OutSlot)]
/// #[publishes(Journalled)]
/// struct Journal;
///
/// /// The body names no broker: the plain bound moves between brokers untouched.
/// #[subscriber(KinesisStream::new("orders"))]
/// async fn record(order: &Order, Out(journal): Out<impl Publisher, Journal>) -> HandlerOutcome {
///     let sent = journal
///         .message(&Journalled { order: order.id })
///         .to("orders.journal")
///         .publish()
///         .await;
///     if sent.is_err() {
///         return HandlerOutcome::retry();
///     }
///     HandlerOutcome::ack()
/// }
///
/// #[ruststream::app]
/// fn app() -> impl App {
///     RustStream::new(AppInfo::new("orders", "0.1.0")).with_broker(KinesisBroker::new(), |b| {
///         b.include(record).out(Journal, KinesisPublish::default()).build();
///     })
/// }
/// # }
/// # fn main() {}
/// ```
#[derive(Debug, Clone, Default, PartialEq, Eq)]
#[must_use]
pub struct KinesisPublish {
    partition_key: Option<String>,
}

impl KinesisPublish {
    /// The partition key every record published through this policy carries unless the call
    /// names its own.
    ///
    /// This is the mount site's default, and the only key a reply can be given: a replying
    /// handler returns a value and never reaches a publish builder, so
    /// [`partition_key`](KinesisPublishSteps::partition_key) has no call site there. A slot the
    /// handler publishes through has one, and a step on that call wins over this. Without either
    /// the records spread across the stream's shards under a process-unique key.
    ///
    /// One key means one shard, and with it one shard's throughput: name it when the records
    /// have to stay mutually ordered, and leave it off otherwise.
    ///
    /// # Examples
    ///
    /// ```
    /// # mod demo {
    /// use ruststream_kinesis::KinesisPublish;
    /// use ruststream_kinesis::prelude::*;
    /// use serde::{Deserialize, Serialize};
    ///
    /// #[derive(Deserialize)]
    /// struct Order {
    ///     id: u64,
    /// }
    ///
    /// #[derive(Serialize, Outgoing)]
    /// struct Journalled {
    ///     order: u64,
    /// }
    ///
    /// /// The slot the journal entries leave through.
    /// #[derive(OutSlot)]
    /// #[publishes(Journalled)]
    /// struct Journal;
    ///
    /// /// The body names no broker: the plain bound moves between brokers untouched.
    /// #[subscriber(KinesisStream::new("orders"))]
    /// async fn record(order: &Order, Out(journal): Out<impl Publisher, Journal>) -> HandlerOutcome {
    ///     let sent = journal
    ///         .message(&Journalled { order: order.id })
    ///         .to("orders.journal")
    ///         .publish()
    ///         .await;
    ///     if sent.is_err() {
    ///         return HandlerOutcome::retry();
    ///     }
    ///     HandlerOutcome::ack()
    /// }
    ///
    /// #[ruststream::app]
    /// fn app() -> impl App {
    ///     RustStream::new(AppInfo::new("orders", "0.1.0")).with_broker(KinesisBroker::new(), |b| {
    ///         // Every journal entry on one shard, mutually ordered.
    ///         let journal = KinesisPublish::default().partition_key("journal-v1");
    ///         b.include(record).out(Journal, journal).build();
    ///     })
    /// }
    /// # }
    /// # fn main() {}
    /// ```
    pub fn partition_key(mut self, key: impl Into<String>) -> Self {
        self.partition_key = Some(key.into());
        self
    }

    /// The key a publisher paired from this policy carries.
    fn key(&self) -> Option<Arc<str>> {
        self.partition_key.as_deref().map(Arc::from)
    }

    /// What the records published through this policy add to their message object in the
    /// generated document.
    ///
    /// The key a single call names is not here: the document describes the mount, and a step on
    /// the publish builder is a property of one record. A policy that names no key says nothing,
    /// so the document carries no field the service did not fix.
    #[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))
    }
}

/// The stream a position publishes to, written onto that position's channel object.
///
/// The policy reads nothing of its own here: a policy holds this broker's settings and never a
/// destination, so the stream is the one the mount site resolved and the runtime handed over -
/// the reply type's `#[outgoing(name)]`, the `reply("dest")` clause, a slot entry's name. The
/// subscription side names its stream on the same object, so both ends of a channel report the
/// Kinesis stream behind it.
#[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))
}

/// What a record published through [`KinesisPublish`] writes into the generated document.
///
/// The partition key is the record's own routing field, which is where the specification puts an
/// ordering key on the brokers it does cover; the specification lists no Kinesis binding at all,
/// so this travels as an `x-` extension.
#[cfg(feature = "asyncapi")]
#[derive(Debug, Serialize)]
struct KinesisMessageBinding<'a> {
    #[serde(rename = "partitionKey")]
    partition_key: &'a str,
}

/// What a position publishing through [`KinesisPublish`] writes into its channel object.
///
/// The stream and nothing else: a publisher holds no read settings, and the document is built
/// before anything connects, so what an existing stream is like is the service's to answer.
#[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)
    }

    /// The partition key is the record's own routing field, so it describes the message and not
    /// the channel the message goes to; the destination therefore has no part in it.
    #[cfg(feature = "asyncapi")]
    fn message_bindings(&self, _channel: &str) -> Bindings {
        self.message_binding()
    }
}

/// The Kinesis publish settings, chained on a mount site after the position that named the
/// policy: `.out_reply(policy)`, `.out_retry(policy)` or `.out(marker, policy)`.
///
/// The publish-side mirror of [`KinesisSubscriberExt`](crate::KinesisSubscriberExt): the
/// framework names the position and the policy, and what this transport does with a record
/// beyond that is the broker's own vocabulary. A call that reaches the publish builder names its
/// own key with [`KinesisPublishSteps::partition_key`]; this settles the key for every record
/// published from the position, including the reply, which has no call site of its own.
///
/// A retry copy keeps the partition key of the record it copies: the copy carries the delivery's
/// `partition-key` header, and that header wins over the key named here, so on
/// `.out_retry(policy)` this setting leaves the copy's key unchanged.
///
/// # Examples
///
/// ```
/// # mod demo {
/// use ruststream_kinesis::prelude::*;
/// use serde::{Deserialize, Serialize};
///
/// #[derive(Deserialize)]
/// struct Order {
///     id: u64,
/// }
///
/// /// The reply type names the stream it goes to; the mount site names how it is published.
/// #[derive(Serialize, Outgoing)]
/// #[outgoing(name = "receipts")]
/// struct Receipt {
///     id: u64,
/// }
///
/// #[subscriber(KinesisStream::new("orders"), reply)]
/// async fn confirm(order: &Order) -> Receipt {
///     Receipt { id: order.id }
/// }
///
/// #[ruststream::app]
/// fn app() -> impl App {
///     RustStream::new(AppInfo::new("orders", "0.1.0")).with_broker(KinesisBroker::new(), |b| {
///         b.include(confirm)
///             .out_reply(Publish::default())
///             .partition_key("tenant-acme");
///     })
/// }
/// # }
/// # fn main() {}
/// ```
///
/// The settings live on the policy the chain named, so a chain that named another broker's
/// policy - or no policy at all - does not have them.
pub trait KinesisPublishSettings: Sized {
    /// The partition key every record published from this position carries unless the call
    /// names its own. The mount-site spelling of [`KinesisPublish::partition_key`].
    // Not `#[must_use]`: on the `with_broker` path a mount chain commits when it drops, so the
    // settled chain is meant to be discarded, exactly as the framework's own steps are.
    #[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))
    }
}

/// The Kinesis publish steps, on the framework's publish builder.
///
/// A step sets one field of [`KinesisPublishOptions`] for the single record the call sends. It is
/// a position on the builder rather than a value wrapping the publisher, so the record still
/// leaves through the entry the mount site named, carrying that entry's codec and transforms and
/// attributed to that entry's slot.
///
/// The bound is on the publisher's options type, so these steps appear on a builder over a
/// Kinesis publisher and on no other broker's.
///
/// # Examples
///
/// ```
/// # mod demo {
/// use ruststream_kinesis::prelude::*;
/// use serde::{Deserialize, Serialize};
///
/// #[derive(Deserialize)]
/// struct Order {
///     id: u64,
/// }
///
/// #[derive(Serialize, Outgoing)]
/// struct Journalled {
///     order: u64,
/// }
///
/// /// The slot the journal entries leave through.
/// #[derive(OutSlot)]
/// #[publishes(Journalled)]
/// struct Journal;
///
/// /// Every entry of one tenant lands on one shard, so they stay in order.
/// #[subscriber(KinesisStream::new("orders"))]
/// async fn record(
///     order: &Order,
///     Out(journal): Out<impl Publisher<Options = KinesisPublishOptions>, Journal>,
/// ) -> HandlerOutcome {
///     let sent = journal
///         .message(&Journalled { order: order.id })
///         .to("orders.journal")
///         .partition_key("tenant-acme")
///         .publish()
///         .await;
///     if sent.is_err() {
///         return HandlerOutcome::retry();
///     }
///     HandlerOutcome::ack()
/// }
///
/// #[ruststream::app]
/// fn app() -> impl App {
///     RustStream::new(AppInfo::new("orders", "0.1.0")).with_broker(KinesisBroker::new(), |b| {
///         b.include(record).out(Journal, Publish::default()).build();
///     })
/// }
/// # }
/// # fn main() {}
/// ```
pub trait KinesisPublishSteps {
    /// Sends this one record under `key` as its partition key.
    ///
    /// The partition key picks the shard, and with it the order the record keeps against its
    /// neighbours. It wins over a `partition-key` header the call wrote by hand and over the key
    /// the mount site named on the policy.
    #[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;

    /// Bytes the test already holds encoded: a serialized type publishes them as they are, so
    /// these checks stay about the headers rather than about a codec.
    #[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")
        );
    }

    /// The step is the call's own answer, so it beats the header spelling of the same answer.
    #[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"));
    }

    /// The portable spelling stays: a message that arrives carrying the header - from another
    /// broker, or from a caller that writes headers by hand - keeps the key it carries.
    #[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")
        );
    }

    /// A record cannot go out without a key, so the bottom of the ladder is a key that spreads
    /// the records rather than no key at all.
    #[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:?}",
        );
    }

    /// The mount site's key is a default, not an override: it applies to every record that names
    /// none, and steps aside for one that does. The policy is the production one, paired against
    /// the broker connected in process - the same type a routes file names.
    #[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")
        );
    }
}