Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Pager rename and refactor #1836

Merged
merged 10 commits into from
Oct 10, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions sdk/core/azure_core/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -51,8 +51,8 @@ pub use typespec_client_core::xml;
pub use typespec_client_core::{
base64, date,
http::{
headers::Header, new_http_client, AppendToUrlQuery, Body, Context, Continuable, HttpClient,
Method, Pageable, Request, RequestContent, StatusCode, Url,
headers::Header, new_http_client, AppendToUrlQuery, Body, Context, HttpClient, Method,
Pager, Request, RequestContent, StatusCode, Url,
},
json, parsing,
sleep::{self, sleep},
Expand Down
4 changes: 1 addition & 3 deletions sdk/cosmos/azure_data_cosmos/examples/cosmos/query.rs
Original file line number Diff line number Diff line change
Expand Up @@ -34,10 +34,8 @@ impl QueryCommand {
container_client.query_items::<serde_json::Value>(&self.query, pk, None)?;

while let Some(page) = items_pager.next().await {
let response = page?;
let response = page?.deserialize_body().await?;
println!("Results Page");
println!(" Query Metrics: {:?}", response.query_metrics);
println!(" Index Metrics: {:?}", response.index_metrics);
println!(" Items:");
for item in response.items {
println!(" * {:#?}", item);
Expand Down
88 changes: 28 additions & 60 deletions sdk/cosmos/azure_data_cosmos/src/clients/container_client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -10,8 +10,9 @@ use crate::{
ItemOptions, PartitionKey, Query, QueryPartitionStrategy,
};

use azure_core::{Context, Request, Response};
use serde::{de::DeserializeOwned, Deserialize, Serialize};
use azure_core::{Context, Pager, Request, Response};
use serde::{de::DeserializeOwned, Serialize};
use typespec_client_core::http::PagerResult;
use url::Url;

#[cfg(doc)]
Expand Down Expand Up @@ -44,13 +45,13 @@ pub trait ContainerClientMethods {
async fn read(
&self,
options: Option<ReadContainerOptions>,
) -> azure_core::Result<azure_core::Response<ContainerProperties>>;
) -> azure_core::Result<Response<ContainerProperties>>;

/// Creates a new item in the container.
///
/// # Arguments
/// * `partition_key` - The partition key of the new item.
/// * `item` - The item to create. The type must implement [`Serialize`] and [`Deserialize`]
/// * `item` - The item to create. The type must implement [`Serialize`] and [`Deserialize`](serde::Deserialize)
/// * `options` - Optional parameters for the request
///
/// # Examples
Expand Down Expand Up @@ -87,14 +88,14 @@ pub trait ContainerClientMethods {
partition_key: impl Into<PartitionKey>,
item: T,
options: Option<ItemOptions>,
) -> azure_core::Result<azure_core::Response<Item<T>>>;
) -> azure_core::Result<Response<Item<T>>>;

/// Replaces an existing item in the container.
///
/// # Arguments
/// * `partition_key` - The partition key of the item to replace.
/// * `item_id` - The id of the item to replace.
/// * `item` - The item to create. The type must implement [`Serialize`] and [`Deserialize`]
/// * `item` - The item to create. The type must implement [`Serialize`] and [`Deserialize`](serde::Deserialize)
/// * `options` - Optional parameters for the request
///
/// # Examples
Expand Down Expand Up @@ -132,7 +133,7 @@ pub trait ContainerClientMethods {
item_id: impl AsRef<str>,
item: T,
options: Option<ItemOptions>,
) -> azure_core::Result<azure_core::Response<Item<T>>>;
) -> azure_core::Result<Response<Item<T>>>;

/// Creates or replaces an item in the container.
///
Expand All @@ -141,7 +142,7 @@ pub trait ContainerClientMethods {
///
/// # Arguments
/// * `partition_key` - The partition key of the item to create or replace.
/// * `item` - The item to create. The type must implement [`Serialize`] and [`Deserialize`]
/// * `item` - The item to create. The type must implement [`Serialize`] and [`Deserialize`](serde::Deserialize)
/// * `options` - Optional parameters for the request
///
/// # Examples
Expand Down Expand Up @@ -178,7 +179,7 @@ pub trait ContainerClientMethods {
partition_key: impl Into<PartitionKey>,
item: T,
options: Option<ItemOptions>,
) -> azure_core::Result<azure_core::Response<Item<T>>>;
) -> azure_core::Result<Response<Item<T>>>;

/// Reads a specific item from the container.
///
Expand Down Expand Up @@ -216,7 +217,7 @@ pub trait ContainerClientMethods {
partition_key: impl Into<PartitionKey>,
item_id: impl AsRef<str>,
options: Option<ItemOptions>,
) -> azure_core::Result<azure_core::Response<Item<T>>>;
) -> azure_core::Result<Response<Item<T>>>;

/// Deletes an item from the container.
///
Expand All @@ -243,7 +244,7 @@ pub trait ContainerClientMethods {
partition_key: impl Into<PartitionKey>,
item_id: impl AsRef<str>,
options: Option<ItemOptions>,
) -> azure_core::Result<azure_core::Response>;
) -> azure_core::Result<Response>;

/// Executes a single-partition query against items in the container.
///
Expand Down Expand Up @@ -304,7 +305,7 @@ pub trait ContainerClientMethods {
query: impl Into<Query>,
partition_key: impl Into<QueryPartitionStrategy>,
options: Option<QueryOptions>,
) -> azure_core::Result<azure_core::Pageable<QueryResults<T>, azure_core::Error>>;
) -> azure_core::Result<Pager<QueryResults<T>>>;
}

/// A client for working with a specific container in a Cosmos DB account.
Expand Down Expand Up @@ -333,7 +334,7 @@ impl ContainerClientMethods for ContainerClient {
#[allow(unused_variables)]
// REASON: This is a documented public API so prefixing with '_' is undesirable.
options: Option<ReadContainerOptions>,
) -> azure_core::Result<azure_core::Response<ContainerProperties>> {
) -> azure_core::Result<Response<ContainerProperties>> {
let mut req = Request::new(self.container_url.clone(), azure_core::Method::Get);
self.pipeline
.send(Context::new(), &mut req, ResourceType::Containers)
Expand All @@ -348,7 +349,7 @@ impl ContainerClientMethods for ContainerClient {
#[allow(unused_variables)]
// REASON: This is a documented public API so prefixing with '_' is undesirable.
options: Option<ItemOptions>,
) -> azure_core::Result<azure_core::Response<Item<T>>> {
) -> azure_core::Result<Response<Item<T>>> {
let url = self.container_url.with_path_segments(["docs"]);
let mut req = Request::new(url, azure_core::Method::Post);
req.insert_headers(&partition_key.into())?;
Expand All @@ -367,7 +368,7 @@ impl ContainerClientMethods for ContainerClient {
#[allow(unused_variables)]
// REASON: This is a documented public API so prefixing with '_' is undesirable.
options: Option<ItemOptions>,
) -> azure_core::Result<azure_core::Response<Item<T>>> {
) -> azure_core::Result<Response<Item<T>>> {
let url = self
.container_url
.with_path_segments(["docs", item_id.as_ref()]);
Expand All @@ -387,7 +388,7 @@ impl ContainerClientMethods for ContainerClient {
#[allow(unused_variables)]
// REASON: This is a documented public API so prefixing with '_' is undesirable.
options: Option<ItemOptions>,
) -> azure_core::Result<azure_core::Response<Item<T>>> {
) -> azure_core::Result<Response<Item<T>>> {
let url = self.container_url.with_path_segments(["docs"]);
let mut req = Request::new(url, azure_core::Method::Post);
req.insert_header(constants::IS_UPSERT, "true");
Expand All @@ -406,7 +407,7 @@ impl ContainerClientMethods for ContainerClient {
#[allow(unused_variables)]
// REASON: This is a documented public API so prefixing with '_' is undesirable.
options: Option<ItemOptions>,
) -> azure_core::Result<azure_core::Response<Item<T>>> {
) -> azure_core::Result<Response<Item<T>>> {
let url = self
.container_url
.with_path_segments(["docs", item_id.as_ref()]);
Expand All @@ -425,7 +426,7 @@ impl ContainerClientMethods for ContainerClient {
#[allow(unused_variables)]
// REASON: This is a documented public API so prefixing with '_' is undesirable.
options: Option<ItemOptions>,
) -> azure_core::Result<azure_core::Response> {
) -> azure_core::Result<Response> {
let url = self
.container_url
.with_path_segments(["docs", item_id.as_ref()]);
Expand All @@ -444,26 +445,9 @@ impl ContainerClientMethods for ContainerClient {
#[allow(unused_variables)]
// REASON: This is a documented public API so prefixing with '_' is undesirable.
options: Option<QueryOptions>,
) -> azure_core::Result<azure_core::Pageable<QueryResults<T>, azure_core::Error>> {
// Represents the raw response model from the server.
// We'll use this to deserialize the response body and then convert it to a more user-friendly model.
#[derive(Deserialize)]
struct QueryResponseModel<M> {
#[serde(rename = "Documents")]
documents: Vec<M>,
}

// We have to manually implement Model, because the derive macro doesn't support auto-inferring type and lifetime bounds.
// See https://github.com/Azure/azure-sdk-for-rust/issues/1803
impl<M: DeserializeOwned> azure_core::Model for QueryResponseModel<M> {
async fn from_response_body(
body: azure_core::ResponseBody,
) -> typespec_client_core::Result<Self> {
body.json().await
}
}

let url = self.container_url.with_path_segments(["docs"]);
) -> azure_core::Result<Pager<QueryResults<T>>> {
let mut url = self.container_url.clone();
url.append_path_segments(["docs"]);
let mut base_req = Request::new(url, azure_core::Method::Post);

base_req.insert_header(constants::QUERY, "True");
Expand All @@ -477,7 +461,7 @@ impl ContainerClientMethods for ContainerClient {
// We have to double-clone here.
// First we clone the pipeline to pass it in to the closure
let pipeline = self.pipeline.clone();
Ok(azure_core::Pageable::new(move |continuation| {
Ok(Pager::from_callback(move |continuation| {
// Then we have to clone it again to pass it in to the async block.
// This is because Pageable can't borrow any data, it has to own it all.
// That's probably good, because it means a Pageable can outlive the client that produced it, but it requires some extra cloning.
Expand All @@ -488,29 +472,13 @@ impl ContainerClientMethods for ContainerClient {
req.insert_header(constants::CONTINUATION, continuation);
}

let resp: Response<QueryResponseModel<T>> = pipeline
let response = pipeline
.send(Context::new(), &mut req, ResourceType::Items)
.await?;

let query_metrics = resp
.headers()
.get_optional_string(&constants::QUERY_METRICS);
let index_metrics = resp
.headers()
.get_optional_string(&constants::INDEX_METRICS);
let continuation_token =
resp.headers().get_optional_string(&constants::CONTINUATION);

let query_response: QueryResponseModel<T> = resp.deserialize_body().await?;

let query_results = QueryResults {
items: query_response.documents,
query_metrics,
index_metrics,
continuation_token,
};

Ok(query_results)
Ok(PagerResult::from_response_header(
response,
&constants::CONTINUATION,
))
}
}))
}
Expand Down
2 changes: 1 addition & 1 deletion sdk/cosmos/azure_data_cosmos/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@
#![cfg_attr(docsrs, feature(doc_cfg_hide))]

pub mod clients;
pub(crate) mod constants;
pub mod constants;
mod options;
mod partition_key;
pub(crate) mod pipeline;
Expand Down
20 changes: 9 additions & 11 deletions sdk/cosmos/azure_data_cosmos/src/models/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,8 @@

//! Model types sent to and received from the Cosmos DB API.

use azure_core::{date::OffsetDateTime, Continuable, Model};
use serde::{Deserialize, Deserializer};
use azure_core::{date::OffsetDateTime, Model};
use serde::{de::DeserializeOwned, Deserialize, Deserializer};

#[cfg(doc)]
use crate::{
Expand Down Expand Up @@ -37,19 +37,17 @@ where

/// A page of query results, where each item is a document of type `T`.
#[non_exhaustive]
#[derive(Clone, Default, Debug)]
#[derive(Clone, Default, Debug, Deserialize)]
pub struct QueryResults<T> {
#[serde(rename = "Documents")]
pub items: Vec<T>,
pub query_metrics: Option<String>,
pub index_metrics: Option<String>,
pub continuation_token: Option<String>,
}

impl<T> Continuable for QueryResults<T> {
type Continuation = String;

fn continuation(&self) -> Option<Self::Continuation> {
self.continuation_token.clone()
impl<T: DeserializeOwned> azure_core::Model for QueryResults<T> {
async fn from_response_body(
body: azure_core::ResponseBody,
) -> typespec_client_core::Result<Self> {
body.json().await
}
}

Expand Down
4 changes: 2 additions & 2 deletions sdk/typespec/typespec_client_core/src/http/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ mod context;
pub mod headers;
mod models;
mod options;
mod pageable;
mod pager;
mod pipeline;
pub mod policies;
pub mod request;
Expand All @@ -19,7 +19,7 @@ pub use context::*;
pub use headers::Header;
pub use models::*;
pub use options::*;
pub use pageable::*;
pub use pager::*;
pub use pipeline::*;
pub use request::{Body, Request, RequestContent};
pub use response::{Model, Response};
Expand Down
Loading