1use 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
74pub struct TcoPoliciesClient {
77    metadata_map: MetadataMap,
78    service_client: Mutex<PoliciesServiceClient<Channel>>,
79}
80
81impl TcoPoliciesClient {
82    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    #[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    #[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    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    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    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}