fix(native): cleanup deleted docs and blobs (#14689)

packages/frontend/native/nbstore/src/lib.rs

@@ -129,6 +129,12 @@ impl DocStoragePool {
     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?;
@@ -485,6 +491,17 @@ impl DocStorage {
     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?;

packages/frontend/native/nbstore/src/storage.rs

@@ -1,6 +1,9 @@
 use std::sync::Arc;
 
-use affine_schema::get_migrator;
+use affine_schema::{
+  get_migrator,
+  import_validation::{V2_IMPORT_SCHEMA_RULES, validate_import_schema, validate_required_schema},
+};
 use memory_indexer::InMemoryIndex;
 use sqlx::{
   Pool, Row,
@@ -49,17 +52,27 @@ impl SqliteDocStorage {
   }
 
   pub async fn validate(&self) -> Result<bool> {
-    let record = sqlx::query("SELECT * FROM _sqlx_migrations ORDER BY installed_on ASC LIMIT 1;")
-      .fetch_optional(&self.pool)
-      .await;
-
-    match record {
-      Ok(Some(row)) => {
-        let name: &str = row.try_get("description")?;
-        Ok(name == "init_v2")
-      }
-      _ => Ok(false),
+    if self.path == ":memory:" {
+      return Ok(validate_required_schema(&self.pool, &V2_IMPORT_SCHEMA_RULES).await?);
     }
+
+    let Ok(pool) = self.open_readonly_pool().await else {
+      return Ok(false);
+    };
+
+    Ok(validate_required_schema(&pool, &V2_IMPORT_SCHEMA_RULES).await?)
+  }
+
+  pub async fn validate_import_schema(&self) -> Result<bool> {
+    if self.path == ":memory:" {
+      return Ok(validate_import_schema(&self.pool, &V2_IMPORT_SCHEMA_RULES).await?);
+    }
+
+    let Ok(pool) = self.open_readonly_pool().await else {
+      return Ok(false);
+    };
+
+    Ok(validate_import_schema(&pool, &V2_IMPORT_SCHEMA_RULES).await?)
   }
 
   pub async fn connect(&self) -> Result<()> {
@@ -159,14 +172,41 @@ impl SqliteDocStorage {
 
     Ok(())
   }
+
+  pub async fn vacuum_into(&self, path: String) -> Result<()> {
+    if self.path == ":memory:" {
+      sqlx::query("VACUUM INTO ?;").bind(path).execute(&self.pool).await?;
+      return Ok(());
+    }
+
+    let pool = self.open_readonly_pool().await?;
+    sqlx::query("VACUUM INTO ?;").bind(path).execute(&pool).await?;
+
+    Ok(())
+  }
+
+  async fn open_readonly_pool(&self) -> Result<Pool<Sqlite>> {
+    let sqlite_options = SqliteConnectOptions::new()
+      .filename(&self.path)
+      .foreign_keys(false)
+      .read_only(true);
+
+    Ok(
+      SqlitePoolOptions::new()
+        .max_connections(1)
+        .connect_with(sqlite_options)
+        .await?,
+    )
+  }
 }
 
 #[cfg(test)]
 mod tests {
-  use std::borrow::Cow;
+  use std::{borrow::Cow, fs, path::Path};
 
   use affine_schema::get_migrator;
   use sqlx::migrate::{Migration, Migrator};
+  use uuid::Uuid;
 
   use super::*;
 
@@ -256,4 +296,105 @@ mod tests {
 
     assert_eq!(checksum, expected_checksum);
   }
+
+  #[tokio::test]
+  async fn vacuum_into_exports_a_compacted_database() {
+    let base = std::env::temp_dir().join(format!("nbstore-vacuum-{}", Uuid::new_v4()));
+    fs::create_dir_all(&base).unwrap();
+
+    let source = base.join("storage.db");
+    let export = base.join("backup.affine");
+
+    let storage = SqliteDocStorage::new(path_string(&source));
+    storage.connect().await.unwrap();
+
+    storage
+      .set_blob(crate::SetBlob {
+        key: "large-blob".to_string(),
+        data: vec![7; 1024 * 1024],
+        mime: "application/octet-stream".to_string(),
+      })
+      .await
+      .unwrap();
+    storage.delete_blob("large-blob".to_string(), true).await.unwrap();
+    storage.checkpoint().await.unwrap();
+
+    let source_len = fs::metadata(&source).unwrap().len();
+    assert!(source_len > 0);
+
+    storage.vacuum_into(path_string(&export)).await.unwrap();
+
+    let export_len = fs::metadata(&export).unwrap().len();
+    assert!(export_len < source_len);
+
+    let exported = SqliteDocStorage::new(path_string(&export));
+    exported.connect().await.unwrap();
+    assert!(exported.list_blobs().await.unwrap().is_empty());
+    exported.close().await;
+    storage.close().await;
+
+    fs::remove_dir_all(base).unwrap();
+  }
+
+  #[tokio::test]
+  async fn validate_import_schema_rejects_unexpected_schema_objects() {
+    let base = std::env::temp_dir().join(format!("nbstore-schema-{}", Uuid::new_v4()));
+    fs::create_dir_all(&base).unwrap();
+
+    let source = base.join("storage.db");
+    fs::File::create(&source).unwrap();
+    let storage = SqliteDocStorage::new(path_string(&source));
+    storage.connect().await.unwrap();
+
+    sqlx::query("CREATE VIEW rogue_view AS SELECT space_id FROM meta")
+      .execute(&storage.pool)
+      .await
+      .unwrap();
+
+    assert!(!storage.validate_import_schema().await.unwrap());
+
+    storage.close().await;
+    fs::remove_dir_all(base).unwrap();
+  }
+
+  #[tokio::test]
+  async fn validate_import_schema_accepts_initial_v2_schema() {
+    let base = std::env::temp_dir().join(format!("nbstore-v2-schema-{}", Uuid::new_v4()));
+    fs::create_dir_all(&base).unwrap();
+
+    let source = base.join("storage.db");
+    let source_path = path_string(&source);
+    let setup_pool = SqlitePoolOptions::new()
+      .max_connections(1)
+      .connect_with(
+        SqliteConnectOptions::new()
+          .filename(&source_path)
+          .create_if_missing(true)
+          .foreign_keys(false),
+      )
+      .await
+      .unwrap();
+
+    let mut migrations = get_migrator().migrations.to_vec();
+    migrations.truncate(1);
+    let migrator = Migrator {
+      migrations: Cow::Owned(migrations),
+      ..Migrator::DEFAULT
+    };
+
+    migrator.run(&setup_pool).await.unwrap();
+    setup_pool.close().await;
+
+    let storage = SqliteDocStorage::new(source_path);
+
+    assert!(storage.validate().await.unwrap());
+    assert!(storage.validate_import_schema().await.unwrap());
+
+    storage.close().await;
+    fs::remove_dir_all(base).unwrap();
+  }
+
+  fn path_string(path: &Path) -> String {
+    path.to_string_lossy().into_owned()
+  }
 }