diff --git a/src/vtab/csvtab.rs b/src/vtab/csvtab.rs index a9f4e4f..e1aa21c 100644 --- a/src/vtab/csvtab.rs +++ b/src/vtab/csvtab.rs @@ -9,7 +9,7 @@ use libc; use {Connection, Error, Result}; use ffi; use types::Null; -use vtab::{declare_vtab, VTab, VTabCursor}; +use vtab::{declare_vtab, escape_double_quote, VTab, VTabCursor}; pub fn load_module(conn: &Connection) -> Result<()> { let aux: Option<()> = None; @@ -29,6 +29,7 @@ struct CSVTab { reader: csv::Reader, offset_first_row: u64, cols: Vec, + eof: bool, } impl VTab for CSVTab { @@ -63,7 +64,7 @@ impl VTab for CSVTab { } else if uc.contains("NO_QUOTE") { reader = reader.quote(0); } else { - cols.push(String::from(arg)); + cols.push(escape_double_quote(String::from(arg))); } } } @@ -102,6 +103,7 @@ impl VTab for CSVTab { reader: reader, offset_first_row: offset_first_row, cols: cols, + eof: false, }; try!(declare_vtab(db, &sql)); Ok(vtab) @@ -137,7 +139,12 @@ impl VTabCursor for CSVTabCursor { unsafe { &mut *(self.base.pVtab as *mut CSVTab) } } - fn filter(&mut self) -> Result<()> { + fn filter(&mut self, + _idx_num: libc::c_int, + _idx_str: *const libc::c_char, + _argc: libc::c_int, + _argv: *mut *mut ffi::sqlite3_value) + -> Result<()> { { let vtab = self.vtab(); try!(vtab.reader.seek(vtab.offset_first_row)); @@ -148,8 +155,9 @@ impl VTabCursor for CSVTabCursor { fn next(&mut self) -> Result<()> { { let vtab = self.vtab(); - if vtab.reader.done() { - return Err(Error::ModuleError(format!("eof"))); + vtab.eof = vtab.reader.done(); + if vtab.eof { + return Ok(()); } vtab.cols.clear(); @@ -163,7 +171,7 @@ impl VTabCursor for CSVTabCursor { } fn eof(&self) -> bool { let vtab = self.vtab(); - vtab.reader.done() + vtab.eof } fn column(&self, ctx: *mut ffi::sqlite3_context, col: libc::c_int) -> Result<()> { use functions::ToResult; @@ -190,3 +198,34 @@ impl From for Error { Error::ModuleError(String::from(err.description())) } } + +#[cfg(test)] +mod test { + use Connection; + use vtab::csvtab; + + #[test] + fn test_csv_module() { + 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 rowid, * FROM vtab").unwrap(); + { + let headers = s.column_names(); + assert_eq!(vec!["rowid", "colA", "colB", "colC"], headers); + } + + let rows = s.query(&[]).unwrap(); + let mut sum = 0; + for row in rows { + let row = row.unwrap(); + let id: i64 = row.get(0); + // println!("{}, {:?}, {:?}, {:?}", id, row.get::(1), row.get::(2), row.get::(3)); + sum = sum + id; + } + assert_eq!(sum, 15); + + db.execute_batch("DROP TABLE vtab").unwrap(); + } +} diff --git a/src/vtab/int_array.rs b/src/vtab/int_array.rs index 17fde9a..adad0a8 100644 --- a/src/vtab/int_array.rs +++ b/src/vtab/int_array.rs @@ -7,27 +7,19 @@ use libc; use {Connection, Error, Result}; use ffi; -use vtab::{declare_vtab, VTab, VTabCursor}; +use vtab::{declare_vtab, escape_double_quote, VTab, VTabCursor}; pub fn create_int_array(conn: &Connection, name: &str) -> Result>>> { let array = Rc::new(RefCell::new(Vec::new())); try!(conn.create_module(name, &INT_ARRAY_MODULE, Some(array.clone()))); - try!(conn.execute_batch(&format!("CREATE VIRTUAL TABLE temp.{0} USING {0}", - escape_quote(name.to_string())))); + try!(conn.execute_batch(&format!("CREATE VIRTUAL TABLE temp.\"{0}\" USING \"{0}\"", + escape_double_quote(name.to_string())))); Ok(array) } pub fn drop_int_array(conn: &Connection, name: &str) -> Result<()> { - conn.execute_batch(&format!("DROP TABLE temp.{0}", escape_quote(name.to_string()))) -} - -fn escape_quote(identifier: String) -> String { - if identifier.contains('"') { - // escape quote by doubling them - identifier.replace("\"", "\"\"") - } else { - identifier - } + conn.execute_batch(&format!("DROP TABLE temp.\"{0}\"", + escape_double_quote(name.to_string()))) } init_module!(INT_ARRAY_MODULE, IntArrayVTab, IntArrayVTabCursor, @@ -85,7 +77,12 @@ impl VTabCursor for IntArrayVTabCursor { fn vtab(&self) -> &mut IntArrayVTab { unsafe { &mut *(self.base.pVtab as *mut IntArrayVTab) } } - fn filter(&mut self) -> Result<()> { + fn filter(&mut self, + _idx_num: libc::c_int, + _idx_str: *const libc::c_char, + _argc: libc::c_int, + _argv: *mut *mut ffi::sqlite3_value) + -> Result<()> { self.i = 0; Ok(()) } diff --git a/src/vtab/mod.rs b/src/vtab/mod.rs index 6ecb36a..b7b2cde 100644 --- a/src/vtab/mod.rs +++ b/src/vtab/mod.rs @@ -37,21 +37,39 @@ use ffi; // \-> if not eof { cursor.column or xrowid } else { cursor.xclose } // +/// Virtual table instance trait. pub trait VTab>: Sized { + /// Create a new instance of a virtual table in response to a CREATE VIRTUAL TABLE statement. + /// The `db` parameter is a pointer to the SQLite database connection that is executing the CREATE VIRTUAL TABLE statement. fn create(db: *mut ffi::sqlite3, aux: *mut libc::c_void, args: &[*const libc::c_char]) -> Result; + /// Determine the best way to access the virtual table. fn best_index(&self, info: *mut ffi::sqlite3_index_info); + /// Create a new cursor used for accessing a virtual table. fn open(&self) -> Result; } +/// Virtual table cursor trait. pub trait VTabCursor>: Sized { + /// Accessor to the associated virtual table. fn vtab(&self) -> &mut V; - fn filter(&mut self) -> Result<()>; + /// Begin a search of a virtual table. + fn filter(&mut self, + idx_num: libc::c_int, + idx_str: *const libc::c_char, + argc: libc::c_int, + argv: *mut *mut ffi::sqlite3_value) + -> Result<()>; + /// Advance cursor to the next row of a result set initiated by `filter`. fn next(&mut self) -> Result<()>; + /// Must return `false` if the cursor currently points to a valid row of data, or `true` otherwise. fn eof(&self) -> bool; + /// Find the value for the `i`-th column of the current row. `i` is zero-based so the first column is numbered 0. + /// May return its result back to SQLite using one of the specified `ctx`. fn column(&self, ctx: *mut ffi::sqlite3_context, i: libc::c_int) -> Result<()>; + /// Return the rowid of row that the cursor is currently pointing at. fn rowid(&self) -> Result; } @@ -98,6 +116,7 @@ impl InnerConnection { } } +/// Declare the schema of a virtual table. pub fn declare_vtab(db: *mut ffi::sqlite3, sql: &str) -> Result<()> { let c_sql = try!(CString::new(sql)); let rc = unsafe { ffi::sqlite3_declare_vtab(db, c_sql.as_ptr()) }; @@ -108,6 +127,16 @@ pub fn declare_vtab(db: *mut ffi::sqlite3, sql: &str) -> Result<()> { } } +/// Escape double-quote (`"`) character occurences by doubling them (`""`). +pub fn escape_double_quote(identifier: String) -> String { + if identifier.contains('"') { + // escape quote by doubling them + identifier.replace("\"", "\"\"") + } else { + identifier + } +} + // FIXME copy/paste from function.rs unsafe extern "C" fn free_boxed_value(p: *mut libc::c_void) { let _: Box = Box::from_raw(mem::transmute(p)); @@ -220,14 +249,14 @@ unsafe extern "C" fn $close(cursor: *mut ffi::sqlite3_vtab_cursor) -> libc::c_in } unsafe extern "C" fn $filter(cursor: *mut ffi::sqlite3_vtab_cursor, - _idx_num: libc::c_int, - _idx_str: *const libc::c_char, - _argc: libc::c_int, - _argv: *mut *mut ffi::sqlite3_value) + idx_num: libc::c_int, + idx_str: *const libc::c_char, + argc: libc::c_int, + argv: *mut *mut ffi::sqlite3_value) -> libc::c_int { use vtab::cursor_error; let cr = cursor as *mut $cursor; - cursor_error(cursor, (*cr).filter()) + cursor_error(cursor, (*cr).filter(idx_num, idx_str, argc, argv)) } unsafe extern "C" fn $next(cursor: *mut ffi::sqlite3_vtab_cursor) -> libc::c_int { use vtab::cursor_error; @@ -242,9 +271,9 @@ unsafe extern "C" fn $column(cursor: *mut ffi::sqlite3_vtab_cursor, ctx: *mut ffi::sqlite3_context, i: libc::c_int) -> libc::c_int { - use vtab::cursor_error; + use vtab::result_error; let cr = cursor as *mut $cursor; - cursor_error(cursor, (*cr).column(ctx, i)) + result_error(ctx, (*cr).column(ctx, i)) } unsafe extern "C" fn $rowid(cursor: *mut ffi::sqlite3_vtab_cursor, p_rowid: *mut ffi::sqlite3_int64) @@ -262,6 +291,7 @@ unsafe extern "C" fn $rowid(cursor: *mut ffi::sqlite3_vtab_cursor, } } +/// Virtual table cursors can set an error message by assigning a string to `zErrMsg`. pub unsafe fn cursor_error(cursor: *mut ffi::sqlite3_vtab_cursor, result: Result) -> libc::c_int { @@ -281,6 +311,7 @@ pub unsafe fn cursor_error(cursor: *mut ffi::sqlite3_vtab_cursor, } } +/// Virtual tables methods can set an error message by assigning a string to `zErrMsg`. pub unsafe fn set_err_msg(vtab: *mut ffi::sqlite3_vtab, err_msg: &str) { if !(*vtab).zErrMsg.is_null() { ffi::sqlite3_free((*vtab).zErrMsg as *mut libc::c_void); @@ -288,22 +319,34 @@ pub unsafe fn set_err_msg(vtab: *mut ffi::sqlite3_vtab, err_msg: &str) { (*vtab).zErrMsg = mprintf(err_msg); } -unsafe fn result_error(ctx: *mut ffi::sqlite3_context, err: Error) -> libc::c_int { +/// To raise an error, the `column` method should use this method to set the error message and return the error code. +pub unsafe fn result_error(ctx: *mut ffi::sqlite3_context, result: Result) -> libc::c_int { use std::error::Error as StdError; - match err { - Error::SqliteFailure(err, s) => { - ffi::sqlite3_result_error_code(ctx, err.extended_code); - if let Some(Ok(cstr)) = s.map(|s| str_to_cstring(&s)) { - ffi::sqlite3_result_error(ctx, cstr.as_ptr(), -1); - } + match result { + Ok(_) => ffi::SQLITE_OK, + Err(Error::SqliteFailure(err, s)) => { + match err.extended_code { + ffi::SQLITE_TOOBIG => { + ffi::sqlite3_result_error_toobig(ctx); + } + ffi::SQLITE_NOMEM => { + ffi::sqlite3_result_error_nomem(ctx); + } + code => { + ffi::sqlite3_result_error_code(ctx, code); + if let Some(Ok(cstr)) = s.map(|s| str_to_cstring(&s)) { + ffi::sqlite3_result_error(ctx, cstr.as_ptr(), -1); + } + } + }; err.extended_code } - _ => { - ffi::sqlite3_result_error_code(ctx, ffi::SQLITE_CORRUPT_VTAB); + Err(err) => { + ffi::sqlite3_result_error_code(ctx, ffi::SQLITE_ERROR); if let Ok(cstr) = str_to_cstring(err.description()) { ffi::sqlite3_result_error(ctx, cstr.as_ptr(), -1); } - ffi::SQLITE_CORRUPT_VTAB + ffi::SQLITE_ERROR } } } diff --git a/test.csv b/test.csv new file mode 100644 index 0000000..708f93f --- /dev/null +++ b/test.csv @@ -0,0 +1,6 @@ +"colA","colB","colC" +1,2,3 +a,b,c +a,"b",c +"a","b","c .. z" +"a","b","c,d"