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

hive_console_sdk/agent/
builder.rs
