performance improvements for SQLite reads (#6204)

This commit is contained in:
Reilly Wood 2022-07-31 23:09:03 -07:00 committed by GitHub
parent 01386f4d58
commit 2ac7a4d48d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -425,18 +425,8 @@ fn open_sqlite_db(path: &Path, call_span: Span) -> Result<Connection, nu_protoco
} }
fn run_sql_query(conn: Connection, sql: &Spanned<String>) -> Result<Value, rusqlite::Error> { fn run_sql_query(conn: Connection, sql: &Spanned<String>) -> Result<Value, rusqlite::Error> {
let mut stmt = conn.prepare(&sql.item)?; let stmt = conn.prepare(&sql.item)?;
let results = stmt.query([])?; prepared_statement_to_nu_list(stmt, sql.span)
let nu_records = results
.mapped(|row| Result::Ok(convert_sqlite_row_to_nu_value(row, sql.span)))
.into_iter()
.collect::<Result<Vec<Value>, rusqlite::Error>>()?;
Ok(Value::List {
vals: nu_records,
span: sql.span,
})
} }
fn read_single_table( fn read_single_table(
@ -444,14 +434,30 @@ fn read_single_table(
table_name: String, table_name: String,
call_span: Span, call_span: Span,
) -> Result<Value, rusqlite::Error> { ) -> Result<Value, rusqlite::Error> {
let mut stmt = conn.prepare(&format!("SELECT * FROM {}", table_name))?; let stmt = conn.prepare(&format!("SELECT * FROM {}", table_name))?;
let results = stmt.query([])?; prepared_statement_to_nu_list(stmt, call_span)
}
fn prepared_statement_to_nu_list(
mut stmt: rusqlite::Statement,
call_span: Span,
) -> Result<Value, rusqlite::Error> {
let column_names = stmt
.column_names()
.iter()
.map(|c| c.to_string())
.collect::<Vec<String>>();
let results = stmt.query([])?;
let nu_records = results let nu_records = results
.mapped(|row| Result::Ok(convert_sqlite_row_to_nu_value(row, call_span))) .mapped(|row| {
Result::Ok(convert_sqlite_row_to_nu_value(
row,
call_span,
column_names.clone(),
))
})
.into_iter() .into_iter()
.collect::<Result<Vec<Value>, rusqlite::Error>>()?; .collect::<Result<Vec<Value>, rusqlite::Error>>()?;
Ok(Value::List { Ok(Value::List {
vals: nu_records, vals: nu_records,
span: call_span, span: call_span,
@ -470,19 +476,9 @@ fn read_entire_sqlite_db(conn: Connection, call_span: Span) -> Result<Value, rus
let table_name: String = row?; let table_name: String = row?;
table_names.push(table_name.clone()); table_names.push(table_name.clone());
let mut rows = Vec::new(); let table_stmt = conn.prepare(&format!("select * from [{}]", table_name))?;
let mut table_stmt = conn.prepare(&format!("select * from [{}]", table_name))?; let rows = prepared_statement_to_nu_list(table_stmt, call_span)?;
let mut table_rows = table_stmt.query([])?; tables.push(rows);
while let Some(table_row) = table_rows.next()? {
rows.push(convert_sqlite_row_to_nu_value(table_row, call_span))
}
let table_record = Value::List {
vals: rows,
span: call_span,
};
tables.push(table_record);
} }
Ok(Value::Record { Ok(Value::Record {
@ -492,19 +488,16 @@ fn read_entire_sqlite_db(conn: Connection, call_span: Span) -> Result<Value, rus
}) })
} }
pub fn convert_sqlite_row_to_nu_value(row: &Row, span: Span) -> Value { pub fn convert_sqlite_row_to_nu_value(row: &Row, span: Span, column_names: Vec<String>) -> Value {
let mut vals = Vec::new(); let mut vals = Vec::with_capacity(column_names.len());
let colnamestr = row.as_ref().column_names().to_vec();
let colnames = colnamestr.iter().map(|s| s.to_string()).collect();
for (i, c) in row.as_ref().column_names().iter().enumerate() { for i in 0..column_names.len() {
let _column = c.to_string();
let val = convert_sqlite_value_to_nu_value(row.get_ref_unwrap(i), span); let val = convert_sqlite_value_to_nu_value(row.get_ref_unwrap(i), span);
vals.push(val); vals.push(val);
} }
Value::Record { Value::Record {
cols: colnames, cols: column_names,
vals, vals,
span, span,
} }