hive_console_sdk/agent/
builder.rs1use 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 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 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 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 pub fn buffer_size(mut self, buffer_size: usize) -> Self {
79 self.buffer_size = buffer_size;
80 self
81 }
82 pub fn connect_timeout(mut self, connect_timeout: Duration) -> Self {
85 self.connect_timeout = connect_timeout;
86 self
87 }
88 pub fn request_timeout(mut self, request_timeout: Duration) -> Self {
91 self.request_timeout = request_timeout;
92 self
93 }
94 pub fn accept_invalid_certs(mut self, accept_invalid_certs: bool) -> Self {
97 self.accept_invalid_certs = accept_invalid_certs;
98 self
99 }
100 pub fn flush_interval(mut self, flush_interval: Duration) -> Self {
103 self.flush_interval = flush_interval;
104 self
105 }
106 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 pub fn retry_policy(mut self, retry_policy: ExponentialBackoff) -> Self {
116 self.retry_policy = retry_policy;
117 self
118 }
119 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
201static 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});
205static 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}