mirror of
https://github.com/toeverything/AFFiNE.git
synced 2026-07-14 08:36:22 +08:00
forehalo
484 lines
12 KiB
Rust
484 lines
12 KiB
Rust
use std::ops::Deref;
|
|
|
|
use chrono::{DateTime, NaiveDateTime};
|
|
use sqlx::{QueryBuilder, Row};
|
|
|
|
use super::storage::{Result, SqliteDocStorage};
|
|
use super::{DocClock, DocRecord, DocUpdate};
|
|
|
|
struct Meta {
|
|
space_id: String,
|
|
}
|
|
|
|
impl SqliteDocStorage {
|
|
pub async fn set_space_id(&self, space_id: String) -> Result<()> {
|
|
// ensure only one record exists in table
|
|
let result = sqlx::query_as!(Meta, "SELECT * FROM meta;")
|
|
.fetch_optional(&self.pool)
|
|
.await?;
|
|
|
|
match result {
|
|
Some(meta) => {
|
|
if meta.space_id != space_id {
|
|
sqlx::query("UPDATE meta SET space_id = $1;")
|
|
.bind(&space_id)
|
|
.execute(&self.pool)
|
|
.await?;
|
|
|
|
sqlx::query("UPDATE updates SET doc_id = $1 WHERE doc_id = $2;")
|
|
.bind(&space_id)
|
|
.bind(&meta.space_id)
|
|
.execute(&self.pool)
|
|
.await?;
|
|
|
|
sqlx::query("UPDATE snapshots SET doc_id = $1 WHERE doc_id = $2;")
|
|
.bind(&space_id)
|
|
.bind(&meta.space_id)
|
|
.execute(&self.pool)
|
|
.await?;
|
|
|
|
sqlx::query("UPDATE clocks SET doc_id = $1 WHERE doc_id = $2;")
|
|
.bind(&space_id)
|
|
.bind(&meta.space_id)
|
|
.execute(&self.pool)
|
|
.await?;
|
|
|
|
sqlx::query("UPDATE peer_clocks SET doc_id = $1 WHERE doc_id = $2;")
|
|
.bind(&space_id)
|
|
.bind(&meta.space_id)
|
|
.execute(&self.pool)
|
|
.await?;
|
|
}
|
|
}
|
|
None => {
|
|
sqlx::query("INSERT INTO meta (space_id) VALUES ($1);")
|
|
.bind(&space_id)
|
|
.execute(&self.pool)
|
|
.await?;
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
pub async fn push_update<Update: AsRef<[u8]>>(
|
|
&self,
|
|
doc_id: String,
|
|
update: Update,
|
|
) -> Result<NaiveDateTime> {
|
|
let mut timestamp = DateTime::from_timestamp_millis(chrono::Utc::now().timestamp_millis())
|
|
.unwrap()
|
|
.naive_utc();
|
|
|
|
let mut tried = 0;
|
|
|
|
// Keep trying with incremented timestamps until success
|
|
loop {
|
|
match self
|
|
.try_insert_update_with_timestamp(&doc_id, update.as_ref(), timestamp)
|
|
.await
|
|
{
|
|
Ok(()) => break,
|
|
Err(e) => {
|
|
if tried > 10 {
|
|
return Err(e);
|
|
}
|
|
|
|
// Increment timestamp by 1ms and retry
|
|
timestamp = timestamp + chrono::Duration::milliseconds(1);
|
|
tried += 1;
|
|
}
|
|
}
|
|
}
|
|
|
|
Ok(timestamp)
|
|
}
|
|
|
|
async fn try_insert_update_with_timestamp(
|
|
&self,
|
|
doc_id: &str,
|
|
update: &[u8],
|
|
timestamp: NaiveDateTime,
|
|
) -> sqlx::Result<()> {
|
|
let mut tx = self.pool.begin().await?;
|
|
|
|
sqlx::query(r#"INSERT INTO updates (doc_id, data, created_at) VALUES ($1, $2, $3);"#)
|
|
.bind(&doc_id)
|
|
.bind(update.as_ref())
|
|
.bind(timestamp)
|
|
.execute(&mut *tx)
|
|
.await?;
|
|
|
|
sqlx::query(
|
|
r#"
|
|
INSERT INTO clocks (doc_id, timestamp) VALUES ($1, $2)
|
|
ON CONFLICT(doc_id)
|
|
DO UPDATE SET timestamp=$2;"#,
|
|
)
|
|
.bind(&doc_id)
|
|
.bind(timestamp)
|
|
.execute(&mut *tx)
|
|
.await?;
|
|
|
|
tx.commit().await?;
|
|
|
|
Ok(())
|
|
}
|
|
|
|
pub async fn get_doc_snapshot(&self, doc_id: String) -> Result<Option<DocRecord>> {
|
|
sqlx::query_as!(
|
|
DocRecord,
|
|
"SELECT doc_id, data, updated_at as timestamp FROM snapshots WHERE doc_id = ?",
|
|
doc_id
|
|
)
|
|
.fetch_optional(&self.pool)
|
|
.await
|
|
}
|
|
|
|
pub async fn set_doc_snapshot(&self, snapshot: DocRecord) -> Result<bool> {
|
|
let result = sqlx::query(
|
|
r#"
|
|
INSERT INTO snapshots (doc_id, data, updated_at)
|
|
VALUES ($1, $2, $3)
|
|
ON CONFLICT(doc_id)
|
|
DO UPDATE SET data=$2, updated_at=$3
|
|
WHERE updated_at <= $3;"#,
|
|
)
|
|
.bind(snapshot.doc_id)
|
|
.bind(snapshot.data.deref())
|
|
.bind(snapshot.timestamp)
|
|
.execute(&self.pool)
|
|
.await?;
|
|
|
|
Ok(result.rows_affected() == 1)
|
|
}
|
|
|
|
pub async fn get_doc_updates(&self, doc_id: String) -> Result<Vec<DocUpdate>> {
|
|
sqlx::query_as!(
|
|
DocUpdate,
|
|
"SELECT doc_id, created_at, data FROM updates WHERE doc_id = ?",
|
|
doc_id
|
|
)
|
|
.fetch_all(&self.pool)
|
|
.await
|
|
}
|
|
|
|
pub async fn mark_updates_merged(
|
|
&self,
|
|
doc_id: String,
|
|
updates: Vec<NaiveDateTime>,
|
|
) -> Result<u32> {
|
|
let mut qb = QueryBuilder::new("DELETE FROM updates");
|
|
|
|
qb.push(" WHERE doc_id = ");
|
|
qb.push_bind(doc_id);
|
|
qb.push(" AND created_at IN (");
|
|
let mut separated = qb.separated(", ");
|
|
updates.iter().for_each(|update| {
|
|
separated.push_bind(update);
|
|
});
|
|
qb.push(");");
|
|
|
|
let query = qb.build();
|
|
|
|
let result = query.execute(&self.pool).await?;
|
|
|
|
Ok(result.rows_affected() as u32)
|
|
}
|
|
|
|
pub async fn delete_doc(&self, doc_id: String) -> Result<()> {
|
|
let mut tx = self.pool.begin().await?;
|
|
|
|
sqlx::query("DELETE FROM updates WHERE doc_id = ?;")
|
|
.bind(&doc_id)
|
|
.execute(&mut *tx)
|
|
.await?;
|
|
|
|
sqlx::query("DELETE FROM snapshots WHERE doc_id = ?;")
|
|
.bind(&doc_id)
|
|
.execute(&mut *tx)
|
|
.await?;
|
|
|
|
sqlx::query("DELETE FROM clocks WHERE doc_id = ?;")
|
|
.bind(&doc_id)
|
|
.execute(&mut *tx)
|
|
.await?;
|
|
|
|
tx.commit().await
|
|
}
|
|
|
|
pub async fn get_doc_clocks(&self, after: Option<NaiveDateTime>) -> Result<Vec<DocClock>> {
|
|
let query = if let Some(after) = after {
|
|
sqlx::query("SELECT doc_id, timestamp FROM clocks WHERE timestamp > $1").bind(after)
|
|
} else {
|
|
sqlx::query("SELECT doc_id, timestamp FROM clocks")
|
|
};
|
|
|
|
let clocks = query.fetch_all(&self.pool).await?;
|
|
|
|
Ok(
|
|
clocks
|
|
.iter()
|
|
.map(|row| DocClock {
|
|
doc_id: row.get("doc_id"),
|
|
timestamp: row.get("timestamp"),
|
|
})
|
|
.collect(),
|
|
)
|
|
}
|
|
|
|
pub async fn get_doc_clock(&self, doc_id: String) -> Result<Option<DocClock>> {
|
|
sqlx::query_as!(
|
|
DocClock,
|
|
"SELECT doc_id, timestamp FROM clocks WHERE doc_id = ?",
|
|
doc_id
|
|
)
|
|
.fetch_optional(&self.pool)
|
|
.await
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use chrono::{DateTime, Utc};
|
|
|
|
use super::*;
|
|
|
|
async fn get_storage() -> SqliteDocStorage {
|
|
let storage = SqliteDocStorage::new(":memory:".to_string());
|
|
storage.connect().await.unwrap();
|
|
|
|
storage
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn set_space_id() {
|
|
let storage = get_storage().await;
|
|
|
|
storage.set_space_id("test".to_string()).await.unwrap();
|
|
|
|
let result = sqlx::query!("SELECT space_id FROM meta;")
|
|
.fetch_one(&storage.pool)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(result.space_id, "test");
|
|
|
|
storage.set_space_id("test2".to_string()).await.unwrap();
|
|
|
|
let result = sqlx::query!("SELECT space_id FROM meta;")
|
|
.fetch_one(&storage.pool)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(result.space_id, "test2");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn set_space_id_with_existing_doc() {
|
|
let storage = get_storage().await;
|
|
|
|
storage.set_space_id("test".to_string()).await.unwrap();
|
|
storage
|
|
.push_update("test".to_string(), vec![0, 0])
|
|
.await
|
|
.unwrap();
|
|
storage
|
|
.set_doc_snapshot(DocRecord {
|
|
doc_id: "test".to_string(),
|
|
data: vec![0, 0].into(),
|
|
timestamp: Utc::now().naive_utc(),
|
|
})
|
|
.await
|
|
.unwrap();
|
|
|
|
storage
|
|
.set_peer_pulled_remote_clock(
|
|
"remote".to_string(),
|
|
"test".to_string(),
|
|
Utc::now().naive_utc(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
storage.set_space_id("new_id".to_string()).await.unwrap();
|
|
|
|
let result = sqlx::query!("SELECT space_id FROM meta;")
|
|
.fetch_one(&storage.pool)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(result.space_id, "new_id");
|
|
|
|
let clocks = storage.get_doc_clocks(None).await.unwrap();
|
|
|
|
assert_eq!(clocks[0].doc_id, "new_id");
|
|
|
|
let clocks = storage
|
|
.get_peer_pulled_remote_clock("remote".to_string(), "new_id".to_string())
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(clocks.doc_id, "new_id");
|
|
|
|
let updates = storage.get_doc_updates("new_id".to_string()).await.unwrap();
|
|
|
|
assert_eq!(updates.len(), 1);
|
|
|
|
let snapshot = storage
|
|
.get_doc_snapshot("new_id".to_string())
|
|
.await
|
|
.unwrap();
|
|
|
|
assert!(snapshot.is_some());
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn push_updates() {
|
|
let storage = get_storage().await;
|
|
|
|
let updates = vec![vec![0, 0], vec![0, 1], vec![1, 0], vec![1, 1]];
|
|
|
|
for update in updates.iter() {
|
|
storage
|
|
.push_update("test".to_string(), update)
|
|
.await
|
|
.unwrap();
|
|
}
|
|
|
|
let result = storage.get_doc_updates("test".to_string()).await.unwrap();
|
|
|
|
assert_eq!(result.len(), 4);
|
|
assert_eq!(
|
|
result.iter().map(|u| u.data.as_ref()).collect::<Vec<_>>(),
|
|
updates
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn get_doc_snapshot() {
|
|
let storage = get_storage().await;
|
|
|
|
let none = storage.get_doc_snapshot("test".to_string()).await.unwrap();
|
|
|
|
assert!(none.is_none());
|
|
|
|
let snapshot = DocRecord {
|
|
doc_id: "test".to_string(),
|
|
data: vec![0, 0].into(),
|
|
timestamp: Utc::now().naive_utc(),
|
|
};
|
|
|
|
storage.set_doc_snapshot(snapshot).await.unwrap();
|
|
|
|
let result = storage.get_doc_snapshot("test".to_string()).await.unwrap();
|
|
|
|
assert!(result.is_some());
|
|
assert_eq!(result.unwrap().data.as_ref(), vec![0, 0]);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn set_doc_snapshot() {
|
|
let storage = get_storage().await;
|
|
|
|
let snapshot = DocRecord {
|
|
doc_id: "test".to_string(),
|
|
data: vec![0, 0].into(),
|
|
timestamp: Utc::now().naive_utc(),
|
|
};
|
|
|
|
storage.set_doc_snapshot(snapshot).await.unwrap();
|
|
|
|
let result = storage.get_doc_snapshot("test".to_string()).await.unwrap();
|
|
|
|
assert!(result.is_some());
|
|
assert_eq!(result.unwrap().data.as_ref(), vec![0, 0]);
|
|
|
|
let snapshot = DocRecord {
|
|
doc_id: "test".to_string(),
|
|
data: vec![0, 1].into(),
|
|
timestamp: DateTime::from_timestamp_millis(Utc::now().timestamp_millis() - 1000)
|
|
.unwrap()
|
|
.naive_utc(),
|
|
};
|
|
|
|
// can't update because it's tempstamp is older
|
|
storage.set_doc_snapshot(snapshot).await.unwrap();
|
|
|
|
let result = storage.get_doc_snapshot("test".to_string()).await.unwrap();
|
|
|
|
assert!(result.is_some());
|
|
assert_eq!(result.unwrap().data.as_ref(), vec![0, 0]);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn get_doc_clocks() {
|
|
let storage = get_storage().await;
|
|
|
|
let clocks = storage.get_doc_clocks(None).await.unwrap();
|
|
|
|
assert_eq!(clocks.len(), 0);
|
|
|
|
for i in 1..5u32 {
|
|
storage
|
|
.push_update(format!("test_{i}"), vec![0, 0])
|
|
.await
|
|
.unwrap();
|
|
}
|
|
|
|
let clocks = storage.get_doc_clocks(None).await.unwrap();
|
|
|
|
assert_eq!(clocks.len(), 4);
|
|
assert_eq!(
|
|
clocks.iter().map(|c| c.doc_id.as_str()).collect::<Vec<_>>(),
|
|
vec!["test_1", "test_2", "test_3", "test_4"]
|
|
);
|
|
|
|
let clocks = storage
|
|
.get_doc_clocks(Some(Utc::now().naive_utc()))
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(clocks.len(), 0);
|
|
|
|
let clock = storage.get_doc_clock("test_1".to_string()).await.unwrap();
|
|
|
|
assert!(clock.is_some());
|
|
assert_eq!(clock.unwrap().doc_id, "test_1");
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn mark_updates_merged() {
|
|
let storage = get_storage().await;
|
|
|
|
let updates = [vec![0, 0], vec![0, 1], vec![1, 0], vec![1, 1]];
|
|
|
|
for update in updates.iter() {
|
|
storage
|
|
.push_update("test".to_string(), update)
|
|
.await
|
|
.unwrap();
|
|
}
|
|
|
|
let updates = storage.get_doc_updates("test".to_string()).await.unwrap();
|
|
|
|
let result = storage
|
|
.mark_updates_merged(
|
|
"test".to_string(),
|
|
updates
|
|
.iter()
|
|
.skip(1)
|
|
.map(|u| u.created_at)
|
|
.collect::<Vec<_>>(),
|
|
)
|
|
.await
|
|
.unwrap();
|
|
|
|
assert_eq!(result, 3);
|
|
|
|
let updates = storage.get_doc_updates("test".to_string()).await.unwrap();
|
|
|
|
assert_eq!(updates.len(), 1);
|
|
}
|
|
}
|