systemprompt-database 0.79.0

PostgreSQL infrastructure for systemprompt.io AI governance. SQLx-backed pool, generic repository traits, and compile-time query verification. Part of the systemprompt.io AI governance pipeline.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96

//! Per-replica role presence: every replica records its `server.role` with
//! each registry heartbeat, so the live replica count per role is readable
//! cluster-wide, including for roles that register no process, and the live
//! members of one role are listable for spreading work across them.
//!
//! Copyright (c) systemprompt.io — Business Source License 1.1.
//! See <https://systemprompt.io> for licensing details.

use systemprompt_identifiers::InstanceId;
use systemprompt_manifest::profile::NodeRole;

use super::liveness::INSTANCE_DEAD_AFTER_SECS;
use super::repo::ServiceRepository;
use crate::error::DatabaseResult;

/// The number of live replicas currently heartbeating one role.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct RoleCount {
    pub role: String,
    pub instances: i64,
}

impl ServiceRepository {
    pub async fn touch_presence(&self, role: NodeRole) -> DatabaseResult<()> {
        sqlx::query!(
            r#"
            INSERT INTO instance_presence (instance_id, role, heartbeat_at)
            VALUES ($1, $2, CURRENT_TIMESTAMP)
            ON CONFLICT (instance_id) DO UPDATE SET
                role = EXCLUDED.role,
                heartbeat_at = EXCLUDED.heartbeat_at
            "#,
            self.instance_id.as_str(),
            role.as_str()
        )
        .execute(&*self.write_pool)
        .await?;
        Ok(())
    }

    pub async fn count_live_instances_by_role(&self) -> DatabaseResult<Vec<RoleCount>> {
        let rows = sqlx::query_as!(
            RoleCount,
            r#"
            SELECT role, COUNT(*) AS "instances!"
            FROM instance_presence
            WHERE heartbeat_at
                  >= CURRENT_TIMESTAMP - make_interval(secs => $1::double precision)
            GROUP BY role
            ORDER BY role
            "#,
            INSTANCE_DEAD_AFTER_SECS as f64
        )
        .fetch_all(&*self.pool)
        .await?;
        Ok(rows)
    }

    pub async fn list_reachable_instances(
        &self,
        role: NodeRole,
        fresh_within_secs: i64,
    ) -> DatabaseResult<Vec<InstanceId>> {
        let rows = sqlx::query_scalar!(
            r#"
            SELECT p.instance_id AS "instance_id!: InstanceId"
            FROM instance_presence p
            JOIN instance_endpoints e ON e.instance_id = p.instance_id
            WHERE p.role = $1
              AND p.heartbeat_at
                  >= CURRENT_TIMESTAMP - make_interval(secs => $2::double precision)
              AND e.heartbeat_at
                  >= CURRENT_TIMESTAMP - make_interval(secs => $2::double precision)
            ORDER BY p.instance_id
            "#,
            role.as_str(),
            fresh_within_secs as f64
        )
        .fetch_all(&*self.write_pool)
        .await?;
        Ok(rows)
    }

    pub async fn delete_stale_presence(&self, older_than_secs: i64) -> DatabaseResult<u64> {
        let result = sqlx::query!(
            r#"
            DELETE FROM instance_presence
            WHERE heartbeat_at < CURRENT_TIMESTAMP - make_interval(secs => $1::double precision)
            "#,
            older_than_secs as f64
        )
        .execute(&*self.write_pool)
        .await?;
        Ok(result.rows_affected())
    }
}