cx_sdk/client/
data_usage.rs

1// Copyright 2024 Coralogix Ltd.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     https://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::{
16    str::FromStr,
17    time::Duration,
18};
19
20use crate::{
21    auth::AuthContext,
22    error::{
23        Result,
24        SdkApiError,
25        SdkError,
26    },
27    metadata::CallProperties,
28    util::make_request_with_metadata,
29};
30
31pub use crate::com::coralogixapis::actions::v2::Action;
32
33pub use cx_api::proto::com::coralogix::datausage::v2::{
34    DateRange,
35    Dimension,
36    ScopesFilter,
37};
38use cx_api::proto::com::coralogix::datausage::v2::{
39    GetDataUsageMetricsExportStatusRequest,
40    GetDataUsageMetricsExportStatusResponse,
41    GetDataUsageRequest,
42    GetDataUsageResponse,
43    GetLogsCountRequest,
44    GetLogsCountResponse,
45    GetSpansCountRequest,
46    GetSpansCountResponse,
47    UpdateDataUsageMetricsExportStatusRequest,
48    UpdateDataUsageMetricsExportStatusResponse,
49    data_usage_service_client::DataUsageServiceClient,
50};
51use tokio::sync::Mutex;
52use tonic::{
53    Streaming,
54    metadata::MetadataMap,
55    transport::{
56        Channel,
57        ClientTlsConfig,
58        Endpoint,
59    },
60};
61
62use crate::CoralogixRegion;
63
64pub use crate::com::coralogix::enrichment::v1::{
65    EnrichmentFieldDefinition,
66    EnrichmentRequestModel as EnrichmentMapping,
67    EnrichmentType,
68    enrichment_type::Type,
69};
70
71const DATA_PLANS_FEATURE_GROUP_ID: &str = "dataplans";
72
73/// A client for the Data Usage API.
74/// Read more at <https://coralogix.com/docs/user-guides/account-management/payment-and-billing/data-usage/>
75pub struct DataUsageClient {
76    metadata_map: MetadataMap,
77    service_client: Mutex<DataUsageServiceClient<Channel>>,
78}
79
80impl DataUsageClient {
81    /// Creates a new client for the Enrichments API.
82    ///
83    /// # Arguments
84    /// * `auth_context` - The [`AuthContext`] to use for authentication.
85    /// * `region` - The [`CoralogixRegion`] to connect to.
86    pub fn new(auth_context: AuthContext, region: CoralogixRegion) -> Result<Self> {
87        let channel: Channel = Endpoint::from_str(region.grpc_endpoint().as_str())?
88            .tls_config(ClientTlsConfig::new().with_native_roots())?
89            .connect_lazy();
90        let request_metadata: CallProperties = (&auth_context.team_level_api_key).into();
91        Ok(Self {
92            metadata_map: request_metadata.to_metadata_map(),
93            service_client: Mutex::new(DataUsageServiceClient::new(channel)),
94        })
95    }
96
97    /// Gets the data usage for the provided filters.
98    ///
99    /// # Arguments
100    /// * `date_range` - The date range to get the data usage for.
101    /// * `resolution` - The resolution to get the data usage for.
102    /// * `aggregate` - The aggregate to get the data usage for.
103    /// * `dimension_filters` - The dimension filters to apply.
104    pub async fn get_data_usage(
105        &self,
106        date_range: Option<DateRange>,
107        resolution: Option<Duration>,
108        aggregate: Vec<i32>,
109        dimension_filters: Vec<Dimension>,
110    ) -> Result<Streaming<GetDataUsageResponse>> {
111        self.service_client
112            .lock()
113            .await
114            .get_data_usage(make_request_with_metadata(
115                GetDataUsageRequest {
116                    date_range,
117                    resolution: resolution.map(|r| r.try_into().unwrap()),
118                    aggregate,
119                    dimension_filters,
120                },
121                &self.metadata_map,
122            ))
123            .await
124            .map(|r| r.into_inner())
125            .map_err(|status| {
126                SdkError::ApiError(SdkApiError {
127                    status,
128                    endpoint: "/com.coralogix.datausage.v2.DataUsageService/GetDataUsage".into(),
129                    feature_group: DATA_PLANS_FEATURE_GROUP_ID.into(),
130                })
131            })
132    }
133
134    /// Gets the span count for the provided filters.
135    ///
136    /// # Arguments
137    /// * `date_range` - The date range to get the span count for.
138    /// * `resolution` - The resolution to get the span count for.
139    /// * `filters` - The filters to apply.
140    pub async fn get_spans_count(
141        &self,
142        date_range: Option<DateRange>,
143        resolution: Option<Duration>,
144        filters: Option<ScopesFilter>,
145    ) -> Result<Streaming<GetSpansCountResponse>> {
146        Ok(self
147            .service_client
148            .lock()
149            .await
150            .get_spans_count(make_request_with_metadata(
151                GetSpansCountRequest {
152                    date_range,
153                    resolution: resolution.map(|r| r.try_into().unwrap()),
154                    filters,
155                },
156                &self.metadata_map,
157            ))
158            .await
159            .map_err(|status| {
160                SdkError::ApiError(SdkApiError {
161                    status,
162                    endpoint: "/com.coralogix.datausage.v2.DataUsageService/GetSpansCount".into(),
163                    feature_group: DATA_PLANS_FEATURE_GROUP_ID.into(),
164                })
165            })?
166            .into_inner())
167    }
168
169    /// Gets the log count for the provided filters.
170    ///
171    /// # Arguments
172    /// * `date_range` - The date range to get the log count for.
173    /// * `resolution` - The resolution to get the log count for.
174    /// * `filters` - The filters to apply.
175    /// * `subsystem_aggregation` - Whether to aggregate subsystems.
176    /// * `application_aggregation` - Whether to aggregate applications.
177    pub async fn get_logs_count(
178        &self,
179        date_range: Option<DateRange>,
180        resolution: Option<Duration>,
181        filters: Option<ScopesFilter>,
182        subsystem_aggregation: Option<bool>,
183        application_aggregation: Option<bool>,
184    ) -> Result<Streaming<GetLogsCountResponse>> {
185        Ok(self
186            .service_client
187            .lock()
188            .await
189            .get_logs_count(make_request_with_metadata(
190                GetLogsCountRequest {
191                    date_range,
192                    resolution: resolution.map(|r| r.try_into().unwrap()),
193                    filters,
194                    subsystem_aggregation,
195                    application_aggregation,
196                },
197                &self.metadata_map,
198            ))
199            .await
200            .map_err(|status| {
201                SdkError::ApiError(SdkApiError {
202                    status,
203                    endpoint: "/com.coralogix.datausage.v2.DataUsageService/GetLogsCount".into(),
204                    feature_group: DATA_PLANS_FEATURE_GROUP_ID.into(),
205                })
206            })?
207            .into_inner())
208    }
209
210    /// Updates the data usage metrics export status.
211    ///
212    /// # Arguments
213    /// * `enabled` - Whether to enable or disable the data usage metrics export.
214    pub async fn update_data_usage_metrics_export_status(
215        &self,
216        enabled: bool,
217    ) -> Result<UpdateDataUsageMetricsExportStatusResponse> {
218        Ok(self
219            .service_client
220            .lock()
221            .await
222            .update_data_usage_metrics_export_status(make_request_with_metadata(
223                UpdateDataUsageMetricsExportStatusRequest {
224                    enabled
225                },
226                &self.metadata_map,
227            ))
228            .await
229            .map_err(|status| {
230                SdkError::ApiError(SdkApiError {
231                    status,
232                    endpoint: "/com.coralogix.datausage.v2.DataUsageService/UpdateDataUsageMetricsExportStatus"
233                        .into(),
234                        feature_group: DATA_PLANS_FEATURE_GROUP_ID.into(),
235                })
236            })?
237            .into_inner())
238    }
239
240    /// Gets the data usage metrics export status.
241    pub async fn get_data_usage_metrics_export_status(
242        &self,
243    ) -> Result<GetDataUsageMetricsExportStatusResponse> {
244        Ok(self
245            .service_client
246            .lock()
247            .await
248            .get_data_usage_metrics_export_status(make_request_with_metadata(
249                GetDataUsageMetricsExportStatusRequest {},
250                &self.metadata_map,
251            ))
252            .await
253            .map_err(|status| {
254                SdkError::ApiError(SdkApiError {
255                    status,
256                    endpoint: "/com.coralogix.datausage.v2.DataUsageService/GetDataUsageMetricsExportStatus"
257                        .into(),
258                        feature_group: DATA_PLANS_FEATURE_GROUP_ID.into(),
259                })
260            })?
261            .into_inner())
262    }
263}
cx_sdk/client/data_usage.rs

cx_sdk/client/
data_usage.rs
