Skip to main content

dbiewlite_core/
lib.rs

1mod types;
2mod sqlite;
3mod duckdb_backend;
4
5pub use types::*;
6pub use sqlite::SqliteBackend;
7pub use duckdb_backend::DuckdbBackend;
8
9use std::io::Write;
10use std::path::Path;
11
12pub enum Database {
13    SQLite(SqliteBackend),
14    DuckDB(DuckdbBackend),
15}
16
17impl Database {
18    pub fn open(path: &str) -> Result<Self, String> {
19        let ext = Path::new(path)
20            .extension()
21            .and_then(|e| e.to_str())
22            .unwrap_or("")
23            .to_lowercase();
24
25        match ext.as_str() {
26            "parquet" | "pq" => {
27                Ok(Database::DuckDB(DuckdbBackend::open_parquet(path)?))
28            }
29            "duckdb" => {
30                Ok(Database::DuckDB(DuckdbBackend::open(path)?))
31            }
32            "sqlite" | "sqlite3" => {
33                Ok(Database::SQLite(SqliteBackend::open(path)?))
34            }
35            "db" => {
36                if is_sqlite_file(path) {
37                    Ok(Database::SQLite(SqliteBackend::open(path)?))
38                } else {
39                    DuckdbBackend::open(path)
40                        .map(Database::DuckDB)
41                        .or_else(|_| SqliteBackend::open(path).map(Database::SQLite))
42                }
43            }
44            _ => {
45                SqliteBackend::open(path)
46                    .map(Database::SQLite)
47                    .or_else(|_| DuckdbBackend::open(path).map(Database::DuckDB))
48            }
49        }
50    }
51
52    pub fn path(&self) -> &str {
53        match self {
54            Database::SQLite(db) => db.path(),
55            Database::DuckDB(db) => db.path(),
56        }
57    }
58
59    pub fn get_info(&self) -> Result<DbInfo, String> {
60        match self {
61            Database::SQLite(db) => db.get_info(),
62            Database::DuckDB(db) => db.get_info(),
63        }
64    }
65
66    pub fn list_tables(&self) -> Result<Vec<TableInfo>, String> {
67        match self {
68            Database::SQLite(db) => db.list_tables(),
69            Database::DuckDB(db) => db.list_tables(),
70        }
71    }
72
73    pub fn list_views(&self) -> Result<Vec<String>, String> {
74        match self {
75            Database::SQLite(db) => db.list_views(),
76            Database::DuckDB(db) => db.list_views(),
77        }
78    }
79
80    pub fn list_indexes(&self) -> Result<Vec<IndexInfo>, String> {
81        match self {
82            Database::SQLite(db) => db.list_indexes(),
83            Database::DuckDB(db) => db.list_indexes(),
84        }
85    }
86
87    pub fn get_schema(&self, table: &str) -> Result<Vec<ColumnInfo>, String> {
88        match self {
89            Database::SQLite(db) => db.get_schema(table),
90            Database::DuckDB(db) => db.get_schema(table),
91        }
92    }
93
94    pub fn query_table(
95        &self,
96        table: &str,
97        limit: usize,
98        offset: usize,
99        sort: Option<Sort>,
100    ) -> Result<QueryResult, String> {
101        match self {
102            Database::SQLite(db) => db.query_table(table, limit, offset, sort),
103            Database::DuckDB(db) => db.query_table(table, limit, offset, sort),
104        }
105    }
106
107    pub fn run_query(&self, sql: &str) -> Result<QueryResult, String> {
108        match self {
109            Database::SQLite(db) => db.run_query(sql),
110            Database::DuckDB(db) => db.run_query(sql),
111        }
112    }
113
114    pub fn get_row_count(&self, table: &str) -> Result<u64, String> {
115        match self {
116            Database::SQLite(db) => db.get_row_count(table),
117            Database::DuckDB(db) => db.get_row_count(table),
118        }
119    }
120
121    /// Writes the whole table, in the order it is being viewed in. The page
122    /// window is deliberately not applied: exporting the fifty rows that happen
123    /// to be on screen out of thousands is not what the button appears to do.
124    ///
125    /// Neither this nor the query export touches anything specific to a
126    /// backend, so both live here rather than being written twice.
127    pub fn export_csv<W: Write>(
128        &self,
129        table: &str,
130        sort: Option<Sort>,
131        writer: &mut W,
132    ) -> Result<(), String> {
133        let sql = format!("SELECT * FROM \"{}\"{}", table, order_clause(&sort));
134        let result = self.run_query(&sql)?;
135        write_csv(&result, writer)
136    }
137
138    /// Re-runs the query rather than being handed its rows, so the export is
139    /// not bounded by whatever the caller happens to be holding.
140    pub fn export_query_csv<W: Write>(&self, sql: &str, writer: &mut W) -> Result<(), String> {
141        let result = self.run_query(sql)?;
142        write_csv(&result, writer)
143    }
144}
145
146fn is_sqlite_file(path: &str) -> bool {
147    use std::io::Read;
148    let mut file = match std::fs::File::open(path) {
149        Ok(f) => f,
150        Err(_) => return false,
151    };
152    let mut header = [0u8; 16];
153    if file.read_exact(&mut header).is_err() {
154        return false;
155    }
156    &header == b"SQLite format 3\0"
157}
158
159#[cfg(test)]
160mod tests {
161    use super::*;
162
163    #[test]
164    fn test_open_parquet() {
165        // Create a DuckDB in-memory, export a parquet, then open it
166        let conn = duckdb::Connection::open_in_memory().unwrap();
167        conn.execute_batch("CREATE TABLE test (id INTEGER, name VARCHAR); INSERT INTO test VALUES (1, 'Alice'), (2, 'Bob');").unwrap();
168        conn.execute_batch("COPY test TO '/tmp/dbiew_test.parquet' (FORMAT PARQUET);").unwrap();
169        drop(conn);
170
171        let db = Database::open("/tmp/dbiew_test.parquet").unwrap();
172        let info = db.get_info().unwrap();
173        assert_eq!(info.engine, "Parquet");
174        assert_eq!(info.table_count, 1);
175
176        let tables = db.list_tables().unwrap();
177        assert_eq!(tables.len(), 1);
178        assert_eq!(tables[0].name, "dbiew_test");
179        assert_eq!(tables[0].row_count, 2);
180
181        let schema = db.get_schema("dbiew_test").unwrap();
182        assert_eq!(schema.len(), 2);
183        assert_eq!(schema[0].name, "id");
184        assert_eq!(schema[1].name, "name");
185
186        let data = db.query_table("dbiew_test", 10, 0, None).unwrap();
187        assert_eq!(data.rows.len(), 2);
188        assert_eq!(data.total_rows, Some(2));
189    }
190
191    #[test]
192    fn test_open_sqlite() {
193        let tmp = "/tmp/dbiew_test.sqlite";
194        let conn = rusqlite::Connection::open(tmp).unwrap();
195        conn.execute_batch("CREATE TABLE IF NOT EXISTS items (id INTEGER PRIMARY KEY, val TEXT); INSERT OR IGNORE INTO items VALUES (1, 'hello');").unwrap();
196        drop(conn);
197
198        let db = Database::open(tmp).unwrap();
199        let info = db.get_info().unwrap();
200        assert_eq!(info.engine, "SQLite");
201        assert!(info.table_count >= 1);
202    }
203
204    #[test]
205    fn test_duckdb_views() {
206        let db = Database::open("../../samples/sensors.duckdb").unwrap();
207        let views = db.list_views().unwrap();
208        println!("Views: {:?}", views);
209        assert!(!views.is_empty(), "Expected at least one view");
210        assert!(views.contains(&"daily_averages".to_string()));
211    }
212}