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