mcp_usage_store/
postgres.rs

1use std::time::Duration;
2use std::{future::Future, result::Result as StdResult};
3
4use mcp_usage_core::TaskAttribution;
5use mcp_usage_tower::{TaskAttributionStore, TaskStoreError, TaskStoreFuture};
6use sqlx::PgPool;
7
8use crate::{StoreConfigError, decode_attribution, encode_attribution, identifier_hash};
9
10/// Schema installation, serialized by a transaction-scoped advisory lock.
11///
12/// `CREATE TABLE IF NOT EXISTS` is not atomic in `PostgreSQL`. Racing sessions
13/// each consult the catalog, all conclude the table is absent, and every loser
14/// fails with a duplicate key violation on `pg_type_typname_nsp_index`. Taking
15/// an advisory lock first serializes the check, so late arrivals find the table
16/// already present and do nothing.
17///
18/// `sqlx::raw_sql` uses the simple query protocol, which `PostgreSQL` executes as
19/// a single implicit transaction. The lock is therefore held across the DDL and
20/// released at commit, with no explicit transaction object to thread through.
21///
22/// The lock key is an arbitrary but stable constant, assembled at compile time.
23/// It only needs to be distinct from other advisory locks the host application
24/// takes on the same database.
25const INSTALL_SQL: &str = concat!(
26    "SELECT pg_advisory_xact_lock(7594834021775991);\n",
27    include_str!("../schema/postgres.sql"),
28    "\nSELECT attribution FROM mcp_usage_task_attribution LIMIT 0;"
29);
30const DEFAULT_OPERATION_TIMEOUT: Duration = Duration::from_secs(2);
31
32/// `PostgreSQL` durable-task attribution with atomic, once-only claims.
33#[derive(Clone)]
34pub struct PostgresTaskStore {
35    pool: PgPool,
36    ttl_seconds: i64,
37    operation_timeout: Duration,
38}
39
40impl std::fmt::Debug for PostgresTaskStore {
41    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
42        formatter
43            .debug_struct("PostgresTaskStore")
44            .field("ttl_seconds", &self.ttl_seconds)
45            .field("operation_timeout", &self.operation_timeout)
46            .finish_non_exhaustive()
47    }
48}
49
50impl PostgresTaskStore {
51    /// Construct a store from an application-owned connection pool.
52    ///
53    /// # Errors
54    ///
55    /// Returns [`StoreConfigError::InvalidTtl`] when the TTL is zero or cannot
56    /// be represented by `PostgreSQL`'s signed interval input.
57    pub fn new(pool: PgPool, ttl: Duration) -> Result<Self, StoreConfigError> {
58        Self::with_timeout(pool, ttl, DEFAULT_OPERATION_TIMEOUT)
59    }
60
61    /// Construct a store with an explicit bound for every database operation.
62    ///
63    /// # Errors
64    ///
65    /// Returns a configuration error for an invalid TTL or zero timeout.
66    pub fn with_timeout(
67        pool: PgPool,
68        ttl: Duration,
69        operation_timeout: Duration,
70    ) -> Result<Self, StoreConfigError> {
71        let ttl_seconds = i64::try_from(ttl.as_secs()).map_err(|_| StoreConfigError::InvalidTtl)?;
72        if ttl_seconds == 0 {
73            return Err(StoreConfigError::InvalidTtl);
74        }
75        if operation_timeout.is_zero() {
76            return Err(StoreConfigError::InvalidTimeout);
77        }
78        Ok(Self {
79            pool,
80            ttl_seconds,
81            operation_timeout,
82        })
83    }
84
85    /// Install the idempotent table and expiry index.
86    ///
87    /// Safe to call concurrently from every instance of a horizontally scaled
88    /// application, which is the situation this store exists for. `CREATE TABLE
89    /// IF NOT EXISTS` is *not* atomic in `PostgreSQL`: racing sessions check the
90    /// catalog, all decide the table is absent, and every loser fails with a
91    /// duplicate key violation on `pg_type_typname_nsp_index`. A transaction
92    /// scoped advisory lock serializes the check so late arrivals simply find
93    /// the table already present and do nothing.
94    ///
95    /// Applications with migration-controlled schemas can apply
96    /// `schema/postgres.sql` themselves and skip this operation.
97    ///
98    /// # Errors
99    ///
100    /// Returns a sanitized schema error without exposing connection details.
101    pub async fn install(&self) -> Result<(), StoreConfigError> {
102        tokio::time::timeout(
103            self.operation_timeout,
104            sqlx::raw_sql(INSTALL_SQL).execute(&self.pool),
105        )
106        .await
107        .map_err(|_| StoreConfigError::Schema)?
108        .map_err(|_| StoreConfigError::Schema)?;
109        Ok(())
110    }
111
112    /// Delete expired task origins and return the affected row count.
113    ///
114    /// # Errors
115    ///
116    /// Returns a sanitized backend error.
117    pub async fn prune_expired(&self) -> Result<u64, TaskStoreError> {
118        self.run(
119            sqlx::query("DELETE FROM mcp_usage_task_attribution WHERE expires_at <= NOW()")
120                .execute(&self.pool),
121        )
122        .await
123        .map(|result| result.rows_affected())
124    }
125
126    async fn run<T, F>(&self, operation: F) -> Result<T, TaskStoreError>
127    where
128        F: Future<Output = StdResult<T, sqlx::Error>>,
129    {
130        tokio::time::timeout(self.operation_timeout, operation)
131            .await
132            .map_err(|_| TaskStoreError::BackendUnavailable)?
133            .map_err(|_| TaskStoreError::BackendUnavailable)
134    }
135}
136
137impl TaskAttributionStore for PostgresTaskStore {
138    fn insert<'a>(
139        &'a self,
140        tenant_id: &'a str,
141        task_id: &'a str,
142        attribution: TaskAttribution,
143    ) -> TaskStoreFuture<'a, ()> {
144        Box::pin(async move {
145            self.run(
146                sqlx::query(
147                    "INSERT INTO mcp_usage_task_attribution \
148                 (tenant_hash, task_hash, attribution, expires_at) \
149                 VALUES ($1, $2, $3, NOW() + ($4 * INTERVAL '1 second')) \
150                 ON CONFLICT (tenant_hash, task_hash) DO NOTHING",
151                )
152                .bind(identifier_hash(tenant_id).to_vec())
153                .bind(identifier_hash(task_id).to_vec())
154                .bind(encode_attribution(attribution).to_vec())
155                .bind(self.ttl_seconds)
156                .execute(&self.pool),
157            )
158            .await?;
159            Ok(())
160        })
161    }
162
163    fn get<'a>(
164        &'a self,
165        tenant_id: &'a str,
166        task_id: &'a str,
167    ) -> TaskStoreFuture<'a, Option<TaskAttribution>> {
168        Box::pin(async move {
169            let row: Option<(Vec<u8>,)> = self
170                .run(
171                    sqlx::query_as(
172                        "SELECT attribution FROM mcp_usage_task_attribution \
173                 WHERE tenant_hash = $1 AND task_hash = $2 AND expires_at > NOW()",
174                    )
175                    .bind(identifier_hash(tenant_id).to_vec())
176                    .bind(identifier_hash(task_id).to_vec())
177                    .fetch_optional(&self.pool),
178                )
179                .await?;
180            row.map(|(attribution,)| decode_attribution(&attribution))
181                .transpose()
182        })
183    }
184
185    fn claim<'a>(
186        &'a self,
187        tenant_id: &'a str,
188        task_id: &'a str,
189    ) -> TaskStoreFuture<'a, Option<TaskAttribution>> {
190        Box::pin(async move {
191            let row: Option<(Vec<u8>,)> = self
192                .run(
193                    sqlx::query_as(
194                        "DELETE FROM mcp_usage_task_attribution \
195                 WHERE tenant_hash = $1 AND task_hash = $2 AND expires_at > NOW() \
196                 RETURNING attribution",
197                    )
198                    .bind(identifier_hash(tenant_id).to_vec())
199                    .bind(identifier_hash(task_id).to_vec())
200                    .fetch_optional(&self.pool),
201                )
202                .await?;
203            row.map(|(attribution,)| decode_attribution(&attribution))
204                .transpose()
205        })
206    }
207
208    fn remove<'a>(&'a self, tenant_id: &'a str, task_id: &'a str) -> TaskStoreFuture<'a, ()> {
209        Box::pin(async move {
210            self.run(
211                sqlx::query(
212                    "DELETE FROM mcp_usage_task_attribution \
213                 WHERE tenant_hash = $1 AND task_hash = $2",
214                )
215                .bind(identifier_hash(tenant_id).to_vec())
216                .bind(identifier_hash(task_id).to_vec())
217                .execute(&self.pool),
218            )
219            .await?;
220            Ok(())
221        })
222    }
223}
224
225#[cfg(test)]
226mod tests {
227    use super::*;
228    use sqlx::postgres::PgPoolOptions;
229
230    #[tokio::test]
231    async fn rejects_zero_operation_timeout_without_contacting_postgres() {
232        let pool = PgPoolOptions::new()
233            .connect_lazy("postgres://usagekit:usagekit@127.0.0.1/usagekit")
234            .unwrap();
235        let error = PostgresTaskStore::with_timeout(pool, Duration::from_secs(60), Duration::ZERO)
236            .unwrap_err();
237
238        assert_eq!(error, StoreConfigError::InvalidTimeout);
239    }
240}
mcp_usage_store/postgres.rs

mcp_usage_store/
postgres.rs
