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

cx_sdk/client/
enrichments.rs
