Skip to main content

etcd_client/rpc/
kv.rs

1//! Etcd KV Operations.
2
3use crate::caller::{ClientCaller, ClientCallerBuilder};
4pub use crate::rpc::pb::etcdserverpb::compare::CompareResult as CompareOp;
5pub use crate::rpc::pb::etcdserverpb::range_request::{SortOrder, SortTarget};
6
7use crate::error::Result;
8use crate::intercept::InterceptedChannel;
9use crate::rpc::pb::etcdserverpb::compare::{CompareTarget, TargetUnion};
10use crate::rpc::pb::etcdserverpb::kv_client::KvClient as PbKvClient;
11use crate::rpc::pb::etcdserverpb::request_op::Request as PbTxnOp;
12use crate::rpc::pb::etcdserverpb::response_op::Response as PbTxnOpResponse;
13use crate::rpc::pb::etcdserverpb::{
14    CompactionRequest as PbCompactionRequest, CompactionRequest,
15    CompactionResponse as PbCompactionResponse, Compare as PbCompare,
16    DeleteRangeRequest as PbDeleteRequest, DeleteRangeRequest,
17    DeleteRangeResponse as PbDeleteResponse, PutRequest as PbPutRequest,
18    PutResponse as PbPutResponse, RangeRequest as PbRangeRequest, RangeResponse as PbRangeResponse,
19    RangeStreamResponse as PbRangeStreamResponse, RequestOp as PbTxnRequestOp,
20    TxnRequest as PbTxnRequest, TxnResponse as PbTxnResponse,
21};
22use crate::rpc::{get_prefix, KeyRange, KeyValue, ResponseHeader};
23use crate::vec::VecExt;
24use std::mem::ManuallyDrop;
25use std::pin::Pin;
26use std::task::{Context, Poll};
27use tokio_stream::Stream;
28use tonic::{IntoRequest, Request, Streaming};
29
30type Client = PbKvClient<InterceptedChannel>;
31
32/// Client for KV operations.
33#[repr(transparent)]
34#[derive(Clone)]
35pub struct KvClient {
36    inner: ClientCaller<Client>,
37}
38
39/// Response for `StreamRange` operation.
40#[cfg_attr(feature = "pub-response-field", visible::StructFields(pub))]
41#[derive(Debug)]
42pub struct GetStreamResponse {
43    stream: Streaming<PbRangeStreamResponse>,
44}
45
46impl GetStreamResponse {
47    #[inline]
48    const fn new(stream: Streaming<PbRangeStreamResponse>) -> Self {
49        Self { stream }
50    }
51
52    /// Fetches the next message from this stream.
53    #[inline]
54    pub async fn message(&mut self) -> Result<Option<GetResponse>> {
55        match self.stream.message().await? {
56            Some(PbRangeStreamResponse {
57                range_response: Some(resp),
58            }) => Ok(Some(GetResponse::new(resp))),
59            _ => Ok(None),
60        }
61    }
62}
63
64impl Stream for GetStreamResponse {
65    type Item = Result<GetResponse>;
66
67    #[inline]
68    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
69        Pin::new(&mut self.get_mut().stream)
70            .poll_next(cx)
71            .map(|t| match t {
72                Some(Ok(PbRangeStreamResponse {
73                    range_response: Some(resp),
74                })) => Some(Ok(GetResponse::new(resp))),
75                Some(Ok(_)) => None,
76                Some(Err(e)) => Some(Err(From::from(e))),
77                None => None,
78            })
79    }
80}
81
82impl KvClient {
83    /// Creates a kv client.
84    #[inline]
85    pub(crate) fn new(builder: ClientCallerBuilder) -> Self {
86        Self {
87            inner: builder.build(Client::new),
88        }
89    }
90
91    /// Limits the maximum size of a decoded message.
92    ///
93    /// Default: `4MB`
94    pub fn max_decoding_message_size(mut self, limit: usize) -> Self {
95        self.inner = self
96            .inner
97            .with(|client| client.max_decoding_message_size(limit));
98        self
99    }
100
101    /// Limits the maximum size of an encoded message.
102    ///
103    /// Default: `usize::MAX`
104    pub fn max_encoding_message_size(mut self, limit: usize) -> Self {
105        self.inner = self
106            .inner
107            .with(|client| client.max_encoding_message_size(limit));
108        self
109    }
110
111    /// Puts the given key into the key-value store.
112    /// A put request increments the revision of the key-value store
113    /// and generates one event in the event history.
114    #[inline]
115    pub async fn put(
116        &mut self,
117        key: impl Into<Vec<u8>>,
118        value: impl Into<Vec<u8>>,
119        options: Option<PutOptions>,
120    ) -> Result<PutResponse> {
121        async fn put_impl(client: &mut Client, req: PutOptions) -> Result<PutResponse> {
122            let resp = client.put(req).await?.into_inner();
123            Ok(PutResponse::new(resp))
124        }
125        self.inner
126            .do_call(options.unwrap_or_default().with_kv(key, value), put_impl)
127            .await
128    }
129
130    /// Gets the key or a range of keys from the store.
131    #[inline]
132    pub async fn get(
133        &mut self,
134        key: impl Into<Vec<u8>>,
135        options: Option<GetOptions>,
136    ) -> Result<GetResponse> {
137        async fn get_impl(client: &mut Client, req: GetOptions) -> Result<GetResponse> {
138            let resp = client.range(req).await?.into_inner();
139            Ok(GetResponse::new(resp))
140        }
141        self.inner
142            .do_call(options.unwrap_or_default().with_key(key.into()), get_impl)
143            .await
144    }
145
146    /// The streaming version of `get`, which was introduced in etcd v3.7.
147    /// It returns the same result as `get`, but the server splits the response into a sequence of
148    /// chunks and streams them to the client. This avoids buffering large ranges entirely in
149    /// memory on either side.
150    ///
151    /// Field population across chunks:
152    ///
153    /// - Kvs - each chunk carries a disjoint slice of the result. Concatenating the `kvs` from
154    ///   every chunk in the order they arrive yields the same key set as a single Range call
155    ///   ([`KvClient::get`]).
156    /// - Header, More, Count - populated **only** on the final chunk, and **only** when the stream
157    ///   completes without error. Earlier chunks leave these fields zero-valued.
158    ///
159    /// Read [`RangeStream`](https://etcd.io/docs/v3.7/learning/api/#rangestream) for more details.
160    #[inline]
161    pub async fn get_stream(
162        &mut self,
163        key: impl Into<Vec<u8>>,
164        options: Option<GetOptions>,
165    ) -> Result<GetStreamResponse> {
166        async fn get_stream_impl(
167            client: &mut Client,
168            req: GetOptions,
169        ) -> Result<GetStreamResponse> {
170            let resp = client.range_stream(req).await?.into_inner();
171            Ok(GetStreamResponse::new(resp))
172        }
173        self.inner
174            .do_call(
175                options.unwrap_or_default().with_key(key.into()),
176                get_stream_impl,
177            )
178            .await
179    }
180
181    /// Deletes the given key or a range of keys from the key-value store.
182    #[inline]
183    pub async fn delete(
184        &mut self,
185        key: impl Into<Vec<u8>>,
186        options: Option<DeleteOptions>,
187    ) -> Result<DeleteResponse> {
188        async fn delete_impl(client: &mut Client, req: DeleteOptions) -> Result<DeleteResponse> {
189            let resp = client.delete_range(req).await?.into_inner();
190            Ok(DeleteResponse::new(resp))
191        }
192        self.inner
193            .do_call(
194                options.unwrap_or_default().with_key(key.into()),
195                delete_impl,
196            )
197            .await
198    }
199
200    /// Compacts the event history in the etcd key-value store. The key-value
201    /// store should be periodically compacted or the event history will continue to grow
202    /// indefinitely.
203    #[inline]
204    pub async fn compact(
205        &mut self,
206        revision: i64,
207        options: Option<CompactionOptions>,
208    ) -> Result<CompactionResponse> {
209        async fn compact_impl(
210            client: &mut Client,
211            req: CompactionOptions,
212        ) -> Result<CompactionResponse> {
213            let resp = client.compact(req).await?.into_inner();
214            Ok(CompactionResponse::new(resp))
215        }
216        self.inner
217            .do_call(
218                options.unwrap_or_default().with_revision(revision),
219                compact_impl,
220            )
221            .await
222    }
223
224    /// Processes multiple operations in a single transaction.
225    /// A txn request increments the revision of the key-value store
226    /// and generates events with the same revision for every completed operation.
227    /// It is not allowed to modify the same key several times within one txn.
228    #[inline]
229    pub async fn txn(&mut self, txn: Txn) -> Result<TxnResponse> {
230        async fn txn_impl(client: &mut Client, txn: Txn) -> Result<TxnResponse> {
231            let resp = client.txn(txn).await?.into_inner();
232            Ok(TxnResponse::new(resp))
233        }
234        self.inner.do_call(txn, txn_impl).await
235    }
236}
237
238/// Options for `Put` operation.
239#[derive(Debug, Default, Clone)]
240#[repr(transparent)]
241pub struct PutOptions(PbPutRequest);
242
243impl PutOptions {
244    /// Set key-value pair.
245    #[inline]
246    fn with_kv(mut self, key: impl Into<Vec<u8>>, value: impl Into<Vec<u8>>) -> Self {
247        self.0.key = key.into();
248        self.0.value = value.into();
249        self
250    }
251
252    /// Creates a `PutOptions`.
253    #[inline]
254    pub const fn new() -> Self {
255        Self(PbPutRequest {
256            key: Vec::new(),
257            value: Vec::new(),
258            lease: 0,
259            prev_kv: false,
260            ignore_value: false,
261            ignore_lease: false,
262        })
263    }
264
265    /// Lease is the lease ID to associate with the key in the key-value store. A lease
266    /// value of 0 indicates no lease.
267    #[inline]
268    pub const fn with_lease(mut self, lease: i64) -> Self {
269        self.0.lease = lease;
270        self
271    }
272
273    /// If prev_kv is set, etcd gets the previous key-value pair before changing it.
274    /// The previous key-value pair will be returned in the put response.
275    #[inline]
276    pub const fn with_prev_key(mut self) -> Self {
277        self.0.prev_kv = true;
278        self
279    }
280
281    /// If ignore_value is set, etcd updates the key using its current value.
282    /// Returns an error if the key does not exist.
283    #[inline]
284    pub const fn with_ignore_value(mut self) -> Self {
285        self.0.ignore_value = true;
286        self
287    }
288
289    /// If ignore_lease is set, etcd updates the key using its current lease.
290    /// Returns an error if the key does not exist.
291    #[inline]
292    pub const fn with_ignore_lease(mut self) -> Self {
293        self.0.ignore_lease = true;
294        self
295    }
296}
297
298impl From<PutOptions> for PbPutRequest {
299    #[inline]
300    fn from(options: PutOptions) -> Self {
301        options.0
302    }
303}
304
305impl IntoRequest<PbPutRequest> for PutOptions {
306    #[inline]
307    fn into_request(self) -> Request<PbPutRequest> {
308        Request::new(self.into())
309    }
310}
311
312/// Response for `Put` operation.
313#[cfg_attr(feature = "pub-response-field", visible::StructFields(pub))]
314#[derive(Debug, Clone)]
315#[repr(transparent)]
316pub struct PutResponse(PbPutResponse);
317
318impl PutResponse {
319    /// Create a new `PutResponse` from pb put response.
320    #[inline]
321    const fn new(resp: PbPutResponse) -> Self {
322        Self(resp)
323    }
324
325    /// Get response header.
326    #[inline]
327    pub fn header(&self) -> Option<&ResponseHeader> {
328        self.0.header.as_ref().map(From::from)
329    }
330
331    /// Takes the header out of the response, leaving a [`None`] in its place.
332    #[inline]
333    pub fn take_header(&mut self) -> Option<ResponseHeader> {
334        self.0.header.take().map(ResponseHeader::new)
335    }
336
337    /// If prev_kv is set in the request, the previous key-value pair will be returned.
338    #[inline]
339    pub fn prev_key(&self) -> Option<&KeyValue> {
340        self.0.prev_kv.as_ref().map(From::from)
341    }
342
343    /// Takes the prev_key out of the response, leaving a [`None`] in its place.
344    #[inline]
345    pub fn take_prev_key(&mut self) -> Option<KeyValue> {
346        self.0.prev_kv.take().map(KeyValue::new)
347    }
348
349    #[inline]
350    pub(crate) fn strip_prev_key_prefix(&mut self, prefix: &[u8]) {
351        if let Some(kv) = self.0.prev_kv.as_mut() {
352            kv.key.strip_key_prefix(prefix);
353        }
354    }
355}
356
357/// Options for `Get` operation.
358#[derive(Debug, Default, Clone)]
359pub struct GetOptions {
360    req: PbRangeRequest,
361    key_range: KeyRange,
362}
363
364impl GetOptions {
365    /// Sets key.
366    #[inline]
367    fn with_key(mut self, key: impl Into<Vec<u8>>) -> Self {
368        self.key_range.with_key(key);
369        self
370    }
371
372    /// Creates a `GetOptions`.
373    #[inline]
374    pub const fn new() -> Self {
375        Self {
376            req: PbRangeRequest {
377                key: Vec::new(),
378                range_end: Vec::new(),
379                limit: 0,
380                revision: 0,
381                sort_order: 0,
382                sort_target: 0,
383                serializable: false,
384                keys_only: false,
385                count_only: false,
386                min_mod_revision: 0,
387                max_mod_revision: 0,
388                min_create_revision: 0,
389                max_create_revision: 0,
390            },
391            key_range: KeyRange::new(),
392        }
393    }
394
395    /// Specifies the range of 'Get'.
396    /// Returns the keys in the range [key, end_key).
397    /// `end_key` must be lexicographically greater than start key.
398    #[inline]
399    pub fn with_range(mut self, end_key: impl Into<Vec<u8>>) -> Self {
400        self.key_range.with_range(end_key);
401        self
402    }
403
404    /// Gets all keys >= key.
405    #[inline]
406    pub fn with_from_key(mut self) -> Self {
407        self.key_range.with_from_key();
408        self
409    }
410
411    /// Gets all keys prefixed with key.
412    #[inline]
413    pub fn with_prefix(mut self) -> Self {
414        self.key_range.with_prefix();
415        self
416    }
417
418    /// Gets all keys.
419    #[inline]
420    pub fn with_all_keys(mut self) -> Self {
421        self.key_range.with_all_keys();
422        self
423    }
424
425    /// Limits the number of keys returned for the request. When limit is set to 0,
426    /// it is treated as no limit.
427    #[inline]
428    pub const fn with_limit(mut self, limit: i64) -> Self {
429        self.req.limit = limit;
430        self
431    }
432
433    /// The point-in-time of the key-value store to use for the range.
434    /// If revision is less or equal to zero, the range is over the newest key-value store.
435    /// If the revision has been compacted, ErrCompacted is returned as a response.
436    #[inline]
437    pub const fn with_revision(mut self, revision: i64) -> Self {
438        self.req.revision = revision;
439        self
440    }
441
442    /// Sets the order for returned sorted results.
443    /// It requires 'with_range' and/or 'with_prefix' to be specified too.
444    #[inline]
445    pub fn with_sort(mut self, target: SortTarget, order: SortOrder) -> Self {
446        if target == SortTarget::Key && order == SortOrder::Ascend {
447            // If order != SortOrder::None, server fetches the entire key-space,
448            // and then applies the sort and limit, if provided.
449            // Since by default the server returns results sorted by keys
450            // in lexicographically ascending order, the client should ignore
451            // SortOrder if the target is SortTarget::Key.
452            self.req.sort_order = SortOrder::None as i32;
453        } else {
454            self.req.sort_order = order as i32;
455        }
456        self.req.sort_target = target as i32;
457        self
458    }
459
460    /// Sets the get request to use serializable member-local reads.
461    /// Get requests are linearizable by default; linearizable requests have higher
462    /// latency and lower throughput than serializable requests but reflect the current
463    /// consensus of the cluster. For better performance, in exchange for possible stale reads,
464    /// a serializable get request is served locally without needing to reach consensus
465    /// with other nodes in the cluster.
466    #[inline]
467    pub const fn with_serializable(mut self) -> Self {
468        self.req.serializable = true;
469        self
470    }
471
472    /// Returns only the keys and not the values.
473    #[inline]
474    pub const fn with_keys_only(mut self) -> Self {
475        self.req.keys_only = true;
476        self
477    }
478
479    /// Returns only the count of the keys in the range.
480    #[inline]
481    pub const fn with_count_only(mut self) -> Self {
482        self.req.count_only = true;
483        self
484    }
485
486    /// Sets the lower bound for returned key mod revisions; all keys with
487    /// lesser mod revisions will be filtered away.
488    #[inline]
489    pub const fn with_min_mod_revision(mut self, revision: i64) -> Self {
490        self.req.min_mod_revision = revision;
491        self
492    }
493
494    /// Sets the upper bound for returned key mod revisions; all keys with
495    /// greater mod revisions will be filtered away.
496    #[inline]
497    pub const fn with_max_mod_revision(mut self, revision: i64) -> Self {
498        self.req.max_mod_revision = revision;
499        self
500    }
501
502    /// Sets the lower bound for returned key create revisions; all keys with
503    /// lesser create revisions will be filtered away.
504    #[inline]
505    pub const fn with_min_create_revision(mut self, revision: i64) -> Self {
506        self.req.min_create_revision = revision;
507        self
508    }
509
510    /// `max_create_revision` is the upper bound for returned key create revisions; all keys with
511    /// greater create revisions will be filtered away.
512    #[inline]
513    pub const fn with_max_create_revision(mut self, revision: i64) -> Self {
514        self.req.max_create_revision = revision;
515        self
516    }
517
518    #[inline]
519    pub(crate) fn key_range_end_mut(&mut self) -> &mut Vec<u8> {
520        &mut self.key_range.range_end
521    }
522}
523
524impl From<GetOptions> for PbRangeRequest {
525    #[inline]
526    fn from(mut options: GetOptions) -> Self {
527        let (key, rang_end) = options.key_range.build();
528        options.req.key = key;
529        options.req.range_end = rang_end;
530        options.req
531    }
532}
533
534impl IntoRequest<PbRangeRequest> for GetOptions {
535    #[inline]
536    fn into_request(self) -> Request<PbRangeRequest> {
537        Request::new(self.into())
538    }
539}
540
541/// Response for `Get` operation.
542#[cfg_attr(feature = "pub-response-field", visible::StructFields(pub))]
543#[derive(Debug, Clone)]
544#[repr(transparent)]
545pub struct GetResponse(PbRangeResponse);
546
547impl GetResponse {
548    /// Create a new `GetResponse` from pb get response.
549    #[inline]
550    const fn new(resp: PbRangeResponse) -> Self {
551        Self(resp)
552    }
553
554    /// Get response header.
555    #[inline]
556    pub fn header(&self) -> Option<&ResponseHeader> {
557        self.0.header.as_ref().map(From::from)
558    }
559
560    /// Takes the header out of the response, leaving a [`None`] in its place.
561    #[inline]
562    pub fn take_header(&mut self) -> Option<ResponseHeader> {
563        self.0.header.take().map(ResponseHeader::new)
564    }
565
566    /// The list of key-value pairs matched by the `Get` request.
567    /// kvs is empty when count is requested.
568    #[inline]
569    pub fn kvs(&self) -> &[KeyValue] {
570        unsafe { &*(self.0.kvs.as_slice() as *const _ as *const [KeyValue]) }
571    }
572
573    /// If `kvs` is set in the request, take the key-value pairs, leaving an empty vector in its place.
574    #[inline]
575    pub fn take_kvs(&mut self) -> Vec<KeyValue> {
576        let kvs = ManuallyDrop::new(std::mem::take(&mut self.0.kvs));
577        unsafe { Vec::from_raw_parts(kvs.as_ptr() as *mut KeyValue, kvs.len(), kvs.capacity()) }
578    }
579
580    #[inline]
581    pub(crate) fn strip_kvs_prefix(&mut self, prefix: &[u8]) {
582        for kv in self.0.kvs.iter_mut() {
583            kv.key.strip_key_prefix(prefix);
584        }
585    }
586
587    /// Indicates if there are more keys to return in the requested range.
588    #[inline]
589    pub const fn more(&self) -> bool {
590        self.0.more
591    }
592
593    /// The number of keys within the range when requested.
594    ///
595    /// Unlike Kvs, it is unaffected by limits and filters (e.g., Min/Max, Create/Modify, Revisions)
596    /// and reflects the full count within the specified range.
597    #[inline]
598    pub const fn count(&self) -> i64 {
599        self.0.count
600    }
601}
602
603/// Options for `Delete` operation.
604#[derive(Debug, Default, Clone)]
605pub struct DeleteOptions {
606    req: PbDeleteRequest,
607    key_range: KeyRange,
608}
609
610impl DeleteOptions {
611    /// Sets key.
612    #[inline]
613    fn with_key(mut self, key: impl Into<Vec<u8>>) -> Self {
614        self.key_range.with_key(key);
615        self
616    }
617
618    /// Creates a `DeleteOptions`.
619    #[inline]
620    pub const fn new() -> Self {
621        Self {
622            req: PbDeleteRequest {
623                key: Vec::new(),
624                range_end: Vec::new(),
625                prev_kv: false,
626            },
627            key_range: KeyRange::new(),
628        }
629    }
630
631    /// `end_key` is the key following the last key to delete for the range [key, end_key).
632    #[inline]
633    pub fn with_range(mut self, end_key: impl Into<Vec<u8>>) -> Self {
634        self.key_range.with_range(end_key);
635        self
636    }
637
638    /// Deletes all keys >= key.
639    #[inline]
640    pub fn with_from_key(mut self) -> Self {
641        self.key_range.with_from_key();
642        self
643    }
644
645    /// Deletes all keys prefixed with key.
646    #[inline]
647    pub fn with_prefix(mut self) -> Self {
648        self.key_range.with_prefix();
649        self
650    }
651
652    /// Deletes all keys.
653    #[inline]
654    pub fn with_all_keys(mut self) -> Self {
655        self.key_range.with_all_keys();
656        self
657    }
658
659    /// If `prev_kv` is set, etcd gets the previous key-value pairs before deleting it.
660    /// The previous key-value pairs will be returned in the delete response.
661    #[inline]
662    pub const fn with_prev_key(mut self) -> Self {
663        self.req.prev_kv = true;
664        self
665    }
666
667    #[inline]
668    pub(crate) fn key_range_end_mut(&mut self) -> &mut Vec<u8> {
669        &mut self.key_range.range_end
670    }
671}
672
673impl From<DeleteOptions> for PbDeleteRequest {
674    #[inline]
675    fn from(mut options: DeleteOptions) -> Self {
676        let (key, rang_end) = options.key_range.build();
677        options.req.key = key;
678        options.req.range_end = rang_end;
679        options.req
680    }
681}
682
683impl IntoRequest<PbDeleteRequest> for DeleteOptions {
684    #[inline]
685    fn into_request(self) -> Request<DeleteRangeRequest> {
686        Request::new(self.into())
687    }
688}
689
690/// Response for `Delete` operation.
691#[cfg_attr(feature = "pub-response-field", visible::StructFields(pub))]
692#[derive(Debug, Clone)]
693#[repr(transparent)]
694pub struct DeleteResponse(PbDeleteResponse);
695
696impl DeleteResponse {
697    /// Create a new `DeleteResponse` from pb delete response.
698    #[inline]
699    const fn new(resp: PbDeleteResponse) -> Self {
700        Self(resp)
701    }
702
703    /// Delete response header.
704    #[inline]
705    pub fn header(&self) -> Option<&ResponseHeader> {
706        self.0.header.as_ref().map(From::from)
707    }
708
709    /// Takes the header out of the response, leaving a [`None`] in its place.
710    #[inline]
711    pub fn take_header(&mut self) -> Option<ResponseHeader> {
712        self.0.header.take().map(ResponseHeader::new)
713    }
714
715    /// The number of keys deleted by the delete request.
716    #[inline]
717    pub const fn deleted(&self) -> i64 {
718        self.0.deleted
719    }
720
721    /// If `prev_kv` is set in the request, the previous key-value pairs will be returned.
722    #[inline]
723    pub fn prev_kvs(&self) -> &[KeyValue] {
724        unsafe { &*(self.0.prev_kvs.as_slice() as *const _ as *const [KeyValue]) }
725    }
726
727    /// If `prev_kvs` is set in the request, take the previous key-value pairs, leaving an empty vector in its place.
728    #[inline]
729    pub fn take_prev_kvs(&mut self) -> Vec<KeyValue> {
730        let kvs = ManuallyDrop::new(std::mem::take(&mut self.0.prev_kvs));
731        unsafe { Vec::from_raw_parts(kvs.as_ptr() as *mut KeyValue, kvs.len(), kvs.capacity()) }
732    }
733
734    #[inline]
735    pub(crate) fn strip_prev_kvs_prefix(&mut self, prefix: &[u8]) {
736        for kv in self.0.prev_kvs.iter_mut() {
737            kv.key.strip_key_prefix(prefix);
738        }
739    }
740}
741
742/// Options for `Compact` operation.
743#[derive(Debug, Default, Clone)]
744#[repr(transparent)]
745pub struct CompactionOptions(PbCompactionRequest);
746
747impl CompactionOptions {
748    /// Creates a `CompactionOptions`.
749    #[inline]
750    pub const fn new() -> Self {
751        Self(PbCompactionRequest {
752            revision: 0,
753            physical: false,
754        })
755    }
756
757    /// The key-value store revision for the compaction operation.
758    #[inline]
759    const fn with_revision(mut self, revision: i64) -> Self {
760        self.0.revision = revision;
761        self
762    }
763
764    /// Physical is set so the RPC will wait until the compaction is physically
765    /// applied to the local database such that compacted entries are totally
766    /// removed from the backend database.
767    #[inline]
768    pub const fn with_physical(mut self) -> Self {
769        self.0.physical = true;
770        self
771    }
772}
773
774impl IntoRequest<PbCompactionRequest> for CompactionOptions {
775    #[inline]
776    fn into_request(self) -> Request<CompactionRequest> {
777        Request::new(self.0)
778    }
779}
780
781/// Response for `Compact` operation.
782#[cfg_attr(feature = "pub-response-field", visible::StructFields(pub))]
783#[derive(Debug, Clone)]
784#[repr(transparent)]
785pub struct CompactionResponse(PbCompactionResponse);
786
787impl CompactionResponse {
788    /// Create a new `CompactionResponse` from pb compaction response.
789    #[inline]
790    const fn new(resp: PbCompactionResponse) -> Self {
791        Self(resp)
792    }
793
794    /// Compact response header.
795    #[inline]
796    pub fn header(&self) -> Option<&ResponseHeader> {
797        self.0.header.as_ref().map(From::from)
798    }
799
800    /// Takes the header out of the response, leaving a [`None`] in its place.
801    #[inline]
802    pub fn take_header(&mut self) -> Option<ResponseHeader> {
803        self.0.header.take().map(ResponseHeader::new)
804    }
805}
806
807/// Transaction comparison.
808#[derive(Debug, Clone)]
809#[repr(transparent)]
810pub struct Compare(PbCompare);
811
812impl Compare {
813    /// Creates a new `Compare`.
814    #[inline]
815    fn new(
816        key: impl Into<Vec<u8>>,
817        cmp: CompareOp,
818        target: CompareTarget,
819        target_union: TargetUnion,
820    ) -> Self {
821        Self(PbCompare {
822            result: cmp as i32,
823            target: target as i32,
824            key: key.into(),
825            range_end: Vec::new(),
826            target_union: Some(target_union),
827        })
828    }
829
830    /// Compares the version of the given key.
831    #[inline]
832    pub fn version(key: impl Into<Vec<u8>>, cmp: CompareOp, version: i64) -> Self {
833        Self::new(
834            key,
835            cmp,
836            CompareTarget::Version,
837            TargetUnion::Version(version),
838        )
839    }
840
841    /// Compares the creation revision of the given key.
842    #[inline]
843    pub fn create_revision(key: impl Into<Vec<u8>>, cmp: CompareOp, revision: i64) -> Self {
844        Self::new(
845            key,
846            cmp,
847            CompareTarget::Create,
848            TargetUnion::CreateRevision(revision),
849        )
850    }
851
852    /// Compares the last modified revision of the given key.
853    #[inline]
854    pub fn mod_revision(key: impl Into<Vec<u8>>, cmp: CompareOp, revision: i64) -> Self {
855        Self::new(
856            key,
857            cmp,
858            CompareTarget::Mod,
859            TargetUnion::ModRevision(revision),
860        )
861    }
862
863    /// Compares the value of the given key.
864    #[inline]
865    pub fn value(key: impl Into<Vec<u8>>, cmp: CompareOp, value: impl Into<Vec<u8>>) -> Self {
866        Self::new(
867            key,
868            cmp,
869            CompareTarget::Value,
870            TargetUnion::Value(value.into()),
871        )
872    }
873
874    /// Compares the lease id of the given key.
875    #[inline]
876    pub fn lease(key: impl Into<Vec<u8>>, cmp: CompareOp, lease: i64) -> Self {
877        Self::new(key, cmp, CompareTarget::Lease, TargetUnion::Lease(lease))
878    }
879
880    /// Sets the comparison to scan the range [key, end).
881    #[inline]
882    pub fn with_range(mut self, end: impl Into<Vec<u8>>) -> Self {
883        self.0.range_end = end.into();
884        self
885    }
886
887    /// Sets the comparison to scan all keys prefixed by the key.
888    #[inline]
889    pub fn with_prefix(mut self) -> Self {
890        self.0.range_end = get_prefix(&self.0.key);
891        self
892    }
893}
894
895/// Transaction operation.
896#[derive(Debug, Clone)]
897#[repr(transparent)]
898pub struct TxnOp(PbTxnOp);
899
900impl TxnOp {
901    /// `Put` operation.
902    #[inline]
903    pub fn put(
904        key: impl Into<Vec<u8>>,
905        value: impl Into<Vec<u8>>,
906        options: Option<PutOptions>,
907    ) -> Self {
908        TxnOp(PbTxnOp::RequestPut(
909            options.unwrap_or_default().with_kv(key, value).into(),
910        ))
911    }
912
913    /// `Get` operation.
914    #[inline]
915    pub fn get(key: impl Into<Vec<u8>>, options: Option<GetOptions>) -> Self {
916        TxnOp(PbTxnOp::RequestRange(
917            options.unwrap_or_default().with_key(key).into(),
918        ))
919    }
920
921    /// `Delete` operation.
922    #[inline]
923    pub fn delete(key: impl Into<Vec<u8>>, options: Option<DeleteOptions>) -> Self {
924        TxnOp(PbTxnOp::RequestDeleteRange(
925            options.unwrap_or_default().with_key(key).into(),
926        ))
927    }
928
929    /// `Txn` operation.
930    #[inline]
931    pub fn txn(txn: Txn) -> Self {
932        TxnOp(PbTxnOp::RequestTxn(txn.into()))
933    }
934}
935
936impl From<TxnOp> for PbTxnOp {
937    #[inline]
938    fn from(op: TxnOp) -> Self {
939        op.0
940    }
941}
942
943/// Transaction of multiple operations.
944#[derive(Debug, Default, Clone)]
945pub struct Txn {
946    req: PbTxnRequest,
947    c_when: bool,
948    c_then: bool,
949    c_else: bool,
950}
951
952impl Txn {
953    /// Creates a new transaction.
954    #[inline]
955    pub const fn new() -> Self {
956        Self {
957            req: PbTxnRequest {
958                compare: Vec::new(),
959                success: Vec::new(),
960                failure: Vec::new(),
961            },
962            c_when: false,
963            c_then: false,
964            c_else: false,
965        }
966    }
967
968    /// Takes a list of comparison. If all comparisons passed in succeed,
969    /// the operations passed into `and_then()` will be executed. Or the operations
970    /// passed into `or_else()` will be executed.
971    #[inline]
972    pub fn when(mut self, compares: impl Into<Vec<Compare>>) -> Self {
973        assert!(!self.c_when, "cannot call when twice");
974        assert!(!self.c_then, "cannot call when after and_then");
975        assert!(!self.c_else, "cannot call when after or_else");
976
977        self.c_when = true;
978
979        let compares = ManuallyDrop::new(compares.into());
980        self.req.compare = unsafe {
981            Vec::from_raw_parts(
982                compares.as_ptr() as *mut PbCompare,
983                compares.len(),
984                compares.capacity(),
985            )
986        };
987
988        self
989    }
990
991    /// Takes a list of operations. The operations list will be executed, if the
992    /// comparisons passed in `when()` succeed.
993    #[inline]
994    pub fn and_then(mut self, operations: impl Into<Vec<TxnOp>>) -> Self {
995        assert!(!self.c_then, "cannot call and_then twice");
996        assert!(!self.c_else, "cannot call and_then after or_else");
997
998        self.c_then = true;
999        self.req.success = operations
1000            .into()
1001            .into_iter()
1002            .map(|op| PbTxnRequestOp {
1003                request: Some(op.into()),
1004            })
1005            .collect();
1006        self
1007    }
1008
1009    /// Takes a list of operations. The operations list will be executed, if the
1010    /// comparisons passed in `when()` fail.
1011    #[inline]
1012    pub fn or_else(mut self, operations: impl Into<Vec<TxnOp>>) -> Self {
1013        assert!(!self.c_else, "cannot call or_else twice");
1014
1015        self.c_else = true;
1016        self.req.failure = operations
1017            .into()
1018            .into_iter()
1019            .map(|op| PbTxnRequestOp {
1020                request: Some(op.into()),
1021            })
1022            .collect();
1023        self
1024    }
1025
1026    #[inline]
1027    pub(crate) fn prefix_with(&mut self, prefix: &[u8]) {
1028        self.req.prefix_with(prefix);
1029    }
1030}
1031
1032impl PbTxnRequest {
1033    fn prefix_with(&mut self, prefix: &[u8]) {
1034        let prefix_op = |op: &mut PbTxnRequestOp| {
1035            if let Some(request) = &mut op.request {
1036                match request {
1037                    PbTxnOp::RequestRange(req) => {
1038                        req.key.prefix_with(prefix);
1039                        req.range_end.prefix_range_end_with(prefix);
1040                    }
1041                    PbTxnOp::RequestPut(req) => {
1042                        req.key.prefix_with(prefix);
1043                    }
1044                    PbTxnOp::RequestDeleteRange(req) => {
1045                        req.key.prefix_with(prefix);
1046                        req.range_end.prefix_range_end_with(prefix);
1047                    }
1048                    PbTxnOp::RequestTxn(req) => {
1049                        req.prefix_with(prefix);
1050                    }
1051                }
1052            }
1053        };
1054
1055        self.compare.iter_mut().for_each(|cmp| {
1056            cmp.key.prefix_with(prefix);
1057            cmp.range_end.prefix_range_end_with(prefix);
1058        });
1059        self.success.iter_mut().for_each(prefix_op);
1060        self.failure.iter_mut().for_each(prefix_op);
1061    }
1062}
1063
1064impl From<Txn> for PbTxnRequest {
1065    #[inline]
1066    fn from(txn: Txn) -> Self {
1067        txn.req
1068    }
1069}
1070
1071impl IntoRequest<PbTxnRequest> for Txn {
1072    #[inline]
1073    fn into_request(self) -> Request<PbTxnRequest> {
1074        Request::new(self.into())
1075    }
1076}
1077
1078/// Transaction operation response.
1079#[cfg_attr(feature = "pub-response-field", visible::StructFields(pub))]
1080#[derive(Debug, Clone)]
1081pub enum TxnOpResponse {
1082    Put(PutResponse),
1083    Get(GetResponse),
1084    Delete(DeleteResponse),
1085    Txn(TxnResponse),
1086}
1087
1088/// Response for `Txn` operation.
1089#[cfg_attr(feature = "pub-response-field", visible::StructFields(pub))]
1090#[derive(Debug, Clone)]
1091#[repr(transparent)]
1092pub struct TxnResponse(PbTxnResponse);
1093
1094impl TxnResponse {
1095    /// Creates a new `Txn` response.
1096    #[inline]
1097    const fn new(resp: PbTxnResponse) -> Self {
1098        Self(resp)
1099    }
1100
1101    /// Transaction response header.
1102    #[inline]
1103    pub fn header(&self) -> Option<&ResponseHeader> {
1104        self.0.header.as_ref().map(From::from)
1105    }
1106
1107    /// Takes the header out of the response, leaving a [`None`] in its place.
1108    #[inline]
1109    pub fn take_header(&mut self) -> Option<ResponseHeader> {
1110        self.0.header.take().map(ResponseHeader::new)
1111    }
1112
1113    /// Returns `true` if the compare evaluated to true or `false` otherwise.
1114    #[inline]
1115    pub const fn succeeded(&self) -> bool {
1116        self.0.succeeded
1117    }
1118
1119    /// Returns responses of transaction operations.
1120    #[inline]
1121    pub fn op_responses(&self) -> Vec<TxnOpResponse> {
1122        self.0
1123            .responses
1124            .iter()
1125            .map(|resp| match resp.response.as_ref().unwrap() {
1126                PbTxnOpResponse::ResponsePut(put) => {
1127                    TxnOpResponse::Put(PutResponse::new(put.clone()))
1128                }
1129                PbTxnOpResponse::ResponseRange(get) => {
1130                    TxnOpResponse::Get(GetResponse::new(get.clone()))
1131                }
1132                PbTxnOpResponse::ResponseDeleteRange(delete) => {
1133                    TxnOpResponse::Delete(DeleteResponse::new(delete.clone()))
1134                }
1135                PbTxnOpResponse::ResponseTxn(txn) => {
1136                    TxnOpResponse::Txn(TxnResponse::new(txn.clone()))
1137                }
1138            })
1139            .collect()
1140    }
1141
1142    #[inline]
1143    pub(crate) fn strip_key_prefix(&mut self, prefix: &[u8]) {
1144        self.0.strip_key_prefix(prefix);
1145    }
1146}
1147
1148impl PbTxnResponse {
1149    fn strip_key_prefix(&mut self, prefix: &[u8]) {
1150        self.responses.iter_mut().for_each(|op| {
1151            if let Some(resp) = &mut op.response {
1152                match resp {
1153                    PbTxnOpResponse::ResponseRange(r) => {
1154                        for kv in r.kvs.iter_mut() {
1155                            kv.key.strip_key_prefix(prefix);
1156                        }
1157                    }
1158                    PbTxnOpResponse::ResponsePut(r) => {
1159                        if let Some(kv) = r.prev_kv.as_mut() {
1160                            kv.key.strip_key_prefix(prefix);
1161                        }
1162                    }
1163                    PbTxnOpResponse::ResponseDeleteRange(r) => {
1164                        for kv in r.prev_kvs.iter_mut() {
1165                            kv.key.strip_key_prefix(prefix);
1166                        }
1167                    }
1168                    PbTxnOpResponse::ResponseTxn(r) => {
1169                        r.strip_key_prefix(prefix);
1170                    }
1171                }
1172            }
1173        });
1174    }
1175}