feat: sqlite subdocument (#2816)

Co-authored-by: Alex Yang <himself65@outlook.com>

packages/native/build.rs

@@ -1,12 +1,22 @@
 use sqlx::sqlite::SqliteConnectOptions;
+use std::fs;
 
 #[tokio::main]
 async fn main() -> Result<(), std::io::Error> {
   dotenv::dotenv().ok();
 
+  // always start with a fresh database to have
+  // latest db schema
+  let db_path = "../../affine.db";
+
+  // check if db exists and then remove file
+  if fs::metadata(db_path).is_ok() {
+    fs::remove_file(db_path)?;
+  }
+
   napi_build::setup();
   let options = SqliteConnectOptions::new()
-    .filename("../../affine.db")
+    .filename(db_path)
     .journal_mode(sqlx::sqlite::SqliteJournalMode::Off)
     .locking_mode(sqlx::sqlite::SqliteLockingMode::Exclusive)
     .create_if_missing(true);

packages/native/index.d.ts

@@ -32,6 +32,11 @@ export interface UpdateRow {
   id: number;
   timestamp: Date;
   data: Buffer;
+  docId?: string;
+}
+export interface InsertRow {
+  docId?: string;
+  data: Uint8Array;
 }
 export class Subscription {
   toString(): string;
@@ -56,8 +61,14 @@ export class SqliteConnection {
   getBlob(key: string): Promise<BlobRow | null>;
   deleteBlob(key: string): Promise<void>;
   getBlobKeys(): Promise<Array<string>>;
-  getUpdates(): Promise<Array<UpdateRow>>;
-  insertUpdates(updates: Array<Uint8Array>): Promise<void>;
+  getUpdates(docId?: string | undefined | null): Promise<Array<UpdateRow>>;
+  getUpdatesCount(docId?: string | undefined | null): Promise<number>;
+  getAllUpdates(): Promise<Array<UpdateRow>>;
+  insertUpdates(updates: Array<InsertRow>): Promise<void>;
+  replaceUpdates(
+    docId: string | undefined | null,
+    updates: Array<InsertRow>
+  ): Promise<void>;
   close(): Promise<void>;
   get isClose(): boolean;
   static validate(path: string): Promise<boolean>;

packages/native/schema/src/lib.rs

@@ -4,7 +4,8 @@
 pub const SCHEMA: &str = r#"CREATE TABLE IF NOT EXISTS "updates" (
   id INTEGER PRIMARY KEY AUTOINCREMENT,
   data BLOB NOT NULL,
-  timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP NOT NULL
+  timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP NOT NULL,
+  doc_id TEXT
 );
 CREATE TABLE IF NOT EXISTS "blobs" (
   key TEXT PRIMARY KEY NOT NULL,

packages/native/src/sqlite/mod.rs

@@ -19,6 +19,13 @@ pub struct UpdateRow {
   pub id: i64,
   pub timestamp: NaiveDateTime,
   pub data: Buffer,
+  pub doc_id: Option<String>,
+}
+
+#[napi(object)]
+pub struct InsertRow {
+  pub doc_id: Option<String>,
+  pub data: Uint8Array,
 }
 
 #[napi]
@@ -53,6 +60,7 @@ impl SqliteConnection {
       .execute(connection.as_mut())
       .await
       .map_err(anyhow::Error::from)?;
+    self.migrate_add_doc_id().await?;
     connection.detach();
     Ok(())
   }
@@ -74,10 +82,14 @@ impl SqliteConnection {
 
   #[napi]
   pub async fn get_blob(&self, key: String) -> Option<BlobRow> {
-    sqlx::query_as!(BlobRow, "SELECT * FROM blobs WHERE key = ?", key)
-      .fetch_one(&self.pool)
-      .await
-      .ok()
+    sqlx::query_as!(
+      BlobRow,
+      "SELECT key, data, timestamp FROM blobs WHERE key = ?",
+      key
+    )
+    .fetch_one(&self.pool)
+    .await
+    .ok()
   }
 
   #[napi]
@@ -100,8 +112,54 @@ impl SqliteConnection {
   }
 
   #[napi]
-  pub async fn get_updates(&self) -> napi::Result<Vec<UpdateRow>> {
-    let updates = sqlx::query_as!(UpdateRow, "SELECT * FROM updates")
+  pub async fn get_updates(&self, doc_id: Option<String>) -> napi::Result<Vec<UpdateRow>> {
+    let updates = match doc_id {
+      Some(doc_id) => sqlx::query_as!(
+        UpdateRow,
+        "SELECT id, timestamp, data, doc_id FROM updates WHERE doc_id = ?",
+        doc_id
+      )
+      .fetch_all(&self.pool)
+      .await
+      .map_err(anyhow::Error::from)?,
+      None => sqlx::query_as!(
+        UpdateRow,
+        "SELECT id, timestamp, data, doc_id FROM updates WHERE doc_id is NULL",
+      )
+      .fetch_all(&self.pool)
+      .await
+      .map_err(anyhow::Error::from)?,
+    };
+    Ok(updates)
+  }
+
+  #[napi]
+  pub async fn get_updates_count(&self, doc_id: Option<String>) -> napi::Result<i32> {
+    let count = match doc_id {
+      Some(doc_id) => {
+        sqlx::query!(
+          "SELECT COUNT(*) as count FROM updates WHERE doc_id = ?",
+          doc_id
+        )
+        .fetch_one(&self.pool)
+        .await
+        .map_err(anyhow::Error::from)?
+        .count
+      }
+      None => {
+        sqlx::query!("SELECT COUNT(*) as count FROM updates WHERE doc_id is NULL")
+          .fetch_one(&self.pool)
+          .await
+          .map_err(anyhow::Error::from)?
+          .count
+      }
+    };
+    Ok(count)
+  }
+
+  #[napi]
+  pub async fn get_all_updates(&self) -> napi::Result<Vec<UpdateRow>> {
+    let updates = sqlx::query_as!(UpdateRow, "SELECT id, timestamp, data, doc_id FROM updates")
       .fetch_all(&self.pool)
       .await
       .map_err(anyhow::Error::from)?;
@@ -109,14 +167,54 @@ impl SqliteConnection {
   }
 
   #[napi]
-  pub async fn insert_updates(&self, updates: Vec<Uint8Array>) -> napi::Result<()> {
+  pub async fn insert_updates(&self, updates: Vec<InsertRow>) -> napi::Result<()> {
     let mut transaction = self.pool.begin().await.map_err(anyhow::Error::from)?;
-    for update in updates.into_iter() {
-      let update = update.as_ref();
-      sqlx::query_as!(UpdateRow, "INSERT INTO updates (data) VALUES ($1)", update)
+    for InsertRow { data, doc_id } in updates {
+      let update = data.as_ref();
+      sqlx::query_as!(
+        UpdateRow,
+        "INSERT INTO updates (data, doc_id) VALUES ($1, $2)",
+        update,
+        doc_id
+      )
+      .execute(&mut *transaction)
+      .await
+      .map_err(anyhow::Error::from)?;
+    }
+    transaction.commit().await.map_err(anyhow::Error::from)?;
+    Ok(())
+  }
+
+  #[napi]
+  pub async fn replace_updates(
+    &self,
+    doc_id: Option<String>,
+    updates: Vec<InsertRow>,
+  ) -> napi::Result<()> {
+    let mut transaction = self.pool.begin().await.map_err(anyhow::Error::from)?;
+
+    match doc_id {
+      Some(doc_id) => sqlx::query!("DELETE FROM updates where doc_id = ?", doc_id)
         .execute(&mut *transaction)
         .await
-        .map_err(anyhow::Error::from)?;
+        .map_err(anyhow::Error::from)?,
+      None => sqlx::query!("DELETE FROM updates where doc_id is NULL",)
+        .execute(&mut *transaction)
+        .await
+        .map_err(anyhow::Error::from)?,
+    };
+
+    for InsertRow { data, doc_id } in updates {
+      let update = data.as_ref();
+      sqlx::query_as!(
+        UpdateRow,
+        "INSERT INTO updates (data, doc_id) VALUES ($1, $2)",
+        update,
+        doc_id
+      )
+      .execute(&mut *transaction)
+      .await
+      .map_err(anyhow::Error::from)?;
     }
     transaction.commit().await.map_err(anyhow::Error::from)?;
     Ok(())
@@ -158,4 +256,22 @@ impl SqliteConnection {
       false
     }
   }
+
+  // todo: have a better way to handle migration
+  async fn migrate_add_doc_id(&self) -> Result<(), anyhow::Error> {
+    // ignore errors
+    match sqlx::query("ALTER TABLE updates ADD COLUMN doc_id TEXT")
+      .execute(&self.pool)
+      .await
+    {
+      Ok(_) => Ok(()),
+      Err(err) => {
+        if err.to_string().contains("duplicate column name") {
+          Ok(()) // Ignore error if it's due to duplicate column
+        } else {
+          Err(anyhow::Error::from(err)) // Propagate other errors
+        }
+      }
+    }
+  }
 }