Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

De-duplicate result handling in Simple API responses #10449

Merged
merged 1 commit into from
Jan 9, 2025
Merged

De-duplicate result handling in Simple API responses #10449

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
122 changes: 54 additions & 68 deletions crates/uv-client/src/registry_client.rs
View file
Open in desktop
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
Loading