cx_sdk/client/
enrichments.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::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    AddEnrichmentsRequest,
32    AddEnrichmentsResponse,
33    AtomicOverwriteEnrichmentsRequest,
34    AtomicOverwriteEnrichmentsResponse,
35    GetCompanyEnrichmentSettingsRequest,
36    GetCompanyEnrichmentSettingsResponse,
37    GetEnrichmentLimitRequest,
38    GetEnrichmentLimitResponse,
39    GetEnrichmentsRequest,
40    GetEnrichmentsResponse,
41    RemoveEnrichmentsRequest,
42    RemoveEnrichmentsResponse,
43    enrichment_service_client::EnrichmentServiceClient,
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
57const ENRICHMENTS_FEATURE_GROUP_ID: &str = "enrichments";
58
59pub use crate::com::coralogix::enrichment::v1::{
60    AwsType,
61    CustomEnrichmentType,
62    EnrichmentFieldDefinition,
63    EnrichmentRequestModel as EnrichmentMapping,
64    EnrichmentType,
65    GeoIpType,
66    SuspiciousIpType,
67    enrichment_type::Type,
68};
69
70/// The Custom Enrichments API client.
71/// Read more at <https://coralogix.com/docs/custom-enrichment-api/>
72pub struct EnrichmentsClient {
73    metadata_map: MetadataMap,
74    service_client: Mutex<EnrichmentServiceClient<Channel>>,
75}
76
77impl EnrichmentsClient {
78    /// Creates a new client for the Enrichments API.
79    ///
80    /// # Arguments
81    /// * `auth_context` - The [`AuthContext`] to use for authentication.
82    /// * `region` - The [`CoralogixRegion`] to connect to.
83    pub fn new(auth_context: AuthContext, region: CoralogixRegion) -> Result<Self> {
84        let channel: Channel = Endpoint::from_str(region.grpc_endpoint().as_str())?
85            .tls_config(ClientTlsConfig::new().with_native_roots())?
86            .connect_lazy();
87        let request_metadata: CallProperties = (&auth_context.team_level_api_key).into();
88        Ok(Self {
89            metadata_map: request_metadata.to_metadata_map(),
90            service_client: Mutex::new(EnrichmentServiceClient::new(channel)),
91        })
92    }
93
94    /// Adds enrichment mappings.
95    ///
96    /// # Arguments
97    /// * `enrichment_mapping` - The enrichment definitions to create.
98    pub async fn add(
99        &self,
100        enrichment_mapping: Vec<EnrichmentMapping>,
101    ) -> Result<AddEnrichmentsResponse> {
102        self.service_client
103            .lock()
104            .await
105            .add_enrichments(make_request_with_metadata(
106                AddEnrichmentsRequest {
107                    request_enrichments: enrichment_mapping,
108                },
109                &self.metadata_map,
110            ))
111            .await
112            .map(|r| r.into_inner())
113            .map_err(|status| {
114                SdkError::ApiError(SdkApiError {
115                    status,
116                    endpoint: "/com.coralogixapis.enrichment.v1.EnrichmentService/AddEnrichments"
117                        .into(),
118                    feature_group: ENRICHMENTS_FEATURE_GROUP_ID.into(),
119                })
120            })
121    }
122
123    /// Updates the mappings for an enrichment.
124    ///
125    /// # Arguments
126    /// * `enrichment_type` - The type of enrichment to update.
127    /// * `field_mappings` - The associated field mappings.
128    pub async fn update(
129        &self,
130        enrichment_mapping: Vec<EnrichmentMapping>,
131    ) -> Result<AtomicOverwriteEnrichmentsResponse> {
132        Ok(self
133            .service_client
134            .lock()
135            .await
136            .atomic_overwrite_enrichments(make_request_with_metadata(
137                AtomicOverwriteEnrichmentsRequest {
138                    request_enrichments: enrichment_mapping,
139                    ..Default::default()
140                },
141                &self.metadata_map,
142            ))
143            .await
144            .map_err(
145                |status| SdkError::ApiError(SdkApiError {
146                    status,
147                    endpoint: "/com.coralogixapis.enrichment.v1.EnrichmentService/AtomicOverwriteEnrichments".into(),
148                    feature_group: ENRICHMENTS_FEATURE_GROUP_ID.into(),
149                },
150            ))?
151            .into_inner())
152    }
153
154    /// Removes all enrichment mappings from the provided id.
155    ///
156    /// # Arguments
157    /// * `ids` - The ids to remove mappings from.
158    pub async fn delete(&self, enrichment_ids: Vec<u32>) -> Result<RemoveEnrichmentsResponse> {
159        Ok(self
160            .service_client
161            .lock()
162            .await
163            .remove_enrichments(make_request_with_metadata(
164                RemoveEnrichmentsRequest { enrichment_ids },
165                &self.metadata_map,
166            ))
167            .await
168            .map_err(|status| {
169                SdkError::ApiError(SdkApiError {
170                    status,
171                    endpoint:
172                        "/com.coralogixapis.enrichment.v1.EnrichmentService/RemoveEnrichments"
173                            .into(),
174                    feature_group: ENRICHMENTS_FEATURE_GROUP_ID.into(),
175                })
176            })?
177            .into_inner())
178    }
179
180    /// Retrieves the all enrichments.
181    pub async fn list(&self) -> Result<GetEnrichmentsResponse> {
182        Ok(self
183            .service_client
184            .lock()
185            .await
186            .get_enrichments(make_request_with_metadata(
187                GetEnrichmentsRequest {},
188                &self.metadata_map,
189            ))
190            .await
191            .map_err(|status| {
192                SdkError::ApiError(SdkApiError {
193                    status,
194                    endpoint: "/com.coralogixapis.enrichment.v1.EnrichmentService/GetEnrichments"
195                        .into(),
196                    feature_group: ENRICHMENTS_FEATURE_GROUP_ID.into(),
197                })
198            })?
199            .into_inner())
200    }
201
202    /// Retrieves the Enrichment limits.
203    pub async fn get_limits(&self) -> Result<GetEnrichmentLimitResponse> {
204        Ok(self
205            .service_client
206            .lock()
207            .await
208            .get_enrichment_limit(make_request_with_metadata(
209                GetEnrichmentLimitRequest {},
210                &self.metadata_map,
211            ))
212            .await
213            .map_err(|status| {
214                SdkError::ApiError(SdkApiError {
215                    status,
216                    endpoint:
217                        "/com.coralogixapis.enrichment.v1.EnrichmentService/GetEnrichmentLimit"
218                            .into(),
219                    feature_group: ENRICHMENTS_FEATURE_GROUP_ID.into(),
220                })
221            })?
222            .into_inner())
223    }
224
225    /// Retrieves the Enrichment settings for a company.
226    pub async fn get_enrichment_settings(&self) -> Result<GetCompanyEnrichmentSettingsResponse> {
227        Ok(self
228            .service_client
229            .lock()
230            .await
231            .get_company_enrichment_settings(make_request_with_metadata(
232                GetCompanyEnrichmentSettingsRequest {},
233                &self.metadata_map,
234            ))
235            .await
236            .map_err(|status| {
237                SdkError::ApiError(SdkApiError {
238                    status,
239                    endpoint:
240                        "/com.coralogixapis.enrichment.v1.EnrichmentService/GetCompanyEnrichmentSettings"
241                            .into(),
242                    feature_group: ENRICHMENTS_FEATURE_GROUP_ID.into(),
243                })
244            })?
245            .into_inner())
246    }
247}
cx_sdk/client/enrichments.rs

cx_sdk/client/
enrichments.rs
