Skip to main content

agentos_vfs/local/
sqlite_metadata_store.rs

1use async_trait::async_trait;
2use rusqlite::{params, Connection};
3use std::collections::BTreeMap;
4use std::path::Path;
5use std::sync::Mutex;
6use vfs::engine::error::{VfsError, VfsResult};
7use vfs::engine::mem::metadata_store::MetadataDump;
8use vfs::engine::mem::InMemoryMetadataStore;
9use vfs::engine::metadata::MetadataStore;
10use vfs::engine::types::{
11    BlockKey, ChunkEdit, ChunkRange, ChunkRef, CreateInodeAttrs, DentryStat, InodeMeta, InodePatch,
12    InodeType, SnapshotId, Storage, Timespec, DEFAULT_CHUNK_SIZE,
13};
14
15pub struct SqliteMetadataStore {
16    connection: Mutex<Connection>,
17    inner: InMemoryMetadataStore,
18}
19
20impl SqliteMetadataStore {
21    pub fn open(path: impl AsRef<Path>) -> VfsResult<Self> {
22        let connection = Connection::open(path)
23            .map_err(|err| VfsError::eio(format!("open SQLite metadata store: {err}")))?;
24        Self::from_connection(connection)
25    }
26
27    pub fn in_memory() -> VfsResult<Self> {
28        let connection = Connection::open_in_memory()
29            .map_err(|err| VfsError::eio(format!("open in-memory SQLite metadata store: {err}")))?;
30        Self::from_connection(connection)
31    }
32
33    fn from_connection(mut connection: Connection) -> VfsResult<Self> {
34        install_schema(&connection)?;
35        let dump = load_dump(&connection)?;
36        let inner = dump
37            .map(InMemoryMetadataStore::from_dump)
38            .unwrap_or_default();
39        if load_dump(&connection)?.is_none() {
40            persist_dump(&mut connection, &inner.dump())?;
41        }
42        Ok(Self {
43            connection: Mutex::new(connection),
44            inner,
45        })
46    }
47
48    pub fn has_schema(&self) -> VfsResult<bool> {
49        let connection = self.connection.lock().expect("sqlite mutex poisoned");
50        let count: i64 = connection
51            .query_row(
52                "SELECT COUNT(*) FROM sqlite_master WHERE type = 'table' AND name IN ('inodes', 'dentries', 'chunks', 'block_refs', 'snapshots')",
53                [],
54                |row| row.get(0),
55            )
56            .map_err(|err| VfsError::eio(format!("inspect SQLite schema: {err}")))?;
57        Ok(count == 5)
58    }
59
60    fn persist(&self) -> VfsResult<()> {
61        let dump = self.inner.dump();
62        let mut connection = self.connection.lock().expect("sqlite mutex poisoned");
63        persist_dump(&mut connection, &dump)
64    }
65}
66
67fn install_schema(connection: &Connection) -> VfsResult<()> {
68    connection
69        .execute_batch(
70            "
71            CREATE TABLE IF NOT EXISTS inodes (
72              ino INTEGER PRIMARY KEY,
73              kind INTEGER NOT NULL,
74              mode INTEGER NOT NULL,
75              uid INTEGER NOT NULL,
76              gid INTEGER NOT NULL,
77              size INTEGER NOT NULL,
78              nlink INTEGER NOT NULL,
79              atime_ns INTEGER NOT NULL, mtime_ns INTEGER NOT NULL,
80              ctime_ns INTEGER NOT NULL, birthtime_ns INTEGER NOT NULL,
81              storage_mode INTEGER NOT NULL,
82              storage_chunk_size INTEGER,
83              inline_content BLOB,
84              symlink_target TEXT
85            );
86            CREATE TABLE IF NOT EXISTS dentries (
87              parent_ino INTEGER NOT NULL,
88              name TEXT NOT NULL,
89              child_ino INTEGER NOT NULL,
90              kind INTEGER NOT NULL,
91              PRIMARY KEY (parent_ino, name)
92            );
93            CREATE INDEX IF NOT EXISTS dentries_parent ON dentries(parent_ino);
94            CREATE TABLE IF NOT EXISTS chunks (
95              ino INTEGER NOT NULL,
96              chunk_index INTEGER NOT NULL,
97              block_key TEXT NOT NULL,
98              len INTEGER NOT NULL,
99              PRIMARY KEY (ino, chunk_index)
100            );
101            CREATE TABLE IF NOT EXISTS block_refs (
102              block_key TEXT PRIMARY KEY,
103              refcount INTEGER NOT NULL
104            );
105            CREATE TABLE IF NOT EXISTS snapshots (
106              snapshot_id INTEGER PRIMARY KEY,
107              root_ino INTEGER NOT NULL,
108              created_ns INTEGER NOT NULL
109            );
110            ",
111        )
112        .map_err(|err| VfsError::eio(format!("install SQLite metadata schema: {err}")))
113}
114
115fn load_dump(connection: &Connection) -> VfsResult<Option<MetadataDump>> {
116    let inode_count: i64 = connection
117        .query_row("SELECT COUNT(*) FROM inodes", [], |row| row.get(0))
118        .map_err(|err| VfsError::eio(format!("count SQLite inodes: {err}")))?;
119    if inode_count == 0 {
120        return Ok(None);
121    }
122
123    let mut inodes = BTreeMap::new();
124    let mut next_ino = 1;
125    let mut statement = connection
126        .prepare(
127            "SELECT ino, kind, mode, uid, gid, size, nlink, atime_ns, mtime_ns, ctime_ns,
128                    birthtime_ns, storage_mode, storage_chunk_size, inline_content, symlink_target
129             FROM inodes",
130        )
131        .map_err(|err| VfsError::eio(format!("prepare inode load: {err}")))?;
132    let rows = statement
133        .query_map([], |row| {
134            let ino: u64 = row.get(0)?;
135            let kind_id: i64 = row.get(1)?;
136            let storage_id: i64 = row.get(11)?;
137            let chunk_size: Option<u32> = row.get(12)?;
138            let inline_content: Option<Vec<u8>> = row.get(13)?;
139            let symlink_target: Option<String> = row.get(14)?;
140            let kind = match kind_id {
141                0 => InodeType::File,
142                1 => InodeType::Directory,
143                _ => InodeType::Symlink,
144            };
145            let storage = match storage_id {
146                1 => Storage::Inline(inline_content.unwrap_or_default()),
147                2 => Storage::Chunked {
148                    chunk_size: chunk_size.unwrap_or(DEFAULT_CHUNK_SIZE),
149                },
150                _ => Storage::None,
151            };
152            Ok(InodeMeta {
153                ino,
154                kind,
155                mode: row.get(2)?,
156                uid: row.get(3)?,
157                gid: row.get(4)?,
158                size: row.get(5)?,
159                nlink: row.get(6)?,
160                atime: ns_to_timespec(row.get(7)?),
161                mtime: ns_to_timespec(row.get(8)?),
162                ctime: ns_to_timespec(row.get(9)?),
163                birthtime: ns_to_timespec(row.get(10)?),
164                storage,
165                symlink_target,
166            })
167        })
168        .map_err(|err| VfsError::eio(format!("load SQLite inodes: {err}")))?;
169    for row in rows {
170        let meta = row.map_err(|err| VfsError::eio(format!("load SQLite inode row: {err}")))?;
171        next_ino = next_ino.max(meta.ino + 1);
172        inodes.insert(meta.ino, meta);
173    }
174
175    let mut dentries = BTreeMap::new();
176    let mut statement = connection
177        .prepare("SELECT parent_ino, name, child_ino FROM dentries")
178        .map_err(|err| VfsError::eio(format!("prepare dentry load: {err}")))?;
179    let rows = statement
180        .query_map([], |row| {
181            Ok((
182                (row.get::<_, u64>(0)?, row.get::<_, String>(1)?),
183                row.get::<_, u64>(2)?,
184            ))
185        })
186        .map_err(|err| VfsError::eio(format!("load SQLite dentries: {err}")))?;
187    for row in rows {
188        let (key, value) =
189            row.map_err(|err| VfsError::eio(format!("load SQLite dentry row: {err}")))?;
190        dentries.insert(key, value);
191    }
192
193    let mut chunks = BTreeMap::new();
194    let mut statement = connection
195        .prepare("SELECT ino, chunk_index, block_key, len FROM chunks")
196        .map_err(|err| VfsError::eio(format!("prepare chunk load: {err}")))?;
197    let rows = statement
198        .query_map([], |row| {
199            let index = row.get::<_, u64>(1)?;
200            Ok((
201                (row.get::<_, u64>(0)?, index),
202                ChunkRef {
203                    index,
204                    key: BlockKey(row.get(2)?),
205                    len: row.get(3)?,
206                },
207            ))
208        })
209        .map_err(|err| VfsError::eio(format!("load SQLite chunks: {err}")))?;
210    for row in rows {
211        let (key, value) =
212            row.map_err(|err| VfsError::eio(format!("load SQLite chunk row: {err}")))?;
213        chunks.insert(key, value);
214    }
215
216    let mut block_refs = BTreeMap::new();
217    let mut statement = connection
218        .prepare("SELECT block_key, refcount FROM block_refs")
219        .map_err(|err| VfsError::eio(format!("prepare block ref load: {err}")))?;
220    let rows = statement
221        .query_map([], |row| Ok((BlockKey(row.get(0)?), row.get::<_, u64>(1)?)))
222        .map_err(|err| VfsError::eio(format!("load SQLite block refs: {err}")))?;
223    for row in rows {
224        let (key, value) =
225            row.map_err(|err| VfsError::eio(format!("load SQLite block ref row: {err}")))?;
226        block_refs.insert(key, value);
227    }
228
229    Ok(Some(MetadataDump {
230        next_ino,
231        inodes,
232        dentries,
233        chunks,
234        block_refs,
235    }))
236}
237
238fn persist_dump(connection: &mut Connection, dump: &MetadataDump) -> VfsResult<()> {
239    let tx = connection
240        .transaction()
241        .map_err(|err| VfsError::eio(format!("begin SQLite metadata transaction: {err}")))?;
242    tx.execute_batch(
243        "
244        DELETE FROM snapshots;
245        DELETE FROM block_refs;
246        DELETE FROM chunks;
247        DELETE FROM dentries;
248        DELETE FROM inodes;
249        ",
250    )
251    .map_err(|err| VfsError::eio(format!("clear SQLite metadata tables: {err}")))?;
252
253    for meta in dump.inodes.values() {
254        let (storage_mode, storage_chunk_size, inline_content) = match &meta.storage {
255            Storage::None => (0, None, None),
256            Storage::Inline(data) => (1, None, Some(data.as_slice())),
257            Storage::Chunked { chunk_size } => (2, Some(*chunk_size), None),
258        };
259        tx.execute(
260            "INSERT INTO inodes
261             (ino, kind, mode, uid, gid, size, nlink, atime_ns, mtime_ns, ctime_ns, birthtime_ns,
262              storage_mode, storage_chunk_size, inline_content, symlink_target)
263             VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
264            params![
265                meta.ino,
266                kind_id(meta.kind),
267                meta.mode,
268                meta.uid,
269                meta.gid,
270                meta.size,
271                meta.nlink,
272                timespec_to_ns(meta.atime),
273                timespec_to_ns(meta.mtime),
274                timespec_to_ns(meta.ctime),
275                timespec_to_ns(meta.birthtime),
276                storage_mode,
277                storage_chunk_size,
278                inline_content,
279                meta.symlink_target,
280            ],
281        )
282        .map_err(|err| VfsError::eio(format!("persist SQLite inode {}: {err}", meta.ino)))?;
283    }
284
285    for ((parent, name), child) in &dump.dentries {
286        let kind = dump
287            .inodes
288            .get(child)
289            .map(|meta| meta.kind)
290            .ok_or_else(|| VfsError::eio(format!("dentry points to missing inode {child}")))?;
291        tx.execute(
292            "INSERT INTO dentries (parent_ino, name, child_ino, kind) VALUES (?, ?, ?, ?)",
293            params![parent, name, child, kind_id(kind)],
294        )
295        .map_err(|err| VfsError::eio(format!("persist SQLite dentry {name}: {err}")))?;
296    }
297
298    for ((ino, index), chunk) in &dump.chunks {
299        tx.execute(
300            "INSERT INTO chunks (ino, chunk_index, block_key, len) VALUES (?, ?, ?, ?)",
301            params![ino, index, chunk.key.0, chunk.len],
302        )
303        .map_err(|err| VfsError::eio(format!("persist SQLite chunk {ino}/{index}: {err}")))?;
304    }
305
306    for (key, refcount) in &dump.block_refs {
307        tx.execute(
308            "INSERT INTO block_refs (block_key, refcount) VALUES (?, ?)",
309            params![key.0, refcount],
310        )
311        .map_err(|err| VfsError::eio(format!("persist SQLite block ref {}: {err}", key.0)))?;
312    }
313
314    tx.commit()
315        .map_err(|err| VfsError::eio(format!("commit SQLite metadata transaction: {err}")))
316}
317
318fn kind_id(kind: InodeType) -> i64 {
319    match kind {
320        InodeType::File => 0,
321        InodeType::Directory => 1,
322        InodeType::Symlink => 2,
323    }
324}
325
326fn timespec_to_ns(time: Timespec) -> i64 {
327    time.sec.saturating_mul(1_000_000_000) + i64::from(time.nsec)
328}
329
330fn ns_to_timespec(ns: i64) -> Timespec {
331    Timespec {
332        sec: ns / 1_000_000_000,
333        nsec: ns.rem_euclid(1_000_000_000) as u32,
334    }
335}
336
337#[async_trait]
338impl MetadataStore for SqliteMetadataStore {
339    async fn resolve(&self, path: &str) -> VfsResult<InodeMeta> {
340        self.inner.resolve(path).await
341    }
342
343    async fn resolve_parent(&self, path: &str) -> VfsResult<(InodeMeta, String)> {
344        self.inner.resolve_parent(path).await
345    }
346
347    async fn lstat(&self, path: &str) -> VfsResult<InodeMeta> {
348        self.inner.lstat(path).await
349    }
350
351    async fn list_dir(&self, ino: u64) -> VfsResult<Vec<DentryStat>> {
352        self.inner.list_dir(ino).await
353    }
354
355    async fn create(
356        &self,
357        parent: u64,
358        name: &str,
359        attrs: CreateInodeAttrs,
360    ) -> VfsResult<InodeMeta> {
361        let result = self.inner.create(parent, name, attrs).await;
362        if result.is_ok() {
363            self.persist()?;
364        }
365        result
366    }
367
368    async fn link(&self, parent: u64, name: &str, target: u64) -> VfsResult<()> {
369        let result = self.inner.link(parent, name, target).await;
370        if result.is_ok() {
371            self.persist()?;
372        }
373        result
374    }
375
376    async fn remove(&self, parent: u64, name: &str) -> VfsResult<Vec<BlockKey>> {
377        let result = self.inner.remove(parent, name).await;
378        if result.is_ok() {
379            self.persist()?;
380        }
381        result
382    }
383
384    async fn rename(
385        &self,
386        src_parent: u64,
387        src: &str,
388        dst_parent: u64,
389        dst: &str,
390    ) -> VfsResult<Vec<BlockKey>> {
391        let result = self.inner.rename(src_parent, src, dst_parent, dst).await;
392        if result.is_ok() {
393            self.persist()?;
394        }
395        result
396    }
397
398    async fn set_attr(&self, ino: u64, patch: InodePatch) -> VfsResult<Vec<BlockKey>> {
399        let result = self.inner.set_attr(ino, patch).await;
400        if result.is_ok() {
401            self.persist()?;
402        }
403        result
404    }
405
406    async fn commit_write(
407        &self,
408        ino: u64,
409        edits: Vec<ChunkEdit>,
410        new_size: u64,
411    ) -> VfsResult<Vec<BlockKey>> {
412        let result = self.inner.commit_write(ino, edits, new_size).await;
413        if result.is_ok() {
414            self.persist()?;
415        }
416        result
417    }
418
419    async fn get_chunks(&self, ino: u64, range: ChunkRange) -> VfsResult<Vec<ChunkRef>> {
420        self.inner.get_chunks(ino, range).await
421    }
422
423    async fn snapshot(&self, root: u64) -> VfsResult<SnapshotId> {
424        self.inner.snapshot(root).await
425    }
426
427    async fn fork(&self, snap: SnapshotId) -> VfsResult<u64> {
428        let result = self.inner.fork(snap).await;
429        if result.is_ok() {
430            self.persist()?;
431        }
432        result
433    }
434
435    async fn gc(&self) -> VfsResult<Vec<BlockKey>> {
436        self.inner.gc().await
437    }
438}