First minimalist test for csvtab implementation.

This commit is contained in:
gwenn 2016-02-11 19:16:05 +01:00
parent 3fbeccdf55
commit 38209ffef3
4 changed files with 123 additions and 38 deletions

View File

@ -9,7 +9,7 @@ use libc;
use {Connection, Error, Result}; use {Connection, Error, Result};
use ffi; use ffi;
use types::Null; 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<()> { pub fn load_module(conn: &Connection) -> Result<()> {
let aux: Option<()> = None; let aux: Option<()> = None;
@ -29,6 +29,7 @@ struct CSVTab {
reader: csv::Reader<File>, reader: csv::Reader<File>,
offset_first_row: u64, offset_first_row: u64,
cols: Vec<String>, cols: Vec<String>,
eof: bool,
} }
impl VTab<CSVTabCursor> for CSVTab { impl VTab<CSVTabCursor> for CSVTab {
@ -63,7 +64,7 @@ impl VTab<CSVTabCursor> for CSVTab {
} else if uc.contains("NO_QUOTE") { } else if uc.contains("NO_QUOTE") {
reader = reader.quote(0); reader = reader.quote(0);
} else { } else {
cols.push(String::from(arg)); cols.push(escape_double_quote(String::from(arg)));
} }
} }
} }
@ -102,6 +103,7 @@ impl VTab<CSVTabCursor> for CSVTab {
reader: reader, reader: reader,
offset_first_row: offset_first_row, offset_first_row: offset_first_row,
cols: cols, cols: cols,
eof: false,
}; };
try!(declare_vtab(db, &sql)); try!(declare_vtab(db, &sql));
Ok(vtab) Ok(vtab)
@ -137,7 +139,12 @@ impl VTabCursor<CSVTab> for CSVTabCursor {
unsafe { &mut *(self.base.pVtab as *mut CSVTab) } 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(); let vtab = self.vtab();
try!(vtab.reader.seek(vtab.offset_first_row)); try!(vtab.reader.seek(vtab.offset_first_row));
@ -148,8 +155,9 @@ impl VTabCursor<CSVTab> for CSVTabCursor {
fn next(&mut self) -> Result<()> { fn next(&mut self) -> Result<()> {
{ {
let vtab = self.vtab(); let vtab = self.vtab();
if vtab.reader.done() { vtab.eof = vtab.reader.done();
return Err(Error::ModuleError(format!("eof"))); if vtab.eof {
return Ok(());
} }
vtab.cols.clear(); vtab.cols.clear();
@ -163,7 +171,7 @@ impl VTabCursor<CSVTab> for CSVTabCursor {
} }
fn eof(&self) -> bool { fn eof(&self) -> bool {
let vtab = self.vtab(); let vtab = self.vtab();
vtab.reader.done() 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;
@ -190,3 +198,34 @@ impl From<csv::Error> for Error {
Error::ModuleError(String::from(err.description())) 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::<i32, Value>(1), row.get::<i32, Value>(2), row.get::<i32, Value>(3));
sum = sum + id;
}
assert_eq!(sum, 15);
db.execute_batch("DROP TABLE vtab").unwrap();
}
}

View File

@ -7,27 +7,19 @@ use libc;
use {Connection, Error, Result}; use {Connection, Error, Result};
use ffi; 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<Rc<RefCell<Vec<i64>>>> { pub fn create_int_array(conn: &Connection, name: &str) -> Result<Rc<RefCell<Vec<i64>>>> {
let array = Rc::new(RefCell::new(Vec::new())); let array = Rc::new(RefCell::new(Vec::new()));
try!(conn.create_module(name, &INT_ARRAY_MODULE, Some(array.clone()))); try!(conn.create_module(name, &INT_ARRAY_MODULE, Some(array.clone())));
try!(conn.execute_batch(&format!("CREATE VIRTUAL TABLE temp.{0} USING {0}", try!(conn.execute_batch(&format!("CREATE VIRTUAL TABLE temp.\"{0}\" USING \"{0}\"",
escape_quote(name.to_string())))); escape_double_quote(name.to_string()))));
Ok(array) Ok(array)
} }
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}", escape_quote(name.to_string()))) conn.execute_batch(&format!("DROP TABLE temp.\"{0}\"",
} escape_double_quote(name.to_string())))
fn escape_quote(identifier: String) -> String {
if identifier.contains('"') {
// escape quote by doubling them
identifier.replace("\"", "\"\"")
} else {
identifier
}
} }
init_module!(INT_ARRAY_MODULE, IntArrayVTab, IntArrayVTabCursor, init_module!(INT_ARRAY_MODULE, IntArrayVTab, IntArrayVTabCursor,
@ -85,7 +77,12 @@ impl VTabCursor<IntArrayVTab> for IntArrayVTabCursor {
fn vtab(&self) -> &mut IntArrayVTab { fn vtab(&self) -> &mut IntArrayVTab {
unsafe { &mut *(self.base.pVtab as *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; self.i = 0;
Ok(()) Ok(())
} }

View File

@ -37,21 +37,39 @@ use ffi;
// \-> if not eof { cursor.column or xrowid } else { cursor.xclose } // \-> if not eof { cursor.column or xrowid } else { cursor.xclose }
// //
/// Virtual table instance trait.
pub trait VTab<C: VTabCursor<Self>>: Sized { pub trait VTab<C: VTabCursor<Self>>: 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, fn create(db: *mut ffi::sqlite3,
aux: *mut libc::c_void, aux: *mut libc::c_void,
args: &[*const libc::c_char]) args: &[*const libc::c_char])
-> Result<Self>; -> Result<Self>;
/// Determine the best way to access the virtual table.
fn best_index(&self, info: *mut ffi::sqlite3_index_info); fn best_index(&self, info: *mut ffi::sqlite3_index_info);
/// Create a new cursor used for accessing a virtual table.
fn open(&self) -> Result<C>; fn open(&self) -> Result<C>;
} }
/// Virtual table cursor trait.
pub trait VTabCursor<V: VTab<Self>>: Sized { pub trait VTabCursor<V: VTab<Self>>: Sized {
/// Accessor to the associated virtual table.
fn vtab(&self) -> &mut V; 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<()>; 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; 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<()>; 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<i64>; fn rowid(&self) -> Result<i64>;
} }
@ -98,6 +116,7 @@ impl InnerConnection {
} }
} }
/// Declare the schema of a virtual table.
pub fn declare_vtab(db: *mut ffi::sqlite3, sql: &str) -> Result<()> { pub fn declare_vtab(db: *mut ffi::sqlite3, sql: &str) -> Result<()> {
let c_sql = try!(CString::new(sql)); let c_sql = try!(CString::new(sql));
let rc = unsafe { ffi::sqlite3_declare_vtab(db, c_sql.as_ptr()) }; 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 // FIXME copy/paste from function.rs
unsafe extern "C" fn free_boxed_value<T>(p: *mut libc::c_void) { unsafe extern "C" fn free_boxed_value<T>(p: *mut libc::c_void) {
let _: Box<T> = Box::from_raw(mem::transmute(p)); let _: Box<T> = 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, unsafe extern "C" fn $filter(cursor: *mut ffi::sqlite3_vtab_cursor,
_idx_num: libc::c_int, idx_num: libc::c_int,
_idx_str: *const libc::c_char, idx_str: *const libc::c_char,
_argc: libc::c_int, argc: libc::c_int,
_argv: *mut *mut ffi::sqlite3_value) argv: *mut *mut ffi::sqlite3_value)
-> libc::c_int { -> libc::c_int {
use vtab::cursor_error; use vtab::cursor_error;
let cr = cursor as *mut $cursor; 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 { unsafe extern "C" fn $next(cursor: *mut ffi::sqlite3_vtab_cursor) -> libc::c_int {
use vtab::cursor_error; use vtab::cursor_error;
@ -242,9 +271,9 @@ unsafe extern "C" fn $column(cursor: *mut ffi::sqlite3_vtab_cursor,
ctx: *mut ffi::sqlite3_context, ctx: *mut ffi::sqlite3_context,
i: libc::c_int) i: libc::c_int)
-> libc::c_int { -> libc::c_int {
use vtab::cursor_error; use vtab::result_error;
let cr = cursor as *mut $cursor; 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, unsafe extern "C" fn $rowid(cursor: *mut ffi::sqlite3_vtab_cursor,
p_rowid: *mut ffi::sqlite3_int64) 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<T>(cursor: *mut ffi::sqlite3_vtab_cursor, pub unsafe fn cursor_error<T>(cursor: *mut ffi::sqlite3_vtab_cursor,
result: Result<T>) result: Result<T>)
-> libc::c_int { -> libc::c_int {
@ -281,6 +311,7 @@ pub unsafe fn cursor_error<T>(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) { pub unsafe fn set_err_msg(vtab: *mut ffi::sqlite3_vtab, err_msg: &str) {
if !(*vtab).zErrMsg.is_null() { if !(*vtab).zErrMsg.is_null() {
ffi::sqlite3_free((*vtab).zErrMsg as *mut libc::c_void); 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); (*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<T>(ctx: *mut ffi::sqlite3_context, result: Result<T>) -> libc::c_int {
use std::error::Error as StdError; use std::error::Error as StdError;
match err { match result {
Error::SqliteFailure(err, s) => { Ok(_) => ffi::SQLITE_OK,
ffi::sqlite3_result_error_code(ctx, err.extended_code); 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)) { if let Some(Ok(cstr)) = s.map(|s| str_to_cstring(&s)) {
ffi::sqlite3_result_error(ctx, cstr.as_ptr(), -1); ffi::sqlite3_result_error(ctx, cstr.as_ptr(), -1);
} }
}
};
err.extended_code err.extended_code
} }
_ => { Err(err) => {
ffi::sqlite3_result_error_code(ctx, ffi::SQLITE_CORRUPT_VTAB); ffi::sqlite3_result_error_code(ctx, ffi::SQLITE_ERROR);
if let Ok(cstr) = str_to_cstring(err.description()) { if let Ok(cstr) = str_to_cstring(err.description()) {
ffi::sqlite3_result_error(ctx, cstr.as_ptr(), -1); ffi::sqlite3_result_error(ctx, cstr.as_ptr(), -1);
} }
ffi::SQLITE_CORRUPT_VTAB ffi::SQLITE_ERROR
} }
} }
} }

6
test.csv Normal file
View File

@ -0,0 +1,6 @@
"colA","colB","colC"
1,2,3
a,b,c
a,"b",c
"a","b","c .. z"
"a","b","c,d"
1 colA colB colC
2 1 2 3
3 a b c
4 a b c
5 a b c .. z
6 a b c,d