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

cx_sdk/client/
data_usage.rs
