hive_console_sdk/agent/
usage_agent.rs

1use graphql_tools::parser::schema::Document;
2use rand::prelude::*;
3use recloser::AsyncRecloser;
4use reqwest_middleware::ClientWithMiddleware;
5use std::{
6    collections::{hash_map::Entry, BTreeMap, HashMap},
7    hash::{Hash, Hasher},
8    sync::Arc,
9    time::Duration,
10};
11use thiserror::Error;
12use tokio_util::sync::CancellationToken;
13use tracing::debug;
14
15use crate::expressions::lib::FromVrlValue;
16use crate::{
17    agent::{buffer::AddStatus, utils::OperationProcessor},
18    expressions::ExecutableProgram,
19    helpers::SharedFifoSet,
20};
21use crate::{
22    agent::{buffer::Buffer, builder::UsageAgentBuilder},
23    expressions::values::boolean::BooleanConversionError,
24};
25use vrl::{compiler::Program as VrlProgram, core::Value as VrlValue, value::KeyString};
26use xxhash_rust::xxh3::Xxh3;
27
28const USAGE_REPORTING_TARGET: &str = "console_sdk::usage_reporting";
29
30#[derive(Debug, Clone, Default)]
31pub enum OperationType {
32    #[default]
33    Query,
34    Mutation,
35    Subscription,
36}
37
38#[derive(Debug, Default, Clone)]
39pub struct ExecutionReport {
40    pub schema: Arc<Document<'static, String>>,
41    pub client_name: Option<String>,
42    pub client_version: Option<String>,
43    pub timestamp: u64,
44    pub duration: Duration,
45    pub ok: bool,
46    pub errors: usize,
47    pub operation_body: String,
48    pub operation_name: Option<String>,
49    pub operation_type: Option<OperationType>,
50    pub persisted_document_hash: Option<String>,
51}
52
53typify::import_types!(schema = "./usage-report-v2.schema.json");
54
55pub struct UsageAgentInner {
56    pub(crate) endpoint: String,
57    pub(crate) buffer: Buffer<ExecutionReport>,
58    pub(crate) processor: OperationProcessor,
59    pub(crate) client: ClientWithMiddleware,
60    pub(crate) flush_interval: Duration,
61    pub(crate) circuit_breaker: AsyncRecloser,
62    pub(crate) exclude: Option<Exclude>,
63    pub(crate) sample_rate: f64,
64    pub(crate) at_least_once: Option<AtLeastOnceSampling>,
65}
66
67#[derive(Debug, Clone)]
68pub enum Exclude {
69    OperationNames(Vec<String>),
70    Expression(Box<VrlProgram>),
71}
72
73#[derive(Debug, Clone, Copy)]
74pub enum SamplingKey {
75    OperationName,
76    OperationType,
77    OperationBody,
78}
79
80pub struct AtLeastOnceSampling {
81    pub(crate) key: Vec<SamplingKey>,
82    pub(crate) seen_hashes: SharedFifoSet,
83}
84
85pub fn non_empty_string(value: Option<String>) -> Option<String> {
86    value.filter(|str| !str.is_empty())
87}
88
89#[derive(Error, Debug)]
90pub enum AgentError {
91    #[error("unable to acquire lock: {0}")]
92    Lock(String),
93    #[error("unable to send report: unauthorized")]
94    Unauthorized,
95    #[error("unable to send report: no access")]
96    Forbidden,
97    #[error("unable to send report: rate limited")]
98    RateLimited,
99    #[error("missing token")]
100    MissingToken,
101    #[error("your access token requires providing a 'target_id' option.")]
102    MissingTargetId,
103    #[error("using 'target_id' with legacy tokens is not supported")]
104    TargetIdWithLegacyToken,
105    #[error("invalid token provided")]
106    InvalidToken,
107    #[error(
108        "invalid target id provided: {0}, it should be either a slug like \"$organizationSlug/$projectSlug/$targetSlug\" or an UUID"
109    )]
110    InvalidTargetId(String),
111    #[error("unable to instantiate the http client for reports sending: {0}")]
112    HTTPClientCreationError(reqwest::Error),
113    #[error("unable to create circuit breaker: {0}")]
114    CircuitBreakerCreationError(#[from] crate::circuit_breaker::CircuitBreakerError),
115    #[error("rejected by the circuit breaker")]
116    CircuitBreakerRejected,
117    #[error("unable to send report: {0}")]
118    Unknown(String),
119    #[error("failed to compile exclude expression: {0}")]
120    ExcludeExpressionCompileError(#[from] crate::expressions::ExpressionCompileError),
121    #[error("failed to execute exclude expression: {0}")]
122    ExcludeExpressionExecutionError(#[from] crate::expressions::ExpressionExecutionError),
123    #[error("failed to convert exclude expression result to boolean: {0}")]
124    ExcludeExpressionResultConversionError(#[from] BooleanConversionError),
125}
126
127pub struct UsageAgentHandle(Option<UsageAgentInner>);
128
129impl UsageAgentHandle {
130    pub(crate) fn new(inner: UsageAgentInner) -> Self {
131        Self(Some(inner))
132    }
133
134    fn inner(&self) -> &UsageAgentInner {
135        self.0.as_ref().expect("UsageAgentHandle used after drop")
136    }
137}
138
139impl Drop for UsageAgentHandle {
140    fn drop(&mut self) {
141        if let Some(inner) = self.0.take() {
142            tokio::spawn(async move {
143                if let Err(e) = inner.flush().await {
144                    tracing::error!(target: USAGE_REPORTING_TARGET, error = ?e, "Failed to flush usage reports during drop");
145                }
146            });
147        }
148    }
149}
150
151pub type UsageAgent = Arc<UsageAgentHandle>;
152
153#[async_trait::async_trait]
154pub trait UsageAgentExt {
155    fn builder() -> UsageAgentBuilder {
156        UsageAgentBuilder::default()
157    }
158    async fn flush(&self) -> Result<(), AgentError>;
159    async fn start_flush_interval(&self, token: &CancellationToken);
160    /// Deprecated: use [`add_report_with_request`] instead.
161    /// This method will be removed in a future version major.
162    #[deprecated(note = "use `add_report_with_request` instead")]
163    async fn add_report(&self, execution_report: ExecutionReport) -> Result<(), AgentError>;
164
165    async fn add_report_with_request(
166        &self,
167        execution_report: ExecutionReport,
168        request: Option<RequestDetails>,
169    ) -> Result<(), AgentError>;
170}
171
172impl UsageAgentInner {
173    fn should_exclude(
174        &self,
175        execution_report: &ExecutionReport,
176        request: Option<&RequestDetails>,
177    ) -> Result<bool, AgentError> {
178        self.exclude.as_ref().map_or(Ok(false), |exclude| {
179            exclude.should_exclude(execution_report, request)
180        })
181    }
182
183    fn should_sample(&self, execution_report: &ExecutionReport) -> bool {
184        if let Some(at_least_once) = &self.at_least_once {
185            let key_hash = at_least_once.resolve_key_hash(execution_report);
186            // Every first distinct report should be sampled
187            if at_least_once.mark_seen(key_hash) {
188                return true;
189            }
190        }
191
192        let sample_rate = self.sample_rate;
193        if sample_rate >= 1.0 {
194            return true;
195        }
196        if sample_rate <= 0.0 {
197            return false;
198        }
199
200        rand::rng().random_bool(sample_rate)
201    }
202
203    fn produce_report(&self, reports: Vec<ExecutionReport>) -> Result<Report, AgentError> {
204        let mut report = Report {
205            size: 0,
206            map: HashMap::new(),
207            operations: Vec::new(),
208            subscription_operations: Vec::new(),
209        };
210
211        // iterate over reports and check if they are valid
212        for op in reports {
213            let operation = self.processor.process(&op.operation_body, &op.schema);
214            match operation {
215                Err(e) => {
216                    tracing::warn!(
217                        target: USAGE_REPORTING_TARGET,
218                        error = ?e,
219                        operation_name = op.operation_name
220                            .clone()
221                            .or_else(|| Some("anonymous".to_string()))
222                            .unwrap(),
223                        phase = "PROCESSING",
224                        "Dropping operation",
225                    );
226
227                    continue;
228                }
229                Ok(operation) => match operation {
230                    Some(operation) => {
231                        let hash = operation.hash;
232
233                        let client_name = non_empty_string(op.client_name);
234                        let client_version = non_empty_string(op.client_version);
235
236                        let metadata: Option<Metadata> =
237                            if client_name.is_some() || client_version.is_some() {
238                                Some(Metadata {
239                                    client: Some(Client {
240                                        name: client_name.unwrap_or_default(),
241                                        version: client_version.unwrap_or_default(),
242                                    }),
243                                })
244                            } else {
245                                None
246                            };
247                        report.operations.push(RequestOperation {
248                            operation_map_key: hash.clone(),
249                            timestamp: op.timestamp,
250                            execution: Execution {
251                                ok: op.ok,
252                                /*
253                                    The conversion from u128 (from op.duration.as_nanos()) to u64 using try_into().unwrap() can panic if the duration is longer than u64::MAX nanoseconds (over 584 years).
254                                    While highly unlikely, it's safer to handle this potential overflow gracefully in library code to prevent panics.
255                                    A safe alternative is to convert the Result to an Option and provide a fallback value on failure,
256                                    effectively saturating at u64::MAX.
257                                */
258                                duration: op
259                                    .duration
260                                    .as_nanos()
261                                    .try_into()
262                                    .ok()
263                                    .unwrap_or(u64::MAX),
264                                errors_total: op.errors.try_into().unwrap(),
265                            },
266                            persisted_document_hash: op
267                                .persisted_document_hash
268                                .map(PersistedDocumentHash),
269                            metadata,
270                        });
271                        if let Entry::Vacant(e) = report.map.entry(ReportMapKey(hash)) {
272                            e.insert(OperationMapRecord {
273                                operation: operation.operation,
274                                operation_name: non_empty_string(op.operation_name),
275                                fields: operation.coordinates,
276                            });
277                        }
278                        report.size += 1;
279                    }
280                    None => {
281                        tracing::debug!(
282                            target: USAGE_REPORTING_TARGET,
283                            phase = "PROCESSING",
284                            "Dropping operation, probably introspection query"
285                        );
286                    }
287                },
288            }
289        }
290
291        Ok(report)
292    }
293
294    async fn send_report(&self, report: Report) -> Result<(), AgentError> {
295        if report.size == 0 {
296            return Ok(());
297        }
298        // Based on https://the-guild.dev/graphql/hive/docs/specs/usage-reports#data-structure
299        let resp_fut = self.client.post(&self.endpoint).json(&report).send();
300
301        let resp = self
302            .circuit_breaker
303            .call(resp_fut)
304            .await
305            .map_err(|e| match e {
306                recloser::Error::Inner(e) => AgentError::Unknown(e.to_string()),
307                recloser::Error::Rejected => AgentError::CircuitBreakerRejected,
308            })?;
309
310        match resp.status() {
311            reqwest::StatusCode::OK => Ok(()),
312            reqwest::StatusCode::UNAUTHORIZED => Err(AgentError::Unauthorized),
313            reqwest::StatusCode::FORBIDDEN => Err(AgentError::Forbidden),
314            reqwest::StatusCode::TOO_MANY_REQUESTS => Err(AgentError::RateLimited),
315            _ => Err(AgentError::Unknown(format!(
316                "({}) {}",
317                resp.status(),
318                resp.text().await.unwrap_or_default()
319            ))),
320        }
321    }
322
323    async fn handle_drained(&self, drained: Vec<ExecutionReport>) -> Result<(), AgentError> {
324        if drained.is_empty() {
325            return Ok(());
326        }
327        let report = self.produce_report(drained)?;
328        self.send_report(report).await
329    }
330
331    async fn flush(&self) -> Result<(), AgentError> {
332        let execution_reports = self.buffer.drain().await;
333
334        self.handle_drained(execution_reports).await?;
335
336        Ok(())
337    }
338}
339
340impl Exclude {
341    fn should_exclude(
342        &self,
343        execution_report: &ExecutionReport,
344        request: Option<&RequestDetails>,
345    ) -> Result<bool, AgentError> {
346        match self {
347            Exclude::OperationNames(operation_names) => Ok(execution_report
348                .operation_name
349                .as_deref()
350                .is_some_and(|operation_name| {
351                    operation_names.iter().any(|name| name == operation_name)
352                })),
353            Exclude::Expression(program) => {
354                let result = program.execute(get_vrl_value_from_execution_report_and_request(
355                    execution_report,
356                    request.cloned(),
357                ))?;
358                bool::from_vrl_value(result).map_err(AgentError::from)
359            }
360        }
361    }
362}
363
364#[derive(Debug, Clone)]
365pub struct RequestDetails {
366    pub method: http::Method,
367    pub url: http::Uri,
368    pub headers: Vec<(String, String)>,
369}
370
371#[async_trait::async_trait]
372impl UsageAgentExt for UsageAgent {
373    async fn flush(&self) -> Result<(), AgentError> {
374        self.inner().flush().await
375    }
376
377    async fn start_flush_interval(&self, token: &CancellationToken) {
378        loop {
379            tokio::select! {
380                _ = token.cancelled() => {
381                    debug!(target: USAGE_REPORTING_TARGET, "Shutting down.");
382                    return;
383                }
384                _ = tokio::time::sleep(self.inner().flush_interval) => {}
385            }
386
387            self.flush()
388                .await
389                .unwrap_or_else(|e| tracing::error!(target: USAGE_REPORTING_TARGET, error = ?e, "Failed to flush usage reports"));
390        }
391    }
392
393    async fn add_report_with_request(
394        &self,
395        execution_report: ExecutionReport,
396        request: Option<RequestDetails>,
397    ) -> Result<(), AgentError> {
398        let inner = self.inner();
399
400        if inner.should_exclude(&execution_report, request.as_ref())? {
401            tracing::debug!(
402                target: USAGE_REPORTING_TARGET,
403                operation_name = execution_report
404                    .operation_name
405                    .as_deref()
406                    .unwrap_or("anonymous"),
407                "Excluding report for operation based on exclude expression evaluation",
408            );
409
410            return Ok(());
411        }
412
413        if !inner.should_sample(&execution_report) {
414            tracing::debug!(
415                target: USAGE_REPORTING_TARGET,
416                operation_name = execution_report
417                    .operation_name
418                    .as_deref()
419                    .unwrap_or("anonymous"),
420                "Sampling dropped report for operation",
421            );
422
423            return Ok(());
424        }
425
426        if let AddStatus::Full { drained } = inner.buffer.add(execution_report).await {
427            inner.handle_drained(drained).await?;
428        }
429
430        Ok(())
431    }
432    async fn add_report(&self, execution_report: ExecutionReport) -> Result<(), AgentError> {
433        self.add_report_with_request(execution_report, None).await
434    }
435}
436
437impl<'req, TBody> From<&'req http::Request<TBody>> for RequestDetails {
438    fn from(req: &'req http::Request<TBody>) -> Self {
439        let mut headers = Vec::with_capacity(req.headers().len());
440        for (name, value) in req.headers().iter() {
441            if let Ok(val_str) = value.to_str() {
442                headers.push((name.to_string(), val_str.to_string()));
443            }
444        }
445
446        RequestDetails {
447            method: req.method().clone(),
448            url: req.uri().clone(),
449            headers,
450        }
451    }
452}
453
454impl From<RequestDetails> for VrlValue {
455    fn from(details: RequestDetails) -> Self {
456        let mut merged_headers: BTreeMap<String, String> = BTreeMap::new();
457        for (header_name, header_value) in details.headers {
458            if let Some(existing_value) = merged_headers.get_mut(&header_name) {
459                existing_value.push_str(", ");
460                existing_value.push_str(&header_value);
461            } else {
462                merged_headers.insert(header_name, header_value);
463            }
464        }
465
466        let headers_value: BTreeMap<KeyString, VrlValue> = merged_headers
467            .into_iter()
468            .map(|(key, value)| (key.into(), value.into()))
469            .collect();
470        let headers_value = VrlValue::Object(headers_value);
471
472        // .request.url
473        let url_value = VrlValue::Object(BTreeMap::from([
474            ("host".into(), details.url.host().unwrap_or_default().into()),
475            ("path".into(), details.url.path().into()),
476            (
477                "port".into(),
478                details
479                    .url
480                    .port_u16()
481                    .map(|p| VrlValue::Integer(p.into()))
482                    .unwrap_or(VrlValue::Null),
483            ),
484        ]));
485
486        // .request
487        VrlValue::Object(BTreeMap::from([
488            ("method".into(), details.method.as_str().into()),
489            ("headers".into(), headers_value),
490            ("url".into(), url_value),
491        ]))
492    }
493}
494
495impl AtLeastOnceSampling {
496    fn resolve_key_hash(&self, report: &ExecutionReport) -> u64 {
497        let mut hasher = Xxh3::new();
498
499        for key in &self.key {
500            let value = match key {
501                SamplingKey::OperationName => {
502                    report.operation_name.as_deref().unwrap_or("anonymous")
503                }
504                SamplingKey::OperationType => match report.operation_type.as_ref() {
505                    None | Some(OperationType::Query) => "query",
506                    Some(OperationType::Mutation) => "mutation",
507                    Some(OperationType::Subscription) => "subscription",
508                },
509                SamplingKey::OperationBody => report.operation_body.as_str(),
510            };
511            value.hash(&mut hasher);
512            0u8.hash(&mut hasher);
513        }
514
515        hasher.finish()
516    }
517
518    /// Marks the given key hash as seen, returning true if it was not already seen.
519    fn mark_seen(&self, key_hash: u64) -> bool {
520        self.seen_hashes.insert(key_hash)
521    }
522}
523
524pub fn get_vrl_value_from_execution_report_and_request(
525    report: &ExecutionReport,
526    request: Option<RequestDetails>,
527) -> VrlValue {
528    let mut map = BTreeMap::from([("default".into(), VrlValue::Boolean(false))]);
529    let mut request_map = BTreeMap::new();
530
531    if let Some(request_details) = request {
532        if let VrlValue::Object(request_object) = VrlValue::from(request_details) {
533            request_map.extend(request_object);
534        }
535    }
536
537    request_map.insert(
538        "operation".into(),
539        VrlValue::Object(BTreeMap::from([
540            (
541                "name".into(),
542                report.operation_name.clone().unwrap_or_default().into(),
543            ),
544            (
545                "type".into(),
546                report
547                    .operation_type
548                    .as_ref()
549                    .map(|operation_type| match operation_type {
550                        OperationType::Query => "query",
551                        OperationType::Mutation => "mutation",
552                        OperationType::Subscription => "subscription",
553                    })
554                    .unwrap_or_default()
555                    .into(),
556            ),
557            ("query".into(), report.operation_body.clone().into()),
558        ])),
559    );
560
561    if let Ok(timestamp_integer) = report.timestamp.try_into() {
562        let timestamp_value = VrlValue::Integer(timestamp_integer);
563        map.insert("timestamp".into(), timestamp_value);
564        request_map.insert("timestamp".into(), VrlValue::Integer(timestamp_integer));
565    }
566
567    map.insert("request".into(), VrlValue::Object(request_map));
568
569    VrlValue::Object(map)
570}
571
572#[cfg(test)]
573mod tests {
574    use std::{sync::Arc, time::Duration};
575
576    use graphql_tools::parser::{parse_query, parse_schema};
577    use reqwest::{
578        header::{AUTHORIZATION, CONTENT_TYPE, USER_AGENT},
579        Method,
580    };
581    use vrl::core::Value as VrlValue;
582    use vrl::value::KeyString;
583
584    use crate::agent::usage_agent::{
585        get_vrl_value_from_execution_report_and_request, ExecutionReport, OperationType, Report,
586        UsageAgent, UsageAgentExt,
587    };
588
589    async fn wait_for_mock(mock: &mockito::Mock) {
590        tokio::time::timeout(Duration::from_secs(2), async {
591            while !mock.matched_async().await {
592                tokio::time::sleep(Duration::from_millis(10)).await;
593            }
594        })
595        .await
596        .expect("mock should be matched after usage agent drop flush");
597    }
598
599    /// Helper to extract a nested VRL value from an Object using string keys.
600    fn vrl_get<'a>(value: &'a VrlValue, keys: &[&str]) -> &'a VrlValue {
601        let mut current = value;
602        for key in keys {
603            match current {
604                VrlValue::Object(map) => {
605                    let ks: KeyString = (*key).into();
606                    current = map
607                        .get(&ks)
608                        .unwrap_or_else(|| panic!("key '{}' not found", key));
609                }
610                _ => panic!("expected Object at key '{}'", key),
611            }
612        }
613        current
614    }
615
616    const CONTENT_TYPE_VALUE: &'static str = "application/json";
617    const GRAPHQL_CLIENT_NAME: &'static str = "Hive Client";
618    const GRAPHQL_CLIENT_VERSION: &'static str = "1.0.0";
619
620    #[tokio::test(flavor = "multi_thread")]
621    async fn should_send_data_to_hive() -> Result<(), Box<dyn std::error::Error>> {
622        let token = "Token";
623
624        let mut server = mockito::Server::new_async().await;
625
626        let server_url = server.url();
627
628        let timestamp = 1625247600;
629        let duration = Duration::from_millis(20);
630        let user_agent = "hive-router-sdk-test";
631
632        let mock = server
633            .mock("POST", "/200")
634            .match_header(AUTHORIZATION, format!("Bearer {}", token).as_str())
635            .match_header(CONTENT_TYPE, CONTENT_TYPE_VALUE)
636            .match_header(USER_AGENT, user_agent)
637            .match_header("X-Usage-API-Version", "2")
638            .match_request(move |request| {
639                let request_body = request.body().expect("Failed to extract body");
640                let report: Report = serde_json::from_slice(request_body)
641                    .expect("Failed to parse request body as JSON");
642                assert_eq!(report.size, 1);
643                let record = report.map.values().next().expect("No operation record");
644                // operation
645                assert!(record.operation.contains("mutation deleteProject"));
646                assert_eq!(record.operation_name.as_deref(), Some("deleteProject"));
647                // fields
648                let expected_fields = vec![
649                    "Mutation.deleteProject",
650                    "Mutation.deleteProject.selector",
651                    "DeleteProjectPayload.selector",
652                    "ProjectSelector.organization",
653                    "ProjectSelector.project",
654                    "DeleteProjectPayload.deletedProject",
655                    "Project.id",
656                    "Project.cleanId",
657                    "Project.name",
658                    "Project.type",
659                    "ProjectType.FEDERATION",
660                    "ProjectType.STITCHING",
661                    "ProjectType.SINGLE",
662                    "ProjectType.CUSTOM",
663                    "ProjectSelectorInput.organization",
664                    "ID",
665                    "ProjectSelectorInput.project",
666                ];
667                for field in &expected_fields {
668                    assert!(
669                        record.fields.contains(&field.to_string()),
670                        "Missing field: {}",
671                        field
672                    );
673                }
674                assert_eq!(
675                    record.fields.len(),
676                    expected_fields.len(),
677                    "Unexpected number of fields"
678                );
679
680                // Operations
681                let operations = report.operations;
682                assert_eq!(operations.len(), 1); // one operation
683
684                let operation = &operations[0];
685                let key = report.map.keys().next().expect("No operation key");
686                assert_eq!(operation.operation_map_key, key.0);
687                assert_eq!(operation.timestamp, timestamp);
688                assert_eq!(operation.execution.duration, duration.as_nanos() as u64);
689                assert_eq!(operation.execution.ok, true);
690                assert_eq!(operation.execution.errors_total, 0);
691                true
692            })
693            .expect(1)
694            .with_status(200)
695            .create_async()
696            .await;
697        let schema: graphql_tools::static_graphql::schema::Document = parse_schema(
698            r#"
699                type Query {
700                    project(selector: ProjectSelectorInput!): Project
701                    projectsByType(type: ProjectType!): [Project!]!
702                    projects(filter: FilterInput): [Project!]!
703                }
704
705                type Mutation {
706                    deleteProject(selector: ProjectSelectorInput!): DeleteProjectPayload!
707                }
708
709                input ProjectSelectorInput {
710                    organization: ID!
711                    project: ID!
712                }
713
714                input FilterInput {
715                    type: ProjectType
716                    pagination: PaginationInput
717                }
718
719                input PaginationInput {
720                    limit: Int
721                    offset: Int
722                }
723
724                type ProjectSelector {
725                    organization: ID!
726                    project: ID!
727                }
728
729                type DeleteProjectPayload {
730                    selector: ProjectSelector!
731                    deletedProject: Project!
732                }
733
734                type Project {
735                    id: ID!
736                    cleanId: ID!
737                    name: String!
738                    type: ProjectType!
739                    buildUrl: String
740                    validationUrl: String
741                }
742
743                enum ProjectType {
744                    FEDERATION
745                    STITCHING
746                    SINGLE
747                    CUSTOM
748                }
749        "#,
750        )?;
751
752        let op: graphql_tools::static_graphql::query::Document = parse_query(
753            r#"
754                mutation deleteProject($selector: ProjectSelectorInput!) {
755                    deleteProject(selector: $selector) {
756                    selector {
757                        organization
758                        project
759                    }
760                    deletedProject {
761                        ...ProjectFields
762                    }
763                    }
764                }
765
766                fragment ProjectFields on Project {
767                    id
768                    cleanId
769                    name
770                    type
771                }
772        "#,
773        )?;
774
775        // Testing async drop
776        {
777            let usage_agent = UsageAgent::builder()
778                .token(token.into())
779                .endpoint(format!("{}/200", server_url))
780                .user_agent(user_agent.into())
781                .build()?;
782
783            let request = http::Request::builder()
784                .method(Method::POST)
785                .uri("http://localhost/graphql")
786                .body(())
787                .unwrap();
788
789            usage_agent
790                .add_report_with_request(
791                    ExecutionReport {
792                        schema: Arc::new(schema),
793                        operation_body: op.to_string(),
794                        operation_name: Some("deleteProject".to_string()),
795                        operation_type: Some(OperationType::Mutation),
796                        client_name: Some(GRAPHQL_CLIENT_NAME.to_string()),
797                        client_version: Some(GRAPHQL_CLIENT_VERSION.to_string()),
798                        timestamp,
799                        duration,
800                        ok: true,
801                        errors: 0,
802                        persisted_document_hash: None,
803                    },
804                    Some((&request).into()),
805                )
806                .await?;
807        }
808
809        wait_for_mock(&mock).await;
810        mock.assert_async().await;
811
812        Ok(())
813    }
814
815    fn make_test_report(
816        operation_name: Option<&str>,
817        operation_type: OperationType,
818        operation_body: &str,
819    ) -> ExecutionReport {
820        let schema: graphql_tools::static_graphql::schema::Document =
821            parse_schema("type Query { hello: String }").unwrap();
822
823        ExecutionReport {
824            schema: Arc::new(schema),
825            operation_body: operation_body.to_string(),
826            operation_name: operation_name.map(|s| s.to_string()),
827            operation_type: Some(operation_type),
828            client_name: Some("test-client".to_string()),
829            client_version: Some("1.0.0".to_string()),
830            timestamp: 1625247600,
831            duration: Duration::from_millis(10),
832            ok: true,
833            errors: 0,
834            persisted_document_hash: None,
835        }
836    }
837
838    fn make_simple_report(
839        operation_name: Option<&str>,
840        operation_type: OperationType,
841    ) -> ExecutionReport {
842        make_test_report(operation_name, operation_type, "query { hello }")
843    }
844
845    fn make_simple_request() -> http::Request<()> {
846        http::Request::builder()
847            .method(Method::GET)
848            .uri("http://localhost/graphql")
849            .body(())
850            .unwrap()
851    }
852
853    #[test]
854    fn vrl_value_contains_operation_name() {
855        let report = make_simple_report(Some("MyQuery"), OperationType::Query);
856        let request = make_simple_request();
857        let value =
858            get_vrl_value_from_execution_report_and_request(&report, Some((&request).into()));
859
860        let name = vrl_get(&value, &["request", "operation", "name"]);
861        assert_eq!(name, &VrlValue::from("MyQuery"));
862    }
863
864    #[test]
865    fn vrl_value_contains_operation_type_query() {
866        let report = make_simple_report(Some("Q"), OperationType::Query);
867        let request = make_simple_request();
868        let value =
869            get_vrl_value_from_execution_report_and_request(&report, Some((&request).into()));
870
871        let op_type = vrl_get(&value, &["request", "operation", "type"]);
872        assert_eq!(op_type, &VrlValue::from("query"));
873    }
874
875    #[test]
876    fn vrl_value_contains_operation_type_mutation() {
877        let report = make_simple_report(Some("M"), OperationType::Mutation);
878        let request = make_simple_request();
879        let value =
880            get_vrl_value_from_execution_report_and_request(&report, Some((&request).into()));
881
882        let op_type = vrl_get(&value, &["request", "operation", "type"]);
883        assert_eq!(op_type, &VrlValue::from("mutation"));
884    }
885
886    #[test]
887    fn vrl_value_contains_operation_type_subscription() {
888        let report = make_simple_report(Some("S"), OperationType::Subscription);
889        let request = make_simple_request();
890        let value =
891            get_vrl_value_from_execution_report_and_request(&report, Some((&request).into()));
892
893        let op_type = vrl_get(&value, &["request", "operation", "type"]);
894        assert_eq!(op_type, &VrlValue::from("subscription"));
895    }
896
897    #[test]
898    fn vrl_value_contains_operation_body() {
899        let report = make_simple_report(Some("Q"), OperationType::Query);
900        let request = make_simple_request();
901        let value =
902            get_vrl_value_from_execution_report_and_request(&report, Some((&request).into()));
903
904        let query = vrl_get(&value, &["request", "operation", "query"]);
905        assert_eq!(query, &VrlValue::from("query { hello }"));
906    }
907
908    #[test]
909    fn vrl_value_contains_request_method() {
910        let report = make_test_report(Some("Q"), OperationType::Query, "query { hello }");
911        let request = http::Request::builder()
912            .method(Method::GET)
913            .uri("http://localhost/graphql")
914            .body(())
915            .unwrap();
916        let value =
917            get_vrl_value_from_execution_report_and_request(&report, Some((&request).into()));
918
919        let method = vrl_get(&value, &["request", "method"]);
920        assert_eq!(method, &VrlValue::from("GET"));
921    }
922
923    #[test]
924    fn vrl_value_contains_url_details() {
925        let report = make_test_report(Some("Q"), OperationType::Query, "query { hello }");
926        let request = http::Request::builder()
927            .method(Method::POST)
928            .uri("http://api.example.com:8080/v1/graphql")
929            .body(())
930            .unwrap();
931        let value =
932            get_vrl_value_from_execution_report_and_request(&report, Some((&request).into()));
933
934        assert_eq!(
935            vrl_get(&value, &["request", "url", "host"]),
936            &VrlValue::from("api.example.com")
937        );
938        assert_eq!(
939            vrl_get(&value, &["request", "url", "port"]),
940            &VrlValue::Integer(8080)
941        );
942        assert_eq!(
943            vrl_get(&value, &["request", "url", "path"]),
944            &VrlValue::from("/v1/graphql")
945        );
946    }
947
948    #[test]
949    fn vrl_value_contains_headers() {
950        let request = http::Request::builder()
951            .method(Method::POST)
952            .uri("http://localhost/graphql")
953            .header("x-custom-header", "custom-value")
954            .header("authorization", "Bearer token123")
955            .body(())
956            .unwrap();
957        let report = make_test_report(Some("Q"), OperationType::Query, "query { hello }");
958        let value =
959            get_vrl_value_from_execution_report_and_request(&report, Some((&request).into()));
960
961        assert_eq!(
962            vrl_get(&value, &["request", "headers", "x-custom-header"]),
963            &VrlValue::from("custom-value")
964        );
965        assert_eq!(
966            vrl_get(&value, &["request", "headers", "authorization"]),
967            &VrlValue::from("Bearer token123")
968        );
969    }
970
971    #[test]
972    fn vrl_value_joins_duplicate_header_values() {
973        let mut request = http::Request::builder()
974            .method(Method::POST)
975            .uri("http://localhost/graphql")
976            .body(())
977            .unwrap();
978
979        request
980            .headers_mut()
981            .append("x-scope", http::HeaderValue::from_static("one"));
982        request
983            .headers_mut()
984            .append("x-scope", http::HeaderValue::from_static("two"));
985
986        let report = make_test_report(Some("Q"), OperationType::Query, "query { hello }");
987        let value =
988            get_vrl_value_from_execution_report_and_request(&report, Some((&request).into()));
989
990        assert_eq!(
991            vrl_get(&value, &["request", "headers", "x-scope"]),
992            &VrlValue::from("one, two")
993        );
994    }
995
996    #[test]
997    fn vrl_value_anonymous_operation_has_empty_name() {
998        let report = make_simple_report(None, OperationType::Query);
999        let request = make_simple_request();
1000        let value =
1001            get_vrl_value_from_execution_report_and_request(&report, Some((&request).into()));
1002
1003        let name = vrl_get(&value, &["request", "operation", "name"]);
1004        assert_eq!(name, &VrlValue::from(""));
1005    }
1006
1007    #[test]
1008    fn vrl_value_has_default_false() {
1009        let report = make_simple_report(Some("Q"), OperationType::Query);
1010        let request = make_simple_request();
1011        let value =
1012            get_vrl_value_from_execution_report_and_request(&report, Some((&request).into()));
1013
1014        let default_val = vrl_get(&value, &["default"]);
1015        assert_eq!(default_val, &VrlValue::Boolean(false));
1016    }
1017
1018    #[tokio::test(flavor = "multi_thread")]
1019    async fn exclude_expression_filters_by_operation_name() -> Result<(), Box<dyn std::error::Error>>
1020    {
1021        let token = "Token";
1022        let mut server = mockito::Server::new_async().await;
1023        let server_url = server.url();
1024
1025        // The mock expects exactly 0 requests because the operation should be excluded
1026        let mock = server
1027            .mock("POST", "/200")
1028            .expect(0)
1029            .with_status(200)
1030            .create_async()
1031            .await;
1032
1033        {
1034            let usage_agent = UsageAgent::builder()
1035                .token(token.into())
1036                .endpoint(format!("{}/200", server_url))
1037                .buffer_size(1) // flush on every report
1038                .exclude_expression(r#".request.operation.name == "ExcludeMe""#.to_string())
1039                .build()?;
1040
1041            // This report should be excluded
1042            let report = make_simple_report(Some("ExcludeMe"), OperationType::Query);
1043            let request = make_simple_request();
1044            usage_agent
1045                .add_report_with_request(report, Some((&request).into()))
1046                .await?;
1047        }
1048
1049        mock.assert_async().await;
1050        Ok(())
1051    }
1052
1053    #[tokio::test(flavor = "multi_thread")]
1054    async fn exclude_expression_allows_non_matching_operations(
1055    ) -> Result<(), Box<dyn std::error::Error>> {
1056        let token = "Token";
1057        let mut server = mockito::Server::new_async().await;
1058        let server_url = server.url();
1059
1060        // This operation should NOT be excluded, so we expect 1 request (via async drop flush)
1061        let mock = server
1062            .mock("POST", "/200")
1063            .expect(1)
1064            .with_status(200)
1065            .create_async()
1066            .await;
1067
1068        {
1069            let usage_agent = UsageAgent::builder()
1070                .token(token.into())
1071                .endpoint(format!("{}/200", server_url))
1072                .exclude_expression(r#".request.operation.name == "ExcludeMe""#.to_string())
1073                .build()?;
1074
1075            let report = make_simple_report(Some("KeepMe"), OperationType::Query);
1076            let request = make_simple_request();
1077            usage_agent
1078                .add_report_with_request(report, Some((&request).into()))
1079                .await?;
1080        }
1081
1082        wait_for_mock(&mock).await;
1083        mock.assert_async().await;
1084        Ok(())
1085    }
1086
1087    #[tokio::test(flavor = "multi_thread")]
1088    async fn exclude_expression_filters_by_operation_type() -> Result<(), Box<dyn std::error::Error>>
1089    {
1090        let token = "Token";
1091        let mut server = mockito::Server::new_async().await;
1092        let server_url = server.url();
1093
1094        let mock = server
1095            .mock("POST", "/200")
1096            .expect(0)
1097            .with_status(200)
1098            .create_async()
1099            .await;
1100
1101        {
1102            let usage_agent = UsageAgent::builder()
1103                .token(token.into())
1104                .endpoint(format!("{}/200", server_url))
1105                .buffer_size(1)
1106                .exclude_expression(r#".request.operation.type == "subscription""#.to_string())
1107                .build()?;
1108
1109            let report = make_simple_report(Some("OnMessage"), OperationType::Subscription);
1110            let request = make_simple_request();
1111            usage_agent
1112                .add_report_with_request(report, Some((&request).into()))
1113                .await?;
1114        }
1115
1116        mock.assert_async().await;
1117        Ok(())
1118    }
1119
1120    #[tokio::test(flavor = "multi_thread")]
1121    async fn exclude_expression_filters_by_header() -> Result<(), Box<dyn std::error::Error>> {
1122        let token = "Token";
1123        let mut server = mockito::Server::new_async().await;
1124        let server_url = server.url();
1125
1126        let mock = server
1127            .mock("POST", "/200")
1128            .expect(0)
1129            .with_status(200)
1130            .create_async()
1131            .await;
1132
1133        {
1134            let usage_agent = UsageAgent::builder()
1135                .token(token.into())
1136                .endpoint(format!("{}/200", server_url))
1137                .buffer_size(1)
1138                .exclude_expression(r#".request.headers."x-internal" == "true""#.to_string())
1139                .build()?;
1140
1141            let request = http::Request::builder()
1142                .method(Method::POST)
1143                .uri("http://localhost/graphql")
1144                .header("x-internal", "true")
1145                .body(())
1146                .unwrap();
1147
1148            let report = make_test_report(Some("Q"), OperationType::Query, "query { hello }");
1149            usage_agent
1150                .add_report_with_request(report, Some((&request).into()))
1151                .await?;
1152        }
1153
1154        mock.assert_async().await;
1155        Ok(())
1156    }
1157
1158    #[tokio::test(flavor = "multi_thread")]
1159    async fn exclude_expression_complex_conditional() -> Result<(), Box<dyn std::error::Error>> {
1160        let token = "Token";
1161        let mut server = mockito::Server::new_async().await;
1162        let server_url = server.url();
1163
1164        // The expression excludes IntrospectionQuery OR any mutation
1165        let exclude_expr = r#"
1166            if (.request.operation.name == "IntrospectionQuery") {
1167                true
1168            } else if (.request.operation.type == "mutation") {
1169                true
1170            } else {
1171                false
1172            }
1173        "#;
1174
1175        let mock = server
1176            .mock("POST", "/200")
1177            .expect(0)
1178            .with_status(200)
1179            .create_async()
1180            .await;
1181
1182        {
1183            let usage_agent = UsageAgent::builder()
1184                .token(token.into())
1185                .endpoint(format!("{}/200", server_url))
1186                .buffer_size(1)
1187                .exclude_expression(exclude_expr.to_string())
1188                .build()?;
1189
1190            let request = make_simple_request();
1191
1192            // Excluded: IntrospectionQuery
1193            let report = make_simple_report(Some("IntrospectionQuery"), OperationType::Query);
1194            usage_agent
1195                .add_report_with_request(report, Some((&request).into()))
1196                .await?;
1197
1198            // Excluded: any mutation
1199            let report = make_simple_report(Some("CreateUser"), OperationType::Mutation);
1200            usage_agent
1201                .add_report_with_request(report, Some((&request).into()))
1202                .await?;
1203        }
1204
1205        mock.assert_async().await;
1206        Ok(())
1207    }
1208
1209    #[tokio::test(flavor = "multi_thread")]
1210    async fn exclude_expression_allows_through_complex_conditional(
1211    ) -> Result<(), Box<dyn std::error::Error>> {
1212        let token = "Token";
1213        let mut server = mockito::Server::new_async().await;
1214        let server_url = server.url();
1215
1216        let exclude_expr = r#"
1217            if (.request.operation.name == "IntrospectionQuery") {
1218                true
1219            } else if (.request.operation.type == "mutation") {
1220                true
1221            } else {
1222                false
1223            }
1224        "#;
1225
1226        // A normal query should NOT be excluded
1227        let mock = server
1228            .mock("POST", "/200")
1229            .expect(1)
1230            .with_status(200)
1231            .create_async()
1232            .await;
1233
1234        {
1235            let usage_agent = UsageAgent::builder()
1236                .token(token.into())
1237                .endpoint(format!("{}/200", server_url))
1238                .exclude_expression(exclude_expr.to_string())
1239                .build()?;
1240
1241            let request = make_simple_request();
1242
1243            let report = make_simple_report(Some("GetUsers"), OperationType::Query);
1244            usage_agent
1245                .add_report_with_request(report, Some((&request).into()))
1246                .await?;
1247        }
1248
1249        wait_for_mock(&mock).await;
1250        mock.assert_async().await;
1251        Ok(())
1252    }
1253
1254    #[tokio::test(flavor = "multi_thread")]
1255    async fn exclude_expression_filters_by_url_path() -> Result<(), Box<dyn std::error::Error>> {
1256        let token = "Token";
1257        let mut server = mockito::Server::new_async().await;
1258        let server_url = server.url();
1259
1260        let mock = server
1261            .mock("POST", "/200")
1262            .expect(0)
1263            .with_status(200)
1264            .create_async()
1265            .await;
1266
1267        {
1268            let usage_agent = UsageAgent::builder()
1269                .token(token.into())
1270                .endpoint(format!("{}/200", server_url))
1271                .buffer_size(1)
1272                .exclude_expression(r#".request.url.path == "/internal/graphql""#.to_string())
1273                .build()?;
1274
1275            let request = http::Request::builder()
1276                .method(Method::POST)
1277                .uri("http://localhost/internal/graphql")
1278                .body(())
1279                .unwrap();
1280
1281            let report = make_test_report(Some("Q"), OperationType::Query, "query { hello }");
1282            usage_agent
1283                .add_report_with_request(report, Some((&request).into()))
1284                .await?;
1285        }
1286
1287        mock.assert_async().await;
1288        Ok(())
1289    }
1290
1291    #[tokio::test(flavor = "multi_thread")]
1292    async fn no_exclude_expression_sends_all_reports() -> Result<(), Box<dyn std::error::Error>> {
1293        let token = "Token";
1294        let mut server = mockito::Server::new_async().await;
1295        let server_url = server.url();
1296
1297        let mock = server
1298            .mock("POST", "/200")
1299            .expect(1)
1300            .with_status(200)
1301            .create_async()
1302            .await;
1303
1304        {
1305            let usage_agent = UsageAgent::builder()
1306                .token(token.into())
1307                .endpoint(format!("{}/200", server_url))
1308                .build()?;
1309
1310            let report = make_simple_report(Some("AnyOp"), OperationType::Query);
1311            let request = make_simple_request();
1312            usage_agent
1313                .add_report_with_request(report, Some((&request).into()))
1314                .await?;
1315        }
1316
1317        wait_for_mock(&mock).await;
1318        mock.assert_async().await;
1319        Ok(())
1320    }
1321
1322    #[test]
1323    fn builder_rejects_invalid_exclude_expression() {
1324        let result = UsageAgent::builder()
1325            .token("Token".into())
1326            .exclude_expression("this is not valid VRL }{".to_string())
1327            .build();
1328
1329        assert!(result.is_err());
1330    }
1331
1332    #[tokio::test(flavor = "multi_thread")]
1333    async fn builder_ignores_empty_exclude_expression() {
1334        let result = UsageAgent::builder()
1335            .token("Token".into())
1336            .exclude_expression("".to_string())
1337            .build();
1338
1339        assert!(result.is_ok());
1340    }
1341}
hive_console_sdk/agent/usage_agent.rs

hive_console_sdk/agent/
usage_agent.rs
