FIX: Add postgres storage adapter for samod

Author: jmoggr

Cargo.lock

@@ -22,6 +22,7 @@ jsonrpsee = "0.24.6"
 jsonrpsee-server = "0.24.6"
 notebook-types = { version = "0.1.0", path = "../notebook-types" }
 qubit = { version = "1.0.0-beta.0", features = ["ts-serde-json", "ts-uuid"] }
+rand = "0.8"
 regex = "1.11.1"
 samod = { git = "https://github.com/alexjg/samod", features = [
   "tokio",

packages/backend/Cargo.toml

@@ -0,0 +1 @@
+pub mod storage;

packages/backend/src/lib.rs

@@ -26,6 +26,7 @@ mod auth;
 mod automerge_json;
 mod document;
 mod rpc;
+mod storage;
 mod user;
 
 use app::AppStatus;
@@ -104,7 +105,7 @@ async fn main() {
 
             // Create samod repo
             let repo = samod::Repo::builder(tokio::runtime::Handle::current())
-                .with_storage(samod::storage::InMemoryStorage::new())
+                .with_storage(storage::PostgresStorage::new(db.clone()))
                 .load()
                 .await;
 

packages/backend/src/main.rs

@@ -0,0 +1,4 @@
+mod postgres;
+pub mod testing;
+
+pub use postgres::PostgresStorage;

packages/backend/src/storage/mod.rs

@@ -0,0 +1,110 @@
+use samod::storage::{Storage, StorageKey};
+use sqlx::PgPool;
+use std::collections::HashMap;
+
+/// A PostgreSQL-backed storage adapter for samod
+///
+/// ## Database Schema
+///
+/// The adapter requires a table with the following structure:
+/// ```sql
+/// CREATE TABLE storage (
+///     key text[] PRIMARY KEY,
+///     data bytea NOT NULL
+/// );
+/// ```
+#[derive(Clone)]
+pub struct PostgresStorage {
+    pool: PgPool,
+}
+
+impl PostgresStorage {
+    pub fn new(pool: PgPool) -> Self {
+        Self { pool }
+    }
+}
+
+impl Storage for PostgresStorage {
+    async fn load(&self, key: StorageKey) -> Option<Vec<u8>> {
+        let key_parts: Vec<String> = key.into_iter().collect();
+
+        let result = sqlx::query_scalar::<_, Vec<u8>>("SELECT data FROM storage WHERE key = $1")
+            .bind(&key_parts)
+            .fetch_optional(&self.pool)
+            .await;
+
+        match result {
+            Ok(data) => data,
+            Err(e) => {
+                tracing::error!("Failed to load from storage: {}", e);
+                None
+            }
+        }
+    }
+
+    async fn load_range(&self, prefix: StorageKey) -> HashMap<StorageKey, Vec<u8>> {
+        let prefix_parts: Vec<String> = prefix.into_iter().collect();
+
+        let result = if prefix_parts.is_empty() {
+            sqlx::query_as::<_, (Vec<String>, Vec<u8>)>("SELECT key, data FROM storage")
+                .fetch_all(&self.pool)
+                .await
+        } else {
+            sqlx::query_as::<_, (Vec<String>, Vec<u8>)>(
+                "SELECT key, data FROM storage WHERE key[1:cardinality($1::text[])] = $1::text[]",
+            )
+            .bind(&prefix_parts)
+            .fetch_all(&self.pool)
+            .await
+        };
+
+        match result {
+            Ok(rows) => {
+                let mut map = HashMap::new();
+                for (key_parts, data) in rows {
+                    if let Ok(storage_key) = StorageKey::from_parts(key_parts) {
+                        map.insert(storage_key, data);
+                    }
+                }
+                map
+            }
+            Err(e) => {
+                tracing::error!("Failed to load range from storage: {}", e);
+                HashMap::new()
+            }
+        }
+    }
+
+    async fn put(&self, key: StorageKey, data: Vec<u8>) {
+        let key_parts: Vec<String> = key.into_iter().collect();
+
+        let result = sqlx::query(
+            "
+            INSERT INTO storage (key, data)
+            VALUES ($1, $2)
+            ON CONFLICT (key) DO UPDATE SET data = $2
+            ",
+        )
+        .bind(&key_parts)
+        .bind(&data)
+        .execute(&self.pool)
+        .await;
+
+        if let Err(e) = result {
+            tracing::error!("Failed to put to storage: {}", e);
+        }
+    }
+
+    async fn delete(&self, key: StorageKey) {
+        let key_parts: Vec<String> = key.into_iter().collect();
+
+        let result = sqlx::query("DELETE FROM storage WHERE key = $1")
+            .bind(&key_parts)
+            .execute(&self.pool)
+            .await;
+
+        if let Err(e) = result {
+            tracing::error!("Failed to delete from storage: {}", e);
+        }
+    }
+}

packages/backend/src/storage/postgres.rs

@@ -0,0 +1,187 @@
+//! Storage adapter testing utilities
+//!
+//! Provides a test suite for any implementation of the `Storage` trait.
+//! Based on the TypeScript `runStorageAdapterTests` from automerge-repo.
+
+#![allow(dead_code)]
+
+use rand::Rng;
+use samod::storage::{Storage, StorageKey};
+use std::future::Future;
+use std::pin::Pin;
+use std::sync::LazyLock;
+
+pub fn payload_a() -> Vec<u8> {
+    vec![0, 1, 127, 99, 154, 235]
+}
+
+pub fn payload_b() -> Vec<u8> {
+    vec![1, 76, 160, 53, 57, 10, 230]
+}
+
+pub fn payload_c() -> Vec<u8> {
+    vec![2, 111, 74, 131, 236, 96, 142, 193]
+}
+
+static LARGE_PAYLOAD: LazyLock<Vec<u8>> = LazyLock::new(|| {
+    let mut vec = vec![0u8; 100000];
+    rand::thread_rng().fill(&mut vec[..]);
+    vec
+});
+
+pub fn large_payload() -> Vec<u8> {
+    LARGE_PAYLOAD.clone()
+}
+
+/// Trait for storage test fixtures
+pub trait StorageTestFixture: Sized + Send {
+    /// The storage type being tested
+    type Storage: Storage + Send + Sync + 'static;
+
+    /// Setup the test fixture
+    fn setup() -> impl std::future::Future<Output = Self> + Send;
+
+    /// Get reference to the storage adapter
+    fn storage(&self) -> &Self::Storage;
+
+    /// Optional cleanup
+    fn teardown(self) -> impl std::future::Future<Output = ()> + Send {
+        async {}
+    }
+}
+
+/// Helper to run a single test with setup and teardown
+async fn run_test<F, TestFn>(test_fn: TestFn)
+where
+    F: StorageTestFixture,
+    TestFn: for<'a> FnOnce(&'a F::Storage) -> Pin<Box<dyn Future<Output = ()> + Send + 'a>> + Send,
+{
+    let fixture = F::setup().await;
+    test_fn(fixture.storage()).await;
+    fixture.teardown().await;
+}
+
+/// Run all storage adapter acceptance tests
+pub async fn run_storage_adapter_tests<F: StorageTestFixture>() {
+    run_test::<F, _>(|a| Box::pin(test_load_should_return_none_if_no_data(a))).await;
+    run_test::<F, _>(|a| Box::pin(test_save_and_load_should_return_data_that_was_saved(a))).await;
+    run_test::<F, _>(|a| Box::pin(test_save_and_load_should_work_with_composite_keys(a))).await;
+    run_test::<F, _>(|a| Box::pin(test_save_and_load_should_work_with_large_payload(a))).await;
+    run_test::<F, _>(|a| Box::pin(test_load_range_should_return_empty_if_no_data(a))).await;
+    run_test::<F, _>(|a| Box::pin(test_save_and_load_range_should_return_all_matching_data(a)))
+        .await;
+    run_test::<F, _>(|a| Box::pin(test_save_and_load_range_should_only_load_matching_values(a)))
+        .await;
+    run_test::<F, _>(|a| Box::pin(test_save_and_remove_should_be_empty_after_removing(a))).await;
+    run_test::<F, _>(|a| Box::pin(test_save_and_save_should_overwrite(a))).await;
+}
+
+// describe("load")
+pub async fn test_load_should_return_none_if_no_data<S: Storage>(adapter: &S) {
+    let actual = adapter
+        .load(StorageKey::from_parts(["AAAAA", "sync-state", "xxxxx"]).unwrap())
+        .await;
+
+    assert_eq!(actual, None);
+}
+
+// describe("save and load")
+pub async fn test_save_and_load_should_return_data_that_was_saved<S: Storage>(adapter: &S) {
+    let key = StorageKey::from_parts(["storage-adapter-id"]).unwrap();
+    adapter.put(key.clone(), payload_a()).await;
+
+    let actual = adapter.load(key).await;
+
+    assert_eq!(actual, Some(payload_a()));
+}
+
+pub async fn test_save_and_load_should_work_with_composite_keys<S: Storage>(adapter: &S) {
+    let key = StorageKey::from_parts(["AAAAA", "sync-state", "xxxxx"]).unwrap();
+    adapter.put(key.clone(), payload_a()).await;
+
+    let actual = adapter.load(key).await;
+
+    assert_eq!(actual, Some(payload_a()));
+}
+
+pub async fn test_save_and_load_should_work_with_large_payload<S: Storage>(adapter: &S) {
+    let key = StorageKey::from_parts(["AAAAA", "sync-state", "xxxxx"]).unwrap();
+    adapter.put(key.clone(), large_payload()).await;
+
+    let actual = adapter.load(key).await;
+
+    assert_eq!(actual, Some(large_payload()));
+}
+
+// describe("loadRange")
+pub async fn test_load_range_should_return_empty_if_no_data<S: Storage>(adapter: &S) {
+    let result = adapter.load_range(StorageKey::from_parts(["AAAAA"]).unwrap()).await;
+
+    assert_eq!(result.len(), 0);
+}
+
+// describe("save and loadRange")
+pub async fn test_save_and_load_range_should_return_all_matching_data<S: Storage>(adapter: &S) {
+    let key_a = StorageKey::from_parts(["AAAAA", "sync-state", "xxxxx"]).unwrap();
+    let key_b = StorageKey::from_parts(["AAAAA", "snapshot", "yyyyy"]).unwrap();
+    let key_c = StorageKey::from_parts(["AAAAA", "sync-state", "zzzzz"]).unwrap();
+
+    adapter.put(key_a.clone(), payload_a()).await;
+    adapter.put(key_b.clone(), payload_b()).await;
+    adapter.put(key_c.clone(), payload_c()).await;
+
+    let result = adapter.load_range(StorageKey::from_parts(["AAAAA"]).unwrap()).await;
+
+    assert_eq!(result.len(), 3);
+    assert_eq!(result.get(&key_a), Some(&payload_a()));
+    assert_eq!(result.get(&key_b), Some(&payload_b()));
+    assert_eq!(result.get(&key_c), Some(&payload_c()));
+
+    let sync_result = adapter
+        .load_range(StorageKey::from_parts(["AAAAA", "sync-state"]).unwrap())
+        .await;
+
+    assert_eq!(sync_result.len(), 2);
+    assert_eq!(sync_result.get(&key_a), Some(&payload_a()));
+    assert_eq!(sync_result.get(&key_c), Some(&payload_c()));
+}
+
+pub async fn test_save_and_load_range_should_only_load_matching_values<S: Storage>(adapter: &S) {
+    let key_a = StorageKey::from_parts(["AAAAA", "sync-state", "xxxxx"]).unwrap();
+    let key_c = StorageKey::from_parts(["BBBBB", "sync-state", "zzzzz"]).unwrap();
+
+    adapter.put(key_a.clone(), payload_a()).await;
+    adapter.put(key_c.clone(), payload_c()).await;
+
+    let actual = adapter.load_range(StorageKey::from_parts(["AAAAA"]).unwrap()).await;
+
+    assert_eq!(actual.len(), 1);
+    assert_eq!(actual.get(&key_a), Some(&payload_a()));
+}
+
+// describe("save and remove")
+pub async fn test_save_and_remove_should_be_empty_after_removing<S: Storage>(adapter: &S) {
+    let key = StorageKey::from_parts(["AAAAA", "snapshot", "xxxxx"]).unwrap();
+    adapter.put(key.clone(), payload_a()).await;
+    adapter.delete(key.clone()).await;
+
+    let range_result = adapter.load_range(StorageKey::from_parts(["AAAAA"]).unwrap()).await;
+    assert_eq!(range_result.len(), 0);
+
+    let load_result = adapter.load(key).await;
+    assert_eq!(load_result, None);
+}
+
+// describe("save and save")
+pub async fn test_save_and_save_should_overwrite<S: Storage>(adapter: &S) {
+    let key = StorageKey::from_parts(["AAAAA", "sync-state", "xxxxx"]).unwrap();
+    adapter.put(key.clone(), payload_a()).await;
+    adapter.put(key.clone(), payload_b()).await;
+
+    let result = adapter
+        .load_range(StorageKey::from_parts(["AAAAA", "sync-state"]).unwrap())
+        .await;
+
+    assert_eq!(result.len(), 1);
+    assert_eq!(result.get(&key), Some(&payload_b()));
+}

packages/backend/src/storage/testing.rs

@@ -0,0 +1,44 @@
+use backend::storage::{PostgresStorage, testing};
+use sqlx::PgPool;
+
+async fn cleanup_test_data(pool: &PgPool) {
+    let _ = sqlx::query("DELETE FROM storage WHERE key[1] = ANY($1)")
+        .bind(&["AAAAA", "BBBBB", "storage-adapter-id"])
+        .execute(pool)
+        .await;
+}
+
+struct PostgresTestFixture {
+    storage: PostgresStorage,
+    pool: PgPool,
+}
+
+impl testing::StorageTestFixture for PostgresTestFixture {
+    type Storage = PostgresStorage;
+
+    async fn setup() -> Self {
+        let database_url =
+            std::env::var("DATABASE_URL").expect("DATABASE_URL must be set for tests");
+
+        let pool = PgPool::connect(&database_url).await.expect("Failed to connect to database");
+
+        cleanup_test_data(&pool).await;
+
+        let storage = PostgresStorage::new(pool.clone());
+
+        Self { storage, pool }
+    }
+
+    fn storage(&self) -> &PostgresStorage {
+        &self.storage
+    }
+
+    async fn teardown(self) {
+        cleanup_test_data(&self.pool).await;
+    }
+}
+
+#[tokio::test]
+async fn postgres_storage_adapter_tests() {
+    testing::run_storage_adapter_tests::<PostgresTestFixture>().await;
+}