Skip to content

Commit d91c978

Browse files
authored
Added as_headers trait to ItemOptions (#2744)
* read me change * adding consistency level option * adding as headers + tests * fixing as headers * fixing as headers test + comments * fixed warnings * fixed errors * adding pr to changelog * changelog fixed * changelog fixed
1 parent 5ec8bd7 commit d91c978

File tree

4 files changed

+205
-28
lines changed

4 files changed

+205
-28
lines changed

sdk/cosmos/azure_data_cosmos/CHANGELOG.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44

55
### Features Added
66
* Added `if_match_etag` to `ItemOptions` ([#2705](https://github.com/Azure/azure-sdk-for-rust/pull/2705))
7+
* Added several more options to `ItemOptions`: `pre_triggers`, `post_triggers`, `session_token`, `consistency_level`, and `indexing_directive` ([#2744](https://github.com/Azure/azure-sdk-for-rust/pull/2744))
78

89
### Breaking Changes
910

sdk/cosmos/azure_data_cosmos/src/clients/container_client.rs

Lines changed: 6 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,6 @@ use crate::{
1212
};
1313

1414
use azure_core::http::{
15-
headers::{self},
1615
request::{options::ContentType, Request},
1716
response::Response,
1817
Method,
@@ -263,9 +262,7 @@ impl ContainerClient {
263262
let options = options.unwrap_or_default();
264263
let url = self.pipeline.url(&self.items_link);
265264
let mut req = Request::new(url, Method::Post);
266-
if !options.enable_content_response_on_write {
267-
req.insert_header(headers::PREFER, constants::PREFER_MINIMAL);
268-
}
265+
req.insert_headers(&options)?;
269266
req.insert_headers(&partition_key.into())?;
270267
req.insert_headers(&ContentType::APPLICATION_JSON)?;
271268
req.set_json(&item)?;
@@ -355,12 +352,7 @@ impl ContainerClient {
355352
let link = self.items_link.item(item_id);
356353
let url = self.pipeline.url(&link);
357354
let mut req = Request::new(url, Method::Put);
358-
if !options.enable_content_response_on_write {
359-
req.insert_header(headers::PREFER, constants::PREFER_MINIMAL);
360-
}
361-
if let Some(etag) = options.if_match_etag {
362-
req.insert_header(headers::IF_MATCH, etag);
363-
}
355+
req.insert_headers(&options)?;
364356
req.insert_headers(&partition_key.into())?;
365357
req.insert_headers(&ContentType::APPLICATION_JSON)?;
366358
req.set_json(&item)?;
@@ -447,12 +439,7 @@ impl ContainerClient {
447439
let options = options.unwrap_or_default();
448440
let url = self.pipeline.url(&self.items_link);
449441
let mut req = Request::new(url, Method::Post);
450-
if !options.enable_content_response_on_write {
451-
req.insert_header(headers::PREFER, constants::PREFER_MINIMAL);
452-
}
453-
if let Some(etag) = options.if_match_etag {
454-
req.insert_header(headers::IF_MATCH, etag);
455-
}
442+
req.insert_headers(&options)?;
456443
req.insert_header(constants::IS_UPSERT, "true");
457444
req.insert_headers(&partition_key.into())?;
458445
req.insert_headers(&ContentType::APPLICATION_JSON)?;
@@ -507,6 +494,7 @@ impl ContainerClient {
507494
let link = self.items_link.item(item_id);
508495
let url = self.pipeline.url(&link);
509496
let mut req = Request::new(url, Method::Get);
497+
req.insert_headers(&options)?;
510498
req.insert_headers(&partition_key.into())?;
511499
self.pipeline
512500
.send(options.method_options.context, &mut req, link)
@@ -543,9 +531,7 @@ impl ContainerClient {
543531
let link = self.items_link.item(item_id);
544532
let url = self.pipeline.url(&link);
545533
let mut req = Request::new(url, Method::Delete);
546-
if let Some(etag) = options.if_match_etag {
547-
req.insert_header(headers::IF_MATCH, etag);
548-
}
534+
req.insert_headers(&options)?;
549535
req.insert_headers(&partition_key.into())?;
550536
self.pipeline
551537
.send(options.method_options.context, &mut req, link)
@@ -619,12 +605,7 @@ impl ContainerClient {
619605
let link = self.items_link.item(item_id);
620606
let url = self.pipeline.url(&link);
621607
let mut req = Request::new(url, Method::Patch);
622-
if !options.enable_content_response_on_write {
623-
req.insert_header(headers::PREFER, constants::PREFER_MINIMAL);
624-
}
625-
if let Some(etag) = options.if_match_etag {
626-
req.insert_header(headers::IF_MATCH, etag);
627-
}
608+
req.insert_headers(&options)?;
628609
req.insert_headers(&partition_key.into())?;
629610
req.insert_headers(&ContentType::APPLICATION_JSON)?;
630611
req.set_json(&patch)?;

sdk/cosmos/azure_data_cosmos/src/constants.rs

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,13 @@ pub const IS_UPSERT: HeaderName = HeaderName::from_static("x-ms-documentdb-is-up
2828
pub const OFFER_THROUGHPUT: HeaderName = HeaderName::from_static("x-ms-offer-throughput");
2929
pub const OFFER_AUTOPILOT_SETTINGS: HeaderName =
3030
HeaderName::from_static("x-ms-cosmos-offer-autopilot-settings");
31+
pub const CONSISTENCY_LEVEL: HeaderName = HeaderName::from_static("x-ms-consistency-level");
32+
pub const PRE_TRIGGER_INCLUDE: HeaderName =
33+
HeaderName::from_static("x-ms-documentdb-pre-trigger-include");
34+
pub const POST_TRIGGER_INCLUDE: HeaderName =
35+
HeaderName::from_static("x-ms-documentdb-post-trigger-include");
36+
pub const SESSION_TOKEN: HeaderName = HeaderName::from_static("x-ms-session-token");
37+
pub const INDEXING_DIRECTIVE: HeaderName = HeaderName::from_static("x-ms-indexing-directive");
3138

3239
pub const QUERY_CONTENT_TYPE: ContentType = ContentType::from_static("application/query+json");
3340

sdk/cosmos/azure_data_cosmos/src/options/mod.rs

Lines changed: 191 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,13 @@
11
// Copyright (c) Microsoft Corporation. All rights reserved.
22
// Licensed under the MIT License.
33

4-
use azure_core::http::{ClientMethodOptions, ClientOptions, Etag};
5-
4+
use crate::constants;
65
use crate::models::ThroughputProperties;
6+
use azure_core::http::headers::{AsHeaders, HeaderName, HeaderValue};
7+
use azure_core::http::{headers, ClientMethodOptions, ClientOptions, Etag};
8+
use std::convert::Infallible;
9+
use std::fmt;
10+
use std::fmt::Display;
711

812
/// Options used when creating a [`CosmosClient`](crate::CosmosClient).
913
#[derive(Clone, Default)]
@@ -43,20 +47,136 @@ pub struct DeleteDatabaseOptions<'a> {
4347
pub method_options: ClientMethodOptions<'a>,
4448
}
4549

50+
/// Specifies consistency levels that can be used when working with Cosmos APIs.
51+
///
52+
/// Learn more at [Consistency Levels](https://learn.microsoft.com/azure/cosmos-db/consistency-levels)
53+
#[derive(Clone)]
54+
pub enum ConsistencyLevel {
55+
ConsistentPrefix,
56+
Eventual,
57+
Session,
58+
BoundedStaleness,
59+
Strong,
60+
}
61+
62+
impl Display for ConsistencyLevel {
63+
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
64+
let value = match self {
65+
ConsistencyLevel::ConsistentPrefix => "ConsistentPrefix",
66+
ConsistencyLevel::Eventual => "Eventual",
67+
ConsistencyLevel::Session => "Session",
68+
ConsistencyLevel::BoundedStaleness => "BoundedStaleness",
69+
ConsistencyLevel::Strong => "Strong",
70+
};
71+
write!(f, "{}", value)
72+
}
73+
}
74+
75+
/// Specifies indexing directives that can be used when working with Cosmos APIs.
76+
#[derive(Clone)]
77+
pub enum IndexingDirective {
78+
Default,
79+
Include,
80+
Exclude,
81+
}
82+
83+
impl Display for IndexingDirective {
84+
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
85+
let value = match self {
86+
IndexingDirective::Default => "Default",
87+
IndexingDirective::Include => "Include",
88+
IndexingDirective::Exclude => "Exclude",
89+
};
90+
write!(f, "{}", value)
91+
}
92+
}
93+
4694
/// Options to be passed to APIs that manipulate items.
4795
#[derive(Clone, Default)]
4896
pub struct ItemOptions<'a> {
4997
pub method_options: ClientMethodOptions<'a>,
98+
/// Triggers executed before the operation.
99+
///
100+
/// See [Triggers](https://learn.microsoft.com/rest/api/cosmos-db/triggers) for more.
101+
pub pre_triggers: Option<Vec<String>>,
102+
/// Triggers executed after the operation.
103+
///
104+
/// See [Triggers](https://learn.microsoft.com/rest/api/cosmos-db/triggers) for more.
105+
pub post_triggers: Option<Vec<String>>,
106+
/// Applies when working with Session consistency.
107+
/// Each new write request to Azure Cosmos DB is assigned a new Session Token.
108+
/// The client instance will use this token internally with each read/query request to ensure that the set consistency level is maintained.
109+
///
110+
/// See [Session Tokens](https://learn.microsoft.com/azure/cosmos-db/nosql/how-to-manage-consistency?tabs=portal%2Cdotnetv2%2Capi-async#utilize-session-tokens) for more.
111+
pub session_token: Option<String>,
112+
/// Used to specify the consistency level for the operation.
113+
///
114+
/// The default value is the consistency level set on the Cosmos DB account.
115+
/// See [Consistency Levels](https://learn.microsoft.com/azure/cosmos-db/consistency-levels)
116+
pub consistency_level: Option<ConsistencyLevel>,
117+
/// Sets indexing directive for the operation.
118+
pub indexing_directive: Option<IndexingDirective>,
50119
/// If specified, the operation will only be performed if the item matches the provided Etag.
51120
///
52-
/// See [Optimistic Concurrency Control](https://learn.microsoft.com/en-us/azure/cosmos-db/nosql/database-transactions-optimistic-concurrency#optimistic-concurrency-control) for more.
121+
/// See [Optimistic Concurrency Control](https://learn.microsoft.com/azure/cosmos-db/nosql/database-transactions-optimistic-concurrency#optimistic-concurrency-control) for more.
53122
pub if_match_etag: Option<Etag>,
54123
/// When this value is true, write operations will respond with the new value of the resource being written.
55124
///
56125
/// The default for this is `false`, which reduces the network and CPU burden that comes from serializing and deserializing the response.
57126
pub enable_content_response_on_write: bool,
58127
}
59128

129+
impl AsHeaders for ItemOptions<'_> {
130+
type Error = Infallible;
131+
type Iter = std::vec::IntoIter<(HeaderName, HeaderValue)>;
132+
133+
fn as_headers(&self) -> Result<Self::Iter, Self::Error> {
134+
let mut headers = Vec::new();
135+
136+
if let Some(pre_triggers) = &self.pre_triggers {
137+
headers.push((
138+
constants::PRE_TRIGGER_INCLUDE,
139+
pre_triggers.join(",").into(),
140+
));
141+
}
142+
143+
if let Some(post_triggers) = &self.post_triggers {
144+
headers.push((
145+
constants::POST_TRIGGER_INCLUDE,
146+
post_triggers.join(",").into(),
147+
));
148+
}
149+
150+
if let Some(session_token) = &self.session_token {
151+
headers.push((constants::SESSION_TOKEN, session_token.into()));
152+
}
153+
154+
if let Some(consistency_level) = &self.consistency_level {
155+
headers.push((
156+
constants::CONSISTENCY_LEVEL,
157+
consistency_level.to_string().into(),
158+
));
159+
}
160+
161+
if let Some(indexing_directive) = &self.indexing_directive {
162+
headers.push((
163+
constants::INDEXING_DIRECTIVE,
164+
indexing_directive.to_string().into(),
165+
));
166+
}
167+
168+
if let Some(etag) = &self.if_match_etag {
169+
headers.push((headers::IF_MATCH, etag.to_string().into()));
170+
}
171+
172+
if !self.enable_content_response_on_write {
173+
headers.push((headers::PREFER, constants::PREFER_MINIMAL));
174+
}
175+
176+
Ok(headers.into_iter())
177+
}
178+
}
179+
60180
/// Options to be passed to [`DatabaseClient::query_containers()`](crate::clients::DatabaseClient::query_containers())
61181
#[derive(Clone, Default)]
62182
pub struct QueryContainersOptions<'a> {
@@ -111,3 +231,71 @@ pub struct ReadDatabaseOptions<'a> {
111231
pub struct ThroughputOptions<'a> {
112232
pub method_options: ClientMethodOptions<'a>,
113233
}
234+
235+
#[cfg(test)]
236+
mod tests {
237+
use super::*;
238+
239+
#[test]
240+
fn item_options_as_headers() {
241+
let item_options = ItemOptions {
242+
pre_triggers: Some(vec!["PreTrigger1".to_string(), "PreTrigger2".to_string()]),
243+
post_triggers: Some(vec!["PostTrigger1".to_string(), "PostTrigger2".to_string()]),
244+
session_token: Some("SessionToken".to_string()),
245+
consistency_level: Some(ConsistencyLevel::Session),
246+
indexing_directive: Some(IndexingDirective::Include),
247+
if_match_etag: Some(Etag::from("etag_value")),
248+
enable_content_response_on_write: false,
249+
..Default::default()
250+
};
251+
252+
let headers_result: Vec<(HeaderName, HeaderValue)> =
253+
item_options.as_headers().unwrap().collect();
254+
255+
let headers_expected: Vec<(HeaderName, HeaderValue)> = vec![
256+
(
257+
constants::PRE_TRIGGER_INCLUDE,
258+
"PreTrigger1,PreTrigger2".into(),
259+
),
260+
(
261+
constants::POST_TRIGGER_INCLUDE,
262+
"PostTrigger1,PostTrigger2".into(),
263+
),
264+
(constants::SESSION_TOKEN, "SessionToken".into()),
265+
(constants::CONSISTENCY_LEVEL, "Session".into()),
266+
(constants::INDEXING_DIRECTIVE, "Include".into()),
267+
(headers::IF_MATCH, "etag_value".into()),
268+
(headers::PREFER, constants::PREFER_MINIMAL),
269+
];
270+
271+
assert_eq!(headers_result, headers_expected);
272+
}
273+
274+
#[test]
275+
fn item_options_empty_as_headers_with_content_response() {
276+
let item_options = ItemOptions::default();
277+
278+
let headers_result: Vec<(HeaderName, HeaderValue)> =
279+
item_options.as_headers().unwrap().collect();
280+
281+
let headers_expected: Vec<(HeaderName, HeaderValue)> =
282+
vec![(headers::PREFER, constants::PREFER_MINIMAL)];
283+
284+
assert_eq!(headers_result, headers_expected);
285+
}
286+
287+
#[test]
288+
fn item_options_empty_as_headers() {
289+
let item_options = ItemOptions {
290+
enable_content_response_on_write: true,
291+
..Default::default()
292+
};
293+
294+
let headers_result: Vec<(HeaderName, HeaderValue)> =
295+
item_options.as_headers().unwrap().collect();
296+
297+
let headers_expected: Vec<(HeaderName, HeaderValue)> = vec![];
298+
299+
assert_eq!(headers_result, headers_expected);
300+
}
301+
}

0 commit comments

Comments
 (0)