309 lines
14 KiB
Rust
309 lines
14 KiB
Rust
/*
|
|
* SPDX-FileCopyrightText: 2020 Stalwart Labs LLC <hello@stalw.art>
|
|
*
|
|
* SPDX-License-Identifier: AGPL-3.0-only OR LicenseRef-SEL
|
|
*/
|
|
|
|
use super::{SqliteStore, into_error};
|
|
use crate::{
|
|
IndexKey, Key, LogKey, SUBSPACE_COUNTER, SUBSPACE_IN_MEMORY_COUNTER, SUBSPACE_QUOTA,
|
|
write::{AssignedIds, Batch, MergeResult, Operation, ValueClass, ValueOp},
|
|
};
|
|
use rusqlite::{OptionalExtension, TransactionBehavior, params};
|
|
use trc::AddContext;
|
|
|
|
impl SqliteStore {
|
|
pub(crate) async fn write(&self, batch: Batch<'_>) -> trc::Result<AssignedIds> {
|
|
let manager = self.conn_pool.clone();
|
|
self.spawn_worker(move || {
|
|
let mut conn = manager.get().map_err(into_error)?;
|
|
|
|
let mut account_id = u32::MAX;
|
|
let mut collection = u8::MAX;
|
|
let mut document_id = u32::MAX;
|
|
let mut change_id = 0u64;
|
|
let trx = conn
|
|
.transaction_with_behavior(TransactionBehavior::Immediate)
|
|
.map_err(into_error)
|
|
.caused_by(trc::location!())?;
|
|
let mut result = AssignedIds::default();
|
|
let has_changes = !batch.changes.is_empty();
|
|
|
|
if has_changes {
|
|
for &account_id in batch.changes.keys() {
|
|
let key = ValueClass::ChangeId.serialize(account_id, 0, 0, 0);
|
|
let change_id = trx
|
|
.prepare_cached(concat!(
|
|
"INSERT INTO n (k, v) VALUES (?, ?) ",
|
|
"ON CONFLICT(k) DO UPDATE SET v = v + ",
|
|
"excluded.v RETURNING v"
|
|
))
|
|
.map_err(into_error)
|
|
.caused_by(trc::location!())?
|
|
.query_row(params![&key, &1i64], |row| row.get::<_, i64>(0))
|
|
.map_err(into_error)
|
|
.caused_by(trc::location!())?;
|
|
result.push_change_id(account_id, change_id as u64);
|
|
}
|
|
}
|
|
|
|
for op in batch.ops.iter_mut() {
|
|
match op {
|
|
Operation::AccountId {
|
|
account_id: account_id_,
|
|
} => {
|
|
account_id = *account_id_;
|
|
if has_changes {
|
|
change_id = result.set_current_change_id(account_id)?;
|
|
}
|
|
}
|
|
Operation::Collection {
|
|
collection: collection_,
|
|
} => {
|
|
collection = u8::from(*collection_);
|
|
}
|
|
Operation::DocumentId {
|
|
document_id: document_id_,
|
|
} => {
|
|
document_id = *document_id_;
|
|
}
|
|
Operation::Value { class, op } => {
|
|
let key = class.serialize(account_id, collection, document_id, 0);
|
|
let table = char::from(class.subspace(collection));
|
|
|
|
match op {
|
|
ValueOp::Set(value) => {
|
|
trx.prepare_cached(&format!(
|
|
"INSERT OR REPLACE INTO {} (k, v) VALUES (?, ?)",
|
|
table
|
|
))
|
|
.map_err(into_error)
|
|
.caused_by(trc::location!())?
|
|
.execute([&key, value])
|
|
.map_err(into_error)
|
|
.caused_by(trc::location!())?;
|
|
}
|
|
ValueOp::SetFnc(set_op) => {
|
|
let value = (set_op.fnc)(&set_op.params, &result)?;
|
|
trx.prepare_cached(&format!(
|
|
"INSERT OR REPLACE INTO {} (k, v) VALUES (?, ?)",
|
|
table
|
|
))
|
|
.map_err(into_error)
|
|
.caused_by(trc::location!())?
|
|
.execute([&key, &value])
|
|
.map_err(into_error)
|
|
.caused_by(trc::location!())?;
|
|
}
|
|
ValueOp::MergeFnc(merge_op) => {
|
|
let merge_result = trx
|
|
.prepare_cached(&format!("SELECT v FROM {} WHERE k = ?", table))
|
|
.map_err(into_error)
|
|
.caused_by(trc::location!())?
|
|
.query_row([&key], |row| {
|
|
Ok((merge_op.fnc)(
|
|
&merge_op.params,
|
|
&result,
|
|
Some(row.get_ref(0)?.as_bytes()?),
|
|
))
|
|
})
|
|
.optional()
|
|
.map_err(into_error)
|
|
.caused_by(trc::location!())?
|
|
.unwrap_or_else(|| {
|
|
(merge_op.fnc)(&merge_op.params, &result, None)
|
|
})?;
|
|
|
|
match merge_result {
|
|
MergeResult::Update(value) => {
|
|
trx.prepare_cached(&format!(
|
|
"INSERT OR REPLACE INTO {} (k, v) VALUES (?, ?)",
|
|
table
|
|
))
|
|
.map_err(into_error)
|
|
.caused_by(trc::location!())?
|
|
.execute([&key, &value])
|
|
.map_err(into_error)
|
|
.caused_by(trc::location!())?;
|
|
}
|
|
MergeResult::Delete => {
|
|
trx.prepare_cached(&format!(
|
|
"DELETE FROM {} WHERE k = ?",
|
|
table
|
|
))
|
|
.map_err(into_error)
|
|
.caused_by(trc::location!())?
|
|
.execute([&key])
|
|
.map_err(into_error)
|
|
.caused_by(trc::location!())?;
|
|
}
|
|
MergeResult::Skip => (),
|
|
}
|
|
}
|
|
ValueOp::AtomicAdd(by) => {
|
|
if *by >= 0 {
|
|
trx.prepare_cached(&format!(
|
|
concat!(
|
|
"INSERT INTO {} (k, v) VALUES (?, ?) ",
|
|
"ON CONFLICT(k) DO UPDATE SET v = v + excluded.v"
|
|
),
|
|
table
|
|
))
|
|
.map_err(into_error)
|
|
.caused_by(trc::location!())?
|
|
.execute(params![&key, *by])
|
|
.map_err(into_error)
|
|
.caused_by(trc::location!())?;
|
|
} else {
|
|
trx.prepare_cached(&format!(
|
|
"UPDATE {table} SET v = v + ? WHERE k = ?"
|
|
))
|
|
.map_err(into_error)
|
|
.caused_by(trc::location!())?
|
|
.execute(params![*by, &key])
|
|
.map_err(into_error)
|
|
.caused_by(trc::location!())?;
|
|
}
|
|
}
|
|
ValueOp::AddAndGet(by) => {
|
|
result.push_counter_id(
|
|
trx.prepare_cached(&format!(
|
|
concat!(
|
|
"INSERT INTO {} (k, v) VALUES (?, ?) ",
|
|
"ON CONFLICT(k) DO UPDATE SET v = v + ",
|
|
"excluded.v RETURNING v"
|
|
),
|
|
table
|
|
))
|
|
.map_err(into_error)
|
|
.caused_by(trc::location!())?
|
|
.query_row(params![&key, &*by], |row| row.get::<_, i64>(0))
|
|
.map_err(into_error)
|
|
.caused_by(trc::location!())?,
|
|
);
|
|
}
|
|
ValueOp::Clear => {
|
|
trx.prepare_cached(&format!("DELETE FROM {} WHERE k = ?", table))
|
|
.map_err(into_error)
|
|
.caused_by(trc::location!())?
|
|
.execute([&key])
|
|
.map_err(into_error)
|
|
.caused_by(trc::location!())?;
|
|
}
|
|
}
|
|
}
|
|
Operation::Index { field, key, set } => {
|
|
let key = IndexKey {
|
|
account_id,
|
|
collection,
|
|
document_id,
|
|
field: *field,
|
|
key: &*key,
|
|
}
|
|
.serialize(0);
|
|
|
|
if *set {
|
|
trx.prepare_cached("INSERT OR IGNORE INTO i (k) VALUES (?)")
|
|
.map_err(into_error)
|
|
.caused_by(trc::location!())?
|
|
.execute([&key])
|
|
.map_err(into_error)
|
|
.caused_by(trc::location!())?;
|
|
} else {
|
|
trx.prepare_cached("DELETE FROM i WHERE k = ?")
|
|
.map_err(into_error)
|
|
.caused_by(trc::location!())?
|
|
.execute([&key])
|
|
.map_err(into_error)
|
|
.caused_by(trc::location!())?;
|
|
}
|
|
}
|
|
Operation::Log { collection, set } => {
|
|
let key = LogKey {
|
|
account_id,
|
|
collection: u8::from(*collection),
|
|
change_id,
|
|
}
|
|
.serialize(0);
|
|
|
|
trx.prepare_cached("INSERT OR REPLACE INTO l (k, v) VALUES (?, ?)")
|
|
.map_err(into_error)
|
|
.caused_by(trc::location!())?
|
|
.execute([&key, set])
|
|
.map_err(into_error)
|
|
.caused_by(trc::location!())?;
|
|
}
|
|
Operation::AssertValue {
|
|
class,
|
|
assert_value,
|
|
} => {
|
|
let key = class.serialize(account_id, collection, document_id, 0);
|
|
let table = char::from(class.subspace(collection));
|
|
|
|
let matches = trx
|
|
.prepare_cached(&format!("SELECT v FROM {} WHERE k = ?", table))
|
|
.map_err(into_error)
|
|
.caused_by(trc::location!())?
|
|
.query_row([&key], |row| {
|
|
Ok(assert_value.matches(row.get_ref(0)?.as_bytes()?))
|
|
})
|
|
.optional()
|
|
.map_err(into_error)
|
|
.caused_by(trc::location!())?
|
|
.unwrap_or_else(|| assert_value.is_none());
|
|
if !matches {
|
|
trx.rollback()
|
|
.map_err(into_error)
|
|
.caused_by(trc::location!())?;
|
|
return Err(trc::StoreEvent::AssertValueFailed
|
|
.into_err()
|
|
.caused_by(trc::location!()));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
trx.commit().map(|_| result).map_err(into_error)
|
|
})
|
|
.await
|
|
}
|
|
|
|
pub(crate) async fn purge_store(&self) -> trc::Result<()> {
|
|
let manager = self.conn_pool.clone();
|
|
self.spawn_worker(move || {
|
|
let conn = manager.get().map_err(into_error)?;
|
|
for subspace in [SUBSPACE_QUOTA, SUBSPACE_COUNTER, SUBSPACE_IN_MEMORY_COUNTER] {
|
|
conn.prepare_cached(&format!("DELETE FROM {} WHERE v = 0", char::from(subspace),))
|
|
.map_err(into_error)
|
|
.caused_by(trc::location!())?
|
|
.execute([])
|
|
.map_err(into_error)
|
|
.caused_by(trc::location!())?;
|
|
}
|
|
|
|
Ok(())
|
|
})
|
|
.await
|
|
}
|
|
|
|
pub(crate) async fn delete_range(&self, from: impl Key, to: impl Key) -> trc::Result<()> {
|
|
let manager = self.conn_pool.clone();
|
|
self.spawn_worker(move || {
|
|
let conn = manager.get().map_err(into_error)?;
|
|
|
|
conn.prepare_cached(&format!(
|
|
"DELETE FROM {} WHERE k >= ? AND k < ?",
|
|
char::from(from.subspace()),
|
|
))
|
|
.map_err(into_error)
|
|
.caused_by(trc::location!())?
|
|
.execute([from.serialize(0), to.serialize(0)])
|
|
.map_err(into_error)
|
|
.caused_by(trc::location!())?;
|
|
|
|
Ok(())
|
|
})
|
|
.await
|
|
}
|
|
}
|