1use itertools::Itertools;
2
3use unitycatalog_common::models::policies::v1::policy_info::Function;
4use unitycatalog_common::models::policies::v1::*;
5use unitycatalog_common::models::{
6 ObjectLabel, Resource, ResourceIdent, ResourceName, ResourceRef,
7};
8
9use super::{RequestContext, SecuredAction};
10use crate::Result;
11pub use crate::codegen::policies::PolicyHandler;
12use crate::policy::{Permission, Policy, process_resources};
13use crate::store::ResourceStore;
14
15fn securable_chain(
20 on_securable_type: &str,
21 on_securable_fullname: &str,
22) -> Vec<(&'static str, String)> {
23 let segments: Vec<&str> = on_securable_fullname.split('.').collect();
24 let levels: &[&str] = match on_securable_type {
25 "tables" => &["tables", "schemas", "catalogs"],
26 "schemas" => &["schemas", "catalogs"],
27 "catalogs" => &["catalogs"],
28 _ => &[],
29 };
30 levels
31 .iter()
32 .enumerate()
33 .filter_map(|(i, level)| {
34 let take = segments.len().checked_sub(i)?;
35 (take > 0).then(|| (*level, segments[..take].join(".")))
36 })
37 .collect()
38}
39
40fn matches_securable(
41 policy: &PolicyInfo,
42 on_securable_type: &str,
43 on_securable_fullname: &str,
44) -> bool {
45 policy.on_securable_type == on_securable_type
46 && policy.on_securable_fullname == on_securable_fullname
47}
48
49fn validate_policy_shape(policy: &PolicyInfo) -> Result<()> {
50 let policy_type = policy
51 .policy_type
52 .as_known()
53 .ok_or_else(|| crate::Error::invalid_argument("unrecognized policy_type"))?;
54 match policy_type {
55 PolicyType::POLICY_TYPE_COLUMN_MASK if policy.match_columns.is_empty() => {
56 Err(crate::Error::invalid_argument(
57 "column mask policies require at least one match_columns entry",
58 ))
59 }
60 PolicyType::POLICY_TYPE_ROW_FILTER
61 if policy
62 .function
63 .as_ref()
64 .and_then(|f| match f {
65 Function::RowFilter(r) => Some(r),
66 _ => None,
67 })
68 .is_none_or(|r| r.function_name.is_empty()) =>
69 {
70 Err(crate::Error::invalid_argument(
71 "row filter policies require row_filter.function_name",
72 ))
73 }
74 PolicyType::POLICY_TYPE_UNSPECIFIED => Err(crate::Error::invalid_argument(
75 "policy_type must be POLICY_TYPE_ROW_FILTER or POLICY_TYPE_COLUMN_MASK",
76 )),
77 _ => Ok(()),
78 }
79}
80
81#[async_trait::async_trait]
82impl<T: ResourceStore + Policy<RequestContext>> PolicyHandler<RequestContext> for T {
83 #[tracing::instrument(skip(self, context), fields(resource_name))]
84 async fn create_policy(
85 &self,
86 request: CreatePolicyRequest,
87 context: RequestContext,
88 ) -> Result<PolicyInfo> {
89 self.check_required(&request, &context).await?;
90 let mut policy = request
91 .policy_info
92 .ok_or_else(|| crate::Error::invalid_argument("policy_info must be provided"))?;
93 policy.on_securable_type = request.on_securable_type;
94 policy.on_securable_fullname = request.on_securable_fullname;
95 validate_policy_shape(&policy)?;
96 tracing::Span::current().record("resource_name", &policy.name);
97
98 let existing = self
100 .list(&ObjectLabel::PolicyInfo, None, None, None)
101 .await?
102 .0;
103 let clashes = existing.into_iter().any(|r| {
104 let Ok(p): std::result::Result<PolicyInfo, _> = r.try_into() else {
105 return false;
106 };
107 matches_securable(&p, &policy.on_securable_type, &policy.on_securable_fullname)
108 && p.name == policy.name
109 });
110 if clashes {
111 return Err(crate::Error::invalid_argument(format!(
112 "policy '{}' already exists on {} '{}'",
113 policy.name, policy.on_securable_type, policy.on_securable_fullname
114 )));
115 }
116
117 Ok(self.create(policy.into()).await?.0.try_into()?)
118 }
119
120 #[tracing::instrument(skip(self, context), fields(resource_name))]
121 async fn delete_policy(
122 &self,
123 request: DeletePolicyRequest,
124 context: RequestContext,
125 ) -> Result<()> {
126 tracing::Span::current().record("resource_name", &request.name);
127 self.check_required(&request, &context).await?;
128 let ident = find_policy_ident(
129 self,
130 &request.on_securable_type,
131 &request.on_securable_fullname,
132 &request.name,
133 )
134 .await?;
135 Ok(self.delete(&ident).await?)
136 }
137
138 #[tracing::instrument(skip(self, context), fields(resource_name))]
139 async fn get_policy(
140 &self,
141 request: GetPolicyRequest,
142 context: RequestContext,
143 ) -> Result<PolicyInfo> {
144 tracing::Span::current().record("resource_name", &request.name);
145 self.check_required(&request, &context).await?;
146 let ident = find_policy_ident(
147 self,
148 &request.on_securable_type,
149 &request.on_securable_fullname,
150 &request.name,
151 )
152 .await?;
153 Ok(self.get(&ident).await?.0.try_into()?)
154 }
155
156 #[tracing::instrument(skip(self, context))]
157 async fn list_policies(
158 &self,
159 request: ListPoliciesRequest,
160 context: RequestContext,
161 ) -> Result<ListPoliciesResponse> {
162 self.check_required(&request, &context).await?;
163
164 let (all, _) = self
165 .list(&ObjectLabel::PolicyInfo, None, None, None)
166 .await?;
167 let mut policies: Vec<PolicyInfo> = all.into_iter().map(|r| r.try_into()).try_collect()?;
168
169 if request.include_inherited.unwrap_or(false) {
170 let chain = securable_chain(&request.on_securable_type, &request.on_securable_fullname);
171 policies.retain(|p| chain.iter().any(|(t, name)| matches_securable(p, t, name)));
172 } else {
173 policies.retain(|p| {
174 matches_securable(
175 p,
176 &request.on_securable_type,
177 &request.on_securable_fullname,
178 )
179 });
180 }
181
182 let mut resources: Vec<Resource> = policies.into_iter().map(Resource::PolicyInfo).collect();
183 process_resources(self, &context, &Permission::Read, &mut resources).await?;
184
185 let max_results = request
186 .max_results
187 .map(|v| v as usize)
188 .unwrap_or(resources.len());
189 let offset = request
190 .page_token
191 .as_deref()
192 .and_then(|t| t.parse::<usize>().ok())
193 .unwrap_or(0);
194 let next_page_token =
195 (offset + max_results < resources.len()).then(|| (offset + max_results).to_string());
196 let page = resources
197 .into_iter()
198 .skip(offset)
199 .take(max_results)
200 .map(|r| r.try_into())
201 .try_collect()?;
202
203 Ok(ListPoliciesResponse {
204 policies: page,
205 next_page_token,
206 ..Default::default()
207 })
208 }
209
210 #[tracing::instrument(skip(self, context), fields(resource_name))]
211 async fn update_policy(
212 &self,
213 request: UpdatePolicyRequest,
214 context: RequestContext,
215 ) -> Result<PolicyInfo> {
216 tracing::Span::current().record("resource_name", &request.name);
217 self.check_required(&request, &context).await?;
218 let ident = find_policy_ident(
219 self,
220 &request.on_securable_type,
221 &request.on_securable_fullname,
222 &request.name,
223 )
224 .await?;
225 let mut policy = request
226 .policy_info
227 .ok_or_else(|| crate::Error::invalid_argument("policy_info must be provided"))?;
228 policy.name = request.name;
230 policy.on_securable_type = request.on_securable_type;
231 policy.on_securable_fullname = request.on_securable_fullname;
232 validate_policy_shape(&policy)?;
233 Ok(self.update(&ident, policy.into()).await?.0.try_into()?)
234 }
235}
236
237async fn find_policy_ident<S: ResourceStore>(
242 store: &S,
243 on_securable_type: &str,
244 on_securable_fullname: &str,
245 name: &str,
246) -> Result<ResourceIdent> {
247 let (all, _) = store
248 .list(&ObjectLabel::PolicyInfo, None, None, None)
249 .await?;
250 all.into_iter()
251 .find_map(|r| {
252 let p: PolicyInfo = r.try_into().ok()?;
253 (matches_securable(&p, on_securable_type, on_securable_fullname) && p.name == name)
254 .then(|| ResourceIdent::policy_info(ResourceName::new([name])))
255 })
256 .ok_or(crate::Error::NotFound)
257}
258
259impl SecuredAction for CreatePolicyRequest {
260 fn resource(&self) -> ResourceIdent {
261 securable_ident(&self.on_securable_type, &self.on_securable_fullname)
262 }
263
264 fn permission(&self) -> &'static Permission {
265 &Permission::Manage
266 }
267}
268
269impl SecuredAction for ListPoliciesRequest {
270 fn resource(&self) -> ResourceIdent {
271 securable_ident(&self.on_securable_type, &self.on_securable_fullname)
272 }
273
274 fn permission(&self) -> &'static Permission {
275 &Permission::Read
276 }
277}
278
279impl SecuredAction for GetPolicyRequest {
280 fn resource(&self) -> ResourceIdent {
281 securable_ident(&self.on_securable_type, &self.on_securable_fullname)
282 }
283
284 fn permission(&self) -> &'static Permission {
285 &Permission::Read
286 }
287}
288
289impl SecuredAction for UpdatePolicyRequest {
290 fn resource(&self) -> ResourceIdent {
291 securable_ident(&self.on_securable_type, &self.on_securable_fullname)
292 }
293
294 fn permission(&self) -> &'static Permission {
295 &Permission::Manage
296 }
297}
298
299impl SecuredAction for DeletePolicyRequest {
300 fn resource(&self) -> ResourceIdent {
301 securable_ident(&self.on_securable_type, &self.on_securable_fullname)
302 }
303
304 fn permission(&self) -> &'static Permission {
305 &Permission::Manage
306 }
307}
308
309fn securable_ident(on_securable_type: &str, on_securable_fullname: &str) -> ResourceIdent {
313 let name = ResourceName::from_naive_str_split(on_securable_fullname);
314 match on_securable_type {
315 "catalogs" => ResourceIdent::catalog(name),
316 "schemas" => ResourceIdent::schema(name),
317 "tables" => ResourceIdent::table(name),
318 _ => ResourceIdent::catalog(ResourceRef::Undefined),
319 }
320}
321
322#[cfg(test)]
323mod tests {
324 use std::sync::Arc;
325
326 use unitycatalog_common::models::catalogs::v1::Catalog;
327 use unitycatalog_common::models::schemas::v1::Schema;
328 use unitycatalog_common::models::tables::v1::{DataSourceFormat, Table, TableType};
329 use unitycatalog_common::services::encryption::{EnvelopeEncryptor, LocalKeyProvider};
330
331 use super::*;
332 use crate::memory::InMemoryResourceStore;
333 use crate::policy::ConstantPolicy;
334 use crate::services::ServerHandler;
335
336 async fn handler() -> ServerHandler<RequestContext> {
337 let encryptor =
338 EnvelopeEncryptor::local(LocalKeyProvider::single("test", vec![0x42; 32]).unwrap());
339 let store = Arc::new(InMemoryResourceStore::new(encryptor));
340 store
341 .create(
342 Catalog {
343 name: "cat".to_string(),
344 ..Default::default()
345 }
346 .into(),
347 )
348 .await
349 .unwrap();
350 store
351 .create(
352 Schema {
353 name: "sch".to_string(),
354 catalog_name: "cat".to_string(),
355 ..Default::default()
356 }
357 .into(),
358 )
359 .await
360 .unwrap();
361 store
362 .create(
363 Table {
364 name: "tbl".to_string(),
365 catalog_name: "cat".to_string(),
366 schema_name: "sch".to_string(),
367 full_name: "cat.sch.tbl".to_string(),
368 table_type: TableType::Managed.into(),
369 data_source_format: DataSourceFormat::Delta.into(),
370 ..Default::default()
371 }
372 .into(),
373 )
374 .await
375 .unwrap();
376 let policy: Arc<dyn Policy<RequestContext>> = Arc::new(ConstantPolicy::default());
377 ServerHandler::try_new_tokio(policy, store).unwrap()
378 }
379
380 fn ctx() -> RequestContext {
381 RequestContext {
382 recipient: crate::policy::Principal::anonymous(),
383 }
384 }
385
386 fn row_filter_policy(
387 name: &str,
388 on_securable_type: &str,
389 on_securable_fullname: &str,
390 ) -> PolicyInfo {
391 PolicyInfo {
392 name: name.to_string(),
393 on_securable_type: on_securable_type.to_string(),
394 on_securable_fullname: on_securable_fullname.to_string(),
395 policy_type: PolicyType::RowFilter.into(),
396 to_principals: vec!["group:analysts".to_string()],
397 function: Some(Function::RowFilter(Box::new(FunctionRef {
398 function_name: "cat.sch.my_filter".to_string(),
399 using: vec![],
400 ..Default::default()
401 }))),
402 ..Default::default()
403 }
404 }
405
406 #[tokio::test]
407 async fn policy_crud_round_trip() {
408 let h = handler().await;
409
410 let created = h
411 .create_policy(
412 CreatePolicyRequest {
413 on_securable_type: "tables".to_string(),
414 on_securable_fullname: "cat.sch.tbl".to_string(),
415 policy_info: Some(row_filter_policy("p1", "", "")).into(),
416 ..Default::default()
417 },
418 ctx(),
419 )
420 .await
421 .unwrap();
422 assert_eq!(created.name, "p1");
423 assert_eq!(created.on_securable_fullname, "cat.sch.tbl");
424
425 let fetched = h
426 .get_policy(
427 GetPolicyRequest {
428 on_securable_type: "tables".to_string(),
429 on_securable_fullname: "cat.sch.tbl".to_string(),
430 name: "p1".to_string(),
431 ..Default::default()
432 },
433 ctx(),
434 )
435 .await
436 .unwrap();
437 assert_eq!(fetched.name, "p1");
438
439 let listed = h
440 .list_policies(
441 ListPoliciesRequest {
442 on_securable_type: "tables".to_string(),
443 on_securable_fullname: "cat.sch.tbl".to_string(),
444 include_inherited: None,
445 max_results: None,
446 page_token: None,
447 ..Default::default()
448 },
449 ctx(),
450 )
451 .await
452 .unwrap();
453 assert_eq!(listed.policies.len(), 1);
454
455 let mut updated_policy = row_filter_policy("p1", "", "");
456 updated_policy.comment = Some("updated".to_string());
457 let updated = h
458 .update_policy(
459 UpdatePolicyRequest {
460 on_securable_type: "tables".to_string(),
461 on_securable_fullname: "cat.sch.tbl".to_string(),
462 name: "p1".to_string(),
463 policy_info: Some(updated_policy).into(),
464 update_mask: None,
465 ..Default::default()
466 },
467 ctx(),
468 )
469 .await
470 .unwrap();
471 assert_eq!(updated.comment.as_deref(), Some("updated"));
472
473 h.delete_policy(
474 DeletePolicyRequest {
475 on_securable_type: "tables".to_string(),
476 on_securable_fullname: "cat.sch.tbl".to_string(),
477 name: "p1".to_string(),
478 ..Default::default()
479 },
480 ctx(),
481 )
482 .await
483 .unwrap();
484
485 let missing = h
486 .get_policy(
487 GetPolicyRequest {
488 on_securable_type: "tables".to_string(),
489 on_securable_fullname: "cat.sch.tbl".to_string(),
490 name: "p1".to_string(),
491 ..Default::default()
492 },
493 ctx(),
494 )
495 .await;
496 assert!(missing.is_err(), "expected NotFound after delete");
497 }
498
499 #[tokio::test]
500 async fn list_include_inherited_unions_ancestor_chain() {
501 let h = handler().await;
502
503 h.create_policy(
504 CreatePolicyRequest {
505 on_securable_type: "catalogs".to_string(),
506 on_securable_fullname: "cat".to_string(),
507 policy_info: Some(row_filter_policy("catalog_policy", "", "")).into(),
508 ..Default::default()
509 },
510 ctx(),
511 )
512 .await
513 .unwrap();
514 h.create_policy(
515 CreatePolicyRequest {
516 on_securable_type: "schemas".to_string(),
517 on_securable_fullname: "cat.sch".to_string(),
518 policy_info: Some(row_filter_policy("schema_policy", "", "")).into(),
519 ..Default::default()
520 },
521 ctx(),
522 )
523 .await
524 .unwrap();
525 h.create_policy(
526 CreatePolicyRequest {
527 on_securable_type: "tables".to_string(),
528 on_securable_fullname: "cat.sch.tbl".to_string(),
529 policy_info: Some(row_filter_policy("table_policy", "", "")).into(),
530 ..Default::default()
531 },
532 ctx(),
533 )
534 .await
535 .unwrap();
536
537 let without_inheritance = h
538 .list_policies(
539 ListPoliciesRequest {
540 on_securable_type: "tables".to_string(),
541 on_securable_fullname: "cat.sch.tbl".to_string(),
542 include_inherited: None,
543 max_results: None,
544 page_token: None,
545 ..Default::default()
546 },
547 ctx(),
548 )
549 .await
550 .unwrap();
551 assert_eq!(without_inheritance.policies.len(), 1);
552 assert_eq!(without_inheritance.policies[0].name, "table_policy");
553
554 let with_inheritance = h
555 .list_policies(
556 ListPoliciesRequest {
557 on_securable_type: "tables".to_string(),
558 on_securable_fullname: "cat.sch.tbl".to_string(),
559 include_inherited: Some(true),
560 max_results: None,
561 page_token: None,
562 ..Default::default()
563 },
564 ctx(),
565 )
566 .await
567 .unwrap();
568 let mut names: Vec<_> = with_inheritance
569 .policies
570 .iter()
571 .map(|p| p.name.as_str())
572 .collect();
573 names.sort();
574 assert_eq!(
575 names,
576 vec!["catalog_policy", "schema_policy", "table_policy"]
577 );
578 let by_name = |n: &str| {
580 with_inheritance
581 .policies
582 .iter()
583 .find(|p| p.name == n)
584 .unwrap()
585 };
586 assert_eq!(by_name("catalog_policy").on_securable_fullname, "cat");
587 assert_eq!(by_name("schema_policy").on_securable_fullname, "cat.sch");
588 assert_eq!(by_name("table_policy").on_securable_fullname, "cat.sch.tbl");
589 }
590
591 #[tokio::test]
592 async fn create_rejects_duplicate_name_on_same_securable() {
593 let h = handler().await;
594 h.create_policy(
595 CreatePolicyRequest {
596 on_securable_type: "tables".to_string(),
597 on_securable_fullname: "cat.sch.tbl".to_string(),
598 policy_info: Some(row_filter_policy("p1", "", "")).into(),
599 ..Default::default()
600 },
601 ctx(),
602 )
603 .await
604 .unwrap();
605
606 let result = h
607 .create_policy(
608 CreatePolicyRequest {
609 on_securable_type: "tables".to_string(),
610 on_securable_fullname: "cat.sch.tbl".to_string(),
611 policy_info: Some(row_filter_policy("p1", "", "")).into(),
612 ..Default::default()
613 },
614 ctx(),
615 )
616 .await;
617 assert!(
618 result.is_err(),
619 "duplicate policy name on same securable must be rejected"
620 );
621 }
622
623 #[tokio::test]
624 async fn create_rejects_column_mask_without_match_columns() {
625 let h = handler().await;
626 let mut policy = row_filter_policy("bad_mask", "", "");
627 policy.policy_type = PolicyType::ColumnMask.into();
628 policy.function = Some(Function::ColumnMask(Box::new(FunctionRef {
629 function_name: "cat.sch.my_mask".to_string(),
630 using: vec![],
631 ..Default::default()
632 })));
633 let result = h
634 .create_policy(
635 CreatePolicyRequest {
636 on_securable_type: "tables".to_string(),
637 on_securable_fullname: "cat.sch.tbl".to_string(),
638 policy_info: Some(policy).into(),
639 ..Default::default()
640 },
641 ctx(),
642 )
643 .await;
644 assert!(
645 result.is_err(),
646 "column mask without match_columns must be rejected"
647 );
648 }
649
650 #[tokio::test]
651 async fn get_unknown_policy_is_not_found() {
652 let h = handler().await;
653 let result = h
654 .get_policy(
655 GetPolicyRequest {
656 on_securable_type: "tables".to_string(),
657 on_securable_fullname: "cat.sch.tbl".to_string(),
658 name: "does-not-exist".to_string(),
659 ..Default::default()
660 },
661 ctx(),
662 )
663 .await;
664 assert!(result.is_err(), "unknown policy must be NotFound");
665 }
666}