Skip to main content

systemprompt_security/authz/repository/
entities.rs

1//! Entity-catalog persistence for access control.
2//!
3//! Copyright (c) systemprompt.io — Business Source License 1.1.
4//! See <https://systemprompt.io> for licensing details.
5
6use 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}