cx_sdk/client/
datasets.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::str::FromStr;
16
17use crate::{
18    auth::AuthContext,
19    error::{
20        Result,
21        SdkApiError,
22        SdkError,
23    },
24    metadata::CallProperties,
25    util::make_request_with_metadata,
26};
27
28pub use crate::com::coralogixapis::actions::v2::Action;
29
30use cx_api::proto::com::coralogix::enrichment::v1::{
31    CreateCustomEnrichmentRequest,
32    CreateCustomEnrichmentResponse,
33    DeleteCustomEnrichmentRequest,
34    DeleteCustomEnrichmentResponse,
35    File,
36    GetCustomEnrichmentRequest,
37    GetCustomEnrichmentResponse,
38    GetCustomEnrichmentsRequest,
39    GetCustomEnrichmentsResponse,
40    UpdateCustomEnrichmentRequest,
41    UpdateCustomEnrichmentResponse,
42    custom_enrichment_service_client::CustomEnrichmentServiceClient,
43    file::Content,
44};
45use tokio::sync::Mutex;
46use tonic::{
47    metadata::MetadataMap,
48    transport::{
49        Channel,
50        ClientTlsConfig,
51        Endpoint,
52    },
53};
54
55use crate::CoralogixRegion;
56
57pub use crate::com::coralogix::enrichment::v1::CustomEnrichment;
58
59const DATASETS_FEATURE_GROUP_ID: &str = "logs";
60
61/// The Custom Enrichments API client.
62/// Read more at <https://coralogix.com/docs/custom-enrichment-api/>
63#[deprecated(note = "This API will change significantly under a new name.")]
64pub struct DatasetClient {
65    metadata_map: MetadataMap,
66    service_client: Mutex<CustomEnrichmentServiceClient<Channel>>,
67}
68
69impl DatasetClient {
70    /// Creates a new client for the Custom Enrichments API.
71    ///
72    /// # Arguments
73    /// * `auth_context` - The [`AuthContext`] to use for authentication.
74    /// * `region` - The [`CoralogixRegion`] to connect to.
75    pub fn new(auth_context: AuthContext, region: CoralogixRegion) -> Result<Self> {
76        let channel: Channel = Endpoint::from_str(region.grpc_endpoint().as_str())?
77            .tls_config(ClientTlsConfig::new().with_native_roots())?
78            .connect_lazy();
79        let request_metadata: CallProperties = (&auth_context.team_level_api_key).into();
80        Ok(Self {
81            metadata_map: request_metadata.to_metadata_map(),
82            service_client: Mutex::new(CustomEnrichmentServiceClient::new(channel)),
83        })
84    }
85
86    /// Creates a new dataset.
87    ///
88    /// # Arguments
89    /// * `name` - The name of the Custom Enrichment.
90    /// * `description` - The description of the Custom Enrichment.
91    /// * `data` - The data of the Custom Enrichment.
92    pub async fn create(
93        &self,
94        name: String,
95        description: Option<String>,
96        data: Vec<u8>,
97    ) -> Result<CreateCustomEnrichmentResponse> {
98        let file = File {
99            name: Some(name.clone()),
100            size: Some(data.len() as u32),
101            content: Some(Content::Binary(data)),
102            extension: Some("csv".into()),
103        };
104        let request = make_request_with_metadata(
105            CreateCustomEnrichmentRequest {
106                file: Some(file),
107                name: Some(name.clone()),
108                description,
109            },
110            &self.metadata_map,
111        );
112        {
113            Ok(self
114                .service_client
115                .lock()
116                .await
117                .create_custom_enrichment(request)
118                .await
119                .map_err(
120                    |status| SdkError::ApiError(SdkApiError {
121                        status,
122                        endpoint: "/com.coralogixapis.enrichment.v1.CustomEnrichmentService/CreateCustomEnrichment".into(),
123                        feature_group: DATASETS_FEATURE_GROUP_ID.into(),
124                    }),
125                )?
126                .into_inner())
127        }
128    }
129
130    /// Replaces the existing dataset identified by its id.
131    ///
132    /// # Arguments
133    /// * `id` - The id of the Custom Enrichment.
134    /// * `name` - The name of the Custom Enrichment.
135    /// * `description` - The description of the Custom Enrichment.
136    /// * `data` - The data of the Custom Enrichment.
137    ///
138    /// Note that the existing dataset will be replaced with a new version of the data.
139    pub async fn replace(
140        &self,
141        id: u32,
142        name: String,
143        description: Option<String>,
144        data: Vec<u8>,
145    ) -> Result<UpdateCustomEnrichmentResponse> {
146        let file = File {
147            name: Some(name.clone()),
148            size: Some(data.len() as u32),
149            content: Some(Content::Binary(data)),
150            extension: Some("csv".into()),
151        };
152        let request = make_request_with_metadata(
153            UpdateCustomEnrichmentRequest {
154                custom_enrichment_id: Some(id),
155                file: Some(file),
156                name: Some(name.clone()),
157                description,
158            },
159            &self.metadata_map,
160        );
161        {
162            Ok(self
163                .service_client
164                .lock()
165                .await
166                .update_custom_enrichment(request)
167                .await
168                .map_err(
169                    |status| SdkError::ApiError(SdkApiError {
170                        status,
171                        endpoint: "/com.coralogixapis.enrichment.v1.CustomEnrichmentService/UpdateCustomEnrichment".into(),
172                        feature_group: DATASETS_FEATURE_GROUP_ID.into(),
173                    }),
174                )?
175                .into_inner())
176        }
177    }
178
179    /// Deletes the dataset identified by its id.
180    ///
181    /// # Arguments
182    /// * `custom_enrichment_id` - The id of the Custom Enrichment.
183    pub async fn delete(
184        &self,
185        custom_enrichment_id: u32,
186    ) -> Result<DeleteCustomEnrichmentResponse> {
187        Ok(self
188            .service_client
189            .lock()
190            .await
191            .delete_custom_enrichment(make_request_with_metadata(
192                DeleteCustomEnrichmentRequest {
193                    custom_enrichment_id: Some(custom_enrichment_id),
194                },
195                &self.metadata_map,
196            ))
197            .await
198            .map_err(
199                |status| SdkError::ApiError(SdkApiError {
200                    status,
201                    endpoint: "/com.coralogixapis.enrichment.v1.CustomEnrichmentService/DeleteCustomEnrichment".into(),
202                    feature_group: DATASETS_FEATURE_GROUP_ID.into(),
203                }),
204            )?
205            .into_inner())
206    }
207
208    /// Retrieves the dataset by id.
209    ///
210    /// # Arguments
211    /// * `id` - The id of the Custom Enrichment.
212    pub async fn get(&self, id: u32) -> Result<GetCustomEnrichmentResponse> {
213        Ok(self
214            .service_client
215            .lock()
216            .await
217            .get_custom_enrichment(make_request_with_metadata(
218                GetCustomEnrichmentRequest { id: Some(id) },
219                &self.metadata_map,
220            ))
221            .await
222            .map_err(
223                |status| SdkError::ApiError(SdkApiError {
224                    status,
225                    endpoint: "/com.coralogixapis.enrichment.v1.CustomEnrichmentService/GetCustomEnrichment".into(),
226                    feature_group: DATASETS_FEATURE_GROUP_ID.into(),
227                }),
228            )?
229            .into_inner())
230    }
231
232    /// Retrieves all datasets.
233    ///
234    /// # Returns
235    /// A list of all Custom Enrichments.
236    pub async fn list(&self) -> Result<GetCustomEnrichmentsResponse> {
237        Ok(self
238            .service_client
239            .lock()
240            .await
241            .get_custom_enrichments(make_request_with_metadata(
242                GetCustomEnrichmentsRequest {},
243                &self.metadata_map,
244            ))
245            .await
246            .map_err(
247                |status| SdkError::ApiError(SdkApiError {
248                    status,
249                    endpoint: "/com.coralogixapis.enrichment.v1.CustomEnrichmentService/GetCustomEnrichments".into(),
250                    feature_group: DATASETS_FEATURE_GROUP_ID.into(),
251                }),
252            )?
253            .into_inner())
254    }
255}