hive_console_sdk/agent/
usage_agent.rs

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

hive_console_sdk/agent/
usage_agent.rs
