1use crate::core::catalog::Catalog;
2use crate::db::context::Context;
3use crate::db::dataset::Dataset;
4use crate::db::driver::Driver;
5use crate::error::MyResult;
6use itertools::Itertools;
7use odbc_api::{Connection, ConnectionOptions, Environment};
8use std::borrow::Cow;
9use std::fs::File;
10use std::io::Read;
11
12pub trait Plugin {
13 fn populate_catalog(&self, database: &Database, catalog: &mut Catalog) -> MyResult<()>;
14
15 fn query_context(&self, database: &Database) -> Context;
16
17 fn get_batch(&self) -> &str;
18}
19
20#[cfg(any(test, debug_assertions))]
21#[derive(Clone, Copy)]
22pub enum Sample {
23 Hugos,
24 Numbers,
25}
26
27pub struct Database<'a> {
28 conn: Connection<'a>,
29 plugin: Box<dyn Plugin>,
30}
31
32impl<'a> Database<'a> {
33 pub fn new(env: &'a Environment, driver: &Driver) -> MyResult<Self> {
34 let odbc = driver.format_odbc();
35 let options = ConnectionOptions::default();
36 let conn = env.connect_with_connection_string(&odbc, options)?;
37 let plugin = driver.create_plugin();
38 let database = Self { conn, plugin };
39 Ok(database)
40 }
41
42 pub fn create_catalog(&self) -> MyResult<Catalog> {
43 let mut catalog = Catalog::new();
44 self.plugin.populate_catalog(self, &mut catalog)?;
45 Ok(catalog)
46 }
47
48 pub fn query_context(&self) -> Context {
49 self.plugin.query_context(self)
50 }
51
52 pub fn get_batch(&self) -> &str {
53 self.plugin.get_batch()
54 }
55
56 pub fn create_dataset<R: Read>(&self, query: &str) -> MyResult<Option<Dataset<'_, R>>> {
57 let cursor = self.conn.execute(query, (), None)?;
58 if let Some(cursor) = cursor {
59 let dataset = Dataset::from_cursor(cursor)?;
60 Ok(dataset)
61 } else {
62 Ok(None)
63 }
64 }
65
66 pub fn fetch_one<F: FnMut(Cow<str>) -> MyResult<()>>(
67 &self,
68 query: &str,
69 mut function: F,
70 ) -> MyResult<()> {
71 self.fetch_records(query, |record| {
72 if let Some(value) = record.into_iter().next() {
73 function(value)?;
74 }
75 Ok(())
76 })?;
77 Ok(())
78 }
79
80 pub fn fetch_two<F: FnMut(Cow<str>, Cow<str>) -> MyResult<()>>(
81 &self,
82 query: &str,
83 mut function: F,
84 ) -> MyResult<()> {
85 self.fetch_records(query, |record| {
86 if let Some((v1, v2)) = record.into_iter().next_tuple() {
87 function(v1, v2)?;
88 }
89 Ok(())
90 })?;
91 Ok(())
92 }
93
94 pub fn fetch_three<F: FnMut(Cow<str>, Cow<str>, Cow<str>) -> MyResult<()>>(
95 &self,
96 query: &str,
97 mut function: F,
98 ) -> MyResult<()> {
99 self.fetch_records(query, |record| {
100 if let Some((v1, v2, v3)) = record.into_iter().next_tuple() {
101 function(v1, v2, v3)?;
102 }
103 Ok(())
104 })?;
105 Ok(())
106 }
107
108 pub fn fetch_four<F: FnMut(Cow<str>, Cow<str>, Cow<str>, Cow<str>) -> MyResult<()>>(
109 &self,
110 query: &str,
111 mut function: F,
112 ) -> MyResult<()> {
113 self.fetch_records(query, |record| {
114 if let Some((v1, v2, v3, v4)) = record.into_iter().next_tuple() {
115 function(v1, v2, v3, v4)?;
116 }
117 Ok(())
118 })?;
119 Ok(())
120 }
121
122 pub fn fetch_five<F: FnMut(Cow<str>, Cow<str>, Cow<str>, Cow<str>, Cow<str>) -> MyResult<()>>(
123 &self,
124 query: &str,
125 mut function: F,
126 ) -> MyResult<()> {
127 self.fetch_records(query, |record| {
128 if let Some((v1, v2, v3, v4, v5)) = record.into_iter().next_tuple() {
129 function(v1, v2, v3, v4, v5)?;
130 }
131 Ok(())
132 })?;
133 Ok(())
134 }
135
136 pub fn fetch_records<F: FnMut(Vec<Cow<str>>) -> MyResult<()>>(
137 &self,
138 query: &str,
139 mut function: F,
140 ) -> MyResult<()> {
141 let dataset = self.create_dataset::<File>(query)?;
142 if let Some(dataset) = dataset {
143 dataset.get_records(64, |record| {
144 function(record)?;
145 Ok(())
146 })?;
147 }
148 Ok(())
149 }
150
151 #[cfg(any(test, debug_assertions))]
152 pub fn with_catalog(self, sample: Sample) -> MyResult<Self> {
153 let create_author = "\
154CREATE TABLE Author
155(
156 Id INTEGER NOT NULL,
157 Forename TEXT(20) NOT NULL,
158 Surname TEXT(20) NOT NULL
159)
160";
161 let create_book = "\
162CREATE TABLE Book
163(
164 Year INTEGER NOT NULL,
165 AuthorId INTEGER NOT NULL,
166 Article TEXT(5) NULL,
167 Title TEXT(50) NOT NULL,
168 Genre TEXT(5) NOT NULL
169)
170";
171 let create_number = "\
172CREATE TABLE Number
173(
174 Label TEXT(10) NOT NULL,
175 Int INTEGER NULL,
176 Dec REAL NULL,
177 LongIntegerValue INTEGER NULL,
178 LongDecimalValue REAL NULL
179)
180";
181 match sample {
182 Sample::Hugos => {
183 self.execute_sql(create_author)?;
184 self.execute_sql(create_book)?;
185 }
186 Sample::Numbers => {
187 self.execute_sql(create_number)?;
188 }
189 }
190 Ok(self)
191 }
192
193 #[cfg(any(test, debug_assertions))]
194 pub fn with_records(self, sample: Sample) -> MyResult<Self> {
195 let zelf = self.with_catalog(sample)?;
196 match sample {
197 Sample::Hugos => {
198 zelf.execute_sql("INSERT INTO Author VALUES (1, 'Alfred', 'Bester')")?;
199 zelf.execute_sql("INSERT INTO Author VALUES (2, 'Robert A', 'Heinlein')")?;
200 zelf.execute_sql("INSERT INTO Author VALUES (3, 'Fritz', 'Leiber')")?;
201 zelf.execute_sql("INSERT INTO Author VALUES (4, 'James', 'Blish')")?;
202 zelf.execute_sql("INSERT INTO Author VALUES (5, 'Walter M', 'Miller')")?;
203 zelf.execute_sql("INSERT INTO Author VALUES (6, 'Philip K', 'Dick')")?;
204 zelf.execute_sql("INSERT INTO Author VALUES (7, 'Clifford D', 'Simak')")?;
205 zelf.execute_sql("INSERT INTO Author VALUES (8, 'Frank', 'Herbert')")?;
206 zelf.execute_sql("INSERT INTO Author VALUES (9, 'Roger', 'Zelazny')")?;
207 zelf.execute_sql("INSERT INTO Author VALUES (10, 'John', 'Brunner')")?;
208 zelf.execute_sql("INSERT INTO Author VALUES (11, 'Ursula K', 'LeGuin')")?;
209 zelf.execute_sql("INSERT INTO Author VALUES (12, 'Larry', 'Niven')")?;
210 zelf.execute_sql("INSERT INTO Author VALUES (13, 'Philip José', 'Farmer')")?;
211 zelf.execute_sql("INSERT INTO Author VALUES (14, 'Isaac', 'Asimov')")?;
212 zelf.execute_sql("INSERT INTO Author VALUES (15, 'Arthur C', 'Clarke')")?;
213 zelf.execute_sql("INSERT INTO Book VALUES (1953, 1, 'The', 'Demolished Man', 'SF')")?;
214 zelf.execute_sql("INSERT INTO Book VALUES (1956, 2, NULL, 'Double Star', 'SF')")?;
215 zelf.execute_sql("INSERT INTO Book VALUES (1958, 3, 'The', 'Big Time', 'SF')")?;
216 zelf.execute_sql("INSERT INTO Book VALUES (1959, 4, 'A', 'Case of Conscience', 'SF')")?;
217 zelf.execute_sql("INSERT INTO Book VALUES (1960, 2, NULL, 'Starship Troopers', 'SF')")?;
218 zelf.execute_sql("INSERT INTO Book VALUES (1961, 5, 'A', 'Canticle for Leibowitz', 'SF')")?;
219 zelf.execute_sql("INSERT INTO Book VALUES (1962, 2, NULL, 'Stranger in a Strange Land', 'SF')")?;
220 zelf.execute_sql("INSERT INTO Book VALUES (1963, 6, 'The', 'Man in the High Castle', 'F')")?;
221 zelf.execute_sql("INSERT INTO Book VALUES (1964, 7, NULL, 'Way Station', 'SF')")?;
222 zelf.execute_sql("INSERT INTO Book VALUES (1965, 3, 'The', 'Wanderer', 'SF')")?;
223 zelf.execute_sql("INSERT INTO Book VALUES (1966, 8, NULL, 'Dune', 'SF/F')")?;
224 zelf.execute_sql("INSERT INTO Book VALUES (1966, 9, NULL, 'And Call Me Conrad', 'SF/F')")?;
225 zelf.execute_sql("INSERT INTO Book VALUES (1967, 2, 'The', 'Moon Is a Harsh Mistress', 'SF')")?;
226 zelf.execute_sql("INSERT INTO Book VALUES (1968, 9, NULL, 'Lord of Light', 'SF/F')")?;
227 zelf.execute_sql("INSERT INTO Book VALUES (1969, 10, NULL, 'Stand on Zanzibar', 'SF')")?;
228 zelf.execute_sql("INSERT INTO Book VALUES (1970, 11, 'The', 'Left Hand of Darkness', 'SF')")?;
229 zelf.execute_sql("INSERT INTO Book VALUES (1971, 12, NULL, 'Ringworld', 'SF')")?;
230 zelf.execute_sql("INSERT INTO Book VALUES (1972, 13, NULL, 'To Your Scattered Bodies Go', 'SF/F')")?;
231 zelf.execute_sql("INSERT INTO Book VALUES (1973, 14, 'The', 'Gods Themselves', 'SF')")?;
232 zelf.execute_sql("INSERT INTO Book VALUES (1974, 15, NULL, 'Rendezvous with Rama', 'SF')")?;
233 }
234 Sample::Numbers => {
235 zelf.execute_sql("INSERT INTO Number VALUES ('Large', -99999, -99999.99999, -99999, -99999.99999)")?;
236 zelf.execute_sql("INSERT INTO Number VALUES ('Medium', -555, -555.555, -555, -555.555)")?;
237 zelf.execute_sql("INSERT INTO Number VALUES ('Small', -1, -1.0, -1, -1.0)")?;
238 zelf.execute_sql("INSERT INTO Number VALUES ('Null', NULL, NULL, NULL, NULL)")?;
239 zelf.execute_sql("INSERT INTO Number VALUES ('Small', 1, 1.0, 1, 1.0)")?;
240 zelf.execute_sql("INSERT INTO Number VALUES ('Medium', 555, 555.555, 555, 555.555)")?;
241 zelf.execute_sql("INSERT INTO Number VALUES ('Large', 99999, 99999.99999, 99999, 99999.99999)")?;
242 }
243 }
244 Ok(zelf)
245 }
246
247 #[cfg(any(test, debug_assertions))]
248 fn execute_sql(&self, query: &str) -> MyResult<()> {
249 self.conn.execute(query, (), None)?;
250 Ok(())
251 }
252}
253
254#[cfg(test)]
255mod tests {
256 use crate::db::database::{Database, Sample};
257 use crate::db::driver::Driver;
258 use crate::error::MyResult;
259 use crate::util::painter::Painter;
260 use odbc_api::Environment;
261 use pretty_assertions::assert_eq;
262 use std::fs::File;
263
264 #[test]
265 fn test_sql_query_returns_columns_and_records() -> MyResult<()> {
266 let expected = vec![
267 vec!["Year", "Title"],
268 vec!["1956", "Double Star"],
269 vec!["1970", "Left Hand of Darkness"],
270 vec!["1967", "Moon Is a Harsh Mistress"],
271 vec!["1960", "Starship Troopers"],
272 vec!["1962", "Stranger in a Strange Land"],
273 ];
274 let query = "\
275SELECT Year, Title FROM Book
276INNER JOIN Author ON Author.Id = Book.AuthorId
277WHERE Surname IN ('Heinlein', 'LeGuin')
278ORDER BY Title, Year
279";
280 let env = Environment::new()?;
281 let database = create_database(&env)?.with_records(Sample::Hugos)?;
282 let records = collect_records(&database, query)?;
283 assert_eq!(records, expected);
284 Ok(())
285 }
286
287 fn create_database(env: &Environment) -> MyResult<Database<'_>> {
288 let driver = Driver::new("Driver=SQLite3;Database=:memory:")?;
289 Database::new(env, &driver)
290 }
291
292 fn collect_records(database: &Database, query: &str) -> MyResult<Vec<Vec<String>>> {
293 let mut records = Vec::new();
294 let dataset = database.create_dataset::<File>(query)?;
295 if let Some(dataset) = dataset {
296 records.push(dataset.get_columns(Painter::None));
297 dataset.get_records(4096, |record| {
298 let record = record.into_iter().map(String::from).collect();
299 records.push(record);
300 Ok(())
301 })?;
302 }
303 Ok(records)
304 }
305}