Route CLI requests across regional deployments (0.7.4)

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "rawtree-cli"
-version = "0.7.3"
+version = "0.7.4"
 edition = "2021"
 repository = "https://github.com/rawtreedb/rawtree-cli"
 

src/client.rs

@@ -1,3 +1,4 @@
+mod discovery;
 mod regions;
 mod routing;
 #[cfg(test)]
@@ -39,11 +40,18 @@ impl ApiClient {
     }
 
     fn unrouted_request(&self, method: reqwest::Method, path: &str) -> RequestBuilder {
-        let request = with_client_header(
-            self.client
-                .request(method, format!("{}{}", self.base_url, path)),
-        );
-        let request = match &self.deployment {
+        self.deployment_request(method, path, self.deployment.as_ref())
+    }
+
+    fn deployment_request(
+        &self,
+        method: reqwest::Method,
+        path: &str,
+        deployment: Option<&(String, String)>,
+    ) -> RequestBuilder {
+        let url = routing::deployment_url(&format!("{}{}", self.base_url, path), deployment);
+        let request = with_client_header(self.client.request(method, url));
+        let request = match deployment {
             Some((deployment, region)) => request
                 .header("X-RawTree-Deployment-Id", deployment)
                 .header("X-RawTree-Region", region),
@@ -135,9 +143,17 @@ impl ApiClient {
     }
 
     pub fn get<T: DeserializeOwned>(&self, path: &str) -> Result<T> {
-        let req = self.routed_request(reqwest::Method::GET, path)?;
-        let resp = req.send().context("failed to connect to server")?;
-        let value: Value = handle_response(resp)?;
+        let url = reqwest::Url::parse(&format!("{}{}", self.base_url, path))?;
+        let organization = url.query_pairs().find(|(key, _)| key == "organization");
+        let value: Value = if url.path() == "/v1/clusters"
+            && self.has_user_credential()
+            && self.deployment.is_none()
+        {
+            self.discover_clusters(organization.as_ref().map(|(_, name)| name.as_ref()), None)?
+        } else {
+            let req = self.routed_request(reqwest::Method::GET, path)?;
+            handle_response(req.send().context("failed to connect to server")?)?
+        };
         self.remember_routing_metadata(path, &value);
         serde_json::from_value(value).context("failed to parse server response")
     }

src/client/discovery.rs

@@ -0,0 +1,163 @@
+use super::{
+    handle_response,
+    regions::Regions,
+    routing::{before_deadline, valid_routing},
+    ApiClient,
+};
+use anyhow::{Context, Result};
+use reqwest::Method;
+use serde::Deserialize;
+use serde_json::{json, Value};
+use std::{collections::HashMap, time::Instant};
+
+// Metadata identifies placement; only the owning catalog supplies authoritative fields.
+#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
+enum CatalogAuthority {
+    Fallback,
+    RegionMetadata,
+    OwningDeployment,
+}
+
+#[derive(Deserialize)]
+struct Organization {
+    organization_id: String,
+    name: String,
+}
+
+#[derive(Deserialize)]
+struct Organizations {
+    organizations: Vec<Organization>,
+}
+
+impl ApiClient {
+    pub(super) fn has_user_credential(&self) -> bool {
+        self.token
+            .as_deref()
+            .is_some_and(|token| !token.starts_with("rt_"))
+    }
+
+    /// Discover each enabled deployment before publishing or caching a complete list.
+    pub(super) fn discover_clusters(
+        &self,
+        organization: Option<&str>,
+        deadline: Option<Instant>,
+    ) -> Result<Value> {
+        let organizations: Organizations = handle_response(
+            before_deadline(
+                self.unrouted_request(Method::GET, "/v1/organizations"),
+                deadline,
+            )?
+            .send()?,
+        )?;
+        let selected: Vec<_> = organizations
+            .organizations
+            .into_iter()
+            .filter(|item| organization.is_none_or(|name| item.name == name))
+            .collect();
+        anyhow::ensure!(
+            organization.is_none() || !selected.is_empty(),
+            "Organization not found."
+        );
+        let mut result = json!({"clusters": []});
+        let mut clusters = Vec::new();
+        let mut positions = HashMap::new();
+        for owner in selected {
+            let regions: Regions = handle_response(
+                before_deadline(
+                    self.unrouted_request(
+                        Method::GET,
+                        &format!(
+                            "/v1/organizations/{}/regions",
+                            urlencoding::encode(&owner.organization_id)
+                        ),
+                    ),
+                    deadline,
+                )?
+                .send()?,
+            )?;
+            let mut deployments = Vec::new();
+            for region in regions.regions.iter().filter(|region| region.enabled) {
+                let pair = region
+                    .deployment_id
+                    .as_ref()
+                    .filter(|id| valid_routing(id, &region.cloud_region))
+                    .map(|id| (id.clone(), region.cloud_region.clone()));
+                if !deployments.contains(&pair) {
+                    deployments.push(pair);
+                }
+            }
+            let regions: HashMap<_, _> = regions
+                .regions
+                .iter()
+                .map(|region| (region.id.as_str(), region))
+                .collect();
+            for deployment in deployments {
+                let response: Value = handle_response(
+                    before_deadline(
+                        self.deployment_request(
+                            Method::GET,
+                            &format!(
+                                "/v1/clusters?organization={}",
+                                urlencoding::encode(&owner.name)
+                            ),
+                            deployment.as_ref(),
+                        ),
+                        deadline,
+                    )?
+                    .send()?,
+                )?;
+                let items = response
+                    .get("clusters")
+                    .and_then(Value::as_array)
+                    .context("Invalid cluster catalog response")?;
+                if organization.is_some() {
+                    result = response.clone();
+                }
+                for item in items {
+                    let mut cluster = item.clone();
+                    let id = item
+                        .get("id")
+                        .and_then(Value::as_str)
+                        .context("Cluster is missing its ID")?;
+                    let region_id = item
+                        .get("region_id")
+                        .and_then(Value::as_str)
+                        .filter(|id| !id.is_empty());
+                    let mut authority = CatalogAuthority::Fallback;
+                    if let Some(region_id) = region_id {
+                        let region = regions.get(region_id).with_context(|| {
+                            format!("Region '{region_id}' for cluster '{id}' was not found.")
+                        })?;
+                        let owns_cluster = match deployment.as_ref() {
+                            Some((id, cloud_region)) => {
+                                region.deployment_id.as_ref() == Some(id)
+                                    && region.cloud_region == *cloud_region
+                            }
+                            None => region.deployment_id.is_none(),
+                        };
+                        authority = if owns_cluster {
+                            CatalogAuthority::OwningDeployment
+                        } else {
+                            CatalogAuthority::RegionMetadata
+                        };
+                        cluster["deployment_id"] = json!(region.deployment_id);
+                        cluster["cloud_region"] = json!(region.cloud_region);
+                    }
+                    match positions.get_mut(id) {
+                        Some((index, previous)) if authority > *previous => {
+                            clusters[*index] = cluster;
+                            *previous = authority;
+                        }
+                        Some(_) => {}
+                        None => {
+                            positions.insert(id.to_owned(), (clusters.len(), authority));
+                            clusters.push(cluster);
+                        }
+                    }
+                }
+            }
+        }
+        result["clusters"] = json!(clusters);
+        Ok(result)
+    }
+}

src/client/routing.rs

@@ -126,12 +126,8 @@ impl ApiClient {
             self.routing.clusters.borrow_mut().clear();
             *self.routing.organizations.borrow_mut() = None;
         }
-        let mut request = self.unrouted_request(method, path);
-        if let Some((deployment, region)) = pair {
-            request = request
-                .header("X-RawTree-Deployment-Id", deployment)
-                .header("X-RawTree-Region", region);
-        }
+        let request =
+            self.deployment_request(method, path, pair.as_ref().or(self.deployment.as_ref()));
         before_deadline(request, deadline)
     }
 
@@ -166,10 +162,8 @@ impl ApiClient {
         let clusters = match cached {
             Some(clusters) => clusters,
             None => {
-                let response: Clusters = handle_response(
-                    before_deadline(self.unrouted_request(Method::GET, &collection), deadline)?
-                        .send()?,
-                )?;
+                let response: Clusters =
+                    serde_json::from_value(self.discover_clusters(Some(&organization), deadline)?)?;
                 self.routing
                     .clusters
                     .borrow_mut()
@@ -192,7 +186,10 @@ impl ApiClient {
     }
 }
 
-fn before_deadline(request: RequestBuilder, deadline: Option<Instant>) -> Result<RequestBuilder> {
+pub(super) fn before_deadline(
+    request: RequestBuilder,
+    deadline: Option<Instant>,
+) -> Result<RequestBuilder> {
     match deadline {
         Some(deadline) => {
             let remaining = deadline
@@ -251,7 +248,7 @@ pub(super) fn valid_routing(deployment: &str, region: &str) -> bool {
     let mut parts = region.rsplit('-');
     let number = parts.next().unwrap_or_default();
     let names: Vec<_> = parts.collect();
-    deployment.len() == 8
+    matches!(deployment.len(), 8 | 32)
         && deployment
             .bytes()
             .all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte))
@@ -264,3 +261,30 @@ pub(super) fn valid_routing(deployment: &str, region: &str) -> bool {
             .iter()
             .all(|part| !part.is_empty() && part.bytes().all(|byte| byte.is_ascii_lowercase()))
 }
+
+/// Only public RawTree API origins use deployment hosts; custom endpoints keep their URL.
+pub(super) fn deployment_url(url: &str, pair: Option<&(String, String)>) -> String {
+    let Some((deployment, region)) =
+        pair.filter(|(deployment, region)| valid_routing(deployment, region))
+    else {
+        return url.to_owned();
+    };
+    let Ok(mut target) = Url::parse(url) else {
+        return url.to_owned();
+    };
+    if !matches!(
+        target.origin().ascii_serialization().as_str(),
+        "https://api.rawtree.com" | "https://api.rawtree.dev"
+    ) {
+        return url.to_owned();
+    }
+    let host = target.host_str().unwrap_or_default().replacen(
+        "api.",
+        &format!("api.{deployment}.{region}.aws."),
+        1,
+    );
+    if target.set_host(Some(&host)).is_err() {
+        return url.to_owned();
+    }
+    target.to_string()
+}