Skip to content

Commit b57ad17

Browse files
charliermarshcharliermarsh
authored
Allow registries to pre-provide core metadata (#15644)
## Summary This PR adds support for the `application/vnd.pyx.simple.v1` content type, similar to `application/vnd.pypi.simple.v1` with the exception that it can also include core metadata for package-versions directly.
1 parent f88aaa8 commit b57ad17

File tree

10 files changed

+457
-28
lines changed

10 files changed

+457
-28
lines changed

‎crates/uv-cache/src/lib.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1002,7 +1002,7 @@ impl CacheBucket {
10021002
Self::Interpreter => "interpreter-v4",
10031003
// Note that when bumping this, you'll also need to bump it
10041004
// in `crates/uv/tests/it/cache_clean.rs`.
1005-
Self::Simple => "simple-v16",
1005+
Self::Simple => "simple-v17",
10061006
// Note that when bumping this, you'll also need to bump it
10071007
// in `crates/uv/tests/it/cache_prune.rs`.
10081008
Self::Wheels => "wheels-v5",

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,11 @@ impl Error {
7575
ErrorKind::BadHtml { source: err, url }.into()
7676
}
7777

78+
/// Create a new error from a `MessagePack` parsing error.
79+
pub(crate) fn from_msgpack_err(err: rmp_serde::decode::Error, url: DisplaySafeUrl) -> Self {
80+
ErrorKind::BadMessagePack { source: err, url }.into()
81+
}
82+
7883
/// Returns `true` if this error corresponds to an offline error.
7984
pub(crate) fn is_offline(&self) -> bool {
8085
matches!(&*self.kind, ErrorKind::Offline(_))
@@ -251,6 +256,12 @@ pub enum ErrorKind {
251256
url: DisplaySafeUrl,
252257
},
253258

259+
#[error("Received some unexpected MessagePack from {}", url)]
260+
BadMessagePack {
261+
source: rmp_serde::decode::Error,
262+
url: DisplaySafeUrl,
263+
},
264+
254265
#[error("Failed to read zip with range requests: `{0}`")]
255266
AsyncHttpRangeReader(DisplaySafeUrl, #[source] AsyncHttpRangeReaderError),
256267

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -204,7 +204,7 @@ impl<'a> FlatIndexClient<'a> {
204204
let unarchived: Vec<File> = files
205205
.into_iter()
206206
.filter_map(|file| {
207-
match File::try_from(file, &base) {
207+
match File::try_from_pypi(file, &base) {
208208
Ok(file) => Some(file),
209209
Err(err) => {
210210
// Ignore files with unparsable version specifiers.

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

Lines changed: 160 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ use tokio::sync::{Mutex, Semaphore};
1515
use tracing::{Instrument, debug, info_span, instrument, trace, warn};
1616
use url::Url;
1717

18-
use uv_auth::Indexes;
18+
use uv_auth::{Indexes, PyxTokenStore};
1919
use uv_cache::{Cache, CacheBucket, CacheEntry, WheelCache};
2020
use uv_configuration::IndexStrategy;
2121
use uv_configuration::KeyringProviderType;
@@ -29,7 +29,7 @@ use uv_normalize::PackageName;
2929
use uv_pep440::Version;
3030
use uv_pep508::MarkerEnvironment;
3131
use uv_platform_tags::Platform;
32-
use uv_pypi_types::{PypiSimpleDetail, ResolutionMetadata};
32+
use uv_pypi_types::{PypiSimpleDetail, PyxSimpleDetail, ResolutionMetadata};
3333
use uv_redacted::DisplaySafeUrl;
3434
use uv_small_str::SmallString;
3535
use uv_torch::TorchStrategy;
@@ -173,6 +173,7 @@ impl<'a> RegistryClientBuilder<'a> {
173173
client,
174174
timeout,
175175
flat_indexes: Arc::default(),
176+
pyx_token_store: PyxTokenStore::from_settings().ok(),
176177
}
177178
}
178179

@@ -202,6 +203,7 @@ impl<'a> RegistryClientBuilder<'a> {
202203
client,
203204
timeout,
204205
flat_indexes: Arc::default(),
206+
pyx_token_store: PyxTokenStore::from_settings().ok(),
205207
}
206208
}
207209
}
@@ -225,6 +227,9 @@ pub struct RegistryClient {
225227
timeout: Duration,
226228
/// The flat index entries for each `--find-links`-style index URL.
227229
flat_indexes: Arc<Mutex<FlatIndexCache>>,
230+
/// The pyx token store to use for persistent credentials.
231+
// TODO(charlie): The token store is only needed for `is_known_url`; can we avoid storing it here?
232+
pyx_token_store: Option<PyxTokenStore>,
228233
}
229234

230235
/// The format of the package metadata returned by querying an index.
@@ -512,7 +517,7 @@ impl RegistryClient {
512517
let result = if matches!(index, IndexUrl::Path(_)) {
513518
self.fetch_local_index(package_name, &url).await
514519
} else {
515-
self.fetch_remote_index(package_name, &url, &cache_entry, cache_control)
520+
self.fetch_remote_index(package_name, &url, index, &cache_entry, cache_control)
516521
.await
517522
};
518523

@@ -553,14 +558,27 @@ impl RegistryClient {
553558
&self,
554559
package_name: &PackageName,
555560
url: &DisplaySafeUrl,
561+
index: &IndexUrl,
556562
cache_entry: &CacheEntry,
557563
cache_control: CacheControl<'_>,
558564
) -> Result<OwnedArchive<SimpleMetadata>, Error> {
565+
// In theory, we should be able to pass `MediaType::all()` to all registries, and as
566+
// unsupported media types should be ignored by the server. For now, we implement this
567+
// defensively to avoid issues with misconfigured servers.
568+
let accept = if self
569+
.pyx_token_store
570+
.as_ref()
571+
.is_some_and(|token_store| token_store.is_known_url(index.url()))
572+
{
573+
MediaType::all()
574+
} else {
575+
MediaType::pypi()
576+
};
559577
let simple_request = self
560578
.uncached_client(url)
561579
.get(Url::from(url.clone()))
562580
.header("Accept-Encoding", "gzip, deflate, zstd")
563-
.header("Accept", MediaType::accepts())
581+
.header("Accept", accept)
564582
.build()
565583
.map_err(|err| ErrorKind::from_reqwest(url.clone(), err))?;
566584
let parse_simple_response = |response: Response| {
@@ -585,17 +603,48 @@ impl RegistryClient {
585603
})?;
586604

587605
let unarchived = match media_type {
588-
MediaType::Json => {
606+
MediaType::PyxV1Msgpack => {
607+
let bytes = response
608+
.bytes()
609+
.await
610+
.map_err(|err| ErrorKind::from_reqwest(url.clone(), err))?;
611+
let data: PyxSimpleDetail = rmp_serde::from_slice(bytes.as_ref())
612+
.map_err(|err| Error::from_msgpack_err(err, url.clone()))?;
613+
614+
SimpleMetadata::from_pyx_files(
615+
data.files,
616+
data.core_metadata,
617+
package_name,
618+
&url,
619+
)
620+
}
621+
MediaType::PyxV1Json => {
622+
let bytes = response
623+
.bytes()
624+
.await
625+
.map_err(|err| ErrorKind::from_reqwest(url.clone(), err))?;
626+
let data: PyxSimpleDetail = serde_json::from_slice(bytes.as_ref())
627+
.map_err(|err| Error::from_json_err(err, url.clone()))?;
628+
629+
SimpleMetadata::from_pyx_files(
630+
data.files,
631+
data.core_metadata,
632+
package_name,
633+
&url,
634+
)
635+
}
636+
MediaType::PypiV1Json => {
589637
let bytes = response
590638
.bytes()
591639
.await
592640
.map_err(|err| ErrorKind::from_reqwest(url.clone(), err))?;
641+
593642
let data: PypiSimpleDetail = serde_json::from_slice(bytes.as_ref())
594643
.map_err(|err| Error::from_json_err(err, url.clone()))?;
595644

596645
SimpleMetadata::from_pypi_files(data.files, package_name, &url)
597646
}
598-
MediaType::Html => {
647+
MediaType::PypiV1Html | MediaType::TextHtml => {
599648
let text = response
600649
.text()
601650
.await
@@ -1089,6 +1138,7 @@ pub struct SimpleMetadata(Vec<SimpleMetadatum>);
10891138
pub struct SimpleMetadatum {
10901139
pub version: Version,
10911140
pub files: VersionFiles,
1141+
pub metadata: Option<ResolutionMetadata>,
10921142
}
10931143

10941144
impl SimpleMetadata {
@@ -1101,7 +1151,7 @@ impl SimpleMetadata {
11011151
package_name: &PackageName,
11021152
base: &Url,
11031153
) -> Self {
1104-
let mut map: BTreeMap<Version, VersionFiles> = BTreeMap::default();
1154+
let mut version_map: BTreeMap<Version, VersionFiles> = BTreeMap::default();
11051155

11061156
// Convert to a reference-counted string.
11071157
let base = SmallString::from(base.as_str());
@@ -1113,19 +1163,65 @@ impl SimpleMetadata {
11131163
warn!("Skipping file for {package_name}: {}", file.filename);
11141164
continue;
11151165
};
1116-
let version = match filename {
1117-
DistFilename::SourceDistFilename(ref inner) => &inner.version,
1118-
DistFilename::WheelFilename(ref inner) => &inner.version,
1166+
let file = match File::try_from_pypi(file, &base) {
1167+
Ok(file) => file,
1168+
Err(err) => {
1169+
// Ignore files with unparsable version specifiers.
1170+
warn!("Skipping file for {package_name}: {err}");
1171+
continue;
1172+
}
11191173
};
1120-
let file = match File::try_from(file, &base) {
1174+
match version_map.entry(filename.version().clone()) {
1175+
std::collections::btree_map::Entry::Occupied(mut entry) => {
1176+
entry.get_mut().push(filename, file);
1177+
}
1178+
std::collections::btree_map::Entry::Vacant(entry) => {
1179+
let mut files = VersionFiles::default();
1180+
files.push(filename, file);
1181+
entry.insert(files);
1182+
}
1183+
}
1184+
}
1185+
1186+
Self(
1187+
version_map
1188+
.into_iter()
1189+
.map(|(version, files)| SimpleMetadatum {
1190+
version,
1191+
files,
1192+
metadata: None,
1193+
})
1194+
.collect(),
1195+
)
1196+
}
1197+
1198+
fn from_pyx_files(
1199+
files: Vec<uv_pypi_types::PyxFile>,
1200+
mut core_metadata: FxHashMap<Version, uv_pypi_types::CoreMetadatum>,
1201+
package_name: &PackageName,
1202+
base: &Url,
1203+
) -> Self {
1204+
let mut version_map: BTreeMap<Version, VersionFiles> = BTreeMap::default();
1205+
1206+
// Convert to a reference-counted string.
1207+
let base = SmallString::from(base.as_str());
1208+
1209+
// Group the distributions by version and kind
1210+
for file in files {
1211+
let file = match File::try_from_pyx(file, &base) {
11211212
Ok(file) => file,
11221213
Err(err) => {
11231214
// Ignore files with unparsable version specifiers.
11241215
warn!("Skipping file for {package_name}: {err}");
11251216
continue;
11261217
}
11271218
};
1128-
match map.entry(version.clone()) {
1219+
let Some(filename) = DistFilename::try_from_filename(&file.filename, package_name)
1220+
else {
1221+
warn!("Skipping file for {package_name}: {}", file.filename);
1222+
continue;
1223+
};
1224+
match version_map.entry(filename.version().clone()) {
11291225
std::collections::btree_map::Entry::Occupied(mut entry) => {
11301226
entry.get_mut().push(filename, file);
11311227
}
@@ -1136,9 +1232,28 @@ impl SimpleMetadata {
11361232
}
11371233
}
11381234
}
1235+
11391236
Self(
1140-
map.into_iter()
1141-
.map(|(version, files)| SimpleMetadatum { version, files })
1237+
version_map
1238+
.into_iter()
1239+
.map(|(version, files)| {
1240+
let metadata =
1241+
core_metadata
1242+
.remove(&version)
1243+
.map(|metadata| ResolutionMetadata {
1244+
name: package_name.clone(),
1245+
version: version.clone(),
1246+
requires_dist: metadata.requires_dist,
1247+
requires_python: metadata.requires_python,
1248+
provides_extras: metadata.provides_extras,
1249+
dynamic: false,
1250+
});
1251+
SimpleMetadatum {
1252+
version,
1253+
files,
1254+
metadata,
1255+
}
1256+
})
11421257
.collect(),
11431258
)
11441259
}
@@ -1177,26 +1292,51 @@ impl ArchivedSimpleMetadata {
11771292

11781293
#[derive(Debug)]
11791294
enum MediaType {
1180-
Json,
1181-
Html,
1295+
PyxV1Msgpack,
1296+
PyxV1Json,
1297+
PypiV1Json,
1298+
PypiV1Html,
1299+
TextHtml,
11821300
}
11831301

11841302
impl MediaType {
11851303
/// Parse a media type from a string, returning `None` if the media type is not supported.
11861304
fn from_str(s: &str) -> Option<Self> {
11871305
match s {
1188-
"application/vnd.pypi.simple.v1+json" => Some(Self::Json),
1189-
"application/vnd.pypi.simple.v1+html" | "text/html" => Some(Self::Html),
1306+
"application/vnd.pyx.simple.v1+msgpack" => Some(Self::PyxV1Msgpack),
1307+
"application/vnd.pyx.simple.v1+json" => Some(Self::PyxV1Json),
1308+
"application/vnd.pypi.simple.v1+json" => Some(Self::PypiV1Json),
1309+
"application/vnd.pypi.simple.v1+html" => Some(Self::PypiV1Html),
1310+
"text/html" => Some(Self::TextHtml),
11901311
_ => None,
11911312
}
11921313
}
11931314

1194-
/// Return the `Accept` header value for all supported media types.
1315+
/// Return the `Accept` header value for all PyPI media types.
11951316
#[inline]
1196-
const fn accepts() -> &'static str {
1317+
const fn pypi() -> &'static str {
11971318
// See: https://peps.python.org/pep-0691/#version-format-selection
11981319
"application/vnd.pypi.simple.v1+json, application/vnd.pypi.simple.v1+html;q=0.2, text/html;q=0.01"
11991320
}
1321+
1322+
/// Return the `Accept` header value for all supported media types.
1323+
#[inline]
1324+
const fn all() -> &'static str {
1325+
// See: https://peps.python.org/pep-0691/#version-format-selection
1326+
"application/vnd.pyx.simple.v1+msgpack, application/vnd.pyx.simple.v1+json;q=0.9, application/vnd.pypi.simple.v1+json;q=0.8, application/vnd.pypi.simple.v1+html;q=0.2, text/html;q=0.01"
1327+
}
1328+
}
1329+
1330+
impl std::fmt::Display for MediaType {
1331+
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1332+
match self {
1333+
Self::PyxV1Msgpack => write!(f, "application/vnd.pyx.simple.v1+msgpack"),
1334+
Self::PyxV1Json => write!(f, "application/vnd.pyx.simple.v1+json"),
1335+
Self::PypiV1Json => write!(f, "application/vnd.pypi.simple.v1+json"),
1336+
Self::PypiV1Html => write!(f, "application/vnd.pypi.simple.v1+html"),
1337+
Self::TextHtml => write!(f, "text/html"),
1338+
}
1339+
}
12001340
}
12011341

12021342
#[derive(Debug, Copy, Clone, Eq, PartialEq, Default)]

0 commit comments

Comments
 (0)