nblm_core/client/
mod.rs

1use std::{sync::Arc, time::Duration};
2
3use reqwest::{Client, Url};
4
5use crate::auth::TokenProvider;
6use crate::error::Result;
7
8mod api;
9mod http;
10mod retry;
11mod url_builder;
12
13pub use self::retry::{RetryConfig, Retryer};
14
15use self::http::HttpClient;
16use self::url_builder::{normalize_endpoint_location, UrlBuilder};
17
18const DEFAULT_TIMEOUT: Duration = Duration::from_secs(30);
19
20pub struct NblmClient {
21    pub(self) http: HttpClient,
22    pub(self) url_builder: UrlBuilder,
23    timeout: Duration,
24}
25
26impl NblmClient {
27    pub fn new(
28        token_provider: Arc<dyn TokenProvider>,
29        project_number: impl Into<String>,
30        location: impl Into<String>,
31        endpoint_location: impl Into<String>,
32    ) -> Result<Self> {
33        let project_number = project_number.into();
34        let location = location.into();
35        let endpoint_location = endpoint_location.into();
36        let base = format!(
37            "https://{}discoveryengine.googleapis.com/v1alpha",
38            normalize_endpoint_location(endpoint_location)?
39        );
40        let parent = format!("projects/{}/locations/{}", project_number, location);
41
42        let client = Client::builder()
43            .user_agent(concat!("nblm-cli/", env!("CARGO_PKG_VERSION")))
44            .timeout(DEFAULT_TIMEOUT)
45            .build()
46            .map_err(crate::error::Error::from)?;
47
48        let retryer = Retryer::new(RetryConfig::default());
49        let http = HttpClient::new(client, token_provider, retryer, None);
50        let url_builder = UrlBuilder::new(base, parent);
51
52        Ok(Self {
53            http,
54            url_builder,
55            timeout: DEFAULT_TIMEOUT,
56        })
57    }
58
59    pub fn with_timeout(mut self, timeout: Duration) -> Self {
60        self.timeout = timeout;
61        // Update the underlying HTTP client's timeout
62        let client = Client::builder()
63            .user_agent(concat!("nblm-cli/", env!("CARGO_PKG_VERSION")))
64            .timeout(timeout)
65            .build()
66            .expect("Failed to rebuild client with new timeout");
67
68        let token_provider = Arc::clone(&self.http.token_provider);
69        let retryer = self.http.retryer.clone();
70        let user_project = self.http.user_project.clone();
71        self.http = HttpClient::new(client, token_provider, retryer, user_project);
72        self
73    }
74
75    pub fn with_retry_config(mut self, config: RetryConfig) -> Self {
76        let client = Client::builder()
77            .user_agent(concat!("nblm-cli/", env!("CARGO_PKG_VERSION")))
78            .timeout(self.timeout)
79            .build()
80            .expect("Failed to rebuild client");
81
82        let token_provider = Arc::clone(&self.http.token_provider);
83        let retryer = Retryer::new(config);
84        let user_project = self.http.user_project.clone();
85        self.http = HttpClient::new(client, token_provider, retryer, user_project);
86        self
87    }
88
89    pub fn with_user_project(mut self, project: impl Into<String>) -> Self {
90        let client = Client::builder()
91            .user_agent(concat!("nblm-cli/", env!("CARGO_PKG_VERSION")))
92            .timeout(self.timeout)
93            .build()
94            .expect("Failed to rebuild client");
95
96        let token_provider = Arc::clone(&self.http.token_provider);
97        let retryer = self.http.retryer.clone();
98        let user_project = Some(project.into());
99        self.http = HttpClient::new(client, token_provider, retryer, user_project);
100        self
101    }
102
103    /// Override API base URL (for tests). Accepts absolute URL. Trims trailing slash.
104    pub fn with_base_url(mut self, base: impl Into<String>) -> Result<Self> {
105        let base = base.into().trim().trim_end_matches('/').to_string();
106        // Basic sanity check: absolute URL
107        let _ = Url::parse(&base).map_err(crate::error::Error::from)?;
108        let parent = self.url_builder.parent.clone();
109        self.url_builder = UrlBuilder::new(base, parent);
110        Ok(self)
111    }
112}
113
114#[cfg(test)]
115mod tests {
116    use super::*;
117
118    #[test]
119    fn with_base_url_accepts_absolute_url() {
120        let provider = Arc::new(crate::auth::StaticTokenProvider::new("test"));
121        let client = NblmClient::new(provider, "123", "global", "us").unwrap();
122        let result = client.with_base_url("http://localhost:8080/v1alpha");
123        assert!(result.is_ok());
124    }
125
126    #[test]
127    fn with_base_url_trims_trailing_slash() {
128        let provider = Arc::new(crate::auth::StaticTokenProvider::new("test"));
129        let client = NblmClient::new(provider, "123", "global", "us")
130            .unwrap()
131            .with_base_url("http://example.com/v1alpha/")
132            .unwrap();
133
134        // Test that URL building works correctly
135        let url = client.url_builder.build_url("/test").unwrap();
136        assert_eq!(url.as_str(), "http://example.com/v1alpha/test");
137    }
138
139    #[test]
140    fn with_base_url_rejects_relative_path() {
141        let provider = Arc::new(crate::auth::StaticTokenProvider::new("test"));
142        let client = NblmClient::new(provider, "123", "global", "us").unwrap();
143        let result = client.with_base_url("/relative/path");
144        assert!(result.is_err());
145    }
146}
nblm_core/client/mod.rs

nblm_core/client/
mod.rs
