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_type: EnrichmentType,
131        field_mappings: Vec<EnrichmentFieldDefinition>,
132    ) -> Result<AtomicOverwriteEnrichmentsResponse> {
133        Ok(self
134            .service_client
135            .lock()
136            .await
137            .atomic_overwrite_enrichments(make_request_with_metadata(
138                AtomicOverwriteEnrichmentsRequest {
139                    enrichment_type: Some(enrichment_type),
140                    enrichment_fields: field_mappings,
141                },
142                &self.metadata_map,
143            ))
144            .await
145            .map_err(
146                |status| SdkError::ApiError(SdkApiError {
147                    status,
148                    endpoint: "/com.coralogixapis.enrichment.v1.EnrichmentService/AtomicOverwriteEnrichments".into(),
149                    feature_group: ENRICHMENTS_FEATURE_GROUP_ID.into(),
150                },
151            ))?
152            .into_inner())
153    }
154
155    /// Removes all enrichment mappings from the provided id.
156    ///
157    /// # Arguments
158    /// * `ids` - The ids to remove mappings from.
159    pub async fn delete(&self, enrichment_ids: Vec<u32>) -> Result<RemoveEnrichmentsResponse> {
160        Ok(self
161            .service_client
162            .lock()
163            .await
164            .remove_enrichments(make_request_with_metadata(
165                RemoveEnrichmentsRequest { enrichment_ids },
166                &self.metadata_map,
167            ))
168            .await
169            .map_err(|status| {
170                SdkError::ApiError(SdkApiError {
171                    status,
172                    endpoint:
173                        "/com.coralogixapis.enrichment.v1.EnrichmentService/RemoveEnrichments"
174                            .into(),
175                    feature_group: ENRICHMENTS_FEATURE_GROUP_ID.into(),
176                })
177            })?
178            .into_inner())
179    }
180
181    /// Retrieves the all enrichments.
182    pub async fn list(&self) -> Result<GetEnrichmentsResponse> {
183        Ok(self
184            .service_client
185            .lock()
186            .await
187            .get_enrichments(make_request_with_metadata(
188                GetEnrichmentsRequest {},
189                &self.metadata_map,
190            ))
191            .await
192            .map_err(|status| {
193                SdkError::ApiError(SdkApiError {
194                    status,
195                    endpoint: "/com.coralogixapis.enrichment.v1.EnrichmentService/GetEnrichments"
196                        .into(),
197                    feature_group: ENRICHMENTS_FEATURE_GROUP_ID.into(),
198                })
199            })?
200            .into_inner())
201    }
202
203    /// Retrieves the Enrichment limits.
204    pub async fn get_limits(&self) -> Result<GetEnrichmentLimitResponse> {
205        Ok(self
206            .service_client
207            .lock()
208            .await
209            .get_enrichment_limit(make_request_with_metadata(
210                GetEnrichmentLimitRequest {},
211                &self.metadata_map,
212            ))
213            .await
214            .map_err(|status| {
215                SdkError::ApiError(SdkApiError {
216                    status,
217                    endpoint:
218                        "/com.coralogixapis.enrichment.v1.EnrichmentService/GetEnrichmentLimit"
219                            .into(),
220                    feature_group: ENRICHMENTS_FEATURE_GROUP_ID.into(),
221                })
222            })?
223            .into_inner())
224    }
225
226    /// Retrieves the Enrichment settings for a company.
227    pub async fn get_enrichment_settings(&self) -> Result<GetCompanyEnrichmentSettingsResponse> {
228        Ok(self
229            .service_client
230            .lock()
231            .await
232            .get_company_enrichment_settings(make_request_with_metadata(
233                GetCompanyEnrichmentSettingsRequest {},
234                &self.metadata_map,
235            ))
236            .await
237            .map_err(|status| {
238                SdkError::ApiError(SdkApiError {
239                    status,
240                    endpoint:
241                        "/com.coralogixapis.enrichment.v1.EnrichmentService/GetCompanyEnrichmentSettings"
242                            .into(),
243                    feature_group: ENRICHMENTS_FEATURE_GROUP_ID.into(),
244                })
245            })?
246            .into_inner())
247    }
248}
cx_sdk/client/enrichments.rs

cx_sdk/client/
enrichments.rs
