Pager rename and refactor (#1836)

Author: analogrelay

sdk/core/azure_core/src/lib.rs

@@ -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},

sdk/cosmos/azure_data_cosmos/examples/cosmos/query.rs

@@ -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);

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

@@ -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)]
@@ -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
@@ -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
@@ -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.
     ///
@@ -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
@@ -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.
     ///
@@ -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.
     ///
@@ -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.
     ///
@@ -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.
@@ -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)
@@ -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())?;
@@ -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()]);
@@ -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");
@@ -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()]);
@@ -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()]);
@@ -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");
@@ -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.
@@ -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,
+                ))
             }
         }))
     }

sdk/cosmos/azure_data_cosmos/src/lib.rs

@@ -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;

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

@@ -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::{
@@ -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
     }
 }
 

sdk/typespec/typespec_client_core/src/http/mod.rs

@@ -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;
@@ -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};