oxia/
requests.rs

1//! Fluent request builders: every [`OxiaClient`] operation returns one of
2//! these; chain options onto it and `.await` it (they implement
3//! [`IntoFuture`]), or call an alternate finisher such as
4//! [`ListBuilder::stream`].
5
6use crate::batcher::WriteOp;
7use crate::client::{OxiaClient, check_status};
8use crate::errors::OxiaError;
9use crate::proto;
10use crate::streams::{ListStream, Notifications, RangeScanStream, SequenceUpdates};
11use crate::types::{ComparisonType, GetResult, PutResult, VERSION_ID_NOT_EXISTS};
12use bytes::Bytes;
13use futures::TryStreamExt;
14use futures::future::BoxFuture;
15use std::future::IntoFuture;
16
17const DEFAULT_SUBSCRIPTION_BUFFER: usize = 100;
18
19/// A pending `put`, created with [`OxiaClient::put`]. Chain options and
20/// `.await` it.
21#[derive(Debug)]
22#[must_use = "builders do nothing unless awaited"]
23pub struct PutBuilder {
24    client: OxiaClient,
25    key: String,
26    value: Bytes,
27    expected_version_id: Option<i64>,
28    partition_key: Option<String>,
29    sequence_key_deltas: Vec<u64>,
30    secondary_indexes: Vec<proto::SecondaryIndex>,
31    ephemeral: bool,
32}
33
34impl PutBuilder {
35    pub(crate) fn new(client: OxiaClient, key: String, value: Bytes) -> Self {
36        PutBuilder {
37            client,
38            key,
39            value,
40            expected_version_id: None,
41            partition_key: None,
42            sequence_key_deltas: Vec::new(),
43            secondary_indexes: Vec::new(),
44            ephemeral: false,
45        }
46    }
47
48    /// Makes the put conditional: it succeeds only if the record's current
49    /// version id matches (compare-and-swap). Fails with
50    /// [`OxiaError::UnexpectedVersionId`] otherwise.
51    pub fn expected_version_id(mut self, version_id: i64) -> Self {
52        self.expected_version_id = Some(version_id);
53        self
54    }
55
56    /// Makes the put conditional on the record not existing yet. Fails with
57    /// [`OxiaError::UnexpectedVersionId`] if it does.
58    pub fn expected_record_not_exists(mut self) -> Self {
59        self.expected_version_id = Some(VERSION_ID_NOT_EXISTS);
60        self
61    }
62
63    /// Routes (and co-locates) the record by `partition_key` instead of by its
64    /// own key. Records sharing a partition key always live on the same shard.
65    pub fn partition_key(mut self, partition_key: impl Into<String>) -> Self {
66        self.partition_key = Some(partition_key.into());
67        self
68    }
69
70    /// Turns the key into a server-generated sequential key: each delta is
71    /// added to the current highest sequence and appended as a zero-padded
72    /// suffix.
73    ///
74    /// Requires [`partition_key`](PutBuilder::partition_key); every delta must
75    /// be greater than zero; and it cannot be combined with
76    /// [`expected_version_id`](PutBuilder::expected_version_id). Awaiting a put
77    /// that violates any of these fails with [`OxiaError::InvalidArgument`].
78    pub fn sequence_key_deltas(mut self, deltas: impl IntoIterator<Item = u64>) -> Self {
79        self.sequence_key_deltas = deltas.into_iter().collect();
80        self
81    }
82
83    /// Additionally indexes the record under `secondary_key` in the named
84    /// secondary index. May be called multiple times.
85    pub fn secondary_index(
86        mut self,
87        index_name: impl Into<String>,
88        secondary_key: impl Into<String>,
89    ) -> Self {
90        self.secondary_indexes.push(proto::SecondaryIndex {
91            index_name: index_name.into(),
92            secondary_key: secondary_key.into(),
93        });
94        self
95    }
96
97    /// Makes the record ephemeral: it is automatically deleted when this
98    /// client's session closes or expires.
99    pub fn ephemeral(mut self) -> Self {
100        self.ephemeral = true;
101        self
102    }
103
104    async fn execute(self) -> Result<PutResult, OxiaError> {
105        let client = self.client;
106        client.ensure_open()?;
107        // Validate sequence-key puts up front (matching the Java client's
108        // `PutOperation`): sequential keys require a partition key, are
109        // incompatible with an expected version, and every delta must be
110        // positive. Negative deltas are already impossible (`u64`).
111        if !self.sequence_key_deltas.is_empty() {
112            if self.partition_key.is_none() {
113                return Err(OxiaError::InvalidArgument(
114                    "sequence_key_deltas requires a partition_key".to_string(),
115                ));
116            }
117            if self.expected_version_id.is_some() {
118                return Err(OxiaError::InvalidArgument(
119                    "sequence_key_deltas cannot be combined with expected_version_id".to_string(),
120                ));
121            }
122            if self.sequence_key_deltas.contains(&0) {
123                return Err(OxiaError::InvalidArgument(
124                    "every sequence delta must be greater than zero".to_string(),
125                ));
126            }
127        }
128        let routing_key = self
129            .partition_key
130            .as_deref()
131            .unwrap_or(&self.key)
132            .to_string();
133        let PutBuilder {
134            key,
135            value,
136            expected_version_id,
137            partition_key,
138            sequence_key_deltas,
139            secondary_indexes,
140            ephemeral,
141            ..
142        } = self;
143        let identity = client.identity().to_string();
144        // Re-resolve the shard and (for ephemeral puts) its session on each
145        // attempt, so a split/merge re-routes to the new shard.
146        let response = client
147            .with_shard_retry(&routing_key, |shard| {
148                let client = client.clone();
149                let key = key.clone();
150                let value = value.clone();
151                let partition_key = partition_key.clone();
152                let sequence_key_deltas = sequence_key_deltas.clone();
153                let secondary_indexes = secondary_indexes.clone();
154                let identity = identity.clone();
155                async move {
156                    let session_id = if ephemeral {
157                        Some(client.session_id_for(shard).await?)
158                    } else {
159                        None
160                    };
161                    let request = proto::PutRequest {
162                        key,
163                        value,
164                        expected_version_id,
165                        session_id,
166                        client_identity: Some(identity),
167                        partition_key,
168                        sequence_key_delta: sequence_key_deltas,
169                        secondary_indexes,
170                        override_version_id: None,
171                        override_modifications_count: None,
172                    };
173                    client.submit_write(shard, request, WriteOp::Put).await
174                }
175            })
176            .await?;
177        check_status(response.status)?;
178        let version = response
179            .version
180            .ok_or_else(|| OxiaError::Decode("missing version in put response".to_string()))?;
181        Ok(PutResult {
182            // The server returns the key only when it generated it
183            // (sequential keys).
184            key: response.key.unwrap_or(key),
185            version: version.into(),
186        })
187    }
188}
189
190impl IntoFuture for PutBuilder {
191    type Output = Result<PutResult, OxiaError>;
192    type IntoFuture = BoxFuture<'static, Self::Output>;
193
194    fn into_future(self) -> Self::IntoFuture {
195        Box::pin(self.execute())
196    }
197}
198
199/// A pending `get`, created with [`OxiaClient::get`]. Chain options and
200/// `.await` it.
201#[derive(Debug)]
202#[must_use = "builders do nothing unless awaited"]
203pub struct GetBuilder {
204    client: OxiaClient,
205    key: String,
206    comparison: ComparisonType,
207    partition_key: Option<String>,
208    include_value: bool,
209    secondary_index_name: Option<String>,
210}
211
212impl GetBuilder {
213    pub(crate) fn new(client: OxiaClient, key: String) -> Self {
214        GetBuilder {
215            client,
216            key,
217            comparison: ComparisonType::Equal,
218            partition_key: None,
219            include_value: true,
220            secondary_index_name: None,
221        }
222    }
223
224    /// Applies a non-exact key comparison (floor/ceiling/lower/higher); the
225    /// returned record's key may then differ from the requested key.
226    pub fn comparison(mut self, comparison: ComparisonType) -> Self {
227        self.comparison = comparison;
228        self
229    }
230
231    /// Restricts the query to the shard owning `partition_key`. Required to
232    /// read records that were written with a partition key.
233    pub fn partition_key(mut self, partition_key: impl Into<String>) -> Self {
234        self.partition_key = Some(partition_key.into());
235        self
236    }
237
238    /// Excludes the value from the response (metadata-only query) when `false`.
239    /// Defaults to `true`.
240    pub fn include_value(mut self, include_value: bool) -> Self {
241        self.include_value = include_value;
242        self
243    }
244
245    /// Queries the named secondary index instead of the primary key space.
246    pub fn use_index(mut self, index_name: impl Into<String>) -> Self {
247        self.secondary_index_name = Some(index_name.into());
248        self
249    }
250
251    async fn execute(self) -> Result<GetResult, OxiaError> {
252        let client = self.client;
253        client.ensure_open()?;
254        let single_shard = is_single_shard_get(
255            self.partition_key.is_some(),
256            self.comparison,
257            self.secondary_index_name.is_some(),
258        );
259        let request = proto::GetRequest {
260            key: self.key.clone(),
261            include_value: self.include_value,
262            comparison_type: self.comparison.to_proto() as i32,
263            secondary_index_name: self.secondary_index_name,
264        };
265        if single_shard {
266            let route_key = self.partition_key.unwrap_or_else(|| self.key.clone());
267            let response = client
268                .with_shard_retry(&route_key, |shard| {
269                    let client = client.clone();
270                    let request = request.clone();
271                    async move { client.submit_get(shard, request).await }
272                })
273                .await?;
274            check_status(response.status)?;
275            GetResult::from_proto(response, Some(self.key))
276        } else {
277            client.broadcast_get(request, self.comparison).await
278        }
279    }
280}
281
282impl IntoFuture for GetBuilder {
283    type Output = Result<GetResult, OxiaError>;
284    type IntoFuture = BoxFuture<'static, Self::Output>;
285
286    fn into_future(self) -> Self::IntoFuture {
287        Box::pin(self.execute())
288    }
289}
290
291/// Whether a get is pinned to a single shard (routed by partition key, or by
292/// the primary key for an exact match) rather than broadcast to every shard.
293///
294/// A partition key always pins the query. Without one, only an exact
295/// ([`Equal`](ComparisonType::Equal)) primary-key lookup owns a single shard;
296/// an inequality comparison (its nearest match may sit on any shard) or a
297/// secondary-index lookup (indexed independently on each shard) must fan out.
298fn is_single_shard_get(
299    has_partition_key: bool,
300    comparison: ComparisonType,
301    has_secondary_index: bool,
302) -> bool {
303    has_partition_key || (comparison == ComparisonType::Equal && !has_secondary_index)
304}
305
306/// A pending `delete`, created with [`OxiaClient::delete`]. Chain options and
307/// `.await` it.
308#[derive(Debug)]
309#[must_use = "builders do nothing unless awaited"]
310pub struct DeleteBuilder {
311    client: OxiaClient,
312    key: String,
313    expected_version_id: Option<i64>,
314    partition_key: Option<String>,
315}
316
317impl DeleteBuilder {
318    pub(crate) fn new(client: OxiaClient, key: String) -> Self {
319        DeleteBuilder {
320            client,
321            key,
322            expected_version_id: None,
323            partition_key: None,
324        }
325    }
326
327    /// Makes the delete conditional: it succeeds only if the record's current
328    /// version id matches. Fails with [`OxiaError::UnexpectedVersionId`]
329    /// otherwise.
330    pub fn expected_version_id(mut self, version_id: i64) -> Self {
331        self.expected_version_id = Some(version_id);
332        self
333    }
334
335    /// Routes the delete to the shard owning `partition_key`. Required for
336    /// records that were written with a partition key.
337    pub fn partition_key(mut self, partition_key: impl Into<String>) -> Self {
338        self.partition_key = Some(partition_key.into());
339        self
340    }
341
342    async fn execute(self) -> Result<(), OxiaError> {
343        let client = self.client;
344        client.ensure_open()?;
345        let routing_key = self
346            .partition_key
347            .as_deref()
348            .unwrap_or(&self.key)
349            .to_string();
350        let request = proto::DeleteRequest {
351            key: self.key,
352            expected_version_id: self.expected_version_id,
353        };
354        let response = client
355            .with_shard_retry(&routing_key, |shard| {
356                let client = client.clone();
357                let request = request.clone();
358                async move { client.submit_write(shard, request, WriteOp::Delete).await }
359            })
360            .await?;
361        check_status(response.status)
362    }
363}
364
365impl IntoFuture for DeleteBuilder {
366    type Output = Result<(), OxiaError>;
367    type IntoFuture = BoxFuture<'static, Self::Output>;
368
369    fn into_future(self) -> Self::IntoFuture {
370        Box::pin(self.execute())
371    }
372}
373
374/// A pending `delete_range`, created with [`OxiaClient::delete_range`]. Chain
375/// options and `.await` it.
376#[derive(Debug)]
377#[must_use = "builders do nothing unless awaited"]
378pub struct DeleteRangeBuilder {
379    client: OxiaClient,
380    min_key_inclusive: String,
381    max_key_exclusive: String,
382    partition_key: Option<String>,
383}
384
385impl DeleteRangeBuilder {
386    pub(crate) fn new(
387        client: OxiaClient,
388        min_key_inclusive: String,
389        max_key_exclusive: String,
390    ) -> Self {
391        DeleteRangeBuilder {
392            client,
393            min_key_inclusive,
394            max_key_exclusive,
395            partition_key: None,
396        }
397    }
398
399    /// Restricts the range deletion to the shard owning `partition_key`
400    /// instead of applying it on every shard.
401    pub fn partition_key(mut self, partition_key: impl Into<String>) -> Self {
402        self.partition_key = Some(partition_key.into());
403        self
404    }
405
406    async fn execute(self) -> Result<(), OxiaError> {
407        let client = self.client;
408        client.ensure_open()?;
409        let request = proto::DeleteRangeRequest {
410            start_inclusive: self.min_key_inclusive,
411            end_exclusive: self.max_key_exclusive,
412        };
413        match self.partition_key {
414            Some(partition_key) => {
415                let response = client
416                    .with_shard_retry(&partition_key, |shard| {
417                        let client = client.clone();
418                        let request = request.clone();
419                        async move {
420                            client
421                                .submit_write(shard, request, WriteOp::DeleteRange)
422                                .await
423                        }
424                    })
425                    .await?;
426                check_status(response.status)
427            }
428            None => client.broadcast_delete_range(request).await,
429        }
430    }
431}
432
433impl IntoFuture for DeleteRangeBuilder {
434    type Output = Result<(), OxiaError>;
435    type IntoFuture = BoxFuture<'static, Self::Output>;
436
437    fn into_future(self) -> Self::IntoFuture {
438        Box::pin(self.execute())
439    }
440}
441
442/// A pending `list`, created with [`OxiaClient::list`].
443///
444/// `.await` it for all matching keys at once (bounded by the request timeout),
445/// or call [`stream`](ListBuilder::stream) for an incremental, ordered
446/// [`ListStream`] without an overall deadline.
447#[derive(Debug)]
448#[must_use = "builders do nothing unless awaited"]
449pub struct ListBuilder {
450    client: OxiaClient,
451    min_key_inclusive: String,
452    max_key_exclusive: String,
453    partition_key: Option<String>,
454    secondary_index_name: Option<String>,
455    include_internal_keys: bool,
456}
457
458impl ListBuilder {
459    pub(crate) fn new(
460        client: OxiaClient,
461        min_key_inclusive: String,
462        max_key_exclusive: String,
463    ) -> Self {
464        ListBuilder {
465            client,
466            min_key_inclusive,
467            max_key_exclusive,
468            partition_key: None,
469            secondary_index_name: None,
470            include_internal_keys: false,
471        }
472    }
473
474    /// Restricts the listing to the shard owning `partition_key`.
475    pub fn partition_key(mut self, partition_key: impl Into<String>) -> Self {
476        self.partition_key = Some(partition_key.into());
477        self
478    }
479
480    /// Lists keys from the named secondary index instead of the primary key
481    /// space.
482    pub fn use_index(mut self, index_name: impl Into<String>) -> Self {
483        self.secondary_index_name = Some(index_name.into());
484        self
485    }
486
487    /// Also returns Oxia's internal keys (hidden by default).
488    pub fn include_internal_keys(mut self, include: bool) -> Self {
489        self.include_internal_keys = include;
490        self
491    }
492
493    fn into_request(self) -> (OxiaClient, Option<String>, proto::ListRequest) {
494        let request = proto::ListRequest {
495            shard: None,
496            start_inclusive: self.min_key_inclusive,
497            end_exclusive: self.max_key_exclusive,
498            secondary_index_name: self.secondary_index_name,
499            include_internal_keys: self.include_internal_keys,
500        };
501        (self.client, self.partition_key, request)
502    }
503
504    /// Opens an incremental, ordered stream of the matching keys.
505    pub async fn stream(self) -> Result<ListStream, OxiaError> {
506        let (client, partition_key, request) = self.into_request();
507        client
508            .open_list_stream(request, partition_key.as_deref())
509            .await
510    }
511
512    async fn execute(self) -> Result<Vec<String>, OxiaError> {
513        let timeout = self.client.request_timeout();
514        let stream = self.stream();
515        tokio::time::timeout(timeout, async move { stream.await?.try_collect().await })
516            .await
517            .map_err(|_| OxiaError::Timeout)?
518    }
519}
520
521impl IntoFuture for ListBuilder {
522    type Output = Result<Vec<String>, OxiaError>;
523    type IntoFuture = BoxFuture<'static, Self::Output>;
524
525    fn into_future(self) -> Self::IntoFuture {
526        Box::pin(self.execute())
527    }
528}
529
530/// A pending `range_scan`, created with [`OxiaClient::range_scan`].
531///
532/// `.await` it for all matching records at once (bounded by the request
533/// timeout), or call [`stream`](RangeScanBuilder::stream) for an incremental,
534/// ordered [`RangeScanStream`] with bounded memory and no overall deadline.
535#[derive(Debug)]
536#[must_use = "builders do nothing unless awaited"]
537pub struct RangeScanBuilder {
538    client: OxiaClient,
539    min_key_inclusive: String,
540    max_key_exclusive: String,
541    partition_key: Option<String>,
542    secondary_index_name: Option<String>,
543    include_internal_keys: bool,
544}
545
546impl RangeScanBuilder {
547    pub(crate) fn new(
548        client: OxiaClient,
549        min_key_inclusive: String,
550        max_key_exclusive: String,
551    ) -> Self {
552        RangeScanBuilder {
553            client,
554            min_key_inclusive,
555            max_key_exclusive,
556            partition_key: None,
557            secondary_index_name: None,
558            include_internal_keys: false,
559        }
560    }
561
562    /// Restricts the scan to the shard owning `partition_key`.
563    pub fn partition_key(mut self, partition_key: impl Into<String>) -> Self {
564        self.partition_key = Some(partition_key.into());
565        self
566    }
567
568    /// Scans the named secondary index instead of the primary key space.
569    pub fn use_index(mut self, index_name: impl Into<String>) -> Self {
570        self.secondary_index_name = Some(index_name.into());
571        self
572    }
573
574    /// Also returns Oxia's internal keys (hidden by default).
575    pub fn include_internal_keys(mut self, include: bool) -> Self {
576        self.include_internal_keys = include;
577        self
578    }
579
580    fn into_request(self) -> (OxiaClient, Option<String>, proto::RangeScanRequest) {
581        let request = proto::RangeScanRequest {
582            shard: None,
583            start_inclusive: self.min_key_inclusive,
584            end_exclusive: self.max_key_exclusive,
585            secondary_index_name: self.secondary_index_name,
586            include_internal_keys: self.include_internal_keys,
587        };
588        (self.client, self.partition_key, request)
589    }
590
591    /// Opens an incremental, ordered stream of the matching records.
592    pub async fn stream(self) -> Result<RangeScanStream, OxiaError> {
593        let (client, partition_key, request) = self.into_request();
594        client
595            .open_range_scan_stream(request, partition_key.as_deref())
596            .await
597    }
598
599    async fn execute(self) -> Result<Vec<GetResult>, OxiaError> {
600        let timeout = self.client.request_timeout();
601        let stream = self.stream();
602        tokio::time::timeout(timeout, async move { stream.await?.try_collect().await })
603            .await
604            .map_err(|_| OxiaError::Timeout)?
605    }
606}
607
608impl IntoFuture for RangeScanBuilder {
609    type Output = Result<Vec<GetResult>, OxiaError>;
610    type IntoFuture = BoxFuture<'static, Self::Output>;
611
612    fn into_future(self) -> Self::IntoFuture {
613        Box::pin(self.execute())
614    }
615}
616
617/// A pending notifications subscription, created with
618/// [`OxiaClient::notifications`]. `.await` it to obtain the
619/// [`Notifications`] stream.
620#[derive(Debug)]
621#[must_use = "builders do nothing unless awaited"]
622pub struct NotificationsBuilder {
623    client: OxiaClient,
624    buffer_size: usize,
625}
626
627impl NotificationsBuilder {
628    pub(crate) fn new(client: OxiaClient) -> Self {
629        NotificationsBuilder {
630            client,
631            buffer_size: DEFAULT_SUBSCRIPTION_BUFFER,
632        }
633    }
634
635    /// Sets the subscription's channel capacity (default 100). When the buffer
636    /// is full, delivery applies backpressure to the per-shard listeners.
637    pub fn buffer_size(mut self, buffer_size: usize) -> Self {
638        self.buffer_size = buffer_size;
639        self
640    }
641}
642
643impl IntoFuture for NotificationsBuilder {
644    type Output = Result<Notifications, OxiaError>;
645    type IntoFuture = BoxFuture<'static, Self::Output>;
646
647    fn into_future(self) -> Self::IntoFuture {
648        Box::pin(async move { self.client.open_notifications(self.buffer_size).await })
649    }
650}
651
652/// A pending sequence-updates subscription, created with
653/// [`OxiaClient::sequence_updates`]. `.await` it to obtain the
654/// [`SequenceUpdates`] stream.
655#[derive(Debug)]
656#[must_use = "builders do nothing unless awaited"]
657pub struct SequenceUpdatesBuilder {
658    client: OxiaClient,
659    key: String,
660    partition_key: String,
661    buffer_size: usize,
662}
663
664impl SequenceUpdatesBuilder {
665    pub(crate) fn new(client: OxiaClient, key: String, partition_key: String) -> Self {
666        SequenceUpdatesBuilder {
667            client,
668            key,
669            partition_key,
670            buffer_size: DEFAULT_SUBSCRIPTION_BUFFER,
671        }
672    }
673
674    /// Sets the subscription's channel capacity (default 100).
675    pub fn buffer_size(mut self, buffer_size: usize) -> Self {
676        self.buffer_size = buffer_size;
677        self
678    }
679}
680
681impl IntoFuture for SequenceUpdatesBuilder {
682    type Output = Result<SequenceUpdates, OxiaError>;
683    type IntoFuture = BoxFuture<'static, Self::Output>;
684
685    fn into_future(self) -> Self::IntoFuture {
686        Box::pin(async move {
687            self.client
688                .open_sequence_updates(self.key, self.partition_key, self.buffer_size)
689                .await
690        })
691    }
692}
693
694#[cfg(test)]
695mod tests {
696    use super::*;
697
698    #[test]
699    fn equal_primary_key_get_is_single_shard() {
700        // The common case: exact get by key, no partition key, no index.
701        assert!(is_single_shard_get(false, ComparisonType::Equal, false));
702    }
703
704    #[test]
705    fn a_partition_key_always_pins_the_shard() {
706        for comparison in [
707            ComparisonType::Equal,
708            ComparisonType::Floor,
709            ComparisonType::Ceiling,
710            ComparisonType::Lower,
711            ComparisonType::Higher,
712        ] {
713            assert!(is_single_shard_get(true, comparison, false));
714            // Even a secondary-index lookup is pinned when a partition key routes it.
715            assert!(is_single_shard_get(true, comparison, true));
716        }
717    }
718
719    #[test]
720    fn inequality_gets_without_partition_key_broadcast() {
721        for comparison in [
722            ComparisonType::Floor,
723            ComparisonType::Ceiling,
724            ComparisonType::Lower,
725            ComparisonType::Higher,
726        ] {
727            assert!(!is_single_shard_get(false, comparison, false));
728        }
729    }
730
731    #[test]
732    fn secondary_index_get_without_partition_key_broadcasts() {
733        // Even an exact match must fan out when it targets a secondary index.
734        assert!(!is_single_shard_get(false, ComparisonType::Equal, true));
735    }
736}
oxia/requests.rs

oxia/
requests.rs
