Skip to content

Commit ab121ee

Browse files
charliermarshcharliermarsh
committed
Allow users to override index cache-control headers
1 parent b046e7f commit ab121ee

File tree

7 files changed

+303
-27
lines changed

7 files changed

+303
-27
lines changed

‎crates/uv-client/src/cached_client.rs‎

Lines changed: 24 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -195,17 +195,19 @@ impl<E: Into<Self> + std::error::Error + 'static> From<CachedClientError<E>> for
195195
}
196196
}
197197

198-
#[derive(Debug, Clone, Copy)]
199-
pub enum CacheControl {
198+
#[derive(Debug, Clone)]
199+
pub enum CacheControl<'a> {
200200
/// Respect the `cache-control` header from the response.
201201
None,
202202
/// Apply `max-age=0, must-revalidate` to the request.
203203
MustRevalidate,
204204
/// Allow the client to return stale responses.
205205
AllowStale,
206+
/// Override the cache control header with a custom value.
207+
Override(&'a str),
206208
}
207209

208-
impl From<Freshness> for CacheControl {
210+
impl From<Freshness> for CacheControl<'_> {
209211
fn from(value: Freshness) -> Self {
210212
match value {
211213
Freshness::Fresh => Self::None,
@@ -259,7 +261,7 @@ impl CachedClient {
259261
&self,
260262
req: Request,
261263
cache_entry: &CacheEntry,
262-
cache_control: CacheControl,
264+
cache_control: CacheControl<'_>,
263265
response_callback: Callback,
264266
) -> Result<Payload, CachedClientError<CallBackError>> {
265267
let payload = self
@@ -292,7 +294,7 @@ impl CachedClient {
292294
&self,
293295
req: Request,
294296
cache_entry: &CacheEntry,
295-
cache_control: CacheControl,
297+
cache_control: CacheControl<'_>,
296298
response_callback: Callback,
297299
) -> Result<Payload::Target, CachedClientError<CallBackError>> {
298300
let fresh_req = req.try_clone().expect("HTTP request must be cloneable");
@@ -469,7 +471,7 @@ impl CachedClient {
469471
async fn send_cached(
470472
&self,
471473
mut req: Request,
472-
cache_control: CacheControl,
474+
cache_control: CacheControl<'_>,
473475
cached: DataWithCachePolicy,
474476
) -> Result<CachedResponse, Error> {
475477
// Apply the cache control header, if necessary.
@@ -481,14 +483,21 @@ impl CachedClient {
481483
http::HeaderValue::from_static("no-cache"),
482484
);
483485
}
486+
CacheControl::Override(value) => {
487+
req.headers_mut().insert(
488+
http::header::CACHE_CONTROL,
489+
http::HeaderValue::from_str(value)
490+
.map_err(|_| ErrorKind::InvalidCacheControl(value.to_string()))?,
491+
);
492+
}
484493
}
485494
Ok(match cached.cache_policy.before_request(&mut req) {
486495
BeforeRequest::Fresh => {
487496
debug!("Found fresh response for: {}", req.url());
488497
CachedResponse::FreshCache(cached)
489498
}
490499
BeforeRequest::Stale(new_cache_policy_builder) => match cache_control {
491-
CacheControl::None | CacheControl::MustRevalidate => {
500+
CacheControl::None | CacheControl::MustRevalidate | CacheControl::Override(_) => {
492501
debug!("Found stale response for: {}", req.url());
493502
self.send_cached_handle_stale(req, cached, new_cache_policy_builder)
494503
.await?
@@ -599,7 +608,7 @@ impl CachedClient {
599608
&self,
600609
req: Request,
601610
cache_entry: &CacheEntry,
602-
cache_control: CacheControl,
611+
cache_control: CacheControl<'_>,
603612
response_callback: Callback,
604613
) -> Result<Payload, CachedClientError<CallBackError>> {
605614
let payload = self
@@ -623,7 +632,7 @@ impl CachedClient {
623632
&self,
624633
req: Request,
625634
cache_entry: &CacheEntry,
626-
cache_control: CacheControl,
635+
cache_control: CacheControl<'_>,
627636
response_callback: Callback,
628637
) -> Result<Payload::Target, CachedClientError<CallBackError>> {
629638
let mut past_retries = 0;
@@ -632,7 +641,12 @@ impl CachedClient {
632641
loop {
633642
let fresh_req = req.try_clone().expect("HTTP request must be cloneable");
634643
let result = self
635-
.get_cacheable(fresh_req, cache_entry, cache_control, &response_callback)
644+
.get_cacheable(
645+
fresh_req,
646+
cache_entry,
647+
cache_control.clone(),
648+
&response_callback,
649+
)
636650
.await;
637651

638652
// Check if the middleware already performed retries

‎crates/uv-client/src/error.rs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -259,6 +259,9 @@ pub enum ErrorKind {
259259
"Network connectivity is disabled, but the requested data wasn't found in the cache for: `{0}`"
260260
)]
261261
Offline(String),
262+
263+
#[error("Invalid cache control header: `{0}`")]
264+
InvalidCacheControl(String),
262265
}
263266

264267
impl ErrorKind {

‎crates/uv-client/src/registry_client.rs‎

Lines changed: 46 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -511,11 +511,18 @@ impl RegistryClient {
511511
format!("{package_name}.rkyv"),
512512
);
513513
let cache_control = match self.connectivity {
514-
Connectivity::Online => CacheControl::from(
515-
self.cache
516-
.freshness(&cache_entry, Some(package_name), None)
517-
.map_err(ErrorKind::Io)?,
518-
),
514+
Connectivity::Online => {
515+
// Check if this index has a custom cache control header for Simple API requests
516+
if let Some(header) = self.index_urls.simple_api_cache_control_for(index) {
517+
CacheControl::Override(header)
518+
} else {
519+
CacheControl::from(
520+
self.cache
521+
.freshness(&cache_entry, Some(package_name), None)
522+
.map_err(ErrorKind::Io)?,
523+
)
524+
}
525+
}
519526
Connectivity::Offline => CacheControl::AllowStale,
520527
};
521528

@@ -571,7 +578,7 @@ impl RegistryClient {
571578
package_name: &PackageName,
572579
url: &DisplaySafeUrl,
573580
cache_entry: &CacheEntry,
574-
cache_control: CacheControl,
581+
cache_control: CacheControl<'_>,
575582
) -> Result<OwnedArchive<SimpleMetadata>, Error> {
576583
let simple_request = self
577584
.uncached_client(url)
@@ -783,11 +790,18 @@ impl RegistryClient {
783790
format!("{}.msgpack", filename.cache_key()),
784791
);
785792
let cache_control = match self.connectivity {
786-
Connectivity::Online => CacheControl::from(
787-
self.cache
788-
.freshness(&cache_entry, Some(&filename.name), None)
789-
.map_err(ErrorKind::Io)?,
790-
),
793+
Connectivity::Online => {
794+
// Check if this index has a custom cache control header for artifact downloads
795+
if let Some(header) = self.index_urls.artifact_cache_control_for(index) {
796+
CacheControl::Override(header)
797+
} else {
798+
CacheControl::from(
799+
self.cache
800+
.freshness(&cache_entry, Some(&filename.name), None)
801+
.map_err(ErrorKind::Io)?,
802+
)
803+
}
804+
}
791805
Connectivity::Offline => CacheControl::AllowStale,
792806
};
793807

@@ -853,11 +867,26 @@ impl RegistryClient {
853867
format!("{}.msgpack", filename.cache_key()),
854868
);
855869
let cache_control = match self.connectivity {
856-
Connectivity::Online => CacheControl::from(
857-
self.cache
858-
.freshness(&cache_entry, Some(&filename.name), None)
859-
.map_err(ErrorKind::Io)?,
860-
),
870+
Connectivity::Online => {
871+
// Check if this index has a custom cache control header for artifact downloads
872+
if let Some(index) = index {
873+
if let Some(header) = self.index_urls.artifact_cache_control_for(index) {
874+
CacheControl::Override(header)
875+
} else {
876+
CacheControl::from(
877+
self.cache
878+
.freshness(&cache_entry, Some(&filename.name), None)
879+
.map_err(ErrorKind::Io)?,
880+
)
881+
}
882+
} else {
883+
CacheControl::from(
884+
self.cache
885+
.freshness(&cache_entry, Some(&filename.name), None)
886+
.map_err(ErrorKind::Io)?,
887+
)
888+
}
889+
}
861890
Connectivity::Offline => CacheControl::AllowStale,
862891
};
863892

@@ -917,7 +946,7 @@ impl RegistryClient {
917946
.get_serde_with_retry(
918947
req,
919948
&cache_entry,
920-
cache_control,
949+
cache_control.clone(),
921950
read_metadata_range_request,
922951
)
923952
.await

‎crates/uv-distribution-types/src/index.rs‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,11 +6,23 @@ use thiserror::Error;
66

77
use uv_auth::{AuthPolicy, Credentials};
88
use uv_redacted::DisplaySafeUrl;
9+
use uv_small_str::SmallString;
910

1011
use crate::index_name::{IndexName, IndexNameError};
1112
use crate::origin::Origin;
1213
use crate::{IndexStatusCodeStrategy, IndexUrl, IndexUrlError, SerializableStatusCode};
1314

15+
/// Cache control configuration for an index.
16+
#[derive(Debug, Clone, Hash, Eq, PartialEq, Ord, PartialOrd, Serialize, Deserialize, Default)]
17+
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
18+
#[serde(rename_all = "kebab-case")]
19+
pub struct IndexCacheControl {
20+
/// Cache control header for Simple API requests.
21+
pub api: Option<SmallString>,
22+
/// Cache control header for file downloads.
23+
pub files: Option<SmallString>,
24+
}
25+
1426
#[derive(Debug, Clone, Hash, Eq, PartialEq, Ord, PartialOrd, Serialize, Deserialize)]
1527
#[cfg_attr(feature = "schemars", derive(schemars::JsonSchema))]
1628
#[serde(rename_all = "kebab-case")]
@@ -104,6 +116,19 @@ pub struct Index {
104116
/// ```
105117
#[serde(default)]
106118
pub ignore_error_codes: Option<Vec<SerializableStatusCode>>,
119+
/// Cache control configuration for this index.
120+
///
121+
/// When set, these headers will override the server's cache control headers
122+
/// for both package metadata requests and artifact downloads.
123+
///
124+
/// ```toml
125+
/// [[tool.uv.index]]
126+
/// name = "my-index"
127+
/// url = "https://<omitted>/simple"
128+
/// cache-control = { api = "max-age=600", files = "max-age=3600" }
129+
/// ```
130+
#[serde(default)]
131+
pub cache_control: Option<IndexCacheControl>,
107132
}
108133

109134
#[derive(
@@ -142,6 +167,7 @@ impl Index {
142167
publish_url: None,
143168
authenticate: AuthPolicy::default(),
144169
ignore_error_codes: None,
170+
cache_control: None,
145171
}
146172
}
147173

@@ -157,6 +183,7 @@ impl Index {
157183
publish_url: None,
158184
authenticate: AuthPolicy::default(),
159185
ignore_error_codes: None,
186+
cache_control: None,
160187
}
161188
}
162189

@@ -172,6 +199,7 @@ impl Index {
172199
publish_url: None,
173200
authenticate: AuthPolicy::default(),
174201
ignore_error_codes: None,
202+
cache_control: None,
175203
}
176204
}
177205

@@ -250,6 +278,7 @@ impl From<IndexUrl> for Index {
250278
publish_url: None,
251279
authenticate: AuthPolicy::default(),
252280
ignore_error_codes: None,
281+
cache_control: None,
253282
}
254283
}
255284
}
@@ -273,6 +302,7 @@ impl FromStr for Index {
273302
publish_url: None,
274303
authenticate: AuthPolicy::default(),
275304
ignore_error_codes: None,
305+
cache_control: None,
276306
});
277307
}
278308
}
@@ -289,6 +319,7 @@ impl FromStr for Index {
289319
publish_url: None,
290320
authenticate: AuthPolicy::default(),
291321
ignore_error_codes: None,
322+
cache_control: None,
292323
})
293324
}
294325
}
@@ -384,3 +415,55 @@ pub enum IndexSourceError {
384415
#[error("Index included a name, but the name was empty")]
385416
EmptyName,
386417
}
418+
419+
#[cfg(test)]
420+
mod tests {
421+
use super::*;
422+
423+
#[test]
424+
fn test_index_cache_control_headers() {
425+
// Test that cache control headers are properly parsed from TOML
426+
let toml_str = r#"
427+
name = "test-index"
428+
url = "https://test.example.com/simple"
429+
cache-control = { api = "max-age=600", files = "max-age=3600" }
430+
"#;
431+
432+
let index: Index = toml::from_str(toml_str).unwrap();
433+
assert_eq!(index.name.as_ref().unwrap().as_ref(), "test-index");
434+
assert!(index.cache_control.is_some());
435+
let cache_control = index.cache_control.as_ref().unwrap();
436+
assert_eq!(cache_control.api.as_deref(), Some("max-age=600"));
437+
assert_eq!(cache_control.files.as_deref(), Some("max-age=3600"));
438+
}
439+
440+
#[test]
441+
fn test_index_without_cache_control() {
442+
// Test that indexes work without cache control headers
443+
let toml_str = r#"
444+
name = "test-index"
445+
url = "https://test.example.com/simple"
446+
"#;
447+
448+
let index: Index = toml::from_str(toml_str).unwrap();
449+
assert_eq!(index.name.as_ref().unwrap().as_ref(), "test-index");
450+
assert_eq!(index.cache_control, None);
451+
}
452+
453+
#[test]
454+
fn test_index_partial_cache_control() {
455+
// Test that cache control can have just one field
456+
let toml_str = r#"
457+
name = "test-index"
458+
url = "https://test.example.com/simple"
459+
cache-control = { api = "max-age=300" }
460+
"#;
461+
462+
let index: Index = toml::from_str(toml_str).unwrap();
463+
assert_eq!(index.name.as_ref().unwrap().as_ref(), "test-index");
464+
assert!(index.cache_control.is_some());
465+
let cache_control = index.cache_control.as_ref().unwrap();
466+
assert_eq!(cache_control.api.as_deref(), Some("max-age=300"));
467+
assert_eq!(cache_control.files, None);
468+
}
469+
}

0 commit comments

Comments
 (0)