From e1e85ab4c8d263d027c11216a2c6f7f1be893d57 Mon Sep 17 00:00:00 2001
From: Frost Ming <mianghong@gmail.com>
Date: Fri, 13 Sep 2024 03:27:54 +0800
Subject: [PATCH] feat(cli): add `--token` option to `self update` command
 (#7279)

<!--
Thank you for contributing to uv! To help us out with reviewing, please
consider the following:

- Does this pull request include a summary of the change? (See below.)
- Does this pull request include a descriptive title?
- Does this pull request include references to any relevant issues?
-->

## Summary

It often reaches the GitHub API rate limit and shows error like `error:
HTTP status client error (403 Forbidden) for url
(https://api.github.com/repos/astral-sh/uv/releases)` when running `uv
self update`.

To bypass this rate limit issue, allow user to pass a GitHub token via
`--token` or `UV_GITHUB_TOKEN` env.

## Test Plan

<!-- How was it tested? -->

---------

Signed-off-by: Frost Ming <me@frostming.com>
---
 Cargo.lock                            |  1 +
 crates/uv-cli/src/lib.rs              |  5 +++++
 crates/uv/Cargo.toml                  |  1 +
 crates/uv/src/commands/self_update.rs | 32 ++++++++++++++++++++++-----
 crates/uv/src/lib.rs                  |  8 +++++--
 5 files changed, 40 insertions(+), 7 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index b57616fd7..26185d0ca 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -4464,6 +4464,7 @@ dependencies = [
  "flate2",
  "fs-err",
  "futures",
+ "http",
  "ignore",
  "indexmap",
  "indicatif",
diff --git a/crates/uv-cli/src/lib.rs b/crates/uv-cli/src/lib.rs
index fb7c3a611..46896ff1c 100644
--- a/crates/uv-cli/src/lib.rs
+++ b/crates/uv-cli/src/lib.rs
@@ -427,6 +427,11 @@ pub enum SelfCommand {
 pub struct SelfUpdateArgs {
     /// Update to the specified version. If not provided, uv will update to the latest version.
     pub target_version: Option<String>,
+
+    /// A GitHub token for authentication.
+    /// A token is not required but can be used to reduce the chance of encountering rate limits.
+    #[arg(long, env = "UV_GITHUB_TOKEN")]
+    pub token: Option<String>,
 }
 
 #[derive(Args)]
diff --git a/crates/uv/Cargo.toml b/crates/uv/Cargo.toml
index 09f4c737c..9a710e16c 100644
--- a/crates/uv/Cargo.toml
+++ b/crates/uv/Cargo.toml
@@ -55,6 +55,7 @@ ctrlc = { workspace = true }
 flate2 = { workspace = true, default-features = false }
 fs-err = { workspace = true, features = ["tokio"] }
 futures = { workspace = true }
+http = { workspace = true }
 indexmap = { workspace = true }
 indicatif = { workspace = true }
 indoc = { workspace = true }
diff --git a/crates/uv/src/commands/self_update.rs b/crates/uv/src/commands/self_update.rs
index edf6f4c43..f6372d72c 100644
--- a/crates/uv/src/commands/self_update.rs
+++ b/crates/uv/src/commands/self_update.rs
@@ -11,10 +11,18 @@ use crate::commands::ExitStatus;
 use crate::printer::Printer;
 
 /// Attempt to update the uv binary.
-pub(crate) async fn self_update(version: Option<String>, printer: Printer) -> Result<ExitStatus> {
+pub(crate) async fn self_update(
+    version: Option<String>,
+    token: Option<String>,
+    printer: Printer,
+) -> Result<ExitStatus> {
     let mut updater = AxoUpdater::new_for("uv");
     updater.disable_installer_output();
 
+    if let Some(ref token) = token {
+        updater.set_github_token(token);
+    }
+
     // Load the "install receipt" for the current binary. If the receipt is not found, then
     // uv was likely installed via a package manager.
     let Ok(updater) = updater.load_receipt() else {
@@ -121,11 +129,25 @@ pub(crate) async fn self_update(version: Option<String>, printer: Printer) -> Re
             )?;
         }
         Err(err) => {
-            return Err(if let AxoupdateError::Reqwest(err) = err {
-                WrappedReqwestError::from(err).into()
+            return if let AxoupdateError::Reqwest(err) = err {
+                if err.status() == Some(http::StatusCode::FORBIDDEN) && token.is_none() {
+                    writeln!(
+                        printer.stderr(),
+                        "{}",
+                        format_args!(
+                            "{}{} GitHub API rate limit exceeded. Please provide a GitHub token via the {} option.",
+                            "error".red().bold(),
+                            ":".bold(),
+                            "`--token`".green().bold()
+                        )
+                    )?;
+                    Ok(ExitStatus::Error)
+                } else {
+                    Err(WrappedReqwestError::from(err).into())
+                }
             } else {
-                err.into()
-            });
+                Err(err.into())
+            };
         }
     }
 
diff --git a/crates/uv/src/lib.rs b/crates/uv/src/lib.rs
index 9360320a3..0ee3a7c27 100644
--- a/crates/uv/src/lib.rs
+++ b/crates/uv/src/lib.rs
@@ -766,8 +766,12 @@ async fn run(cli: Cli) -> Result<ExitStatus> {
         }
         #[cfg(feature = "self-update")]
         Commands::Self_(SelfNamespace {
-            command: SelfCommand::Update(SelfUpdateArgs { target_version }),
-        }) => commands::self_update(target_version, printer).await,
+            command:
+                SelfCommand::Update(SelfUpdateArgs {
+                    target_version,
+                    token,
+                }),
+        }) => commands::self_update(target_version, token, printer).await,
         Commands::Version { output_format } => {
             commands::version(output_format, &mut stdout())?;
             Ok(ExitStatus::Success)