dbdash/src/db.rs

use std::{path::Path, time::Duration};

use native_tls::TlsConnector;
use postgres::{Client, Config};
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<Vec<Option<String>>>,
    pub has_next: bool,
}

pub struct QueryResult {
    pub headers: Vec<String>,
    pub page: Page,
}

impl Database {
    pub fn connect(connection: &str) -> Result<Self> {
        Self::connect_postgres(connection.parse()?)
    }

    pub fn connect_postgres(mut config: Config) -> Result<Self> {
        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';",
        )?;
        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<Self> {
        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<Vec<Relation>> {
        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<Vec<Column>> {
        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<QueryResult> {
        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::<Vec<_>>();
        let aliases = (0..headers.len())
            .map(|index| format!("c{index}"))
            .collect::<Vec<_>>();
        let projection = aliases
            .iter()
            .map(|alias| format!("query_result.{alias}::text"))
            .collect::<Vec<_>>()
            .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<Page> {
        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::<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 = 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 page_offset(page: usize) -> Result<i64> {
    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 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::<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 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(())
    }
}