systemprompt_database/repository/service/
repo.rs

1//! Async repository over the `services` registry table.
2//!
3//! Copyright (c) systemprompt.io — Business Source License 1.1.
4//! See <https://systemprompt.io> for licensing details.
5
6use std::sync::Arc;
7
8use sqlx::PgPool;
9
10use super::model::{CreateServiceInput, ServiceConfig};
11use crate::DbPool;
12use crate::error::DatabaseResult;
13
14#[derive(Debug, Clone)]
15pub struct ServiceRepository {
16    pool: Arc<PgPool>,
17    write_pool: Arc<PgPool>,
18}
19
20impl ServiceRepository {
21    pub fn new(db: &DbPool) -> DatabaseResult<Self> {
22        let pool = db.pool_arc()?;
23        let write_pool = db.write_pool_arc()?;
24        Ok(Self { pool, write_pool })
25    }
26
27    pub async fn find_service_by_name(&self, name: &str) -> DatabaseResult<Option<ServiceConfig>> {
28        let row = sqlx::query!(
29            r#"
30            SELECT name, module_name, status, pid, port, binary_mtime,
31                   created_at::text as "created_at!", updated_at::text as "updated_at!"
32            FROM services
33            WHERE name = $1
34            "#,
35            name
36        )
37        .fetch_optional(&*self.pool)
38        .await?;
39
40        Ok(row.map(|r| ServiceConfig {
41            name: r.name,
42            module_name: r.module_name,
43            status: r.status,
44            pid: r.pid,
45            port: r.port,
46            binary_mtime: r.binary_mtime,
47            created_at: r.created_at,
48            updated_at: r.updated_at,
49        }))
50    }
51
52    pub async fn list_all_agent_service_names(&self) -> DatabaseResult<Vec<String>> {
53        let rows = sqlx::query!(r#"SELECT name FROM services WHERE module_name = 'agent'"#)
54            .fetch_all(&*self.pool)
55            .await?;
56        Ok(rows.into_iter().map(|r| r.name).collect())
57    }
58
59    pub async fn list_mcp_services(&self) -> DatabaseResult<Vec<ServiceConfig>> {
60        let rows = sqlx::query!(
61            r#"
62            SELECT name, module_name, status, pid, port, binary_mtime,
63                   created_at::text as "created_at!", updated_at::text as "updated_at!"
64            FROM services
65            WHERE module_name = 'mcp'
66            ORDER BY name
67            "#
68        )
69        .fetch_all(&*self.pool)
70        .await?;
71        Ok(rows
72            .into_iter()
73            .map(|r| ServiceConfig {
74                name: r.name,
75                module_name: r.module_name,
76                status: r.status,
77                pid: r.pid,
78                port: r.port,
79                binary_mtime: r.binary_mtime,
80                created_at: r.created_at,
81                updated_at: r.updated_at,
82            })
83            .collect())
84    }
85
86    pub async fn create_service(&self, input: CreateServiceInput<'_>) -> DatabaseResult<()> {
87        let port_i32 = i32::from(input.port);
88        sqlx::query!(
89            r#"
90            INSERT INTO services (name, module_name, status, port, binary_mtime)
91            VALUES ($1, $2, $3, $4, $5)
92            ON CONFLICT (name) DO UPDATE SET
93              module_name = EXCLUDED.module_name,
94              status = EXCLUDED.status,
95              port = EXCLUDED.port,
96              binary_mtime = EXCLUDED.binary_mtime,
97              updated_at = CURRENT_TIMESTAMP
98            "#,
99            input.name,
100            input.module_name,
101            input.status,
102            port_i32,
103            input.binary_mtime
104        )
105        .execute(&*self.write_pool)
106        .await?;
107        Ok(())
108    }
109
110    pub async fn update_service_status(
111        &self,
112        service_name: &str,
113        status: &str,
114    ) -> DatabaseResult<()> {
115        sqlx::query!(
116            r#"UPDATE services SET status = $1, updated_at = CURRENT_TIMESTAMP WHERE name = $2"#,
117            status,
118            service_name
119        )
120        .execute(&*self.write_pool)
121        .await?;
122        Ok(())
123    }
124
125    pub async fn delete_service(&self, service_name: &str) -> DatabaseResult<()> {
126        sqlx::query!(r#"DELETE FROM services WHERE name = $1"#, service_name)
127            .execute(&*self.write_pool)
128            .await?;
129        Ok(())
130    }
131
132    pub async fn update_service_pid(&self, service_name: &str, pid: i32) -> DatabaseResult<()> {
133        sqlx::query!(
134            r#"UPDATE services SET pid = $1, updated_at = CURRENT_TIMESTAMP WHERE name = $2"#,
135            pid,
136            service_name
137        )
138        .execute(&*self.write_pool)
139        .await?;
140        Ok(())
141    }
142
143    pub async fn clear_service_pid(&self, service_name: &str) -> DatabaseResult<()> {
144        sqlx::query!(
145            r#"UPDATE services SET pid = NULL, updated_at = CURRENT_TIMESTAMP WHERE name = $1"#,
146            service_name
147        )
148        .execute(&*self.write_pool)
149        .await?;
150        Ok(())
151    }
152
153    pub async fn list_all_running_services(&self) -> DatabaseResult<Vec<ServiceConfig>> {
154        let rows = sqlx::query!(
155            r#"
156            SELECT name, module_name, status, pid, port, binary_mtime,
157                   created_at::text as "created_at!", updated_at::text as "updated_at!"
158            FROM services
159            WHERE status = 'running'
160            ORDER BY name
161            "#
162        )
163        .fetch_all(&*self.pool)
164        .await?;
165        Ok(rows
166            .into_iter()
167            .map(|r| ServiceConfig {
168                name: r.name,
169                module_name: r.module_name,
170                status: r.status,
171                pid: r.pid,
172                port: r.port,
173                binary_mtime: r.binary_mtime,
174                created_at: r.created_at,
175                updated_at: r.updated_at,
176            })
177            .collect())
178    }
179
180    pub async fn count_running_services(&self, module_name: &str) -> DatabaseResult<usize> {
181        let row = sqlx::query!(
182            r#"SELECT COUNT(*) as "count!" FROM services WHERE module_name = $1 AND status = 'running'"#,
183            module_name
184        )
185        .fetch_one(&*self.pool)
186        .await?;
187        Ok(usize::try_from(row.count).unwrap_or(0))
188    }
189
190    pub async fn mark_service_crashed(&self, service_name: &str) -> DatabaseResult<()> {
191        sqlx::query!(
192            r#"UPDATE services SET status = 'error', pid = NULL, updated_at = CURRENT_TIMESTAMP WHERE name = $1"#,
193            service_name
194        )
195        .execute(&*self.write_pool)
196        .await?;
197        Ok(())
198    }
199
200    pub async fn update_service_stopped(&self, service_name: &str) -> DatabaseResult<()> {
201        sqlx::query!(
202            r#"UPDATE services SET status = 'stopped', pid = NULL, updated_at = CURRENT_TIMESTAMP WHERE name = $1"#,
203            service_name
204        )
205        .execute(&*self.write_pool)
206        .await?;
207        Ok(())
208    }
209
210    pub async fn list_running_services_with_pid(&self) -> DatabaseResult<Vec<ServiceConfig>> {
211        self.list_all_running_services().await
212    }
213
214    pub async fn list_services_by_type(
215        &self,
216        module_name: &str,
217    ) -> DatabaseResult<Vec<ServiceConfig>> {
218        let rows = sqlx::query!(
219            r#"
220            SELECT name, module_name, status, pid, port, binary_mtime,
221                   created_at::text as "created_at!", updated_at::text as "updated_at!"
222            FROM services
223            WHERE module_name = $1
224            ORDER BY name
225            "#,
226            module_name
227        )
228        .fetch_all(&*self.pool)
229        .await?;
230        Ok(rows
231            .into_iter()
232            .map(|r| ServiceConfig {
233                name: r.name,
234                module_name: r.module_name,
235                status: r.status,
236                pid: r.pid,
237                port: r.port,
238                binary_mtime: r.binary_mtime,
239                created_at: r.created_at,
240                updated_at: r.updated_at,
241            })
242            .collect())
243    }
244
245    pub async fn cleanup_stale_entries(&self) -> DatabaseResult<u64> {
246        let result = sqlx::query!(
247            r#"
248            DELETE FROM services
249            WHERE status IN ('error', 'crashed')
250               OR (status = 'running' AND pid IS NULL)
251            "#
252        )
253        .execute(&*self.write_pool)
254        .await?;
255        Ok(result.rows_affected())
256    }
257}
systemprompt_database/repository/service/repo.rs

systemprompt_database/repository/service/
repo.rs
