mirror of
				https://github.com/isar/rusqlite.git
				synced 2025-11-01 06:18:54 +08:00 
			
		
		
		
	Ensure there is one csv reader by cursor.
This commit is contained in:
		| @@ -3,6 +3,8 @@ extern crate csv; | |||||||
| use std::ffi::CStr; | use std::ffi::CStr; | ||||||
| use std::fs::File; | use std::fs::File; | ||||||
| use std::mem; | use std::mem; | ||||||
|  | use std::path::Path; | ||||||
|  | use std::result; | ||||||
| use std::str; | use std::str; | ||||||
| use libc; | use libc; | ||||||
|  |  | ||||||
| @@ -27,10 +29,21 @@ init_module!(CSV_MODULE, CSVTab, CSVTabCursor, | |||||||
| struct CSVTab { | struct CSVTab { | ||||||
|     /// Base class |     /// Base class | ||||||
|     base: ffi::sqlite3_vtab, |     base: ffi::sqlite3_vtab, | ||||||
|     reader: csv::Reader<File>, |     filename: String, | ||||||
|  |     has_headers: bool, | ||||||
|  |     delimiter: u8, | ||||||
|  |     quote: u8, | ||||||
|     offset_first_row: u64, |     offset_first_row: u64, | ||||||
|     cols: Vec<String>, | } | ||||||
|     eof: bool, |  | ||||||
|  | impl CSVTab { | ||||||
|  |     fn reader(&self) -> result::Result<csv::Reader<File>, csv::Error> { | ||||||
|  |         csv::Reader::from_file(&self.filename).map(|reader| { | ||||||
|  |             reader.has_headers(self.has_headers) | ||||||
|  |                   .delimiter(self.delimiter) | ||||||
|  |                   .quote(self.quote) | ||||||
|  |         }) | ||||||
|  |     } | ||||||
| } | } | ||||||
|  |  | ||||||
| impl VTab<CSVTabCursor> for CSVTab { | impl VTab<CSVTabCursor> for CSVTab { | ||||||
| @@ -47,33 +60,43 @@ impl VTab<CSVTabCursor> for CSVTab { | |||||||
|             c_filename = &c_filename[1..c_filename.len() - 1]; |             c_filename = &c_filename[1..c_filename.len() - 1]; | ||||||
|         } |         } | ||||||
|         let filename = try!(str::from_utf8(c_filename)); |         let filename = try!(str::from_utf8(c_filename)); | ||||||
|         let mut reader = try!(csv::Reader::from_file(filename)).has_headers(false); // TODO flexible ? |         if !Path::new(filename).exists() { | ||||||
|  |             return Err(Error::ModuleError(format!("file '{}' does not exist", filename))); | ||||||
|  |         } | ||||||
|  |         let mut vtab = CSVTab { | ||||||
|  |             base: Default::default(), | ||||||
|  |             filename: String::from(filename), | ||||||
|  |             has_headers: false, | ||||||
|  |             delimiter: b',', | ||||||
|  |             quote: b'"', | ||||||
|  |             offset_first_row: 0, | ||||||
|  |         }; | ||||||
|         let mut cols: Vec<String> = Vec::new(); |         let mut cols: Vec<String> = Vec::new(); | ||||||
|  |  | ||||||
|         let args = &args[4..]; |         let args = &args[4..]; | ||||||
|         for c_arg in args { |         for c_arg in args { | ||||||
|             let c_slice = unsafe { CStr::from_ptr(*c_arg).to_bytes() }; |             let c_slice = unsafe { CStr::from_ptr(*c_arg).to_bytes() }; | ||||||
|             if c_slice.len() == 1 { |             if c_slice.len() == 1 { | ||||||
|                 reader = reader.delimiter(c_slice[0]); |                 vtab.delimiter = c_slice[0]; | ||||||
|             } else if c_slice.len() == 3 && c_slice[0] == b'\'' { |             } else if c_slice.len() == 3 && c_slice[0] == b'\'' { | ||||||
|                 reader = reader.delimiter(c_slice[1]); |                 vtab.delimiter = c_slice[1]; | ||||||
|             } else { |             } else { | ||||||
|                 let arg = try!(str::from_utf8(c_slice)); |                 let arg = try!(str::from_utf8(c_slice)); | ||||||
|                 let uc = arg.to_uppercase(); |                 let uc = arg.to_uppercase(); | ||||||
|                 if uc.contains("HEADER") { |                 if uc.contains("HEADER") { | ||||||
|                     reader = reader.has_headers(true); |                     vtab.has_headers = true; | ||||||
|                 } else if uc.contains("NO_QUOTE") { |                 } else if uc.contains("NO_QUOTE") { | ||||||
|                     reader = reader.quote(0); |                     vtab.quote = 0; | ||||||
|                 } else { |                 } else { | ||||||
|                     cols.push(escape_double_quote(arg).into_owned()); |                     cols.push(escape_double_quote(arg).into_owned()); | ||||||
|                 } |                 } | ||||||
|             } |             } | ||||||
|         } |         } | ||||||
|  |  | ||||||
|         let mut offset_first_row = 0; |         if vtab.has_headers { | ||||||
|         if reader.has_headers { |             let mut reader = try!(vtab.reader()); | ||||||
|             let headers = try!(reader.headers()); |             let headers = try!(reader.headers()); | ||||||
|             offset_first_row = reader.byte_offset(); |             vtab.offset_first_row = reader.byte_offset(); | ||||||
|             // headers ignored if cols is not empty |             // headers ignored if cols is not empty | ||||||
|             if cols.is_empty() { |             if cols.is_empty() { | ||||||
|                 cols = headers; |                 cols = headers; | ||||||
| @@ -99,13 +122,6 @@ impl VTab<CSVTabCursor> for CSVTab { | |||||||
|             } |             } | ||||||
|         } |         } | ||||||
|  |  | ||||||
|         let vtab = CSVTab { |  | ||||||
|             base: Default::default(), |  | ||||||
|             reader: reader, |  | ||||||
|             offset_first_row: offset_first_row, |  | ||||||
|             cols: cols, |  | ||||||
|             eof: false, |  | ||||||
|         }; |  | ||||||
|         try!(declare_vtab(db, &sql)); |         try!(declare_vtab(db, &sql)); | ||||||
|         Ok(vtab) |         Ok(vtab) | ||||||
|     } |     } | ||||||
| @@ -113,7 +129,7 @@ impl VTab<CSVTabCursor> for CSVTab { | |||||||
|     fn best_index(&self, _info: *mut ffi::sqlite3_index_info) {} |     fn best_index(&self, _info: *mut ffi::sqlite3_index_info) {} | ||||||
|  |  | ||||||
|     fn open(&self) -> Result<CSVTabCursor> { |     fn open(&self) -> Result<CSVTabCursor> { | ||||||
|         Ok(CSVTabCursor::new()) |         Ok(CSVTabCursor::new(try!(self.reader()))) | ||||||
|     } |     } | ||||||
| } | } | ||||||
|  |  | ||||||
| @@ -122,15 +138,21 @@ impl VTab<CSVTabCursor> for CSVTab { | |||||||
| struct CSVTabCursor { | struct CSVTabCursor { | ||||||
|     /// Base class |     /// Base class | ||||||
|     base: ffi::sqlite3_vtab_cursor, |     base: ffi::sqlite3_vtab_cursor, | ||||||
|  |     reader: csv::Reader<File>, | ||||||
|     /// Current cursor position |     /// Current cursor position | ||||||
|     row_number: usize, |     row_number: usize, | ||||||
|  |     cols: Vec<String>, | ||||||
|  |     eof: bool, | ||||||
| } | } | ||||||
|  |  | ||||||
| impl CSVTabCursor { | impl CSVTabCursor { | ||||||
|     fn new() -> CSVTabCursor { |     fn new(reader: csv::Reader<File>) -> CSVTabCursor { | ||||||
|         CSVTabCursor { |         CSVTabCursor { | ||||||
|             base: Default::default(), |             base: Default::default(), | ||||||
|  |             reader: reader, | ||||||
|             row_number: 0, |             row_number: 0, | ||||||
|  |             cols: Vec::new(), | ||||||
|  |             eof: false, | ||||||
|         } |         } | ||||||
|     } |     } | ||||||
| } | } | ||||||
| @@ -147,23 +169,22 @@ impl VTabCursor<CSVTab> for CSVTabCursor { | |||||||
|               _argv: *mut *mut ffi::sqlite3_value) |               _argv: *mut *mut ffi::sqlite3_value) | ||||||
|               -> Result<()> { |               -> Result<()> { | ||||||
|         { |         { | ||||||
|             let vtab = self.vtab(); |             let offset_first_row = self.vtab().offset_first_row; | ||||||
|             try!(vtab.reader.seek(vtab.offset_first_row)); |             try!(self.reader.seek(offset_first_row)); | ||||||
|         } |         } | ||||||
|         self.row_number = 0; |         self.row_number = 0; | ||||||
|         self.next() |         self.next() | ||||||
|     } |     } | ||||||
|     fn next(&mut self) -> Result<()> { |     fn next(&mut self) -> Result<()> { | ||||||
|         { |         { | ||||||
|             let vtab = self.vtab(); |             self.eof = self.reader.done(); | ||||||
|             vtab.eof = vtab.reader.done(); |             if self.eof { | ||||||
|             if vtab.eof { |  | ||||||
|                 return Ok(()); |                 return Ok(()); | ||||||
|             } |             } | ||||||
|  |  | ||||||
|             vtab.cols.clear(); |             self.cols.clear(); | ||||||
|             while let Some(col) = vtab.reader.next_str().into_iter_result() { |             while let Some(col) = self.reader.next_str().into_iter_result() { | ||||||
|                 vtab.cols.push(String::from(try!(col))); |                 self.cols.push(String::from(try!(col))); | ||||||
|             } |             } | ||||||
|         } |         } | ||||||
|  |  | ||||||
| @@ -171,21 +192,19 @@ impl VTabCursor<CSVTab> for CSVTabCursor { | |||||||
|         Ok(()) |         Ok(()) | ||||||
|     } |     } | ||||||
|     fn eof(&self) -> bool { |     fn eof(&self) -> bool { | ||||||
|         let vtab = self.vtab(); |         self.eof | ||||||
|         vtab.eof |  | ||||||
|     } |     } | ||||||
|     fn column(&self, ctx: *mut ffi::sqlite3_context, col: libc::c_int) -> Result<()> { |     fn column(&self, ctx: *mut ffi::sqlite3_context, col: libc::c_int) -> Result<()> { | ||||||
|         use functions::ToResult; |         use functions::ToResult; | ||||||
|         let vtab = self.vtab(); |         if col < 0 || col as usize >= self.cols.len() { | ||||||
|         if col < 0 || col as usize >= vtab.cols.len() { |  | ||||||
|             return Err(Error::ModuleError(format!("column index out of bounds: {}", col))); |             return Err(Error::ModuleError(format!("column index out of bounds: {}", col))); | ||||||
|         } |         } | ||||||
|         if vtab.cols.is_empty() { |         if self.cols.is_empty() { | ||||||
|             unsafe { Null.set_result(ctx) }; |             unsafe { Null.set_result(ctx) }; | ||||||
|             return Ok(()); |             return Ok(()); | ||||||
|         } |         } | ||||||
|         // TODO Affinity |         // TODO Affinity | ||||||
|         unsafe { vtab.cols[col as usize].set_result(ctx) }; |         unsafe { self.cols[col as usize].set_result(ctx) }; | ||||||
|         Ok(()) |         Ok(()) | ||||||
|     } |     } | ||||||
|     fn rowid(&self) -> Result<i64> { |     fn rowid(&self) -> Result<i64> { | ||||||
| @@ -211,6 +230,7 @@ mod test { | |||||||
|         csvtab::load_module(&db).unwrap(); |         csvtab::load_module(&db).unwrap(); | ||||||
|         db.execute_batch("CREATE VIRTUAL TABLE vtab USING csv('test.csv', HAS_HEADERS)").unwrap(); |         db.execute_batch("CREATE VIRTUAL TABLE vtab USING csv('test.csv', HAS_HEADERS)").unwrap(); | ||||||
|  |  | ||||||
|  |         { | ||||||
|             let mut s = db.prepare("SELECT rowid, * FROM vtab").unwrap(); |             let mut s = db.prepare("SELECT rowid, * FROM vtab").unwrap(); | ||||||
|             { |             { | ||||||
|                 let headers = s.column_names(); |                 let headers = s.column_names(); | ||||||
| @@ -222,11 +242,27 @@ mod test { | |||||||
|             for row in rows { |             for row in rows { | ||||||
|                 let row = row.unwrap(); |                 let row = row.unwrap(); | ||||||
|                 let id: i64 = row.get(0); |                 let id: i64 = row.get(0); | ||||||
|             // println!("{}, {:?}, {:?}, {:?}", id, row.get::<i32, Value>(1), row.get::<i32, Value>(2), row.get::<i32, Value>(3)); |  | ||||||
|                 sum = sum + id; |                 sum = sum + id; | ||||||
|             } |             } | ||||||
|             assert_eq!(sum, 15); |             assert_eq!(sum, 15); | ||||||
|  |         } | ||||||
|  |         db.execute_batch("DROP TABLE vtab").unwrap(); | ||||||
|  |     } | ||||||
|  |  | ||||||
|  |     #[test] | ||||||
|  |     fn test_csv_cursor() { | ||||||
|  |         let db = Connection::open_in_memory().unwrap(); | ||||||
|  |         csvtab::load_module(&db).unwrap(); | ||||||
|  |         db.execute_batch("CREATE VIRTUAL TABLE vtab USING csv('test.csv', HAS_HEADERS)").unwrap(); | ||||||
|  |  | ||||||
|  |         { | ||||||
|  |             let mut s = db.prepare("SELECT v1.rowid, v1.* FROM vtab v1 NATURAL JOIN vtab v2 \ | ||||||
|  |                                     WHERE v1.rowid < v2.rowid") | ||||||
|  |                           .unwrap(); | ||||||
|  |  | ||||||
|  |             let row = s.query(&[]).unwrap().next().unwrap().unwrap(); | ||||||
|  |             assert_eq!(row.get::<i32, i32>(0), 2); | ||||||
|  |         } | ||||||
|         db.execute_batch("DROP TABLE vtab").unwrap(); |         db.execute_batch("DROP TABLE vtab").unwrap(); | ||||||
|     } |     } | ||||||
| } | } | ||||||
|   | |||||||
| @@ -26,8 +26,7 @@ pub fn create_int_array(conn: &Connection, name: &str) -> Result<Rc<RefCell<Vec< | |||||||
| /// If not done explicitly by the application, the virtual table will be dropped implicitly | /// If not done explicitly by the application, the virtual table will be dropped implicitly | ||||||
| /// by the system when the database connection is closed. | /// by the system when the database connection is closed. | ||||||
| pub fn drop_int_array(conn: &Connection, name: &str) -> Result<()> { | pub fn drop_int_array(conn: &Connection, name: &str) -> Result<()> { | ||||||
|     conn.execute_batch(&format!("DROP TABLE temp.\"{0}\"", |     conn.execute_batch(&format!("DROP TABLE temp.\"{0}\"", escape_double_quote(name))) | ||||||
|                                 escape_double_quote(name))) |  | ||||||
| } | } | ||||||
|  |  | ||||||
| init_module!(INT_ARRAY_MODULE, IntArrayVTab, IntArrayVTabCursor, | init_module!(INT_ARRAY_MODULE, IntArrayVTab, IntArrayVTabCursor, | ||||||
|   | |||||||
| @@ -194,7 +194,7 @@ unsafe extern "C" fn $create(db: *mut ffi::sqlite3, | |||||||
|             let boxed_vtab: *mut $vtab = Box::into_raw(Box::new(vtab)); |             let boxed_vtab: *mut $vtab = Box::into_raw(Box::new(vtab)); | ||||||
|             *pp_vtab = boxed_vtab as *mut ffi::sqlite3_vtab; |             *pp_vtab = boxed_vtab as *mut ffi::sqlite3_vtab; | ||||||
|             ffi::SQLITE_OK |             ffi::SQLITE_OK | ||||||
|         } |         }, | ||||||
|         Err(Error::SqliteFailure(err, s)) => { |         Err(Error::SqliteFailure(err, s)) => { | ||||||
|             if let Some(s) = s { |             if let Some(s) = s { | ||||||
|                 *err_msg = mprintf(&s); |                 *err_msg = mprintf(&s); | ||||||
| @@ -237,7 +237,7 @@ unsafe extern "C" fn $open(vtab: *mut ffi::sqlite3_vtab, | |||||||
|                 set_err_msg(vtab, &err_msg); |                 set_err_msg(vtab, &err_msg); | ||||||
|             } |             } | ||||||
|             err.extended_code |             err.extended_code | ||||||
|         } |         }, | ||||||
|         Err(err) => { |         Err(err) => { | ||||||
|             set_err_msg(vtab, err.description()); |             set_err_msg(vtab, err.description()); | ||||||
|             ffi::SQLITE_ERROR |             ffi::SQLITE_ERROR | ||||||
|   | |||||||
		Reference in New Issue
	
	Block a user