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) instance_id: InstanceId,
27}
28
29impl ServiceRepository {
30    pub fn new(db: &DbPool, instance_id: InstanceId) -> Self {
31        let pool = db.pool();
32        let write_pool = db.write_pool();
33        Self {
34            pool,
35            write_pool,
36            instance_id,
37        }
38    }
39
40    pub const fn instance_id(&self) -> &InstanceId {
41        &self.instance_id
42    }
43
44    pub async fn find_service_by_name(
45        &self,
46        name: &ServiceName,
47    ) -> DatabaseResult<Option<ServiceConfig>> {
48        let row = sqlx::query_as!(
49            ServiceRow,
50            r#"
51            SELECT instance_id as "instance_id: InstanceId", name, module_name, status, pid, port, binary_mtime,
52                   heartbeat_at::text as "heartbeat_at!",
53                   created_at::text as "created_at!", updated_at::text as "updated_at!"
54            FROM services
55            WHERE instance_id = $1 AND name = $2
56            "#,
57            self.instance_id.as_str(),
58            name.as_str()
59        )
60        .fetch_optional(&*self.pool)
61        .await?;
62        row.map(ServiceRow::into_config).transpose()
63    }
64
65    pub async fn create_service(&self, input: CreateServiceInput<'_>) -> DatabaseResult<()> {
66        let port_i32 = i32::from(input.port);
67        sqlx::query!(
68            r#"
69            INSERT INTO services (instance_id, name, module_name, status, port, binary_mtime)
70            VALUES ($1, $2, $3, $4, $5, $6)
71            ON CONFLICT (instance_id, name) DO UPDATE SET
72              module_name = EXCLUDED.module_name,
73              status = EXCLUDED.status,
74              port = EXCLUDED.port,
75              binary_mtime = EXCLUDED.binary_mtime,
76              heartbeat_at = CURRENT_TIMESTAMP,
77              updated_at = CURRENT_TIMESTAMP
78            "#,
79            self.instance_id.as_str(),
80            input.name.as_str(),
81            input.module_name.as_str(),
82            input.status.as_str(),
83            port_i32,
84            input.binary_mtime
85        )
86        .execute(&*self.write_pool)
87        .await?;
88        Ok(())
89    }
90
91    pub async fn upsert_service_process(
92        &self,
93        input: UpsertServiceProcessInput<'_>,
94    ) -> DatabaseResult<()> {
95        let port_i32 = i32::from(input.port);
96        sqlx::query!(
97            r#"
98            INSERT INTO services (instance_id, name, module_name, pid, port, status, updated_at)
99            VALUES ($1, $2, $3, $4, $5, $6, CURRENT_TIMESTAMP)
100            ON CONFLICT (instance_id, name) DO UPDATE SET
101              pid = EXCLUDED.pid,
102              port = EXCLUDED.port,
103              status = EXCLUDED.status,
104              heartbeat_at = CURRENT_TIMESTAMP,
105              updated_at = CURRENT_TIMESTAMP
106            "#,
107            self.instance_id.as_str(),
108            input.name.as_str(),
109            input.module_name.as_str(),
110            input.pid,
111            port_i32,
112            input.status.as_str()
113        )
114        .execute(&*self.write_pool)
115        .await?;
116        Ok(())
117    }
118
119    pub async fn update_service_status(
120        &self,
121        service_name: &ServiceName,
122        status: ServiceStatus,
123    ) -> DatabaseResult<()> {
124        sqlx::query!(
125            r#"UPDATE services SET status = $1, updated_at = CURRENT_TIMESTAMP
126               WHERE instance_id = $2 AND name = $3"#,
127            status.as_str(),
128            self.instance_id.as_str(),
129            service_name.as_str()
130        )
131        .execute(&*self.write_pool)
132        .await?;
133        Ok(())
134    }
135
136    pub async fn delete_service(&self, service_name: &ServiceName) -> DatabaseResult<()> {
137        sqlx::query!(
138            r#"DELETE FROM services WHERE instance_id = $1 AND name = $2"#,
139            self.instance_id.as_str(),
140            service_name.as_str()
141        )
142        .execute(&*self.write_pool)
143        .await?;
144        Ok(())
145    }
146
147    pub async fn update_service_pid(
148        &self,
149        service_name: &ServiceName,
150        pid: i32,
151    ) -> DatabaseResult<()> {
152        sqlx::query!(
153            r#"UPDATE services SET pid = $1, updated_at = CURRENT_TIMESTAMP
154               WHERE instance_id = $2 AND name = $3"#,
155            pid,
156            self.instance_id.as_str(),
157            service_name.as_str()
158        )
159        .execute(&*self.write_pool)
160        .await?;
161        Ok(())
162    }
163
164    pub async fn update_service_port(
165        &self,
166        service_name: &ServiceName,
167        port: u16,
168    ) -> DatabaseResult<()> {
169        let port_i32 = i32::from(port);
170        sqlx::query!(
171            r#"UPDATE services SET port = $1, updated_at = CURRENT_TIMESTAMP
172               WHERE instance_id = $2 AND name = $3"#,
173            port_i32,
174            self.instance_id.as_str(),
175            service_name.as_str()
176        )
177        .execute(&*self.write_pool)
178        .await?;
179        Ok(())
180    }
181
182    pub async fn clear_service_pid(&self, service_name: &ServiceName) -> DatabaseResult<()> {
183        sqlx::query!(
184            r#"UPDATE services SET pid = NULL, updated_at = CURRENT_TIMESTAMP
185               WHERE instance_id = $1 AND name = $2"#,
186            self.instance_id.as_str(),
187            service_name.as_str()
188        )
189        .execute(&*self.write_pool)
190        .await?;
191        Ok(())
192    }
193
194    pub async fn mark_service_crashed(&self, service_name: &ServiceName) -> DatabaseResult<()> {
195        sqlx::query!(
196            r#"UPDATE services SET status = 'error', pid = NULL, updated_at = CURRENT_TIMESTAMP
197               WHERE instance_id = $1 AND name = $2"#,
198            self.instance_id.as_str(),
199            service_name.as_str()
200        )
201        .execute(&*self.write_pool)
202        .await?;
203        Ok(())
204    }
205
206    pub async fn update_service_stopped(&self, service_name: &ServiceName) -> DatabaseResult<()> {
207        sqlx::query!(
208            r#"UPDATE services SET status = 'stopped', pid = NULL, updated_at = CURRENT_TIMESTAMP
209               WHERE instance_id = $1 AND name = $2"#,
210            self.instance_id.as_str(),
211            service_name.as_str()
212        )
213        .execute(&*self.write_pool)
214        .await?;
215        Ok(())
216    }
217}
systemprompt_database/repository/service/repo.rs

systemprompt_database/repository/service/
repo.rs
