lean_ctx/core/providers/
postgres.rs

1//! PostgreSQL provider — database schema introspection via `psql`.
2//!
3//! Extracts table/column definitions from `information_schema` to make
4//! database structure available as context. Uses `psql` CLI to avoid
5//! adding a native PG driver dependency.
6//!
7//! Configuration via environment variables:
8//!   - `DATABASE_URL`: Full connection string (e.g., "postgres://user:pass@host/db")
9//!   - Or individual: `PGHOST`, `PGPORT`, `PGDATABASE`, `PGUSER`, `PGPASSWORD`
10
11use crate::core::providers::{ContextProvider, ProviderItem, ProviderParams, ProviderResult};
12
13pub struct PostgresProvider {
14    available: bool,
15}
16
17impl Default for PostgresProvider {
18    fn default() -> Self {
19        Self::new()
20    }
21}
22
23impl PostgresProvider {
24    pub fn new() -> Self {
25        let available =
26            std::env::var("DATABASE_URL").is_ok() || std::env::var("PGDATABASE").is_ok();
27        Self { available }
28    }
29}
30
31impl ContextProvider for PostgresProvider {
32    fn id(&self) -> &'static str {
33        "postgres"
34    }
35
36    fn display_name(&self) -> &'static str {
37        "PostgreSQL"
38    }
39
40    fn supported_actions(&self) -> &[&str] {
41        &["schemas", "tables"]
42    }
43
44    fn execute(&self, action: &str, params: &ProviderParams) -> Result<ProviderResult, String> {
45        if !self.available {
46            return Err("PostgreSQL not configured (need DATABASE_URL or PGDATABASE)".into());
47        }
48        match action {
49            "schemas" | "tables" => list_tables(params),
50            _ => Err(format!("Unsupported action: {action}")),
51        }
52    }
53
54    fn cache_ttl_secs(&self) -> u64 {
55        300
56    }
57
58    fn requires_auth(&self) -> bool {
59        true
60    }
61
62    fn is_available(&self) -> bool {
63        self.available
64    }
65}
66
67fn list_tables(params: &ProviderParams) -> Result<ProviderResult, String> {
68    let schema = params.state.as_deref().unwrap_or("public");
69    let limit = params.limit.unwrap_or(50);
70
71    let query = format!(
72        "SELECT table_name, column_name, data_type, is_nullable \
73         FROM information_schema.columns \
74         WHERE table_schema = '{schema}' \
75         ORDER BY table_name, ordinal_position \
76         LIMIT {limit_cols};",
77        limit_cols = limit * 20, // ~20 columns per table avg
78    );
79
80    let mut cmd = std::process::Command::new("psql");
81
82    if let Ok(url) = std::env::var("DATABASE_URL") {
83        cmd.arg(&url);
84    }
85
86    let output = cmd
87        .args(["-t", "-A", "-F", "|", "-c", &query])
88        .output()
89        .map_err(|e| format!("Failed to run psql: {e}"))?;
90
91    if !output.status.success() {
92        let stderr = String::from_utf8_lossy(&output.stderr);
93        return Err(format!("psql error: {stderr}"));
94    }
95
96    let stdout = String::from_utf8_lossy(&output.stdout);
97    let mut tables: std::collections::BTreeMap<String, Vec<String>> =
98        std::collections::BTreeMap::new();
99
100    for line in stdout.lines() {
101        let parts: Vec<&str> = line.split('|').collect();
102        if parts.len() >= 3 {
103            let table = parts[0].trim();
104            let col = parts[1].trim();
105            let dtype = parts[2].trim();
106            let nullable = parts.get(3).map_or("", |s| s.trim());
107
108            let null_marker = if nullable == "YES" { "?" } else { "" };
109            tables
110                .entry(table.to_string())
111                .or_default()
112                .push(format!("  {col}: {dtype}{null_marker}"));
113        }
114    }
115
116    let items: Vec<ProviderItem> = tables
117        .iter()
118        .take(limit)
119        .map(|(table, columns)| {
120            let body = format!("{schema}.{table}\n{}", columns.join("\n"));
121            ProviderItem {
122                id: table.clone(),
123                title: format!("{schema}.{table}"),
124                state: Some("active".into()),
125                author: None,
126                created_at: None,
127                updated_at: None,
128                url: None,
129                labels: vec![schema.to_string()],
130                body: Some(body),
131            }
132        })
133        .collect();
134
135    Ok(ProviderResult {
136        provider: "postgres".into(),
137        resource_type: "schemas".into(),
138        items,
139        total_count: Some(tables.len()),
140        truncated: tables.len() > limit,
141    })
142}
143
144#[cfg(test)]
145mod tests {
146    use super::*;
147
148    #[test]
149    fn postgres_provider_unavailable_without_env() {
150        std::env::remove_var("DATABASE_URL");
151        std::env::remove_var("PGDATABASE");
152
153        let provider = PostgresProvider::new();
154        assert!(!provider.is_available());
155        assert_eq!(provider.id(), "postgres");
156        assert!(provider.requires_auth());
157    }
158
159    #[test]
160    fn postgres_provider_supported_actions() {
161        let provider = PostgresProvider::new();
162        assert!(provider.supported_actions().contains(&"schemas"));
163        assert!(provider.supported_actions().contains(&"tables"));
164    }
165}
lean_ctx/core/providers/postgres.rs

lean_ctx/core/providers/
postgres.rs
