De-duplicate result handling in Simple API responses (#10449)

## Summary

See: https://github.com/astral-sh/uv/pull/10432#issuecomment-2581084234
This commit is contained in:
Charlie Marsh 2025-01-09 17:45:13 -05:00 committed by GitHub
parent 22222e945f
commit c5e536f0ec
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 54 additions and 68 deletions

View File

@ -253,80 +253,36 @@ impl RegistryClient {
// If we're searching for the first index that contains the package, fetch serially. // If we're searching for the first index that contains the package, fetch serially.
IndexStrategy::FirstIndex => { IndexStrategy::FirstIndex => {
for index in it { for index in it {
match self.simple_single_index(package_name, index).await { if let Some(metadata) = self
Ok(metadata) => { .simple_single_index(package_name, index, capabilities)
.await?
{
results.push((index, metadata)); results.push((index, metadata));
break; break;
} }
Err(err) => match err.into_kind() {
// The package could not be found in the remote index.
ErrorKind::WrappedReqwestError(url, err) => match err.status() {
Some(StatusCode::NOT_FOUND) => {}
Some(StatusCode::UNAUTHORIZED) => {
capabilities.set_unauthorized(index.clone());
}
Some(StatusCode::FORBIDDEN) => {
capabilities.set_forbidden(index.clone());
}
_ => return Err(ErrorKind::WrappedReqwestError(url, err).into()),
},
// The package is unavailable due to a lack of connectivity.
ErrorKind::Offline(_) => {}
// The package could not be found in the local index.
ErrorKind::FileNotFound(_) => {}
err => return Err(err.into()),
},
};
} }
} }
// Otherwise, fetch concurrently. // Otherwise, fetch concurrently.
IndexStrategy::UnsafeBestMatch | IndexStrategy::UnsafeFirstMatch => { IndexStrategy::UnsafeBestMatch | IndexStrategy::UnsafeFirstMatch => {
let fetches = futures::stream::iter(it) // TODO(charlie): Respect concurrency limits.
results = futures::stream::iter(it)
.map(|index| async move { .map(|index| async move {
match self.simple_single_index(package_name, index).await { let metadata = self
Ok(metadata) => Ok(Some((index, metadata))), .simple_single_index(package_name, index, capabilities)
Err(err) => match err.into_kind() { .await?;
// The package could not be found in the remote index. Ok((index, metadata))
ErrorKind::WrappedReqwestError(url, err) => match err.status() { })
Some(StatusCode::NOT_FOUND) => Ok(None), .buffered(8)
Some(StatusCode::UNAUTHORIZED) => { .filter_map(|result: Result<_, Error>| async move {
capabilities.set_unauthorized(index.clone()); match result {
Ok(None) Ok((index, Some(metadata))) => Some(Ok((index, metadata))),
} Ok((_, None)) => None,
Some(StatusCode::FORBIDDEN) => { Err(err) => Some(Err(err)),
capabilities.set_forbidden(index.clone());
Ok(None)
}
_ => Err(ErrorKind::WrappedReqwestError(url, err).into()),
},
// The package is unavailable due to a lack of connectivity.
ErrorKind::Offline(_) => Ok(None),
// The package could not be found in the local index.
ErrorKind::FileNotFound(_) => Ok(None),
err => Err(err.into()),
},
} }
}) })
.buffered(8); .try_collect::<Vec<_>>()
.await?;
futures::pin_mut!(fetches);
while let Some(result) = fetches.next().await {
match result {
Ok(Some((index, metadata))) => {
results.push((index, metadata));
}
Ok(None) => continue,
Err(err) => return Err(err),
}
}
} }
} }
@ -346,11 +302,14 @@ impl RegistryClient {
/// ///
/// The index can either be a PEP 503-compatible remote repository, or a local directory laid /// The index can either be a PEP 503-compatible remote repository, or a local directory laid
/// out in the same format. /// out in the same format.
///
/// Returns `Ok(None)` if the package is not found in the index.
async fn simple_single_index( async fn simple_single_index(
&self, &self,
package_name: &PackageName, package_name: &PackageName,
index: &IndexUrl, index: &IndexUrl,
) -> Result<OwnedArchive<SimpleMetadata>, Error> { capabilities: &IndexCapabilities,
) -> Result<Option<OwnedArchive<SimpleMetadata>>, Error> {
// Format the URL for PyPI. // Format the URL for PyPI.
let mut url: Url = index.clone().into(); let mut url: Url = index.clone().into();
url.path_segments_mut() url.path_segments_mut()
@ -377,11 +336,38 @@ impl RegistryClient {
Connectivity::Offline => CacheControl::AllowStale, Connectivity::Offline => CacheControl::AllowStale,
}; };
if matches!(index, IndexUrl::Path(_)) { let result = if matches!(index, IndexUrl::Path(_)) {
self.fetch_local_index(package_name, &url).await self.fetch_local_index(package_name, &url).await
} else { } else {
self.fetch_remote_index(package_name, &url, &cache_entry, cache_control) self.fetch_remote_index(package_name, &url, &cache_entry, cache_control)
.await .await
};
match result {
Ok(metadata) => Ok(Some(metadata)),
Err(err) => match err.into_kind() {
// The package could not be found in the remote index.
ErrorKind::WrappedReqwestError(url, err) => match err.status() {
Some(StatusCode::NOT_FOUND) => Ok(None),
Some(StatusCode::UNAUTHORIZED) => {
capabilities.set_unauthorized(index.clone());
Ok(None)
}
Some(StatusCode::FORBIDDEN) => {
capabilities.set_forbidden(index.clone());
Ok(None)
}
_ => Err(ErrorKind::WrappedReqwestError(url, err).into()),
},
// The package is unavailable due to a lack of connectivity.
ErrorKind::Offline(_) => Ok(None),
// The package could not be found in the local index.
ErrorKind::FileNotFound(_) => Ok(None),
err => Err(err.into()),
},
} }
} }