Clean up DB refactor access (#192)

src-tauri/yaak-models/src/commands.rs

@@ -1,7 +1,7 @@
 use crate::error::Result;
-use crate::manager::QueryManagerExt;
+use crate::query_manager::QueryManagerExt;
 use crate::models::AnyModel;
-use crate::queries_legacy::UpdateSource;
+use crate::util::UpdateSource;
 use tauri::{Runtime, WebviewWindow};
 
 #[tauri::command]
@@ -9,9 +9,8 @@ pub(crate) async fn upsert<R: Runtime>(
     window: WebviewWindow<R>,
     model: AnyModel,
 ) -> Result<String> {
-    let queries = window.queries().connect().await?;
     let id = match model {
-        AnyModel::HttpRequest(r) => queries.upsert(&r, &UpdateSource::from_window(&window))?.id,
+        AnyModel::HttpRequest(r) => window.db().upsert(&r, &UpdateSource::from_window(&window))?.id,
         _ => todo!(),
     };
 

src-tauri/yaak-models/src/connection_or_tx.rs

@@ -0,0 +1,25 @@
+use r2d2::PooledConnection;
+use r2d2_sqlite::SqliteConnectionManager;
+use rusqlite::{Connection, Statement, ToSql, Transaction};
+
+pub enum ConnectionOrTx<'a> {
+    Connection(PooledConnection<SqliteConnectionManager>),
+    Transaction(&'a Transaction<'a>),
+}
+
+impl<'a> ConnectionOrTx<'a> {
+    pub(crate) fn resolve(&self) -> &Connection {
+        match self {
+            ConnectionOrTx::Connection(c) => c,
+            ConnectionOrTx::Transaction(c) => c,
+        }
+    }
+
+    pub(crate) fn prepare(&self, sql: &str) -> rusqlite::Result<Statement<'_>> {
+        self.resolve().prepare(sql)
+    }
+
+    pub(crate) fn execute(&self, sql: &str, params: &[&dyn ToSql]) -> rusqlite::Result<usize> {
+        self.resolve().execute(sql, params)
+    }
+}

src-tauri/yaak-models/src/db_context.rs

@@ -1,38 +1,40 @@
+use crate::connection_or_tx::ConnectionOrTx;
 use crate::error::Error::RowNotFound;
-use crate::error::Result;
-use crate::manager::DbContext;
 use crate::models::{AnyModel, ModelType, UpsertModelInfo};
-use crate::queries_legacy::{generate_model_id, ModelChangeEvent, ModelPayload, UpdateSource};
+use crate::util::{generate_model_id, ModelChangeEvent, ModelPayload, UpdateSource};
 use rusqlite::OptionalExtension;
 use sea_query::{
     Asterisk, Expr, IntoColumnRef, IntoIden, IntoTableRef, OnConflict, Query, SimpleExpr,
     SqliteQueryBuilder,
 };
 use sea_query_rusqlite::RusqliteBinder;
+use tokio::sync::mpsc;
 
-pub(crate) const MAX_HISTORY_ITEMS: usize = 20;
+pub struct DbContext<'a> {
+    pub(crate) tx: mpsc::Sender<ModelPayload>,
+    pub(crate) conn: ConnectionOrTx<'a>,
+}
 
 impl<'a> DbContext<'a> {
     pub(crate) fn find_one<'s, M>(
         &self,
         col: impl IntoColumnRef,
         value: impl Into<SimpleExpr>,
-    ) -> Result<M>
+    ) -> crate::error::Result<M>
     where
         M: Into<AnyModel> + Clone + UpsertModelInfo,
     {
         match self.find_optional::<M>(col, value) {
-            Ok(Some(v)) => Ok(v),
-            Ok(None) => Err(RowNotFound),
-            Err(e) => Err(e),
+            Some(v) => Ok(v),
+            None => Err(RowNotFound),
         }
     }
 
-    pub fn find_optional<'s, M>(
+    pub(crate) fn find_optional<'s, M>(
         &self,
         col: impl IntoColumnRef,
         value: impl Into<SimpleExpr>,
-    ) -> Result<Option<M>>
+    ) -> Option<M>
     where
         M: Into<AnyModel> + Clone + UpsertModelInfo,
     {
@@ -41,11 +43,13 @@ impl<'a> DbContext<'a> {
             .column(Asterisk)
             .cond_where(Expr::col(col).eq(value))
             .build_rusqlite(SqliteQueryBuilder);
-        let mut stmt = self.conn.prepare(sql.as_str())?;
-        Ok(stmt.query_row(&*params.as_params(), M::from_row).optional()?)
+        let mut stmt = self.conn.prepare(sql.as_str()).expect("Failed to prepare query");
+        stmt.query_row(&*params.as_params(), M::from_row)
+            .optional()
+            .expect("Failed to run find on DB")
     }
 
-    pub fn find_all<'s, M>(&self) -> Result<Vec<M>>
+    pub fn find_all<'s, M>(&self) -> crate::error::Result<Vec<M>>
     where
         M: Into<AnyModel> + Clone + UpsertModelInfo,
     {
@@ -63,7 +67,7 @@ impl<'a> DbContext<'a> {
         col: impl IntoColumnRef,
         value: impl Into<SimpleExpr>,
         limit: Option<u64>,
-    ) -> Result<Vec<M>>
+    ) -> crate::error::Result<Vec<M>>
     where
         M: Into<AnyModel> + Clone + UpsertModelInfo,
     {
@@ -88,7 +92,7 @@ impl<'a> DbContext<'a> {
         Ok(items.map(|v| v.unwrap()).collect())
     }
 
-    pub fn upsert<M>(&self, model: &M, source: &UpdateSource) -> Result<M>
+    pub fn upsert<M>(&self, model: &M, source: &UpdateSource) -> crate::error::Result<M>
     where
         M: Into<AnyModel> + From<AnyModel> + UpsertModelInfo + Clone,
     {
@@ -112,7 +116,7 @@ impl<'a> DbContext<'a> {
         other_values: Vec<(impl IntoIden + Eq, impl Into<SimpleExpr>)>,
         update_columns: Vec<impl IntoIden>,
         source: &UpdateSource,
-    ) -> Result<M>
+    ) -> crate::error::Result<M>
     where
         M: Into<AnyModel> + From<AnyModel> + UpsertModelInfo + Clone,
     {
@@ -147,7 +151,11 @@ impl<'a> DbContext<'a> {
         Ok(m)
     }
 
-    pub(crate) fn delete<'s, M>(&self, m: &M, update_source: &UpdateSource) -> Result<M>
+    pub(crate) fn delete<'s, M>(
+        &self,
+        m: &M,
+        update_source: &UpdateSource,
+    ) -> crate::error::Result<M>
     where
         M: Into<AnyModel> + Clone + UpsertModelInfo,
     {

src-tauri/yaak-models/src/lib.rs

@@ -1,6 +1,6 @@
 use crate::commands::{delete, upsert};
-use crate::manager::QueryManager;
-use crate::queries_legacy::ModelChangeEvent;
+use crate::query_manager::QueryManager;
+use crate::util::ModelChangeEvent;
 use log::info;
 use r2d2::Pool;
 use r2d2_sqlite::SqliteConnectionManager;
@@ -19,11 +19,13 @@ use tokio::sync::mpsc;
 
 mod commands;
 
+mod connection_or_tx;
+mod db_context;
 pub mod error;
-pub mod manager;
+pub mod query_manager;
 pub mod models;
 pub mod queries;
-pub mod queries_legacy;
+pub mod util;
 pub mod render;
 
 pub struct SqliteConnection(pub Mutex<Pool<SqliteConnectionManager>>);

src-tauri/yaak-models/src/manager.rs

@@ -1,108 +0,0 @@
-use crate::error::Result;
-use crate::queries_legacy::ModelPayload;
-use r2d2::{Pool, PooledConnection};
-use r2d2_sqlite::SqliteConnectionManager;
-use rusqlite::{Connection, Statement, ToSql, Transaction, TransactionBehavior};
-use std::sync::Arc;
-use tauri::{Manager, Runtime};
-use tokio::sync::{mpsc, Mutex};
-
-pub trait QueryManagerExt<'a, R> {
-    fn queries(&'a self) -> &'a QueryManager;
-}
-
-impl<'a, R: Runtime, T: Manager<R>> QueryManagerExt<'a, R> for T {
-    fn queries(&'a self) -> &'a QueryManager {
-        let qm = self.state::<QueryManager>();
-        qm.inner()
-    }
-}
-
-#[derive(Clone)]
-pub struct QueryManager {
-    pool: Arc<Mutex<Pool<SqliteConnectionManager>>>,
-    events_tx: mpsc::Sender<ModelPayload>,
-}
-
-impl QueryManager {
-    pub(crate) fn new(
-        pool: Pool<SqliteConnectionManager>,
-        events_tx: mpsc::Sender<ModelPayload>,
-    ) -> Self {
-        QueryManager {
-            pool: Arc::new(Mutex::new(pool)),
-            events_tx,
-        }
-    }
-
-    pub async fn connect(&self) -> Result<DbContext> {
-        let conn = self.pool.lock().await.get()?;
-        Ok(DbContext {
-            tx: self.events_tx.clone(),
-            conn: ConnectionOrTx::Connection(conn),
-        })
-    }
-
-    pub async fn with_conn<F, T>(&self, func: F) -> Result<T>
-    where
-        F: FnOnce(&DbContext) -> Result<T>,
-    {
-        let conn = self.pool.lock().await.get()?;
-        let db_context = DbContext {
-            tx: self.events_tx.clone(),
-            conn: ConnectionOrTx::Connection(conn),
-        };
-        func(&db_context)
-    }
-
-    pub async fn with_tx<F, T>(&self, func: F) -> Result<T>
-    where
-        F: FnOnce(&DbContext) -> Result<T>,
-    {
-        let mut conn = self.pool.lock().await.get()?;
-        let tx = conn.transaction_with_behavior(TransactionBehavior::Immediate)?;
-
-        let db_context = DbContext {
-            tx: self.events_tx.clone(),
-            conn: ConnectionOrTx::Transaction(&tx),
-        };
-
-        match func(&db_context) {
-            Ok(val) => {
-                tx.commit()?;
-                Ok(val)
-            }
-            Err(e) => {
-                tx.rollback()?;
-                Err(e)
-            }
-        }
-    }
-}
-
-pub enum ConnectionOrTx<'a> {
-    Connection(PooledConnection<SqliteConnectionManager>),
-    Transaction(&'a Transaction<'a>),
-}
-
-impl<'a> ConnectionOrTx<'a> {
-    pub(crate) fn resolve(&self) -> &Connection {
-        match self {
-            ConnectionOrTx::Connection(c) => c,
-            ConnectionOrTx::Transaction(c) => c,
-        }
-    }
-
-    pub(crate) fn prepare(&self, sql: &str) -> rusqlite::Result<Statement<'_>> {
-        self.resolve().prepare(sql)
-    }
-
-    pub(crate) fn execute(&self, sql: &str, params: &[&dyn ToSql]) -> rusqlite::Result<usize> {
-        self.resolve().execute(sql, params)
-    }
-}
-
-pub struct DbContext<'a> {
-    pub(crate) tx: mpsc::Sender<ModelPayload>,
-    pub(crate) conn: ConnectionOrTx<'a>,
-}

src-tauri/yaak-models/src/models.rs

@@ -3,7 +3,7 @@ use crate::models::HttpRequestIden::{
     Authentication, AuthenticationType, Body, BodyType, CreatedAt, Description, FolderId, Headers,
     Method, Name, SortPriority, UpdatedAt, Url, UrlParameters, WorkspaceId,
 };
-use crate::queries_legacy::UpdateSource;
+use crate::util::UpdateSource;
 use chrono::{NaiveDateTime, Utc};
 use rusqlite::Row;
 use sea_query::{enum_def, IntoIden, IntoTableRef, SimpleExpr};

src-tauri/yaak-models/src/queries/batch.rs

@@ -1,8 +1,8 @@
 use crate::error::Result;
-use crate::manager::DbContext;
 use crate::models::{Environment, Folder, GrpcRequest, HttpRequest, WebsocketRequest, Workspace};
-use crate::queries_legacy::{BatchUpsertResult, UpdateSource};
+use crate::util::{BatchUpsertResult, UpdateSource};
 use log::info;
+use crate::db_context::DbContext;
 
 impl<'a> DbContext<'a> {
     pub fn batch_upsert(

src-tauri/yaak-models/src/queries/cookie_jars.rs

@@ -1,7 +1,7 @@
+use crate::db_context::DbContext;
 use crate::error::Result;
-use crate::manager::DbContext;
 use crate::models::{CookieJar, CookieJarIden};
-use crate::queries_legacy::UpdateSource;
+use crate::util::UpdateSource;
 
 impl<'a> DbContext<'a> {
     pub fn get_cookie_jar(&self, id: &str) -> Result<CookieJar> {

src-tauri/yaak-models/src/queries/environments.rs

@@ -1,11 +1,11 @@
 use crate::error::Result;
-use crate::manager::DbContext;
 use crate::models::{Environment, EnvironmentIden, UpsertModelInfo};
-use crate::queries_legacy::UpdateSource;
+use crate::util::UpdateSource;
 use log::info;
 use sea_query::ColumnRef::Asterisk;
 use sea_query::{Cond, Expr, Query, SqliteQueryBuilder};
 use sea_query_rusqlite::RusqliteBinder;
+use crate::db_context::DbContext;
 
 impl<'a> DbContext<'a> {
     pub fn get_environment(&self, id: &str) -> Result<Environment> {

src-tauri/yaak-models/src/queries/folders.rs

@@ -1,10 +1,10 @@
+use crate::db_context::DbContext;
 use crate::error::Result;
-use crate::manager::DbContext;
 use crate::models::{
     Folder, FolderIden, GrpcRequest, GrpcRequestIden, HttpRequest, HttpRequestIden,
     WebsocketRequest, WebsocketRequestIden,
 };
-use crate::queries_legacy::UpdateSource;
+use crate::util::UpdateSource;
 
 impl<'a> DbContext<'a> {
     pub fn get_folder(&self, id: &str) -> Result<Folder> {

src-tauri/yaak-models/src/queries/grpc_connections.rs

@@ -1,11 +1,11 @@
 use crate::error::Result;
-use crate::manager::DbContext;
 use crate::models::{GrpcConnection, GrpcConnectionIden, GrpcConnectionState};
-use crate::queries::base::MAX_HISTORY_ITEMS;
-use crate::queries_legacy::UpdateSource;
+use crate::util::UpdateSource;
 use log::debug;
 use sea_query::{Expr, Query, SqliteQueryBuilder};
 use sea_query_rusqlite::RusqliteBinder;
+use crate::db_context::DbContext;
+use crate::queries::MAX_HISTORY_ITEMS;
 
 impl<'a> DbContext<'a> {
     pub fn get_grpc_connection(&self, id: &str) -> Result<GrpcConnection> {

src-tauri/yaak-models/src/queries/grpc_events.rs

@@ -1,7 +1,7 @@
+use crate::db_context::DbContext;
 use crate::error::Result;
-use crate::manager::DbContext;
 use crate::models::{GrpcEvent, GrpcEventIden};
-use crate::queries_legacy::UpdateSource;
+use crate::util::UpdateSource;
 
 impl<'a> DbContext<'a> {
     pub fn get_grpc_events(&self, id: &str) -> Result<GrpcEvent> {

src-tauri/yaak-models/src/queries/grpc_requests.rs

@@ -1,11 +1,11 @@
+use crate::db_context::DbContext;
 use crate::error::Result;
-use crate::manager::DbContext;
 use crate::models::{GrpcRequest, GrpcRequestIden};
-use crate::queries_legacy::UpdateSource;
+use crate::util::UpdateSource;
 
 impl<'a> DbContext<'a> {
-    pub fn get_grpc_request(&self, id: &str) -> Result<Option<GrpcRequest>> {
-        self.find_optional(GrpcRequestIden::Id, id)
+    pub fn get_grpc_request(&self, id: &str) -> Result<GrpcRequest> {
+        self.find_one(GrpcRequestIden::Id, id)
     }
 
     pub fn list_grpc_requests(&self, workspace_id: &str) -> Result<Vec<GrpcRequest>> {
@@ -26,7 +26,7 @@ impl<'a> DbContext<'a> {
         id: &str,
         source: &UpdateSource,
     ) -> Result<GrpcRequest> {
-        let request = self.get_grpc_request(id)?.unwrap();
+        let request = self.get_grpc_request(id)?;
         self.delete_grpc_request(&request, source)
     }
 

src-tauri/yaak-models/src/queries/http_requests.rs

@@ -1,11 +1,11 @@
+use crate::db_context::DbContext;
 use crate::error::Result;
-use crate::manager::DbContext;
 use crate::models::{HttpRequest, HttpRequestIden};
-use crate::queries_legacy::UpdateSource;
+use crate::util::UpdateSource;
 
 impl<'a> DbContext<'a> {
-    pub fn get_http_request(&self, id: &str) -> Result<Option<HttpRequest>> {
-        self.find_optional(HttpRequestIden::Id, id)
+    pub fn get_http_request(&self, id: &str) -> Result<HttpRequest> {
+        self.find_one(HttpRequestIden::Id, id)
     }
 
     pub fn list_http_requests(&self, workspace_id: &str) -> Result<Vec<HttpRequest>> {
@@ -26,7 +26,7 @@ impl<'a> DbContext<'a> {
         id: &str,
         source: &UpdateSource,
     ) -> Result<HttpRequest> {
-        let http_request = self.get_http_request(id)?.unwrap();
+        let http_request = self.get_http_request(id)?;
         self.delete_http_request(&http_request, source)
     }
 

src-tauri/yaak-models/src/queries/http_responses.rs

@@ -1,12 +1,12 @@
 use crate::error::Result;
-use crate::manager::DbContext;
 use crate::models::{HttpResponse, HttpResponseIden, HttpResponseState};
-use crate::queries::base::MAX_HISTORY_ITEMS;
-use crate::queries_legacy::UpdateSource;
+use crate::util::UpdateSource;
 use log::{debug, error};
 use sea_query::{Expr, Query, SqliteQueryBuilder};
 use sea_query_rusqlite::RusqliteBinder;
 use std::fs;
+use crate::db_context::DbContext;
+use crate::queries::MAX_HISTORY_ITEMS;
 
 impl<'a> DbContext<'a> {
     pub fn get_http_response(&self, id: &str) -> Result<HttpResponse> {

src-tauri/yaak-models/src/queries/key_values.rs

@@ -1,11 +1,11 @@
 use crate::error::Result;
-use crate::manager::DbContext;
 use crate::models::{KeyValue, KeyValueIden};
-use crate::queries_legacy::{ModelChangeEvent, ModelPayload, UpdateSource};
+use crate::util::{ModelChangeEvent, ModelPayload, UpdateSource};
 use log::error;
 use sea_query::Keyword::CurrentTimestamp;
 use sea_query::{Asterisk, Cond, Expr, OnConflict, Query, SqliteQueryBuilder};
 use sea_query_rusqlite::RusqliteBinder;
+use crate::db_context::DbContext;
 
 impl<'a> DbContext<'a> {
     pub fn list_key_values_raw(&self) -> Result<Vec<KeyValue>> {

src-tauri/yaak-models/src/queries/mod.rs

@@ -1,4 +1,3 @@
-mod base;
 mod batch;
 mod cookie_jars;
 mod environments;
@@ -18,3 +17,5 @@ mod websocket_events;
 mod websocket_requests;
 mod workspace_metas;
 mod workspaces;
+
+const MAX_HISTORY_ITEMS: usize = 20;

src-tauri/yaak-models/src/queries/plugin_key_values.rs

@@ -1,5 +1,5 @@
+use crate::db_context::DbContext;
 use crate::error::Result;
-use crate::manager::DbContext;
 use crate::models::{PluginKeyValue, PluginKeyValueIden};
 use sea_query::Keyword::CurrentTimestamp;
 use sea_query::{Asterisk, Cond, Expr, OnConflict, Query, SqliteQueryBuilder};

src-tauri/yaak-models/src/queries/plugins.rs

@@ -1,13 +1,13 @@
+use crate::db_context::DbContext;
 use crate::error::Result;
-use crate::manager::DbContext;
 use crate::models::{Plugin, PluginIden};
-use crate::queries_legacy::UpdateSource;
+use crate::util::UpdateSource;
 
 impl<'a> DbContext<'a> {
     pub fn get_plugin(&self, id: &str) -> Result<Plugin> {
         self.find_one(PluginIden::Id, id)
     }
-    
+
     pub fn list_plugins(&self) -> Result<Vec<Plugin>> {
         self.find_all()
     }
@@ -20,7 +20,7 @@ impl<'a> DbContext<'a> {
         let plugin = self.get_plugin(id)?;
         self.delete_plugin(&plugin, source)
     }
-    
+
     pub fn upsert_plugin(&self, plugin: &Plugin, source: &UpdateSource) -> Result<Plugin> {
         self.upsert(plugin, source)
     }

src-tauri/yaak-models/src/queries/settings.rs

@@ -1,22 +1,24 @@
+use crate::db_context::DbContext;
 use crate::error::Result;
-use crate::manager::DbContext;
 use crate::models::{Settings, SettingsIden};
-use crate::queries_legacy::UpdateSource;
+use crate::util::UpdateSource;
 
 impl<'a> DbContext<'a> {
-    pub fn get_or_create_settings(&self, source: &UpdateSource) -> Result<Settings> {
-        let id = "default";
-        if let Some(s) = self.find_optional::<Settings>(SettingsIden::Id, id)? {
-            return Ok(s);
+    pub fn get_or_create_settings(&self, source: &UpdateSource) -> Settings {
+        let id = "default".to_string();
+
+        if let Some(s) = self.find_optional::<Settings>(SettingsIden::Id, &id) {
+            return s;
         };
 
         self.upsert(
             &Settings {
-                id: id.to_string(),
+                id,
                 ..Default::default()
             },
             source,
         )
+        .expect("Failed to upsert settings")
     }
 
     pub fn upsert_settings(&self, settings: &Settings, source: &UpdateSource) -> Result<Settings> {

src-tauri/yaak-models/src/queries/sync_states.rs

@@ -1,10 +1,10 @@
 use crate::error::Result;
-use crate::manager::DbContext;
 use crate::models::{SyncState, SyncStateIden, UpsertModelInfo};
-use crate::queries_legacy::UpdateSource;
+use crate::util::UpdateSource;
 use sea_query::{Asterisk, Cond, Expr, Query, SqliteQueryBuilder};
 use sea_query_rusqlite::RusqliteBinder;
 use std::path::Path;
+use crate::db_context::DbContext;
 
 impl<'a> DbContext<'a> {
     pub fn get_sync_state(&self, id: &str) -> Result<SyncState> {

src-tauri/yaak-models/src/queries/websocket_connections.rs

@@ -1,8 +1,8 @@
+use crate::db_context::DbContext;
 use crate::error::Result;
-use crate::manager::DbContext;
 use crate::models::{WebsocketConnection, WebsocketConnectionIden, WebsocketConnectionState};
-use crate::queries::base::MAX_HISTORY_ITEMS;
-use crate::queries_legacy::UpdateSource;
+use crate::queries::MAX_HISTORY_ITEMS;
+use crate::util::UpdateSource;
 use log::debug;
 use sea_query::{Expr, Query, SqliteQueryBuilder};
 use sea_query_rusqlite::RusqliteBinder;

src-tauri/yaak-models/src/queries/websocket_events.rs

@@ -1,10 +1,10 @@
+use crate::db_context::DbContext;
 use crate::error::Result;
-use crate::manager::DbContext;
 use crate::models::{
     WebsocketEvent,
     WebsocketEventIden,
 };
-use crate::queries_legacy::UpdateSource;
+use crate::util::UpdateSource;
 
 impl<'a> DbContext<'a> {
     pub fn get_websocket_event(&self, id: &str) -> Result<WebsocketEvent> {

src-tauri/yaak-models/src/queries/websocket_requests.rs

@@ -1,11 +1,11 @@
+use crate::db_context::DbContext;
 use crate::error::Result;
-use crate::manager::DbContext;
 use crate::models::{WebsocketRequest, WebsocketRequestIden};
-use crate::queries_legacy::UpdateSource;
+use crate::util::UpdateSource;
 
 impl<'a> DbContext<'a> {
-    pub fn get_websocket_request(&self, id: &str) -> Result<Option<WebsocketRequest>> {
-        self.find_optional(WebsocketRequestIden::Id, id)
+    pub fn get_websocket_request(&self, id: &str) -> Result<WebsocketRequest> {
+        self.find_one(WebsocketRequestIden::Id, id)
     }
 
     pub fn list_websocket_requests(&self, workspace_id: &str) -> Result<Vec<WebsocketRequest>> {
@@ -26,7 +26,7 @@ impl<'a> DbContext<'a> {
         id: &str,
         source: &UpdateSource,
     ) -> Result<WebsocketRequest> {
-        let request = self.get_websocket_request(id)?.unwrap();
+        let request = self.get_websocket_request(id)?;
         self.delete_websocket_request(&request, source)
     }
 

src-tauri/yaak-models/src/queries/workspace_metas.rs

@@ -1,10 +1,10 @@
+use crate::db_context::DbContext;
 use crate::error::Result;
-use crate::manager::DbContext;
 use crate::models::{Workspace, WorkspaceMeta, WorkspaceMetaIden};
-use crate::queries_legacy::UpdateSource;
+use crate::util::UpdateSource;
 
 impl<'a> DbContext<'a> {
-    pub fn get_workspace_meta(&self, workspace: &Workspace) -> Result<Option<WorkspaceMeta>> {
+    pub fn get_workspace_meta(&self, workspace: &Workspace) -> Option<WorkspaceMeta> {
         self.find_optional(WorkspaceMetaIden::WorkspaceId, &workspace.id)
     }
 
@@ -13,7 +13,7 @@ impl<'a> DbContext<'a> {
         workspace: &Workspace,
         source: &UpdateSource,
     ) -> Result<WorkspaceMeta> {
-        let workspace_meta = self.get_workspace_meta(workspace)?;
+        let workspace_meta = self.get_workspace_meta(workspace);
         if let Some(workspace_meta) = workspace_meta {
             return Ok(workspace_meta);
         }

src-tauri/yaak-models/src/queries/workspaces.rs

@@ -1,10 +1,10 @@
+use crate::db_context::DbContext;
 use crate::error::Result;
-use crate::manager::DbContext;
 use crate::models::{
     Folder, FolderIden, GrpcRequest, GrpcRequestIden, HttpRequest, HttpRequestIden,
     WebsocketRequest, WebsocketRequestIden, Workspace, WorkspaceIden,
 };
-use crate::queries_legacy::UpdateSource;
+use crate::util::UpdateSource;
 
 impl<'a> DbContext<'a> {
     pub fn get_workspace(&self, id: &str) -> Result<Workspace> {

src-tauri/yaak-models/src/query_manager.rs

@@ -0,0 +1,124 @@
+use crate::connection_or_tx::ConnectionOrTx;
+use crate::db_context::DbContext;
+use crate::error::Result;
+use crate::util::ModelPayload;
+use r2d2::Pool;
+use r2d2_sqlite::SqliteConnectionManager;
+use rusqlite::TransactionBehavior;
+use std::sync::{Arc, Mutex};
+use tauri::{Manager, Runtime};
+use tokio::sync::mpsc;
+
+pub trait QueryManagerExt<'a, R> {
+    fn db(&'a self) -> DbContext<'a>;
+    fn with_db<F, T>(&'a self, func: F) -> T
+    where
+        F: FnOnce(&DbContext) -> T;
+    fn with_tx<F, T>(&'a self, func: F) -> Result<T>
+    where
+        F: FnOnce(&DbContext) -> Result<T>;
+}
+
+impl<'a, R: Runtime, M: Manager<R>> QueryManagerExt<'a, R> for M {
+    fn db(&'a self) -> DbContext<'a> {
+        let qm = self.state::<QueryManager>();
+        qm.inner().connect_2()
+    }
+
+    fn with_db<F, T>(&'a self, func: F) -> T
+    where
+        F: FnOnce(&DbContext) -> T,
+    {
+        let qm = self.state::<QueryManager>();
+        qm.inner().with_conn(func)
+    }
+
+    fn with_tx<F, T>(&'a self, func: F) -> Result<T>
+    where
+        F: FnOnce(&DbContext) -> Result<T>,
+    {
+        let qm = self.state::<QueryManager>();
+        qm.inner().with_tx(func)
+    }
+}
+
+#[derive(Clone)]
+pub struct QueryManager {
+    pool: Arc<Mutex<Pool<SqliteConnectionManager>>>,
+    events_tx: mpsc::Sender<ModelPayload>,
+}
+
+impl QueryManager {
+    pub(crate) fn new(
+        pool: Pool<SqliteConnectionManager>,
+        events_tx: mpsc::Sender<ModelPayload>,
+    ) -> Self {
+        QueryManager {
+            pool: Arc::new(Mutex::new(pool)),
+            events_tx,
+        }
+    }
+
+    pub fn connect_2(&self) -> DbContext {
+        let conn = self
+            .pool
+            .lock()
+            .expect("Failed to gain lock on DB")
+            .get()
+            .expect("Failed to get a new DB connection from the pool");
+        DbContext {
+            tx: self.events_tx.clone(),
+            conn: ConnectionOrTx::Connection(conn),
+        }
+    }
+
+    pub fn with_conn<F, T>(&self, func: F) -> T
+    where
+        F: FnOnce(&DbContext) -> T,
+    {
+        let conn = self
+            .pool
+            .lock()
+            .expect("Failed to gain lock on DB for transaction")
+            .get()
+            .expect("Failed to get new DB connection from the pool");
+
+        let db_context = DbContext {
+            tx: self.events_tx.clone(),
+            conn: ConnectionOrTx::Connection(conn),
+        };
+
+        func(&db_context)
+    }
+
+    pub fn with_tx<F, T>(&self, func: F) -> Result<T>
+    where
+        F: FnOnce(&DbContext) -> Result<T>,
+    {
+        let mut conn = self
+            .pool
+            .lock()
+            .expect("Failed to gain lock on DB for transaction")
+            .get()
+            .expect("Failed to get new DB connection from the pool");
+        let tx = conn
+            .transaction_with_behavior(TransactionBehavior::Immediate)
+            .expect("Failed to start DB transaction");
+
+        let db_context = DbContext {
+            tx: self.events_tx.clone(),
+            conn: ConnectionOrTx::Transaction(&tx),
+        };
+
+        match func(&db_context) {
+            Ok(val) => {
+                tx.commit()?;
+                Ok(val)
+            }
+            Err(e) => {
+                tx.rollback()?;
+                Err(e)
+            }
+        }
+    }
+}

src-tauri/yaak-models/src/util.rs

@@ -1,12 +1,15 @@
 use crate::error::Result;
-use crate::manager::QueryManagerExt;
-use crate::models::{AnyModel, Environment, Folder, GrpcRequest, HttpRequest, ModelType, WebsocketRequest, Workspace, WorkspaceIden};
+use crate::models::{
+    AnyModel, Environment, Folder, GrpcRequest, HttpRequest, ModelType, WebsocketRequest,
+    Workspace, WorkspaceIden,
+};
 use chrono::{NaiveDateTime, Utc};
 use log::warn;
 use nanoid::nanoid;
 use serde::{Deserialize, Serialize};
 use tauri::{AppHandle, Listener, Runtime, WebviewWindow};
 use ts_rs::TS;
+use crate::query_manager::QueryManagerExt;
 
 pub fn generate_model_id(model: ModelType) -> String {
     let id = generate_id();
@@ -131,7 +134,7 @@ pub async fn get_workspace_export_resources<R: Runtime>(
         },
     };
 
-    let db = app_handle.queries().connect().await?;
+    let db = app_handle.db();
     for workspace_id in workspace_ids {
         data.resources.workspaces.push(db.find_one(WorkspaceIden::Id, workspace_id)?);
         data.resources.environments.append(&mut db.list_environments(workspace_id)?);