Skip to main content

blvm_sdk/module/
database.rs

1//! Module database utilities
2//!
3//! Each module has its own database at `{data_dir}/db/`. Use this helper to open it.
4//! Modules use the same format as the node by default (via MODULE_CONFIG_DATABASE_BACKEND).
5
6use anyhow::Result;
7use blvm_node::storage::database::{create_database, default_backend, Database, DatabaseBackend};
8use std::path::Path;
9use std::sync::Arc;
10use tracing::warn;
11
12fn parse_backend(s: &str) -> DatabaseBackend {
13    match s.to_lowercase().as_str() {
14        "redb" => DatabaseBackend::Redb,
15        "rocksdb" => DatabaseBackend::RocksDB,
16        "sled" => DatabaseBackend::Sled,
17        "tidesdb" => DatabaseBackend::TidesDB,
18        // "auto" and anything unrecognised: let the caller fall through to the default selection
19        _ => default_backend(),
20    }
21}
22
23/// Returns true for backends that support fully dynamic `open_tree()` without any pre-declaration.
24///
25/// - **Sled**: any tree name is created on demand.
26/// - **TidesDB**: uses `get_or_create_cf`, fully dynamic.
27/// - **Redb**: only works for a hard-coded set of node/module tables (`schema`, `items`, etc.).
28///   Arbitrary tree names return an error — use Sled or TidesDB for general module storage.
29/// - **RocksDB**: column families must be declared at database open time.
30fn supports_dynamic_trees(backend: DatabaseBackend) -> bool {
31    matches!(backend, DatabaseBackend::Sled | DatabaseBackend::TidesDB)
32}
33
34/// Best available backend for module databases (prefers fully-dynamic backends).
35///
36/// Returns Sled as the preferred target; `create_database` will surface an error if the
37/// feature is not compiled into blvm-node, and the caller falls through to TidesDB.
38fn best_module_backend() -> DatabaseBackend {
39    // Prefer Sled: lightweight, embedded, fully dynamic open_tree().
40    // If Sled is not compiled in, create_database will return an error and the caller
41    // falls back to TidesDB. The node default (rocksdb) is intentionally NOT used here.
42    DatabaseBackend::Sled
43}
44
45/// Open the module's database at `{data_dir}/db/`.
46///
47/// Module databases need to create named trees on demand (`open_tree()`). Only **Sled** and
48/// **TidesDB** support fully dynamic tree creation. Redb only supports a fixed set of
49/// pre-declared tables; RocksDB requires all column families to be declared at open time.
50///
51/// Backend selection (first match wins):
52/// 1. `MODULE_CONFIG_DATABASE_BACKEND` env var (`sled`, `tidesdb`, `redb`, `rocksdb`, `auto`)
53/// 2. `MODULE_DATABASE_BACKEND` env var (same values)
54/// 3. Auto-select best available: Sled → TidesDB → fallback with a warning
55///
56/// If the resolved backend does not support dynamic trees the function warns and falls back
57/// to the best available dynamic backend. If no dynamic backend is compiled in, it returns
58/// an error.
59///
60/// # Example
61/// ```ignore
62/// let db = open_module_db(module_data_dir)?;
63/// let tree = db.open_tree("my_tree")?;
64/// tree.insert(b"key", b"value")?;
65/// ```
66pub fn open_module_db<P: AsRef<Path>>(data_dir: P) -> Result<Arc<dyn Database>> {
67    let db_path = data_dir.as_ref().join("db");
68    std::fs::create_dir_all(&db_path)?;
69
70    let explicit = std::env::var("MODULE_CONFIG_DATABASE_BACKEND")
71        .or_else(|_| std::env::var("MODULE_DATABASE_BACKEND"))
72        .ok();
73
74    let backend = explicit
75        .as_deref()
76        .map(parse_backend)
77        .unwrap_or_else(best_module_backend);
78
79    if supports_dynamic_trees(backend) {
80        // Explicitly requested or auto-selected a good module backend.
81        return create_database(&db_path, backend, None).map(Arc::from);
82    }
83
84    // The resolved backend does not support fully-dynamic open_tree().
85    if explicit.is_some() {
86        // User explicitly asked for this backend — warn but still try to fall back so
87        // existing module processes don't hard-fail on start-up.
88        warn!(
89            "MODULE_DATABASE_BACKEND={:?} does not support dynamic open_tree(). \
90             Module databases require Sled or TidesDB. \
91             Set MODULE_CONFIG_DATABASE_BACKEND=sled (or tidesdb) to remove this warning. \
92             Note: Redb only supports pre-declared tables (schema, items); \
93             RocksDB requires all column families at open time.",
94            backend
95        );
96    } else {
97        // Auto-detected a static backend (e.g. rocksdb is the node default).
98        warn!(
99            "Default backend {:?} does not support dynamic open_tree(); \
100             auto-selecting best available module backend (sled or tidesdb). \
101             Set MODULE_CONFIG_DATABASE_BACKEND=sled to suppress this warning.",
102            backend
103        );
104    }
105
106    // Try Sled then TidesDB as fallbacks.
107    create_database(&db_path, DatabaseBackend::Sled, None)
108        .or_else(|_| create_database(&db_path, DatabaseBackend::TidesDB, None))
109        .map(Arc::from)
110        .map_err(|_| {
111            anyhow::anyhow!(
112                "No module-compatible database backend available. \
113                 Backend {:?} requires pre-declared trees/column families. \
114                 Recompile with the 'sled' or 'tidesdb' feature on blvm-node, \
115                 or set MODULE_CONFIG_DATABASE_BACKEND=sled.",
116                backend
117            )
118        })
119}
120
121/// Schema version key (stored in the schema tree).
122const SCHEMA_VERSION_KEY: &[u8] = b"schema_version";
123
124/// Context passed to migration functions. Provides put/get/delete against the module's schema tree
125/// and access to the database for opening other trees.
126///
127/// Migrations run locally in the module process. Use `put`/`get`/`delete` for schema metadata.
128/// For application data migrations, use `open_tree` to open and migrate other trees.
129#[derive(Clone)]
130pub struct MigrationContext {
131    tree: Arc<dyn blvm_node::storage::database::Tree>,
132    db: Arc<dyn Database>,
133}
134
135impl MigrationContext {
136    /// Create a new MigrationContext wrapping the schema tree and database.
137    pub fn new(tree: Arc<dyn blvm_node::storage::database::Tree>, db: Arc<dyn Database>) -> Self {
138        Self { tree, db }
139    }
140
141    /// Insert a key-value pair into the schema tree.
142    pub fn put(&self, key: &[u8], value: &[u8]) -> Result<()> {
143        self.tree.insert(key, value)
144    }
145
146    /// Get a value by key from the schema tree.
147    pub fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>> {
148        self.tree.get(key)
149    }
150
151    /// Remove a key from the schema tree.
152    pub fn delete(&self, key: &[u8]) -> Result<()> {
153        self.tree.remove(key)
154    }
155
156    /// Open a named tree for application data migrations.
157    pub fn open_tree(&self, name: &str) -> Result<Box<dyn blvm_node::storage::database::Tree>> {
158        self.db.open_tree(name)
159    }
160}
161
162/// A single up migration step.
163pub type MigrationUp = fn(&MigrationContext) -> Result<()>;
164
165/// A single down migration step (for rollback).
166pub type MigrationDown = fn(&MigrationContext) -> Result<()>;
167
168/// Migration pair: (version, up, optional down for rollback).
169pub type Migration = (u32, MigrationUp, Option<MigrationDown>);
170
171/// Run pending up migrations. Opens the "schema" tree, reads current version, runs each migration
172/// with version > current in order, then updates schema_version.
173///
174/// # Example
175/// ```ignore
176/// let db = open_module_db(data_dir)?;
177/// run_migrations(&db, &[(1, up_initial, Some(down_initial)), (2, up_add_cache, None)])?;
178/// ```
179pub fn run_migrations(db: &Arc<dyn Database>, migrations: &[(u32, MigrationUp)]) -> Result<()> {
180    run_migrations_with_down(
181        db,
182        &migrations
183            .iter()
184            .map(|(v, u)| (*v, *u, None))
185            .collect::<Vec<_>>(),
186    )
187}
188
189/// Run pending up migrations. Supports optional down migrations for rollback.
190pub fn run_migrations_with_down(db: &Arc<dyn Database>, migrations: &[Migration]) -> Result<()> {
191    let tree = db.open_tree("schema")?;
192    let tree = Arc::from(tree);
193    let ctx = MigrationContext::new(tree, Arc::clone(db));
194
195    let current: u32 = ctx
196        .get(SCHEMA_VERSION_KEY)?
197        .and_then(|v| String::from_utf8(v).ok())
198        .and_then(|s| s.parse().ok())
199        .unwrap_or(0);
200
201    let mut pending: Vec<_> = migrations
202        .iter()
203        .filter(|(v, _, _)| *v > current)
204        .copied()
205        .collect();
206    pending.sort_by_key(|(v, _, _)| *v);
207
208    for (version, up, _down) in pending {
209        up(&ctx)?;
210        ctx.put(SCHEMA_VERSION_KEY, version.to_string().as_bytes())?;
211    }
212
213    Ok(())
214}
215
216/// Rollback migrations down to `target_version` (exclusive). Runs down migrations in reverse
217/// order for each applied version > target_version. Requires down functions to be provided.
218///
219/// # Example
220/// ```ignore
221/// run_migrations_down(&db, &[(1, up_initial, Some(down_initial)), (2, up_add_cache, Some(down_cache))], 1)?;
222/// // Rolls back from 2 to 1 (runs down_cache only).
223/// ```
224pub fn run_migrations_down(
225    db: &Arc<dyn Database>,
226    migrations: &[Migration],
227    target_version: u32,
228) -> Result<()> {
229    let tree = db.open_tree("schema")?;
230    let tree = Arc::from(tree);
231    let ctx = MigrationContext::new(tree, Arc::clone(db));
232
233    let current: u32 = ctx
234        .get(SCHEMA_VERSION_KEY)?
235        .and_then(|v| String::from_utf8(v).ok())
236        .and_then(|s| s.parse().ok())
237        .unwrap_or(0);
238
239    if current <= target_version {
240        return Ok(());
241    }
242
243    let mut to_rollback: Vec<_> = migrations
244        .iter()
245        .filter(|(v, _, d)| *v > target_version && *v <= current && d.is_some())
246        .copied()
247        .collect();
248    to_rollback.sort_by_key(|(v, _, _)| std::cmp::Reverse(*v));
249
250    for (version, _up, down) in to_rollback {
251        if let Some(down_fn) = down {
252            down_fn(&ctx)?;
253        } else {
254            anyhow::bail!("Migration version {} has no down function", version);
255        }
256    }
257
258    ctx.put(SCHEMA_VERSION_KEY, target_version.to_string().as_bytes())?;
259    Ok(())
260}