From cdb8497cebb72a81160c41ac6aa49f34e005033a Mon Sep 17 00:00:00 2001
From: Frost Ming <me@frostming.com>
Date: Wed, 11 Sep 2024 11:25:03 +0800
Subject: [PATCH] feat(cli): add `--token` option to `self update` command

Signed-off-by: Frost Ming <me@frostming.com>
---
 crates/uv-cli/src/lib.rs              |  4 ++++
 crates/uv/src/commands/self_update.rs | 10 +++++++++-
 crates/uv/src/lib.rs                  |  8 ++++++--
 3 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/crates/uv-cli/src/lib.rs b/crates/uv-cli/src/lib.rs
index daf02811e819..ed7833430cd6 100644
--- a/crates/uv-cli/src/lib.rs
+++ b/crates/uv-cli/src/lib.rs
@@ -427,6 +427,10 @@ 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>,
+
+    /// Specify a github token for authentication during the update process.
+    #[arg(long, env = "UV_GITHUB_TOKEN")]
+    pub token: Option<String>,
 }
 
 #[derive(Args)]
diff --git a/crates/uv/src/commands/self_update.rs b/crates/uv/src/commands/self_update.rs
index edf6f4c43b3a..9fb23f6bf9ba 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(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 {
diff --git a/crates/uv/src/lib.rs b/crates/uv/src/lib.rs
index 9360320a3acc..0ee3a7c2720f 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)