cx_sdk/client/
enrichments.rs1use 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
70pub struct EnrichmentsClient {
73 metadata_map: MetadataMap,
74 service_client: Mutex<EnrichmentServiceClient<Channel>>,
75}
76
77impl EnrichmentsClient {
78 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 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 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 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 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 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 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}