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    pub fn export_csv<W: Write>(&self, table: &str, writer: &mut W) -> Result<(), String> {
122        match self {
123            Database::SQLite(db) => db.export_csv(table, writer),
124            Database::DuckDB(db) => db.export_csv(table, writer),
125        }
126    }
127}
128
129fn is_sqlite_file(path: &str) -> bool {
130    use std::io::Read;
131    let mut file = match std::fs::File::open(path) {
132        Ok(f) => f,
133        Err(_) => return false,
134    };
135    let mut header = [0u8; 16];
136    if file.read_exact(&mut header).is_err() {
137        return false;
138    }
139    &header == b"SQLite format 3\0"
140}
141
142#[cfg(test)]
143mod tests {
144    use super::*;
145
146    #[test]
147    fn test_open_parquet() {
148        // Create a DuckDB in-memory, export a parquet, then open it
149        let conn = duckdb::Connection::open_in_memory().unwrap();
150        conn.execute_batch("CREATE TABLE test (id INTEGER, name VARCHAR); INSERT INTO test VALUES (1, 'Alice'), (2, 'Bob');").unwrap();
151        conn.execute_batch("COPY test TO '/tmp/dbiew_test.parquet' (FORMAT PARQUET);").unwrap();
152        drop(conn);
153
154        let db = Database::open("/tmp/dbiew_test.parquet").unwrap();
155        let info = db.get_info().unwrap();
156        assert_eq!(info.engine, "Parquet");
157        assert_eq!(info.table_count, 1);
158
159        let tables = db.list_tables().unwrap();
160        assert_eq!(tables.len(), 1);
161        assert_eq!(tables[0].name, "dbiew_test");
162        assert_eq!(tables[0].row_count, 2);
163
164        let schema = db.get_schema("dbiew_test").unwrap();
165        assert_eq!(schema.len(), 2);
166        assert_eq!(schema[0].name, "id");
167        assert_eq!(schema[1].name, "name");
168
169        let data = db.query_table("dbiew_test", 10, 0, None).unwrap();
170        assert_eq!(data.rows.len(), 2);
171        assert_eq!(data.total_rows, Some(2));
172    }
173
174    #[test]
175    fn test_open_sqlite() {
176        let tmp = "/tmp/dbiew_test.sqlite";
177        let conn = rusqlite::Connection::open(tmp).unwrap();
178        conn.execute_batch("CREATE TABLE IF NOT EXISTS items (id INTEGER PRIMARY KEY, val TEXT); INSERT OR IGNORE INTO items VALUES (1, 'hello');").unwrap();
179        drop(conn);
180
181        let db = Database::open(tmp).unwrap();
182        let info = db.get_info().unwrap();
183        assert_eq!(info.engine, "SQLite");
184        assert!(info.table_count >= 1);
185    }
186
187    #[test]
188    fn test_duckdb_views() {
189        let db = Database::open("../../samples/sensors.duckdb").unwrap();
190        let views = db.list_views().unwrap();
191        println!("Views: {:?}", views);
192        assert!(!views.is_empty(), "Expected at least one view");
193        assert!(views.contains(&"daily_averages".to_string()));
194    }
195}