systemprompt_cli/commands/admin/setup/
common.rs

1//! Shared setup types: Postgres connection configuration.
2//!
3//! Copyright (c) systemprompt.io — Business Source License 1.1.
4//! See <https://systemprompt.io> for licensing details.
5
6use anyhow::Result;
7use rand::distr::Alphanumeric;
8use rand::{RngExt, rng};
9use sqlx::postgres::PgPoolOptions;
10use std::net::ToSocketAddrs;
11use std::time::Duration;
12use systemprompt_logging::CliService;
13
14#[derive(Debug, Clone)]
15pub struct PostgresConfig {
16    pub host: String,
17    pub port: u16,
18    pub user: String,
19    pub password: String,
20    pub database: String,
21}
22
23impl PostgresConfig {
24    pub fn database_url(&self) -> String {
25        format!(
26            "postgres://{}:{}@{}:{}/{}",
27            self.user, self.password, self.host, self.port, self.database
28        )
29    }
30}
31
32pub fn generate_password() -> String {
33    let mut rng = rng();
34    (0..16)
35        .map(|_| rng.sample(Alphanumeric))
36        .map(char::from)
37        .collect()
38}
39
40pub fn detect_postgresql(host: &str, port: u16) -> bool {
41    let addr = format!("{}:{}", host, port);
42    let socket_addrs = match addr.to_socket_addrs() {
43        Ok(addrs) => addrs.collect::<Vec<_>>(),
44        Err(e) => {
45            tracing::debug!(host = %host, port = %port, error = %e, "Failed to resolve socket address");
46            return false;
47        },
48    };
49
50    for socket_addr in socket_addrs {
51        if std::net::TcpStream::connect_timeout(&socket_addr, Duration::from_secs(3)).is_ok() {
52            return true;
53        }
54    }
55
56    false
57}
58
59pub async fn test_connection(config: &PostgresConfig) -> bool {
60    let Ok(pool) = PgPoolOptions::new()
61        .max_connections(1)
62        .acquire_timeout(Duration::from_secs(5))
63        .connect(&config.database_url())
64        .await
65    else {
66        return false;
67    };
68
69    let result = sqlx::query_scalar!("SELECT 1 as one")
70        .fetch_one(&pool)
71        .await
72        .is_ok();
73    pool.close().await;
74    result
75}
76
77pub async fn enable_extensions(config: &PostgresConfig) -> Result<()> {
78    let pool = match PgPoolOptions::new()
79        .max_connections(1)
80        .acquire_timeout(Duration::from_secs(5))
81        .connect(&config.database_url())
82        .await
83    {
84        Ok(pool) => pool,
85        Err(e) => {
86            CliService::warning(&format!("Could not enable extensions: {}", e));
87            return Ok(());
88        },
89    };
90
91    let extensions = ["uuid-ossp", "unaccent", "pg_trgm"];
92
93    for ext in extensions {
94        let sql = format!("CREATE EXTENSION IF NOT EXISTS \"{}\"", ext);
95        if let Err(e) = sqlx::query(sqlx::AssertSqlSafe(sql)).execute(&pool).await {
96            CliService::warning(&format!("Could not create extension '{}': {}", ext, e));
97        }
98    }
99
100    pool.close().await;
101    Ok(())
102}
systemprompt_cli/commands/admin/setup/common.rs

systemprompt_cli/commands/admin/setup/
common.rs
