convert non-success status codes to errors in the cosmos pipeline

Author: analogrelay

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

@@ -7,10 +7,13 @@ mod signature_target;
 use std::sync::Arc;
 
 pub use authorization_policy::AuthorizationPolicy;
-use azure_core::http::{
-    request::{options::ContentType, Request},
-    response::Response,
-    ClientOptions, Context, Method, PagerState, RawResponse,
+use azure_core::{
+    error::HttpError,
+    http::{
+        request::{options::ContentType, Request},
+        response::Response,
+        ClientOptions, Context, Method, PagerState, RawResponse,
+    },
 };
 use futures::TryStreamExt;
 use serde::de::DeserializeOwned;
@@ -64,7 +67,7 @@ impl CosmosPipeline {
         resource_link: ResourceLink,
     ) -> azure_core::Result<RawResponse> {
         let ctx = ctx.with_value(resource_link);
-        self.pipeline.send(&ctx, request).await
+        self.pipeline.send(&ctx, request).await?.success().await
     }
 
     pub async fn send<T>(
@@ -75,7 +78,7 @@ impl CosmosPipeline {
     ) -> azure_core::Result<Response<T>> {
         self.send_raw(ctx, request, resource_link)
             .await
-            .map(|r| r.into())
+            .map(Into::into)
     }
 
     pub fn send_query_request<T: DeserializeOwned + Send>(
@@ -105,7 +108,7 @@ impl CosmosPipeline {
                     req.insert_header(constants::CONTINUATION, continuation);
                 }
 
-                let resp = pipeline.send(&ctx, &mut req).await?;
+                let resp = pipeline.send(&ctx, &mut req).await?.success().await?;
                 let page = FeedPage::<T>::from_response(resp).await?;
 
                 Ok(page.into())
@@ -191,3 +194,39 @@ pub(crate) fn create_base_query_request(url: Url, query: &Query) -> azure_core::
     request.set_json(query)?;
     Ok(request)
 }
+
+trait ResponseExt {
+    async fn success(self) -> azure_core::Result<Self>
+    where
+        Self: Sized;
+}
+
+impl<T, F> ResponseExt for Response<T, F> {
+    async fn success(self) -> azure_core::Result<Self> {
+        if self.status().is_success() {
+            Ok(self)
+        } else {
+            RawResponse::from(self).success().await.map(Into::into)
+        }
+    }
+}
+
+impl ResponseExt for RawResponse {
+    async fn success(self) -> azure_core::Result<Self> {
+        if self.status().is_success() {
+            Ok(self)
+        } else {
+            let http_error = HttpError::new(self).await;
+            let status = http_error.status();
+            let error_kind = azure_core::error::ErrorKind::http_response(
+                status,
+                http_error.error_code().map(ToOwned::to_owned),
+            );
+            Err(azure_core::Error::full(
+                error_kind,
+                http_error,
+                format!("Request failed with status code: {}", status),
+            ))
+        }
+    }
+}

sdk/cosmos/azure_data_cosmos/tests/framework/test_account.rs

@@ -147,7 +147,12 @@ impl TestAccount {
         // We COULD choose not to delete them and instead validate that they were deleted, but this is what I've gone with for now.
         for id in ids {
             println!("Deleting left-over database: {}", &id);
-            cosmos_client.database_client(&id).delete(None).await?;
+            cosmos_client
+                .database_client(&id)
+                .delete(None)
+                .await?
+                .success()
+                .await?;
         }
         Ok(())
     }

sdk/cosmos/azure_data_cosmos/tests/framework/test_data/mod.rs

@@ -66,7 +66,12 @@ pub async fn create_database(
 ) -> azure_core::Result<DatabaseClient> {
     // The TestAccount has a unique context_id that includes the test name.
     let db_name = account.unique_db("TestData");
-    let response = match cosmos_client.create_database(&db_name, None).await {
+    let response = match cosmos_client
+        .create_database(&db_name, None)
+        .await?
+        .success()
+        .await
+    {
         // The database creation was successful.
         Ok(props) => props,
         Err(e) if e.http_status() == Some(StatusCode::Conflict) => {

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

@@ -3,12 +3,18 @@
 
 //! HTTP responses.
 
-use crate::http::{headers::Headers, DeserializeWith, Format, JsonFormat, StatusCode};
+use crate::{
+    error::HttpError,
+    http::{headers::Headers, DeserializeWith, Format, JsonFormat, StatusCode},
+};
 use bytes::Bytes;
 use futures::{Stream, StreamExt};
 use serde::de::DeserializeOwned;
 use std::{fmt, marker::PhantomData, pin::Pin};
-use typespec::error::{ErrorKind, ResultExt};
+use typespec::{
+    error::{ErrorKind, ResultExt},
+    Error,
+};
 
 #[cfg(not(target_arch = "wasm32"))]
 pub type PinnedStream = Pin<Box<dyn Stream<Item = crate::Result<Bytes>> + Send>>;
@@ -77,6 +83,24 @@ pub struct Response<T, F = JsonFormat> {
 }
 
 impl<T, F> Response<T, F> {
+    pub async fn success(self) -> crate::Result<Self> {
+        let status = self.status();
+        if status.is_success() {
+            Ok(self)
+        } else {
+            let http_error = HttpError::new(self.raw).await;
+            let error_kind = ErrorKind::http_response(
+                status,
+                http_error.error_code().map(std::borrow::ToOwned::to_owned),
+            );
+            Err(Error::full(
+                error_kind,
+                http_error,
+                format!("server returned an error response: {status}"),
+            ))
+        }
+    }
+
     /// Get the status code from the response.
     pub fn status(&self) -> StatusCode {
         self.raw.status()