anki/rslib/src/backend/dbproxy.rs
2020-03-29 12:12:35 +10:00

156 lines
4.1 KiB
Rust

// Copyright: Ankitects Pty Ltd and contributors
// License: GNU AGPL, version 3 or later; http://www.gnu.org/licenses/agpl.html
use crate::err::Result;
use crate::storage::SqliteStorage;
use rusqlite::types::{FromSql, FromSqlError, ToSql, ToSqlOutput, ValueRef};
use rusqlite::OptionalExtension;
use serde_derive::{Deserialize, Serialize};
#[derive(Deserialize)]
#[serde(tag = "kind", rename_all = "lowercase")]
pub(super) enum DBRequest {
Query {
sql: String,
args: Vec<SqlValue>,
first_row_only: bool,
},
Begin,
Commit,
Rollback,
ExecuteMany {
sql: String,
args: Vec<Vec<SqlValue>>,
},
}
#[derive(Serialize)]
#[serde(untagged)]
pub(super) enum DBResult {
Rows(Vec<Vec<SqlValue>>),
None,
}
#[derive(Serialize, Deserialize, Debug)]
#[serde(untagged)]
pub(super) enum SqlValue {
Null,
String(String),
Int(i64),
Double(f64),
Blob(Vec<u8>),
}
impl ToSql for SqlValue {
fn to_sql(&self) -> std::result::Result<ToSqlOutput<'_>, rusqlite::Error> {
let val = match self {
SqlValue::Null => ValueRef::Null,
SqlValue::String(v) => ValueRef::Text(v.as_bytes()),
SqlValue::Int(v) => ValueRef::Integer(*v),
SqlValue::Double(v) => ValueRef::Real(*v),
SqlValue::Blob(v) => ValueRef::Blob(&v),
};
Ok(ToSqlOutput::Borrowed(val))
}
}
impl FromSql for SqlValue {
fn column_result(value: ValueRef<'_>) -> std::result::Result<Self, FromSqlError> {
let val = match value {
ValueRef::Null => SqlValue::Null,
ValueRef::Integer(i) => SqlValue::Int(i),
ValueRef::Real(v) => SqlValue::Double(v),
ValueRef::Text(v) => SqlValue::String(String::from_utf8_lossy(v).to_string()),
ValueRef::Blob(v) => SqlValue::Blob(v.to_vec()),
};
Ok(val)
}
}
pub(super) fn db_command_bytes(ctx: &SqliteStorage, input: &[u8]) -> Result<String> {
let req: DBRequest = serde_json::from_slice(input)?;
let resp = match req {
DBRequest::Query {
sql,
args,
first_row_only,
} => {
if first_row_only {
db_query_row(ctx, &sql, &args)?
} else {
db_query(ctx, &sql, &args)?
}
}
DBRequest::Begin => {
ctx.begin_trx()?;
DBResult::None
}
DBRequest::Commit => {
ctx.commit_trx()?;
DBResult::None
}
DBRequest::Rollback => {
ctx.rollback_trx()?;
DBResult::None
}
DBRequest::ExecuteMany { sql, args } => db_execute_many(ctx, &sql, &args)?,
};
Ok(serde_json::to_string(&resp)?)
}
pub(super) fn db_query_row(ctx: &SqliteStorage, sql: &str, args: &[SqlValue]) -> Result<DBResult> {
let mut stmt = ctx.db.prepare_cached(sql)?;
let columns = stmt.column_count();
let row = stmt
.query_row(args, |row| {
let mut orow = Vec::with_capacity(columns);
for i in 0..columns {
let v: SqlValue = row.get(i)?;
orow.push(v);
}
Ok(orow)
})
.optional()?;
let rows = if let Some(row) = row {
vec![row]
} else {
vec![]
};
Ok(DBResult::Rows(rows))
}
pub(super) fn db_query(ctx: &SqliteStorage, sql: &str, args: &[SqlValue]) -> Result<DBResult> {
let mut stmt = ctx.db.prepare_cached(sql)?;
let columns = stmt.column_count();
let res: std::result::Result<Vec<Vec<_>>, rusqlite::Error> = stmt
.query_map(args, |row| {
let mut orow = Vec::with_capacity(columns);
for i in 0..columns {
let v: SqlValue = row.get(i)?;
orow.push(v);
}
Ok(orow)
})?
.collect();
Ok(DBResult::Rows(res?))
}
pub(super) fn db_execute_many(
ctx: &SqliteStorage,
sql: &str,
args: &[Vec<SqlValue>],
) -> Result<DBResult> {
let mut stmt = ctx.db.prepare_cached(sql)?;
for params in args {
stmt.execute(params)?;
}
Ok(DBResult::None)
}