feat: add read-only PostgreSQL inspection TUI

src/db.rs

@@ -0,0 +1,286 @@
+use std::time::Duration;
+
+use native_tls::TlsConnector;
+use postgres::{Client, Config};
+use postgres_native_tls::MakeTlsConnector;
+
+use crate::Result;
+
+pub const PAGE_SIZE: usize = 100;
+
+pub struct Database {
+    client: Client,
+}
+
+pub struct Relation {
+    pub oid: u32,
+    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<Vec<Option<String>>>,
+    pub has_next: bool,
+}
+
+impl Database {
+    pub fn connect(connection: &str) -> Result<Self> {
+        let mut config: Config = connection.parse()?;
+        if config.get_connect_timeout().is_none() {
+            config.connect_timeout(Duration::from_secs(5));
+        }
+        let tls = MakeTlsConnector::new(TlsConnector::new()?);
+        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';",
+        )?;
+        Ok(Self { client })
+    }
+
+    pub fn relations(&mut self) -> Result<Vec<Relation>> {
+        let rows = self.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 {
+                    oid: row.get(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<Vec<Column>> {
+        let rows = self.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",
+            &[&relation.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 page(&mut self, relation: &Relation, columns: &[Column], page: usize) -> Result<Page> {
+        let projection = columns
+            .iter()
+            .map(|column| format!("t.{}::text", quote_identifier(&column.name)))
+            .collect::<Vec<_>>()
+            .join(", ");
+        let primary_key = columns
+            .iter()
+            .filter(|column| column.primary_key)
+            .map(|column| format!("t.{}", quote_identifier(&column.name)))
+            .collect::<Vec<_>>()
+            .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 = i64::try_from(page)?
+            .checked_mul(PAGE_SIZE as i64)
+            .ok_or("page offset overflow")?;
+        let result = self
+            .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 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 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::<Config>()?
+            .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)?;
+        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 error = db
+            .client
+            .batch_execute(&format!("DELETE FROM {schema}.empty_table"))
+            .unwrap_err();
+        assert_eq!(
+            error.code(),
+            Some(&postgres::error::SqlState::READ_ONLY_SQL_TRANSACTION)
+        );
+        Ok(())
+    }
+}