systemprompt_database/repository/service/
repo.rs

1//! Async repository over the `services` registry table.
2//!
3//! Rows are keyed by `(instance_id, name)`: every replica registers, judges
4//! and reaps only the processes it spawned itself. The single cross-instance
5//! statement is `delete_dead_instances` (in `maintenance`), which reaps rows
6//! whose heartbeat stopped, so a replica that vanished without cleanup is
7//! garbage-collected by the scheduler rather than by the next node to boot.
8//!
9//! Copyright (c) systemprompt.io — Business Source License 1.1.
10//! See <https://systemprompt.io> for licensing details.
11
12use std::sync::Arc;
13
14use sqlx::PgPool;
15use systemprompt_identifiers::{InstanceId, ServiceName};
16use systemprompt_manifest::services::ServiceStatus;
17
18use super::model::{CreateServiceInput, ServiceConfig, ServiceRow, UpsertServiceProcessInput};
19use crate::DbPool;
20use crate::error::DatabaseResult;
21
22#[derive(Debug, Clone)]
23pub struct ServiceRepository {
24    pub(super) pool: Arc<PgPool>,
25    pub(super) write_pool: Arc<PgPool>,
26    pub(super) session_pool: Arc<PgPool>,
27    pub(super) instance_id: InstanceId,
28}
29
30impl ServiceRepository {
31    pub fn new(db: &DbPool, instance_id: InstanceId) -> Self {
32        let pool = db.pool();
33        let write_pool = db.write_pool();
34        let session_pool = db.session_pool();
35        Self {
36            pool,
37            write_pool,
38            session_pool,
39            instance_id,
40        }
41    }
42
43    pub const fn instance_id(&self) -> &InstanceId {
44        &self.instance_id
45    }
46
47    pub async fn find_service_by_name(
48        &self,
49        name: &ServiceName,
50    ) -> DatabaseResult<Option<ServiceConfig>> {
51        let row = sqlx::query_as!(
52            ServiceRow,
53            r#"
54            SELECT instance_id as "instance_id: InstanceId", name, module_name, status, pid, port, binary_mtime,
55                   heartbeat_at::text as "heartbeat_at!",
56                   created_at::text as "created_at!", updated_at::text as "updated_at!"
57            FROM services
58            WHERE instance_id = $1 AND name = $2
59            "#,
60            self.instance_id.as_str(),
61            name.as_str()
62        )
63        .fetch_optional(&*self.pool)
64        .await?;
65        row.map(ServiceRow::into_config).transpose()
66    }
67
68    pub async fn create_service(&self, input: CreateServiceInput<'_>) -> DatabaseResult<()> {
69        let port_i32 = i32::from(input.port);
70        sqlx::query!(
71            r#"
72            INSERT INTO services (instance_id, name, module_name, status, port, binary_mtime)
73            VALUES ($1, $2, $3, $4, $5, $6)
74            ON CONFLICT (instance_id, name) DO UPDATE SET
75              module_name = EXCLUDED.module_name,
76              status = EXCLUDED.status,
77              port = EXCLUDED.port,
78              binary_mtime = EXCLUDED.binary_mtime,
79              heartbeat_at = CURRENT_TIMESTAMP,
80              updated_at = CURRENT_TIMESTAMP
81            "#,
82            self.instance_id.as_str(),
83            input.name.as_str(),
84            input.module_name.as_str(),
85            input.status.as_str(),
86            port_i32,
87            input.binary_mtime
88        )
89        .execute(&*self.write_pool)
90        .await?;
91        Ok(())
92    }
93
94    pub async fn upsert_service_process(
95        &self,
96        input: UpsertServiceProcessInput<'_>,
97    ) -> DatabaseResult<()> {
98        let port_i32 = i32::from(input.port);
99        sqlx::query!(
100            r#"
101            INSERT INTO services (instance_id, name, module_name, pid, port, status, updated_at)
102            VALUES ($1, $2, $3, $4, $5, $6, CURRENT_TIMESTAMP)
103            ON CONFLICT (instance_id, name) DO UPDATE SET
104              pid = EXCLUDED.pid,
105              port = EXCLUDED.port,
106              status = EXCLUDED.status,
107              heartbeat_at = CURRENT_TIMESTAMP,
108              updated_at = CURRENT_TIMESTAMP
109            "#,
110            self.instance_id.as_str(),
111            input.name.as_str(),
112            input.module_name.as_str(),
113            input.pid,
114            port_i32,
115            input.status.as_str()
116        )
117        .execute(&*self.write_pool)
118        .await?;
119        Ok(())
120    }
121
122    pub async fn update_service_status(
123        &self,
124        service_name: &ServiceName,
125        status: ServiceStatus,
126    ) -> DatabaseResult<()> {
127        sqlx::query!(
128            r#"UPDATE services SET status = $1, updated_at = CURRENT_TIMESTAMP
129               WHERE instance_id = $2 AND name = $3"#,
130            status.as_str(),
131            self.instance_id.as_str(),
132            service_name.as_str()
133        )
134        .execute(&*self.write_pool)
135        .await?;
136        Ok(())
137    }
138
139    pub async fn delete_service(&self, service_name: &ServiceName) -> DatabaseResult<()> {
140        sqlx::query!(
141            r#"DELETE FROM services WHERE instance_id = $1 AND name = $2"#,
142            self.instance_id.as_str(),
143            service_name.as_str()
144        )
145        .execute(&*self.write_pool)
146        .await?;
147        Ok(())
148    }
149
150    pub async fn update_service_pid(
151        &self,
152        service_name: &ServiceName,
153        pid: i32,
154    ) -> DatabaseResult<()> {
155        sqlx::query!(
156            r#"UPDATE services SET pid = $1, updated_at = CURRENT_TIMESTAMP
157               WHERE instance_id = $2 AND name = $3"#,
158            pid,
159            self.instance_id.as_str(),
160            service_name.as_str()
161        )
162        .execute(&*self.write_pool)
163        .await?;
164        Ok(())
165    }
166
167    pub async fn update_service_port(
168        &self,
169        service_name: &ServiceName,
170        port: u16,
171    ) -> DatabaseResult<()> {
172        let port_i32 = i32::from(port);
173        sqlx::query!(
174            r#"UPDATE services SET port = $1, updated_at = CURRENT_TIMESTAMP
175               WHERE instance_id = $2 AND name = $3"#,
176            port_i32,
177            self.instance_id.as_str(),
178            service_name.as_str()
179        )
180        .execute(&*self.write_pool)
181        .await?;
182        Ok(())
183    }
184
185    pub async fn clear_service_pid(&self, service_name: &ServiceName) -> DatabaseResult<()> {
186        sqlx::query!(
187            r#"UPDATE services SET pid = NULL, updated_at = CURRENT_TIMESTAMP
188               WHERE instance_id = $1 AND name = $2"#,
189            self.instance_id.as_str(),
190            service_name.as_str()
191        )
192        .execute(&*self.write_pool)
193        .await?;
194        Ok(())
195    }
196
197    pub async fn mark_service_crashed(&self, service_name: &ServiceName) -> DatabaseResult<()> {
198        sqlx::query!(
199            r#"UPDATE services SET status = 'error', pid = NULL, updated_at = CURRENT_TIMESTAMP
200               WHERE instance_id = $1 AND name = $2"#,
201            self.instance_id.as_str(),
202            service_name.as_str()
203        )
204        .execute(&*self.write_pool)
205        .await?;
206        Ok(())
207    }
208
209    pub async fn update_service_stopped(&self, service_name: &ServiceName) -> DatabaseResult<()> {
210        sqlx::query!(
211            r#"UPDATE services SET status = 'stopped', pid = NULL, updated_at = CURRENT_TIMESTAMP
212               WHERE instance_id = $1 AND name = $2"#,
213            self.instance_id.as_str(),
214            service_name.as_str()
215        )
216        .execute(&*self.write_pool)
217        .await?;
218        Ok(())
219    }
220}
systemprompt_database/repository/service/repo.rs

systemprompt_database/repository/service/
repo.rs
