mas_storage_pg/
policy_data.rs

1// Copyright 2025, 2026 Element Creations Ltd.
2// Copyright 2025 New Vector Ltd.
3//
4// SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-Element-Commercial
5// Please see LICENSE files in the repository root for full details.
6
7//! A module containing the PostgreSQL implementation of the policy data
8//! storage.
9
10use async_trait::async_trait;
11use mas_data_model::{Clock, PolicyData, UlidExt as _};
12use mas_storage::policy_data::PolicyDataRepository;
13use rand::Rng;
14use serde_json::Value;
15use sqlx::{PgConnection, types::Json};
16use ulid::Ulid;
17use uuid::Uuid;
18
19use crate::{DatabaseError, ExecuteExt};
20
21/// An implementation of [`PolicyDataRepository`] for a PostgreSQL connection.
22pub struct PgPolicyDataRepository<'c> {
23    conn: &'c mut PgConnection,
24}
25
26impl<'c> PgPolicyDataRepository<'c> {
27    /// Create a new [`PgPolicyDataRepository`] from an active PostgreSQL
28    /// connection.
29    #[must_use]
30    pub fn new(conn: &'c mut PgConnection) -> Self {
31        Self { conn }
32    }
33}
34
35struct PolicyDataLookup {
36    policy_data_id: Uuid,
37    created_at: chrono::DateTime<chrono::Utc>,
38    data: Json<Value>,
39}
40
41impl From<PolicyDataLookup> for PolicyData {
42    fn from(value: PolicyDataLookup) -> Self {
43        PolicyData {
44            id: value.policy_data_id.into(),
45            created_at: value.created_at,
46            data: value.data.0,
47        }
48    }
49}
50
51#[async_trait]
52impl PolicyDataRepository for PgPolicyDataRepository<'_> {
53    type Error = DatabaseError;
54
55    #[tracing::instrument(
56        name = "db.policy_data.get",
57        skip_all,
58        fields(
59            db.query.text,
60        ),
61        err,
62    )]
63    async fn get(&mut self) -> Result<Option<PolicyData>, Self::Error> {
64        let row = sqlx::query_as!(
65            PolicyDataLookup,
66            r#"
67            SELECT policy_data_id, created_at, data
68            FROM policy_data
69            ORDER BY policy_data_id DESC
70            LIMIT 1
71            "#
72        )
73        .traced()
74        .fetch_optional(&mut *self.conn)
75        .await?;
76
77        let Some(row) = row else {
78            return Ok(None);
79        };
80
81        Ok(Some(row.into()))
82    }
83
84    #[tracing::instrument(
85        name = "db.policy_data.set",
86        skip_all,
87        fields(
88            db.query.text,
89        ),
90        err,
91    )]
92    async fn set(
93        &mut self,
94        rng: &mut (dyn Rng + Send),
95        clock: &dyn Clock,
96        data: Value,
97    ) -> Result<PolicyData, Self::Error> {
98        let created_at = clock.now();
99        let id = Ulid::from_datetime_with_rng(created_at, rng);
100
101        sqlx::query!(
102            r#"
103            INSERT INTO policy_data (policy_data_id, created_at, data)
104            VALUES ($1, $2, $3)
105            "#,
106            Uuid::from(id),
107            created_at,
108            data,
109        )
110        .traced()
111        .execute(&mut *self.conn)
112        .await?;
113
114        Ok(PolicyData {
115            id,
116            created_at,
117            data,
118        })
119    }
120
121    #[tracing::instrument(
122        name = "db.policy_data.prune",
123        skip_all,
124        fields(
125            db.query.text,
126        ),
127        err,
128    )]
129    async fn prune(&mut self, keep: usize) -> Result<usize, Self::Error> {
130        let res = sqlx::query!(
131            r#"
132            DELETE FROM policy_data
133            WHERE policy_data_id IN (
134                SELECT policy_data_id
135                FROM policy_data
136                ORDER BY policy_data_id DESC
137                OFFSET $1
138            )
139            "#,
140            i64::try_from(keep).map_err(DatabaseError::to_invalid_operation)?
141        )
142        .traced()
143        .execute(&mut *self.conn)
144        .await?;
145
146        Ok(res
147            .rows_affected()
148            .try_into()
149            .map_err(DatabaseError::to_invalid_operation)?)
150    }
151}
152
153#[cfg(test)]
154mod tests {
155    use mas_data_model::clock::MockClock;
156    use mas_storage::policy_data::PolicyDataRepository;
157    use rand::SeedableRng;
158    use rand_chacha::ChaChaRng;
159    use serde_json::json;
160    use sqlx::PgPool;
161
162    use crate::policy_data::PgPolicyDataRepository;
163
164    #[sqlx::test(migrator = "crate::MIGRATOR")]
165    async fn test_policy_data(pool: PgPool) {
166        let mut rng = ChaChaRng::seed_from_u64(42);
167        let clock = MockClock::default();
168        let mut conn = pool.acquire().await.unwrap();
169        let mut repo = PgPolicyDataRepository::new(&mut conn);
170
171        // Get an empty state at first
172        let data = repo.get().await.unwrap();
173        assert_eq!(data, None);
174
175        // Set some data
176        let value1 = json!({"hello": "world"});
177        let policy_data1 = repo.set(&mut rng, &clock, value1.clone()).await.unwrap();
178        assert_eq!(policy_data1.data, value1);
179
180        let data_fetched1 = repo.get().await.unwrap().unwrap();
181        assert_eq!(policy_data1, data_fetched1);
182
183        // Set some new data
184        clock.advance(chrono::Duration::seconds(1));
185        let value2 = json!({"foo": "bar"});
186        let policy_data2 = repo.set(&mut rng, &clock, value2.clone()).await.unwrap();
187        assert_eq!(policy_data2.data, value2);
188
189        // Check the new data is fetched
190        let data_fetched2 = repo.get().await.unwrap().unwrap();
191        assert_eq!(data_fetched2, policy_data2);
192
193        // Prune until the first entry
194        let affected = repo.prune(1).await.unwrap();
195        let data_fetched3 = repo.get().await.unwrap().unwrap();
196        assert_eq!(data_fetched3, policy_data2);
197        assert_eq!(affected, 1);
198
199        // Do a raw query to check the other rows were pruned
200        let count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM policy_data")
201            .fetch_one(&mut *conn)
202            .await
203            .unwrap();
204        assert_eq!(count, 1);
205    }
206}
mas_storage_pg/policy_data.rs

mas_storage_pg/
policy_data.rs
