Integrate Salsa into document tracking for server (#148)

Justfile

@@ -20,9 +20,18 @@ nox SESSION *ARGS:
 bumpver *ARGS:
     uv run --with bumpver bumpver {{ ARGS }}
 
+check *ARGS:
+    cargo check {{ ARGS }}
+
 clean:
     cargo clean
 
+clippy:
+    cargo clippy --all-targets --all-features --fix -- -D warnings
+
+fmt *ARGS:
+    cargo +nightly fmt {{ ARGS }}
+
 # run pre-commit on all files
 lint:
     @just --fmt

crates/djls-server/Cargo.toml

@@ -14,6 +14,7 @@ djls-templates = { workspace = true }
 
 anyhow = { workspace = true }
 pyo3 = { workspace = true }
+salsa = { workspace = true }
 serde = { workspace = true }
 serde_json = { workspace = true }
 tokio = { workspace = true }
@@ -23,3 +24,6 @@ percent-encoding = "2.3"
 
 [build-dependencies]
 djls-dev = { workspace = true }
+
+[lints]
+workspace = true

crates/djls-server/src/db.rs

@@ -0,0 +1,22 @@
+use salsa::Database;
+
+#[salsa::db]
+#[derive(Clone, Default)]
+pub struct ServerDatabase {
+    storage: salsa::Storage<Self>,
+}
+
+impl ServerDatabase {
+    /// Create a new database from storage
+    pub fn new(storage: salsa::Storage<Self>) -> Self {
+        Self { storage }
+    }
+}
+
+impl std::fmt::Debug for ServerDatabase {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_struct("ServerDatabase").finish_non_exhaustive()
+    }
+}
+
+impl Database for ServerDatabase {}

crates/djls-server/src/documents.rs

@@ -3,7 +3,20 @@ use std::collections::HashMap;
 use anyhow::anyhow;
 use anyhow::Result;
 use djls_project::TemplateTags;
-use tower_lsp_server::lsp_types::*;
+use salsa::Database;
+use tower_lsp_server::lsp_types::CompletionItem;
+use tower_lsp_server::lsp_types::CompletionItemKind;
+use tower_lsp_server::lsp_types::CompletionResponse;
+use tower_lsp_server::lsp_types::DidChangeTextDocumentParams;
+use tower_lsp_server::lsp_types::DidCloseTextDocumentParams;
+use tower_lsp_server::lsp_types::DidOpenTextDocumentParams;
+use tower_lsp_server::lsp_types::Documentation;
+use tower_lsp_server::lsp_types::InsertTextFormat;
+use tower_lsp_server::lsp_types::MarkupContent;
+use tower_lsp_server::lsp_types::MarkupKind;
+use tower_lsp_server::lsp_types::Position;
+use tower_lsp_server::lsp_types::Range;
+use tower_lsp_server::lsp_types::TextDocumentContentChangeEvent;
 
 #[derive(Debug, Default)]
 pub struct Store {
@@ -19,54 +32,42 @@ impl Store {
         }
     }
 
-    pub fn handle_did_open(&mut self, params: DidOpenTextDocumentParams) -> Result<()> {
-        let document = TextDocument::new(
-            params.text_document.uri.to_string(),
-            params.text_document.text,
-            params.text_document.version,
-            params.text_document.language_id,
-        );
+    pub fn handle_did_open(&mut self, db: &dyn Database, params: &DidOpenTextDocumentParams) {
+        let uri = params.text_document.uri.to_string();
+        let version = params.text_document.version;
 
-        self.add_document(document);
+        let document = TextDocument::from_did_open_params(db, params);
 
-        Ok(())
+        self.add_document(document, uri.clone());
+        self.versions.insert(uri, version);
     }
 
-    pub fn handle_did_change(&mut self, params: DidChangeTextDocumentParams) -> Result<()> {
+    pub fn handle_did_change(
+        &mut self,
+        db: &dyn Database,
+        params: &DidChangeTextDocumentParams,
+    ) -> Result<()> {
         let uri = params.text_document.uri.as_str().to_string();
         let version = params.text_document.version;
 
         let document = self
-            .get_document_mut(&uri)
+            .get_document(&uri)
             .ok_or_else(|| anyhow!("Document not found: {}", uri))?;
 
-        for change in params.content_changes {
-            if let Some(range) = change.range {
-                document.apply_change(range, &change.text)?;
-            } else {
-                // Full document update
-                document.set_content(change.text);
-            }
-        }
+        let new_document = document.with_changes(db, &params.content_changes, version);
 
-        document.version = version;
+        self.documents.insert(uri.clone(), new_document);
         self.versions.insert(uri, version);
 
         Ok(())
     }
 
-    pub fn handle_did_close(&mut self, params: DidCloseTextDocumentParams) -> Result<()> {
+    pub fn handle_did_close(&mut self, params: &DidCloseTextDocumentParams) {
         self.remove_document(params.text_document.uri.as_str());
-
-        Ok(())
     }
 
-    fn add_document(&mut self, document: TextDocument) {
-        let uri = document.uri.clone();
-        let version = document.version;
-
-        self.documents.insert(uri.clone(), document);
-        self.versions.insert(uri, version);
+    fn add_document(&mut self, document: TextDocument, uri: String) {
+        self.documents.insert(uri, document);
     }
 
     fn remove_document(&mut self, uri: &str) {
@@ -78,6 +79,7 @@ impl Store {
         self.documents.get(uri)
     }
 
+    #[allow(dead_code)]
     fn get_document_mut(&mut self, uri: &str) -> Option<&mut TextDocument> {
         self.documents.get_mut(uri)
     }
@@ -88,13 +90,14 @@ impl Store {
     }
 
     #[allow(dead_code)]
-    pub fn get_documents_by_language(
-        &self,
+    pub fn get_documents_by_language<'db>(
+        &'db self,
+        db: &'db dyn Database,
         language_id: LanguageId,
-    ) -> impl Iterator<Item = &TextDocument> {
+    ) -> impl Iterator<Item = &'db TextDocument> + 'db {
         self.documents
             .values()
-            .filter(move |doc| doc.language_id == language_id)
+            .filter(move |doc| doc.language_id(db) == language_id)
     }
 
     #[allow(dead_code)]
@@ -109,17 +112,18 @@ impl Store {
 
     pub fn get_completions(
         &self,
+        db: &dyn Database,
         uri: &str,
         position: Position,
         tags: &TemplateTags,
     ) -> Option<CompletionResponse> {
         let document = self.get_document(uri)?;
 
-        if document.language_id != LanguageId::HtmlDjango {
+        if document.language_id(db) != LanguageId::HtmlDjango {
             return None;
         }
 
-        let context = document.get_template_tag_context(position)?;
+        let context = document.get_template_tag_context(db, position)?;
 
         let mut completions: Vec<CompletionItem> = tags
             .iter()
@@ -135,7 +139,7 @@ impl Store {
                     documentation: tag.doc().as_ref().map(|doc| {
                         Documentation::MarkupContent(MarkupContent {
                             kind: MarkupKind::Markdown,
-                            value: doc.to_string(),
+                            value: (*doc).to_string(),
                         })
                     }),
                     insert_text: Some(match context.closing_brace {
@@ -158,89 +162,110 @@ impl Store {
     }
 }
 
-#[derive(Clone, Debug)]
+#[salsa::input(debug)]
 pub struct TextDocument {
+    #[return_ref]
     uri: String,
+    #[return_ref]
     contents: String,
+    #[return_ref]
     index: LineIndex,
     version: i32,
     language_id: LanguageId,
 }
 
 impl TextDocument {
-    fn new(uri: String, contents: String, version: i32, language_id: String) -> Self {
+    pub fn from_did_open_params(db: &dyn Database, params: &DidOpenTextDocumentParams) -> Self {
+        let uri = params.text_document.uri.to_string();
+        let contents = params.text_document.text.clone();
+        let version = params.text_document.version;
+        let language_id = LanguageId::from(params.text_document.language_id.as_str());
+
         let index = LineIndex::new(&contents);
-        Self {
-            uri,
-            contents,
-            index,
-            version,
-            language_id: LanguageId::from(language_id),
+        TextDocument::new(db, uri, contents, index, version, language_id)
+    }
+
+    pub fn with_changes(
+        self,
+        db: &dyn Database,
+        changes: &[TextDocumentContentChangeEvent],
+        new_version: i32,
+    ) -> Self {
+        let mut new_contents = self.contents(db).to_string();
+
+        for change in changes {
+            if let Some(range) = change.range {
+                let index = LineIndex::new(&new_contents);
+
+                if let (Some(start_offset), Some(end_offset)) = (
+                    index.offset(range.start).map(|o| o as usize),
+                    index.offset(range.end).map(|o| o as usize),
+                ) {
+                    let mut updated_content = String::with_capacity(
+                        new_contents.len() - (end_offset - start_offset) + change.text.len(),
+                    );
+
+                    updated_content.push_str(&new_contents[..start_offset]);
+                    updated_content.push_str(&change.text);
+                    updated_content.push_str(&new_contents[end_offset..]);
+
+                    new_contents = updated_content;
+                }
+            } else {
+                // Full document update
+                new_contents.clone_from(&change.text);
+            }
         }
-    }
 
-    pub fn apply_change(&mut self, range: Range, new_text: &str) -> Result<()> {
-        let start_offset = self
-            .index
-            .offset(range.start)
-            .ok_or_else(|| anyhow!("Invalid start position: {:?}", range.start))?
-            as usize;
-        let end_offset = self
-            .index
-            .offset(range.end)
-            .ok_or_else(|| anyhow!("Invalid end position: {:?}", range.end))?
-            as usize;
-
-        let mut new_content = String::with_capacity(
-            self.contents.len() - (end_offset - start_offset) + new_text.len(),
-        );
-
-        new_content.push_str(&self.contents[..start_offset]);
-        new_content.push_str(new_text);
-        new_content.push_str(&self.contents[end_offset..]);
-
-        self.set_content(new_content);
-
-        Ok(())
-    }
-
-    pub fn set_content(&mut self, new_content: String) {
-        self.contents = new_content;
-        self.index = LineIndex::new(&self.contents);
+        let index = LineIndex::new(&new_contents);
+        TextDocument::new(
+            db,
+            self.uri(db).to_string(),
+            new_contents,
+            index,
+            new_version,
+            self.language_id(db),
+        )
     }
 
     #[allow(dead_code)]
-    pub fn get_text(&self) -> &str {
-        &self.contents
+    pub fn get_text(self, db: &dyn Database) -> String {
+        self.contents(db).to_string()
     }
 
     #[allow(dead_code)]
-    pub fn get_text_range(&self, range: Range) -> Option<&str> {
-        let start = self.index.offset(range.start)? as usize;
-        let end = self.index.offset(range.end)? as usize;
-
-        Some(&self.contents[start..end])
+    pub fn get_text_range(self, db: &dyn Database, range: Range) -> Option<String> {
+        let index = self.index(db);
+        let start = index.offset(range.start)? as usize;
+        let end = index.offset(range.end)? as usize;
+        let contents = self.contents(db);
+        Some(contents[start..end].to_string())
     }
 
-    pub fn get_line(&self, line: u32) -> Option<&str> {
-        let start = self.index.line_starts.get(line as usize)?;
-        let end = self
-            .index
+    pub fn get_line(self, db: &dyn Database, line: u32) -> Option<String> {
+        let index = self.index(db);
+        let start = index.line_starts.get(line as usize)?;
+        let end = index
             .line_starts
             .get(line as usize + 1)
             .copied()
-            .unwrap_or(self.index.length);
+            .unwrap_or(index.length);
 
-        Some(&self.contents[*start as usize..end as usize])
+        let contents = self.contents(db);
+        Some(contents[*start as usize..end as usize].to_string())
     }
 
     #[allow(dead_code)]
-    pub fn line_count(&self) -> usize {
-        self.index.line_starts.len()
+    pub fn line_count(self, db: &dyn Database) -> usize {
+        self.index(db).line_starts.len()
     }
 
-    pub fn get_template_tag_context(&self, position: Position) -> Option<TemplateTagContext> {
-        let line = self.get_line(position.line)?;
+    pub fn get_template_tag_context(
+        self,
+        db: &dyn Database,
+        position: Position,
+    ) -> Option<TemplateTagContext> {
+        let line = self.get_line(db, position.line)?;
         let char_pos: usize = position.character.try_into().ok()?;
         let prefix = &line[..char_pos];
         let rest_of_line = &line[char_pos..];
@@ -252,7 +277,7 @@ impl TextDocument {
 
             let closing_brace = if rest_trimmed.starts_with("%}") {
                 ClosingBrace::FullClose
-            } else if rest_trimmed.starts_with("}") {
+            } else if rest_trimmed.starts_with('}') {
                 ClosingBrace::PartialClose
             } else {
                 ClosingBrace::None
@@ -279,7 +304,7 @@ impl LineIndex {
         let mut pos = 0;
 
         for c in text.chars() {
-            pos += c.len_utf8() as u32;
+            pos += u32::try_from(c.len_utf8()).unwrap_or(0);
             if c == '\n' {
                 line_starts.push(pos);
             }
@@ -307,7 +332,7 @@ impl LineIndex {
         let line_start = self.line_starts[line];
         let character = offset - line_start;
 
-        Position::new(line as u32, character)
+        Position::new(u32::try_from(line).unwrap_or(0), character)
     }
 }
 

crates/djls-server/src/lib.rs

@@ -1,3 +1,4 @@
+mod db;
 mod documents;
 mod queue;
 mod server;

crates/djls-server/src/queue.rs

@@ -98,7 +98,7 @@ impl Queue {
                             if let Err(e) = task_future.await {
                                 // Log the error if the task failed.
                                 // TODO: Integrate with a proper logging framework.
-                                eprintln!("Task failed: {}", e);
+                                eprintln!("Task failed: {e}");
                             }
                         } else {
                             // Channel closed, implies all senders (Queue instances)
@@ -209,7 +209,7 @@ mod tests {
             queue
                 .submit(async move {
                     sleep(Duration::from_millis(5)).await;
-                    println!("Processing task {}", i);
+                    println!("Processing task {i}");
                     counter_clone.fetch_add(1, Ordering::SeqCst);
                     Ok(())
                 })
@@ -262,7 +262,7 @@ mod tests {
                     })
                     .await
                     .expect("Submit should succeed");
-                println!("Submitted task {}", i);
+                println!("Submitted task {i}");
             }));
         }
         for task in tasks {
@@ -278,11 +278,13 @@ mod tests {
         });
 
         match tokio::time::timeout(Duration::from_millis(500), submit_task).await {
-            Ok(Ok(_)) => {
+            Ok(Ok(())) => {
                 println!("Successfully submitted 33rd task");
             }
-            Ok(Err(e)) => panic!("Submit failed unexpectedly: {}", e),
-            Err(_) => panic!("Submit timed out, likely blocked due to backpressure not resolving"),
+            Ok(Err(e)) => panic!("Submit failed unexpectedly: {e}"),
+            Err(timeout_err) => panic!(
+                "Submit timed out: {timeout_err}, likely blocked due to backpressure not resolving"
+            ),
         }
 
         sleep(Duration::from_millis(200)).await;
@@ -318,7 +320,7 @@ mod tests {
         sleep(Duration::from_millis(200)).await;
 
         let final_count = counter.load(Ordering::SeqCst);
-        println!("Final count after shutdown: {}", final_count);
+        println!("Final count after shutdown: {final_count}");
         assert!(final_count <= 2);
     }
 

crates/djls-server/src/server.rs

@@ -2,7 +2,26 @@ use std::sync::Arc;
 
 use tokio::sync::RwLock;
 use tower_lsp_server::jsonrpc::Result as LspResult;
-use tower_lsp_server::lsp_types::*;
+use tower_lsp_server::lsp_types::CompletionOptions;
+use tower_lsp_server::lsp_types::CompletionParams;
+use tower_lsp_server::lsp_types::CompletionResponse;
+use tower_lsp_server::lsp_types::DidChangeConfigurationParams;
+use tower_lsp_server::lsp_types::DidChangeTextDocumentParams;
+use tower_lsp_server::lsp_types::DidCloseTextDocumentParams;
+use tower_lsp_server::lsp_types::DidOpenTextDocumentParams;
+use tower_lsp_server::lsp_types::InitializeParams;
+use tower_lsp_server::lsp_types::InitializeResult;
+use tower_lsp_server::lsp_types::InitializedParams;
+use tower_lsp_server::lsp_types::MessageType;
+use tower_lsp_server::lsp_types::OneOf;
+use tower_lsp_server::lsp_types::SaveOptions;
+use tower_lsp_server::lsp_types::ServerCapabilities;
+use tower_lsp_server::lsp_types::ServerInfo;
+use tower_lsp_server::lsp_types::TextDocumentSyncCapability;
+use tower_lsp_server::lsp_types::TextDocumentSyncKind;
+use tower_lsp_server::lsp_types::TextDocumentSyncOptions;
+use tower_lsp_server::lsp_types::WorkspaceFoldersServerCapabilities;
+use tower_lsp_server::lsp_types::WorkspaceServerCapabilities;
 use tower_lsp_server::Client;
 use tower_lsp_server::LanguageServer;
 
@@ -19,6 +38,7 @@ pub struct DjangoLanguageServer {
 }
 
 impl DjangoLanguageServer {
+    #[must_use]
     pub fn new(client: Client) -> Self {
         Self {
             client,
@@ -86,6 +106,7 @@ impl LanguageServer for DjangoLanguageServer {
         })
     }
 
+    #[allow(clippy::too_many_lines)]
     async fn initialized(&self, _params: InitializedParams) {
         self.client
             .log_message(
@@ -143,7 +164,7 @@ impl LanguageServer for DjangoLanguageServer {
                     session.project().map(|p| {
                         (
                             p.path().display().to_string(),
-                            session.settings().venv_path().map(|s| s.to_string()),
+                            session.settings().venv_path().map(std::string::ToString::to_string),
                         )
                     })
                 };
@@ -153,8 +174,7 @@ impl LanguageServer for DjangoLanguageServer {
                         .log_message(
                             MessageType::INFO,
                             &format!(
-                                "Task: Starting initialization for project at: {}",
-                                path_display
+                                "Task: Starting initialization for project at: {path_display}"
                             ),
                         )
                         .await;
@@ -163,7 +183,7 @@ impl LanguageServer for DjangoLanguageServer {
                         client
                             .log_message(
                                 MessageType::INFO,
-                                &format!("Using virtual environment from config: {}", path),
+                                &format!("Using virtual environment from config: {path}"),
                             )
                             .await;
                     }
@@ -184,8 +204,7 @@ impl LanguageServer for DjangoLanguageServer {
                                 .log_message(
                                     MessageType::INFO,
                                     &format!(
-                                        "Task: Successfully initialized project: {}",
-                                        path_display
+                                        "Task: Successfully initialized project: {path_display}"
                                     ),
                                 )
                                 .await;
@@ -195,8 +214,7 @@ impl LanguageServer for DjangoLanguageServer {
                                 .log_message(
                                     MessageType::ERROR,
                                     &format!(
-                                        "Task: Failed to initialize Django project at {}: {}",
-                                        path_display, e
+                                        "Task: Failed to initialize Django project at {path_display}: {e}"
                                     ),
                                 )
                                 .await;
@@ -221,7 +239,7 @@ impl LanguageServer for DjangoLanguageServer {
             self.client
                 .log_message(
                     MessageType::ERROR,
-                    &format!("Failed to submit project initialization task: {}", e),
+                    &format!("Failed to submit project initialization task: {e}"),
                 )
                 .await;
         } else {
@@ -236,57 +254,47 @@ impl LanguageServer for DjangoLanguageServer {
     }
 
     async fn did_open(&self, params: DidOpenTextDocumentParams) {
-        let result = self
-            .with_session_mut(|session| session.documents_mut().handle_did_open(params.clone()))
-            .await;
-
-        if let Err(e) = result {
-            eprintln!("Error handling document open: {}", e);
-            return;
-        }
-
         self.client
             .log_message(
                 MessageType::INFO,
                 &format!("Opened document: {:?}", params.text_document.uri),
             )
             .await;
+
+        self.with_session_mut(|session| {
+            let db = session.db();
+            session.documents_mut().handle_did_open(&db, &params);
+        })
+        .await;
     }
 
     async fn did_change(&self, params: DidChangeTextDocumentParams) {
-        let result = self
-            .with_session_mut(|session| session.documents_mut().handle_did_change(params.clone()))
-            .await;
-
-        if let Err(e) = result {
-            eprintln!("Error handling document change: {}", e);
-            return;
-        }
-
         self.client
             .log_message(
                 MessageType::INFO,
                 &format!("Changed document: {:?}", params.text_document.uri),
             )
             .await;
+
+        self.with_session_mut(|session| {
+            let db = session.db();
+            let _ = session.documents_mut().handle_did_change(&db, &params);
+        })
+        .await;
     }
 
     async fn did_close(&self, params: DidCloseTextDocumentParams) {
-        let result = self
-            .with_session_mut(|session| session.documents_mut().handle_did_close(params.clone()))
-            .await;
-
-        if let Err(e) = result {
-            eprintln!("Error handling document close: {}", e);
-            return;
-        }
-
         self.client
             .log_message(
                 MessageType::INFO,
                 &format!("Closed document: {:?}", params.text_document.uri),
             )
             .await;
+
+        self.with_session_mut(|session| {
+            session.documents_mut().handle_did_close(&params);
+        })
+        .await;
     }
 
     async fn completion(&self, params: CompletionParams) -> LspResult<Option<CompletionResponse>> {
@@ -294,7 +302,9 @@ impl LanguageServer for DjangoLanguageServer {
             .with_session(|session| {
                 if let Some(project) = session.project() {
                     if let Some(tags) = project.template_tags() {
+                        let db = session.db();
                         return session.documents().get_completions(
+                            &db,
                             params.text_document_position.text_document.uri.as_str(),
                             params.text_document_position.position,
                             tags,
@@ -324,7 +334,7 @@ impl LanguageServer for DjangoLanguageServer {
                     *session.settings_mut() = new_settings;
                 }
                 Err(e) => {
-                    eprintln!("Error loading settings: {}", e);
+                    eprintln!("Error loading settings: {e}");
                 }
             })
             .await;

crates/djls-server/src/session.rs

@@ -1,15 +1,43 @@
 use djls_conf::Settings;
 use djls_project::DjangoProject;
+use salsa::StorageHandle;
 use tower_lsp_server::lsp_types::ClientCapabilities;
 
+use crate::db::ServerDatabase;
 use crate::documents::Store;
 
-#[derive(Debug, Default)]
+#[derive(Default)]
 pub struct Session {
     client_capabilities: Option<ClientCapabilities>,
     project: Option<DjangoProject>,
     documents: Store,
     settings: Settings,
+
+    /// A thread-safe Salsa database handle that can be shared between threads.
+    ///
+    /// This implements the insight from [this Salsa Zulip discussion](https://salsa.zulipchat.com/#narrow/channel/145099-Using-Salsa/topic/.E2.9C.94.20Advice.20on.20using.20salsa.20from.20Sync.20.2B.20Send.20context/with/495497515)
+    /// where we're using the `StorageHandle` to create a thread-safe handle that can be
+    /// shared between threads. When we need to use it, we clone the handle to get a new reference.
+    ///
+    /// Usage:
+    /// ```rust,ignore
+    /// // Use the StorageHandle in Session
+    /// let db_handle = StorageHandle::new(None);
+    ///
+    /// // Clone it to pass to different threads
+    /// let db_handle_clone = db_handle.clone();
+    ///
+    /// // Use it in an async context
+    /// async_fn(move || {
+    ///     // Get a database from the handle
+    ///     let storage = db_handle_clone.into_storage();
+    ///     let db = ServerDatabase::new(storage);
+    ///
+    ///     // Use the database
+    ///     db.some_query(args)
+    /// });
+    /// ```
+    db_handle: StorageHandle<ServerDatabase>,
 }
 
 impl Session {
@@ -19,11 +47,12 @@ impl Session {
             project: None,
             documents: Store::new(),
             settings: Settings::default(),
+            db_handle: StorageHandle::new(None),
         }
     }
 
-    pub fn client_capabilities(&self) -> &Option<ClientCapabilities> {
-        &self.client_capabilities
+    pub fn client_capabilities(&self) -> Option<&ClientCapabilities> {
+        self.client_capabilities.as_ref()
     }
 
     pub fn client_capabilities_mut(&mut self) -> &mut Option<ClientCapabilities> {
@@ -53,4 +82,21 @@ impl Session {
     pub fn settings_mut(&mut self) -> &mut Settings {
         &mut self.settings
     }
+
+    /// Get the raw database handle from the session
+    ///
+    /// Note: In most cases, you'll want to use `db()` instead to get a usable
+    /// database instance directly.
+    pub fn db_handle(&self) -> &StorageHandle<ServerDatabase> {
+        &self.db_handle
+    }
+
+    /// Get a database instance directly from the session
+    ///
+    /// This creates a usable database from the handle, which can be used
+    /// to query and update data in the database.
+    pub fn db(&self) -> ServerDatabase {
+        let storage = self.db_handle.clone().into_storage();
+        ServerDatabase::new(storage)
+    }
 }