tusk/src-tauri/src/lib.rs

mod commands;
mod error;
mod mcp;
mod models;
mod state;
mod utils;

use models::settings::{AppSettings, DockerHost};
use state::AppState;
use std::sync::Arc;
use tauri::Manager;

pub fn run() {
    let shared_state = Arc::new(AppState::new());

    let _ = tauri::Builder::default()
        .plugin(tauri_plugin_shell::init())
        .plugin(tauri_plugin_dialog::init())
        .manage(shared_state)
        .setup(|app| {
            let state = app.state::<Arc<AppState>>().inner().clone();
            let data_dir = app
                .path()
                .app_data_dir()
                .map_err(|e| Box::new(e) as Box<dyn std::error::Error>)?;
            let connections_path = data_dir.join("connections.json");

            // Read app settings
            let settings_path = data_dir.join("app_settings.json");

            let settings = if settings_path.exists() {
                std::fs::read_to_string(&settings_path)
                    .ok()
                    .and_then(|data| serde_json::from_str::<AppSettings>(&data).ok())
                    .unwrap_or_default()
            } else {
                AppSettings::default()
            };

            // Apply docker host from settings
            let docker_host = match settings.docker.host {
                DockerHost::Remote => settings.docker.remote_url.clone(),
                DockerHost::Local => None,
            };

            let mcp_enabled = settings.mcp.enabled;
            let mcp_port = settings.mcp.port;

            // Set docker host synchronously (state is fresh, no contention)
            let state_for_setup = state.clone();
            tauri::async_runtime::block_on(async {
                *state_for_setup.docker_host.write().await = docker_host;
            });

            if mcp_enabled {
                let shutdown_rx = state.mcp_shutdown_tx.subscribe();
                let mcp_state = state.clone();
                tauri::async_runtime::spawn(async move {
                    *mcp_state.mcp_running.write().await = true;
                    if let Err(e) = mcp::start_mcp_server(
                        mcp_state.clone(),
                        connections_path,
                        mcp_port,
                        shutdown_rx,
                    )
                    .await
                    {
                        log::error!("MCP server error: {}", e);
                    }
                    *mcp_state.mcp_running.write().await = false;
                });
            }

            Ok(())
        })
        .invoke_handler(tauri::generate_handler![
            // connections
            commands::connections::get_connections,
            commands::connections::save_connection,
            commands::connections::delete_connection,
            commands::connections::test_connection,
            commands::connections::connect,
            commands::connections::switch_database,
            commands::connections::disconnect,
            commands::connections::set_read_only,
            commands::connections::get_read_only,
            commands::connections::get_db_flavor,
            // queries
            commands::queries::execute_query,
            // schema
            commands::schema::list_databases,
            commands::schema::list_schemas,
            commands::schema::list_tables,
            commands::schema::list_views,
            commands::schema::list_functions,
            commands::schema::list_indexes,
            commands::schema::list_sequences,
            commands::schema::get_table_columns,
            commands::schema::get_table_constraints,
            commands::schema::get_table_indexes,
            commands::schema::get_completion_schema,
            commands::schema::get_column_details,
            commands::schema::get_table_triggers,
            commands::schema::get_schema_erd,
            // data
            commands::data::get_table_data,
            commands::data::update_row,
            commands::data::insert_row,
            commands::data::delete_rows,
            // export
            commands::export::export_csv,
            commands::export::export_json,
            // management
            commands::management::get_database_info,
            commands::management::create_database,
            commands::management::drop_database,
            commands::management::list_roles,
            commands::management::create_role,
            commands::management::alter_role,
            commands::management::drop_role,
            commands::management::get_table_privileges,
            commands::management::grant_revoke,
            commands::management::manage_role_membership,
            commands::management::list_sessions,
            commands::management::cancel_query,
            commands::management::terminate_backend,
            // history
            commands::history::add_history_entry,
            commands::history::get_history,
            commands::history::clear_history,
            // saved queries
            commands::saved_queries::list_saved_queries,
            commands::saved_queries::save_query,
            commands::saved_queries::delete_saved_query,
            // ai
            commands::ai::get_ai_settings,
            commands::ai::save_ai_settings,
            commands::ai::list_ollama_models,
            commands::ai::generate_sql,
            commands::ai::explain_sql,
            commands::ai::fix_sql_error,
            commands::ai::generate_validation_sql,
            commands::ai::run_validation_rule,
            commands::ai::suggest_validation_rules,
            commands::ai::generate_test_data_preview,
            commands::ai::insert_generated_data,
            commands::ai::get_index_advisor_report,
            commands::ai::apply_index_recommendation,
            // snapshot
            commands::snapshot::create_snapshot,
            commands::snapshot::restore_snapshot,
            commands::snapshot::list_snapshots,
            commands::snapshot::read_snapshot_metadata,
            // lookup
            commands::lookup::entity_lookup,
            // docker
            commands::docker::check_docker,
            commands::docker::list_tusk_containers,
            commands::docker::clone_to_docker,
            commands::docker::start_container,
            commands::docker::stop_container,
            commands::docker::remove_container,
            // settings
            commands::settings::get_app_settings,
            commands::settings::save_app_settings,
            commands::settings::get_mcp_status,
        ])
        .run(tauri::generate_context!())
        .inspect_err(|e| {
            log::error!("Tauri application error: {}", e);
        });
}