Skip to main content

systemprompt_database/lifecycle/
validation.rs

1//! Pre-flight validation helpers used by the boot path and tests.
2//!
3//! Copyright (c) systemprompt.io — Business Source License 1.1.
4//! See <https://systemprompt.io> for licensing details.
5
6use crate::error::{DatabaseResult, RepositoryError};
7use crate::services::{Database, DatabaseProvider};
8
9pub async fn validate_database_connection(db: &dyn DatabaseProvider) -> DatabaseResult<()> {
10    db.test_connection()
11        .await
12        .map_err(|e| RepositoryError::Connection(Box::new(e)))
13}
14
15pub async fn validate_write_pool_is_primary(db: &Database) -> DatabaseResult<()> {
16    let result = db
17        .write()
18        .query_raw(&"SELECT pg_is_in_recovery() as in_recovery")
19        .await?;
20
21    let in_recovery = result
22        .first()
23        .and_then(|row| row.get("in_recovery"))
24        .and_then(serde_json::Value::as_bool)
25        .ok_or_else(|| {
26            RepositoryError::Internal(
27                "Failed to determine whether the write pool is a primary".to_owned(),
28            )
29        })?;
30
31    if !in_recovery {
32        return Ok(());
33    }
34
35    Err(RepositoryError::invalid_state(if db.has_write_pool() {
36        "`database_write_url` points at a read-only standby. Writes, migrations and \
37         LISTEN/NOTIFY all require the primary — point it at the primary and restart"
38    } else {
39        "`database_url` points at a read-only standby and no `database_write_url` is set, so \
40         the write pool falls back to it. Set `database_write_url` (or `DATABASE_WRITE_URL` \
41         with the env secrets source) to the primary and restart"
42    }))
43}
44
45#[derive(Debug, Clone, Copy, PartialEq)]
46pub struct ReplicaStatus {
47    pub in_recovery: bool,
48    pub replay_lag_secs: Option<f64>,
49}
50
51pub async fn replica_status(db: &dyn DatabaseProvider) -> DatabaseResult<ReplicaStatus> {
52    let result = db
53        .query_raw(
54            &"SELECT pg_is_in_recovery() AS in_recovery, CASE WHEN pg_is_in_recovery() THEN \
55              EXTRACT(EPOCH FROM (now() - pg_last_xact_replay_timestamp()))::double precision \
56              ELSE NULL END AS lag_secs",
57        )
58        .await?;
59    let row = result.first().ok_or_else(|| {
60        RepositoryError::Internal("replica status probe returned no row".to_owned())
61    })?;
62    let in_recovery = row
63        .get("in_recovery")
64        .and_then(serde_json::Value::as_bool)
65        .ok_or_else(|| {
66            RepositoryError::Internal("replica status probe lacks in_recovery".to_owned())
67        })?;
68    let replay_lag_secs = row.get("lag_secs").and_then(serde_json::Value::as_f64);
69    Ok(ReplicaStatus {
70        in_recovery,
71        replay_lag_secs,
72    })
73}
74
75pub async fn validate_table_exists(
76    db: &dyn DatabaseProvider,
77    table_name: &str,
78) -> DatabaseResult<bool> {
79    let result = db
80        .query_raw_with(
81            &"SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_schema = \
82              'public' AND table_name = $1) as exists",
83            &[&table_name],
84        )
85        .await?;
86
87    result
88        .first()
89        .and_then(|row| row.get("exists"))
90        .and_then(serde_json::Value::as_bool)
91        .ok_or_else(|| {
92            RepositoryError::Internal(format!(
93                "Failed to check table existence for '{table_name}'"
94            ))
95        })
96}
97
98pub async fn validate_column_exists(
99    db: &dyn DatabaseProvider,
100    table_name: &str,
101    column_name: &str,
102) -> DatabaseResult<bool> {
103    let result = db
104        .query_raw_with(
105            &"SELECT EXISTS (SELECT 1 FROM information_schema.columns WHERE table_schema = \
106              'public' AND table_name = $1 AND column_name = $2) as exists",
107            &[&table_name, &column_name],
108        )
109        .await?;
110
111    result
112        .first()
113        .and_then(|row| row.get("exists"))
114        .and_then(serde_json::Value::as_bool)
115        .ok_or_else(|| {
116            RepositoryError::Internal(format!(
117                "Failed to check column existence for '{table_name}.{column_name}'"
118            ))
119        })
120}