Skip to content

Commit

Permalink
De-duplicate result handling in Simple API responses (#10449)
Browse files Browse the repository at this point in the history
## Summary

See: #10432 (comment)
  • Loading branch information
charliermarsh authored Jan 9, 2025
1 parent 22222e9 commit c5e536f
Showing 1 changed file with 54 additions and 68 deletions.
122 changes: 54 additions & 68 deletions crates/uv-client/src/registry_client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -253,80 +253,36 @@ impl RegistryClient {
// If we're searching for the first index that contains the package, fetch serially.
IndexStrategy::FirstIndex => {
for index in it {
match self.simple_single_index(package_name, index).await {
Ok(metadata) => {
results.push((index, metadata));
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()),
},
};
if let Some(metadata) = self
.simple_single_index(package_name, index, capabilities)
.await?
{
results.push((index, metadata));
break;
}
}
}

// Otherwise, fetch concurrently.
IndexStrategy::UnsafeBestMatch | IndexStrategy::UnsafeFirstMatch => {
let fetches = futures::stream::iter(it)
// TODO(charlie): Respect concurrency limits.
results = futures::stream::iter(it)
.map(|index| async move {
match self.simple_single_index(package_name, index).await {
Ok(metadata) => Ok(Some((index, 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()),
},
}
let metadata = self
.simple_single_index(package_name, index, capabilities)
.await?;
Ok((index, metadata))
})
.buffered(8);

futures::pin_mut!(fetches);

while let Some(result) = fetches.next().await {
match result {
Ok(Some((index, metadata))) => {
results.push((index, metadata));
.buffered(8)
.filter_map(|result: Result<_, Error>| async move {
match result {
Ok((index, Some(metadata))) => Some(Ok((index, metadata))),
Ok((_, None)) => None,
Err(err) => Some(Err(err)),
}
Ok(None) => continue,
Err(err) => return Err(err),
}
}
})
.try_collect::<Vec<_>>()
.await?;
}
}

Expand All @@ -346,11 +302,14 @@ impl RegistryClient {
///
/// The index can either be a PEP 503-compatible remote repository, or a local directory laid
/// out in the same format.
///
/// Returns `Ok(None)` if the package is not found in the index.
async fn simple_single_index(
&self,
package_name: &PackageName,
index: &IndexUrl,
) -> Result<OwnedArchive<SimpleMetadata>, Error> {
capabilities: &IndexCapabilities,
) -> Result<Option<OwnedArchive<SimpleMetadata>>, Error> {
// Format the URL for PyPI.
let mut url: Url = index.clone().into();
url.path_segments_mut()
Expand All @@ -377,11 +336,38 @@ impl RegistryClient {
Connectivity::Offline => CacheControl::AllowStale,
};

if matches!(index, IndexUrl::Path(_)) {
let result = if matches!(index, IndexUrl::Path(_)) {
self.fetch_local_index(package_name, &url).await
} else {
self.fetch_remote_index(package_name, &url, &cache_entry, cache_control)
.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()),
},
}
}

Expand Down

0 comments on commit c5e536f

Please sign in to comment.