fix(hybrid-cloud): Adds region fan-out to organizations list command (#1860)

GabeVillalobos · web-flow · commit 5898da86969f · 2023-12-12T09:28:50.000-08:00
diff --git a/src/api.rs b/src/api.rs
@@ -252,6 +252,8 @@ pub enum ApiErrorKind {
     RequestFailed,
     #[error("could not compress data")]
     CompressionFailed,
+    #[error("region overrides cannot be applied to absolute urls")]
+    InvalidRegionRequest,
 }
 
 #[derive(Debug, thiserror::Error)]
@@ -391,28 +393,65 @@ impl Api {
     /// URL is just a path then it's relative to the configured API host
     /// and authentication is automatically enabled.
     pub fn request(&self, method: Method, url: &str) -> ApiResult<ApiRequest> {
-        let mut handle = self.pool.get().unwrap();
-        handle.reset();
-        if !self.config.allow_keepalive() {
-            handle.forbid_reuse(true).ok();
-        }
-        let mut ssl_opts = curl::easy::SslOpt::new();
-        if self.config.disable_ssl_revocation_check() {
-            ssl_opts.no_revoke(true);
+        let (url, auth) = self.resolve_base_url_and_auth(url, None)?;
+        self.construct_api_request(method, &url, auth)
+    }
+
+    /// Like `request`, but constructs a new `ApiRequest` using the base URL
+    /// plus a region prefix for requests that must be routed to a region.
+    pub fn region_request(
+        &self,
+        method: Method,
+        url: &str,
+        region: &Region,
+    ) -> ApiResult<ApiRequest> {
+        let (url, auth) = self.resolve_base_url_and_auth(url, Some(region))?;
+        self.construct_api_request(method, &url, auth)
+    }
+
+    fn resolve_base_url_and_auth(
+        &self,
+        url: &str,
+        region: Option<&Region>,
+    ) -> ApiResult<(String, Option<&Auth>)> {
+        if is_absolute_url(url) && region.is_some() {
+            return Err(ApiErrorKind::InvalidRegionRequest.into());
         }
-        handle.ssl_options(&ssl_opts)?;
+
         let (url, auth) = if is_absolute_url(url) {
             (Cow::Borrowed(url), None)
         } else {
+            let host_override = region.map(|rg| rg.url.as_str());
+
             (
-                Cow::Owned(match self.config.get_api_endpoint(url) {
+                Cow::Owned(match self.config.get_api_endpoint(url, host_override) {
                     Ok(rv) => rv,
                     Err(err) => return Err(ApiError::with_source(ApiErrorKind::BadApiUrl, err)),
                 }),
                 self.config.get_auth(),
             )
         };
 
+        Ok((url.into_owned(), auth))
+    }
+
+    fn construct_api_request(
+        &self,
+        method: Method,
+        url: &str,
+        auth: Option<&Auth>,
+    ) -> ApiResult<ApiRequest> {
+        let mut handle = self.pool.get().unwrap();
+        handle.reset();
+        if !self.config.allow_keepalive() {
+            handle.forbid_reuse(true).ok();
+        }
+        let mut ssl_opts = curl::easy::SslOpt::new();
+        if self.config.disable_ssl_revocation_check() {
+            ssl_opts.no_revoke(true);
+        }
+        handle.ssl_options(&ssl_opts)?;
+
         if let Some(proxy_url) = self.config.get_proxy_url() {
             handle.proxy(&proxy_url)?;
         }
@@ -431,7 +470,7 @@ impl Api {
         let env = self.config.get_pipeline_env();
         let headers = self.config.get_headers();
 
-        ApiRequest::create(handle, &method, &url, auth, env, headers)
+        ApiRequest::create(handle, &method, url, auth, env, headers)
     }
 
     /// Convenience method that performs a request using DSN as authentication method.
@@ -445,7 +484,7 @@ impl Api {
         // We resolve an absolute URL to skip default authentication flow.
         let url = self
             .config
-            .get_api_endpoint(url)
+            .get_api_endpoint(url, None)
             .map_err(|err| ApiError::with_source(ApiErrorKind::BadApiUrl, err))?;
 
         let mut request = self
@@ -1443,11 +1482,19 @@ impl Api {
     }
 
     /// List all organizations associated with the authenticated token
-    pub fn list_organizations(&self) -> ApiResult<Vec<Organization>> {
+    /// in the given `Region`. If no `Region` is provided, we assume
+    /// we're issuing a request to a monolith deployment.
+    pub fn list_organizations(&self, region: Option<&Region>) -> ApiResult<Vec<Organization>> {
         let mut rv = vec![];
         let mut cursor = "".to_string();
         loop {
-            let resp = self.get(&format!("/organizations/?cursor={}", QueryArg(&cursor)))?;
+            let current_path = &format!("/organizations/?cursor={}", QueryArg(&cursor));
+            let resp = if let Some(rg) = region {
+                self.region_request(Method::Get, current_path, rg)?.send()?
+            } else {
+                self.get(current_path)?
+            };
+
             if resp.status() == 404 || (resp.status() == 400 && !cursor.is_empty()) {
                 if rv.is_empty() {
                     return Err(ApiErrorKind::ResourceNotFound.into());
@@ -1466,6 +1513,22 @@ impl Api {
         Ok(rv)
     }
 
+    pub fn list_available_regions(&self) -> ApiResult<Vec<Region>> {
+        let resp = self.get("/users/me/regions/")?;
+        if resp.status() == 404 {
+            // This endpoint may not exist for self-hosted users, so
+            // returning a default of [] seems appropriate.
+            return Ok(vec![]);
+        }
+
+        if resp.status() == 400 {
+            return Err(ApiErrorKind::ResourceNotFound.into());
+        }
+
+        let region_response = resp.convert::<RegionResponse>()?;
+        Ok(region_response.regions)
+    }
+
     /// List all monitors associated with an organization
     pub fn list_organization_monitors(&self, org: &str) -> ApiResult<Vec<Monitor>> {
         let mut rv = vec![];
@@ -2947,3 +3010,14 @@ impl fmt::Display for ProcessedEventTag {
         Ok(())
     }
 }
+
+#[derive(Clone, Debug, Deserialize)]
+pub struct Region {
+    pub name: String,
+    pub url: String,
+}
+
+#[derive(Clone, Debug, Deserialize)]
+pub struct RegionResponse {
+    pub regions: Vec<Region>,
+}
diff --git a/src/commands/organizations/list.rs b/src/commands/organizations/list.rs
@@ -1,7 +1,8 @@
 use anyhow::Result;
 use clap::{ArgMatches, Command};
+use log::debug;
 
-use crate::api::Api;
+use crate::api::{Api, Organization};
 use crate::utils::formatting::Table;
 
 pub fn make_command(command: Command) -> Command {
@@ -10,7 +11,22 @@ pub fn make_command(command: Command) -> Command {
 
 pub fn execute(_matches: &ArgMatches) -> Result<()> {
     let api = Api::current();
-    let mut organizations = api.list_organizations()?;
+
+    // Query regions available to the current CLI user
+    let regions = api.list_available_regions()?;
+
+    let mut organizations: Vec<Organization> = vec![];
+    debug!("Available regions: {:?}", regions);
+
+    // Self-hosted instances won't have a region instance or prefix, so we
+    // need to check before fanning out.
+    if regions.len() > 1 {
+        for region in regions {
+            organizations.append(&mut api.list_organizations(Some(&region))?)
+        }
+    } else {
+        organizations.append(&mut api.list_organizations(None)?)
+    }
 
     organizations.sort_by_key(|o| o.name.clone().to_lowercase());
 
diff --git a/src/config.rs b/src/config.rs
@@ -270,10 +270,13 @@ impl Config {
     }
 
     /// Returns the API URL for a path
-    pub fn get_api_endpoint(&self, path: &str) -> Result<String> {
-        let base = self.get_base_url()?;
+    pub fn get_api_endpoint(&self, path: &str, base_url_override: Option<&str>) -> Result<String> {
+        let base: &str = base_url_override
+            .unwrap_or(self.get_base_url()?)
+            .trim_end_matches('/');
         let path = path.trim_start_matches('/');
         let path = path.trim_start_matches("api/0/");
+
         Ok(format!("{}/api/0/{}", base, path))
     }
 
@@ -788,16 +791,26 @@ mod tests {
 
         assert_eq!(
             config
-                .get_api_endpoint("/organizations/test-org/chunk-upload/")
+                .get_api_endpoint("/organizations/test-org/chunk-upload/", None)
                 .unwrap(),
             "https://sentry.io/api/0/organizations/test-org/chunk-upload/"
         );
 
         assert_eq!(
             config
-                .get_api_endpoint("/api/0/organizations/test-org/chunk-upload/")
+                .get_api_endpoint("/api/0/organizations/test-org/chunk-upload/", None)
                 .unwrap(),
             "https://sentry.io/api/0/organizations/test-org/chunk-upload/"
         );
+
+        assert_eq!(
+            config
+                .get_api_endpoint(
+                    "/api/0/organizations/test-org/chunk-upload/",
+                    Some("https://us.sentry.io/")
+                )
+                .unwrap(),
+            "https://us.sentry.io/api/0/organizations/test-org/chunk-upload/"
+        );
     }
 }
diff --git a/tests/integration/organizations/list.rs b/tests/integration/organizations/list.rs
@@ -1,3 +1,5 @@
+use mockito::server_url;
+
 use crate::integration::{mock_endpoint, register_test, EndpointOptions};
 
 #[test]
@@ -6,6 +8,21 @@ fn command_organizations_list() {
         EndpointOptions::new("GET", "/api/0/organizations/?cursor=", 200)
             .with_response_file("organizations/get-organizations.json"),
     );
+
+    let region_response = format!(
+        r#"{{
+            "regions": [{{
+                "name": "monolith",
+                "url": "{}"
+            }}]
+        }}"#,
+        server_url(),
+    );
+
+    let _mock_regions = mock_endpoint(
+        EndpointOptions::new("GET", "/api/0/users/me/regions/", 200)
+            .with_response_body(region_response),
+    );
     register_test("organizations/organizations-list.trycmd");
 }
 
