cx_sdk/client/
tco_policies.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
28use cx_api::proto::com::coralogix::quota::v1::{
29    CreatePolicyRequest,
30    CreatePolicyResponse,
31    DeletePolicyRequest,
32    DeletePolicyResponse,
33    GetCompanyPoliciesRequest,
34    GetCompanyPoliciesResponse,
35    GetPolicyRequest,
36    GetPolicyResponse,
37    UpdatePolicyRequest,
38    UpdatePolicyResponse,
39    policies_service_client::PoliciesServiceClient,
40    update_policy_request,
41};
42use tokio::sync::Mutex;
43use tonic::{
44    metadata::MetadataMap,
45    transport::{
46        Channel,
47        ClientTlsConfig,
48        Endpoint,
49    },
50};
51
52pub use cx_api::proto::com::coralogix::quota::v1::{
53    ArchiveRetention,
54    LogRules,
55    Rule,
56    RuleTypeId,
57    SourceType,
58    SpanRules,
59    create_policy_request::SourceTypeRules,
60};
61
62fn convert_source_types(a: SourceTypeRules) -> update_policy_request::SourceTypeRules {
63    match a {
64        SourceTypeRules::LogRules(r) => update_policy_request::SourceTypeRules::LogRules(r),
65        SourceTypeRules::SpanRules(r) => update_policy_request::SourceTypeRules::SpanRules(r),
66    }
67}
68
69use crate::CoralogixRegion;
70
71const TCO_FEATURE_GROUP_ID: &str = "tco";
72
73/// The TCO client.
74/// Read more at <https://coralogix.com/docs/tco-tracing-policy-grpc-api/>
75pub struct TcoPoliciesClient {
76    metadata_map: MetadataMap,
77    service_client: Mutex<PoliciesServiceClient<Channel>>,
78}
79
80impl TcoPoliciesClient {
81    /// Creates a new client for the TCO.
82    ///
83    /// # Arguments
84    /// * `auth_context` - The [`AuthContext`] to use for authentication.
85    /// * `region` - The [`CoralogixRegion`] to connect to.
86    pub fn new(auth_context: AuthContext, region: CoralogixRegion) -> Result<Self> {
87        let channel: Channel = Endpoint::from_str(region.grpc_endpoint().as_str())?
88            .tls_config(ClientTlsConfig::new().with_native_roots())?
89            .connect_lazy();
90        let request_metadata: CallProperties = (&auth_context.team_level_api_key).into();
91        Ok(Self {
92            metadata_map: request_metadata.to_metadata_map(),
93            service_client: Mutex::new(PoliciesServiceClient::new(channel)),
94        })
95    }
96
97    /// Creates a new TCO policy.
98    ///
99    /// # Arguments
100    /// * `name` - The name of the policy.
101    /// * `description` - The description of the policy.
102    /// * `priority` - The priority of the policy.
103    /// * `application_rule` - The application [`Rule`] of the policy.
104    /// * `subsystem_rule` - The subsystem [`Rule`] of the policy.
105    /// * `archive_retention` - The [`ArchiveRetention`] of the policy.
106    /// * `source_type_rules` - The [`SourceTypeRules`] of the policy.
107    #[allow(clippy::too_many_arguments)]
108    pub async fn create(
109        &self,
110        name: Option<String>,
111        description: Option<String>,
112        priority: i32,
113        application_rule: Option<Rule>,
114        subsystem_rule: Option<Rule>,
115        archive_retention: Option<ArchiveRetention>,
116        source_type_rules: Option<SourceTypeRules>,
117    ) -> Result<CreatePolicyResponse> {
118        let request = make_request_with_metadata(
119            CreatePolicyRequest {
120                name,
121                description,
122                priority,
123                application_rule,
124                subsystem_rule,
125                archive_retention,
126                source_type_rules,
127            },
128            &self.metadata_map,
129        );
130        self.service_client
131            .lock()
132            .await
133            .create_policy(request)
134            .await
135            .map(|r| r.into_inner())
136            .map_err(|status| {
137                SdkError::ApiError(SdkApiError {
138                    status,
139                    endpoint: "/com.coralogixapis.quota.v1.PoliciesService/CreatePolicy".into(),
140                    feature_group: TCO_FEATURE_GROUP_ID.into(),
141                })
142            })
143    }
144
145    /// Creates a new TCO policy.
146    ///
147    /// # Arguments
148    /// * `id` - The id of the policy.
149    /// * `name` - The name of the policy.
150    /// * `description` - The description of the policy.
151    /// * `priority` - The priority of the policy.
152    /// * `application_rule` - The application [`Rule`] of the policy.
153    /// * `subsystem_rule` - The subsystem [`Rule`] of the policy.
154    /// * `archive_retention` - The [`ArchiveRetention`] of the policy.
155    /// * `source_type_rules` - The [`SourceTypeRules`] of the policy.
156    #[allow(clippy::too_many_arguments)]
157    pub async fn update(
158        &self,
159        id: String,
160        name: Option<String>,
161        description: Option<String>,
162        priority: i32,
163        application_rule: Option<Rule>,
164        subsystem_rule: Option<Rule>,
165        archive_retention: Option<ArchiveRetention>,
166        enabled: Option<bool>,
167        source_type_rules: Option<SourceTypeRules>,
168    ) -> Result<UpdatePolicyResponse> {
169        let request = make_request_with_metadata(
170            UpdatePolicyRequest {
171                id: Some(id),
172                name,
173                description,
174                priority,
175                application_rule,
176                subsystem_rule,
177                archive_retention,
178                enabled,
179                source_type_rules: source_type_rules.map(convert_source_types),
180            },
181            &self.metadata_map,
182        );
183        self.service_client
184            .lock()
185            .await
186            .update_policy(request)
187            .await
188            .map(|r| r.into_inner())
189            .map_err(|status| {
190                SdkError::ApiError(SdkApiError {
191                    status,
192                    endpoint: "/com.coralogixapis.quota.v1.PoliciesService/UpdatePolicy".into(),
193                    feature_group: TCO_FEATURE_GROUP_ID.into(),
194                })
195            })
196    }
197
198    /// Deletes a TCO policy.
199    ///
200    /// # Arguments
201    /// * `id` - The id of the policy to delete.
202    pub async fn delete(&self, id: String) -> Result<DeletePolicyResponse> {
203        let request =
204            make_request_with_metadata(DeletePolicyRequest { id: Some(id) }, &self.metadata_map);
205        self.service_client
206            .lock()
207            .await
208            .delete_policy(request)
209            .await
210            .map(|r| r.into_inner())
211            .map_err(|status| {
212                SdkError::ApiError(SdkApiError {
213                    status,
214                    endpoint: "/com.coralogixapis.quota.v1.PoliciesService/DeletePolicy".into(),
215                    feature_group: TCO_FEATURE_GROUP_ID.into(),
216                })
217            })
218    }
219
220    /// Retrieves a TCO policy.
221    ///
222    /// # Arguments
223    /// * `id` - The id of the policy to retrieve.
224    pub async fn get(&self, id: String) -> Result<GetPolicyResponse> {
225        let request =
226            make_request_with_metadata(GetPolicyRequest { id: Some(id) }, &self.metadata_map);
227
228        self.service_client
229            .lock()
230            .await
231            .get_policy(request)
232            .await
233            .map(|r| r.into_inner())
234            .map_err(|status| {
235                SdkError::ApiError(SdkApiError {
236                    status,
237                    endpoint: "/com.coralogixapis.quota.v1.PoliciesService/GetPolicy".into(),
238                    feature_group: TCO_FEATURE_GROUP_ID.into(),
239                })
240            })
241    }
242
243    /// Retrieves a list of TCO policies.
244    ///
245    /// # Arguments
246    /// * `source_type` - The [`SourceType`] of the policies to retrieve.
247    /// * `enabled_only` - Whether to retrieve only enabled policies.
248    pub async fn list(
249        &self,
250        source_type: SourceType,
251        enabled_only: bool,
252    ) -> Result<GetCompanyPoliciesResponse> {
253        let request = make_request_with_metadata(
254            GetCompanyPoliciesRequest {
255                enabled_only: Some(enabled_only),
256                source_type: Some(source_type.into()),
257            },
258            &self.metadata_map,
259        );
260
261        self.service_client
262            .lock()
263            .await
264            .get_company_policies(request)
265            .await
266            .map(|r| r.into_inner())
267            .map_err(|status| {
268                SdkError::ApiError(SdkApiError {
269                    status,
270                    endpoint: "/com.coralogixapis.quota.v1.PoliciesService/GetCompanyPolicies"
271                        .into(),
272                    feature_group: TCO_FEATURE_GROUP_ID.into(),
273                })
274            })
275    }
276}