hive_console_sdk/agent/
builder.rs

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

hive_console_sdk/agent/
builder.rs
