Add method to retrieve all users from Firebase Auth

src/auth/credential/api_auth_token.rs

@@ -52,9 +52,7 @@ impl ApiAuthTokenManager {
             "grant_type=urn%3Aietf%3Aparams%3Aoauth%3Agrant-type%3Ajwt-bearer&assertion={jwt}"
         );
 
-        let url = format!(
-            "https://{GOOGLE_AUTH_TOKEN_HOST}{GOOGLE_AUTH_TOKEN_PATH}"
-        );
+        let url = format!("https://{GOOGLE_AUTH_TOKEN_HOST}{GOOGLE_AUTH_TOKEN_PATH}");
 
         let res = self
             .http_client

src/auth/mod.rs

@@ -1,11 +1,13 @@
+use std::collections::HashMap;
+
 use anyhow::Context;
 use reqwest::Response;
 use serde::{de::DeserializeOwned, Deserialize, Serialize};
 
 use crate::{
     auth::{
         error::AuthApiErrorResponse,
-        models::{UpdateUserBody, UpdateUserValues},
+        models::{BatchGetResponse, UpdateUserBody, UpdateUserValues},
     },
     error::FirebaseError,
     ServiceAccount,
@@ -26,6 +28,7 @@ pub struct FirebaseAuthClient {
     api_url: String,
     user_token_manager: UserTokenManager,
     api_auth_token_manager: ApiAuthTokenManager,
+    project_id: String,
 }
 
 impl FirebaseAuthClient {
@@ -36,35 +39,52 @@ impl FirebaseAuthClient {
             .context("Failed to create HTTP client")?;
 
         let credential_manager = ApiAuthTokenManager::new(service_account.clone());
+        let project_id = service_account.project_id.clone();
         let token_handler = UserTokenManager::new(service_account, client.clone());
 
         Ok(Self {
             user_token_manager: token_handler,
             client,
             api_url: "https://identitytoolkit.googleapis.com/v1".to_string(),
             api_auth_token_manager: credential_manager,
+            project_id,
         })
     }
 
     fn url(&self, path: impl AsRef<str>) -> String {
         format!("{}{}", self.api_url, path.as_ref())
     }
 
-    /// Creates a new `POST` request builder with the `Authorization` header set
-    /// to an authorized admin access token.
-    async fn auth_post(
-        &self,
-        url: impl AsRef<str>,
-    ) -> Result<reqwest::RequestBuilder, FirebaseError> {
+    fn project_url(&self, path: impl AsRef<str>) -> String {
+        format!(
+            "{}/projects/{}{}",
+            self.api_url,
+            &self.project_id,
+            path.as_ref()
+        )
+    }
+
+    async fn get_access_token(&self) -> Result<String, FirebaseError> {
         let access_token = self
             .api_auth_token_manager
             .get_access_token()
             .await
             .map_err(|e| {
-                tracing::error!("Failed to get access token: {}", e);
+                tracing::error!("Failed to get access token: {e}");
                 e
             })?;
 
+        Ok(access_token)
+    }
+
+    /// Creates a new `POST` request builder with the `Authorization` header set
+    /// to an authorized admin access token.
+    async fn auth_post(
+        &self,
+        url: impl AsRef<str>,
+    ) -> Result<reqwest::RequestBuilder, FirebaseError> {
+        let access_token = self.get_access_token().await?;
+
         let builder = self
             .client
             .post(url.as_ref())
@@ -73,6 +93,22 @@ impl FirebaseAuthClient {
         Ok(builder)
     }
 
+    /// Creates a new `GET` request builder with the `Authorization` header set
+    /// to an authorized admin access token.
+    async fn auth_get(
+        &self,
+        url: impl AsRef<str>,
+    ) -> Result<reqwest::RequestBuilder, FirebaseError> {
+        let access_token = self.get_access_token().await?;
+
+        let builder = self
+            .client
+            .get(url.as_ref())
+            .header("Authorization", format!("Bearer {access_token}"));
+
+        Ok(builder)
+    }
+
     /// Decodes an ID token and returns its claims. Only succeeds if the token
     /// is valid. The token is valid if it:
     ///
@@ -306,6 +342,125 @@ impl FirebaseAuthClient {
         Ok(user)
     }
 
+    /// Retrieve all users in Firebase Auth. By default, this will send multiple batch requests to
+    /// the Firebase Auth API sequentially until all users have been retrieved. If you have
+    /// thousands of users, this can take a while. See
+    /// [`get_all_users_concurrently`](Self::get_all_users_concurrently) if you want these requests
+    /// to be sent concurrently by `N` workers.
+    #[tracing::instrument(name = "Get all users", skip_all)]
+    pub async fn get_all_users(&self) -> Result<Vec<User>, FirebaseError> {
+        self.get_all_users_concurrently(1).await
+    }
+
+    /// Retrieve all users in Firebase Auth using `num_workers` concurrent workers. See
+    /// [`get_all_users`](Self::get_all_users) for more info.
+    ///
+    /// The number of workers should be between 1 and 62. If you pass a number outside
+    /// this range, it will be clamped.
+    ///
+    /// Keep in mind that, at the time of writing, the Identity Toolkit API limit is 500
+    /// requests/second per service account. This method will fire off multiple requests
+    /// simultaneously. See the up-to-date limits in the [official docs](https://firebase.google.com/docs/auth/limits).
+    #[tracing::instrument(name = "Get all users concurrently", skip_all)]
+    pub async fn get_all_users_concurrently(
+        &self,
+        num_workers: u8,
+    ) -> Result<Vec<User>, FirebaseError> {
+        let base_url = self.project_url("/accounts:batchGet");
+
+        fn make_pagination_url(
+            base_url: &str,
+            max_results: usize,
+            next_page_token: Option<&str>,
+        ) -> String {
+            format!(
+                "{base_url}?maxResults={max_results}{}",
+                next_page_token
+                    .map(|token| format!("&nextPageToken={token}"))
+                    .unwrap_or_else(|| { "".to_string() })
+            )
+        }
+
+        let legal = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"
+            .chars()
+            .collect::<Vec<_>>();
+
+        let num_workers = (num_workers as usize).clamp(1, legal.len());
+
+        let chunk_size = legal.len() / num_workers;
+        let mut ranges = Vec::new();
+
+        for i in 0..num_workers {
+            let min = legal[i * chunk_size];
+            let max = if i < num_workers - 1 {
+                legal[(i + 1) * chunk_size]
+            } else {
+                (*legal.last().unwrap() as u8 + 1) as char
+            };
+
+            ranges.push((min, max));
+        }
+
+        async fn run_worker(
+            client: &FirebaseAuthClient,
+            range: (char, char),
+            base_url: &str,
+        ) -> Result<Vec<User>, FirebaseError> {
+            const MAX_RESULTS: usize = 1000;
+            let (min, max) = (range.0.to_string(), range.1.to_string());
+
+            let mut all_users = Vec::new();
+            let mut next_page_token = Some(min);
+            loop {
+                let url = make_pagination_url(base_url, MAX_RESULTS, next_page_token.as_deref());
+
+                let res = client
+                    .auth_get(url)
+                    .await?
+                    .header("Content-Type", "application/json")
+                    .send()
+                    .await
+                    .context("Failed to send get all users request")?;
+
+                if !res.status().is_success() {
+                    return Err(response_error("Failed to get all users", res).await);
+                }
+
+                let res_body: BatchGetResponse =
+                    res.json().await.context("Failed to read response JSON")?;
+
+                if let Some(mut users) = res_body.users {
+                    all_users.append(&mut users);
+                }
+
+                next_page_token = res_body.next_page_token.map(|t| t.to_string());
+
+                if next_page_token.is_none() || next_page_token.as_deref() >= Some(&max) {
+                    break;
+                }
+            }
+
+            Ok(all_users)
+        }
+
+        let futures = ranges
+            .into_iter()
+            .map(|range| run_worker(self, range, &base_url));
+
+        let worker_results = futures::future::try_join_all(futures).await?;
+
+        let mut deduped = HashMap::new();
+        for users in worker_results {
+            for user in users {
+                deduped.insert(user.uid.clone(), user);
+            }
+        }
+
+        let all_users = deduped.into_values().collect();
+
+        Ok(all_users)
+    }
+
     /// Creates a new user in Firebase Auth using the email/password provider.
     ///
     /// # Examples

src/auth/models/mod.rs

@@ -11,6 +11,13 @@ pub(crate) struct GetAccountInfoResponse {
     pub users: Option<Vec<User>>,
 }
 
+#[derive(Debug, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub(crate) struct BatchGetResponse {
+    pub users: Option<Vec<User>>,
+    pub next_page_token: Option<String>,
+}
+
 #[derive(Debug, Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub struct User {