systemprompt_security/authz/repository/
entities.rs1use std::collections::HashMap;
7
8use super::AccessControlRepository;
9use crate::authz::error::AuthzResult;
10use crate::authz::types::{EntityKind, EntityRow};
11
12impl AccessControlRepository {
13 pub async fn get_entity(
14 &self,
15 entity_type: EntityKind,
16 entity_id: &str,
17 ) -> AuthzResult<Option<EntityRow>> {
18 let row = sqlx::query_as!(
19 EntityRow,
20 r#"
21 SELECT entity_type AS "kind: EntityKind", entity_id AS id, default_included, source
22 FROM access_control_entities
23 WHERE entity_type = $1 AND entity_id = $2
24 "#,
25 entity_type.as_str(),
26 entity_id,
27 )
28 .fetch_optional(&*self.pool)
29 .await?;
30 Ok(row)
31 }
32
33 pub async fn list_entities_bulk(
34 &self,
35 entity_type: EntityKind,
36 entity_ids: &[String],
37 ) -> AuthzResult<HashMap<String, EntityRow>> {
38 if entity_ids.is_empty() {
39 return Ok(HashMap::new());
40 }
41 let rows = sqlx::query_as!(
42 EntityRow,
43 r#"
44 SELECT entity_type AS "kind: EntityKind", entity_id AS id, default_included, source
45 FROM access_control_entities
46 WHERE entity_type = $1 AND entity_id = ANY($2)
47 "#,
48 entity_type.as_str(),
49 entity_ids,
50 )
51 .fetch_all(&*self.pool)
52 .await?;
53 Ok(rows.into_iter().map(|row| (row.id.clone(), row)).collect())
54 }
55
56 pub async fn upsert_entity(
57 &self,
58 entity_type: EntityKind,
59 entity_id: &str,
60 default_included: bool,
61 source: &str,
62 ) -> AuthzResult<()> {
63 sqlx::query!(
64 r#"
65 INSERT INTO access_control_entities (entity_type, entity_id, default_included, source)
66 VALUES ($1, $2, $3, $4)
67 ON CONFLICT (entity_type, entity_id) DO UPDATE
68 SET default_included = EXCLUDED.default_included,
69 source = EXCLUDED.source,
70 updated_at = NOW()
71 "#,
72 entity_type.as_str(),
73 entity_id,
74 default_included,
75 source,
76 )
77 .execute(&*self.write_pool)
78 .await?;
79 Ok(())
80 }
81
82 pub async fn upsert_entities(
83 &self,
84 entity_type: EntityKind,
85 ids: &[&str],
86 default_included: bool,
87 source: &str,
88 ) -> AuthzResult<()> {
89 if ids.is_empty() {
90 return Ok(());
91 }
92 let ids_owned: Vec<String> = ids.iter().map(|id| (*id).to_owned()).collect();
93 sqlx::query!(
94 r#"
95 INSERT INTO access_control_entities (entity_type, entity_id, default_included, source)
96 SELECT $1, id, $3, $4
97 FROM UNNEST($2::text[]) AS id
98 ON CONFLICT (entity_type, entity_id) DO UPDATE
99 SET default_included = EXCLUDED.default_included,
100 source = EXCLUDED.source,
101 updated_at = NOW()
102 "#,
103 entity_type.as_str(),
104 &ids_owned,
105 default_included,
106 source,
107 )
108 .execute(&*self.write_pool)
109 .await?;
110 Ok(())
111 }
112
113 pub async fn list_entities(&self, entity_type: EntityKind) -> AuthzResult<Vec<EntityRow>> {
114 let rows = sqlx::query_as!(
115 EntityRow,
116 r#"
117 SELECT entity_type AS "kind: EntityKind", entity_id AS id, default_included, source
118 FROM access_control_entities
119 WHERE entity_type = $1
120 ORDER BY entity_id
121 "#,
122 entity_type.as_str(),
123 )
124 .fetch_all(&*self.pool)
125 .await?;
126 Ok(rows)
127 }
128}