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    )]
62    async fn get(&mut self) -> Result<Option<PolicyData>, Self::Error> {
63        let row = sqlx::query_as!(
64            PolicyDataLookup,
65            r#"
66            SELECT policy_data_id, created_at, data
67            FROM policy_data
68            ORDER BY policy_data_id DESC
69            LIMIT 1
70            "#
71        )
72        .traced()
73        .fetch_optional(&mut *self.conn)
74        .await?;
75
76        let Some(row) = row else {
77            return Ok(None);
78        };
79
80        Ok(Some(row.into()))
81    }
82
83    #[tracing::instrument(
84        name = "db.policy_data.set",
85        skip_all,
86        fields(
87            db.query.text,
88        ),
89    )]
90    async fn set(
91        &mut self,
92        rng: &mut (dyn Rng + Send),
93        clock: &dyn Clock,
94        data: Value,
95    ) -> Result<PolicyData, Self::Error> {
96        let created_at = clock.now();
97        let id = Ulid::from_datetime_with_rng(created_at, rng);
98
99        sqlx::query!(
100            r#"
101            INSERT INTO policy_data (policy_data_id, created_at, data)
102            VALUES ($1, $2, $3)
103            "#,
104            Uuid::from(id),
105            created_at,
106            data,
107        )
108        .traced()
109        .execute(&mut *self.conn)
110        .await?;
111
112        Ok(PolicyData {
113            id,
114            created_at,
115            data,
116        })
117    }
118
119    #[tracing::instrument(
120        name = "db.policy_data.prune",
121        skip_all,
122        fields(
123            db.query.text,
124        ),
125    )]
126    async fn prune(&mut self, keep: usize) -> Result<usize, Self::Error> {
127        let res = sqlx::query!(
128            r#"
129            DELETE FROM policy_data
130            WHERE policy_data_id IN (
131                SELECT policy_data_id
132                FROM policy_data
133                ORDER BY policy_data_id DESC
134                OFFSET $1
135            )
136            "#,
137            i64::try_from(keep).map_err(DatabaseError::to_invalid_operation)?
138        )
139        .traced()
140        .execute(&mut *self.conn)
141        .await?;
142
143        Ok(res
144            .rows_affected()
145            .try_into()
146            .map_err(DatabaseError::to_invalid_operation)?)
147    }
148}
149
150#[cfg(test)]
151mod tests {
152    use mas_data_model::clock::MockClock;
153    use mas_storage::policy_data::PolicyDataRepository;
154    use rand::SeedableRng;
155    use rand_chacha::ChaChaRng;
156    use serde_json::json;
157    use sqlx::PgPool;
158
159    use crate::policy_data::PgPolicyDataRepository;
160
161    #[sqlx::test(migrator = "crate::MIGRATOR")]
162    async fn test_policy_data(pool: PgPool) {
163        let mut rng = ChaChaRng::seed_from_u64(42);
164        let clock = MockClock::default();
165        let mut conn = pool.acquire().await.unwrap();
166        let mut repo = PgPolicyDataRepository::new(&mut conn);
167
168        // Get an empty state at first
169        let data = repo.get().await.unwrap();
170        assert_eq!(data, None);
171
172        // Set some data
173        let value1 = json!({"hello": "world"});
174        let policy_data1 = repo.set(&mut rng, &clock, value1.clone()).await.unwrap();
175        assert_eq!(policy_data1.data, value1);
176
177        let data_fetched1 = repo.get().await.unwrap().unwrap();
178        assert_eq!(policy_data1, data_fetched1);
179
180        // Set some new data
181        clock.advance(chrono::Duration::seconds(1));
182        let value2 = json!({"foo": "bar"});
183        let policy_data2 = repo.set(&mut rng, &clock, value2.clone()).await.unwrap();
184        assert_eq!(policy_data2.data, value2);
185
186        // Check the new data is fetched
187        let data_fetched2 = repo.get().await.unwrap().unwrap();
188        assert_eq!(data_fetched2, policy_data2);
189
190        // Prune until the first entry
191        let affected = repo.prune(1).await.unwrap();
192        let data_fetched3 = repo.get().await.unwrap().unwrap();
193        assert_eq!(data_fetched3, policy_data2);
194        assert_eq!(affected, 1);
195
196        // Do a raw query to check the other rows were pruned
197        let count: i64 = sqlx::query_scalar("SELECT COUNT(*) FROM policy_data")
198            .fetch_one(&mut *conn)
199            .await
200            .unwrap();
201        assert_eq!(count, 1);
202    }
203}
mas_storage_pg/policy_data.rs

mas_storage_pg/
policy_data.rs
