AFFiNE-Mirror/packages/frontend/native/nbstore/src/lib.rs

pub mod blob;
pub mod blob_sync;
pub mod doc;
pub mod doc_sync;
pub mod error;
pub mod indexer;
pub mod indexer_sync;
pub mod pool;
pub mod storage;

#[cfg(not(feature = "use-as-lib"))]
use affine_common::napi_utils::to_napi_error;
use chrono::NaiveDateTime;
use napi::bindgen_prelude::*;
use napi_derive::napi;
use pool::{Ref, SqliteDocStoragePool};
use storage::SqliteDocStorage;

#[cfg(feature = "use-as-lib")]
type Result<T> = anyhow::Result<T>;

#[cfg(not(feature = "use-as-lib"))]
type Result<T> = napi::Result<T>;

#[cfg(not(feature = "use-as-lib"))]
impl From<error::Error> for napi::Error {
  fn from(err: error::Error) -> Self {
    to_napi_error(err, napi::Status::GenericFailure)
  }
}

#[cfg(feature = "use-as-lib")]
pub type Data = Vec<u8>;

#[cfg(not(feature = "use-as-lib"))]
pub type Data = Uint8Array;

#[napi(object)]
pub struct DocUpdate {
  pub doc_id: String,
  pub timestamp: NaiveDateTime,
  #[napi(ts_type = "Uint8Array")]
  pub bin: Data,
}

#[napi(object)]
pub struct DocRecord {
  pub doc_id: String,
  #[napi(ts_type = "Uint8Array")]
  pub bin: Data,
  pub timestamp: NaiveDateTime,
}

#[derive(Debug)]
#[napi(object)]
pub struct DocClock {
  pub doc_id: String,
  pub timestamp: NaiveDateTime,
}

#[derive(Debug)]
#[napi(object)]
pub struct DocIndexedClock {
  pub doc_id: String,
  pub timestamp: NaiveDateTime,
  pub indexer_version: i64,
}

#[napi(object)]
pub struct SetBlob {
  pub key: String,
  #[napi(ts_type = "Uint8Array")]
  pub data: Data,
  pub mime: String,
}

#[napi(object)]
pub struct Blob {
  pub key: String,
  #[napi(ts_type = "Uint8Array")]
  pub data: Data,
  pub mime: String,
  pub size: i64,
  pub created_at: NaiveDateTime,
}

#[napi(object)]
pub struct ListedBlob {
  pub key: String,
  pub size: i64,
  pub mime: String,
  pub created_at: NaiveDateTime,
}

#[napi]
pub struct DocStoragePool {
  pool: SqliteDocStoragePool,
}

#[napi]
impl DocStoragePool {
  #[napi(constructor)]
  pub fn new() -> Result<Self> {
    Ok(Self {
      pool: SqliteDocStoragePool::default(),
    })
  }

  async fn get(&self, universal_id: String) -> Result<Ref<SqliteDocStorage>> {
    Ok(self.pool.get(universal_id).await?)
  }

  #[napi]
  /// Initialize the database and run migrations.
  pub async fn connect(&self, universal_id: String, path: String) -> Result<()> {
    self.pool.connect(universal_id, path).await?;
    Ok(())
  }

  #[napi]
  pub async fn disconnect(&self, universal_id: String) -> Result<()> {
    self.pool.disconnect(universal_id).await?;
    Ok(())
  }

  #[napi]
  pub async fn checkpoint(&self, universal_id: String) -> Result<()> {
    self.pool.get(universal_id).await?.checkpoint().await?;
    Ok(())
  }

  #[napi]
  pub async fn vacuum_into(&self, universal_id: String, path: String) -> Result<()> {
    self.pool.get(universal_id).await?.vacuum_into(path).await?;
    Ok(())
  }

  #[napi]
  pub async fn crawl_doc_data(&self, universal_id: String, doc_id: String) -> Result<indexer::NativeCrawlResult> {
    let result = self.get(universal_id).await?.crawl_doc_data(&doc_id).await?;
    Ok(result)
  }

  #[napi]
  pub async fn set_space_id(&self, universal_id: String, space_id: String) -> Result<()> {
    self.get(universal_id).await?.set_space_id(space_id).await?;
    Ok(())
  }

  #[napi]
  pub async fn push_update(&self, universal_id: String, doc_id: String, update: Uint8Array) -> Result<NaiveDateTime> {
    Ok(self.get(universal_id).await?.push_update(doc_id, update).await?)
  }

  #[napi]
  pub async fn get_doc_snapshot(&self, universal_id: String, doc_id: String) -> Result<Option<DocRecord>> {
    Ok(self.get(universal_id).await?.get_doc_snapshot(doc_id).await?)
  }

  #[napi]
  pub async fn set_doc_snapshot(&self, universal_id: String, snapshot: DocRecord) -> Result<bool> {
    Ok(self.get(universal_id).await?.set_doc_snapshot(snapshot).await?)
  }

  #[napi]
  pub async fn get_doc_updates(&self, universal_id: String, doc_id: String) -> Result<Vec<DocUpdate>> {
    Ok(self.get(universal_id).await?.get_doc_updates(doc_id).await?)
  }

  #[napi]
  pub async fn mark_updates_merged(
    &self,
    universal_id: String,
    doc_id: String,
    updates: Vec<NaiveDateTime>,
  ) -> Result<u32> {
    Ok(
      self
        .get(universal_id)
        .await?
        .mark_updates_merged(doc_id, updates)
        .await?,
    )
  }

  #[napi]
  pub async fn delete_doc(&self, universal_id: String, doc_id: String) -> Result<()> {
    self.get(universal_id).await?.delete_doc(doc_id).await?;
    Ok(())
  }

  #[napi]
  pub async fn get_doc_clocks(&self, universal_id: String, after: Option<NaiveDateTime>) -> Result<Vec<DocClock>> {
    Ok(self.get(universal_id).await?.get_doc_clocks(after).await?)
  }

  #[napi]
  pub async fn get_doc_clock(&self, universal_id: String, doc_id: String) -> Result<Option<DocClock>> {
    Ok(self.get(universal_id).await?.get_doc_clock(doc_id).await?)
  }

  #[napi]
  pub async fn get_doc_indexed_clock(&self, universal_id: String, doc_id: String) -> Result<Option<DocIndexedClock>> {
    Ok(self.get(universal_id).await?.get_doc_indexed_clock(doc_id).await?)
  }

  #[napi]
  pub async fn set_doc_indexed_clock(
    &self,
    universal_id: String,
    doc_id: String,
    indexed_clock: NaiveDateTime,
    indexer_version: i64,
  ) -> Result<()> {
    self
      .get(universal_id)
      .await?
      .set_doc_indexed_clock(doc_id, indexed_clock, indexer_version)
      .await?;
    Ok(())
  }

  #[napi]
  pub async fn clear_doc_indexed_clock(&self, universal_id: String, doc_id: String) -> Result<()> {
    self.get(universal_id).await?.clear_doc_indexed_clock(doc_id).await?;
    Ok(())
  }

  #[napi(async_runtime)]
  pub async fn get_blob(&self, universal_id: String, key: String) -> Result<Option<Blob>> {
    Ok(self.get(universal_id).await?.get_blob(key).await?)
  }

  #[napi]
  pub async fn set_blob(&self, universal_id: String, blob: SetBlob) -> Result<()> {
    self.get(universal_id).await?.set_blob(blob).await?;
    Ok(())
  }

  #[napi]
  pub async fn delete_blob(&self, universal_id: String, key: String, permanently: bool) -> Result<()> {
    self.get(universal_id).await?.delete_blob(key, permanently).await?;
    Ok(())
  }

  #[napi]
  pub async fn release_blobs(&self, universal_id: String) -> Result<()> {
    self.get(universal_id).await?.release_blobs().await?;
    Ok(())
  }

  #[napi]
  pub async fn list_blobs(&self, universal_id: String) -> Result<Vec<ListedBlob>> {
    Ok(self.get(universal_id).await?.list_blobs().await?)
  }

  #[napi]
  pub async fn get_peer_remote_clocks(&self, universal_id: String, peer: String) -> Result<Vec<DocClock>> {
    Ok(self.get(universal_id).await?.get_peer_remote_clocks(peer).await?)
  }

  #[napi]
  pub async fn get_peer_remote_clock(
    &self,
    universal_id: String,
    peer: String,
    doc_id: String,
  ) -> Result<Option<DocClock>> {
    Ok(
      self
        .get(universal_id)
        .await?
        .get_peer_remote_clock(peer, doc_id)
        .await?,
    )
  }

  #[napi]
  pub async fn set_peer_remote_clock(
    &self,
    universal_id: String,
    peer: String,
    doc_id: String,
    clock: NaiveDateTime,
  ) -> Result<()> {
    self
      .get(universal_id)
      .await?
      .set_peer_remote_clock(peer, doc_id, clock)
      .await?;
    Ok(())
  }

  #[napi]
  pub async fn get_peer_pulled_remote_clocks(&self, universal_id: String, peer: String) -> Result<Vec<DocClock>> {
    Ok(
      self
        .get(universal_id)
        .await?
        .get_peer_pulled_remote_clocks(peer)
        .await?,
    )
  }

  #[napi]
  pub async fn get_peer_pulled_remote_clock(
    &self,
    universal_id: String,
    peer: String,
    doc_id: String,
  ) -> Result<Option<DocClock>> {
    Ok(
      self
        .get(universal_id)
        .await?
        .get_peer_pulled_remote_clock(peer, doc_id)
        .await?,
    )
  }

  #[napi]
  pub async fn set_peer_pulled_remote_clock(
    &self,
    universal_id: String,
    peer: String,
    doc_id: String,
    clock: NaiveDateTime,
  ) -> Result<()> {
    self
      .get(universal_id)
      .await?
      .set_peer_pulled_remote_clock(peer, doc_id, clock)
      .await?;
    Ok(())
  }

  #[napi]
  pub async fn get_peer_pushed_clocks(&self, universal_id: String, peer: String) -> Result<Vec<DocClock>> {
    Ok(self.get(universal_id).await?.get_peer_pushed_clocks(peer).await?)
  }

  #[napi]
  pub async fn get_peer_pushed_clock(
    &self,
    universal_id: String,
    peer: String,
    doc_id: String,
  ) -> Result<Option<DocClock>> {
    Ok(
      self
        .get(universal_id)
        .await?
        .get_peer_pushed_clock(peer, doc_id)
        .await?,
    )
  }

  #[napi]
  pub async fn set_peer_pushed_clock(
    &self,
    universal_id: String,
    peer: String,
    doc_id: String,
    clock: NaiveDateTime,
  ) -> Result<()> {
    self
      .get(universal_id)
      .await?
      .set_peer_pushed_clock(peer, doc_id, clock)
      .await?;
    Ok(())
  }

  #[napi]
  pub async fn clear_clocks(&self, universal_id: String) -> Result<()> {
    self.get(universal_id).await?.clear_clocks().await?;
    Ok(())
  }

  #[napi]
  pub async fn set_blob_uploaded_at(
    &self,
    universal_id: String,
    peer: String,
    blob_id: String,
    uploaded_at: Option<NaiveDateTime>,
  ) -> Result<()> {
    self
      .get(universal_id)
      .await?
      .set_blob_uploaded_at(peer, blob_id, uploaded_at)
      .await?;
    Ok(())
  }

  #[napi]
  pub async fn get_blob_uploaded_at(
    &self,
    universal_id: String,
    peer: String,
    blob_id: String,
  ) -> Result<Option<NaiveDateTime>> {
    let result = self
      .get(universal_id)
      .await?
      .get_blob_uploaded_at(peer, blob_id)
      .await?;

    Ok(result)
  }

  #[napi]
  pub async fn fts_add_document(
    &self,
    id: String,
    index_name: String,
    doc_id: String,
    text: String,
    index: bool,
  ) -> Result<()> {
    let storage = self.pool.get(id).await?;
    storage.fts_add(&index_name, &doc_id, &text, index).await?;
    Ok(())
  }

  #[napi]
  pub async fn fts_flush_index(&self, id: String) -> Result<()> {
    let storage = self.pool.get(id).await?;
    storage.flush_index().await?;
    Ok(())
  }

  #[napi]
  pub async fn fts_index_version(&self) -> Result<u32> {
    Ok(SqliteDocStorage::index_version())
  }

  #[napi]
  pub async fn fts_delete_document(&self, id: String, index_name: String, doc_id: String) -> Result<()> {
    let storage = self.pool.get(id).await?;
    storage.fts_delete(&index_name, &doc_id).await?;
    Ok(())
  }

  #[napi]
  pub async fn fts_get_document(&self, id: String, index_name: String, doc_id: String) -> Result<Option<String>> {
    let storage = self.pool.get(id).await?;
    Ok(storage.fts_get(&index_name, &doc_id).await?)
  }

  #[napi]
  pub async fn fts_search(
    &self,
    id: String,
    index_name: String,
    query: String,
  ) -> Result<Vec<indexer::NativeSearchHit>> {
    let storage = self.pool.get(id).await?;
    Ok(storage.fts_search(&index_name, &query).await?)
  }

  #[napi]
  pub async fn fts_get_matches(
    &self,
    id: String,
    index_name: String,
    doc_id: String,
    query: String,
  ) -> Result<Vec<indexer::NativeMatch>> {
    let storage = self.pool.get(id).await?;
    Ok(storage.fts_get_matches(&index_name, &doc_id, &query).await?)
  }
}

#[napi]
pub struct DocStorage {
  storage: SqliteDocStorage,
}

#[napi]
impl DocStorage {
  #[napi(constructor, async_runtime)]
  pub fn new(path: String) -> Self {
    Self {
      storage: SqliteDocStorage::new(path),
    }
  }

  #[napi]
  pub async fn validate(&self) -> Result<bool> {
    Ok(self.storage.validate().await?)
  }

  #[napi]
  pub async fn validate_import_schema(&self) -> Result<bool> {
    Ok(self.storage.validate_import_schema().await?)
  }

  #[napi]
  pub async fn vacuum_into(&self, path: String) -> Result<()> {
    self.storage.vacuum_into(path).await?;
    Ok(())
  }

  #[napi]
  pub async fn set_space_id(&self, space_id: String) -> Result<()> {
    self.storage.connect().await?;
    self.storage.set_space_id(space_id).await?;
    self.storage.close().await;
    Ok(())
  }
}

#[cfg(all(test, not(feature = "use-as-lib")))]
mod tests {
  use super::error;

  #[test]
  fn napi_error_mapping_preserves_reason() {
    let err: napi::Error = error::Error::InvalidOperation.into();
    assert_eq!(err.status, napi::Status::GenericFailure);
    assert!(err.reason.contains("Invalid operation"));
  }

  #[test]
  fn napi_error_mapping_connection_in_progress() {
    let err: napi::Error = error::Error::ConnectionInProgress.into();
    assert_eq!(err.status, napi::Status::GenericFailure);
    assert!(err.reason.contains("Connection in progress"));
  }
}