snapshot of https://github.com/toeverything/AFFiNE.git (canary, 2026-09-16; blobs > 131072B omitted)

packages/frontend/native/nbstore/src/indexer/persistence.rs

@@ -0,0 +1,193 @@
+use std::sync::atomic::Ordering;
+
+use memory_indexer::{Document, MemoryIndex, TermsAggregation, Value};
+use sqlx::Row;
+
+use super::{
+  DocIndexedClock, NBSTORE_INDEXER_VERSION, NativeIndexAggregateResult, NativeIndexBucket, NativeIndexDocument,
+  NativeIndexQuery, NativeIndexSearchOptions, NativeIndexSearchResult, SqliteDocStorage, error::Result, string_values,
+};
+
+impl SqliteDocStorage {
+  pub async fn init_index(&self) -> Result<()> {
+    let snapshots = sqlx::query("SELECT index_name, data FROM idx_snapshots WHERE index_name IN ('doc', 'block')")
+      .fetch_all(&self.pool)
+      .await?;
+    let mut corrupted = false;
+    for row in snapshots {
+      let name: String = row.get("index_name");
+      let data: Vec<u8> = row.get("data");
+      let table = self.indexes.table(&name)?;
+      if let Ok(index) = MemoryIndex::from_checkpoint(table.schema.clone(), &data) {
+        *table.index.write().await = index;
+      } else {
+        corrupted = true;
+      }
+    }
+    if corrupted {
+      self.indexes.reset().await;
+      let mut tx = self.pool.begin().await?;
+      sqlx::query("DELETE FROM idx_snapshots").execute(&mut *tx).await?;
+      sqlx::query("DELETE FROM indexer_sync").execute(&mut *tx).await?;
+      tx.commit().await?;
+    }
+    sqlx::query("DELETE FROM idx_snapshots WHERE index_name NOT IN ('doc', 'block')")
+      .execute(&self.pool)
+      .await?;
+    Ok(())
+  }
+
+  pub async fn flush_index(&self) -> Result<()> {
+    let mut checkpoints = Vec::new();
+    for (name, table) in self.indexes.tables() {
+      let index = table.index.read().await;
+      if index.has_unpersisted_changes() {
+        checkpoints.push((name, index.checkpoint()?));
+      }
+    }
+    if checkpoints.is_empty() {
+      return Ok(());
+    }
+    let mut tx = self.pool.begin().await?;
+    for (name, checkpoint) in &checkpoints {
+      sqlx::query("INSERT OR REPLACE INTO idx_snapshots (index_name, data) VALUES (?, ?)")
+        .bind(name)
+        .bind(&checkpoint.bytes)
+        .execute(&mut *tx)
+        .await?;
+    }
+    tx.commit().await?;
+    for (name, checkpoint) in checkpoints {
+      self
+        .indexes
+        .table(name)?
+        .index
+        .write()
+        .await
+        .mark_checkpoint_persisted(checkpoint.sequence)?;
+    }
+    Ok(())
+  }
+
+  pub async fn commit_indexed_clocks(&self, clocks: &[DocIndexedClock]) -> Result<()> {
+    let mut checkpoints = Vec::new();
+    for (name, table) in self.indexes.tables() {
+      checkpoints.push((name, table.index.read().await.checkpoint()?));
+    }
+    let mut tx = self.pool.begin().await?;
+    for (name, checkpoint) in &checkpoints {
+      sqlx::query("INSERT OR REPLACE INTO idx_snapshots (index_name, data) VALUES (?, ?)")
+        .bind(name)
+        .bind(&checkpoint.bytes)
+        .execute(&mut *tx)
+        .await?;
+    }
+    for clock in clocks {
+      sqlx::query(
+        r#"INSERT INTO indexer_sync (doc_id, indexed_clock, indexer_version)
+           VALUES ($1, $2, $3)
+           ON CONFLICT(doc_id)
+           DO UPDATE SET indexed_clock=$2, indexer_version=$3"#,
+      )
+      .bind(&clock.doc_id)
+      .bind(clock.timestamp)
+      .bind(clock.indexer_version)
+      .execute(&mut *tx)
+      .await?;
+    }
+    tx.commit().await?;
+    for (name, checkpoint) in checkpoints {
+      self
+        .indexes
+        .table(name)?
+        .index
+        .write()
+        .await
+        .mark_checkpoint_persisted(checkpoint.sequence)?;
+    }
+    self.indexes.ready.store(true, Ordering::Release);
+    Ok(())
+  }
+
+  pub fn index_version() -> u32 {
+    NBSTORE_INDEXER_VERSION
+  }
+
+  pub async fn index_upsert(&self, table: &str, document: NativeIndexDocument) -> Result<()> {
+    let table = self.indexes.table(table)?;
+    let mut index_document = Document::new(document.id);
+    for field in document.fields {
+      let field_id = table.field(&field.field)?;
+      index_document.add_values(field_id, field.values.into_iter().map(Value::String));
+    }
+    table.index.write().await.upsert(index_document)?;
+    Ok(())
+  }
+
+  pub async fn index_delete(&self, table: &str, id: &str) -> Result<()> {
+    self.indexes.table(table)?.index.write().await.delete(id);
+    Ok(())
+  }
+
+  pub async fn index_search(
+    &self,
+    table: &str,
+    query: NativeIndexQuery,
+    options: NativeIndexSearchOptions,
+  ) -> Result<NativeIndexSearchResult> {
+    self.indexes.ensure_ready()?;
+    let table = self.indexes.table(table)?;
+    let query = table.compile_query(query)?;
+    let result = table
+      .index
+      .read()
+      .await
+      .search(&query, table.compile_options(options)?)?;
+    Ok(NativeIndexSearchResult {
+      total: result.total as u32,
+      hits: result.hits.into_iter().map(|hit| table.hit(hit)).collect(),
+    })
+  }
+
+  pub async fn index_aggregate(
+    &self,
+    table: &str,
+    query: NativeIndexQuery,
+    field: &str,
+    limit: u32,
+    offset: u32,
+    hits: Option<NativeIndexSearchOptions>,
+  ) -> Result<NativeIndexAggregateResult> {
+    self.indexes.ensure_ready()?;
+    let table = self.indexes.table(table)?;
+    let query = table.compile_query(query)?;
+    let result = table.index.read().await.aggregate(
+      &query,
+      TermsAggregation {
+        field: table.field(field)?,
+        limit: limit as usize,
+        offset: offset as usize,
+        top_hits: hits.map(|options| table.compile_options(options)).transpose()?,
+      },
+    )?;
+    Ok(NativeIndexAggregateResult {
+      total: result.total as u32,
+      buckets: result
+        .buckets
+        .into_iter()
+        .map(|bucket| NativeIndexBucket {
+          key: string_values(vec![bucket.key]).pop().unwrap_or_default(),
+          count: bucket.count as u32,
+          score: bucket.max_score as f64,
+          hits: bucket.hits.into_iter().map(|hit| table.hit(hit)).collect(),
+        })
+        .collect(),
+    })
+  }
+
+  pub async fn index_delete_by_query(&self, table: &str, query: NativeIndexQuery) -> Result<u32> {
+    let table = self.indexes.table(table)?;
+    let query = table.compile_query(query)?;
+    Ok(table.index.write().await.delete_by_query(&query)? as u32)
+  }
+}