hocuspocus_extension_sqlite/
sqlite.rs

1use anyhow::Result;
2use async_trait::async_trait;
3use hocuspocus_extension_database::types::{FetchContext, StoreContext};
4use hocuspocus_extension_database::DatabaseExtension;
5use sqlx::SqlitePool;
6
7pub struct SqliteDatabase {
8    pool: SqlitePool,
9}
10
11impl SqliteDatabase {
12    pub async fn connect(database_url: &str) -> Result<Self> {
13        let pool = SqlitePool::connect(database_url).await?;
14        sqlx::query(
15            "CREATE TABLE IF NOT EXISTS documents (name TEXT PRIMARY KEY, state BLOB NOT NULL, updated_at INTEGER NOT NULL)"
16        ).execute(&pool).await?;
17        Ok(Self { pool })
18    }
19
20    pub fn from_pool(pool: SqlitePool) -> Self {
21        Self { pool }
22    }
23}
24
25#[async_trait]
26impl DatabaseExtension for SqliteDatabase {
27    async fn fetch(&self, ctx: FetchContext) -> Result<Option<Vec<u8>>> {
28        let bytes: Option<Vec<u8>> =
29            sqlx::query_scalar("SELECT state FROM documents WHERE name = ?")
30                .bind(&ctx.document_name)
31                .fetch_optional(&self.pool)
32                .await?;
33        Ok(bytes)
34    }
35
36    async fn store(&self, ctx: StoreContext<'_>) -> Result<()> {
37        sqlx::query(
38            "INSERT INTO documents(name, state, updated_at) VALUES(?, ?, ?)\n             ON CONFLICT(name) DO UPDATE SET state = excluded.state, updated_at = excluded.updated_at"
39        )
40            .bind(&ctx.document_name)
41            .bind(ctx.state)
42            .bind(ctx.updated_at_millis)
43            .execute(&self.pool)
44            .await?;
45        Ok(())
46    }
47}