Skip to main content

hive_console_sdk/agent/
builder.rs

1use std::{sync::Arc, time::Duration};
2
3use async_dropper_simple::AsyncDropper;
4use recloser::AsyncRecloser;
5use reqwest::header::{HeaderMap, HeaderValue};
6use reqwest_middleware::ClientBuilder;
7use reqwest_retry::RetryTransientMiddleware;
8use std::sync::LazyLock;
9
10use crate::agent::buffer::Buffer;
11use crate::agent::usage_agent::{non_empty_string, AgentError, UsageAgent, UsageAgentInner};
12use crate::agent::utils::OperationProcessor;
13use crate::circuit_breaker;
14use retry_policies::policies::ExponentialBackoff;
15
16pub struct UsageAgentBuilder {
17    token: Option<String>,
18    endpoint: String,
19    target_id: Option<String>,
20    buffer_size: usize,
21    connect_timeout: Duration,
22    request_timeout: Duration,
23    accept_invalid_certs: bool,
24    flush_interval: Duration,
25    retry_policy: ExponentialBackoff,
26    user_agent: Option<String>,
27    circuit_breaker: Option<AsyncRecloser>,
28}
29
30pub static DEFAULT_HIVE_USAGE_ENDPOINT: &str = "https://app.graphql-hive.com/usage";
31
32impl Default for UsageAgentBuilder {
33    fn default() -> Self {
34        Self {
35            endpoint: DEFAULT_HIVE_USAGE_ENDPOINT.to_string(),
36            token: None,
37            target_id: None,
38            buffer_size: 1000,
39            connect_timeout: Duration::from_secs(5),
40            request_timeout: Duration::from_secs(15),
41            accept_invalid_certs: false,
42            flush_interval: Duration::from_secs(5),
43            retry_policy: ExponentialBackoff::builder().build_with_max_retries(3),
44            user_agent: None,
45            circuit_breaker: None,
46        }
47    }
48}
49
50fn is_legacy_token(token: &str) -> bool {
51    !token.starts_with("hvo1/") && !token.starts_with("hvu1/") && !token.starts_with("hvp1/")
52}
53
54impl UsageAgentBuilder {
55    /// Your [Registry Access Token](https://the-guild.dev/graphql/hive/docs/management/targets#registry-access-tokens) with write permission.
56    pub fn token(mut self, token: String) -> Self {
57        if let Some(token) = non_empty_string(Some(token)) {
58            self.token = Some(token);
59        }
60        self
61    }
62    /// For self-hosting, you can override `/usage` endpoint (defaults to `https://app.graphql-hive.com/usage`).
63    pub fn endpoint(mut self, endpoint: String) -> Self {
64        if let Some(endpoint) = non_empty_string(Some(endpoint)) {
65            self.endpoint = endpoint;
66        }
67        self
68    }
69    /// A target ID, this can either be a slug following the format “$organizationSlug/$projectSlug/$targetSlug” (e.g “the-guild/graphql-hive/staging”) or an UUID (e.g. “a0f4c605-6541-4350-8cfe-b31f21a4bf80”). To be used when the token is configured with an organization access token.
70    pub fn target_id(mut self, target_id: String) -> Self {
71        if let Some(target_id) = non_empty_string(Some(target_id)) {
72            self.target_id = Some(target_id);
73        }
74        self
75    }
76    /// A maximum number of operations to hold in a buffer before sending to Hive Console
77    /// Default: 1000
78    pub fn buffer_size(mut self, buffer_size: usize) -> Self {
79        self.buffer_size = buffer_size;
80        self
81    }
82    /// A timeout for only the connect phase of a request to Hive Console
83    /// Default: 5 seconds
84    pub fn connect_timeout(mut self, connect_timeout: Duration) -> Self {
85        self.connect_timeout = connect_timeout;
86        self
87    }
88    /// A timeout for the entire request to Hive Console
89    /// Default: 15 seconds
90    pub fn request_timeout(mut self, request_timeout: Duration) -> Self {
91        self.request_timeout = request_timeout;
92        self
93    }
94    /// Accepts invalid SSL certificates
95    /// Default: false
96    pub fn accept_invalid_certs(mut self, accept_invalid_certs: bool) -> Self {
97        self.accept_invalid_certs = accept_invalid_certs;
98        self
99    }
100    /// Frequency of flushing the buffer to the server
101    /// Default: 5 seconds
102    pub fn flush_interval(mut self, flush_interval: Duration) -> Self {
103        self.flush_interval = flush_interval;
104        self
105    }
106    /// User-Agent header to be sent with each request
107    pub fn user_agent(mut self, user_agent: String) -> Self {
108        if let Some(user_agent) = non_empty_string(Some(user_agent)) {
109            self.user_agent = Some(user_agent);
110        }
111        self
112    }
113    /// Retry policy for sending reports
114    /// Default: ExponentialBackoff with max 3 retries
115    pub fn retry_policy(mut self, retry_policy: ExponentialBackoff) -> Self {
116        self.retry_policy = retry_policy;
117        self
118    }
119    /// Maximum number of retries for sending reports
120    /// Default: ExponentialBackoff with max 3 retries
121    pub fn max_retries(mut self, max_retries: u32) -> Self {
122        self.retry_policy = ExponentialBackoff::builder().build_with_max_retries(max_retries);
123        self
124    }
125    pub(crate) fn build_agent(self) -> Result<UsageAgentInner, AgentError> {
126        let mut default_headers = HeaderMap::new();
127
128        default_headers.insert("X-Usage-API-Version", HeaderValue::from_static("2"));
129
130        let token = match self.token {
131            Some(token) => token,
132            None => return Err(AgentError::MissingToken),
133        };
134
135        let mut authorization_header = HeaderValue::from_str(&format!("Bearer {}", token))
136            .map_err(|_| AgentError::InvalidToken)?;
137
138        authorization_header.set_sensitive(true);
139
140        default_headers.insert(reqwest::header::AUTHORIZATION, authorization_header);
141
142        default_headers.insert(
143            reqwest::header::CONTENT_TYPE,
144            HeaderValue::from_static("application/json"),
145        );
146
147        let mut reqwest_agent = reqwest::Client::builder()
148            .danger_accept_invalid_certs(self.accept_invalid_certs)
149            .connect_timeout(self.connect_timeout)
150            .timeout(self.request_timeout)
151            .default_headers(default_headers);
152
153        if let Some(user_agent) = &self.user_agent {
154            reqwest_agent = reqwest_agent.user_agent(user_agent);
155        }
156
157        let reqwest_agent = reqwest_agent
158            .build()
159            .map_err(AgentError::HTTPClientCreationError)?;
160        let client = ClientBuilder::new(reqwest_agent)
161            .with(RetryTransientMiddleware::new_with_policy(self.retry_policy))
162            .build();
163
164        let mut endpoint = self.endpoint;
165
166        match self.target_id {
167            Some(_) if is_legacy_token(&token) => return Err(AgentError::TargetIdWithLegacyToken),
168            Some(target_id) if !is_legacy_token(&token) => {
169                let target_id = validate_target_id(&target_id)?;
170                endpoint.push_str(&format!("/{}", target_id));
171            }
172            None if !is_legacy_token(&token) => return Err(AgentError::MissingTargetId),
173            _ => {}
174        }
175
176        let circuit_breaker = if let Some(cb) = self.circuit_breaker {
177            cb
178        } else {
179            circuit_breaker::CircuitBreakerBuilder::default()
180                .build_async()
181                .map_err(AgentError::CircuitBreakerCreationError)?
182        };
183
184        let buffer = Buffer::new(self.buffer_size);
185
186        Ok(UsageAgentInner {
187            endpoint,
188            buffer,
189            processor: OperationProcessor::new(),
190            client,
191            flush_interval: self.flush_interval,
192            circuit_breaker,
193        })
194    }
195    pub fn build(self) -> Result<UsageAgent, AgentError> {
196        let agent = self.build_agent()?;
197        Ok(Arc::new(AsyncDropper::new(agent)))
198    }
199}
200
201// Target ID regexp for validation: slug format
202static SLUG_REGEX: LazyLock<regex_automata::meta::Regex> = LazyLock::new(|| {
203    regex_automata::meta::Regex::new(r"^[a-zA-Z0-9-_]+\/[a-zA-Z0-9-_]+\/[a-zA-Z0-9-_]+$").unwrap()
204});
205// Target ID regexp for validation: UUID format
206static UUID_REGEX: LazyLock<regex_automata::meta::Regex> = LazyLock::new(|| {
207    regex_automata::meta::Regex::new(
208        r"^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}$",
209    )
210    .unwrap()
211});
212
213fn validate_target_id(target_id: &str) -> Result<&str, AgentError> {
214    let trimmed_s = target_id.trim();
215    if trimmed_s.is_empty() {
216        Err(AgentError::InvalidTargetId("<empty>".to_string()))
217    } else {
218        if SLUG_REGEX.is_match(trimmed_s) {
219            return Ok(trimmed_s);
220        }
221        if UUID_REGEX.is_match(trimmed_s) {
222            return Ok(trimmed_s);
223        }
224        Err(AgentError::InvalidTargetId(format!(
225            "Invalid target_id format: '{}'. It must be either in slug format '$organizationSlug/$projectSlug/$targetSlug' or UUID format 'a0f4c605-6541-4350-8cfe-b31f21a4bf80'",
226            trimmed_s
227        )))
228    }
229}