support first_row_only

This commit is contained in:
Damien Elmes 2020-03-06 10:11:05 +10:00
parent e14c5e4745
commit db1508e27c
3 changed files with 44 additions and 5 deletions

View File

@ -67,8 +67,7 @@ class DBProxy:
self.mod = True self.mod = True
assert ":" not in sql assert ":" not in sql
# fetch rows # fetch rows
# fixme: first_row_only return self._backend.db_query(sql, args, first_row_only)
return self._backend.db_query(sql, args)
# Query shortcuts # Query shortcuts
################### ###################

View File

@ -386,8 +386,12 @@ class RustBackend:
def restore_trash(self): def restore_trash(self):
self._run_command(pb.BackendInput(restore_trash=pb.Empty())) self._run_command(pb.BackendInput(restore_trash=pb.Empty()))
def db_query(self, sql: str, args: List[ValueForDB]) -> List[DBRow]: def db_query(
return self._db_command(dict(kind="query", sql=sql, args=args)) self, sql: str, args: List[ValueForDB], first_row_only: bool
) -> List[DBRow]:
return self._db_command(
dict(kind="query", sql=sql, args=args, first_row_only=first_row_only)
)
def db_execute_many(self, sql: str, args: List[List[ValueForDB]]) -> List[DBRow]: def db_execute_many(self, sql: str, args: List[List[ValueForDB]]) -> List[DBRow]:
return self._db_command(dict(kind="executemany", sql=sql, args=args)) return self._db_command(dict(kind="executemany", sql=sql, args=args))

View File

@ -4,6 +4,7 @@
use crate::err::Result; use crate::err::Result;
use crate::storage::StorageContext; use crate::storage::StorageContext;
use rusqlite::types::{FromSql, FromSqlError, ToSql, ToSqlOutput, ValueRef}; use rusqlite::types::{FromSql, FromSqlError, ToSql, ToSqlOutput, ValueRef};
use rusqlite::OptionalExtension;
use serde_derive::{Deserialize, Serialize}; use serde_derive::{Deserialize, Serialize};
#[derive(Deserialize)] #[derive(Deserialize)]
@ -12,6 +13,7 @@ pub(super) enum DBRequest {
Query { Query {
sql: String, sql: String,
args: Vec<SqlValue>, args: Vec<SqlValue>,
first_row_only: bool,
}, },
Begin, Begin,
Commit, Commit,
@ -68,7 +70,17 @@ impl FromSql for SqlValue {
pub(super) fn db_command_bytes(ctx: &StorageContext, input: &[u8]) -> Result<String> { pub(super) fn db_command_bytes(ctx: &StorageContext, input: &[u8]) -> Result<String> {
let req: DBRequest = serde_json::from_slice(input)?; let req: DBRequest = serde_json::from_slice(input)?;
let resp = match req { let resp = match req {
DBRequest::Query { sql, args } => db_query(ctx, &sql, &args)?, 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 => { DBRequest::Begin => {
ctx.begin_trx()?; ctx.begin_trx()?;
DBResult::None DBResult::None
@ -86,6 +98,30 @@ pub(super) fn db_command_bytes(ctx: &StorageContext, input: &[u8]) -> Result<Str
Ok(serde_json::to_string(&resp)?) Ok(serde_json::to_string(&resp)?)
} }
pub(super) fn db_query_row(ctx: &StorageContext, 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: &StorageContext, sql: &str, args: &[SqlValue]) -> Result<DBResult> { pub(super) fn db_query(ctx: &StorageContext, sql: &str, args: &[SqlValue]) -> Result<DBResult> {
let mut stmt = ctx.db.prepare_cached(sql)?; let mut stmt = ctx.db.prepare_cached(sql)?;
let columns = stmt.column_count(); let columns = stmt.column_count();