lean_ctx/core/providers/
postgres.rs

1// SPDX-License-Identifier: Apache-2.0
2//! PostgreSQL provider — database schema introspection via `psql`.
3//!
4//! Extracts table/column definitions from `information_schema` to make
5//! database structure available as context. Uses `psql` CLI to avoid
6//! adding a native PG driver dependency.
7//!
8//! Configuration via environment variables:
9//!   - `DATABASE_URL`: Full connection string (e.g., "postgres://user:pass@host/db")
10//!   - Or individual: `PGHOST`, `PGPORT`, `PGDATABASE`, `PGUSER`, `PGPASSWORD`
11
12use crate::core::providers::{ContextProvider, ProviderItem, ProviderParams, ProviderResult};
13
14const MAX_TABLES: usize = 100;
15const MAX_STDOUT_BYTES: usize = 1_048_576;
16const MAX_STDERR_BYTES: usize = 16_384;
17const QUERY_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(10);
18
19pub struct PostgresProvider {
20    available: bool,
21}
22
23impl Default for PostgresProvider {
24    fn default() -> Self {
25        Self::new()
26    }
27}
28
29impl PostgresProvider {
30    pub fn new() -> Self {
31        let available = super::credentials::var("DATABASE_URL").is_ok()
32            || super::credentials::var("PGDATABASE").is_ok();
33        Self { available }
34    }
35}
36
37impl ContextProvider for PostgresProvider {
38    fn id(&self) -> &'static str {
39        "postgres"
40    }
41
42    fn display_name(&self) -> &'static str {
43        "PostgreSQL"
44    }
45
46    fn supported_actions(&self) -> &[&str] {
47        &["schemas", "tables"]
48    }
49
50    fn execute(&self, action: &str, params: &ProviderParams) -> Result<ProviderResult, String> {
51        if !self.available {
52            return Err("PostgreSQL not configured (need DATABASE_URL or PGDATABASE)".into());
53        }
54        match action {
55            "schemas" | "tables" => list_tables(params),
56            _ => Err(format!("Unsupported action: {action}")),
57        }
58    }
59
60    fn cache_ttl_secs(&self) -> u64 {
61        300
62    }
63
64    fn requires_auth(&self) -> bool {
65        true
66    }
67
68    fn is_available(&self) -> bool {
69        self.available
70    }
71}
72
73/// Validates a PostgreSQL identifier before it is interpolated into SQL.
74///
75/// The schema name comes from provider params (agent-controlled), and the query
76/// is executed via `psql -c`, so parameterized queries are not available.
77/// A strict identifier whitelist (`[A-Za-z_][A-Za-z0-9_$]*`, max 63 bytes — the
78/// PostgreSQL `NAMEDATALEN` limit) makes injection impossible.
79fn validate_pg_identifier(name: &str) -> Result<(), String> {
80    let valid_start = name
81        .chars()
82        .next()
83        .is_some_and(|c| c.is_ascii_alphabetic() || c == '_');
84    let valid_rest = name
85        .chars()
86        .skip(1)
87        .all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '$');
88    if name.is_empty() || name.len() > 63 || !valid_start || !valid_rest {
89        return Err(format!(
90            "Invalid PostgreSQL schema identifier: {name:?} (allowed: [A-Za-z_][A-Za-z0-9_$]*, max 63 chars)"
91        ));
92    }
93    Ok(())
94}
95
96fn list_tables(params: &ProviderParams) -> Result<ProviderResult, String> {
97    let schema = params.state.as_deref().unwrap_or("public");
98    validate_pg_identifier(schema)?;
99    let limit = params.limit.unwrap_or(50);
100    if !(1..=MAX_TABLES).contains(&limit) {
101        return Err(format!(
102            "PostgreSQL table limit must be between 1 and {MAX_TABLES}"
103        ));
104    }
105    let column_limit = limit * 20; // bounded before multiplication
106
107    let query = format!(
108        "SELECT table_name, column_name, data_type, is_nullable \
109         FROM information_schema.columns \
110         WHERE table_schema = '{schema}' \
111         ORDER BY table_name, ordinal_position \
112         LIMIT {column_limit};",
113    );
114
115    let mut cmd = psql_command()?;
116    cmd.args([
117        "-X",
118        "-w",
119        "-v",
120        "ON_ERROR_STOP=1",
121        "-t",
122        "-A",
123        "-F",
124        "|",
125        "-c",
126        &query,
127    ])
128    .env("PGCLIENTENCODING", "UTF8")
129    .stdin(std::process::Stdio::null());
130    let stdout = run_psql(&mut cmd, QUERY_TIMEOUT)?;
131    parse_tables(&stdout, schema, limit, column_limit)
132}
133
134/// A `psql` command connected through `DATABASE_URL` (environment or store).
135/// The address is handed over only as `PG*` variables of the child, never as
136/// an argument: argv is visible to every user in the process list.
137pub(in crate::core) fn psql_command() -> Result<std::process::Command, String> {
138    match super::credentials::var("DATABASE_URL") {
139        Ok(url) => psql_command_for(&url),
140        Err(_) => Ok(std::process::Command::new("psql")),
141    }
142}
143
144/// A `psql` command connected to `url` through `PG*` variables only.
145pub(in crate::core) fn psql_command_for(url: &str) -> Result<std::process::Command, String> {
146    let mut command = std::process::Command::new("psql");
147    command.envs(connection_env(url)?);
148    Ok(command)
149}
150
151/// `postgres[ql]://[user[:password]@][host[:port][,…]][/dbname][?sslmode=…]`
152/// as libpq environment variables. Errors never quote the address.
153fn connection_env(url: &str) -> Result<Vec<(&'static str, String)>, String> {
154    let decode = |part: &str| {
155        urlencoding::decode(part)
156            .map(std::borrow::Cow::into_owned)
157            .map_err(|_| "DATABASE_URL has an invalid percent-encoding".to_string())
158    };
159    let rest = url
160        .trim()
161        .strip_prefix("postgresql://")
162        .or_else(|| url.trim().strip_prefix("postgres://"))
163        .ok_or_else(|| "DATABASE_URL must start with postgres:// or postgresql://".to_string())?;
164    let (rest, query) = rest.split_once('?').unwrap_or((rest, ""));
165    let (authority, dbname) = rest.split_once('/').unwrap_or((rest, ""));
166    let (userinfo, hosts) = authority
167        .rsplit_once('@')
168        .map_or((None, authority), |(u, h)| (Some(u), h));
169
170    let mut env = Vec::new();
171    if let Some(userinfo) = userinfo {
172        let (user, password) = userinfo
173            .split_once(':')
174            .map_or((userinfo, None), |(u, p)| (u, Some(p)));
175        if !user.is_empty() {
176            env.push(("PGUSER", decode(user)?));
177        }
178        if let Some(password) = password.filter(|p| !p.is_empty()) {
179            env.push(("PGPASSWORD", decode(password)?));
180        }
181    }
182    if !hosts.is_empty() {
183        let (mut names, mut ports) = (Vec::new(), Vec::new());
184        for host in hosts.split(',') {
185            // `[::1]:5432` keeps its brackets out of PGHOST.
186            let (name, port) = if let Some(v6) = host.strip_prefix('[') {
187                let (name, tail) = v6
188                    .split_once(']')
189                    .ok_or_else(|| "DATABASE_URL has an unclosed [host]".to_string())?;
190                (name, tail.strip_prefix(':').unwrap_or(""))
191            } else {
192                host.rsplit_once(':').unwrap_or((host, ""))
193            };
194            names.push(decode(name)?);
195            ports.push(port.to_string());
196        }
197        env.push(("PGHOST", names.join(",")));
198        if ports.iter().any(|p| !p.is_empty()) {
199            env.push(("PGPORT", ports.join(",")));
200        }
201    }
202    if !dbname.is_empty() {
203        env.push(("PGDATABASE", decode(dbname)?));
204    }
205    for (key, value) in query.split('&').filter_map(|pair| pair.split_once('=')) {
206        if key == "sslmode" {
207            env.push(("PGSSLMODE", decode(value)?));
208        }
209    }
210    Ok(env)
211}
212
213/// Shared bounded execution for catalog introspection and the pgvector backend.
214pub(in crate::core) fn run_psql(
215    command: &mut std::process::Command,
216    timeout: std::time::Duration,
217) -> Result<String, String> {
218    let captured = crate::core::process_capture::run_with_output_limits(
219        command,
220        Some(timeout),
221        MAX_STDOUT_BYTES,
222        MAX_STDERR_BYTES,
223    )
224    .map_err(|_| "PostgreSQL command failed or exceeded its output bound".to_owned())?;
225    if captured.timed_out {
226        return Err("PostgreSQL command exceeded its deadline".to_owned());
227    }
228    if captured.cancelled || !captured.output.status.success() {
229        // Database diagnostics can echo connection credentials, SQL or row data.
230        return Err("PostgreSQL command failed; database diagnostics withheld".to_owned());
231    }
232    String::from_utf8(captured.output.stdout)
233        .map_err(|_| "PostgreSQL returned invalid UTF-8".to_owned())
234}
235
236fn parse_tables(
237    stdout: &str,
238    schema: &str,
239    limit: usize,
240    column_limit: usize,
241) -> Result<ProviderResult, String> {
242    let mut tables: std::collections::BTreeMap<String, Vec<String>> =
243        std::collections::BTreeMap::new();
244    let mut column_count = 0;
245    for line in stdout.lines() {
246        if line.is_empty() {
247            continue;
248        }
249        let parts: Vec<&str> = line.split('|').collect();
250        if parts.len() == 4
251            && !parts[0].trim().is_empty()
252            && !parts[1].trim().is_empty()
253            && !parts[2].trim().is_empty()
254            && matches!(parts[3].trim(), "YES" | "NO")
255        {
256            column_count += 1;
257            let table = parts[0].trim();
258            let col = parts[1].trim();
259            let dtype = parts[2].trim();
260            let nullable = parts.get(3).map_or("", |s| s.trim());
261
262            let null_marker = if nullable == "YES" { "?" } else { "" };
263            tables
264                .entry(table.to_string())
265                .or_default()
266                .push(format!("  {col}: {dtype}{null_marker}"));
267        } else {
268            return Err("PostgreSQL returned an invalid schema record".to_owned());
269        }
270    }
271
272    let items: Vec<ProviderItem> = tables
273        .iter()
274        .take(limit)
275        .map(|(table, columns)| {
276            let body = format!("{schema}.{table}\n{}", columns.join("\n"));
277            ProviderItem {
278                id: table.clone(),
279                title: format!("{schema}.{table}"),
280                state: Some("active".into()),
281                author: None,
282                created_at: None,
283                updated_at: None,
284                url: None,
285                labels: vec![schema.to_string()],
286                body: Some(body),
287                ..Default::default()
288            }
289        })
290        .collect();
291
292    let truncated = column_count >= column_limit || tables.len() > limit;
293    Ok(ProviderResult {
294        provider: "postgres".into(),
295        resource_type: "schemas".into(),
296        items,
297        total_count: (!truncated).then_some(tables.len()),
298        truncated,
299    })
300}
301
302#[cfg(test)]
303mod tests {
304    use super::*;
305
306    #[test]
307    fn postgres_provider_unavailable_without_env() {
308        let _env_lock = crate::core::data_dir::test_env_lock();
309        crate::test_env::remove_var("DATABASE_URL");
310        crate::test_env::remove_var("PGDATABASE");
311
312        let provider = PostgresProvider::new();
313        assert!(!provider.is_available());
314        assert_eq!(provider.id(), "postgres");
315        assert!(provider.requires_auth());
316    }
317
318    #[test]
319    fn postgres_provider_supported_actions() {
320        let provider = PostgresProvider::new();
321        assert!(provider.supported_actions().contains(&"schemas"));
322        assert!(provider.supported_actions().contains(&"tables"));
323    }
324
325    // P0-5 (#417): schema names are agent-controlled and interpolated into SQL —
326    // only strict identifiers may pass.
327    #[test]
328    fn valid_pg_identifiers_pass() {
329        for ok in ["public", "my_schema", "_internal", "Schema1", "a$b"] {
330            assert!(validate_pg_identifier(ok).is_ok(), "{ok} should be valid");
331        }
332    }
333
334    #[test]
335    fn sql_injection_payloads_are_rejected() {
336        for evil in [
337            "public' UNION SELECT usename, passwd, '', '' FROM pg_shadow --",
338            "public'; DROP TABLE users; --",
339            "a\"b",
340            "a b",
341            "a;b",
342            "schema\n--",
343            "",
344            "1starts_with_digit",
345        ] {
346            assert!(
347                validate_pg_identifier(evil).is_err(),
348                "{evil:?} must be rejected"
349            );
350        }
351    }
352
353    #[test]
354    fn overlong_identifier_is_rejected() {
355        let too_long = "a".repeat(64);
356        assert!(validate_pg_identifier(&too_long).is_err());
357        let max_ok = "a".repeat(63);
358        assert!(validate_pg_identifier(&max_ok).is_ok());
359    }
360
361    #[test]
362    fn injection_via_params_state_fails_closed() {
363        let params = ProviderParams {
364            state: Some("public' OR '1'='1".into()),
365            ..Default::default()
366        };
367        let err = list_tables(&params).unwrap_err();
368        assert!(err.contains("Invalid PostgreSQL schema identifier"));
369    }
370
371    #[test]
372    fn table_limits_reject_overflow_and_empty_requests_before_dispatch() {
373        for limit in [0, MAX_TABLES + 1, usize::MAX] {
374            let params = ProviderParams {
375                limit: Some(limit),
376                ..Default::default()
377            };
378            assert!(list_tables(&params).unwrap_err().contains("table limit"));
379        }
380    }
381
382    #[test]
383    fn schema_column_cap_is_reported_as_incomplete_not_a_total() {
384        let result = parse_tables("orders|id|integer|NO\n", "public", 1, 1).unwrap();
385        assert!(result.truncated);
386        assert_eq!(result.total_count, None);
387        let complete = parse_tables("orders|id|integer|NO\n", "public", 1, 20).unwrap();
388        assert!(!complete.truncated);
389        assert_eq!(complete.total_count, Some(1));
390        assert_eq!(
391            complete.items[0].body.as_deref(),
392            Some("public.orders\n  id: integer")
393        );
394        for invalid in [
395            "orders|partial\n",
396            "orders|id||NO\n",
397            "orders|id|  |NO\n",
398            "orders|id|integer|\n",
399            "orders|id|integer|UNKNOWN\n",
400        ] {
401            assert!(parse_tables(invalid, "public", 1, 20).is_err());
402        }
403    }
404
405    #[cfg(unix)]
406    #[test]
407    fn process_failure_does_not_disclose_database_diagnostics() {
408        let mut command = std::process::Command::new("/bin/sh");
409        command.args(["-c", "printf 'private-diagnostic-marker' >&2; exit 1"]);
410        let error = run_psql(&mut command, QUERY_TIMEOUT).unwrap_err();
411        assert!(error.contains("diagnostics withheld"));
412        assert!(!error.contains("private-diagnostic-marker"));
413    }
414
415    #[cfg(unix)]
416    #[test]
417    fn process_capture_rejects_oversized_stdout_and_stderr() {
418        for script in [
419            "/bin/dd if=/dev/zero bs=1048577 count=1 2>/dev/null",
420            "/bin/dd if=/dev/zero bs=16385 count=1 1>&2 2>/dev/null",
421        ] {
422            let mut command = std::process::Command::new("/bin/sh");
423            command.args(["-c", script]);
424            assert!(
425                run_psql(&mut command, QUERY_TIMEOUT)
426                    .unwrap_err()
427                    .contains("output bound")
428            );
429        }
430    }
431
432    #[cfg(unix)]
433    #[test]
434    fn process_capture_has_a_deadline_and_rejects_invalid_utf8() {
435        let mut command = std::process::Command::new("/bin/sleep");
436        command.arg("20");
437        let started = std::time::Instant::now();
438        let error = run_psql(&mut command, std::time::Duration::from_millis(20)).unwrap_err();
439        assert!(error.contains("deadline"));
440        assert!(started.elapsed() < std::time::Duration::from_secs(5));
441        let mut command = std::process::Command::new("/bin/sh");
442        command.args(["-c", "printf '\\377'"]);
443        assert!(
444            run_psql(&mut command, QUERY_TIMEOUT)
445                .unwrap_err()
446                .contains("invalid UTF-8")
447        );
448    }
449}
lean_ctx/core/providers/postgres.rs

lean_ctx/core/providers/
postgres.rs
