Eliminate untracked singleton

Author: rsdy

server/bleep/src/db.rs

@@ -1,36 +1,30 @@
-use std::{fs, path::Path};
+use std::{fs, path::Path, sync::Arc};
 
-use anyhow::{anyhow, Context, Result};
-use once_cell::sync::OnceCell;
+use anyhow::{Context, Result};
 use sqlx::SqlitePool;
 use tracing::{debug, warn};
 
 use crate::Configuration;
 
-static POOL: OnceCell<SqlitePool> = OnceCell::new();
+pub type SqlDb = Arc<SqlitePool>;
 
-pub async fn init(config: &Configuration) -> Result<()> {
+pub async fn init(config: &Configuration) -> Result<SqlitePool> {
     fs::create_dir_all(&config.data_dir)?;
     let data_dir = config.data_dir.to_string_lossy();
 
-    let pool = match connect(&data_dir).await {
-        Ok(pool) => pool,
+    match connect(&data_dir).await {
+        Ok(pool) => Ok(pool),
         Err(e) => {
             warn!(
                 ?e,
                 "encountered DB error while migrating, recreating database..."
             );
             reset(&data_dir)?;
-            connect(&data_dir)
+            Ok(connect(&data_dir)
                 .await
-                .context("failed to recreate database")?
+                .context("failed to recreate database")?)
         }
-    };
-
-    POOL.set(pool)
-        .map_err(|_| anyhow!("database was already initialized!"))?;
-
-    Ok(())
+    }
 }
 
 async fn connect(data_dir: &str) -> Result<SqlitePool> {
@@ -46,8 +40,3 @@ fn reset(data_dir: &str) -> Result<()> {
     let bk_path = db_path.with_extension("db.bk");
     std::fs::rename(db_path, bk_path).context("failed to backup old database")
 }
-
-pub async fn get() -> Result<&'static SqlitePool> {
-    POOL.get()
-        .ok_or(anyhow!("database pool was not initialized"))
-}

server/bleep/src/lib.rs

@@ -13,6 +13,7 @@
 #[cfg(any(bench, test))]
 use criterion as _;
 
+use db::SqlDb;
 #[cfg(any(bench, test))]
 use git_version as _;
 
@@ -91,6 +92,9 @@ pub struct Application {
     /// Main cookie encryption keypair
     cookie_key: axum_extra::extract::cookie::Key,
 
+    /// SQL database for persistent storage
+    sql: SqlDb,
+
     /// Analytics backend -- may be unintialized
     pub analytics: Option<Arc<analytics::RudderHub>>,
 }
@@ -113,7 +117,7 @@ impl Application {
         let config = Arc::new(config);
         debug!(?config, "effective configuration");
 
-        db::init(&config).await?;
+        let sqlite = db::init(&config).await?.into();
 
         // Initialise Semantic index if `qdrant_url` set in config
         let semantic = match config.qdrant_url {
@@ -153,6 +157,7 @@ impl Application {
             sync_queue: SyncQueue::start(config.clone()),
             cookie_key: config.source.initialize_cookie_key()?,
             credentials: config.source.initialize_credentials()?.into(),
+            sql: sqlite,
             repo_pool,
             analytics,
             semantic,

server/bleep/src/webserver/answer.rs

@@ -27,7 +27,7 @@ use tracing::{debug, info, trace, warn};
 use super::middleware::User;
 use crate::{
     analytics::{EventData, QueryEvent},
-    db,
+    db::SqlDb,
     query::parser::{self, SemanticQuery},
     repo::RepoRef,
     Application,
@@ -103,7 +103,7 @@ pub(super) async fn _handle(
         thread_id: params.thread_id,
     };
 
-    let mut conversation = Conversation::load(&conversation_id)
+    let mut conversation = Conversation::load(&app.sql, &conversation_id)
         .await?
         .unwrap_or_else(|| Conversation::new(params.repo_ref.clone()));
 
@@ -200,7 +200,7 @@ pub(super) async fn _handle(
         }
 
         // Storing the conversation here allows us to make subsequent requests.
-        conversation.store(conversation_id).await?;
+        conversation.store(&ctx.app.sql, conversation_id).await?;
     };
 
     let thread_stream = futures::stream::once(async move {
@@ -959,9 +959,8 @@ impl Conversation {
         Ok(())
     }
 
-    async fn store(self, id: ConversationId) -> Result<()> {
+    async fn store(self, db: &SqlDb, id: ConversationId) -> Result<()> {
         info!("writing conversation {}-{}", id.user_id, id.thread_id);
-        let db = db::get().await?;
         let mut transaction = db.begin().await?;
 
         // Delete the old conversation for simplicity. This also deletes all its messages.
@@ -1010,9 +1009,7 @@ impl Conversation {
         Ok(())
     }
 
-    async fn load(id: &ConversationId) -> Result<Option<Self>> {
-        let db = db::get().await?;
-
+    async fn load(db: &SqlDb, id: &ConversationId) -> Result<Option<Self>> {
         let (user_id, thread_id) = (id.user_id.clone(), id.thread_id.to_string());
 
         let row = sqlx::query! {
@@ -1021,7 +1018,7 @@ impl Conversation {
             user_id,
             thread_id,
         }
-        .fetch_optional(db)
+        .fetch_optional(db.as_ref())
         .await?;
 
         let row = match row {

server/bleep/src/webserver/answer/conversations.rs

@@ -1,14 +1,14 @@
 use axum::{
-    extract::{Path, Query},
+    extract::{Path, Query, State},
     response::IntoResponse,
     Extension, Json,
 };
 use reqwest::StatusCode;
 
 use crate::{
-    db,
     repo::RepoRef,
     webserver::{self, middleware::User, Error, ErrorKind},
+    Application,
 };
 
 use super::{Conversation, ConversationId};
@@ -28,9 +28,9 @@ pub(in crate::webserver) struct List {
 pub(in crate::webserver) async fn list(
     Extension(user): Extension<User>,
     Query(query): Query<List>,
+    State(app): State<Application>,
 ) -> webserver::Result<impl IntoResponse> {
-    let db = db::get().await?;
-
+    let db = app.sql.as_ref();
     let user_id = user.0.ok_or_else(|| Error::user("missing user ID"))?;
 
     let conversations = if let Some(repo_ref) = query.repo_ref {
@@ -71,8 +71,9 @@ pub(in crate::webserver) struct Delete {
 pub(in crate::webserver) async fn delete(
     Query(params): Query<Delete>,
     Extension(user): Extension<User>,
+    State(app): State<Application>,
 ) -> webserver::Result<()> {
-    let db = db::get().await?;
+    let db = app.sql.as_ref();
     let user_id = user.0.ok_or_else(|| Error::user("missing user ID"))?;
 
     let result = sqlx::query! {
@@ -94,9 +95,10 @@ pub(in crate::webserver) async fn delete(
 pub(in crate::webserver) async fn thread(
     Path(thread_id): Path<uuid::Uuid>,
     Extension(user): Extension<User>,
+    State(app): State<Application>,
 ) -> webserver::Result<impl IntoResponse> {
     let user_id = user.0.ok_or_else(|| Error::user("missing user ID"))?;
-    let conversation = Conversation::load(&ConversationId { thread_id, user_id })
+    let conversation = Conversation::load(&app.sql, &ConversationId { thread_id, user_id })
         .await?
         .ok_or_else(|| Error::new(ErrorKind::NotFound, "thread was not found"))?;