//! Tauri-specific extensions for yaak-models. //! //! This module provides the Tauri plugin initialization and extension traits //! that allow accessing QueryManager and BlobManager from Tauri's Manager types. use tauri::plugin::TauriPlugin; use tauri::{Emitter, Manager, Runtime, State}; use tauri_plugin_dialog::{DialogExt, MessageDialogKind}; use yaak_models::blob_manager::BlobManager; use yaak_models::db_context::DbContext; use yaak_models::error::Result; use yaak_models::models::{AnyModel, GraphQlIntrospection, GrpcEvent, Settings, WebsocketEvent}; use yaak_models::query_manager::QueryManager; use yaak_models::util::UpdateSource; /// Extension trait for accessing the QueryManager from Tauri Manager types. pub trait QueryManagerExt<'a, R> { fn db_manager(&'a self) -> State<'a, QueryManager>; fn db(&'a self) -> DbContext<'a>; fn with_tx(&'a self, func: F) -> Result where F: FnOnce(&DbContext) -> Result; } impl<'a, R: Runtime, M: Manager> QueryManagerExt<'a, R> for M { fn db_manager(&'a self) -> State<'a, QueryManager> { self.state::() } fn db(&'a self) -> DbContext<'a> { let qm = self.state::(); qm.inner().connect() } fn with_tx(&'a self, func: F) -> Result where F: FnOnce(&DbContext) -> Result, { let qm = self.state::(); qm.inner().with_tx(func) } } /// Extension trait for accessing the BlobManager from Tauri Manager types. pub trait BlobManagerExt<'a, R> { fn blob_manager(&'a self) -> State<'a, BlobManager>; fn blobs(&'a self) -> yaak_models::blob_manager::BlobContext; } impl<'a, R: Runtime, M: Manager> BlobManagerExt<'a, R> for M { fn blob_manager(&'a self) -> State<'a, BlobManager> { self.state::() } fn blobs(&'a self) -> yaak_models::blob_manager::BlobContext { let manager = self.state::(); manager.inner().connect() } } // Commands for yaak-models use tauri::WebviewWindow; #[tauri::command] pub(crate) fn models_upsert( window: WebviewWindow, model: AnyModel, ) -> Result { use yaak_models::error::Error::GenericError; let db = window.db(); let blobs = window.blob_manager(); let source = &UpdateSource::from_window_label(window.label()); let id = match model { AnyModel::CookieJar(m) => db.upsert_cookie_jar(&m, source)?.id, AnyModel::Environment(m) => db.upsert_environment(&m, source)?.id, AnyModel::Folder(m) => db.upsert_folder(&m, source)?.id, AnyModel::GrpcRequest(m) => db.upsert_grpc_request(&m, source)?.id, AnyModel::HttpRequest(m) => db.upsert_http_request(&m, source)?.id, AnyModel::HttpResponse(m) => db.upsert_http_response(&m, source, &blobs)?.id, AnyModel::KeyValue(m) => db.upsert_key_value(&m, source)?.id, AnyModel::Plugin(m) => db.upsert_plugin(&m, source)?.id, AnyModel::Settings(m) => db.upsert_settings(&m, source)?.id, AnyModel::WebsocketRequest(m) => db.upsert_websocket_request(&m, source)?.id, AnyModel::Workspace(m) => db.upsert_workspace(&m, source)?.id, AnyModel::WorkspaceMeta(m) => db.upsert_workspace_meta(&m, source)?.id, a => return Err(GenericError(format!("Cannot upsert AnyModel {a:?})"))), }; Ok(id) } #[tauri::command] pub(crate) fn models_delete( window: WebviewWindow, model: AnyModel, ) -> Result { use yaak_models::error::Error::GenericError; let blobs = window.blob_manager(); // Use transaction for deletions because it might recurse window.with_tx(|tx| { let source = &UpdateSource::from_window_label(window.label()); let id = match model { AnyModel::CookieJar(m) => tx.delete_cookie_jar(&m, source)?.id, AnyModel::Environment(m) => tx.delete_environment(&m, source)?.id, AnyModel::Folder(m) => tx.delete_folder(&m, source)?.id, AnyModel::GrpcConnection(m) => tx.delete_grpc_connection(&m, source)?.id, AnyModel::GrpcRequest(m) => tx.delete_grpc_request(&m, source)?.id, AnyModel::HttpRequest(m) => tx.delete_http_request(&m, source)?.id, AnyModel::HttpResponse(m) => tx.delete_http_response(&m, source, &blobs)?.id, AnyModel::Plugin(m) => tx.delete_plugin(&m, source)?.id, AnyModel::WebsocketConnection(m) => tx.delete_websocket_connection(&m, source)?.id, AnyModel::WebsocketRequest(m) => tx.delete_websocket_request(&m, source)?.id, AnyModel::Workspace(m) => tx.delete_workspace(&m, source)?.id, a => return Err(GenericError(format!("Cannot delete AnyModel {a:?})"))), }; Ok(id) }) } #[tauri::command] pub(crate) fn models_duplicate( window: WebviewWindow, model: AnyModel, ) -> Result { use yaak_models::error::Error::GenericError; // Use transaction for duplications because it might recurse window.with_tx(|tx| { let source = &UpdateSource::from_window_label(window.label()); let id = match model { AnyModel::Environment(m) => tx.duplicate_environment(&m, source)?.id, AnyModel::Folder(m) => tx.duplicate_folder(&m, source)?.id, AnyModel::GrpcRequest(m) => tx.duplicate_grpc_request(&m, source)?.id, AnyModel::HttpRequest(m) => tx.duplicate_http_request(&m, source)?.id, AnyModel::WebsocketRequest(m) => tx.duplicate_websocket_request(&m, source)?.id, a => return Err(GenericError(format!("Cannot duplicate AnyModel {a:?})"))), }; Ok(id) }) } #[tauri::command] pub(crate) fn models_websocket_events( app_handle: tauri::AppHandle, connection_id: &str, ) -> Result> { Ok(app_handle.db().list_websocket_events(connection_id)?) } #[tauri::command] pub(crate) fn models_grpc_events( app_handle: tauri::AppHandle, connection_id: &str, ) -> Result> { Ok(app_handle.db().list_grpc_events(connection_id)?) } #[tauri::command] pub(crate) fn models_get_settings(app_handle: tauri::AppHandle) -> Result { Ok(app_handle.db().get_settings()) } #[tauri::command] pub(crate) fn models_get_graphql_introspection( app_handle: tauri::AppHandle, request_id: &str, ) -> Result> { Ok(app_handle.db().get_graphql_introspection(request_id)) } #[tauri::command] pub(crate) fn models_upsert_graphql_introspection( app_handle: tauri::AppHandle, request_id: &str, workspace_id: &str, content: Option, window: WebviewWindow, ) -> Result { let source = UpdateSource::from_window_label(window.label()); Ok(app_handle.db().upsert_graphql_introspection(workspace_id, request_id, content, &source)?) } #[tauri::command] pub(crate) fn models_workspace_models( window: WebviewWindow, workspace_id: Option<&str>, ) -> Result { let db = window.db(); let mut l: Vec = Vec::new(); // Add the settings l.push(db.get_settings().into()); // Add global models l.append(&mut db.list_workspaces()?.into_iter().map(Into::into).collect()); l.append(&mut db.list_key_values()?.into_iter().map(Into::into).collect()); l.append(&mut db.list_plugins()?.into_iter().map(Into::into).collect()); // Add the workspace children if let Some(wid) = workspace_id { l.append(&mut db.list_cookie_jars(wid)?.into_iter().map(Into::into).collect()); l.append(&mut db.list_environments_ensure_base(wid)?.into_iter().map(Into::into).collect()); l.append(&mut db.list_folders(wid)?.into_iter().map(Into::into).collect()); l.append(&mut db.list_grpc_connections(wid)?.into_iter().map(Into::into).collect()); l.append(&mut db.list_grpc_requests(wid)?.into_iter().map(Into::into).collect()); l.append(&mut db.list_http_requests(wid)?.into_iter().map(Into::into).collect()); l.append(&mut db.list_http_responses(wid, None)?.into_iter().map(Into::into).collect()); l.append(&mut db.list_websocket_connections(wid)?.into_iter().map(Into::into).collect()); l.append(&mut db.list_websocket_requests(wid)?.into_iter().map(Into::into).collect()); l.append(&mut db.list_workspace_metas(wid)?.into_iter().map(Into::into).collect()); } let j = serde_json::to_string(&l)?; Ok(escape_str_for_webview(&j)) } fn escape_str_for_webview(input: &str) -> String { input .chars() .map(|c| { let code = c as u32; // ASCII if code <= 0x7F { c.to_string() // BMP characters encoded normally } else if code < 0xFFFF { format!("\\u{:04X}", code) // Beyond BMP encoded a surrogate pairs } else { let high = ((code - 0x10000) >> 10) + 0xD800; let low = ((code - 0x10000) & 0x3FF) + 0xDC00; format!("\\u{:04X}\\u{:04X}", high, low) } }) .collect() } /// Initialize database managers as a plugin (for initialization order). /// Commands are in the main invoke_handler. /// This must be registered before other plugins that depend on the database. pub fn init() -> TauriPlugin { tauri::plugin::Builder::new("yaak-models-db") .setup(|app_handle, _api| { let app_path = app_handle.path().app_data_dir().unwrap(); let db_path = app_path.join("db.sqlite"); let blob_path = app_path.join("blobs.sqlite"); let (query_manager, blob_manager, rx) = match yaak_models::init_standalone(&db_path, &blob_path) { Ok(result) => result, Err(e) => { app_handle .dialog() .message(e.to_string()) .kind(MessageDialogKind::Error) .blocking_show(); return Err(Box::from(e.to_string())); } }; app_handle.manage(query_manager); app_handle.manage(blob_manager); // Forward model change events to the frontend let app_handle = app_handle.clone(); tauri::async_runtime::spawn(async move { for payload in rx { app_handle.emit("model_write", payload).unwrap(); } }); Ok(()) }) .build() }