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 sqlx::PgConnection;
9
10use super::AccessControlRepository;
11use crate::authz::error::AuthzResult;
12use crate::authz::types::{EntityKind, EntityRow};
13
14impl AccessControlRepository {
15    pub async fn get_entity(
16        &self,
17        entity_type: EntityKind,
18        entity_id: &str,
19    ) -> AuthzResult<Option<EntityRow>> {
20        let row = sqlx::query_as!(
21            EntityRow,
22            r#"
23            SELECT entity_type AS "kind: EntityKind", entity_id AS id, default_included, source
24            FROM access_control_entities
25            WHERE entity_type = $1 AND entity_id = $2
26            "#,
27            entity_type.as_str(),
28            entity_id,
29        )
30        .fetch_optional(&*self.pool)
31        .await?;
32        Ok(row)
33    }
34
35    pub async fn list_entities_bulk(
36        &self,
37        entity_type: EntityKind,
38        entity_ids: &[String],
39    ) -> AuthzResult<HashMap<String, EntityRow>> {
40        if entity_ids.is_empty() {
41            return Ok(HashMap::new());
42        }
43        let rows = sqlx::query_as!(
44            EntityRow,
45            r#"
46            SELECT entity_type AS "kind: EntityKind", entity_id AS id, default_included, source
47            FROM access_control_entities
48            WHERE entity_type = $1 AND entity_id = ANY($2)
49            "#,
50            entity_type.as_str(),
51            entity_ids,
52        )
53        .fetch_all(&*self.pool)
54        .await?;
55        Ok(rows.into_iter().map(|row| (row.id.clone(), row)).collect())
56    }
57
58    pub async fn upsert_entity(
59        &self,
60        entity_type: EntityKind,
61        entity_id: &str,
62        default_included: bool,
63        source: &str,
64    ) -> AuthzResult<()> {
65        sqlx::query!(
66            r#"
67            INSERT INTO access_control_entities (entity_type, entity_id, default_included, source)
68            VALUES ($1, $2, $3, $4)
69            ON CONFLICT (entity_type, entity_id) DO UPDATE
70            SET default_included = EXCLUDED.default_included,
71                source = EXCLUDED.source,
72                updated_at = NOW()
73            "#,
74            entity_type.as_str(),
75            entity_id,
76            default_included,
77            source,
78        )
79        .execute(&*self.write_pool)
80        .await?;
81        Ok(())
82    }
83
84    // Why: leaves an existing row's `default_included` alone, so a write that
85    // only needs the FK satisfied cannot widen access. `upsert_entity`
86    // overwrites the flag; this does not.
87    pub async fn ensure_entity(
88        &self,
89        entity_type: EntityKind,
90        entity_id: &str,
91        source: &str,
92    ) -> AuthzResult<()> {
93        sqlx::query!(
94            r#"
95            INSERT INTO access_control_entities (entity_type, entity_id, default_included, source)
96            VALUES ($1, $2, false, $3)
97            ON CONFLICT (entity_type, entity_id) DO NOTHING
98            "#,
99            entity_type.as_str(),
100            entity_id,
101            source,
102        )
103        .execute(&*self.write_pool)
104        .await?;
105        Ok(())
106    }
107
108    // Why: `access_control_rules` FKs onto this table ON DELETE CASCADE, so a
109    // pruned entity takes its grants with it. Only call this for a kind whose
110    // caller can enumerate the complete real set — an empty `keep` deletes
111    // every row of the kind, which is why `gateway_entities` refuses one.
112    pub async fn reconcile_entities(
113        &self,
114        entity_type: EntityKind,
115        keep: &[&str],
116        default_included: bool,
117        source: &str,
118    ) -> AuthzResult<u64> {
119        let mut tx = self.write_pool.begin().await?;
120        upsert_entities_on(&mut tx, entity_type, keep, default_included, source).await?;
121        let keep_owned: Vec<String> = keep.iter().map(|id| (*id).to_owned()).collect();
122        let res = sqlx::query!(
123            r#"
124            DELETE FROM access_control_entities
125            WHERE entity_type = $1
126              AND entity_id <> ALL($2::text[])
127            "#,
128            entity_type.as_str(),
129            &keep_owned,
130        )
131        .execute(&mut *tx)
132        .await?;
133        tx.commit().await?;
134        Ok(res.rows_affected())
135    }
136
137    pub async fn list_entities(&self, entity_type: EntityKind) -> AuthzResult<Vec<EntityRow>> {
138        let rows = sqlx::query_as!(
139            EntityRow,
140            r#"
141            SELECT entity_type AS "kind: EntityKind", entity_id AS id, default_included, source
142            FROM access_control_entities
143            WHERE entity_type = $1
144            ORDER BY entity_id
145            "#,
146            entity_type.as_str(),
147        )
148        .fetch_all(&*self.pool)
149        .await?;
150        Ok(rows)
151    }
152}
153
154async fn upsert_entities_on(
155    conn: &mut PgConnection,
156    entity_type: EntityKind,
157    ids: &[&str],
158    default_included: bool,
159    source: &str,
160) -> AuthzResult<()> {
161    if ids.is_empty() {
162        return Ok(());
163    }
164    let ids_owned: Vec<String> = ids.iter().map(|id| (*id).to_owned()).collect();
165    sqlx::query!(
166        r#"
167        INSERT INTO access_control_entities (entity_type, entity_id, default_included, source)
168        SELECT $1, id, $3, $4
169        FROM UNNEST($2::text[]) AS id
170        ON CONFLICT (entity_type, entity_id) DO UPDATE
171        SET default_included = EXCLUDED.default_included,
172            source = EXCLUDED.source,
173            updated_at = NOW()
174        "#,
175        entity_type.as_str(),
176        &ids_owned,
177        default_included,
178        source,
179    )
180    .execute(conn)
181    .await?;
182    Ok(())
183}
systemprompt_security/authz/repository/entities.rs

systemprompt_security/authz/repository/
entities.rs
