Encapsulate middleware switching in MiddlewareStack

tdejager · konstin · commit 7b88899bf712 · 2024-05-14T11:47:37.000+02:00
diff --git a/crates/uv-client/src/base_client.rs b/crates/uv-client/src/base_client.rs
@@ -8,6 +8,7 @@ use std::env;
 use std::fmt::Debug;
 use std::ops::Deref;
 use std::path::Path;
+use std::slice::Iter;
 use std::sync::Arc;
 use tracing::debug;
 use uv_auth::AuthMiddleware;
@@ -20,69 +21,76 @@ use crate::linehaul::LineHaul;
 use crate::middleware::OfflineMiddleware;
 use crate::Connectivity;
 
-#[derive(Clone, Default)]
-pub struct CustomMiddleware {
-    pub middleware: Vec<Arc<dyn Middleware>>,
+/// Newtype to implement [`Debug`] on [`Middleware`].
+#[derive(Clone)]
+pub struct MiddlewareStack(Vec<Arc<dyn Middleware>>);
+
+impl MiddlewareStack {
+    /// Use a custom middleware stack, skipping the default retry and auth layers.
+    // This function exists for rustlib users.
+    pub fn custom(middleware_stack: Vec<Arc<dyn Middleware>>) -> Self {
+        Self(middleware_stack)
+    }
+
+    /// Initialize the middleware stack, using an [`ExponentialBackoff`] with the given number of
+    /// retries and an [`AuthMiddleware`] layer with the given keyring provider.
+    pub fn new(retries: u32, keyring: KeyringProviderType) -> Self {
+        let retry_policy = ExponentialBackoff::builder().build_with_max_retries(retries);
+        let retry_strategy = RetryTransientMiddleware::new_with_policy(retry_policy);
+
+        let auth_middleware = AuthMiddleware::new().with_keyring(keyring.to_provider());
+
+        Self(vec![Arc::new(retry_strategy), Arc::new(auth_middleware)])
+    }
 }
-impl Debug for CustomMiddleware {
+
+impl Default for MiddlewareStack {
+    /// Initialize the default middleware stack, consisting of an [`ExponentialBackoff`] retry
+    /// strategy and an [`AuthMiddleware`] layer not using the keyring.
+    fn default() -> Self {
+        Self::new(3, KeyringProviderType::default())
+    }
+}
+
+impl Debug for MiddlewareStack {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         f.debug_struct("CustomMiddleware").finish()
     }
 }
 
+impl<'a> IntoIterator for &'a MiddlewareStack {
+    type Item = &'a Arc<dyn Middleware>;
+    type IntoIter = Iter<'a, Arc<dyn Middleware>>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.0.iter()
+    }
+}
+
 /// A builder for an [`BaseClient`].
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, Default)]
 pub struct BaseClientBuilder<'a> {
-    keyring: KeyringProviderType,
     native_tls: bool,
-    retries: u32,
     connectivity: Connectivity,
     client: Option<Client>,
     markers: Option<&'a MarkerEnvironment>,
     platform: Option<&'a Platform>,
-    custom_middleware: CustomMiddleware,
-}
-
-impl Default for BaseClientBuilder<'_> {
-    fn default() -> Self {
-        Self::new()
-    }
+    middleware_stack: MiddlewareStack,
 }
 
 impl BaseClientBuilder<'_> {
     pub fn new() -> Self {
-        Self {
-            keyring: KeyringProviderType::default(),
-            native_tls: false,
-            connectivity: Connectivity::Online,
-            retries: 3,
-            client: None,
-            markers: None,
-            platform: None,
-            custom_middleware: CustomMiddleware::default(),
-        }
+        Self::default()
     }
 }
 
 impl<'a> BaseClientBuilder<'a> {
-    #[must_use]
-    pub fn keyring(mut self, keyring_type: KeyringProviderType) -> Self {
-        self.keyring = keyring_type;
-        self
-    }
-
     #[must_use]
     pub fn connectivity(mut self, connectivity: Connectivity) -> Self {
         self.connectivity = connectivity;
         self
     }
 
-    #[must_use]
-    pub fn retries(mut self, retries: u32) -> Self {
-        self.retries = retries;
-        self
-    }
-
     #[must_use]
     pub fn native_tls(mut self, native_tls: bool) -> Self {
         self.native_tls = native_tls;
@@ -108,14 +116,8 @@ impl<'a> BaseClientBuilder<'a> {
     }
 
     #[must_use]
-    pub fn add_middleware<M: Middleware>(mut self, middleware: M) -> Self {
-        self.custom_middleware.middleware.push(Arc::new(middleware));
-        self
-    }
-
-    #[must_use]
-    pub fn custom_middleware(mut self, custom_middleware: CustomMiddleware) -> Self {
-        self.custom_middleware = custom_middleware;
+    pub fn middleware_stack(mut self, middleware_stack: MiddlewareStack) -> Self {
+        self.middleware_stack = middleware_stack;
         self
     }
 
@@ -187,26 +189,10 @@ impl<'a> BaseClientBuilder<'a> {
         let client = match self.connectivity {
             Connectivity::Online => {
                 let mut client = reqwest_middleware::ClientBuilder::new(client.clone());
-
-                // Use custom middleware instead if provided
-                if !self.custom_middleware.middleware.is_empty() {
-                    for middleware in &self.custom_middleware.middleware {
-                        client = client.with_arc(middleware.clone());
-                    }
-                    client.build()
-                } else {
-                    // Initialize the retry strategy.
-                    let retry_policy =
-                        ExponentialBackoff::builder().build_with_max_retries(self.retries);
-                    let retry_strategy = RetryTransientMiddleware::new_with_policy(retry_policy);
-                    let client = client.with(retry_strategy);
-
-                    // Initialize the authentication middleware to set headers.
-                    let client =
-                        client.with(AuthMiddleware::new().with_keyring(self.keyring.to_provider()));
-
-                    client.build()
+                for middleware in &self.middleware_stack {
+                    client = client.with_arc(middleware.clone());
                 }
+                client.build()
             }
             Connectivity::Offline => reqwest_middleware::ClientBuilder::new(client.clone())
                 .with(OfflineMiddleware)
diff --git a/crates/uv-client/src/lib.rs b/crates/uv-client/src/lib.rs
@@ -1,4 +1,4 @@
-pub use base_client::{BaseClient, BaseClientBuilder};
+pub use base_client::{BaseClient, BaseClientBuilder, MiddlewareStack};
 pub use cached_client::{CacheControl, CachedClient, CachedClientError, DataWithCachePolicy};
 pub use error::{BetterReqwestError, Error, ErrorKind};
 pub use flat_index::{FlatIndexClient, FlatIndexEntries, FlatIndexError};
diff --git a/crates/uv-client/src/registry_client.rs b/crates/uv-client/src/registry_client.rs
@@ -2,13 +2,11 @@ use std::collections::BTreeMap;
 use std::fmt::Debug;
 use std::path::Path;
 use std::str::FromStr;
-use std::sync::Arc;
 
 use async_http_range_reader::AsyncHttpRangeReader;
 use futures::{FutureExt, TryStreamExt};
 use http::HeaderMap;
 use reqwest::{Client, Response, StatusCode};
-use reqwest_middleware::Middleware;
 use serde::{Deserialize, Serialize};
 use tokio::io::AsyncReadExt;
 use tokio_util::compat::{FuturesAsyncReadCompatExt, TokioAsyncReadCompatExt};
@@ -24,10 +22,9 @@ use platform_tags::Platform;
 use pypi_types::{Metadata23, SimpleJson};
 use uv_cache::{Cache, CacheBucket, WheelCache};
 use uv_configuration::IndexStrategy;
-use uv_configuration::KeyringProviderType;
 use uv_normalize::PackageName;
 
-use crate::base_client::CustomMiddleware;
+use crate::base_client::MiddlewareStack;
 use crate::base_client::{BaseClient, BaseClientBuilder};
 use crate::cached_client::CacheControl;
 use crate::html::SimpleHtml;
@@ -40,31 +37,27 @@ use crate::{CachedClient, CachedClientError, Error, ErrorKind};
 pub struct RegistryClientBuilder<'a> {
     index_urls: IndexUrls,
     index_strategy: IndexStrategy,
-    keyring: KeyringProviderType,
     native_tls: bool,
-    retries: u32,
     connectivity: Connectivity,
     cache: Cache,
     client: Option<Client>,
     markers: Option<&'a MarkerEnvironment>,
     platform: Option<&'a Platform>,
-    custom_middleware: CustomMiddleware,
+    middleware_stack: MiddlewareStack,
 }
 
 impl RegistryClientBuilder<'_> {
     pub fn new(cache: Cache) -> Self {
         Self {
             index_urls: IndexUrls::default(),
             index_strategy: IndexStrategy::default(),
-            keyring: KeyringProviderType::default(),
             native_tls: false,
             cache,
             connectivity: Connectivity::Online,
-            retries: 3,
             client: None,
             markers: None,
             platform: None,
-            custom_middleware: CustomMiddleware::default(),
+            middleware_stack: MiddlewareStack::default(),
         }
     }
 }
@@ -82,24 +75,12 @@ impl<'a> RegistryClientBuilder<'a> {
         self
     }
 
-    #[must_use]
-    pub fn keyring(mut self, keyring_type: KeyringProviderType) -> Self {
-        self.keyring = keyring_type;
-        self
-    }
-
     #[must_use]
     pub fn connectivity(mut self, connectivity: Connectivity) -> Self {
         self.connectivity = connectivity;
         self
     }
 
-    #[must_use]
-    pub fn retries(mut self, retries: u32) -> Self {
-        self.retries = retries;
-        self
-    }
-
     #[must_use]
     pub fn native_tls(mut self, native_tls: bool) -> Self {
         self.native_tls = native_tls;
@@ -131,8 +112,8 @@ impl<'a> RegistryClientBuilder<'a> {
     }
 
     #[must_use]
-    pub fn add_middleware<M: Middleware>(mut self, middleware: M) -> Self {
-        self.custom_middleware.middleware.push(Arc::new(middleware));
+    pub fn middleware_stack(mut self, middleware_stack: MiddlewareStack) -> Self {
+        self.middleware_stack = middleware_stack;
         self
     }
 
@@ -153,11 +134,9 @@ impl<'a> RegistryClientBuilder<'a> {
         }
 
         let client = builder
-            .retries(self.retries)
             .connectivity(self.connectivity)
             .native_tls(self.native_tls)
-            .custom_middleware(self.custom_middleware)
-            .keyring(self.keyring)
+            .middleware_stack(self.middleware_stack)
             .build();
 
         let timeout = client.timeout();
@@ -859,9 +838,10 @@ impl MediaType {
     }
 }
 
-#[derive(Debug, Copy, Clone, Eq, PartialEq)]
+#[derive(Debug, Copy, Clone, Eq, PartialEq, Default)]
 pub enum Connectivity {
     /// Allow access to the network.
+    #[default]
     Online,
 
     /// Do not allow access to the network.
diff --git a/crates/uv/src/commands/pip/compile.rs b/crates/uv/src/commands/pip/compile.rs
@@ -26,7 +26,9 @@ use pypi_types::Metadata23;
 use requirements_txt::EditableRequirement;
 use uv_auth::store_credentials_from_url;
 use uv_cache::Cache;
-use uv_client::{BaseClientBuilder, Connectivity, FlatIndexClient, RegistryClientBuilder};
+use uv_client::{
+    BaseClientBuilder, Connectivity, FlatIndexClient, MiddlewareStack, RegistryClientBuilder,
+};
 use uv_configuration::{
     Concurrency, ConfigSettings, Constraints, IndexStrategy, NoBinary, NoBuild, Overrides,
     PreviewMode, SetupPyStrategy, Upgrade,
@@ -113,7 +115,7 @@ pub(crate) async fn pip_compile(
     let client_builder = BaseClientBuilder::new()
         .connectivity(connectivity)
         .native_tls(native_tls)
-        .keyring(keyring_provider);
+        .middleware_stack(MiddlewareStack::new(3, keyring_provider));
 
     // Read all requirements from the provided sources.
     let RequirementsSpecification {
@@ -273,7 +275,7 @@ pub(crate) async fn pip_compile(
         .connectivity(connectivity)
         .index_urls(index_locations.index_urls())
         .index_strategy(index_strategy)
-        .keyring(keyring_provider)
+        .middleware_stack(MiddlewareStack::new(3, keyring_provider))
         .markers(&markers)
         .platform(interpreter.platform())
         .build();
diff --git a/crates/uv/src/commands/pip/install.rs b/crates/uv/src/commands/pip/install.rs
@@ -22,7 +22,8 @@ use pypi_types::Yanked;
 use uv_auth::store_credentials_from_url;
 use uv_cache::Cache;
 use uv_client::{
-    BaseClientBuilder, Connectivity, FlatIndexClient, RegistryClient, RegistryClientBuilder,
+    BaseClientBuilder, Connectivity, FlatIndexClient, MiddlewareStack, RegistryClient,
+    RegistryClientBuilder,
 };
 use uv_configuration::{
     Concurrency, ConfigSettings, Constraints, IndexStrategy, NoBinary, NoBuild, Overrides,
@@ -98,7 +99,7 @@ pub(crate) async fn pip_install(
     let client_builder = BaseClientBuilder::new()
         .connectivity(connectivity)
         .native_tls(native_tls)
-        .keyring(keyring_provider);
+        .middleware_stack(MiddlewareStack::new(3, keyring_provider));
 
     // Read all requirements from the provided sources.
     let RequirementsSpecification {
@@ -287,7 +288,7 @@ pub(crate) async fn pip_install(
         .connectivity(connectivity)
         .index_urls(index_locations.index_urls())
         .index_strategy(index_strategy)
-        .keyring(keyring_provider)
+        .middleware_stack(MiddlewareStack::new(3, keyring_provider))
         .markers(&markers)
         .platform(interpreter.platform())
         .build();
diff --git a/crates/uv/src/commands/pip/sync.rs b/crates/uv/src/commands/pip/sync.rs
@@ -13,7 +13,9 @@ use platform_tags::Tags;
 use pypi_types::Yanked;
 use uv_auth::store_credentials_from_url;
 use uv_cache::Cache;
-use uv_client::{BaseClientBuilder, Connectivity, FlatIndexClient, RegistryClientBuilder};
+use uv_client::{
+    BaseClientBuilder, Connectivity, FlatIndexClient, MiddlewareStack, RegistryClientBuilder,
+};
 use uv_configuration::{
     Concurrency, ConfigSettings, IndexStrategy, NoBinary, NoBuild, PreviewMode, Reinstall,
     SetupPyStrategy,
@@ -74,7 +76,7 @@ pub(crate) async fn pip_sync(
     let client_builder = BaseClientBuilder::new()
         .connectivity(connectivity)
         .native_tls(native_tls)
-        .keyring(keyring_provider);
+        .middleware_stack(MiddlewareStack::new(3, keyring_provider));
 
     // Read all requirements from the provided sources.
     let RequirementsSpecification {
@@ -213,7 +215,7 @@ pub(crate) async fn pip_sync(
         .connectivity(connectivity)
         .index_urls(index_locations.index_urls())
         .index_strategy(index_strategy)
-        .keyring(keyring_provider)
+        .middleware_stack(MiddlewareStack::new(3, keyring_provider))
         .markers(venv.interpreter().markers())
         .platform(venv.interpreter().platform())
         .build();
diff --git a/crates/uv/src/commands/pip/uninstall.rs b/crates/uv/src/commands/pip/uninstall.rs
@@ -8,7 +8,7 @@ use tracing::debug;
 use distribution_types::{InstalledMetadata, Name, Requirement, UnresolvedRequirement};
 use pep508_rs::UnnamedRequirement;
 use uv_cache::Cache;
-use uv_client::{BaseClientBuilder, Connectivity};
+use uv_client::{BaseClientBuilder, Connectivity, MiddlewareStack};
 use uv_configuration::{KeyringProviderType, PreviewMode};
 use uv_fs::Simplified;
 use uv_interpreter::{PythonEnvironment, Target};
@@ -36,7 +36,7 @@ pub(crate) async fn pip_uninstall(
     let client_builder = BaseClientBuilder::new()
         .connectivity(connectivity)
         .native_tls(native_tls)
-        .keyring(keyring_provider);
+        .middleware_stack(MiddlewareStack::new(3, keyring_provider));
 
     // Read all requirements from the provided sources.
     let spec =
diff --git a/crates/uv/src/commands/venv.rs b/crates/uv/src/commands/venv.rs

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-pub use base_client::{BaseClient, BaseClientBuilder};`
	`1`	`+pub use base_client::{BaseClient, BaseClientBuilder, MiddlewareStack};`
`2`	`2`	`pub use cached_client::{CacheControl, CachedClient, CachedClientError, DataWithCachePolicy};`
`3`	`3`	`pub use error::{BetterReqwestError, Error, ErrorKind};`
`4`	`4`	`pub use flat_index::{FlatIndexClient, FlatIndexEntries, FlatIndexError};`