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::str::FromStr;
7
8use super::AccessControlRepository;
9use crate::authz::error::AuthzResult;
10use crate::authz::types::{EntityKind, EntityRow};
11
12impl AccessControlRepository {
13    /// `Ok(None)` means the entity has no catalog row at all (publish-pipeline
14    /// bootstrap gap) — the resolver turns this into
15    /// [`crate::authz::DenyReason::UnknownEntity`].
16    pub async fn get_entity(
17        &self,
18        entity_type: EntityKind,
19        entity_id: &str,
20    ) -> AuthzResult<Option<EntityRow>> {
21        let row = sqlx::query!(
22            r#"
23            SELECT entity_type, entity_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
33        let Some(row) = row else {
34            return Ok(None);
35        };
36        Ok(Some(EntityRow {
37            kind: EntityKind::from_str(&row.entity_type)?,
38            id: row.entity_id,
39            default_included: row.default_included,
40            source: row.source,
41        }))
42    }
43
44    /// Overwrites `default_included` and `source` on conflict so the most
45    /// recent bootstrap pass wins — the publish pipeline is the source of
46    /// truth and runs ahead of YAML grant ingestion.
47    pub async fn upsert_entity(
48        &self,
49        entity_type: EntityKind,
50        entity_id: &str,
51        default_included: bool,
52        source: &str,
53    ) -> AuthzResult<()> {
54        sqlx::query!(
55            r#"
56            INSERT INTO access_control_entities (entity_type, entity_id, default_included, source)
57            VALUES ($1, $2, $3, $4)
58            ON CONFLICT (entity_type, entity_id) DO UPDATE
59            SET default_included = EXCLUDED.default_included,
60                source = EXCLUDED.source,
61                updated_at = NOW()
62            "#,
63            entity_type.as_str(),
64            entity_id,
65            default_included,
66            source,
67        )
68        .execute(&*self.write_pool)
69        .await?;
70        Ok(())
71    }
72
73    /// One statement for the whole batch, instead of `ids.len()` awaits of
74    /// [`Self::upsert_entity`]; all rows share one `default_included` and
75    /// `source`.
76    pub async fn upsert_entities(
77        &self,
78        entity_type: EntityKind,
79        ids: &[&str],
80        default_included: bool,
81        source: &str,
82    ) -> AuthzResult<()> {
83        if ids.is_empty() {
84            return Ok(());
85        }
86        let ids_owned: Vec<String> = ids.iter().map(|id| (*id).to_owned()).collect();
87        sqlx::query!(
88            r#"
89            INSERT INTO access_control_entities (entity_type, entity_id, default_included, source)
90            SELECT $1, id, $3, $4
91            FROM UNNEST($2::text[]) AS id
92            ON CONFLICT (entity_type, entity_id) DO UPDATE
93            SET default_included = EXCLUDED.default_included,
94                source = EXCLUDED.source,
95                updated_at = NOW()
96            "#,
97            entity_type.as_str(),
98            &ids_owned,
99            default_included,
100            source,
101        )
102        .execute(&*self.write_pool)
103        .await?;
104        Ok(())
105    }
106
107    pub async fn list_entities(&self, entity_type: EntityKind) -> AuthzResult<Vec<EntityRow>> {
108        let rows = sqlx::query!(
109            r#"
110            SELECT entity_type, entity_id, default_included, source
111            FROM access_control_entities
112            WHERE entity_type = $1
113            ORDER BY entity_id
114            "#,
115            entity_type.as_str(),
116        )
117        .fetch_all(&*self.pool)
118        .await?;
119
120        let mut out = Vec::with_capacity(rows.len());
121        for row in rows {
122            out.push(EntityRow {
123                kind: EntityKind::from_str(&row.entity_type)?,
124                id: row.entity_id,
125                default_included: row.default_included,
126                source: row.source,
127            });
128        }
129        Ok(out)
130    }
131}