use crate::error::{TuskError, TuskResult}; use crate::models::management::*; use crate::state::AppState; use crate::utils::escape_ident; use sqlx::Row; use tauri::State; #[tauri::command] pub async fn get_database_info( state: State<'_, AppState>, connection_id: String, ) -> TuskResult> { let pools = state.pools.read().await; let pool = pools .get(&connection_id) .ok_or(TuskError::NotConnected(connection_id))?; let rows = sqlx::query( "SELECT d.datname, \ pg_catalog.pg_get_userbyid(d.datdba) AS owner, \ pg_catalog.pg_encoding_to_char(d.encoding) AS encoding, \ d.datcollate, \ d.datctype, \ COALESCE(t.spcname, 'pg_default') AS tablespace, \ d.datconnlimit, \ pg_catalog.pg_size_pretty(pg_catalog.pg_database_size(d.datname)) AS size, \ pg_catalog.shobj_description(d.oid, 'pg_database') AS description \ FROM pg_catalog.pg_database d \ LEFT JOIN pg_catalog.pg_tablespace t ON d.dattablespace = t.oid \ WHERE NOT d.datistemplate \ ORDER BY d.datname", ) .fetch_all(pool) .await .map_err(TuskError::Database)?; let databases = rows .iter() .map(|row| DatabaseInfo { name: row.get("datname"), owner: row.get("owner"), encoding: row.get("encoding"), collation: row.get("datcollate"), ctype: row.get("datctype"), tablespace: row.get("tablespace"), connection_limit: row.get("datconnlimit"), size: row.get("size"), description: row.get("description"), }) .collect(); Ok(databases) } #[tauri::command] pub async fn create_database( state: State<'_, AppState>, connection_id: String, params: CreateDatabaseParams, ) -> TuskResult<()> { if state.is_read_only(&connection_id).await { return Err(TuskError::ReadOnly); } let pools = state.pools.read().await; let pool = pools .get(&connection_id) .ok_or(TuskError::NotConnected(connection_id))?; let mut sql = format!("CREATE DATABASE {}", escape_ident(¶ms.name)); if let Some(ref owner) = params.owner { sql.push_str(&format!(" OWNER {}", escape_ident(owner))); } if let Some(ref template) = params.template { sql.push_str(&format!(" TEMPLATE {}", escape_ident(template))); } if let Some(ref encoding) = params.encoding { sql.push_str(&format!(" ENCODING '{}'", encoding.replace('\'', "''"))); } if let Some(ref tablespace) = params.tablespace { sql.push_str(&format!(" TABLESPACE {}", escape_ident(tablespace))); } if let Some(limit) = params.connection_limit { sql.push_str(&format!(" CONNECTION LIMIT {}", limit)); } sqlx::query(&sql) .execute(pool) .await .map_err(TuskError::Database)?; Ok(()) } #[tauri::command] pub async fn drop_database( state: State<'_, AppState>, connection_id: String, name: String, ) -> TuskResult<()> { if state.is_read_only(&connection_id).await { return Err(TuskError::ReadOnly); } let pools = state.pools.read().await; let pool = pools .get(&connection_id) .ok_or(TuskError::NotConnected(connection_id))?; // Terminate active connections to the target database let terminate_sql = format!( "SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = '{}' AND pid <> pg_backend_pid()", name.replace('\'', "''") ); sqlx::query(&terminate_sql) .execute(pool) .await .map_err(TuskError::Database)?; let drop_sql = format!("DROP DATABASE {}", escape_ident(&name)); sqlx::query(&drop_sql) .execute(pool) .await .map_err(TuskError::Database)?; Ok(()) } #[tauri::command] pub async fn list_roles( state: State<'_, AppState>, connection_id: String, ) -> TuskResult> { let pools = state.pools.read().await; let pool = pools .get(&connection_id) .ok_or(TuskError::NotConnected(connection_id))?; let rows = sqlx::query( "SELECT r.rolname, \ r.rolsuper, \ r.rolcanlogin, \ r.rolcreatedb, \ r.rolcreaterole, \ r.rolinherit, \ r.rolreplication, \ r.rolconnlimit, \ r.rolpassword IS NOT NULL AS password_set, \ r.rolvaliduntil::text, \ COALESCE(( \ SELECT array_agg(g.rolname ORDER BY g.rolname) \ FROM pg_catalog.pg_auth_members m \ JOIN pg_catalog.pg_roles g ON m.roleid = g.oid \ WHERE m.member = r.oid \ ), ARRAY[]::text[]) AS member_of, \ COALESCE(( \ SELECT array_agg(m2.rolname ORDER BY m2.rolname) \ FROM pg_catalog.pg_auth_members am \ JOIN pg_catalog.pg_roles m2 ON am.member = m2.oid \ WHERE am.roleid = r.oid \ ), ARRAY[]::text[]) AS members, \ pg_catalog.shobj_description(r.oid, 'pg_authid') AS description \ FROM pg_catalog.pg_roles r \ WHERE r.rolname !~ '^pg_' \ ORDER BY r.rolname", ) .fetch_all(pool) .await .map_err(TuskError::Database)?; let roles = rows .iter() .map(|row| RoleInfo { name: row.get("rolname"), is_superuser: row.get("rolsuper"), can_login: row.get("rolcanlogin"), can_create_db: row.get("rolcreatedb"), can_create_role: row.get("rolcreaterole"), inherit: row.get("rolinherit"), is_replication: row.get("rolreplication"), connection_limit: row.get("rolconnlimit"), password_set: row.get("password_set"), valid_until: row.get("rolvaliduntil"), member_of: row.get("member_of"), members: row.get("members"), description: row.get("description"), }) .collect(); Ok(roles) } #[tauri::command] pub async fn create_role( state: State<'_, AppState>, connection_id: String, params: CreateRoleParams, ) -> TuskResult<()> { if state.is_read_only(&connection_id).await { return Err(TuskError::ReadOnly); } let pools = state.pools.read().await; let pool = pools .get(&connection_id) .ok_or(TuskError::NotConnected(connection_id))?; let mut sql = format!("CREATE ROLE {}", escape_ident(¶ms.name)); let mut options = Vec::new(); options.push(if params.login { "LOGIN" } else { "NOLOGIN" }); options.push(if params.superuser { "SUPERUSER" } else { "NOSUPERUSER" }); options.push(if params.createdb { "CREATEDB" } else { "NOCREATEDB" }); options.push(if params.createrole { "CREATEROLE" } else { "NOCREATEROLE" }); options.push(if params.inherit { "INHERIT" } else { "NOINHERIT" }); options.push(if params.replication { "REPLICATION" } else { "NOREPLICATION" }); if let Some(ref password) = params.password { options.push("PASSWORD"); // Will be appended separately sql.push_str(&format!(" {}", options.join(" "))); sql.push_str(&format!(" '{}'", password.replace('\'', "''"))); } else { sql.push_str(&format!(" {}", options.join(" "))); } if let Some(limit) = params.connection_limit { sql.push_str(&format!(" CONNECTION LIMIT {}", limit)); } if let Some(ref valid_until) = params.valid_until { sql.push_str(&format!( " VALID UNTIL '{}'", valid_until.replace('\'', "''") )); } if !params.in_roles.is_empty() { let roles: Vec = params.in_roles.iter().map(|r| escape_ident(r)).collect(); sql.push_str(&format!(" IN ROLE {}", roles.join(", "))); } sqlx::query(&sql) .execute(pool) .await .map_err(TuskError::Database)?; Ok(()) } #[tauri::command] pub async fn alter_role( state: State<'_, AppState>, connection_id: String, params: AlterRoleParams, ) -> TuskResult<()> { if state.is_read_only(&connection_id).await { return Err(TuskError::ReadOnly); } let pools = state.pools.read().await; let pool = pools .get(&connection_id) .ok_or(TuskError::NotConnected(connection_id))?; let mut options = Vec::new(); if let Some(login) = params.login { options.push(if login { "LOGIN".to_string() } else { "NOLOGIN".to_string() }); } if let Some(superuser) = params.superuser { options.push(if superuser { "SUPERUSER".to_string() } else { "NOSUPERUSER".to_string() }); } if let Some(createdb) = params.createdb { options.push(if createdb { "CREATEDB".to_string() } else { "NOCREATEDB".to_string() }); } if let Some(createrole) = params.createrole { options.push(if createrole { "CREATEROLE".to_string() } else { "NOCREATEROLE".to_string() }); } if let Some(inherit) = params.inherit { options.push(if inherit { "INHERIT".to_string() } else { "NOINHERIT".to_string() }); } if let Some(replication) = params.replication { options.push(if replication { "REPLICATION".to_string() } else { "NOREPLICATION".to_string() }); } if let Some(ref password) = params.password { options.push(format!("PASSWORD '{}'", password.replace('\'', "''"))); } if let Some(limit) = params.connection_limit { options.push(format!("CONNECTION LIMIT {}", limit)); } if let Some(ref valid_until) = params.valid_until { options.push(format!( "VALID UNTIL '{}'", valid_until.replace('\'', "''") )); } if !options.is_empty() { let sql = format!( "ALTER ROLE {} {}", escape_ident(¶ms.name), options.join(" ") ); sqlx::query(&sql) .execute(pool) .await .map_err(TuskError::Database)?; } if let Some(ref new_name) = params.rename_to { let sql = format!( "ALTER ROLE {} RENAME TO {}", escape_ident(¶ms.name), escape_ident(new_name) ); sqlx::query(&sql) .execute(pool) .await .map_err(TuskError::Database)?; } Ok(()) } #[tauri::command] pub async fn drop_role( state: State<'_, AppState>, connection_id: String, name: String, ) -> TuskResult<()> { if state.is_read_only(&connection_id).await { return Err(TuskError::ReadOnly); } let pools = state.pools.read().await; let pool = pools .get(&connection_id) .ok_or(TuskError::NotConnected(connection_id))?; let sql = format!("DROP ROLE {}", escape_ident(&name)); sqlx::query(&sql) .execute(pool) .await .map_err(TuskError::Database)?; Ok(()) } #[tauri::command] pub async fn get_table_privileges( state: State<'_, AppState>, connection_id: String, schema: String, table: String, ) -> TuskResult> { let pools = state.pools.read().await; let pool = pools .get(&connection_id) .ok_or(TuskError::NotConnected(connection_id))?; let rows = sqlx::query( "SELECT grantee, table_schema, table_name, privilege_type, \ is_grantable = 'YES' AS is_grantable \ FROM information_schema.role_table_grants \ WHERE table_schema = $1 AND table_name = $2 \ ORDER BY grantee, privilege_type", ) .bind(&schema) .bind(&table) .fetch_all(pool) .await .map_err(TuskError::Database)?; let privileges = rows .iter() .map(|row| TablePrivilege { grantee: row.get("grantee"), table_schema: row.get("table_schema"), table_name: row.get("table_name"), privilege_type: row.get("privilege_type"), is_grantable: row.get("is_grantable"), }) .collect(); Ok(privileges) } #[tauri::command] pub async fn grant_revoke( state: State<'_, AppState>, connection_id: String, params: GrantRevokeParams, ) -> TuskResult<()> { if state.is_read_only(&connection_id).await { return Err(TuskError::ReadOnly); } let pools = state.pools.read().await; let pool = pools .get(&connection_id) .ok_or(TuskError::NotConnected(connection_id))?; let privs = params.privileges.join(", "); let object_type = params.object_type.to_uppercase(); let object_ref = escape_ident(¶ms.object_name); let role_ref = escape_ident(¶ms.role_name); let sql = if params.action.to_uppercase() == "GRANT" { let grant_option = if params.with_grant_option { " WITH GRANT OPTION" } else { "" }; format!( "GRANT {} ON {} {} TO {}{}", privs, object_type, object_ref, role_ref, grant_option ) } else { format!( "REVOKE {} ON {} {} FROM {}", privs, object_type, object_ref, role_ref ) }; sqlx::query(&sql) .execute(pool) .await .map_err(TuskError::Database)?; Ok(()) } #[tauri::command] pub async fn manage_role_membership( state: State<'_, AppState>, connection_id: String, params: RoleMembershipParams, ) -> TuskResult<()> { if state.is_read_only(&connection_id).await { return Err(TuskError::ReadOnly); } let pools = state.pools.read().await; let pool = pools .get(&connection_id) .ok_or(TuskError::NotConnected(connection_id))?; let role_ref = escape_ident(¶ms.role_name); let member_ref = escape_ident(¶ms.member_name); let sql = if params.action.to_uppercase() == "GRANT" { format!("GRANT {} TO {}", role_ref, member_ref) } else { format!("REVOKE {} FROM {}", role_ref, member_ref) }; sqlx::query(&sql) .execute(pool) .await .map_err(TuskError::Database)?; Ok(()) }