use std::{path::Path, time::Duration}; use native_tls::{Certificate, TlsConnector}; use postgres::{Client, Config, config::SslMode}; use postgres_native_tls::MakeTlsConnector; use crate::Result; mod sqlite; pub const PAGE_SIZE: usize = 100; pub struct Database { backend: Backend, display_name: String, } #[allow( clippy::large_enum_variant, reason = "ponytail: one connection per app; box only if storing backends in bulk" )] enum Backend { Postgres(Client), Sqlite(rusqlite::Connection), } pub struct Relation { pub id: i64, pub schema: String, pub name: String, pub label: String, pub kind: String, physical_order: bool, } pub struct Column { pub name: String, pub label: String, pub data_type: String, pub nullable: bool, pub primary_key: bool, pub default: String, } pub struct Page { pub rows: Vec>>, pub has_next: bool, } pub struct QueryResult { pub headers: Vec, pub page: Page, } impl Database { pub fn connect(connection: &str, root_cert: Option<&Path>) -> Result { Self::connect_postgres(connection.parse()?, root_cert) } pub fn connect_postgres(mut config: Config, root_cert: Option<&Path>) -> Result { if config.get_connect_timeout().is_none() { config.connect_timeout(Duration::from_secs(5)); } let tls = MakeTlsConnector::new(postgres_tls(config.get_ssl_mode(), root_cert)?); let mut client = config.connect(tls)?; // Views can invoke functions: enforce read-only on the server, not just in the UI. client.batch_execute( "SET default_transaction_read_only = on; SET statement_timeout = '5s';", )?; let display_name: String = client.query_one("SELECT current_database()", &[])?.get(0); Ok(Self { backend: Backend::Postgres(client), display_name: display_text(&display_name), }) } pub fn open_sqlite(path: &Path) -> Result { Ok(Self { backend: Backend::Sqlite(sqlite::open(path)?), display_name: display_text( &path .file_name() .unwrap_or(path.as_os_str()) .to_string_lossy(), ), }) } pub fn display_name(&self) -> &str { &self.display_name } pub fn name(&self) -> &'static str { match self.backend { Backend::Postgres(_) => "PostgreSQL", Backend::Sqlite(_) => "SQLite", } } pub fn relations(&mut self) -> Result> { let client = match &mut self.backend { Backend::Postgres(client) => client, Backend::Sqlite(connection) => return sqlite::relations(connection), }; let rows = client.query( "SELECT c.oid, n.nspname, c.relname, c.relkind::text FROM pg_catalog.pg_class c JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace WHERE c.relkind IN ('r', 'p', 'v', 'm', 'f') AND n.nspname NOT IN ('pg_catalog', 'information_schema') AND n.nspname !~ '^pg_(toast|temp)' AND pg_catalog.has_schema_privilege(n.oid, 'USAGE') AND pg_catalog.has_table_privilege(c.oid, 'SELECT') ORDER BY n.nspname, c.relname", &[], )?; Ok(rows .into_iter() .map(|row| { let schema: String = row.get(1); let name: String = row.get(2); let kind: &str = row.get(3); Relation { id: i64::from(row.get::<_, u32>(0)), label: display_text(&format!("{schema}.{name}")), schema, name, physical_order: matches!(kind, "r" | "p" | "m"), kind: match kind { "v" => "view", "m" => "materialized view", "f" => "foreign table", "p" => "partitioned table", _ => "table", } .into(), } }) .collect()) } pub fn columns(&mut self, relation: &Relation) -> Result> { let client = match &mut self.backend { Backend::Postgres(client) => client, Backend::Sqlite(connection) => return sqlite::columns(connection, relation), }; let oid = u32::try_from(relation.id)?; let rows = client.query( "SELECT a.attname, pg_catalog.format_type(a.atttypid, a.atttypmod), NOT a.attnotnull, EXISTS (SELECT FROM pg_catalog.pg_index i WHERE i.indrelid = a.attrelid AND i.indisprimary AND a.attnum = ANY(i.indkey)), COALESCE(pg_catalog.pg_get_expr(d.adbin, d.adrelid), '') FROM pg_catalog.pg_attribute a LEFT JOIN pg_catalog.pg_attrdef d ON d.adrelid = a.attrelid AND d.adnum = a.attnum WHERE a.attrelid = $1 AND a.attnum > 0 AND NOT a.attisdropped ORDER BY a.attnum", &[&oid], )?; Ok(rows .into_iter() .map(|row| { let name: String = row.get(0); Column { label: display_text(&name), name, data_type: display_text(row.get(1)), nullable: row.get(2), primary_key: row.get(3), default: display_text(row.get(4)), } }) .collect()) } pub fn query(&mut self, sql: &str) -> Result { let sql = sql.trim(); let sql = sql.strip_suffix(';').unwrap_or(sql).trim_end(); let client = match &mut self.backend { Backend::Postgres(client) => client, Backend::Sqlite(connection) => return sqlite::query(connection, sql), }; let mut transaction = client.build_transaction().read_only(true).start()?; transaction.batch_execute("SET LOCAL statement_timeout = '5s'")?; let statement = transaction.prepare(&format!( "SELECT * FROM (\n{sql}\n) AS query_result LIMIT 0" ))?; let headers = statement .columns() .iter() .map(|column| display_text(column.name())) .collect::>(); let aliases = (0..headers.len()) .map(|index| format!("c{index}")) .collect::>(); let projection = aliases .iter() .map(|alias| format!("query_result.{alias}::text")) .collect::>() .join(", "); let alias_list = if aliases.is_empty() { String::new() } else { format!("({})", aliases.join(", ")) }; let result = transaction.query( &format!( "SELECT {projection} FROM (\n{sql}\n) AS query_result{alias_list} LIMIT {}", PAGE_SIZE + 1 ), &[], )?; let has_next = result.len() > PAGE_SIZE; let rows = result .into_iter() .take(PAGE_SIZE) .map(|row| { (0..headers.len()) .map(|index| row.get::<_, Option<&str>>(index).map(display_text)) .collect() }) .collect(); Ok(QueryResult { headers, page: Page { rows, has_next }, }) } pub fn page(&mut self, relation: &Relation, columns: &[Column], page: usize) -> Result { let client = match &mut self.backend { Backend::Postgres(client) => client, Backend::Sqlite(connection) => { return sqlite::page(connection, relation, columns, page); } }; let projection = columns .iter() .map(|column| format!("t.{}::text", quote_identifier(&column.name))) .collect::>() .join(", "); let primary_key = columns .iter() .filter(|column| column.primary_key) .map(|column| format!("t.{}", quote_identifier(&column.name))) .collect::>() .join(", "); let order = if !primary_key.is_empty() { format!(" ORDER BY {primary_key}") } else if relation.physical_order { " ORDER BY t.tableoid, t.ctid".into() } else { // ponytail: keyless views have no stable row identity; use keyset paging if a key is provided. String::new() }; // ponytail: OFFSET is simple but slows on deep pages; switch to keyset paging if that matters. let sql = format!( "SELECT {projection} FROM {}.{} AS t{order} LIMIT $1 OFFSET $2", quote_identifier(&relation.schema), quote_identifier(&relation.name), ); let offset = page_offset(page)?; let result = client.query(&sql, &[&((PAGE_SIZE + 1) as i64), &offset])?; let has_next = result.len() > PAGE_SIZE; let rows = result .into_iter() .take(PAGE_SIZE) .map(|row| { (0..columns.len()) .map(|index| row.get::<_, Option<&str>>(index).map(display_text)) .collect() }) .collect(); Ok(Page { rows, has_next }) } } fn postgres_tls(mode: SslMode, root_cert: Option<&Path>) -> Result { let mut tls = TlsConnector::builder(); if let Some(path) = root_cert { if mode != SslMode::Require { return Err("A CA bundle requires Require TLS (sslmode=require); plaintext fallback is not allowed.".into()); } let pem = std::fs::read(path) .map_err(|error| format!("Cannot read CA bundle: {error}"))?; let certificates = Certificate::stack_from_pem(&pem) .map_err(|error| format!("Invalid PEM CA bundle: {error}"))?; if certificates.is_empty() { return Err("CA bundle must contain at least one PEM CA certificate.".into()); } // An explicit bundle replaces system roots; certificate and hostname checks stay enabled. tls.disable_built_in_roots(true); for certificate in certificates { tls.add_root_certificate(certificate); } } Ok(tls.build()?) } fn page_offset(page: usize) -> Result { i64::try_from(page)? .checked_mul(PAGE_SIZE as i64) .ok_or_else(|| "page offset overflow".into()) } fn quote_identifier(value: &str) -> String { format!("\"{}\"", value.replace('"', "\"\"")) } /// Database text must never emit terminal control sequences. pub fn display_text(value: &str) -> String { let mut output = String::with_capacity(value.len()); for character in value.chars() { if character.is_control() { output.extend(character.escape_debug()); } else { output.push(character); } } output } #[cfg(test)] mod tests { use super::*; #[test] fn explicit_ca_bundle_fails_closed_on_insecure_modes_and_invalid_files() -> Result<()> { let bundle = tempfile::NamedTempFile::new()?; for mode in [SslMode::Prefer, SslMode::Disable] { assert!(postgres_tls(mode, Some(bundle.path())).is_err()); } for contents in ["", "not a PEM certificate", "-----BEGIN CERTIFICATE-----\ninvalid\n-----END CERTIFICATE-----"] { std::fs::write(bundle.path(), contents)?; assert!(postgres_tls(SslMode::Require, Some(bundle.path())).is_err()); } let missing = bundle.path().with_extension("missing"); assert!(postgres_tls(SslMode::Require, Some(&missing)).is_err()); Ok(()) } #[test] fn sqlite_identity_is_filename_without_terminal_controls() -> Result<()> { let directory = tempfile::tempdir()?; let path = directory.path().join("accounts\nprivate.sqlite"); drop(rusqlite::Connection::open(&path)?); let db = Database::open_sqlite(&path)?; assert_eq!(db.display_name(), "accounts\\nprivate.sqlite"); Ok(()) } #[test] fn identifiers_and_terminal_controls_are_safe() { assert_eq!( quote_identifier("a\"; DROP TABLE x;--"), "\"a\"\"; DROP TABLE x;--\"" ); assert_eq!(display_text("é🙂\n\t\r\u{1b}[2J"), "é🙂\\n\\t\\r\\u{1b}[2J"); } #[test] #[ignore = "requires DBDASH_TEST_DATABASE_URL pointing at a disposable PostgreSQL database"] fn inspect_postgres_read_only() -> Result<()> { let connection = std::env::var("DBDASH_TEST_DATABASE_URL")?; let mut fixture = connection .parse::()? .connect(MakeTlsConnector::new(TlsConnector::new()?))?; let schema = format!("dbdash_test_{}", std::process::id()); fixture.batch_execute(&format!("CREATE SCHEMA {schema}"))?; struct Cleanup(Client, String); impl Drop for Cleanup { fn drop(&mut self) { let _ = self .0 .batch_execute(&format!("DROP SCHEMA {} CASCADE", self.1)); } } let mut cleanup = Cleanup(fixture, schema.clone()); cleanup.0.batch_execute(&format!( "CREATE TABLE {schema}.\"odd\"\"table\" ( id integer PRIMARY KEY, \"odd\"\"column\" text, amount numeric NOT NULL DEFAULT 1.25, details jsonb, tags text[], bytes bytea, happened timestamp); INSERT INTO {schema}.\"odd\"\"table\" SELECT n, CASE WHEN n = 1 THEN NULL WHEN n = 2 THEN E'line\\nnext' ELSE 'NULL' END, 1.25, '{{\"ok\": true}}', ARRAY['a', 'b'], '\\x00ff', TIMESTAMP '2020-01-01 12:00:00' FROM generate_series(1, 101) n; CREATE VIEW {schema}.sample_view AS SELECT id FROM {schema}.\"odd\"\"table\"; CREATE TABLE {schema}.empty_table (id integer);" ))?; let mut db = Database::connect(&connection, None)?; let relations = db.relations()?; let table = relations .iter() .find(|r| r.schema == schema && r.name == "odd\"table") .unwrap(); assert!( relations .iter() .any(|r| r.schema == schema && r.kind == "view") ); assert!(!relations.iter().any(|r| r.schema == "pg_catalog")); let columns = db.columns(table)?; assert!(columns[0].primary_key); assert!(!columns[0].nullable); assert_eq!(columns[1].name, "odd\"column"); assert!(columns[1].nullable); assert_eq!(columns[2].data_type, "numeric"); assert_eq!(columns[2].default, "1.25"); let first = db.page(table, &columns, 0)?; assert!(first.has_next); assert_eq!(first.rows.len(), 100); assert_eq!( first.rows[0], vec![ Some("1".into()), None, Some("1.25".into()), Some("{\"ok\": true}".into()), Some("{a,b}".into()), Some("\\x00ff".into()), Some("2020-01-01 12:00:00".into()) ] ); assert_eq!(first.rows[1][1].as_deref(), Some("line\\nnext")); assert_eq!(first.rows[2][1].as_deref(), Some("NULL")); let second = db.page(table, &columns, 1)?; assert!(!second.has_next); assert_eq!(second.rows[0][0].as_deref(), Some("101")); let empty = relations .iter() .find(|r| r.schema == schema && r.name == "empty_table") .unwrap(); let columns = db.columns(empty)?; let page = db.page(empty, &columns, 0)?; assert!(page.rows.is_empty()); assert!(!page.has_next); let query = db.query(&format!( " \nWITH data AS (SELECT * FROM {schema}.\"odd\"\"table\") SELECT * FROM data ORDER BY id; \n" ))?; assert_eq!(query.headers[1], "odd\"column"); assert_eq!(query.page.rows, first.rows); assert!(query.page.has_next); let duplicate = db.query("SELECT 1 AS same, NULL::text AS same")?; assert_eq!(duplicate.headers, ["same", "same"]); assert_eq!(duplicate.page.rows, [vec![Some("1".into()), None]]); assert!(!duplicate.page.has_next); let empty = db.query("SELECT 1 AS empty WHERE false")?; assert_eq!(empty.headers, ["empty"]); assert!(empty.page.rows.is_empty()); assert!(!empty.page.has_next); let controls = db.query("SELECT E'\\033[2J' AS \"line\nheader\" -- trailing comment")?; assert_eq!(controls.headers, ["line\\nheader"]); assert_eq!(controls.page.rows[0][0].as_deref(), Some("\\u{1b}[2J")); assert!(db.query("SELECT 1; SELECT 2").is_err()); assert!( db.query(&format!("DELETE FROM {schema}.empty_table RETURNING id")) .is_err() ); assert!( db.query(&format!( "WITH deleted AS (DELETE FROM {schema}.empty_table RETURNING id) SELECT * FROM deleted" )) .is_err() ); db.query("SELECT set_config('default_transaction_read_only', 'off', false)")?; let settings = db.query("SELECT current_setting('default_transaction_read_only')")?; assert_eq!(settings.page.rows[0][0].as_deref(), Some("on")); if let Backend::Postgres(client) = &mut db.backend { client.batch_execute( "SET default_transaction_read_only = off; SET statement_timeout = 0", )?; } let settings = db.query( "SELECT current_setting('transaction_read_only'), current_setting('statement_timeout')", )?; assert_eq!( settings.page.rows[0], [Some("on".into()), Some("5s".into())] ); let Backend::Postgres(client) = &mut db.backend else { unreachable!() }; client.batch_execute("SET default_transaction_read_only = on")?; let error = client .batch_execute(&format!("DELETE FROM {schema}.empty_table")) .unwrap_err(); assert_eq!( error.code(), Some(&postgres::error::SqlState::READ_ONLY_SQL_TRANSACTION) ); Ok(()) } }