union_square 0.1.0

A proxy/wire-tap service for making LLM calls and recording everything that happens in a session for later analysis and test-case extraction
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

use crate::{Error, Result};
use sqlx::{PgPool, Row};

/// Database connection pool wrapper
pub struct Database {
    pool: PgPool,
}

impl Database {
    pub fn new(pool: PgPool) -> Self {
        Self { pool }
    }

    pub fn pool(&self) -> &PgPool {
        &self.pool
    }

    /// Health check for the database connection
    pub async fn health_check(&self) -> Result<()> {
        let row = sqlx::query("SELECT 1 as health_check")
            .fetch_one(&self.pool)
            .await?;

        let health_check: i32 = row.try_get("health_check")?;

        if health_check == 1 {
            Ok(())
        } else {
            Err(Error::application("Database health check failed"))
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    #[ignore = "requires database connection"]
    async fn test_database_health_check() {
        let pool = PgPool::connect("postgres://postgres:password@localhost:5432/union_square")
            .await
            .expect("Failed to connect to database");

        let db = Database::new(pool);
        let result = db.health_check().await;
        assert!(result.is_ok());
    }
}