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>(
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 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 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}