cx_sdk/client/
data_usage.rs

Help
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    pub async fn get_logs_count(
176        &self,
177        date_range: Option<DateRange>,
178        resolution: Option<Duration>,
179        filters: Option<ScopesFilter>,
180    ) -> Result<Streaming<GetLogsCountResponse>> {
181        Ok(self
182            .service_client
183            .lock()
184            .await
185            .get_logs_count(make_request_with_metadata(
186                GetLogsCountRequest {
187                    date_range,
188                    resolution: resolution.map(|r| r.try_into().unwrap()),
189                    filters,
190                },
191                &self.metadata_map,
192            ))
193            .await
194            .map_err(|status| {
195                SdkError::ApiError(SdkApiError {
196                    status,
197                    endpoint: "/com.coralogix.datausage.v2.DataUsageService/GetLogsCount".into(),
198                    feature_group: DATA_PLANS_FEATURE_GROUP_ID.into(),
199                })
200            })?
201            .into_inner())
202    }
203
204    /// Updates the data usage metrics export status.
205    ///
206    /// # Arguments
207    /// * `enabled` - Whether to enable or disable the data usage metrics export.
208    pub async fn update_data_usage_metrics_export_status(
209        &self,
210        enabled: bool,
211    ) -> Result<UpdateDataUsageMetricsExportStatusResponse> {
212        Ok(self
213            .service_client
214            .lock()
215            .await
216            .update_data_usage_metrics_export_status(make_request_with_metadata(
217                UpdateDataUsageMetricsExportStatusRequest {
218                    enabled
219                },
220                &self.metadata_map,
221            ))
222            .await
223            .map_err(|status| {
224                SdkError::ApiError(SdkApiError {
225                    status,
226                    endpoint: "/com.coralogix.datausage.v2.DataUsageService/UpdateDataUsageMetricsExportStatus"
227                        .into(),
228                        feature_group: DATA_PLANS_FEATURE_GROUP_ID.into(),
229                })
230            })?
231            .into_inner())
232    }
233
234    /// Gets the data usage metrics export status.
235    pub async fn get_data_usage_metrics_export_status(
236        &self,
237    ) -> Result<GetDataUsageMetricsExportStatusResponse> {
238        Ok(self
239            .service_client
240            .lock()
241            .await
242            .get_data_usage_metrics_export_status(make_request_with_metadata(
243                GetDataUsageMetricsExportStatusRequest {},
244                &self.metadata_map,
245            ))
246            .await
247            .map_err(|status| {
248                SdkError::ApiError(SdkApiError {
249                    status,
250                    endpoint: "/com.coralogix.datausage.v2.DataUsageService/GetDataUsageMetricsExportStatus"
251                        .into(),
252                        feature_group: DATA_PLANS_FEATURE_GROUP_ID.into(),
253                })
254            })?
255            .into_inner())
256    }
257}
cx_sdk/client/data_usage.rs

cx_sdk/client/
data_usage.rs
